Reproduce “Twilight: Adaptive Attention Sparsity with Hierarchical Top-$p$ Pruning” on OpenResearch

About this paper

Chaofan Lin, Jiaming Tang, Shuo Yang, Hanshuo Wang, Tian Tang, Boyu Tian, Ion Stoica, Song Han and 1 more

Twilight is an adaptive attention sparsity framework that leverages hierarchical Top-p pruning to optimize large language model inference. It dynamically adjusts the Key-Value cache budget, achieving substantial speedups in self-attention (up to 15.8x) and end-to-end decoding (up to 3.9x) with minimal to no accuracy degradation on various benchmarks.

Key insights

• Dynamically adjusting the KV cache budget using top-p pruning is more effective than fixed top-k methods, as it inherently adapts to 'focused' vs. 'diffuse' attention distributions.
• A two-stage hierarchical pruning approach allows leveraging the strengths of existing token selection algorithms while adding the adaptivity of top-p for final budget determination.
• Efficient hardware-aware implementations, such as 4-bit quantized SpGEMV and optimized parallel binary search, are crucial for translating algorithmic improvements into practical performance gains.

Summary by alphaXiv.