KV cache, paged, flash and sparse attention

The Saturday before Diwali, Crumb’s WhatsApp assistant fell over. Since the summer it had run on Crumb’s own GPU server (one 24 GB graphics card hosting an open 8-billion-parameter Llama model), because Meera did not want customers’ addresses and phone numbers leaving Crumb’s machines. On a normal day it handled eight or ten conversations at once with replies in about two seconds.

At 11 a.m. on Saturday, forty customers were chatting at once about hampers. Replies took 25 seconds, then the server logged CUDA out of memory and dropped every conversation in progress.

“The model is 16 GB,” Arjun said, staring at the dashboard. “The card is 24. We have 8 GB spare. How is a few kilobytes of chat text running us out of memory?”

“It isn’t the text,” said Priya. “It’s what the model remembers about the text so it doesn’t have to reread it. Today we open up how a model actually runs, token by token, and where the memory and time really go.”

Why writing one token at a time is wasteful

Recall how a decoder-only model writes, from Module 30. It reads all the tokens so far, runs them through every layer, and the last position’s output gives a probability for every possible next token (a chunk of text, roughly three quarters of a word). It picks one, appends it, and goes again. Each loop is a full forward pass through the network.

Inside each layer, self-attention gives every token a query (what am I looking for?), a key (what do I offer?) and a value (what do I hand over?). The newest token’s query is compared with the keys of every earlier token, and their values are blended.

Here is the waste. When the model writes token 501, it needs the keys and values of tokens 1 to 500. When it writes token 502, it needs the keys and values of tokens 1 to 501. Tokens 1 to 500 have not changed. Thanks to the causal mask (each token only sees earlier ones), their keys and values never change. Yet a naive loop recomputes all of them, every step.

Arjun’s first attempt: recompute everything

Arjun had written exactly that loop for his tiny GPT in Module 30: feed the whole sequence in, take the last position’s prediction, append, repeat. To see the cost on a real model, he timed GPT-2 (a small 124-million-parameter model) writing 200 tokens after a 261-token prompt, on his laptop, with Hugging Face’s generate and the cache switched off:

import time, torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()
ids = tok("Dear customer, thank you for ordering from Crumb and Co. " * 20, return_tensors="pt")

for use_cache in (True, False):
    start = time.perf_counter()
    with torch.no_grad():
        model.generate(**ids, max_new_tokens=200, min_new_tokens=200, do_sample=False,
                       use_cache=use_cache, pad_token_id=tok.eos_token_id)
    print(f"use_cache={use_cache}: {time.perf_counter() - start:.1f} s")

use_cache=True: 4.2 s
use_cache=False: 31.3 s

Same model, same output, seven and a half times slower without the cache. And the gap grows with length.

Count the work. With a 1,000-token prompt and 200 new tokens, the naive loop runs tokens through the network $1000+1001+\cdots +1199\approx 220,000$ times. With a cache, it runs the 1,000 prompt tokens once and then each new token once: 1,200. That is about 180 times less work for the expensive layers.

Priya’s order book

“Picture Shanti at the Deccan counter taking a long phone order,” Priya said. “Every time the customer adds an item, does she re-read the entire order from the top before writing the next line?”

“No, she’s written it down. She glances at her notes.”

“That notebook is the KV cache. For each token, each layer computes a key and a value once and writes them down. When the next token arrives, it computes only its own query, key and value, adds its key and value to the notebook, and compares its query against everything in the notebook.”

“So the notebook grows by one line per token.”

“Per token, per layer. And that’s your out-of-memory error. Forty long conversations means forty very thick notebooks.”

Following one new token through the cache

Take the sentence so far, “Your cake is”, and the model about to write “late”. Look at one attention layer.

Before this step, the cache for this layer holds three keys and three values, computed when those tokens were processed:

Cached	Key	Value
Your	$k_1$	$v_1$
cake	$k_2$	$v_2$
is	$k_3$	$v_3$

Step 1. The newest token, is, has been processed a moment ago. Its query $q_3$ is compared with $k_1,k_2,k_3$ from the cache. Three dot products, a softmax, a blend of $v_1,v_2,v_3$. Nothing about Your or cake is recomputed.

