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.loader import is_available
+from astrai.extension.attention_ops import (
+    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``
-    with ``page_size=1`` (each token slot is a single-token "page").
-    The ``req_to_token`` table serves directly as the page table.
+    Decode path: writes K/V to the flat pool, then calls
+    ``attn_paged_decode`` with req_to_token + kv_indptr.
 
-    Prefill path: writes K/V to cache, gathers full-sequence K/V via
-    indirect indexing (same as TorchNativeBackend), then calls
-    ``attn_prefill``.
+    Prefill path: writes K/V to the flat pool, then calls
+    ``attn_paged_prefill`` with ragged-batch support via qo_indptr +
+    kv_indptr.
 
-    Training path (``kv_cache is None``): calls ``attn_prefill`` directly
-    on the projected q/k/v.
+    ``kv_cache is None`` (training) is not handled — use
+    ``TorchNativeBackend`` for training.
 
-    Falls back to ``TorchNativeBackend`` for any path where the
-    corresponding CUDA kernel is not available.
+    Raises ``RuntimeError`` if the required 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"):
-            return self._fallback.fwd_decode(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
+        if kv_cache is None:
+            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
         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:
-            page_table = kv_cache.page_table
-        else:
-            page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
+        kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
+        kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
 
-        k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
-        v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
-
-        if q.size(0) == 1:
-            mask = None
-        elif kv_cache.decode_mask is not None:
+        mask = None
+        if b > 1 and 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,
-            page_table,
-            k_cache,
-            v_cache,
-            page_size=1,
-            kv_len=max_len,
+            q_3d,
+            kv_cache.k_buffer[layer_id],
+            kv_cache.v_buffer[layer_id],
+            kv_cache.req_to_token,
+            kv_cache.req_pool_indices,
+            kv_indptr,
+            kv_cache.max_len,
             mask=mask,
             is_causal=is_causal,
         )
-
-        out = out.flatten(2)
-        return out
+        return out.unsqueeze(1).flatten(2)
 
     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"):
-                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
-            )
+            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
         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
-        indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-        pos_mask = (
-            torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
-        )
-        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]
+        b = q.size(0)
+        q_len = q.size(1)
 
-        out = attn_prefill(q, k_full, v_full, mask=attn_mask, is_causal=is_causal)
-        return out.flatten(2)
+        kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
+        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,
-    kv_len: int,
+    req_to_token: torch.Tensor,
+    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)
-        page_table: [batch, max_pages] (int64)
-        k_cache: [n_pages, page_size, n_kv_heads, head_dim] (bf16)
+        q: [batch, n_heads, head_dim] (bf16, 3D — no seq dim)
+        k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
         v_cache: same as k_cache
-        page_size: tokens per page
-        kv_len: actual sequence length per request
-        mask: 2D [batch, kv_len] or 3D [batch, 1, kv_len] (bool, True=keep)
+        req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
+        req_pool_indices: [batch] (int64) — rows into req_to_token
+        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,
-        kv_len,
+        req_to_token,
+        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] = {}