Decodifica Speculativa in Batch Eseguita Correttamente

Abstract

Il decoding speculativo accelera l'inferenza dei LLM utilizzando un piccolo modello draft per proporre più token che un modello target verifica in parallelo. Estendere questa idea ai batch è essenziale per il serving in produzione, ma introduce il problema dei tensori irregolari: le sequenze nello stesso batch accettano numeri diversi di token draft, rompendo l'allineamento a destra e corrompendo gli ID posizionali, le maschere di attenzione e lo stato della KV-cache. Dimostriamo che diverse implementazioni batch esistenti violano l'equivalenza dell'output – il requisito fondamentale per cui il decoding speculativo deve produrre sequenze di token identiche alla generazione autoregressiva standard. Queste violazioni si verificano proprio a causa di una gestione impropria del problema dei tensori irregolari. In risposta, noi (1) caratterizziamo i requisiti di sincronizzazione che garantiscono la correttezza, (2) presentiamo un decoding speculativo batch orientato alla correttezza, EQSPEC, che rivela come il riallineamento assorba il 40% dell'overhead, e (3) introduciamo EXSPEC, che mantiene un pool scorrevole di sequenze e forma dinamicamente gruppi di stessa lunghezza, per ridurre l'overhead di riallineamento preservando i miglioramenti di velocità speculativi per sequenza. Sul dataset SpecBench, attraverso le coppie target/draft Vicuna-7B/68M, Qwen3-8B/0.6B e GLM-4-9B/0.6B, il nostro approccio raggiunge un miglioramento di produttività fino a 3 volte con dimensione del batch 8 rispetto alla dimensione del batch 1, con un efficiente scaling fino alla dimensione del batch 8, mantenendo il 95% di equivalenza dell'output. Il nostro metodo non richiede kernel personalizzati e si integra agevolmente con gli stack di inferenza esistenti. Il nostro codice è disponibile su 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.

Decodifica Speculativa in Batch Eseguita Correttamente

Batch Speculative Decoding Done Right

Abstract

Support