Source-linked AI summary
PCoMoE: Shifting MoE Inference from Monolithic Expert Selection to Fine-Grained Path Composition
Ziyan Gan, Fangxin Liu, Chenyang Guan, Junjie Wang, Ning Yang, Haomin Li, Xiang Li, Siran Yang, Jiamang Wang, Lin Qu, Zongwu Wang, Li Jiang, Haibing Guan
TL;DR
MoE inference is limited by treating complete experts as atomic units, leaving intra-expert redundancy underexplored. PCoMoE composes reusable internal paths with compatibility-aware pruning and hardware-conscious execution. It reports up to a 1.31× serving speedup and up to 10% higher downstream-task accuracy than vanilla MoE baselines.
Problem
Existing MoE inference frameworks manage, schedule, or prune complete experts, fixing the optimization boundary before fine-grained intra-expert redundancy can be exploited.
Method
PCoMoE decomposes experts into composable sub-transformations, selects compatible paths with layer-wise pruning, and executes them through source-grouped compute reuse.
Results
Up to 1.31× serving speedup and up to 10% downstream-task accuracy improvement are reported over vanilla MoE baselines.
Takeaways & Limitations
PCoMoE shifts MoE inference from monolithic expert selection to fine-grained path composition while improving throughput and downstream representation fidelity.
Takeaways & Limitations
The current system primarily optimizes autoregressive decoding, while distributed and prefill-stage integrations remain future directions.
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) architectures scale Large Language Model (LLM) capacity efficiently by activating a sparse subset of experts per token. However, modern MoE inference remains heavily constrained by the rigid, whole-expert abstraction. Existing frameworks manage, schedule, or prune experts as atomic execution units, which fixes the optimization boundary too early and leaves fine-grained intra-expert computational redundancy underexplored. In this work, we present PCoMoE, a path-compositional execution framework that shifts MoE inference from coarse-grained expert selection to fine-grained path composition. PCoMoE incorporates a path-level formulation of expert computation, a compatibility-aware layer-wise pruning strategy to suppress low-value path combinations, and a hardware-friendly execution engine to exploit reusable sub-expert structures under strictly bounded overheads. Experimental results demonstrate that PCoMoE achieves up to a 1.31x end-to-end inference speedup while enhancing model accuracy by 10%. The code is available at https://github.com/gzyyy0/PCoMoE
1 Introduction
MoE inference traditionally optimizes routing and complete experts, leaving fine-grained redundancy inside experts underused. PCoMoE instead composes internal paths and reports higher efficiency with improved accuracy.
- Motivation: MoE architectures scale capacity by activating only a sparse subset of experts for each token.This decouples total parameter count from per-token activated computation.
- Optimization gap: Existing optimization frameworks preserve a whole-expert boundary, overlooking fine-grained computational redundancy within individual experts.Inter-expert approaches manage routing decisions and complete experts rather than internal sub-structures.
- Optimization gap: Treating experts as computationally atomic restricts optimization and forces a binary choice between executing or skipping an entire expert.This can create representation bottlenecks or unnecessary computation.
- PCoMoE: PCoMoE decomposes experts into composable sub-structures and uses compatibility-aware routing with layer-wise pruning to suppress low-value path combinations.Its hardware-friendly engine groups dispatches by source and shares expansion-side operations under bounded overheads.
- Results: 1.31× serving speedup and up to 10% accuracy improvement are reported over vanilla MoE baselines.The evaluation reports both end-to-end serving acceleration and downstream-task accuracy gains.
2 Background
A vanilla MoE layer routes tokens to top-k complete experts and executes a dispatch–execute–combine pipeline. Because experts contain structured gated FFNs and system orchestration costs are substantial, whole-expert optimization leaves finer-grained opportunities constrained.
- Vanilla MoE execution: A vanilla MoE layer uses a router and sparse feed-forward experts, selecting the top-k experts for each token hidden representation.The selected experts are denoted E(h).
- Vanilla MoE execution: Routing weights combine the selected expert transformations after tokens are dispatched, grouped for batched execution, and scattered back.This forms a rigid dispatch-execute-combine pipeline.
- Expert structure: Each expert is a structured feed-forward module rather than an atomic computational primitive, commonly instantiated as a SwiGLU-style gated FFN.Representative open-source models include Mixtral and Qwen-MoE.
- Expert structure: SwiGLU-style experts use parallel gate and up projections, element-wise activation fusion, and a final down projection.The gate modulates activation thresholds, the up branch extracts expanded features, and the down branch returns them to model dimension.
- System bottlenecks: End-to-end MoE latency includes routing, dispatch, packing, kernel launches, scattering, and weighted aggregation, while existing optimizations retain whole experts as execution units.This rigid boundary restricts finer-grained optimization opportunities.
3 Motivation
Vanilla MoE optimization stops at whole-expert routing, although experts contain finer-grained computational structure and layer-dependent redundancy. Whole-expert pruning can therefore trade computation for severe and uneven quality loss, motivating structured intra-expert composition.
- Coarse-Grained Routing: Whole-expert routing overlooks redundancy inside indivisible expert modules.Vanilla MoE exposes computation only through monolithic expert boundaries, fixing the optimization boundary too early.
- Whole-Expert Reduction: Single-layer whole-expert pruning causes highly divergent perplexity degradation across layers.Some layers tolerate deleting an expert, whereas others incur substantial degradation, so no fixed global pruning rule is safe.
- Whole-Expert Reduction: Global repetition of locally tolerable expert reductions compounds into severe end-to-end quality drops.The degradation trend persists across models even though its absolute magnitude differs.
- Intra-Expert Operator Asymmetry: Separately composing Gate, Up, and Down increases routing decisions, while grouping Gate and Up exposes a structured reuse boundary.The grouped boundary reflects asymmetric computational costs and preserves a practical middle ground between monolithic execution and unconstrained composition.
- Intra-Expert Operator Asymmetry: Expansion-side substitution is structurally more tolerant than projection-side substitution in Mixtral profiling.Alternative expansion sources produce lower representation error in 23 of 32 layers under independent pairing and in 25 of 32 layers under a shared-source constraint.
- Design Challenges: Fine-grained composition must address combinatorial path growth, runtime control overhead, and layer-varying compositional stability.These constraints jointly require path-space regulation, low-overhead dispatch, and quality preservation.
4 Design
PCoMoE replaces monolithic expert execution with compositional paths formed from reusable expansion and projection operators. Compatibility-aware pruning selects a compact active set, while source-grouped execution reuses expansion computation to bound serving overhead.
- Compositional Path Formulation: PCoMoE decomposes each expert into expansion-side and projection-side operators, creating an n × n compositional path space.Diagonal paths preserve original expert computations, while off-diagonal paths combine an expansion source with a projection target.
- Compositional Path Formulation: A compositional path combines expansion source Ui with projection target Dj, expanding trajectories from n physical experts to n^2 paths.The diagonal path Pi,i is the original expert; off-diagonal paths synthesize new trajectories without adding weights.
- Compositional Path Formulation: PCoMoE preserves vanilla MoE as the diagonal-path special case while activating off-diagonal paths selectively.Restricting execution to diagonal paths exactly recovers the original weighted expert summation.
- Compatibility-Aware Path Routing: Compatibility-aware gating scores paths using routing priors and learned source-to-target compatibility, then prunes low-value trajectories.The resulting active mask is layer-specific and suppresses unoptimized off-diagonal routes while preserving baseline behavior through diagonal initialization.
- Hardware-Efficient Path Execution: Active paths are executed through routing followed by source-grouped reuse rather than isolated path computation.Tokens sharing an expansion operator reuse one compute footprint before dispatching to their projection-side destinations.
- Hardware-Efficient Path Execution: Expansion evaluations scale with unique selected sources |S^(ℓ)(h)| instead of total active paths |Π^(ℓ)(h)|.This structural sharing converts compositional flexibility into deterministic latency reductions.
5 Evaluation
PCoMoE is evaluated across three MoE models and five downstream benchmarks, with matched fine-tuning and single-GPU profiling. It improves accuracy and decoding throughput while showing that compatibility-aware composition and hardware-aligned execution are important to its gains.
- Evaluation Setup: Evaluations cover three MoE models, five benchmarks, matched LoRA fine-tuning, and single-GPU profiling under a unified environment.The models span different routing granularities, and macro average is used as the aggregate representation fidelity metric.
- Quality and Efficiency: Qwen1.5-MoE macro accuracy rises from 67.90 to 73.70, a +5.80-point gain over Vanilla, while matched fine-tuning adds +2.87 points over Vanilla-FT.For Mixtral-8x7B and DeepSeek-V2-Lite, PCoMoE improves over Vanilla by +2.10 and +2.46 points, and over matched Vanilla-FT by +1.61 and +2.14 points.
- Quality and Efficiency: 1.305× decode and 1.294× end-to-end speedups are achieved over vanilla baselines, including 26.50 to 34.58 tokens/s decoding on Mixtral-8x7B.Prefill remains identical because the execution optimizations target autoregressive decoding.
- Design Ablation: 1.238× Decode speedup with a 2.87-point gain over Vanilla-FT shows learned path compatibility is critical to the quality–efficiency balance.Frozen compatibility biases increase throughput but substantially degrade accuracy to 55.50, while Router-Only FT does not reproduce PCoMoE’s improvement.
- Execution Ablation: 1.73× cumulative acceleration, or 191.4 tokens/s, results when path routing, fusion, and source-grouped dispatch are combined.Individually, these components yield 1.15×, 1.28×, and 1.45× speedups over Base PCoMoE, showing the role of hardware-aligned scheduling.
- Pruning Sensitivity: Structural pruning begins at the third block, while expansion-side reuse saves 20.6% on average and varies from 15% to 27% across layers.The first two blocks retain native configuration because early sparsification is lossy; layer-wise budgets respond to local compatibility.
6 Conclusion
PCoMoE shifts MoE inference from monolithic expert selection to fine-grained sub-transformation composition. By combining path decomposition, compatibility-guided gating, and source-grouped compute reuse, it improves throughput and downstream representation fidelity over vanilla MoE baselines.
- Conclusion: PCoMoE shifts MoE inference from monolithic expert selection to fine-grained sub-transformation composition.The framework preserves native MoE layer interfaces while exploiting intra-expert computational redundancy.
- Conclusion: Structural path decomposition, compatibility-guided gating, and source-grouped compute reuse jointly support improved inference throughput and downstream representation fidelity.These mechanisms retain the vanilla layer boundary while enabling compositional execution.
Limitations
The evaluation demonstrates PCoMoE on standard SwiGLU-style MoE structures, while broader deployment requires adapting operator boundaries, extending optimization beyond decoding, and integrating offline calibration into compilation pipelines.
- Scope: PCoMoE is demonstrated using standard SwiGLU-style MoE structures, with alternative expert internals requiring adapted operator boundaries.The underlying path-composition principle is described as generic, but this extension is not part of the demonstrated setting.
- Deployment Scope: The current system primarily optimizes autoregressive decoding; distributed and prefill-stage optimizations remain complementary future directions.This defines the present optimization scope rather than a claim about those unintegrated settings.
- Deployment Scope: Offline calibration adds zero inference-time overhead, but broader deployment would benefit from integrating it into automated end-to-end model compilation pipelines.The compilation integration is identified as future work to streamline deployments.