Triton Gluon AMD 目标 API 全解析：CDNA/RDNA 代际矩阵核、Buffer 原子操作与张量描述符

Triton Gluon AMD 目标 API 全解析：CDNA/RDNA 代际矩阵核、Buffer 原子操作与张量描述符

【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton

Gluon 是 Triton 提供的底层 GPU 编程模型，它把布局（Layout）、共享内存、Warp 特化与目标相关特性直接暴露给开发者，让高级 Kernel 以“牺牲便利换取控制力”的方式获得更高的性能上限。本文以 docs/gluon/api/amd.rst 为骨架，系统梳理 Gluon 中面向 AMD GPU 的全部目标相关 API，涵盖通用布局（AMDMFMALayout、AMDWMMALayout、slice、warp_pipeline_stage）与 CDNA 3/4/5、RDNA 3/4 各代际的矩阵乘法、Buffer 加载/原子操作、异步拷贝与张量描述符（TDM）能力。读完本文，你将掌握在 Gluon 中为 AMD 各代 GPU 编写高性能 Kernel 的完整 API 地图与底层实现依据。

Gluon AMD 模块的代码组织

在进入各 API 细节之前，先理解模块结构。Gluon AMD 相关代码位于仓库的python/triton/experimental/gluon/language/amd/目录，其包入口 python/triton/experimental/gluon/language/amd/init.py 统一导出以下内容：

• 通用布局类：AMDMFMALayout、AMDWMMALayout
• 按代际划分的子模块：cdna3、cdna4、cdna5、rdna3、rdna4，以及别名的gfx1250（即 CDNA 5 对应的 GPU 代号）
• 通用辅助操作：slice、warp_pipeline_stage、get_scaled_upcast_fp4_scale_layout

值得注意的是，gfx1250子包通过 python/triton/experimental/gluon/amd/gfx1250/init.py 直接from ..cdna5 import *转发，印证了文档中 "CDNA 5" 与 gfx1250 架构的对应关系；而顶层python/triton/experimental/gluon/amd/包则对应triton.experimental.gluon.amd，提供TensorDescriptor等宿主端辅助类型（见 python/triton/experimental/gluon/amd/cdna5.py）。

API 参考文档本身（amd.rst）划分为两大块：Common APIs（跨代际通用）与GPU Generations（按架构代际展开的 toctree），与源码的目录划分完全对应。

通用 API：跨代际共享的布局与流水线原语

amd.rst的 "Common APIs" 一节列出了四个通用符号，它们不属于某一个特定 GPU 代际，而是在 CDNA/RDNA 全系列上共同使用。

AMDMFMALayout：CDNA 矩阵核（MFMA）指令布局

AMDMFMALayout定义于 python/triton/experimental/gluon/language/amd/_layouts.py，是一个冻结的DistributedLayout子类，描述 AMD 原生矩阵核单元（MFMA）指令在 CTA 内的数据分布。其字段及约束如下：

从源码的verify()可以看出该布局的硬性校验规则：instr_shape的 M/N 两维只能取四种固定组合，element_bitwidth限定为 32 或 64。_to_ir通过builder.get_amd_mfma_layout(...)把布局转换为 IR，mangle()方法生成布局的唯一名称（如MFMA_3_..._MFMA），用于缓存与代码生成阶段识别。

AMDWMMALayout：WMMA 指令布局

AMDWMMALayout（python/triton/experimental/gluon/language/amd/_layouts.py）描述 AMD WMMA（矩阵乘法累加）指令的数据分布，同样继承DistributedLayout：

与 MFMA 布局不同，WMMA 布局的version覆盖了 RDNA 与 CDNA 两个产品线：RDNA 3/4 与 CDNA 5 均使用该布局描述 WMMA 指令。这解释了为何rdna3、rdna4、cdna5三个子模块都提供wmma操作（见下文）。

slice：寄存器内的张量切片

slice是amd.rst中另一个跨代际通用操作，实现在 python/triton/experimental/gluon/language/amd/slice.py：

slice(source, shape, offsets) -> tensor

它从source中提取指定shape、起始偏移为offsets的纯寄存器视图。关键特性在于：

