Décodage Spéculatif par Lots Bien Réalisé

papers.abstract

Le décodage spéculatif accélère l'inférence des LLM en utilisant un petit modèle d'ébauche pour proposer plusieurs jetons qu'un modèle cible vérifie en parallèle. Étendre cette idée au traitement par lots est essentiel pour le déploiement en production, mais cela introduit le problème des tenseurs irréguliers : les séquences d'un même lot acceptent un nombre différent de jetons d'ébauche, brisant l'alignement à droite et corrompant les ID de position, les masques d'attention et l'état du cache KV. Nous montrons que plusieurs implémentations existantes de traitement par lots violent l'équivalence de sortie – l'exigence fondamentale que le décodage spéculatif doit produire des séquences de jetons identiques à la génération autorégressive standard. Ces violations se produisent précisément en raison d'une gestion inadéquate du problème des tenseurs irréguliers. En réponse, nous (1) caractérisons les exigences de synchronisation qui garantissent l'exactitude, (2) présentons un décodage spéculatif par lots axé sur la correction, EQSPEC, qui révèle que le réalignement représente 40 % de la surcharge, et (3) introduisons EXSPEC, qui maintient un pool glissant de séquences et forme dynamiquement des groupes de même longueur, pour réduire la surcharge de réalignement tout en préservant les accélérations spéculatives par séquence. Sur l'ensemble de données SpecBench, avec les paires cible/ébauche Vicuna-7B/68M, Qwen3-8B/0.6B et GLM-4-9B/0.6B, notre approche atteint jusqu'à 3 fois l'amélioration du débit pour une taille de lot de 8 par rapport à une taille de lot de 1, avec une mise à l'échelle efficace jusqu'à la taille de lot 8, tout en maintenant 95 % d'équivalence de sortie. Notre méthode ne nécessite aucun noyau personnalisé et s'intègre parfaitement aux piles d'inférence existantes. Notre code est disponible à l'adresse 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.

Décodage Spéculatif par Lots Bien Réalisé

Batch Speculative Decoding Done Right

papers.abstract

Support