Struct PagedKvCache 

pub struct PagedKvCache { /* private fields */ }

Paged KV cache.

Memory is allocated in fixed-size pages (blocks) of PAGE_SIZE tokens. Pages grow on demand — short sequences don’t waste memory for unused context positions. The assemble_* methods reconstruct contiguous slices for attention computation.

Implementations

impl PagedKvCache

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

Create a new paged KV cache.

Unlike the contiguous cache, this does NOT pre-allocate all memory. Pages are allocated on demand as tokens are processed.

pub fn page_size(&self) -> usize

Returns the page size (tokens per page).

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 total_pages(&self) -> usize

Returns total number of allocated pages across all layers.

pub fn memory_bytes(&self) -> usize

Returns total memory usage in bytes (approximate).

pub fn clear(&mut self)

Reset the cache, freeing all pages.

pub fn shrink_to_fit(&mut self)

Shrink allocated pages to fit the current sequence length. Useful after trimming context.

impl PagedKvCache

Extended paged KV cache operations (not part of the base trait).

pub fn get_keys_into(&self, layer: usize, buf: &mut Vec<f32>) -> ArchResult<()>

Copy all cached keys for a layer into a contiguous buffer.

This is the multi-page alternative to get_keys(). The caller provides a reusable buffer to avoid repeated allocation.

pub fn get_values_into( &self, layer: usize, buf: &mut Vec<f32>, ) -> ArchResult<()>

Copy all cached values for a layer into a contiguous buffer.

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

Read a specific token’s key data from the cache.

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

Read a specific token’s value data from the cache.

pub fn iter_keys<F>(&self, layer: usize, f: F) -> ArchResult<()>

where F: FnMut(usize, &[f32]),

Iterate over key tokens for a layer, calling f for each (pos, key_data).

pub fn iter_values<F>(&self, layer: usize, f: F) -> ArchResult<()>

where F: FnMut(usize, &[f32]),

Iterate over value tokens for a layer.

Trait Implementations

impl KvCacheAccess for PagedKvCache

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]), ) -> ArchResult<()>

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]), ) -> ArchResult<()>

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