PaperlayerView exact revision on arXiv

Source-linked AI summary

HDA-MoE: Hybrid Parallelism and Dynamic, Adaptive Scheduling for Mixture-of-Experts with 3D Near-Memory Processing

Haochen Huang, Shuzhang Zhong, Shengxuan Qiu, Zhe Zhang, Shuangchen Li, Cong Li, Dimin Niu, Hongzhong Zheng, Guangyu Sun, Runsheng Wang, Meng Li

arXiv:2609.08682v1cs.AR

TL;DR

Sparse and dynamic MoE routing complicates inference on distributed 3D NMP, where TP and EP have different communication and load-balance trade-offs. HDA-MoE combines hybrid placement with runtime scheduling and hardware-aware gating. It reports speedups over TP, EP, a compute-balanced Hybrid TP-EP baseline, and HD-MoE.

  • Problem

    Sparse MoE activation creates irregular memory accesses and memory bottlenecks, while dynamic expert activation makes static deployment insufficient on memory-constrained 3D NMP.

    Source: §I · p. 1

  • Method

    HDA-MoE combines TP and EP placement with runtime scheduling and hardware-aware gating for MoE inference on 3D NMP.

    Source: §I · pp. 1–2

  • Results

    HDA-MoE achieves speedups of 1.1×–3.4× over TP, 1.1×–1.5× over EP, 1.1×–3.7× over the Hybrid TP-EP compute-balanced baseline, and 1.1×–1.3× over HD-MoE.

    Source: B. End-to-End Performance · p. 9

  • Takeaways & Limitations

    The reported comparisons support HDA-MoE as an effective approach to MoE inference on 3D NMP relative to the evaluated parallelization strategies.

    Source: §VII · p. 12

Abstract

from arXiv · show

Mixture-of-Experts (MoE) architectures have become a key technique for scaling Large Language Models (LLMs), enabling high model capacity with reduced computational cost. However, this efficiency comes at the expense of increased memory capacity and bandwidth demands. Recent 3D Near-Memory Processing (NMP) architectures, which vertically integrate memory and compute through hybrid bonding, provide high internal bandwidth and energy efficiency, making them attractive for accelerating MoE inference. Nevertheless, the distributed memory and compute organization of NMP systems introduces new challenges for mapping MoE workloads. Existing parallelization strategies, such as Tensor Parallelism (TP) and Expert Parallelism (EP), suffer from either high communication costs or unbalanced computation utilization, leading to inferior efficiency. In addition, the dynamic routing behavior of MoE models further complicates efficient deployment. To address these challenges, we present HDA-MoE, a framework that optimizes MoE execution on NMP architectures through hybrid parallel deployment and runtime scheduling. HDA-MoE integrates an offline hybrid parallel mapping algorithm with an online dynamic and adaptive scheduling mechanism to reduce communication overhead while improving computation utilization. Experimental results show that HDA-MoE achieves a speedup of 1.1x--3.4x over TP, 1.1x--1.5x over EP, 1.1x--3.7x over the Hybrid TP-EP compute-balanced baseline, and 1.1x--1.3x over HD-MoE. Source code is available at https://github.com/PKU-SEC-Lab/HDA-MoE-TCAD26.

I. INTRODUCTION

Sparse MoE routing strains memory and bandwidth, while distributed 3D NMP makes expert placement a trade-off between communication and workload balance. HDA-MoE combines hybrid placement with runtime scheduling and hardware-aware gating to address these challenges.

I. INTRODUCTION

  • On 3D NMP, TP balances computation but incurs heavy all-reduce communication, whereas EP reduces communication but risks workload imbalance.Dynamic expert activation also makes static deployment insufficient, and replication is less practical under NMP memory constraints.Source: §I · p. 1
  • HDA-MoE combines EP for low-activation experts with TP-style partitioning for frequent experts to reduce communication and improve utilization.Its placement framework jointly optimizes workload balance and communication overhead.Source: §I · pp. 1–2
  • A unified performance model estimates MoE computation and communication costs to guide deployment optimization.Experiments report improvements in TBT and speedup over strong baseline methods.Source: §I · p. 2
  • Runtime scheduling targets latency from short-term expert-activation fluctuations under a given routing pattern.The framework also uses online dynamic placement to handle dynamic activation patterns.Source: §I · p. 2
  • Hardware-aware gating aims to reduce routing-induced computation and communication bottlenecks with limited model perturbation.This journal version adds hardware-aware gating and analysis of low-impact expert substitutability.Source: §I · p. 2

