AI Engineer Exhaustive
HardPlusDeepSeek-V3

Sparse MoE FFN

DeepSeek-V3

Hard

Full sparse MoE feedforward block

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

After the first three dense layers, every DeepSeek-V3 feed-forward is a DeepSeekMoE block: one always-on shared SwiGLU plus 256 routed SwiGLUs of which 8 fire per token. That is the compute that makes 671B total parameters collapse to about 37B activated (including attention and embeddings). This note is the assembly of pieces covered elsewhere — sigmoid top-k routing, auxiliary-loss-free bias, shared expert — into one residual FFN. The paper's design goals here are fine-grained experts (many small FFNs instead of few wide ones), isolation of shared knowledge, balanced load without token dropping, and node-limited communication so DualPipe can hide all-to-all.

How it works

Residual form (paper eq. 12). With \mathbf{u}_t the pre-normed FFN input:

\mathbf{h}_t'=\mathbf{u}_t+\sum_{i=1}^{N_s}\operatorname{FFN}_i^{(s)}(\mathbf{u}_t)+\sum_{i=1}^{N_r}g_{i,t}\,\operatorname{FFN}_i^{(r)}(\mathbf{u}_t).

On 671B: N_s=1, N_r=256, K_r=8, each expert inner width 2048, hidden d=7168. Official Block adds the \mathbf{u}_t skip outside MoE; the module returns shared + routed only.

Each expert is SwiGLU, same algebra as the dense prefix MLP:

\operatorname{FFN}(\mathbf{u})=\mathrm{W}_2\bigl(\operatorname{SiLU}(\mathrm{W}_1\mathbf{u})\odot\mathrm{W}_3\mathbf{u}\bigr),\quad \mathrm{W}_1,\mathrm{W}_3\in\mathbb{R}^{2048\times 7168},\;\mathrm{W}_2\in\mathbb{R}^{7168\times 2048}.

Routed copies are local Linears (expert-parallel shards). The shared path is one MLP of inner width N_s\cdot 2048.

Routing. Affinities s_{i,t}=\operatorname{Sigmoid}(\mathbf{u}_t^{\top}\mathbf{e}_i). Top-K_r is taken on s_{i,t}+b_i (and after group masking: 8 groups, keep 4, group score = sum of top-2 when bias is present). Gates $$ are the unbiased selected s_{i,t} renormalized over the K_r winners, then multiplied by route_scale=2.5. Tokens that miss a local expert contribute 0 to that shard; shards all_reduce the routed sum.

Packed forward (official). Flatten (B,S,d)\to(T,d). Gate returns weights (T,K_r) and indices (T,K_r). For each local expert i, gather the tokens where indices==i, run the SwiGLU, scale by the matching weights, and scatter-add into a zero buffer $$. Compute z=\operatorname{FFN}^{(s)}(\mathbf{u}) on all T tokens. Return (y+z) reshaped. No token is dropped: empty experts are skipped, but every token still receives shared + its K_r routed outputs.

What is specified vs. inferred. Equations (12)–(16), N_s,N_r,K_r,M=4, and “no token-dropping” are in the report. Expert-parallel index ranges, fused shared width, and route_scale come from inference/model.py and config_671B.json. The \gamma bias update and \mathcal{L}_{\mathrm{Bal}} are training-only and absent from the inference file.

Activated FFN FLOPs per token are one shared SwiGLU plus eight routed SwiGLUs, not 257. Attention is MLA and is outside this block.

Official code

inference/model.py: Block.ffn is MLP when layer_id < n_dense_layers (3) else MoE. MoE owns Gate, an ModuleList of Expert (or None for non-local ids), and shared_experts. Expert.forward is the SwiGLU above. 671B numbers: inference/configs/config_671B.json. Distributed ranks set experts_start_idx = rank * (N_r / world_size).

Watch-outs

Sources