Sand.ai MagiAttention v1.1.0:重新定义分布式Attention性能标杆

Sand.ai发布MagiAttention v1.1,通过原生算子组件和系统级优化,重新定义分布式Attention性能上限,加速多模态大模型训练。

原文标题:Sand.ai重磅更新MagiAttention,正在定义分布式Attention性能新标杆

原文作者:数据派THU

冷月清谈:

Sand.ai 发布 MagiAttention v1.1.0,重新定义 Hopper 与 Blackwell 两代架构分布式 Attention 的性能上限。该版本引入了基于 Flash-Attention 4 的 FFA_FA4 后端,完成了对 Blackwell 架构的初步适配,支持灵活掩码和高效分块稀疏生成,并利用 R2P 指令级加速。MagiAttention 还构建了原生 Group Collective 通信内核,重塑了分布式 Attention 在节点内外的数据交换范式。此外,通过 Dispatch Solver 和 Adaptive Multi-Stage Overlap 等系统级协同优化,实现了负载均衡和极致重叠。目前,MagiAttention v1.1.0 已在 Magi-1 等大规模视频生成模型训练中得到实证,并在各大厂中被广泛应用于多模态大模型训练。

怜星夜思:

1、MagiAttention v1.1.0 针对 Blackwell 架构的适配,具体体现在哪些方面?除了文章中提到的灵活掩码支持、高效分块稀疏生成和 R2P 指令级加速外,还有哪些潜在的优化方向?
2、文章提到 MagiAttention 构建了原生 Group Collective 通信内核,解决了 RDMA 带宽瓶颈。相比传统的 AlltoAll-v 方案,这种方案有哪些优势和局限性?在实际应用中,如何选择合适的通信方案?
3、MagiAttention 在大规模视频生成模型 Magi-1 中得到了实证,那么除了视频生成外,MagiAttention 还有哪些潜在的应用场景?在这些场景中,MagiAttention 能发挥哪些优势?

原文内容

图片
来源:机器之心
本文约1500字,建议阅读5分钟
开启人工智能处理复杂多模态任务的新篇章。


2025 年 4 月,Sand.ai 开源了 MagiAttention v1.0.0,定义了下一代分布式 Attention 的全新设计和系统框架。历经一年的深耕,日前Sand.ai 正式发布:MagiAttention v1.1.0,以更成熟的原生算子组件,重新定义 Hopper 与 Blackwell 两代架构分布式 Attention 的性能上限。
  • 代码仓库:https://github.com/SandAI-org/MagiAttention

  • 使用文档:https://sandai-org.github.io/MagiAttention/docs

  • 技术博客:https://sandai-org.github.io/MagiAttention/docs/main/blog/magi_attn


Blackwell 新架构适配:释放下一代算力红利

为了在下一代硬件上延续 Flex-Flash-Attention (FFA) 的灵活性,研发团队引入了基于 Flash-Attention 4 的 FFA_FA4 后端,完成了对 Blackwell 架构的初步适配:

  • 灵活掩码支持引入 HSTU Function 表达,在无需大幅改动 FA4 内核主体逻辑的前提下,实现了对任意掩码的无缝兼容,且性能损耗 < 5%。

  • 高效分块稀疏生成开发了高效分块掩码生成算子,大大降低了 FlexAttention 实例化完整掩码的性能开销和显存风险。

  • R2P 指令级加速我们利用 Register-to-Predicate 技巧,将复杂的边界检查映射为单条硬件指令,大幅降低了指令周期。

  • 极致显存和延迟针对超长序列,对掩码元数据进行 CSR 压缩,对内核启动进行 FFI 加速,确保显存与延迟的双重极致。


原生 Group Collective 原语:突破 RDMA 带宽瓶颈

跨机通信效率决定了分布式扩展上限,受 DeepEP 启发,Sand.ai 构建了原生 Group Collective 通信内核,彻底重塑了分布式 Attention 在节点内外的数据交换范式:

  • 算子级融合研发团队将数据重排直接融合进通信算子,显著降低了访存和额外拷贝开销。

  • RDMA 传输去重以 "NVLink 替代冗余 RDMA" 传输,实现节点间单次物理交换与节点内高效转发和规约,跨机通信量降低数倍,性能远超传统 AlltoAll-v 方案。


系统级协同优化:负载均衡与多阶段重叠

