把 397B MoE 榨進 TPU7x：Google 拿到 4.7 倍加速

把397B 參數的模型放上一顆推論晶片，多數團隊的第一個動作是抄現成的 tensor parallelism 設定，再慢慢調 batch size。Qwen3.5-397B 不吃這一套——它是 397B 參數裡只有 17B 會在某個 token 的前向傳播裡被真正算到的 MoE，官方自己給的比例是 4.3%。60 層裡不是清一色的注意力機制，而是 3 層 Gated DeltaNet 線性注意力配 1 層標準 GQA 的固定節奏，兩種層各自掛著 512 個路由專家（top-10 選取）加 1 個永遠執行的共享專家。Google 這篇 systems engineering playbook 要講的，是這個結構逼著平行策略、kernel 設計、集體通訊全部重寫一輪的過程。

稀疏在哪裡：397B 參數、17B 啟動

稀疏度這件事在論文摘要裡常常只是一行，但它決定了每一步 forward pass 實際要搬多少 byte。GDN 層用 64 個線性注意力的 Value head 配 16 個 Query／Key head，head dimension 是 128；GQA 層反過來，是 32 個 query head 對上「只有」2 個 key-value head（Nkv=2）。這個「只有 2」的數字，是後面整章要處理的引信。512 個路由專家，每個 intermediate expert dimension 是 1024，每次前向傳播用 soft-max 機率選出 top 10，再加一個不管路由結果、每次都跑的共享專家路徑。換算下來，模型看起來像 397B，實際算力需求卻貼著 17B 走——但記憶體 footprint 仍然要放下整個 397B。這個「算得少、放得多」的落差，是 compute-bound 與 memory-bound 兩條優化路徑分岔的起點。

那個「永遠執行」的共享專家還有一層意義：不管路由結果怎麼分佈，這個專家在每個 device 上都固定被算到，等於是幫每次前向傳播墊了一個穩定的計算底噪，不會因為路由不均衡而完全空轉。

硬體那一邊，Ironwood（TPU v7x）每顆晶片有 2 個 Tensor Core、每個 Tensor Core 再配 2 個 MXU（Matrix Execution Unit），時脈 2.2 GHz，FP8 尖峰算力 4,614 TFLOPS，掛 192GB HBM。這組數字不是拿來炫耀的規格表——它是後面所有 roofline 計算的分母。模型的稀疏結構決定了「該搬多少資料、該算多少乘加」，硬體規格決定了「搬跟算各自能多快」，兩邊對不齊的地方，就是接下來要動刀的地方。

這個模型的上下文視窗原生是 262,144 個 token，搭配 YaRN RoPE scaling 可以再拉到超過 1,010,000 個 token。原文沒有明講兩者的因果，但把長上下文的需求跟「60 層裡 3/4 是 GDN 線性注意力」擺在一起看，方向是合理的：線性注意力的狀態大小不會隨序列長度等比例膨脹，GQA 的傳統 KV cache 才會——這也是為什麼超長上下文的推論場景，架構上常常偏好線性注意力。把這幾個數字放在一起看：17B 的活躍參數、2 個 KV head、512 選 10 的路由，任何一個現成推論框架的預設設定都沒辦法直接套用。接下來三節要拆的，分別是這個結構怎麼逼平行策略讓步、怎麼逼 kernel 重寫、怎麼逼通訊協定換一種打包方式。

平行策略的物理極限：2 個 KV head 撞上 TP=8

傳統推論框架的預設答案是 8-way tensor parallelism——把每一層的權重矩陣沿著某個維度切成 8 份，分給 8 個裝置。這招在大多數 attention 層上行得通，但 Qwen3.5-397B 的 GQA 層只有 2 個 KV head。切成 8 份，平均每個裝置分到 2/8=0.25 個 head——這在硬體上根本不存在，官方原文寫得很直接：「fractional head sharding（2/8 = 0.25 heads per device）, which is physically impossible on hardware」。

團隊換了一套切法：attention 部分不切，直接在 8 個裝置上各自複製一份完整權重（Data Parallelism，DP=8），維持每個裝置本地 KV cache 的一致性；MoE 部分反過來切，512 個路由專家平均分配到 8 個裝置，每台裝 64 個（Expert Parallelism，EP=8）。兩種平行各自負責模型的一半，而不是硬逼同一套切法套用在整個模型上。

這裡的 Data Parallelism 用詞跟一般認知的 DP 有一點微妙差異——訓練場景的 DP 通常是把不同的訓練資料切給不同副本、各自算完梯度再同步；這裡的 DP=8 只是單純的權重複製，8 個裝置各自服務不同的請求，彼此不需要同步任何梯度或狀態，是推論場景下更簡化的用法。Expert Parallelism 的切法則更接近直覺——每個裝置就是拿到一個固定子集的專家，永久常駐，不隨請求變動；一個 token 該去哪個專家，純粹看路由結果落在哪個裝置的子集裡，不需要額外的一輪資料搬移決策。

DP=8 底下還有一個容易被忽略的細節：8 份複製的權重雖然完全一樣，但各自維護的本地 KV cache 卻對應不同的請求批次——換句話說，複製的是「能力」，不是「狀態」。這跟傳統 tensor parallelism 把同一個請求的計算拆給 8 個裝置合作完成，是完全不同的思路：那邊是分工，這邊是各自獨立值班。

這個混合設計省下的，是硬套 tensor parallelism 才會出現的記憶體代價——如果 attention 跟 MoE 一起套用舊的 TP 切法，400GB 級的參數 footprint 會在每個節點重複放一份。DP+EP 把這份重複拿掉，同樣的硬體才有餘裕把並發規模往上推，而不是被記憶體先卡死。

值得算一下這筆帳：397B 參數裡，絕大多數躺在 512 個路由專家裡，GDN 與 GQA 的權重相對只佔一小塊。DP=8 把這一小塊 attention 權重複製 8 份，代價可以接受；如果連 512 個專家的權重也跟著複製 8 份，那才是真正會撐爆 HBM 的地方——這正是 EP=8 存在的理由：把體積最大的那塊留給切分，體積最小的那塊留給複製。

Kernel 手術台：RPA、Grouped GEMM v2、GDN 融合

拓樸問題解決之後，剩下的效能差距要靠 kernel 本身去摳。Ragged Page Attention（RPA）原本的 KV block size 是 16 個 token——block 切得太細，每次讀取要處理的索引數量暴增，官方的說法是造成「massive indexing overhead」，讓 VPU（Vector Processing Unit）反覆停頓等資料。把 block size 拉到 256（對應 --block-size=256 這個設定），在 concurrency-512 下把 decode step latency 從 428µs 壓到 283µs，單一個參數帶來 33.8% 的 kernel 級加速。

拖動手把切換 block size · 16 對 256 兩個真實設定值

block=256（優化後）：decode step latency 283µs（↓33.8%）。連續讀取的區塊變大，索引開銷被攤薄，VPU 管線才有機會填滿。

這不是單純調大一個數字就有效的巧合——block 越小，同一段連續記憶體被切成越多獨立的定址單位，VPU 要花更多時間算位址、更少時間做乘加；放大到 256，連續讀取的區塊變大，索引開銷被攤薄，VPU 才有機會把管線填滿。

MoE 層的 SwiGLU 前饋網路也挨了一刀：Grouped GEMM（GMM）v2 把 gating 與 up-projection 兩個矩陣乘法用 dual DMA read 融合進單一個 tile，再搭配 512-subchannel 的 FP8 量化，消掉了原本會發生的暫存器溢位（VREG spills），讓 VPU 的向量運算吞吐量直接翻倍。同一個 kernel 還加了動態邊界切片，讓每個專家實際拿到的 token 數量不一時，padding 也能壓到最低。

MoE 的路由天生會不均衡：512 個專家裡，熱門的幾個會被大量 token 選中，冷門的幾個可能整個 batch 都沒人選。Grouped GEMM v2 的動態邊界切片，就是在處理這種不均衡——每個專家實際收到的 token 數量不同，傳統做法是全部 padding 到最大值，GMM v2 讓每個專家按自己實際收到的數量計算，不平白浪費算力在填充的 0 上。

