Résumé Technique : ParisKV

Énoncé du Problème

Les grands modèles de langage (LLM) étendent rapidement leurs fenêtres de contexte à des millions de jetons, créant une pression sans précédent sur l'efficacité de l'inférence. Le cache KV, qui stocke les clés et les valeurs pour le décodage autorégressif, domine rapidement l'empreinte mémoire et la latence à mesure que le contexte croît. Le décodage à long contexte est intrinsèquement limité par la mémoire, car chaque étape nécessite la lecture d'un grand volume de vecteurs KV, avec des coûts de bande passante évoluant linéairement avec la longueur du contexte.

Bien que les méthodes d'attention éparse ou sélective tentent d'atténuer cela en ne portant attention qu'à un sous-ensemble Top-k de jetons passés pertinents, les approches actuelles de récupération du cache KV font face à trois défis critiques dans les scénarios de longue génération sensibles à la latence :

1. Compromis Vitesse–Qualité : La récupération approximative légère (ex: regroupement grossier, quantification à faible nombre de bits) sacrifie souvent le rappel au profit de la vitesse. Récupérer la précision nécessite généralement d'augmenter le budget de récupération, ce qui érode les bénéfices de la parcimonie.
2. Dérive de Décodage (Decoding Drift) : Les centroïdes appris par regroupement (clustering) des clés de préremplissage (prefill) deviennent obsolètes au fur et à mesure que la génération progresse. À mesure que les clés de décodage s'accumulent, la distribution change, provoquant un décalage entre les centroïdes du préremplissage et la véritable distribution des clés. Cela conduit à un effondrement sévère du rappel lors de longues étapes de décodage dans les méthodes antérieures.
3. Récupération Limitée par le CPU : Lorsque les caches KV complets sont déchargés vers la mémoire CPU pour économiser l'espace GPU, la récupération implique souvent une recherche côté CPU et un transfert de données CPU-vers-GPU. Cela introduit des erreurs d'approximation et fait de l'orchestration par le CPU le facteur dominant de la latence de bout en bout.

Méthodologie : Le Framework ParisKV

ParisKV est un algorithme et un système co-conçus pour le matériel (hardware-aware) dédiés à la récupération du cache KV, permettant un décodage rapide et un rappel stable sous des entrées et/ou des générations longues (jusqu'à l'échelle du million de jetons), même en présence de dérive de distribution.

Innovations Algorithmiques Fondamentales

1. Centroïdes Analytiques Robustes à la Dérive :
Contrairement aux méthodes précédentes qui apprennent des centroïdes à partir des clés de préremplissage, ParisKV transforme les requêtes et les clés dans un espace stable et indépendant des données.

• Normalisation et Rotation : Les clés et les requêtes sont normalisées $\ell_2$ et soumises à une rotation orthogonale aléatoire partagée (implémentée via SRHT). Cela les projette sur une hypersphère unité.
• Centroïdes Analytiques : Au lieu d'apprendre des centroïdes, ParisKV définit un ensemble de centroïdes uniformément distribués sur l'hypersphère unité basés sur des motifs de signe (ex: $\pm 1/\sqrt{m}$). Comme ces centroïdes sont indépendants des données et uniformément distribués, toute nouvelle clé générée lors du décodage sera toujours proche d'au moins un centroïde, empêchant le problème du « centroïde obsolète ».

2. Pipeline de Récupération Natif GPU à Deux Étapes :
ParisKV implémente un pipeline de récupération du grossier au fin entièrement sur le GPU pour éviter les goulots d'étranglement du CPU.

• Génération de Candidats Grossière (Basée sur les Collisions) : Les clés rotées sont divisées en $B$ sous-espaces. Chaque direction de sous-espace est assignée au centroïde analytique le plus proche. Pour une requête, le système sélectionne les top-$\rho$ centroïdes les plus similaires dans chaque sous-espace. Les clés dans ces compartiments sélectionnés reçoivent des votes (avec des poids hiérarchisés). Le système somme les votes à travers les sous-espaces pour générer un score entier grossier, réduisant l'espace de recherche à un petit ensemble de candidats (ex: top 5–10 %).
• Reclassement des Candidats (RSQ-IP) : Les candidats sont reclassés à l'aide d'un estimateur de produit scalaire calibré sans accéder aux clés en pleine précision.
  • Quantification : Chaque direction de sous-espace est encodée à l'aide d'un code à 4 bits (1 bit de signe + 3 bits de magnitude) dérivé des priors Beta induits par la rotation.
  • Calibration : Un facteur d'échelle léger ($w_{i,b}$) est précalculé durant la phase de préremplissage pour corriger les erreurs d'alignement de quantification.
  • Estimation : Le système estime le score d'attention brut $\langle k, q \rangle$ en utilisant la formule : $\|q\|_2 \sum_{b=1}^B w_{i,b} \langle v_{i,b}, \tilde{q}_b \rangle$, où $v_{i,b}$ est la direction déquantifiée et $\tilde{q}_b$ est la requête rotée. Cela permet un reclassement de haute fidélité en utilisant uniquement des métadonnées compactes.

Innovations au Niveau Système

1. Déchargement par Adressage Virtuel Unifié (UVA) :
ParisKV décharge le cache KV en pleine précision vers la mémoire CPU mais exploite l'UVA pour permettre aux noyaux GPU de récupérer directement uniquement les paires Top-k de KV sélectionnées sur demande. Cela contourne les opérations memcpy explicites et la surcharge de planification côté CPU, permettant au GPU d'accéder à la mémoire CPU comme s'il s'agissait de mémoire locale.

2. Mécanisme de Mise à Jour par Fenêtre Glissante :
Pour supporter une génération illimitée, le cache KV est géré en quatre régions : Sink (jetons précoces sur GPU), Local (jetons récents sur GPU), Retrieval (jetons historiques déchargés) et un Buffer de Mise à Jour. Lorsque le buffer se remplit, les jetons locaux les plus anciens sont évincés vers la zone Retrieval (déchargés sur CPU), et les jetons du buffer sont promus vers la zone Local. Cela garantit que les métadonnées de récupération restent fraîches tout en préservant un accès rapide au contexte récent et à l'historique à long terme.

Contributions Clés

1. Centroïdes Analytiques Robustes à la Dérive : Une transformation indépendante des données (normalisation + rotation) qui maintient un rappel de récupération stable même sous des changements de distribution importants à l'échelle du million de jetons, éliminant le besoin de centroïdes appris qui se dégradent avec le temps.
2. Pipeline Natif GPU de Grossier à Fin : Un système de récupération combinant l'élagage par collision et le reclassement calibré (RSQ-IP) qui permet une sélection Top-k de haute fidélité directement sur des codes compressés sans recherche côté CPU.
3. Déchargement Scalable Basé sur l'UVA : Une implémentation permettant la récupération de KV sur demande avec une intervention minimale du CPU, permettant au système de passer à des contextes de millions de jetons là où l'attention complète manque de mémoire.
4. Évaluation Étendue : Performance démontrée sur des tâches de raisonnement à génération longue (GPQA-Diamond, MATH500, AIME25) et de compréhension de contexte long (LongBench-V2, RULER).

Résultats Expérimentaux

Les auteurs ont évalué ParisKV sur des benchmarks de raisonnement à génération longue (GPQA-Diamond, MATH500, AIME25) et de compréhension de contexte long (LongBench-V2, RULER) à travers plusieurs familles de modèles (Qwen-3-8B, DeepSeek-R1-Llama-8B, etc.).

Précision :

• ParisKV égale ou surpasse la qualité de l'attention complète sur les benchmarks de long-input et de longue génération.
• Sur AIME25, il améliore la précision pass@8 de +40 à +77 points par rapport aux bases de référence de pointe MagicPIG et PQCache.
• Dans 7 cas sur 9, ParisKV égale ou dépasse la précision de l'attention complète, surpassant de manière significative les bases de référence de récupération antérieures.

Efficacité :

• Débit (Throughput) : ParisKV atteint jusqu'à 2,8× plus de débit que l'attention complète dans sa plage d'exécution.
• Latence à l'Échelle : À l'échelle du million de jetons (1M tokens), ParisKV réduit la latence de décodage de 17× par rapport à MagicPIG et de 44× par rapport à PQCache.
• Scalabilité : Alors que l'attention complète tombe en erreur de mémoire (OOM) à des tailles de batch $\ge 4$ pour 128K contextes et $\ge 2$ pour 256K contextes, ParisKV passe à des tailles de batch de 8 et 5, respectivement. Il reste exécutable à 384K et 1024K contextes là où l'attention complète échoue complètement.

Signification et Revendications

L'article affirme que ParisKV répond aux défis fondamentaux de la dérive de distribution et de la latence de récupération dans l'inférence LLM à long contexte. En transformant les clés sur une hypersphère unité et en employant des centroïdes analytiques indépendants des données, le système élimine l'obsolescence des centroïdes qui affecte les méthodes existantes.

Les auteurs soutiennent qu'une récupération de cache KV soigneusement conçue peut simultanément améliorer l'efficacité et la précision pour l'inférence à long contexte, même en présence de dérive. Cela ouvre de nouvelles possibilités pour le déploiement de LLM dans des applications de génération longue sensibles à la latence (telles que des tâches de raisonnement complexe) où les méthodes précédentes échouaient soit à passer à l'échelle, soit subissaient une dégradation sévère de la précision. Le travail démontre qu'il est possible d'atteindre une qualité quasi sans perte tout en conservant les avantages de vitesse de l'attention éparse, même lorsque le cache KV est déchargé vers la mémoire CPU.