Google 與 Meta 推出 TorchTPU，為 Google TPU 提供原生 PyTorch 後端

簡而言之

TorchTPU 是 Google 與 Meta 合作開發的開源項目，為 Google TPU 推出原生且以原則為本的 PyTorch 後端。它讓模型可開箱即用地執行，並簡化遷移至 TPU 的流程；透過 eager 執行最佳化、整合 torch.compile 及原生硬件加速，提供具業界競爭力的推理和訓練效能。

推出 TorchTPU

Google TPU 一直成功用於在 Google 內外大規模訓練及部署大型 ML 模型。TPU 專為 ML 而設，因此在效能、能源消耗和成本方面都帶來顯著優勢。雖然 TPU 歷來與 TensorFlow 和 JAX 生態系統緊密相關，但我們很高興推出 TorchTPU。這是 Google 於 2025 年年中啟動的新計劃，並與 Meta 的 PyTorch 團隊合作，旨在讓 PyTorch 原生地在 TPU 上運行。

要在 Python 框架以下的層級編寫 Google TPU 程式，開發者會使用 XLA 編譯器堆疊。2020 年，Google 推出首個基於 Lazy Tensors、可在 TPU 上運行的 PyTorch 編譯版本 PyTorch/XLA。這項初期整合雖然奠定了基礎，卻難以跟上 PyTorch 由社羣推動的快速演進。尤其是，PyTorch/XLA 必須侵入式修改原始碼，明確指定可編譯區域；編譯延遲嚴重；不支援 PyTorch 原生的命令式（「eager」）執行模式；而且支援 PyTorch 生態系統中新興分散式運算原語的進度落後。與 PyTorch/XLA 和 TorchAX 等以往將 PyTorch 硬套進 JAX 運作模式的方案不同，TorchTPU 從一開始便以原則為本，打造 PyTorch 的 TPU 後端：eager 執行是首要考量，現有 PyTorch 工作負載只需極少修改程式碼即可運行（例如 .to("cuda") → .to("tpu")）。

具體而言，TorchTPU 啟動之初，我們便刻意制定了一系列設計原則：

• PyTorch Citizenship:
  • 執行階段不依賴 JAX；
  • 提供完善的 eager 模式；
  • 以 Torch Dynamo 為基礎（而非 Lazy Tensor）的 Torch.compile()，效能與 JAX compile 相若；
  • 廣泛支援 PyTorch 分散式運算慣用模式（DDP、FSDP、DTensor 等），並在 PyTorch 堆疊的最低層提供 TPU 集合通信操作（例如 Process Group API）；
  • 與 PyTorch 生態系統無縫整合；
  • 促進社羣參與。
• Focus on the User:
  • 將工作負載移植至 TPU，應該只需將 model.to("cuda") 改為 model.to("tpu")。 （僅作這項簡單變更，無法達到 TPU 的峯值效能，因為模型的張量形狀、分片方式和運算可能偏向 GPU 架構；不過，模型應該可以開箱即用地運行。）
  • 針對執行階段或編譯器錯誤，提供易於理解且可採取相應行動的訊息；
  • 整合 PyTorch 社羣已在使用的工具和框架（例如 torch.profiler、DTensor、Ray）。
• Improve the State of the Art:
  • 打造能在大規模工作負載上與 GPU 上的 PyTorch 競爭的方案；
  • 為 eager 執行引入 JIT 編譯方案，並透過執行階段最佳化，縮小逐個運算子執行與編譯執行之間的效能差距；
  • 透過引入有界動態性、避免因形狀變更而觸發 XLA 重新編譯，以及探索可預先編譯的動態 TPU kernel，改進 TPU 程式設計堆疊。

TorchTPU 正在積極開發中，已用於部分規模最大的 Google Cloud 客戶的推理和訓練工作負載，並以開源專案形式在這個 GitHub 儲存庫提供。截至本文撰寫時，我們已：

• 支援即時執行和編譯兩種模式，包括全新的 DeferAndFuse 即時執行模式；
• 支援約 450 個 Aten 算子（略過具有結構式或複合式委派的算子），包括主要深度學習架構所需所有 Core ATen 算子的前向及反向傳播；
• 支援透過 JAX/Pallas 和 Helion DSL 編寫自訂核心；
• Validated out-of-box execution on the 5,000 most downloaded Hugging Face models, including:
  • 大型語言模型：Qwen、DeepSeek、Gemma、LLama、Mistral、GPT、OPT、BLOOM、Falcon、StableML
  • 音訊／語音：Whisper、Wav2Vec2、HuBERT、MMS
  • 視覺：DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet、DETR
  • 文字編碼器及搜尋：BERT、ALBERT & Electra、RoBERTa、DeBERTa、T5、MT5、Flan-T5
• 已驗證以 DTensor 為基礎的分片；
• 已驗證多主機推理（vLLM 和 SGLang）及微調／訓練工作負載（TorchTitan），並以 Ray 進行協調；
• 已驗證與 vLLM、SGLang、TorchTitan、Ray 和 TorchRec 的整合。

本文其餘部分將探討 TorchTPU 架構的主要方面及一些初步結果。

架構深入探討

TPU 入門

Google TPU 是專為執行深度學習基礎所需的大規模張量運算而設計的專用專用集成電路。從硬件架構角度來看，TPU 晶片包含（視乎世代而定）：

• 一個或多個TensorCore，每個均由一個或多個矩陣乘法單元（MXU）、一個向量單元（VPU）和一個純量單元組成（參閲此連結）。視乎 TPU 世代，MXU 可包含 256×256 個乘加器，並排列成脈動陣列。
• 一個或多個SparseCore，這些是用於稀疏運算（例如 gather、scatter、embedding）和集體通信操作的專用單元。它們可以與 TensorCore 並行運作，讓計算與計算、計算與通信互相重疊。
• 大量高頻寬記憶體（HBM）——TPU 8i 配備 288 GB，頻寬為 8.6 TB/s，可用於儲存機器學習數據並與 CPU 共用。
• 雙向晶片間互連（ICI）——TPU 8t 和 8i 的傳輸速度為 19.2 TBps，可讓用戶在配置的 TPU Slices 內快速傳輸資料。TPU Slices 是 TPU Pod 的邏輯子組，由預先配置數量的 TPU 晶片組成，並採用 3D torus（TPU 7、TPUv 8t）或 Boardfly（TPU 8i）配置互連。
• 雙向數據中心網絡（DCN）——TPU 8t 的傳輸速度為 400 GBps，可讓不同 Slice 之間的 TPU 晶片互相通信。

TPU 透過兩個不同入口，使用 XLA 編譯器堆疊進行程式設計：StableHLO（用於機器學習運算的較高層級 IR）和 Mosaic（用於自訂核心的較低層級IR）。實際的 TPU ISA 屬於專有技術，可能會變更，亦未公開提供。實際上，可以把 StableHLO（尤其是 Mosaic）視為相當於 PTX。下圖展示了高層次概覽。

C++ 程式可透過建立 PjRT Client 存取本機 TPU。其 API 提供各種方法，可將 MLIR（StableHLO 或 Mosaic）編譯成 TPU 可執行檔、啟動該可執行檔，以及向 TPU 複製資料或從 TPU 複製資料。請注意，TPU 的執行是非同步的；CPU 可等待相關的 PjRTBuffer，直至 TPU 結果準備就緒。

TPU 程式設計模型的主要特點：

