借助NVIDIA CUDA Tile IR后端推进OpenAI Triton的GPU编程

作者：Jie Xin和Jonathan Bentz

NVIDIA CUDA Tile是基于 GPU 的编程模型，其设计目标是为 NVIDIA Tensor Cores 提供可移植性，从而释放 GPU 的极限性能。CUDA Tile 的一大优势是允许开发者基于其构建自定义的 DSL。

本文介绍了 NVIDIA 正在将 CUDA Tile 集成为OpenAI Triton的后端的相关工作。OpenAI Triton 是一个开源的 Python DSL，专门用于编写 GPU 上的 DL kernels。它支持分块计算，可将数据与计算任务划分为较小的 Block。Triton 内置基于 MLIR 的编译器，能够生成PTX代码，这使得即使不具备 CUDA 经验的研究人员也能编写出高效的 GPU 代码。

什么是CUDA Tile和CUDA Tile IR？

CUDA Tile 是对 CUDA 编程模型的扩展，对 Tile 编程提供原生支持。CUDA Tile 在CUDA 13.1版本首次推出，标志着 GPU 编程的一次重要演进。与要求开发者基于 SIMT 模型、以单个线程为单位进行思考不同，基于 Tile 的模型可使开发者在更高抽象层级上表达计算。

开发者只需操作数据块（即 Tile），编译器和运行时系统便会自动处理线程调度、硬件映射与资源分配。这不仅显著降低了编程复杂度，也为编译器的大幅优化创造了条件。

CUDA Tile IR是基于 MLIR 的中间表示形式及编译器基础设施。CUDA Tile 的整个开发流程由 CUDA Tile IR 规范驱动，该规范明确定义了在 NVIDIA GPU 上进行 Tile 计算所需的形式化语义、操作和类型系统。

什么是Triton-to-TileIR？

Triton-to-TileIR后端是 Triton 的一个桥接层，使其能够以 CUDA Tile IR（而非 PTX）作为目标代码。它扩展了 Triton 编译器生态，使开发者能将使用 OpenAI Triton 编写的 GPU kernel 编译并运行在新推出的 CUDA Tile IR 后端之上，使开发者无需重写代码，便可无缝利用现代硬件能力。

随着 GPU 编程从传统的 SIMT 模型演进到基于 Tile 的抽象，这一集成让开发者既能保留 Triton 易用的 Python 语法，又能获得对 Tensor Cores 的原生 TileIR 支持以及更强的架构可移植性。

Triton-to-TileIR 降低了这些新能力的使用门槛。值得注意的是，Triton 本质上是基于 Tile 的编程语言，其理念是让开发者以数据块（即 Tile）为单位计算，而非单个线程。这一设计思想与 CUDA Tile IR 十分契合。

这种一致性使得 Triton 能够采用一条更直接的后端编译路径：Triton-to-TileIR 不再将 Triton 的 Tile 级抽象转换为线程级的 SIMT 代码，而是保持其 Tile 级语义，并直接编译至原生支持 Tile 粒度计算的 CUDA Tile IR。

Triton 现有用户无需学习新语言或重写代码，即可体验 CUDA Tile IR 带来的性能优势。仅需设置环境变量，便可将编译流程从原有的 PTX 后端切换至 CUDA Tile IR 后端，从而获得更优性能与面向未来的架构兼容性。

此外，Triton 用户将能够根据每个 kernel 的具体需求，在应用中灵活选择使用 PTX 后端或 CUDA Tile IR 后端。

Triton-to-TileIR开发路线图

Triton-to-TileIR 目前是 triton-lang 团队下的一个孵化项目，目前正处于积极开发阶段。该代码库旨在作为协作平台，共同实现并优化 CUDA Tile IR 后端，并为其未来并入 Triton 主分支奠定基础。

其开发路线图主要包括以下几项关键技术工作流：

核心转换基础设施：实现 MLIR 语言转换机制，将 Triton 操作映射到对应的 CUDA Tile IR 操作。

测试与验证：构建完整的测试套件，以验证包括控制流、内存访问模式、数值精度等各类边缘情况转换过程中的语义正确性。

性能基准测试：设立性能基线，在矩阵乘法、卷积、逐元素操作及规约等多种操作中，系统性地对比 kernel 分别经由 TileIR 与 PTX 编译后的性能差异。

开源项目集成：与开源社区紧密协作，推进 CUDA Tile IR 后端在 Helion 等相关开源项目中的支持。

如何使用Triton-to-TileIR

目前，Triton-to-TileIR 仅支持从源代码编译，暂无预编译的二进制包，因此您需要在本地自行从源代码构建该项目。

前提条件：

CUDA版本：CUDA 13.1 或更高

GPU架构：NVIDIA Blackwell架构 GPU（如GeForce RTX 5080）；后续 CUDA 版本预计将增加对早期 GPU 架构的支持

从源码构建

在满足以上条件后，可用以下代码构建项目：

# Clone the repository
git clone https://github.com/triton-lang/Triton-to-tile-IR.git
cd Triton-to-tile-IR
 
# Build and install
# Specific build instructions should be followed according to the project's README
pip install -e .

