BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding

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$.

• The paper introduces BlockPilot, an instance-adaptive policy that selects optimal block sizes for diffusion-based speculative decoding, significantly improving inference efficiency.
• It formulates block size selection as a supervised classification task over a localized interval, reducing the search space and adaptation overhead.
• Empirical results show notable speedups and increased acceptance lengths across various benchmarks and LLM scales.

BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding

Introduction and Motivation

Autoregressive LLMs exhibit significant computational inefficiency due to sequential token-by-token generation. Speculative decoding strategies have enabled substantial acceleration by introducing parallel candidate drafting and verification, often employing lightweight models as drafters. Diffusion-based speculative decoding further improves inference throughput by generating multi-token blocks in parallel through block-level diffusion models, as in DFlash (Chen et al., 5 Feb 2026). However, all prior block-diffusion speculative decoding methods rely on a fixed block size during inference, which fails to adapt to input-specific characteristics relevant to generation stability and quality.

The paper "BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding" (2606.31315) challenges the conventional design by demonstrating that the optimal block size for speculative decoding is input-dependent and exhibits strong local structure near the training block size. The authors propose BlockPilot, a plug-and-play, instance-adaptive policy learning approach for input-conditioned block size selection, yielding state-of-the-art speculative decoding efficiency across a comprehensive set of benchmarks and LLM scales.

Analysis of Block Size Variability and Locality

Through an empirical block size sweep on multiple datasets and models, the authors establish that the acceptance length (average number of tokens successfully verified and accepted per speculative decoding round) exhibits significant sample-level variability. Notably, many inputs achieve higher efficiency with block sizes distinct from the default training-time setting.

Figure 1: The proportion of samples for which the optimal block size $B^*$ matches or mismatches the training block size $B$, underscoring the inadequacy of globally fixed decoding strategies.

Despite this variability, the optimal block size distribution shows strong locality: for nearly all samples, $B^*$ remains concentrated within a narrow interval centered on the training block size $B$, typically within $[B-3, B+3]$. This observation reduces the adaptive block selection search space and permits efficient formulation of the adaptation problem as a structured classification task.

BlockPilot: Instance-Adaptive Policy Formulation

BlockPilot predicts the sample-optimal block size using the last-token predictive distribution from the target LLM's prefilling stage as a compact proxy for decoding state and context uncertainty. The block size selection is implemented as a supervised classification problem over the local block size interval. The policy network—a lightweight 2-layer MLP—operates on the prefilling token distribution and outputs the block size choice for the entire subsequent speculative decoding process.

The BlockPilot inference pipeline does not modify the backbone LLM, the diffusion draft model, or the speculative verification protocol. The block size decision is computed once per input, resulting in negligible additional latency (≈7 ms) compared to overall LLM inference, and can be seamlessly integrated into any diffusion-based speculative decoding system.

Figure 2: Schematic of the BlockPilot inference flow, from prefilling and contextual feature extraction through block size selection and parallel drafting with diffusion models.

Empirical Evaluation and Numerical Results

Comprehensive benchmarking is conducted on both general-purpose and code-specialized LLMs (Qwen3, Llama-3.1, Qwen3-Coder), across Math (GSM8K, MATH-500, AIME24), Code (HumanEval, MBPP, SWE-Bench), and Chat (MT-Bench) tasks. Comparisons are made with state-of-the-art diffusion-based (DFlash) and non-diffusion speculative decoding baselines (EAGLE-3).

The empirical results indicate consistent improvements in both speedup and acceptance length over all baselines, for all model and decoding configurations:

• On Qwen3-4B under temperature $T=1$, BlockPilot achieves an average speedup of 4.20$\times$ with an acceptance length ($\tau$) of 5.92, exceeding the best fixed-block DFlash by a substantial margin.
• Similar gains are observed for Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-Coder-30B-A3B, with maximum speedups ranging from 3.25$\times$ to 4.66$\times$.
• Gains are robust under both deterministic and stochastic generation regimes, and cover all benchmarked domains (math, code, dialogue).

  Figure 3: Model-wise speedup comparison under $B$0, demonstrating BlockPilot's consistent superiority across all settings.

These results directly demonstrate the suboptimality of global fixed block-size policies and the marked efficiency gains made possible by adaptive approaches.

Practical and Theoretical Implications

The demonstration that decoding efficiency is governed by a learnable, sample-dependent policy component—rather than purely by model or hardware architecture—has significant consequences:

• BlockPilot's instance-adaptive policy framework complements advances in block diffusion modeling and hardware acceleration, providing an orthogonal pathway for real-world inference speedup.
• The structured local decision formulation, leveraging the locality of optimal block sizes, minimizes adaptation overhead and permits scalability to even large-scale LLM deployments.
• From a theoretical perspective, the findings motivate further research into input-conditional inference strategies, policy learning for LLM decoding hyperparameters, and task-adaptive hardware utilization.

BlockPilot's design also suggests future avenues in exploiting richer decoding-state representations (e.g., from deeper attention heads), more sophisticated policy architectures (contextual bandits, reinforcement learning), or continuous block size selection. Automated, data-driven adaptation of other speculative decoding hyperparameters (e.g., candidate diversity, verification depth) may further increase responsiveness and utilization.

Conclusion

BlockPilot introduces a principled, lightweight strategy for input-adaptive block size selection in diffusion-based speculative decoding. By identifying and exploiting the structured variability and locality of optimal block size per sample, BlockPilot achieves superior decoding efficiency with negligible overhead, outperforming state-of-the-art baselines across multiple models and benchmarks. This work establishes decoding policy as a key leverage point for efficient LLM inference and provides a foundation for broader research at the intersection of inference optimization, policy learning, and LLM deployment (2606.31315).