PyTorch官宣：告别CUDA，GPU推理迎来Triton加速新时代 | 最新快讯

　　新智元报道

　　编辑：乔杨 Frey

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

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

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

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

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

　　原文地址：CUDA-Free Inference for LLMs | PyTorch

　　那么，让我们紧跟 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 内核有一系列不同的配置和实现，包括：

　　- AMD Flash

　　- 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 推理。

　　参考资料：

　　CUDA-Free Inference for LLMs | PyTorch

来自: 网易科技