FlashMask-3: attenzione distribuita efficiente ed espressiva che tiene conto delle maschere

Abstract degli autori

L’addestramento di grandi modelli linguistici e multimodali con contesti lunghi combina maschere di attenzione strutturate (per documenti, prefix-LM, a domanda condivisa, a finestra scorrevole, a blocchi sparsi) e parallelismo del contesto su più nodi, rendendo la maschera un collo di bottiglia dominante. La maschera sparsa per colonne di FlashMask comprime queste maschere da $O(n^2)$ a $O(n)$ mediante quattro vettori dei confini lungo l’asse delle query, coprendo un’ampia gamma di configurazioni strutturate. Implementato per le GPU Ampere, FlashMask non sfrutta la pipeline TMA/WGMMA di FlashAttention-3, con specializzazione dei warp, su Hopper, né supporta l’esecuzione distribuita. La maschera sparsa per colonne è indicizzata in base alle posizioni delle chiavi, ma memorizza i confini degli intervalli lungo l’asse delle query. Con qualsiasi assegnazione del contesto parallelo su più chunk, riduce la localizzazione della maschera per ciascun rank a un’operazione \emph{clip--shift--max} sui vettori dei confini che, come si può dimostrare, non richiede comunicazione ed è compatibile con il kernel per una singola GPU. FlashMask-3 introduce un kernel per Hopper che tiene conto della maschera e specializza i warp in tre ruoli sulla pipeline TMA/WGMMA di FlashAttention-3, oltre a un’estensione distribuita che tiene conto della maschera: localizzazione della maschera senza comunicazione, un bilanciatore del carico sensibile alla sparsità (HCC-IPO) e una sovrapposizione tra calcolo e comunicazione che tiene conto della maschera e riduce le comunicazioni in base a essa. L’all-gather della fase forward di quest’ultima usa un instradamento gerarchico che tiene conto della topologia. Su una singola H100, con 12 maschere e contesti da 4K a 128K, aumenta i TFLOPs raggiunti rispetto a FlashMask, FlexAttention e MagiAttention rispettivamente del 45–141%, del 34–71% e del 2–44%. Con gradi di parallelismo del contesto da 4 a 32 e sequenze da 8K per dispositivo, raggiunge un’efficienza della sovrapposizione tra calcolo e comunicazione del 31–75% e aumenta il throughput rispetto a Megatron-Core CP ring, Megatron-Core CP all-gather e MagiAttention rispettivamente di $1.80\times$--$3.71\times$, $1.71\times$--$2.92\times$ e $1.09\times$--$1.46\times$. Nell’addestramento end-to-end di un MoE 118B-A14B con contesto da 128K su 128 GPU H100, aumenta il throughput del 7,3–11,0% rispetto a una baseline Dual Chunk all-gather, con le fasi forward+backward dell’attenzione più veloci del 24,5–44,9%.