在服务大量短请求时，嵌入模型推理往往难以兼顾效率——这是搜索、检索和推荐系统中的常见模式。在 MongoDB 旗下的 Voyage AI，我们将这些短请求称为查询（queries），其他请求则称为文档（documents）。查询通常必须以极低延迟（通常为 100–300 毫秒）提供服务。

查询通常较短，其 token 长度分布高度偏斜。因此，查询推理往往受内存限制而非计算限制。查询流量波动很大，因此自动扩缩容过于缓慢。总之，顺序服务大量短请求的效率极低。

在这篇博客文章中，我们探讨如何利用批处理（batching）来更高效地服务查询。我们首先讨论现代推理引擎中的填充移除（padding removal），这是一项实现有效批处理的关键技术。然后，我们介绍形成批次和选择合适批次大小的实用策略。最后，我们详细介绍实现细节并分享由此带来的性能提升：GPU（图形处理器） 推理延迟降低 50%——尽管使用的 GPU 数量减少了 3 倍。

填充移除使有效批处理成为可能

鉴于查询流量的模式，一个直接的想法是：我们能否对它们进行批处理以提高推理效率？vLLM 和 SGLang 等推理引擎支持的填充移除使高效批处理成为可能。

大多数推理引擎接受形式为 (B, S) 的请求，其中 B 是批次中的序列数量，S 是最大序列长度。序列应填充至最大序列长度，以便张量对齐。但这种便利是有代价的：填充 token 不执行任何有用工作，却仍然消耗计算和内存带宽，因此延迟随 B × S 而非实际 token 数量增长。在服务大量短请求时，这会浪费大量计算资源，并可能推高尾部延迟。填充移除和变长处理通过将所有活跃序列拼接成一个长度为 T = Σtoken_count_i 的长“超级序列”来解决这个问题，其中 token_count_i 是序列 i 中的 token 数量。vLLM 和 SGLang 等推理引擎可以处理这个组合序列。注意力掩码和位置索引确保每个序列只关注自身的 token。因此，推理时间现在取决于 T 而非 B × S，使 GPU 工作与实际需求对齐。

方案：基于 token 计数的批处理

在 Voyage AI，我们提出并构建了基于 token 计数的批处理，即按批次中的总 token 数（Σtoken_count_i）对查询（短请求）进行批处理，而非按总请求数或任意时间窗口。

在服务大量短请求时，时间窗口批处理效率低下。时间窗口批处理会根据流量在批次填充不足和过度填充之间摇摆。短窗口保持低延迟，但产生填充不足的小批次；长窗口提高利用率，但增加排队延迟。流量具有突发性，因此单一窗口大小会在填充不足和过度填充之间振荡。时间窗口批处理引入了资源利用率的可变性，导致系统在内存受限和计算受限操作之间切换。请求计数批处理也有类似问题。

Token 计数批处理使批次大小（批次中的总 token 数）与实际所需计算量对齐。当许多查询几乎同时到达时，我们按 token 计数将它们分组，使 GPU 在单次前向传播中处理更大的组合工作负载。根据我们的实验，基于 token 计数的批处理摊销了固定成本，降低了每请求延迟和成本，并提高了吞吐量和模型浮点运算利用率（MFU）。

最优批次大小是多少？

我们对查询推理的推理延迟与 token 计数分析显示出清晰的模式：延迟在达到某个阈值（饱和点）之前大致持平，之后变为近似线性。对于小请求，固定的每请求开销（如 GPU 调度、内存移动、池化和归一化等）占主导地位，延迟几乎保持不变；超过该点后，延迟随 token 计数增长。阈值（饱和点）取决于模型架构、推理引擎和 GPU 等因素。对于我们在 A100 上运行的 voyage-3 模型，该阈值约为 600 个 token。

基于推理延迟与 token 数量的数据，我们可以分析 FLOPs 利用率（MFU）与 token 数量的关系，以及吞吐量与 token 数量的关系，如下图所示。我们观察到，模型 FLOPs 利用率（MFU）和吞吐量随 token 数量近似线性增长，直至达到饱和点。我们的大多数查询推理都处于内存受限区域，远未达到饱和点。

