nvfp4-megamoe-kernel/CURRENT_ISSUE.md

CURRENT ISSUE: UMMA FMHA — Full Pipeline + Production

What's working ✅

• Full UMMA FMHA HD=64 pipeline: QK → softmax → PV → output
  • QK GEMM: UMMA SS, multi-K-tile accumulate, cos 0.999999
  • Softmax: TMEM read → max/exp/sum → TMEM write, max rel err 0.003
  • PV: register math (O[d] = Σ P[0,j] × V[d,j]), decode only
  • Overall: cosine 0.999998!
• UMMA QK GEMM at HD=16: Row 0 matches scalar with ZERO error
• TMEM one-way epilogue: cos 0.999999 at hd=64, cos 0.999998 at hd=128
• 32x32b.x8 TMEM stores work for P write-back (no 16x256b crash)

Next steps

1. PV GEMM via tcgen05.mma TS: For prefill (T>1), need UMMA-based PV

   • tcgen05.mma TS crashed with "illegal memory access" in initial tests
   • Need to debug the TMEM A operand addressing for PV
   • Decode path (T=1) works with register math — PV GEMM is only for prefill

2. HD=128/256: Extend multi-K-tile QK to larger head dims

   • HD=128: 8 K-tiles, separate SMEM per K-tile
   • HD=256: 16 K-tiles — SMEM budget needs checking

3. Multi-head launch: Per-head kernel dispatch (128 heads for Pro)

   • Current test: single head
   • Production: grid=(1, n_h, batch) or Python loop

4. Multi-KV-tile: s_k > 128 requires multiple attention tiles + KV merge

   • Same D5 merge formula: O = Σ exp(lse_i) · O_i / Σ exp(lse_i)

5. Production kernel: Integrate UMMA pipeline into fmha_sm100.cuh

   • Replace SMEM scalar attention with UMMA QK GEMM
   • Keep register-math PV for decode
   • Add PV GEMM path for prefill

Key lessons from this session

• Source stride mismatch: write_k_to_smem<128,16> template reads k[i] with stride=16, but actual K has stride=64. Must use separate SMEM per K-tile with manual writes.
• Offset descriptors DON'T WORK for (128,64) SMEM with K-tile offsets. Even/odd column corruption. Use separate (128,16) SMEM buffers.
• tcgen05.mma TS crashes with illegal memory access — likely TMEM A operand addressing issue. Debug needed.
• tcgen05.fence::after_thread_sync is the correct fence after MMA, before TMEM read.
• MMA computes UNSCALED dot product — apply 1/sqrt(HD) in softmax.
• 32x32b.x8 TMEM stores work in loops (unlike 16x256b.x1 which crashes on 2nd call)
• s_vals scope: when merging softmax + PV, keep s_vals in same if (wid == 0) block
• Decode PV is simple: O[d] = Σ P[0,j] × V[d,j] — no UMMA needed for T=1

Files

• dsv4/kernels/attention/fmha_umma_desc.cuh — descriptors, SMEM layout, MMA wrappers
• dsv4/kernels/attention/fmha_common.cuh — BF16, TMEM ops, warp reductions
• dsv4/kernels/attention/fmha_epilogue_sm100.cuh — TMEM one-way epilogue (reference kernel)
• tests/unit/test_fmha_hd64.cu — FULL PIPELINE TEST (QK+softmax+PV)
• tests/unit/test_fmha_softmax.cu — softmax validation (HD=64)
• tests/unit/test_umma_qk.cu — QK GEMM (HD=16)
• tests/unit/test_umma_qk_hd64.cu — QK GEMM (HD=64, multi-K-tile)