西風 發自 凹非寺
量子位 | 公眾號 QbitAI

無需CUDA代碼，給H100加速33%-50%

Flash Attention、Mamba作者之一Tri Dao的新作火了。

他和兩位普林斯頓CS博士生提出了一個名叫QuACK的新SOL內存綁定內核庫，借助CuTe-DSL，完全用Python寫，一點CUDA C++代碼都沒用到。

在帶寬3TB/s的H100上，它的速度比像PyTorch的torch.compile、Liger這類已經過深度優化的庫還要快33%-50%。

Tri Dao表示，讓內存密集型的內核達到“光速”并非什么神秘技巧，只需把幾個細節處理到位就行。

• 我很喜歡Phil Tillet對不同工具在生產力和性能方面各有取舍的觀點，比如torch compile、triton、CUDA、PTX。
• 但CuTe-DSL以及類似的基于Python的DSL或許能改變這一局面，雖然目前還處于早期階段。而且，說不定很快我們就能讓大語言模型來生成這些內核了！

新作一經發出，吸引不少大佬關注。

英偉達CUTLASS團隊資深架構師Vijay轉發，自夸他們團隊做的CuTe-DSL把各種細節都打磨得很好，由此像Tri Dao這樣的專家能夠讓GPU飛速運行。

同時他還預告今年會有更多相關內容推出。

同樣被吸引而來的還有PyTorch團隊成員Horace He，上來就夸贊“太酷了，尤其對于長序列來說”。

不過他還指出，在處理長度不超過約16384的序列時，PyTorch的torch.compile的性能數據能較輕松地得到優化，更接近理想狀態。

接著給出了幾點優化torch.compile性能的建議：

• 默認情況下，若用不同形狀數據測試，torch.compile會生成動態形狀內核，可通過設置dynamic=False關閉該行為。
  進行更多自動調優操作能進一步提升性能。
  torch.compile在生成無循環的持久化歸約內核上較保守，可通過啟用多內核（設置
  (TORCHINDUCTOR_MULTI_KERNEL=1)來讓其自動調優。

最后他表示，還是不可否認QuACK是一項非常出色的工作，而且它也是CuTe-DSL一個很好的教學示例。

Tri Dao也作出了回應，“太棒了，這正是我們想要的，我們會試試這些方法，然后更新圖表”。

食用指南

QuACK作者們寫了一篇教程來介紹具體做法，里面的代碼可以直接使用。

讓內存密集型內核達到“光速”

想讓GPU在模型訓練和推理時都高速運轉，就得同時優化兩種內核：一種是計算密集型（比如矩陣乘法、注意力機制），另一種是內存密集型（像逐元素運算、歸一化、損失函數計算）

其中，矩陣乘法和注意力機制已經是優化得相當到位了。所以作者這次把重點放在內存密集型內核上——這類內核大部分時間都耗在內存訪問（輸入輸出）上，真正用來計算的時間反而不多。

只要搞懂并利用好現代加速器的線程和內存層級結構，就能讓這些內核的速度逼近“理論極限”。而且多虧了最新的CuTe-DSL，不用寫CUDA C或C++代碼，在順手的Python環境里就能做到這一點。

內存密集型的內核有個特點：它的算術強度（也就是浮點運算量FLOPs和傳輸字節數的比值）很小。一旦內核的算術強度落到內存密集型的范疇，它的吞吐量就不再由每秒能完成多少浮點運算決定，而是看每秒能傳輸多少字節了。

在這類內存密集型的內核里，逐元素的激活操作處理起來相對簡單。因為每個元素的計算互不干擾，天生就適合完全并行處理。

不過，像softmax、RMSNorm這些深度學習算子中，還經常用到“歸約”操作，需要對所有值進行聚合。

并行的結合性歸約算法會執行O(log (歸約維度數))輪的部分歸約，這些歸約在不同空間的線程間進行，而作者對 GPU內存層級的了解將在此過程中發揮作用。

并行最大歸約：

接下來，作者將介紹如何利用GPU的內存層級結構來實現高效的歸約內核。

作為示例，使用CuTe DSL實現了大語言模型里常用的三個內核：RMSNorm、softmax和交叉熵損失

目標是達到硬件的最大吞吐量，即 “GPU光速吞吐量”，而這需要兩個關鍵要素：1）全局內存的合并加載/存儲；2）硬件感知的歸約策略。

此外，作者還將解釋集群歸約，以及它如何助力超大規模歸約任務，這是英偉達GPU從Hopper架構（H100）開始引入的一個較新特性。

然后，詳細講解這些關鍵要素的細節，并闡述它們如何幫助編寫“光速”內核。

