PaperlayerView exact revision on arXiv

Source-linked AI summary

Launch-Bound and Substitutable: Why Three Inference Optimizations Fail to Pay Off in Mixture-of-Experts Models

Gokulakannan Sakthivel, Jerry Wu, Amogh Rajendra, Giriprasad Radhakrishnan

arXiv:2608.26612v1cs.PF

TL;DR

MoE routing is data-dependent, leaving the practical effects of standard inference optimizations and routing changes insufficiently examined. The paper measures system performance and routing perturbations across three MoE models, using expert-selection interventions to separate routing from weight-error effects. It finds that fused kernels, compilation, and quantization provide limited end-to-end or quality benefits because execution is launch-bound and experts are substitutable.

  • Problem

    The practical effects of inference optimizations under data-dependent MoE routing, including how quantization changes expert selection, remain insufficiently examined.

  • Method

    The paper evaluates system metrics and routing perturbations across three MoE models, then replays quantized expert choices with full-precision weights to isolate routing effects.

  • Results

    Fused kernels reach 5.6x–9.0x in isolation but 0.999x end to end, while zero graph breaks are three times slower than eager and routing explains 2.7% of INT4 degradation.

  • Takeaways & Limitations

    MoE inference is launch-bound, and the experts are substitutable, so isolated kernel speedups and routing fidelity do not reliably translate into end-to-end performance or output quality.

  • Takeaways & Limitations

    The causal replay’s 2.7% attribution is limited to OLMoE, while the broader results use one quantization library and uncontrolled capacity comparisons.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) models route each token to a few of many expert networks, and that routing is data- dependent in a way standard inference optimizations do not ex- pect. This paper measures what three of them actually deliver on OLMoE-1B-7B, DeepSeek-V2-Lite, and Qwen3-30B-A3B. Fused Triton kernels reach 5.6x to 9.0x in isolation but 0.999x end to end against a measured 1.07x ceiling, because the model spends its time waiting on roughly a thousand kernel launches per forward pass rather than on the arithmetic those kernels improve. INT4 quantization changes on average 0.53 of the eight selected experts per token position, yet replaying exactly those changed routes through full-precision weights reproduces only 2.7% of the quality loss, which makes the experts substitutable rather than specialized. Removing all 23 torch.compile graph breaks, the step prior work treats as the structural fix, makes the model three times slower. A fourth result ties the three together: leaving the routers in FP16 lowers drift by 20% while raising loss, so routing fidelity and output quality are separable objectives. Every number recomputes from committed per-token route dumps.

I. INTRODUCTION … B. Custom GPU Kernels

MoE inference optimizations struggle because token-varying routing violates assumptions behind fixed-shape kernels, small-perturbation quantization, and graph compilation. This work measures system effects and routing changes across three differently configured MoE models, including custom kernels around routing.

  • I. INTRODUCTION: Dynamic, data-dependent routing selects different experts per token, undermining fixed-shape kernels, small-perturbation quantization assumptions, and torch.compile’s operating assumptions.Conditional, token-varying activation is presented as the root cause of the optimization difficulties studied.
  • II. LITERATURE SURVEY: Prior deployment work schedules around per-device forward-pass cost through heterogeneous-device collaboration, serverless expert placement, or overlapped expert computation and communication.The cited approaches are MoE-squared, Liu et al.’s serverless distribution, and SP-MoE.
  • I. INTRODUCTION: The study evaluates latency, throughput, memory, and four routing metrics while intervening with forced full-precision routing to test whether drift causes quality changes.The routing metrics are Routing Similarity, Jaccard Drift, Overlap@k, and Selection Shift.
  • I. INTRODUCTION: OLMoE-1B-7B, DeepSeek-V2-Lite, and Qwen3-30B-A3B vary in expert-pool size, selection density, layers, and gate-probability normalization.OLMoE uses 64 experts and top-8 routing across 16 MoE layers; DeepSeek-V2-Lite uses 64 routed plus 2 shared experts and top-6 across 26 layers; Qwen3 uses 128 experts and top-8 across 48 layers.
  • I. INTRODUCTION: 2.7% of INT4’s quality loss is explained by routing changes in a causal route-replay intervention, while eliminating all 23 torch.compile graph breaks costs a factor of three in latency.The contribution list also describes a control verifying that the intervention machinery is neutral.
  • A. MoE Architecture: Top-k gating enables capacity scaling without proportional compute, but the resulting conditional, token-varying activation creates the optimization difficulties examined here.The three studied models inherit this sparse activation design and differ in how aggressively they sparsify.
  • B. Custom GPU Kernels: The custom-kernel target is RMSNorm and gate softmax around routing, because data-dependent dispatch produces output shapes that standard Triton tiling cannot express.The approach follows fusion work such as FlashAttention and MegaBlocks but narrows the target to dense operations surrounding the routing decision.

