perf: offer prompt lookup on the server without slowing plain decode #2255

Part of #2166.

Depends on #2229 (it reuses that issue's per-row submit/draw entries only if convenient; the shared verify round below is independent).

Split out of #2229, whose Part 2 this was. The maintainer asked that offering prompt lookup on the server must not slow plain decode, so this issue carries speed gates next to the correctness criteria. Implement everything in one PR, run the local gate once at the end, and keep pushes to a minimum.

Problem / Background

PromptLookupDrafter (src/lib/mlxcel-core/src/speculative/prompt_lookup_drafter.rs, DrafterKind::PromptLookup at src/lib/mlxcel-core/src/drafter/mod.rs:152) runs only in mlxcel generate --prompt-lookup (run_prompt_lookup, src/commands/generate.rs:2024, over DirectEngine::generate_with_drafter at src/lib/mlxcel-core/src/engine/speculative.rs:145, which uses Engine::verify / Engine::commit_appends at speculative.rs:53 and :73). The server cannot offer it: SpeculativeDispatch::resolve (src/server/speculative_dispatch.rs:223) returns Disabled when no --draft-model is set (:227), before it reads --draft-kind, and with a draft model it rejects prompt-lookup as unrecognised (:247). The burst arms in src/server/batch/speculative_burst.rs are per-family adapters around MtpGenerator / DFlashGenerator, and is_kind_specific (speculative_dispatch.rs:399) knows only MTP and DFlash.

Current Behavior

Speculative dispatch and the lookahead pipeline are mutually exclusive today. BatchScheduler::lookahead_params (src/server/batch/scheduler/decode_tick.rs:330) returns None whenever should_dispatch_speculative() (src/server/batch/scheduler/config.rs:623, which is is_kind_specific()), so a naive prompt-lookup arm that set that flag would make every plain token synchronous. The retired CLI loop measured the lookahead overlap at about 12 percent of decode time. DirectEngine already solves this for the CLI: pipelined_plain_round (src/lib/mlxcel-core/src/engine/speculative_plain.rs:45) pipelines one-token rounds once Drafter::pipelines_plain_rounds (drafter/mod.rs:1001, prompt-lookup impl at prompt_lookup_drafter.rs:198) says a proposal is unlikely, and when a lookup finds a proposal while a step is in flight it submits no next forward, reads the in-flight token, calls retract_draft (prompt_lookup_drafter.rs:202), and verifies synchronously on the next round, so the switch costs no forward.

Proposed Solution

