JetSpec:通过并行树绘制技术突破投机性解码的扩展限制
JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting
摘要
推测解码技术通过同时生成并验证多个令牌来加速自回归大语言模型的处理过程,但存在一种限制:只有当接受率较高且生成成本较低时,增加生成预算才能提升处理速度。这种限制难以突破,因为传统的基于头部的推测解码方法面临着因果性与效率之间的矛盾问题。自回归生成机制能够产生符合路径条件的候选树结构,从而在更长的输入长度下保持较高的接受率,但树的深度越大,生成成本也越高。双向块扩散生成机制则可以在一次操作中生成所有位置的信息,但不同分支的边际信息可能会形成相互矛盾的树结构,导致资源浪费和接受率下降。我们提出了JetSpec这一基于头部的推测解码框架,它结合了单向生成效率与分支级因果条件化机制。JetSpec通过从冻结的目标模型中提取隐藏状态来训练因果并行生成头,从而生成与目标模型自回归结构相一致的候选树结构。这使得JetSpec能够利用更大的生成预算获得更长的接受前缀,并实现更高的端到端处理速度。在数学、编程和对话等领域针对密集模型和MoE Qwen3模型的测试中,JetSpec始终优于其他双向头部和树结构推测解码方法。在H100 GPU上,JetSpec在MATH-500任务中实现了9.64倍的处理速度提升,在开放式对话任务中则实现了4.58倍的提升;通过vLLM集成后,在真实的服务环境中还能进一步降低延迟。我们的代码和模型可访问于https://github.com/hao-ai-lab/JetSpec。
English Abstract
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.