xserv/crates/xserv-model/src/paged_kv_cache.rs

//! Paged KV cache: vLLM-style block-based KV cache with O(1) allocation
//! and indirection via per-sequence block tables.
//!
//! Physical layout per layer:
//!   K pool: [total_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16
//!   V pool: same
//!
//! Logical view per sequence: a list of physical block ids. Token at logical
//! position p lives in block_ids[p / BLOCK_SIZE] at slot (p % BLOCK_SIZE).

use crate::config::ModelConfig;
use xserv_cuda::{GpuBuffer, PinnedBuffer};
use xserv_tensor::{DType, Tensor};

pub const BLOCK_SIZE: usize = 16;

/// Stack-based block allocator: O(1) alloc/free.
pub struct BlockAllocator {
    free_stack: Vec<u32>,
    total: usize,
}

impl BlockAllocator {
    pub fn new(total_blocks: usize) -> Self {
        // Reserve block 0 as a sentinel "null" block (never allocated).
        // Free list contains [total-1, total-2, ..., 1] so pop returns 1 first.
        // total_blocks==0 means "disabled" (e.g. swap off): empty free list.
        let mut free_stack = Vec::with_capacity(total_blocks.saturating_sub(1));
        for b in (1..total_blocks).rev() {
            free_stack.push(b as u32);
        }
        Self {
            free_stack,
            total: total_blocks,
        }
    }

    pub fn alloc(&mut self) -> Option<u32> {
        self.free_stack.pop()
    }

    pub fn free(&mut self, block: u32) {
        debug_assert!((block as usize) < self.total && block != 0);
        self.free_stack.push(block);
    }

    pub fn free_count(&self) -> usize {
        self.free_stack.len()
    }

    pub fn total(&self) -> usize {
        self.total
    }

    pub fn can_alloc(&self, n: usize) -> bool {
        self.free_stack.len() >= n
    }
}

/// Where a sequence's KV blocks currently live.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Location {
    Gpu,
    Cpu,
}

/// Per-sequence state held in the cache.
#[derive(Clone)]
pub struct SeqState {
    /// Block ids into the GPU pool when `location == Gpu`, or into the CPU
    /// (pinned host) pool when `location == Cpu`.
    pub block_ids: Vec<u32>,
    pub seq_len: usize,
    pub location: Location,
}

pub struct PagedKVCache {
    // [layer]: GpuBuffer of size total_blocks * nkv * BLOCK_SIZE * hd * elem_size
    k_pools: Vec<GpuBuffer>,
    v_pools: Vec<GpuBuffer>,

    // CPU (pinned host) swap pools, same per-layer layout as the GPU pools but
    // sized for `cpu_total_blocks`. Empty when swap is disabled.
    cpu_k_pools: Vec<PinnedBuffer>,
    cpu_v_pools: Vec<PinnedBuffer>,
    cpu_allocator: BlockAllocator,

    // Bytes occupied by one block within a single layer pool:
    //   num_kv_heads * BLOCK_SIZE * head_dim * elem_size.
    block_bytes: usize,

    allocator: BlockAllocator,
    seq_states: Vec<Option<SeqState>>,

    // GPU-resident per-sequence metadata. Uploaded each step via sync_to_gpu().
    // block_table_gpu: i32 [max_seqs, max_blocks_per_seq]
    // context_lens_gpu: i32 [max_seqs]
    block_table_gpu: GpuBuffer,
    context_lens_gpu: GpuBuffer,
    // Host-side staging mirroring the GPU buffers above.
    block_table_host: Vec<i32>,
    context_lens_host: Vec<i32>,

    // Config
    num_layers: usize,
    num_kv_heads: usize,
    head_dim: usize,
    elem_size: usize,
    dtype: DType,
    device: u32,
    max_seqs: usize,
    max_blocks_per_seq: usize,
}

impl PagedKVCache {
    /// Bytes occupied by all KV blocks for ONE physical block across the whole
    /// model (both K and V, all layers). Use this to size pools against VRAM.
    pub fn bytes_per_block(config: &ModelConfig, dtype: DType) -> usize {
        2 * config.num_layers()
            * config.num_kv_heads()
            * BLOCK_SIZE
            * config.head_dim()
            * dtype.size_bytes()
    }

