You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
On NVIDIA and AMD devices, this is the architecture of the MoE computation flow (#8715):
[input_hidden_states]
|
v
TopK.forward ->`select_experts` / `triton_kernels.routing` / bypass
|
V
[TopKOutput]
|
v
FusedMoE.forward -> Dispatcher.dispatch -> DeepEP / PPLX / bypass
||| v
| [DispatchOutput]
||| v
| quant_mothod.apply -> MoeRunner.forward -
||||| v
|| pre-permute + grouped_gemm + post-permute
|||||--------------------------------
| v
| [CombineInput]
||| v
| Dispatcher.combine -> DeepEP / PPLX / bypass
|||---------------------
v
[final_hidden_states]
There are a few custom kernels to make this run very fast. To concretely show the computations and time breakdown, I profiled a forward pass of Qwen3-30B-A3B on an NVIDIA B300 using this command to launch the server and the next command to send a request to it:
Let's analyze the custom kernels. The first custom kernel that is launched is topkGatingSoftmax. This custom kernel was written to fuse 3 separate operations:
The second custom kernel that is launched is Count_and_sort_expert_tokens_kernel. This kernel has a very simple task: sort an integer array by each element's integer value. The input to this kernel an integer array (technically a Torch.Tensor) and the output is the sorted integer array. The input array is the top_indices from the topkGatingSoftmax step. These indices are the expert ID that each token is sent to. Thus, if the config sets topk=4 (meaning each token is sent to 4 experts), then top_indices will have shape (num_tokens, 4) because there are 4 expert IDs assigned for each token. Count_and_sort_expert_tokens_kernel will then flatten top_indices into an array of shape num_tokens x 4, sort the array so tokens that are sent to the same expert are next to each other. An example output is:
Token 0 is sent to Expert 0, Token 1 is sent to Expert 0, etc.
The third custom kernel that is launched is Fused_moe_kernel. This kernel is launched twice: once for the gate_up_proj and once for the down_proj. In between each launch, Act_and_mul_kernel is launched and applied to the intermediate tensor.
Now, let's explore how the MoE layer is computed in SGLang currently (it just uses MLX-LM under the hood), which is our baseline that we are attempting to optimize to achieve higher performance.
I profiled mlx-community/Qwen1.5-MoE-A2.7B-4bit using this command on my MacBook Pro M1, 16Gb RAM, Sequoia 15.6:
Before analyzing the trace, I looked through the MLX-LM and MLX repo to understand the code and know what to look for. For Qwen1.5-MoE-A2.7B, MLX-LM uses a class called SwitchGLU for the MoE layer.
Here is the implementation of SwitchLinear in MLX-LM:
The main expert computation (equivalent of Fused_moe) is the MLX kernel gather_mm.
Now analyzing the trace:
We can see the MoE computations:
affine_qmv_float16_t_gs_64_b_4_batch_0 (scores = router(token)) -> block_softmax_precise_float16 (probabilities = softmax(scores)) -> carg_block_sort_float16_uint32_bn32_tn4 (top_weights, top_indices = topk(probabilities, k=4) -> and then the expert computation, activation, and expert computation.
To optimize further, we need to 1) fuse kernels 2) optimize kernels. Thus, now that we know what kernels are used, we need to study the implementation of the kernels and figure out if there are any optimizations that can be done inside the kernel implementation to make it faster and if we can fuse any kernels to make inference faster.
Profiled Qwen1.5 MoE A2.7B-4bit decode on M4 Pro. MoE cost splits are as follows:
Findings:
K=1408 fails the K%512 check in gather_qmv (quantized.cpp:984), so the
routed down projection falls to the slower kernel variant.
Wrote a custom mx.fast.metal_kernel for the non-fast case. Output matches
MLX's quantized_matmul on production shapes (K=1408, N=2048, 60 experts, top-k 4).
Next: benchmark, fp16, integrate into the MLX MoE forward.
quantized.cpp gates the qmv and gather_qmv fast paths on K % 512, but qmv_fast_impl derives its block size from bits, which is 256 at 6 and 8 bits. The gate has been stricter than the kernel since March 2024. fp_qmv_fast_impl shares it, so MXFP8 at 8 bits hit the same wall.
ml-explore/mlx#3965 derives the alignment from bits. 186 correctness cases pass, Metal shader validation is clean, gather_qmm_bench is unchanged. On an M4 Pro, 8 bit K=768 gains about 18% and K=1280 about 10%, with every 4 bit control flat.
Fixed: Qwen1.5-MoE-A2.7B K=1408 and Qwen3-30B-A3B K=768, at 6 and 8 bits.
Still broken: 4 bit, where the kernel really does need 512 alignment. Most Apple Silicon serving is 4 bit, so both models stay on the slow gather_qmv path. Keeping this open.
Three options for 4 bit: a custom mx.fast.metal_kernel for the unaligned case, padding K, or relaxing the kernel to handle a tail. My prototype for the first is correct but 3.3x slower than the slow path it replaces.
ml-explore/mlx#3965 merged on 8/7; it was still open when the previous comment posted.
K=1408 eligibility by bit width, from the merged qmv_fast_k_alignment gate, confirmed on an M4 Pro by direct dispatch readout rather than timing (GEMV timing at small K is too variant to carry a conclusion, per the merging maintainer on ml-explore/mlx#3965):
bits
alignment required
K=1408 eligible
2
512
no
3
512
no
4
512
no
5
512
no
6
256
no
8
256
no
1408 is 128 x 11; 128 falls below both alignments.
Correction: the previous comment listed Qwen1.5-MoE-A2.7B (K=1408) alongside Qwen3-30B-A3B (K=768) as fixed at 6 and 8 bits. Only K=768 is fixed; K=1280 also benefits, verified on the pre merge branch. K=1408 was never eligible, before or after the merge; the merged build's dispatch, including gather_qmv, confirms it.
Qwen1.5-MoE-A2.7B keeps paying the slow gather_qmv path at every bit width and still needs a dedicated fix: a custom kernel, K padding, or a kernel that tolerates a tail.
Activity
jonahbernard commentedon Apr 7, 2026
On NVIDIA and AMD devices, this is the architecture of the MoE computation flow (#8715):
[input_hidden_states] | v TopK.forward -> `select_experts` / `triton_kernels.routing` / bypass | V [TopKOutput] | v FusedMoE.forward -> Dispatcher.dispatch -> DeepEP / PPLX / bypass | | | v | [DispatchOutput] | | | v | quant_mothod.apply -> MoeRunner.forward - | | | | | v | | pre-permute + grouped_gemm + post-permute | | | | |-------------------------------- | v | [CombineInput] | | | v | Dispatcher.combine -> DeepEP / PPLX / bypass | | |--------------------- v [final_hidden_states]There are a few custom kernels to make this run very fast. To concretely show the computations and time breakdown, I profiled a forward pass of
Qwen3-30B-A3Bon an NVIDIA B300 using this command to launch the server and the next command to send a request to it:CUDA_VISIBLE_DEVICES=7 PYTHONPATH=/my-sglang/sglang/python nsys profile -t cuda,nvtx -o triton_moe_server_profile \ python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-30B-A3B \ --moe-runner-backend triton \ --disable-cuda-graph \ --disable-piecewise-cuda-graph \ --port 30004Here is the full trace file. Please unzip it and open it with Nsight Systems.
nsys-report-bac2.nsys-rep.zip
I added NVTX markers around the 3 main MoE components (router, topk, and fused_moe):
Here are the 3 NVTX markers that I added to get these labels:
In these 3 components, we see that 9 kernels are launched (8 unique kernels are used):
Router:
Top-K:
Fused_MoE:
Out of these 8 unique kernels, only 5 of them are custom kernels written by the SGLang team for fast inference:
topkGatingSoftmax
Count_and_sort_expert_tokens_kernel
Fused_moe_kernel
Moe_align_block_size_kernel
Act_and_mul_kernel
The other 3 are auto-generated kernels at inference time:
nvjet_sm103_tst_32x64_64x16_4x1_v_bz_splitK_TNN
splitKreduce_kernel
triton_per_fused_copy__mul_sum_0
Let's analyze the custom kernels. The first custom kernel that is launched is
topkGatingSoftmax. This custom kernel was written to fuse 3 separate operations:scores = router(token)
probabilities = softmax(scores)
top_weights, top_indices = topk(probabilities, k=some integer)
The second custom kernel that is launched is
Count_and_sort_expert_tokens_kernel. This kernel has a very simple task: sort an integer array by each element's integer value. The input to this kernel an integer array (technically a Torch.Tensor) and the output is the sorted integer array. The input array is thetop_indicesfrom thetopkGatingSoftmaxstep. These indices are the expert ID that each token is sent to. Thus, if the config setstopk=4(meaning each token is sent to 4 experts), thentop_indiceswill have shape (num_tokens, 4) because there are 4 expert IDs assigned for each token.Count_and_sort_expert_tokens_kernelwill then flattentop_indicesinto an array of shapenum_tokens x 4, sort the array so tokens that are sent to the same expert are next to each other. An example output is:In plain English:
Token 0 is sent to Expert 0, Token 1 is sent to Expert 0, etc.
The third custom kernel that is launched is
Fused_moe_kernel. This kernel is launched twice: once for thegate_up_projand once for thedown_proj. In between each launch,Act_and_mul_kernelis launched and applied to the intermediate tensor.Now, let's explore how the MoE layer is computed in SGLang currently (it just uses MLX-LM under the hood), which is our baseline that we are attempting to optimize to achieve higher performance.
I profiled
mlx-community/Qwen1.5-MoE-A2.7B-4bitusing this command on my MacBook Pro M1, 16Gb RAM, Sequoia 15.6:This profile took ~40 minutes to complete.
Before analyzing the trace, I looked through the MLX-LM and MLX repo to understand the code and know what to look for. For Qwen1.5-MoE-A2.7B, MLX-LM uses a class called
SwitchGLUfor the MoE layer.Here is the implementation of
SwitchLinearin MLX-LM:The main expert computation (equivalent of
Fused_moe) is the MLX kernelgather_mm.Now analyzing the trace:
We can see the MoE computations:
affine_qmv_float16_t_gs_64_b_4_batch_0 (scores = router(token)) -> block_softmax_precise_float16 (probabilities = softmax(scores)) -> carg_block_sort_float16_uint32_bn32_tn4 (top_weights, top_indices = topk(probabilities, k=4) -> and then the expert computation, activation, and expert computation.
To optimize further, we need to 1) fuse kernels 2) optimize kernels. Thus, now that we know what kernels are used, we need to study the implementation of the kernels and figure out if there are any optimizations that can be done inside the kernel implementation to make it faster and if we can fuse any kernels to make inference faster.
jlee5814 commentedon Apr 14, 2026
Update on #22283:
Profiled Qwen1.5 MoE A2.7B-4bit decode on M4 Pro. MoE cost splits are as follows:
Findings:
K=1408 fails the
K%512check ingather_qmv(quantized.cpp:984), so therouted down projection falls to the slower kernel variant.
Wrote a custom
mx.fast.metal_kernelfor the non-fast case. Output matchesMLX's
quantized_matmulon production shapes (K=1408, N=2048, 60 experts, top-k 4).Next: benchmark, fp16, integrate into the MLX MoE forward.
[Apple Silicon] [MLX] Fuse SwiGLU activation into gate gather_qmv for…
github-actions commentedon Jun 14, 2026
This issue has been automatically closed due to inactivity. Please feel free to reopen it if needed.
jlee5814 commentedon Jun 16, 2026
FusedSwitchUpGate:up_proj+gate_proj([Apple Silicon] [MLX] Fuse SwitchGLU up_proj + gate_proj #24712)SwiGLUactivation into gategather_qmv: ([Apple Silicon] [MLX] Fuse SwiGLU activation into gate gather_qmv for SwitchGLU MoE blocks #26188)jlee5814 commentedon Aug 4, 2026
quantized.cppgates theqmvandgather_qmvfast paths onK % 512, butqmv_fast_implderives its block size from bits, which is 256 at 6 and 8 bits. The gate has been stricter than the kernel since March 2024.fp_qmv_fast_implshares it, so MXFP8 at 8 bits hit the same wall.ml-explore/mlx#3965 derives the alignment from bits. 186 correctness cases pass, Metal shader validation is clean,
gather_qmm_benchis unchanged. On an M4 Pro, 8 bitK=768gains about 18% andK=1280about 10%, with every 4 bit control flat.Fixed: Qwen1.5-MoE-A2.7B K=1408 and Qwen3-30B-A3B K=768, at 6 and 8 bits.
Still broken: 4 bit, where the kernel really does need 512 alignment. Most Apple Silicon serving is 4 bit, so both models stay on the slow
gather_qmvpath. Keeping this open.Three options for 4 bit: a custom
mx.fast.metal_kernelfor the unaligned case, padding K, or relaxing the kernel to handle a tail. My prototype for the first is correct but 3.3x slower than the slow path it replaces.jlee5814 commentedon Aug 10, 2026
ml-explore/mlx#3965merged on 8/7; it was still open when the previous comment posted.K=1408 eligibility by bit width, from the merged
qmv_fast_k_alignmentgate, confirmed on an M4 Pro by direct dispatch readout rather than timing (GEMV timing at small K is too variant to carry a conclusion, per the merging maintainer onml-explore/mlx#3965):1408 is 128 x 11; 128 falls below both alignments.
Correction: the previous comment listed Qwen1.5-MoE-A2.7B (K=1408) alongside Qwen3-30B-A3B (K=768) as fixed at 6 and 8 bits. Only K=768 is fixed; K=1280 also benefits, verified on the pre merge branch. K=1408 was never eligible, before or after the merge; the merged build's dispatch, including
gather_qmv, confirms it.Qwen1.5-MoE-A2.7B keeps paying the slow
gather_qmvpath at every bit width and still needs a dedicated fix: a custom kernel, K padding, or a kernel that tolerates a tail.