AI Research Library
Causal Masks, KV Caches and Generation
Autoregressive generation remains the same model only when the mask fixes a legal past and the cache reuses that past without silently rewriting it.
Rajesh Ranjan Mahapatra
© 2026 Rajesh Ranjan Mahapatra. 30 min read
Causal Masks, KV Caches and Generation
Autoregressive generation remains the same model only when the mask fixes a legal past and the cache reuses that past without silently rewriting it.
TLDR
- A causal mask is part of the model's statistical contract: token t may depend on earlier tokens, never on the answer it is being trained to predict.
- A KV cache can preserve that contract exactly because past keys and values are fixed once their tokens and positions have been processed.
- The cache removes repeated projection and attention work during decoding, but trades compute for memory capacity, bandwidth and lifecycle complexity.
- Eviction, quantisation, wrong positions or stale cross-request reuse can change logits; a smaller cache is therefore not automatically the same model.
- Release decoding optimisations only after classifying them as exact or approximate, proving parity, and measuring prefill, token latency, throughput and memory separately.
Three decoders continue the same sentence
A service receives the prompt, “The vault code is stored in”. Three workers hold identical model weights. They all produce the continuation “the sealed register”, and they all appear correct.
The first worker recomputes every layer over the complete prefix before choosing each new token. The second computes the prompt once, stores the attention keys and values created at every layer, then processes one new token at a time. The third behaves like the second but discards most old cache entries whenever memory is tight.
One successful continuation cannot tell us whether these workers executed the same conditional distribution. The first two should agree, apart from bounded numerical differences, because the second only reuses quantities the first recomputes. The third may agree on this prompt and diverge on another because it has changed which past states remain available.
The causal mask defines the model's information boundary; an exact KV cache changes the execution plan, not that boundary. Once a cache is pruned, quantised, mis-positioned, mixed across requests or reused after the model changes, the speed-up becomes a modelling claim that requires task-level evidence.
Figure 1. The triangle that makes next-token learning legal
- Rows are queries made by positions 1 to 6.
- Columns are keys offered by positions 1 to 6.
- Green cells are legal past or current positions.
- Hatched cells would leak a future target.
- During one-token cached decode, only the newest legal row is evaluated.
Part I
The mask defines the legal past
A decoder-only language model assigns a probability to a sequence by decomposing one difficult joint prediction into a chain of next-token predictions. For tokens x1 to xT, the factorisation is:
T is sequence length, xt is the token at position t, and each factor may use only earlier tokens. If a future token enters the condition, the model is trained on information unavailable at generation time.
Training nevertheless processes many positions in parallel. The apparent contradiction disappears once computation and information are separated. Matrix operations for all positions can run together, while a triangular mask removes every attention edge that points forward. The original Transformer paper explicitly uses this masking to preserve the autoregressive property, setting illegal attention logits to negative infinity before softmax (Vaswani et al.).
For one attention head, the operation can be written as:
Q, K and V are query, key and value matrices; dk is key width; M is zero for legal edges and negative infinity for forbidden future edges. Making one cell of M legal gives that query access to a future representation.
The mask is therefore more than a speed option or a display convention. It identifies which conditional model is being learned. A bidirectional encoder, a prefix language model and a left-to-right decoder can share attention machinery while representing different dependency structures because their masks differ.
Parallel training is sometimes described as if the model “generates every token at once”. That merges two different timelines. On the computational timeline, a device evaluates many query rows together. On the dependency timeline, row t still conditions only on rows before it. The mask lets hardware parallelise lawful conditional predictions without granting any prediction access to its future target.
The distinction matters for debugging. A low training loss can survive a mask defect because leaked targets make the optimisation problem easier. Generation then exposes the mismatch: the future columns are absent, so the model is asked to operate under a dependency graph it did not learn. A robust test should therefore compare the implemented mask against the intended factorisation, verify target shifting, and include a synthetic sequence whose answer becomes trivial if even one future edge opens.
The mask also shapes gradient routes. During training, a future representation that enters an earlier query can influence the loss through attention and then through shared parameters. Removing the leaked edge at inference does not merely remove a convenience; it removes a path the optimiser may have relied upon. Training-time legality and decode-time legality must be the same contract.
Thought experiment: the completed exam
Give two students the same partially completed arithmetic worksheet. Every row contains a question followed by its answer. Student A may look only at previous rows while filling the next answer. Student B receives the complete sheet and may look one cell to the right, where the target answer already appears.
Both can be trained in parallel across rows. Vary only future visibility. Student A must exploit the prefix; Student B can copy. If the evaluation later presents a blank final answer, Student B's training loss was evidence of leakage, not predictive competence. The causal mask plays the role of the opaque strip covering unearned answers.
Now restore the mask but remove the one-position shift between inputs and targets. The representation at row t can contain token xt while the loss also asks for xt. The diagonal becomes another route to the answer. The mask and target alignment jointly enforce the factorisation. Testing only that a triangular tensor exists is insufficient.
Causality here is statistical, not metaphysical and not a security permission. The mask says which token positions can influence another inside the forward pass. It does not establish whether a prompt is trusted, whether a tool call is authorised or whether cached state belongs to the same user.
flowchart TD A["Three decoders continue the<br/>same sentence"] --> B["Part I The mask defines the<br/>legal past"] B --> C["Part II The cache is exact reuse"] C --> D["Part III Faster generation has<br/>a fidelity boundary"] D -. "re-test" .-> A classDef first fill:#0a1424,stroke:#D95545,stroke-width:2px,color:#e6f3ff classDef step fill:#0a1424,stroke:#3f7da0,stroke-width:1px,color:#e6f3ff class A first class B,C,D step
Part II
The cache is exact reuse
After the prompt has been processed, generation samples one token, appends it, then asks for another distribution. A naive worker sends the whole growing prefix through every layer on every step. Earlier positions produce the same per-layer keys and values each time because a causal model prevents later tokens from changing their representations.
That invariance creates the cache. Each attention layer stores the keys and values already computed for legal past positions. On the next step it computes a query, key and value for the new token, appends the new key and value, and lets the new query attend over the stored sequence. Queries are usually not retained because a query serves the position that issued it once; its key and value may be consulted by every later query.
Figure 2. Time-unrolled cache cutaway
The equality is local to every layer. Suppose the cached worker has processed tokens 1 through t. At layer l, its stored keys and values equal those from a full causal forward pass over the same prefix. The next token's layer input is also equal by induction from layer l minus one. The new query sees the same past keys and values in the same order, so it produces the same attention output. Repeating this argument through all layers gives the same next-token logits, subject to floating-point implementation differences.
What kind of “same” is being claimed?
Teams often test only that two decoders emit the same visible string. That can be too weak or unnecessarily strong. A small logit change may leave greedy output unchanged while changing sampled continuations. Conversely, two mathematically equivalent kernels can differ in their last floating-point bits because reductions occur in another order. The release claim should name its equivalence level.
| Equivalence level | What must match | What it supports | What it does not prove |
|---|---|---|---|
| Bitwise | Every stored value and output bit | Strict replay on a fixed stack | Portability across kernels or hardware |
| Logit tolerance | Next-token logits within a declared numerical bound | Implementation parity under floating-point variation | Identical sampled tokens without controlled randomness |
| Distributional | Sampling probabilities or a justified divergence bound | Claims about stochastic generation | Equal task outcomes on rare cases |
| Sequence | Greedy or seeded sampled token sequence | Regression testing of a decoding policy | Unchanged probability mass below the chosen path |
| Outcome | Task metric inside a declared envelope | Use-case acceptance of an approximation | That the underlying model execution is exact |
For an ordinary exact cache, tolerance-bounded logit parity is the most revealing default. It tests the mechanism before sampling hides small differences. Sequence tests then exercise cache growth, stopping and random-state handling. Task tests remain necessary for numerical kernels and essential for any approximation, but they should not replace the lower-level comparison that localises a defect.
The identity of reusable state is stricter than a human-readable prompt. Two requests may display the same characters while using different token IDs, adapters, rotary scaling, attention windows or model revisions. Each difference can change keys and values. Cache identity is therefore a tuple of the processed token prefix and every configuration item that can affect the per-layer state, not simply a string hash.
A cache is valid because past per-layer attention state is invariant to future tokens, not because “the model remembers”. The cache is workflow state for a particular model execution. It should not be confused with episodic memory, enterprise knowledge or an evidence ledger.
Minimal worked example: four tokens, one new row
Take a four-token prompt. A full causal attention matrix contains ten legal score cells: one in row one, two in row two, three in row three and four in row four. The final row alone determines the prompt's next-token logits after the remaining layer operations. Once the first four tokens are cached, processing token five requires five score cells, all in one row. Recomputing the five-token prefix would evaluate fifteen legal cells, including ten whose outputs are unchanged.
Continue for many steps and the difference compounds. If P is prompt length and G is the number of sampled tokens, a deliberately naive full-recompute plan evaluates a growing triangle before each sample. Cached decoding evaluates the prompt triangle once, then one row for each later token:
Each count is per layer and query head. Increasing G makes full recomputation grow cubically in this intentionally wasteful execution plan, while cached attention score work grows quadratically. Actual kernels also spend time on projections, feed-forward layers, memory movement, batching and scheduling.
Figure 3. Repeated attention work separates rapidly
The bill moves from compute to memory
Exact reuse is not free. For a conventional cache, approximate storage in bytes is:
L is layer count, T cached tokens, Hkv key-value heads, dh head width, and b bytes per element. The leading 2 stores both keys and values. Doubling context or concurrent sequences doubles this state.
A synthetic 32-layer model with eight KV heads, head width 128, an 8,192-token context and two-byte elements needs exactly 1 GiB of KV state per sequence before allocator overhead. With 32 KV heads it needs 4 GiB. This is why multi-query attention and grouped-query attention reduce serving memory and bandwidth by sharing key-value heads, while introducing architecture-level quality trade-offs that must be settled during model design (Shazeer; Ainslie et al.).
Figure 4. One request can carry a gigabyte of transient model state
Part III
Faster generation has a fidelity boundary
“Uses a KV cache” covers several unlike interventions. One stores exact tensors in a different layout. Another shares an identical prefix across requests. A third keeps two-bit approximations. A fourth drops old tokens. A fifth lets a small model propose tokens and asks the target model to accept or reject them. These changes can all reduce cost, yet they carry different claims about the output distribution.
The first release question is not “How much faster?” but “What, if anything, has changed in the target model's conditional distribution?”
Figure 5. The fidelity boundary for decoding shortcuts
Exact state, exact result
Ordinary KV reuse, exact attention kernels, paged storage and exact prefix reuse. Layout or scheduling changes; the same logical keys and values reach the same queries.
Approximate proposal, exact acceptance
Speculative decoding may use a cheaper draft model, but target-model verification and corrected sampling preserve the target distribution when the algorithm's conditions hold.
Approximate state, evaluated result
KV quantisation, eviction, sliding windows and sparse selection can alter logits. They may be excellent engineering choices, but quality equivalence is empirical and workload-specific.
Wrong state, invalid result
Stale weights, wrong positions, mismatched masks, adapter changes or cross-request contamination violate cache identity. A plausible output does not repair the execution contract.
Exact changes still need parity tests
FlashAttention is an IO-aware exact attention algorithm: it changes tiling and memory traffic rather than approximating the dense result (Dao et al.). PagedAttention changes how KV blocks are allocated and shared so dynamic sequences waste less memory (Kwon et al.). Prefix caching reuses blocks when token prefixes match, reducing repeated prefill work; both vLLM and TensorRT-LLM document this behaviour (vLLM; TensorRT-LLM).
These are semantic-preserving intentions, not immunity from implementation defects. Parity must cover token IDs, model and adapter revision, position IDs or rotary offsets, attention window, KV head mapping, cache dtype, padding convention and any mask that changes by layer. Sampling tests must also hold the random stream constant if exact token-sequence comparison is required.
Speculative decoding illustrates a subtler exactness argument. A cheap model proposes several tokens; the target model evaluates them in parallel, accepts a valid prefix and corrects the sampling distribution when needed. Under the published algorithm, the target distribution is unchanged even though proposals are approximate (Leviathan et al.). Removing verification would cross the fidelity boundary.
The strongest objection: outcomes matter more than internal equality
A service owner can reasonably object that users do not consume logits. If a quantised or pruned cache meets the task's quality, latency and cost thresholds, why privilege exact execution? For many bounded workloads, the objection is correct. Exactness is not a moral property, and an approximation can be the better architecture.
The distinction still matters because evidence transfers differently. Algebraic or algorithmic equivalence supports new prompts within the mechanism's domain without retesting every semantic case. An approximate cache earns only the envelope actually evaluated. Long-range copying, code closure, numerical consistency and rare safety-relevant dependencies may respond differently even when an average benchmark barely moves. Calling both paths “the same model” silently expands evidence beyond its tested boundary.
An observation could weaken this paper's practical recommendation. Suppose an approximation is evaluated across representative and adversarial strata, its divergence remains below a decision-relevant threshold, and later incidents show no concentration in the altered state. In that use case, insisting on exact cache parity would add little. The correct conclusion would be to keep the exact-versus-approximate label for traceability while governing release primarily through the demonstrated outcome envelope.
Thought experiment: the perfect and forgetful stenographers
Two stenographers listen to the same hearing. Both avoid replaying the entire recording before typing the next sentence. The perfect stenographer keeps an exact index of every earlier statement. The forgetful stenographer retains only the opening oath and the last ten minutes.
For a routine exchange about the current witness, their transcripts match. Now vary one causal feature: ask whether today's amount contradicts a figure stated two hours earlier. The perfect index preserves the relevant dependency. The sliding record cannot recover it. The second method is faster and may be adequate for local dialogue, but it no longer implements unrestricted attention over the full hearing.
StreamingLLM shows that naive sliding-window eviction can fail and that retaining initial attention-sink tokens can recover much of the behaviour in streaming settings (Xiao et al.). KIVI reports large memory and throughput gains from asymmetric two-bit KV quantisation with near-baseline aggregate quality on evaluated models and tasks (Liu et al.). Those results support useful approximations; they do not turn eviction or quantisation into algebraic identities.
A clean negative control is a request that asks for one next-token score and stops. Both workers must process the prompt once; there is no second decode step on which stored K and V can repay their allocation. A cache may still be produced by a generic runtime, but it has no causal opportunity to save repeated work. Likewise, reranking a fixed list with a bidirectional encoder uses a different dependency structure and should not inherit a decoder-cache argument by analogy.
Another boundary appears when the system is dominated by moving cache state rather than using it. Offloading a large cache to host memory, transmitting it between disaggregated workers or restoring it after pre-emption can cost more than recomputation for short remaining outputs. The right comparison is avoided work minus allocation, lookup, transfer and synchronisation cost under the actual residual generation length.
The companion NumPy lab supplies a negative control. It reduces theoretical attention score pairs by 75 times for one synthetic workload, yet its tiny cached implementation can be slower than full recomputation because Python calls and small matrix operations dominate. On a longer synthetic prefix, the crossover can reverse. Fewer operations are a mechanism, not a wall-clock guarantee.
A 2025 production-oriented evaluation reached a related warning at larger scale: KV compression can lower memory while failing to improve throughput in existing serving paths, and altered generations can become longer, increasing end-to-end latency (Gao et al.). Aggregate accuracy alone can miss sample-specific failures. The proper unit of evidence is the configured model, kernel, scheduler, hardware and workload.
Figure 6. Exactness occupies a narrow basin
stateDiagram-v2 [*] --> Claim state "Figure 1. The triangle that<br/>makes next-token learning<br/>legal" as Claim state "Thought experiment, the<br/>completed exam" as Test state "Figure 2. Time-unrolled cache<br/>cutaway" as Revision Claim --> Test: expose mechanism Test --> Decision: evidence survives Test --> Revision: rival survives Revision --> Claim: narrow or rebuild Decision --> [*]
Part IV
Operate the cache as governed state
At service scale, the cache is not one tensor attached to one request. It is a changing pool of blocks shared by batches with different prompt lengths, generated lengths and arrival times. Prefill and decode also stress hardware differently: prefill processes many prompt tokens in parallel, while decode repeatedly processes one token per active sequence. Systems work on chunked prefill, continuous batching and phase disaggregation because optimising one phase can damage the latency of the other (Agrawal et al.; Zhong et al.).
Production-shaped worked scenario
Consider a synthetic internal assistant serving 24 concurrent conversations with the 32-layer, eight-KV-head configuration from Figure 4. At the maximum 8,192-token context, raw KV state alone can reach 24 GiB. Model weights, activations, allocator metadata and temporary kernel buffers remain outside that number.
The service uses paged blocks so a short conversation does not reserve its maximum length. A common system instruction is eligible for prefix reuse. Each request then appends private conversation tokens. The scheduler admits new prefills in chunks while protecting the token cadence of active decodes. When pressure rises, it may queue a request, offload blocks, lower admission or invoke an explicitly evaluated approximate-cache policy.
Paged allocation solves a physical-memory problem: variable-length sequences no longer need one contiguous maximum-sized reservation, and shared prefixes can refer to common blocks. It does not decide whether two logical requests are entitled to share those blocks or whether their state was produced by the same model configuration. Memory indirection removes fragmentation; it does not remove semantic identity.
The architecture decision is not “turn caching on”. It is a lifecycle with typed boundaries:
Figure 7. A cache block has a request lifecycle
Prefix reuse deserves particular care. Equality should be evaluated over token IDs and every model-state attribute that affects K and V. Text equality before tokenisation is not enough; neither is a hash that ignores model revision or adapter selection. Official vLLM code exposes prefix-cache reset for flows where weights change, which makes the invalidation dependency concrete (vLLM API).
| Identity field | Why it belongs | Failure if omitted |
|---|---|---|
| Model and weight revision | Projection matrices create K and V | Old tensors are interpreted under new weights |
| Adapter or fine-tune selection | LoRA and related changes can alter layer outputs | Two configured models share an invalid prefix |
| Prefix token IDs and length | K and V are functions of the actual tokens and order | Text-normalisation or boundary differences collide |
| Position contract | Absolute, rotary or relative position state affects attention | Correct tokens occupy the wrong geometry |
| Attention policy by layer | Full, sliding and hybrid layers expose different histories | A block is reused outside its legal window |
| KV representation | Head grouping, dtype and quantisation define storage semantics | Readers decode bytes under the wrong layout or scale |
| Tenant and policy scope | Permission to share is separate from computational equality | Private state crosses an application boundary |
A cache handle should bind to this identity rather than trusting caller-supplied position or length alone. When a worker resumes, it can compare the handle's declared token range with the actual cache pages, reject gaps or duplicate appends, and rebuild from the authoritative token sequence when uncertainty remains. Recompute is often safer than guessing whether partially restored state is coherent.
Cross-request reuse is also an authority boundary. A cache manager may know that two prefixes are computationally identical while the application knows they belong to different tenants, classifications or retention policies. Computational sameness does not grant permission to share state. Tenancy and entitlement should either partition the cache namespace or be included in a policy check before reuse. This is practitioner synthesis rather than a claim that one public engine supplies the complete control.
Measure four clocks, not one
A single “latency” number hides different causal paths. Time to first token includes queueing and prefill. Time per output token reflects decode cadence. Throughput depends on batching and admission. Memory headroom determines how many sequences can coexist before eviction or queueing. A cache change can improve one and worsen another.
| Measure | Primary question | Typical confounder | Required slice |
|---|---|---|---|
| Time to first token | How quickly can the prompt be admitted and prefetched? | Queueing, prefix-hit rate, chunking | Prompt length and hit or miss |
| Time per output token | How stable is the decode cadence? | Batch churn, KV bandwidth, prefill interference | Active batch and context length |
| Goodput or throughput | How much demand meets the latency objective? | Unbounded batching can inflate raw throughput | Declared service-level thresholds |
| Peak and live KV bytes | How many sequences fit without pressure action? | Fragmentation and shared-page accounting | Allocator state and concurrency |
| Parity or quality | Did the optimisation alter logits, tokens or task outcomes? | Sampling randomness and aggregate masking | Deterministic logits plus task strata |
Executable artefact: cached-decoder parity benchmark
The companion programme implements a two-layer causal decoder with synthetic NumPy weights. It compares naive full-prefix generation with per-layer KV reuse, asserts identical generated tokens and near-identical logits, then introduces a wrong-position negative case that must diverge. It also reports deterministic score-pair counts and machine-dependent timings.
Open the complete runnable Python benchmark
Save as kv_cache_parity_benchmark.py, install NumPy, and run python3 kv_cache_parity_benchmark.py.
"""Exact-versus-broken KV-cache laboratory for a tiny causal decoder.
Python 3.10+ and NumPy 1.24+.
All weights, tokens, dimensions and timings are synthetic teaching fixtures.
"""
from __future__ import annotations
from dataclasses import dataclass
from time import perf_counter
from typing import List
import copy
import numpy as np
Array = np.ndarray
@dataclass
class LayerWeights:
wq: Array
wk: Array
wv: Array
wo: Array
w1: Array
w2: Array
class TinyDecoder:
def __init__(
self,
*,
vocab_size: int = 29,
d_model: int = 24,
n_heads: int = 3,
n_layers: int = 2,
max_positions: int = 512,
seed: int = 7,
) -> None:
if d_model % n_heads:
raise ValueError("d_model must be divisible by n_heads")
self.vocab_size = vocab_size
self.d_model = d_model
self.n_heads = n_heads
self.head_dim = d_model // n_heads
rng = np.random.default_rng(seed)
scale = 1.0 / np.sqrt(d_model)
self.token_embed = rng.normal(0.0, scale, (vocab_size, d_model))
self.position_embed = rng.normal(0.0, scale, (max_positions, d_model))
self.layers: List[LayerWeights] = []
for _ in range(n_layers):
self.layers.append(
LayerWeights(
wq=rng.normal(0.0, scale, (d_model, d_model)),
wk=rng.normal(0.0, scale, (d_model, d_model)),
wv=rng.normal(0.0, scale, (d_model, d_model)),
wo=rng.normal(0.0, scale, (d_model, d_model)),
w1=rng.normal(0.0, scale, (d_model, 2 * d_model)),
w2=rng.normal(0.0, scale, (2 * d_model, d_model)),
)
)
self.unembed = rng.normal(0.0, scale, (d_model, vocab_size))
def _split_heads(self, x: Array) -> Array:
# [time, d_model] -> [heads, time, head_dim]
return x.reshape(x.shape[0], self.n_heads, self.head_dim).transpose(1, 0, 2)
@staticmethod
def _softmax(scores: Array) -> Array:
shifted = scores - np.max(scores, axis=-1, keepdims=True)
exp = np.exp(shifted)
return exp / np.sum(exp, axis=-1, keepdims=True)
def _full_attention(self, x: Array, layer: LayerWeights) -> tuple[Array, Array, Array]:
q = self._split_heads(x @ layer.wq)
k = self._split_heads(x @ layer.wk)
v = self._split_heads(x @ layer.wv)
scores = (q @ k.transpose(0, 2, 1)) / np.sqrt(self.head_dim)
t = x.shape[0]
causal = np.triu(np.ones((t, t), dtype=bool), k=1)
scores = np.where(causal[None, :, :], -np.inf, scores)
weights = self._softmax(scores)
mixed = weights @ v
merged = mixed.transpose(1, 0, 2).reshape(t, self.d_model)
return merged @ layer.wo, k, v
def _incremental_attention(
self,
x: Array,
layer: LayerWeights,
cache: dict[str, Array],
) -> tuple[Array, dict[str, Array]]:
# x has one row. No explicit triangular mask is needed because the cache
# contains only legal past positions plus the current position.
q = self._split_heads(x @ layer.wq)
k_new = self._split_heads(x @ layer.wk)
v_new = self._split_heads(x @ layer.wv)
k_all = np.concatenate([cache["k"], k_new], axis=1)
v_all = np.concatenate([cache["v"], v_new], axis=1)
scores = (q @ k_all.transpose(0, 2, 1)) / np.sqrt(self.head_dim)
weights = self._softmax(scores)
mixed = weights @ v_all
merged = mixed.transpose(1, 0, 2).reshape(1, self.d_model)
return merged @ layer.wo, {"k": k_all, "v": v_all}
@staticmethod
def _block(x: Array, attn: Array, layer: LayerWeights) -> Array:
x = x + attn
ff = np.tanh(x @ layer.w1) @ layer.w2
return x + ff
def full_forward(self, tokens: list[int]) -> Array:
positions = np.arange(len(tokens))
x = self.token_embed[np.asarray(tokens)] + self.position_embed[positions]
for layer in self.layers:
attn, _, _ = self._full_attention(x, layer)
x = self._block(x, attn, layer)
return x @ self.unembed
def prefill(self, tokens: list[int]) -> tuple[Array, list[dict[str, Array]]]:
positions = np.arange(len(tokens))
x = self.token_embed[np.asarray(tokens)] + self.position_embed[positions]
caches: list[dict[str, Array]] = []
for layer in self.layers:
attn, k, v = self._full_attention(x, layer)
caches.append({"k": k.copy(), "v": v.copy()})
x = self._block(x, attn, layer)
return (x @ self.unembed)[-1], caches
def decode_one(
self,
token: int,
*,
position: int,
caches: list[dict[str, Array]],
) -> tuple[Array, list[dict[str, Array]]]:
x = self.token_embed[[token]] + self.position_embed[[position]]
updated: list[dict[str, Array]] = []
for layer, cache in zip(self.layers, caches, strict=True):
attn, new_cache = self._incremental_attention(x, layer, cache)
updated.append(new_cache)
x = self._block(x, attn, layer)
return (x @ self.unembed)[0], updated
def generate_full(model: TinyDecoder, prompt: list[int], steps: int) -> list[int]:
seq = prompt.copy()
logits = model.full_forward(seq)[-1]
for _ in range(steps):
token = int(np.argmax(logits))
seq.append(token)
logits = model.full_forward(seq)[-1]
return seq
def generate_cached(
model: TinyDecoder, prompt: list[int], steps: int
) -> tuple[list[int], float]:
seq = prompt.copy()
logits, caches = model.prefill(seq)
max_delta = float(np.max(np.abs(logits - model.full_forward(seq)[-1])))
for _ in range(steps):
token = int(np.argmax(logits))
seq.append(token)
logits, caches = model.decode_one(token, position=len(seq) - 1, caches=caches)
reference = model.full_forward(seq)[-1]
max_delta = max(max_delta, float(np.max(np.abs(logits - reference))))
return seq, max_delta
def score_pair_counts(prompt_tokens: int, generated_tokens: int) -> tuple[int, int]:
"""Attention score pairs per layer and query head for two execution plans.
Full recompute runs a complete causal forward pass before every sampled token.
Cached decode prefills once, then evaluates one new query row per later token.
"""
full = sum((prompt_tokens + j) * (prompt_tokens + j + 1) // 2 for j in range(generated_tokens))
cached = prompt_tokens * (prompt_tokens + 1) // 2
cached += sum(prompt_tokens + j for j in range(1, generated_tokens))
return full, cached
def timed_generation(model: TinyDecoder, prompt: list[int], steps: int) -> tuple[float, float]:
start = perf_counter()
generate_full(model, prompt, steps)
full_ms = (perf_counter() - start) * 1000
start = perf_counter()
generate_cached(model, prompt, steps)
cached_ms = (perf_counter() - start) * 1000
return full_ms, cached_ms
def main() -> None:
model = TinyDecoder()
prompt = [2, 5, 1, 8, 3, 13, 21, 1]
steps = 32
full_sequence = generate_full(model, prompt, steps)
cached_sequence, max_delta = generate_cached(model, prompt, steps)
assert full_sequence == cached_sequence
assert max_delta < 1e-10
# Negative case: reuse the exact cache but assign the next token the wrong
# absolute position. The causal history is intact, yet the logits change.
next_token = full_sequence[len(prompt)]
good_prefix = prompt + [next_token]
good_logits = model.full_forward(good_prefix)[-1]
_, caches = model.prefill(prompt)
wrong_logits, _ = model.decode_one(
next_token,
position=len(prompt) + 1, # should be len(prompt)
caches=copy.deepcopy(caches),
)
broken_delta = float(np.max(np.abs(good_logits - wrong_logits)))
assert broken_delta > 1e-4
full_pairs, cached_pairs = score_pair_counts(128, 128)
small_full_ms, small_cached_ms = timed_generation(model, prompt, 32)
long_full_ms, long_cached_ms = timed_generation(model, prompt * 32, 64)
print(f"positive: identical generated tokens = {full_sequence == cached_sequence}")
print(f"positive: max absolute logit delta = {max_delta:.3e}")
print(f"negative: wrong-position logit delta = {broken_delta:.3e}")
print(f"score pairs, full recompute = {full_pairs:,}")
print(f"score pairs, cached decode = {cached_pairs:,}")
print(f"pair-count reduction = {full_pairs / cached_pairs:.1f}x")
print(f"small workload, full/cached = {small_full_ms:.1f}/{small_cached_ms:.1f} ms")
print(f"long workload, full/cached = {long_full_ms:.1f}/{long_cached_ms:.1f} ms")
print("timings are machine-dependent; parity and pair counts are the assertions")
if __name__ == "__main__":
main()
A release gate for cache and decoding changes
The decision instrument below separates semantic proof from performance evidence. A team may release an approximate method, but it should do so under an explicit quality envelope rather than an “optimisation” label that hides changed behaviour.
Figure 8. Three rails must reach release together
- Name the intervention. Separate tensor reuse, memory layout, scheduling, proposal-and-verify, quantisation and eviction.
- Declare the equivalence claim. Exact logits, tolerance-bounded logits, identical greedy tokens or bounded task-level degradation are different contracts.
- Build cache identity. Include every attribute that can change K, V, mask or position.
- Run the positive parity case. Compare cached and uncached execution across layers, lengths, batches and prefix hits.
- Run negative cases. Wrong position, stale adapter, truncated block, changed mask and cross-request handle must be rejected or visibly diverge.
- Measure phase-specific performance. Report time to first token, token latency, goodput, live bytes and pressure actions.
- Slice quality where approximation enters. Long-range retrieval, exact copying, code, numbers and task-specific dependencies can fail differently.
- Test invalidation and recovery. A model update, worker restart or cache-manager failover must not revive incompatible state.
Optional derivation: why only the newest query row is needed
At generation step t, logits for positions before t have already served their purpose. Causality guarantees that adding token t cannot alter earlier hidden states. Future predictions need earlier positions only through each layer's keys and values. Therefore the runtime can retain those K and V tensors, compute the current token's query, key and value, and produce only the newest hidden state. Architectures with recurrent, convolutional, sliding-window or state-space layers need layer-specific state rules; a generic full-attention cache assumption is not sufficient.
Compact glossary
- Causal mask
- An attention constraint that blocks a position from reading future positions.
- Prefill
- The parallel processing of the supplied prompt that creates first-token logits and initial KV state.
- Decode
- The serial generation phase that extends the sequence, commonly one sampled token per iteration.
- KV cache
- Per-layer keys and values retained from already processed positions for later queries.
- Prefix cache
- KV blocks reused across requests that share an eligible token prefix and cache identity.
- Semantic parity
- Evidence that an execution change preserves the declared logits, tokens, distribution or bounded task behaviour.
Release speed-ups only after semantic parity
The causal mask and the KV cache solve different problems. The mask defines the conditional model by preventing future information from entering a next-token prediction. The cache exploits a consequence of that definition: once a past position has been processed, its per-layer keys and values do not change when future tokens arrive.
That distinction changes architecture review. Exact reuse, exact kernels and verified speculative methods can be assessed as alternative execution plans. Quantisation, eviction and restricted windows must be assessed as changed models or changed inference policies. Position mistakes, stale revisions and mixed request state are defects, even when the generated prose looks reasonable.
The decision is to make every decoding optimisation declare its fidelity class, cache identity and invalidation rule before performance claims are accepted. Then prove the stated parity, inject the failure cases, and measure prefill, token cadence, throughput and memory under the same workload. Faster generation is valuable. Knowing whether it is still the same generation is the engineering result.
sequenceDiagram participant Q as Question participant M as Mechanism participant E as Evidence participant D as Decision Q->>M: Three decoders continue the same sentence M->>E: Figure 1. The triangle that makes next-token learning legal E-->>M: compare rival explanations E->>D: Sources D-->>Q: revise scope or act
Sources
- Primary Ashish Vaswani et al., Attention Is All You Need. Transformer attention, decoder masking and autoregressive information flow.
- Primary Noam Shazeer, Fast Transformer Decoding: One Write-Head Is All You Need. Memory-bandwidth motivation for multi-query attention.
- Primary Joshua Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query quality and speed trade-off.
- Primary Reiner Pope et al., Efficiently Scaling Transformer Inference. Large-model inference partitioning and KV-memory implications.
- Primary Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Exact attention through IO-aware tiling.
- Primary Woosuk Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention. Paged KV allocation and sharing.
- Primary Yaniv Leviathan, Matan Kalman and Yossi Matias, Fast Inference from Transformers via Speculative Decoding. Distribution-preserving proposal and verification.
- Primary Guangxuan Xiao et al., Efficient Streaming Language Models with Attention Sinks. Failure of naive windows and attention-sink retention.
- Primary Zirui Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. Evaluated low-bit KV compression.
- Primary Lianmin Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs. RadixAttention and prefix-state reuse.
- Primary Amey Agrawal et al., Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve. Chunked prefill and decode scheduling.
- Primary Yinmin Zhong et al., DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving. Separate phase objectives and interference.
- Primary limitation study Wei Gao et al., Rethinking Key-Value Cache Compression Techniques for Large Language Model Serving. Production-oriented performance and sample-level limitations.
- Official implementation Hugging Face Transformers, Caching. Cache shapes, updates, inference use and position handling.
- Official implementation PyTorch, Scaled dot-product attention. Causal-mask interface and execution contract.
- Official implementation vLLM, Automatic Prefix Caching. Reuse of KV blocks for shared prefixes.
- Official implementation vLLM, KV cache manager API. Prefix-cache reset and lifecycle operations.
- Official implementation NVIDIA TensorRT-LLM, KV cache reuse. Page sharing for repeated prompt prefixes.