当前位置: 首页 » 资讯 » 科技头条 » 正文

摩尔线程吴庆:MUSA三大创新落地,助力大规模集群训练效率提升

IP属地 中国·北京 编辑:钟景轩 爱集微 时间:2025-08-02 00:16:48

强大的芯片有效算力是作为“AI工厂”的核心竞争力。而作为GPU的 “灵魂”,其软件能力深刻决定了硬件潜力的释放程度,特别是驱动与全栈系统软件,是解锁大规模集群极致性能效率的关键。

7月25日,在WAIC2025期间举办的摩尔线程技术分享日上,摩尔线程GPU计算软件开发总监吴庆以《MUSA软件栈助力KUAE集群释放无限潜能》为主题,分享了摩尔线程MUSA软件栈方面的创新实践。

以下为吴庆的演讲实录精华:

MUSA驱动和运行时库,KUAE集群坚实底座

推理场景对Kernel延时非常敏感。极致高效的MUSA驱动和运行时库,具有高效率的任务调度能力,使用户的Kernel launch开销极致缩减。


(1)即时任务下发:通过软硬协同,将核函数启动延迟降低到业界平均水平的1/2。举个例子方便大家理解,比如现在刚好是仲夏,外面气温非常高,我们的车子停在外面暴晒,那如何让大家一上车即时就享受到舒适的温度?现在普遍的解决方案是软硬协同,车要支持远程控制,还没有到车上的时候就通过手机提前把车上的空调打开,这样一上车以后马上就有清凉舒适的温度。这个原理类似,我们把很多调度开销前置,上一个Kernel执行结束之前,就可以把下一个Kernel要做的配置提前做好,这样就最大限度地降低Kernel跟Kernel之间的调度延时开销。


(2)批量任务下发:批量下发计算和通信任务,将近千次下发开销优化为单次,减少GPU等待时间。为方便大家理解,我再类比一下,车上座椅调节,调舒适位子有上下前后,包括靠背俯仰角度,有很多调节步骤对应很多Kernel,但我们可以调节一次、打包,变成一键记忆功能,下次再上车,不需要再重复那个步骤了。这就类似于批量任务下发,一系列操作打包、我们一次kick下去。Graph capture阶段对应第一次调节,后面再launch一个Graph的时候就类似于我们唤醒座椅某个记忆配置。Graph这种批量任务下发的收益取决于Graph里面Kernel本身执行时间、调度空泡的占比,Kernel执行得越快,调度空泡在端到端总开销里占比越高,批量任务下发的端到端收益就越高,根据不同的应用场景可达到百分之几十甚至好几倍的性能收益。


(3)引擎间依赖解析:GPU是多引擎可配置的,不同引擎之间要做依赖解析、交互同步,平湖支持硬件解析引擎间的依赖做同步,不再回到host,任务流之间的依赖解析延时可以大幅降低至1.5μs,优于业界头部算力卡。


(4)调优工具接口MUPTI:专门用于性能分析和事件追踪,帮助开发者优化MUSA应用程序,可基于MUPTI开发第三方的各种丰富工具。

(5)GPU错误转存功能(GPU Core Dump,GCD):高效定位疑难Bug,不需要去调度大量的集群和人力反复去复现“昙花一现”的错误,用于在GPU程序崩溃或发生错误时生成详细的信息(类似CPU程序Core Dump),精确保存错误,提供一个log,极大提升MUSA核函数崩溃、非法显存地址访问、硬件错误等问题定位的效率。

MUSA算子库生态日趋完善,性能与广度兼备


算子库的效率直接关乎分布式集群的训练效率。MUSA算子库不仅追求极致性能,还兼容覆盖广度和兼容性。

MUSA算子库生态会提供三大算子库:极致性能muDNN、易用MUTLASS、MUSA AI Tensor Engine开源推理算子库。

muDNN,是一款极致性能的开箱即用标准算子库,完整覆盖常见的前向和反向算子。其中最主要的特性是支持完整的XMMA,支持Tensor Core全精度及所有量化模式,以及常用的神经网络算子操作。


一般来说,矩阵乘的效率大部分可以做到90%以上,摩尔线程muDNN矩阵乘算子效率可以做到98%。Flash Attention算子因为多了softmax等操作,基于国际一流水平厂商的Flash Attention 3论文里提到的效率大约是75%,muDNN的Flash Attention可以做到95%。

影响FP8 GEMM算子效率的重要因素是scale模式,scale是为了避免精度损失做的一个精细化的缩放因子,常见的有Per-Tensor、Per-Block,DeepSeek V3用的是Per-Block。

