ferrox_core/

cache.rs

//! A per-layer KV cache, growable one position at a time during decode.
//! Two growth strategies exist:
//!
//! - `with_pool`: PagedAttention-style block allocation. Many caches
//!   (one per concurrent request, typically) draw fixed-size blocks
//!   from one shared, bounded `KvBlockPool` instead of each
//!   independently pre-committing to a worst-case context length.
//!   Growth happens in fixed block-sized quanta, and a cache's blocks
//!   return to the shared pool when it's dropped, so the pool's free
//!   count is a real, live admission-control signal a caller can check
//!   before accepting a new request. This is the block-*allocation*
//!   half of PagedAttention; it does not (yet) change how attention
//!   reads a cache -- `k`/`v` are still read as one contiguous slice
//!   per sequence (see `Decoder::forward_token`/`forward_batch`), just
//!   backed by capacity that grows in block-sized steps instead of
//!   Rust's default exponential `Vec` growth. Wiring this into
//!   `ferrox-server` as live per-request admission control via
//!   `FERROX_KV_POOL_BLOCKS`/`FERROX_KV_POOL_BLOCK_SIZE`.

use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};

/// Returned by `KvCache::push` (and `with_pool`) when a pool-backed
/// cache needs another block but its shared `KvBlockPool` has none
/// free. Caches built with `new`/`with_capacity` never return this --
/// their growth is unconditional, matching their pre-paging behavior
/// exactly.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct KvPoolExhausted;

impl std::fmt::Display for KvPoolExhausted {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "KV cache block pool exhausted: no free blocks remain")
    }
}

impl std::error::Error for KvPoolExhausted {}

/// A bounded pool of fixed-size KV-cache blocks (in positions) shared
/// across many `KvCache` instances, typically one pool per server
/// process. Each `KvCache::with_pool` acquires one block up front and
/// one more each time it grows past its currently held capacity;
/// `free_blocks` is therefore a live, accurate admission-control
/// signal -- a caller can check it before accepting a new request
/// rather than discovering exhaustion only after committing memory.
pub struct KvBlockPool {
    block_size: usize,
    total_blocks: usize,
    free_blocks: usize,
}

impl KvBlockPool {
    /// `block_size` positions per block, `total_blocks` blocks in the
    /// whole shared budget (so `block_size * total_blocks` positions
    /// total, across however many caches draw from this pool at once).
    pub fn new(block_size: usize, total_blocks: usize) -> Self {
        assert!(block_size > 0, "block_size must be positive");
        KvBlockPool {
            block_size,
            total_blocks,
            free_blocks: total_blocks,
        }
    }

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

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

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

    /// Re-budget the pool.
    ///
    /// The pool is an *accounting* budget, not an allocator: each
    /// `KvCache` owns its own buffer and this counts how many blocks
    /// the deployment has promised. So a resize is arithmetic, with one
    /// rule that is not.
    ///
    /// Shrinking below what is currently held is REFUSED and the pool
    /// is left exactly as it was. `free_blocks` would have to go
    /// negative to represent it, and the alternative -- clamping it to
    /// zero -- silently over-promises: the caches already holding those
    /// blocks do not give them back, so every later `try_acquire`
    /// would be deciding against a budget that does not describe the
    /// memory in use.
    ///
    /// Returns the number of blocks currently held when it refuses, so
    /// the caller can say what the floor actually is rather than making
    /// the operator find it by being rejected.
    pub fn resize(&mut self, total_blocks: usize) -> Result<(), usize> {
        let in_use = self.total_blocks - self.free_blocks;
        if total_blocks < in_use {
            return Err(in_use);
        }
        self.free_blocks = total_blocks - in_use;
        self.total_blocks = total_blocks;
        Ok(())
    }

    fn try_acquire(&mut self, n: usize) -> bool {
        if n <= self.free_blocks {
            self.free_blocks -= n;
            true
        } else {
            false
        }
    }

    fn release(&mut self, n: usize) {
        self.free_blocks = (self.free_blocks + n).min(self.total_blocks);
    }
}

struct PooledState {
    pool: Arc<Mutex<KvBlockPool>>,
    block_size: usize,
    blocks_held: usize,
}

pub struct KvCache {
    pub n_kv_heads: usize,
    pub head_dim: usize,
    pub k: Vec<f32>, // [seq_len, n_kv_heads, head_dim], flattened
    pub v: Vec<f32>,
    pub seq_len: usize,
    /// The capacity (in positions) this cache was pre-allocated for,
    /// if any. `None` for caches built with `new` or `with_pool`.
    planned_capacity: Option<usize>,
    /// `Some` for caches built with `with_pool`; tracks the shared
    /// pool and how many blocks this cache currently holds, so its
    /// blocks can be returned on drop.
    pool_state: Option<PooledState>,
}

/// Cloning a pool-backed cache detaches the clone from pool accounting
/// (its `k`/`v`/`seq_len` data is copied normally, but the clone does
/// not hold or later release any blocks itself) -- mirroring how
/// `ferrox-models::prefix_cache` already uses `KvCache::clone` to fork
/// a cached prefix into a new, independent request's cache. Only the
/// original cache's blocks are released, exactly once, when it drops.
impl Clone for KvCache {
    fn clone(&self) -> Self {
        KvCache {
            n_kv_heads: self.n_kv_heads,
            head_dim: self.head_dim,
            k: self.k.clone(),
            v: self.v.clone(),
            seq_len: self.seq_len,
            planned_capacity: self.planned_capacity,
            pool_state: None,
        }
    }
}

impl Drop for KvCache {
    fn drop(&mut self) {
        if let Some(state) = &self.pool_state {
            if let Ok(mut pool) = state.pool.lock() {
                pool.release(state.blocks_held);
            }
        }
    }
}

impl KvCache {
    pub fn new(n_kv_heads: usize, head_dim: usize) -> Self {
        KvCache {
            n_kv_heads,
            head_dim,
            k: Vec::new(),
            v: Vec::new(),
            seq_len: 0,
            planned_capacity: None,
            pool_state: None,
        }
    }

    /// Pre-allocates storage for up to `max_seq_len` positions, so
    /// `push` never triggers a reallocation-and-copy during decode.
    /// Use this when the maximum context length is known ahead of time
    pub fn with_capacity(n_kv_heads: usize, head_dim: usize, max_seq_len: usize) -> Self {
        let elems_per_position = n_kv_heads * head_dim;
        KvCache {
            n_kv_heads,
            head_dim,
            k: Vec::with_capacity(max_seq_len * elems_per_position),
            v: Vec::with_capacity(max_seq_len * elems_per_position),
            seq_len: 0,
            planned_capacity: Some(max_seq_len),
            pool_state: None,
        }
    }

