nvfp4-megamoe-kernel/CURRENT_ISSUE.md

CURRENT_ISSUE.md — PV GEMM for Prefill

Status: ✅ PV via SS MMA WORKING (cosine 0.9997)

What we proved:

1. tcgen05.mma TS (TMEM A operand) CANNOT be used — the 32x32b store format doesn't match the TS MMA's A-fragment layout (Layout A). Even though the isolated test_mma_ts.cu works with uniform data, non-uniform data produces garbage because the TMEM column mapping differs.

2. PV via SS MMA with SMEM-P is the correct path — softmax writes P to SMEM in canonical (128, 16) layout, then PV uses SS MMA with A=P(SMEM) × B=V(SMEM) → C=O(TMEM). No TMEM layout issues because both operands are in SMEM with layouts we control.

3. Per-K-tile P fill — the (128, 128) canonical P with K-tile offsets has an accumulation bug (first 8 output values lose one K-tile). The fix: reuse a single (128, 16) sPk buffer, filling it from s_p_vals for each K-tile.

4. V canonical layout bug — the original V load code had the MN and K axes swapped in the canonical formula. d (HD dim = MN) was mapped to g_k and lr (sequence position = K) was mapped to g_mn. Fix: swap to g_mn = d/8, g_k = lr/8, llr = d%8, lc = lr%8. This didn't affect the constant-value tests.

5. SMEM allocation bug — the original test_fmha_smem_p.cu had TILE_SZ*2 treating BF16 counts as bytes, causing half the needed SMEM. Fix: use TILE_SZ*sizeof(bf16_t) or TILE_SZ*2 with proper byte accounting.

Key design decisions:

• PV SS MMA scale factor: ~1.0 (not 0.5 like QK). The MMA with BLOCK_MN_A=128, BLOCK_MN_B=16 produces output at approximately 1× the raw dot product.
• Per-K-tile sPk fill: For decode T=1, only row 0 needs filling (16 values per K-tile). For prefill T>1, all 128 rows need filling from s_p_vals.
• s_p_vals shared memory: 128 floats (512 bytes) to share softmax output across warps. Avoids recomputing P for each K-tile.

Files:

• tests/unit/test_fmha_v5.cu — Full FMHA HD=16 with PV SS MMA (SMEM-P, per K-tile fill)
• tests/unit/test_pv_ss_128.cu — PV SS MMA with (128,128) P (accumulation debug, shows the bug)
• tests/unit/test_pv_ss.cu — Minimal PV SS MMA (A=128×16, B=16×16, standalone)
• tests/unit/test_ss_ts_sequence.cu — SS+TS sequence test (proved SS+TS coexistence works)

Next steps:

1. Extend to HD=64, HD=128, HD=256 — the per-K-tile approach scales naturally
2. Prefill T>1 — fill all 128 rows of sPk, not just row 0
3. Multi-head support — per-head launch or head-packed M
4. Production kernel — integrate into fmha_sm100.cuh
5. Benchmark — compare SS-PV vs register-math PV at decode T=1