Résumé Technique : RDKV - Allocation de Bits pour la Distorsion-Taux dans l'Éviction et la Quantification Jointes du Cache KV

1. Énoncé du Problème

Les Modèles de Langage de Grande Taille (LLM) font face à un goulot d'étranglement critique lors de l'inférence en contexte long : la bande passante et la capacité mémoire. La taille du cache Key-Value (KV) croît linéairement avec la longueur de la séquence, nécessitant la rélecture de dizaines de gigaoctets de données depuis la mémoire haute bande passante (HBM) hors puce vers la mémoire sur puce à chaque étape de décodage. Cela rend l'inférence limitée par la mémoire plutôt que par le calcul.

Les stratégies de compression existantes relèvent généralement de deux catégories isolées :

1. Éviction : Notation des tokens ou des canaux et rejet complet des moins bien classés (binaire : garder ou jeter). Cela ne laisse aucune échelle de précision pour les tokens d'importance modérée.
2. Quantification : Conservation de tous les tokens mais réduction de leur largeur de bits. Cela conserve même les tokens négligeables, gaspillant ainsi la mémoire.

Bien que des hybrides récents tentent de combiner ces approches, ils reposent souvent sur des scores de routage heuristiques ou des ratios pré-alloués déterminés par recherche hors ligne, échouant à traiter l'éviction et la quantification comme un problème d'optimisation unifié.

2. Méthodologie : RDKV

L'article propose RDKV (compression de cache KV par distorsion-taux), un cadre qui formule la compression du cache KV comme un problème de distorsion-taux. Dans cette perspective, l'éviction et la quantification ne sont pas des actions séparées mais deux extrémités d'une seule courbe d'allocation de bits.

2.1 Formulation Théorique

RDKV déduit un poids de distorsion pour chaque unité de cache (token dans le cache V, canal dans le cache K) basé sur la perturbation induite par son retrait ou sa quantification sur le calcul de l'attention.

• Cache V (Tokens) : Le poids de distorsion $w_t$ pour un token est la distance de variation totale de la distribution d'attention après éviction. Il est calculé comme la somme des scores d'attention reçus par le token à travers toutes les positions de requête : $w_t = \sum_{\tau} a_{\tau,t}$.
• Cache K (Canaux) : Le poids de distorsion $w_c$ pour un canal est la norme spectrale de la déviation des logits causée par la mise à zéro de ce canal. Il est calculé comme le produit des normes des colonnes de requête et de clé correspondantes : $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Ces poids servent de coefficients de coût dans une fonction objectif de distorsion-taux :
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{sous contrainte} \quad \sum_{u} b_u \leq B$$
où $b_u$ est la largeur de bits assignée à l'unité $u$, et $B$ est le budget total de bits.

2.2 Pipeline d'Allocation de Bits

L'allocation est résolue par remplissage d'eau inverse (relaxation de Lagrange) adaptée aux largeurs de bits discrètes supportées par le matériel $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Étape 1 (Pondération) : Calcul des poids de tokens ($w_t$) et des poids de canaux ($w_c$) une seule fois après l'étape de préremplissage en utilisant le cache complet non compressé.
2. Étape 2 (Allocation de Tokens) : Résolution d'un Problème du Sac à Dos à Choix Multiples (MCKP) pour le cache V. Les tokens sont assignés des largeurs de bits depuis $\{0, 2, 4, 8, 16\}$. Une largeur de bits de 0 correspond à l'éviction.
3. Étape 3 (Allocation de Canaux) : Résolution du MCKP pour le cache K, mais uniquement pour les tokens conservés à l'Étape 2. Les canaux sont assignés des largeurs de bits indépendamment de la largeur de bits V du token, permettant à un seul token d'avoir des composantes K et V de précision mixte.

2.3 Décompression Efficace Emboîtée (TriZone)

Pour réaliser des économies réelles de mémoire et de latence, RDKV introduit une disposition de cache TriZone qui fusionne la déquantification dans le noyau d'attention, évitant ainsi le besoin de dépaqueter les données en FP16 avant le traitement.

• Zone A : Flux de cache ancien empaqueté contenant les lignes K et V conservées avec des largeurs de bits $\{2, 4, 8\}$. Celles-ci sont regroupées par largeur de bits pour permettre des chemins de déquantification uniformes.
• Zone B : Lignes V en FP16 pour les tokens assignés à une précision de 16 bits.
• Zone C : Lignes K et V en FP16 pour les tokens nouvellement générés.

Le noyau d'attention calcule les logits et les sorties en combinant ces trois zones, la déquantification étant fusionnée directement dans la multiplication matricielle, éliminant ainsi les surcoûts supplémentaires de bande passante mémoire.

3. Contributions Clés

1. Formulation Distorsion-Taux : L'article présente la compression du cache KV comme un problème de distorsion-taux unifié, prouvant que la solution optimale mélange intrinsèquement des unités de taux zéro (éviction) et de taux fini (quantification). Cela élimine le besoin de ratios heuristiques ou de recherches hors ligne.
2. Allocation Discrète par Remplissage d'Eau Inverse : Les auteurs instancient l'allocation comme un problème de sac à dos discret soluble par relaxation de Lagrange, assignant des largeurs de bits optimales depuis $\{0, 2, 4, 8, 16\}$ basées sur des poids de distorsion dérivés de la dynamique de l'attention.
3. Disposition Empaquetée TriZone : Une implémentation consciente du matériel qui fusionne la déquantification dans le noyau d'attention, convertissant les économies théoriques de largeur de bits en réductions réelles de la mémoire de pointe et de la latence de décodage.

4. Résultats Expérimentaux

Les expériences ont été menées sur LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B et Qwen2.5-72B à travers LongBench, RULER et InfiniteBench.

• Précision : RDKV surpasse la meilleure ligne de base évaluée (Ada-SnapKV) de 9,1 % en moyenne. Sur LongBench, il récupère 97,81 % de la précision du cache complet avec seulement 2,48 % de rétention de cache.
• Performance en Contexte Long :
  • RULER : À un contexte de 128K, RDKV atteint 66,95 % de précision contre 58,43 % pour la ligne de base la plus forte (AdaKV), soit un écart de 8,5 points.
  • InfiniteBench : RDKV se classe premier parmi les méthodes de compression avec un score moyen de 39,46, surpassant AdaKV de 1,40 point.
  • Aiguille dans une Botte de Foin : RDKV maintient une précision de récupération quasi uniforme (moyenne de 0,99) à travers les profondeurs, tandis que les lignes de base basées uniquement sur l'éviction échouent significativement aux profondeurs intermédiaires.
• Efficacité :
  • Latence : Avec un contexte de 128K, RDKV réalise une accélération de décodage de 4,5× par rapport au décodage FlashAttention-2 avec cache complet.
  • Mémoire : Il réalise une réduction de 1,9× de la mémoire de pointe à un contexte de 128K et exécute avec succès un contexte de 256K sur un seul GPU A100 64 Go, là où FullKV échoue par manque de mémoire (OOM).

5. Importance et Revendications

L'article revendique que traiter l'éviction et la quantification conjointement au sein d'un cadre d'optimisation unique est supérieur aux approches par étapes ou basées sur des ratios. En dérivant les scores d'allocation directement de la distorsion induite sur le calcul de l'attention, RDKV évite le « piège binaire » des méthodes d'éviction exclusive (qui rejettent les tokens modérés) et le « gaspillage » des méthodes de quantification exclusive (qui conservent les tokens négligeables).

Les auteurs soulignent que leur méthode :

• Ne nécessite aucun réentraînement du modèle.
• Est agnostique de l'architecture, performant de manière cohérente à travers différentes tailles et familles de modèles.
• Fournit un compromis Pareto-optimal entre latence et précision, surpassant strictement les lignes de base sur les deux métriques aux longueurs de contexte élevées.
• Démontre que l'« optimum continu » de la théorie de la distorsion-taux, lorsqu'il est discrétisé pour le matériel, produit naturellement une stratégie de bits mixtes qui surpasse les hybrides actuels de l'état de l'art.

L'article conclut que cette approche conjointe améliore la précision sans sacrifier l'efficacité, permettant l'inférence en contexte long sur du matériel aux ressources contraintes.