perf: use int32 paged KV indices

- store page-table, request-row, and cache-location indices as int32
- preserve CUDA graph replay with bit-exact logits and KV cache coverage
- improve B=1 decode latency by 1-6% across 1K-32K contexts on L20

astrai/extension/attention_backend.py

@@ -531,8 +531,8 @@ class CudaBackend(AttentionBackend):
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
         loc = kv_cache.out_cache_loc[:, 0]
-        kv_cache.k_buffer[layer_id].index_copy_(0, loc, k[:, 0])
+        kv_cache.k_buffer[layer_id, loc] = k[:, 0]
-        kv_cache.v_buffer[layer_id].index_copy_(0, loc, v[:, 0])
+        kv_cache.v_buffer[layer_id, loc] = v[:, 0]
 
         q_3d = q.squeeze(1)
 
@@ -566,12 +566,8 @@ class CudaBackend(AttentionBackend):
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
         loc = kv_cache.out_cache_loc.reshape(-1)
-        kv_cache.k_buffer[layer_id].index_copy_(
+        kv_cache.k_buffer[layer_id, loc] = k.reshape(-1, k.size(2), k.size(3))
-            0, loc, k.reshape(-1, k.size(2), k.size(3))
+        kv_cache.v_buffer[layer_id, loc] = v.reshape(-1, v.size(2), v.size(3))
-        )
-        kv_cache.v_buffer[layer_id].index_copy_(
-            0, loc, v.reshape(-1, v.size(2), v.size(3))
-        )
 
         b = q.size(0)
         q_len = q.size(1)

astrai/extension/attention_ops.py

@@ -113,8 +113,8 @@ def attn_paged_decode(
         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
-        req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
+        req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
-        req_pool_indices: [batch] (int64) — rows into req_to_token
+        req_pool_indices: [batch] (int32) — rows into req_to_token
         kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
         mask: 2D [batch, max_context_len] (bool, True=keep) or None
         is_causal: apply causal mask
@@ -163,8 +163,8 @@ def attn_paged_prefill(
         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_to_token: [num_reqs, max_context_len] (int32)
-        req_pool_indices: [batch] (int64)
+        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
         mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None

astrai/inference/cache/buffer.py

@@ -27,7 +27,7 @@ class ReqToTokenPool:
         self.size = size
         self.max_context_len = max_context_len
         self.req_to_token = torch.zeros(
-            (size, max_context_len), dtype=torch.long, device=device
+            (size, max_context_len), dtype=torch.int32, device=device
         )
         self.free_slots = list(range(size))
         self._lock = threading.Lock()

astrai/inference/cache/pool.py

@@ -115,6 +115,8 @@ class PagePool:
 
         self.contiguous = n_tokens is None
         self.n_tokens = max_batch_size * max_seq_len if self.contiguous else n_tokens
+        if self.n_tokens > torch.iinfo(torch.int32).max:
+            raise ValueError("KV cache token count exceeds the int32 slot index limit")
 
         self._storage = KVStorage(
             self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
@@ -124,7 +126,10 @@ class PagePool:
         if self.contiguous:
             for i in range(max_batch_size):
                 self._req_pool.req_to_token[i] = torch.arange(
-                    i * max_seq_len, (i + 1) * max_seq_len, device=device
+                    i * max_seq_len,
+                    (i + 1) * max_seq_len,
+                    dtype=torch.int32,
+                    device=device,
                 )
             self._strategy: AllocationStrategy = ContiguousStrategy()
         else:
@@ -184,7 +189,7 @@ class PagePool:
             kvp_buf[: b + 1] += inc_buf[: b + 1]
         else:
             rpi_buf[:b].copy_(
-                torch.tensor(req_indices, dtype=torch.long, device=device)
+                torch.tensor(req_indices, dtype=torch.int32, device=device)
             )
             sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
             kvp_buf[: b + 1].zero_()

astrai/inference/workspace.py

@@ -74,7 +74,7 @@ class InferenceWorkspace:
         # when the Executor passes this workspace).  Stable addresses make the
         # decode forward CUDA-graph capturable.
         self.req_pool_indices = torch.empty(
-            (max_batch_size,), dtype=torch.long, device=device
+            (max_batch_size,), dtype=torch.int32, device=device
         )
         self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
         self.kv_indptr = torch.empty(
@@ -85,7 +85,7 @@ class InferenceWorkspace:
         )
         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.long, device=device
+            (max_batch_size, 1), dtype=torch.int32, device=device
         )
 
         # Per-step position IDs (must be at a fixed address for CUDA-graph capture).

csrc/kernels/attn_common.h

@@ -55,8 +55,8 @@ struct AttentionParams {
     int mask_l_stride;
 
     // Paged K/V addressing
-    const int64_t* __restrict__ req_to_token;     // [num_reqs, max_context_len]
+    const int* __restrict__ req_to_token;         // [num_reqs, max_context_len]
-    const int64_t* __restrict__ req_pool_indices; // [batch]
+    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
     int max_context_len; // req_to_token stride (dim 1)

csrc/kernels/attn_entry_utils.cuh

@@ -160,8 +160,9 @@ inline void attn_pack_paged_decode_params(
     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_to_token.dtype() == torch::kInt32, "req_to_token must be int32");
-    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kInt32,
+                "req_pool_indices must be int32");
     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]");
@@ -184,8 +185,8 @@ inline void attn_pack_paged_decode_params(
     p.k_ptr = (const T*)k_cache.data_ptr();
     p.v_ptr = (const T*)v_cache.data_ptr();
     p.q_ptr = (const T*)q.data_ptr();
-    p.req_to_token = req_to_token.data_ptr<int64_t>();
+    p.req_to_token = req_to_token.data_ptr<int>();
-    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    p.req_pool_indices = req_pool_indices.data_ptr<int>();
     p.kv_indptr = kv_indptr.data_ptr<int>();
     p.qo_indptr = nullptr;
     p.max_context_len = (int)req_to_token.size(1);
@@ -239,8 +240,9 @@ inline void attn_pack_paged_prefill_params(
     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_to_token.dtype() == torch::kInt32, "req_to_token must be int32");
-    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kInt32,
+                "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(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
@@ -267,8 +269,8 @@ inline void attn_pack_paged_prefill_params(
     p.k_ptr = (const T*)k_cache.data_ptr();
     p.v_ptr = (const T*)v_cache.data_ptr();
     p.q_ptr = (const T*)q.data_ptr();
-    p.req_to_token = req_to_token.data_ptr<int64_t>();
+    p.req_to_token = req_to_token.data_ptr<int>();
-    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    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.max_context_len = (int)req_to_token.size(1);

csrc/kernels/attn_kv_source.cuh

@@ -33,7 +33,7 @@ using bf16 = __nv_bfloat16;
 // Hoisted per-(batch, kv_head) addressing context.
 struct KVContext {
     int kv_base;          // contig: batch*kv_b_stride + kv_head*kv_h_stride
-    int64_t req_idx;      // paged: req_pool_indices[batch]
+    int req_idx;          // paged: req_pool_indices[batch]
     int64_t rtt_stride;   // paged: max_context_len
     int64_t pool_stride;  // paged: kv_head * HEAD_DIM
     int64_t head_off;     // paged: kv_head * HEAD_DIM
@@ -177,9 +177,9 @@ struct PagedKV {
     }
     HOST_DEV_FORCEINLINE KVAddr kv_addr(
         const AttentionParams<bf16>& p, const KVContext& c, int kc, int d, bool valid) {
-        const int64_t slot = valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : 0;
+        const int slot = valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : 0;
         const bool ok = valid && (slot >= 0);
-        const int64_t gmem_off = slot * c.pool_stride + c.head_off + d;
+        const int64_t gmem_off = (int64_t)slot * c.pool_stride + c.head_off + d;
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], ok};
     }
 };

csrc/tests/attn_paged_test.cu

@@ -18,7 +18,7 @@ struct PagedPrefillDispatch { AttentionParams<bf16>& p; template<int H> void ope
 // 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* req_to_token, const int* 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)
@@ -27,7 +27,7 @@ static void cpu_paged_decode_ref(
     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];
+        int req_idx = req_pool_indices[b];
         #pragma omp parallel for schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
             int kv_h = h / n_rep;
@@ -35,7 +35,7 @@ static void cpu_paged_decode_ref(
             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];
+                int 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] *
@@ -66,7 +66,7 @@ static void cpu_paged_decode_ref(
 //   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* req_to_token, const int* req_pool_indices,
     const int* kv_indptr, const int* qo_indptr,
     const bool* mask, int mask_l_stride, int mask_kv_stride,
     int B, int Hq, int Hkv, int D, int max_ctx_len, int causal,
@@ -78,7 +78,7 @@ static void cpu_paged_prefill_ref(
         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];
+        int req_idx = req_pool_indices[b];
         #pragma omp parallel for collapse(2) schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
             for (int qi = 0; qi < q_len; qi++) {
@@ -89,7 +89,7 @@ static void cpu_paged_prefill_ref(
                 for (int kj = 0; kj < lim; kj++) {
                     if (mask && !mask[b * mask_l_stride * mask_kv_stride
                                      + qi * mask_kv_stride + kj]) continue;
-                    int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
+                    int 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] *
@@ -149,14 +149,14 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_rtt, *d_rpi;
     int *d_kvi;
     float *d_op, *d_ml;
     cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
@@ -181,7 +181,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     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* h_rtt = (int*)malloc(sz_rtt);
     int next_slot = 0;
     for (int r = 0; r < num_reqs; r++)
         for (int p = 0; p < max_ctx; p++) {
@@ -191,7 +191,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     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);
+    int* h_rpi = (int*)malloc(sz_rpi);
     for (int b = 0; b < B; b++) h_rpi[b] = b;
     cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
 
@@ -278,15 +278,15 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_rtt, *d_rpi;
     int *d_kvi;
     bool *d_mask;
     float *d_op, *d_ml;
@@ -312,7 +312,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
     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* h_rtt = (int*)malloc(sz_rtt);
     int next_slot = 0;
     for (int r = 0; r < num_reqs; r++)
         for (int p = 0; p < max_ctx; p++) {
@@ -321,7 +321,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
         }
     cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
 
-    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    int* h_rpi = (int*)malloc(sz_rpi);
     for (int b = 0; b < B; b++) h_rpi[b] = b;
     cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
 
@@ -417,13 +417,13 @@ static int run_prefill_test(int B, int Hq, int Hkv,
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_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);
@@ -446,7 +446,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     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* h_rtt = (int*)malloc(sz_rtt);
     int next_slot = 0;
     for (int r = 0; r < num_reqs; r++)
         for (int p = 0; p < max_ctx; p++) {
@@ -455,7 +455,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
         }
     cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
 
-    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    int* h_rpi = (int*)malloc(sz_rpi);
     for (int b = 0; b < B; b++) h_rpi[b] = b;
     cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
 
@@ -546,14 +546,14 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_rtt, *d_rpi;
     int *d_kvi, *d_qoi;
     bool *d_mask;
     cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
@@ -577,7 +577,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     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* h_rtt = (int*)malloc(sz_rtt);
     int next_slot = 0;
     for (int r = 0; r < num_reqs; r++)
         for (int p = 0; p < max_ctx; p++) {
@@ -586,7 +586,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
         }
     cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
 
-    int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
+    int* h_rpi = (int*)malloc(sz_rpi);
     h_rpi[0] = 0;
     cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
 
@@ -673,14 +673,14 @@ static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_rtt, *d_rpi;
     int *d_kvi;
     float *d_op, *d_ml;
     cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
@@ -696,12 +696,12 @@ static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
     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);
+    int* h_rtt = (int*)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);
+    int* h_rpi = (int*)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);
@@ -749,13 +749,13 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
 
     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_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
-    size_t sz_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rpi = (size_t)B * sizeof(int);
     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_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);
