XQuant : franchir le mur de la mémoire pour l'inférence des grands modèles linguistiques grâce à la rématérialisation du cache KV

Bien que l’inférence des grands modèles linguistiques (LLM) se soit imposée comme une charge de calcul essentielle pour de nombreuses applications ultérieures, son exécution efficace reste un défi en raison de la taille importante de la mémoire requise et des besoins élevés en bande passante. Parallèlement, au cours des dernières décennies, les capacités de calcul ont progressé de manière constante, dépassant à la fois la capacité mémoire et la bande passante, une tendance qui demeure manifeste dans les architectures GPU modernes et aggrave encore le problème de l’inférence des LLM. À cet égard, de nouveaux algorithmes émergent, qui échangent une augmentation du calcul contre une réduction du nombre d’opérations mémoire. Dans cette optique, nous proposons XQuant, qui exploite cette tendance pour permettre une réduction d’un ordre de grandeur de la consommation mémoire grâce à une quantification à faible précision, offrant des gains significatifs en précision par rapport aux méthodes d’état de l’art de quantification du cache KV. Nous y parvenons en quantifiant et en mettant en cache les activations d’entrée des couches, notées X, au lieu d’adopter la méthode classique de mise en cache KV, puis en recalculant dynamiquement les clés (Keys) et les valeurs (Values) pendant l’inférence. Cette approche permet une économie immédiate de 2× en mémoire par rapport à la mise en cache KV. Grâce à XQuant, nous obtenons jusqu’à 7,7× d’économie mémoire avec une dégradation de moins de 0,1 en perplexité par rapport à la référence en FP16. En outre, notre méthode exploite le fait que les valeurs de X sont similaires à travers les couches. En s’appuyant sur cette observation, nous introduisons XQuant-CL, qui exploite la similarité entre couches des embeddings X afin d’obtenir une compression extrême. Sur divers modèles, XQuant-CL atteint jusqu’à 10× d’économie mémoire par rapport à la référence FP16 avec une dégradation de seulement 0,01 en perplexité, et jusqu’à 12,5× avec une dégradation de seulement 0,1 en perplexité. XQuant exploite ainsi la croissance rapide des capacités de calcul des plateformes matérielles pour éliminer le goulot d’étranglement mémoire, tout en dépassant les méthodes d’état de l’art de quantification du cache KV et en atteignant une précision quasi-équivalente à celle du FP16 sur une large gamme de modèles.