Command Palette
Search for a command to run...
Kascade:一种面向长上下文LLM推理的实用稀疏注意力方法
Kascade:一种面向长上下文LLM推理的实用稀疏注意力方法
Dhruv Deshmukh Saurabh Goyal Nipun Kwatra Ramachandran Ramjee
Abstract
在长上下文大语言模型(LLM)推理中,注意力机制是导致延迟的主要因素,而这类工作负载在推理模型与检索增强生成(RAG)系统中正变得日益普遍。我们提出Kascade,一种无需训练的稀疏注意力方法,其利用了若干已知观察结果:1)经过softmax后的注意力权重本质上具有稀疏性;2)在相邻层之间,高权重键(key)的标识具有稳定性。Kascade在一组少量的锚点层(anchor layers)中精确计算Top-k注意力索引,随后在中间可复用层(reuse layers)中重用这些索引。锚点层通过动态规划目标算法自动选择,该目标旨在最大化开发集上的跨层相似性,从而实现对多种模型的简便部署。该方法在预填充(prefill)与解码(decode)注意力阶段均引入了高效的实现约束(如基于tile的操作),确保计算效率。Kascade中的Top-k选择与复用机制具备头感知(head-aware)特性,我们的实验表明,这一设计对于保持高精度至关重要。在H100 GPU上,Kascade相较于FlashAttention-3基线,在解码注意力阶段实现最高达4.1倍的加速,在预填充注意力阶段实现最高2.2倍的加速,同时在LongBench和AIME-24等长上下文基准测试中,性能几乎完全匹配密集注意力(dense attention)的精度水平。