Implementing DeepSeekMoE

The previous five sections built up the DeepSeekMoE idea one piece at a time: why sparse conditional compute, how a router decides, why fine-grained experts outperform fat ones, why a small always-on shared block stabilizes the whole thing, and how the routed pool gets sharded across hundreds of GPUs. Now we wire it all together. This section is the assembly — one complete nn.Module that takes a batched transformer hidden state in, runs every part of DeepSeekMoE, and hands the right tensor back to the next transformer layer.

What you should walk away with. A mental model of the full DeepSeekMoE forward pass, the exact tensor shapes at every stage, a from-scratch NumPy reference you can step through with a debugger, and a production-style PyTorch block that scales — with the same code structure — from a toy on your laptop to the 671B-parameter DeepSeek-V3 on a real cluster.

From Five Ideas to One Block

Before code, a quick stocktake. A DeepSeekMoE block is built out of exactly five ingredients, and every one of them was earned in a previous section of this chapter:

The implementation is just these five things glued in the right order. Most of the engineering complexity sits in the glue: the right tensor shapes, the right flatten/reshape sandwich, the right place to put the softmax, and — for real clusters — the right all-to-all hook. The forward pass equation, by contrast, is astonishingly simple.

The Real Problem: Plumbing the Whole Forward Pass

Each previous section showed its piece in isolation. The router lived in a tiny NumPy snippet. Shared experts had their own toy. Fine-grained decomposition was a parameter-count argument. None of those snippets was a thing you could drop into a transformer and train. The implementation problem is exactly that: turn five mental models into one batched, autograd-traceable, shard- ready module.

Three things make this hard in practice. First, tensor shapes change rapidly: $(B, T, D)$ in, $(N, D)$ after flattening, $(N, E_r)$ after the router, $(N, k)$ after top-k, then per-expert mini-batches of shape $(M, D)$ where $M$ changes every step. Getting the indexing right is most of the bug surface. Second, the shared and routed paths combine by addition, not by a global softmax — a subtle but load-bearing detail that has been gotten wrong in more than one published reimplementation. Third, the same forward code has to work for a 1-GPU toy and a 256-GPU cluster — the only difference being where the all-to-all wrappers appear.

Intuition: A Tiny Scheduler Around Two Pools

Strip the autograd machinery away and a DeepSeekMoE block is a scheduler. For each token, it makes one decision: which routed experts get to see this token. The shared experts do not need a decision — they always run. So the forward loop is:

1. Always-on path. Run every shared expert on every token. Sum their outputs.
2. Scoring. The router computes one logit per routed expert. Cheap matmul, no FFNs touched yet.
3. Selection. Top-$k$ picks the winners. The other $N_r - k$ routed experts go to sleep for this token.
4. Gating. Softmax over the survivors gives convex gates that sum to 1.
5. Specialist path. Each surviving routed expert runs on the tokens that picked it. Multiply by the gate, add into the running output.
6. Return. The accumulated $y$ is the block's output. Pass to the next transformer layer as if nothing interesting happened.

The Complete DeepSeekMoE Forward Equation

Pull every piece together and the block computes a single tensor from a single tensor:

$y(x) = \underbrace{\sum_{i=1}^{N_s} E_s^{(i)}(x)}_{\text{always-on shared}} + \underbrace{\sum_{i \in \mathcal{T}(x)} g_i(x) \cdot E_r^{(i)}(x)}_{\text{gated routed}}$

The notation is the same we have used throughout the chapter, repeated here for completeness. $x \in \mathbb{R}^{d}$ is the input token hidden state. $E_s^{(i)}: \mathbb{R}^{d} \to \mathbb{R}^{d}$ and $E_r^{(i)}: \mathbb{R}^{d} \to \mathbb{R}^{d}$ are individual expert FFNs — typically $E(x) = W_2 \, \sigma(W_1 x)$ with $\sigma$ a GELU. The top-k set $\mathcal{T}(x) = \mathrm{TopK}(W_r x,\, k)$ contains the $k$ indices with the largest router scores. The gates $g_i(x) = \exp(s_i) / \sum_{j \in \mathcal{T}(x)} \exp(s_j)$ come from a softmax restricted to the survivors — $g_i = 0$ for $i \notin \mathcal{T}(x)$.

What the equation does not say

Three things are conspicuously absent from the formula, and each absence encodes an architectural decision:

