Windowed-MTP: Removing the Full-Context Draft-KV Tax at Million-Token Context

📄 arXiv: 2607.21535v1 📥 PDF

作者: Alagappan Valliappan

分类: cs.LG, cs.CL, cs.PF

发布日期: 2026-07-23

备注: 25 pages, 2 figures, 11 tables

💡 一句话要点

提出Windowed-MTP以解决百万标记上下文中的全上下文草稿KV成本问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 多标记预测 滑动窗口 注意力机制 计算效率 自然语言生成

📋 核心要点

  1. 现有的多标记预测草稿头在百万标记上下文中面临高计算成本，影响生成效率。
  2. 论文提出Windowed-MTP，通过滑动窗口机制优化草稿的注意力计算，降低成本。
  3. 在多个架构上实验表明，Windowed-MTP在每步解码中成本降低28%至44%，提升了生成速度。

📝 摘要（中文）

论文提出了一种名为Windowed-MTP的方法，旨在解决在百万标记上下文中，现有多标记预测草稿头的全注意力机制导致的高成本问题。现有方法在每个草稿步骤中对整个KV缓存进行全注意力计算，随着上下文的增加，计算成本呈线性增长，最终影响生成效率。Windowed-MTP通过引入滑动窗口和注意力汇聚机制，仅对草稿的注意力进行优化，保持全注意力验证不变，从而显著降低了每步解码的成本，提升了生成速度，同时不影响生成质量。

🔬 方法详解

问题定义：论文要解决的问题是现有多标记预测草稿头在百万标记上下文中，因全注意力机制导致的高计算成本。现有方法在每个草稿步骤中对整个KV缓存进行全注意力计算，随着上下文的增加，计算成本线性增长，影响生成效率。

核心思路：论文的核心解决思路是引入滑动窗口机制，仅对草稿的注意力进行优化，同时保持全注意力验证不变。这样设计的目的是在不影响生成质量的前提下，显著降低计算成本。

技术框架：整体架构包括一个滑动窗口机制和注意力汇聚模块。滑动窗口负责限制草稿的KV工作集大小，而注意力汇聚模块则确保全注意力验证的完整性。

关键创新：最重要的技术创新点在于通过窗口化草稿的注意力计算，降低了草稿的KV缓存使用，减少了约99%的KV条目，同时保持了生成质量和接受的输出分布。

关键设计：关键设计包括滑动窗口的大小设置和注意力汇聚的实现，确保在不影响接受长度的情况下，优化每步解码的计算成本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Windowed-MTP在多个架构上实现了每步解码成本降低28%至44%的显著提升，且在匹配接受的端到端解码延迟方面也有相应改善，证明了其在高效生成中的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言生成、对话系统和文本摘要等任务，能够显著提升大规模语言模型的生成效率和响应速度。未来，Windowed-MTP可能在实时应用场景中发挥重要作用，推动更高效的AI系统发展。

📄 摘要（原文）

Speculative decoding accelerates autoregressive generation by having a cheap draft propose tokens that a target verifies in parallel. Frontier models increasingly ship a built-in Multi-Token-Prediction (MTP/NEXTN) draft head under the assumption that the draft is negligibly cheap. At million-token context this breaks: an MTP draft head typically runs full attention over the entire KV cache at every draft step, so its read grows linearly with context and comes to dominate the draft cost -- precisely where speculation is most valuable. The effect compounds with draft length (a deep native draft can turn net-negative, slower than no speculation) and sharpens under hybrid/linear-attention targets, where cheaper verification leaves the draft's full-attention read exposed. We apply a StreamingLLM-style sliding window plus attention sink to the draft's attention only (Windowed-MTP), leaving full-attention verification intact. It is training-free, drop-in, and lossless by construction: the full-attention target still decides every accepted token, so windowing changes only which tokens are proposed, never which are accepted. It bounds the draft's KV working set to a constant, dropping ~99% of KV entries at 1M. Across three architecture families (Qwen GDN-MoE 35B/122B and a Mamba2-hybrid NoPE 120B) at 1M context on a single GPU in SGLang, windowing cuts the per-decode-step cost over the shipping native MTP draft by +28% to +44%, an input-invariant margin that widens with context. Since per-token latency is this cost divided by acceptance length, at matched acceptance end-to-end decode latency improves by the same amount, and more where windowing also lifts acceptance, while preserving the target's verified output distribution. Finally, the unread draft KV -- 7.7-11% of total KV at 1M -- is reclaimed via a compact ring buffer at no acceptance or quality cost.