🤖 machine learning

StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation

StreamKL introduit une nouvelle primitive GPU fusionnée qui élimine les goulots d'étranglement quadratiques de la mémoire et des E/S de la distillation d'attention en diffusant les tuiles de requêtes-clés en un seul passage, réalisant des accélérations significatives et réduisant l'empreinte mémoire de $O(N_QN_K)$ à $O(1)$ pour permettre la distillation de contexte long sur un seul GPU.

Auteurs originaux : Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

Publié 2026-06-19

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'enseigner à un petit étudiant rapide (un « modèle étudiant ») à penser exactement comme un professeur brillant et lent (un « modèle enseignant »). Dans le monde de l'IA, on fait cela en comparant la façon dont ils « prêtent tous deux attention » à différentes parties d'une histoire ou d'une phrase. Ce processus est appelé Distillation d'Attention.

Pour effectuer cette comparaison, l'ordinateur calcule un nombre spécifique appelé Divergence KL. Considérez cela comme un « score de distance » qui vous indique à quel point l'attention de l'étudiant diffère de celle de l'enseignant. Le but est de rendre ce score aussi petit que possible.

Le Problème : L'« Explosion de la Mémoire »

L'article explique que faire cette comparaison pour de longues histoires (comme un roman de 100 000 mots) est actuellement un cauchemar pour la mémoire de l'ordinateur.

Voici l'analogie :
Imaginez que vous avez deux tableaux blancs géants, un pour l'attention de l'enseignant et un pour l'attention de l'étudiant. Pour les comparer, l'ancienne méthode nécessite d'écrire chaque paire possible de mots sur ces tableaux.

Si vous avez 64 000 mots, vous devez écrire 64 000 × 64 000 paires. Cela représente plus de 4 milliards de nombres.
Faire cela nécessite un tableau blanc si vaste qu'il ne tient pas dans la mémoire principale de l'ordinateur (HBM). C'est comme essayer de stocker une bibliothèque de livres dans une boîte à chaussures.
Parce que l'ordinateur ne peut pas contenir toute l'image à la fois, il doit découper l'histoire en petits morceaux, les traiter, puis les rassembler. C'est lent, comme essayer de lire un livre en regardant une lettre à la fois et en l'écrivant avant de passer à la suivante.

La Solution : StreamKL (L'approche par « Flux »)

Les auteurs ont créé un nouvel outil appelé StreamKL. Au lieu d'écrire tout sur un tableau blanc géant d'abord, StreamKL utilise une astuce ingénieuse pour calculer le « score de distance » à la volée, comme un tapis roulant.

L'Analogie Créative : La Chaîne de Montage d'une Usine
Imaginez une usine où vous comparez deux tapis roulants de produits (l'attention de l'enseignant et l'attention de l'étudiant).

L'Ancienne Méthode : Vous arrêtez la ligne, vous déversez chaque produit sur un immense entrepôt (HBM), vous les mesurez tous, puis vous nettoyez. Cela occupe tout l'entrepôt et c'est lent.
La Méthode StreamKL : Vous laissez les produits circuler sur le tapis roulant. À mesure que chaque paire d'articles passe devant un capteur (la puce GPU), vous les comparez instantanément, calculez la différence et jetez le résultat dans une petite poche (SRAM) avant l'arrivée de la paire suivante. Vous ne stoppez jamais la ligne, et vous n'avez jamais besoin d'un entrepôt. Vous n'avez besoin que d'une poche.

Comment cela fonctionne (Le Tour de Magie)

L'article décrit deux parties principales de ce tour de magie :

La Passe Avant (Calcul du Score) : Les chercheurs ont inventé une nouvelle formule mathématique qui permet à l'ordinateur de mettre à jour le « score de distance » de manière incrémentale. À mesure qu'il parcourt les données, il garde un décompte courant de seulement quelques nombres (comme un maximum courant et une somme) au lieu de la liste entière. Cela signifie qu'il peut gérer des histoires de n'importe quelle longueur sans manquer de mémoire.
La Passe Arrière (Apprendre des Erreurs) : Lorsque l'ordinateur doit apprendre de l'indice pour améliorer l'étudiant, il doit généralement revenir en arrière pour consulter les données. L'ancienne méthode sauvegarde toute la liste géante de données pour y revenir. StreamKL est plus intelligent : il jette la liste mais se souvient de quelques « clés secrètes » (appelées valeurs LSE). Lorsqu'il doit revenir en arrière, il utilise ces clés pour reconstruire la partie spécifique de la donnée dont il a besoin, sur le moment même, calcule la leçon, puis l'oublie à nouveau. C'est comme se souvenir de la recette d'un gâteau pour pouvoir cuire une part dès que l'on veut la goûter, plutôt que de cuire tout le gâteau et de le stocker dans le réfrigérateur.

Les Résultats : Vitesse et Espace

L'article a testé cela sur des GPU NVIDIA puissants (H200 et A100) avec des contextes très longs (jusqu'à 512 000 mots).

