Decodificación Especulativa por Lotes Hecha Correctamente

Batch Speculative Decoding Done Right

October 26, 2025

Autores: Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang

cs.AI

Resumen

El descodificado especulativo acelera la inferencia de LLM utilizando un modelo borrador pequeño para proponer múltiples tokens que un modelo objetivo verifica en paralelo. Extender esta idea a lotes es esencial para el servicio en producción, pero introduce el problema del tensor irregular: las secuencias en el mismo lote aceptan diferentes cantidades de tokens borradores, rompiendo la alineación derecha y corrompiendo los ID de posición, las máscaras de atención y el estado de la caché KV. Demostramos que varias implementaciones existentes de lotes violan la equivalencia de salida, el requisito fundamental de que el descodificado especulativo debe producir secuencias de tokens idénticas a la generación autoregresiva estándar. Estas violaciones ocurren precisamente debido al manejo inadecuado del problema del tensor irregular. En respuesta, (1) caracterizamos los requisitos de sincronización que garantizan la corrección, (2) presentamos un descodificado especulativo por lotes centrado en la corrección, EQSPEC, que expone que la realineación consume el 40% de la sobrecarga, y (3) introducimos EXSPEC, que mantiene un grupo deslizante de secuencias y forma grupos de misma longitud dinámicamente, para reducir la sobrecarga de realineación preservando las aceleraciones especulativas por secuencia. En el conjunto de datos SpecBench, a través de pares objetivo/borrador Vicuna-7B/68M, Qwen3-8B/0.6B y GLM-4-9B/0.6B, nuestro enfoque logra una mejora de rendimiento de hasta 3 veces con un tamaño de lote de 8 en comparación con un tamaño de lote de 1, con un escalado eficiente hasta tamaño de lote 8, manteniendo un 95% de equivalencia de salida. Nuestro método no requiere núcleos personalizados y se integra limpiamente con las pilas de inferencia existentes. Nuestro código está disponible en https://github.com/eBay/spec_dec.

English

Speculative decoding speeds up LLM inference by using a small draft model to propose multiple tokens that a target model verifies in parallel. Extending this idea to batches is essential for production serving, but it introduces the ragged tensor problem: sequences in the same batch accept different numbers of draft tokens, breaking right-alignment and corrupting position IDs, attention masks, and KV-cache state. We show that several existing batch implementations violate output equivalence-the fundamental requirement that speculative decoding must produce identical token sequences to standard autoregressive generation. These violations occur precisely due to improper handling of the ragged tensor problem. In response, we (1) characterize the synchronization requirements that guarantee correctness, (2) present a correctness-first batch speculative decoding EQSPEC that exposes realignment as consuming 40% of overhead, and (3) introduce EXSPEC, which maintains a sliding pool of sequences and dynamically forms same-length groups, to reduce the realignment overhead while preserving per-sequence speculative speedups. On the SpecBench dataset, across Vicuna-7B/68M, Qwen3-8B/0.6B, and GLM-4-9B/0.6B target/draft pairs, our approach achieves up to 3times throughput improvement at batch size 8 compared to batch size 1, with efficient scaling through batch size 8, while maintaining 95% output equivalence. Our method requires no custom kernels and integrates cleanly with existing inference stacks. Our code is available at https://github.com/eBay/spec_dec.