JetSpec: Rompiendo el techo de escalado de la decodificación especulativa con drafting paralelo de árboles

Resumen

La decodificación especulativa (DS) acelera los grandes modelos de lenguaje (LLMs) autorregresivos al redactar múltiples tokens y verificarlos en paralelo, pero enfrenta una limitación de escalado: aumentar el presupuesto de borrador mejora la velocidad solo cuando la aceptación sigue siendo alta y la sobrecarga de borrado se mantiene baja. Este techo ha sido difícil de superar porque los métodos previos de DS basados en cabezales enfrentan un dilema de causalidad-eficiencia. Los borradores autorregresivos producen candidatos condicionados por trayectoria que son efectivos para la decodificación especulativa en árbol con mayor longitud de aceptación, pero su costo de borrado crece con la profundidad del árbol. Los borradores de difusión de bloques bidireccionales generan todas las posiciones en una sola pasada, pero sus marginales independientes de ramas pueden formar árboles individualmente plausibles pero mutuamente inconsistentes, desperdiciando presupuesto y reduciendo la aceptación. Proponemos JetSpec, un marco de DS basado en cabezales que combina la eficiencia de borrado en una sola pasada hacia adelante con el condicionamiento causal por rama. JetSpec entrena un cabezal de borrador paralelo causal sobre estados ocultos fusionados del modelo objetivo congelado, produciendo árboles candidatos cuyas puntuaciones se alinean con la factorización autorregresiva del modelo objetivo. Esto permite que JetSpec convierta presupuestos de borrador más grandes en prefijos aceptados más largos y una mayor aceleración de extremo a extremo. En puntos de referencia de matemáticas, codificación y chat en modelos Qwen3 densos y MoE, JetSpec supera consistentemente a las líneas base de DS basadas en cabezales bidireccionales y en árbol. En GPUs H100, JetSpec logra una aceleración de hasta 9.64x en MATH-500 y 4.58x en cargas de trabajo conversacionales de final abierto, con mayores ganancias de latencia demostradas mediante la integración con vLLM bajo cargas de servicio realistas. Nuestro código y modelos están disponibles en 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.