perf(moe): batched decode scales 1.55x against 3.25x on dense #1616

Problem / Background

Under continuous batching on M5 Max both dense models scale close to linearly to B=4 while the MoE model does not. qwen3-30b-a3b-4bit reaches only 1.55x aggregate at B=4, per-request decode falls 61% (169.6 to 65.5 tok/s), and TTFT rises 2.5x (149.1 to 379.2 ms). Four concurrent clients buy 55% more aggregate throughput while each one sees a 61% slower stream.

mlxcel 0.6.0 at commit b2ff1eee, MacBook Pro M5 Max 128GB, macOS 26.6.2, MLX pin 9a795735. scripts/bench_serving_concurrency.py against one fresh mlxcel-server per model, pinned to --parallel 4 --max-batch-prefill 4 at --prompt-tokens 512 --max-tokens 128. Source CSV: benchmarks/metal_m5max_batch_2026-09-04.csv. N concurrent streaming clients occupy N decode slots, so the concurrency level is the effective decode batch size. 0 fail across all nine cells.

This is the MoE decode path, not the scheduler: qwen2.5-0.5b-bf16 is in the trio to isolate scheduler overhead and scales fine, only the MoE model collapses. Levels run ascending and share the prompt, so B=2 and B=4 start with a warm prompt cache while B=1 pays the cold prefill; that is why dense TTFT falls from B=1 to B=2. MoE TTFT rises monotonically despite the same advantage, which strengthens the result. SSM, hybrid and mixed-cache families are excluded because the server serializes their slots, so a B-ladder over them measures nothing.

Current Behavior

The fused MoE decode kernel is a B=1-only path. src/models/qwen3_moe.rs:446-452 gates it on mlxcel_core::array_shape(&x_flat)[0] == 1 && fused_moe_enabled(), where x_flat is [batch * seq_len, hidden]. Fifteen families carry the identical gate (grep -rn 'array_shape(&x_flat)\[0\] == 1' src/models/, plus b * s == 1 at src/models/gemma4.rs:1045).

The scheduler supplies a batch dimension as soon as there is more than one slot: dispatch_sync_decode at src/server/batch/scheduler/decode_tick.rs:269-277 uses per-sequence decode_single_step only when seq_ids.len() <= 1 || !self.model.supports_batching(), otherwise execute_batched_decode builds the input as [b, 1] at decode_tick.rs:747. So B=1 runs the fused kernel and B>=2 silently falls back to SwitchGLU::forward plus moe_weighted_sum (gather_qmm) on all 48 layers, every tick.

SwitchGLU::forward_fused_kernel repeats the restriction independently at src/models/switch_layers.rs:914-917 (x_elems != din returns None), and the kernel has no batch dimension at all: run_fused_moe_two_kernel at src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:1957 declares outputs Shape{k, dff} and Shape{k, din} over grids (32, round_up(dff, sgy), k) and (32, round_up(din, sgy), k). Relaxing the call-site gate alone would not reach a batched kernel.

The Dff cap is not the cause: models/qwen3-30b-a3b-4bit/config.json is 48 all-MoE layers (decoder_sparse_step 1), 128 experts, top_k 8, moe_intermediate_size 768, 4-bit group 64, so Dff 768 sits far under FUSED_MOE_MAX_DFF_METAL (4096, src/models/switch_layers.rs:106). Only the token-count gate declines it.

Leading hypothesis for the 61% per-request loss, to be measured and not assumed: neither fallback branch deduplicates expert ids across the cohort. SwitchGLU::forward (switch_layers.rs:761-786) sorts only when n_tokens * top_k >= 64, so at top_k 8 the B=1/2/4 steps all stay unsorted (8/16/32 slots) and B=8 is the first to cross; gather_sort (switch_layers.rs:1022-1044) sorts by expert id but still emits one row per (token, k) slot. Expert-plane traffic therefore grows with B * top_k while a dense model reads its weights once per step regardless of B.

Proposed Solution

Attribute first, with the mlxcel-gpu-profiling hooks already applied in docs/benchmark_results/moe-decode-gap-investigation.md: confirm gather_qmm is reached on every layer at B=4, split per-step cost into router, expert GEMV and combine at B=1 against B=4, and determine whether the 2.5x TTFT growth is prefill cohorts serialized behind MoE decode (instrument src/server/batch/prefill_cohort.rs) or just a longer decode tick. Then, in expected-payoff order, with the profile deciding what is worth building:

• Give the fused kernel a token dimension: extend run_fused_moe_two_kernel to outputs [n_tokens, k, dff] and [n_tokens, k, din] with grid z of k * n_tokens, sum over k per token, and relax the call-site gates from == 1 to a measured n_tokens cap resolved the way fused_moe_max_dff_from resolves the Dff cap. At n_tokens == 1 the dispatch stays bit-identical to today.
• Deduplicate expert ids across the cohort so each expert plane is read once per step rather than once per (token, k) slot.
• Re-tune the do_sort constant 64 in SwitchGLU::forward against batch size rather than total slots, so the B=8 crossing becomes a measured choice.

