Skip to content

MoE top_k override (--moe-top-k)

Reduces the number of experts activated per token in Mixture-of-Experts models like Qwen3-30B-A3B, trading a small amount of quality for meaningfully higher decode throughput.

Status: opt-in flag. Default behaviour is unchanged. Quality numbers below are for Qwen3-30B-A3B-4bit on M4 Max 128 GB — verify on your model before shipping this to production workloads.

What it does

Qwen3-30B-A3B is trained with top_k=8 — every token picks 8 out of 128 experts. On Apple Silicon at batch=1 decode the expert matmul (SwitchGLU) is the single biggest chunk of each layer's compute, and that cost scales roughly linearly with top_k. Lowering top_k at inference time has been shown (LExI 2025, Lynx 2024) to preserve most of the trained quality while cutting decode time materially.

--moe-top-k N iterates every layer of the loaded model, and on each layer that has .mlp.switch_mlp (i.e. a sparse-MoE block) sets top_k = N. Dense layers and dense models are untouched — the flag is a no-op for them.

Usage

# Server
vllm-mlx serve mlx-community/Qwen3-30B-A3B-4bit \
  --continuous-batching \
  --moe-top-k 4

# Bench
vllm-mlx bench mlx-community/Qwen3-30B-A3B-4bit --moe-top-k 4

The flag is rejected if N is greater than the model's trained top_k (it only makes sense to lower, never to raise).

Measured impact

Decode throughput (M4 Max 128 GB, batch=1, greedy)

top_k tok/s vs baseline
8 (baseline) 126.5 —
6 136.1 +7.6%
5 140.3 +10.9%
4 147.3 +16.5%

Quality (Qwen3-30B-A3B-4bit, lm-evaluation-harness, MLX backend)

top_k MMLU (acc) GSM8K (exact match) Δ vs baseline
8 TBD TBD —
6 TBD TBD TBD
5 TBD TBD TBD
4 TBD TBD TBD

MMLU: 200 randomly-selected samples, 0-shot. GSM8K: 100 randomly-selected samples, 0-shot, exact-match strict.

These numbers are directional — full suites are larger and would shift the absolute accuracy but not the relative delta between configs by much.

Greedy output parity

With top_k=4 on the 4-bit checkpoint we observed identical first 16 generated tokens vs the baseline across every probe prompt we tried. This suggests top_k=4 does not change the argmax in the early decode steps — the model is internally robust to dropping half its activated experts.

At top_k=3 or lower quality would start to degrade visibly (not measured here; inferred from LExI paper), so the flag is intentionally not lowered below 1 at the config validation layer but the recommended floor for production is top_k=4.

When to use it, when not to

Use it when: - You run a Qwen3 MoE (or compatible: Qwen3.5 MoE, Gemma-MoE) and single-user decode throughput is your bottleneck. - You have a workload where a small quality drop is acceptable in exchange for a visible latency improvement. - You're deploying on memory-bandwidth-bound hardware (M-series Apple Silicon) where expert gather dominates per-step decode time.

Skip it when: - You serve dense models — flag is a no-op, adds nothing. - You care about top-1% leaderboard accuracy on eval suites. - You run long chain-of-thought / "thinking mode" generations where the quality cliff may be steeper than 0-shot MMLU suggests.

Stacking with other optimizations

This flag composes with quantization. On Qwen3-30B-A3B-4bit our measured stack is:

  • 4-bit + top_k=8: 126.5 tok/s (baseline)
  • 4-bit + top_k=4: 147.3 tok/s (+16.5%)
  • 3-bit + top_k=8: 138.6 tok/s (+9.6%)
  • 3-bit + top_k=6: 147.1 tok/s (+16.3%) — quality divergence measurable
  • 3-bit + top_k=4: 157.3 tok/s (+24%) — output quality breaks (model answered a different question in our smoke test)

3-bit + top_k=4 compounded the numerical error past the point where the argmax is stable. Stick to at most one aggressive knob: either 4-bit + top_k=4 or 3-bit + top_k=6. Both give approximately the same tok/s (~147) with very different quality profiles.

Internals

  • Patch helper: vllm_mlx.scheduler.apply_moe_top_k_override(model, k)
  • Applied in Scheduler.__init__ after the model is loaded.
  • Tests: tests/test_moe_top_k.py — covers dense models, mixed architectures, and validation paths.

References