十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Triton Gluon AMD CDNA4 API 完全指南:面向 gfx950 的 MFMA 缩放矩阵乘、异步复制与 MX 格式转换原语

Triton Gluon AMD CDNA4 API 完全指南:面向 gfx950 的 MFMA 缩放矩阵乘、异步复制与 MX 格式转换原语 Triton Gluon AMD CDNA4 API 完全指南面向 gfx950 的 MFMA 缩放矩阵乘、异步复制与 MX 格式转换原语【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton本篇技术指南围绕 Triton 实验性 Gluon 前端中面向 AMD CDNA4gfx950架构的硬件原语模块triton.experimental.gluon.language.amd.cdna4展开系统讲解该模块暴露的异步全局到共享复制async copy、缓冲内存访问buffer load/store、缓冲原子操作、MFMA 矩阵乘与缩放矩阵乘mfma_scaled、E8M0 缩放格式的升/降采样转换scaled_upcast/scaled_downcast以及共享内存填充布局计算等 API 的语义、参数约束与底层实现。读完本文你将掌握在 Gluon 方言中编写面向 CDNA4 的高性能内核所需的全部硬件级原语用法并理解其在 Triton 编译流水线中的落点。1. 模块定位CDNA4 在 Gluon AMD 语言家族中的角色该 API 的权威入口是 docs/gluon/api/amd.cdna4.rst它通过 Sphinxautosummary自动生成triton.experimental.gluon.language.amd.cdna4模块的文档索引其中async_copy使用 docs/_templates/autosummary/gluon-module.rst 模板展开为子模块页面。模块源码位于 python/triton/experimental/gluon/language/amd/cdna4/init.py其设计有两个显著特征继承复用通过from ..cdna3 import *全量继承 CDNA3 的公开 API再通过__all__ [*__cdna3_all, ...]将其一并导出。因此buffer_load、buffer_store、mfma等基础操作在 CDNA4 上可直接使用无需重复实现。CDNA4 专属扩展在继承之上新增了mfma_scaled、scaled_upcast、scaled_downcast、load_shared_fp4_repacked、get_mfma_scale_layout、compute_efficient_padded_shared_layout六个 CDNA4 专属接口以及async_copy子模块。架构版本的对应关系在 python/triton/experimental/gluon/language/amd/_layouts.py 中明确给出AMDMFMALayout.version从 1 到 4 依次对应 gfx908、gfx90a、gfx942CDNA3与 gfx950CDNA4。测试套件 python/test/gluon/test_frontend.py 也印证了这一点HIP_TARGET_CDNA3 GPUTarget(hip, gfx942, 64)HIP_TARGET_CDNA4 GPUTarget(hip, gfx950, 64)。2.async_copy异步全局到共享内存复制async_copy是 CDNA4 模块唯一的子模块源码见 python/triton/experimental/gluon/language/amd/cdna4/async_copy.py导出 5 个接口global_load_to_shared、buffer_load_to_shared、commit_group、wait_group、load_shared_relaxed。2.1 两种异步加载入口global_load_to_shared与buffer_load_to_shared都能将全局内存数据不经寄存器直接异步搬入共享内存但寻址方式不同取舍也截然不同接口寻址方式优势代价global_load_to_shared张量指针每线程 64 位指针灵活支持 64 位索引范围寄存器压力更高无硬件越界掩码buffer_load_to_shared标量基址 32 位偏移张量性能更好支持硬件越界掩码严格限制为 32 位偏移两个接口都遵循相同的调用协议调用后必须commit_group()提交、wait_group(n)等待之后才能安全读取dest。且两者都要求与load、store、buffer_load、buffer_store顺序完成源码明确警告与这些操作交错会损害性能。硬件约束不满足则 LLVM 降级失败源码在global_load_to_shared与buffer_load_to_shared的 docstring 中共同列出三条对指针/偏移布局每线程大小 × 每元素位宽必须是 128 或 32推荐 128 位以获得理想性能对dest的写入必须合并coalesced若dest做了 swizzle只能在 warp 边界内 swizzle。原因在于底层硬件指令为每个线程使用独立寄存器保存全局地址但整个 warp 共享同一个本地内存地址寄存器。2.2 提交与等待commit_group/wait_groupcommit_group()提交一组待处理的异步操作之后可用wait_group等待。wait_group(num_outstanding0)阻塞直到未完成的提交组数 ≤num_outstanding。注意即使num_outstanding为 0未提交的异步操作也会被等待。它被实现为_semantic.builder.create_async_wait_group(num_outstanding)。2.3load_shared_relaxed免等待共享加载load_shared_relaxed(smem, layout)从共享内存加载张量并在结果句柄上设置ttg.amdg.syncedViaAsyncWait属性见 async_copy.py给后端编译器额外提示避免在加载前发出不必要的等待指令——适用于开发者已自行通过wait_group完成同步的场景。典型用法来自 python/test/gluon/test_core.py 的test_amd_direct_load_to_sharedgluon.jit def kernel(smem, a_ptr, b_ptr, blocked): offsets ... cdna4_async_copy.buffer_load_to_shared(smem, a_ptr, offsets) # 或 global_load_to_shared cdna4_async_copy.commit_group() cdna4_async_copy.wait_group(0) a cdna4_async_copy.load_shared_relaxed(smem, blocked) ttgl.store(b_ptr offsets, a)3. 缓冲内存访问buffer_load/buffer_store这两个接口继承自 CDNA3实现见 python/triton/experimental/gluon/language/amd/cdna3/init.py通过标量基址指针 32 位偏移张量访问全局内存是 CDNA 架构上比张量指针寻址更高效的访存方式buffer_load(ptr, offsets, maskNone, otherNone, cacheNone)加载数据到寄存器。cache为可选缓存修饰符经_str_to_load_cache_modifier解析返回与offsets同形状、元素类型为指针元素类型的张量。buffer_store(stored_value, ptr, offsets, maskNone, cacheNone)将张量写回全局内存。参数顺序为先值后指针与常规 store 相反需特别注意cache经_str_to_store_cache_modifier解析。共同校验规则_verify_buffer_opsptr必须是标量指针offsets必须是distributed_type且元素类型为int32或uint32若提供other掩码替换值则mask必填。4. 缓冲原子操作buffer_atomic_*CDNA4 模块继承并重导出了 CDNA3 的全部原子操作buffer_atomic_add、buffer_atomic_and、buffer_atomic_max、buffer_atomic_min、buffer_atomic_or、buffer_atomic_xchg、buffer_atomic_xor统一签名见 cdna4/init.pybuffer_atomic_add(ptr, offsets, value, maskNone, semNone, scopeNone)mask布尔张量mask[i] 0的元素跳过原子操作sem内存语义描述符默认None表示acq_relscope内存同步范围默认None映射到gpuAMDGPU 中称为agent返回值全局内存中该位置操作前的旧值read-modify-write 语义。类型支持矩阵由 cdna3/init.py 的_verify_element_type_and_dispatch_op定义关键规则and/or/xor/xchg仅支持int32/int64max/minint32/int64映射为有符号smax/sminuint32/uint64映射为无符号umax/umin浮点类型不支持adduint32/uint64走整数iaddfloat16/float32/float64走faddbfloat16的fadd仅 CDNA4 支持assert arch cdna4, Buffer atomic fadd with bf16 is only supported on CDNA4 for now.。这是 CDNA4 相对 CDNA3 在原子操作上的唯一差异模块顶部注释也明确写道The cdna4 version additionally supportsfaddwithbf16.5. 矩阵核心运算mfma与mfma_scaled5.1 基础mfmamfma(a, b, acc)计算a * b acc使用 AMD 原生矩阵核心单元实现于 CDNA3cdna3/init.pyacc必填底层复用语义层的dot操作输入精度取knobs.language.fp32_default。5.2 缩放 MFMAmfma_scaledmfma_scaled(a, a_scale, a_format, b, b_scale, b_format, acc)是 CDNA4 专属的微缩放microscaling矩阵乘计算语义为c a * a_scale b * b_scale acc其格式遵循 OCP 组织发布的 OCP Microscaling Formats (MX) Specification 标准E8M0 缩放 低精度数据目前仅在 CDNA4 硬件上受支持。完整实现见 cdna4/init.py。参数格式约束a_format/b_format格式含义e2m1FP4两值打包进一个字节e4m3FP8e5m2FP8布局校验编译期断言acc的布局必须是AMDMFMALayout实例a、b的布局必须是DotOperandLayout且其parent必须与acc的 MFMA 布局一致非法的a_format/b_format会在编译期直接断言失败。底层调用链mfma_scaled→_mma_scaledamd/_ops.py→ 语义层dot_scaled(..., fast_mathFalse, lhs_k_packTrue, rhs_k_packTrue, out_dtypettgl.float32)。_mma_scaled内部完成两项关键工作scale 形状推导K 维经unpack_factore2m1 为 2其余为 1换算后除以scale_factorA 操作数缩放在最后一维B 操作数缩放形状还需交换最后两维默认 scale 生成a_scale/b_scale均可为None——此时按无缩放语义生成默认值fp4e2m1默认uint8值0x7Ffp8 默认float8e4nv值1.0注意 fp4×fp4 场景要求两个操作数使用相同的 scale dtype。scale_factor 由semantic.deduce_scale_factor推导固定为 32见下文get_mfma_scale_layout的断言。5.3get_mfma_scale_layout推导 scale 张量布局get_mfma_scale_layout(dot_operand_layout, shape, scale_factor32)返回缩放张量所需的DistributedLinearLayout。它是constexpr_function最终调用triton._C.libtriton.gluon_ir中的get_amd_mfma_scale_layout。约束与行为只支持scale_factor 32assert scale_factor 32, Only scale factor 32 is supported for CDNA4 Scaled MFMA即每 32 个数据元素共享一个 E8M0 缩放字节依据dot_operand_layout.operand_index、父 MFMA 布局的instr_shape[0]M 维、tiles_per_warp与warps_per_cta计算。测试中python/test/gluon/test_core.py的用法示例a_layout: ttgl.constexpr ttgl.DotOperandLayout(operand_index0, parentmfma_layout, k_width16) a_scale_layout: ttgl.constexpr ttgl.amd.cdna4.get_mfma_scale_layout(a_layout, [M, K // 32]) b_layout: ttgl.constexpr ttgl.DotOperandLayout(operand_index1, parentmfma_layout, k_width16) b_scale_layout: ttgl.constexpr ttgl.amd.cdna4.get_mfma_scale_layout(b_layout, [N, K // 32])6. MX 格式缩放转换scaled_upcast/scaled_downcast这两个接口处理E8M0 缩放字节与 FP4/FP8 数据之间的双向转换是块缩放block-scaled推理与训练内核的核心数据通路。6.1scaled_upcast(src, scale, elem_type, axisNone)将 FP4/FP8 张量与原始 E8M0 scale 载荷折叠进 CDNA4 的原生 scaled-upcast 指令cdna4/init.py。scale必须用int8或uint8承载原始 E8M0 载荷且必须已具备展开后的输出形状与 scaled-upcast 结果布局elem_type必须是fp16或bf16FP4 输入必须显式提供axis打包解包所在维度且输出轴长度是输入的 2 倍每个字节含两个 FP4 值。底层实现_scaled_upcastamd/_ops.py对两种情形分支FP8 输入float8e4nv/float8e5逐元素转换axis必须为None且scale的形状与布局必须与src完全一致走create_scaled_upcast_fp8FP4 输入int8/uint8打包走create_scaled_upcast_fp4并校验轴范围、形状整除关系。6.2scaled_downcast(input, scale, format, axis-1)将 FP16 / BF16 / FP32 除以原始 E8M0 scale 载荷量化为低精度 MX 格式amd/_ops.py。format决定目标类型与打包行为format目标类型打包行为e4m3/e5m2FP8逐元素输出形状与布局不变e2m1打包 FP4沿axis每两个相邻值打包进一个字节偶数元素→低半字节下一元素→高半字节该轴长度减半——恰是scaled_upcast的逆操作scale 布局硬性约束_check_scaled_downcast_scale_shapeamd/_ops.pyscale 必须沿axis紧凑每个缩放块对应连续多个输入元素不支持每元素一个展开 scale每个缩放块必须覆盖8 的倍数个连续输入元素FP4 情形下scale 轴的划分以打包后的输出尺寸为基准。7. 共享内存侧优化原语load_shared_fp4_repacked与compute_efficient_padded_shared_layout7.1load_shared_fp4_repacked(mem_desc, layout)将共享内存中的 M/N 打包 FP4 字节加载为K 打包的 MFMA dot operand 布局cdna4/init.py。约束与行为实现见 amd/_ops.pymem_desc必须是shared_memory_descriptor元素类型为int8/uint8打包 FP4 载荷秩必须为 2 或 3layout必须是DotOperandLayout且operand_index∈ {0, 1}其parent必须为AMDMFMALayout目标形状由源形状与操作数索引自动推导_infer_fp4_repacked_shape对 operand 0K 维减半、打包维翻倍对 operand 1 同理但作用于倒数第二维底层调用builder.create_local_load_packed_transposed生成一条本地加载 打包转置指令。7.2compute_efficient_padded_shared_layout(dot_operand_layout, shape, dtype, is_k_contigTrue)计算避免 bank conflict 的高效填充共享布局是 CDNA4 数据搬移流水线pipeline中为 MFMA 操作数定制共享内存布局的关键工具。约束cdna4/init.pydot_operand_layout的 parent 必须是MFMA v4CDNA4布局assert parent.version 4shape为共享内存 tile 形状operand A 用[BM, BK]operand B 用[BK, BN]dtype位宽仅支持 {4, 8, 16}对打包 FP4每字节两个值应传ttgl.uint8——因为 LDS 层 4 位数据与 8 位共享相同的填充模式is_k_contig指示共享内存中 K 是否为连续维默认True。返回None的常见原因源码 docstring 明示k_width不在 {4, 8, 16}元素位宽不在 {4, 8, 16}或 MFMA 指令形状与 kWidth 组合不在底层算法支持范围内。该函数最终调用triton._C.libtriton.gluon_ir的compute_amd_efficient_padded_shared_layout传入 operand 索引、k_width、父布局版本、warps_per_cta、指令形状、转置标志、tiles_per_warp、元素位宽、cga_layout、目标形状、元素位宽与 K 连续性共 11 个参数。8. 实战形态块缩放矩阵乘内核中的组合运用这些 CDNA4 原语在 python/tutorials/10-block-scaled-matmul.py 中得到完整实战演示。该教程的block_scaled_matmul_kernel_cdna4内核L524展示了完整的数据通路缩放布局推导用get_mfma_scale_layout为 A、B 操作数分别计算 scale 布局A 为[BLOCK_M, K//32]B 为[N, K//32]类形状加载与转换buffer_load加载数据convert_layout转换到 dot operand 布局打包 FP4 处理通过load_shared_fp4_repacked将共享内存中的 M/N 打包 FP4 数据转为 K 打包的 MFMA 操作数布局缩放计算mfma_scaled完成a * a_scale b * b_scale acc。同时该教程提供了配套的shuffle_scales_cdna4L604主机侧函数说明 CDNA4 的 scale 张量在全局内存中需要按mfma_nonkdim16 或 32进行特定重排shuffle后才能被mfma_scaled直接消费——这是硬件布局要求的直接体现。入口处通过is_hip_cdna4()target.arch gfx950判断运行平台且当前教程仅支持mxfp4格式见 L763。9. 验证与测试支撑CDNA4 专属行为在测试套件中有系统性覆盖python/test/gluon/test_core.pytest_amd_direct_load_to_sharedL1826skipif not is_hip_cdna4覆盖global_load_to_shared/buffer_load_to_sharedcommit_group/wait_group/load_shared_relaxed全流程MFMA 测试L1924-L1991对 CDNA4 目标强制mfma版本为 4scaled MFMA 测试L1994参数化覆盖e2m1/e4m3/e5m2三种格式组合并演示get_mfma_scale_layout的调用方式python/test/gluon/test_frontend.py 导入cdna4_async_copy通过 IR filecheck 验证异步复制指令的生成形态测试目标定义HIP_TARGET_CDNA4 GPUTarget(hip, gfx950, 64)L60再次确认 CDNA4 gfx950。10. 小结triton.experimental.gluon.language.amd.cdna4模块是 Triton Gluon 前端面向 AMD gfx950 硬件能力的完整封装层以继承 CDNA3 的方式复用缓冲访存与基础 MFMA同时以六个专属接口与async_copy子模块覆盖 CDNA4 的差异化能力——原生缩放 MFMAmfma_scaledget_mfma_scale_layout、E8M0 缩放格式双向转换scaled_upcast/scaled_downcast、打包 FP4 的布局重排load_shared_fp4_repacked、共享内存填充布局优化compute_efficient_padded_shared_layout、硬件异步复制流水async_copy以及 bf16 的缓冲原子fadd。掌握这些原语即可在 Gluon 方言中直接驱动 CDNA4 的 MX 格式块缩放计算路径构建生产级的高性能内核。【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表