Decodificação Especulativa em Lote Feita Corretamente

Resumo

A descodificação especulativa acelera a inferência de LLMs utilizando um pequeno modelo de rascunho para propor múltiplos tokens que um modelo alvo verifica em paralelo. Estender esta ideia para lotes (batches) é essencial para a produção em servidores, mas introduz o problema do tensor irregular: sequências no mesmo lote aceitam números diferentes de tokens de rascunho, quebrando o alinhamento à direita e corrompendo IDs de posição, máscaras de atenção e o estado da KV-cache. Demonstramos que várias implementações existentes de lotes violam a equivalência de saída – o requisito fundamental de que a descodificação especulativa deve produzir sequências de tokens idênticas à geração autoregressiva padrão. Estas violações ocorrem precisamente devido ao manuseamento inadequado do problema do tensor irregular. Em resposta, nós (1) caracterizamos os requisitos de sincronização que garantem a correção, (2) apresentamos uma descodificação especulativa em lote com prioridade à correção, o EQSPEC, que expõe o realinhamento como consumidor de 40% da sobrecarga, e (3) introduzimos o EXSPEC, que mantém um conjunto deslizante de sequências e forma dinamicamente grupos de mesmo comprimento, para reduzir a sobrecarga de realinhamento, preservando os ganhos de velocidade especulativa por sequência. No conjunto de dados SpecBench, através dos pares alvo/rascunho Vicuna-7B/68M, Qwen3-8B/0.6B e GLM-4-9B/0.6B, a nossa abordagem atinge até 3 vezes a melhoria de throughput no tamanho de lote 8 em comparação com o tamanho de lote 1, com escalagem eficiente até ao tamanho de lote 8, mantendo 95% de equivalência de saída. O nosso método não requer kernels personalizados e integra-se facilmente com as pilhas de inferência existentes. O nosso código está disponível em 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ção Especulativa em Lote Feita Corretamente

Batch Speculative Decoding Done Right

Resumo

Support