MoE routing from scratch

Mixture-of-Experts is one of those ideas that sounds too good to be true: have a model with 8x the parameters but only activate a fraction of them per token, getting the capacity of a huge model with the compute cost of a small one. Mixtral 8x7B has 46.7B total parameters but only activates 12.9B per token (2 of 8 experts per layer). DeepSeek-V3 has 671B parameters but activates around 37B per token.

The magic that makes this work is the router, a small network that decides which experts process each token. Today I am building one from scratch to understand the design decisions and failure modes.

The basic top-k router

The simplest MoE router is a linear layer followed by a top-k selection. For each token, it produces a score for every expert, picks the top k, and combines the experts' outputs weighted by the softmax of the selected scores.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TopKRouter(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)
        self.top_k = top_k
        self.num_experts = num_experts

    def forward(self, x):
        # x shape: (batch * seq_len, hidden_size)
        logits = self.gate(x)  # (batch * seq_len, num_experts)

        # Select top-k experts per token
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)

        # Softmax over selected experts only
        top_k_weights = F.softmax(top_k_logits, dim=-1)

        return top_k_weights, top_k_indices

This is essentially what Mixtral uses. The gate is a single linear layer with shape (hidden_size, num_experts), which for Mixtral is (4096, 8). That is just 32,768 parameters per router, trivially small compared to the expert FFNs.

The load balancing problem

The naive router has a critical failure mode: expert collapse. Without any balancing signal, the router tends to converge on sending most tokens to a small subset of experts. Some experts become heavily loaded while others sit idle. This wastes memory (you are storing expert weights that never get used) and hurts model quality (the model is not using its full capacity).

During training, this is solved with an auxiliary load balancing loss. The Switch Transformer paper introduced this formulation:

def load_balancing_loss(router_logits, top_k_indices, num_experts):
    """Auxiliary loss to encourage uniform expert utilization."""
    # f_i: fraction of tokens routed to expert i
    # P_i: average router probability for expert i

    num_tokens = router_logits.shape[0]
    probs = F.softmax(router_logits, dim=-1)  # (num_tokens, num_experts)

    # One-hot encode which experts were selected
    expert_mask = F.one_hot(top_k_indices, num_experts).sum(dim=1)  # (num_tokens, num_experts)

    # f_i: fraction of tokens assigned to each expert
    f = expert_mask.float().mean(dim=0)  # (num_experts,)

    # P_i: mean router probability for each expert
    P = probs.mean(dim=0)  # (num_experts,)

    # Loss: sum of f_i * P_i (minimized when both are uniform = 1/num_experts)
    aux_loss = num_experts * (f * P).sum()
    return aux_loss

The loss is minimized when every expert receives an equal fraction of tokens (f_i = 1/N) and has equal average probability (P_i = 1/N). It is scaled by a coefficient (typically 0.01 to 0.1) and added to the main language modeling loss.

Expert capacity and token dropping

During inference, load imbalance has a direct performance consequence. If you process all experts in parallel (as you would on multiple GPUs), the overall latency is determined by the most loaded expert. If expert 3 gets 30% of tokens while expert 7 gets 5%, you are wasting 83% of expert 7's compute capacity.

The concept of expert capacity formalizes this. Each expert has a maximum number of tokens it can process in one batch:

def compute_expert_capacity(num_tokens, num_experts, top_k, capacity_factor=1.25):
    """Maximum tokens per expert per batch."""
    # If perfectly balanced, each expert gets (num_tokens * top_k) / num_experts
    balanced_load = (num_tokens * top_k) / num_experts
    capacity = int(balanced_load * capacity_factor)
    return capacity

# Example: 4096 tokens, 8 experts, top_k=2, capacity_factor=1.25
# balanced_load = (4096 * 2) / 8 = 1024
# capacity = 1280 tokens per expert

Tokens that exceed an expert's capacity are dropped and passed through a residual connection without expert processing. The capacity factor (typically 1.0 to 1.5) controls the tradeoff between wasted compute and dropped tokens. A factor of 1.0 means no headroom and frequent drops. A factor of 1.5 means 50% extra capacity but more wasted compute when experts are underloaded.

The permutation and unpermutation step

The most intricate part of MoE inference is the token permutation. After routing, you need to group tokens by their assigned expert, process each group through the corresponding expert FFN, and then scatter the results back to their original positions. This is essentially a gather-scatter pattern:

def moe_forward(x, router, experts):
    """Full MoE forward pass with gather-scatter."""
    weights, indices = router(x)  # (num_tokens, top_k), (num_tokens, top_k)

    # Permute: group tokens by expert
    output = torch.zeros_like(x)
    for k in range(router.top_k):
        for expert_id in range(len(experts)):
            # Find tokens assigned to this expert at this top-k position
            mask = (indices[:, k] == expert_id)
            if mask.any():
                expert_input = x[mask]
                expert_output = experts[expert_id](expert_input)
                output[mask] += weights[mask, k].unsqueeze(-1) * expert_output

    return output

This naive loop is slow. Optimized implementations use custom CUDA kernels to perform the permutation as a batched index-select, process all tokens for each expert as a single batched matmul, and scatter back with a fused kernel. Libraries like Megablocks implement this with block-sparse matrix multiplication, treating the entire MoE layer as a single sparse operation.

DeepSeek's innovations: shared experts and fine-grained routing

DeepSeek-V2 and V3 introduced two important routing innovations:

• Shared experts. Some experts are always activated for every token, providing a stable baseline computation. The router only selects from the remaining "routed" experts. This prevents the degradation that occurs when important tokens are routed to weak experts.
• Fine-grained routing. Instead of 8 large experts with top-2, DeepSeek-V3 uses 256 small experts with top-8. More experts with more selections means better load distribution (the law of large numbers helps) and more specialized expert behavior.

DeepSeek-V3 also replaces the auxiliary loss with a bias term added to the router logits during training. Each expert has a learnable bias that can be adjusted to push tokens toward underutilized experts without distorting the actual routing probabilities used for weighting.

Inference implications

For inference engineering, the router creates unique challenges:

• Memory. All expert weights must be in memory even though only k are activated per token. A Mixtral 8x7B model needs memory for all 8 experts (~46.7B params) but only does compute for 2.
• Batching. Different tokens in a batch may route to different experts, creating irregular computation patterns. This is hard to optimize on GPUs which prefer uniform workloads.
• Expert parallelism. The natural parallelism strategy is to place different experts on different GPUs and use all-to-all communication to route tokens. This is tomorrow's topic.

The router is just a linear layer, but it creates a systems problem that spans load balancing, communication patterns, and memory management. Building it from scratch makes you appreciate why MoE inference is a different discipline from dense model inference.

Tomorrow: expert parallelism simulation, where we put different experts on different GPUs and work through the all-to-all communication pattern.