A3D-MoE: 3D-Integrated MoE Accelerator

• The paper introduces a co-designed accelerator that leverages 3D heterogeneous integration and an adaptive GEMV-GEMM systolic array to tackle irregular computation in MoE serving.
• It utilizes a hardware resource-aware operation fusion scheduler and MoE score-aware HBM access reduction to optimize throughput and reduce latency.
• End-to-end evaluations show 1.8× to 2× latency reduction and up to 4× energy savings compared to current state-of-the-art MoE serving systems.

Searching arXiv for the specified paper and closely related work to ground the article. A3D-MoE is a hardware-software co-designed accelerator for serving fine-grained Mixture-of-Experts (MoE) LLMs efficiently through 3D heterogeneous integration, a 3D-Adaptive GEMV-GEMM-ratio systolic array, a hardware resource-aware operation fusion scheduler, and MoE score-aware HBM access reduction with even-odd expert placement (Huang et al., 25 Jul 2025). It is motivated by the observation that sparse MoE activation reduces the number of activated parameters per token, yet end-to-end inference remains memory- and latency-heavy because fine-grained expert selection, continuous batching, and long decode phases create highly irregular computation and memory-access patterns. In this design, the central target is not generic transformer acceleration, but the specific irregularity of MoE serving under mixed prefill and decode workloads, where variable GEMV-GEMM ratios, non-fused attention and MoE execution, and expert-weight traffic to HBM remain dominant inefficiencies (Huang et al., 25 Jul 2025).

1. Problem formulation and architectural rationale

A3D-MoE is formulated around three bottlenecks in fine-grained MoE inference. First, runtime workloads exhibit irregular GEMV/GEMM ratios: prefill is GEMM-heavy, decode is GEMV-heavy, and many low-arithmetic-intensity GEMM cases behave almost like GEMV. Second, conventional scheduling does not fuse attention operations with MoE operations effectively, particularly in settings such as Sarathi-Serve chunked prefill, where mixed prefill and decode execution creates bubbles and idle GEMM resources. Third, expert weights still have to be fetched from HBM or DRAM, and fine-grained expert usage yields poor reuse, so bandwidth demand and energy remain high despite sparse activation (Huang et al., 25 Jul 2025).

The paper’s conceptual claim is therefore narrower and more technical than the common shorthand that “MoE is efficient.” Its position is that MoE reduces FLOPs, but not necessarily memory movement. For a token $x$, routing is defined as

$G(x) = \text{softmax}(xW_g),$

with sparse expert set $S(x)$, and output

$\text{MoE}(x) = \sum_{j \in S(x)} G_j(x)\,E_j(x),$

where each expert is

$E_j(x) = (\text{ACT}(xW_G)\odot (xW_U))W_D.$

The implication drawn in the paper is that sparse conditional computation does not by itself solve bandwidth pressure, low reuse, or workload-shape instability; accelerator design must explicitly target those phenomena (Huang et al., 25 Jul 2025).

This framing distinguishes A3D-MoE from dense-transformer accelerators and from software-only serving optimizations. It also helps explain why the proposal combines packaging, dataflow, scheduling, and precision-aware memory access rather than treating these as separable optimizations.

2. 3D heterogeneous integration and system organization

The system architecture uses 3D heterogeneous integration to tightly couple compute and memory. Its stack consists of a compute logic die at the bottom, an HBM base logic die in the middle, and HBM DRAM dies stacked on top. The design uses fine-grained HBM with TSV-based vertical connectivity and bumpless hybrid bonding, and it integrates a V-Cache-like stacked SRAM on the memory side (Huang et al., 25 Jul 2025).

The paper attributes three benefits to this organization relative to a 2.5D interposer design. Vertical stacking provides very wide, dense, low-latency interfaces between compute and HBM; in the modeled setup, the system achieves 9600 GB/s per HBM. Direct movement through TSVs reduces routing area and runtime NoC energy by avoiding large 2D on-chip traversal. The design also eliminates SerDes overhead by replacing conventional interposer-style interfaces with direct vertical links. Type-2 SRAM, described as V-Cache, stores reused weights and activations closer to compute and is particularly important for small GEMM/GEMV-like fragments with limited reuse (Huang et al., 25 Jul 2025).

On the compute side, the accelerator contains two structures: a normal systolic array for conventional GEMM execution and a 3D-Adaptive GEMV-GEMM-ratio systolic array that can operate in GEMM or GEMV-like modes. This heterogeneous compute composition is essential to the paper’s argument. A fixed accelerator balance is treated as structurally mismatched to fine-grained MoE serving because decode-heavy, prefill-heavy, and mixed regimes stress different parts of the machine at different times (Huang et al., 25 Jul 2025).

A recurring misconception is that A3D-MoE is primarily a packaging paper. The system description does not support that reading. Packaging is one pillar, but the proposal is explicitly a full-stack accelerator design in which 3D integration is coupled to runtime reconfiguration, operation fusion, and score-aware memory reduction.

