PaperlayerView exact revision on arXiv

Source-linked AI summary

DeaMoE: Efficient MoE Structure for Fast Small-Batch Decoding

Zewen Jin, Shen Fu, Zeping Duan, Shannon Wang, Weihao Wu, Chengjie Tang, Congkun Ai, Ping Gong, Zijian Dai, Youhui Bai, Cheng Li

arXiv:2608.14385v1cs.LGcs.AI

TL;DR

Small-batch MoE decoding is bottlenecked by expert weight loading, while existing optimizations do not explicitly target this issue. DeaMoE groups related experts around shared parameters and uses two-stage routing, delivering speedups in end-to-end serving and microbenchmarks.

  • Problem

    Small-batch MoE decoding remains relatively underexplored despite expert weight loading dominating latency in interactive serving.

  • Method

    DeaMoE groups related experts into departments with shared parameters and uses customized two-stage routing to reduce redundant loading.

  • Results

    DeaMoE delivers consistent decoding speedups, including up to 1.33× end-to-end TPOT speedup for the pre-trained 7B model on A40 and 2.00× peak microbenchmark speedup for DeepSeek-V3 on A40.

  • Takeaways & Limitations

    DeaMoE shows clear advantages for small-batch MoE decoding across end-to-end serving and microbenchmarks on Ampere and Hopper GPUs.

  • Takeaways & Limitations

    The paper focuses on small-batch decoding latency rather than objectives such as model compactness, training scalability, or large-batch throughput.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) models have been widely adopted in real-time interactive applications such as coding assistants, real-time audio-video interaction systems. To meet the extremely low response latency requirements of these scenarios, practitioners commonly employ small-batch decoding, under which MoE inference becomes memory-bound and is severely bottlenecked by expert weight loading. However, this bottleneck has received limited attention, and existing solutions such as post-training weight compression or fine-grained expert design during pre-training either degrade model accuracy or introduce additional computation and communication overhead. To tackle this issue, we propose DeaMoE, a decoding-efficient MoE architecture, in which the experts are grouped into several departments, and the experts belonging to the same department share most parameters since they come from the same professional field, and additionally each expert contains a few private parameters to reflect its uniqueness. Moreover, we design customized two-stage routing strategy for DeaMoE to avoid redundant loading, under which DeaMoE greatly improves the efficiency during LLM decoding. Compared with vanilla MoE, DeaMoE reduces per-step loaded weights by up to 50.9% and achieves up to 1.33 end-to-end TPOT speedup for the pre-trained 7B model on A40, and up to 2.00x and 1.97x peak speedup for DeepSeek-V3 on A40 and H100 in microbenchmarks.

1. Introduction

Small-batch decoding makes MoE inference memory-bound because expert weight loading dominates, while existing remedies risk quality loss or depend on unavailable serving-time parallelism. DeaMoE addresses this with department-shared expert parameters and two-stage routing, preserving model quality while improving decoding speed.

  • Motivation: Small-batch decoding prioritizes latency but makes MoE inference memory-bound by bottlenecking on expert weight loading.Larger batches improve throughput but introduce unacceptable delays in real-time applications.
  • Limitations of Existing Approaches: Existing expert compression and fine-grained expert designs risk degrading model quality or rely on serving-time parallelism unavailable in small-batch decoding.The introduction identifies this expert-loading bottleneck as largely overlooked by existing MoE optimizations.
  • DeaMoE Architecture: DeaMoE groups experts into departments that share most parameters while retaining small private parameter sets, reducing expert parameter redundancy.The shared parameters form a common backbone, while private parameters capture each expert’s uniqueness.
  • Two-Stage Routing: DeaMoE’s two-stage routing first selects departments and then routes tokens to experts, improving weight reuse across tokens within each decoding step.Each department processes its collected tokens uniformly before assigning them to individual experts for differentiated processing.
  • Evaluation: 1.33× TPOT speedup is achieved for end-to-end vLLM serving of the 7.3B model on A40, while DeepSeek-V3 reaches 2.00× on A40 and 1.97× on H100 in microbenchmarks.A 7.3B DeaMoE model and a budget-matching standard MoE counterpart were pretrained over 110B tokens, with DeaMoE preserving model quality on pre-training loss and a dozen downstream benchmarks.

