diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 8a11ca4..b712fa2 100644 --- a/astrai/extension/attention_backend.py +++ b/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) diff --git a/astrai/extension/attention_ops.py b/astrai/extension/attention_ops.py index adca269..05cc8d8 100644 --- a/astrai/extension/attention_ops.py +++ b/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 diff --git a/astrai/inference/cache/buffer.py b/astrai/inference/cache/buffer.py index 5f72ace..30a8d83 100644 --- a/astrai/inference/cache/buffer.py +++ b/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() diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index 99e7b6f..eebf4f4 100644 --- a/astrai/inference/cache/pool.py +++ b/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_() diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index 7fc08e7..ca2ea0d 100644 --- a/astrai/inference/workspace.py +++ b/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). diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index 1d2d78a..c29a4b3 100644 --- a/csrc/kernels/attn_common.h +++ b/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) diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 53d228a..0e56604 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/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(); - p.req_pool_indices = req_pool_indices.data_ptr(); + p.req_to_token = req_to_token.data_ptr(); + p.req_pool_indices = req_pool_indices.data_ptr(); p.kv_indptr = kv_indptr.data_ptr(); 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(); - p.req_pool_indices = req_pool_indices.data_ptr(); + p.req_to_token = req_to_token.data_ptr(); + p.req_pool_indices = req_pool_indices.data_ptr(); p.kv_indptr = kv_indptr.data_ptr(); p.qo_indptr = qo_indptr.data_ptr(); p.max_context_len = (int)req_to_token.size(1); diff --git a/csrc/kernels/attn_kv_source.cuh b/csrc/kernels/attn_kv_source.cuh index 748352e..8501c2d 100644 --- a/csrc/kernels/attn_kv_source.cuh +++ b/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& 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}; } }; diff --git a/csrc/tests/attn_paged_test.cu b/csrc/tests/attn_paged_test.cu index d6438f1..ad05e7f 100644 --- a/csrc/tests/attn_paged_test.cu +++ b/csrc/tests/attn_paged_test.cu @@ -18,7 +18,7 @@ struct PagedPrefillDispatch { AttentionParams& p; template 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); diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index b691d7b..a775c05 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/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() diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index e8d9cd7..0d4a559 100644 --- a/tests/inference/test_cache.py +++ b/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():