• TensorCore 裝置可視為一個具備一個或多個大型矩陣乘法單元（MXU）的單執行緒處理器。因此，要達至最高效能，模型的形狀必須足夠寬，才能充分發揮 MXU 的潛力。這與 GPU 不同：在 GPU 上，平行啟動多個較小的核心，有助盡量提高硬件使用率。
• XLA 編譯器可指示 TensorCore 將稀疏及集體運算卸載至 SparseCore，以更有效率及／或平行的方式執行。這可以自動進行，亦可以透過用戶在已編譯的 StableHLO 模組中提供 MLIR 註釋來進行。
• TPU 裝置只能由一個用戶空間程序存取。為了實現與 PyTorch 生態系統的無縫整合，我們的分散式 TorchTPU 工作負載採用每個裝置一個程序的模型，與 JAX 的每部主機一個程序模型不同。不過，由於 TPU 只能由一個用戶空間程序存取，RL 共置無法開箱即用；我們正研究放寬或繞過這項 TPU 限制的方法。
• StableHLO 遵循 SSA 範式，因此不允許原地修改，即寫入操作總會產生新值（至少在概念上如此）。這對 TorchTPU 尤其重要，因為 TorchTPU 透過視圖分解及反演，提供熟悉的 Pytorch「view」語義，包括修改操作，詳情見本文稍後的説明。
• XLA 編譯器和 PjRT 執行環境假設張量形狀在編譯時已知。我們認為這對 PyTorch 工作負載構成限制，因此一直在加入對有界動態性的支援：形狀維度的上限在編譯時已知，但實際維度只在執行時才知道。
• 集體運算會以 StableHLO 運算的形式降階至 XLA，而非使用 NCCL 等獨立程式庫。

易用性與自動載入機制

我們在設計 TorchTPU 時已考慮易用性。在已安裝 TorchTPU wheel 的 TPU 電腦上，用戶匯入 PyTorch 後，只需作簡單修改，透過數行程式碼即可在 TPU 上執行 PyTorch 程式碼，這有賴 PyTorch 的自動載入機制：

下表列出並列比較。

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda")
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.device("tpu")
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator()
    dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

數值精度

TorchTPU 為 TPU 硬件調整了標準 PyTorch 精度 API。與所有高效能機器學習架構一樣，TPU 會針對其特定世代偏好的原生數據類型進行最佳化，以發揮最高效能和能源效益。TPU 原生的矩陣乘法格式為 bfloat16 (BF16)，兼具高動態範圍與高效吞吐量。TPU 矩陣乘法單元 (MXU) 專為使用 BF16 輸入進行運算，並以 32 位浮點數作中間累加而設計，因此這種配置是 TPU 高效能運算的基石。至於非收縮類操作（例如逐元素加法、啟動函數、層正規化和 softmax），TorchTPU 會透過 TPU 向量執行單元 (VPU) 執行 FP32 張量操作，而 VPU 原生支援完整的 IEEE-754 FP32 算術運算。

TorchTPU 會將 PyTorch 精度模式對應至底層 XLA/StableHLO 配置，以發揮這項硬件能力：

• medium：乘法採用較低精度（將輸入轉換為 MXU 原生的 BF16），累加則以 FP32 執行，以維持穩定性。TorchTPU 預設使用此模式，以充分發揮 TPU 硬件加速效能；單次 BF16 乘法配合 FP32 累加，對幾乎所有深度學習工作負載而言都能提供足夠精度。
• high：使用 3 遍算法，將輸入拆分為高位和低位部分，以計算三個交叉項；計算成本約為 3 倍，並可恢復大部分 FP32 精度。
• highest：使用 6 遍算法，計算所有交叉項，以近似完整 IEEE FP32 精度；MXU 的計算開銷約為 6 倍。

對絕大部分深度學習工作負載，我們建議使用 medium（預設）精度。在 TorchTPU 中明確選用 high 或 highest 會觸發效能警告（TORCH_WARN_ONCE）。只有在明確發現與精度相關的收斂問題時，才應改用 high 或 highest，因為這些模式會帶來顯著的效能開銷。除了全域設定外，TorchTPU 亦提供執行緒區域的內容管理器（torch.tpu.precision），以便精細控制適用範圍（例如用於對數值較敏感的層）。

以下表格列出並列比較。

Eager 執行、StableHLO 降階及 JIT 編譯

按照這份 PyTorch 指引，我們擴充了 PyTorch 分派器，加入一個新的 C++ 後端，並以 PrivateUse1 註冊。此外，我們使用標準的 at::Tensor 機制，定義張量儲存空間的裝置專屬實作（使用 c10::StorageImpl 內的 c10::UniqueVoidPtr）。

TorchTPU 的 eager 執行有一項獨特之處，就是可以透過多種模式將 Aten 運算派發至 TPU：

• 即時派發，實際上就是逐個運算派發；
• 延後至自動判定的物化點（例如 .to("cpu") 陳述式），並與截至該點延後的其他運算一併彙集和編譯；
• 以及介乎兩者之間的模式。

這種彈性無疑會令設計更複雜，但我們從一開始便刻意加入這項彈性，因為我們希望試驗不同策略，提升 eager 效能，使其勝過基本的逐個運算派發模式。（此外，正如本文稍後所述，這種彈性也方便處理 view 語義和 torch.compile。）

總括而言，TorchTPU 提供 3 種 eager 模式，可透過環境變數或上下文管理器設定：

• Strict（預設）：每次派發一個 Aten 運算，並由 TPU 與 CPU 非同步執行。這相當於預期在 GPU 上以 eager 模式執行 PyTorch 時的行為。
• Debug（export TORCH_TPU_LAUNCH_BLOCKING=1）：每次派發一個 Aten 運算，但 TPU 與 CPU 同步執行。這種模式速度較慢，主要用於除錯；其效果相當於在 GPU 上使用 CUDA_LAUNCH_BLOCKING，並額外支援為 TorchTPU 啟用界限檢查。
• DeferAndFuse（export TORCH_TPU_DEFER_AND_FUSE=1）：將多個 Aten 運算延後處理，並融合成較大的區塊。這是 TorchTPU 引入的一種全新 eager／compiled 混合模式，速度可比 Strict Eager 快 50% 或以上，同時減少 HBM 峯值記憶體用量（運算融合可減少 HBM 中的暫存空間）。

介紹完 TorchTPU eager 模式的基本概念後，以下簡述 PyTorch 分派器呼叫回呼時會發生甚麼事。假設我們使用 DeferAndFuse 模式，並以 at::mul.out 為例：

• Create a Deferred Op instance that contains:
  • A StableHLO builder object, which when invoked returns a StableHLO function implementing the semantics of aten::mul, as shown below (note that any implicit broadcasting and dtype promotion are handled elsewhere)

    func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> {
      %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16>
      return %0 : tensor<4x16xbfloat16>
    }

  • 可選的靜態運算參數映射表（例如 aten::native_batch_norm 的 training 和 momentum），用於計算快取鍵（下文會討論）；
• 返回新的 at::Tensor，其中 c10::StorageImpl 會（透過間接引用）指向該 Deferred Op 實例。

收到 PyTorch 分派器傳來、對應於 mytensor.to("cpu")、mytensor.item() 或類似運算的回呼後，TorchTPU 會：

1. 分析目前為止累積、並導致計算 mytensor 的延後運算區域 R；
2. 進行活躍性分析，判斷除了產生 mytensor 的運算外，R 內還有哪些運算需要物化。值得注意的是，具有副作用的運算（例如 collective 運算）一律會保留，其呼叫順序亦然；至於在 R 以外不再存活的運算，由於 Python 的記憶體管理採用參照計數，因此可在不同工作負載副本之間以確定性方式識別。
3. 為要物化的運算圖 G 建立編譯快取鍵，以便日後快速且唯一地識別該運算圖；
4. Use the computed cache keys to check in a Compilation Cache if graph G has been already compiled.
   • 若有，便在 TPU 上派發 G 的可執行程式；
   • 若沒有，則使用執行緒集區以非同步方式編譯 G，然後在 TPU 上派發執行。
5. 等待 TPU 上 mytensor 的值就緒，然後將其從 TPU HBM 複製至與目的地 at::Tensor 關聯的 CPU 儲存空間。

表面看來，上述方法似乎效率不高，因為只有到達 .to("cpu") 陳述式時，CPU 才會開始執行 TPU。對於實際工作負載，我們可以預期，在最初幾次迭代／步驟填滿編譯快取後，CPU 執行時只需進行 PyTorch 追蹤及運算子派發，因此始終會領先 TPU 執行；此時 TPU 正忙於進行耗時的數值運算。儘管如此，為了加快多次重複運算子序列的物化速度（例如一疊 Transformer 區塊），我們加入了「重複運算子啟發式」，用以偵測這類重複序列，並在到達 TPU 至 CPU 傳輸之前觸發物化點。

