cs.ARMar 29, 2026
SaveExpert Streaming: Accelerating Low-Batch MoE Inference via Multi-chiplet Architecture and Dynamic Expert Trajectory Scheduling
Organizations: AI Chip Center for Emerging Smart Systems, Hong Kong SAR, China · The Hong Kong University of Science and Technology, Hong Kong SAR, China · Shanghai UniVista Industrial Software Group Co., Ltd., Shanghai, China
Abstract
Mixture-of-Experts is a promising approach for edge AI with low-batch inference. Yet, on-device deployments often face limited on-chip memory and severe workload imbalance; the prevalent use of offloading further incurs off-chip memory access bottlenecks. Moreover, MoE sparsity and dynamic gating shift distributed strategies toward much finer granularity and introduce runtime scheduling considerations. Recently, high die-to-die bandwidth chiplet interconnects have created new opportunities for multi-chiplet systems to address workload imbalance and offloading bottlenecks with fine-grained scheduling. In this paper, we propose Fully Sharded Expert Data Parallelism, a parallelization paradigm specifically architected for low-batch MoE inference on multi-chiplet accelerators. FSE-DP attains adaptive computation-communication overlap and balanced load by orchestrating fine-grained, complementary expert streams along dynamic trajectories across high-bandwidth D2D links. The attendant dataflow complexity is tamed by a minimal, hardware-amenable set of virtualization rules and a lightweight scheduling algorithm. Our approach achieves 1.22 to 2.00 times speedup over state-of-the-art baselines and saves up to 78.8 percent on-chip memory.
Figures & tables
Fig. 1: Typical template of (a) multi-chiplet-based AI accelerator. (b) Mixture-of-Experts network.
Fig. 2: (a) Shapes for different models (b,c) Long-tail effect of MoE models under different batch sizes. The figure shows the number of tokens processed in a specific layer for DeepSeek-MoE-16B [ 48 ] on the Wikitext-2 dataset [ 41 ] and Qwen3-30B-A3B [ 75 ] on WinoGrande [ 49 ] . Experts on the x-axis are sorted by the number of tokens they process; the y-axis gives the token count per expert. The long-tail effect is more pronounced at smaller token numbers. R denotes different requests.
Fig. 3: Fully sharded expert–data parallelism (example: four chiplets compute expert 1). The figure illustrates how a single expert is computed within one MCM: expert 1 is evenly sharded into slices across chiplets (E1-S1–E1-S4). black tokens denote the tokens that activate expert 1, while non-highlighted (gray/black) tokens are other buffered tokens on that chiplet that do not activate expert 1. Before computing expert 1, tokens are redispatched across chiplets to balance the number of black tokens per chiplet for load balancing. denotes a request, and Seq denotes the activated token sequence stored on a chiplet. During expert computation, chiplets can exchange data in two equivalent ways to cover all black tokens: (a) keep token sequences fixed and circulate expert-1 slices so each slice visits the chiplets holding black tokens; or (b) keep expert 1 slices fixed and circulate black-token sequences so they visit all chiplets containing slices of expert 1.
Fig. 4: Micro-slice flow for overlapping D2D communication and computation (example: chiplet 1 while computing expert 1). Each expert slice (E1-S1–E1-S4) is further partitioned into micro-slices (M1–M4). (a) Baseline micro-slice overlap. In each step, chiplet 1 computes its local sequence using the current micro-slice, while concurrently receiving the next micro-slice from a neighbor chiplet and sending the just-computed micro-slice to the next chiplet; arrows indicate the D2D transfers that overlap with the compute stage. The weight buffer show the micro-slice storage on chiplet 1 over time, where each row is one micro-slice-sized buffer slot and each column is a time step. A colored cell indicates that the slot stores a micro-slice from a specific expert slice (see the color legend), and the numeral in the cell is the micro-slice index within that slice. Blank cells are free slots, and bold numerals mark the micro-slice being computed in that step. (b) Eager micro-slice usage. Chiplet 1 immediately forwards the micro-slice under computation and, in the next step, computes the most recently received micro-slice, so each micro-slice quickly traverses all chiplets and can be discarded earlier, reducing average weight-buffer occupancy.
Fig. 5: Demonstration of paired-load policy and token buffering policy.
Fig. 6: Flow fusion of DDR load and D2D micro-slice flow (example: chiplet 1 while computing expert 1 paired expert 4). This figure extends Figure 4 (b) by fusing off-chip DDR loading with the on-chip D2D micro-slice circulation, illustrated on chiplet 1 while co-executing two experts (E1 in red and E4 in purple under paired-load). DDR loads of the locally assigned micro-slices (e.g., E1-S1-M* and E4-S1-M*) are pipelined with computation and D2D receive/send of micro-slices from other chiplets. DDR load latency is assumed to be one micro-slice compute step. The expert-buffer grid is read as in Figure 4 : the first two rows cache DDR-loaded micro-slices for E1 and E4, respectively, while the remaining rows serve as a shared staging area for interleaved D2D-received micro-slices from both experts. Bold numerals indicate the micro-slices being computed at each time step.
Fig. 7: Virtualization demonstration under micro-slice streaming (example: 4 chiplets, ring expert trajectory). Each panel shows the on-chip storage state of the four chiplets at three consecutive time steps ( – ): each block is a stored micro-slice (labeled by expert and index, e.g., E1-M1), and arrows indicate next-hop forwarding along the ring trajectory. (a) Baseline case where expert 1 is evenly sharded across chiplets; achieving efficient overlap with explicit buffer assignment requires careful placement (as in Figure 6 ). (b) Expert 1 is unevenly distributed across chiplets, yet the computation proceeds correctly once micro-slices start streaming along the trajectory. For simplicity, (a) and (b) assume a ring topology in which micro-slices move from Chiplet 1 to 2, 2 to 3, and so on, to introduce the background of virtualization.(c) There experts execute concurrently with highly uneven on-chip arrangement. Despite the seemingly irregular trajectories, virtualization rules abstract away these runtime details and still realize a smooth flow. The key message is that the hardware scheduler need not reason about per-chiplet implementation: as long as micro-slices follow the prescribed trajectory, correct and efficient execution emerges without fine-grained bookkeeping.
Fig. 8: Scheduler hardware design. The scheduler is implemented in the IO die for task allocation.
| Component | Specification |
|---|---|
| DDR3-1600 4 25.6 GB/s, 800 MHz | |
| NoP & DDR | 2D Mesh, Multiple UCIe D2D IPs: 288 GB/s |
| 4-24 Gbps/pin, FDI-to-FDI Latency: 4.02 ns | |
| Samsung 5nm 1P13M CMOS, 800 MHz | |
| Compute Die | 2048 MACs, 0.675-0.9 V, 736.5 -2187 mW |
| 2.69mm 4.72mm, 4.865 TOPS |
TABLE I: Hardware and Model Configurations for Evaluation
Fig. 9: Single MoE layer latency with different models, datasets and input token counts.
Fig. 10: Chiplet floorplan and package photo
Fig. 11: Utilization fluctuation during inference of one layer.
Fig. 12: On-chip memory usage of different models.
Fig. 13: Activity timeline of expert trajectories across chiplets under FSE-DP (paired load). Qwen3-MoE, C4 with 256 input tokens, a runtime snapshot segment.
Fig. 14: End-to-end throughput comparison across model–dataset combinations. +10% means paired-load policy with 10% token buffering slackness.
Fig. 15: Ablation study on key design knobs of FSE-DP.
Fig. 16: DSE. Qwen3-MoE-A3B, C4, 64 input tokens. The star is the position of our test chip. Much larger scales are preferred in practice. (a) Fix D2D bandwidth at 288 GB/s. (b) Fix buffer size at 14MB.
Fig. 17: Granularity sensitivity. Latency heatmap over on-chip expert weight storage size and micro-slice number, evaluated on (a) Phi-3.5 and (b) Qwen3-MoE-A3B using C4.
Fig. 18: Scalability (utilization) evaluation based on Qwen3-MoE-A3B, C4.
More in this topic
Mixture-of-Experts Inference
Mixture-of-experts inference studies how to run sparsely activated expert models efficiently while preserving their quality.
See all 175 papers