Struct KvCache 

pub struct KvCache { /* private fields */ }

Simple contiguous KV cache implementation.

Stores key and value tensors for all layers in contiguous FP32 buffers. Each layer has a separate key buffer and value buffer, sized for the maximum context length.

Implementations

impl KvCache

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

Allocate a new KV cache.

Arguments

• num_layers - Number of transformer layers.
• max_seq_len - Maximum context length.
• kv_dim - KV dimension per token (num_kv_heads * head_dim).

pub fn clear(&mut self)

Reset the cache, clearing all stored KV pairs.

pub fn max_seq_len(&self) -> usize

Returns the maximum sequence length.

pub fn kv_dim(&self) -> usize

Returns the KV dimension per token.

pub fn num_layers(&self) -> usize

Returns the number of layers.

pub fn advance(&mut self)

Advance the sequence position by one token.

pub fn restore_from_snapshot( &mut self, keys: &[Vec<f32>], values: &[Vec<f32>], seq_len: usize, )

Restore from a prefix cache snapshot.

Copies the provided per-layer key/value data into internal buffers and sets seq_len to the snapshot’s length. The caller must ensure that keys.len() == values.len() == num_layers and that each inner vec has seq_len * kv_dim elements.

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

Truncate the KV cache to n tokens.

After this call seq_len() returns n (clamped to the current seq_len if n is already beyond it — truncate never extends the cache). The underlying buffers are not zeroed; the truncated region is simply considered invalid and will be overwritten on the next store_kv call.

This is the low-level primitive for speculative-decoding rollback: the target engine calls truncate(divergence_pos) after rejecting a draft token, then continues generating from divergence_pos.

pub fn snapshot(&self) -> KvCacheSnapshot

Capture a snapshot of the current KV state.

Only the data up to seq_len * kv_dim is copied per layer, keeping the snapshot compact.

pub fn to_payload(&self) -> KvStatePayload

Build a serializable crate::snapshot::KvStatePayload from the current state.

pub fn restore_from_payload( &mut self, payload: &KvStatePayload, ) -> RuntimeResult<()>

Restore cache state from a crate::snapshot::KvStatePayload.

Validates that layer count and dimensions match the cache configuration, then restores the key/value buffers and sequence length.

Trait Implementations

impl KvCacheAccess for KvCache

fn seq_len(&self) -> usize

Get the current sequence length (number of cached tokens).

fn store_kv( &mut self, layer: usize, key: &[f32], value: &[f32], ) -> ArchResult<()>

Store key and value tensors for a layer at the current position.

fn get_keys(&self, layer: usize) -> ArchResult<&[f32]>

Retrieve all cached keys for a layer up to the current sequence length.

fn get_values(&self, layer: usize) -> ArchResult<&[f32]>

Retrieve all cached values for a layer up to the current sequence length.

fn advance(&mut self)

Advance the cache position by one token.

fn kv_dim(&self) -> usize

KV dimension per token (num_kv_heads * head_dim).

fn for_each_key( &self, layer: usize, f: &mut dyn FnMut(usize, &[f32]), ) -> Result<(), ArchError>

Iterate over every cached key token for layer, calling f(pos, key_data).

fn for_each_value( &self, layer: usize, f: &mut dyn FnMut(usize, &[f32]), ) -> Result<(), ArchError>

Iterate over every cached value token for layer, calling f(pos, value_data).