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

一、论文概述

项目	内容
标题	Predict, Reuse, and Repair: Accelerating Dynamic Sparse Attention for Long-Context LLM Decoding
作者	Tianyu Wang, Gourav Rattihalli, Aditya Dhakal, Junbo Li, Zhiwei Ren, Dejan Milojicic, Longfei Shangguan
机构	1 University of Pittsburgh, Pittsburgh, PA, USA; 2 HPE Labs, Milpitas, CA, USA
论文	arXiv:2606.30389
代码	github.com/Tianyu9748/Incremental_FlashAttention
发布	2026-06-29 (cs.LG)
许可	BSD-3-Clause

核心贡献：

1. 发现 DSA（Dynamic Sparse Attention）引入的新瓶颈：selection-to-attention dependency——每个 decoding step 中 selection 严格串行在 attention 之前，占 per-token 生成延迟高达 41%
2. 提出 PRR（Predict, Reuse, Repair），一个 correctness-preserving 的投机注意力运行时，利用 DSA 选择的时序局部性（~68% 块重复率），在 selection 进行时并行执行 speculative attention
3. 设计轻量级 EMA-based predictor，通过 prefill 阶段的 importance score 轨迹进行 prompt-adaptive 超参数校准，预测精度达 ~91%（优化后 ~98%）
4. 实现基于 FlashAttention online-softmax recurrence 的定制化 CUDA kernel，支持 incremental repair——将 missed blocks 的贡献精确合并到已有 accumulator，无 accuracy loss
5. 在 6 个 LLM 和 5 个 long-context benchmarks 上的评估表明：相比 Quest 平均加速 1.42×，相比 InfLLM-v2 平均加速 1.56×，同时保持相同下游任务准确率

二、核心思想

问题定义

DSA（如 Quest、InfLLM-V2）通过在每个 decoding step 只选择 top-K KV blocks 来加速长上下文 LLM 推理。但这引入了一个新的关键路径瓶颈：selection-to-attention dependency。

具体而言，每个 decoding step 包含三个阶段：

1. Selection：运行 compressed attention，识别 top-K 个最相关的 KV block
2. Attention：仅对选中的 top-K blocks 执行稀疏 attention
3. FFN：标准前馈层

由于 block identities 在 selection 完成前未知，attention 严格串行等待 selection 结束。随着 context length 增长，这一依赖变得日益昂贵——selection 占 per-token generation latency 高达 41%。

解决方案概述

PRR 的核心洞察：DSA 的 block 选择具有显著的时序稳定性（temporal locality）。实验测量表明，约 68% 的 selected blocks 在连续 decoding steps 间被重复选择。

PRR 利用这一特性，采用三步策略：

1. Predict：使用轻量级 EMA predictor 预测当前 step 的 top-K blocks，提前执行 speculative attention
2. Reuse：speculative attention 与 selection 并行执行，覆盖预测命中块的 attention 计算
3. Repair：当真实 top-K 集合 A 已知后，对 missed blocks M = A \ P 执行 incremental attention repair，将贡献精确合并到已有 accumulator

最终 attention 输出覆盖完整的 true selected set A，与标准 DSA 语义等价，accuracy 无损。

三、技术架构

整体框架

PRR 由三个核心组件构成：

组件	功能	关键技术
EMA Predictor (§4.1)	预测下一个 decoding step 的 top-K blocks	双指数平滑 + prompt-adaptive 超参数校准
Incremental Repair Kernel (§4.3)	将 missed blocks 精确合并到 partial attention	FlashAttention online-softmax recurrence
Speculation Budget (§4.4)	控制 speculative work 不超过 critical path	profile-guided dynamic budget ratio δ

Motivation

Selection-to-Attention Dependency (§2.1)

DSA 每个 decoding step 包含三个阶段：selection → attention → FFN。Selection 阶段需要运行 compressed attention 来识别 top-K KV blocks。由于 block identities 在 selection 完成前不可知，attention 严格串行等待。

Temporal Locality of Block Selection (§2.2)

在五个 long-context benchmarks 上测量连续 decoding steps 间的 block overlap rate：

关键发现：两种 DSA 方法在所有 benchmark 上均稳定重选超过 65% 的 blocks，平均重叠率约 68%。

GPU Headroom (§2.2)

关键发现：即使在高 context 长度下，GPU SM、L2-BW 和 DRAM-BW 利用率均低于 40%，说明有充足的空闲资源执行 speculative attention 而不会与 critical path 竞争。

核心公式

EMA Predictor (§4.1)

PRR 使用训练-free 的轻量级 EMA predictor，直接操作 DSA 产生的 block importance scores。

令 $IS_i^t$ 为 decoding step $t$ 中 block $i$ 的真实 importance score。在 step $t$ 的 selection 阶段运行之前，PRR 使用历史分数预测 $\widehat{IS}_i^t$。

预测阶段（extrapolation）：

$$\widehat{IS}_{i}^{t}=\ell_{i}^{t-1}+\gamma v_{i}^{t-1} \tag{2}$$

其中 $\gamma \in [0, 1]$ 控制如何激进地外推近期趋势。

更新阶段（update）：

$$\begin{aligned} v_{i}^{t} &= \beta\bigl(IS_{i}^{t}-\ell_{i}^{t-1}\bigr)+(1-\beta)v_{i}^{t-1}, \\ \ell_{i}^{t} &= \alpha IS_{i}^{t}+(1-\alpha)\ell_{i}^{t-1}. \end{aligned} \tag{3}$$

其中 $\alpha \in (0, 1]$ 控制 level 跟随新分数的速度，$\beta \in (0, 1]$ 控制 trend estimate 适应分数变化的速度。

当 block $i$ 首次进入 KV cache 时（step $\tau$），初始化 $\ell_i^\tau = IS_i^\tau$，$v_i^\tau = 0$。

Prompt-Adaptive Hyperparameter Calibration (§4.2)

固定超参数 $(\alpha, \beta, \gamma)$ 对不同 prompt 效果差异大。PRR 在 prefill 阶段进行 per-prompt 校准：

令 $T_p$ 为 prefill token 数量，$IS \in \mathbb{R}^{T_p \times N_p}$ 为 prefill 阶段产生的 importance score 矩阵（row $\phi$ 是 prefill token $\phi$ 处的 compressed-attention score vector）。

对候选超参数 $\theta = (\alpha, \beta, \gamma)$，模拟 predictor 过 $IS$，产生 predicted selection set $P_\phi(\theta)$，测量 score-weighted hit rate：

$$H\bigl(\theta\,;\,IS\bigr)\;=\;\frac{1}{T_{p}}\sum_{\phi}\frac{\sum_{b\in P_{\phi}(\theta)\cap A_{\phi}}IS_{\phi,b}}{\sum_{b\in A_{\phi}}IS_{\phi,b}} \tag{4}$$

其中 $A_\phi$ 是 prefill token $\phi$ 处的 ground-truth top-K block selection set，$IS_{\phi,b}$ 是 block $b$ 在 token $\phi$ 处的 importance score。

最优超参数选择：

$$\theta^{\star}=\arg\max_{\theta}H\bigl(\theta;IS\bigr) \tag{5}$$

搜索网格：$\alpha \in [0.2, 0.8]$（步长 0.2）、$\beta \in [0.1, 0.5]$（步长 0.1）、$\gamma \in [0, 0.75]$（步长 0.25），共 80 个候选。搜索与 prefill 阶段重叠，仅增加 0.06 ms 额外延迟。

Incremental Repair via Online Softmax (§4.3)

Mispredictions 不可避免（$M = A \setminus P \neq \emptyset$），忽略 $M$ 会产生不精确的 softmax 输出且误差随 decoding steps 累积。

PRR 实现定制化 CUDA kernel，基于 FlashAttention 的 online-softmax recurrence 进行 incremental repair。给定 speculative attention 的输出 $(O^{\text{spec}}, \ell^{\text{spec}}, m^{\text{spec}})$ 和 missed blocks 的 KV 条目，第二个 kernel 应用以下 recurrence：

$$\begin{aligned} m^{(t+1)} &= \max(m^{(t)}, \tilde{m}) \\ \ell^{(t+1)} &= e^{m^{(t)}-m^{(t+1)}}\,\ell^{(t)} + e^{\tilde{m}-m^{(t+1)}}\,\tilde{\ell} \\ O^{(t+1)} &= \frac{e^{m^{(t)}-m^{(t+1)}}\,\ell^{(t)}}{\ell^{(t+1)}}\,O^{(t)} + \frac{e^{\tilde{m}-m^{(t+1)}}}{\ell^{(t+1)}}\sum_{j}e^{S_{t+1,j}-\tilde{m}}\,V_{t+1,j} \end{aligned} \tag{6}$$

其中 $\tilde{m}, \tilde{\ell}$ 是 missed blocks 的 running max 和 log-sum-exp，$S_{t+1,j}$ 和 $V_{t+1,j}$ 是 missed blocks 的 attention scores 和 values。

关键实现细节：

• Speculative kernel 额外存储 output $O$、denominator $\ell$ 和 running maximum $m$ 到 HBM
• 修改是 surgical 的：inner-loop arithmetic 不变，额外存储仅为每 head 两个 scalars
• Missed blocks 以 tiled 方式 stream through on-chip memory，不 materialize intermediate score matrix
• Both FLOPs 和 memory traffic 仅随 $|M|$ 缩放

Speculation Budget Optimization (§4.4)

增大 speculation budget 可覆盖更多 true selected blocks并减少 missed blocks，但也会增加 speculative KV movement 和 attention computation，可能延长 critical path。

PRR 将预算约束表述为优化问题：

$$\min\;|M|=|A\setminus P|\quad\text{s.t.}\quad|P|\leq\delta|A| \tag{1}$$

其中 $\delta$ 是 speculation budget ratio。PRR 采用 profile-guided 方法：对每个 model-hardware-DSA 配置进行一次 offline profiling sweep，测量 selection latency 和 speculative-attention latency 在不同 context lengths 和 $\delta$ 下的值，动态选择 $\delta$ 使 speculative attention latency 适配当前 selection window。

Profiling 的 context lengths：4K, 8K, 16K, 32K, 64K。

设计空间探索

PRR 解决两个核心挑战：

1. Temporal locality alone is not sufficient：盲目重用前一步的 top-K block indices 仅获得 modest hit rate。需要更智能的 predictor。
2. Speculation must preserve DSA result：最终 attention 必须在 true DSA-selected set 上计算。现有引擎（vLLM、SGLang）和 attention kernels（FlashAttention）不支持将 missed blocks 增量合并到 partial attention result。

四、核心创新

创新点	说明	理论/实验依据
Selection-to-attention dependency 识别	首次将 DSA 的 latency bottleneck 定位为 selection→attention 串行依赖	§2.1: selection 占 per-token 延迟 41%
EMA-based block predictor	训练-free 的轻量级 predictor，直接操作 DSA importance scores	Eq. (2)-(3): 双指数平滑，~68% base overlap
Prompt-adaptive calibration	Prefill 阶段 per-prompt 超参数搜索，无需在线训练	Eq. (4)-(5): grid search over 80 candidates, 0.06ms overhead
Incremental repair kernel	基于 FlashAttention online-softmax 的定制化 CUDA kernel	Eq. (6): exact combine, no accuracy loss
Critical-path-aware budget	Profile-guided dynamic δ 确保 speculative work 不超出 selection window	Fig. 5: balanced timeline

五、代码实现分析

项目结构

Incremental_FlashAttention/
├── hopper/                          # FA-3 (CUTLASS 3.x, SM90 TMA+WGMMA+WS) — 主要 sparse 工作区
│   ├── mainloop_fwd_sm90_tma_gmma_ws.hpp   # Forward mainloop with sparse dispatch (~131KB)
│   ├── flash_fwd_kernel_sm90.h             # Forward kernel with sparse_block_table support
│   ├── flash_fwd_combine.cu                # Combine kernel for speculative+sparse merge
│   ├── flash_fwd_combine_kernel.h          # Combine kernel header
│   ├── flash_attn_interface.py             # Python entry point
│   └── sparse_flash_attn_3/               # Renamed package namespace
├── csrc/flash_attn/              # FA-2 kernels (CUTLASS 2.x) — Ampere fallback
├── sparse_flash_attn_2/          # FA-2 Python package (renamed)
├── tests/                        # Correctness tests
├── benchmarks/                   # Performance benchmarks
├── setup.py                      # Build configuration
└── instruction.md               # Build instructions

关键实现文件

文件	说明	大小
hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp	FA-3 forward mainloop，含 load_sparse() + mma_sparse() sparse block dispatch	131KB
hopper/flash_fwd_kernel_sm90.h	Forward kernel template，支持 sparse_block_table 参数	29KB
hopper/flash_fwd_combine.cu	Exact combine kernel，合并 speculative 和 repair partials	929B
hopper/flash_fwd_combine_kernel.h	Combine kernel header，online-softmax recurrence 实现	26KB

API 接口

# 1. Sparse-only forward
from flash_attn_interface import flash_attn_with_sparse_block_table
out, lse = flash_attn_with_sparse_block_table(q, k, v, sparse_block_table, block_size=128)

# 2. Speculative + sparse with exact combine
from flash_attn_interface import flash_attn_speculative_sparse
out, lse = flash_attn_speculative_sparse(q, k_spec, v_spec, k, v, sparse_block_table, causal_spec=True)

硬件要求

• FA-3 sparse path: NVIDIA Hopper GPU (H100/H800, SM90)
• FA-2 path: Ampere (SM80) 兼容
• CUDA: ≥ 12.3 (推荐 12.8)
• Precision: FP16/BF16 (half precision)

六、实验结果

实验设置

配置项	值
GPU	NVIDIA H100, CUDA 12.8, TP=2
Models	GLM-4-9B, GLM-Z1-9B, DeepSeek-R1-8B, Llama3-8B, Qwen3-14B, Qwen3-32B
DSA Methods	Quest, InfLLM-v2
Benchmarks	LongBench, InfiniteBench, RULER, AIME, MATH500
Backend	FlashInfer’s BlockSparseAttention

DSA Speedup over Serial Execution (Table I)

Model	LongBench	InfiniteBench	RULER	AIME	MATH500	Avg.
Quest
GLM-4	1.50×	1.40×	1.45×	1.30×	1.42×	1.42×
GLM-Z1	1.48×	1.39×	1.43×	1.33×	1.45×	1.42×
DeepSeek-R1	1.39×	1.42×	1.47×	1.35×	1.39×	1.41×
Llama3	1.36×	1.40×	1.41×	1.39×	1.34×	1.38×
Qwen3-14B	1.44×	1.27×	1.41×	1.44×	1.45×	1.40×
Qwen3-32B	1.26×	1.34×	1.27×	1.27×	1.26×	1.28×
InfLLM-v2
GLM-4	1.64×	1.38×	1.58×	1.62×	1.59×	1.56×
GLM-Z1	1.61×	1.35×	1.61×	1.64×	1.55×	1.55×
DeepSeek-R1	1.30×	1.48×	1.61×	1.52×	1.28×	1.44×
Llama3	1.32×	1.45×	1.55×	1.55×	1.30×	1.43×
Qwen3-14B	1.56×	1.43×	1.51×	1.53×	1.53×	1.51×
Qwen3-32B	1.35×	1.27×	1.29×	1.32×	1.35×	1.31×

关键发现： PRR 在 Quest 上平均加速 1.35×，在 InfLLM-v2 上平均加速 1.47×。较小模型（9B）增益更大，Qwen3-32B 增益较低但仍显著。

Temporal Locality Overlap Rate (Table IV)

Model	LongBench	InfiniteBench	RULER	AIME	MATH500	Avg.
Quest (GLM-4)	98.65%	96.86%	98.36%	98.15%	98.25%	98.05%
InfLLM-v2 (GLM-4)	97.86%	97.14%	96.46%	98.15%	98.01%	97.52%

Ablation Study (Table II)

Benchmark	S1 (Reuse prev)	S2 (+Repair Kernel)	S3 (Full PRR)
LongBench	1.02×	1.34×	1.64×
InfiniteBench	1.00×	1.19×	1.38×
RULER	1.01×	1.34×	1.58×
AIME	1.02×	1.34×	1.62×
MATH500	1.01×	1.31×	1.59×
Avg.	1.01×	1.30×	1.56×

关键发现：

• S1（简单重用前一步 selection）几乎无改善——单个 missed block 触发 full recomputation 抵消了 speculation 收益
• S2（引入 repair kernel）将 avg speedup 提升至 1.30×
• S3（完整 PRR）达到 1.56× avg speedup，EMA predictor 将 overlap rate 从 91% 提升至 98%

Kernel Speed Comparison (Table III)

Block Size	1024 tokens	2048	4096	6144	8192	Avg.
16	1.17×	1.28×	2.00×	2.83×	3.68×	2.19×
32	1.17×	1.38×	2.09×	2.78×	3.69×	2.22×
64	1.17×	1.38×	2.26×	2.91×	3.64×	2.27×
128	1.17×	1.41×	2.23×	3.00×	3.71×	2.31×
256	1.03×	1.11×	2.00×	2.75×	3.64×	2.11×

关键发现： PRR 的 sparse kernel 在大多数 block size 和 token budget 组合下显著优于 FlashInfer 的 BlockSparseAttention。Token budget 越大，speedup 越高（最大 3.71× @ 8192 tokens）。

三大增益机制

1. EMA predictor 达到 ~91% average top-K overlap rate，使大部分 speculative attention work 被利用
2. Customized sparse kernel 大幅超越 BlockSparseAttention（avg 2.2× headroom），使 PRR 能在 speculation 中包含更多 blocks，将 overlap rate 从 91% 提升至 98% 且不增加延迟
3. Incremental repair 而非 full recomputation，保留 speculation 的 latency savings

七、相关工作

KV Cache Retrieval for DSA

• KV-dropping methods（如 Sinkhorn Attention 等）：通过丢弃不重要的 KV tokens 节省内存，但有信息丢失
• KV retrieval approaches（LouisKV, AsyncSpade）：将完整 KV cache 保存在 CPU 内存，按需 fetch query-relevant tokens。与 PRR 正交
• LouisKV / AsyncSpade：通过预测 query state prefetch relevant tokens，但近似估计可能 drop important tokens 并降低质量。PRR 基于 temporal locality of historical importance scores，且 repair 保证 zero accuracy degradation

Temporal Locality for KV Cache Management

• FlexiCache：exploit temporal stability across heads——stable heads 仅保留 top-K KV pages on GPU，unstable heads 保留全部。与 PRR 互补且可组合：PRR 的 EMA predictor 可指导 FlexiCache stable heads 内的页面保留策略

Dynamic Sparse Attention

• Quest：基于 compressed attention 的 top-K block selection
• InfLLM-v2：基于 inference-time KV page eviction 的 DSA
• NSA（Neural Sparse Attention）：需从头训练的 trainable DSA，PRR 未包含（需巨大计算资源）

八、总结

核心贡献

1. 瓶颈识别：首次将 DSA 的 selection-to-attention dependency 识别为新的 critical-path bottleneck，selection 占 per-token 延迟高达 41%
2. PRR 运行时：提出 correctness-preserving 的 speculate-reuse-repair runtime，打破 selection→attention 串行依赖
3. EMA Predictor：轻量级训练-free predictor，通过 prefill 阶段 prompt-adaptive 校准，达到 ~98% overlap rate
4. Incremental Repair Kernel：基于 FlashAttention online-softmax recurrence 的定制化 CUDA kernel，exact combine，无 accuracy loss
5. 显著性能提升：6 个 LLM × 5 个 benchmark 上，Quest 平均 1.42×、InfLLM-v2 平均 1.56× 加速

局限性

1. 仅评估 training-free DSA：仅评估 Quest 和 InfLLM-v2，未包含 trainable 方法如 NSA（需巨大计算资源从头训练）
2. 单一 GPU 架构：仅针对 NVIDIA Hopper (SM90) 和 half precision 优化 kernel。可扩展至 Blackwell 和其他低精度（如 fp8）
3. 仅 batch size 1：仅评估 batch size 1 场景。未来可探索 batch 内多请求间的 coordination
4. 主文仅展示 GLM-4-9B：其他模型的完整结果在 Appendix A 中

未来方向

1. 扩展 kernel optimization 至其他 GPU 架构（Blackwell）和精度（fp8）
2. 集成 trainable DSA mechanisms（如 NSA）
3. 探索 batch 内多请求的 stage coordination 以提升 decoding throughput
4. 与 FlexiCache 等 temporal locality 方法组合

九、参考资源

• 论文: arXiv:2606.30389
• 代码: github.com/Tianyu9748/Incremental_FlashAttention
• FlashAttention: github.com/Dao-AILab/flash-attention
• FlashInfer: github.com/flashinfer-ai/flashinfer
• Quest: Tang et al. “Quest: Quantile-based Efficient Streaming Attention for Long-context LLMs”
• InfLLM-v2: Zhao et al. “InfLLM: Infinite Context Length Streaming Attention for Large Language Models”

附图索引

编号	文件名	说明
Figure 1	figures/prr-dsa/figure-01-teaser.svg	Standard DSA vs PRR: serial selection→attention→FFN vs parallel speculative attention + incremental repair
Figure 2	figures/prr-dsa/figure-02-latency-breakdown.svg	(a) DSA 三阶段与时序相似性；(b) Quest/InfLLM-V2 延迟分解
Figure 3	figures/prr-dsa/figure-03-hit-rate.png	各 benchmark 下连续 selection 的块重叠率 (~68%)
Figure 4	figures/prr-dsa/figure-04-utilization.svg	SM/L2-BW/DRAM-BW 利用率 profiling (<40%)
Figure 5	figures/prr-dsa/figure-05-choose-alpha1.svg	Speculation budget ratio δ 对执行时间线的影响
Figure 6	figures/prr-dsa/figure-06-ema-calibration.svg	PRR 整体设计：prefill calibration + decode prediction + repair

附表格索引

编号	说明
Table I	DSA decoding speedup across models and benchmarks (Quest + InfLLM-v2)
Table II	Ablation: S1/S2/S3 speedup across benchmarks for GLM-4-9B + InfLLM-v2
Table III	PRR sparse kernel vs BlockSparseAttention speedup (block sizes × token budgets)
Table IV	EMA prediction overlap rate (%) for GLM-4-9B
Table V (Appendix)	Temporal locality across all 6 LLMs and benchmarks
Table VI (Appendix)	GPU utilization across batch sizes
Table VII (Appendix)	EMA overlap rate across all LLMs
Table VIII (Appendix)	Full ablation across stages, models, and benchmarks