GDN 層的 Causal Conv1D 原本是獨立的一步：算完中間結果、寫回 HBM、下一步再讀出來接著算。融合後的版本把歷史 token 狀態直接留在 TPU 的 VPU 暫存器裡，用 register-level 的 sliding window 演算法邊算邊維護，原本一次前向要跑的 6 次 HBM 往返直接消失。

路由本身也有一顆專門的引擎可以借：TPU 的 SparseCore 天生擅長不規則的間接定址，團隊把 token 到專家的路由運算整段搬到 SparseCore 上執行，繞開原本會出現在 HBM 裡的大量填充、未對齊中間張量——原文的說法是「Bypasses heavily-padded, unaligned intermediate tensors in HBM, saving massive memory bandwidth」。MoE 路由的本質就是「決定誰去哪裡」的稀疏 gather/scatter，交給為此設計的硬體單元，比硬塞進通用矩陣運算單元划算。

decode 階段還有一個更直接的手法：把多個獨立的 decode stream 打包進同一個 Pallas kernel 呼叫，而不是每個請求各自發一次 kernel。好處官方列了三項——攤薄 VPU 指令派發的延遲、打斷循序請求之間的資料相依停頓、改善記憶體對齊。單一使用者的請求延遲看起來沒變，但當 TPU 上同時跑著上百個並發請求時，這種打包能把每個 kernel 呼叫的固定成本攤到更多實際工作量上。

這也解釋了為什麼 RPA 與 Batched RPA 的效果在高並發時特別明顯：並發數低的時候，系統沒有足夠多的獨立 decode stream 可以打包，固定開銷分攤不到太多請求上；並發數一拉高，打包空間變大，單位請求分攤到的固定成本才會真正下降。

GDN 的遞迴公式本身也被重新推導過：把方程式重新整理之後，原本必要的一次 rank-1 矩陣乘法被完全省略——官方原文是「completely skipped the expensive post-rank-1 matrix multiplication」。同一輪改動裡，遞迴狀態變數的精度從 Float32 降到 BFloat16，VPU 的向量運算吞吐量因此再翻一倍。這類改動的風險是精度：降精度換速度的交易，通常需要額外的驗證才敢上線——這點留到最後一節談。

前面提到的 Conv1D 融合後來又被推進一步：一個完全融合的 Pallas kernel 把 causal 1D convolution 跟整個 GDN recurrent block 合併成一次呼叫，中間序列與遞迴狀態全部留在區域暫存器裡，原文的講法是「Completely bypasses the need to read and write intermediate activation tensors to VMEM or HBM」——不只省了 HBM 往返，連 VMEM 這一層的讀寫都被跳過。搭配這個改動的，是一個把 GDN 遞迴狀態與 GQA 的 KV cache 對齊排列的自訂記憶體佈局，目的是減少 padding 與碎片化，「reclaiming critical HBM headroom」。

勾選切換融合前後 · 6 層 HBM 往返對 1 個融合 kernel

勾選查看融合後的版本

① Conv1D 算輸出 → 寫 HBM

② 讀 HBM → 準備 gate 輸入

③ 算 gate → 寫 HBM

④ 讀 HBM → 準備 recurrent state

⑤ 更新 state → 寫 HBM

⑥ 讀 HBM → 進下一層 matmul

全部留在 VPU 暫存器裡做完，0 次 HBM 往返

集體通訊瘦身：All-Gather 與 Reduce-Scatter

EP=8 把專家攤開到 8 個裝置之後，每個 token 算完要送去哪個專家、算完又要收回來，都要靠集體通訊（collective communication）。原本的路由需要三次獨立的 All-Gather：一次收 token 的 hidden state、一次收專家索引、一次收 gating 權重。團隊把後兩者「stack, bitcast, and pack」成一個 single dense 32-bit 整數陣列，跟第一次 All-Gather 合併，三次變兩次，routing metadata 的 collective latency 直接砍半。

MoE 算完之後的 Reduce-Scatter 也拆成兩段：先在同一顆晶片內部走共享記憶體，速度是跨晶片路徑的 6 倍；再用遞迴倍增（recursive doubling）的方式跨晶片、透過 ICI 互連完成剩下的聚合。整批資料切成 2 到 4 個 micro-batch，讓 DMA 搬運跟計算重疊——搬第 i 批資料的同時，晶片還在算第 i-1 批，通訊不再是純粹等待的死時間。