    /// Acquires up front however many blocks from `pool` are needed to
    /// cover `max_seq_len` positions (at least one, even if
    /// `max_seq_len` is `0`), so a caller that knows its worst-case
    /// sequence length ahead of time (as `ferrox-server` does: prompt
    /// length + `max_tokens`) never needs to acquire another block
    /// mid-decode. This matters beyond performance: `push` growing past
    /// its currently held capacity can fail if the pool is exhausted by
    /// *other* requests by then, and callers like
    /// `ferrox_models::Decoder::forward_token` treat `push` as
    /// infallible for non-pooled caches -- a pooled cache that
    /// under-reserves at construction and then fails to grow later
    /// would violate that assumption and panic mid-decode. Sizing to
    /// `max_seq_len` up front turns that into an admission-control
    /// decision made once, honestly, before any generation work starts,
    /// exactly mirroring `with_capacity`'s worst-case pre-allocation --
    /// just drawn from a shared pool instead of a private allocation.
    /// Returns `Err(KvPoolExhausted)` without mutating anything if the
    /// pool doesn't have that many blocks free.
    pub fn with_pool(
        n_kv_heads: usize,
        head_dim: usize,
        pool: Arc<Mutex<KvBlockPool>>,
        max_seq_len: usize,
    ) -> Result<Self, KvPoolExhausted> {
        let block_size = pool.lock().unwrap().block_size();
        let blocks_needed = max_seq_len.div_ceil(block_size).max(1);
        if !pool.lock().unwrap().try_acquire(blocks_needed) {
            return Err(KvPoolExhausted);
        }
        let elems_per_position = n_kv_heads * head_dim;
        Ok(KvCache {
            n_kv_heads,
            head_dim,
            k: Vec::with_capacity(blocks_needed * block_size * elems_per_position),
            v: Vec::with_capacity(blocks_needed * block_size * elems_per_position),
            seq_len: 0,
            planned_capacity: None,
            pool_state: Some(PooledState {
                pool,
                block_size,
                blocks_held: blocks_needed,
            }),
        })
    }

    /// Appends one position's key/value vectors (each
    /// `n_kv_heads * head_dim` long) to the cache. For pool-backed
    /// caches, this may need to acquire another block first; if the
    /// shared pool has none free, no data is appended and
    /// `Err(KvPoolExhausted)` is returned. Caches built with `new` or
    /// `with_capacity` always return `Ok`.
    pub fn push(&mut self, k_step: &[f32], v_step: &[f32]) -> Result<(), KvPoolExhausted> {
        assert_eq!(k_step.len(), self.n_kv_heads * self.head_dim);
        assert_eq!(v_step.len(), self.n_kv_heads * self.head_dim);

        let elems_per_position = self.n_kv_heads * self.head_dim;
        if let Some(state) = &mut self.pool_state {
            let capacity_positions = self.k.capacity() / elems_per_position;
            if self.seq_len == capacity_positions {
                if !state.pool.lock().unwrap().try_acquire(1) {
                    return Err(KvPoolExhausted);
                }
                state.blocks_held += 1;
                self.k.reserve_exact(state.block_size * elems_per_position);
                self.v.reserve_exact(state.block_size * elems_per_position);
            }
        }

        self.k.extend_from_slice(k_step);
        self.v.extend_from_slice(v_step);
        self.seq_len += 1;
        Ok(())
    }

    /// Advance length by `n` positions without storing real K/V values
    /// (zero-fill). Used when Metal owns the KV plane and the host cache
    /// only needs matching `seq_len` for sync checks.
    pub fn advance_len(&mut self, n: usize) -> Result<(), KvPoolExhausted> {
        if n == 0 {
            return Ok(());
        }
        let elems_per_position = self.n_kv_heads * self.head_dim;
        let zeros = vec![0f32; elems_per_position];
        for _ in 0..n {
            self.push(&zeros, &zeros)?;
        }
        Ok(())
    }

    /// Returns this cache's blocks to its shared pool immediately
    /// (rather than waiting for `Drop`) and detaches it from pool
    /// accounting; a no-op for caches that aren't pool-backed, and
    /// idempotent if called more than once.
    pub fn release_to_pool(&mut self) {
        if let Some(state) = self.pool_state.take() {
            if let Ok(mut pool) = state.pool.lock() {
                pool.release(state.blocks_held);
            }
        }
    }

    pub fn clear(&mut self) {
        self.k.clear();
        self.v.clear();
        self.seq_len = 0;
    }

    /// Rolls the cache back to exactly `new_seq_len` positions,
    /// discarding everything after. Used to reject speculatively
    /// decoded draft tokens that turned out wrong: their K/V were
    /// already pushed during batched verification, and rejection means
    /// removing them so the next real decode step continues from the
    /// last *accepted* position, not the last *attempted* one.
    pub fn truncate(&mut self, new_seq_len: usize) {
        assert!(
            new_seq_len <= self.seq_len,
            "truncate target {new_seq_len} must not exceed current seq_len {}",
            self.seq_len
        );
        let elems_per_position = self.n_kv_heads * self.head_dim;
        self.k.truncate(new_seq_len * elems_per_position);
        self.v.truncate(new_seq_len * elems_per_position);
        self.seq_len = new_seq_len;
    }

    /// Bytes currently resident for this cache's K and V buffers
    /// combined (actual allocated capacity, not just used length) --
    /// the number that matters for "does this fit in the context
    /// budget,"
    pub fn allocated_bytes(&self) -> usize {
        (self.k.capacity() + self.v.capacity()) * std::mem::size_of::<f32>()
    }

    /// True if this cache was pre-allocated via `with_capacity` and
    /// has not yet grown past that planned capacity (i.e. `push` has
    /// never had to reallocate). Useful for tests/diagnostics
    /// confirming the pre-allocation path actually avoided reallocs.
    pub fn is_within_planned_capacity(&self) -> bool {
        match self.planned_capacity {
            Some(cap) => {
                self.seq_len <= cap
                    && self.k.capacity() >= self.seq_len * self.n_kv_heads * self.head_dim
            }
            None => false,
        }
    }
}

/// The other half of PagedAttention that `KvBlockPool`/`KvCache::with_pool`
/// deliberately don't implement (see this module's doc comment): real,
/// *shared* physical block storage that many sequences' block tables can
/// address into, instead of each `KvCache` still owning its own private,
/// contiguous `Vec`. `KvBlockPool` only ever bounds a *count* of blocks
/// each cache may grow to; `PagedKvStore` is the actual backing memory,
/// and a sequence's `PagedKvCache` holds a block table (an ordered list
/// of block IDs into this shared store) instead of owning K/V data
/// directly. This is what makes non-contiguous-block reads during
/// attention (`causal_gqa_attention_paged`, in `attention.rs`) possible
/// at all -- `causal_gqa_attention`'s existing contiguous-slice read
/// pattern has no way to express "position 37 lives in block 12, cached
/// out of order relative to block 5."
pub struct PagedKvStore {
    block_size: usize,
    n_kv_heads: usize,
    head_dim: usize,
    k: Vec<f32>, // [total_blocks * block_size, n_kv_heads, head_dim], flattened
    v: Vec<f32>,
    free_block_ids: Vec<usize>,
}

impl PagedKvStore {
    pub fn new(block_size: usize, total_blocks: usize, n_kv_heads: usize, head_dim: usize) -> Self {
        assert!(block_size > 0, "block_size must be positive");
        let elems_per_block = block_size * n_kv_heads * head_dim;
        PagedKvStore {
            block_size,
            n_kv_heads,
            head_dim,
            k: vec![0.0; total_blocks * elems_per_block],
            v: vec![0.0; total_blocks * elems_per_block],
            // Pushed in descending order so `pop()` hands out ascending
            // block IDs -- not load-bearing for correctness (any free ID
            // works), just makes manual debugging/inspection saner.
            free_block_ids: (0..total_blocks).rev().collect(),
        }
    }

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

    pub fn free_block_count(&self) -> usize {
        self.free_block_ids.len()
    }

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

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

    fn acquire_block(&mut self) -> Option<usize> {
        self.free_block_ids.pop()
    }

    fn release_block(&mut self, id: usize) {
        self.free_block_ids.push(id);
    }

    fn elems_per_block(&self) -> usize {
        self.block_size * self.n_kv_heads * self.head_dim
    }

