Large language models generate code and text token by token—autoregressive, one token at a time. GPUs are powerful, but this serial generation leaves most of their compute idle.
Speculative decoding addresses this: a small model quickly 'guesses' a sequence of tokens, and the large model verifies them in parallel. EAGLE-3 is a top method in this direction, but its drafter is also autoregressive—guessing 8 tokens takes 8 passes, so cost grows linearly. The drafter must stay shallow (1 Transformer layer) and guesses are imprecise.
DFlash's key insight: a diffusion model can do the guessing, and guess a whole block at once.
But there's a trap: a naive diffusion drafter performs poorly, because it doesn't 'understand' context the way the target model does.
DFlash's core principle: the target knows best.
The target LLM's hidden layers contain implicit, rich information about future tokens—not an explicit 'next word', but contextual representations. DFlash extracts these hidden states and injects them into every layer of the drafter.
Not just a hint at the input layer—at every layer. This is KV Injection: the target's Keys and Values are written directly into the drafter's KV cache, so the drafter 'sees' the target's reasoning at each layer.
What does this change?
EAGLE-3's problem: context enters at the input layer and dilutes with depth, so adding layers yields diminishing returns. DFlash injects context at every layer, so acceptance length (tokens accepted per cycle) grows effectively with drafter depth. A 5-layer drafter beats a 1-layer one significantly—and latency is even lower, because the diffusion model generates an entire block of tokens in parallel.
A concrete number: a 5-layer DFlash generating 16 tokens is faster than 1-layer EAGLE-3 generating 8 tokens.
The three-layer architecture
1. Feature Fusion
After the target finishes prefill, DFlash samples hidden states from 5 uniformly spaced layers (layer 2 through the third-from-last). Shallow layers carry local information; deep layers carry global semantics. These are concatenated and passed through a lightweight projection layer (FC + RMSNorm) to produce a compact target context feature.
2. KV Injection
This is what distinguishes DFlash: instead of feeding the context feature as an embedding, it is injected into the K and V projections of every layer. The analogy: an ordinary drafter only sees your first sentence; DFlash's drafter senses your entire train of thought while writing each paragraph.
3. Block Diffusion
Traditional diffusion models denoise over many iterative steps. DFlash needs only one step—because the target's hidden features serve as conditioning, multi-step guessing is unnecessary. One block (e.g., 16 tokens) is generated in a single forward pass.
Training strategy
1. Random sampling of masked blocks: block starts are randomly sampled anchor tokens rather than uniform splits—mirroring real inference, where each draft block starts from the previous cycle's verified bonus token. 2. Loss weighting: earlier tokens in a block matter more, since 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—a lightweight adapter, not a standalone model.
An intuition check (Feynman style)
Imagine an exam. A top student (the target model) has already worked out the answer. You (the drafter) must quickly write a draft for them to check.
Ordinary methods have you write from your own understanding, often wrong. DFlash lets the top student 'transmit' their thinking—not the answer, but the reasoning. You sense their train of thought as you write each character, so your draft closely matches what they had in mind, verification is fast, and overall speed improves.
Design philosophy
DFlash's brilliance 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 a diffusion model an excellent guesser?"
It doesn't compete with autoregressive models on generation quality—it only needs to guess fast and accurately; verification is left to the target. This is the essence of a 'parasitic' architecture: not replacement, but symbiosis. The diffusion model provides parallel speed; the autoregressive model provides quality guarantees; KV Injection tightly couples them.
Smart architecture design.
Technical details
- Layer extraction: 5 layers uniformly sampled from layer 2 to the third-from-last of the target model
- Projection: lightweight FC + RMSNorm fusing cross-layer information
- Sparse attention during training: bidirectional attention within a block, no information flow across blocks (Flex Attention)
- Parameter scale: drafter is only ~2B parameters (27B target version), sharing the target's embedding and LM head
- Training data: ~800K samples (Nemotron + CodeAlpaca), with responses regenerated by the target model