Perplexity 工程团队近日发布了题为《GPU 上的快速嵌入》的技术博客，首次详细介绍了支撑其嵌入模型 pplx-embed 的推理基础设施。在 AI 搜索产品中，检索质量主要受两个因素制约：一是嵌入模型本身的质量，二是如何低成本地在海量索引上运行该模型。Perplexity 的这篇文章聚焦后者，说明了为 pplx-embed 以及搜索、计算机和 API 平台中的排序模型所构建的服务架构。
Perplexity 将嵌入服务划分为两种流量模式：批处理嵌入和在线嵌入。批处理嵌入发生在构建或重建向量数据库时，此时吞吐量决定成本；在线嵌入发生在查询时，短查询需要极低的延迟。排序任务则介于两者之间：向量检索后需要为大批量文档打分，需要在吞吐和延迟之间取得平衡。
架构上的关键决策是：Perplexity 并没有为嵌入服务单独开发一个推理引擎。由于嵌入模型本质上是小型 Transformer，批处理嵌入类似于计算密集的 prefill 阶段，而在线嵌入通常只包含几个 token，类似于内存密集的 decode 阶段。因此，开发团队直接复用了其 LLM 推理栈中的 prefill 和 decode 内核。
整个请求链路由三个组件构成。Ivy 是一个 Rust 编写的 HTTP 网关，负责 CPU 侧的工作，包括 JSON 解析、分词、输入模板处理和批次切分，并将请求转换为自定义的 gRPC 协议。它还能将大批量请求拆分为更小的块，并在各个副本间进行负载均衡，从而解决生产环境中负载不均衡的问题。Tulip 是推理服务器接口，使用 Rust、tokio 和 tonic 构建，负责任务调度和批次组装，然后将其分发给引擎。ROSE 则是推理引擎，主要使用 Python，提供内核、层和模型定义，管理 CUDA 图，并向 Tulip 暴露 step() 函数。
调度器的设计刻意保持简单：Tulip 采用先到先服务的方式处理请求，同时积累批次。这种设计源于一个观察：对于小尺寸嵌入模型，在 Perplexity 使用的序列长度范围内，稠密层的线性计算成本往往高于注意力机制带来的平方级成本。因此，延迟大致与 token 数成正比，而不是序列数。当批次规模足以让 GPU 饱和（在亚十亿参数模型上大约为 512 个 token）时，继续增加序列数并不会带来额外的效率收益。
为了减少小批量场景下的 CPU 端内核启动开销，Perplexity 为所有嵌入模型构建了整模型 CUDA 图，将数十次内核启动合并为一次驱动调用。由于嵌入模型较小，GPU 计算超过启动开销的临界点通常出现在数千 token 和数十序列的批次规模。部分注意力实现依赖动态主机端输入，会阻碍整图捕获；Perplexity 为此将改动上游提交给 FlashInfer，解决了该问题。CUDA 图需要针对每种配置捕获，因此 token 数会按 64 或 256 的倍数进行填充。这仍然可能产生数千张图，每模型需要数分钟捕获时间。Perplexity 采用懒捕获策略：每种配置先进行一次 eager 预热，在第二次命中时才触发捕获和重放。这虽然会带来启动阶段的 p99 延迟代价，但将数分钟的 eager 工作分散到数小时内。
另一个关键组件是 LazyTensor。它跟踪一个页锁定主机缓冲区，并关联一次 cudaMemcpyAsync 操作和一个 CUDA 事件。step() 不会阻塞等待 GPU 完成，而是立即返回一个 LazyTensor，使得 Rust 异步任务可以在等待批次 N 完成的同时，CPU 继续入队批次 N+1。这种设计实现了 CPU 端的批准备与 GPU 执行的重叠。
内核层面依然存在优化空间。ROSE 支持多种适用于非均匀输入的注意力后端，包括 FlashInfer 2、FlashInfer 3 和 FlashAttention 4。Perplexity 的报告表明，FlashAttention 4 在大多数情况下更快，但在基于 Qwen 的超长序列模型中，FlashInfer 3 表现更好，因此后端选择是逐个案例进行的。值得注意的是，当 ROSE 用于嵌入模型时，不会实例化 KV 缓存，而是调度到非均匀注意力实现以避免填充。
在基准测试方面，Perplexity 与 vLLM v0.22.0 在 BF16 精度下进行了对比，使用真实权重和评估数据，并通过预热运行验证余弦相似度差异不超过 0.1%。测试涵盖四个场景：低延迟嵌入（batch 1，128/512/4096 token）、低延迟排序（batch 5/25/50，512 token）、高吞吐嵌入（batch 100，四个并发进程）以及高并发嵌入（1 到 16 个并发请求，包含 Ivy 的 tokenization 和网络开销）。测试结果显示，该架构在延迟与吞吐上均具有竞争力。
总体而言，Perplexity 的嵌入服务栈以复用 LLM 内核为核心，用三个内部服务协同完成请求处理，并通过 CUDA Graph、懒捕获和 LazyTensor 实现了显著的性能提升。这些组件目前属于内部基础设施，外部用户可通过 Perplexity 的 Embeddings API 使用 pplx-embed 能力。