    /// One position's K (or V) row within block `id` at `offset` (0-based
    /// within the block) -- `[n_kv_heads * head_dim]` long. Used by
    /// `causal_gqa_attention_paged` to read attention inputs directly out
    /// of shared physical storage via a block table, and by
    /// `PagedKvCache::push` to write a new position into it.
    pub fn k_row(&self, id: usize, offset: usize) -> &[f32] {
        let elems_per_position = self.n_kv_heads * self.head_dim;
        let start = id * self.elems_per_block() + offset * elems_per_position;
        &self.k[start..start + elems_per_position]
    }

    pub fn v_row(&self, id: usize, offset: usize) -> &[f32] {
        let elems_per_position = self.n_kv_heads * self.head_dim;
        let start = id * self.elems_per_block() + offset * elems_per_position;
        &self.v[start..start + elems_per_position]
    }

    fn k_row_mut(&mut self, id: usize, offset: usize) -> &mut [f32] {
        let elems_per_position = self.n_kv_heads * self.head_dim;
        let start = id * self.elems_per_block() + offset * elems_per_position;
        &mut self.k[start..start + elems_per_position]
    }

    fn v_row_mut(&mut self, id: usize, offset: usize) -> &mut [f32] {
        let elems_per_position = self.n_kv_heads * self.head_dim;
        let start = id * self.elems_per_block() + offset * elems_per_position;
        &mut self.v[start..start + elems_per_position]
    }
}

/// Returned when a `PagedKvCache` needs another block but its
/// `PagedKvStore` has none free -- the paged-storage analog of
/// `KvPoolExhausted`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PagedStoreExhausted;

impl std::fmt::Display for PagedStoreExhausted {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "paged KV store exhausted: no free blocks remain")
    }
}

impl std::error::Error for PagedStoreExhausted {}

/// One sequence's view into a shared `PagedKvStore`: a block table
/// (which physical blocks this sequence's positions live in, in order)
/// plus how many positions have been written so far. Unlike `KvCache`,
/// this holds no K/V data itself -- every read and write goes through
/// the shared store.
#[derive(Debug, Clone, Default)]
pub struct PagedKvCache {
    block_table: Vec<usize>,
    seq_len: usize,
}

impl PagedKvCache {
    pub fn new() -> Self {
        PagedKvCache {
            block_table: Vec::new(),
            seq_len: 0,
        }
    }

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

    pub fn block_table(&self) -> &[usize] {
        &self.block_table
    }

    /// Appends one position's key/value vectors, acquiring a new block
    /// from `store` first if the current tail block is full (or none
    /// held yet). Mirrors `KvCache::push`'s signature/semantics exactly,
    /// just against shared storage instead of a private buffer.
    pub fn push(
        &mut self,
        store: &mut PagedKvStore,
        k_step: &[f32],
        v_step: &[f32],
    ) -> Result<(), PagedStoreExhausted> {
        let block_size = store.block_size();
        let offset_in_block = self.seq_len % block_size;
        // Index by position rather than taking the tail block: a
        // sequence that pre-reserved (see `reserve`) already holds the
        // block this position belongs in, and appending another would
        // both leak a block and write the row in the wrong place.
        let block_index = self.seq_len / block_size;
        if block_index >= self.block_table.len() {
            let id = store.acquire_block().ok_or(PagedStoreExhausted)?;
            self.block_table.push(id);
        }
        let block_id = self.block_table[block_index];
        store
            .k_row_mut(block_id, offset_in_block)
            .copy_from_slice(k_step);
        store
            .v_row_mut(block_id, offset_in_block)
            .copy_from_slice(v_step);
        self.seq_len += 1;
        Ok(())
    }

    /// Appends a block the caller already owns, without taking one from
    /// the store.
    ///
    /// This is how a sliding window recycles. A block whose positions
    /// have fallen behind the window is never read again -- the paged
    /// attention kernel indexes `block_table[t / block_size]` only for
    /// `t >= seq_len - window` -- so its storage can back a *later*
    /// position instead of being handed back and re-acquired. The table
    /// keeps its absolute-position indexing and simply names the same
    /// physical block at two indices: the stale one, which nothing
    /// reads, and the live one.
    ///
    /// That aliasing is the reason this is a separate method rather than
    /// a flag on [`Self::reserve`]. A caller that recycles owns the
    /// obligation to release each distinct block exactly once, and to
    /// have established that the donor index really is out of window --
    /// neither of which this type can check for itself.
    pub fn append_block(&mut self, block_id: usize) {
        self.block_table.push(block_id);
    }

    /// Releases every block this sequence holds back to `store`. Must be
    /// called explicitly (there's no `Drop` here, since dropping needs a
    /// `&mut PagedKvStore` this type doesn't own a reference to) --
    /// mirrors `KvCache::release_to_pool`, just not automatic.
    ///
    /// Each *distinct* block once: a table that has recycled through
    /// [`Self::append_block`] names one block at more than one index, and
    /// releasing per index would put the same id on the free list twice,
    /// after which two sequences are handed the same memory.
    pub fn release(&mut self, store: &mut PagedKvStore) {
        let mut seen: Vec<usize> = Vec::new();
        for id in self.block_table.drain(..) {
            if !seen.contains(&id) {
                seen.push(id);
                store.release_block(id);
            }
        }
        self.seq_len = 0;
    }

    /// How many *additional* blocks appending `n_new` positions would
    /// take from `store`, given what this sequence already holds.
    ///
    /// Counted against held CAPACITY rather than against `seq_len`, so
    /// it is right in both cases. The tail block is usually part-full,
    /// so the answer is never simply `n_new / block_size`: positions
    /// that land in a block already held cost nothing. And a sequence
    /// that pre-reserved (see [`Self::reserve`]) holds blocks beyond
    /// its length, which a `seq_len`-only sum would ask for twice.
    ///
    /// Callers that must not fail part-way through a write check this
    /// against [`PagedKvStore::free_block_count`] before touching
    /// anything.
    pub fn blocks_needed_for(&self, store: &PagedKvStore, n_new: usize) -> usize {
        let held_capacity = self.block_table.len() * store.block_size();
        let unused = held_capacity.saturating_sub(self.seq_len);
        n_new.saturating_sub(unused).div_ceil(store.block_size())
    }

    /// Takes the blocks `n_new` more positions will need, without
    /// advancing `seq_len`.
    ///
    /// This is what makes a multi-layer append all-or-nothing. The
    /// check and the taking happen together, so every later
    /// [`Self::push`] writes into a block this sequence already owns
    /// and cannot fail. Reserving and then not filling is harmless: the
    /// blocks are this sequence's until it releases, and `seq_len`
    /// still says how far it really got.
    pub fn reserve(
        &mut self,
        store: &mut PagedKvStore,
        n_new: usize,
    ) -> Result<(), PagedStoreExhausted> {
        let need = self.blocks_needed_for(store, n_new);
        if need > store.free_block_count() {
            return Err(PagedStoreExhausted);
        }
        for _ in 0..need {
            let id = store
                .acquire_block()
                .expect("checked against free_block_count immediately above");
            self.block_table.push(id);
        }
        Ok(())
    }

