KV caches, GQA/MQA, and long context
KV caching removes repeated history compute by turning context length into a first-class memory budget.
Long-context generation breaks latency and cost targets.
This chapter moves from separate prefill from decode to enforce a context policy. We will make the idea concrete in a small code example before placing it back inside a production system with p99 first-token and cost/token SLOs.
Bound admitted tokens, evict or offload explicitly, and test that cached and uncached logits match at every generated position.
Verify the behavior illustrated in kv_budget.py, then measure decision utility, slice reliability, and system latency.
Maximizing aggregate tokens/s destroys p99 first-token latency and harms a rare-language quality slice. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
Without a cache, generating token 1,001 recomputes attention keys and values for the first 1,000 tokens even though they have not changed. A KV cache trades that repeated compute for memory, which is why long context and high concurrency can exhaust a server before model weights do.
During prefill, every prompt token produces a key and value at every decoder layer. Decode adds one new key-value pair per step and attends over the stored history. Cache size therefore grows linearly with layers, live sequences, sequence length, KV heads, head width, and bytes per scalar.
Multi-query attention shares one set of keys and values across all query heads; grouped-query attention shares them within groups. Both reduce cache traffic compared with full multi-head attention, but they are architectural choices fixed during training. Paged allocation can reduce fragmentation, yet it does not remove the underlying bytes or the need for an admission policy.
- sampling
- KV layout
- GQA/MQA
- quantization
Follow the reasoning before touching the code.
- 01
Separate prefill from decode
Profile prompt processing and token-by-token generation independently because their arithmetic intensity and latency behavior differ.
- 02
Budget cache bytes
Calculate cache memory from model shape and the actual mix of concurrent context lengths, then reserve headroom for weights and runtime workspaces.
- 03
Enforce a context policy
Bound admitted tokens, evict or offload explicitly, and test that cached and uncached logits match at every generated position.
def kv_gib(layers, batch, tokens, kv_heads, head_dim, bytes_per_value=2):
keys_and_values = 2
total = (keys_and_values * layers * batch * tokens
* kv_heads * head_dim * bytes_per_value)
return total / 1024**3
full_mha = kv_gib(32, 16, 8192, 32, 128)
gqa = kv_gib(32, 16, 8192, 8, 128)
print(f'full MHA: {full_mha:.1f} GiB')
print(f'8-head GQA: {gqa:.1f} GiB')Estimate KV-cache memory before choosing concurrency
What to notice
- Reducing KV heads from 32 to 8 cuts cache storage by four for the same layer, batch, context, and precision.
- The calculation is a lower-bound planning number; allocator fragmentation, metadata, activations, and workspaces still require headroom.
YOUR TURN Pause and predict before revealing the answer.
Keep the GQA model fixed and double both context length and concurrent sequences. By what factor does cache memory grow?
Need a hint?
Both batch and token count are independent multiplicative terms.
Check your reasoning
It grows by four: doubling concurrency multiplies by two and doubling context multiplies by another two.
Lessons 1 in this course.
Connect the chapter to the course’s quantitative spine.
KV caching removes repeated history compute by turning context length into a first-class memory budget. The kv-cache memory estimate below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
KV-cache memory estimate
| Symbol | Meaning / shape / unit |
|---|---|
L | decoder layers |
B | active sequences |
S | cached tokens |
H_kv | KV heads |
q | bytes per scalar |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: L, B, S, H_kv, q.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c16(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
GPU memory ledger
Change microbatch size and calculate a simplified training-memory budget.
Assumption: 3.2 GB weights+grads+optimizer, 0.34 GB activations per sample, 24 GB device, 10% reserve.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| block 1 | 1.09 GB activations | resident |
| block 2 | 2.18 GB activations | resident |
| block 3 | 3.26 GB activations | resident |
| block 4 | 4.35 GB activations | resident |
| block 5 | 5.44 GB activations | resident |
Trace the complete operating path.
Carry the chapter’s central idea into the system: KV caching removes repeated history compute by turning context length into a first-class memory budget. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Separate prefill from decode” into a written invariant and a failing test.
- 02
Promote kv_budget.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Enforce a context policy” as the release checkpoint: Bound admitted tokens, evict or offload explicitly, and test that cached and uncached logits match at every generated position.
- 04
Serve or materialize behind explicit p99 first-token and cost/token SLOs.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
Maximizing aggregate tokens/s destroys p99 first-token latency and harms a rare-language quality slice. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
LLM inference capacity plan
Benchmark cache, batching, concurrency, sampling, and one quantized path on mixed traces.
- Benchmark contract
- Profiler/bottleneck report
- Quality gates
- Capacity and cost recommendation
Read primary material with a purpose.
Return to the opening failure.
Choose optimizations from workload traces with quality gates, tail latency, memory headroom, and cost per token.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce a capacity plan using efficient attention, quantization, batching, and modern routing.