本系列文章「PyTorch 效能剖析」的第一部分，我們使用 torch.add(torch.matmul(x, w), b) 來學習如何閱讀 PyTorch 剖析器追蹤。我們也討論了其他相關主題，例如 CPU 調度鏈、啟動開銷、開銷限制與計算限制之間的差異，以及 torch.compile 的一些內部機制。
在第二部分（即本篇文章），我們將更進一步。我們將手寫的矩陣乘法與加法組合替換為 nn.Linear（設定 bias=True），這是每個深度學習模型都會使用的基本建構區塊。接著，我們將三個 nn.Linear 層堆疊起來，並在中間加入一個激活函數，以形成一個多層感知器 (MLP) 區塊。
本篇文章的程式碼位於：02_linear.py、03_simple_mlp.py 和 03_kernels_mlp.py。如同之前，建議您在閱讀時另開分頁查看這些程式碼。我們使用 NVIDIA A100-SXM4-80GB GPU 來執行這些腳本，您可以在 Hugging Face 的基礎設施上，透過 Spaces 的開發模式輕鬆設定 GPU 並實驗這些腳本，或者使用 Hugging Face Jobs 管線來執行。
在開始之前，快速回顧兩個我們將反覆依賴的概念：
GPU 核心（kernel）是一個在 GPU 多個執行緒上並行執行的程式。
CPU 負責排程並啟動這些核心。您在剖析器追蹤中看到的大部分 PyTorch 開銷，都來自於這種排程工作。
從 matmul-add 到 Linear
nn.Linear 是一個模組包裝器，它封裝了我們在第一部分中已經剖析過的矩陣乘法和加法操作。唯一的區別在於，它將權重（weight）和偏置（bias）作為參數擁有，並公開了 PyTorch 使用者所熟悉的 forward 方法。
# bias=True 將真正模擬我們在系列第一部分中看到的乘法和加法操作
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)
手邊的操作可以寫成：
y = x @ w.T + b
其中 x 是輸入，w 是權重，b 是偏置。讓我們執行 02_linear.py 並檢查剖析結果。
uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b traces
trace-util 是一個實用工具，它會將您的追蹤同步到 Hugging Face 儲存桶，然後在您的終端機上提供 Preffeto URL。
圖 1 顯示了線性層 forward 呼叫的剖析器追蹤。我們以與先前追蹤類似的排程設定（wait=1、warmup=1 和 active=3）追蹤了線性層的 forward 呼叫。這就是為什麼我們在 CPU 和 GPU 軌道上看到三個「Profile Steps」。
轉置操作在做什麼？
如果我們放大剖析器追蹤，如圖 2 所示，我們會注意到在 aten::addmm（乘法和加法）操作之前有一個 aten::t（轉置）操作。我們已經可以判斷 nn.Linear 會先轉置權重參數，然後再與輸入相乘。這就是我們看到 aten::t 操作的原因。
值得注意的是，aten::t 並沒有真正複製或重新組織資料：它只是在 CPU 上重寫張量元資料（形狀和步長）以表示轉置後的矩陣。它不會在 GPU 上啟動核心。您可以透過兩種方式驗證這一點：查看追蹤中的 GPU 軌道，或者檢查剖析器表格中 aten::t 行及其在 CUDA 上花費的時間。
為什麼沒有單獨的乘法和加法核心？
如圖 3 所示，線性層的調度鏈中沒有 aten::add（偏置加法）。這是因為偏置加法已經透過所謂的「後處理 (epilogue)」融合到矩陣乘法核心中。
後處理是 GEMM（通用矩陣乘法）核心在最後，即將結果寫回 HBM（高頻寬記憶體，即 GPU 的主記憶體）之前，所做的一個小型計算。添加偏置、應用激活函數或按常數縮放都是典型的後處理。後處理的目的是避免第二次載入或寫入 HBM，因為記憶體流量會使操作變得昂貴。
nn.Linear 呼叫 torch.nn.functional.linear，後者又呼叫 aten::linear。aten::linear 會檢查輸入，注意到傳入了偏置，然後調度 aten::addmm(bias, x, weight)，而不是單獨執行矩陣乘法和加法。addmm 計算：
out = x @ weight.T + bias
在 GPU 上執行的 cuBLAS GEMM 核心內建了偏置加法變體，這就是 aten::addmm 所選擇的核心。加法從未作為一個單獨的核心出現，因為它是矩陣乘法核心寫回操作的一部分，這正是後處理的意義。
此時值得注意一個微妙之處。您在第一部分 --compile 模式下看到的那個核心 (addmm)，正是 eager 模式的 nn.Linear 已經使用的核心。對於 torch.compile 來說，這裡沒有什麼可以再融合的了，這也是我們接下來要驗證的。
--compile 能幫助單一 Linear 層嗎？
讓我們編譯 forward 呼叫並查看剖析器追蹤。（剖析器追蹤在下一節中視覺化）
uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile
uvx trace-util traces -b traces
如果您比較單一 nn.Linear forward 呼叫的 eager 模式和編譯模式追蹤，您會發現：
GPU 上是相同的 cuBLAS GEMM 核心。
CPU 上是相同的 aten::addmm 操作。
CPU 軌道上有一些編譯模式特有的額外行。
這值得深思。當模型感覺緩慢時，人們通常會反射性地使用 torch.compile。對於單一帶偏置的 GEMM 操作，compile 幾乎無事可做。這不是一個錯誤，只是 compile 需要多於一個操作才能進行可能的融合。讓我們透過查看 MLP 來證明這一點。
轉置操作去哪了？核心佈局與前置操作
仔細閱讀這兩種追蹤（eager 模式與編譯模式）的讀者會注意到，eager 模式的 CPU 調度鏈比編譯模式的更長。
eager 模式下 aten::linear 內部的 CPU 調度鏈是 aten::t 後跟 aten::addmm（圖 4）。要理解 aten::t 實際做了什麼，我們需要快速了解一下步長（strides）和視圖（views）。
張量將其資料儲存為記憶體中一個扁平、連續的數字序列。形狀（shape）和步長是位於該序列之上的元資料，它們告訴 PyTorch 如何遍歷它：步長為 (s0, s1) 意味著「移動一行需要 s0 個元素，移動一列需要 s1 個元素」。改變元資料，您就會得到相同原始資料的不同視圖，而無需複製：
>>> M = torch.tensor([[0, 1],
... [2, 3],
... [4, 5]])
>>> M.shape, M.stride()
(torch.Size([3, 2]), (2, 1)) # 每行兩步，每列一步
>>> T = M.t() # 轉置
>>> T.shape, T.stride()
(torch.Size([2, 3]), (1, 2)) # 形狀和步長交換，資料未動
>>> T
tensor([[0, 2, 4],
[1, 3, 5]])
>>> T.flatten() # 強制實體化，資料被重新排序
tensor([0, 2, 4, 1, 3, 5])
M.t() 沒有移動任何數字。它返回了一個步長交換的新視圖，因此現在逐行讀取它會以轉置的順序遍歷原始緩衝區 0, 1, 2, 3, 4, 5。底層資料是相同的；只有元資料不同。
這正是 aten::t 在線性層內部所做的事情：它不分配新的張量或複製任何資料，它只是產生一個權重張量，其步長被重寫為視圖。
如圖 5 所示，compile 並沒有移除一個 GPU 核心：它移除了調度該視圖的 CPU 開銷。Inductor 在編譯時追蹤了視圖鏈，計算了一次最終的步長，並發出了一個帶有這些硬編碼步長的直接 aten::addmm 呼叫。幾微秒的 CPU 工作消失了，而 GPU 執行著相同的數學運算。
正如預期，當輸入資料違反編譯器預先計算的步長時，它將拋出錯誤。
如果您查看兩種追蹤中的 GPU 軌道，每個 forward 呼叫都只有一個核心，而且兩次都是相同的核心：
cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
如果沒有轉置核心運行，是誰教導 GEMM 以轉置順序讀取權重矩陣的呢？答案就在核心的名稱中。請看後綴：
cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
^^
那個 tn 是佈局描述符。cuBLAS 和 CUTLASS 會為每種輸入佈局組合預編譯一個單獨的核心二進位檔。
n（非轉置）和 t（轉置）描述了核心在內部迴圈期間如何遍歷其輸入。調度器的任務是查看輸入步長，決定哪種後綴組合匹配，然後選擇正確的預編譯核心。
剖析器追蹤中的核心名稱是核心身份的雜湊傾印。如果兩次運行顯示相同的核心名稱，則 GPU 正在執行相同的工作。如果它們不同（例如 _tn_ 與 _nn_、bf16 與 fp16，或 s16816gemm 與 s161616gemm），則 GPU 正在執行不同的工作，並且調度器選擇了不同的分支。學習閱讀這個名稱是比較追蹤時最有用的習慣之一。
堆疊三個 Linear 層：MLP
在本節中，我們將剖析一個多層感知器 (MLP)。為了使其更有趣，我們將剖析一個帶有 GeGLU 激活變體的前饋網路（這在實踐中被廣泛使用）。這也是我們向深度學習研究史上最偉大的句子之一致敬的方式（圖 6）。
class SimpleGeGLUMLP(nn.Module):
def __init__(self, dim, hidden):
super().__init__()
self.gate_proj = nn.Linear(dim, hidden, bias=False)
self.up_proj = nn.Linear(dim, hidden, bias=False)
self.down_proj = nn.Linear(hidden, dim, bias=False)
def forward(self, x):
g = self.gate_proj(x)
u = self.up_proj(x)
h = F.gelu(g, approximate="tanh")
m = h * u
y = self.down_proj(m)
return y
您可以在這裡找到完整的腳本：03_simple_mlp.py。執行方式如下：
uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072
uvx trace-util traces -b traces
在打開追蹤之前，讓我們一起思考應該會看到什麼。forward 函數執行了相當多的計算，但其中大部分我們已經很熟悉。
我們應該預期會看到三個 aten::linear 調度，每個 nn.Linear 層一個。我們還應該預期會啟動兩個逐點（pointwise）核心，一個用於 GeLU，一個用於乘法。在查看之前形成這種預期是剖析過程中最有用的習慣：您閱讀追蹤是為了確認或推翻猜測，而不是從頭開始形成猜測。
從圖 7 我們可以為自己感到驕傲，因為我們的直覺是正確的。每個 forward 傳遞（一個 mlp_fwd）中，GPU 恰好運行 5 個核心。圖 8 突顯了線性投影 CPU 軌道中看到的「佔用率查詢 (occupancy query)」。
下表總結了每個操作的 CPU 操作、GPU 核心和啟動方式。其中，三個 nn.Linear 層（gate_proj、up_proj、down_proj）都對應到 aten::linear CPU 操作，並啟動了 ampere_bf16_s16816gemm 類型的 GPU 核心，伴隨著佔用率查詢。
而 gelu 和 h * u 則分別對應 aten::gelu 和 aten::mul CPU 操作，並啟動了 vectorized_elementwise_kernel 類型的 GPU 核心。
三個 GEMM 各自執行一個額外的 cudaOccupancyMaxActiveBlocksPerMultipr