DFlash: Block Diffusion for Flash Speculative Decoding

In plain words

Language models usually write one word-piece at a time, and each step waits for the one before it. Speculative decoding speeds this up with a helper: a small model guesses several upcoming pieces, and the big model checks the whole guess in a single pass, keeping the part it agrees with.

DFlash changes the helper. Instead of guessing one piece after another, it fills in a block of blank slots all at once, the way diffusion models fill in masked text. Its main trick is that it gets to look at the big model's internal state while it guesses, not only the last word. Because the big model still checks every piece, the final text is the one it would have produced anyway, only faster.

The authors report that this is several times faster than plain decoding and clearly faster than a strong earlier helper method.

TL;DR

• Speculative decoding with a non-autoregressive drafter: one parallel forward pass proposes a block of tokens §4.1.
• The draft is a small block-diffusion transformer; target-model hidden features enter every draft layer as extra keys and values §4.1 app. A.3.
• Reported on Qwen3, LLaMA-3.1 and others with the paper's own draft models; author-reported, Transformers and SGLang backends §5.
• Official code is MIT-licensed, but it contains inference only: the training recipe cannot be checked against code model.py.

Problem

Autoregressive decoding is sequential and memory-bound, so latency and GPU utilisation suffer §1. Speculative decoding with a small drafter helps, but the strongest autoregressive drafter the paper cites (EAGLE-3) still drafts one token per step, which the authors say caps speedups at about 2–3× §1. Diffusion drafters avoid that serial cost, but previously proposed diffusion drafters used very large models (e.g. 7B parameters), and small autoregressive imitations of parallel drafting were limited by weak capacity §1. The question is whether a small, parallel drafter can still propose high-quality tokens §1.

Assumptions: the target model is a fixed autoregressive LLM whose hidden states are available; verification is done by that target §1 §3.1.

Method

Inference (§4.1)

1. Context features. The target runs a normal prefill over the prompt and emits the first token. Hidden states are taken from a fixed set of layers spaced evenly from shallow to deep, concatenated, and projected once to one compact context feature. The paper calls this the target context feature §4.1. The projection is a linear map followed by RMS normalisation app. A.3.
2. KV injection. The context feature is not added to the draft's input. In every draft layer it is projected into keys and values, alongside the draft tokens' keys and values. Queries come only from the draft tokens. Context entries skip the query projection, the output projection, attention updates and the feed-forward block, so they act as extra memory for the draft §4.1 app. A.3 fig. 2. Stored context entries are reused in later drafting rounds §4.1.
3. Parallel block drafting. Each round starts from an anchor token, which is the target's last accepted or bonus token. The anchor is followed by block-size minus one mask tokens. The masked positions are embedded with the target's embedding table and decoded together in one draft forward pass §4.1 §4.2.
4. Verification. The target runs one forward pass over the whole block. The accepted prefix and the target's own next token (the bonus token) are kept. The new anchor is the bonus token §3.1 §4.1.
5. Lossless. Verification by the target keeps the output distribution of the target model, which the authors describe as lossless §1.

Drafting cost does not grow with the block size for moderate blocks, because the draft runs once per round rather than once per token §3.2 fig. 3. Deeper drafts give higher acceptance length but cost more per round §3.2 tab. 6.

Training (§4.2)

1. Features from the clean sequence. The frozen target runs once over the full clean sequence, prompt plus response, to extract and fuse all features. The draft is trained on these features §4.2 fig. 4. Training responses are generated by the target model §5.
2. Anchor-based blocks. Random response positions are chosen as anchors. Each anchor starts a block, and the remaining block positions are masked. The draft learns to predict the next block-size minus one tokens in parallel, matching inference §4.2. The paper states that standard block diffusion training divides the response uniformly and masks random positions inside each block; DFlash replaces this §4.2.
3. Anchor count and resampling. A fixed number of anchors is drawn per sequence, and the anchor positions are resampled at every epoch. This bounds training cost on long sequences and augments the data §4.2 app. A.1.
4. Packed blocks with a sparse mask. All blocks of a sequence are concatenated into one sequence. Tokens attend bidirectionally inside their own block and to the context features, and attention across blocks is disallowed. This lets many blocks train in one forward and backward pass with FlexAttention §4.2 fig. 4.
5. Loss weighting. The cross-entropy for a token at position k of a block is multiplied by an exponentially decaying weight, so early positions count more, because an early error invalidates every later token in the draft §4.2 eq. 4.
6. Frozen embeddings and head. The draft shares the target's token embedding and output head and keeps them frozen. Only the draft's transformer layers are trained §4.2.

Ablations: random anchor sampling beats standard uniform block construction on all three tasks tested tab. 13; KV injection beats input-only fusion in both autoregressive and block-diffusion drafting tab. 9; loss decay converges faster and to a higher acceptance length fig. 5 §4.2.

