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

Update CURRENT_ISSUE: ALL HD=16/64/128/256 PASS cos 0.999997+

Browse Source 
Documented Layout D N=64 bug and N=16 sub-tile workaround.
This commit is contained in:
biondizzle
2026-05-28 16:03:05 +00:00
parent 01319d7247
commit a18d9c1584
1 changed files with 27 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
CURRENT_ISSUE.md
View File

@@ -1,30 +1,36 @@
# CURRENT_ISSUE.md — PV GEMM for Prefill
# CURRENT_ISSUE.md — FMHA Raw CUDA Stage D
## Status: HD=16 ✅ (cos 0.9997), HD=64 ✅ (cos 0.931, BF16 accumulation precision)
## Status: HD=16/64/128/256 ALL PASS ✅ (cos 0.999997+)
### What works:
- **HD=16**: Full pipeline QK(SS) → softmax → PV(SS, 8 K-tiles) → epilogue. **Cosine 0.9997**.
- **HD=64**: Full pipeline with BLOCK_MN_B=64. **Cosine 0.931** — V canonical layout verified correct, error is BF16 accumulation precision in PV MMA (kind::f16 uses BF16 dot products, reference uses FP32).
- Both require `cudaFuncSetAttribute` for SMEM >48KB (HD=64 uses 52.9 KB).
- **HD=16**: Full pipeline QK(SS) → softmax → PV(SS, N=16) → epilogue. **Cosine 0.999998**.
- **HD=64**: Full pipeline with 4 N=16 PV sub-tiles. **Cosine 0.999997**.
- **HD=128**: Full pipeline with 8 N=16 PV sub-tiles. **Cosine 0.999997**.
- **HD=256**: Full pipeline with 16 N=16 PV sub-tiles. **Cosine 0.999997**.
### Key architectural decisions:
1. **PV via SS MMA with SMEM-P** (NOT TS MMA). TS MMA's A-fragment TMEM layout (Layout A) is incompatible with 32x32b stores. SS MMA with both operands in SMEM avoids this.
2. **Per-K-tile P fill** into reusable (128,16) buffer from shared `s_p_vals[128]`. The (128,128) canonical P with K-tile offsets has an accumulation bug.
3. **V canonical layout**: `g_mn=d/8, g_k=lr/8, llr=d%8, lc=lr%8` where d=HD(MN), lr=seq_pos(K). The original code had MN/K swapped.
4. **PV MMA scale**: ~1.0 for both BLOCK_MN_B=16 and 64. QK MMA scale is also 1.0 (NOT 0.5 as assumed earlier).
5. **SMEM >48KB**: Must opt in with `cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)`.
1. **PV via SS MMA with N=16 sub-tiles** (NOT N=64 or N=128). The `tcgen05.mma` with `make_idesc(128, N)` for N≠16,128 has a Layout D bug that skips TMEM columns (8 missing columns at N=64: cols 32-35 and 48-51).
2. **Q/K loaded one K-tile at a time** (reusing single SMEM buffer). This keeps SMEM at ~25KB regardless of HD.
3. **TMEM offset `tb + n*16`** for each PV N-sub-tile. The MMA with N=16 writes 16 columns starting at the C operand address.
4. **MMA scale is 1.0** for both QK and PV (NOT 0.5 as initially assumed).
### Files:
- `test_fmha_v5.cu` — HD=16 full pipeline (cos 0.9997) ✅
- `test_fmha_hd64_smem_p.cu` — HD=64 full pipeline (cos 0.931) 🚧
- `test_pv_ss_128.cu` — PV SS MMA with (128,128) P (shows accumulation bug)
- `test_pv_ss_b64.cu` — PV SS MMA with B=(64,16) BLOCK_MN=64 (works)
- `test_ss_ts_sequence.cu` — SS+TS coexistence test
- `test_pv_ss.cu` — Minimal PV SS MMA (A=128×16, B=16×16)
- `tests/unit/test_fmha_gen.cu` — Generalized kernel + test (parameterized on HD_VAL)
- `tests/unit/test_fmha_hd{16,64,128,256}_gen.cu` — Wrapper files
- `tests/unit/test_fmha_hd64_n16_v2.cu` — Standalone HD=64 test (first to pass)
- `tests/unit/test_tmem_zero_pv.cu` — TMEM Layout D N=64 bug proof
- `tests/unit/test_tmem_all_lanes.cu` — All-lane TMEM dump for N=64
### Next steps:
1. **HD=64 precision**: Investigate FP32-accumulation MMA variant (kind::f32) for PV
2. **HD=128, HD=256**: Extend the pipeline (more QK/PV K-tiles, larger V)
3. **Prefill T>1**: Fill all 128 rows of sPk from s_p_vals
4. **Multi-head**: Per-head launch or head-packed M dimension
5. **Production kernel**: Extract into `fmha_sm100.cuh` with proper warp specialization
1. **Production kernel**: Extract the generalized kernel into `fmha_sm100_tc.cuh` with proper 6-warp specialization
2. **Prefill T>1**: Fill all 128 rows of sPk (not just row 0), enable multi-row softmax
3. **Multi-head**: Per-head launch or head-packed M dimension
4. **CUDA graph**: Integration with production.py
5. **Benchmarking**: Measure actual throughput vs CuTeDSL kernel
### Layout D N=64 Bug (documented for NVIDIA):
- `tcgen05.mma.cta_group::1.kind::f16` with `make_idesc(128, 64)` writes to only 56 out of 64 expected TMEM columns
- Missing columns: 32, 33, 34, 35, 48, 49, 50, 51
- These columns contain zero after the MMA, causing output positions d=32-35 and d=48-51 to be wrong
- Workaround: use N=16 sub-tiles (4 × make_idesc(128, 16)) with different TMEM offsets
- This bug does NOT affect N=16 or N=128 — only intermediate N values