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

Hao Zhang, Yiming Hu, Yong Wang, Mingqiao Mo, Xin Xiao, Xiangxiang Chu

BlockPilot dynamically selects optimal block sizes for diffusion-based speculative decoding to maximize inference speedup.

How can we dynamically adjust the number of tokens (block size) a diffusion-based draft model proposes to maximize inference speedup for each specific input?

Diffusion-based speculative decoding accelerates inference by generating token blocks in parallel, but existing methods rely on a fixed block size that is often suboptimal for diverse input samples. BlockPilot treats block size as a learnable control variable, using a lightweight predictor to select the optimal size for each input based on the prefilling stage's predictive distribution. This adaptive approach consistently outperforms fixed-block strategies, achieving up to 4.20× speedup on Qwen3-4B without modifying the underlying model architecture.

Paper Primer

Speculative decoding relies on a draft model to propose tokens that a target model verifies in parallel. While diffusion-based drafters improve parallelism, they typically use a static block size inherited from training, which fails to account for the fact that some inputs are highly predictable and support large blocks, while others accumulate errors quickly and require smaller, more conservative blocks.

BlockPilot is a plug-and-play policy network: it maps the target model's final prefilling-stage predictive distribution to an optimal block size chosen from a local candidate interval. The core insight is that optimal block sizes exhibit a "local interval property," concentrating near the training-time configuration, which allows the system to treat selection as a simple, low-dimensional classification task rather than an expensive online search.

BlockPilot consistently improves inference speedup and acceptance length across diverse model scales and benchmarks.

On Qwen3-4B at temperature T=1, the method achieves a 4.20× speedup and an average acceptance length of 5.92, outperforming the strongest fixed-block baseline (DFlash).

Why is a fixed block size suboptimal for speculative decoding?

Inputs vary in semantic constraints and predictability; larger blocks improve efficiency for constrained continuations but cause error accumulation for less predictable trajectories, making a single static block size misaligned with input-level characteristics.

Does this method require retraining the target or draft models?

No. The predictor is a lightweight, separate multilayer perceptron that runs once after the prefilling stage, allowing it to integrate into existing frameworks without modifying the generation pipeline or the original model weights.

Decoding policy is a critical, learnable dimension of inference efficiency; by treating block size as a sample-adaptive decision, practitioners can extract significant speedups from existing diffusion-based speculative decoding frameworks.

Introduction to Speculative Decoding

Fixed block sizes limit diffusion‑based speculative decoding, motivating an adaptive solution.

Speculative decoding speeds up inference by letting a lightweight draft model generate candidate tokens in parallel, which the target model then verifies losslessly. Diffusion‑based variants further boost parallelism by generating whole token blocks via block‑level diffusion, achieving state‑of‑the‑art performance. However, all existing approaches fix the block size during inference, ignoring that the optimal size varies with each input.

It uses a fast draft model to propose a batch of candidate tokens, while the target model verifies each candidate in parallel, allowing multiple tokens to be accepted per step without changing the final distribution.

A dLLM treats token generation as a denoising diffusion process, enabling it to produce an entire block of tokens in a single forward pass.

A fixed block size treats decoding as a static hyperparameter, but the predictability of the next tokens differs across samples, making a single size suboptimal. This mismatch limits the throughput gains of diffusion‑based speculative decoding. BlockPilot addresses the issue by predicting an instance‑adaptive block size from the prefilling representation.

**Figure 1.** Diffusion-based speculative decoding with a dLLM draft model. The dLLM proposes a block of tokens in parallel, while the target LLM verifies the block and accepts the longest consistent prefix.

The bottleneck of fixed‑size speculative decoding is its inability to adapt block size to input‑specific predictability.

The BlockPilot Inference Pipeline

Methodology introduces BlockPilot Predictor to adaptively select block size per input, improving speculative decoding throughput.

Using a fixed block size $B$ during speculative decoding ignores the fact that token predictability varies across inputs. This mismatch leads to suboptimal acceptance lengths and wasted parallel compute. The method must therefore adapt $B$ per sample to recover efficiency.

A tiny neural network consumes the predictive distribution of the last token after prefilling and instantly outputs the block size that will maximize the expected acceptance length for that sample.

Compute logits for each candidate: $o_3=0.45$, $o_4=0.70$, $o_5=0.30$ (higher logit means higher predicted acceptance).

Apply softmax: $P(3)=0.25$, $P(4)=0.60$, $P(5)=0.15$.

Take the argmax → $\hat{B}=4$.

Run speculative decoding with block size 4; the measured acceptance length $\tau(4)=3.2$ tokens per step.

If we had mistakenly used $B=5$, $\tau(5)$ would drop to $2.1$, showing the predictor avoided a costly overshoot.

The predictor selects the block size that balances parallelism (larger $B$) against verification failures (smaller $\tau$), achieving the highest effective throughput.

How does BlockPilot differ from a simple heuristic that always picks the training block size $B$?

Heuristics ignore the per‑sample predictive distribution and therefore cannot react to variations in context difficulty. BlockPilot reads the full distribution, evaluates the expected acceptance length for each candidate, and picks the size that maximizes that expectation, which can be smaller or larger than the training $B$ depending on the input.

Acceptance length is the average number of tokens that survive verification for a given block size; it quantifies how much parallel work actually contributes to the final output.

Why is acceptance length more informative than raw latency when choosing a block size?

Latency $L(B)$ combines both compute costs and acceptance length, but the compute terms $T_{\text{draft}}$ and $T_{\text{verify}}$ are nearly constant across $B$. Therefore variations in $L(B)$ are driven almost entirely by $\tau(B)$. Maximizing $\tau$ directly improves throughput, while latency is just a derived metric.

