Résumé Technique : Pré-remplissage Creux Guidé par Oracle pour les Couches Full/GQA dans les Modèles Hybrides à Long Contexte

1. Énoncé du Problème

Le pré-remplissage (prefill) de long contexte reste un goulot d'étranglement computationnel significatif pour le service des grands modèles de langage (LLM). Bien que les architectures modernes adoptent de plus en plus de conceptions hybrides (combinant des composants locaux, creux, linéaires et récurrents), les couches de pleine attention ou de Grouped-Query Attention (GQA) persistent souvent comme le mécanisme principal pour la récupération et la fusion d'indices à longue portée. Bien que la GQA réduise la mémoire du cache KV en partageant les têtes de clés-valeurs, elle n'élimine pas la recherche dense sur la dimension séquentielle des jetons historiques.

La question diagnostique centrale abordée par ce rapport technique est la suivante : sous une granularité de support et un budget top-k fixes, quelle part de la recherche d'attention dense dans les couches full/GQA existantes est réellement nécessaire pour préserver le comportement au niveau de la tâche ?

Les méthodes actuelles d'attention creuse (ex. : Quest, DuoAttention, MInference, NSA) font face à un défi d'interprétation : une baisse de qualité lors de l'introduction de la parcimonie peut provenir d'un budget insuffisant, d'une mauvaise sélection de l'indexeur, d'un partage de support trop grossier ou de différences de mise en œuvre au niveau du runtime. Sans une référence contrôlée, il est difficile de distinguer la faisabilité théorique d'un support creux des limitations pratiques d'un indexeur ou d'un système d'exécution.

2. Méthodologie

2.1 L'Oracle Top-k de Masse d'Attention

Les auteurs introduisent un oracle top-k de masse d'attention comme un instrument scientifique contrôlé, et non comme une méthode d'accélération déployable. Son but est de séparer la « faisabilité du budget creux » de « l'erreur de l'indexeur » et des « effets de réalisation du runtime ».

• Mécanisme : Pour chaque couche et position de requête, l'oracle calcule la distribution d'attention dense complète. Il agrège ensuite les scores d'attention à travers les têtes de requête pour former une distribution d'importance de jeton moyennée par tête ($P_{\ell,t}$).
• Sélection : Sur la base d'un budget top-k spécifié ($k$) et d'une granularité de support (par requête ou partagée par bloc de sélection), l'oracle sélectionne les indices de jetons top-$k$ de cette distribution.
• Exécution : Le modèle recompute la sortie d'attention en se restreignant uniquement à cet ensemble de support sélectionné, en utilisant les requêtes, clés et valeurs originales. Crucialement, l'oracle utilise l'attention dense uniquement pour sélectionner le support ; la trajectoire suivante du modèle est calculée de manière creuse.
• Granularité : L'étude évalue deux granularités :
  • Par requête ($b_{sel}=1$) : Chaque position de requête sélectionne son propre support.
  • Partagée par bloc de sélection ($b_{sel}=64$) : Un ensemble de support unique est sélectionné pour un bloc de 64 requêtes consécutives afin d'amortir les coûts système.

2.2 Indexeur Effondré par Têtes via Distillation KL

Guidé par l'oracle, les auteurs dérivent un indexeur pour l'inférence pour les noyaux GQA gelés.

• Principe de conception : Le modèle de base reste inchangé. Un indexeur auxiliaire est entraîné pour prédire la distribution d'importance des jetons ($P_{\ell,t}$) définie par l'oracle.
• Objectif d'entraînement : L'indexeur est entraîné via la distillation de divergence KL à partir de la distribution de masse d'attention dense.
• Architecture : Les auteurs proposent une conception effondrée par têtes (head-collapsed). Au lieu de maintenir des canaux de score séparés pour chaque tête de requête ou groupe KV, l'indexeur effondre l'espace de score auxiliaire à travers les têtes, prédisant une distribution d'importance de jeton partagée. Cela réduit les activations de projection et élimine le besoin d'une dimension de tête auxiliaire explicite.
• Inférence : L'indexeur score les clés candidates, sélectionne les $k$ indices top, et la couche GQA creuse recompute l'attention uniquement sur ces jetons sélectionnés. Le budget creux ($k$) reste ajustable lors de l'inférence sans réentraînement.

2.3 Décomposition de l'Évidence

Le papier catégorise explicitement les résultats en trois types distincts d'évidences pour éviter de confondre la faisabilité avec la préparation au déploiement :

1. Écart Oracle (Oracle Gap) : La différence entre l'attention dense et l'attention creuse de l'oracle (utilisant le support sélectionné par l'oracle). Cela teste si la famille de support est théoriquement viable.
2. Écart Indexeur (Indexer Gap) : La dégradation supplémentaire lors du remplacement du support de l'oracle par un indexeur distillé.
3. Écart de Réalisation (Realization Gap) : La dégradation supplémentaire introduite par les contraintes du runtime (ex. : partage de bloc de sélection, fusion de noyaux).

3. Résultats Clés

3.1 Faisabilité de l'Oracle (Écart Oracle)

