Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization

Two phases, two bottlenecks

Generation (7.6 Decoder-only: GPT and autoregressive generation) has a prefill phase — the prompt's $n$ tokens run through the model in one parallel pass, a big matrix multiplication that keeps GPU compute busy — and a decode phase that produces one token per step. Decode is a different beast: for each token the model must read all its weights from GPU memory (and the whole KV cache) to do only a small amount of arithmetic. Decode is therefore memory-bandwidth-bound, not compute-bound, and most of inference engineering is about doing more useful work per byte read.

The KV cache

Without caching, step $t$ recomputes keys and values for all $t$ previous tokens — $O(t)$ redundant work per step, $O(T^2)$ over a generation. Caching every layer's $\mathbf{K}$ and $\mathbf{V}$ makes each step compute only the new token's projections and attend over the cache. The price is memory:

where the 2 counts keys and values, $L$ is the number of layers, $n_{kv}$ the number of key/value heads, $d_{\text{head}}$ the head size and $b$ the bytes per element (2 for fp16/bf16).

Worked example. Llama 2 7B: $L=32$, $n_{kv}=32$, $d_{\text{head}}=128$, fp16 → $2\cdot 32\cdot 32\cdot 128\cdot 2 = 524{,}288$ bytes = 0.5 MB per token. A 4,096-token context costs 2 GB; a batch of 16 such sequences costs 32 GB — more than the 13.5 GB of weights. Llama 3 8B uses GQA with $n_{kv}=8$: 128 KB per token, 4× less, so a 128k-token context fits in 16 GB.

Arithmetic intensity: why batching matters

A decode step for a batch of $B$ sequences performs about $2NB$ FLOPs and reads about $2N$ bytes of fp16 weights (plus the KV cache). Its arithmetic intensity — FLOPs per byte — is therefore about $B$. An H100 can do roughly 990 TFLOP/s of dense bf16 math and move 3.35 TB/s from memory, a ridge point near 300 FLOPs per byte. Below that, the chip waits on memory: at batch size 1 you use well under 1% of its compute. Reading 14 GB of weights at 3.35 TB/s takes 4.2 ms, so a 7B fp16 model can produce at most about 240 tokens/s for a single user — and the same 4.2 ms could serve 200 users. Serving systems therefore batch aggressively, with continuous batching (admitting new requests at every step as others finish) and the PagedAttention scheme of vLLM, which stores the KV cache in fixed-size blocks like virtual-memory pages so that memory is not wasted on pre-allocated maximum lengths.

Shrinking the cache: MQA and GQA

Multi-query attention shares one key/value head across all query heads; grouped-query attention (GQA) is the middle ground, sharing each K/V head among a group of query heads (Llama 2 70B and all Llama 3 models: 8 KV heads for 32–128 query heads). The cache shrinks by $n_{\text{heads}}/n_{kv}$ with a small quality cost, and the K/V projection matrices shrink too. Checkpoints trained with full multi-head attention can be converted to GQA with brief "uptraining".

FlashAttention: exact attention without the $n\times n$ matrix

Standard attention writes the $n\times n$ score matrix to GPU memory, reads it back for the softmax, writes the probabilities, reads them again to multiply by $\mathbf{V}$ — memory traffic in $O(n^2)$. FlashAttention computes the same result by tiling: it loads blocks of $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ into fast on-chip SRAM and uses the online softmax trick — maintaining a running maximum $m$ and running sum $\ell$ so that partial results from different key blocks can be rescaled and combined exactly:

The full matrix is never materialized, memory becomes $O(n)$, and wall-clock speedups of 2–4× follow because attention is memory-bound. FlashAttention-2 improved parallelism and work partitioning; FlashAttention-3 exploits the asynchronous units of Hopper GPUs. It is an exact algorithm — no approximation — and it is what makes long contexts trainable at all.

Speculative decoding

Decode is bandwidth-bound, so verifying several tokens in one forward pass costs about the same as generating one. Speculative decoding uses a small draft model to propose $\gamma$ tokens, then the target model scores all of them in a single pass. Each proposed token $x$ is accepted with probability $\min\big(1,\ p(x)/q(x)\big)$, where $p$ is the target's distribution and $q$ the draft's; on the first rejection a replacement is sampled from the normalized $\max(0,\ p-q)$. This procedure provably reproduces the target's exact output distribution. If each token is accepted with probability $\alpha$, the expected number of tokens produced per target pass is

so $\alpha=0.8$, $\gamma=4$ gives $3.36$ tokens per pass — a 2–3× speedup with no change in output quality. Variants draft from the model's own early layers (Medusa, EAGLE) or from n-gram lookup in the prompt.

Quantization: fewer bytes per weight

If bandwidth is the bottleneck, halve the bytes. Quantization stores weights (and sometimes activations or the KV cache) in low precision. Uniform integer quantization maps a tensor with scale $s$ and zero point $z$:

Per-channel or per-group scales (for example groups of 128 weights) keep the rounding error small. LLM.int8() handles the few outlier feature dimensions in fp16 and the rest in int8; GPTQ quantizes to 4 bits one column at a time while compensating the error using second-order (Hessian) information; AWQ rescales channels by activation statistics before rounding; NF4 is the 4-bit normal-float format used by QLoRA (8.2 Fine-tuning and PEFT). Memory for a 7B model: 14 GB in fp16, 7 GB in int8, about 3.5–4 GB in 4-bit — which is why a 7B model runs on a laptop. Perplexity typically rises very little at 8 bits and modestly at 4 bits; below 4 bits quality drops quickly, although ternary "1.58-bit" training is an active research line.

Putting it together

A single-user latency estimate: time per token ≈ (weight bytes + KV bytes) / memory bandwidth. Every row of the table attacks one of those terms or the number of sequential steps.

Where this shows up later

• The 4-bit NF4 quantization is half of QLoRA — 8.2 Fine-tuning and PEFT; serving stacks (vLLM, TensorRT-LLM, llama.cpp) appear in the LLM project — 8.13 Project: fine-tune a small LLM with LoRA and ship a chat app.
• Latency and cost budgets for RAG and agent systems are set by these numbers — 9.9 Production: latency, cost, caching, guardrails, security and privacy, 10.11 Building production agents.
• Sliding-window and sub-quadratic attention are the subject of 7.12 Beyond vanilla attention: state space models, long context and hybrids; the hardware vocabulary (bandwidth, FLOP/s, memory hierarchy) is 4.12 Compute, hardware and scale basics.