epic: one batch-native engine for CLI and server decode #2166

Problem / Background

The CLI and the server share no decode loop. mlxcel generate, mlxcel run and the chat REPL decode through CxxGenerator (src/lib/mlxcel-core/src/generate.rs:1914, :2233, :2416, :2683, reached through MlxInferenceSession in src/lib/mlxcel-core/src/session.rs:47-60). mlxcel-server, mlxcel serve and --no-batch decode through BatchScheduler (src/server/model_worker.rs:1103-1115: the legacy worker is the scheduler at max_batch_size=1). ADR 0004 and src/backend/mod.rs:37-43 record this split as deferred work.

What is shared today: the single-row model forward, the per-token sampler (sample_token_optimized[_with_state]), build_sampling_config (src/execution/sampling.rs:87), EOS merging, the streaming detokenizer, and the MTP/DFlash round loops.

What is implemented twice:

• The post-sample finish step (EOS, history push, stop strings, generation and context bounds, max_tokens, loop detection, cache-clear cadence): five server copies and four CLI loops.
• The per-row sampling step, the SamplerState lifecycle, and the token-history ordering invariant.
• Prefill chunking, padding and trim: chunk 2048 on the CLI (generate.rs:282) vs 512 on the server (src/server/config.rs:1172), and two trim hooks (trim_internal_caches at generate.rs:652, trim_sequence_state at generate.rs:675).
• Sampling-default resolution, token-bias composition, KV-mode application, reasoning splitting, and DFlash target setup.
• Attention kernels: the CLI's dense cache runs MLX fused SDPA, while pool-backed server decode runs mlxcel's paged kernels (src/models/qwen3.rs:354 branches on cache.is_paged_backed()).

Consequences seen so far:

• fix(generate): penalties in CxxGenerator see a history one token stale #2090 hit only the CLI, and fix(generate): sample each penalized step against the token just read #2117 had to fix it in four loops. fix(speculative): Qwen 3.5 MTP output is not byte-identical to classic decode, and the exactness probe does not catch it #1201 needed a server-side twin in fix(server): gate tile-aligned prefill on the model, not on the hardware alone #1205. feat(scheduler): sequence-aware pad trim for model-owned families #2140 had to change both trim hooks.
• DRY defaults differed (docs: explain CLI DRY runs without sequence breakers #1118, fix(server): wire --dry-sequence-breaker through to the sampler #1128), so tuning on mlxcel run did not carry to the server. The repeat_last_n default still differs (feat(inference): complete b10621 sampling and grammar semantics #1436: full history on the CLI, 64 on the server). fix(video): align the CLI and server video-frames fallback #1777 found the CLI joining prompt parts with a different separator.
• Stop-string support (fix(server): honor string stop sequences on the MLX serving path #1466) touched five server emission sites.
• fix(cuda): Qwen3 greedy decoding is not reproducible run to run on GB10 #2128 and fix(cuda): add MLXCEL_SDPA_DETERMINISTIC for reproducible decode #2162 affected only the CLI, because only the CLI runs cuDNN SDPA.
• A server prompt-cache hit gives different greedy output than an uncached run of the same request.
• docs/CONTINUOUS_BATCHING.md:65-79 disclaims token exactness between the server and mlxcel generate.
• src/bin/bench_decode.rs:381 measures only CxxGenerator, so the decode path the server runs has no decode benchmark.

Goal

Follow the llama.cpp design (maintainer decision, 2026-10-07). Build one batch-native engine in the core: a single sequence is a batch of one, one KV abstraction has dense and paged storage behind it, and one implementation does sampling and finishing. The server runs on that engine. mlxcel run and chat become clients of the in-process server engine, and a thin raw-completion generate (the llama-completion analog) also runs on the engine.

Reference design, verified on llama.cpp HEAD 65840ed (2026-10-06):

• The core forward is batch-native: llama_decode(llama_batch) with multi-sequence KV, where a single sequence is B=1.
• Sampling with history lives in one sampler object (common_sampler_* over libllama's llama_sampler chain), and every tool uses it.
• llama-cli is a thin client of llama-server: tools/cli/CMakeLists.txt links llama-server-impl, tools/cli/cli-server.h runs llama_server() in a thread on a random local port, and tools/cli/cli-context.cpp posts to /v1/chat/completions over SSE.
• History: cli: new CLI experience ggml-org/llama.cpp#17824 (2025-12) moved the CLI onto the server context, and the old loop became llama-completion for raw use. cli : move to HTTP-based implementation ggml-org/llama.cpp#24948 (2026-07) made the CLI an HTTP client, with --server-base for a remote server.

Rejected alternative: routing the CLI through today's scheduler as-is. The CLI would only inherit the server path's different numerics (paged kernels, 512-token chunks) and the scheduler's per-step overhead. Unifying the engine first is what makes "same input, same output" hold.

Program invariants:

• Once Phase 5 lands, greedy output for the same request under MLXCEL_SDPA_DETERMINISTIC=1 is identical across the CLI and the server engine (B=1, dense and paged).
• No phase may regress single-stream decode throughput beyond the threshold Phase 0 sets.
• Every phase ends integrated into the real code paths. No phase may leave a standalone module unused.

Out of scope: new model families, kernel tuning, the router, OpenXLA/IREE.

Sub-issues

The ### Phase N headings below are execution waves for /epic-impl, which accepts only positive integers. Each item names its program phase (0, 1, 2, 3, 4a, 4b, 5, 6), plus a final ROCm verification phase. Program Phases 1, 2 and 3 run in parallel after Phase 0. Phase 4a follows Phase 3. Phase 4b needs 4a, 1 and 2. Phase 5 follows 4b, and Phase 6 follows 5. The final ROCm verification (wave 7) follows Phase 6 and runs on the gfx1151 ROCm host, because the epic is implemented on a non-ROCm machine.

Phase 1

• test: ADR 0007 plus CLI/server parity and B=1 decode benchmark harness #2167 Phase 0: ADR 0007, parity harness, B=1 decode benchmark, defaults table

Phase 2

• refactor: one post-sample finish step for every CLI and server decode site #2168 (depends on test: ADR 0007 plus CLI/server parity and B=1 decode benchmark harness #2167) Phase 1: one post-sample finish step
• refactor: one per-row sampling step and sampler-state lifecycle #2169 (depends on test: ADR 0007 plus CLI/server parity and B=1 decode benchmark harness #2167) Phase 2: one per-row sampling step
• refactor: one prefill planner and one pad-trim hook for CLI and server #2170 (depends on test: ADR 0007 plus CLI/server parity and B=1 decode benchmark harness #2167) Phase 3: one prefill plan and one trim hook

Phase 3

• refactor: one KV storage abstraction with dense and paged backends #2171 (depends on refactor: one prefill planner and one pad-trim hook for CLI and server #2170) Phase 4a: one KV storage abstraction

Phase 4

• refactor: batch-native engine step API with BatchScheduler on top #2172 (depends on refactor: one KV storage abstraction with dense and paged backends #2171, refactor: one post-sample finish step for every CLI and server decode site #2168, refactor: one per-row sampling step and sampler-state lifecycle #2169) Phase 4b: engine step API, scheduler on it

Phase 5

• refactor: route mlxcel run and chat through the in-process server engine #2173 (depends on refactor: batch-native engine step API with BatchScheduler on top #2172) Phase 5: run and chat on the in-process server engine

Phase 6

• refactor: run mlxcel generate on the engine and retire CxxGenerator #2176 (depends on refactor: route mlxcel run and chat through the in-process server engine #2173) Phase 6: raw generate on the engine, retire CxxGenerator

Phase 7

• test(rocm): verify the unified engine on ROCm after the batch-native engine epic #2192 (depends on refactor: run mlxcel generate on the engine and retire CxxGenerator #2176) ROCm verification of the unified engine on the gfx1151 host

ROCm constraints

Not verified on the implementing machine. The epic is implemented on a non-ROCm host, so ROCm-only tests (files under #![cfg(feature = "rocm")], and tests that skip when the backend has no kernel port) do not run there. All ROCm verification is deferred to #2192.

Behavior to preserve. Phases 2, 4a, 4b and 0 must not drop these while moving code (#2192 lists the tests that pin each):

• Sampler routing reads the port tables: rejection_sample_supported() returns has_kernel_port(rejection_ports()), with Gumbel and rejection HIP ports (PR update(rocm): port the Gumbel-max and rejection samplers to HIP #2100, src/lib/mlx-cpp/turbo/sampling_rejection.cpp, src/lib/mlxcel-core/src/sampling.rs).
• The lang-bias counter gate's thread-scoped override (PR fix(sampling): scope the lang-bias counter arming to the test thread #2188, src/lib/mlxcel-core/src/lang_bias_counters.rs).
• SliceUpdate does not donate a still-referenced source (PR fix(rocm): stop SliceUpdate donating a still-referenced source #2070, src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md item 24, tests/rocm_slice_update_source.rs).
• Paged decode: PagedDecodeBackend::Rocm and the HIP paged-attention ports (PR update(rocm): port paged attention (v1, v2 partial, merge) to HIP #2103, layers.rs, turbo/paged_attention_hip.h); 64-bit pool offsets and the PagedDecodePlan check that sends oversize batches to gather (PR fix(paged-attention): use 64-bit pool offsets in the fused kernels #2186, paged_v2/plan.rs); the wave64 hold in port_for (PR fix(rocm): refuse wave32-only kernel ports on wave64 devices #2179, turbo/kernel_port.cpp).
• Fused add-RMSNorm and RoPE-append on by default on ROCm via cfg!(feature = "rocm") on FUSED_*_DEFAULT, the Llama 3.1 bypass notice only on an explicit truthy value, and dest_layout for dense and paged RoPE-append outputs (PR update(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm #2189, layers.rs, turbo/fused_rope_append.cpp).
• The bf16 qmm route select_qmm_route and the dense-prefill predicate quantized_matmul_runs_dequant_gemm (PR fix(rocm): route large bf16 qmm to hipBLASLt and gate dense prefill #2085), and the expert-batched gather for bf16/f16 affine and bf16 mxfp4 (PRs fix(rocm): fix and enable the expert-batched gather_qmm prefill kernel #2112, perf(rocm): batch sorted mxfp4 gather_qmm prefill by expert #2164), all in patches-rocm/mlx/backend/rocm/quantized/qmm.hip. Phase 3's chunk change moves which route and gather path each prefill reaches.
• MLX_ROCM_MAX_INFLIGHT_MB and the ROCm MLXCEL_CACHE_LIMIT default (PR fix(rocm): bound in-flight batch memory and enforce the cache limit #2084, patches-rocm/mlx/backend/rocm/device.cpp, src/execution/runtime.rs), plus the fix(rocm): cover the in-flight working set in the pre-load memory estimate #2155 pre-load estimate.
• The SSM update HIP port and MLXCEL_SSM_KERNEL (PR update(rocm): port ssm_update_kernel to HIP and read the port table #2099), and the fused MoE port with ROCm MLXCEL_FUSED_MOE_SGY default 2 (PR update(rocm): port the fused MoE decode kernels to HIP #2098), both in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.
• The Gemma 3 lookahead rewind on ROCm (PR perf(scheduler): restore Gemma 3 decode lookahead via ring rewind #2182, src/lib/mlxcel-core/src/cache/decode_undo.rs).

Sequencing with in-flight ROCm work.

• fix(rocm): cover the in-flight working set in the pre-load memory estimate #2155 (touches src/bin/bench_decode.rs) lands before Phase 0's bench changes, or Phase 0 rebases on it.
• perf(rocm): find why batched server decode is far below single-stream #2156 (touches src/server/batch/scheduler/run_loop.rs) lands before Phase 4b, and its conclusion is a constraint on 4b.
• perf(rocm): recalibrate the joint top-k + top-p rejection vocab cap #2157 (rejection vocab cap that routes joint top-k + top-p sampling) lands before Phase 2.
• fix(models): stop model loaders from printing progress to stdout #1794 (loader progress in src/models/ and src/loading/ moves off stdout; src/commands/ output stays as is) has no ordering constraint with this epic: it does not touch src/commands/, and its model-file edits are loader prints, not the attention or MTP code the phases move. Either order rebases cleanly.
• Phase 0 keeps the bench_decode output lines that scripts/bench_decode.sh and scripts/rocm_decode_profile.py parse (Decode: / Prefill:, [Memory] from PR fix(rocm): bound in-flight batch memory and enforce the cache limit #2084, [phase] marks from PR perf(rocm): profile gfx1151 decode per kernel and rank the #1814 ports #2086), or updates both scripts in the same PR.

ROCm parity definition. ADR 0007 (Phase 0) should state what the parity invariant means on ROCm: MLXCEL_SDPA_DETERMINISTIC is read only by the CUDA overlay (src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp), and ROCm dense decode (MLX SDPA) and paged decode (HIP paged ports) differ in reduction order. #2192 measures it.

Acceptance Criteria

• All sub-issues completed and merged
• The Phase 0 parity harness shows greedy output identical across mlxcel generate, mlxcel run and the server (B=1 dense and paged) under MLXCEL_SDPA_DETERMINISTIC=1
• Single-stream decode stays within the Phase 0 threshold of the pre-epic CxxGenerator baseline, at short and long context
• docs/CONTINUOUS_BATCHING.md and the CLI docs describe a single engine, and the token-exactness disclaimer is replaced
• test(rocm): verify the unified engine on ROCm after the batch-native engine epic #2192 ROCm verification passes (make verify-rocm green on the final main, ROCm parity recorded)

Technical Considerations

ADR 0007 (Phase 0) supersedes ADR 0004's deferral of the KV, batching and scheduler abstraction, and it names the engine API. The decision table on CLI vs server defaults in ADR 0007 must be confirmed by the maintainer before Phase 5 relies on it. Per project policy, verify each phase on a real checkpoint and take performance measurements once, at the end of a phase, on a quiet host.

Two mechanisms already on main constrain the phases. Model-owned families can rewind decode appends (LanguageModel::supports_decode_lookahead_rewind and rewind_decode_appends, src/lib/mlxcel-core/src/generate.rs:698 and :712, backed by the RotatingKVCache undo log in src/lib/mlxcel-core/src/cache/decode_undo.rs, #2182), which is what keeps Gemma 3 on the scheduler's lookahead pipeline; Phases 1, 3, 4a, 4b and 6 must preserve it. With the prompt cache on, classic serving splits a chat prefill at the history boundary (capture_history_boundary_snapshot, src/server/batch/scheduler/prompt_cache.rs:1088, #1143), and the B=1 Gemma 4 MTP burst mirrors that split through history_boundary_split (:1212, #2185) because the partition alone changes greedy output. The #2185 parity harness (docs/benchmark_results/data/gemma4-mtp-cuda-gb10-2026-10-07/harness/) is prior art for Phase 0, and the split is a lead for Phase 3's prompt-cache divergence.

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