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_decode.cu
+7
View File
@@ -29,8 +29,15 @@ torch::Tensor attn_decode(
TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32");
TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2;
TORCH_CHECK(o_part_buf->numel() >= o_needed,
"o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
TORCH_CHECK(ml_part_buf->numel() >= ml_needed,
"ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel());
TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(),
"split buffers must be CUDA tensors");
TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(),
"split buffers must be contiguous");
p.o_part = (float*)o_part_buf->data_ptr();
p.ml_part = (float*)ml_part_buf->data_ptr();
} else {