Inside SGLang: DFlash and DFlash 2 Implementation

DFlash and DFlash 2 describe how a small model predicts a block of future tokens in one pass. Implementing that idea in a serving engine means coordinating two models, two sets of attention state, an allocator, and a scheduler whose CPU work may run ahead of GPU results.

This walkthrough follows one request through SGLang’s implementation. The goal is to make the source readable: which object owns each piece of state, which functions move it, and which invariants have to survive every draft–verify cycle.

We start with one GPU, greedy decoding, and a four-position block: one anchor plus three draft predictions. Then we add DFlash 2, sampling, batching, and overlap scheduling.

The code references are pinned to SGLang commit f5ccc4d, inspected October 4, 2026. The structure follows Aleksa Gordić’s Inside vLLM: build a mental model first, then attach the implementation to it. The command near the end is a configuration example; no GPU benchmark was run for this post.

The first distinction: algorithm, worker, and model

Three names appear repeatedly, and they describe different things:

Name	What it selects
`SpeculativeAlgorithm.DFLASH`	The serving algorithm and its preparation/dispatch hooks
`DFlashWorkerV2`	The coordinator that drives both target and draft execution
`DFlashDraftModel` / `DFlash2DraftModel`	The checkpoint’s neural-network implementation

Both generations use the DFLASH algorithm and DFlashWorkerV2. The V2 suffix belongs to SGLang’s speculative scheduling infrastructure; it does not mean “DFlash 2 only.” The checkpoint architecture selects the model, and the worker checks whether that model has a candidate_selector.

The worker also serves both overlap and synchronous scheduling. Disabling overlap changes how results reach the next iteration, not the underlying DFlash model.

The whole loop before opening a file

Scheduler: choose requests and prepare cache capacity
                         |
                         v
DFlashWorkerV2.forward_batch_generation(batch)
    |
    +-- EXTEND / prefill
    |     target forward + first sampled token
    |     capture prompt features -> write draft context K/V
    |     return anchor + computed lengths
    |
    +-- DECODE
          prepare [anchor, MASK, MASK, MASK]
          draft model: one parallel backbone pass
          choose [draft_1, draft_2, draft_3]
          target verifies [anchor, draft_1, draft_2, draft_3]
          accept prefix + sample correction/bonus
          write valid target features into draft context K/V
          return output tokens + next anchor + new lengths
                         |
                         v
Scheduler: process output, stopping conditions, metrics, cache lifetime
                         |
                         +--> next iteration

Drafting and verification happen inside the same worker call. During prefill, the worker populates context and returns an anchor; the next decode call produces the draft.

This is an important boundary when comparing engines. Follow SGLang’s call sites rather than assuming proposals are always produced at the end of a separate target-model step.

The five files to read first

Displayed paths are relative to python/sglang/. The links open the full source files at the pinned commit.

File	Question it answers	Entry point
srt/speculative/dflash_worker_v2.py	What happens during one cycle?	`DFlashWorkerV2.forward_batch_generation()`
srt/models/dflash.py	What does the draft network compute?	`DFlashDraftModel.forward()`
srt/speculative/dflash_info_v2.py	What survives between cycles, and who reserves slots?	`DFlashDraftInputV2.prepare_for_decode()`
srt/speculative/dflash_info.py	How does a proposal become a target batch?	`DFlashVerifyInput.prepare_for_verify()`
srt/speculative/dflash_utils.py	How are checkpoint settings and acceptance interpreted?	`parse_dflash_draft_config()`, `compute_dflash_correct_drafts_and_bonus()`

The worker is the best starting point, but its constructor is not the best first read. Open forward_batch_generation() and follow its three decode comments: draft, target verify, and materialize committed verify-input tokens. Return to initialization after those stages make sense.

Startup: how the scheduler gets a DFlash worker

The launch options establish the target model, draft checkpoint, and block size. _handle_dflash(), in srt/arg_groups/speculative_hook.py, validates the combination, resolves the draft attention backend, and infers the block size when it is omitted. DFlash does not use EAGLE’s autoregressive draft-step count or tree top-k as its proposal strategy.

The relevant construction path is:

Scheduler.init_model_worker()
    init_tp_model_worker()              -> target worker / target ModelRunner
    maybe_init_draft_worker()
        SpeculativeAlgorithm.create_worker()
            -> DFlashWorkerV2 class
        DFlashWorkerV2(..., target_worker=self.tp_worker)
            build_draft_tp_worker()
                -> draft TpModelWorker / draft ModelRunner / draft model
    init_memory_pools()
    init_all_attention_backends()
    init_all_cuda_graphs()
    self.model_worker = self.draft_worker