拆成兩段的理由很直接：同一顆晶片內部的記憶體頻寬，天生比走 ICI 互連跨晶片快上一個數量級——先在近的地方把能合併的合併掉，才把真正必須跨晶片的那一小份資料送出去，是分散式系統裡常見的「先本地聚合、再跨節點聚合」的思路，只是這裡的節點是晶片。

這五項改動沒有一個是孤立生效的：RPA 影響的是 decode 階段每一步的延遲，GMM v2 與 GDN 融合影響的是 compute 端的浮點利用率，All-Gather 與 Reduce-Scatter 影響的是 EP=8 切分之後才會出現的通訊開銷。少了 DP+EP 的平行策略，後面這一整組 kernel 優化根本沒有適用的場景——這是為什麼工程順序要先定拓樸、才動 kernel。把這五個改動攤開來看，各自解決的瓶頸並不重疊：

五個 kernel／通訊層優化各自的機制與量化效果，數字皆取自原文。

技術	機制	量化效果
RPA block size 調整	KV block 從 16 拉到 256，消除 indexing overhead 造成的 VPU 停頓	decode latency 428µs→283µs（↓33.8%）
Grouped GEMM v2	SwiGLU 的 gating／up-projection 融合單一 tile，512-subchannel FP8 量化	VPU 向量吞吐量翻倍
GDN Causal Conv1D 融合	register-level sliding window，歷史狀態留在 VPU 暫存器	消除 6 次 HBM 往返
All-Gather 合併	三次獨立 collective 打包成兩次 32-bit 整數陣列傳輸	routing metadata collective latency 減半
Reduce-Scatter 分層	intra-chip 共享記憶體＋inter-chip 遞迴倍增，2–4 個 micro-batch pipeline	intra-chip 段快 6 倍

Roofline 記分板：兩個效率百分比

Roofline model 的邏輯很直接：算出硬體在某個工作負載下理論上「最多」能跑多快，分成兩種瓶頸——如果乘加運算比資料搬運慢，問題出在算力，是 compute-bound；如果資料搬運比乘加運算慢，問題出在頻寬，是 memory-bound。Prefill 階段要重算整段輸入的 attention，乘加運算量隨輸入長度膨脹，天生偏 compute-bound；decode 階段每一步只算一個新 token，但要把整份 KV cache 從 HBM 掃過一遍，天生偏 memory-bound。這也是為什麼同一顆晶片、同一個模型，prefill 跟 decode 會撞上兩種不同形狀的天花板。

所有拓樸與 kernel 調整最後都要接受同一套檢驗：roofline model 算出的理論上限，實測吞吐量摸到了幾成。Prefill-heavy（8K 輸入／1K 輸出）是 compute-bound 場景，在 concurrency 64 下，理論上限是 5,170 tokens/s/chip（未打折）、排程打折後是 4,500；實測 3,707，官方自己算出來的效率是 82.4%。Decode-heavy（1K 輸入／8K 輸出）是 memory-bound 場景，理論上限 978（未打折）、打折後 850；實測 677，效率 79.6%。兩個場景撞到的天花板不一樣——一個被 FP8 算力卡住，一個被 HBM 頻寬卡住——但摸到的比例意外地接近。

切換 discounted／undiscounted 兩種天花板 · 2 場景 × 3 根柱

對比現實上限（discounted） 對比理論上限（undiscounted）

Prefill-heavy 效率82.4%

Decode-heavy 效率79.6%

discounted 與 undiscounted 這兩個版本的差距，本身就是一個工程判斷：排程開銷、實際 batch 組成、記憶體對齊的損耗，都會讓真正能拿到的算力低於晶片規格表上的峰值。把 82.4% 這個數字念成「還有 17.6% 的手術空間」，還是念成「已經逼近實際可行的極限」，取決於你拿哪一個分母——這正是為什麼原文同時給出兩種 roofline，而不是只挑一個好看的數字。

文章用的「concurrency」指的是同時掛在系統上的並發請求數——64 到 512 這個範圍，對應的是一台推論伺服器在真實流量下可能同時服務的使用者數量。並發數越高，排隊等 batch 湊滿的時間通常拉長，但攤到每個 chip 的吞吐量效率也可能提升，這是為什麼同一組 kernel 優化在不同並發檔位下的效果不會完全一樣。

