Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures
Authors: Shuqing Luo
Ye Han
Pingzhi Li
Jiayin Qin
Jie Peng
Yang Katie Zhao
Yu Cao
Tianlong Chen
Venue: NeurIPS 2025 spotlight ()
Keywords: MoE
Chiplet Architecture
Post-Training Efficiency
Primary area: infrastructure
TL;DR: We propose a framework for efficient MoE post-training on 3.5D Wafer-scale chiplets.

Abstract: Mixture-of-Experts (MoE) architecture offers enhanced efficiency for Large Language Models (LLMs) with modularized computation, yet its inherent sparsity poses significant hardware deployment challenges, including memory locality issues, communication overhead, and inefficient computing resource utilization. Inspired by the modular organization of the human brain, we propose $\texttt{Mozart}$, a novel algorithm-hardware co-design framework tailored for efficient training of MoE-based LLMs on 3.5D wafer-scale chiplet architectures. On the algorithm side, $\texttt{Mozart}$ exploits the inherent modularity of chiplets and introduces: 
($1$) an expert allocation strategy that enables efficient on-package all-to-all communication, and ($2$) a fine-grained scheduling mechanism that improves communication-computation overlap through streaming tokens and experts. On the architecture side, $\texttt{Mozart}$ adaptively co-locates heterogeneous modules on specialized chiplets with a 2.5D NoP-Tree topology and hierarchical memory structure.
Evaluation across three popular MoE models demonstrates significant efficiency gains, enabling more effective parallelization and resource utilization for large-scale modularized MoE-LLMs.

Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures
Shuqing Luot! Han Yet², Pingzhi Litl, Jiayin Qin Jie Peng'. Yang (Katie) Zhao², Yu (Kevin) Cao², Tianlong Chen
I University of North Carolina at Chapel Hill Unnivrriity of Minnesota Twin Cities
Abstract
Mixture-of-Experts (MoE) architecture offers enhanced efficiency for Large Lan- guage Models (LLMs) with modularized computation, yet its inherent sparsity poses significant hardware deployment challenges, including memory locality issues, communication overhead, and inefficient computing resource utilization. Inspired by the modular organization of the human brain, we propose Mozart, a novel sagorithm-harwaarn co-design framework tailored for efficient training MoE-based LLMs on 3.5D wafer-scale chiplet architectures. On the algorithm side, Mozart exploits the inherent modularity of chiplets and introduces: (1) an expert al- location strategy that enables efficient on on-package all-to-all communication, and (2) fine-gr une-grained scheduling mechanism that improves communication communication-computaaii overlap through streaming tokens and experts. On the architecture side, Mozart adaptively co-locates heterogeneous modules on specialized chiplets with a 2.5D NoP-Tree topology and hierarchical memory structure. Evaluation across three pop- ular MoE models demonstrates significant efficiency gains, enabling more effective parallelization and resource utilization for large-scale modularized MoE-LLMs.
Introduction
The human brain, known for its cognitive efficiency and modular organization, has long inspired the design of large-scale computational systems [9, 14, 27]. It comprises specialized modules that handle distinct tasks, ranging from memory-intensive to computation-heavy operations, while maintaining low-latency coordination with adjacent regions [1, 3, 13]. This modularity enables efficient, scalable, and flexible processing [2, 17], which is principle increasingly adopted in deep learning systems such as Large Language Models (LLMs) [21, 35].
Meanwhile, recent advances in LLMMs, particularly Mixture-of- fixture-of-Experts (MoEs), reflect similar modu- lar principles by dynamically activating specialized sub-networks based on input. However, the scale and heterogeneity of MoE-LLMs pose significant challenges for conventional hardware platforms [21] (e.g., traditional GPUs or CPUs), including photoreticle-limite
scalability [18] and transistor scaling limits [36], as well as poor memory locality [12], high inter-module communication overhead [15], and inefficient resource utilization [25] due to dynamic and uneven computational workloads.
2.5D/3.5D heterogeneous chiplet-based architectures have gained popularity due to their scalability and modularity to meet the demands of the aforementioned LLM-related workloads, including MoEs [32, 46, 8, 20]. Typically in 2.5D designs, multiple chiplets are interconnected via Network- on-Package (NoP) through an interposer [7, 34, 26, 37], redueing the area and cost overhead of monolithic integration. To further boost inter-chiplet bandwidth, 3D integration techniques such
"Correspondence Tianlong Chen <tiamlongdcs. edu.. Equal Contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Shared Experts
Routed Experts
Attention
Dense FFN
Embedding LM Head
LLMMEE-BB-7B0925
93.1 %
6.9B
DeepSeek-MoE
91.3%
16.4B
Qwen3-30B-A3B
95.0%
30.5B
Qwen3-235B-A225
96.6%
235B
Parameters 0%
25%
75%
100%
Figure 1: Parameter distribution in modern MoE-LLMs across various scales. The routed experts module constitutes over 90% of the total parameters these architectures.
as Abrough-Silicoon Vias (TSVs) are employed along the vertical direction. However, prior works tend to neglect wafer-scale integration [8, 20, 46] and largely adopt coarse- grained, static workload partitioning strategies [32] that assume dense and uniform computation, tiling the model workload to different chiplets without incorporating system-level coordination and optimization [23]. None of these works consider or explicitly address the system-level design challenges posed by MoE- LLMs with fine- ene-giaieed modularity, which often results in excessive inter-chiplet communication and inefficient resource utilization.
Given the growing need for system-level coordination in modular LLM deployment, we propose Mozart, an algorithm-haun hardware co-design framework for efficiently mapping MoE-LLMs onto 3.5D wafer-scale chiplet platforms. Our contributions summarized as follows:
To efficiently train MoE-based LLMs, propose Mozart, comprehensive framework with optimization techniques including: (1) a specialized expert layout strategy placing frequently co-activated experts the same or adjacent chiplets, targeted at balanced workload across multiple chiplets, (2) a communication -efficient all-to-all strategy utilizing expert collaboration, and (3) a fiine-griined scheduling strategy for improved communication- computation overlap using streaming tokens and experts.
To better accommodate the modularized structure of MoE, we design wafer-scale 3.5D chiplet architecture, featuring tightly integrated 3D logic-on-memory stacks, 2D NoP-tree interconnect, and two-level memory hierarchy. It supports our proposed optimizations with low-latency on-chip activation reuse, communication-Ba aware expert clustering, and inter- leaved execution of communication and computation tailored for sparse MoE computing.
Mozart achieves over 1.9x acceleration compared to baseline methods when evaluated across three popular open-source MoE-LLMs with varying scales, demonstrating its potential to enhance parallelization efficiency and optimize resource utilization for the post-training deployment of large-scale modularized MoE-LLMs.
2
Related Works
Modularized LLMs. Modularized LLMs, also known as Mixture-of-Expprtt (MoB) [35], demon- strate exceptional efficiency on scaling model capacity, allowing significant parameter growth without proportional computational costs. This efficiency derives from replacing traditional dense feed- forward layers with sparse modularized components, where sophisticated routing mechanisms se- lectively direct input tokens to appropriate expert subnetworks. Models like Mixtral-8x77 [16] demonstrate how activating just two experts per token per layer can leverage a substantially larger parameter space, matching the performance of dense counterparts while dramatically reducing active parameter requirements. The architecture was further refined in DeepSeek-No introduced finegrained experts and shared experts improve specialization and parameter efficiency. The expert specialization phenomenon where routing networks learn to direct specific input patterns to dedicated experts- enhances processing proficiency [4, 22, 42). Complementing this, expert collaboration, the strategic co-activation of multiple experts for processing complex inputs, has recently minimized communication overhead through optimized expert placement and routing algo- rithms [24, 47]. In our work, we leverage these expert specialization and collaboration principles to enhance training efficiency specifically for 2.5D/3.5D wafer-scale chiplet architectures, where physical hardware modularity naturally complements the logical modularity of MoE systems.
[4, 10], which
2.5D/3.5D Chiplet for MI. Workloads. Chiplet-based architectures have emerged as a promising solution to support the growing computational demands of large-scale neural networks and LLMs.

---

