今日阅读 tritonBLAS 相关内容

https://t.co/RPI44VAXOt https://t.co/vQc5EV7LZ8

基于 Triton 的分析方法用于 GEMM 内核参数选择

来源：https://arxiv.org/html/2512.04226

tritonBLAS：基于 Triton 的分析方法用于 GEMM 内核参数选择

Ryan Swann, Muhammad Osama, Xiaohu Guo, Bryant Nelson, Lixun Zhang, Alex Brown, Yen Ong, Ali Yazdani, Sean Siddens, Ganesh Dasika, Alex Underwood
所属机构：Advanced Micro Devices, Inc. {ryan.swann, muhammad.osama, xiaohu.guo, bryant.nelson, lixun.zhang, alex.brown, yen.ong, ali.yazdani, sean.siddens, ganesh.dasika, alexander.underwood}@amd.com

摘要

我们推出 tritonBLAS，一个快速且确定性的分析模型。它利用缓存层次结构等架构参数以及代码和数据的相对位置，来生成高性能的 GPU GEMM 内核。tritonBLAS 明确地对架构拓扑、矩阵形状和算法分块行为之间的关系进行建模，以在无需运行时自动调优的情况下预测近最优配置。基于此模型，我们在 Triton 中开发并实现了一个轻量级的 GEMM 框架。

我们在现代 GPU 上针对多样化的 GEMM 问题规模评估了 tritonBLAS 的性能。tritonBLAS 达到了自动调优解决方案性能的 95% 以上，同时将自动调优时间缩减为零。这使得 tritonBLAS 成为生产环境 HPC 和 ML 工作负载中经验调优的实用即插即用替代方案。

索引关键词：

分析性分块，通用矩阵乘法，GPU，ROCm， Triton

I 引言

