TPU 的 4x4x4 到底在乘什麼？一篇看懂 TPU Slice 與 GPU 設計哲學

一開始在學習 TPU 的架構時，我對於規格說明上切片 (slice) 是 4x4x4、8x16、16x16 這類的數字感到困惑……當下冒出一連串疑問：為什麼是「幾乘幾」？這些數字是矩陣乘法的維度嗎？為什麼 GPU 的世界好像不用管這些，TPU 卻要把拓撲搞得這麼複雜？

之前在 AI 有意識嗎？為什麼 ChatGPT 這麼好用卻常一本正經地出錯 ¹ 拆解過 LLM 的本質：一台用矩陣運算做機率預測的機器。這篇想往硬體再挖一層：當矩陣大到一顆晶片裝不下時，Google 是怎麼把一個數學問題，變成一座機房的接線問題？

結論在前：TPU slice 的「幾乘幾」（例如 4x4x4 或 8x16）指的是晶片在機房網路中的物理排列拓撲，不是矩陣乘法在數學上的切分方式。這兩者是兩個不同層次的問題，中間由 XLA 編譯器負責對應。而 TPU 之所以選擇 2D/3D 環面（torus）拓撲，而不是像 GPU 叢集那樣蓋一座昂貴的交換機網路，本質上是一個成本與複雜度的取捨：Google 用軟體與編譯器的複雜度，換取硬體規模化時的極致性價比。

TL;DR：四個核心判斷

作為我個人的閱讀與消化筆記，本篇討論以下內容：

• 深度學習的本質是矩陣乘法：神經網路在前向與反向傳播時，多數計算都可化為矩陣運算，這也正是 TPU 把整顆晶片押注在「脈動陣列」上的原因。
• Slice 的「幾乘幾」是晶片接線圖，不是矩陣切分：4x4x4 描述的是 64 顆晶片的 3D 網格排列；矩陣怎麼切、切給誰算，是 XLA 編譯器把邏輯網格「貼」到物理拓撲上的結果。
• TPU 不是「完全沒有交換機」：機櫃內的 cube 用銅纜直連，cube 之間其實靠 OCS（光學電路交換）串接。真正省下的是昂貴的電子封包交換機，OCS 與光學元件只佔整體系統成本不到 5%。
• GPU 的複雜拓撲沒有消失：GPU 叢集把複雜度藏進硬體帳單（NVSwitch、InfiniBand 胖樹等拓撲架構）裡，而這些硬體元件在資料中心往往成本高昂。TPU 則省下這部分硬體開銷，把複雜度交給 XLA 編譯器與工程師。

為什麼深度學習離不開矩陣乘法？

在以神經網路構建的深度學習中，每一層本質上是「一組輸入乘上一組權重，再加總」的線性變換，再接一個非線性啟動函數：每個神經元把上一層的輸出各自乘上權重、相加、過一個非線性函數。當你把一整層神經元、一整批（batch）資料同時考慮進來，這些成千上萬個「乘法加總」就能自然地打包成一個矩陣乘法。

---
config:
  layout: dagre
  flowchart:
    curve: linear
    nodeSpacing: 15
    rankSpacing: 110
---
flowchart LR
    subgraph IN["輸入層（上一層的輸出）"]
        x1(("x1"))
        x2(("x2"))
        x3(("x3"))
    end
    subgraph HID["隱藏層神經元"]
        h1(("h1"))
        h2(("h2"))
        h3(("h3"))
        h4(("h4"))
    end
    subgraph OUT["輸出層（傳給下一層）"]
        y1(("y1"))
        y2(("y2"))
    end
    x1 --> h1 & h2 & h3 & h4
    x2 --> h1 & h2 & h3 & h4
    x3 --> h1 & h2 & h3 & h4
    h1 --> y1 & y2
    h2 --> y1 & y2
    h3 --> y1 & y2
    h4 --> y1 & y2

圖中每一條箭頭都代表一次「乘上權重」：一個神經元的輸出，就是所有進入箭頭的加權總和再過一個非線性函數。把一整層的箭頭合起來看，就是輸入向量乘上權重矩陣；再把一整批資料疊起來，就是一個完整的矩陣乘法 X × W。

打包成矩陣有兩個巨大的工程好處：

