基于令牌计数的分批处理：提升嵌入模型推理效率的实战方案
MongoDB Blog2025/12/18 23:00机翻/自动摘要/自动分类
5 阅读
内容评分
技术含量
8/10
营销水分
4/10
摘要
本文提出基于令牌计数的分批处理方法，用于优化嵌入模型的查询推理效率。通过填充移除技术，将短查询合并为超级序列，减少内存浪费和计算延迟。实验表明，该方法在降低GPU推理延迟、提升吞吐量和资源利用率方面效果显著，尤其适用于高突发性流量场景。
正文
在处理大量短查询时，嵌入模型的推理效率常常受限于内存而非计算能力。MongoDB Voyage AI团队将这类短请求称为‘查询’，而较长的请求称为‘文档’。由于查询通常较短且令牌长度分布不均，传统的分批处理方式（如基于时间窗口或请求数量）效率低下，且导致GPU资源浪费和延迟增加。为此，团队提出了一种基于令牌计数的分批处理方法，通过将所有活跃序列合并为一个‘超级序列’，并利用填充移除技术（如vLLM和SGLang支持）来优化内存和计算资源的使用。这种方法不仅减少了填充令牌带来的额外开销，还使推理时间与实际令牌数量成正比，而非批次大小与最大序列长度的乘积。实验结果显示，该方法显著提升了GPU利用率和吞吐量，同时降低了延迟。例如，GPU推理延迟降低50%，吞吐量提升8倍，且在流量高峰时保持更稳定的延迟表现。