ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: harden attention kernel boundaries

Browse Source 
- 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:
ViperEkura
2026-08-09 14:53:24 +08:00
parent 596c35fd71
commit d565d44c43
8 changed files with 50 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_mma_utils.cuh
+8 -7
View File
@@ -199,8 +199,9 @@ __device__ inline void mma_softmax_tile(
int maxc0, int maxc1,
int qrow0, int qrow1,
int mask_b_stride, int mask_h_stride, int mask_q_stride,
int mask_batch, int mask_head,
int mask_batch, int mask_head0, int mask_head1,
const bool* __restrict__ mask,
bool valid0, bool valid1,
float Sacc[Traits::NC8][4],
float Oacc[Traits::DN8][4],
float& m0, float& m1,
@@ -210,16 +211,16 @@ __device__ inline void mma_softmax_tile(
int tid4 = lane & 3;
float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
int mask_base0 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow0 * mask_q_stride;
int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride;
int mask_base0 = mask_batch * mask_b_stride + mask_head0 * mask_h_stride + qrow0 * mask_q_stride;
int mask_base1 = mask_batch * mask_b_stride + mask_head1 * mask_h_stride + qrow1 * mask_q_stride;
#pragma unroll
for (int n8 = 0; n8 < Traits::NC8; n8++) {
int cc = kv0 + n8 * 8 + 2 * tid4;
int c1 = cc + 1;
bool b0 = (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
bool b1 = (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
bool b2 = (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
bool b3 = (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
bool b0 = !valid0 || (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
bool b1 = !valid0 || (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
bool b2 = !valid1 || (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
bool b3 = !valid1 || (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
float s2 = b2 ? -FLT_MAX : Sacc[n8][2];