Prior works such as Maestro [20], Cambricon-LLM [46], and ScalePoM [8] primarily focus on scle chiplet designs. Maestro adopts a 3D memory-on-loge structure to coordinate multiple small-scale systolic arrays for inference acceleration. Cambricon-LLLM integrates Neural Processing Units (NPUs) with flash-based chiplets for energy-effcient on-device inference, while ScalePoM explores hierarchical power delivery for the chiplets. However, these works mainly focus on LLM inference and do not consider wafer-scale integration, limiting their scalability.
In contrast, FRED [32] explores wafer-scale integration by leveraging high-bandwidth interconnects and in-newwork collective communication to accelerate LLM training. Nonetheless, it largely relies on coarse- -gaained, static workload partitioning strategies that assume dense and uniform computation. When applied to sparse and modular models such as MoEs, such strategies result in inefficient resource utilization and increased inter-chiplet communication. To overcome these limitations, we propose 3.5D heterogeneous chiplet architecture that combines vertical stacking with 2.5D NoP- Tree interconnects, providing high-bandwidth, energy-efficient communication while maintaining architectural modularity. Built upon this hardware foundation, we introduce a fine-grained modular partitioning and communication-awan scheduling framework tailored for the post-training process of sparse workloads like MoE. By aligning expert activation patterns with the chiplet topology. our design reduces redundant data movement and significantly improves system throughput under modular model execution.
3 Preliminary
3.1 Mixture-of-Experts
Formulation. Given an input token embedding æ, the output of MoE layer can be formulated as the weighted sum of outputs from the Ne experts (Eo,
MoE(x)
R(x),-E,
(1)
where R(x); the output of small gating network R(-) for the i-th expert. For each token, the MoE layer aggregates the output of & experts, determined by the indices of the top-k highest routing scores, derived from the Softmax value of gating function g(-). which usually a single linear layer:
R(x) top-x(Seflmax((g) (x))).))
(2)
Expert Parallelism Pipeline. Expert parallelism [11, 21] has been demonstrated to be the most efficient distributed training technique for MoE models, where different experts are scattered on different parallel units and the workloads are dispatched to each unit during both forward and backward pass. Specifically, typical MoE pipeline with expert parallelism in the forward pass can be formulated as Dispatch >> All-to-All -> Expert Computing -> All-to-All -> Combine [15].
3.2 Analyzing Expert Activation Prior
Mozart focuses on efficient post-training of MoE-LLMs chiplet systems. Before deployment, we first analyze the empirical prior of the routing policy and then develop scheduling algorithms to enhance post-training efficiency. Given an instruction tuning dataset, we first the prefilling stage of inference on it to get the routing choice of a large token batch B, and next we compute 2 metries:
Analyzing Workload Distribution across Individual Experts. We construct vector V with Ne elements to quantify the workload distribution across individual experts, where
N.
V; = Σ I((()), #CB
V;= = V.
Σ 1=0
VJ-
Analyzing Collaboration Pattern across Paired Experts. To uncover co-activation patterns among experts single MoE layer, we construct graph & with No nodes. This graph is represented by an adjacency matrix C E NN,,N, where Ci2j denotes the edge value between nodes i and j. We further normalize it with the maximum edge value confine all entries in the matrix to the interval [0, 1]:
Cij Σ (R(()) # R(w), # #0,, Ptj -Cas/ ZEB
max OSiJSN.-1
Cij-
3

---

(a) Expert Clustering Allocation Algorithms
(b) 3.5D Chiplet System Architecture
Atterion
Shared
Export
Roused
Expert
Expert Cusserig ningiing
Expert Alccation
Central (Atamioo) Chiplet
Parpphen (MoE) Chipket
Bwitch
DRAM
Sinige Tranafumee
P
P
Decoder Layer
Pg
Figure 2: Algorithm- Hardwarr Co-Design Diagram of Mozart. Mozart provides an algorithm- hardware co-design approach, and we present both the algorithum-level expert clustering & allocation schemes in the left part, and the architecture- 3.5D chiplet system the right part. The MoE- LLM parameters are modularized in each decoder layer and mapped to the individual chiplets.
3.3 Efficient All-to-All Communication
All-to-All communication is key bottleneck expert parallelism, as necessitates synchronization across all parallel units-a process often constrained by communication bandwidth. Reducing the data volume in such communication can effectively reduce end-to-end latency. this paper, we quantify the communication complexity using the average number of replications per token in the Disputch stage, denoted as Cr. We prove the Appendix that CT is the least upper bound for the ratio between the actual data volume during all-to-all communication and the number of tokens in training step. In standard expert parallelism frameworks [12], CT under top-k routing. However, if two co- activattedd experts for token are assigned to the same parallel unit (e.g., GPU in modern data centers or a chiplet in Mozart), only one replica would be required. therefore reducing Cr. By optimizing expert layout to increase the likelihood of such co-location, CT can be further minimized, thereby lowering the overhead of all-to-all communication.
Methodology
4.1 Overview of Mozart
We detail the design principles and methodology of Mozart with an overview in Figure 2, which addresses key bottlenecks of the post-training process of MoE-LLMs on chiplet systems through algorithm-bardware co-design.
On the algorithm side, we first profile the instruction tuning dataset using the pre-trained model, then apply strategic optimizations to improve post-training efficiency: Expert Clustering and Allocation: We cluster individual experts using the collaboration pattern prior, and map these clusters to chiplets using the workload distribution prior. aiming at balancing workload across MoE chiplet groups. More details are provided in Sec. 4.2. @ Fine-grained Scheduling: To overlap the DRAM communication overhead with on-chip computing, we propose streaming both the expert loading process and expert computing process of tokens using fine-grained scheduling, following the expert layout derived from the clustering and allocation algorithms. More details are provided in Sec. 4.3.
From the hardware side, propose 3.5D wafer-scale chiplet architecture featuring: 2.5D NoP-Tree Topology: We propose the 2.5D NoP-tree interconnect in Mozart that organizes attention chiplets central dispatchers and expert chiplets leaves. Switches enable in-network MoE aggre- gation, reducing communication latency and bandwidth cost. @ Hierarchical Memory Structure: Mozart introduces a two-level memory structure with model weights stored in distributed DRAM and activations cached in local SRAM. To further reduce data access latency, we adopt a logic-on-mant on-mmmory 3D integration, where each compute chiplet vertically stacks a compute die with an SRAM die via hybrid bonding. This tightly coupled design enables fast local access to intermediate results, such as activations, and aligns well with their temporal reuse patterns. More details are provided in Sec. 4.4.
4.2 Expert Collaboration for Efficient On-Package All-to-All Communication
Although every expert may be activated, the activation and co-activation patterns are not exactly balanced in practice. We take the pro- filing results on Alpaca [40] using DeepSeek- MoE [4] as example. At the final layer, some experts are sensibly activated more fre- quently (long horizontal bar in Figure 3), and
Expert Actustion Freg ency Layae Deep
Co-Activaioo Patten at Find Layer
Expert
3.50
45
3.00
Freeaency
Expart
Iigure 3: Left: Activation frequency for pre-
trained DeepSeek-MoE, indicating expert special- ization. Right: Co-activation pattern for pre-trained DeepScck-Mot. indicating expert collaboration.

---

some expert pairs are also activated more fre- quently (dark-colored blocks in Figure 3), mo- tivating us to specialize the expert layout on chiplets for balanced workload distribution during post-training. This clustered layout can also reduce the all-to-all communication volume, which is synchronous and cannot be overlapped with computation. To determine the expert placement on chiplets, we implement a 2-stage approach follows.
Stage-1: Expert Clustering. We cluster individual experts as candidates for expert-chiplet assign- ment, aiming at enhancing intra-cluster collaboration while minimizing inter-cluster collaboration. Intra-cluster collaboration defined as the average co-activation frequency among all expert pairs in single cluster, whereas inter-cluster collaboration represents the average co-activation frequency between all expert pairs across distinct clusters. Inspired by the farthest point sampling algorithm in point cloud learning [31], we implement the clustering shown in Algorithm 1.
Stage-2: Expert Cluster Allocation. Since our 3.5D chiplet architecture (Figure 2) allocates DRAM chip for group of MoE chiplets interconnected with switch, balanced workload distribution across these groups becomes critical. To achieve this, we formalize the cluster-chiplet assignment as binary integer programming problem. Let Ng denotes the number of groups (asserting No can be divided by N₂) and a binary matrix M {0, 1,,, N,,NN represents the cluster-group assignment, our optimization objective is formulated as:
Na
N.
minlMV- - Vanal, S.L.
MR.ii =1, VOS.S N. and Σ MR = AO No (5)
t=
j=0
where Vauw is an auxiliary vector with Ng elements, each one equals to 1 1/Ng.
Algorithm 1 Expert Clustering.
Require: Adjacent matrix C E No for graph G, number of chiplets N₂ (also the number of clusters).
Initialize expert clustering result C with N₂ empty lists. for N. do then
Find the 2 most highly co-activated experts, and push them into Lje)-
else
Find unselected expert with the lowest co-activation frequency with the experts in C.
Push it into Lje]-
end if
while len(C(e) < N./N, do Assert N₂ can be divided by No.
Find unselected expert with the highest average co-activation frequency with the experts in Lio)- Push it into Lja)-
end while end for return C.
4.3 Fine-grained Scheduling with Streaming Tokens and Experts
Of-paskage
Orspackage
112
2
1
i
Attention
Chipiet
MoE Ope.
Am Ope.
Onpadkage:
Highy- actissted
Load Param
Otspaskage
2
Chipal
Load
Compute.
On-package:
Lese-
Save
Chipler
Router:
Sync.
Altention Computing
Discatch
Expert Computing
Combine
Attentio Computing
Al-to-All
Layer
AlHto-All
Figure 4: Fine-Grained scheduling pipeline in the forward pass. The streaming tokens, marked with the execution order, can effectively overlap the computation (purple blocks) and DRAM com- munication (pink blocks, saving activations). We present 3 types of chiplets in the training pipeline, including attention chiplet, highly-activan chiplet, and less- asstvveiie chiplet. Since the 2 MoE chiplets share the same DRAM I/O, the highly activated experts should be first loaded to the chiplet for better communication-come compputtiiionnoooputtaiiion overlap.
5

---

The sheer size of parameters in MoE-LLMs necessitates storing them in DRAM and dynamically loading layers to chiplets for computation. However, this approach incurs significant communica- tion overhead compared to on-chip processing. To mitigate this bottleneck and enhance training parallelism, we propose a fine-grained scheduling scheme through streaming experts and tokens:
Streaming Experts Since multiple MoE chiplets within group share the same DRAM, their concurrent memory accesses require serialization. To optimize parallelism, we strategically schedule communication order by ranking expert clusters: using profiled workload distribution V, we quantify the importance of an expert cluster using the aggregated per-expert workloads, and prioritize the loading order of expert clusters with heavier computational workload first.
Streaming Tokens Partitioning the global token batch into streaming tokens ) enables overlapping DRAM communication (for saving activations during backward passes) with on-chip computation. To be specifie: (I) For the attention module, all tokens are partitioned into streaming attention tokens; (2) For the MoE module, the workload of each expert is partitioned into streaming expert tokens, and different experts on the same chiplet are computed sequentially.
Fine-Grained Scheduling The huge routed experts (Figure 1) results in significant communication overhead between DRAM and MoE chiplets, so we overlap it with on-chip computations using fine- grained scheduling in Figure 4, which mainly occurs in 2 aspects: (I) Loading highly-activater cluster & Attention computing; (2) Loading activaeti cluster & Highly-actinn activated cluster computing.
4.4 Wafer-Scale Chiplet Architecture
DRAM
Logic Tile
SRAM Tile
Switch
Hybrid Bonding
Silicon Intergoser
2.5D Direct Signaling
Package Substrate
TSV
Figure 5: The overall 3.5D chiplet architecture in Mozart. The hardware architecture implements three-layer hierarchical tree topology, comprising central attention node, switch nodes, and peripheral MoE nodes. The two-tier dies are connected face-to-face.
Motivation of Mozart Architecture Each MoE-LLM decoder block generally involves 2 stages: (1) the attention module and router network, and (2) expert computation with structural sparsity. This heterogencous and dynamic exccution flow poses significant challenges to system-leve scheduling and communication. To tackle them, we propose a wafer-scale 3.5D) chiplet architecture that integrates heterogeneous compute and memory resources: (1) To improve memory locality, we design hierarchical memory system aligned with the temporal reuse patterns of MoE-LLMs, enabling frequently reused data such as activations be cached closer to the computing unit using the 3D logic-on-memory stack, thereby reducing costly accesses to off-chip DRAM; (2) To reduce communication overhead across chiplets, we co-locate frequently co-activated experts onto the same chiplet based on profiling of activation patterns, significantly reducing costly inter-chiplet transfers during both forward and backward passes; (3) To enhance compute resource utilization, we disaggregate memory- bound and compute- compute-bound components onto specialized chiplets with matching bandwidth and compute capabilities, and apply fine-gra ane-grained, pipeline-awe -wwrre scheduling to balance load and overlap communication with computation.
Physical Layout To efficiently support the computation flow, we introduce a wafer-scale chiplet architecture that combines 3.5D integration with 2.5D -oP-tree interconnect, as shown in Figure 5, 0 3D Chiplet Stack: Each computing chiplet integrates logic die and an SRAM die in a vertical stack using hybrid bonding, supporting either attention or MoE-based FFN operations. The SRAM layer serves as a fast buffer for intermediate data with frequent reads and writes. Leveraging the short vertical interconnects enabled by 3D integration, this tightly coupled logic-on-memory chiplet offers significantly higher bandwidth and lower latency compared to conventional 2D designs, as illustrated in Figure 5. @ 2.5D NoP- Tree Topology: The inter-chiplet network adopts a 2.5D NoP-Tree

---

topology [19] that disaggregates attention and MoE operations across the network. Memory-bound attention chiplets are placed near the center of the tree as dispatching nodes with higher DRAM bandwidth, while compute-bound expert-cluster chiplets are organized as leaf nodes to execute MoF feed-forward computations with more computing resources. The system comprises 16 MoE chiplets, partitioned into 4 switch-comnected groups, each containing 4 expert-cluster chiplets with pre-defined placement. Moreover, the switch modules are equipped with in-network compute capabilities to aggregate MoE outputs locally, significantly reducing inter-chiplet communication and improving pipeline throughput. @ Memory Hierarchy: We design a two-level memory hierarchy for Mozart. Model weights are stored in DRAM distributed around the core the wafer. Every four expert clusters share a dedicated DRAM I/O interface. The DRAM connects to SRAM on the attention chiplet and switches for the MoE groups, enabling weight transformation from off-chip to the computing unit. Given that weights are relatively static during one iteration of training and exhibit low temporal access locality, they suited for off-chip DRAM storage. In contrast, activations highly transient and frequently accessed during the computation pipeline. Therefore, they are cached in the local SRAM die under each computing die minimize access latency and support rapid data exchange during the process.
Algorithm-to- -o-addaare Mapping To efficiently execute the MoE-LLM models on the proposed chiplet-based architecture, we map its major computational components, including attention and MoE-expert layers, onto specialized chiplets with coordinated dataflows and scheduling strategies. 0 Dataflow of the Training Process: During each training step, the system processes 32 samples (sequences), divided into 4 serially executed micro-batches of size 8. A weight-streaming strategy is adopted, where only one transformer block's 's weights are loaded at time, in response to area and interconnect constraints of the wafer-scale architecture. QKV V projection and multi-head attention score computation are mapped to multiple systolic arrays (SAs). SAs are grouped into tiles, each integrating local adder tree to aggregate partial sums and reduce intermediate communication. These partial results are transmitted from the compute die and stored in the underlying SRAM die via hybrid bonding. After attention completes, activations are routed through the NoP-Tree network to a switch module with for token-wiss wiiee routing and reduction. Tokens are dispatched to selected experts for FFN. Local aggregation is performed within each expert-cluster chiplets, followed by global aggregation via the switch. The aggregated expert outputs routed back to the attention module to continue processing the subsequent transformer blocks. During backpropagation, gradients follow the reverse path, with parameter updates performed locally on attention and expert chiplets before being written Pipeline: To improve compute
back to DRAM. @ Scheduling for Camputing-Commuunt throughput and resource utilization. the system adopts fine ane-grained pipeline scheduling. Leveraging the temporal locality in expert selection across adjacent training steps. frequently activated weights are prefetched onto expert chiplets ahead of token routing, reducing memory stalls. At runtime, each switch group coordinates micro- picro-batch-level pipelined execution. Based on the received activation load per token, chiplets within group sequentially fetch weights from DRAM via the shared switch. While one micro-batch undergoes FFN computation, the next micro-batch 's weights are concurrently loaded, enabling overlapped execution of compute and memory access. A similar strategy is applied during backpropagation to hide communication latency and sustain high training throughput.
5 Experiments
5.1
Algorithmic Setup
Table 1: Configurations of three pre-trained MoF-LLMMs used in our experiments.
# Total
# Activated Parameters
# Routed Experts
Shared Experts
Hidden Size
Model
# +Layers
Routing
Parameters
Qwen3-30B-Ae 333 [41]
30.5B
3.3B
128
0
2048
48
top
OLMoE-1B-78-025 [28]
6,92B
1.3B
64
0
2048
16
top-8
deepseek moe- 16b base [4]
16.4B
2.7B
64
2
2048
28
top-6
Our experiments include three MoE models with various architectures: Qwen3-30B-ARe [44, 45], OLMoE ---044 [28], and deepseek-moe- -6bbbaseb [4]. Details of them are summarized in Table 1. We use Alpaca [40]. an instruction tuning dataset of 52K samples, for all our experiments Our evaluation includes latency and energy as metrics to indicate the real-world impact of our designs. We use NVIDIA A100 80G GPU servers and Py PoTch for our profiling and simulation experiments.
7

