当前位置：首页 > news >正文

TileLang 与 Triton，AMD 显卡上自定义高性能算子的开发笔记

news 2026/6/26 5:30:28

为什么在 ROCm 7.x 时代还要手写算子？

在大模型推理日益普及的今天，很多开发者习惯了直接调用 PyTorch 或 vLLM 现成的接口。但在 AMD Instinct GPU 上，尤其是面对 ROCm 7.x 这样快速迭代的生态，通用算子往往无法完全榨干硬件性能。特别是在处理非标准维度、特殊量化格式或自定义注意力机制时，官方库的支持可能存在滞后。这时候，掌握利用 Triton 或 TileLang 编写自定义 Kernel 的能力，就成了区分“调包侠”和“系统专家”的分水岭。

最近我在 DevCloud 上折腾 MI300X 时，发现一个有趣的现象：默认的矩阵乘法在特定 batch size 下，显存带宽利用率竟然只有理论值的 60% 左右。排查后发现，问题出在内存访问模式不对齐，导致大量的 Load/Store 指令浪费在了无效的数据搬运上。与其等待社区更新，不如自己动手优化。本文将分享如何利用 Triton 和 TileLang 在 AMD 架构上编写高性能算子，重点解决内存对齐问题，并对比优化前后的性能差异。

环境准备与工具链选型

工欲善其事，必先利其器。在开始编写代码前，确保你的开发环境已经就绪。我推荐使用 Ubuntu 22.04 LTS，这是目前对 ROCm 7.x 支持最稳定的发行版。安装完官方驱动后，务必运行rocm-smi确认显卡状态正常，并通过rocminfo记下你的 GPU 架构代码（例如 MI300X 对应的是gfx942）。这一步至关重要，后续编译参数全靠它。

对于编程语言的选择，Triton依然是首选。它的 Python 嵌入式 DSL 让编写 GPU Kernel 变得像写 NumPy 一样直观，且 ROCm 后端在 7.x 版本中已经相当成熟。而TileLang作为新兴的张量编程语言，虽然在生态丰富度上稍逊一筹，但在描述复杂的分块（Tiling）策略和内存层级管理上有着独特的语法优势，特别适合需要极致控制的场景。本次实践主要基于 Triton 展开，因为它更容易上手且社区资料更多，但核心优化思路对 TileLang 同样适用。

你需要安装与 ROCm 7.x 匹配的 Triton 版本。注意，不要直接使用 pip 上的通用包，最好从源码编译或寻找专门针对 AMD 构建的 wheel 包，以确保 HIP 后端被正确启用。验证安装是否成功的最快方法是运行一个简单的向量加法测试，如果能顺利输出结果且rocprof能看到对应的 Kernel 启动记录，说明环境没问题。

诊断性能瓶颈：内存访问的对齐陷阱

在优化之前，我们先得知道“慢”在哪里。通过rocprof分析默认实现的性能剖析报告，我发现了一个典型问题：非合并内存访问（Uncoalesced Memory Access）。

在矩阵乘法 $C = A \times B$ 中，如果线程块（Thread Block）内的线程读取全局显存时，地址不是连续的，硬件就无法将多次小请求合并为一次大请求。在 AMD CDNA 架构上，这会导致显存事务数激增，带宽利用率直线下降。特别是在处理非 2 的幂次维度，或者使用了特殊的 Padding 策略时，这种情况尤为常见。

举个例子，假设我们按行优先存储矩阵 A，但在 Kernel 设计中让线程按列去读取数据。这就好比去图书馆借书，本来可以一次抱走一排，结果非要一本一本跑断腿。解决这个问题的核心在于重排数据加载逻辑，确保相邻线程读取相邻内存地址。

实战：使用 Triton 编写优化的矩阵乘法

下面是一个基于 Triton 优化的矩阵乘法 Kernel 示例。这段代码的核心在于精心设计的BLOCK_SIZE和指针算术运算，以确保内存访问的对齐。

import triton import triton.language as tl import torch @triton.jit def matmul_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, ): pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + (pid % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m # 计算当前线程块负责的块起始位置 offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) # 关键优化：构建指针时确保步长对齐 a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # 加载数据，利用 mask 防止越界，同时保证合并访问 a = tl.load(a_ptrs, mask=(offs_am[:, None] < M) & (offs_k[None, :] < K - k * BLOCK_SIZE_K), other=0.0) b = tl.load(b_ptrs, mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N), other=0.0) accumulator += tl.dot(a, b) a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk c = accumulator.to(tl.float16) offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_bn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_bn[None, :] < N) tl.store(c_ptrs, c, mask=c_mask)

这段代码中有几个细节值得注意：

Swizzling 策略：通过GROUP_SIZE_M引入了一种简单的线程块调度优化，这有助于改善 L2 缓存的命中率，减少显存冲突。
Mask 加载：tl.load中的 mask 不仅是为了安全，更是为了告诉编译器哪些线程是活跃的，从而生成更高效的指令序列。
常量表达式：BLOCK_SIZE使用tl.constexpr，这让编译器能在编译期展开循环，极大减少运行时开销。

编译与性能剖析

代码写完后，编译过程由 Triton 自动完成，但我们需要指定正确的架构。在运行脚本前，设置环境变量：

export PYTORCH_ROCM_ARCH=gfx942

然后在 Python 脚本中调用 Kernel 时，传入合适的 block size。对于 MI300X，经过多次实验，BLOCK_SIZE_M=128,BLOCK_SIZE_N=128,BLOCK_SIZE_K=32通常能取得不错的效果。

为了验证优化效果，我编写了一个简单的 Benchmark 脚本，对比了原生 PyTorchtorch.mm和上述自定义 Kernel 在不同矩阵规模下的表现。结果显示，在 $4096 \times 4096$ 的矩阵乘法中，自定义 Kernel 的吞吐量提升了约18%，显存带宽利用率从 60% 提升到了 78% 左右。更重要的是，在非标准维度（如 $3500 \times 3500$）下，优化后的 Kernel 表现更加稳定，没有出现明显的性能抖动。

使用rocprof再次分析，可以看到GLOBAL_MEM_LOAD和GLOBAL_MEM_STORE的指令效率显著提高，不再出现大量的碎片化事务。这证明了针对内存访问模式的优化是行之有效的。