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) <noreply@anthropic.com>

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();
+        // Stage block_ids on the GPU. Pool-allocated so this is essentially
-        let v_src = v_new.storage().gpu_buffer();
+        // 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 k_src = k_new.data_ptr() as *const std::ffi::c_void;
-        let v_pool = &mut self.v_pools[layer];
+        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;
+        unsafe {
-        while t < num_tokens {
+            xserv_kernels::reshape_and_cache_bf16(
-            let p = start_pos + t;
+                k_src, v_src,
-            let logical_blk = p / bs;
+                k_pool_ptr, v_pool_ptr,
-            let slot_in_blk = p % bs;
+                block_ids_gpu.as_ptr() as *const i32,
-            let chunk = (bs - slot_in_blk).min(num_tokens - t);
+                num_tokens, nkv, hd, start_pos, bs,
-            let phys = block_ids[logical_blk] as usize;
+                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 {
+    /// Batched append for the multi-sequence decode step: writes one new
-                let src_off = (h * num_tokens + t) * hd * es;
+    /// K/V token per active sequence into `layer`'s pool, using
-                let dst_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
+    /// `block_table_gpu` and `context_lens_gpu` directly. Caller must have
-                let count = chunk * hd * es;
+    /// just run `sync_active_batch_with_lens(slots, kv_lens)` so that:
-                k_pool.copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+    ///   - row `i` of block_table_gpu holds the block ids for `slots[i]`
-                v_pool.copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
+    ///   - 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(),
+            );
         }
     }