摩尔线程开源TileLang-MUSA:高效GPU算子开发 代码精简90%
创始人
2026-02-11 10:19:12
0

快科技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等复杂模型架构的跨算子调度与全局优化;

完善调试和性能分析工具链,为开发者提供全流程支持。

相关内容

热门资讯

神州答卷丨文旅新潮涌三晋   新华网太原2月11日电 题:文旅新潮涌三晋  山西,现有全国重点文物保护单位总数全国排名第一,拥...
北京市属公园新春攻略来了!百项... 2月11日,北京青年报记者从市公园管理中心获悉,北京市属公园及中国园林博物馆新春游园环境已全部布置完...
中国模型为何会在AI视频上领跑 直到这次字节的Seedance2.0出圈,很多人才第一次真正意识到,中国模型在 AI 视频这条赛道上...
AI看病遇上真人会“掉链子”? 当大语言模型在医师资格考试中都能取得高分时,人们曾期待它能成为贴身的“AI健康助手”。然而《自然·医...
年味香港,开心开运!解锁专属亲... 当新年的第一缕阳光洒向维多利亚港,一场充满欢笑与惊喜的亲子之旅,正等待着您的家庭启程。香港,这座中西...
2026马跃新春:全球达人游中... 2026马年新春将至,入境游市场迎来高质量发展新阶段,海外游客从观光打卡转向文化深度体验,“反向春运...
海岛游、山水游、城市观光游 春... 记者从交通运输部海事局了解到,今年春运,全国水路旅客发送量预计3300万人次。春运一周以来,水上出行...
请到广东过大年!新春登广州塔,... 广州塔旅游区以“新春登广州塔,好运一整年,马上有福!”为主题,精心策划了一系列涵盖传统民俗、潮流体验...
观山湖灯会庙市璀璨启幕 105... 本文转自:人民网-贵州频道观山湖灯会庙市璀璨启幕。华灯初上,年味渐浓。随着暮色浸染天际,2月10日晚...