[pierreablin]

• Independent of Fused Metal kernels for swiglu experts combination, faster single-token MoE experts #2053

• kernels/softmax_topk.py: softmax routing after the router matmul (softmax + top-k with an optional renormalization) in one fused kernel.

• expert_routing now also takes DeepSeek-V3's group-limited selection: n_group, topk_group, and the score of the masked groups' experts.

• We use these new kernels for the models that route with them, each only under the conditions where the kernel computes its routing; otherwise we use the previous path. To break ties, the router picks the lowest index expert.

• Tests: softmax_topk and the groups against mlx ops, a Qwen3-MoE decode step, and 34 MoE blocks against their generic routing.

• LFM2: routing weights are now in float32 instead of bf16 in order to route it through the expert_routing kernel

Benchmark

mlx_lm.benchmark --model <MODEL> -p 128 -g 256 -n 3 -b 1

gives the following headlines on m5 ultra:

Headline numbers

• mlx-community/Ling-mini-2.0-4bit: 330.8 → 412.3 (1.25×)
• LiquidAI/LFM2.5-8B-A1B-MLX-4bit: 328.1 → 388.4 (1.18×)
• mlx-community/Qwen3-30B-A3B-4bit: 156.3 → 166.1 (1.06×)

There is also a mild speedup at batch size 8.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude found the optimization, wrote the kernels, and the benchmarks.

[pierreablin]

re ping @nastya236 ! ;)