Résumé Technique : NOSA : Attention Éparse Native et Déchargeable (Native and Offloadable Sparse Attention)

Énoncé du Problème

Le débit de décodage des grands modèles de langage (LLM) est de plus en plus contraint par les limitations de la mémoire GPU, principalement en raison des exigences de stockage du cache Clé-Valeur (KV). Bien que le déchargement (offloading) du cache KV (garder le contexte redondant sur le CPU et récupérer des sous-ensembles vers le GPU) atténue la pression sur la mémoire, les approches existantes souffrent d'un compromis critique :

1. Déchargement sans entraînement (Training-Free Offloading) : Des méthodes comme ShadowKV et InfLLM sélectionnent des sous-ensembles épars sans entraînement. Cela crée un décalage entraînement-inférence, où les motifs épars non vus lors de l'entraînement entraînent une dégradation significative des performances, particulièrement dans les scénarios de génération longue où les erreurs de sélection s'accumulent.
2. Attention éparse entraînable (Trainable Sparse Attention) : Les méthodes qui apprennent des motifs épars (par exemple, InfLLMv2) évitent ce décalage mais sont incompatibles avec un déchargement efficace. Sans contraintes, ces méthodes peuvent récupérer des blocs KV arbitraires du CPU, provoquant un overhead de communication PCIe massif qui annule les gains de débit.

Les solutions actuelles ne parviennent pas à atteindre simultanément une haute efficacité de décodage (via le déchargement) et une haute qualité de génération (via la cohérence de l'entraînement).

Méthodologie

Les auteurs proposent NOSA (Native and Offloadable Sparse Attention), un mécanisme d'attention éparse entraînable explicitement conçu pour le déchargement du cache KV, couplé à NOSI, un système d'inférence spécialisé.

1. Algorithme : NOSA

NOSA introduit une contrainte de localité pendant l'entraînement pour garantir que les motifs d'attention éparse restent stables à travers les étapes de décodage, minimisant ainsi les transferts de données CPU-GPU.

• Mécanisme de double sélection : NOSA décompose la sélection KV en deux composantes :
  • Sélection sensible à la requête ($\Gamma_q$) : Utilise la requête pour récupérer des informations spécifiques, potentiellement distantes (préservant le rappel).
  • Sélection indépendante de la requête ($\Gamma_e$) : Utilise une tête d'éviction (inspirée de DMA) pour attribuer des scores d'importance indépendamment de la requête. Cette composante impose une politique d'éviction stricte : si un jeton n'est pas sélectionné à l'étape $t$, il ne peut pas être sélectionné à une étape future ultérieure $t' > t$.
• Garantie de localité : En allouant un budget fixe $k_e$ à la sélection indépendante de la requête, NOSA garantit mathématiquement une localité minimale $\gamma(t) \geq k_e/k$. Cela assure que la majorité des blocs KV sélectionnés à l'étape $t$ ont également été sélectionnés à l'étape $t-1$, permettant de les maintenir en cache sur le GPU.
• Implémentation par blocs : Pour maximiser l'utilisation de la bande passante PCIe, NOSA opère sur des blocs de jetons plutôt que sur des jetons individuels. Il emploie une stratégie de pooling à deux étapes (pooling moyen au sein des sous-blocs, suivi d'un pooling max au niveau du bloc) pour calculer les scores de sélection.
• ED-DMA (Exp-Delayed DMA) : Pour gérer les problèmes de précision numérique dans la tête d'éviction, NOSA diffère l'opération exponentielle jusqu'à l'étape de calcul de l'attention, sélectionnant les blocs sur la base de scores pré-exponentiels.

2. Système : NOSI

Les auteurs observent que les implémentations vanilla (par exemple, Hugging Face Transformers) introduisent des surcoûts excessifs de lancement de noyaux (kernel launch overheads) qui masquent les véritables goulots d'étranglement de communication : NOSI est un système d'inférence personnalisé construit pour exposer et optimiser ces goulots d'étranglement :

