Set Block Decoding es un Acelerador de Inferencia para Modelos de Lenguaje.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

Autores: Itai Gat, Heli Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, Yaron Lipman

cs.AI

Resumen

Los modelos de lenguaje de predicción de siguiente token autoregresivos ofrecen capacidades potentes, pero enfrentan desafíos significativos en su implementación práctica debido a los altos costos computacionales y de memoria durante la inferencia, particularmente en la etapa de decodificación. Introducimos Set Block Decoding (SBD), un paradigma simple y flexible que acelera la generación al integrar la predicción estándar de siguiente token (NTP) y la predicción de tokens enmascarados (MATP) dentro de una única arquitectura. SBD permite al modelo muestrear múltiples tokens futuros, no necesariamente consecutivos, en paralelo, una distinción clave respecto a métodos de aceleración previos. Esta flexibilidad permite el uso de solucionadores avanzados de la literatura de difusión discreta, ofreciendo aceleraciones significativas sin sacrificar precisión. SBD no requiere cambios arquitectónicos ni hiperparámetros adicionales de entrenamiento, mantiene la compatibilidad con el almacenamiento exacto de KV-caching y puede implementarse mediante el ajuste fino de modelos existentes de predicción de siguiente token. Al ajustar Llama-3.1 8B y Qwen-3 8B, demostramos que SBD permite una reducción de 3 a 5 veces en el número de pasos hacia adelante requeridos para la generación, logrando el mismo rendimiento que el entrenamiento equivalente con NTP.

English

Autoregressive next token prediction language models offer powerful capabilities but face significant challenges in practical deployment due to the high computational and memory costs of inference, particularly during the decoding stage. We introduce Set Block Decoding (SBD), a simple and flexible paradigm that accelerates generation by integrating standard next token prediction (NTP) and masked token prediction (MATP) within a single architecture. SBD allows the model to sample multiple, not necessarily consecutive, future tokens in parallel, a key distinction from previous acceleration methods. This flexibility allows the use of advanced solvers from the discrete diffusion literature, offering significant speedups without sacrificing accuracy. SBD requires no architectural changes or extra training hyperparameters, maintains compatibility with exact KV-caching, and can be implemented by fine-tuning existing next token prediction models. By fine-tuning Llama-3.1 8B and Qwen-3 8B, we demonstrate that SBD enables a 3-5x reduction in the number of forward passes required for generation while achieving same performance as equivalent NTP training.