Inference pipeline with BlockPilot Predictor

Table 1 shows that the BlockPilot predictor adds only 0.32 B parameters and 7.34 ms latency, a negligible overhead compared with the multi‑billion‑parameter backbone models.

**Figure 4.** Overview of the BlockPilot inference pipeline. Given an input sequence, the target LLM performs prefilling and produces the predictive distribution of the last token, which serves as a compact representation of the decoding state. This distribution is then fed into a lightweight block size predictor to determine an instance-specific block size. Based on the predicted block size, the diffusion-based draft model generates a block of draft tokens in parallel.

Experimental Setup

We describe models, baselines, metrics, and implementation details for evaluating BlockPilot.

We evaluate BlockPilot on four LLMs—Qwen3‑4B, Qwen3‑8B, Llama‑3.1‑8B‑Instruct, and Qwen3‑Coder‑30B‑A3B—covering both general‑purpose and code‑specialized domains, across three task categories: mathematical reasoning (GSM8K, MATH‑500, AIME24), code generation (HumanEval, MBPP, SWE‑Bench), and conversational generation (MT‑Bench).

Baselines include standard autoregressive decoding, the classic speculative decoding with an autoregressive drafter, EAGLE‑3, and the diffusion‑based SOTA method DFlash (with varying block sizes).

We report efficiency using Average Acceptance Length $\\tau$—the mean number of draft tokens accepted per cycle—and Speedup Ratio, the factor by which inference time improves over standard autoregressive decoding.

Experiments run on PyTorch with Hugging Face Transformers on NVIDIA H100 80 GB GPUs; the training set combines ShareGPT, WSC, and COPA, training for 100 epochs with Adam (lr 1e‑5). The predictor network has two hidden layers of size 2048 and uses $k$ = 2.

Performance and Ablation Studies

BlockPilot delivers up to 0.24× speedup over the strongest fixed‑block baseline.

Diffusion‑based speculative decoding fixes the draft block size, ignoring per‑sample predictability. BlockPilot instead selects the block size adaptively for each input step.

BlockPilot improves speedup by 0.24× on Qwen3‑8B (temperature = 0) compared to the strongest fixed‑block baseline DFlash(16).

On Qwen3‑8B, speedup rises from 4.42× (DFlash(16)) to 4.66× (BlockPilot).

**Figure 2.** Speedup comparison across models under temperature $T = 1$. Our method achieves the highest acceleration across all settings. Here, DFlash($n$) denotes DFlash with block size $n$.

**Figure.** (Left) Bar chart showing the proportion of cases where $B^* = B$ (dark blue) versus $B^* \neq B$ (tan) across five datasets: GSM8K, MT-Bench, AIME24, MBPP, and Alpaca. (Middle and Right) Histograms showing the distribution of values relative to $B$, with a superimposed fitting curve, illustrating different probability density patterns.

Ablation studies (Tables 3–5) confirm that a two‑layer MLP with hidden size 2048, a candidate interval radius k = 2, and raw prefilling distributions as predictor input together yield the best speedup and acceptance length.

Theoretical Analysis

Analyzes the theoretical basis of adaptive block size and its practical limits.

We first formalize acceptance length as a prefix‑survival process, then examine how block size interacts with predictability and finally bound the regret of imperfect block‑size prediction.

The event that the first $i$ drafted tokens are all accepted is denoted $A_i$.

Conditioning on the prefix that has already survived, we define the per‑position acceptance probability.

Increasing the block size $b$ extends the summation horizon but also changes the conditional probabilities $q_j$, creating a trade‑off between more candidate positions and harder verification.

Next we separate the influence of sample predictability from the degradation caused by deviating from the training block size.

If predictability is roughly constant across positions, the expression collapses to a geometric series.

The optimal block size for a given sample is the one that maximizes $\tau$ over a candidate set $\mathcal{B}$.

Finally we bound the loss incurred when the predictor selects a sub‑optimal block size.

Empirically, the predictor’s lightweight design suffices because the regret bound guarantees that near‑optimal choices already yield most of the possible speedup.

Performance evaluation on Llama‑3.1‑8B‑Instruct and Qwen3‑Coder‑30B‑A3B shows consistent speedups (up to 4.12×) and competitive acceptance lengths across temperature settings.

Limitations stem from the offline data‑construction cost: evaluating five block sizes for a 32 B model takes ≈ 25 s per training sample, which, while acceptable offline, motivates future work on more efficient search strategies.

Related Work and Conclusion

We situate BlockPilot among prior speculative decoding and diffusion approaches, and highlight DFlash as the leading diffusion‑based method.

Speculative decoding reduces the sequential bottleneck of autoregressive generation by drafting multiple tokens in parallel and verifying them with a larger target model. Diffusion Language Models (dLLMs) replace autoregression with parallel masked token prediction, but fixed‑length blocks limit their efficiency. Recent diffusion‑based speculative decoding methods, such as TiDAR and DiffuSpec, still suffer from high memory and latency costs, whereas DFlash achieves state‑of‑the‑art draft quality and acceptance length.

DFlash improves block‑diffusion speculative decoding by injecting the target model’s hidden states into the diffusion drafter, which raises draft quality and extends the acceptance length.

The conclusion emphasizes that block size is the dominant factor for inference efficiency, varies per sample, and can be predicted from the prefilling distribution. BlockPilot leverages this insight to select an adaptive block size per step, demonstrating that the decoding policy—not just model architecture—drives performance gains.

**Table.** Performance comparison of different methods across various benchmarks and models, measured by speedup and $\tau$ (tau).

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers