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

csrc/kernels/attn_decode.cu

@@ -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

@@ -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,
+                                            p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
-                                           batch, 0,
+                                            batch, q_head0 + gid, q_head0 + gid + 8,
-                                           p.mask,
+                                            p.mask,
-                                           Sacc, Oacc, m0, m1, l0, l1, lane);
+                                            va, vb,
+                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
     };

csrc/kernels/attn_dispatchers.cuh

@@ -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

@@ -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

@@ -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_base0 = mask_batch * mask_b_stride + mask_head0 * 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_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 b0 = !valid0 || (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
-        bool b1 = (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
+        bool b1 = !valid0 || (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
-        bool b2 = (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
+        bool b2 = !valid1 || (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
-        bool b3 = (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
+        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

@@ -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(1) == q.size(1), "out_buf heads must match q");
-        TORCH_CHECK(out_buf->size(2) >= q.size(2), "out_buf head_dim too small");
+        TORCH_CHECK(out_buf->size(2) == q.size(2), "out_buf head_dim must match q");
-        O = out_buf.value().slice(0, 0, q.size(0))
+        TORCH_CHECK(q.is_contiguous(),
-                            .slice(1, 0, q.size(1))
+                    "q must be contiguous when out_buf is provided");
-                            .slice(2, 0, q.size(2));
+        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

@@ -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

@@ -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