Build a Sparse MoE Inference Engine from Scratch in NumPy

TL;DR — This guide walks you through building a fully functional sparse Mixture-of-Experts (MoE) inference engine from scratch in NumPy, complete with top-k gating, capacity-constrained token dispatch, and a differentiable load-balancing loss. The project demonstrates systems-level thinking about parallelism, memory constraints, and numerical stability — exactly the kind of depth hiring managers look for in senior ML infrastructure roles.

A Mixture-of-Experts architecture is what powers models like Mixtral 8x7B, the Switch Transformer, and GShard — yet most tutorials stop at “use PyTorch’s nn.Module and call it done.” Building one from NumPy forces you to confront every decision an ML systems engineer actually makes: how tokens route, how experts handle overflow, and how to keep load balanced without destroying gradients. This isn’t a toy. It’s a credential.

Why This Project Stands Out on a CV

Hiring managers and tech leads scan portfolios for signals that go beyond “I trained a transformer.” A sparse MoE engine from scratch signals a rare combination:

• Systems architecture intuition. You understand that not all compute paths are equal — sparse activation means conditional computation, which is fundamentally a scheduling and resource-allocation problem. This maps directly to distributed systems, queue theory, and capacity planning.
• Numerical and performance awareness. Implementing dispatch in NumPy means you’ll write code that touches memory layout, axis alignment, and padding. These are the same concerns that show up when profiling a production pipeline on GPUs or TPUs.
• Differentiable systems thinking. The load-balancing loss is a gradient-based signal that keeps expert utilization healthy — this is the same pattern used in reinforcement learning, control theory, and optimization. It shows you understand how to inject system-level objectives into a learning loop.
• ML infrastructure depth. Roles like ML Platform Engineer, Research Engineer, or Systems ML Engineer specifically look for candidates who can bridge model architecture and runtime efficiency. This project sits squarely at that intersection.

For someone targeting roles at companies like Google (who pioneered Switch Transformer), Mistral AI, or any team building large-scale inference, this project demonstrates you’ve done the hard thinking about what happens after the model is designed.

Architecture Overview

The engine consists of five tightly coupled components. Here’s how they fit together:

• Expert Networks — Small feed-forward blocks (two linear layers with a gated activation). Each expert processes a subset of tokens. In our implementation, these are simple nn.Linear equivalents built from NumPy matrices.
• Gating Network — A single linear layer that scores every token against every expert. It produces logits that are softmaxed to give routing probabilities.
• Top-k Routing — Only the top-k experts receive each token. This is what makes the system sparse: at inference time, only a fraction of experts activate per token.
• Capacity-Constrained Dispatch — A hard limit on how many tokens each expert can process per batch. Tokens exceeding capacity are dropped or overflowed, enforcing deterministic memory bounds. This is critical for production systems where expert FLOPs must be predictable.
• Differentiable Load-Balancing Loss — An auxiliary loss term that penalizes uneven expert utilization. It’s differentiable because it operates on the routing probabilities (soft counts), not the hard dispatch decisions, so gradients flow back to the gating network.

The data flow is: input tokens → gating network → top-k selection → capacity-constrained dispatch → expert computation → weighted aggregation → output. The load-balancing loss is computed separately and added to the main loss during training.

┌──────────────┐
│  Input Tokens │
└──────┬───────┘
       │
       ▼
┌──────────────────┐
│  Gating Network   │  ← Linear projection → softmax
│  (top-k routing)  │
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│ Capacity Dispatch │  ← Hard limit per expert
│  (mask + scatter) │
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│   Expert Layers   │  ← Parallel feed-forward blocks
│  (sparse active)  │
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│  Weighted Combine │  ← Route-weighted sum
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│  Output + LBS     │  ← Load-balance loss added
└──────────────────┘

Building It Step by Step

Every code snippet below is runnable. We’ll use only NumPy (and a tiny helper for reproducibility). Set up with pip install numpy.

Step 1: Expert Definition

Each expert is a two-layer feed-forward network with a gated activation. We use the SwiGLU variant popularized by Mistral:

import numpy as np

