Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU

面向云TPU的长上下文多模态嵌入推理的高精度企业级解决方案 - Google Developers Blog

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

面向云TPU的长上下文多模态嵌入推理的高精度企业级解决方案

2026年8月26日

Anthony Su

软件工程师

Injae Kwak

产品经理

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

什么是嵌入模型？其用途是什么？

在现代AI架构中，嵌入模型作为连接原始非结构化数据与下游智能推理的基础转换器。简单来说，嵌入模型将文本、图像和音频等各类数据输入转换为稠密向量数学表示。这些高维数值数组捕捉语义关系，支撑着语义搜索、推荐系统、意图分类、个性化内容发现和向量聚类等关键企业级功能。

图1。嵌入模型的企业级功能。

为了理解嵌入模型的实际工作原理，考虑一个标准的向量搜索查询。当用户在语义搜索引擎中查询"cat"时，模型会将该标记映射到稠密坐标空间。在这个向量空间中，"cat"与"feline"或"dog"之间的数学距离较短，产生高相似度评分；相反，"hat"或"car"等词尽管在正字法上相似，但会映射到距离较远的坐标。

通过vLLM-TPU与GKE实现无缝弹性扩展

虽然部署小型文本嵌入模型用于原型应用较为直接，但将流水线扩展到服务数百万查询时会引入不同的生产瓶颈。我们最常见的瓶颈包括：访问加速器的弹性容量以实现与动态流量波动的无缝计算资源扩展，以及提升成本/性能效率。

为克服这些扩展和容量限制，Google Cloud已将原生TPU支持集成到vLLM中——这是行业标准的、高度优化且广受欢迎的开源LLM服务引擎。通过将vLLM标准化为TPU服务，可实现架构层面的真正弹性。工程团队可通过直接与其它XPU实例一起分配TPU节点，动态扩展服务容量。

图2。当主要TPU预留资源完全使用时，服务基础设施会自动回退到次级GPU现货或按需池，而不会中断正在进行的推理流量。

通过利用Google Kubernetes Engine（GKE）中的自定义计算类等原语，组织可以根据严格的优先级规则自动执行节点自动扩展，当先前资源不可用时，可跨不同容量类型或加速器进行扩展。

在TPU上构建高精度嵌入支持

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

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

挑战 A：硬件安全的张量对齐

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

挑战 B：Materialization 强化、TPU 延迟加载与编译预热

vLLM 依赖于 TPU 上的延迟加载机制，以最小化服务器冷启动延迟并减少主机内存峰值。为消除延迟张量转换过程中的模型初始化失败，我们在反量化流程中引入了属性提升，使权重加载完全兼容 vLLM 的 TPU 延迟加载器，从而实现零失败初始化。

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

挑战 C：长上下文 StepPool 架构

超长上下文需要在池化层使用分块预填充，以防止高带宽内存（HBM）耗尽，从而避免跨步骤边界的状态丢失风险。我们设计了一种混合 StepPool 架构，并将元数据迁移至 CachedRequestState，确保池化状态在各步骤间正确累积并存活于请求抢占场景中。

TPU 上的 vLLM 嵌入示例

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

from vllm import LLM

# 使用 vLLM 原生的池化运行器在云 TPU 上初始化 Qwen3-Embedding-8B
llm = LLM(
    model="Qwen/Qwen3-Embedding-8B",
    runner="pooling",             # 启用密集池化输出
    tensor_parallel_size=2,       # 跨 TPU 拓扑网格分片
    max_model_len=16384,
    max_num_batched_tokens=512,
    dtype="bfloat16",
    trust_remote_code=True
)

# 跨输入提取密集向量嵌入
prompts = ["Enterprise-grade semantic retrieval on TPUs with vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding

Python

已复制

金标准精度与数值一致性

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

为评估在 TPU 上生成的密集嵌入向量（v Tpu ）与在 XPU 上生成的参考基线向量（v Ref ）之间的数值一致性，我们计算它们的余弦相似度：

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

有关生成成对计算的逐步说明，请查看 GitHub 上的官方 AI-Hypercomputer Qwen3-Embedding-8B 配方。

Qwen3-Embedding-8B（7K+ 令牌，TPU 与 CPU 基线）

Table1. 在保持严格数值对齐的前提下，TPU Ironwood 上运行 qwen-3-embedding-8b（bf16，16K+序列长度，TP=4）实现了令人印象深刻的吞吐量，达到 83,996 个总令牌/秒和 5.13 请求数/秒。

Qwen3-VL-Embedding-8B（15K+ 令牌，TPU 与 XPU 基准对比）

Table2. vLLM-TPU 仅对多模态预填充的文本部分进行分块处理。

公共配方与资源探索

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

入门指南：

• Qwen3-Embedding-8B TPU 配方：在 AI-Hypercomputer/Qwen3-Embedding-8B 中探索文本嵌入配方

• Qwen3-VL-Embedding-8B TPU 配方：在 AI-Hypercomputer/Qwen3-VL-Embedding-8B 中探索多模态嵌入配方

• vLLM TPU 引擎：在 vllm-project/tpu-inference 中了解更多关于 TPU 上 vLLM 框架的信息

• Google Cloud TPU 门户：在 cloud.google.com/tpu 上预置云 TPU 实例并查看硬件规格

致谢

本文中强调的工程成就和跨硬件优化得以实现，得益于 Google Cloud 产品、工程团队与 vLLM 社区的卓越协作。

posted in:

• AI

• Cloud

• 公告

• 最佳实践

• 学习

• 探索

• 影响力

上一篇

下一篇

相关文章

列表

AI

案例研究

社区

使用深度学习和 Keras 解码宇宙信号

2026 年 8 月 27 日

移动

网络

公告

解决方案

介绍 Credentio：Google 开源的 C++ 库用于 C2PA 内容凭证

2026 年 8 月 13 日

云

最佳实践

使用 Google 的代理开发工具包构建零信任 AI 代理

2026 年 8 月 17 日

操作指南

学习

如何在 ADK 中评估实时与语音代理

2026 年 8 月 24 日

导航点