A. Mixture-of-Experts Models

The section introduces sparse MoE execution and 3D NMP, then explains the TP–EP trade-off and routing behaviors that motivate HDA-MoE. It also begins the performance model with a computation-latency formulation based on node workload.

A. Mixture-of-Experts Models

  • MoE activates only a few experts per token to reduce computation, while dynamically routed experts create irregular activation patterns that complicate deployment.Shared experts activate for every token and can be assigned statically.Source: A. Mixture-of-Experts Models · p. 2

B. Near-Memory Processing Architectures

  • 3D NMP stacks DRAM and logic dies to provide high bandwidth with moderate compute density; Hybrid Bonding DRAM uses fine-pitch vertical links for low-latency, high-bandwidth communication.The cited work presents this architecture as suitable for memory-intensive LLM workloads.Source: B. Near-Memory Processing Architectures · p. 2

C. Distributed Inference Strategies

  • TP partitions experts across processing elements but requires bandwidth-intensive all-reduce, whereas EP assigns complete experts and can suffer compute imbalance and irregular communication.HDA-MoE combines the strategies to balance computation and communication.Source: C. Distributed Inference Strategies · p. 2

III. MOTIVATION

  • Routing exhibits long-term activation skew and co-activation affinity, short-term changes with high next-layer predictability, and hardware-cost mismatch at near-boundary expert choices.These patterns motivate hybrid placement, runtime scheduling, and hardware-aware gating that focuses on lower-impact boundary choices.Source: §III · pp. 2–3

IV. PERFORMANCE MODELING

  • The authors develop a performance model to estimate computation and communication costs for MoE inference on NMP architectures.The supplied section begins by describing the computation-latency component.Source: §IV · p. 3, A. Computation Overhead Modeling · p. 3

A. Computation Overhead Modeling

  • Computation latency is determined by the most heavily loaded node, whose workload depends on the fraction of each expert assigned to it and the number of activating tokens.The model treats expert placement as continuous, allowing partial distribution across nodes while introducing a balance-versus-communication trade-off.Source: A. Computation Overhead Modeling · p. 3

B. Communication Overhead Modeling

The communication model combines topology-aware discrete-event simulation with a calibrated node-traffic approximation for deployment optimization. Validation compares the model with both the simulator and ASTRA-sim.

B. Communication Overhead Modeling

  • The discrete-event simulator generates tasks from token–expert mappings, builds topology-specific paths, and schedules transfers while tracking directed-link occupancy.A priority queue processes the earliest-ready hop; occupied links delay and reschedule tasks.Source: B. Communication Overhead Modeling · pp. 3–4
  • The simulator supports Mesh, Torus, and Fat-tree by adapting path construction and link-bandwidth models to each topology.It tracks occupancy and bandwidth per directed link, including independently assigned bandwidths in Fat-tree.Source: B. Communication Overhead Modeling · p. 3
  • The optimization proxy estimates communication latency from node traffic, using endpoint-side bandwidth and a fitted scaling factor.The simulator instead uses the bandwidth of each traversed link; the scaling factor is refit when topology or bandwidth settings change.Source: B. Communication Overhead Modeling · p. 4
  • R2 exceeds 0.9 across evaluated scenarios for the estimated and simulated communication latencies; Fig. 4 reports R2 = 0.96.The latency estimate is calibrated from simulated samples using linear regression.Source: B. Communication Overhead Modeling · p. 4, Figure 4 · p. 4
  • Within 3.2% relative error, the performance model aligns with ASTRA-sim across evaluated ring all-reduce settings.The paper reports this comparison as validation of the model’s accuracy.Source: B. Communication Overhead Modeling · p. 4, Table I · p. 4

V. HDA-MOE DESIGN

HDA-MoE combines offline hybrid placement and physical mapping with online decisions at runtime. Its two-stage deployment optimization balances computation and communication while addressing node load and link congestion.