    /// Create a new paged cache.
    /// - `total_blocks`: total number of physical GPU blocks across all sequences.
    /// - `cpu_total_blocks`: physical blocks in the pinned-host swap pool (0 = swap off).
    /// - `max_seqs`: max number of concurrent sequences (slots), incl. swapped.
    /// - `max_blocks_per_seq`: capacity of the block table per slot
    ///   (must be >= ceil(max_seq_len / BLOCK_SIZE)).
    pub fn new(
        config: &ModelConfig,
        total_blocks: usize,
        cpu_total_blocks: usize,
        max_seqs: usize,
        max_blocks_per_seq: usize,
        dtype: DType,
        device: u32,
    ) -> Self {
        Self::new_tp(
            config,
            config.num_kv_heads(),
            total_blocks,
            cpu_total_blocks,
            max_seqs,
            max_blocks_per_seq,
            dtype,
            device,
        )
    }

    /// Like `new`, but with an explicit `num_kv_heads` — under tensor parallelism
    /// each rank only stores its `num_kv_heads / world` heads, so the pool is
    /// sized for the local head count, not the model's full count.
    #[allow(clippy::too_many_arguments)]
    pub fn new_tp(
        config: &ModelConfig,
        num_kv_heads: usize,
        total_blocks: usize,
        cpu_total_blocks: usize,
        max_seqs: usize,
        max_blocks_per_seq: usize,
        dtype: DType,
        device: u32,
    ) -> Self {
        assert!(
            total_blocks >= 2,
            "need at least 2 blocks (one is sentinel)"
        );
        let num_layers = config.num_layers();
        let head_dim = config.head_dim();
        let elem_size = dtype.size_bytes();
        let block_bytes = num_kv_heads * BLOCK_SIZE * head_dim * elem_size;
        let pool_bytes = total_blocks * block_bytes;

        let mut k_pools = Vec::with_capacity(num_layers);
        let mut v_pools = Vec::with_capacity(num_layers);
        for _ in 0..num_layers {
            let mut k = GpuBuffer::alloc(pool_bytes).expect("alloc paged K pool");
            let mut v = GpuBuffer::alloc(pool_bytes).expect("alloc paged V pool");
            k.zero().unwrap();
            v.zero().unwrap();
            k_pools.push(k);
            v_pools.push(v);
        }

        // Pinned-host swap pools (one per layer, mirroring the GPU layout).
        let mut cpu_k_pools = Vec::new();
        let mut cpu_v_pools = Vec::new();
        if cpu_total_blocks >= 2 {
            let cpu_pool_bytes = cpu_total_blocks * block_bytes;
            for _ in 0..num_layers {
                cpu_k_pools
                    .push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
                cpu_v_pools
                    .push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
            }
        }
        let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 {
            cpu_total_blocks
        } else {
            0
        });

        let block_table_gpu =
            GpuBuffer::alloc(max_seqs * max_blocks_per_seq * std::mem::size_of::<i32>())
                .expect("alloc block table");
        let context_lens_gpu =
            GpuBuffer::alloc(max_seqs * std::mem::size_of::<i32>()).expect("alloc context lens");

        let block_table_host = vec![0i32; max_seqs * max_blocks_per_seq];
        let context_lens_host = vec![0i32; max_seqs];

        let seq_states = (0..max_seqs).map(|_| None).collect();

