Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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

About

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

Tianyu Wang, Gourav Rattihalli, Aditya Dhakal, Junbo Li, Zhiwei Ren, Dejan Milojicic, Longfei Shangguan• 2026

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningMATH500
Speedup1.59
28
Complex ReasoningAIME
Decoding Speedup1.64
12
General LLM accelerationLongBench, InfiniteBench, RULER, AIME, MATH500 Average
Average Decoding Speedup1.56
12
Long-context UnderstandingLongBench
Decoding Speedup1.64
12
Long-context UnderstandingInfiniteBench
Decoding Speedup1.48
12
Long-context UnderstandingRULER
Decoding Speedup1.61
12
Showing 6 of 6 rows

Other info

Follow for update