feat: unify attention backend with multi-dim mask support

- Add attention() functional entry delegating to active backend
- GQA/MLA forward calls attention() instead of inline cache/SDPA
- CUDA kernels support 2D/3D/4D mask via mask_h_stride field
- CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens
- KVCache.max_len precomputed in bind_tasks to avoid GPU sync
- batch==1 decode short-circuits mask=None
- Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask
- 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native

astrai/extension/__init__.py

@@ -20,6 +20,7 @@ from astrai.extension.attention_backend import (
     AttentionBackend,
     CudaBackend,
     TorchNativeBackend,
+    attention,
     attn_backend,
     get_backend,
 )
@@ -35,6 +36,7 @@ __all__ = [
     "AttentionBackend",
     "CudaBackend",
     "TorchNativeBackend",
+    "attention",
     "attn_backend",
     "get_backend",
     "attn_decode",

astrai/extension/attention_backend.py

@@ -113,6 +113,37 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
     )
 
 
+def attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    kv_cache: Optional[KVCache] = None,
+    layer_id: int = 0,
+    attn_mask: Optional[Tensor] = None,
+    is_causal: bool = False,
+) -> Tensor:
+    """Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
+
+    Delegates to the active backend (set via ``with attn_backend(...)``).
+    Handles KV cache I/O, GQA head expansion, and causal masking so the
+    caller only needs to provide projected q/k/v.
+
+    Args:
+        q: [batch, q_len, n_heads, head_dim] (blhd)
+        k: [batch, q_len, n_kv_heads, head_dim] (blhd)
+        v: [batch, q_len, n_kv_heads, head_dim] (blhd)
+        kv_cache: cache dataclass, or None for training (no cache).
+        layer_id: transformer layer index for buffer access.
+        attn_mask: pre-built attention mask (SDPA-compatible).
+        is_causal: whether to apply causal masking.
+
+    Returns:
+        [batch, q_len, n_heads * head_dim]
+    """
+    backend = get_backend()
+    return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
+
+
 class AttentionBackend(ABC):
     """Abstract base for attention computation strategies.
 
@@ -307,13 +338,19 @@ class CudaBackend(AttentionBackend):
         kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
         kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
 
-        max_len = kv_cache.seq_lens.max().item()
+        seq_lens = kv_cache.seq_lens
+        max_len = kv_cache.max_len
 
         page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
 
         k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
         v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
 
+        if q.size(0) == 1:
+            mask = None
+        else:
+            mask = torch.arange(max_len, device=q.device)[None, :] < seq_lens[:, None]
+
         out = attn_paged_decode(
             q,
             page_table,
@@ -321,7 +358,7 @@ class CudaBackend(AttentionBackend):
             v_cache,
             page_size=1,
             kv_len=max_len,
-            mask=None,
+            mask=mask,
             is_causal=is_causal,
         )
 
@@ -354,7 +391,7 @@ class CudaBackend(AttentionBackend):
         kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
         kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
 
-        max_len = kv_cache.seq_lens.max()
+        max_len = kv_cache.max_len
         indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
         pos_mask = (
             torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]

astrai/inference/core/cache.py

@@ -202,6 +202,7 @@ class KVCache:
         req_pool_indices: [batch_size] — row indices into req_to_token
         seq_lens: [batch_size] — per-request total sequence lengths
         out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
+        max_len: max(seq_lens) as Python int — avoids GPU sync in decode
     """
 
     k_buffer: Tensor
@@ -210,6 +211,7 @@ class KVCache:
     req_pool_indices: Tensor
     seq_lens: Tensor
     out_cache_loc: Tensor
+    max_len: int = 0
 
 
 class PagePool:
@@ -439,6 +441,7 @@ class PagePool:
             req_pool_indices=req_pool_indices,
             seq_lens=seq_lens_t,
             out_cache_loc=out_cache_loc,
+            max_len=max(seq_lens),
         )
 
     # ---- internals ----

astrai/model/components/__init__.py

@@ -1,4 +1,4 @@
-from astrai.model.components.attention import GQA, MLA, repeat_kv
+from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
@@ -21,5 +21,4 @@ __all__ = [
     "RotaryEmbedding",
     "apply_rotary_emb",
     "get_rotary_emb",
-    "repeat_kv",
 ]

astrai/model/components/attention.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
+from astrai.extension import attention
 from astrai.factory import BaseFactory
 from astrai.inference.core.cache import KVCache
 from astrai.model.components.linear import Linear
