Search papers, labs, and topics across Lattice.
This paper introduces PRR, a runtime system that enhances dynamic sparse attention (DSA) for long-context LLM decoding by predicting and reusing attention blocks, thereby addressing the latency bottleneck caused by serialized selection dependencies. By leveraging temporal locality and a lightweight EMA-based predictor, PRR allows for speculative attention calculations while waiting for the final selection, followed by an efficient repair mechanism to incorporate any missed blocks. The results demonstrate a significant reduction in per-token decoding latency by up to 40% without compromising accuracy on downstream tasks.
PRR slashes decoding latency by up to 40% in long-context LLMs while maintaining accuracy, revolutionizing the efficiency of dynamic sparse attention.
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