3. Adaptive compute substrate and unified 3D dataflow

The main compute innovation is the 3D-Adaptive GEMV-GEMM-ratio systolic array. The motivation is that fine-grained MoE serving produces arbitrary and time-varying GEMV/GEMM ratios that fixed hardware cannot handle efficiently. In conventional systolic arrays, GEMV-like execution leads to large PE idleness and high pipeline-fill cost relative to useful work, and low-arithmetic-intensity GEMM can become effectively memory-bound. A3D-MoE addresses this by making the 3D-systolic array run-time reconfigurable and able to switch between modes as components become idle (Huang et al., 25 Jul 2025).

The array supports three execution modes: 3D-GEMM mode for high-arithmetic-intensity matrix operations, 3D-GEMV mode for vector-like or low-arithmetic-intensity operations, and 3D-GEMV-VCache mode to exploit stacked SRAM reuse. This is paired with a unified 3D dataflow that uses vertical TSVs to move weights, inputs, and outputs while relying on spatial interleaving to reduce pipeline fill and skew overhead. In the input-stationary GEMM case, for example, A3D-MoE loads input and weight blocks in parallel through TSVs and then corrects output ordering with a de-interleaving circuit before writing back to V-Cache or HBM (Huang et al., 25 Jul 2025).

The same logic extends to GEMV. Rather than collapsing accumulation into a single-cycle SIMD-style critical path, the 3D design spreads computation across cycles while preserving efficient vertical data movement. For low-arithmetic-intensity GEMM/GEMV-like cases, the architecture can decompose low-intensity GEMM into multiple GEMVs and use V-Cache for reuse. This suggests that the paper interprets “GEMM versus GEMV” less as a rigid operator taxonomy than as a workload-shape continuum that the hardware should accommodate dynamically.

The transformer primitives used to ground this design are standard but operationally important in the paper. Multi-head attention is written as

$$Q_i = XW_i^Q,\quad K_i = XW_i^K,\quad V_i = XW_i^V,$$

$$S_i = \text{softmax}(Q_iK_i^T/\sqrt{d_k}),\quad O_i = S_iV_i,$$

$\text{MHA}(X) = \text{Concat}(O_1,\ldots,O_h)W^O,$

and the FFN as

$$\text{FFN}(X) = \big(\text{ACT}(XW_G)\odot (XW_U)\big)W_D.$$

These expressions are used to argue that real MoE serving alternates among compute patterns whose arithmetic intensity and parallelism properties vary sharply over time (Huang et al., 25 Jul 2025).

4. Scheduling, operation fusion, and HBM traffic reduction

The hardware resource-aware operation fusion scheduler, abbreviated HR-OFS, addresses the second major bottleneck: poor interaction between attention and MoE execution. In mixed serving regimes, especially with test-time computation and Sarathi-Serve chunked prefill, conventional systems wait until all attention work completes before starting MoE, which wastes time and compute. HR-OFS instead fuses and reorders operations to maximize and balance compute and bandwidth utilization (Huang et al., 25 Jul 2025).

Its policy is based on ranking tokens or expert paths by arithmetic intensity and then scheduling QKV generation, prefill attention, decode attention, and High-AI, Mid-AI, or Low-AI MoE accordingly. In decode-dominated cases, it prioritizes QKV generation for high-AI tokens, completes prefill attention, and then executes High-AI MoE and decode attention concurrently. In prefill-dominated cases, it prioritizes low-AI tokens during QKV generation and then schedules Low-AI MoE as soon as decoding finishes and low-AI prefill tokens are ready, so compute-bound prefill attention overlaps with memory-bound MoE. The scheduler begins from layer 4 because expert prediction in early layers is less accurate. In ablation, HR-OFS alone reduces 99th-percentile token-between-token latency by 1.42× to 1.86× across the evaluated models and settings (Huang et al., 25 Jul 2025).

The third major optimization, MoE-HBMR-EOP, targets memory traffic directly. After top-$K$ expert selection, the accelerator uses MoE scores to determine whether a selected expert should be fetched in full BF-16 or in reduced FP-8 form. For experts with scores below a threshold, only FP-8 data is fetched from HBM. The threshold is set empirically at 0.45 after min-max normalization because accuracy is preserved at that level. Since BF-16 and FP-8 exponents are not directly aligned, the paper introduces an offline profiling-based lookup mechanism that analyzes exponent distributions, observes that most exponent values fit in 4 bits, stores exponent-range information in a shared regular data map, and stores rare outliers separately in an outlier map. At runtime, the HBM or logic die reconstructs the correct FP-8 exponent with negligible overhead (Huang et al., 25 Jul 2025).