Key equations:

L = \frac{T_{\text{draft}} + T_{\text{verify}}}{\tau} eq. 1

w_k = \exp\!\left(-\frac{k-1}{\gamma}\right) eq. 4

Equation 1 gives per-token latency, where τ counts accepted tokens including the bonus token §3.1. Equations 2 and 3 give autoregressive and parallel drafting costs §3.2. The KV projection equations in the appendix are unnumbered app. A.3.

Pseudocode for one greedy decoding round, from the paper's inference description §4.1 §3.1:

# One DFlash decoding loop, greedy case  [§4.1]
# Input:  target LLM, draft model, prompt token ids, max_new_tokens
# Output: generated token ids

def dflash_decode(target, draft, prompt, max_new_tokens):
    # Prefill: target gives the first token and its hidden states  [§4.1]
    first_token, hidden = target.prefill(prompt)
    # Fuse the chosen layers' features into one context feature  [app. A.3]
    context = draft.fuse(select_layers(hidden))    # (n_ctx, d)
    generated = [first_token]
    anchor = first_token
    while len(generated) < max_new_tokens:
        # Anchor plus B-1 mask tokens, embedded by the target  [§4.2]
        noise = target.embed([anchor] + [MASK] * (block_size - 1))
        # One parallel draft pass; keep only the mask positions  [§4.1]
        drafts = argmax(target.head(draft(noise, context)[1:]))
        # Target scores the whole block in one forward pass  [§3.1]
        posterior, verified = target.verify([anchor] + drafts)
        # Count drafts that match the target's greedy choice  [§3.1]
        accepted = count_leading_matches(drafts, posterior[:-1])
        # Keep the accepted drafts and the target's bonus token  [§3.1]
        bonus = posterior[accepted]
        generated += drafts[:accepted] + [bonus]
        # Add verified features to the context for the next round  [§4.1]
        context = concat(context, verified[:accepted + 1])
        anchor = bonus
    return generated[:max_new_tokens]

Pseudocode for one draft layer, with the context entering only as keys and values app. A.3:

# One draft layer with target features as extra KV  [app. A.3]
# Input:  draft_states (B, d), context (n_ctx, d); B = block size
# Output: draft_states (B, d)

def draft_layer(draft_states, context, layer):
    # Queries come from the draft tokens only  [app. A.3]
    queries = layer.w_q(draft_states)
    # Context and draft tokens both give keys and values  [app. A.3]
    keys = concat(layer.w_k(context), layer.w_k(draft_states))
    values = concat(layer.w_v(context), layer.w_v(draft_states))
    # Attention over the context-plus-block sequence  [§4.1]
    mixed = attention(queries, keys, values)
    # Residual update of the draft tokens only  [app. A.3]
    hidden = draft_states + layer.w_o(mixed)
    return hidden + layer.feed_forward(hidden)

Pseudocode for the training loss, from the paper's description §4.2 eq. 4 app. A.1:

# One training step of the draft model  [§4.2] [app. A.1]
# Input:  clean sequence = prompt + target-generated response
# Output: scalar loss

def training_loss(target, draft, sequence, response, block_size):
    # Frozen target forward over the clean sequence  [§4.2]
    hidden = target.hidden_states(sequence)
    context = draft.fuse(select_layers(hidden))
    # Random anchors in the response, a fixed count per sequence  [app. A.1]
    anchors = sample(response, num_anchors)
    # Each block: clean anchor, masked remainder  [§4.2]
    blocks = [make_block(sequence, a, block_size) for a in anchors]
    # Pack blocks into one sequence; attention stays inside each block [§4.2]
    preds = draft(pack(blocks), context, block_mask(blocks))
    # Decay the weight of later block positions  [eq. 4]
    weights = exp(-(block_positions - 1) / gamma)
    return weighted_cross_entropy(preds, targets, weights)

Hyperparameters

Results

Author-reported. Setup: Transformers backend, thinking mode disabled, maximum 2048 generated tokens, NVIDIA H200 unless noted, speedup over the autoregressive baseline §5.1 tab. 1. Each cell is speedup / average acceptance length τ.

Table 1 tab. 1: Qwen3 models, thinking disabled

Parenthesised numbers are the EAGLE-3 draft-tree size, or the DFlash block size tab. 1 caption. The EAGLE-3 checkpoints for Qwen3 are from AngelSlim, a third party, and the LLaMA comparison uses the official EAGLE-3 checkpoint §5.

Table 2 tab. 2: Qwen3 models, thinking enabled (DFlash only)

Table 3 tab. 3: SGLang throughput (tok/s), FA4 backend, single B200. Each block gives baseline throughput, DFlash throughput, the speedup and average τ.

Table 4 tab. 4: long-context acceptance length, Qwen3.5-27B base draft vs long-context fine-tune (τ; columns: base / long)

