西風(fēng) 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

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

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

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

在帶寬3TB/s的H100上，它的速度比像PyTorch的torch.compile、Liger這類已經(jīng)過深度優(yōu)化的庫(kù)還要快33%-50%。

Tri Dao表示，讓內(nèi)存密集型的內(nèi)核達(dá)到“光速”并非什么神秘技巧，只需把幾個(gè)細(xì)節(jié)處理到位就行。

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

新作一經(jīng)發(fā)出，吸引不少大佬關(guān)注。

英偉達(dá)CUTLASS團(tuán)隊(duì)資深架構(gòu)師Vijay轉(zhuǎn)發(fā)，自夸他們團(tuán)隊(duì)做的CuTe-DSL把各種細(xì)節(jié)都打磨得很好，由此像Tri Dao這樣的專家能夠讓GPU飛速運(yùn)行。

同時(shí)他還預(yù)告今年會(huì)有更多相關(guān)內(nèi)容推出。

同樣被吸引而來的還有PyTorch團(tuán)隊(duì)成員Horace He，上來就夸贊“太酷了，尤其對(duì)于長(zhǎng)序列來說”。

不過他還指出，在處理長(zhǎng)度不超過約16384的序列時(shí)，PyTorch的torch.compile的性能數(shù)據(jù)能較輕松地得到優(yōu)化，更接近理想狀態(tài)。

接著給出了幾點(diǎn)優(yōu)化torch.compile性能的建議：

• 默認(rèn)情況下，若用不同形狀數(shù)據(jù)測(cè)試，torch.compile會(huì)生成動(dòng)態(tài)形狀內(nèi)核，可通過設(shè)置dynamic=False關(guān)閉該行為。
  進(jìn)行更多自動(dòng)調(diào)優(yōu)操作能進(jìn)一步提升性能。
  torch.compile在生成無循環(huán)的持久化歸約內(nèi)核上較保守，可通過啟用多內(nèi)核（設(shè)置
  (TORCHINDUCTOR_MULTI_KERNEL=1)來讓其自動(dòng)調(diào)優(yōu)。

最后他表示，還是不可否認(rèn)QuACK是一項(xiàng)非常出色的工作，而且它也是CuTe-DSL一個(gè)很好的教學(xué)示例。

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

食用指南

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

讓內(nèi)存密集型內(nèi)核達(dá)到“光速”

想讓GPU在模型訓(xùn)練和推理時(shí)都高速運(yùn)轉(zhuǎn)，就得同時(shí)優(yōu)化兩種內(nèi)核：一種是計(jì)算密集型（比如矩陣乘法、注意力機(jī)制），另一種是內(nèi)存密集型（像逐元素運(yùn)算、歸一化、損失函數(shù)計(jì)算）

其中，矩陣乘法和注意力機(jī)制已經(jīng)是優(yōu)化得相當(dāng)?shù)轿涣恕Ｋ宰髡哌@次把重點(diǎn)放在內(nèi)存密集型內(nèi)核上——這類內(nèi)核大部分時(shí)間都耗在內(nèi)存訪問（輸入輸出）上，真正用來計(jì)算的時(shí)間反而不多。

只要搞懂并利用好現(xiàn)代加速器的線程和內(nèi)存層級(jí)結(jié)構(gòu)，就能讓這些內(nèi)核的速度逼近“理論極限”。而且多虧了最新的CuTe-DSL，不用寫CUDA C或C++代碼，在順手的Python環(huán)境里就能做到這一點(diǎn)。

內(nèi)存密集型的內(nèi)核有個(gè)特點(diǎn)：它的算術(shù)強(qiáng)度（也就是浮點(diǎn)運(yùn)算量FLOPs和傳輸字節(jié)數(shù)的比值）很小。一旦內(nèi)核的算術(shù)強(qiáng)度落到內(nèi)存密集型的范疇，它的吞吐量就不再由每秒能完成多少浮點(diǎn)運(yùn)算決定，而是看每秒能傳輸多少字節(jié)了。

在這類內(nèi)存密集型的內(nèi)核里，逐元素的激活操作處理起來相對(duì)簡(jiǎn)單。因?yàn)槊總€(gè)元素的計(jì)算互不干擾，天生就適合完全并行處理。

不過，像softmax、RMSNorm這些深度學(xué)習(xí)算子中，還經(jīng)常用到“歸約”操作，需要對(duì)所有值進(jìn)行聚合。

并行的結(jié)合性歸約算法會(huì)執(zhí)行O(log (歸約維度數(shù)))輪的部分歸約，這些歸約在不同空間的線程間進(jìn)行，而作者對(duì) GPU內(nèi)存層級(jí)的了解將在此過程中發(fā)揮作用。

