首页 > AI教程 >基于Token计数的批处理：更快更便宜的查询嵌入推理

基于Token计数的批处理：更快更便宜的查询嵌入推理

来源：互联网 2026-07-16 06:15:07

针对短文本嵌入推理的内存瓶颈问题，提出基于Token计数的批处理策略，结合填充移除技术，按总token数而非请求数量分组，有效对齐GPU计算量。实验表明，GPU用量减少3倍，推理延迟降低50%，吞吐量最高提升8倍，显著优化资源利用率与成本。

在搜索引擎、RAG（检索增强生成）系统或推荐系统的底层，一个常见场景是需要处理大量短文本的嵌入推理请求。在MongoDB Voyage AI，这类短请求被称为“queries”（查询），其他类型的请求则称为“documents”（文档）。查询请求有一个硬性要求：延迟必须非常低，通常在100-300毫秒以内。基于Token计数的批处理：更快更便宜的查询嵌入推理

问题的关键在于，这些查询通常很短，且长度分布极不均衡。这导致一个核心矛盾：当在GPU上逐个处理这些短请求时，计算资源无法被充分利用。系统大部分时间花费在数据等待和搬运等“杂活”上，真正用于计算的时间很少。也就是说，此时的推理过程属于“内存瓶颈型”（memory-bound），而非“计算瓶颈型”（compute-bound）。更棘手的是，查询流量的突发性很强，像心跳一样一阵阵到来，自动扩缩容根本来不及应对。如果按照传统方式一个个顺序处理，效率自然低下。本文探讨如何利用“批处理”（batching）技术优雅解决这一问题。首先介绍现代推理引擎中的关键技术——填充移除（Padding Removal），正是它让高效的批处理成为可能。接着分享实用的批处理策略以及如何选择最优的批大小。最后展示具体的实现细节和最终效果：尽管GPU用量减少了3倍，但推理延迟却降低了50%。

填充移除：让批处理变高效的魔法

面对查询流量“短而急”的特性，一个自然的想法是：既然单个请求无法喂饱GPU，不妨一次多塞几个，形成批次（Batch）一同处理，从而提高效率。然而传统做法有一个“潜规则”：大多数推理引擎要求输入张量的形状为 (B, S)，其中B是批次大小，S是批次中最长的那个序列长度。所有比最长序列短的请求，都需要在其后补上“填充符”（Paddings），使它们长度一致，以便GPU进行并行计算。这个“填充”操作正是效率的敌人。填充符不含任何有效信息，却消耗计算和内存带宽。这意味着，如果处理100个短请求，每个只有10个有效token，但有一个请求是100个token，那么所有请求都必须被“充气”到100个token长度，总共消耗100×100=10000个token的计算量。而理想状态下，只需处理10×100=1000个有效token。这一巨大浪费正是延迟居高不下的主要原因。解决方案是“填充移除”（Padding Removal）和“变长处理”（Variable-length Processing）。其思路是：不再将所有请求硬撑到同一长度，而是将它们首尾相连，拼接成一个长度为 T = Σtoken_count_i 的“超级长序列”。像vLLM、SGLang这类现代推理引擎能优雅地处理这个拼接后的长序列。通过精心设计的注意力掩码（Attention Masks）和位置索引（Position Indices），引擎确保计算时序列中的每个片段只与自己原本的“邻居”交互，互不干扰。这样一来，推理时间只与总有效token数 T 挂钩，不再浪费在 B × S 的虚高数字上。

核心方案：基于Token计数的批处理

正是基于这一技术，Voyage AI提出并实现了“基于Token计数的批处理”（Token-count-based Batching）。核心思路很简单：不再按请求数量来凑批，而是按批内所有请求的总token数来凑批。 相比之下，传统的“时间窗口批处理”问题很多：窗口太短，只能捕获两三个请求，批次太小，浪费了GPU的每次启动开销；窗口太长，等来的请求太多，虽提高了GPU利用率，但排队等待时间又拉长了延迟。而“请求数量批处理”也有类似短板。面对突发流量，这两种方式都很难找到平衡，要么“吃不饱”，要么“吃撑了”。基于Token计数的批处理：更快更便宜的查询嵌入推理

