perf: precompute ragged Q tile scheduling

astrai/extension/backend/attention.py

@@ -610,6 +610,8 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_pool_indices,
             kv_cache.kv_indptr,
             kv_cache.qo_indptr,
+            kv_cache.q_tile_to_batch,
+            kv_cache.q_tile_to_index,
             attn_mask,
             is_causal=is_causal,
         )

astrai/extension/ops/attention.py

@@ -150,6 +150,8 @@ def attn_paged_prefill(
     req_pool_indices: torch.Tensor,
     kv_indptr: torch.Tensor,
     qo_indptr: torch.Tensor,
+    q_tile_to_batch: torch.Tensor,
+    q_tile_to_index: torch.Tensor,
     mask: Optional[torch.Tensor] = None,
     is_causal: bool = False,
 ) -> torch.Tensor:
@@ -167,6 +169,8 @@ def attn_paged_prefill(
         req_pool_indices: [batch] (int32)
         kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
         qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
+        q_tile_to_batch: [num_q_tiles] (int32) — request index per Q tile
+        q_tile_to_index: [num_q_tiles] (int32) — local Q tile index per request
         mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
         is_causal: apply causal mask
 
@@ -183,6 +187,8 @@ def attn_paged_prefill(
         req_pool_indices,
         kv_indptr,
         qo_indptr,
+        q_tile_to_batch,
+        q_tile_to_index,
         mask,
         causal_offset=causal_offset,
     )

astrai/inference/cache/buffer.py

@@ -99,6 +99,8 @@ class KVCache:
     max_len: int = 0
     kv_indptr: Optional[Tensor] = None
     qo_indptr: Optional[Tensor] = None
+    q_tile_to_batch: Optional[Tensor] = None
+    q_tile_to_index: Optional[Tensor] = None
     decode_o_part: Optional[Tensor] = None
     decode_ml_part: Optional[Tensor] = None
     decode_out: Optional[Tensor] = None

astrai/inference/cache/pool.py

@@ -25,7 +25,7 @@ from astrai.inference.cache.strategy import (
     RadixCache,
     TaskCacheState,
 )
-from astrai.inference.workspace import InferenceWorkspace
+from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
 
 # Re-export everything so existing ``from astrai.inference.cache import ...``
 # continues to work unchanged after the file split.
@@ -217,6 +217,21 @@ class PagePool:
                 torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(0)
             )
             qo_indptr = workspace.qo_indptr[: b + 1]
