OpenAI Triton编程入门指南

英伟达

CUDA编程变得愈发复杂，这一点饱受诟病，也让人不禁感叹天下苦英伟达久矣。借着大模型时代的契机，近期开发社区里Triton之风盛行，网络上也有不少优质回答在介绍Triton入门。Triton火热的关键在于，这种基于Tile的编程范式能以较少的代码量实现接近CUDA的程序性能，就连PyTorch官方都把它纳入了torch.compile后端全家桶。不过，对于AI Infra开发者来讲，Triton真的那么好用吗？举个例子，就像下图展示的那样，在Triton里构建一个FlashAttention内核大概得写700行Python代码（要是用CUDA则需要大概7000行C++代码）。我们能看到，开发者要编写这样的Python代码可不容易，不但得深入领会FlashAttention优化算法，还得熟练运用Block编程技能，像是Program的拆分、Block的排列、Index的计算、DRAM/SRAM的转换之类的。这在一定程度上意味着开发者得重新学习一套新的硬件抽象知识，其难度说不定跟学习GPU差不多。鉴于新兴硬件以及GPU新特性不断涌现，日后Triton要在通用性（Generality）和专一性（Specificity）之间达成平衡也变得更难了。对于Triton而言，会不会还没挑战CUDA的地位就先走上失败的老路？

拖了很久，最近总算有空到网络分享我们的新工作Mirage了。Mirage项目以SuperOptimization技术（ABS/2405.05751">https://arxiv.org/ABS/2405.05751）为基础，它能在用户不用编写CUDA或者Triton代码的情况下，自动为PyTorch生成GPU内核算子，从而实现更优的性能。例如，对于FlashAttention算子，用户仅需写几行Python代码描述注意力计算过程，不必了解GPU编程细节，如下：写起来是不是很容易？Mirage能够自动搜索可能的Attention GPU内核实现方式。其搜索范围既包含现有的手动设计的注意力内核，像FlashAttention和FlashDecoding，还涵盖了在某些场景下比当前手写版本快达3.5倍的其他实现。Mirage生成的GPU内核可直接对PyTorch张量进行操作，而且能在PyTorch程序里直接调用。Mirage提供了一种新编程范式，和CUDA/Triton编程比有三个主要优势。

GPU计算的内核函数会以单程序多数据（SPMD）的形式，同时在多个流处理器（SM）上运行。GPU内核（Kernel）依靠由线程块（Thread Block）构成的网格结构来安排计算工作，每个线程块都在一个单独的SM上运行。每个线程块还包含多个线程（Thread），这些线程用于对单个的数据元素进行计算。GPU有着复杂的内存层次结构，以适配这种复杂的处理结构。每个线程都有自己的寄存器文件（Register File），这样能快速访问数据。线程块内的所有线程能够访问一块公共的共享内存（Shared Memory），这有利于线程之间高效地进行数据交换和集体操作。内核中的所有线程都可以访问分配给整个GPU的大容量设备内存（Device Memory）。

AI

Mirage运用Graph来对GPU内核进行描述，Graph具有多个层次，其分别代表内核、线程块以及线程级别的计算。大致而言，Kernel Graph、Thread Block Graph和Thread Graph依次代表整个GPU、一个流处理器（SM）以及一个CUDA/tensor核心上的计算。如果读者对Graph的细节感兴趣，可以查看：https://mirage - project.readthedocs.io/en/latest/mugraph.html。

上图呈现了Mirage的工作流程。针对输入的PyTorch程序，Mirage的Graph生成器会自动搜索功能等同于输入程序的其他Graph，其搜索范围包含内核、线程块和线程级别的各类GPU优化。全部生成的Graph都会被送至等价性验证器，该验证器自动检验每个Graph是否和目标程序等价。接着，Graph转译器把所有通过验证的Graph转译为CUDA内核。最终，Mirage会返回其中性能最优的CUDA内核。多个LLM/GenAI基准测试结果表明，Mirage生成的内核往往比现有的手写或编译器生成的内核快1.2到2.5倍。接下来，本文将以LLM中的Transformer架构为例，展示当前系统所缺失的几项GPU程序优化技术。