Step 2. The model’s output predicts late. Now late goes through the layers as a single token. It produces $q_4$, $k_4$ and $v_4$. The cache appends $k_4$ and $v_4$ (four rows now), and $q_4$ attends over all four.

Step 3. Repeat for every layer, for every new token.

In code, for one attention head, with and without the notebook:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
d = 16
W_q, W_k, W_v = (torch.randn(d, d) / d**0.5 for _ in range(3))

def attend(q, K, V):
    """One query attending over all the keys and values it may see."""
    return F.softmax(q @ K.T / d**0.5, dim=-1) @ V

x = torch.randn(6, d)                       # embeddings for a 6-token sequence

# Without a cache: recompute every key and value of the prefix at each step
outs_no_cache, rows_no_cache = [], 0
for t in range(1, 7):
    K, V = x[:t] @ W_k, x[:t] @ W_v
    rows_no_cache += t
    outs_no_cache.append(attend(x[t - 1] @ W_q, K, V))

# With a cache: compute each token's key and value once, append, reuse
K_cache, V_cache, outs_cache = [], [], []
for t in range(1, 7):
    new = x[t - 1]
    K_cache.append(new @ W_k)               # only the newest token's key...
    V_cache.append(new @ W_v)               # ...and value are computed
    outs_cache.append(attend(new @ W_q, torch.stack(K_cache), torch.stack(V_cache)))

print("same outputs:", torch.allclose(torch.stack(outs_no_cache), torch.stack(outs_cache)))
print("key/value rows computed:", rows_no_cache, "without cache,", len(K_cache), "with")

same outputs: True
key/value rows computed: 21 without cache, 6 with

Identical answers, a fraction of the work. For 6 tokens the saving is 21 against 6; for 1,200 tokens it is about 720,000 against 1,200.

This also explains two numbers you have seen on every provider’s pricing page and latency chart since Module 21. Processing the prompt, called prefill, handles all prompt tokens in one big parallel pass and fills the cache. Writing the reply, called decode, then goes one token at a time, each step reading the whole cache. Prefill is limited by raw arithmetic; decode is limited by how fast the GPU can read the cache and weights from memory. That is why the first token takes a moment and then the rest stream steadily.

Opening the hood: how big is the notebook?

For every token, every layer stores one key and one value per KV head, each $d_{head}$ numbers long, at some number of bytes per number. So:

$$\text{KV bytes per token}=2\times n_{layers}\times n_{kvheads}\times d_{head}\times \text{bytes per number}$$

The 2 is for keys plus values. Crumb’s model is in the Llama 3 8B family: 32 layers, 8 KV heads of 128 numbers each, stored as 16-bit numbers (2 bytes):

$$2\times 32\times 8\times 128\times 2=131,072\text{ bytes}=128\text{ KB per token}$$

Now Arjun’s Saturday:

• A 2,000-token conversation (system prompt, menu, chat history): $2000\times 128\text{ KB}\approx 250$ MB of cache.
• 8 GB spare on the card, minus a working margin, holds about 25 such conversations.
• Forty customers at once: out of memory.

Notice the “8 KV heads”. The model has 32 query heads but shares each key-value head among four of them, a design called grouped-query attention (GQA). An older design with 32 separate KV heads, like Llama 2 7B, needs four times as much: 512 KB per token. GQA exists almost entirely to shrink this notebook. Keep that in mind; the DeepSeek case study later in this module takes the idea much further.

Paged attention: stop reserving empty tables

Arjun’s first fix was to cap conversations. His serving code reserved cache space for each conversation’s maximum length, 8,192 tokens (1 GB), up front, so a chat could never run out mid-reply. Most chats used under 1,000 tokens. Most of each reservation sat empty.

“That’s a restaurant that seats every party at a twelve-person table, in case more friends show up,” said Priya. “You turn away customers with half the chairs empty.”

Two kinds of waste were happening:

• Over-reservation: space held for tokens that never arrived.
• Fragmentation: as chats of different lengths start and finish, free memory ends up in scattered gaps too small for a new full-size reservation.

Paged attention, introduced by the vLLM project in 2023, borrows the trick operating systems use for computer memory. Split the cache into small fixed-size blocks, say 16 tokens each. Give each conversation a block table, a list of which blocks hold its tokens, in order. Blocks need not sit next to each other in memory. Hand out a new block only when a conversation fills its last one.