@@ -769,12 +769,12 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
     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);
+    int* h_rtt = (int*)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);
+    int* h_rpi = (int*)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);

tests/extension/test_backend_equivalence.py

@@ -8,8 +8,11 @@ import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
 from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.inference.workspace import InferenceWorkspace
 from tests.extension.conftest import D, skip_no_kernel
+from tests.helpers import FakeTokenizer
 
 
 def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
@@ -176,12 +179,69 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
 
 
+@skip_no_kernel
+def test_decode_cuda_graph_replay_is_exact(cuda_model):
+    """INT32 cache indices must remain graph-capturable and replay exactly."""
+    model, _ = cuda_model
+    device = "cuda"
+    prompt_ids = [1, 2, 3, 4, 5, 6, 7, 8]
+    cache = PagePool(
+        n_layers=2,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=1,
+        max_seq_len=64,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    task_cache = _mk_task_cache(cache)
+    ws = _ws(cache)
+    task_cache.task_alloc("t1", prompt_ids)
+
+    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+    position_ids = torch.arange(len(prompt_ids), device=device).unsqueeze(0)
+    input_mask = torch.ones(1, len(prompt_ids), dtype=torch.bool, device=device)
+
+    with attn_backend(ATTN_BACKEND.CUDA), torch.inference_mode():
+        model(
+            input_ids,
+            input_mask=input_mask,
+            position_ids=position_ids,
+            kv_cache=task_cache.bind(["t1"], ws, start_pos=0),
+        )
+
+        task_cache.task_extend("t1", len(prompt_ids))
+        kv_cache = task_cache.bind(["t1"], ws)
+        assert kv_cache.req_to_token.dtype == torch.int32
+        assert kv_cache.req_pool_indices.dtype == torch.int32
+        assert kv_cache.out_cache_loc.dtype == torch.int32
+
+        decode_args = {
+            "input_ids": torch.tensor([[9]], dtype=torch.long, device=device),
+            "input_mask": torch.ones(1, 1, 64, dtype=torch.bool, device=device),
+            "position_ids": torch.tensor([[len(prompt_ids)]], device=device),
+            "kv_cache": kv_cache,
+        }
+        graph = CudaGraphContext(enabled=True)
+        graph.forward(model, key=(1,), **decode_args)
+        graph.forward(model, key=(1,), **decode_args)
+        first = graph.forward(model, key=(1,), **decode_args)["logits"].clone()
+        slot = kv_cache.out_cache_loc[0, 0]
+        first_k = kv_cache.k_buffer[:, slot].clone()
+        first_v = kv_cache.v_buffer[:, slot].clone()
+
+        second = graph.forward(model, key=(1,), **decode_args)["logits"].clone()
+        torch.cuda.synchronize()
+
+    assert graph.has_graph((1,))
+    torch.testing.assert_close(second, first, rtol=0, atol=0)
+    torch.testing.assert_close(kv_cache.k_buffer[:, slot], first_k, rtol=0, atol=0)
+    torch.testing.assert_close(kv_cache.v_buffer[:, slot], first_v, rtol=0, atol=0)
+
+
 @skip_no_kernel
 def test_run_batch_cuda_matches_torch_greedy(cuda_model):
     """Greedy decode (temperature=0) should produce identical tokens."""
-    from astrai.inference.scheduler import InferenceScheduler
-    from tests.helpers import FakeTokenizer
-
     model, _ = cuda_model
     tokenizer = FakeTokenizer()
 

tests/inference/test_cache.py

@@ -176,6 +176,7 @@ def test_page_pool_task_cacheable_ids_excludes_unmaterialized_tail():
 
 def test_req_to_token_pool_alloc_free():
     pool = ReqToTokenPool(4, 128, torch.device("cpu"))
+    assert pool.req_to_token.dtype == torch.int32
     slots = pool.alloc(2)
     assert len(slots) == 2
     assert len(pool.free_slots) == 2
@@ -279,8 +280,10 @@ def test_page_pool_contiguous_bind_tasks_prefill():
     task_cache.task_alloc("t2", list(range(10)))
     kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0)
     assert kv.out_cache_loc.shape == (2, 10)
+    assert kv.out_cache_loc.dtype == torch.int32
     assert kv.seq_lens.tolist() == [10, 10]
     assert kv.req_pool_indices.shape == (2,)
+    assert kv.req_pool_indices.dtype == torch.int32
 
 
 def test_page_pool_contiguous_bind_tasks_decode():