+            tile_batches = []
+            tile_indices = []
+            for batch, q_len in enumerate(q_lens):
+                n_tiles = (q_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS
+                tile_batches.extend([batch] * n_tiles)
+                tile_indices.extend(range(n_tiles))
+            n_tiles = len(tile_batches)
+            workspace.q_tile_to_batch[:n_tiles].copy_(
+                torch.tensor(tile_batches, dtype=torch.int32, device=device)
+            )
+            workspace.q_tile_to_index[:n_tiles].copy_(
+                torch.tensor(tile_indices, dtype=torch.int32, device=device)
+            )
+            q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
+            q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
             decode_o_part = decode_ml_part = decode_out = None
         else:
             # ---- decode: out_cache_loc is a single column (last position) ----
@@ -226,6 +241,7 @@ class PagePool:
             out_cache_loc = ocl_buf[:b].reshape(-1)
             workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
             qo_indptr = workspace.qo_indptr[: b + 1]
+            q_tile_to_batch = q_tile_to_index = None
             decode_o_part = getattr(workspace, "decode_o_part", None)
             decode_ml_part = getattr(workspace, "decode_ml_part", None)
             decode_out = getattr(workspace, "decode_out", None)
@@ -240,6 +256,8 @@ class PagePool:
             max_len=max(seq_lens),
             kv_indptr=kv_indptr,
             qo_indptr=qo_indptr,
+            q_tile_to_batch=q_tile_to_batch,
+            q_tile_to_index=q_tile_to_index,
             decode_o_part=decode_o_part,
             decode_ml_part=decode_ml_part,
             decode_out=decode_out,

astrai/inference/workspace.py

@@ -10,6 +10,7 @@ import torch
 from torch import Tensor
 
 _MAX_SPLITS = 32
+Q_TILE_ROWS = 64
 
 
 class InferenceWorkspace:
@@ -83,6 +84,13 @@ class InferenceWorkspace:
         self.qo_indptr = torch.empty(
             (max_batch_size + 1,), dtype=torch.int32, device=device
         )
+        max_q_tiles = max_batch_size * ((max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS)
+        self.q_tile_to_batch = torch.empty(
+            (max_q_tiles,), dtype=torch.int32, device=device
+        )
+        self.q_tile_to_index = torch.empty(
+            (max_q_tiles,), dtype=torch.int32, device=device
+        )
         self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
         self.out_cache_loc = torch.empty(
             (max_batch_size, 1), dtype=torch.int32, device=device

csrc/kernels/attn_common.h

@@ -13,7 +13,7 @@ enum TensorLayout : int {
 //   - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).
 //   - Paged (SGLang-style): flat pool [size, kv_head, head_dim] + req_to_token.
 // Each kernel selects the addressing via a KVSource policy (see
-// attn_kv_source.cuh); a given call only touches the fields of one mode, so
+// attn_layout_policies.cuh); a given call only touches the fields of one mode, so
 // this is a POD shared by both paths rather than two parallel structs that
 // drift out of sync.
 template<typename T, typename AT = float>
@@ -59,6 +59,9 @@ struct AttentionParams {
     const int* __restrict__ req_pool_indices;     // [batch]
     const int* __restrict__ kv_indptr;             // [batch + 1]
     const int* __restrict__ qo_indptr;             // [batch + 1] or nullptr for decode
+    const int* __restrict__ q_tile_to_batch;        // [num_q_tiles], prefill only
+    const int* __restrict__ q_tile_to_index;        // [num_q_tiles], prefill only
+    int num_q_tiles;
     int max_context_len; // req_to_token stride (dim 1)
 
     // Decode split-KV workspace

csrc/kernels/attn_decode_split_kv.cuh

@@ -2,7 +2,7 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
-#include "attn_kv_source.cuh"
+#include "attn_layout_policies.cuh"
 #include "attn_warp_utils.cuh"
 constexpr int DC_CHUNK = 64;
 

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -2,7 +2,7 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
-#include "attn_kv_source.cuh"
+#include "attn_layout_policies.cuh"
 #include "attn_mma_utils.cuh"
 #include "attn_warp_utils.cuh"
 

csrc/kernels/attn_dispatchers.cuh

@@ -3,7 +3,7 @@
 // No torch dependency; pure CUDA.
 //
 // The paged and contiguous kernels are unified by the KVSource policy
-// (ContigKV / PagedKV from attn_kv_source.cuh), so each launcher struct
+// (ContigKV / PagedKV from attn_layout_policies.cuh), so each launcher struct
 // below is templated on KV and the paged dispatch is just the same launcher
 // instantiated with PagedKV.  Only the grid/split math differs, and that is
 // covered by KV::host_q_len / KV::host_kv_len.
@@ -11,7 +11,7 @@
 #include <cuda_runtime.h>
 #include <algorithm>
 #include "attn_warp_utils.cuh"
-#include "attn_kv_source.cuh"
+#include "attn_layout_policies.cuh"
 #include "attn_prefill_split_q.cuh"
 #include "attn_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
@@ -80,31 +80,32 @@ template <> struct PrefillConfigMap<128, true>  : PrefillKernelConfig<32> {};
 template <> struct PrefillConfigMap<256, false> : PrefillKernelConfig<16> {};
 template <> struct PrefillConfigMap<256, true>  : PrefillKernelConfig<16> {};
 
-template <typename KV>
+template <typename QSchedule, typename KV>
 struct PrefillLauncherMMA {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
     static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
         using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
         using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
         constexpr int ROWS = Traits::BR * Config::WARPS;
-        dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
+        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
-                  KV::kPaged ? 1 : p.batch);
+                  QSchedule::host_grid_batch(p));
         dim3 block(Traits::NUM_THREADS);
-        attn_prefill_split_q_mma_kernel<Traits, KV, IsCausal, HasMask>
+        attn_prefill_split_q_mma_kernel<Traits, QSchedule, KV, IsCausal, HasMask>
             <<<grid, block, 0, stream>>>(p);
     }
 };
 #endif
 
-template <typename KV>
+template <typename QSchedule, typename KV>
 struct PrefillLauncherScalar {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
     static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
-        constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32;
+        constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 64, P_BC = 32;
-        dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
+        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
-                  KV::kPaged ? 1 : p.batch);
+                  QSchedule::host_grid_batch(p));
         dim3 block(G, ROWS);
-        attn_prefill_split_q_kernel_t<HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>
+        attn_prefill_split_q_kernel_t<HEAD_DIM, QSchedule, KV, G, ROWS, P_BC,
+                                      IsCausal, HasMask>
             <<<grid, block, 0, stream>>>(p);
     }
 };
@@ -115,12 +116,14 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p, cudaStream_t strea
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
+    using Launcher = PrefillLauncherMMA<DenseQSchedule, ContigKV>;
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
-                         PrefillLauncherMMA<ContigKV>::template launch,
+                         Launcher::template launch,
                          HEAD_DIM, p, stream);
 #else
+    using Launcher = PrefillLauncherScalar<DenseQSchedule, ContigKV>;
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
-                         PrefillLauncherScalar<ContigKV>::template launch,
+                         Launcher::template launch,
                          HEAD_DIM, p, stream);
 #endif
 }
@@ -131,12 +134,14 @@ static inline void dispatch_paged_prefill(AttentionParams<bf16>& p, cudaStream_t
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
+    using Launcher = PrefillLauncherMMA<PackedQSchedule, PagedKV>;
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
-                         PrefillLauncherMMA<PagedKV>::template launch,
+                         Launcher::template launch,
                          HEAD_DIM, p, stream);
 #else
+    using Launcher = PrefillLauncherScalar<PackedQSchedule, PagedKV>;
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
-                         PrefillLauncherScalar<PagedKV>::template launch,
+                         Launcher::template launch,
                          HEAD_DIM, p, stream);
 #endif
 }

csrc/kernels/attn_entry_utils.cuh

