Lesson 13.1 · 27 min

LLM Inference Optimization: The Full Landscape

Why does a chatbot that fits easily on one GPU suddenly run out of memory when eight people paste in long documents at the same time?

In short: An LLM writes one token at a time, and to avoid redoing work it stores a Key and a Value vector for every past token in every layer: the KV cache. That cache grows with context length and with the number of users, and it quickly becomes the main memory cost of serving. This lesson shows where the memory goes and walks through the four big families of KV cache compression (quantization, token eviction, sharing keys and values across heads, and low-rank compression) and when to pick each.

What an LLM is and how it writes text

A large language model (LLM) is a neural network trained to answer one question over and over: given the text so far, what is the next token? A token is a small piece of text, often a word or part of a word. “unbelievable” might be split into “un”, “believ” and “able”.

To write a reply, the model runs in a loop. It reads the prompt, predicts a probability for every token in its vocabulary, picks one, appends it to the text, and runs again. This loop is called autoregressive generation: each new token depends on all the tokens before it. A 300-token answer needs 300 trips through the model.

Think of it like a writer who keeps notes Imagine a writer who, before adding each new word, must glance back over everything written so far. If they had to re-read the whole document from scratch each time, a long letter would take forever. If they keep good notes about each earlier word, each glance is quick. The KV cache is those notes; inference optimization is largely about keeping the notes small and the glances fast.

Throughout this lesson we will follow one running example: a support chatbot built on an 8-billion-parameter model. Customers paste long logs and contracts, so prompts of 30,000 tokens are common, and we want to serve many customers on one GPU.

Inference means using a trained model to produce outputs, as opposed to training it. Inference optimization is the set of tricks that make this cheaper and faster: smaller numbers (quantization), smarter memory management, batching users together, and guessing tokens ahead of time. Later lessons in this module cover batching, paged memory and speculative decoding. This lesson focuses on the single biggest memory consumer during serving, the KV cache, and the ways we shrink it.

What attention does

Inside each Transformer layer, attention is the step where a token gathers information from the tokens before it. Every token is turned into three vectors by multiplying it with three learned weight matrices:

• Query (Q): what this token is looking for.
• Key (K): a label describing what this token offers, used to match against queries.
• Value (V): the actual content this token passes on if it is chosen.

Most models also split attention into several heads. Each head has its own smaller Q, K and V and can learn a different pattern, for example one head tracks the previous token and another tracks the subject of the sentence. A model with 32 heads and a head size of 128 has 32 separate K and V vectors per token per layer.

What the KV cache is

Here is the key observation. When the model generates token number 501, the keys and values of tokens 1 to 500 are exactly the same as they were in the previous step, because a causal model never lets earlier tokens look at later ones. Recomputing them would be pure waste.

So we store them. The KV cache is a buffer that holds the key and value vectors of every token processed so far, for every layer and every head. At each new step the model only computes Q, K and V for the newest token, appends its K and V to the cache, and lets the new query attend over the whole cache.

One decoding step with a KV cache

1. Embed the new token: Only the token generated in the previous step enters the model, not the whole text.
2. Project to Q, K, V: In each layer, multiply the token's vector by the three weight matrices to get its query, key and value.
3. Append to the cache: Write the new K and V to the end of this layer's cache. The cache grows by one row per layer per step.
4. Attend over the cache: Score the new query against every cached key, softmax, and take the weighted sum of cached values.
5. Predict the next token: After the last layer, turn the result into probabilities and pick a token. Repeat.

Pause and think: Pause and predict: if we double the number of layers but keep everything else equal, what happens to the KV cache size?

It doubles. Every layer keeps its own keys and values for every token, so the cache is proportional to the number of layers.

Why the KV cache becomes huge

The cache size follows a simple multiplication. Every factor in it is large for modern models and real workloads.

Plug in our chatbot's 8B-class shape: 32 layers, 32 KV heads, head size 128, FP16. Per token that is 2 × 32 × 32 × 128 × 2 bytes = 524,288 bytes, or 512 KiB per token. A single 32,000-token conversation needs about 15.6 GiB. Eight such users need about 125 GiB, more than the memory of a single 80 GB GPU, and the model weights (about 16 GB in FP16) still have to fit too.

The memory problem also becomes a speed problem. During generation, each step must read the entire cache from GPU memory to compute attention. More cache bytes means more time moving data, so a smaller cache often makes each token faster as well as letting more users share the GPU.

What KV cache compression is

KV cache compression is any technique that stores the keys and values in fewer bytes while keeping the model's answers as close as possible to the uncompressed version. Look at the size formula again: each factor is a lever we can pull.

Some of these can be applied to an existing model at serving time (quantization, many eviction methods). Others change the architecture and must be built in when the model is trained or adapted (head sharing, most low-rank schemes). That difference matters a lot in practice, because most teams serve models they did not train.

The four approaches

Approaches stack These families multiply. A GQA model with 8 KV heads and an FP8 cache already uses 8× less KV memory than a 32-head FP16 model, before any eviction.

