摩爾線程開源TileLang-MUSA,以高效算子開發(fā)釋放全功能GPU計算潛力

來源: 摩爾線程官網(wǎng)官微

  近日,摩爾線程正式開源TileLang-MUSA項目,實現(xiàn)對TileLang編程語言的完整支持。該項目已成功在摩爾線程多代全功能GPU上完成功能驗證與特性開發(fā),旨在通過高層抽象與編譯器優(yōu)化,大幅降低開發(fā)門檻,為國產(chǎn)算力平臺提供更高效的AI與高性能計算開發(fā)體驗。

  TileLang:重新定義GPU算子編程的語言

  TileLang是一款基于張量分塊(Tiling)抽象的高性能AI算子編程語言,屬于領域特定語言(DSL)。它采用聲明式語法與類Python前端,使開發(fā)者能夠以接近數(shù)學公式的形式描述計算意圖,并由編譯器自動完成循環(huán)優(yōu)化、內(nèi)存調(diào)度與代碼生成,在保持底層性能的同時大幅降低GPU及異構(gòu)計算平臺的編程復雜度。

  在實際應用中,TileLang通過三大核心作用顯著提升GPU計算的開發(fā)效率:

  通過高級抽象降低開發(fā)門檻,開發(fā)者無需深入底層硬件知識即可生成高性能代碼;

  具備跨平臺能力,實現(xiàn)“一次編寫、多架構(gòu)運行”,有效解決多元算力生態(tài)的適配難題;

  編譯器自動執(zhí)行Layout推導、線程映射、Warp特化、流水線排布、內(nèi)存優(yōu)化等復雜優(yōu)化,在保障性能的同時提升開發(fā)效率。

  TileLang已廣泛應用于多個關鍵領域,例如:在AI與機器學習中,研究人員可用其快速定義新型算子(如注意力機制);在科學計算中,可便捷移植大型數(shù)值模擬程序至不同硬件平臺;對硬件廠商而言,TileLang可作為構(gòu)建芯片軟件生態(tài)的基礎工具鏈。

  在產(chǎn)業(yè)實踐中,DeepSeek-V3的研發(fā)已采用TileLang進行算子快速原型設計與性能驗證,證明了其在大規(guī)模模型訓練中的實戰(zhàn)價值。摩爾線程開源的TileLang-MUSA項目,正是這一技術理念的產(chǎn)業(yè)落地——通過提供高效開發(fā)工具鏈降低創(chuàng)新門檻,推動國產(chǎn)算力應用生態(tài)的繁榮發(fā)展。

  TileLang-MUSA:連接前沿語法與國產(chǎn)算力的橋梁

  摩爾線程此次開源的 TileLang-MUSA項目,旨在充分釋放全功能GPU的性能潛力。它提供了一種介于底層匯編與高層DSL之間的“中間層”抽象,在保留硬件控制力的同時,顯著降低了編程復雜度。具體特性如下:

  ▼ 廣泛的硬件架構(gòu)覆蓋:TileLang-MUSA已在摩爾線程多代全功能GPU上完成功能驗證與打通,包括訓推一體全功能智算卡 MTT S5000和MTT S4000,展現(xiàn)了良好的硬件兼容性。

  ▼ 核心計算特性的深度映射:項目團隊實現(xiàn)了TileLang高層語義到摩爾線程GPU底層MUSA架構(gòu)的精準映射。

  Tensor Core 加速:編譯器能夠自動調(diào)用MUSA的MMA(矩陣乘累加)指令,充分發(fā)揮硬件張量核心的峰值計算能力;

  Tile-Level Pipeline:自動處理從全局內(nèi)存(Global Memory)到共享內(nèi)存(Shared Memory)再到寄存器(Registers)的多級數(shù)據(jù)搬運,利用MUSA異步拷貝指令掩蓋訪存延遲;

  Warp級并行優(yōu)化:完整支持Warp Specialization特性。

  目前,基于MUSA架構(gòu)的TileLang原生算子單元測試覆蓋率已超過80%,為大規(guī)模應用提供了可靠保障。

  代碼示例:體驗“零門檻”算子遷移

  開發(fā)者在完成環(huán)境配置后,可保留原有的import tilelang習慣,通過Cython編譯后端直接在MUSA環(huán)境中運行TileLang代碼。

  以下是一個基于TileLang-MUSA的高性能矩陣乘法實現(xiàn)示例:通過簡潔明了的語法描述計算邏輯,編譯器能夠自動生成高度優(yōu)化的MUSA內(nèi)核代碼。依托TileLang豐富且經(jīng)過深度調(diào)優(yōu)的原語算子庫(primitives),開發(fā)者借助TileLang-MUSA不僅能夠顯著提升編碼效率,更能直接調(diào)用底層硬件的計算潛力,實現(xiàn)媲美手工優(yōu)化性能的矩陣運算。

  import tilelangimport tilelang.language as Timport torchfrom tilelang.primitives.gemm.base import GemmWarpPolicy def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"): @T.prim_funcdef 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()

  *代碼示例

  實測表現(xiàn):兼顧開發(fā)效率與運行性能

  TileLang-MUSA的設計理念是讓開發(fā)者“寫得快”且“跑得快”。在實際算子開發(fā)實踐中,這一目標已得到驗證。

  以大語言模型中至關重要的FlashAttention-3和GEMM(通用矩陣乘)算子為例,在摩爾線程MTT S5000上的測試結(jié)果顯示:

  ▼ 開發(fā)效率倍增:相較手寫MUSA C++代碼,使用TileLang-MUSA的代碼量減少了約90%,且代碼邏輯更加清晰,極大降低了開發(fā)與維護成本。

  ▼ 性能媲美手寫:得益于TileLang-MUSA優(yōu)秀的編譯器優(yōu)化,生成的算子性能在典型配置下,Gemm最高可達到手寫優(yōu)化版本的95%, FlashAttention-3可以達到手寫優(yōu)化版本的85%。

  ▼ 自動化調(diào)優(yōu):借助TileLang-MUSA的Auto-tuning機制,開發(fā)者可以在MUSA架構(gòu)的全功能GPU上快速搜索最優(yōu)的分塊策略(Tile Size)和流水線級數(shù),輕松超越未經(jīng)深度優(yōu)化的基準實現(xiàn)。

  TileLang-MUSA的推出,為開發(fā)者帶來了三大核心價值:

  ▼ 無縫遷移:TileLang用戶可以近乎零成本地將算子邏輯遷移至摩爾線程GPU;

  ▼ 降低門檻:TileLang-MUSA為不熟悉MUSA指令集的AI工程師提供了高層次的開發(fā)入口,大幅降低了國產(chǎn)GPU的應用門檻;

  ▼ 賦能大模型:TileLang-MUSA通過支持FlashAttention等關鍵算子的高效開發(fā),將有力加速大語言模型等前沿AI應用在國產(chǎn)算力平臺上的部署與落地。

  TileLang-MUSA不僅驗證了“高層語言+國產(chǎn)GPU”這一技術路線的可行性,更為構(gòu)建開放、易用的國產(chǎn)算力開發(fā)生態(tài)提供了關鍵工具支撐。

  未來展望:打造基于MUSA的深度學習統(tǒng)一平臺

  TileLang-MUSA的開源是摩爾線程構(gòu)建國產(chǎn)算力生態(tài)的關鍵一步。以此為起點,摩爾線程將持續(xù)推進平臺與生態(tài)建設,致力于打造一個覆蓋從單算子到完整大模型的國產(chǎn)算力統(tǒng)一加速平臺:持續(xù)進行性能優(yōu)化,開發(fā)更多MUSA架構(gòu)定制擴展,使生成代碼性能穩(wěn)定達到手寫優(yōu)化版本的90%以上;深度集成SGLang等主流AI框架,實現(xiàn)訓練與推理場景的端到端無縫加速;從單算子優(yōu)化延伸至Transformer、MoE等復雜模型架構(gòu)的跨算子調(diào)度與全局優(yōu)化;同時完善調(diào)試和性能分析工具鏈,為開發(fā)者提供全流程支持。

  這一切努力,最終都將匯聚于一個核心目標:構(gòu)建一個開放、高效、充滿生命力的國產(chǎn)算力生態(tài),讓創(chuàng)新的想法在堅實的基石上自由生長。

關注同花順財經(jīng)(ths518),獲取更多機會

0

+1
  • 北信源
  • 兆易創(chuàng)新
  • 科森科技
  • 卓翼科技
  • 天融信
  • 吉視傳媒
  • 御銀股份
  • 中油資本
  • 代碼|股票名稱 最新 漲跌幅