LServe: Servicio eficiente de LLM para secuencias largas con atención dispersa unificada

LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention

February 20, 2025

Autores: Shang Yang, Junxian Guo, Haotian Tang, Qinghao Hu, Guangxuan Xiao, Jiaming Tang, Yujun Lin, Zhijian Liu, Yao Lu, Song Han

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) han demostrado un potencial notable en el procesamiento de secuencias largas. Sin embargo, servir de manera eficiente estos modelos de contexto extenso sigue siendo un desafío debido a la complejidad computacional cuadrática de la atención en la etapa de prefilling y a la gran huella de memoria de la caché KV en la etapa de decodificación. Para abordar estos problemas, presentamos LServe, un sistema eficiente que acelera el servicio de LLMs de secuencias largas mediante atención híbrida dispersa. Este método unifica diferentes patrones de dispersión estructurada, compatibles con hardware, tanto para la atención en prefilling como en decodificación, en un único marco de trabajo, donde los cálculos en tokens menos importantes se omiten en bloques. LServe demuestra la compatibilidad de la dispersión estática y dinámica en la atención de LLMs de contexto extenso. Este diseño permite aceleraciones multiplicativas al combinar estas optimizaciones. Específicamente, convertimos la mitad de las cabezas de atención en cabezas de transmisión casi gratuita tanto en la etapa de prefilling como en la de decodificación. Además, descubrimos que solo se requiere un número constante de páginas KV para preservar las capacidades de contexto extenso, independientemente de la longitud del contexto. Luego, diseñamos una política jerárquica de selección de páginas KV que poda dinámicamente las páginas KV basándose en la similitud centrada en la consulta. En promedio, LServe acelera el prefilling de LLMs hasta 2.9 veces y la decodificación entre 1.3 y 2.1 veces en comparación con vLLM, manteniendo la precisión en contextos extensos. El código está disponible en https://github.com/mit-han-lab/omniserve.

English

Large language models (LLMs) have shown remarkable potential in processing long sequences, yet efficiently serving these long-context models remains challenging due to the quadratic computational complexity of attention in the prefilling stage and the large memory footprint of the KV cache in the decoding stage. To address these issues, we introduce LServe, an efficient system that accelerates long-sequence LLM serving via hybrid sparse attention. This method unifies different hardware-friendly, structured sparsity patterns for both prefilling and decoding attention into a single framework, where computations on less important tokens are skipped block-wise. LServe demonstrates the compatibility of static and dynamic sparsity in long-context LLM attention. This design enables multiplicative speedups by combining these optimizations. Specifically, we convert half of the attention heads to nearly free streaming heads in both the prefilling and decoding stages. Additionally, we find that only a constant number of KV pages is required to preserve long-context capabilities, irrespective of context length. We then design a hierarchical KV page selection policy that dynamically prunes KV pages based on query-centric similarity. On average, LServe accelerates LLM prefilling by up to 2.9x and decoding by 1.3-2.1x over vLLM, maintaining long-context accuracy. Code is released at https://github.com/mit-han-lab/omniserve.