Struct KVCache 

pub struct KVCache { /* private fields */ }

KV-cache for efficient autoregressive generation.

Stores the key and value tensors from previous positions so they don’t need to be recomputed at each generation step. Each layer has its own cache entry.

Shapes

• keys[i]: [batch, num_kv_heads, seq_len, head_dim]
• values[i]: [batch, num_kv_heads, seq_len, head_dim]

Implementations

impl KVCache

pub fn new(n_layers: usize) -> Self

Create a new empty cache for the given number of layers.

pub fn seq_len(&self) -> Result<usize>

Current sequence length from the cache (0 if empty).

Errors

Returns MIError::Model if a cached tensor has an unexpected shape.

pub fn is_empty(&self) -> bool

Whether the cache is empty (no layers have been populated).

pub const fn n_layers(&self) -> usize

Number of layers in the cache.

pub fn clear(&mut self)

Clear all cached tensors.

pub fn layer_mut( &mut self, layer: usize, ) -> Result<(&mut Option<Tensor>, &mut Option<Tensor>)>

Get mutable references to the cache entry for a specific layer.

Returns (&mut Option<Tensor>, &mut Option<Tensor>) for (key, value).

Errors

Returns MIError::Hook if layer is out of range.

pub fn memory_usage(&self) -> usize

Estimate memory usage in bytes.

Returns the total memory used by all cached tensors.

pub fn trim_to(&mut self, max_seq_len: usize) -> Result<bool>

Trim the cache to keep only the last max_seq_len tokens.

Useful for memory-constrained scenarios with long sequences. Returns Ok(true) if trimming occurred, Ok(false) if no trimming was needed.

Errors

Returns MIError::Model if tensor operations fail.

pub fn enforce_memory_limit(&mut self, max_bytes: usize) -> Result<bool>

Check if cache exceeds memory limit and trim if needed.

Trims to ~75% of current length if memory limit is exceeded. Returns Ok(true) if trimming occurred.

Errors

Returns MIError::Model if tensor operations fail.

Trait Implementations

impl Clone for KVCache

fn clone(&self) -> KVCache

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for KVCache

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

Formats the value using the given formatter.

impl Default for KVCache

fn default() -> Self

Returns the “default value” for a type.