C. Quantization … III. METHODOLOGY

The methodology tests quantization and graph compilation on real MoE checkpoints rather than assuming standard optimization heuristics transfer. It directly measures routing changes, causal quality impact, and whether eliminating graph breaks improves latency.

  • C. Quantization: INT4 cuts weight memory roughly four-fold against FP16, enabling a 30B-parameter MoE model to fit on one 80GB device instead of two.The compression also raises the batch size that fits alongside the KV cache.
  • C. Quantization: Quantization can change the selected expert set because gate-logit perturbations interact with discrete top-k selection.Perturbations below the smallest logit margin change nothing, while larger ones can change the expert set entirely.
  • C. Quantization: The study measures routing-topology change directly and then tests whether that change matters, extending beyond prior perplexity or dense-layer outlier analyses.This motivation distinguishes the paper’s routing-stability analysis from the cited prior work.
  • D. Graph Compilation: torch.compile typically gains 10 to 50% on dense transformers, but MoE routing creates graph breaks from CUDA-tensor conditionals and dynamic-shape operators.TorchDynamo traces Python bytecode and TorchInductor lowers the traced result.
  • D. Graph Compilation: A zero-graph-break configuration is three times slower on the real checkpoint, falsifying the assumption that minimizing graph breaks guarantees maximum speed.The benchmark re-measures this implication rather than assuming it from prior framing.
  • E. Why We Chose This Approach: The three optimization families are studied independently under identical model configurations and hardware so their results are directly comparable.Each sub-study uses its own measurement methodology: kernel microbenchmarks with Amdahl analysis, quantization routing metrics with causal intervention, and compilation tracing with latency sweeps.
  • E. Why We Chose This Approach: The design cross-checks the sub-studies, linking kernel launch-bound scaling with the broader evaluation of MoE optimization behavior.The supplied methodology passage identifies launch-bound scaling as measured in Sub-study 1.

A. Models Under Study · B. Platform and Software · C. Sub-study 1: Triton Kernel Engineering