• 结果保留 source 的分布式布局，因此切片过程不产生跨线程数据搬移（无 layout conversion）；
• 切片范围与偏移必须与 source 布局的 CTA 平铺对齐，否则校验失败；
• 参数校验要求shape、offsets与 source 同秩，且offsets[i] + shape[i] <= source.shape[i]。

底层通过builder.create_extract_slice生成 IR，类型保持不变（distributed_type(source.dtype, shape, source.type.layout)）。该操作在分块 GEMM 中常用于把共享/寄存器缓冲区切成多个 dot 子块。

warp_pipeline_stage：Warp 流水线阶段标记

warp_pipeline_stage（python/triton/experimental/gluon/language/amd/warp_pipeline.py）是一个上下文管理器（with语句），用于在 Gluon Kernel 内标记 Warp 流水线的不同阶段。同一个块内的所有操作属于同一个流水线簇，作为一个整体与其他阶段协调执行。

@gluon.jit def warp_pipelined_kernel(a_ptr, b_ptr, c_ptr, K: gl.constexpr): acc = 0.0 for k in range(0, K): # Stage 0: 预取 tile。 with gl.amd.warp_pipeline_stage("load", priority=3): a = gl.load(a_ptr + k) b = gl.load(b_ptr + k) # Stage 1: 准备 MFMA 操作数。 with gl.amd.warp_pipeline_stage("prep"): a_tile = a # 转换为所需 dot 操作数布局 b_tile = b # Stage 2: 计算。 with gl.amd.warp_pipeline_stage("compute", priority=0): acc += a_tile * b_tile gl.store(c_ptr, acc)

参数说明（来自源码 docstring）：

• label：可选字符串（如"load"、"compute"），仅用于诊断标识，不影响程序语义；
• priority：可选整数，表示执行该阶段的 Warp 的调度优先级，取值范围 0（最低）到 3（最高），与 AMDs_setprio指令的操作数范围一致。省略时若循环内其他阶段显式设置了优先级，则复位为 0；否则不发射优先级指令。

需要强调的是，priority只是对硬件调度器的性能提示，实际效果取决于不同 Warp 指令流之间的动态交互，应谨慎使用。退出with块时，语义层通过builder.create_warp_pipeline_border(marker, prio)插入流水线边界标记。

GPU 代际划分：CDNA 与 RDNA 的能力矩阵

amd.rst的 "GPU Generations" 一节通过 toctree 链接了五个代际子页面，每个代际都有独立的 API 参考文件（均位于 docs/gluon/api/）：

从源码可见，cdna4通过from ..cdna3 import *完整继承 CDNA 3 的 Buffer 能力并叠加新特性（见 python/triton/experimental/gluon/language/amd/cdna4/init.py），体现了代际间的向后兼容设计。

CDNA 3：Buffer 加载/存储与原子 RMW

CDNA 3 模块（python/triton/experimental/gluon/language/amd/cdna3/init.py）提供三类核心能力。

buffer_load / buffer_store：基于标量基址的全局内存访问

buffer_load与buffer_store是 AMD 特有的内存操作：通过一个标量基址指针（ptr）加一个 offsets 张量访问全局内存，而非传递指针张量。buffer_load直接把数据加载进寄存器：

buffer_load(ptr, offsets, mask=None, other=None, cache=None)

参数约定（源码 docstring 与校验逻辑）：

• ptr：指向标量的全局内存基址指针；
• offsets：offsets 张量，元素类型必须是int32或uint32；
• mask：可选谓词张量，用于条件加载；若提供other则mask必填；
• other：可选标量/张量，作为被 mask 掉元素（predicated-off）的填充值；
• cache：可选缓存修饰符（cache modifier）字符串，映射到底层CACHE_MODIFIER。

buffer_store(stored_value, ptr, offsets, mask=None, cache=None)的签名与之对称，且要求 stored_value 与 offsets 可广播，广播后 offsets 形状不得改变。

原子 RMW 操作族：buffer_atomic_*

CDNA 3 提供完整的全局内存原子操作族，全部基于_buffer_atomic_rmw_impl实现：buffer_atomic_max、buffer_atomic_min、buffer_atomic_add、buffer_atomic_and、buffer_atomic_or、buffer_atomic_xor、buffer_atomic_xchg。每个原子操作返回操作前的旧值，mask[i] == 0的元素被跳过（不执行原子）。

