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

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

/// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
/// appends new K/V via D2D copy at offset (no CPU round-trip).
pub struct GpuKVCache {
    // Per layer: contiguous GPU buffer for K and V
    // Layout: [num_kv_heads, max_seq_len, head_dim] — contiguous per head
    k_bufs: Vec<GpuBuffer>,
    v_bufs: Vec<GpuBuffer>,
    // Per layer: pre-allocated staging buffers for get_kv_len output.
    // Size: num_kv_heads * max_seq_len * head_dim * elem_size (max possible output).
    // Avoids cudaMalloc/cudaFree on every get_kv_len call.
    k_staging: Vec<GpuBuffer>,
    v_staging: Vec<GpuBuffer>,
    seq_len: usize,
    max_seq_len: usize,
    num_kv_heads: usize,
    head_dim: usize,
    elem_size: usize,
    dtype: DType,
    device: u32,
}

impl GpuKVCache {
    pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType, device: u32) -> Self {
        let num_layers = config.num_layers();
        let num_kv_heads = config.num_kv_heads();
        let head_dim = config.head_dim();
        let elem_size = dtype.size_bytes();
        let buf_size = num_kv_heads * max_seq_len * head_dim * elem_size;

        let mut k_bufs = Vec::with_capacity(num_layers);
        let mut v_bufs = Vec::with_capacity(num_layers);
        let mut k_staging = Vec::with_capacity(num_layers);
        let mut v_staging = Vec::with_capacity(num_layers);
        for _ in 0..num_layers {
            let mut k = GpuBuffer::alloc(buf_size).expect("alloc KV cache K");
            let mut v = GpuBuffer::alloc(buf_size).expect("alloc KV cache V");
            k.zero().unwrap();
            v.zero().unwrap();
            k_bufs.push(k);
            v_bufs.push(v);
            k_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging K"));
            v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
        }

        Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
    }

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

    /// Append new K/V tensors for a given layer.
    /// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
    /// `write_pos` is the sequence position to write at (caller manages this).
    pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
        assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
        let es = self.elem_size;
        let hd = self.head_dim;
        let max_s = self.max_seq_len;
        let nh = self.num_kv_heads;

        let k_src = k_new.storage().gpu_buffer();
        let v_src = v_new.storage().gpu_buffer();

        for h in 0..nh {
            let src_off = h * new_tokens * hd * es;
            let dst_off = (h * max_s + write_pos) * hd * es;
            let count = new_tokens * hd * es;
            self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
            self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
        }
    }

    pub fn advance_seq_len(&mut self, new_tokens: usize) {
        self.seq_len += new_tokens;
    }

    /// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
    pub fn get_kv(&mut self, layer: usize) -> (Tensor, Tensor) {
        let sl = self.seq_len;
        self.get_kv_len(layer, sl)
    }

    pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
        let hd = self.head_dim;
        let nh = self.num_kv_heads;
        let es = self.elem_size;
        let max_s = self.max_seq_len;

        // Copy each head's valid portion into pre-allocated staging buffers.
        // Split borrows: staging (mut) vs cache (shared) are separate struct fields,
        // so the borrow checker allows simultaneous &mut staging + &cache.
        let out_size = nh * sl * hd * es;
        let k_stg = &mut self.k_staging[layer];
        let k_buf = &self.k_bufs[layer];
        let v_stg = &mut self.v_staging[layer];
        let v_buf = &self.v_bufs[layer];
        for h in 0..nh {
            let src_off = (h * max_s) * hd * es;
            let dst_off = (h * sl) * hd * es;
            let count = sl * hd * es;
            k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
            v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
        }
        // Grab raw pointers before dropping the mutable borrows
        let k_ptr = k_stg.as_mut_ptr();
        let v_ptr = v_stg.as_mut_ptr();

        // Create Tensors that borrow from the staging buffers (no cudaMalloc/cudaFree).
        // Safety: staging buffers are owned by GpuKVCache and outlive the returned Tensors
        // in practice (Tensors are consumed within the same forward pass before the next
        // get_kv_len call overwrites the staging buffer).
        let shape = &[1usize, nh, sl, hd];
        let k = unsafe {
            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
        };
        let v = unsafe {
            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
        };
        (k, v)
    }
}

/// Create a Tensor from a GpuBuffer (takes ownership).
unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
    use xserv_tensor::storage::Storage;
    use xserv_tensor::shape::contiguous_strides;
    use smallvec::SmallVec;

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