feat: SGLang-style paged attention kernels replace page-table path

- PagedAttentionParams uses flat KV pool + req_to_token + kv_indptr/qo_indptr instead of page_table
- MMA split-KV decode and split-Q prefill kernels with indirect ragged-batch addressing
- Prefill kernel accepts 4D mask (causal-aware); decode kernel supports 2D mask
- CudaBackend is inference-only: kv_cache=None raises, no torch fallback
- benchmark.py: required --ckpt, --backend/--compare options
- Parallel build isolates build-temp/build-lib per subprocess
- Standalone test covers decode/prefill with mask, 27 cases pass

astrai/extension/attention_backend.py

@@ -38,8 +38,10 @@ import torch
 import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension.attention_ops import attn_paged_decode, attn_prefill
+from astrai.extension.attention_ops import (
-from astrai.extension.loader import is_available
+    attn_paged_decode,
+    attn_paged_prefill,
+)
 from astrai.inference.core.cache import KVCache
 
 _current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
@@ -307,24 +309,19 @@ _default_backend = TorchNativeBackend()
 class CudaBackend(AttentionBackend):
     """CUDA kernel backend with direct KV cache access.
 
-    Decode path: writes K/V to cache, then calls ``attn_paged_decode``
+    Decode path: writes K/V to the flat pool, then calls
-    with ``page_size=1`` (each token slot is a single-token "page").
+    ``attn_paged_decode`` with req_to_token + kv_indptr.
-    The ``req_to_token`` table serves directly as the page table.
 
-    Prefill path: writes K/V to cache, gathers full-sequence K/V via
+    Prefill path: writes K/V to the flat pool, then calls
-    indirect indexing (same as TorchNativeBackend), then calls
+    ``attn_paged_prefill`` with ragged-batch support via qo_indptr +
-    ``attn_prefill``.
+    kv_indptr.
 
-    Training path (``kv_cache is None``): calls ``attn_prefill`` directly
+    ``kv_cache is None`` (training) is not handled — use
-    on the projected q/k/v.
+    ``TorchNativeBackend`` for training.
 
-    Falls back to ``TorchNativeBackend`` for any path where the
+    Raises ``RuntimeError`` if the required kernel is not available.
-    corresponding CUDA kernel is not available.
     """
 