---

Table 2: Hardware metries of the three MoE-LLMs used in experiments. The number of inter-chiplet links is computed based on the chiplet area for 2.5D signaling. The link counts are calculated as the product of horizontal and vertical link numbers for the 3D stack. The total area encompasses not only chiplets but also off-chip components such as DRAM.
Memory
2.5D
3D Hybrid Bonding /Link
Model
(DRAMStack SRAM/Tile)
Direct Signaling /Link
Area (mim) )
Power (kW)
Cap. (M8)
BW (GB/s)
BW (GBB))
Pitch (jm)
BW (GBB/)
Pitch (jmm)
Quen3-308-A36
14175
8192822.261
2568-32
0.125
50
0.125
OLMoB-18 11--1------00 7B-0924
10200
3.55
81928:2.255
2568:32
0.125
50
0.125
decpsonk--z0e- bib--bae
11230
3.19
819282.226
2568:32
0.125
50
0.125
50
5.2
Hardware Setup
The overall Mozart architecture comprises 16 expert-cluster chiplets for MoE computation, organized clusters, as well one dedicated attention chiplet. Each MoE/attention chiplet has 36-100 tiles, with 16 Systolic Arrays (SAs) in one tile and 256-576 Processing Elements (PEs) in one SA. Off-chip memory is provided by HBM2-based DRAM [29], with shared across expert-cluster groups (one per group) and 2 exclusively connected to the attention chiplet to provide high-bandwidth. We implement the logic dies, SRAM dies, inter-chiplet interconnects and switches in Verilog, and synthesize the gate-level netlist using Synopsys Design Compiler [38] targeting 28nm technology. The typical power consumption as reported by Synopsys PrimePower [39] based on the generated gate-level netlist. To evaluate the performance of Mozart, we further develop cycle- accurate simulator, whose runtime and power ouiputs are validated against the Verilog simulation results to ensure accuracy. For wolldworld implementation, we adjust hardware configurations for all three algorithmic baselines with FP16 precision to meet key 3.5D chiplet process constraints. We simulate all the design under 1GHz clock frequency. Detailed configurations for the three models are summarized in Table 2.
into 4
5.3 Experimental Results
Effectiveness of the Optimization Techniques Table summarizes four configurations of Mozart used to evaluate the effectiveness of our proposed algorithm-side methodologies, including one baseline without any optimizations and three variants that incrementally incorporate the optimization methods described in Sections 4.2 and 4.3. The simulation results demonstrate that our proposed 3 op- timization techniques can jointly reduce the end-to-er post-training latency, with 1.92x speedup for owen3-30B--a --B---3---------- A3B-Base, x.37x for DLMoE-1B-77 --2492 and 2.17x x for deepSeek-MoEL -Mk--66E-6B 1----66B-Bass We further provide Table to demonstrate the correlation between all-to-all communication data volume and the end-to-end training latency, where Mozart-A, B, and C present different data volumes during all-to-all communication, which is positively correlated to latency.
Study on the Impact of Sequence Length and DRAM bandwidth training latency increases as the sequence length per batch grows from 128 to 512. Although the number of batches decreases accordingly, each micro-batch carries longer sequences and heavier computation loads, which, when executed sequentially, become more constrained by communication bandwidth. This trend is particularly pronounced in the baseline design without any optimizations, where latency rises from 3.88s at length 128 to 7.64s at 512. In contrast, Mozart-C consistently achieves the lowest latency across all sequence lengths and exhibits reduced sensitivity to longer sequences, achieving a speedup of 2.34x at sequence length 5 512 and 1.47x at length 128 compared to the baseline. This improvement stems from its architecture that enables efficient communication- computation overlap and alleviates communication congestion through expert-aware layout and routing, which together mitigate the latency increase caused by longer and heavier micro-batches.
shown in Figure 6(b), the
Table 3: Configurations for different settings used in our experiments.
Method
Baseline
Mozart-B Mozart-A
Mozart-C
Optimization Technique
Specialized Expert Layout on Chiplets (Section 4.2)
X
Efficient All-to-All Communication (Section 4.2)
x
X
/
/
Communication Commutttioonoooaaaiiiiiioooooon Overlap (Section 4.3)
/
When comes to study of DRAM bandwidth depicted in Figure 6(C), all configurations achieve lower latency with HBM2 (256GB/s) [29] compared to SSD (15.8GB/s) [43] due to its higher memory
8

---

bandwidth. Notably, the relative speedup from Mozart optimizations becomes higher with HBM2 than SSD. This can be attributed to the domination of latency caused by DRAM-based expert weight streaming when using SSD, which remains the bottleneck even after optimization. Since MoR computation accounts for only a small portion of the overall training time, pipelining and token-level scheduling have limited impact when memory access is slow. Furthermore, the communication cost reduced by all-to-all optimization is only about one-third of the streaming latency, making the total gain under SSD more constrained. In contrast, with HBM2, faster streaming allows better utilization of compute-communicatin overlap, enabling the co-design techniques in Mozart to take full effect.
Baseline
Mozart-A
Basalino
Mozart-A
Baseline
Mozart-A
Normalized Latency
Mozart-8
Mozart-CC
Nornalized Latency
Nozart-6
Mozart-C
Normalized Latency
Mozart-8
Mozart-C
6.75
0.75
0.75
0.5
0.5
252
25
0.20
Owen3
OLMoE deepseek-moe
512
256
128
SSD
HBM2
Model
Sequence Length
DRAM
(b)
Figure 6: Experimental results on the chiplet system of Mozart. We report the average training latency per step for 1k iterations, with the micro batch size for streaming attention/expert tokens set to 8. (a) Study on the proposed optimization results (Sequence Length=256, DRAM-HBM2), with max latency of 4.87 s. (b) Study on the impact of the sequence length on Dwen3-30B-AS3 Model (DRAM=HBM2), with max latency of 7.65 s. (c) Study on the impact of the DRAM bandwidth on Qwen3-30B-A3B Model (Sequence Length=256), with a max latency of 9.17 s.
Table 4: The corelation between all-to-all communication complexity CT and end-to-end latency. C is calculated by averaging both the training iterations and the MoE layers for each setting.
Method
Qwen3-308-
A
OLMoB 1B- 7B-0929
DeepSeck NoB-16B----- -0Be-aae
Metric
Mozart-A
Mozart-B
Mozart-C
Yozart-A
Mozart-B
Mozart-C
Momart-A
Mozart-B
Mozart-C
Normalized Latency
0.73
0.59
0.52
0.63
0.48
0.422
0.67
0.46
Cr
8
6,58
5.77
8
6.84
5.63
6
5.56
4.32
5.4 Further Investigation
Q1: Is Mozart memory-bound or crapputing-boundd A: Memory-bound. This is because our proposed 3.5D) chiplet architecture in Mozart well-parallelize the MoE computation workload. While this design successfully eliminates the computational bottleneck associated with heavy MoE operations, the system's overall latency becomes constrained by the sequential MoE weight loading process. This fundamental limitation persists because weight loading throughput cannot be sub- stantially improved without hardware resource upgrades. Consequently, Mozart's 's performance is primarily governed by this unavoidable sequential bottleneck inherent to current hardware constraints.
coniniiaationCoopuaatitii Overlap > Efficient All-to-All Communication Specialized Expert Layout on Chiplets. The key insights are: O The communication overhead between DRAM and chiplets is the main bottleneck, and applying it on the baseline can offer 1.33x acceleration on Qwen3-305
Q2: Which algorithmic deigns are more critical in Mozart? A:
A3B-Base, 1.58x x on and 1.49x x on D 1-t--MoE-----B--- 16---e.B-BBse.... @ The all-to-all communication overhead is a secondary bottleneck for training latency, since it requires synchronization across all the chiplets and is constrained by the on-package bandwidth. Our specialized expert layout on chiplets can further reduce the data volume during all-to-all communication, as we illustrated in Table 4.
1B-OE--1--B--92-
Q3: Is Mozart compatible with existing efficient training algorithms? A: Yes, it is compatible with parfieter-efficient fine-tuning methods such as LoRA, QLoRA, etc. Mozart' architecture and scheduling mechanisms are designed to work orthogonally to these methods, as they primarily focus on different optimization goals. While PEFT methods reduce the total trainable parameters, Mozart optimizes the physical deployment of MoE workloads on chiplet architectures.
6 Conclusion and Limitations
We present Mozart, an algorithm- haadwware co-design framework for efficient post-training of MoE- L.MSs on chiplet systems. By jointly optimizing expert allocation, fine-grained scheduling, and hetero- geneous chiplet mapping on a 3.5D wafer-scale architecture, Mozart significantly improves communi- cation efficiency and hardware utilization, enabling scalable and efficient deployment of modularized workload. While Mozart demonstrates 1.92x performance improvement for Qwen3-300 -0ee -8---33----BBB-

---

2.37x for OLMoE-1B- 1B-7--99--- and 2.17x x for chitectures, two limitations remain. First, the attention modules are assigned to an individual chiplet, which may lead to suboptimal latency due to limited resources. This can be further tackled with data or tensor parallelism. Second, the switches can become performance bottlenecks under high commu- nication demand. While Mozart currently tries to reduce end-to-end latency through fine-grained scheduling, further improvements may potentially be achieved by allocating more chiplet area to switch resources and increasing bandwidth to achieve low-latency communication.
-ooe---MM--6B-B-- on 3.5D wafer-scale ar-
Acknowledgement
This research was partially funded by the National Institutes of Health (NIH) under award IR01EB037101-01. The views and conclusions contained this document are those of the au- thors and should not be interpreted of the NIH. Tianlong Chen was also partially supported by the Amazon Research Award. We thank Nikhil Kumar Cherukuri for the valuable discussions throughout this project.
representing the official policies, either expressed or implied,
References
[1] E. Bullmore and O. Sporns. Complex brain networks: graph theoretical analysis of structural and functional systems. Nature reviews neuroscience, 10(3):186-198, 2009.
[2] E. Bullmore and O. Sporns. The economy of brain network organization. Nature reviews neuroscience, 13(5):336-44) 2012.
[3] G. Buzsáki. Rhythms of the Brain. Oxford university press, 2006.
[4] D. Dai, C. Deng, C. Zhao, R. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y. Wu, al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-exxprts language models. arXiv preprini urXiv:2401.066 06066, 2024.
[5] T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. The Twelfth International Conference Learning Representations, 2024.
[6] Dao, D. Fu, Ermon, Rudra, and C. Ré. Flashattention: Fast and memory- exact attention with io-awareness. Advances in neural information processing systems, 35:16344 16359, 2022.
[7] D. Das Sharma, G. Pasdast, Z. Qian, and K. Aygun. Universal chiplet interconnect express (ucie): An open industry standard for innovations with chiplets package level. IEEE Transactions on Components, Packaging and Manufacturing Technology, 2022.
[8] Y. Dong, X. Liu, X. Hao, Y. Liang, R. Huang, Ye, and T. Jia. Hierarchical power co- optimization and management for Ilm chiplet designs. In Proceedings of the 43rd IEEE/ACM International Conference on Computer-Aided Design, ICCAD '24 2025.
[9] W. Duch. Brain-inspired conscious computing architecture. The Journal of mind and behavior, pages 1-21, 2005.
[10] D.-A. et al. Deepseck-v3 technical report, 2025.
[11] W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, (3(100): 2022.
[12] T. Gale, D. Narayanan, C. Young, and M. 7aharia. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems, 5:288-304, 2023.
[13] C. L. Gallen and M. D' D'Esposito. Brain modularity: a biomarker of intervention-relatee plasticity. Trends in cognitive sciences, 23(4):293-304, 2019.
[14] Goodfellow, Bengio, A. Courville, and Bengio. Deep learning, volume MIT press Cambridge, 2016.
10

---

[15] C. Hwang. W. Cui, Y. Xiong, Z. Yang, Z. Liu, H. Hu, Z. Wang, R. Salas, J. Jose, P. Ram, et al. Tutel: Adaptive -ixture-of-experts scale. Proceedings of Machine Learning and Systems, 5:269-287, 2023.
[16] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. I. Casas, E. B. Hanna, Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.0400 2024.
[17] N. Kashtan and U. Alon. Spontaneous evolution of modularity and network motifs. Proceedings of the National Academy of Sciences, 102(39):13733-133))
2005.
[18] M. Khairy, V. Nikiforov, D. Nellans, and T. G. Rogers. Locality-centric data and threadblock management for massive gpus. In 2020 53rd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), pages 1022-1036, 2020.
[19] G. Krishnan, S. K. Mandal, C. Chakrabarti, J.-S. Seo, U. Y. Ogras, and Y. Cao. Impact of on-chip interconnect on in-memory acceleration of deep neural networks. 18(2), 2022.
[20] H. T. Kung, B. McDanel, Q. Zhang. X. Dong, and C. C. Chen. Maestro: memory-on- logic architecture for coordinated parallel use of many systolic arrays. In 2019 IEEE 30th International Conference on ipplicattionspeciffe Systems, Architectures and Processors (ASAP). volume 2160-052X, pages 42-50, 2019.
[21] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. (GS Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021.
[22] P. Li, Z. Zhang, P. Yadav, Y.-L. Sung. Y. Cheng, M. Bansal, and T. Chen. Merge, then compress: Demystify efficient smoe with hints from its routing policy, 2024.
[23] S. Liu, R. M. Radway, X. Wang, J. Kwon, C. Trippel, P. Levis, S. Mitra, and H.-S. P. Wong. Future of memory: Massive, diverse, tightly integrated with compute from device to software. In 2024 IEEE International Electron Devices Meeting (IEDM), pages 1-4, 2024.
[24] Luo, P. J. Peng, Y. Zhao, Y. Cao. Y. Cheng, and T. Chen. Occult: Optimizing collaborative communications across experts for accelerated parallel moe training and inference. In Forty- second International Conference on Machine Learning, 2025.
[25] S. Luo, J. Peng, P. Li, H. Wang, and T. Chen. Hexa-moe: Efficient and heterogeneous-aweee training for mixture-of of-experts, 2025.
[26] R. Mahajan, R. Sankman, N. Patel, D.-W. Kim, K. Aygun, Z. Qian, Y. Mekonnen, Salama, S. Sharan, D. Iyengar, and D. Mallik. Embedded multi-die interconnect bridge (emib) high density, high bandwidth packaging interconnect. 2016 IEEE 66th Electronic Components and Technology Conference (ECTC), pages 557-565, 2016.
[27] A. Mehonic and A. J. Kenyon. Brain-inspire computing needs master plan. Nature, 604(7905):255-266) 260, 2022.
[28] N. Muennighoff, L. Soldaini, D. Groeneveld, K. Lo, J. Morrison, S. Min, W. Shi, P. Walsh, O. Tafjord, N. Lambert, Y. Gu, S. Arora, A. Bhagia, D. Schwenk, D. Wadden, A. Wettig, B. Hui, T. Dettmers, D. Kiela, A. Farhadi, N. A. Smith, P. W. Koh, A. Singh, and H. Hajishirzi. Olmoe: Open mixture-of-expprts language models, 2024.
[29] S. Naffziger, N. Beck, T. Burd, K. Lepak, G. H. Loh, M. Subramony, and S. White. Pioncering chiplet technology and design for the amd epycTM and ryzen TM processor families Industrial product. 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), pages 57-70, 2021.
[30] T. OLMo, P. Walsh, L. Soldaini, D. Groeneveld, K. Lo, S. Arora, A. Bhagia, Y. Gu, S. Huang, M. Jordan, et al. 2 olmo 2 furious. arXiv preprint arXiv:2501.000 20066006666.... 2024.
[31] C. R. Qi. L. Yi, Su, and I.. Guibas. Pointnet++:: Deep hierarchical feature learning on point sets in metric space. Advances in neural information processing systems, 30, 2017.
11

