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_paged_decode.cu
+14 -5
View File
@@ -27,12 +27,14 @@ torch::Tensor attn_paged_decode(
torch::Tensor O;
if (out_buf.has_value() && out_buf->defined()) {
TORCH_CHECK(out_buf->dtype() == q.dtype(), "out_buf dtype must match q");
TORCH_CHECK(out_buf->is_cuda() && out_buf->is_contiguous(),
"out_buf must be a contiguous CUDA tensor");
TORCH_CHECK(out_buf->size(0) >= q.size(0), "out_buf batch too small");
TORCH_CHECK(out_buf->size(1) >= q.size(1), "out_buf heads too small");
TORCH_CHECK(out_buf->size(2) >= q.size(2), "out_buf head_dim too small");
O = out_buf.value().slice(0, 0, q.size(0))
.slice(1, 0, q.size(1))
.slice(2, 0, q.size(2));
TORCH_CHECK(out_buf->size(1) == q.size(1), "out_buf heads must match q");
TORCH_CHECK(out_buf->size(2) == q.size(2), "out_buf head_dim must match q");
TORCH_CHECK(q.is_contiguous(),
"q must be contiguous when out_buf is provided");
O = out_buf.value().slice(0, 0, q.size(0));
} else {
O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
}
@@ -43,8 +45,15 @@ torch::Tensor attn_paged_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 {