編譯快取

TorchTPU 使用編譯快取，避免不必要的 XLA 編譯。與現代電腦儲存系統相似，此快取由多個互相協作的層級組成：

• 第 1 層：程序本機的記憶體內快取，用於在同一程序內共用編譯結果。這是三個層級中速度最快的。
• 第 2 層：機器本機的共享記憶體，可在同一主機上的不同程序之間共用。這適用於多裝置訓練／推理。
• 第 3 層：由遠端檔案系統（例如 Google Cloud Storage）支援的持久快取，可在不同主機、程式及工作重新啟動後重用。這適用於在 Colab 和數據中心等執行環境中執行 TPU 程式。

第 1 層快取一律啟用。預設情況下，第 2 層會在分散式程式中啟用，在單裝置程式中則停用。用戶可按需要覆寫此設定。第 3 層預設停用，但可透過指定共享儲存路徑明確啟用。

快取中的每個項目都是 StableHLO 計算圖的已編譯 TPU 可執行檔。在第 1 層快取中，每個可執行檔都是 C++ 物件，可直接排程到 TPU 裝置執行。由於物件不能直接在不同程序之間共用，第 2 層和第 3 層快取會以序列化二進位字串形式儲存可執行檔。

基於效能考慮，我們不會使用 StableHLO 圖本身作為快取鍵。相反，我們會將計算、編譯器選項及主要元件版本編成指紋，並轉換為三個 64 位元整數，再以這個元組作為快取鍵：

• 無形狀鍵，用以記錄計算圖的結構，但不包含張量輸入／輸出的大小；
• 維度鍵，用以記錄張量大小；以及
• 編譯選項鍵，用以記錄傳遞給 XLA 編譯器的選項，以及 TorchTPU、libtpu 和 PjRt runtime 的版本。

無形狀鍵和維度鍵合稱為圖鍵。理論上，不同圖、編譯器選項及元件版本之間可能出現快取鍵碰撞，但我們認為 192 位元的鍵空間足夠大，實際上不會發生這類碰撞。

檢視語義

PyTorch 的操作語義與 StableHLO 語義之間其中一個明顯差異，在於處理可變性及共享記憶體檢視的方式：

• PyTorch 假設每個 Tensor 都是指向線性位元組陣列的指標，並在 Tensor 層級附有步幅中繼資料，用以描述如何將邏輯索引元組轉換為該陣列中的線性偏移量。多個 Tensor 可以共用同一個底層位元組陣列，因此可以透過任何一個 Tensor 執行原地寫入，而寫入結果也會在其他 Tensor 中顯示。由於附有步幅中繼資料，分派時只需更新該中繼資料，便可在 CPU 上低成本地執行檢視操作。雖然大多數這類檢視都屬於少數幾種常見操作，但也可以使用 as_strided 運算子定義任意複雜的模式。
• 相較之下，StableHLO 是純函數式語言，所有張量值在建立後都不可變。實體配置對用戶而言是抽象的，並由編譯器作出的分塊決策所決定，因此不支援直接操作步幅。此外，StableHLO 不支援原地寫入；相反，可以透過 dynamic_update_slice 操作，將一組新值合併至現有張量的子切片中，藉此選擇性地更新內容。

在有寫入操作的情況下，要在 TPU 上實現正確的檢視語義並不簡單；不過，TorchTPU 透過一些機制，包括檢視分解、檢視反演和延遲操作，成功實現了這一點。這座語義橋樑讓 TorchTPU 原生支援 PyTorch 的一些模式，否則這些模式將無法實現，例如供分散式訓練使用的串接參數陣列。

串流與 SparseCore

最佳化模型效能時，PyTorch 用戶可能會希望同時執行以下 3 類操作：

• 計算操作；
• 集體操作；
• 主機與裝置之間的資料傳輸。

在 GPU 上，通常會使用不同的命令佇列來達成這一點，讓裝置將工作分派至不同的硬件單元（SM、DMA 等）。在 PyTorch 中，裝置命令佇列以串流的形式提供，並常見於用戶程式碼及框架程式碼中（例如 FSDPv2）。

TorchTPU 也實作了 PyTorch Stream API，但並非所有情況都能實現 GPU 上可用的相同功能。由於 TPU TensorCore 本質上是單執行緒執行處理器，因此每個 TensorCore 一次只能分派一項計算。相反，GPU 包含許多相同的計算核心（SM），可以平行運作。幸好，TPU TensorCore 配有一個或多個 DMA 引擎，以及一個或多個 SparseCore 引擎。DMA 引擎可在 TensorCore 執行操作時，同時用於主機與裝置之間的資料傳輸，這一點並不意外。更值得注意的是，SparseCore 引擎是獨立單元，專為稀疏操作（例如 gather/scatter）而最佳化，也支援集體操作。原則上，它可以支援完整的 StableHLO 操作集；但要注意的是，在稀疏運算方面，其效能會勝過 TensorCore，而在密集運算方面則不及 TensorCore，因為 TensorCore 內的 MXU 正是為此所長。為了讓用戶全面掌控工作的執行位置，TorchTPU 串流可以手動固定至 SparseCore；無論在效能還是操作支援完整性方面，這種串流都不會與在 TensorCore 上執行的預設串流相同。串流也可用來協調非同步資料傳輸（使用 DMA 引擎），方式與 GPU 上相似。

編譯模式

我們的首要目標是維持與其他 torch.compile() 後端一致的用戶體驗，因此共用相同的編譯器前端，即 Dynamo。根據這份 PyTorch 指引，我們會註冊一個新的 TorchDynamo 後端，使用方法如下：

compiled_model = torch.compile(model, backend="tpu")

Dynamo 會將 torch.fx.GraphModule 及一系列範例輸入張量傳遞給我們的後端，以呼叫該後端。在這個階段，我們會如同處於即時執行模式般追蹤所提供的計算圖，但所有運算子都會延後執行，並且不會在 TPU 上實際執行。這讓我們可以為整個 FX 計算圖生成並編譯單一 StableHLO 程式。最後，我們的 Dynamo 後端會傳回一個 Python 可呼叫物件，用以包裝對生成的 TPU 可執行程式的呼叫。請注意，作為一項暫時限制，在 TPU 機器上呼叫 torch.compile 時，我們會自動設定 dynamic=False，因為支援 SymInts 需要有界動態性（下文會討論），但目前尚未預設啟用。

有界動態性

TPU 與 GPU 的一項主要差異是，XLA 傳統上會使用靜態形狀來生成經最佳化的硬件佈局和融合核心。實際上，這並不是嚴重限制；當張量維度（例如批次大小或序列長度）在工作負載執行期間有所變化時，開發者通常會預先配置較大的張量，並在計算前後手動填充或切片，以解決問題。

為避免手動修改程式碼及重新編譯的開銷，TorchTPU 正積極開發透過 XLA 有界動態性原生支援動態形狀的功能。使用有界動態性時，動態緩衝區會在主要模型模組內配置至最大上限，並且只需編譯一次計算圖。計算會使用有界緩衝區在 TPU 上執行，而 XLA 會追蹤邏輯維度大小，並自動處理填充和切片，不會觸發重新編譯。

核心

我們區分兩類核心：

1. PyTorch 框架提供的核心，例如 torch.nn.functional.scaled_dot_product_attention。
2. 用戶提供的自訂核心，用於加速模型的特定部分；這類核心通常是特定於模型的；

對於 1)，我們起初採用一種做法：以 Pallas 編寫 TPU 核心，然後連同符號形狀匯出至 Mosaic，並在 TorchTPU 中加入一個步驟，在呼叫核心時針對其形狀進行特化。這種機制的靈活性不及預期，因此我們正試驗一種新的 Mosaic 核心編寫方式，以支援動態形狀。

對於 2)，我們已開發三種方案：

