blog地址：https://pytorch.org/blog/accelerating-triton/ triton kernel 地址：https://github.com/foundation-model-stack/foundation-model-stack/blob/triton/triton/kernels/gptq/splitk_dequant_gemm.py#L51

加速 GPTQ 的 Triton Dequantization Kernel

太长不看版

利用第一性原理方法，我们展示了一个逐步过程，旨在将当前的Triton GPTQ kernel加速3倍（core GPTQ）和6倍（AutoGPTQ）。例如：在典型的Llama风格推理输入上，将处理时间从275微秒降低到47微秒。我们的目标是提供一个有用的模板，用于加速任何给定的Triton kernel。我们提供了Triton和GPTQ量化及反量化过程的背景信息，展示了合并内存访问对改善共享和全局内存吞吐量的影响，强调了为减少线程束停滞以提高总体吞吐量所做的更改，并概述了将Triton kernel集成到PyTorch代码中的方法。长期来看，我们希望我们的Triton kernel能够超越现有的CUDA原生GPTQ kernel。

1.0 对Triton的介绍

Triton框架提供了一种硬件无关的方式对GPU编程，目前支持NVIDIA和AMD，并正在进行对其他硬件供应商的支持。Triton现已成为PyTorch 2.0的主要组成部分，torch.compile将eager PyTorch代码分解并重新组装成大比例的Triton kernel，并配以PyTorch连接代码。

随着Triton被更广泛地采用，程序员需要了解如何系统地遍历Triton堆栈（从高层Python到底层SASS），以解决性能瓶颈，达到比torch.compile生成的Triton kernel更快。

在这篇文章中，我们将介绍Triton编程语言的一些核心概念，如何识别GPU kernel中常见的性能限制因素，并同时调优一个用于AutoGPTQ的量化kernel，该kernel可用于高吞吐量推理应用。

GPTQ量化和反量化简介

GPTQ（https://arxiv.org/abs/2210.17323） 是一种量化算法，能够通过近似二阶信息（Hessian逆矩阵）有效地将超大型（175B+）LLM压缩为4位整数表示。AutoGPTQ (https://github.com/AutoGPTQ/AutoGPTQ) 是一个建立在GPTQ基础上的框架，允许快速反量化并推理或者服务使用GPTQ量化的LLM。

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

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

我们因此存储4位权重以及每组权重的比例和零点元信息。

要"反量化"这些权重，我们执行以下操作：

然后进行 矩阵乘法 ，将反量化的权重与该线性层的dense输入特征矩阵相乘。

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

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

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

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

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

1. L2 优化
2. 向量化加载
3. 线程束停滞

让我们逐一讨论每个主题，进行适当的更改，并观察其对我们Triton kernel的相应影响。这个Triton kernel是一个fuse反量化kernel，它将Packed的int32权重（我们将其称为B矩阵，一个INT32权重对应了8个INT4的权重）反量化为FP16数据类型，以FP16模式与激活张量（称为A矩阵）进行矩阵乘法，然后将结果存储回矩阵C。

我们在 https://github.com/foundation-model-stack/foundation-model-stack/blob/triton/triton/kernels/gptq/splitk_dequant_gemm.py#L51 这里看到的 // 8就是把 INT32 权重对应到INT4权重上。

上述过程被称为W4A16量化。请记住，我们描述的过程可以且应该用于开发任何GPU kernel，因为这些是任何未优化kernel中常见的瓶颈。

3.0 L2 优化

这种优化已经存在于AutoGPTQ kernel中，但我们想专门讨论一下，以帮助读者更好地理解Triton中线程块的映射和执行顺序是如何处理的。因此，我们将逐步介绍一个朴素映射，然后是一个更优化的映射，以观察其相应的影响。

让我们从朴素地构建我们kernel开始，首先是从全局内存进行"线性"加载，然后将其与更优化的"swizzled"加载进行比较。线性与交错决定了我们在GPU上工作网格的执行顺序。让我们看看Nvidia Nsight Compute工具在朴素情况下提供的关于我们kernel共享内存访问模式的提示：

