Mixture of Experts in vLLM - From First Principles to fused_moe

Hello all. This post explains one mixture-of-experts (MoE) layer from first principles, and then shows how vLLM runs it. The LLM inference systems post has a short MoE section about memory capacity and expert parallelism. This post goes inside one layer on one GPU. It answers four questions:

• What is an expert, and what does one token do with the experts?
• What is a (token, expert) pair, and why does a batch need pairs?
• How do the pairs become GEMMs, and what are the GEMM shapes?
• How does vLLM’s fused_moe code run all of this without copying the input?

The running example is Qwen3-30B-A3B. Every shape in the post comes from its config.json. The vLLM code is quoted from commit e90158d912 on main, as of 2026-10-13. vLLM changes quickly, so line numbers can move.

The post has two parts. The first four sections are the math of the layer, with no kernel details. The fifth section shows how vLLM runs the same math. If you know MoE already, start at the vLLM section.

The plan:

• One expert, and one token
• A batch: (token, expert) pairs
• Group the pairs by expert, and the GEMM shapes
• Put the results back, and combine them per token
• How vLLM runs the layer: two index lists and no copies

Let’s get started.

One expert, and one token

An expert is a small SwiGLU MLP

A dense transformer layer has one feed-forward network (FFN) after attention. An MoE layer replaces this one FFN with many smaller FFNs. Each small FFN is an expert. The idea of a sparsely gated MoE layer comes from Shazeer et al., 2017.

In Qwen3-30B-A3B, each expert is a SwiGLU MLP (Shazeer, 2020). The config gives these sizes:

One expert does three steps on one token vector \(x\) of size 2048:

1. Multiply \(x\) by the gate weights and by the up weights. Each matrix is \([768, 2048]\). vLLM stores the two matrices as one matrix, W13, with shape \([1536, 2048]\). So one GEMM gives both halves, \(h\) of size 1536.
2. Compute \(a = \text{SiLU}(h_{\text{gate}}) \odot h_{\text{up}}\). Here \(h_{\text{gate}}\) is the first 768 values of \(h\) and \(h_{\text{up}}\) is the last 768 values. This step has no weights.
3. Multiply \(a\) by the down weights W2, with shape \([2048, 768]\). The result \(y\) has size 2048 again.

The weight count of one expert and of one layer:

One token uses only 8 experts

A router decides which experts a token uses. The router is one small linear layer, \(W_r\), with shape \([128, 2048]\). It gives 128 scores, one for each expert. Then the layer does these steps:

1. Apply softmax to the 128 scores.
2. Keep the 8 largest. The result is 8 expert ids and 8 weights for the token.
3. Rescale the 8 weights so that they add up to 1. The config sets norm_topk_prob: true, and vLLM passes it as renormalize (qwen3_moe.py, line 210).
4. Run the 8 chosen experts on the token.
5. Add the 8 expert outputs, each multiplied by its weight.

Here \(\text{id}_j\) and \(w_j\) are the expert id and the weight in slot \(j\) of the token, and \(f_e\) is expert \(e\). The other 120 experts do no work for this token.

This is the purpose of the MoE design. Each token computes with 8 experts, but the layer stores 128. The model card states 30.5B parameters in total and 3.3B activated. I computed both numbers from the config: 30.53B in total and 3.35B active, with the embedding and the output head included. The difference between the two numbers is almost all expert weights:

References

• Qwen3-30B-A3B model card and config.json
• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, 2017
• Shazeer, GLU Variants Improve Transformer, 2020
• vLLM, vllm/model_executor/models/qwen3_moe.py (the router GateLinear at line 170, renormalize at line 210)

A batch: (token, expert) pairs

A GPU does not run one token at a time. It runs a batch of \(M\) tokens. The router gives each token its own 8 experts, so different tokens choose different experts. The router output for the batch is two tensors:

• topk_ids, with shape \([M, 8]\): the 8 expert ids of each token.
• topk_weights, with shape \([M, 8]\): the 8 weights of each token.

Each cell of topk_ids is one job: “run expert \(e\) on token \(t\)”. From now on, this post calls one cell a pair. A batch has \(M \times 8\) pairs.

Each pair has a number, the pair id \(p\). The pair id is the position of the cell in topk_ids, read row by row:

Here \(t\) is the token number and \(j\) is the slot in the token’s row, from 0 to top_k minus 1.

The rest of the post uses a small example, so that every pair fits in a figure. The example has \(M = 3\) tokens, 4 experts and top_k = 2. The expert choices are made up. The real layer has 128 experts and top_k = 8.

The pairs are in token order, so neighboring pairs often use different experts. One GEMM multiplies all of its rows by one weight matrix. Each expert has a different weight matrix. So the layer cannot run the pairs in token order as one GEMM. It must first sort the pairs by expert.

Group the pairs by expert, and the GEMM shapes

Group by expert

For each expert \(e\), collect the pairs that use \(e\). Then stack the input rows of those tokens into one matrix \(X_e\). The number of rows of \(X_e\) is \(n_e\), the number of pairs for expert \(e\).