Rejected, so it is not relitigated: forcing supports_batching() false for MoE restores the fused kernel per sequence but serializes B forward passes per tick, so it cannot beat the 1.55x it replaces and pushes TTFT higher. Do not take that route without a measurement contradicting this.

Scope

In scope: src/models/switch_layers.rs, the fused kernel in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp and its bridge at src/lib/mlxcel-core/src/lib.rs:1720, the per-family call-site gates, and any decode_tick.rs change the profile justifies.

Out of scope: the B=1 MoE gap vs mlx-lm (#268), prefill-side MoE work, CUDA batched quantized GEMM (#725), and SSM/hybrid families whose slots the server serializes.

Implementation Notes

• Reuse: the decline-and-fall-back contract (forward_fused_kernel returns None, caller runs SwitchGLU::forward plus moe_weighted_sum), fused_moe_max_dff_from as the backend-aware runtime cap pattern, MLXCEL_FUSED_MOE=0, and the MLXCEL_FUSED_MOE_PARITY_CHECK=1 probe.
• Constraints: greedy B=1 output stays byte-identical on qwen3-30b-a3b-4bit; fused_moe_parity_tests in mlxcel-core pins the kernel against the all-f32 dense reference and must cover any new batched shape; the bug: Gemma 4 26B A4B multi-turn corruption from numerically unfaithful fused decode-MoE kernel #886 discipline of f32 partials with a single final rounding must carry over, since early per-expert rounding previously caused 5-13% relative error.
• Edge cases: B changes tick to tick as sequences join and leave, so the JIT template must not respecialize per B (keep templating on k/Din/Dff/bits/group_size and pass n_tokens as a runtime grid extent); mixed widths must keep working (dots.llm1 is gate/up 4-bit with down 6-bit, and 6-bit down needs Dff divisible by 16); the GeGLU variant fused_moe_geglu_kernel (gemma4) needs the same treatment or an explicit decline; a cohort routing to fully disjoint experts is the worst case for plane traffic and belongs in the measurement.
• Error handling: unsupported configurations keep returning None and fall back to gather_qmm rather than failing. The bridge crosses as UniquePtr<MlxArray> rather than Result, so a shape mistake surfaces as an uncatchable C++ terminate at the first routed forward; validate shapes on the Rust side before dispatch.

Acceptance Criteria

• A recorded profile attributes the B=4 MoE gap to named code paths and ops rather than asserting a cause, committed under docs/benchmark_results/.
• Either B=4 aggregate scaling for qwen3-30b-a3b-4bit measurably exceeds 1.55x on the same harness and host, or a documented explanation of why MoE cannot batch like dense here lands in the same place.
• Dense B=4 rows unchanged within noise: llama-3.1-8b-4bit at or above roughly 3.2x, qwen2.5-0.5b-bf16 at or above roughly 3.1x.
• MoE B=1 decode unregressed: greedy output byte-identical and decode throughput within noise of 169.6 tok/s.
• cargo test --release -p mlxcel-core fused_moe_parity_tests passes, extended to the batched shape if a batched kernel lands.
• The gain is visible end to end through mlxcel-server under scripts/bench_serving_concurrency.py, not only in a microbenchmark, with a fresh CSV under benchmarks/.
• docs/benchmark_results/model_tests_m5max.md updated at the existing "Batched serving (B = 1/2/4)" section (line 501) with the new numbers and the attribution.

Verification

cargo build --release --features metal,accelerate
cargo test --release -p mlxcel-core fused_moe_parity_tests
cargo clippy --all-targets -- -D warnings
cargo fmt --check

# B=1 parity guard: greedy output must be byte-identical before and after.
./target/release/mlxcel generate -m models/qwen3-30b-a3b-4bit -p "Hello" -n 128 --temp 0

# Re-run the ladder that produced the numbers above, same flags and host.
python3 scripts/bench_serving_concurrency.py --parallel 4 --max-batch-prefill 4 --prompt-tokens 512 --max-tokens 128 --concurrency 1,2,4

A pass is: parity and lint gates green, B=1 greedy output byte-identical to the pre-change run, dense B=4 scaling unchanged within noise, and either a higher B=4 aggregate for qwen3-30b-a3b-4bit or a written explanation backed by the profile.

Technical Considerations

Related: #268 produced the fused single-token kernel this issue finds unreachable at B>1; #725 covered small-M batched decode quantized GEMM on CUDA; #628 and #632 set the serving defaults and lookahead decode pipeline this ladder exercises.