TidalDecode : Décodage LLM Rapide et Précis avec Attention Éparse Persistante de Position

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

October 7, 2024

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

cs.AI

Résumé

Les grands modèles de langage (LLM) ont entraîné des avancées significatives dans diverses tâches de traitement automatique du langage naturel (TALN), les modèles à long contexte gagnant en importance pour traiter des entrées étendues. Cependant, l'augmentation de la taille du cache clé-valeur (KV) requise par les architectures Transformer intensifie les contraintes de mémoire, en particulier lors de la phase de décodage, créant un goulot d'étranglement significatif. Les mécanismes d'attention épars existants conçus pour résoudre ce goulot d'étranglement présentent deux limitations : (1) ils échouent souvent à identifier de manière fiable les jetons les plus pertinents pour l'attention, et (2) ils négligent la cohérence spatiale de la sélection de jetons à travers les couches Transformer consécutives, ce qui peut entraîner une dégradation des performances et des frais généraux substantiels dans la sélection de jetons. Cet article présente TidalDecode, un algorithme et un système simples mais efficaces pour un décodage rapide et précis des LLM grâce à une attention épars persistante en position. TidalDecode exploite la cohérence spatiale des jetons sélectionnés par les méthodes d'attention éparses existantes et introduit quelques couches de sélection de jetons qui effectuent une attention complète pour identifier les jetons avec les scores d'attention les plus élevés, tandis que toutes les autres couches effectuent une attention éparses avec les jetons pré-sélectionnés. Cette conception permet à TidalDecode de réduire considérablement les frais généraux de la sélection de jetons pour l'attention éparses sans sacrifier la qualité des résultats générés. L'évaluation sur un ensemble diversifié de LLM et de tâches montre que TidalDecode correspond étroitement aux performances génératives des méthodes à attention complète tout en réduisant la latence de décodage des LLM jusqu'à 2,1 fois.

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.