Decoding the Skew: Distribution-Aware MoE Inference with Adaptive Kernel Dispatch

2026-07-25

AI Summary

This work addresses the inefficiencies in existing Mixture-of-Experts (MoE) inference systems, which suffer from high padding overhead, low memory reuse, and degraded kernel efficiency due to neglecting runtime routing distributions. To overcome these limitations, we propose DA-MoE, the first distribution-aware, GPU-resident MoE inference framework. DA-MoE introduces an Effective Experts metric and employs inverse Dirichlet modeling to generate controllable routing distributions. It further integrates offline tuning with online matching to enable adaptive fused kernel scheduling without CPU–GPU synchronization. Evaluated on HumanEval-X execution traces, DA-MoE achieves geometric mean latency speedups of 1.16× and 1.29× on DeepSeek-V3 and Kimi K2, respectively, with peak speedups reaching 1.40× and 1.56×.

Abstract

Mixture-of-Experts (MoE) inference consists of sparse expert GEMMs whose shapes vary with the runtime routing distribution. Existing serving systems typically select fused-MoE kernels using static token-count buckets, ignoring the per-expert routing distribution that determines tile padding, memory reuse, and kernel efficiency. We introduce a distribution-aware framework for modeling and benchmarking MoE inference. The framework combines the compact Effective Experts metric with a Dirichlet-based reverse-modeling procedure that generates controllable routing distributions for systematic hardware studies. Using it, we show that the best fused-MoE kernel changes with routing skew and token count. We further present DA-MoE, a GPU-resident kernel-dispatch runtime for NVIDIA GPUs that matches the live routing histogram to offline-tuned distributions and selects a near-optimal fused-MoE kernel without CPU--GPU synchronization. On HumanEval-X serving traces, DA-MoE improves geomean fused-MoE latency by 1.16X on DeepSeek-V3 and 1.29X on Kimi K2, with peak speedups of 1.40X and 1.56X.

Problem

• Mixture-of-Experts
• routing distribution
• kernel selection
• inference efficiency
• distribution skew

Innovation

• Mixture-of-Experts
• distribution-aware inference
• adaptive kernel dispatch
• Effective Experts
• GPU runtime