隨著加速計算不斷提升 AI 和科學計算各個領域的應用程序性能,人們對 GPU 優化技術的興趣也越來越濃厚,以確保應用程序獲得盡可能好的性能。作為應用程序開發者,有很多方法可以對軟件堆棧上下進行 GPU 編程。在本文中,我們將介紹堆棧的一些不同級別,并深入探討最低級別:手寫并行線程執行 (PTX) 代碼。
加速計算軟件堆棧
現在,您無需編寫特定于 GPU 的代碼,即可使用 GPU 完成大量工作。庫開發者和軟件工程師已經為您完成了底層工作。例如,您可以在堆棧中進行高級工作,使用藍圖構建完整的 AI 工作流。或者,您可以在 PyTorch 等框架中開發應用程序,您可以在其中指定模型、適當的 GPU 代碼和庫來自動執行程序。
您還可以使用全套 NVIDIA CUDA-X 庫開發應用程序,其中包括量子計算、數據處理、物理 AI、基因測序、邊緣計算、藥物研發等領域特定的庫。如果這些特定于域的庫不包含您所需的所有功能,您可以使用 OpenACC 等編譯器指令對 GPU 進行編程,也可以使用 libcu++ 對 C++ stdpar 算法等庫和 C++ 標準庫進行編程。
在上述所有情況下,您不是在編寫特定于 GPU 的代碼,而是依賴于由專家工程師精心設計、實施和優化的庫或編譯器指令。
但是,在某些情況下,您可能必須實現自己的 GPU 代碼,因為當前不存在用于滿足您所需功能的庫。然后,您可以進一步向下移動堆棧,并直接使用高級語言 (例如 C++、Fortran 和 Python) 編寫 CUDA GPU 代碼。
最后,在極少數情況下,開發者可能會選擇更深入地使用 PTX 直接編寫代碼中對性能極為敏感的部分。與大多數性能優化技術一樣,您期望的控制越多,在堆棧中提取性能所需的越低。應謹慎考慮這一權衡:除了增加開發和調試的復雜性之外,手寫低級代碼帶來的性能提升可能無法移植到其他 GPU 架構。
正如我們在之前的博文中所展示的,PTX 是 GPU 的匯編語言。直接編寫 PTX 是一種非常先進的優化技術,對于大多數開發者來說并不是必需的,因此應將其視為萬不得已的工具。不過,在某些情況下,通過編寫 PTX 實現的精細控制可直接提高特定應用程序的性能。這些情況通常發生在應用程序對性能非常敏感的部分,其中性能提升的每一部分都具有顯著的優勢。所有可用的 PTX 指令都在 PTX ISA 文檔中。
在這篇博客文章中,我們將深入探討一個示例,其中使用手寫 PTX 來提高某些 AI 模型實現中出現的重要算法的性能。
編寫 PTX
在進入示例之前,我們將列出一些在應用中包含手寫 PTX 代碼的方法。換句話說,原則上是如何做到這一點。以下示例展示了真實場景并顯示了性能變化。
內聯 PTX
在代碼中包含 PTX 的一種標準方法是使用內聯 PTX。這是我們將在下面展示的方法,有關語法和語義的詳細信息請參閱文檔。這與在 CPU 上編寫組裝代碼非常相似。
cuda::ptx 命名空間
在代碼中包含 PTX 的另一個選項是使用 libcu++,其中包含命名空間 cuda::ptx,可提供直接映射到 PTX 指令的函數。這有助于在 C++ 應用程序中輕松使用特定的 PTX 指令。有關 cuda::ptx 命名空間的更多信息,請參閱 cuda::ptx 命名空間文檔。
CUTLASS 示例
為了說明如何手動編寫 PTX 代碼,我們將使用線性代數的特定示例。一般來說,如果您的運算可以表示為線性代數運算 (例如 GEMM) ,則推薦使用 NVIDIA CUBLAS 在 GPU 上運行。CUBLAS 已經針對許多矩陣的大小和形狀進行了高度優化,并且具有多個數值精度可供選擇。
有時,CUBLAS 中的功能無法完全表達您想要執行的操作,或者您希望在 GEMM 之前或之后直接執行計算。有時,您可以通過將其他運算與 GEMM 運算融合來提高性能,而不是先調用一些函數,然后調用 CUBLAS 和更多函數。這有很多好處,因為融合內核可能會實現更多優化,例如更高效地使用數據。
這正是 NVIDIA CUTLASS 庫的用武之地。CUTLASS 包含一系列 CUDA C++ 模板抽象,用于在 CUDA 內的各個級別和規模上實現高性能矩陣乘法 (GEMM) 和相關計算。由于 CUTLASS 支持對 GEMM 和類似 GEMM 的操作進行更多的控制和自定義,因此與 CUBLAS 相比,CUTLASS 涉及的開發者代碼略多一些。
CUTLASS 包含大量手寫 PTX,因為它在每個 GPU 架構上都以最佳性能進行設計。這使得 CUTLASS 成為說明手寫 PTX 實際應用的現成示例。
GEMM 以及 top_k 和 softmax
我們將演示的特定運算是 GEMM 與 top_k 和 softmax 算法的融合。這是運行混合專家神經網絡時的常見操作。我們將重點介紹 NVIDIA Hopper 架構。由于這是一種常用的操作,因此 CUTLASS 已經為此提供了一個帶有一些內聯 PTX 的特殊內核,因此可以直接演示 CUTLASS 如何將手寫 PTX 融入其高性能 GPU 代碼。
在這篇博文中,我們使用:
- 版本CUTLASS 的 3.9.2
- NVIDIA GH200 GPU
- 驅動程序版本 570.140
- CUDA 工具包版本 12.8
按照 CUTLASS 網站上的構建說明,我們使用構建選項 -DCUTLASS_NVCC_ARCHS = 90a 進行 cmake,以確保啟用 Hopper 架構的完整功能集。CUTLASS 資源庫中有許多示例展示了最新架構上的各種功能。完成 cmake 后,我們導航至構建目錄 (例如,build/examples/ 61_hopper_gemm_with_topk_and_softmax) ,以便構建和運行示例代碼。
我們執行 make 來構建代碼,可執行文件已構建完畢并可隨時運行。該應用程序接受一些不同的選項作為輸入,包括矩陣大小 m、n 和 k、容錯 epsilon,以及為生成 GFlop/s 基準測試數據而運行的迭代次數。
選擇 m = 1024、n = 8 (默認) 、k = 4096、迭代 = 1000000 以及 epsilon 為 1e-4,即可獲得以下輸出。在此基準測試中,類似于 LLM 執行,m 是 token 的數量,n 是專家的數量,k 是專家的嵌入維度,top_k 的值為 2 (在測試代碼中進行硬編碼) 。
$ ./61_hopper_gemm_with_topk_and_softmax --m=1024 --k=4096 --iterations=1000000 --eps=1e-4
Disposition: Passed Relative error: 1.52478e-05
Problem Size: 1024x8x4096x1
Avg runtime: 0.011765 ms
GFLOPS: 5704.11
在此基準測試示例中,性能為 5704 GFlop/s。我們將令牌數量 ( m 個參數) 最多改為 16384,并生成以下性能表。
米 | GFlop/ s |
1024 | 5704 |
2048 | 9551 |
4096 | 14569 |
8192 | 19794 |
16384 | 21476 |
表 1。基準測試代碼的性能,包括將內聯 PTX 用于 top_k 和 softmax 函數
刪除內聯 PTX
此基準測試示例將 GEMM 與 top_k 和 softmax 融合在一起,并調用在 top_ k 函數中使用內聯 PTX 的函數,前提是 k 的值為 2 或 4。(注意:`k` 不同于上述矩陣維度 k。)
在某些條件下,它還使用 inline PTX 作為 softmax 函數。此外,當不符合特定條件時,top_k 和 softmax 都有使用 CUDA C++ 編寫的后備例程。您可以直接更改 top_k 和 softmax 的內部函數,以注釋掉 PTX 函數的調用,并運行備用 CUDA C++ 代碼。在本例中,這將使我們能夠量化手寫 PTX 的價值。
為從此示例中刪除內聯 PTX,我們編輯了 cutlass/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp 文件,以注釋掉內聯 PTX 函數的使用。在此文件的頂部附近,您會看到一些使用內聯 PTX 編寫的函數,這些函數以“top_2”和“top_4”開頭。例如,以下是您將遇到的第一個 PTX 功能。
CUTLASS_DEVICE
Array<float, 2> top_2_reduce_scalar(Array<float, 2> a, float scalar) {
Array<float, 2> out;
asm volatile(
"{\n"
" .reg .f32 mx;\n"
" .reg .pred p;\n"
" max.f32 mx, %3, %4;\n"
" setp.gtu.f32 p, %2, %4;\n"
" selp.f32 %1, mx, %2, p;\n"
" selp.f32 %0, %2, %4, p;\n"
"}\n" : "=f"(out[0]), "=f"(out[1]) : "f"(a[0]), "f"(a[1]), "f"(scalar));
return out;
}
無需了解此代碼的每個細節。主要目的是展示一個簡短的內聯 PTX 功能的外觀示例。
在源代碼下方,您還會看到 softmax 函數。這些是我們將省略的 PTX 函數,以查看性能變化。
在同一源文件中,您可以找到作為 if 語句的一部分調用這些函數的位置。我們只需注釋掉調用內聯函數的 if 語句,然后保留語句的 else 部分。這將省略對內聯 PTX 函數的調用,轉而執行使用 CUDA C++ 編寫的代碼。
例如,有一個名為 add_element_to_desc_sorted_array
的函數,如果 k = 2 或 k = 4,則會分別調用 top_2 或 top_4 PTX 函數,或者調用該算法的 C++ 實現。此函數的代碼如下:
void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) { if constexpr (N == 2 && is_same_v<Element, float >) { a = top_2_reduce_scalar(a, b); } else if constexpr (N == 4 && is_same_v<Element, float >) { a = top_4_reduce_scalar(a, b); } else { // slower generic path with branching, slower, and can cause register spill CUTLASS_PRAGMA_UNROLL for ( int k = 0; k < N; ++k) { if (a[k] < b) { // Shift down CUTLASS_PRAGMA_UNROLL for ( int l = N - 1; l > k; --l) { a[l] = a[l-1]; } a[k] = b; break ; } } } } |
為確定手寫 PTX 函數的效果,我們將對這些 PTX 函數的調用注釋掉,并僅允許代碼執行 C++ 版本的算法,如下所示:
void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) { /* BEGIN COMMENT if constexpr (N == 2 && is_same_v<Element, float>) { a = top_2_reduce_scalar(a, b); } else if constexpr (N == 4 && is_same_v<Element, float>) { a = top_4_reduce_scalar(a, b); } else { END COMMENT */ // slower generic path with branching, slower, and can cause register spill CUTLASS_PRAGMA_UNROLL for ( int k = 0; k < N; ++k) { if (a[k] < b) { // Shift down CUTLASS_PRAGMA_UNROLL for ( int l = N - 1; l > k; --l) { a[l] = a[l-1]; } a[k] = b; break ; } } //} COMMENT THE END OF THE ELSE } |
我們對函數 merge_desc_sorted_arrays
和 masked_softmax
進行類似更改,以刪除 if/else
語句,從而從本示例中刪除手寫 PTX 函數 $ ./61_hopper_gemm_with_topk_and_softmax --m=1024 --k=4096 --iterations=1000000 --eps=1e-4
Disposition: Passed Relative error: 1.52478e-05
Problem Size: 1024x8x4096x1
Avg runtime: 0.011765 ms
GFLOPS: 5704.11
0、tg_ 11、tg_ 12、tg_ 13 和 tg_ 14。
以下是性能結果。
米 | GFlop/ s |
1024 | 4998 |
2048 | 8376 |
4096 | 13267 |
8192 | 17885 |
16384 | 20066 |
表 2。對于 top_k
和 tg_ 16 函數,基準測試代碼在無內聯 PTX 的情況下的性能
將這些結果與表 1 中的結果進行比較后,您會發現當使用手寫 PTX 而非 CUDA C++ 代碼時,性能將從 7% 提高到 14%。帶回家所傳達的信息并不是此處具體顯示的絕對性能提升,而是在某些精心選擇的情況下,手寫 PTX 可以帶來性能提升。應仔細分析性能和可移植性權衡,以確定在您的應用中包含手寫 PTX 的可行性。
這是一個經過高度優化的示例代碼,我們之所以選擇它,是因為它具有由 NVIDIA CUTLASS 工程師編寫的手寫 PTX,顯示出顯著的性能提升。
此示例強化了以下指導:在絕大多數情況下,開發者應將 PTX 的手寫工作交給 CUTLASS、CUBLAS 和其他 GPU 庫的開發者,并在這些庫的基礎上進行構建。
總結
在本文中,我們展示了一個示例,展示了 CUTLASS 如何使用手寫 PTX 來提高某些 AI 模型中使用的特定融合 GEMM 操作的性能。我們不想讓人覺得每個開發者都應該編寫 PTX。絕大多數開發者不需要這樣做。手動編寫 PTX 應該是萬不得已的工具。
盡管如此,手寫 PTX 是一項可供所有開發者使用的技術。這是一種先進的專業技術,如果使用得當,可以成為高級 GPU 程序員工具箱中的另一種工具。
這是 CUDA 平臺的一大優勢,開發者可以在任何適合自己的級別使用 NVIDIA 堆棧,從應用程序級別一直到編寫組合代碼 (PTX) ,以及介于兩者之間的任何級別。
致謝
感謝以下 NVIDIA 貢獻者:Ali Hassani
?