Capacity and dispatch

After routing, each token has a top_k list of expert indices and weights. The dispatch layer has to:

  1. Optionally cap how many tokens each expert can receive (capacity_factor).

  2. Actually run each expert on its assigned tokens — either in a Python loop (sequential) or via one batched torch._grouped_mm call (grouped GEMM).

  3. Weight-combine the expert outputs back into one tensor per token.

Capacity and dispatch both live in kempnerforge/model/moe.py.

Capacity factor

[model]
moe_capacity_factor = 0.0   # 0 = no cap, >0 = cap per expert

With capacity_factor > 0, each expert can only accept a bounded number of tokens per forward:

# kempnerforge/model/moe.py — _apply_capacity
capacity = max(1, ceil(num_tokens * top_k / num_experts * capacity_factor))

At capacity_factor = 1.0, each expert gets exactly the average load. At 1.25, each expert gets 25% headroom above the average — the canonical Switch Transformer setting. At 0.0 (default), there is no cap: every token goes to its top-k choices regardless of load.

Drop policy

Tokens routed to an expert past capacity are dropped from that expert — their routing weight is set to 0:

# kempnerforge/model/moe.py — _apply_capacity (trimmed)
for k in range(top_k):
    for e in range(num_experts):
        assigned = (indices[:, k] == e).nonzero(as_tuple=True)[0]
        if assigned.numel() <= capacity:
            continue
        drop = assigned[capacity:]
        weights[drop, k] = 0.0

“First capacity tokens wins” — tokens are kept in sequence order, later tokens get dropped. If a token’s only chosen expert is overloaded, the token contributes nothing to the MoE output — but the residual connection still carries it through unchanged:

# kempnerforge/model/transformer.py — TransformerBlock.forward
x = x + self.mlp(self.mlp_norm(x))   # MoE output + residual

So a dropped token is not an error — it passes through the layer as if it skipped the MLP entirely.

When to use capacity

With EP: sometimes necessary. All-to-all buffers are sized for the worst-case token distribution. Without a cap, one rank can overflow its dispatch buffer when many tokens happen to prefer experts on that rank. A capacity factor of 1.25 bounds the buffer size predictably.

Without EP: usually leave at 0. All experts live on the same rank and the sequential/grouped path handles any distribution. Setting a cap just costs throughput.

Caveats: the inner loop is O(num_experts × top_k) Python and non-vectorized. At num_experts = 64 × top_k = 8 = 512 Python iterations per forward, it starts to show on profile. Okay for the scale MoE runs we’ve tested; watch for it at 128+ experts.

Dispatch: two compute paths

The dispatch happens in MoEMLP._local_forward. There are two paths guarded by a feature check:

# kempnerforge/model/moe.py — module-level
_HAS_GROUPED_MM = hasattr(torch, "_grouped_mm")
_GROUPED_MM_DTYPES = {torch.bfloat16, torch.float16}

torch._grouped_mm exists in recent PyTorch (2.5+) but requires bf16/fp16 inputs — the meta registration rejects fp32.

Path A: grouped GEMM (bf16/fp16)

The fast path when available. Tokens are sorted by expert and the per-expert counts become an int32 prefix sum that tells the kernel where each group starts, so no expert is padded to another’s size:

# kempnerforge/model/moe.py — grouped_expert_forward (trimmed)
offs = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)        # (E,)
if is_swiglu:
    hidden = F.silu(_grouped_mm(x_sorted, gate_w, offs)) * _grouped_mm(x_sorted, up_w, offs)
else:
    hidden = activation(_grouped_mm(x_sorted, up_w, offs))
out = _grouped_mm(hidden, down_w, offs)                               # (total_tokens, dim)

_grouped_mm is one ragged torch._grouped_mm(x, w, offs=offs) call on CUDA: group i (rows offs[i-1]:offs[i]) is multiplied by w[i], and empty groups are fine. Three calls for SwiGLU (gate + up + down), two for a standard MLP. Activation memory and compute scale with total_tokens = num_tokens × top_k, whatever the routing looks like — a hot expert costs no more than a balanced one. The counts stay on the device; the dispatch never syncs with the host.