Économies de Mémoire : StreamKL a réduit la mémoire supplémentaire nécessaire d'un mode « quadratique » (explosion vers les téraoctets) à un mode « constant » (restant minuscule). Il est passé de 512 Go de mémoire nécessaires pour un contexte de 64k à presque rien de plus. Cela permet à un seul GPU de gérer des tâches qui nécessitaient auparavant un supercalculateur ou qui étaient impossibles.
Vitesse : Parce qu'il n'a pas besoin d'écrire et de relire de massives quantités de données, il est incroyablement rapide.
- Dans certains tests, il était 43 fois plus rapide que la méthode standard pour calculer le score.
- Dans la phase d'apprentissage, il était 14 fois plus rapide.

Résumé

StreamKL est une nouvelle façon d'enseigner aux modèles d'IA comment prêter attention. Il résout le problème de « l'épuisement de la mémoire » lors du traitement de textes longs en empêchant l'ordinateur d'écrire toute la liste de comparaison. Au lieu de cela, il fait circuler les données à travers un pipeline étroit et efficace, calculant le résultat instantanément. Cela rend possible l'entraînement et l'exécution de modèles d'IA sur des ordinateurs uniques qui étaient auparavant trop volumineux pour être gérés.

Résumé Technique : StreamKL

Énoncé du Problème

La distillation d'attention, une technique utilisée pour entraîner une distribution d'attention à correspondre à une autre en minimisant la divergence de Kullback-Leibler (KL), est cruciale pour la distillation de connaissances, l'élagage de modèles (pruning), l'apprentissage continu et l'entraînement de modèles de langage à large échelle (LLM) à attention parcimonieuse. Cependant, les implémentations existantes font face à de graves goulots d'étranglement de scalabilité lors de la manipulation de longueurs de contexte ( $N_Q, N_K$ ).

Les approches actuelles « matérialisent » les deux distributions d'attention ( $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ ) dans la mémoire à haute bande passante (HBM) avant de calculer la réduction KL élément par élément. Cela entraîne :

Une empreinte mémoire quadratique : Le stockage de deux matrices $N_Q \times N_K$ nécessite une mémoire de $O(N_Q N_K)$ . Pour un contexte de 64K jetons avec 32 têtes en BF16, cela consomme à lui seul environ 512 Go, dépassant la capacité de 141 Go d'un seul GPU NVIDIA H200.
Des coûts d'E/S quadratiques : L'écriture de ces distributions dans l'HBM et leur lecture pour la réduction KL dominent le temps d'exécution, progressant de manière quadratique avec la longueur du contexte.
L'inefficacité du découpage (chunking) : Les solutions de contournement courantes comme le traitement par blocs échangent la latence contre la mémoire ; les petits blocs sous-utilisent le parallélisme du GPU, tandis que les grands blocs déclenchent des erreurs de type Out-of-Memory (OOM).

Le calcul d'attention standard a affronté des problèmes similaires, que FlashAttention a résolus via des noyaux fusionnés et en ligne (online). Cependant, appliquer cela à la distillation d'attention est non trivial car la divergence KL couple deux distributions indépendantes via un terme de différence de logit pondéré, rendant une formulation en un seul passage (one-pass) mathématiquement complexe.

Méthodologie : StreamKL

StreamKL introduit le premier primitif GPU fusionné en un seul passage pour la divergence KL d'attention, éliminant la matérialisation $O(N_Q N_K)$ des distributions intermédiaires.

1. Formulation en ligne (Passage Forward)

StreamKL dérive un nouveau schéma de mise à jour en ligne pour la réduction KL couplée de deux distributions. En reformulant la divergence KL par ligne, les auteurs expriment l'objectif en termes de cinq scalaires courants par ligne :

$m_1, m_2$ : Maximums courants pour les deux distributions de logits.
$l_1, l_2$ : Sommes non normalisées d'exponentielles (similaires au softmax en ligne).
$acc$ : Un accumulateur pour la différence de logit pondérée.

L'objectif reformulé est :
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
où $\text{LSE}_t = m_t + \log l_t$ .

Cette formulation permet au noyau de faire circuler des tuiles de requêtes-clés (query-key tiles) à travers la mémoire SRAM sur puce, en mettant à jour ces cinq scalaires de manière incrémentale. Crucialement, l'accumulateur $acc$ est un vecteur de taille $B_Q$ (taille de bloc), et non une matrice complète, réduisant l'overhead de rescaling arithmétique par rapport au FlashAttention standard. L'algorithme gère le masquage causal en sautant entièrement les tuiles de clés au-delà de la frontière causale.

2. Stratégie de Recomputation (Passage Backward)

Pour éviter de stocker les matrices quadratiques $P_1$ et $P_2$ pendant la rétropropagation, StreamKL emploie une stratégie de recomputation :

