FlashAttention-4 comble les limitations de performance des cœurs tensoriels 4 bits (FP4) de Blackwell en introduisant de nouveaux chemins pour l'inférence non causale et l'entraînement causal. La méthode, appelée Direct-P, mappe directement les scores aux probabilités FP4 afin de contourner les surcoûts de conversion softmax qui dominent généralement lorsque les produits matriciels diminuent.

  • Direct-P atteint jusqu'à 2,13 fois le débit avant (forward throughput) bfloat16 (BF16) sur un NVIDIA GB200 pour l'inférence non causale.
  • Le chemin causal transmet la quantification directe vers l'arrière, reconstruisant les probabilités à partir des requêtes et clés quantifiées sauvegardées en utilisant des opérandes de gradient en virgule flottante 8 bits (FP8).
  • Cette approche causale accélère une mise à jour complète sur un seul GPU de 8 milliards de paramètres jusqu'à 1,14 fois.
  • L'entraînement distribué apparié conserve les probabilités et valeurs FP8, car chaque trajectoire d'entraînement de probabilité/valeur MXFP4 testée diverge.

Les auteurs considèrent cela important car cela permet une attention FP4 efficace sensible au matériel sur les architectures Blackwell, en surmontant les dépendances qui empêchent les accélérations automatiques des cœurs tensoriels FP4.