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:
@@ -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);
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user