---

[32] S. Rashidi, W. Won, S. Srinivasan, P. Gupta, and T. Krishna. Fred: Flexible reduction- distribution interconnect and communication implementation for wafer-scale distributed training of dnn models, 2024.
[33] J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances Neural Information Processing Systems, (1166658-668888 2024.
[34] Y. S. Shao, J. Clemons, R. Venkatesan, B. Zimmer, M. Fojtik, N. Jiang, B. Keller, A. Klinefelter, N. Pinckney, Raina, S. G. Tell, Y. Zhang. W. J. Dally, J. Emer, C. Gray, Khailany. and S. W. Keckler. Simba: Scaling deep-learning inference with multi-chip-mouule- casp--modullbaaee based architecture. In Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture, MICRO '52, page 14-27, New York, NY, USA, 2019. Association for Computing Machinery.
[35] N. Shazeer, Mirhoseini, K. Maziarz, Davis, Q. Le, G. Hinton, and J. Dean. Outra- geously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 066388 2017.
[36] D. Stow, Y. Xie, T. Siddiqua, and G. H. Loh. Cost-effective design of scalable high-performance systems using active and passive interposers. In 2017 IEEE/ACM International Conference on Computer-Aided Design (ICCAD), pages 728-735, 2017.
[37] D. Stow, Y. Xie, T. Siddiqua, and G. H. Loh. Cost-effective design of scalable high-performance systems using active and passive interposers. In 2017 IEEF/ACM International Conference on Computer-Aided Design (ICCAD), pages 728-735, 2017.
(38] Synopsys. Design compiler: Concurrent timing, area, power, and test optimization. attps://wwe https:://ww.eynnopyss comccoo//mmm// www./mpelmmentattttoooooooooeeee limmmmmmettationnaaaaee andtattioo--nn----eeee http://awww.aiooteettttsttteeeee tests dc-ultra. html. accessed 2024.
[39] Synopsys. Primepower: Rtl to signoff power analysis. https://www. http://wwwwwwoss........ com/ implementation and ottp://nnnaaaaiooaaaaiooo/o/a/// 22.
.ttll, 2024. Accessed: 2024-11-
[40] R. Taori, Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto. Stanford alpaca: An Instruction-folloiii llama model. h coomtatsu-lale stanford_ale _alaaca, 2023.
[41] Q. Team. Qwen3: Think deeper, act faster, April 2025.
[42] T. Wei, B. Zhu, L. Zhao, C. Cheng, B. Li, W. Lü, P. Cheng. J. Zhang, X. Zhang, L. Zeng, X. Wang, Y. Ma, R. Hu, Yan, H. Fang, and Zhou. Skywork-moe: A deep dive into training techniques for mixture-of-experts language models, 2024.
[43] G. Yadgar, M. Gabel, S. Jaffer, and B. Schroeder. Ssd-based workload characteristics and their performance implications. ACM Trans. Storuge, 17(1), Jan. 2021.
[44] A. Yang. B. Yang, B. Hui, Zheng, B. Yu, C. Zhou. C. Li. Li, D. Liu, Huang, G. Dong. H. Wei, H. Lin, J. Tang, J. Wang, J. Yang, J. Tu, J. Zhang, J. Ma, J. Xu, Zhou, J. Bai, J. He, Lu, K. Chen, K. Yang, M. Li, M. Xue, N. Ni, Zhang, Wang, R. Peng, R. Men, R. Gao, R. Lin, S. Wang, S. Bai, S. Tan, T. Zhu, T. Li, T. Liu, W. Ge. X. Deng. X. Zhou,
J. Lin, K. Dang,
X. Ren, X. Zhang, X. Wei, Ren, Y. Fan, Y. Yao, Y. Zhang, Y. Wan, Y. Chu, Y. Liu, Z. Cui,
Z. Zhang, and Z. Fan. Qwen2 technical report. arXiv preprint arXiv:2407.1066
2024.
[45] A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang. H. Wei, H. Lin, Yang, Tu, J. Zhang, J. Yang, Yang. J. Zhou, Lin, K. Dang, K. Lu, K. Bao, K. Yang, L. Yu, M. Li, M. Xue, P. Zhang, Q Zhu, R. Men, R. Lin, T. Li, T. Xia, X. Ren, X. Ren, Y. Fan, Su, Y. Zhang, Y. Wan, Y. Liu, 7. Cui, Z. Zhang, and Qiu. Qwen2.5 technical report, arXiv preprint 2111v22412.15115. 2024.
12

---

[46] Z. Yu, S. Liang, T. Ma, Y. Cai, Z. Nan, D. Huang, X. Song, Y. Hao, J. Zhang, T. Zhi, Y. Zhao, Z. Du, X. Hu, Q. Guo, and Chen. Cambricon-IIm: A chiplet-based hybrid architecture for on-device inference of 70b Ilm. In 2024 57th IEEF/ACM International Symposium on Microachitectuure (MICRO), pages 1488, 2024.
[47] M. Zhang. P. Li, J. Peng, M. Qiu, and T. Chen. Advancing moe efficiency: A collaboration constrained routing (c2r) strategy for better expert parallelism design, 2025.
13

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA]-
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (After eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because would be too computationally expensive" " or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
the checklist one of the factors in their evaluation.
"[No] " or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification elaborate, All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: Claim 1: Mozart is an aggorithm-hardwee co-design approach aiming at
efficient post-training of MoE-LLM wafer-scale chiplet systems. We show this in Sec- tion 4.1, and demonstrate the correlation between algorithm designs and chiplet architecture in Figure 2. Claim 2: Mozart specializes the expert layout chiplets simultancously balance the workload for MoE chiplets and minimize the data volume in all-to-all communi- cation. We show this in Section 4.2. We provide the motivations in Figure 3 and formulate the procedures using Algorithm and Equation 5. Claim 3: Mozart tackles the heavy communication overhead between DRAM and MoE chiplets with fine-grained scheduling. We illustrate this in Figure 4. Please see Section 4.3. Claim 4: Mozart designs novel 3.5D chiplet architecture and the soorrthh--t-hhrdwaar mapping strategy. We show the chiplet architecture in Figure 5. Please see Section 4.4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
14

---

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations are mentioned briefly throughout the paper and discussed in detail in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (And correct) proof?
Answer: [Yes]
Justification: Please see Section 3.3 and Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
15

---

4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all code and data necessary to reproduce every experimental result that we describe in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes].
Justification:
We provide all experimental code and data. Please see the Appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines hattps://nips.com cc. Baaali//guides/CoeeSubbbaa www.ammmissooooooioi...o.... for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new soun-source benchmark).
16

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //nips. //wpss.ccppublll../ wwtpcccwppppppppe.e.eeeeoc wtwiic//ggid///eeeee /wt////oobbbisssseeeeeeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (Appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide sufficient details to make sense of the results in the core paper. Full details are provided in the available code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA].
Justification: We follow previous works and conduct simulation experiments and ablation studies, without including error bars or similar information.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
17

---

8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: We employ NVIDIA A100 80G servers for profiling, with results shown in Figure 3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: https://neut //tpprss///uii....// httpppplc//tcisssiiiiiieeapaaaee
Answer: [Yes]
Justification:
We follow the guidelines of the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discuss broader impacts of Mozart the Appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
18

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA].
Justification:
Our work does not involve the release of new models or data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The LLMs, datasets, and codebase used in our work comply with open-source licenses and can be used for scientific research.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi wwwoee. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA].
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
19

---

Researchers should communicate the details of the dataset/codeemmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA].
Justification: This paper does not involve data annotation crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA].
Justification: We do not employ LLM to play part in the core methodology, scientific rigorousness, or originality of the research in this paper.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
20

---

Please refer to our LLM policy (nttps:/heuripss.c/// htt.r//s...///ofeeeerec////////1
for what should or should be described.
21

---

A Code
Please find the code base for this paper here: https//gtthu..com/UNNTESSSLLL
B More Experimental Results
Sequence Length=128
Normalized Latency
Baseline
0.75
Nozart-A
Mozart-B
0.25
Mozart-C
SSD
HBM2
SSD
HBM2
SSD
HBM2
Qwen3
OLMoE
deepseek-moe
Figure 7: Normalized Latency Comparison for 3 MoE-LLMs with Sequence length 128. The
max wall-clock latency here 7.61 (Qwen3 model with baseline method using SSD for DRAM),
Sequence Length=256
Normalized Latency
Baseline
0.75
Nozart-A
0.5
Mozart-B
0.25
Mozart-C
SSD
HBM2
SSD
HBM2
SSD
HBM2
Qwen3
OLMoE
deepseek-moe
Figure 8: Normalized Latency Comparison for 3 MoE-LLMs with Sequence length 256. The
max wall-clock latency here 9.17 (Qwen3 model with baseline method using SSD for DRAM).
Sequence Length=512
Normalized Latency
Baseline
0.75
Mozart-A
0.5
Mozart-B
0.25
Nozart-C
SSD
HBM2
SSD
HBM2
SSD
HBM2
Qwen3
OLMoE
deepseek-moe
Figure 9: Normalized Latency Comparison for MoE-LLMs with Sequence length 512, The max wall-clock latency here is 13.03 s (Qwen3 model with bascline method using SSD for DRAM).
We provide comprehensive numerical latency results for all configurations, including sequence length (128, 256, 512), 4 methods (Mozart Baseline, A, B, and C), and DRAM (SSD and HBM2). Results comparison visualizations are provided in Figure 7, 8, and 9.
C Motivation Explanations
C.1
Why Attention Memory-Bound and FFN Compute-Boun
The chiplet architecture in Mozart utilized the fact that in typical decoder layer in modern large language models, the Attention module memory- and the FFN module is computation-bound We demonstrate it using profiling experiments the OLMo-2 model series [30]. The experiment settings are:
We examine single decoder layer, and collect the wall-clock latency and the FLOPSs for both attention and FFN modules.
22

---

The results are collected through running the forward pass, i.e., the prefilling stage of model inference, and the results are normalized for easier comparison.
We fix the batch size to and test the sequence length of 512, 1024, and 2048.
We select OLMo-2 models with scales, including IB, 7B, 13B, and 32B.
The profiling experiments are visualized in Figure 10 (IB), Figure 11 (7B), Figure 10 (13B), and Figure 13 (32B). We can find that, the FFN module counts for more FLOPs but less wall-clock latency. It is because the Attention module is memory-bound and the FFN module is computation-
The FFN module counts for more FL.OPs because it contains the computation task for it is mainly composed of large matrix multiplication, which is easy parallelize. Therefore, the wall-clock latency of can be lower than attention.
model parameters. But
The Attention module requires frequent memory access operations, which is demonstrated by the Wlash-Attention series [6, 5, 33]. Although it contains fewer model parameters, the computation tasks here are difficult to parallelize. Therefore, the attention module counts for more wall-clock latency.
OLMo-2-0425-1
OLMo-2
Length
OLMo-2 0425
2048)
No nommated Latency
Normalzzed FLOPs
Nommstoee Laneney
to
Nomes FLOPs
4.
So
Norral FLOPs
to
20
2.0
20
10
12
Abantton FFN
Adamion FFN
Adansoon FFN
Adamson FRN
Attenton PFN
Abantion FFN
Figure 10: Profiling results on latency & FLOPs for Attention & FFN using OLMo-2-0425-11 1B.
(e) 2: 2-11247B Sequence Langh- 6612
(8) OLMe 2 1-124-78 Sequence Longth 1024)
[8] CLMo 1-14-7B Sequence Longth 2048)
Nomatte Lateroy
Normalzodd FLOPs
Nomma aad PLOPPS
Nomalozd FLOPs
3.50
3.
10
Normstized
1.0
10
20
10
FFN
Adaomion
Abontion
Figure 11: Profiling results on latency FLOPs for Attention
FFN using OLMo-2-1124-7B
20Mo-2-11224 19B Sequence
(b) OLMo 1-124
(c) O 1-114
Nomeazer Latrey
5g
Normalieed FLOPs
Normarzad FLDrs
4.0
Nomalued Latency
Normatzzdd FLOPs
2.
152
Normuizzer
Ro
20
10
1.0
Altotion FFN
Ademmtin FFN
Adantion FPN
Adomtonn FRN
Attenton PFN
Abtention FFN
Figure 12: Profiling results on latency & FLOPs for Attention & FFN using OLMo-2-1124-13B
C.2
Challenges for Mixture-of-Expert Computation
We present 3 challenges for MoE computation in the abstract part of this paper, including memory locality issues, communication overhead, and insufficient computing resource utilization. Our algorithm-hardware co-design scheme in Mozart tries to solve these challenges with joint efforts. demonstrate these challenges through fine-tuning an OLMoE- 1B-7B model with 4-way expert parallelism, with batch size 8 on each GPU and sequence length 512. We use MegaBlocks [12], standard expert parallelism framework, for the MoE modules, and use data parallelism for the attention modules. We employ the dropless MoE implementation. The training speed is 2-3 iterations per second, and we monitor the behavior each GPU with an interval of 0.1 s. We take 3 fragments for visualization, as shown in Figure 14, 15, and 16, which demonstrate that both the GPU power and the memory consumption show high dynamism. These phenomena can explain 2 challenges:
Memory Locality Issues: Since the workload for each expert changes dynamically, the activation tensors should be frequently allocated and freed, leading to severe memory management issues.
23

---

