別猜瓶頸：TPU 的 Roofline 怎麼讀

看到 step time 卡住的時候，多數工程師的第一反應是猜——是不是 collective communication 太慢？還是 HBM 頻寬不夠？還是純粹算力不足？Google 在 7 月底的一篇文章裡介紹了 accelerator-microbenchmarks 這套已公開的工具（GitHub：AI-Hypercomputer/accelerator-microbenchmarks），主張把這個猜測的過程換成量測。文章開場把方法論講得很直白：「Determining the empirical performance of TPUs requires moving beyond product specifications to granular, real-world measurements.」——TPU 的真實效能，得靠拆到更細分項的實測，不能只看規格表上的峰值。這篇文章要講的是這套量測方法怎麼運作：量什麼、怎麼分類瓶頸，以及一個真實案例裡，量出來的結果怎麼變成訓練 step time 減少 21.2% 的具體改動。

step time 卡住的時候，你通常是在猜

Google 團隊在文章裡舉了一個具體案例：「In one example, we analyzed the performance tuning of a 110B Mixture-of-Experts (MoE) training workload deployed on a 4x4x4 TPU 7x configuration, leveraging microbenchmarks to establish hardware baselines and drive optimization for complex, sparse workloads.」110B 參數的 MoE 模型，訓練跑在 4x4x4 拓撲的 TPU 7x（也就是 Ironwood）上。案例的重點不是「這個模型有多大」，是團隊怎麼從「step time 不夠快」這句抱怨，走到「routing 卡在 memory-bound、forward pass 卡在 compute-bound」這句診斷——中間靠的不是經驗法則，是把整條 pipeline 拆成五個可以獨立量測的面向。

文章把這個階段稱作「establish hardware baselines」——在真正動手最佳化之前，先確認這批硬體在理想情況下能跑到多快、五個面向各自的數字長什麼樣子。沒有這個基準線，後面看到的任何進步都無法判斷是真的把瓶頸解決了，還是量測本身的雜訊。

這裡說的 step time，指的是訓練迴圈跑完一個 batch、更新一次參數所花的時間——forward pass 算出預測、backward pass 算出梯度、optimizer 更新權重，一個 step 走完這整圈才算數。step time 越短，單位時間內能跑的 step 越多，同樣的訓練預算下能看到更多資料、迭代更多次。這也是為什麼 21.2% 這個數字直接對應訓練成本與時間：同一批硬體、同一個訓練目標，step time 少兩成，整個訓練排程也跟著縮短。

MoE 模型稀疏在什麼地方，要先說清楚：一般的稠密模型每個 token 都會經過全部參數；MoE 模型每個 token 只會被路由到少數幾個 expert，大部分參數對這個 token 來說是閒置的。這帶來的後果是，同一個模型裡不同運算路徑的效能特性天差地遠——dense-core 的計算量固定、行為可預期，routing 決定哪些 expert 被啟用，運算量隨輸入變動，這正是為什麼團隊要先做細分項量測，而不是直接看一個模型級的平均數字。

規格表給的是峰值，答不出瓶頸在哪

一顆晶片的 TFLOPS 峰值規格，量的是理想條件下、MXU 全速運轉、資料永遠準時送到的上限——這是這類規格數字的通用性質，不是這篇文章獨有的主張。真實訓練 pipeline 很少符合這個理想條件：collective communication 要等、HBM 頻寬有限、host 端搬資料要花時間。同一顆晶片跑同一個模型，某個 kernel 可能貼著峰值跑，另一個 kernel 可能只用到峰值的三成——峰值數字本身分不出這兩種情況。

峰值數字量不出這件事，跟 MXU 這種硬體的排列方式有關：systolic array 是由一整排乘加運算單元排成陣列，資料像水流一樣從一端流進去，經過整排單元、從另一端流出來，每個單元同時做一次乘加。這種排列只有在陣列被完全填滿、資料源源不絕流入時才跑得動峰值——運算矩陣的形狀得配合陣列大小，形狀不對，陣列填不滿，實際吃到的算力自然比峰值低。這是硬體排列方式決定的限制，不是軟體寫得好不好的問題。

