biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update CURRENT_ISSUE: HD=16 done, HD=64 in progress

Browse Source
This commit is contained in:
biondizzle
2026-05-28 15:16:19 +00:00
parent 6b9b06647a
commit 6896d1aebb
1 changed files with 21 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
CURRENT_ISSUE.md
View File

@@ -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