并行最大歸約：

接下來，作者將介紹如何利用GPU的內(nèi)存層級(jí)結(jié)構(gòu)來實(shí)現(xiàn)高效的歸約內(nèi)核。

作為示例，使用CuTe DSL實(shí)現(xiàn)了大語(yǔ)言模型里常用的三個(gè)內(nèi)核：RMSNorm、softmax和交叉熵?fù)p失

目標(biāo)是達(dá)到硬件的最大吞吐量，即 “GPU光速吞吐量”，而這需要兩個(gè)關(guān)鍵要素：1）全局內(nèi)存的合并加載/存儲(chǔ)；2）硬件感知的歸約策略。

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

然后，詳細(xì)講解這些關(guān)鍵要素的細(xì)節(jié)，并闡述它們?nèi)绾螏椭帉憽肮馑佟眱?nèi)核。

GPU內(nèi)存層級(jí)結(jié)構(gòu)

在寫內(nèi)核代碼前，得先搞明白現(xiàn)代GPU的內(nèi)存層級(jí)是怎么回事。這里以Hopper架構(gòu)的GPU（比如 H100）為例進(jìn)行說明。

Hopper架構(gòu)的GPU里，CUDA的執(zhí)行分為四個(gè)層級(jí)：線程（threads）、線程塊（thread blocks）、新引入的線程塊集群（thread block cluster）以及完整網(wǎng)格（the full grid）。

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

H100的線程集群允許最多16個(gè)運(yùn)行在相鄰SM上的線程塊，通過分布式共享內(nèi)存（DSMEM）結(jié)構(gòu)讀取、寫入彼此的共享內(nèi)存并執(zhí)行原子操作。這一過程通過低延遲的集群屏障進(jìn)行協(xié)調(diào)，從而避免了代價(jià)高昂的全局內(nèi)存往返傳輸。

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

H100中的內(nèi)存層級(jí)結(jié)構(gòu)：

Hopper GPU的執(zhí)行粒度與內(nèi)存層級(jí)之間的對(duì)應(yīng)關(guān)系：

每個(gè)內(nèi)存層級(jí)的訪問延遲和帶寬都不一樣。

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

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

所以，為了把硬件性能榨干，設(shè)計(jì)內(nèi)存密集型的內(nèi)核時(shí)，得順著內(nèi)存層級(jí)來

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

H100中內(nèi)存訪問的延遲：

硬件感知的加載與存儲(chǔ)策略

寫內(nèi)核代碼時(shí)，第一個(gè)要解決的問題就是“怎么加載輸入數(shù)據(jù)、存儲(chǔ)結(jié)果”。對(duì)于受內(nèi)存限制的內(nèi)核來說，HBM的3.35 TB/s通常是瓶頸，這意味著需要在加載和存儲(chǔ)策略上做到極致優(yōu)化。

在啟動(dòng)內(nèi)核之前，首先會(huì)通過特定的線程-值布局（TV-layout）對(duì)輸入數(shù)據(jù)進(jìn)行分區(qū)。這決定了每個(gè)線程怎么加載和處理歸約維度上的值。

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

合并內(nèi)存訪問：

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

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

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

下面是用Python CuTe DSL寫的加載操作代碼片段，為了看著簡(jiǎn)單，這里省略了數(shù)據(jù)類型轉(zhuǎn)換和掩碼謂詞的相關(guān)代碼。

硬件感知的歸約策略

當(dāng)每個(gè)線程持有一個(gè)小的輸入向量后，就可以開始對(duì)它們進(jìn)行歸約了。每次歸約都需要進(jìn)行一次或多次完整的行掃描。

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

因此，歸約策略應(yīng)遵循這種硬件內(nèi)存層級(jí)

一旦部分結(jié)果存留在內(nèi)存金字塔的較高層級(jí)中，就立即對(duì)其進(jìn)行聚合，只將經(jīng)過本地歸約后的值傳遞到下一個(gè)內(nèi)存層級(jí)。

作者會(huì)按照下表從頂層到低層對(duì)值進(jìn)行歸約，并且每一步都只在對(duì)應(yīng)的內(nèi)存層級(jí)中進(jìn)行加載和存儲(chǔ)操作。

不同內(nèi)存層級(jí)中的歸約策略：

1、線程級(jí)歸約（讀寫寄存器）

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

2、Warp級(jí)歸約（讀寫寄存器）

warp是由32個(gè)連續(xù)線程組成的固定組，每周期會(huì)執(zhí)行相同的指令。（同步的）warp歸約允許同一Warp內(nèi)的每個(gè)線程通過專用的洗牌（shuffle）網(wǎng)絡(luò)，在一個(gè)周期內(nèi)讀取另一個(gè)線程的寄存器。經(jīng)過蝶式warp歸約后，同一warp中的每個(gè)線程都會(huì)得到歸約后的值。

作者定義了一個(gè)輔助函數(shù)warp_reduce，用于以“蝶式”歸約順序執(zhí)行Warp歸約。關(guān)于warp級(jí)原語(yǔ)的詳細(xì)解釋，讀者可參考Yuan和Vinod撰寫的CUDA博客“Using CUDA Warp-Level Primitives”。

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

3、線程塊級(jí)歸約（讀寫共享內(nèi)存）

一個(gè)線程塊通常包含多個(gè)（在H100中最多32個(gè)）warp。在線程塊歸約中，每個(gè)參與歸約的warp中的第一個(gè)線程會(huì)將該warp的歸約結(jié)果寫入共享內(nèi)存中預(yù)先分配的歸約緩沖區(qū)。

在經(jīng)過線程塊級(jí)同步（屏障）確保所有參與的warp都完成寫入后，每個(gè)warp的首線程會(huì)從歸約緩沖區(qū)中讀取數(shù)據(jù)，并在本地計(jì)算出線程塊的歸約結(jié)果。

4、集群歸約（讀寫分布式共享內(nèi)存）

線程塊集群是Hopper架構(gòu)中新增的執(zhí)行層級(jí)，由一組相鄰的線程塊（最多16個(gè)）組成。同一集群內(nèi)的線程塊通過分布式共享內(nèi)存（DSMEM）進(jìn)行通信，這種內(nèi)存有專門的高速SM間網(wǎng)絡(luò)支持。

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

分布式共享內(nèi)存：

在集群歸約中，作者首先把當(dāng)前warp的歸約結(jié)果通過專用的SM間網(wǎng)絡(luò)（也就是DSMEM），發(fā)送到其他對(duì)等線程塊的共享內(nèi)存緩沖區(qū)里。

隨后，每個(gè)warp從其本地歸約緩沖區(qū)中獲取所有warp的值，并對(duì)這些值進(jìn)行歸約。

這里還得用到一個(gè)內(nèi)存屏障用來統(tǒng)計(jì)數(shù)據(jù)到達(dá)的數(shù)量，以避免過早訪問本地共享內(nèi)存（否則會(huì)導(dǎo)致非法內(nèi)存訪問的錯(cuò)誤）。

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

NCU 性能分析（Softmax 內(nèi)核）

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

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

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

該方案的內(nèi)存工作負(fù)載圖：

作者還拿自己的實(shí)現(xiàn)與torch.compile（PyTorch 2.7.1 版本）進(jìn)行了對(duì)比。

首先，獲取了torch.compile生成的Triton內(nèi)核代碼。

該內(nèi)核實(shí)現(xiàn)softmax時(shí)包含2次全局內(nèi)存加載（計(jì)算行最大值和部分指數(shù)和時(shí)各加載1次，以及最終的softmax 值）和1次存儲(chǔ)。

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

torch.compile生成的Triton內(nèi)核（調(diào)優(yōu)配置部分省略）

內(nèi)存吞吐量

作者對(duì)RMSNorm、softmax和交叉熵?fù)p失這幾個(gè)內(nèi)核做了基準(zhǔn)測(cè)試。測(cè)試仍在配備HBM3顯存的1塊NVIDIA H100 80GB GPU和Intel Xeon Platinum 8468 CPU上進(jìn)行。

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

基準(zhǔn)對(duì)比方案如下：

• Torch.compile（PyTorch 2.7.1版本）：使用默認(rèn)編譯模式。
• Liger 內(nèi)核 v0.5.10版本 ：只測(cè)了RMSNorm 和 softmax，且歸約維度最多到65k（因?yàn)樗壳安恢С指蟮木S度）。
• cuDNN v9.10.1版本：只測(cè)了RMSNorm內(nèi)核。

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

歸約維度262k時(shí)，F(xiàn)P32的softmax吞吐量能到3.01TB/s，而torch.compile只有1.89TB/s，快了近50%。對(duì)于這3個(gè)內(nèi)核，當(dāng)歸約維度≥65k 時(shí)，該實(shí)現(xiàn)方案顯著優(yōu)于所有基準(zhǔn)對(duì)比方案。

多個(gè)內(nèi)核的模型內(nèi)存吞吐量：

