Fast MoE Inference via Predictive Prefetching and Expert Replication
作者: Ankit Jyothish, Ali Jannesari, Aishwarya Sarkar, Joseph Zuber
分类: cs.LG
发布日期: 2026-05-12
💡 一句话要点
提出基于预测预取和专家复制的快速MoE推理方法，提升GPU利用率
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 混合专家模型 MoE 推理加速 动态专家复制 GPU利用率 负载均衡 大型语言模型
📋 核心要点
- MoE推理面临GPU利用率低、负载不均衡和高延迟等问题，主要原因是专家激活的稀疏性导致多个token等待同一专家。
- 提出动态专家复制策略，预测过载专家并复制，使复制专家能并发处理token批次，从而提升并行性并减少GPU空闲时间。
- 实验表明，该方法在大型MoE模型上实现了接近100%的GPU利用率，推理速度提升高达3倍，同时保持了接近原始模型的性能。
📝 摘要（中文）
混合专家(MoE)架构已成为最先进的大型语言模型(LLM)中的基本构建块，它提高了LLM中特定领域的专业知识，并在不按比例增加计算开销的情况下扩展了模型容量。然而，MoE推理通常会受到次优GPU利用率、负载不平衡和延迟升高的影响，这些问题源于多个token等待同一专家进行计算，这是由专家激活的稀疏性引起的。为了解决这些挑战，我们提出了一种动态专家复制策略，该策略预测哪些专家可能过载，并为即将到来的token批次复制它们。复制的专家跨层并发处理批次token，从而提高并行性，缩短GPU空闲时间，并显着加快推理速度。在大型MoE模型（包括Switch-base-128和Switch-base-256）上进行的实验评估表明，我们的方法实现了接近完全的GPU利用率（约100%），从而使推理速度提高了3倍，同时保留了基线架构约90-95%的性能。
🔬 方法详解
问题定义：MoE模型推理过程中，由于专家激活的稀疏性，导致负载不均衡，部分专家成为瓶颈，造成GPU利用率低下和推理延迟增加。现有方法难以有效解决专家负载不均衡的问题，导致GPU资源浪费和推理效率降低。
核心思路：通过预测哪些专家可能过载，并动态地复制这些专家，使得多个专家可以并行处理不同的token批次，从而缓解负载不均衡的问题，提高GPU利用率，并最终加速推理过程。核心在于预测和动态复制，以适应不同token批次的专家需求。
技术框架：该方法主要包含两个阶段：预测阶段和复制阶段。在预测阶段，模型会分析当前批次的token，预测哪些专家可能会过载。在复制阶段，根据预测结果，动态地复制这些专家，并将token分配给不同的专家副本进行并行处理。整体流程是：输入token -> 专家预测 -> 专家复制（如果需要）-> token分配 -> 并行计算 -> 输出结果。
关键创新：该方法的核心创新在于动态专家复制策略，它能够根据token的实际需求，自适应地调整专家的数量，从而更好地利用GPU资源。与静态专家分配或简单的负载均衡方法相比，该方法能够更有效地应对专家负载不均衡的问题。
关键设计：具体的预测方法和复制策略是关键设计。预测方法可能基于历史负载信息、token特征或其他相关因素。复制策略需要考虑复制的数量、副本的放置位置以及token的分配方式。此外，还需要考虑复制带来的额外开销，并进行权衡，以确保最终的性能提升。
🖼️ 关键图片
📊 实验亮点
在Switch-base-128和Switch-base-256等大型MoE模型上的实验结果表明，该方法能够实现接近100%的GPU利用率，推理速度提升高达3倍，同时保持了基线架构约90-95%的性能。这些结果表明该方法在提高MoE模型推理效率方面具有显著优势。
🎯 应用场景
该研究成果可广泛应用于各种基于MoE架构的大型语言模型推理加速，例如文本生成、机器翻译、对话系统等。通过提高GPU利用率和降低推理延迟，可以显著提升用户体验，并降低部署成本。该方法还有潜力应用于其他稀疏激活的神经网络模型，具有重要的实际应用价值和未来发展前景。
📄 摘要（原文）
The Mixture of Experts (MoE) architecture has become a fundamental building block in state-of-the-art large language models (LLMs), improving domain-specific expertise in LLMs and scaling model capacity without proportionally increasing their computational overhead. However, MoE inference often suffers from suboptimal GPU utilization, load imbalance, and elevated latency arising from multiple tokens waiting on the same experts for their computation which arises from sparsity of expert activation. To address these challenges, we propose a dynamic expert replication strategy that predicts which experts are likely to be overloaded and replicates them for upcoming batches of tokens. The replicated experts process batch tokens concurrently across layers, which leads to improved parallelism, shorter GPU idle time, and significantly faster inference. Experimental evaluations conducted on large-scale MoE models, including Switch-base-128 and Switch-base-256, demonstrate that our method achieves near-complete GPU utilization (approx 100%), leading to upto 3x improvement in inference speed while preserving approximately 90-95% of the performance of baseline architectures