1. 平行化：矩陣乘法裡的大量乘加運算彼此獨立，可以交給成千上萬個運算單元同時做，完美匹配 GPU / TPU 這類高度平行運算的硬體。
2. 資料重用：矩陣乘法中同一筆資料會被多次使用，硬體可以讓資料在晶片內部流動、重複利用，減少昂貴的記憶體存取。在現代晶片中，搬動資料的能耗往往遠高於計算本身。

從 Transformer 到今天的 LLM，底層跑的幾乎都是同一件事：把 attention、前饋網路全部化約成一連串巨大的矩陣乘法。

flowchart LR
    A["神經網路的一層<br>數千個神經元 × 數千維輸入"] --> B["打包成矩陣乘法<br>輸入矩陣 X × 權重矩陣 W"]
    B --> C["高度平行的硬體<br>GPU / TPU 一次算數萬個乘加"]
    C --> D["資料在晶片內流動重用<br>減少記憶體搬運與能耗"]

TPU 的心臟：脈動陣列 MXU

TPU（Tensor Processing Unit）是 Google 為此打造的專用晶片，主要任務是執行矩陣運算，也就是大量的乘法與累加。它內建數千甚至上萬個彼此直連的乘加器，組成一個大型的實體運算陣列。

它的核心是 MXU（Matrix Multiply Unit，矩陣乘法單元），採用一種叫「脈動陣列」（systolic array）的設計 ²。

可以把脈動陣列想像成一條精心編排的工廠流水線：權重先被「預載」進一個二維的運算單元網格裡固定不動，輸入資料像水流一樣從側面一波波流進去，每個運算單元做完「乘法 + 累加」就把結果往下傳。資料在網格裡規律地「脈動」，不需要反覆讀寫記憶體——這正是它比通用處理器省電又高效的原因。

TPU 主機會將資料串流至饋入（infeed）佇列。TPU 再從饋入佇列載入資料，並將其儲存到高頻寬記憶體（HBM）。計算完成後，TPU 會把結果寫入饋出（outfeed）佇列；接著，TPU 主機從饋出佇列讀取結果，並儲存在主機記憶體中。若要執行矩陣運算，TPU 會將參數從 HBM 載入矩陣乘法單元（MXU）。

各世代 TPU 的 MXU 規格與拓撲如下 ^{2 3 4 5 6 7 8 9}：

這張表有兩個值得注意的細節：

• MXU 的邊長從 v2 到 v5p 一直維持 128x128，到了第六代 Trillium（v6e）才擴大為 256x256，並由第七代 Ironwood（TPU7x）延續。邊長翻倍，意味著每個時脈週期的乘加運算量變成四倍 ^{2 7 10}。
• 這個數字對寫程式的人有實際意義：要餵飽 MXU，矩陣維度最好是 128 的倍數（v6e / TPU7x 上是 256），否則硬體會自動補零（padding），等於花同樣的時間算一堆零 ^{10 11}。

再看最新兩代的走向：

• 第七代 Ironwood（TPU7x，2025）延續 v5p 的 3D torus + OCS 路線，單一 superpod 直上 9,216 顆晶片、FP8 峰值算力 42.5 ExaFLOPS，更大規模則交給資料中心網路（DCN）串接多個 superpod ^{8 12 13}。
• 2026 年 Cloud Next 發表的第八代則首次拆成兩顆專用晶片：訓練用的 TPU 8t 與低延遲推論用的 TPU 8i；8i 則針對 AI agent 工作負載引入新的「Boardfly」拓撲，用直連光路把晶片間通訊的跳數從 16 跳降到 7 跳 ^{9 14}。

訓練與推論的通訊模式不同，最佳拓撲也開始分道揚鑣，拓撲不是教條，而是跟著工作負載演化的成本工程，這正是本文主軸的最好註腳。

Slice 是什麼？晶片的「座位表」，不是矩陣的切法

理解了單顆晶片，接下來就是本文最容易搞混的地方。當你在 Google Cloud 租用一個 TPU slice，例如 v5e 的 8x16，這個「8x16」的意思是：128 顆晶片在機房裡排成 8 行 16 列的二維網格，每顆晶片只和上下左右的鄰居用高速通道（ICI，晶片間高速互連）直連 ⁵。

