FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration

TL;DR

FlexDraft introduces a flexible, lossless speculative decoding framework that adapts to batch size variations, improving inference speed and accuracy in large language models through attention tuning and bonus-guided calibration.

Contribution

It proposes novel techniques—Attention Tuning, Bonus-guided Calibration, and Flex Decoding—to enhance parallel speculative decoding efficiency and robustness without quality loss.

Findings

Achieves high throughput at large batch sizes without quality degradation.

Effectively mitigates draft verification mismatch caused by bonus token uncertainty.

Maintains high draft quality with minimal additional training or computational overhead.

Abstract

Speculative decoding accelerates memory-bound LLM inference without quality degradation by using a fast drafter to propose multiple candidate tokens and the target model to verify them in parallel. However, conventional sequential speculative decoding suffers from mutual waiting between drafting and verification, and repeated exchange of intermediate states further increases memory access overhead. Parallel speculative decoding addresses this limitation by performing drafting and verification within a single target forward pass, allowing future drafts to be prepared while current candidates are being verified. Although effective at small batch sizes, existing parallel speculative decoding methods either require costly continual pretraining with quality degradation or suffer from low acceptance rates. More importantly, this paradigm inherently suffers from uncertainty in both the bonus token and the accepted length, leading to draft verification mismatch and causing throughput gains to collapse at large batch sizes. To address these limitations, we introduce FlexDraft, a lossless speculative decoding framework that flexibly adapts to varying batch sizes through three key designs. (1) Attention Tuning enables block diffusion drafting by tuning only the attention projectors of the final few layers on mask tokens, while keeping the autoregressive path frozen to preserve the target distribution and produce high quality drafts with minimal trainable parameters. (2) Bonus-guided Calibration uses a lightweight MLP conditioned on the resolved bonus token to calibrate draft logits, mitigating draft verification mismatch caused by bonus token uncertainty. (3) Flex Decoding dynamically switches between parallel draft and verify at small batch sizes and sequential draft then verify at large batch sizes, and adjusts verification length based on draft confidence to eliminate redundant computation.