Kernel 函式
3 秒看懂
GPU Kernel 是執行在圖形處理器(GPU)上的最小編譯單元,在深度學習中它把矩陣乘法、卷積、啟用函式等演算法對映到成千上萬個執行緒並行執行。其本質是一段裝置端程式碼,由主機端(CPU)呼叫,在 GPU 的流多處理器(SM)上以 SIMT 方式排程。理解 Kernel 就是理解大型模型訓練/推論效能的起點:一次 Transformer 的 forward pass 背後是數百個精心最佳化的 Kernel 在 flash attention、矩陣乘、layer normalization 等關鍵路徑上協同工作。
3 分鐘產業解釋
在 “CPU 負責排程、GPU 負責密集計算” 的異構計算範式裡,Kernel 是 GPU 能夠理解並執行的程式單元。深度學習架構(PyTorch、TensorFlow)並不直接產生 GPU 指令,而是通過具備高度工程化特性的底層庫——如 NVIDIA 的 cuBLAS、cuDNN、CUTLASS,或者 OpenAI 自研的 Triton 語言——將高層運算元自動或手動分解為一系列 Kernel 呼叫。
一個典型的 GPU Kernel 在邏輯上描述一個 網格-塊-執行緒 的三級計算組織。比如一個 [M, K] × [K, N] 的矩陣乘法,網格會被劃分為若干執行緒塊,每個執行緒塊計算輸出矩陣的一個子塊;塊內執行緒通過共享記憶體和暫存器協作,對全域性記憶體的資料進行分片和複用,以最大化頻寬利用。
產業上,GPU Kernel 效能直接影響大型模型訓練的 MFU(模型浮點運算利用率)和推論延遲。A100/H100 等高階 GPU 的理論峰值算力很高,但如果沒有像 FlashAttention 那樣通過 Kernel 融合、重計算、IO 感知等最佳化,實際利用率可能只有 20–30%。因此 Kernel 最佳化工程師 已成為各大 AI 實驗室爭奪的關鍵人才,他們需要精通 CUDA C++、PTX/SASS 指令、GPU 記憶體層次,甚至需要針對單一運算元手寫彙編級 Kernel。
15 分鐘專家深入
把對 GPU Kernel 的理解提升到專家級,需要吃透三個維度:執行模型、記憶體模型、以及與深度學習編譯棧的互動。
執行模型:SIMT 與 warp 排程
GPU 採用單指令多執行緒(SIMT)架構。一個 Kernel 啟動時會配置網格和執行緒塊尺寸,硬體上將每個執行緒塊分配到某個 SM。在 NVIDIA 架構中,一個執行緒塊的所有執行緒以 32 個執行緒為一組(一個 warp)進行排程。warp 內的執行緒在同一時刻執行相同指令,處理不同資料。如果一個 warp 遇到分支且執行緒走向不同,硬體會序列執行兩條路徑,掩蔽部分執行緒,導致執行緒束髮散,這是 Kernel 最佳化的頭號大敵。
Kernel 的並行粒度選擇(blockDim 和 gridDim)需要匹配 SM 的可用暫存器檔案和共享記憶體容量。每個 SM 能駐留的執行緒塊和 warp 數量受限於暫存器壓力、共享記憶體用量等。更高的佔用率(occupancy)可以隱藏記憶體延遲,但若 Kernel 本身需要大量暫存器儲存中間變數,過度追求 occupancy 反會因暫存器溢位到 L1 快取或全域性記憶體而嚴重拖慢速度。最佳化工程師會通過 __launch_bounds__ 或編譯選項向編譯器提示最大執行緒數,並利用 Nsight Compute 分析實際瓶頸。
記憶體模型:層次與延遲隱藏
GPU 記憶體體系從片上到片外依次是:暫存器(1 cycle 延遲,但每個執行緒私有)、共享記憶體(SRAM,約 20–40 cycles,塊內執行緒共享)、L1 快取/紋理快取/常量快取、L2 快取(片上,數百 cycles)、全域性記憶體(HBM/GDDR,最高頻寬但延遲可達 600–800 cycles)。Kernel 的關鍵最佳化幾乎全是圍繞這一記憶體層次展開的資料搬運與計算的重疊。
以矩陣乘法 Kernel 為例,經典策略是每個執行緒塊將輸入矩陣的子塊從全域性記憶體通過 coalesced 方式載入到共享記憶體(保證相鄰執行緒訪問連續地址),再利用共享記憶體上的資料塊進行乘累加,結果寫回暫存器,最終將暫存器中的部分和寫出到全域性記憶體。同時,每個執行緒塊會使用雙緩衝(double buffering)技術,在計算當前分片的同時,非同步預取下一個分片的資料(從全域性記憶體到共享記憶體),從而掩蓋記憶體延遲。在較新的架構中,可以直接使用非同步複製指令(如 cp.async)將資料直接從全域性記憶體載入到共享記憶體,而無需浪費計算指令。
FlashAttention 系列 Kernel 的突破正在於通過分塊(tiling)與重計算(recomputation) 將 softmax 的歸約操作完全限定在 SRAM 內,避免頻繁讀寫 HBM 的 attention 矩陣。這種 IO-aware Kernel 設計已成為大型模型推論/訓練的標準範式。
深度學習編譯器與 Kernel 生成
手寫 CUDA Kernel 雖然能榨乾效能,但開發成本極高、跨架構遷移困難。近年來,TVM、Halide、Triton、MLIR 等工具鏈應運而生,希望自動生成高效能 Kernel。OpenAI 的 Triton 語言允許程式設計師用類 Python 語法描述塊級別的資料流,編譯器自動完成執行緒級對映、共享記憶體分配、warp 級最佳化和硬體的細粒度調優。Triton Kernel 在很多場景下能達到與手寫 cuBLAS 相當的效能,極大降低了 Kernel 門檻。
在編譯棧中,Kernel 通常先被編譯成中間表示(如 LLVM IR 或 MLIR 的 GPU 方言),再翻譯為 PTX(並行執行緒執行)偽彙編,最後由 GPU 驅動編譯為機器碼(SASS)。PTX 具有一定的前向相容性,而 SASS 則與具體架構繫結。專家級最佳化時常直接閱讀 SASS 指令,確認是否存在非預期的暫存器溢位、過剩的 load/store 指令或 bank conflict 等。
技術原理
目的:拆解 GPU Kernel 的執行時機制,聚焦對深度學習至關重要的幾個硬核細節。
1. 執行緒層級與索引對映
global_idx = blockIdx * blockDim + threadIdx;
這是一個典型的一維對映。對於二維矩陣計算,通常將 threadIdx 對映到輸出矩陣的行/列:
row = blockIdx.y * blockDim.y + threadIdx.y;
col = blockIdx.x * blockDim.x + threadIdx.x;
Kernel 啟動時,這些內建變數由硬體自動填充。
2. 記憶體請求合併(Coalescing)
全域性記憶體頻寬能否跑滿,關鍵在於 warp 內的所有執行緒在一個指令內訪問的地址是否落在同一個 32 位元組、64 位元組或 128 位元組對齊的段內。當 coalesced 條件滿足,一次記憶體事務即可為整個 warp 服務 32/64/128 位元組;否則會拆分為多個事務,頻寬利用率斷崖下降。深度學習的特徵圖張量通常以 NCHW/NHWC 格式儲存,不同的版面配置會顯著影響 convolve Kernel 的 coalescing 效率,由此催生了各種 layout 重排 Kernel。
3. 共享記憶體 bank conflict
共享記憶體被劃分為多個 bank(NVIDIA 普遍為 32 個 bank),若同一 warp 內的多個執行緒訪問相同 bank 的不同地址(且寬度為 4 位元組),將發生 bank conflict,導致請求序列化。最佳化 Kernel 常用的一種技巧是填充(padding) 共享記憶體陣列的維度,例如將 float sA[32](#src-32)[32](#src-32) 改成 float sA[32](#src-32)[33](#src-33),以打破 bank 衝突的模式。在矩陣乘法 Kernel 中,bank conflict 的消解對效能影響可達 10–30%。
4. 計算與資料移動的重疊(Compute-Memory Overlap)
除雙緩衝外,現代 GPU 支援非同步記憶體複製引擎。Kernel 可以手動插入 __pipeline_memcpy_async 等指令,將下一批資料從全域性記憶體載入到共享記憶體,而當前 warp 立即進入計算階段,不等待複製完成。當計算用到了這批資料時,通過 __pipeline_wait_prior 等同步原語確保資料已到達。這是實現接近 roofline 上限效能的核心手法。
5. 算術強度與 roofline 模型
Kernel 的算術強度(FLOP/Byte)決定了其效能瓶頸在算力還是頻寬。例如逐元素啟用函式 Kernel 的算術強度極低(每個元素數條指令,但僅讀寫一次),幾乎永遠頻寬受限。而大矩陣乘法的算術強度可達數百,是算力受限。通過 Kernel 融合(將 element-wise 操作合併到一個 Kernel 中,減少中間結果的全域性記憶體讀寫),可以顯著提升低強度運算元的有效頻寬利用,這是 XLA、TorchInductor 等編譯器大力推行的最佳化。
技術演進史
- 2006 年之前:GPGPU 萌芽,開發者通過圖形 API(OpenGL/DirectX)偽裝為渲染操作進行計算,Kernel 概念為 shader 程式。
- 2006 年:NVIDIA 釋出 CUDA 1.0,提出 grid/block/thread 執行模型和 CUDA C 語言,Kernel 函式以
__global__宣告,大規模通用平行計算成為可能。 - 2010–2012 年:Fermi 架構引入可快取全域性記憶體(L1/L2),共享記憶體與 L1 可配置。cuBLAS、cuFFT 等庫將常用 Kernel 模板化,深度學習尚在萌芽。
- 2014–2016 年:Maxwell/Pascal 架構大幅提升能效。cuDNN 針對卷積、池化等提供預最佳化 Kernel。Kernel 自動調優(auto-tuning)出現,如 CUDA 的
cudaOccupancyMaxPotentialBlockSize等工具。 - 2017–2018 年:Volta 架構引入 Tensor Core,支援混合精度矩陣乘累加。需要專門編寫 Warp-level Matrix Multiply-Accumulate (WMMA) 指令或使用 nvcuda::wmma 的 C++ 封裝來呼叫 Tensor Core Kernel。深度學習進入 Transformer 時代,注意力矩陣的記憶體瓶頸凸顯。
- 2019–2021 年:Turing/Ampere 增強了 Tensor Core 的稀疏化和細粒度結構化稀疏支援。Triton 語言逐漸成熟,簡化 Kernel 開發。
- 2022 年至今:Hopper 架構引入 TMA(Tensor Memory Accelerator)和 FP8 支援,Kernel 編寫更依賴非同步複製和新的程式設計模型。FlashAttention (2022) 通過分塊+重計算 Kernel 震撼業界,說明極致 IO-aware Kernel 能帶來數量級的提速和節省記憶體。CUTLASS 3.x 將 Kernel 抽象為可組合的模板,加速了針對新硬體的手工最佳化。同時,PyTorch 2.0 的 TorchInductor 用 Triton 編譯器後端自動生成 Kernel,推高了自動化水平。
技術路線對比
| 維度 | 手寫 CUDA Kernel | cuBLAS/cuDNN 庫 | Triton 語言 | 編譯器自動生成 (TorchInductor/TVM) |
|---|---|---|---|---|
| 效能潛力 | 最高(可榨乾硬體特性) | 高(廠商精細調優) | 很高(接近手寫) | 中–高(視模板與調優程度) |
| 開發效率 | 低(需要深入 GPU 微架構) | 高(呼叫 API) | 中(塊級別程式設計) | 高(Python 前端,自動最佳化) |
| 可移植性 | 差(繫結 CUDA 生態) | 中(同庫跨代,不同庫需適配) | 較高(支援 NVIDIA/AMD 等) | 高(IR 可接多後端) |
| 典型應用場景 | 新一代注意力機制、奇特啟用函式 | 標準卷積、GEMM、歸一化 | 研究型運算元、快速迭代的定製融合 | 圖級別最佳化、常見運算元組合 |
| 學習曲線 | 陡峭 | 平緩 | 中等 | 平緩 |
| 硬體利用率上限 | 極近 Roofline | 接近 Roofline | 接近 Roofline | 常為 60–80% 廠商峰值 |
注:資料為定性評估,無具體對標基準。
上下游
- 上游(硬體與驅動):GPU 架構設計(NVIDIA、AMD、Intel)、CUDA 驅動、PTX/SASS 編譯器。硬體團隊定義的 ISA 和記憶體層次直接決定 Kernel 能用到什麼指令和頻寬上限。
- 中游(底層庫與編譯器):cuBLAS、cuDNN、CUTLASS、Thrust、cuSPARSE、Triton、MLIR 的 GPU 方言。它們是封裝 Kernel 的中介軟體,將高層運算元對映到最佳化後的 Kernel 集合。
- 下游(架構與模型):PyTorch、TensorFlow、JAX、PaddlePaddle 等架構,以及 GPT、Llama、Stable Diffusion 等模型。架構的排程器通過運算元分發機制,將使用者的
torch.matmul呼叫路由到最佳的底層 Kernel 實現。 - 下游的末端:大型模型訓練/推論服務,比如 vLLM 架構中通過 PagedAttention Kernel 管理 KV 快取,顯著提升推論吞吐。這些自定義 Kernel 直接面向業務指標調優。
關鍵指標
評估 GPU Kernel 效能的常用指標(未注出處為通用度量):
- 執行時間/延遲:單次 Kernel 呼叫的耗時,通常用納秒或微秒。
- 吞吐量(TFLOPS):實際達到的浮點運算次數,與理論峰值的比值即 MFU。
- 頻寬利用率(%):實際資料傳輸速率與裝置頻寬的比例。
- 佔用率(Occupancy, %):每個 SM 上活躍 warp 與理論最大 warp 數的比值,反應延遲隱藏能力,但不是越高越好。
- 暫存器用量/共享記憶體用量:每個執行緒所用暫存器數量、每執行緒塊所用共享記憶體量,直接影響 occupancy 和溢位風險。
- 算術強度(FLOP/Byte):決定 Kernel 為頻寬受限還是算力受限。
- Bank conflict 率:可通過分析工具獲得,影響共享記憶體頻寬。
供需與市場資料
由於檢索未返回具體資料,以下為基於公開資訊的定性描述:
- 人才需求:AI 基礎設施領域對 GPU Kernel 工程師的需求三年內呈指數增長,尤其在大型模型公司、雲端廠商和 AI 晶片初創企業。職位涵蓋 CUDA 開發、高效能運算(HPC)以及 ML 編譯器工程師。據行業觀察,資深 GPU Kernel 最佳化專家的薪酬位於 AI 人才金字塔頂端。
- 工具鏈市場:NVIDIA 的 CUDA 生態仍佔據絕對主導,其編譯器與庫是事實標準。Triton 在頂尖 AI 實驗室滲透率快速上升,並有部分企業將其用於自研加速硬體。開源編譯器(如 MLIR)的進展可能降低市場集中度。
- 硬體供給:NVIDIA 資料中心 GPU 供應緊張長期持續,每代新架構的釋出都會引發大量 Kernel 重寫和適配需求,從而拉動對 Kernel 開發的投資。
- 行業趨勢:隨著模型規模的增速超越單晶片視訊記憶體增長,手工融合 Kernel 從可選項變為必選項。從 FlashAttention 到各類自定義融合運算元的需求激增,Kernel 已成為大型模型公司的核心技術壁壘之一。
(具體人才規模、市佔率數字未充分揭露)
代表公司與資本對映
- NVIDIA:GPU Kernel 生態建置者,CUDA、cuBLAS、cuDNN 等庫的提供者,持續通過新指令集(Tensor Core, TMA)定義 Kernel 最佳化的硬體基礎。
- AMD:ROCm 生態中的 HIP 可轉換 CUDA 程式碼為運行於 AMD GPU 的 Kernel,正加大在 AI 架構相容性上的投入。
- Intel:通過 oneAPI 的 DPC++ 和 SYCL 瞄準 Kernel 可移植性,其 GPU 產品在資料中心嶄露頭角。
- AI 晶片初創公司(如 Groq、Tenstorrent、Graphcore 等):很多采用專用架構或資料流架構,其 Kernel 概念與傳統 GPU 差異顯著,需要專門的軟體棧,其成功很大程度上取決於編譯器和 Kernel 庫的完備度。
- 核心實驗室:OpenAI(Triton 語言與 FlashAttention 的原創者)、Meta(PyTorch 及 TorchInductor)、Google(JAX/XLA)、Microsoft(DeepSpeed 中的自定義 Kernel)。這些組織掌握大量最佳化過的 Kernel IP,並以此作為技術護城河。
- 資本關注點:GPU Kernel 能力強的公司往往在訓練成本、推論吞吐上有碾壓性優勢,因而獲得更高估值。VC 在審視 AI Infra 創業專案時,會重點考察團隊是否具備自研高效能 Kernel 的能力。
投資邏輯
- 算力利用率是成本核心:GPU 硬體成本巨大,訓練一個大型模型需要成千上萬塊 GPU。Kernel 最佳化能將 MFU 從 30% 提升至 60% 甚至更高,相當於同等硬體下節省一半的訓練時間和電費,具備直接的投資回報邏輯。
- 鎖定生態護城河:定製 Kernel 一旦與模型結構深度繫結,遷移到其他架構或硬體的成本極高。這使某些大型模型公司的訓練/推論系統具有排他性優勢,有利於其 SaaS/API 業務的毛利率。
- 下游推論降本:推論場景對延遲極度敏感,手寫 Kernel 在 Batch Size=1 的低吞吐場景下效果遠超市面通用庫。能夠提供極致端到端延遲的公司將在企業級 AI 應用中佔得先機。
- 警惕硬體依賴性:當前大部分 Kernel 最佳化系針對 NVIDIA GPU 設計。若未來 AI 硬體格局走向多元化,過度依賴單一平台的手工 Kernel 可能成為遷移負擔。編譯器自動生成技術路徑的價值因此凸顯,相關工具鏈(如 Modular/MLIR)值得關注。
- 人才密集型企業價值:頂級 Kernel 工程團隊稀缺,擁有這類團隊的公司具備長期的技術迭代能力,應視為核心資產。
常見誤讀糾偏
誤讀 1:“Kernel 就是運算元(operator),兩者可以互換。”
糾正:在深度學習語境下,運算元(如卷積、softmax)是數學定義,而 Kernel 是特定硬體上實現該運算元的可執行程式碼。一個運算元可能對應多個 Kernel(不同輸入形狀、資料型別、最佳化策略),執行時由架構的分發器動態選擇。兩者不是同一抽象層級。
誤讀 2:“GPU 的執行緒是獨立並行的,越多越快。”
糾正:GPU 執行緒的執行不是完全獨立的,它們以 warp 為單位鎖步執行。過多的執行緒塊可能導致每個執行緒的暫存器資源減少,引發暫存器溢位到慢速記憶體,反而拖慢效能。而且,執行緒塊尺寸過小可能造成 SM 內控制開銷增大,無法充分利用硬體。合理的 Kernel 設計需要精確權衡並行粒度與資源佔用。
學習路徑
基礎階段
- 掌握 C/C++ 和計算機體系結構基礎(快取、記憶體層次)。
- 學習《CUDA C++ Programming Guide》,理解 grid/block/thread、共享記憶體、同步原語。
- 在 Nsight Systems 中執行官方 matrix multiplication 示例,觀察 Kernel 時間線。
進階階段
- 閱讀 CUTLASS 原始碼,理解 tiled GEMM 的軟體流水線和 Tensor Core 對映。
- 完成 Triton 官方教程,手寫一個融合的 LayerNorm+Dropout Kernel並比較與 PyTorch 的效能。
- 使用 Nsight Compute 分析 Kernel 的 roofline、bank conflict 和 occupancy。
高階階段
- 攻讀 NVIDIA PTX ISA 參考手冊,逐步從 PTX 角度理解編譯器的最佳化決策。
- 深入研究 FlashAttention-2 的 CUDA 實現,體會如何通過分塊避免 HBM 讀寫全域性 attention 矩陣。
- 結合 MLIR 的 GPU 方言,嘗試從高層 IR 生成 Kernel。
- 參與開源專案(如 vLLM、FastTransformer)貢獻定製 Kernel,獲取實戰經驗。
推薦書籍/資料:
- Programming Massively Parallel Processors (Kirk & Hwu)
- NVIDIA GPU 技術大會(GTC)關於 CUDA 的演講錄影
- CUDA Toolkit Documentation
- 論文 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
一句話總結
GPU Kernel 是將深度學習的數學描述轉化為 GPU 並行機器的最後“一公里”,是大型模型效能和成本的關鍵控制點,理解它等於掌握 AI 基礎設施的效能命脈。
延伸閱讀與來源
- CUDA C++ Programming Guide, NVIDIA.
- CUTLASS: CUDA Templates for Linear Algebra Subroutines and Solvers, NVIDIA, github.com/NVIDIA/cutlass.
- Dao, T., et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022.
- Tillet, P., et al., Triton: an intermediate language and compiler for tiled neural network computations, MAPS 2019.
- Lattner, C., et al., MLIR: Scaling Compiler Infrastructure for Domain Specific Computation, IEEE 2021.
- 相關公司官網與公開發布的白皮書。
注:因檢索未返回資料,所有具體規格和數字均為定性描述或基於公開領域知識,未作硬量宣告。