Skip to content

필사 모드: Triton Gluon — 把编译器藏起来的布局,重新用手写出来的语言

中文
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言 — Triton 卡壳的地方不是「太慢」

Triton 的吸引力很明显。用 Python 写一个 tile 级别的内核,编译器就会替你处理布局、内存分配、数据搬移、同步。多数情况下,这是一笔划算的交易 — 不用降到 CUDA C++ 就能拿到能用的性能。

问题出在这笔交易崩掉的那一刻。Gluon 教程的开篇精准地描述了这种局面:Triton 编译器对很大范围的内核都能生成高效代码,但也可能输给手工调优的底层代码。而一旦发生这种情况 — 用教程自己的话说 — "由于所有细节都被隐藏了,用户几乎没有办法有意义地改善性能"。

这就是核心所在。Triton 的局限不是"慢",而是一旦卡住,就没有可以走出去的门。抽象工作得好的时候是恩赐,一旦出岔子就是一堵墙。到目前为止,站在这堵墙前只有一个选项 — 用 CUDA C++ 从头重写。

Gluon 试图创造出第三个选项。

Gluon 到底是什么 — 同一套编译器栈,反过来的契约

官方文档的定义很短。Gluon 是 Triton 的底层 GPU 编程模型,它"直接暴露布局、共享内存、warp 特化和目标特定的功能,让高级内核可以用便利性换取控制权"。

重要的是,Gluon 不是一个独立的项目。正如教程所说,Gluon 和 Triton 用的是同一套编译器栈。两者都实现了基于 tile 的 SPMD 编程模型,也都是共享同一个前端和 JIT 基础设施的 Python DSL。主机(host)侧的代码也几乎一样 — 调用内核的方式、指定网格的方式,甚至 triton.autotune 都原样可用。

变化的是设备端代码的写法。照搬教程自己的总结:Gluon 改变的是设备代码怎么写,主机代码的变化只不过是内核多了几个超参数而已。

最简单的 Gluon 内核,和 Triton 的看起来没有区别。

import torch
import triton
from triton.experimental import gluon
from triton.experimental.gluon import language as gl


@gluon.jit
def copy_scalar_kernel(in_ptr, out_ptr):
    value = gl.load(in_ptr)
    gl.store(out_ptr, value)

@gluon.jit 代替 @triton.jit,用 gl 代替 tl。到这一步为止,看起来只是换了个名字。

留意一下那个 import 路径 — triton.experimental。本文最后还会回到这个话题。

第一个内核令人失望 — 666 GB/s

教程很诚实地从一次失败开始。它写了一个逐个复制标量的 memcpy 内核,对 XBLOCK 做自动调优,在 GB200 上复制 8 GB 数据。教程源码里以注释形式留下的结果是这样的。

Time:        24.00 ms
Throughput: 666.24 GB/s

自动调优器选出的最优 XBLOCK 是 2048,结果大约是 666 GB/s。教程对这个数字的评价是"远低于 GPU 8 TB/s 的峰值带宽" — 大致是峰值的 8% 左右。

原因看一眼内核就明白了。

@gluon.jit
def memcpy_kernel(in_ptr, out_ptr, xnumel, XBLOCK: gl.constexpr):
    pid = gl.program_id(0)
    start = pid * XBLOCK
    end = min(start + XBLOCK, xnumel)
    for i in range(start, end):
        value = gl.load(in_ptr + i)
        gl.store(out_ptr + i, value)

每个 program(也就是 GPU 上的 CTA,即线程块)一次只复制一个元素。这段代码完全没用上 GPU 的并行性。要一次搬动多个元素,就得加载 tile 而不是标量,而要加载 tile — 就必须选一个布局

这正是 Gluon 和 Triton 分道扬镳的地方。Triton 里本该由编译器替你选的东西,在 Gluon 里要由你自己写。

布局 — Gluon 交还给你的东西

布局规定了张量的各个元素如何分配给线程块内的各个线程。这种分配沿着 GPU 的层级结构展开 — 线程块、warp、lane,以及每个 lane 各自的寄存器。