MagiAttention 的卓越性能不仅源于算子端的极致打磨,更得益于系统级的全栈协同调度和全场景通用的启发式算法:

  • Dispatch Solver基于最小堆贪心算法,实现序列的细粒度分配,保证任意掩码下设备的计算负载均衡,避免 “短板效应” 拖垮分布式整体性能。

  • Adaptive Multi-Stage Overlap突破静态流水线限制,自适应调整流水线阶段,通过最小化调度开销实现极致重叠,为超长序列训练提供线性扩展保障。


实测表现与应用

目前,MagiAttention v1.1.0 已在 Magi-1 等大规模视频生成模型训练中得到实证,也在各大厂中被 “悄悄” 应用于多模态大模型训练。为了验证 MagiAttention 在真实长文训练中的表现,Sand.ai 也给出了细致的 Benchmark 结果:

  • 算子层面:支持灵活掩码并维持 SOTA 性能


H100/B200 Varlen Causal 掩码下前反向的内核算子性能对比

  • 分布式层面:重塑超长序列的扩展性曲线


H100/B200 Varlen Causal 掩码下前反向的分布式性能对比

结语与未来展望

自去年 v1.0 发布以来,Sand.ai 收到的社区反馈让研发团队更加坚定:只有将底层算力压榨到极限,才能开启人工智能处理复杂多模态任务的新篇章。MagiAttention v1.1.0,是 Sand.ai 向这一愿景迈进的关键一步。Sand.ai 相信,强大的模型能力必须建立在普惠且极致的技术基石之上。

编辑:文婧



关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。




新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU



视频生成只是 Attention 的一个应用方向,Attention 机制本身在很多领域都有用武之地。例如,自然语言处理里面的机器翻译、文本摘要,Attention可以帮助模型更好地理解上下文关系。还有图像识别,Attention 可以让模型更关注图像中的关键区域。MagiAttention 如果能把分布式 Attention 的性能优势发挥出来,应该能在这些领域提高模型训练效率。

楼上两位说的有道理!我补充一点,可以关注一下 Blackwell 对混合精度计算的支持情况。MagiAttention 如果能更好地利用混合精度计算,应该也能带来性能提升。

Group Collective 的核心在于优化了跨机通信的数据交换方式,减少了冗余传输,这比传统的 AlltoAll-v肯定高效得多。优点很明显,降低了跨机通信量,性能更好。但局限性可能在于实现复杂度较高,需要深度定制。选择合适的方案需要综合考虑集群规模、网络拓扑、数据分布等因素,可能需要根据实际情况进行调整。

我理解 Group Collective 的优势在于它更聪明地组织了数据传输。AlltoAll-v 就像是所有人互相发送邮件,有很多重复内容;而 Group Collective 就像是有个邮局,先把邮件分好类,再统一发送。但是,如果数据量很小,或者网络延迟很高,Group Collective 的优势可能就不明显了。这时候,简单的 AlltoAll-v 反而可能更快。

我觉得在需要处理长序列数据的场景,MagiAttention 应该都能派上用场。比如,基因序列分析、金融时序数据预测等等。这些场景的数据量通常很大,计算复杂度也很高,MagiAttention 的高性能和分布式特性能够减轻计算压力。

从一个普通用户的角度来看,我希望 MagiAttention 能够帮助我更好地理解长篇小说!现在的AI阅读理解能力还是有点弱,如果能用上 MagiAttention,让 AI 更好地理解上下文,那我就可以直接问 AI 一些复杂的问题,不用自己啃书了。

从我做硬件的角度来看,Blackwell 的进步不仅仅是算力,更重要的是架构上的革新。我觉得 MagiAttention 要想充分发挥 Blackwell 的潜力,需要深入理解 Blackwell 架构的特性。除了文章中提到的几点,我觉得还可以关注以下几个方面:1. 内存访问优化,Blackwell 的内存墙问题是否得到改善?2. 片上网络优化,数据传输效率如何?3. 新的指令集支持,是否能加速 Attention 计算?

从系统设计的角度来看,通信方案的选择要trade-off很多东西。Group Collective 降低了跨机通信量,但增加了节点内的计算和数据重排开销。AlltoAll-v 简单粗暴,但容易造成网络拥塞。实际应用中,需要找到一个平衡点。可以考虑根据数据量大小、节点间延迟等因素,动态选择通信方案。

文章里提到了三个关键点:灵活掩码支持(引入 HSTU Function),高效分块稀疏生成(降低实例化完整掩码的开销),以及 R2P 指令级加速(降低指令周期)。我觉得未来的优化方向可以考虑和 Blackwell 架构的 Tensor Core 更紧密结合,进一步提升计算密度。另外,Blackwell 的 NVLink 带宽更高,可以探索更激进的通信优化策略。