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
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};
}
};