TidalDecode: Fast and Accurate LLM Decoding with Position Persistent
  Sparse Attention
Authors: name: Lijie Yang; name: Zhihao Zhang; name: Zhuofu Chen; name: Zikun Li; name: Zhihao Jia
Published: 2024-10-07 14:30:27

Abstract: 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.