PaperlayerView exact revision on arXiv

Source-linked AI summary

DFlash: Block Diffusion for Flash Speculative Decoding

Jian Chen, Yesheng Liang, Zhijian Liu

arXiv:2602.06036v2cs.CL

TL;DR

Autoregressive decoding is sequential and slow, while diffusion models generate in parallel but often sacrifice accuracy. DFlash uses a lightweight diffusion drafter conditioned on target-model features and reports higher speedups than EAGLE-3 across evaluated settings.

  • Problem

    Autoregressive models provide strong performance but suffer sequential latency, whereas diffusion models enable parallel generation at potential accuracy cost.

  • Method

    DFlash uses a lightweight block diffusion model for parallel speculative drafting, conditioning it on hidden features from the target autoregressive model.

  • Results

    DFlash achieves a 4.9× average speedup over the autoregressive baseline under greedy decoding and a 2.4× improvement over EAGLE-3.

  • Takeaways & Limitations

    DFlash positions diffusion models as lightweight specialized drafters whose parallelism and target-model conditioning support faster speculative decoding while preserving verified output quality.

  • Takeaways & Limitations

    Adaptive block-size scheduling is left for future work, although reducing block size may improve speedup in compute-bound serving settings.

Abstract

from arXiv · show

Autoregressive large language models (LLMs) deliver strong performance but require inherently sequential decoding, leading to high inference latency and poor GPU utilization. Speculative decoding mitigates this bottleneck by using a fast draft model whose outputs are verified in parallel by the target LLM; however, existing methods still rely on autoregressive drafting, which remains sequential and limits practical speedups. Diffusion LLMs offer a promising alternative by enabling parallel generation, but current diffusion models typically underperform compared with autoregressive models. In this paper, we introduce DFlash, a speculative decoding framework that employs a lightweight block diffusion model for parallel drafting. By generating draft tokens in a single forward pass and conditioning the draft model on context features extracted from the target model, DFlash enables efficient drafting with high-quality outputs and higher acceptance rates. Experiments show that DFlash achieves over 6x lossless acceleration across a range of models and tasks, delivering up to 2.5x higher speedup than the state-of-the-art speculative decoding method EAGLE-3.

1. Introduction

LLM inference is bottlenecked by sequential autoregressive generation, while existing speculative decoding still drafts autoregressively. DFlash addresses this trade-off with a lightweight block-diffusion drafter conditioned on target-model features, achieving strong practical speedups.

  • Autoregressive LLM decoding is slow, memory-bound, and poorly utilizes GPUs because each token depends on the full preceding context.
  • Speculative decoding verifies draft tokens in parallel, but methods such as EAGLE-3 retain sequential autoregressive drafting.
  • Diffusion models enable parallel generation, yet existing models often sacrifice quality or require many denoising steps that reduce inference speed.
  • DFlash uses a lightweight block-diffusion model for parallel drafting and conditions it on hidden features extracted from the target model.

2. Related Work

Related work combines speculative verification with alternative drafting strategies, including diffusion and feature-conditioned approaches. Existing diffusion drafters face quality, memory, latency, or losslessness limitations that motivate DFlash.

  • Speculative decoding verifies proposed token sequences in parallel, while EAGLE methods exploit feature-level context from a frozen target model.
  • Most existing speculative decoding methods still use autoregressive drafting, which remains sequential and limits speedups.
  • Block diffusion models denoise masked tokens block-by-block, combining parallel generation with autoregressive structure and improving practical cache compatibility.
  • TiDAR supports parallel diffusion thinking and sequential autoregressive talking, but its final generation quality is not yet lossless.
  • Large diffusion drafters such as DiffuSpec and SpecDiff-2 can achieve long acceptance lengths, but their memory and drafting costs often offset practical speedups.

3. Preliminaries

