refactor: template combine kernel, fix mask bug, unify dispatch

- Template combine kernel, share macros, extract entry_utils helpers
- Fix mask indexing (pass stride not pre-multiplied base)
- Remove !p.use_mask — MMA handles mask

csrc/kernels/attn_decode.cu

@@ -38,7 +38,7 @@ template <int HEAD_DIM>
 static void dispatch_decode(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     int G = p.q_head / p.kv_head;
-    if (!p.use_mask && G >= 1 && G <= 16) {
+    if (G >= 1 && G <= 16) {
         launch_mma_decode<HEAD_DIM, 32>(p);
         return;
     }
@@ -65,7 +65,7 @@ torch::Tensor attn_decode(
     auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
     p.o = (bf16*)O_view.data_ptr();
 
-    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_decode<D>(p); });
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
     return O;
 }
 

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -64,7 +64,6 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
 
     // KV: stride-based base — [batch, kv_head, kv_len, head_dim]
     const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    const int mask_batch_base = batch * p.mask_b_stride;
     const int tiles_total = (p.kv_len + BC - 1) / BC;
     const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
     const int ti_begin = split * tiles_per_split;
@@ -125,7 +124,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                     0, 0,
-                                    mask_batch_base, 0,
+                                    p.mask_b_stride, 0,
                                     batch,
                                     p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);

csrc/kernels/attn_entry_utils.cuh

@@ -5,10 +5,6 @@
 
 using bf16 = __nv_bfloat16;
 