@@ -227,6 +227,8 @@ inline void attn_pack_paged_prefill_params(
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
     torch::Tensor qo_indptr,
+    torch::Tensor q_tile_to_batch,
+    torch::Tensor q_tile_to_index,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -237,6 +239,7 @@ inline void attn_pack_paged_prefill_params(
     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_tile_to_batch.is_cuda() && q_tile_to_index.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");
@@ -245,6 +248,10 @@ inline void attn_pack_paged_prefill_params(
                 "req_pool_indices must be int32");
     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(q_tile_to_batch.dtype() == torch::kInt32,
+                "q_tile_to_batch must be int32");
+    TORCH_CHECK(q_tile_to_index.dtype() == torch::kInt32,
+                "q_tile_to_index 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]");
@@ -261,6 +268,10 @@ inline void attn_pack_paged_prefill_params(
     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]");
+    TORCH_CHECK(q_tile_to_batch.dim() == 1 && q_tile_to_index.dim() == 1,
+                "Q tile mappings must be 1D");
+    TORCH_CHECK(q_tile_to_batch.size(0) == q_tile_to_index.size(0),
+                "Q tile mappings must have equal length");
 
     p.q_l_stride = (int)q.stride(0);
     p.q_h_stride = (int)q.stride(1);
@@ -273,6 +284,9 @@ inline void attn_pack_paged_prefill_params(
     p.req_pool_indices = req_pool_indices.data_ptr<int>();
     p.kv_indptr = kv_indptr.data_ptr<int>();
     p.qo_indptr = qo_indptr.data_ptr<int>();
+    p.q_tile_to_batch = q_tile_to_batch.data_ptr<int>();
+    p.q_tile_to_index = q_tile_to_index.data_ptr<int>();
+    p.num_q_tiles = (int)q_tile_to_batch.size(0);
     p.max_context_len = (int)req_to_token.size(1);
 
     p.causal_offset = (int)causal_offset;

csrc/kernels/attn_layout_policies.cuh

@@ -3,13 +3,10 @@
 #include "attn_common.h"
 
 // ============================================================================
-// KVSource policies — the single dimension along which the paged and
+// Attention layout policies keep Q scheduling independent from K/V storage.
-// non-paged attention kernels differ.  Each kernel is templated on one of
+// DenseQSchedule / PackedQSchedule map blocks to Q tiles; ContigKV / PagedKV
-// these (ContigKV / PagedKV) and stays fully generic: the policy owns every
+// resolve logical K/V positions to physical addresses. This lets the shared
-// place where "where does K/V live" and "what is this request's seq_len"
+// kernels compose Q layout and K/V storage without coupling the two concerns.
-// are answered.  All methods are __host__ __device__ so the same policy
-// serves both the device kernels (addressing, seq_len) and the host-side
-// launchers (grid / split computation).
 //
 //   ContigKV:  K/V are dense [batch, kv_head, kv_len, head_dim] tensors.
 //              Params fields used: k, v, kv_stride_*, kv_len, q_len,
@@ -25,11 +22,76 @@
 // (e.g. the req_pool_indices global read) element-by-element.
 // ============================================================================
 
-// Every policy method is static + callable from both host and device code.
+#define HOST_FORCEINLINE static __host__ __forceinline__
+#define DEVICE_FORCEINLINE static __device__ __forceinline__
 #define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__
 
 using bf16 = __nv_bfloat16;
 
+// ============================================================================
+// Q scheduling policies
+//
+// Map CUDA blocks to request-local Q tiles independently of K/V storage.
+// Dense tensors encode the request in blockIdx.z; packed ragged tensors use
+// a compact precomputed work map indexed by blockIdx.x.
+// ============================================================================
+
+struct DenseQSchedule {
+    HOST_FORCEINLINE int host_q_blocks(
+        const AttentionParams<bf16>& p, int rows) {
+        return (p.q_len + rows - 1) / rows;
+    }
+
+    HOST_FORCEINLINE int host_grid_batch(
+        const AttentionParams<bf16>& p) {
+        return p.batch;
+    }
+
+    DEVICE_FORCEINLINE void map_block(
+        const AttentionParams<bf16>&, int& batch, int& q_tile) {
+        batch = blockIdx.z;
+        q_tile = blockIdx.x;
+    }
+
+    DEVICE_FORCEINLINE int q_len(
+        const AttentionParams<bf16>& p, int) {
+        return p.q_len;
+    }
+
+    DEVICE_FORCEINLINE int q_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return batch * p.q_b_stride + q_head * p.q_h_stride;
+    }
+};
+
+struct PackedQSchedule {
+    HOST_FORCEINLINE int host_q_blocks(
+        const AttentionParams<bf16>& p, int) {
+        return p.num_q_tiles;
+    }
+
+    HOST_FORCEINLINE int host_grid_batch(
+        const AttentionParams<bf16>&) {
+        return 1;
+    }
+
+    DEVICE_FORCEINLINE void map_block(
+        const AttentionParams<bf16>& p, int& batch, int& q_tile) {
+        batch = p.q_tile_to_batch[blockIdx.x];
+        q_tile = p.q_tile_to_index[blockIdx.x];
+    }
+
+    DEVICE_FORCEINLINE int q_len(
+        const AttentionParams<bf16>& p, int batch) {
+        return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
+    }
+
+    DEVICE_FORCEINLINE int q_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride;
+    }
+};
+
 // Hoisted per-(batch, kv_head) addressing context.
 struct KVContext {
     int kv_base;          // contig: batch*kv_b_stride + kv_head*kv_h_stride
@@ -56,59 +118,40 @@ struct KVAddr {
 struct ContigKV {
     static constexpr bool kPaged = false;
 
-    // host-side length hooks (grid + split computation in the launchers)
+    HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
-    HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& p, int rows) {
-        return (p.q_len + rows - 1) / rows;
-    }
-    template <int ROWS>
-    HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>&,
-                                         int flat_tile, int grid_batch,
-                                         int& batch, int& q_tile) {
-        batch = grid_batch;
-        q_tile = flat_tile;
-        return true;
-    }
-    HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
         return p.kv_len;
     }
 
-    // prefill: element offset of the request's Q rows (kernel adds qrow*q_l_stride)
-    HOST_DEV_FORCEINLINE int q_base(
-        const AttentionParams<bf16>& p, int batch, int q_head) {
-        return batch * p.q_b_stride + q_head * p.q_h_stride;
-    }
     // decode: same offset (q_len == 1, so there is no row stride component)
-    HOST_DEV_FORCEINLINE int q_decode_base(
+    DEVICE_FORCEINLINE int q_decode_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
         return batch * p.q_b_stride + q_head * p.q_h_stride;
     }
 
-    HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int) {
         return p.kv_len;
     }
