跳到主内容
@wquguru
精选90meng shao技巧与观点

Perplexity 详解:通过 CUDA Graph 与异步追踪优化嵌入推理

每次搜索的第一步:Perplexity 如何把嵌入模型的推理开销压到极限 ?

原文
发到 X
推荐理由

这是一篇极具复用价值的工程实践指南,详细拆解了从架构设计到 Kernel 优化的完整链路,特别是 CUDA Graph 惰性捕获和异步追踪的具体实现,做 RAG 或向量搜索的同学可以直接参考落地。

每次搜索的第一步:Perplexity 如何把嵌入模型的推理开销压到极限 ?

Perplexity 的每一个答案,都始于嵌入与排序模型从 EB 级索引中挑出最相关的结果。

Perplexity 拆解了它背后的服务基础设施:当 GPU 计算已趋于最优,加速空间在 GPU 之外,CPU 调度、kernel 启动开销、主机-设备同步、请求路由。团队通过三层架构(Ivy / Tulip / ROSE)加两项关键技术(整模型 CUDA Graph + LazyTensor 异步追踪)来吃掉这些开销,在延迟和吞吐上优于 vLLM v0.22.0。 https://www.perplexity.ai/hub/blog/fast-embeddings-on-gpus

问题定义:嵌入推理的两种流量形态

Perplexity 做了一个很关键的类比,把嵌入推理映射到 LLM 推理的已知模式上: · 批量嵌入(建索引/重索引):大批量长文档,追求最大吞吐以降低成本。计算特征等同于 LLM 的 Prefill 阶段,属于算力受限(compute-bound)。 · 重排打分(向量检索之后):中等批量,需要在吞吐与延迟之间取得平衡。同样对应 Prefill,瓶颈介于算力与带宽之间。 · 在线嵌入(用户查询):只有几个 token,追求最低延迟。计算特征等同于 LLM 的 Decode 阶段,属于显存带宽受限(memory-bound)。

这个映射的工程意义在于:他们不需要为嵌入模型单独写一套 kernel。LLM 服务栈里已经调优过的 prefill 和 decode kernel 可以直接复用,pplx-embed 和 Qwen3.5 的 LLM decoding 走的是同一套 kernel。

三层架构:Ivy → Tulip → ROSE

Ivy(Rust HTTP 网关) 负责所有 CPU 侧预处理:JSON 解析、分词、输入模板、批次切分,再转成自定义 gRPC 协议发给下游。把这些从推理实例剥离出来的好处是,分词和格式化参数可以独立调整,不用碰昂贵的 GPU 实例。Ivy 还承担负载均衡:把大批量请求切块分发到多个副本,避免单个副本被大请求压垮。另外他们自研的 unigram 分词器已全量上线,对延迟改善显著(这是另一篇博文的内容)。

Tulip(Rust gRPC 推理服务器) 基于 tokio + tonic,做调度和批处理。调度策略刻意做得极简:先来先服务,请求在 GPU 忙时自然累积,空闲时打包发出。

这个"简单"是有理论依据的。他们观察到:对小型嵌入模型、在实际服务的序列长度下,dense 层的线性开销主导,attention 的二次开销可以忽略。所以延迟基本与 token 总数成正比,与序列数无关。一旦批次达到约 512 token(十亿参数以下的模型),GPU 就已饱和,再塞更多序列没有收益。既然如此,复杂的调度算法带来的边际价值就很小。

ROSE(Runtime-Optimized Serving Engine,Python) 真正的模型实现层,定义 kernel、层、模型结构、前向传播和 CUDA Graph 管理。它原本为 LLM 服务而建,通过一个 step() 函数与 Tulip 桥接——接收一个 batch,返回一个指向 GPU 上计算结果的引用(而非结果本身,这点后面很关键)。

嵌入模型与 LLM 的差异被限定在 attention 层:嵌入不需要 KV cache,改用支持 ragged(变长非填充)输入的 attention kernel 变体,避免 padding 浪费。dense 层则完全一致,因为 token 向量本来就是独立处理的。

核心技术一:整模型 CUDA Graph + 惰性捕获

