Commit Graph

18 Commits

Author SHA1 Message Date
d1c1eaeddc clean: remove debug prints, multirow kernel complete with multi-tile KV merge 2026-05-28 23:57:31 +00:00
c65baabcc9 fix: V tile copy — V is (HD, SK_TOTAL) so tile columns are not contiguous 2026-05-28 23:55:52 +00:00
869460a932 debug: add LSE verification and merge debug prints 2026-05-28 23:54:30 +00:00
2f2259395e fix: always normalize in kernel, correct KV merge with normalized O + LSE 2026-05-28 23:53:44 +00:00
914f76d30c multirow: add normalize flag, un-norm + LSE output, multi-tile KV merge test 2026-05-28 23:51:23 +00:00
ca5cf0e517 test: add multi-head and batched prefill tests for multirow kernel 2026-05-28 23:48:53 +00:00
0220e51d18 fix: typo cudaErrorCudaSuccess -> cudaSuccess 2026-05-28 23:18:21 +00:00
468614a4e2 fmha_multirow: non-interleaved design — softmax first, then PV
KEY FIX: TMEM is shared between QK output (S) and PV output (O).
Cannot interleave softmax reads with PV writes because PV overwrites S.

New flow:
1. QK GEMM → S in TMEM
2. Softmax: read ALL S from TMEM, compute P in registers
   - Pass 1: row_max (4 warps, 32x32b.x8)
   - Pass 2: exp, sum, store P in p_vals[SK_TILE] registers
3. PV GEMM: write P to sPk per K-tile, accumulate O in TMEM
4. Epilogue: read O from TMEM, normalize, write GMEM

P in registers: each lane holds float p_vals[128] = 512 bytes.
Register budget: 128 lanes × 512B = 64KB (within B200 256KB register file).
2026-05-28 23:17:43 +00:00
c768abed95 test: softmax-only kernel (QK + row_max, no PV) 2026-05-28 23:15:36 +00:00
d840fbbf85 test: clean multirow test with proper SMEM calc 2026-05-28 23:10:49 +00:00
f2124b9378 fix: SMEM calc in decode test 2026-05-28 23:08:54 +00:00
58ff781388 test: simplified decode kernel for debugging multirow 2026-05-28 23:08:33 +00:00
4cfb707405 fix: correct SMEM size calculation in multirow test 2026-05-28 22:53:46 +00:00
863a030c3b fmha_multirow: rewrite with 32x32b.x8 only, no s_p_vals, row_page addressing
- Kill 64KB s_p_vals buffer — P is streamed per K-tile through sPk
- All TMEM ops use 32x32b.x8 exclusively (16x256b.x1 crashes on 2nd call)
- T>32: 4 softmax warps use row_page offset in TMEM address (row<<16)
- Lane l in warp w handles row w*32+l
- Two-pass softmax: pass 1 row_max, pass 2 exp/sum interleaved with PV
- PV: N=16 sub-tiles, SS MMA sPk(128,16) × sV(16,16) → TMEM
- Epilogue: 32x32b.x8 TMEM read, normalize, BF16 → GMEM
- SMEM budget: ~14KB (well within 232KB)
2026-05-28 22:52:52 +00:00
08694b8136 Fix multi-row softmax v3: 32x32b.x8 with per-lane per-row (no wmax/wsum), per-row sRowMax/sRowSum arrays 2026-05-28 20:10:13 +00:00
bf4dfd131b Fix nvcc goto-bypasses-init: move var decls before goto targets 2026-05-28 20:04:59 +00:00
2b09d4f2ef Fix nvcc goto-bypasses-init in multi-row test 2026-05-28 20:04:45 +00:00
d8b421ccee Multi-row FMHA kernel (Milestone 4): T>1 prefill support with 4-warp parallel softmax 2026-05-28 20:04:29 +00:00