MoE decoding speeds up with multi-engine dataflow on tensor accelerators
A Multi-Engine Dataflow for MoE Decoding on Scratchpad-Based Tensor Accelerators
Hardware Architecture
Summary
Mixture-of-Experts (MoE) models use many separate 'experts' to process data, but when running on special hardware called tensor accelerators, moving expert information slows things down. The authors present CARDAN, a way to represent expert data so that some parts are shared and some are unique, allowing better use of multiple compute engines at the same time. This reduces waiting times and speeds up the decoding, meaning the model can generate results faster without losing accuracy. Tests on different MoE models show consistent speed improvements and maintained or better prediction quality.
What this means in practice
- •For machine learning engineers: Deploy MoE models on tensor accelerators with faster, more efficient decoding at small and moderate batch sizes.
- •For hardware accelerator designers: Design tensor accelerator dataflows that overlap data transfer and computation for MoE workloads using shared and private expert representations.
Authors
Bin Ma, Wenjie Fan, Dong Li
Abstract
Mixture-of-Experts (MoE) decoding on scratchpad-based tensor accelerators (STA) is dominated by moving expert weights while the compute engines sit idle. This traffic is hard to hide, because the experts are known only after routing, and hard to shrink without losing quality or adding critical-path work. We present CARDAN, which represents each expert-weight matrix as a vector-quantized component plus a shared-basis low-rank component and co-designs this representation with a multi-engine decoding dataflow. The representation separates expert-common from expert-private work, so the dataflow overlaps DMA with computation on several engines. Across five MoE families on AWS Trainium3, CARDAN matches or improves BF16-teacher perplexity across all five models and speeds up batch-one decoding by 1.15-1.31x over AWS dense MoE megakernels, rising to 1.7x at batch size 16.