The study fixes FP16 checkpoints, software versions, hardware stages, and measurement procedures before testing fused Triton kernels across three differently routed MoE models. Its design isolates implementation effects while preserving routing drift against a genuine baseline.

  • A. Models Under Study: Three open MoE checkpoints with deliberately different routing configurations were loaded from pinned FP16 revisions to establish reference routing.Quantized variants came from those same checkpoints rather than separately published quantized releases.
  • A. Models Under Study: The gate implementation determines quantization exposure: OLMoE and Qwen use nn.Linear routers, whereas DeepSeek-V2-Lite keeps its gate as an FP16 nn.Parameter.DeepSeek therefore experiences upstream perturbation only, while OLMoE also has gate-weight error.
  • A. Models Under Study: Each model was quantized at load time by bitsandbytes from its own FP16 checkpoint, keeping one algorithm and codepath across operating points.This avoids confounding model comparisons with different calibration data or quantization algorithms.
  • B. Platform and Software: Every measurement ran on Modal in a single-A100 serverless container, with persistent volumes retaining model caches and completed results.Containers did not outlive their stages, preventing measurements from inheriting a device warmed by a previous stage.
  • B. Platform and Software: The three drift runs, replay intervention, and both kernel stages used A100 SXM4, while the 17-configuration sweep used A100 PCIe.The SXM4 platform is quoted at 2039 GB/s HBM2e bandwidth, 312 TFLOP/s peak FP16 tensor throughput, and a 156 FLOP/byte ridge point.
  • B. Platform and Software: Transformers version 5 or later silently quantizes attention alone because MoE experts become packed three-dimensional parameters beyond that version.On Qwen3-30B-A3B, this produced zero replaced Linear layers, so images were pinned below version 5 and expert-weight replacement became a hard failure condition.
  • C. Sub-study 1: Triton Kernel Engineering: Two fused Triton kernels target RMSNorm and router softmax, with FusedRMSNorm reducing stock PyTorch’s three HBM-reading and writing kernels to one SRAM-resident pass.The kernels were installed by monkey-patching, and the softmax path intercepted only tensors shaped [batch, num_experts].
  • C. Sub-study 1: Triton Kernel Engineering: 1000 warm-up and 5000 timed iterations defined isolated medians, whereas end-to-end tests used 50 warm-up and 200 timed forward passes.Peak memory was measured with torch.cuda.max_memory_allocated, and the Amdahl ceiling used the optimized-operation fraction f and isolated speedup s.

D. Sub-study 2: Routing Drift and Its Causal Role

This sub-study measures routing drift against FP16 references and tests its causal role through full-precision replay of quantized routing decisions. It also normalizes drift across models with different expert-selection widths.

  • Measurement: 119,952 token positions per configuration were recorded from every MoE gate across 100 MMLU-style prompts at each precision.Hooks captured selected expert indices by layer and token position, using FP16 as the reference for Table 4 metrics.
  • Drift normalization: 2/(k+1) measures a single swapped expert, equaling 0.222 at top-8 and 0.286 at top-6.Raw Jaccard drift is therefore not directly comparable across models with different k; multiplying selection shift by k estimates swapped experts per token.
  • Causal intervention: The central causal test replays quantized expert selections through the FP16 model while keeping every weight at full precision.If routing changes cause output damage, this replay should reproduce most of the quantized model’s quality loss.
  • Causal intervention: The FP16-route control must exactly return the baseline, separating genuine routing effects from artifacts introduced by the replay hooks.Quality is evaluated using negative log-likelihood over the recorded token positions.

E. Sub-study 3: Graph Breaks and Compilation … 1) Isolated Performance:

The study measures graph breaks on the full OLMoE checkpoint and compares five identically shaped inference configurations. Isolated Triton kernels deliver substantial RMSNorm speedups but only a 2.01x gain for gate softmax, with both kernels remaining memory-limited.

  • E. Sub-study 3: Graph Breaks and Compilation: Graph breaks were counted on the real 16-layer OLMoE checkpoint at full hidden size using torch._dynamo.explain.Tracing does not execute kernels, so the count is device-independent and requires no GPU.
  • E. Sub-study 3: Graph Breaks and Compilation: Five configurations were timed at identical shapes: eager, eager with Triton kernels, default torch.compile, dynamic-shape capture, and dynamic-shape capture with kernels.The cuDNN attention backend was disabled because Inductor’s tensor layouts make it fail outright.
  • 1) Isolated Performance:: 5.62x to 8.98x: fused RMSNorm speedup rose with hidden dimension from 512 to 4096.At hidden dimension 4096, the kernel reached 1620 GB/s, versus 9% for the stock implementation.
  • 1) Isolated Performance:: 1620 GB/s: the fused RMSNorm kernel reached 79% of the SXM4 part’s 2039 GB/s at hidden dimension 4096.Fusing the two passes removes a third of the memory traffic.
  • 1) Isolated Performance:: 2.01x: fused gate softmax reduced runtime from 0.0173 ms baseline to 0.0086 ms.The gain does not improve with size because a 64-wide row dispatches a single warp with most lanes idle and occupancy near 45%.
  • 1) Isolated Performance:: Below 1.0 FLOP/byte: both gate-softmax kernels had arithmetic intensities far under the 156 FLOP/byte ridge point.Neither kernel was anywhere near compute-limited.

