GPU 上的快速嵌入

快速且准确的搜索对 Perplexity 的方方面面至关重要，从 Search、Computer 到我们的 API 平台。在幕后，繁重的工作由嵌入和排序模型完成，它们帮助我们的系统为给定的查询识别出最相关的结果。通过训练和部署我们自己的模型（例如 pplx-embed），我们实现了顶尖的质量和延迟表现。

本文深入探讨了 Perplexity 针对此类特殊模型提供的推理基础设施。我们讨论了高效满足 AI 原生搜索推理需求的各项技术，在为我们的百亿亿次（Exabyte 级）搜索索引提供支持的同时，能够快速对模型进行原型设计与评估。这些技术共同拓展了搜索质量与效率的帕累托前沿（Pareto frontier），使我们能够以最低的成本和延迟，为 Agent 和用户提供最优质的搜索结果。

在典型的搜索设置中，索引文档会使用嵌入模型映射到高维向量空间中并存储在向量数据库里。通过使用相同的模型对查询进行嵌入，可以通过寻找最接近查询向量的向量来定位相似的文档。这就为推理引擎带来了两种不同的流量模式：

批量嵌入（Batch Embedding）：在构建、扩充或重新索引数据库时，必须将大批量的文档嵌入到向量空间中，从而最大化吞吐量以将成本降至最低。

在向量搜索之后，必须对大批量的文档进行评分，在吞吐量和延迟之间取得平衡。

在线嵌入（Online Embedding）：在查询数据库时，必须对简短的查询进行嵌入以进行查找，从而将延迟降至最低。

我们构建了推理基础设施，以便在各个用例之间尽可能多地复用公共组件。由于我们通常使用小型 Transformer 模型来生成嵌入，因此我们与 LLM 推理代码共享了大部分实现：批量嵌入类似于计算密集型的预填充（Prefill），而通常针对少量 Token 运行的在线嵌入在计算上则类似于内存密集型的解码（Decode）。因此，我们复用了经过优化的预填充和解码算子（Kernel）来提供嵌入模型服务。这样一来，我们便能以极少的额外工程工作量实现巨大的批量推理吞吐量，同时为在线嵌入工作负载保持低延迟。

我们通过标准化 API 对外及对内（通过我们的 API 平台）暴露推理服务。在底层，处理嵌入请求涉及多个服务：

Ivy 是 Perplexity 服务调用的 Rust HTTP 网关。

它处理请求的 CPU 端工作，例如 JSON 解析、分词、输入模板化和批次拆分，并将请求转换为自定义的 gRPC 协议以供下游服务器使用。这种分离使我们能够配置有关分词和输入格式的某些参数，而不必触碰较重的推理实例。

Tulip 是推理服务器接口。

这是一个使用 Rust、tokio 和 `tonic` 实现的 gRPC 服务器。Tulip 接收 gRPC 推理请求，并处理调度和批处理。然后，它将批次发送到 ROSE 引擎，并将完成的响应返回给客户端。

**ROSE**（运行时优化推理引擎，Runtime-Optimized Serving Engine）用于实现模型推理。

它主要用 Python 定义，为各种模型提供算子（Kernel）、层和定义。ROSE 实现了模型的前向传递，同时提供了专门针对嵌入（Embedding）优化的 CUDA 图管理。它通过一个 step() 函数与 Tulip 桥接，该函数接收一个批次并返回它在加速器上执行的计算的引用。

Transformer 模型以及底层的 Hopper/Blackwell 架构都是成熟的技术，因此在 GPU 端进行嵌入推理时，各推理引擎的实现方案已基本收敛至接近最优。即便如此，我们发现在运行时以及将模型端到端暴露给客户端的测试框架（Harness）中，仍然存在进一步优化的空间。特别是，我们发现通过精细管理 CUDA 图，并在原生 Rust 引擎中构建 LazyTensor 抽象来异步追踪 GPU 端的结果，可以有效降低延迟。我们在 Tulip 中实现了这些功能，使其能够与 ROSE 的模型实现进行高效交互。

我们设计 Tulip 的初衷，是让它成为模型服务层上极其轻量级的接口。它通过 Tokio 异步任务处理传入的请求，维护一个跟踪请求的池，并从中调度批次以分发给加速器。Tulip 中的调度机制非常简单：在 Tulip 分发工作或等待结果时，请求会不断累积。从累积的请求中，系统会按照先来先服务的原则挑选序列来通过模型运行。