OLM---0022--32 Langth
(6)
OLM6-203355-288 (Squuence Length
6:Nk
Normalted Laauncy
Nomelted FLLDTs
L2ee PLOPP
3.0
No PLCOP
1.0
5.0
1.0
4.0
1.0
3.18
3.11
Normaltend
3:12
Nerratzz
2.12
%
10
1.
10
1.0
Attention FFN
Adomtion FFN
Atamnoon FPN
Attontion
FFN
Adantion FFN
Attedoon FFN
Figure 13: Profiling results on latency & FLOPs Attention FFN using OLMo-2-0325-322.
GPJ Power Consumption Over Time
Time
GPU Time
6000
8000
Figure 14: GPU Behavior Monitor at Time Step 6k-7k.
Insufficient Computing Resource Utilization: The reason for this challenge lies in 2 aspects: (1) the dynamism of workload leads to dynamism of GPU power, and (2) the dynamism of workload restricts the training batch size to avoid out-of-memory error, which also constrains the utilization of GPU computing resources.
The all-to-all communication issues have been explained in Tutel [15], which is significant bottle- neck for training MoE models at scale, consuming up to 40% of the total runtime.
D Meassuring All-to-All Communication Complexity with Cy
We propose to measure the all-to-all communication data volume in Section 3.3 using the average replication times of each token, denoted as C7. We prove that CT is the least upper bound of the ratio between actual all-to-all communication data volume and the total number of tokens. We take single all-to-all communication in D-way way expert parallelism as example, and denote the original tokens as {S:) For single token E S; on device i, we denote the number of replications for transmitting from device i device j as N/ (L), i.e., token i on device activates N/ (1) experts preserved on device j. In the standard expert parallel framework, given top-k routing, we have
P-1 E x((1) - V, and VOSiS
(6)
For the actual all-to-all communication data volume:
D-1 MΣCN +
D-1
D-1
Σ N((())
ΣΣ
N/ N((1 + Σ NE(1))
1=0 rES, 1=
5=
1=0 rES, ==
#=
D-1
D-1
Σ
(((
(7)
1=0 rES;
1=0
D-1
S
The 2 inequalities in Equation 7 are reached when

---

GPJ Power Co Consumption on Over Time
6000
5300
DACO
0500
2000
Hurorry Over Time
S77133
Time
Figure 15: GPU Behavior Monitor at Time Step 8k-9k.
GPU Power Consumprion Over Time
Timm
1O000
X,2000
3odce
304ce
20002
21900
Figure 16: GPU Behavior Monitor
Time Step 10k-1 1k.
The first one is achieved when (t) 0 for all 0 < i S D - and t E S;, i.e., no token would activate the experts kept on the device where the token is originally kept.
The second one is achieved for standard expert parallelism, iee, making k replications for each token in the dispatch stage under top-k routing.
The first inequality cannot be utilized for communication efficiency, since it is data-dependent and lask-dependent. While the second inequality can be leveraged by employing our proposed strategy in Section 3.3.
E Impact Statement
As the paper's primary innovation is efficiently deploying the post-training process of MoE-based large language models on the chiplet-based system, it by itself doesn't It pose any obvious risks. The potential for negative societal impact depends on the specific MoE-LLMs. We strongly recommend these models be used in compliance with all ethical standards appropriate to the domain in which it is targeted to deployed.
25


Reviews:
[{"review_id": "nLdumd9zrz", "reviewer": "Reviewer_Bfvg", "summary": "This paper introduces Mozart, a software-hardware co-design framework for the efficient post-training of MOE LLMs. This work is motivated by the hardware challenges posed by the sparse and modular nature of MoE models especially memory locality and communication overhead. The proposed strategy tries to addresses expert routing and communication overheads in MoE-style LLMs, On the algorithmic side, Mozart implements an expert allocation strategy that clusters experts based on their co-activation patterns to minimize inter-chiplet communication and balance workloads. It also implements a fine-grained scheduling mechanism to improve the overlap between communication and computation. On the hardware side, Mozart designs a specialized 3.5D wafer-scale chiplet architecture with a hierarchical memory system and a 2.5D NoP-Tree interconnect topology. Using cycle-accurate simulations, the paper shows that Mozart achieves speedups (over 1.9x) across three popular MoE models compared to baseline designs.\n\nDespite its specialization to MoE-style LLMs and modular hardware assumptions, the paper presents valuable design principles for future accelerators. Its latency-aware expert scheduling, communication modeling, and hardware-algorithm co-design offer practical insights for chiplet-based and sparse-aware systems. These contributions are particularly timely as AI hardware trends shift toward modular packaging, inter-die communication optimization, and efficient support for dynamic, sparse workloads. While generalization beyond MoEs is limited, the paper nonetheless informs a growing class of AI accelerators targeting scalable, fine-grained model execution.", "questions": "1.\tThe expert clustering approach (Algorithm 1) is based on a greedy, farthest-point-sampling method. Have you investigated or considered alternative graph clustering algorithms (e.g., spectral clustering, METIS) on the expert co-activation graph especially that the used algorithm is computationally intensive? It would be good to know if other methods could yield a better workload balance or further reduce communication volume.\n2.\tCan you elaborate on how the system's performance and communication patterns would be affected if you were to apply data or tensor parallelism across multiple chiplets for the attention computation? How would this interact with the NoP-Tree topology? Would it be compatible with existing TP approaches?\n3.\tThe expert allocation strategy is based on a static analysis of expert co-activation patterns from an instruction tuning dataset. How stable are these patterns during a long pre-training run? Would a dynamic approach, where expert clusters are periodically re-evaluated and re-allocated, be necessary or beneficial?\n4.\tAll-to-all communication is identified as a key bottleneck. Your approach reduces the data volume by co-locating experts. How does this compare to other techniques for optimizing all-to-all, such as using hierarchical communication algorithms or topology-aware communication libraries?\n5.\tYour proposed hardware–software co-design is tightly coupled to the modular and sparse nature of MoE architectures. I'm not expecting a complete solution here, but could you share how you might begin to extend or adapt your approach to support a broader range of model types—such as dense LLMs, hybrid MoE + dense architectures, or even non-transformer models like state-space models (e.g., Mamba)? In particular, how might your expert clustering, routing strategies, or inter-chiplet communication optimizations evolve (or be abstracted) to support models with different compute and dataflow characteristics?", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and organized. \n- The authors address a timely topic.  MoE architectures are increasingly being adopted to scale LLMs to trillion-parameters. Hence addressing their communication overhead and resource utilization bottleneck is an important and practical research direction.\n- The concept of an algorithm-hardware co-design frameworks specifically designed for MoE models on emerging 3.5D wafer-scale architectures is a strong contribution. \n- The work is well motivated by a clear analysis of the system-level challenges. \n- The of a cycle-accurate simulator that has been validated against Verilog simulations is also a strong point. \n- The emprirical evaluation is generally well explained and done.\n\nWeaknesses/Areas of improvements:\n- The hardware evaluation is entirely simulation-based. While this is a critical first step in any new architecture research, the performance evaluations and gains depend heavily on the realization of the 3.5D wafer-scale system. I recommend that the authors include a detailed discussion about the assumptions made in the simulator and if there are any challenges to translate their design and simulation into a real physical hardware/chip.  \n- This statement is partially correct: “Prior works tend to neglect wafer-scale integration [7, 19, 42] and largely adopt coarse-grained, static workload partitioning strategies [29].”. Wafer-scale integration (e.g., Cerebras WSE) is rare and only done by a few companies due to extreme cost, yield, and packaging challenges. So although most existing work do not consider wafer-scal, this is rather a of practicality than oversight. Coarse-grained, static partitioning is common in chiplet designs: models are statically sharded across chiplets based on layers or tensor slices. That said, some recent works that are still evolving (e.g., FlexGen, Megatron-DeepSpeed) do explore dynamic routing or load-balancing. ”\n- The statement in the paper: “Prior works lack system-level coordination and optimization for MoE” uses reference 22, Future of Memory (Liu et al., IEDM 2024). This reference is not a MoE-LLM workload partitioning paper strategy. This work advocates for tighter integration of diverse memory systems but does not propose specific coordination or workload tiling strategies for MoE-LLMs or chiplet systems. So, this citation does not directly support the claim about workload tiling or coordination strategies for MoE.\n- Some key relevant works are missing from the related work discussion:\n   - GShard (Lepikhin et al., ICLR 2021): A foundational MoE system that introduced routing + automatic sharding across devices. Crucial for historical context. This is the first large-scale MoE at Google.\n   - BASE Layers (Lewkowycz et al., NeurIPS 2022): A hybrid sparse-dense expert design that maintains performance while reducing memory. This should be compared and contrasted Mozart’s memory-aware scheduling.\n    - MoEfication (Zhou et al., ICML 2022): Shows conversion of dense models to MoE post-training. This is related to the post-training MoE optimization approaches as proposed by Mozart.\n   - TaskMoE / Expert Tuning (Zhang et al., 2022): Specialization for multi-task models. \n   - Tutel & DeepSpeed-MoE need to be more discussed, especially in comparison to Mozart’s routing, scheduling, and clustering.\n- Suggestions for improving the related work section:\n   - Cite GShard and discuss how Mozart differs in routing granularity, post-training design, and hardware assumptions.\n   - Include Cerebras WSE to contrast wafer-scale monolithic vs. modular chiplet-based wafer-scale.\n   - Add a table or paragraph contrasting Mozart with FRED, Cambricon-LLM, Maestro, etc., clarifying MoE-awareness and system-level scheduling gaps.\n   - Consider adding a  short discussion of commercial accelerators (AMD MI300, Intel PV) for practical relevance.\n   - Cite MoEfication to better support the post-training and fine-grained modularity motivation.\n- The clustering method used in Equation 5 is based on the k-clustering with pairwise distances.This algorithm is computationally expensive (O(n²)), and no mention is made of approximation methods (e.g., hierarchical clustering, METIS). \n- The mathematical formulation in Equation 6 does not provide constraints on the number of experts per token (e.g., top-2) of the token capacity per expert (e.g., max tokens per batch). This needs to be clarified by the authors. \n- In Equation 7, runtime score for scheduling, how is the load defined? Is it measured in flops, tokens, latency? What is the utilization measured against? (e.g., theoretical peak). The formulation remains abstract and hard to reproduce without further clarification by the authors. \n- The proposed co-design strategy effectively addresses expert routing and communication overheads in MoE-style LLMs, however, it is highly specialized to this sparse model family. The system assumes hard top-k expert gating and token-level dispatch, which do not generalize to dense autoregressive models (e.g., GPT-3, LLaMA), diffusion models, or state-space architectures (e.g., Mamba) that lack modular sparse execution patterns. Additionally, the hardware assumptions—modular chiplet-based designs with a Network-on-Package interconnect—limit applicability to a narrow class of accelerators. Real-world systems such as monolithic GPUs, wafer-scale engines (e.g., Cerebras), and mobile AI chips follow vastly different architectural and memory paradigms. The lack of runtime adaptability in expert scheduling and coarse treatment of memory hierarchies further constrain its deployability. A discussion on these generalization limitations is absent.\n- Regarding the empirical evaluation. There is a lack of baselines beyong static scheduling. The main baseline is a static MoE routing strategy or uniform tiling. It would be good to add comparisons against DeepSpeed-MoE or Tutel routing strategies and Existing expert placement heuristics (e.g., token-balanced sharding, load-aware dispatch). The existing evaluation is not enough to claim that Mozart significantly outperforms state-of-the-art MoE runtimes.\n- Limited System Scaling Study. The paper does not analyze how Mozart behaves under increasing number of chiplets, varying interconnect topologies (e.g., mesh vs. ring) and realistic memory capacity constraints. Such scaling or sensitivity analysis can showcase if the benefits hold at larger or smaller deployment scales.\n- The paper’s results report mean latency or throughput, but does not show variance or standard deviation, confidence intervals, and sensitivity to input token distribution (e.g., random vs. clustered token-to-expert assignment). \n- While latency and utilization are discussed, the paper completely omits energy or area modeling. Given the focus on chiplet hardware, this is a notable omission—MoE communication patterns can be power-intensive.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer Bfvg for the insightful and detailed comments. To address your questions, we provide pointwise responses below.\n\n**[Weakness 1. Assumptions of the simulator]**\n\nMozart is simulated on a tree-structured chiplet system with 27 heterogeneous chiplets: 1 for attn, 16 for FFN experts (4 clusters x 4 chiplets), 4 for switches, and 6 for DRAMs, organized in a balanced quadtree of depth 2.\n\nEach compute chiplet employs 3D hybrid bonding, stacking high-density distributed SRAM beneath the compute logic to provide significantly higher bandwidth and capacity than conventional 2D integration. Inspired by wafer-scale systems like Cerebras, Mozart offers sufficient on-chip SRAM to hold expert parameters and pre-activated activations, eliminating intra-chiplet caching bottlenecks. For area efficiency, expert chiplets adopt a weight-streaming approach: only the parameters of the activated experts are prefetched from the local cluster DRAM in each training cycle. The top logic layer integrates weight-stationary PEs for GEMM operations and standard non-linear functions.\n\nEach inter-chiplet link provides 256 GB/s bandwidth and incurs a 1 ns per-hop latency. Switch chiplets support in-network computing, enabling efficient collective operations (e.g., token scatter/gather and gradient aggregation) required for fwd and bwd passes. Our simulation models both computation and inter-chiplet NoC routing at cycle level to evaluate token latency under warm-up and steady-state conditions.\n\n**[Weakness 2. Partially correct statement]**\n\nWe will refine it following your constructive comments.\n\n**[Weakness 3. Incorrect citation on Future of Memory]**\n\nWe apologize for the omission, and will replace [22] with Cambricon-LLM (Yu et al, ISCA 2024) and Hecaton (Huang et al, arxiv 2024).\n\n**[Weakness 4 & 5. Related works and discussions]**\n\nWe will add a subsection after *Modularized LLMs* to incorporate these fundamental works on GPU-deployment of MoE-LLM.\n\n**[Weakness 6. Computation complexity of expert clustering]**\n\nThe clustering problem is solved offline on the profiled data. Since all the latest MoE-LLMs (even the largest Kimi K2) use only a handful of routed experts (i.e, 384) per layer, the expert assignment problem is compact enough to be solved exactly with commercial solvers such as Gurobi.\n\n**[Weakness 7. TopK]**\n\nWe follow the dropless MoE design in MegaBlocks without the notion of expert capacity. We will claim it in our revised version.\n\n**[Weakness 8. Runtime score]**\n\nWe use Equation (7) to demonstrate the effectiveness of measuring all-to-all communication volume with our proposed $C_{\\mathcal{T}}$. The workload here is the tokens before MoE computing. It is not a runtime scheduling scheme, which lies in Section 4.3.\n\n**[Weakness 9 & Question 5. Generalization]**\n\nMozart targets deploying MoE-LLMs on chiplet systems, as sparse scaling has been widely adopted in the latest LLMs (e.g, DeepSeek, Qwen3, and Kimi K2), but studies on specialized hardware are very few.\n\nFrom the hardware side, modularization enables flexible composition of heterogeneous components on chiplet systems for diverse applications, which is a key trend in AI accelerator design. This modular design is also validated in practice by a 3D SoC partition and composition Compiler platform from Synopsys, which has been adopted by industry to accelerate AI system development. Mozart adopts a modular design with 2 specialized chiplets for memory-bound attn and compute-bound FFN, respectively, enabling flexible workload-to-chiplet mapping by decoupling operators based on their characteristics.\n\nThough tailored for MoE, Mozart is composable and generalizable. Dense models are also FFN-dominant, and overlapping DRAM communication is also critical. FFN can be decomposed into parallel subtasks mapped to multiple chiplets. We can remove the routing designs, evenly allocate FFN shards on previous MoE chiplets, and replace all-to-all communication with scatter & gather. Although the NoP-Tree in Mozart is physically implemented at the chiplet level, its concept can serve as a logical topology in GPU or wafer-scale systems. For example, GPUs can emulate tree-based expert routing in software using NCCL, while Cerebras WSE, as a wafer-scale monolithic, allows logical NoP-Tree topology via programmable routing. This aligns with FRED (ISCA 2025), which advocates tree-like interconnects over rigid 2D meshes to better support flexible, non-blocking collective communication.\n\nOverall, although Mozart is optimized for MoE deployment, its core principles, modular compute decomposition, communication-aware scheduling, and logical topology design, are highly extensible.\n\n**[Weakness 10. Baselines]**\n\nWe omitted to claim that we follow dropless MoE in MegaBlocks to enable seamless adaptation to pre-trained MoE-LLMs. Different routing strategies or expert layout heuristics for load balancing can be easily integrated into Mozart.\n\n**[Weakness 11. System scaling]**\n\nMozart is an empirical optimal solution for MoE-LLMs because:\n- Mozart supports moderate scaling by expanding expert chiplet clusters and switch nodes. However, as the number of MoE chiplets grows excessively, token replication and attn-to-expert communication workloads also grow, leading to higher all-to-all burdens. Furthermore, assigning a certain number of additional chiplets to the attn module can reduce its computation latency; however, this introduces extra routing complexity within the attention group. Therefore, a trade-off must be made between computing latency and routing overhead to maintain system efficiency.\n- We apply in-network computing switches to establish inter-chiplet connections. The tree topology with such switches is inherently well suited for essential reduce/scatter operations in MoE workloads, while sacrificing extra area as the switch chiplet is nearly as large as a computing die.\n- We apply mesh to establish intra-chiplet connections, since it is feasible for regular data movement in GEMM.\n\n**[Weakness 12. Comprehensive results]**\n\nFor more numerical results, we will investigate factors that potentially affect simulation performance and update the results & visualizations. For sensitivity to input token distribution, we will examine more post-training datasets from multiple domains under different expert layouts in our revised version.\n\n**[Weakness 13. Energy & Area modeling]**\n\n|Qwen3-30B-A3B|Area (mm^2)|Energy/Step (J)|\n|-|-|-|\n|Attn|474|2842|\n|MoE|6562|3470|\n|Switches & Interconnect|5839|1320|\n|DRAM|1300|1541|\n\nAttn module has a smaller footprint than interconnect or MoE chiplets due to per-chiplet area constraints, limiting local caching & computational capacity and leading to more frequent intra-chiplet data movement and higher energy consumption regarding its area. Moreover, most energy is consumed by computing modules rather than DRAM or interconnect, confirming the effectiveness of our communication optimization, as switch energy remains low despite large-scale sparse expert activation.\n\n**[Questions 1. Expert Clustering]**\n\nSince profiling and expert assignment are conducted offline before deployment, the efficiency of these algorithms is not very critical; thus, we can try to solve the exact results for the binary programming task. We will try to solve the expert allocation task using alternative formulations in our future work.\n\n**[Question 2. DP/TP for Attention Layer]**\n\nThough DP and TP are not considered in Mozart (stated in \"Conclusion and Limitations\"), they can be easily adopted when attn requires more resources beyond an individual chiplet with area constraints. First, the NoP-Tree topology with hierarchical routing and in-network aggregation at switch chiplets naturally supports parallel patterns such as all-gather & reduce-scatter. Distributing attention heads (TP) or input tokens (DP) across multiple smaller attention chiplets would increase area usage and introduce additional inter-chiplet communication overhead. However, the increased computing resources may help reduce computing latency. The overall performance would reflect a trade-off among factors such as per-die area, the number of sub-chiplets, and routing complexity. Second, our chiplet partition is compatible with existing TP frameworks like Megatron-LM, as they typically partition the weight tensors within each model layer and perform frequent all-gather and reduce-scatter operations. Mozart’s NoP-Tree topology with in-network computing switches natively supports such communication primitives, offering efficient collective capabilities. Conversely, the mesh structure we employed in the single computing die is compatible with GEMM and GEMV computations with ring or torus logical topology. The dual-topology design enables Mozart to optimize both intra-chiplet computation and inter-chiplet communication overhead.\n\n**[Question 3. Expert allocation]**\n\nSince Mozart only targets post-training, it is rational to assume that the collaboration pattern would not change much during tuning. But for pre-training, the expert layout at each layer should be dynamically adjusted to continuously benefit from the efficient all-to-all communication it brings. It is complex, as it requires recording routing choices continuously, fast expert allocation algorithms, and dynamic expert layout updating, which is beyond the scope of this paper. We will try to make a more systematic resolution with algorithm-system co-design as our future work.\n\n**[Question 4. All-to-all communication]**\n\nMozart has conducted these 2 aspects, including hierarchical routing (first determine the routed chiplet cluster, then determine the specific chiplet) and topology-aware communication (the 2.5D NoP-tree topology) to accommodate this refactored routing. It can be seamlessly applied to any MoE-LLMs as it does not modify the routing policy of the pretrained model.\n\n**[Question 5. Extension to more LLMs]**\n\nPlease refer to *Weakness 9*."}, {"review_id": "une3zq3mhE", "reviewer": "Reviewer_6v38", "summary": "This paper presents a 3.5-D hardware architecture that leverages new in-package semiconductor technologies to tightly integrate small chiplets at a large scale, connected together using a hierarchical network on-chip. It is labeled as 3.5-D (instead of 2.5-D) because each chiplet has 3D integration of SRAM (memory) on top of a logic die as well. On the algorithmic side, the paper presents a number of architecture-specific optimizations, specifically to optimize the post-training fine-tuning of mixture-of-expert architectures such as Deepseek, Olmo, and others. The proposed optimizations are: (1) clustering experts and pins them to the suitable chiplet in the system hierarchy to minimize overall communication, and (2) pipelines memory access to hide computation latency.", "questions": "The main missing component of this paper is to compare to any existing hardware solution to enable some understanding of the impact of this work. Are you able to compare, for example, to 8xA100?", "limitations": "N/A", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: the presented architecture is very interesting and combines many new ideas to tackle large-model fine-tuning. Furthermore, there are RTL simulated areas with details down to the pitch of TSVs suggesting a realistic hardware architecture. Although the focus of the paper is on the mapping of MOE model mapping onto the device. On that front, the presented ideas are interesting and logical, but:\n\nWeaknesses:\n1- There are no comparisons, even high-level, rough comparisons to any available hardware. This makes it very difficult to appreciate the presented hardware architecture, its uitility, or efficiency compared to available solutions such as GPUs.\n2- Validation: Barring tape-out, it is conventional in the hardware space to map a small design onto an FPGA or a cycle-accurate simulator to validate the cycle-accurate assumptions made in the paper.\n3- MOE mapping methods are somewhat straightforward. If I assume the novelty of the paper is on the hardware side, I find that this is not the right venue for presentation, if the novelty is in specifically mapping MOE to this new hardware architecture, I find that the presented methods are practically sound and logical but somewhat straightforward.", "was_revised": true, "final_justification": "Additional comparisons to GPUs, while rough, strengthen the paper. However, I am not totally convinced of the evaluation and methodology yet.", "rebuttal": "We thank reviewer 6v38 for acknowledging our work, which “is interesting and logical, combines many new ideas”. We supplement the omitted runtime comparison and the novelty of mapping to address your concerns.\n\n**[Weakness 1 & Questions. Comparison with A100]**\n\nThanks for this constructive suggestion. We conducted comparisons on an advanced server with 8x NVIDIA A100 GPUs (each with 80GB HBM) to mimic $\\texttt{Mozart}$, which features 4 chiplet clusters and has a similar total power consumption. Any two GPUs are interconnected with 12 NVLink in this server, providing a bidirectional bandwidth of around 400 GB/s (as verified through peer-to-peer communication using send & recv in torch.distributed). The CPU and memory configurations of this server are:\n\n|CPU Config|Values|\n|-|-|\n|CPU(s):|152|\n|On-line CPU(s) list:|0-151|\n|Model name:|Intel(R) Xeon(R) Platinum 8368 CPU @ 2.40GHz|\n|Thread(s) per core:|2|\n|Core(s) per socket:|38|\n|Socket(s):|2|\n|CPU(s) scaling MHz:|31%|\n|NUMA node0 CPU(s):|0-37,76-113|\n|NUMA node1 CPU(s):|38-75,114-151|\n\n|Memory Config|total|used|free|shared|buff/cache|available|\n|-|-|-|-|-|-|-|\n|Mem:|1.0Ti|158Gi|835Gi|33Gi|53Gi|849Gi|\n|Swap:|9Gi|4.2Gi|5.8Gi|\n\nWe adopt PyTorch’s Fully Sharded Data Parallel (FSDP)—the $de facto$ standard for fine-tuning open-source LLMs on the Hugging Face ecosystem—as our GPU baseline, leveraging its ZeRO-style parameter sharding and optimization strategies for advanced runtime efficiency and scalability. We set the local batch size to 4 in GPU baselines to keep consistent with the global batch size in $\\texttt{Mozart}$. We also report the average training latency per step for 1k iterations. We examine both offloading and non-offloading configurations, which are both commonly employed techniques when fine-tuning large models with FSDP. We utilize FlashAttention 2 to accelerate the attention module in all experiments, thereby minimizing the influence of attention computation overhead.\n\n*Sequence Length 128*\n|Model|Qwen3|DeepSeek|OLMoE|\n|-|-|-|-|\n|$\\texttt{Mozart}$ (HBM2)|2.65|0.57|1.48|\n|PyTorch FSDP|6.16 (2.32x)|1.14 (2x)|2.48 (1.68x)|\n|PyTorch FSDP (Offload)|28.36 (10.70x)|5.75 (10.09x)|13.39 (9.05x)|\n\n*Sequence Length 256*\n|Model|Qwen3|OLMoE|DeepSeek|\n|-|-|-|-|\n|Mozart (HBM 2)|2.75|0.77|1.56|\n|PyTorch FSDP|6.42 (2.33x)|1.23 (1.60x)|2.65 (1.70x)|\n|PyTorch FSDP (Offload)|28.70 (10.44x)|5.85 (7.60x)|13.78 (8.83x)|\n\n*Sequence Length 512*\n|Model|Qwen3|OLMoE|DeepSeek|\n|-|-|-|-|\n|Mozart (HBM 2)|3.26|1.04|2.46|\n|PyTorch FSDP|6.80 (2.09x)|1.73 (1.66x)|3.20 (1.30x)|\n|PyTorch FSDP (Offload)|29.64 (9.09x)|6.26 (6.02x)|14.21 (5.78x)|\n\nCompared to the GPU baselines, $\\texttt{Mozart}$ mainly features the following properties:\n- ***Improved efficiency***: $\\texttt{Mozart}$ consistently surpasses GPU counterparts across all the baselines with a significant reduction in training latency.\n- ***Better scalability***: $\\texttt{Mozart}$ demonstrates a higher speedup with Qwen3-30B-A3B model on both offloading and non-offloading settings, which contains much larger parameters than OLMoE-1B-7B and DeepSeek-MoE-16B.\n\n**[Weakness 2. Cycle-accurate simulator]**\n\nThank you for the valuable feedback. While we have not yet mapped the full Mozart system onto an FPGA or performed tape-out, we have implemented the logic chiplets, SRAM chiplets, inter-chiplet interconnects, and switches in Verilog and synthesized them with Synopsys Design Compiler under 28nm technology. Furthermore, we have developed a cycle-level simulator, whose key components have been validated against Verilog simulation results in terms of runtime and power to ensure accuracy. Therefore, our simulation methodology provides high fidelity, following the common practices.  As part of our future work, we plan to perform a full-system FPGA validation of Mozart and eventually proceed to tape-out. For simulation details, please refer to Section 5.2 of our paper.\n\n**[Weakness 3. Novelty]**\n\nWe respectfully disagree with the assessment that the contribution is either “hardware-only” or “straightforward.” Our paper focuses on algorithm-hardware co-design rather than pure hardware or algorithm innovations. Therefore, our innovations lie in 3 aspects, including hardware-side, algorithm-side, and co-design-side, listed as follows:\n\n*Hardware-side:*\n\nThe hardware innovations of $\\texttt{Mozart}$ are proposed to tackle the heterogeneity of attention and FFN modules, along with the heavy communication overhead in MoE-LLMs. The insights are listed as follows:\n- ***3.5D vertical logic-on-memory stack***: To improve memory locality, we design a hierarchical memory system aligned with the temporal reuse patterns of MoE-LLMs, enabling frequently reused data to be cached closer to the computing unit using the 3D logic-on-memory stack, thereby reducing costly accesses to off-chip DRAM; \n- ***2.5D NoP-Tree topology with switches & DRAMs***: Mozart adopts a 2.5D NoP-Tree interconnect where attention chiplets sit at the root and MoE chiplets form the leaves. Each group of four MoE chiplets shares a DRAM channel and connects through a switch chiplet with local aggregation support, reducing inter-group bandwidth contention. By co-locating frequently activated experts on the same chiplet and applying scheduling-aware dataflow, the system minimizes memory and communication bottlenecks to improve throughput.\n\n*Algorithm-side:*\n\nOur algorithm-side innovations focus on expert-chiplet mapping and fine-grained streaming schedules. The insights of $\\texttt{Mozart}$ are listed as:\n- ***Advancing communication efficiency with expert activation & collaboration awareness***: Our expert-chiplet mapping scheme does not simply follow conventional expert parallelism, i.e, evenly allocating experts on chiplets without considering the profiled routing priors. As we pointed out in Section 3.3, expert placement can affect the all-to-all communication volume, which is a critical bottleneck in expert parallelism. Therefore, we develop novel expert clustering & allocation algorithms to specialize in grouped expert layout for communication efficiency, which is demonstrated as $\\texttt{Mozart-B}$ & $\\texttt{Mozart-C}$ in Figure 6.\n- ***Streaming schedules for improved parallelism***: Our fine-grained streaming schedules stem from the heavy DRAM communication volume in MoE deployment, as well as the all-to-all communication overhead. We designed streaming tokens & experts as fine-grained scheduling algorithms to overlap the high communication budgets. The effectiveness is demonstrated in $\\texttt{Baseline}$ & $\\texttt{Mozart-A}$ in Figure 6, where the performance improvement can be more significant under longer input sequences.\n\n*Co-design-side:*\n\nTo efficiently deploy MoE-LLMs, we proposed more tightly-coupled designs to integrate the efficiency-oriented algorithm innovations and the sparsity-driven hardware hierarchy. The insights are listed as follows:\n- ***Expert allocation algorithm & NoP-tree topology***: We proposed a 2.5D NoP-tree interconnect specialized for the MoE deployment. The motivation lies in tackling the heterogeneity of the Attention and FFN modules with our proposed grouped expert-chiplet assignment for communication efficiency. We validate the effectiveness in Table 4, where our coupled design can reduce both all-to-all communication volume and latency.\n- ***Streaming schedule algorithm & Memory hierarchy***: We assign a DRAM chip for both the attention chiplet and MoE chiplet groups to alleviate the heavy burden of DRAM communication through specialized storage, and employ streaming algorithms to overlap off-package communications with on-package computations fully."}, {"review_id": "jgKjZ06wiU", "reviewer": "Reviewer_X4mS", "summary": "Mozart is a novel and comprehensive hardware-software co-design framework for training MoE models on 3.5D wafer-scale chiplet architectures. On the software part, Mozart optimizes the expert placement to reduce inter-chiplet communication, and it uses streaming to improve the communication-compute overlap. On the hardware part, Mozart can place heterogeneous components on each chiplet, and it connects them via a 2.5D NoP-Tree.", "questions": "Overall, I enjoy reading this paper and have learned a lot!\n\nSee weakness 1, the authors can discuss more on why the proposed overlap scheme is better than existing works for this hardware.\n\nTo support larger models, how is the proposed solution compatible with existing distributed training optimizations like model parallelism?", "limitations": "The paper is great! I can only observe one limitation.\n\nSee weakness 2, as LLMs scale quickly, how does a Mozart hardware manufactured today support future larger models?\nIf we support that by adapting model parallelism on top of your design, will your performance benefits degrade?\nAnd if so, by how much?\nDiscussing these questions would help address the limitations on hardware scalability.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n1. This paper focuses on popular topics like MoE, wafer-scale, and 3D architecture.\n\n2. This paper proposes a comprehensive solution for placing the experts based on the hardware architecture.\n\n3. The evaluation shows impressive performance benefits.\n\nWeakness:\n\n1. The compute-communication overlap seems similar to existing works like DualPipe from DeepSeek and TileLink from ByteDance.\n\n2. This paper did not discuss that, as model sizes scale quickly, what will happen if the size of an expert exceeds the capacity of one chiplet group. Specifically, when each expert is mapped to multiple chiplet groups, how will you handle the parallelism?", "was_revised": true, "final_justification": "All concerns solved in the rebuttal.\nThis paper well reveal a lot of new knowledge to the community.\nHowever, you may consider directly touching the points in your rebuttal next time, without using too many lengthy descriptions and fancy words.\nFor the paper itself, I am still positive about it.", "rebuttal": "We sincerely thank reviewer X4mS for acknowledging our work, especially “enjoy reading this paper and have learned a lot”. To address your concerns, we provide a pointwise response below:\n\n**[Weakness 1 & Questions 1. Communication-computation overlap scheme]**\n\nWe appreciate your insightful comments on introducing the relevant research. Although both DualPipe and TileLink are tailored for large-scale GPU clusters, they share some similar design principles with our work, while $\\texttt{Mozart}$ also encompasses some specialized mechanisms beyond their scope. We will present a comprehensive comparison, along with elaborating our design innovations as follows:\n\n*Different design principles*\n\n- ***Communication type***: Both TileLink and DualPipe consider overlapping the communication across GPUs, which is typically constrained by the limited bandwidth. Mozart tries to overlap the communication between DRAM and chiplets, which is the dominant bottleneck on wafer-scale chiplet systems.\n- ***Overlapping strategies***: In $\\texttt{Mozart}$, the DRAM communication for MoE parameters is overlapped with both the computation for streaming tokens and the hierarchical all-to-all communication across different chiplets, while the DRAM communication for attention parameters is overlapped with the computation for streaming tokens. In TileLink, the cross-device communication is mainly overlapped with tiled computing tasks that are executed alternately. In DualPipe, the cross-device communication is mainly overlapped with computation operations from both fwd and bwd passes (e.g, in fwd pass all-to-all communication can be overlapped with attention computation for other micro batches). Since the architecture hierarchy and performance bottleneck of GPU clusters and our chiplet systems are quite varied, our overlapping strategies are highly specialized for chiplet systems to alleviate the dominant bottleneck.\n\n*Innovations of $\\texttt{Mozart}$*\n\n- ***Streaming tokens***: Both DualPipe and TileLink used the concept of micro-batch, while $\\texttt{Mozart}$ introduced a similar notion as streaming tokens, which differs from these previous notions. Our motivation is to utilize the heterogeneity of attention and FFN modules on chiplet systems. The streaming tokens are separated for attention and FFN at each layer, isolated with synchronization (as shown in Figure 4). For attention, all the tokens are streamed into evenly distributed micro-batches and processed sequentially. For FFN, the tokens are streamed for each chiplet after the hierarchical all-to-all communication, aiming at overlapping the activation saving communication (between chiplet and DRAM) with expert computation. Different MoE chiplets can be executed in parallel to adhere to the distributed MoE DRAMs.\n- ***Streaming experts***: Different from previous works, Mozart harnesses the profiled routing priors of the pre-trained MoE-LLMs to promote the overlap between off-package DRAM communication and on-package operations, through the proposed streaming experts. Since multiple MoE chiplets share the same DRAM, we load the clustered experts in sequential order following the profiled probability of being routed, which offers an additional benefit for runtime latency. \n\nIn a word, our work stems from the fundamental performance bottlenecks of distributed chiplet systems. By harnessing the profiled expert activation & co-activation priors, we effectively enhance the system runtime efficiency in the process of post-training for modern MoE-based LLMs, aiming at providing a generic framework for deployment of all the large-scale sparse models, in the era that sparse scaling has almost become the industrial standard, as the latest frontier models like DeepSeek, Kimi K2, Llama-4, and Qwen3 all possess this architecture. Compared to the pure system-level optimizations, we explore a novel perspective of algorithm-hardware co-design as an early-stage attempt for efficient MoE deployment on domain-specific hardware. \n\n**[Questions 2. Compatibility with distributed training optimizations]**\n\nModel parallelism techniques for large models mainly include pipeline parallel (PP), tensor parallel (TP), and expert parallel (EP). Currently, the trend for scaling LLMs is sparsely scaling the FFN module with MoE, with successful applications on DeepSeek, Qwen3, and Kimi K2. These large models feature the same characteristic, where the attention module only consumes a very small proportion of the total parameters (also demonstrated in Figure 1 on models with smaller scales). We list the potential applications of these model parallelism techniques below:\n- ***PP***: In $\\texttt{Mozart}$, we mainly consider the pipelining for attention and MoE within the same layer. While for inter-layer pipelining, $\\texttt{Mozart}$ architecture is also compatible by incorporating multiple replicas of individual $\\texttt{Mozart}$, and partitioning the integral pipeline evenly across them. The bubble preventing techniques (Zero bubble pipeline parallelism, Qi et al, ICLR 2024) are also compatible in this extended architecture.\n- ***TP***: Tensor parallelism is usually used on embeddings, lm_head, and attention modules. For $\\texttt{Mozart}$, these can all be placed on the attention chiplet for computation, loaded in sequential order from the attention DRAM.\n- ***EP***: $\\texttt{Mozart}$ mainly targets efficient expert parallelism on chiplet systems. For larger models with higher sparsity degree, we can extend the MoE chiplets, e.g, adopting 16 MoE chiplet clusters if extending routed experts from 64 to 256.\n\nIn a word, $\\texttt{Mozart}$ architecture is compatible and scalable for larger-scale distributed training of MoE-LLMs through the integration of multiple model parallelism techniques. \n\n**[Weakness 2 & Limitations. Scalability of $\\texttt{Mozart}$]**\n\nThank you for the insightful question. As demonstrated in our experiments, the current $\\texttt{Mozart}$ architecture is capable of supporting the typical expert sizes found in mainstream MoE models, meeting the demands of most practical applications today. To address the challenges brought by future model scaling, we envision two main directions for extending support. \n- ***Increasing the number of experts***: When model scaling is achieved by adding more experts, we can expand the number of switch groups at the hardware level to accommodate more expert chiplets. On the algorithmic side, we can also cluster or merge infrequently activated experts to limit the total number of experts, thereby avoiding communication congestion and maintaining efficient scheduling. \n- ***Growing the size of a single expert***: When an individual expert’s parameter size exceeds the capacity of a single chiplet group, we can adopt intra-expert model parallelism by partitioning the expert’s MLP module across multiple chiplet groups using tensor parallelism locally. Since $\\texttt{Mozart}$ adopts a hierarchical NoP-Tree topology, where switch chiplets natively support in-network aggregation and scheduling (e.g., reduce-scatter and all-gather operations), this architecture is well-suited for executing experts in parallel across chiplets. In summary, while $\\texttt{Mozart}$’s current design assumes a one-expert-per-group configuration, the underlying NoP-Tree interconnect and communication mechanisms are inherently scalable, making the architecture well-equipped to support expert scaling with tremendous sparsity and the continued evolution of large-scale MoE-based large language models."}]

Author rebuttal:

Meta-review: (a) Summarize the scientific claims and findings of the paper based on your own reading and characterizations from the reviewers.

This paper introduces Mozart, a software-hardware co-design framework for the efficient post-training of Mixture of Experts (MoE) LLMs. This work is motivated by the hardware challenges posed by the sparse and modular nature of MoE models especially memory locality and communication overhead. The proposed strategy tries to addresses expert routing and communication overheads in MoE-style LLMs, On the algorithmic side, Mozart implements an expert allocation strategy that clusters experts based on their co-activation patterns to minimize inter-chiplet communication and balance workloads. It also implements a fine-grained scheduling mechanism to improve the overlap between communication and computation. On the hardware side, Mozart designs a specialized 3.5D wafer-scale chiplet architecture with a hierarchical memory system and a 2.5D NoP-Tree interconnect topology.  It is labeled as 3.5-D (instead of 2.5-D) because each chiplet has 3D integration of SRAM (memory) on top of a logic die as well. On the algorithmic side, the paper presents a number of architecture-specific optimizations, specifically to optimize the post-training fine-tuning of mixture-of-expert architectures such as Deepseek, Olmo, and others. The proposed optimizations are: (1) clustering experts and pins them to the suitable chiplet in the system hierarchy to minimize overall communication, and (2) pipelines memory access to hide computation latency. Using cycle-accurate simulations, the paper shows that Mozart achieves speedups (over 1.9x) across three popular MoE models compared to baseline designs.
Despite its specialization to MoE-style LLMs and modular hardware assumptions, the paper presents valuable design principles for future accelerators. Its latency-aware expert scheduling, communication modeling, and hardware-algorithm co-design offer practical insights for chiplet-based and sparse-aware systems. These contributions are particularly timely as AI hardware trends shift toward modular packaging, inter-die communication optimization, and efficient support for dynamic, sparse workloads. While generalization beyond MoEs is limited, the paper nonetheless informs a growing class of AI accelerators targeting scalable, fine-grained model execution.

(b) What are the strengths of the paper?

The paper is well written and organized.
The authors address a timely topic. MoE architectures are increasingly being adopted to scale LLMs to trillion-parameters. Hence addressing their communication overhead and resource utilization bottleneck is an important and practical research direction.
The concept of an algorithm-hardware co-design frameworks specifically designed for MoE models on emerging 3.5D wafer-scale architectures is a strong contribution.
The work is well motivated by a clear analysis of the system-level challenges.
The of a cycle-accurate RTL simulator, with details down to the pitch of TSVs suggesting a realistic hardware architecture, that has been validated against Verilog simulations is also a strong point.
The empirical evaluation is generally well explained and done.

(c) What are the weaknesses of the paper? What might be missing in the submission?

The hardware evaluation is entirely simulation-based. While this is a critical first step in any new architecture research, the performance evaluations and gains depend heavily on the realization of the 3.5D wafer-scale system. However, this is a minor weakness at this point in the system's development. It is acceptable to have cycle-accurate simulations at this stage, as these provide significant evidence for the ultimate successful deployment of the physical hardware.

Some key relevant works are missing from the related work discussion:
GShard (Lepikhin et al., ICLR 2021): A foundational MoE system that introduced routing + automatic sharding across devices. Crucial for historical context. This is the first large-scale MoE at Google.
BASE Layers (Lewkowycz et al., NeurIPS 2022): A hybrid sparse-dense expert design that maintains performance while reducing memory. This should be compared and contrasted Mozart’s memory-aware scheduling.
MoEfication (Zhou et al., ICML 2022): Shows conversion of dense models to MoE post-training. This is related to the post-training MoE optimization approaches as proposed by Mozart.
TaskMoE / Expert Tuning (Zhang et al., 2022): Specialization for multi-task models.
Tutel & DeepSpeed-MoE need to be more discussed, especially in comparison to Mozart’s routing, scheduling, and clustering.

The clustering method used in Equation 5 is based on the k-clustering with pairwise distances. This algorithm is computationally expensive (O(n²)), and no mention is made of approximation methods (e.g., hierarchical clustering, METIS).
The mathematical formulation in Equation 6 does not provide constraints on the number of experts per token (e.g., top-2) of the token capacity per expert (e.g., max tokens per batch). This needs to be clarified by the authors.

The proposed co-design strategy effectively addresses expert routing and communication overheads in MoE-style LLMs, however, it is highly specialized to this sparse model family. The system assumes hard top-k expert gating and token-level dispatch, which do not generalize to dense autoregressive models (e.g., GPT-3, LLaMA), diffusion models, or state-space architectures (e.g., Mamba) that lack modular sparse execution patterns. Additionally, the hardware assumptions—modular chiplet-based designs with a Network-on-Package interconnect—limit applicability to a narrow class of accelerators. Real-world systems such as monolithic GPUs, wafer-scale engines (e.g., Cerebras), and mobile AI chips follow vastly different architectural and memory paradigms. The lack of runtime adaptability in expert scheduling and coarse treatment of memory hierarchies further constrain its deployability. A discussion on these generalization limitations is absent.

Regarding the empirical evaluation. There is a lack of baselines beyond static scheduling. The main baseline is a static MoE routing strategy or uniform tiling. It would be good to add comparisons against DeepSpeed-MoE or Tutel routing strategies and Existing expert placement heuristics (e.g., token-balanced sharding, load-aware dispatch). The existing evaluation is not enough to claim that Mozart significantly outperforms state-of-the-art MoE runtimes.

Limited System Scaling Study. The paper does not analyze how Mozart behaves under increasing number of chiplets, varying interconnect topologies (e.g., mesh vs. ring) and realistic memory capacity constraints. Such scaling or sensitivity analysis can showcase if the benefits hold at larger or smaller deployment scales. 

While latency and utilization are discussed, the paper completely omits energy or area modeling. Given the focus on chiplet hardware, this is a notable omission—MoE communication patterns can be power-intensive.

The compute-communication overlap seems similar to existing works like DualPipe from DeepSeek and TileLink from ByteDance.

(d) Provide the most important reasons for your decision to accept/reject. For spotlights or orals explain why the paper stands out (other than by high scores or popularity trends).

As the title of the conference is Neural Information Processing *Systems*, it is good to see papers exploring the forefront of the computational substrate for these systems. One of the drivers of the current renaissance of machine learning has been the massive increase in computational capacity due to high-performance GPU systems. Wafer-scale systems seem to be the next frontier that will drive the next wave of development in machine learning. The paper presents, and justifies/validates, a new chiplet based wafer scale substrate supporting the next generation of machine learning systems.

(e) Summarize the discussion and changes during the rebuttal period. What were the points raised by the reviewers? How were each of these points addressed by the authors? How did you weigh in each point in your final decision? 

The rebuttal discussion was extensive, and the authors did a lot of work in addressing the issues raised by the reviewers. 

One key issue that was clarified in the rebuttal discussion was the contributions made on the hardware side. As doing FPGA, ASIC or full custom VLSI is premature (and expensive) at this stage, the detailed cycle-accurate RTL simulations are sufficient to establish the benefits and feasibility of the approach. 

As requested by reviewer 6v38 the authors provided extensive comparisons on an advanced server with 8x NVIDIA A100 GPUs (each with 80GB HBM) to mimic Mozart, which features 4 chiplet clusters and has a similar total power consumption. This comparison was non-trivial and must have involved a great deal of additional work on the part of the authors.
Decision: Accept (spotlight)