为了解决这个问题，我们可以使用一种称为"tile-swizzling"的方法。这种方法的想法是以更友好于L2缓存的顺序启动我们的线程块。

让我们后退一步，熟悉一下Triton的一些语义，并做一个简单的CUDA类比来更好地理解这个概念。Triton kernel启动"程序"。这些所谓的程序映射到CUDA中的线程块概念，它是Triton kernel中并行性的基本单位。每个程序都有一个关联的"pid"，程序中的所有线程都保证执行相同的指令。

如果您对"pid"进行简单的线性映射到输出矩阵C的2D网格位置，Triton程序将以朴素的方式分布到您的SM上。

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

红色高亮的代码是朴素的"线性"tile排序，绿色高亮的代码是"交错"tile排序。这种启动方式改善了局部性。这里有一个可视化图帮助来更好地理解这一点。

在合并这个更改后，ncu profiler不再抱怨未合并的内存访问。让我们看看我们的内存吞吐量如何变化：

这个更改在一个简单的加载存储 kernel 上进行了测试。查看分析器中的GPU速度统计部分，我们还看到简单加载kernel的内存吞吐量增加了112.07%，这正是我们通过此优化想要达到的目标。

再次强调，这种优化已经存在于AutoGPTQ kernel中，但它是每个Triton kernel程序员在开始编写任何令人兴奋的反量化或矩阵乘法逻辑之前，都必须在kernel开头编写的样板逻辑。因此，重要的是要理解：

• 这种映射并不是唯一的
• Triton不会自动为程序员处理这种优化，必须仔细考虑以确保您的kernel以最佳方式处理共享内存访问

对于那些刚接触Triton的人来说，这些并不明显，因为大部分共享内存访问优化都是由Triton编译器处理的。然而，在编译器无法处理的情况下，重要的是要能够理解我们可以使用哪些工具和方法来影响内存行为。

4.0 向量化加载

现在，回到我们未优化 kernel 的原始问题。我们想要优化 kernel 的全局内存访问模式。从Nvidia Nsight计算工具的详细页面中，我们看到以下注释，其中分析器抱怨未合并的全局内存访问。让我们深入研究一下未优化内存读取的SASS（汇编）代码加载：

这个加载操作导致了32个16位宽的全局加载操作，这并不理想。

我们希望以向量化的方式进行全局内存加载，以使其产生最少的加载指令。为了解决这个问题，我们可以给Triton编译器一些帮助。

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

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

加载现在通过4个每个128位宽的全局加载操作来执行，而不是32个16位的全局加载操作。这意味着减少了28个内存获取指令，更重要的是实现了合并的内存访问。这可以从单个线程不再访问连续内存地址的事实中看出，而没有编译器提示时则是这种行为。

在隔离的加载操作中，结果是73倍的加速，在将其整合到完整的反量化kernel中后，我们能够看到另外6%的加速。这是朝着正确方向迈出的又一步！

5.0 线程束停滞

现在将所有更改重新应用到我们的完整反量化kernel中，我们看到以下性能限制因素：线程束停滞。

这些线程束停滞主要由"长记分板"停滞引起，占总数的92.63%。

从高层次来看，长记分板停滞（https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference） 发生在线程束为了进入issued(发射)状态必须的数据还没准备好的时候。换句话说，GPU是吞吐量机器，我们需要用计算指令来隐藏加载指令的延迟。通过加载更多数据并重新安排脚本中加载指令的位置，我们可以解决这个问题。

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

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

我们如何改进这一点？

我们希望能够增加我们的内存吞吐量，这样当线程束发出指令时，我们不会等待加载被存储到SRAM中以便用于计算。我们尝试了多个参数（如流水线阶段数和线程束数），其中对k维度的块大小增加2倍产生了最大影响。