Struct MultiHeadKvCache 

pub struct MultiHeadKvCache { /* private fields */ }

Multi-head KV cache for transformer attention.

Maintains one KvCacheCompressor per attention head.

Implementations

impl MultiHeadKvCache

pub fn new(num_heads: usize, config: &KvCacheConfig) -> Result<Self>

Create a new multi-head KV cache.

Arguments

• num_heads — number of attention heads
• config — shared configuration applied to all heads

Errors

• Any error from KvCacheCompressor::new

pub fn push_token(&mut self, keys: &[&[f32]], values: &[&[f32]]) -> Result<()>

Push a token’s KV vectors for all heads simultaneously.

keys and values must each have exactly num_heads slices of length head_dim.

Errors

• DimensionMismatch, NonFiniteInput

pub fn attention_scores(&self, queries: &[&[f32]]) -> Result<Vec<Vec<f32>>>

Compute attention scores for each head given per-head query vectors.

queries must have exactly num_heads slices of length head_dim.

Returns one score vector per head.

Errors

• DimensionMismatch, NonFiniteInput

pub fn len(&self) -> usize

Number of tokens stored (same for all heads).

pub fn is_empty(&self) -> bool

Returns true if no tokens have been stored.

pub fn clear(&mut self)

Clear all heads.

Trait Implementations

impl Clone for MultiHeadKvCache

fn clone(&self) -> MultiHeadKvCache

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for MultiHeadKvCache

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

Formats the value using the given formatter.