Sort your prompts: a free speed-up for MoE batch inference

TL;DR. In a mixture-of-experts model, each decode step has to read the weights of every expert that any request in the batch routes to. If you run a big offline job (evals, synthetic data, labelling) over prompts from several domains, mixing them in each batch makes the model read more experts than it needs to. Just sorting the prompts before calling llm.generate, so that similar requests end up decoding together, made unmodified vLLM +12–17% faster on Qwen3-30B-A3B and +7–12% faster on gpt-oss-120b (one AMD MI300X). On unlabeled chat traffic, grouping by a 0.6B embedding model still gives +5% on Qwen and up to +4% on gpt-oss.". No model change, no kernel change, no custom scheduler: one argsort.

Real recorded routing of Qwen3-30B-A3B. Each cell is one (layer, expert) pair; lit = its weights are read at this decode step, for a batch of 32 requests. Left: random order. Right: the same requests sorted by domain.

Why batch composition matters for MoEs

During decoding, a GPU spends most of its time reading weights, not computing. For a dense model the bytes read per step are fixed: all weights, once. For an MoE it depends on the batch. Every expert chosen by at least one token in the batch must be streamed from HBM once per layer; experts nobody picked cost nothing.

Qwen3-30B-A3B has 48 layers × 128 experts, and each token picks 8. A single request touches 8 experts per layer. A batch of 32 random requests touches about 94 of 128, because their choices barely overlap. But expert choice is not random: code tokens prefer some experts, French tokens others. Put 32 code requests together and they overlap a lot more.

Each domain uses its own subset of experts

We counted which experts each domain's generated tokens pick in one layer. Each domain concentrates on its own group: its 32 favourite experts (a quarter of the layer) receive 68–80% of all its routing decisions, against 25% if routing were uniform. Code, math, prose, French, Arabic and Chinese each light up a different corner.

Decode time follows the number of experts touched

We measured that link directly. With a vLLM-style fused-MoE Triton kernel on real routing, the MoE layer time is a straight line in the number of distinct experts the batch touches: about 2.9 µs per extra expert, which is just 9.4 MB of expert weights at ~3.3 TB/s.

How many experts does a batch touch?

We recorded the routing of every prompt and generated token for 768 prompts across 6 domains: code (MBPP), math (GSM8K), creative writing (no_robots), and French, Arabic and Chinese (Aya). Routing is computed per token, so we can replay any batch composition exactly, offline.

• With 128 fine-grained experts (Qwen3-30B-A3B), one-domain batches read 19–25% fewer experts per step than mixed batches, at every batch size from 8 to 128.
• With only 32 coarse experts (gpt-oss-20b), the gap is 8–11%: there is less room to specialise.
• The savings grow with depth. Later layers are the most domain-specialised.

Does it survive end to end?

Fewer bytes in the MoE layers don't automatically mean faster generation: attention, the KV cache, sampling and scheduling don't care about expert locality. So we measured in stock vLLM 0.27 (ROCm, CUDA graphs, its own fused-MoE kernels). Setup:

• 768 prompts, each generating exactly 256 tokens (ignore_eos=True), greedy decoding.
• The two runs do identical work. The only difference is the order of the prompt list.
• vLLM's offline engine is first-come first-served, so the order decides which requests share a decode batch.

outs = llm.generate(shuffled_prompts, sp)        # baseline
outs = llm.generate(sorted_by_domain_prompts, sp)  # same prompts, sorted

Two GPUs, expert parallel. With the experts split across both MI300X (tensor_parallel_size=2, enable_expert_parallel=True), sorting still gives +12.9% at batch 64 (3,440 → 3,885 tok/s) and +8.7% at batch 128 (6,835 → 7,431 tok/s). With experts sharded, each GPU only waits for its own distinct experts, and grouping lowers that count on the busiest GPU too (~22% fewer in replay).

Each cell is 3 alternating repeats. Run-to-run spread is below 1% for Qwen and up to 4% for gpt-oss-120b. Every sorted run beat every shuffled run, except for gpt-oss on UltraChat at batch 128, where the gain is within noise. For Qwen, the gains are about half of what the bytes-read model predicts (+20–27%). That is expected, because only the MoE part of each step gets cheaper.

