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_len16384在 Ironwood TPU 上Google 报告的一个 16K 序列、BF16、TP4 配置达到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然后比较向量。importnumpyasnpdefcosine(a,b):anp.asarray(a)bnp.asarray(b)returnfloat(np.dot(a,b)/(np.linalg.norm(a)*np.linalg.norm(b)))Gateassertcosine(reference,candidate)0.999文本模型可以先用 0.999 作为高精度目标。多模态可按业务设不同阈值。为什么不能只看TopK一致率两个 Embeddingcosine 0.995在普通搜索里可能 TopK 还差不多。但当你的向量库有几千万条、而且很多候选距离很近时小偏差也可能改变排名。所以我会同时测Vector Parity TopK Overlap NDCG Business Recall一套完整精度回归embedding_gate:vector_cosine:p50:0.9999p01:0.9990top10_overlap:min:0.98retrieval_recall:regression_max:0.005不要因为吞吐翻倍就接受检索质量悄悄下降。Google公开配置里有几个值得照抄的测试维度示意配置fromvllmimportLLM llmLLM(modelQwen/Qwen3-Embedding-8B,runnerpooling,tensor_parallel_size2,max_model_len16384,max_num_batched_tokens512,dtypebfloat16,trust_remote_codeTrue,)真正需要调的是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个Bucket0—1K 1K—4K 4K—8K 8K—16K每个 Bucket 独立看P50 P95 Throughput Queue Time不要把短文本和长文档平均在一起。否则“平均延迟80ms”可能掩盖15K输入P95 4秒生产路由可以按长度分池例如Short Pool: 4K Long Pool: 4KKubernetesembeddingPools:short:max_tokens:4096replicas:8long:max_tokens:16384replicas:4流量由 Gateway 按 Token Count 路由。这样长请求不会把短查询堵住。Tensor Parallel为什么会遇到AlignmentTPU 的矩阵执行单元对切分维度有对齐要求。当词表矩阵做 Tensor Parallel Sharding 时如果维度不能整除拓扑就可能出现 All-Gather 对齐问题。Google 的做法是统一做硬件安全的 Vocabulary Padding。这类优化提醒我们跨硬件迁移 不是换一个device参数尤其是高维 Pooling Model。Compilation Pre-warming解决的是发布抖动JAX/XLA 类后端常见第一次请求触发编译线上就会出现新Pod Ready ↓ 第一批请求进入 ↓ P95突然几十倍更稳的是启动 →加载权重 →预热常见Shape →编译Cache稳定 →ReadyKubernetes Readiness 不应该只测HTTP 200而要等Model Loaded Compile Warmed Memory Stable才放量。长Context还要测Preemption故意把队列打满长请求执行一半 ↓ 被调度器抢占 ↓ 恢复最后向量必须和无抢占 Reference 一致。这是测试 StepPool / Cached State 最关键的一类 Case。一个Preemption TestInput: 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 precisionVector 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 才算从“能跑”进入“能稳定服务”。