nvfp4-megamoe-kernel/CURRENT_ISSUE.md

CURRENT_ISSUE.md — PV GEMM for Prefill

Status: HD=16 ✅ (cos 0.9997), HD=64 🚧 (cos 0.931, ~0.4% PV MMA error)

HD=16 — COMPLETE

Full pipeline: QK(SS, 1 K-tile) → softmax(TMEM→SMEM) → PV(SS, 8 K-tiles) → epilogue. Cosine 0.9997.

HD=64 — IN PROGRESS

• Pipeline runs end-to-end but PV MMA with BLOCK_MN_B=64 has ~0.4% systematic error
• Register-math PV with same QK+softmax output matches reference exactly → QK+softmax is correct
• The 0.4% error is specifically in the PV SS MMA with V=(64,16) BLOCK_MN=64
• Must opt into >48KB shared memory: cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)
• Alternative approach to try: BLOCK_MN_B=16 with 4 PV N-tiles per K-tile (avoids (64,16) V layout)
• QK MMA scale is 1.0 (NOT 0.5) — confirmed by comparing with test_fmha_hd64.cu register-math PV

Key findings (all from today):

1. tcgen05.mma TS (TMEM A) NOT usable — 32x32b store ≠ Layout A. Use SS MMA with SMEM-P instead.
2. V canonical layout bug — MN/K axes were swapped. Fix: g_mn=d/8, g_k=lr/8, llr=d%8, lc=r%8
3. Per-K-tile P fill — (128,128) canonical with K-tile offsets has accumulation bug. Use single (128,16) buffer.
4. SMEM >48KB needs opt-in on SM100
5. PV SS MMA scale: ~1.0 for both BLOCK_MN_B=16 and BLOCK_MN_B=64

Next steps:

1. Fix HD=64 PV MMA: try BLOCK_MN_B=16 with 4 N-tiles, or debug (64,16) canonical layout
2. Extend to HD=128, HD=256 — the per-K-tile approach scales naturally
3. Prefill T>1 — fill all 128 rows of sPk
4. Multi-head support — per-head launch or head-packed M
5. Production kernel — integrate into fmha_sm100.cuh