    /// Installs a block table the caller allocated, with `seq_len`
    /// positions already computed in it.
    ///
    /// This is how a sequence starts life on top of a cached prefix:
    /// the blocks are somebody else's, already full, and this sequence
    /// appends past them. `seq_len` MUST be a whole number of blocks,
    /// because the first append writes at `seq_len` and a shared block
    /// must never be written -- another sequence is attending over it.
    /// A ragged length would put that write inside the last shared
    /// block, corrupting a prefix every other holder is reading.
    pub fn adopt_blocks(&mut self, block_table: Vec<usize>, seq_len: usize, block_size: usize) {
        assert_eq!(
            seq_len % block_size,
            0,
            "an adopted prefix must end on a block boundary, or the first \
             append writes into a block another sequence is reading"
        );
        assert!(
            seq_len / block_size <= block_table.len(),
            "block table too short for the adopted length"
        );
        self.block_table = block_table;
        self.seq_len = seq_len;
    }

    /// Copies this sequence's KV out of the shared store into a plain
    /// contiguous [`KvCache`].
    ///
    /// This is what lets the batched prefill path run *unchanged* over
    /// paged storage. Its fast arm hands `cache.k` / `cache.v` to a
    /// blocked kernel that reads them as flat slices, and a block table
    /// cannot be expressed that way. Rather than maintain a second
    /// prefill kernel that reads through the table -- a copy that could
    /// drift from the one every other model path uses -- the pages are
    /// materialised once per layer, the existing kernel runs, and the
    /// new rows go back with [`Self::append_contiguous`].
    ///
    /// The cost is one `seq_len * n_kv_heads * head_dim` copy per layer
    /// per prefill call, against matmuls that dominate prefill. Decode
    /// still reads through the block table and copies nothing, which is
    /// where page sharing actually pays.
    pub fn to_contiguous(&self, store: &PagedKvStore) -> KvCache {
        let elems_per_position = store.n_kv_heads * store.head_dim;
        let mut cache = KvCache::with_capacity(store.n_kv_heads, store.head_dim, self.seq_len);
        cache.k.reserve_exact(self.seq_len * elems_per_position);
        cache.v.reserve_exact(self.seq_len * elems_per_position);
        for pos in 0..self.seq_len {
            let block_id = self.block_table[pos / store.block_size];
            let offset = pos % store.block_size;
            cache.k.extend_from_slice(store.k_row(block_id, offset));
            cache.v.extend_from_slice(store.v_row(block_id, offset));
        }
        cache.seq_len = self.seq_len;
        cache
    }

    /// Appends `count` positions' worth of contiguous K/V rows, the
    /// inverse of [`Self::to_contiguous`].
    ///
    /// Blocks are reserved for the whole append *before* the first row
    /// is written, so a store that cannot hold the request refuses it
    /// having changed nothing. Writing rows until the store runs dry
    /// would leave the sequence with a `seq_len` that disagrees with
    /// the model's own idea of how far it has got, which is not a
    /// recoverable state.
    pub fn append_contiguous(
        &mut self,
        store: &mut PagedKvStore,
        k: &[f32],
        v: &[f32],
        count: usize,
    ) -> Result<(), PagedStoreExhausted> {
        let elems_per_position = store.n_kv_heads * store.head_dim;
        assert_eq!(k.len(), count * elems_per_position, "k row count");
        assert_eq!(v.len(), count * elems_per_position, "v row count");
        if self.blocks_needed_for(store, count) > store.free_block_count() {
            return Err(PagedStoreExhausted);
        }
        for i in 0..count {
            let lo = i * elems_per_position;
            let hi = lo + elems_per_position;
            self.push(store, &k[lo..hi], &v[lo..hi])
                .expect("blocks reserved above, so no push here can exhaust the store");
        }
        Ok(())
    }
}

/// Per-layer [`PagedKvStore`]s that many concurrent requests share.
///
/// # Why a lock per layer, and why two phases
///
/// `ferrox-server` runs generation on `spawn_blocking` with, in its own
/// words, "no I/O and no shared lock". `KvBlockPool` survives that
/// because it only bounds a *count*: each `KvCache` owns a private
/// `Vec`, and the pool mutex is taken briefly at acquire and release,
/// never during a forward. A `PagedKvStore` is the opposite -- it IS
/// the backing memory -- so sharing one across concurrent requests
/// needs an answer to "who may touch these bytes when".
///
/// The answer the API already implies: attention takes
/// `&PagedKvStore` and only `push` takes `&mut`. So the accesses split
/// cleanly into many concurrent readers and one short exclusive write
/// per position, which is exactly an `RwLock` -- and one per LAYER
/// rather than one for the whole model, so two requests contend only
/// when both are writing the same layer at the same instant.
///
/// A caller must therefore take the write guard for the push alone and
/// drop it before attending under a read guard. Holding the write
/// guard across attention would serialise the expensive half and give
/// back a global lock with extra steps. Nothing breaks in the gap: a
/// sequence's block table and length are its own, and another
/// request's push in between only touches blocks it exclusively holds.
///
/// # Deadlock
///
/// [`Self::write_all`] is the one place several layers are held at
/// once, and it takes them in ascending layer order. Every caller
/// getting the same order is what makes that safe; there is no other
/// multi-layer acquisition in the codebase, and a new one must follow
/// the same rule.
///
/// # Poisoning
///
/// A panic while holding a store leaves the KV mid-write, which is not
/// recoverable state, but it is also not *unsound* -- the bytes are
/// plain `f32`. Poison is stepped over with `into_inner`, matching how
/// `ferrox-server` already treats its pool mutex: a poisoned lock
/// should not turn one request's panic into a permanently dead server.
pub struct SharedPagedKv {
    layers: Vec<RwLock<PagedKvStore>>,
    /// Guarded separately from the layers, and always taken BEFORE
    /// them, never while a layer guard is held. That one-way order is
    /// what keeps group allocation and the per-layer push paths from
    /// deadlocking against each other.
    groups: Mutex<GroupTable>,
}

impl SharedPagedKv {
    /// One store per layer, each with `blocks_per_layer` blocks.
    pub fn new(
        n_layers: usize,
        block_size: usize,
        blocks_per_layer: usize,
        n_kv_heads: usize,
        head_dim: usize,
    ) -> Self {
        SharedPagedKv {
            layers: (0..n_layers)
                .map(|_| {
                    RwLock::new(PagedKvStore::new(
                        block_size,
                        blocks_per_layer,
                        n_kv_heads,
                        head_dim,
                    ))
                })
                .collect(),
            groups: Mutex::new(GroupTable::default()),
        }
    }

    /// Wraps stores the caller built, for tests and for callers that
    /// size layers differently.
    pub fn from_stores(stores: Vec<PagedKvStore>) -> Self {
        SharedPagedKv {
            layers: stores.into_iter().map(RwLock::new).collect(),
            groups: Mutex::new(GroupTable::default()),
        }
    }

    pub fn layer_count(&self) -> usize {
        self.layers.len()
    }

