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.