通过预测复用修复机制，显著加速长文本推理中的动态稀疏注意力
Predict, Reuse, and Repair: Accelerating Dynamic Sparse Attention for Long-Context LLM Decoding
- 预测未来可能用到的键值块，提前推测计算以减少等待时间
- 在真实选择结果返回前完成部分计算，降低每令牌解码延迟40%
- 适合需要高效长上下文处理的大模型应用，如文档生成与分析
动态稀疏注意力（DSA）通过仅关注每个查询相关的前K个键值块来加速长上下文大模型解码，但其选择与注意力计算之间的串行依赖成为新的延迟瓶颈。本文提出PRR运行时机制，利用DSA选择的时间局部性，预测可能使用的块，在选择过程进行时推测注意力计算，并在真实选择结果确定后增量修复遗漏块。PRR采用轻量级EMA预测器、基于性能分析的推测预算策略，确保推测工作不阻塞关键路径，并使用FlashAttention基修复内核，通过在线软最大统计信息将遗漏块融入部分注意力状态。在多个长上下文基准和典型DSA方法上，PRR将每令牌解码延迟降低最高达40%，同时保持下游任务准确率。代码开源：https://github.com/Tianyu9748/Incremental_FlashAttention
原文摘要 · Abstract (English)
Dynamic sparse attention (DSA) accelerates long-context LLM decoding by attending to only the top-K KV blocks relevant to each query, but it introduces a serialized selection-to-attention dependency that emerges as a new latency bottleneck. We present PRR, a speculate-reuse-repair runtime that exploits temporal locality in DSA selections to predict likely blocks, speculate the attention over them while selection is in flight, and incrementally repair missed blocks once the true selected set is known. PRR uses a lightweight EMA-based predictor, a profiling-guided speculation budget that keeps speculative work off the critical path, and a FlashAttention-based repair kernel that folds missed blocks into the partial attention state using online-softmax statistics. Across long-context benchmarks and representative DSA methods, PRR reduces per-token decoding latency by up to 40% while preserving downstream task accuracy. Github: https://github.com/Tianyu9748/Incremental_FlashAttention
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。