KV Cache¶

Table of Contents¶

• What Is a KV Cache?
• Memory Growth
• Three Structural Cache Layouts
• Architecture Selection
• Per-Sequence State and CachePool
• Recurrent and Hybrid Models
• Attention Dispatch and In-Place Decode Writes
• Paged Decode Storage and Kernels
• Matrix-Absorbed MLA Decode
• Speculative Rewind

What Is a KV Cache?¶

During autoregressive decoding, every attention layer reuses the keys (K) and values (V) produced for earlier tokens. A KV cache stores those tensors so each decode step needs to project only the new token instead of recomputing the entire prefix.

Without caching, repeatedly processing an ever-growing prefix makes generation prohibitively expensive. With caching, the amount of stored state grows with the visible context, while the token projection work for each ordinary decode step stays constant. Attention still reads the relevant cached prefix or window, so total attention cost is not literally constant.

For a typical grouped-query attention layer, K and V have shape [batch, num_kv_heads, sequence, head_dim]. Query-head count can be larger than KV-head count and must not be used for cache sizing.

Memory Growth¶

For layers with the same dimensions, approximate cache growth per token and sequence is:

2 × num_layers × num_kv_heads × head_dim × bytes_per_element

The factor 2 accounts for both K and V. For heterogeneous architectures, sum the per-layer K and V sizes instead. Batch size, quantized-cache metadata, alignment, allocator slack, paged block granularity, and temporary tensors add overhead beyond this estimate.

The standard cache grows buffers in allocation steps so it does not reallocate on every token. Its allocated size can therefore exceed the filled length reported by offset.

Three Structural Cache Layouts¶

src/lib/mlxcel-core/src/cache.rs defines three attention-cache layouts. These layouts are separate from the user-facing KV storage mode (for example FP16, uniform quantization, or TurboQuant) and from the server's dense-versus-paged storage backend.

Standard KVCache¶

KVCache retains the full visible prefix. It is used by conventional full-attention layers and provides trim() so speculative decoding can rewind unaccepted tokens. Buffers grow in chunks and writes use in-place slice updates where supported.

RotatingKVCache¶

RotatingKVCache implements sliding-window attention. Once the logical window is full, the oldest entries are replaced and the cache exposes a linearized recent window to attention. Memory is bounded by the configured window plus implementation allocation slack; it is not guaranteed to equal exactly window_size × bytes_per_token at every instant.

Representative users include Gemma ¾, ExaOne 4 and ExaOne MoE, GPT-OSS, Ministral 3, Mellum, Muse Glimmer, RecurrentGemma, Step 3.5, and AFMoE. Some of these mix rotating and full-attention layers.

ChunkedKVCache¶

ChunkedKVCache tracks both a global token offset and the start of a bounded visible chunk. It is used for Llama 4's interleaved/chunked attention layers; full-attention layers in the same model continue to use standard caches.

Jamba does not use ChunkedKVCache for its Mamba layers. It keeps model-owned Mamba recurrent state and uses standard KVCache only for its attention layers.

Architecture Selection¶

Each model's make_caches() implementation selects the layout per attention layer. Hybrid models may return different layouts for different layers.

The exact mix is an implementation property of the loaded architecture, not a global command-line choice.

Per-Sequence State and CachePool¶

Both the CLI and the server keep cache state in a CachePool owned by the decode Engine. mlxcel generate runs on DirectEngine, which opens one sequence in that pool per generation and closes it afterwards. The server assigns each request a monotonically increasing SequenceId and uses CachePool to track per-sequence cache sets, prompt lengths, offsets, and allocation time.

SequenceId is a u64; wrapping is practically unreachable during a reasonable server lifetime, but it is not a mathematical “never wraps” guarantee.

CachePool itself is deliberately not Sync. Callers that share it across threads must provide external Mutex/RwLock synchronization. Allocation enforces max_sequences, release removes completed state, and memory_usage_bytes() sums the allocated K/V buffers. The server's admission queue, rather than CachePool alone, determines whether a saturated request waits or receives a busy response.

For architectures with model-owned state, the pool holds lightweight placeholders for scheduler bookkeeping while the model keeps the real state itself, keyed by SequenceId (LanguageModel::sequence_state_layout, prepare_sequence_state, sync_sequence_storage). The scheduler admits several such sequences to one batched forward pass only when the model reports supports_batching(). A model that does not report it is clamped to one sequence.

