2 Solving the inference bottleneck with the key-value cache
This chapter introduces the inference bottleneck in autoregressive language models and explains why it matters specifically at generation time. A model like GPT-2 or DeepSeek produces text one token at a time, repeatedly feeding its own output back into the input. Although only the newest token determines the next prediction, the entire sequence is reprocessed at each step, so the same calculations are done over and over as the context grows. By walking through the transformer pipeline from embeddings to attention, feed-forward layers, and finally logits, the chapter shows that the last token’s representation is the only one needed for the next prediction, making the repeated recomputation in earlier positions clearly wasteful.
The key-value cache solves this inefficiency by storing the key and value representations from previous tokens so they do not have to be recomputed every time a new token is generated. During inference, only the new token’s query, key, and value are calculated fresh; the older keys and values are retrieved from memory, appended to the cache, and reused for attention. This transforms the per-token cost of decoding from quadratic growth to roughly linear growth, producing major speedups in practice. However, the chapter also emphasizes that the cache is not free: its memory footprint grows with layers, heads, sequence length, and batch size, and for large models or long contexts it can become a severe deployment bottleneck.
To reduce that memory burden, the chapter presents Multi-Query Attention and Grouped-Query Attention as the first generation of cache-saving architectural ideas. MQA makes all heads share a single key and value projection, dramatically shrinking the cache but also weakening the model’s ability to let different heads specialize. GQA softens this compromise by letting heads share keys and values only within groups, preserving more expressivity while still lowering memory use. Together, these methods define the performance-versus-memory trade-off that motivates the next step in the book: DeepSeek’s approach, which aims not merely to reduce the number of heads, but to compress the information inside the attention mechanism itself.
The four-stage roadmap for building the DeepSeek model. Stage 1 establishes the Key-Value Cache Foundation, which we will cover in this chapter.
In the autoregressive generation loop, the model's output from one step is appended to the input for the next, progressively extending the context.
High-level overview of the Transformer block architecture. The diagram shows the data flow from initial input tokens ("The next day is bright") through embeddings, multi-head attention, and feed-forward layers, producing the logits used for next-token prediction.
The input embedding matrix X is projected into three new matrices: Query, Key, and Value. Each projection is a matrix multiplication by a unique, learned weight matrix.
The dot product of the query and transposed key matrices produces the attention scores matrix. Each element in this matrix represents the relevance of one token to another.
The attention weights are multiplied by the value matrix to produce the final context matrix. Each row is now a context-aware representation of the original token.
Parallel projections in multi-head attention. The input embedding X is projected into separate query, key, and value matrices for each attention head.
Each attention head computes its own unique attention score matrix in parallel.
The attention scores for each head are processed independently to create the final attention weights.
Each head produces its own context matrix, which represents its unique, context-aware view of the input sequence.
The context matrices from all heads are concatenated to form a single, richer matrix, and then passed through a final projection layer.
The journey from the final context matrix to the logits matrix. The output layer projects each context-aware vector into a long vector of scores, one for each word in the vocabulary.
The final prediction step. The logits vector for the last token is converted into a probability distribution, and the token with the highest probability is selected as the output.
Full attention calculation for the input sequence "The next day is."
The full attention calculation for this new, 5-token input.
A plot comparing the quadratic (O(n²)) growth of computations in uncached autoregressive inference versus the ideal linear (O(n)) growth.
Out of the entire context matrix, only the final row corresponding to the newest token is required to predict the next token in the sequence.
The context vector for "bright" is calculated by multiplying the attention weights for "bright" with the full Value matrix.
The attention weights for "bright" are derived from its query vector and the key vectors of all tokens in the sequence.
The three essential computations for a new token. The input embedding for "bright" is projected to create its unique query, key, and value vectors.
The formula for calculating the KV cache size and its application to well-known models.
A comparison of model parameter count versus KV cache size for different GPT-3 model variants. The dashed line shows that as models become larger, their memory requirements for the KV cache grow at a similar, steep rate.
Standard Multi-Head Attention (MHA). Each of the four attention heads has its own distinct Key and Value weight matrices, indicated by the different colors. This allows each head to specialize and learn different patterns.
Multi-Query Attention (MQA). All four heads still have unique Query projections, but they now share a single, common Key and Value projection, indicated by the uniform light blue and yellow colors.
Grouped-Query Attention (GQA). The four attention heads are divided into two groups. Within Group 1 (light blue/light yellow), Head 1 and Head 2 share the same Key and Value projections. Within Group 2 (dark blue/dark yellow), Head 3 and Head 4 share a different, distinct set of Key and Value projections.
Summary
- Autoregressive generation, where each new token is appended to the input, results in the re-processing of the entire sequence at every step in a naive implementation.
- This repeated computation leads to a quadratic (O(n²)) complexity problem, which makes generating long sequences of text computationally impractical.
- The key-value (KV) cache optimizes inference by storing the key and value matrices of past tokens, avoiding redundant calculations and reducing the cost of generating each new token from quadratic to linear (O(n)) time (with total generation going from cubic to quadratic).
- While the KV cache dramatically accelerates computation, it introduces a severe memory bottleneck, as its size grows proportionally with sequence length, number of layers, and attention heads.
- Multi-Query Attention (MQA) drastically reduces the KV cache's memory footprint by forcing all attention heads to share a single key and value projection, but this significantly degrades model performance by preventing head specialization.
- Grouped-Query Attention (GQA) provides a tunable compromise by having groups of attention heads share key and value projections, allowing a balance between memory savings and model expressivity.
- Architectures like MQA and GQA fundamentally operate by reducing the number of unique key/value pairs, establishing an inherent trade-off between memory efficiency and the expressive power of the model.
Build a DeepSeek Model (From Scratch) ebook for free
