Cached and Uncached Logits Must Match

Treat cache equivalence as a numerical contract. Compare full-prefix and cached logits across every controlled prefix, declare dtype-aware tolerance, and use wrong-position and wrong-layer interventions to test the checker.

A KV cache is an algebraic reuse of earlier key and value projections. It is not an approximation in this chapter. For every valid prefix, compare the final logits row from:

the registered model running the complete prefix with a causal mask;
sequential cached execution processing each prefix token once.

Compare Every Prefix, Not One Final Sequence

The executable audit checks lengths 1 through 4 for all four corpus patterns. The largest absolute logit difference is

4.7684\times10^{-6},

below the frozen float32 tolerance $10^{-5}$ .

Prefix length	Cache shape per K or V tensor per layer	Expected comparison
1	$(1,2,1,2)$	one-token full and cached logits
2	$(1,2,2,2)$	final row for both paths
3	$(1,2,3,2)$	final row for both paths
4	$(1,2,4,2)$	final row for both paths

Small differences arise because the full and incremental implementations group floating-point operations differently. The tolerance is a tested regression boundary, not a universal constant for every dtype, device, or kernel.

Test Intermediate State when Final Logits Fail

Locate the first mismatch in this order:

token plus position record;
layer 1 normalized current record;
layer 1 new K/V and appended cache;
layer 1 attention and residual output;
corresponding layer 2 values;
final normalization;
tied readout logits.

Final equality alone can occasionally hide offsetting errors. Intermediate comparisons make the implementation easier to trust.

Q1. Interpret an equivalence table

Prefixes of lengths 1 and 2 differ by at most $10^{-7}$ , but length 3 differs by 0.8 immediately after position addition. What should you inspect first?

Choose the first check

Inspect the position index used for the third token.

Change the sampling temperature and top-k value.

Replace layer 2's cache before checking earlier state.

Select one choice, then check.

Hint

Follow the first differing intermediate, not the final logits.

Solution

Inspect whether the third input used position row 2 in both paths.

Run the Complete Audit

Download the generation and KV-cache script. Keep the Chapter 7 training script in the same directory, install PyTorch using its official selector, and run:

python generate-with-kv-cache.py

The script recreates the trained checkpoint, checks every prefix, runs two deliberate failures, reports cache memory, and performs greedy and seeded sampled generation.