-    HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int causal_offset(
-        return p.q_len;
+        const AttentionParams<bf16>& p, int, int) {
-    }
-    HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
         return p.causal_offset;
     }
     // decode: exclusive bound of the single query's attend range
-    HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int) {
         return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1);
     }
 
     template <int HEAD_DIM>
-    HOST_DEV_FORCEINLINE KVContext make_ctx(
+    DEVICE_FORCEINLINE KVContext make_ctx(
         const AttentionParams<bf16>& p, int batch, int kv_head) {
         KVContext c = {};
         c.kv_base = batch * p.kv_b_stride + kv_head * p.kv_h_stride;
         return c;
     }
-    HOST_DEV_FORCEINLINE int resolve_token(
+    DEVICE_FORCEINLINE int resolve_token(
         const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
         return valid ? kc : -1;
     }
-    HOST_DEV_FORCEINLINE KVAddr kv_addr_from_token(
+    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
         const AttentionParams<bf16>& p, const KVContext& c, int token, int d) {
         const bool valid = token >= 0;
         const int safe_token = valid ? token : 0;
@@ -123,58 +166,30 @@ struct ContigKV {
 struct PagedKV {
     static constexpr bool kPaged = true;
 
-    HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& p, int rows) {
+    HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
-        // sum(ceil(q_len[b] / rows)) <= ceil(total_q / rows) + batch - 1.
-        return (p.q_len + rows - 1) / rows + p.batch - 1;
-    }
-    template <int ROWS>
-    HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>& p,
-                                         int flat_tile, int,
-                                         int& batch, int& q_tile) {
-        int tile_base = 0;
-        for (int b = 0; b < p.batch; ++b) {
-            int len = p.qo_indptr[b + 1] - p.qo_indptr[b];
-            int tiles = (len + ROWS - 1) / ROWS;
-            if (flat_tile < tile_base + tiles) {
-                batch = b;
-                q_tile = flat_tile - tile_base;
-                return true;
-            }
-            tile_base += tiles;
-        }
-        return false;
-    }
-    HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
         return p.max_context_len;
     }
 
-    // prefill: Q rows start at qo_indptr[batch] (ragged batch base)
-    HOST_DEV_FORCEINLINE int q_base(
-        const AttentionParams<bf16>& p, int batch, int q_head) {
-        return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride;
-    }
     // decode: Q is [batch, q_head, head_dim], so batch is the outer row
-    HOST_DEV_FORCEINLINE int q_decode_base(
+    DEVICE_FORCEINLINE int q_decode_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
         return batch * p.q_l_stride + q_head * p.q_h_stride;
     }
 
-    HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
         return p.kv_indptr[batch + 1] - p.kv_indptr[batch];
     }
-    HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int causal_offset(
-        return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
+        const AttentionParams<bf16>& p, int batch, int q_len) {
-    }
+        return kv_len(p, batch) - q_len;
-    HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
-        return kv_len(p, batch) - q_len(p, batch);
     }
     // decode: the query is the last token, so [0, seq_len) IS its causal range
-    HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
+    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
         return kv_len(p, batch);
     }
 
     template <int HEAD_DIM>
