After studying this material, you should be able to:

1. Explain what Mixture-of-Experts (MoE) layers are and why they create bottlenecks in distributed training
2. Understand the tradeoffs between push-based vs pull-based GPU communication
3. Describe how computation and communication can be overlapped efficiently
4. Explain the megakernel design pattern and why it eliminates CPU-GPU synchronization overhead
5. Analyze how minibatch/macrobatch sizing affects GPU utilization

---

Step-by-Step Teaching

Step 1: What is a Mixture-of-Experts (MoE) Layer?

The Core Idea

Instead of every token passing through the same neural network, MoE routes each token to only a subset of specialized sub-networks called experts.

Token → Router → [Expert 3, Expert 7] → Weighted Sum → Output
                  (only 2 of 256 experts activated)

The Math

For each token x, the MoE layer computes:

MoE(x) = SharedExpert(x) + Σ(router_weight_i × RoutedExpert_i(x))
                             i ∈ top-k selected experts

Each routed expert runs a standard Feed-Forward Network (FFN):

Expert_i(x) = Down_i( SwiGLU( Up_i(x), Gate_i(x) ) )

Where:

• Up projection: expands dimension from d_model → d_expert
• Gate projection: runs in parallel with Up
• SwiGLU: activation function combining both
• Down projection: compresses back d_expert → d_model

Key Parameters

---

Step 2: Why Does MoE Create a Distributed Computing Problem?

Expert Parallelism (EP)

With 256 experts and 64 GPUs, each GPU holds only 4 experts. This is called Expert Parallelism (EP).

GPU 0: Experts 0-3
GPU 1: Experts 4-7
GPU 2: Experts 8-11
...
GPU 63: Experts 252-255

The Token Routing Problem

A token on GPU 0 might be assigned to Expert 7 (on GPU 1) and Expert 200 (on GPU 50). This means:

Step 1: DISPATCH  → Send token from GPU 0 → GPU 1 and GPU 50
Step 2: COMPUTE   → Run FFN on GPU 1 and GPU 50
Step 3: COMBINE   → Send results back to GPU 0
Step 4: WEIGHTED SUM → Combine expert outputs

Why This Is Slow (The Bottleneck)

Running dispatch → compute → combine sequentially wastes time:

Timeline (naive):
|--Dispatch--|--Compute--|--Combine--|
             ↑           ↑
         GPU idle    GPU idle
         waiting     waiting

Communication can take as long as computation itself, so sequential execution wastes up to 50% of time.

---

Step 3: The Solution — Overlapping Communication and Computation

The Pipelining Idea

Instead of waiting for ALL tokens to arrive before computing, process tokens in chunks (minibatches):

Timeline (pipelined):
Comms SMs: |--Dispatch chunk 1--|--Dispatch chunk 2--|--Combine chunk 1--|--Combine chunk 2--|
Comp  SMs:                      |----FFN chunk 1----|----FFN chunk 2----|

This is called inter-SM overlapping: different groups of Streaming Multiprocessors (SMs) handle different tasks simultaneously.

Two Types of SMs

Total SMs on GPU
├── Comms SMs (~1/3 of SMs)  → Handle NVLink transfers (dispatch/combine)
└── Comp SMs  (~2/3 of SMs)  → Handle FFN computation (tensor cores)

Key insight: TMA (Tensor Memory Accelerator) loads can saturate NVLink bandwidth using fewer than 1/3 of SMs, leaving the majority free for computation.

---

Step 4: Push vs Pull Communication — Choosing the Right Direction

This is one of MoK's most important innovations. There are two ways GPU A can transfer data to GPU B:

Push-Based Communication

GPU A (has data) → actively WRITES → GPU B's memory

• GPU A controls the transfer
• Fewer total bytes transferred
• Problem: GPU B must wait for signals from up to 71 peers before it knows data arrived
• Problem: Requires complex scheduling table with 3 columns: {src_index, dst_rank, dst_index}

Pull-Based Communication

GPU B (needs data) → actively READS → from GPU A's memory

• GPU B controls the transfer
• More total bytes (protocol overhead in both directions)
• Advantage: No cross-GPU signaling needed — GPU B knows data arrived when the load completes
• Advantage: Simple 2-column schedule: {src_rank, src_index}

The Bandwidth Paradox

You might expect push to be faster (fewer bytes), but:

NVLink has SEPARATE lanes for each direction:
→ direction: 900 GB/s
← direction: 900 GB/s
Total: 1800 GB/s (only if BOTH directions are busy)

Push sends almost everything in one direction → other direction sits idle → wastes half the bandwidth

Pull splits traffic between both directions → both lanes stay busy → up to 29% higher bandwidth utilization under expert imbalance

Signaling Latency Comparison

MoK's Hybrid Strategy

MoK cleverly mixes both approaches to get the best of each:

Bonus: The schedule table is built once and reused for all 4 operations, taking less than 3% of total MoE runtime.

---

Step 5: Choosing the Right Minibatch Size

The Granularity Tradeoff

Too fine-grained (tiny chunks):
|D|C|D|C|D|C|D|C|  ← Many small barriers, tensor cores never fully saturate

Too coarse-grained (huge chunks):
|----Dispatch all----|----Compute all----|----Combine all----|  ← Long idle periods

Just right:
|--D--|--D--|--D--|
      |--C--|--C--|--C--|
            |--Cb-|--Cb-|--Cb-|

The Wave Concept

A wave = one round of concurrent execution across ALL SMs on the GPU.

• If a GEMM has too few tokens → partial wave → some SMs sit idle
• Target: at least 2 full waves per GEMM per minibatch

Why 2 waves?

1. First wave fully saturates tensor cores
2. Second wave's matrix multiplications overlap with first wave's epilogue (SwiGLU, next GEMM)

The Formula

On Blackwell GPUs, each SM works on a 128×128 output tile for optimal tensor core utilization.

For up + gate projections (run in parallel):

tokens_per_minibatch × d_ff ≥ 2 × num_SMs × 128²

For down projection (runs alone):

tokens_per_minibatch × d_model ≥ 2 × num_SMs × 128²

Combined requirement:

tokens_per_minibatch ≥ 2 × num_SMs × 128² / min(d_model, d_ff)

Worked Example (Kimi 2.5)

d_model = 7168
d_ff    = 2048  ← smaller, so this is the bottleneck
num_SMs = 132   (Blackwell GPU)

tokens_per_minibatch ≥ 2 × 132 × 128² / 2048
                     ≥ 2 × 132 × 16384 / 2048
                     ≥ 2 × 132 × 8
                     ≥ 2112

This matches the benchmark data showing performance peaks around 2048-2560 tokens:

---

Step 6: Eliminating CPU-GPU Synchronization with Ring Buffers

The Dynamism Problem

The router decides at runtime how many tokens go to each GPU. You don't know in advance.

Existing solutions and their problems:

Why CPU Sync Is Especially Bad on NVL72

GB300 NVL72 architecture:
├── 72 × GPU (very fast)
└── 72 × Grace CPU (integrated, relatively slow)

Problem: GPU streams catch up to CPU work → GPU sits completely idle
         waiting for CPU to push the next kernel launch

The Ring Buffer (Macrobatch) Solution

Instead of one large buffer (mostly empty), use a fixed-size circular ring buffer:

Ring Buffer (few hundred MB):
┌─────┬─────┬─────┬─────┬─────┬─────┐
│ S0  │ S1  │ S2  │ S3  │ S4  │ S5  │  ← Slots
└─────┴─────┴─────┴─────┴─────┴─────┘
  ↑                             ↑
 Write                         Read
 head                          head

Key principle: Drain each slot as early as possible so it can be reused.

Dispatch-Combine Interleaving

Without interleaving (bad):

Macrobatch 1: |--Dispatch all--|--Compute all--|--Combine all--|
Macrobatch 2:                                                  |--Dispatch--|...
              ↑ Long gap between macrobatches!

With interleaving (good):

Macrobatch 1 Dispatch: |--D1--|--D2--|--D3--|
Macrobatch 1 Compute:         |--C1--|--C2--|--C3--|
Macrobatch 1 Combine:                |--Cb1-|--Cb2-|--Cb3-|
Macrobatch 2 Dispatch:               |--D1--|--D2--|--D3--|  ← reuses same buffer slots!

Combine for macrobatch 1 and dispatch for macrobatch 2 use the same ring buffer slots, so the buffer is refilled as soon as it's emptied.

Reversed Ring for Backward Pass

During backward, some activations saved in the ring buffer get overwritten and must be replayed (recomputed).

Naive order (ascending macrobatches):

Forward:  MB1 → MB2 → MB3 → MB4
Ring:     [MB4 overwrites MB1, MB2 data]
Backward: Must replay MB1, MB2 from scratch

Reversed ring (descending macrobatches):

Forward:  MB4 → MB3 → MB2 → MB1
Ring:     [MB1 is last written, ring is full or contains all tokens]
Backward: Minimal replay needed

---

Step 7: The Megakernel Design

What Is a Megakernel?

Traditional GPU programming launches many separate kernels:

CPU launches: [Dispatch kernel] → [FFN kernel] → [Combine kernel] → ...
              ↑ overhead        ↑ overhead       ↑ overhead

A megakernel fuses everything into one persistent kernel:

CPU launches: [Single megakernel — runs entire MoE layer]
              SMs internally coordinate via shared counters

Why MoK Needs a Megakernel

1. Minibatching/macrobatching without megakernel = many kernel launches = high overhead
2. Inter-SM overlapping requires exact SM allocation — multiple CUDA streams with green contexts proved unreliable; software partitioning gives exact guarantees

Cluster Launch Control (CLC) for RDMA Overlap

During training, inter-rack communication (InfiniBand/RoCE for FSDP all-gather) must overlap with the megakernel.

Without CLC:
|----MoK megakernel----|----IB transfer----|  ← serialized, slow

With CLC:
|----MoK megakernel----|
              |--IB--|  ← megakernel yields to higher-priority stream

CLC is a Blackwell hardware feature that allows the persistent megakernel to yield SMs to higher-priority streams without fully terminating.

---

Step 8: MXFP8 Support and Determinism

MXFP8 Precision

MoK supports two precision modes:

MXFP8 overhead: Tensors must be quantized before tensor cores can use them.

MoK minimizes this by:

1. Pre-quantizing weights with an optimized kernel (done once, reused)
2. Fusing activation quantization into dispatch, GEMMs, and SwiGLU

Exception: Shared expert stays in BF16 for training stability.

Full Determinism

MoK fixes the order of floating point operations so:

Same input → Bitwise-identical output
             regardless of hardware scheduling

This matters for:

• Internal ablations (comparing runs fairly)
• On-policy RL post-training (reproducibility critical)

---

Step 9: Putting It All Together — The Complete MoK Forward Pass

┌─────────────────────────────────────────────────────────────┐
│                    MoK Megakernel                           │
│                                                             │
│  Comms SMs:                                                 │
│  1. Build schedule (pull-based, <3% of runtime)             │
│  2. Pull dispatch chunk 1 from remote GPUs                  │
│  3. Signal comp SMs: "chunk 1 ready"                        │
│  4. Pull dispatch chunk 2...                                │
│  5. (After all dispatches) Push combine chunk 1             │
│  6. Push combine chunk 2...                                 │
│                                                             │
│  Comp SMs:                                                  │
│  1. Run shared expert FFN (overlaps with first dispatch)    │
│  2. Wait for signal: "chunk 1 ready"                        │
│  3. Run FFN (Up+Gate → SwiGLU → Down) on chunk 1           │
│  4. Signal comms SMs: "chunk 1 ready for combine"           │
│  5. Run FFN on chunk 2...                                   │
│                                                             │
│  Ring Buffer: Slots freed as combine drains them            │
│  Macrobatch interleaving: Next macrobatch dispatches        │
│                           reuse freed slots immediately     │
└─────────────────────────────────────────────────────────────┘

---

Step 10: Results Summary

MoE Layer Benchmarks (Single NVL72 Rack, EP=64)

End-to-End Production Training (512 GPUs, Multiple NVL72 Racks)

---

Concept Map Summary

MoE Bottleneck
│
├── Problem: Communication as slow as computation
│   └── Solution: Overlap via inter-SM pipelining
│
├── Problem: Push vs Pull tradeoff
│   └── Solution: Hybrid (pull dispatch, push combine)
│       ├── Better bandwidth utilization (29% higher)
│       └── 5.8x lower signaling latency
│
├── Problem: Minibatch size affects GPU utilization
│   └── Solution: Formula-based sizing (≥2 waves per GEMM)
│
├── Problem: Dynamic token counts require CPU sync
│   └── Solution: Ring buffer (macrobatch)
│       └── Dispatch-combine interleaving keeps buffer flowing
│
├── Problem: Multiple kernel launches = overhead
│   └── Solution: Megakernel with CLC for RDMA overlap
│
└── Result: 1.41x end-to-end speedup, 2.37x MoE layer speedup