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