Draft-model fallback for N-gram proposer misses (ngram_fallback = "draft")

[danielwinterw]

Draft: fall back to a draft model when the N-gram proposer misses

speculative.ngram_fallback = "draft" (alongside a configured draft model):
when the suffix/cache proposer has no candidates, the pipelined verify-window
path proposes from the draft model instead of degrading to one token per
round trip — the failure mode that dominates freeform text on high-RTT links.

• The draft session syncs incrementally to committed context + optimistic
  suffix (DraftRunner::sync_to_context); a full re-prefill only happens on
  divergence. Accepted fallback proposals extend the sync prefix, so
  consecutive fallback refills are cheap.
• Drafting feeds the same candidate pipeline at both the seed and refill
  points; the classic serial draft loop stays disabled while fallback drives
  the pipeline, and depth-1 setups keep classic behavior.
• Telemetry: speculative_fallback_draft_{proposals,tokens,ms}.

Draft because: unit-tested and suite-green, but not yet exercised
end-to-end with a real draft model — needs a freeform-workload bench run
(the N-gram-friendly re-emit workload never misses, so the fallback path
stays cold there).

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review.

[i386]

The config plumbing and gating look solid (validated enum, resolver bails when draft is set without an N-gram strategy, default false keeps all existing paths identical, and the classic serial draft loop correctly stays authoritative at depth 1).

One blocking ask: the new incremental-sync logic in DraftRunner is subtle and entirely untested. synced bookkeeping is now load-bearing for KV correctness across fallback proposals (prefix-extension vs full reset on divergence, plus propose pushing each accepted current), and a mistake here silently corrupts draft proposals at runtime. Please add unit coverage for at least:

1. sync_to_context extends incrementally when the target has the synced tokens as a prefix (assert no reset/prefill of the whole context),
2. sync_to_context falls back to a full reset on divergence,
3. one end-to-end path exercising ngram_fallback_draft = true with an N-gram miss (asserting fallback_draft_proposals > 0 and that output still matches the non-fallback run), or — if a model-backed test isn't feasible in CI — a test at the draft_runner level proving propose-after-sync produces the same tokens as propose-after-reset for the same context.

Minor, non-blocking: in the first fallback site the budget is ... .min(native_mtp_remaining) .max(2) — the .max(2) can exceed a remaining budget of 1; worth a comment or clamp if a downstream invariant depends on it.

[i386]

The .max(2) here can raise the budget above native_mtp_remaining when only 1 token of run-ahead budget is left — the floor wins over the cap. If a downstream invariant assumes the proposal fits the remaining window this overshoots by one; either clamp after the min (budget.min(remaining).max(1) semantics) or document that exceeding the remaining budget by one is benign here.

Also flagging inline: the two sync_to_context call sites (here and the refill path extending with optimistic_suffix()) are the new untested KV-bookkeeping path from my review — please add the DraftRunner prefix/divergence unit tests before merge.

[danielwinterw]

Both fixed in 16c2de0 (now c6038f2 after rebase).

The floor: it now applies before the remaining-window cap (.max(2).min(native_mtp_remaining)), and the fallback is skipped entirely when fewer than two tokens remain, so it falls through to the serial path instead of overshooting by one. I went with the clamp rather than documenting the overshoot as benign — the downstream window arithmetic assumes the proposal fits, and a one-token slack there is not worth the reasoning.

The sync bookkeeping: extracted DraftSyncState/DraftSyncPlan so the load-bearing decision is separable from the session I/O and testable without a model. Eight unit tests cover incremental extension when the target has the synced tokens as a prefix, the exact-sync no-op, divergence resetting (both a differing token and a synced run longer than the target, which is the rejected-proposal case), propose's accepted tokens joining the prefix so the next sync extends rather than resets, a rejected proposal step forcing a reset, record_reset replacing the prefix, and the Extend range indexing the caller's context slice correctly.

[i386]

Reviewed at origin/feat/ngram-draft-fallback against its base feat/runahead-verify-windows. Reviewing it as a draft — the "needs a freeform bench" caveat is the right call and I'm not treating the missing E2E run as a finding.

