Verifying cached-prefix equivalence

Reuse of a cached prefix either produces the same computation a full pass would have produced or it does not, and nothing about a language model's output tells you which. A cache that returns nearly the right keys and values raises no exception; it produces slightly different logits, which produce a different token, which produces a different continuation, and the result is a fluent answer to a question nobody asked. Neither a health check nor a latency graph nor a human reading the output can distinguish a generation that has quietly forked from the one you would have got without the cache.

This chapter builds the instrument that can. It is an equivalence test between two paths through the same model: the cached path, which prefills a prefix, retains the per-layer key and value tensors, and then feeds a single appended token against them, and the full path, which runs the whole sequence through in one go. The two are supposed to produce identical logits for that appended token. Writing the test is worth doing even against a serving engine you did not implement, because it is the only check that answers "is my reuse exact" without appealing to anyone's documentation — and because when it fails, the diagnostic value is in which prefix lengths fail, which points at three specific bugs that have nothing to do with the cache itself.

A cached step and a full pass, side by side#

Two of the three bugs this test catches are about position, so it is worth establishing why a transformer has to be told about position at all. Attention is a weighted sum over the positions a token is allowed to look at, and a weighted sum does not care what order its terms arrive in. Shuffle the tokens of a sentence, feed them through unmodified attention, and every position sees the same collection of keys and values it saw before, with identical arithmetic. A model built that way would treat "the refund precedes the return" and "the return precedes the refund" as the same input, so something has to inject order into the vectors themselves.

Rotary position embeddings do it by rotation. Each query and key vector is split into pairs of components, and every pair is rotated through an angle proportional to the token's absolute index in the sequence, with different pairs rotating at different rates. A dot product between two rotated vectors depends on the difference between their angles, so the score between position i and position j ends up carrying i − j, which is the relative distance the model wants. The property to hold on to for the rest of this chapter is narrower than the mechanism. The rotation is applied using a number — the absolute index — that the surrounding code has to supply, and that number changes the vector. A key rotated for position 0 is not the key for position 512 with a label attached. It is a different vector, and nothing inside it records which angle it was given.

Strip attention down until nothing is left but the part that can be wrong. One head, rotary position embeddings — RoPE, which encodes where a token sits by rotating its query and key vectors through an angle proportional to that absolute position — a cache that is a pair of growing tensors, and two entry points into the same block: one that consumes a whole sequence and one that consumes a single appended token against a cache.

import torch
import math

D = 64  # head dim

def rope(x, positions, base=10000.0):
    """x: [seq, D] -> rotary-embedded. positions: [seq] of absolute indices."""
    half = D // 2
    inv = base ** (-torch.arange(0, half, dtype=torch.float32) / half)
    ang = positions[:, None].float() * inv[None, :]          # [seq, half]
    cos, sin = ang.cos(), ang.sin()
    a, b = x[..., :half], x[..., half:]
    return torch.cat([a * cos - b * sin, a * sin + b * cos], dim=-1)

def attend(q, k, v, causal):
    scores = (q @ k.transpose(-1, -2)) / math.sqrt(D)        # [q_len, k_len]
    if causal:
        q_len, k_len = scores.shape
        offset = k_len - q_len                               # keys already cached
        mask = torch.arange(k_len)[None, :] > (torch.arange(q_len)[:, None] + offset)
        scores = scores.masked_fill(mask, float("-inf"))
    return scores.softmax(dim=-1) @ v

The offset line is where most implementations go wrong, and writing the mask this way rather than as a lower-triangular constant is the point of showing it. During a full pass, q_len == k_len and offset is zero, giving the familiar triangle. During a cached step, q_len is one and k_len is the cache length plus one, so offset is the number of tokens already resident and the single query row is allowed to see all of them plus itself. A mask written as torch.tril(torch.ones(q_len, k_len)) is correct in the first case and silently wrong in the second — it lets the appended token attend to exactly one position, the first one in the sequence.

That last number is worth checking rather than accepting, because it is surprising enough to look like a typo. torch.tril keeps the entries on and below the main diagonal of a matrix and zeroes everything above it. In a cached step the mask has one row, for the single query, and as many columns as there are keys. The main diagonal of a one-row matrix is the single entry in column zero, so every column from one onwards sits above the diagonal and is zeroed. The appended token is therefore permitted to attend to the first token of the sequence and to nothing else — not to the two thousand tokens in between, and not even to itself. It will still produce a logit, that logit will still name a token, and the model will still continue fluently from whatever it names, which is why this belongs in a test rather than in a runtime assertion nobody wrote.

The block itself then has two callers over one body:

class Block(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.wq, self.wk, self.wv = (torch.nn.Linear(D, D, bias=False) for _ in range(3))

    def forward(self, x, cache=None, start_pos=0):
        """x: [n_new, D]. cache: (K, V) each [n_cached, D] or None."""
        n_new = x.shape[0]
        pos = torch.arange(start_pos, start_pos + n_new)

        q = rope(self.wq(x), pos)
        k = rope(self.wk(x), pos)          # RoPE applied BEFORE the cache write
        v = self.wv(x)

        if cache is not None:
            k = torch.cat([cache[0], k], dim=0)
            v = torch.cat([cache[1], v], dim=0)

        out = attend(q, k, v, causal=True)
        return out, (k, v)

And the test is three lines of intent:

def check(model, ids, prefix_len, atol=1e-2):
    full = model(ids[: prefix_len + 1])[-1]              # logits for the appended token
    _, cache = model(ids[:prefix_len])                   # prefill, keep K/V
    stepped = model(ids[prefix_len : prefix_len + 1],
                    cache=cache, start_pos=prefix_len)[-1]
    delta = (full - stepped).abs().max().item()
    assert full.argmax() == stepped.argmax(), f"argmax diverged, max|Δ|={delta:.2e}"
    assert delta < atol, f"max|Δ|={delta:.2e}"

The load-bearing assertion is the one on argmax — the index of the highest-scoring logit, which is the token a greedy decode emits. Raw logit deltas of a few thousandths are the expected state of the world, because attention kernels pick tile sizes based on shape and floating-point addition is not associative. The argmax is the quantity that decides whether the two paths produce the same text.

The non-associativity claim is load-bearing and easy to wave past, so it is worth spending three sentences on. A floating-point number carries a fixed number of significant digits, and every addition rounds its result to fit them. Add a large number to a small one and the small one's least significant digits fall off the end, where adding the two small ones together first would have let them survive into the total. (a + b) + c and a + (b + c) are therefore different computations on a real machine rather than two spellings of one.

That is what connects rounding to tile size. A GPU kernel summing a long row of products splits the row across threads and combines the partial sums in whatever order the tile layout dictates, and the tile layout is chosen from the shape of the matrices it was handed. A single query row against a long run of keys is a different shape from a full square of queries against those same keys, so it gets a different layout, a different summation order, and a total that differs in the last few digits. Nothing is broken in that case. Two paths are computing the same quantity with different rounding, which is exactly what a few thousandths of a logit looks like.

The assertion therefore belongs on the argmax and the raw delta belongs on a warning light. Sampling never reads a logit's absolute value; it reads the ordering, and at temperature 0 only the position of the maximum. A difference of one part in a thousand changes an ordering only when two candidates were already within one part in a thousand of each other, which does happen, and is why atol stays as a secondary check rather than being dropped. A test written with atol=1e-6 is asserting bit-identity across shapes, which no correct engine provides, and a team that watches it fail on green builds learns to ignore the one signal that would have caught a real divergence.

Reading a divergence#

When check fails, the instinct is to suspect the cache. That instinct is almost always wrong, because the cache is the one component that does nothing interesting — it is a concatenation, or in a real engine a memcpy into a preallocated block. It has no arithmetic to get wrong. What is wrong is one of the three things that behave differently depending on whether a position arrives inside a batch of 512 or on its own.

The instinct is tempting for a nameable reason, which is that the cache is the thing that changed. A test that passed before reuse was added and fails after it reads naturally as an indictment of the new component. But the change did not introduce any arithmetic. It introduced a second shape through arithmetic that was already there, and every branch the cached path exercises was previously only ever run with q_len == k_len. What the test is finding is code that was always wrong for a case nobody had run.

The positional row deserves its own paragraph because it is the most common and the most seductive. The cached path knows only that it has one new token; it does not know, unless you tell it, that this token sits at absolute position 512. Pass start_pos=0 by omission and RoPE rotates the appended token's query and key as though it were the beginning of the sequence. At a prefix of 64 tokens the resulting rotation error is small and the argmax survives, so the test passes and you ship. At 512 the phase error is large enough to reorder the top logits.

The layer-order row is subtler. Some implementations cache the key before rotation, on the reasoning that positions are absolute and can be applied on read, which is a legitimate design — but only if the read path applies the rotation with the same absolute indices, on every layer, including whichever layer someone later optimised by hand. Mixing the two conventions across layers gives a model that matches at layer 0 and drifts thereafter, which reads as "deep layers are unstable" and sends people looking at numerics.

Choosing the lengths to sweep#

A single call to check at a comfortable prefix length is close to worthless, because each of the three bug classes has a length regime where it hides. Sweeping is cheap — the model is small, the assertions are fast — so sweep deliberately rather than picking round numbers.

The two ends of that table catch different failures. Very short prefixes expose index arithmetic, because at a cache length of one an error of one position is an error of a hundred percent and the argmax will flip. Very long prefixes expose anything phase-dependent, because a rotary error that shifts a key by half a radian is invisible in the top logit at position 64 and decisive at position 8,192. The middle of the table is where most test suites live and where the least is learned.

The reason for that asymmetry falls straight out of how the rotation is built. The angle applied to a pair of components is the position multiplied by a rate, so an error in the position is multiplied by that same rate. Default start_pos to zero at a prefix of 64 and the fastest-rotating pairs are off by a small angle, which perturbs the scores without reordering them. Do it at a prefix of 8,192 and the same pairs are off by an angle large enough to point the key somewhere unrelated. The wrongness in the vector grows with the wrongness in the index, the wrongness in the index is the cache length, and so the defect is mildest in exactly the regime where a comfortable test sits. A single check at 200 tokens is a check on the case where the bug is smallest.

Run the sweep in CI against a small model, not a large one. The properties being tested are structural — mask offsets, position indices, per-layer ordering — and they do not depend on parameter count, so a two-layer randomly-initialised model with a 64-dimension head verifies the same logic in half a second that a 70B checkpoint would verify in ten minutes.

Checking an engine you did not write#

Most teams do not implement attention; they configure vLLM, SGLang, TensorRT-LLM, or a hosted API. The same equivalence question applies, and the same test shape works, with the cache disabled standing in for the full recompute.

import itertools, requests

def logprob_of_completion(prompt, engine_url, n=8):
    r = requests.post(engine_url + "/v1/completions", json={
        "model": MODEL, "prompt": prompt, "max_tokens": n,
        "temperature": 0.0, "logprobs": 1,
    })
    c = r.json()["choices"][0]
    return c["text"], c["logprobs"]["token_logprobs"]

# Server A launched with --enable-prefix-caching, server B without it.
for prefix_len in (1, 17, 200, 2048, 8192):
    p = build_prompt(prefix_len)
    ta, la = logprob_of_completion(p, SERVER_WITH_CACHE)
    _      = logprob_of_completion(p, SERVER_WITH_CACHE)     # ensure a hit
    tb, lb = logprob_of_completion(p, SERVER_NO_CACHE)
    drift = max(abs(x - y) for x, y in zip(la, lb))
    print(f"{prefix_len:>6}  text_match={ta == tb}  max|Δlogprob|={drift:.2e}")

Two servers rather than two request flags, because a per-request bypass on the same process may share allocator state and mask exactly the bug you are hunting. Greedy decoding rather than sampling, so any text difference is a real difference. And comparing the first eight tokens rather than one, because a divergence that starts at token four is a divergence — a cached path that agrees on the first token and then drifts is telling you something about how the newly-written blocks are being read back, not about the original prefill.

Against a hosted API you get less. You cannot disable the provider's caching directly, you rarely get logprobs, and you cannot control batching, so the strongest available assertion is a behavioural one: send the same greedy request many times, some hitting cache and some not according to the cache_read_input_tokens field the response reports, and check that the completions agree. That is a weaker test than the one above and it is still worth running once, because it converts "the provider says reuse is exact" into something you have observed.

Keying, bit-identity and the serving stack#

The test does not establish that the right prefix was reused. Every check in this chapter compares two computations over the same token IDs, so it verifies the mechanics of reuse and says nothing about whether the entry you loaded belonged to this request. That is a keying question, it is a different failure with a different signature, and it is the subject of chapter 4, on what a prefix cache keys on, and chapter 8, on entries keyed wrongly.

It also does not promise bit-identical output across runs, and you should not write it as though it did. A batch of 8 sequences reduces in a different order than a batch of 1, so differences of a few thousandths of a logit are expected, and a test with atol=1e-6 will fail on a correct engine and teach the team to ignore it. The tolerance on the raw delta is a secondary signal for spotting drift before it flips a decision; the argmax assertion is the one that decides whether the chapter's question has been answered.

The last thing it does not do is survive contact with a real serving stack unchanged. Everything so far has lived inside one process, where the cache is a Python tuple that exists for the duration of one generation and is garbage-collected when the request ends.

Equivalence, then, is a property you verify rather than one you assume, and it is cheap enough that there is no good reason to assume it. Prefill a prefix, step one token against the retained tensors, compare the logits with a full pass, and assert on the argmax; do it at zero, one, two, sixteen, sixty-four, two hundred, two thousand and eight thousand tokens, against a two-layer model in CI. A failure at that point is a causal-mask offset, a positional index, or a per-layer ordering mistake, in that order of likelihood, and the table above maps the observed length profile onto which of the three it is.

What the test cannot do is follow the tensors out of the process. A serving layer keeps them resident after the response has been sent, hands them to a request that arrives twenty minutes later from a different user, and has to decide when to throw them away — which turns a local data structure into a store with a lifecycle, a lookup key, a memory budget, and a line on the invoice.