Struct Cache 

pub struct Cache { /* private fields */ }

Abstraction over cosine and sine tables, kv-caching and attention masking.

Implementations

impl Cache

pub fn new( use_kv_cache: bool, dtype: DType, config: &Config, device: &Device, ) -> Result<Self>

Creates a new cache instance with the provided configuration. Set use_kv_cache to false to disable kv-caching.

pub fn with_kv_cache(&self) -> bool

Return true if kv-caching is enabled.

pub fn cosine( &self, index_pos: usize, seq_len: usize, device: &Device, ) -> Result<Tensor>

Return the cached cosine value for the given position and sequence length. When device differs from the cache’s own device, the result is copied to that device (enables multi-GPU workers).

pub fn sine( &self, index_pos: usize, seq_len: usize, device: &Device, ) -> Result<Tensor>

Return the cached sine value for the given position and sequence length.

pub fn mask(&mut self, seq_len: usize, device: &Device) -> Result<Tensor>

Get the attention mask for the given sequence length.

pub fn process_kv( &mut self, block_idx: usize, k: Tensor, v: Tensor, ) -> Result<(Tensor, Tensor)>

Process the input k and v by either generating their cache entry or applying a previously cached one.

pub fn get_recurrent_state(&self, block_idx: usize) -> Option<&Tensor>

Get the recurrent state for a linear attention layer.

pub fn set_recurrent_state(&mut self, block_idx: usize, state: Tensor)

Set the recurrent state for a linear attention layer.

pub fn get_conv_state(&self, block_idx: usize) -> Option<&Tensor>

Get the conv state for a linear attention layer.

pub fn set_conv_state(&mut self, block_idx: usize, state: Tensor)

Set the conv state for a linear attention layer.

pub fn as_new(&self) -> Self

Return a copy of this cache with the same state but new kv table.

pub fn clear(&mut self)

Clear the cache.

Trait Implementations

impl Clone for Cache

fn clone(&self) -> Cache

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Cache

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

Formats the value using the given formatter.