对短请求进行批处理可以将推理从内存受限转变为计算受限。如果我们选择图 3 中的饱和点作为批大小（批次中的总 token 数量），则可以在延迟和吞吐量/MFU 之间实现平衡与优化。

队列设计：实现基于 token 数量的批处理

基于 token 数量的批处理需要一个不仅仅是简单 FIFO 投递的数据系统。该系统必须为每个请求附加一个估算的 token_count，能够跨待处理请求进行窥探，然后原子性地认领一个子集，使其总 token 数符合最优批大小（ Σtoken_count_i ≤ optimal_batch_size ）。如果没有这些原语，我们要么填充不足 GPU——浪费固定开销——要么填充过度并导致尾部延迟飙升。

像 RabbitMQ 和 Kafka 这样的通用消息代理在持久性、扇出和投递方面表现出色，但它们的批处理调节参数是消息数量/字节数，而非 token 数。RabbitMQ 的预取是基于请求数量的，且消息是推送给消费者的，因此没有高效的方式按 Σtoken_count_i 来窥探和批处理请求。Kafka 在分区内按字节/消息进行批处理；token 数量随文本和分词器而变化，因此没有高效的方式按 Σtoken_count_i 来批处理请求。

因此，有两种切实可行的路径来实现基于 token 数量的批处理。一种是在 Kafka/RabbitMQ 前面放置一个轻量级聚合器，按 token 数量消费批次，然后将批次分发给模型服务器。另一种是使用天然支持快速窥探 + 条件批处理的存储——例如，配合 Lua 脚本的 Redis。在我们的实现中，我们使用 Redis，因为它允许我们在单次 Lua 脚本调用中原子性地“弹出至最优批大小”并设置每项的 TTL。无论我们选择哪种方式，核心要求都是一样的：队列必须让我们的系统能够看到多个待处理项，按 Σtoken_count_i 进行批处理，并原子性地认领它们，以保持利用率稳定和延迟可预测。

我们的系统将每个嵌入查询请求入队到 Redis 列表中，格式如下：

模型服务器原子性地调用 Lua 脚本以获取一批请求，直到达到最优批大小。Redis 丢失数据的概率非常低。在极少数情况下如果确实发生，用户可能会收到 503 Service Unavailable 错误，只需重试即可。当 QPS 较低时，批次仅被部分填充，GPU 利用率仍然较低，但延迟仍有改善。

结果

我们在 Voyage-3-Large 模型服务上运行了一次生产实验，将我们的新流水线（查询批处理 + vLLM）与旧流水线（无批处理 + Hugging Face Inference）进行对比。我们发现 GPU 推理延迟降低了 50%——尽管使用的 GPU 数量减少了 3 倍。

我们逐步将 7 个以上的模型接入上述查询批处理方案，并看到了以下结果（请注意，这些结果基于我们对“新”和“旧”流水线的特定实现，不一定具有普适性）：

• vLLM 为我们大多数模型将 GPU 推理时间减少了最多约 20 ms。
• GPU 利用率和 MFU 提高，反映出填充减少、每批次开销摊销更好，以及推理更接近计算受限状态。
• 通过基于 token 数量的批处理，吞吐量提升最高达 8 倍。
• 在资源争用下，由于排队时间减少，部分模型服务器的 P90 端到端延迟下降了 60+ ms。
• 即使在 GPU 数量更少的情况下，流量高峰期间 P90 端到端延迟也更加稳定。

总之，将填充移除与基于 token 计数的批处理相结合，可以提高吞吐量并降低延迟，同时在短查询嵌入推理过程中提高资源利用率并降低运营成本。

后续步骤

要了解更多关于 Voyage AI 最先进的嵌入模型和重排序器如何帮助您构建准确、可靠的语义搜索和 AI 应用，请参阅 Voyage AI by MongoDB 页面。

非常感谢 Tengyu Ma、Andrew Whitaker、Ken Hong、Angel Lim 和 Andrew Gaut 在此过程中提供的有益讨论，以及 Akshat Vig、Murat Demirbas 和 Stan Halka 的深思熟虑的审阅和反馈。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：MongoDB Blog

原创地址：https://www.mongodb.com/company/blog/engineering/token-count-based-batching-faster-cheaper-embedding-inference-for-queries