论文

BASE:使用预测移除误差批量感知专家选择以实现高效 MoE 解码

BASE: Batch-Aware Selection of Experts Using Predicted Removal Error for Efficient MoE Decoding

模型推理推理加速

摘要

大语言模型的服务成本越来越高。在大规模服务系统中,自回归解码通常会因将模型权重从加速器高带宽内存传输到片上 SRAM 而遇到瓶颈。专家混合 (MoE) 模型通过仅激活每个词元的一小部分专家来减少计算,但这种稀疏性不会直接转化为批量解码。不同的要求选择不同的专家;因此,许多并发请求的组合活动集可以跨越专家库的很大一部分,并且需要传输更多的专家权重。大多数专家缩减技术独立地为每个词元做出保留决策,因此不能解决这种批量级扩展的问题。最近,批处理感知方法尝试跨并发请求协调专家的使用,并重用已为批处理获取的专家。然而,他们的选择标准主要基于路由器排名或校准期间收集的专家统计数据。因此,这些标准并不直接与因删除专家而导致的输出错误相关,也没有捕获专家的贡献在推理时跨词元的变化。相反,我们根据移除专家对 MoE 层输出的改变程度来对专家进行排名。为了在服务期间应用此标准,我们在校准期间训练一个轻量级线性预测器,该预测器估计每个传入词元的专家移除成本,并开发自定义 GPU 内核以进行成本预测和专家选择。在三种 MoE 架构中,BASE 无需重新训练即可改善质量与效率的权衡。在 Qwen3-30B-A3B 上,在专家预算紧张的情况下,在可比吞吐量的情况下,它的平均准确度比最强基线提高了 29.5 个点。在专家预算较高的情况下,它比密集推理快 60%,同时保持在 0.4 个精度点以内。