2) Share of the Forward Pass: … 1) Overall Drift:

The optimized kernels cannot improve end-to-end inference because their targeted operations occupy only 7.87% of the forward pass and launch overhead dominates. Quantization also changes routing substantially, chiefly at marginal expert ranks, while repeated runs at fixed precision remain identical.

  • 2) Share of the Forward Pass:: 7.87% of the forward pass is occupied by RMSNorm and router softmax, yielding a 1.072x Amdahl ceiling.RMSNorm accounts for 7.70% and router softmax for 0.17%.
  • 3) End to End:: 0.999x is the patched model’s end-to-end result at seq 512, batch 4, the only configuration directly comparable to the 1.072x ceiling.At seq 128, batch 1, the patched model is 21.5% slower; the largest configuration reaches 1.027x.
  • 3) End to End:: 32-fold more work raises baseline latency only from 246.8 ms to 383.4 ms, a factor of 1.55, indicating launch-dominated rather than arithmetic-bound execution.The comparison spans 128 to 4096 tokens.
  • B. Routing Drift Results: 0.53 of eight selected experts change per token position under INT4, versus 0.22 under INT8, and their drift intervals are far from overlapping.The drift estimates use OLMoE’s native top-8 routing and intervals bootstrapped over prompts rather than token rows.
  • 1) Overall Drift:: 0.1142 Jaccard drift at top-8 corresponds to a 0.0660 selection shift across eight slots, typically replacing one expert among eight.A minority of affected tokens replace two experts.
  • 1) Overall Drift:: The changed experts are overwhelmingly ranked 8th or 9th, where the logit margin is thinnest, while the top few ranks are essentially never disturbed.This describes the routing changes associated with quantization drift.
  • 1) Overall Drift:: Repeated passes at a fixed precision are bit-identical, separating quantization-induced routing changes from nondeterministic rerouting.The supplied passage explicitly reports bit-identical repeated passes at fixed precision.

2) Layer Coverage: · 3) Which Knobs Matter: · 4) Does Drift Cause the Damage?:

Across layer coverage and quantization settings, drift rises or shifts substantially without tracking quality loss, while routing replay attributes only 2.7% of degradation to expert-choice changes. The results instead identify quantization weight error as the dominant source of damage and show that experts are substitutable.

  • 2) Layer Coverage:: Drift rises monotonically from 0.0626 at two layers to 0.1142 at sixteen, while almost all NLL penalty is paid by layer four.Quantizing only the first N of 16 layers makes coverage a dial: drift keeps increasing, but quality saturates early.
  • 2) Layer Coverage:: The coverage dial separates drift from quality, providing the first evidence that drift is not the mechanism of damage.The two quantities do not move together as more layers are quantized.
  • 3) Which Knobs Matter:: FP4 drifts 34% more than NF4 at identical bit width, making the 4-bit data type the dominant quantization knob.NF4’s advantage reflects its better fit to normally distributed weights; double quantization and FP32 compute change drift only in the fourth decimal place.
  • 3) Which Knobs Matter:: The INT8 outlier threshold changes drift tenfold, from +0.0050 to +0.0535, while the library default of 6.0 gives the lowest loss of the four.The threshold is the activation magnitude above which a column is held in FP16 instead of quantized.
  • 3) Which Knobs Matter:: An activation-threshold setting moves routing by 8% and output quality by an order of magnitude, undermining routing as the damage mechanism.The effect is not monotonic across the four threshold settings.
  • 4) Does Drift Cause the Damage?:: The control reproduces the FP16 baseline to all printed digits, confirming that the replay machinery adds nothing of its own.This establishes the validity of the intervention used to separate routing effects from weight-error effects.
  • 4) Does Drift Cause the Damage?:: Forcing full-precision weights to follow INT4’s expert choices costs +0.002332 NLL versus INT4’s +0.086642, assigning 2.7% of degradation to routing and 97.3% to weight error.The replay uses INT4’s expert choices with full-precision weights, directly isolating routing from quantization error.
  • 4) Does Drift Cause the Damage?:: Experts are substitutable: sending a token to its 9th-best expert instead of its 8th-best costs very little, so drift is real but almost entirely harmless on its own.This reframes the sub-study: expert-choice changes explain little of the observed quality degradation.

