由 Less Wright，Adnan Hoque（IBM）著

TL;DR

通过采用第一性原理方法，我们展示了加速当前 Triton GPTQ 内核（核心 GPTQ）3 倍和 6 倍（AutoGPTQ）的逐步过程。例如：在典型的 Llama 风格推理输入上，从 275us 减少到 47us。目标是提供一个加速任何给定 Triton 内核的有用模板。我们提供了关于 Triton 和 GPTQ 量化和反量化的背景介绍，展示了合并内存访问对提高共享和全局内存吞吐量的影响，突出了为提高总吞吐量而进行的减少 warp 停滞的更改，并概述了将 Triton 内核集成到 PyTorch 代码中的方法。从长远来看，我们希望我们的 Triton 内核能够超越现有的 CUDA 原生 GPTQ 内核。

Fig 1: Performance benchmarking the optimized AutoGTPQ kernel vs the current AutoGPTQ kernel on H100

图 1：在 H100 上对优化后的 AutoGTPQ 内核与当前 AutoGPTQ 内核的性能基准测试

Fig 2: Performance benchmarking the newly optimized AutoGTPQ kernel vs the current AutoGPTQ kernel on A100

图 2：在 A100 上对新优化后的 AutoGTPQ 内核与当前 AutoGPTQ 内核的性能基准测试

Fig 3: Even with these improvements, there remains a gap between our optimized Triton kernel and the CUDA native AutoGTPQ kernel on A100.

即使有这些改进，我们的优化 Triton 内核与 A100 上的 CUDA 原生 AutoGTPQ 内核之间仍然存在差距。更多内容即将揭晓……

1.0 Triton 简介

Triton 框架提供了一种硬件无关的编程和针对 GPU 的方法，目前支持 NVIDIA 和 AMD，正在推进对其他硬件供应商的支持。Triton 现在是 PyTorch 2.0 的支柱，因为 torch.compile 将急切 PyTorch 分解并重新组装成高比例的 Triton 内核，PyTorch 连接代码。

随着 Triton 的广泛应用，程序员了解如何系统地遍历 Triton 堆栈（从高级 Python 到低级 SASS）以解决性能瓶颈，对于优化超越 torch.compile 生成的内核的算法的 GPU 效率至关重要。

在本文中，我们将介绍 Triton 编程语言的一些核心概念，如何识别 GPU 内核中的常见性能限制因素，并在并行中调整 AutoGPTQ 中使用的量化内核，该内核可用于高吞吐量推理应用。

GPTQ 量化和去量化简介

GPTQ 是一种量化算法，能够通过近似二阶信息（Hessian 逆）将超大型（175B+）LLMs有效地压缩到 int4 位表示，AutoGPTQ 是基于 GPTQ 构建的框架，允许快速去量化以及量化LLMs的推理/服务。

作为 AutoGPTQ 堆栈的一部分，他们提供了一个 Triton GPTQ 内核来处理模型推理的去量化。

INT 量化的基本过程如下，涉及确定缩放和零点，然后使用缩放和零点计算量化的 4 位权重：

The basic process for INT quantization

因此，我们将 4 位权重以及每个权重组的缩放和零点元信息一起存储。

要对这些权重进行“反量化”，我们执行以下操作：

To ‘dequant’ these weights

然后继续用这个线性层的密集输入特征矩阵与反量化后的权重进行矩阵乘法。

2.0 识别瓶颈 - 优化矩阵乘法

事实上，编写一个快速的矩阵乘法内核并非易事。一个简单实现的矩阵乘法很少能在高度并行的机器（如 GPU）上达到峰值吞吐量性能。因此——我们需要以分层的方式处理我们的 GPU 中的计算和内存子系统，以确保我们最大限度地利用每种资源。

我们通过运行未经优化的 Triton 内核，并通过 Nvidia Nsight Compute 工具，记录一些重要的指标和警告来开始我们的优化过程：

some important metrics and warnings

图 xy（待补充）

some important metrics and warnings

我们首先注意到计算和内存吞吐量都很低，分别为 7.40%和 21.19%（见图 xy）。鉴于对于典型的推理矩阵问题规模，我们处于内存限制状态，我们将尝试通过针对我们 A100 GPU 内存子系统的代码更改来优化内核。

本篇文章将涵盖三个主题：

  1. L2 优化
  2. 向量化加载
  3. 螺旋停转

让我们逐一分析每个主题，进行适当的修改，并观察其对我们的 Triton 内核的相应影响。这个 Triton 内核是一个融合的解量化内核，它将打包的 int32 权重张量（我们将称之为 B 矩阵）解量化为 int4 权重，以 FP16 模式与激活张量（称为 A 矩阵）进行矩阵乘法，然后将结果存储回矩阵 C。

上述内容被称为 W4A16 量化。请注意，我们描述的过程可以也应该用于任何 GPU 内核的开发，因为这些是任何未优化内核中的常见瓶颈。

3.0 L2 优化

这种优化已经在 AutoGPTQ 内核中存在，但我们想为此专门划分一个章节，帮助读者更好地理解如何在 Triton 中处理线程块的映射和执行顺序。因此，我们将逐步介绍一种简单的映射方法，然后是一种更优化的映射方法，以观察其对应的影响。

让我们以“线性”方式从全局内存加载我们的内核，然后将其与更优化的“交错”加载进行比较。线性与交错决定了我们在 GPU 上工作网格的执行顺序。让我们看看 Nvidia Nsight Compute 工具在简单情况下关于我们的内核共享内存访问模式的提示：

the hints from the Nvidia Nsight Compute Tool

为了解决这个问题，我们可以使用一种称为“瓦片交错”的方法。这种方法的想法是以更友好的 L2 缓存顺序启动我们的线程块。

让我们退一步，熟悉一下 Triton 的语义，并通过一个简单的 CUDA 类比来更好地理解这个概念。Triton 内核启动“程序”。这些所谓的程序对应于 CUDA 中的线程块的概念，它是 Triton 内核中的基本并行单元。每个程序都与其关联一个“pid”，并且程序中的所有线程都保证执行相同的指令。

如果您对输出矩阵 C 的二维网格位置进行简单的“pid”线性映射，Triton 程序将以一种直观的方式分布到您的 SM 上。

这个二维网格位置由 Triton 中的 pid_m 和 pid_n 确定。当我们分配我们的工作网格时，我们希望利用 GPU 的 L2 缓存中的数据和缓存局部性。为了在 Triton 中实现这一点，我们可以进行以下更改：

To do this in Triton

红色高亮的代码是直观的“线性”瓦片排序，绿色高亮的代码是“交错”瓦片排序。这种类型的启动促进了局部性的感觉。以下是一个帮助理解的视觉图。

a sense of locality

在合并此更改后，分析器不再抱怨未合并的内存访问。让我们看看我们的内存吞吐量是如何变化的：

how our memory throughput has changed

此更改已在简单的加载存储内核上进行测试。查看分析器的 GPU 光速统计部分，我们还可以看到简单加载内核的内存吞吐量提高了 112.07%，这正是我们通过此优化所追求的。再次强调，这种优化已经在 AutoGPTQ 内核中存在，但这是每个 Triton 内核程序员在内核开始编写任何令人兴奋的量化和矩阵乘法逻辑之前必须编写的样板逻辑。因此，重要的是要理解：

  1. 这种映射不是唯一的

  2. Triton 不会自动为程序员处理此类优化，因此必须仔细思考以确保您的内核能够最优地处理共享内存访问

对于初识 Triton 的人来说，这些可能并不明显，因为大部分的共享内存访问优化都是由 Triton 编译器处理的。然而，在这些情况下，如果编译器没有处理，那么理解我们有哪些工具和方法可以影响内存行为就变得非常重要。

4.0 向量化加载

现在，回到我们未优化的内核的原始抱怨。我们希望优化内核的全局内存访问模式。从 Nvidia Nsight Compute 工具的详细信息页面，我们看到以下注释，其中分析器抱怨未归一化的全局内存访问。

让我们深入挖掘，看看未优化的内存读取的 SASS（汇编）代码加载：

an unoptimized memory read

此次加载操作导致产生了 32 次全局加载操作，这些操作宽度为 16 位。这并不理想。

我们希望以矢量化方式执行全局内存加载，以使加载指令数量最少。为了解决这个问题，我们可以为 Triton 编译器提供一些帮助。

code block

上方的绿色高亮行充当编译器提示。它告诉编译器这些元素在内存中是连续的，并且这个加载操作可以被合并。

让我们看看添加这些行后汇编代码的效果。

the effect in assembly after adding these lines

现在负载操作由 4 个 128 位宽的全局负载操作组成，而不是 32 个 32 位宽的全局负载操作。这意味着减少了 28 条内存读取指令，并且更重要的是，实现了内存访问的合并。这可以从单个线程不再访问连续内存地址这一事实中看出，这是在没有编译器提示时的行为。

