MoE GroupGEMM padding under expert skew
MoE GroupGEMM padding under expert skew cannot be predicted from total tokens alone. Two layers can route the same number of tokens and need different amounts of padded tile work.
The input is the full per-expert token histogram. For a tile width q, each active expert rounds its token count up to a multiple of q.
Smaller tiles reduce that rounding. Larger tiles can reuse expert weights better, so the tile with the least padding is not automatically the fastest kernel.
MoE GroupGEMM padding under expert skew
An MoE router sends each token to a small subset of experts. The grouped expert kernel receives a vector of token counts rather than one rectangular batch.
The DA-MoE paper shows why that vector matters. Current tuning often chooses a fused kernel from total-token buckets while ignoring how those tokens are distributed.
Tensor Core kernels execute fixed MMA tiles. When an expert receives fewer token columns than a tile covers, the remaining columns do no useful model work.
Large tiles reduce instruction count and improve weight reuse. They also waste more compute on small expert matrices.
The best choice moves with batch size and routing skew. One static kernel per total-token bucket ignores that difference.
Fanout's MoE batch and EP guide derives the average assignments per expert. Padding requires the individual counts.
Calculate padded token columns
Let t_i be the tokens assigned to expert i, and let q be the kernel's token tile width.
The executed token columns for expert i are q times the ceiling of t_i divided by q. An inactive expert with t_i equal to zero contributes zero.
Sum those rounded counts across experts. Call the result P. The useful token assignments are U, the sum of the unrounded t_i values.
P minus U is the padded token count. The padding fraction is (P minus U) divided by P.
This is a planning model for one tiled dimension. Real fused kernels also differ in K tiles, pipeline depth, load path, quantization, and scheduling.
Still, the calculation catches a common error: using only U and assuming every routing histogram with that total presents the same work.
Work a near-uniform routing example
Take eight local experts with token counts [17, 17, 17, 17, 16, 16, 16, 16]. The layer has 132 useful expert assignments.
With q equal to 8, each 17 rounds to 24 and each 16 stays 16. P is 160, so 28 token columns are padding.
The padding fraction is 28 divided by 160, or 17.5%.
With q equal to 16, each 17 rounds to 32 and each 16 stays 16. P becomes 192, with 60 padded columns and a 31.25% padding fraction.
With q equal to 32, all eight experts consume one 32-column tile. P is 256, so 124 columns are padding, or 48.4%.
The total token count did not change. Only the tile changed, and the useful share of executed columns fell from 82.5% to 51.6%.
This is why small batches with broad routing often favor small token tiles.
Contrast a concentrated histogram
Now keep U at 132 but route [100, 32, 0, 0, 0, 0, 0, 0]. Only two experts are active.
At q equal to 8, 100 rounds to 104 and 32 stays 32. P is 136, so padding is 2.9%.
At q equal to 32, 100 rounds to 128 and 32 stays 32. P is 160, so padding is 17.5%.
The 32-column tile still wastes more compute than the 8-column tile. Its penalty is much smaller than in the near-uniform example because most work sits behind one hot expert.
That hot expert can reuse its weight tiles across many token columns. The larger tile now has a better chance of repaying its padding through fewer loads and higher arithmetic intensity.
The paper observed this shift on Blackwell. Higher routing concentration moved the preferred Tile-N toward larger values, while small broad batches favored smaller tiles.
Skew does not always help. A hot expert can create a slow rank, and several tiny nonzero experts can still force extra tiles.
Summarize skew without losing the histogram
DA-MoE defines effective experts as the square of total assignments divided by the sum of squared per-expert counts.
The metric equals the full expert count under uniform routing and approaches one when a single expert carries nearly all tokens.
For the near-uniform example, effective experts is almost eight. For [100, 32, 0, 0, 0, 0, 0, 0], it is about 1.58.
This scalar is useful for grouping similar workloads and generating controlled test distributions.
It is not a replacement for exact rounding. Two histograms with the same effective-expert score can straddle different tile boundaries and produce different P values.
Use effective experts to choose candidate kernels or benchmark buckets. Use the original token counts to calculate padding and validate the selected tactic.
The NVIDIA hardware-friendly design guide supplies the related average: global concurrency times top-k divided by expert count.
Record both. The average describes expected matrix height, while effective experts describes concentration.
Grouped execution does not require padded experts
GroupedGEMM is broader than one padding strategy. A kernel can accept variable per-expert M sizes and schedule only valid tiles.
Meta's Llama 4 GroupedGEMM design accepts an m_sizes vector with one entry per group.
Its persistent kernel builds work partitions from each expert's actual M and output width. Zero-sized experts can skip weight loads, and padded tokens at the end can be excluded.
That removes host-side rectangular padding between expert batches. It does not remove padding inside a fixed MMA tile when an expert's remaining rows or columns do not fill the tile.
Keep those layers of padding separate in profiling:
- capacity padding inserted before the grouped operation
- tail padding inside the selected MMA tile
- communication padding used by an all-to-all or graph-safe buffer
A single "padding percent" counter can mix all three and point tuning at the wrong component.
The cuBLAS grouped GEMM API groups different matrix sizes in one launch.
The matrix list stays ragged even though launch overhead is shared.
Padding is not the whole cost model
The DA-MoE paper distinguishes physical throughput from effective throughput.
Physical throughput counts all MMA work, including zero-padded columns. Effective throughput scales it by the fraction that contributes to valid outputs.
A large tile can report higher physical throughput while delivering lower effective throughput because the Tensor Cores process more zeros.
The reverse can happen in a memory-bound case. A small tile saves padded arithmetic but reloads the same expert weights more often.
Kernel choice must price at least four terms: padded MMA work, expert-weight traffic, launch and scheduling overhead, and communication around the expert computation.
Software pipeline depth and the memory load path matter too. The paper measured more than a 10% spread among configurations sharing the same Tile-N.
That result rules out a selector based only on the padding formula. Use it to filter and explain candidates, then benchmark the candidates on the deployed GPU.
Rank tails can dominate layer latency
Expert parallelism distributes experts across ranks. Routing skew therefore creates both a kernel-shape problem and a placement problem.
Suppose one rank owns the 100-token expert while another owns two 16-token experts. Their local padding, weight traffic, and execution time differ even when global padding looks acceptable.
The layer waits for the required dispatch, expert work, and combine path. A fast aggregate kernel does not hide the slowest rank.
Fanout's EPLB guide explains how placement can separate hot experts or replicate them.
EPLB can reduce the rank tail without changing the router's global histogram. Tile selection can improve local kernels without fixing a rank that owns several hot experts.
Profile padding and latency per rank, then report the maximum beside the sum. Aggregate effective throughput alone can hide a repeated tail straggler.
The NVIDIA wide-EP report also ties GroupGEMM gains to fewer local experts and more aggregate bandwidth.
Those system gains come with all-to-all cost. A tile improvement is not an end-to-end win if the wider placement moves the bottleneck to communication.
Adapt kernels without adding a CPU stall
A runtime selector needs the routing histogram after the router produces it on the GPU.
Copying that histogram to the CPU, choosing a kernel, and launching it again can erase the gain for small decode batches.
DA-MoE sorts and normalizes the per-expert count vector, matches it against offline-tuned exemplars, and dispatches a kernel through conditional CUDA Graphs.
Its reported HumanEval-X traces improved fused-MoE geomean latency by 1.16 times on DeepSeek-V3 and 1.29 times on Kimi K2.
Peak improvements were larger, but they belong to the tested models, hardware, kernels, and traces. They are not a generic reward for computing effective experts.
The useful design lesson is the control path. Distribution-aware selection must remain on the GPU or make its host round trip cheap enough to measure.
Add guardrails for neighboring histograms. If small count changes switch kernels every layer, dispatch can become noisy even when either candidate performs similarly.
Build a routing-shaped benchmark
Collect per-layer, per-rank token counts from real prefill and decode traces. Preserve the correlation between total tokens, layer position, and routing skew.
For every candidate tile or fused kernel, calculate P, padding fraction, active experts, and effective experts.
Then measure kernel latency, effective throughput, physical throughput, HBM bytes, L2 hit rate, and achieved Tensor Core utilization.
Keep the same expert placement while comparing kernels. Otherwise an EPLB move can look like a tile-selection gain.
Run a separate placement sweep after choosing a small kernel set. Report slowest-rank time and all-to-all duration with the local GroupGEMM measurements.
Include low-token decode buckets. Under one tested distribution, the paper favored Tile-N 8 at 32 and 64 tokens, then Tile-N 64 at 1,024 tokens.
Also include concentrated and broad histograms at the same U. Random uniform routing is not enough to tune a production selector.
To calculate MoE GroupGEMM padding under expert skew, round every active expert count to the candidate tile, then sum.
Use that number to expose wasted work. Choose the kernel only after measuring the counterweight from weight reuse, pipeline configuration, rank tails, and communication.