作者認(rèn)為，在輸入規(guī)模≥65k時(shí)的優(yōu)異性能得益于成功利用了H100中的集群歸約

當(dāng)輸入數(shù)據(jù)量大到把SM的寄存器和共享內(nèi)存都占滿時(shí)，如果不用集群歸約，就只能換成在線算法（比如在線softmax）；不然的話，寄存器里的數(shù)據(jù)會(huì)大量“溢出”，導(dǎo)致吞吐量顯著下降。

舉個(gè)例子，作者觀察到，當(dāng)使用Liger softmax內(nèi)核時(shí)，輸入規(guī)模從32k漲到65k，吞吐量就從約3.0 TB/s掉到了2.0 TB/s左右。

作者用NCU（Nsight Compute）工具分析了它的內(nèi)存負(fù)載圖和SASS代碼，發(fā)現(xiàn)當(dāng)每個(gè)SM要加載65k數(shù)據(jù)時(shí)，SM的資源被耗盡，結(jié)果就是大量寄存器溢出，還會(huì)頻繁往HBM里回寫數(shù)據(jù)，這才拖慢了速度。

Liger softmax內(nèi)核在批量維度為16k、歸約維度為65k且數(shù)據(jù)類型為FP32時(shí)的內(nèi)存工作負(fù)載：

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

但集群歸約能讓多個(gè)SM協(xié)同工作，共享各自的資源，相當(dāng)于組成一個(gè)“超級(jí)”SM（靠DSMEM實(shí)現(xiàn)）。

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

由于作者對(duì)硬件知識(shí)有清晰的理解，即使使用常規(guī)的3遍掃描softmax算法，也能輕松充分利用所有內(nèi)存層級(jí)的每一個(gè)字節(jié)，實(shí)現(xiàn)“光速”級(jí)別的吞吐量。

總結(jié)

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

但這種效率是以針對(duì)每個(gè)算子甚至每個(gè)輸入形狀進(jìn)行調(diào)優(yōu)為代價(jià)的，這自然在效率與開發(fā)成本之間形成了一種權(quán)衡。

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

根據(jù)作者使用CuTe-DSL的經(jīng)驗(yàn)，它既具備Python的開發(fā)效率，又擁有CUDA C++的控制能力和性能。

作者認(rèn)為，高效的GPU內(nèi)核開發(fā)流程是可以自動(dòng)化的。

例如，RMSNorm中的輸入張量TV布局、加載/存儲(chǔ)策略以及歸約輔助函數(shù)，可直接應(yīng)用于softmax內(nèi)核并仍能達(dá)到相近的吞吐量。

此外，CuTe DSL將為開發(fā)者或基于CuTe DSL運(yùn)行的其他代碼生成應(yīng)用提供靈活的GPU內(nèi)核開發(fā)能力。

• 目前，將大語(yǔ)言模型應(yīng)用于自動(dòng)生成GPU內(nèi)核是一個(gè)活躍的研究方向，未來，或許只需調(diào)用“LLM.compile” 就能生成高度優(yōu)化的GPU內(nèi)核。

作者簡(jiǎn)介

這項(xiàng)工作作者有三位。

Wentao Guo

Wentao Guo目前是普林斯頓大學(xué)計(jì)算機(jī)科學(xué)專業(yè)的博士生，師從Tri Dao。

在這之前，他在康奈爾大學(xué)獲得了計(jì)算機(jī)科學(xué)的本科和碩士學(xué)位。

Ted Zadouri

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

此前Ted Zadouri曾在英特爾實(shí)習(xí)，也曾在Cohere研究大語(yǔ)言模型的參數(shù)高效微調(diào)。

Tri Dao

Tri Dao目前是普林斯頓大學(xué)計(jì)算機(jī)科學(xué)助理教授，還是生成式AI初創(chuàng)公司Together AI的首席科學(xué)家。

他因提出一系列優(yōu)化Transformer模型注意力機(jī)制的工作而聞名學(xué)界。

其中最有影響力的，是其作為作者之一提出了Mamba架構(gòu)，這一架構(gòu)在語(yǔ)言、音頻和基因組學(xué)等多種模態(tài)中都達(dá)到了SOTA性能。

尤其在語(yǔ)言建模方面，無論是預(yù)訓(xùn)練還是下游評(píng)估，Mamba-3B模型都優(yōu)于同等規(guī)模的Transformer模型，并能與兩倍于其規(guī)模的Transformer模型相媲美。

另外他還參與發(fā)表了FlashAttention1-3版本，F(xiàn)lashAttention被廣泛用于加速Transformers，已經(jīng)使注意力速度提高了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