JetSpec: Rompendo o Teto de Escalabilidade da Decodificação Especulativa com Rascunho Paralelo em Árvore

Resumo

A decodificação especulativa (SD) acelera Modelos de Linguagem de Grande Porte (LLMs) autorregressivos ao rascunhar múltiplos tokens e verificá-los em paralelo, mas enfrenta uma limitação de escalabilidade: aumentar o orçamento de rascunho melhora a velocidade apenas quando a aceitação permanece alta e a sobrecarga de rascunho se mantém baixa. Esse teto tem sido difícil de romper porque métodos SD baseados em cabeça anteriores enfrentam um dilema entre causalidade e eficiência. Rascunhadores autorregressivos produzem candidatos condicionados por caminho que são eficazes para a decodificação especulativa em árvore com maior comprimento de aceitação, mas seu custo de rascunho cresce com a profundidade da árvore. Rascunhadores de difusão em bloco bidirecionais geram todas as posições em uma única passagem, mas suas marginais agnósticas a ramos podem formar árvores individualmente plausíveis, porém mutuamente inconsistentes, desperdiçando orçamento e reduzindo a aceitação. Propomos o JetSpec, uma estrutura SD baseada em cabeça que combina a eficiência de rascunho em uma passagem com o condicionamento causal entre ramos. O JetSpec treina uma cabeça de rascunho paralela causal sobre estados ocultos fundidos do modelo alvo congelado, produzindo árvores candidatas cujas pontuações se alinham com a fatoração autorregressiva do modelo alvo. Isso permite que o JetSpec converta orçamentos de rascunho maiores em prefixos aceitos mais longos e maior aceleração ponta a ponta. Em benchmarks de matemática, codificação e bate-papo nos modelos Qwen3 densos e MoE, o JetSpec supera consistentemente as linhas de base SD baseadas em cabeça bidirecional e em árvore. Em GPUs H100, o JetSpec atinge até 9,64x de aceleração no MATH-500 e 4,58x em cargas de trabalho conversacionais de livre andamento, com ganhos adicionais de latência demonstrados por meio da integração com vLLM sob cargas de serviço realistas. Nosso código e modelos estão disponíveis em 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.