Mixtral › Expert Execution
Language AI HandbookAll chapters Mixtral › Expert ExecutionOne section · Sections in this chapter5 / 15 ↓
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
With routing decisions made, the MoE layer must execute the selected experts and combine their outputs. The naive implementation loops over tokens, executing each token's two selected experts sequentially, but this is inefficient for GPU execution because it creates many small matrix operations that underutilize GPU parallelism. GPUs reach their performance by executing thousands of operations simultaneously. Small matrix multiplications cannot saturate the available compute resources.
Think of GPU compute as a factory floor with thousands of workstations operating in parallel. Dense matrix multiplication is like feeding a long continuous assembly line where every workstation is always busy. Routing tokens to experts is like suddenly assigning different subsets of workstations to different tasks: some workstations sit idle waiting for their expert's tokens to arrive, while others are temporarily overwhelmed. The key challenge in efficient MoE execution is reorganizing the work so that all workstations stay busy as much as possible, which requires batching tokens by expert rather than processing them in their original order.
The key insight is that the output of the MoE layer does not depend on the order in which tokens are processed by experts. Each token's output is computed independently using only that token's hidden state and the selected experts' weights. We can therefore collect all tokens assigned to expert into a single batch, run one large matrix multiplication, and scatter the results back to their original positions. This reordering preserves correctness while creating the large, regular matrix operations that GPUs execute efficiently.
A more practical approach processes all tokens assigned to each expert in parallel: group all tokens routed to expert 1, process them in a single batched operation, then move to expert 2, and so on. This expert-centric iteration creates larger matrix multiplications that GPUs can execute efficiently. Instead of computing 20 tokens one at a time through expert 1, we batch all 20 into a single matrix multiplication. The ordering of tokens does not matter for the expert computation since experts process tokens independently, so we can reorganize the computation for efficiency without affecting correctness:
import torch
import torch.nn as nn
import torch.nn.functional as F
class MixtralRouter(nn.Module):
def __init__(self, hidden_dim: int, num_experts: int, top_k: int = 2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
def forward(self, hidden_states: torch.Tensor):
router_logits = self.gate(hidden_states)
top_k_logits, selected_experts = torch.topk(
router_logits, self.top_k, dim=-1
)
router_weights = F.softmax(top_k_logits, dim=-1)
return router_weights, selected_experts
class MixtralExpert(nn.Module):
def __init__(self, hidden_dim: int, ffn_dim: int):
super().__init__()
self.gate_proj = nn.Linear(hidden_dim, ffn_dim, bias=False)
self.up_proj = nn.Linear(hidden_dim, ffn_dim, bias=False)
self.down_proj = nn.Linear(ffn_dim, hidden_dim, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class MixtralMoELayer(nn.Module):
def __init__(
self, hidden_dim: int, ffn_dim: int, num_experts: int, top_k: int = 2
):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.router = MixtralRouter(hidden_dim, num_experts, top_k)
self.experts = nn.ModuleList(
[MixtralExpert(hidden_dim, ffn_dim) for _ in range(num_experts)]
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
batch_size, seq_len, hidden_dim = hidden_states.shape
router_weights, selected_experts = self.router(hidden_states)
hidden_flat = hidden_states.view(-1, hidden_dim)
output = torch.zeros_like(hidden_flat)
for expert_idx in range(self.num_experts):
expert_mask = selected_experts == expert_idx
if not expert_mask.any():
continue
batch_idx, seq_idx, topk_idx = torch.where(expert_mask)
flat_idx = batch_idx * seq_len + seq_idx
expert_input = hidden_flat[flat_idx]
expert_weights = router_weights[batch_idx, seq_idx, topk_idx]
expert_output = self.experts[expert_idx](expert_input)
output.index_add_(
0, flat_idx, expert_output * expert_weights.unsqueeze(-1)
)
return output.view(batch_size, seq_len, hidden_dim)This implementation iterates over experts rather than tokens, which is more efficient when tokens are unevenly distributed across experts. Each expert processes all its assigned tokens in a single batched forward pass, creating matrix multiplications large enough to utilize GPU compute effectively.
Let's verify the MoE layer works correctly:
import torch
torch.manual_seed(42)
moe_layer = MixtralMoELayer(hidden_dim=256, ffn_dim=512, num_experts=8, top_k=2)
test_input = torch.randn(2, 16, 256)
with torch.no_grad():
test_output = moe_layer(test_input)Input shape: torch.Size([2, 16, 256]) Output shape: torch.Size([2, 16, 256]) Output magnitude (mean abs): 0.0637 Expert utilization (tokens assigned): Expert 0: 16 ████████ Expert 1: 11 █████ Expert 2: 7 ███ Expert 3: 7 ███ Expert 4: 5 ██ Expert 5: 2 █ Expert 6: 10 █████ Expert 7: 6 ███
The shape preservation (input and output both have shape [2, 16, 256]) confirms the MoE layer acts as a drop-in replacement for a standard FFN. The output magnitude checks that values remain reasonable rather than degenerate. The expert-utilization histogram reveals load distribution: with 32 total tokens (2 batches × 16 sequence length) and top-2 routing, we expect 64 total expert assignments. The visualization shows whether routing is balanced or whether certain experts dominate. In a well-trained model, this distribution should be relatively balanced, though some variation is expected if experts specialize in different input types.