The same call runs on CPU through torch’s reference kernel, so the unit tests exercise this exact dispatch.

Path B: sequential loop (fp32)

Fallback when grouped GEMM can’t run:

# kempnerforge/model/moe.py — MoEMLP._local_forward (trimmed)
for i in range(self.num_experts):
    mask = (indices == i).any(dim=-1)
    if not mask.any():
        continue
    expert_input   = x_flat[mask]
    expert_output  = self.experts[i](expert_input)
    weight_for_i   = (weights * (indices == i).float()).sum(dim=-1)
    output[mask]  += weight_for_i[mask].unsqueeze(-1) * expert_output

One kernel launch per expert (or more, per matmul). Fine for num_experts ≤ 8, painful at 32+.

Rule of thumb: if you’re in fp32, you’re not doing real MoE training. The fallback exists for unit-testing and debugging the grouped path without a CUDA device. Production runs use bf16 (default) or fp16, both of which hit the grouped path.

Packed experts

Opt-in via moe_packed_experts = true. Replaces the nn.ModuleList of experts with three packed nn.Parameter tensors of shape (num_experts, dim, hidden):

# kempnerforge/model/moe.py — MoEMLP.__init__ (trimmed)
self.up_w   = nn.Parameter(stack(e.up_proj.weight.t()   for e in experts))
self.down_w = nn.Parameter(stack(e.down_proj.weight.t() for e in experts))
self.gate_w = nn.Parameter(stack(e.gate_proj.weight.t() for e in experts))  # SwiGLU

Grouped GEMM consumes these directly via grouped_expert_forward_packed — no per-step torch.stack over the ModuleList, and FSDP2 wraps three tensors instead of 3 × E.

Measured impact (see Benchmark § MoE Expert Packing):

Experts

Unpacked tok/s

Packed tok/s

Delta

8

48,521

50,972

+5.1%

16

26,994

36,860

+36.5%

64

1,796

2,204

+22.7%

The speedup grows with num_experts — the per-step torch.stack cost (in unpacked mode) scales with E, so eliminating it matters more as E grows. Peak memory is approximately the same.

Default is off because EP integration (slicing packed weights along the expert axis) is newer than the unpacked path. Flip it on for MoE runs with ≥ 16 experts.

Output combine

Both paths end in a weighted scatter-add back to the token axis:

# kempnerforge/model/moe.py — _local_forward (grouped path, trimmed)
expert_out = expert_out * sorted_weights.unsqueeze(-1)
output = torch.zeros(num_tokens, dim, dtype=x_flat.dtype, device=x_flat.device)
output.scatter_add_(0, sorted_token_ids.unsqueeze(-1).expand_as(expert_out), expert_out)

Dropped tokens (weight=0) contribute 0 here. Tokens with multiple experts (top-k > 1) get a sum of weighted outputs — the same token appears in top_k sorted-positions, and each contribution is added.

After _local_forward returns, the shared expert (if any) is added:

# kempnerforge/model/moe.py — MoEMLP.forward
if self.shared_expert is not None:
    output = output + self.shared_expert(x_flat)

Shared expert runs on every token, dense-style. It’s additive, not part of the top-k selection.

EP branching

When ep_world_size > 1, MoEMLP.forward bypasses _local_forward and calls ep_dispatch_and_compute from kempnerforge/distributed/expert_parallel.py:

# kempnerforge/model/moe.py — MoEMLP.forward (trimmed)
if self.ep_world_size > 1:
    output = ep_dispatch_and_compute(
        x_flat, weights, indices, self,
        self.ep_group, self.local_expert_start,
        self.num_local_experts, self.ep_world_size,
        gradient_scale=self.gradient_scale,
    )
else:
    output = self._local_forward(x_flat, weights, indices)

The capacity factor still applies before this branch — the weights tensor handed to ep_dispatch_and_compute has dropped-token weights already zeroed. See Expert parallelism § Dispatch / combine flow for the all-to-all mechanics.

See also