最常见的一种是 BlockedLayout

gl.BlockedLayout(
    size_per_thread=[2, 4],
    threads_per_warp=[16, 2],
    warps_per_cta=[2, 2],
    order=[1, 0],
)

读法是这样的。

  • size_per_thread=[2, 4] — 每个线程拥有一块连续的 2x4 子 tile,存放在该线程的寄存器里。
  • threads_per_warp=[16, 2] — 线程 tile 按这个方式堆叠起来构成一个 warp tile。三者的乘积必须与每个 warp 的线程数相符(NVIDIA 硬件上是 32)。
  • warps_per_cta=[2, 2] — warp tile 按这个方式堆叠起来构成整个线程块。
  • order=[1, 0] — 先沿哪个维度做 tile 化。这是行优先(row-major)。

三者逐元素相乘就得到了块的形状 — 这里是 [64, 16]。

order 为什么重要,只要画出单个线程的 tile 就立刻能看出来。当 order=[1, 0] 时,寄存器沿内层维度递增。

[[T:0, T:1, T:2, T:3],
 [T:4, T:5, T:6, T:7]]

order=[0, 1] 时,同一个 2x4 tile 就变成了这样。

[[T:0, T:2, T:4, T:6],
 [T:1, T:3, T:5, T:7]]

同样的元素、同样的线程,排布方式不同。而正如后面会看到的,就这一处差异,能把带宽拉开好几倍。

把布局明确写出来,还有一个附带的好处 — 寄存器预算变得可以算出来。顺着教程的例子往下走,用上面那个布局(块 [64, 16])处理一个 128x128 的 f32 张量,会把块 tile 化成 [2, 8],每个线程每个块要用 8 个寄存器,整个张量总共用 8 * 16 = 128 个寄存器。在 Triton 里,这个数字一直藏在编译器背后。

需要留意的陷阱也随之而来。如果张量比块小,就会被广播。一个 32x8 的 f32 张量只有 256 个元素,但由于块的 tile 化方式不变,每个 program 都要用 64 * 16 = 1024 个物理寄存器来存它。一个小张量硬是吃掉了 4 倍的寄存器 — 这是那种没有显式布局就很难察觉的成本。

用数字看布局

从这里开始,都是记录在 Triton 仓库教程源码(python/tutorials/gluon/02-layouts.py)里的注释测量值。是上游开发者在 GB200 上实测后写进源码的,并非我自己复现的结果。

把 1D memcpy 重写成真正的 tile 加载,固定 XBLOCK=2048num_warps=4,只改变 size_per_thread(下文中的 R)。

gl.BlockedLayout(
    size_per_thread=[R],
    threads_per_warp=[32],
    warps_per_cta=[4],
    order=[0],
)

记录下来的结果。

R=1   6.574 TB/s
R=2   6.476 TB/s
R=4   6.474 TB/s
R=8   6.502 TB/s
R=16  6.214 TB/s

能读出两件事。第一,从标量循环的 0.666 TB/s 变成了 6.574 TB/s — 把两个记录下来的数值直接相除,大约是 10 倍(这是我自己的算术,不是上游的说法)。tile 加载加上布局带来的提升,量级就是这么大。

第二,把 R 调大并不会更好。R=1 最快,R=16 最慢。教程从这里开始打开 SASS,直接数 LDG.E/STG.E 指令的条数。这很能说明 Gluon 的性格 — 暴露布局并不意味着给你"更好的默认值",而是意味着把选择权连同验证这个选择的责任一起交给你。"加宽就会更快"这种直觉,在这里是错的。

布局选错,性能就会崩

到了 2D,故事变得更尖锐。

一个 8 GB、32K x 64K 的张量,用 XBLOCK=1YBLOCK=2048,配合一个设计成表现得像 1D 中获胜的 R=1 那样的布局来复制。

layout = gl.BlockedLayout([1, 1], [1, 32], [1, 4], [1, 0])