• (Production ready) The user writes their TPU kernel in Pallas and integrates it with PyTorch via a TorchTPU API, as follows

  # Define a simple vector addition Pallas kernel
  def add_vectors_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]
  
  # Wrap the Pallas kernel into a JAX function
  @torch.tpu.pallas.jax_op("pallas::add_vectors")
  def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
    out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
    return pl.pallas_call(add_vectors_pallas_kernel, out_shape)(x, y)
  
  # Use it like a standard PyTorch function
  x = torch.tensor([1.0, 2.0, 3.0], device="tpu")
  y = torch.tensor([4.0, 5.0, 6.0], device="tpu")
  z = add_vectors(x, y)

• （可供正式環境使用）用戶使用來自 Tokamax 的核心程式。Tokamax 是一個高效能 TPU 核心程式庫，現已支援與 PyTorch 無縫互通。Tokamax 為多種關鍵運算子提供高度最佳化的實作，包括混合專家模型的構建元件，例如分組矩陣乘法和 token 收集／散佈；注意力機制變體，例如分頁注意力、gated deltanet、Kimi delta attention 和多頭潛在注意力；以及 TopK 計算。運算子組合會隨時間演進，以支援在最新硬件加速器上執行的最新開源模型（例如在推出時支援 TPU v8i）。這些核心程式由穩健、整合完善的基礎設施支援，例如自動調校（例如搜尋分塊大小）、以 XProf 為基礎且開銷低的基準測試，以及回歸測試，以減輕操作上的複雜性。核心程式本身以 JAX/Pallas 編寫，而程式庫的基礎設施則採用 torch.tpu.pallas.jax_op API，以透明方式管理互通性，並提供原生 PyTorch API，方便開發人員使用。
• (Experimental) The user rewrites their GPU kernel in Helion, which enables targeting multiple accelerators with the same PyTorch-based DSL, as shown below

  # Define the kernel directly in Helion
  @helion.kernel(backend="pallas")
  def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    # Device execution loop over automatically partitioned tiles
    for tile in hl.tile(out.size()):
      out[tile] = x[tile] + y[tile]
    return out
  
  # Use it like a standard PyTorch function
  x = torch.tensor([1.0, 2.0, 3.0], device="tpu")
  y = torch.tensor([4.0, 5.0, 6.0], device="tpu")
  z = add_vectors(x, y)

分散式支援

為確保與現有 PyTorch 分散式 API 和框架相容，我們決定在 Process Group API 層級實作 TPU 集合通訊，方法是繼承 c10d::Backend；這樣便能在即時執行和編譯模式下攔截集合運算，並無縫支援 DTensor、FSDPv2、SimpleFSDP、FairScale 等 PyTorch 分散式框架。與 GPU 上的 PyTorch 類似，我們採用（並實際強制執行）一個程序對應一個裝置的方式。從實作角度來看，TPU 堆疊沒有提供獨立的 TPU 對 TPU 通訊程式庫；相反，TPU 集合運算會表示為 StableHLO 運算（stablehlo.all_gather、stablehlo.all_scatter、stablehlo.all_to_all 等），這與上述機制自然配合。

在 Process Group 層級攔截集合運算，幾乎已讓我們大功告成，但仍有幾項 TPU 相對於 GPU 的明顯特殊之處：

