PagedAttention

PagedAttention est un algorithme d'attention et un gestionnaire de mémoire destiné au service de grands modèles de langage (LLM). Il stocke le cache KV de chaque requête dans de petits blocs de taille fixe, placés n'importe où dans la mémoire du GPU, et les suit au moyen d'une table de blocs propre à chaque requête, comme un système d'exploitation gère les pages de mémoire virtuelle. Woosuk Kwon, Zhuohan Li et leurs collègues de l'UC Berkeley et d'universités partenaires l'ont présenté avec le moteur de service vLLM dans Efficient Memory Management for Large Language Model Serving with PagedAttention, présenté au SOSP 2023. La suppression de la fragmentation due aux allocations contiguës du cache permet à un GPU de conserver simultanément beaucoup plus de séquences ; l'article rapporte un débit 2 à 4 fois supérieur à celui de FasterTransformer et Orca pour la même latence.

Le problème de mémoire résolu#

Pendant la génération de texte, un transformeur conserve les vecteurs de clés et de valeurs de chaque jeton antérieur afin de ne jamais avoir à les recalculer. Ce cache augmente d'une entrée par jeton et par couche ; une seule séquence du modèle LLaMA à 13 milliards de paramètres peut occuper jusqu'à 1,7 Go. La longueur de sortie étant inconnue à l'arrivée d'une requête, les premiers systèmes de service réservaient, pour chaque requête, une région contiguë dimensionnée pour la longueur maximale possible. La majeure partie de cette région restait vide, et les espaces entre les régions étaient trop petits pour être réutilisés. L'équipe vLLM a mesuré que ces systèmes gaspillaient 60 à 80 % de la mémoire du cache KV, ce qui limitait la taille de lot bien avant l'épuisement des ressources de calcul.

Fonctionnement de la pagination#

PagedAttention divise le cache en blocs contenant chacun les clés et les valeurs d'un nombre fixe de jetons, 16 par défaut dans vLLM. Une requête ne reçoit un nouveau bloc que lorsque le précédent est plein ; elle l'obtient à l'emplacement disponible et sa table de blocs enregistre l'ordre. Le noyau du mécanisme d'attention suit cette table pour rassembler les clés et les valeurs bloc par bloc ; le cache d'une séquence n'a donc jamais besoin d'être contigu. Le gaspillage se limite à la partie inutilisée du dernier bloc de chaque séquence, soit moins de 4 % selon les mesures de vLLM.

Cet exemple PyTorch stocke les clés d'une requête dans des blocs dispersés d'un pool partagé, puis les relit à l'aide de sa table de blocs :

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Le programme affiche la table de blocs [41, 7, 19] et keys intact: True, puis 8 emplacements inutilisés pour la requête paginée, contre 2 008 pour une région réservée à 2 048 jetons.

Partage des blocs entre les séquences#

Les tables de blocs permettent également à plusieurs séquences de pointer vers les mêmes blocs physiques. L'échantillonnage parallèle, la recherche par faisceau et les requêtes qui commencent par la même invite système réutilisent une seule copie du préfixe partagé, avec un compteur de références pour chaque bloc. Lorsqu'une séquence doit écrire dans un bloc encore utilisé par d'autres, le moteur le copie d'abord : cette copie sur écriture est une technique empruntée aux systèmes d'exploitation.

Les auteurs de vLLM rapportent que ce partage réduit jusqu'à 55 % la mémoire utilisée par l'échantillonnage parallèle et la recherche par faisceau, ce qui permet d'obtenir jusqu'à 2,2 fois plus de débit. Le hachage des blocs complets étend cette approche aux requêtes ; c'est ainsi que vLLM implémente la mise en cache des invites.

Techniques connexes#

PagedAttention détermine où se trouvent les clés et les valeurs mises en cache ; cette méthode se combine avec des techniques qui réduisent la quantité à stocker ou accélèrent la lecture :

• Attention flash : découpe le calcul de l'attention en tuiles afin que la matrice complète des scores ne soit jamais écrite dans la mémoire principale du GPU. Les moteurs actuels exécutent des noyaux de type FlashAttention qui lisent les blocs paginés.
• Traitement par lots continu : admet une nouvelle requête dès qu'une autre se termine ; la pagination remet aussitôt les blocs de la séquence terminée dans le pool libre, libérant de la place pour la nouvelle requête.
• Attention à requêtes groupées : partage les têtes de clés et de valeurs entre les têtes de requêtes, ce qui réduit le contenu de chaque bloc.
• Quantification du cache FP8 : stocke les clés et les valeurs sur 8 bits, réduisant approximativement de moitié la mémoire du cache par rapport aux formats 16 bits.

L'interface LLM d'Ultralytics transmet les arguments des requêtes sans les modifier ; un appel Chat Completions avec n=3 envoyé par son base_url à un serveur vLLM local renvoie donc trois échantillons qui partagent les blocs mis en cache de l'invite, et seules leurs continuations divergentes utilisent de la mémoire supplémentaire.