把鏡頭拉遠到整條時間軸：2026 年 4 月到 6 月之間，同樣的 512 concurrency 檔位，decode-heavy 提速約 3.1 倍、prefill-heavy 約 4.7 倍。原文沒有給這兩個倍數對應的絕對吞吐量，只給了倍數本身——對照的基準，是同一個團隊三個月前還沒套上這整套 playbook 的版本。兩個效率百分比與兩個加速倍數，量的是同一件事的兩種切面：前者是「離物理天花板還差多少」，後者是「跟自己的舊版本比，差了多少」。

數值驗證：省下來的效能有沒有偷跑

所有降精度、融合、跳過運算的改動，都有一個共同的疑問：數字還準不準？團隊為此另外加了一層檢查——「incorporated a dedicated Numerical Verification Layer to audit accumulation precision」。這層驗證不是跑一次測試就結案，而是持續比對 Pallas 降階後的實際計算路徑，跟高精度 Float32 參考路徑之間的差距。

FP8 矩陣乘法通常會在較高精度的累加器裡做加總，再把結果轉回低精度——這是業界處理低精度推論的標準做法，單一乘加的誤差可以接受，但成千上萬次加總疊在一起，誤差會被放大。Numerical Verification Layer 要盯的，正是這個累加路徑有沒有在某個環節悄悄引入偏移。

結果是「verified that our Pallas-lowered gating weights maintain zero deviation from the high-precision Float32 reference path」——路由用的 gating 權重，在降精度、融合、演算法重寫之後，跟原始 Float32 版本的計算結果沒有偏差。路由這件事對精度特別敏感：gating 網路算出的分數只要有一點點漂移，就可能讓 top-10 的排名洗牌，連帶影響哪個專家被選中——這也是為什麼數值驗證特別鎖定 gating weights，而不是籠統地說「整體精度沒問題」。這句話撐起了前面所有 kernel 手術的正當性：如果每一刀下去都要犧牲一點準確度才能換速度，這篇文章談的就不是同一個模型；數值驗證層存在的意義，是讓「快」跟「對」不必二選一。

把數值驗證擺在這裡看，前面 82.4%、79.6% 這兩個效率數字，以及 3.1 倍、4.7 倍這兩個加速倍數，才站得住腳——如果降精度的路徑本身就有偏差，那麼「更快」有可能只是「算得比較粗」，兩者在數字上很難分辨。有驗證層盯著 gating weights 零偏差，至少可以排除這一種混淆。

decode-heavy 情境下，單一個 token step 跨越全部 60 層的總執行延遲，官方算出來是 16.36 毫秒。這個數字是所有上述改動疊加後的總帳——RPA 的 block size、GMM v2 的融合、GDN 的暫存器內運算、集體通訊的兩段式 Reduce-Scatter，每一項各自省下的時間，最後都收斂成這一行每步驟的毫秒數。如果讀者要在自己的叢集上做類似的精度下修，這篇文章示範的不是「大膽下修就對了」，而是「每次下修都配一個可以量測偏差的驗證層」——沒有這一層，前面所有 kernel 手術都只是賭注。

這套系統工程流程被團隊自己定位成一種方法論，而不是一次性的調參紀錄：「Rather than tackling models as monolithic systems, we decompose them into self-contained, independent building blocks（such as Batched RPA, Grouped GEMMs, and SparseCore unpermutation）accompanied by hardware-aware cost models.」每個 building block 各自有自己的 cost model，換一個新模型上來，先套進去看哪一塊卡住，而不必從頭重跑一輪 trial-and-error。

這套方法論解鎖了什麼：把模型拆成 Batched RPA、Grouped GEMMs、SparseCore unpermutation 這類彼此獨立的 building block，每一塊配上自己的硬體感知 cost model——團隊自己的說法是，這樣才不必「每個新模型家族都重新花幾個月做 trial-and-error」。下一個 397B 級的 MoE 模型上 Ironwood，要重複的是查 roofline、換 kernel、調集體通訊這三個步驟，而不是從頭再走一次。