biondizzle
32f7fa7bce
- MMA produces exactly 4× scalar reference for all output values - SMEM data verified correct, descriptor values correct - 4× persists across different N, warp counts, padding - TMEM multi-store bug documented (16x256b.x1 crashes on 2nd store) - Layout D read with 32x32b.x8 works - Next: study CUTLASS FMHA TMEM output layout to fix 4× factor
53 lines
2.5 KiB
Markdown
53 lines
2.5 KiB
Markdown
# CURRENT ISSUE: UMMA QK GEMM — 4× Scaling Bug
|
||
|
||
## What's working
|
||
- UMMA SMEM descriptors: K-major NONE, LBO=BLOCK_MN*16, SBO=128
|
||
- SMEM canonical layout: column-major interleaving of 8×8 BF16 core matrices
|
||
- Q and K SMEM data verified EXACT match with originals
|
||
- tcgen05.mma produces non-zero output — descriptor and data layout are valid
|
||
- TMEM Layout D read with tcgen05.ld.32x32b.x8 works (no crash)
|
||
- TMEM alloc/dealloc works
|
||
|
||
## The 4× Bug
|
||
MMA output is exactly 4× the scalar reference for ALL output values.
|
||
- S[0,0] MMA = 0.1529, scalar = 0.0382, ratio = 4.0000
|
||
- Persists with different N in idesc (8, 32, 128)
|
||
- Persists with 4 warp leaders calling MMA (vs 1 thread)
|
||
- Persists with 8KB zero padding between Q and K in SMEM
|
||
|
||
### Root cause hypothesis
|
||
The MMA with cta_group::1 and M=128 uses 4 "warpgroups" internally (Layout D).
|
||
The TMEM output is written in a format where each warpgroup contributes to
|
||
different rows. When we read with 32x32b.x8 (warp 0, rows 0-31), we get
|
||
the correct S[0,0] but multiplied by 4 because the MMA accumulates contributions
|
||
from all 4 warpgroups into the same TMEM columns.
|
||
|
||
Alternatively: the TMEM Layout D has a specific column mapping that we're not
|
||
accounting for. The MMA output columns might not correspond 1:1 with the
|
||
attention score columns.
|
||
|
||
### How to fix
|
||
1. Study CUTLASS FMHA Python reference (fmha.py on B200) for TMEM output layout
|
||
2. Check if the 4× factor is a known issue with single-CTA MMA
|
||
3. Try M=64 (2 warpgroups) — should give 2× if warpgroup count is the cause
|
||
4. Look at gau-nernst's GEMM example to see how he reads the MMA output
|
||
5. Check if the MMA output needs to be divided by the number of warpgroups
|
||
|
||
## TMEM multi-store bug
|
||
Calling tcgen05.st.16x256b.x1.b32 more than once causes "misaligned address".
|
||
- Single store: works
|
||
- 2+ stores: crash (even with fence+sync between them)
|
||
- CUTLASS uses different TMEM store atoms (St32x32bOp)
|
||
- Need to investigate: is 16x256b.x1 not meant for multiple stores?
|
||
|
||
## Files
|
||
- `dsv4/kernels/attention/fmha_umma_desc.cuh` — descriptor construction, write_smem_*
|
||
- `tests/unit/test_umma_qk.cu` — UMMA QK GEMM test (HD=16, SK=128)
|
||
- `tests/unit/test_tmem_cols.cu` — TMEM multi-store debug test
|
||
|
||
## Key references
|
||
- gau-nernst tcgen05 tutorial: https://gau-nernst.github.io/tcgen05/
|
||
- CUTLASS SM100 UMMA: include/cute/arch/mma_sm100_umma.hpp
|
||
- CUTLASS InstrDescriptor: include/cute/arch/mma_sm100_desc.hpp
|
||
- CUTLASS FMHA reference on B200: /root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py
|