通过推测未来专家减少显存加载延迟,提升MoE模型推理速度
Speculating Experts Accelerates Inference for Mixture-of-Experts
- 用当前模型表示推测下一时刻将用的专家
- 实现14%的输出令牌时间降低,且任务准确率不变
- 适合部署在显存受限的推理场景,如边缘设备
混合专家(MoE)模型通过稀疏激活实现大模型扩容,但在内存受限的推理中,专家权重需从CPU加载到GPU,导致性能瓶颈。本文提出一种专家预取机制,利用当前模型内部表示推测未来将被选中的专家,使内存传输与计算重叠。实验表明,内部表示可可靠预测未来专家。推测执行通常保持下游任务准确率,避免重复加载真实路由选择的专家,从而提升计算-内存重叠效率。集成至优化推理引擎后,相比按需从CPU加载,可实现最高14%的每输出令牌时间(TPOT)降低。对于推测效果不佳的MoE,进一步引入轻量级估计器提高预测命中率,减少性能下降。代码已开源。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) models have gained popularity as a means of scaling the capacity of large language models (LLMs) while maintaining sparse activations and reduced per-token compute. However, in memory-constrained inference settings, expert weights must be offloaded to CPU, creating a performance bottleneck from CPU-GPU transfers during decoding. We propose an expert prefetching scheme that leverages currently computed internal model representations to speculate future experts, enabling memory transfers to overlap with computation. Across multiple MoE architectures, we demonstrate that future experts can be reliably predicted by these internal representations. We also demonstrate that executing speculated experts generally maintains downstream task accuracy, thus preserving more effective compute-memory overlap by eliminating the need to re-fetch true router-selected experts. Integrated into an optimized inference engine, our approach achieves up to 14\% reduction in time per output token (TPOT) over on-demand loading of experts from CPU memory. For MoEs where speculative execution alone yields suboptimal accuracy, we further examine lightweight estimators that improve expert prediction hit rates, thereby reducing performance degradation. Our code is released in open-source at https://github.com/axonn-ai/yalis/tree/offload_prefetch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。