    /// Shared access to one layer, for attention.
    pub fn read(&self, layer: usize) -> RwLockReadGuard<'_, PagedKvStore> {
        self.layers[layer]
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    /// Exclusive access to one layer, for a push. Hold it for the push
    /// and nothing else -- see the type docs.
    pub fn write(&self, layer: usize) -> RwLockWriteGuard<'_, PagedKvStore> {
        self.layers[layer]
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    /// Every layer at once, in ascending order, so a multi-layer append
    /// is atomic against other requests.
    ///
    /// This is what makes "all layers advance or none do" hold under
    /// concurrency rather than only single-threaded: checking free
    /// space and then appending are separate steps, and without the
    /// guards spanning both, another request can take the blocks in
    /// between and leave this one half-written.
    ///
    /// Ascending order is the deadlock rule; see the type docs.
    pub fn write_all(&self) -> Vec<RwLockWriteGuard<'_, PagedKvStore>> {
        self.layers
            .iter()
            .map(|l| l.write().unwrap_or_else(|poisoned| poisoned.into_inner()))
            .collect()
    }

    /// Free blocks in one layer, for admission control. A snapshot: by
    /// the time a caller acts on it another request may have taken
    /// them, which is why the append itself re-checks under the guard.
    pub fn free_blocks(&self, layer: usize) -> usize {
        self.read(layer).free_block_count()
    }

    /// Takes one block from EVERY layer as a single group, refcount 1.
    ///
    /// All layers or none: a group that existed in some layers and not
    /// others could not answer "which block holds position p in layer
    /// l", which is the only question it exists to answer.
    pub fn acquire_group(&self) -> Option<PageGroup> {
        let mut guards = self.write_all();
        if guards.iter().any(|s| s.free_block_count() == 0) {
            return None;
        }
        let blocks: Vec<usize> = guards
            .iter_mut()
            .map(|s| {
                s.acquire_block()
                    .expect("checked every layer under these same guards")
            })
            .collect();
        let mut groups = self
            .groups
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        Some(PageGroup(groups.insert(blocks)))
    }

    /// One more holder of `group`.
    ///
    /// Called when a second sequence adopts a cached prefix. Without
    /// it, the first sequence to finish frees pages the second is
    /// still attending over -- a use-after-free that shows up as
    /// another conversation's tokens rather than as a crash.
    pub fn retain_group(&self, group: PageGroup) {
        let mut groups = self
            .groups
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        groups.retain(group.0);
    }

    /// One fewer holder. At zero the blocks go back to their layers.
    ///
    /// Returns whether this was the last holder, so a caller can assert
    /// on it rather than guess.
    pub fn release_group(&self, group: PageGroup) -> bool {
        let blocks = {
            let mut groups = self
                .groups
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            match groups.release(group.0) {
                Some(blocks) => blocks,
                None => return false,
            }
        };
        // The groups lock is dropped before the layer guards are taken,
        // so the lock order is always groups-then-layers and never the
        // reverse. See the type docs on deadlock.
        let mut guards = self.write_all();
        for (store, block) in guards.iter_mut().zip(blocks) {
            store.release_block(block);
        }
        true
    }

    /// Which block in each layer this group owns, indexed by layer.
    pub fn group_blocks(&self, group: PageGroup) -> Vec<usize> {
        let groups = self
            .groups
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        groups.blocks(group.0).to_vec()
    }

    /// How many holders `group` has. Zero means it does not exist.
    pub fn group_refs(&self, group: PageGroup) -> u32 {
        let groups = self
            .groups
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        groups.refs(group.0)
    }

    /// Groups that could still be allocated, bounded by the layer with
    /// the fewest free blocks: a group needs one from each.
    pub fn free_groups(&self) -> usize {
        (0..self.layers.len())
            .map(|l| self.free_blocks(l))
            .min()
            .unwrap_or(0)
    }
}

/// A handle to one block in every layer.
///
/// The unit of sharing between sequences, and the only thing small
/// enough to be what a radix prefix cache stores: that cache maps a
/// token prefix to ONE index per token, while a position's KV lives in
/// `n_layers` different blocks. A group is the name for all of them.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PageGroup(pub u32);

/// Group ids, their per-layer blocks, and how many holders each has.
#[derive(Debug, Default)]
struct GroupTable {
    /// Indexed by group id. `None` for an id currently on the free list.
    blocks: Vec<Option<Vec<usize>>>,
    refs: Vec<u32>,
    free_ids: Vec<u32>,
}

impl GroupTable {
    fn insert(&mut self, blocks: Vec<usize>) -> u32 {
        if let Some(id) = self.free_ids.pop() {
            self.blocks[id as usize] = Some(blocks);
            self.refs[id as usize] = 1;
            return id;
        }
        self.blocks.push(Some(blocks));
        self.refs.push(1);
        (self.blocks.len() - 1) as u32
    }

    fn retain(&mut self, id: u32) {
        let refs = &mut self.refs[id as usize];
        assert!(*refs > 0, "cannot retain group {id}, which has no holders");
        *refs += 1;
    }

    /// Drops one holder, returning the blocks to free only when the
    /// last one goes.
    fn release(&mut self, id: u32) -> Option<Vec<usize>> {
        let refs = &mut self.refs[id as usize];
        assert!(*refs > 0, "double free of group {id}");
        *refs -= 1;
        if *refs > 0 {
            return None;
        }
        // The id is reusable now, but only after the blocks are out:
        // handing the id back while it still named blocks would let a
        // later `acquire_group` believe it owns them too.
        let blocks = self.blocks[id as usize]
            .take()
            .expect("a group with holders always has blocks");
        self.free_ids.push(id);
        Some(blocks)
    }

    fn blocks(&self, id: u32) -> &[usize] {
        self.blocks[id as usize]
            .as_deref()
            .expect("group has no blocks; it was already released")
    }

