From cc4bd4cfe51815985ebe1c7422208dc5fa424964 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sat, 30 May 2026 12:50:28 +0800 Subject: [PATCH] paged-kv: kernel-based scatter + fix data_ptr offset bug 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) --- crates/xserv-model/src/paged_kv_cache.rs | 95 ++++++++++++++++++------ 1 file changed, 73 insertions(+), 22 deletions(-) diff --git a/crates/xserv-model/src/paged_kv_cache.rs b/crates/xserv-model/src/paged_kv_cache.rs index f125483..5871abc 100644 --- a/crates/xserv-model/src/paged_kv_cache.rs +++ b/crates/xserv-model/src/paged_kv_cache.rs @@ -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 = 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::(); + 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(), + ); } }