Repository navigation
perf(moe): fused MoE expert decode kernel, correctness-validated (#268 step 2a) - #275
Merged
Merged
Conversation
…step 2a) Adds a single-launch Metal kernel for the single-token decode expert path: for the K selected experts, in-kernel affine dequant-GEMV of gate/up, swiglu, dequant-GEMV of down, and the score-weighted combine, all in one launch. Built via fast::metal_kernel (the same JIT path as ssm_update_kernel). Handles power-of-2 bits (4/8); callers fall back to the SwitchGLU + moe_weighted_sum path for 6-bit, non-affine, mismatched bits/group_size across gate/up/down, non-single-token, or oversized threadgroup memory, so dots.llm1's 6-bit experts and every other config stay on the proven path. Wired into qwen3_moe behind MLXCEL_FUSED_MOE (default off; single-token only), plus the integration point on the shared SwitchGLU for the dots.llm1/nemotron path. On qwen3-30b-a3b the kernel output is byte-identical to the reference for both short and 40-token greedy generations, validating the dequant/GEMV/swiglu/combine math. Correctness-first: one threadgroup per token, so it is ~100x slower than MLX's multi-threadgroup gather_qmm (0.46 vs 47 tok/s) and stays off by default. Step 2b parallelizes it (multi-threadgroup tiling + SIMD reduction) to turn the validated correctness into the decode win; the design and validation gate are in docs/benchmark_results/fused-moe-decode-kernel-design.md. No default behavior change.
inureyes
added
type:performance
This was referenced Jun 14, 2026
perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE)
#268
Closed
7 of 8 tasks
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
…lup#268 step 2a) (lablup#275) Adds a single-launch Metal kernel for the single-token decode expert path: for the K selected experts, in-kernel affine dequant-GEMV of gate/up, swiglu, dequant-GEMV of down, and the score-weighted combine, all in one launch. Built via fast::metal_kernel (the same JIT path as ssm_update_kernel). Handles power-of-2 bits (4/8); callers fall back to the SwitchGLU + moe_weighted_sum path for 6-bit, non-affine, mismatched bits/group_size across gate/up/down, non-single-token, or oversized threadgroup memory, so dots.llm1's 6-bit experts and every other config stay on the proven path. Wired into qwen3_moe behind MLXCEL_FUSED_MOE (default off; single-token only), plus the integration point on the shared SwitchGLU for the dots.llm1/nemotron path. On qwen3-30b-a3b the kernel output is byte-identical to the reference for both short and 40-token greedy generations, validating the dequant/GEMV/swiglu/combine math. Correctness-first: one threadgroup per token, so it is ~100x slower than MLX's multi-threadgroup gather_qmm (0.46 vs 47 tok/s) and stays off by default. Step 2b parallelizes it (multi-threadgroup tiling + SIMD reduction) to turn the validated correctness into the decode win; the design and validation gate are in docs/benchmark_results/fused-moe-decode-kernel-design.md. No default behavior change.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Step 2a of the fused decode-MoE kernel (#268): a correctness-validated single-launch Metal kernel for the single-token decode expert path. Off by default; the optimization that makes it fast is step 2b.
What it does
One launch computes
out[h] = sum_k scores[k] * down_k( silu(gate_k(x)) * up_k(x) )[h]for the K selected experts: in-kernel affine dequant-GEMV of gate/up, swiglu, dequant-GEMV of down, and the score-weighted combine. Built viafast::metal_kernel(same JIT path asssm_update_kernel).Multi-bit handling (the non-4-bit point)
The kernel handles power-of-2 bits (4 and 8) via a
bitstemplate arg with a clean shift-unpack. For everything else it returnsNoneand the caller uses the provenSwitchGLU+moe_weighted_sumpath: 6-bit (dots.llm1'sdown_proj), non-affine, mismatched bits/group_size across gate/up/down, non-single-token, or threadgroup memory over the 32 KiB Metal limit. So no config ever loses correctness; 6-bit just isn't accelerated yet (MLX packs non-power-of-2 bits differently; that unpack is a 2b/follow-up item).Validation
Wired into qwen3_moe behind
MLXCEL_FUSED_MOE(default off, single-token only). Onqwen3-30b-a3b(uniform 4-bit) the kernel output is byte-identical to the reference for both short and 40-token greedy generations.Why it's off by default
Correctness-first: one threadgroup per token, so a single token has no parallelism to fill the GPU. It runs ~100x slower than MLX's multi-threadgroup
gather_qmm(0.46 vs 47 tok/s). Step 2b parallelizes it (multi-threadgroup tiling + SIMD reduction). The correctness (the hard part for a quantized fused kernel) is now proven, so 2b is pure optimization against a fixed reference.Scope
Issue #268 stays open (2b is the perf win). No default behavior change; the kernel only runs under the env flag. Builds on the design in
docs/benchmark_results/fused-moe-decode-kernel-design.md.