Prove cached execution against the full prefix

TL;DR: Cached execution must preserve each query's visible prefix and position. With zero scores and values [2, 4, 8, 16, 32], queries at positions three and four return 7.5 and 12.4; restarting their mask at row zero returns 2 and 3.

Where you are. You have a state map. Now establish a small numerical reference before introducing scheduling changes.

Choose inputs that expose the boundary

Text that looks plausible is a weak cache test. A wrong mask can still produce a fluent continuation. Start with an attention calculation whose answer you can derive independently, then graduate to logits from a small model.

Our synthetic attention has one head and scalar values. Every query-key score is zero, so softmax assigns equal weight to every allowed key. Values at positions zero through four are 2, 4, 8, 16 and 32. The first three positions are already cached; the new chunk contains queries at absolute positions three and four.

The full-prefix implementation computes all five queries. Compare its final two rows with the two rows from cached execution. Their shapes differ, so comparing row zero with row zero would test different positions. Explicitly select the corresponding absolute positions.

Make the wrong implementation fail

A local two-row causal mask that starts at position zero admits one key for its first row and two keys for its second. On this fixture it returns 2 and 3. The failure is large and deterministic. It does not need a large model, random weights or a statistical tolerance.

def visible_means(values, positions):
    return [sum(values[:p + 1]) / (p + 1) for p in positions]

values = [2, 4, 8, 16, 32]
full_prefix = visible_means(values, range(5))
cached_chunk = visible_means(values, [3, 4])
wrong_origin = visible_means(values, [0, 1])
assert cached_chunk == full_prefix[-2:] == [7.5, 12.4]
assert wrong_origin == [2.0, 3.0]
print(cached_chunk, wrong_origin)

This is a mask oracle, not a complete attention implementation. The next test should use nonzero scores, several KV heads and a deliberately poisoned unused tail. If an invalid padded value changes the answer, the read boundary is wrong even when the nominal sequence length looks correct.

Separate position state from stored tensors

Cache reset has at least two responsibilities: release or invalidate the stored prefix and reset whatever determines future positions. Clearing K and V while leaving a position counter advanced can make an independent second request start at the previous request's offset. Conversely, resetting the position counter while retaining old KV can make new queries read an incompatible prefix.

For a model comparison, disable dropout and use evaluation mode. Compare corresponding logits with a justified numerical tolerance before comparing sampled text. Different kernel shapes can change rounding near a tie. A failed exact token match does not identify whether the cause was state corruption, sampling randomness or a small numerical difference. Preserve the first diverging logit row so you can investigate the right layer.

Do this before moving on

Run the small fixture, then change the cached prefix length from three to one. The new chunk starts at absolute position one, whose reference output is (2 + 4) / 2 = 3. Write a matrix with empty cache, one-token cache, a multistep append and reset between two independent requests.

For each case, specify which rows are compared and what input would make a wrong position visible. A completed matrix contains both a passing reference and a deliberately incorrect mask or reset that it rejects. Use the source pointer to find the corresponding counters in a real implementation.

Go deeper

• Attention with a KV cache defines the computation whose execution paths you are comparing.
• Paged KV memory adds physical layout changes without changing the logical prefix contract.
• Attention correctness and dispatch extends this check to grouped heads, poisoned padding and dispatch decisions.

Key takeaways

• Compare queries at matching absolute positions.
• Use a fixture whose wrong result you can predict before running it.
• Reset position state as well as stored cache tensors between independent requests.

Check yourself

Say or write your answer first, numbers included, then reveal. Producing it from memory is the skill an interviewer is checking.

1. 1Why does the wrong-origin fixture return 2 for its first query?

   Answer it, then reveal

2. 2A cache reset clears K and V but leaves absolute positions advanced. Is a second independent prompt correctly reset?

   Yes, empty tensors guarantee isolationNo, position state also belongs to the request