驱动之家 02-11
摩尔线程开源TileLang-MUSA:高效GPU算子开发 代码精简90%
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_keji1.html

 

快科技 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, A_shared )

                T.copy ( B, B_shared )

                T.gemm ( A_shared, B_shared, C_local, policy=T.GemmWarpPolicy.Square )

            T.copy ( C_local, C )

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

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

宙世代

宙世代

ZAKER旗下Web3.0元宇宙平台

一起剪

一起剪

ZAKER旗下免费视频剪辑工具

相关标签

摩尔 ai gpu 开源 编程
相关文章
评论
没有更多评论了
取消

登录后才可以发布评论哦

打开小程序可以发布评论哦

12 我来说两句…
打开 ZAKER 参与讨论