perf(server): port lookahead async_eval pipelining into BatchScheduler decode #632

Part of #623. Can proceed in parallel with Phase 3; rebase measurements on the post-#625 baseline.

Context

The CLI generation loop overlaps graph construction with GPU execution: it builds step n+1's graph (reshape -> forward -> fused sample) before reading back step n's token, schedules it with ffi::async_eval_pair, and only then syncs the current token with item_i32 (src/lib/mlxcel-core/src/generate.rs:1157-1260; first token at :1246; MLXCEL_FORCE_SYNC opt-out at :1227). A dedicated thread-local generation stream supports the overlap (src/lib/mlxcel-core/src/streams.rs, installed at generate.rs:1059).

The server scheduler has none of this. Both decode paths synchronize every step:

• decode_single_step (src/server/batch/scheduler.rs:4699): forward at :4724, then synchronous eval + item_i32 at :4778-4779.
• execute_batched_decode (scheduler.rs:4300): forward at :4363, fused batch sample at :4384-4388, then a host copy of the [B] token ids.

On CUDA, where kernel-launch and graph-build latency is higher than Metal, this makes server decode structurally slower than CLI decode for identical models. Overhead-bound models (<= 2B: 67-75% of roofline) lose the most; the effect also scales the whole batched path since one tick's build time is serialized against the GPU.

Scope

Bring lookahead pipelining to the scheduler's decode, for both B=1 and batched decode, without breaking continuous-batching semantics (admission, completion, chunked-prefill interleave, cancellation).

Implementation plan

1. Mechanics (mirror the CLI loop): after sampling step n's token array, immediately build step n+1's graph feeding the token array back as the next [B,1] input WITHOUT reading it to host, schedule with async_eval (async_eval_pair when logprobs are also needed), and only then extract step n's ids to host for stop checks and streaming. The sampled-token MLX array is a valid device-side input; the CLI loop already relies on this.
2. Batch-topology invalidation: lookahead is only valid while batch membership, masks, and per-sequence sampler state are unchanged. Invalidate (discard the prebuilt step and fall back to a synchronous step) when the next tick's decide_action is not Decode with the identical id set: admission of a queued request (scheduler.rs:2616-2621), any sequence finishing (EOS/stop/length), chunked-prefill interleave (:2593-2601), preemption, or --max-kv-size trimming that mutates cache shape. Simplest correct policy: compute the lookahead only when the queue is empty and no sequence is within 1 token of a stop condition that is knowable host-side; otherwise run synchronously. Measure how often lookahead engages (add a counter to BatchObservability).
3. Stop-condition subtlety: stop-string and EOS checks need host tokens, and a lookahead step may compute one token past a stop. That is the same speculative-waste tradeoff the CLI loop accepts (it trims after the fact); the scheduler must additionally make sure the extra step does not corrupt per-sequence KV: one extra decode step appends one KV position, so on stop, trim the cache by the overshoot (KVCache::trim, src/lib/mlxcel-core/src/cache.rs:2112) before the slot is reused or donated to the prompt cache.
4. Sampler-state ordering: penalties/DRY state updates (SamplerState, src/lib/mlxcel-core/src/sampling.rs:273-287) must see token n before sampling token n+1. The CLI loop already threads this; reuse the same structure rather than reimplementing.
5. Keep MLXCEL_FORCE_SYNC=1 as the global kill switch (already respected in the CLI loop; honor it here too).
6. Cancellation: the lazy disconnect check in finalize_completed (scheduler.rs:4886-4924) is unaffected, but verify a cancelled sequence with a prebuilt lookahead step aborts cleanly.

Acceptance criteria

• B=1 server decode within 5% of CLI decode tok/s on GB10 for llama-3.1-8b-4bit AND qwen2.5-0.5b-bf16 (the overhead-bound case is the one that proves the win; measure via perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624 telemetry and mlxcel-bench-decode as the CLI reference).
• Batched decode (B=4) aggregate tok/s improves vs the synchronous baseline; number recorded.
• Greedy token streams byte-identical to the synchronous path (run with and without MLXCEL_FORCE_SYNC=1) including stop-string cases that trigger mid-stream and multi-sequence runs with staggered completions.
• Lookahead-engagement counter exposed in /metrics; no KV corruption: prompt-cache donation after an overshoot-trim verified by a cache-hit test.
• cargo test --features cuda passes; scheduler unit tests extended for the invalidation cases (admission mid-lookahead, stop mid-lookahead, preemption).

Validation

cargo build --release --features cuda
./target/release/mlxcel-bench-decode --model ./models/qwen2.5-0.5b-bf16 --max-tokens 512          # CLI reference
./target/release/mlxcel-server --model ./models/qwen2.5-0.5b-bf16 &                                # server side
python3 scripts/bench_serving_concurrency.py --concurrency 1
MLXCEL_FORCE_SYNC=1 python3 scripts/bench_serving_concurrency.py --concurrency 1                   # equivalence check

References

• CLI pipelined loop: src/lib/mlxcel-core/src/generate.rs:1131,1157-1260.
• Bridge: async_eval/async_eval_all/async_eval_pair (src/lib/mlxcel-core/src/lib.rs:1719-1793).
• Scheduler decode: src/server/batch/scheduler.rs:4259-4400,4699-4790; tick policy :2585-2622.