ICML'24 | Quest：不用扫描全部 KV Cache，Query-aware Sparsity 如何加速长上下文推理

August 13, 2026

2026   ·   LLM   KV Cache   Attention   长上下文   论文解读  

ICML’24 | Quest：不用扫描全部 KV Cache，Query-aware Sparsity 如何加速长上下文推理

原文：Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

1. 前言

长上下文 LLM 的一个直觉误区是：上下文窗口能做到 128K、1M token，模型就会“均匀地使用”这么长的上下文。实际生成时，当前问题往往只需要历史中的一小部分信息。

全注意力的问题在于，它每生成一个 token，都要让当前 Query 和整个历史 KV Cache 做匹配。历史越长，读取和计算越重；而真正有用的 token 可能只占很小一部分。

Quest 的问题非常直接：

如果每个 Query 只需要少数几个 KV page，能不能先找出这些 page，再只对它们做 attention？

论文给出的答案是可以，而且关键不在于固定保留“最近 token”或“最早 token”，而在于让当前 Query 决定当前该看哪里。这就是 Query-aware sparsity。

2. 背景：Attention 到底在算什么

2.1 一次 attention 就像一次相关性检索

对当前 token 来说，模型会产生一个 Query $q$；历史 token 各自有 Key $k_i$ 和 Value $v_i$。注意力大致是：

\[Attention(q,K,V)=softmax(qK^T)V\]

第一步 $qK^T$ 是相关性打分：当前问题和哪些历史 token 更相关？第二步用这些分数加权 Value，把信息取回来。

因此，Attention 也可以看成一个“以 Query 为条件的检索器”。问题是传统实现会把所有 Key 都扫一遍，哪怕最后只有少数 token 的分数真正重要。

2.2 KV Cache 为什么拖慢 Decode

Prefill 阶段可以并行处理 prompt；Decode 阶段每次只生成一个 token，却要重复读取不断增长的历史 K/V。这个阶段经常不是算力不够，而是内存带宽不够：GPU 的计算单元在等待 KV 从显存中流过来。

把历史 KV 切成 page/block 后，系统可以把 page 当成检索单元。Quest 的目标不是预测每一个 token 的精确 attention 分数，而是先判断哪些 page 值得保留，然后对入选 page 做精确 attention。

2.3 为什么常见的稀疏策略不够

常见方法各有合理性，但都带着固定偏好：

  • Sliding window 只看最近 token，长距离引用会丢失；
  • StreamingLLM 保留开头的 attention sink 和最近窗口，但中间的重要事实未必在这里；
  • H2O/TOVA 等方法根据历史累计分数或特定规则淘汰，未必适合下一个 Query；
  • 固定比例稀疏虽然简单，但不同 Query 需要的预算和位置可能完全不同。

论文 Figure 2 用一个很小的例子说明了这一点：同一段文本里，当前 Query 变成不同 token 后，真正重要的历史位置也会变。稀疏性必须 query-aware，而不是只看位置或历史平均分数。

Query 改变时的注意力位置

3. Quest 的核心：先找重要 page，再做精确 attention

3.1 为什么要用 page，而不是逐 token 选择

逐 token 选择看似精细，但会产生大量索引和不连续访存。Quest 把连续的历史 token 分成 page，每个 page 内有一段 Key/Value。这样做有三个好处：

  1. page 是 GPU 友好的连续内存块；
  2. 选择的索引数量从 token 数降到 page 数；
  3. page 粒度可以在选择开销和稀疏收益之间调节。

page 越大，索引开销越小，但一个 page 里混入无关 token 的概率越高；page 越小，选择更精细，但 metadata 和 gather 开销会上升。

3.2 用 page 的 Key 统计量近似 page 重要性

对于一个 page，Quest 不必先把 page 内所有 Key 都拿出来做完整 attention。它对 page 内 Key 做 min/max pooling，得到一个很小的摘要。

对当前 Query $q$，摘要可以快速给出该 page 的相关性上界或近似分数。然后把 page 按分数排序，选择 Top-K page。最后才读取入选 page 的完整 KV，执行精确 attention。

可以把整个过程想成搜索引擎：先用索引筛掉明显不相关的文档，再对候选文档逐字阅读。索引阶段允许是近似的，但候选集必须尽量不要漏掉真正相关的 page。

Quest 的 page 选择流程

3.3 Top-K page 预测不是另一个完整 attention

Quest 的关键优化是让 page selection 足够便宜。它利用 page 级别的 Key 摘要完成粗排，再把少量候选 page 的完整 KV 交给 attention kernel。

最终计算近似为：

\[Attention(q,K,V) \approx Attention(q,K_{TopK},V_{TopK})\]