• 不同的初始化程序。假設主機連接了多個裝置。使用 GPU 時，由 torchrun 啟動的程序可以看到所有已連接的裝置，即使它只會使用其中一個（例如使用 device=f"cuda:{local_rank}"），並可在呼叫 torch.distributed.init_process_group() 前將模型移至該裝置。然而，使用 TPU 時，程序只能看到一個裝置，亦不知道主機連接了多少個裝置。要正確初始化 TPU 執行階段，程序必須在呼叫 torch.distributed.init_process_group 前知道：(a) 是否打算以分散式方式執行；(b) TPU 數量（world size）；以及 (c) TPU 類型（因為這會影響拓撲）。
• Default Process Group Knowledge. Consider two subgroups (4, 0, 1, 3) and (5, 2, 6, 7) running an all_reduce among themselves. Note that torch.distributed.all_reduce takes as input only the subgroup in which the collective op must execute. In contrast, stablehlo.all_reduce requires information about all subgroups (i.e., PyTorch’s default process group), even those the current process is not part of, as shown in the following MLIR snippet:

  %0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({...

• MPMD Support. Historically, the TPU stack was built with SPMD in mind. Thanks to tools such as Shardy, distributed JAX programs are written for a single host, with no communication code included, and then are sharded at massive scale with user-provided directives, while the XLA compiler is fully responsible for inserting collective ops at the right points while maximizing the overlap between computation and communication. Unfortunately, the PyTorch ecosystem doesn’t adopt a clean SPMD approach (a notable exception is SimpleFSDP) and different processes may execute different ML ops leading to a collective op for reasons such as:
  • 由用戶插入、以程序 rank 為條件的語句；
  • 當張量的某個維度不能被裝置數量整除時，DTensor 對共享張量進行不均等切分。

  當包含集合運算的 TPU StableHLO 模組彼此不相同時，可能會發生不可預測的死鎖或記憶體損毀（這是因為產生的 TPU 程式碼中，各 rank 的同步旗標 ID 或緩衝區偏移量可能不一致）。TorchTPU 會按以下方式處理：

  • （預設）Skinny Collectives。預設情況下，TorchTPU 會將每個 collective 操作各自編譯成獨立的「精簡」StableHLO 模組，確保各個 rank 的緩衝區和同步旗標 ID 一致。這種做法可行，但相較於下一個項目所述的做法，可能會造成明顯的效能損失（例如端到端效能下降 30-40%），因為這無法讓 XLA 編譯器重疊運算與通訊。目前正努力將 Skinny Collectives 放在固定於 SparseCore 的獨立 stream 上執行，讓它們可以與 TensorCore 上的運算重疊。我們預期這項最佳化可彌補 Skinny Collectives 造成的大部分效能損失。
  • （可選）@spmd_safe 裝飾器。如果用戶知道某段程式碼是 SPMD（例如，沒有以 torch.distributed.get_rank() 為條件的 if 陳述式），便可使用 TorchTPU 的 @spmd_safe 裝飾器標記該段程式碼。在 torch.compile 下，這會移除 TorchTPU 原本會在 collective 操作周圍插入的圖形中斷（見上文的 Skinny Collectives），從而產生包含內嵌 collective 操作的大型 StableHLO 模組。XLA 編譯器可以透過重新安排 collective 操作的位置，有效地最佳化該模組，盡量重疊運算與通訊；在搭配 SimpleFSDP 或手動放置 collective 操作時，已證明這能帶來顯著的效能提升。這種做法的一項挑戰是，MPMD 程式碼可能不易察覺；即使程式碼並非嚴格的 SPMD，用戶仍可能使用 @spmd_safe 裝飾器，令情況變得脆弱。
  • (Optional) Compilation Handshake (RFC). This mechanism runs when a given torch.compile statement triggers compilation of a module containing communication operations. It initiates a consensus round across all replicas to verify that the compiled FX graphs are identical, which in turn allows replicas to disable graph breaks around collectives for the given FX graphs, similarly to the @spmd_safe mechanism described above.The handshake stage is controlled by an environment variable TORCH_TPU_SPMD_HANDSHAKE_STAGE which can be set to:
    • OFF – 不執行握手（目前的預設設定）；
    • COMPILE_STAGE – 開始編譯 FX 計算圖之前執行握手；
    • DISPATCH_STAGE – 開始執行已編譯的 FX 計算圖之前執行握手。只有在用戶預期會針對特定 rank 重新編譯時，才應使用此模式，因為其開銷高於上一個選項。

分佈式初始化

在 TorchTPU 中初始化分佈式工作負載，遵循 PyTorch 分佈式的標準語義；但 TorchTPU 需要設定兩個 TPU 專用變數，而標準的 PyTorch torchrun 並不知道如何設定：

• TORCH_TPU_TOPOLOGY：TPU 切片的實體網格幾何配置（例如，2x2x1 切片配置使用「2,2,1」）。
• TORCH_TPU_SLICEBUILDER_ADDRESSES：用於建構 TPU 晶片間互連（ICI）網絡網格的開放連接埠／位址清單（每部裝置一個；例如「localhost:8471,localhost:8472,…」）。

另一種做法是在 Python 中使用 torch.tpu.distributed.set_tpu_launch_env()，它可以在 Google Cloud 環境中設定這些變數。以下展示啟動 TorchTPU 分佈式工作負載的範例程式碼。

A. 工作程序實作（train.py）

import os
import torch
import torch.distributed as dist

def worker_main():
  # 1. Initialize distributed process group on TPU using the tpu_dist backend
  dist.init_process_group(backend="tpu_dist")

  # 2. Query rank, world size, and construct target TPU device
  rank = dist.get_rank()
  world_size = dist.get_world_size()
  device = torch.device("tpu")

  print(f"[Rank {rank}/{world_size}] Initialized on {device}")

  # 3. Perform a distributed collective operation across TPU workers
  x = torch.ones(4, 4, device=device) * (rank + 1)
  dist.all_reduce(x, op=dist.ReduceOp.SUM)

  print(f"[Rank {rank}] Reduced sum: {x[0, 0].item()}")

  # 4. Clean up process group resources
  dist.destroy_process_group()

if __name__ == "__main__":
  worker_main()

B. 使用 torchrun 啟動工作負載（CLI）

# Single host launch (spawns 1 process per TPU chip on an 8-chip node)
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=8 train.py

最後請注意，指定節點的預設 process group 就是該節點的切片。

最佳裝置網格映射

由於 TPU 的連接方式獨特，必須將模型的平行策略（TC、CP、FSDP、DP）與 TPU 硬件的實體互連拓撲對齊，才能盡量提高頻寬並降低網絡延遲。請注意，TPU 使用分層網絡架構：

• ICI：透過直接的光纖／銅線互連，在 Torus 網格中連接晶片。提供每秒數十至數百 gigabytes/sec 的頻寬。
• DCN：用於連接不同 TPU 切片的標準主機層級乙太網絡。其頻寬遠低於 ICI，延遲則高得多。

一個好做法是使用 TorchTPU 的 topology_aware_mesh() API，並將通訊量最大的並行維度對應至速度最快、連續性最高的實體 ICI 維度；也就是按網絡密集度由低至高排列網格維度。以下表格列出一些例子：

以下是程式碼範例：

from torch.distributed.device_mesh import DeviceMesh

# Single-Slice 2D Mesh Example (DP + TP)

# Logical mesh shape: (dp=4, tp=2) on an 8-chip slice ordered from least
# communication-intensive to most communication-intensive:
mesh_shape = (4, 2)
dim_names = ("dp", "tp")

# Computes optimal rank layout aligned with physical ICI rings
optimal_rank_tensor = torch.tpu.topology_aware_mesh(
    mesh_shape=mesh_shape,
)

# Instantiate the DeviceMesh
mesh_2d = DeviceMesh("tpu", mesh=optimal_rank_tensor,
                     mesh_dim_names=dim_names, preserve_rank_order=True)

dp_mesh = mesh_2d["dp"]
tp_mesh = mesh_2d["tp"]

作為未來工作的一部分，我們會探索提供 API，讓用戶更直接地自訂網格，以便在特定情況下達至最佳集合通訊效能。

分散式 Pallas 核心

TorchTPU 支援自訂 Pallas 核心，可透過 ICI 上的硬件 RDMA 直接進行晶片間通訊，無須主機 CPU 參與。這讓你可以編寫高效的點對點通訊核心（例如環形注意力、自訂全互換／洗牌、流水線並行位移或邊界交換），並將它們作為原生 PyTorch 函式呼叫。以下範例展示了這一點：

# 1. Pallas Kernel Body: Asynchronously sends local buffer to neighbor rank (rank + 1)
def pallas_send_to_next_kernel(x_ref, y_ref, send_sem, recv_sem):
  rank = jax.lax.axis_index("x")
  world_size = jax.lax.axis_size("x")
  dest_id = (rank + 1) % world_size
  # Trigger hardware Remote DMA over TPU interconnect
  async_copy = pltpu.make_async_remote_copy(
      src_ref=x_ref,
      dst_ref=y_ref,
      send_sem=send_sem,
      recv_sem=recv_sem,
      device_id=(dest_id,),
  )
  async_copy.start()
  async_copy.wait()

# 2. JAX Wrapper with ShardMap and DMA Semaphores
def build_jax_fn(mesh):
  def local_pallas_call(x):
    return pl.pallas_call(
        pallas_send_to_next_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid_spec=pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            in_specs=[pl.BlockSpec(x.shape)],
            out_specs=pl.BlockSpec(x.shape),
            scratch_shapes=(
                pltpu.SemaphoreType.DMA,
                pltpu.SemaphoreType.DMA,
            ),  # send & recv semaphores
        ),
    )(x)
  return jax.shard_map(
      local_pallas_call,
      mesh=mesh,
      in_specs=P("x"),
      out_specs=P("x"),
      check_vma=False,
  )

# 3. Create the TorchTPU callable wrapper
def get_send_to_next_op(world_size: int):
  mesh = torch.tpu.pallas.get_pallas_mesh(axis_names=("x",), mesh_shape=(world_size,))
  jax_fn = build_jax_fn(mesh)
  return torch.tpu.pallas.jax_op(
      "pallas::send_to_next",
      build_jax_fn(mesh),
      mesh=mesh,
      input_partition_specs=[P("x")],
      output_partition_specs=[P("x")],
  )

下表概述如何編寫具通訊功能的 Pallas 核心。

處理設備映射不一致問題。以往在 TorchTPU 中，使用低階晶片間操作（例如 lax.axis_index 和 async_remote_dma／make_async_remote_copy）的 Pallas 核心，與 JAX 一樣，會根據原始硬件設備 ID 而非 PyTorch rank 路由通訊。當硬件設備 ID 與 rank 順序不一致時，環形通訊和跨設備 DMA 會路由至錯誤的實體設備，而在隔離子羣組中執行則會因分區數量不匹配而失敗。

TorchTPU 透過強制採用邏輯定址，解決了這些不一致問題：

• Rank 對齊映射：TorchTPU 現在會將 Pallas 核心直接映射至 PyTorch rank，而非原始實體設備 ID，確保透過 get_pallas_mesh API 進行一致的路由。此 API 會將 Pallas 中的分區和軸映射至 PyTorch Process Group 中的 rank 順序。
• 邏輯設備定址：TorchTPU 現在會強制採用 DeviceIdType.LOGICAL 和標準化設備配置（da(0, p) = p），確保每個邏輯 rank 一致地映射至正確的實體目的地。

統一的 ProcessGroup 整合。TorchTPU 現在提供取得 rank 對齊 JAX Meshes 的機制，讓核心沿用 PyTorch ProcessGroup 拓撲。這樣一來，無論是在完整全域環境還是隔離子羣組配置中，都能高效執行，無須手動修補狀態或造成計算圖中斷。

開發模式

TorchTPU 由 Google 與社羣協作，以開放方式開發，並優先重視透明度和參與。這種開發模式鼓勵更廣泛的 PyTorch 生態系及業界夥伴作出貢獻。

我們預期 TorchTPU 可服務各類用戶，包括研究機構和業界領導者。我們亦計劃將其整合至 Hugging Face 等平台，讓各種開源模型開箱即用地支援 TPU。

為支援開發者，TorchTPU 提供全面的文件，內容包括 (請參閲 http://torchtpu.readthedocs.io/)：

• 教學：循序漸進的學習路徑，包括 VM 配置和模型移植。
• 實作方案：針對問題提供解決方案，例如效能最佳化和核心函式實作。
• 參考資料：技術規格及 API 對應。
• 解説：深入探討架構和硬件基礎。

文件使用可直接執行的筆記本編寫。

此專案採用 Bazel 管理建置和測試，並持續完善測試框架，計劃採用 Pytest。CI/CD 基礎設施正整合至 PyTorch 官方 CI HUD，以維持高可靠性並促進社羣貢獻。

TorchTPU 的設計具備可擴充性，可與 PyTorch 生態系的重要組件無縫整合。它支援透過 Pallas 和 Helion 開發自訂核心函式，並透過 vLLM、SGLang 和 TorchTitan 等專案，為分散式訓練和推理工作負載提供一級支援。這種模組化設計讓用戶可沿用現有的 PyTorch 慣用寫法，同時針對 TPU 專用硬件架構進行最佳化。

案例研究

案例研究 1：Hugging Face 模型涵蓋範圍

TorchTPU 與 PyTorch 生態系開箱即用地相容，讓開發者無須修改程式碼或重新撰寫，即可在 Google Cloud TPU 上執行開源模型。為驗證這項承諾，我們使用 Hugging Face 模型建立自動化測試框架，涵蓋 NLP、視覺、音訊、多模態及生成式擴散等領域。這條持續整合管道會以 CPU 和 GPU 基準為參照，評測 TorchTPU 的涵蓋範圍、穩定性及數值等效性。我們持續在從 v5e 開始的各代 TPU 硬件上，以 eager 和編譯模式評測約 5,000 個 Hugging Face 模型，涵蓋數百種架構。

我們採用穩健的測試方法，以語料分層和均衡的架構選擇為基礎。這套方法透過兩層篩選，在反映實際採用情況與廣泛涵蓋不同架構之間取得平衡，避免重複測試細微變體，同時確保全面評測正式環境中的實際權重。第一層優先選取社羣採用率、下載增長速度和趨勢分數名列前茅的模型，例如 Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformers，確保社羣最關注的模型均能正常運作。第二層則從第一層中挑選一個精簡子集，確保從超過 500 個不同模型系列中，至少選取 1 個代表模型。我們利用第二層模型，針對全面涵蓋 PyTorch 運算子和編譯器模式的集合，與 GPU 基準比較，以驗證數值準確度。

驗證流程由第 1 階段開始，重點是進行冒煙測試及 lowering 驗證，以確保執行階段穩定性。在這個初始階段，流程會以合成權重建立模型、驗證 eager mode 下的前向執行，並確認圖追蹤、將 FX 圖輸出為 StableHLO 圖，以及 XLA 編譯均可順利完成，不會出現錯誤（例如當機或例外）。確認執行階段穩定後，第 2 階段會使用特定領域的評估指標，對照既定基準評估數值正確性。對於因果 LLM，會透過困惑度衡量數學保真度（困惑度 = e^L = exp((1/N) Σ_i=1^N L_i)）；若 TPU 與基準困惑度的相對差異維持在 0.10 以內，測試便算通過。對於視覺、音訊及編碼器架構中的高維特徵空間，我們採用雙閾值相對 MSE 指標（相對 MSE = MSE(Y_TPU, Y_Baseline) / (Var(Y_Baseline) + 10⁻⁷)）；相對 MSE 必須 ≤ 0.10，或絕對 MSE 必須 < 0.05，以避免基準輸出接近常數時造成影響。

下表概述截至本文撰寫時的模型覆蓋結果：

針對有界動態性支援（SymInts），我們在 839 個模型上評估正確性（這些模型是來自 5K 語料庫的因果 LLM 清單），並確認其中 834 個（99%）的數值正確。

案例研究 2：開箱即用效能

TorchTPU 致力成為真正原生的 PyTorch TPU 後端。我們認為，這也代表執行純 PyTorch、與裝置無關的程式碼時，應有良好效能。以下以兩個例子展示這一點。

第一個例子收錄於 TorchTPU 程式碼庫，是我們以標準 PyTorch 模組及運算（例如 nn.Linear、nn.RMSNorm、F.scaled_dot_product_attention 及 torch.distributed 集體操作）編寫的 Muse Glimmer 30B 模型純 PyTorch 自包含實作。我們刻意沒有使用任何模型專用的自訂核心（但仍會使用 PyTorch 內部核心）。這個例子的效能如下：

• Inference (prefill) performance on a single TPUv7 chip at
  • 序列長度為 5,120 時，約 25.7k TPS，MFU 為 59%；
  • 序列長度為 10,240 時，約 22.4k TPS，MFU 為 51%。

  請注意，這個例子使用 FSDP（以原生 PyTorch 集體操作實作），world size 為 2，因為單一 TPUv7 晶片會向 PyTorch 提供兩個裝置，分別對應兩個內置 TensorCore。

• Training on 4 TPUv7 chips (i.e. a full single host), using FSDP, at
  • 每個 TPU 的效能：序列長度為 5,120 時，約 5.6k TPS，MFU 為 38%；
  • 每個 TPU 的效能：序列長度為 10,240 時，約 4.6k TPS，MFU 為 32%。

  在第二個例子中，我們啟用activation checkpointing，只將部分前向傳遞激活值儲存在 HBM 中；進行反向傳遞計算時，則按需要重新計算／重新物化未儲存的中間激活值。

這些示例的一項關鍵特點是，同一段程式碼除了指定加速器類型（例如，指定為 device="tpu" 而非 device="cuda"）之外，無須作任何實質改動，便可在 TPU 和 GPU 上執行；而且採用 torch.compile 和 FSDP 等標準 PyTorch 最佳化及分片技術時，在 TPU 上也有出色表現。

為進一步驗證 TorchTPU 在純 PyTorch 程式碼上的效能，我們亦採用了 TorchTitan 專案中一個由外部開發的 Muse Glimmer 30B 實作版本。該版本原本為 GPU/CUDA 編寫。只需作少量修改便可配合 TorchTPU 使用（我們只把 TorchTPU 尚未支援的 FlexAttention 替換為原生 F.scaled_dot_product_attention，並套用 torch.compile 和 SimpleFSDP），移植至 TorchTPU 的版本便可達到：

• Training on 4 TPUv7 chips:
  • 每個 TPU 的約 5.3k TPS，MFU 為 36%，序列長度為 5,120；以及
  • 每個 TPU 的約 4.2k TPS，MFU 為 29%，序列長度為 10,240。

總括而言，所有這些示例均表明，純 PyTorch 模型實作，甚至原本為 GPU 編寫的模型，只需作簡單的 PyTorch 修改並採用標準 PyTorch 最佳化，無須使用自訂核心等進階、裝置專屬的反覆調校技術，仍可在 TPU 上有足夠出色的表現。我們認為，這是促進 TorchTPU 普及的一項重要特點。當然，對於追求最高效能的生產環境工作負載，我們絕對建議採用以 Pallas kernels 為基礎的最佳方案，以及 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU 等生產級技術堆疊。

個案研究 3：將訓練擴展至 Pod 規模

我們使用 TorchTitan 驗證 TorchTPU 可擴展至 1 個 TPU v7 POD。這項測試選用了 Qwen3 MoE 235B。作為初步測試，我們確認在 64 個 TPU v7 的配置中，TorchTitan 配合 TorchTPU 執行該模型的效能，與 JAX 配合 MaxText 相若（PyTorch/TorchTitan 為 1828 TPS，JAX/MaxText 為 1937 TPS）。

以下數據按以下設定計算：

• 模型：qwen3_moe_235b_scan
• 本機批次大小（即每個 TensorCore）：16，序列長度：4096
• 並行方式：單一 slice 執行時使用 FSDP（透過 SimpleFSDP），多個 slice 執行時使用 HSDP
• 精度：BF16，並採用 FP32 混合精度歸約（梯度歸約：f32、RoPE：f32、優化器狀態：f32）
• 核心（Pallas）：Tokamax Splash Attention（帶來顯著效能提升）、融合交叉熵損失核心，以及 MoE 層的分組矩陣乘法
• 掃描最佳化：94 個解碼器層以一個 scan 迴圈處理堆疊權重
• 卸載解碼器層輸入的激活值（作為掃描最佳化的一部分）
• 以流水線方式安排集合通信，讓集合通信與相鄰層的計算重疊（scan 的一部分）
• 完整激活值檢查點
• 使用 TorchTitan 的 Graph Trainer 合併執行前向及反向傳播
• 專家不平衡：使用 TorchTitan 的輔助損失以防止 MoE 不平衡
• 訓練數據集：Hugging Face C4 en 數據集，並採用自訂數據集分片

在所有拓撲配置下，我們的數據均證實弱擴展性：整體計算量按線性比例擴展；如下所示，整體每秒 token 數亦線性增加，而 MFU 基本維持不變。

我們的數據亦顯示，在所測試的不同並行度下，損失曲線幾乎完全一致。

在多切片 HSDP 配置中，前向傳播不涉及切片之間的通訊，每個切片各自處理一部分批次資料。不過在反向傳播中，切片會先計算第 t 層的梯度，然後使用 all-reduce（透過 DCN，並採用 Megascale-XLA）將所有切片的梯度加總。接着，計算會繼續處理第 t-1 層的梯度。得益於 XLA，TorchTPU 能將第 t+1 層的 all-reduce 與第 t 層的計算完全重疊執行，如下圖所示。這種重疊執行是每個 TensorCore 的 TPS 從 1 個切片增至 8 個切片時仍幾乎維持不變的主要原因之一，如上表所示。

案例研究 4：在 vLLM 上使用 Qwen3.5 進行 LLM 推理（SemiAnalysis InferenceX）

大型語言模型（LLM）推理，尤其是大型混合專家（MoE）架構的推理，是 AI 雲端基建中要求最高、成本最關鍵的工作負載之一。在 SemiAnalysis 發表的一項獨立基準測試評估中（TPU Inference Externalization Full Steam Ahead — InferenceX），Google 的 TPU v7 配合 vLLM-TorchTPU 技術堆疊，展現出具業界競爭力的效能成本比。這項基準測試由第三方審核 TorchTPU 軟件在大型混合專家（MoE）工作負載上的就緒程度，並證明面向 TPU 的新 vLLM 技術堆疊，在向社羣公開提供的 TPU 系統上具備出色的成本效益和高吞吐量。

在此基礎上，Google 最近提交了首批 Qwen 3.5 在 Ironwood TPU 上的 AgentX 效能結果。AgentX 涵蓋代理式工作負載的四個要素：多輪對話、長上下文、高前綴重用率，以及子代理突發工作負載。TPU 已證明，即將推出的 vLLM-TorchTPU 技術堆疊能夠成功處理代理式工作負載。TPU v7 僅需 4 顆晶片，便能全面領先帕累托前沿。

• 在 InferenceX 8K/1K 測試中，整體每百萬個 token 的成本最低，為 2.01 美分（低於我們在 9 月首次亮相時的 3.6 美分）；
• 在 AgentX 測試中，整體每百萬個 token 的成本最低，為 0.83 美分，並在互動性（用戶回應速度）方面達到業界最先進水平。全面領先帕累托前沿。

這些基準測試結果反映 TPU 系統在真實應用中的強大推理服務效率，突顯在 Google 自行設計的 ML 矽晶片上部署大型模型具有巨大的經濟可行性。

要取得這些具競爭力的效能和 TCO 結果，必須透過自訂 Pallas 核函數和針對性的執行階段調校，解決硬件與軟件配合上的關鍵瓶頸。下表列出幾個例子：

這些結果突顯現代服務基礎設施中的 TPU 最佳化，展示 TCO 和軟硬件協同如何決定實際服務效率。最終，採用 TorchTPU 後端的新 vLLM-TPU，讓工程團隊能夠在公開提供的 TPU 基礎設施上，使用標準 PyTorch 慣用寫法部署現成的開放權重模型。

案例研究 5：使用 TorchTitan 微調 LLM

TorchTPU 原生整合 TorchTitan——PyTorch 的大型分散式訓練參考平台——為密集型大型語言模型（MuseGlimmer 30B、Llama 3 8B–405B）、稀疏混合專家模型（Qwen 3 235B MoE）及生成式擴散 Transformer（FLUX.1-dev 12B）提供統一且高效能的訓練環境。在 TPU v6e (Trillium) 和 TPU v7x (Ironwood) 叢集上以 BF16 精度進行評估，多切片執行透過 Google Kubernetes Engine (GKE) 上的 torchrun 無縫編排。TorchTitan 可作為開源專案在此GitHub 儲存庫找到。

分散式策略與分片架構。高效的多加速器訓練，需要模型並行策略與 TPU 實體互連拓撲無縫配合。為此，我們採用 2D+ DeviceMesh 對應方式，將張量並行和專家並行等通訊密集型維度，對應至連續、高頻寬的晶片間互連 (ICI)；同時透過較慢的跨切片數據中心網絡 (DCN)，傳輸複製的數據並行資料。這項策略透過整合 SimpleFSDP 實現，後者提供符合 SPMD 的操作。為進一步減低延遲，我們設定 XLA 執行階段，將 All-Gather 和 Reduce-Scatter 等集合操作卸載至晶片上的 SparseCore 引擎，從而有效地讓通訊與運算密集的 TensorCore MXU 運算重疊執行。最後，我們採用 HLO 差異比對工具和 @spmd_safe 註解，確保執行結果具決定性並避免同步死鎖，同時移除 torch.compile 下不必要的計算圖中斷。

效能與多節點擴展。透過逐層系統化最佳化，TorchTitan on TorchTPU 為 FLUX 模型帶來出色效能和多節點擴展效率。以開箱即用的 eager 基準為例，步驟時間為 11.49s（0.26 樣本/秒）；採用 DeferAndFuse eager 模式後，步驟延遲降至 9.95s（加速 1.15 倍）。加入 SimpleFSDP 分片後，執行速度進一步提升，步驟時間降至 6.52s（加速 1.76 倍）；整合自訂 Pallas SplashAttention 核心後，步驟時間則降至 5.55s（加速 2.07 倍）。啟用完整的 torch.compile(backend=”tpu”) 後，效能大幅提升至 3.02s（加速 3.80 倍）；再套用針對目標模型的 peephole 改寫，最終最佳化步驟時間為 2.78s（1.08 樣本/秒）。在多主機叢集中，TorchTitan 展現接近理想的弱擴展和強擴展能力：在 FLUX.1-dev 上，使用 SimpleFSDP，從單一 v6e-8 節點（8 個晶片）擴展至 v6e-256 切片（64 個節點共 256 個晶片），每個樣本的步驟時間僅由 0.92s 輕微增加至 1.00s（總步驟時間為 3.01s），展現其在光學 ICI 網絡上的近線性效率。在更大規模下，TorchTitan 在 TPU v7x (Ironwood) 上以 BF16 執行 Llama 3 405B，使用 1,024 個晶片和 8K 上下文窗口，達到每個晶片每秒 500 tokens 和 50+% MFU。同樣地，在 8,192 個 TPU v7x 晶片上評估 Qwen3 235B MoE，亦展現出高效的擴展能力，在 BF16 精度及 4K 序列長度下達到 28% MFU。

主要架構最佳化。為了在不同模型架構中維持峯值運算密度並消除硬件瓶頸，我們採用多項系統及核心最佳化措施：

• SimpleFSDP: 適合編譯器的 FSDP 版本。FSDP2 透過模組掛鈎注入 all-gather 和 reduce-scatter 操作，令 torch.compile 拆分為逐個操作的派發；SimpleFSDP 則將它們表示為追蹤圖中的一般操作。這樣可讓 XLA 看見這些集合操作，TPU 編譯器便能在其周圍進行融合、讓它們與 MXU 運算重疊，並可將它們卸載至 SparseCore。
• Scan Higher-Order Op (torch._higher_order_ops.scan)：以一個迴圈，使用堆疊的權重執行重複的 Transformer 區塊。隨著深度增加，編譯時間和已聚集參數的記憶體用量仍維持不變。每層的集合操作位於迴圈內，因此 XLA 可讓它們與運算重疊：隱藏延遲的排程器會在每層運算期間隱藏該層 all-gather 的延遲，集合操作管線器則可選擇預取下一層的 all-gather；在反向傳播時，它會讓某層的 reduce-scatter 在下一層運算期間完成。
• Pallas SplashAttention Kernel：以自訂 Pallas 核心取代缺乏向量 tile 對齊的預設注意力實作。該核心支援完整的分組查詢注意力（GQA）、偏置和 dropout，令擴散注意力區塊的速度提升 +17.5%。
• Fused Pallas Cross-Entropy Loss: Avoids out-of-memory (OOM) faults when materializing [B×S×V] logits for large vocabularies by porting Tokamax’s online reduction kernel to compute loss in-place.
• Graph trainer: Captures the full training step, forward, loss, backward and optimizer update, as a single graph that compiles into one XLA program. This gives the TPU compiler a global view for fusion, collective scheduling and memory planning, and lets TPU-specific passes decide per activation whether to save, recompute or offload to host, and donate optimizer buffers so state is updated in place rather than double-buffered.
• Activation Offloading: The activations that the backward pass needs can be moved to host memory after the forward and copied back just before backward uses them, which frees HBM for larger models and longer sequences. The graph trainer inserts the offload and reload ops into the traced graph. Under scan, they sit inside the loop body, so saved activations move one layer at a time. XLA turns these copies into async transfers, and in backward copies the next layer’s activations back while the current layer computes. This is essential for the largest models, such as Llama3-405B and Qwen3-235B, to fit at their target sequence lengths.

Ongoing development for TorchTitan on TorchTPU targets four strategic priorities: expanding list of models optimized for TPU; integrating FP8 and FP4 low-bit quantization via torchao; and supporting end-to-end Reinforcement Learning (RL) by coupling training loops with vLLM/SGLang sampling backends in VeRL.

Case Study 6: Hardware-Accelerated Recommendation Models with TorchRec and SparseCore

Recommendation and ranking architectures—such as Meta’s Deep Learning Recommendation Model (DLRM)—face critical memory and compute bottlenecks due to high-cardinality embedding tables. On standard GPU clusters, managing these tables may require complex software-based sharding via DistributedModelParallel (DMP) and separate optimizer updates, which often introduce latency. To bridge this gap, Google collaborated with the Meta PyTorch team to upstream a native TPU backend directly into the TorchRec repository, allowing recommendation models to bypass host-side compute bottlenecks and tap directly into TPU SparseCore (SC) hardware.

• Architectural Integration. The resulting implementation creates a seamless data pipeline where high-cardinality lookups are accelerated by dedicated hardware. A multi-threaded C++ PyBind backend runs on the host CPU to parallelize partition mapping and index padding, outputting structured batches that map directly to the TPU SparseCore units. Within the TPU, we implemented In-Place Optimizer Fusion: rather than materializing gradient tensors (.grad) in High Bandwidth Memory (HBM)—a process that traditionally complicates memory management—we map backpropagation directly to SparseDenseMatmulGrad. This ensures parameter updates (SGD, AdaGrad, or Adam) occur in-place, eliminating gradient memory overhead entirely. Finally, low-level ATen operators, such as SparseDenseMatmulOp (for forward lookups) and RaggedAllToAll (for optical ICI cross-chip collectives), provide hardware-accelerated kernels that execute directly on the silicon.
• Performance & Results. When evaluated on an industry-standard DLRM v2 model across a distributed 32-chip TPU v7x (Ghostfish) cluster, this integration achieved 100% performance parity with Google’s highly optimized JAX baseline, reaching an identical 18.3ms step latency. Notably, this performance matching was established using baseline SparseCore kernels before applying advanced optimizations. Convergence and mathematical fidelity were equally robust, yielding identical loss trajectories compared to the JAX implementation.

Active development focuses on further saturating bandwidth through TensorCore/SparseCore pipelining—interleaving MLP execution with embedding lookups—and advanced table-stacking techniques to unify memory buffers and reduce kernel overhead. Additionally, ongoing work aims to expose lower-level SparseCore embedding lookup operations (e.g. SparseMapRow, SparseGather etc) via ATen kernels, enabling users to implement custom lookup algorithms and advanced sharding strategies. This integration solidifies the TPU as a high-efficiency engine for the next generation of recommendation architectures.

Future Work

With TorchTPU now publicly available on GitHub, this initial open-source release represents only the beginning of our development journey. As we engage directly with the broader PyTorch community, our roadmap for the next 6 to 12 months focuses on production hardening, ecosystem expansion, and continually integrating community feedback to bridge the gap between TPU hardware capabilities and PyTorch-native developer workflows. We are prioritizing four strategic pillars to drive this evolution:

• Foundational Hardening & User Velocity: We are refining the “True Eager” experience by building a robust library of precompiled, dynamic kernels for ATen operators. We are also making XLA Bounded Dynamism fully production-ready, allowing the torch.compile stack to handle SymInts without extra recompilation, and streamlining our GKE packaging to lower the barrier for deployment across PyTorch versions. Furthermore, we are continuously expanding operator coverage and advancing our quantization support—specifically targeting FP8 and FP4 formats for next-generation TPUv8 architectures, and implementing nested compiler regions to reduce compilation time.
• Pod-Scale Distributed Training & Advanced Training: Building on our successful POD-scale training validation of Qwen3, we are expanding to additional model architecture and sharding techniques at and beyond that scale. We are enabling end-to-end Reinforcement Learning (RL) training loops by tightly coupling scalable training with high-throughput vLLM/SGLang sampling backends orchestrated via Ray Core. To enhance distributed stability, we are also validating a consensus-based compilation handshake to safely optimize collective operations with compute without fragile @spmd_safe annotations.
• Emerging Architectures, Inference & Serving: We are committing to proactively support cutting-edge model architectures as they emerge in the open-source ecosystem—including reasoning models, multimodal and generative diffusion models (e.g., Wan 2.2, Flux Klein), and novel Mixture-of-Experts (MoE) variants. Our efforts focus on optimizing the serving stack in vLLM and SGLang for low-latency production, incorporating native FP8/FP4 quantization across future TPU generations, and ensuring day-one TPU support for new architectural paradigms.
• High-Efficiency Communication & Diagnostics: We are optimizing collective operations within and across TPU slices and adding zero-copy DMA to enable hardware-assisted memory offloading. To support debugging at scale, we are transitioning to a WebGPU-based Trace Viewer 2.0 and deploying “flight recorder” continuous profiling (XProfz) to identify intermittent bottlenecks in production environments.

TorchTPU is built on the belief that a successful ecosystem thrives through transparent, collaborative engineering. As we open-source our four foundational repositories (TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan) at the October 2026 PyTorch Conference, we actively seek feedback on model architecture prioritization and workload sequencing. We invite developers, researchers, and partners to collaborate with us via GitHub or our active TPU Builders program.

†The TorchTPU Team

Claudio Basile, Michael Voznesensky, Gunhyun Park, Ben Wilson, Parker Holloway, Zhanyong Wan, Vlad Belous, Manish Khanna, Paul Ganssle, Kevin Gleason, Andrew Dame, Brendan Folie, Bhavya Bahl, Haifeng Jin, Pedro Goncalves Mokarzel, Arian Arfain, Luke Boyer, Chi Zhang, Yiling Yuan, Adam Mainz, Aditya Patel, Huangze Tang, Jay Gan, Lizzie Liang, Will Froom, Tushar Agarwal, Karlo Basioli, Theotime Combes, Emilio Cota, Minos Papadopoulos, Tudor Pipernea, Yarong Mu, Qi Zhou, Yaoshiang Ho, Aleksey Vlasenko, Alexander Brauckmann, Ziyin Huang, Jack Oh, Jerry Hsu, Tanvi Bajpai, Jialei Chen, Wenjun Gu, Chintan Shah, Chris Chan, Sayce Falk, Chris Jones, Rob Torres, Kat Ko, Lee Howes, Freya Azad, Andi Gavrilescu, Joe Pamer, Bill Jia

Plus many, many others…

The post Introducing TorchTPU appeared first on PyTorch.