JetSpec: Преодоление потолка масштабирования спекулятивного декодирования с помощью параллельного построения деревьев-черновиков

Аннотация

Спекулятивное декодирование (SD) ускоряет авторегрессионные большие языковые модели (LLM) за счет формирования нескольких токенов-кандидатов (черновиков) и их параллельной верификации, однако оно сталкивается с ограничением масштабируемости: увеличение бюджета черновика повышает скорость только при сохранении высокой вероятности принятия и низких накладных расходах на формирование черновика. Этот потолок было трудно преодолеть, поскольку предыдущие методы SD на основе головы сталкиваются с дилеммой причинности и эффективности. Авторегрессионные черновики порождают обусловленные путями кандидаты, эффективные для деревьев спекулятивного декодирования с большей длиной принятия, но стоимость их формирования растет с глубиной дерева. Двунаправленные блок-диффузионные черновики генерируют все позиции за один проход, однако их независимые от ветвей маргинальные вероятности могут образовывать по отдельности правдоподобные, но взаимно несовместимые деревья, что приводит к потере бюджета и снижению вероятности принятия. Мы предлагаем JetSpec — фреймворк SD на основе головы, который сочетает эффективность формирования черновика за один прямой проход с причинным обусловливанием по ветвям. JetSpec обучает причинную параллельную голову черновика на объединенных скрытых состояниях от замороженной целевой модели, порождая деревья-кандидаты, оценки которых согласуются с авторегрессионной факторизацией целевой модели. Это позволяет JetSpec преобразовывать более крупный бюджет черновика в более длинные принятые префиксы и более высокое сквозное ускорение. На тестах по математике, программированию и диалогам на плотных и MoE моделях Qwen3 JetSpec последовательно превосходит базовые методы SD с двунаправленной головой и деревом. На графических процессорах H100 JetSpec достигает ускорения до 9.64× на MATH-500 и 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.