• Performance Proche de la Densité : À travers les checkpoints de la famille Qwen (Qwen3-8B, Qwen3.5/3.6 27B/35B-A3B), l'oracle par requête reste à moins de 1 point de pourcentage de l'attention dense aux contextes les plus longs évalués (jusqu'à 128K).
• Sensibilité au Budget : Dans un balayage Qwen3.5-9B de 4K à 100K, l'écart de l'oracle est resté à moins de 0,48 point de la densité sous un calendrier top-k configuré (256–2048), avec une parcimonie allant de 87,9 % à 96,0 %.
• Impact de la Granularité : Bien que les oracles par requête soient robustes, les oracles partagés par bloc de sélection ($b_{sel}=64$) montrent un effet de « falaise ». À 128K, les budgets élevés (top-k=2048) maintiennent une qualité proche de la densité, mais les budgets faibles (top-k=256) provoquent un effondrement significatif sur les tâches de type recherche (retrieval). Cela indique que le partage de support nécessite une mise à l'échelle du budget sensible à la longueur.

3.2 Qualité de l'Indexeur Distillé (Écart Indexeur)

• Préservation de la Qualité : En utilisant des indexeurs distillés séparément pour Qwen3.5-0.8B et Qwen3.5-9B, les auteurs rapportent aucune dégradation globale matérielle sur les tranches de validation tenues à l'écart (contextes 16K/32K) sur RULER, VideoMME, MMLongBench-Doc et LongBench-v2.
• Écarts Macro : Les écarts macro rapportés (Distillé moins Dense) étaient de +2,04 pour Qwen3.5-0.8B et +1,13 pour Qwen3.5-9B. Les auteurs traitent ces écarts positifs comme une préservation de la qualité dans la résolution d'évaluation, et non comme une amélioration intrinsèque par rapport à la densité.
• Par Requête vs Fusionné : Dans une validation pure RULER-32K, le chemin distillé par requête était à moins de 0,6 point de la densité, tandis que le chemin fusionné (partagé par bloc de sélection) présentait un écart de 5,0 points, soulignant le coût de l'amortissement compatible avec les systèmes.

3.3 Réserve de Temps de Calcul (TTFT)

• Accélérations de l'Indexeur Distillé : Les mesures sur carte unique montrent des accélérations du TTFT pour le pré-remplissage creux de l'indexeur distillé :
  • Qwen3.5-0.8B (NPU) : Jusqu'à 1,71× d'accélération à 256K de contexte.
  • Qwen3.5-9B (GPU) : Jusqu'à 1,93× d'accélération à 256K de contexte.
  • Les accélérations augmentent avec la longueur du contexte, car le coût relatif des couches de non-attention et des surcharges diminue.
• Tests de Stress de l'Indexeur Non-Entraîné : L'exécution du chemin creux avec des poids d'indexeur aléatoirement initialisés (non entraînés) a révélé un potentiel de réserve encore plus élevé (jusqu'à 3,44× sur NPU à 128K), indiquant que le chemin d'exécution creux lui-même est efficace, bien que la qualité de la sortie ne soit pas validée dans cette configuration.

4. Signification et Revendications

Le papier se positionne comme un rapport technique de diagnostic plutôt que comme la revendication d'un système creux entièrement déployé et sans perte. Ses principales contributions sont :

1. Recadrage du Problème : Il cadre le pré-remplissage creux pour les couches GQA comme une question de réductibilité guidée par l'oracle. En établissant ce qui est théoriquement possible avec la masse d'attention dense, il isole les sources d'erreur dans les systèmes creux pratiques.
2. Définition de l'Oracle : L'oracle top-k de masse d'attention sert de « référence contrôlée » pour séparer la faisabilité du budget de la performance de l'indexeur et de la réalisation du runtime.
3. Évidence Empirique : Il fournit la première preuve que, pour les checkpoints GQA de la famille Qwen, l'attention dense identifie souvent un support compact et dépendant de la requête qui préserve le comportement au niveau de la tâche, même à des longueurs extrêmes (128K–256K).
4. Séparation Méthodologique : Le rapport sépare explicitement trois affirmations distinctes :
   • Faisabilité de l'Oracle : Le support creux est viable (prouvé par les lignes de l'oracle).
   • Qualité de l'Indexeur Distillé : Les indexeurs à squelette gelé peuvent approximer ce support avec une perte de qualité minimale (prouvé par les lignes distillées).
   • Réserve de Temps de Calcul : Le chemin de service creux offre des réductions significatives de TTFT (prouvé par les lignes de stress et les lignes de TTFF distillées).

Limites et Portée :
Les auteurs sont modestes quant à leurs revendications, notant qu'il s'agit d'une « première version ». Ils précisent explicitement que :

• Il n'existe pas de frontière unique « densité–oracle–qualité-latence parfaitement correspondante » où toutes les métriques sont mesurées sur la même configuration exacte.
• Les indexeurs distillés ont été validés principalement à 16K/32K, tandis que les plus grandes accélérations de TTFT ont été mesurées à 128K/256K sans vérifications de qualité appariées.
• La conception de l'indexeur effondré par têtes est présentée comme un choix orienté système, et non comme un optimum prouvé par ablation contre des variantes à têtes étendues.
• L'évidence est actuellement concentrée sur la famille Qwen ; la validation croisée avec d'autres familles est laissée aux travaux futurs.

En résumé, le papier démontre que l'attention dense est souvent un « surplus » (overkill) pour la recherche de long contexte dans les couches GQA et qu'un indexeur effondré par têtes et distillé peut efficacement approximer le support nécessaire, préservant la qualité tout en débloquant des accélérations significatives du TTFT, à condition que la granularité du support et le budget soient gérés avec soin.