fix(server): speculative decoding mis-attributes decode time to prompt_ms in native timings #1592

Problem / Background

On the speculative burst path (--model-draft, MTP or DFlash), the native /completion response's timings block reports almost the entire request wall time under prompt_ms and 0 to 1 ms under predicted_ms. predicted_per_second is therefore meaningless for any drafted request (95000 tok/s on the measurement below), and the per-token predicted_per_token_ms is equally wrong. The llama-server timings contract that docs/llama-server-compat.md pins to b10621 has prompt_ms cover prompt evaluation only and predicted_ms cover token generation measured from the first token (docs/llama-server-compat.md:111, which is also why predicted_per_second divides by predicted_n - 1).

This predates #1588. The same request measured before that PR showed prompt_ms = 4075.0 and predicted_ms = 0.0; #1588 only added the draft_* acceptance counters (documented in docs/speculative-acceptance.md, "Per-request acceptance on the response itself"), which is what made the 48 verify rounds visible on the response and exposed where the time is going.

Evidence

Measured 2026-09-02 on an Apple M1 Ultra, release build at main 2fddec101, mlxcel-server -m models/mlx/qwen3.5-4b-4bit, POST /completion with {"prompt":"Write a short paragraph about the history of the bicycle.","n_predict":96,"temperature":0}.

Interpretation: on the drafted request, drafter warmup plus every verify round lands in prompt_ms, and the decode phase reports the few microseconds it takes to replay the finished token vector.

Current Behavior

• src/server/model_worker.rs:3549-3552: StreamingDecodeState::on_token stamps first_token_time when the first decoded token passes through it. finish_with_cache (src/server/model_worker.rs:3633-3654) derives prompt_eval_ms = first_token_time - start, and build_generation_result_with_cache (src/server/model_worker.rs:3386-3404) sets generation_only_ms = elapsed_ms - prompt_eval_ms. SequenceInfo::take_generation_result (src/server/batch/sequence.rs:595-616) passes self.created_at as start.
• On the classic path this is right: src/server/batch/scheduler/prefill.rs:1456 feeds the prefill's first token into on_token the moment prefill returns, so prompt_eval_ms is queue wait plus prefill and everything after it is decode.
• On the B=1 burst path, try_run_burst_b1 (src/server/batch/speculative_burst.rs:830) runs the whole request to completion before any token reaches the sequence. For DFlash, run_dflash_burst (src/server/batch/speculative_burst.rs:1552) calls WorkerDrafterSlot::ensure_loaded (:1608, the lazy disk load, whose duration goes only into a tracing::info! at :295-301), takes prefill_start = Instant::now() (:1645), runs run_dflash_on_qwen35 (:1746, target prefill plus every verify round), and computes prefill_time_ms = total_burst - decode_time_ms (:1696-1697). For MTP, run_mtp_burst (:1094) does the same with ensure_loaded at :1183 and takes prefill_time_ms / decode_time_ms from MtpGenerator's own timers (:1437-1438, produced at src/lib/mlxcel-core/src/speculative/mtp/generator.rs:1449-1455).
• Only after the burst returns does finalize_burst_success (src/server/batch/speculative_burst.rs:2026-2038) call stream_burst_tokens (:2097), which calls seq.mark_first_token() (:2112) and replays the finished token vector through seq.decode_state.on_token (:2149) in a tight loop. Both first_token_time stamps therefore land after drafter load, prefill, and all verify rounds, so prompt_eval_ms absorbs everything and generation_only_ms is the replay loop.
• finalize_burst_success receives the burst's measured split as _prefill_time_ms and _decode_time_ms (:2031-2032) and discards both. The batched window (try_run_burst_batched, :2391) passes literal 0.0, 0.0 (:2500-2506) because the batched round loops report prefill_time_ms: 0.0 (src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs:575, src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs:507).
• Route side, every consumer reads the same wrong figures: the non-streaming route uses result.prompt_eval_ms / result.generation_only_ms (src/server/routes/native_completion.rs:323-324), the streaming final frame does the same (:1034-1035), and the streaming per-token predicted_ms is StreamPrefill::predicted_ms() (:806-810), measured from when the route observed the GenerateEvent::Prefill snapshot that mark_first_token emits (src/server/batch/sequence.rs:411-425), which on the burst path is also after the whole burst. Timings::new (src/server/types/native_completion.rs:134-162) then divides 95 intervals by 1 ms.
• The tick-cooperative MTP slice path (begin_slice_session, src/server/batch/speculative_slice.rs:459-495) streams slice 0 right after prefill plus seed, so its predicted_ms covers the later rounds; its prompt_ms still includes the drafter load that BatchScheduler performs at slice start (src/server/batch/scheduler.rs:140-166).

Proposed Solution