單一個 MFU 數字（實際打到的 FLOPs 除以硬體理論峰值）也回答不了這個問題——MFU 低，可能是算力真的不夠，也可能是 collective 通訊卡住讓 MXU 空等、或是 HBM 頻寬跟不上餵資料。同一個 MFU 讀數背後可以是三種完全不同的原因，這正是為什麼不能只看一個彙總指標，得往下拆到子系統層級。accelerator-microbenchmarks 把這個問題拆成五個各自獨立的量測面向，對應 TPU pipeline 裡五個不同的硬體子系統。網路面向量 collective 通訊：「Key tests include all-gather, all-reduce, reduce-scatter, and all-to-all. Metrics include throughput (GB/s), and latency (seconds).」運算面向量 GEMM：「The primary operation is General Matrix Multiply (GEMM/Matmul), measured in TFLOPs and Model FLOPs Utilization (MFU).」記憶體面向只看一個數字：「The core metric is memory bandwidth measured in GB/s.」host transfer 量的是 CPU 與加速器之間的搬運：「Host-to-Device (H2D) and Device-to-Host (D2H) transfers, measured in GB/s」。最後是針對推論路徑的 RPAv3：「Ragged-Paged Attention (RPAv3): Focuses on benchmarks tailored for transformer-based models, specifically measuring throughput for primitives like Attention (tokamax splash) and Batch Matrix Multiplication (BMM).」五個面向各量各的，五個數字擺在一起，才看得出瓶頸落在哪一塊。

面向	量什麼	指標
collective 通訊	all-gather · all-reduce · reduce-scatter · all-to-all	throughput（GB/s）、latency（秒）
運算（GEMM）	General Matrix Multiply／Matmul	TFLOPs、Model FLOPs Utilization（MFU）
HBM 頻寬	memory bandwidth	GB/s
host transfer	Host-to-Device／Device-to-Host	GB/s
RPAv3	Attention（tokamax splash）、Batch Matrix Multiplication（BMM）	throughput

從這五個面向各自量什麼，可以合理推出「數字不好看」時該往哪裡查：collective 通訊的吞吐量低、延遲高，通常代表 sharding topology 或 overlap 策略有問題；GEMM 的 TFLOPs 或 MFU 低，指向 kernel 實作或 operand 形狀不對；HBM 頻寬量出來遠低於硬體上限，代表資料存取模式該調；host transfer 的 H2D／D2H 數字差，通常是餵資料的 input pipeline 拖慢了訓練；RPAv3 的吞吐量低，則是推論路徑本身的 attention 或 BMM 實作效率不夠。五個面向各自指向不同的修法，這也是為什麼不能只看一個總分。

H2D 指的是資料從 CPU 端的 host 記憶體搬進加速器；D2H 是反方向，把運算結果搬回 host。這條路徑通常是 input pipeline 的最後一段——前面的資料前處理、shuffle、batch 組裝都在 CPU 上做，做完才透過 H2D 送進 TPU。這段搬運如果慢，TPU 端就算算力、頻寬都沒問題，也會因為等資料而閒置。

這五個面向裡，只有 collective 通訊同時報 throughput 與 latency 兩個數字。文章沒有解釋為什麼，合理的推測是：訓練 pipeline 裡的 collective op 常常在等待而不是在跑，吞吐量高不代表等待時間短，一次 all-reduce 如果延遲高，其他晶片一樣得閒置等它做完。GEMM 或 HBM 頻寬這類面向主要在乎搬了多少、算了多少，因為這些操作通常連續發生、可以互相重疊，單次延遲不是主要問題。

