Set Block Decoding est un accélérateur d'inférence pour modèles de langage.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

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

cs.AI

papers.abstract

Les modèles de langage à prédiction de tokens suivants de type autoregressif offrent des capacités puissantes mais rencontrent des défis importants dans leur déploiement pratique en raison des coûts élevés en calcul et en mémoire lors de l'inférence, en particulier pendant l'étape de décodage. Nous introduisons le Set Block Decoding (SBD), un paradigme simple et flexible qui accélère la génération en intégrant la prédiction standard de tokens suivants (NTP) et la prédiction de tokens masqués (MATP) au sein d'une seule architecture. Le SBD permet au modèle d'échantillonner plusieurs tokens futurs, pas nécessairement consécutifs, en parallèle, une distinction clé par rapport aux méthodes d'accélération précédentes. Cette flexibilité permet l'utilisation de solveurs avancés issus de la littérature sur la diffusion discrète, offrant des accélérations significatives sans sacrifier la précision. Le SBD ne nécessite aucune modification architecturale ni hyperparamètre d'entraînement supplémentaire, maintient la compatibilité avec le KV-caching exact, et peut être mis en œuvre par le fine-tuning de modèles existants de prédiction de tokens suivants. En effectuant un fine-tuning de Llama-3.1 8B et Qwen-3 8B, nous démontrons que le SBD permet une réduction de 3 à 5 fois du nombre de passes avant nécessaires pour la génération tout en atteignant des performances équivalentes à celles de l'entraînement NTP correspondant.

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.