GPU 上的兩個數字
你已經知道如何在 Python 中執行 c = a + b，而且做過上千次了。這裡我們要討論的是在 PyTorch 張量中的運作方式。張量就是一個數值陣列。把它放到 GPU 上，代表這個陣列存在 GPU 的記憶體中，而不是一般的 CPU 記憶體。當 a 和 b 都是 GPU 上的張量時，這一行程式碼執行得飛快，快到你根本不會去想它。
現在把它縮小。假設 a 和 b 是兩個單一的浮點數，都放在 GPU 上。同樣一行程式碼：c = a + b。晶片上實際執行了什麼？
答案是 kernel。在這個領域中，kernel 是一個小程式，讓 GPU 在某一塊資料上執行。不是你筆電開機時的那個 OS kernel，也不是線性代數課本裡的數學 kernel。這個詞被重複使用很多次，那不是你的錯。在 GPU 的世界裡，kernel 就是：一個小函式，被告知要立刻、平行地在你交給它的資料上執行。
到這篇文章結束時，你將能夠看著一段 PyTorch 程式碼，數出 GPU 會執行多少個 kernel。這聽起來像是一個小技巧，事實上也的確是，但這是打破「GPU 是黑盒子」那種感覺的第一步——那種當你的模型跑很慢卻不知道原因的感覺。所有你能問的關於 GPU 效能的問題，最終都會回到「哪個 kernel 跑了，它們在做什麼」。所以我們從這裡開始。
你的第一個 kernel
讓我們把 a 和 b 稍微變大一點。這次是長度為 8 的張量。仍然是同一行 PyTorch 程式碼：c = a + b。
當你執行這行程式碼時，你的 CPU（實際執行 Python 的機器）對 GPU 說：嘿，去執行這個程式，處理這筆資料。這個指令稱為 launch（啟動）。被啟動的是一個 kernel：一個準備好執行的程式。啟動本身很便宜，每次幾微秒。真正昂貴的是圍繞在啟動周圍的部分（資料送到 GPU、結果傳回來），這才是我們要計算的成本。
在 kernel 內部，實際工作是由稱為 threads（執行緒）的小型工作者完成的。GPU 有成千上萬個這樣的執行緒可用。對於我們的長度 8 加法，8 個執行緒接手工作：執行緒 0 處理元素 0，執行緒 1 處理元素 1，依此類推直到執行緒 7。每個執行緒執行相同的小程式：讀取 a 的一個元素，讀取對應的 b 元素，將它們相加，然後將結果寫入 c。
（實際上，GPU 會以固定大小的群組啟動執行緒，稱為 warp，在 NVIDIA 顯卡上總是 32 個執行緒，當你的陣列無法均分時，會遮蔽掉多餘的執行緒。目前暫時可以忽略這個細節。）
所以我們有一行 PyTorch 程式碼、一次啟動、一個 kernel、8 個執行緒執行 8 次加法。現在讓我們計算實際上穿越晶片的資料量。為了執行加法，每個執行緒需要它的 a 元素和 b 元素。那就是 8 次 a 的讀取和 8 次 b 的讀取。然後每個執行緒將結果寫入 c。那就是 8 次寫入。
這些讀寫操作是針對位於 GPU 晶片旁邊的大型記憶體。在資料中心顯卡（A100、H100）上，這種記憶體稱為 HBM（高頻寬記憶體）。在消費級顯卡（RTX 4090）以及人們實際上可能用來測試的 Colab 風格 T4 上，則稱為 GDDR。無論哪種，都是放在晶片旁邊的快速記憶體，我們就稱之為 GPU 記憶體。它很快，但並非免費，每次存取都有成本。
一個 kernel = 一次啟動 = 一次對資料的遍歷。 無論 kernel 內部做什麼，其邊界上的讀寫操作（到 GPU 記憶體提取輸入、寫回輸出）都是成本所在。這就是整體的樣貌。
當張量變大時，這些都不會改變。同樣的 PyTorch 程式碼，同樣的一個 kernel，只是執行緒更多。如果 a 和 b 各有 100 萬個元素，GPU 會啟動同一個 kernel，但使用更大群的執行緒。運算量隨之增加，位元組移動量也隨之增加，但思維模型不變。一行程式碼，一個 kernel。
兩個運算之間發生什麼事
1c = (a + b).relu()
你對 Python 夠熟悉，知道這是兩個運算：一個加法，然後一個 relu，串聯在一起。在直譯器中，這是兩個函式呼叫。在 GPU 上，在 eager 模式下的 PyTorch 中，這是兩個 kernel 啟動：一個用於加法，一個用於 relu。到目前為止，並不令人意外。
真正有趣的是兩個 kernel 之間發生的事。
當加法完成時，它必須把結果放到某處。那個某處就是 GPU 記憶體。加法將一個完整的中間陣列（稱之為 tmp）寫入記憶體。然後過了一會兒，relu 啟動，它的第一個工作就是從記憶體中讀回同一個 tmp 陣列。它讀取整個陣列，對每個元素應用 relu，然後將結果寫入 c。
計算這兩個 kernel 的記憶體流量：
- 加法： 讀取
a，讀取b，寫入tmp。三次陣列大小的傳輸。 - relu： 讀取
tmp，寫入c。再多兩次。
總共五次陣列大小的傳輸。相比之下，前一節中單獨的長度 8 加法只有三次。在鏈中加入 .relu() 不僅僅花費了 relu 的計算成本，還讓你付出了一次完整的陣列往返 GPU 記憶體的成本，因為 tmp 必須被寫出，讓下一個 kernel 才能讀回來。
沒有任何東西被快取。tmp 並沒有留在暫存器或快速的本地快取中。它被送到 GPU 記憶體（較慢且較遠的那種）然後立刻回來。這兩個 kernel 彼此不認識。它們必須透過它們都知道如何溝通的中介來傳遞資料：GPU 記憶體。
為什麼 PyTorch 要這樣做？因為在 eager 模式下，當你寫 a + b 時，PyTorch 會立刻執行它。它不知道你接下來要呼叫 .relu()。每個運算在其 Python 行執行時就被派發。沒有計畫，沒有前瞻。每個運算獨立存在，產生一個真實的陣列，並透過記憶體交給接下來的運算。
這是你要記住的模式。 你 PyTorch 程式碼中的每一個中間值都會被實際寫入 GPU 記憶體，然後被下一個運算讀回來。每一個都是。這就是「kernel 計數」真正衡量的東西。每一個額外的 kernel 都代表你的資料必須多一次往返 GPU 記憶體。
融合：兩個運算，一個 kernel
想像一個 kernel 一次完成所有工作：讀取它的 a 元素，讀取它的 b 元素，相加，對結果應用 relu（全部在 kernel 內部進行，使用每個執行緒的小型私有暫存空間，永遠不會離開晶片），然後才將最終值寫入 c。中間的 (a + b) 仍然存在，但只在 kernel 內部，在每個執行緒的私有暫存空間中。它從未被寫入 GPU 記憶體。tmp 作為一個真實的陣列，根本不存在。
現在計算傳輸次數。讀取 a：每個元素 1 次。讀取 b：每個元素 1 次。寫入 c：每個元素 1 次。三次陣列大小的傳輸。 與兩個 kernel 的版本數學相同，但少了兩次往返。
在長度 8 時，這只是個四捨五入的誤差。沒人在乎。在長度 100 萬或 1 億時，那些額外的往返就會佔據執行時間的很大一部分，而且實際時鐘時間也會反映出來。為什麼記憶體流量會如此主導，是這個系列文章第 2 篇的主題，所以我在這裡不深入探討原因。現在的重點只是：相同的數學，更少的往返，實際執行更快。
這個技巧——將原本會是分開的 kernel 的運算組合成一個單一的 kernel，使得中間值不必訪問 GPU 記憶體——有一個名字。它叫做 fusion（融合）。就是這個詞。就是這個概念。
現在尷尬的部分來了。手動撰寫那個組合的 kernel 看起來很簡單，就是加法加 relu。兩個運算。中間一行「計算」。但真實的 PyTorch 程式碼有數十個運算串聯在一起，每個都有自己的形狀、資料類型和廣播規則。編寫一個能正確處理所有這些情況的融合 kernel 是真正的工程工作。你通常不會手動編寫這些常規的元素級 kernel。
好消息是：PyTorch 提供了一個工具，可以自動為你重寫這類情況。它叫做 torch.compile。
你幾乎肯定看過 torch.compile 這樣使用：
1model = torch.compile(model)
一行程式碼。網路上有人告訴你這樣做會讓程式更快。以下是用白話文說明它實際做了什麼：torch.compile 不會像 eager 模式那樣一次一個運算地執行你的操作，而是捕捉你的函式所執行的張量運算，尋找合併的機會，並產生最佳化的程式碼。後續符合相同假設的呼叫可以重複使用該程式碼。
我們在紙上做的融合（加法與 relu，共用一個 kernel，tmp 從不接觸記憶體）正是 torch.compile 會自動對你的程式碼做的事情，只要運算夠簡單。當人們說 torch.compile「讓 PyTorch 更快」時，這正是他們意思中的一大部分。
對於 torch.compile 無法自行融合的情況（它無法識別的自訂運算、不常見的歸約、奇怪的記憶體佈局），仍然需要有人手動編寫 kernel。這就是 Triton 和 CUDA 等工具的用途。另一篇文章再談。
親眼看看
以上所有內容都是在紙上計數 kernel。現在是時候在真實的 GPU 上計數了。如果你有任何配備 CUDA GPU 的機器（桌機、Colab 筆記本、雲端實例），你可以在幾分鐘內自己執行看看。
使用的工具是 torch.profiler。它內建在 PyTorch 中。它所做的就是記錄 GPU 在執行程式碼時實際做了什麼，並回傳一個你可以閱讀的表格。
步驟 1：eager 版本。
將兩行運算包在一個函式中，以便我們可以呼叫：
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() 只是為了確保 GPU 在我們讀取時間之前已完成工作。GPU 工作是非同步執行的，如果沒有同步，有時你會測量到啟動開銷，而不是實際的 kernel 工作。
步驟 2：讀取輸出。
你的實際分析器輸出可能會有比預期更多的行。會混雜一些記憶體配置和 PyTorch 簿記行。我們關心的行是 CUDA kernel，也就是 GPU 實際執行的函式。尋找名稱中包含 kernel 的行。這兩個 kernel 看起來大致像這樣：
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
確切的模板名稱會隨著 PyTorch 版本而變化（relu 通常顯示為 threshold，因為那是底層的運算，而 add 有時顯示為 CUDAFunctor_add）。不要試圖解析整個名稱。只要數行數就好。兩行。兩個 kernel。 一個用於加法，一個用於 relu。完全符合我們在上一節中說的。
步驟 3：編譯版本。
改變一行程式碼。將函式包在 torch.compile 中：
1compiled = torch.compile(add_relu)
現在有一個你在執行之前需要知道的陷阱。第一次呼叫 torch.compile 編譯過的函式會很慢，有時異常地慢，因為那是 torch.compile 做它的工作的時候：分析你的程式碼，找出要融合的內容，產生融合的 kernel。如果你分析第一次呼叫，你測量到的是編譯步驟，而不是 kernel。所以模式是：先呼叫一次來暖機，丟棄結果，然後再分析。
1compiled(a, b) # 暖機，丟棄結果2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
步驟 4：再次讀取輸出。
這次：
1triton_poi_fused_add_relu_0 1 14us
一行。一個 kernel。 名稱甚至告訴你它做了什麼：融合的加法與 relu。與之前相同的數學運算，一次啟動而不是兩次。
你剛剛做到了這篇文章一直在說的事情，用一句話總結：你要求 PyTorch 將兩個運算組合成一個 kernel，觀察了分析器，並確認計數從兩個變成一個。融合，在真實環境中，在你的機器上。
如果你想看到更明顯的效果，可以在不同的張量大小下試試看。在長度 100 時，兩個版本都跑得非常快，差異會消失在雜訊中。在長度 1000 萬或 1 億時，編譯版本開始明顯領先，因為我們減少的往返在這種規模下是實際工作的一部分。
計數 kernel 不再是抽象的建議。你現在有方法可以驗證了。
總結一下
以下是全部內容的濃縮版。
你的 PyTorch 程式碼在 GPU 上執行時，會轉換成一系列 kernel。每個 kernel 是一次啟動、一次對資料的遍歷、一次往返 GPU 記憶體以提取輸入和寫出輸出。簡單的運算變成一個 kernel。多個運算串聯預設變成每個運算一個 kernel，中間值在記憶體中往返。torch.compile 可以為你融合簡單的運算鏈，使這些中間值永遠不接觸記憶體。較少的 kernel 通常意味著較少的記憶體流量。而較少的記憶體流量通常意味著更快。
