MEMORY OPTIMIZATION

Module 2 established the waste: a naive allocator reserves max_seq_len contiguously per sequence, and the vLLM authors measured that only 20.4% to 38.2% of allocated KV memory actually held tokens. Between 60% and 80% was fragmentation.

Operating systems solved exactly this problem in the 1960s, and the solution transfers almost without modification.

The mechanism. Divide the KV cache into fixed-size blocks — typically 16 tokens' worth. Each sequence gets a block table mapping its logical positions to physical block numbers. Blocks are allocated on demand as the sequence grows, from a shared pool, and returned when it finishes. Physical blocks belonging to one sequence need not be adjacent.

sequence A, 35 tokens, block size 16:
  block table: [ 7, 2, 19 ]          3 blocks = 48 slots, 35 used

  physical pool:
    block 0  [ free    ]
    block 1  [ seq B   ]
    block 2  [ seq A   ]  <- logical positions 16..31
    block 3  [ seq C   ]
    ...
    block 7  [ seq A   ]  <- logical positions 0..15
    ...
    block 19 [ seq A   ]  <- logical positions 32..34, 13 slots spare

Waste drops to at most one partial block per sequence — on average half a block, so 8 tokens. Against 128 KiB/token on Llama-3-8B that is about 1 MiB per sequence, versus gigabytes under reservation. vLLM reports waste below 4%.

The attention kernel has to change to match: instead of reading a contiguous [seq_len, head_dim] tensor, it gathers through the block table. That indirection costs a little, and the gain vastly outweighs it.

Copy-on-write falls out for free, and this is the elegant part. Two sequences sharing a prefix can point their block tables at the same physical blocks. Nothing is copied. When one of them writes to a shared block — which only happens at the boundary block where they diverge — that block is copied and the writer's table updated. Parallel sampling with n=4 from one prompt therefore costs one copy of the prompt's cache, not four. Beam search gets the same treatment.

Choosing the block size is a real trade-off. Small blocks (8) minimize internal fragmentation but mean more block-table entries and more indirection per attention call. Large blocks (32+) reduce overhead but waste more on partial blocks and make prefix sharing coarser — two sequences must share a whole block to share anything. 16 is the common default and is a reasonable compromise rather than a derived optimum.

Module 2 introduced the idea; here is how it is actually built.

The naive version handles one fixed system prompt: hash it, cache its KV blocks, reuse them. Useful, but it only catches the case you anticipated.

RadixAttention — SGLang's contribution — generalises it. Maintain a radix tree (a compressed trie) whose edges are token sequences and whose nodes hold the corresponding KV blocks. For each incoming request, walk the tree to find the longest matching prefix, reuse those blocks, and prefill only the remainder. Insert the new suffix as you go. Evict by LRU when memory is short.

This catches sharing you did not design for:

The agentic case is the one that changes the economics. An agent taking 20 steps re-sends its growing history every step, so without prefix caching the total prefill work is quadratic in step count. With it, each step prefills only the newest action and observation. For long-running agents this is not an optimization, it is the difference between viable and not.

Three implementation realities:

RoPE positions must be handled. Cached keys were rotated by their absolute position. Reuse a prefix at a different offset without accounting for this and attention is silently wrong — no crash, just degraded output. This is the most common prefix-caching bug.

Block granularity limits matching. With 16-token blocks you can only share whole blocks, so two prompts differing at token 5 share nothing even though the first 5 tokens match. Prefixes that matter should be block-aligned, which is one reason system prompts are often padded.

Eviction is a genuine trade-off, not a detail. Cached prefixes occupy memory that running sequences need. Keeping a prefix costs you batch size now against a possible hit later. LRU is the usual answer and it is a heuristic; the right policy depends on your traffic's reuse distribution, which most operators have never measured.

There is also a security consideration worth naming: cache is shared derived state across users. Hits require an exact token-prefix match, so content does not leak directly, but hit/miss timing can in principle reveal whether another user has submitted a given prefix. For most deployments this is acceptable; for multi-tenant systems with sensitive prompts it is worth an explicit decision rather than a default.

Paging and prefix caching stop you wasting memory. Architecture changes how much you need in the first place.

