6 hours ago|
AI

KV Cache and How It Works

A technical guide to key-value caching in large language model inference

Blog Image

KV Cache and How It Works

A technical guide to key-value caching in large language model inference

Tom Hermann · AI Strategy and Adoption · September 2026

Introduction

Autoregressive language models generate text one token at a time. Each new token attends to every token that came before it. Without an optimisation, that property forces the model to recompute key and value projections for the entire context at every decoding step. The cost grows quadratically with sequence length and quickly dominates inference time.

The key-value cache, commonly called the KV cache, stores those projections after they are computed once. Subsequent steps reuse the stored tensors and compute projections only for the newest token. The technique is the principal reason modern transformers can generate long responses at interactive latency. This article explains the mechanism in detail, from attention arithmetic through tensor layout, memory scaling, grouped-query attention, and the engineering trade-offs that appear in production systems.

Self-attention without a cache

Scaled dot-product attention maps a query matrix Q against a key matrix K and a value matrix V. For a single attention head the formula is:

Attention(Q, K, V) = softmax(Q Kᵀ / √d) V

Here d is the head dimension. In a decoder-only transformer used for generation, Q, K and V are linear projections of the hidden states of the current sequence. Causal masking ensures that position i cannot attend to positions greater than i.

During training the full sequence is known, so Q, K and V are computed in one parallel pass. During inference the sequence is built incrementally. At step t the model must produce a distribution over the next token given tokens 1 to t. Naively that means projecting all t hidden states into K and V, forming a t by t attention matrix, and discarding everything except the last row. Almost all of that work repeats work already done at steps 1 to t−1.

What the KV cache stores

After the first pass over a prompt, every transformer layer has already produced keys and values for every prompt token. Those tensors do not depend on future tokens under causal attention. They can therefore be retained.

The cache is a pair of tensors per layer (sometimes per attention group). Typical layout is:

•        K cache: shape [batch, num_kv_heads, seq_len, head_dim]

•        V cache: shape [batch, num_kv_heads, seq_len, head_dim]

Some implementations store the sequence dimension first or pack heads into a single leading dimension. The invariant is the same: each past position contributes one key vector and one value vector per key-value head.

Queries are not cached. A query is needed only for the token being generated. Caching Q would waste memory and would be unused, because past queries never participate in later attention rows under causal decoding.

Prefill and decode

Prefill

Prefill is the processing of the prompt. The model receives all prompt tokens at once. Hidden states are computed in parallel across the prompt length. Each layer writes the full K and V tensors into the cache. Attention during prefill is still causal, so the implementation uses a causal mask or a flash-attention kernel that respects causality. Prefill is compute-bound on modern GPUs because matrix multiplications are large and regular.

Decode

Decode is the token-by-token loop. At step t the model embeds only the latest token, runs it through the stack, and at each attention layer:

•        Projects the new hidden state into a single query vector q, a single key k, and a single value v.

•        Appends k and v to the layer cache, increasing seq_len by one.

•        Computes attention between q and the entire cached K, then applies the resulting weights to the entire cached V.

The attention matrix at this step is 1 by t rather than t by t. The expensive projections of past tokens disappear. Decode is typically memory-bandwidth bound: each step streams the growing cache from GPU memory into the compute units.

A worked numerical sketch

Consider one layer, one head, head dimension 4, and a prompt of three tokens. After prefill the cache holds:

K = [k1, k2, k3],  V = [v1, v2, v3]

Each ki and vi is a vector of length 4. The model now generates the fourth token. It computes q4, k4, v4 from the new hidden state. The cache becomes:

K = [k1, k2, k3, k4],  V = [v1, v2, v3, v4]

Scores are the four dot products q4·k1, q4·k2, q4·k3, q4·k4, scaled by 1/√4. Softmax yields four weights. The output of the head is the weighted sum of v1 to v4. No earlier key or value is recomputed. The same pattern repeats at every later step and at every layer.

Memory footprint

Cache size scales linearly with sequence length, number of layers, number of key-value heads, and head dimension. For floating-point 16 storage the bytes required are approximately:

2 × num_layers × num_kv_heads × seq_len × head_dim × 2 bytes

The leading factor of two accounts for keys and values. A 70 billion parameter model with 80 layers, 8 key-value heads, head dimension 128, and a 32,768 token context consumes on the order of 40 gigabytes of cache at FP16 for a single sequence. That figure often exceeds the weight memory of the attention projections themselves and is the reason long-context serving is constrained by GPU capacity rather than by arithmetic throughput.

