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

什么是嵌入模型，主要用于哪些场景？

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

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

要理解嵌入模型在实践中的工作方式，可以考虑一个标准的向量搜索查询。当用户在语义搜索引擎中查询“cat”一词时，模型会将该 token 映射到稠密坐标空间。在这个向量空间中，“cat”与“feline”或“dog”之间的数学距离较短，因此具有较高的相似度分数；相反，“hat”或“car”等词尽管在拼写上看似相近，却会被映射到距离更远的坐标。

借助 vLLM-TPU 与 GKE 实现无缝弹性扩缩容

为原型应用部署小型文本嵌入模型较为直接，但当管道需要服务数百万次查询时，就会面临不同类型的生产环境瓶颈。我们最常见的问题包括：获取加速器的弹性容量，以便随动态流量波动无缝扩缩算力资源，以及提升成本/性能效率。

为克服这些扩缩容与容量约束，Google Cloud 已将原生 TPU 支持集成到 vLLM 中。vLLM 是业界标准、高度优化且广受欢迎的开源 LLM 服务引擎。基于 vLLM 统一 TPU 服务方案可为架构提供真正的弹性。工程团队能够通过直接配置 TPU 节点并配合其他 XPU 实例，动态扩缩服务容量。

图 2. 如果主 TPU 预留容量被完全使用，服务基础设施会自动回退到次级 GPU Spot 或按需资源池，而不会中断传入的推理流量。

借助 Google Kubernetes Engine（GKE）中的 Custom Compute Classes 等基础能力，组织可以基于严格的优先级规则自动进行节点扩缩容：当先前的容量类型或加速器不可用时，可扩容至不同的容量类型或加速器。

在 TPU 上实现高精度嵌入支持的工程实践

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

为将高维向量池化模型适配 TPU 硬件拓扑，我们选取 Qwen3 Embedding 系列作为目标工程模型，并在 TPU 上针对 vLLM 框架完成了多项关键优化。

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

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

挑战 B：实例化加固、TPU 懒加载与编译预热

vLLM 在 TPU 上依赖懒加载机制，以尽量缩短服务器冷启动延迟并降低主机内存峰值。为消除懒张量转换期间的模型初始化失败问题，我们在反量化管道中引入了属性提升（attribute promotion），使权重加载完全兼容 vLLM 的 TPU 懒加载器，实现零失败初始化。

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

挑战 C：长上下文 StepPool 架构

超长上下文要求池化层采用 Chunked Prefill，以防止高带宽内存（HBM）耗尽，但这会带来步骤边界之间状态丢失的风险。我们设计了混合式 StepPool，并将元数据迁移至 CachedRequestState，确保池化状态能够跨步骤正确累积，并在请求抢占后继续保留。

TPU 上的 vLLM 嵌入示例

下面是一个演示如何在 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

Python

Golden Reference 精度与数值一致性

为验证企业级精度，我们开展了严格的数学一致性评估，针对多语言和多模态数据集，将 TPU 输出与其他 XPU Golden Reference 输出进行比较。

为评估 TPU（vTpu）生成的稠密嵌入向量与 XPU（vRef）生成的参考基线向量之间的数值对齐情况，我们计算两者的余弦相似度：

当余弦相似度接近 1.0 时——文本的目标质量通过阈值为 ≥0.999，多模态输入为 ≥0.995——表明不同硬件后端之间具有近乎完美的数值一致性。这证实了在 vLLM-TPU 技术栈中实施的优化能够保持 Golden Reference 精度，而不会牺牲准确性。

如需查看生成成对计算结果的分步说明，请参阅 GitHub 上官方 AI-Hypercomputer Qwen3-Embedding-8B Recipes。

Qwen3-Embedding-8B（7K+ Token，TPU 对比 CPU 基线）

表 1. 在保持严格数值对齐的同时，TPU Ironwood 上运行 qwen-3-embedding-8b（bf16、16K+ 序列长度、TP=4）实现了 83,996 总 token/s 的吞吐量和 5.13 请求/s。

表 2. vLLM-TPU 仅对多模态 Prefill 中的文本部分进行分块。

公开配方与资源探索

为帮助开发者复现我们的数值一致性评估，并在 Google Cloud TPU 上快速部署嵌入工作负载，我们已在 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 配置 Cloud TPU 实例并查看硬件规格。

致谢

本文重点介绍的工程成果与跨硬件优化，得益于 Google Cloud 产品团队、工程团队及 vLLM 社区之间的卓越协作。

What is an Embedding Model and What is it Used For?

In modern AI architectures, embedding models serve as the foundational translators bridging raw unstructured data and downstream intelligent reasoning. Simply put, embedding models translate inputs of various types of data — including text, images, and audio — into dense vector math. These high-dimensional numeric arrays capture semantic relationships, powering critical enterprise capabilities such as semantic search, recommender systems, intent classification, personalized content discovery, and vector clustering.

Figure 1. Embedding models enterprise capabilities.

To understand how embedding models work in practice, consider a standard vector search query. When a user queries a semantic search engine for the word "cat", the model maps the token into a dense coordinate space. In this vector space, the mathematical distance between "cat" and "feline" or "dog" is short, yielding high similarity scores; conversely, terms like "hat" or "car" map to distant coordinates despite their orthographic similarity.

Seamless Elasticity with vLLM-TPU & GKE

While deploying small text embedding models for prototype applications is straightforward, scaling pipelines to serve millions of queries introduces different types of production bottlenecks. The most common ones we see are accessing elastic capacity of accelerators to seamlessly scale compute resources alongside dynamic traffic fluctuations and improving cost/performance efficiency.

To overcome these scaling and capacity constraints, Google Cloud has integrated native TPU support into vLLM — the industry-standard, highly optimized and popular open-source LLM serving engine. Standardizing on vLLM for TPU serving provides architecture true elasticity. Engineering teams can scale serving capacity up and down dynamically by provisioning TPU nodes directly alongside other XPU instances.

Figure 2. If primary TPU reservations are fully utilized, the serving infrastructure automatically falls back to secondary GPU spot or on-demand pools without interrupting incoming inference traffic.

By taking advantage of primitives like Custom Compute Classes in Google Kubernetes Engine (GKE), organizations can automate node autoscaling based on strict priority rules to scale up across different capacity types or accelerators if the previous one isn’t available.

Engineering High-Precision Embedding Support on TPUs

Serving next-generation embedding models in production demands processing ultra-long sequence contexts - ranging from 4K+ tokens for text workloads up to 15K+ tokens for multimodal text-and-image inputs. Crucially, enterprise applications require that these embeddings maintain strict mathematical parity and high precision across heterogeneous hardware backends compared to reference.

To bring high-dimensional vector pooling models to TPU hardware topologies, we took the Qwen3 Embedding model series as the target engineering models and engineered several key optimizations of the vLLM framework on TPU.

Challenge A: Hardware-Safe Tensor Alignment

TPU Matrix Execution Units (MXUs) impose strict divisibility constraints when sharding vocabulary matrices across topology meshes via Tensor Parallelism (TP). We implemented a unified, hardware-safe vocabulary padding strategy that guarantees exact tensor alignment during All-Gather execution.

Challenge B: Materialization Hardening, TPU Lazy-Loading & Compilation Pre-warming

vLLM relies on lazy-loading mechanisms on TPUs to minimize server cold-start latencies and reduce host memory peaks. To eliminate model initialization failures during lazy tensor transformations, we introduced attribute promotion within the unquantization pipeline, making weight loading fully compatible with vLLM’s TPU lazy-loader for zero-failure initialization.

Furthermore, to eliminate runtime JIT compilation latencies and avoid compilation traps in multi-processes deployments, we implemented sharding-aware pre-warming to lock JAX/XLA compilation caches prior to inference and stabilize the production pipelines and rollouts.

Challenge C: Long-Context StepPool Architecture

Ultra-long contexts require Chunked Prefill in the pooling layer to prevent High Bandwidth Memory (HBM) exhaustion, creating risk of state loss across step boundaries. We engineered a hybrid StepPool and migrated metadata to CachedRequestState , ensuring pooling states correctly accumulate across steps and survive request preemptions.

vLLM Embedding Sample on TPU

Below is a minimal example demonstrating how to initialize Qwen3-Embedding-8B on TPU. For complete setup scripts and environment deployment steps, refer to the official AI-Hypercomputer Qwen3-Embedding-8B Recipes on GitHub:

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

Python

Golden-Reference Precision & Numerical Parity

To certify enterprise-grade precision, we conducted rigorous mathematical parity evaluations comparing TPU outputs against other XPU golden references across multi-language and multimodal datasets.

To evaluate numerical alignment between dense embedding vectors generated on TPUs (vTpu) and reference baseline vectors generated on XPUs (vRef), we calculate their cosine similarity :

A cosine similarity score approaching 1.0 (with a target quality pass threshold of ≥0.999 for text and ≥0.995 for multimodal inputs) demonstrates near-perfect numerical parity across hardware backends. This confirms that optimizations implemented on the vLLM-TPU stack maintain golden-reference precision without sacrificing accuracy.

For step-by-step instructions on generating pairwise calculations, check out the official AI-Hypercomputer Qwen3-Embedding-8B Recipes on GitHub.

Qwen3-Embedding-8B (7K+ Tokens, TPU vs. CPU Baseline)

Table1. While maintaining strict numerical alignment, serving qwen-3-embedding-8b (bf16, 16K+ sequence length, TP=4) on TPU Ironwood achieved an impressive throughput of 83,996 total token/s and 5.13 req/s.

Table2. vLLM-TPU only chunks the text portion of multimodal prefill.

Public Recipes & Explore Resources

To help developers reproduce our numerical parity evaluations and rapidly deploy embedding workloads on Google Cloud TPUs, we have open-sourced official setup and execution recipes on the AI-Hypercomputer Public Repository.

Get Started:

*** Qwen3-Embedding-8B TPU Recipe:** Explore text embedding recipes at AI-Hypercomputer/Qwen3-Embedding-8B

*** Qwen3-VL-Embedding-8B TPU Recipe:** Explore multimodal embedding recipes at AI-Hypercomputer/Qwen3-VL-Embedding-8B

*** vLLM TPU Engine:** Explore more on vLLM framework on TPU at vllm-project/tpu-inference

*** Google Cloud TPU Portal:** Provision Cloud TPU instances and explore hardware specs at cloud.google.com/tpu

Acknowledgments

The engineering achievements and cross-hardware optimizations highlighted in this post were made possible through the incredible collaboration across Google Cloud Product, Engineering and the vLLM community.