Recall the formula: 2 × n_layers × n_kv_heads × head_dim × dtype_bytes. Two of those terms are architectural choices.

MHA → MQA → GQA attack n_kv_heads, and Module 2 has the numbers: on Llama-3-70B, MHA needs 2.5 MiB/token, GQA-8 needs 320 KiB, MQA needs 40 KiB. GQA at 4–8 groups is where the field settled, because MQA's single shared KV head costs measurably more quality than the memory is worth for most models.

MLA — multi-head latent attention, from DeepSeek-V2 — attacks the problem differently and more aggressively. Instead of reducing the number of KV heads, it compresses K and V jointly into a single low-rank latent vector per token, and caches only that.

standard GQA, per token per layer:
    K: n_kv_heads x head_dim  = 8 x 128 = 1024 values
    V: n_kv_heads x head_dim  = 8 x 128 = 1024 values
                                          ---- 2048 values

MLA, per token per layer:
    c_kv: kv_lora_rank        = 512 values      (compressed K and V together)
    k_rope: rope dim          =  64 values      (decoupled positional part)
                                 ---- 576 values

At inference, the per-head K and V are reconstructed from the latent by up-projection. The trick that makes this cheap is absorption: the up-projection matrices can be folded into the adjacent query and output projections algebraically, so you never actually materialize the full K and V — you compute attention directly against the latent.

The decoupled RoPE part exists because rotation does not commute with the low-rank projection. Applying RoPE to a compressed latent and then decompressing does not give the same answer as compressing rotated keys, so MLA keeps a small separate rotary component. It is an implementation-forced wrinkle rather than an elegant piece of design, and it is where most of the complexity of MLA lives.

DeepSeek report a KV reduction of roughly 93% versus MHA, and — the part that made people pay attention — better benchmark quality than the MHA baseline, not worse. The plausible explanation is that the low-rank bottleneck acts as a regularizer. Whether that generalises beyond their models is not yet established.

The catch for all three: these are training-time decisions. You cannot convert a deployed MHA model to MLA. GQA has an uptraining recipe (mean-pool the KV heads within a group, fine-tune on ~5% of original compute), but MLA requires training from scratch. If you are choosing a model to serve, its attention architecture is one of the most consequential facts about its inference cost, and it is fixed before you get involved.

Recall the two regimes from Module 2. In the weight-dominated regime — small batch, short context, single-user local inference — weight bytes are your step time, so halving them roughly halves your latency.

Weight-only quantization stores weights at low precision and dequantizes to fp16 on the fly for the matmul. The arithmetic still happens in fp16; the saving is purely in bytes moved, which is exactly what you want when you are memory-bound.

GPTQ (Frantar et al., 2022) quantizes layer by layer, using approximate second-order information from a small calibration set to choose quantization points that minimize the layer's output error. It works down to 3–4 bits on large models with modest perplexity loss.

AWQ (Lin et al., 2023) starts from the observation that not all weights matter equally: a small fraction of channels — around 1% — are salient, identifiable from activation magnitude rather than weight magnitude. Scale those channels up before quantizing so they land on finer grid points. Simple, no backpropagation, and generally more robust than GPTQ at 4 bits.

Weight+activation quantization is a different game. SmoothQuant (Xiao et al., 2022) addresses the fact that activations have severe outliers that make them hard to quantize, by migrating the difficulty into the weights via a per-channel scaling that cancels out mathematically. This enables int8 for both operands, which speeds up the arithmetic too — worth it only when you are compute-bound, which for decode you are not.

FP8 on Hopper and later is the pragmatic modern default for datacenter serving: hardware support, roughly half the bytes, and quality loss small enough that many deployments accept it without extensive validation. FP4 on Blackwell extends the same idea further.

A rough guide:

Two cautions. First, the speedup is not the compression ratio — you also move activations and KV cache, and dequantization costs some arithmetic. 4-bit weights typically give 2.5–3.5×, not 4×. Second, and more important: in the cache-dominated regime, weight quantization does much less than you expect. At batch 32 with 8k context on Llama-3-8B, the KV cache is 2.3× the weights, so even eliminating weight bytes entirely would cut your step time by only about 30%. Check which regime you are in before choosing.