create_worker() returns a worker class; maybe_init_draft_worker() constructs it. Later, the scheduler dispatches generation to self.model_worker, which is the DFlash coordinator when speculation is enabled.

Who owns what?

Expression in the DFlash worker	Meaning
`self._target_worker` / `self.target_worker`	The target `TpModelWorker`
`self.model_runner`	The target runner, assigned from `target_worker.model_runner`
`self._draft_worker`	A separate `TpModelWorker` built for the draft
`self.draft_model_runner`	The draft runner
`self.draft_model`	The actual draft network; its optional selector is `self.selector`

self.model_runner is an easy name to misread: being inside a draft worker does not make it the draft runner.

In the ordinary non-windowed path, alloc_memory_pool() shares the target’s request-to-token mapping and allocator state with the draft worker. The K/V tensors remain separate. A slot number can address target K/V in one pool and draft K/V in another. With draft windowing enabled, the draft uses a private compact mapping over the shared global index space.

Configure feature capture before capturing graphs

The drafter needs features from selected target layers. This must be wired into the target before its forward graph is captured.

srt/model_executor/model_runner_components/spec_aux_hidden_state.py resolves the draft’s required target_layer_ids. During attention setup, configure_aux_hidden_state_capture() calls the target model’s set_dflash_layers_to_capture() hook. The target architecture implements the actual capture.

If target features are missing or have the wrong width, start here and at project_target_hidden(). The number of captured target layers is not necessarily the number of draft Transformer layers.

The state objects crossing the boundaries

Before following tensors, separate the scheduler’s description of work from the runner’s description of a forward pass.

Object	Role
`ScheduleBatch`	Active requests, request-pool rows, lengths, allocation metadata, sampling settings
`DFlashDraftInputV2`	Next anchor in `bonus_tokens`, computed lengths in `new_seq_lens`, preparation/relay metadata
`ForwardBatch`	Flattened model inputs, positions, cache write locations, attention and capture metadata
`DFlashVerifyInput`	A proposed linear chain and the metadata needed to build a verification `ForwardBatch`
`GenerationBatchResult`	Output token buffer, valid counts, next draft input, and auxiliary outputs

DFlashDraftInputV2 has Eagle-shaped fields such as topk_p, topk_index, and hidden_states. For this path, make_draft_input_v2() fills those with empty tensors. Their presence does not mean DFlash passes an EAGLE proposal tree or a large target-feature tensor between iterations. The useful payload is the anchor and length state; the context has already been written into draft K/V.

Step 1: prefill the target and materialize draft context

The extend branch of forward_batch_generation() calls:

batch_output = self.target_worker.forward_batch_generation(
    batch,
    pp_proxy_tensors=pp_proxy_tensors,
    capture_hidden_mode=CaptureHiddenMode.FULL,
)

The target processes the prompt, returns logits/features, and samples the first output token. The worker then calls _append_target_hidden_to_draft_kv_by_loc() with those features, batch.out_cache_loc, and the corresponding positions.

The conversion is:

target features: [number of new context tokens, concatenated feature width]
                             |
                             v
DFlashDraftModel.project_target_hidden(): fc + hidden_norm
                             |
                             v
projected context: [number of new context tokens, draft hidden width]
                             |
                             v
for each draft layer:
    context K/V projection -> K normalization -> K RoPE -> cache write

Read _append_target_hidden_sequential() to understand the math. It calls kv_proj_only(), apply_k_norm(), apply_k_rope(), and set_kv_buffer() for each layer. Fused paths optimize the same materialization work and support prefix-valid writes without first packing variable-length rows.

The context is not forwarded through the complete draft Transformer stack. Each layer receives projected target context as its own K/V source. Separately, the query block will go through attention and MLP layers during drafting.

Prefill writes the context before returning. The code calls out radix-cache safety: a reusable target prefix needs the corresponding draft context ready. The sampled first token becomes bonus_tokens; it is the next anchor, and its own target hidden state has not yet been computed.

The worker also clears logits_output.hidden_states after consuming the features. This avoids sending a large feature tensor through the result-copy path.

Step 2: reserve cache capacity before acceptance is known

Before the decode worker runs, the scheduler calls ScheduleBatch.prepare_for_decode(). Speculation delegates preparation through spec_prepare_for_decode() to the algorithm-specific state.

In DFlashDraftInputV2.prepare_for_decode(), the allocation path uses:

block_size = int(get_spec().speculative_num_draft_tokens)
reserve = 2 * block_size

The reservation calculation is page-aligned and based on request allocation state. The 2 * block_size headroom accommodates preparation running ahead under overlap; it does not mean the target verifies two blocks in one forward.

There are three different lengths to keep separate:

Length	Meaning
Committed/computed prefix	Positions whose model state is valid
Verification extent	The prefix plus the current block being evaluated
Reserved extent	Capacity the allocator has made available, possibly farther ahead

The worker gathers write slots from the prepared request-to-token table. Allocating slots is not accepting tokens. Increasing a reservation must not make rejected or uncomputed positions visible to attention.

Step 3: prepare the draft block and its cache locations

Let the computed prefix have length L = 10, batch size be B = 1, and block size be S = 4. The sampled anchor belongs at position 10.

The worker builds these tensors:

Tensor	Shape	Example / purpose
`block_ids`	`[B, S]`	`[anchor, MASK, MASK, MASK]`
`positions_2d`	`[B, S]`	`[10, 11, 12, 13]`
`verify_out_cache_loc_2d`	`[B, S]`	Physical cache slots for those logical positions
`draft_hidden`	`[B, S, H_d]`	Draft hidden states after the backbone
`draft_next`	`[B, S - 1]`	Three proposed future token IDs

_prepare_dflash_draft_block_unchecked() invokes a Triton preparation kernel. Conceptually, for request-pool row r and block offset j:

logical position = prefix length + j
physical slot    = req_to_token[r, logical position]
input id         = anchor if j == 0 else mask_token_id

Positions tell RoPE where the token belongs. Cache locations tell the memory pool where to write it. They are not interchangeable, even when a tiny example makes both look like consecutive integers.

The kernel prepares the three outputs together. The eager fallback does the same work using ordinary tensor operations. Persistent buffers keep allocation and graph-replay overhead down.

Step 4: run the draft backbone once

The worker constructs a draft ForwardBatch and calls:

draft_out = self.draft_model_runner.forward(forward_batch)

The draft batch uses ForwardMode.TARGET_VERIFY as part of the block execution plumbing, even though the executing model is the drafter. A forward-mode name alone does not identify which model is running. Check the runner and capture_hidden_mode at the call site.

Inside DFlashDraftModel.forward(), input embeddings enter the layer stack; each DFlashDecoderLayer applies attention and an MLP with residual/norm handling; a final normalization produces hidden states.

There are two K/V sources inside draft attention:

cached prefix K/V      <- projected target features, materialized earlier
current block Q/K/V    <- anchor/mask representations in the draft forward

DFlashAttention.forward() computes current-block Q/K/V, normalizes Q/K, applies RoPE, and dispatches attention through the configured backend. In ordinary full-attention DFlash, query positions can interact within the block. Checkpoint-specific sliding-window/causality settings can change that pattern.

The model returns a LogitsProcessorOutput with hidden states and next_token_logits=None. Token selection is a separate worker/head operation. In the plain DFlash path, the worker selects from draft_hidden[:, 1:, :] using the target vocabulary head. The anchor row is not a new prediction slot.

Step 5: turn the proposal into a target verification batch

The worker assembles:

draft input:      [anchor, MASK,    MASK,    MASK]
target input:     [anchor, draft_1, draft_2, draft_3]
target logit row: [p_1,    p_2,     p_3,     p_4]

The target is causal. p_1, produced at the anchor, verifies draft_1; p_2 verifies draft_2; p_3 verifies draft_3. If all three survive, p_4 supplies a bonus token.

DFlashVerifyInput.prepare_for_verify() sets up the target batch, including ForwardMode.TARGET_VERIFY, the proposed input IDs, and CaptureHiddenMode.FULL. It initializes attention metadata or loads the batch into an eligible target CUDA graph. The target worker then runs with is_verify=True.

This is a linear chain, not an EAGLE branching tree. topk=1 in DFlashVerifyInput describes that chain; it is unrelated to DFlash 2’s candidate-selector top-K.

Grammar-constrained requests add a chain-based grammar mask. The worker also applies target-logit adjustments before acceptance. Those operations change what the verifier is allowed to accept, so they belong before _accept_block().

Step 6: acceptance, output, and computed state

Consider this greedy cycle:

computed prefix length:   10
old anchor:               "of"
proposals:                ["France", "is", "London"]
target choices:           ["France", "is", "Paris", ...]

Two drafts match; the third does not. _accept_block() produces:

accept_len  = 2                 # accepted draft tokens only
commit_lens = 3                 # accepted drafts + one target token
bonus       = "Paris"           # also used for a correction at first rejection
out_tokens  = ["France", "is", "Paris", padding/ignored tail]
new_seq_len = 10 + 3 = 13

