refactor: batch-native engine step API with BatchScheduler on top #2172

Part of #2166. Phase 4b. Depends on #2171 (Phase 4a, KV storage), #2168 (Phase 1, finish step) and #2169 (Phase 2, sampling step), and through 4a on #2170 (Phase 3, prefill plan). Phase 5 builds on this.

Problem / Background

BatchScheduler mixes two jobs: deciding what runs (admission, tick policy, preemption, prompt cache, handoff), and running it (model forwards, sampling, finishing). Because the execution half lives in the scheduler, the CLI cannot reuse it without inheriting the scheduler. The single-sequence LanguageModel::forward is also a separate path from forward_batched_with_context, so "B=1" on the server and "single sequence" on the CLI are different code. llama.cpp has one batch-native entry, llama_decode(llama_batch), where a single sequence is B=1. This phase gives mlxcel the equivalent.

Current Behavior

The scheduler calls model forwards directly:

• src/server/batch/scheduler/decode_tick.rs:682, :686, :874 and :1448.
• src/server/batch/scheduler/prefill.rs:649, :837, :856, :1008, :1027 and :1248.
• src/server/batch/scheduler/prompt_cache.rs:982 and :1160.

It calls samplers directly at six sites in decode_tick.rs and three in prefill.rs. The single-sequence and batched forwards are separate LanguageModel methods (src/lib/mlxcel-core/src/generate.rs:481 and :1011-1051).

