Struct PagedKvCache 

pub struct PagedKvCache { /* private fields */ }

Paged KV cache across all attention layers.

Implementations

impl PagedKvCache

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

Allocate the KV cache pool on the GPU.

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

Allocate a free block. Returns block index.

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

Zero all KV data in a block across all layers. Prevents stale KV data from previous sequences from leaking into new sequences via paged attention reads beyond the current seq_len.

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

DIAGNOSTIC (ATLAS_KV_POISON): fill a freshly-allocated block with 0xFF (a NaN bit-pattern in both bf16 0xFFFF and fp8-e4m3 0xFF) instead of zero. Any KV region that decode/attention reads but prefill never wrote then yields deterministic NaN rather than plausible-but-wrong zeros. Used to falsify the “unwritten fresh tail block” hypothesis: if cache-ON output goes NaN under poison while cache-OFF stays clean, a fresh block is being read unwritten; if both stay clean, fresh KV is fully written and the run-to-run nondeterminism originates elsewhere (scratch/scan).

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

Try to allocate a free block without failing. Returns None if exhausted.

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

Increment reference count on a block (for prefix cache sharing).

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

Decrement reference count. Returns true if block was freed (count hit 0).

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

Free a previously allocated block (decrements ref, frees if count hits 0).

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

Free all blocks in a block table.

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

Return a block to the free pool directly, bypassing ref counting. Used by eviction: the radix tree already removed its reference.

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

Current reference count for a block.

pub fn num_free_blocks(&self) -> usize

Number of free blocks.

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

Get K cache pointer for a layer and block.

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

Get V cache pointer for a layer and block.

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

DEBUG (env-gated): PER-LAYER K and V fingerprint over blocks, emitting (sum, ssq, sabs) for each attention layer so a localized divergence can’t cancel in a global sum. Splits the block list at boundary_idx: blocks [0, boundary_idx) are the REUSED-PREFIX region (carried over from a prior turn’s prefill) and [boundary_idx, end) are the RECOMPUTED-SUFFIX region. Each region gets its own per-layer line so we can localize the FIRST layer/region where chained (ON) differs from cold (OFF). Only valid for BF16 KV (the experiment uses --kv-cache-dtype bf16); non-BF16 layers are skipped with a one-shot warning.

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

DEBUG (env-gated): per-LOGICAL-BLOCK K/V fingerprint for ONE layer, walking blocks in block_table order. Emits (logical_idx, physical_block, k_ssq, v_ssq) per block so a per-position aliasing / reordering bug (identical region SUM but wrong block→position mapping) is visible. BF16 only.

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

Get the full K cache pool pointer for a layer (for paged decode kernel).

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

Get the full V cache pool pointer for a layer.

pub fn cache_stride(&self) -> usize

Cache stride in elements (for FP8/BF16 kernels that need explicit stride). Same for all layers (element count is dtype-independent).

pub fn block_stride_bytes(&self) -> usize

Block stride in bytes (for NVFP4 kernels), using the uniform dtype.

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

Block stride in bytes for a specific attention layer. For symmetric dtypes returns the K stride (which equals V). For asymmetric dtypes returns the K-side stride; use v_block_stride_bytes_for_layer for the V-side stride explicitly.

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

K-side block stride in bytes for a specific attention layer. Same as block_stride_bytes_for_layer; named for clarity in asym call sites.

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

V-side block stride in bytes for a specific attention layer. Differs from K-side only for asymmetric KV cache dtypes.

pub fn nvfp4_data_bytes(&self) -> usize

NVFP4 data section size in bytes per block (uniform).

pub fn turbo4_data_bytes(&self) -> usize

Turbo4 data section bytes (same layout as NVFP4: 4-bit packed).

pub fn turbo3_data_bytes(&self) -> usize

Turbo3 data section bytes (3-bit packed).

pub fn turbo2_data_bytes(&self) -> usize

Turbo2 data section bytes (2-bit packed).

pub fn turbo8_data_bytes(&self) -> usize

Turbo8 data section bytes (FP8 E4M3 per element).

pub fn turbo4_scale_bytes(&self) -> usize

Turbo4 scale section bytes (same layout as NVFP4: FP8 per-group).

pub fn config(&self) -> &KvCacheConfig

Cache configuration (read-only). Used by attention layers to query the --high-speed-swap HBM-shrink cap (cache_blocks_per_seq).

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

Effective KV cache dtype for a specific attention layer.

pub fn block_size(&self) -> usize

pub fn num_blocks(&self) -> usize

pub fn dtype(&self) -> KvCacheDtype

pub fn num_layers(&self) -> usize

Number of attention layers.

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

Read K and V data for one block at one layer from GPU to host.

Returns (k_data, v_data) sized to each side’s block stride (which may differ for asymmetric dtypes).

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

Write K and V data for one block at one layer from host to GPU.

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

Compute how many blocks can fit given available GPU memory. Accounts for mixed dtypes when layer_dtypes is set.

Trait Implementations

impl ModelResource<dyn GpuBackend> for PagedKvCache

Release both pools of every layer.

Each layer allocates its K and V pools separately, so freeing per layer is correct. The block bookkeeping (free_blocks, block_ref_counts) is host state indexing into those pools — cleared with them so a released cache cannot hand out a block into freed memory.

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.