Struct KvCache 

pub struct KvCache { /* private fields */ }

Per-layer KV cache storing FP32 key and value vectors.

Implementations

impl KvCache

pub fn new( num_layers: usize, num_kv_heads: usize, head_dim: usize, max_seq_len: usize, ) -> Self

Create a new KV cache.

pub fn seq_len(&self) -> usize

Current number of cached tokens.

pub fn max_seq_len(&self) -> usize

Maximum sequence length.

pub fn store_key(&mut self, layer: usize, head: usize, pos: usize, key: &[f32])

Store a key vector for a specific layer, head, and position.

pub fn store_value( &mut self, layer: usize, head: usize, pos: usize, value: &[f32], )

Store a value vector for a specific layer, head, and position.

pub fn keys_for(&self, layer: usize, head: usize, seq_len: usize) -> &[f32]

Get all cached keys for a layer and head up to seq_len.

Returns a slice of [seq_len × head_dim] in row-major order.

pub fn values_for(&self, layer: usize, head: usize, seq_len: usize) -> &[f32]

Get all cached values for a layer and head up to seq_len.

pub fn advance(&mut self)

Advance the sequence position by one token.

pub fn clear(&mut self)

Reset the cache (clear all stored KV pairs).

pub fn memory_bytes(&self) -> usize

Total memory used by this cache in bytes.

pub fn utilization_ratio(&self) -> f64

Utilization ratio: fraction of cache capacity currently used.

Returns a value in [0.0, 1.0].

pub fn num_layers(&self) -> usize

Number of layers in this cache.

pub fn num_kv_heads(&self) -> usize

Number of KV heads per layer.

pub fn head_dim(&self) -> usize

Head dimension.

pub fn set_seq_len(&mut self, n: usize)

Manually set the cached sequence length.

Used by the prefix-cache integration when restoring previously computed KV blocks: after inject_block writes the block contents, the consumer must call this to advertise the number of valid positions to subsequent attention computations.

n is clamped to max_seq_len.

pub fn extract_block( &self, layer: usize, start_pos: usize, block_size: usize, ) -> (Vec<f32>, Vec<f32>)

Extract one prefix-cache block worth of KV for a single layer.

Reads block_size consecutive positions starting at start_pos for every KV head in layer and returns them in [head][pos_in_block][dim] order, packed as a flat Vec<f32> of length num_kv_heads * block_size * head_dim.

Mirrors the layout used by crate::prefix_cache::CacheBlock.

Returns (keys, values). If the requested range exceeds max_seq_len, the trailing positions are returned as zeros.

pub fn inject_block( &mut self, layer: usize, start_pos: usize, block_size: usize, keys: &[f32], values: &[f32], )

Inject a previously extracted block back into the cache for a single layer.

keys and values must have the same [head][pos_in_block][dim] layout produced by extract_block; they are expected to be of length num_kv_heads * block_size * head_dim. Positions outside max_seq_len are silently skipped.

Trait Implementations

impl Debug for KvCache

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.