本文依據 CS149 Fall 2025 版。 這是 Stanford CS149 導讀系列第 12 篇，對應 10 月 21 日的第 9 講 Efficiently Evaluating DNNs on GPUs: Transformers and ConvNets，官方投影片 PDF 共 75 頁。
Fall 2025 的錄影只放在 Stanford Canvas。最接近的公開錄影是 2023 Lecture 10: Efficiently Evaluating DNNs on GPUs。只能當部分補充：對照 2023 年同一講的投影片（68 頁），卷積轉 GEMM、分塊矩陣乘、implicit GEMM、融合與 attention 融合都在；2025 版新增了 Triton、ThunderKittens 兩頁、逐列 softmax 的融合例子、ThunderKittens 版 Flash-Attention，也拿掉了 2023 版「用更好的網路拓撲（ResNet、MobileNet）省算力」那一段。本文以 2025 投影片為準。整門課的公開程度是 A3（足以自學），缺口列在系列總覽。
上一篇 PA3 在 GPU 上畫圓，這一篇換成跑神經網路。橋接點是：一個卷積層其實就是一大堆內積，而內積要快，看的還是第 7 篇講過的 arithmetic intensity。
開場：你其實已經懂大半了
投影片第 3–10 頁是「你已經知道、而且該記住的東西」：
- 管線化讓資料搬移與計算重疊（第 4–5 頁）。投影片問兩個問題：這程式受算力還是頻寬限制？重疊要付多少晶片上的儲存？提示是 double buffering：正在算的資料和正在搬的資料，都要有緩衝區放。
- roofline 曲線（第 6–8 頁）：橫軸是 arithmetic intensity，縱軸是吞吐量。左邊斜線是頻寬受限區，右邊平線是算力受限區。算力變強但記憶體不變，轉折點往右移，更多程式會落進頻寬受限區。
- loop fusion（第 9 頁）：
E = D + (A + B) * C分三個迴圈做，每個迴圈兩次讀、一次寫、一次運算，整體 arithmetic intensity 是 1/3；融合成一個迴圈後是四次讀、一次寫、三次運算，變成 3/5。
第 10 頁的總結有四條，最值得記的是後兩條：硬體變快，程式更可能被頻寬卡住；提高程式的 arithmetic intensity，程式更可能變成受算力限制。
第 11 頁接著說：懂了上一頁，你幾乎就懂了現代 AI 軟體面效能最佳化需要知道的一切。剩下的要等下一講，大致是兩件事：搬資料很耗能，以及晶片上拿來做儲存的面積就不能拿來算，所以緩衝區越小越好。
最小的 DNN 介紹：它就是一個電路
第 12–23 頁是給沒學過深度學習的人的速成：
- 一個神經元是 n 個輸入、n+1 個參數（權重加 bias），先算加權和，再過一個非線性函數，例如 ReLU
max(0, x)（第 14 頁）。投影片給了兩種解讀：計算上「它就是一個電路」，機器學習上是一個二元分類器。 - 全連接層可以寫成矩陣乘向量，再逐元素套 ReLU（第 16 頁）。
- 卷積層是局部連接、而且整層共用同一組權重（第 19 頁）。第 17 頁先用一段 C 程式讓你猜它在做什麼，答案是 3×3 模糊。第 20 頁的梯度偵測濾波器說明：一個濾波器就是一個圖樣「偵測器」，輸出像素的大小就是它對周圍區域的「反應」。
- 一層通常同時套很多濾波器（第 21 頁舉 96 個 11×11×3 的濾波器），接著 ReLU、再 pool（取 2×2 區域的最大值），資料量因 pooling 變小（第 23 頁）。
卷積層怎麼做才快
直接寫：七層迴圈
第 25 頁的直接實作是七層巢狀迴圈：batch 中每張圖、每個輸出像素（y、x）、每個濾波器，再對輸入通道與濾波器的空間範圍（y、x）累加。投影片特別點出這裡資料重用很多：濾波器權重在卷積時一直被重用，輸入值也被不同濾波器重用。有重用，就有機會提高 arithmetic intensity。
改寫成矩陣乘：explicit GEMM
第 26–29 頁示範 explicit GEMM（也常叫 im2col）：把每個輸出像素需要的 3×3 輸入排成矩陣的一列，整張圖就變成一個 (W×H) × 9 的矩陣；乘上 9 × 濾波器數的權重矩陣，就一次算完所有濾波器。多個輸入通道時，每列變成 9 × 通道數個元素。
為什麼要這樣繞？第 32 頁回答：GEMM 是全連接層、卷積層、transformer attention 共同的核心 kernel（第 30–31 頁先說明 attention 裡也全是矩陣乘）。GEMM 已經有非常成熟的高效能實作，改寫成 GEMM 就能直接用。
代價寫在第 33 頁：要用現成函式庫，就得把輸入矩陣實際寫出來。對卷積層來說，這會讓 DRAM 流量變成 R×S 倍（R×S 是濾波器的空間大小），也要一大塊額外儲存。
GEMM 本身怎麼寫快：分塊
第 34 頁的三層迴圈矩陣乘加了 #pragma omp parallel for，問題在哪？arithmetic intensity 很低，沒有利用存取 A 與 B 時的時間局部性。
第 35 頁的解法是 blocking（分塊）：一次只算 C 的一小塊，讓對應的 A、B 小塊在計算期間都留在快取裡。投影片留了一個自我測驗：BLOCKSIZE 是不是越大越好？為什麼？（提示：三塊要同時放得進快取。）
第 36 頁把分塊推廣到多層記憶體階層：外層依 L2 大小分塊、內層依 L1 分塊，最內層還有一層為暫存器分塊（投影片註明沒畫出來）。第 37–39 頁再加上 SIMD，比較三種向量化方式：
- 向量化
i迴圈：順便改善 B 的空間局部性，但工作集變大 SIMD_WIDTH 倍，而且在 B 上仍是大步跳著走。 - 若
i維度很小，先把 B 的小塊轉置到暫存緩衝區，再向量化最內層的內積。 - 把 A 與 C 的小塊預先轉置，一次算一個 SIMD_WIDTH × SIMD_WIDTH 的區塊，最內層迴圈的各項彼此不相依。
第 40 頁用 MobileNet 各層的尺寸表提醒：同一個網路的不同層，矩陣形狀差很多，最好的排程也不同。 投影片的評語是「對函式庫作者來說很痛苦」。
不要把整個矩陣寫出來：implicit GEMM
第 42–43 頁的改良是 implicit GEMM：不把完整的卷積矩陣寫到 DRAM，而是每次只在 GPU 晶片上的共享記憶體裡組出一小塊，再交給調校好的共享記憶體 GEMM 程式碼算。這樣不需要額外的晶片外儲存，也不增加 DRAM 流量。投影片指向 NVIDIA 的 CUTLASS。
寫分塊程式的工具
第 44–47 頁介紹三個工具，都是為了讓你不必從零手寫分塊 kernel：
| 工具 | 投影片怎麼說 |
|---|---|
| CUTLASS | 自己實作高效能 DNN 層的基本元件，適合 cuDNN 沒特別調過的特殊尺寸；有共享記憶體 GEMM、warp 層級 GEMM、分塊載入的 iterator、tensor reduction 等 |
| Triton | 語言層級支援 tensor 的載入與儲存：把資料以「塊」為單位載入 GPU 共享記憶體，再對整塊做資料平行運算；第 46 頁的完整矩陣乘範例有兩層分塊 |
| ThunderKittens | 以 tile 為單位的 CUDA 程式庫，目標是讓進階開發者（投影片說「CS149 程度的人」）寫分塊程式更有生產力；支援非同步載入／儲存 tile 與進階記憶體布局 |
第 51–54 頁補上另一條路：直接用晶片廠商的函式庫。投影片提到 AWS NKI（系列第 15 篇的 PA4 會用到），以及 cuDNN 的卷積可以選不同演算法。
為什麼「更多工作」比較快
第 48–49 頁拿 V100 回顧（80 個 SM、6 MB L2、900 GB/s 的 HBM）：GPU 要「很多平行工作」才填得滿。投影片比較兩種卷積輸出量：N=1、64×64 輸出時只有約 52 萬個輸出（2 MB）；N=32、256×256 輸出時有 256M（2²⁸，約 2.7 億）個輸出（1 GB）。batch 越大、圖越大，越容易把機器餵飽。
融合：別把中間結果寫回記憶體
conv + scale/bias + max pool
第 56 頁描述一串常見的運算：Conv → Scale/Bias → Max Pool。投影片要你想像：把 1 GB 的卷積輸出寫回記憶體，讀回來只為了逐元素縮放，再讀一次做 pooling——頻寬成本很驚人。
更好的做法：
- 逐元素的 scale + bias，可以在卷積算出每個元素的當下順手做掉（第 57 頁把它寫進卷積迴圈的最後一行）。
- max pool 的輸出，只要每個 2×2 區域的卷積輸出算完就能得到。
第 57 頁留了課堂練習：怎麼把 max pool 也融合進來？提示是「怎麼把輸出像素的那兩層迴圈分塊」。
逐列 softmax
第 58–59 頁的第二個例子是對矩陣每一列算 softmax。直接照定義寫要分好幾步（先求最大值 m(x)、算 f(x) = e^(x−m(x))、加總 l(x)、再相除），每一步都把整個 M×N 矩陣讀寫一遍。投影片算出直接寫法讀 5MN + 2M、寫 3MN + 2M 個元素；融合版「每列載入一次、算完整個 softmax、存回去」只讀 MN、寫 MN，前提是一列的工作集放得進晶片上的儲存。
融合版 attention
第 60–64 頁是這一講的高潮。attention 的輸入是三個 N×d 矩陣 Q、K、V（N 是序列長度，d 是嵌入維度）：先算 S = QKᵀ（N×N），逐列 softmax 得 P（N×N），再算 O = PV（N×d）。第 61 頁點出麻煩：N 在長序列時可達數千，直接寫要 N² 的空間。
第 63 頁：softmax 為什麼可以分塊算
把一列 x 切成兩段 x⁽¹⁾、x⁽²⁾。整列的最大值是 m(x) = max(m(x⁽¹⁾), m(x⁽²⁾))。每段先各自用自己的最大值算 f 與 l，合併時乘上修正因子 e^(m(x⁽ⁱ⁾) − m(x)) 就對了：
- f(x) = [ e^(m(x⁽¹⁾)−m(x)) · f(x⁽¹⁾), e^(m(x⁽²⁾)−m(x)) · f(x⁽²⁾) ]
- l(x) = e^(m(x⁽¹⁾)−m(x)) · l(x⁽¹⁾) + e^(m(x⁽²⁾)−m(x)) · l(x⁽²⁾)
所以不用先看到整列，也能逐段累積出正確的 softmax。
第 64 頁的融合版演算法：對每個 Q 的區塊 Qᵢ 與每個 K、V 的區塊，載入 Qᵢ、Kⱼᵀ、Vⱼ、Oᵢ，算 Sᵢⱼ = QᵢKⱼᵀ，逐列算 m、f、l，再把 PᵢⱼVⱼ 乘上適當的縮放後累加到 Oᵢ。投影片列出三個效果：
- 省記憶體：永遠不寫出 N² 的矩陣。
- 省頻寬：一次讀三個區塊、做兩次矩陣乘加上幾次列加總，累加到留在快取裡的 O 區塊，arithmetic intensity 高。
- 代價：比原版多一點計算，因為每一步都要重新縮放之前累加的 O。
第 65 頁展示用 ThunderKittens 寫的 Flash-Attention。原理只講到這裡就夠用；想看 FlashAttention 的完整推導與後續版本，站上有 CMU 11-868 的 FlashAttention 導讀。系列第 18 篇的 PA5 會要你真的在 H100 上寫一版。
框架裡的融合
第 66–69 頁說明融合在框架裡的演進：早期是函式庫作者手寫幾個固定的「融合運算」（投影片舉 TensorFlow）；後來 cuDNN 的 backend 由編譯器產生新的實作，把多個運算融成一個節點，不必透過記憶體傳中間結果；現在則有很多以編譯器自動排程 DNN 運算的嘗試，投影片列了 torch.compile。
其他招數與收尾
第 70 頁：低精度。16 位元與 8 位元的權重和中間值已經很常見，正在往 4 位元走，極端情況是 1 位元。
第 71 頁把最佳化手段分三類：
- 更好的演算法：手動設計更有效率的模型（深度、濾波器寬度與數量、stride），也常用自動搜尋找有效率的拓撲。
- 軟體最佳化：把關鍵運算排程好——分塊、融合。通常由人手動調，但有很多研究想自動化。
- 近似：壓縮模型，例如降低位元精度。
第 72–74 頁用兩個問題收尾。GPU 為什麼適合跑 DNN？高 arithmetic intensity 的矩陣運算吃得到 GPU 大量的 FLOPS，而且有 cuDNN 這類高度最佳化的 kernel 庫。GPU 為什麼可能不是最理想的平台？提示是：真的需要一顆通用處理器嗎？第 75 頁預告下一講：用專用硬體加速 DNN 推論與訓練，列出 TPU、Apple Neural Engine、Cerebras、SambaNova、有 Tensor Core 的 Ampere GPU 等。
這一講留給你的三個習慣
- 看到一個 DNN 層，先問它能不能變成 GEMM。 能的話，你就能借用最成熟的那套最佳化。
- 看到兩個相鄰的運算，先問中間結果有沒有必要回到 DRAM。 沒必要就融合。
- 看到一個「要先看完整列才能算」的運算，想想能不能像 softmax 一樣分段累積。 FlashAttention 就是這麼來的。
今晚可以做的事：照第 9 頁的方法，算一個你常用的層序列（例如 linear → bias → GELU）不融合與融合時各讀寫多少元素、做多少運算，比較兩者的 arithmetic intensity。
延伸閱讀：想看 Triton 怎麼實際寫分塊 kernel，接著讀站上的 CS336 Kernels 與 Triton 導讀；想看同一套觀念在 LLM 系統課裡的樣子，讀 CMU 11-868 GPU 程式設計與加速。
系列導覽：上一篇 PA3 + Written 2：CUDA circle renderer｜下一篇 L10 硬體專用化與 DNN 加速器設計｜系列總覽
參考資料
Loading...