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

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 int64_t* __restrict__ req_pool_indices; // [batch]
+    const int* __restrict__ req_to_token;         // [num_reqs, max_context_len]
+    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_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(req_to_token.dtype() == torch::kInt32, "req_to_token must be int32");
+    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_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    p.req_to_token = req_to_token.data_ptr<int>();
+    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_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
+    TORCH_CHECK(req_to_token.dtype() == torch::kInt32, "req_to_token must be int32");
+    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_pool_indices = req_pool_indices.data_ptr<int64_t>();
+    p.req_to_token = req_to_token.data_ptr<int>();
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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_rpi = (size_t)B * sizeof(int64_t);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    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);