@@ -12,17 +13,6 @@ from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import apply_rotary_emb
 
 
-def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    bs, slen, n_heads, head_dim = x.shape
-    if n_rep == 1:
-        return x
-    return (
-        x[:, :, :, None, :]
-        .expand(bs, slen, n_heads, n_rep, head_dim)
-        .reshape(bs, slen, n_heads * n_rep, head_dim)
-    )
-
-
 class AttnFactory(BaseFactory[nn.Module]):
     pass
 
@@ -86,29 +76,7 @@ class GQA(nn.Module):
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
 
-        if kv_cache is not None:
+        sdqa_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
-            kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k
-            kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v
-
-            max_len = kv_cache.seq_lens.max()
-            indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-            pos_mask = (
-                torch.arange(max_len, device=x.device)[None, :]
-                < kv_cache.seq_lens[:, None]
-            )
-            indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
-            k = kv_cache.k_buffer[self.layer_id, indices]
-            v = kv_cache.v_buffer[self.layer_id, indices]
-
-        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
-
-        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
-        sdqa_out = (
-            F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
-            .permute(0, 2, 1, 3)
-            .contiguous()
-            .flatten(2)
-        )
 
         if self.use_gated_attention:
             sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
@@ -203,28 +171,7 @@ class MLA(nn.Module):
             q = self.q_norm(q)
             k = self.k_norm(k)
 
-        if kv_cache is not None:
+        attn_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
-            kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k
-            kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v
-
-            max_len = kv_cache.seq_lens.max()
-            indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-            pos_mask = (
-                torch.arange(max_len, device=x.device)[None, :]
-                < kv_cache.seq_lens[:, None]
-            )
-            indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
-            k = kv_cache.k_buffer[self.layer_id, indices]
-            v = kv_cache.v_buffer[self.layer_id, indices]
-
-        q = q.permute(0, 2, 1, 3)
-        k = k.permute(0, 2, 1, 3)
-        v = v.permute(0, 2, 1, 3)
-
-        attn_out = F.scaled_dot_product_attention(
-            q, k, v, attn_mask, is_causal=is_causal
-        )
-        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
 
         if self.use_gated_attention:
             attn_out = attn_out * F.sigmoid(self.gate(x))

csrc/kernels/attn_common.h

@@ -19,9 +19,11 @@ struct AttentionParams {
     // KV strides (K and V share the same layout — only base pointers differ)
     int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
 
-    // Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len]
+    // Mask: 2D [batch, kv_len], 3D [batch, q_len, kv_len],
-    int mask_b_stride;   // = kv_len (both 2D and 3D)
+    // or 4D [batch, n_heads, q_len, kv_len] (head dim broadcasts when stride=0)
-    int mask_q_stride;   // 2D: 0 (all q rows share); 3D: kv_len
+    int mask_b_stride;   // batch stride
+    int mask_h_stride;   // head stride (0 = broadcast across heads)
+    int mask_q_stride;   // q stride (0 = all q rows share)
 
     const T* __restrict__ q;
     const T* __restrict__ k;
@@ -52,8 +54,9 @@ struct PagedAttentionParams {
     // Q strides (layout-agnostic)
     int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
 
-    // Mask strides (2D or 3D)
+    // Mask strides (2D, 3D, or 4D)
     int mask_b_stride;
+    int mask_h_stride;
     int mask_q_stride;
 
     const T* __restrict__ q;

csrc/kernels/attn_decode_split_kv.cuh

@@ -24,7 +24,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
 
     // KV: [batch, kv_head, kv_len, head_dim] — stride-based base
     int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    int mask_base = batch * p.mask_b_stride;
+    int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
 
     float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
 

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -109,8 +109,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
                                            0, 0,
-                                           p.mask_b_stride, 0,
+                                           p.mask_b_stride, 0, 0,
-                                           batch,
+                                           batch, 0,
                                            p.mask,
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 

csrc/kernels/attn_entry_utils.cuh

@@ -44,6 +44,9 @@ inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
 }
 
 // ---- Shared mask packing ----
+// Accepts 2D [batch, kv_len], 3D [batch, q_len, kv_len],
+// or 4D [batch, n_heads, q_len, kv_len].
+// Head/q dimensions with size 1 broadcast (stride set to 0).
 template <typename P>
 inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
     if (p.use_mask) {
@@ -54,18 +57,26 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
         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_h_stride = 0;
             p.mask_q_stride = 0;
         } else if (m.dim() == 3) {
-            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
+            TORCH_CHECK(m.size(1) == 1 || 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);
+            p.mask_h_stride = 0;
+            p.mask_q_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
+        } else if (m.dim() == 4) {
+            TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.q_len, "mask q_len mismatch");
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
+            p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
         } else {
-            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
+            TORCH_CHECK(false, "mask must be 2D, 3D, or 4D");
         }
         p.mask = m.data_ptr<bool>();
     } else {
         p.mask = nullptr;
         p.mask_b_stride = 0;
+        p.mask_h_stride = 0;
         p.mask_q_stride = 0;
     }
 }

