JetSpec : briser le plafond de mise à l'échelle du décodage spéculatif avec la rédaction d'arbres parallèle

Résumé

Le décodage spéculatif (DS) accélère les grands modèles de langage (LLM) autorégressifs en générant plusieurs jetons sous forme de brouillons et en les vérifiant en parallèle, mais il se heurte à une limitation de passage à l'échelle : augmenter le budget de brouillon n'améliore la vitesse que si le taux d'acceptation reste élevé et la surcharge de rédaction faible. Ce plafond a été difficile à briser car les méthodes de DS antérieures basées sur des têtes sont confrontées à un dilemme causalité-efficacité. Les générateurs de brouillons autorégressifs produisent des candidats conditionnés par le chemin, efficaces pour le décodage spéculatif arborescent avec une longueur d'acceptation plus élevée, mais leur coût de rédaction croît avec la profondeur de l'arbre. Les générateurs de brouillons par diffusion de blocs bidirectionnelle génèrent toutes les positions en un seul passage, mais leurs marginales indépendantes des branches peuvent former des arbres individuellement plausibles mais mutuellement inconsistants, gaspillant le budget et réduisant l'acceptation. Nous proposons JetSpec, un cadre de DS basé sur des têtes qui combine l'efficacité de rédaction en un seul passage avant avec un conditionnement causal par branche. JetSpec entraîne une tête de brouillon parallèle causale sur les états cachés fusionnés du modèle cible gelé, produisant des arbres candidats dont les scores s'alignent sur la factorisation autorégressive du modèle cible. Cela permet à JetSpec de convertir des budgets de brouillon plus importants en préfixes acceptés plus longs et en une accélération de bout en bout plus élevée. Sur des benchmarks de mathématiques, de codage et de conversation, avec les modèles Qwen3 denses et MoE, JetSpec surpasse systématiquement les lignes de base de DS à têtes bidirectionnelles et arborescentes. Sur des GPU H100, JetSpec atteint une accélération allant jusqu'à 9,64× sur MATH-500 et 4,58× sur des charges conversationnelles ouvertes, avec des gains de latence supplémentaires démontrés via l'intégration vLLM sous des charges de service réalistes. Notre code et nos modèles sont disponibles à l'adresse 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.