In the cache-dominated regime, quantizing the cache is the lever that matters. It halves (fp8/int8) or quarters (int4) the denominator of the batch equation directly.

The naive approach — quantize everything per-tensor — degrades quality more than you would expect from the bit count. The reason is a real structural asymmetry, documented in the KIVI work and others:

Keys have strong per-channel outliers. Certain channels of the key vectors consistently carry much larger magnitudes than others, across tokens. Quantizing a key tensor per-tensor means those outlier channels dominate the scale and every other channel is crushed into a few quantization levels.

Values do not show this pattern. Value distributions are comparatively uniform across channels.

So the correct treatment is asymmetric:

K: quantize PER CHANNEL   -- each channel gets its own scale, so an outlier
                             channel does not ruin the others

V: quantize PER TOKEN     -- each token's value vector gets its own scale

Per-channel quantization of K has an awkward consequence: the cache grows one token at a time, so a per-channel scale computed over the tokens seen so far changes as you append. Implementations handle this by quantizing in groups and keeping a small recent window in full precision.

Practical guidance:

• fp8 KV is close to free on hardware with fp8 support. It is a reasonable default for production serving and doubles your batch capacity.
• int8 KV with the per-channel/per-token split is also reliable.
• int4 KV is where quality starts to show, and it shows specifically on long-context tasks — which is a problem, because long context is exactly when you wanted the savings. Test on retrieval and long-context benchmarks, not perplexity.

There is a related family that reduces the cache without changing its precision: eviction and compression.

• StreamingLLM observed that models allocate large attention weight to the first few tokens regardless of content — "attention sinks" — and that keeping those plus a sliding window lets a model stream indefinitely without collapse. Cheap and effective for streaming; it does discard the middle, so it is not suitable when you need recall over the full context.
• H2O keeps "heavy hitter" tokens identified by accumulated attention score, plus recent tokens. Better recall than a pure window.
• SnapKV compresses the prompt's cache at the end of prefill by observing which positions the final query attends to.

All of these are lossy in a way that is hard to bound. You cannot know at eviction time which token some future query will need. They work well on benchmarks that resemble their design assumptions and can fail sharply outside them. Treat them as workload-specific tools, not defaults.

Every technique in this module except paging and prefix caching costs something. The discipline that separates good work from bad here is measurement, and the field is not good at it.

Why perplexity is insufficient. Perplexity is the average negative log-likelihood over a corpus. It is dominated by common tokens in ordinary contexts, and a quantized model reproduces those nearly perfectly. Damage concentrates in the tail: rare tokens, unusual contexts, long dependencies, precise arithmetic. A 4-bit model can match fp16 perplexity to two decimal places and fail visibly at 20-step reasoning.

There is a second problem: perplexity is measured on short sequences, usually 2k–4k tokens. If your concern is a quantized KV cache, that is exactly the regime where it looks fine.

What to measure instead, roughly in order of value per unit effort:

1. Long-context retrieval. Needle-in-a-haystack and, better, RULER — which tests multi-needle, aggregation and tracing rather than a single lookup. This is where KV quantization damage shows first and clearest.
2. Multi-step reasoning. GSM8K, MATH, or your own chain-of-thought tasks. Errors compound across steps, so small per-token degradation becomes visible.
3. Your actual task. A few hundred real prompts with a human or model-graded comparison against the fp16 baseline. Unglamorous and the most informative thing on this list.
4. Output distribution divergence. KL divergence between quantized and baseline logits on a fixed prompt set. Sensitive, cheap, and catches problems before they show in task metrics.
5. Standard benchmark suites. MMLU and friends. Useful as a regression check, weak as a quality signal — they are multiple-choice and forgiving.

The comparison must be controlled. Same prompts, same sampling parameters, same seed where possible. It is remarkably common to see quantization "evaluated" against a baseline that used different sampling settings, which measures nothing.

Test at the operating point. If you serve at 32k context, evaluating at 2k tells you almost nothing about KV quantization. If you serve at batch 64, evaluate under batching — batch-dependent kernel selection can itself shift outputs, as Module 3 discussed.

The honest position on most published quantization results is that they demonstrate the method does not catastrophically fail, not that it is free. Whether it is free for your workload is a question only your evaluation can answer.