Résumé Technique : ThriftAttention

Énoncé du Problème

L'inférence efficace des grands modèles de langage (LLM) dans des charges de travail à contexte long est actuellement limitée par le coût computationnel quadratique et le trafic mémoire du mécanisme d'attention. Bien que l'architecture Blackwell de NVIDIA introduise des Tensor Cores natifs en FP4 offrant un débit arithmétique 4 fois supérieur à celui du FP16, leur utilisation pour l'attention via une quantification uniforme (par exemple, SageAttention3) entraîne une dégradation significative de la qualité, particulièrement à mesure que les longueurs de séquence augmentent. À l'inverse, les méthodes d'éparpillement existantes qui tentent d'atténuer ce problème en supprimant entièrement des blocs KV nécessitent souvent des ratios d'éparpillement agressifs (suppression de plus de 75 % des blocs) pour égaler la latence du FP4, conduisant à des pertes de performance irrécupérables. Il existe une tension fondamentale entre l'efficacité de l'inférence et la qualité de la sortie dans les contextes longs.

Méthodologie : ThriftAttention

Les auteurs proposent ThriftAttention, un mécanisme d'attention sélectif à précision mixte sans entraînement, conçu pour offrir une qualité proche du FP16 avec une efficacité d'inférence en FP4. L'idée centrale sous-tendant cette méthode est que l'erreur de quantification n'est pas uniformément distribuée ; plutôt, l'impact de l'erreur est hautement non uniforme et se concentre dans un petit nombre de paires de blocs requête-clé où les scores d'attention sont les plus élevés et les plus fonctionnellement pertinents pour la distribution de sortie finale.

L'algorithme se déroule en deux étapes :

1. Sélection Heuristique de Blocs :

   • Les tenseurs de requête ($Q$) et de clé/valeur ($K, V$) sont partitionnés en blocs.
   • Une heuristique légère calcule un score d'importance pour chaque paire de blocs requête-clé $(i, j)$ en utilisant le produit scalaire des moyennes des tokens : $\hat{S}_{ij} = \bar{q}_i \cdot \bar{k}_j$, où $\bar{q}_i$ et $\bar{k}_j$ sont les vecteurs moyens des blocs respectifs.
   • Les $k$ paires de blocs ayant les scores les plus élevés sont identifiées pour un calcul haute précision.

2. Calcul à Précision Mixte et Fusion :

   • Chemin FP16 : Les $k$ paires de blocs sélectionnées sont calculées en précision FP16 complète.
   • Chemin FP4 : Les blocs restants sont calculés en utilisant la quantification FP4 (utilisant le format micro-échelle NVFP4 pris en charge sur les GPU Blackwell).
   • Fusion en Ligne : Les deux chemins sont calculés et fusionnés via une procédure de softmax en ligne (similaire à FlashAttention-2) pour produire une sortie unique. Le chemin FP4 utilise un schéma de quantification à deux niveaux pour les blocs de probabilité, adapté de SageAttention3, afin de maintenir la stabilité numérique lors de l'accumulation.

L'implémentation est réalisée sous la forme d'un seul noyau CUDA fusionné. Pour gérer la pression sur les registres, le noyau traite d'abord les blocs non sélectionnés via le chemin FP4, puis utilise une routine FP16 séparée pour les blocs promus, en chargeant les tuiles de requête FP16 uniquement lorsque nécessaire. Les warps ne contenant aucun bloc top-$k$ contournent entièrement le chemin FP16 pour éviter des chargements mémoire inutiles.

Contributions Clés

• Attention Sélective à Précision Mixte : L'introduction de ThriftAttention, qui calcule les interactions de blocs les plus critiques en FP16 tout en conservant le reste en FP4. Les auteurs notent qu'il s'agit du premier travail à utiliser des formats sous-octet de cette manière spécifique à précision mixte pour le calcul d'attention.
• Approche Sans Entraînement : La méthode ne nécessite aucun réentraînement ou affinage du modèle, reposant uniquement sur un mécanisme de sélection heuristique lors de l'inférence.
• Évaluation Complète : L'approche est évaluée sur diverses familles de modèles (Llama, Qwen, Ministral) et des benchmarks à contexte long (LongBench-v1, HELMET, RULER, PG-19), démontrant son efficacité sur des architectures et des longueurs de contexte variées.

Résultats

• Récupération de Qualité : Avec un budget de blocs FP16 de 5 % (calculant uniquement 5 % des blocs requête-clé en FP16), ThriftAttention récupère en moyenne 89,1 % de l'écart de performance entre les attentions FP4 et FP16. Cette récupération augmente à 91,8 % et 92,4 % pour des budgets de 10 % et 25 % respectivement.
• Gains d'Efficacité :
  • Préremplissage : Atteint un accélération de noyau allant jusqu'à 1,7× par rapport à FlashAttention-2 (FP16) et des améliorations constantes de bout en bout du préremplissage (environ 1,2× à 131k de contexte).
  • Décodage : Atteint une accélération de 3× à 5,5× par rapport à FlashAttention-2, se traduisant par une accélération de génération de bout en bout d'environ 2× par rapport à l'attention FP16 complète à des longueurs de contexte de 131k.
• Mise à l'Échelle de la Longueur de Contexte : L'avantage de ThriftAttention croît avec la longueur de séquence. Alors que l'attention FP4 uniforme se dégrade considérablement à mesure que la longueur de contexte augmente (par exemple, la rétention FP4 chute de 50 % à 8k à 32 % à 131k pour Llama3.1-8B), l'amélioration relative de ThriftAttention par rapport au FP4 augmente, atteignant 2,2× à 131k.
• Comparaison avec l'Éparpillement : Lorsqu'elles sont comparées pour un calcul équivalent (FLOPs équivalents FP16), ThriftAttention surpasse significativement les bases d'attention éparse à l'inférence (par exemple, Quest, Sparse Top-k). Les auteurs attribuent cela au fait que les méthodes éparse rejettent complètement des blocs, entraînant des erreurs de queue, tandis que ThriftAttention conserve toutes les interactions en FP4, les dégradant progressivement via la quantification plutôt que par suppression.

Importance et Revendications

L'article revendique que ThriftAttention offre une voie pratique vers l'inférence à contexte long qui approche la latence du FP4 tout en préservant une qualité proche du FP16. En identifiant que l'erreur de quantification est concentrée dans des interactions spécifiques à haute importance, la méthode atténue la dégradation systématique de la qualité observée dans les attentions à faible bit uniformes sur des contextes longs.

Les auteurs positionnent ce travail comme une solution à la « tension fondamentale » entre efficacité et qualité dans l'inférence de l'ère Blackwell. Ils soulignent que préserver la prise en charge complète de l'attention en basse précision (FP4), plutôt que d'éparpiller agressivement et de calculer uniquement un sous-ensemble en haute précision, produit des résultats supérieurs avec des budgets de calcul équivalents. L'approche vise spécifiquement les GPU Blackwell grand public et de centre de données pour exploiter les Tensor Cores FP4 natifs, bien que les auteurs notent les limitations actuelles concernant l'empreinte mémoire du cache KV (augmentée de 28 % en raison de la mise en cache double) et l'accent actuel mis sur l'inférence plutôt que sur l'entraînement.