Skip to content

Architecture

Context

How DFlash 2 works: the one-pass block-diffusion draft inherited from DFlash (Z Lab, ICML 2026), the KV-injection conditioning that keeps a 2B-parameter drafter accurate, and the two DFlash 2 additions — the candidate path selector and the two-tap dynamic convolution — each fixing one measured loss mode of parallel drafting.

Speculative Decoding Context

A speculation cycle has two phases: a small drafter proposes a block of future tokens, and the frozen target LLM verifies the whole block in one forward pass, accepting the longest valid prefix (speculative decoding fundamentals cover the mechanism). Rejection sampling makes the scheme lossless. Speedup reduces to acceptance length x (verification cost / (draft cost + verification cost)): how many tokens survive each cycle, and how cheap the guessing was.

Autoregressive drafters (EAGLE-3, native MTP) need K sequential forward passes to propose K tokens, so their draft cost grows linearly with block size and they run shallow drafters. UCSD's TPU v5p measurements added the hardware punchline, K-flat verification: on datacenter accelerators, verifying 1024 drafted tokens costs nearly the same as verifying 16 (weight loading dominates, not attention math). Draft quality, not width, is the binding constraint — which is exactly the constraint DFlash 2 attacks.

Component Breakdown

Component Role Key Facts
Target LLM (frozen) Verifier; also the context expert e.g. Qwen3.8-27B, Muse-Glimmer-30B; shares embed_tokens and lm_head with the drafter
DFlash drafter backbone One-pass block denoiser ~2B params (Qwen3.8-27B drafter), 5 layers, sliding-window attention (1024-token window), own attention shape independent of target
Block attention mask Non-causal in-block attention Position 0 holds the real anchor token; positions 1..K-1 start as MASK and are denoised in parallel
KV injection path Conditioning on target context Target hidden states pass through the drafter's own KV projection into every draft layer's KV cache — not just the input embedding (the EAGLE-3 way)
Candidate path selector (DFlash 2) Coherence without autoregression — fixes incoherent independent picks (detail below) Top-16 candidates per position, low-rank bilinear pair scoring; +2.0M params, +0.6% cycle latency
Two-tap dynamic convolution (DFlash 2) Suffix-decay fix (detail below) Two taps around every attention and MLP sublayer, mixing each position with its predecessor; +16.5M params (+3%), +0.7% cycle latency
DFlashWorker / DFlashDraftModel SGLang runtime Draft worker drives the scheduler and wraps the target worker for verification passes (PR #22077, then Spec V2 in PR #23000)
Speculators library vLLM runtime Connects the drafter to target hidden states inside the vLLM path; swap EAGLE-3 for DFlash by config only
TPU proposer (tpu-inference) JAX runtime Dual-cache: target keeps paged KV (Pallas kernels), drafter uses static on-device JAX arrays (PRs #1868-1870)

Draft-Verify Cycle

flowchart TB
    subgraph CTX["Context Capture (prefill / last verify)"]
        T0["Target LLM layers<br/>(frozen weights)"]
        HP["Target hidden states<br/>of context tokens"]
        T0 --> HP
    end

    subgraph DRAFT["Drafter — single forward pass"]
        KV["KV injection:<br/>draft KV projection of hidden states<br/>into every draft layer"]
        BB["5-layer block-diffusion backbone<br/>sliding window 1024, non-causal block mask"]
        CONV["Two-tap dynamic convolution<br/>before/after each attn + MLP sublayer"]
        HEAD["Shared frozen lm_head<br/>logits + top-16 candidates per position"]
        KV --> BB --> CONV --> HEAD
    end

    subgraph SEL["Path Selection (sequential, tiny)"]
        SC["Pairwise scorer<br/>U_t(b) + bilinear match A(a), B(b)<br/>gated by context H(h_t)"]
        WALK["Greedy walk or sampling<br/>from last verified token"]
        SC --> WALK
    end

    subgraph VER["Verification"]
        TV["Target LLM forward pass<br/>over draft block + next token"]
        RS["Rejection sampling<br/>(lossless, target distribution)"]
        TV --> RS
    end

    HP --> KV
    HEAD --> SC
    WALK -->|"draft block (K tokens)"| TV
    RS -->|"accepted prefix + bonus token"| CTX
    RS -->|"committed output"| OUT["Detokenizer / client"]

One Cycle, Step by Step

sequenceDiagram
    autonumber
    participant S as SGLang Scheduler
    participant DW as DFlashWorker
    participant DR as DFlash Drafter (2B)
    participant PS as Path Selector (+2M)
    participant T as Target LLM (frozen)

    S->>DW: forward_batch_generation(batch)
    DW->>DR: inject target hidden states (KV projection)
    DR-->>DW: whole block predicted in ONE pass<br/>top-16 candidates per position
    DW->>PS: candidate lists + last verified token
    PS-->>DW: one coherent path (K tokens)
    DW->>T: verify draft block in parallel
    T-->>DW: target logits for every position
    DW->>DW: rejection sampling (lossless)
    DW-->>S: accepted prefix + 1 bonus token

The counterintuitive SGLang detail (useful when reading traces): the draft worker is the one that talks to the scheduler — it wraps and calls the target worker when drafts are ready, not the other way around.

How It Works: The Two DFlash 2 Fixes

Problem 1 — Independent picks are incoherent (selection headroom)

Predicting every position independently means each pick is locally plausible but jointly incoherent — two neighbors may both pick "decoding", and the stutter dies at verification. The evidence that this is a selection problem, not a prediction problem, is in DFlash's own candidate lists (5-layer Qwen3-4B drafter on GSM8K, conditioned on all earlier positions being right):

Metric Pos 0 1 2 3 4 5 6 Acceptance length
Recall@1 85.4% 80.3% 79.4% 78.3% 77.5% 75.9% 72.9% 4.27
Recall@16 99.5% 97.3% 94.8% 92.6% 90.8% 89.4% 87.8% 6.79

An oracle picking from the top-16 would reach 6.79 vs 4.27 — 2.5 tokens of pure selection headroom. DFlash 2's path selector harvests part of it: keep the top-16 candidates per position and score every adjacent pair,

S_t(a, b) = U_t(b) + < A(a) ⊙ H(h_t), B(b) >

where U_t(b) is the drafter's own logit for candidate b, A and B are compact 256-dimensional token embeddings, and H(h_t) is a context gate deciding which parts of the predecessor-match count — in essence a low-rank bilinear attention over adjacent candidates. All pairs score in one parallel shot (no extra backbone or LM-head pass); only the final walk over precomputed scores is sequential (greedy follows the best successor; sampling draws from the same scores; rejection sampling restores the exact target distribution).

Results vs the DSpark-style sequential correction (5-layer Qwen3-4B, GSM8K; overheads relative to plain DFlash):

Method Added params Latency Acceptance @ T=0 @ T=1
DFlash — — 4.27 3.78
+ DSpark correction +77.8M +9.6% 4.49 4.08
+ path selector (DFlash 2) +2.0M +0.6% 4.61 4.25

Roughly 40x fewer parameters and 16x lower latency overhead than the correction-head approach, with higher acceptance. Choosing is cheaper than predicting.

Problem 2 — Suffix decay is local (backbone, not selection)

Even the oracle decays down the block (99.5% → 87.8% across positions in the table above): the candidates themselves run out. Depth fixes it indiscriminately — 3-, 5-, and 15-layer drafters are identical at position 0 and fan apart down the block — but ten extra layers cost +15.2% cycle latency and erase DFlash's efficiency edge.

DFlash's attention shows where the capacity is needed. It has two jobs — read the context before the block, and model dependencies inside it — but the block's share of attention falls from 30% in layer 1 to 8% in layer 5, with the remainder concentrating in a shrinking handful of heads. DFlash 2 splits the jobs: a dedicated module takes the within-block work. A two-tap dynamic depthwise convolution is inserted before and after each attention and feed-forward sublayer:

Conv_k(x)_t = k_t,0 ⊙ x_t + k_t,1 ⊙ x_{t-1}

Each coefficient combines a learned base kernel with a small correction computed from the current hidden state (one correction shared per 16 channels). The first position reads the last verified token's representation; every later position reads its predecessor. Information crosses the block while all positions still compute in parallel — and the module is block-local and stateless, so attention, the LM head, and verification are untouched.

Outcome: the convolutions buy what ten extra layers bought (overheads in the tables above) — five-layer DFlash + convolutions nearly matches 15-layer DFlash on suffix decay. Average within-block attention across layers 4-5 falls from 9.4% to 0.5%: the module has absorbed the local work, and attention returns to reading context. Suffix decay is mostly a local problem.

Combined System

Selector + convolutions together add 1.3% to the draft-verify cycle latency. Per-request mean acceptance length (lossless rejection sampling; default sampling per model, block size 8 for Qwen3.8-27B, 16 for Muse Glimmer):

Target MTP DFlash DSpark DFlash 2
Qwen3.5-4B (mean of 5 tasks) 4.54 4.92 5.49 5.97
Qwen3.8-27B (vs native MTP) 4.28 — 3.62 4.80
Muse-Glimmer-30B (vs official DFlash) — 4.44 4.48 5.70

On Qwen3.5-4B that is +1.05 tokens per pass over DFlash (+21%) and +0.48 over DSpark; on both launch models DFlash 2 averages more than a full token ahead of DSpark. On MATH-500 the gain is visible position by position: DFlash 2 holds ~86% conditional acceptance to the last position while every baseline ends the block 6-9 points lower.

Internals Worth Knowing (Serving Path)

  • Immediate materialization of the draft KV projection. Target latents are projected by the drafter ahead of the draft forward pass rather than stored, preserving KV space and radix-cache prefix sharing. SGLang implements this with a layer-batched linear projection plus a fused Triton kernel for norm+RoPE post-processing.
  • Overlap scheduling (SGLang Spec V2). Host-side cleanup of batch N-1 and KV allocation for batch N overlap with GPU work, cutting host-device synchronization; combining DFlash with Spec V2 improved throughput >33% (11.4 → 15.3 ktok/s, Qwen3-8B on B200 at concurrency 32).
  • TPU dual-cache. DFlash's non-causal block diffusion is incompatible with paged attention, so the tpu-inference port runs the target on paged KV (Pallas kernels) and the drafter on static on-device JAX arrays; a metadata rework fixed "sequence length inflation" where draft state drifted from the target's accepted-token count.

Benchmarks

H200 (SGLang, Qwen3.8-27B, block size 8 = 7 draft tokens)

FlashAttention 3 for target and draft; model-default sampling (temperature 1.0, top-p 0.95, top-k 20), xhigh reasoning effort, 4096 max new tokens. Throughput = output tokens / wall time, speedup vs autoregressive:

Concurrency GSM8K MATH-500 HumanEval MBPP MT-Bench
1 236.1 tok/s (3.43x) 230.7 (3.34x) 214.6 (3.11x) 226.9 (3.29x) 184.0 (2.67x)
8 1,328.7 (2.84x) 1,368.3 (2.85x) 1,291.5 (2.67x) 1,328.0 (2.78x) 1,090.2 (2.27x)
32 1,922.5 (1.45x) 1,951.8 (1.30x) 1,799.0 (1.16x) 1,886.8 (1.25x) 1,525.3 (1.01x)

Honest scaling picture: native MTP at concurrency 32 falls below 1x on several tasks (0.77-0.94x); DFlash 2 degrades most gracefully but still compresses toward parity. All speculative methods proposed 7 draft tokens per step in this comparison.

NVIDIA Blackwell (TensorRT-LLM, DFlash v1 drafter)

gpt-oss-120b on 8x DGX B300 (Blackwell Ultra), SPEED-Bench: >15x higher throughput at 500-600 tok/s/user interactivity vs autoregressive, 1.5x higher than EAGLE-3 at the same point; batch size 1 more than doubles interactivity. Interactivity speedups at matched concurrency — gpt-oss-120b: DFlash 2.3x avg vs EAGLE-3 1.7x; Llama-3.1-8B: DFlash 2.8x vs 2.2x. Single-GPU: Gemma-4-31B on vLLM/B300 up to 5.8x (MATH-500); Qwen3-8B on SGLang/B200 5.1x (MATH-500).

Google TPU v5p (JAX, vLLM tpu-inference)

Average 3.13x tokens/s across datasets (peak ~6x on math; MATH-500 8.02 → 1.40 ms/token, Qwen3-4B, K=16, greedy). Head-to-head, Llama-3.1-8B out-of-the-box checkpoints: DFlash 2.29x vs EAGLE-3 1.30x end-to-end serving speedup. Scaling theory: K=16 already captures >90% of the theoretical maximum; K from 16 to 128 would add less than one accepted token per step; improving per-position acceptance is 2-3x more valuable than growing K.

Ablations (LMSYS, Qwen3-4B-class drafters, acceptance / speedup)

Configuration GSM8K HumanEval MT-Bench
EAGLE-3 (5-layer) 4.2 / 2.1x 4.3 / 2.2x 3.1 / 1.4x
DFlash (both techniques) 4.2 / 3.3x 4.0 / 3.2x 3.0 / 2.2x
DFlash diffusion only 3.5 / 2.9x 3.5 / 2.9x 2.6 / 2.0x
DFlash injection only 4.8 / 2.4x 4.6 / 2.3x 3.4 / 1.5x

Diffusion drafting buys speed at equal acceptance; KV injection buys acceptance that survives draft depth. Together they dominate.

Source Discrepancies

  • NVIDIA's "15x" is throughput at matched interactivity on an 8-GPU Pareto curve, not a batch-1 latency speedup; the blog's 2.7-3.4x (H200, batch 1) and NVIDIA's number describe different operating points. Provenance caveats (release-date and attribution discrepancies) are annotated in the index Sources.

Sources