# 3D-MoE: Complementarity-Guided Hardware-Software Co-Design for MoE Decode on 3D-Stacked NMP

> 3D-MoE exploits expert complementarity—trace-level load cancellation in MoE decode—to achieve 4.63x layer speedups over expert parallelism with only 2.11% area overhead.

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

## Summary

## Summary (Overview)

- **Problem**: Mixture-of-Experts (MoE) inference on 3D-stacked Near-Memory Processing (NMP) suffers from a fundamental tension between expert weight locality (preserved by Expert Parallelism, EP) and load balance under highly skewed, time-varying expert loads.
- **Key Insight**: The paper identifies and validates **expert complementarity** — a trace-level load-cancellation structure in MoE decode where domain-specialized token routing creates offsetting expert-load fluctuations across decode steps. This structural signal enables EP-friendly load balancing without sacrificing weight locality.
- **Proposed System**: **3D-MoE**, a hardware–software co-designed system combining (1) offline complementarity-guided expert grouping and locality-aware mapping, (2) a lightweight **Collaborative Memory-Network (CMN)** hardware fast path for residual intra-group hotspot assistance, and (3) tiered cost-aware online remapping for workload drift robustness.
- **Results**: 3D-MoE achieves an average **4.63× MoE-layer speedup** over EP baselines, **1.58×–1.86× end-to-end throughput** gains over the strongest NMP baselines, and **5.54×–7.88× speedups over GPU** baselines, with only **2.11% additional logic-die area** overhead.

---

## Introduction and Theoretical Foundation

### MoE Decode Challenges

Mixture-of-Experts (MoE) architectures have become mainstream for scaling Large Language Models (LLMs): each token activates only a small expert subset, decoupling model capacity from inference cost. However, this sparsity makes autoregressive decoding highly memory-bound:

- Each expert receives few tokens per step, producing small GEMMs with poor data reuse.
- MoE layers account for **83–92% of end-to-end latency** at batch size 64.
- Expert loads are highly skewed and dynamic — the hottest expert can receive **over 20× the step-average token count**.

### 3D-Stacked NMP as the Target Platform

3D-stacked DRAM Near-Memory Processing (NMP) provides the bandwidth that MoE decode demands. Each NMP core integrates a GEMM engine, on-die buffers, a channel controller, and a NoC router, tightly coupled to an HBM channel. This tiled organization creates a **severe cost asymmetry**:

> Local DRAM accesses are fast and bandwidth-rich, whereas remote accesses traverse the network-on-chip (NoC) and consume shared bandwidth.

### The Locality–Balance Conflict

| Strategy | Strengths | Weaknesses |
|----------|-----------|------------|
| **Expert Parallelism (EP)** | Preserves expert weight locality; aggregates routed tokens | Suffers from hotspot-induced imbalance |
| **Tensor Parallelism (TP)** | Smooths load across cores | NoC overhead reaches 52% of expert execution time at 32 cores; weakens local reuse |
| **Dynamic EP methods** | Partially alleviate hotspots | Hot-expert prediction accuracy <50%; weight migration consumes NoC bandwidth |

### Key Observation: Expert Complementarity

Although MoE training includes load-balancing objectives, pre-trained models develop **structured expert preferences** reflected in domain-skewed activation patterns. As the request mix changes, demand shifts across expert subsets, inducing **offsetting variations** in decode-step loads. The paper defines this trace-level load cancellation as **expert complementarity**.

---

## Methodology

### 1. Complementarity Detection and Validation

**Trace Construction**: For each layer, expert $e$ is represented by its routed-token count $F_{e,d}$ at decode step $d$ under batched serving.

**Statistical Test**: Complementarity is quantified using Pearson correlation between expert traces. To avoid model-independent thresholds, an **empirical null distribution** is constructed by circularly shifting one expert trace relative to the other:

> A pair is significantly complementary when its observed correlation falls below the lower α-quantile of this null (the 5th percentile for α = 0.05).

**Validation Results**:

| Model | Complementary | Positive | Random |
|-------|---------------|----------|--------|
| DeepSeek-V2-Lite | 0.19 | 0.67 | 0.43 |
| Qwen3-30B-A3B | 0.25 | 0.71 | 0.46 |
| GLM-4.5-Air | 0.21 | 0.69 | 0.44 |

Complementary pairs show significantly negative correlation (low cosine similarity), confirming association with domain-shifted expert subsets.

### 2. Offline Grouping and Mapping (ILP Formulation)

**Expert Grouping Objective** — hierarchical multi-objective ILP:

**Objective 1**: Minimize inter-group range of computation cost:

$$\text{Minimize} \sum_{d \in D} \left( \max_{g \in G} T_{g,d} - \min_{g \in G} T_{g,d} \right)$$

**Objective 2**: Minimize inter-group range of collaboration cost:

$$\text{Minimize} \sum_{d \in D} \left( \max_{g \in G} M_{g,d} - \min_{g \in G} M_{g,d} \right)$$

**Objective 3**: Minimize total collaboration cost:

$$\text{Minimize} \sum_{d \in D} \sum_{g \in G} M_{g,d}$$

**Computation cost** for group $g$ at step $d$:

$$T_{g,d} = \sum_{e \in E} B_{e,g} \cdot D_{e,g} \cdot F_{e,d}$$

where $B_{e,g} \in \{0,1\}$ indicates expert-to-group assignment and $D_{e,g} \in [0,1]$ is the workload fraction.

**Pairwise load-offset score** (transferable workload):

$$C_{e_1,e_2,d} = \begin{cases} \min(|F_{e_1,d} - AL_d|, |F_{e_2,d} - AL_d|) & \text{if } (F_{e_1,d} - AL_d)(F_{e_2,d} - AL_d) < 0 \\ 0 & \text{otherwise} \end{cases}$$

where $AL_d = \frac{1}{|E|} \sum_{e \in E} F_{e,d}$ is the step-average expert load.

**Expert Mapping Objective**:

$$\text{Minimize} \sum_{e_1 \in E} \sum_{e_2 > e_1} \left[ \sum_{g \in G} B_{e_1,g} B_{e_2,g} \sum_{d \in D} C_{e_1,e_2,d} + \lambda_{route} A_{e_1,e_2} \right] R_{e_1,e_2}$$

where $A_{e_1,e_2}$ is co-activation affinity (normalized co-activation frequency under top-k routing) and $R_{e_1,e_2}$ is communication cost between experts.

**Constraints**:

$$\forall e \in E, \quad \sum_{g \in G} D_{e,g} = 1$$

$$\forall g_1, g_2 \in G, \quad \sum_{e \in E} B_{e,g_1} = \sum_{e \in E} B_{e,g_2}$$

$$\forall n \in N, \quad \sum_{e \in E} P_{e,n} \leq P_{max}$$

The two stages are solved sequentially with Gurobi, terminating at a 3% relative MIP gap or after three minutes.

### 3. Collaborative Memory-Network (CMN) Hardware

CMN provides a dedicated "load once, broadcast to many" path for residual intra-group hotspots:

- **Data Forwarding Unit (DFU)**: Integrated into the 3D memory controller after the Read Return Buffer; forwards each hot expert's weight tile to both the host core's local Weight Buffer and the RCL path.
- **Data Receiving Unit (DRU)**: Captures incoming streams on helper cores and writes to local Weight Buffers.
- **Reconfigurable Collaboration Link (RCL)**: Reuses existing NoC links in two phase-aligned modes — standard packet transport for token dispatch/gather, and temporary local broadcast for hot-expert collaboration. A lightweight **Bus Switch** and distributed **Session Arbiter** detach selected links from the packet pipeline to form a temporary bus-like broadcast path.

**Hardware motivation** (hotspot collaboration benchmark):

| #Collab Cores | 2 | 4 | 8 | 16 |
|---------------|-----|-----|-----|------|
| Ideal latency (ms) | 0.84 | 0.42 | 0.21 | 0.105 |
| On-demand remote fetch (ms) | 1.08 | 0.69 | 0.45 | 0.36 |
| Single-read + NoC broadcast (ms) | 0.94 | 0.56 | 0.33 | 0.20 |
| Intra-group TP (ms) | 1.00 | 0.64 | 0.42 | 0.30 |
| **CMN (ms)** | **0.88** | **0.46** | **0.25** | **0.14** |

Fixed router stages contribute 63.2% of average flit latency (vs. only 7.8% from dynamic waiting), motivating the bypass of per-hop packet processing.

### 4. Online Scheduling and Remapping

**Intra-Group CMN Scheduling** (Algorithm 1): Greedy heuristic that pairs overloaded source cores with underloaded helpers, constructs link-disjoint collaboration domains (max 16 cores per stack), and generates ordered broadcast schedules after accounting for CMN cost.

**Tiered Online Remapping**:
- **Tier-1**: Redistributes load fractions among existing expert replicas (migration-free, runs before every MoE layer).
- **Tier-2**: Adds/removes expert replicas (requires one weight transfer).
- **Tier-3**: Swaps experts across groups (requires two weight migrations).

**Cost-aware trigger**: EMA-based prediction with $\alpha = 0.3$; Tier-2/3 candidates checked every eight decode steps or when group-level M2M exceeds steady-state by 20% for three consecutive steps.

---

## Empirical Validation / Results

### Experimental Setup