-// ---------------------------------------------------------------------------
-// Shared dispatch helpers — eliminates duplication across .cu entry files.
-// ---------------------------------------------------------------------------
-
 inline int compute_num_splits(int base_blocks, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
@@ -16,21 +12,20 @@ inline int compute_num_splits(int base_blocks, int tiles_total) {
     return std::max(1, std::min(n, std::min(tiles_total, 32)));
 }
 
-// Dispatch head_dim to a generic lambda FN (callable as FN.operator()<D>()).
+// Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
-template <typename Fn>
+// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
-inline void dispatch_head_dim(int hd, Fn&& fn) {
+//   Expands to: fn<32>(arg); fn<64>(arg); etc.
-    switch (hd) {
+#define DISPATCH_HEAD_DIM(hd, fn, arg) \
-        case 32:  fn.template operator()<32>();  break;
+    switch (hd) { \
-        case 64:  fn.template operator()<64>();  break;
+        case 32:  fn<32>(arg); break; \
-        case 128: fn.template operator()<128>(); break;
+        case 64:  fn<64>(arg); break; \
-        case 256: fn.template operator()<256>(); break;
+        case 128: fn<128>(arg); break; \
-        default:
+        case 256: fn<256>(arg); break; \
-            TORCH_CHECK(false, "unsupported head_dim ", hd,
+        default: \
-                         " (supported: 32, 64, 128, 256)");
+            TORCH_CHECK(false, "unsupported head_dim ", hd, \
-    }
+                         " (supported: 32, 64, 128, 256)"); \
     }
 
-// Allocate split-KV partial buffers and wire into params.
 template<typename P>
 inline void alloc_split_partials(P& p) {
     auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
@@ -40,10 +35,48 @@ inline void alloc_split_partials(P& p) {
     p.ml_part = (float*)ml_part.data_ptr();
 }
 
-// ---------------------------------------------------------------------------
+// ---- Shared Q-dims + strides extraction ----
-// Param packing — fills AttentionParams from torch tensors with validation.
+template <typename P>
-// ---------------------------------------------------------------------------
+inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
+    if (layout == 1) q = q.transpose(1, 2);
+    p.batch = (int)q.size(0);
+    p.q_head = (int)q.size(1);
+    p.q_len = (int)q.size(2);
+    p.head_dim = (int)q.size(3);
+    p.q_stride_b = (int)q.stride(0);
+    p.q_stride_h = (int)q.stride(1);
+    p.q_stride_l = (int)q.stride(2);
+    p.q_stride_d = (int)q.stride(3);
+}
 
+// ---- Shared mask packing ----
+template <typename P>
+inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
+    if (p.use_mask) {
+        auto m = mask.value();
+        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
+        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
+        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
+        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
+        if (m.dim() == 2) {
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = 0;
+        } else if (m.dim() == 3) {
+            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = (int)m.stride(1);
+        } else {
+            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
+        }
+        p.mask = m.data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+        p.mask_b_stride = 0;
+        p.mask_q_stride = 0;
+    }
+}
+
+// ---- attn_pack_params (contiguous KV) ----
 template<typename T>
 inline void attn_pack_params(
     torch::Tensor q,
@@ -52,7 +85,7 @@ inline void attn_pack_params(
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    int64_t layout,   // 0 = b h l d, 1 = b l h d
+    int64_t layout,
     AttentionParams<T>& p
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
@@ -64,26 +97,14 @@ inline void attn_pack_params(
     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");
 
-    // Normalize to b h l d view (zero-copy transpose if user passed b l h d)
+    extract_q_dims_and_strides(q, layout, p);
-    if (layout == 1) {
+
-        q = q.transpose(1, 2);
+    if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
-        k = k.transpose(1, 2);
-        v = v.transpose(1, 2);
-    }
 
-    p.batch = (int)q.size(0);
-    p.q_head = (int)q.size(1);
-    p.q_len = (int)q.size(2);
-    p.head_dim = (int)q.size(3);
     p.kv_head = (int)k.size(1);
     p.kv_len = (int)k.size(2);
     TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
 
-    // Strides (layout-agnostic: works for b h l d and b l h d)
-    p.q_stride_b = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_l = (int)q.stride(2);
-    p.q_stride_d = (int)q.stride(3);
     p.kv_stride_b = (int)k.stride(0);
     p.kv_stride_h = (int)k.stride(1);
     p.kv_stride_l = (int)k.stride(2);
@@ -100,34 +121,10 @@ inline void attn_pack_params(
     p.o_part = nullptr;
     p.ml_part = nullptr;
 
-    if (p.use_mask) {
+    pack_mask(mask, p);
-        auto m = mask.value();
-        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
-        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
-        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
-        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
-        if (m.dim() == 2) {
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = 0;
-        } else if (m.dim() == 3) {
-            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = (int)m.stride(1);
-        } else {
-            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
-        }
-        p.mask = m.data_ptr<bool>();
-    } else {
-        p.mask = nullptr;
-        p.mask_b_stride = 0;
-        p.mask_q_stride = 0;
-    }
 }
 
-// ---------------------------------------------------------------------------
+// ---- attn_pack_paged_params ----
-// Param packing for paged attention.
-// ---------------------------------------------------------------------------
-
 template<typename T>
 inline void attn_pack_paged_params(
     torch::Tensor q,
@@ -139,7 +136,7 @@ inline void attn_pack_paged_params(
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    int64_t layout,   // 0 = b h l d, 1 = b l h d
+    int64_t layout,
     PagedAttentionParams<T>& p
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
@@ -151,15 +148,8 @@ inline void attn_pack_paged_params(
     TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
     TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
 
-    // Normalize Q to b h l d view if user passed b l h d
+    extract_q_dims_and_strides(q, layout, p);
-    if (layout == 1) {
-        q = q.transpose(1, 2);
-    }
 
-    p.batch = (int)q.size(0);
-    p.q_head = (int)q.size(1);
-    p.q_len = (int)q.size(2);
-    p.head_dim = (int)q.size(3);
     p.kv_head = (int)k_cache.size(2);
     p.kv_len = (int)kv_len;
     p.page_size = (int)page_size;
@@ -171,12 +161,6 @@ inline void attn_pack_paged_params(
                 "k_cache dim 1 must equal page_size, got ",
                 k_cache.size(1), " vs ", page_size);
 
-    // Q strides
-    p.q_stride_b = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_l = (int)q.stride(2);
-    p.q_stride_d = (int)q.stride(3);
-
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
@@ -189,26 +173,5 @@ inline void attn_pack_paged_params(
     p.o_part = nullptr;
     p.ml_part = nullptr;
 
-    if (p.use_mask) {
+    pack_mask(mask, p);
-        auto m = mask.value();
-        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
-        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
-        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
-        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
-        if (m.dim() == 2) {
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = 0;
-        } else if (m.dim() == 3) {
-            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
-            p.mask_b_stride = (int)m.stride(0);
-            p.mask_q_stride = (int)m.stride(1);
-        } else {
-            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
-        }
-        p.mask = m.data_ptr<bool>();
-    } else {
-        p.mask = nullptr;
-        p.mask_b_stride = 0;
-        p.mask_q_stride = 0;
-    }
 }

csrc/kernels/attn_paged_decode.cu

@@ -34,7 +34,7 @@ template <int HEAD_DIM>
 static void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     int G = p.q_head / p.kv_head;
-    if (!p.use_mask && G >= 1 && G <= 16 && p.page_size >= 32) {
+    if (G >= 1 && G <= 16 && p.page_size >= 32) {
         launch_paged_mma_decode<HEAD_DIM, 32>(p);
         return;
     }
@@ -62,7 +62,7 @@ torch::Tensor attn_paged_decode(
     auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
     p.o = (bf16*)O_view.data_ptr();
 
-    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_paged_decode<D>(p); });
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
     return O;
 }
 

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -51,7 +51,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    const int mask_batch_base = batch * p.mask_b_stride;
     const int tiles_total = (p.kv_len + BC - 1) / BC;
     const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
     const int ti_begin = split * tiles_per_split;
@@ -119,7 +118,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                     0, 0,
-                                    mask_batch_base, 0,
+                                    p.mask_b_stride, 0,
                                     batch,
                                     p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);

csrc/kernels/attn_prefill.cu

@@ -47,7 +47,7 @@ torch::Tensor attn_prefill(
     auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
     p.o = (bf16*)O_view.data_ptr();
 
-    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_prefill<D>(p); });
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
     return O;
 }
 

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -94,7 +94,6 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int block_max_kv =
         blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
     const int has_mask = p.use_mask && p.mask;
-    const int mask_batch_base = batch * p.mask_b_stride;
 
     // Last active tile: block-level causal bound (all warps in the block share
     // the K/V load, so the prefetch range is the block max, not per-warp).
@@ -165,7 +164,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
                                         : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
                                     qr0, qr1,
-                                    mask_batch_base, p.mask_q_stride,
+                                    p.mask_b_stride, p.mask_q_stride,
                                     batch,
                                     p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);