需要注意的是，具体的构建步骤可能随项目更新而变化。可查阅Triton-to-TileIR README文件与构建指南，获取针对您系统架构的详细配置说明、依赖管理指引及排障方法。

验证Tile IR编译

随后，您可以通过运行向量加法教程，并确认其是否调用 Tile IR 后端，来验证安装是否成功：

# Navigate to the tutorial directory
cd python/tutorials
 
# Run the vector addition example with Tile IR enabled
export ENABLE_TILE=1
python 01-vector-add.py

当 Tile IR 后端处于激活状态时，Triton 会使用 .tileIR 文件扩展来缓存编译后的 kernel，而不是 SIMT 后端所使用的标准 cubin 文件。请检查以下缓存文件：

# Find the Triton cache directory (typically in ~/.triton/cache)

Triton-to-TileIR的局限性

尽管 Triton-to-TileIR 展现了一定的潜力，但目前仍处于早期开发阶段，对部分操作的支持尚不完善且存在一些临时性的性能问题。

暂不支持的操作

目前，Tile IR 后端尚未完全覆盖 Triton 支持的所有操作。建议查阅“当前暂未支持或功能未完全覆盖的操作清单”。

随着 CUDA 版本的不断迭代，Triton CUDA Tile IR 后端也将在兼容性与功能上实现同步增强。

Tensor-of-pointer性能下降

在 CUDA 13.1 的 Tile IR 后端上，Triton 的“tensor-of-pointer”模式（即用 pointer 构成的 tensors 来描述内存访问模式）目前性能暂未达到最优水平。这是一个暂时性的性能局限。对于受影响的工作负载，可以采取以下措施：

将关键操作暂时切换回 SIMT 后端执行

关注后续版本更新，相关问题将在新版本中优化

优化代码，改用 TMA load/store API

关于最后提到的代码优化，许多 kernel 中加载的 tensors 本身具有连续的 Tile 布局以及明确的形状和步幅。因此，无需在 kernel 内部显式构建 tensor-of-pointers。相反，可以直接将这些布局信息传递给 TMA load/store API，从而提升 Tile IR 后端性能。

以下为典型 tensor-of-pointers 模式代码：

# Before: tensor-of-pointer style
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
 
a_ptrs = a_ptr + (offs_m[:, None] * stride_am
                  + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk
                  + offs_n[None, :] * stride_bn)
 
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)

a_ptrs 中的每个元素都是在 kernel 中计算出的显式 pointer，即使该 Tile 本身是连续的，并且其布局可以通过（shape,strides,block_shape）完全描述。

使用 TMA，相同的操作可以重写为：

desc_a = tl.make_tensor_descriptor(
    a,                             # base pointer
    shape=(M, K),
    strides=(stride_am, stride_ak),
    block_shape=(BLOCK_M, BLOCK_K) # tile size
)
desc_b = tl.make_tensor_descriptor(
    b, shape=(K, N),
    strides=(stride_bk, stride_bn),
    block_shape=(BLOCK_K, BLOCK_N)
)
 
 
offs_m = pid_m * BLOCK_M
offs_n = pid_n * BLOCK_N
 
a_tile = desc_a.load([offs_m, 0])       # [BLOCK_M, BLOCK_K]
b_tile = desc_b.load([0, offs_n])       # [BLOCK_K, BLOCK_N]
desc_c.store([offs_m, offs_n], acc)     # TMA-backed store

进一步了解Triton-to-TileIR

Triton-to-TileIR 项目代表了 GPU 编程演进的关键一步，它弥合了开发生产力与硬件效率之间的差距。通过将 Triton 易用、面向 Tile 的编程模型与 CUDA Tile IR 虚拟指令集连接起来，该集成有望为机器学习从业者和 GPU 开发者带来性能提升、可移植性并为未来做好准备。

TileIR 后端为现有 Triton 开发者提供了一条平滑的升级路径，仅需修改极少的代码即可体验下一代 GPU 架构。同时，它也向更广阔的 GPU 编程生态展现了语言设计者与硬件供应商深度协同的战略价值：在不牺牲高级抽象快速迭代能力的前提下，更简单地获取先进硬件功能。

随着该项目从孵化阶段逐步走向生产就绪，其如何推动 Triton 的采用率、如何重塑基于 Tile 的 GPU 编程范式，将成为值得持续关注的焦点。最终的衡量标准也将非常直观：不具备深厚 GPU 专业知识的研究人员能否 NVIDIA GPU 上编写出接近最优性能的 Triton 代码。

如需了解更多详情，可查阅triton-lang/Triton-to-tile-IR的 GitHub 库以及《CUDA Tile IR 后端性能调优提示》文档。

关于作者

Jie Xin 是 NVIDIA 的计算架构师。他毕业于华中理工大学，主修计算机科学。专注于编译器、深度学习框架和 GPU 架构。

Jonathan Bentz 领导 NVIDIA 的 CUDA 技术营销工程团队，其团队专注于创建和提供引人入胜的内容，并与 CUDA 开发者建立联系。Jonathan 拥有爱荷华州立大学化学博士学位和计算机科学硕士学位。