Now the only waste is the unfilled part of each conversation’s final block: at most 15 tokens. The vLLM authors reported that earlier serving systems wasted roughly 60 to 80% of their KV cache memory this way, against under 4% with paging, and that fitting more conversations into the same memory raised throughput severalfold in their benchmarks.

There is a bonus. Every Crumb conversation starts with the same 600-token system prompt. With block tables, many conversations can point at the same physical blocks for that shared beginning, copying only if one of them would change a block. That is called prefix sharing, and it is closely related to the prompt caching you used through provider APIs in Module 21.

Arjun moved the assistant to vLLM, which uses paged attention by default. The same card now held over 60 typical conversations.

Flash attention: stop carrying everything back and forth

Paging solved how much memory. Long conversations raised a second problem: time. Attention’s score grid is $n\times n$ for $n$ tokens, per head, per layer. At 32,000 tokens that grid has about a billion entries; at 2 bytes each, that is 2 GB, for one head in one layer.

The surprising part is that the arithmetic is not the slow bit. Moving data is. A GPU has two kinds of memory:

• HBM (high-bandwidth memory): large, tens of gigabytes, but each trip to it is comparatively slow.
• SRAM: tiny, a few hundred kilobytes next to each group of compute cores, but extremely fast.

Standard attention computes the whole score grid, writes it out to HBM, reads it back to apply softmax, writes the weights out, and reads them back again to blend the values. The GPU spends much of its time waiting on those round trips.

“It’s the kitchen,” said Priya. “The walk-in cold store is huge but far away. The counter is tiny but right there. Standard attention carries every ingredient out of the cold store, puts it back, fetches it again for the next step. Flash attention brings one tray at a time, does every step for that tray on the counter, and keeps only a running tally.”

Flash attention (Tri Dao and colleagues, 2022) splits $Q$, $K$ and $V$ into tiles small enough to fit in SRAM, and processes them tile by tile without ever writing the full score grid to HBM. The hard part is softmax, which normally needs the whole row before it can divide by the total. The trick is an online softmax: keep a running maximum and running total per row, and rescale as each new tile arrives.

Follow one row of four scores, [1, 2, 3, 0], arriving in two tiles. (Subtracting the maximum before exponentiating is standard practice, to keep numbers from overflowing.)

• Tile 1, scores [1, 2]: running max $m=2$. Running total $\ell =e^{1-2}+e^{2-2}=0.368+1=1.368$.
• Tile 2, scores [3, 0]: the new max is 3. The old total was measured against max 2, so rescale it by $e^{2-3}=0.368$: $1.368\times 0.368=0.503$. Add the new tile: $e^{3-3}+e^{0-3}=1+0.050$. New total $\ell =1.553$.

Computing it all at once against the max of 3 gives $e^{-2}+e^{-1}+e^0+e^{-3}=1.553$. Identical. The blended output is rescaled the same way as it accumulates. So flash attention is exact, not an approximation, and its memory use grows linearly with length instead of with the square. Its authors reported several-fold speedups on long sequences, and later versions (FlashAttention-2 and 3) went further on newer GPUs.

You rarely write this yourself. In PyTorch, one function call picks the fastest available kernel, including flash attention on supported GPUs:

import torch
import torch.nn.functional as F

B, H, T, D = 1, 8, 1024, 64                  # batch, heads, tokens, head size
q, k, v = (torch.randn(B, H, T, D) for _ in range(3))

out = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # mask built in
print(out.shape)                             # torch.Size([1, 8, 1024, 64])

Compare that one line with the scores-mask-softmax-multiply code you wrote in Module 30. Same maths; a kernel that respects the memory hierarchy. vLLM, Hugging Face Transformers and most serving stacks use these kernels automatically when the hardware supports them.

Sparse attention: stop looking at everything

Flash attention computes full attention efficiently. Sparse attention changes the question: does every token need to look at every other token?

Mostly, words care about their neighbourhood. Sparse attention lets each token attend to a chosen subset instead of all earlier tokens. Common patterns:

