Summary (Overview)

• CAP (Communication-Aware Assignment and Pruning) is a framework for efficient Mixture-of-Experts (MoE) inference that jointly optimizes expert placement and pruning while explicitly considering communication volume, computation, and accuracy.
• The framework consists of three components: (1) co-activation-driven expert placement to reduce communication by grouping frequently co-activated experts, (2) communication-computation trade-off adjustment that generates a spectrum of placements with different communication/load-balance trade-offs, and (3) communication-aware expert pruning that removes routing destinations with low accuracy contribution and high communication cost.
• CAP achieves 1.23×–1.86× throughput improvement over DeepSeek EPLB and sequential placement in vLLM across single-node (8-GPU) and multi-node settings.
• The framework demonstrates strong hardware adaptivity: it selects different optimal placements depending on whether the hardware is communication-constrained, balanced, or communication-rich.
• CAP preserves better model accuracy at the same target speedup under lossy acceleration compared to naive pruning approaches, and shows strong stability and transferability of expert groupings across workloads.

---

Introduction and Theoretical Foundation

Background

The Mixture-of-Experts (MoE) architecture has become increasingly important for scaling large language models (LLMs) while keeping per-token computation manageable. In MoE models, the dense FFN in each transformer layer is replaced by a set of experts, and a gating network selects the top-k experts for each token. The MoE layer output is computed as:

$$\operatorname{MoE} (x) = \sum_ {i = 1} ^ {k} \operatorname{Gate} _ {i} (x) \cdot \operatorname{Expert} _ {i} (x)\tag{1}$$

Key Challenge: Communication in Expert Parallelism

MoE inference systems combine expert parallelism (experts distributed across devices) with data parallelism (dense layers replicated). When tokens are routed to experts on different devices, all-to-all communication is required. Critically, the communication volume for a single token depends on the number of distinct devices it must visit, not the number of experts it activates—since experts co-located on the same device share a single token copy.

Key insight: Prior MoE models (e.g., Mixtral 8x7b with only 8 experts) left no room for placement-based communication optimization. Modern models like Qwen3-30B-A3B (128 experts) and DeepSeek-V2-Lite (64 experts) have large expert counts, making communication-aware placement crucial.

Limitations of Prior Work

• Expert placement methods (e.g., DeepSeek EPLB) primarily target load balance but overlook communication volume.
• Expert pruning methods reduce computation but fail to consider how pruning decisions affect communication patterns.
• The optimal trade-off between communication and load balance varies significantly across hardware platforms.

---

Methodology

1. Co-Activation Driven Expert Placement

The approach models expert co-activation using a weighted graph where edge weights represent co-activation probabilities $p_{ij}$ (the ratio of times experts $i$ and $j$ are selected together to the times either is selected).

Communication objective:

$$\mathcal {C} (\pi) = \sum_ {i < j} p_{ij} d (\pi (i), \pi (j))\tag{2}$$

where $d(\pi(i), \pi(j))$ is the communication cost between GPUs hosting experts $i$ and $j$:

$$d (\pi (i), \pi (j)) = \left\{ \begin{array}{l l} 0, & \pi (i) = \pi (j), \\ 1, & \pi (i) \neq \pi (j) \text { and } \operatorname{node} (\pi (i)) = \operatorname{node} (\pi (j)), \\ c, & \operatorname{node} (\pi (i)) \neq \operatorname{node} (\pi (j)) \end{array} \right.\tag{3}$$

Algorithm (two-stage greedy):

1. Stage 1 (GPU-level): Select $k$ experts with highest weighted degrees as seeds for $k$ GPU groups. Iteratively grow groups in round-robin fashion, assigning each unassigned expert to the group with which it has the highest cumulative co-activation weight.
2. Stage 2 (Node-level): Treat GPU groups as coarse-grained units, compute inter-group co-activation strength, and apply the same grouping principle to assign GPU groups to nodes.

2. Communication-Computation Trade-off Adjustment

Load imbalance objective:

$$\mathcal {B} (\pi) = \frac {\max_ {g} L_ {g}}{\frac {1}{k} \sum_ {g} L_ {g}}\tag{4}$$

where $L_g$ is the number of tokens processed by GPU $g$.

Combined communication-load objective:

$$\min_ {\pi} \lambda \mathcal {C} (\pi) + (1 - \lambda) \mathcal {B} (\pi)\tag{5}$$

where $\lambda \in [0, 1]$ controls the trade-off.

The Adjustment Approach starts from the communication-oriented placement and iteratively swaps experts between high-load and low-load devices, generating a placement spectrum with gradually increasing communication cost and improving load balance. This spectrum forms a Pareto frontier over the default sequential placement.

3. Communication-Aware Expert Pruning

Pruning is performed at the device level rather than expert level, since removing an entire routing destination yields much larger communication reduction.

Formulation:

$$\min_ {x_ {d} \in \{0, 1 \}} \sum_ {d} c_ {d} x_ {d} \quad \text{s.t.} \quad \sum_ {d} a_ {d} x_ {d} \geq 1 - p\tag{6}$$

where:

• $a_d$ = accuracy contribution of device $d$ (sum of routing scores of candidate experts on that device)
• $c_d$ = communication cost (1 for intra-node, larger constant $c$ for inter-node)
• $p$ = user-specified pruning threshold

A greedy solver retains devices in descending order of the ratio $a_d / c_d$ until the target accuracy $1-p$ is satisfied.

4. Overall Pipeline (4 Steps)

1. Run inference with communication-aware pruning enabled; collect token-level routing traces.
2. Perform co-activation-driven placement on pruned routing data (pruning changes co-activation patterns).
3. Apply communication-computation trade-off adjustment to generate a placement spectrum.
4. Use lightweight profiling to select the most suitable placement for the target hardware.

---

Empirical Validation / Results

Testbed Configurations

End-to-End Performance

• CAP consistently outperforms baselines across all testbeds, with up to 1.86× throughput improvement on the comm.-constrained node.
• DeepSeek EPLB performs similarly to or worse than Default on the comm.-constrained machine—highlighting that load balance alone is insufficient when communication dominates.

Placement Spectrum Analysis

• Comm.-constrained node: Communication-minimal placement is optimal.
• Balanced node: Mid-spectrum placement (e.g., Placement 3) achieves best performance—20.4% improvement over unadjusted placement, 23.4% over EPLB, and 16% over Default.
• Comm.-rich node: Most load-balanced placements deliver the best performance.

Pruning Effectiveness

Table I: Accuracy under different pruning thresholds (lower is better for PPL, higher is better for others)

Key result: CAP with $p = 0.2$ achieves the same 33% latency reduction as naive pruning with $p = 0.3$, preserving better accuracy at the same speedup.

Inter-node Pruning in Clusters

Table II: Impact of inter-node communication cost in multi-node pruning

Setting $c = 5$ reduces cross-node communication; the benefit is significant only when inter-node bandwidth is limited (Cluster B: 13.35% latency reduction).

Stability and Transferability

• In-workload stability: Communication cost variation < 2.9%, load imbalance variation < 5.8% across windows.
• Cross-workload transferability (Tables III & IV): Groupings remain stable across HumanEval, LMSYS Chat, and Arxiv datasets.

Table III: Cross-workload communication cost (Device per Token)

Table IV: Cross-workload load imbalance (Max Load/Avg. Load)

Optimization Overhead

• Routing statistics stabilize after ~8,000 tokens.
• Total offline optimization: < 6 minutes (36.72s for trace collection, ~80s for initial placement, 9.1s for 7 adjusted placements, ~240s for profiling).
• Grouping generation is once per model; placement profiling is once per machine.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Communication-aware placement formulation: The paper formalizes expert placement as a capacity-constrained weighted graph partitioning problem with a communication objective ($\mathcal{C}(\pi)$) that accounts for both intra-node and inter-node communication costs.

2. Communication-load multi-objective optimization: The combined objective $\min_\pi \lambda \mathcal{C}(\pi) + (1-\lambda)\mathcal{B}(\pi)$ explicitly models the inherent trade-off between communication minimization and load balance, which prior work treated as separate concerns.

3. Device-level pruning formulation: The pruning objective (Eq. 6) reframes expert pruning as a cost-constrained discrete selection problem at the device level, demonstrating that pruning decisions should consider communication topology, not just accuracy contribution.

Practical Implications

1. Hardware-adaptive optimization: CAP demonstrates that the optimal placement strategy depends critically on hardware characteristics. Communication-constrained systems benefit most from communication-minimal placement, while communication-rich systems favor load-balanced placement.

2. Load balance is insufficient: The failure of EPLB on comm.-constrained hardware shows that load-balance-oriented optimization alone can degrade performance when communication dominates.

3. Pruning as communication optimization: Communication-aware pruning achieves the same speedup with less accuracy loss by eliminating entire routing destinations, especially those involving expensive inter-node communication.

4. Practical deployment: The low offline overhead (< 6 minutes) and strong cross-workload transferability make CAP practical for real-world deployment without frequent regrouping.

---

Conclusion

CAP addresses a critical gap in MoE inference optimization by jointly considering computation, communication, and accuracy. The framework's three components work in a unified pipeline:

1. Co-activation-driven placement reduces communication by grouping frequently co-activated experts.
2. Trade-off adjustment generates a spectrum of placements balancing communication and load, with hardware-aware selection.
3. Communication-aware pruning reduces both computation and communication while preserving accuracy.

Key findings:

• Communication volume is a first-order concern in modern large-expert-count MoE models.
• The optimal communication-load trade-off is hardware-dependent, requiring adaptive selection.
• Device-level pruning that considers communication topology outperforms expert-level pruning in both speedup and accuracy preservation.

Future directions suggested by this work include:

• Extending the framework to handle heterogeneous hardware configurations.
• Exploring dynamic (runtime) placement adaptation as routing patterns evolve.
• Investigating interaction between communication-aware placement and other optimization dimensions (e.g., memory management, scheduling).