-    HOST_DEV_FORCEINLINE KVContext make_ctx(
+    DEVICE_FORCEINLINE KVContext make_ctx(
         const AttentionParams<bf16>& p, int batch, int kv_head) {
         KVContext c = {};
         c.req_idx = p.req_pool_indices[batch];
@@ -183,11 +198,11 @@ struct PagedKV {
         c.head_off = (int64_t)kv_head * HEAD_DIM;
         return c;
     }
-    HOST_DEV_FORCEINLINE int resolve_token(
+    DEVICE_FORCEINLINE int resolve_token(
         const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
         return valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : -1;
     }
-    HOST_DEV_FORCEINLINE KVAddr kv_addr_from_token(
+    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
         const AttentionParams<bf16>& p, const KVContext& c, int slot, int d) {
         const bool valid = slot >= 0;
         const int safe_slot = valid ? slot : 0;
@@ -195,30 +210,3 @@ struct PagedKV {
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
     }
 };
-
-// ---- Q-block mapping ----
-// Contiguous grids map directly to (batch, q_tile). Paged grids flatten the
-// ragged Q tiles, so one thread resolves the request and broadcasts it.
-template <int ROWS, typename KV>
-__device__ __forceinline__ bool map_q_block(
-    const AttentionParams<bf16>& p, int& batch, int& q_tile) {
-    if constexpr (!KV::kPaged) {
-        batch = blockIdx.z;
-        q_tile = blockIdx.x;
-        return true;
-    } else {
-        __shared__ int mapped_batch;
-        __shared__ int mapped_q_tile;
-
-        if ((threadIdx.x | threadIdx.y) == 0) {
-            mapped_batch = -1;
-            KV::template map_q_tile<ROWS>(
-                p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
-        }
-        __syncthreads();
-
-        batch = mapped_batch;
-        q_tile = mapped_q_tile;
-        return batch >= 0;
-    }
-}

csrc/kernels/attn_paged_prefill.cu

@@ -9,6 +9,8 @@ torch::Tensor attn_paged_prefill(
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
     torch::Tensor qo_indptr,
+    torch::Tensor q_tile_to_batch,
+    torch::Tensor q_tile_to_index,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale
@@ -19,7 +21,8 @@ torch::Tensor attn_paged_prefill(
     AttentionParams<bf16> p;
     attn_pack_paged_prefill_params(q, k_cache, v_cache,
                                      req_to_token, req_pool_indices,
-                                    kv_indptr, qo_indptr, mask,
+                                     kv_indptr, qo_indptr,
+                                     q_tile_to_batch, q_tile_to_index, mask,
                                     causal_offset, scale, p);
 
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
@@ -39,6 +42,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("req_pool_indices"),
         py::arg("kv_indptr"),
         py::arg("qo_indptr"),
+        py::arg("q_tile_to_batch"),
+        py::arg("q_tile_to_index"),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,

csrc/kernels/attn_prefill_split_q.cuh

@@ -2,7 +2,7 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
-#include "attn_kv_source.cuh"
+#include "attn_layout_policies.cuh"
 
 using bf16 = __nv_bfloat16;
 
@@ -32,13 +32,13 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) {
     }
 }
 
-template <int HEAD_DIM, typename KV, int G, int ROWS, int P_BC, bool IsCausal, bool HasMask>
+template <int HEAD_DIM, typename QSchedule, typename KV, int G, int ROWS, int P_BC,
+          bool IsCausal, bool HasMask>
 __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     constexpr int DPT = HEAD_DIM / G;
 
     int batch, q_tile;
-    if (!map_q_block<ROWS, KV>(p, batch, q_tile))
+    QSchedule::map_block(p, batch, q_tile);
-        return;
 
     int q_head = blockIdx.y;
     int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
@@ -47,8 +47,8 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
 
     // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
     const int seq_len    = KV::kv_len(p, batch);
-    const int q_len      = KV::q_len(p, batch);
+    const int q_len      = QSchedule::q_len(p, batch);
-    const int causal_off = KV::causal_offset(p, batch);
+    const int causal_off = KV::causal_offset(p, batch, q_len);
     const int kv_head = q_head / (p.q_head / p.kv_head);
     const KVContext kctx = KV::template make_ctx<HEAD_DIM>(p, batch, kv_head);
 
@@ -56,7 +56,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
 
     // Q: stride-based load [batch, q_head, q_len, head_dim]
-    const int q_base = KV::q_base(p, batch, q_head);
+    const int q_base = QSchedule::q_base(p, batch, q_head);
     float qreg[DPT];
     if (q_row < q_len) {
         int q_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride;

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -2,7 +2,7 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
-#include "attn_kv_source.cuh"
+#include "attn_layout_policies.cuh"
 #include "attn_mma_utils.cuh"
 
 // Tensor-core prefill flash attention (raw mma.sync PTX), unified across
@@ -16,7 +16,7 @@
 // dead branches in the inner compute loop (FA2-style).
 //
 // Traits = KernelTraits<HEAD_DIM, BC, WARPS=4, STAGES=2>.
-template <typename Traits, typename KV, bool IsCausal, bool HasMask>
+template <typename Traits, typename QSchedule, typename KV, bool IsCausal, bool HasMask>
 __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int warp = threadIdx.x / 32;
     const int lane = threadIdx.x % 32;
@@ -25,15 +25,14 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
 
     const int q_head = blockIdx.y;
     int batch, q_tile;
-    if (!map_q_block<Traits::BR * Traits::WARPS, KV>(p, batch, q_tile))
+    QSchedule::map_block(p, batch, q_tile);
-        return;
     const int kv_head = q_head / (p.q_head / p.kv_head);
     const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
 
     // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
     const int seq_len    = KV::kv_len(p, batch);
-    const int q_len      = KV::q_len(p, batch);
+    const int q_len      = QSchedule::q_len(p, batch);
-    const int causal_off = KV::causal_offset(p, batch);
+    const int causal_off = KV::causal_offset(p, batch, q_len);
     const KVContext kctx = KV::template make_ctx<Traits::HEAD_DIM>(p, batch, kv_head);
 
     // Static shared memory: double-buffered K/V (no sQ — Q goes direct
@@ -42,7 +41,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
 
     // Load Q fragments straight from global into mma A-operand layout.
-    const int q_base = KV::q_base(p, batch, q_head);
+    const int q_base = QSchedule::q_base(p, batch, q_head);
     const int qra = qrow0 + gid;
     const int qrb = qrow0 + gid + 8;
     const bool va = qra < q_len, vb = qrb < q_len;
@@ -138,7 +137,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     // ---- 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 = KV::q_base(p, batch, q_head);
+    const int o_base = QSchedule::q_base(p, batch, q_head);
     #pragma unroll
     for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;

csrc/tests/attn_paged_test.cu

@@ -12,6 +12,26 @@
 struct PagedDecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_decode<H>(p, 0); } };
 struct PagedPrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_prefill<H>(p, 0); } };
 
+static int make_q_tile_mapping(const std::vector<int>& q_lens,
+                               int** d_batch, int** d_tile) {
+    constexpr int ROWS = 64;
+    std::vector<int> h_batch;
+    std::vector<int> h_tile;
+    for (int b = 0; b < (int)q_lens.size(); ++b) {
+        int n_tiles = (q_lens[b] + ROWS - 1) / ROWS;
+        for (int tile = 0; tile < n_tiles; ++tile) {
+            h_batch.push_back(b);
+            h_tile.push_back(tile);
+        }
+    }
+    size_t bytes = h_batch.size() * sizeof(int);
+    cudaMalloc(d_batch, bytes);
+    cudaMalloc(d_tile, bytes);
+    cudaMemcpy(*d_batch, h_batch.data(), bytes, cudaMemcpyHostToDevice);
+    cudaMemcpy(*d_tile, h_tile.data(), bytes, cudaMemcpyHostToDevice);
+    return (int)h_batch.size();
+}
+
 // ---- 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]
@@ -483,6 +503,9 @@ static int run_prefill_test(int B, int Hq, int Hkv,
                              nullptr, 0, 0,
                              B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref);
 
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
     // Kernel launch
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
@@ -497,6 +520,8 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = 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.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
     p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
     dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
@@ -523,6 +548,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     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_qtb); cudaFree(d_qti);
     return pass ? 0 : 1;
 }
 
