ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: use int32 paged KV indices

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-15 13:17:06 +08:00
parent b5afe3d7a4
commit 3fb4b8ab13
11 changed files with 135 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+4 -8
View File
@@ -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
+4 -4
View File
@@ -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
Vendored
+1 -1
View File
@@ -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
Vendored
+7 -2
View File
@@ -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
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+10 -8
View File
@@ -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
+3 -3
View File
@@ -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
+36 -36
View File
@@ -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);
tests/extension/test_backend_equivalence.py
+63 -3
View File
@@ -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
+3
View File
@@ -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():