diff --git a/CURRENT_ISSUE.md b/CURRENT_ISSUE.md index fc322dbe..d873182f 100644 --- a/CURRENT_ISSUE.md +++ b/CURRENT_ISSUE.md @@ -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