TidalDecode: Быстрая и точная декодирование LLM с постоянным положением разреженного внимания.

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

October 7, 2024

Авторы: Lijie Yang, Zhihao Zhang, Zhuofu Chen, Zikun Li, Zhihao Jia

cs.AI

Аннотация

Большие языковые модели (LLM) привели к значительным достижениям в различных задачах обработки естественного языка (NLP), причем модели с учетом длинного контекста стали более популярными для обработки расширенных входных данных. Однако увеличение размера кэша ключ-значение (KV), необходимого для архитектур Transformer, усиливает ограничения памяти, особенно во время фазы декодирования, что создает значительное узкое место. Существующие механизмы разреженного внимания, разработанные для решения этой проблемы, имеют два ограничения: (1) они часто не могут надежно определить наиболее релевантные токены для внимания, и (2) они не учитывают пространственную согласованность выбора токенов через последовательные слои Transformer, что может привести к снижению производительности и значительным накладным расходам на выбор токенов. В данной статье представлен TidalDecode - простой, но эффективный алгоритм и система для быстрого и точного декодирования LLM с помощью разреженного внимания с постоянным положением. TidalDecode использует пространственную согласованность токенов, выбранных существующими механизмами разреженного внимания, и вводит несколько слоев выбора токенов, которые выполняют полное внимание для определения токенов с наивысшими оценками внимания, в то время как все остальные слои выполняют разреженное внимание с предварительно выбранными токенами. Такая концепция позволяет TidalDecode значительно сократить накладные расходы на выбор токенов для разреженного внимания, не жертвуя качеством сгенерированных результатов. Оценка на разнообразном наборе LLM и задач показывает, что TidalDecode близок по производительности генерации к методам полного внимания, снижая задержку декодирования LLM до 2,1 раза.

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.