Struct QuantizedKVCache 

pub struct QuantizedKVCache {
    pub k_data: Vec<Vec<u8>>,
    pub v_data: Vec<Vec<u8>>,
    pub k_scales: Vec<Vec<f32>>,
    pub v_scales: Vec<Vec<f32>>,
    pub format: KVCacheFormat,
    pub seq_len: usize,
    pub max_seq_len: usize,
    pub num_kv_heads: usize,
    pub head_dim: usize,
    pub num_layers: usize,
}

Quantized KV cache that stores K/V in reduced precision

Fields

k_data: Vec<Vec<u8>>

Quantized key data per layer - raw bytes in the chosen format

v_data: Vec<Vec<u8>>

Quantized value data per layer

k_scales: Vec<Vec<f32>>

Per-head scale factors for INT8 (one per head per position per layer) Layout: [layer][head * max_seq_len + pos]

v_scales: Vec<Vec<f32>>

format: KVCacheFormat

Storage format

seq_len: usize

Current sequence length

max_seq_len: usize

num_kv_heads: usize

head_dim: usize

num_layers: usize

Implementations

impl QuantizedKVCache

pub fn new( num_layers: usize, num_kv_heads: usize, max_seq_len: usize, head_dim: usize, format: KVCacheFormat, ) -> Self

Create a new quantized KV cache with the given dimensions and format

pub fn reset(&mut self)

Reset the cache for a new sequence

pub fn remaining_capacity(&self) -> usize

Get remaining capacity

pub fn is_full(&self) -> bool

Check if cache is full

pub fn memory_usage(&self) -> usize

Get memory usage in bytes

pub fn write_kv( &mut self, layer: usize, pos: usize, k_data: &[f32], v_data: &[f32], )

Write quantized K/V for one position

k_data and v_data are [num_kv_heads * head_dim] each

pub fn read_k(&self, layer: usize, head: usize, pos: usize) -> Vec<f32>

Dequantize and return key for one head at one position

pub fn read_v(&self, layer: usize, head: usize, pos: usize) -> Vec<f32>

Dequantize and return value for one head at one position

pub fn read_k_range( &self, layer: usize, head: usize, start_pos: usize, end_pos: usize, ) -> Vec<f32>

Dequantize key range for one head

Returns [end_pos - start_pos, head_dim] as flat vec

pub fn read_v_range( &self, layer: usize, head: usize, start_pos: usize, end_pos: usize, ) -> Vec<f32>

Dequantize value range for one head

Returns [end_pos - start_pos, head_dim] as flat vec

pub fn shift_left(&mut self, amount: usize)

Shift cache left by amount positions (for sliding window)

pub fn truncate(&mut self, new_len: usize)

Truncate cache to a specific length