Token计数批处理的好处显而易见：它将批大小（即总token数）与GPU实际所需完成的计算量对齐。当大量短查询几乎同时到达时，根据它们的token总数进行分组，让GPU一次性处理更大的有效负载，从而摊薄每次处理的固定开销。试验数据表明，这种方法能有效降低单个请求的延迟和成本，并显著提升吞吐量和模型利用率（MFU）。

最优的批大小是多少？

任何优化方案都不能拍脑袋。需要回答一个关键问题：到底多大的批（总Token数）才是最合适的？对自己模型的推理延迟进行细致的性能分析后，结果揭示了一个清晰的模式：在某个阈值以下，延迟几乎是一条平缓的直线。 基于Token计数的批处理：更快更便宜的查询嵌入推理

这意味着，对于较小的请求，固定开销（如GPU调度、内存搬运、最后的池化和归一化等）占据主导地位，因此延迟几乎恒定。当总token数超过这个“饱和点”（Saturation Point）后，延迟开始线性增长。对于voyage-3模型在A100上的表现，饱和点大约是600个token。那么，饱和点就是最优批大小。在饱和点处，能在延迟不显著增加的前提下，最大化MFU和吞吐量。就像开车，在保持舒适（低延迟）的前提下，尽量把车速提到最高（高吞吐量）。基于Token计数的批处理：更快更便宜的查询嵌入推理

队列设计：为Token计数批处理而生的“智能调度器”

实现这一方案需要一个更智能的数据中间件，不能只是简单的FIFO（先进先出）队列。该系统必须具备以下能力：首先，能预估每个请求的token数量；其次，能“查看”队列中待处理的所有请求；最后，能原子性地“抓取”一组请求，使得它们的总token数恰好接近最优批大小。像RabbitMQ、Kafka这类通用消息中间件，在其他方面很强大，但在这种场景下显得有些“水土不服”。它们的批处理调优参数大多是消息条数或字节数，而非token数。很难在Kafka或RabbitMQ中优雅实现“凑满600个token就打包”的逻辑。因此有两条实际可行的路径。一是，在Kafka/RabbitMQ前面放一个轻量级聚合器，由它负责按token数消费和打包，再发送给模型服务器。二是，直接用像Redis这样的存储系统，它天然支持快速的“窥探”和条件批处理操作。在实现中，选择了后者，因为可以用Lua脚本在Redis里原子性地执行“弹出一个批次的数据，直到达到最优批大小”，并且还能设置每个请求的TTL（生存时间）。基于Token计数的批处理：更快更便宜的查询嵌入推理

系统流程如下：每个嵌入查询请求被放入一个Redis列表中。然后，模型服务器调用一个Lua脚本，由脚本负责原子性地从列表中取出请求，直到累加的token数达到最优批大小。Redis本身数据丢失的概率极低，万一发生，用户只会收到503错误，重试即可。

效果：一石三鸟，延迟、吞吐、成本全赢了

理论讲得再好，不如看看实战效果。在Voyage-3-Large模型的生产环境中，用新方案（查询批处理 + vLLM）和旧方案（无批处理 + Hugging Face推理）进行了一次A/B测试。结果令人振奋：尽管GPU用量减少了3倍，但推理延迟却降低了50%。 随后将这套基于查询批处理的方案推广到了7个模型，并观察到以下显著变化： - vLLM的引入，使大部分模型的GPU推理时间缩减了约20毫秒。 - GPU利用率和MFU显著提高，反映了填充的减少、每批固定开销被更好摊薄，以及推理过程从“内存瓶颈型”向“计算瓶颈型”的转变。 - 通过基于Token计数的批处理，系统吞吐量提升了**最高8倍**。 - 在资源争抢严重的情况下，一些模型服务器的P90端到端延迟降低了整整60毫秒。 - 即使使用更少的GPU，P90端到端延迟在流量高峰时也表现得更加稳定。总而言之，结合填充移除和基于Token计数的批处理，是一个被实践证明有效的策略。它不仅显著提升了短查询嵌入推理的吞吐量和延迟表现，更重要的是，极大优化了资源利用率，直接降低了运营成本。

侠游戏发布此文仅为了传递信息，不代表侠游戏网站认同其观点或证实其描述