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 {
csrc/kernels/attn_decode_split_kv_mma.cuh
+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);
};
csrc/kernels/attn_dispatchers.cuh
+5 -1
View File
@@ -80,7 +80,7 @@ template <typename KV>
struct PrefillLauncherScalar {
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
constexpr int G = 8, ROWS = 32, P_BC = 32;
constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32;
int q_len = KV::host_q_len(p);
dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
dim3 block(G, ROWS);
@@ -161,6 +161,10 @@ struct DecodeLauncherScalar {
int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit
dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
dim3 block(32, g);
cudaFuncSetAttribute(
attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem);
attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>
<<<grid, block, smem, stream>>>(p);
}
csrc/kernels/attn_entry_utils.cuh
+8 -1
View File
@@ -106,14 +106,17 @@ inline void attn_pack_params(
TORCH_CHECK(v.dtype() == torch::kBFloat16);
TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
extract_q_dims_and_strides(q, layout, p);
if (layout == BLHD) k = k.transpose(1, 2), v = v.transpose(1, 2);
p.kv_head = (int)k.size(1);
p.kv_len = (int)k.size(2);
TORCH_CHECK(p.q_head % p.kv_head == 0,
"q_head must be divisible by kv_head");
TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1,
"Q/K/V head_dim must be contiguous");
p.kv_stride_b = (int)k.stride(0);
p.kv_stride_h = (int)k.stride(1);
@@ -170,6 +173,8 @@ inline void attn_pack_paged_decode_params(
p.head_dim = (int)q.size(2);
p.kv_head = (int)k_cache.size(1);
TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
"Q/K/V head_dim must be contiguous");
TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
@@ -252,6 +257,8 @@ inline void attn_pack_paged_prefill_params(
p.kv_head = (int)k_cache.size(1);
p.batch = (int)req_pool_indices.size(0);
TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
"Q/K/V head_dim must be contiguous");
TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]");
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];
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 {
csrc/kernels/attn_prefill_split_q_mma.cuh
+2 -1
View File
@@ -123,8 +123,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
qr0, qr1,
p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
batch, q_head,
batch, q_head, q_head,
p.mask,
va, vb,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
csrc/tests/attn_test.cu
+1
View File
@@ -323,6 +323,7 @@ int main() {
// ---- PREFILL ----
{
const int configs[][7] = {
{1,2,1,64,128,32,0}, // scalar fallback D=32
{1,2,1,64,128,64,0}, // tiny: B,Hq,Hk,q,kv,D,causal
{1,32,4,512,512,128,0}, // standard
{1,32,4,128,256,128,0}, // medium