# CAP: Communication-Aware Placement and Pruning for Efficient Mixture-of-Experts Inference

> CAP jointly optimizes expert placement and pruning for MoE inference, achieving up to 1.86x throughput gains by prioritizing communication over load balance on bandwidth-limited hardware.

- **Source:** [arXiv](https://arxiv.org/abs/2607.05116)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/SLMlah
- **Whiteboard:** https://picx.dev/p/SLMlah/image

## Summary

## 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

| Node | Hardware | Communication Characteristics |
|------|----------|-------------------------------|
| Node A | 8× RTX 3090 | Comm.-constrained (no NVLink, CPU-routed) |
| Node B | 8× A100 | Comm.-comp. balanced (PCIe P2P) |
| Node C | 8× H100 (NVLink) | Comm.-rich |
| Cluster A | 2 nodes × 8 H100, 8×400 Gbps IB | High inter-node bandwidth |
| Cluster B | 2 nodes × 8 H100, 200 Gbps | Low inter-node bandwidth |

### 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)**

| Metric | p = 0 | p = 0.2 | p = 0.3 |
|--------|-------|---------|---------|
| PPL ↓ | 17.36 | 18.17 (+0.81) | 19.12 (+1.76) |
| GSM8K ↑ | 0.8863 | 0.8590 (-2.73 pp) | 0.8514 (-3.49 pp) |
| HumanEval ↑ | 0.8902 | 0.8841 (-0.61 pp) | 0.8598 (-3.04 pp) |
| MMLU ↑ | 0.7154 | 0.7161 (+0.07 pp) | 0.6874 (-2.80 pp) |

**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**

| Cluster | Config | Latency (ms) | Node per Token |
|---------|--------|--------------|----------------|
| A (8×400 Gbps) | c = 1 | 1434.68 | 1.8875 |
| A (8×400 Gbps) | c = 5 | 1447.89 (+0.92%) | 1.8337 (-2.85%) |
| B (200 Gbps) | c = 1 | 2285.37 | 1.8875 |
| B (200 Gbps) | c = 5 | 1980.26 (-13.35%) | 1.8337 (-2.85%) |

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)**

| Target \ Source | HumanEval | LMSYS Chat | Arxiv | Default |
|-----------------|-----------|------------|-------|---------|
| HumanEval | 4.1753 | 4.2481 | 4.4817 | 5.3497 |
| LMSYS Chat | 4.4444 | 4.1405 | 4.4223 | 5.3315 |
| Arxiv | 4.4930 | 4.2832 | 4.3443 | 5.3408 |

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

| Target \ Source | HumanEval | LMSYS Chat | Arxiv | Default |
|-----------------|-----------|------------|-------|---------|
| HumanEval | 1.6772 | 1.8069 | 1.6629 | 1.4372 |
| LMSYS Chat | 1.5620 | 1.4675 | 1.5008 | 1.3303 |
| Arxiv | 1.8140 | 1.9201 | 1.6741 | 1.5972 |

### 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).

---

_Markdown view of https://picx.dev/p/SLMlah, served by PicX — AI-generated visual whiteboard summaries of research papers._