這些量測不是紙上談兵——repo 裡的 collective benchmark 直接用 JAX 的 shard_map 實作：「psum, all_gather, all_to_all, reduce_scatter (using shard_map)」，跑起來的是真的 collective op，數字不是模擬出來的。

RPAv3 名字裡的 ragged 與 paged 各自對應一個工程問題：推論的時候，同一個 batch 裡每個請求已經生成的 token 數不一樣長，把它們硬湊成同一個矩陣形狀會浪費算力，「ragged」就是允許每個請求維持自己實際的長度；「paged」借用作業系統分頁記憶體的概念，把 KV cache 切成固定大小的頁面，需要的時候才分配，減少長度不一時造成的記憶體碎片。文章把 RPAv3 獨立列成一個量測面向，量的是 Attention（tokamax splash 實作）與 Batch Matrix Multiplication 的吞吐量。

前面四個面向——collective、GEMM、HBM 頻寬、host transfer——量的是訓練時期的行為；RPAv3 量的是推論服務階段的行為。同一顆晶片、同一個模型，訓練與推論的瓶頸型態通常不一樣：訓練時 batch size 大、運算連續，容易撞上 compute-bound 或 network-bound；推論時每個請求長度不同、常常小 batch，記憶體存取型態更破碎，memory-bound 出現的機率更高。文章沒有說明它為什麼不把 RPAv3 併進 GEMM 或 collective 一起算；合理的推測是推論路徑的存取模式跟訓練路徑差得夠遠，混在一起量會失去意義。

Roofline：一張圖把瓶頸分成三種

量完五個面向的數字之後，下一步是分類——文章用的是 Roofline 模型，把結果歸到三種瓶頸類型之一。compute-bound 的判準是：「The model is hitting the MXU plateau. Tuning memory access will yield no benefit.」如果瓶頸真的在算力天花板，花時間調記憶體存取模式是白工。memory-bound 的判準是：「Performance can be improved by optimizing data locality (VMEM) or reducing HBM traffic.」這裡用的是「can be improved」，不是保證——調 VMEM 資料局部性或減少 HBM traffic 有機會改善，但不一定，得看實際卡在哪個環節。network-bound 的判準是：「Stalls in ICI or DCN prevent the chips from reaching their compute potential.」晶片間（ICI）或資料中心內（DCN）的通訊停等，讓算力潛力打不出來。

這三種分類不是互斥的選項，而是分別檢查：compute-bound 要看是不是打到 MXU 天花板，memory-bound 要看 HBM 頻寬夠不夠，network-bound 要看 ICI／DCN 有沒有停等。三個問題可以對同一個模型的不同部分分別問——110B MoE 案例就是最好的例子，forward pass 的答案是 compute-bound，routing 與 attention 的答案是 memory-bound，同一次訓練裡兩個答案都成立。

滑鼠移到縮寫上看完整說明 · 6 個縮寫

這幾個縮寫會反覆出現：MXU矩陣乘法運算單元，TPU 內執行 GEMM 的硬體核心。、VMEM晶片上的向量記憶體，比 HBM 快、但容量小得多。、ICIinter-chip interconnect，同一個 pod 內晶片之間的高速連結。、DCNdata-center network，跨機櫃、跨 pod 的網路層。、MFUModel FLOPs Utilization——實際打到的 FLOPs 除以硬體理論峰值。，以及SOLspeed-of-light，某個 kernel 在理想條件下能跑到的效能上限；文章用「佔 SOL 的百分比」表示量測結果離這個上限有多近。。

Roofline 的三種瓶頸分類

• compute-boundMXU plateau

  「The model is hitting the MXU plateau. Tuning memory access will yield no benefit.」——瓶頸在算力天花板，調記憶體沒用。

• memory-boundVMEM／HBM

  「Performance can be improved by optimizing data locality (VMEM) or reducing HBM traffic.」——調 VMEM 局部性或減少 HBM traffic，有機會改善。

• network-boundICI／DCN

  「Stalls in ICI or DCN prevent the chips from reaching their compute potential.」——晶片間或跨機櫃的通訊停等，打不出算力潛力。

