Google 開源 TPU Sync：在 TPU 之間高速搬移 KV 快取的新軟體庫

KV 快取（鍵值快取）是語言模型在讀取提示詞時建立的工作記憶：對於每一個已處理的詞元（token），模型都會保存中間數值，這樣在產生每個新詞時就不必重新計算。對話或文件越長，這個快取就越大，而它通常存放在加速器晶片上稀少且昂貴的高頻寬記憶體（HBM）中。

Google 近日發布了 TPU Sync，這是一個開源的原生 C++ 函式庫，用於在其 TPU（張量處理單元）之間搬移並管理這種快取。根據 TechBytes 的報導，它概述了 Google Cloud 部落格上由工程師 Fenghui Zhang 和 Weiren Yu 撰寫的文章，這個函式庫還會在強化學習流程中同步模型參數。SemiAnalysis 指出，TPU Sync 以前名為 TPU-raiden，而 Open Source For You 則以 TPU Raiden 之名介紹該專案，並稱其採用 Apache 2.0 授權。

為什麼 KV 快取需要搬移？

Open Source For You 解釋說，語言模型的推論分為兩個階段：預填充（prefill），模型讀取整個提示詞；解碼（decode），模型逐個詞元地產生回答。分離式服務把這兩個階段放在不同的硬體池上執行，因此預填充結束後，KV 快取必須被送到負責解碼的晶片上。SemiAnalysis 指出，沒有傳輸函式庫的話，TPU 的外部服務只能停留在聚合模式，儘管 Google 多年來一直在內部為 Gemini 使用分離式服務。

第二個需求是卸載。SemiAnalysis 解釋說，當一顆晶片的 HBM 裝不下所有使用者的快取時，資料必須從 HBM 移到主機的 DRAM，再移回來。根據 The Next Platform 在 2026 年 11 月 9 日的報導，Google AI 與運算基礎設施副總裁暨總經理 Mark Lohmeyer 在電話會議上表示，在代理型工作負載中，TPU 有很大一部分時間處於閒置狀態。他說，TPU Sync 會在代理暫停時把工作記憶移出晶片，在代理恢復時再還原，這樣加速器就能騰出來處理其他任務，並降低每個代理型工作負載的成本。The Next Platform 還補充說，由於 TPU Sync 是軟體，Google 可以把它回溯移植到更早幾代的 TPU 系統上。

Google 宣稱的效能如何？

以下數字來自日期為 2026 年 11 月 9 日的 TechBytes，該媒體將其歸於 Google Cloud 部落格。我們未能讀到 Google 的原文，因此這些是廠商公布、由二手媒體轉述的數字。

• 根據 TechBytes，在參數規模為 2350 億至 4800 億的開放權重混合專家（Mixture of Experts）模型上，鍵值傳輸延遲降低超過 90%。
• 根據 TechBytes，首個詞元時間（TTFT，即回答的第一個字出現前的等待時間）在前綴快取和分離式服務的基準測試中，與標準配置相比平均下降 90%。
• 根據 TechBytes，raw_transfer 元件在 Ironwood TPU v7x 主機上，裝置與主機之間雙向的聚合吞吐量最高可達 223 GB/s，而 PCIe Gen5 的物理上限為 256 GB/s。
• 根據 TechBytes，在資料中心網路上，該軟體在一張 200 Gbit/s 網卡上可持續約 192 Gbit/s，在兩張網卡並採用 NUMA 感知條帶化時可達 350 Gbit/s。

TPU Sync 內部是如何運作的？

據 TechBytes 介紹，TPU Sync 依靠保持資料配置的直接記憶體存取（DMA，一種無需處理器參與即可複製記憶體的硬體機制），並搭配原地緩衝區別名技術。這個函式庫直接與 JAX 的 PJRT 緩衝區和 PyTorch 張量協作，從而避免了 Python 序列化和中間拷貝，而且傳輸發生在 XLA 編譯圖之外。SemiAnalysis 把同一思路描述為零拷貝傳輸，即透過提取 PJRT 緩衝區的原生硬體描述符來實現。

據 TechBytes 表示，執行階段由三個模組組成。全域鍵值儲存相當於一個登錄表，把提示詞雜湊對應到實體儲存區塊，並負責最近最少使用的淘汰。叢集控制器負責執行規劃和租約管理。本地快取管理器追蹤裝置上各區塊的狀態，管理鎖定的主機記憶體，並在 /dev/shm 的共享記憶體中保留一份熱啟動副本。記憶體層級涵蓋 HBM、鎖定的主機 DRAM、本地 SSD 以及 Google Cloud Managed Lustre。

在整合方面，TechBytes 列出了透過 TorchTPU 面向 PyTorch 的官方綁定、面向 JAX 的官方綁定、面向 vLLM、SGLang 和 llm-d 服務堆疊的連接器，以及面向 Pathways 和多控制器 JAX 流程的底層原語。

有哪些限制，成熟度如何？

TechBytes 表示，HBM 與主機 DRAM 之間的傳輸已可用於生產，而向 Managed Lustre 等共享儲存和外部 POSIX 後端的分層儲存仍處於實驗性驗證階段。強化學習的權重同步模組同樣是實驗性的，要達到最高網路吞吐量，需要兩張 200 Gbit/s 網卡並正確設定 NUMA 感知條帶化。Open Source For You 的文章沒有標註日期，它的說法更保守，稱該專案仍在積極開發中，尚不適合一般的生產使用，因此不同來源對成熟度的描述並不一致。

SemiAnalysis 提供了競爭視角。其 InferenceX 預覽版把聚合模式的 TPU v7 與聚合模式的輝達 B200 和 B300 作比較，並稱輝達的 GB200 和 GB300 NVL72 系統目前在分離式服務曲線的一部分上領先。待基於 TPU-Sync 的分離式服務在外部技術堆疊中最佳化後，將推出與分離式 TPU v7 的對比，SemiAnalysis 預期 TPU 上的 Mooncake Store 支援將建立在 TPU-Sync 原語之上。

這對企業意味著什麼？

對於在 Google Cloud 上執行或租用長上下文助理和代理的台灣企業來說，搬移 KV 快取變成了需要設定的事，而不是需要自己建構的事。實際步驟是：用自己的提示詞測試與 vLLM 或 SGLang 的連接器，測量啟用前後的首個詞元時間，並在依賴之前確認哪些記憶體層級仍被標註為實驗性。使用強化學習訓練的團隊應在自己的訓練器和取樣器拓撲上驗證權重同步模組。已公布的任何數字都不能取代在您自身工作負載上的基準測試。