Decoding the Skew: Distribution-Aware MoE Inference with Adaptive Kernel Dispatch
This paper introduces a distribution-aware framework for modeling and benchmarking Mixture-of-Experts (MoE) inference, showing that the best fused-MoE kernel changes with routing skew and token count, and presents DA-MoE, a GPU-resident kernel-dispatch runtime that improves geomean fused-MoE latency.
The paper introduces a distribution-aware framework for MoE inference and a GPU-resident runtime, DA-MoE, to improve its performance.
Before reading this…
Applications
- →MoE inference in machine learning models
To understand this paper, make sure you know these concepts first:
- Understanding of Mixture-of-Experts (MoE) inference and serving systemsfind papers →
Abstract
More Like ThisMixture-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.