Attention and the KV cache, visually

One layer generating token n + 1 with n tokens already in context. With a cache the step projects one new Query, Key and Value, reads the n cached Key/Value pairs, scores the Query against all n + 1 Keys, and appends the new pair. That is "read five, compute one" when n = 5. Switch to Without a cache and every Key and Value is projected again, n of them for nothing. Step through with ◀ ▶ or Play.

The chart adds this up over a whole generation of N tokens. Without a cache the K/V projections total 1 + 2 + … + N = N(N+1)/2, which is O(N²). With one they total N. The table gives the ratio, (N+1)/2. The cache does not remove the quadratic attention: step i still reads i − 1 cached pairs. That is why decode is limited by memory bandwidth.

Pick a model and an attention layout. The diagram shows which query heads share which K/V heads, and the cache size updates as you change things:

Bytes per token = 2 × layers × KV heads × head dimension × bytes per value (MLA: layers × latent × bytes). For Llama 3 70B in FP16 that is 2 × 80 × 8 × 128 × 2 = 327,680 bytes = 320 KiB = 0.3125 MiB ≈ 0.328 MB. At 8,192 tokens that is 2.68 GB (2.5 GiB), and at 131,072 tokens 42.9 GB (40 GiB). Full MHA would be 8 times larger.

Every size is shown in both units. Decimal units: kB, MB, GB, TB = 10³, 10⁶, 10⁹, 10¹² bytes. Binary units: KiB, MiB, GiB, TiB = 2¹⁰, 2²⁰, 2³⁰, 2⁴⁰ bytes. The book’s "0.31 MB per token", "2.5 GB" and "40 GB" are binary figures (MiB, GiB) labelled as MB and GB. Exactly 8,000 and 128,000 tokens give 2.62 GB and 41.9 GB.

The same workload runs through two allocators over the same pool of KV-cache slots, shown side by side:

Each cell is one block. Strong colour means stored tokens, pale colour means reserved but empty. The table counts running and waiting requests, slots held against slots actually written, the waste, and how many decode steps it takes to finish everything. A shared prefix, such as a common system prompt, is stored once in purple blocks that every block table points at. Copying prefixes is what vLLM’s automatic prefix caching and SGLang’s RadixAttention avoid. When the paged pool runs out, the newest request is preempted and recomputed later, as vLLM does.

A toy attention head (N tokens, head dimension d, random Q, K and V from a seed) worked tile by tile in FlashAttention-2 order. The outer loop runs over blocks of Br query rows, and the inner loop over blocks of Bc key/value columns. The score-matrix picture shows the tile currently in on-chip SRAM. Earlier tiles have been folded into the result and thrown away; none of them is ever written to HBM.

The table traces one query row through the online softmax. Each row keeps a running maximum m, a running sum ℓ and an unnormalised output O. When a block raises the maximum, the sums so far are multiplied by e^(m_old − m_new)^. After the last block the result matches standard attention, apart from about 10⁻¹⁶ of rounding difference.

The second panel compares HBM traffic for one head at real size. Unfused standard attention moves 4Nd + 4N² elements: it writes the score matrix, reads it back, writes the probabilities and reads them again. FlashAttention-2 moves 2Nd + 2Nd·⌈N/Br⌉ + N. Both are still quadratic in N; the paper’s bound is Θ(N²d²/M). FlashAttention’s traffic is smaller by about 2·Br/d, and the extra memory falls from N² to N.

KV-cache size from 1,000 to 10 million tokens on a log–log chart, against the memory of one GPU, an 8-GPU node and an NVL72 rack (≈13.5 TB). The compression options stack:

The tool reports whether the cache fits in the memory you chose beside the weights, how many such sequences fit, and how long one decode step takes to read the cache.

Prefill attention costs 4N²d FLOPs per layer per head: 2N²d for QKᵀ and 2N²d for PV, or half that with a causal mask. At 1M tokens and d = 128 this is 5.1 × 10¹⁴. The tool also gives the whole-model total, the context length at which attention FLOPs overtake the 2 × parameters × N of the weight matmuls (about 54,000 tokens for Llama 3 70B), and the time at the target’s dense peak.

The memory and throughput figures are in one table on the tab, with a vendor source for each row, as of October 2026. Throughput is dense BF16; NVIDIA’s headline numbers assume 2:4 sparsity and are twice as high.

Related: softmax and attention (the attention equation as a heat map), AI hardware calculator (weights plus cache against GPU memory, the roofline, decode speed and cost), floating-point inspector (FP8, INT4 and block-scaled formats).

Companion to the KV cache, FlashAttention and PagedAttention section (Part III), Attention and the transformer, mapped to hardware (Part IX), the 70B worked example (Part IV) and Long-context inference of _AI Acceleration: From Transistors to Trillion-Parameter Models.

Links can be written by hand: /tools/attention-kv/?<parameters>#<tab>, where the tab is decode, heads, paged, flash or long. An example= link with no #tab opens on that example’s tab. Share link builds a link from whatever is on screen. Token counts in ctx, N and seq accept 131072, 128k (= 128,000) and 1M (= 1,000,000). A context window called "128K" is usually 131,072 tokens, so write that number when you mean it. No parameter needs a `. In a link a ` means a space, but the token fields ignore spaces, so N=1e+6 still reads as 1,000,000.

Presets for example=: