JetSpec：以平行樹狀草稿突破推測解碼的規模天花板

摘要

推测解码（Speculative Decoding, SD）通过草拟多个token并并行验证，加速了自回归大语言模型（LLMs），但面临扩展性瓶颈：只有在接受率保持较高且草拟开销较低时，增加草稿预算才能提升速度。这一上限难以突破，因为此前基于预测头的SD方法面临因果性与效率的困境。自回归草拟器生成路径条件化的候选序列，这些序列在树状推测解码中能实现更高的接受长度，但其草拟成本随树深度增长。双向块扩散草拟器可一次性生成所有位置，但其与分支无关的边缘分布可能产生单个合理但彼此不一致的树状结构，浪费预算并降低接受率。我们提出JetSpec，一种基于预测头的SD框架，将单次前向草拟效率与分支级因果条件相结合。JetSpec在冻结目标模型的融合隐藏状态上训练一个因果并行草拟头，生成的候选树得分与目标模型的自回归分解一致。这使得JetSpec能够将更大的草稿预算转化为更长的已接受前缀和更高的端到端加速。在密集型和MoE Qwen3模型上的数学、编程及聊天基准测试中，JetSpec始终优于双向预测头和树状SD基线。在H100 GPU上，JetSpec在MATH-500上实现了高达9.64倍加速，在开放式对话工作负载上达到4.58倍加速，并通过vLLM集成在实际服务负载下进一步降低了延迟。我们的代码和模型已公开在https://github.com/hao-ai-lab/JetSpec。

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.