Skip to main content

Module kv_cache

Module kv_cache 

Source
Expand description

KV cache primitives for a single sequence within a single transformer layer.

Stores one contiguous, pre-allocated buffer for all tokens’ keys and one for all values (shape: max_tokens_per_seq * num_heads * hidden_dim_per_head floats each), instead of one small heap allocation per token. The previous design called resize_with on a Vec<KVCacheEntry> where every entry independently allocated its own keys/values Vec<f32> — for the default 8192-token max sequence length, a full-length initialize() performed up to 16,384 separate small heap allocations. A single upfront allocation eliminates that, keeps a sequence’s KV data in one contiguous cache-friendly region, and enables read_range — a single batched read over a token span, the shape attention actually needs, instead of N separate per-token reads.

Reads and writes are guarded by an RwLock rather than a Mutex: real attention reads happen far more often than writes (every position reads all previous positions’ KV every decode step; only one write happens per step), so letting reads proceed concurrently with each other matters.

Zero-copy read paths (with_read_kv, with_read_range, read_range_into) allow inspecting or copying cached keys and values directly without owned heap allocations.

§Read API: prefer the zero-copy closure forms

with_read_kv / with_read_range are the intended read API — they hand out borrowed slices under the read lock and allocate nothing. The owned forms read_kv / read_range exist only as a reference implementation and as the equivalence oracle for the zero-copy paths’ tests; they are #[doc(hidden)] and deliberately not part of the public surface.

Measured on the reference host (cargo bench -p ghostlink-core --bench kv_cache, wide config), the difference is not a micro-optimization:

operationownedzero-copy
read_range_10244.57 ms12.8 ns
decode_step_256141.9 ms1.26 ms
concurrent_4_readers_2561.63 ms221 ns

A naive wiring of the owned paths into a decode loop costs 100x-350,000x. Use the closure forms.

This module has no current caller in runtime.rs — Ghostlink delegates actual model execution to an external inference engine (llama-server / Ollama, see ghost-link::native_engine) rather than computing attention itself. It stays self-contained as a ready-to-use primitive for a future local execution path, so it can evolve without affecting the stable public API exported from lib.rs.

Structs§

KVCacheConfig
KVCacheEntry
Owned copy of one token’s cached key/value vectors, returned by read_kv. Small and fixed-size (config.token_width() floats each) — not a copy of the whole cache.
KVSpan
Borrowed slice view over a contiguous token range’s cached keys and values.
LayerKvCache
Cheap to clone (an Arc around the shared, lock-guarded buffer) — matches the shared-state pattern used elsewhere in this crate (e.g. ClusterState).

Constants§

DEFAULT_HIDDEN_DIM_PER_HEAD
DEFAULT_MAX_TOKENS_PER_SEQ
DEFAULT_NUM_HEADS