智能体/生成式 AI

基于 CUDA Tile IR 后端的 OpenAI Triton 推动 GPU 编程发展

NVIDIA CUDA Tile 是一种基于 GPU 的编程模型,旨在实现 NVIDIA Tensor Core 的可移植性,从而充分发挥 GPU 的性能潜力。CUDA Tile 的一大优势在于,您可以在其上构建自己的领域特定语言(DSL)。

本文将介绍 NVIDIA 如何将 CUDA Tile 集成为 OpenAI Triton 的后端。OpenAI Triton 支持平铺计算,这是一种将数据与计算任务划分为较小块的技术。Triton 包含一个基于 MLIR 的编译器,能够生成 PTX 代码,使缺乏 CUDA 经验的研究人员也能编写高效的 GPU 程序。

什么是 CUDA Tile 和 CUDA Tile IR?

CUDA Tile 扩展了 CUDA 编程模型,为图块编程提供一级支持。CUDA Tile 由 CUDA 13.1 引入,标志着 GPU 编程的一次范式转变。基于图块的模型允许在更高层次的抽象中表达计算,而无需开发者按照 SIMT 模型中的单个线程进行思考。

您只需在数据块(图块)上指定操作,编译器和运行时系统便会自动处理线程调度、硬件映射与资源分配。这一设计既能降低编程的复杂性,又能支持更激进的编译器优化。

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 内核编译并运行在新推出的 CUDA Tile IR 后端上。它将高级编程语言(Triton)与 NVIDIA 新一代 GPU 编程模型连接起来,为开发者提供了一条无需重写代码即可利用现代硬件功能的平滑路径。

随着 GPU 编程不断从传统的 SIMT 模型向基于图块的抽象演进,这种集成使开发者既能受益于 Triton 简洁易用的 Python 语法,又能获得对 Tensor Core 和架构可移植性的 TileIR 原生支持。

Triton-to-TileIR 实现了这些新功能的普及。值得注意的是,Triton 本身本质上是一种基于图块的编程语言,开发者使用数据块(图块)来表达计算,而非以单个线程为单位,这在概念上与 CUDA Tile IR 保持一致。

这提供了一个直接的后端编译路径:Triton 的图块级抽象不编译为线程级 SIMT 代码,而是保留图块级语义,并直接编译为 CUDA Tile IR,使后者能够原生支持图块粒度的计算。

现有的 Triton 用户社区无需学习新语言或重写现有代码,即可充分发挥 CUDA Tile IR 的优势。通过简单的环境变量配置,便可将编译流程从 PTX 后端切换至 CUDA Tile IR 后端,从而获得更高性能以及面向未来的架构兼容性。

Triton 用户将能够为其应用程序中的每个内核选择所使用的后端(PTX 后端或 CUDA Tile IR 后端)。

Triton-to-TileIR 的开发路线图

作为 Triton-lang 组织内部的孵化器项目,Triton-to-TileIR 正在积极开发中。该仓库作为协作空间,用于在集成到主 Triton 编译器之前实现并完善 CUDA Tile IR 后端。

开发路线图可能包含多个技术工作流程,例如:

  1. 核心转换基础设施: 实施 MLIR 方言转换模式,将 Triton 运算映射到 CUDA Tile IR 的等效表示
  2. 测试与验证: 构建全面的测试套件,验证转换过程的语义正确性,覆盖控制流、内存访问模式及数值精度中的边界情况
  3. 性能基准测试: 建立性能评估体系,针对不同算子(如矩阵乘法、卷积、元素级运算、归约等),对比 TileIR 编译生成的内核与 PTX 编译内核的执行效率
  4. 开源项目集成: 协同开源社区,推动 CUDA Tile IR 后端在开源项目(例如 Helion)中的深度集成与支持

如何使用 Triton-to-TileIR

Triton-to-TileIR 目前仅支持从源码进行编译。预编译的二进制文件不可用,需在本地环境中自行从源代码构建项目。

预备知识:

  • CUDA 版本: CUDA 13.1 或更高版本
  • GPU 架构: NVIDIA Blackwell GPU(例如 GeForce RTX 5080);先前的 GPU 架构将在后续发布的 CUDA 版本中启用

从源代码构建

在满足前提条件后,从源代码克隆项目并进行构建:

# 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 文件扩展名来缓存已编译的内核,而不是 SIMT 后端所使用的标准 .cubin 文件。请检查以下缓存文件:

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

Triton-toTileIR 的局限性

虽然 Triton-to-TileIR 带来了前景广阔的新可能性,但该项目仍处于较早的开发阶段,存在一些已知的限制,例如不支持的操作和临时的性能问题。

不支持的操作

并非所有 Triton 支持的操作都已在 Tile IR 后端实现。可进一步了解尚未支持或未完全支持的操作与功能

随着 CUDA 持续发布新版本,Triton CUDA Tile IR 后端的兼容性将逐步提升。

指针张量降级导致性能欠佳

Triton 中的“指针张量”模式(其中张量由用于描述内存访问模式的指针构成)表明,在使用 CUDA 13.1 时,Tile IR 后端的性能表现不佳。这一性能问题是暂时的。对于受影响的工作负载,您可以:

  • 对于某些关键操作,暂时回退到 SIMT 后端
  • 等待后续项目版本中的优化
  • 优化代码以采用 TMA 加载/存储 API

关于最后一点,优化代码以采用 TMA 加载/存储 API:内核中加载的许多张量具有连续的图块以及明确的形状和步长。因此,无需在核函数内实现指针张量,而可将这些布局信息传递给 TMA 加载/存储 API,从而使 Tile IR 后端实现更优的性能。

例如,典型的指针张量模式可能如下所示:

# 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 中的每个元素都是核函数中计算的显式指针,尽管图块本身是连续的,其布局仍可通过 (shapestridesblock_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 编程发展的重要一步,有效弥合了开发效率与硬件性能之间的鸿沟。通过为 CUDA Tile IR 虚拟指令集提供 Triton 可访问的、面向图块的编程模型,该集成有望为机器学习从业者和 GPU 开发者带来更优的性能、更强的可移植性以及更好的未来适应能力。

对于已采用 Triton 的开发者,TileIR 后端将提供一种利用新一代 GPU 架构的途径,同时所需代码修改极少。对于更广泛的 GPU 编程生态系统而言,此次合作展现了语言设计者与硬件供应商之间的战略协作如何创造复合优势——在不牺牲支持快速创新的高级抽象的前提下,使高级硬件功能得以实现。

随着项目的成熟以及从孵化阶段向生产就绪阶段的演进,我们将非常关注集成对 Triton 采用情况的影响,以及基于图块的 GPU 编程整体发展路径。最终的成功标准十分明确:GPU 专业知识有限的研究人员是否能够编写出在 NVIDIA GPU 上接近最优性能运行的 Triton 代码。

如需了解更多信息,请查阅 Triton-lang/Triton-to-tile-IR GitHub 仓库以及 CUDA Tile IR 后端的性能调优建议

 

标签