Struct PagedKvCache 

pub struct PagedKvCache { /* private fields */ }

2026-09-25: Paged KV cache across all attention layers: the pools and the block free list with per-block reference counts.

Implementations

impl PagedKvCache

pub fn debug_kv_checksum_per_layer( &self, blocks: &[u32], boundary_idx: usize, gpu: &dyn GpuBackend, stream: u64, tag: &str, )

2026-09-25: Diagnostic. Synchronises stream, then logs a K and V (sum, ssq, sabs) line per BF16 layer for each of two regions of blocks: [0, boundary_idx) (“prefix”) and the rest (“suffix”). Per-layer lines keep a local divergence from cancelling in a global sum. Non-BF16 layers are skipped, with a warning when layer 0 is one; a block whose copy fails is skipped. The model engine calls it only under METRALE_SSM_SAVE_DUMP.

pub fn debug_kv_per_block( &self, layer_idx: usize, blocks: &[u32], gpu: &dyn GpuBackend, stream: u64, tag: &str, )

2026-09-25: Diagnostic. Synchronises stream, then logs one line per entry of blocks, in order: logical index, physical block, K and V sum of squares for layer layer_idx. Unlike a region sum this shows a wrong block-to-position mapping. Does nothing for a non-BF16 layer.

impl PagedKvCache

pub fn new( config: KvCacheConfig, num_blocks: usize, gpu: &dyn GpuBackend, ) -> Result<Self>

2026-09-25: Allocate a K pool and a V pool of num_blocks blocks for every attention layer. All blocks start free with count 0.

pub fn alloc_block(&mut self) -> Result<u32>

2026-09-25: Pop a free block and set its count to 1. Errors when the free list is empty.

pub fn zero_block( &self, block_idx: u32, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: Zero the block’s K and V bytes in every layer, enqueued on stream.

pub fn poison_block( &self, block_idx: u32, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: Diagnostic twin of zero_block that fills with 0xFF. 0xFFFF is a BF16 NaN and 0xFF an FP8 E4M3 NaN, so attention over a slot nothing wrote yields NaN instead of zeros. The model engine calls it instead of zero_block under METRALE_KV_POISON.

pub fn try_alloc_block(&mut self) -> Option<u32>

2026-09-25: alloc_block, returning None when the free list is empty.

pub fn inc_ref(&mut self, block_idx: u32)

2026-09-25: Add one reference to a block. There is no check that the block is allocated.

pub fn dec_ref(&mut self, block_idx: u32) -> bool

2026-09-25: Drop one reference. Returns true if the count reached 0 and the block went back on the free list. At 0 refs a debug build panics; a release build logs an error and returns false.

pub fn free_block(&mut self, block_idx: u32)

2026-09-25: dec_ref, discarding the result.

pub fn free_blocks(&mut self, block_table: &[u32])

2026-09-25: free_block on every entry of block_table.

pub fn return_evicted_block(&mut self, block_idx: u32)

2026-09-25: Release the prefix cache’s own reference on a block its eviction returned: one decrement, and a push onto the free list only on a 1 to 0 transition. At 0 refs it logs a warning and does nothing.

pub fn ref_count(&self, block_idx: u32) -> u32

2026-09-25: The block’s reference count.

pub fn num_free_blocks(&self) -> usize

2026-09-25: Length of the free list.

pub fn k_cache_ptr(&self, layer_idx: usize, block_idx: u32) -> DevicePtr

2026-09-25: Device address of a block in a layer’s K pool.

pub fn v_cache_ptr(&self, layer_idx: usize, block_idx: u32) -> DevicePtr

2026-09-25: Device address of a block in a layer’s V pool.

pub fn k_pool_ptr(&self, layer_idx: usize) -> DevicePtr

2026-09-25: Base address of a layer’s K pool.

pub fn v_pool_ptr(&self, layer_idx: usize) -> DevicePtr

2026-09-25: Base address of a layer’s V pool.

pub fn cache_stride(&self) -> usize

2026-09-25: KvCacheConfig::cache_stride_elements: elements per block and side from the global dims, not from layer_dims.

pub fn block_stride_bytes(&self) -> usize

2026-09-25: KvCacheConfig::block_bytes: uniform dtype, global dims.

pub fn block_stride_bytes_for_layer(&self, layer_idx: usize) -> usize

2026-09-25: The K-side block stride of attention layer layer_idx, which is also the V-side stride for a symmetric format.

pub fn k_block_stride_bytes_for_layer(&self, layer_idx: usize) -> usize

2026-09-25: The K-side block stride of attention layer layer_idx.

pub fn v_block_stride_bytes_for_layer(&self, layer_idx: usize) -> usize

2026-09-25: The V-side block stride of attention layer layer_idx.

pub fn nvfp4_data_bytes(&self) -> usize

2026-09-25: KvCacheConfig::nvfp4_data_bytes.

pub fn turbo4_data_bytes(&self) -> usize

2026-09-25: KvCacheConfig::turbo4_data_bytes.

pub fn turbo3_data_bytes(&self) -> usize

2026-09-25: KvCacheConfig::turbo3_data_bytes.

pub fn turbo2_data_bytes(&self) -> usize

2026-09-25: KvCacheConfig::turbo2_data_bytes.

pub fn turbo8_data_bytes(&self) -> usize

2026-09-25: KvCacheConfig::turbo8_data_bytes.

pub fn turbo4_scale_bytes(&self) -> usize

2026-09-25: KvCacheConfig::turbo4_scale_bytes.

pub fn config(&self) -> &KvCacheConfig

2026-09-25: The configuration the cache was built with.

pub fn dtype_for_layer(&self, layer_idx: usize) -> KvCacheDtype

2026-09-25: The format of attention layer layer_idx.

pub fn block_size(&self) -> usize

pub fn num_blocks(&self) -> usize

pub fn dtype(&self) -> KvCacheDtype

pub fn num_layers(&self) -> usize

2026-09-25: KvCacheConfig::num_layers.

pub fn read_block( &self, layer_idx: usize, block_idx: u32, gpu: &dyn GpuBackend, ) -> Result<(Vec<u8>, Vec<u8>)>

2026-09-25: Copy one block of one layer to the host, as (k, v), each the length of its side’s block stride.

pub fn write_block( &self, layer_idx: usize, block_idx: u32, k_data: &[u8], v_data: &[u8], gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Copy k_data and v_data from the host to the start of one block of one layer.

pub fn compute_num_blocks( config: &KvCacheConfig, available_bytes: usize, ) -> Result<usize>

2026-09-25: How many block indices fit in available_bytes, at block_bytes_kv_all_layers each. Errors if that is zero.

Trait Implementations

impl ModelResource<dyn GpuBackend> for PagedKvCache

2026-09-25: Release both pools of every layer, and clear the free list and ref counts with them, so a released cache cannot hand out a block of freed memory: alloc_block then fails. Every pool is freed even after an error; the first error is returned.

fn label(&self) -> &'static str

2026-09-25: Human name, used to attribute a failure in Teardown::release_all’s error.

fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()>

2026-09-25: Release everything this owns. Must be idempotent.