JetSpec: Durchbrechen der Skalierungsobergrenze des spekulativen Decodierens mit parallelem Baum-Drafting

Zusammenfassung

Spekulative Dekodierung (SD) beschleunigt autoregressive große Sprachmodelle (LLMs), indem mehrere Token entworfen und parallel verifiziert werden, stößt jedoch auf eine Skalierungsgrenze: Eine Erhöhung des Entwurfsbudgets verbessert die Geschwindigkeit nur dann, wenn die Akzeptanz hoch und der Entwurfsaufwand gering bleibt. Diese Grenze war schwer zu durchbrechen, da frühere kopfbasierte SD-Verfahren mit einem Kausalitäts-Effizienz-Dilemma konfrontiert sind. Autoregressive Entwerfer erzeugen pfadbedingte Kandidaten, die für die baumspekulative Dekodierung mit höherer Akzeptanzlänge effektiv sind, aber ihre Entwurfskosten wachsen mit der Baumtiefe. Bidirektionale Blockdiffusions-Entwerfer generieren alle Positionen in einem Durchgang, aber ihre zweigagnostischen Randverteilungen können individuell plausible, aber gegenseitig inkonsistente Bäume bilden, was Budget verschwendet und die Akzeptanz verringert. Wir schlagen JetSpec vor, ein kopf-basiertes SD-Framework, das die Effizienz eines einmaligen Vorwärtsentwurfs mit einer zweigweisen kausalen Bedingung kombiniert. JetSpec trainiert einen kausalen parallelen Entwurfskopf über fusionierte versteckte Zustände des eingefrorenen Zielmodells, wodurch Kandidatenbäume erzeugt werden, deren Bewertungen mit der autoregressiven Faktorisierung des Zielmodells übereinstimmen. Dies ermöglicht JetSpec, größere Entwurfsbudgets in längere akzeptierte Präfixe und eine höhere End-to-End-Beschleunigung umzuwandeln. In Benchmarks zu Mathematik, Programmierung und Chat auf dichten und MoE-Qwen3-Modellen übertrifft JetSpec durchgängig bidirektionale Kopf- und baumbasierte SD-Baselines. Auf H100-GPUs erreicht JetSpec eine bis zu 9,64-fache Beschleunigung auf MATH-500 und eine 4,58-fache auf offenen Konversationsworkloads, mit weiteren Latenzgewinnen durch vLLM-Integration unter realistischen Serverlasten. Unser Code und unsere Modelle sind verfügbar unter 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.