Roofline 模型本身是效能工程裡的通用工具（這是背景知識，不是這篇文章的獨有主張）：橫軸是 arithmetic intensity，愈往右代表每搬一個 byte 能做的運算愈多；縱軸是實際打到的效能。這個模型的核心公式很簡單：可達到的效能等於算力天花板與（arithmetic intensity 乘以記憶體頻寬）兩者的較小值。intensity 低的時候，第二項比較小，效能被記憶體頻寬卡住，這時候不管硬體算力多強都沒用；intensity 高到一定程度，第二項超過算力天花板，效能就被天花板卡住，這時候記憶體再快也沒用。兩條線的交叉點，就是圖上那個從斜率轉平的轉折點。低 intensity 的 kernel 沿著記憶體頻寬的斜率往上爬，直到撞上算力天花板——這條天花板，就是前面說的 compute-bound 判準。三種瓶頸類型分別對應圖上不同的位置：靠著斜率爬升代表 memory-bound，貼著天花板代表 compute-bound；network-bound 則是另一種情況——理論上有機會貼近天花板，卻被 ICI 或 DCN 的停等硬生生攔在半路，看起來像是天花板被人為壓低了一截。

這套框架不是 TPU 專屬的——GPU、CPU、甚至網路硬體的效能分析都在用同一套邏輯，差別只在天花板由哪個硬體規格決定、斜率由哪種記憶體頻寬決定。放在 TPU 上，天花板是 MXU 的峰值算力，斜率是 HBM 頻寬；換成別的硬體，同樣的圖形換一組數字，判讀邏輯不變。

110B MoE 在 4x4x4 Ironwood 上：兩種瓶頸同時出現

回到文章開頭那個 110B MoE 案例。Ironwood（TPU7x）的硬體本身就對某些運算形狀設了限制：「For example, on the Ironwood (TPU7x) architecture, the 256x256 systolic array imposes a physical constraint on operand shapes.」256x256 的 systolic array 是 compute-bound 天花板的具體來源之一：operand 形狀不對，MXU 沒辦法全速運轉，天花板自然比理論峰值低。

團隊在這個 110B MoE 案例裡量到兩種瓶頸同時發生：「The forward-pass and dense-core operations were primarily compute-bound, reaching 1.85 PFLOPS.」forward pass 與 dense-core 運算貼著算力天花板跑，量到 1.85 PFLOPS；同一個模型裡，MoE 架構特有的 routing 與 attention 原語卻卡在 memory-bound 區，只打到 SOL 的三成到六成。同一個模型、同一次訓練，兩個部分落在 Roofline 圖上完全不同的位置。

文章沒有另外拆解 backward pass 的 Roofline 分類——量到的是 forward pass 與 dense-core 部分落在 compute-bound。訓練一個 step 除了 forward，還有 backward 與 optimizer update 兩個階段，這兩段各自的瓶頸類型，文章沒有公佈對應的數字，也是這篇案例目前沒有涵蓋到的部分。

合理的推測是，這個落差跟 MoE 架構本身的計算型態有關：dense-core 的矩陣乘法是規則的大矩陣運算，資料存取模式固定，容易把 MXU 填滿；routing 決定每個 token 要送去哪幾個 expert，屬於 gather／scatter 一類的不規則存取，attention 也要頻繁讀寫 KV cache，這類存取模式對 HBM 頻寬的依賴比對算力的依賴更重。文章本身沒有把這個機制寫出來，只公佈了量測結果——但落在哪個象限，跟這兩類運算存取記憶體的方式很不一樣，是說得通的。

110B MoE 案例的兩個量測點在 Roofline 圖上的位置

