Sampling, KV cache, and efficient generation
Efficient generation couples a declared sampling policy with a position-correct KV cache whose logits match the uncached model.
Domain autocomplete needs a model whose every component the team can explain.
This chapter moves from define the sampling policy to test parity, then throughput. We will make the idea concrete in a small code example before placing it back inside a production system with one-device educational build.
Compare cached and full-prefix logits for every generated token, including batch growth and termination. Only then measure prefill and decode latency separately.
Verify the behavior illustrated in cached_decode.py, then measure decision utility, slice reliability, and system latency.
Without the causal mask, each position sees its target and validation loss becomes impossibly good. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
Training predicts every position in parallel; generation reveals one new token at a time. That change in execution makes policy and caching central.
Greedy decoding chooses the largest logit, while temperature and top-k or top-p sampling reshape the probability distribution before drawing a token. These are product-policy choices: they change diversity, repeatability, and failure modes without changing weights. Fix random seeds only for controlled tests, not as a substitute for evaluating the policy.
Without a cache, each decode step recomputes keys and values for the entire prefix. A KV cache stores past keys and values per layer, so the model projects only the new token while still attending over all cached positions. Cache memory grows with sequence length, layers, heads, and batch size; cached and uncached logits should match token by token before performance claims.
- next-token loss
- KV cache
Follow the reasoning before touching the code.
- 01
Define the sampling policy
Specify temperature, truncation rule, stop tokens, maximum length, and seed behavior. Evaluate deterministic and stochastic paths separately.
- 02
Cache only reusable state
At each layer, append the new token’s key and value while computing its query fresh. Track absolute or rotary position consistently across steps.
- 03
Test parity, then throughput
Compare cached and full-prefix logits for every generated token, including batch growth and termination. Only then measure prefill and decode latency separately.
import torch
@torch.no_grad()
def generate(model, first_token, steps, temperature=0.8):
token = first_token
generated = [token]
cache = None
for _ in range(steps):
logits, cache = model.decode_step(token, cache)
probabilities = torch.softmax(logits / temperature, dim=-1)
token = torch.multinomial(probabilities, num_samples=1)
generated.append(token)
return torch.cat(generated, dim=1)
# decode_step projects only the new token and appends keys/values to the cache.An illustrative decode loop carries a per-layer KV cache between token steps.
What to notice
- The generated token becomes the only new input on the next iteration; past context is supplied through `cache`.
- Lower temperature sharpens the same logits before sampling, but it does not remove randomness unless the policy becomes greedy.
YOUR TURN Pause and predict before revealing the answer.
Why does KV caching improve decode compute but not make attention memory constant?
Need a hint?
Each new query still needs keys and values for every retained past position.
Check your reasoning
Past projections are reused rather than recomputed, but their key and value tensors must remain stored. Cache memory grows linearly with retained sequence length, and each decode step still reads that growing history.
Lessons 1, 2, 3, 4 in this course.
Connect the chapter to the course’s quantitative spine.
Efficient generation couples a declared sampling policy with a position-correct KV cache whose logits match the uncached model. The scaled dot-product attention 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.
Scaled dot-product attention
| Symbol | Meaning / shape / unit |
|---|---|
Q, K, V | query, key, value matrices |
d_k | head width |
M | causal or padding mask |
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: Q, K, V, d_k, M.
- 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_c14(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.
Attention temperature lab
Change score temperature and inspect the exact softmax distribution.
Assumption: Fixed scores [1.2, 0.3, −0.4, 1.8].
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| token 1 | 1.2 | 0.2915 |
| token 2 | 0.3 | 0.1185 |
| token 3 | -0.4 | 0.0589 |
| token 4 | 1.8 | 0.5311 |
Trace the complete operating path.
Carry the chapter’s central idea into the system: Efficient generation couples a declared sampling policy with a position-correct KV cache whose logits match the uncached model. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Define the sampling policy” into a written invariant and a failing test.
- 02
Promote cached_decode.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Test parity, then throughput” as the release checkpoint: Compare cached and full-prefix logits for every generated token, including batch growth and termination. Only then measure prefill and decode latency separately.
- 04
Serve or materialize behind explicit one-device educational build.
- 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
Without the causal mask, each position sees its target and validation loss becomes impossibly good. 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.
Decoder-only transformer from scratch
Specify tokenizer, attention blocks, causal training, unit tests, sampling, and KV-cache parity without a transformer module.
- Shape/mask tests
- Tiny-batch overfit
- Held-out evaluation
- Cached/uncached parity
Read primary material with a purpose.
Return to the opening failure.
Version tokenizer, vocabulary, special tokens, architecture, weights, prompt template, and generation policy together.
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 tokenizer, attention, rope, residual stack, training loop, generation, and kv cache.