        Self {
            k_pools,
            v_pools,
            cpu_k_pools,
            cpu_v_pools,
            cpu_allocator,
            block_bytes,
            allocator: BlockAllocator::new(total_blocks),
            seq_states,
            block_table_gpu,
            context_lens_gpu,
            block_table_host,
            context_lens_host,
            num_layers,
            num_kv_heads,
            head_dim,
            elem_size,
            dtype,
            device,
            max_seqs,
            max_blocks_per_seq,
        }
    }

    pub fn num_layers(&self) -> usize {
        self.num_layers
    }
    pub fn num_kv_heads(&self) -> usize {
        self.num_kv_heads
    }
    pub fn head_dim(&self) -> usize {
        self.head_dim
    }
    pub fn dtype(&self) -> DType {
        self.dtype
    }
    pub fn max_seqs(&self) -> usize {
        self.max_seqs
    }
    pub fn max_blocks_per_seq(&self) -> usize {
        self.max_blocks_per_seq
    }
    pub fn free_blocks(&self) -> usize {
        self.allocator.free_count()
    }
    pub fn total_blocks(&self) -> usize {
        self.allocator.total()
    }

    pub fn k_pool(&self, layer: usize) -> &GpuBuffer {
        &self.k_pools[layer]
    }
    pub fn v_pool(&self, layer: usize) -> &GpuBuffer {
        &self.v_pools[layer]
    }
    pub fn block_table_gpu(&self) -> &GpuBuffer {
        &self.block_table_gpu
    }
    pub fn context_lens_gpu(&self) -> &GpuBuffer {
        &self.context_lens_gpu
    }

    pub fn seq_len(&self, slot: usize) -> usize {
        self.seq_states[slot]
            .as_ref()
            .map(|s| s.seq_len)
            .unwrap_or(0)
    }

    pub fn is_slot_free(&self, slot: usize) -> bool {
        self.seq_states[slot].is_none()
    }

    /// Register a new sequence at `slot`. Allocates the first block.
    /// Returns Err(()) if no slot or no blocks are available.
    pub fn register_sequence(&mut self, slot: usize) -> Result<(), &'static str> {
        if slot >= self.max_seqs {
            return Err("slot out of range");
        }
        if self.seq_states[slot].is_some() {
            return Err("slot already in use");
        }
        let block = self.allocator.alloc().ok_or("out of blocks")?;
        self.seq_states[slot] = Some(SeqState {
            block_ids: vec![block],
            seq_len: 0,
            location: Location::Gpu,
        });
        Ok(())
    }

    /// Free all blocks for `slot` and clear the slot. Frees from whichever pool
    /// (GPU or CPU) the sequence currently lives in.
    pub fn free_sequence(&mut self, slot: usize) {
        if let Some(state) = self.seq_states[slot].take() {
            let alloc = match state.location {
                Location::Gpu => &mut self.allocator,
                Location::Cpu => &mut self.cpu_allocator,
            };
            for b in state.block_ids {
                alloc.free(b);
            }
        }
    }

    /// Number of blocks needed to hold `seq_len + new_tokens` tokens, beyond
    /// what is currently allocated for `slot`.
    pub fn additional_blocks_needed(&self, slot: usize, new_tokens: usize) -> usize {
        let state = self.seq_states[slot].as_ref().expect("unregistered slot");
        let cur = state.block_ids.len();
        let needed_total = (state.seq_len + new_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;
        if needed_total > cur {
            needed_total - cur
        } else {
            0
        }
    }

    /// Pre-allocate enough physical blocks in `slot` to cover positions
    /// `[0, end_pos)`. Call once before the per-layer append loop so that
    /// every layer's append uses the same block table.
    pub fn ensure_capacity(&mut self, slot: usize, end_pos: usize) {
        let state = self.seq_states[slot].as_mut().expect("unregistered slot");
        let needed_total = (end_pos + BLOCK_SIZE - 1) / BLOCK_SIZE;
        while state.block_ids.len() < needed_total {
            let b = self
                .allocator
                .alloc()
                .expect("out of blocks (caller must check)");
            assert!(
                state.block_ids.len() < self.max_blocks_per_seq,
                "block table overflow"
            );
            state.block_ids.push(b);
        }
    }

    /// Append `num_tokens` of K/V into the paged pool for `slot` at logical
    /// position `start_pos`. Caller must have called `ensure_capacity(slot, start_pos + num_tokens)`
    /// first (or accept that this method may also extend block list).
    /// Does NOT touch `seq_len`. Call `advance_seq_len(slot, num_tokens)` after
    /// every layer has been written.
    ///
    /// `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,
        layer: usize,
        k_new: &Tensor,
        v_new: &Tensor,
        num_tokens: usize,
        start_pos: usize,
    ) {
        if num_tokens == 0 {
            return;
        }
        // Make sure blocks exist for the target range.
        self.ensure_capacity(slot, start_pos + num_tokens);

        let nkv = self.num_kv_heads;
        let hd = self.head_dim;
        let bs = BLOCK_SIZE;

        // 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_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;

        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,
                xserv_cuda::current_stream_raw(),
            );
        }
        // 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).
    }

    /// 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]);

        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,
                xserv_cuda::current_stream_raw(),
            );
        }
    }

    /// Advance the logical seq_len after append_tokens for ALL layers has completed.
    pub fn advance_seq_len(&mut self, slot: usize, num_tokens: usize) {
        let state = self.seq_states[slot].as_mut().expect("unregistered slot");
        state.seq_len += num_tokens;
    }

    /// Roll a registered sequence back to `new_len` tokens.
    ///
    /// This only changes cache metadata and frees whole physical blocks that are
    /// no longer reachable. Bytes inside retained blocks are left untouched; the
    /// logical `seq_len` prevents attention from reading them, and later writes
    /// to the same positions overwrite them.
    pub fn truncate_sequence(&mut self, slot: usize, new_len: usize) -> Result<(), &'static str> {
        if slot >= self.max_seqs {
            return Err("truncate_sequence: slot out of range");
        }
        let state = self.seq_states[slot]
            .as_mut()
            .ok_or("truncate_sequence: empty slot")?;
        if new_len > state.seq_len {
            return Err("truncate_sequence: cannot extend");
        }

        let needed_blocks = ((new_len + BLOCK_SIZE - 1) / BLOCK_SIZE).max(1);
        while state.block_ids.len() > needed_blocks {
            let block = state.block_ids.pop().expect("checked len");
            match state.location {
                Location::Gpu => self.allocator.free(block),
                Location::Cpu => self.cpu_allocator.free(block),
            }
        }
        state.seq_len = new_len;
        Ok(())
    }

    /// Refresh the host-side block table + context lens from `seq_states`,
    /// then upload to GPU. Call once per decode step before the paged kernel.
    pub fn sync_to_gpu(&mut self) {
        let stride = self.max_blocks_per_seq;
        for slot in 0..self.max_seqs {
            let row = &mut self.block_table_host[slot * stride..(slot + 1) * stride];
            row.fill(0);
            let len = match &self.seq_states[slot] {
                Some(s) => {
                    for (i, b) in s.block_ids.iter().enumerate() {
                        row[i] = *b as i32;
                    }
                    s.seq_len as i32
                }
                None => 0,
            };
            self.context_lens_host[slot] = len;
        }

        self.upload_metadata();
    }

    /// Pack the given active slots into rows 0..slots.len() of block_table_gpu
    /// and context_lens_gpu, then upload. Used by paged decode where the kernel
    /// iterates over `batch` active sequences in order.
    pub fn sync_active_batch_to_gpu(&mut self, slots: &[usize]) {
        let lens: Vec<i32> = slots
            .iter()
            .map(|&s| self.seq_states[s].as_ref().unwrap().seq_len as i32)
            .collect();
        self.sync_active_batch_with_lens(slots, &lens);
    }

    /// Like sync_active_batch_to_gpu but uses caller-supplied kv_lens (number
    /// of valid K/V tokens to attend over per active row). Useful when the
    /// kv_len for the current step differs from the cached seq_len (e.g.
    /// before advance_seq_len has run).
    pub fn sync_active_batch_with_lens(&mut self, slots: &[usize], kv_lens: &[i32]) {
        assert_eq!(slots.len(), kv_lens.len());
        assert!(
            slots.len() <= self.max_seqs,
            "active batch exceeds max_seqs"
        );
        let stride = self.max_blocks_per_seq;
        for row in &mut self.block_table_host {
            *row = 0;
        }
        for cl in &mut self.context_lens_host {
            *cl = 0;
        }
        for (i, &slot) in slots.iter().enumerate() {
            let s = self.seq_states[slot]
                .as_ref()
                .expect("unregistered slot in active batch");
            let row = &mut self.block_table_host[i * stride..(i + 1) * stride];
            for (j, b) in s.block_ids.iter().enumerate() {
                row[j] = *b as i32;
            }
            self.context_lens_host[i] = kv_lens[i];
        }
        self.upload_metadata();
    }

    fn upload_metadata(&mut self) {
        let bt_bytes = unsafe {
            std::slice::from_raw_parts(
                self.block_table_host.as_ptr() as *const u8,
                self.block_table_host.len() * std::mem::size_of::<i32>(),
            )
        };
        self.block_table_gpu.copy_from_host(bt_bytes).unwrap();

        let cl_bytes = unsafe {
            std::slice::from_raw_parts(
                self.context_lens_host.as_ptr() as *const u8,
                self.context_lens_host.len() * std::mem::size_of::<i32>(),
            )
        };
        self.context_lens_gpu.copy_from_host(cl_bytes).unwrap();
    }

    /// Materialize a contiguous K/V tensor for a sequence at `layer`, shaped
    /// [1, num_kv_heads, seq_len, head_dim]. Used for prefill, where Flash
    /// Attention 2 expects contiguous K/V.
    ///
    /// Allocates from the cached allocator; the returned Tensors own their storage.
    pub fn gather_kv_contiguous(&self, slot: usize, layer: usize) -> (Tensor, Tensor) {
        let state = self.seq_states[slot].as_ref().expect("unregistered slot");
        let sl = state.seq_len;
        let nkv = self.num_kv_heads;
        let hd = self.head_dim;
        let es = self.elem_size;
        let bs = BLOCK_SIZE;

        let out_bytes = nkv * sl * hd * es;
        let mut k_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather K");
        let mut v_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather V");

        let k_pool = &self.k_pools[layer];
        let v_pool = &self.v_pools[layer];

        let mut p = 0usize;
        while p < sl {
            let logical_blk = p / bs;
            let slot_in_blk = p % bs;
            let chunk = (bs - slot_in_blk).min(sl - p);
            let phys = state.block_ids[logical_blk] as usize;

            for h in 0..nkv {
                let src_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
                let dst_off = (h * sl + p) * hd * es;
                let count = chunk * hd * es;
                k_dst
                    .copy_from_device_at(k_pool, src_off, dst_off, count)
                    .unwrap();
                v_dst
                    .copy_from_device_at(v_pool, src_off, dst_off, count)
                    .unwrap();
            }
            p += chunk;
        }

        let shape = &[1usize, nkv, sl, hd];
        let k = unsafe { tensor_from_owned_buf(k_dst, shape, self.dtype, self.device) };
        let v = unsafe { tensor_from_owned_buf(v_dst, shape, self.dtype, self.device) };
        (k, v)
    }

    // ----- Swapping (vLLM-style preemption to pinned host memory) -----

    pub fn free_cpu_blocks(&self) -> usize {
        self.cpu_allocator.free_count()
    }
    pub fn swap_enabled(&self) -> bool {
        !self.cpu_k_pools.is_empty()
    }

    pub fn is_swapped(&self, slot: usize) -> bool {
        matches!(
            self.seq_states[slot].as_ref().map(|s| s.location),
            Some(Location::Cpu)
        )
    }

    /// Number of physical blocks currently held by `slot` (in either pool).
    pub fn block_count(&self, slot: usize) -> usize {
        self.seq_states[slot]
            .as_ref()
            .map(|s| s.block_ids.len())
            .unwrap_or(0)
    }

    /// Whether a swapped sequence at `slot` can be brought back (enough free GPU blocks).
    pub fn can_swap_in(&self, slot: usize) -> bool {
        self.allocator.can_alloc(self.block_count(slot))
    }

    /// Whether the GPU sequence at `slot` can be evicted (enough free CPU blocks).
    pub fn can_swap_out(&self, slot: usize) -> bool {
        self.cpu_allocator.can_alloc(self.block_count(slot))
    }

    /// Evict `slot`'s KV from GPU to pinned host memory and free its GPU blocks.
    /// The slot stays registered (location = Cpu); the sequence is paused.
    pub fn swap_out(&mut self, slot: usize) -> Result<(), &'static str> {
        let state = self.seq_states[slot]
            .as_ref()
            .ok_or("swap_out: empty slot")?;
        if state.location == Location::Cpu {
            return Ok(());
        }
        let gpu_ids = state.block_ids.clone();
        let n = gpu_ids.len();
        if !self.cpu_allocator.can_alloc(n) {
            return Err("swap_out: CPU pool full");
        }

        let cpu_ids: Vec<u32> = (0..n)
            .map(|_| self.cpu_allocator.alloc().expect("checked can_alloc"))
            .collect();

        let bb = self.block_bytes;
        for layer in 0..self.num_layers {
            for i in 0..n {
                let g_off = gpu_ids[i] as usize * bb;
                let c_off = cpu_ids[i] as usize * bb;
                self.k_pools[layer]
                    .copy_to_host_at(
                        &mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb],
                        g_off,
                        bb,
                    )
                    .unwrap();
                self.v_pools[layer]
                    .copy_to_host_at(
                        &mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb],
                        g_off,
                        bb,
                    )
                    .unwrap();
            }
        }

        for b in gpu_ids {
            self.allocator.free(b);
        }
        let state = self.seq_states[slot].as_mut().unwrap();
        state.block_ids = cpu_ids;
        state.location = Location::Cpu;
        Ok(())
    }

    /// Bring `slot`'s KV back from host to GPU and free its CPU blocks.
    pub fn swap_in(&mut self, slot: usize) -> Result<(), &'static str> {
        let state = self.seq_states[slot]
            .as_ref()
            .ok_or("swap_in: empty slot")?;
        if state.location == Location::Gpu {
            return Ok(());
        }
        let cpu_ids = state.block_ids.clone();
        let n = cpu_ids.len();
        if !self.allocator.can_alloc(n) {
            return Err("swap_in: GPU pool full");
        }

        let gpu_ids: Vec<u32> = (0..n)
            .map(|_| self.allocator.alloc().expect("checked can_alloc"))
            .collect();

        let bb = self.block_bytes;
        for layer in 0..self.num_layers {
            for i in 0..n {
                let g_off = gpu_ids[i] as usize * bb;
                let c_off = cpu_ids[i] as usize * bb;
                self.k_pools[layer]
                    .copy_from_host_at(
                        &self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb],
                        g_off,
                        bb,
                    )
                    .unwrap();
                self.v_pools[layer]
                    .copy_from_host_at(
                        &self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb],
                        g_off,
                        bb,
                    )
                    .unwrap();
            }
        }

        for b in cpu_ids {
            self.cpu_allocator.free(b);
        }
        let state = self.seq_states[slot].as_mut().unwrap();
        state.block_ids = gpu_ids;
        state.location = Location::Gpu;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tiny_config() -> ModelConfig {
        serde_json::from_value(serde_json::json!({
            "model_type": "qwen3",
            "hidden_size": 8,
            "intermediate_size": 16,
            "num_attention_heads": 1,
            "num_key_value_heads": 1,
            "num_hidden_layers": 1,
            "vocab_size": 32,
            "max_position_embeddings": 64
        }))
        .unwrap()
    }

    #[test]
    fn truncate_sequence_frees_whole_blocks_and_keeps_slot_registered() {
        if xserv_cuda::device::set_device(0).is_err() {
            eprintln!("skipping CUDA-backed PagedKVCache test: device 0 unavailable");
            return;
        }

        let config = tiny_config();
        let mut cache = PagedKVCache::new(&config, 5, 0, 1, 4, DType::BF16, 0);

        assert_eq!(
            cache.truncate_sequence(1, 0),
            Err("truncate_sequence: slot out of range")
        );
        assert_eq!(
            cache.truncate_sequence(0, 0),
            Err("truncate_sequence: empty slot")
        );

        cache.register_sequence(0).unwrap();
        cache.ensure_capacity(0, BLOCK_SIZE * 3 + 1);
        cache.advance_seq_len(0, BLOCK_SIZE * 3 + 1);
        assert_eq!(cache.seq_len(0), BLOCK_SIZE * 3 + 1);
        assert_eq!(cache.block_count(0), 4);
        assert_eq!(cache.free_blocks(), 0);

        cache.truncate_sequence(0, BLOCK_SIZE + 1).unwrap();
        assert_eq!(cache.seq_len(0), BLOCK_SIZE + 1);
        assert_eq!(cache.block_count(0), 2);
        assert_eq!(cache.free_blocks(), 2);

        cache.truncate_sequence(0, BLOCK_SIZE).unwrap();
        assert_eq!(cache.seq_len(0), BLOCK_SIZE);
        assert_eq!(cache.block_count(0), 1);
        assert_eq!(cache.free_blocks(), 3);

        cache.truncate_sequence(0, 0).unwrap();
        assert_eq!(cache.seq_len(0), 0);
        assert_eq!(cache.block_count(0), 1);
        assert_eq!(cache.free_blocks(), 3);
        assert_eq!(
            cache.truncate_sequence(0, 1),
            Err("truncate_sequence: cannot extend")
        );
    }
}

unsafe fn tensor_from_owned_buf(
    buf: GpuBuffer,
    shape: &[usize],
    dtype: DType,
    device: u32,
) -> Tensor {
    use smallvec::SmallVec;
    use xserv_tensor::shape::contiguous_strides;
    use xserv_tensor::storage::Storage;

    let storage = Storage::cuda(buf, device);
    Tensor::from_storage(
        storage,
        SmallVec::from_slice(shape),
        contiguous_strides(shape),
        0,
        dtype,
    )
}