Week 14 — Mixture of Experts (MoE)

Learning Objectives

Key Concepts

Setup & Configuration

# Verify GPU memory: need 2× H100 80GB
nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader

# Download Mixtral weights (~93 GB disk space required)
python -c "
from huggingface_hub import snapshot_download
snapshot_download('mistralai/Mixtral-8x7B-Instruct-v0.1')
"

# Quick memory estimate before launching
python -c "
# 46.7B params * 2 bytes BF16 / 1e9 = 93.4 GB
print(f'Expected weight memory: {46.7e9 * 2 / 1e9:.1f} GB')
print(f'Per GPU with TP=2: {46.7e9 * 2 / 2 / 1e9:.1f} GB + KV cache')
"

Experiments

vllm serve mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --tensor-parallel-size 2 \
  --port 8000 \
  --disable-log-requests &
sleep 60  # Mixtral takes longer to load

for rate in 1 2 4 8; do
  python benchmarks/benchmark_serving.py \
    --backend vllm \
    --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
    --dataset-name sharegpt \
    --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
    --num-prompts 100 \
    --request-rate $rate \
    --port 8000 \
    2>&1 | tee results_moe_tp2_rr${rate}.txt
done
kill %1

# Dense baseline: Llama-3.1-8B with TP=2
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --tensor-parallel-size 2 \
  --port 8010 \
  --disable-log-requests &
sleep 30

python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 100 --request-rate 4 \
  --port 8010 \
  2>&1 | tee results_dense_tp2_rr4.txt
kill %1

# MoE: Mixtral-8x7B with TP=2 (same request rate)
vllm serve mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --tensor-parallel-size 2 \
  --port 8011 \
  --disable-log-requests &
sleep 60

python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 100 --request-rate 4 \
  --port 8011 \
  2>&1 | tee results_moe_tp2_rr4_compare.txt
kill %1

# expert_probe.py — log expert selection frequencies
from vllm import LLM, SamplingParams
import torch, collections

llm = LLM(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    tensor_parallel_size=2,
)

# Patch router to capture selections
expert_counts = collections.defaultdict(lambda: collections.Counter())

def hook_router(layer_idx, expert_indices):
    for idx in expert_indices.flatten().tolist():
        expert_counts[layer_idx][idx] += 1

prompts = ["Explain quantum computing in simple terms."] * 50
outputs = llm.generate(prompts, SamplingParams(max_tokens=128))

# Print expert load balance per layer
for layer, counts in sorted(expert_counts.items()):
    total = sum(counts.values())
    print(f"Layer {layer}:", {k: f"{v/total:.1%}" for k,v in counts.most_common(8)})

# SGLang server for Mixtral
python -m sglang.launch_server \
  --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --tp 2 \
  --port 8002 &
sleep 90

python -m sglang.bench_serving \
  --backend sglang \
  --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 100 --request-rate 4 \
  --port 8002 \
  2>&1 | tee results_moe_sglang_tp2.txt
kill %1

Experiment Results

Hardware

Mixtral-8x7B-Instruct-v0.1 runs on 2× NVIDIA H200 141GB HBM3e at TP=2. Dense 8B comparisons run on 2× H100/A100/H200. All models in BF16.

⭐ Mixtral-8x7B MoE — Offline Throughput (TP=2, 2× H200)

Mixtral reads ~6× more weight bytes per decode step than Llama-8B despite only activating 2/8 experts. Decode is memory-bandwidth bound, so the 3.6× throughput gap directly reflects the total-parameter ratio, not the activated-parameter ratio.

⭐ Mixtral-8x7B — Latency (TP=2, bs=1, 2× H200)

At bs=1, Mixtral is only 1.57× slower (9.21 vs 5.87 ms/tok) — the gap narrows because activated FLOPs (12.9B) are closer to Llama's 8B. At higher batch sizes, memory-bandwidth dominates and the gap widens toward 3.6×.

⭐ Mixtral-8x7B — Online Serving (TP=2, 2× H200)

At 1 req/s: TTFT ~1s, ITL 8.1 ms (near bs=1 decode latency). At 4 req/s: TTFT 1.7s, ITL 13.2 ms as batching increases concurrent decodes.

Llama-3.1-8B Throughput: TP=1 vs TP=2 (offline, 150 prompts)

An 8B dense model fits comfortably in one 80 GB H100, so adding a second GPU doesn't unlock any compute the first one was missing — the only effect is NCCL all-reduce overhead, which roughly cancels the (negligible) extra parallel work. The TP=2 numbers being identical to TP=1 is the expected and correct result for a model this small. The lesson: you only TP when memory or peak GEMM tile size forces you to.

Llama-3.1-8B Latency vs Batch Size (TP=2)

Per-token latency only grows from 7.01 ms → 8.49 ms when batch jumps 1 → 32 (a 1.21× cost) yet aggregate throughput jumps 26.4×. This is the canonical 'decode is memory-bandwidth bound' shape — the model weights are read from HBM only once per step regardless of batch size, so amortizing across more concurrent requests is essentially free. The lesson is to push batch size up until KV cache fills, not until per-token latency degrades.

Charts

Figure 1: Left — TP=1 vs TP=2 throughput (no speedup at 8B). Right — Per-token latency vs batch size on 2× H100 (sub-linear growth = memory-bandwidth limited).

Expected vs Actual

Metrics to Collect

Source Code Reading

# Simplified MoE forward (from mixtral.py)
def forward(self, hidden_states):
    batch_size, seq_len, hidden_dim = hidden_states.shape
    hidden_states = hidden_states.view(-1, hidden_dim)  # [T, H]

    # Step 1: Router — compute logits for all 8 experts
    router_logits = self.gate(hidden_states)  # [T, 8]
    routing_weights = F.softmax(router_logits, dim=-1)

    # Step 2: Top-2 selection
    routing_weights, selected_experts = topk(routing_weights, k=2)
    routing_weights /= routing_weights.sum(dim=-1, keepdim=True)

    # Step 3: Token permutation — group tokens by expert
    # Step 4: Batched expert computation (fused CUDA kernel)
    # Step 5: Weighted sum of expert outputs
    return final_hidden_states

Core Source Code Walkthrough

Below are real excerpts from vllm-continuum that implement the concepts you measured. Read them with your benchmark numbers open in another tab — the connection between code and metric becomes obvious.

vllm/model_executor/models/mixtral.py:102 — MixtralMoE.forward

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    orig_shape = hidden_states.shape
    hidden_states = hidden_states.view(-1, self.hidden_size)
    # router_logits: (num_tokens, n_experts)
    router_logits, _ = self.gate(hidden_states)
    final_hidden_states = self.experts(hidden_states, router_logits)
    return final_hidden_states.view(orig_shape)

Each token gets routed to top-k experts (k=2 for Mixtral). The gate linear layer produces routing scores; self.experts (a FusedMoE op) dispatches each token to its assigned experts and computes their outputs. The key property: only k of the 8 experts run per token, so compute is ~k/n × dense — but ALL experts must still be loaded into HBM (no sparsity in memory).

Written Analysis — Reference Answers

Mixtral was unavailable in our offline HF cache (see Hardware callout above), so the answers below combine arithmetic on Mixtral's published architecture with our measured Llama-3.1-8B TP=1 vs TP=2 numbers and the 70B TP=2 numbers from Week 13.

Q1: Mixtral memory vs Llama-3.1-70B memory

Q2: Is expert routing load-balanced?

Q3: Expert parallelism (EP) vs tensor parallelism (TP) for MoE

Q4: vLLM vs SGLang on Mixtral