简单的调度机制源于对模型性能的一个观察：对于小型嵌入模型，在我们服务的序列长度下，我们注意到稠密层的线性成本主导了注意力机制的二次成本。因此，延迟主要与 Token 数量成正比，而与序列数量无关。因此，一旦批大小足以使 GPU 饱和（在参数量不到 10 亿的模型上大约为 512 个 Token），向其中塞入更多序列并不能提高效率。

为了与模型进行高效交互，Tulip 依靠 CUDA 图和延迟结果追踪（Lazy result tracking）来重叠 GPU 和 CPU 的工作，并充分利用可用资源。

运行模型的前向传递涉及 CPU 端和 GPU 端的工作。CPU 负责调度批次并使用适当的参数启动算子，而 GPU 则执行相关的矩阵乘法、注意力、归一化（Norm）或激活算子。对于训练和重新索引（Reindexing）等高吞吐量工作负载，由于批大小和 GPU 端延迟都很大，CPU 端的开销可以忽略不计。然而，在较小的批大小下，CPU 端的工作量可能会超过 GPU 端。

为了缓解开销，与其启动独立的算子，不如构建一个 CUDA 图来捕获通过对 CUDA 驱动程序进行单次调用即可启动前向传递的所有算子所需的元数据。这就消除了为能够捕获 CUDA 图的配置重新运行昂贵的 Python 和 PyTorch 代码的需求。

针对每个模型，我们都会追踪一个拐点，以确定 GPU 执行成本高于 CPU 端算子（Kernel）启动成本的最低 Token 数量。由于嵌入模型较小，我们发现这个拐点出现在包含数千个 Token 和数十个序列的批次中。某些注意力机制实现依赖于动态的主机端输入来配置算子启动，这阻碍了全模型预填充/稠密 CUDA 图的使用。我们将相关更改合入上游（Upstream），以便在我们的推理引擎中启用它们。

为了解决开销问题，我们为所有嵌入模型构建了全模型 CUDA 图，并使 CPU 工作与 GPU 工作重叠。由于 CUDA 图最大程度地减少了 CPU 端的开销，一旦启动了一个图，我们便有空闲时间在下一个批次可用时立即启动并将其执行排队。待处理批次的结果由 LazyTensor 追踪，这允许 Rust 中的异步任务进行阻塞，直到前一项批次执行完成。CUDA 图通过确保我们不会被算子启动的成本拖累来助力低延迟服务，并在高吞吐量情况下促进更好的调度，因为它们释放了 CPU 以便更早地处理下一批次的工作。

必须为每种不同的配置捕获 CUDA 图，对于嵌入来说，这意味着每个序列数和 Token 数组合都需要一个图。由于这个网格非常庞大，我们将 Token 数填充（Pad）到 64 或 256 的倍数的存储桶（Buckets）中。这仍然会产生数千个图，对于典型模型而言，捕获这些图可能需要几分钟时间。捕获成本来自两个方面：首先是必须执行贪婪（Eager）前向传递以编译算子并为需要它们的各种算子设置缓冲区，随后是重新执行 Python 代码的捕获运行。

我们在引擎提供服务时通过延迟捕获 CUDA 图来缓解启动成本。我们追踪每种配置，并确保在第二次命中时触发图捕获和重播之前，先通过贪婪（Eager）预热运行。随后对同一图配置的所有执行都将通过 CUDA 图重播进行。延迟图捕获会对启动期间的 p99 延迟产生影响；然而，它对于将数分钟的贪婪工作分散到多个小时内非常有价值。更快的启动时间使我们能够更好地扩展和管理嵌入部署。

通过 CUDA，GPU 工作是异步的。由于异步启动算子会将其排队到一个流（Stream）中，主机代码必须显式同步才能读出生成的向量。为了促进更高程度的并行性，并能够在等待前一个设备端任务完成的同时启动未来的批次，我们依靠 LazyTensor 抽象来追踪值。

LazyTensor 通过一个事件来追踪页面锁定内存（Page-locked memory）中的主机缓冲区以及一个 cudaMemcpyAsync 操作，该操作用于从设备复制数据。它在同一流（Stream）上启动前向传递后触发。由于复制操作必须等待流上的所有先前算子执行完毕，因此关联的事件既能追踪前向传递的完成，又能追踪 CPU 上结果的可用性。

我们在 ROSE 编码器引擎中利用 LazyTensor 来重叠 GPU 和 CPU 的工作。每个 step() 调用不再是运行 CUDA 图并等待其完成，而是由 step() 返回一个 LazyTensor 来异步追踪其结果。结合 CUDA 图，这有助于我们实现低延迟和更好的吞吐量。

