The answer to "why can't I just set batch = 154?"

Lesson 8 · the KV cache · the constraint Lesson 2 deliberately left out

Lesson 2 ended with a promise: batching to ~154 makes FP16 decode compute-bound, and something stops you. This is that something. It also forces an honest correction to Lesson 2's model — at realistic context lengths, the KV cache, not the weights, is most of your decode traffic.

Why the cache exists at all

At decode time you generate one token. Attention for that token needs its query vector against the keys and values of every previous token. Those keys and values are a pure function of tokens you already processed — so recomputing them each step would make generation quadratic for no reason.

So you keep them. That is the entire idea, and it is the same instinct you have about memoizing an expensive pure function. The cost is that the memoization table is enormous, lives in HBM, and grows with every token you emit.

KV cache size: bytes held per token, per sequence.

Listen first: bytes per token equals two — one for keys and one for values — times the number of layers, times the number of key-value heads, times the head dimension, times bytes per element. For Llama 2 7B in half precision that is two times thirty-two layers times thirty-two key-value heads times a head dimension of one hundred twenty-eight times two bytes, which comes to five hundred and twenty-four thousand two hundred eighty-eight bytes — five hundred and twelve kibibytes for a single token of a single sequence. Multiply by a four-thousand-token context and one conversation is holding over two gigabytes of cache.

kv_bytes_per_token = 2 · n_layers · n_kv_heads · head_dim · bytes_per_element
                     ^ K and V

Llama 2 7B, fp16:  2 · 32 · 32 · 128 · 2  =  524,288 B  =  512 KiB  per token
                   × 4096 tokens of context  =  2.1 GB  for ONE sequence

The number that ends the argument

An A100 has 80 GB. Llama 2 7B in fp16 takes 14 GB of it, and you need a few GB for activations and workspace. That leaves roughly 62 GB for cache. At 2.1 GB per sequence:

Maximum batch ≈ 28. Lesson 2 said you needed 154 to stop wasting the GPU's math units. You are not short by a little — you are short by 5×, and no amount of scheduling cleverness changes it. The KV cache is the binding constraint on serving throughput.

Now flip one knob. Drop the context from 4096 to 512 and the same card holds 230 sequences — comfortably past 154. This is why serving systems care so much about context length, and why "supports 128K context" and "serves cheaply" are in direct tension.

The correction to Lesson 2

Lesson 2's simulator counted only weight traffic, and said so. Here is what it was leaving out. Weights are read once per step and shared by the whole batch. The KV cache is read once per step per sequence. So as you batch up, KV traffic grows linearly while weight traffic stays flat:

Decode traffic per step: weights are amortised across the batch, the cache is not.

Listen first: per decode step the weight traffic is just the weight bytes, paid once no matter how large the batch. The cache traffic is the per-token cache size times the context length times the batch size — it scales with every sequence you add. Once you have filled the card, the cache occupies essentially all the free memory, so cache traffic is roughly the whole of HBM and weight traffic is a small fraction of it. For 7B at four thousand context that split is about eighty-one percent cache to nineteen percent weights.

per decode step:
    weight traffic = weight_bytes                     # ONCE, shared by the batch
    cache traffic  = kv_per_token · context · batch   # PER SEQUENCE, grows with batch

fill the card and the cache is ~all your free HBM, so it dominates:
    7B fp16, ctx 4096, batch 28  →  14 GB weights  +  60 GB cache  =  81% cache

Which reorders your optimisation priorities. Quantizing weights from fp16 to INT4 was a 4× win in Lesson 2's model; at long context it saves 10.5 GB out of a 74 GB step — worth having, but the bigger lever has become the cache itself.

What the architects did about it: grouped-query attention

Llama 2 70B has 80 layers to 7B's 32 — but only 8 key-value heads against 64 query heads. Its KV cache is 320 KiB per token, smaller than the 7B model's 512 KiB, despite being ten times the model. That is grouped-query attention: many query heads share one KV head, and the cache shrinks by the sharing factor.

It is worth sitting with that inversion. A 10× larger model with a smaller cache. The architecture was changed specifically to relieve the constraint this lesson is about.

Micro-world: find your real batch ceiling

What fits on one A100 80GB

4 GB is reserved for activations and workspace. Everything else is weights plus cache — and the cache is what runs out.

Model

Context length

Weight precision

KV cache precision

80 GB of HBM, allocated

weights KV cache reserved / unused

Table view

Defaults: 7B, fp16, 4096 context. Max batch 28 against a target of 154. There is the answer to Lesson 2's cliffhanger, in one number.
Drop the context to 512. Batch ceiling jumps past the target and decode finally becomes compute-bound. Short conversations are cheap per token; long ones are not, and now you can say exactly why.
Pick the 70B model with fp16 weights. It does not fit at all — 140 GB of weights on an 80 GB card. Switch to INT4 and it fits with room for 30 sequences. For big models quantization is not an optimisation, it is the price of admission.
Compare 13B and 70B KV-per-token. The 70B model's cache is smaller, because of 8 KV heads instead of 40. That is GQA, and it is the single biggest architectural concession to this constraint.
Quantize the KV cache to INT8 at 4096 context. The batch ceiling doubles. Now compare that to what quantizing the weights bought you at the same context — at long context the cache is the bigger lever, which reverses Lesson 2's advice.

Check yourself

Cold-recall defense (one breath):
Decode attends over every previous token, so K and V are cached instead of recomputed: 2 · layers · kv_heads · head_dim · bytes per token per sequence — 512 KiB/token for 7B fp16, so 2.1 GB for one 4096-token sequence. On an 80 GB card that caps the batch near 28 when Lesson 2 wanted 154: the cache, not the scheduler, is the throughput ceiling. Weights are read once per step and shared; the cache is read per sequence, so once the card is full the cache is ~80% of decode traffic. Levers: shorter context, GQA (70B's 8 KV heads make its cache smaller than 7B's), quantized KV, paged KV.

Where streaming softmax comes back

Lesson 7's online softmax is what makes this survivable on the compute side. Attention scores are a context × context matrix during prefill — at 4096 tokens that is 16M scores per head per layer, and materialising it in HBM would dwarf everything here. Because softmax can be computed with a running max and a rescaled running sum, FlashAttention processes the scores in tiles that stay in on-chip memory and never writes the matrix out at all. The numerical trick and the memory constraint are the same story told twice.

Primary source

Read Efficient Memory Management for Large Language Model Serving with PagedAttention (the vLLM paper) §3 — the section measuring how much KV memory real serving systems waste to fragmentation and over-reservation (~15 min). It is the paper that turned this constraint into an operating-system problem. For grouped-query attention itself, the Llama 2 paper §2.2 is the primary reference.

💬 I'm your teacher for this — ask me followups any time. The obvious sequel: "how does paged attention actually reclaim that memory?" — that is vLLM's whole contribution and would make a good next lesson.