Skip to main content

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 原生支持已适配 MP22MP31 两代架构无需重写代码即可跨代际部署
硬件特性抽象支持 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
note

请先准备好可运行 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.3develop 分支(v0.1.9+musa.1)的主要差异如下:

方向release_v0.1.8_musa.3develop / 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.copyT.tma_copy 上 cache policy hint 的适用范围
GEMM / SQMMA / WMMA支持 MP31 SQMMA、WMMA、FMA lowering 及 TileLang-MUSA GEMM 示例继续完善 GEMM lowering、SQMMA 指令选择和 wg_wait 等同步控制;补充 MP22 FP16 GEMM M16N16K16M8N32K16 等形状测试
Pipeline / warp-specialize支持基础 pipeline、barrier 和 producer 线程控制完善 producer-consumer warp-specialize 与 pipeline planning,增强 async copy、mbarrier、fallback barrier 等路径稳定性
语言与算子 API支持主流 TileLang 语法和 MUSA 扩展接口新增或完善 T.tma_copyT.tma_storeT.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.3v0.1.9+musa.1 的 MUSA 扩展差异。

能力项v0.1.8+musa.3v0.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_ssSQMMA -> WMMA -> FMA 优先级 lowering,gemm_rrWMMA -> FMA 优先级 lowering保持支持,并补充 wg_wait=-1T.wait_sqmma() 的使用说明
SQMMA layout 标注支持 make_sqmma_swizzled_layout 配合 T.annotate_layout 标注 shared operand layout保持支持,并补充 allow_buffer_regionallow_reannotation 使用场景
PassConfig支持控制 SQMMA/WMMA lowering 和安全访问相关 pass补充 TL_DISABLE_INDEX_TYPE_PROMOTION 等 index 类型提升控制说明
accelerated_opsrelease 文档中未单独成文新增 MUSA accelerated ops 文档,包含 T.mul_half_float_to_bfloat16_x4

Copy 扩展

force_async_copy=True 用于显式请求将 global -> sharedT.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_policyouter_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 -> FMAA/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)
caution

关闭安全访问相关 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 必须是 float16x4rhs 必须是 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 参数。布尔变量通常接受 1trueyeson 作为启用值。

变量用途
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_SOURCEdump 生成的 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",
},
)

名词解释

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