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

  1. 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.
  2. A KV cache can preserve that contract exactly because past keys and values are fixed once their tokens and positions have been processed.
  3. The cache removes repeated projection and attention work during decoding, but trades compute for memory capacity, bandwidth and lifecycle complexity.
  4. Eviction, quantisation, wrong positions or stale cross-request reuse can change logits; a smaller cache is therefore not automatically the same model.
  5. 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

1→11→21→31→41→51→6 2→12→22→32→42→52→6 3→13→23→33→43→53→6 4→14→24→34→44→54→6 5→15→25→35→45→55→6 6→16→26→36→46→56→6
  • 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.
Conceptual figure. The lower triangle is the causal information boundary described in the Transformer decoder. No measured quantities are shown. The current-position diagonal is legal because the representation at that position contains the preceding token after the usual right shift.

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:

p(x₁:T) = ∏ₜ₌₁ᵀ p(xₜ | x₁, …, xₜ₋₁)

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:

Attention(Q,K,V) = softmax(QKᵀ / √dₖ + M)V

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.

Published evidence The Transformer formulation masks subsequent positions and shifts decoder outputs so a prediction at position i depends only on known outputs before i. Current PyTorch attention APIs also treat causality as an explicit execution hint, warning that an incorrect hint can produce incorrect execution (PyTorch documentation).

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.

The argument in one passChart · flow
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
A structural route through Causal Masks, KV Caches and Generation, from Three decoders continue the same sentence to Sources. Reading use: follow the relation, then test it against the surrounding evidence and limits.

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

Full recomputation compared with cached incremental decoding The upper path repeatedly recomputes triangular attention over a growing prefix. The lower path prefills once and then appends one key and value column for each new token while computing only one query row. Full recomputation Exact KV cache prompt + token 1 + token 2 prefill once one new row one new row append K,V append K,V Same legal past. Different amount of repeated work.
Signature figure. Conceptual execution cutaway, not a timing measurement. Triangles represent all legal query-key pairs recomputed by a naive full pass. Indigo rows represent the single current query evaluated after prefill; narrow columns are newly appended keys and values.

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 levelWhat must matchWhat it supportsWhat it does not prove
BitwiseEvery stored value and output bitStrict replay on a fixed stackPortability across kernels or hardware
Logit toleranceNext-token logits within a declared numerical boundImplementation parity under floating-point variationIdentical sampled tokens without controlled randomness
DistributionalSampling probabilities or a justified divergence boundClaims about stochastic generationEqual task outcomes on rare cases
SequenceGreedy or seeded sampled token sequenceRegression testing of a decoding policyUnchanged probability mass below the chosen path
OutcomeTask metric inside a declared envelopeUse-case acceptance of an approximationThat 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:

Cfull = Σⱼ₌₀ᴳ⁻¹ (P+j)(P+j+1)/2 Ccache = P(P+1)/2 + Σⱼ₌₁ᴳ⁻¹ (P+j)

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

Synthetic attention score-pair counts for cached and full-recompute generation A coral curve for full recomputation rises much faster than a teal curve for cached decoding as output length grows from 1 to 128 tokens after a 128-token prompt. 1326496128 00.6M1.2M1.8M2.4M generated tokens after a 128-token prompt full recompute: 2,446,656 cached: 32,640 75.0× pair-count difference
Synthetic operation counts from the companion artefact for P = 128 and up to G = 128. Counts cover attention score pairs per layer and query head, not total FLOPs or measured latency. The cached curve is dashed as well as teal.

The bill moves from compute to memory

Exact reuse is not free. For a conventional cache, approximate storage in bytes is:

B = 2 × L × T × Hkv × dh × b

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

KV cache anatomy for a synthetic 32-layer grouped-query model Thirty-two layer strata each contain key and value bands. A side ruler shows 8192 token columns, 8 key-value heads and 128 dimensions at two bytes, totalling one gibibyte per sequence. 32 per-layer cache strata KV eight visible bands stand for 32 layers tokens8,192 KV heads8 head width128 element size2 bytes 1 GiB per active sequence before allocator overhead
Worked scenario with synthetic dimensions. The 1 GiB value follows directly from the displayed formula and uses binary units. It is not a claim about a named model or deployment.

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

Practitioner synthesis from primary systems and compression work. Categories express the logical burden of proof, not a ranking. Exact methods can still differ numerically across kernels, hardware and floating-point order.

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.

Failure boundary Exact KV reuse ceases to help when the workload does not repeatedly decode from the same causal prefix, when state cannot be retained, or when cache movement and bookkeeping cost exceed avoided computation. Encoder-only inference, one-shot scoring and tiny low-concurrency workloads can sit outside the favourable region.

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

Failure surface for cache identity and positional correctness A contour surface is safe only where both cache identity and positional correctness are high. Stale models, mixed requests and wrong rotary positions move execution into higher logit-delta regions. parity basin plausible output, changed logits wrong position stale or mixed cache cache identity correctness → position and mask correctness →
Illustrative failure surface, not measured data. The safe region requires matching model state, token prefix, position system and attention policy. Hatched coral indicates risk as well as colour.
How the claim changes under pressureChart · state
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 --> [*]
The revision loop begins with Figure 1. The triangle that makes… and keeps counterexamples inside the reasoning path. Reading use: follow the relation, then test it against the surrounding evidence and limits.

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

