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 更新

安装

配置 pip 软件源

请根据使用场景，选择以下任一种方式配置 pip 软件源：

为当前 shell 会话或 CI 任务配置软件源：

export PIP_INDEX_URL=https://dl.mthreads.com/repo/api/pypi/pypi/simple

将软件源写入 pip 配置文件，并查看当前配置：

python -m pip config set global.index-url \
  https://dl.mthreads.com/repo/api/pypi/pypi/simple
python -m pip config list

不修改 pip 配置，在安装时指定软件源：

python -m pip uninstall -y tilelang-musa tilelang_musa
python -m pip install 'tilelang_musa==0.1.8+musa.3' \
  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple

备注

请勿在同一次安装中同时使用摩尔线程 Python 包源和公开 PyPI 镜像源。安装普通第三方包时，请切换软件源后单独安装。

安装 Python 包

使用环境变量或 pip 配置文件完成配置后，可通过 pip 直接安装 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()

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 路径时使用：

配置项	说明
`tilelang.PassConfigKey.TL_DISABLE_SQMMA`	禁用 gemm_ss lowering 到 SQMMA
`tilelang.PassConfigKey.TL_DISABLE_PH1_WMMA`	禁用 PH1 WMMA lowering
`tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS`	关闭后续自动补上的通用 global memory 防越界保护
`tilelang.PassConfigKey.TL_DISABLE_SAFE_COPY_PREDICATION`	关闭 `T.copy` lowering 过程中自动生成的安全谓词
`tilelang.PassConfigKey.TL_DISABLE_SAFE_ROBUST_COPY_PREDICATION`	关闭 MUSA robust async copy 最终附带的安全谓词
`tilelang.PassConfigKey.TL_DISABLE_INDEX_TYPE_PROMOTION`	关闭 index 表达式自动提升到更宽整数类型

pass_configs = {
    tilelang.PassConfigKey.TL_DISABLE_SQMMA: True,
    tilelang.PassConfigKey.TL_DISABLE_SAFE_COPY_PREDICATION: True,
}
kernel = tilelang.compile(program, target="musa", pass_configs=pass_configs)

警告

关闭安全访问相关 PassConfig 后，如果访问范围、tail tile 或 flattened offset 没有由用户 guard、launch geometry 或 robust descriptor 保证，可能导致越界访问或错误结果。

Accelerated Ops

accelerated_ops 用于暴露 MUSA target 上可以高效执行的计算 pattern，定位更接近“显式请求使用某类硬件友好的 fused/accelerated 计算形式”。例如 T.mul_half_float_to_bfloat16_x4(x, y) 表示将 float16x4 * float32x4 的乘法结果转换为 bfloat16x4。普通表达式 T.Cast("bfloat16", A[i] * B[i]) 在满足 pattern 和类型约束时，也可能由编译器自动识别并优化生成对应操作。

该接口在 v0.1.9+musa.1 中补充为 MUSA accelerated ops 文档的一部分。显式调用时需要传入 4-lane vector 表达式，推荐使用 slice 或 T.Ramp 写法；lhs 必须是 float16x4，rhs 必须是 float32x4，返回值为 bfloat16x4。

C[offset : offset + 4] = T.mul_half_float_to_bfloat16_x4(
    A[offset : offset + 4],
    B[offset : offset + 4],
)

开发者辅助工具

TileLang-MUSA 正在建设 tilelang-musa-skills 辅助开发能力集合。开发者可使用相关 skill 辅助完成 MUSA kernel 编写、CUDA/TileLang kernel 迁移、编译与运行问题定位、正确性验证、benchmark、性能分析、性能优化以及前向/反向/autograd kernel 测试等工作。

相关 skill 正在开源中，敬请期待。

调试诊断

调试打印

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

也可以通过 pass config 将 IR dump 到固定目录：

kernel = tilelang.compile(
    program,
    target="musa",
    pass_configs={
        "tl.enable_dump_ir": True,
        "tl.dump_ir_path": "/tmp/tilelang_ir",
    },
)

常用环境变量

TileLang-MUSA 支持通过环境变量控制缓存、编译输出、MUSA 源码 dump 和默认 JIT 参数。布尔变量通常接受 1、true、yes、on 作为启用值。

变量	用途
`MUSA_HOME` / `MUSA_PATH`	指定 MUSA 安装目录，默认会尝试从 `mcc` 路径或 `/usr/local/musa` 探测
`TILELANG_CACHE_DIR`	指定 kernel cache 根目录，默认 `~/.tilelang/cache`
`TILELANG_DISABLE_CACHE`	全局禁用 kernel cache
`TILELANG_PRINT_DEVICE_COMPILE_COMMAND`	打印 `mcc` 等设备编译命令
`TILELANG_REPLACE_MUSAC`	使用指定 MUSA C 源码替换生成源码并重新编译 wrapper library
`TILELANG_MCC_EXTRA_ARGS`	追加到 MUSA wrapper library 的 `mcc` 参数，多个参数用逗号分隔
`TILELANG_DUMP_KERNEL_SOURCE`	dump 生成的 MUSA kernel source，并输出对应编译命令
`TILELANG_TARGET`	设置默认编译 target，例如 `musa`
`TILELANG_VERBOSE`	开启默认 verbose 编译输出

示例：

TILELANG_TARGET=musa TILELANG_VERBOSE=1 python your_script.py
TILELANG_PRINT_DEVICE_COMPILE_COMMAND=1 python your_script.py
TILELANG_MCC_EXTRA_ARGS="-save-temps,-v" python your_script.py

dump MUSA kernel source：

TILELANG_DUMP_KERNEL_SOURCE=/tmp/kernel.mu python your_script.py

如果希望把生成的 MUSA kernel source 和目标产物固定输出到目录，可以通过 pass config 设置：

kernel = tilelang.compile(
    program,
    target="musa",
    pass_configs={
        "musa.kernels_output_dir": "/tmp/tilelang_musa_kernels",
    },
)

名词解释

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

TileLang-MUSA

概述

什么是 TileLang-MUSA

关键特性

架构原理

软硬件架构图

版本与环境

环境要求

版本说明

版本差异概览

安装

配置 pip 软件源

安装 Python 包

从源码安装（可选）

Elementwise Add

编程接口

语法

类型系统

控制流

不支持语法

性能分析

示例代码

GEMM

更多 Kernel

TileLang-MUSA 扩展（v0.1.8+musa.3 / v0.1.9+musa.1）

Copy 扩展

GEMM 扩展

SQMMA Layout 标注

PassConfig 常用选项

Accelerated Ops

开发者辅助工具

调试诊断

调试打印

查看 MUSA 代码

查看 IR

常用环境变量

名词解释

这篇文章对您有帮助吗？

概述​

什么是 TileLang-MUSA​

关键特性​

架构原理​

软硬件架构图​

版本与环境​

环境要求​

版本说明​

版本差异概览​

安装​

配置 pip 软件源​

安装 Python 包​

从源码安装（可选）​

Elementwise Add​

编程接口​

语法​

类型系统​

控制流​

不支持语法​

性能分析​

示例代码​

GEMM​

更多 Kernel​

TileLang-MUSA 扩展（v0.1.8+musa.3 / v0.1.9+musa.1）​

Copy 扩展​

GEMM 扩展​

SQMMA Layout 标注​

PassConfig 常用选项​

Accelerated Ops​

开发者辅助工具​

调试诊断​

调试打印​

查看 MUSA 代码​

查看 IR​

常用环境变量​

名词解释​

这篇文章对您有帮助吗？

概述

什么是 TileLang-MUSA

关键特性

架构原理

软硬件架构图

版本与环境

环境要求

版本说明

版本差异概览

安装

配置 pip 软件源

安装 Python 包

从源码安装（可选）

Elementwise Add

编程接口

语法

类型系统

控制流

不支持语法

性能分析

示例代码

GEMM

更多 Kernel

TileLang-MUSA 扩展（v0.1.8+musa.3 / v0.1.9+musa.1）

Copy 扩展

GEMM 扩展

SQMMA Layout 标注

PassConfig 常用选项

Accelerated Ops

开发者辅助工具

调试诊断

调试打印

查看 MUSA 代码

查看 IR

常用环境变量

名词解释