无需CUDA代码给H100加速33%-50%,Flash Attention作者新作火了

西风 发自 凹非寺
量子位 | 公众号 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

(文:量子位)

发表评论