一句话:提出专家感知的多批次流水线调度与预取机制,消除 MoE 卸载推理中的流水线气泡,吞吐最高提升 85.12 倍。
混合专家(MoE)以稀疏结构支撑万亿参数语言模型而不显著增加计算开销,但庞大的参数量使推理受限于 GPU 显存,需借助 CPU 与磁盘卸载并并行化 I/O 与计算。由于单个专家的计算时间往往小于其加载时间,流水线中产生大量气泡。本文提出 Klotski 推理引擎,采用专家感知的多批次流水线范式:通过批处理延长当前层的计算时间,使其与下一层的加载时间重叠。针对更多批次会激活更多专家、从而延长加载时间并加剧气泡的问题,Klotski 依据不同批次下专家的异构计算/IO 需求与激活模式编排推理顺序,以平衡二者时间;并设计约束敏感的 I/O-计算规划器与相关性感知的专家预取器,使调度适配不同硬件环境与模型。实验表明,Klotski 相比现有最优技术取得更优的吞吐-延迟权衡,吞吐提升最高达 85.12 倍。
Mixture of Experts (MoE), with its distinctive sparse structure, enables the scaling of language models up to trillions of parameters without significantly increasing computational costs. However, the substantial parameter size presents a challenge for inference, as the expansion in GPU memory cannot keep pace with the growth in parameters. Although offloading techniques utilise memory from the CPU and disk and parallelise the I/O and computation for efficiency, the computation for each expert in MoE models is often less than the I/O, resulting in numerous bubbles in the pipeline. Therefore, we propose Klotski, an efficient MoE inference engine that significantly reduces pipeline bubbles through a novel expert-aware multi-batch pipeline paradigm. The proposed paradigm uses batch processing to extend the computation time of the current layer to overlap with the loading time of the next layer. Although this idea has been effectively applied to dense models, more batches may activate more experts in the MoE, leading to longer loading times and more bubbles. Thus, unlike traditional approaches, we balance computation and I/O time and minimise bubbles by orchestrating their inference orders based on their heterogeneous computation and I/O requirements and activation patterns under different batch numbers. Moreover, to adapt to different hardware environments and models, we design a constraint-sensitive I/O-compute planner and a correlation-aware expert prefetcher for a schedule that minimises pipeline bubbles. Experimental results demonstrate that Klotski achieves a superior throughput-latency trade-off compared to state-of-the-art techniques, with throughput improvements of up to 85.12x.
@inproceedings{FangHHLC0YZ25,
title = {K
<scp>lotski</scp>
: Efficient Mixture-of-Expert Inference via Expert-Aware Multi-Batch Pipeline},
author = {Zhiyuan Fang and Yuegui Huang and Zicong Hong and Yufeng Lyu and Wuhui Chen and Yue Yu and Fan Yu and Zibin Zheng},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3676641.3716261}
}