（實際運行負載的範例可以參考 GKE AI 系列：在 TPU 上用 JAX 跑第一個訓練工作 ¹⁵）

2D Torus 拓撲示意：以 4×4 slice 為例，每個方塊是一顆 TPU 晶片；實線是和上下左右鄰居的 ICI 直連，虛線是每行每列頭尾相接的「環繞」連線（8×16 同理，只是網格更大）

上圖用 4x4 縮小示意：每顆晶片只和上下左右的鄰居直連，而每一行、每一列的頭尾再用環繞連線相接，這就是下一段要講的「環面」（torus）。

這條晶片間的高速通道叫 ICI（Inter-Chip Interconnect）。3D 拓撲的晶片（v4、v5p）有 6 個 ICI 連接埠，朝 ±X、±Y、±Z 六個方向各接一個鄰居；2D 拓撲的晶片（v5e、v6e）則有 4 個，接上下左右 ^{4 11}。網格的頭尾再繞回來相接，就成了「環面」（torus），像把一張棋盤捲成甜甜圈，讓最邊緣的晶片也不會變成孤島。

這樣設計的首要考量是降低延遲，但也連帶幾個好處：

• 每顆晶片的「出入口」固定，布線成本可控：每顆晶片只需要幾個 ICI 埠（2D 是 4 個、3D 是 6 個），硬體設計能標準化量產；擴到更大規模時，不必把每顆晶片都接到昂貴的交換機（例如 Nvidia 採用 NVSwitch + 特殊的 NVLink 專線達到 GPU-GPU 之間高速的通訊）。
• 通訊距離隨維度下降：在環面上，任兩點之間的最遠距離大約是各維度長度的一半相加；從 1D 拉長線，改成 2D/3D 後，最壞情況的跳數（hop）會顯著下降，集體通訊（collectives）更容易跑出穩定帶寬和低延遲。
• 集體通訊有「天然路徑」：All-Reduce、All-Gather 這類操作可以沿著環面執行 ring / mesh 演算法，把資料像接力一樣繞一圈（或繞多圈）就完成同步，對深度學習這種規律工作負載特別友善。

但這樣的設計也意味軟體層通常需要管理 TPU 複雜的拓撲網路。在 GPU 胖樹（Fat-Tree）或全連接網路裡，你可以比較不用管哪些 GPU 彼此相鄰。但在 TPU slice 的 torus 上，「並行切分」必須對齊晶片的物理相鄰關係。一旦沒對齊，大量通訊就會被塞進長距離路徑、導致任務無法同步，甚至寫出 deadlock 的應用程式，最直接的影響就是 Job 卡住了。

把整個階層由小到大排開：

flowchart TD
    A["TPU 晶片<br>內含 TensorCore 與 MXU"] --> B["Tray 主機板<br>4 顆晶片配 1 台 CPU 主機"]
    B --> C["機櫃 Rack<br>4x4x4 Cube = 64 顆晶片<br>銅纜直連成 3D 網格"]
    C --> D["OCS 光學電路交換<br>把多個 Cube 接成任意大小的 Slice"]
    D --> E["Pod<br>v4: 4096 顆 / v5p: 8960 顆 / v7: 9216 顆"]
    E --> F["跨 Pod：資料中心網路 DCN<br>Multislice 訓練"]

以 TPU v4 為例：基本單位是一個機櫃裡的 4x4x4 cube（64 顆晶片），cube 內部全部用便宜的銅纜直連；要組成更大的 slice 時，再由 OCS 把多個 cube 的光纖接起來。slice 的形狀可以是 4 的倍數的各種組合，例如 4x4x12、4x8x8，甚至不必是 2 的冪次 ^{4 16}。

順帶一提，v4 的命名如 v4-128 指的是 128 個 TensorCore（每顆晶片有 2 個），實際上就是 64 顆晶片、正好一個 cube。

TPU v4 常見的 slice 拓撲形狀：從 2×2×1（4 顆晶片）到一整個 4×4×4 cube（64 顆晶片）

所以，回到開頭的問題：slice 的「幾乘幾」是實際矩陣乘法的分布嗎？

不是。它比較像是晶片的物理座位表。