_verify_element_type_and_dispatch_op揭示了平台相关约束（源码第 37~71 行）：

• 支持的元素类型：fp16、fp32、bf16、fp64、int32、int64、uint32、uint64；
• 逻辑操作and/or/xor/xchg：仅支持int32/int64；
• max/min：有符号整型与 fp64 走smax/smin，无符号整型走umax/umin，fp16/fp32/bf16 不支持；
• add：无符号整型走iadd，fp16/fp32/fp64 走fadd，bf16 的 fadd 仅 CDNA 4 支持（arch == "cdna4"断言）；
• sem与scope两个可选参数分别控制内存语义与同步范围，默认语义为acq_rel，默认 scope 映射到 GPU（AMD 侧即agent，详见 LLVM 文档中 gfx942 的内存模型）。

mfma 与 scaled_upcast

mfma(a, b, acc)调用 AMD 原生矩阵核单元计算a * b + acc，内部经由语义层的dot操作实现，acc为必填项，累加器类型即返回类型。

scaled_upcast(src, scale, elem_type, axis=None)用于把 fp4/fp8 张量升位，并将原始 E8M0 缩放因子载荷（raw E8M0 payload）折叠进 CDNA 3 的 scaled-upcast 操作中。CDNA 3 上该操作走软件模拟路径，不使用原生硬件 scaled-upcast 指令。scale张量必须用int8/uint8承载 E8M0 载荷，且已具备展开后的输出形状与布局；elem_type必须是fp16或bf16。

CDNA 4：异步拷贝与 Scaled MFMA

CDNA 4 在完整继承 CDNA 3 能力（含全部 buffer 原子操作）的基础上，新增了异步数据移动与低精度缩放矩阵乘法。

async_copy：全局到共享的异步拷贝

CDNA 4 的异步拷贝模块实现在 python/triton/experimental/gluon/language/amd/cdna4/async_copy.py，包含五个 API：

两种 load-to-shared 操作都是异步的，访问dest前必须先调用wait_group。源码还明确指出一个硬件细节：底层指令为每个线程使用独立的全局地址寄存器、但整个 Warp 共享同一份本地地址寄存器，因此使用时必须满足以下条件，否则 LLVM 降级会失败：

• 对ptr/offsets布局，每线程字节数 × 每元素位数必须为 128 或 32（理想性能建议 128 位）；
• 对dest的写入必须合并（coalesced）；
• 若dest使用 swizzle，只能在 Warp 边界内 swizzle。

mfma_scaled 与缩放布局

mfma_scaled(a, a_scale, a_format, b, b_scale, b_format, acc)计算c = a * a_scale @ b * b_scale + acc，操作数采用 OCP MicroScaling（MX）规范定义的格式。可用格式：e2m1、e4m3、e5m2（源码断言强制校验）。前置条件：acc布局必须是AMDMFMALayout，且 a/b 的布局必须是DotOperandLayout且 parent 与累加器布局一致。

配套的get_mfma_scale_layout(dot_operand_layout, shape, scale_factor=32)返回缩放张量的分布式线性布局，目前仅支持 scale_factor=32。

compute_efficient_padded_shared_layout

compute_efficient_padded_shared_layout(dot_operand_layout, shape, dtype, is_k_contig=True)是一个constexpr函数，为即将以 padding 方式写入共享内存的 dot 操作数计算避免 bank 冲突的高效 padded 共享布局。约束要点：

• 父布局必须是 AMDMFMALayout v4（即 CDNA4）；
• shape为共享内存 tile 形状：A 操作数[BM, BK]，B 操作数[BK, BN]；
• dtype位宽仅支持 {4, 8, 16}；打包的 fp4（每字节两个值）应传ttgl.uint8，因为 LDS 层 4 位与 8 位共用 padding 模式；
• is_k_contig指示 K 是否为共享内存中的连续维；
• 当输入超出支持集合时返回None（常见原因：k_width不在 {4, 8, 16}、元素位宽不在 {4, 8, 16}，或指令形状与 kWidth 组合未被底层算法覆盖）。

fp4/fp8 缩放通道

CDNA 4 还提供scaled_upcast、scaled_downcast与load_shared_fp4_repacked。后者把共享内存中 M/N 打包的 fp4 字节加载为 K 打包的 MFMA dot 操作数布局；源共享内存描述符须为int8/uint8打包的 fp4 值，目标形状由源形状与layout中的操作数索引推导。

