TidalDecode: Schnelles und präzises LLM-Decodieren mit positionspersistenter spärlicher Aufmerksamkeit

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

October 7, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben signifikante Fortschritte in verschiedenen NLP-Aufgaben vorangetrieben, wobei Langkontextmodelle an Bedeutung gewinnen, um erweiterte Eingaben zu verarbeiten. Die zunehmende Größe des Schlüssel-Wert (KV)-Caches, die von Transformer-Architekturen benötigt wird, verstärkt jedoch die Speicherbeschränkungen, insbesondere während der Dekodierungsphase, was einen signifikanten Engpass schafft. Bestehende spärliche Aufmerksamkeitsmechanismen, die entwickelt wurden, um diesen Engpass zu bewältigen, haben zwei Einschränkungen: (1) sie scheitern oft daran, die relevantesten Token für die Aufmerksamkeit zuverlässig zu identifizieren, und (2) sie übersehen die räumliche Kohärenz der Tokenauswahl über aufeinanderfolgende Transformer-Schichten hinweg, was zu Leistungsverschlechterung und erheblichem Mehraufwand bei der Tokenauswahl führen kann. Dieses Papier stellt TidalDecode vor, einen einfachen, aber effektiven Algorithmus und ein System für schnelles und präzises LLM-Decodieren durch positionspersistente spärliche Aufmerksamkeit. TidalDecode nutzt die räumliche Kohärenz der von bestehenden spärlichen Aufmerksamkeitsmethoden ausgewählten Token und führt einige Tokenauswahlschichten ein, die eine vollständige Aufmerksamkeit durchführen, um die Tokens mit den höchsten Aufmerksamkeitswerten zu identifizieren, während alle anderen Schichten spärliche Aufmerksamkeit mit den vorab ausgewählten Tokens durchführen. Dieses Design ermöglicht es TidalDecode, den Mehraufwand bei der Tokenauswahl für spärliche Aufmerksamkeit wesentlich zu reduzieren, ohne die Qualität der generierten Ergebnisse zu beeinträchtigen. Die Evaluation an einer vielfältigen Gruppe von LLMs und Aufgaben zeigt, dass TidalDecode die generative Leistung von Methoden mit voller Aufmerksamkeit nahezu erreicht, während die LLM-Decodierungs-Latenz um bis zu 2,1-fach reduziert wird.

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.