• Flag: mlxcel-server and mlxcel serve accept --draft-kind prompt-lookup with no --draft-model (a draft model with it is SpeculativeDispatchError::InvalidKind: "prompt-lookup takes no draft model"). --draft-block-size N sets PromptLookupConfig::max_draft; every other field is PromptLookupConfig::default() (src/lib/mlxcel-core/src/speculative/prompt_lookup.rs:131). Add SpeculativeDispatch::PromptLookup { config: PromptLookupConfig }, check it before the draft_model_path early return, list prompt-lookup in the unrecognised-kind message, and print speculative=prompt-lookup (max_draft=N, policy=P) from summary() (speculative_dispatch.rs:318). is_kind_specific() stays false for PromptLookup, so should_dispatch_speculative() stays false and the burst path is never entered; the scheduler checks a separate prompt_lookup_config() accessor instead. mlxcel generate keeps --prompt-lookup; its --draft-kind rejection message (src/cli/speculative_args.rs:177) is unchanged.
• Shared round: factor the verify round out of DirectEngine::speculate (speculative.rs:213) into one Engine-level function that both callers use: given a sequence's StepRow, its current token and a proposal, run Engine::verify, sample or argmax every position through the per-row chain in order, commit_appends the accepted prefix plus the target's token, unwind_appends the rest, and return the emitted tokens with their RowOutcomes and the accepted count. Also factor the eligibility rule (SpeculativeRunError::NotTrimmable, SamplerFeedbackState, NeedsHiddenStates, plus prompt_lookup_unsupported_reason at prompt_lookup.rs:243) into one predicate the scheduler calls per sequence.
• Scheduler arm: at most one verify round per scheduler tick per prompt-lookup sequence, not a run-to-completion burst. At prefill completion an eligible sequence (text-only, no structured output, no logprobs, no mirostat / adaptive-p / extended chain, the shared eligibility predicate true) gets a PromptLookupDrafter on its SequenceInfo, bound and primed with prefill_from_target_hidden (prompt_lookup_drafter.rs:129). On each decode tick a row whose drafter proposes nothing, or whose verify round the gate below declines, joins the regular Engine::step batch; a row with an accepted-by-gate proposal (capped so it never passes max_tokens) runs the shared round on its own, its tokens go through the same finish step and streaming as every other row, and accept_verified_tokens reports the outcome. An ineligible request decodes plainly with a debug log naming the reason, as should_burst_for_sequence (speculative_burst.rs:398) does.
• Lookahead stays on for proposal-free ticks (replaces perf: pipeline penalty and DRY requests in the engine's B=1 client #2229's "keep the pipeline off under speculative dispatch"). Mirror pipelined_plain_round at batch level in run_decode_tick (decode_tick.rs:284) and pipelined_steady_decode (:581): (1) while a lookahead step is in flight, before submitting the step n+1 prime, ask every gated-in prompt-lookup row's drafter for a proposal on its committed tokens; (2) if none proposes, prime and pipeline exactly as today; (3) if any proposes, submit no prime, read and commit step n, retract_draft those proposals, and run the next tick synchronously: proposing rows run the shared verify round, every other row takes the synchronous batched step; (4) re-prime (maybe_prime_lookahead, :478) only once no row proposes and every prompt-lookup row's pipelines_plain_rounds() is true. Where a prime has already been submitted when a proposal appears (admission, completion), tear it down with discard_lookahead (:456) / apply_lookahead_trim (:400) and the existing lookahead_teardown_positions rules. lookahead_params drops its blanket speculative check only for PromptLookup; MTP and DFlash keep it.
• Verify gate (verify rounds must not lower batched throughput): each proposing row runs its own verify forward, so with many concurrent rows a tick would cost one forward per proposing row plus the batched decode. A row runs a verify round only when (a) its drafter's governor grants a non-zero budget (DraftGovernor::budget, prompt_lookup.rs:354, already 0 when the Graded or Gated policy sees recent acceptance below its cutoff), and (b) the tick's active decode batch is at most a new server option --prompt-lookup-max-batch N, whose default is the largest batch at which measurement (c) below shows no aggregate loss. A row declined by (b) does not call draft_block that tick (or retracts the call), so the governor does not record a miss, and its drafter keeps observing committed tokens so it resumes when the batch shrinks.
• Paged storage: a multi-token verify forward cannot take the paged single-token kernel and goes through the pool intercept and gather (docs/adr/0001-paged-attention-gather-vs-fused-kernel.md measured batch-4 gather at 2 to 3 times contiguous SDPA cost past 4096 tokens). Measure verify cost on paged storage at 8192 tokens (measurement (d)). If it erases the gain, skip verify rounds for paged sequences past a measured context length (a named constant next to the gate, with the measurement cited in its doc comment); if it does not, document in the module docs that no limit applies and why.
• Rejected: a run-to-completion burst like the DFlash arm. It holds the worker for the whole request and stalls every concurrent row (the head-of-line block perf(server): tick-cooperative speculative decoding to remove the burst head-of-line block #734 removed for MTP). The engine entries work on the scheduler's own pooled sequence, so the round's state (the drafter and its counters) lives on SequenceInfo across ticks.
• Rejected: batching several prompt-lookup verify rounds into one forward (out of scope here; the verify gate limits the per-row cost instead).

Scope

In scope: src/server/speculative_dispatch.rs (+ tests), the server CLI args for --prompt-lookup-max-batch, src/server/batch/scheduler/ (decode tick, lookahead gate, SequenceInfo), the shared verify round and eligibility predicate in src/lib/mlxcel-core/src/engine/speculative.rs, scripts/bench_serving_concurrency.py (a --prompt-style plain|copy option: copy wraps the synthetic paragraph from build_prompt as "Repeat the following paragraph exactly: ..."), results under docs/benchmark_results/.

Out of scope: llama-server's --spec-type ngram-* spellings (still rejected by src/cli/spec_compat_args.rs); batched verify forwards; vision and structured-output requests (they decode plainly).

Implementation Notes

• Reuse: the verify round and acceptance accounting of speculate, moved rather than copied, so mlxcel generate --prompt-lookup and the server run the same code; pipelined_plain_round's switch rule; the scheduler's step_rows::step_row, apply_row_outcomes, apply_lookahead_trim and discard_lookahead; PromptLookupStats (prompt_lookup.rs:575) for the reported counts.
• Edge cases: a proposal that would pass max_tokens is truncated (as speculate does); a token inside an accepted block that hits EOS, a stop string or the context bound ends the row there and unwinds the positions after it; a request cancelled mid-round is finalized after the round; a --max-kv-size trim runs before the verify forward as before a step; a request admitted while a lookahead is in flight follows the existing admission teardown before its first lookup.
• Error handling: a verify forward error fails only that row (FinishReason::Error, fix(server): catch MLX C++ aborts at the decode FFI boundary so a backend throw fails the request, not the process #822) and unwinds its appends; a failed lookahead teardown follows fail_desynchronized_sequence as today.
• Greedy exactness on real checkpoints: a multi-token verify forward can round differently from the one-token step at a near-tie (the prompt-lookup module's "Greedy exactness" section); the real-checkpoint equality check runs under MLXCEL_SDPA_DETERMINISTIC=1 on a prompt whose reply copies the prompt text, and a mismatch must be shown to be such a near-tie (logit gap printed) or it is a defect.

Acceptance Criteria

• SpeculativeDispatch::resolve tests (src/server/speculative_dispatch_tests.rs) for --draft-kind prompt-lookup alone (resolves to PromptLookup, should_dispatch_speculative() false), with a draft model (error), and with --draft-block-size 4 (max_draft == 4).
• Scheduler test on a stub model whose reply repeats the prompt: a prompt-lookup sequence emits the same tokens as plain decode, accepted proposals are counted, and a concurrent plain request keeps getting a token every tick while the prompt-lookup request decodes.
• Scheduler test: with the flag on and no row proposing, lookahead_params returns Some and the tick is pipelined (fails on current main, where speculative dispatch disables it); with a proposing row the next tick is synchronous and no prime was submitted for it.
• Scheduler test: with an active batch above --prompt-lookup-max-batch, no verify round runs and the drafter's miss count is unchanged.
• Real checkpoint: mlxcel-server -m models/mlx/qwen3-1.7b-4bit --draft-kind prompt-lookup returns the same greedy text as the same server without the flag for a copy-heavy prompt, under MLXCEL_SDPA_DETERMINISTIC=1, with draft_n_accepted > 0 and draft_kind == "prompt-lookup" in the /completion response's speculative timings (SpeculativeTimings, src/server/types/native_completion.rs:66).
• mlxcel generate --prompt-lookup produces the same output as before the refactor (speculative_plain_tests.rs and the prompt-lookup tests pass unchanged).
• Speed, measured once at the end of the implementation run on a quiet host by the orchestrator (together with perf: pipeline penalty and DRY requests in the engine's B=1 client #2229's measurement), 3 to 5 interleaved rounds, Qwen3-1.7B 4-bit: (a) plain chat prompt with little copying, B=1, flag on vs off: decode no worse than -1.0 percent; (b) copy-heavy prompt, B=1: decode speedup and accepted tokens per verify reported; (c) bench_serving_concurrency.py at concurrency 4 and 8, flag on vs off, --prompt-style plain and --prompt-style copy: aggregate tok/s with the flag no worse than without beyond the round spread, and the chosen --prompt-lookup-max-batch default justified by these numbers; (d) 8192-token prompt on paged storage (--decode-storage-backend paged): verify cost reported and whether the paged limit engages. Results committed under docs/benchmark_results/.
• Integrated: the flag works through mlxcel-server and mlxcel serve end to end, and --help lists prompt-lookup and --prompt-lookup-max-batch.

Verification

cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features cuda -- -D warnings
gpu-lock run --tag test -- cargo test --profile test-fast --features cuda -p mlxcel-core --lib engine:: -- --test-threads=1
gpu-lock run --tag test -- cargo test --profile test-fast --features cuda --lib -- --test-threads=1 server::speculative_dispatch server::batch
gpu-lock run --tag verify -- make verify-test-cuda

cargo build --release --features cuda --bin mlxcel --bin mlxcel-server
export MLXCEL_SDPA_DETERMINISTIC=1
./target/release/mlxcel-server -m models/mlx/qwen3-1.7b-4bit --draft-kind prompt-lookup --port 18080 &
./target/release/mlxcel-server -m models/mlx/qwen3-1.7b-4bit --port 18081 &
# POST the copy-heavy prompt with temperature 0 to /completion on both ports, diff the content, check timings.draft_n_accepted > 0

# Measurement, once, at the end, on a quiet host, by the orchestrator
gpu-lock run --tag serve-bench -- python3 scripts/bench_serving_concurrency.py --port 18080 --concurrency 4,8 --max-tokens 256 --prompt-style copy
gpu-lock run --tag serve-bench -- python3 scripts/bench_serving_concurrency.py --port 18081 --concurrency 4,8 --max-tokens 256 --prompt-style copy
# repeat with --prompt-style plain, B=1 (--concurrency 1), and --prompt-tokens 8192 against servers started with --decode-storage-backend paged; interleave on/off rounds

The diff must print nothing, and each speed criterion must hold across the interleaved rounds.