邏輯網格 vs 物理拓撲：XLA 怎麼把矩陣「貼」到晶片上

訓練大模型時，我們會在邏輯層把運算切分成不同的並行維度，最常見的兩種：

• 資料並行（Data Parallelism）：不同晶片處理不同批次的資料，每輪結束後同步梯度。
• 模型／張量並行（Tensor Parallelism）：單一權重矩陣太大，切成好幾塊分給不同晶片，各算一部分。

在 JAX 裡，這通常表達成一個邏輯網格（mesh），例如把 128 顆晶片宣告成 Mesh(('data', 'model'), (8, 16))：8 路資料並行 × 16 路張量並行 ¹⁷。

關鍵在於：邏輯網格和物理拓撲是兩張不同的圖，由 XLA 編譯器負責把前者「貼」到後者上。因為 ICI 只連鄰居，距離遠的晶片通訊必須一跳一跳穿過中間的晶片 ¹¹，編譯器會盡量把「通訊最頻繁的邏輯維度」對應到「物理上相鄰的晶片軸」：

flowchart LR
    subgraph L["邏輯層（數學問題）"]
        A["巨大的矩陣乘法"] --> B["切分成並行維度<br>data = 8, model = 16"]
    end
    subgraph P["物理層（機房接線）"]
        C["8x16 晶片網格<br>ICI 只連上下左右鄰居"]
    end
    B -->|XLA 編譯器對應| C
    C --> D["每顆晶片內部<br>仍由 128x128 或 256x256 的 MXU 運算"]

以 8x16 的 v5e slice 配上 data = 8, model = 16 的邏輯網格為例：編譯器會把 8 路資料並行對應到物理網格的「列」，16 路張量並行對應到「行」。這樣一來，張量並行需要的高頻寬通訊（每一層 forward/backward 都要交換部分結果）就發生在物理相鄰的晶片之間，而資料並行相對低頻的梯度同步（All-Reduce）也能沿著環面高效繞行。

總結這一段：

• Slice 的乘法（如 8x16）：晶片網路的物理接線結構。
• 矩陣的切分：數學上的張量分片（sharding）維度。
• 實際運作：XLA 把數學切分「貼」到物理拓撲上，追求最短通訊路徑；而不論外層怎麼切，每顆晶片內部最終仍是由 MXU 的脈動陣列吞吐運算。

為什麼不排成一條線就好？1D 拓撲的物理極限

看到這裡你可能會想：搞 2D、3D 拓撲讓工程師算座位表，有夠複雜，為什麼不簡化成一條線（1D）？

答案很直接：一維在物理上撐不起數千顆晶片的規模。想像把 4096 顆晶片排成一條直線：頭尾兩顆晶片要通訊，資料得穿過中間幾千顆晶片，就算首尾相連成環，最遠也要跳 2048 步。改排成 16x16x16 的 3D 環面呢？任兩顆晶片的最遠距離立刻縮到 24 步（每個維度最多繞 8 步）。維度越高，晶片間的平均通訊距離越短，因此 TPU 的設計依賴純粹的幾何學，達到最高運算和成本效益。

那為什麼 GPU 叢集「感覺」不用管拓撲？因為 GPU 也不是 1D，而是把拓撲藏了起來。

比如常見的 NVIDIA DGX 的解決方案，其伺服器內部用 NVSwitch 把 8 顆 GPU 全連接，讓它們像一顆巨大的 GPU；跨伺服器則用 InfiniBand 交換機蓋出胖樹（Fat-Tree）拓撲，例如 DGX SuperPOD 參考架構就是用 Quantum-2 InfiniBand 組成的三層胖樹 ¹⁸。對軟體來說，所有 GPU 彷彿都掛在同一個大交換中樞下，邏輯上是「平的」。

GPU 叢集網路簡化示意：每台 DGX 伺服器內部用 NVSwitch 把 8 顆 GPU 全連接；跨伺服器則由 Leaf 與 Spine 層的 InfiniBand 交換機組成胖樹（Fat-Tree）拓撲，實際的 SuperPOD 使用三層胖樹。

想當然，每種架構都存在一度程度的 trade-off，這類架構最直接的代價是錢。胖樹拓撲要維持任兩點之間的無阻塞頻寬，交換機、光模組、線材的數量隨規模暴增。據估計，網路設備往往佔 GPU 叢集總成本的 15–20%，規模越大比例越可觀 ¹⁹。

