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) 캐시 크기는 특히 디코딩 단계에서 메모리 제약을 심화시키는 중대한 병목 현상을 초래합니다. 이 병목 현상을 해결하기 위해 설계된 기존의 희소 어텐션 메커니즘은 두 가지 제한이 있습니다: (1) 주로 가장 관련성 높은 토큰을 신뢰할 수 없이 식별하는 경우가 많으며, (2) 연속된 트랜스포머 레이어 간에 토큰 선택의 공간적 일관성을 간과하여 성능 저하와 토큰 선택에서 상당한 오버헤드를 초래할 수 있습니다. 본 논문은 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.