fix: harden attention kernel boundaries

- fix scalar prefill head_dim=32 out-of-bounds via G=4 dispatch - fix MMA decode 4D mask head indexing and invalid-row mask access - add q_head/kv_head divisibility and head-dim contiguity checks - validate split-KV scratch and decode out_buf layout in bindings - set max dynamic shared memory for scalar decode D=256 - cover scalar prefill D=32 in pure C test
2026-08-09 14:53:24 +08:00
parent 596c35fd71
commit d565d44c43
8 changed files with 50 additions and 19 deletions
@@ -123,8 +123,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
            mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
                                               qr0, qr1,
                                               p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
-                                               batch, q_head,
+                                               batch, q_head, q_head,
                                               p.mask,
+                                               va, vb,
                                               Sacc, Oacc, m0, m1, l0, l1, lane);

            mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);