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 利用相邻 decode 步之间 Top-K 的时间相关性，在 Blackwell 上为稀疏注意力 decode 设计数据感知的精确 Top-K 算法，相对 radix-select kernel 平均加速 1.88×，端到端 TPOT 最高提升 7.52%。

核心观点

  • 在长上下文 LLM 服务中，即使 indexer 和 attention kernel 已高度优化，每次 decode query 的 Top-K 选择仍是显著延迟瓶颈。
  • 相邻 decode 步之间的 Top-K 存在强时间相关性：offset+1 shift 后约 60% 的 Top-K token 保持不变，仅约 40% 变化。
  • 该现象与 DeepSeek Sparse Attention (DSA) indexer 分数的 Toeplitz / RoPE 结构相关。
  • 基于此先验可将精确 Top-K 转化为"猜测-验证-精修"问题，在保证 bit-exact 输出的同时大幅降低开销。

图 2 图 3

方法

GVR（Guess-Verify-Refine）面向 NVIDIA Blackwell 的数据感知精确 Top-K：

  • Guess：用上一步的 Top-K 作为预测信号，并计算 pre-indexed 统计量。
  • Verify：通过 secant-style counting 在 1-2 次 global pass 内收敛到合法阈值，用 ballot-free collector 验证候选。
  • Refine：在 shared memory 中完成精确选择，保证与原算法 bit-exact 一致。
  • 替换原 decode 阶段按序列长度派发 insert sort / radix sort / multi-CTA 的 Top-K kernel。

图 1

实验

  • 集成到 TensorRT-LLM 的 DSA stack，模型为 DeepSeek-V3.2。
  • 工作负载来自 SWE-bench 派生的 LongSeqTasks（如 swe_bench_64k.jsonl）。
  • 基线：生产环境的 radix-select kernel。
  • 部署配置：TEP8 min-latency，测试 100K 上下文，并涵盖 speculative decoding。
  • 指标：单算子加速比、端到端 TPOT、Top-K 输出位级一致性。

结果

  • 单算子相对 radix-select 平均加速 1.88×，单层单步最高 2.42×。
  • TEP8 min-latency 下 100K 上下文端到端 TPOT 最高提升 7.52%；上下文越长收益越大。
  • speculative decoding 下收益变小但仍为正。
  • 保持与原 Top-K 完全一致的 bit-exact 输出。

为什么重要

对 long-context LLM serving，Top-K 已成为 decode 延迟新瓶颈。GVR 证明可用数据感知 + 时间相关性在 Blackwell 上进一步压榨，且不牺牲精度；对部署 DSA 类稀疏注意力的推理栈是直接可落地的优化点。

与已有工作的关系

  • 依托 DeepSeek-V3.2 的 DSA 稀疏注意力与 RoPE 结构。
  • 替换 TensorRT-LLM 中现有的 radix-select / insert sort / multi-CTA Top-K 派发策略。
  • 与 speculative decoding 推理加速方向正交互补。

尚未回答的问题

  • 是否能推广到非 DSA、非 Toeplitz/RoPE 结构的稀疏注意力 decoder？
  • 在 Hopper 等非 Blackwell 架构上的适配与收益？
  • prefill 阶段或 batch 较大场景下时间相关性是否仍然成立？
  • 与其他 indexer / KV 压缩方案协同的端到端收益边界？

原始摘要（中文翻译）

稀疏注意力 decoder 依赖精确 Top-K 选择来为每个 query token 选取最重要的 key-value 条目。在长上下文 LLM serving 中，该 Top-K 阶段每次 decode query 都要运行一次，即使 indexer 与 attention kernel 已经被高度优化，它仍会成为显著的延迟瓶颈。我们提出 Guess-Verify-Refine (GVR)，一种面向 NVIDIA Blackwell 的数据感知精确 Top-K 算法，用于稀疏注意力 decoding。GVR 利用相邻 decode 步之间的时间相关性：它使用上一步的 Top-K 作为预测信号，计算 pre-indexed 统计量，通过 secant-style counting 在 1-2 次全局遍历内收敛到合法阈值，用 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 上下文下端到端 TPOT 最高提升 7.52%，上下文越长收益越大，在 speculative decoding 下收益较小但仍为正。尽管当前实现并在 Blackwell 上基于 TensorRT-LLM DSA 栈验证，相同原理可能推广到其他 decode 阶段 Top-K 表现出时间稳定性的稀疏注意力 decoder。