通用矩阵乘法 (GEMM) 驱动着 AI、ML 和 HPC 工作负载的性能[1 (https://arxiv.org/html/2512.04226#bib.bib21),2 (https://arxiv.org/html/2512.04226#bib.bib22)]。像 GPT-4 这样的大型语言模型 (LLMs) 依赖于每个前向传播需要执行数千到数百万次 GEMM 操作的 Transformer 网络[3 (https://arxiv.org/html/2512.04226#bib.bib23)]。配备专用矩阵核心的 GPU 主导着 GEMM 加速。尽管该算法表面上简单，但为 GPU 开发高效的 GEMM 内核和库仍然是一项艰巨的任务。编写快速的 GEMM 内核需要精心捕捉 GPU 的计算和内存层次结构，协调数据在寄存器、共享内存和缓存之间的移动，并确保矩阵核心持续获得数据供给以避免停顿。

实现一个快速内核只是问题的一部分。构建一个完整的库不仅需要一组优化的内核，还需要在运行时为用户的问题选择或即时 (JIT) 编译111即时 (JIT) 编译指的是在运行时根据用户提供的问题参数生成并编译专门的机器代码的策略。这使库无需依赖预编译变体即可根据特定的矩阵形状和硬件特性定制内核实现。合适的内核配置。这些配置决策包括分块大小、循环展开深度、循环调度策略以及改善 L2 缓存命中率的数据重排的选择。因此，库的有效性在很大程度上取决于内核的选择和编排，而不仅仅是单个内核的原始效率。

Triton[4 (https://arxiv.org/html/2512.04226#bib.bib1)]已成为编写高性能 GPU 内核的高效编程模型。它简化了 GPU 编程的许多复杂性，例如线程同步、向量化和内存合并，同时仍然暴露了性能调优所需的必要控制。因此，Triton 使得编写 GEMM 内核变得显著容易。然而，Triton 及类似框架仍然依赖于自动调优来选择内核参数。这些基于搜索的技术需要大量的编译和基准测试，消耗大量时间和资源，并且限制了调优配置在不同硬件代际之间的可移植性（参见第 II-B 节 (https://arxiv.org/html/2512.04226#S2.SS2)）。

我们推出 tritonBLAS，一个基于 Triton 构建的一流 GEMM 库，它消除了对自动调优的依赖。tritonBLAS 不使用经验搜索，而是使用一种分析性能模型，该模型捕捉了 GPU 架构与算法分块选择之间的交互。该模型使我们能够确定性地预测近最优配置，并通过 JIT 编译生成专门的内核，而无需运行调优阶段。通过将 Triton 的可编程性与有原则的分析模型相结合，我们为生产环境的 HPC 和 ML 工作负载提供了一种替代自动调优的实用方案。

本文的贡献如下：

• •一个用于 Triton GEMM 的确定性、零自动调优选择器。我们利用此模型实现了 tritonBLAS，一个基于 Triton 的 GEMM 库，它用基于架构参数的配置选择取代了经验性自动调优。
• •具有轻量级校准的架构可移植性。该模型仅由可测量的硬件比率（带宽、指令延迟和矩阵核心形状）参数化，从而可以通过基于微基准测试的校准重新定位到新的 GPU 代。
• •广泛的评估证明了接近自动调优的性能。在 150,000 种 GEMM 形状和真实 LLM 工作负载上，tritonBLAS 达到了详尽自动调优性能的 94.7%，同时消除了调优开销，并在内存受限场景中匹配或超越了供应商库。

II 背景与相关工作

II-A 通用矩阵乘法 (GEMM)

通用矩阵乘法 (GEMM) 定义为 C=αAB+βC，其中 C 是 M×N 大小的输出矩阵，A 和 B 是输入矩阵，大小分别为 M×K 和 K×N，α 和 β 是标量值。GPU 上的高效 GEMM 实现利用了 GPU 的内存和计算层次结构，通过层次化的分块结构来最大化局部性以及可用并行性的利用率之间的平衡[5 (https://arxiv.org/html/2512.04226#bib.bib6),6 (https://arxiv.org/html/2512.04226#bib.bib7)]。GEMM 分块使用“输出驻留”数据流进行调度，其中归约维度 K 随时间累积到输出缓冲区中，输出分块在处理单元之间并行化（参见算法 1 (https://arxiv.org/html/2512.04226#alg1) 和 2 (https://arxiv.org/html/2512.04226#alg2)）。GEMM 的输出分块维度同时影响并行性和数据局部性，并受内存层次带宽和容量以及计算资源形状和吞吐量等架构特性影响。这意味着任何问题规模和形状 (M, N, K) 的 GEMM 可能都需要不同的分块大小才能在特定架构上达到峰值性能。

算法 1 计算矩阵 C=A×B 1: 输入 A∈Rm×k, B∈Rk×n 2: 输出 C∈Rm×n 3: for i=1,…,m do 4: for j=1,…,n do 5: C[i,j] ← 0 6: for ℓ=1,…,k do 7: C[i,j] ← C[i,j] + A[i,ℓ] × B[ℓ,j] 8: end for 9: end for 10: end for

算法 2 分块矩阵乘法调度 1: M,N,K: 问题的维度; M_T,N_T,K_T: 空间和时间上的分块大小 2: 部分结果累积在 C 中 3: function ScheduleTile(i,j,ℓ) ▷ 在 CU 索引 (i,j,ℓ) 处计算一个分块 4: for m ← 1 to M_T do ▷ 遍历一个块中的波 5: for n ← 1 to N_T do ▷ 遍历一个块中的波 6: for k ← 1 to K_T do ▷ 遍历一个波中的时间 7: C[i][j][m][n] += A[i][ℓ][m][k] × B[j][ℓ][n][k] 8: end for 9: end for 10: end for 11: end function 12: for i ← 1 to ⌈M/M_T⌉ do ▷ 在 M 上遍历 CU（空间） 13: for j ← 1 to ⌈N/N_T⌉ do ▷ 在 N 上遍历 CU（空间） 14: for ℓ ← 1 to ⌈K/K_T⌉ do ▷ 遍历时间分块 15: ScheduleTile(i,j,ℓ) 16: end for 17: end for 18: end for

一个说明输出分块大小如何影响并行性的例子是当矩阵大小为 M=256, N=256 和 K=8192 时。分块大小 (M_T×N_T×K_T)=(16×16×256) 会产生 256 个输出分块 (M/M_T × N/N_T = 256/16 × 256/16)，这些分块在一个拥有 256 个计算单元 (CU) 的 GPU 上并行化，在典型情况下每个 CU 负责生成输出矩阵的一个分块[6 (https://arxiv.org/html/2512.04226#bib.bib7)]。考虑如果我们选择了分块大小 (M_T×N_T×K_T)=(256×256×128)，结果是一个输出分块 (M/M_T × N/N_T = 256/256 × 256/256)：整个问题将在单个 CU 上运行，导致计算资源未充分利用。

单独优化局部性会导致分块中的元素被重用 M_T×N_T×K_T 次 MAC 操作。例如，一个 16×16×256 的分块大小总共被重用 65,536 次，而 256×256×128 分块中的元素被重用 8,388,608 次。这是因为分块参数改变了一次加载到较低级别内存层次结构中的数据量，即一次加载到较低级别缓存的数据可以用于多个矩阵乘法操作。分块配置也是递归的；大多数 GPU 都有一些针对给定数据类型的硬件预定义的“矩阵指令”，这些指令将作为递归基本情况的固定大小输入分块进行处理。

当代方法使用自动调优作为解决方案来构建 GEMM 库，以解决并行性和局部性之间微妙的平衡。对 GEMM 问题进行基准测试以确定峰值性能配置。虽然自动调优在调优情况下能产生 100% 效率的解决方案，但它需要 O(P×M×N×K) 的时间，其中 P 是分块配置的数量，当添加不同的批处理或分组 GEMM 维度时复杂性进一步增加，并且不是一种可扩展的方法。

II-B 自动调优与 Triton

Triton 是一种领域特定语言和编译器框架，专为编写高性能 GPU 内核而设计，特别是那些用于深度学习和张量计算工作负载的内核[4 (https://arxiv.org/html/2512.04226#bib.bib1)]。Triton 旨在通过一个高级 Python 接口简化 GPU 编程，同时提供与手工优化内核相当的性能。Triton 内核用 Python 编写，并通过即时 (JIT) 编译以针对特定 GPU 架构。这使开发者可以专注于内核的算法结构，同时将许多底层优化细节留给编译器。

为了在不同输入规模和硬件平台上实现性能可移植性，Triton 包含一个自动调优系统，这是现代 GPU 编程框架中常见的机制，例如 TVM[7 (https://arxiv.org/html/2512.04226#bib.bib2)]、TensorRT[8 (https://arxiv.org/html/2512.04226#bib.bib3)]、TensorFlow[9 (https://arxiv.org/html/2512.04226#bib.bib4)] 和 TorchInductor[10 (https://arxiv.org/html/2512.04226#bib.bib5)]。自动调优指的是自动探索实现选择空间的过程，例如分块大小、内存分块策略、循环展开因子和 warp 调度参数，以发现为特定目标设备和输入形状提供最佳运行时性能的配置。

在 Triton 中，自动调优机制通过 @triton.autotune 装饰器暴露。用户使用 triton.Config 对象定义一组候选配置，每个配置指定编译时参数，如分块大小、warp/wavefront 数量和流水线阶段数。此外，用户提供一组依赖于输入的调优键（例如张量大小或形状），这些键决定了何时应触发新的调优搜索。当某个 Triton 内核首次使用特定键调用时，框架会编译所有候选配置，并使用基于事件的机制在目标 GPU 上对其进行基准测试，以测量执行延迟。产生最短执行时间的配置被选为最优配置，并在内部缓存，以便在后续使用相同调优键的调用中重用。

然而，自动调优过程引入了多种形式的开销。首先，初始调优阶段会产生编译和测量成本，当配置空间很大时，这可能非常显著。其次，自动调优器的准确性受限于用户手动提供的候选配置空间的质量和完整性。Triton 目前使用对用户定义配置的穷举搜索，这限制了可扩展性和对具有各种内核形状和大小的应用程序的适应性。最后，在运行时对每个配置进行基准测试的需求使得这种方法对于许多应用程序来说是不适用和不切实际的，例如：

• •动态张量形状和大小随运行而变化，
• •需要低延迟的实时约束，如在线推理或控制系统，以及
• •已经具有极高成本和功耗的应用程序，如 LLM 训练运行。

GEMM 内核配置、并行性、局部性和性能之间的复杂底层关系早已被研究[11 (https://arxiv.org/html/2512.04226#bib.bib12),12 (https://arxiv.org/html/2512.04226#bib.bib11),13 (https://arxiv.org/html/2512.04226#bib.bib10),14 (https://arxiv.org/html/2512.04226#bib.bib9)]。许多研究基于架构形状和瓶颈分析性地推导出理想的分块配置。我们将此概念扩展到高度并行且复杂的 GPU 架构，并以 GEMM 为例。在接下来的两节中，我们提供了一个分析模型，用于在 AMD Instinct™ MI300X 加速器上捕捉这些复杂性。

II-C GPU GEMM 的分析性能模型

最近的 GPU 优化系统越来越多地采用分析或半分析模型来降低经验性自动调优的成本。例如，CUTLASS[15 (https://arxiv.org/html/2512.04226#bib.bib24)] 和 cuBLAS[16 (https://arxiv.org/html/2512.04226#bib.bib28)] 使用手工设计的启发式方法来评估线程块分块基于内存流量、寄存器压力和张量核心利用率的可行性。这些规则对于 CUDA 模板效果很好，但它们是闭源的，并且与 NVIDIA 特有的执行模型紧密耦合，无法很好地移植到 Triton。同样，Ansor 和 AutoTVM[17 (https://arxiv.org/html/2512.04226#bib.bib26),18 (https://arxiv.org/html/2512.04226#bib.bib25)] 结合了分析性重用估计和学习到的校正项来修剪大的搜索空间；然而，它们对数据驱动拟合的依赖使得它们不太适合动态形状的工作负载或需要确定性、零样本预测的场景。

一些工作还扩展了屋顶线模型以更好地捕捉 GPU 内核中的局部性。例如，DeLTA[19 (https://arxiv.org/html/2512.04226#bib.bib27)] 引入了一种局部性感知的屋顶线分析，用于…