• Fusion de noyaux (Kernel Fusion) : Fusionne les calculs de la tête d'éviction, la division QKV et les opérations de pooling en noyaux Triton uniques pour réduire l'overhead de lancement.
• Organisation de la mémoire optimisée : Organise le cache KV en unités de taille de bloc minimale et utilise des noyaux CUDA personnalisés pour calculer efficacement les différences d'ensembles pour l'échange de blocs (block swapping), évitant les exigences de mémoire contiguë.
• Communication à large bande passante : Implémente des noyaux de transfert hôte-vers-périphérique personnalisés utilisant l'adressage virtuel unifié (UVA) pour atteindre une bande passante PCIe proche du pic, surmontant les problèmes de fragmentation des implémentations PyTorch standards.

Contributions Clés

1. Analyse de faisabilité : L'article identifie que si l'attention éparse entraînable (comme InfLLMv2) possède une localité inhérente, elle est insuffisante pour surmonter les limitations de la bande passante PCIe sans contraintes explicites.
2. Algorithme NOSA : Un nouveau mécanisme d'attention éparse entraînable qui limite explicitement le volume des transferts KV CPU-GPU en imposant une contrainte de localité via une politique d'éviction indépendante de la requête.
3. Système NOSI : Un système d'inférence adapté à NOSA qui élimine les surcoûts logiciels, révélant la communication comme le véritable goulot d'étranglement et permettant un déchargement à haut débit.
4. Évaluation complète : Benchmarks étendus sur des modèles de 1B, 3B et 8B démontrant que NOSA surpasse à la fois les bases de comparaison de déchargement sans entraînement et les autres méthodes d'attention éparse entraînables.

Résultats Expérimentaux

Les auteurs ont évalué NOSA sur des modèles de 1B, 3B et 8B à travers des tâches générales, des entrées à contexte long et des tâches de raisonnement à génération longue.

• Performance : NOSA surpasse systématiquement les bases de comparaison de déchargement sans entraînement (ShadowKV, InfLLM, ArkVale) et les bases entraînables (InfLLMv2, DMA) sur les benchmarks de contexte long (LongBench, HELMET). Il maintient un rappel élevé et des capacités de raisonnement là où les méthodes sans entraînement se dégradent significativement à cause du décalage entraînement-inférence.
• Débit :
  • NOSA atteint un débit de décodage jusqu'à 5,04× supérieur à FullAttn (grâce à des tailles de lots plus grandes rendues possibles par le déchargement).
  • Il atteint un débit jusqu'à 1,92× supérieur à InfLLMv2 (grâce à la réduction de la communication par les contraintes de localité).
  • Il atteint un débit jusqu'à 1,83× supérieur à ShadowKV (grâce à une meilleure efficacité dans la gestion des motifs épars et aux optimisations système).
• Longue Génération : Dans les tâches de raisonnement à génération longue, NOSA maintient une faible perplexité, tandis que les méthodes de déchargement sans entraînement présentent une accumulation d'erreurs et une explosion de la perplexité.
• Ablation : L'étude confirme que le motif de sélection par blocs est critique pour l'efficacité PCIe et que l'implémentation de la tête d'éviction ED-DMA produit les meilleures performances.

Signification et Revendications

L'article affirme que NOSA représente un passage nécessaire des approches "sans entraînement" ou "entraînables non contraintes" vers une co-conception algorithme-système. En contraignant explicitement le volume des transferts KV pendant l'entraînement, NOSA comble le fossé entre le besoin d'efficacité mémoire (déchargement) et le besoin de qualité de génération (cohérence de l'entraînement).

Les auteurs déclarent que leur travail permet le déploiement de LLM dotés de capacités de contexte long sur du matériel à mémoire limitée sans sacrifier la qualité du raisonnement. Ils soulignent que leur approche réduit les coûts de calcul et d'énergie en améliant l'efficacité du décodage, facilitant potentiellement un déploiement plus large des applications de LLM. Le code est publié pour encourager la recherche ultérieure sur les mécanismes de déchargement natifs.