AI Engineer Exhaustive
MediumPlusDeepSeek-V3

MoE Router

DeepSeek-V3

Medium

Top-k expert selection with gating

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Each DeepSeekMoE token must pick a handful of routed experts out of a large pool and assign them mixing weights. V3 keeps the DeepSeek-V2 fine-grained expert layout but changes the affinity: scores are sigmoid, not softmax over all N_r experts, and the gates are a renormalization of the selected scores only. Selection itself can be biased for load balance (separate note) and is further restricted to a few expert groups so a token travels to at most M nodes. On the 671B model that is $$ routed experts, K_r=8 activated, 8 groups, 4 groups kept.

How it works

For FFN input \mathbf{u}_t\in\mathbb{R}^{d} and expert centroid \mathbf{e}_i\in\mathbb{R}^{d} the paper defines

s_{i,t}=\operatorname{Sigmoid}(\mathbf{u}_t^{\top}\mathbf{e}_i).

Unbiased top-K_r (the form written before the bias section) zeros every non-winner:

g'_{i,t}=\begin{cases} s_{i,t}, & s_{i,t}\in\operatorname{Topk}(\{s_{j,t}\}_{j=1}^{N_r},K_r),\\ 0, & \text{otherwise,} \end{cases} \qquad g_{i,t}=\frac{g'_{i,t}}{\sum_{j=1}^{N_r}g'_{j,t}}.

Because the denominator sums only the K_r survivors, the gates form a distribution over the chosen experts, not over the whole pool. That is the V2→V3 change the paper calls out.

Load-balance bias (training) adds b_i before Topk but leaves s_{i,t} untouched when forming g'_{i,t}. See the load-balancing note; the router must implement both tensors.

Node-limited / grouped routing. Experts are partitioned into n_{\text{groups}} groups (8 on 671B, matching 8 nodes of 32 experts during the paper's EP layout). The paper picks at most $$ nodes by the sum of the highest K_r/M=2 affinity scores on each node. Official inference does the same in Gate: reshape scores to (tokens, n_groups, experts_per_group); if a bias exists, group_scores = topk(2).sum; otherwise amax. Keep n_limited_groups groups, mask the rest to -\infty, then take the token-level top-K_r.

After gather, sigmoid routers renormalize the K_r original scores and multiply route_scale. The 671B config sets route_scale to $$. That scale is in the official JSON, not written as a numbered equation in the report; treat it as a released hyper-parameter.

Shapes for a packed token batch T: scores (T,N_r), indices (T,K_r), weights (T,K_r). Shared experts are not in this router; they run on every token outside the gate.

Official code

Class Gate in inference/model.py. self.weight is (N_r,d); linear(x, weight) produces logits (no Linear bias). score_func is "sigmoid" on 671B ("softmax" remains as a branch for smaller ModelArgs defaults). Returns (weights, indices). Group counts and route_scale are in inference/configs/config_671B.json: n_routed_experts=256, n_activated_experts=8, n_expert_groups=8, n_limited_groups=4, score_func="sigmoid", route_scale=2.5. The paper's training-only \gamma bias update is not in this file.

Watch-outs

Sources