arXiv:2606.00144cs.LGcs.AI2026-06被引 1

让稀疏推理的草案模型更懂接受率,提升长文本生成速度

BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

论文配图:BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding
图 1 · 摘自论文原文
  • 训练时让草案模型接触多种稀疏缓存配置,学会适配不同内存预算
  • 在4K~16K上下文长度下,端到端速度提升最高达6.55倍,接受率稳定
  • 无需额外推理组件,适合资源受限的长文本生成场景

推测解码通过草案模型并行提出多个词元,由验证器统一校验以加速自回归生成。在资源受限场景下,草案模型使用稀疏键值缓存以控制峰值GPU内存和端到端延迟,而验证器保留完整缓存。中长上下文推理(4K–16K)在实际应用中常见。然而,朴素的稀疏/全量推测解码随上下文增长出现稀疏与全量不匹配,导致接受率迅速下降。本文提出BudgetDraft,一种面向中长上下文推理的多视角稀疏训练方法。草案模型在训练时暴露于多种采样的KV缓存预算,并学习将每个稀疏视图对齐至统一的全缓存教师目标。BudgetDraft结合全缓存分支的接受率感知损失与稀疏缓存分支的多视角损失,训练出一个单一的、对预算鲁棒的草案模型,无需额外推理组件即可恢复各稀疏度下的接受率。在PG-19、LongBench和LWM上的实验表明,该方法在4K、8K和16K上下文长度下相比自回归推理分别实现最高6.55倍、4.46倍和2.10倍的端到端加速,同时保持内存友好性。

原文摘要 · Abstract (English)

Speculative decoding speeds up autoregressive decoding by using a drafter to propose multiple tokens that a verifier validates in parallel. In resource-constrained deployments, the drafter uses a sparse KV cache to limit peak GPU memory and end-to-end latency under a fixed KV budget, while the verifier keeps a full KV cache. Mid-to-long context inference (4K--16K context length) is common in real applications. However, naive sparse/full speculative decoding suffers from the sparse/full mismatch as context length grows, causing the acceptance rate to drop quickly. We propose BudgetDraft, a multi-view sparse training method for sparse drafting in mid-to-long inference. The drafter is exposed to multiple sampled KV budgets during training and learns to align each sparse view with one shared full-cache teacher target. BudgetDraft combines an acceptance-aware loss on a full-cache branch with a multi-view loss on a sparse-cache branch, producing a single budget-robust drafter that recovers acceptance across sparsity levels without extra inference-time components. Experimental results on PG-19, LongBench, and LWM show that BudgetDraft achieves up to 6.55x, 4.46x, and 2.10x end-to-end speedup vs AR at 4K, 8K, and 16K context lengths, while keeping the inference pipeline memory-friendly.

推理加速稀疏缓存推测解码

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。