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 算子。

TileLang-MUSA 通过 分层编程模型 + 硬件感知编译 + 自动调优 三位一体方案实现 TileLang 在 MUSA 架构的深度适配：

三层编程接口：

• 入门级：无需了解硬件，用高层张量表达式编写算法；
• 开发者级：使用预定义 Tile Library（如 copy, gemm, reduce）快速构建 kernel；
• 专家级：直接操作线程原语（warp、barrier、shared memory），进行极致优化。

硬件抽象与自动映射：

• T.gemm 根据不同的架构自动转换为 SQMMA 指令或 WMMA 指令，
• T.copy 自动转换为 TME（Tensor Memory Extension）指令或 global↔shared 专用搬运指令，
• mbarrier 原语在 MP31 上编译为 ASYNC barrier，支持异步流水。

端到端编译链：源码 → Tile IR → MUSA ASM → 可执行 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 等模型
✅ 高性能	GEMM 达手写最佳实现的 95%，FA3 达到手写最佳实现的 85%；具体数值受测试条件影响	接近手写汇编性能
✅ 自动调优	内置 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.1.0 版本（需要 MUSA SDK 包含 mccl 和 muDNN）
• 准备一个可执行 torch_musa 的 Python 3.10 环境
• 从源码安装 TileLang 时会一并编译所需的 TVM 组件，无需手动单独安装 TVM

note

请先准备好可运行 torch_musa 的 Python 环境。torch_musa 的安装，参见 Torch Musa。

安装

从源码安装

以下命令默认在已经准备好的 torch_musa Python 3.10 环境中执行。

git clone https://github.com/tile-ai/tilelang-musa.git --recursive
cd tilelang-musa
pip install -r ./requirements-dev.txt
export USE_MUSA=1
pip install -e . -v --no-build-isolation

执行 TileLang 程序

以下示例使用当前推荐的 @tilelang.jit 写法。

示例 1：Elementwise Add

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 代码，以及结果比对正确

已经测试过的用例在 musa_tests 目录下，可以选取单个 UT 用 Python 执行

TileLang 语法与API

语法

关于具体语法信息，可参考 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

API

性能分析

支持 TileLang Profiler 性能分析工具。

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')

TileLang 示例代码

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()

Flash Attention3

参考 musa_tests/flash_attention/example_mha_fwd_bhsd.py

TileLang MUSA 扩展

关联接口	扩展内容	使用指导
T.kernel(producer_threads=128)	新增 producer_threads 参数	开启 warp_specialize 特性后，针对 producer 中存在 SIMT Copy，可以通过 producer_threads 指定生产者线程数(默认和 consumer 线程数相等)
T.copy(force_async_copy=True)	新增 force_async_copy 参数	在不使用 TME 特性的前提下，可以通过 force_async_copy 使能 ldlms 特性

调试诊断

调试打印

TileLang print op

TileLang 提供了丰富的 print op 用于打印 kernel 内部变量，参考

def debug_print_buffer(M=16, N=16, dtype='float16'):

    @T.prim_func
    def program(Q: T.Tensor((M, N), dtype)):
        with T.Kernel(4, 4, 2, threads=128 * 2) as (bx, by, bz):
            shared_buf = T.alloc_shared([M, N], dtype)
            T.print(shared_buf)

    jit_kernel = tilelang.compile(program, target='musa')
    profiler = jit_kernel.get_profiler()
    profiler.run_once()

def _manual_device_assert_triggered():

    @T.prim_func
    def program():
        with T.Kernel(threads=128):
            tid = T.get_thread_binding()
            T.device_assert(tid > 0, 'Assertion Trigger !')

    jit_kernel = tilelang.compile(program, target='musa')
    profiler = jit_kernel.get_profiler()
    profiler.run_once()

查看 MUSA 代码

TileLang 的编译过程会将 TileLang DSL 转换为 MUSA C 代码，然后调用 MUSA 后端编译器编译成二进制文件，可以通过 get_kernel_source 或者 show_source 接口打印 MUSA C 代码，参考

kernel = tilelang.compile(program, execution_backend='cython')
print(kernel.get_kernel_source())

kernel = sparse_mla_fwd(
    heads,
    dim,
    tail_dim,
    topk,
    kv_group,
    sm_scale,
    is_casual,
    block_I=block_I,
    num_stages=num_stages,
    threads=threads,
)
kernel.show_source()

查看 IR

TileLang 使用 TVM TIR。如果想要查看 Pass 前后的 IR，可以使用 TVM Instruments 机制；源码安装完成后可直接使用，无需再单独安装 TVM。

from tvm.ir.instrument import PrintAfterAll

kernel = tilelang.compile(
    program, execution_backend='cython', instruments=[PrintAfterAll()]
)

性能说明

以下数据为参考测试结果，实际性能取决于目标架构、数据类型、算子形状、编译配置以及运行时环境。建议基于本文档示例或随发布包提供的 benchmark 在目标环境中实测。

性能实测数据（MP31 架构，FP16）

算子	TileLang-MUSA
GEMM (4096×4096)	95% asm（参考测试结果）
FlashAttention (seq=2048)	85% asm（参考测试结果）

名词解释

术语	说明
MUSA	Metaverse Computing Unified System Architecture，摩尔线程统一系统架构
Tile	计算分块单位，通常为 32×32 或 128×128
Warp	MUSA 中的基本执行单元，每个 Warp 含 32 个线程
Tensor Core	MUSA 上的专用矩阵计算单元，支持 FP16/BF16 MMA 指令