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
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. R 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 4× 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 ( t=1 – 3 ): 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.
Mixture-of-Experts (MoE) inference on memory-constrained edge platforms is often dominated by expert-weight movement rather than computation. Although each token activates only a small subset of experts, routing is input-dependent, so all experts must remain accessible and memory movement dominates latency. Existing approaches such as uniform pruning and expert dropping overlook two properties of MoE execution: expert combinations are highly skewed in frequency, and experts contribute asymmetrically within each combination. We present CRISP, a combination-aware framework that co-designs pruning and runtime scheduling for efficient MoE serving on memory-constrained edge platforms. CRISP reorders neurons by importance so one stored expert can support multiple execution widths, assigns differentiated execution to hot, warm, and cold combinations based on frequency and expert role, and exploits heterogeneous widths to overlap expert loading with computation. Across Mixtral-8x7B and DeepSeek-MoE-16B, CRISP improves throughput by up to 2.1× at the same memory budget while limiting accuracy loss to 3.6% relative to full-width inference. Our results highlight that efficient edge MoE serving requires jointly optimizing which expert capacity is preserved and how expert execution is scheduled.
Congwei Zhang, Youming Tao, Yifei Zou +6
School of Electrical Engineering and Computer Science, TU Berlin, Berlin, Germany · School of Computer Science and Technology, Shandong University, Qingdao, China · Department of Data Science, City University of Hong Kong, Hong Kong, China
Despite the computational efficiency of MoE models, the excessive memory footprint and I/O overhead inherent in multi-expert architectures pose formidable challenges for real-time inference on resource-constrained edge platforms. While existing static methods struggle with a rigid latency-accuracy trade-off, we observe that expert importance is highly skewed and depth-dependent. Motivated by these insights, we propose DyMoE, a dynamic mixed-precision quantization framework designed for high-performance edge inference. Leveraging insights into expert importance skewness and depth-dependent sensitivity, DyMoE introduces: (1) importance-aware prioritization to dynamically quantize experts at runtime; (2) depth-adaptive scheduling to preserve semantic integrity in critical layers; and (3) look-ahead prefetching to overlap I/O stalls. Experimental results on commercial edge hardware show that DyMoE reduces Time-to-First-Token (TTFT) by 3.44x-22.7x and up to a 14.58x speedup in Time-Per-Output-Token (TPOT) compared to state-of-the-art offloading baselines, enabling real-time, accuracy-preserving MoE inference on resource-constrained edge devices.
Yuegui Huang, Zhiyuan Fang, Weiqi Luo +3
Sun Yat-sen University, Guangzhou, China · Tencent, Shenzhen, China
Mixture-of-Experts (MoE) architectures are evolving towards finer granularity to improve parameter efficiency. However, existing MoE designs face an inherent trade-off between the granularity of expert specialization and hardware execution efficiency. We propose OmniMoE, a system-algorithm co-designed framework that pushes expert granularity to its logical extreme. OmniMoE introduces vector-level Atomic Experts, enabling scalable routing and execution within a single MoE layer, while retaining a shared dense MLP branch for general-purpose processing. Although this atomic design maximizes capacity, it poses severe challenges for routing complexity and memory access. To address these, OmniMoE adopts a system-algorithm co-design: (i) a Cartesian Product Router that decomposes the massive index space to reduce routing complexity from O(N) to O(sqrt(N)); and (ii) Expert-Centric Scheduling that inverts the execution order to turn scattered, memory-bound lookups into efficient dense matrix operations. Validated on seven benchmarks, OmniMoE (with 1.7B active parameters) achieves 50.9% zero-shot accuracy across seven benchmarks, outperforming coarse-grained (e.g., DeepSeekMoE) and fine-grained (e.g., PEER) baselines. Crucially, OmniMoE reduces inference latency from 73ms to 6.7ms (a 10.9-fold speedup) compared to PEER, demonstrating that massive-scale fine-grained MoE can be fast and accurate. Our code is open-sourced at https://github.com/flash-algo/omni-moe.
Jingze Shi, Zhangyang Peng, Yizhang Zhu +3
The Hong Kong University of Science and Technology (Guangzhou) · Beijing Academy of Artificial Intelligence