44 lines
1.7 KiB
Markdown
44 lines
1.7 KiB
Markdown
# CURRENT_ISSUE.md — FMHA 6-Warp Specialization
|
|
|
|
## Status: Milestone 1 COMPLETE ✅ (cos 0.999997+ at HD=16/64/128/256)
|
|
|
|
### What works:
|
|
- **6-warp kernel**: Warps 0-3 softmax/epilogue, Warp 4 MMA, Warp 5 data staging
|
|
- **All HD values**: HD=16/64/128/256 pass with cos 0.999997+
|
|
- **Warp role separation**: MMA and data loading on separate warps
|
|
- **CTA-wide sync**: __syncthreads() between phases
|
|
|
|
### Architecture:
|
|
```
|
|
Warp 0-3 (tid 0-127): Softmax + correction + epilogue
|
|
- Read S from TMEM → softmax → write P to SMEM
|
|
- After PV: read O from TMEM → BF16 → GMEM
|
|
- T=1 decode: only warp 0 processes row 0
|
|
Warp 4 (tid 128-159): MMA
|
|
- tcgen05.mma SS for QK (N=128) and PV (N=16 sub-tiles)
|
|
- TMEM alloc/dealloc
|
|
Warp 5 (tid 160-191): Data staging
|
|
- Load Q/K/V from GMEM to SMEM (canonical layout)
|
|
- Fill sPk from s_p_vals
|
|
```
|
|
|
|
### Next milestones:
|
|
1. **TMA loads** (Milestone 2): Replace direct GMEM reads with cp.async.bulk.tensor
|
|
- Requires CUtensorMap creation on host
|
|
- mbarrier synchronization
|
|
2. **Pipeline overlap** (Milestone 3): Double-buffer K/V loads
|
|
- Load next K/V while computing current QK
|
|
- mbarrier producer-consumer sync between warp 5 and warp 4
|
|
3. **Multi-row softmax** (Milestone 4): Process all 128 rows (prefill T>1)
|
|
4. **Multi-head launch** (Milestone 5): grid=(1, n_h, batch)
|
|
5. **Production integration** (Milestone 6): Hook into production.py
|
|
|
|
### Files:
|
|
- `dsv4/kernels/attention/fmha_6warp.cuh` — 6-warp kernel
|
|
- `tests/unit/test_fmha_6warp.cu` — Test harness
|
|
- `tests/unit/test_fmha_6warp_hd{16,64,128,256}.cu` — HD-specific wrappers
|
|
|
|
### Layout D N=64 Bug (documented for NVIDIA):
|
|
- tcgen05.mma with make_idesc(128, 64) skips TMEM cols 32-35, 48-51
|
|
- Workaround: N=16 sub-tiles with TMEM offset n*16
|