csrc/kernels/attn_mma_utils.cuh

@@ -192,8 +192,8 @@ __device__ inline void mma_softmax_tile(
     int kv0,
     int maxc0, int maxc1,
     int qrow0, int qrow1,
-    int mask_b_stride, int mask_q_stride,
+    int mask_b_stride, int mask_h_stride, int mask_q_stride,
-    int mask_batch,
+    int mask_batch, int mask_head,
     const bool* __restrict__ mask,
     float Sacc[Traits::NC8][4],
     float Oacc[Traits::DN8][4],
@@ -204,8 +204,8 @@ __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 + qrow0 * mask_q_stride;
+    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 + qrow1 * mask_q_stride;
+    int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride;
     #pragma unroll
     for (int n8 = 0; n8 < Traits::NC8; n8++) {
         int cc = kv0 + n8 * 8 + 2 * tid4;

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -31,7 +31,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
     int ch_begin = split * chunks_per_split;
     int ch_end = min(chunks_total, ch_begin + chunks_per_split);
 
-    const int mask_base = batch * p.mask_b_stride;
+    const int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
 
     for (int ci = ch_begin; ci < ch_end; ci++) {
         int chunk_start = ci * PDC_CHUNK;

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -110,8 +110,8 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
                                            0, 0,
-                                           p.mask_b_stride, 0,
+                                           p.mask_b_stride, 0, 0,
-                                           batch,
+                                           batch, 0,
                                            p.mask,
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 

csrc/kernels/attn_prefill_split_q.cuh

@@ -64,7 +64,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
 
     // KV: stride-based base
     int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    int mask_batch_base = batch * p.mask_b_stride;
+    int mask_batch_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
     int tiles   = (p.kv_len + P_BC - 1) / P_BC;
     int tt      = G * ROWS;
     int lid     = row * G + gpos;

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -114,8 +114,8 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
                                  : p.kv_len;
             mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
                                                qr0, qr1,
-                                               p.mask_b_stride, p.mask_q_stride,
+                                               p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
-                                               batch,
+                                               batch, q_head,
                                                p.mask,
                                                Sacc, Oacc, m0, m1, l0, l1, lane);
 

csrc/tests/test_utils.cuh

@@ -103,6 +103,7 @@ inline void set_default_strides(P& p) {
     p.kv_stride_l = p.head_dim;
     p.kv_stride_d = 1;
     p.mask_b_stride = p.kv_len;
+    p.mask_h_stride = 0;
     p.mask_q_stride = 0;
 }
 
@@ -114,6 +115,7 @@ inline void set_default_paged_strides(P& p) {
     p.q_stride_l  = p.head_dim;
     p.q_stride_d  = 1;
     p.mask_b_stride = p.kv_len;
+    p.mask_h_stride = 0;
     p.mask_q_stride = 0;
 }
 

tests/extension/conftest.py

@@ -0,0 +1,35 @@
+"""Shared fixtures for extension tests."""
+
+import pytest
+import torch
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.extension import is_available
+from astrai.model.transformer import AutoRegressiveLM
+
+CUDA_AVAILABLE = torch.cuda.is_available() and is_available("attn_paged_decode")
+skip_no_cuda = pytest.mark.skipif(
+    not CUDA_AVAILABLE, reason="CUDA not available or kernels not built"
+)
+
+D = 64
+CFG = dict(
+    vocab_size=1000,
+    hidden_size=128,
+    num_attention_heads=2,
+    num_key_value_heads=1,
+    intermediate_size=256,
+    max_position_embeddings=64,
+    num_hidden_layers=2,
+    rms_norm_eps=1e-5,
+    attn_type="gqa",
+    ffn_type="mlp",
+)
+
+
+@pytest.fixture
+def cuda_model():
+    config = AutoRegressiveLMConfig(**CFG)
+    model = AutoRegressiveLM(config).to(device="cuda", dtype=torch.bfloat16)
+    model.eval()
+    return model, config

tests/extension/test_backend.py

@@ -0,0 +1,45 @@
+"""Backend selection and context-manager switching tests.
+
+These tests do not require CUDA — they only check that the active
+backend is correctly set and restored.
+"""
+
+import pytest
+
+from astrai.extension import (
+    ATTN_BACKEND,
+    CudaBackend,
+    TorchNativeBackend,
+    attn_backend,
+    get_backend,
+)
+
+
+def test_default_backend_is_torch_native():
+    backend = get_backend()
+    assert isinstance(backend, TorchNativeBackend)
+
+
+def test_attn_backend_context_with_enum():
+    with attn_backend(ATTN_BACKEND.CUDA):
+        assert isinstance(get_backend(), CudaBackend)
+    assert isinstance(get_backend(), TorchNativeBackend)
+
+
+def test_attn_backend_context_with_class():
+    with attn_backend(CudaBackend):
+        assert isinstance(get_backend(), CudaBackend)
+    assert isinstance(get_backend(), TorchNativeBackend)
+
+
+def test_attn_backend_context_with_instance():
+    custom = CudaBackend()
+    with attn_backend(custom):
+        assert get_backend() is custom
+    assert isinstance(get_backend(), TorchNativeBackend)
+
+
+def test_cudabackend_is_context_manager():
+    with CudaBackend():
+        assert isinstance(get_backend(), CudaBackend)
+    assert isinstance(get_backend(), TorchNativeBackend)

tests/extension/test_backend_equivalence.py

@@ -0,0 +1,199 @@
+"""Numerical equivalence between TorchNativeBackend and CudaBackend.
+
+Covers training forward, inference prefill, inference decode (mixed
+seq_lens with padding mask), and end-to-end scheduler.run_batch.
+"""
+
+import torch
+
+from astrai.extension import ATTN_BACKEND, attn_backend
+from astrai.inference.core.cache import PagePool
+from tests.extension.conftest import D, skip_no_cuda
+
+
+@skip_no_cuda
+def test_training_forward_matches_torch(cuda_model):
+    """Training forward (kv_cache=None) should produce identical logits."""
+    model, _ = cuda_model
+    input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
+
+    with torch.no_grad():
+        out_torch = model(input_ids)
+    with attn_backend(ATTN_BACKEND.CUDA):
+        with torch.no_grad():
+            out_cuda = model(input_ids)
+
+    diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
+    assert diff == 0.0, f"Training forward diff {diff} should be 0"
+
+
+@skip_no_cuda
+def test_prefill_with_kv_cache_matches_torch(cuda_model):
+    """Inference prefill with KV cache should match torch backend."""
+    model, _ = cuda_model
+    prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
+    max_len = max(len(p) for p in prompt_ids)
+    batch = len(prompt_ids)
+
+    device = "cuda"
+    input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
+    input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
+    position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
+    for i, p in enumerate(prompt_ids):
+        input_ids[i, : len(p)] = torch.tensor(p, device=device)
+        input_mask[i, : len(p)] = True
+        position_ids[i, : len(p)] = torch.arange(len(p), device=device)
+
+    cache = PagePool(
+        n_layers=2,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=4,
+        max_seq_len=64,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+
+    cache.task_alloc("t1", prompt_ids[0])
+    cache.task_alloc("t2", prompt_ids[1])
+    kv1 = cache.bind_tasks(
+        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
+    )
+    with torch.inference_mode():
+        out_torch = model(
+            input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
+        )
+
+    cache.task_free("t1")
+    cache.task_free("t2")
+    cache.task_alloc("t1", prompt_ids[0])
+    cache.task_alloc("t2", prompt_ids[1])
+    kv2 = cache.bind_tasks(
+        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
+    )
+    with attn_backend(ATTN_BACKEND.CUDA):
+        with torch.inference_mode():
+            out_cuda = model(
+                input_ids,
+                input_mask=input_mask,
+                kv_cache=kv2,
+                position_ids=position_ids,
+            )
+
+    for i, p in enumerate(prompt_ids):
+        d = (
+            (
+                out_torch["logits"][i, : len(p)].float()
+                - out_cuda["logits"][i, : len(p)].float()
+            )
+            .abs()
+            .max()
+            .item()
+        )
+        assert d == 0.0, f"Prefill diff for sample {i}: {d}"
+
+
+@skip_no_cuda
+def test_decode_mixed_seq_lens_matches_torch(cuda_model):
+    """Decode with mixed seq_lens in batch — padding mask must produce correct output."""
+    model, _ = cuda_model
+    device = "cuda"
+
+    prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
+    cache = PagePool(
+        n_layers=2,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=4,
+        max_seq_len=64,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+
+    # Prefill to populate cache
+    max_len = max(len(p) for p in prompt_ids)
+    batch = len(prompt_ids)
+    input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
+    input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
+    position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
+    for i, p in enumerate(prompt_ids):
+        input_ids[i, : len(p)] = torch.tensor(p, device=device)
+        input_mask[i, : len(p)] = True
+        position_ids[i, : len(p)] = torch.arange(len(p), device=device)
+
+    cache.task_alloc("t1", prompt_ids[0])
+    cache.task_alloc("t2", prompt_ids[1])
+    kv = cache.bind_tasks(
+        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
+    )
+    with torch.inference_mode():
+        model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
+
+    # Decode step — seq_lens are 9 and 7 (after extending)
+    dec_ids = torch.tensor([[99], [98]], dtype=torch.long, device=device)
+    dec_pos = torch.tensor([[8], [6]], dtype=torch.long, device=device)
+    total_len = 9
+    dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
+
+    kv_t = cache.bind_tasks(["t1", "t2"], [9, 7], device)
+    with torch.inference_mode():
+        out_torch = model(
+            dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
+        )
+
+    kv_c = cache.bind_tasks(["t1", "t2"], [9, 7], device)
+    with attn_backend(ATTN_BACKEND.CUDA):
+        with torch.inference_mode():
+            out_cuda = model(
+                dec_ids, input_mask=dec_mask, kv_cache=kv_c, position_ids=dec_pos
+            )
+
+    diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
+    assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
+
+
+@skip_no_cuda
+def test_run_batch_cuda_matches_torch_greedy(cuda_model):
+    """Greedy decode (temperature=0) should produce identical tokens."""
+    from astrai.inference.core.scheduler import InferenceScheduler
+    from tests.helpers import FakeTokenizer
+
+    model, _ = cuda_model
+    tokenizer = FakeTokenizer()
+
+    prompts = [[1, 2, 3, 4, 5], [10, 11, 12, 13, 14, 15, 16]]
+
+    sched = InferenceScheduler(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=4,
+        max_seq_len=64,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    out_torch = sched.run_batch(prompts, max_tokens=5, temperature=0.0)
+    sched.stop()
+
+    cache_cuda = PagePool(
+        n_layers=2,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=4,
+        max_seq_len=64,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    sched2 = InferenceScheduler(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=4,
+        max_seq_len=64,
+        device="cuda",
+        dtype=torch.bfloat16,
+        cache=cache_cuda,
+    )
+    with attn_backend(ATTN_BACKEND.CUDA):
+        out_cuda = sched2.run_batch(prompts, max_tokens=5, temperature=0.0)
+    sched2.stop()
+
+    assert out_torch == out_cuda, f"Torch={out_torch} != CUDA={out_cuda}"

tests/extension/test_kernel_mask.py

@@ -0,0 +1,74 @@
+"""Kernel-level mask dimension support (2D, 3D, 4D)."""
+
+import torch
+
+from tests.extension.conftest import D, skip_no_cuda
+
+
+@skip_no_cuda
+def test_kernel_accepts_2d_mask():
+    """Kernel should accept 2D mask [batch, kv_len]."""
+    from astrai.extension.attention_ops import attn_prefill
+
+    batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
+    kv_len = 8
+    q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, kv_len, dtype=torch.bool, device="cuda")
+    mask[:, 4:] = False
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+def test_kernel_accepts_3d_mask():
+    """Kernel should accept 3D mask [batch, q_len, kv_len]."""
+    from astrai.extension.attention_ops import attn_prefill
+
+    batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
+    kv_len = 8
+    q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda")
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+def test_kernel_accepts_4d_mask():
+    """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
+    from astrai.extension.attention_ops import attn_prefill
+
+    batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
+    kv_len = 8
+    q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
+    mask[:, :, :, 4:] = False
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+def test_4d_mask_matches_no_mask_when_all_true():
+    """A 4D all-True mask should produce the same output as no mask."""
+    from astrai.extension.attention_ops import attn_prefill
+
+    batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
+    kv_len = 8
+    q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+
+    out_no_mask = attn_prefill(q, k, v, mask=None, is_causal=False)
+    mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
+    out_with_mask = attn_prefill(q, k, v, mask=mask, is_causal=False)
+
+    diff = (out_no_mask.float() - out_with_mask.float()).abs().max().item()
+    assert diff == 0.0, f"4D all-True mask diff: {diff}"