JetSpec突破推测解码速度瓶颈,实现更高吞吐与更低延迟。
JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

- 采用头级并行草稿机制,兼顾效率与分支因果一致性。
- 在数学与对话任务中最高提速9.64倍,显著优于现有方法。
- 适合追求高并发推理性能的模型部署场景。
推测解码(SD)通过并行起草多个词元并验证来加速自回归大语言模型,但存在扩展瓶颈:只有当接受率高且起草开销低时,增加草稿预算才能提升速度。此前基于头的SD方法面临因果性与效率的矛盾。自回归草稿器生成路径相关候选,适用于树形推测解码,但成本随树深度上升;双向块扩散草稿器可一次性生成所有位置,但其不区分分支的边缘分布会导致各分支虽各自合理却相互矛盾,浪费预算并降低接受率。我们提出JetSpec,一种基于头的SD框架,将单向起草效率与分支级因果条件结合。JetSpec在冻结目标模型的融合隐藏状态上训练因果并行草稿头,生成与目标模型自回归分解对齐的候选树。这使得更大草稿预算可转化为更长的有效前缀和更高的端到端加速比。在密集型与MoE Qwen3模型上,跨数学、编程与对话基准测试中,JetSpec持续优于双向头与树形SD基线。在H100 GPU上,对MATH-500实现最高9.64倍加速,对开放对话任务达4.58倍加速,并在真实服务负载下通过vLLM集成进一步降低延迟。代码与模型已开源。
原文摘要 · Abstract (English)
Speculative decoding (SD) accelerates autoregressive Large Language Models (LLMs) by drafting multiple tokens and verifying them in parallel, but it faces a scaling limitation: increasing the draft budget improves speed only when acceptance remains high and drafting overhead stays low. This ceiling has been difficult to break because prior head-based SD methods face a causality-efficiency dilemma. Autoregressive drafters produce path-conditioned candidates that are effective for tree speculative decoding with higher acceptance length, but their drafting cost grows with tree depth. Bidirectional block-diffusion drafters generate all positions in one pass, but their branch-agnostic marginals can form individually plausible yet mutually inconsistent trees, wasting budget and reducing acceptance. We propose JetSpec, a head-based SD framework that combines one-forward drafting efficiency with branch-wise causal conditioning. JetSpec trains a causal parallel draft head over fused hidden states from the frozen target model, producing candidate trees whose scores align with the target model's autoregressive factorization. This enables JetSpec to convert larger draft budgets into longer accepted prefixes and higher end-to-end speedup. Across math, coding, and chat benchmarks on dense and MoE Qwen3 models, JetSpec consistently outperforms bidirectional-head and tree-based SD baselines. On H100 GPUs, JetSpec achieves up to 9.64x speedup on MATH-500 and 4.58x on open-ended conversational workloads, with further latency gains demonstrated through vLLM integration under realistic serving loads. Our code and models are available at https://github.com/hao-ai-lab/JetSpec.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。