JetSpec: Het doorbreken van het schaalplafond van speculatieve decodering met parallelle boomdrafting

Samenvatting

Speculatieve decodering (SD) versnelt autoregressieve Grote Taalmodellen (LLM's) door meerdere tokens te genereren en deze parallel te verifiëren, maar kent een schaalbaarheidsbeperking: het verhogen van het conceptbudget versnelt alleen wanneer de acceptatie hoog blijft en de conceptoverhead laag blijft. Dit plafond is moeilijk te doorbreken omdat eerdere op koppen gebaseerde SD-methoden een causaliteit-efficiëntiedilemma kennen. Autoregressieve conceptgeneratoren produceren pad-geconditioneerde kandidaten die effectief zijn voor boom-speculatieve decodering met een hogere acceptatielengte, maar hun conceptkosten stijgen met de boomdiepte. Bidirectionele blokdiffusie-conceptgeneratoren genereren alle posities in één doorgang, maar hun tak-agnostische marginalen kunnen afzonderlijk plausibele maar onderling inconsistente bomen vormen, wat budget verspilt en de acceptatie vermindert. Wij stellen JetSpec voor, een op koppen gebaseerd SD-raamwerk dat de efficiëntie van eenmalige voorwaartse conceptgeneratie combineert met takgewijze causale conditionering. JetSpec traint een causale parallelle conceptkop over gefuseerde verborgen toestanden van het bevroren doelmodel, waardoor kandidaatbomen worden geproduceerd waarvan de scores overeenkomen met de autoregressieve factorisatie van het doelmodel. Hierdoor kan JetSpec grotere conceptbudgetten omzetten in langere geaccepteerde voorvoegsels en een hogere end-to-end-versnelling. Op wiskunde-, codeer- en chatbenchmarks met dichte en MoE Qwen3-modellen presteert JetSpec consequent beter dan bidirectionele kop- en boomgebaseerde SD-basislijnen. Op H100-GPU's behaalt JetSpec tot 9,64x versnelling op MATH-500 en 4,58x op open conversatiewerklasten, met verdere latentiewinst aangetoond via vLLM-integratie onder realistische serverbelastingen. Onze code en modellen zijn beschikbaar op 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.