1. Carry a real prefill-end instant out of the burst. Replace BurstSuccess::{prefill_time_ms, decode_time_ms} (src/server/batch/speculative_burst.rs:980-981) with prefill_end: Instant, the instant the target prefill finished and verify round 0 began; finalize_burst_success is the only consumer of the two floats and it discards them today.
   • DFlash B=1: prefill_end = prefill_start + Duration::from_secs_f64(prefill_time_ms / 1000.0) from the existing prefill_start (:1645) and the existing prefill_time_ms = total_burst - decode_time_ms (:1697); both already exclude ensure_loaded (:1608), and DFlashTargetRun.decode_time_ms is the round loop's decode_elapsed (src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs:889), which spans every verify round. If an explicit stamp is preferred, capture Instant::now() inside run_dflash_on_qwen35 right after the target prefill forward returns; the two differ only by the gap between prefill return and round 0.
   • MTP B=1: take let gen_start = Instant::now() after the bind succeeds and before the generator runs, and derive prefill_end = gen_start + Duration::from_secs_f64(stats.prefill_time_ms / 1000.0) (:1437), since MtpGenerator measures prefill and decode separately.
   • Batched window: stamp Instant::now() in run_mtp_burst_batched (:2573) / run_dflash_burst_batched (:2727) right after the batched target prefill returns and before round 0, return it through the per-window result, and pass that one instant for every row (all rows share the prefill), so try_run_burst_batched stops passing 0.0, 0.0 at :2500-2506.
2. Let the sequence accept a supplied first-token instant. Add SequenceInfo::mark_first_token_at(&mut self, at: Instant) in src/server/batch/sequence.rs and make mark_first_token() delegate to it with Instant::now(), keeping the idempotent guard. Add StreamingDecodeState::stamp_first_token(&mut self, at: Instant) in src/server/model_worker.rs that sets first_token_time only when it is None; on_token keeps its own stamp for the classic path. mark_first_token_at calls self.decode_state.stamp_first_token(at) so the streamed PrefillStats.prompt_ms and the final prompt_eval_ms derive from one instant, which is the fix(server): separate native completion and embedding routes from OpenAI routes #1441 invariant that the two frames cannot disagree.
3. finalize_burst_success calls seq.mark_first_token_at(prefill_end) before stream_burst_tokens; the guard makes the replay's mark_first_token() a no-op. predicted_ms then equals done - prefill_end, which is every verify round plus the sub-millisecond replay. The slice path already stamps at the end of slice 0 and needs only item 4.
4. Exclude drafter cold-load time from prompt_ms. Change WorkerDrafterSlot::ensure_loaded (:283) to return Result<Option<Duration>, String>: Some(elapsed) when this call loaded from disk, None when the drafter was already resident. Update its four call sites (:1183, :1608, :2607, src/server/batch/scheduler.rs:140). Add SequenceInfo::timing_origin_shift: Duration (zero in every constructor: src/server/batch/scheduler/admission.rs:527 area and the two in src/server/batch/scheduler/handoff.rs:388 / :507), set by the driver that paid the load; in the batched window set it on every row, since every row waited. mark_first_token_at, report_prefill_progress, and take_generation_result use self.created_at + self.timing_origin_shift as the origin, which removes the load from prompt_ms, prompt_eval_ms, and generation_time_ms while leaving generation_only_ms unchanged. Rejected alternative: a new timings.draft_load_ms key. It is a once-per-process cost that says nothing about the request, and it would need a b10621 manifest entry (compat/llama-server/b10621/routes.json) plus docs; the existing Drafter loaded (kind=.., N ms) log line at :298-301 stays the record of it.

Scope

In scope: src/server/batch/speculative_burst.rs (BurstSuccess, the three drivers, finalize_burst_success, ensure_loaded), src/server/batch/speculative_burst_tests.rs, src/server/batch/sequence.rs, src/server/model_worker.rs (StreamingDecodeState), src/server/batch/scheduler.rs (slice-start load), src/server/batch/scheduler/admission.rs and handoff.rs (field init), docs/speculative-acceptance.md and docs/llama-server-compat.md cross-links.