@@ -619,6 +645,10 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
                              h_mask, q_len, q_len,
                              B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref);
 
+    std::vector<int> q_lens(B, q_len);
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
     p.head_dim = HEAD_DIM;
@@ -632,6 +662,8 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = 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.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
     p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
     dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
@@ -659,6 +691,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     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);
+    cudaFree(d_qtb); cudaFree(d_qti);
     return pass ? 0 : 1;
 }
 
@@ -786,6 +819,10 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
     for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len;
     cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
 
+    std::vector<int> q_lens(B, q_len);
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
     p.head_dim = HEAD_DIM;
@@ -798,6 +835,8 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
     p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = 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.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
     p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
     auto launch = [&]() {
@@ -827,6 +866,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
     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);
+    cudaFree(d_qtb); cudaFree(d_qti);
 }
 
 int main() {

docs/developer/architecture.md

@@ -1437,7 +1437,7 @@ classDiagram
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
-| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, attn_paged_prefill, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
+| **astrai.extension** | `backend` policy package, `ops` kernel-wrapper package, AttentionBackend, TorchNativeBackend, CudaBackend, FlashAttnBackend, attention, attn_backend, ATTN_BACKEND, apply_rotary_emb, is_available | Stable API over attention/rotary execution policy and optional CUDA kernels |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
@@ -1468,7 +1468,7 @@ classDiagram
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
-5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (cuda > flash > torch priority; `ASTR_BACKEND` env var overrides default; `TorchNativeBackend` fallback). Rotary embedding auto-dispatches to CUDA kernel when available, else torch complex multiply.
+5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. `astrai.extension.backend` owns attention/rotary dispatch, fallback, and KV cache policy; it calls the stateless compiled-kernel wrappers in `astrai.extension.ops`. Attention uses cuda > flash > torch priority unless explicitly selected by `ASTR_BACKEND` or `attn_backend()`. Rotary embedding auto-dispatches to the CUDA op when supported, else torch complex multiply.
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (`MmapStore`/`JsonlStore`) loads data with explicit `_length` and multi-segment `_data`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata; `CheckpointCallback` performs rank-0 training saves, with extra state saved as `{key}.pt`
@@ -1476,4 +1476,4 @@ classDiagram
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-08-02
+> Document Update Time: 2026-08-16

docs/developer/cuda_kernels.md

@@ -19,9 +19,10 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
 | Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) |
 | Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) |
 
-> The paged and non-paged paths are ONE kernel templated on a `KVSource`
+> The paged and non-paged paths share one kernel body. Prefill is templated on
-> policy (`ContigKV` / `PagedKV` in `attn_kv_source.cuh`); there are no
+> an independent Q schedule (`DenseQSchedule` / `PackedQSchedule`) and KV
-> separate `attn_paged_*.cuh` files anymore.
+> source (`ContigKV` / `PagedKV`); decode only needs the KV source. There are
+> no separate `attn_paged_*.cuh` files.
 
 ### Rotary Embedding Kernel
 
@@ -81,6 +82,113 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
 
 Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all five kernel targets in parallel via `cmake --build -j N`.
 
