摘要（TL;DR）
通过采用第一性原理方法，我们展示了将现有 Triton GPTQ 内核加速 3 倍（核心 GPTQ）和 6 倍（AutoGPTQ）的详细过程。示例：在典型的 Llama 类推理输入上，运行时间从 275us 缩短至 47us。我们的目标是为加速任何给定的 Triton 内核提供一个有用的模板。我们提供了关于 Triton 和 GPTQ 量化与反量化过程的背景信息，展示了合并内存访问（coalesced memory access）对改善共享内存和全局内存吞吐量的影响，重点介绍了为减少线程束停顿（warp stalling）以提高总吞吐量所做的更改，并概述了如何将 Triton 内核集成到 PyTorch 代码中。从长远来看，我们希望通过我们的 Triton 内核超越现有的 CUDA 原生 GPTQ 内核。
图 1：在 H100 上对优化后的 AutoGPTQ 内核与现有 AutoGPTQ 内核进行性能基准测试
图 2：在 A100 上对新优化后的 AutoGPTQ 内核与现有 AutoGPTQ 内核进行性能基准测试
图 3：即使有了这些改进，我们的优化 Triton 内核与 A100 上的 CUDA 原生 AutoGPTQ 内核之间仍存在差距。敬请期待后续进展……
1.0 Triton 简介
Triton 框架提供了一种硬件无关的 GPU 编程和定位方式，目前同时支持 NVIDIA 和 AMD，对其他硬件供应商的支持也在进行中。Triton 现已成为 PyTorch 2.0 的核心支柱，因为 torch.compile 会将渴望模式（eager PyTorch）代码分解，并将其重组为包含高比例 Triton 内核以及 PyTorch 连接代码的形式。
随着 Triton 得到更广泛的采用，程序员必须了解如何系统地深入分析 Triton 堆栈（从高级 Python 到低级 SASS），以解决性能瓶颈，从而为那些超出 torch.compile 生成内核范围的算法优化 GPU 效率。
在本文中，我们将介绍 Triton 编程语言的一些核心概念，如何识别 GPU 内核中的常见性能限制因素，并同步优化 AutoGPTQ 中使用的高吞吐量推理量化内核。
GPTQ 量化与反量化简介
GPTQ 是一种量化算法，能够通过近似二阶信息（海森矩阵的逆）有效地将超大型（175B+）大语言模型压缩为 int4 位表示。AutoGPTQ 是一个建立在 GPTQ 之上的框架，允许对已通过 GPTQ 量化的大语言模型进行快速反量化和推理/服务。
作为 AutoGPTQ 堆栈的一部分，他们提供了一个 Triton GPTQ 内核来处理模型推理时的反量化。
INT 量化的基本过程如下所示，包括确定缩放因子（scale）和零点（zero point），然后利用缩放因子和零点计算出量化后的 4bit 权重。
因此，我们为每组权重存储 4 Bit 权重以及缩放因子和零点的元信息。
要对这些权重进行“反量化”，我们执行以下操作：
然后继续将反量化后的权重与此线性层的稠密输入特征矩阵进行矩阵乘法运算。
2.0 识别瓶颈 – 优化矩阵乘法
事实证明，编写一个快速的矩阵乘法内核并非易事。在像 GPU 这样高度并行的机器上，天真（naively）实现的矩阵乘法几乎不可能达到峰值吞吐量。因此，我们需要以分层方式处理 GPU 中的计算和内存子系统，以确保最大限度地利用每项资源。
我们通过 Nvidia Nsight Compute 工具运行未优化的 Triton 内核，并记录一些重要的指标和警告，以此开始我们的优化过程。
图 xy（待定）
我们首先注意到计算吞吐量和内存吞吐量都很低，分别为 7.40% 和 21.19%（图 xy）。已知对于典型的推理矩阵问题规模，我们处于内存受限（memory bound）状态，因此我们将尝试通过应用针对 A100 GPU 内存子系统的代码更改来优化内核。
本文将涵盖以下三个主题：
- L2 缓存优化
- 向量化加载
- 线程束停顿（Warp Stalling）
让我们逐一讲解每个主题，做出相应的更改，并观察它们对我们的 Triton 内核产生的具体影响。此 Triton 内核是一个融合反量化内核，它将打包的 int32 权重（我们称之为 B 矩阵）张量反量化为 int4 权重，在 FP16 模式下与激活张量（称之为 A 矩阵）进行矩阵乘法运算，然后将结果存储回矩阵 C。
上述过程被称为 W4A16 量化。请记住，我们描述的过程可以也应该用于开发任何 GPU 内核，因为这些是任何未优化内核中的常见瓶颈。
3.0 L2 缓存优化
此项优化已存在于 AutoGPTQ 内核中，但我们希望专门用一节来介绍它，以帮助读者更好地理解 Triton 中如何处理线程块（thread blocks）的映射和执行顺序。因此，我们将逐步介绍天真的映射以及更优的映射，以观察其相应的影响。
让我们从全局内存的“线性”加载开始，天真地构建我们的内核，然后将其与更优化的“混洗”（swizzled）加载进行比较。线性 vs 混洗决定了 GPU 上工作网格的执行顺序。让我们看看 Nvidia Nsight Compute 工具针对我们在天真情况下共享内存访问模式提供的提示。
为了解决这个问题，我们可以使用一种称为“瓦片混洗”（tile-swizzling）的方法。该方法的核心思想是以一种更利于 L2 缓存的方式启动线程块。
让我们退后一步，熟悉一些 Triton 语义，并做一个简单的 CUDA 类比来更好地理解这个概念。Triton 内核启动“程序”。这些所谓的程序映射到 CUDA 中的线程块概念，它是 Triton 内核中并行性的基本单位。每个程序都有一个关联的“pid”，并且保证程序中的所有线程都在执行相同的指令。
如果你对“pid”到输出矩阵 C 的 2D 网格位置进行简单的线性映射，Triton 程序将以天真的方式分配到你的流多处理器（SM）上。
这个 2D 网格位置由 Triton 中的 pid_m 和 pid_n 决定。我们希望在分配工作网格时利用 GPU L2 缓存中的数据和缓存局部性。为此，我们可以在 Triton 中进行以下更改：
红色高亮显示的代码是天真的“线性”瓦片排序，绿色高亮显示的代码是“混洗”瓦片排序。这种启动方式促进了局部性。以下是帮助更好理解的示意图。
加入此更改后，分析器不再报告未合并（uncoalesced）的内存访问问题。让我们看看内存吞吐量发生了怎样的变化。
此更改在一个简单的加载/存储内核上进行了测试。查看分析器中的“GPU 光速统计”部分，我们还看到简单加载内核的内存吞吐量增加了 112.07%，这正是我们通过此项优化想要达到的目标。同样，此优化已存在于 AutoGPTQ 内核中，但这是每个 Triton 内核程序员在内核编写初期、在处理任何激动人心的反量化或矩阵乘法逻辑之前都必须编写的样板逻辑。因此，理解以下几点非常重要：
- 这种映射不是唯一的。
- Triton 不会自动为程序员处理此类优化，必须深思熟虑以确保你的内核能够最佳地处理共享内存访问。
对于 Triton 新手来说，这些并不显眼，因为大部分共享内存访问优化是由 Triton 编译器处理的。然而，在编译器无法处理的情况下，能够了解有哪些可用工具和方法来影响内存行为非常重要。
4.0 向量化加载
现在，回到我们未优化内核最初的问题。我们想要优化内核的全局内存访问模式。从 Nvidia Nsight Compute 工具的详细信息页面中，我们看到以下提示，分析器抱怨存在未合并的全局内存访问。
让我们深入研究并查看未优化内存读取的 SASS（汇编）代码加载部分。
此加载操作导致了 32 次 16 位宽的全局加载操作。这不是最优的。
我们希望以向量化的方式执行全局内存加载，使其产生的加载指令数量最少。为了解决这个问题，我们可以给 Triton 编译器提供一些帮助。
上面绿色高亮的行充当了编译器提示。它告诉编译器这些元素在内存中是连续的，并且此加载操作可以合并。
让我们看看添加这些行后在汇编代码层面的效果。
加载现在以 4 次全局加载操作执行，每次 128 位宽，而不是 32 次 16 位全局加载操作。这意味着减少了 28 条内存获取指令，重要的是实现了合并内存访问。这可以从单个线程不再访问连续内存地址这一事实看出，而没有编译器提示时，行为正是如此。
最终效果是在隔离加载操作中实现了 73 倍的加速，而在将其合并到完整的反量化内核后，我们又看到了 6% 的加速。这是朝着正确方向迈出的又一步！
5.0 线程束停顿（Warp Stalling）
现在将所有更改放回完整的反量化内核中，我们看到了以下性能限制因素：线程束停顿。
这些线程束停顿主要由“长记分板”（Long Scoreboard）停顿引起，占总数的 92.63%。
从高层来看，长记分板停顿发生在一个线程束（warp）需要的数据尚未准备好，因此无法进入“已发布”（issued）状态时。换句话说，GPU 是吞吐量机器，我们需要用计算指令隐藏加载指令的延迟。通过加载更多数据并重新安排脚本中加载指令的位置，我们可以解决这个问题。
在理想情况下，每个线程束调度器应该能够在每个时钟周期发布 1 条指令。注意：A100 GPU 上的每个 SM 都有 4 个线程束调度器。
然而，我们的内核存在瓶颈，且在 AutoGPTQ Triton 内核根据其预设认为最优的块大小下，有 4.4 个周期处于停顿状态。
我们如何改进这一点？
我们希望能够提高内存吞吐量，从而增加线程束发布指令时，无需等待加载数据存储到 SRAM 即可参与计算的几率。我们尝试了多个参数（如流水线阶段数和线程束数量），其中影响最大的是将 k 维度的块大小增加了 2 倍。
这些更改对计算和内存吞吐量产生了直接影响。
我们还看到，在执行量化权重移位和缩放的步骤中，长记分板等待时间显著下降，这正是我们最初在源代码中确定的瓶颈。虽然该行代码处仍存在停顿，但只有 68% 是由长记分板停顿引起的，相比最初的 92% 有所改善。理想情况下，我们不希望观察到任何停顿，所以这里仍有工作要做，但长记分板引起的停顿减少告诉我们，我们的数据现在能以比原始内核更高的频率，在线程束想要执行指令时准备好（在 L1TEX 内存中）。
相应的效果是我们内核的执行时间提升了 1.4 倍。
6.0 结果
通过有条不紊地解决所有这些问题领域，我们最终得到的内核在 Nvidia A100 GPU 上比使用 AutoGPTQ 开箱即用的 Triton 内核快 6 倍。
以相关的 Llama 推理样本数据点为例，我们开发的 Triton 内核执行反量化和矩阵乘法需要 47us，而 AutoGPTQ 内核在相同矩阵大小下需要 275us。
通过复制这种分步方法，应该可以在其他内核中获得类似的加速，并有助于加深对常见 GPU 瓶颈及其解决方法的理解。
需要注意的是，尽管 AutoGPTQ Triton 内核的性能有所提升，但我们尚未缩小与 AutoGPTQ 中现有的 exllamaV2 CUDA 内核之间的差距。
需要进行更多研究来了解我们如何进一步优化此内核，以达到等效的自定义 CUDA 内核性能。
总结与未来工作
Triton 通过允许在比 CUDA 编程更高的抽象级别上进行底层 GPU 优化来扩展 PyTorch，最终结果是添加优化的 Triton 内核可以帮助 PyTorch 模型运行得更快。
我们在本文中的目标是展示一个加速 GPTQ 反量化内核的示例，并提供实现这些加速的模板工作流。
对于未来的工作，矩阵乘法的 SplitK 工作分解是我们打算研究的一个潜在加速方向。
将自定义 Triton 内核集成到 PyTorch 中
鉴于上述加速效果，一个常见的问题是如何在给定的 PyTorch 代码库中使用自定义内核。
Triton 内核将至少包含两部分——由 Triton 编译器编译的实际 Triton 内核代码。
除了实际的内核代码外，还有一个 Python 包装器，它可能会也可能不会继承 PyTorch 的 autograd 类——具体取决于它是否需要支持反向传播（例如用于训练目的，还是仅用于推理目的）。
你只需像使用任何其他 Python / PyTorch 函数一样，将该 Python 类导入到想要使用它的 PyTorch 代码中即可。
在这种情况下，只需导入并使用 ‘fast_qlinear’，即可调用底层 Triton 内核，并将上述加速效果应用到你的 PyTorch 模型中。
致谢
感谢 IBM Research 的 Jamie Yang 和 Hao Yu 在收集这些结果时提供的技术指导。