吴庆打了个比方,Per-Tensor scale是整个张量共用一个缩放因子,相当于一面墙刷一个颜色,很好刷;Per-Block scale则是每个小的矩阵块共用一个缩放因子,相当于一面墙有很多小方块,每个方块单独一个颜色,刷墙的复杂度要高很多,会比一个颜色刷一面墙效率低10%-20%。

摩尔线程通过软硬协同的深度算法优化、指令编排,做到了Per-Block跟Per-Tensor FP8 GEMM计算效率几乎相当,差距不到2%。

MUTLASS,是一个高性能的线性代数模板库,可极大降低在MUSA环境中自定义算子的开发工作量,相当于提供了一个可供二次开发的模板。

MUTLASS已在Github上开源,支持平湖架构所有特性,包括Warp-level MMA、WarpSquad-level MMA、TensorMemoryEngine、AsyncBarrier等原语全面支持,同时提供高性能矩阵乘、卷积算子实现,可在Kernel里调用模板库就可以进行二次开发,摩尔线程也提供Attention最优化的示例,给用户打样,可以参考提供的最佳实践,去开发自定义的各种Attention的变种。

MUSA AI Tensor Engine,是摩尔线程即将发布的面向大语言模型的开源推理算子库,可帮助开发者快速搭建自定义推理引擎。 该算子库提供用户友好的Python API,进一步降低开发者使用门槛,无需花太多时间学习MUSA C,只要会Python就可以把MTX这个库用起来;未来会支持丰富的后端,包括模板库,对标PTX虚拟指令集去开发的极致算子库,也支持类Triton的DSL。

高性能集合通信库及通算并行,避免通信抢占计算核心资源

大模型分布式训练的主要瓶颈是通信。摩尔线程在MTT S5000上支持MTLINK 2.0,配备GPU异步通信引擎,实现了高性能通信技术与计算高度并行。


最近十几年,随着模型规模增加,分布式集群规模也在增加,通信瓶颈会越来越明显。DeepSeek V3报告中呼吁硬件厂商提供的芯片能不能做一个单独的通信硬件来卸载通信,避免通信跟计算去抢SM计算单元。

平湖设计阶段远早于DeepSeek V3发布的时间,彼时摩尔线程已洞察到行业痛点,并提供了解决方案——在GPU上增加了一个异步通信引擎。

MTT S5000服务器拓扑每节点有8张GPU,通过MTLINK 2.0全互连,每张GPU与其他7张GPU都有直连总线。

每个GPU上配置的异步通信引擎是原生支持丰富的reduce操作,支持常用的reduce操作有ADD、MIN、MAX(累加、最小值、最大值)等操作,支持Float、FP16、BF16等数据类型,可直接通过MTLink进行C2C跨系统的数据通信。

其ACE通过Zero Copy技术进一步提升性能。一般的集合通信会把数据buffer拷贝到通信buffer里,再去进行C2C通信。Zero Copy则避免了在本地上的D2D的拷贝,直接把数据buffer做跨卡通信。

摩尔线程在衡量集合通信库性能的时候主要是两大核心指标:通信延迟、通信带宽。


基于全互联(FC8)拓扑的高效通信算法,可大幅降低通信延迟。全互联拓扑中GPU两两互连,在All Reduce场景中可以把其他GPU的数据一次性拿过来,理论上FC8算法相比Ring算法,7步通信可变成1步完成。在单机8卡All Reduce延迟场景的实测也差不多,Ring算法大概是53us,用FC8算法可降低至7.8us,将近缩减至1/7,显著提升了模型推理性能。


带宽方面,Ring算法与FC8算法效率相当,以Ring算法为例,摩尔线程把FC8拆成了7个通信环,每个通信环结合通信Kernel的优化,基于MTLink2.0 + FC8拓扑,节点内充分利用全互联的拓扑优势,可做到单机8卡All Reduce带宽利用率接近85%,与国际第一的厂商在软件效率上对齐,实现高效的Scale-up通信。

高效Scale-out跨节点通信方面,摩尔线程结合通信库做了网卡拓扑感知与自适应优化,每个GPU都能在PCIe拓扑上找到延迟和带宽效率最高的网卡端口进行RDMA通信,还额外开发了一个插件Smart NIC Adaptor,在配置4张双口网卡时,可实现和8张单口网卡一样的RDMA通信效率。

经测试,其All Reduce带宽可以做到194GB/s,RDMA通信带宽利用率达到97%,而一般RDMA通信带宽利用率约在80%-85%。


除了延迟和带宽之外,摩尔线程还做了异步通信引擎的通信优化。

MPC是计算核心。传统做法是通信也走MPC,就要分走一部分的计算单元。ACE则可以把通信部分卸载到单独的通信引擎上,不跟MPC抢计算资源,同时在无数据依赖时完全做到overlap,通过优化,结合MT Transformer Engine,在Llama模型上端到端计算通信性能可提升10%左右。

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。