在当今的机器学习模型里，像LayerNorm、RMSNorm、GroupNorm和BatchNorm这样的归一化（Normalization）操作被广泛运用。现有的机器学习编译器往往会在单独的内核中启动归一化层，这是因为归一化包含归约和广播操作，很难与其他计算融合。不过，Mirage察觉到，多数归一化层经过适当的代数变换后，是能够和后续的线性层（如MatMul）融合的。

Mirage发现的自定义内核利用RMSNorm里除法与MatMul中乘法的可交换性，把除法挪到MatMul之后。这种变换功能上等效，还避免了中间张量Y的实例化。该内核性能比单独运行这两个操作要快1.5至1.7倍。

LoRA在预训练模型的微调方面应用广泛，可使其适配特定领域与任务。其适配器一般插入模型线性层，会带来额外矩阵乘法。当前系统往往为原始矩阵乘法和LoRA中的两个矩阵乘法分别启动独立内核，这就造成了较高的内核启动成本。

如上图所示，Mirage找到了一个能把三个矩阵乘法和后续加法融合成单个内核的内核。其通过将计算重构成两个线程块级别的矩阵乘法达成这一目的，还利用了代数变换（此处未详细说明），其中两个拼接操作无计算过程，靠更新GPU共享内存中的张量偏移量完成。Mirage发现的这个内核比现有系统所用内核快1.6倍。

许多大型语言模型（如LLAMA - 2、LLAMA - 3及其变体）都在使用Gated MLP层。其输入张量X会与两个权重矩阵相乘，然后组合输出结果得到最终结果。Mirage发现了一个内核，这个内核能执行两个矩阵乘法、SiLU激活以及后续的逐元素乘法，这减少了内核启动开销和对设备内存的访问。

当下，多数大型语言模型（LLM）以注意力及其变体为基础。尽管现有系统往往会提供高度优化的注意力实现方式，像FlashAttention、FlashInfer和FlexAttention，不过要支持注意力变体常常需要新的自定义内核。下面将通过两个例子，展示Mirage是怎样为非传统注意力计算找到自定义GPU内核的。不少近期的LLM架构，像Chameleon、ViT - 22B等，都在LLaMA架构里引入了QK - Norm以减轻训练时数值发散的问题。QK - Norm会在注意力之前给Query和Key向量运用LayerNorm层。现有的注意力实现并不支持这些额外的归一化层，而且这些归一化层还得作为独立内核启动。

注意力存在在其之前和/或之后引入计算的变体形式。这些计算可与注意力融合以提升GPU性能，但这需要自定义内核。针对带有QK - Norm的注意力，Mirage找到了能融合计算的上述内核，如此就可避免在GPU设备内存中生成中间结果。这个自定义内核还对注意力进行了GPU现有的优化，让性能提升了1.7至2.5倍。

MLA（Multi - Head Latent Attention）是另一种常用的注意力变体。它把注意力的KV Cache压缩成一个向量，从而削减存储KV Cache的内存成本。并且，在注意力之前会引入两个线性层，所示。和QK - Norm一样，现有的注意力实现并不支持这些额外的归一化层，也需要作为独立内核启动。不过，Mirage能够把线性层和注意力融合成一个单独的自定义内核。Mirage项目的长远目标是，让未来的AI开发者不必学习CUDA、Triton这类复杂的GPU编程语言。开发者只需明确所需的数学操作，就能在GPU上顺利构建AI模型。借助Mirage的SuperOptimization技术，各类计算任务可自动转化为经过高度优化的GPU实现形式。随着大型语言模型（LLM）和其他生成式AI应用的快速发展，在众多实际部署场景下，高效的GPU支持不可或缺，降低GPU编程难度、提升程序效率也变得日益关键。