Speculative-decoding speedup depends on balancing drafting cost, verification cost, and accepted tokens per cycle. Diffusion drafting changes this balance by generating blocks in parallel, enabling deeper drafters and improved acceptance-length efficiency.

  • 3.1. Speculative Decoding Speedup: Speculative decoding drafts γ tokens with a smaller model and verifies them in parallel using the target model.
  • 3.1. Speculative Decoding Speedup: Speedup is η = Ltarget/L, where cycle latency combines drafting and verification costs and τ is expected accepted tokens, including the bonus token.
  • 3.1. Speculative Decoding Speedup: Speedup improves by increasing expected acceptance length τ or reducing drafting overhead Tdraft.
  • 3.2. Autoregressive vs. Diffusion Drafting: Autoregressive drafting costs Tdraft = γ · tstep, so cost grows linearly with the speculation budget γ.
  • 3.2. Autoregressive vs. Diffusion Drafting: Diffusion drafters generate all γ tokens in one parallel pass, making Tdraft largely insensitive to γ for moderate block sizes.
  • 3.2. Autoregressive vs. Diffusion Drafting: A five-layer DFlash drafter generating 16 tokens has lower latency and higher acceptance length than one-layer EAGLE-3 generating 8 tokens.

4. Method

DFlash combines target-model conditioning with block diffusion to draft future tokens in parallel using a lightweight model. Its training aligns this drafting process with speculative decoding through KV injection, masked-block construction, efficient long-context sampling, and early-position loss weighting.

  • Context conditioning: DFlash extracts and fuses hidden representations from multiple target-model layers, then uses the resulting context feature to condition the draft model.The target performs a prefill pass, and a lightweight projection layer combines features sampled from shallow to deep layers.
  • Context conditioning: DFlash injects the fused target context directly into every draft layer’s Key and Value projections, preserving conditioning across drafting iterations.The projected features are stored in the draft model’s KV cache and reused across iterations.
  • Parallel diffusion drafting: All masked positions in each block are decoded in parallel in a single forward pass, reducing drafting latency and improving GPU utilization relative to autoregressive drafting.Block-level diffusion allows parallel generation even with deeper draft models.
  • Training alignment: DFlash samples response anchors, masks the remaining block positions, and trains the draft model to predict the next block size −1 tokens in parallel.This construction matches inference, where drafting conditions on a clean target-produced bonus token.
  • Training alignment: Training concatenates blocks and uses sparse attention to permit within-block bidirectional attention while preventing information flow across different blocks.The design supports joint processing of multiple draft blocks in one forward and backward pass.
  • Training efficiency and optimization: DFlash bounds long-context training cost by fixing masked-block counts and randomizing anchor positions, while exponentially weighting earlier block positions to improve acceptance length.The weighting emphasizes early predictions because errors there invalidate subsequent tokens; γ controls the decay rate.
  • Training efficiency and optimization: The draft model shares the target model’s frozen token embedding and language-modeling head, updating only draft Transformer layers to remain lightweight.This design reduces trainable parameters and aligns the diffusion adapter with the target representation space.

5. Experiments

DFlash is evaluated across diverse models, tasks, decoding settings, serving conditions, and design choices. It consistently delivers substantial speedups, while ablations show how conditioning, model depth, hidden features, block size, and long-context adaptation affect performance.

  • Models and Evaluations: Experiments cover LLaMA-3.1 and Qwen3 models across math, code, and chat benchmarks, including realistic SGLang serving setups.The evaluation also includes greedy and nongreedy decoding, reasoning mode, concurrency levels, and long-context conditions.
  • Instruct Models: DFlash consistently outperforms EAGLE-3 across tasks and settings, achieving 4.9× speedup over baseline and 2.4× improvement over EAGLE-3 under greedy decoding.Under nongreedy sampling, it maintains 4.1× speedup over baseline and a 2.2× improvement over EAGLE-3.
  • Reasoning Models: DFlash maintains roughly 4.5× and 3.9× speedups over baseline for reasoning models with thinking mode enabled.The draft models are trained on target-model outputs containing reasoning traces.
  • Serving Performance: DFlash provides speedups across Qwen3 models and concurrency levels from 1 to 32, reaching up to 5.1× on Qwen3-8B in SGLang.These experiments use a single B200 GPU with the FlashAttention-4 backend and scheduling overlap.
  • Long Context Adaptation: Fine-tuning preserves or improves acceptance length beyond 4K context, whereas the base drafter degrades as context length grows.This indicates that target features remain representative at long contexts and that lightweight adaptation can learn longer-range patterns.
  • Ablations: Five target hidden features outperform three, while the five-layer draft model can achieve higher overall speedup than an eight-layer model because of lower drafting cost.KV injection also improves acceptance length over input fusion, and block diffusion drafts multiple tokens in parallel.
  • Training–Inference Block Size: Larger training block sizes generalize well to smaller inference block sizes, but smaller-trained models do not generalize equivalently in the reverse direction.This asymmetry could support dynamic block-size scheduling, although adaptive scheduling is left for future work.

