FlashAttention-2: riscrittura completa, circa 2x più veloce, con MQA/GQA e head-dim fino a 256
In una frase Il 17 luglio 2023 Tri Dao pubblica FlashAttention-2: circa 2x più veloce di FlashAttention, dal 50 al 73% dei FLOPs teorici su A100 (contro 25-40%), parallelismo anche sulla lunghezza di sequenza, supporto a MQA/GQA e head dimension fino a 256.
Un anno dopo FlashAttention, Tri Dao ne pubblica una riscrittura completa. La versione originale non sfruttava al massimo la GPU: il lavoro veniva diviso solo per batch e per teste di attenzione, e con sequenze lunghe e batch piccoli molte unità di calcolo restavano ferme.
FlashAttention-2 divide il lavoro anche lungo la sequenza, riduce le operazioni che non sono moltiplicazioni di matrici (lente sulle GPU moderne) e organizza meglio il lavoro tra i gruppi di thread. Supporta head dimension fino a 256 e le varianti di attenzione MQA e GQA, usate per ridurre la memoria della KV cache.
Il risultato è circa il doppio della velocità rispetto a FlashAttention.
Aziende
Princeton University, Stanford University
Tool
FlashAttention-2, CUDA, PyTorch
Tag
Fonti