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.v_buffer[layer_id].index_copy_(0, loc, v[:, 0])
+        kv_cache.k_buffer[layer_id, loc] = k[:, 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_(
-            0, loc, k.reshape(-1, k.size(2), k.size(3))
-        )
-        kv_cache.v_buffer[layer_id].index_copy_(
-            0, loc, v.reshape(-1, v.size(2), v.size(3))
-        )
+        kv_cache.k_buffer[layer_id, 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))
 
         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_pool_indices: [batch] (int64) — rows into req_to_token
+        req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
+        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_pool_indices: [batch] (int64)
+        req_to_token: [num_reqs, max_context_len] (int32)
+        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).