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