Milestone 5 of 7

Generate with and without a cache

Verify prefix logits, trace cache positions and shapes, and generate with declared policies and stopping rules.

Freeze the trained checkpoint and prove the incremental path preserves its logits before using it to generate tokens.

Milestone goal

Implement full-prefix and cached inference, compare every valid controlled prefix, record cache shapes and positions, then run one greedy and one seeded sampling policy with explicit stopping reasons.

Establish the uncached reference

For each prefix of lengths 1 through 4 in all four corpus patterns:

run the ordinary causal decoder on the complete prefix;
retain only its final logits row;
record the prefix IDs, decoded tokens, logits, and model mode.

This is the reference computation. Do not compare only selected tokens: two logit rows can share an argmax while differing substantially elsewhere.

Build the cache one consumed token at a time

At layer $\ell$ and consumed length $t$ , store

K^{(\ell)},V^{(\ell)}:(B,h,t,d_h).

For the frozen model, each K or V tensor has shape $(1,2,t,2)$ . A new token:

receives absolute position index $t$ ;
creates one query, key, and value per layer and head;
appends K and V exactly once along the sequence axis;
lets the new query attend over the $t+1$ valid cache entries;
produces one new logits row.

Each layer owns its cache. Each generation request owns a separate list of layer caches, token IDs, position state, random generator, and stopping state.

Prove equivalence at every prefix

For each prefix, calculate

\epsilon=\max_i|z^{full}_i-z^{cached}_i|.

The checked float32 CPU run has maximum error $4.7684\times10^{-6}$ across all 16 prefixes, below its declared tolerance $10^{-5}$ . The tolerance is a regression boundary for this environment, not a universal constant for every dtype or kernel.

Also assert every layer's cache shape and length. Include two deliberate failures: resetting every position to zero should produce a large mismatch, and feeding one layer another layer's cache should also fail. A test that never rejects a broken cache is weak evidence.

Keep selection policy outside the model

After equivalence passes, apply:

greedy decoding with a fixed tie rule;
temperature 1, top- $k=3$ sampling with seed 41.

Record original logits, transformed probabilities, chosen token, and random state. The checked greedy trace is

<bos> A B D <eos>

and the checked seeded trace is

<bos> Y X Y <eos>

These traces test execution and random-state control. They do not measure general language quality.

Distinguish returned tokens from consumed inputs

The decoder accepts positions 0 through 3. Its position-3 logits can return a fifth token such as <eos>, but that returned token cannot be fed back at position 4. Record eos and context as different stopping reasons.

Question. Set the next position

The frozen model's cache contains three consumed input positions. Which absolute position index should the next input token receive?

Your answer

Compute it first, then check your number.

Hint

The cache length equals the next zero-based index.

Solution

The next input uses position index 3, the model's last valid input position.

Diagnose the first mismatch

First differing value	Likely contract
token plus position record	position index or token ownership
new layer-1 K/V	projection, normalization, or axis order
appended cache length	missing or duplicate append
layer 1 matches, layer 2 differs	wrong layer's cache
logits match, selected token differs	policy, tie rule, or RNG state
one request changes another	cross-request cache contamination

Acceptance gate

Continue only when:

all 16 controlled prefixes pass full-versus-cached logit comparison;
every layer's K/V axes and length are asserted after every consumed token;
wrong-position and wrong-layer tests fail by more than tolerance;
greedy and seeded policies are reproducible from the recorded state;
every run records whether it stopped for EOS or context capacity.

Deliverable: prefix-equivalence table, per-step cache trace, deliberate-failure results, policy configuration, two token traces, probabilities, and stopping reasons.