单卡高效运行MoE模型,吞吐量提升8-31倍
MoE-Gen: High-Throughput MoE Inference on a Single GPU with Module-Based Batching
- 按模块分批处理,动态调度大批次提升显存利用率
- 优化各模块批大小,实现计算与通信充分重叠
- 适合大规模离线推理,尤其对DeepSeek、Mixtral等模型效果显著
本文提出MoE-Gen,一种面向单卡部署的高吞吐MoE推理系统。现有系统依赖模型级或连续分批策略,原为交互式推理设计,导致MoE的核心模块(注意力与专家模块)批大小过小,严重降低吞吐。为此,我们引入模块级分批:在主机内存累积令牌,动态向GPU批量推送大批次以最大化利用。同时,针对MoE中每个模块优化批大小,使GPU计算与通信完全重叠,进一步提升吞吐。实验表明,相较于采用模型级分批的先进系统(FlexGen、MoE-Lightning、DeepSpeed),MoE-Gen在离线推理任务中实现8-31倍吞吐提升;相比连续分批系统(如vLLM、Ollama),在DeepSeek和Mixtral等主流MoE模型上亦有更显著优势。代码已开源:https://github.com/EfficientMoE/MoE-Gen
原文摘要 · Abstract (English)
This paper presents MoE-Gen, a high-throughput MoE inference system optimized for single-GPU execution. Existing inference systems rely on model-based or continuous batching strategies, originally designed for interactive inference, which result in excessively small batches for MoE's key modules-attention and expert modules-leading to poor throughput. To address this, we introduce module-based batching, which accumulates tokens in host memory and dynamically launches large batches on GPUs to maximize utilization. Additionally, we optimize the choice of batch sizes for each module in an MoE to fully overlap GPU computation and communication, maximizing throughput. Evaluation demonstrates that MoE-Gen achieves 8-31x higher throughput compared to state-of-the-art systems employing model-based batching (FlexGen, MoE-Lightning, DeepSpeed), and offers even greater throughput improvements over continuous batching systems (e.g., vLLM and Ollama) on popular MoE models (DeepSeek and Mixtral) across offline inference tasks. MoE-Gen's source code is publicly available at https://github.com/EfficientMoE/MoE-Gen
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。