GPU內存層級結構

在寫內核代碼前，得先搞明白現代GPU的內存層級是怎么回事。這里以Hopper架構的GPU（比如 H100）為例進行說明。

Hopper架構的GPU里，CUDA的執行分為四個層級：線程（threads）、線程塊（thread blocks）、新引入的線程塊集群（thread block cluster）以及完整網格（the full grid）。

單個線程是在流式多處理器（SM）里，以32個線程一組的“warp”形式運行的；每個線程塊擁有一塊192-256 KB的統一共享內存（SMEM），同一線程塊內的所有warp都可訪問該內存。

H100的線程集群允許最多16個運行在相鄰SM上的線程塊，通過分布式共享內存（DSMEM）結構讀取、寫入彼此的共享內存并執行原子操作。這一過程通過低延遲的集群屏障進行協調，從而避免了代價高昂的全局內存往返傳輸。

內存的每個層級都有用于本地歸約的讀寫原語。因此，作者將在CuTe DSL中開發一個通用的歸約模板，使H100在256-262k的歸約維度范圍內始終達到“光速”吞吐量。

H100中的內存層級結構：

Hopper GPU的執行粒度與內存層級之間的對應關系：

每個內存層級的訪問延遲和帶寬都不一樣。

比如，訪問線程自己的寄存器也就幾納秒，訪問共享內存大概要 10-20納秒。再往上，訪問L2緩存的延遲就會飆升到150-200納秒，最后訪問DRAM（主存）得花約400納秒。

帶寬方面，訪問寄存器能達到100 TB/s，訪問共享內存（SMEM）約為20-30 TB/s，訪問L2緩存是5-10 TB/s。對于受內存限制的內核來說，H100的HBM3顯存帶寬（3.35TB/s）往往是性能瓶頸。

所以，為了把硬件性能榨干，設計內存密集型的內核時，得順著內存層級來

最好將大部分本地歸約操作分配在較高的內存層級上，只將少量經過本地歸約后的值傳遞到下一個內存層級。Chris Fleetwood在博客里對A100（不含線程塊集群）的內存訪問延遲進行了類似的說明，而H100則在共享內存（SMEM）和全局內存（GMEM）之間增加了一個額外的內存層級抽象。

H100中內存訪問的延遲：

硬件感知的加載與存儲策略

寫內核代碼時，第一個要解決的問題就是“怎么加載輸入數據、存儲結果”。對于受內存限制的內核來說，HBM的3.35 TB/s通常是瓶頸，這意味著需要在加載和存儲策略上做到極致優化。

在啟動內核之前，首先會通過特定的線程-值布局（TV-layout）對輸入數據進行分區。這決定了每個線程怎么加載和處理歸約維度上的值。

由于每個線程都要從全局內存（GMEM）加載數據，所以得想辦法確保每次加載操作在硬件上連續地傳輸最大數量的bits。這種技術通常被稱為內存合并（memory coalescing）或全局內存的合并訪問（coalesced access to global memory），CUDA最佳實踐指南對這一概念進行了更詳細的解釋。

合并內存訪問：

在H100中，這意味著每個線程處理的數據量得是128bits的倍數， 即4xFP32或者8xBF16。因此，對于FP32來說，這會將4次加載和存儲操作組合（或“向量化”）為一次內存事務，從而最大化吞吐量。

具體操作上，作者會異步地將數據從全局內存（GMEM）加載到共享內存（SMEM），然后將加載操作向量化到寄存器中。等歸約出最終結果后，就直接存回全局內存。

有時候，還可以把輸入數據從全局內存或共享內存重新讀到寄存器，這樣能減少寄存器的占用，避免數據“溢出”。

下面是用Python CuTe DSL寫的加載操作代碼片段，為了看著簡單，這里省略了數據類型轉換和掩碼謂詞的相關代碼。

硬件感知的歸約策略

當每個線程持有一個小的輸入向量后，就可以開始對它們進行歸約了。每次歸約都需要進行一次或多次完整的行掃描。

回想一下，從內存層級的頂層到低層，訪問延遲逐漸增加，而帶寬逐漸減少。

因此，歸約策略應遵循這種硬件內存層級

一旦部分結果存留在內存金字塔的較高層級中，就立即對其進行聚合，只將經過本地歸約后的值傳遞到下一個內存層級。

作者會按照下表從頂層到低層對值進行歸約，并且每一步都只在對應的內存層級中進行加載和存儲操作。

不同內存層級中的歸約策略：

1、線程級歸約（讀寫寄存器）

