TileLang-MUSA
TileLang-MUSA 是 TileLang 在摩尔线程(Moore Threads)MUSA 架构上的深度适配,包含 DSL 编程语言、JIT 编译器、运行时系统及预优化算子库。
概述
什么是 TileLang-MUSA
TileLang-MUSA 是 Tile Language(tile-lang)在 MUSA 架构上的后端适配,下文统一简称 TileLang。Tile Language 是一种简洁的基于 Tile 的领域特定语言,旨在提升高性能 GPU/CPU 算子的开发效率,相关算子包括 GEMM、Dequant GEMM、FlashAttention、LinearAttention 等。TileLang 采用 Python 风格语法,并且基于 TVM 构建底层编译器,使开发者能够在保持开发效率的同时利用实现先进性能所需的底层优化。
除了深度兼容 TileLang 本身的 API 和语言特性之外,TileLang-MUSA 还扩展了一些 MUSA 专属能力,帮助开发者更高效地编写高性能 MUSA 算子。
TileLang-MUSA 通过“分层编程模型 + 硬件感知编译 + 自动调优”三位一体方案,实现 TileLang 在 MUSA 架构上的深度适配:

三层编程接口:
- 入门级:无需了解硬件,用高层张量表达式编写算法;
- 开发者级:使用预定义 Tile Library(如 copy, gemm, reduce)快速构建 kernel;
- 专家级:直接操作线程原语(warp、barrier、shared memory),进行极致优化。
硬件抽象与自动映射:
T.gemm根据目标架构自动转换为 SQMMA、WMMA 或 FMA 路径;T.copy自动转换为 TME(Tensor Memory Extension)或 global↔shared 专用搬运路径;mbarrier原语在 MP31 上可编译为 ASYNC barrier ,支持异步流水。
端到端编译链:源码 → Tile IR → MUSA 后端代码 → 可执行 kernel;复用现有 MUSA 工具链,兼容 nvrtc 式 JIT 流程。
关键特性
| 特性 | 说明 | 开发者价值 |
|---|---|---|
| ✅ 分层编程模型 | Beginner / Developer / Expert 三级接口 | 按需选择抽象层级,兼顾易用性与控制力 |
| ✅ MUSA 原生支持 | 已适配 MP22 与 MP31 两代架构 | 无需重写代码即可跨代际部署 |
| ✅ 硬件特性抽象 | 支持 TME、ASYNC、TCE、LDLMS 等 MP31 特有能力 | 自动利用高带宽异步搬运与 Tensor Core |
| ✅ 丰富算子库 | 覆盖 Elementwise、Reduction、GEMM、FA、MoE 等 50+ 算子 | 快速构建 Transformer、LLM 等模型 |
| ✅ 高性能 | 随仓库提供 benchmark 用例,覆盖 GEMM、FA、Sparse MLA、GDN 等典型 kernel | 便于在目标环境中对齐正确性和实测性能 |
| ✅ 自动调优 | 内置 autotuner 与 layout profiler | 自动搜索最优 tile size、内存布局、循环展开策略 |
架构原理
软硬件架构图

