Struct PagedKVCache 

pub struct PagedKVCache<B: Backend> { /* private fields */ }

MLC-style paged KV cache.

Per layer, K and V live in fixed arena tensors of shape [num_pages, n_kv, page_size, head_dim], allocated lazily on the layer’s first use. A single-sequence page table maps logical pages to physical page ids drawn from a free-page allocator (the struct is shaped so per-sequence tables can be added later).

attention() writes the new K/V into page slots (one slice_assign per touched page), gathers the active pages into a contiguous [1, n_kv, total, head_dim] window and runs the standard matmul path. Steady-state decode therefore writes a single slot and gathers — the Phase 1 O(seq) cat-rewrite per token is gone. (A fused no-gather CubeCL kernel is a later task.)

Implementations

impl<B: Backend> PagedKVCache<B>

pub fn new(num_layers: usize, config: CacheConfig) -> Self

Creates an empty paged cache for num_layers layers. Arena tensors are allocated lazily on first use of each layer.

pub fn num_free_pages(&self) -> usize

Number of free pages in the arena.

Trait Implementations

impl<B: Backend> KVCache<B> for PagedKVCache<B>

fn popn(&mut self, n: usize) -> usize

Rolls back the last n cached tokens, freeing trailing pages that become fully unused. K/V content of popped positions is left in the arena but is never read (writes always cover seq_len.. densely).

fn attention( &mut self, layer: usize, q: Tensor<B, 4>, k: Tensor<B, 4>, v: Tensor<B, 4>, pos: usize, scale: f64, ) -> Tensor<B, 4>

Appends seq new positions of K/V for layer and computes attention of q against the full cached window (past + new).

fn seq_len(&self) -> usize

Total cached sequence length.

fn reset(&mut self)

Drops all cached state (session reset).

fn pages_used(&self) -> Option<usize>

Pages currently allocated to the sequence (paged cache only).