📢 转载信息
原文链接：https://news.mydrivers.com/1/1103/1103648.htm
原文作者：上方文Q
摩尔线程TileLang-MUSA开源：高效GPU算子开发 代码精简90%
快科技2月11日消息，近日，摩尔线程正式开源TileLang-MUSA项目，实现对TileLang编程语言的完整支持。
该项目已成功在摩尔线程多代全功能GPU上完成功能验证与特性开发，旨在通过高层抽象与编译器优化，大幅降低开发门槛，为国产算力平台提供更高效的AI与高性能计算开发体验。
▼ TileLang-MUSA开源地址：
https://github.com/MooreThreads/tilelang_musa
TileLang是一款基于张量分块（Tiling）抽象的高性能AI算子编程语言，属于领域特定语言（DSL）。
它采用声明式语法与类Python前端，使开发者能够以接近数学公式的形式描述计算意图，并由编译器自动完成循环优化、内存调度与代码生成，在保持底层性能的同时大幅降低GPU及异构计算平台的编程复杂度。
在实际应用中，TileLang通过三大核心作用显著提升GPU计算的开发效率：
- 通过高级抽象降低开发门槛，开发者无需深入底层硬件知识即可生成高性能代码；
- 具备跨平台能力，实现“一次编写、多架构运行”，有效解决多元算力生态的适配难题；
- 编译器自动执行Layout推导、线程映射、Warp特化、流水线排布、内存优化等复杂优化，在保障性能的同时提升开发效率。
TileLang已广泛应用于多个关键领域，例如：
- 在AI与机器学习中，研究人员可用其快速定义新型算子（如注意力机制）；
- 在科学计算中，可便捷移植大型数值模拟程序至不同硬件平台；
- 对硬件厂商而言，TileLang可作为构建芯片软件生态的基础工具链。
在产业实践中，DeepSeek-V3的研发已采用TileLang进行算子快速原型设计与性能验证，证明了其在大规模模型训练中的实战价值。
摩尔线程此次开源的 TileLang-MUSA项目，旨在充分释放摩尔线程全功能GPU的性能潜力。它提供了一种介于底层汇编与高层DSL之间的“中间层”抽象，在保留硬件控制力的同时，显著降低了编程复杂度。
具体特性如下：
广泛的硬件架构覆盖：
TileLang-MUSA已在摩尔线程多代全功能GPU上完成功能验证与打通，包括训推一体全功能智算卡 MTT S5000和MTT S4000，展现了良好的硬件兼容性。
核心计算特性的深度映射：
项目团队实现了TileLang高层语义到摩尔线程GPU底层MUSA架构的精准映射。
- Tensor Core 加速：编译器能够自动调用MUSA的MMA（矩阵乘累加）指令，充分发挥硬件张量核心的峰值计算能力；
- Tile-Level Pipeline：自动处理从全局内存（Global Memory）到共享内存（Shared Memory）再到寄存器（Registers）的多级数据搬运，利用MUSA异步拷贝指令掩盖访存延迟；
- Warp级并行优化：完整支持Warp Specialization特性。
目前，基于MUSA架构的TileLang原生算子单元测试覆盖率已超过80%，为大规模应用提供了可靠保障。
开发者在完成环境配置后，可保留原有的import tilelang习惯，通过Cython编译后端直接在MUSA环境中运行TileLang代码。
以下是一个基于TileLang-MUSA的高性能矩阵乘法实现示例：通过简洁明了的语法描述计算逻辑，编译器能够自动生成高度优化的MUSA内核代码。依托TileLang丰富且经过深度调优的原语算子库（primitives），开发者借助TileLang-MUSA不仅能够显著提升编码效率，更能直接调用底层硬件的计算潜力，实现媲美手工优化性能的矩阵运算。
import tilelang
import tilelang.language as T
import torch
from tilelang.primitives.gemm.base import GemmWarpPolicy
def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
@T.prim_func
def matmul_kernel(
A: T.Tensor((M, K), dtype),
B: T.Tensor((K, N), dtype),
C: T.Tensor((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 matmul_kernel
def main():
M, N, K, BLOCK_M, BLOCK_N, BLOCK_K = 8192, 7168, 16384, 256, 256, 64
device = "musa"
A = torch.randn((M, K), dtype=torch.float16, device=device)
B = torch.randn((K, N), dtype=torch.float16, device=device)
program = matmul(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, dtype="float16", accum_dtype="float32")
kernel = tilelang.compile(
program,
out_idx=-1,
target="musa",
execution_backend="cython",
verbose=True,
)
ref_out = torch.mm(A, B)
C = kernel(A, B)
torch.testing.assert_close(ref_out.to(torch.float16), C.to(torch.float16), rtol=1.25e-1, atol=1.25e-1)
if __name__ == "__main__":
main()
TileLang-MUSA的设计理念是让开发者“写得快”且“跑得快”。在实际算子开发实践中，这一目标已得到验证。
以大语言模型中至关重要的FlashAttention-3和GEMM（通用矩阵乘）算子为例，在摩尔线程MTT S5000上的测试结果显示：
开发效率倍增：
相较手写MUSA C++代码，使用TileLang-MUSA的代码量减少了约90%，且代码逻辑更加清晰，极大降低了开发与维护成本。
性能媲美手写：
得益于TileLang-MUSA优秀的编译器优化，生成的算子性能在典型配置下，Gemm最高可达到手写优化版本的95%, FlashAttention-3可以达到手写优化版本的85%。
自动化调优：
借助TileLang-MUSA的Auto-tuning机制，开发者可以在MUSA架构的全功能GPU上快速搜索最优的分块策略（Tile Size）和流水线级数，轻松超越未经深度优化的基准实现。
TileLang-MUSA的推出，为开发者带来了三大核心价值：
无缝迁移：
TileLang用户可以近乎零成本地将算子逻辑迁移至摩尔线程GPU；
降低门槛：
TileLang-MUSA为不熟悉MUSA指令集的AI工程师提供了高层次的开发入口，大幅降低了国产GPU的应用门槛；
赋能大模型：
TileLang-MUSA通过支持FlashAttention等关键算子的高效开发，将有力加速大语言模型等前沿AI应用在国产算力平台上的部署与落地。
TileLang-MUSA的开源是摩尔线程构建国产算力生态的关键一步。
以此为起点，摩尔线程将持续推进平台与生态建设，致力于打造一个覆盖从单算子到完整大模型的国产算力统一加速平台：
- 持续进行性能优化，开发更多MUSA架构定制扩展，使生成代码性能稳定达到手写优化版本的90%以上；
- 深度集成SGLang等主流AI框架，实现训练与推理场景的端到端无缝加速；
- 从单算子优化延伸至Transformer、MoE等复杂模型架构的跨算子调度与全局优化；
- 完善调试和性能分析工具链，为开发者提供全流程支持。
🚀 想要体验更好更全面的AI调用？
欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。
评论区