FlashMoBA

FlashMoBA是一種最佳化注意力機制，是一種硬體感知的CUDA核心，旨在使小塊MoBA變得實用且高效。它採用三個融合核心以減少HBM往返次數，用Flash TopK替換了Top-k選擇過程，實現了完全並行化。在性能上，比原始MoBA快7.4倍，記憶體占用減少6.1倍，能擴展到512K序列而不溢出。在長序列和小塊場景下，相比FlashAttention-2可實現最高14.7倍加速。

基本介紹

• 開發團隊：麻省理工學院和英偉達團隊
• 技術特點：硬體感知的CUDA核心
• 提出時間：2025年11月18日
• 類型：最佳化注意力機制

開發背景

FlashMoBA由麻省理工學院和英偉達團隊的韓松等人於2025年11月18日提出。旨在解決原始MoBA在小塊配置下效率低、記憶體占用高，在長序列時易記憶體溢出的問題。

最佳化後的FlashMoBA能擴展到512K序列而不溢出。在序列長度為64K、批次為128的配置下，FlashMoBA比原始MoBA快7.4倍，記憶體占用減少6.1倍。

技術原理

FlashMoBA是一種硬體感知的CUDA核心，旨在使小塊MoBA變得實用且高效。它採用了三個融合核心，以最大限度地減少HBM往返次數。

關鍵算法改進將原始MoBA中顯式生成完整分數矩陣並串列處理的Top-k選擇瓶頸替換為完全並行化的Flash TopK。整個流水線在批次和注意力頭之間完全並行化，消除了原始的性能瓶頸。

反向傳播利用了FlashAttention-2的記憶體高效設計，並實現為三個核心的序列。

通過上述最佳化，FlashMoBA相比原始MoBA實現了顯著的加速和記憶體節省；在N=64K且B=128的配置下，FlashMoBA比原始MoBA快7.4倍，記憶體占用減少6.1倍。原始MoBA在128K序列就會記憶體溢出，而FlashMoBA能擴展到512K。

性能表現

FlashMoBA在延遲和記憶體占用方面的表現優於原始MoBA實現。在序列長度64K、批次大小128的配置下，FlashMoBA的處理速度比原始MoBA快7.4倍，記憶體占用減少6.1倍。原始MoBA在128K序列長度時會發生記憶體溢出，而FlashMoBA能擴展到512K。隨著序列長度增加和塊大小減小，其優勢進一步擴大，在長序列條件下相比FlashAttention-2可實現最高14.7倍的加速。在64K序列下，FlashMoBA的前向傳播時間為49毫秒，而FlashAttention-2在相同設定下為99毫秒。