CDNA 5：张量描述符（TDM）、WMMA 与分区共享布局

CDNA 5（gfx1250）是当前 AMD 侧能力最丰富的代际，对应 python/triton/experimental/gluon/language/amd/cdna5/ 下的五个模块。

tdm：硬件张量描述符与异步数据移动

tdm.py提供 CDNA 5 的 Tensor Descriptor（TDM）机制，核心符号包括tensor_descriptor、tensor_descriptor_type、make_tensor_descriptor、update_tensor_descriptor、async_load、async_load_fused、async_wait、prefetch、async_store、async_scatter、async_gather。

make_tensor_descriptor(base, shape, strides, block_shape, layout)的关键约束（源码第 93~120 行）：

• ndim必须为 1~5 维，strides、block_shape与shape等秩；
• base必须是指针类型；
• layout仅支持PaddedSharedLayout、SwizzledSharedLayout、PartitionedSharedLayout三种共享布局；
• CDNA 5 TDM 张量描述符仅支持零填充（zero padding）。

tensor_descriptor_type的类型定义同时携带block_type、shape_type、strides_type与layout，其mangle()生成形如TD{...}TD的唯一类型名。宿主端triton.experimental.gluon.amd.TensorDescriptor（python/triton/experimental/gluon/amd/cdna5.py）则负责把torch.Tensor包装成描述符（from_tensor），并校验末维连续（strides[-1] == 1）、block shape 合法性与 padding 必须为"zero"。

wmma 与 wmma_scaled

CDNA 5 的wmma(a, b, acc)通过_wmma(3, a, b, acc, ...)分发（version=3 即 gfx1250）。wmma_scaled提供缩放 WMMA，格式与mfma_scaled相同（e2m1/e4m3/e5m2），但额外包含严格的指令形状约束：

• 任一操作数为e2m1时，其 WMMA 布局的instr_shape必须是[16, 16, 64]或[32, 16, 64]；
• 累加器布局的instr_shape必须为[16, 16, 128]；当 a、b 均为e2m1时还额外允许[32, 16, 128]；
• scale 张量 dtype 到格式的映射：float8e4nv→e4m3，int8/uint8→e8m0；合法的 (a_format, b_format, a_scale_format, b_scale_format) 组合由源码中_valid_dtype_combinations集合穷举定义。

配套的get_wmma_scale_layout(dot_operand_layout, shape, scale_factor=32)是 constexpr 函数，scale_factor仅支持 16 或 32，返回基于AMDWMMALayout父布局推导的缩放布局。

PartitionedSharedLayout 与 make_partitioned_dot_layouts

PartitionedSharedLayout（python/triton/experimental/gluon/language/amd/cdna5/_layouts.py）是一种共享内存布局，把张量切分到多个物理共享内存分区中，从而减少 LDS 分区冲突（bank conflict）。字段：num_partitions、num_groups（每个 group 含num_partitions片）、partition_dim（切分维度）、partition_layout（每片的内层布局，如 Swizzled 或 Padded）。校验要求num_partitions与num_groups均为 2 的幂。

make_partitioned_dot_layouts(block_m, block_n, original_layout_a, original_layout_b, num_warps, instr_shape, a_transposed=False, b_transposed=False, slice_m=None, slice_n=None, transposed=True)是面向 CDNA 5 GEMM 的布局推导函数，返回(shared_layout_a, shared_layout_b, wmma_layout)三元组。其关键约定：

• block_m/block_n是共享操作数缓冲区的 tile 尺寸，M 维被切为 2 个分区、每分区至少 2 条指令宽，N 维切为 2 个分区、每分区至少 1 条指令宽；
• num_warps目前仅支持 4 或 8；
• original_layout_a/original_layout_b的连续轴位于 tile 的第 1 维（内存序），因此线性分量统一按order=[1, 0]构建（TDM 额外要求order=[rank-1, ..., 0]）；
• slice_m/slice_n为切分后 dot 的尺寸，默认等于block_m/block_n；transposed=True时返回的 WMMA 结果布局逻辑指令 tile 为N x M。