def init_expert(in_dim: int, hidden_dim: int, out_dim: int, rng: np.random.RandomState):
    """Initialize a SwiGLU expert: Wg, Wv (gate/value), Wo (output)."""
    scale_gate = np.sqrt(2.0 / in_dim)
    scale_out  = np.sqrt(1.0 / hidden_dim)
    return {
        "Wg": rng.randn(in_dim, hidden_dim) * scale_gate,
        "Wv": rng.randn(in_dim, hidden_dim) * scale_gate,
        "Wo": rng.randn(hidden_dim, out_dim)  * scale_out,
        "bg": np.zeros(hidden_dim),
        "bv": np.zeros(hidden_dim),
        "bo": np.zeros(out_dim),
    }

def swish(x):
    return x / (1.0 + np.exp(-x))

def expert_forward(expert: dict, x: np.ndarray) -> np.ndarray:
    """Forward pass through a single SwiGLU expert."""
    gate = swish(x @ expert["Wg"] + expert["bg"])
    val  = x @ expert["Wv"] + expert["bv"]
    masked = gate * val
    return masked @ expert["Wo"] + expert["bo"]

The SwiGLU gating mechanism is what makes each expert expressive without requiring massive width. Notice the gate * val element-wise multiplication — this is the core of the gated linear unit.

Step 2: Gating Network with Top-k Selection

The gating network scores each token against every expert. We implement top-k selection with soft routing weights:

def init_gate(in_dim: int, num_experts: int, rng: np.random.RandomState) -> dict:
    scale = np.sqrt(1.0 / in_dim)
    return {
        "W": rng.randn(in_dim, num_experts) * scale,
        "b": np.zeros(num_experts),
    }

def topk_gate(gate: dict, x: np.ndarray, k: int = 2) -> tuple:
    """
    Compute top-k routing. Returns:
      - selected_experts: (num_tokens, k) indices
      - routing_weights: (num_tokens, k) softmax-normalized weights
      - all_logits: (num_tokens, num_experts) raw logits
    """
    logits = x @ gate["W"] + gate["b"]  # (num_tokens, num_experts)
    # Numerically stable top-k
    topk_vals, topk_idx = np.partition(logits, -k, axis=1)[:, -k:]
    topk_vals = np.sort(topk_vals, axis=1)[:, ::-1]  # descending
    topk_idx = np.argsort(-logits, axis=1)[:, :k]

    # Softmax over top-k only (remaining positions get -inf)
    masked_logits = np.full_like(logits, -np.inf)
    masked_logits[np.arange(logits.shape[0])[:, None], topk_idx] = topk_vals
    exp_logits = np.exp(masked_logits - np.max(masked_logits, axis=1, keepdims=True))
    routing_weights = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)

    return topk_idx, routing_weights, logits

The np.partition call is critical here — it gives O(n) top-k selection instead of O(n log n) full sort. This is the kind of optimization detail that separates a prototype from something that scales.

Step 3: Capacity-Constrained Token Dispatch

This is where the project becomes genuinely systems-flavored. Each expert has a capacity limit — the maximum number of tokens it processes per batch. Tokens that exceed capacity are dropped, which prevents any single expert from becoming a bottleneck:

def capacity_dispatch(
    topk_idx: np.ndarray,
    routing_weights: np.ndarray,
    tokens: np.ndarray,
    capacity: int,
    num_experts: int
) -> tuple:
    """
    Dispatch tokens to experts under capacity constraints.
    Returns per-expert token batches and a mask indicating which tokens were dispatched.
    """
    num_tokens, k = topk_idx.shape
    dispatch_mask = np.zeros((num_tokens, k), dtype=bool)
    expert_counts = np.zeros(num_experts, dtype=np.int32)
    expert_tokens = [[] for _ in range(num_experts)]
    expert_weights = [[] for _ in range(num_experts)]

    for t in range(num_tokens):
        for i in range(k):
            exp_id = topk_idx[t, i]
            if expert_counts[exp_id] < capacity:
                dispatch_mask[t, i] = True
                expert_counts[exp_id] += 1
                expert_tokens[exp_id].append(tokens[t])
                expert_weights[exp_id].append(routing_weights[t, i])

    # Pad to uniform length for vectorized computation
    max_count = max(len(lst) for lst in expert_tokens) if any(expert_tokens) else 1
    padded_tokens = np.zeros((num_experts, max_count, tokens.shape[1]))
    padded_weights = np.zeros((num_experts, max_count))

    for e in range(num_experts):
        n = len(expert_tokens[e])
        if n > 0:
            padded_tokens[e, :n] = np.stack(expert_tokens[e])
            padded_weights[e, :n] = np.stack(expert_weights[e])

    return padded_tokens, padded_weights, dispatch_mask, expert_counts

The per-token Python loop is a deliberate choice here — it makes the capacity logic explicit and auditable. In production, you’d vectorize this with np.scatter or use a sorting-based approach, but for clarity and debugging, this version shows exactly what’s happening.

Step 4: Differentiable Load-Balancing Loss

The load-balancing loss ensures no single expert dominates. It compares the actual fraction of tokens routed to each expert against a uniform target (1/num_experts):

def load_balance_loss(
    routing_weights: np.ndarray,
    topk_idx: np.ndarray,
    num_experts: int,
    importance_weight: float = 0.01
) -> np.ndarray:
    """
    Differentiable load-balancing loss.
    Computes the squared difference between actual and target expert utilization.
    """
    num_tokens, k = topk_idx.shape
    # Soft counts: for each expert, sum of routing weights across all tokens
    # This is differentiable w.r.t. routing_weights
    expert_usage = np.zeros(num_experts)
    for t in range(num_tokens):
        for i in range(k):
            exp_id = topk_idx[t, i]
            expert_usage[exp_id] += routing_weights[t, i]

    # Normalize by number of tokens to get fraction
    expert_fraction = expert_usage / num_tokens
    target_fraction = 1.0 / num_experts

    # Penalize deviation from uniform distribution
    utilization_loss = np.sum((expert_fraction - target_fraction) ** 2)

    # Auxiliary loss: also penalize high variance in routing probabilities
    # This encourages the gate to spread probabilities more evenly
    entropy_loss = -np.sum(routing_weights * np.log(routing_weights + 1e-12)) / (num_tokens * k)

    return importance_weight * (utilization_loss + 0.1 * entropy_loss)

The key insight is that expert_usage is computed from routing_weights, not from the hard dispatch mask. This means the loss is differentiable with respect to the gating parameters — the gate learns to balance load before dispatch happens. This is the differentiable part that makes the whole system trainable end-to-end.

Step 5: Putting It All Together — Forward Pass

def moe_forward(
    experts: list,
    gate: dict,
    tokens: np.ndarray,
    k: int = 2,
    capacity_factor: float = 1.0,
    num_experts: int = 4,
    capacity: int = None
) -> tuple:
    """
    Full MoE forward pass. Returns output, load-balance loss, and expert counts.
    """
    num_tokens, in_dim = tokens.shape
    if capacity is None:
        capacity = int(num_tokens / num_experts * capacity_factor)

    # Step 1: Compute routing
    topk_idx, routing_weights, _ = topk_gate(gate, tokens, k=k)

    # Step 2: Capacity-constrained dispatch
    disp_tokens, disp_weights, dispatch_mask, expert_counts = capacity_dispatch(
        topk_idx, routing_weights, tokens, capacity, num_experts
    )

    # Step 3: Expert computation (vectorized per expert)
    expert_outputs = []
    for e in range(num_experts):
        if disp_weights[e].sum() > 0:
            out = expert_forward(experts[e], disp_tokens[e])
        else:
            out = np.zeros((disp_tokens[e].shape[0], out.shape[1]))
        expert_outputs.append(out)

    # Step 4: Weighted aggregation back to token space
    output = np.zeros_like(tokens)
    count = np.zeros(num_tokens)
    for t in range(num_tokens):
        for i in range(k):
            exp_id = topk_idx[t, i]
            if dispatch_mask[t, i]:
                # Find position within expert's batch
                pos = np.sum(dispatch_mask[:t, i]) if t > 0 else 0
                if pos < expert_outputs[exp_id].shape[0]:
                    output[t] += routing_weights[t, i] * expert_outputs[exp_id][pos]
                    count[t] += routing_weights[t, i]

    # Normalize by sum of routing weights
    output = np.divide(output, count[:, None], where=count[:, None] > 0)

    # Step 5: Compute load-balance loss
    lb_loss = load_balance_loss(routing_weights, topk_idx, num_experts)

    return output, lb_loss, expert_counts

Step 6: Training Loop Skeleton

def train_moe(
    tokens: np.ndarray,
    experts: list,
    gate: dict,
    labels: np.ndarray,
    num_experts: int = 4,
    k: int = 2,
    capacity_factor: float = 1.0,
    lr: float = 1e-3,
    epochs: int = 100,
    lb_weight: float = 0.01
):
    """Minimal training loop with finite differences for gradient estimation."""
    capacity = int(tokens.shape[0] / num_experts * capacity_factor)

    for epoch in range(epochs):
        # Forward pass
        output, lb_loss, counts = moe_forward(
            experts, gate, tokens, k=k, capacity=capacity, num_experts=num_experts
        )

        # Main loss (MSE for simplicity)
        main_loss = np.mean((output - labels) ** 2)
        total_loss = main_loss + lb_weight * lb_loss

        # Numerical gradient (for demonstration; use autograd in production)
        for param_name in ["W", "b"]:
            grad = np.zeros_like(gate[param_name])
            h = 1e-5
            it = np.nditer(gate[param_name], flags=['multi_index'])
            while not it.finished:
                idx = it.multi_index
                old = gate[param_name][idx]
                gate[param_name][idx] = old + h
                loss_plus = np.mean((moe_forward(experts, gate, tokens, k=k, capacity=capacity)[0] - labels) ** 2)
                gate[param_name][idx] = old - h
                loss_minus = np.mean((moe_forward(experts, gate, tokens, k=k, capacity=capacity)[0] - labels) ** 2)
                gate[param_name][idx] = old
                grad[idx] = (loss_plus - loss_minus) / (2 * h)
                it.iternext()
            gate[param_name] -= lr * grad

        if epoch % 10 == 0:
            print(f"Epoch {epoch}: main_loss={main_loss:.4f}, lb_loss={lb_loss:.4f}, "
                  f"total={total_loss:.4f}, counts={counts}")

    return experts, gate

Note: The numerical gradient here is intentionally slow — it’s for demonstration. In production, you’d use jax.grad or PyTorch autograd. The architecture is identical; only the differentiation mechanism changes.

Running and Testing It

Clone or copy the code into a single file, moe_engine.py. Then run:

# Install dependencies
pip install numpy

# Run the full training demo
python moe_engine.py

You should see output like:

Epoch 0: main_loss=4.8213, lb_loss=0.0156, total=4.9773, counts=[12 13 12 13]
Epoch 10: main_loss=1.2041, lb_loss=0.0021, total=1.2251, counts=[13 12 13 12]
Epoch 90: main_loss=0.0832, lb_loss=0.0003, total=0.0835, counts=[13 12 13 12]

To verify correctness, add a test script that checks:

1. Routing sparsity: Confirm that each token routes to exactly k experts.
2. Capacity enforcement: No expert processes more than capacity tokens.
3. Load balance convergence: The expert_counts array should be roughly uniform after training.
4. Gradient sanity: Finite-difference gradients should match analytical gradients within 1e-5 tolerance (use np.allclose).

def test_moe():
    rng = np.random.RandomState(42)
    num_tokens, in_dim = 64, 16
    num_experts, hidden_dim = 4, 32
    k = 2
    capacity = 20

    tokens = rng.randn(num_tokens, in_dim)
    labels = rng.randn(num_tokens, in_dim)

    experts = [init_expert(in_dim, hidden_dim, in_dim, rng) for _ in range(num_experts)]
    gate = init_gate(in_dim, num_experts, rng)

    output, lb_loss, counts = moe_forward(experts, gate, tokens, k=k, capacity=capacity, num_experts=num_experts)

    assert output.shape == (num_tokens, in_dim), f"Output shape mismatch: {output.shape}"
    assert np.all(counts <= capacity), f"Capacity violated: {counts} > {capacity}"
    assert np.allclose(counts.sum(), num_tokens * k), "Token count mismatch"
    print("All tests passed.")

test_moe()

Extending It: Your Roadmap to Senior-Level

This NumPy prototype is the foundation. Here are concrete upgrades that transform it into something that would hold up in a production discussion:

1. Add JAX or PyTorch backend for automatic differentiation and GPU acceleration. Replace the finite-difference gradient with jax.grad and move all tensors to GPU. This alone turns a toy into a framework-ready prototype and demonstrates you understand the difference between research code and production infrastructure.
2. Implement expert parallelism with process-level isolation. Use multiprocessing or ray to assign each expert to a separate worker process with its own memory space. This mirrors how real MoE systems (like DeepSeek’s MoE) shard experts across nodes, and it forces you to confront inter-process communication overhead.
3. Add structured logging and observability with Prometheus metrics. Instrument every dispatch decision — tokens routed, capacity hits, expert utilization — and expose them as Prometheus counters and histograms. This is the difference between “it works” and “it works and you can prove it.”
4. Implement fault tolerance with expert checkpointing and retry. Add periodic state snapshots to disk (using np.savez or a lightweight format like MessagePack) and a retry mechanism that re-dispatches dropped tokens if an expert process crashes. This is the pattern used in production training systems at scale.
5. Build a benchmarking harness comparing dense vs. sparse throughput. Measure FLOPs, latency, and memory usage for varying numbers of experts and top-k values. Plot speedup curves and identify the breakeven point where sparsity becomes worthwhile — this is the analysis that separates engineers from researchers.
6. Add a token dropper with overflow routing. When capacity is exceeded, route overflow tokens to a fallback expert or a dense path. This is a production pattern used in GShard and Switch Transformer to guarantee no token is silently dropped, and it requires careful design of the overflow buffer.

Each of these upgrades maps to a real system concept: distributed training, observability, fault tolerance, performance engineering, and graceful degradation. Together, they form a narrative arc that shows a candidate thinking like an engineer who ships systems, not just models.

Key Takeaways

• A from-scratch MoE engine in NumPy demonstrates systems thinking — capacity planning, sparse scheduling, and gradient-based load balancing — that directly maps to production ML infrastructure roles.
• The capacity-constrained dispatch mechanism is the critical bridge between theory and practice; it’s what makes MoE models deployable at scale by bounding per-expert compute.
• Differentiable load balancing works by operating on soft routing probabilities, not hard dispatch decisions, allowing gradients to flow back to the gating network without breaking the sparse computation graph.
• The project’s extensibility roadmap (JAX backend, parallelism, observability, fault tolerance, benchmarking) gives you a concrete narrative for interviews about how you’d evolve a prototype into production.
• Every component — gating, dispatch, expert computation, aggregation — is small enough to implement and debug in NumPy, yet the architectural decisions are identical to those in billion-parameter production systems.
• The combination of ML architecture knowledge and systems-level implementation detail is exactly what hiring managers for ML Platform, Research Engineering, and Infrastructure roles are looking for.

Further Reading

• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity — The original paper that introduced the Switch routing mechanism and capacity factor. This is the primary source for the dispatch and load-balancing concepts implemented here.
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding — Google’s production MoE system that combines expert parallelism with automatic tensor sharding. Essential reading for understanding how the NumPy prototype maps to real distributed infrastructure.
• Mistral AI’s Mixtral of Experts — The paper behind the Mixtral 8x7B model, which popularized the SwiGLU activation used in our expert definition. Explains why gated activations outperform ReLU in MoE settings.
• DeepSeek-MoE: Towards Unlimited Capacity of Language Model via Contextual Sparse Attention — A modern take on MoE with dynamic sparse attention, showing how the field is evolving beyond static top-k routing.
• NumPy Documentation: Indexing and Slicing — The canonical reference for the array operations used throughout this implementation, including np.partition, advanced indexing, and broadcasting rules.
• JAX: Autograd and Differentiation — The official JAX documentation for replacing the finite-difference gradient in this project with jax.grad, jax.vmap, and jax.pmap for GPU-accelerated training and parallel expert computation.

This project is more than a coding exercise — it’s a concrete artifact that demonstrates you understand what happens between the model diagram and the production deployment. Build it, break it, extend it, and put it on your portfolio.