Recurrent and Hybrid Models¶

Mamba, Mamba2, RWKV, and hybrid architectures carry recurrent state whose size does not grow with the full context like a standard KV cache. This state is model-owned and keyed by sequence id, and Mamba, Mamba2, Jamba, and RWKV7 report supports_batching() with BatchedDecodeKind::Native, so the server decodes several of their sequences in one batched forward pass with row isolation. Families that do not report supports_batching() run through a sequential single-slot worker.

Dense transformer models may opt into padded prefill on eligible hardware and then trim outer caches to the real prompt length. A family that owns its attention caches trims its own state through LanguageModel::trim_state; Gemma 3 does this and takes padded prefill like a dense model. Recurrent hybrids such as NemotronH, and AFMoE, Llama 4, Gemma 4, DeepSeek V4, Bailing linear, and Qwen3-Next, report supports_padded_prefill() == false, so normal generation does not pad their prefill. MLXCEL_FORCE_PADDED_PREFILL forces tile alignment on any hardware and MLXCEL_NO_PADDED_PREFILL turns it off; both apply to the CLI and the server, and MLXCEL_NO_PADDED_PREFILL wins when both are set. The trim_internal_caches() hook remains a defensive model-specific mechanism for any path that must discard excess internal state; it should not be read as evidence that NemotronH normally uses padded prefill.

The server's decode lookahead pipelines a model-owned family only when the model can unwind its speculative single-token appends exactly (supports_decode_lookahead_rewind()). Gemma 3 logs the rows each decode write overwrites in its sliding layers, so it decodes on the lookahead pipeline when every sliding layer is FP16; other model-owned families stay on synchronous decode.

Attention Dispatch and In-Place Decode Writes¶

The cache, not the model, decides how attention runs. A model calls KVCache::attend (or cache::attend_batched for a batch of rows), and the cache chooses dense SDPA or the paged decode kernels from how its state is stored. RotatingKVCache::attend and ChunkedKVCache::attend apply the same dispatch while keeping their window semantics.

Single-token FP16 decode writes the new K/V row into the cache buffer in place instead of rebuilding the cache with a slice update. This is on by default; MLXCEL_KV_INPLACE_WRITE=0 (or false, off, no) returns to the slice-update write.

Paged Decode Storage and Kernels¶

The server can store per-sequence KV state in fixed-size pages instead of a single dense allocation. This makes block reuse, eviction, preemption, and prefix sharing practical under a bounded --kv-cache-budget. Paged storage is independent of whether an architecture's attention layer is full, rotating, or chunked; wrappers preserve the model's visible-window semantics.

The production paged decode path can use fused paged attention, sparse page-indirection decode, and opt-in cascade attention for groups that share a sufficiently long whole-page prefix. Operational controls include MLXCEL_PAGED_DECODE_SPLITS, MLXCEL_SPARSE_PAGED_ATTENTION, MLXCEL_SPARSE_PAGED_MIN_SPARSITY, MLXCEL_CASCADE_ATTENTION, MLXCEL_CASCADE_MIN_MEMBERS, and MLXCEL_CASCADE_MIN_SHARED_PAGES. See Performance Tuning before overriding defaults.

Matrix-Absorbed MLA Decode¶

DeepSeek-family multi-head latent attention can cache a compressed latent instead of expanded per-head K/V. MLXCEL_MLA_ABSORBED=1 enables this lower-footprint path. Absorption also covers multi-token steps: mlxcel_core::mla::use_absorbed_mla compares the per-call cost of absorption against materializing the window, so a speculative verify block or a chunked-prefill continuation behind a long cache runs absorbed, while a cold prefill materializes and stays numerically identical to the baseline. MLXCEL_MLA_SPLIT_KV=1 additionally splits a long latent range and requires absorbed MLA. Both are off by default.

Speculative Rewind¶

Speculative decoding proposes several tokens and verifies them with the target model. When verification accepts only a prefix, cache state beyond the accepted position must be removed. KVCache::trim(n) rewinds the logical offset, and the speculative round loops (the engine's token-only loop for prompt lookup and full-model drafters, and the MTP and DFlash loops) trim target and draft state as their rounds require. Other cache layouts and model-owned drafter paths preserve their own position/window invariants rather than assuming every cache is a grow-only dense tensor.