PyTorch官宣 GPU推理迎来Triton加速新时代 告别CUDA

新智元报道

:乔杨 Frey

【新智元导读】 用英伟达的GPU,但可以不用CUDA?PyTorch官宣,借助OpenAI开发的Triton语言编写内核来加速LLM推理,可以实现和CUDA类似甚至更佳的性能。

试问,有多少机器学习小白曾被深度学习框架和CUDA的兼容问题所困扰?

又有多少开发者曾因为频频闪烁的警报「CUDA版本必须与安装的PyTorch匹配!!!」而企图炸键盘?

无论是TensorFlow还是Pytorch,GPU和CUDA搭配的概念早已深入骨髓。

如果我说,就在昨天,有款为LLM「量身定做」的CUDA-free推理上新了!你激不激动?

原文地址:https://pytorch.org/blog/cuda-free-inference-for-llms/?hss_channel=tw-776585502606721024

那么,让我们紧跟Pytorch的官方技术博客,一探究竟!看看它是如何将「自由」变为现实!

GPU的好搭子CUDA

CUDA(Compute Unified Device Architecture)到底是何方神物?为何被视为GPU的好搭子,LLMs的「利器」?

它是由英伟达开发的用于并行计算平台和应用程序的编程API,让开发者能通过GPU开展高性能计算,包括:

1. 多个能并行处理任务的核心,实现多线程

2. 多种高效管理GPU内存的方法,如全局内存、共享内存和常量内存

3. 创建并管理多条并行线程,提高数据处理效率

4. 编译器、调试器和性能分析工具组成的工具链,,帮助开发者优化代码

简而言之,CUDA使GPU加速LLM训练变为现实,大幅缩短了训练时间。

100%的Triton内核

Pytorch最近发表了一篇技术博客,他们以两个模型——Llama3-8B和IBM的Granite-8B Code为例,100%使用Triton内核实现了FP16推理。

Granite-8B Code是由IBM开发的一种仅限解码器的代码模型,专为代码生成任务设计。

仓库地址:https://huggingface.co/ibm-granite/granite-8b-code-base-4k

值得注意的是,PyTorch指出他们实现了F16推理,也就是使用半精度浮点计算。

FP32单精度浮点数

F16半精度浮点数

相对于FP32,使用FP16可以将位数减少一半,因而减少了所需内存,允许使用更大的模型或更大的批大小,且数据传输速度更快。

与F32相比,英伟达GPU提供的FP16将算术吞吐量提高了8倍,大幅加快了数学受限层的训练速度。

此外,PyTorch团队还着重强调,计算全部是依赖OpenAI的Triton语言执行的。

Triton是一种用于编写高效自定义深度学习基元的语言和编译器。

Triton的开发者致力于建立一个开源环境,以比CUDA更高效地编写代码,同时也期望它比现有的特定领域语言(domain-specific language)更具灵活性。

论文:https://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf

仓库:https://github.com/triton-lang/triton

团队发现,在英伟达H100上使用Triton内核训练模型,性能可达CUDA内核的76%~78%,在A100上也能达到62%~82%。

既然相比CUDA有一定的性能损失,那为什么要全部使用Triton语言?

PyTorch团队称,Triton实现了LLM在GPU上的「可移植性」,能跨越多个不同个品牌的硬件,如英伟达、AMD、英特尔等。

此外,它还在Python中为GPU编程提供了更高的「抽象层」,使开发者有机会编写自定义的具备更高性能的内核。

最终,通过在H100和A100上使用Llama3-8B和Granite-8B的Triton和CUDA变体,并进行推理阶段的基准测试,PyTorch团队证实了,Triton内核能实现CUDA-Free的计算,且生成token的吞吐量有显著提升。

内核架构

以Llama3为例,经典的Transformer块由一般由以下部分组成:

其中涉及的核心操作包括:

- RMS归一化

- 矩阵乘法:融合QKV矩阵

- 旋转位置编码(RoPE)

- Flash Attention

- 矩阵乘法:投影为为输出矩阵

- RMS归一化

- 矩阵乘法:融合门控+向上投影

- 激活函数SiLU

- 逐元素(element-wise)矩阵乘法

- 矩阵乘法:向下投影

这些操作中都需要一个或多个GPU内核进行计算,虽然不同的Transformer模型的执行细节可能有所不同,但核心操作是类似的。

例如,与Llama 3不同,IBM的Granite 8B Code模型在MLP层中使用了bias,此类更改确实需要对内核的修改。

将这些Transformer块堆叠在一起,再连接编码层,就组成了一个经典的Transformer模型。

模型推理

这些架构代码都会包含在model.py文件中,在PyTorch的eager执行模式下,C会启动CUDA内核执行这些代码。

为了让Llama3-8B和Granite-8B模型100%用Triton语言实现端到端推理,我们需要手写Triton内核(kernel),或利用torch.compile模块自动生成。

对于较小的操作,比如 RMS归一化、RoPE、SiLU函数和element-wise矩阵乘法,torch.compile可以自动生成Triton内核。

使用Nsight等工具即可对这些内核进行观察,如下图所示,自动生成的内核显示为QKV乘法和flash attention之前的深绿色方块:

使用torch.compile跟踪 Llama3-8B,显示CUDA内核

通过Nsight的跟踪信息可以观察到,在Llama3-8B中,占端到端延迟80%的两个主要操作是矩阵乘法和注意力内核,而且它们依旧由CUDA内核操作。

为了进一步提升性能,我们开始手写Triton内核来替换上述两个操作。

手写Triton内核

矩阵乘法

对于线性层中的矩阵乘法,编写一个自定义的 FP16 Triton GEMM (General Matrix-Matrix Multiply)内核,执行通用的矩阵-矩阵乘法,其中利用了SplitK进行工作分解。

为了实现最佳性能,还使用了穷举搜索来调整SplitK GEMM内核。

因为每个线性层的权重矩阵都有不同的形状,如果要获得最佳性能,就需要针对每种矩阵形状调整Triton内核。

Granite-8B和Llama3-8B的线性层权重矩阵规格如下:

调整每个线性层后,相比未调整的Triton内核,可以实现1.2倍的端到端加速。

Flash Attention

Triton的flash attention内核有一系列不同的配置和实现,包括:

- OpenAI Flash

- Dao AI Lab Flash

- XFormers Flash

- PyTorch FlexAttention

首先,采用eager模式,之后用torch.compile的标准方法进行编译,并对文本生成质量进行评估;

上表总结了第2~5个内核「开箱即用」时的表现。

这些结果表明,如果目标是构建一个端到端的生产级内核,那么拥有一个能跑基准测试的内核还远远不够。

后续测试中使用AMD flash attention内核,因为它可以通过torch.compile进行编译,且在eager和compile模式下都有清晰的输出。

为了满足torch.compile与AMD flash attention内核的兼容性,我们需要自定义torch运算符,主要包括以下两步:

1. 将函数包装到PyTorch自定义运算符中

2. 在运算符中添加一个FakeTensor Kernel,给定flash输入张量的形状(q、k 和 v),它可以提供一种计算flash内核输出形状的方法

将模型中的运算换为Triton的自定义内核后,就能成功地进行编译和运行,Nsight跟踪信息如下图所示:

对比图5可以发现,图6就是100%使用Triton内核的前向计算。

基准测试

基准测试中使用Granite-8B和Llama3-8B模型,在英伟达H100和A100上进行单GPU运行,并定义了两种不同的配置:

Triton内核配置使用:

1. Triton SplitK GEMM

2. AMD Triton Flash Attention

CUDA 内核配置使用:

1. cuBLAS GEMM

2. cuDNN Flash Attention - 缩放点积注意力 (SDPA)

在典型的推理设置下,eager和torch编译模式的吞吐量和token间延迟如下:

批大小=2,输入序列长度=512,输出序列长度=25

Triton模型在H100上的性能最高可达CUDA模型的78%,在A100上的性能最高可达82%。两者间性能的差距可能源于矩阵乘法和flash attention的内核延迟,下一节将详细讨论。

微基准测试

解码延迟时间对比,输入是任意提示,批大小=1,提示长度=44

将端到端推理中的各部分进行单独对比,我们注意到以下两点:

1. Triton的matmul内核比CUDA慢1.2~1.4倍

2. AMD的Triton Flash Attention内核比CUDA SDPA慢1.6倍

这些结果表明,需要进一步提升GEMM和Flash Attention等关键原语的内核性能。

比如最近提出的FlashAttention-3、FlexAttention等工作提供了更好的方法来利用底层硬件,有希望在此基础上为Triton进一步加速。

将 FlexAttention与SDPA和AMD 的 Triton Flash内核进行比较,微基准测试结果显示,Flex有望被用于上下文更长、解码规模更大的问题场景。

英伟达H100 SXM5 80GB上的FlexAttention内核基准测试

未来展望

接下来,我们期望进一步优化矩阵乘法(matmuls),以更充分地利用硬件。

比如使用不同的工作分解方法(类似StreamK的持久内核技术),以加快基于Triton的方法。

我们还期望继续探索FlexAttention和FlashAttention-3,进一步缩小Triton和CUDA间的差距。

以上的实验只针对FP16精度,但早前的研究表明,与cuBLAS FP8 GEMM相比,FP8 Triton GEMM内核表现更好。因此接下来的工作还会探讨端到端FP8 LLM推理。

参考资料:

https://pytorch.org/blog/cuda-free-inference-for-llms/?utm_content=306418723&utm_medium=social&utm_source=twitter&hss_channel=tw-776585502606721024

全球首发第二代东方屏 最好的2K屏旗舰来了!一加13十月亮相
DeepSeek分享沉淀多年的高性能深度学习架构 用60%成