Google Cloud TPU 原生集成 vLLM：长上下文Embedding 推理优化实践

Google Cloud TPU 原生集成 vLLM：长上下文Embedding 推理优化实践 | 前端进阶之旅

Full translation

完整中文译文

什么是 Embedding 模型及其用途

在现代 AI 架构中，embedding 模型作为基础翻译器，桥接原始非结构化数据与下游智能推理。简而言之，embedding 模型将各种类型的数据输入——包括文本、图像和音频——转化为密集向量数学表示。这些高维数值数组捕获语义关系，为关键的企业级能力提供支撑，如语义搜索、推荐系统、意图分类、个性化内容发现和向量聚类。

要理解 embedding 模型在实际中如何运作，可以看一个标准的向量搜索查询例子。当用户在语义搜索引擎中查询单词"cat"时，模型会将该 token 映射到一个密集的坐标空间中。在这个向量空间中，"cat"与"feline"或"dog"的数学距离很短，获得较高的相似度得分；相反，诸如"hat"或"car"这样的词尽管拼写相似，却被映射到遥远的坐标位置。

vLLM-TPU 与 GKE 带来的无缝弹性扩展

虽然为原型应用部署小型文本 embedding 模型非常简单，但将流水线扩展到服务数百万级查询时会引入不同类型的生产瓶颈。最常见的有：访问加速器的弹性容量以随动态流量波动无缝扩展计算资源，以及提升成本/性能效率。

为克服这些扩展和容量约束，Google Cloud 将原生 TPU 支持集成到了 vLLM 中——这是行业标准、高度优化且流行的开源 LLM 服务引擎。以 vLLM 作为 TPU 服务的标准化方案提供了真正的架构弹性。工程团队可以通过直接配置 TPU 节点与其他 XPU 实例一起，动态地扩展或缩减服务容量。

通过利用 Google Kubernetes Engine（GKE）中的 Custom Compute Classes 等原语，组织可以基于严格的优先级规则自动执行节点自动扩缩容，在先前类型不可用时跨不同容量类型或加速器进行扩展。

在 TPU 上工程化高精度 Embedding 支持

在生产环境中服务下一代 embedding 模型需要处理超长序列上下文——文本工作负载从 4K+ token，到多模态文本加图像输入的 15K+ token 不等。关键在于，企业应用要求这些 embedding 在异构硬件后端上与参考实现保持严格的数学一致性和高精度。

为将高维向量聚类模型引入 TPU 硬件拓扑，我们以 Qwen3 Embedding 模型系列作为目标工程模型，并对 vLLM 框架在 TPU 上进行了几项关键优化。

挑战 A：硬件安全的 Tensor 对齐

TPU 矩阵执行单元（MXU）在通过 Tensor Parallelism（TP）将词汇矩阵分片到拓扑网格时施加了严格的整除约束。我们实现了一种统一的、硬件安全的词汇填充策略，保证在 All-Gather 执行期间精确的 tensor 对齐。

挑战 B：物化强化、TPU 懒加载与编译预热

vLLM 在 TPU 上依赖懒加载机制来最小化服务器冷启动延迟并降低主机内存峰值。为消除懒加载 tensor 转换过程中的模型初始化失败，我们在内反量化流水线中引入了属性提升，使权重加载与 vLLM 的 TPU 懒加载器完全兼容，实现零故障初始化。

此外，为消除运行时 JIT 编译延迟并避免多进程部署中的编译陷阱，我们实现了分片感知的预热机制，在推理之前锁定 JAX/XLA 编译缓存，从而稳定生产流水线和发布过程。

挑战 C：长上下文 StepPool 架构

超长上下文需要在聚类层中使用 Chunked Prefill 来防止高带宽内存（HBM）耗尽，但这会在步骤边界产生状态丢失风险。我们设计了一种混合 StepPool 并将元数据迁移到 CachedRequestState，确保聚类状态在多个步骤间正确累积并经受住请求抢占。

TPU 上的 vLLM Embedding 示例

以下是展示如何在 TPU 上初始化 Qwen3-Embedding-8B 的最小示例。完整的设置脚本和环境部署步骤请参阅 GitHub 上的官方 AI-Hypercomputer Qwen3-Embedding-8B Recipes：

from vllm import LLM

# Initialize Qwen3-Embedding-8B on Cloud TPU using vLLM's native pooling runner
llm = LLM(
    model="Qwen/Qwen3-Embedding-8B",
    runner="pooling",             # Enables dense pooling output
    tensor_parallel_size=2,       # Sharded across TPU topology mesh
    max_model_len=16384,
    max_num_batched_tokens=512,
    dtype="bfloat16",
    trust_remote_code=True
)

# Extract dense vector embeddings across inputs
prompts = ["Enterprise-grade semantic retrieval on TPUs with vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding

金标准精度与数值一致性

为认证企业级精度，我们进行了严格的数学一致性评估，在多语言和多模态数据集上将 TPU 输出与其他 XPU 金标准参考进行对比。

为评估在 TPU 上生成的密集 embedding 向量（vTpu）与在 XPU 上生成的参考基线向量（vRef）之间的数值对齐程度，我们计算它们的余弦相似度：

余弦相似度得分接近 1.0（文本输入的目标质量通过阈值为 ≥0.999，多模态输入为 ≥0.995）表明跨硬件后端实现了近乎完美的数值一致性。这证实了在 vLLM-TPU 堆栈上实现的优化在保持金标准精度的同时没有牺牲准确性。

有关生成分对计算的分步说明，请参阅官方 AI-Hypercomputer Qwen3-Embedding-8B Recipes。

Qwen3-Embedding-8B（7K+ Tokens，TPU vs. CPU 基线）

Qwen3-VL-Embedding-8B（15K+ Tokens，TPU vs. XPU 基线）

公开 Recipes 与探索资源

为帮助开发者复现我们的数值一致性评估，并在 Google Cloud TPU 上快速部署 embedding 工作负载，我们在 AI-Hypercomputer 公开仓库中开源了官方设置和执行 recipes。

Qwen3-Embedding-8B TPU Recipe：在 AI-Hypercomputer/Qwen3-Embedding-8B 探索文本 embedding recipes
Qwen3-VL-Embedding-8B TPU Recipe：在 AI-Hypercomputer/Qwen3-VL-Embedding-8B 探索多模态 embedding recipes
vLLM TPU Engine：在 vllm-project/tpu-inference 了解更多 vLLM 框架在 TPU 上的信息
Google Cloud TPU Portal：在 cloud.google.com/tpu 配置 Cloud TPU 实例并探索硬件规格

本文中展示的工程成就和跨硬件优化得益于 Google Cloud 产品、工程团队与 vLLM 社区之间的出色协作。