Practitioner-synthesis lifecycle for the synthetic service. It extends the state-management ideas in paged and prefix-caching systems with explicit identity, tenancy and invalidation controls. It is not a report of a client deployment.

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 fieldWhy it belongsFailure if omitted
Model and weight revisionProjection matrices create K and VOld tensors are interpreted under new weights
Adapter or fine-tune selectionLoRA and related changes can alter layer outputsTwo configured models share an invalid prefix
Prefix token IDs and lengthK and V are functions of the actual tokens and orderText-normalisation or boundary differences collide
Position contractAbsolute, rotary or relative position state affects attentionCorrect tokens occupy the wrong geometry
Attention policy by layerFull, sliding and hybrid layers expose different historiesA block is reused outside its legal window
KV representationHead grouping, dtype and quantisation define storage semanticsReaders decode bytes under the wrong layout or scale
Tenant and policy scopePermission to share is separate from computational equalityPrivate 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.

MeasurePrimary questionTypical confounderRequired slice
Time to first tokenHow quickly can the prompt be admitted and prefetched?Queueing, prefix-hit rate, chunkingPrompt length and hit or miss
Time per output tokenHow stable is the decode cadence?Batch churn, KV bandwidth, prefill interferenceActive batch and context length
Goodput or throughputHow much demand meets the latency objective?Unbounded batching can inflate raw throughputDeclared service-level thresholds
Peak and live KV bytesHow many sequences fit without pressure action?Fragmentation and shared-page accountingAllocator state and concurrency
Parity or qualityDid the optimisation alter logits, tokens or task outcomes?Sampling randomness and aggregate maskingDeterministic 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.

Assumptions and expected interpretation The lab uses greedy decoding, learned absolute position embeddings, dense attention and no dropout. Passing the positive case demonstrates the cache logic for this specimen, not all runtimes. The negative case shows that an intact tensor cache is still invalid when its position contract is wrong.
positive: identical generated tokens = True positive: max absolute logit delta = 1.776e-15 negative: wrong-position logit delta = 3.682e+00 score pairs, full recompute = 2,446,656 score pairs, cached decode = 32,640 pair-count reduction = 75.0x small-workload timing may favour full recompute longer-workload timing may favour cached decode timings are machine-dependent; parity and pair counts are the assertions
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

Decision instrument. Green cells are required pass evidence, saffron cells demand workload-specific thresholds, and the coral invalidation cell is a hard stop if stale state can survive a relevant change. No performance numbers are implied.
  1. Name the intervention. Separate tensor reuse, memory layout, scheduling, proposal-and-verify, quantisation and eviction.
  2. Declare the equivalence claim. Exact logits, tolerance-bounded logits, identical greedy tokens or bounded task-level degradation are different contracts.
  3. Build cache identity. Include every attribute that can change K, V, mask or position.
  4. Run the positive parity case. Compare cached and uncached execution across layers, lengths, batches and prefix hits.
  5. Run negative cases. Wrong position, stale adapter, truncated block, changed mask and cross-request handle must be rejected or visibly diverge.
  6. Measure phase-specific performance. Report time to first token, token latency, goodput, live bytes and pressure actions.
  7. Slice quality where approximation enters. Long-range retrieval, exact copying, code, numbers and task-specific dependencies can fail differently.
  8. 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.

The evidence conversationChart · sequence
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
Question, mechanism, evidence and decision exchange the burdens developed in Causal Masks, KV Caches and Generation. Reading use: follow the relation, then test it against the surrounding evidence and limits.

Sources

  1. Primary Ashish Vaswani et al., Attention Is All You Need. Transformer attention, decoder masking and autoregressive information flow.
  2. Primary Noam Shazeer, Fast Transformer Decoding: One Write-Head Is All You Need. Memory-bandwidth motivation for multi-query attention.
  3. Primary Joshua Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query quality and speed trade-off.
  4. Primary Reiner Pope et al., Efficiently Scaling Transformer Inference. Large-model inference partitioning and KV-memory implications.
  5. Primary Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Exact attention through IO-aware tiling.
  6. Primary Woosuk Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention. Paged KV allocation and sharing.
  7. Primary Yaniv Leviathan, Matan Kalman and Yossi Matias, Fast Inference from Transformers via Speculative Decoding. Distribution-preserving proposal and verification.
  8. Primary Guangxuan Xiao et al., Efficient Streaming Language Models with Attention Sinks. Failure of naive windows and attention-sink retention.
  9. Primary Zirui Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. Evaluated low-bit KV compression.
  10. Primary Lianmin Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs. RadixAttention and prefix-state reuse.
  11. Primary Amey Agrawal et al., Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve. Chunked prefill and decode scheduling.
  12. Primary Yinmin Zhong et al., DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving. Separate phase objectives and interference.
  13. 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.
  14. Official implementation Hugging Face Transformers, Caching. Cache shapes, updates, inference use and position handling.
  15. Official implementation PyTorch, Scaled dot-product attention. Causal-mask interface and execution contract.
  16. Official implementation vLLM, Automatic Prefix Caching. Reuse of KV blocks for shared prefixes.
  17. Official implementation vLLM, KV cache manager API. Prefix-cache reset and lifecycle operations.
  18. Official implementation NVIDIA TensorRT-LLM, KV cache reuse. Page sharing for repeated prompt prefixes.