5) Drift, Gate Divergence, and Quality: · 6) Router Exemption:

Drift predicts quality loss nearly as well as gate KL divergence, but their high collinearity and route-replay intervention show that drift tracks damage without causing it. Router exemption reduces drift while worsening NLL, separating routing fidelity from output quality.

  • 5) Drift, Gate Divergence, and Quality:: +0.907 Pearson and +0.891 Spearman correlations show drift predicts dNLL quality loss across 16 configurations.Gate KL divergence predicts the same target with +0.855 Pearson and +0.900 Spearman correlations.
  • 5) Drift, Gate Divergence, and Quality:: +0.978 Pearson correlation makes drift and gate KL divergence 98% collinear, preventing the sweep from separating them.The two predictors therefore provide nearly redundant explanations of the observed quality relationship.
  • 5) Drift, Gate Divergence, and Quality:: 2.7% of the effect is attributed by causal route replay, so drift tracks gate perturbation and weight-error damage without causing the damage.The correlation and intervention are therefore compatible rather than contradictory.
  • 6) Router Exemption:: Table 12 quantizes everything except the 16 router modules to test whether protecting routers improves routing fidelity.This directly evaluates router exemption as an intervention rather than relying only on correlations.
  • 6) Router Exemption:: 20% drift reduction at NF4, from 0.1140 to 0.0910, coincides with worse NLL, from +0.0872 to +0.0979.The same inversion also appears at INT8.
  • 6) Router Exemption:: Exempting the routers removes about one fifth of drift, while upstream perturbations account for the other four fifths.The result attributes only about a fifth of drift and upstream perturbation to the router’s own quantization.
  • 6) Router Exemption:: Routing fidelity and output quality are separable objectives that can move in opposite directions under quantization.Protecting routers lowers drift but worsens NLL in the tested configuration.

7) Across Architectures: · 8) Task Accuracy:

Across architectures, corrected expert-swap counts show that finer selection granularity drifts more, but capacity and router precision confound cross-model comparisons. Task accuracy shows no distinguishable drop at the available evaluation budget, so the paper reports NLL for causal analysis; routing fidelity and output quality can nevertheless diverge.

  • 7) Across Architectures:: 0.4686 versus 0.5281 corrected swaps per token overturn the raw-Jaccard comparison of DeepSeek and OLMoE, because top-6 selection inflates the same physical event.DeepSeek’s raw Jaccard is 0.1303 versus OLMoE’s 0.1142, but cross-model claims must use the corrected quantity.
  • 7) Across Architectures:: Qwen’s 128 experts at top-8 drift most on every measure, consistent with finer selection granularity producing more drift.The passage states that all six intervals within a precision are disjoint.
  • 7) Across Architectures:: Spearman correlation between drift and quality-loss ordering is +1.00 at INT8 but only +0.50 at INT4.The rankings establish an ordering rather than a fit, and capacity is not controlled.
  • 7) Across Architectures:: 2.7% of INT4’s NLL increase is recovered by replaying its routes through FP16 weights, leaving 97.3% attributable to weight error.The control returns exactly zero, confirming that the replay machinery is neutral.
  • 7) Across Architectures:: DeepSeek’s unquantized gate confounds its apparent advantage, making that advantage an implausible basis for architectural conclusions.The INT4 disagreement falls precisely on the 7B-versus-30B pair, while the unquantized gate is a plausible mechanism for DeepSeek’s advantage.
  • 8) Task Accuracy:: No task-accuracy drop in Table 14 is distinguishable from zero at the evaluation budget, with the largest reaching about 0.8 sigma.NLL over 119,952 token positions has lower variance than a few hundred multiple-choice outcomes, motivating its use for causal results.
  • 8) Task Accuracy:: At both precisions, exempting the 16 router modules from quantization lowers Jaccard drift while raising NLL, separating routing fidelity from output quality.The opposing movements occur when everything except the router modules is quantized.