Outputs. Reordering doesn't change what each request computes, but it changes which requests share a batch, and in bf16 that can flip a near-tie in greedy decoding. That holds for any change of batch composition. To put a number on it, we compared the sorted run against a shuffled run, and one shuffled run against a different shuffled run: sorting changes outputs about as much as reshuffling does.

These differences come from ordinary batch nondeterminism: none of these requests computes anything different. If you need bit-exact outputs, you need batch-invariant kernels, whatever order you use.

No labels? Use the router, or a tiny embedding model

Real traffic doesn't come with domain labels. There are two label-free ways to group it.

1. Router fingerprint. After prefill, every prompt has a vector of how often it used each expert in each layer. k-means on these vectors recovers our 6 domains with 100% purity, and it gives the same savings as the true labels. It needs access to the router, though.
2. Text embedding. Embed each prompt with Qwen3-Embedding-0.6B and run k-means. This needs nothing from the MoE, runs before inference, and works for any model.

On the multi-domain workload, the embedding clusters are as good as the labels up to batch 32. On UltraChat, where no domains exist, clustering still cuts experts read by 7–12%. End to end in vLLM: grouping UltraChat with the embedding model gives +5.1-5.5% throughput on Qwen3-30B-A3B at every batch size, and +2–4% on gpt-oss-120b (nothing measurable at batch 128). On the 6-domain set, embedding clusters give +13.0% at batch 64, slightly above the true labels (+12.6%).

The whole recipe:

from sort_prompts import expert_friendly_order

order = expert_friendly_order(prompts, n_clusters=16)   # embed + k-means + stable argsort
outs = llm.generate([prompts[i] for i in order], sampling_params)
outs = [o for _, o in sorted(zip(order, outs))]          # restore the original order

What if I can't sort the whole dataset?

A streaming job or a server can't see all requests in advance. It can, however, reorder within a window of the next W requests, which bounds how long any request waits.

Measured in vLLM at batch 32, sorting every chunk of 128 prompts gives +11.7%, every chunk of 256 gives +14.8%, and sorting the full 768 gives +17.3%. A reorder window of 4–8 batches keeps most of the gain.

Natural output lengths

With EOS enabled, sorting also groups requests of similar length (math answers are short, stories are long), which changes how the batch drains. With natural lengths (EOS on, up to 512 tokens, batch 32), the sorted job finishes 16.6% sooner in wall-clock time: 1,810 → 2,104 output tok/s. The gain survives the realistic setting.

Where this doesn't help

• Dense models. Every weight is read at every step anyway.
• Prefill-heavy work. Thousands of prompt tokens touch nearly every expert regardless of order. This is a decode effect.
• Coarse MoEs. With few experts (gpt-oss-20b, 32 experts), there's less to gain.
• Very large batches. Once even a one-domain batch touches almost every expert, there's nothing left to save, and the step becomes compute-bound.
• Latency-sensitive online serving. Holding requests back to group them costs queueing time. The natural home is offline batch jobs, or a scheduler that picks similar requests from its existing waiting queue.

Related work

The underlying observation, that decode cost tracks the number of distinct experts and that similar requests share experts, is not new, and others have built systems on it:

• ELDR (Choi et al., 2026) routes requests to decode workers by a prefill "expert signature" in disaggregated serving. They report 17–21% fewer activated experts for same-domain batches and 5.9–13.9% lower TPOT in vLLM, which agrees with our numbers.
• Bambhaniya et al., 2026 groups micro-batches by expert activation to cut multi-node all-to-all traffic.
• Sem-MoE schedules requests onto the expert-parallel ranks that host their likely experts.
• ExpertFlow groups tokens by predicted routes in an offloading setting.
• Read-ME uses a decoupled router to do expert-aware batching, but needs a model change.
• Lynx, Opportunistic Expert Activation and XShare reduce distinct experts per batch by changing routing. That is complementary, but it changes outputs.

Our point is the minimal version of this: no new system, no predictor, no model change. One sort before an offline job, measured end to end in unmodified vLLM.

Reproduce

Everything (routing capture, replay simulator, Triton kernel, vLLM benchmark, figures) is in the repo. Hardware: AMD Instinct MI300X; software: vLLM 0.27.1, ROCm 10, PyTorch 2.12.