推論基礎設施
vLLM 在 Ironwood TPU 支援 16K embedding,StepPool 保存分段 prefill 狀態
Google 與 vLLM 為 Qwen3 embedding 模型加入張量對齊、編譯預熱及可跨分段累積的 pooling 狀態。官方在四顆 Ironwood TPU 上報告約 8.4 萬 token/s,但尚未提供同成本 GPU 對照或真實檢索品質測試。

Google Cloud 公開 vLLM-TPU 的長上下文 embedding 支援與可重現部署配方,目標模型包括文字版 Qwen3-Embedding-8B 及多模態 Qwen3-VL-Embedding-8B。這次工作的重點不是新增模型權重,而是讓原本以自回歸生成為核心的 serving 引擎,在 Ironwood TPU 上可靠執行 pooling、16K 級輸入與多模態 prefill。
第一個改動處理 TPU MXU 的張量整除限制:模型詞彙矩陣經 tensor parallel 切分時,vLLM 會加入硬件安全的 padding,避免 All-Gather 因形狀未對齊失敗。第二個改動強化 lazy loading,並在服務接收請求前預熱具分片感知的 JAX/XLA 編譯快取,降低多程序部署於首批請求才觸發編譯的風險。
更關鍵的是混合式 StepPool。長序列為免 HBM 在 prefill 階段耗盡,必須把輸入切成多段;但 embedding 要對整條序列完成 pooling,若中間聚合狀態只跟隨單一步驟,請求遭搶占或恢復時便可能遺失。新版把相關 metadata 移入 `CachedRequestState`,使聚合狀態能跨 chunk 累積並在 preemption 後延續。多模態版本目前只切分文字部分,影像路徑仍是需要觀察的容量邊界。
官方配方使用單一 2×2×1 Ironwood 節點、四顆 TPU、BF16、TP=4 及 16,384-token 隨機輸入。Google 部落格列出 83,996 token/s、5.13 request/s;GitHub 配方表則是 84,038.27 token/s、同樣 5.13 request/s,細微差異顯示數字可能來自不同跑次。精度驗證採 TPU 與參考後端向量的 cosine similarity,門檻為文字 0.999、多模態 0.995;這能檢查後端數值一致性,卻不能證明檢索品質、延遲尾部或成本優於 GPU。工程團隊下一步應以自身長度分布、圖文比例及 GKE 擴縮行為重跑測試。