TileLang-MUSA算子开发 | 摩尔线程 GPU 学习中心

一、TileLang-MUSA：高性能算子的 DSL 之路

当 PyTorch 标准算子和 torch.compile 仍无法压榨出 GPU 的全部性能时，就要手写 kernel。但手写 MUSA C++ kernel 门槛高、调优复杂、迁移困难。TileLang-MUSA 给出了第三条路：用 Python DSL 写算子，编译器自动展开为高性能 MUSA 代码。

TileLang 本身基于 TVM，是一种面向 Tile 编程抽象的领域特定语言。摩尔线程维护的 tilelang_musa 仓库把它对接到了 MUSA 后端，面向 MTT S5000 / S4000 / M1000。公开的基准数据：

GEMM：可达手写 MUSA kernel 的 95% 性能。
FlashAttention-3：可达手写实现的 85% 性能。
Auto-tuning：几分钟内自动在数十种 tile 配置中挑选最优。

💡

什么叫 Tile 编程？

现代 GPU kernel 的核心抽象是 tile：把大矩阵切成若干小块（如 128×128），每个 block 处理一个 tile，tile 内部再切成 warp/thread tile。TileLang 把这些层级用 Python 显式表达，让开发者只关心"形状 + 数据流"，其余交给编译器。

二、环境安装

bash

# 前置：MUSA SDK >= 4.3.2, torch_musa >= 2.9.0
pip install torch_musa torchada

# 安装 tilelang_musa（推荐源码）
git clone https://github.com/MooreThreads/tilelang_musa.git
cd tilelang_musa
pip install -e .

# 验证
python -c "import tilelang; print(tilelang.__version__)"

三、第一个 TileLang Kernel：GEMM

下面这段代码实现 C = A @ B，A: (M, K)，B: (K, N)，C: (M, N)，block size 128×128，K 方向每次加载 32。

python

import tilelang
import tilelang.language as T
import torch, torchada

@tilelang.jit(target="musa")
def gemm(M, N, K, block_M=128, block_N=128, block_K=32, dtype="float16"):
    @T.prim_func
    def main(
        A: T.Buffer((M, K), dtype),
        B: T.Buffer((K, N), dtype),
        C: T.Buffer((M, N), dtype),
    ):
        # 划分 grid：每个 block 处理 (block_M x block_N) 的 C 子块
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            # 共享内存 tile
            As = T.alloc_shared((block_M, block_K), dtype)
            Bs = T.alloc_shared((block_K, block_N), dtype)
            # 寄存器累加器（float32 精度）
            Cr = T.alloc_fragment((block_M, block_N), "float32")
            T.clear(Cr)

            # K 维度流水线
            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                T.copy(A[by * block_M, ko * block_K], As)
                T.copy(B[ko * block_K, bx * block_N], Bs)
                T.gemm(As, Bs, Cr)

            T.copy(Cr, C[by * block_M, bx * block_N])
    return main

# 生成 kernel 并绑定 PyTorch
M, N, K = 1024, 1024, 1024
kernel = gemm(M, N, K)

A = torch.randn(M, K, device="cuda", dtype=torch.float16)
B = torch.randn(K, N, device="cuda", dtype=torch.float16)
C = torch.empty(M, N, device="cuda", dtype=torch.float16)
kernel(A, B, C)

ref = A.float() @ B.float()
print("最大误差:", (C.float() - ref).abs().max().item())

这段代码没有任何 MUSA 汇编，也没有 shared memory 手工 swizzle —— 编译器根据你的 tile size 自动生成最优布局、双缓冲、tensor core 指令。

四、Auto-tuning：自动选最优 tile

同一个算法，不同 shape 下最优 tile 大小不同。TileLang 提供 autotune 装饰器，在指定配置空间内自动枚举并选最快：

python

import itertools
from tilelang import autotune

configs = [
    {"block_M": m, "block_N": n, "block_K": k, "num_stages": s}
    for m, n, k, s in itertools.product([64, 128], [64, 128], [32, 64], [2, 3, 4])
]

@autotune(configs=configs, keys=["M", "N", "K"])
@tilelang.jit(target="musa")
def tuned_gemm(M, N, K, block_M, block_N, block_K, num_stages, dtype="float16"):
    @T.prim_func
    def main(A: T.Buffer((M,K), dtype), B: T.Buffer((K,N), dtype), C: T.Buffer((M,N), dtype)):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            As = T.alloc_shared((block_M, block_K), dtype)
            Bs = T.alloc_shared((block_K, block_N), dtype)
            Cr = T.alloc_fragment((block_M, block_N), "float32"); T.clear(Cr)
            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
                T.copy(A[by*block_M, ko*block_K], As)
                T.copy(B[ko*block_K, bx*block_N], Bs)
                T.gemm(As, Bs, Cr)
            T.copy(Cr, C[by*block_M, bx*block_N])
    return main

# 首次调用会自动跑完整 config 空间，挑选最优
tuned_gemm(1024, 1024, 1024)(A, B, C)

✅

Auto-tuning 缓存

TileLang 会把每组 (M,N,K) 对应的最优配置保存到本地缓存（默认 ~/.cache/tilelang）。生产环境建议把缓存提前生成好、打进镜像，避免第一次请求拉胯。

五、FlashAttention 实战骨架

FlashAttention 的核心是在线 softmax + 分块重算。TileLang 下的结构大致如下：

python

@tilelang.jit(target="musa")
def flash_attention(batch, heads, seq, dim, block_M=128, block_N=64, dtype="float16"):
    @T.prim_func
    def main(
        Q: T.Buffer((batch, heads, seq, dim), dtype),
        K: T.Buffer((batch, heads, seq, dim), dtype),
        V: T.Buffer((batch, heads, seq, dim), dtype),
        O: T.Buffer((batch, heads, seq, dim), dtype),
    ):
        with T.Kernel(T.ceildiv(seq, block_M), heads, batch, threads=128) as (bx, by, bz):
            Qs = T.alloc_shared((block_M, dim), dtype)
            Ks = T.alloc_shared((block_N, dim), dtype)
            Vs = T.alloc_shared((block_N, dim), dtype)
            Or = T.alloc_fragment((block_M, dim), "float32")
            m_i = T.alloc_fragment((block_M,), "float32")   # row max
            l_i = T.alloc_fragment((block_M,), "float32")   # row sum
            T.clear(Or); T.fill(m_i, -1e30); T.clear(l_i)

            T.copy(Q[bz, by, bx*block_M], Qs)

            for kb in T.Pipelined(T.ceildiv(seq, block_N), num_stages=2):
                T.copy(K[bz, by, kb*block_N], Ks)
                T.copy(V[bz, by, kb*block_N], Vs)
                S = T.alloc_fragment((block_M, block_N), "float32")
                T.gemm(Qs, Ks, S, transpose_B=True)
                # 在线 softmax: 更新 m_i / l_i / Or
                T.online_softmax_update(S, m_i, l_i, Or, Vs)

            T.copy(Or / l_i, O[bz, by, bx*block_M])
    return main

真正的细节远比这段骨架复杂：causal mask、row-broadcast 指令、accumulator dtype 等都需要额外处理。推荐直接阅读 tilelang_musa/examples/flash_attention/ 的完整实现。

六、把 TileLang kernel 接入 PyTorch

编译出来的 kernel 就是一个普通 Python callable，可以封装成 torch.autograd.Function 支持训练：

python

class TLGemm(torch.autograd.Function):
    kernel = gemm(1024, 1024, 1024)

    @staticmethod
    def forward(ctx, A, B):
        C = torch.empty(A.size(0), B.size(1), device=A.device, dtype=A.dtype)
        TLGemm.kernel(A, B, C)
        ctx.save_for_backward(A, B)
        return C

    @staticmethod
    def backward(ctx, dC):
        A, B = ctx.saved_tensors
        dA = dC @ B.T
        dB = A.T @ dC
        return dA, dB

# 使用
y = TLGemm.apply(x1, x2)

七、TileLang-MUSA 的未来

根据 3/23 周报与 tilelang_musa 仓库的 Issue 规划：

SGLang 集成：将 TileLang 版 FlashAttention 作为 SGLang-MUSA 的默认 attention 后端。
MoE GroupGEMM：针对 DeepSeek / MiniMax 的 MoE expert 加速。
训练算子：覆盖 LayerNorm backward、RMSNorm、Rotary embedding 等必备训练算子。

八、DeepSeek-V4 TileKernels Day-0 复用（2026-05-11 新增）

2026 年 4 月 25 日，TileLang-MUSA 正式进入 TileLang 官方主线，同时实现了对 DeepSeek 在 V4 发布时新放出的 TileKernels 算子库的 Day-0 无缝支持。这是国内 GPU 生态里少见的"上游主线 + 前沿模型同步"的双重落地。

8.1 三层技术路径

摩尔线程在官方报道（2026-05-01 新闻）中公开了 DeepSeek-V4 推理适配的系统化工程路径，可以总结成下面三层：

第一层：MUSA ↔ CUDA 全栈兼容
vLLM / Megatron-LM / Transformers 等主流框架直接跑

第二层：复用 TileKernels 开源算子库
替代 PyTorch 标准融合算子，吃掉 DeepSeek-V4 的 FP8 / Sparse Attention 设计

第三层：AI Agent + TileLang 快速补齐模型特异算子
用 AI 辅助写 TileLang DSL，分钟级补 missing kernel

图：DeepSeek-V4 在 MTT S5000 上 Day-0 跑通的三层技术路径

8.2 把 TileKernels 直接拿来用

DeepSeek 在 V4 发布时把 TileKernels 同步开源，这是一组用 TileLang 写的、面向 V4 自身架构（MoE、Sparse Attention、FP8 Linear）调优过的算子。Day-0 适配意味着：当 V4 模型权重发布时，TileLang-MUSA 已经把全部 kernel 原样编译跑通，没有等待期。

python

# 复用 DeepSeek-V4 官方 TileKernels（在 tilelang_musa 已合入主线之后）
import tilelang
import tilelang.musa  # 注册 MUSA 后端
from tilekernels.attention import flash_sparse_attention_v4
from tilekernels.linear  import fp8_block_scaled_linear

# 直接当作 PyTorch 算子使用，调度到 MTT S5000 的 TME/TCE 单元
out = flash_sparse_attention_v4(
    q, k, v,
    sparse_pattern="deepseek_v4",
    fp8_quant=True,
)
proj = fp8_block_scaled_linear(out, w_fp8, scale_w, block=128)

✅

编译优化关键点

FlagOS 团队在 DeepSeek-V4 适配过程里做了两件事：一是通过精细化 shape 对齐让 FP8 / Sparse Attention 的计算 shape 满足 MTT S5000 的 TME / TCE（张量访存与计算引擎）要求；二是结合 MUSA_ENABLE_SQMMA 加速 tl.dot 矩阵乘。这两条经验都已经回流到 TileLang-MUSA 主线，业务方只需要把 TileKernels import 进来就能吃到。

8.3 MUSA 与 CUDA 的全栈兼容在这次起了什么作用

DeepSeek-V4 带来的是架构层面的升级：它第一次采用了 FP4+FP8 的混合精度设计，同时在总参量 284B / 激活参量 13B 的 MoE 上叠加了 Sparse Attention，最大上下文长度被拉到百万级。这意味着传统的 BF16 计算路径、标准融合 attention、以及不考虑稀疏性的 GEMM kernel 都不再适用。

摩尔线程这次能跳一上来就 Day-0 跑通，依靠的正是三层路径中间那一层的复用能力：MUSA 生态对 CUDA 生态的全栈兼容 + TileLang DSL 的跨后端特性。具体表现为：TileKernels 上游发布的是面向 CUDA 后端的 TileLang 源码，TileLang-MUSA 主线部分提供了等价的 MUSA 后端，同一份 DSL 源码不需要改动即可同时生成 CUDA 与 MUSA 可执行代码。这也是由它背后的 MUSA TME（张量访存引擎）与 TCE（张量计算引擎）、以及 MUSA对 CUDA 运行时 API 的 95%+ 改造率决定的。

从工程视角看，这意味着未来再出现类似 DeepSeek-V4 级别的架构升级时，MUSA 上的适配周期不再需要从零重写 kernel，而是可以复用上游适配成果。这对于面向生产环境的团队、复现论文结果、跟踪最新模型能力的团队是重要的可预测性保障。

8.4 AI Agent + TileLang 补 missing kernel

新模型出来时永远会有少数算子未覆盖。摩尔线程在 V4 适配里走的是"AI Agent + TileLang"的路径：

把待实现算子的数学定义（来自模型论文 / 参考实现）丢给 LLM。

让模型按 TileLang DSL 风格生成初始版本（block / tile / pipeline 等都暴露给 LLM 作为模板）。

用 TileLang 内置的 numerical check + Auto-tuning 自动校验精度并搜索最优配置。

通过后注册到 PyTorch dispatcher，立即可被模型代码调用。

💡

为什么必须用 DSL

用 AI Agent 直接产生纯 MUSA C++ kernel 的成功率很低（缺乏目标硬件资料），而 TileLang 是一种对硬件细节做了抽象的 DSL：Agent 只需要在 tile + pipeline 层面思考，编译器把后端工作交给 TileLang-MUSA 完成。这是 "DSL + AI" 跑得通而 "C++ + AI" 跑不通的关键。

⚠️

学习曲线提示

TileLang 的 Python DSL 看着像普通代码，但语义是编译期的 —— 你写的循环、分配都是在生成代码，运行时只会跑一次编译好的 kernel。对 tile、shared memory、pipeline 等概念不熟悉的读者，建议先完成模块三（MUSA 编程入门）的学习。

✅

本课小结

TileLang-MUSA 为摩尔线程 GPU 带来了高生产力 + 高性能的算子开发方式。GEMM 可达手写 95%，FlashAttention-3 可达 85%，配合 Auto-tuning 能在几分钟内跑完 shape 空间搜索。本周它进入 TileLang 官方主线，同时实现了对 DeepSeek-V4 TileKernels 算子库的 Day-0 无缝支持，意味着业务方可以复用官方调优过的 FP8 / Sparse Attention kernel。仓库：MooreThreads/tilelang_musa。