Out of scope: the draft_n / draft_n_accepted / draft_rounds / draft_kind fields themselves (shipped in #1588); any new timings key; the classic path's prompt_ms origin at created_at (queue wait included, as #1441 decided); per-request TTFT observability for the burst path (record_request_completion is only called from src/server/batch/scheduler/decode_tick.rs:1651, the classic path, so the burst never feeds that histogram; separate issue if wanted).

Implementation Notes

• Reuse: the burst drivers already measure the split; do not add a second timer around the round loop. mark_first_token's idempotent guard is what keeps the replay from re-stamping; do not add a flag for it.
• Constraints: the classic path's numbers must not move: prefill.rs:1456 still stamps through on_token and timing_origin_shift stays zero there. The streamed Prefill snapshot and the final frame must still agree (fix(server): separate native completion and embedding routes from OpenAI routes #1441). Clamp prompt_ms at zero if a driver ever hands in an instant earlier than the shifted origin rather than panicking on Instant subtraction.
• Edge cases: burst declined to classic (BurstOutcome::DeclineToClassic) after a cold load: the classic path serves the request, but the load happened on it, so leave the shift applied and prompt_ms still excludes the load. Burst error after load: the error event carries no timings, nothing to do. Zero verify rounds (EOS at the first bonus): prefill_end is the burst end, predicted_ms is legitimately about 0 with predicted_n <= 1, and Timings::new's zero-divisor guard yields 0 rather than a huge rate. Warm drafter (second request onward): ensure_loaded returns None, the shift stays zero, and prompt_ms is prefill plus queue wait exactly as on the classic path.
• Error handling: unchanged; ensure_loaded errors still map to BurstOutcome::Error and the client sees the same error event.

Acceptance Criteria

• On the DFlash request in the Evidence table (n_predict = 96, temperature = 0), the non-streaming /completion response reports predicted_ms within a few percent of the wall time of the 48 verify rounds, and prompt_ms no larger than the classic path's prompt_ms for the same prompt plus scheduling jitter (tens of ms, not thousands); predicted_per_second is in the same order of magnitude as the classic path's value, not 95000.
• The first drafted request after server start and the second one report prompt_ms of the same order; the drafter cold-load time no longer appears in any timings field and still appears in the Drafter loaded log line.
• With stream: true, the first timings frame's prompt_ms equals the final frame's prompt_ms, and the final frame's predicted_ms matches the non-streaming value for the same request within jitter.
• MTP B=1 (a Qwen 3.5 MTP or Gemma 4 Inkling checkpoint), the batched window (two concurrent drafted requests), and the classic path (no --model-draft) all report prompt_ms bounded by their prefill and predicted_ms > 0 whenever predicted_n > 1; the classic path's numbers are unchanged from before the fix.
• Unit test in src/server/batch/speculative_burst_tests.rs using make_test_sequence() (:101) and a prefill_end a known duration before now, driving finalize_burst_success (or mark_first_token_at plus take_generation_result if the BurstContext is too heavy to build), asserting on the GenerateEvent::Done result that prompt_eval_ms <= prefill_duration + tolerance and generation_only_ms >= decode_duration - tolerance, and on the GenerateEvent::Prefill snapshot that its prompt_ms equals prompt_eval_ms.
• Unit test with the existing synthetic drafter pair (MockMtpTarget at :544, TrackingMockDrafter at :685) driving drive_mtp_generator (src/server/batch/speculative_burst.rs:1493) where prefill_and_seed sleeps a known duration and each verify round sleeps another, asserting the derived prefill_end falls inside the prefill window and that the resulting predicted_ms > 0.
• Unit test that a timing_origin_shift of D ms reduces prompt_eval_ms and generation_time_ms by D and leaves generation_only_ms unchanged, and that ensure_loaded on an already-resident slot returns Ok(None).
• docs/speculative-acceptance.md "Per-request acceptance on the response itself" (:179) gains one paragraph stating what prompt_ms and predicted_ms cover on a drafted request and linking the timings note in docs/llama-server-compat.md; that note links back.
• The PR body records before/after numbers for the two Evidence rows, re-measured on the same machine, checkpoint pair, and build flags.
• Integrated into the real code flow: the fix lands in the burst drivers and the sequence bookkeeping, not as a route-side post-hoc correction.

Verification

cargo build --release --features metal,accelerate
cargo test --workspace --profile test-fast --features metal,accelerate
cargo test --profile test-fast --features metal,accelerate speculative_burst_tests::
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --all -- --check

# Manual: classic vs drafted, same prompt, non-streaming and streaming.
./target/release/mlxcel-server -m models/mlx/qwen3.5-4b-4bit &
curl -s localhost:8080/completion -d '{"prompt":"Write a short paragraph about the history of the bicycle.","n_predict":96,"temperature":0}' | jq .timings
./target/release/mlxcel-server -m models/mlx/qwen3.5-4b-4bit --model-draft models/mlx/qwen3.5-4b-dflash &
curl -s localhost:8080/completion -d '{"prompt":"Write a short paragraph about the history of the bicycle.","n_predict":96,"temperature":0}' | jq .timings
curl -s localhost:8080/completion -d '{"prompt":"Write a short paragraph about the history of the bicycle.","n_predict":96,"temperature":0}' | jq .timings
curl -sN localhost:8080/completion -d '{"prompt":"Write a short paragraph about the history of the bicycle.","n_predict":96,"temperature":0,"stream":true}' | grep -o '"timings":{[^}]*}' | sed -n '1p;$p'

Pass looks like: on the drafted server, prompt_ms on the order of the classic 73 ms (first request included, once the drafter is resident on the second call it should be indistinguishable), predicted_ms on the order of the drafted request's wall time minus prefill, predicted_per_second in the low hundreds, and draft_rounds unchanged at 48.

Technical Considerations

Related: #1588 (acceptance counters and docs/speculative-acceptance.md), #1441 (the Prefill snapshot and the streamed-equals-final invariant), #1314 (per-request acceptance reporting), #734 (tick-cooperative MTP slice path, which already stamps at the end of slice 0).