État sauvegardé : Seules les valeurs KL finales et les vecteurs Log-Sum-Exp ( $\text{LSE}_1, \text{LSE}_2$ ) sont stockés en HBM ( $O(N_Q)$ ).
Régénération : Pendant le passage backward, les probabilités d'attention sont régénérées tuile par tuile en utilisant les LSE sauvegardés et les matrices $Q, K$ originales ( $P_i = e^{S_i - \text{LSE}}$ ).
Calcul du gradient : Le noyau calcule les gradients pour deux configurations :
- Configuration 1 : $P_1$ fixé, optimisation de $P_2$ (distillation standard).
- Configuration 2 : $P_2$ fixé, optimisation de $P_1$ .
Fusion de noyau : Les auteurs conçoivent un noyau backward fusionné qui visite chaque paire de tuiles $QK$ exactement une fois, réduisant ainsi le trafic HBM. Pour gérer le conflit entre l'accumulation de $dQ$ (selon $K$ ) et de $dK$ (selon $Q$ ), le noyau utilise des opérations atomiques pour $dQ$ lorsque le tampon tient dans le cache L2, et bascule vers des noyaux séparés lorsque la contention devient trop élevée.

3. Optimisations du Noyau GPU

Variante Split-K : Pour les charges de travail avec peu de requêtes (ex: décodage autorégressif où $N_Q=1$ ), une variante Split-K partitionne la dimension des clés entre les blocs de threads pour saturer les processeurs de flux (SM).
Accélération TMA : Sur les GPU Hopper, l'implémentation exploite le Tensor Memory Accelerator (TMA) pour des copies de mémoire groupées asynchrones, masquant la latence mémoire.
Stabilité Numérique : Le calcul du rapport logarithmique dans la Configuration 2 est effectué à l'aide de différences de logits et de LSE sauvegardés plutôt que par des logs de probabilités directs afin d'éviter l'instabilité numérique ( $-\infty$ ou NaN).

Contributions Clés

Identification des Goulots d'Étranglement : Le papier identifie les goulots d'étranglement de mémoire et d'E/S quadratiques qui empêchent la divergence KL d'attention de passer à l'échelle pour de longs contextes.
Reformulation en Ligne : Il introduit la première reformulation en ligne de la divergence KL d'attention, permettant un primitif forward entièrement fusionné et un passage backward basé sur la recomputation qui élimine les intermédiaires quadratiques.
Implémentation Efficace : Les auteurs conçoivent et implémentent des noyaux GPU efficaces avec des optimisations dédiées (Split-K, TMA, fusion adaptative) pour les passages forward et backward.
Scalabilité : StreamKL réduit l'empreinte HBM supplémentaire de la distillation d'attention de $O(N_Q N_K)$ à $O(1)$ , permettant la distillation de contextes longs sur un seul GPU.

Résultats Expérimentaux

Évalué sur des GPU NVIDIA A100 et H200 pour des longueurs de contexte allant de 4K à 512K :

Efficacité Mémoire : StreamKL maintient une empreinte mémoire plate, quel que soit la longueur du contexte. En revanche, les bases de référence (PyTorch, torch.compile, FLA) croissent de manière quadratique et atteignent des erreurs OOM entre 16K et 64K. À 64K, StreamKL utilise plus de 16 000 fois moins d'HBM que torch.compile/FLA et plus de 135 000 fois moins que PyTorch.
Latence Forward :
- StreamKL atteint un accélération de 13 à 18× par rapport à PyTorch et de 2,4 à 3,7× par rapport à torch.compile/FLA dans les contextes non-causaux.
- Sous masquage causal, les accélérations s'élargissent à 26–43× par rapport à PyTorch et 4,2–7,1× par rapport aux bases de référence, car StreamKL saute les calculs masqués tandis que les bases de référence matérialisent toujours les matrices complètes.
Latence Backward :
- Dans la Configuration 1, StreamKL est 2,6 à 4,0× plus rapide que PyTorch (non-causal) et 5,3 à 9,9× plus rapide (causal).
- Dans la Configuration 2, StreamKL atteint une accélération de 3,8 à 6,5× par rapport à PyTorch (non-causal) et de 9,1 à 14,0× (causal).
- Comparé à torch.compile, StreamKL est compétitif ou plus rapide, particulièrement dans les contextes causaux où le coût de recomputation est divisé par deux en sautant les tuiles masquées.
Ablation : La variante Split-K fournit jusqu'à 8,4× d'accélération pour les charges de travail à faible $N_Q$ . Le noyau backward fusionné est jusqu'à 2,0× plus rapide que les noyaux séparés lorsque $N_Q$ est petit, justifiant la stratégie de sélection adaptative.

Signification

Le papier affirme que StreamKL est la première solution capable de soutenir des charges de travail de distillation d'attention avec des longueurs de contexte de 64K et au-delà sur un seul GPU. En éliminant la matérialisation quadratique des distributions d'attention, StreamKL supprime le principal goulot d'étranglement système empêchant la mise à l'échelle de la distillation d'attention dans l'entraînement des LLM modernes (ex: pour les modèles à attention parcimonieuse comme DeepSeek V3.2 et GLM-5). Il démontre que les primitives fusionnées et en ligne peuvent être étendues avec succès au-delà de l'attention standard vers des distributions couplées complexes comme la divergence KL.