ChunkAttention : Auto-attention efficace avec cache KV conscient du préfixe et partitionnement en deux phases
ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition
February 23, 2024
Auteurs: Lu Ye, Ze Tao, Yong Huang, Yang Li
cs.AI
Résumé
L'auto-attention est un composant essentiel des grands modèles de langage (LLM), mais constitue une source significative de latence lors de l'inférence pour les longues séquences. Dans les scénarios de service multi-locataires de LLM, le coût en calcul et en opérations mémoire de l'auto-attention peut être optimisé en exploitant la probabilité que plusieurs requêtes de LLM partagent des prompts système dans leurs préfixes. Dans cet article, nous présentons ChunkAttention, un module d'auto-attention conscient des préfixes, capable de détecter les préfixes de prompts correspondants à travers plusieurs requêtes et de partager leurs tenseurs clé/valeur en mémoire à l'exécution pour améliorer l'utilisation mémoire du cache KV. Cela est réalisé en divisant les tenseurs clé/valeur monolithiques en morceaux plus petits et en les structurant dans un arbre de préfixes auxiliaire. Par conséquent, sur la base du cache KV structuré en arbre de préfixes, nous concevons un noyau d'auto-attention efficace, où un algorithme de partition en deux phases est implémenté pour améliorer la localité des données lors du calcul de l'auto-attention en présence de prompts système partagés. Les expériences montrent que ChunkAttention peut accélérer le noyau d'auto-attention de 3,2 à 4,8 fois par rapport à l'implémentation de pointe, avec une longueur de prompt système variant de 1024 à 4096.
English
Self-attention is an essential component of large language models(LLMs) but a
significant source of inference latency for long sequences. In multi-tenant
LLMs serving scenarios, the compute and memory operation cost of self-attention
can be optimized by using the probability that multiple LLM requests have
shared system prompts in prefixes. In this paper, we introduce ChunkAttention,
a prefix-aware self-attention module that can detect matching prompt prefixes
across multiple requests and share their key/value tensors in memory at runtime
to improve the memory utilization of KV cache. This is achieved by breaking
monolithic key/value tensors into smaller chunks and structuring them into the
auxiliary prefix tree. Consequently, on top of the prefix-tree based KV cache,
we design an efficient self-attention kernel, where a two-phase partition
algorithm is implemented to improve the data locality during self-attention
computation in the presence of shared system prompts. Experiments show that
ChunkAttention can speed up the self-attention kernel by 3.2-4.8times
compared to the start-of-the-art implementation, with the length of the system
prompt ranging from 1024 to 4096.