A. Overview

  • At runtime, dynamic scheduling adjusts how routed tokens use expert replicas, while hardware-aware gating adjusts low-impact expert choices under a shared latency model.The two mechanisms act respectively on hardware supply and software demand, without changing routing scores in dynamic scheduling.Source: A. Overview · pp. 4–5

B. Automated Hybrid Parallelism via Node-Link Balance

  • The Node-Link Balance strategy first optimizes logical expert placement for computation balance and lower communication volume, then maps clusters to physical nodes to alleviate link congestion.The mapping uses model configuration, workload profile, and hardware setting as inputs.Source: B. Automated Hybrid Parallelism via… · p. 5

1) Node Balance Optimization via Linear Programming:

  • The Node Balance linear program jointly models computation bottlenecks and approximated communication cost when optimizing expert placement across logical clusters.It abstracts away physical routing and uses configured bandwidth in the communication proxy.Source: 1) Node Balance Optimization via… · p. 5
  • Continuous P_ic values let the optimizer choose per-expert EP-style placement or TP-style partitioning instead of requiring manual mode labels.Z_ic indicates whether expert i is placed on cluster c; Y_gc indicates whether cluster c stores an expert from co-activated group g.Source: 1) Node Balance Optimization via… · p. 5
  • The placement constraints model communication at co-activated-group granularity and limit per-node expert-weight storage to respect memory capacity.The memory-load parameter controls permitted excess over uniform average: tighter bounds enforce more balanced storage, while looser bounds give the latency objective more placement freedom.Source: 1) Node Balance Optimization via… · pp. 5–6
  • The LP minimizes node-level inference overhead to produce a globally coordinated logical placement that balances computation and communication for subsequent physical mapping.The objective includes computation latency and twice the calibrated communication estimate, accounting for symmetric all-to-all dispatch and gather.Source: 1) Node Balance Optimization via… · p. 6

2) Link Balancing via Bayesian Optimization:

HDA-MoE maps logical clusters to physical nodes to reduce link congestion and communication tail latency, then schedules experts at runtime using predicted priorities. Its placement balances compute and communication, while runtime scheduling anticipates expert hotspots.

2) Link Balancing via Bayesian Optimization:

  • Node-Link Balance remaps logical clusters onto physical nodes, spreading synchronization traffic over less congested links while preserving computation balance.Node Balance first splits overloaded experts to remove compute hotspots, but can leave synchronization traffic concentrated on specific paths.Source: 2) Link Balancing via Bayesian… · p. 6, Figure 6 · p. 6

C. Dynamic Scheduling

  • Runtime scheduling combines expert-priority prediction, cost-aware pre-broadcasting, and communication-efficient token dispatch.Source: C. Dynamic Scheduling · p. 6

1) Priority Detection and Computation Prediction:

  • Applying the next layer’s gate to the previous layer’s hidden states predicts next-layer computation hotspots; priority scores rank experts by expected latency benefit from pre-broadcasting.The highest-priority expert on the most congested node is repeatedly pre-broadcast within the previous layer’s runtime budget, with metadata piggybacked on the lightweight exchange.Source: 1) Priority Detection and Computation… · pp. 6–7

2) Optimal Broadcast Chunk Size:

HDA-MoE chooses pre-broadcast chunks to balance hop count against per-hop traffic, then dispatches tokens to existing expert copies. Hardware-aware gating adds computation and communication penalties to routing scores.

2) Optimal Broadcast Chunk Size:

  • Larger pre-broadcast chunks reduce hop count and latency, while smaller chunks reduce per-hop traffic but increase transmission delay.The paper models this chunk-size trade-off with the traditional α–β communication model.Source: 2) Optimal Broadcast Chunk Size · p. 7
  • The pre-broadcast latency lower bound is tight when chunk size is chosen optimally under the runtime-window constraint.The constraint is the number of pre-broadcast iterations allowed within the runtime window, denoted k.Source: 2) Optimal Broadcast Chunk Size · p. 7

3) Communication-Efficient Dispatch:

  • Dispatch selects the least-loaded node that already holds a token’s routed experts, reducing workload imbalance without extra data movement.The runtime rule is greedy: global optimization would require solving a discrete scheduling problem under current link and compute queues each window.Source: 3) Communication-Efficient Dispatch · p. 7, Figure 7 · p. 7

D. Hardware-aware Gating

  • Hardware-aware gating adds marginal computation and communication costs to expert routing scores, with typically negative penalty weights controlling the efficiency–fidelity trade-off.It complements pre-broadcasting when runtime bandwidth and budget cannot cover multiple simultaneous hotspots.Source: 1) Objective · p. 7, D. Hardware-aware Gating · p. 7

2) Computation Penalty:

The computation penalty estimates how selecting a candidate expert changes maximum node load. Hardware-aware gating also uses communication cost, and the evaluation includes four MoE models.

2) Computation Penalty:

  • The computation penalty estimates the marginal increase in maximum node compute load from selecting each candidate expert.Experts that increase the maximum node load receive a penalty.Source: 2) Computation Penalty · p. 8

3) Communication Penalty:

  • Deployment dispersion serves as a lightweight proxy for dispatch and aggregation communication cost.The Figure 8 example shows gating redirecting low-impact tokens from a high-cost boundary choice toward available experts on less-loaded nodes.Source: 3) Communication Penalty · p. 8, Figure 8 · p. 8

1) Models:

  • Evaluation covers Mixtral, DeepSeek, Qwen2, and Qwen3.5, whose 256 routed experts test larger expert spaces.Source: 1) Models · p. 8

2) Baselines:

The evaluation compares HDA-MoE with TP, EP, a compute-balanced hybrid, and HD-MoE across latency settings and workloads. HDA-MoE outperforms these baselines and maintains low TBT across mesh sizes and alternative network topologies.

2) Baselines:

  • Baselines are TP, EP, a compute-balanced hybrid using EP across mesh regions and TP within them, and HD-MoE without hardware-aware gating.Source: 2) Baselines · p. 8

3) Evaluation Metrics:

  • The reported metrics are normalized TBT, normalized to TP, and a MoE latency breakdown separating node-local computation from inter-node communication.Source: 3) Evaluation Metrics · p. 8

4) Dataset:

  • MT-Bench measures latency under conversational workloads, while GSM8K, HumanEval, ARC-E, and ARC-C measure model quality.Source: 4) Dataset · p. 8

5) Hardware-aware Gating Parameters:

  • Computation and communication penalty weights are tuned per model and compute/bandwidth setting, then reused across topologies with the same hardware setting.The two penalty magnitudes are not comparable because their underlying cost terms have different model-dependent scales.Source: 5) Hardware-aware Gating Parameters · pp. 8–9, Table IV · p. 9

6) Implementation and Offline Optimization:

  • LP finishes within minutes for Mixtral, DeepSeek, and Qwen3.5, while Qwen2 needs longer; Bayesian Optimization stabilizes within about 45 of its 70 iterations.The offline search typically takes several hours; Qwen2 takes longer because its co-activations span more distinct expert groups, increasing placement variables.Source: 6) Implementation and Offline… · p. 9, Figure 9 · p. 9

B. End-to-End Performance

  • End-to-end evaluation spans 2.5 TFLOPS/75 GB/s, 5 TFLOPS/50 GB/s, and 10 TFLOPS/25 GB/s settings, plus varied mesh sizes and network topologies.The topology comparison holds 32 endpoints and 50 GB/s constant; each node has 625 GB/s memory bandwidth.Source: B. End-to-End Performance · p. 9
  • HDA-MoE achieves 1.1×–3.4× over TP, 1.1×–1.5× over EP, 1.1×–3.7× over the Hybrid TP-EP compute-balanced baseline, and 1.1×–1.3× over HD-MoE.It maintains low TBT across mesh sizes and outperforms all baselines on Torus and Fat-tree.Source: B. End-to-End Performance · p. 9

C. Latency Breakdown

HDA-MoE reduces communication and compute-tail latency relative to TP and EP, and hardware-aware gating further improves over HD-MoE. The evaluation reports comparable accuracy and examines how routing changes affect model outputs, including model-dependent expert substitutability.

