Sand.ai重磅更新MagiAttention，定義分布式Attention性能新標(biāo)桿

機(jī)器之心編輯部

2025 年 4 月，Sand.ai 開源了 MagiAttention v1.0.0，定義了下一代分布式 Attention 的全新設(shè)計(jì)和系統(tǒng)框架。歷經(jīng)一年的深耕，今天Sand.ai正式發(fā)布：MagiAttention v1.1.0，以更成熟的原生算子組件，重新定義 Hopper 與 Blackwell 兩代架構(gòu)分布式 Attention 的性能上限

• 代碼倉(cāng)庫(kù)：https://github.com/SandAI-org/MagiAttention
• 使用文檔：https://sandai-org.github.io/MagiAttention/docs
• 技術(shù)博客：https://sandai-org.github.io/MagiAttention/docs/main/blog/magi_attn

Blackwell 新架構(gòu)適配：釋放下一代算力紅利

為了在下一代硬件上延續(xù) Flex-Flash-Attention (FFA) 的靈活性，研發(fā)團(tuán)隊(duì)引入了基于 Flash-Attention 4 的 FFA_FA4 后端，完成了對(duì) Blackwell 架構(gòu)的初步適配：

• 靈活掩碼支持：引入 HSTU Function 表達(dá)，在無需大幅改動(dòng) FA4 內(nèi)核主體邏輯的前提下，實(shí)現(xiàn)了對(duì)任意掩碼的無縫兼容，且性能損耗 < 5%。
• 高效分塊稀疏生成：開發(fā)了高效分塊掩碼生成算子，大大降低了 FlexAttention 實(shí)例化完整掩碼的性能開銷和顯存風(fēng)險(xiǎn)。
• R2P 指令級(jí)加速：我們利用 Register-to-Predicate 技巧，將復(fù)雜的邊界檢查映射為單條硬件指令，大幅降低了指令周期。
• 極致顯存和延遲：針對(duì)超長(zhǎng)序列，對(duì)掩碼元數(shù)據(jù)進(jìn)行 CSR 壓縮，對(duì)內(nèi)核啟動(dòng)進(jìn)行 FFI 加速，確保顯存與延遲的雙重極致。

原生 Group Collective 原語：突破 RDMA 帶寬瓶頸

跨機(jī)通信效率決定了分布式擴(kuò)展上限，受 DeepEP 啟發(fā)，Sand.ai 構(gòu)建了原生 Group Collective 通信內(nèi)核，徹底重塑了分布式 Attention 在節(jié)點(diǎn)內(nèi)外的數(shù)據(jù)交換范式：

• 算子級(jí)融合：研發(fā)團(tuán)隊(duì)將數(shù)據(jù)重排直接融合進(jìn)通信算子，顯著降低了訪存和額外拷貝開銷。
• RDMA 傳輸去重：以 "NVLink 替代冗余 RDMA" 傳輸，實(shí)現(xiàn)節(jié)點(diǎn)間單次物理交換與節(jié)點(diǎn)內(nèi)高效轉(zhuǎn)發(fā)和規(guī)約，跨機(jī)通信量降低數(shù)倍，性能遠(yuǎn)超傳統(tǒng) AlltoAll-v 方案。

系統(tǒng)級(jí)協(xié)同優(yōu)化：負(fù)載均衡與多階段重疊

MagiAttention 的卓越性能不僅源于算子端的極致打磨，更得益于系統(tǒng)級(jí)的全棧協(xié)同調(diào)度和全場(chǎng)景通用的啟發(fā)式算法：

• Dispatch Solver: 基于最小堆貪心算法，實(shí)現(xiàn)序列的細(xì)粒度分配，保證任意掩碼下設(shè)備的計(jì)算負(fù)載均衡，避免 “短板效應(yīng)” 拖垮分布式整體性能。
• Adaptive Multi-Stage Overlap: 突破靜態(tài)流水線限制，自適應(yīng)調(diào)整流水線階段，通過最小化調(diào)度開銷實(shí)現(xiàn)極致重疊，為超長(zhǎng)序列訓(xùn)練提供線性擴(kuò)展保障。

實(shí)測(cè)表現(xiàn)與應(yīng)用

目前，MagiAttention v1.1.0 已在 Magi-1 等大規(guī)模視頻生成模型訓(xùn)練中得到實(shí)證，也在各大廠中被 “悄悄” 應(yīng)用于多模態(tài)大模型訓(xùn)練。為了驗(yàn)證 MagiAttention 在真實(shí)長(zhǎng)文訓(xùn)練中的表現(xiàn)，Sand.ai 也給出了細(xì)致的 Benchmark 結(jié)果：

• 算子層面：支持靈活掩碼并維持 SOTA 性能

H100/B200 Varlen Causal 掩碼下前反向的內(nèi)核算子性能對(duì)比

• 分布式層面：重塑超長(zhǎng)序列的擴(kuò)展性曲線

H100/B200 Varlen Causal 掩碼下前反向的分布式性能對(duì)比

結(jié)語與未來展望

自去年 v1.0 發(fā)布以來，Sand.ai 收到的社區(qū)反饋?zhàn)屟邪l(fā)團(tuán)隊(duì)更加堅(jiān)定：只有將底層算力壓榨到極限，才能開啟人工智能處理復(fù)雜多模態(tài)任務(wù)的新篇章。MagiAttention v1.1.0，是 Sand.ai 向這一愿景邁進(jìn)的關(guān)鍵一步。Sand.ai 相信，強(qiáng)大的模型能力必須建立在普惠且極致的技術(shù)基石之上。

了解更多信息，歡迎訪問 Sand.ai 官網(wǎng)：https://sand.ai