结果是孤立负载操作的速度提高了 73 倍，在将其纳入完整的去量化内核后，我们还看到了 6%的速度提升。这是朝着正确方向迈出的又一步！

5.0 线程阻塞

performance limiter, warp stalling

现在将所有更改放回我们的完整去量化内核中，我们看到了以下性能限制：线程阻塞。

这些 warp 停滞主要是由“长计分板”引起的，占总数的 92.63%。

从高层次来看，当 warp 需要的数据尚未准备好以进入“已发布”状态时，就会发生长计分板停滞。换句话说，GPU 是吞吐量机器，我们需要通过计算指令来隐藏加载指令的延迟。通过加载更多数据并重新排列脚本中的加载指令，我们可以解决这个问题。

在理想情况下，每个 warp 调度器在每个时钟周期都能发布 1 条指令。注意 - A100 GPU 上的每个 SM 都有 4 个 warp 调度器。

然而——我们的内核存在瓶颈，在 AutoGPTQ Triton 内核认为最优的块大小下，内核在停滞状态下花费了 4.4 个周期。

我们如何改进这一点？

我们希望提高内存吞吐量，以便在 warp 发出指令时，我们不会等待将数据存储到 SRAM 中，以便用于计算。我们尝试了多个参数（例如流水线阶段数量和 warp 数量），其中影响最大的是将 k 维度的块大小增加 2 倍。

这些更改对计算和内存吞吐量都有即时的提升。

an immediate impact on both compute and memory throughput

我们还看到在移位和缩放量化权重时的长计分板等待时间显著下降，这是我们识别出的源代码中的原始瓶颈。尽管在这个地方仍然存在停滞，但只有 68%的停滞是由长计分板停滞引起的，而最初是 92%。理想情况下，我们不应该观察到任何停滞，所以这里还有工作要做，但长计分板引起的停滞数量减少告诉我们，我们的数据此时已经准备好被 warp 想要执行的指令以更高的频率使用（在 L1TEX 内存中），比原始内核要高。

1.4x speedup in the execution time of our kernel

我们内核的执行时间相应提高了 1.4 倍。

6.0 结果

通过有系统地解决所有这些问题区域，我们得到的内核在 Nvidia A100 GPU 上的速度比使用 AutoGPTQ 提供的 Triton 内核快 6 倍。

以一个相关的 Llama 推理样本数据点为例，我们开发的 Triton 内核进行去量化矩阵乘法需要 47 微秒，而 AutoGPTQ 内核进行相同矩阵大小的操作需要 275 微秒。

通过复制这种逐步的方法，应该可以在其他内核中实现类似的速度提升，并有助于建立对常见 GPU 瓶颈及其解决方法的了解。

需要注意的是，尽管 AutoGPTQ Triton 内核的性能已经有所提升，但我们仍未缩小与 AutoGPTQ 中现有的 exllamaV2 CUDA 内核之间的差距。

需要更多研究来了解我们如何进一步优化这个内核，以匹配等效的定制 CUDA 内核性能。

摘要及未来工作

Triton 通过允许在比 CUDA 编程更高的抽象级别进行底层 GPU 优化，从而扩展了 PyTorch，使得添加优化的 Triton 内核可以帮助 PyTorch 模型运行得更快。

本文的目标是展示加速 GPTQ 去量化内核的示例，并提供如何实现加速的模板工作流程。

对于未来的工作，我们将研究矩阵乘法中 SplitK 工作分解的潜在加速方法。

将自定义 Triton 内核集成到 PyTorch 中

根据上述加速度，一个常见的问题是如何在给定的 PyTorch 代码库中实际使用自定义内核。

Triton 内核将包含至少两个部分——实际的 Triton 内核代码，该代码将由 Triton 编译器编译：

the actual Triton kernel code which will be compiled by the Triton compiler

除了实际的内核代码外，还有一个 Python 包装器，它可能或可能不继承 PyTorch 的 autograd 类——这取决于它是否将支持反向传播（即用于训练目的或仅用于推理目的）。

您只需将 Python 类导入到您想要使用它的 PyTorch 代码中，就像导入任何其他 Python / PyTorch 函数一样。

import the python class into your PyTorch code

在这种情况下，只需导入并使用‘fast_qlinear’，就会调用底层 Triton 内核，并应用我们上面展示的速度提升来加速您的 PyTorch 模型。

致谢

感谢 IBM Research 的 Jamie Yang 和 Hao Yu 在收集这些结果中的技术指导。