这里的近似成立依赖于一个条件：Top-K page 覆盖真正有贡献的历史位置。论文的实验显示，在很多长依赖任务中，Query-aware 的选择可以用很小的 KV budget 保持接近 full attention 的效果。

4. 系统实现：算法和 kernel 必须一起设计

如果只在 Python 层做 page 索引，然后把挑出来的 KV 拼接成一个新张量，索引和内存拷贝很可能吃掉稀疏带来的收益。因此 Quest 还需要专门的 GPU 实现。

一个高效实现通常包含三步：

  1. 对 page 摘要计算与 Query 的相关性；
  2. 用并行 reduction 找出 Top-K page；
  3. 在不产生大规模中间张量的情况下，对入选 page 做 attention。

这也是长上下文优化经常被低估的地方：理论上只算 1/16 的 KV，不代表 wall-clock 就是 1/16。gather、排序、kernel launch 和非连续访存都要付钱。Quest 的贡献是把 query-aware page selection 做成了可实际运行的 self-attention 实现，而不是只提出一个离线剪枝指标。

5. 准确率：稀疏预算不是越小越好

Quest 的稀疏 attention 通常会给每层一个 KV budget，同时保留模型前两层的 full cache。原因很直观：早期层的表示和全局信息传播更敏感，把它们完整保留可以降低误差；后面的层再使用 query-aware sparsity。

实验中，Quest 在不同 budget 下都明显优于 H2O、TOVA 和 StreamingLLM。一个很直观的结果是，在较小 page budget 下，Quest 仍能保持较高的检索准确率，而其他方法很快掉到接近零的命中水平。

不同 KV budget 下的检索准确率

对长文本中的事实问答，这种差异尤其重要。固定保留最近窗口的方法可能把答案所在的中间位置丢掉；Quest 每一步重新根据 Query 选择，因此能把远距离但当前相关的信息找回来。

6. 速度：真正要看端到端，而不是只看稀疏比例

论文报告 Quest 最高可以获得 7.03 倍的 self-attention 加速，并带来最高 2.23 倍的 decode 阶段端到端延迟下降。与已有稀疏方法在相同准确率目标下比较，Quest 的 self-attention 延迟最高降低 4.5 倍。

不同上下文长度下的 attention 延迟

这几个数字要正确理解：7.03 倍是 attention 子模块的最高加速，2.23 倍是整个 decode 阶段的最高延迟收益，二者不是同一个指标。端到端收益更小是正常的，因为每一步还包含 QKV projection、FFN、采样、调度和其他固定开销。

在相同准确率约束下，Quest 相比基线在 GovReport 上最高有 3.82 倍推理速度提升，在 TriviaQA 上最高有 4.54 倍提升。换句话说，Quest 的优势不只是“可以更激进地丢 KV”，而是同样的效果下需要更少的 KV 访问。

相同准确率约束下的效率比较

7. 一个具体例子：为什么 Query-aware 比最近窗口更合理

假设上下文里有这样一段内容：

A is B. C is D. A is …

当当前 Query 需要回答和 A 相关的问题，答案可能在很早的 B 附近；当 Query 需要回答和 D 相关的问题，重要位置又会变成 D 所在 page。最近窗口策略只知道“最后一段最重要”，但 Quest 会用当前 Query 重新计算 page 的相关性。

这也是我觉得 Quest 最值得记住的抽象：长上下文并不等于每次都需要长计算。模型的上下文窗口可以很长，但每个 decode step 的有效感受野可能是稀疏的，而且这个稀疏模式由问题决定。

8. 局限性和工程代价

Quest 不是无条件替代 full attention。

首先，page 选择本身需要计算。如果 page 数很少，或者上下文很短，选择开销可能抵消收益。其次，page 太大时会引入无关 token，page 太小时又会增加索引和访存开销。再次，某些任务确实需要聚合大量上下文，稀疏预算太小会损失信息。

此外，模型和硬件实现必须配合。不同 GPU、不同 batch size、不同 KV layout 下，Top-K、gather 和 attention kernel 的开销不同。论文展示的是一套有效的 query-aware 实现，但部署时仍然需要按真实 workload 调 page size 和 budget。

9. 总结

Quest 的一句话版本是：把 Attention 当成 Query 驱动的 KV 检索，先在 page 粒度找 Top-K，再对这些 page 做精确 attention。

它解决了长上下文 decode 中“每一步都扫描完整 KV Cache”的浪费，并且把准确率、page 粒度和 GPU kernel 放在一起考虑。后来的 KV 分层和预取工作，正是在这个方向上继续追问：既然只需要少数 KV page，这些 page 能不能还没用到时就提前搬到 GPU？

原文链接：Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

Flag Counter