大规模搜索服务中的GPU嵌入推理与批处理优化实践

做搜索服务的人这两年大概率逃不开一个话题怎么在检索链路里塞入向量召回、怎么把用户查询和候选文档做embedding、怎么在延迟预算内把模型推理跑完。Perplexity这类AI驱动的大规模搜索结果服务本质上就是把传统倒排索引的绿色通道变成了一条GPU推理流水线。很多人以为这只是“索引里多了一列向量”实际上整个服务从架构形态到资源规划都不一样了——每个query进来要做query嵌入推理百万级候选文档要提前做文档嵌入推理一次搜索可能要触发成千上万次小张量的矩阵运算这些活全部压到GPU上之后批处理策略就成了吞吐和延迟之间的胜负手。这篇文章我想从实际落地角度聊聊大规模搜索结果服务里GPU嵌入推理与批处理这件事。适合正打算给自己的搜索系统加向量召回、或者在已有rerank服务上压吞吐的同行参考也适合对Perplexity这类产品背后的工程实现感兴趣、想搞明白“它们为什么能这么搜”的人。我不会复述官方文档而是把我在真实服务里踩过、调过、优化过的细节拿出来讲。1. 从Perplexity说起为什么搜索引擎需要GPU嵌入推理1.1 嵌入推理到底是什么嵌入推理说人话就是让模型把文本变成一串固定长度的浮点数。这串数字经过训练之后被设计成一种“语义坐标”——语义相近的文本在高维空间里离得近语义无关的文本离得远。搜索引擎拿到一个query同样做一次嵌入推理得到查询向量然后再到提前建好的文档向量库里用矩阵运算把最相近的一批文档找出来。这个过程跟传统BM25最大的区别是BM25匹配的是字面词项嵌入匹配的是语义概念。用户搜“怎么给猫洗澡”传统搜索能召回包含“猫”“洗澡”字样的页面但嵌入搜索还能召回那些通篇在讲“宠物清洁护理”但没出现完整query词的页面。Perplexity这类产品的体验感远超老一代搜索引擎靠的就是这个语义泛化能力。那为什么非要“GPU”嵌入推理因为Transformer架构的嵌入模型核心计算是矩阵乘法和注意力机制。这类运算在CPU上属于“能跑但活很重”的类型。一次推理还好但搜索引擎的流量是按毫秒和并发算的一秒钟哪怕只有几百个查询每个查询都要过一遍TransformerCPU很快就会被吃满延迟也会一路飙升。1.2 为什么必须上GPU一次线上事故的教训我之前维护过一套纯CPU做query嵌入的服务。模型是12层的MiniLM参数量大概3300万单条文本的嵌入推理延迟在CPU上大约是30到50毫秒。听起来还能接受是不是问题出在并发上。搜索系统的查询不是排队来的是一波一波涌进来的。一旦QPS超过50CPU就忙不过来了排队延迟直接把P99推到1秒以上搜索结果页刷不出来。那次事故之后我把推理迁到了GPU。同样的模型在T4上单条延迟大约4到6毫秒吞吐翻了将近十倍。更关键的是GPU天然适合并行处理批量小张量——把几十条文本拼成一个batch喂给模型总耗时往往只比单条多一点点。这就为批处理埋下了伏笔CPU是“一个人做很多事”GPU是“一堆人同时做类似的事”搜索这种海量短文本推理场景天生就是GPU的主场。2. 大规模搜索结果服务的整体设计与思路拆解2.1 服务链路全景从query到结果页先画一下大规模搜索结果服务的完整链路让后面讲细节时大家都有个上下文。整个服务大致分四段入口层接收query做query理解。包括语言识别、分词、改写。Perplexity这类产品通常会在这里把query做一次改写或扩展生成多个更利于检索的子查询。召回层双路召回。一路走传统稀疏倒排索引快速拿回字面匹配的文档另一路走向量召回把query做嵌入推理从向量索引里拿回语义相近的文档。两路结果做合并。向量召回这一步就是嵌入推理的主战场之一。排序层粗排给候选集做快速打分精排则会上更强的模型做rerank。现在的趋势是粗排精排都在向量化两个阶段都有嵌入计算需求。生成层这是Perplexity和传统搜索最不一样的地方。从召回和排序选出的高相关文档片段作为上下文喂给一个大语言模型让它组织成一段自然语言回答并附上引用来源。这一层也需要GPU推理但我们这篇主要聚焦嵌入与批处理生成层的细节先略过。在这个链路里GPU嵌入推理出现在两个位置一是query侧实时推理延迟敏感二是文档侧离线批量处理吞吐敏感。两侧的特点完全不同优化思路也不同。2.2 批处理的核心思想与其刷新更快不如算得更省搜索服务的候选文档量是非常大的。一个中等规模的索引可能有上亿篇文档每篇文档要切成若干段落每段都要做嵌入。如果一段一条的串行推理就算GPU再强上亿次推理也要跑到天荒地老。批处理的核心思想很简单把大量独立的推理请求攒在一起拼成大batch喂给模型。GPU的算力在batch足够大的时候才能真正跑满而单次推理的额外开销比如kernel launch被分摊到每条样本上平均成本大幅下降。打个比方。你有一辆大巴车如果每次只载一个人跑一趟油钱和司机工资摊下来很亏。批处理就是把几十上百个乘客装上车一起跑人均成本一下子就下来了。GPU是辆大巴车你要想尽办法让它满载。但批处理不等于“无脑拼batch”。batch太大模型显存会炸batch太小GPU利用率不够。而且不同请求的文本长度差异很大拼在一起要填充padding导致一堆无效计算。所以真正的工程优化是动态且精细的。2.3 在线推理 vs 离线批处理两套完全相反的逻辑在动手之前必须把在线和离线两条线分开讲因为它们看上去都是GPU嵌入推理优化目标却完全相反。离线文档向量化是典型的吞吐敏感场景。上亿篇文档切段之后可能是几千万甚至上亿条短文本。目标只有一个在尽可能短的时间里用尽可能低的成本把这些文本全部变成向量。这里可以接受较大的batch size比如128甚至512可以把模型精度从fp16降到int8来换取吞吐可以跑很久甚至可以断点续跑。在线query嵌入是典型的延迟敏感场景。用户在搜索框里敲下回车几百毫秒内就要出结果。query嵌入这条链路的延迟预算通常只有几十毫秒。这里batch size不能大因为等不起攒batch的时间模型精度尽量保持高一点因为精度下降会导致召回质量下降。两种场景可以在同一套基础设施上跑但要分开做成两条服务管道否则一次突发的离线跑批任务会把在线查询的GPU资源挤占掉P99瞬间崩盘。后面有一章我会讲怎么混部资源隔离。3. 核心细节解析与实操要点3.1 嵌入模型选型与向量维度权衡模型选型是所有决策里最难回退的一步。嵌入模型的门派很多有老牌的BERT系MiniLM、Sentence-BERT有专门为检索优化的BGE系列、GTE系列还有多语言模型、开源指令微调的嵌入模型。选择时最核心的三个维度向量维度维度越高表达能力越强但存储成本和检索耗时越高。768维和384维的向量在百万级文档的索引里检索耗时会差一倍以上。模型大小MiniLM约120MB和BGE-large约1.3GB的开销差了十倍。模型越大表示能力越强但推理延迟和显存占用也越高。语言覆盖如果你的搜索服务是多语言的必须选多语言模型否则跨语言语义检索就是空谈。我的建议是先跑评测别拍脑袋。拿一批真实query和带标注的相关文档对比各模型的Recall10。如果384维的小模型能达到768维大模型95%以上的召回率果断选小的。Perplexity这类产品能保持低延迟体验很重要的一个原因就是把每个环节的模型做到“够用就好”。另外需要注意一个细节嵌入模型id的输入长度上限。很多嵌入模型是512 token上限超过会被截断。搜索结果里的网页正文远不止512 token这就需要在文本切段环节把每段控制在模型允许范围内还要让切分段保留语义完整性通常是按句子边界切。3.2 GPU推理引擎与显存估算选好模型之后接下来是推理引擎。主流选项有三个HuggingFace Transformers、ONNX Runtime、TensorRT。还有一个近年很火的vLLM主要用于大模型生成嵌入场景用前三个更常见。三者的取舍我用表格说明推理引擎上手难度延迟表现吞吐表现适用场景Transformers低中低快速验证、调试ONNX Runtime中中高中高生产落地、跨平台TensorRT高低高极致性能、固定shape我实际用下来生产环境首选ONNX Runtime它和PyTorch的兼容性好导出一条命令而且自带图优化比如算子融合不用像TensorRT那样花大量时间调engine。TensorRT性能确实最好但模型结构一改就得重新导出、重新调参适合稳定不动的模型。显存估算是另一个容易踩坑的地方。公式大约是总显存 模型权重显存 激活显存 推理框架运行时显存模型权重显存参数量×每个参数的比特数。比如3300万参数的MiniLMfp32就是3300万×4字节≈126MBfp16直接对半≈63MBint8再对半≈32MB。注意这里没有算上CUDA context和cuDNN的缓存实测下来为了安全应该把计算值乘以1.3到1.5。激活显存取决于batch size和输入长度。一个能被batch size为64、序列长度512填满的T416GB显存模型用int8时大概还能剩70%左右的空间所以T4跑这种模型非常充裕。但如果换成千亿参数模型那就是另一套算法了。在搜索嵌入场景模型普遍在1亿参数以下显存一般不是瓶颈反而是batch太大导致的算子中间结果溢出才是实际会遇到的坑。3.3 批处理调度动态batching的完整实现在搜索服务里请求到达不是匀速的而是忽高忽低。静态batch比如固定每32条一批在低峰期会让请求白白排队等候高峰期batch又不够大、GPU吃不饱。动态batching才是工程上的正解。动态batching的思路是设一个最大等待时间比如8毫秒和一个最大batch上限比如128。请求进来后先进入队列等待两个条件之一满足就触发推理队列里的请求数达到上限或者第一个请求已经在队列里等了8毫秒。这样低峰期延迟可控高峰期吞吐最大化。我实现过一版动态batching调度器核心逻辑可以简化成下面这段伪代码class DynamicBatcher: def __init__(self, max_batch128, max_wait_ms8, infer_fnNone): self.max_batch max_batch self.max_wait_ms max_wait_ms self.infer_fn infer_fn self.queue Queue() self._stop False def submit(self, texts): 提交单个请求返回一个future future Future() self.queue.put((texts, future)) return future def run(self): 调度主循环 while not self._stop: # 先取第一条然后等待攒batch first, first_future self.queue.get() batch [(first, first_future)] deadline time.time() self.max_wait_ms / 1000.0 while len(batch) self.max_batch and time.time() deadline: try: item self.queue.get(timeout0.001) batch.append(item) except Empty: break # 合并请求并推理 texts [t for t, _ in batch] vectors self.infer_fn(texts) for i, (_, fut) in enumerate(batch): fut.set_result(vectors[i])这套调度器在实际压测中拿到了不错的收益在T4上跑MiniLM模型单流串行推理的QPS大约80左右换上动态batching后QPS可以拉到600以上P99延迟没有明显恶化因为等待窗口被控制在8毫秒以内。代码里的关键点在于合并请求后要保证futures一一对应返回并且推理异常要有超时和重试机制。线上环境我还加了“队列积压熔断”——如果队列里的请求超过500条说明下游推理已经跟不上了此时直接放弃等待、返回降级结果强于让所有请求排队到超时。3.4 缓存被低估的批处理前置层批处理负责让GPU吃饱和低延迟但真正的第一道加速器应该是缓存。搜索场景有个很大的特点热门query的重复率很高。在Perplexity的架构里缓存也大量存在有人查过“今天天气如何”另一个人大概率也会查类似的话。缓存要做两件事精确匹配缓存完全相同的query文本直接命中返回之前计算好的向量。可以用redis或者内存LRU实现。语义近似缓存这个高级一些。新query和某个已缓存query的向量距离很近时直接复用。这要求缓存索引本身也是一个向量检索库本质上“用搜索服务来加速搜索服务”。我第一次做缓存时忽略了缓存值的过期策略结果上游文档更新后旧向量和新向量混在一个索引里召回质量出现肉眼可见的抖动。后来给每个缓存项加了文档版本号版本不匹配直接强制穿透重算问题才解决。4. 实操过程与核心环节实现这一章把前面提到的所有方案串起来走一遍可落地的实操流程。环境以PyTorch 2.x ONNX Runtime Triton Inference Server为例这是我现在生产环境用的组合。4.1 环境准备与模型导出首先是PyTorch的GPU环境。别只看装个torch就完事CUDA、cuDNN和PyTorch版本之间的兼容矩阵非常关键。以PyTorch 2.1为例对应的CUDA 12.1安装命令pip install torch2.1.2 torchvision0.16.2 torchaudio2.1.2 --index-url https://download.pytorch.org/whl/cu121装完之后先跑一行验证GPU是否真的可用import torch print(torch.cuda.is_available()) # True print(torch.cuda.get_device_name(0)) # Tesla T4 之类踩坑提示torch.cuda.is_available()返回True不代表一切正常最好再跑一次实际张量运算确认没有算子兼容问题。接着把训练好或者选定的嵌入模型从PyTorch格式导出为ONNX。导出时要把动态轴设置好否则后续想换batch size就麻烦了。以MiniLM为例from transformers import AutoTokenizer, AutoModel import torch model AutoModel.from_pretrained(sentence-transformers/all-MiniLM-L6-v2) model.eval() dummy_input torch.randint(0, 30522, (1, 128)) # (batch, seq_len) torch.onnx.export( model, dummy_input, minilm.onnx, input_names[input_ids], output_names[sentence_embedding], dynamic_axes{ input_ids: {0: batch, 1: seq_len}, sentence_embedding: {0: batch} }, opset_version17, )这里有一个容易忽略的点很多库里的嵌入模型本身输出不止一个张量你要确定取的是哪个输出。SentenceTransformer系列通常取最后一层隐藏层的mean pooling结果。如果导出的是原始模型结构pooling需要在导出前或者推理后自己做千万别把last_hidden_state当成最终向量用。4.2 Triton Inference Server部署推理服务ONNX模型导出之后用NVIDIA Triton Inference Server来部署是最稳的路径。Triton原生支持动态batching不用自己手写调度器。配置文件长这样name: embedding platform: onnxruntime_onnx max_batch_size: 256 input [ { name: input_ids data_type: TYPE_INT64 dims: [128] } ] output [ { name: sentence_embedding data_type: TYPE_FLOAT32 dims: [384] } ] dynamic_batching { preferred_batch_size: [16, 32, 64, 128] max_queue_delay_microseconds: 5000 }这段配置里preferred_batch_size和max_queue_delay_microseconds就是Triton版的动态batching参数。前者告诉调度器优先凑哪些batch大小后者限制请求最大在队列里待多久。5毫秒的延迟上限在实测里对在线搜索的P99影响可以控制在10毫秒以内。Triton的启动命令tritonserver --model-repository/models --gpu-memory-fraction0.5注意--gpu-memory-fraction如果你在同一块GPU上还要跑别的推理服务这个参数能限制Triton的显存占用防止模型加载时把显存一次性吃光。但如果这是专卡专用建议设成0.9以上让Triton有足够显存做cuda缓存管理。客户端调用Triton的Python代码import tritonclient.http as httpclient import numpy as np client httpclient.InferenceServerClient(urllocalhost:8000) input_ids np.array([[101, 1, 1, ..., 102]], dtypenp.int64) inp httpclient.InferInput(input_ids, input_ids.shape, INT64) inp.set_data_from_numpy(input_ids) result client.infer(embedding, [inp]) embedding result.as_numpy(sentence_embedding)这里有个隐蔽的坑Triton对输入张量shape的约束非常严格。如果你在配置里指定了dims: [128]那所有请求的序列长度都必须固定128。实际搜索场景每条文本长度不一所以要么在预处理时统一padding到固定长度要么在配置里把维度设成-1表示动态。前者实现简单但浪费算力后者要用到Triton的max_batch_size和dynamic shape功能配置更复杂但效率更高。4.3 离线文档向量化的完整跑批流程文档向量化的跑批不像在线推理那样有现成的Triton配置可用因为它还涉及文档的读取、切段、去重、过滤以及结果写回。我的流程大致分四步文档读取与切段把网页/PDF/文档切分成512 token以内的段落。切分点尽量放在句子末尾。从实践看段落重叠overlap设成50-100 token能明显改善跨段语义割裂的问题。批量请求把切好的段落按128条一组打包通过HTTP/gRPC发给Triton的嵌入服务。这里要注意并发控制别一次性把所有段落全压进队列否则要么客户端内存爆掉要么把服务端打挂。我是按文档维度做分片每个分片一个worker线程。向量校验拿到向量之后做一个快速校验检查shape对不对、有没有NaN/Inf值。亲测这个步骤能拦住很多“索引里出现零向量”的诡异Bug。写回向量库向量和文档ID一起写入向量数据库。这里根据你的规模选型百万级用FAISS单机够了千万级以上要上Milvus、Elasticsearch插件这类分布式方案。整个Pipeline用Airflow编排每天全量跑一次增量每五分钟跑一批。某次版本发布后我们发现增量任务老是跑到一半就失败排查半天发现是切段逻辑改了之后生成了几条比模型max_length还长的段落导致Triton返回400。从此之后所有段落生成完一律本地先用tokenizer检查长度超长直接丢弃再也没出过问题。4.4 在线查询链路的嵌入实时推理在线侧query嵌入走Triton的HTTP/gRPC接口。查询阶段注意三点第一query改写后要重算向量。Perplexity这类产品通常会对query做改写但改写出的N个子查询必须逐条做嵌入推理不能复用原query的向量。我见过有同学图省事直接用原始query向量去召回改写后子查询的相关文档召回率掉了将近10个点。第二在线batch size不能大。虽然Triton支持动态batching但在线场景请求不能等太久。max_queue_delay_microseconds建议设在3-8毫秒之间超过这个时间宁可单条推也不要让用户在结果页前多等几十毫秒。第三超时与降级。在线推理服务必须设置超时我通常设100毫秒。一旦Triton响应超时整个向量召回分支放弃只走传统的BM25召回保证搜索结果页永不空白。这个降级策略在后端抖动时是救命的。4.5 向量索引与检索服务嵌入推理完成之后向量本身是死的关键在检索。拿FAISS做示例影响检索效果和速度的一是索引类型二是检索参数。索引类型这里有个典型取舍索引类型检索速度内存占用召回精度适用规模IndexFlatIP暴力扫描慢高100%百万级以下IndexIVFFlat快中高调好nprobe千万级IndexIVFPQ极快低中低亿级我实际用下来千万级文档用IVFFlatnprobe设为16-32时召回率和延迟都能兼顾。如果直接上PQ压缩召回率下降通常不止1个点除非你的业务对完全召回没有刚性需求否则不建议默认选PQ。再补充一个容易被忽略的点向量索引是需要训练和更新的。IVF索引要先用一批样本来训练聚类中心如果文档分布发生明显漂移旧的聚类中心就不准了。我们的做法是每天凌晨用全量向量重新训练一次索引训练期间新写入的向量进增量缓冲区搜索时合并查两处结果。5. 常见问题与排查技巧实录5.1 显存不足与OOM排查现象Triton启动后一切正常跑着跑着突然报CUDA out of memory服务自动重启。排查步骤先看是不是并发请求过多导致激活显存暴涨。nvidia-smi能看实时显存占用如果发现显存被逐步吃满先降max_batch_size把256改成128试试。排查模型权重有没有被重复加载。容器化部署时同一个模型如果副本开太多每个副本都占一份显存。用--gpu-memory-fraction限制每副本的显存上限。确认是否开启了CUDA缓存自动清理。PyTorch和ONNX Runtime会缓存一些显存块以加速后续分配平时看着占用高是正常的但搭配Triton时它会在request结束后自动清理如果你手动设了PYTORCH_NO_CUDA_CACHE反而可能导致频繁分配碎片化。踩坑心得有一段时间我为了控制显存把Triton的max_batch_size和--gpu-memory-fraction同时调小结果吞吐下降一半不说显存碎片化还更严重。后来把内存分数调回0.9只用batch size控制并发问题才消失。5.2 推理延迟突刺怎么定位现象P99延迟从平稳的15毫秒突然冲到200毫秒持续几秒后恢复。第一排查目标永远是“有没有别的任务在抢GPU”。离线跑批任务大多不做资源隔离一个大的批处理任务占满GPU后在线查询的推理请求就得排队。我在混部架构里加了独占调度策略给在线服务预留了一整块GPU离线任务只能使用剩余的份额P99突刺基本消失。第二排查目标是预处理耗时。嵌入服务在真正推理之前要做tokenization和padding这一步用的是CPU。如果文本长度分布极不均衡比如一批长文本突然进来导致padding数量暴涨预处理时间会上来。在调度层限制单batch内文本最大长度比把超长文本单独放一批能显著平滑延迟。第三排查目标是Triton的模型实例数量。ONNX Runtime一次性只能处理一个请求的batch如果你的QPS超过了单实例的吞吐极限Triton会自动扩容新实例但这个扩容过程耗时在几十毫秒量级正好是延迟尖刺的来源。提前用instance_group配置好足够的实例数量避免运行中扩容。5.3 批处理参数调优速查表批处理的参数调优是门手艺活我把调过的参数和心得整理成一张表参数影响面调整倾向备注max_batch_size吞吐、显存从64开始每翻倍压测一次超过128后收益递减max_queue_delay延迟、吞吐在线3-8ms离线20-50ms越长吞吐越高但延迟越糟preferred_batch_size稳定性设为2的幂次配合模型优化效果最好文本padding长度算力利用率按业务文本长度分布选P90别按最大长度给所有文本装箱最后一条最容易被忽视。搜索结果的文档长度是长尾分布的大部分是300 token以内少数长文有2000 token。如果统一按512 padding短文本大部分算力都浪费在填充token上。我是先统计线上文本长度分布把长度阈值按90分位切成256 token超长的截断到256精确的关键信息如果因截断丢失再靠进阶的段落筛选逻辑保底。事实证明这类“按数据分布定参数”的做法比任何优化技巧都提效明显。6. 最后再聊聊我踩过的几次坑我对嵌入推理批处理最深的体会是很多问题不是模型或者GPU带来的而是服务化那一层没有做好。格式转换、超时控制、并发模型、资源隔离这些“不性感”的工程环节才是决定体验上限的地方。Perplexity这类产品能在产品质量上拉开差距除了模型选得好更大的功劳是工程层面把每条链路的延迟和资源控制到了极致。我自己踩过的坑包括刚开始用Triton时没配动态batching请求一多就100% CPU打满后来配了动态batching但没设max_queue_delay低峰期体验还好高峰期延迟直接无法接受再后来在多个服务混部时没做显存限制一次离线跑批把在线推理服务的显存挤爆线上搜索挂了十几分钟。这几件事让我彻底明白了一个道理GPU推理服务的批处理从来不是“把请求拼一起”这么简单你要在吞吐、延迟、稳定性之间找一个可持续的平衡点。如果你现在正打算给自己的服务接入GPU嵌入推理我的建议是先从一个小规模、可观测的POC开始把Triton动态batching的四个关键参数跑一遍压测记录数据再决定怎么调。别一上来就追求极致性能稳定性永远是第一位的。