Three facts follow from the grouping:

• The same token can appear in up to top_k groups. Its input row is the same in each group.
• The row counts add up to the number of pairs: \(\sum_e n_e = M \times \text{top\_k}\).
• The row counts are not equal, because they depend on the router’s choices. An expert can get many rows, a few rows or none.

For the real layer, \(\sum_e n_e = 8M\) and \(E = 128\). So the mean row count is \(8M / 128 = M / 16\):

These are means. The router does not spread tokens evenly, so the real counts vary around the mean.

Each expert is now an ordinary GEMM

After the grouping, each expert runs the same three steps as in the first section, but on \(n_e\) rows at a time:

In GEMM terms, with \(M_{\text{gemm}}\) rows, inner dimension \(K_{\text{gemm}}\) and \(N_{\text{gemm}}\) output columns:

So each GEMM step of the layer is 128 GEMMs, one for each expert. The GEMMs have different row counts, but all have the same \(K_{\text{gemm}}\) and \(N_{\text{gemm}}\). A kernel that runs many such GEMMs in one launch is a grouped GEMM. If you stack the outputs of all experts, GEMM 1 gives \(8M\) rows of 1536 values. This is the reason for the factor of 8 in the intermediate tensors later in the post.

The mean row counts in the first table of this section also show a cost of MoE in decode. If \(n_e\) is about 1, then each expert reads all 4,718,592 of its weights from memory to make about one output row. In prefill, the same weight read serves about 256 rows.

Put the results back, and combine them per token

After GEMM 2, the results are in expert order: \(Y_e\) holds the \(n_e\) output rows of expert \(e\), one for each of its pairs. But the layer must return one vector per token, in token order. Two steps do this:

1. Send each output row back to the slot of its pair. The slots form a tensor of shape \([M, \text{top\_k}, 2048]\). Row \(t\) of the slots holds the top_k results of token \(t\).
2. For each token, multiply each result by its router weight and add the top_k results.

The result is \([M, 2048]\), the same shape as the input \(x\). This is the sum from the first section, \(y = \sum_j w_j f_{\text{id}_j}(x)\), done for every token of the batch at once.

So the math of the layer has three steps: group the pairs by expert, run the GEMMs, and put the results back per token. The next section shows how vLLM does the grouping and the putting back without moving any input rows.

How vLLM runs the layer: two index lists and no copies

This section follows the Triton path in fused_moe.py. The function fused_experts_impl (line 1652) runs one MoE layer. For bf16 and FP8 weights, it calls one Triton kernel, fused_moe_kernel (line 298), for both GEMMs. vLLM also has other MoE backends, for example CUTLASS and DeepGEMM kernels. This post does not cover them.

Why vLLM does not build X_e

The simple way to group by expert is to copy rows. Copy every pair’s input row into a new buffer in expert order, run the GEMMs, and copy the outputs back. Each copy writes \(M \times 8\) rows and the next step reads them again. vLLM avoids both copies. It builds two small index lists. The GEMM kernel uses them as addresses. So it reads each input row from its original place, and it writes each output row to its final place.

The two index lists

The function moe_align_block_size builds the index lists from topk_ids:

• sorted_token_ids: the pair ids, sorted by expert. Each expert’s group is padded up to a multiple of BLOCK_SIZE_M, the number of rows that one GEMM tile handles. A pad entry holds the id \(M \times \text{top\_k}\), which is not a real pair id. So the kernel can skip it.
• expert_ids: one expert id for each block of BLOCK_SIZE_M entries in sorted_token_ids.

The padding makes every block belong to exactly one expert. Thus each GEMM tile multiplies its rows by one weight matrix, which is what a GEMM needs.

The figure below uses two short names. The GEMM 1 input is qhidden_states in the code and qhidden in the figures. It has one row per token. If the model uses FP8, then this tensor is the quantized copy of \(x\). For a bf16 model, it is \(x\) itself. The GEMM 1 output is intermediate_cache1 in the code and cache1 in this post. It has one row per pair.

BLOCK_SIZE_M is 4 in the example so that the figure stays small. In vLLM, try_get_optimal_moe_config (line 1419) gives the value. It loads a tuned config file for the number of experts, the expert size and the weight dtype. From this file, it takes the entry with the batch size nearest to \(M\). If no file exists, then it uses a default config.

The order of the pairs inside one expert’s group is not fixed, because the align kernel fills each group in parallel. The figure shows one possible order.

One GEMM 1 program, step by step

The Triton grid has one program for each combination of a block and a tile of output columns. Take the program for block 0 in the figure:

1. It reads expert_ids[0], which is e0 (line 421). So its weight operand is W13 of expert e0.
2. It reads its 4 entries of sorted_token_ids: p1, p2 and two pads.
3. It masks the pads. The mask is offs_token < num_valid_tokens (line 419), and num_valid_tokens is \(M \times \text{top\_k} = 6\). A pad holds the id 6, so the mask removes it.
4. For each real pair, it loads the input row p // top_k (line 471). This is the integer division of the pair id by top_k.
5. It multiplies the rows by the weight tiles of e0. All rows of the block share the same weight tiles.
6. It writes each result to row p of the output (line 606).

