KV Cache & FlashAttention

Attention compares every position with every earlier position. During training, and during the first pass over a prompt, the whole sequence is available at once, so all those comparisons happen together.

Generating text is different. The model produces one token at a time, and each new token needs to look back at the keys and values — the vectors that say what earlier positions offer and pass along — of every token before it. Without a shortcut, token number 1,000 would recompute the keys and values of the first 999 tokens, in every layer, at every step. The arithmetic is identical every time.

The KV cache removes the repetition. Keep the keys and values already computed, and append exactly one new key and one new value per layer per token.

The cache size has a closed-form formula: 2 × layers × KV heads × head dim × bytes per element per token. Each factor is plain. Layers: how many repeated blocks the model stacks. KV heads: how many parallel key/value projections it has. Head dim: how many numbers are in each key or value vector. Bytes per element: 2 for the usual 16-bit format. The leading 2 counts keys and values separately.

Work an example. A 7-billion-parameter model with 32 layers, 32 key/value heads, and 128 numbers per head, stored in 16 bits, needs 2 × 32 × 32 × 128 × 2 = 524,288 bytes per token — 512 kilobytes for a single token. At 128,000 tokens that is 64 gigabytes for one conversation, while the model's weights take about 14 gigabytes. At long context the cache is several times larger than the model it serves.

That multiplication is the whole lesson: cache cost grows in a straight line with the number of tokens, and the slope is set by the projection shapes.

Three architecture choices shrink the slope. Grouped-query attention lets several query heads share one key/value projection; at 8 query heads per key/value head, the cache drops 8-fold. Multi-query attention goes all the way to one key/value head for every query head. Multi-head latent attention compresses keys and values into a small shared vector and expands them again when needed, buying 7 to 14 times less memory.

Number format is the other lever. Storing the cache in 8-bit floats instead of 16-bit halves the bytes with under 1 percent accuracy loss on validated setups, which is why it is a common serving default. Together, grouped-query attention and 8-bit storage regularly turn 64 gigabytes of cache into 8.

In the decode panel, watch the cache grow one token at a time. Each append is cheap: two small projections per layer. The real cost is moving bytes. Decoding is memory-bound: every generated token streams the weights and the whole cache through the chip just to compute one new row of attention.

That is why batching raises throughput so much: one read of the cache serves several sequences at once. It is also why a paged allocator matters. Sequences start short and grow, so the cache should be handed out in blocks rather than one large buffer reserved for the worst case.

FlashAttention attacks the other half of the problem: the intermediate table produced while computing attention. The straightforward implementation writes the full n×n table of scores into the chip's main memory (HBM), softens each row, and multiplies by V. Main memory can hold those bytes, but moving them is the bottleneck, and a large intermediate that is written once and read once is exactly what you do not want to push through it.

FlashAttention never stores the table. It slices the queries, keys, and values into blocks that fit in the chip's fast on-board memory (SRAM), computes the scores for one block, folds the result into a running summary, and throws the block away. The final answer is essentially identical to exact attention; only the memory traffic changes.

The running summary is a pair of numbers: m, the largest score seen so far, and l, the sum of exponentials so far. When a new block arrives, its largest score may exceed m. The old sum is then out of date and must be rescaled: l_new = l_old · exp(m_old − m_new) + Σ exp(s_i − m_new), where the sum runs over the scores s_i in the new block.

The partial output is rescaled by the same ratio. Because m only ever grows, a single multiply corrects everything that was computed against the old maximum. Every number a block needs lives in fast registers, so no row ever has to be revisited. That is the online softmax at the heart of every memory-efficient attention kernel.

The trade changes with the phase. Processing a prompt happens all at once, so each weight read pays off across many tokens and the work becomes compute-bound. Tiled attention wins there, because the score table never leaves fast memory. Decoding, one token at a time, is bandwidth-bound on the cache, so the wins come from grouped-query attention, 8-bit cache storage, and dropping entries instead.

A practical serving stack usually runs all three: grouped-query attention in the architecture, 8-bit storage for the cache, and a paged allocator that keeps memory compact as sequences grow. None of them changes what attention computes.

Long-answer reasoning models broke the old assumption that the prompt is the problem. A 2,000-token question can trigger 100,000 tokens of thinking, so the cache keeps growing while the model works. Two families of fixes exist.

Eviction keeps a budget. A sliding window drops the oldest tokens. Attention sinks always keep the first few, which carry outsized weight. Heavy-hitter policies keep the tokens that have received the most attention so far. Compression rewrites entries instead: store fewer bits, merge similar keys, or skip layers. Both families trade a measurable accuracy loss for a fixed memory ceiling.

The problems below compute bytes per token, the cache for a whole sequence, one append step, the online-softmax rescale, and the savings from grouped-query attention and tiled attention. It is all arithmetic, so every number can be checked against the formulas above.