[Apple Silicon] [MLX] Fuse MoE combine multiply-reduce into one Metal kernel

[jlee5814]

Motivation

The reference MoE combine, (y * scores[..., None]).sum(axis=-2), runs a broadcast multiply and a reduction with a transient materialized between them. Fusing both into one Metal kernel saves a full read of y, the transient round trip, and one dispatch per MoE layer.

This PR began as the kernel plus its verification suite. Review exposed a real coverage bug in the eligibility guard (rank 3 assumption against a rank 4 production call site), which the fix generalized. The performance question required measured serving numbers, which required an integration path. The abstraction ask captured the JIT declaration pattern this kernel now shares with fused_swiglu.py. A founder-level review of the abstraction and its test suite followed; this revision applies that review's fixes.

Modifications

Guard generalization and bf16 scores

• SGLang never flattens batch and sequence upstream, so every real combine call site is rank 4 y with rank 3 scores (prefill [1, L, 4, 2048], batched decode [R, 1, 4, 2048]). The original rank 3 guard would have fallen back on every real forward pass; caught in review. Both target models run bf16 activations, and the original _SCORES_DTYPES rejected bf16 scores.
• The contract is now scores.shape == y.shape[:-1] with rank at least 3. Leading dims flatten into the kernel's row dimension before dispatch and the output reshapes back. bf16 scores are accepted.

Default off integration

• SGLANG_MLX_FUSE_MOE_COMBINE = EnvBool(False) in environ.py, plus patch_moe_combine_with_fused, a subclass swap on Qwen2MoeSparseMoeBlock and Qwen3MoeSparseMoeBlock, gated in _load_model beside the existing SwiGLU patch. Per call eligibility stays inside fused_combine, which falls back inline, so patching never changes results. Merging changes nothing for any user; the flag flip is a separate, evidence backed decision.

metal_jit abstraction

• hardware_backend/mlx/metal_jit.py: kernel files declare source once via the @metal_jit.kernel(...) decorator on a MetalJitOp subclass. The base class owns a can_fuse to dispatch_fused or dispatch_fallback template with typed NotFusable routing. Compile failures degrade via a memoized guard mirroring jit_kernel/moe_fused_gate.py's can_use_moe_fused_gate pattern (catch, log once, memoize, fall back) instead of crashing the forward pass.
• MetalJitKernel still owns the per dtype key compile cache and warm dedupe underneath, unchanged in behavior.
• Both kernels migrate as the first consumers. Name templates keep kernel names byte identical, so profiler labels and shader cache keys do not change; the SwiGLU warm set folds into warm_once.
• Patch registration stays in model_runner. The env flags' one reader should remain the load path, and the patch functions encode mlx_lm model class knowledge, not Metal JIT knowledge.

can_fuse is now abstract, per review

