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
2.5 KiB
2.5 KiB
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
- Study CUTLASS FMHA Python reference (fmha.py on B200) for TMEM output layout
- Check if the 4× factor is a known issue with single-CTA MMA
- Try M=64 (2 warpgroups) — should give 2× if warpgroup count is the cause
- Look at gau-nernst's GEMM example to see how he reads the MMA output
- 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