perf: stop server prompt-lookup rounds that cost more than they emit #2264

Problem / Background

#2255 (PR #2262) added --draft-kind prompt-lookup to mlxcel-server. Its criterion (a) required plain chat at B=1 with the flag on to be no worse than -1.0 % against the flag off. The end-of-run measurement missed it: -3.49 % (paired per-round range -3.73 to -3.22, every round negative, per question -0.6 to -6.1 %), with about 25 to 45 percent of proposed tokens accepted (for example 7 of 30, 40 of 112).

Setup: GB10, Qwen3-1.7B 4-bit, server defaults (--parallel 4, so the lone row is on paged storage), streamed /v1/chat/completions through the chat template (Qwen3 reasons first, so the reply copies little), 8 questions x 5 rounds, 256 tokens, flag on and off alternating. Report: docs/benchmark_results/server-prompt-lookup-gb10-2026-10-11.md, section (a). Scripts and data: docs/benchmark_results/data/server-prompt-lookup-gb10-2026-10-11/ (run3.sh and m_chat.py for (a), run.sh for (b) to (d), run2.sh for the paged 1024 to 4096 follow-up, analyze.py for the tables). Both land in a PR shortly; if they are not on main when you start, base on that PR.

The same flag gains +53 to +67 % on copy-heavy prompts (b) and 1.12 to 2.8x on long copy prompts (d). The fix must keep those.

Current Behavior

• Governor. DraftGovernor (src/lib/mlxcel-core/src/speculative/prompt_lookup.rs:299-440) runs DraftPolicy::Gated on CUDA builds (DraftPolicy::default_for_backend, :182-191). budget() (:354) returns GATED_NARROW_DRAFT (2) until a 0.5-weight EMA of accepted proposals reaches GATED_FULL_AT (1.5), then max_draft (7). record() (:379) pauses only on zero-accept rounds: at once while on probation, otherwise after MISSES_BEFORE_COOLDOWN (3) in a row. A row that keeps landing 1 of 2 or 2 of 7 never pauses. While paused, PromptLookupDrafter::draft_block (src/lib/mlxcel-core/src/speculative/prompt_lookup_drafter.rs:151-207) runs a ShadowProbe, and shadow_confirmed() resumes on probation after SHADOW_CONFIRM (2) proposed tokens come true. These rules were tuned on the CLI, where a verify round costs only its forward (width 3 = 1.76 plain steps, width 8 = 3.91 on Qwen3-1.7B; doc comment at prompt_lookup.rs:148-165).
• Server cost of a proposal (src/server/batch/scheduler/prompt_lookup.rs). With step n in flight, run_prompt_lookup_tick (:319) asks every gated-in row. If any proposes, it retracts the asks and commit_lookahead_without_prime (:390) commits step n with no prime (rule 3). The next tick is synchronous (prompt_lookup_sync_tick, :362). A prime is submitted again only when nobody proposed and prompt_lookup_blocks_prime (:231) finds every gated-in drafter's pipelines_plain_rounds() true, which needs plain_streak > PLAIN_ROUNDS_BEFORE_PIPELINE (three proposal-free asks; prompt_lookup_drafter.rs:47,209-211) unless the governor is paused. So each proposal costs its verify forward plus the lost overlap of the switch tick and the synchronous ticks before re-priming, and at 25 to 45 percent acceptance a round emits fewer tokens than that.
• Config. The CLI (src/commands/generate.rs:2049) and the server (prompt_lookup.rs:525,586) both call PromptLookupDrafter::new(config). The server's config comes from SpeculativeDispatch::resolve_prompt_lookup (src/server/speculative_dispatch.rs:369), which starts from PromptLookupConfig::default().

Proposed Solution

Add a server-only break-even rule: a row stops proposing when its recent emitted tokens per verify round (accepted + 1) fall below what a round costs on the server. The pause must go through the governor's existing paused state (drafting = false), because that gives both required properties for free: probes_while_paused() makes pipelines_plain_rounds() true at once, so the following ticks pipeline, and the shadow probe keeps checking, so a copy-heavy stretch resumes proposals.

Recommended shape (another design is fine if it meets every acceptance criterion and leaves the CLI unchanged): add an opt-in field to PromptLookupConfig (prompt_lookup.rs:116), for example min_tokens_per_round: Option<f64>, None in Default so mlxcel generate --prompt-lookup behaves exactly as today, and set it in resolve_prompt_lookup. In DraftGovernor keep an EMA of accepted + 1 per drafted round next to accepted_ema, and in record() pause once at least two drafted rounds back the EMA and it is below the threshold. Derive the threshold from the per-width cost table plus the server's lost overlap, confirm it on the (a) and (b) reruns, and document the value and its evidence in the constant's doc comment, as PROMPT_LOOKUP_PAGED_CONTEXT_LIMIT does.

