GPU嵌入推理与批处理：AI搜索服务性能优化实践

最近在折腾一套语义搜索服务的时候我又把 Perplexity 的公开技术分享翻出来看了一遍。Perplexity 这个产品看起来是个“搜索框 大模型回答”但真正撑住它每天海量查询的是背后那条大规模搜索结果服务链路。而这条链路里最容易被低估、又最影响成本与响应速度的恰恰是 GPU 嵌入推理和批处理这两个环节。我之前做搜索系统改造时第一版嵌入模型直接跑在 CPU 上上线初期 QPS 不高还看不出问题等流量一上来P99 直接从 80ms 飙到 800ms整条检索链路全卡在向量化这一步。后来把嵌入推理迁到 GPU并且重写了批处理逻辑吞吐才真正上去。这篇文章我就顺着 Perplexity 这类 AI 搜索服务的技术路线把 GPU 嵌入推理、批处理、显存规划、性能调优这些事一次讲透。1. 大规模搜索服务为什么对嵌入推理如此敏感1.1 从关键词匹配到语义匹配嵌入推理成为搜索链路的前置依赖传统搜索引擎的核心是倒排索引你输入“GPU 推理”系统按词匹配返回包含这几个词的网页。这套方案成熟、快但有个天然缺陷——它不理解语义。搜“显卡计算加速”的时候它不知道你其实想找的是 GPU 相关的内容。AI 搜索产品之所以叫 AI 搜索本质上是把“文本匹配”换成了“向量匹配”。用户 query 进来之后第一步不是查倒排索引而是先把 query 送入嵌入模型变成一个高维向量再拿这个向量去向量数据库里做相似度检索。这就带来一个结构性变化嵌入推理从“可选的后处理”变成了“搜索链路的前置依赖”。一旦这一步阻塞后面的大模型生成、重排序全部无从谈起。Perplexity 的做法本质上也是这套路只不过规模更大、要求更苛刻。它每天的查询量级决定了嵌入推理不可能用 CPU 慢慢磨必须是 GPU 集群来扛。而且嵌入推理还出现在多个环节不是只有 query 需要向量化候选文档、知识库切片、历史对话上下文都可能需要实时向量化。链路越长GPU 嵌入推理的压力就越大。1.2 查询流量特性为什么搜索场景的推理负载并不均匀很多人以为搜索的流量曲线是平滑的实际完全不是。真实搜索流量的特点是白天高、夜间低热点事件到来时瞬间洪峰。更麻烦的是query 的语义分布极不均衡——高频 query 可能占总量的 30%而这些高频 query 的文本长度、语言分布都有明显偏向。这对嵌入推理意味着什么意味着 GPU 资源的分配不能按平均值来算。如果按峰值流量准备 GPU平时就是巨大的浪费如果按均值准备热点一来直接打爆。Perplexity 这类服务之所以强调批处理本质上就是因为流量波动越大批处理能带来的吞吐增益就越明显。还有一个容易忽略的点query 文本长度差异极大。有人搜“GPU”两个 token 完事有人会粘贴一整段话进来。嵌入模型对输入长度是敏感的不同长度的样本混合在一起GPU 的计算效率完全不同。这就引出了后面要讲的批处理策略——好的批处理不是简单把请求堆在一起而是要考虑长度分布、显存占用、计算密度的匹配。2. GPU 嵌入推理服务的设计模型、显存与并发模型2.1 嵌入模型的选择与显存预算计算做嵌入服务第一个绕不开的问题是模型选型。Perplexity 早期公开的资料里提到过它用多个模型做检索和重排具体细节不一定完全公开但整个行业用的无非是几个方向开源的有 BGE、E5、GTE、text-embedding 系列闭源的有 OpenAI、Cohere 的 API。自建 GPU 嵌入服务模型选择主要看三个指标维度、最大输入长度、推理时延。维度直接决定向量数据库的存储成本和检索带宽。768 维和 3072 维存储差四倍检索耗时也可能差一倍。最大输入长度决定你能不能处理长文档。很多搜索场景需要把整篇网页切片后向量化模型最长只能吃 512 token 的话切片策略就要跟着调整。推理时延决定单卡吞吐。测试时要看的是“batch size 拉满之后的平均时延”而不是单条数据的处理速度。显存预算这块有一个很实用的估算公式。以 BGE-large 这类模型为例参数量大约 335MFP16 精度下模型权重占用约 670MB。但这只是“模型本身”真正跑推理时还有几块开销开销项说明典型占比模型权重FP16 权重的显存占用约 1/4激活值前向计算时每层产生的中间结果随 batch 增长CUDA context框架初始化占用的基础显存约 500MB - 1GB推理引擎缓存TensorRT、ONNX Runtime 的临时缓存不定实际上一张 24GB 的 4090 或者 A5000跑 BGE-large 这种规模的模型batch size 开到 64 左右比较稳。如果你用的是更大的模型比如 7B 参数的嵌入模型现在确实有人这么干显存直接按 FP16 权重 14GB 起步算基本就是一张 A10 或者 A100 才能带得动。2.2 单卡起步还是多卡并行真实业务的决策路径很多团队一上来就想着多卡并行我的建议是能单卡先单卡。嵌入推理和 LLM 推理不一样它不需要 KV Cache也不存在自回归的串行依赖单卡的并发吞吐其实非常可观。以 BGE-large-zh 为例A100 上单卡跑动态批处理实测吞吐做到几千 QPS 是常态。对一个日请求量百万级的搜索服务来说一张卡绰绰有余。真正的瓶颈往往不在单卡算力而在下游的向量检索——你向量化再快向量数据库扛不住也没用。那什么时候必须上多卡两种情况一是单卡吞吐确实不够而且已经确认瓶颈在 embedding 服务本身二是需要多副本做高可用单点故障会直接导致搜索服务不可用。Perplexity 这种体量肯定是要多副本的但它也不会为了追求“看起来分布式”而盲目上多卡。多卡并行还有一个容易踩坑的地方——数据并行 vs 模型并行。嵌入模型体量通常不大正常用数据并行就够了也就是每张卡放一份完整模型请求按一定策略分摊到不同卡上。千万别一上来就搞张量并行、流水线并行那是给大模型用的嵌入模型这样做纯属徒增通信开销。3. 批处理机制把 GPU 吞吐拉到极限的关键3.1 嵌入推理批处理的数学基础GPU 和 CPU 的本质区别在于CPU 适合做“少量复杂任务”GPU 适合做“大量简单并行计算”。嵌入模型的本质是 Transformer 的前向计算包含大量矩阵乘法天然适合 GPU。但如果你一个一个请求地喂给 GPU每个请求之间的 GPU 计算单元大量空闲利用率可能只有 5% 到 10%。批处理解决的就是这个问题。一次前向计算同时处理 N 个样本矩阵乘法从 [1, D] 变成 [N, D]GPU 的计算单元被充分占满单样本的平均计算成本大幅下降。这里有一个关键参数吞吐与 batch size 的关系不是线性的而是近似收益递减。batch size 从 1 提到 8吞吐提升非常明显可能直接翻 4 倍。从 8 提到 32还能再涨一截但幅度变小。从 32 提到 64提升已经很有限反而显存占用和单请求延迟在增加。所以实际调优时不是无脑把 batch size 调大而是要找到一个“吞吐平缓区”。这个平缓区因模型和 GPU 而异需要压测来确定。我一般在 batch size 为 16、32、64、128 时各跑一轮压测画一张吞吐曲线从中挑拐点位置。3.2 动态批处理队列实现细节与参数调优生产环境的嵌入服务不能像离线任务那样手动指定 batch size因为线上请求是不断到达的、稀疏且不规律的。如果每次来一个请求就推理一次GPU 浪费如果非要攒够 64 个请求才开始推理那第一个请求的等待时间可能长达数秒完全不可接受。正确做法是动态批处理dynamic batching。核心机制如下请求进入队列后设置一个最大等待窗口比如 10ms 或 20ms。在这个窗口内新到的请求不断加入当前批次。一旦窗口到期或者达到最大 batch size立即把当前批次送入 GPU 执行。窗口大小是个值得细调的参数。窗口太短batch 攒不满吞吐上不去窗口太长尾延迟爆炸。Perplexity 这类大规模服务的经验是在线服务的默认等待窗口往往取 5-20ms 之间具体取决于 QPS 和可接受的整体延迟。有意思的是动态批处理在高 QPS 和低 QPS 场景下的表现逻辑完全不同。高 QPS 时队列里几乎总有积压请求窗口还没到期批次就满了等待时间接近于零。低 QPS 时每个请求可能都要等满整个窗口这时候就得权衡是让单条请求多等 10ms 换取更高的吞吐还是直接牺牲吞吐保证低延迟。我的经验是在线搜索服务优先保障延迟低峰期宁可让 batch 小一点也别让用户多等。另外动态批处理还需要考虑长度分桶length bucketing。嵌入模型是定长 padding 的如果 batch 里一条 500 token 的长文本和一堆 10 token 的短文本混在一起所有短文本都要被 padding 到 500 token计算量白白浪费一大截。实测中这种 padding 浪费可能让吞吐缩水 30% 以上。处理方式是按输入长度分桶把长度相近的请求放在一个 batch 里短文本和长文本走不同批次虽然实现复杂度高一点收益却非常直接。4. 性能调优实录一次完整的接入压测与回归4.1 压测方案与关键指标纸上谈兵没用。我把自己最近做的一次嵌入服务压测过程完整贴出来大家可以照着这个思路去验证自己系统。压测环境配置项参数GPUNVIDIA A1024GB模型BGE-large-zh约 335M 参数FP16推理框架ONNX Runtime CUDA EP动态批处理窗口10ms压测工具自研 Python 异步客户端 Locust输入样本平均长度 128 token最大 512 token关键指标就两个平均吞吐QPS和 P99 延迟。给每个指标设了目标线目标吞吐 2000 QPSP99 不超过 100ms。压测时我特意分了三轮第一轮 batch size 固定为 8不开动态批处理作为基线。第二轮固定 batch size 32同样不开动态批处理。第三轮开启动态批处理窗口 10msbatch 上限 64。4.2 实测结果与参数调整三轮结果如下场景平均吞吐P99 延迟batch8无动态780 QPS23msbatch32无动态1520 QPS38ms动态批处理上限 642310 QPS71ms看到数据我的第一反应是“动态批处理的 P99 怎么比固定 batch 高了这么多”。细看日志才发现问题不在批处理本身而在于动态批处理引入了排队延迟——高峰期排队的请求比固定 batch 场景多尾部请求等得更久。这里有一个非常重要的调优原则追求吞吐不能以牺牲 P99 为代价。哪怕平均延迟很好看只要 P99 破线用户侧的体感就是“卡顿”。之后的调整方向有两个把动态批处理窗口从 10ms 改到 5ms代价是吞吐降了约 8%P99 从 71ms 降到 49ms。给队列加了一个最大长度限制超出后对请求执行“降级处理”——直接单样本推理不进批处理队列。这样洪峰时个别请求会牺牲一点计算效率但延迟可控。最终我们把参数定为窗口 5msbatch 上限 48队列最大积压 256。上线后高峰期吞吐稳定在 2100 QPS 左右P99 约 55ms完全满足业务要求。这个压测过程给我们的教训是不要只看平均吞吐P99 和尾延迟才是搜索服务的生命线。动态批处理参数不是一个固定值你必须要根据实际流量曲线反复压测找到吞吐和延迟都满足要求的平衡点。5. 经验总结与常见坑GPU假忙、批处理失效、显存估算错误5.1 GPU 利用率很低但请求就是慢排查路径嵌入服务上线一段时间后我们遇到过一个诡异问题。监控面板上 GPU 利用率只有 7%但接口平均延迟却高达 300ms。按常理想GPU 应该很闲才对怎么会慢成这样排查链路是这样的第一步看 CPU 负载。发现 CPU 跑满了但 CPU 是给预处理用的数据很快就交给 GPU 了。第二步看 GPU 的显存占用和计算利用率。显存占用很稳定但 SMI 里显示“Compute”和“Memory”都不忙。第三步看推理框架日志。发现每一批数据进入 GPU 之前都要等 CPU 端的数据预处理做 padding、tokenize、转 tensor——这一步是单线程的完全成了瓶颈。问题根源是数据预处理 pipeline 没有和推理解耦。传统的同步流程是“接请求 → 预处理 → 推理 → 返回”遇到慢的预处理GPU 就空转等待。解决办法是引入生产者-消费者模型预处理线程和推理线程分离预处理完的数据放进缓冲区推理线程拿到数据立即执行。这样 GPU 利用率直接涨到 60% 以上延迟也降了一半。这个坑非常隐蔽如果你的嵌入服务 GPU 利用率低先别怀疑 GPU 性能排查一下 CPU 预处理是不是没跟上。5.2 批处理在长尾查询下的退化问题动态批处理也不是万能的。我在低 QPS 时段观察到一个现象batch 里经常只有 1-2 条请求批处理几乎失效GPU 利用率和没开批处理时一样低。这本质上是一个流量密度问题。低峰期 QPS 只有几十动态批处理窗口内攒不到足够多的请求批处理自然没法生效。Perplexity 这种全球性产品因为流量足够大可能对这个问题感受不明显但做垂直领域搜索服务的团队大概率会碰到。我的建议是提前做好弹性伸缩策略低峰期把服务缩容到单卡只保证最低可用性高峰期再扩容。另一种优化思路是引入“延迟容忍队列”——非实时性的请求比如离线文档入库的向量化任务可以在低峰期复用 GPU 算力把空闲窗口填满。这个方案我们在实际中也验证过高峰期不抢资源低峰期能提升约 40% 的整体吞吐。5.3 显存估算错误导致线上 OOM最后一个常见的坑是显存估算误差。很多人按“模型权重 激活值”算显存但忽略了推理引擎本身的显存开销。有一次我们把嵌入模型从 PyTorch 换成 TensorRT 加速想着 TRT 优化后显存应该更省结果一部署就 OOM。查监控发现 TRT 的显存不是一次性分配的而是根据 batch size 动态增长。我们设置的 batch 上限是 128整个推理引擎预留了超大显存块和系统中的其他 CUDA context 撞在一起直接打爆显存。解决方式是两层一是限制 batch 上限到 64换回 ONNX Runtime二是在推理引擎初始化时设置 GPU memory pool 的 preallocation 比例不要让框架一次性把显存全部占了。这个案例想说明的是——显存估算不能只看模型本身要看推理框架的完整显存占用曲线。任何一个框架在初始化、预热、峰值 batch 三个阶段的显存占用都可能完全不同。生产环境上线前务必用最大 batch size 跑一次真实的压测而不是纸上算一算就拍板上线。搜索引擎里嵌入推理这件事看起来只是一个小环节实际上直接决定了整体链路的质量和成本。Perplexity 能支撑大规模搜索结果服务GPU 嵌入推理和批处理功不可没。希望这篇从原理到实践的拆解能帮你在做自己搜索服务时少走一些弯路把钱花在刀刃上。