Working the roofline
for Kimi K3 on GB200

Three roofs, three different byte counts

A kernel can be limited by HBM. Distributed work can instead be limited by NVLink or the scale-out network. Divide work by the bandwidth of the resource that actually carries its bytes. In particular, a fast NVLink roof does not make cross-domain traffic disappear.

The roofline lab

The lab is an accounting model, not a launch-config generator. Every run has 512 GPUs and TP1. Changing rows or microbatches changes the global batch; changing CP changes the number of independent data replicas. Memory feasibility still needs its own budget.

Stored parameters are not activated FLOPs

For Y = XW, with X [M,K] and W [K,N], forward costs 2MNK FLOPs. Backward has two more contractions of the same size: dX = dYW^T and dW = X^TdY. Three contractions give 6MNK. Summed over learned matrices, this becomes 6N_activeD_tokens. Here N is activated matrix weights, not the entire 2.780T checkpoint.

Routed expert weights activated / token
  = 92 layers * 16 experts * 3 matrices * 3584 * 3072
  = 48,620,371,968

All learned active matrices = 104,174,092,288
Learned-matrix training work = 625.04 GFLOP / token

The familiar 104.1895B active-parameter count also includes small non-matmul weights. It is useful for a conventional metric, but those weights do not each cause a dense matmul. The input embedding lookup and text-only vision path are excluded from per-token matrix compute; the LM head is included.

Attention adds work without adding a weight matrix

Equal documents make S_eff easy to see: one 4K document gives 4,096; four independent 1K documents give 1,024. Two 2K rows have the same tokenwise projection work as one 4K row, but roughly half the MLA contractions. Flattening [rows, context] to T does not remove document boundaries. CP shortens local token arrays, not the logical context attended by MLA.

EP is payload and control plane

K3 compresses the routed stream from 7,168 to 3,584 channels before dispatch. A BF16 latent row is 7,168 bytes, not 14,336. The shared experts stay on the 7,168-wide dense path and are not dispatched. With full checkpointing, the route is traversed three times: forward, recomputed forward, and the two reverse communication operations in backward.

No checkpoint: 4 payload exchanges / layer / microbatch
Full checkpoint: 6 payload exchanges / layer / microbatch

Assignment budget, uniform routing:
  remote GPU copies / token = top_k * (1 - 1 / EP)
  cross-domain copies      = top_k * (1 - local_EP_ranks / EP)

Payload bytes = exchanges * local_tokens * latent_width * 2 * remote_copies

One-copy-per-assignment is a conservative payload budget. If the dispatcher sends a token once per destination rank and sums results before returning them, multiple chosen experts on that rank can share a copy. The second accounting option models this ideal deduplication under uniform sampling without replacement; it does not assert that every backend or phase implements it. Metadata, alignment and routing weights are additional.

The extra HybridEP all-gather is not an expert weight gather

The 4K reference exchanges a Boolean routing map [4,096, 896] before dispatch: 3.5 MiB locally, 224 MiB gathered across EP64, of which 220.5 MiB came from other ranks. The original and recomputed forward each do it. On a 12-MoE-layer stage, 64 microbatches × 12 layers × 2 = 1,536 metadata all-gathers per update. Small inputs can thus imply a substantial global output and a repeated dependency.

FSDP counts depend on the lifecycle, not just the acronym

At PP8/EP64, routed experts have expert FSDP1: there is no routed-weight FSDP gather. Dense weights are sharded over 64 ranks. A BF16 ring all-gather receives 2P(1 − 1/n) bytes per rank; an FP32 reduce-scatter has the corresponding 4P(1 − 1/n) budget. Multiply by the actual number of collectives.

For a hierarchical collective, a domain with a of n shards need only import P(1 − a/n) parameters once, then share them locally. Dividing its import by a models the per-rank share of an a-NIC aggregate budget. This is an ideal domain-level network bound, not the measured wire volume of NCCL's chosen algorithm. NIC sharing or a different rank layout changes it.

Long-context CP scenarios keep the logical attention length, DP count and local-token shape straight, but CP attention communication is not priced here. KDA state/prefix exchange, MLA KV/head redistribution and backend-specific CP collectives require their own trace-derived ledger. The 64K preset is a shape and compute comparison, not a complete long-context step-time prediction.

Ordinary 1F1B has a bubble, not a ban on overlap