Batching multiplies the cost by batch size. Continuous batching keeps many sequences alive at different lengths, so production servers allocate cache pages dynamically rather than reserving the maximum length for every slot.

Multi-head, multi-query and grouped-query attention

Classic multi-head attention uses the same number of query heads and key-value heads. Multi-query attention (MQA) shares a single K and V head across all query heads. Grouped-query attention (GQA) sits between the two: query heads are partitioned into groups, and each group shares one K and one V head.

GQA is now common in open-weight models because it shrinks the cache by the ratio of query heads to key-value heads, often 4× or 8×, with a modest quality cost. During decode the single new query set is still compared against the reduced K cache. Implementation-wise the cache tensors simply have fewer head slots; the append-and-attend algorithm does not change.

Implementation notes

Frameworks expose the cache as an object that travels with the forward pass. In conceptual terms the attention module accepts past_key_values, concatenates the new k and v along the sequence axis, runs attention, and returns the updated pair. Concatenation on every step would copy the whole cache; production kernels therefore preallocate a maximum-length buffer and write the new slice in place, tracking the valid length with a position index or a slot mapping.

Rotary position embeddings (RoPE) are applied to queries and keys. Implementations either store keys already rotated, or store unrotated keys and apply rotation at attention time using the absolute positions of cached tokens. The first choice saves compute. The second choice simplifies some cache-reuse patterns such as prefix sharing across requests.

Attention masks during decode are usually implicit. If the cache buffer is larger than the current length, the kernel must ignore padding slots, either with an explicit mask or by passing the true sequence length into a flash-attention or paged-attention kernel.

Optimisations used in serving stacks

Paged attention, popularised by vLLM, splits the cache into fixed-size blocks and maps logical token positions to physical blocks through a page table. Sequences no longer need contiguous reservation of the maximum context. Blocks can be shared when several requests share a prompt prefix, which is common in system-prompt heavy applications.

Cache quantisation stores K and V in 8-bit or 4-bit formats with per-channel scales. Because decode is bandwidth bound, reduced precision often improves tokens per second even after the extra dequantisation arithmetic. Quality impact is usually small if the quantisation is applied per head or per channel rather than per tensor.

Prefix caching and prompt caching persist the prefill result for repeated system prompts or retrieved documents. Speculative decoding still uses a KV cache; the draft model and the target model each maintain their own, and rejected draft tokens require a rollback of the target cache to the last accepted position.

Limits and failure modes

The cache does not reduce the quadratic cost of prefill. Very long prompts remain expensive to ingest even when decode is cheap. Eviction policies for caches that exceed device memory (sliding windows, attention sinks) change the semantics of attention: tokens that fall out of the window can no longer be attended to, which can degrade tasks that require exact recall of early instructions.

Incorrect position handling is a frequent source of bugs. If RoPE is applied with the wrong absolute index after a cache reuse or a speculative rollback, later tokens attend with a corrupted geometry and generation quality collapses even though shapes remain valid. Another class of error is failing to isolate caches across concurrent sequences in a batch, which leaks keys from one user request into another.

Why the cache dominates inference design

Weight memory is fixed. Activation memory for a single token is small. The KV cache is the only major tensor that grows with conversation length and with concurrency. Throughput, maximum context, and cost per token in a serving fleet are therefore cache-allocation problems as much as they are model-architecture problems. Architectural choices such as GQA, sliding-window attention, and linear-complexity variants exist largely to shrink or bound that tensor.

For practitioners building generation pipelines the operational rule is simple. Prefill once, store K and V, decode with append-only updates, and treat cache lifetime as part of the request lifecycle. Everything else, from kernel choice to quantisation to paging, is an engineering refinement of that rule.

Conclusion

KV caching converts repeated full-sequence projection and attention into a constant-time projection of one token plus a linear scan of stored keys and values. The idea is a direct consequence of causal self-attention: past keys and values are independent of the future. Understanding tensor layout, prefill versus decode, memory scaling, and grouped-query reductions is sufficient to reason about latency, GPU memory, and serving architecture for contemporary language models.


#KVcache#key-valuecache, #transformer #inference, self-attention, autoregressive decoding, #LLM #inferenceoptimisation.

Share on:

0 comments

No comments yet

Your Views Please!

Your email address will not be published. Required fields are marked *
Please Login to Comment

You need to be logged in to post a comment on this blog post.

Login Sign Up

You may also like