6. Conclusion

DFlash uses diffusion models only for speculative drafting, combining their parallelism with autoregressive verification to accelerate inference without changing the target model’s output quality. The paper also proposes lightweight, target-conditioned diffusion drafters as a practical development paradigm for efficient speculative decoding.

  • DFlash confines diffusion models to drafting, exploiting parallel generation while avoiding the quality degradation associated with standalone diffusion language models.
  • Conditioning the diffusion drafter on rich target-model context enables high acceptance rates and pushes inference speed beyond prior speculative decoding methods.
  • DFlash suggests developing diffusion language models as lightweight, specialized drafters for fast and accurate block prediction rather than competing in end-to-end generation quality.Speculative verification supports aggressive reduction in denoising steps while preserving output-quality guarantees within the decoding framework.

Impact Statement

The paper targets more efficient LLM inference through system- and algorithm-level speculative-decoding optimization. Its expected impacts remain within the benefits and known limitations of existing LLM deployments.

  • DFlash reduces LLM inference and serving costs through system- and algorithm-level optimization without altering model capabilities or intended use cases.
  • The authors foresee no significant new ethical risks beyond those already associated with large language models in general.

A.1. Training Implementation

DFlash training uses fixed optimization and data-processing procedures, with target hidden features computed either online or precomputed offline. Offline caching reduces computational overhead during draft-model optimization.

  • Draft models are trained for 6 epochs with AdamW, a 6 × 10^-4 learning rate, gradient clipping at 1.0, and cosine scheduling with 0.04 warmup.
  • Training uses maximum sequence lengths of 3072 tokens, or 4096 for Qwen3-Coder, with 512 anchor positions randomly sampled per sequence.
  • Target hidden features can be computed online during training or precomputed and cached offline to reduce optimization overhead.

A.2. Diffusion Drafter without Target Feature

Without target context features, the 5-layer block diffusion draft model achieves only modest acceptance length and speedup.

  • The 5-layer block diffusion draft model omits target context features.
  • It achieves only modest acceptance length and speedup.

A.3. KV Injection Mechanism and Memory Overhead

DFlash conditions its diffusion drafter on projected target-model features through KV injection, adding limited memory and activation overhead. Evaluations show gains across models and serving frameworks, while loss decay improves convergence and random anchor sampling improves drafting outcomes.

  • KV Injection Mechanism: DFlash concatenates hidden states from selected target layers and projects them once into the draft hidden dimension.
  • KV Injection Mechanism: Projected target features are shared across draft layers and provide additional key-value entries for masked-block draft tokens.
  • KV Injection Mechanism: Target features bypass the draft model’s Q projection, output projection, self-attention update, and FFN.
  • Memory Overhead: The shared projection is the only extra parameterized component, and the memory overhead is small relative to the target model.For Qwen3.5-35B-A3B, projection input and output activations require about 40 MB and 8 MB at sequence length 2048; decoding temporary activation is below 400 KB for block size 16.
  • Evaluation: DFlash consistently improves over native MTP where available and scales across larger Qwen3.5, Qwen3-Coder, and GPT-OSS models on SGLang.On vLLM with Qwen3.5-9B, it achieves strong speedup at low and medium concurrency while retaining throughput gains at high concurrency.
  • Training Ablations: Applying position-dependent loss decay leads to faster and better convergence than uniform token weighting within each draft block.The ablation emphasizes early-token accuracy during training.
  • Training Ablations: Randomly sampling anchor tokens during training leads to higher acceptance length and better speedup.The comparison uses three-layer draft models, five target hidden features, block size 16, and the 100K training data.
Loading 2602.06036v2…