14. Mixture-of-Experts Top-2 Routing Layer

Problem

Sparse mixture-of-experts (MoE) layers scale model capacity without scaling per-token compute: instead of running every token through one big feed-forward block, a small router picks a handful of "experts" per token from a larger pool, and only those experts run. Top-2 gating (Shazeer et al. 2017's Sparsely-Gated MoE; used in Switch Transformer and Mixtral-style models) routes each token to its two highest-scoring experts and combines their outputs weighted by a softmax computed over just those two logits, not the full expert pool.

Given x (shape (batch, d_model)), gate_logits (shape (batch, num_experts)), and expert_weights (shape (num_experts, d_model, d_model), each a plain linear map with no bias), select each token's top-2 experts, combine their outputs by a locally-renormalized softmax, and return the combined output plus which experts were picked and their weights.

Source: src/14_moe_top2_routing.py

def moe_top2_forward(x, gate_logits, expert_weights) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

>>> import torch
>>> torch.manual_seed(0)
>>> x = torch.tensor([[1.0, 0.0, 0.0]])
>>> gate_logits = torch.tensor([[1.0, 5.0, 0.5, 2.0]])
>>> expert_weights = torch.randn(4, 3, 3)
>>> out, idx, w = moe_top2_forward(x, gate_logits, expert_weights)
>>> idx
tensor([[1, 3]])

Step-by-Step Approach

1. Select each token's top-2 experts with torch.topk(gate_logits, k=2, dim=-1), which returns both the top-2 logit values and the expert indices they belong to, already ranked highest first.
2. Softmax only the two selected logits with torch.softmax(topk_logits, dim=-1) — this is what makes it "top-2 gating" instead of a full soft mixture over every expert, and is the detail most easily gotten wrong (softmaxing the full gate_logits row first and then indexing into it gives the wrong weights, since the unselected experts would still pull normalization mass).
3. Gather each token's two chosen expert weight matrices in a single indexing operation: expert_weights[topk_idx] broadcasts the (batch, 2) index tensor into a (batch, 2, d_model, d_model) gathered weight tensor.
4. Apply each chosen expert's linear map to its token's input with torch.einsum("bnji,bi->bnj", chosen_W, x), avoiding ever materializing a per-token dense pass through all num_experts experts.
5. Combine the two expert outputs by their gating weights: broadcast-multiply topk_weights.unsqueeze(-1) against the expert outputs and sum over the "which of the 2 experts" dimension.

The key insight is that "top-2 gating" specifically means the softmax normalization happens after selection, over only the selected experts' logits — this is what keeps a token's total routing weight summing to exactly 1 across its two chosen experts regardless of how many total experts exist, and is the detail that distinguishes true top-k MoE gating from a "soft mixture with hard top-k masking" that superficially looks similar but computes different combination weights.

Reference solution

import torch


def moe_top2_forward(x, gate_logits, expert_weights):
    # pick each token's top-2 experts by router score; topk already returns
    # them ranked highest-first, which we rely on for tests/inspection
    topk_logits, topk_idx = torch.topk(gate_logits, k=2, dim=-1)  # (batch, 2)
    # softmax over ONLY the 2 selected logits -- not the full expert pool --
    # is what makes this "top-2 gating" rather than a full soft mixture
    topk_weights = torch.softmax(topk_logits, dim=-1)  # (batch, 2)

    # gather each token's two chosen expert weight matrices in one indexing op
    chosen_W = expert_weights[topk_idx]  # (batch, 2, d_model, d_model)
    # apply y = x @ W^T per selected expert without materializing a (batch,
    # num_experts, ...) tensor for experts that weren't even selected
    expert_out = torch.einsum("bnji,bi->bnj", chosen_W, x)  # (batch, 2, d_model)
    # combine the two expert outputs by their gating weights
    output = (topk_weights.unsqueeze(-1) * expert_out).sum(dim=1)  # (batch, d_model)

    return output, topk_idx, topk_weights

Key Functions & Tricks

• torch.topk(gate_logits, k=2, dim=-1) — selects both which experts and their raw scores in one call, pre-ranked.
• torch.softmax(topk_logits, dim=-1) applied post-selection — the defining detail of top-k gating versus a full soft mixture.
• expert_weights[topk_idx] — advanced/fancy indexing that turns a (batch, 2) index tensor into a per-token gathered batch of weight matrices without a Python loop over tokens.
• torch.einsum("bnji,bi->bnj", chosen_W, x) — applies a per-token, per-expert-slot linear map in one call; the index letters make explicit which dims are batched (b), which are the "slot" dim (n), and which are contracted (i).
• Tensor.unsqueeze(-1) then broadcast-multiply and .sum(dim=1) — the standard weighted-combination pattern once per-slot outputs and per-slot weights are aligned.

How to Recognize This Pattern