    fn refs(&self, id: u32) -> u32 {
        self.refs.get(id as usize).copied().unwrap_or(0)
    }
}

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

    /// `blocks_needed_for` is the reservation the whole no-partial-write
    /// guarantee rests on, and it is wrong in two opposite directions
    /// that fail very differently.
    ///
    /// UNDER-counting is the dangerous one: `append_contiguous` reserves
    /// on this answer and then pushes with an `expect`, so too small a
    /// number panics part-way through a layer -- exactly the corrupted
    /// state the reservation exists to prevent. Over-counting merely
    /// refuses a request that would have fitted.
    ///
    /// Both mistakes are one edit away. Flooring instead of ceiling
    /// under-counts whenever the append does not land on a block
    /// boundary; ignoring the part-full tail over-counts whenever a
    /// sequence is mid-block, which after the first token is almost
    /// always. Neither shows up when the numbers happen to divide
    /// evenly, so the cases here are chosen so that they do not.
    #[test]
    fn blocks_needed_for_accounts_for_the_part_full_tail_block() {
        let mut store = PagedKvStore::new(/* block_size = */ 4, 64, 1, 1);
        let mut cache = PagedKvCache::new();
        let row = [1.0f32];
        // Real pushes rather than poking `seq_len`: the count is
        // against blocks this sequence HOLDS, so a length with no
        // blocks behind it is a state that cannot occur and would only
        // let the test agree with an arithmetic nothing produces.
        let advance = |cache: &mut PagedKvCache, store: &mut PagedKvStore, n: usize| {
            for _ in 0..n {
                cache.push(store, &row, &row).unwrap();
            }
        };

        // Empty: a whole-block boundary, and a remainder that a floor
        // would round away.
        assert_eq!(cache.blocks_needed_for(&store, 0), 0);
        assert_eq!(cache.blocks_needed_for(&store, 1), 1);
        assert_eq!(cache.blocks_needed_for(&store, 4), 1);
        assert_eq!(cache.blocks_needed_for(&store, 5), 2, "5 into 4s needs 2");

        // One position in: three slots free in the tail, so appending up
        // to three costs NOTHING. Ignoring the tail would say 1.
        advance(&mut cache, &mut store, 1);
        assert_eq!(cache.blocks_needed_for(&store, 3), 0, "fits in the tail");
        assert_eq!(cache.blocks_needed_for(&store, 4), 1);
        assert_eq!(cache.blocks_needed_for(&store, 8), 2);

        // The awkward case: 1 free in the tail, 6 to append. 5 spill
        // over 4-wide blocks, so 2. A floor gives 1 and a tail-blind
        // ceil gives 2 for the wrong reason, so this pins the shape.
        advance(&mut cache, &mut store, 2); // seq_len = 3
        assert_eq!(cache.blocks_needed_for(&store, 6), 2);
        assert_eq!(cache.blocks_needed_for(&store, 5), 1);

        // Tail exactly full: no free slots, so this behaves like empty.
        advance(&mut cache, &mut store, 1); // seq_len = 4
        assert_eq!(cache.blocks_needed_for(&store, 1), 1);
        assert_eq!(cache.blocks_needed_for(&store, 4), 1);

        // A RESERVED block is capacity this sequence already holds, so
        // it must not be asked for twice. Counting from `seq_len` alone
        // would say 1 here and take a second block for positions the
        // reservation already covers.
        cache.reserve(&mut store, 4).unwrap();
        assert_eq!(
            cache.blocks_needed_for(&store, 4),
            0,
            "a reserved block is already held"
        );
        assert_eq!(cache.blocks_needed_for(&store, 5), 1);
    }

    /// A group takes one block from every layer, and gives them all
    /// back together.
    ///
    /// All-or-nothing is the point: a group holding blocks in some
    /// layers and not others cannot answer "which block holds position
    /// p in layer l", which is the only question it exists for.
    #[test]
    fn a_group_takes_one_block_from_every_layer_and_returns_them_together() {
        let kv = SharedPagedKv::new(3, 2, 4, 1, 1);
        assert_eq!(kv.free_groups(), 4);

        let g = kv.acquire_group().expect("4 groups available");
        let blocks = kv.group_blocks(g);
        assert_eq!(blocks.len(), 3, "one block per layer");
        for l in 0..3 {
            assert_eq!(kv.free_blocks(l), 3, "layer {l} gave up exactly one");
        }
        assert_eq!(kv.free_groups(), 3);

        assert!(kv.release_group(g), "sole holder, so this frees it");
        for l in 0..3 {
            assert_eq!(kv.free_blocks(l), 4, "layer {l} got its block back");
        }
        assert_eq!(kv.free_groups(), 4);
    }

    /// A group survives until its LAST holder releases it.
    ///
    /// This is what makes prefix sharing safe. Two sequences off one
    /// system prompt hold the same pages; if the first to finish freed
    /// them, the second would keep attending over blocks the store had
    /// already handed to somebody else -- surfacing as another
    /// conversation's tokens, not as a crash.
    #[test]
    fn a_group_shared_by_two_holders_survives_the_first_release() {
        let kv = SharedPagedKv::new(2, 2, 2, 1, 1);
        let g = kv.acquire_group().unwrap();
        let blocks = kv.group_blocks(g);
        kv.retain_group(g);
        assert_eq!(kv.group_refs(g), 2);

        assert!(
            !kv.release_group(g),
            "one holder remains, so nothing is freed"
        );
        assert_eq!(kv.group_refs(g), 1);
        assert_eq!(kv.free_blocks(0), 1, "the blocks are still held");
        assert_eq!(kv.group_blocks(g), blocks, "and still name the same blocks");

        assert!(kv.release_group(g), "last holder frees it");
        assert_eq!(kv.group_refs(g), 0);
        assert_eq!(kv.free_blocks(0), 2);
    }

    /// Exhaustion is per group, bounded by the tightest layer.
    ///
    /// A layer with one block left caps the whole pool at one more
    /// group however much room the others have, because a group needs
    /// one block from each.
    #[test]
    fn group_capacity_is_bounded_by_the_layer_with_the_fewest_blocks() {
        let kv = SharedPagedKv::from_stores(vec![
            PagedKvStore::new(2, 5, 1, 1),
            PagedKvStore::new(2, 1, 1, 1),
        ]);
        assert_eq!(kv.free_groups(), 1, "layer 1 has only one block");

        let g = kv.acquire_group().expect("one group fits");
        assert_eq!(kv.free_groups(), 0);
        assert!(
            kv.acquire_group().is_none(),
            "layer 1 is empty, so no group can be formed"
        );
        // The refused attempt must not have taken layer 0's block.
        assert_eq!(kv.free_blocks(0), 4, "a refused group leaks nothing");
        kv.release_group(g);
        assert_eq!(kv.free_blocks(0), 5);
    }

    /// A released id is reused, with a refcount that starts over.
    #[test]
    fn a_released_group_id_is_reused_with_a_fresh_refcount() {
        let kv = SharedPagedKv::new(1, 2, 2, 1, 1);
        let first = kv.acquire_group().unwrap();
        kv.retain_group(first);
        assert_eq!(kv.group_refs(first), 2);
        kv.release_group(first);
        kv.release_group(first);
        assert_eq!(kv.group_refs(first), 0, "gone, not merely decremented");

        let second = kv.acquire_group().unwrap();
        assert_eq!(second, first, "the id is reused");
        assert_eq!(
            kv.group_refs(second),
            1,
            "a reused id must not inherit the old count"
        );
        assert_eq!(kv.group_blocks(second).len(), 1);
        assert_eq!(kv.free_blocks(0), 1);
    }

    /// Reading a group after its last holder released it PANICS rather
    /// than answering with stale blocks.
    ///
    /// This is the observable half of clearing the entry on release,
    /// and the reason it is `take` rather than `clone`: a caller still
    /// holding a `PageGroup` after releasing it is exactly the bug
    /// refcounting exists to prevent, and blocks that now belong to
    /// somebody else are the worst possible answer -- the caller reads
    /// another sequence's KV and nothing says so.
    ///
    /// Written after sabotage showed the previous test here passed with
    /// `clone` in place of `take`: `insert` overwrites the entry on
    /// reuse, so a stale entry was never reachable through the path
    /// that test took. This one reaches it.
    #[test]
    #[should_panic(expected = "already released")]
    fn reading_a_released_group_panics_rather_than_returning_stale_blocks() {
        let kv = SharedPagedKv::new(2, 2, 2, 1, 1);
        let g = kv.acquire_group().unwrap();
        assert!(kv.release_group(g));
        let _ = kv.group_blocks(g);
    }

    /// Releasing a group nobody holds is a bug, not a no-op.
    ///
    /// Silently ignoring it would let a double release return the same
    /// blocks to the store twice, after which two sequences are handed
    /// the same page and both write it.
    #[test]
    #[should_panic(expected = "double free of group")]
    fn releasing_a_group_twice_panics_rather_than_freeing_it_twice() {
        let kv = SharedPagedKv::new(1, 2, 2, 1, 1);
        let g = kv.acquire_group().unwrap();
        assert!(kv.release_group(g));
        kv.release_group(g);
    }

    /// A recycled block backs a later position without the store ever
    /// being asked for another one, and the later position's writes are
    /// what a read at that position returns.
    ///
    /// This is the whole sliding-window mechanism in miniature. Blocks
    /// of two, four positions, and only two blocks in the store: without
    /// recycling, position 2 has nowhere to go.
    #[test]
    fn a_recycled_block_backs_a_later_position_without_touching_the_store() {
        let mut store = PagedKvStore::new(2, 2, 1, 2);
        let mut cache = PagedKvCache::new();
        cache.push(&mut store, &[1.0, 1.0], &[1.0, 1.0]).unwrap();
        cache.push(&mut store, &[2.0, 2.0], &[2.0, 2.0]).unwrap();
        assert_eq!(store.free_block_count(), 1, "one block per position pair");

        // Positions 0..2 have fallen behind a window of two. Their block
        // backs positions 2..4 instead, and the store is untouched.
        let recycled = cache.block_table()[0];
        cache.append_block(recycled);
        assert_eq!(
            store.free_block_count(),
            1,
            "recycling must not take a block from the store"
        );
        cache.push(&mut store, &[3.0, 3.0], &[3.0, 3.0]).unwrap();
        assert_eq!(cache.seq_len(), 3);
        assert_eq!(
            cache.block_table(),
            &[recycled, recycled],
            "the same block at the stale index and the live one"
        );

        // Reading position 2 sees the new row. Position 0's row is gone,
        // which is exactly what "behind the window" means -- the kernel
        // never indexes it.
        let flat = cache.to_contiguous(&store);
        assert_eq!(&flat.k[4..6], &[3.0, 3.0], "position 2 reads its own row");
        assert_eq!(
            &flat.k[0..2],
            &[3.0, 3.0],
            "position 0 now reads the recycled row, and nothing may read it"
        );
    }

    /// Releasing an aliased table hands each block back ONCE.
    ///
    /// Per index instead of per distinct block would put the recycled id
    /// on the free list twice, and the next two acquisitions would hand
    /// two sequences the same memory -- which does not fail, it
    /// interleaves two conversations' KV.
    #[test]
    fn releasing_a_recycled_table_gives_each_block_back_once() {
        // Exactly one block in the store, so "handed back twice" is
        // observable as a second acquisition succeeding.
        let mut store = PagedKvStore::new(2, 1, 1, 2);
        let mut cache = PagedKvCache::new();
        cache.push(&mut store, &[1.0, 1.0], &[1.0, 1.0]).unwrap();
        let held = cache.block_table()[0];
        cache.append_block(held);
        cache.append_block(held);

        let free_before = store.free_block_count();
        cache.release(&mut store);
        assert_eq!(
            store.free_block_count(),
            free_before + 1,
            "three table entries naming one block are one block back"
        );
        // And the store agrees: it can hand out that block once.
        assert!(store.acquire_block().is_some());
        assert!(store.acquire_block().is_none());
    }

    #[test]
    fn a_gathered_sequence_round_trips_through_the_store() {
        let mut store = PagedKvStore::new(2, 8, 2, 2);
        let mut cache = PagedKvCache::new();
        // Five positions over blocks of two: the tail block is half
        // full, which is where an off-by-one in the gather shows up.
        let rows: Vec<[f32; 4]> = (0..5)
            .map(|i| {
                let b = i as f32 * 10.0;
                [b + 1.0, b + 2.0, b + 3.0, b + 4.0]
            })
            .collect();
        for r in &rows {
            cache.push(&mut store, r, r).unwrap();
        }

        let flat = cache.to_contiguous(&store);
        assert_eq!(flat.seq_len, 5);
        assert_eq!(flat.k.len(), 5 * 4);
        for (i, r) in rows.iter().enumerate() {
            assert_eq!(&flat.k[i * 4..(i + 1) * 4], r, "position {i} k");
            assert_eq!(&flat.v[i * 4..(i + 1) * 4], r, "position {i} v");
        }

        // And appending those same rows back onto a fresh sequence
        // reproduces the store's view of them exactly.
        let mut rebuilt = PagedKvCache::new();
        let mut store2 = PagedKvStore::new(2, 8, 2, 2);
        rebuilt
            .append_contiguous(&mut store2, &flat.k, &flat.v, 5)
            .unwrap();
        let again = rebuilt.to_contiguous(&store2);
        assert_eq!(again.k, flat.k);
        assert_eq!(again.v, flat.v);
        assert_eq!(again.seq_len, flat.seq_len);
    }

    #[test]
    fn push_grows_seq_len_and_stores_values() {
        let mut cache = KvCache::new(2, 2);
        cache
            .push(&[1.0, 2.0, 3.0, 4.0], &[5.0, 6.0, 7.0, 8.0])
            .unwrap();
        assert_eq!(cache.seq_len, 1);
        cache
            .push(&[9.0, 10.0, 11.0, 12.0], &[13.0, 14.0, 15.0, 16.0])
            .unwrap();
        assert_eq!(cache.seq_len, 2);
        assert_eq!(cache.k.len(), 2 * 2 * 2);
        assert_eq!(cache.k[4], 9.0);
    }

    #[test]
    #[should_panic]
    fn push_wrong_size_panics() {
        let mut cache = KvCache::new(2, 2);
        let _ = cache.push(&[1.0, 2.0], &[1.0, 2.0]); // too short
    }

    #[test]
    fn clear_resets_state() {
        let mut cache = KvCache::new(1, 1);
        cache.push(&[1.0], &[2.0]).unwrap();
        cache.clear();
        assert_eq!(cache.seq_len, 0);
        assert!(cache.k.is_empty());
    }

    #[test]
    fn truncate_rolls_back_to_exact_length_preserving_earlier_data() {
        let mut cache = KvCache::new(2, 2);
        cache
            .push(&[1.0, 2.0, 3.0, 4.0], &[10.0, 20.0, 30.0, 40.0])
            .unwrap();
        cache
            .push(&[5.0, 6.0, 7.0, 8.0], &[50.0, 60.0, 70.0, 80.0])
            .unwrap();
        cache
            .push(&[9.0, 9.0, 9.0, 9.0], &[90.0, 90.0, 90.0, 90.0])
            .unwrap();
        assert_eq!(cache.seq_len, 3);

        cache.truncate(1);
        assert_eq!(cache.seq_len, 1);
        assert_eq!(cache.k, vec![1.0, 2.0, 3.0, 4.0]);
        assert_eq!(cache.v, vec![10.0, 20.0, 30.0, 40.0]);
    }

    #[test]
    fn truncate_to_current_length_is_a_no_op() {
        let mut cache = KvCache::new(1, 2);
        cache.push(&[1.0, 2.0], &[3.0, 4.0]).unwrap();
        cache.truncate(1);
        assert_eq!(cache.seq_len, 1);
        assert_eq!(cache.k, vec![1.0, 2.0]);
    }

    #[test]
    fn truncate_to_zero_empties_the_cache() {
        let mut cache = KvCache::new(1, 2);
        cache.push(&[1.0, 2.0], &[3.0, 4.0]).unwrap();
        cache.truncate(0);
        assert_eq!(cache.seq_len, 0);
        assert!(cache.k.is_empty());
        assert!(cache.v.is_empty());
    }

    #[test]
    #[should_panic]
    fn truncate_beyond_current_length_panics() {
        let mut cache = KvCache::new(1, 2);
        cache.push(&[1.0, 2.0], &[3.0, 4.0]).unwrap();
        cache.truncate(5);
    }

    #[test]
    fn push_after_truncate_continues_correctly() {
        let mut cache = KvCache::new(1, 1);
        cache.push(&[1.0], &[10.0]).unwrap();
        cache.push(&[2.0], &[20.0]).unwrap();
        cache.push(&[3.0], &[30.0]).unwrap(); // this one will be "rejected"
        cache.truncate(2);
        cache.push(&[99.0], &[990.0]).unwrap(); // real continuation after rejection
        assert_eq!(cache.seq_len, 3);
        assert_eq!(cache.k, vec![1.0, 2.0, 99.0]);
        assert_eq!(cache.v, vec![10.0, 20.0, 990.0]);
    }

    #[test]
    fn with_capacity_preallocates_and_never_reallocates_within_plan() {
        let n_kv_heads = 4;
        let head_dim = 8;
        let max_seq_len = 16;
        let mut cache = KvCache::with_capacity(n_kv_heads, head_dim, max_seq_len);

        let expected_elems = max_seq_len * n_kv_heads * head_dim;
        assert!(cache.k.capacity() >= expected_elems);
        assert!(cache.v.capacity() >= expected_elems);

        let step = vec![0.5f32; n_kv_heads * head_dim];
        let k_ptr_before = cache.k.as_ptr();
        for _ in 0..max_seq_len {
            cache.push(&step, &step).unwrap();
        }
        let k_ptr_after = cache.k.as_ptr();
        assert_eq!(
            k_ptr_before, k_ptr_after,
            "pushing exactly up to the planned capacity must not reallocate"
        );
        assert!(cache.is_within_planned_capacity());
    }

    #[test]
    fn allocated_bytes_reflects_preallocated_capacity_not_just_used_length() {
        let cache = KvCache::with_capacity(4, 8, 100);
        // 100 positions * 4 kv_heads * 8 head_dim * 2 (k+v) * 4 bytes/f32
        let expected_min = 100 * 4 * 8 * 2 * 4;
        assert!(
            cache.allocated_bytes() >= expected_min,
            "allocated_bytes={} expected_min={expected_min}",
            cache.allocated_bytes()
        );
        // Nothing has been pushed yet, but the memory is already reserved.
        assert_eq!(cache.seq_len, 0);
    }

    #[test]
    fn grow_as_you_go_cache_reports_not_within_planned_capacity() {
        let mut cache = KvCache::new(2, 2);
        cache
            .push(&[1.0, 2.0, 3.0, 4.0], &[5.0, 6.0, 7.0, 8.0])
            .unwrap();
        assert!(
            !cache.is_within_planned_capacity(),
            "a cache built with `new` has no plan to be within"
        );
    }

    #[test]
    fn with_pool_acquires_one_block_and_reports_it_in_free_blocks() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 10)));
        let cache = KvCache::with_pool(2, 2, pool.clone(), 0).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 9);
        assert_eq!(cache.seq_len, 0);
    }

    #[test]
    fn with_pool_fails_without_mutating_the_pool_when_exhausted() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 0)));
        let result = KvCache::with_pool(2, 2, pool.clone(), 0);
        assert!(result.is_err());
        assert_eq!(pool.lock().unwrap().free_blocks(), 0);
    }

    #[test]
    fn push_acquires_additional_blocks_as_the_cache_crosses_block_boundaries() {
        let block_size = 2;
        let pool = Arc::new(Mutex::new(KvBlockPool::new(block_size, 10)));
        let mut cache = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 9);

        // First block holds `block_size` = 2 positions; pushing them
        // must not need a second block.
        cache.push(&[1.0], &[1.0]).unwrap();
        cache.push(&[2.0], &[2.0]).unwrap();
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            9,
            "filling exactly the first block must not acquire a second one"
        );

        // The third position crosses into a second block.
        cache.push(&[3.0], &[3.0]).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 8);
        assert_eq!(cache.seq_len, 3);
        assert_eq!(cache.k, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn push_returns_pool_exhausted_and_leaves_state_unchanged_when_no_blocks_remain() {
        let block_size = 1;
        let pool = Arc::new(Mutex::new(KvBlockPool::new(block_size, 1)));
        let mut cache = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 0);

        cache.push(&[1.0], &[1.0]).unwrap(); // fills the one held block

        let before_k = cache.k.clone();
        let result = cache.push(&[2.0], &[2.0]);
        assert_eq!(result, Err(KvPoolExhausted));
        assert_eq!(cache.seq_len, 1, "a failed push must not change seq_len");
        assert_eq!(cache.k, before_k, "a failed push must not append data");
    }

    #[test]
    fn dropping_a_pooled_cache_returns_its_blocks_to_the_pool() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(1, 2)));
        {
            let mut cache = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
            cache.push(&[1.0], &[1.0]).unwrap(); // fills the first (only held) block
            cache.push(&[2.0], &[2.0]).unwrap(); // crosses into a second block
            assert_eq!(pool.lock().unwrap().free_blocks(), 0);
        }
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            2,
            "both blocks held by the dropped cache must return to the pool"
        );
    }

    #[test]
    fn release_to_pool_is_explicit_and_idempotent() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 5)));
        let mut cache = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 4);

        cache.release_to_pool();
        assert_eq!(pool.lock().unwrap().free_blocks(), 5);

        cache.release_to_pool(); // no-op, must not over-release
        assert_eq!(pool.lock().unwrap().free_blocks(), 5);

        drop(cache); // must not release again either
        assert_eq!(pool.lock().unwrap().free_blocks(), 5);
    }

    #[test]
    fn two_pooled_caches_share_one_bounded_budget() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(1, 1)));
        let cache_a = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
        let cache_b = KvCache::with_pool(1, 1, pool.clone(), 0);
        assert!(
            cache_b.is_err(),
            "a second concurrent request must not be admitted when the shared budget is full"
        );

        drop(cache_a);
        let cache_c = KvCache::with_pool(1, 1, pool, 0);
        assert!(
            cache_c.is_ok(),
            "once the first request's cache is dropped, its budget must become available again"
        );
    }

    #[test]
    fn cloning_a_pooled_cache_detaches_the_clone_from_pool_accounting() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 3)));
        let original = KvCache::with_pool(1, 1, pool.clone(), 0).unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 2);

        let clone = original.clone();
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            2,
            "cloning must not acquire additional blocks"
        );
        assert_eq!(clone.k, original.k);

        drop(clone);
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            2,
            "dropping a detached clone must not release the original's blocks"
        );

        drop(original);
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            3,
            "dropping the original must release its blocks exactly once"
        );
    }

    /// A resize is arithmetic, and the one rule that is not: shrinking
    /// past what is held is refused, and the pool is left exactly as it
    /// was.
    ///
    /// Clamping to zero instead would silently over-promise -- the
    /// caches holding those blocks do not give them back, so every
    /// later acquire would decide against a budget that does not
    /// describe the memory in use. This test fails under that clamp.
    #[test]
    fn a_pool_refuses_to_shrink_below_what_is_already_held() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 10)));
        let held = KvCache::with_pool(2, 4, Arc::clone(&pool), 24).expect("blocks");
        let in_use = {
            let p = pool.lock().unwrap();
            p.total_blocks() - p.free_blocks()
        };
        assert!(in_use > 0, "the fixture must actually hold blocks");

        let mut p = pool.lock().unwrap();
        assert_eq!(p.resize(in_use - 1), Err(in_use));
        assert_eq!(p.total_blocks(), 10, "a refused resize changes nothing");
        assert_eq!(p.free_blocks(), 10 - in_use);

        // Down to exactly what is held is legal, and leaves nothing free.
        assert_eq!(p.resize(in_use), Ok(()));
        assert_eq!(p.free_blocks(), 0);
        drop(p);
        drop(held);
    }

    /// Growing hands the new blocks to the free list without disturbing
    /// what is held, which is the whole point of a live re-split.
    #[test]
    fn growing_a_pool_adds_to_what_is_free_and_not_to_what_is_held() {
        let pool = Arc::new(Mutex::new(KvBlockPool::new(4, 8)));
        let held = KvCache::with_pool(2, 4, Arc::clone(&pool), 16).expect("blocks");
        let mut p = pool.lock().unwrap();
        let in_use = p.total_blocks() - p.free_blocks();

        assert_eq!(p.resize(32), Ok(()));
        assert_eq!(p.total_blocks(), 32);
        assert_eq!(p.free_blocks(), 32 - in_use);
        drop(p);
        drop(held);
    }
}