A Scheduling Framework for Efficient MoE Inference on Edge GPU-NDP Systems

TL;DR

The paper tackles the memory and compute bottlenecks of deploying large MoE models at the edge by introducing a scheduling framework that jointly uses tensor parallelism across NDP-DIMMs, a load-balancing strategy between GPU and NDP units, and dataset-free pre-fetching to reduce activation delays. The approach enables single-batch MoE inference on GPU-NDP DIMM systems, achieving up to 2.56x end-to-end latency reductions and a 2.41x average speedup over prior state-of-the-art. Key contributions include a two-stage tensor partitioning scheme, analytical load-balancing formulas, and a three-step runtime prefetching pipeline that adapts to actual workload patterns without calibration data. The framework significantly improves MoE inference efficiency in resource-constrained edge environments, enabling cost-effective and scalable edge deployments.

Abstract

Mixture-of-Experts (MoE) models facilitate edge deployment by decoupling model capacity from active computation, yet their large memory footprint drives the need for GPU systems with near-data processing (NDP) capabilities that offload experts to dedicated processing units. However, deploying MoE models on such edge-based GPU-NDP systems faces three critical challenges: 1) severe load imbalance across NDP units due to non-uniform expert selection and expert parallelism, 2) insufficient GPU utilization during expert computation within NDP units, and 3) extensive data pre-profiling necessitated by unpredictable expert activation patterns for pre-fetching. To address these challenges, this paper proposes an efficient inference framework featuring three key optimizations. First, the underexplored tensor parallelism in MoE inference is exploited to partition and compute large expert parameters across multiple NDP units simultaneously towards edge low-batch scenarios. Second, a load-balancing-aware scheduling algorithm distributes expert computations across NDP units and GPU to maximize resource utilization. Third, a dataset-free pre-fetching strategy proactively loads frequently accessed experts to minimize activation delays. Experimental results show that our framework enables GPU-NDP systems to achieve 2.41x on average and up to 2.56x speedup in end-to-end latency compared to state-of-the-art approaches, significantly enhancing MoE inference efficiency in resource-constrained environments.

Figures (11)

• Figure 1: The rising memory gap between consumer-grade GPU VRAM and large MoE models motivates GPU-NDP systems as a promising cost-effective edge solution. Our scheduling framework enables efficient MoE inference on such systems via tensor parallelism, load balancing, and dataset-free pre-fetching.
• Figure 2: Overview of GPU-NDP DIMM System.
• Figure 3: Comparison between MoE workflow scheduling in the decode stage. The proposed scheduling framework is based on (b) tensor parallelism and supplemented by (c) load balancing and (d) dataset-free pre-fetching.
• Figure 4: The computing process at the MoE layer with our introduced tensor parallelism. The left side is the common structure of the MoE Expert FFN, and the right side is the deployment of one expert to three computing devices.
• Figure 5: MoE workflow scheduling in the prefill stage within our framework.