TPU 的省錢哲學：ICI + OCS

如果你以為 TPU 完全不用交換機，只要把晶片綁在一起就能擴充，那並不完全正確。更精確的說法是 ⁴：

1. Cube 內（64 顆晶片）：晶片透過 ICI 以被動銅纜直連，組成 4x4x4 的 3D 網格；這一層確實沒有交換機，也就是所謂的 glueless 設計。
2. Cube 之間：靠 OCS（Optical Circuit Switch，光學電路交換）串接；一個 v4 Pod 會用 48 台 OCS 把 64 個 cube（共 4096 顆晶片）連成完整系統。

OCS 和 InfiniBand 這類「電子封包交換機」是完全不同的東西。一般的交換機像郵件分揀中心：每一份經過的資料封包都要被拆開來看一下收件地址，才能決定往哪送，這需要晶片運算，也會增加延遲和耗電。

OCS 則簡單得多：它內部是一片可以轉動角度的微型鏡子陣列（MEMS），工作原理就是把光反射到另一條光纖上，像老式電話總機的接線生把 A 的線直接插到 B 的孔上。光路一旦接好，資料穿過 OCS 時就像走在一條直通的光纖上，全程不拆封包、不看內容。這帶來三個關鍵優勢 ⁴：

• 便宜省電：OCS 與底層光學元件佔整體系統成本不到 5%、功耗不到 3%。
• 拓撲可重組：不用重新拉線，重新設定光路就能把 slice 組成不同形狀，甚至支援「twisted torus」這種變形拓撲，把環面的接線「扭」一下，讓相距較遠的晶片少繞幾步路，對 all-to-all 通訊更友善；使用者要 4x4x12 就給 4x4x12。
• 容錯與調度：某個 cube 故障時，光路直接繞過它接上備援 cube；調度器也不必找「連續」的空閒晶片，任選幾個 cube 拼起來就是一個 slice。

把兩種設計哲學放在一起看：

所以 TPU 把拓撲「攤在你面前」不是為了整工程師，而是一個非常清楚的取捨：用軟體與編譯器的複雜度，換取硬體規模化時的性價比與彈性。

這個取捨之所以成立，是因為深度學習訓練的通訊模式非常「有規律」：每一輪計算，晶片幾乎都在跟固定的鄰居交換資料（例如同步梯度），像工廠流水線一樣重複同樣的動作，而環面拓撲最擅長的正是這種「跟鄰居說話」的工作。反過來說，TPU 的設計也因此高度針對深度學習訓練優化：如果搬到通用運算上，誰跟誰通訊、什麼時候通訊都難以預測，資料往往得在網格上繞遠路，這套設計就未必划算了。

結語

回到開頭那個困惑：TPU 的 slice 為什麼是 4x4x4？一路拆解下來會發現，答案不在數學課本裡，而在機房的接線方式與建置成本。本篇嘗試把重點整理成以下幾點：

• 「幾乘幾」是晶片的座位表，不是矩陣的切法。4x4x4 說的是 64 顆晶片在機櫃裡怎麼排、怎麼接線；矩陣要怎麼切、誰算哪一塊，是 XLA 編譯器把邏輯網格「貼」到這張座位表上的結果。
• TPU 省的是交換機的錢。cube 內用銅纜直連、cube 之間靠 OCS 轉光路；以 v4 為例，OCS 加上光模組、光纖等整套光學元件，只佔系統成本不到 5%。GPU 叢集則據估計花 15–20% 的預算在 InfiniBand 胖樹等網路設備上（NVSwitch 則算在伺服器本身的帳單裡），把拓撲藏起來，換取軟體好寫。
• 複雜度不會消失，只會轉移。Google 省下的硬體錢，變成編譯器與工程師要扛的功課：矩陣維度要對齊 128（或 256），並行切分要貼齊物理拓撲，不然輕則 MXU 空轉，重則 job 直接卡死。

希望這篇文章能為你提供 TPU 架構設計取捨的視角，在評估 AI 基礎設施或閱讀規格表時，能做出更務實的判斷。

參考資源