结果是 6.260 TB/s — 比 1D 慢 5%。教程把原因归结为 2D 算术更复杂等因素,认为这个程度还算说得过去。

但如果把输入张量转置呢?同一个内核,同一个布局。

0.774 TB/s

用教程自己的话说,性能"崩塌(craters)"了。原因很简单 — 内层维度不再连续,合并访存(coalescing)也就随之消失。32 个 lane 只有读取相邻地址时才能被合并成一次事务,但在转置后的张量上,每个 lane 指向的地址相隔甚远。

修复方法同样简单。改变块大小,把布局也转置过来。

layout = gl.BlockedLayout([1, 1], [32, 1], [4, 1], [0, 1])
# XBLOCK=2048, YBLOCK=1
6.590 TB/s

从 0.774 到 6.590。内核的算法一个字符都没变,变的只是布局描述符。

这里教程补充的一个观察很有意思。修复后的版本(6.590 TB/s)比 1D memcpy(6.574 TB/s)还要略快一点。转置和未转置两种情况下,每个 program 访问内存的方式其实是一样的,却仍然出现了差距,教程把这归因于 program 在 GPU 上被调度到哪里,也就是数据局部性的问题 — 比如 TLB 虚拟地址转换的缓存效果,又比如 H100 上 L2 缓存被划分成互相通信的多个分区。在这个层面,即便是"相同的访问模式",也不代表相同的性能。

由此得到的实务教训是:最优布局取决的不是内核本身,而是张量在全局内存里的 stride。所以教程最终收敛到了这样一个辅助函数。

def get_layout_for_gmem_access(tensor, num_warps):
    if len(tensor.shape) == 1:
        return gl.BlockedLayout([1], [32], [num_warps], [0])

    assert len(tensor.shape) == 2, "only 1D and 2D tensors are supported"
    assert 1 in tensor.stride(), "expected at least 1 contiguous dimension"
    if tensor.stride(1) == 1:
        return gl.BlockedLayout([1, 1], [1, 32], [1, num_warps], [1, 0])
    else:
        return gl.BlockedLayout([1, 1], [32, 1], [num_warps, 1], [0, 1])

布局不是硬编码成常量,而是从运行时的 stride 推导出来的。Gluon 内核如果既要灵活又要快,就需要这样的代码。

赢过 PyTorch 的地方,以及它的代价

Gluon 的 2D memcpy 什么时候真正值回票价,教程也讲到了。对一个 8 GB 的张量,把隔行跳过的视图(非连续张量)复制成连续张量 — 这正是 PyTorch 的 x.contiguous() 做的事。

2D memcpy:                6.258 TB/s
torch.Tensor.contiguous:  2.946 TB/s
2D memcpy (transposed):   6.398 TB/s

教程自己的说法是"已经比 PyTorch 的实现快了 2 倍以上"。这是 Gluon 真正体现价值的一类场景 — 输入或输出带有奇特布局的时候。不过有一个附带条件:1D memcpy 只有在输入和输出同时是连续内存块的视图时才能用。

那如果输入和输出的布局正好相反呢?教程也测了这种情况。输入沿 dim 1 连续,输出沿 dim 0 连续,张量大小是 32K x 32K。

2D memcpy (order=[1, 0]): 0.978 TB/s
2D memcpy (order=[0, 1]): 1.674 TB/s

不管选哪个布局都很糟糕 — 满足了一个,另一个就会崩掉。解法是在加载和存储时使用不同的布局,中间再做一次转换。

@gluon.jit
def memcpy_2d_inout_kernel(in_ptr, out_ptr, ..., layout_in: gl.constexpr, layout_out: gl.constexpr, ...):
    ...
    value = gl.load(in_ptr + in_offsets, mask=mask_in)

    # Use `gl.convert_layout` to perform layout conversions.
    value = gl.convert_layout(value, layout_out)

    gl.store(out_ptr + out_offsets, value, mask=mask_out)
