KV cacheThe KV cache

Autoregression computes only one new token per step, all thanks to caching prior tokens' K/V - that is the KV cache (the "memory" L03/L04 kept mentioning). This lesson digs into llama-kv-cache: how cells are managed, how to shift when the context fills, how multiple sequences coexist, and the various variants for long context. It closes Part 4 and is the physical basis for a large model "remembering the earlier text".

Why close with the KV cache? Because it is a VRAM heavyweight and the key to real capabilities like long context and concurrency. L17 said the context holds it, L18 said the batch writes into it - this lesson takes that "memory" itself fully apart.

Why a KV cache is needed

First, what it solves. Attention has "the current token" look at "all earlier tokens", needing each historical token's K (key) and V (value). Without a cache, generating each new word recomputes all prior tokens' K/V - and as length grows a bit, the recompute grows quadratically.

With a cache, each prior token's K/V is computed once and stored, and each step only computes the new token's K/V, appends it to the cache, and reads all historical K/V for attention. So per-step compute drops from "recompute the whole segment" to "compute one" - exactly what L04 proved: caching and recompute are numerically identical, but an order of magnitude faster.

This also directly explains L03's "why decode is fast": decode adds only one token per step, reusing history via the KV cache, so a step does just one token's forward. In short, without the KV cache there is no practical autoregressive generation - it is the key piece turning "computable in theory" into "actually runnable".

A quick recap of what K/V are (L04/L11): in attention, the current token uses its own Q (query) to dot-product with each historical token's K (key), computing "who to attend to", then weight-sums each historical token's V (value) by those weights. So "looking at history" needs exactly each historical token's K and V - cache them and you avoid recomputing each step.

The prefill phase (L03) is exactly filling the cache with the whole prompt's K/V at once: feed the prompt's tens-to-hundreds of tokens as a batch (L18), compute their K/V in parallel, and write them all into cells. Once filled, the cache holds the whole prompt's memory, and later when decode generates word by word, each new word reads this history directly without recomputing the prompt.

A note on where the letters "K/V" come from: a database-style "key-value" analogy - K is like an index (Q matches against it to find relevant positions), V like the retrieved content (summed by match weights). Do not press the analogy too hard, but it helps you remember: caching K/V is caching "each historical position's retrievable content".

Feel the "quadratic vs linear" gap concretely: generating the 1000th token, without a cache you would recompute all prior 999 tokens' K/V, each step slower than the last, quadratic overall; with a cache, step 1000 is like step 10, both computing just one new token, linear overall. For long conversations of thousands of tokens, this gap is the line between "runnable" and "frozen".

How cells are managed: pos and seq_id

Inside, the KV cache is a grid of cells, each storing one position's K/V. Managing them is llama_kv_cells: it records each cell's position pos and the sequence seq_id it belongs to; the cache also has a rolling write pointer head, marking where to write next.

// simplified from src/llama-kv-cells.h / src/llama-kv-cache.h
class llama_kv_cells {            // manages the grid of cells
    std::vector<llama_pos> pos;    // each cell's position
    // each cell also records: which seq_ids it belongs to
};
class llama_kv_cache : llama_memory_i {  // src/llama-kv-cache.h
    llama_kv_cells_vec v_cells;   // actual storage(multi-sequence capable)
    // head(): rolling write pointer
};

Note llama_kv_cache derives from llama_memory_i - this base interface matters, and we return to it for the variants. Each cell carries pos because attention's rope (L16) and causal mask both need to know "which position this K/V is"; each cell carries seq_id to support multiple sequences (below).

Each decode step, after build_attn (L16) computes the new token's K/V, it writes them into the cell head points to and advances head by one; for attention, it reads historical K/V from all cells belonging to this sequence. So the KV cache is not passive storage but an active memory that grows every step and is read every step. Draw a few consecutive steps and "only one new cell" becomes obvious:

Cells and tokens are one-to-one: the i-th token in a sequence occupies some cell in the cache, storing its (one per layer, actually) K and V. So "how big the cache is" roughly equals "how many tokens' K/V it can remember" - which is the meaning of n_ctx (context length, L17): the upper bound on cells the cache can hold.

