arXiv: 2604.22312 · PDF

作者: Long Cheng, Ritchie Zhao, Timmy Liu, Mindy Li, Xianjie Qiao, Kefeng Duan, Yu-Jung Chen, Xiaoming Chen, Bita Darvish Rouhani, June Yang

单位: NVIDIA

主分类: cs.DC · 全部: cs.AR, cs.DC, cs.PF

命中关键词: llm, rag, serving, speculative decoding, attention, latency

TL;DR

GVR 是面向 Blackwell GPU 的数据感知精确 Top-K 算法，利用解码步间时间相关性加速 DeepSeek Sparse Attention 的 Top-K 选择，单算子平均提速 1.88×，端到端 TPOT 最多提升 7.52%。

核心观点

  • Top-K 选择在长上下文稀疏注意力解码中是不可忽视的延迟瓶颈。
  • 相邻解码步的 Top-K 存在显著时间相关性，可作为预测信号。
  • 这种时间稳定性与 DSA indexer 分数的 Toeplitz / RoPE 结构直接相关。
  • 提出 Guess-Verify-Refine 范式，在保持 bit-exact 输出的前提下显著加速。

图 3

上图说明：当 query 前进一位时，相对位置整体偏移 +1，约 60% 的 Top-K token 仍然保留，只有约 40% 更换；attention score 的 Toeplitz 结构解释了为何大相对位置的峰值不会单调衰减，这正是 GVR 可利用的时间稳定性来源。

方法

GVR 把精确 Top-K 拆成三步：Guess 用上一步 Top-K 和预索引统计给出阈值预测；Verify 通过 secant 式计数在 1–2 次 global pass 中收敛到合法阈值，并用 ballot-free collector 收集候选；Refine 在 shared memory 中完成精确选择。整套算法替换了 TensorRT-LLM 原有的 dispatch 逻辑。

图 1

上图为改造前的原始 decode-stage Top-K dispatch：invokeIndexerTopKDecode 根据序列长度在 insert / radix / multi-CTA split 三种 kernel 间切换，GVR 正是替换这一部分。

实验

在集成进 TensorRT-LLM 的真实 DeepSeek-V3.2 工作负载上评测，baseline 为生产级 radix-select kernel；端到端部署采用 TEP8 min-latency 配置，并覆盖 100K context、speculative decoding 等场景；使用 SWE-bench 派生的 LongSeqTasks 数据（swe_bench_64k）验证 Top-K 时间重叠率。

图 2

上图测量了 DeepSeek-V3.2 各层相邻解码步的原始 Top-K 命中率，是最严格口径下的时间持久性证据，支持 GVR 的预测假设。

结果

单算子层面：平均 1.88× 加速，单层单步最高 2.42×，输出与 baseline bit-exact 一致。端到端：100K context 下 TPOT 最高提升 7.52%，上下文越长收益越大；speculative decoding 场景下增益较小但仍为正。

为什么重要

为长上下文 LLM serving 中常被忽视的 Top-K 阶段给出可落地的加速方案，适配 Blackwell 与 TensorRT-LLM DSA stack，思路可推广到任何 decode 阶段 Top-K 具备时间稳定性的稀疏注意力解码器。

与已有工作的关系

构建在 DeepSeek Sparse Attention (DSA)、DeepSeek-V3.2 indexer 与 RoPE 的结构性质之上；替代 TensorRT-LLM 中基于 insert / radix / multi-CTA 的传统 Top-K 路径；与稀疏注意力、speculative decoding 等长上下文推理优化方向互补。

尚未回答的问题

  • 在非 DSA、无明显 Toeplitz 结构的稀疏注意力上是否仍有效？
  • 极短上下文或高 churn 场景下的最差情况表现如何？
  • 跨硬件（Hopper、消费级 GPU）和 prefill 阶段的可迁移性？
  • 与 speculative decoding 深度结合时收益受限的原因与优化空间？

原始摘要（中文翻译）

稀疏注意力解码器依赖精确的 Top-K 选择来为每个 query token 挑选最重要的 key-value 条目。在长上下文 LLM serving 中，这个 Top-K 阶段在每次 decode query 时都会运行一次，即便 indexer 与 attention kernel 都已高度优化，它仍会成为一个显著的延迟瓶颈。我们提出 Guess-Verify-Refine (GVR)，一种面向 NVIDIA Blackwell 上稀疏注意力解码的数据感知精确 Top-K 算法。GVR 利用连续 decode 步之间的时间相关性：它把上一步的 Top-K 作为预测信号，计算预索引统计，通过 secant 式计数在 1–2 次 global pass 中收敛到合法阈值，用 ballot-free collector 验证候选，并在 shared memory 中完成精确选择。我们将这种行为与 DeepSeek Sparse Attention (DSA) indexer 分数的 Toeplitz / RoPE 结构联系起来，并在集成进 TensorRT-LLM 的真实 DeepSeek-V3.2 工作负载上验证了设计。GVR 相对生产环境的 radix-select kernel 实现了平均 1.88× 的单算子加速，单层单步最高 2.42×，同时保持 bit-exact 的 Top-K 输出。在受控的 TEP8 min-latency 部署下，它在 100K context 上最多将端到端 TPOT 提升 7.52%，上下文越长增益越大，而在 speculative decoding 下增益较小但仍为正。尽管当前实现与验证是在 Blackwell 上的 TensorRT-LLM DSA stack 中完成的，同样的原理可能扩展到任何 decode 阶段 Top-K 具备时间稳定性的稀疏注意力解码器。