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
This commit is contained in:
2026-08-09 14:53:24 +08:00
parent 596c35fd71
commit d565d44c43
8 changed files with 50 additions and 19 deletions
+5 -4
View File
@@ -107,10 +107,11 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
int maxc = IsCausal ? KV::decode_attend_len(p, batch) : seq_len;
mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
0, 0,
p.mask_b_stride, 0, 0,
batch, 0,
p.mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
batch, q_head0 + gid, q_head0 + gid + 8,
p.mask,
va, vb,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
};