這張圖想說的是同一件事的兩種讀法：forward pass 落在天花板上，代表這部分的效能已經頂到硬體極限，除非換更快的晶片，不然沒有太多再進步的空間；routing 與 attention 原語落在斜率區，代表這部分還有調整空間——沿著斜率往右移，不用換硬體，靠 kernel 或存取模式的調整就有機會把效能往上推。同一個 Roofline 分類，對應到完全不同的下一步動作，這正是分類本身的價值所在。

把這兩個數字放在一起看，能大致理解為什麼最後的改善是 21.2%，而不是更高或更低：forward pass 已經貼著天花板，能擠出的空間有限；routing 與 attention 原語還在 SOL 三到六成，代表這裡才是真正有空間可以拉的地方。合理的推測是，介入的重心會落在後者——把落後的部分拉起來，比在已經頂到天花板的部分硬擠，報酬更高。文章沒有把 21.2% 分攤到個別介入上，只公佈了合計後的結果。

分類清楚之後，團隊照 Roofline 診斷去做介入，結果是：「Microbenchmark results directly guided interventions that reduced the training step time by 21.2%.」量測結果直接導向的改動，把訓練 step time 砍了 21.2%。想自己重現這類量測，repo 裡就有現成的執行方式，用 bazel 帶一份 YAML config 跑：

// configs/sample.yaml 節錄
global:
  warmup_tries: 2
  num_runs: 5
  dtype: "bfloat16"

benchmarks:
  - name: hbm_bandwidth
    sweep:
      size:
        start: 1024
        end: 8192
        multiplier: 2

bazel run //src/accelerator_microbenchmarks:main -- --config configs/sample.yaml

跑完之後，repo 輸出兩份檔案：「detailed.json and summary.csv」——前者是每一筆量測的明細，後者是彙整過的摘要，兩份都是把數字餵進 Roofline 圖之前的原始輸入。

README 只列出 warmup_tries 與 num_runs 兩個參數名稱，沒有解釋語意；照 benchmark 的一般慣例推測，warmup_tries 是讓 JAX 先把運算圖編譯完、把資料搬進暫存，跳過第一次執行常有的編譯延遲，num_runs 則是正式量測要跑幾次，取多次結果才不會被單次的雜訊誤導。size 的 sweep 範圍（1024 到 8192，倍增）則是讓同一個 benchmark 在不同資料量下各跑一次，畫出頻寬隨資料量變化的曲線，而不是只看單一個點。

合理的起手式是先跑 HBM 頻寬：這是設定最簡單、跑最快的一項，config 裡只要換 size 的 sweep 範圍，幾分鐘內就能拿到這批硬體的頻寬曲線，作為後面比對 collective 與 GEMM 數字時的基準。collective 通訊牽涉到多晶片協調，GEMM 要挑對 kernel 與資料型別，設定起來都比單一晶片的 HBM 頻寬測試複雜，適合排在後面——但這是使用上的合理猜測，文章本身沒有規定跑的順序。

量出瓶頸之後，能轉的三個旋鈕

分類只是第一步，後面要對症下藥。針對 compute-bound、或懷疑實作效率不夠的 kernel，文章給的做法是換實作比較：「You can compare standard JAX operations against specialized Pallas kernels or Tokamax Splash Attention.」拿標準 JAX op 跟手寫的 Pallas kernel 或 Tokamax Splash Attention 對比，看哪個實作能把同一個運算塞進更接近天花板的效能。

Pallas 是 JAX 生態裡讓工程師直接寫 TPU kernel 的框架，比標準 JAX op 更能控制記憶體存取模式與並行方式；Tokamax Splash Attention 是針對 attention 運算寫的專門 kernel，對應的是 flash-attention 那類「減少 HBM 讀寫次數」的最佳化思路。標準 JAX op 泛用但不見得貼近硬體特性，專門 kernel 換來的是更接近天花板的效能，代價是維護成本更高。