2D memcpy (in/out layouts): 4.814 TB/s

从 0.978 涨到了 4.814 — 但依然够不到 6.5 TB/s 这个档位。教程也没有回避原因。convert_layout 需要用到共享内存,而共享内存在 GPU 上是稀缺资源。用它来做布局转换会拉低占用率(occupancy)、缩短最大流水线深度,从而拖累性能。教程的结论是"在这种情况下,转换成本是不可避免的,而且比低效全局内存访问的成本要便宜得多"。

也就是说,4.814 TB/s 不是一次胜利,而是矬子里拔将军式的选择。像这样把成本变得看得见,正是 Gluon 的性格所在。

发布时间线 — 走到了哪一步

Gluon 走到了哪一步,可以从发布说明里追踪出来。以下日期是从 GitHub 发布 API 确认过的。

版本日期Gluon 相关内容
v3.4.02025-07-30"Gluon Framework Comprehensive Enhancement" — static_assert、TensorDescriptor 内核参数、异步 TMA、张量内存、同步屏障、split/join/reshape 与规约
v3.5.02025-10-21Hopper WGMMA + 异步等待 (#7300, #7313)、向 Gluon 暴露 libdevice (#7890)、公开 API 文档字符串 (#7323)
v3.6.02026-01-212CTA 模式的初步支持 (#8644, #8653)、num_ctas 实现 (#8602)、面向 AMD gfx1250 的 async_copy (#8622)、实验性的 Triton → Gluon 转换器 (#8417)
v3.7.02026-05-072CTA 端到端支持、AMD warp 流水线 (#8586, #8975, #8980)、面向 gfx1250 的 4/8-warp Stream-K Gluon 内核 (#9370)、TDM L2 预取 (#9086, #9148)、细粒度集群屏障 (#9206)、sm103 上的 tcgen05.ld.red (#9151)、局部 scatter/gather (#8480)
v3.7.12026-06-18补丁发布 — 按发布说明所言"没有新功能或 API 变更",修复了 2 个回归问题

能读出一条清晰的脉络。2025 年年中骨架落地,进入 2026 年之后,逐步扩展到多 CTA 和最新硬件特性(Blackwell 的 tcgen05、集群屏障)。而 v3.7.1 是一次纯粹的 bug 修复发布,这也意味着截至最新的稳定版本,Gluon 的表面已经安静了大约一个月。

3.6.0 里的 Triton → Gluon 转换器值得单独提一句。听起来像是一个能让迁移变轻松的工具,但对应 PR(#8417,2025-10-11 合并)的说明文字把预期精准地压了下去:"这不是用于生产环境的,它只是让你能把 Triton 内核转换成一个朴素(naive)的 Gluon 版本,而这个 Gluon 版本无疑会慢上不少。"

一句很诚实的话。而且它概括了 Gluon 的本质 — 迁移到 Gluon 这件事本身不会带来任何性能提升。性能是在迁移之后,当你把布局选对了才会出现的。

不是 NVIDIA 专属 — AMD gfx950 与 MI355

上面那张表里有一点很显眼,那就是 AMD 相关条目反复出现。面向 gfx1250 的 async_copy、AMD warp 流水线、面向 gfx1250 的 Stream-K 内核、AMD 集群屏障。Gluon 从设计上就不是一个只面向 NVIDIA 的逃生通道。

AMD 也在推动这件事。2026 年 5 月 22 日,ROCm 博客发布了一篇由 Lixun Zhang、Jason Furmanek、Peng Sun、Emad Barsoum 撰写的文章"From Naive to Near-Peak: Building High-Performance GEMM Kernels with Gluon"。这是一篇教程,在基于 gfx950(CDNA4)的 MI350/MI355 上,用 Gluon 把一个 GEMM 内核从 v0 一路优化到 v9。

以下数字全部是 AMD 自己的厂商自测数据。 博客给出的测量条件是一块 MI355、ROCm 7.0,以及用 gfx950-tutorial-v0.1 标签构建的 Triton。

内核形状结果
FP16 GEMM v0 (naive)4096x4096x8192520 TFLOPS,MFMA 效率 25%
FP16 GEMM v9 (优化后)4096x4096x81921489 TFLOPS,MFMA 效率 98.75%
BF8 GEMM4096x4096x163843257 TFLOPS,MFMA 效率 99.72%
MXFP4 GEMM4096x4096x327685255 TFLOPS,MFMA 效率 92.41%

这里有一处容易被误读的地方,需要挑明。从 520 到 1489 TFLOPS 大约 3 倍的提升,并不意味着 Gluon 比 Triton 快 3 倍。 v0 和 v9 全都是 Gluon 内核 — 它们只是同一个仓库(ROCm/gfx950-gluon-tutorials)里、从 v0_naivev9_beyond_hotloop,一段为教学目的设计的优化旅程的两端而已。朴素的 Gluon 内核跑得慢是意料之中的事(正是前面那个转换器 PR 所说的那样),而这 3 倍讲的是"把布局和调度做对了能到什么程度",不是不同语言之间的比较。

而且博客自己也划清了界限 — 明确表示这不是 hipBLASLt 之类生产级库的替代品。就我核实过的正文而言,里面没有和 hipBLASLt 的直接对比数字。也就是说,"Gluon 内核赢过了厂商库"这种说法,并不是从这份材料里得出来的。

在方法论上,这篇博客也有诚实的一面。它把 MFMA 效率而不是 TFLOPS 作为主要指标,给出的理由是 MFMA 效率与时钟频率无关,跨次运行的可复现性也好于原始 TFLOPS。相比那些只顾着甩出随 boost 时钟起伏的 TFLOPS 数字的厂商材料,这算是高出一筹的态度。(当然,MFMA 效率 99% 本身并不代表工作负载就一定快 — 它的意思是矩阵引擎在热循环中几乎没有空闲时间,而这一点只有在内核本身用的是正确算法时才有意义。)

布局的底层 — 线性布局

再往下深入一层,就是线性布局(linear layout)。正如教程所说,Gluon 并没有一种正规(canonical)的布局表示 — 多个布局可以描述同一种元素映射。举个例子,下面这两者是等价的。

gl.BlockedLayout([1], [32], [4], [0])
gl.SliceLayout(1, gl.BlockedLayout([1, 1], [32, 1], [4, 1], [1, 0]))

如果你知道它们是等价的,就可以用 gl.convert_layout(x, layout, assert_trivial=True) 让编译器帮你验证这个转换是不是真的免费(除了寄存器重排之外没有别的开销)。这是一个很符合"显式语言"气质的 API。

而所有 Gluon 布局都可以用线性布局来表示。教程把它称为"表达力最强、最强大的表示方式",同时也补充说它"相对少见,且可能难以理解"。背后的理论已经发表成论文 — Keren Zhou、Mario Lezcano、Jeff Niu、Phil Tillet、Thomas Raoux 等人所著的"Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using F2"(arXiv:2505.23819,2025 年 5 月 28 日提交,2026 年 3 月 6 日修订)。这种方法把布局建模成 GF(2) 上的二元矩阵,从而避开了既有系统的二次复杂度。

对实践者来说,要点是这样的 — 多数情况下 BlockedLayoutSliceLayout 就够用,线性布局在你真正需要之前都可以不了解。但一旦你需要零成本的 split/join/reshape/permute,答案就在这里。

诚实的权衡 — 什么时候不该用它

现在到了该诚实的部分。

它依然是 experimental 的。 import 路径是 triton.experimental.gluon,在仓库里也放在 python/triton/experimental/gluon 下。官方的 Gluon 概览文档里并没有另外给出关于实验状态或稳定性的警告文字,但路径本身就已经说明它不承诺 API 稳定。从 3.4.0 一路到 3.7.0,每个版本 API 都在持续增加和变动的这段历史,也和这种解读吻合。如果把生产内核放在这上面,就要做好每次升级 Triton 版本都可能被打破的心理准备。

它要求的知识量不一样。 教程的开篇说得很直白 — 要写 Gluon 内核,需要对 GPU 硬件和 GPU 编程的多个方面有更深的理解。如果不知道每个 warp 有 32 个线程、不知道合并访存为什么取决于内层维度、不知道共享内存 bank 冲突是什么、不知道寄存器压力如何影响占用率,就填不出 BlockedLayout 的四个参数。Triton 把这些都藏起来,不是偷懒,而是一种服务。

迁移到 Gluon 本身不是免费的性能。 正如转换器 PR 所说 — 朴素迁移过去的 Gluon 会明显更慢。而前面看到的 R 扫描也表明,直觉上的选择(把宽度调大)反而可能更糟。如果你打算把自动化编译器原本做的事情接过来自己做,就得每次都用基准测试证明自己确实做得更好。

整篇教程讲的都是 memcpy。 本文引用的 GB200 数字,是上游开发者实测后记录在源码里的,并非我自己复现的结果。而且它针对的是受内存带宽限制的 memcpy — 如果你的内核是受计算限制的,布局扫描不会给你带来 10 倍提升。

归纳一下,判断标准是这样的。

用了值回成本的情况

  • 已经用 Triton 写好了,并且通过性能分析用测量数据确认编译器生成的代码远远达不到硬件极限。
  • 瓶颈在布局、数据搬移或同步上,在 Triton 这个层面已经没有可以拨动的旋钮。
  • 需要直接用到最新的硬件特性(Blackwell 的 tcgen05、Hopper WGMMA、2CTA、TMA、CDNA4 的 MXFP4)。
  • 内核生命周期长、跑得频繁,足以摊销调优所花的时间和每次升版所需的维护成本。

属于过度的情况

  • 还没有在 Triton 上做过性能分析。(这是大多数情况。)
  • torch.compile 生成的代码已经够用。
  • 内核受计算限制,布局不是瓶颈。
  • 厂商库(cuBLAS、hipBLASLt、cuDNN)已经覆盖了这个运算 — 就连 AMD 那篇博客也明确表示自己的教程不是 hipBLASLt 的替代品。
  • 团队里没人懂 GPU 微架构。

结语

Gluon 不是"更好的 Triton"。它是另一份契约

Triton 的契约是"细节我来处理,你只管写算法"。多数情况下这是一笔好交易,往后大多数内核也应该继续留在 Triton 里。Gluon 的契约是"细节你自己处理,但作为交换,卡住的时候有一扇门,而不是一堵墙"。

这个光谱很重要。过去 Triton 和 CUDA C++ 之间隔着一道悬崖 — 编译器给不出你要的东西时,你就得把语言和整个生态系统一并换掉。Gluon 在这道悬崖上架了一段台阶。你留在同一套编译器栈、同一个前端、同一个 JIT、同一个自动调优器里,只往下走你真正需要的那一部分。而这段台阶正同时架在 NVIDIA 和 AMD 两边这件事 — AMD 亲自写了 gfx950 教程,几乎每份发布说明里都能看到 AMD 的条目 — 或许才是更重要的消息。

不过请守住顺序。先用 Triton 写,先做性能分析,先确认编译器确实输了,然后再拿出 Gluon。把 0.774 TB/s 和 6.590 TB/s 分开的不是算法,而是一行布局描述符,这个事实是把双刃剑 — 选对了能拿到 10 倍,选错了也会丢掉 10 倍。Triton 让你可以不做这个选择。Gluon 把这个选择交还给了你。它究竟是一份礼物还是一张账单,取决于你有没有真的打开过 profiler。

参考资料

현재 단락 (1/170)

Triton 的吸引力很明显。用 Python 写一个 tile 级别的内核,编译器就会替你处理布局、内存分配、数据搬移、同步。多数情况下,这是一笔划算的交易 — 不用降到 CUDA C++ 就能拿到...

작성 글자: 0원문 글자: 11,286작성 단락: 0/170