在生成式 AI 時代，充分發揮 GPU 的潛力對于訓練更好的模型和大規模服務用戶至關重要。通常，這些模型的層由于細微的修改而無法表示為現成的庫運算，而 DL 編譯器通常會放棄最后幾個百分點的優化，以實現其部署。
為了向 NVIDIA CUDA 開發者提供更大限度地提高 DL 和 HPC 內核性能所需的功率和控制水平，自 2017 年以來，我們一直在 CUTLASS 上構建和迭代。
借助新的 Python 接口，它現已進入下一階段的開發階段。重新設計 CUTLASS 3.x 時引入的基本抽象概念在 Python 和 CUTLASS 4.0 中直接公開。在本文中，我們討論了 CUTLASS 3.x 背后的設計原則、其核心后端庫、CUDA 張量和空間微核 (CuTe) ，以及利用 CuTe 關鍵功能的優化示例。
來自 CUTLASS 3.x 的亮點
CUTLASS 3 引入了 CuTe，這是一個基于布局概念的新庫，作為描述和操作線程和數據的統一且可組合的抽象。通過將布局提升為編程模型的一級公民，CuTe 的使用大大簡化了線程數據的組織。CuTe 以易于理解和靜態可檢查的方式向開發者展示索引邏輯，同時保持與 CUTLASS 2.x 中相同的高性能和 Tensor Core 運算覆蓋率。
除了這種更有意義的布局方法之外，CUTLASS 3 的目標與所有之前版本的 CUTLASS 相同，即通過圍繞最新硬件功能開發直觀的編程模型，幫助 CUDA 開發者編寫高性能 GPU 線性代數核函數。在這一新的主要迭代中，我們強調了以下幾點：
- 能夠在庫設計中自定義任何層，同時保持與其他層的可組合性，從而提高開發者的工作效率并更清晰地分離運動部件
- 編譯時檢查，以確保內核結構的正確性。這可以保證，如果編譯，它將正確運行，否則將使用可操作的靜態 assert 消息。
- 通過更少的命名類型和更平滑的學習曲線，以及自定義 Hook 的單入口點來減少 API 表面積。
- NVIDIA Hopper H100 和 NVIDIA Blackwell B200 性能出色，可使用 WGMMA (適用于 Hopper) 或 UMMA (適用于 Blackwell) 、Tensor Memory Accelerator for Hopper (TMA) 和線程塊集群等功能。
CuTe
CUTLASS 3.x 的核心是 CuTe，這是一個用于描述和操作線程和數據張量的新庫。CuTe 由兩部分組成：強大的布局表示和作用于這些布局的運算代數。
CuTe 的布局表示采用原生分層式，自然支持靜態和動態信息，并用于表示多維張量。相同的布局表示用于描述數據張量和線程張量。在多個獨立資源中使用相同的詞匯類型顯示了 CuTe 布局概念的廣泛適用性。
基于這種表征能力，CuTe 提供了形式化的布局代數，使用戶能夠根據簡單的已知布局構建復雜的布局，或將一個布局分割到另一個布局。這使得程序員可以專注于其算法的邏輯描述，而 CuTe 可以為其進行機械記帳。借助這些工具，用戶可以快速設計、實施和修改密集線性代數算法。
與之前的任何 GPU 編程模型不同，線程和數據張量的功能組合消除了 GPU 編程中最復雜的障礙之一，即始終將大量線程映射到其所運行的數據。一旦描述了獨立于將要操作的數據布局的線程布局，CuTe 的布局代數就可以跨線程對數據進行分區，而無需手動實施復雜的后分區迭代方案。
CuTe 布局和張量
有關布局和張量的更多 CuTe 文檔可在其專用文檔目錄中找到。
CuTe 提供 Layout 和 Tensor objects，可緊湊地封裝數據的類型、形狀、內存空間和布局，同時為用戶執行復雜的索引。
Layout<Shape,Stride>提供Shape內邏輯坐標與使用Stride計算出的索引之間的映射。(請參見圖 1 示例)Shape定義了一個或多個坐標空間以及它們之間的映射。Stride定義了將坐標轉換為索引的索引圖。- T
ensor<Engine,Layout>通過迭代器提供Layout的合成。迭代器可以是指向全局內存、共享內存、寄存器內存或任何其他提供隨機訪問偏移和解引用的數據的指針。
Shape 和 tg_ 11 函數可操作多種矩陣類型，以創建索引值得強調的是，CuTe 中的布局是分層的，并受張量代數中折疊張量運算的啟發。如圖所示，分層形狀和步長可實現遠超簡單行主和列主的布局表示。同時，分層布局仍然可以像正常張量一樣訪問 (例如，所示的邏輯 2D 坐標) ，因此這些更高級的數據布局在算法開發中被抽象化。
CUTLASS 3.x 中的 CuTe
CUTLASS 3.x 使用單一詞匯表類型 (cute::Layout) ，從而實現簡化、形式化和統一的布局表示，幫助用戶輕松編寫超快的內核。
用于轉換和分區的 CuTe 布局
CuTe 布局支持將功能合成作為核心運算。功能合成可用于轉換另一個布局的形狀和順序。如果我們有一個具有坐標 (m,n) 的數據布局，而我們想改用坐標 (tg_ 15) ，則我們會使用描述映射的布局來編寫數據布局 (tg_ 16) – > (tg_ 17) 。
結果是具有坐標 (thread_idx,value_idx) 的數據布局，我們可以使用該布局輕松訪問每個線程的每個值！
例如，考慮 4 × 8 的數據布局。此外，假設我們要為 4 × 8 數據的每個坐標分配線程和值。我們編寫一個“電視布局”，記錄特定的分區模式，然后在數據布局和電視布局之間執行功能合成。
如圖所示，合成會對數據進行排列和重塑，以便在結果的每一行中排列每個線程的值。只需使用線程索引對結果進行切片，即可完成分區。
更直觀的分區模式視圖是電視布局的反向視圖。
此布局顯示了從 4 × 8 數據布局中的每個坐標到線程和值的映射。可以記錄任意分區模式，并將其應用于任意數據布局。
有關 CuTe 布局代數的其他文檔可在 GitHub 上找到。
CuTe 矩陣乘積累加原子
原子是必須協同參與執行硬件加速數學運算或復制運算的最小線程和數據集合。
Atom 將 PTX 指令與有關線程形狀和安排的元數據以及必須參與該指令的值相結合。此元數據表示為 CuTe TV 布局，然后可用于劃分輸入和輸出數據的任意張量。用戶通常不必擴展此層，因為我們將為新架構提供 CuTe 原子的實現。
上圖顯示了 SM70_8x8x4_F32F16F16F32_NT 指令及其關聯的 MMA_Traits 元數據。在左側，映射 (thread_id,value_id) -> coord 的 TV 布局記錄在特征中，在右側，通過 tg_ 24 映射可視化特征。可以使用以下命令生成右側圖像
print_latex(make_tiled_mma(cute::SM70_8x8x4_F32F16F16F32_NT{}))
有關矩陣乘積累加 (MMA) 原子的其他 CuTe 文檔位于 GitHub 上。
CuTe 平鋪 MMA
平鋪式 MMA 和平鋪式文案分別是 MMA 原子和拷貝原子的平鋪式。我們將此級別稱為“平鋪”，因為它在原子之上構建更大的運算，就像將單個圖塊組合在一起以構建可重復使用的馬賽克組件一樣。這些平鋪在線程和數據之間重現原子，并且可能存在原子的排列和交錯。
此層最類似于 CUTLASS 2.x 中 MMA 指令的線程束級平鋪；但是，它會從參與操作的所有線程的角度查看平鋪，并將此概念推廣到復制操作。此層的目的是基于大量硬件加速的數學和數據移動操作構建可組合的 GPU 微內核，每個操作都可能具有自己的線程和數據內部布局。平鋪的 MMA 和平鋪的 Copy 類型通過一個統一的 API 來劃分數據，從而呈現所有這些硬件加速的 CuTe 原子。
例如，CuTe 可能會提供一個 MMA 原子，用戶可以針對固定的 M、N 和 K 維度在單個線程束中調用該原子。然后，我們可以使用 CuTe 運算 make_tiled_mma 將此原子轉換為適用于整個線程塊的運算，以處理更大的 M、N 和 K 維度。在上一節中，我們已經看到了 Tiled MMA 的一個示例，即 SM70_8x8x4_F32F16F16F32_NT 的 1x1x1 平鋪。
下圖顯示了另外兩個使用相同 SM70_8x8x4_F32F16F16F32_NT 原子的平鋪 MMA。在左側，其中四個原子組合成 2 × 2 行的主要布局，以生成 16x16x4 的單經 MMA。在右側，其中四個原子是 2 × 2 行的主要布局，以產生 16x16x4 的單曲面 MMA，然后行 (M) 和列 (N) 被排列以交錯這些原子。這兩種方法都會產生可應用于任何數據布局的分區模式，如下節所示。
CuTe GEMM 和主回路
借助與架構無關的平鋪 API，用戶可以構建通往 GEMM 外部循環的一致接口，其中包含來自原子層的內部循環。
Tensor gA = . . . // Tile of 64x16 gmem for A
Tensor gB = . . . // Tile of 96x16 gmem for B
Tensor gC = . . . // Tile of 64x96 gmem for C
// 64x16 static-layout padded row-major smem for A
Tensor sA = make_tensor(make_smem_ptr<TA>(smemAptr),
Layout<Shape < _64,_16>,
Stride<Int<17>, _1>>{});
// 96x16 static-layout interleaved col-major smem for B
Tensor sB = make_tensor(make_smem_ptr<TB>(smemBptr),
Layout<Shape <Shape <_32, _3>,_16>,
Stride<Stride< _1,_512>,_32>>{});
// Partition tensors across threads according to the TiledMMA
ThrMMA thr_mma = tiled_mma.get_slice(thread_idx);
Tensor tCsA = thr_mma.partition_A(sA); // (MMA, MMA_M, MMA_K) smem
Tensor tCsB = thr_mma.partition_B(sB); // (MMA, MMA_N, MMA_K) smem
Tensor tCgC = thr_mma.partition_C(gC); // (MMA, MMA_M, MMA_N) gmem
// Make register tensors the same shape/layout as above
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA, MMA_M, MMA_K) rmem
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA, MMA_N, MMA_K) rmem
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA, MMA_M, MMA_N) rmem
// COPY from smem to rmem thread-level partitions
cute::copy(tCsA, tCrA);
cute::copy(tCsB, tCrB);
// CLEAR rmem thread-level partition (accumulators)
cute::clear(tCrC);
// GEMM on rmem: (V,M,K) x (V,N,K) => (V,M,N)
cute::gemm(tiled_mma, tCrA, tCrB, tCrC);
// Equivalent to
// for(int k = 0; k < size<2>(tCrA); ++k)
// for(int m = 0; m < size<1>(tCrC); ++m)
// for(int n = 0; n < size<2>(tCrC); ++n)
// tiled_mma.call(tCrA(_,m,k), tCrB(_,n,k), tCrC(_,m,n));
// AXPBY from rmem to gmem thread-level partitions
cute::axpby(alpha, tCrC, beta, tCgC);
// Equivalent to
// for(int i = 0; i < size(tCrC); ++i)
// tCgC(i) = alpha * tCrC(i) + beta * tCgC(i)對于上述代碼，現在有許多關于計算和復制指令的時間交錯的決策需要做出
- 僅將 rmem 分配為
A: (MMA,MMA_M)以及 tg_ 33 和 tg_ 34Tensors，并在每次 k-block 迭代時復制到其中。 - 考慮 gmem 的多個 K 圖塊，并在每次 K 圖塊迭代中復制到 smem。
- 將上述復制階段與計算階段異步重疊。
- 通過尋找更好的 smem 布局來優化，從而改進 smem – > rmem 文案的訪問模式。
- 通過為 gmem – > smem 復制找到高效的 TiledCopy 分區模式來進行優化。
這些問題被視為“時間微核”的一部分，而非 CuTe 提供的“空間微核”。通常，有關管線和 CuTe 張量指令執行的決策將由 CUTLASS 級別做出，并將在本系列的下一部分中進行討論。
總結
總之，CuTe 通過抽象出張量布局和線程映射的低級細節，并為現代 NVIDIA GPU 上的密集線性代數提供統一的代數接口，使開發者能夠編寫更具可讀性、可維護性和高性能的 CUDA 代碼。
有關更多信息，您可以在 GitHub 上下載軟件，閱讀我們的文檔，或加入我們的開發者論壇進行更深入的討論。
致謝
感謝 Jack Kosaian、Mark Hoemmen、Haicheng Wu 和 Matt Nicely 為本文做出的貢獻。特別感謝 Jay Shah、Paul VanKoughnett 和 Rya Asai 的 Colfax International 團隊。
?
