PCoMoE: Shifting MoE Inference from Monolithic Expert Selection to Fine-Grained Path Composition

Traditional MoE acceleration relies on eliminating entire expert executions. Since the vanilla MoE interface exposes computation via top-$k$ routing, a baseline replaces top-$k$ routing with top-$(k-1)$ routing to decrease the number of activated experts per token. We sweep individual layers to profile the perplexity impact of this single-layer pruning.

As shown in Figure 3, imposing the same expert-level pruning across different layers yields highly divergent quality degradation. Certain layers exhibit minimal perplexity increases, whereas others incur substantial degradation that exceeds the average trend, proving that whole-expert reduction is neither uniformly safe nor uniformly detrimental. This variation demonstrates that inter-expert redundancy is highly non-homogeneous and layer-dependent. While specific layers tolerate complete expert deletion, this local redundancy cannot be generalized into a fixed global pruning rule. Multi-layer reduction compounds these errors.

As illustrated in Figure 4, local errors compound catastrophically across deep networks when expert reduction is extended globally. Although the absolute degradation scales differently across models, the underlying trend remains consistent: locally tolerable approximations aggregate into severe end-to-end quality drops. Consequently, repeated whole-expert reduction rapidly depletes safe redundancy, proving that monolithic expert boundaries are too coarse for optimal efficiency-quality trade-offs.

The central challenge is therefore not merely adjusting the number of reduced experts, but breaking the whole-expert abstraction to exploit finer computational granularities.

To bypass these bottlenecks, we analyze the internal computational structures of individual experts. A vanilla SwiGLU expert maps inputs through parallel gate and up projections on the expansion side (comprising two-thirds of total expert FLOPs) before a final down projection back to the hidden dimension. Directly treating these operators as independent, unconstrained composition units (Figure 5(a)) introduces prohibitive control-flow overheads. Instead, the intrinsic expansion-projection boundary (Figure 5(b)) defines a structured middle ground. In a top-2 Mixtral architecture, reusing the expansion side eliminates approximately one-third of total expert computation, whereas reusing the projection side saves only one-sixth, concentrating potential efficiency gains on the expansion side.

Profiling operator substitution on Mixtral confirms that this compute-dense expansion side also exhibits superior structural tolerance for representation sharing. Substituting the gate and up projections with an alternative source expert yields lower representation error in 23 out of 32 layers under optimal independent pairing. Under a stricter constraint where all experts within a layer share a single replacement source, 25 out of 32 layers still favor expansion-side substitution. This dual computational and structural asymmetry demonstrates that the expansion-projection boundary serves as an optimal abstraction, avoiding the accuracy drops of whole-expert dropping while circumscribing the irregular execution overheads of arbitrary operator splitting.

However, translating this theoretical boundary into serving gains introduces three critical co-design challenges: (1) the path search space scales combinatorially compared to vanilla expert selection, (2) additional control logic or dynamic indexing can easily offset computational savings under strict latency constraints, and (3) compositional stability varies heavily across layers, threatening feature fidelity. Addressing these hardware-software trade-offs requires an execution framework that simultaneously regulates the path composition space, minimizes runtime dispatch overheads, and guarantees model quality.

PCoMoE redefines the computational granularity of the MoE layer by decoupling SwiGLU experts along their structural asymmetry boundaries. The expansion-side operator aggregates the compute-dominant gate and up projections, while the projection-side operator maps intermediate activations back to the hidden dimension. Formulating their ordered combination as a compositional path concentrates optimization capabilities precisely where computational reuse potential is maximized, shifting the execution paradigm from monolithic expert selection to flexible path composition.

For an MoE layer containing $n$ physical experts, the computation of an individual expert $e$ is decoupled into an expansion-side function $U_{e}(h)$ and a projection-side function $D_{e}(z)$. We express the expansion side as

and the projection side as

which reconstructs the original vanilla expert transformation via

We define the pools of reusable expansion and projection operators across the layer as $\mathcal{U}=\{U_{1},\ldots,U_{n}\}$ and $\mathcal{D}=\{D_{1},\ldots,D_{n}\}$, respectively. A compositional path $P_{i,j}(h)$ represents the ordered composition of an expansion source $i$ and a projection target $j$:

This formulation expands the available computational trajectories from $n$ physical modules to an $n^{2}$ compositional space. The diagonal path $P_{i,i}$ represents the original vanilla expert $F_{i}$, while an off-diagonal path $P_{i,j}$ ($i\neq j$) synthesizes a new computational trajectory combining the expansion side of expert $i$ with the projection side of expert $j$.

This optimization landscape is purely architectural and introduces zero parameter or storage overheads. All routes are dynamically synthesized by recombining existing physical operators without creating new weights. Restricting execution to diagonal paths naturally reduces PCoMoE back to the vanilla MoE layer, whereas off-diagonal paths serve as selectively activated candidates to maximize efficiency-quality trade-offs.

Consequently, the MoE layer output generalizes from expert-level routing to path-level aggregation. Let $\Pi(h)\subseteq\{1,\ldots,n\}\times\{1,\ldots,n\}$ denote the active compositional paths selected for a token hidden state $h$, and let $\beta_{i,j}(h)$ represent the corresponding path weight. The PCoMoE layer output is formulated as