C. torch.compile Results · 1) Graph Breaks: · 2) Compile Benchmark:

The real checkpoint contains 23 data-dependent graph breaks, but eliminating them makes torch.compile substantially slower rather than improving inference. The benchmark also shows that dispatch fusion cannot amortize launches enough for the Triton kernels to pay off.

  • 1) Graph Breaks:: 23 graph breaks occur in the real 16-layer checkpoint, including 16 at torch.nonzero in expert dispatch.The remaining breaks are nearby conditionals on CUDA tensors.
  • 1) Graph Breaks:: capture_dynamic_output_shape_ops removes all 23 graph breaks, making the break count fully controllable.The resulting zero-break configuration enables a direct benchmark of what happens when graph breaks reach zero.
  • 2) Compile Benchmark:: 18% overhead at sequence length 512 results from compiling at default settings.The benchmark times five configurations at two identical shapes.
  • 2) Compile Benchmark:: 39% overhead at sequence length 512 and 67% at sequence length 1024 result from removing every graph break.The first forward pass rises from one second to nearly forty minutes.
  • 2) Compile Benchmark:: 3× slower than eager is the zero-graph-break result, contradicting the usual emphasis on reporting break counts.Unbacked symbolic shapes force Inductor to handle any size, with guards and fallbacks contributing to the cost.
  • 2) Compile Benchmark:: 3× loss from fusing the dispatch means the hypothesis that fusion would amortize launches fails before the Triton kernels can pay.The combined configuration is not worth reporting separately.
  • 2) Compile Benchmark:: Top-k correction reverses the INT4 ordering between DeepSeek-V2-Lite and OLMoE because top-6 selection inflates the same physical event.The figure compares raw Jaccard drift with expected swapped experts per token.

V. CONCLUSION

Across three models, none of the measured optimizations improves MoE inference end to end because expert dispatch is launch-bound, while routing fidelity and output quality remain separable. The conclusion therefore prioritizes restructuring dispatch over kernel fusion, graph-break removal, or router protection, subject to single-model and limited-correlation evidence.

  • Launch-bound inference: 5.6x to 9.0x isolated kernel gains return 0.999x end to end against a real 1.072x ceiling because inference waits on roughly a thousand sequential launches.Thirty-two times the tokens cost 1.55 times the time, indicating launch-bound execution rather than arithmetic-bound execution.
  • Launch-bound inference: A threefold compilation regression shows that removing graph breaks does not rescue launch-bound expert dispatch.The compiler’s own attempt to fuse dispatch caused the regression.
  • Substitutable experts: 0.53 of eight selected experts change under INT4 on average, but replaying those changes with full-precision weights reproduces only 2.7% of the loss.This causal replay was run on OLMoE alone, so the 2.7% attribution is a single-model result.
  • Substitutable experts: 20% lower drift from exempting routers from quantization makes the model worse, separating routing fidelity from output quality.Drift correlates with quality at +0.907 but is 98% collinear with gate divergence, so correlation alone does not establish mechanism.
  • Implications and future work: The binding constraint is expert-dispatch launch structure, making a single batched operation more promising than kernel gains, routing protection, or graph-break reduction alone.The proposed direction is a single grouped matrix multiplication rather than a loop.
  • Limitations: The cross-model ordering rests on three points with capacity uncontrolled, while compiler timings use cuDNN attention disabled.The drift-quality correlation also spans 16 configurations from one checkpoint and one prompt set.
Loading 2608.26612v1…