Trait KVCache

Source

pub trait KVCache<B: Backend>: Send {
    // Required methods
    fn attention_opts(
        &mut self,
        layer: usize,
        q: Tensor<B, 4>,
        k: Tensor<B, 4>,
        v: Tensor<B, 4>,
        pos: usize,
        scale: f64,
        window: Option<usize>,
    ) -> Tensor<B, 4>;
    fn seq_len(&self) -> usize;
    fn reset(&mut self);

    // Provided methods
    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> { ... }
    fn popn(&mut self, n: usize) -> usize { ... }
    fn pages_used(&self) -> Option<usize> { ... }
    fn page_stats(&self) -> Option<PageStats> { ... }
}

Expand description

Per-layer key/value storage that owns the attention computation.

Tensors are 4-D [batch=1, heads, seq, head_dim]; q has n_q heads while k/v have n_kv heads (GQA expansion happens inside the implementation, as does causal masking).

Required Methods§

Source

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

KVCache::attention with an optional sliding-window span (Gemma local layers): when Some(w), query at absolute position p attends only keys in (p - w, p] — older keys stay cached but are masked out. None = full causal attention (Llama-family behavior).

Source

fn seq_len(&self) -> usize

Total cached sequence length.

Source

fn reset(&mut self)

Drops all cached state (session reset).

Provided Methods§

Source

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).

pos is the absolute position of the first new token and must equal KVCache::seq_len on entry (dense contiguous appends). scale is the attention logit scale (1/sqrt(head_dim)). Returns the attention output [1, n_q, seq, head_dim].

Source

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

Rolls back the last n cached tokens, returning how many were actually dropped. Caches that cannot roll back (the contiguous baseline) return 0 — callers gate prefix reuse on a nonzero result.

Source