perf(speculative): prompt lookup slows prose decoding 3-21% on GB10

[inureyes]

Problem / Background

Prompt lookup (mlxcel generate --prompt-lookup, PR #2074) speeds up replies that copy their prompt but still costs throughput on replies with little to copy. The PR measured break-even (1.01x) on free writing on an M4 Pro; on GB10 the same kind of reply runs below plain decoding.

Measured 2026-10-02 on GB10 (CUDA), release build of the PR #2074 head after its hardening commit (a1fa9e7); greedy, whole-call tok/s including prefill, median of 3 runs, run-to-run spread mostly under 2%. Speedup over plain decoding:

Model	edit_fn	edit_json	summary	write	story (800 tok)
Qwen3-1.7B-4bit	1.10x	1.37x	0.79x	1.06x	0.84x
Qwen3-4B-4bit	1.21x	1.55x	1.32x	0.93x	0.87x
Qwen3-8B-4bit	1.36x	1.81x	1.50x	0.97x	0.90x
Llama-3.1-8B-Instruct-4bit	1.51x	1.94x	0.92x	0.96x	0.92x

Current Behavior

The losing rows share few drafted rounds, almost all rejected (story: acceptance 0.07 to 0.18, tokens per forward 1.02 to 1.06; 35 to 73 drafted rounds of 371 to 748). Each drafted round pays three costs in PromptLookupGenerator::generate (src/lib/mlxcel-core/src/speculative/prompt_lookup.rs):

• A verify block wider than one token, run synchronously: the host reads the argmax before building the next round (prompt_lookup.rs:761-831).
• A broken pipeline: when a proposal appears while a plain step is in flight, the loop submits nothing new, reads the in-flight token synchronously, emits it, and looks the proposal up again next round (prompt_lookup.rs:724-746).
• A slow restart: plain_streak resets to 0, and pipelining resumes only once it exceeds PLAIN_ROUNDS_BEFORE_PIPELINE (2), so two synchronous plain rounds follow every drafted round (prompt_lookup.rs:254, prompt_lookup.rs:747-759).

On GB10 a forced-synchronous plain step costs 17% on Qwen3-1.7B (MLXCEL_FORCE_SYNC=1: 180 to 150 tok/s), and the 1.7B story works out to about 2.4 plain steps lost per drafted round. The DraftGovernor constants (MISSES_BEFORE_COOLDOWN 3, cooldown MIN_COOLDOWN 4 doubling to MAX_COOLDOWN 32, budget 2 * ema + 2; prompt_lookup.rs:240-322) and the pipelining threshold were chosen from a 36-point grid on an M4 Pro, where verify and sync are relatively cheaper.

Proposed Solution

Candidate directions, each to be measured on both Metal and CUDA before any default changes:

• Keep the pipeline running when a proposal appears while a step is in flight: verify from the in-flight token instead of draining it.
• Require a minimum proposal length or recent acceptance before breaking the pipeline.
• Backend-specific governor thresholds derived from a measured verify-to-step cost ratio.

Decision criterion: ship the direction (or combination) that meets every acceptance criterion below on both backends; if more than one does, prefer the smaller change to the round loop.

Scope

In scope: the decode round loop and DraftGovernor in src/lib/mlxcel-core/src/speculative/prompt_lookup.rs, its tests in prompt_lookup_tests.rs, and the benchmark harness and results under docs/benchmark_results/.

Out of scope: first-use kernel costs per verify width, which made short replies much worse (0.65x) and were fixed in PR #2074 by warming every verify width before the timed run (warm_up_verify_widths, prompt_lookup.rs:544). Server-side prompt lookup is also out of scope.

Implementation Notes

• Reuse: the existing in_flight / async_eval pipelining in the same loop and the PromptLookupStats counters; keep MLXCEL_FORCE_SYNC as the kill switch (prompt_lookup.rs:696).
• Constraints: cache trim accounting (trim_caches with draft.len() - accepted) must stay exact if a verify starts from an in-flight token; samplers that need the emitted history (needs_history) stay synchronous.
• Edge cases: a proposal on the last allowed token (remaining <= 1), EOS read from the in-flight token, and --prompt-lookup-no-adaptive.

Acceptance Criteria

• On GB10, the write and story rows above reach at least 0.98x of plain decoding for all four models, with the edit and summary gains preserved within 5%.
• The same matrix on an Apple Silicon machine shows no regression against PR feat(speculative): add prompt-lookup decoding to mlxcel generate #2074's numbers.
• Greedy output stays byte-identical to the pre-change prompt-lookup output wherever it was identical to plain decoding.
• Unit tests in prompt_lookup_tests.rs cover any changed governor or pipelining rule.
• Measured once at the end on a quiet machine; the harness and raw numbers committed under docs/benchmark_results/.

Verification

make release-cuda          # GB10; `make release` on Apple Silicon
make test-fast-cuda FILTER="-p mlxcel-core --lib prompt_lookup"
./target/release/mlxcel generate --model <model> --temp 0 -n 400 -p "<prompt>"
./target/release/mlxcel generate --model <model> --temp 0 -n 400 -p "<prompt>" --prompt-lookup

Run the five prompts on all four models, median of 3, and diff the greedy outputs against the pre-change prompt-lookup outputs.

Refs #2074

[inureyes]

PR #2092 (merged as c6268bc3) adds DraftPolicy::Gated (default on CUDA builds; Metal and ROCm keep PR #2074's governor as graded). Record: docs/benchmark_results/prompt-lookup-governor-gb10-2026-10-02.md.

Against the acceptance criteria, on GB10 (final matrix at 34c627d7, greedy, median of 3, plain/graded/gated interleaved):

• Story rows reach 0.98x-1.00x on all four models (graded 0.83x-0.93x).
• Write rows: 0.98x-1.02x on Qwen3-1.7B, Qwen3-4B and Llama-3.1-8B; Qwen3-8B stays at 0.97x (0.96x-0.97x across three matrices, graded 0.97x-0.98x). A round trace attributes it to a short prompt fragment the reply restates; a stricter full-block threshold left it at 0.97x and broke parity on two rows, so it was reverted.
• Edit and summary gains kept: worst change vs graded is -4% (Qwen3-4B edit_json), most rows rise.
• Apple Silicon: not measured (no host reachable). The Metal default makes exactly PR feat(speculative): add prompt-lookup decoding to mlxcel generate #2074's decisions, so its numbers stand by construction; --prompt-lookup-policy gated is there to measure.
• Greedy parity: every reply graded keeps byte-identical to plain decoding stays identical.
• Unit tests cover the gated rules and probe alignment (39 pass, mutation-checked).
• Harness and raw results committed (load median 3.1 on a shared host; documented).

Left open for a decision: accept the Qwen3-8B email at 0.97x, or try the remaining lever (verify from the in-flight token, drop the two synchronous rounds after a drafted round), which needs its own parity pass because it moves verify-block starts.