• Sliding window: each token sees only the previous $w$ tokens, say 4,096. Cost per token is fixed, so total cost grows linearly with length. Mistral 7B used a sliding window, and some recent models interleave sliding-window layers with occasional full-attention layers.
• Global tokens: a few special positions (such as the start of the text) see everything and are seen by everything, acting as hubs.
• Strided patterns: each token sees its neighbours plus every 64th token further back, so information can hop across long distances in a few layers.

Information still travels beyond the window, indirectly. With a window of 4,096 and 32 layers, token 100,000 can be influenced by token 1 through a chain of hops, the way news spreads along a queue of people each talking only to their neighbours. But indirect is not the same as direct. Facts far away can get blurred, which is the trade-off: sparse attention saves compute and cache memory at some risk to long-range recall. Unlike flash attention, it is an approximation of full attention, so quality must be measured, not assumed.

Where it breaks

Prefill spikes. A customer who pastes a 30-page PDF triggers a huge prefill that briefly hogs the GPU and slows everyone else’s replies. Serving systems split long prefills into chunks and interleave them with other users’ decode steps.

Cache eviction. When memory runs out, a server must pause or drop some conversations. vLLM can pre-empt a request and recompute its cache later. Arjun’s original setup crashed, which was the real Saturday failure.

Quality versus speed. Paged and flash attention are exact; the model’s answers do not change. Sparse attention and cache compression are approximations whose quality must be checked on your own tasks, with the evaluation habits from Module 23.

Try it yourself

1. Using the formula, compute the KV cache per token for a model with 28 layers, 4 KV heads, a head size of 128 and 16-bit numbers. How many 4,000-token conversations fit in 10 GB?
2. Run the GPT-2 timing script with max_new_tokens=50 and then 400. Does the speedup from caching grow or shrink with length?

Show the answer

1. $2\times 28\times 4\times 128\times 2=57,344$ bytes, about 56 KB per token. A 4,000-token conversation needs about 230 MB, so 10 GB holds roughly 43 of them.
2. It grows. Without a cache, each new token reprocesses an ever-longer prefix, so total work grows roughly with the square of the length; with a cache it grows roughly linearly. Longer outputs widen the gap.

Common mistakes

• Thinking the KV cache stores text. It stores keys and values for every token in every layer: hundreds of kilobytes per token in a mid-sized model, far more than the text itself.
• Assuming free memory equals capacity. Weights, the cache and working buffers all compete. Leave headroom, and cap context length and concurrency.
• Confusing flash attention with sparse attention. Flash attention computes exact full attention faster. Sparse attention deliberately skips some pairs and may change answers.
• Pre-allocating the maximum length per request. Use a server with paged attention, such as vLLM, instead of rolling your own.
• Ignoring prefill. Long prompts cost time before the first token appears, and slow other users down too.

Recap

• Generating text recomputes nothing about past tokens if their keys and values are cached; the KV cache turns quadratic work into roughly linear work.
• Cache size per token is $2\times$ layers $\times$ KV heads $\times$ head size $\times$ bytes; at scale, the cache, not the weights, often limits concurrency.
• Paged attention stores the cache in small blocks with a block table, wasting almost nothing and letting conversations share common prefixes.
• Flash attention computes exact attention tile by tile in fast on-chip memory with an online softmax, never writing the full score grid out.
• Sparse attention lets each token look at a subset of others, trading some long-range accuracy for linear cost.

On Monday, Meera asked what had happened on Saturday. Arjun had his answer ready.

“The model keeps a notebook for every conversation, so it doesn’t reread the whole chat for each word it writes. Those notebooks are big, about 128 KB per word, and we were reserving a full-size notebook for every customer even when they only wrote two lines. Forty customers, forty fat empty notebooks, no room left. Now the server hands out pages as conversations grow, like a binder. Same card, more than sixty conversations.”

“So we don’t need a bigger computer?”

“Not for this Diwali. But there’s another lever.” He glanced at Priya. “The model itself takes 16 GB. If we could store each of its numbers in fewer bits, it would shrink, and so would the bill.”

“Quantisation,” said Priya. “Next lesson. And a way to teach a small model to imitate a big one.”

Go deeper

• vLLM documentation, including how PagedAttention and prefix caching work.
• PyTorch tutorials, including scaled dot-product attention.
• Attention Is All You Need, for the attention these tricks accelerate.
• Hugging Face docs, see the Transformers guides on caching and generation.