AdaFlash: Adaptive Speculative Decoding via On-Policy Distilled Diffusion Drafters

Abstract: Speculative decoding, in which a lightweight draft model first generates a draft sequence that is then verified in parallel by the target model, has become a prevalent paradigm for accelerating LLM inference. Recent work such as DFlash further boosts drafting efficiency by leveraging diffusion drafters, whose parallel denoising mechanism enables draft generation in a single forward pass. In this work, we uncover a central pitfall of diffusion drafters: bidirectional attention is a double-edged sword. On one hand, it endows the model with parallel generation and global contextual modeling capabilities; on the other hand, this inherent global dependency introduces high variance at both the domain-level and the token-level: acceptance rates fluctuate substantially across different domains, and draft token quality also varies heterogeneously at different token positions. To tackle this issue, we propose AdaFlash framework, comprising two components: (i) an on-policy distillation (OPD) algorithm with reverse-KL divergence tailored for diffusion drafters, bringing stable convergence and effectively reducing domain-level variance; and (ii) an adaptive length head that dynamically adjusts the candidate sequence length on the fly, substantially lowering the verification cost of the target model and handling token-level variance. Experiments demonstrate that AdaFlash consistently improves speedup rate during deployment, with especially significant gains in high-concurrency scenarios, achieving up to approximately 66% higher throughput than previous state-of-the-art methods.

• The paper introduces AdaFlash, combining reverse-KL on-policy distillation with entry-wise divergence clipping to improve one-step diffusion drafters while preserving the target model’s output distribution.
• An online adaptive length head predicts acceptance rates and truncates verification dynamically, delivering up to 5.3× speedup over autoregressive decoding and 66% higher throughput than prior methods.
• Experiments across Qwen models, reasoning benchmarks, concurrency levels, sampling settings, GPUs, and NPUs show that AdaFlash maintains positive speedups at high concurrency where fixed-length speculative decoding can underperform standard decoding.

Background and motivation

Speculative decoding accelerates LLM inference by having a lightweight draft model propose a candidate sequence that a larger target model verifies in parallel, with rejection sampling guaranteeing that the output distribution is identical to the target's. The speedup is governed by the draft-to-target inference time ratio $\rho$, the candidate length $k$, and the acceptance rate. A recent line of work replaces autoregressive (AR) drafters with diffusion LLMs (dLLMs); in particular, DFlash introduced a one-step diffusion drafter that generates an entire $k$-token block in a single forward pass, reducing the per-token drafting overhead by roughly a factor of $k$ relative to AR drafters.

The paper's central observation is that the bidirectional attention enabling this parallel generation is a double-edged sword: the resulting global dependency introduces high variance in draft quality at two distinct levels. Domain-level variance manifests as large fluctuations in acceptance rate across task domains—for Qwen3-8B, DFlash's average accepted length $\tau$ ranges from 3.40 on ShareGPT to 7.09 on MathQA (a ratio of roughly $2.1\times$), whereas EAGLE-3 varies only between 3.89 and 4.60 (ratio $\approx 1.2\times$). Token-level variance manifests as heterogeneous per-position acceptance probabilities within a single draft: the per-position acceptance probability decreases monotonically along the draft sequence, and two nearly identical prefixes differing by one token can yield drastically different acceptance lengths. A static system with a fixed drafter and constant candidate length consequently suffers degraded acceptance on out-of-distribution queries and wastes target-model compute verifying tokens unlikely to be accepted.

Method

AdaFlash addresses the two variance levels with two components, supported by serving-infrastructure changes.

On-policy distillation (OPD) with reverse-KL. The framework exploits feedback that speculative decoding provides naturally: at each round, the target model returns its output distributions over candidate positions, forming a draft–feedback–adapt loop that keeps the training distribution aligned with the inference distribution. Because bidirectional attention in diffusion drafters produces high-entropy distributions, and speculative decoding only requires covering the target's high-probability modes, the authors adopt reverse-KL rather than mode-covering forward KL. The loss mixes reverse-KL with a hard-label cross-entropy on the target's top-1 token, $$\ell_{\mathrm{OPD}} = \alpha\,\ell_{\mathrm{hard}} + (1-\alpha)\,\ell_{\mathrm{rkl}}^{\mathrm{clip}}$$, where the hard-label term provides a low-variance gradient signal. A key technical contribution is entry-wise divergence clipping: the reverse-KL at each draft position sums per-vocabulary-entry terms, and a few entries where the drafter places non-negligible mass but the target does not produce large ratios that dominate the gradient. Bounding each entry's contribution at a threshold $\delta$ before summation suppresses these outlier gradients while preserving the remaining signal.

Adaptive length head. A lightweight head attached to the drafter predicts the acceptance rate $\widehat{a}_t \in [0,1]$ from per-position features via linear projection, SiLU activation, mean pooling, and a sigmoid output; the verification length is then $k$0, and only the first $k$1 draft tokens are verified. The head is trained with an MSE loss against the ground-truth acceptance rate observed at each round—supervision that costs nothing extra since verification outcomes are already available. Its gradients are detached from the drafter so that the MSE and reverse-KL objectives remain decoupled. Notably, the head is continuously updated online, co-adapting with the evolving drafter; this distinguishes the approach from the concurrent DSpark, which uses an offline-trained per-position confidence head requiring post-hoc calibration.

