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).
    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 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();

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

        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;

            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();
            }

            t += chunk;
        }
    }

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

    /// 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(())
    }
}

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,
    )
}