Qwen3-Embedding上TPU：16K长上下文怎么稳住
Embedding 服务最容易被低估的性能问题，不是单条文本有多快。
而是：
输入突然从1K tokens
变成15K tokens以后，
系统还能不能稳定批处理？
Google 8月26日公开了 vLLM 在 Cloud TPU 上服务 Qwen3 Embedding 系列的一组工程实现。目标场景已经不是短文本向量化，而是：
文本：4K+ tokens
多模态：15K+ tokens
公开的 Qwen3-Embedding-8B 配置支持 max_model_len=16384；在 Ironwood TPU 上，Google 报告的一个 16K+ 序列、BF16、TP=4 配置达到 83,996 total token/s、5.13 req/s，同时使用跨硬件向量余弦相似度校验精度，文本阈值设为 ≥0.999，多模态阈值为 ≥0.995。
真正值得拆的是它为长上下文做的三件事：
Tensor Alignment
Compilation Pre-warming
Chunked Prefill + StepPool
为什么Embedding也会把HBM打爆
很多人对长上下文的第一反应是：
LLM KV Cache
但 Embedding Pooling 同样需要保存长序列中间状态、Pooling State 和 Batch Metadata。
如果 15K tokens 一次性 Prefill：
请求多
+
序列长
+
Batch大
HBM 压力会迅速上升。
所以需要：
Chunked Prefill
把一次长输入拆成多个 Step。
Chunked Prefill最危险的不是慢，而是状态丢失
假设 16K tokens 分成：
4K
4K
4K
4K
每一段执行完都要把 Pooling 所需状态保留下来。
如果请求被 Preempt、Reschedule 或 Batch 重排以后中间状态没有恢复，最终向量可能看起来“有结果”，但数学上已经偏了。
这比直接报错更危险。
所以长Embedding一定要做Golden Parity
不要只测：
API返回200
同一条输入分别跑：
Reference Backend
Candidate Backend
然后比较向量。
import numpy as np
def cosine(a, b):
a = np.asarray(a)
b = np.asarray(b)
return float(
np.dot(a, b) /
(
np.linalg.norm(a)
* np.linalg.norm(b)
)
)
Gate：
assert cosine(
reference,
candidate
) >= 0.999
文本模型可以先用 0.999 作为高精度目标。
多模态可按业务设不同阈值。
为什么不能只看TopK一致率
两个 Embedding：
cosine = 0.995
在普通搜索里可能 TopK 还差不多。
但当你的向量库有几千万条、而且很多候选距离很近时，小偏差也可能改变排名。
所以我会同时测：
Vector Parity
TopK Overlap
NDCG
Business Recall
一套完整精度回归
embedding_gate:
vector_cosine:
p50: 0.9999
p01: 0.9990
top10_overlap:
min: 0.98
retrieval_recall:
regression_max: 0.005
不要因为吞吐翻倍，就接受检索质量悄悄下降。
Google公开配置里有几个值得照抄的测试维度
示意配置：
from vllm import LLM
llm = LLM(
model="Qwen/Qwen3-Embedding-8B",
runner="pooling",
tensor_parallel_size=2,
max_model_len=16384,
max_num_batched_tokens=512,
dtype="bfloat16",
trust_remote_code=True,
)
真正需要调的是：
tensor_parallel_size
max_model_len
max_num_batched_tokens
它们共同决定：
吞吐
内存
调度
长序列稳定性
max_num_batched_tokens不是越大越好
Batch Token 上限大，吞吐可能提高，但也可能出现：
HBM上涨
尾延迟上涨
长请求挤压短请求
所以压测至少分：
512
1024
2048
4096
看：
Tokens/s
Req/s
P95
HBM
OOM
我会按长度建4个Bucket
0—1K
1K—4K
4K—8K
8K—16K
每个 Bucket 独立看：
P50
P95
Throughput
Queue Time
不要把短文本和长文档平均在一起。
否则“平均延迟80ms”可能掩盖：
15K输入P95 = 4秒
生产路由可以按长度分池
例如：
Short Pool:
4K
Kubernetes：
embeddingPools:
short:
max_tokens: 4096
replicas: 8
long:
max_tokens: 16384
replicas: 4
流量由 Gateway 按 Token Count 路由。
这样长请求不会把短查询堵住。
Tensor Parallel为什么会遇到Alignment
TPU 的矩阵执行单元对切分维度有对齐要求。
当词表矩阵做 Tensor Parallel Sharding 时，如果维度不能整除拓扑，就可能出现 All-Gather 对齐问题。
Google 的做法是统一做硬件安全的 Vocabulary Padding。
这类优化提醒我们：
跨硬件迁移
不是换一个device参数
尤其是高维 Pooling Model。
Compilation Pre-warming解决的是发布抖动
JAX/XLA 类后端常见：
第一次请求触发编译
线上就会出现：
新Pod Ready
↓
第一批请求进入
↓
P95突然几十倍
更稳的是：
启动
→加载权重
→预热常见Shape
→编译Cache稳定
→Ready
Kubernetes Readiness 不应该只测：
HTTP 200
而要等：
Model Loaded
+
Compile Warmed
+
Memory Stable
才放量。
长Context还要测Preemption
故意把队列打满：
长请求执行一半
↓
被调度器抢占
↓
恢复
最后向量必须和无抢占 Reference 一致。
这是测试 StepPool / Cached State 最关键的一类 Case。
一个Preemption Test
Input:
15,500 tokens
Baseline:
单请求运行
Stress:
并发64
触发Preemption
Check:
cosine >= 0.999
如果只做性能压测，不做向量一致性，这类 Bug 很难发现。
混合硬件弹性也要测结果一致
Google 的架构方向允许 GKE 在 TPU Capacity 不足时回退到其他 GPU Spot / On-demand Pool。
这意味着同一个服务可能出现：
请求A → TPU
请求B → GPU
所以硬件间 Parity 必须是发布门禁。
否则用户可能发现：
同样文本
上午一个向量
下午另一个向量
进而造成索引和查询 Embedding 不一致。
最危险的是Index和Query跑不同后端
比如：
离线索引：
GPU
在线Query：
TPU
如果两边向量偏差明显，Recall 会直接下降。
所以必须记录：
embedding_model_version
backend_profile
precision
Vector Metadata建议增加
{
"model": "qwen3-embedding-8b",
"model_version": "v1",
"dtype": "bf16",
"backend_profile": "tpu-ironwood-tp4",
"embedding_schema": "emb-v7"
}
后续迁移可以知道哪些向量需要重建。
一个真正实用的压测矩阵
Length:
1K / 4K / 8K / 16K
Concurrency:
1 / 8 / 32 / 64
Batch Tokens:
512 / 1024 / 2048
Backend:
TPU / GPU
共：
4 × 4 × 3 × 2
=
96组
每组输出：
Req/s
Token/s
P95
HBM
OOM
Cosine
Recall
这样才知道最优点在哪。
长上下文 Embedding 服务真正难的地方，是：
性能问题和精度问题会绑在一起。
Chunked Prefill、TP Sharding、JIT Compilation、Preemption 都可能提升系统效率，也可能制造很隐蔽的向量偏差。
所以发布标准可以压成两句话：
先证明向量没变，
再证明系统变快。
以及：
短请求和长请求分开看，
不要被平均值骗了。
做到这两点，16K Embedding 才算从“能跑”进入“能稳定服务”。
更多企业级 AI 应用、Agent、RAG 与模型工程化内容，我会继续整理在 智元界：
https://www.zyentor.com/