黃千慈
強(qiáng)大的芯片有效算力是作為“AI工廠”的核心競爭力。而作為GPU的“靈魂”,其軟件能力深刻決定了硬件潛力的釋放程度,特別是驅(qū)動與全棧系統(tǒng)軟件,是解鎖大規(guī)模集群極致性能效率的關(guān)鍵。
7月25日,在WAIC2025期間舉辦的摩爾線程技術(shù)分享日上,摩爾線程GPU計算軟件開發(fā)總監(jiān)吳慶以《MUSA軟件棧助力KUAE集群釋放無限潛能》為主題,分享了摩爾線程MUSA軟件棧方面的創(chuàng)新實踐。
以下為吳慶的演講實錄精華:
MUSA驅(qū)動和運行時庫,KUAE集群堅實底座
推理場景對Kernel延時非常敏感。極致高效的MUSA驅(qū)動和運行時庫,具有高效率的任務(wù)調(diào)度能力,使用戶的Kernellaunch開銷極致縮減。
(1)即時任務(wù)下發(fā):通過軟硬協(xié)同,將核函數(shù)啟動延遲降低到業(yè)界平均水平的1/2。舉個例子方便大家理解,比如現(xiàn)在剛好是仲夏,外面氣溫非常高,我們的車子停在外面暴曬,那如何讓大家一上車即時就享受到舒適的溫度?現(xiàn)在普遍的解決方案是軟硬協(xié)同,車要支持遠(yuǎn)程控制,還沒有到車上的時候就通過手機(jī)提前把車上的空調(diào)打開,這樣一上車以后馬上就有清涼舒適的溫度。這個原理類似,我們把很多調(diào)度開銷前置,上一個Kernel執(zhí)行結(jié)束之前,就可以把下一個Kernel要做的配置提前做好,這樣就最大限度地降低Kernel跟Kernel之間的調(diào)度延時開銷。
(2)批量任務(wù)下發(fā):批量下發(fā)計算和通信任務(wù),將近千次下發(fā)開銷優(yōu)化為單次,減少GPU等待時間。為方便大家理解,我再類比一下,車上座椅調(diào)節(jié),調(diào)舒適位子有上下前后,包括靠背俯仰角度,有很多調(diào)節(jié)步驟對應(yīng)很多Kernel,但我們可以調(diào)節(jié)一次、打包,變成一鍵記憶功能,下次再上車,不需要再重復(fù)那個步驟了。這就類似于批量任務(wù)下發(fā),一系列操作打包、我們一次kick下去。Graphcapture階段對應(yīng)第一次調(diào)節(jié),后面再launch一個Graph的時候就類似于我們喚醒座椅某個記憶配置。Graph這種批量任務(wù)下發(fā)的收益取決于Graph里面Kernel本身執(zhí)行時間、調(diào)度空泡的占比,Kernel執(zhí)行得越快,調(diào)度空泡在端到端總開銷里占比越高,批量任務(wù)下發(fā)的端到端收益就越高,根據(jù)不同的應(yīng)用場景可達(dá)到百分之幾十甚至好幾倍的性能收益。
(3)引擎間依賴解析:GPU是多引擎可配置的,不同引擎之間要做依賴解析、交互同步,平湖支持硬件解析引擎間的依賴做同步,不再回到host,任務(wù)流之間的依賴解析延時可以大幅降低至1.5μs,優(yōu)于業(yè)界頭部算力卡。
(4)調(diào)優(yōu)工具接口MUPTI:專門用于性能分析和事件追蹤,幫助開發(fā)者優(yōu)化MUSA應(yīng)用程序,可基于MUPTI開發(fā)第三方的各種豐富工具。
(5)GPU錯誤轉(zhuǎn)存功能(GPUCoreDump,GCD):高效定位疑難Bug,不需要去調(diào)度大量的集群和人力反復(fù)去復(fù)現(xiàn)“曇花一現(xiàn)”的錯誤,用于在GPU程序崩潰或發(fā)生錯誤時生成詳細(xì)的信息(類似CPU程序CoreDump),精確保存錯誤,提供一個log,極大提升MUSA核函數(shù)崩潰、非法顯存地址訪問、硬件錯誤等問題定位的效率。
MUSA算子庫生態(tài)日趨完善,性能與廣度兼?zhèn)?/p>
算子庫的效率直接關(guān)乎分布式集群的訓(xùn)練效率。MUSA算子庫不僅追求極致性能,還兼容覆蓋廣度和兼容性。
MUSA算子庫生態(tài)會提供三大算子庫:極致性能muDNN、易用MUTLASS、MUSAAITensorEngine開源推理算子庫。
muDNN,是一款極致性能的開箱即用標(biāo)準(zhǔn)算子庫,完整覆蓋常見的前向和反向算子。其中最主要的特性是支持完整的XMMA,支持TensorCore全精度及所有量化模式,以及常用的神經(jīng)網(wǎng)絡(luò)算子操作。
一般來說,矩陣乘的效率大部分可以做到90%以上,摩爾線程muDNN矩陣乘算子效率可以做到98%。FlashAttention算子因為多了softmax等操作,基于國際一流水平廠商的FlashAttention3論文里提到的效率大約是75%,muDNN的FlashAttention可以做到95%。
影響FP8GEMM算子效率的重要因素是scale模式,scale是為了避免精度損失做的一個精細(xì)化的縮放因子,常見的有Per-Tensor、Per-Block,DeepSeekV3用的是Per-Block。
吳慶打了個比方,Per-Tensorscale是整個張量共用一個縮放因子,相當(dāng)于一面墻刷一個顏色,很好刷;Per-Blockscale則是每個小的矩陣塊共用一個縮放因子,相當(dāng)于一面墻有很多小方塊,每個方塊單獨一個顏色,刷墻的復(fù)雜度要高很多,會比一個顏色刷一面墻效率低10%-20%。
摩爾線程通過軟硬協(xié)同的深度算法優(yōu)化、指令編排,做到了Per-Block跟Per-TensorFP8GEMM計算效率幾乎相當(dāng),差距不到2%。
MUTLASS,是一個高性能的線性代數(shù)模板庫,可極大降低在MUSA環(huán)境中自定義算子的開發(fā)工作量,相當(dāng)于提供了一個可供二次開發(fā)的模板。
MUTLASS已在Github上開源,支持平湖架構(gòu)所有特性,包括Warp-levelMMA、WarpSquad-levelMMA、TensorMemoryEngine、AsyncBarrier等原語全面支持,同時提供高性能矩陣乘、卷積算子實現(xiàn),可在Kernel里調(diào)用模板庫就可以進(jìn)行二次開發(fā),摩爾線程也提供Attention最優(yōu)化的示例,給用戶打樣,可以參考提供的最佳實踐,去開發(fā)自定義的各種Attention的變種。
MUSAAITensorEngine,是摩爾線程即將發(fā)布的面向大語言模型的開源推理算子庫,可幫助開發(fā)者快速搭建自定義推理引擎。該算子庫提供用戶友好的PythonAPI,進(jìn)一步降低開發(fā)者使用門檻,無需花太多時間學(xué)習(xí)MUSAC,只要會Python就可以把MTX這個庫用起來;未來會支持豐富的后端,包括模板庫,對標(biāo)PTX虛擬指令集去開發(fā)的極致算子庫,也支持類Triton的DSL。
高性能集合通信庫及通算并行,避免通信搶占計算核心資源
大模型分布式訓(xùn)練的主要瓶頸是通信。摩爾線程在MTTS5000上支持MTLINK2.0,配備GPU異步通信引擎,實現(xiàn)了高性能通信技術(shù)與計算高度并行。
最近十幾年,隨著模型規(guī)模增加,分布式集群規(guī)模也在增加,通信瓶頸會越來越明顯。DeepSeekV3報告中呼吁硬件廠商提供的芯片能不能做一個單獨的通信硬件來卸載通信,避免通信跟計算去搶SM計算單元。
平湖設(shè)計階段遠(yuǎn)早于DeepSeekV3發(fā)布的時間,彼時摩爾線程已洞察到行業(yè)痛點,并提供了解決方案——在GPU上增加了一個異步通信引擎。
MTTS5000服務(wù)器拓?fù)涿抗?jié)點有8張GPU,通過MTLINK2.0全互連,每張GPU與其他7張GPU都有直連總線。
每個GPU上配置的異步通信引擎是原生支持豐富的reduce操作,支持常用的reduce操作有ADD、MIN、MAX(累加、最小值、最大值)等操作,支持Float、FP16、BF16等數(shù)據(jù)類型,可直接通過MTLink進(jìn)行C2C跨系統(tǒng)的數(shù)據(jù)通信。
其ACE通過ZeroCopy技術(shù)進(jìn)一步提升性能。一般的集合通信會把數(shù)據(jù)buffer拷貝到通信buffer里,再去進(jìn)行C2C通信。ZeroCopy則避免了在本地上的D2D的拷貝,直接把數(shù)據(jù)buffer做跨卡通信。
摩爾線程在衡量集合通信庫性能的時候主要是兩大核心指標(biāo):通信延遲、通信帶寬。
基于全互聯(lián)(FC8)拓?fù)涞母咝ㄐ潘惴?,可大幅降低通信延遲。全互聯(lián)拓?fù)渲蠫PU兩兩互連,在AllReduce場景中可以把其他GPU的數(shù)據(jù)一次性拿過來,理論上FC8算法相比Ring算法,7步通信可變成1步完成。在單機(jī)8卡AllReduce延遲場景的實測也差不多,Ring算法大概是53us,用FC8算法可降低至7.8us,將近縮減至1/7,顯著提升了模型推理性能。
帶寬方面,Ring算法與FC8算法效率相當(dāng),以Ring算法為例,摩爾線程把FC8拆成了7個通信環(huán),每個通信環(huán)結(jié)合通信Kernel的優(yōu)化,基于MTLink2.0+FC8拓?fù)?,?jié)點內(nèi)充分利用全互聯(lián)的拓?fù)鋬?yōu)勢,可做到單機(jī)8卡AllReduce帶寬利用率接近85%,與國際第一的廠商在軟件效率上對齊,實現(xiàn)高效的Scale-up通信。
高效Scale-out跨節(jié)點通信方面,摩爾線程結(jié)合通信庫做了網(wǎng)卡拓?fù)涓兄c自適應(yīng)優(yōu)化,每個GPU都能在PCIe拓?fù)渖险业窖舆t和帶寬效率最高的網(wǎng)卡端口進(jìn)行RDMA通信,還額外開發(fā)了一個插件SmartNICAdaptor,在配置4張雙口網(wǎng)卡時,可實現(xiàn)和8張單口網(wǎng)卡一樣的RDMA通信效率。
經(jīng)測試,其AllReduce帶寬可以做到194GB/s,RDMA通信帶寬利用率達(dá)到97%,而一般RDMA通信帶寬利用率約在80%-85%。
除了延遲和帶寬之外,摩爾線程還做了異步通信引擎的通信優(yōu)化。
MPC是計算核心。傳統(tǒng)做法是通信也走M(jìn)PC,就要分走一部分的計算單元。ACE則可以把通信部分卸載到單獨的通信引擎上,不跟MPC搶計算資源,同時在無數(shù)據(jù)依賴時完全做到overlap,通過優(yōu)化,結(jié)合MTTransformerEngine,在Llama模型上端到端計算通信性能可提升10%左右。
來源:紅網(wǎng)
作者:詹夢順
編輯:續(xù)開濟(jì)
本文為紅辣椒評論 原創(chuàng)文章,僅系作者個人觀點,不代表紅網(wǎng)立場。轉(zhuǎn)載請附原文出處鏈接和本聲明。