+## Python Extension Architecture
+
+The Python extension package separates low-level kernel bindings from execution
+policy:
+
+```text
+astrai/extension/
+├── __init__.py             # Stable public API
+├── loader.py               # Optional compiled-module discovery and loading
+├── ops/
+│   ├── attention.py        # Stateless attention kernel wrappers
+│   └── rotary.py           # Stateless rotary kernel wrapper
+└── backend/
+    ├── attention.py        # Backend selection, KV cache I/O, and fallback
+    └── rotary.py           # Per-call CUDA/torch rotary dispatch
+```
+
+The dependency direction is one-way:
+
+```text
+model / inference
+       |
+       v
+extension public API
+       |
+       v
+backend policy  --->  ops wrappers  --->  loader  --->  compiled .so
+       |
+       +----------->  torch / flash-attn fallback
+```
+
+`ops` must not import `backend`. This keeps direct kernel bindings independent
+of model, cache, fallback, and backend-selection policy.
+
+### Ops Layer
+
+`astrai.extension.ops` is the low-level boundary around compiled extensions:
+
+- Wrappers are stateless and map Python arguments to pybind or
+  `torch.library.custom_op` calls.
+- Wrappers validate kernel availability and raise `RuntimeError` when a
+  requested extension was not built.
+- Wrappers do not choose another implementation, gather KV cache entries, or
+  decide whether an input is supported by a backend.
+- Tests that specifically exercise a compiled kernel may import from
+  `astrai.extension.ops`.
+
+For example, `attn_prefill(...)` means "run this CUDA kernel" rather than "run
+attention using the best available implementation":
+
+```python
+from astrai.extension.ops import attn_prefill
+
+output = attn_prefill(q, k, v, mask=mask, is_causal=True)
+```
+
+If the kernel is unavailable, this call fails. Callers that need fallback and
+capability dispatch must use the public `attention(...)` entry point instead.
+
+### Backend Layer
+
+`astrai.extension.backend` owns execution policy:
+
+- It selects CUDA, FlashAttention, or torch-native attention.
+- It checks per-call constraints such as dtype, shape, head dimension, cache
+  availability, and installed optional dependencies.
+- It owns KV cache writes and reads because those operations differ by backend.
+- It provides torch fallbacks and raises when an explicitly requested backend
+  cannot handle a call.
+- Rotary dispatch follows the same boundary without a backend class: the
+  policy layer chooses the fused op for supported inference calls and otherwise
+  uses the autograd-compatible torch implementation.
+
+Normal model and inference code should import the stable API from
+`astrai.extension`:
+
+```python
+from astrai.extension import ATTN_BACKEND, attention, attn_backend
+
+output = attention(q, k, v, kv_cache=cache, layer_id=layer_id, fwd="decode")
+
+with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
+    output = attention(q, k, v)
+```
+
+The package root re-exports the supported high-level API and selected direct
+kernel wrappers. Internal code should use `astrai.extension.backend` only when
+it needs a backend type or policy implementation, and `astrai.extension.ops`
+only when it deliberately requires one exact kernel.
+
+### Placement Rules
+
+When extending this package:
+
+| Change | Location |
+|--------|----------|
+| Add a pybind call for a compiled kernel | `astrai/extension/ops/` |
+| Add argument translation required by the compiled ABI | `astrai/extension/ops/` |
+| Add capability checks or implementation selection | `astrai/extension/backend/` |
+| Add a torch or third-party fallback | `astrai/extension/backend/` |
+| Add attention KV cache behavior | `astrai/extension/backend/attention.py` |
+| Expose a supported user-facing symbol | `astrai/extension/__init__.py` |
+
+Imports belong at module scope. Optional dependencies such as `flash_attn` may
+use a module-level guarded import. Type-only imports that would create a runtime
+cycle belong under `TYPE_CHECKING`.
+
 ## Attention Backend
 
 `astrai/extension/backend/attention.py` provides the backend abstraction:
@@ -102,7 +210,11 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
-`CudaBackend` falls back to `FlashAttnBackend` (when flash-attn is installed and supports the input dtype) or `TorchNativeBackend` otherwise.
+The `attention(...)` policy entry point falls back to `FlashAttnBackend` (when
+flash-attn is installed and supports the call) or `TorchNativeBackend` when the
+automatically selected CUDA backend cannot handle an input. An explicit
+`ASTR_BACKEND` or `attn_backend(...)` selection is strict and raises instead of
+silently switching implementations.
 
 ### Rotary Backend
 
