One-kernel MoE combine for Qwen3.5/3.6 MoE prefill on the invariant lane (opt-in, depends on #549)

[jvmenen]

Depends on #549. This branch is #549 plus one commit; only the last commit (perf(prefill): one-kernel MoE combine for Qwen3.5/3.6 MoE on the invariant lane (opt-in)) is new here. Once #549 is merged, the diff reduces to that commit.

Summary

Opt-in MTPLX_A3B_MOE_PREFILL_COMBINE=1: on the batch-invariant prefill lane from #549, the tail of mlx-lm's Qwen3NextSparseMoeBlock (Qwen3.5/3.6 MoE) runs as one kernel. The routed experts' sorted output, its inverse permutation, the routing scores and the gated shared expert go straight into the existing Flash-Next combine kernel (mtplx/kernels/qwen4_moe_prefill_combine.py) instead of materializing three [rows, top_k, hidden] tensors. Bit-identical to the stock block. Default off, and only installed together with the lane.

Motivation

After the routed experts, the stock block unsorts the expert outputs to token order, multiplies by the scores and sums over top_k, each step writing a [rows, top_k, hidden] BF16 tensor, and then adds the shared expert. On the invariant lane every prefill forward from the lane's sorted minimum (128 tokens at 256 experts, top-8) already runs the experts expert-sorted, so the sorted output and its inverse permutation are available anyway. Flash-Next's prefill combine kernel does unsort, weighting, sum and shared add in one pass with the same rounding as the stock tail, and is generic in top_k and hidden.

Change

• mtplx/batch_invariant_prefill.py: BatchInvariantSwitchGLU gains sorted_route_applies(x, top_k) (prefill phase and at least the sorted minimum, i.e. no padding) and sorted_experts(x, indices): the stock sorted chain (_gather_sort, up/gate/down with sorted_indices=True) without the final unsort, returning the expert-sorted outputs and the inverse permutation. __call__ is unchanged apart from sharing the minimum computation.
• New mtplx/a3b_moe_prefill_combine.py: a subclass of the stock Qwen3NextSparseMoeBlock. When the lane's SwitchGLU reports sorted_route_applies, it computes the routing exactly as the stock block does, runs the stock shared expert and gate, sorted_experts, and the combine kernel; otherwise it calls the stock block. Installed by class swap; the parameter tree is unchanged.
• mtplx/runtime.py _install_batch_invariant_lane: installs the combine only when the switch is on and the lane was admitted; the lane's install report gains moe_prefill_combine_blocks.
• mtplx/kernel_selfcheck.py: new bitwise self-check lane a3b_moe_prefill_combine (128 rows, top-8 of hidden 2048, the lane's narrowest sorted forward at that geometry) against the stock unsort-weight-sum-add tail. Any difference or failure turns the route off for the process.
• Narrower forwards, decode, verify and the lone final prompt token keep the stock block.

How to enable: MTPLX_BATCH_INVARIANT_PREFILL=1 MTPLX_A3B_MOE_PREFILL_COMBINE=1. Without the lane the switch does nothing.

Tests

• New tests/test_a3b_moe_prefill_combine.py (Metal, 15 tests): the block is bit-identical to the stock block at 256 experts, top-8, at 128, 331, 2,048 and 4,096 rows; whole-prefill logits, hidden and cache bit-identical with and without the combine on a tiny Qwen3.5-MoE (tests/a3b_tiny_synth.py, the same helper as in the draft-head history PR); forwards below the sorted minimum, decode and the final token take the stock block; install only on the lane and only with the switch; off by default; the self-check validates the kernel at this geometry, and a failed self-check turns the route off.
• tests/test_batch_invariant_prefill.py passes.
• Full suite with MTPLX_CONFIG=/nonexistent: 9,494 passed, 67 skipped, 1 failed. The failure is test_laguna_model.py::test_laguna_s_2_1_ar_route_skips_qwen_performance_hooks, which fails the same way on main on a Mac with less than 85.3 GiB (addressed in Tests no longer read the user's ~/.mtplx/config.toml or depend on the Mac's memory size #535). Ruff: new files clean, no new findings in the changed files.

🤖 Generated with Claude Code