C. Latency Breakdown

  • Memory access is a small latency component; HDA-MoE reduces TP all-reduce communication and EP compute-tail latency, and hardware-aware gating further improves over HD-MoE.The breakdown separates computation, communication, and memory access.Source: C. Latency Breakdown · p. 9, Figure 13 · p. 10

D. Accuracy Evaluation

E. Routing Fidelity and Expert Substitutability Analysis

  • Non-Top-1 boundary replacements achieve 0.949 average similarity across four models, while Top-1 replacements are more sensitive.Fig. 14 indexes replacements by selected rank and after-top-k candidate rank.Source: E. Routing Fidelity and Expert… · p. 10, Figure 14 · p. 10
  • Routing changes have limited impact overall, but boundary-replacement sensitivity varies by model, with Mixtral more sensitive than DeepSeek, Qwen2, and Qwen3.5.The authors assess perturbation using routing retention and distributional measures; Table IX reports results across settings.Source: E. Routing Fidelity and Expert… · p. 11, Table IX · p. 11

F. Ablation Study

The ablations show how node placement and topology-aware link mapping address compute imbalance and communication congestion. Node Balance improves performance and load distribution, while Link Balance further reduces communication latency and congestion.

1) Node Balancing:

  • Node Balance achieves 1.0×–3.0× speedup over TP/EP and 1.5× over compute-balanced hybrid by improving compute balance and reducing communication volume.Fig. 15 evaluates Node Balance speedup.Source: 1) Node Balancing · p. 11
  • Node Balance reduces EP’s compute-tail latency by 2.0× on average, mitigating MoE load skew.Fig. 16 isolates the effect on compute latency.Source: 1) Node Balancing · p. 11, Figure 16 · p. 11
  • Compared with EP, Node Balance produces a more even per-node compute and communication load distribution and removes severe hotspots.Fig. 17 visualizes node-level loads before and after balancing.Source: 1) Node Balancing · p. 11, Figure 17 · p. 11

2) Link Balancing:

  • Topology-aware Link Balance achieves an average 1.2× communication-latency reduction versus Node Balance alone and reduces congestion by distributing mesh traffic more evenly.It also reduces communication latency compared with TP and the hybrid baseline.Source: 2) Link Balancing · p. 11, Figure 18 · p. 11, Figure 19 · p. 11

3) Memory-Capacity Constraint:

HDA-MoE’s final evaluation reports speedups over TP, EP, the hybrid TP-EP baseline, and HD-MoE. The accompanying studies examine memory balance, dynamic placement, hardware-aware gating, and gating-parameter sensitivity.

3) Memory-Capacity Constraint:

  • Tighter memory-capacity constraints decrease the max/mean DRAM-load ratio while normalized TBT stays close to the unconstrained case.The constraint balances expert-weight loads across DRAM nodes without sacrificing the end-to-end advantage.Source: 3) Memory-Capacity Constraint · pp. 11–12

4) Dynamic Placement Strategy:

  • Dynamic placement averages 1.15× and 1.25× speedup with 2 and 5 pre-broadcast experts, respectively, especially on math and coding traces that differ from reasoning.The comparison uses MT-Bench traces and contrasts static placement derived from reasoning traces with runtime-adaptive placement.Source: 4) Dynamic Placement Strategy · p. 12, Figure 21 · p. 12

5) Hardware-aware Gating:

  • Hardware-aware gating remains effective at small batches, where pre-broadcasting weakens, and combining the approaches improves speedup; accuracy stays around 70%, comparable to original gating.The evaluation uses GSM8K with DeepSeek across batch sizes.Source: 5) Hardware-aware Gating · p. 12, Figure 22 · p. 12

G. Layer-wise Sensitivity of Gating Parameters

H. Sensitivity to Gating Parameters

VII. CONCLUSION

  • HDA-MoE achieves 1.1×–3.4× speedup over TP, 1.1×–1.5× over EP, 1.1×–3.7× over the Hybrid TP-EP baseline, and 1.1×–1.3× over HD-MoE.The framework combines hybrid expert placement, runtime scheduling, and routing adaptation.Source: §VII · p. 12
Loading 2609.08682v1…