Paper 2606.31315
BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding
- Published
- Jun 2026
- Research lab
- Independent
- Citations
- 0
- GitHub
- Not linked
01 In brief
Summary
BlockPilot introduces a sample-adaptive policy for diffusion-based speculative decoding, addressing the inefficiency of fixed block sizes.
The authors observe that the optimal block size varies across samples but concentrates near the training block size, enabling a lightweight classification approach.
BlockPilot predicts the block size from the last token's predictive distribution after prefilling, using a two-layer MLP.
This method integrates seamlessly into existing frameworks, requiring only a single prediction per sample.
Experiments on Qwen3-4B, Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-Coder-30B-A3B show consistent speedups over fixed-block baselines, achieving up to 4.20x speedup and an acceptance length of 5.92 on Qwen3-4B at temperature 1.
The method's overhead is minimal, with the predictor adding only 0.32B parameters and 7.34 ms latency.
02 From the paper
Abstract
Speculative decoding accelerates inference by using a lightweight draft model to generate candidate tokens in parallel, and are then verified by the target model, enabling lossless acceleration. Recently, diffusion-based speculative decoding further improves parallelism by generating multiple tokens per forward pass via block-level diffusion, achieving state-of-the-art (SOTA) performance. However, existing methods adopt a fixed inference block size and assume a uniform optimal decoding strategy across all inputs. In this paper, we show that this assumption is suboptimal, as the optimal block size varies across samples and plays a critical role in speculative decoding performance. Moreover, these values exhibit a clear local structure, concentrating around the training block size, which reduces the problem to a low-dimensional and structured decision space. Based on these insights, we propose BlockPilot, a sample-adaptive policy that predicts the optimal block size from the prefilling representation. Specifically, we formulate block size selection as a lightweight policy learning problem and propose an instance-adaptive decision mechanism that predicts the optimal block size based on the representation of the prefilling stage. The prediction is performed only once after prefilling, allowing for seamless integration. Extensive experiments demonstrate that our method is plug-and-play, introduces minimal overhead, and consistently improves efficiency, achieving an acceptance length of 5.92 and a 4.20$\times$ speedup on Qwen3-4B under temperature $T=1$.