How does Grok-1 use a mixture-of-experts design?

Grok-1’s mixture-of-experts design is a case study in scaling transformer capacity without proportional compute cost. At 314 billion total parameters but only 86 billion active per forward pass, it demonstrates how sparse activation patterns can deliver dense-model quality at a fraction of the FLOPs. This post breaks down the routing mechanics, expert specialization, and the practical constraints you hit when serving a model like this in production.

The high-level architecture

Grok-1 follows the standard MoE transformer pattern: replace the dense feed-forward network (FFN) in every other transformer block with a sparse MoE layer. The model has 64 layers total, with MoE layers at every other layer (layers 1, 3, 5… 63), giving 32 MoE layers. Each MoE layer contains 8 experts, and the router selects 2 experts per token (top-2 routing with no auxiliary loss during training, per the release notes).

# Simplified Grok-1 MoE layer structure
class GrokMoELayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_experts = 8
        self.top_k = 2
        self.hidden_size = config.hidden_size  # 6144
        self.expert_dim = config.expert_dim    # 32768 (≈ 5.33x hidden)
        
        # Router: projects to logits over 8 experts
        self.router = nn.Linear(self.hidden_size, self.num_experts, bias=False)
        
        # 8 independent experts, each a SwiGLU FFN
        self.experts = nn.ModuleList([
            SwiGLUExpert(self.hidden_size, self.expert_dim)
            for _ in range(self.num_experts)
        ])
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [batch, seq_len, hidden]
        batch_seq = x.shape[0] * x.shape[1]
        x_flat = x.view(batch_seq, -1)
        
        # Router logits: [batch*seq, 8]
        router_logits = self.router(x_flat)
        
        # Top-2 selection
        weights, indices = torch.topk(router_logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)  # normalize selected weights
        
        # Dispatch to experts and combine
        return self._dispatch_and_combine(x_flat, weights, indices).view_as(x)

The total parameter count breaks down roughly as: attention parameters (shared across all tokens) + 32 layers × 8 experts × expert_params. Each expert is a SwiGLU FFN with input 6144 and hidden 32768, yielding ~8.4M params per expert. 256 experts × 8.4M ≈ 2.15B expert params. The remaining ~312B are in the attention blocks and embeddings — but only the 2 active experts per layer per token contribute to FLOPs.

Routing: top-2 without auxiliary loss

Most MoE implementations (Switch Transformer, Mixtral) use an auxiliary load-balancing loss to prevent router collapse. Grok-1 notably omits this. The xAI team reported that their training recipe — specifically the combination of expert capacity factor, batch size, and learning rate schedule — kept expert utilization balanced without explicit regularization.

# Grok-style routing: pure top-k, no aux loss
def grok_router_forward(
    hidden_states: torch.Tensor,  # [tokens, hidden]
    router_weight: torch.Tensor,  # [num_experts, hidden]
    top_k: int = 2,
    capacity_factor: float = 1.25
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Returns: (dispatch_mask, combine_weights, router_logits)
    dispatch_mask: [tokens, top_k, num_experts] one-hot
    combine_weights: [tokens, top_k] softmax weights
    """
    router_logits = hidden_states @ router_weight.T  # [tokens, 8]
    
    # Top-k selection
    topk_weights, topk_indices = torch.topk(router_logits, top_k, dim=-1)
    topk_weights = F.softmax(topk_weights, dim=-1)  # [tokens, 2]
    
    # Build dispatch mask for scatter/gather
    dispatch_mask = F.one_hot(topk_indices, num_classes=8).float()  # [tokens, 2, 8]
    
    # Optional: capacity-aware dropping (training only)
    if capacity_factor < float('inf'):
        dispatch_mask = _apply_capacity(dispatch_mask, capacity_factor)
    
    return dispatch_mask, topk_weights, router_logits

The absence of auxiliary loss simplifies the training objective but shifts burden to hyperparameter tuning. In practice, xAI used a capacity factor of 1.25 (each expert gets 1.25 × tokens/8 capacity) and dropped overflow tokens during training. At inference, there’s no dropping — you route every token to its top-2 experts and compute both.

Expert implementation: SwiGLU at scale

Each expert is a standard SwiGLU feed-forward block. The expansion ratio (32768/6144 ≈ 5.33) is higher than dense transformers (typically 4x), which compensates for the fact that only 2/8 experts fire per token.

class SwiGLUExpert(nn.Module):
    def __init__(self, hidden_size: int, expert_dim: int):
        super().__init__()
        self.w_gate = nn.Linear(hidden_size, expert_dim, bias=False)
        self.w_up = nn.Linear(hidden_size, expert_dim, bias=False)
        self.w_down = nn.Linear(expert_dim, hidden_size, bias=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [tokens, hidden]
        gate = F.silu(self.w_gate(x))
        up = self.w_up(x)
        return self.w_down(gate * up)

Critical implementation detail: the gate and up projections can be fused into a single matmul for memory bandwidth efficiency.

# Fused SwiGLU: single matmul for gate+up
class FusedSwiGLUExpert(nn.Module):
    def __init__(self, hidden_size: int, expert_dim: int):
        super().__init__()
        self.w_gate_up = nn.Linear(hidden_size, 2 * expert_dim, bias=False)
        self.w_down = nn.Linear(expert_dim, hidden_size, bias=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up = self.w_gate_up(x)  # [tokens, 2*expert_dim]
        gate, up = gate_up.chunk(2, dim=-1)
        return self.w_down(F.silu(gate) * up)

With 8 experts per layer, you have 8 independent weight matrices for w_gate_up (6144 × 65536) and 8 for w_down (32768 × 6144). These don’t share parameters — that’s the point. The specialization emerges from routing gradients.

Training dynamics: expert specialization

Without an auxiliary loss, what prevents all tokens from routing to the same 2 experts? The Grok-1 paper notes that specialization emerges naturally at scale. Early in training, routing is near-random. As experts differentiate, the router learns to send tokens to experts that minimize loss for those token types.

Empirically (from the release and community analysis), experts tend to specialize along syntactic and semantic boundaries:

• Experts 0, 1: code and structured data
• Experts 2, 3: natural language reasoning
• Experts 4, 5: mathematical notation
• Experts 6, 7: multilingual / low-resource patterns

But this is soft specialization — every expert sees every domain, just with different frequency. The router logits for a given token typically show a clear top-2 with a gap to the 3rd expert, but the 3rd and 4th experts often have non-trivial probability mass.

# Analyzing expert utilization (post-training)
@torch.no_grad()
def analyze_expert_utilization(model, dataloader, num_experts=8):
    expert_counts = torch.zeros(num_experts)
    total_tokens = 0
    
    for batch in dataloader:
        # Hook router logits at each MoE layer
        for layer_idx in range(0, 64, 2):  # MoE layers
            router_logits = model.get_router_logits(layer_idx, batch)
            top2 = router_logits.topk(2, dim=-1).indices  # [batch, seq, 2]
            expert_counts += top2.flatten().bincount(minlength=num_experts)
            total_tokens += top2.numel()
    
    utilization = expert_counts / total_tokens * num_experts / 2  # normalized
    return utilization  # should be ~1.0 per expert if balanced

In practice, Grok-1 achieves within 5-10% balance across experts on diverse corpora. The capacity factor during training (1.25) provides a buffer that prevents catastrophic forgetting for under-utilized experts.

Inference: the real engineering challenge

Serving Grok-1 is where the MoE design earns its keep — or causes headaches. The 86B active parameters fit on 2× H100 (80GB) with 4-bit quantization, or 4× A100 (80GB) in FP16. But the memory footprint of all 314B parameters matters for loading and expert placement.

Expert parallelism vs. tensor parallelism

With 8 experts per layer, you have natural expert parallelism (EP) degree 8. But Grok-1’s 32 MoE layers means you need to decide how to shard across devices.

# Two viable sharding strategies for 8 GPUs
# Strategy A: Expert parallel (EP=8, TP=1)
# Each GPU holds 1 expert per MoE layer + all attention params
# All-to-all communication per MoE layer

# Strategy B: Tensor parallel (TP=8, EP=1) 
# Each GPU holds 1/8 of every expert + 1/8 of attention
# All-reduce per matmul, no cross-device routing

# Strategy C: Hybrid (EP=4, TP=2) - often optimal for 8 GPUs
# 4 experts per GPU, each expert split across 2 GPUs
# Reduces all-to-all volume by 2x vs pure EP

The all-to-all communication in expert parallelism is the bottleneck. For each MoE layer, you need to scatter input tokens to their assigned experts and gather outputs. With 32 MoE layers, that’s 32 all-to-alls per forward pass.

# Pseudocode for EP forward pass (simplified)
def moe_layer_ep_forward(x, router, experts, ep_group):
    # x: [local_batch, seq, hidden] on this rank
    router_logits = router(x)
    topk_weights, topk_indices = router_logits.topk(2, dim=-1)
    topk_weights = topk_weights.softmax(dim=-1)
    
    # Build dispatch: which tokens go to which expert (global expert ID)
    # Then all-to-all to send tokens to correct expert ranks
    dispatched_x, dispatched_weights, expert_indices = all_to_all_dispatch(
        x, topk_weights, topk_indices, ep_group
    )
    
    # Local expert computation
    expert_outputs = []
    for local_expert_idx, expert in enumerate(experts):
        mask = (expert_indices == local_expert_idx)
        if mask.any():
            expert_input = dispatched_x[mask]
            expert_weight = dispatched_weights[mask]
            expert_outputs.append(expert(expert_input) * expert_weight.unsqueeze(-1))
    
    # Combine local outputs
    local_output = torch.cat(expert_outputs, dim=0) if expert_outputs else zeros_like(x)
    
    # All-to-all gather back to original token order
    output = all_to_all_gather(local_output, ep_group)
    return output

The communication volume per all-to-all: batch × seq × hidden × 2 bytes (send 2 experts’ worth, receive 2). For batch=4, seq=4096, hidden=6144, FP16: ~400 MB per all-to-all. 32 layers = ~12 GB shuffled per token generation step. On NVLink (600 GB/s), that’s ~20ms of pure communication — significant but manageable.

KV cache considerations

MoE doesn’t change KV cache mechanics — attention is still dense. But the larger hidden size (6144 vs. 4096 in Llama-70B) means larger KV cache per layer. 64 layers × 2 (K+V) × 6144 × 2 bytes = ~3 MB per token per sequence. At 8K context, that’s 24 GB just for KV cache. Plan accordingly.

Tradeoffs: when MoE wins and when it doesn’t

The memory wall is real. You need ~600 GB VRAM just for weights in FP16. Quantization to 4-bit (AWQ/GPTQ) brings this to ~160 GB — feasible on 8× H100 or 4× H100 with offloading. But you’re still loading 314B params from disk at startup, which takes minutes on NVMe.

The batch throughput penalty is subtle but measurable. Dense matmuls achieve near-peak FLOPs utilization. MoE’s scattered expert computation fragments the work: each expert gets batch × seq / 4 tokens on average (2 active experts / 8 total). Small expert batches underutilize tensor cores.

# Expert batch size analysis
def expert_batch_size(global_batch, seq_len, num_experts=8, top_k=2):
    tokens_per_expert = global_batch * seq_len * top_k / num_experts
    return tokens_per_expert

# Examples:
# batch=4, seq=4096 → 4096 tokens/expert (good utilization)
# batch=1, seq=4096 → 1024 tokens/expert (moderate)
# batch=1, seq=512  → 128 tokens/expert (poor - tensor core underutilization)

For latency-critical serving (batch=1), MoE can be slower than a dense model of equivalent active params because the expert matmuls are too small to saturate the GPU. This is why Grok-1 shines at high batch / high throughput scenarios, not at interactive single-stream latency.

The decisive takeaway

Grok-1’s MoE design proves that 314B parameters with 86B active is a Pareto improvement over dense 70-100B models for throughput-oriented workloads — but only if you have the memory bandwidth and multi-GPU infrastructure to amortize the communication overhead. The routing mechanism is deliberately simple (top-2, no aux loss), relying on scale and capacity factor to achieve balance. Expert specialization emerges without explicit supervision.

For engineers evaluating MoE for their own systems: start with a smaller MoE (e.g., Mixtral 8×7B or 8×22B) to validate your expert parallelism pipeline before committing to Grok-1 scale. The communication patterns, capacity planning, and quantization behavior transfer directly. And if your serving profile is low-batch, high-latency-sensitivity, a dense model of equivalent active params will likely outperform on tail latency — the MoE tax isn’t worth it there.