Résumé technique : OSCAR - Rotation offline consciente de la covariance spectrale pour la quantification de cache KV en 2 bits

1. Énoncé du problème

L'inférence à contexte long dans les modèles de langage de grande taille (LLM) est de plus en plus limitée par l'empreinte mémoire du cache Key-Value (KV). À mesure que la longueur du contexte, la taille du lot et la profondeur du modèle augmentent, le cache consomme une mémoire GPU significative, limitant le débit et l'évolutivité. Bien que la compression du cache KV vers une précision ultra-faible (par exemple, INT2) offre une réduction de mémoire de 8×, elle se heurte à un défi critique de précision : les activations KV contiennent des valeurs aberrantes sévères au niveau des canaux.

Sous une quantification à faible nombre de bits, ces valeurs aberrantes dominent l'échelle de quantification, forçant la majorité des valeurs « normales » vers quelques niveaux de quantification effectifs et dégradant sévèrement la qualité de l'attention. Les méthodes existantes basées sur la rotation (par exemple, les transformations de Hadamard) tentent de lisser ces valeurs aberrantes en redistribuant l'énergie à travers les canaux. Cependant, ces méthodes sont aveugles aux données ; elles lissent la distribution brute du cache mais échouent à s'aligner sur les directions spécifiques que les mécanismes d'attention en aval consomment réellement. À INT2, où seuls quatre niveaux de quantification existent, ce désalignement provoque l'effondrement des méthodes de rotation naïves, entraînant une précision proche de zéro sur les tâches de raisonnement.

2. Méthodologie : OSCAR

Les auteurs proposent OSCAR (Offline Spectral Covariance-Aware Rotation), une méthode qui estime les structures de covariance conscientes de l'attention hors ligne pour dériver des rotations fixes et des seuils de clipping. Contrairement aux travaux antérieurs qui minimisent l'erreur de reconstruction euclidienne brute ($\|K - \hat{K}\|_F$), OSCAR minimise la distorsion dans les scores et les sorties d'attention en aval.

2.1 Motivation théorique

L'article soutient que la cible de rotation optimale doit être dérivée des statistiques du calcul d'attention lui-même :

• Pour les Clés ($K$) : La distorsion dans les logits d'attention est contrôlée par la covariance de la requête $Q^\top Q$, et non simplement par la covariance de la clé $K^\top K$.
• Pour les Valeurs ($V$) : La distorsion dans la sortie d'attention est pondérée par les scores d'attention $S$, conduisant à une covariance cible de $V^\top S^\top S V$.

2.2 Conception de l'algorithme

OSCAR fonctionne en deux phases :

A. Phase de calibration offline :

1. Estimation de la covariance : En utilisant un petit ensemble de données de calibration, la méthode calcule des covariances cibles empiriques :
   • $C_Q = \frac{1}{N} \sum q_n^\top q_n$ (Cible de clé consciente de la requête).
   • $C_S = \frac{1}{N} V^\top S^\top S V$ (Cible de valeur consciente du score).
2. Dérivation de la rotation : La méthode effectue une décomposition en valeurs propres sur ces matrices pour trouver les bases propres optimales ($U_Q$ et $U_S$).
3. Construction de la rotation composite : Les matrices de rotation finales ($R_K, R_V$) sont construites comme le produit de trois composantes :
   • $U_Q/U_S$ : La base propre consciente de l'attention qui aligne les canaux avec les directions d'importance.
   • $H_{Had}$ : Une transformation de Hadamard qui égalise la métrique d'importance à travers les canaux et disperse l'énergie des valeurs aberrantes.
   • $P_{br}$ : Une permutation de réversion de bits qui entrelace les canaux à variance élevée et faible pour garantir des plages dynamiques équilibrées au sein des groupes de quantification.
   • Formule : $R_K = U_Q H_{Had} P_{br}$ et $R_V = U_S H_{Had} P_{br}$.
4. Seuils de clipping : Les seuils de clipping par couche sont estimés à partir des activations de calibration rotées pour contrôler les valeurs aberrantes.

B. Phase de service en ligne :