实现逻辑说明:
- 编译路径:基于
codegen_musa.cc实现,新增MUSA TARGET,复用 LLVM 兼容流程; - 原语映射:
copy(src, dst)→ MP31:TME 指令;MP22:专用 LD/ST 指令;gemm(a, b)→ MP31:TCE(Tensor Core Engine);MP22:FMA 序列;mbarrier→ MP31:ASYNC barrier,支持异步数据就绪通知;
- 内存模型:显式 shared memory 分配 + 自动 bank conflict 优化。
版本与环境
环境要求
运行 TileLang 需要准备以下环境:
- 支持 MUSA 架构的 GPU(如 MTTS series)
- 安装 MUSA SDK 5.2.0 版本(需要 MUSA SDK 包含 MCCL 和 muDNN)
- 准备一个可执行
torch_musa的 Python 3.10 环境 - 从源码安装 TileLang 时会一并编译所需的 TVM 组件,无需手动单独安装 TVM
请先准备好可运行 torch_musa 的 Python 环境。torch_musa 的安装,参见 Torch Musa。
版本说明
当前 MUSA SDK 5.2.0 对应的 TileLang-MUSA release 版本为 v0.1.8+musa.3;开源版本和内部 develop 版本为 v0.1.9+musa.1,该版本基于 TileLang 0.1.9 适配 MUSA 平台。
版本差异概览
release_v0.1.8_musa.3 与 develop 分支(v0.1.9+musa.1)的主要差异如下:
| 方向 | release_v0.1.8_musa.3 | develop / v0.1.9+musa.1 |
|---|---|---|
| 基础版本 | 基于 TileLang 0.1.8,MUSA patch 版本为 musa.3 | 基于 TileLang 0.1.9,MUSA patch 版本为 musa.1,同步更多上游 TileLang 语法和运行时能力 |
| MUSA copy / TME | 支持 MUSA T.copy、robust copy、async copy 和 TME 相关能力 | 完善 TMA/TME copy 路径,补充 T.tma_copy 语义;手动 barrier 使用 T.tma_copy(..., barrier=...),T.copy 不再接受 barrier= 参数 |
| TME cache hint | 支持 descriptor 形式 TME load/store cache hint | 明确 T.copy 与 T.tma_copy 上 cache policy hint 的适用范围 |
| GEMM / SQMMA / WMMA | 支持 MP31 SQMMA、WMMA、FMA lowering 及 TileLang-MUSA GEMM 示例 | 继续完善 GEMM lowering、SQMMA 指令选择和 wg_wait 等同步控制;补充 MP22 FP16 GEMM M16N16K16、M8N32K16 等形状测试 |
| Pipeline / warp-specialize | 支持基础 pipeline、barrier 和 producer 线程控制 | 完善 producer-consumer warp-specialize 与 pipeline planning,增强 async copy、mbarrier、fallback barrier 等路径稳定性 |
| 语言与算子 API | 支持主流 TileLang 语法和 MUSA 扩展接口 | 新增或完善 T.tma_copy、T.tma_store、T.transpose、warp vote、buffer pointer、min blocks per SM、函数属性等语言能力 |
| MUSA 模板与算子 | 覆盖 atomic、copy、GEMM、reduce 等 MUSA 模板 | 增加 packed vector、FP4、MP22 MMA 扩展、MUSA fused/accelerated pattern 等模板和 intrinsic 支持 |
| JIT / cache / runtime | 支持源码安装、JIT 编译、kernel cache 和 MUSA 后端 codegen | 改进 cache 原子写入、后端 CMake、日志、C host codegen、Cython wrapper 和 TVM FFI 相关路径 |
| 测试与 benchmark | 已有 GEMM、FA、DSA、TileKernels 等测试和 benchmark | 扩展 testing 覆盖面,新增 MP22/MP31 MUSA kernel、language、transform、cache、autotune、profiler 等测试,并补充 FA、DSA、GDN、Sparse MLA 等 benchmark 更新 |
安装
安装 Python 包
如已按 安装 MUSA SDK 5.2.0 完成 Python 包安装准备,可直接安装当前 TileLang-MUSA 版本。
python -m pip uninstall -y tilelang-musa tilelang_musa
python -m pip install 'tilelang_musa==0.1.8+musa.3'
安装完成后仍通过 import tilelang 使用。
python - <<'PY'
import tilelang
print("tilelang:", tilelang.__version__)
PY
从源码安装(可选)
如需开发、调试或使用未发布特性,可在已经准备好的 torch_musa Python 3.10 环境中从源码安装。
git clone https://github.com/tile-ai/tilelang-musa.git --recursive
cd tilelang-musa
python -m pip install -r ./requirements-dev.txt
export MUSA_HOME=/usr/local/musa
export USE_MUSA=1
python -m pip install -e . -v --no-build-isolation
安装完成后可以 执行以下命令验证 Python 包是否可导入:
python -c "import tilelang; print(tilelang.__version__)"
Elementwise Add
以下示例使用当前推荐的 @tilelang.jit 写法。
elementwise_add.py
import tilelang
import tilelang.language as T
import torch
# 禁用缓存
tilelang.disable_cache()
@tilelang.jit
def elementwise_add(A, B, num_per_thread=8, threads=256, dtype="float32"):
N = T.const("N")
A: T.Tensor[(N,), dtype]
B: T.Tensor[(N,), dtype]
C = T.empty((N,), dtype)
# 设置 grid size 和 block thread number
with T.Kernel(T.ceildiv(N, threads * num_per_thread), threads=threads) as b_x:
# thread block 级别执行代码
for i, j in T.Parallel(threads, num_per_thread):
offsets = (b_x * threads + i) * num_per_thread
C[offsets + j] = A[offsets + j] + B[offsets + j]
return C
def ref_program(x, y):
return x + y
N = 4096
# 编译 TileLang kernel
kernel = elementwise_add.compile(N=N)
# 打印 MUSA C 代码
print(kernel.get_kernel_source())
# Run kernel and compare with torch reference.
a = torch.randn(N, dtype=torch.float32, device='musa')
b = torch.randn(N, dtype=torch.float32, device='musa')
c = kernel(a, b)
torch.testing.assert_close(c, ref_program(a, b), rtol=1e-2, atol=1e-2)
使用 python elementwise_add.py 执行上面代码,可以看到生成的 MUSA C 代码,以及结果比对正确
已经测试过的用例在 testing 目录下,可以选取单个 UT 用 Python 执行
编程接口
语法
关于具体语法信息,可参考 TileLang 官网 https://tilelang.com/programming_guides/language_basics.html。TileLang-MUSA 已支持绝大多数 TileLang 语法。
类型系统
目前除了 float6 float4 之外,其余都支持。具体的类型系统,请参见 TileLang 官方文档。
控制流
支持全部 控制流语句,参考:
for i in T.serial(N):
...
for i, j in T.Parallel(M, N):
C[i, j] = A[i, j] + B[i, j]
for k in T.Pipelined(T.ceildiv(K, BK), num_stages=3):
# overlap copy/compute across stages
...
不支持语法
TileLang-MUSA 不支持 TileLang 中与 NVIDIA 硬件特性强相关的 IR,例如:
alloc_tmem
alloc_descriptor
alloc_wgmma_desc
alloc_tcgen05_smem_desc
alloc_tcgen05_instr_desc
TileLang-MUSA 也不支持以下 IR:
set_max_nreg
inc_max_nreg
dec_max_nreg
no_set_max_nreg
annotate_producer_reg_dealloc
annotate_consumer_reg_alloc
disable_warp_group_reg_alloc
性能分析
支持 TileLang Profiler 性能分析工具。建议先使用 testing 目录中的用例验证正确性,再使用 benchmark 目录中的脚本在目标设备上测量 latency、throughput 或带宽。
kernel = tilelang.compile(
program,
out_idx=-1,
target=TARGET,
execution_backend='cython',
)
profiler = kernel.get_profiler()
latency_ms = profiler.do_bench(
n_warmup=n_warmup,
n_repeat=n_repeat,
backend=backend,
quantiles=quantiles,
return_mode=return_mode,
)
print(f'Latency: {latency_ms} ms')
示例代码
GEMM
import tilelang
import tilelang.language as T
import torch
tilelang.disable_cache()
@tilelang.jit
def matmul(A, B, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
M, N, K = T.const("M N K")
A: T.Tensor[[M, K], dtype]
B: T.Tensor[[K, N], dtype]
C = T.empty((M, N), dtype)
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=512) as (bx, by):
T.use_swizzle(panel_size=4, order="col")
A_shared = T.alloc_shared((block_M, block_K), dtype)
B_shared = T.alloc_shared((block_K, block_N), dtype)
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(C_local)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
T.copy(A[by * block_M, k * block_K], A_shared)
T.copy(B[k * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local, policy=T.GemmWarpPolicy.Square)
T.copy(C_local, C[by * block_M, bx * block_N])
return C
def main():
M, N, K, BLOCK_M, BLOCK_N, BLOCK_K = 1024, 1024, 1024, 256, 256, 64
kernel = matmul.compile(
M=M,
N=N,
K=K,
block_M=BLOCK_M,
block_N=BLOCK_N,
block_K=BLOCK_K,
)
A = torch.randn((M, K), dtype=torch.float16, device="musa")
B = torch.randn((K, N), dtype=torch.float16, device="musa")
C = kernel(A, B)
ref_out = torch.mm(A, B)
torch.testing.assert_close(ref_out.to(torch.float32), C.to(torch.float32), rtol=1.25e-1, atol=1.25e-1)
if __name__ == "__main__":
main()
更多 Kernel
Flash Attention3、Sparse MLA、GDN 等更多 kernel 示例代码,参见 TileLang-MUSA 仓库的 testing 目录;对应性能测试和端到端形状参考可查看 benchmark 目录。
TileLang-MUSA 扩展(v0.1.8+musa.3 / v0.1.9+musa.1)
以下从接口和使用语义角度,对比 v0.1.8+musa.3 与 v0.1.9+musa.1 的 MUSA 扩展差异。
| 能力项 | v0.1.8+musa.3 | v0.1.9+musa.1 |
|---|---|---|
T.Kernel(producer_threads=...) | 支持通过 producer_threads 指定 warp-specialize 场景下的 producer 线程数 | 保持支持 |
T.copy(force_async_copy=True) | 支持显式请求 global→shared async copy,用于在不使用 TME 的场景下使能 ldlms 路径 | 保持支持,并补充 async copy 触发条件说明 |
T.copy(src_robust_desc=desc) | 支持通过 robust descriptor 描述源地址有效字节范围 | 保持支持,可与 force_async_copy=True 组合 |
| 手动 TME/TMA barrier | 可通过 copy 相关路径指定 barrier,手动控制 shared barrier 同步 | 调整为 T.tma_copy(..., barrier=barrier);T.copy 不再接受 barrier= 参数 |
| TME cache policy hint | 支持 descriptor 形式 TME load/store 的 inner_cache_policy / outer_cache_policy / eviction_policy | 进一步明确 T.copy(...) 和 T.tma_copy(...) 的适用范围 |
T.gemm(...) lowering | 支持 gemm_ss 按 SQMMA -> WMMA -> FMA 优先级 lowering,gemm_rr 按 WMMA -> FMA 优先级 lowering | 保持支持,并补充 wg_wait=-1 与 T.wait_sqmma() 的使用说明 |
| SQMMA layout 标注 | 支持 make_sqmma_swizzled_layout 配合 T.annotate_layout 标注 shared operand layout | 保持支持,并补充 allow_buffer_region、allow_reannotation 使用场景 |
| PassConfig | 支持控制 SQMMA/WMMA lowering 和安全访问相关 pass | 补充 TL_DISABLE_INDEX_TYPE_PROMOTION 等 index 类型提升控制说明 |
accelerated_ops | release 文档中未单独成文 | 新增 MUSA accelerated ops 文档,包含 T.mul_half_float_to_bfloat16_x4 |
Copy 扩展
force_async_copy=True 用于显式请求将 global -> shared 的 T.copy lowering 成 async copy。使用时源 buffer 必须是 global,目标 buffer 必须是 shared,且源和目标 dtype 保持一致;单条 async copy 的字节数通常需要是 4、8 或 16 字节,可以通过 T.vectorized(...) 或连续 BufferRegion slice 写法形成对应宽度的 copy。
for v in T.vectorized(4):
T.copy(src_global[v], dst_shared[v], force_async_copy=True)
src_robust_desc 用于描述源地址的有效字节范围,适合源侧可能越界的 copy 场景。descriptor 通过 T.make_robust_desc(addr, size_bytes) 创建,其中 addr 通常由 T.address_of(...) 获取。
robust_desc = T.make_robust_desc(T.address_of(src_global[1]), 8)
T.copy(src_global[tid], dst_shared[tid], src_robust_desc=robust_desc)
对于能够 lowering 成 TME/TMA load 的搬运,可以手动指定 shared barrier,并在消费目标 shared buffer 前显式完成 barrier arrive / wait。v0.1.9+musa.1 中该能力通过 T.tma_copy(..., barrier=barrier) 使用;T.tma_copy 的 load 侧只发起数据搬运,不会像普通 T.copy 那样自动完成同步,且 T.copy 不再接受 barrier= 参数。
barrier = T.alloc_barrier(128)
T.tma_copy(src[0], tile, barrier=barrier)
T.barrier_arrive(barrier)
T.barrier_wait(barrier, 0)
descriptor 形式的 MUSA TME load/store 支持 cache policy hint。inner_cache_policy 和 outer_cache_policy 可选值包括 "cache_none"、"cache_once"、"cache_normal"、"cache_persist";也可以使用兼容写法 eviction_policy="evict_first" 等。eviction_policy 不能和 inner_cache_policy / outer_cache_policy 同时设置。
T.copy(
A_global[0:block_m, 0:block_n],
A_shared,
inner_cache_policy="cache_once",
outer_cache_policy="cache_normal",
)
barrier = T.alloc_barrier(128)
T.tma_copy(
A_global[0:block_m, 0:block_n],
A_shared,
barrier=barrier,
inner_cache_policy="cache_once",
outer_cache_policy="cache_normal",
)
GEMM 扩展
MUSA 后端会根据 operand scope、shape、dtype 和线程布局选择合适的 GEMM lowering 路径。A/B operand 都是 shared buffer 时属于 gemm_ss,优先级为 SQMMA -> WMMA -> FMA;A/B operand 都是 local fragment 时属于 gemm_rr,优先级为 WMMA -> FMA,MUSA SQMMA 不支持 gemm_rr。
A_shared = T.alloc_shared((block_M, block_K), dtype)
B_shared = T.alloc_shared((block_K, block_N), dtype)
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.gemm(A_shared, B_shared, C_local, policy=T.GemmWarpPolicy.Square)
对于 SQMMA GEMM,可以通过 wg_wait=-1 表示发起 SQMMA 后不立即等待。使用该模式时,必须在读取、copy 或继续计算 accumulator 前调用 T.wait_sqmma();T.wait_sqmma() 与 T.wait_wgmma() 等价。
T.gemm(A_shared, B_shared, C_local, wg_wait=-1)
T.wait_sqmma()
SQMMA Layout 标注
tilelang.layout.make_sqmma_swizzled_layout(buffer, continuity=None, k_major=True) 用于为 MUSA SQMMA shared operand 生成 swizzled layout,并通过 T.annotate_layout(...) 绑定到 shared buffer 或 BufferRegion。该接口只生成 layout 描述,不会自动搬运或重排 shared memory 中已有数据。
KV_shared = T.alloc_shared((block_m, block_k), T.float16)
T.annotate_layout({
KV_shared: tilelang.layout.make_sqmma_swizzled_layout(KV_shared, k_major=True)
})
当需要给 BufferRegion 标注 layout 时,需要设置 allow_buffer_region=True;当同一个 shared buffer 在不同阶段需要重新标注 layout 时,可以使用 allow_reannotation=True。
PassConfig 常用选项
以下选项适合在用户已明确保证访问范围或希望控制 lowering 路径时使用: