The Problem
Large language models write code token by token — autoregressive, token by token. GPUs are powerful, but this serial generation leaves them mostly idle with poor utilization.
Speculative Decoding was proposed as a fix: a small model quickly "guesses" a batch of tokens, and the large model verifies them in parallel. Correct guesses save time. EAGLE-3 is a top method in this direction, but its drafter is also autoregressive — guessing 8 tokens requires 8 forward passes, so the cost grows linearly. It can only afford a shallow drafter (1 Transformer layer), which guesses poorly.
DFlash's key insight: a diffusion model can do the guessing, and it can guess a whole block at once.
But there's a trap: a generic diffusion model drafts poorly, because the small model doesn't "understand" context the way the target model does.
Core Insight: "The Target Knows Best"
The target model's hidden layers contain implicit, rich information about future tokens — not an explicit "next word," but contextual signals. DFlash extracts these hidden states and injects them into every layer of the drafter.
This is KV Injection: the target model's Keys and Values are inserted directly into the drafter's KV cache, so the drafter perceives the target's "reasoning" at every layer.
Why it matters: In EAGLE-3, context features enter at the input layer and get diluted with depth, so adding drafter layers yields diminishing returns. With per-layer injection, the acceptance length grows effectively with drafter depth. A 5-layer DFlash generating 16 tokens is faster than 1-layer EAGLE-3 generating 8 tokens — because the diffusion drafter generates a whole block in parallel.
The Three-Layer Architecture
1. Feature Fusion. After the target finishes prefill, DFlash samples hidden states from 5 uniformly spaced layers (shallow layers carry local information, deep layers carry global semantics), concatenates them, and passes them through a lightweight projection layer to get a compact target context feature.
2. KV Injection. The differentiator. The context feature isn't just fed as an input embedding — it's injected into the K and V projections of every layer. The drafter "senses" the target model's internal state throughout its depth.
3. Block Diffusion. Traditional diffusion needs many iterative denoising steps. DFlash needs only one, because the target's hidden features serve as strong conditioning. An entire block (e.g., 16 tokens) is generated in a single forward pass — parallel and efficient.
Training Strategy
1. Random sampling of masked blocks: instead of uniform partitioning, anchor tokens are randomly sampled as block starts — mimicking real inference, where each draft block begins with a verified bonus token from the previous cycle.
2. Loss weighting: earlier tokens in a block matter more because errors cascade. Exponentially decaying weights: w_k = exp(-(k-1)/γ).
3. Shared embedding & LM head: the drafter reuses the target's embedding and output head, training only the intermediate Transformer layers — making the drafter a lightweight adapter, not a standalone model.
Design Philosophy
DFlash's strength isn't using a diffusion model — it's asking the right question: not "how do we make a diffusion model generate better?" but "how do we make it an excellent guesser?" It doesn't compete with autoregressive generation on quality; it only needs to guess fast and accurately, leaving verification to the autoregressive model.
That's the essence of a "parasitic" architecture — not replacement, but symbiosis. The diffusion model provides parallel speed; the autoregressive model provides quality guarantees; KV Injection couples them tightly.
Technical Details
- Layer extraction: 5 layers uniformly sampled from layer 2 to the third-from-last layer of the target model
- Projection layer: lightweight FC + RMSNorm fusing cross-layer information
- Sparse attention training: bidirectional attention within a block, no cross-block information flow (Flex Attention)
- Parameter scale: drafter is only 2B parameters (for the 27B version), sharing the target's embedding and LM head
- Training data: ~800K samples (Nemotron + CodeAlpaca), with responses regenerated by the target model