Skand Tandon
Advanced

Expert parallelism simulation

Expert parallelism places each MoE expert on a separate GPU. Tokens must travel to their assigned expert and back, creating an all-to-all communication pattern. Simulating this reveals the bandwidth constraints and load imbalance that dominate MoE serving cost.

2025-12-31 advanced 8 min read

Yesterday I built a MoE router from scratch. Today I am simulating what happens when you distribute those experts across GPUs and actually run the all-to-all communication. This is where MoE inference gets interesting, because the communication pattern is fundamentally different from tensor parallelism and the load balancing problem becomes a systems engineering challenge.

Expert parallelism vs tensor parallelism

In tensor parallelism (TP), every GPU holds a shard of every layer and they collectively compute the output through AllReduce. The communication pattern is uniform: every GPU sends the same amount of data in every step.

In expert parallelism (EP), each GPU holds a subset of experts. Token routing is data-dependent: the router decides, based on the token's hidden state, which expert (and therefore which GPU) should process it. The communication pattern is all-to-all: each GPU sends a different number of tokens to each other GPU, depending on the routing decisions.

# Expert parallelism layout for Mixtral 8x7B on 8 GPUs
# GPU 0: Expert 0  (attention layers are replicated on all GPUs)
# GPU 1: Expert 1
# GPU 2: Expert 2
# ...
# GPU 7: Expert 7

# For each MoE layer:
# 1. All GPUs compute attention (replicated)
# 2. Router computes expert assignments (replicated)
# 3. All-to-all: tokens travel to their assigned expert's GPU
# 4. Each GPU runs its local expert on received tokens
# 5. All-to-all: results travel back to the originating GPU
# 6. Combine expert outputs with router weights

Simulating the all-to-all

Let me simulate the communication pattern for a batch of 4,096 tokens with 8 experts on 8 GPUs, top-2 routing. Each token goes to 2 experts, so there are 8,192 total token-expert assignments.

import numpy as np

np.random.seed(42)
num_tokens = 4096
num_experts = 8
num_gpus = 8
top_k = 2
hidden_size = 4096
bytes_per_element = 2  # FP16

# Simulate router decisions (random for now)
# In practice, routing is learned and not uniform
assignments = np.random.randint(0, num_experts, size=(num_tokens, top_k))

# Build the communication matrix: comm[src_gpu][dst_gpu] = num_tokens
# Assume tokens are distributed evenly across GPUs initially
tokens_per_gpu = num_tokens // num_gpus  # 512

comm_matrix = np.zeros((num_gpus, num_gpus), dtype=int)
for gpu_id in range(num_gpus):
    start = gpu_id * tokens_per_gpu
    end = start + tokens_per_gpu
    for token_idx in range(start, end):
        for k in range(top_k):
            expert_id = assignments[token_idx, k]
            dst_gpu = expert_id  # expert i lives on GPU i
            comm_matrix[gpu_id, dst_gpu] += 1

print("Communication matrix (tokens sent from GPU[row] to GPU[col]):")
print(comm_matrix)
# With random routing, each cell is approximately:
# tokens_per_gpu * top_k / num_experts = 512 * 2 / 8 = 128 tokens

With perfectly balanced random routing, each GPU sends about 128 tokens to each other GPU. The total data per GPU is 128 tokens * hidden_size * bytes_per_element = 128 * 4096 * 2 = 1 MB per direction per pair. Across all 8 GPUs, that is about 7 MB sent and 7 MB received per GPU per MoE layer.

The bandwidth cost

Mixtral 8x7B has 32 transformer layers, each with an MoE FFN. Two all-to-all operations per layer (dispatch and combine) means 64 all-to-all operations per forward pass.

# Total communication per forward pass
tokens_per_a2a = num_tokens * top_k  # 8192 token dispatches
bytes_per_token = hidden_size * bytes_per_element  # 8192 bytes = 8 KB
bytes_per_a2a = tokens_per_a2a * bytes_per_token  # 67 MB

# Two all-to-all per layer, 32 layers:
total_bytes = bytes_per_a2a * 2 * 32  # ~4.3 GB per forward pass

# On NVLink (450 GB/s bidirectional within a node):
nvlink_time = total_bytes / 450e9  # ~9.5 ms

# On InfiniBand (400 Gb/s = 50 GB/s between nodes):
ib_time = total_bytes / 50e9  # ~86 ms

This is the fundamental constraint of expert parallelism: all-to-all communication is expensive, especially across nodes. Within a single NVLink-connected node (like a DGX H100), the 9.5 ms communication overhead is manageable. Across InfiniBand between nodes, 86 ms is devastating for latency-sensitive workloads.

Why EP stays intra-node

This is why most MoE deployments use expert parallelism within a single node (8 GPUs connected by NVLink) and tensor parallelism between nodes if needed. The all-to-all pattern is uniquely sensitive to interconnect bandwidth because every GPU needs to communicate with every other GPU, unlike AllReduce where the communication can be pipelined in a ring.

Load imbalance in practice

The simulation above assumed random (uniform) routing. Real routing is not uniform. Some experts are more popular than others, and popularity varies by input. Let me simulate a skewed distribution:

# Simulate skewed routing: experts 0, 1 get 2x the load
expert_probs = np.array([0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
# Normalize for top-2 selection
assignments_skewed = np.array([
    np.random.choice(num_experts, size=top_k, replace=False, p=expert_probs)
    for _ in range(num_tokens)
])

# Count tokens per expert
expert_loads = np.zeros(num_experts, dtype=int)
for token_idx in range(num_tokens):
    for k in range(top_k):
        expert_loads[assignments_skewed[token_idx, k]] += 1

print("Expert loads (skewed):", expert_loads)
# Expert 0: ~1640, Expert 1: ~1640, Experts 2-7: ~820 each

# Compute time is determined by the most loaded expert
max_load = expert_loads.max()  # ~1640
min_load = expert_loads.min()  # ~820

# Load imbalance ratio:
imbalance = max_load / (num_tokens * top_k / num_experts)  # ~1.6x
# GPU 0 and 1 take 60% longer than the balanced case

A 1.6x load imbalance means the most loaded GPU takes 60% longer than it would under perfect balance. Every other GPU sits idle waiting for it to finish. The effective throughput of the entire system drops by 37.5%.

Strategies for managing EP load imbalance

EP for DeepSeek-V3's 256 experts

DeepSeek-V3 uses 256 routed experts with top-8 selection. On a single 8-GPU node, you cannot give each expert its own GPU. Instead, each GPU holds 32 experts. The all-to-all now routes tokens to the correct GPU (based on which GPU holds the target expert), and then each GPU dispatches tokens locally across its 32 experts.

# DeepSeek-V3 EP layout on 8 GPUs
experts_per_gpu = 256 // 8  # 32 experts per GPU
top_k = 8  # each token selects 8 experts

# With 256 experts and top-8, load is naturally more balanced
# (law of large numbers with more experts)
# Expected tokens per expert per batch:
# num_tokens * top_k / num_experts = 4096 * 8 / 256 = 128 tokens/expert

# But each GPU handles 32 experts, so:
# Expected tokens per GPU = 128 * 32 = 4096 token-expert pairs per GPU
# This is naturally balanced because each GPU aggregates many experts

The fine-grained routing of DeepSeek-V3 is not just a modeling choice. It is a systems choice: more experts means each GPU aggregates more of them, which smooths out load imbalance through statistical averaging. This is why fine-grained MoE is becoming the standard architecture.

Expert parallelism is the natural distribution strategy for MoE, but it introduces a communication pattern (all-to-all) and a load balancing problem that dense models never face. Simulating the traffic and imbalance before deploying saves you from discovering these constraints in production.

Tomorrow: dynamic disaggregation with NVIDIA Dynamo, where we explore splitting prefill and decode onto separate GPU pools and dynamically adjusting the split based on workload.