Struct KvCache 

pub struct KvCache {
    pub layers: Vec<Option<KvLayer>>,
    pub cached_len: usize,
    pub n_kv_heads: usize,
    pub head_dim: usize,
    pub max_seq_len: Option<usize>,
    pub dtype: Option<DType>,
    /* private fields */
}

Backend-erased KV cache. Replaces lazy::KVCache<B> and lazy::LlamaKVCache. Indexed by global layer index; layers[i] is None if layer i hasn’t been populated yet.

Pipeline-parallel friendly: when a model is split across devices, each device’s cache holds entries only for its layer range; other indices are None. The model code queries cache.layer(global_i) regardless of which device owns it; the storage Arc inside KvLayer carries the device identity via its fuel_memory::BackendStorage variant.

Fields

layers: Vec<Option<KvLayer>>

cached_len: usize

n_kv_heads: usize

head_dim: usize

max_seq_len: Option<usize>

Pre-allocated capacity along the sequence axis. Some(n) when constructed via KvCache::with_capacity — every layer’s K + V storage holds [1, n_kv_heads, n, head_dim] of zeros and the forward path writes into it via Op::WriteSlice. None for the legacy KvCache::with_dims grow-by-replacement constructor.

dtype: Option<DType>

Dtype the pre-allocated buffers were allocated with. Some(dt) in the KvCache::with_capacity path, None in the legacy with_dims path (which leaves dtype to whatever the first inserted layer specifies).

Implementations

impl KvCache

pub fn with_dims(n_layers: usize, n_kv_heads: usize, head_dim: usize) -> Self

Legacy grow-by-replacement constructor. Layers start None; the caller calls Self::set_layer each forward step with a freshly allocated KvLayer whose storage is sized to the current cached length. Subsequent steps replace the layer entirely with a larger buffer.

Replaced by Self::with_capacity when the caller knows the generation’s max sequence length up-front (the common case for autoregressive decoding).

pub fn with_capacity( n_layers: usize, n_kv_heads: usize, head_dim: usize, max_seq_len: usize, dtype: DType, device: &Device, ) -> Result<Self>

Pre-allocated KV cache. Every layer’s K + V storage is allocated up-front as a [1, n_kv_heads, max_seq_len, head_dim] zero buffer on device with dtype. Subsequent forward steps write fresh K/V slabs into these buffers via Op::WriteSlice rather than growing-by-replacement.

Memory footprint: n_layers * 2 * n_kv_heads * max_seq_len * head_dim * dtype_size bytes. For Llama-7B at max_seq_len=4096, bf16: 32 * 2 * 8 * 4096 * 128 * 2 ≈ 512 MiB.

Returns Err if any per-layer allocation fails (e.g. CUDA OOM) or the requested device hasn’t been wired up in crate::pipelined_bridge (Vulkan / Metal — D2H/H2D for those still goes through the legacy executor).

pub fn slot_storage( &self, layer_idx: usize, slot: KvSlot, ) -> Option<Arc<RwLock<Storage>>>

Borrow the K-or-V storage Arc for a given layer + slot. Used by the forward path to bind cache storage to per-step Const nodes via InferenceContext::insert. Returns None if the layer isn’t populated (with_dims path before first Self::set_layer) or layer_idx is out of range.

pub fn bump_version(&mut self, layer_idx: usize, slot: KvSlot)

Bump a slot’s monotonic version counter. Called after every successful Op::WriteSlice targeting the slot — the forward path invokes this once per (layer, slot) per step.

The version is consumed by the future multi-device coherence protocol (placeholder today — see AuthorityState).

pub fn n_layers(&self) -> usize

pub fn layer(&self, li: usize) -> Option<&KvLayer>

pub fn layer_mut(&mut self, li: usize) -> Option<&mut KvLayer>

pub fn alloc_id(&self) -> KvAllocId

Identity of the storage this cache currently owns — the value a held DecodeSession’s validity key carries so a cache swap cannot be mistaken for a continuation. See KvAllocId for why re-allocation re-mints it and rewinding does not.

