FLAMIA: attenzione con microscaling flash per un addestramento rapido e causalmente sicuro su Blackwell

Abstract degli autori

L’attenzione continua a rappresentare un costo computazionale significativo nell’addestramento dei Transformer, soprattutto con contesti lunghi. Le GPU NVIDIA Blackwell offrono moltiplicazioni di matrici a bassa precisione e ad alto throughput, ma per sfruttarle nell’attenzione occorre affrontare l’overhead della quantizzazione, la fuga di informazioni causali dovuta all’uso di scale condivise e le incoerenze numeriche tra i passaggi forward e backward. Presentiamo FLAMIA, una famiglia di kernel di attenzione per Blackwell con percorsi di addestramento MXFP8 e misti NVFP4/FP8. Il nostro approccio combina una scalatura a blocchi di token, che riutilizza operandi quantizzati in diversi orientamenti delle matrici, misure di sicurezza mirate contro le fughe indotte dalle scale e calcoli coerenti nei passaggi forward e backward. La quantizzazione fusa e le ottimizzazioni dei kernel specifiche per l’architettura riducono l’overhead, mentre la scalatura locale ai documenti e una pianificazione consapevole del carico di lavoro permettono di gestire direttamente documenti di lunghezza variabile in layout impacchettati, senza reimpacchettarli. In diversi carichi di lavoro causali con documenti impacchettati, il percorso MXFP8 accelera la chiamata combinata di attenzione forward e backward, inclusa la quantizzazione, di fattori pari a $1.10$--$1.40\times$ su B200 e $1.06$--$1.42\times$ su B300 rispetto alla baseline valutata più veloce per ciascun carico di lavoro. Convalidiamo la stabilità dell’addestramento e il mantenimento della qualità nel preaddestramento di modelli in stile Llama. Questi risultati mostrano come la progettazione congiunta della quantizzazione, degli aspetti numerici dell’addestramento e della pianificazione dei documenti permetta di tradurre il throughput dell’hardware a bassa precisione in un’accelerazione pratica dell’attenzione.