Top-K Routing › Combining Expert Outputs
Language AI HandbookAll chapters Top-K Routing › Combining Expert OutputsOne section · Sections in this chapter6 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Once we have selected experts and computed their weights, tokens must be routed through the selected experts and their outputs combined. Unlike a standard feed-forward layer where all tokens pass through the same parameters simultaneously, MoE dynamically dispatches different tokens to different experts before gathering and combining the results. The dispatch and gather operations are the engineering bottleneck of MoE, and efficient implementation is necessary to achieve the theoretical compute savings.
The basic challenge is that the routing decisions are different for every token. Token A might go to experts 2 and 5, while token B goes to experts 1 and 7, and token C goes to experts 2 and 3. We cannot simply pass all tokens through all experts and then zero out the wrong contributions, because that would defeat the purpose of sparse routing. We need to physically route each token only to its selected experts, and we need to do this in a way that exploits GPU parallelism rather than serializing over individual tokens.
The Naive (But Slow) Approach
The simplest approach processes tokens individually:
def combine_expert_outputs_simple(
hidden_states, expert_weights, expert_indices, experts
):
"""
Simple (but slow) expert combination by iterating over tokens.
Args:
hidden_states: Input tensor (batch, seq, hidden_dim)
expert_weights: Weights for selected experts (batch, seq, k)
expert_indices: Indices of selected experts (batch, seq, k)
experts: List of expert modules
Returns:
combined_output: Weighted combination of expert outputs (batch, seq, hidden_dim)
"""
batch_size, seq_len, hidden_dim = hidden_states.shape
k = expert_weights.shape[-1]
output = torch.zeros_like(hidden_states)
for b in range(batch_size):
for s in range(seq_len):
for i in range(k):
expert_idx = expert_indices[b, s, i].item()
weight = expert_weights[b, s, i]
expert_output = experts[expert_idx](
hidden_states[b, s].unsqueeze(0)
)
output[b, s] += weight * expert_output.squeeze(0)
return outputThis explicit loop makes the logic clear: for each token, iterate through its K selected experts, compute each expert's output, and accumulate them weighted by the routing weights. However, this approach is extremely slow because it processes tokens sequentially. Modern GPUs are designed to process thousands of operations simultaneously through SIMD parallelism, and this triple-nested loop structure forces strictly sequential execution that underutilizes the hardware. For a sequence of 2048 tokens with K=2, this requires 4096 individual forward passes through expert networks, each processing a single vector. The GPU sits mostly idle, waiting for each sequential operation to complete.
The Efficient Batched Approach
The efficient approach groups tokens by their assigned experts and processes all tokens assigned to the same expert simultaneously:
def combine_expert_outputs_efficient(
hidden_states, expert_weights, expert_indices, experts
):
"""
Efficient expert combination using batched operations.
Groups tokens by expert and processes each expert's tokens in parallel.
"""
batch_size, seq_len, hidden_dim = hidden_states.shape
num_experts = len(experts)
k = expert_weights.shape[-1]
# Flatten batch and sequence dimensions
flat_hidden = hidden_states.view(-1, hidden_dim) # (batch*seq, hidden)
flat_weights = expert_weights.view(-1, k) # (batch*seq, k)
flat_indices = expert_indices.view(-1, k) # (batch*seq, k)
# Initialize output
output = torch.zeros_like(flat_hidden)
# Process each expert position (first selected, second selected, etc.)
for pos in range(k):
pos_indices = flat_indices[:, pos] # Which expert for this position
pos_weights = flat_weights[:, pos] # Weight for this position
# Process each expert's assigned tokens
for expert_idx in range(num_experts):
# Find tokens assigned to this expert at this position
mask = pos_indices == expert_idx
if mask.any():
expert_input = flat_hidden[mask]
expert_output = experts[expert_idx](expert_input)
output[mask] += pos_weights[mask].unsqueeze(-1) * expert_output
return output.view(batch_size, seq_len, hidden_dim)This version groups all tokens assigned to each expert and processes them in a single batched forward pass through that expert. The key insight is that even though different tokens go to different experts, all tokens assigned to the same expert can be processed together in one operation. The outer loop over K expert positions and the inner loop over experts means we make exactly forward passes through expert networks, but each pass processes all relevant tokens in parallel. For a model with 8 experts, K=2, and 2048 tokens, this means 16 forward passes through expert networks instead of 4096. Each of those 16 passes processes a batch of tokens rather than a single token, fully exploiting GPU parallelism.
The mask tensor identifies which tokens are assigned to a given expert at a given selection position. The flat_hidden[mask] operation gathers the relevant token representations into a contiguous batch, the expert processes them all simultaneously, and the result is scattered back to the correct positions in the output tensor. This gather-process-scatter pattern is the basic building block of efficient sparse computation on GPUs.
Complete MoE Layer
Let's create a complete working implementation combining routing and expert computation:
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
"""Simple feed-forward expert network."""
def __init__(self, hidden_dim, intermediate_dim):
super().__init__()
self.w1 = nn.Linear(hidden_dim, intermediate_dim)
self.w2 = nn.Linear(intermediate_dim, hidden_dim)
def forward(self, x):
return self.w2(F.gelu(self.w1(x)))
class Router(nn.Module):
"""Gating network that produces routing logits."""
def __init__(self, hidden_dim, num_experts):
super().__init__()
self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
def forward(self, x):
return self.gate(x)import torch.nn as nn
class TopKMoELayer(nn.Module):
"""Complete top-K MoE layer combining routing and expert computation."""
def __init__(self, hidden_dim, intermediate_dim, num_experts, k=2):
super().__init__()
self.num_experts = num_experts
self.k = k
self.router = Router(hidden_dim, num_experts)
self.experts = nn.ModuleList(
[Expert(hidden_dim, intermediate_dim) for _ in range(num_experts)]
)
def forward(self, hidden_states):
batch_size, seq_len, hidden_dim = hidden_states.shape
# Get routing decisions
router_logits = self.router(hidden_states)
expert_weights, expert_indices = top_k_routing(router_logits, self.k)
# Combine expert outputs
output = combine_expert_outputs_efficient(
hidden_states, expert_weights, expert_indices, self.experts
)
return output, router_logits # Return logits for auxiliary lossesThe layer returns both the output and the raw router logits. The logits are passed up to the enclosing transformer layer so they can be used in computing auxiliary load balancing losses. This is why it is important to separate the routing computation from the expert computation in the implementation: we need access to the intermediate values even after the forward pass is complete.
Let's test the complete layer:
# Create MoE layer
hidden_dim = 256
intermediate_dim = 512
num_experts = 8
k = 2
moe_layer = TopKMoELayer(hidden_dim, intermediate_dim, num_experts, k)
# Create sample input
batch_size, seq_len = 4, 16
x = torch.randn(batch_size, seq_len, hidden_dim)
# Forward pass
output, router_logits = moe_layer(x)Input shape: torch.Size([4, 16, 256]) Output shape: torch.Size([4, 16, 256]) Router logits shape: torch.Size([4, 16, 8]) Total expert parameters: 2,103,296 Parameters used per token: 525,824
The output confirms that our MoE layer maintains the expected dimensions while activating only K experts per token. The total expert parameters scale with the number of experts, but each token only uses K experts worth of computation. This separation between total parameter count and per-token computational cost is the central value proposition of MoE architectures: you get the representational power of a large model at the per-token cost of a small one.