Kascade通过复用关键注意力索引,显著加速长文本推理且保持高精度。
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
- 基于锚层计算并复用高权重键的Top-k索引,无需训练
- 解码阶段提速4.1倍,编码阶段提速2.2倍,精度接近全量注意力
- 适用于各类长文本任务,部署简单,支持Tile级高效实现
注意力是长上下文大模型推理中的主要延迟来源,而这一负载在推理模型和RAG中日益普遍。我们提出Kascade,一种无需训练的稀疏注意力方法,利用两个已知特性:1)后Softmax注意力本身具有内在稀疏性;2)相邻层中高权重键的身份稳定。Kascade在少量锚层中精确计算Top-k索引,并在中间复用层重复使用。锚层通过动态规划目标算法自动选择,以最大化跨层相似性,便于在不同模型间部署。该方法满足高效实现约束(如块级操作),适用于预填充和解码阶段。其索引选择具备头感知特性,实验表明这对高精度至关重要。在H100 GPU上,与FlashAttention-3基线相比,解码阶段最高提速4.1倍,预填充阶段提速2.2倍,同时在LongBench和AIME-24等长文本基准上紧密匹配密集注意力精度。
原文摘要 · Abstract (English)
Attention is the dominant source of latency during long-context LLM inference, an increasingly popular workload with reasoning models and RAG. We propose Kascade, a training-free sparse attention method that leverages known observations such as 1) post-softmax attention is intrinsically sparse, and 2) the identity of high-weight keys is stable across nearby layers. Kascade computes exact Top-k indices in a small set of anchor layers, then reuses those indices in intermediate reuse layers. The anchor layers are selected algorithmically, via a dynamic-programming objective that maximizes cross-layer similarity over a development set, allowing easy deployment across models. The method incorporates efficient implementation constraints (e.g. tile-level operations), across both prefill and decode attention. The Top-k selection and reuse in Kascade is head-aware and we show in our experiments that this is critical for high accuracy. Kascade achieves up to 4.1x speedup in decode attention and 2.2x speedup in prefill attention over FlashAttention-3 baseline on H100 GPUs while closely matching dense attention accuracy on long-context benchmarks such as LongBench and AIME-24.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。