Table 5 tab. 5: LLaMA-3.1-8B-Instruct, SGLang with FlashInfer, single B200. Speedups per concurrency level; baseline row gives absolute tok/s.

Table 6 tab. 6: number of draft layers (block 16, 5 hidden features; speedup / τ)

Table 7 tab. 7: number of target hidden features (block 16, 3 draft layers; speedup / τ)

Table 8 tab. 8: training versus inference block size (8 layers, 5 features; speedup / τ)

Table 9 tab. 9: target-feature conditioning (Qwen3-4B, 5-layer drafts, block 8; τ / speedup)

Table 10 tab. 10: draft without target features (appendix; speedup / τ; no context features, 5 layers)

Table 11 tab. 11: more target models, SGLang, concurrency 8 (τ / speedup; MTP is the native multi-token prediction the paper compares against app. A.4)

Table 12 tab. 12: vLLM, Qwen3.5-9B (DFlash tok/s, speedup in parentheses)

Table 13 tab. 13: anchor sampling vs standard blocks (3 layers, 5 features, block 16; speedup / τ)

Figures, described only fig. 1 fig. 2 fig. 3 fig. 4 fig. 5: fig. 1 plots speedup against autoregressive decoding for DFlash and EAGLE-3 on Qwen3-8B with the Transformers backend. fig. 2 shows the inference design, with context features injected into each draft layer's key-value cache. fig. 3 compares draft latency of 1, 3 and 5-layer DFlash with a 1-layer EAGLE-3. fig. 4 is the training attention pattern: clean prompt and response, anchors, mask tokens, and the blocked cross-block attention. fig. 5 shows acceptance length per epoch with and without loss decay.

Paper vs code

Code: https://github.com/z-lab/dflash at code_commit 07ebd93, MIT licence; the paper links this repository [abstract]. Files read: dflash/model.py and README.md in full, dflash/cli.py for defaults, and dflash/benchmark.py by search only.