每個線程會在本地對多個向量化加載的值進行歸約。作者使用TensorSSA.reduce函數，其中需要傳入一個可結合的歸約算子op、歸約前的初始值init_val，以及歸約維度reduction_profile。

2、Warp級歸約（讀寫寄存器）

warp是由32個連續線程組成的固定組，每周期會執行相同的指令。（同步的）warp歸約允許同一Warp內的每個線程通過專用的洗牌（shuffle）網絡，在一個周期內讀取另一個線程的寄存器。經過蝶式warp歸約后，同一warp中的每個線程都會得到歸約后的值。

作者定義了一個輔助函數warp_reduce，用于以“蝶式”歸約順序執行Warp歸約。關于warp級原語的詳細解釋，讀者可參考Yuan和Vinod撰寫的CUDA博客“Using CUDA Warp-Level Primitives”。

蝶式warp歸約（Butterfly warp reduction），也稱為 “xor warp shuffle”：

3、線程塊級歸約（讀寫共享內存）

一個線程塊通常包含多個（在H100中最多32個）warp。在線程塊歸約中，每個參與歸約的warp中的第一個線程會將該warp的歸約結果寫入共享內存中預先分配的歸約緩沖區。

在經過線程塊級同步（屏障）確保所有參與的warp都完成寫入后，每個warp的首線程會從歸約緩沖區中讀取數據，并在本地計算出線程塊的歸約結果。

4、集群歸約（讀寫分布式共享內存）

線程塊集群是Hopper架構中新增的執行層級，由一組相鄰的線程塊（最多16個）組成。同一集群內的線程塊通過分布式共享內存（DSMEM）進行通信，這種內存有專門的高速SM間網絡支持。

在同一集群中，所有線程都可通過DSMEM訪問其他SM的共享內存，其中共享內存的虛擬地址空間在邏輯上分布于集群內的所有線程塊。DSMEM可通過簡單的指針直接訪問。

分布式共享內存：

在集群歸約中，作者首先把當前warp的歸約結果通過專用的SM間網絡（也就是DSMEM），發送到其他對等線程塊的共享內存緩沖區里。

隨后，每個warp從其本地歸約緩沖區中獲取所有warp的值，并對這些值進行歸約。

這里還得用到一個內存屏障用來統計數據到達的數量，以避免過早訪問本地共享內存（否則會導致非法內存訪問的錯誤）。

把整個歸約流程串起來看：首先做線程級歸約，然后在同一個warp內聚合線程級歸約的結果（即warp級歸約），接著根據歸約維度的數量，在每個線程塊或線程塊集群上進一步傳遞歸約后的值。

NCU 性能分析（Softmax 內核）

作者在配備HBM3顯存（DRAM峰值吞吐量=3.35 TB/s）的NVIDIA H100 上，對批量維度為16K、歸約維度為 131K的softmax內核做了性能測試。內存工作負載圖由Nsight Compute生成。

配置是：線程塊集群大小為4，每個線程塊有256個線程，輸入數據類型為FP32。加載和存儲操作都做了向量化處理，每條指令一次搬運128 bits數據（也就是4 FP32值）

最終測出來的DRAM吞吐量也就是顯存帶寬利用率達到了3.01TB/s，相當于DRAM峰值吞吐量的89.7%。除了共享內存（SMEM）外，還高效利用了分布式共享內存（DSMEM）。

該方案的內存工作負載圖：

作者還拿自己的實現與torch.compile（PyTorch 2.7.1 版本）進行了對比。

首先，獲取了torch.compile生成的Triton內核代碼。

該內核實現softmax時包含2次全局內存加載（計算行最大值和部分指數和時各加載1次，以及最終的softmax 值）和1次存儲。

在這種情況下，盡管該Triton內核仍能使硬件的DRAM吞吐量跑滿，但額外的1次不必要加載會導致Triton 內核的有效模型內存吞吐量（約2.0 TB/s）僅為本文作者實現方案的三分之二（約3.0 TB/s）

torch.compile生成的Triton內核（調優配置部分省略）

內存吞吐量

作者對RMSNorm、softmax和交叉熵損失這幾個內核做了基準測試。測試仍在配備HBM3顯存的1塊NVIDIA H100 80GB GPU和Intel Xeon Platinum 8468 CPU上進行。

測試中使用的批量大小范圍為8k-32k，歸約維度范圍為256-262k（256×1024），輸入數據類型為FP32和 BF16。

基準對比方案如下：

• Torch.compile（PyTorch 2.7.1版本）：使用默認編譯模式。
• Liger 內核 v0.5.10版本 ：只測了RMSNorm 和 softmax，且歸約維度最多到65k（因為它目前不支持更大的維度）。
• cuDNN v9.10.1版本：只測了RMSNorm內核。