源码注释还给出了 4-Warp 情况的布局推导图解（warp_bases=[[2, 1], [1, 0]]、reg_bases=[[2, 0]]）：该布局让 w0/w1、w2/w3 在单条指令内读取不同的 A/B 操作数数据块，这是避免 LDS 分区冲突的必要前提；8-Warp 情况下则把重复 tile 交给 w4 而非让 w0 重复执行。

cluster、mbarrier 与 async_copy

• cluster.arrive()/cluster.wait()（python/triton/experimental/gluon/language/amd/cdna5/cluster.py）：用于同步同一 Cluster 内多个 CTA 的簇屏障操作。arrive 与 wait 必须成对出现，先 wait 后 arrive 或连续多次 arrive 都是未定义行为。
• mbarrier：提供 mbarrier 相关操作（含mbarrier_arrive等）。
• async_copy：CDNA 5 继续提供异步拷贝能力（包含mbarrier_arrive等相关实现）。

此外 CDNA 5 也继承了buffer_load/buffer_store（从 cdna3 重导出），并提供load_shared_fp4_repacked（面向 WMMA 的 fp4 重打包加载）与scaled_upcast/scaled_downcast。其中 CDNA 5 的scaled_upcast对 fp4 输入支持两种 scale 布局：扩展 scale（每输出值一个 scale，如 fp4 字节[M, K/2]产生输出[M, K]与 scale[M, K]）与紧凑 scale（每个 scale block 一个 scale，如axis=1且 32 元素 block 时 scale 为[M, K/32]）；fp8 输入则要求axis=None且 scale 与 src 同形状同布局。

RDNA 3 与 RDNA 4：精简的 WMMA 入口

rdna3与rdna4子模块（python/triton/experimental/gluon/language/amd/rdna3/init.py、python/triton/experimental/gluon/language/amd/rdna4/init.py）结构完全相同，各自只导出一个wmma操作，内部通过_wmma(1, ...)（RDNA3）与_wmma(2, ...)（RDNA4）按 version 分发到对应硬件的 WMMA 指令。这对应amd.rdna3.rst/amd.rdna4.rst文档中仅列出wmma一项的事实——消费级 RDNA 架构在 Gluon AMD 目标中的公开面即矩阵乘累加原语，更丰富的数据移动与缩放能力集中在 CDNA 系列。

如何在项目中使用这些 API

所有上述 API 的命名空间为triton.experimental.gluon.language.amd（通常以gl.amd别名导入）。典型导入方式：

from triton.experimental import gluon from triton.experimental.gluon import language as gl # 通用布局与操作 layout_a = gl.amd.AMDMFMALayout(version=3, instr_shape=[16, 16, 16], transposed=False, warps_per_cta=[4, 1]) tile = gl.amd.slice(tensor, shape=[64, 64], offsets=[0, 0]) # CDNA 代际能力 from triton.experimental.gluon.language.amd.cdna4 import buffer_load_to_shared, commit_group, wait_group from triton.experimental.gluon.language.amd.cdna5 import make_tensor_descriptor, make_partitioned_dot_layouts

选择哪个代际模块，取决于目标 GPU：CDNA 数据中心卡（gfx908/gfx90a/gfx942/gfx950/gfx1250）选cdna3/cdna4/cdna5，消费级 RDNA 卡（gfx1100/gfx1101/gfx1200/gfx1201）选rdna3/rdna4。布局版本号与 GPU 的对应关系在 python/triton/experimental/gluon/language/amd/_layouts.py 的两个布局类 docstring 中有明确标注，可作为硬编码依据。

结语

Gluon 的 AMD 目标 API 呈现出清晰的分层设计：AMDMFMALayout/AMDWMMALayout提供跨代际统一的矩阵核布局抽象，buffer_*族提供统一的全局内存访问与原子操作，而cdna4的异步拷贝、cdna5的 TDM 与分区共享布局则逐代叠加硬件新特性。结合 amd.rst 及其五个代际子页面（amd.cdna3.rst、amd.cdna4.rst、amd.cdna5.rst、amd.rdna3.rst、amd.rdna4.rst）与上述源码路径，你可以按代际查表、按需组合，为 AMD GPU 写出既能发挥硬件指令能力、又受编译器严格校验保护的 Gluon Kernel。若需进一步探索完整用法，可参考 Gluon 总览 与 Gluon 教程目录 所列的示例与教程。

【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton