Top-K Routing › Routing Implementation
Language AI HandbookAll chapters Top-K Routing › Routing ImplementationOne section · Sections in this chapter5 / 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
Implementing top-K routing correctly requires handling several details: selecting the top K indices, computing normalized weights, dealing with numerical stability, and letting gradient flow. Each of these sounds simple in isolation, but combining them into a reliable, efficient implementation requires care. Bugs in the routing implementation can silently degrade model quality in ways that are difficult to debug, so it pays to understand what each step is doing and why.
The core routing operation takes router logits and produces both the selected expert indices and their corresponding weights. Let's build this step by step. A reliable implementation must be efficient during both training and inference, and must handle edge cases like ties in routing scores gracefully. PyTorch's torch.topk provides an efficient GPU-accelerated top-K selection that handles most of the heavy lifting.
The Core Routing Function
The routing function itself is remarkably concise once you have the right ingredients:
import torch
import torch.nn.functional as F
def top_k_routing(router_logits, k=2):
"""
Select top-k experts and compute their normalized weights.
Args:
router_logits: Tensor of shape (batch_size, seq_len, num_experts)
k: Number of experts to select per token
Returns:
expert_weights: Normalized weights for selected experts (batch, seq, k)
expert_indices: Indices of selected experts (batch, seq, k)
"""
# Get top-k logits and their indices
top_k_logits, expert_indices = torch.topk(router_logits, k, dim=-1)
# Normalize weights using softmax over only the selected experts
expert_weights = F.softmax(top_k_logits, dim=-1)
return expert_weights, expert_indicesThe torch.topk function efficiently finds the K largest values along the expert dimension. It returns both the values and their indices in a single operation. This makes the selection efficient even for large numbers of experts. We then apply softmax only to these selected logits. This keeps the weights sum to 1 across the K active experts. This two-step process, first selecting then normalizing, is needed for obtaining proper convex combination weights.
Notice that the softmax is applied to top_k_logits (shape: batch, seq, k), not to router_logits (shape: batch, seq, num_experts). This is the localized softmax we discussed in the theory section. If we had called F.softmax(router_logits, dim=-1) and then indexed into the result, we would get probabilities that don't sum to 1 over the selected pair. The order of operations matters here.
Let's verify this works as expected:
# Simulate routing for a small batch
batch_size, seq_len, num_experts = 2, 4, 8
router_logits = torch.randn(batch_size, seq_len, num_experts)
# Apply top-2 routing
weights, indices = top_k_routing(router_logits, k=2)Router logits shape: torch.Size([2, 4, 8]) Expert weights shape: torch.Size([2, 4, 2]) Expert indices shape: torch.Size([2, 4, 2]) For first token in batch: All router logits: [1.9269150495529175, 1.4872841835021973, 0.9007171988487244, -2.1055214405059814, 0.6784184575080872, -1.2345449924468994, -0.043067481368780136, -1.6046669483184814] Selected expert indices: [0, 1] Normalized weights: [0.6081710457801819, 0.3918289244174957] Weights sum to: 1.0000
The output shows that for each token, we get exactly K = 2 expert indices and their corresponding normalized weights that sum to 1. The two selected indices correspond to the two highest-valued entries in the logit vector for that token, and the weights reflect the relative magnitudes of those logits after softmax normalization.
Preserving the Full Probability Distribution
For top-1 routing specifically, we often want to preserve the routing probability distribution even though only one expert is selected. The reason is subtle: load balancing auxiliary losses (which we cover in the next chapter) need the full routing probabilities rather than only the argmax. If we discard the full distribution, we cannot compute the auxiliary losses that keep routing balanced during training.
The implementation needs both discrete forward-pass selection and the full probability distribution for computing auxiliary penalties:
def top_1_routing_with_probs(router_logits):
"""
Top-1 routing that preserves full probability distribution.
Returns:
expert_index: Index of selected expert (batch, seq)
routing_probs: Full softmax probabilities (batch, seq, num_experts)
"""
# Full probability distribution for auxiliary losses
routing_probs = F.softmax(router_logits, dim=-1)
# Select the argmax expert
expert_index = torch.argmax(router_logits, dim=-1)
return expert_index, routing_probsThis design pattern, where the routing function returns auxiliary information beyond just the selected experts, is common in practical MoE implementations. The router_logits themselves are also frequently returned from the MoE layer, since auxiliary losses typically operate on the logits directly rather than on the derived probabilities.