• No training code. The repository contains inference, the CLI, benchmarks and an MLX port. It has no loss, anchor sampler, optimiser or data pipeline. A search for loss, anchor, AdamW, learning-rate and decay terms in dflash/*.py found none. The training recipe (anchor sampling, eq. 4 and the settings in app. A.1) therefore cannot be checked against code. model.py covers inference only.
• Context layers match the text. Paper: five layers between the second and the third-to-last layer §5. Code: start index 1, end index L-3, spaced evenly model.py. extract_context_feature adds one because hidden state 0 is the embedding output model.py. If a checkpoint config sets target_layer_ids, that list overrides the formula model.py.
• KV injection matches, with extras. Code projects target features and draft states with the same key and value weights and concatenates them model.py. It also applies per-head RMS normalisation to queries and keys model.py and rotary embeddings to context keys at their absolute positions model.py. The paper's appendix equations include neither app. A.3.
• Context fusion matches. A linear map to the hidden size, then RMS normalisation model.py matches app. A.3.
• Embedding and head. The code uses the target's input embeddings multiplied by input_embedding_scale model.py, and its output head multiplied by output_multiplier with optional tanh soft-capping model.py. The defaults are 1.0 and no capping model.py. The paper mentions neither mechanism §4.2.
• Drafted tokens match. The anchor's output is dropped and the B-1 mask positions are decoded, greedily at temperature 0 model.py. This matches §4.2.
• Attention pattern. Full-attention draft layers are non-causal, so the block attends bidirectionally to itself and the context model.py §4.1. The code also supports per-layer sliding-window attention from layer_types model.py. The paper does not describe it.
• Acceptance length counts the bonus token. produced = accepted + 1 is what the benchmark averages model.py, matching the definition in §3.1.
• Sampling rule. At temperature 0, a draft is accepted while it equals the target's argmax; the bonus token is the target's argmax after the accepted prefix model.py. At temperature above 0, speculative rejection sampling with residual resampling is used model.py. Top-p and top-k apply to both draft and target distributions model.py. The paper's text does not spell out this rule.
• Maximum length. The CLI defaults to 2048 new tokens and greedy decoding cli.py, matching the setting in tab. 1.
• DFlash 2 is a different model. DFlash2DraftModel adds grouped dynamic causal convolutions and a candidate selector model.py. Neither appears in this paper. Do not implement them as DFlash.

Pitfalls

• Headline numbers are best cases. "Over 6x" in the abstract and "6.1x" in the introduction are single-task peaks on MATH-500 §1 tab. 1. Seven-task greedy averages are 4.9x for both Qwen3 models, and temperature-1 averages are 4.2x and 4.0x tab. 1. Chat (MT-Bench) speedups are much lower, 2.85x and 2.75x at greedy tab. 1.
• The EAGLE-3 ratio is attributed to the wrong tree size. The text gives 2.4x for greedy decoding "over EAGLE-3 (16)" and 2.2x at temperature 1, without a tree size §5.1. The numbers match the tree-size-60 averages in tab. 1: 4.91/2.08 and 4.86/2.02 for greedy, and 4.24/1.93 and 4.03/1.88 at temperature 1. Against tree size 16 the ratios are about 2.7x and 2.8x (greedy) and 2.5x and 2.4x (temperature 1).
• Reasoning averages. §5.2 says "roughly 4.5x and 3.9x" §5.2. The three-task averages in tab. 2 are about 4.4x (greedy) and about 3.8x (temperature 1). The paper does not say which temperature the rounded figures refer to.
• Naive-draft range. §4.1 says draft models without context features give speedups "typically around 2–3x" §4.1. tab. 10 shows 2.65x to 3.73x.
• Identical τ in Table 3. Qwen3-4B and Qwen3-8B on MATH-500 both list τ = 8.01 tab. 3. This may be a copy error; the paper gives no way to check it.
• Serving tables are not comparable with Table 1. tab. 1 uses Transformers on H200; tab. 3 uses SGLang with FA4 on B200 §5 §5.3. Do not compare speedups across these tables.
• Speedups fall with concurrency. Qwen3-4B on MATH-500 drops from 4.8x at one request to 2.9x at 32 tab. 3. The paper says large blocks cost more to verify when compute-bound §5.5.4.
• Block size transfers in one direction only. A model trained at 16 works at 8 but a model trained at 8 does not work well at 16 §5.5.4 tab. 8. The paper leaves adaptive block-size scheduling for future work.
• Token count per round. §3.2 says a 5-layer draft "generating 16 tokens" §3.2. §4.2 and the code draft B-1 = 15 tokens per round, plus the anchor §4.2 model.py. Use γ = B-1 in eq. 1–3 unless the paper says otherwise.
• Index k in eq. 4 is ambiguous. The paper does not say whether k = 1 is the anchor or the first masked position §4.2 eq. 4. With no training code in the repository, choose one and record it.
• Loss normalisation not stated. The paper does not say whether the weighted cross-entropy is summed or averaged over positions, blocks or the batch §4.2.
• Block packing not fully specified. The paper says attention across blocks is disallowed §4.2 fig. 4 but not whether blocks may overlap in response positions, or how many blocks fit a sequence beyond the fixed anchor count app. A.1.
• Context KV cache grows with length. The appendix quantifies the projection and per-round activations only app. A.3. The context entries are stored in each draft layer's KV cache for the whole sequence §4.1 model.py, and that memory is not quantified. The 42 MB figure is for Qwen3.5-35B-A3B, not the Qwen3 models app. A.3.
• Feature indexing. Hidden states from Hugging Face-style outputs include the embedding output at index 0. Index target layers with the offset in model.py, not directly.
• Positional encoding of context features. The paper does not say how context entries are positioned. The code gives them the absolute positions of their tokens model.py.
• Baseline checkpoints. EAGLE-3 checkpoints for Qwen3 come from AngelSlim, a third party §5. The LLaMA comparison uses the official EAGLE-3 checkpoint, and the paper trains DFlash on the same data for that comparison §5.5.1.
• Long-context fine-tune. The learning rate and batch size are not stated §5.4. The base draft degrades beyond 4K tokens tab. 4.
• Models beyond Qwen3 and LLaMA. Qwen3.5, Qwen3-Coder-Next and GPT-OSS results appear only in the serving and vLLM tables app. A.4 tab. 11 tab. 12. Their training recipes are not described beyond the general settings above.
• Lossless claim. The paper calls the output lossless §1. Greedy decoding gives exactly the target's greedy output, and sampling uses rejection sampling in the code model.py. The paper's sections read here do not give that argument in full.

Lineage

• Builds on: Speculative decoding (Leviathan et al.), which introduced the draft-and-verify framework §1.
• Builds on: EAGLE-3, arXiv:2503.01840, the strongest autoregressive drafter the paper compares against. DFlash keeps the idea of using target features but injects them as keys and values in every draft layer rather than fusing them into the input §1 §4.1.
• Builds on: block diffusion LMs, arXiv:2503.09573. DFlash modifies their training by sampling anchors as block starts rather than masking uniformly §4.2.
• Builds on: Samragh et al., arXiv:2507.11851, which the paper cites for the observation that target hidden states encode information about future tokens §1 §4.1.
• Related, not compared: PARD, arXiv:2504.18583, small autoregressive models trained to imitate parallel drafting §1. Diffusion-based drafters such as DiffuSpec (arXiv:2510.02358) and SpecDiff-2 (arXiv:2511.00606) are not compared because no open-source implementation was available §1 §5.
• Compared against: EAGLE-3 (tree sizes 16 and 60) and autoregressive decoding §5.