針對 network-bound、或需要重新分配運算的情況，做法是調拓撲：「By measuring collective performance across different topologies (e.g., 2x2x2 vs 4x2x1), you can optimize Fully Sharded Data Parallel (FSDP) parameters.」同一批晶片，用 2x2x2 還是 4x2x1 的拓撲排列，collective 通訊的效能會不一樣，量出差異之後才知道 FSDP 的分片參數該怎麼調。

2x2x2 與 4x2x1 這兩種寫法，描述的是同一批晶片在三維 mesh 上怎麼切：每個維度的晶片數量不同，collective 通訊要走的實體連線長度與跳數也不同，量出來的 all-gather、all-reduce 吞吐量自然不一樣。FSDP 的分片維度如果跟通訊效能差的那個維度對齊，就會把慢的那條路徑用得最頻繁——這正是為什麼要先量再決定怎麼切，而不是憑經驗選一個看起來合理的拓撲。

針對 memory-bound、尤其是 HBM 頻寬吃緊的情況，還有第三個旋鈕：要不要把算過的中間結果存在記憶體裡，還是乾脆丟掉、要用的時候重算一次。「Microbenchmarks help calculate the exact trade-off between compute cycles and memory pressure, guiding decisions on when to store calculated data in memory vs. 'rematerialize' (recalculate) it.」這個 trade-off 沒有標準答案：存記憶體省算力、但吃頻寬；重算省頻寬、但吃算力。量出實際的 compute cycle 與 memory pressure 數字，才知道哪邊划算。

memory-bound 除了 rematerialization，還有另一個方向是判準裡提到的 data locality：VMEM 是晶片上、離 MXU 最近的一層記憶體，比 HBM 快一個數量級以上，但容量小得多。優化 data locality 的意思是盡量讓運算需要的資料留在 VMEM 裡重複使用，減少來回 HBM 搬資料的次數——同一份資料如果要用兩次，寧可留在 VMEM 裡等第二次用，也不要用完就丟、下次再從 HBM 重新搬一次。這跟 rematerialization 是兩個獨立的旋鈕：一個決定資料要不要留著，一個決定留著的話放在哪一層記憶體。

三個旋鈕對應三種不同的瓶頸診斷，這也是這套方法論真正的重點：110B MoE 這個案例本身是單一個案（in one example），21.2% 的 step time 改善也是這一次訓練跑出來的結果，不是可以直接套用到任何 workload 的固定倍率。真正可以複製的是流程本身——量五個面向、用 Roofline 分類、再挑對應的旋鈕，而不是看到 step time 慢就照著別人的改法抄一遍。對正在跑大型模型訓練的團隊來說，這套流程可以直接套進自己的 pipeline：先把 collective、GEMM、HBM、host transfer、RPAv3 這五個面向的數字量出來，畫進同一張 Roofline 圖，看每個 kernel 落在哪一區，再決定值不值得花時間換 kernel、調 sharding topology，或是接受 rematerialization 的計算換記憶體代價——量測本身不貴，貴的是分類錯了、把力氣花在調一個其實沒有卡住的環節。

這套流程對負責訓練基礎設施的團隊（而不是模型研究本身）特別有用：模型研究關心的是準確率、loss 曲線；基礎設施團隊關心的是同樣的模型能不能用更少的晶片、更短的時間跑完，而這正是 Roofline 分類直接回答的問題。兩邊的目標不衝突，但用的是完全不同的一套指標。

這篇文章由 Junjie Qian、Chi Shuen Lee、Yu-Hsuan（Amy）Lin、Haixiong（Sean）Wang 四人具名，發表於 2026 年 7 月 30 日的 Google Developers Blog，文章介紹的 accelerator-microbenchmarks repo 是公開的——文章裡引用的每個數字，理論上都可以拿這個 repo 的 config 重新跑一次去對照。

Take-away：瓶頸不是用猜的——先照五個面向量出數字，再用 Roofline 把每個 kernel 歸到 compute-bound、memory-bound、network-bound 其中一類，最後才決定要換 kernel、調 sharding，還是拿 rematerialization 換頻寬。