gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paged-kv: kernel-based scatter + fix data_ptr offset bug

Browse Source 
Replace the Rust cudaMemcpy loop in append_tokens() with the new
reshape_and_cache kernel. Add append_tokens_batched() for the decode
path using the batched variant.

Fix: use data_ptr() instead of storage().gpu_buffer().as_ptr() so that
tensor offset is respected. The old code silently read from storage base
(element 0) instead of the tensor's logical start, which produced wrong
results when K/V tensors were narrow() views into a fused QKV buffer.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-30 12:50:28 +08:00
parent 13ae3de69e
commit cc4bd4cfe5
1 changed files with 73 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
crates/xserv-model/src/paged_kv_cache.rs
View File

@@ -305,6 +305,10 @@ impl PagedKVCache {
/// `k_new`, `v_new`: GPU tensors with logical shape
/// [1, num_kv_heads, num_tokens, head_dim]
/// stored contiguously (head-major, then tokens, then dim).
///
/// Implementation: a single `reshape_and_cache` kernel per call. The
/// previous Rust loop fired `num_tokens * num_kv_heads` cudaMemcpys per
/// layer (≈290k for a 1024-token Qwen3 prefill across 36 layers).
pub fn append_tokens(
&mut self,
slot: usize,
@@ -318,36 +322,83 @@ impl PagedKVCache {
// Make sure blocks exist for the target range.
self.ensure_capacity(slot, start_pos + num_tokens);
let block_ids = self.seq_states[slot].as_ref().unwrap().block_ids.clone();
let nkv = self.num_kv_heads;
let hd = self.head_dim;
let es = self.elem_size;
let bs = BLOCK_SIZE;
let k_src = k_new.storage().gpu_buffer();
let v_src = v_new.storage().gpu_buffer();
// Stage block_ids on the GPU. Pool-allocated so this is essentially
// free after the first call (same bucket every step).
let block_ids: Vec<i32> = self.seq_states[slot].as_ref().unwrap()
.block_ids.iter().map(|&b| b as i32).collect();
let bytes = block_ids.len() * std::mem::size_of::<i32>();
let mut block_ids_gpu = xserv_cuda::allocator::cached_alloc(bytes)
.expect("alloc append block_ids");
let block_ids_bytes = unsafe {
std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes)
};
block_ids_gpu.copy_from_host(block_ids_bytes).expect("upload block_ids");
let k_pool = &mut self.k_pools[layer];
let v_pool = &mut self.v_pools[layer];
let k_src = k_new.data_ptr() as *const std::ffi::c_void;
let v_src = v_new.data_ptr() as *const std::ffi::c_void;
let k_pool_ptr = self.k_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
let v_pool_ptr = self.v_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
let mut t = 0usize;
while t < num_tokens {
let p = start_pos + t;
let logical_blk = p / bs;
let slot_in_blk = p % bs;
let chunk = (bs - slot_in_blk).min(num_tokens - t);
let phys = block_ids[logical_blk] as usize;
unsafe {
xserv_kernels::reshape_and_cache_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
block_ids_gpu.as_ptr() as *const i32,
num_tokens, nkv, hd, start_pos, bs,
std::ptr::null_mut(),
);
}
// block_ids_gpu drops here; the launch on the null stream will have
// finished consuming it before any subsequent op alloc()s the same
// bucket (null stream is sequential).
}
for h in 0..nkv {
let src_off = (h * num_tokens + t) * hd * es;
let dst_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
let count = chunk * hd * es;
k_pool.copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
v_pool.copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
}
/// Batched append for the multi-sequence decode step: writes one new
/// K/V token per active sequence into `layer`'s pool, using
/// `block_table_gpu` and `context_lens_gpu` directly. Caller must have
/// just run `sync_active_batch_with_lens(slots, kv_lens)` so that:
/// - row `i` of block_table_gpu holds the block ids for `slots[i]`
/// - context_lens_gpu[i] == seq_len(slots[i]) + 1 (the kv_len **after**
/// this step — i.e., the new token will be written at index kv_len-1)
///
/// `k_new`, `v_new`: GPU tensors, contiguous, BF16, shape
/// `[batch, num_kv_heads, head_dim]`.
///
/// Like `append_tokens`, this does **not** touch `seq_len`. Call
/// `advance_seq_len(slot, 1)` for each slot after every layer has been
/// written.
pub fn append_tokens_batched(
&mut self,
layer: usize,
k_new: &Tensor,
v_new: &Tensor,
batch: usize,
) {
if batch == 0 { return; }
let nkv = self.num_kv_heads;
let hd = self.head_dim;
debug_assert_eq!(k_new.shape(), &[batch, nkv, hd]);
debug_assert_eq!(v_new.shape(), &[batch, nkv, hd]);
t += chunk;
let k_src = k_new.data_ptr() as *const std::ffi::c_void;
let v_src = v_new.data_ptr() as *const std::ffi::c_void;
let k_pool_ptr = self.k_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
let v_pool_ptr = self.v_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
let bt_ptr = self.block_table_gpu.as_ptr() as *const i32;
let cl_ptr = self.context_lens_gpu.as_ptr() as *const i32;
unsafe {
xserv_kernels::reshape_and_cache_batched_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
bt_ptr, cl_ptr,
batch, nkv, hd, BLOCK_SIZE, self.max_blocks_per_seq,
std::ptr::null_mut(),
);
}
}