Triton-MUSA
Triton-MUSA 是一套面向 MUSA 后端适配的 Triton 语言与编译器栈,用于使用 Python 编写并运行面向 MUSA 设备的 GPU Kernel。当前 MUSA 已完成 Triton 3.1、Triton 3.2 和 Triton 3.6 三个版本的适配,其中 Triton 3.6 适配已在 MUSA SDK 5.2.0 中正式推出。
概述
什么是 Triton-MUSA
Triton-MUSA(MUSA 后端适配的 Triton 语言与编译器栈)是 Triton 语言和编译器栈的 MUSA 后端适配,允许用户使用 Triton DSL 以 Python 方式编写 GPU Kernel,并将其编译为能够在 MUSA 设备上运行的后端产物。它是一个面向 MUSA 硬件的、以 Python 为中心的 Kernel 编程环境和编译工具链。
现代工作负载经常需要自定义 Kernel 来实现算子融合、带宽优化、特殊布局支持或硬件特性利用。Triton-MUSA 缩小了开发效率与运行性能之间的差距:用户可以使用 Triton 风格的 Python Kernel 表达计算逻辑,同时由 MUSA 后端完成目标相关的编译、运行和调优支持。
当前 Triton-MUSA 的版本适配可以理解为三个阶段:
| 适配版本 | 发布状态 | 后端形态 | 主要定位 |
|---|---|---|---|
| Triton 3.1 | 已适配 | MTGPU 后端 | 支持 MUSA 设备上的基础 Triton Kernel 编写、编译、运行、benchmark 和验证流程 |
| Triton 3.2 | 已适配 | MTGPU 后端 | 在 Triton 3.1 适配基础上增强 MUSA-only 构建、AOT、benchmark 接口、WMMA/SQMMA/TME 测试和后端稳定性 |
| Triton 3.6 | 已随 MUSA SDK 5.2.0 推出 | MUSA 后端 | 面向 Triton 3.6 的 MUSA 后端重构,补充 TensorDescriptor、target info、descriptor/TME、WMMA/SQMMA、AOT/JIT、更多 dtype 和更完整测试覆盖 |
本文档中已有的基础使用方式、安装验证、最小 Kernel、benchmark 和调试排查内容,可以视为 Triton 3.1 与 Triton 3.2 适配现状。 后续涉及 TensorDescriptor、TME、WMMA/SQMMA、MUSA-only 构建、更多 dtype 和更完整测试覆盖的内容,主要对应 Triton 3.6 适配方向。
当前 Triton-MUSA 3.1、3.2 和 3.6 适配版本都会调用 MUSA SDK 中的后端编译器完成目标代码生成,因此对应 Python 发布包需要与 MUSA SDK 版本匹配使用。MUSA SDK 5.3.0 起,Triton-MUSA 计划不再与 MUSA SDK 版本强绑定,届时请以对应版本发布说明为准。
版本适配状态
| 对比项 | Triton 3.1 适配 | Triton 3.2 适配 | Triton 3.6 适配 |
|---|---|---|---|
| 版本基础 | 基于 Triton 3.1.0,Wheel 版本带 MUSA 后缀 | 基于 Triton 3.2.0,Wheel 版本带 MUSA 后缀 | 基于 Triton 3.6.0,面向 MUSA SDK 5.2.0 发布 |
| SDK 依赖 | 需要与对应 MUSA SDK 版本匹配 | 需要与对应 MUSA SDK 版本匹配 | 需要与 MUSA SDK 5.2.0 版本组合匹配;MUSA SDK 5.3.0 起计划解除 SDK 版本强绑定 |
| 构建与发布 | 完成 MUSA 后端基础接入 | 强化 MUSA-only 构建路径,减少非 MUSA 工具链依赖 | 面向 Triton 3.6 工程结构升级,强化 MUSA-only 构建和发布流程 |
| 后端组织 | 以 MTGPU 后端形态承载 MUSA 编译、运行和测试能力 | 延续 MTGPU 后端形态,同时增强编译选项、运行时 benchmark、AOT 和缓存稳定性 | 升级为更清晰的 MUSA 后端形态,系统性增强编译、运行和测试覆盖 |
| 支持架构 | 支持 PH1(MP_31)和 QY2(MP_22) | 支持 PH1(MP_31)和 QY2(MP_22) | 支持 PH1(MP_31)和 QY2(MP_22),其中 QY2 支持相对 PH1 较 弱 |
| 编程模型 | 支持 @triton.jit、tl.load、tl.store、tl.arange、tl.dot、autotune 等基础 Triton 写法 | 在 3.1 基础上补充更多语言语义、模块映射、类型转换和矩阵路径修复 | 增加 tl.target_info.is_musa()、tl.target_info.musa_capability_geq()、TensorDescriptor.from_tensor、tl.load_tensor_descriptor / tl.store_tensor_descriptor 等接口 |
| 矩阵与数据搬运 | 已具备 WMMA、SQMMA、TME 相关基础验证;其中 WMMA 需要通过 ENABLE_MUSA_MMA=1 显式开启,SQMMA 需要通过 MUSA_ENABLE_SQMMA=1 显式开启 | 扩展 PH1 WMMA、SQMMA K16、TME load/store、inplace、transpose、FP8/BF16/TF32 等场景,满足条件时由后端自动选择矩阵加速路径 | 对 descriptor/TME、WMMA/SQMMA、rank-3 dot、int8、FP8/BF16/TF32、转置 operand、accumulator 布局等路径做更完整覆盖,满足条件时由后端自动选择矩阵加速路径 |
| AOT/JIT | 以 JIT 为主,具备 MUSA 后端编译与运行基础 | 增加面向 musa target 的 AOT 编译测试和工具链验证 | AOT 与 JIT 适配到 Triton 3.6 工具链,支持显式 MUSA target 编译、运行时 target 检测和后端产物生成 |
| 测试覆盖 | 覆盖基础语言、runtime、operator、SQMMA/TME 和 WMMA 场景 | 新增 AOT、PH1 WMMA、更多 SQMMA benchmark、inplace 和回归测试 | 补充 target info、TensorDescriptor、dot path coverage、MUSA 后端工具链和 MLIR 级测试 |
关键特性
Triton-MUSA 提供以下能力: