This document is relevant for: Trn2, Trn3

DFlash Speculative Decoding in the vLLM Neuron Framework#

What is DFlash Speculative Decoding?#

This document focuses on how DFlash differs from the EAGLE3 design. It assumes familiarity with the shared speculative decoding flow, including target verification, rejection sampling, and scheduler integration.

DFlash is a block-parallel draft method. Like EAGLE3, it consumes auxiliary hidden states from intermediate target layers and uses the shared target verification and rejection-sampling flow. EAGLE3 feeds those states into its draft head and generates drafts through dependent autoregressive iterations. DFlash instead projects the states directly into each draft layer’s KV cache, then evaluates a fixed [bonus, mask, ..., mask] query block in one non-causal forward. Each query position attends to the injected context and the rest of the block, allowing the borrowed target LM head to produce all K draft tokens in parallel.

Net trade-off: DFlash swaps EAGLE3’s recurrent decoder iterations for one wider, heavier block evaluation plus a KV-injection precompute. The balance is workload-dependent.

API Usage and Configuration#

Enabling DFlash Speculative Decoding#

DFlash is selected through speculative_config:

speculative_config = {
    "method": "dflash",
    "model": "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat",
    "num_speculative_tokens": 9,
}

Pass this to LLM(..., speculative_config=speculative_config) or vllm serve --speculative-config. For serving commands, benchmarking, and acceptance metrics, see the DFlash tutorial.

Configuration Parameters#

speculative_config fields:

Parameter

Required

Description

method

Yes

Must be "dflash".

model

Yes

Path or HuggingFace ID for the DFlash draft checkpoint. The checkpoint must be paired with the target model — it defines the target-layer mapping and relies on the target’s vocabulary, token embedding, and LM head.

num_speculative_tokens

Yes

Number of draft tokens to propose per step. Must equal the checkpoint’s draft width, block_size - 1, because one block slot is reserved for the bonus token. For the reference checkpoint, block_size = 10, so this must be 9.

Architecture#

This section describes how DFlash speculative decoding is architected within the vLLM Neuron plugin and how the components connect.

Draft Model#

In contrast to EAGLE3’s single-decoder-layer drafter, the DFlash draft model contains:

  • A small Qwen3 decode layer stack (five layers for the reference checkpoint) regardless of target model family

  • An FC layer that projects the concatenated auxiliary hidden states into one hidden-size context representation per token before KV projection

  • Token embedding and LM head aliased from the target model

Component Overview#

The scheduler, target verification path, and rejection sampler are shared with EAGLE3. DFlash changes the proposal path after target verification:

Target forward
  ├── concatenated auxiliary hidden states ───────────────────┐
  └── target token predictions                                │
               │                                              │
               ▼                                              │
       Rejection sampler                                      │
               │                                              │
               └── accepted-token output ─────────────────────┤
                                                              ▼
                                                       DFlashProposer
                                                              │
                                                              ▼
                                                       Compiled proposal graph
                                                         1. Correct positions and sequence lengths
                                                         2. Build paged slots and [bonus, mask, ..., mask]
                                                         3. Project auxiliary states into draft-layer KV
                                                         4. Run the non-causal block decode over the injected target KV context
                                                         5. Apply the target LM head
                                                              │
                                                              ▼
                                                       K drafts for the next target step

Key source files:

File

Role

vllm_neuron/vllm/spec_decode/dflash.py

DFlashProposer: loads and compiles the draft model, prepares optional async inputs, and invokes the compiled graph

vllm_neuron/model/qwen3/dflash_model.py

DFlashQwen3ForCausalLM: prepares proposal tensors on device, injects context KV, and runs the block forward

vllm_neuron/model/qwen3/config.py

DFlashConfig: block size, mask token id, target layer ids

vllm_neuron/nn/rejection_sampler.py

On-device greedy rejection sampler (shared with EAGLE3)

vllm_neuron/vllm/worker/neuron_model_runner.py

Integration: orchestrates target verification, rejection, and draft proposal; aliases target weights; supplies and remaps async correction inputs

vllm_neuron/model/registry.py

Model registry (maps "DFlashDraftModel" to DFlashQwen3ForCausalLM)

Implementation Details#

DFlashProposer#

The DFlashProposer class (dflash.py) uses the same load, compile, graph capture, and warmup lifecycle as EagleProposer.

On the execution path, propose() moves inputs to the draft device, supplies shape-stable no-op tensors when prior async state is unavailable, and dispatches one compiled DFlashQwen3ForCausalLM.forward() call.

DFlash Draft Model Architecture#

The draft model (dflash_model.py) is a compact Qwen3-based drafter.

DFlashQwen3DecoderLayer:

A standard Qwen3 decoder layer (attention + MLP with pre-norms), but its attention runs in a non-causal decode mode: self_attn.forward_decode() is called with an explicit non-causal attention_mask.

DFlashQwen3Model (backbone):

  • A small, checkpoint-configured number of DFlashQwen3DecoderLayers form the decoder stack.

  • precompute_and_store_context_kv() combines the concatenated target auxiliary states and projects them into the paged KV cache for every draft layer.

  • forward() embeds the [bonus, mask, ..., mask] query block with the target’s embedding table and evaluates it against the injected context using non-causal attention.

DFlashQwen3ForCausalLM (wrapper):

Its forward() captures four DFlash-specific stages in one compiled call:

  • Async state correction: Correct optimistic positions and sequence lengths using previous- and current-step rejection counts. See Async Rejection and Position Correction.

  • Context and block placement: Map each target auxiliary-state row to its paged KV position and place the K + 1 query-block positions after each post-verification sequence.

  • Query block construction: Use the last valid rejection-sampler output as the bonus token and append K mask tokens. See Draft Block Construction.

  • Block proposal: Inject the target auxiliary states into each draft layer’s KV cache, evaluate the query block with non-causal attention, apply the target LM head, and assemble the bonus-prefixed and drafts-only outputs.

def forward(...):
    # Correct async positions and sequence lengths, then build proposal inputs.
    (input_ids, positions, draft_attn_metadata, keep_idx,
     context_positions, context_slot_mapping,
     bonus_token_ids) = self._prepare_proposal_inputs(...)

    # Inject target auxiliary states into every draft layer's KV cache.
    self.model.precompute_and_store_context_kv(
        target_hidden_states, context_positions, context_slot_mapping)

    # Evaluate and sample all draft positions.
    hidden_states = self.model(input_ids, positions, draft_attn_metadata)
    drafts_hidden = torch.index_select(hidden_states, 0, keep_idx)
    drafts_only = self.compute_draft_tokens(drafts_hidden).view(
        bonus_token_ids.shape[0], -1)
    draft_token_ids = torch.cat([bonus_token_ids[:, None], drafts_only], dim=1)
    return draft_token_ids, drafts_only

_prepare_proposal_inputs() implements the first three stages as device tensor operations that execute as part of the drafter forward NEFF.

compute_draft_tokens() applies the borrowed LM head (ColumnParallelLinear) and greedy sampling. When debug_logits_dir is configured, forward() also returns raw logits for capture.

Target Auxiliary Hidden State Extraction#

DFlash consumes auxiliary hidden states from several intermediate target layers. The layer ids come from the draft checkpoint config (DFlashConfig.target_layer_ids, e.g. [1, 8, 15, 22, 29]), read either directly or from a nested dflash_config.

There is a crucial off-by-one between the checkpoint’s convention and Neuron’s capture point:

  • DFlash checkpoints store HuggingFace output hidden-state layer ids — the state after layer i.

  • The Neuron target captures pre-layer states — the state before a layer.

So the runner must request each stored id plus one. This is handled in DFlashProposer.get_target_aux_hidden_state_layers():

def get_target_aux_hidden_state_layers(self) -> tuple[int, ...]:
    target_layer_ids = self._target_layer_ids_from_config(...)  # e.g. [1,8,15,22,29]
    return tuple(t + 1 for t in target_layer_ids)               # -> (2,9,16,23,30)

For a reduced target used only for profiling, if those configured capture layers do not fit, DFlash instead captures the same number of evenly spaced pre-layer states. The target must still have at least that many layers, but note that performance will regress due to differences with the config’s original intended target layer indices.

The runner passes these layers to the target model’s aux-capture interface (the same interface EAGLE3 uses, neuron_model_runner.py:1330-1348). During the target forward, the captured states are concatenated inside the target NEFF and returned as a single [num_ctx, n_target * hidden_size] device tensor named aux_hidden_states.

Context KV Injection (Precompute)#

precompute_and_store_context_kv() is the KV Injection step that improves the quality of DFlash block-parallel draft tokens by enriching the draft model’s context using target aux hidden states:

  1. Combine and normalize the concatenated auxiliary states: ctx = hidden_norm(fc(context_states)) — [num_ctx, hidden_size].

  2. Project to K/V for all layers at once. At bind_kv_cache time the drafter stacks every layer’s K and V projection weights (this rank’s KV heads only) into one _fused_kv_weight. A single matmul (linear(ctx, _fused_kv_weight)) produces the K/V for every draft layer, then each layer’s keys are normalized and RoPE-encoded for their context positions.

  3. Scatter into the paged cache. For each layer, index_put_ writes the rotated keys and values into k_cache / v_cache at the block/pos indices derived from context_slot_mapping.

After precompute, the draft layers’ KV caches hold a projected representation of the target’s context at the correct paged positions, ready for the block query pass to attend to.

Draft Block Construction#

Each request’s query block has the fixed form [bonus, mask, mask, …, mask] of length block_size:

  • Slot 0 = the bonus token — the most recently accepted target token, extracted from raw_sampled_token_ids. This anchors the block on real context.

  • Slots 1 … block_size-1 = the mask token (mask_token_id, 128002 for the reference checkpoint) — placeholders whose output hidden states become the draft predictions.

_prepare_proposal_inputs() builds this flattened [num_reqs * block_size] tensor. The backbone embeds it with the borrowed target embedding and computes RoPE from the block positions (seq_lens + [0 … block_size-1]).

After the block forward, keep_idx selects positions 1 … block_size-1 of each request (dropping the bonus slot), yielding block_size - 1 draft hidden states per request. This block_size - 1 is the draft width; the reference checkpoint’s block_size = 10 gives nine draft positions, hence num_speculative_tokens = 9.

Paged KV Integration#

DFlash uses the same paged vLLM KV cache paradigm as the target model. Both report their specs via get_kv_spec(), vLLM allocates block tables spanning target + draft layers, and bind_kv_cache() binds the pre-allocated tensors to each draft layer’s attention (and builds the fused KV weight used by precompute).

The model-side _paged_slots() helper maps absolute positions to cache slots for both the injected context and the query block:

block_numbers = clamp(positions // block_size, max=max_blocks_per_seq - 1)
block_ids     = block_table_flat[req_ids * max_blocks_per_seq + block_numbers]
slot          = block_ids * block_size + (positions % block_size)

_prepare_proposal_inputs() derives each target row’s request id from last_token_indices, handles trailing padded rows, and maps the context rows used by KV injection. It uses the same helper for the fixed query block and packages those slots into per-layer attention metadata.

Async Rejection and Position Correction#

With async scheduling, the scheduler prepares the next step assuming full acceptance while the previous rejection output may still be an on-device future. DFlash passes the resulting optimistic positions and rejection tensors to _prepare_proposal_inputs(), which applies two corrections:

  • _correct_target_positions_for_async() counts how many of the previous step’s drafts were actually valid (prev_sampled_token_ids != -1 and < vocab_size), computes prev_num_rejected = (1 + prev_num_draft_tokens) - prev_valid_count, and subtracts that from target_positions per request.

  • _current_seq_lens() derives each request’s post-verify sequence length by removing rejected positions from the current step, using cur_num_draft_tokens when available.

The runner supplies the previous/current-step tensors through _build_async_spec_kwargs(), which returns real data when a prior step exists and no-op identity tensors otherwise. At batch-composition changes, the runner remaps the previous step’s drafts by request id. DFlash applies the correction and sequence-length calculation inside the compiled proposal graph. See the async scheduling design for the broader mechanism.

Adding DFlash Support for a New Target Model#

DFlash’s Qwen3-based drafter architecture is target-agnostic, but expanding DFlash support to other target models requires the following:

  1. Auxiliary hidden state capture on the target. The target model must support capturing pre-layer hidden states at requested layer indices and returning them concatenated (the same capability EAGLE3 relies on). See the EAGLE3 design for the target-side interface.

  2. A trained DFlash checkpoint. The checkpoint must:

    • Provide target_layer_ids within the valid range of target’s number of layers.

    • Share the target’s vocabulary, so the borrowed embedding and LM head line up.

Because training a draft model is the largest part of this work, start with the checkpoints published by the DFlash authors.

This document is relevant for: Trn2, Trn3