Infrastructure. The system, built on SGLang, decouples an inference server from a training server via a shared replay buffer: on-policy trajectories are logged during inference, and the training server asynchronously updates drafter and length-head weights, which are hot-swapped into GPU memory between scheduling steps without touching the target model. The serving engine is modified to accept variable candidate lengths, packing requests into a compact verification batch of total length $k$2, with an exponential moving average estimate of the number of admissible requests to manage GPU memory under dynamic batch shapes.

Experimental results

The evaluation spans eight benchmarks (including long-sequence MATH-500 and AIME25 with 32K-token outputs in thinking mode), three target models (Qwen3-8B, Qwen3-Coder-30B-A3B, and Qwen3.5-9B with Gated DeltaNet), and concurrency levels $k$3. Contenders include standard AR decoding, EAGLE-3, DFlash, and OSD (online distillation with forward KL applied to DFlash).

The headline results support the paper's claims directly. At $k$4 with Qwen3-8B, AdaFlash achieves an average speedup of $k$5 versus $k$6 for OSD, $k$7 for DFlash, and $k$8 for EAGLE-3, with accepted length on MathQA rising from 7.09 (DFlash) to 9.44 (OSD) to 9.83 (AdaFlash)—confirming that reverse-KL on-policy distillation outperforms forward-KL distillation. The high-concurrency regime is where the adaptive length head proves decisive: at $k$9, AdaFlash sustains $k$0 speedup while OSD ($k$1) and DFlash ($k$2) fall below standard AR decoding, because fixed-length verification wastes GPU compute when few draft tokens are accepted. The paper reports up to approximately 66% higher throughput than prior state of the art and up to $k$3 speedup over AR decoding. Ablations isolate each component: with the length head enabled, divergence clipping and the mixture OPD loss each add measurable gains (speedup at $k$4 rising from $k$5 to $k$6 to $k$7 on GSM8K); conversely, replacing adaptive truncation with a fixed verification length of 11 collapses the $k$8 speedup to $k$9, and enabling online updates of the length head restores $k$0, showing that co-adaptation with the shifting drafter is essential.

Hyperparameter analysis shows the mixing coefficient $k$1 is optimal (pure reverse-KL at $k$2 and pure hard-label at $k$3 both regress), the clipping threshold $k$4 best balances robustness and convergence (aggressive clipping at $k$5 reduces throughput), and the length scale factor $k$6 yields stable speedups between $k$7 and $k$8 over $k$9, with $\tau$0 best.

Additional experiments broaden the evidence: on Qwen3.5-9B, AdaFlash achieves $\tau$1 average speedup at $\tau$2 versus $\tau$3 for OSD and $\tau$4 for the model's native multi-token prediction; long-sequence results on MATH-500 ($\tau$5 at $\tau$6) and AIME25 ($\tau$7) confirm the gains extend to extended reasoning; results under temperature-1 sampling retain $\tau$8 speedup at $\tau$9 despite reduced accepted length ($2.1\times$0 vs. $2.1\times$1 greedy); and a from-scratch drafter trained on Qwen3-1.7B reaches $2.1\times$2 speedup, substantially outperforming EAGLE-3's $2.1\times$3 under the same budget. Cross-domain evaluation on Huawei Ascend 910C NPUs shows similar gains, indicating hardware generality.

Limitations and open questions

The paper concedes several constraints. On Qwen3.5-9B, gains on ShareGPT narrow or reverse slightly at $2.1\times$4, which the authors attribute to SGLang's immature support for Gated DeltaNet preventing fully efficient variable-length verification scheduling—a limitation they expect to close as framework support matures, but which currently undermines the high-concurrency advantage on that architecture. The cross-domain experiments show that offline OPD on a mixed-domain corpus (PerfectBlend) improves over the offline DFlash drafter but still lags in-domain online adaptation, confirming that the method's benefits depend on continuous adaptation to the deployment distribution rather than conferring inherent distributional robustness. The adaptive length head's predicted acceptance rate is a scalar aggregate, discarding per-position information; whether a finer-grained formulation would improve truncation decisions remains unexamined. The clipping threshold $2.1\times$5 and mixing coefficient $2.1\times$6 were tuned empirically on GSM8K, and no theoretical account is given for how these choices transfer across domains or model scales. Finally, the evaluation is limited to Qwen-family targets and greedy or temperature-1 decoding; behavior under intermediate temperatures, other model families, and multi-turn agentic workloads is left open.

Conclusion

AdaFlash identifies high variance—domain-level and token-level—as the central deficiency of one-step diffusion drafters in speculative decoding, and addresses it with reverse-KL on-policy distillation with entry-wise divergence clipping and an online-updated adaptive length head, realized through an asynchronous training–inference serving pipeline. The empirical results are consistent across dense and mixture-of-experts targets, long-sequence reasoning, sampling regimes, from-scratch training, and both GPU and NPU backends, with the most distinctive contribution being the restoration of positive speedup at high concurrency where fixed-length diffusion and AR speculative methods fall below standard autoregressive decoding.