Even-odd expert placement complements this scheme at the memory layout level. FP-8 data are stored in odd rows of the DRAM cell array and remaining BF-16 data in even rows, allowing the HBM interface to remain fully utilized whether the system fetches FP-8 or BF-16 data. The reported effect is a 1.35× to 1.44× reduction in DRAM access count (Huang et al., 25 Jul 2025).

A common misconception is that these mechanisms merely compress weights. The paper presents them more specifically as score-aware, routing-conditioned bandwidth control and bus-utilization optimization. Precision reduction is only one part of the broader HBM-access design.

5. Evaluation methodology and quantitative results

A3D-MoE is evaluated on three fine-grained MoE models: OLMoE-1B-7B, DeepSeek-V2-Lite, and Qwen1.5-MoE-A2.7B (Huang et al., 25 Jul 2025). Expert-usage statistics are profiled using MMLU, MATH-500, and LiveCodeBench. Serving behavior is simulated with batch size, prefill length $G(x) = \text{softmax}(xW_g),$0, and decode length $G(x) = \text{softmax}(xW_g),$1, with the note that $G(x) = \text{softmax}(xW_g),$2 to reflect test-time compute conditions. The main baselines are NeuPIM and Duplex. The evaluation toolchain combines NeuroSim, Maestro, and Ramulator, with Verilog implementation and synthesis in a 7 nm PDK and thermal analysis using Ansys (Huang et al., 25 Jul 2025).

The headline quantitative results are end-to-end. A3D-MoE reports 1.8× to 2× latency reduction, 1.44× to 1.8× throughput improvement, and 2× to 4× energy reduction relative to the state of the art (Huang et al., 25 Jul 2025). The paper further specifies that throughput improves by up to 1.6× to 1.8× compared with NeuPIM and about 1.2× to 1.44× compared with Duplex, while energy is about 1.9× lower than Duplex and about 3.4× lower energy than NeuPIM in some settings, with average reductions around 2× vs Duplex and 4× vs NeuPIM (Huang et al., 25 Jul 2025).

These gains are distributed across the three principal mechanisms rather than attributed solely to 3D integration. HR-OFS produces significant token-latency reductions on its own, while MoE-HBMR-EOP reduces DRAM access count by 1.35× to 1.44×. The combined interpretation advanced by the paper is that compute utilization, attention–MoE overlap, and expert-weight traffic must be optimized jointly to obtain the reported end-to-end improvements (Huang et al., 25 Jul 2025).

Thermal constraints are treated explicitly. Because 3D integration raises thermal challenges, the paper evaluates operation with and without liquid cooling and limits operation so DRAM stays below 95°C. This detail is important because it constrains the practical operating envelope of the proposed vertical stack and prevents the performance claims from being read as thermally unconstrained peak numbers (Huang et al., 25 Jul 2025).

6. Trade-offs, assumptions, and relation to adjacent work

A3D-MoE relies on several assumptions that shape its applicability. HR-OFS and MoE-HBMR-EOP depend on profiling and statistics, including expert-usage probability, exponent distribution, and score thresholding, so the design assumes that these distributions are sufficiently stable to guide scheduling and compression decisions. The scheduler begins from layer 4 because early-layer expert prediction accuracy is low. The architecture also requires a 3D heterogeneous package, custom TSV-based interconnect, V-Cache integration, adaptive compute modes, and lookup/recovery logic for reduced-precision expert reads, so the reported gains come with substantial implementation complexity (Huang et al., 25 Jul 2025).

In the broader literature, A3D-MoE is best read as an inference-centric counterpart to other MoE-oriented hardware co-design efforts. “Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures” addresses training rather than serving and argues that sparse expert execution must be matched with expert-aware placement, balanced cluster allocation, and micro-batch-aware streaming on a topology that mirrors the model’s modular structure (Luo et al., 7 Mar 2026). A plausible implication is that both works converge on the same systems principle: MoE irregularity must be made hardware-native, but they operationalize that principle differently—A3D-MoE through 3D heterogeneous integration and adaptive inference dataflow, Mozart through 3.5D wafer-scale chiplet modularity and training-time expert placement.

The name should also not be conflated with “A3D: Agentic AI flow for autonomous Accelerator Design,” which is not a learned MoE model with a router gating tokens to different experts, but a modular multi-agent orchestration system for accelerator design automation (Nallathambi et al., 14 May 2026). That distinction matters because “A3D” in one context denotes an autonomous design flow, whereas in A3D-MoE it denotes a 3D heterogeneous integration strategy for MoE inference acceleration.

Taken together, the paper’s main contribution is a specific systems claim: efficient fine-grained MoE serving requires co-design across packaging, compute substrate, runtime scheduling, and HBM access policy. The architecture is not presented as a universal LLM accelerator, but as a targeted response to dynamic workload imbalance, attention/MoE scheduling conflicts, and memory bandwidth and energy costs in fine-grained MoE LLMs (Huang et al., 25 Jul 2025).