fix(speculative): draft KV cache rewind trims two entries too many, starving the drafter #994

Problem

SpeculativeGenerator::generate rewinds the draft model's KV cache by two entries too many on every round that rejects. The shortfall is never recovered, so the drafter's context falls monotonically behind the emitted sequence and its proposals degrade until speculation stops paying for itself.

In src/lib/mlxcel-core/src/speculative/mod.rs, step 4 of the decode loop:

let rejected = draft_tokens.len() - accepted;
if rejected > 0 {
    trim_caches(&mut self.main_caches, rejected as i32);
    let draft_trim = (rejected + 1) as i32;
    trim_caches(&mut self.draft_caches, draft_trim.min(draft_tokens.len() as i32));
}

Why rejected + 1 is wrong

The draft loop forwards num_draft single tokens: current_token, d_0, ..., d_{n-2} (the input to iteration j is the token iteration j-1 produced). So the draft cache gains n entries per round.

At the end of the round the generator has emitted d_0..d_{a-1} (the a accepted draft tokens) plus one replacement r, and r becomes the next round's current_token. For the next round's first draft forward of r to be correctly conditioned, the draft cache must hold everything up to and including d_{a-1}, that is a + 1 new entries.

Correct trim: n - (a + 1) = rejected - 1. Actual trim: min(rejected + 1, n).

For num_draft = 4 the cache ends each round behind by:

At accepted == 0 the draft cache does not advance at all: n forwards followed by a trim of min(n + 1, n) == n.

Consequences

1. Positive feedback loop. The further behind the draft cache is, the worse the drafter's proposals, the lower the acceptance, the more rounds per emitted token, and the faster it falls behind. Once acceptance is low the drafter is effectively predicting from a frozen prefix.
2. Depressed acceptance. Measured per-position acceptance on mlx-community/Meta-Llama-3.1-8B-Instruct-4bit with a Llama-3.2-1B-Instruct drafter, --num-draft-tokens 4 --temp 0.7, is 0.12 to 0.46 depending on trajectory. Published speculative-decoding acceptance for a comparable pair is 0.6 to 0.8. The gap is consistent with the drafter conditioning on a truncated prefix.
3. Uninterpretable A/B measurements. Because the drift rate depends on the acceptance rate, two runs that diverge early end up in very different regimes. Two 200-token runs on the same pair produced closed_form_sum_min of 0.221 and 0.396, a 1.8x spread on a quantity that should be a property of the model pair. Any speculative benchmark on this path is dominated by that variance.

Evidence

draft_cache_falls_behind_the_main_cache_every_round_pre_existing_defect in src/lib/mlxcel-core/src/speculative/distribution_tests.rs pins the current behavior: with a drafter that never agrees with the target, after 24 emitted tokens the main cache holds 24 entries and the draft cache is still frozen at 1 (the prompt). The test asserts the wrong values deliberately, so a fix produces a failing test rather than silence.

Scope notes

• Predates Distribution-preserving speculative-decoding acceptance (chain speculative sampling) #902 and is independent of it. The trim arithmetic is untouched by that change, and it affects the sampler-match, greedy-argmax and stochastic acceptance rules identically.
• Fixing it changes token streams at every temperature, including temperature 0, because the drafter will propose differently. It therefore needs its own before/after parity story and cannot ride along in an acceptance-rule PR.
• The all-accepted case is separately off by one in the other direction: with a == n the cache holds current_token, d_0..d_{n-2} and is missing d_{n-1}, while the next round forwards the bonus token. Worth resolving in the same pass.

Acceptance criteria

• Draft cache trim corrected so the draft cache holds exactly the prompt plus every emitted token before the next round's current_token, in all of the all-reject, mixed and all-accept regimes.
• A test asserts draft-cache offset exactly (replacing the current defect-pinning test), in the same two deterministic regimes the main-cache test already uses.
• Measured per-position acceptance on a Llama 8B / 1B pair at temperature 0.7 recorded before and after, using MLXCEL_SPECULATIVE_ACCEPT_DIAG=1 so closed_form_sum_min is reported alongside the measured rate.
• Temperature-0 output difference acknowledged and justified in the PR body, with the new stream verified against a drafterless greedy reference.