Speculative MoE: Communication Efficient Parallel MoE Inference with Speculative Token and Expert Pre-scheduling
作者: Yan Li, Pengfei Zheng, Shuang Chen, Zewei Xu, Yuanhao Lai, Yunfei Du, Zhengang Wang
发布时间: 2025-03-10
来源: arxiv
研究方向: 大规模语言模型（LLM）的并行推理优化
主要内容
本文研究了大规模语言模型（LLM）中MoE（Mixture of Experts）架构的并行推理优化，针对现有MoE推理框架在通信开销方面的瓶颈，提出了一种名为Speculative MoE（s-MoE）的优化方法。该方法通过预测并预调度token和expert，有效地减少了通信开销，从而提升了MoE推理的效率。
主要贡献
1. 提出了一种名为Speculative MoE（s-MoE）的优化方法，通过预测并预调度token和expert，有效地减少了通信开销。
2. s-MoE包含两个关键机制：Speculative Token Reshuffling（s-TS）和Speculative Expert Pre-grouping（s-EG），分别用于预测token的expert路由路径和预聚类expert。
3. s-MoE在DeepSpeed-MoE和SGLang等现有MoE推理框架上进行了实现，并取得了显著的性能提升。
4. 实验结果表明，s-MoE在不同数据集、模型和硬件上，相比于DeepSpeed-MoE，在TTFT、TPOT和p90-TBT延迟约束下，吞吐量分别提升了1.58x-2.34x、1.04x-2.34x和1.37x-5.98x。
研究方法
1. Speculative Token Reshuffling（s-TS）
2. Speculative Expert Pre-grouping（s-EG）
3. 概率模型预测token-expert路由路径
4. 平衡token-expert共聚类
5. 高效的内核实现、通信合并、调度去重、缓存查找表等系统级优化
实验结果
实验结果表明，s-MoE在不同数据集、模型和硬件上，相比于DeepSpeed-MoE，在TTFT、TPOT和p90-TBT延迟约束下，吞吐量分别提升了1.58x-2.34x、1.04x-2.34x和1.37x-5.98x。在DeepSeek-V2和Mixtral-8x7B等典型MoE模型上，s-MoE均取得了显著的性能提升。
未来工作
未来工作将包括将s-MoE扩展到更多LLM推理框架，以及进一步优化s-MoE的预测模型和聚类算法，以进一步提升性能。