• Disposition mixte de précision : Le système maintient un cache KV paginé où les tokens « puits » (initiaux) et « récents » sont conservés en pleine précision BF16. Le contexte historique intermédiaire est stocké en INT2 après application de la rotation OSCAR fixe et du clipping.
• Intégration du noyau : Les auteurs ont développé un noyau Triton fusionné personnalisé pour le framework de service SGLang. Ce noyau gère le dépaquetage, la déquantification, la rotation inverse et le calcul d'attention pour l'historique INT2, garantissant la compatibilité avec l'attention paginée et la mise en cache des préfixes sans introduire de trafic mémoire supplémentaire.

3. Contributions clés

1. Identification de la cible manquante : L'article identifie que les rotations génériques dispersent les valeurs aberrantes mais échouent à INT2 car elles ne s'alignent pas sur l'attention en aval. La cible de rotation doit être induite par les statistiques d'attention ($Q^\top Q$ et $V^\top S^\top S V$), et non par la reconstruction brute du cache.
2. Cadre OSCAR : Un système de quantification de cache KV en INT2 entièrement déployable qui utilise une calibration offline pour dériver des rotations conscientes de l'attention. L'analyse théorique montre que ces cibles sont optimales sous une hypothèse de substitut d'erreur figée.
3. Système prêt pour la production : Un noyau d'attention INT2 personnalisé intégré à SGLang qui prend en charge les dispositions de cache KV paginées et la mise en cache des préfixes. Il réalise une intégration transparente avec les frameworks de service modernes (SGLang, vLLM) tout en maintenant un débit élevé.

4. Résultats expérimentaux

Les auteurs ont évalué OSCAR sur des modèles de raisonnement (Qwen3-4B, Qwen3-8B, Qwen3-32B) et un modèle de pointe (GLM-4.7, 358 milliards de paramètres) à travers cinq benchmarks (AIME25, GPQA, HumanEval, LiveCodeBench, MATH500) et la récupération à contexte long (RULER-NIAH jusqu'à 128k).

• Préservation de la précision :

  • Sur Qwen3-4B-Thinking et Qwen3-8B, OSCAR réduit l'écart de précision par rapport au BF16 à 3,78 et 1,42 points, respectivement. En revanche, la rotation naïve INT2 s'effondre presque à zéro de précision.
  • Sur Qwen3-32B et GLM-4.7, OSCAR reste effectivement à parité avec le BF16 (dans une marge de 0,02–0,27 points), tandis que d'autres bases de référence INT2 subissent une dégradation significative.
  • Dans la récupération à contexte long (RULER-NIAH), OSCAR maintient des performances robustes jusqu'à 128k tokens, alors que les méthodes INT2 naïves et uniquement Hadamard s'effondrent à des longueurs beaucoup plus courtes (par exemple, 16k).

• Performance système :

  • Mémoire : Réduit l'empreinte mémoire du cache KV d'environ 8× par rapport au BF16.
  • Débit : Améliore le débit jusqu'à 7× pour de grandes tailles de lot sous le même budget mémoire.
  • Latence : Accélère le décodage de taille de lot 1 jusqu'à 3× par rapport au BF16 en raison de la réduction de la surcharge de bande passante mémoire.
  • Évolutivité : Démontre la capacité de gérer jusqu'à 28 requêtes simultanées avec des entrées de 100k tokens sur un seul H100, un régime où les bases de référence BF16 et INT4 manquent de mémoire.

5. Importance et affirmations

L'article affirme qu'OSCAR comble avec succès le fossé entre la promesse théorique du caching KV ultra-bas bit (INT2) et sa déployabilité pratique. En déplaçant la cible de rotation de la reconstruction de données brutes vers la covariance consciente de l'attention, OSCAR récupère la précision perdue par la quantification naïve tout en conservant les avantages massifs de mémoire et de débit du stockage en 2 bits.

Les auteurs soulignent que leur approche n'est pas simplement une amélioration algorithmique mais une solution au niveau système. En fournissant un noyau personnalisé qui s'intègre à l'attention paginée et à la mise en cache des préfixes, ils démontrent que la quantification INT2 peut être utilisée dans les frameworks de service de production (comme SGLang) sans sacrifier les fonctionnalités d'efficacité (comme la réutilisation des préfixes) sur lesquelles repose le service moderne des LLM. Le travail suggère que pour le raisonnement à contexte long, la clé d'une quantification réussie réside dans l'alignement de la stratégie de compression avec les demandes statistiques spécifiques du mécanisme d'attention.