1. No global softmax across shared and routed. The two sums are added independently. Shared experts are not in the gating competition; they are part of the architectural baseline.
2. No bias term on the router. DeepSeek's router uses a plain linear layer with no bias. The auxiliary-loss-free balancing in chapter 6 will add a per-expert bias to logits only for selection, not for the gates — a subtle distinction that matters at scale.
3. No capacity drop. The equation assumes every token gets its k winners. Real training uses a capacity factor that occasionally drops tokens when a single expert overflows; we covered that in 5.5.

FLOP accounting

A dense FFN of width $d_{ff}$ costs roughly $2 \cdot d \cdot d_{ff}$ FLOPs per token. A DeepSeekMoE block with $N_s$ shared experts, $N_r$ routed experts, and top-$k$ routing costs:

$\mathrm{FLOPs}(x) \approx (N_s + k) \cdot 2 d \cdot d_{ff} + N_r \cdot d$

The first term is the FFNs that actually run. The second is the router's matvec — negligible because $d \ll d_{ff}$. Compare against the $(N_s + N_r) \cdot 2 d \cdot d_{ff}$ parameters held in memory and you see the whole DeepSeekMoE bargain in one ratio:$\mathrm{active}/\mathrm{total} = (N_s + k) / (N_s + N_r)$. For DeepSeek-V3 that is $(1 + 8) / (1 + 256) \approx 3.5\%$ — and that is why a 671B-parameter model trains and serves at the cost of a ~37B-parameter dense model.

Manual Numerical Walkthrough

We push one toy token through a complete block with $d = 4$, $N_s = 2$, $N_r = 8$, $k = 2$. Every number below is computed by hand so you can verify any step against your own reimplementation.

Visualizing the Forward Pass

The tracer below walks the full block one stage at a time. Pick a token at the top right; the shared (violet) experts always fire, while the routed (green) pool only lights up the two experts the router selects for that domain. Use Play to auto-advance through the seven stages or step manually with the arrows.

Three things to watch as you switch tokens. First, the violet column is invariant — the same two shared experts fire for every token, with nearly the same output magnitude. That is the baseline doing its job. Second, the green grid changes its highlighted cells: the router picks different routed experts for code vs. math vs. history. The dimmed cells stayed in memory but consumed zero FLOPs. Third, the final output $y$ is just the sum of the two boxed accumulators — never a softmax, never a weighted average across the two pools.

Plain Python: A Complete DeepSeekMoE Block

Here is the entire block in NumPy. No autograd, no batching, no GPU — just one class with one forward method that you can step through with pdb. The shape of every tensor matches the equation we just wrote.

(8, 4)

For logits = [3.2, 0.4, 0.7, 2.6, 0.1, 1.4, 0.8, 0.9] with k=2, topk = [0, 3].

1import numpy as np
2
3class DeepSeekMoENumPy:
4    """One full DeepSeekMoE block in plain NumPy. No autograd, no shortcuts."""
5
6    def __init__(self, d, d_ff, n_shared, n_routed, k, seed=0):
7        self.d, self.d_ff = d, d_ff
8        self.n_shared, self.n_routed, self.k = n_shared, n_routed, k
9        rng = np.random.default_rng(seed)
10
11        # Two pools of FFNs, stored separately on purpose.
12        self.W1_s = rng.standard_normal((n_shared, d_ff, d)) * 0.05
13        self.W2_s = rng.standard_normal((n_shared, d, d_ff)) * 0.05
14        self.W1_r = rng.standard_normal((n_routed, d_ff, d)) * 0.05
15        self.W2_r = rng.standard_normal((n_routed, d, d_ff)) * 0.05
16
17        # Router maps d -> n_routed (NEVER n_shared + n_routed).
18        self.W_router = rng.standard_normal((n_routed, d)) * 0.05
19
20    def _ffn(self, W1, W2, x):
21        # Single-expert forward: GELU-ish (use ReLU here for clarity).
22        return W2 @ np.maximum(0, W1 @ x)
23
24    def forward(self, x):
25        # ---- (A) shared block: unconditional, no router ----
26        y = np.zeros(self.d)
27        for i in range(self.n_shared):
28            y += self._ffn(self.W1_s[i], self.W2_s[i], x)
29
30        # ---- (B) router scores ONLY the routed pool ----
31        logits = self.W_router @ x                  # (n_routed,)
32        topk = np.argsort(-logits)[: self.k]        # indices of k winners
33
34        # ---- (C) softmax restricted to the survivors ----
35        z = logits[topk] - logits[topk].max()
36        gates = np.exp(z) / np.exp(z).sum()         # (k,), sums to 1
37
38        # ---- (D) run routed experts and accumulate gated outputs ----
39        for g, i in zip(gates, topk):
40            y += g * self._ffn(self.W1_r[i], self.W2_r[i], x)
41
42        return y, {"chosen": topk, "gates": gates, "logits": logits}
43
44block = DeepSeekMoENumPy(d=4, d_ff=8, n_shared=2, n_routed=8, k=2)
45x = np.array([0.8, -0.3, 0.5, 0.2])
46y, info = block.forward(x)
47print("chosen:", info["chosen"], "gates:", info["gates"].round(3))
48print("y:", y.round(3))