• A founder-level review found FusedGateQmvSiluMulKernel was a live example of the risk in an optimistic base default: it inherited can_fuse returning True unconditionally instead of declaring its own eligibility, relying entirely on dispatch_fused raising NotFusable before any observable effect to stay safe. MetalJitOp.can_fuse now raises NotImplementedError (metal_jit.py:123-125, matching this codebase's existing abstract-method idiom of a bare raise NotImplementedError, used already by dispatch_fused/dispatch_fallback on the same class; no abc.ABC is used anywhere in this tree). The class docstring states the rule directly: every op declares its own predicate, the base declares none.
• FusedMoeCombineKernel.can_fuse already had a real structural predicate (fused_combine.py:100-112); unchanged.
• FusedGateQmvSiluMulKernel.can_fuse (new, fused_swiglu.py:246-263) delegates to the same _check_eligibility regime check dispatch_fused/dispatch_fallback already share, so all three stay in sync by construction. This does not change the op's observable behavior: for structurally ineligible inputs, dispatch() now short-circuits straight to dispatch_fallback instead of first attempting dispatch_fused (which raised the identical diagnostic anyway); for eligible inputs whose Metal compile fails, can_fuse cannot see that and the routing is unchanged. SwiGLU's asymmetry is preserved exactly: dispatch_fallback still always raises the diagnostic chain rather than computing, since its real fallback needs wrapper layer context the op does not receive; only its predicate became explicit.
• Three test_metal_jit.py stub classes that exercised dispatch_fused/dispatch_fallback routing without overriding can_fuse (relying on the removed default) now declare an explicit can_fuse returning True, so they keep testing what they were written to test.

Tolerance test fix, per review

• The same review found the bf16 1 ULP, zero headroom bound in test_fused_combine.py could reach tens of thousands of "ULP" at production scale. Root cause, found this session: the huge values are a metric artifact, not a kernel regression. Near zero, fp16/bf16 use absolute (subnormal-adjacent) steps far finer than the ULP ceiling's relative-error assumption, and the two paths can even land on opposite sides of zero (e.g. 0.0 vs -1e-7), where the bit-pattern-distance metric's own documented same-sign assumption breaks down. allclose(atol=tol) already bounds these elements correctly.
• Fix: the integer-ULP ceiling is now checked only on elements at or above a 1e-3 ground-truth-magnitude floor (_ULP_FLOOR, test_fused_combine.py), where the bit-pattern-distance metric is meaningful. A live 100-seed-per-shape sweep across every shape in _SHAPES plus the real production prefill shape (1, 268, 4, 2048) found max ULP at or above the floor was 0-1 in every case, for both dtypes; the ceilings (_ROUNDED_MAX_ULP) are now 3 for both fp16 and bf16, roughly 3x headroom over that observed max. The check is not deleted or loosened past what the evidence supports.

Accuracy Tests

• test_fused_combine.py: 102 of 102 matrix cases fire the fused path. fp16 and bf16 scores are bit exact against the fp32 reference; fp32 scores are asserted within narrowing tolerance and a floor-masked integer-ULP ceiling (see Modifications). [hardware verified, this run] the updated bound was demonstrated seed robust: test_scores_fp32_correctly_rounded and the production-shape case in test_leading_dims_match_reference were run across 7 shifted seed offsets (0 through 60000) on this hardware, all passing, with observed ULP at or above the floor staying in {0, 1} throughout. Patch idempotency on real, small dimension mlx_lm blocks is a no op with bit identical outputs.
• test_metal_jit.py: cache identity, name formatting, warm dedupe, decorator validation, and the dispatch template's routing contract (fused, fallback, NotFusable propagation, now including can_fuse as a required override), including the compile failure memoization guard.
• test_fused_swiglu.py: regime diagnostics through the wrapper, dtype membership fallback (not just dtype agreement), warm skip logging, warmup enumeration, and (new) the explicit can_fuse predicate.
• 47 tests across the three files at head (17 + 10 + 20, pytest --collect-only) [hardware verified, this run]. Unchanged from before this revision: the fixes changed test bodies and one helper signature, not test count.
• [hardware verified, this run] Full suite: 45 passed, 2 skipped (opt-in, HF-checkpoint-gated, not a hardware skip), 0 failed, both before and after this revision's changes.
• [hardware verified, this run] The three NotFusable safety sub claims from the prior review (forced Metal compile failure on an eligible input: no raise to caller, bit exact output, no partial/corrupted cache state) were re-run against the real fused_combine() entry point after the can_fuse change and still hold; the equivalent check on FusedGateQmvSiluMulKernel's own compile-failure path (now gated by its own can_fuse) also holds.

Speed Tests and Profiling

Original end to end measurement [author reported, not reverified end to end], mlx-community/Qwen1.5-MoE-A2.7B-Chat-4bit on a 24 GB Apple Silicon machine, interleaved on/off server launches, warmup discarded, temperature 0, medians with min and max:

config	decode bs=1 tok/s	per step ms	decode bs=4 tok/s	prefill L=268 tok/s
unfused	131.02 [130.61, 131.67]	7.632	201.20 [199.13, 203.11]	758.4 [756.0, 815.8]*
fused	131.80 [131.64, 132.11]	7.587	203.56 [201.82, 205.55]	818.2 [809.8, 820.1]

* One unfused launch at 9.7 GB free excluded as memory pressure (732.7 tok/s).

Decode bs=1 is +0.6 percent, decode bs=4 is +1.2 percent, prefill is +0.0 to +1.5 percent depending on which launches are compared.

Isolated op measurement, decode bs=1 [hardware verified, this run]: re-ran the fused-vs-fallback dispatch check on this hardware after the can_fuse change. dispatch() still routes through exactly one metal_jit.get("fused_moe_combine", ...) call on the fused path, output still bit exact against an independent reference, and the fused path still measures faster than the fallback (paired timing this run: +2.4%, a simpler protocol than the prior review's bootstrapped interleaved A/B, so not a direct comparison to the fuller table below).

Prior isolated op measurement [prior review, not rerun this session], Apple M4 Pro, 24 GB, bootstrap 95 percent confidence interval:

shape	fused (us)	fallback (us)	speedup	95% CI excludes zero
decode bs=1	17.87	18.80	+9.6%	yes
decode bs=4	18.39	19.54	+7.4%	yes
prefill L=268	34.83	190.94	+82.2%	yes

The two end-to-end/isolated tables are not directly comparable: the isolated measurement strips out every other per token cost (attention, routing, the expert GEMMs, sampling, KV bookkeeping), so it upper bounds what the combine op itself could contribute, while the end-to-end measurement is what actually reaches a user after dilution by the rest of the forward pass. The metal_jit registry promotion and the can_fuse abstraction change are both same-process refactors with no dispatch-count or numerics change on the fused path, confirmed this run; neither table's numbers are affected by them.

Known Limitations

• warmup_specs declares each op's supported variant set and has no consumer yet. The working warm path (SwiGLU only; combine pays its first dispatch JIT cost today) warms instantiated shapes derived from the live model. aot.py reconciles the two in a follow up PR.
• The patch time can_fuse(switch_mlp) free function in fused_swiglu.py is unrelated to MetalJitOp.can_fuse; a rename is follow up material. This is a naming collision only; both FusedMoeCombineKernel and FusedGateQmvSiluMulKernel now declare their own MetalJitOp.can_fuse explicitly (see Modifications), so the base class no longer has an optimistic default for a future op to silently inherit.
• SwiGLU's dispatch_fallback intentionally raises rather than computes, since its real fallback needs wrapper layer context the op does not receive. Unifying that is follow up scope.
• Nothing counts fallback occurrences on either op (no logging on FusedMoeCombineKernel's fallback path, a once-per-process warning on SwiGLU's). This was flagged by review as a real but lower-leverage gap for this pass and is not addressed here.
• The register_mlx_ci addition to ci_register.py this branch carries is now byte identical to what [Apple Silicon] [CI] Move the MLX lane to the check-changes + pr-gate composite #30121 independently merged to main; a real merge auto-resolves it with no manual edit, confirmed this run (git hash-object matches on both sides). Attempting the full merge this run surfaced a separate, unrelated conflict: an upstream commit (#32448, merged after this branch forked) relocated test_fused_swiglu.py to test/registered/ and rewrote its CI-harness registration, colliding with this branch's own, differently-evolved copy of that file at its old path. That conflict is outside what the prior review named and was not resolved unilaterally in this pass; it is reported here as the actual remaining blocker to a clean merge, and needs a codeowner call on how to reconcile the two files' content, not just their location.
• Mergeable state at time of writing needs the test_fused_swiglu.py relocation above resolved before this branch can merge cleanly against current main; the ci_register.py duplication that previously blocked it is confirmed moot.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. No user facing documentation exists yet for the MLX backend's fusion flags; out of scope for this PR, flagged here rather than silently skipped.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS (yeahdongcn for this path) and any other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
4. After green CI and required approvals, ask Merge Oncalls or a maintainer with write permission to merge.

---

CI States

Latest PR Test (Base): ❌ Run #31957166728
Latest PR Test (Extra): ❌ Run #31957166641

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 975232d23f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[yeahdongcn]

What is the e2e perf gain?

[yeahdongcn]

Could we introduce some abstraction here for declaring JIT kernels on Apple Silicon? Since we already have python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py, it seems we're starting to establish common patterns for things like registration, warmup, and related initialization. It would be good to capture those patterns so future JIT kernels can follow a consistent approach.

[jlee5814]

Added metal_jit.py in ab47947.

MetalJitKernel owns the source, per dtype lazy compile, warmup dedupe, and name mangling. Both kernels migrate as first consumers; kernel names stay byte identical. Patching stays in model_runner: the patch functions encode mlx-lm model knowledge, not Metal JIT knowledge. Tested in test_metal_jit.py, plus double patch is a no op.

[jlee5814]

Registry follow-up landed as well; see the thread on metal_jit.py for details. Declaring a new JIT kernel is now source plus one registration call.

[yeahdongcn]

can_fuse() only accepts y.ndim == 3 and scores.ndim == 2, why? The mlx-lm Qwen3-MoE combine site has y shaped like [batch, seq, top_k, hidden] and scores shaped like [batch, seq, top_k] before y = (y * scores[..., None]).sum(axis=-2).

[jlee5814]

Fixed in a34b78.

SGLang never flattens upstream; the guard was written against a shape that never arrives. Contract is now scores.shape == y.shape[:-1], rank >= 3, leading dims flattened; bf16 bit exact. 102 matrix cases pass, including the serving shapes.

[jlee5814]

What is the e2e perf gain?

Measured on Qwen1.5-MoE-A2.7B-Chat-4bit, eight interleaved launch pairs, medians [min, max]. Decode bs=1: 131.02 [130.61, 131.67] vs 131.80 [131.64, 132.11] tok/s, +0.6%, one saved dispatch plus transient per MoE layer. bs=4: +1.2%, spreads overlap. Prefill L=268: +0.0 to +1.5% memory matched, but the unfused baseline hit a discrete slow mode (756 to 758 tok/s, 4 of 8 launches) that fused never entered in 8 of 8. Flag ships default off.

[yeahdongcn]

MetalJitKernel is a useful step, but I do not think it fully removes the repeated pattern yet. Each fused kernel still has to declare a module-level _KERNEL = MetalJitKernel(...) and keep a local _get_kernel() wrapper mostly for caching/test seams.

Since we are starting to add more MLX JIT kernels, can we make metal_jit.py own a small registry/factory instead? For example, kernels could be registered by name/template and looked up through one shared API that owns dtype-key caching, name formatting, warmup bookkeeping, and the test seam. Then fused_combine.py and fused_swiglu.py would only define the kernel source plus call the registry, instead of each file carrying the same _KERNEL/_get_kernel pattern.

[jlee5814]

Confirmed. A registry would remove the per-file _KERNEL + _get_kernel pair. One thing not visible from the diff: the wrapper was the dispatch tests' spy seam, not leftover pattern. The registry subsumes it: tests patch the registry lookup, a cleaner single seam. Will add the name keyed registry, migrate both kernels to source plus registration, and move the seam.

[jlee5814]

Resolved in 4e1201f.

Kernels reduce to source plus registration, with the registry owning dtype caching, naming, warmup and the test seam. d684d3d adds the dispatch template on top per the follow-up sketch.

[yeahdongcn]

Although these tests have been added, could we also make sure they're actually executed by the CI in #30121?

[jlee5814]

Yes, model free and in stage A's explicit list, so they run on every PR once #30121 lands.

[yeahdongcn]

How about something like this?

@metal_jit.kernel(
    name="fused_moe_combine",
    name_template="fused_moe_combine_y{0}_s{1}",
    input_names=["y", "scores"],
    output_names=["out"],
)
class FusedMoeCombineKernel(MetalJitOp):
    source = _KERNEL_SOURCE

    def dispatch(...): ...

    def warmup_specs(model) -> Iterable[WarmupSpec]: ...

[jlee5814]

@metal_jit.kernel plus MetalJitOp only, SwiGLU migrated, register deleted. warmup_specs stubbed per your request, real specs in the aot.py follow up. Sources byte identical, tests pass.

[yeahdongcn]

I just DMed you on Slack. Maybe we can have a quick discussion?

[jlee5814]

SwiGLU's can_fuse is newly written (the module level one is a patch time check), and its fallback re-raises the existing diagnostics since the nn path needs wrapper context. Deferred: docstrings, can_patch rename, NotFusable.

[yeahdongcn]

  for module_name, class_name, fused_call in (
      ("mlx_lm.models.qwen2_moe", "Qwen2MoeSparseMoeBlock", _fused_qwen2_moe_call),
      ("mlx_lm.models.qwen3_moe", "Qwen3MoeSparseMoeBlock", _fused_qwen3_moe_call),
  ):
      try:
          module = importlib.import_module(module_name)
      except ImportError:
          continue
      targets.append((getattr(module, class_name), fused_call))

Looks better.

[jlee5814]

This has been adopted.

[yeahdongcn]

It seems not necessary.

[jlee5814]

This has been dropped.

[yeahdongcn]

Should we have the new register_mlx_ci?

[jlee5814]

Fixed.

Both files carry register_mlx_ci(suite="stage-a-unit-test-mlx") alongside register_cpu_ci per the lane convention. The marker is backported verbatim from #30121.

[yeahdongcn]

Can we do it in this PR?

[jlee5814]

Done.

Both ops enumerate WarmupSpecs from declared dtype support; aot.py consumes them in the follow up. Also fixed a gap this surfaced: eligibility checked dtype agreement but not membership, so matched fp32 could reach the unvalidated kernel; such dtypes now fall back.

[jlee5814]

/tag-and-rerun-ci