2. Background and Motivation

MoE models serve latency-sensitive interactive applications, but small-batch decoding trades away throughput and becomes dominated by expert-weight loading. Existing optimizations leave this bottleneck underexplored, motivating a pre-training-first structure that reuses shared expert parameters while preserving specialization.

  • Latency-sensitive applications: Token-level response latencies of tens of milliseconds make decoding latency a primary target in real-time and interactive applications.
  • Small-batch trade-off: Large batches improve throughput but violate real-time latency constraints, whereas small batches reduce latency while severely degrading throughput.
  • Expert-loading bottleneck: Small-batch decoding is dominated by expert-parameter loading because few tokens activate multiple experts with limited weight reuse across tokens.
  • Expert-loading bottleneck: 44 MB of FP8 weights per DeepSeek-V3 expert approaches the H100 GPU’s L2-cache limit, worsening loading challenges as expert scale increases.
  • Limits of prior work: Existing MoE optimizations largely target compactness, training scalability, or large-batch throughput rather than small-batch decoding latency.
  • Motivation for DeaMoE: DeaMoE instead redesigns the MoE during pre-training to reuse shared backbone parameters while retaining lightweight expert-specific transformations.

3. DeaMoE: New Model Structure

DeaMoE groups experts into departments that share large projection matrices while retaining expert-specific transforms, then routes tokens in two stages to reuse department computations and reduce repeated weight loading. The design also favors soft group-limited routing, SiLU interfaces, and identity initialization to preserve efficiency and model quality.

  • Model structure: DeaMoE groups experts into departments, sharing gate, up, and down projection matrices within each department while adding three private sub-matrices per expert.Experts across different departments remain independent because their areas of responsibility differ.
  • Routing and decoding: Two-stage routing first aggregates tokens by department for shared computation, then dispatches them to experts for private-parameter transforms, reducing repeated department-parameter loading.The method retains the router and final-output functions while collecting department-specific token groups before expert computation.
  • Loading efficiency: The loading ratio decreases at small token counts and increases as the batch grows, producing a non-monotonic relationship between batch size and per-step expert-parameter loading.At very small batches, limited cross-token sharing keeps normalized loading relatively high; reuse improves as token counts increase from that regime.
  • Training design: SiLU is applied at all three gate, up, and down interfaces to avoid consecutive linear maps between department backbones and expert-specific transforms.The paper considers SiLU, GeLU, and RMSNorm before selecting SiLU.
  • Training design: Identity initialization makes each expert initially equivalent to its department backbone, allowing expert-specific behavior to emerge during training rather than imposing specialization through random noise.The paper reports that identity initialization achieves lower training loss in ablation results.
  • Routing design: Exact department coverage degrades model quality, whereas soft group-limited top-k routing better preserves the quality–efficiency trade-off.The soft constraint caps the number of activated departments as an upper bound rather than requiring every token to cover exactly that number.

4. Accuracy Evaluation after Pre-Training

After identical 7.3B-parameter pre-training on the same 110B-token corpus, DeaMoE preserves downstream accuracy and language-modeling quality relative to the standard MoE baseline. It outperforms the baseline on six of ten tasks and achieves lower perplexity on PTB and WikiText-2 while remaining comparable on WikiText-103.

  • Evaluation Scope: The evaluation compares DeaMoE with a standard MoE baseline using downstream benchmark accuracy and language-modeling perplexity after pre-training.The downstream suite covers commonsense reasoning, natural language inference, and multi-choice question answering tasks.
  • Experimental Setup: Both models use 7.3B parameters and are matched in total parameter count, per-token FLOPs, corpus, optimizer settings, and training configuration.Pre-training uses RedPajama-v1 with 110B tokens and identical software and hardware configurations.
  • Downstream Accuracy: DeaMoE achieves comparable performance across ten downstream classification and reasoning tasks, outperforming the baseline on six.The gains include SIQA, BoolQ, WinoGrande, and all three ANLI subsets: R1, R2, and R3.
  • Language Modeling Perplexity: DeaMoE achieves lower perplexity on PTB and WikiText-2 while remaining comparable to the baseline on WikiText-103.These results indicate that the architecture preserves language-modeling quality under the same pre-training budget and inference FLOPs consumption.

5. Evaluation for Inference Speedups

DeaMoE improves decoding speed in end-to-end serving and per-layer microbenchmarks, with gains strongest when expert weight loading dominates. The benefits persist across A40 and H100 GPUs but are smaller for small-expert models on H100 due to cache reuse.

  • End-to-end serving: On A40, DeaMoE achieves 1.02×–1.33× latency speedup across evaluated batch sizes in end-to-end serving.The evaluation measures average Time Per Output Token (TPOT) across varying concurrency levels using Baseline-7B and its DeaMoE counterpart.
  • End-to-end serving: DeaMoE reduces memory traffic per decoding step, translating lower expert weight loading into measurable small-batch latency and throughput gains.This advantage alleviates the A40 bandwidth bottleneck and is observed in small-batch scenarios.
  • Microbenchmarks: On A40, peak speedups approach 2.00× for DeepSeek-V3 and 1.74× for Qwen3-235B-A22B across microbenchmarks.DeaMoE provides clear and stable speedups for all three evaluated models, with the largest improvements in large-expert models.
  • Microbenchmarks: On H100, DeaMoE continues to deliver substantial speedups for DeepSeek-V3 and Qwen3-235B-A22B.The microbenchmarks measure per-layer execution time across batch sizes from 4 to 128 on both A40 and H100 GPUs.
  • Microbenchmarks: For Baseline-7B on H100, gains are smaller and mildly regress at very small and very large batch sizes because expert weights can be reused from the large L2 cache.This reduces baseline weight-loading cost and limits optimization headroom for small-expert models.
  • Overall findings: DeaMoE’s effectiveness correlates with expert weight-loading dominance: large, sparsely reused expert matrices yield consistent gains, while small-expert configurations on high-bandwidth hardware benefit less.When weight movement is not the bottleneck, the relative advantage becomes less pronounced.

6. Conclusion

DeaMoE improves MoE performance in LLM decoding by grouping similar, parameter-redundant experts into departments. Experts in each department share common parameters, while a customized two-stage router first assigns tokens to departments and then to experts.

  • 6. Conclusion: DeaMoE groups all experts into several departments to address parameter redundancy caused by similarity among experts.The paper observes that some experts in an MoE layer exhibit a high degree of similarity.
  • 6. Conclusion: Experts within the same department share most parameters because they capture shared knowledge within the same functional category.This shared structure is combined with differentiated expert processing.
  • 6. Conclusion: The customized two-stage routing strategy first allocates tokens to departments and then assigns them to corresponding experts.Each department uniformly processes all collected tokens at once before experts perform further differentiated processing.

A. Ablation Study

A controlled ablation on a smaller 5B-scale model supports DeaMoE’s choices of SiLU activation, identity initialization, and soft group-limited routing. Identity initialization lowers loss, while strict department coverage performs worse than the default soft constraint.

  • Experimental setup: The ablation trains a smaller 5B-scale model on 12.4B RedPajama-V1 tokens, varying one design choice at a time while keeping the remaining setup unchanged.Loss is averaged over the last 200 optimization steps.
  • Non-linear operator: SiLU achieves the best last-200-step average loss among the evaluated non-linear operators, outperforming GeLU, RMSNorm, and no operator.The comparison covers variants No.1 to No.4 under the same training setup.
  • Expert matrix initialization: 2.504 loss with identity initialization improves on 2.530 with the alternative initialization for expert-specific matrices.The comparison uses variants No.1 and No.5 and averages loss over the last 200 steps.
  • Department Top-k: 2.519 loss from strict department coverage is worse than 2.504 from the default soft constraint.Strict routing activates exactly k_dept departments per token, whereas the default uses soft group-limited routing.
  • Overall findings: Overall, the ablations support SiLU activation, identity initialization for expert transforms, and soft group-limited routing as DeaMoE’s adopted design choices.The paper states that these components preserve model quality while enabling decoding-time efficiency gains.
Loading 2608.14385v1…