@@ -127,6 +239,47 @@ mask:      2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
 
 Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`.
 
+### Q Scheduling and KV Addressing
+
+Prefill separates Q work scheduling from KV storage:
+
+- `DenseQSchedule` maps a rectangular grid directly with
+  `batch = blockIdx.z` and `q_tile = blockIdx.x`.
+- `PackedQSchedule` consumes a compact work map for a packed
+  `[total_q, q_heads, head_dim]` tensor.
+- `ContigKV` and `PagedKV` only provide KV lengths and translate logical KV
+  positions into physical addresses. They do not schedule Q blocks.
+
+For ragged Q lengths `[70, 10, 130]` and 64 rows per Q tile, cache binding
+builds:
+
+```text
+qo_indptr       = [0, 70, 80, 210]
+q_tile_to_batch = [0, 0, 1, 2, 2, 2]
+q_tile_to_index = [0, 1, 0, 0, 1, 2]
+```
+
+Paged prefill launches:
+
+```text
+grid.x = num_q_tiles  # 6, exactly the valid ragged work items
+grid.y = q_heads
+grid.z = 1
+```
+
+Each block resolves its request and request-local tile in O(1):
+
+```cpp
+batch = q_tile_to_batch[blockIdx.x];
+q_tile = q_tile_to_index[blockIdx.x];
+```
+
+The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent
+`qo_indptr` / `kv_indptr` entries for that request's Q and KV lengths. This
+avoids the previous per-block linear scan over the batch, shared-memory
+broadcast, mapping barrier, and upper-bound grid with potentially invalid
+blocks.
+
 ## Standalone Testing
 
 Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example:
@@ -174,7 +327,7 @@ csrc/
 │   ├── attn_decode_split_kv_mma.cuh     # Split-KV + MMA variant (contig + paged)
 │   ├── attn_prefill_split_q.cuh         # Split-Q variant (contig + paged via KVSource)
 │   ├── attn_prefill_split_q_mma.cuh     # Split-Q + MMA variant (contig + paged)
-│   ├── attn_kv_source.cuh               # KVSource policies (ContigKV / PagedKV)
+│   ├── attn_layout_policies.cuh          # Q schedules and KVSource policies
 │   ├── attn_dispatchers.cuh             # Kernel dispatch macros + KV-templated launchers
 │   ├── attn_entry_utils.cuh             # Entry point helpers
 │   ├── attn_mma_utils.cuh               # MMA utilities
@@ -187,4 +340,4 @@ csrc/
 
 Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16

docs/developer/internals.md

@@ -176,12 +176,19 @@ Three-layer separation (SGLang-inspired):
 
 ### Attention Backend
 
+The extension package separates mechanism from policy:
+
+- `astrai/extension/ops/` contains stateless wrappers that invoke one exact compiled kernel and fail when it is unavailable.
+- `astrai/extension/backend/` owns capability checks, implementation selection, fallback, and KV cache I/O.
+- Model and inference code use the stable `astrai.extension` API instead of selecting ops directly.
+
 Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`):
 
-- **`CudaBackend`** (default): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `FlashAttnBackend` when dtype unsupported.
+- **`CudaBackend`** (default when supported): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool).
 - **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
 - **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
-- Default priority: cuda > flash > torch. Set `ASTR_BACKEND=cuda|torch_native|flash` to override.
+- The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
+- `ASTR_BACKEND=cuda|torch_native|flash` and `attn_backend(...)` are explicit selections; incompatible calls raise instead of silently changing backend.
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
 
@@ -196,6 +203,8 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`.
 
+Direct imports from `astrai.extension.ops` are reserved for low-level kernel tests and code that intentionally requires a specific compiled implementation. They do not provide fallback.
+
 ## Mask Algorithm Internals
 
 ### Template mode (`template: true`)
@@ -253,4 +262,4 @@ total_steps          = (batches_per_replica // grad_accum_steps) * n_epoch
 
 This accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset.
 
-> Document Update Time: 2026-08-02
+> Document Update Time: 2026-08-16

docs/guides/inference.md

@@ -61,6 +61,14 @@ Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k`
 
 ## Attention Backend
 
+Inference code calls the policy API exported by `astrai.extension`. The
+extension implementation is split into two layers:
+
+- `astrai.extension.backend` owns capability checks, backend selection,
+  fallback, and KV cache I/O.
+- `astrai.extension.ops` contains direct wrappers around compiled CUDA kernels;
+  these wrappers raise if a kernel is unavailable and do not fall back.
+
 Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the model via `AttentionBackend` ABC:
 
 ```
@@ -70,8 +78,9 @@ AttentionBackend (ABC)
   └── TorchNativeBackend   SDPA + indirect KV cache gather (always-available fallback)
 ```
 
-Default priority: cuda > flash > torch.  Set ``ASTR_BACKEND=cuda|torch_native|flash``
+Default priority is cuda > flash > torch. Automatic selection may choose a
-to override.
+compatible fallback for a particular call. Set
+`ASTR_BACKEND=cuda|torch_native|flash` to require one backend process-wide.
 
 Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
@@ -82,12 +91,20 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
+Environment and context selections are strict: if the selected backend cannot
+handle the call, inference raises an error rather than silently switching.
+
 `CudaBackend` decode path: writes K/V to cache, then calls `attn_paged_decode` with `page_size=1` — the `req_to_token` table serves directly as the page table, each token slot is a single-token "page". No explicit K/V gather needed.
 
 `CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
 Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`.
 
+This fallback is performed by the public `attention(...)` policy entry point
+only when no backend was explicitly selected. Import from
+`astrai.extension.ops` only for direct kernel tests or when failure on a missing
+kernel is the intended behavior.
+
 ### Rotary Embedding Backend
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches:
@@ -329,4 +346,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16

tests/inference/test_cache.py

@@ -286,6 +286,16 @@ def test_page_pool_contiguous_bind_tasks_prefill():
     assert kv.req_pool_indices.dtype == torch.int32
 
 
+def test_page_pool_bind_tasks_builds_compact_q_tile_mapping():
+    pool = _make_contiguous_pool(max_batch_size=3, max_seq_len=256)
+
+    kv = pool.bind_tasks([0, 1, 2], [70, 10, 130], _ws(pool), start_pos=0)
+
+    assert kv.qo_indptr.tolist() == [0, 70, 80, 210]
+    assert kv.q_tile_to_batch.tolist() == [0, 0, 1, 2, 2, 2]
+    assert kv.q_tile_to_index.tolist() == [0, 1, 0, 0, 1, 2]
+
+
 def test_page_pool_contiguous_bind_tasks_decode():
     pool = _make_contiguous_pool()
     task_cache = _make_task_cache(pool)