作者基于CuTe DSL的實現方案，在歸約維度大于4k時，內存吞吐量一般能穩定在3TB/s左右（差不多是峰值的90%）

歸約維度262k時，FP32的softmax吞吐量能到3.01TB/s，而torch.compile只有1.89TB/s，快了近50%。對于這3個內核，當歸約維度≥65k 時，該實現方案顯著優于所有基準對比方案。

多個內核的模型內存吞吐量：

作者認為，在輸入規模≥65k時的優異性能得益于成功利用了H100中的集群歸約

當輸入數據量大到把SM的寄存器和共享內存都占滿時，如果不用集群歸約，就只能換成在線算法（比如在線softmax）；不然的話，寄存器里的數據會大量“溢出”，導致吞吐量顯著下降。

舉個例子，作者觀察到，當使用Liger softmax內核時，輸入規模從32k漲到65k，吞吐量就從約3.0 TB/s掉到了2.0 TB/s左右。

作者用NCU（Nsight Compute）工具分析了它的內存負載圖和SASS代碼，發現當每個SM要加載65k數據時，SM的資源被耗盡，結果就是大量寄存器溢出，還會頻繁往HBM里回寫數據，這才拖慢了速度。

Liger softmax內核在批量維度為16k、歸約維度為65k且數據類型為FP32時的內存工作負載：

Liger softmax內核匯編代碼中的寄存器溢出（LDL指令）：

但集群歸約能讓多個SM協同工作，共享各自的資源，相當于組成一個“超級”SM（靠DSMEM實現）。

假設單個SM僅能處理32k輸入，那么一個大小為16的集群將允許處理50萬（0.5M）輸入，而無需從全局內存（GMEM）重新加載數據。

由于作者對硬件知識有清晰的理解，即使使用常規的3遍掃描softmax算法，也能輕松充分利用所有內存層級的每一個字節，實現“光速”級別的吞吐量。

總結

作者通過實踐證明，只要精心手工編寫CuTe內核，就能把硬件里所有內存層級的潛力都榨干，實現“光速”級別的內存吞吐量。

但這種效率是以針對每個算子甚至每個輸入形狀進行調優為代價的，這自然在效率與開發成本之間形成了一種權衡。

Phil Tillet（Triton的作者）在他的演講中用這張圖很好地闡述了這一點。

根據作者使用CuTe-DSL的經驗，它既具備Python的開發效率，又擁有CUDA C++的控制能力和性能。

作者認為，高效的GPU內核開發流程是可以自動化的。

例如，RMSNorm中的輸入張量TV布局、加載/存儲策略以及歸約輔助函數，可直接應用于softmax內核并仍能達到相近的吞吐量。

此外，CuTe DSL將為開發者或基于CuTe DSL運行的其他代碼生成應用提供靈活的GPU內核開發能力。

• 目前，將大語言模型應用于自動生成GPU內核是一個活躍的研究方向，未來，或許只需調用“LLM.compile” 就能生成高度優化的GPU內核。

作者簡介

這項工作作者有三位。

Wentao Guo

Wentao Guo目前是普林斯頓大學計算機科學專業的博士生，師從Tri Dao。

在這之前，他在康奈爾大學獲得了計算機科學的本科和碩士學位。

Ted Zadouri

Ted Zadouri同樣是普林斯頓大學計算機科學專業的博士生，本碩分別讀自加州大學歐文分校、加州大學洛杉磯分校

此前Ted Zadouri曾在英特爾實習，也曾在Cohere研究大語言模型的參數高效微調。

Tri Dao

Tri Dao目前是普林斯頓大學計算機科學助理教授，還是生成式AI初創公司Together AI的首席科學家。

他因提出一系列優化Transformer模型注意力機制的工作而聞名學界。

其中最有影響力的，是其作為作者之一提出了Mamba架構，這一架構在語言、音頻和基因組學等多種模態中都達到了SOTA性能。

尤其在語言建模方面，無論是預訓練還是下游評估，Mamba-3B模型都優于同等規模的Transformer模型，并能與兩倍于其規模的Transformer模型相媲美。

另外他還參與發表了FlashAttention1-3版本，FlashAttention被廣泛用于加速Transformers，已經使注意力速度提高了4-8倍。

GitHub鏈接：https://github.com/Dao-AILab/quack/blob/main/media/2025-07-10-membound-sol.md
[1]https://x.com/__tensorcore__/status/1943374119208169858
[2]https://x.com/tri_dao/status/1943372100774900056