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

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:

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

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

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

Hardware Overhead

---

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.