Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

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:

[Weakness 1 & Questions 1. Communication-computation overlap scheme]

We 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 $`Mozart`$ also encompasses some specialized mechanisms beyond their scope. We will present a comprehensive comparison, along with elaborating our design innovations as follows:

Different design principles

• 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.
• Overlapping strategies: In $`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.

Innovations of $`Mozart`$

• Streaming tokens: Both DualPipe and TileLink used the concept of micro-batch, while $`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.
• 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.

In 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.

[Questions 2. Compatibility with distributed training optimizations]

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

• PP: In $`Mozart`$, we mainly consider the pipelining for attention and MoE within the same layer. While for inter-layer pipelining, $`Mozart`$ architecture is also compatible by incorporating multiple replicas of individual $`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.
• TP: Tensor parallelism is usually used on embeddings, lm_head, and attention modules. For $`Mozart`$, these can all be placed on the attention chiplet for computation, loaded in sequential order from the attention DRAM.
• EP: $`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.

In a word, $`Mozart`$ architecture is compatible and scalable for larger-scale distributed training of MoE-LLMs through the integration of multiple model parallelism techniques.

[Weakness 2 & Limitations. Scalability of $`Mozart`$]

Thank you for the insightful question. As demonstrated in our experiments, the current $`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.

• 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.
• 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 $`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 $`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.

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.

[Weakness 1 & Questions. Comparison with A100]

Thanks for this constructive suggestion. We conducted 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. 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:

We 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 $`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.

Sequence Length 128

Sequence Length 256

Sequence Length 512

Compared to the GPU baselines, $`Mozart`$ mainly features the following properties:

• Improved efficiency: $`Mozart`$ consistently surpasses GPU counterparts across all the baselines with a significant reduction in training latency.
• Better scalability: $`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.

[Weakness 2. Cycle-accurate simulator]

Thank 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.

[Weakness 3. Novelty]

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

Hardware-side:

The hardware innovations of $`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:

• 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;
• 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.

Algorithm-side:

Our algorithm-side innovations focus on expert-chiplet mapping and fine-grained streaming schedules. The insights of $`Mozart`$ are listed as:

• 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 $`Mozart-B`$ & $`Mozart-C`$ in Figure 6.
• 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 $`Baseline`$ & $`Mozart-A`$ in Figure 6, where the performance improvement can be more significant under longer input sequences.

Co-design-side:

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

• 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.
• 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.

We sincerely thank reviewer Bfvg for the insightful and detailed comments. To address your questions, we provide pointwise responses below.

[Weakness 1. Assumptions of the simulator]

Mozart 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.

Each 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.

Each 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.

[Weakness 2. Partially correct statement]

We will refine it following your constructive comments.

[Weakness 3. Incorrect citation on Future of Memory]

We apologize for the omission, and will replace [22] with Cambricon-LLM (Yu et al, ISCA 2024) and Hecaton (Huang et al, arxiv 2024).

[Weakness 4 & 5. Related works and discussions]

We will add a subsection after Modularized LLMs to incorporate these fundamental works on GPU-deployment of MoE-LLM.

[Weakness 6. Computation complexity of expert clustering]

The 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.

[Weakness 7. TopK]

We follow the dropless MoE design in MegaBlocks without the notion of expert capacity. We will claim it in our revised version.

[Weakness 8. Runtime score]

We 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.

[Weakness 9 & Question 5. Generalization]

Mozart 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.

From 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.

Though 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.

Overall, although Mozart is optimized for MoE deployment, its core principles, modular compute decomposition, communication-aware scheduling, and logical topology design, are highly extensible.

[Weakness 10. Baselines]

We 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.

[Weakness 11. System scaling]

Mozart is an empirical optimal solution for MoE-LLMs because:

• 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.
• 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.
• We apply mesh to establish intra-chiplet connections, since it is feasible for regular data movement in GEMM.

[Weakness 12. Comprehensive results]

For 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.

[Weakness 13. Energy & Area modeling]

Attn 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.

[Questions 1. Expert Clustering]

Since 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.

[Question 2. DP/TP for Attention Layer]

Though 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.

[Question 3. Expert allocation]

Since 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.

[Question 4. All-to-all communication]

Mozart 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.

[Question 5. Extension to more LLMs]

Please refer to Weakness 9.