LServe: Serviço Eficiente de LLM para Sequências Longas com Atenção Esparsa 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

Resumo

Modelos de linguagem de grande escala (LLMs) demonstraram um potencial notável no processamento de sequências longas, mas a eficiência no atendimento desses modelos de contexto longo continua sendo um desafio devido à complexidade computacional quadrática da atenção na etapa de pré-preenchimento e à grande pegada de memória do cache KV na etapa de decodificação. Para abordar esses problemas, introduzimos o LServe, um sistema eficiente que acelera o atendimento de LLMs de sequência longa por meio de atenção esparsa híbrida. Esse método unifica diferentes padrões de esparsidade estruturada, amigáveis ao hardware, para a atenção tanto no pré-preenchimento quanto na decodificação, em um único framework, onde os cálculos em tokens menos importantes são pulados em blocos. O LServe demonstra a compatibilidade da esparsidade estática e dinâmica na atenção de LLMs de contexto longo. Esse design permite acelerações multiplicativas ao combinar essas otimizações. Especificamente, convertemos metade das cabeças de atenção em cabeças de streaming quase gratuitas tanto na etapa de pré-preenchimento quanto na de decodificação. Além disso, descobrimos que apenas um número constante de páginas KV é necessário para preservar as capacidades de contexto longo, independentemente do comprimento do contexto. Em seguida, projetamos uma política hierárquica de seleção de páginas KV que poda dinamicamente as páginas KV com base na similaridade centrada na consulta. Em média, o LServe acelera o pré-preenchimento de LLMs em até 2,9x e a decodificação em 1,3-2,1x em relação ao vLLM, mantendo a precisão de contexto longo. O código está disponível em 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.