DraftSyncState is the good part of this PR. Splitting the plan out from the session I/O and unit-testing it separately is exactly right for something where a wrong answer is silent KV corruption rather than an error, and the test list covers the cases that matter (divergence, over-long session after a rejection, proposal steps joining the prefix).

Three things.

1. The seed path trusts an invariant the refill path is careful about

The refill site does this correctly:

let mut sequence = context_tokens.clone();
sequence.extend_from_slice(pipeline.optimistic_suffix());
if let Some(&last) = sequence.last() {
    draft.sync_to_context(&sequence)?;
    let draft_tokens = draft.propose(last, budget)?;

sync_to_context materializes all but the last token, and propose decodes from that exact last token. Self-consistent.

The seed site doesn't:

draft.sync_to_context(&context_tokens)?;
...
let draft_tokens = draft.propose(current, budget)?;

This is only correct if current == *context_tokens.last(). It is today — current starts as prompt_token_ids.last() and the two context_tokens.push(current) sites keep them in step — but that's an invariant maintained across ~1100 lines of decode loop plus fused_first_decode, held together by nothing but convention. If it ever slips, DraftSyncState records a prefix the session didn't materialize and every subsequent proposal is quietly drawn from a corrupt KV, which is the exact failure your own doc comment calls out:

claiming a prefix extension the session has not materialized silently corrupts every later proposal.

Cheapest fix is to just not depend on it — propose from *context_tokens.last() the way the refill path does. If you'd rather keep current for clarity, a debug_assert_eq!(context_tokens.last(), Some(&current)) before the sync makes the slip loud in test builds instead of silent in production.

2. ngram_fallback = "draft" with no draft model is a silent no-op

The resolver bails when the N-gram strategy is missing:

"draft" => {
    if config.ngram.is_none() {
        bail!("skippy speculative ngram_fallback = \"draft\" requires an N-gram strategy");
    }
    true
}

but nothing checks for a draft model, and the runtime gate is ... && draft_guard.is_some() && .... So a config with ngram_fallback = "draft" and no draft_model starts cleanly, reports nothing, and never once takes the fallback path — the operator gets baseline behaviour and a telemetry counter stuck at zero, with no way to tell that from "the proposer never missed."

SpeculativeConfig already carries draft_model, so the resolver can check it in the same place it checks ngram. Same for the verify_window.depth() > 1 requirement — a depth-1 config with ngram_fallback = "draft" is also a silent no-op, and that one is at least worth a warning since the comment says depth 1 is deliberate.

3. Small: propose records the step before it's known to have happened

for _ in 0..max_tokens {
    self.synced.record_proposal_step(current);
    current = self.session.decode_step(current)...?;

If decode_step fails, synced.tokens now claims a token the session may not hold. Unreachable in practice — the error aborts the request and the next request's reset_to_context clears the state — but the ordering is free to fix and removes the need to reason about it:

let stepped_from = current;
current = self.session.decode_step(current)...?;
self.synced.record_proposal_step(stepped_from);

---

Nothing here is structural. Once the bench run lands I'd want (1) resolved before merge; (2) and (3) are cheap enough to fold in now.

[danielwinterw]

All three in 2d47d28.

1. Seed path invariant. Took the cheap fix rather than depending on it — the seed path now proposes from *context_tokens.last(), the same token sync_to_context deliberately leaves unmaterialized, so the two sites are self-consistent for the same reason the refill path is. Kept a debug_assert_eq! against current as well, so if the decode loop's invariant ever does slip it's loud in tests rather than quietly gone.

2. Silent no-op. Both cases now bail in the resolver: ngram_fallback = "draft" requires speculative.draft_model, and requires depth > 1. I went with an error rather than a warning for depth 1 — a config that asks for a feature the depth-1 path can't provide is a config bug, and a warning in a log nobody reads reproduces the exact problem you're describing. Two resolver tests cover the messages.

3. propose ordering. Fixed as suggested — records stepped_from after decode_step returns.

Bench still owed before this leaves draft; that's the freeform-workload run, not an E2E smoke.