The greedy kernel compares candidates[:, 1:] against the target’s preceding rows. It stops at the first mismatch, selects the target token at accept_len, and builds the padded output buffer. Only out_tokens[:commit_lens] is valid for each request.

There is a one-position shift between output and cache state:

What becomes valid	Tokens	Positions
Newly computed target inputs	`of`, `France`, `is`	10, 11, 12
Newly emitted outputs	`France`, `is`, `Paris`	11, 12, 13
Next anchor, not yet target-computed	`Paris`	13

The old anchor was already emitted in the previous cycle, so it must not be emitted again. The new correction is emitted now, but its own target state will be computed next cycle.

If no draft survives, the cycle still emits one target correction and computes the old anchor’s state. If all three survive, it emits three drafts plus a bonus and commits four input positions. These two cases are useful for checking off-by-one errors.

The result field has a naming trap: GenerationBatchResult.accept_lens receives commit_lens, including the target token. It is not the worker’s draft-only accept_len.

Step 7: commit the context and hand off the next anchor

Target verification computed features for the entire candidate block, including the rejected suffix. The worker calls _append_target_hidden_to_draft_kv_by_loc() with the dense [B, S] cache locations and per-request commit_lens. Prefix-valid writes keep only the live verification-input rows.

In the example, draft context grows with the features for [of, France, is]. The rejected London features must not become future context. The correction Paris has no target features yet.

The target may have written speculative suffix state into reserved slots during verification. Logical lengths and subsequent overwrite/commit behavior prevent that suffix from being treated as a valid continuation. Reservation is not a promise that every stored row is live.

For hybrid attention/recurrent targets, _update_target_mamba_state_after_verify() commits the corresponding recurrent-state boundary. Changing cache accounting without tracing that state can break hybrid models even when attention-only models work.

Finally, the worker consumes and clears the target feature output, builds next_draft_input from bonus and new_seq_lens, and returns GenerationBatchResult.

DFlash 2: replace selection, retain the serving loop

DFlash2DraftModel inherits the draft backbone and adds CandidateSelector. Configured grouped convolutions are constructed by the base model and wrap both attention and MLP sublayers.

DFlashGroupedConv has prepare() and finish() methods: mix local input positions, run the sublayer, then mix its output. The coefficients combine learned base kernels with content-dependent corrections. The operation stays within each request’s block and is evaluated across positions together.

After the single backbone pass, _propose_selector_block() follows this path:

draft_hidden[:, 1:, :]                   [B, S-1, H_d]
    compute_candidates()
candidate_ids + unary logits            [B, S-1, K]
    CandidateSelector.build_lattice()
transition scores                       [B, S-1, K, K]
    CandidateSelector.sample_path()
draft_next                              [B, S-1]
realized q_rows                         [B, S-1, K]

K counts candidate tokens per position. It does not increase the number of query positions or turn target verification into a K-way tree.

The score for predecessor p and successor c combines the successor’s unary logit with a learned low-rank compatibility term:

\[\operatorname{score}(p,c;h) = U(c) + \langle A[p] \odot W_h h,\ B[c] \rangle.\]

For the first prediction, the predecessor is the anchor. For later predictions, it is a candidate from the preceding position. The lattice scores are built together; the walk then selects the row corresponding to the predecessor actually chosen.

The walk is sequential selection, but there is no new Transformer pass per token. _selector_walk_kernel() runs one Triton program per request, looping through positions with the current K scores in registers. It greedily chooses or samples a successor; it does not search for the globally highest-scoring full path.

Sampling: preserve the proposal distribution the walk actually used

Greedy decoding needs the longest prefix matching target choices. Lossless sampled verification additionally needs the proposal distribution q for each selected prefix.

For DFlash 2, q_rows is conditioned on the predecessor actually selected by the walk. Substituting the original independent unary distribution would describe a different drafter.

_selector_sampling_accept() scatters the K candidate probabilities into a reusable dense vocabulary buffer and calls accept_sampling(), shared with DSpark. It clears those candidate entries afterward because the next cycle may use different IDs.

At each proposed token y, standard rejection sampling accepts with probability $\min(1,p(y)/q(y))$. At the first rejection, the correction comes from the normalized positive residual $[p-q]_+$. After full acceptance, the target supplies a bonus token.

The worker has separate branches for selector sampling, plain-DFlash sampling verification, and greedy verification. The plain draft branch in this snapshot proposes argmax tokens; a sampling target does not automatically mean the drafter samples too. Platform support also gates the available branches. Inspect _accept_block() and _validate_phase1_sampling_support() together when changing sampling behavior.

Back in the scheduler: GPU continuation and CPU output

The outer synchronous path is straightforward:

event_loop_normal()
    get_next_batch_to_run()
    run_batch()
        self.model_worker.forward_batch_generation()
        install result.next_draft_input in batch.spec_info
        update GPU sequence lengths
    process_batch_result()

process_batch_result_decode(), in srt/managers/scheduler_components/batch_result_processor.py, receives the copied result. _resolve_spec_v2_tokens() slices each request’s fixed-stride output row by its valid accept_lens, updates committed request bookkeeping and speculative metrics, and handles grammar state. The surrounding result-processing path appends output tokens, checks stop conditions, and manages request completion and cache lifetime.

With overlap enabled, Scheduler.run_batch() adds a GPU relay:

before worker: resolve_forward_inputs(batch, future_map)
    gather last iteration's anchor / lengths for the current request rows

inside worker: on_publish(new_seq_lens)
    publish lengths after acceptance, before draft-context append finishes

after worker: _relay_forward_payload(..., batch_result)
    stash next_draft_input's anchor in FutureMap

next iteration: resolve the relay for its current request rows

FutureMap and RelayPayload live in srt/managers/overlap_utils.py. The relay is indexed by request-pool slots, allowing the batch to be filtered or merged between iterations without relying on old row ordering.

Publishing lengths early lets scheduler preparation overlap later GPU work. It does not mean the next draft may read unfinished K/V: the forward stream orders the worker’s cache append before subsequent forward work. Result copies have their own synchronization boundary.

CPU metadata can lag by an iteration. When debugging length inconsistencies, identify whether a value is the CPU committed count, a published GPU count, or an allocator bound before comparing them.

Important invariants to check while reading or changing code

Anchor ownership: the anchor has been emitted but has not yet been target-forwarded. Its row is input context, not another emitted token.
Feature ownership: only valid target-input features become draft context; padded or rejected rows must not write unrelated slots.
Length ownership: accept_len, commit_lens, verification extent, and reserved capacity describe different quantities.
Distribution ownership: sampled selector verification uses the realized conditional q, with target-logit constraints applied before acceptance.
Buffer lifetime: graph outputs and persistent buffers can be overwritten next iteration; consume features and respect stream dependencies before reuse.

These are the connections to follow before optimizing a kernel. An acceptance improvement is only useful when the engine commits the correct state and preserves the intended target distribution.

Which files matter for a specific problem?

Symptom or change	Start here	What to trace
Wrong checkpoint behavior or missing target features	`spec_aux_hidden_state.py`, `models/dflash.py`	Layer selection, feature width, projection, model architecture
Wrong token count or duplicate output	`_accept_block()`, `_resolve_spec_v2_tokens()`	Anchor shift, draft-only versus total valid counts
Long-context, prefix-reuse, or windowing errors	`dflash_info_v2.py`, worker cache append/compact helpers	Reserved slots, committed lengths, physical mapping
DFlash 2 sampling differs from intended target distribution	`_propose_selector_block()`, `_selector_sampling_accept()`	Realized q rows, residual correction, logit adjustments
Low throughput despite good acceptance	`_maybe_build_draft_sampler()`, graph/selection and K/V materialization paths	Graph versus eager execution, vocabulary top-K, synchronization, context writes

For correctness scenarios, test/registered/spec/dflash/test_dflash.py exercises real server requests, constraints, log probabilities, and cache behavior. Tests under test/registered/unit/spec/ cover narrower implementation concerns. This walkthrough identifies those checks; it does not claim they were executed here.

A launch configuration and a reading route

For a compatible installation, the DFlash 2 checkpoint card gives this serving configuration:

python -m sglang.launch_server \
  --model-path Qwen/Qwen3.8-27B \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path incoai/Qwen3.8-27B-DFlash2 \
  --speculative-num-draft-tokens 8

Here 8 includes the anchor: seven future tokens are proposed. Use a compatible target/draft pair and inspect the pinned argument hook for the supported platform/backend combinations.

To read the source in about 45–60 minutes:

10 minutes: open forward_batch_generation() and identify prefill, draft, verify, acceptance, and context append.
10 minutes: follow prepare_for_decode() and the block-preparation kernel; write down logical positions and physical slots for L=10, S=4.
10 minutes: follow DFlashDraftModel.forward() and _append_target_hidden_sequential(); distinguish query computation from context materialization.
10 minutes: trace the two-accepted-drafts example through _accept_block() and the scheduler’s output slicing.
5–20 minutes: read the selector or overlap relay, depending on whether the next question is model behavior or serving behavior.

Keep the four-position example beside the code. At each boundary, identify the executing model, the valid token positions, the output count, and who owns the next anchor.