为什么需要? 一次前向传播是 CPU 与 GPU 的交替配合:CPU 负责调度和 kernel 启动,GPU 负责矩阵乘法、attention、norm、激活。大批量时 GPU 执行时间长,CPU 开销可忽略;但小批量时,Python/PyTorch 的 kernel 启动开销会反过来超过 GPU 实际计算时间。嵌入模型本身很小,这个"拐点"来得很早——文中说在数千 token、几十个序列的批次规模上 CPU 开销才被 GPU 时间盖过。也就是说,在线查询这种几十 token 的请求,绝大部分时间在等 CPU。

做法 CUDA Graph 把整个前向传播的所有 kernel 启动元数据一次性捕获,之后一次驱动调用即可回放,跳过所有 Python 代码。Perplexity 为所有嵌入模型构建整模型级别的图。有些 attention 实现依赖主机侧动态输入来配置 kernel 启动,会破坏图捕获,他们把修复上游到了相关 kernel 库。

代价与缓解 CUDA Graph 必须按配置逐一捕获,嵌入场景下配置 = 序列数 × token 数的组合。即便把 token 数填充到 64 或 256 的倍数分桶,仍然有数千张图,全量捕获要几分钟。每次捕获要跑两遍:一次 eager 预热(编译 kernel、分配缓冲区),一次捕获运行。

解决方案是惰性捕获:不在启动时预捕获,而是在服务过程中,每个配置第一次命中走 eager,第二次命中触发捕获并回放,之后全部走回放。代价是启动初期 p99 延迟会抖动,收益是把几分钟的冷启动开销摊到几小时里,实例可以更快上线,扩缩容更灵活。

这是一个典型的工程权衡:牺牲部分冷启动期的尾延迟,换取部署敏捷性。

核心技术二:LazyTensor 异步结果追踪

问题 CUDA 操作是异步入队的,但主机代码要读取结果就必须显式同步等待。如果 step() 每次都"启动图 → 等待完成 → 返回结果",那么 GPU 执行期间 CPU 是空闲的,反之亦然,两者无法重叠。

做法 LazyTensor 封装三样东西:一块页锁定(pinned)主机内存、一个 cudaMemcpyAsync 设备到主机的拷贝操作、一个 CUDA event。它在前向传播启动后立即在同一个 stream 上入队。因为 stream 内顺序执行,这个拷贝必然等待前面所有 kernel 完成,所以关联的 event 同时标志"前向传播完成"和"结果已在 CPU 可读"。

step() 返回的就是这个 LazyTensor。Tulip 的 Rust 异步任务可以在它上面 await,而与此同时 CPU 已经在准备并入队下一个批次。结合 CUDA Graph 极低的启动开销,GPU 几乎不会因为等 CPU 而空转。

两项技术的协同 CUDA Graph 解决"CPU 启动太慢",LazyTensor 解决"CPU 等 GPU 时无事可做"。前者压低单次延迟,后者提升流水线利用率。在高吞吐场景下,CUDA Graph 释放的 CPU 时间正好被用来提前调度下一批次,这是文中强调的双重收益。

Kernel 层面:多后端按场景选择

虽然主题是"kernel 之外",文章也承认 kernel 仍然重要。ROSE 集成了 FlashInfer 2、FlashInfer 3 和 FlashAttention 4 来实现 ragged attention。

总体上 FlashAttention 4 更快,但在 Qwen 系模型的超长序列上 FlashInfer 3 反超。因为性能随 attention head 数量和维度变化,他们保留多套配置,按模型逐案选择而非一刀切。

基准测试

对比对象是 vLLM v0.22.0,BF16 精度,真实模型权重和评估集输入,预热后验证余弦相似度偏差在 0.1% 以内(保证是同等质量下比速度)。四组测试覆盖了前文定义的全部流量形态: · 低延迟嵌入:batch 1、完全串行、序列长度 128/512/4096——对应在线查询 · 低延迟打分:batch 5/25/50、512 token——对应重排 · 高吞吐嵌入:batch 100、4 个并发进程、序列长度 512/1024/4096——对应建索引 · 高并发嵌入:512 token、batch 1、并发 1~16——这组把 Ivy 分词和 Ivy-Tulip 网络开销也计入,是最接近端到端生产的指标

从测试设计可以看出他们最关注的是尾延迟(p99/max)而非仅均值,这与搜索产品的用户体验直接相关。

更进一步:量化金融体系

看懂新闻只是起点——沿量化金融路径,把它变成能交付的工程能力

进入量化体系 →

关联讨论

同一事件的更多信源

相似阅读

另一事件,读法相近