我们对原本为 LLM 服务构建的 ROSE 引擎进行了调整，使其同样能够处理嵌入模型的执行。为了最大限度地减少支持嵌入模型所需的开发工作，ROSE 在 LLM 和嵌入之间积极复用代码。例如，pplx-embed 服务和 Qwen3.5 LLM 解码都使用相同的算子。这种代码共享使我们能够轻松提供最初从 LLM 微调而来的嵌入模型，用于原型设计、评估和生产推理。

对于稠密层（Dense layers），由于 Token 向量是独立处理的，嵌入和 LLM 推理完全相同。在注意力层中，差异通过以下方式处理：添加对参差不齐输入（Ragged inputs）的支持，以及 LLM 所需的分页预填充和解码设置。在提供嵌入模型服务时，我们不会实例化 KV 缓存，而是分发到支持参差不齐格式（Ragged format）的各种注意力算子，以避免填充（Padding）。配套的转换和校准例程也与 LLM 共享。

我们的推理 HTTP 代理层 Ivy 在性能中也扮演着重要角色。由于生产环境中的请求负载各不相同，将单个请求路由到单个副本可能会导致负载不均衡。Ivy 将大批量请求拆分为多个数据块（Chunk），并在各个副本之间进行负载均衡，从而提高资源利用率并平滑延迟。我们最近在 Ivy 中全面推出的自研一元语法（Unigram）分词器工作，大幅降低了相比现成分词器的延迟。

ROSE 支持多种注意力后端。不同的算子可能适合特定的问题规模。随着时间的推移，我们集成了 FlashInfer 2、FlashInfer 3 和 FlashAttention 4 算子来实现参差不齐的注意力机制（Ragged attention）。

总体而言，我们发现 FlashAttention 4 的速度更快。然而，在序列长度非常长的情况下，FlashInfer 3 在 Qwen 系列模型上的表现优于前者。由于性能和调优会随注意力机制头的数量和维度而变化，我们保留了对多种配置的支持，并在提供推理服务时针对具体情况做出决策。

我们以 vLLM v0.22.0 为基准进行测试，在评估数据集派生的实际模型权重和输入上以 BF16 精度运行推理。所有计时运行前均进行预热运行，以验证余弦相似度的偏差在 0.1% 以内。

我们报告了预分词请求批大小为 1、完全顺序请求、序列长度分别为 128、512 和 4096 个 Token 时的运行时间。

预分词的请求批大小为 5、25 和 50，序列长度为 512 个 Token。

请求批大小为 100，4 个并发进程提交请求，序列长度分别为 512、1024 和 4096 个 Token。

序列长度为 512，批大小为 1，但我们发送 1、2、4、8 和 16 个并发请求。此基准测试还包括通过 Ivy 进行的分词成本，以及 Ivy 与 Tulip 之间的网络开销。

由 Ivy、Tulip 和 ROSE 组成的推理基础设施使我们能够以更低的延迟和更高的吞吐量为 Perplexity 提供嵌入服务，与现成（Off-the-shelf）解决方案相比，在降低成本的同时实现了更准确的搜索。

通过专注于特定模型并掌控整个技术栈，我们获得了在性能与灵活性之间取得有效平衡所需的自由，将高度可重用且高性能的 Rust 原语与更通用的 Python 建模代码结合在一起。许多开源推理引擎（如 vLLM、SGLang 和 TokenSpeed）正在将 Rust 和 C++ 等语言整合到其技术栈中。在过去两年中，我们持续投入 Rust，并在性能和可维护性方面获得了巨大回报。通过与我们的 LLM 服务技术栈共享大部分嵌入实现，我们在吞吐量上也取得了提升，而无需在嵌入模型的维护上投入大量的工程精力。

随着模型的演进，我们将继续改进技术栈的每一层，以减少 CPU 密集型和 GPU 锁定的延迟。我们在 Ivy 和 Tulip 内采用的自定义 gRPC 协议使我们能够调整通信以减少网络延迟，而 ROSE 则为提高计算吞吐量奠定了基础。此外，随着生态系统中对无 GIL 线程 Python（Free-threaded Python）支持的增长，我们将能够进一步改善 Python 与 Rust 的互操作性，从而降低开销。

面向搜索的嵌入（Embeddings for Search）

Tulips、Roses 和一些 Ivy

关注算子之外（Paying Attention Beyond the Kernel）

Tulip

CUDA 图管理

延迟张量（Lazy Tensors）

ROSE

Ivy

……但算子（Kernel）依然至关重要

基准测试

低延迟嵌入（p50 / p90 / p99 / 最大值 ms）

低延迟评分（p50 / p90 / p99 / 最大值 ms）

高吞吐量嵌入（emb/s）

高并发嵌入（p50 / p90 / p99 / 最大值 ms）

结论与未来工作

参考文献