diff --git a/CURRENT_ISSUE.md b/CURRENT_ISSUE.md index d6f28d4e..cd62b506 100644 --- a/CURRENT_ISSUE.md +++ b/CURRENT_ISSUE.md @@ -1,32 +1,28 @@ # CURRENT_ISSUE.md β€” PV GEMM for Prefill -## Status: βœ… PV via SS MMA WORKING (cosine 0.9997) +## Status: HD=16 βœ… (cos 0.9997), HD=64 🚧 (cos 0.931, ~0.4% PV MMA error) -### 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. +### HD=16 β€” COMPLETE +Full pipeline: QK(SS, 1 K-tile) β†’ softmax(TMEMβ†’SMEM) β†’ PV(SS, 8 K-tiles) β†’ epilogue. Cosine 0.9997. -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. +### 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 -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) +### 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. **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 +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