"Reading only this sequence's cells" during attention is worth clarifying: the cache may mix several sequences' cells, but the current token should see only its own sequence's K/V at positions before itself. The former is filtered by seq_id, the latter by the causal mask (L11's soft_max_ext). The two filters together ensure "sequences do not mix flavors, and each token sees only the past".

Emphasize "one per layer": a transformer has dozens of layers, each with its own attention, each caching its own set of K/V. So a token's "memory" is actually a collection of dozens of K/V (one per layer). This is why a deeper model's KV cache is especially big - layer count multiplies straight into the cache size (the "layer count" in Dig-deeper 1's product is exactly this).

Sequence operations and context shift

The KV cache is not just "write into" - it can be edited. A set of sequence operations lets you remove, copy, keep, or shift a sequence's K/V, corresponding to the public C API llama_memory_seq_* (obtaining the memory object via llama_get_memory).

# pseudocode: sequence operations (public C API)
mem = llama_get_memory(ctx)
llama_memory_seq_rm (mem, seq, p0, p1)      # remove the [p0,p1) span of KV
llama_memory_seq_add(mem, seq, p0, p1, d)   # context shift: shift pos by d

The most worthwhile thing here is context shift. When a conversation grows past n_ctx and the cache fills, what then? One way: drop the oldest span of K/V, shift the remaining tokens' pos forward as a whole (a negative seq_add), and free tail space to keep generating. This avoids "stop when full" without recomputing the whole segment - it just slides the memory window forward a bit.

seq_cp (copy a sequence) has a clever use: shared prefix. Say a long system prompt must generate several different answers at once - compute it once into sequence 0, then seq_cp it to sequences 1, 2, 3... so several sequences share the same prefix's KV without each recomputing it. A common server compute-saving trick.

Back to that rename: from llama_kv_self_* to llama_memory_seq_* is not just a name change but a conceptual abstraction upgrade. Early on there was only one "KV cache", so the API was called kv; later came recurrent, a "not-KV but still memory" thing, so the name was lifted to the more general memory. One rename records this engine's growth from "supporting only transformers" to "also accommodating other architectures".

The p0, p1 common in those sequence operations are a position range: many operations act not on a whole sequence but on the "positions p0 to p1" span - removing only a span, shifting only a span. This "operate by position interval" granularity lets the engine do many tricks - rolling back only the last few tokens (undo), shifting only a middle span. Making memory editable by interval is where the flexibility comes from.

Multiple sequences and variants

Because each cell carries seq_id, one KV cache can hold multiple sequences at once: their cells mix in the same cache, distinguished by seq_id, each doing its own attention (seeing only its own sequence's cells). This is the memory basis for L18's multi-sequence and a server's concurrency.

Why so many variants? Because while standard full-attention's KV cache makes compute linear, its memory still grows linearly with length - the longer the context, the more VRAM. Sliding window (iSWA) keeps only a recent window, recurrent uses fixed state, hybrid mixes - all different trade-offs to save memory for long context. They all implement the same base interface llama_memory_i, so they can be swapped wholesale with the rest of the engine unchanged.

Emphasize multi-sequence "no mixing" once more, as it is the key to concurrency: three sequences' cells crowd into the same cache, but when each sequence does attention, the seq_id filter lets it see only its own cells, as if the cache held only it. So one physical cache is logically split into mutually-invisible portions. This "physically shared, logically isolated" is the same memory-saving philosophy as L17's context sharing weights.

The recurrent variant goes further: it corresponds to non-transformer architectures like Mamba, using a fixed-size state to summarize "all history so far", the state size not changing with context length at all. This fundamentally sidesteps the KV cache's growth-with-length problem, at the cost of the state being "compressed history" - not able to look back precisely at each token as full attention can. Folding it too into llama_memory_i is exactly this abstraction's power - even architectures whose "fundamental memory mechanism differs" can plug into the same engine.

Put the KV cache back into the whole inference chain: L18's batch feeds a new token in -> L16's build_attn computes its K/V, writes them into the KV cache, and reads history back out -> after computing, updates the cache and advances one step. The KV cache is that memory read and written every step, the state core that "links past and future" in the autoregressive loop. All the earlier components, in the end, revolve around it.

One more thing worth knowing: the KV cache's memory layout works together with optimizations like flash attention (mentioned in L11) - laying K/V tidily in memory is what lets the attention kernel read block by block and compute as it reads efficiently. So the KV cache is not just about "fitting"; how it is laid out also directly affects how fast attention computes. Storage layout and compute kernel are a mutually-accommodating pair at this layer.