Predict, Reuse, and Repair: Accelerating Dynamic Sparse Attention for Long-Context LLM Decoding

2026-06-29

AI Summary

This work addresses the significant latency incurred by dynamic sparse attention (DSA) during long-context LLM decoding, which stems from the serial dependency between query processing and key-value (KV) block selection. To overcome this bottleneck, the authors propose PRR, a runtime system that introduces speculative execution and incremental correction mechanisms into DSA for the first time. Leveraging temporal locality, PRR employs a lightweight exponential moving average (EMA) predictor to speculatively select KV blocks and compute attention in parallel. Upon revealing the true KV selection, it efficiently corrects any omissions using a profiling-guided speculation budget and an online Softmax-statistics-driven incremental FlashAttention repair kernel. Experiments demonstrate that PRR reduces per-token decoding latency by up to 40% across diverse DSA methods and long-context benchmarks, without compromising downstream task accuracy.

Abstract

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

Problem

• dynamic sparse attention
• latency bottleneck
• long-context LLM decoding
• selection-to-attention dependency

Innovation

• Dynamic Sparse Attention
• Speculative Execution
• Incremental Repair
• FlashAttention
• Temporal Locality