paperbot · PL 论文追踪

全部论文 待读 收藏

P2 K <scp>lotski</scp> : Efficient Mixture-of-Expert Inference via Expert-Aware Multi-Batch Pipeline

Klotski:基于专家感知多批次流水线的高效混合专家推理
ASPLOS Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 22025匹配度 2/10 · 该论文聚焦MoE推理引擎的流水线调度与I/O优化,属于机器学习系统/体系结构方向,并未涉及用户核心关注的静态分析、形式化验证、浮点程序分析或贝叶斯程序分析,也未与LLM/Agent作为工具改进程序分析产生交叉。引用 34 inbox
Zhiyuan Fang, Yuegui Huang, Zicong Hong, Yufeng Lyu, Wuhui Chen, Yue Yu, Fan Yu, Zibin Zheng
crossref: 2026-09-12dblp: 2026-09-12arxiv: 2026-09-12

AI 速览

一句话:提出专家感知的多批次流水线调度与预取机制,消除 MoE 卸载推理中的流水线气泡,吞吐最高提升 85.12 倍。

混合专家(MoE)以稀疏结构支撑万亿参数语言模型而不显著增加计算开销,但庞大的参数量使推理受限于 GPU 显存,需借助 CPU 与磁盘卸载并并行化 I/O 与计算。由于单个专家的计算时间往往小于其加载时间,流水线中产生大量气泡。本文提出 Klotski 推理引擎,采用专家感知的多批次流水线范式:通过批处理延长当前层的计算时间,使其与下一层的加载时间重叠。针对更多批次会激活更多专家、从而延长加载时间并加剧气泡的问题,Klotski 依据不同批次下专家的异构计算/IO 需求与激活模式编排推理顺序,以平衡二者时间;并设计约束敏感的 I/O-计算规划器与相关性感知的专家预取器,使调度适配不同硬件环境与模型。实验表明,Klotski 相比现有最优技术取得更优的吞吐-延迟权衡,吞吐提升最高达 85.12 倍。

问题
MoE 模型参数规模巨大而 GPU 显存增长有限,虽可通过卸载到 CPU/磁盘并行 I/O 与计算,但专家计算时间常短于加载时间,导致流水线中出现大量气泡。
方法
提出 Klotski 引擎,采用专家感知的多批次流水线范式,按专家的异构计算/IO 需求与激活模式编排推理顺序,并配合约束敏感的 I/O-计算规划器与相关性感知的专家预取器来最小化气泡。
评估
在多种硬件环境与模型上开展实验,与当前最优技术比较吞吐-延迟权衡,吞吐提升最高达 85.12 倍。
原文摘要(Abstract)

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.

链接与引用

DOI 原文 · arXiv · DBLP

BibTeX
@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}
}