Under balanced stage costs, an ideal ordinary 1F1B schedule has compute-bubble overhead (PP − 1) / M relative to useful scheduled work. At PP8 and M64 that is 10.94% overhead, or 9.86% of the compute-plus-bubble interval. Each optimizer update fills and drains its own pipeline; a previous warmup update does not remove those structural bubbles.

More microbatches help amortize the bubble, but there is a tradeoff. At fixed global batch, each microbatch gets smaller, potentially lowering GEMM efficiency. In this lab the row shape stays fixed when M changes, so the global batch grows. Science determines whether that larger update is appropriate.

The PP packet is bigger than [T, 7168]

K3 sends (hidden_TD, delta_TND). With one stage per rank, a microbatch visits each receiving rank for the first time, so the residual cache cannot omit the earlier blocks. At the default 12-layer split, outgoing block counts are 2, 4, 6, 8, 10, 12 and 14. The final hop carries one hidden stream plus 14 residual streams: 15 × 56 MiB = 840 MiB forward, and a matching gradient payload backward.

Interleaving with multiple virtual stages can shrink a stage's gathered working set, revisit a rank with a populated residual cache, and offer different scheduling opportunities. It also adds stage boundaries and communication. It is not modeled by simply dividing the 1F1B bubble or memory by VP. The lab deliberately keeps one contiguous stage per rank.

Muon is a matrix workload of its own

FP32 momentum is persistent training history. Newton–Schulz uses a BF16 working matrix. For one expert matrix [3,584, 3,072], orient the working matrix as X [r,c] with r = 3,072 and c = 3,584. A standard iteration forms A = XX^T, B = bA + c_coefA², then updates X using BX and a scaled X. The leading contractions cost 4r²c + 2r³ FLOPs per iteration.

Five iterations cost about 0.966 TFLOP per logical expert matrix. An interior PP rank with 12 MoE layers and 14 experts per layer has 12 × 14 × 3 = 504 such matrices: about 487 TFLOP per update, before the smaller non-routed matrices, redistribution, normalization, pointwise work and communication. That is a workload estimate, not an optimizer timing measurement.

Fused gate/up storage still contains two logical expert matrices. Muon orthogonalizes those matrices separately; using one wider fused storage shape would give a different Newton–Schulz cost and a different update.

A long CPU optimizer span can include waiting for the pipeline's last backward or a gradient-norm collective. The reference's complete exposed optimizer tail was small compared with its clean step; it does not support blaming the throughput gap on Muon alone. More tokens per update amortize one optimizer update, but also change optimization and memory requirements.

A roofline gap is a question, not a diagnosis

The fixed 512-GB200 HybridEP references below completed 12 updates with FP32 resident/master weights and optimizer states, BF16 compute, full checkpointing, PP8/FSDP64/EP64, 64 microbatches, and no CUDA graphs. Clean throughput excludes the profiled update and trace-export disturbance. These measurements stay fixed when you change the lab.

MFU here uses a fixed 2,500-TFLOP/s dense denominator and the documented TorchTitan estimator, not a 5,000-TFLOP/s sparse ceiling. The 4K data are packed C4, so using a full 4K document for every row in the calculator is an attention-work upper scenario, not an audited document-length histogram. The learned-routing reference also changes HybridEP to blocking mode; its slowdown is not a controlled routing-imbalance-only measurement.

A concrete discrepancy: metadata wait versus byte service

In the inspected middle-stage trace, all 1,536 same-shape routing-map all-gathers were effectively exposed relative to useful model compute. Their kernel residency times varied far more than a constant-byte model predicts. NCCL residency includes peer waiting and contention; it is not a direct bandwidth sample.

Metadata calls overlapping FSDP all-gathers had a 31.464-ms median, versus 1.029 ms without that overlap. This is evidence of correlated communication delay, not proof of one cause: rank arrival skew, contention and earlier compute stalls can coincide. Summing overlapping NCCL bars would double-count wall time; measure interval unions and the wait before the consumer can start.

The lab omits CP-specific communication, optimizer service from its main step sketch, backend chunk work, pointwise kernels, CPU/launch cost, collective startup, imbalance, allocator effects and non-ideal overlap. Its resource-plus-bubble number is an optimistic diagnostic sketch, not a speedup promise, validated prediction, or complete training speed of light.