33 lines
2.6 KiB
Markdown
33 lines
2.6 KiB
Markdown
# CURRENT_ISSUE.md — PV GEMM for Prefill
|
||
|
||
## Status: ✅ PV via SS MMA WORKING (cosine 0.9997)
|
||
|
||
### 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.
|
||
|
||
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.
|
||
|
||
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)
|
||
|
||
### 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
|