TidalDecode: Decodificação LLM Rápida e Precisa com Atenção Esparsa Persistente de Posição

TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention

October 7, 2024

Autores: Lijie Yang, Zhihao Zhang, Zhuofu Chen, Zikun Li, Zhihao Jia

cs.AI

Resumo

Os grandes modelos de linguagem (LLMs) impulsionaram avanços significativos em diversas tarefas de Processamento de Linguagem Natural (NLP), com modelos de contexto longo ganhando destaque para lidar com entradas extensas. No entanto, o aumento do tamanho do cache chave-valor (KV) exigido pelas arquiteturas Transformer intensifica as restrições de memória, especialmente durante a fase de decodificação, criando um gargalo significativo. Mecanismos de atenção esparsa existentes projetados para lidar com esse gargalo têm duas limitações: (1) frequentemente falham em identificar de forma confiável os tokens mais relevantes para a atenção e (2) ignoram a coerência espacial da seleção de tokens entre camadas consecutivas do Transformer, o que pode levar à degradação de desempenho e a um overhead substancial na seleção de tokens. Este artigo apresenta o TidalDecode, um algoritmo e sistema simples, porém eficaz, para decodificação rápida e precisa de LLMs por meio de atenção esparsa persistente à posição. O TidalDecode aproveita a coerência espacial dos tokens selecionados por métodos de atenção esparsa existentes e introduz algumas camadas de seleção de tokens que realizam atenção total para identificar os tokens com os maiores escores de atenção, enquanto todas as outras camadas realizam atenção esparsa com os tokens pré-selecionados. Esse projeto permite que o TidalDecode reduza substancialmente o overhead de seleção de tokens para atenção esparsa sem sacrificar a qualidade dos resultados gerados. A avaliação em um conjunto diversificado de LLMs e tarefas mostra que o TidalDecode se equipara de perto ao desempenho generativo de métodos de atenção total, reduzindo a latência de decodificação de LLMs em até 2,1 vezes.

English

Large language models (LLMs) have driven significant advancements across diverse NLP tasks, with long-context models gaining prominence for handling extended inputs. However, the expanding key-value (KV) cache size required by Transformer architectures intensifies the memory constraints, particularly during the decoding phase, creating a significant bottleneck. Existing sparse attention mechanisms designed to address this bottleneck have two limitations: (1) they often fail to reliably identify the most relevant tokens for attention, and (2) they overlook the spatial coherence of token selection across consecutive Transformer layers, which can lead to performance degradation and substantial overhead in token selection. This paper introduces TidalDecode, a simple yet effective algorithm and system for fast and accurate LLM decoding through position persistent sparse attention. TidalDecode leverages the spatial coherence of tokens selected by existing sparse attention methods and introduces a few token selection layers that perform full attention to identify the tokens with the highest attention scores, while all other layers perform sparse attention with the pre-selected tokens. This design enables TidalDecode to substantially reduce the overhead of token selection for sparse attention without sacrificing the quality of the generated results. Evaluation on a diverse set of LLMs and tasks shows that TidalDecode closely matches the generative performance of full attention methods while reducing the LLM decoding latency by up to 2.1x.