Triton MUSA

Triton MUSA 是一套面向 MUSA 后端适配的 Triton 语言与编译器栈，用于使用 Python 编写并运行面向 MUSA 设备的高性能 GPU Kernel。

概述

什么是 Triton MUSA

Triton MUSA（MUSA 后端适配的 Triton 语言与编译器栈）是 Triton 语言和编译器栈的 MUSA 后端适配。triton_musa 允许用户使用 Triton DSL 以 Python 方式编写 GPU Kernel，并将其编译为能够在 MUSA 设备上运行的二进制产物。它是一个面向 MUSA 硬件的、以 Python 为中心的 Kernel 编程环境和编译工具链。现代工作负载经常需要自定义 Kernel 来实现算子融合、带宽优化、特殊布局支持或硬件特性利用。纯框架算子往往不够，而手写底层 Kernel 的开发和维护成本又很高。triton_musa 缩小了开发效率与运行性能之间的差距。用户可以用 Python 编写 Kernel，而不是直接编写底层设备代码，同时仍然能够获得优化后的后端代码生成结果。

关键特性

triton_musa 提供以下能力：

基于 @triton.jit 的 Python Kernel 编程方式
MUSA / MTGPU 后端支持，并且在多个驱动同时可用时优先选择 MTGPU 驱动
从 Triton Kernel 到后端二进制产物的 JIT 编译流程
对 tl.load、tl.store、tl.arange、tl.dot 以及 launch grid 等标准 Triton 语言构造的支持
通过 triton.autotune、triton.Config、num_warps 和 num_stages 进行性能调优
通过 triton.musa_testing.do_bench 进行 MUSA 性能测试
面向 MUSA 的 AOT 编译路径，可生成 C 启动代码和嵌入式 mubin
面向 MUSA 架构特性的矩阵加速和数据搬运路径，包括 TME、WMMA、WGMMA 风格的 warp-group 矩阵计算模型以及 SQMMA
通过 MLIR 和 LLVM IR dump 开关进行编译器调试

适用读者

本文档适用于：

希望在 MUSA 设备上编写和运行 Triton Kernel 的用户
需要调试、压测或集成 MUSA 定向 Kernel 的开发者

工作原理

工作流程

主要组成部分

前端用户使用 Triton 装饰器和语言 API 在 Python 中编写 Kernel。
编译器 Triton 将 Kernel Lower 到多个中间表示，执行优化 Pass，最终为 MUSA 生成后端代码。
运行时 Runtime 负责选择当前激活的后端驱动并启动已编译的 Kernel。
后端 MUSA 后端提供目标相关的代码生成、启动桩代码以及 AOT 集成能力。

驱动选择

在当前版本中，如果运行时检测到多个驱动同时处于激活状态，会优先选择 MTGPU 驱动。用户通常不需要手动选择驱动，只需要确认 MUSA Runtime、Driver 和 torch_musa 环境可用。

架构

从用户视角看，整个栈可以理解为：

使用 Python 编写 Kernel
通过 JIT 或 AOT 完成编译
在 device="musa" 上分配 Tensor 并发起运行
与 PyTorch 参考实现进行正确性校验
对性能进行测试和调优

环境与前置条件

在使用 triton_musa 之前，需要准备以下环境：

Linux 环境
Python 3.10
正确安装并可用的 MUSA Runtime 和 Driver
torch_musa
与当前 Python 版本、操作系统和硬件环境匹配的 Triton MUSA Wheel 包

推荐先执行以下检查：

python - <<'PY'
import torch
import torch_musa
import triton
print("triton version:", triton.__version__)
print("has torch.musa:", hasattr(torch, "musa"))
print("musa available:", hasattr(torch, "musa") and torch.musa.is_available())
PY

如果还想进一步查看当前 Triton 目标信息：

python - <<'PY'
import torch_musa
from triton.runtime.driver import driver
target = driver.active.get_current_target()
print("backend:", target.backend)
print("arch:", target.arch)
print("warp_size:", target.warp_size)
PY

快速开始

准备安装包

请从产品发布渠道获取与当前环境匹配的 Triton MUSA Wheel 包。Wheel 文件名通常包含 Python ABI、平台和架构信息，例如 cp310、linux_x86_64 等标签。安装前请确认：

Wheel 包对应当前 Python 3.10 环境
torch_musa、MUSA Runtime 和 Driver 版本互相匹配
当前 Python 环境中没有会覆盖该包的上游 triton 安装

安装 Wheel 包

python -m pip install --upgrade pip
python -m pip install ./triton-<version>-<python>-<platform>.whl

如果当前环境中已经安装过其他 Triton 包，建议先确认来源：

python -m pip show triton

如需替换旧版本，可以先卸载后重新安装发布的 Wheel 包：

python -m pip uninstall -y triton
python -m pip install ./triton-<version>-<python>-<platform>.whl

首次功能验证

安装完成后，先验证 import 与 target 检测是否正常：

python - <<'PY'
import triton
import torch
import torch_musa
from triton.runtime.driver import driver

print("triton:", triton.__version__)
print("musa available:", torch.musa.is_available())
print("target:", driver.active.get_current_target())
PY

主要操作与 API

面向用户最常用的 API 包括：

@triton.jit 将 Python 函数标记为 Triton Kernel。
triton.cdiv 用于 launch grid 计算的向上取整除法辅助函数。
triton.autotune 与 triton.Config 用于配置搜索与运行时调优。
triton.testing.do_bench 通用性能测试辅助函数。
triton.musa_testing.do_bench MUSA 专用性能测试辅助函数。
triton.compile 编程方式调用的编译入口。
python -m triton.tools.compile 支持 musa 目标的 AOT 编译工具。
triton.tools.experimental_descriptor.create_1d_tma_descriptor 创建一维 TME/TMA descriptor。
triton.tools.experimental_descriptor.create_2d_tma_descriptor 创建二维 TME/TMA descriptor，常用于 GEMM 的矩阵 tile 搬运。
triton.tools.experimental_descriptor.create_3d_tma_descriptor 创建三维 TME/TMA descriptor。

常用 Triton 语言 API 包括：

tl.program_id
tl.arange
tl.load
tl.store
tl.dot
tl.zeros
tl.constexpr
tl._experimental_descriptor_load
tl._experimental_descriptor_store

示例代码

最小 Vector Add 示例

下面是一个可直接保存并运行的最小示例。请将示例保存为 .py 文件后执行，@triton.jit 需要读取 Kernel 函数源码，不建议从标准输入或临时交互片段中定义 Kernel。

import torch
import torch_musa
import triton
import triton.language as tl


@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(output_ptr + offsets, x + y, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output


x = torch.rand(98432, device="musa", dtype=torch.float32)
y = torch.rand(98432, device="musa", dtype=torch.float32)
out = add(x, y)
ref = x + y
torch.testing.assert_close(out, ref)
print("vector add passed")

Benchmark 示例

ms = triton.musa_testing.do_bench(lambda: add(x, y))
print(f"triton add latency: {ms:.4f} ms")

面向 MUSA 的高级示例入口

如果发布包或配套材料中随附高级示例，可以优先参考其中的 descriptor-based load/store、GEMM benchmark、WMMA、SQMMA 和 AOT 代码生成示例。未随包提供示例时，可先基于本文档的最小 Vector Add 示例确认安装和运行环境。

关键性能数据

Matmul (4096*4096) 性能达到手写汇编的 72%。

备注

该数值仅做参考。Triton MUSA 的实际性能取决于算子形状、数据类型、目标架构、Wheel 包版本以及调优参数。建议按照本文档中的验证和 benchmark 方法在目标环境中实测。

测试与验证

用户侧推荐以安装验证、最小 Kernel 正确性验证和性能基准验证为主。

安装验证

python - <<'PY'
import triton
import torch
import torch_musa
from triton.runtime.driver import driver

print("triton:", triton.__version__)
print("musa available:", torch.musa.is_available())
print("target:", driver.active.get_current_target())
PY

Kernel 正确性验证

运行本文档中的最小 Vector Add 示例，并确认输出包含：

vector add passed

性能验证

在正确性验证通过后，再运行 Benchmark 示例，确认 triton.musa_testing.do_bench 能返回稳定的耗时数据。

验证检查项：

Import 成功
torch.musa.is_available() 返回 True
当前激活目标后端为 musa
Kernel 输出与 PyTorch 参考实现一致
Benchmark 可以正常运行
开启编译器 dump 选项时，能够生成预期的中间结果

调试与故障排查

常用环境变量

这些开关只在需要时启用：

# 始终强制重新编译 Kernel
TRITON_ALWAYS_COMPILE=1

# Dump Triton 中所有 MLIR 中间结果
MLIR_ENABLE_DUMP=1

# Dump Triton 中所有 LLVM IR 中间结果
LLVM_IR_ENABLE_DUMP=1

# Dump MUSA 流程最终 LLVM IR
MUSA_ENABLE_DUMP_LLVM=1

# 用自定义文件替换最终 LLVM IR
cp replace.llir /usr/local/musa/mtgpu/bitcode/replace.llir
MUSA_REPLACE_IR=1

另外一些 Triton 已支持的常用调试开关：

TRITON_INTERPRET=1
TRITON_PRINT_AUTOTUNING=1
MLIR_ENABLE_TIMING=1
LLVM_ENABLE_TIMING=1

常见问题

现象	可能原因	建议处理方式
`torch.musa.is_available()` 为 `False`	Runtime 或 Driver 未准备好	检查 MUSA 安装和 `torch_musa` 环境
当前后端不是 `musa`	Runtime 选择了其他后端	检查环境并打印 `driver.active.get_current_target()`
安装 Wheel 失败	Wheel 标签与 Python 版本或系统平台不匹配	更换与当前环境匹配的 Wheel 包
`import triton` 后版本不符合预期	环境中已有其他 `triton` 包	使用 `python -m pip show triton` 确认安装位置，必要时卸载后重装发布包
IR dump 没变化	命中了缓存，复用了已有编译产物	开启 `TRITON_ALWAYS_COMPILE=1`，必要时清理 Triton cache
上游 Triton 示例里仍然使用 `cuda`	上游示例面向 CUDA 后端	将示例中的 device 字符串和环境依赖改成 MUSA 对应写法

建议的调试流程

推荐顺序：

先验证设备是否可用
再运行一个最小 Kernel
与 PyTorch 参考结果做对比
在正确版本上进行性能测试
只有在需要看编译器行为时才开启 MLIR / LLVM dump
如问题仍然存在，记录 Wheel 版本、MUSA Driver/Runtime 版本、torch_musa 版本和最小复现代码后反馈给支持人员

高级主题

AOT 编译

如果安装包包含 triton.tools.compile 模块，可以通过 python -m triton.tools.compile 使用面向 musa 的 AOT 编译能力。它可以读取用户自己的 Triton Kernel Python 文件，并生成自包含的 C 启动代码以及嵌入式后端二进制数据。

用于 AOT 编译的 Python 文件建议只包含 Kernel 定义和必要的 import，避免在文件顶层执行设备分配或 Kernel launch。--signature 需要按照 Kernel 参数顺序填写；如果 Kernel 中包含 tl.constexpr meta 参数，也需要在对应位置提供常量值。

典型流程：

查询当前激活目标：

python - <<'PY'
import torch_musa
from triton.runtime.driver import driver
t = driver.active.get_current_target()
print(t.backend, t.arch, t.warp_size)
PY

使用显式目标参数编译 Kernel：

python -m triton.tools.compile \
  --kernel-name add_kernel \
  --signature "*fp32,*fp32,*fp32,i32,1024" \
  --grid 1,1,1 \
  --target-backend musa \
  --target-arch <arch> \
  --target-warp-size <warp_size> \
  path/to/kernel.py

MUSA 架构特性概览

Triton MUSA 的用户层编程模型尽量保持 Triton 风格：普通访存使用 tl.load / tl.store，矩阵计算使用 tl.dot，性能调优通过 block size、num_warps、num_stages 和 autotune 配置完成。面向具体 MUSA 架构的 TME、WMMA、WGMMA 风格矩阵计算模型和 SQMMA 由后端在编译阶段选择和降级，用户通常不需要直接写底层指令。

这些能力是否生效取决于当前硬件架构、数据类型、矩阵 tile shape、访存布局以及 Wheel 包版本。建议先通过下面的方式确认当前目标架构：

python - <<'PY'
import torch_musa
from triton.runtime.driver import driver
t = driver.active.get_current_target()
print("backend:", t.backend)
print("arch:", t.arch)
print("warp_size:", t.warp_size)
PY

TME 与 Descriptor Load/Store

TME（Tensor Memory Engine）用于以 descriptor 描述张量 tile 的形状、跨度、元素大小和基地址，再由硬件执行更适合矩阵 tile 的数据搬运。Triton MUSA 中，descriptor-based 工作流通常由两部分组成：

Host 侧创建 descriptor：create_1d_tma_descriptor、create_2d_tma_descriptor 或 create_3d_tma_descriptor
Kernel 侧通过 tl._experimental_descriptor_load / tl._experimental_descriptor_store 读写 descriptor 描述的 tile

这种方式常用于 GEMM、batched GEMM、带特殊布局的矩阵读写，尤其适合与 tl.dot 组成高性能矩阵 Kernel。descriptor 接口属于高级接口，函数名中带有 _experimental 的部分以后可能随版本演进调整；生产代码建议以当前 Wheel 包的发布说明和随包示例为准。

WMMA、WGMMA 与 SQMMA

WMMA、WGMMA 和 SQMMA 都属于矩阵乘加相关的硬件加速路径，但关注层级不同：

WMMA 通常表示 warp 级矩阵计算路径，适合较传统的矩阵 tile 计算。
WGMMA 表示 warp-group 级矩阵计算模型，常见于上游 Triton/Hopper 相关材料。对 Triton MUSA 用户来说，它更适合作为理解“多 warp 协作完成矩阵乘加”的概念；实际 MUSA 后端会根据目标架构选择对应的矩阵加速 lowering。
SQMMA 是 MUSA/MTGPU 后端面向特定架构的矩阵加速路径。对满足条件的 tl.dot，后端可以在编译阶段将其降级到 SQMMA 相关实现，从而利用硬件矩阵乘加能力。

实际开发时，推荐从标准 tl.dot 写法开始，再结合 TME descriptor、合适的 block size 和 autotune 搜索优化性能。不是所有 dtype、tile shape 和 num_warps 组合都会触发同一种矩阵加速路径；如果需要对比或定位矩阵 lowering，可以临时使用 DISABLE_SQMMA=1 或 DISABLE_WMMA=1 关闭对应路径进行 A/B 验证。

Triton MUSA

概述

什么是 Triton MUSA

关键特性

适用读者

工作原理

工作流程

主要组成部分

驱动选择

架构

环境与前置条件

快速开始

准备安装包

安装 Wheel 包

首次功能验证

主要操作与 API

示例代码

最小 Vector Add 示例

Benchmark 示例

面向 MUSA 的高级示例入口

关键性能数据

测试与验证

安装验证

Kernel 正确性验证

性能验证

调试与故障排查

常用环境变量

常见问题

建议的调试流程

高级主题

AOT 编译

MUSA 架构特性概览

TME 与 Descriptor Load/Store

WMMA、WGMMA 与 SQMMA

相关文档

这篇文章对您有帮助吗？

概述​

什么是 Triton MUSA​

关键特性​

适用读者​

工作原理​

工作流程​

主要组成部分​

驱动选择​

架构​

环境与前置条件​

快速开始​

准备安装包​

安装 Wheel 包​

首次功能验证​

主要操作与 API​

示例代码​

最小 Vector Add 示例​

Benchmark 示例​

面向 MUSA 的高级示例入口​

关键性能数据​

测试与验证​

安装验证​

Kernel 正确性验证​

性能验证​

调试与故障排查​

常用环境变量​

常见问题​

建议的调试流程​

高级主题​

AOT 编译​

MUSA 架构特性概览​

TME 与 Descriptor Load/Store​

WMMA、WGMMA 与 SQMMA​

相关文档​

这篇文章对您有帮助吗？

概述

什么是 Triton MUSA

关键特性

适用读者

工作原理

工作流程

主要组成部分

驱动选择

架构

环境与前置条件

快速开始

准备安装包

安装 Wheel 包

首次功能验证

主要操作与 API

示例代码

最小 Vector Add 示例

Benchmark 示例

面向 MUSA 的高级示例入口

关键性能数据

测试与验证

安装验证

Kernel 正确性验证

性能验证

调试与故障排查

常用环境变量

常见问题

建议的调试流程

高级主题

AOT 编译

MUSA 架构特性概览

TME 与 Descriptor Load/Store

WMMA、WGMMA 与 SQMMA

相关文档