芯東西(公眾號:aichip001)
作者 | ?ZeR0
編輯 | ?漠影
芯東西2月10日報道,今天,摩爾線程宣布近日開源TileLang-MUSA項目,實現(xiàn)對TileLang編程語言的完整支持,旨在充分釋放全功能GPU的性能潛力,大幅降低國產(chǎn)GPU開發(fā)門檻。
此前DeepSeek-V3大模型的研發(fā)已采用TileLang進(jìn)行算子快速原型設(shè)計與性能驗證,證明了其在大規(guī)模模型訓(xùn)練中的實戰(zhàn)價值。
TileLang_MUSA項目已在摩爾線程多代全功能GPU上完成功能驗證與特性開發(fā),驗證了“高層語言+國產(chǎn)GPU”技術(shù)路線的可行性。

開源地址:https://github.com/MooreThreads/tilelang_musa
一、重新定義GPU算子編程,“零門檻”實現(xiàn)算子遷移
TileLang是一款基于張量分塊(Tiling)抽象的高性能AI算子編程語言,屬于領(lǐng)域特定語言(DSL),采用聲明式語法與類Python前端,使開發(fā)者能以接近數(shù)學(xué)公式的形式描述計算意圖,并由編譯器自動完成循環(huán)優(yōu)化、內(nèi)存調(diào)度與代碼生成,在保持底層性能的同時,大幅降低GPU及異構(gòu)計算平臺的編程復(fù)雜度。
在實際應(yīng)用中,TileLang通過高級抽象降低開發(fā)門檻,具備跨平臺能力,實現(xiàn)“一次編寫、多架構(gòu)運(yùn)行”,編譯器自動執(zhí)行Layout推導(dǎo)、線程映射、Warp特化、流水線排布、內(nèi)存優(yōu)化等復(fù)雜優(yōu)化,在保障性能的同時提升GPU計算的開發(fā)效率,已廣泛應(yīng)用于AI與機(jī)器學(xué)習(xí)、科學(xué)計算等領(lǐng)域。
TileLang-MUSA項目提供了一種介于底層匯編與高層DSL之間的“中間層”抽象,
在保留硬件控制力的同時,顯著降低編程復(fù)雜度。
該項目硬件兼容性良好,已在摩爾線程多代全功能GPU上完成功能驗證與打通,包括訓(xùn)推一體全功能智算卡MTT S5000和MTT S4000。
團(tuán)隊還實現(xiàn)了TileLang高層語義到摩爾線程GPU底層MUSA架構(gòu)的精準(zhǔn)映射。
- Tensor Core加速:編譯器能夠自動調(diào)用MUSA的MMA(矩陣乘累加)指令,充分發(fā)揮硬件張量核心的峰值計算能力;
- Tile-Level Pipeline:自動處理從全局內(nèi)存(Global Memory)到共享內(nèi)存(Shared Memory)再到寄存器(Registers)的多級數(shù)據(jù)搬運(yùn),利用MUSA異步拷貝指令掩蓋訪存延遲;
- Warp級并行優(yōu)化:完整支持Warp Specialization特性。
目前,基于MUSA架構(gòu)的TileLang原生算子單元測試覆蓋率已超過80%,為大規(guī)模應(yīng)用提供了可靠保障。
在完成環(huán)境配置后,開發(fā)者可保留原有的import tilelang習(xí)慣,通過Cython編譯后端直接在MUSA環(huán)境中運(yùn)行TileLang代碼。
二、實測兼顧開發(fā)效率與運(yùn)行性能,代碼量減少約90%
在實際算子開發(fā)實踐中,TileLang-MUSA實現(xiàn)了讓開發(fā)者“寫得快”且“跑得快”。
以大語言模型中重要的FlashAttention-3和GEMM(通用矩陣乘)算子為例,在摩爾線程MTT S5000上的測試結(jié)果顯示:
(1)開發(fā)效率倍增:相較手寫MUSA C++代碼,使用TileLang-MUSA的代碼量減少了約90%,且代碼邏輯更加清晰,極大降低了開發(fā)與維護(hù)成本。
(2)性能媲美手寫:得益于編譯器優(yōu)化,生成的算子性能在典型配置下,Gemm最高可達(dá)手寫優(yōu)化版本的95%,F(xiàn)lashAttention-3可達(dá)手寫優(yōu)化版本的85%。
(3)自動化調(diào)優(yōu):借助TileLang-MUSA的Auto-tuning機(jī)制,開發(fā)者可在MUSA架構(gòu)的全功能GPU上快速搜索最優(yōu)的分塊策略(Tile Size)和流水線級數(shù),輕松超越未經(jīng)深度優(yōu)化的基準(zhǔn)實現(xiàn)。
TileLang-MUSA的推出,使TileLang用戶能近乎零成本地將算子邏輯遷移至摩爾線程GPU,還為不熟悉MUSA指令集的AI工程師提供了高層次的開發(fā)入口,并通過支持FlashAttention等關(guān)鍵算子的高效開發(fā),加速大語言模型等前沿AI應(yīng)用在國產(chǎn)算力平臺上的部署與落地。
結(jié)語:計劃打造基于MUSA的深度學(xué)習(xí)統(tǒng)一平臺
TileLang-MUSA的開源是摩爾線程構(gòu)建國產(chǎn)算力生態(tài)的關(guān)鍵一步,摩爾線程計劃持續(xù)推進(jìn)平臺與生態(tài)建設(shè),致力于打造一個覆蓋從單算子到完整大模型的國產(chǎn)算力統(tǒng)一加速平臺。
其計劃包括深度集成SGLang等主流AI框架,實現(xiàn)Transformer、MoE等復(fù)雜模型架構(gòu)的跨算子調(diào)度與全局優(yōu)化,完善調(diào)試和性能分析工具鏈,以及通過性能優(yōu)化,開發(fā)更多MUSA架構(gòu)定制擴(kuò)展,使生成代碼性能穩(wěn)定達(dá)到手寫優(yōu)化版本的90%以上。
這將為構(gòu)建一個開放、易用的國產(chǎn)算力開發(fā)生態(tài)提供工具支撐。