Dropless MoE, in depth: sort and permute dispatch, grouped and block-sparse GEMM, variable all-to-all, worst-case memory and stragglers

What dropping costs and what dropless changes

With capacity-based routing, each expert gets C = capacity_factor * tokens * k / num_experts slots. Tokens beyond C for an expert are dropped for that expert. Three costs follow. First, quality: a dropped token gets no update from that expert on the forward pass and contributes no gradient to it on the backward pass, and the drops concentrate on exactly the experts the router likes most. Second, padding: when an expert receives fewer than C tokens, the remaining slots are padded and still multiplied, so a capacity factor of 1.25 spends up to 25 percent more expert FLOPs than the routed work. Third, coupling: results depend on batch composition, because whether your token is dropped depends on which other tokens shared its batch. That makes evaluation and inference behave differently from training unless you are careful.

Raising the capacity factor trades drops for padding; there is no setting that removes both. Dropless escapes the trade by letting each expert's row count be whatever the router produced. The MegaBlocks paper (Gale, Narayanan, Young and Zaharia, MLSys 2023) made this practical on GPUs by expressing the expert computation as block-sparse matrix multiplication, and reported end-to-end training speedups of up to 40 percent over Tutel-trained MoEs and 2.4x over dense models trained with Megatron-LM. Dropless is now the default in Megatron-Core: its MoE documentation states that the default strategy neither drops nor pads tokens, and that setting --moe-expert-capacity-factor switches to capacity-based dropping. DeepSeek-V3 also reports that it drops no tokens in training or inference.

The dropless forward pass, step by step

Take a flattened batch of T tokens with hidden size d, E experts and top-k routing. The router produces, per token, k expert ids and k combine weights. That gives T times k assignments, each one a (token, expert, weight) triple.

Step one is to sort the assignments by expert id. A stable argsort over the flattened expert ids yields a permutation, and a bincount yields how many rows each expert owns. Step two is the permute: gather the token rows in sorted order so each expert's inputs sit in one contiguous slice. Step three computes all experts at once on those slices. Step four is the unpermute: scatter each output row back to its source token, multiply by the combine weight, and sum the k contributions per token.

Nothing in this pipeline has a capacity. The permuted tensor has exactly T times k rows, every one of which is real work. What changes from step to step is how those rows are partitioned among experts, and that partition is data, not shape.

A reference implementation

The reference below is correct and readable, and deliberately not fast. It is the version to test optimised kernels against.

import torch

def dropless_moe(x, router_w, experts, k=2):
    """x: [T, d]; router_w: [d, E]; experts: list of E modules d -> d."""
    probs = (x @ router_w).softmax(dim=-1)            # [T, E]
    top_w, top_e = probs.topk(k, dim=-1)              # [T, k]
    top_w = top_w / top_w.sum(dim=-1, keepdim=True)   # renormalise over chosen k

    flat_e = top_e.reshape(-1)                        # [T*k]
    order = flat_e.argsort(stable=True)               # group assignments by expert
    src_tok = order // k                              # token each sorted row came from
    counts = torch.bincount(flat_e, minlength=len(experts))

    xs = x.index_select(0, src_tok)                   # permute: [T*k, d]
    outs = [experts[e](chunk)                         # variable rows per expert
            for e, chunk in enumerate(xs.split(counts.tolist()))]
    ys = torch.cat(outs) * top_w.reshape(-1)[order].unsqueeze(-1)

    y = torch.zeros_like(x)
    y.index_add_(0, src_tok, ys)                      # unpermute and combine
    return y, counts

Two details matter. The stable sort keeps the token order within each expert deterministic, which makes runs reproducible and debugging tractable. And counts.tolist() copies the counts to the host, which forces the CPU to wait for the GPU. In a training step that stall serialises the launch queue at every MoE layer. Production implementations keep counts on the device and hand them straight to a grouped kernel, or overlap the copy with other work. When a profile shows idle gaps before each expert block, look for exactly this sync.

Grouped GEMM and block-sparse kernels

The per-expert Python loop launches E small matrix multiplies, several per expert once you count the up, gate and down projections. With 64 or 256 fine-grained experts, launch overhead and under-filled GPUs dominate. Two kernel strategies fix this.

Grouped GEMM. One kernel launch computes a list of independent multiplications, each with its own row count, sharing a weight layout. The kernel reads a small array of per-group offsets (the cumulative counts) and maps thread blocks to tiles across all groups. CUTLASS provides grouped GEMM, and Megatron-Core exposes it for MoE through --moe-grouped-gemm. Waste comes only from the last partial tile of each group.

Block-sparse formulation. MegaBlocks views the whole expert layer as one large matrix multiply whose weight is block-diagonal: the stacked expert weights, with each expert's row range sized to its token count, rounded up to a block multiple. Its kernels then compute only the non-zero blocks, using a sparse layout rebuilt each step from the counts. The rounding is the only padding, and it is bounded by one block per expert rather than by a capacity factor.

Both approaches share the property that matters: the work done is proportional to the routed rows plus a small per-expert rounding term. Which one wins depends on the library you are already using; few teams write either kernel themselves. On the inference side, engines implement MoE as fused kernels that likewise align each expert's rows to the kernel tile instead of dropping, so serving is dropless in practice even for models trained with capacity limits.