-    def __init__(self):
-        self._fallback = TorchNativeBackend()
-
     def fwd_decode(
         self,
         q: Tensor,
@@ -335,47 +332,34 @@ class CudaBackend(AttentionBackend):
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> Tensor:
-        if kv_cache is None or not is_available("attn_paged_decode"):
+        if kv_cache is None:
-            return self._fallback.fwd_decode(
+            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
 
         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.max_len
+        b = q.size(0)
+        q_3d = q.squeeze(1)
 
-        if kv_cache.page_table is not None:
+        kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
-            page_table = kv_cache.page_table
+        kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
-        else:
-            page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
 
-        k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
+        mask = None
-        v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
+        if b > 1 and kv_cache.decode_mask is not None:
-
-        if q.size(0) == 1:
-            mask = None
-        elif kv_cache.decode_mask is not None:
             mask = kv_cache.decode_mask
-        else:
-            mask = (
-                torch.arange(max_len, device=q.device)[None, :]
-                < kv_cache.seq_lens[:, None]
-            )
 
         out = attn_paged_decode(
-            q,
+            q_3d,
-            page_table,
+            kv_cache.k_buffer[layer_id],
-            k_cache,
+            kv_cache.v_buffer[layer_id],
-            v_cache,
+            kv_cache.req_to_token,
-            page_size=1,
+            kv_cache.req_pool_indices,
-            kv_len=max_len,
+            kv_indptr,
+            kv_cache.max_len,
             mask=mask,
             is_causal=is_causal,
         )
-
+        return out.unsqueeze(1).flatten(2)
-        out = out.flatten(2)
-        return out
 
     def fwd_prefill(
         self,
@@ -388,32 +372,34 @@ class CudaBackend(AttentionBackend):
         is_causal: bool = False,
     ) -> Tensor:
         if kv_cache is None:
-            if is_available("attn_prefill"):
+            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
-                out = attn_prefill(q, k, v, mask=attn_mask, is_causal=is_causal)
-                return out.flatten(2)
-            return self._fallback.fwd_prefill(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
-
-        if not is_available("attn_prefill"):
-            return self._fallback.fwd_prefill(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
 
         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.max_len
+        b = q.size(0)
-        indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
+        q_len = q.size(1)
-        pos_mask = (
-            torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
-        )
-        indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
-        k_full = kv_cache.k_buffer[layer_id, indices]
-        v_full = kv_cache.v_buffer[layer_id, indices]
 
-        out = attn_prefill(q, k_full, v_full, mask=attn_mask, is_causal=is_causal)
+        kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
-        return out.flatten(2)
+        kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
+
+        qo_indptr = torch.arange(b + 1, dtype=torch.int32, device=q.device) * q_len
+
+        q_flat = q.reshape(b * q_len, q.size(2), q.size(3))
+
+        out = attn_paged_prefill(
+            q_flat,
+            kv_cache.k_buffer[layer_id],
+            kv_cache.v_buffer[layer_id],
+            kv_cache.req_to_token,
+            kv_cache.req_pool_indices,
+            kv_indptr,
+            qo_indptr,
+            attn_mask,
+            q_len,
+            is_causal=is_causal,
+        )
+        return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2)
 
 
 _BACKEND_REGISTRY: dict[ATTN_BACKEND, type[AttentionBackend]] = {

astrai/extension/attention_ops.py

@@ -92,39 +92,94 @@ def attn_prefill(
 
 def attn_paged_decode(
     q: torch.Tensor,
-    page_table: torch.Tensor,
     k_cache: torch.Tensor,
     v_cache: torch.Tensor,
-    page_size: int,
+    req_to_token: torch.Tensor,
-    kv_len: int,
+    req_pool_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    max_seq_len: int,
     mask: Optional[torch.Tensor] = None,
     is_causal: bool = False,
 ) -> torch.Tensor:
-    """Paged GQA decode attention (q_len == 1, direct page-table access).
+    """SGLang-style paged decode (q_len == 1, flat KV pool).
+
+    Reads K/V directly from a flat pool [size, kv_head, head_dim] via
+    req_to_token indirect indexing.  Each request has its own seq_len
+    (from kv_indptr), eliminating padding waste.
 
     Args:
-        q: [batch, 1, n_heads, head_dim] (blhd, bf16)
+        q: [batch, n_heads, head_dim] (bf16, 3D — no seq dim)
-        page_table: [batch, max_pages] (int64)
+        k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
-        k_cache: [n_pages, page_size, n_kv_heads, head_dim] (bf16)
         v_cache: same as k_cache
-        page_size: tokens per page
+        req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
-        kv_len: actual sequence length per request
+        req_pool_indices: [batch] (int64) — rows into req_to_token
-        mask: 2D [batch, kv_len] or 3D [batch, 1, kv_len] (bool, True=keep)
+        kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
+        max_seq_len: max per-request seq_len (Python int, for split computation)
+        mask: 2D [batch, max_seq_len] (bool, True=keep) or None
         is_causal: apply causal mask
 
     Returns:
-        [batch, 1, n_heads, head_dim] (blhd, bf16)
+        [batch, n_heads, head_dim] (bf16, 3D)
     """
     _check_available("attn_paged_decode")
-    causal_offset = (kv_len - 1) if is_causal else -1
+    causal_offset = 0 if is_causal else -1
     return _modules["attn_paged_decode"].attn_paged_decode(
         q,
-        page_table,
         k_cache,
         v_cache,
-        page_size,
+        req_to_token,
-        kv_len,
+        req_pool_indices,
+        kv_indptr,
+        max_seq_len,
         mask=mask,
         causal_offset=causal_offset,
-        layout=TensorLayout.BLHD,
+    )
+
+
+def attn_paged_prefill(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    qo_indptr: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    max_q_len: int = 0,
+    is_causal: bool = False,
+) -> torch.Tensor:
+    """SGLang-style paged prefill (ragged batch, flat KV pool).
+
+    Reads K/V directly from a flat pool [size, kv_head, head_dim] via
+    req_to_token.  Supports ragged batches: each request has its own
+    q_len and kv_len, addressed via qo_indptr and kv_indptr.
+
+    Args:
+        q: [total_q, n_heads, head_dim] (bf16, 3D — flattened across requests)
+        k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
+        v_cache: same as k_cache
+        req_to_token: [num_reqs, max_context_len] (int64)
+        req_pool_indices: [batch] (int64)
+        kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
+        qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
+        mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
+        max_q_len: max per-request q_len (Python int, for grid computation)
+        is_causal: apply causal mask
+
+    Returns:
+        [total_q, n_heads, head_dim] (bf16, 3D)
+    """
+    _check_available("attn_paged_prefill")
+    causal_offset = 0 if is_causal else -1
+    return _modules["attn_paged_prefill"].attn_paged_prefill(
+        q,
+        k_cache,
+        v_cache,
+        req_to_token,
+        req_pool_indices,
+        kv_indptr,
+        qo_indptr,
+        mask,
+        max_q_len,
+        causal_offset=causal_offset,
     )

astrai/extension/loader.py

@@ -11,7 +11,13 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode", "rotary_emb"]
+KERNEL_NAMES = [
+    "attn_decode",
+    "attn_prefill",
+    "attn_paged_decode",
+    "attn_paged_prefill",
+    "rotary_emb",
+]
 
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}

csrc/build.py

@@ -72,4 +72,5 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
 register("attn_decode")
 register("attn_prefill")
 register("attn_paged_decode")
+register("attn_paged_prefill")
 register("rotary_emb")

csrc/kernels/attn_common.h

@@ -43,35 +43,55 @@ struct AttentionParams {
     AT* __restrict__ ml_part;
 };
 
+// ---- PagedAttentionParams ----
+// SGLang-style indirect params over a shared KV pool.
+// k_cache/v_cache: [size, kv_head, head_dim] (bare buffers, no gather).
+// req_to_token:    [num_reqs, max_context_len] token -> slot.
+// req_pool_indices:[batch] rows of the current batch into req_to_token.
+// kv_indptr:       [batch+1] prefix sum of per-request seq_lens (device).
+// qo_indptr:       [batch+1] prefix sum of per-request q_len (prefill) or
+//                  nullptr for decode (q_len == 1 everywhere).
 template<typename T, typename AT = float>
 struct PagedAttentionParams {
     int batch;
     int q_head;
     int kv_head;
-    int q_len;
-    int kv_len;
     int head_dim;
+    int num_splits;
     int use_mask;
-    int causal_offset;
+    int causal_offset;   // -1 = non-causal; >=0 = causal (per-request offset
+                         //   computed inside kernel from kv_indptr/qo_indptr)
     float scale;
 
-    int num_splits;
+    // Q: [total_q, q_head, head_dim] (3D flattened — no batch dim).
-    int page_size;
+    // For decode total_q == batch (q_len=1 per request).
-    int max_pages;
+    // For prefill total_q == qo_indptr[batch].
+    int q_stride_l, q_stride_h, q_stride_d;
 
-    // Q strides (layout-agnostic)
+    // Q: [total_q, q_head, head_dim]
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
+    const T* __restrict__ q;
 
-    // Mask strides (2D, 3D, or 4D)
+    // Flat KV pool: [size, kv_head, head_dim]
+    const T* __restrict__ k_cache;
+    const T* __restrict__ v_cache;
+
+    // Indexing
+    const int64_t* __restrict__ req_to_token;      // [num_reqs, max_context_len]
+    const int64_t* __restrict__ req_pool_indices;   // [batch]
+    const int* __restrict__ kv_indptr;               // [batch+1]
+    const int* __restrict__ qo_indptr;               // [batch+1] or nullptr (decode)
+    int max_context_len;  // req_to_token stride (dim 1)
+    int max_seq_len;      // max per-request seq_len (host-side, for split computation)
+    int total_q;          // total Q tokens across all requests (host-side, for grid)
+    int max_q_len;        // max per-request q_len (host-side, for prefill grid)
+
+    // Mask: [batch, max_seq_len] (decode) or [batch, 1, q_len, kv_len]
+    // (prefill, optional).  mask_h_stride/mask_q_stride are 0 when those
+    // dims are size 1 (broadcast).
     int mask_b_stride;
     int mask_h_stride;
     int mask_q_stride;
-
-    const T* __restrict__ q;
-    const T* __restrict__ k_cache;
-    const T* __restrict__ v_cache;
     const bool* __restrict__ mask;
-    const int64_t* __restrict__ page_table;
 
     T* __restrict__ o;
     AT* __restrict__ o_part;

csrc/kernels/attn_dispatchers.cuh

@@ -12,6 +12,7 @@
 #include "attn_prefill_split_q_mma.cuh"
 #include "attn_decode_split_kv_mma.cuh"
 #include "attn_paged_decode_split_kv_mma.cuh"
+#include "attn_paged_prefill_split_q_mma.cuh"
 #endif
 
 // Cached SM count — cudaDeviceGetAttribute is a host-side call that was
@@ -145,18 +146,18 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
 }
 
 // ======================================================================
-// Paged Decode
+// Paged Decode (SGLang-style: flat pool + req_to_token + kv_indptr)
 // ======================================================================
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int group_size) {
+static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
     int G = p.q_head / p.kv_head;
     constexpr int MAX_G = 16;
     constexpr int BC = 16;
     int num_passes = (G + MAX_G - 1) / MAX_G;
-    int tiles_total = (p.kv_len + BC - 1) / BC;
+    int tiles_total = (p.max_seq_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total, 2);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head * num_passes, tiles_total, 2);
     constexpr int STAGES = 2;
     using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
     dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
@@ -166,10 +167,10 @@ static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int gr
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
 static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size) {
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
+    int chunks_total = (p.max_seq_len + PDC_CHUNK - 1) / PDC_CHUNK;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
     size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
-    int g = min(group_size, 32);  // cap at 32 to respect 1024-thread limit
+    int g = min(group_size, 32);
     dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
     dim3 block(32, g);
     paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
@@ -182,10 +183,37 @@ static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
     int group_size = p.q_head / p.kv_head;
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, 0);
 #else
     DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
 #endif
 
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
+
+// ======================================================================
+// Paged Prefill (SGLang-style: flat pool + ragged batch)
+// ======================================================================
+
+#ifndef ASTRAI_NO_MMA
+template <int HEAD_DIM, bool IsCausal, bool HasMask>
+static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p) {
+    constexpr int WARPS = 4;
+    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
+    using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
+    int max_q_tiles = (p.max_q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS);
+    dim3 grid(max_q_tiles, p.q_head, p.batch);
+    dim3 block(Traits::NUM_THREADS);
+    paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
+}
+#endif
+
+template <int HEAD_DIM>
+static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p) {
+    bool is_causal = (p.causal_offset >= 0);
+    bool has_mask = (p.use_mask && p.mask);
+
+#ifndef ASTRAI_NO_MMA
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p);
+#endif
+}

csrc/kernels/attn_entry_utils.cuh

@@ -134,54 +134,178 @@ inline void attn_pack_params(
     pack_mask(mask, p);
 }
 
-// ---- attn_pack_paged_params ----
+// ---- attn_pack_paged_decode_params ----
+// SGLang-style: flat KV pool + req_to_token indexing + variable
+// seq_lens via kv_indptr.  Q is [batch, q_head, head_dim] (q_len=1 per req).
 template<typename T>
-inline void attn_pack_paged_params(
+inline void attn_pack_paged_decode_params(
     torch::Tensor q,
-    torch::Tensor page_table,
     torch::Tensor k_cache,
     torch::Tensor v_cache,
-    int64_t page_size,
+    torch::Tensor req_to_token,
-    int64_t kv_len,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    int64_t max_seq_len,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    int64_t layout,
     PagedAttentionParams<T>& p
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
 
-    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda() && kv_indptr.is_cuda());
     TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
     TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
     TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
-    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
+    TORCH_CHECK(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
-    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
+    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
+    TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
+    TORCH_CHECK(q.dim() == 3, "q must be 3D [batch, q_head, head_dim]");
 
-    extract_q_dims_and_strides(q, layout, p);
+    p.batch = (int)q.size(0);
-
+    p.q_head = (int)q.size(1);
-    p.kv_head = (int)k_cache.size(2);
+    p.head_dim = (int)q.size(2);
-    p.kv_len = (int)kv_len;
+    p.kv_head = (int)k_cache.size(1);
-    p.page_size = (int)page_size;
+    TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
-    p.max_pages = (int)page_table.size(1);
-
-    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
-    TORCH_CHECK(k_cache.size(1) == page_size,
+    TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
-                "k_cache dim 1 must equal page_size, got ",
+
-                k_cache.size(1), " vs ", page_size);
+    p.q_stride_l = (int)q.stride(0);
+    p.q_stride_h = (int)q.stride(1);
+    p.q_stride_d = (int)q.stride(2);
+
+    p.k_cache = (const T*)k_cache.data_ptr();
+    p.v_cache = (const T*)v_cache.data_ptr();
+    p.q = (const T*)q.data_ptr();
+    p.req_to_token = req_to_token.data_ptr<int64_t>();
+    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    p.kv_indptr = kv_indptr.data_ptr<int>();
+    p.qo_indptr = nullptr;
+    p.max_context_len = (int)req_to_token.size(1);
+    p.max_seq_len = (int)max_seq_len;
+    p.total_q = p.batch;  // decode: 1 Q token per request
+    p.max_q_len = 1;
 
     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);
 
-    p.page_table = page_table.data_ptr<int64_t>();
+    if (p.use_mask) {
+        auto m = mask.value();
+        TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
+        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
+        p.mask_b_stride = (int)m.stride(0);
+        p.mask_h_stride = 0;
+        p.mask_q_stride = 0;
+        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;
+    }
+
+    p.o = nullptr;
+    p.o_part = nullptr;
+    p.ml_part = nullptr;
+}
+
+// ---- attn_pack_paged_prefill_params ----
+// SGLang-style: flat KV pool + req_to_token + ragged batch via qo_indptr.
+// Q is [total_q, q_head, head_dim] (flattened across all requests).
+template<typename T>
+inline void attn_pack_paged_prefill_params(
+    torch::Tensor q,
+    torch::Tensor k_cache,
+    torch::Tensor v_cache,
+    torch::Tensor req_to_token,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    torch::Tensor qo_indptr,
+    c10::optional<torch::Tensor> mask,
+    int64_t max_q_len,
+    int64_t causal_offset,
+    double scale,
+    PagedAttentionParams<T>& p
+) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+
+    TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda());
+    TORCH_CHECK(kv_indptr.is_cuda() && qo_indptr.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
+    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
+    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
+    TORCH_CHECK(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
+    TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be int32");
+    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
+    TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
+    TORCH_CHECK(q.dim() == 3, "q must be 3D [total_q, q_head, head_dim]");
+
+    p.q_head = (int)q.size(1);
+    p.head_dim = (int)q.size(2);
+    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(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]");
+    TORCH_CHECK(qo_indptr.size(0) == p.batch + 1, "qo_indptr must be [batch+1]");
+
+    p.q_stride_l = (int)q.stride(0);
+    p.q_stride_h = (int)q.stride(1);
+    p.q_stride_d = (int)q.stride(2);
+
     p.k_cache = (const T*)k_cache.data_ptr();
     p.v_cache = (const T*)v_cache.data_ptr();
     p.q = (const T*)q.data_ptr();
+    p.req_to_token = req_to_token.data_ptr<int64_t>();
+    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    p.kv_indptr = kv_indptr.data_ptr<int>();
+    p.qo_indptr = qo_indptr.data_ptr<int>();
+    p.max_context_len = (int)req_to_token.size(1);
+    p.total_q = (int)q.size(0);  // prefill: flattened Q across all requests
+    p.max_q_len = (int)max_q_len;
+    // max_seq_len is unused by the prefill path (decode uses it for split
+    // computation); fill with max_q_len only to keep the POD struct defined.
+    p.max_seq_len = p.max_q_len;
+
+    p.causal_offset = (int)causal_offset;
+    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
+    if (p.use_mask) {
+        auto m = mask.value();
+        TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
+        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
+        if (m.dim() == 2) {
+            TORCH_CHECK(m.size(1) <= p.max_context_len, "mask kv_len mismatch");
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_h_stride = 0;
+            p.mask_q_stride = 0;
+        } else if (m.dim() == 4) {
+            TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch");
+            TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch");
+            TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_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 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;
+    }
+    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
+
     p.o = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
-
-    pack_mask(mask, p);
 }

csrc/kernels/attn_paged_decode.cu

@@ -3,23 +3,23 @@
 
 torch::Tensor attn_paged_decode(
     torch::Tensor q,
-    torch::Tensor page_table,
     torch::Tensor k_cache,
     torch::Tensor v_cache,
-    int64_t page_size,
+    torch::Tensor req_to_token,
-    int64_t kv_len,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    int64_t max_seq_len,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
-    double scale,
+    double scale
-    int64_t layout
 ) {
     PagedAttentionParams<bf16> p;
-    attn_pack_paged_params(q, page_table, k_cache, v_cache,
+    attn_pack_paged_decode_params(q, k_cache, v_cache,
-                           page_size, kv_len, mask, causal_offset, scale, layout, p);
+                                   req_to_token, req_pool_indices, kv_indptr,
+                                   max_seq_len, mask, causal_offset, scale, p);
 
-    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
+    auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
-    auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
+    p.o = (bf16*)O.data_ptr();
-    p.o = (bf16*)O_view.data_ptr();
 
     alloc_split_partials(p);
     DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
@@ -30,14 +30,14 @@ torch::Tensor attn_paged_decode(
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("attn_paged_decode", &attn_paged_decode,
         py::arg("q"),
-        py::arg("page_table"),
         py::arg("k_cache"),
         py::arg("v_cache"),
-        py::arg("page_size"),
+        py::arg("req_to_token"),
-        py::arg("kv_len"),
+        py::arg("req_pool_indices"),
+        py::arg("kv_indptr"),
+        py::arg("max_seq_len"),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,
-        py::arg("layout") = (int64_t)BHLD,
+        "SGLang-style paged decode: flat KV pool + req_to_token + kv_indptr.");
-        "Paged GQA decode — split-KV with direct page-table access.");
 }

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -5,6 +5,8 @@
 #include "attn_warp_utils.cuh"
 constexpr int PDC_CHUNK = 64;
 
+// Scalar paged decode (fallback for sm < 80, no tensor cores).
+// Reads K/V from flat pool via req_to_token indexing.
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
 __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
     int batch = blockIdx.x / p.kv_head;
@@ -15,8 +17,11 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
     int lane = threadIdx.x;
     int hd_per_thread = p.head_dim / 32;
 
+    const int seq_len = p.kv_indptr[batch + 1] - p.kv_indptr[batch];
+    const int64_t req_idx = p.req_pool_indices[batch];
+
     float q_reg[8];
-    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
+    int q_off = batch * p.q_stride_l + q_head * p.q_stride_h
               + lane * hd_per_thread * p.q_stride_d;
     #pragma unroll
     for (int i = 0; i < hd_per_thread; i++)
@@ -26,16 +31,19 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
 
     extern __shared__ __align__(16) bf16 k_smem[];
 
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
+    int chunks_total = (seq_len + PDC_CHUNK - 1) / PDC_CHUNK;
     int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
     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 + q_head * p.mask_h_stride;
+    const int mask_base = batch * p.mask_b_stride;
+    const int64_t pool_stride = (int64_t)p.kv_head * p.head_dim;
+    const int64_t head_off = (int64_t)kv_head * p.head_dim;
+    const int64_t rtt_stride = (int64_t)p.max_context_len;
 
     for (int ci = ch_begin; ci < ch_end; ci++) {
         int chunk_start = ci * PDC_CHUNK;
-        int this_chunk = min(PDC_CHUNK, p.kv_len - chunk_start);
+        int this_chunk = min(PDC_CHUNK, seq_len - chunk_start);
 
         int total = this_chunk * p.head_dim;
         for (int i = threadIdx.y * 32 + lane; i < total;
@@ -43,14 +51,9 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             int s = i / p.head_dim;
             int d_dim = i % p.head_dim;
             int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
+            int64_t slot = p.req_to_token[req_idx * rtt_stride + pos];
-            int page_offset = pos % p.page_size;
+            if (slot >= 0) {
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
+                int64_t off = slot * pool_stride + head_off + d_dim;
-            if (phys_page >= 0) {
-                int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
-                            + (int64_t)page_offset * p.kv_head * p.head_dim
-                            + (int64_t)kv_head * p.head_dim
-                            + d_dim;
                 k_smem[i] = p.k_cache[off];
             } else {
                 k_smem[i] = __float2bfloat16(0.0f);
@@ -85,17 +88,13 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             d = d * alpha + beta;
 
             int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
+            int64_t slot = p.req_to_token[req_idx * rtt_stride + pos];
-            int page_offset = pos % p.page_size;
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
             if (masked) {
                 #pragma unroll
                 for (int i = 0; i < hd_per_thread; i++)
                     acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
-            } else if (phys_page >= 0) {
+            } else if (slot >= 0) {
-                int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
+                int64_t v_base = slot * pool_stride + head_off;
-                               + (int64_t)page_offset * p.kv_head * p.head_dim
-                               + (int64_t)kv_head * p.head_dim;
                 #pragma unroll
                 for (int i = 0; i < hd_per_thread; i++)
                     acc_reg[i] = fmaf(acc_reg[i], alpha,
@@ -148,6 +147,6 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
     }
 
     float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
+    int o_off = batch * p.q_stride_l + q_head * p.q_stride_h + d * p.q_stride_d;
     p.o[o_off] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -5,12 +5,16 @@
 #include "attn_mma_utils.cuh"
 #include "attn_warp_utils.cuh"
 
-// Paged split-KV tensor-core decode via GQA head-packing.
+// SGLang-style split-KV tensor-core decode.
-// Reads K/V directly from the page pool through a page table — one tile
-// (BC=32) fits within a single page (page_size >= 32), so the page-table
-// lookup happens once per tile for cp.async.
 //
-// IsCausal and HasMask are compile-time bools.
+// Reads K/V directly from a flat pool [size, kv_head, head_dim] via
+// req_to_token indexing — no gather, no page-table dimension.
+// Each batch element has its own seq_len (from kv_indptr), eliminating
+// padding waste: short sequences only process the tiles they own.
+//
+// For decode (q_len=1), causal masking is implicit — each request attends
+// to [0, seq_len) which is exactly its valid range.  The IsCausal flag
+// is accepted for dispatch uniformity but does not change maxc.
 template <typename Traits, bool IsCausal, bool HasMask>
 __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
     const int lane = threadIdx.x;
@@ -22,6 +26,10 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     const int batch = blockIdx.y;
     const int split = blockIdx.z;
 
+    // Per-request seq_len from device-side kv_indptr — no padding.
+    const int seq_len = p.kv_indptr[batch + 1] - p.kv_indptr[batch];
+    const int64_t req_idx = p.req_pool_indices[batch];
+
     constexpr int MAX_G = 16;
     const int G_total = p.q_head / p.kv_head;
     const int g_begin = pass * MAX_G;
@@ -38,13 +46,14 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     }
     __syncwarp();
 
-    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
+    const int q_base = batch * p.q_stride_l + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
     unsigned Qa[Traits::KD][4];
-    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
+    load_q_mma_frags<Traits::KD>(p.q + q_base,
-                                  qra, qrb, va, vb, tid4, Qa);
+                                   p.q_stride_h, p.q_stride_d,
+                                   qra, qrb, va, vb, tid4, Qa);
 
     float Oacc[Traits::DN8][4];
     #pragma unroll
@@ -52,19 +61,19 @@ __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 tiles_total = (p.kv_len + Traits::BC - 1) / Traits::BC;
+    const int tiles_total = (seq_len + Traits::BC - 1) / Traits::BC;
     const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
     const int ti_begin = split * tiles_per_split;
     const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
 
-    const int64_t page_stride = (int64_t)p.page_size * p.kv_head * Traits::HEAD_DIM;
+    // Flat pool stride: [size, kv_head, head_dim] — contiguous.
-    const int64_t pos_stride  = (int64_t)p.kv_head * Traits::HEAD_DIM;
+    const int64_t pool_stride = (int64_t)p.kv_head * Traits::HEAD_DIM;
     const int64_t head_off    = (int64_t)kv_head * Traits::HEAD_DIM;
+    const int64_t rtt_stride  = (int64_t)p.max_context_len;
 
-    // ---- Load tile lambda: paged addressing ----
+    // ---- Load tile lambda: SGLang addressing ----
-    // Unified per-element page-table lookup. When page_size >= BC, all
+    // slot = req_to_token[req_idx * max_context_len + kc]
-    // elements in a tile share the same page, so the lookup is redundant
+    // gmem = k_cache[slot * pool_stride + head_off + d]
-    // but harmless (L1-cached). This avoids a branch on page_size.
     auto load_tile = [&](int ti, int buf) {
         int kv0 = ti * Traits::BC;
         bf16* dK = sK + buf * Traits::BC * Traits::LD;
@@ -74,16 +83,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
              i += Traits::NUM_THREADS * Traits::VEC) {
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
-            bool valid = (kc < p.kv_len);
+            bool valid = (kc < seq_len);
             if constexpr (HasMask) {
                 valid = valid && p.mask[batch * p.mask_b_stride + kc];
             }
-            int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
+            int64_t slot = valid ? p.req_to_token[req_idx * rtt_stride + kc] : 0;
-            valid = valid && (phys_page >= 0);
+            valid = valid && (slot >= 0);
-            int page_off = kc % p.page_size;
+            int64_t gmem_base = slot * pool_stride + head_off;
-            int64_t gmem_base = (int64_t)phys_page * page_stride
-                              + (int64_t)page_off * pos_stride
-                              + head_off;
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
             cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
             cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
@@ -91,10 +97,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         cp_async_commit();
     };
 
-    // ---- Multi-stage cp.async pipeline ----
-    // Prologue loads STAGES tiles; each loop iteration waits only for the
-    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
-    // tile loads stay in flight and overlap with the current tile's compute.
     constexpr int STAGES = Traits::STAGES;
     const int ntiles = ti_end - ti_begin;
 
@@ -111,8 +113,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
             Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
             Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
 
-        int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        // For decode, maxc = seq_len regardless of IsCausal — the valid
-        mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
+        // range [0, seq_len) IS the causal range (query is the last token).
+        mma_softmax_tile<Traits, HasMask>(kv0, seq_len, seq_len,
                                            0, 0,
                                            p.mask_b_stride, 0, 0,
                                            batch, 0,
@@ -136,7 +139,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
                 load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
         }
     } else {
-        // Fewer tiles than stages: load all, wait for all, process.
         for (int i = 0; i < ntiles; i++)
             load_tile(ti_begin + i, i);
         cp_async_wait_group<0>();
@@ -145,6 +147,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
             process_tile(it, it);
     }
 
+    // ---- write partials ----
     auto split_slot = [&](int h) -> size_t {
         size_t bh = (size_t)batch * p.q_head + h;
         return bh * MAX_SPLITS + split;

csrc/kernels/attn_paged_prefill.cu

@@ -0,0 +1,45 @@
+#include "attn_dispatchers.cuh"
+#include "attn_entry_utils.cuh"
+
+torch::Tensor attn_paged_prefill(
+    torch::Tensor q,
+    torch::Tensor k_cache,
+    torch::Tensor v_cache,
+    torch::Tensor req_to_token,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    torch::Tensor qo_indptr,
+    c10::optional<torch::Tensor> mask,
+    int64_t max_q_len,
+    int64_t causal_offset,
+    double scale
+) {
+    PagedAttentionParams<bf16> p;
+    attn_pack_paged_prefill_params(q, k_cache, v_cache,
+                                    req_to_token, req_pool_indices,
+                                    kv_indptr, qo_indptr, mask,
+                                    max_q_len, causal_offset, scale, p);
+
+    auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
+    p.o = (bf16*)O.data_ptr();
+
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p);
+    C10_CUDA_CHECK(cudaGetLastError());
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("attn_paged_prefill", &attn_paged_prefill,
+        py::arg("q"),
+        py::arg("k_cache"),
+        py::arg("v_cache"),
+        py::arg("req_to_token"),
+        py::arg("req_pool_indices"),
+        py::arg("kv_indptr"),
+        py::arg("qo_indptr"),
+        py::arg("mask") = py::none(),
+        py::arg("max_q_len"),
+        py::arg("causal_offset") = -1,
+        py::arg("scale") = 0.0,
+        "SGLang-style paged prefill: flat KV pool + ragged batch.");
+}

csrc/kernels/attn_paged_prefill_split_q_mma.cuh

@@ -0,0 +1,164 @@
+#pragma once
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+#include "attn_mma_utils.cuh"
+
+// SGLang-style split-Q tensor-core prefill.
+//
+// Reads K/V directly from a flat pool [size, kv_head, head_dim] via
+// req_to_token — no gather, no temporary tensor.  Supports ragged batches:
+// each request has its own q_len and kv_len, addressed via qo_indptr and
+// kv_indptr.
+//
+// Grid: (max_q_tiles, q_head, batch) — one batch element per blockIdx.z.
+// Blocks beyond a request's q_len exit early after writing sentinel-free
+// no-ops.  This avoids the binary-search approach and guarantees every Q
+// token is covered, even when q_len < BR*WARPS (e.g. decode-like prefill).
+//
+// Q layout: [total_q, q_head, head_dim] (3D, flattened across requests).
+// O layout: same as Q.
+//
+// IsCausal is a compile-time bool.  When true, each Q row qi (within its
+// request) attends to [0, causal_offset_b + qi + 1) where
+// causal_offset_b = kv_len_b - q_len_b (position of first Q token).
+template <typename Traits, bool IsCausal, bool HasMask>
+__global__ void paged_attn_prefill_split_q_mma_kernel(PagedAttentionParams<bf16> p) {
+    const int warp = threadIdx.x / 32;
+    const int lane = threadIdx.x % 32;
+    const int gid = lane >> 2;
+    const int tid4 = lane & 3;
+
+    const int q_head = blockIdx.y;
+    const int req_b  = blockIdx.z;
+    const int qrow0  = (blockIdx.x * Traits::WARPS + warp) * Traits::BR;
+
+    const int seq_len    = p.kv_indptr[req_b + 1] - p.kv_indptr[req_b];
+    const int q_len      = p.qo_indptr[req_b + 1] - p.qo_indptr[req_b];
+    const int causal_off = seq_len - q_len;
+    const int64_t req_idx = p.req_pool_indices[req_b];
+
+    // No per-warp early exit — all warps must participate in __syncthreads.
+    // Warps beyond q_len get zero-filled Q frags (va=vb=false) and skip output.
+    const int kv_head = q_head / (p.q_head / p.kv_head);
+
+    __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
+    __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
+
+    // Q base: offset by qo_indptr[req_b] to get absolute token address.
+    const int q_base = p.qo_indptr[req_b] * p.q_stride_l + q_head * p.q_stride_h;
+    const int qra = qrow0 + gid;
+    const int qrb = qrow0 + gid + 8;
+    const bool va = qra < q_len, vb = qrb < q_len;
+    unsigned Qa[Traits::KD][4];
+    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
+                                   qra, qrb, va, vb, tid4, Qa);
+
+    float Oacc[Traits::DN8][4];
+    #pragma unroll
+    for (int j = 0; j < Traits::DN8; j++)
+        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 int64_t pool_stride = (int64_t)p.kv_head * Traits::HEAD_DIM;
+    const int64_t head_off    = (int64_t)kv_head * Traits::HEAD_DIM;
+    const int64_t rtt_stride  = (int64_t)p.max_context_len;
+
+    const int tiles = (seq_len + Traits::BC - 1) / Traits::BC;
+    const int qr0 = qrow0 + gid;
+    const int qr1 = qrow0 + gid + 8;
+
+    // Causal tile-skip (dead code when IsCausal == false)
+    const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
+    const int block_max_kv =
+        blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
+        + causal_off;
+
+    int t_end = tiles - 1;
+    if constexpr (IsCausal) {
+        int bt = block_max_kv / Traits::BC;
+        if (bt < t_end) t_end = bt;
+    }
+
+    // ---- Load tile lambda: SGLang addressing ----
+    auto load_tile = [&](int ti, int buf) {
+        int kv0 = ti * Traits::BC;
+        bf16* dK = sK + buf * Traits::BC * Traits::LD;
+        bf16* dV = sV + buf * Traits::BC * Traits::LD;
+        #pragma unroll
+        for (int i = threadIdx.x * Traits::VEC; i < Traits::TOTAL;
+             i += Traits::NUM_THREADS * Traits::VEC) {
+            int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
+            int kc = kv0 + r;
+            bool valid = kc < seq_len;
+            int64_t slot = valid ? p.req_to_token[req_idx * rtt_stride + kc] : 0;
+            valid = valid && (slot >= 0);
+            int64_t gmem_base = slot * pool_stride + head_off;
+            int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
+            cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
+            cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
+        }
+        cp_async_commit();
+    };
+
+    // ---- Prologue + main loop (FA2-style double-buffer) ----
+    load_tile(0, 0);
+
+    for (int ti = 0; ti <= t_end; ti++) {
+        int buf = ti & 1;
+
+        cp_async_wait_group<0>();
+        __syncthreads();
+        if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
+
+        const bf16* bK = sK + buf * Traits::BC * Traits::LD;
+        const bf16* bV = sV + buf * Traits::BC * Traits::LD;
+        int kv0 = ti * Traits::BC;
+
+        if (!IsCausal || kv0 <= max_kv) {
+            float Sacc[Traits::NC8][4];
+            mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
+
+            #pragma unroll
+            for (int n8 = 0; n8 < Traits::NC8; n8++)
+                Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
+                Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
+
+            int maxc0 = IsCausal ? min(seq_len, causal_off + qr0 + 1)
+                                 : seq_len;
+            int maxc1 = IsCausal ? min(seq_len, causal_off + qr1 + 1)
+                                 : seq_len;
+            // HasMask: mask[batch, q_head, qi, kc] — kc is request-local.
+            mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
+                                              qr0, qr1,
+                                              p.mask_b_stride, p.mask_h_stride,
+                                              p.mask_q_stride,
+                                              req_b, q_head,
+                                              p.mask,
+                                              Sacc, Oacc, m0, m1, l0, l1, lane);
+
+            mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
+        }
+    }
+
+    // ---- write output: packed bf16x2 stores ----
+    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
+    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
+    const int o_base = p.qo_indptr[req_b] * p.q_stride_l + q_head * p.q_stride_h;
+    #pragma unroll
+    for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
+        int d = dn8 * 8 + 2 * tid4;
+        if (qr0 < q_len) {
+            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
+                                                       Oacc[dn8][1] * rl0);
+            *reinterpret_cast<__nv_bfloat162*>(
+                &p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
+        }
+        if (qr1 < q_len) {
+            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
+                                                       Oacc[dn8][3] * rl1);
+            *reinterpret_cast<__nv_bfloat162*>(
+                &p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
+        }
+    }
+}

csrc/tests/attn_paged_decode_test.cu

@@ -1,308 +0,0 @@
-// Compile:
-//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
-//        --extra-device-vectorization csrc/tests/attn_paged_decode_test.cu \
-//        -o /tmp/test_paged && /tmp/test_paged
-
-#include <cstring>
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-static void gather_kv_cpu(
-    const bf16* h_k_pool, const bf16* h_v_pool,
-    const int64_t* h_pt, int B, int Hkv, int kv_len,
-    int page_size, int head_dim,
-    bf16* h_k, bf16* h_v)
-{
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    size_t page_stride = (size_t)page_size * Hkv * head_dim;
-    for (int b = 0; b < B; b++) {
-        for (int pos = 0; pos < kv_len; pos++) {
-            int log_pg = pos / page_size;
-            int pg_off = pos % page_size;
-            int phys = (int)h_pt[b * max_pages + log_pg];
-            for (int h = 0; h < Hkv; h++) {
-                size_t src_base = (size_t)phys * page_stride
-                                + (size_t)pg_off * Hkv * head_dim
-                                + h * head_dim;
-                size_t dst_base = ((size_t)b * Hkv + h) * kv_len * head_dim
-                                  + (size_t)pos * head_dim;
-                memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16));
-                memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16));
-            }
-        }
-    }
-}
-
-template <int HEAD_DIM>
-static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int causal, int seed) {
-    printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d causal=%d ... ",
-           B, Hq, Hkv, kv_len, page_size, HEAD_DIM, causal);
-    fflush(stdout);
-
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    int n_phys_pages = B * max_pages;
-    int max_splits = 32;
-
-    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
-    size_t sz_o  = sz_q;
-    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
-    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
-    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
-    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
-
-    bf16 *d_q, *d_o_paged;
-    bf16 *d_k_pool, *d_v_pool;
-    int64_t* d_pt;
-    float *d_op, *d_ml;
-
-    cudaMalloc(&d_q, sz_q);
-    cudaMalloc(&d_o_paged, sz_o);
-    cudaMalloc(&d_k_pool, sz_kv);
-    cudaMalloc(&d_v_pool, sz_kv);
-    cudaMalloc(&d_pt, sz_pt);
-    cudaMalloc(&d_op, sz_op);
-    cudaMalloc(&d_ml, sz_ml);
-
-    srand(seed);
-    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
-
-    bf16* h_q = (bf16*)malloc(sz_q);
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
-        h_q[i] = __float2bfloat16(rnd());
-    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
-
-    bf16* h_k_pool = (bf16*)malloc(sz_kv);
-    bf16* h_v_pool = (bf16*)malloc(sz_kv);
-    size_t ps = (size_t)page_size * Hkv * HEAD_DIM;
-    for (int pg = 0; pg < n_phys_pages; pg++) {
-        for (int off = 0; off < page_size; off++) {
-            for (int h = 0; h < Hkv; h++) {
-                for (int d = 0; d < HEAD_DIM; d++) {
-                    float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d));
-                    size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM
-                                 + h * HEAD_DIM + d;
-                    h_k_pool[idx] = __float2bfloat16(v);
-                    h_v_pool[idx] = __float2bfloat16(v * 0.3f);
-                }
-            }
-        }
-    }
-    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
-    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
-
-    int64_t* h_pt = (int64_t*)malloc(sz_pt);
-    int next_pg = 0;
-    for (int b = 0; b < B; b++)
-        for (int p = 0; p < max_pages; p++)
-            h_pt[b * max_pages + p] = next_pg++;
-    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
-
-    bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
-    bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
-    gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont);
-
-    float* h_q_f = (float*)malloc((size_t)B * Hq * HEAD_DIM * sizeof(float));
-    float* h_k_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
-    float* h_v_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
-    for (int i = 0; i < B * kv_len * Hkv * HEAD_DIM; i++) {
-        h_k_f[i] = bf2f(h_k_cont[i]);
-        h_v_f[i] = bf2f(h_v_cont[i]);
-    }
-
-    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
-    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv,
-                      1, kv_len, HEAD_DIM, causal ? 0 : -1);
-
-    PagedAttentionParams<bf16> p;
-    p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
-    p.kv_len = kv_len; p.head_dim = HEAD_DIM;
-    p.use_mask = 0; p.causal_offset = causal ? 0 : -1;
-    set_default_paged_strides(p);
-    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.page_size = page_size; p.max_pages = max_pages;
-    p.page_table = d_pt;
-    p.k_cache = d_k_pool; p.v_cache = d_v_pool;
-    p.q = d_q; p.mask = nullptr; p.o = d_o_paged;
-    p.o_part = d_op; p.ml_part = d_ml;
-
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
-    cudaDeviceSynchronize();
-
-    bf16* h_o_bf16 = (bf16*)malloc(sz_o);
-    cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost);
-    float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
-        h_o_paged[i] = __bfloat162float(h_o_bf16[i]);
-
-    float max_abs_err = 0.0f, max_rel_err = 0.0f;
-    int bad_idx = -1;
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
-        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
-        if (e > max_abs_err) { max_abs_err = e; bad_idx = i; }
-        float rel = e / fmaxf(fabsf(h_o_ref[i]), 1e-8f);
-        if (rel > max_rel_err) max_rel_err = rel;
-    }
-
-    const float atol = 0.01f, rtol = 0.01f;
-    bool pass = true;
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
-        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
-        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
-    }
-
-    if (pass) {
-        printf("PASS (max_abs_err=%.4e max_rel_err=%.4e)\n", max_abs_err, max_rel_err);
-    } else {
-        int b = bad_idx / (Hq * HEAD_DIM);
-        int h = (bad_idx / HEAD_DIM) % Hq;
-        int d = bad_idx % HEAD_DIM;
-        printf("FAIL (max_abs_err=%.4e max_rel_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
-               max_abs_err, max_rel_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
-        printf("  ref[0..7]:");
-        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
-            printf(" %.4f", h_o_ref[i]);
-        printf("\n  got[0..7]:");
-        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
-            printf(" %.4f", h_o_paged[i]);
-        printf("\n");
-    }
-
-    free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt);
-    free(h_k_cont); free(h_v_cont);
-    free(h_q_f); free(h_k_f); free(h_v_f);
-    free(h_o_ref); free(h_o_bf16); free(h_o_paged);
-    cudaFree(d_q); cudaFree(d_o_paged);
-    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
-    cudaFree(d_op); cudaFree(d_ml);
-
-    return pass ? 0 : 1;
-}
-
-struct TestCase {
-    int head_dim;
-    int B, Hq, Hkv, kv_len, page_size, causal, seed;
-};
-
-static const TestCase TESTS[] = {
-    {128, 1, 1, 1, 8, 128, 0, 1},
-    {128, 1, 4, 4, 128, 128, 0, 2},
-    {128, 2, 4, 4, 256, 128, 0, 3},
-    {128, 1, 4, 1, 64, 64, 0, 4},
-    {128, 1, 8, 2, 64, 128, 0, 5},
-    {128, 2, 16, 4, 128, 128, 0, 6},
-    {64, 1, 4, 2, 32, 128, 0, 7},
-    {256, 1, 2, 1, 16, 128, 0, 8},
-    {32, 1, 4, 2, 32, 64, 0, 9},
-    {128, 3, 8, 2, 256, 128, 0, 10},
-    {128, 2, 32, 8, 512, 128, 0, 11},
-    {128, 1, 16, 2, 256, 128, 0, 12},
-    {128, 2, 32, 4, 512, 128, 0, 13},
-    {128, 2, 8, 2, 128, 128, 1, 14},   // causal
-};
-
-static int dispatch_test(const TestCase& tc) {
-    int r = 0;
-    dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
-        r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.causal, tc.seed);
-    });
-    return r;
-}
-
-template <int HEAD_DIM>
-static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    int n_phys_pages = B * max_pages;
-    int max_splits = 32;
-
-    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
-    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
-    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
-    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
-    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
-
-    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
-    int64_t* d_pt;
-    float *d_op, *d_ml;
-    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
-    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
-    cudaMalloc(&d_pt, sz_pt);
-    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
-
-    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
-    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
-    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
-    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
-    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
-    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
-
-    int64_t* h_pt = (int64_t*)malloc(sz_pt);
-    int next_pg = 0;
-    for (int b = 0; b < B; b++)
-        for (int p = 0; p < max_pages; p++)
-            h_pt[b * max_pages + p] = next_pg++;
-    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
-    free(h_pt);
-
-    PagedAttentionParams<bf16> pa;
-    pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
-    pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
-    pa.use_mask = 0; pa.causal_offset = -1;
-    set_default_paged_strides(pa);
-    pa.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    pa.page_size = page_size; pa.max_pages = max_pages;
-    pa.page_table = d_pt;
-    pa.k_cache = d_k_pool; pa.v_cache = d_v_pool;
-    pa.q = d_q; pa.mask = nullptr; pa.o = d_o;
-    pa.o_part = d_op; pa.ml_part = d_ml;
-
-    const int WARMUP = 10, ITERS = 100;
-    auto launch = [&]() {
-        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(pa); });
-    };
-    double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
-    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
-    double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-    BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
-
-    char cfg[64];
-    snprintf(cfg, sizeof(cfg),
-             "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
-             B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
-    print_bench_row(cfg, r);
-
-    free(tmp);
-    cudaFree(d_q); cudaFree(d_o);
-    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
-    cudaFree(d_op); cudaFree(d_ml);
-}
-
-static void bench() {
-    printf("\n===== PAGED DECODE BENCH =====\n");
-    print_bench_header();
-    bench_config<128>(1, 32, 4, 512, 128);
-    bench_config<128>(1, 32, 4, 1024, 128);
-    bench_config<128>(1, 32, 4, 2048, 128);
-    bench_config<128>(1, 32, 4, 4096, 128);
-    bench_config<128>(16, 32, 4, 2048, 128);
-    bench_config<128>(32, 32, 4, 1024, 128);
-}
-
-int main() {
-    int n = sizeof(TESTS) / sizeof(TESTS[0]);
-    int fail = 0;
-    printf("=== Paged Decode vs CPU reference (%d cases) ===\n\n", n);
-
-    for (int i = 0; i < n; i++) {
-        fail += dispatch_test(TESTS[i]);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("\nFAILED (%d/%d tests failed)\n", fail, n);
-        return fail;
-    }
-    printf("\nAll %d tests passed!\n", n);
-    bench();
-    return 0;
-}

csrc/tests/attn_paged_test.cu

@@ -0,0 +1,939 @@
+// Compile:
+//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
+//        --extra-device-vectorization csrc/tests/attn_paged_test.cu \
+//        -o /tmp/test_paged && /tmp/test_paged
+
+#include <cstring>
+#include <vector>
+#include "test_utils.cuh"
+#include "../kernels/attn_dispatchers.cuh"
+
+// ---- CPU reference: paged decode with variable seq_lens ----
+// Q: [B, Hq, D], K/V pool: [pool_size, Hkv, D]
+// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
+// kv_indptr: [B+1].  mask: [B, max_seq_len] bool (True=keep) or NULL.
+static void cpu_paged_decode_ref(
+    const float* Q, const float* K_pool, const float* V_pool,
+    const int64_t* req_to_token, const int64_t* req_pool_indices,
+    const int* kv_indptr, const bool* mask, int mask_b_stride,
+    int B, int Hq, int Hkv, int D, int max_ctx_len,
+    float* O)
+{
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hkv;
+    for (int b = 0; b < B; b++) {
+        int seq_len = kv_indptr[b + 1] - kv_indptr[b];
+        int64_t req_idx = req_pool_indices[b];
+        for (int h = 0; h < Hq; h++) {
+            int kv_h = h / n_rep;
+            float mv = -INFINITY, sv = 0.0f;
+            float accum[256] = {0.0f};
+            for (int kj = 0; kj < seq_len; kj++) {
+                if (mask && !mask[b * mask_b_stride + kj]) continue;
+                int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
+                float dot = 0.0f;
+                for (int d = 0; d < D; d++)
+                    dot += Q[(b * Hq + h) * D + d] *
+                           K_pool[slot * Hkv * D + kv_h * D + d];
+                dot *= scale;
+                float nm = fmaxf(mv, dot);
+                float a = expf(mv - nm);
+                float be = expf(dot - nm);
+                sv = sv * a + be;
+                for (int d = 0; d < D; d++)
+                    accum[d] = accum[d] * a +
+                               V_pool[slot * Hkv * D + kv_h * D + d] * be;
+                mv = nm;
+            }
+            float inv = 1.0f / sv;
+            for (int d = 0; d < D; d++)
+                O[(b * Hq + h) * D + d] = accum[d] * inv;
+        }
+    }
+}
+
+// ---- CPU reference: paged prefill with ragged batch ----
+// Q: [total_q, Hq, D], K/V pool: [pool_size, Hkv, D]
+// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
+// kv_indptr: [B+1], qo_indptr: [B+1].
+// mask: [B, max_q_len, max_seq_len] bool (True=keep, q-local + kv-local
+//   positions) or NULL.  Used only when causal==0 to apply an arbitrary
+//   attention mask on top of the (unused) causal logic.
+static void cpu_paged_prefill_ref(
+    const float* Q, const float* K_pool, const float* V_pool,
+    const int64_t* req_to_token, const int64_t* req_pool_indices,
+    const int* kv_indptr, const int* qo_indptr,
+    const bool* mask, int mask_q_stride, int mask_kv_stride,
+    int B, int Hq, int Hkv, int D, int max_ctx_len, int causal,
+    float* O)
+{
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hkv;
+    for (int b = 0; b < B; b++) {
+        int seq_len = kv_indptr[b + 1] - kv_indptr[b];
+        int q_len = qo_indptr[b + 1] - qo_indptr[b];
+        int causal_off = seq_len - q_len;
+        int64_t req_idx = req_pool_indices[b];
+        for (int h = 0; h < Hq; h++) {
+            int kv_h = h / n_rep;
+            for (int qi = 0; qi < q_len; qi++) {
+                float mv = -INFINITY, sv = 0.0f;
+                float accum[256] = {0.0f};
+                int lim = causal ? min(seq_len, causal_off + qi + 1) : seq_len;
+                for (int kj = 0; kj < lim; kj++) {
+                    if (mask && !mask[b * mask_q_stride * mask_kv_stride
+                                     + qi * mask_kv_stride + kj]) continue;
+                    int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
+                    float dot = 0.0f;
+                    for (int d = 0; d < D; d++)
+                        dot += Q[(qo_indptr[b] + qi) * Hq * D + h * D + d] *
+                               K_pool[slot * Hkv * D + kv_h * D + d];
+                    dot *= scale;
+                    float nm = fmaxf(mv, dot);
+                    float a = expf(mv - nm);
+                    float be = expf(dot - nm);
+                    sv = sv * a + be;
+                    for (int d = 0; d < D; d++)
+                        accum[d] = accum[d] * a +
+                                   V_pool[slot * Hkv * D + kv_h * D + d] * be;
+                    mv = nm;
+                }
+                float inv = 1.0f / sv;
+                for (int d = 0; d < D; d++)
+                    O[(qo_indptr[b] + qi) * Hq * D + h * D + d] = accum[d] * inv;
+            }
+        }
+    }
+}
+
+// ======================================================================
+// DECODE TEST
+// ======================================================================
+template <int HEAD_DIM>
+static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
+                            int causal, int seed) {
+    // Variable seq_lens per request
+    srand(seed);
+    std::vector<int> seq_lens(B);
+    for (int b = 0; b < B; b++)
+        seq_lens[b] = 8 + rand() % (max_seq - 8);
+    int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
+    int max_ctx = max_sl + 16;
+
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    printf("DECODE B=%d Hq=%d Hkv=%d D=%d seqs=[", B, Hq, Hkv, HEAD_DIM);
+    for (int b = 0; b < B; b++) printf("%d%s", seq_lens[b], b < B-1 ? "," : "");
+    printf("] causal=%d ... ", causal);
+    fflush(stdout);
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    // req_to_token: assign unique slots per request (scattered, not contiguous)
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    // req_pool_indices: pick B random request rows
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    // kv_indptr: prefix sum of seq_lens
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    // CPU reference
+    float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
+                            nullptr, 0,
+                            B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
+
+    // Kernel launch
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = B;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.02f, rtol = 0.02f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
+    else      printf("FAIL (max_err=%.4e)\n", max_err);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// DECODE WITH MASK TEST (regression: 2D mask on mixed seq_lens)
+// ======================================================================
+template <int HEAD_DIM>
+static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
+                                int seed) {
+    srand(seed);
+    std::vector<int> seq_lens(B);
+    for (int b = 0; b < B; b++)
+        seq_lens[b] = 8 + rand() % (max_seq - 8);
+    int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
+    int max_ctx = max_sl + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    printf("DECODE-MASK B=%d Hq=%d Hkv=%d D=%d max_sl=%d ... ", B, Hq, Hkv, HEAD_DIM, max_sl);
+    fflush(stdout);
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_mask = (size_t)B * max_sl * sizeof(bool);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi;
+    bool *d_mask;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_mask, sz_mask);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    // Mask: keep first half of each request's kv range, drop the rest —
+    // exercises the HasMask path with per-request seq_len.
+    bool* h_mask = (bool*)malloc(sz_mask);
+    for (int b = 0; b < B; b++)
+        for (int k = 0; k < max_sl; k++)
+            h_mask[b * max_sl + k] = (k < seq_lens[b]) && (k % 2 == 0);
+    cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
+
+    float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
+                            h_mask, max_sl,
+                            B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
+
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = B;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
+    p.causal_offset = -1; p.use_mask = 1;
+    p.mask = d_mask; p.mask_b_stride = max_sl;
+    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.02f, rtol = 0.02f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
+    else      printf("FAIL (max_err=%.4e)\n", max_err);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_mask);
+    cudaFree(d_op); cudaFree(d_ml);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// PREFILL TEST
+// ======================================================================
+template <int HEAD_DIM>
+static int run_prefill_test(int B, int Hq, int Hkv,
+                             std::vector<int>& q_lens,
+                             std::vector<int>& kv_lens,
+                             int causal, int seed) {
+    int total_q = 0;
+    int max_sl = 0;
+    for (int b = 0; b < B; b++) {
+        total_q += q_lens[b];
+        max_sl = max(max_sl, kv_lens[b]);
+    }
+    int max_ctx = max_sl + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    printf("PREFILL B=%d Hq=%d Hkv=%d D=%d q_lens=[", B, Hq, Hkv, HEAD_DIM);
+    for (int b = 0; b < B; b++) printf("%d%s", q_lens[b], b < B-1 ? "," : "");
+    printf("] kv_lens=[");
+    for (int b = 0; b < B; b++) printf("%d%s", kv_lens[b], b < B-1 ? "," : "");
+    printf("] causal=%d ... ", causal);
+    fflush(stdout);
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+
+    srand(seed);
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0;
+    for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_lens[b];
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    // CPU reference
+    float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
+                             nullptr, 0, 0,
+                             B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref);
+
+    // Kernel launch
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = total_q;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
+    int max_ql = 0;
+    for (int b = 0; b < B; b++) max_ql = max(max_ql, q_lens[b]);
+    p.max_q_len = max_ql;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.02f, rtol = 0.02f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
+    else      printf("FAIL (max_err=%.4e)\n", max_err);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_qoi); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// PREFILL WITH MASK TEST (regression: 4D causal mask on single request)
+// ======================================================================
+template <int HEAD_DIM>
+static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
+    srand(seed);
+    int B = 1;
+    int total_q = q_len;
+    int seq_len = q_len;  // pure prefill: kv_len == q_len
+    int max_ctx = seq_len + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    printf("PREFILL-MASK Hq=%d Hkv=%d D=%d q_len=%d ... ", Hq, Hkv, HEAD_DIM, q_len);
+    fflush(stdout);
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_mask = (size_t)B * q_len * q_len * sizeof(bool);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    bool *d_mask;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+    cudaMalloc(&d_mask, sz_mask);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    h_rpi[0] = 0;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0; h_kvi[1] = seq_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0; h_qoi[1] = q_len;
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    // 4D causal mask [B, 1, q_len, q_len], True=keep.
+    bool* h_mask = (bool*)malloc(sz_mask);
+    for (int qi = 0; qi < q_len; qi++)
+        for (int kj = 0; kj < q_len; kj++)
+            h_mask[qi * q_len + kj] = (kj <= qi);
+    cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
+
+    float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
+    // CPU ref with causal=0 so it consults the mask (not the causal flag).
+    cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
+                             h_mask, q_len, q_len,
+                             B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref);
+
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = total_q;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = q_len;
+    p.max_q_len = q_len;
+    p.causal_offset = -1; p.use_mask = 1;
+    p.mask = d_mask; p.mask_b_stride = q_len * q_len;
+    p.mask_h_stride = 0; p.mask_q_stride = q_len;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.02f, rtol = 0.02f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
+    else      printf("FAIL (max_err=%.4e)\n", max_err);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_qoi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+    cudaFree(d_mask);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// BENCH
+// ======================================================================
+template <int HEAD_DIM>
+static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
+    int max_ctx = seq_len + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B;
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++)
+            h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = B;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = seq_len;
+    p.causal_offset = 0; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    auto launch = [&]() {
+        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
+    };
+    // Decode: q_len=1, query is the last token → attends to all [0, seq_len).
+    // FLOPs = 2 * (QK^T + PV) = 4 * B * Hq * seq_len * D.
+    double flops = 4.0 * B * Hq * (double)seq_len * HEAD_DIM;
+    // HBM: K+V read (Q/O negligible for decode).
+    size_t nKV = (size_t)B * Hkv * seq_len * HEAD_DIM;
+    double bytes = 2.0 * nKV * sizeof(bf16);
+    BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
+
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "DEC B=%2d Hq=%2d Hk=%d kv=%4d D=%3d",
+             B, Hq, Hkv, seq_len, HEAD_DIM);
+    print_bench_row(cfg, r);
+
+    free(tmp); free(h_rtt); free(h_rpi); free(h_kvi);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
+}
+
+template <int HEAD_DIM>
+static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int causal) {
+    int total_q = B * q_len;
+    int max_ctx = kv_len + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B;
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+
+    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++)
+            h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0;
+    for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len;
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    PagedAttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM; p.total_q = total_q;
+    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
+    p.max_context_len = max_ctx; p.max_seq_len = kv_len;
+    p.total_q = total_q; p.max_q_len = q_len;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    auto launch = [&]() {
+        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
+    };
+    // FLOPs = 2 * (QK^T + PV) = 4 * effective_qk_pairs * Hq * D.
+    // Non-causal: effective = q_len * kv_len.
+    // Causal: Q row qi attends to [0, causal_off + qi + 1) where
+    //   causal_off = kv_len - q_len.  Total KV accesses per request:
+    //   sum_{qi=0}^{q_len-1} (kv_len - q_len + qi + 1)
+    //   = q_len * (kv_len - q_len) + q_len * (q_len + 1) / 2.
+    double eff_kv;
+    if (causal) {
+        eff_kv = (double)q_len * (kv_len - q_len)
+               + (double)q_len * (q_len + 1) / 2.0;
+    } else {
+        eff_kv = (double)q_len * kv_len;
+    }
+    double flops = 4.0 * B * Hq * eff_kv * HEAD_DIM;
+    // HBM: Q read + K read + V read + O write.
+    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
+    size_t nQ  = (size_t)total_q * Hq * HEAD_DIM;
+    double bytes = (2.0 * nQ + 2.0 * nKV) * sizeof(bf16);
+    BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PRE B=%d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d c=%d",
+             B, Hq, Hkv, q_len, kv_len, HEAD_DIM, causal);
+    print_bench_row(cfg, r);
+
+    free(tmp); free(h_rtt); free(h_rpi); free(h_kvi); free(h_qoi);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+}
+
+int main() {
+    int fail = 0;
+
+    // ===== DECODE TESTS =====
+    printf("=== Paged Decode Tests ===\n\n");
+    fail += run_decode_test<128>(1, 32, 4, 512, 0, 1);
+    fail += run_decode_test<128>(1, 32, 4, 1024, 0, 2);
+    fail += run_decode_test<128>(4, 32, 4, 512, 0, 3);
+    fail += run_decode_test<128>(8, 32, 4, 1024, 0, 4);
+    fail += run_decode_test<128>(4, 32, 8, 2048, 0, 5);
+    fail += run_decode_test<128>(1, 16, 1, 256, 0, 6);
+    fail += run_decode_test<128>(2, 8, 2, 512, 1, 7);
+    fail += run_decode_test<64>(1, 4, 2, 256, 0, 8);
+    fail += run_decode_test<256>(1, 2, 1, 256, 0, 9);
+    fail += run_decode_test<128>(16, 32, 4, 2048, 0, 10);
+    fail += run_decode_test<128>(32, 32, 4, 1024, 0, 11);
+
+    // Decode with 2D mask (regression: mixed seq_lens + HasMask)
+    fail += run_decode_mask_test<128>(2, 8, 2, 256, 30);
+    fail += run_decode_mask_test<128>(4, 32, 4, 512, 31);
+    fail += run_decode_mask_test<64>(2, 4, 2, 128, 32);
+
+    if (fail) { printf("\nFAILED decode tests\n"); return fail; }
+
+    // ===== PREFILL TESTS =====
+    printf("\n=== Paged Prefill Tests ===\n\n");
+    // Single request, pure prefill (q_len == kv_len)
+    {
+        std::vector<int> ql = {512};
+        std::vector<int> kl = {512};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 20);
+    }
+    {
+        std::vector<int> ql = {1024};
+        std::vector<int> kl = {1024};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 21);
+    }
+    {
+        std::vector<int> ql = {2048};
+        std::vector<int> kl = {2048};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 22);
+    }
+    // Ragged batch: different q_lens and kv_lens
+    {
+        std::vector<int> ql = {128, 256, 64};
+        std::vector<int> kl = {128, 256, 64};
+        fail += run_prefill_test<128>(3, 32, 4, ql, kl, 1, 23);
+    }
+    {
+        std::vector<int> ql = {64, 128, 256, 32};
+        std::vector<int> kl = {64, 128, 256, 32};
+        fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 24);
+    }
+    // Extend: kv_len > q_len (append to existing cache)
+    {
+        std::vector<int> ql = {64, 128};
+        std::vector<int> kl = {256, 512};
+        fail += run_prefill_test<128>(2, 32, 4, ql, kl, 1, 25);
+    }
+    // Non-causal
+    {
+        std::vector<int> ql = {256, 128};
+        std::vector<int> kl = {256, 128};
+        fail += run_prefill_test<128>(2, 32, 4, ql, kl, 0, 26);
+    }
+    // Single token (q_len=1 per request, like decode but via prefill path)
+    {
+        std::vector<int> ql = {1, 1, 1, 1};
+        std::vector<int> kl = {128, 256, 64, 512};
+        fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 27);
+    }
+    // D=64
+    {
+        std::vector<int> ql = {128, 64};
+        std::vector<int> kl = {128, 64};
+        fail += run_prefill_test<64>(2, 4, 2, ql, kl, 1, 28);
+    }
+    // D=256
+    {
+        std::vector<int> ql = {128, 64};
+        std::vector<int> kl = {128, 64};
+        fail += run_prefill_test<256>(2, 2, 1, ql, kl, 1, 29);
+    }
+
+    // Prefill with 4D causal mask (regression: single-request mask path)
+    fail += run_prefill_mask_test<128>(32, 4, 512, 40);
+    fail += run_prefill_mask_test<128>(32, 4, 1024, 41);
+    fail += run_prefill_mask_test<64>(4, 2, 256, 42);
+
+    if (fail) { printf("\nFAILED prefill tests\n"); return fail; }
+    printf("\nAll tests passed!\n");
+
+    // ===== BENCH =====
+    printf("\n===== PAGED DECODE BENCH =====\n");
+    print_bench_header();
+    bench_decode<128>(1, 32, 4, 512);
+    bench_decode<128>(1, 32, 4, 1024);
+    bench_decode<128>(1, 32, 4, 2048);
+    bench_decode<128>(1, 32, 4, 4096);
+    bench_decode<128>(4, 32, 4, 2048);
+    bench_decode<128>(16, 32, 4, 2048);
+    bench_decode<128>(32, 32, 4, 1024);
+
+    printf("\n===== PAGED PREFILL BENCH =====\n");
+    print_bench_header();
+    bench_prefill<128>(1, 32, 4, 512, 512, 0);
+    bench_prefill<128>(1, 32, 4, 1024, 1024, 0);
+    bench_prefill<128>(1, 32, 4, 2048, 2048, 0);
+    bench_prefill<128>(1, 32, 4, 2048, 2048, 1);
+    bench_prefill<128>(4, 32, 4, 2048, 2048, 1);
+    bench_prefill<128>(1, 32, 4, 4096, 4096, 1);
+
+    return 0;
+}

scripts/tools/benchmark.py

@@ -12,9 +12,14 @@ from astrai.model import AutoModel
 
 _DTYPES = ["bfloat16", "float16", "float32"]
 _CACHES = ["contiguous", "paged"]
-DEFAULT_CKPT = str(Path(__file__).resolve().parents[2] / "ckpt_bucket" / "kami-15bt")
+_BACKENDS = ["cuda", "torch_native"]
 CACHE_MAX_SEQ = 2048
 
+_BACKEND_MAP = {
+    "cuda": ATTN_BACKEND.CUDA,
+    "torch_native": ATTN_BACKEND.TORCH_NATIVE,
+}
+
 
 class BenchmarkResult:
     def __init__(
@@ -42,12 +47,14 @@ class GenerationBenchmark:
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         cache_type: str = "contiguous",
+        backend: ATTN_BACKEND = ATTN_BACKEND.CUDA,
     ):
         self.device = device
         self.dtype = dtype
         self.cache_type = cache_type
         self.model = model
         self.config = config
+        self.backend = backend
 
     def _make_pool(self, batch_size: int) -> PagePool:
         return PagePool(
@@ -82,7 +89,7 @@ class GenerationBenchmark:
         kv_cache = pool.bind_tasks(
             task_ids, [prompt_len] * batch_size, self.device, start_pos=0
         )
-        with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
+        with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
                 input_mask=input_mask,
@@ -105,7 +112,7 @@ class GenerationBenchmark:
             total_len, device=self.device
         )
         kv_cache = pool.bind_tasks(task_ids, [seq_len + 1] * batch_size, self.device)
-        with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
+        with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
                 input_mask=input_mask,
@@ -121,6 +128,11 @@ class GenerationBenchmark:
     ) -> BenchmarkResult:
         import time
 
+        pool = self._make_pool(batch_size)
+        task_ids = [f"bench_prefill_{i}" for i in range(batch_size)]
+        for tid in task_ids:
+            pool.task_alloc(tid, list(range(prompt_length)))
+
         input_ids = torch.randint(
             0, self.config.vocab_size, (batch_size, prompt_length), device=self.device
         )
@@ -129,16 +141,32 @@ class GenerationBenchmark:
             .unsqueeze(0)
             .expand(batch_size, -1)
         )
+        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
+            prompt_length, device=self.device
+        )
+        kv_cache = pool.bind_tasks(
+            task_ids, [prompt_length] * batch_size, self.device, start_pos=0
+        )
 
         for _ in range(3):
-            with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
+            with torch.inference_mode(), attn_backend(self.backend):
-                self.model(input_ids, position_ids=position_ids)
+                self.model(
+                    input_ids,
+                    input_mask=input_mask,
+                    kv_cache=kv_cache,
+                    position_ids=position_ids,
+                )
 
         torch.cuda.synchronize()
         t0 = time.perf_counter()
         for _ in range(num_trials):
-            with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
+            with torch.inference_mode(), attn_backend(self.backend):
-                self.model(input_ids, position_ids=position_ids)
+                self.model(
+                    input_ids,
+                    input_mask=input_mask,
+                    kv_cache=kv_cache,
+                    position_ids=position_ids,
+                )
         torch.cuda.synchronize()
         elapsed = time.perf_counter() - t0
         tokens = batch_size * prompt_length * num_trials
@@ -208,6 +236,17 @@ def print_benchmark_result(result: BenchmarkResult) -> None:
 @click.option(
     "--cache", type=click.Choice(_CACHES), default="contiguous", help="KV cache type."
 )
+@click.option(
+    "--backend",
+    type=click.Choice(_BACKENDS),
+    default="cuda",
+    help="Attention backend.",
+)
+@click.option(
+    "--compare",
+    is_flag=True,
+    help="Run both backends and print side-by-side speed comparison.",
+)
 @click.option("--batch_size", type=int, default=4, help="Batch size.")
 @click.option("--prompt_length", type=int, default=512, help="Prompt length.")
 @click.option("--gen_length", type=int, default=128, help="Generation length.")
@@ -216,13 +255,16 @@ def print_benchmark_result(result: BenchmarkResult) -> None:
 @click.option("--decode_only", is_flag=True, help="Decode benchmark only.")
 @click.option(
     "--ckpt",
-    default=DEFAULT_CKPT,
+    required=True,
+    type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
     help="Checkpoint directory.",
 )
 def benchmark_command(
     device: str,
     dtype: str,
     cache: str,
+    backend: str,
+    compare: bool,
     batch_size: int,
     prompt_length: int,
     gen_length: int,
@@ -244,32 +286,38 @@ def benchmark_command(
     model.to(device=device, dtype=dtype_map[dtype])
     model.eval()
 
-    bench = GenerationBenchmark(
+    backends = _BACKENDS if compare else [backend]
-        model=model,
-        config=config,
-        device=device,
-        dtype=dtype_map[dtype],
-        cache_type=cache,
-    )
 
-    click.secho(f"Benchmark: device={device} dtype={dtype}", bold=True)
+    for name in backends:
-
+        bench = GenerationBenchmark(
-    if not decode_only:
+            model=model,
-        result = bench.run_prefill_benchmark(
+            config=config,
-            batch_size=batch_size,
+            device=device,
-            prompt_length=prompt_length,
+            dtype=dtype_map[dtype],
-            num_trials=num_trials,
+            cache_type=cache,
+            backend=_BACKEND_MAP[name],
         )
-        print_benchmark_result(result)
 
-    if not prefill_only:
+        click.secho(
-        result = bench.run_decoding_benchmark(
+            f"Benchmark: device={device} dtype={dtype} backend={name}", bold=True
-            batch_size=batch_size,
-            prompt_length=prompt_length,
-            gen_length=gen_length,
-            num_trials=num_trials,
         )
-        print_benchmark_result(result)
+
+        if not decode_only:
+            result = bench.run_prefill_benchmark(
+                batch_size=batch_size,
+                prompt_length=prompt_length,
+                num_trials=num_trials,
+            )
+            print_benchmark_result(result)
+
+        if not prefill_only:
+            result = bench.run_decoding_benchmark(
+                batch_size=batch_size,
+                prompt_length=prompt_length,
+                gen_length=gen_length,
+                num_trials=num_trials,
+            )
+            print_benchmark_result(result)
 
 
 if __name__ == "__main__":

setup.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import subprocess
 import sys
 import warnings
@@ -89,14 +90,30 @@ if _should_build():
                 super().build_extensions()
                 return
 
+            # Each subprocess gets its own build-temp / build-lib so the
+            # ninja files (build.ninja, .ninja_log) never race.  The built
+            # .so files are then collected into the parent's build_lib so the
+            # normal setuptools copy steps (inplace / editable wheel) work.
             names = [e.name for e in self.extensions]
             env = {**os.environ, "BUILD_PARALLEL": "1"}
+            base = os.path.join("build", "parallel")
+            os.makedirs(base, exist_ok=True)
             procs = {}
             for i in range(0, len(names), max_workers):
                 batch = names[i : i + max_workers]
                 for name in batch:
                     e = {**env, "ASTRAI_BUILD_SINGLE_EXT": name}
-                    cmd = [sys.executable, __file__, "build_ext", "--inplace"]
+                    tag = name.replace(".", "_")
+                    subdir = os.path.join(base, tag)
+                    cmd = [
+                        sys.executable,
+                        __file__,
+                        "build_ext",
+                        "--build-temp",
+                        os.path.join(subdir, "temp"),
+                        "--build-lib",
+                        os.path.join(subdir, "lib"),
+                    ]
                     procs[name] = subprocess.Popen(
                         cmd, env=e, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
                     )
@@ -108,6 +125,19 @@ if _should_build():
                             f"parallel build failed for {name} "
                             f"(exit {procs[name].returncode})"
                         )
+                    self._collect_extensions(
+                        os.path.join(base, name.replace(".", "_"), "lib")
+                    )
+
+        def _collect_extensions(self, sub_lib):
+            src = os.path.join(sub_lib, "astrai", "extension", "lib")
+            if not os.path.isdir(src):
+                return
+            dst = os.path.join(self.build_lib, "astrai", "extension", "lib")
+            os.makedirs(dst, exist_ok=True)
+            for f in os.listdir(src):
+                if f.endswith(".so"):
+                    shutil.copy2(os.path.join(src, f), os.path.join(dst, f))
 
     cmdclass["build_ext"] = ParallelBuildExtension
 

tests/extension/test_backend_equivalence.py

@@ -13,18 +13,26 @@ from tests.extension.conftest import D, skip_no_kernel
 
 @skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
-    """Training forward (kv_cache=None) should produce identical logits."""
+    """Training forward (kv_cache=None) should produce identical logits.
+
+    CudaBackend is inference-only: it raises when kv_cache is None.  Training
+    must use TorchNativeBackend (the default).  Verify the torch path is
+    stable and that CudaBackend rejects the training path explicitly.
+    """
+    import pytest
+
     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()
+    with pytest.raises(RuntimeError, match="does not support training"):
-    assert diff == 0.0, f"Training forward diff {diff} should be 0"
+        with attn_backend(ATTN_BACKEND.CUDA):
+            with torch.no_grad():
+                model(input_ids)
+
+    assert out_torch["logits"].shape[0] == 2
 
 
 @skip_no_kernel