Two structural observations before we move on. First, the router's output dim is $N_r$ on line 16 — not $N_s + N_r$. If you find yourself wanting to score the shared experts too, you are off the architecture. Second, the only difference between the shared loop (line 25) and the routed loop (line 37) is that the routed loop iterates over chosen indices with gates, while the shared loop iterates over every shared expert with no gate. That single line of asymmetry is the whole point of the shared-vs-routed split.

Sanity check via degenerate cases. Set $N_s = 0$: the block collapses to vanilla MoE (section 5.1). Set $k = 0$ and $N_s > 0$: the block becomes a stacked dense FFN. Set $N_r = 0$: the block is a dense FFN. DeepSeekMoE is the principled interpolation across all three corners.

PyTorch: Production-Style Implementation

Now the production-style PyTorch version. Two things change from the NumPy reference: weights are stored as packed $(E, d_{ff}, d)$ tensors so we can issue grouped GEMMs, and the routed loop uses the gather / run / scatter trick that keeps GPU matmuls large. Everything else is one-to-one with the math.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class DeepSeekMoEBlock(nn.Module):
6    def __init__(self, d_model: int, d_ff: int,
7                 n_shared: int, n_routed: int, k: int):
8        super().__init__()
9        self.d_model, self.k = d_model, k
10        self.n_shared, self.n_routed = n_shared, n_routed
11
12        # Stacked weights for batched matmul. (E, d_ff, d) and (E, d, d_ff).
13        self.W1_s = nn.Parameter(torch.empty(n_shared, d_ff, d_model))
14        self.W2_s = nn.Parameter(torch.empty(n_shared, d_model, d_ff))
15        self.W1_r = nn.Parameter(torch.empty(n_routed, d_ff, d_model))
16        self.W2_r = nn.Parameter(torch.empty(n_routed, d_model, d_ff))
17        for w in (self.W1_s, self.W2_s, self.W1_r, self.W2_r):
18            nn.init.kaiming_uniform_(w, a=5 ** 0.5)
19
20        self.router = nn.Linear(d_model, n_routed, bias=False)
21
22    def _run_pool(self, W1, W2, x):
23        # x: (N, d) — run every expert in the pool on every row.
24        # Returns (E, N, d) so the caller decides how to combine.
25        h = torch.einsum("efd,nd->enf", W1, x)            # (E, N, d_ff)
26        h = F.gelu(h)
27        return torch.einsum("edf,enf->end", W2, h)        # (E, N, d)
28
29    def forward(self, x: torch.Tensor):
30        B, T, D = x.shape
31        x_flat = x.reshape(B * T, D)                      # (N, D)
32
33        # (A) Shared path: stacked GEMM, then sum across the shared axis.
34        shared_out = self._run_pool(self.W1_s, self.W2_s, x_flat)  # (E_s, N, D)
35        y = shared_out.sum(dim=0)                                  # (N, D)
36
37        # (B) Router scores the routed pool.
38        logits = self.router(x_flat)                      # (N, E_r)
39        topk_w, topk_idx = logits.topk(self.k, dim=-1)    # (N, k), (N, k)
40        gates = F.softmax(topk_w, dim=-1)                 # (N, k)
41
42        # (C) Loop over routed experts, gather tokens, run, scatter.
43        for i in range(self.n_routed):
44            mask = (topk_idx == i)                        # (N, k)
45            if not mask.any():
46                continue
47            tok_rows, slot = mask.nonzero(as_tuple=True)
48            g = gates[tok_rows, slot].unsqueeze(-1)       # (M, 1)
49            # Single packed FFN call for all M tokens of expert i.
50            h = F.gelu(x_flat[tok_rows] @ self.W1_r[i].T) # (M, d_ff)
51            out = h @ self.W2_r[i].T                      # (M, D)
52            y[tok_rows] += g * out
53
54        return y.reshape(B, T, D), {
55            "logits": logits, "topk_idx": topk_idx, "gates": gates,
56        }

Three implementation details that real DeepSeek code does differently and you should know about:

1. Grouped GEMM instead of a Python loop. The for i in range(self.n_routed) loop is fine at 8 experts but a throughput killer at 256. Real implementations use torch._scaled_grouped_mm (or DeepGEMM, the kernel DeepSeek released) to issue all the routed FFNs as one batched call. The Python becomes ~10 lines shorter and ~3× faster.
2. Permute, then call, then unpermute. A common pattern in high-performance MoE code: sort the token batch so all tokens for expert 0 come first, then all for expert 1, then call a single grouped GEMM with a length-per-expert array, then unpermute back. This avoids any per-expert indexing and lets the kernel use contiguous memory.
3. All-to-all wraps the routed loop. On a multi-GPU cluster, the permute step also dispatches tokens to the GPU that owns each routed expert. The forward pass becomes: shared FFN locally → permute + all-to-all → routed grouped GEMM on remote GPU → all-to-all back + unpermute → add to shared output. The shared path runs entirely on-device with no comms — section 5.5 showed why.

What Changes at Massive Scale

Take the same DeepSeekMoEBlock above, change three integers, and you have a layer of DeepSeek-V3:

Four operational concerns emerge once the integers grow.

Memory: optimizer states dominate

At a 4B-parameter block, the weights themselves are 8 GB in bf16. AdamW adds moment buffers — first and second moments, both in fp32 for stability — and you are at 24 GB per block before activations. FSDP / ZeRO shards the optimizer states across data-parallel replicas; expert parallelism additionally shards the routed weights themselves across $E_p$ GPUs. The shared block stays replicated everywhere — it is small enough that the cross-GPU traffic to shard it would cost more than the memory it saves.

Communication: all-to-all is the new attention

Every routed-expert pool larger than one node lives behind an all-to-all collective. Each token has to travel to whichever GPU owns its $k$ chosen experts, run, and return. With 256 experts sharded across 32 GPUs, each forward pass moves on the order of $N \cdot k \cdot D$ bytes per direction. Bandwidth is the bottleneck — not FLOPs. DeepSeek uses InfiniBand-tuned NCCL kernels and the DeepEP library to overlap the all-to-all with the shared-block compute, hiding most of the latency behind work the shared experts were going to do anyway.

Numerical precision: fp8 wherever you can

DeepSeek-V3 runs the routed expert GEMMs in fp8 — the router stays in bf16, the softmax stays in fp32, the FFN matmuls drop to fp8. Per-tile scaling and an fp32 accumulate keep the fp8 numerics stable. The mixed-precision sandwich saves ~40% of memory traffic on the dominant cost (routed FFNs) without measurable loss damage. Shared experts can stay in bf16 — they handle every token, so any fp8 noise compounds faster there.

Load balance: the elephant in the next chapter

The forward pass above assumes the router picks roughly equal numbers of tokens per expert. It does not. Without intervention, training collapses into routing one or two experts hard while the rest starve. The classical fix — an auxiliary load-balance loss — actively hurts model quality at scale. DeepSeek invented the auxiliary-loss-free bias-term trick that fixes the imbalance without contaminating the gradient signal, and that is the entire subject of chapter 6. The block we just built is the right shape for that fix to slot into; you will see exactly where the bias term lives when we get there.

Engineering Reality and Gotchas

Three failure modes you will absolutely hit if you reimplement DeepSeekMoE without reading the next chapter first:

1. Routing collapse on warmup. Cold-started routers produce near-uniform logits, so top-k picks effectively random experts for the first few hundred steps. The lucky experts get more gradient, become slightly better targets, get even more tokens — a feedback loop that ends with two experts owning the entire batch and the other 254 idle. Chapter 6 covers the bias-term fix; in the meantime, kaiming-uniform on the router (not zero init) and a short LR warmup buy you stability.
2. Activation checkpointing recomputes routing. If you activation-checkpoint the MoE block (which you almost certainly will, to fit training), the backward pass re-runs the router and top-k. That is fine mathematically — the same logits produce the same top-k — but you must seed any stochastic routing identically on the recompute, and you must avoid any non-determinism (e.g., atomic gather/scatter) that would change which experts claim which tokens.
3. Mixed precision around softmax. Run softmax in fp32 even when everything around it is bf16. The exp() is highly sensitive to small logit shifts — fp16 softmax loses gates entirely when one logit dominates. The cost is negligible; the bug it prevents is silent and devastating.

The one sentence to carry forward: DeepSeekMoE is a sum of two pools, gated only on one side. Every other piece of complexity — fine-grained decomposition, expert parallelism, fp8 GEMMs, bias-term balancing — is an optimization layered on top of that single additive equation. Once that equation is in your fingers, the rest of the chapter (and the rest of the book) is engineering, not architecture.