Step 4 is the grouping of Group the pairs by expert, done as address arithmetic. The input qhidden has one row per token, not one row per pair. So the program must convert a pair id back to its token number. The pair id formula is \(p = t \times \text{top\_k} + j\) with \(0 \le j < \text{top\_k}\). So integer division by top_k gives \(t\) back:

These are the two tokens that chose e0, the same rows as \(X_{e0}\) in the grouping figure. The program reads them in place. Token t0 is also in the e2 block, as pair p0. The e2 program also computes 0 // 2 = 0, so it reads the same row 0 of qhidden. The row exists once in memory, and two programs read it.

Step 6 is the putting back of Put the results back. The output cache1 has one row per pair, because each pair gives a different result. Token t0 has two results, one from e2 and one from e0, so they cannot share a row. The program writes each result directly to row \(p\). So the output is in pair order at once, with no copy.

Pair order gives the per-token view with no copy

Row \(p = t \times \text{top\_k} + j\), so the top_k rows of one token are next to each other in memory. For this reason, vLLM can view the same buffer as \([M, \text{top\_k}, N]\) with no copy (line 1742):

intermediate_cache1 = cache13[: M * top_k_num * N].view(M, top_k_num, N)

In this view, row \(t\) holds the top_k results of token \(t\) side by side. These are the pair slots of Put the results back.

The rest of the layer stays in pair order

After GEMM 1, every step works on rows in pair order:

• SiLU and multiply (line 1811) works on each row by itself. It needs no expert information. It turns cache1, \([M \times 8, 1536]\), into cache2, \([M \times 8, 768]\). Row \(p\) stays row \(p\).
• GEMM 2 uses the same sorted_token_ids and expert_ids, so each block again uses one expert’s weights, now W2. Its input cache2 already has one row per pair. So vLLM calls the same kernel with top_k = 1, and the read address p // 1 is just p. The call passes not apply_router_weight_on_input as mul_routed_weight (line 1837). If apply_router_weight_on_input is false, which is the default, then the kernel multiplies each output row by topk_weights[p] (line 592). It writes row \(p\) of cache3, \([M \times 8, 2048]\).
• moe_sum (line 1850) views cache3 as \([M, 8, 2048]\) and adds the 8 rows of each token. The output is \([M, 2048]\) in token order.

So the expert grouping exists only in the two index lists. Only two places change the row order:

One more detail in the same function: intermediate_cache1 and intermediate_cache3 are views of one buffer, cache13 (lines 1742 and 1743). GEMM 2 can write over cache1, because SiLU and multiply has already read it into cache2. This reuse saves one large allocation.

Small batches skip the sort

For very small batches, vLLM does not call moe_align_block_size. The function _prepare_expert_assignment makes this decision (lines 1559 to 1567). If expert parallelism is off and the weights are not group-quantized int8 or int4, then it checks this condition:

For Qwen3-30B-A3B, the condition is \(8M \times 4 \le 128\), so \(M \le 4\). If the condition is true, then sorted_token_ids is None and expert_ids is topk_ids flattened. The kernel then puts each pair in its own block: block \(i\) holds only pair \(i\), and its expert is topk_ids at position \(i\). A block of BLOCK_SIZE_M rows then has one real row, and the kernel masks the other rows.

A code comment calls the factor 4 a heuristic margin, so that the pairs use only a small fraction of the experts. The code does not give more detail. The next two sentences are my reading of the reason, not a statement from the code. If most experts get at most one pair, then a sorted block also holds only one real row. So the sort gives almost no gain.

References

• vLLM, vllm/model_executor/layers/fused_moe/fused_moe.py: fused_moe_kernel (line 298), the pad mask (line 419), the expert of a block (line 421), the input row p // top_k (line 471), the router weight (line 592), the output row p (line 606), _prepare_expert_assignment (line 1541), fused_experts_impl (line 1652), the cache views (lines 1742 and 1743), the activation (line 1811), the GEMM 2 call (line 1837) and moe_sum (line 1850)
• vLLM, vllm/model_executor/layers/fused_moe/moe_align_block_size.py (the docstring of moe_align_block_size)

Wrapping up

An MoE layer replaces one FFN with many small expert FFNs and a router that picks top_k experts for each token. A batch of \(M\) tokens gives \(M \times \text{top\_k}\) (token, expert) pairs. The math groups the pairs by expert, runs one GEMM per expert with \(n_e\) rows, and adds each token’s top_k results. vLLM does the grouping and the putting back with two index lists, sorted_token_ids and expert_ids. GEMM 1 reads row p // top_k of the input and writes row p. So the tensors between GEMM 1 and moe_sum are in pair order, and no step copies the input.

The same structure explains a cost of MoE in decode. A small batch gives about one row per expert. So each expert reads all of its weights from memory to make about one row of output.

If you find a mistake anywhere in here, please let me know and I’ll fix it.