32. Sparse MoE Token Dispatch & Combine (Capacity-Limited)

Problem

Once a router has decided which expert each token goes to, a real MoE layer's forward pass still has to physically move tokens to their assigned expert's compute and back — this is the "dispatch/combine" plumbing that sits underneath every production sparse-MoE implementation. Because GPU kernels want fixed-size batches, each expert is given a fixed capacity: if more tokens than capacity get routed to one expert, the overflow tokens (in original token order) are simply dropped for this layer — a real tradeoff between perfect routing fidelity and predictable compute.

Given each token's already-computed top-1 expert assignment and gate weight, dispatch tokens into per-expert buffers of size at most capacity (first-come-first-served by original token order), apply each expert's linear transform, and combine the results back into token order, weighted by the gate weight. Dropped tokens get an all-zero output row.

Source: src/32_sparse_moe_dispatch_combine.py

def sparse_moe_dispatch_combine(
    x: torch.Tensor, expert_idx: torch.Tensor, gate_weight: torch.Tensor,
    expert_weights: torch.Tensor, capacity: int,
) -> torch.Tensor: ...

>>> x = torch.randn(8, 4)
>>> expert_idx = torch.tensor([0, 0, 0, 1, 1, 2, 0, 1])
>>> gate_weight = torch.rand(8)
>>> expert_weights = torch.randn(3, 4, 4)
>>> sparse_moe_dispatch_combine(x, expert_idx, gate_weight, expert_weights, capacity=2).shape
torch.Size([8, 4])

Step-by-Step Approach

1. One-hot encode the assignment: onehot = F.one_hot(expert_idx, num_experts).float(), shape (T, E).
2. Compute each token's position within its expert's queue via a cumulative sum: (onehot.cumsum(dim=0) - onehot) * onehot, summed over the expert axis — this counts, for each token, how many earlier tokens were assigned to the same expert.
3. Tokens whose position is >= capacity are dropped: keep_mask = position_in_expert < capacity.
4. Build a one-hot dispatch tensor of shape (T, E, C): for each kept token, set dispatch[t, expert_idx[t], position_in_expert[t]] = 1 — this is literally "the address" of where in the expert-major buffer that token's data goes.
5. Gather tokens into a per-expert fixed-size batch with torch.einsum("tec,td->ecd", dispatch, x), giving (num_experts, capacity, d_model), then apply every expert's linear transform in one torch.bmm call.
6. Combine back to token order: multiply dispatch by the gate weight, then torch.einsum("tec,ecd->td", combine, expert_outputs) scatters each expert-slot's output to its original token. Dropped tokens have an all-zero dispatch row and so automatically get an all-zero output.

The key insight is that dispatch and combine are the same one-hot tensor used in two different einsum contractions (gather in one direction, scatter in the other) — there's no need for separate gather/scatter index bookkeeping once that one-hot "address matrix" exists.

Reference solution

import torch
import torch.nn.functional as F


def sparse_moe_dispatch_combine(
    x: torch.Tensor, expert_idx: torch.Tensor, gate_weight: torch.Tensor,
    expert_weights: torch.Tensor, capacity: int,
) -> torch.Tensor:
    num_tokens, d_model = x.shape
    num_experts = expert_weights.shape[0]

    onehot = F.one_hot(expert_idx, num_experts).float()   # (T, E) one-hot assignment
    # position of each token within its expert's queue, counting only
    # earlier tokens assigned to the SAME expert (exclusive cumulative count)
    position_in_expert = ((onehot.cumsum(dim=0) - onehot) * onehot).sum(dim=1).long()  # (T,)
    keep_mask = position_in_expert < capacity              # overflow tokens dropped

    # dispatch[t, e, c] = 1 iff token t is kept AND assigned to expert e AND
    # lands in that expert's slot c -- a one-hot "address" for every kept token
    dispatch = torch.zeros(num_tokens, num_experts, capacity)
    keep_idx = keep_mask.nonzero(as_tuple=True)[0]
    dispatch[keep_idx, expert_idx[keep_idx], position_in_expert[keep_idx]] = 1.0

    # gather: (T,E,C) x (T,d_model) -> (E,C,d_model), i.e. build each
    # expert's fixed-size input batch from scattered tokens in one einsum
    expert_inputs = torch.einsum("tec,td->ecd", dispatch, x)     # (E, C, d_model)
    expert_outputs = torch.bmm(expert_inputs, expert_weights)    # per-expert FFN, batched over E

    # combine: same dispatch tensor (now weighted by gate prob) scatters
    # each expert-slot's output back to its original token position; dropped
    # tokens have an all-zero dispatch row, so they get exactly zero output
    combine = dispatch * gate_weight.view(-1, 1, 1)
    output = torch.einsum("tec,ecd->td", combine, expert_outputs)  # (T, d_model)
    return output

Key Functions & Tricks

• F.one_hot(expert_idx, num_experts) — converts integer expert assignments into a (T, E) indicator matrix, the basis for every subsequent vectorized step.
• onehot.cumsum(dim=0) — a running per-expert token count along the token axis; subtracting onehot itself makes it exclusive (position before this token).
• torch.einsum("tec,td->ecd", dispatch, x) — the gather step: for each expert/slot, sums over tokens weighted by the one-hot dispatch tensor, which is equivalent to "pick out exactly the one assigned token" since each column of dispatch has at most one 1.
• torch.bmm(expert_inputs, expert_weights) — applies each expert's own weight matrix to its own capacity-sized batch in a single batched matmul over the expert axis.
• tensor.nonzero(as_tuple=True) — converts a boolean keep-mask into integer indices suitable for advanced (fancy) indexing when scattering into dispatch.

How to Recognize This Pattern