pub fn device_location(&self) -> Option<DeviceLocation>

Where this cache’s storage lives, or None when that is not provable. See the field’s doc — None is a refusal, never a CPU default.

pub fn set_layer(&mut self, li: usize, layer: KvLayer)

Install a layer’s storage, replacing whatever was there.

This re-mints Self::alloc_id: the storage Arc a held plan baked into its base_cache for this layer is no longer the storage this cache reads, so any plan welded to the old one must be rebuilt. It also clears Self::device_location: the storage arrives from the caller, so where it lives is no longer something this cache measured. That forfeits DecodeSession::rebind_kv for the rest of this cache’s life — deliberately, because an unprovable residency spliced into a baked plan is the GAP-014 failure one layer down.

pub fn clear(&mut self)

Drop every layer; reset cached_len to zero. Use between independent generations (the cache’s K/V shapes are tied to a specific prompt prefix).

Re-mints Self::alloc_id — every storage Arc is released, so this is an allocation change, not a rewind. Contrast Self::truncate_to, which moves only cached_len and therefore must preserve the id.

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

Shrink the cache to new_len cached positions. Speculative- decoding’s reject path calls this after a draft batch is rejected by the target model.

Phase E.3.0 limitation: this only updates cached_len; the underlying storages and their layouts are not shrunk. The grow-by-replacement pattern (today) means the next forward pass would still see the old layout’s full extent. Phase E.3.2 pairs this with pre-allocated buffers + Op::WriteSlice so truncate becomes a pure metadata update — the trailing rows of the pre-allocated buffer simply stop being read.

Deliberately does NOT touch Self::alloc_id. Rewinding the conversation is not re-allocating the buffers: a held plan is still welded to exactly the right storage, so invalidating here would forfeit plan reuse on every rejected draft batch and buy no correctness. The id names the allocation, never the conversation — see KvAllocId.

pub fn write_bytes( &self, layer_idx: usize, slot: KvSlot, start: usize, n: usize, bytes: &[u8], device: &Device, ) -> Result<()>

Write host-supplied K/V bytes into layer layer_idx’s slot buffer at position range [start, start+n) — the KvCache-side sibling of DeviceKvPool::write_block_bytes.

Builds a throwaway Op::Const(bytes) → Op::WriteSlice graph and realizes it ONCE into the SAME buffer Arc this cache already holds — it does not replace the Arc, so unlike Self::set_layer this does not touch Self::alloc_id or Self::device_location: a DecodeSession already welded to this cache’s buffers (via base_cache) sees the write transparently on its next token, with no rebind and no rebuild — the write lands in the exact buffer the held plan’s Op::WriteSlice nodes already target.

bytes is n * n_kv_heads * head_dim * dtype.size_in_bytes() raw little-endian element bytes, row-major [1, n_kv_heads, n, head_dim] — HEAD-MAJOR (position is axis 2). This is NOT the same layout as DeviceKvPool’s position-major [block_size, Hkv, D] block layout: a caller moving bytes from the paged pool into a KvCache must transpose, a plain memcpy silently produces wrong attention.

Err (never a panic): this cache has no fixed dtype/capacity (built via Self::with_dims rather than Self::with_capacity — this primitive is with_capacity-only, since the legacy grow-by-replacement path has no stable buffer to write into); layer_idx is out of range or unpopulated; start + n exceeds the allocated max_seq_len; or bytes.len() doesn’t match the expected size for n positions.

Trait Implementations

impl KvRebindSource for KvCache

fn alloc_id(&self) -> KvAllocId

Identity of this allocation. See KvAllocId.

fn rebind_location(&self) -> Option<DeviceLocation>

Where the storage lives, or None when that is not provable — which DecodeSession::rebind_kv treats as a refusal, never as a default.

fn rebind_slot( &self, layer: usize, which: usize, ) -> Option<Arc<RwLock<Storage>>>

Storage for slot which (0 or 1) of layer — the two halves of the plan’s kv_nodes[layer] pair, in the order the builder bound them.