ChunkAttention: Atención Autónoma Eficiente con Caché KV Consciente de Prefijos y Partición Bifásica
ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition
February 23, 2024
Autores: Lu Ye, Ze Tao, Yong Huang, Yang Li
cs.AI
Resumen
La autoatención es un componente esencial de los modelos de lenguaje grandes (LLMs, por sus siglas en inglés), pero también una fuente significativa de latencia en la inferencia para secuencias largas. En escenarios de servicio multiinquilino de LLMs, el costo computacional y de operaciones de memoria de la autoatención puede optimizarse utilizando la probabilidad de que múltiples solicitudes de LLM compartan prefijos de sistema en sus secuencias iniciales. En este artículo, presentamos ChunkAttention, un módulo de autoatención consciente de prefijos que puede detectar prefijos de sistema coincidentes en múltiples solicitudes y compartir sus tensores de clave/valor en memoria durante la ejecución para mejorar la utilización de la memoria de la caché KV. Esto se logra dividiendo los tensores monolíticos de clave/valor en fragmentos más pequeños y organizándolos en un árbol de prefijos auxiliar. En consecuencia, sobre la base de la caché KV estructurada en árbol de prefijos, diseñamos un núcleo de autoatención eficiente, donde se implementa un algoritmo de partición en dos fases para mejorar la localidad de los datos durante el cálculo de la autoatención en presencia de prefijos de sistema compartidos. Los experimentos muestran que ChunkAttention puede acelerar el núcleo de autoatención entre 3.2 y 4.8 veces en comparación con la implementación más avanzada, con longitudes de prefijos de sistema que van desde 1024 hasta 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.