- **Simulator**: Cycle-accurate simulator built on Ramulator 2.0 (HBM3 timing) and BookSim (NoC traffic), with RTL-derived CMN costs.
- **Models**: DeepSeek-V2-Lite (16B, 64 experts), Qwen3-30B-A3B (31B, 128 experts), GLM-4.5-Air (106B, 128 experts).
- **Baselines**: GPU (vLLM), EP, TP, HD-MoE, PS (EPLB-based), WS (work stealing).

### Generalization of Offline Grouping

- **Cross-dataset reuse**: Ours-Reuse lowers group-level M2M by **20–24%** over History-Reuse and closes **86%** of the gap to Oracle-Reopt.
- **Cross-window stability**: Only 9% degradation across successive 256-step windows; remapping further reduces later-window M2M.
- **Cross-batch stability**: Stable from batch size 32 to 768, reducing M2M by 22% over History-Reuse.

### MoE Layer Decode Performance

| Metric | Speedup |
|--------|---------|
| vs. GPU | 5.87× average |
| vs. EP | 4.63× average |
| vs. HD-MoE | 1.72× average |

3D-MoE achieves **83% average compute resource utilization**, more than 4.25× that of static deployment strategies.

### End-to-End Performance

- **3D-MoE outperforms the strongest NMP baseline by 1.58×–1.86×** across all workloads.
- vs. GPU: **5.54×, 7.88×, and 6.95×** speedups across the three models.
- **Decode-token latency (TBT)**: Reduces p50/p90/p99 TBT over HD-MoE by **45.3%/43.7%/26.0%** on DeepSeek-V2-Lite.

### Ablation Study

| Configuration | Speedup (relative to Random) |
|---------------|------------------------------|
| +G/M-Coarse | 1.35× |
| +G/M-Steps | 1.72× (vs. +G/M-Coarse) |
| +Hardware | 1.77× |
| +Remapping | 1.11× (additional, dynamic workloads) |

### Hardware Overhead

| Component | Area (mm²) | Area ratio | Peak power (W) |
|-----------|-----------|------------|----------------|
| DFU | 0.56 | 0.94% | 0.172 |
| RCL | 0.156 | 0.28% | 0.033 |
| DRU | 0.26 | 0.42% | 0.080 |
| Look-Up Table | 0.28 | 0.47% | 0.052 |
| **Total** | **1.256** | **2.11%** | **0.337** |

---

## Theoretical and Practical Implications

### Theoretical Significance

1. **Complementarity as a deployment primitive**: The paper establishes that MoE routing skew is not purely random — domain-specialized routing creates persistent, statistically detectable load-cancellation structures that can be exploited for load balancing without explicit domain labels.

2. **Hierarchical balance approach**: The three-level strategy (complementarity-guided grouping → CMN-assisted intra-group collaboration → online remapping) decomposes an intractable global load-balancing problem into tractable sub-problems at different time scales.

3. **Hardware–software co-design principle**: The paper demonstrates that a lightweight, specialized hardware path (CMN) can be dramatically more effective than general-purpose NoC mechanisms for structured communication patterns (one-to-many weight streaming), achieving near-ideal latency with minimal area overhead.

### Practical Implications

1. **EP remains the right backbone for NMP**: The results validate that EP's locality benefits can be retained while mitigating imbalance through complementarity-guided placement, avoiding TP's communication overhead.

2. **One-time offline optimization suffices**: A single offline plan (1–3 hours preprocessing) generalizes across datasets, time windows, and batch sizes, making deployment practical.

3. **Scalability**: 3D-MoE scales to larger models (e.g., 236B DeepSeek-V2 on eight NMP devices) with 3.51× throughput speedup over GPU, and remains beneficial even with low expert counts (8-expert Llama model: 1.2× gain).

---

## Conclusion

This paper presents **3D-MoE**, a hardware–software co-designed system for MoE decode on 3D-stacked NMP. The key contributions are:

1. **Identification and validation of expert complementarity** — a trace-level load-cancellation structure in MoE decode that is stable across datasets, time windows, and batch sizes.

2. **Complementarity-guided offline grouping and mapping** — formulated as hierarchical ILP stages that balance inter-group load while accounting for intra-group collaboration and routing costs.

3. **Collaborative Memory-Network (CMN)** — a lightweight hardware fast path that reads each hot expert's weight tile once and broadcasts it to helper cores over a reconfigured local path, avoiding repeated DRAM reads and per-hop NoC processing.

4. **Tiered online remapping** — cost-aware adjustments that maintain robustness under workload drift with minimal overhead (0.87% exposed end-to-end overhead).

**Key results**: 4.63× average MoE-layer speedup over EP, 1.58×–1.86× end-to-end throughput gains over the strongest NMP baselines, 5.54×–7.88× over GPU, with only 2.11% additional logic-die area.

**Future directions** implied by the work include: extending complementarity analysis to training-time expert assignment, exploring cross-device complementarity for larger-scale deployments, and investigating whether similar load-cancellation structures exist in other sparse model architectures beyond MoE.

---

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