The decode lookahead pipeline also reaches model state directly. lookahead_params (decode_tick.rs:372) admits a model-owned family only when LanguageModel::supports_decode_lookahead_rewind() is true (Gemma 3 since #2182). The prime forward in lookahead_forward (:669) runs inside a DecodeLookaheadAppendScope (:677) so a rewinding model logs the rows its speculative writes overwrite. Teardown in apply_lookahead_trim (:464-512) trims pool caches and, for a model-owned family, calls LanguageModel::rewind_decode_appends(seq_id, n). A failed rewind finishes that row with an error through fail_desynchronized_sequence (:518-526), and the callers drop it from the synchronous re-dispatch (without_failed, :529).

Proposed Solution

Add an engine module in the core, src/lib/mlxcel-core/src/engine/. Use the names ADR 0007 (#2167) fixes. If the ADR leaves names open, use these:

pub struct Engine<M: LanguageModel> { /* model, KV pool, per-sequence RowSampler and finish state */ }
impl<M: LanguageModel> Engine<M> {
    pub fn open(&mut self, req: SequenceParams) -> Result<SeqHandle, EngineError>;
    pub fn prefill(&mut self, seq: SeqHandle, piece: &PrefillPiece) -> Result<PrefillOutcome, EngineError>;
    pub fn step(&mut self, batch: &[SeqHandle]) -> Result<Vec<StepOutcome>, EngineError>;
    pub fn close(&mut self, seq: SeqHandle);
}
pub struct StepOutcome { pub seq: SeqHandle, pub token: i32, pub finish: Option<FinishCause> }

step composes the contracts from the earlier phases, restated here so this issue stands alone:

• Forward: through the Phase 4a KV abstraction (refactor: one KV storage abstraction with dense and paged backends #2171). KVCache::attend(&mut self, q, k, v, scale, mask) and its batched counterpart choose dense SDPA or paged kernels inside the cache. Models do not branch on storage.
• Sampling: through the Phase 2 RowSampler (refactor: one per-row sampling step and sampler-state lifecycle #2169). It owns SamplerState creation (needs_token_history() || needs_sampler_feedback_state()), the fix(generate): penalties in CxxGenerator see a history one token stale #2090 history-ordering invariant, the structured-output mask hook, the thinking-budget override (returning (sampled, final)), and compose_token_bias. Its fused_eligible selects the fused [B, vocab] chain.
• Finishing: through the Phase 1 finish_step(FinishInput, &mut dyn FinishHooks) -> Option<FinishCause> (refactor: one post-sample finish step for every CLI and server decode site #2168). FinishCause is Eos | StopSequence | StructuredStop | Length | ContextExhausted | RepetitionLoop, applied in the order EOS, push, stop string, generation bound, structured stop, max_tokens, context bound, repetition loop, cache-clear cadence.
• Prefill: through the Phase 3 PrefillPlan (refactor: one prefill planner and one pad-trim hook for CLI and server #2170). The scheduler passes one PrefillPiece { range, padded_len } at a time so chunked prefill still interleaves with decode, and trim_state(seq: Option<SequenceId>, excess) runs after padding.

A single sequence is a batch of one. LanguageModel::forward becomes a provided method that delegates to the batched entry with one row, so no family keeps a separate single-row decode path.

BatchScheduler is rebuilt on Engine. These stay in the scheduler: admission, tick policy, preemption, mixed prefill/decode steps (ADR 0005), prompt-cache lookup, donate and adopt, disaggregated handoff, and the batched speculative generators. These come from the engine: model execution, sampling and finishing. The legacy --no-batch worker (src/server/model_worker.rs:1103-1115) is the same scheduler at max_batch_size=1, so it is covered automatically.

Scope

In scope: the engine module, delegating the single-row forward, rebuilding the scheduler's execution half, the legacy worker, and tests.

Out of scope: the CLI (Phase 5), retiring CxxGenerator (Phase 6), and changing scheduler policy.

Implementation Notes

• Reuse: the existing CachePool, the speculative round loops and the prompt-cache code. Wrap them, do not reimplement them. The engine exposes the hooks the prompt cache needs (snapshot_sequence_state and restore_sequence_state, generate.rs:788-843).
• Constraints: the scheduler's lookahead pipeline (decode_tick.rs:563-798, maybe_prime_lookahead through pipelined_steady_decode) must keep its overlap, so step needs a split submit and collect form, or a lookahead variant. That form must also carry the speculative-append contract above: the prime forward runs inside DecodeLookaheadAppendScope, and a teardown unwinds up to DECODE_LOOKAHEAD_MAX_SPECULATIVE_APPENDS (2) appends per row, through the pool trim or rewind_decode_appends for a model-owned family. No per-step allocation proportional to batch size beyond what exists today.
• Error handling: EngineError per sequence becomes FinishReason::Error(msg) for that row only, matching today's fix(server): catch MLX C++ aborts at the decode FFI boundary so a backend throw fails the request, not the process #822 behavior. One failed row never aborts the batch. A failed model-owned rewind is the same kind of row-only error, and it overrides a Length or Stop recorded on the same tick so the desynchronized state is not donated (perf(scheduler): restore Gemma 3 decode lookahead via ring rewind #2182).

Acceptance Criteria

• rg -n '\.forward[a-z_]*\(|sample_token_[a-z_]*\(' src/server/batch/scheduler --glob '!*_tests.rs' returns nothing, so the scheduler no longer calls model forwards or samplers directly (PR refactor: batch-native engine step API with BatchScheduler on top #2217; the grep's only hits are plan.forwarded_len(), a PrefillPlan accessor)
• LanguageModel::forward delegates to the batched entry. A test asserts that single-row forward and one-row forward_batched_with_context give bitwise-identical logits on a real checkpoint (PR refactor: batch-native engine step API with BatchScheduler on top #2217: implemented at the engine instead of inverting the trait, see ADR 0009 and the comment below; src/models/single_row_batch_parity_tests.rs pins the bitwise equality for Qwen3 and Llama 3)
• The server tests and the continuous-batching benches are green (PR refactor: batch-native engine step API with BatchScheduler on top #2217: scheduler test modules green on GB10; benches not run, per the epic's end-of-run measurement rule)
• Under MLXCEL_SDPA_DETERMINISTIC=1, the Phase 0 harness shows server B=1 dense identical to a direct Engine B=1 run (PR refactor: batch-native engine step API with BatchScheduler on top #2217: d:engine arm, identical on Qwen3-1.7B and Llama-3.2-1B for both cases)
• Batched and single-stream serving throughput stay within the Phase 0 threshold
• The engine is the execution path the server actually runs, with no parallel unused module (PR refactor: batch-native engine step API with BatchScheduler on top #2217)
• Phase 4a deferral (comment below): Gemma 3, Llama 4 and model_owned move onto attend / attend_batched so rg -n 'is_paged_backed\(\)|is_paged_decode\(\)' src/models returns nothing, with the Gemma 3 decode-undo log kept on RotatingKVCache::update (PR refactor: batch-native engine step API with BatchScheduler on top #2217: KvAttention, RotatingKVCache::attend and ChunkedKVCache::attend write through update_and_fetch, cache::attend_batched_rows replaces the model_owned helpers; recorded in ADR 0009)
• Phase 4a deferral (comment below): qwen3_5, qwen3_moe and muse_glimmer drop the ignored DecodeBatchContext argument (PR refactor: batch-native engine step API with BatchScheduler on top #2217: DecodeBatchContext and the core DecodeStorageBackend are removed, the trait entry is forward_batched_with_ids)

Verification

make verify                                   # macOS gate
make verify-fmt verify-test-cuda              # Linux/CUDA gate
cargo test --profile test-fast --features cuda -p mlxcel-core engine
cargo test --profile test-fast --features cuda --lib server::batch
MLXCEL_SDPA_DETERMINISTIC=1 make <Phase 0 parity target> MODEL=<Qwen3-1.7B-4bit path>

References verified against main at fc21784 (2026-10-07).