Code: a KV cache budget for our chatbot

Let us compute the cache for eight users at 32,000 tokens each, and apply each compression idea in turn. The numbers come straight from the size formula.

kv_budget.py

# KV cache size for one long chat, and what each compression idea saves
layers, q_heads, head_dim = 32, 32, 128   # a 7B/8B-class model shape
context, batch = 32_000, 8                # 32k tokens, 8 users at once
def kv_bytes(kv_heads, bytes_per_num, tokens, dim=head_dim):
# 2 = one Key + one Value vector per head, per layer, per token
return 2 * layers * kv_heads * dim * bytes_per_num * tokens * batch
GB = 1024**3
plans = {
"baseline FP16":              kv_bytes(32, 2,   context),
"quantize to INT8":           kv_bytes(32, 1,   context),
"quantize to INT4":           kv_bytes(32, 0.5, context),
"evict: keep 4k tokens":      kv_bytes(32, 2,   4_000),
"share heads: GQA 8 KV":      kv_bytes(8,  2,   context),
"low-rank: 128 -> 32 dims":   kv_bytes(32, 2,   context, dim=32),
"GQA 8 + INT8 combined":      kv_bytes(8,  1,   context),
}
base = plans["baseline FP16"]
per_token = kv_bytes(32, 2, 1) / batch
print(f"KV bytes per token (FP16): {per_token/1024:.0f} KiB")
for name, b in plans.items():
print(f"{name:26s} {b/GB:6.1f} GB  ({base/b:4.1f}x smaller)")

Output:

KV bytes per token (FP16): 512 KiB
baseline FP16               125.0 GB  ( 1.0x smaller)
quantize to INT8             62.5 GB  ( 2.0x smaller)
quantize to INT4             31.2 GB  ( 4.0x smaller)
evict: keep 4k tokens        15.6 GB  ( 8.0x smaller)
share heads: GQA 8 KV        31.2 GB  ( 4.0x smaller)
low-rank: 128 -> 32 dims     31.2 GB  ( 4.0x smaller)
GQA 8 + INT8 combined        15.6 GB  ( 8.0x smaller)

The baseline needs 125 GiB, which does not fit on one 80 GB GPU. GQA plus INT8 brings it to about 16 GiB, which fits comfortably next to the 16 GB of weights. The memory arithmetic is exact, but the quality cost of each plan is not visible here: that is what the next lesson measures.

Pause and think: Our calculation says eviction to 4k tokens and GQA+INT8 both give 8× savings. Are they equally good choices?

No. GQA+INT8 keeps information about every token, slightly blurred. Eviction to 4k permanently drops 28k tokens per user, so any question about the dropped part of the document can no longer be answered accurately. Same bytes, very different risk.

Comparison of the approaches

When to use which one

• You serve someone else's model and need memory now: turn on KV cache quantization (FP8 or INT8). It is the lowest-risk change.
• You are choosing a model: prefer one with GQA or MLA. You get most of the savings for free.
• You have endless streams (a voice assistant that never resets, a log monitor): use eviction with attention sinks plus a recent window, and accept that old details fade.
• Your app depends on exact recall from long documents (contracts, code): avoid aggressive eviction; quantize instead and buy more memory or use paged memory management.

Our support chatbot We pick a GQA model (8 KV heads), enable an FP8 KV cache in the serving engine, and skip eviction because customers often ask about details buried deep in their pasted logs. Memory for eight 32k conversations drops from about 125 GiB to about 16 GiB.

Common mistake Judging compression by memory saved or by a short benchmark only. Quality losses from eviction and aggressive quantization often show up only on long-context retrieval tasks. Always evaluate on your real long prompts, including questions about the middle and start of the context.

When not to bother: for short chats (a few hundred tokens) with few concurrent users, the cache is small compared with the weights. Weight quantization or a smaller model will save more than KV tricks.

Worked example, step by step

So far we asked how big the cache is for a fixed number of users. In practice we ask the reverse question: how many users fit on the GPU we already have? Let us work it out by hand for our support chatbot on one GPU with 80 GiB of memory. We turn the size formula around: users = free memory ÷ cache per user.

From GPU memory to a user count

1. Subtract the weights: The FP16 weights take about 16 GiB. That leaves 80 − 16 = 64 GiB for KV caches. Real servers also keep some spare room for temporary buffers, so treat 64 GiB as an upper limit.
2. Find the cost of one token: Full multi-head FP16: 2 × 32 × 32 × 128 × 2 bytes = 512 KiB. With GQA (8 KV heads) it is 4 times smaller: 128 KiB. With GQA and a 1-byte cache it is 64 KiB.
3. Find the cost of one user: Each user holds 32,000 tokens. Full FP16: 512 KiB × 32,000 ≈ 15.6 GiB. GQA: 128 KiB × 32,000 ≈ 3.9 GiB. GQA with 1-byte numbers: ≈ 1.95 GiB.
4. Divide and round down: 64 ÷ 15.6 → 4 users. 64 ÷ 3.9 → 16 users. 64 ÷ 1.95 → 32 users. We always round down: a user whose cache only half fits cannot be served.
5. Remember the output tokens: The 32,000 tokens must include the answer we are about to write. A 31,500-token prompt with a 500-token reply fills the whole slot by the end of the reply.