Dropless under expert parallelism

Under expert parallelism the experts live on different GPUs and tokens travel to them by all-to-all, as described in MoE all-to-all communication and expert parallelism. With capacity routing, every message has the same size, C rows per expert, so the all-to-all is uniform and can be launched without coordination. Dropless breaks that.

Each rank now sends a different number of rows to every other rank, so the collective becomes a variable-split all-to-all. Before the payload moves, the ranks must exchange their split sizes: a small all-to-all of counts, then the large one of rows. In PyTorch the large exchange is torch.distributed.all_to_all_single with input_split_sizes and output_split_sizes lists, which are host integers. That is a second source of device-to-host sync, and it sits on the critical path of every MoE layer, forward and backward.

The count exchange also turns imbalance into time. In a uniform all-to-all every rank finishes its expert GEMMs at about the same moment. In a dropless one, the rank hosting the most popular expert receives the most rows and becomes the straggler that every other rank waits for at the return all-to-all. Dropless does not remove the cost of imbalance; it moves it from lost tokens to lost time. That is why dropless systems still run balance losses or bias-based balancing, and why Megatron-Core's --moe-token-dispatcher-type alltoall path is typically paired with them.

Memory: planning for the worst rank

Capacity routing has a fixed activation footprint per expert, C rows times d. Dropless has a footprint that depends on the router. In the worst case one expert receives one row from every token in the batch, and under expert parallelism the rank holding it can receive rows from every token on every rank in the expert-parallel group. The bound is T times EP-size rows for one expert, where T is tokens per rank, against an average of T times k divided by experts per rank.

Nobody provisions for the literal worst case; you provision for an observed high percentile plus headroom, and you make overflow survivable. Practical controls are:

• Record the maximum per-rank received rows per layer every step, and alert on its trend, not just on out-of-memory crashes.
• Keep balance regularisation on, so the tail of the received-rows distribution stays close to the mean.
• Use activation recomputation for expert MLPs; their intermediate activations scale with received rows and dominate the spike.
• Leave allocator headroom, because variable shapes fragment the caching allocator. Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True reduces fragmentation from shapes that change every step.
• Keep a capacity-factor fallback configured but disabled, so a run that starts hitting out-of-memory can be resumed with a generous cap instead of restarted from scratch.

Worked example: one skewed batch

Take one rank with 4096 tokens, 8 experts and top-2 routing: 8192 assignments, 1024 per expert if perfectly balanced. Suppose the router produces loads of 1600, 1300, 1100, 1000, 900, 800, 700 and 792.

Capacity factor 1.0. C is 1024. Experts 0, 1 and 2 overflow by 576, 276 and 76, so 928 assignments, 11.3 percent, are dropped. The five under-full experts pad up to 1024, so the GEMMs process 8 times 1024 = 8192 rows, of which 7264 are real work.

Dropless with grouped GEMM. All 8192 rows are computed, nothing is dropped and nothing is padded except partial tiles. Compute is identical to the capacity case on a single device.

Dropless with block rounding. As an illustration, assume 128-row blocks. Rounding each expert up gives 1664, 1408, 1152, 1024, 1024, 896, 768 and 896, a total of 8832 rows, 7.8 percent above the routed work. The overhead depends on the block size your kernel uses and shrinks as per-expert counts grow.

Dropless under expert parallelism, one expert per rank. The rank holding expert 0 computes 1600 rows while a balanced rank computes 1024, so the layer takes about 1.56 times as long as a balanced one, because everyone waits for the slowest rank. That 1.56 is the number to drive down with balancing; the 11.3 percent drop rate is the number dropless removes.

Failure modes

• Out-of-memory spikes late in training. Router collapse onto a few experts inflates one rank's received rows. Watch per-rank maximum rows and the balance loss together; a rising balance loss usually precedes the crash.
• Host syncs everywhere. Counts copied to the CPU for splits or Python loops stall the launch queue at every layer. Profile for gaps before expert blocks.
• Non-determinism. An unstable sort or atomic scatter-adds in the unpermute make runs irreproducible. Use stable sorts and deterministic combine paths when debugging loss spikes.
• Fragmentation. Changing shapes leave holes in the caching allocator, so reserved memory grows while allocated memory does not.
• Train-serve mismatch. A model trained with drops and served dropless sees tokens at inference that its experts never trained on in that context. Evaluate with the routing mode you will serve.
• Graph capture breaks. CUDA graphs and some compilers need static shapes. Dropless layers either stay outside the captured region or are padded to a fixed upper bound for capture, which reintroduces a cap.

Trade-offs

What to do next

1. Run the reference implementation against your optimised MoE layer on random inputs and compare outputs and gradients to tolerance.
2. Log per-expert counts, drop rate (if any) and per-rank maximum received rows for every layer, every step.
3. Profile one training step and remove host syncs from count handling and split sizes.
4. Measure straggler cost: time of the slowest expert-parallel rank divided by the median.
5. Set allocator options and activation recomputation, then size headroom from the 99.9th percentile of received rows.
6. Evaluate checkpoints with the routing mode you will serve, and keep a disabled capacity-factor fallback ready.