Token计数批处理：查询嵌入推理速度与成本双优化

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

那么，问题来了。这些查询通常很短，而且它们的长度分布极其不均衡。这就导致了一个核心矛盾：当我们在GPU上逐个处理这些短请求时，计算资源根本“吃不饱”。系统的大部分时间都花在了等待数据、搬运数据这些“杂活”上，真正用于计算的时间少得可怜。换句话说，这时的推理过程是“内存瓶颈型”（memory-bound）的，而不是“计算瓶颈型”（compute-bound）的。更麻烦的是，查询流量就像心跳一样，一阵一阵地来，突发性极强，想靠自动扩缩容来应对根本来不及。如果像传统做法那样一个个顺次处理，效率自然低下。

在这篇文章里，我们来聊聊如何利用“批处理”（batching）技术来优雅地解决这个问题。我们会先探讨现代推理引擎里一个至关重要的技术——填充移除（Padding Removal），正是它让高效的批处理成为可能。接着，我会分享一些实用的批处理策略和如何选择最优的批大小。最后，我们来看一下具体的实现细节和最终效果：尽管GPU的用量减少了3倍，但推理延迟却降低了50%。

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

面对查询流量这种“短而急”的特性，一个很自然的想法冒出来了：既然一个请求喂不饱GPU，那咱们能不能一次多塞几个，搞个“团购”呢？如果把多个短请求捆在一起，形成一个批次（Batch），一起喂给GPU，效率不就上来了吗？

想法很美好，但传统做法有个“潜规则”。大多数推理引擎在处理请求时，会要求你提供一个形状为 (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数）与GPU实际所需完成的计算量对齐。当大量短查询几乎同时到达时，我们根据它们的token总数进行分组，让GPU一次性处理更大的有效负载，从而摊薄了每次处理的固定开销。试验数据表明，这种方法能有效降低单个请求的延迟和成本，并显著提升吞吐量和模型利用率（MFU）。

最优的批大小是多少？

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

这意味着，对于小的请求来说，固定开销（如GPU调度、内存搬运、最后的池化和归一化等）占据了主导地位，所以我们看到的延迟几乎恒定。而当总token数超过这个“饱和点”（Saturation Point）后，延迟就开始线性增长了。对于我们的voyage-3模型在A100上的表现来说，这个饱和点大约是600个token。

那么，这个饱和点就是我们梦寐以求的最优批大小。我们在饱和点处，能在延迟不显著增加的前提下，最大化MFU和吞吐量。这就像开车一样，在保持舒适（低延迟）的前提下，尽量把车速提到最高（高吞吐量）。

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

实现这个方案，我们需要一个更智能的数据中间件，它不能只是个简单的FIFO（先进先出）队列。这个系统必须具备以下能力：首先，能预估每个请求的token数量；其次，能“头盔”队列中待处理的所有请求；最后，能原子性地“抓取”一组请求，使得它们的总token数恰好接近我们的最优批大小。

像RabbitMQ、Kafka这些通用的消息中间件，虽然在其他方面很强大，但在这种场景下显得有些“水土不服”。它们的批处理调优参数大多是消息条数或字节数，而不是token数。我们很难在Kafka或RabbitMQ里优雅地实现“凑满600个token就打包”的逻辑。

因此，我们有两条实际可行的路径。一是，在Kafka/RabbitMQ前面放一个轻量级的聚合器，由它来负责按token数消费和打包，然后再发送给模型服务器。二就是，直接用像Redis这样的存储系统，它天然支持快速的“窥探”和条件批处理操作。在我们的实现中，我们选择了后者，因为可以用Lua脚本在Redis里原子性地执行“弹出一个批次的数据，直到达到最优批大小”，并且还能设置每个请求的TTL（生存时间）。

我们系统的流程是这样的：每个嵌入查询请求，都会被放入一个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计数的批处理，是一个被实践证明有效的策略。它不仅显著提升了短查询嵌入推理的吞吐量和延迟表现，更重要的是，它极大地优化了资源利用率，直接降低了运营成本。这才是工程师们真正应该追求的“优雅”方案。

来源：互联网