JetSpec:基于因果并行草稿头的推测解码框架
阅读原文· arxiv.orgJetSpec 是一种头部驱动推测解码框架,通过在冻结目标模型的融合隐藏状态上训练因果并行草稿头,生成与自回归因子分解对齐的候选树,从而将更大草稿预算转化为更长接受前缀和更高端到端加速。在密集和 MoE Qwen3 模型的数学、编码及聊天基准测试中,JetSpec 一致优于双向头和树形基线。在 H100 GPU 上,MATH-500 达 9.64 倍加速,开放对话达 4.58 倍;经 vLLM 集成在现实服务负载下进一步降低延迟。
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.