Efficient Inference & Compression

Training happens once; inference happens every time someone uses the model. Generating text is slow for a reason that has little to do with arithmetic: to produce each token the GPU (graphics processing unit) must read every weight of the model, and often the whole history of the conversation, from memory. Reading memory, not computing, is the bottleneck.

This page covers the techniques that make generation cheaper:

• caching keys and values, so that nothing is computed twice;
• decoding strategies that search instead of sampling;
• speculative decoding, which uses a small model to guess ahead;
• compressing the weights by quantization, pruning and low-rank adapters, so that open-weight models fit on local hardware.

The KV cache

A language model generates one token at a time, and each new token attends to all the previous ones. Recomputing the keys and values of the whole prefix at every step would repeat the same work again and again. But in a causal model the keys and values of past tokens never change, so they can be stored: the KV (key-value) cache. Each step computes the query, key and value of the new token only, appends the key and value to the cache, and attends over the cache.

Example. Generating $T = 1{,}000$ tokens. Without a cache, step $t$ recomputes the keys and values of all $t$ tokens so far: $1 + 2 + \dots + 1000 = 500{,}500$ computations. With a cache, each step computes one: $1{,}000$.

A tiny trace, one head with $d = 2$ (queries and keys equal to the input, values equal to the input with the second number doubled). Tokens $(1, 0)$, $(0, 1)$ and $(1, 1)$ arrive one by one:

step	new token	cache after the step	output
1	$(1, 0)$	1 key, 1 value	$(1, 0)$: the only token attends to itself
2	$(0, 1)$	2 keys, 2 values	$(0.33, 1.34)$
3	$(1, 1)$	3 keys, 3 values	$(0.75, 1.50)$

Each output is the same as the corresponding row of full causal attention over the three tokens, computed from scratch.

The cache is what limits batch size and context length (Pope et al., 2023). Its size is

where $L$ is the number of layers, $G$ the number of key/value heads of size $d$, $T$ the number of cached tokens, and the 2 is for keys and values; all per sequence.

Example. A model with 32 layers, 32 key/value heads of size 128, with values stored in BF16 (brain floating point, 2 bytes): $2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288$ bytes, 512 KiB (kibibytes, 1 KiB is 1,024 bytes) per token. For one sequence of 8,192 tokens: $4$ GiB (gibibytes, 1 GiB is $2^{30}$ bytes).

With grouped-query attention (GQA, section 7) and only 8 key/value heads, the cache of the same sequence is 4 times smaller: 1 GiB.

This is why grouped-query attention (section 7) and sliding windows (section 10) matter so much at inference time. Serving systems such as vLLM go further and allocate the cache in fixed-size pages, like virtual memory, so that many sequences of different lengths share GPU memory without fragmentation (Kwon et al., 2023).

Beam search

Greedy decoding picks the most likely token at each step, but the most likely sequence may start with a less likely token. Beam search keeps the $k$ best partial sequences (the beam) by total log-probability. At each step every sequence in the beam is extended with every possible next token, and only the $k$ best of all these candidates survive. Sequences that emit an end-of-sequence token are moved to a list of finished hypotheses.

Example. Two steps. At step 1 the model gives A probability $0.6$ and B probability $0.4$.

• After A, the next token is "continue" with probability $0.4$, or two other tokens with $0.3$ each.
• After B, the next token is "end" with probability $0.9$.

Greedy takes A (0.6) and then "continue" (0.4): probability $0.24$. Beam search with $k = 2$ keeps both A and B after step 1, extends both, and finds B then "end": $0.4 \cdot 0.9 = 0.36$, better than any path through A. Greedy could never find it, because it threw B away at the first step.

Beam search was the standard decoder for machine translation (Sutskever et al., 2014). For open-ended generation it tends to produce bland, repetitive text (section 5), so chat models sample instead; beam search remains useful when one correct output is expected.

Speculative decoding

Generating with a large model is limited by memory bandwidth, not arithmetic: scoring several proposed tokens in one forward pass costs little more than generating one. Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) exploits this. A small draft model proposes $\gamma$ (gamma) tokens, cheaply; the large target model scores all of them in one pass; then the proposals are accepted or rejected so that the output follows exactly the target model's distribution:

1. For each proposed token $x$, in order, accept it with probability $\min\!\left(1, \dfrac{p(x)}{q(x)}\right)$, where $p$ is the target's and $q$ the draft's probability of that token at that position.
2. At the first rejection, stop, and sample a replacement from the residual distribution $p'(x) = \mathrm{norm}\big(\max(0,\ p(x) - q(x))\big)$, where $\mathrm{norm}$ divides by the total so that it sums to 1.
3. If all $\gamma$ proposals are accepted, sample one extra token from the target's next distribution, which the same forward pass already computed.

Example. One position, three tokens. Target $p = (0.2, 0.5, 0.3)$, draft $q = (0.5, 0.3, 0.2)$.

• The draft proposes token 0. It is accepted with probability $\min(1, 0.2/0.5) = 0.4$. The target does not like it as much as the draft does.
• The draft proposes token 1: accepted with probability $\min(1, 0.5/0.3) = 1$.
• If a proposal is rejected, the replacement comes from $\max(0, p - q) = (0, 0.2, 0.1)$, normalized to $(0, 0.667, 0.333)$: only the tokens that the target likes more than the draft.

Why the result is exactly $p$: the probability that the draft proposes token 1 and it is accepted is $0.3 \cdot 1 = 0.3$. The total probability of acceptance is $0.2 + 0.3 + 0.2 = 0.7$, so with probability $0.3$ there is a rejection, and the replacement is token 1 with probability $0.667$: $0.3 \cdot 0.667 = 0.2$. Together $0.3 + 0.2 = 0.5 = p(1)$. The same works for every token.

When the draft agrees with the target the speedup is up to $\gamma + 1$ tokens per target pass, and the output distribution is the same as sampling from the target alone: speculative decoding changes the speed, never the result.

Quantization

A 7B model (7 billion parameters) in BF16 takes 14 GB (gigabytes). Storing weights with fewer bits shrinks it and speeds up memory-bound generation. The simplest scheme is symmetric absmax quantization (as used by Dettmers et al., 2022): for a group of weights, the scale is $s = \max|w| / (2^{b-1} - 1)$, each weight is stored as the integer $\mathrm{round}(w / s)$ in $[-(2^{b-1}-1),\ 2^{b-1}-1]$, and is recovered as $s \cdot q$. Here $b$ is the number of bits. With $b = 8$ the integers go up to 127, with $b = 4$ only up to 7. The rounding error of every weight is at most $s/2$.

Example. A group of 4 weights $w = (0.4, -1.2, 0.05, 0.9)$ with $b = 4$ bits.

1. The largest absolute value is $1.2$ and $q_{\max} = 2^3 - 1 = 7$, so $s = 1.2 / 7 = 0.171$.
2. $w / s = (2.33, -7, 0.29, 5.25)$, rounded: $q = (2, -7, 0, 5)$.
3. Back: $s \cdot q = (0.343, -1.2, 0, 0.857)$. The errors are $(0.057, 0, 0.05, 0.043)$, all below $s/2 = 0.086$.

With $b = 8$ ($q_{\max} = 127$, $s = 0.0094$) the same weights come back as $(0.397, -1.2, 0.047, 0.898)$: much closer.

A single outlier weight makes the scale large and the precision of every other weight in its group poor.

Example. In a group of 8 weights, seven are small, such as $0.1$ or $0.3$, and one is $10$. With 4 bits, $s = 10/7 = 1.43$: every small weight is below $s/2 = 0.71$, so it rounds to 0. The group keeps the outlier and loses everything else.

So 4-bit schemes use small groups (for example 32 or 128 weights per scale) and store one scale per group. The price is a little extra storage.

Example. 7 billion weights at 4 bits, in groups of 128 with a 2-byte scale each: $7 \times 10^9 \cdot 0.5 = 3.5$ GB for the weights plus $5.5 \times 10^7$ scales of 2 bytes, $0.11$ GB: $3.6$ GB in total, against $14$ GB in BF16.

More advanced methods choose the rounding to minimize the error of each layer's output rather than of each weight: GPTQ (Frantar et al., 2023) uses second-order information, AWQ (Activation-aware Weight Quantization; Lin et al., 2024) protects the weights that multiply large activations. QLoRA's 4-bit NormalFloat (Dettmers et al., 2023) places the 16 levels where normally distributed weights are densest.

