model: paged KV cache with CPU swap pool, decode graph, qwen3 updates

- paged_kv_cache: new block-paged KV cache; adds a pinned-host swap pool with
  a second BlockAllocator, per-sequence Location {Gpu,Cpu}, and lossless
  swap_out/swap_in (block-granular D2H/H2D) for vLLM-style preemption.
  bytes_per_block helper exposes per-block cost for VRAM-based sizing.
- decode_graph: CUDA-graph decode path.
- qwen3/gpt2/kv_cache: paged prefill/decode forward + related updates.
- tokenizer/bins: BPE updates, new xserv-chat CLI, bench-qwen3 tweaks.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/xserv-model/src/qwen3.rs

@@ -6,6 +6,7 @@ use xserv_tensor::{DType, Device, Tensor};
 use crate::config::ModelConfig;
 use crate::gpt2::KVCache;
 use crate::kv_cache::GpuKVCache;
+use crate::paged_kv_cache::PagedKVCache;
 
 pub struct Qwen3 {
     pub config: ModelConfig,
@@ -255,6 +256,196 @@ impl Qwen3 {
         matmul_2d(&x, &self.lm_head_t) // [B, vocab_size]
     }
 
+    /// Paged decode: process one token per sequence using a shared paged KV cache.
+    ///
+    /// tokens:    [B] one token per sequence
+    /// positions: [B] current logical position (BEFORE this step) per sequence
+    /// seq_slots: [B] slot ids in `paged_cache`
+    pub fn forward_decode_paged(
+        &self,
+        tokens: &[u32],
+        positions: &[usize],
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let batch = tokens.len();
+        assert_eq!(positions.len(), batch);
+        assert_eq!(seq_slots.len(), batch);
+        assert!(batch > 0);
+
+        let num_heads = self.config.num_heads();
+        let num_kv_heads = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        // Ensure all slots have enough physical blocks for this token, then
+        // upload block tables + context_lens once for the whole forward (the
+        // tables are identical across layers; only the layer's K/V pool changes).
+        let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
+        for (b, &slot) in seq_slots.iter().enumerate() {
+            paged_cache.ensure_capacity(slot, positions[b] + 1);
+        }
+        paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens);
+
+        let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
+        let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
+        let max_blocks = paged_cache.max_blocks_per_seq();
+
+        // Batched embedding: [B, hidden]
+        let mut x = embedding(&self.embed_tokens, tokens);
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let q_all = matmul_2d(&normed, &layer.q_proj_wt);
+            let k_all = matmul_2d(&normed, &layer.k_proj_wt);
+            let v_all = matmul_2d(&normed, &layer.v_proj_wt);
+
+            let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
+            for b in 0..batch {
+                let q_row = row_view(&q_all, b);
+                let k_row = row_view(&k_all, b);
+                let v_row = row_view(&v_all, b);
+
+                let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim);
+                let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim);
+                let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim);
+
+                let q = head_rmsnorm(&q, &layer.q_norm, eps);
+                let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+                let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                let pos = [positions[b] as u32];
+                rope_inplace(&q, &self.rope_cache, &pos);
+                rope_inplace(&k, &self.rope_cache, &pos);
+
+                let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim);
+                let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim);
+
+                paged_cache.append_tokens(seq_slots[b], layer_idx, &k, &v, 1, positions[b]);
+
+                let q_flat = xserv_kernels::merge_heads_gpu(&q, 1, num_heads, head_dim);
+                q_rows.push(q_flat);
+            }
+
+            let q_batched_2d = concat_rows(&q_rows);
+            // q_batched_2d: [B, num_heads * head_dim]. Memory is [B, H, D] —
+            // a plain reshape view to [B, H, 1, D] is what the paged kernel expects.
+            let q_4d = q_batched_2d.reshape(&[batch, num_heads, 1, head_dim]);
+
+            let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
+            let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
+
+            let attn_out = xserv_kernels::paged_decode_attention(
+                &q_4d,
+                k_pool_ptr,
+                v_pool_ptr,
+                bt_ptr,
+                cl_ptr,
+                batch,
+                num_heads,
+                num_kv_heads,
+                head_dim,
+                max_blocks,
+            );
+
+            // attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
+            // Plain reshape is a view; merge_heads_gpu would incorrectly swap B<->H.
+            let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        // Advance logical seq_len now that all layers have been written.
+        for &slot in seq_slots {
+            paged_cache.advance_seq_len(slot, 1);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
+    /// Paged prefill: write a sequence of `new_tokens` K/V into the paged
+    /// cache for `slot`, run flash attention via gathered contiguous K/V.
+    /// Returns logits [new_tokens, vocab_size].
+    pub fn forward_prefill_paged(
+        &self,
+        token_ids: &[u32],
+        slot: usize,
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let new_tokens = token_ids.len();
+        let pos_offset = paged_cache.seq_len(slot);
+        let num_heads = self.config.num_heads();
+        let num_kv_heads = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
+
+        // Pre-allocate enough blocks and bump seq_len up-front so per-layer
+        // gather_kv_contiguous returns the freshly written K/V range.
+        paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
+        paged_cache.advance_seq_len(slot, new_tokens);
+
+        let mut x = embedding(&self.embed_tokens, token_ids);
+        let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            let q = matmul_2d(&normed, &layer.q_proj_wt);
+            let k = matmul_2d(&normed, &layer.k_proj_wt);
+            let v = matmul_2d(&normed, &layer.v_proj_wt);
+
+            let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
+            let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
+
+            let q = head_rmsnorm(&q, &layer.q_norm, eps);
+            let k = head_rmsnorm(&k, &layer.k_norm, eps);
+
+            let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+            rope_inplace(&q, &self.rope_cache, &positions);
+            rope_inplace(&k, &self.rope_cache, &positions);
+            let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
+            let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
+
+            // Write into paged pool at the original (pre-advance) position.
+            paged_cache.append_tokens(slot, layer_idx, &k, &v, new_tokens, pos_offset);
+
+            // Gather contiguous K/V for the full sequence (seq_len already includes new_tokens).
+            let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
+            let attn_out = flash_attention(&q, &k_full, &v_full, true);
+
+            let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
+            let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
+
+            let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
+            let residual = x_new.clone();
+
+            let gate = matmul_2d(&normed, &layer.gate_proj_wt);
+            let up = matmul_2d(&normed, &layer.up_proj_wt);
+            let hidden_states = xserv_kernels::silu_mul(&gate, &up);
+            let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            x = add_any(&residual, &down);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul_2d(&x, &self.lm_head_t)
+    }
+
     /// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
     pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
         let new_tokens = token_ids.len();
@@ -320,6 +511,40 @@ impl Qwen3 {
         let x = rmsnorm(&x, &self.norm, eps);
         matmul_2d(&x, &self.lm_head_t)
     }
+
+    /// Extract weight pointers for CUDA Graph capture.
+    pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
+        self.layers.iter().map(|l| crate::decode_graph::LayerWeightPtrs {
+            input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
+            q_proj_wt: l.q_proj_wt.data_ptr() as *const std::ffi::c_void,
+            k_proj_wt: l.k_proj_wt.data_ptr() as *const std::ffi::c_void,
+            v_proj_wt: l.v_proj_wt.data_ptr() as *const std::ffi::c_void,
+            o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
+            q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
+            k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
+            post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
+            gate_proj_wt: l.gate_proj_wt.data_ptr() as *const std::ffi::c_void,
+            up_proj_wt: l.up_proj_wt.data_ptr() as *const std::ffi::c_void,
+            down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
+        }).collect()
+    }
+
+    /// Get pointers needed for CUDA Graph capture.
+    pub fn graph_capture_ptrs(&self) -> (
+        *const std::ffi::c_void,  // norm weight
+        *const std::ffi::c_void,  // lm_head_t
+        *const std::ffi::c_void,  // embed_tokens
+        *const std::ffi::c_void,  // rope cos
+        *const std::ffi::c_void,  // rope sin
+    ) {
+        (
+            self.norm.data_ptr() as *const std::ffi::c_void,
+            self.lm_head_t.data_ptr() as *const std::ffi::c_void,
+            self.embed_tokens.data_ptr() as *const std::ffi::c_void,
+            self.rope_cache.cos.as_ptr() as *const std::ffi::c_void,
+            self.rope_cache.sin.as_ptr() as *const std::ffi::c_void,
+        )
+    }
 }
 
 // --- Helpers ---