Dropless MoE Is a Runtime Problem as Much as a Model Design

Dropless MoE keeps every routed token instead of enforcing a fixed expert capacity. That removes token loss and avoids padding, but it exposes the runtime to variable-length expert batches. The router can produce a different token count for every expert on every step, so the central problem is not merely sparse arithmetic. It is coordinating irregular memory movement, matrix multiplication, synchronization, and numerical scaling without turning dynamic metadata into host-side control flow.

The proposed grouped GEMM path addresses the local compute problem. Rather than launching one matrix multiplication per expert or padding all experts to a common capacity, it supplies a collection of expert shapes to a single implementation. This can reduce launch overhead and eliminate invalid padded work. The boundary condition is important: grouped execution does not make small or badly imbalanced matrices efficient automatically. If a few experts receive most tokens, Tensor Core occupancy, workspace use, and tail latency can still be governed by the worst shapes.

Avoiding device-to-host copies of expert token counts is equally significant in JAX. Host-visible dynamic shapes can disrupt CUDA graph capture and trigger recompilation or synchronization. A runtime that keeps routing metadata on device preserves a more stable execution graph, but it also increases the need for correct bounds, overflow detection, and memory planning. The supplied MaxText configuration therefore includes a ragged buffer factor and periodic overflow checks; those controls should be treated as operational safeguards, not proof that arbitrary routing skew is free.

Expert parallelism moves tokens to the GPUs holding their selected experts and later restores their original placement. Fusing dispatch and combine can reduce intermediate writes and synchronization. Token deduplication may reduce network bytes when Top-K routes send the same token to multiple experts on a destination scope. Its benefit depends on routing locality and parallelism layout: if selected experts are widely distributed, replication and metadata management can offset some savings. Communication measurements must therefore report payload volume, message count, and overlap, rather than only end-to-end step time.

MXFP8 changes the tradeoff from pure kernel efficiency to numerical policy. Block-scaled low precision can increase arithmetic throughput and reduce movement, but routing creates many independently sized expert groups whose activation distributions may differ. Group-aware scaling must be validated for overflow, underflow, accumulation precision, and convergence. A faster short run is insufficient evidence; the relevant test is whether loss curves, routing balance, gradient statistics, and final task quality remain comparable to a BF16 reference under the same optimizer and data order.

Host offloading exchanges device memory pressure for host bandwidth, latency, and synchronization. It is useful when saved activations prevent a run from fitting, but it can become harmful when PCIe or fabric transfers occupy the critical path. The right decision depends on activation lifetime, rematerialization cost, host-memory capacity, and whether transfers overlap with independent work. XLA multistreaming has a similar boundary: concurrent collectives can hide latency, yet stream concurrency can raise contention, memory consumption, or deadlock risk if replica groups and dependencies are mis-modeled.

A reproducible evaluation should isolate each contribution. Start with a small dropless MoE whose routing and outputs are checked against a simple reference implementation. Record per-step latency after compilation, compile time, peak device and host memory, expert token histograms, grouped-GEMM time, dispatch and combine time, collective time, network bytes, and achieved overlap. Compare padded capacity-based, dropless baseline, grouped GEMM only, communication fusion only, precision changes only, and the complete stack. Then repeat across balanced and adversarial routing distributions, different Top-K values, sequence lengths, expert counts, node counts, and failure or overflow conditions. Report hardware, software revisions, seeds, warm-up policy, profiler traces, convergence metrics, and cost assumptions before accepting the vendor’s throughput claims as transferable.