Pruning

Pruning removes weights. Magnitude pruning (Han et al., 2015) sets the weights with the smallest absolute values to zero and, in its original form, fine-tunes the rest to recover. SparseGPT (Frantar and Alistarh, 2023) prunes very large language models to 50% sparsity in one shot, without retraining. Zeros only save time if the hardware or the kernels can skip them, which is easier with structured patterns (whole rows, or 2 of every 4 weights) than with zeros scattered at random.

Example. The matrix $\begin{pmatrix} 0.5 & -0.1 & 0.3 \\ 0.05 & -0.8 & 0.2 \end{pmatrix}$ pruned to 50% sparsity: of its 6 weights, the 3 with the smallest absolute value ($0.05$, $0.1$ and $0.2$) become 0, leaving $\begin{pmatrix} 0.5 & 0 & 0.3 \\ 0 & -0.8 & 0 \end{pmatrix}$.

Low-rank adapters, merged

Fine-tuning a whole model changes billions of weights. LoRA (Low-Rank Adaptation; Hu et al., 2021) freezes the weight matrix $W$ of shape $(d_{\text{in}}, d_{\text{out}})$ and learns only a small correction: a low-rank update $\Delta W = AB$, with $A$ of shape $(d_{\text{in}}, r)$, $B$ of shape $(r, d_{\text{out}})$ and $r \ll d$, scaled by $\alpha / r$ (the names of $A$ and $B$ follow this page's row-vector convention; the paper writes the transposed product):

where $x$ is the input, $r$ is the rank (a small number such as 8) and $\alpha$ (alpha) is a scaling constant. $B$ starts at zero, so training starts from the unchanged model. Only $A$ and $B$ are trained (section 12 uses it for fine-tuning).

Example. For a $4096 \times 4096$ matrix, $W$ has $16{,}777{,}216$ weights. With $r = 8$, $A$ and $B$ have $2 \cdot 4096 \cdot 8 = 65{,}536$: $0.4\%$ of the matrix.

For inference the update can be merged into the weight once, $W' = W + \frac{\alpha}{r} AB$, so the adapted model runs at exactly the cost of the original one. Many adapters can also be kept separate and swapped on top of the same base model.

Example. $W = \begin{pmatrix} 1 & 2 \\ 3 & 4 \end{pmatrix}$, $A = \begin{pmatrix} 1 \\ 0 \end{pmatrix}$, $B = \begin{pmatrix} 0.5 & 0.5 \end{pmatrix}$, $\alpha = 2$, $r = 1$. Then $AB = \begin{pmatrix} 0.5 & 0.5 \\ 0 & 0 \end{pmatrix}$, and $W' = W + 2 \cdot AB = \begin{pmatrix} 2 & 3 \\ 3 & 4 \end{pmatrix}$. For $x = (1, 1)$: the adapted layer gives $xW + 2 \cdot (xA)B = (4, 6) + (1, 1) = (5, 7)$, and the merged one $xW' = (5, 7)$: the same.

References

• R. Pope et al., Efficiently Scaling Transformer Inference, MLSys, 2023. The KV cache and memory-bound decoding.
• W. Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP, 2023. vLLM.
• I. Sutskever, O. Vinyals, Q. V. Le, Sequence to Sequence Learning with Neural Networks, arXiv, 2014. Left-to-right beam search decoding.
• Y. Leviathan, M. Kalman, Y. Matias, Fast Inference from Transformers via Speculative Decoding, ICML, 2023.
• C. Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, arXiv, 2023.
• T. Dettmers, M. Lewis, Y. Belkada, L. Zettlemoyer, LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, NeurIPS, 2022. Absmax quantization.
• E. Frantar, S. Ashkboos, T. Hoefler, D. Alistarh, GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers, ICLR, 2023.
• J. Lin et al., AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration, MLSys, 2024.
• T. Dettmers, A. Pagnoni, A. Holtzman, L. Zettlemoyer, QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS, 2023.
• S. Han, J. Pool, J. Tran, W. J. Dally, Learning both Weights and Connections for Efficient Neural Networks, NIPS, 2015.
• E. Frantar, D. Alistarh, SparseGPT: Massive Language Models Can be Accurately Pruned in One-Shot, ICML, 2023.
• E. J. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, arXiv, 2021.