• Resume churn. Prose confirms a two-token shadow probe often, so a break-even pause that resumes on the first confirmation pays a switch tick each time and may not recover the loss. Under the server rule, resuming must take more evidence than the CLI's (a longer server shadow probe, several confirmations in a row, or a minimum paused span; pick by the (a) and (b) numbers).
• Re-prime wait. Evaluate a shorter server-side PLAIN_ROUNDS_BEFORE_PIPELINE (it exists for CLI edits that miss for a token or two per field). If you change it, make it server-only through the same config path; either way, state the measured reason in the PR.
• Rejected: skipping draft_block in the scheduler for low-yield rows. A row that is not asked records nothing and runs no shadow probe (I5 in the module docs), so it would never resume on a copy stretch.

Scope

In scope: DraftGovernor and PromptLookupConfig (and PromptLookupDrafter if the re-prime wait changes), resolve_prompt_lookup and its startup summary, src/server/batch/scheduler_prompt_lookup_tests.rs, the governor unit tests, docs/speculative-acceptance.md, the prompt-lookup section of docs/CONTINUOUS_BATCHING.md, CHANGELOG.md, and a new benchmark report under docs/benchmark_results/.

Out of scope: the CLI's policy; DraftPolicy::Graded (non-CUDA default, unmeasured here; the server rule may leave it unchanged); DEFAULT_PROMPT_LOOKUP_MAX_BATCH (2) and PROMPT_LOOKUP_PAGED_CONTEXT_LIMIT (4096), unless the rerun shows they must move.

Implementation Notes

• Reuse: the governor's paused state and ShadowProbe; in the scheduler tests, the cycle_llama fixture and the decode_lookahead assertions of proposal_free_ticks_stay_pipelined and a_proposal_makes_the_next_tick_synchronous_without_a_prime.
• Constraints: invariants I1 to I8 in the scheduler/prompt_lookup.rs module docs stay as they are. PromptLookupConfig is public and built with ..PromptLookupConfig::default() in tests, so the new field needs a default that changes nothing. A break-even pause counts in paused_rounds, and the draft_n / draft_n_accepted timings stay correct.
• Edge cases: adaptive = false ignores the rule; a row with fewer drafted rounds than the minimum keeps today's behavior; at B=2 each row decides on its own evidence, but any proposing row still makes the tick synchronous for both, so the rule matters for (c) too.
• Error handling: no new failure paths; the rule is host-side arithmetic.

Acceptance Criteria

• Scheduler unit test in scheduler_prompt_lookup_tests.rs: a row whose verify rounds land below the threshold stops proposing and the following ticks are pipelined (decode_lookahead.is_some()); a copy-heavy row keeps verifying. The low-yield half fails on current main.
• Governor unit test: with the server threshold set, low-yield rounds pause; with PromptLookupConfig::default(), the same sequence follows today's Gated behavior.
• mlxcel generate --prompt-lookup behavior unchanged.
• (a) rerun with the same scripts (run3.sh / m_chat.py): flag on vs off no worse than -1.0 %.
• (b) copy-heavy B=1 and (d) paged/dense long copy keep at least 90 percent of their measured gains (+53 to +67 %; 2.69x, 2.79x and 1.42x paged at 1024 / 2048 / 4096; 2.10x dense at 8192).
• (c) concurrency 2 with the chat-template workloads (bench_serving_concurrency.py --prompt-style plain|copy): on no worse than off beyond the round spread.
• Measurements run once at the end, on a quiet GB10 under gpu-lock, and are committed under docs/benchmark_results/.
• One PR, one verification pass at the end, as few pushes as possible.

Verification

cargo fmt -- --check
cargo clippy -p mlxcel --lib --tests --features cuda -- -D warnings
cargo test --profile test-fast --features cuda --lib scheduler_prompt_lookup_tests -- --test-threads=1
cargo test --profile test-fast --features cuda -p mlxcel-core --lib speculative::prompt_lookup -- --test-threads=1
cargo build --release --features cuda --bin mlxcel-server
# end-of-run measurement (WT_PREFIX / MODELS as in the scripts; quiet host)
NEW_DATA=docs/benchmark_results/data/server-prompt-lookup-break-even-gb10-YYYY-MM-DD
mkdir -p "$NEW_DATA" && cp docs/benchmark_results/data/server-prompt-lookup-gb10-2026-10-11/{run.sh,run2.sh,run3.sh,m_chat.py,m_client.py,analyze.py,pick.sh} "$NEW_DATA"/
gpu-lock run --tag pl-break-even -- bash "$NEW_DATA"/run3.sh
gpu-lock run --tag pl-break-even -- bash "$NEW_DATA"/run.sh
gpu-lock run --tag pl-break-even -- bash "$NEW_DATA"/run2.sh

The scripts write their output next to themselves (D=$(dirname "$0")), so copy them into a new dated data directory first and run them from there, or the 2026-10-11 data is overwritten. Drop block A of run.sh (the #2229 single-stream arms, which need the base and prev builds); only its (b) to (d) blocks are needed. A pass is every test green and the new report's (a) to (d) tables meeting the thresholds above, measured against a flag-off server from the same build.