Two things stand out. First, the user count moves in whole steps of the compression factor: 4, then 16, then 32. Second, the weights are a fixed cost that compression of the cache never touches. If the context were only 2,000 tokens, each GQA user would need about 0.24 GiB, and the 16 GiB of weights would be the larger part of the bill for the first 60 or so users.

Practice: try it yourself

We will build a toy KV cache with a fixed row budget and run two eviction policies over the same ten tokens. The model is tiny and made up, so the byte counts are small, but the accounting is the same as in the real formula. Watch which tokens survive, and whether the order number A17 is still there at the end.

practice_kv_eviction.py

# Toy KV cache with a row budget: which tokens survive each eviction policy?
LAYERS, KV_HEADS, HEAD_DIM, BYTES = 4, 2, 8, 2    # a tiny made-up model, FP16
ROW_BYTES = 2 * LAYERS * KV_HEADS * HEAD_DIM * BYTES   # K + V for one token
BUDGET = 6                                         # max tokens we may keep
def sliding_window(cache, budget):
# keep only the most recent tokens
return cache[-budget:]
def sinks_plus_window(cache, budget, sinks=2):
# keep the first few tokens (attention sinks) plus the most recent ones
if len(cache) <= budget:
return cache
return cache[:sinks] + cache[-(budget - sinks):]
tokens = ["<s>", "Order", "A17", "is", "late", "and", "the", "box", "was", "wet"]
print("bytes per cached token:", ROW_BYTES)
for name, policy in [("sliding window", sliding_window),
("sinks + window", sinks_plus_window)]:
cache = []
for pos, tok in enumerate(tokens):
cache.append((pos, tok))           # decode step: append one row
cache = policy(cache, BUDGET)      # evict if we are over budget
kept = [tok for _, tok in cache]
print(f"{name}: kept {kept}")
print(f"  bytes = {len(cache) * ROW_BYTES}, "
f"order id still cached: {'A17' in kept}")
print("no eviction: bytes =", len(tokens) * ROW_BYTES)

Output:

bytes per cached token: 256
sliding window: kept ['late', 'and', 'the', 'box', 'was', 'wet']
bytes = 1536, order id still cached: False
sinks + window: kept ['<s>', 'Order', 'the', 'box', 'was', 'wet']
bytes = 1536, order id still cached: False
no eviction: bytes = 2560

Both policies end at the same 1,536 bytes, and both have lost the order number. Now change it:

• Set sinks=3 in sinks_plus_window. Predict first: which token is now protected, and which recent token do we lose to pay for it?
• Set BYTES = 1 and BUDGET = 10 (a 1-byte cache with no eviction). Predict the final byte count and compare it with 1,536. Which plan is smaller here, and which one still holds A17?
• Add twenty more words to tokens and run the original settings. Predict the bytes for each policy and for the no-eviction line. Which of the three numbers keeps growing?

Pause and think: In the practice run both eviction policies use exactly the same number of bytes. Why can they still give different answers in a real model?

Bytes measure how much we keep, not what we keep. The two policies hold different rows, so the new query attends over different keys and values. Equal memory does not mean equal information.

Pause and think: A 1-byte cache with no eviction beat the 6-row budget on memory for our ten tokens. Why does eviction still win for a stream that never ends?

A quantized cache is smaller per token but still grows by one row per token forever. A row budget puts a hard cap on the cache, so its size stops growing no matter how long the stream runs. Quantization changes the slope; eviction sets a ceiling.

Key takeaways

• LLMs generate one token at a time and cache each past token's keys and values to avoid recomputing them.
• KV cache size = 2 × layers × kv_heads × head_dim × bytes × tokens × batch, so it grows linearly with context and users.
• For long contexts the KV cache, not the weights, is usually what limits memory and decode speed.
• Quantization, eviction, head sharing and low-rank compression each shrink a different factor, and they stack.
• Eviction is the only one that throws information away; test it on your real long-context questions.

Key terms

• Autoregressive generation: Producing text one token at a time, each token conditioned on all previous ones.
• KV cache: Stored key and value vectors of all processed tokens, for every layer and head, reused at each decoding step.
• Attention head: One of several parallel attention computations in a layer, each with its own query, key and value projections.
• Quantization: Representing numbers with fewer bits plus a scale, trading a little accuracy for memory.
• Token eviction: Dropping cached keys and values of selected tokens to keep the cache within a budget.
• Grouped-query attention (GQA): An attention design where groups of query heads share a single key/value head.
• Multi-head Latent Attention (MLA): A low-rank attention design that caches a compact latent vector per token and expands it into keys and values.

← 12.8 Cursor: Inside an AI-Native Code Editor · 13.2 Prefill vs Decode: Two Distinct Phases of LLM Inference →