53 lines
2.9 KiB
Markdown
53 lines
2.9 KiB
Markdown
# CURRENT ISSUE: UMMA FMHA — Full Pipeline + Production
|
||
|
||
## What's working ✅
|
||
- **Full UMMA FMHA HD=64 pipeline**: QK → softmax → PV → output
|
||
- QK GEMM: UMMA SS, multi-K-tile accumulate, cos 0.999999
|
||
- Softmax: TMEM read → max/exp/sum → TMEM write, max rel err 0.003
|
||
- PV: register math (O[d] = Σ P[0,j] × V[d,j]), decode only
|
||
- **Overall: cosine 0.999998!**
|
||
- **UMMA QK GEMM at HD=16**: Row 0 matches scalar with ZERO error
|
||
- **TMEM one-way epilogue**: cos 0.999999 at hd=64, cos 0.999998 at hd=128
|
||
- **32x32b.x8 TMEM stores** work for P write-back (no 16x256b crash)
|
||
|
||
## Next steps
|
||
1. **PV GEMM via tcgen05.mma TS**: For prefill (T>1), need UMMA-based PV
|
||
- tcgen05.mma TS crashed with "illegal memory access" in initial tests
|
||
- Need to debug the TMEM A operand addressing for PV
|
||
- Decode path (T=1) works with register math — PV GEMM is only for prefill
|
||
|
||
2. **HD=128/256**: Extend multi-K-tile QK to larger head dims
|
||
- HD=128: 8 K-tiles, separate SMEM per K-tile
|
||
- HD=256: 16 K-tiles — SMEM budget needs checking
|
||
|
||
3. **Multi-head launch**: Per-head kernel dispatch (128 heads for Pro)
|
||
- Current test: single head
|
||
- Production: grid=(1, n_h, batch) or Python loop
|
||
|
||
4. **Multi-KV-tile**: s_k > 128 requires multiple attention tiles + KV merge
|
||
- Same D5 merge formula: O = Σ exp(lse_i) · O_i / Σ exp(lse_i)
|
||
|
||
5. **Production kernel**: Integrate UMMA pipeline into fmha_sm100.cuh
|
||
- Replace SMEM scalar attention with UMMA QK GEMM
|
||
- Keep register-math PV for decode
|
||
- Add PV GEMM path for prefill
|
||
|
||
## Key lessons from this session
|
||
- **Source stride mismatch**: `write_k_to_smem<128,16>` template reads k[i] with stride=16, but actual K has stride=64. Must use separate SMEM per K-tile with manual writes.
|
||
- **Offset descriptors DON'T WORK** for (128,64) SMEM with K-tile offsets. Even/odd column corruption. Use separate (128,16) SMEM buffers.
|
||
- **tcgen05.mma TS crashes** with illegal memory access — likely TMEM A operand addressing issue. Debug needed.
|
||
- **tcgen05.fence::after_thread_sync** is the correct fence after MMA, before TMEM read.
|
||
- **MMA computes UNSCALED dot product** — apply 1/sqrt(HD) in softmax.
|
||
- **32x32b.x8 TMEM stores** work in loops (unlike 16x256b.x1 which crashes on 2nd call)
|
||
- **s_vals scope**: when merging softmax + PV, keep s_vals in same `if (wid == 0)` block
|
||
- **Decode PV is simple**: O[d] = Σ P[0,j] × V[d,j] — no UMMA needed for T=1
|
||
|
||
## Files
|
||
- `dsv4/kernels/attention/fmha_umma_desc.cuh` — descriptors, SMEM layout, MMA wrappers
|
||
- `dsv4/kernels/attention/fmha_common.cuh` — BF16, TMEM ops, warp reductions
|
||
- `dsv4/kernels/attention/fmha_epilogue_sm100.cuh` — TMEM one-way epilogue (reference kernel)
|
||
- `tests/unit/test_fmha_hd64.cu` — FULL PIPELINE TEST (QK+softmax+PV)
|
||
- `tests/unit/test_fmha_softmax.cu` — softmax validation (HD=64)
|
||
- `tests/unit/test_umma_qk.cu` — QK GEMM (HD=16)
|
||
- `tests/unit/test_umma_qk_hd64.cu` — QK GEMM (HD=64, multi-K-tile)
|