DFlash

This post records my notes on DFlash.

Background

1. Diffusion LLMs

Diffusion models originated in image generation. Successive rounds of the forward process add noise to an image until it approaches Gaussian noise; the reverse process is trained to remove that noise step by step and thereby learn image generation.

Applying this method to LLMs produces a generation paradigm different from autoregression: text is created through successive diffusion steps.

I may expand the dLLM discussion in an earlier post when I have time.

2. Speculative Decoding

Speculative decoding uses a lightweight draft model to generate candidate tokens and a target model to verify them. It is an elegant form of lossless acceleration.

Representative earlier methods include Medusa, which generates in parallel with multiple heads, and EAGLE, arguably the ceiling of autoregressive drafting. See my earlier posts for details.

3. Discussion

As a parallel model, Medusa is limited because draft tokens after the first lack the preceding drafted context, constraining the acceptance rate.

EAGLE instead uses autoregression, feeding the predicted token and hidden state into a small draft model to improve acceptance. It can be viewed as reaching the ceiling of autoregressive methods.

Why?

In autoregressive speculative decoding, speedup $\eta$ is proportional to acceptance rate $\alpha$: $\eta = \frac{\alpha \gamma t_{target}}{T_{draft} + t_{target}}$. Since $T_{\text{draft}} = \gamma \cdot t_{step}$, we obtain $\eta = \frac{\alpha \gamma t_{target}}{\gamma \cdot t_{step} + t_{target}}$.

Improving speedup requires reducing per-step generation time or increasing the expected accepted draft length. These goals generally trade off against one another, which fundamentally limits autoregressive methods.

To generate drafts faster, DFlash borrows from dLLM generation. It is parallel—all results come from one forward pass—yet the generated draft tokens are not logically independent.

Block diffusion

The most important component of DFlash is its draft model, implemented with block diffusion.

Besides applying diffusion by placing and then removing a fixed number of masks, the diagram highlights the crucial feature: KV-cache injection.

KV injection

As in ordinary speculative decoding, while the target model prefills the sequence, the draft model performs its own prefill to build a cache.

The draft and target models have different QKV matrices, so they cannot share a KV cache; each model retains its own.

To share the more informative target-model representation with the small model, DFlash extracts hidden states from several target layers, compresses them through an FFN to the draft model’s hidden dimension $H_d$, projects them with the draft model’s KV matrices, and prepends the result to the draft model’s own KV cache as context.

This adds computation and cache overhead, plus some cost during later attention. However, Q computation, the $O_i$ projection, and the following FFN are unaffected. The acceptance-rate benefit must outweigh these costs.

limitation

DFlash appears to achieve better acceleration through diffusion. For long contexts, however, it must read $N$ positions once per denoising step. The resulting memory bottleneck sharply reduces speedup and can make it slower than the vanilla approach.

Its memory footprint also becomes large, and performance under concurrency is not especially strong.

Sparse attention for the draft model—or even a cache-free draft model—may therefore have substantial potential for long contexts.