The vanilla MoE layer remains a strict subset of this generalized formulation, where setting $\Pi(h)=\{(e,e)\mid e\in\mathcal{E}(h)\}$ and $\beta_{e,e}(h)=\alpha_{e}(h)$ exactly recovers the original weighted expert summation. Although this formulation unlocks a fine-grained optimization landscape, evaluating the entire $n^{2}$ space introduces prohibitive computational overheads. PCoMoE treats this matrix strictly as a candidate space, using a compatibility-driven pruning mechanism to isolate a compact, high-value active path set for execution.

The compositional path formulation expands the architectural search landscape into an $n\times n$ routing matrix. To suppress noisy off-diagonal recompositions and bound runtime dispatch complexity, PCoMoE introduces a compatibility-aware path gating mechanism (Figure 7) that factors each trajectory score into a source-side routing prior, a target-side routing prior, and a learned pairwise compatibility bias. For a compositional path $P_{i,j}^{(\ell)}$ in layer $\ell$, the gating function computes the composite routing score as

where $g_{\mathrm{base}}^{(\ell)}(h)\in\mathbb{R}^{n}$ denotes the base expert-routing signal, $B^{(\ell)}=[b_{i,j}^{(\ell)}]\in\mathbb{R}^{n\times n}$ represents the learned compatibility bias, $g_{\mathrm{tgt}}^{(\ell)}(h)\in\mathbb{R}^{n}$ provides a projection-side prior, and $\lambda$ scales this target-side preference. The base routing signal anchors the optimization to the pretrained vanilla checkpoints by broadcasting across source indices, while the target prior regularizes projection choices. Crucially, the compatibility matrix $B^{(\ell)}$ parameterizes the structural alignment between expansion source $U_{i}$ and projection target $D_{j}$, transforming the raw Cartesian product into a filtered execution topology.

PCoMoE initializes this compatibility matrix to strictly prioritize vanilla computational trajectories. Diagonal entries are zero-initialized via $b_{i,i}^{(\ell)}=0$ to safeguard baseline expert behavior, whereas off-diagonal elements receive a negative offset with bounded random noise:

This negative prior prevents unoptimized off-diagonal trajectories from disrupting early training phases while breaking structural symmetry. During fine-tuning, this bias is dynamically converted into a layer-specific active mask $M^{(\ell)}\in\{0,1\}^{n\times n}$. Rather than applying a post-hoc, one-shot pruning step, PCoMoE iteratively tightens a structural sparsity criterion at progressive training intervals to filter low-value off-diagonal routes while hard-coding $m_{i,i}^{(\ell)}=1$. This continuous regularization allows the remaining routing weights to co-adapt with the contracting active path set, preserving model convergence stability.

The active mask operates independently across layers on routing eligibility rather than on physical parameters, enabling each layer to autonomously converge to its optimal path density based on local feature sensitivity. Because no physical weights are deleted, suppressed paths are merely excluded from the active runtime dispatch graph. During inference, inactive paths are masked out prior to gating, ensuring that the online router evaluates only the optimized, compact set of active trajectories induced by $M^{(\ell)}$, thereby successfully minimizing control-flow overheads.

To translate the pruned path space into physical serving gains, PCoMoE decouples its runtime execution into a routing phase and a source-grouped reuse phase (Figure 8).

Compositional Path Routing. During inference, the framework isolates active trajectories by masking the $n\times n$ routing matrix with the pruned path set $\mathcal{A}^{(\ell)}=\{(i,j)\mid m_{i,j}^{(\ell)}=1\}$. The router selects the active path subset via

collapsing dense evaluation into a sparse search. PCoMoE materializes path indices, source-to-target layouts, and grouped dispatch metadata offline. Based on this cached layout, the fusion optimization combines online filtering, top-$k$ selection, and dispatch preparation into a single lightweight indexing operation, avoiding reconstruction of the full $n\times n$ path space. Because all diagonal paths remain active, the routing graph automatically falls back to vanilla expert execution under low-confidence scenarios, preserving the native MoE interface.

Source-Grouped Compute Reuse. Following routing, PCoMoE aggregates active trajectories by their expansion-side sources to eliminate redundant tensor evaluations. This source-grouped dispatch changes the dispatch unit from individual paths $(i,j)$ to source groups, allowing paths that share the same expansion-side operator to be scheduled together. Denoting the unique selected sources as $\mathcal{S}^{(\ell)}(h)=\{i\mid\exists j,(i,j)\in\Pi^{(\ell)}(h)\}$ and their associated projection targets as $\mathcal{T}_{i}^{(\ell)}(h)=\{j\mid(i,j)\in\Pi^{(\ell)}(h)\}$, the final layer output is formulated as:

By executing each compute-heavy expansion operator $U_{i}(h)$ exactly once per unique source, expansion evaluations scale with the active source count $|\mathcal{S}^{(\ell)}(h)|$ instead of the total path budget $|\Pi^{(\ell)}(h)|$. This hardware-software co-design converts structural sharing into deterministic latency reductions.