摘要
利用第一性原理方法，我们展示了加速当前 Triton GPTQ 内核（核心 GPTQ 提速 3 倍，AutoGPTQ 提速 6 倍）所采取的循序渐进的过程。例如：在典型的 Llama 风格推理输入上，时间从 275 微秒缩短到 47 微秒。目标是为加速任何 Triton 内核提供有用的模板。我们提供了 Triton 和 GPTQ 量化与反量化过程的背景知识，展示了合并内存访问对提高共享内存和全局内存吞吐量的影响，重点介绍了为减少 Warp 停滞以提高总吞吐量所做的更改，并概述了如何将 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 是一种量化算法，能够通过近似二阶信息（Hessian 逆）有效地将超大 (175B+) LLM 压缩到 int4 位表示。 AutoGPTQ 是一个基于 GPTQ 构建的框架，允许对已使用 GPTQ 量化的 LLM 进行快速反量化和推理/服务。
作为 AutoGPTQ 技术栈的一部分，他们提供了一个 Triton GPTQ 内核来处理模型推理的反量化过程。
INT 量化的基本过程如下所示，包括确定 scale 和 zero point，然后使用 scale 和 zero point 计算量化的 4 位权重
因此，我们为每组权重存储 4 位权重以及 Scale 和 ZeroPoint 的元信息。
要对这些权重进行“反量化”，我们执行以下操作
然后继续将反量化后的权重与此线性层的密集输入特征矩阵进行矩阵乘法。
2.0 识别瓶颈 - 优化矩阵乘法
事实证明，构建快速矩阵乘法内核并非易事。一个朴素实现的矩阵乘法在 GPU 这种高度并行机器上很少能达到峰值吞吐量性能。因此，我们需要分层处理 GPU 中的计算和内存子系统，以确保我们最大程度地利用每种资源。
我们通过 Nvidia Nsight Compute 工具运行未优化的 Triton 内核，记录一些重要的指标和警告，从而开始我们的优化过程。
图 xy (待办)
我们首先注意到计算吞吐量和内存吞吐量都很低，分别为 7.40% 和 21.19%（图 xy）。考虑到对于典型的推理矩阵问题规模，我们处于内存限制区域 (memory bound regime)，我们将尝试通过应用针对 A100 GPU 内存子系统的代码更改来优化内核。
本文将涵盖的三个主题是
- L2 优化
- 向量化加载
- Warp 停滞
让我们逐个主题深入，进行适当的更改，并查看其对我们的 Triton 内核的相应影响。这个 Triton 内核是一个融合的反量化内核，它将打包的 int32 权重（我们将其称为 B 矩阵）张量反量化为 int4 权重，在 FP16 模式下与激活张量（称为 A 矩阵）执行矩阵乘法，然后将结果存储回矩阵 C。
以上称为 W4A16 量化。请记住，我们描述的过程可以而且应该用于任何 GPU 内核的开发，因为这些是任何未优化内核中的常见瓶颈。
3.0 L2 优化
此优化已存在于 AutoGPTQ 内核中，但我们想专门用一节来帮助读者更好地理解 Triton 中线程块的映射和执行顺序是如何处理的。因此，我们将逐步介绍一种朴素映射，然后再介绍一种更优化的映射，以查看其相应影响。
让我们朴素地构建我们的内核，从全局内存的“线性”加载开始，然后将其与更优化的“交错式”（swizzled）加载进行比较。线性与交错式加载决定了我们的工作网格在 GPU 上的执行顺序。让我们看一下 Nvidia Nsight Compute 工具在朴素情况下就我们的内核共享内存访问模式提供的提示
为了解决这个问题，我们可以使用一种称为“tile-swizzling”（瓦片交错）的方法。该方法的思想是以对 L2 缓存更友好的顺序启动我们的线程块。
让我们退一步，熟悉一下 Triton 的一些语义，并做一个简单的 CUDA 类比来更好地理解这个概念。Triton 内核会启动“程序”。这些所谓的程序对应于 CUDA 中的线程块 (Thread Block) 的概念，它是 Triton 内核中并行的基本单位。每个程序都关联一个“pid”，程序中的所有线程都保证执行相同的指令。
如果您将“pid”简单地线性映射到输出矩阵 C 的 2D 网格位置，Triton 程序将以朴素的方式分布到您的 SM 上。
这个 2D 网格位置由 Triton 中的 pid_m 和 pid_n 确定。当我们在 GPU 上分配工作网格时，我们希望利用 L2 缓存中的数据和缓存局部性。在 Triton 中，我们可以进行以下更改来实现这一点
红色突出显示的代码是朴素的“线性”瓦片排序，绿色突出显示的代码是“交错式”瓦片排序。这种启动方式能提升局部性。这里有一个可视化图例帮助更好地理解。
整合此更改后，分析器不再抱怨非合并的内存访问。让我们看看我们的内存吞吐量是如何变化的
此更改在一个简单的加载存储内核上进行了测试。查看分析器中 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 停滞
现在将所有更改放回我们的完整反量化内核中，我们看到了以下性能限制因素：Warp 停滞。
这些 Warp 停滞主要由“长记分牌”（Long Scoreboard）停滞引起，占总数的 92.63%。
从高层次来看，长记分牌停滞当一个 warp 需要尚未准备好的数据才能进入“已发出”（issued）状态时，就会发生长记分牌停滞。换句话说，GPU 是吞吐量机器，我们需要用计算指令来隐藏加载指令的延迟。通过加载更多数据并重新安排脚本中加载指令的位置，我们可以解决这个问题。
在理想情况下，每个 warp 调度器每时钟周期都能发出 1 条指令。注意 - A100 GPU 上的每个 SM 都有 4 个 warp 调度器。
然而，我们的内核存在瓶颈，并且在 AutoGPTQ Triton 内核根据其预设认为最优的块大小下，在停滞状态中花费了 4.4 个时钟周期。
如何改进这一点？
我们希望能够提高内存吞吐量，以便增加当 warp 发出指令时，我们不必等待加载的数据存储到 SRAM 中以供计算使用的机会。我们尝试了多个参数（例如流水线阶段数和 warp 数），其中影响最大的是将 k 维度上的块大小增加 2 倍。
这些更改立即对计算和内存吞吐量产生了影响。
我们还看到在对量化权重进行移位和缩放的步骤中，长记分牌等待时间显著下降，这是我们最初在源代码中识别出的瓶颈。尽管在此行仍有停滞，但只有 68% 是由长记分牌停滞引起，而原始比例为 92%。理想情况下，我们不会观察到任何停滞，所以此处仍有待改进，但长记分牌停滞数量的减少表明，此时我们的数据已准备好被 warp 希望执行的指令在 L1TEX 内存中使用，频率高于原始内核。
相应的效果是我们的内核执行时间加快了 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 编译器编译
除了实际的内核代码外，还有一个 Python 包装器，它可能会或可能不会继承 PyTorch autograd 类 - 这取决于它是否支持反向传播（即用于训练目的还是仅用于推理目的）。
您只需将该 Python 类导入到您的 PyTorch 代码中您想要使用它的地方，就像使用任何其他 Python / PyTorch 函数一样。
在本例中，简单地导入并使用‘fast_qlinear’将调用底层 Triton 内核，并将我们上面展示的加速应用于您的 PyTorch 模型。
致谢
感谢来自 IBM Research 的 Jamie Yang 和 Hao Yu 在收集这些结果过程中提供的技术指导。
