Understanding
DFlash

# EAGLE-3: generates K draft tokens in K sequential steps
# Each step depends on the previous — cannot be parallelized
draft_tokens = []
current_input = fused_features  # 4096-dim vector from target model

for step in range(K):   # K = 7 → 7 sequential decoder runs
    logits = eagle_decoder_layer(current_input)  # 1 transformer layer
    next_token = sample(logits)
    draft_tokens.append(next_token)
    current_input = embed(next_token)  # feed back for next step
    # ↑ each iteration WAITS for the previous one — pure sequential

# Result: 7 tokens from 7 sequential passes through 1 decoder layer
# DFlash will replace this entire loop with ONE parallel pass

vllm serve meta-llama/Llama-3.3-70B-Instruct \
  --speculative-config '{"model": "yuhuili/EAGLE3-LLaMA3.3-Instruct-70B",
    "num_speculative_tokens": 3,
    "method": "eagle3",
    "draft_tensor_parallel_size": 1}'
# Draft head is tiny (~277 MB), co-deployed on same GPU as target
# num_speculative_tokens: start at 3-5 for most tasks

import torch

# Forward pass with output_hidden_states=True
# This returns the hidden state tensor at every layer — no extra compute!
outputs = big_model(
    input_ids=token_ids,                # [batch, seq_len]
    output_hidden_states=True          # the magic flag
)

# outputs.hidden_states is a tuple of tensors, one per layer:
# hidden_states[0]  → shape [batch, seq_len, 4096]  (embedding layer)
# hidden_states[8]  → shape [batch, seq_len, 4096]  (layer 8 = early)
# hidden_states[16] → shape [batch, seq_len, 4096]  (layer 16 = mid)
# hidden_states[32] → shape [batch, seq_len, 4096]  (layer 32 = final)

num_layers = len(outputs.hidden_states)            # e.g. 33 (0 to 32)

# Sample uniformly across layers (DFlash samples more than EAGLE-3)
step = num_layers // 6
sampled = [outputs.hidden_states[i][0, -1, :]   # last token position
           for i in range(0, num_layers, step)]   # sample every 5th layer
# sampled: list of 6 tensors, each [4096]

stacked = torch.stack(sampled, dim=0)             # [6, 4096]
fused = feature_projector(stacked)                # learned FC → [4096]
# fused is the distilled "essence" of what the big model knows about context

########################################
# Standard transformer attention
########################################
d = 64  # head dimension
Q = linear_q(hidden)       # [seq, d] — "what am I looking for?"
K = linear_k(hidden)       # [seq, d] — "what do I have as keys?"
V = linear_v(hidden)       # [seq, d] — "what info do I carry?"

scores = softmax(Q @ K.T / d**0.5)  # [seq, seq] attention weights
output = scores @ V                   # [seq, d] attended output

########################################
# DFlash attention with KV injection
########################################
# The big model's fused features are projected to K and V space
K_from_big = linear_k_inject(fused_features)   # [1, d]
V_from_big = linear_v_inject(fused_features)   # [1, d]

# Concatenate draft model's own KV with injected KV from big model
K_combined = torch.cat([linear_k(hidden), K_from_big], dim=0)  # [seq+1, d]
V_combined = torch.cat([linear_v(hidden), V_from_big], dim=0)  # [seq+1, d]

# Now every draft token attends to BOTH draft context AND big model knowledge
scores = softmax(Q @ K_combined.T / d**0.5)  # [seq, seq+1]
output = scores @ V_combined                    # [seq, d]
# This is done at EVERY draft layer — not just the first

# CAUSAL mask — standard autoregressive (EAGLE-3's decoder uses this)
# 1 = can attend, 0 = blocked
causal = torch.tensor([
    [1, 0, 0, 0],  # pos 0: sees only itself
    [1, 1, 0, 0],  # pos 1: sees 0,1
    [1, 1, 1, 0],  # pos 2: sees 0,1,2
    [1, 1, 1, 1],  # pos 3: sees all
])
# Each position must be generated one at a time — SEQUENTIAL

# BIDIRECTIONAL (non-causal) mask — DFlash draft model uses this
bidirectional = torch.tensor([
    [1, 1, 1, 1],  # draft[?] 0 sees all draft positions
    [1, 1, 1, 1],  # draft[?] 1 sees all draft positions
    [1, 1, 1, 1],  # draft[?] 2 sees all draft positions
    [1, 1, 1, 1],  # draft[?] 3 sees all draft positions
])
# All positions computed at once — ONE FORWARD PASS for 16 tokens

##########################################################
# STEP 1: Run big model forward pass
# Identical to normal generation — just add output_hidden_states=True
##########################################################
big_outputs = big_model(
    input_ids=current_tokens,           # [1, seq_len]
    use_cache=True,                   # reuse KV cache from previous tokens
    output_hidden_states=True          # ← DFlash needs these
)
# big_outputs.hidden_states: tuple of (num_layers+1) tensors
# each tensor shape: [1, seq_len, hidden_size]
last_token_logits = big_outputs.logits[0, -1]   # [vocab_size]

##########################################################
# STEP 2: Extract hidden states from multiple layers
##########################################################
all_hidden = big_outputs.hidden_states            # tuple of num_layers tensors
num_layers = len(all_hidden)                     # e.g. 33 for 32-layer model

# Sample uniformly across all layers (more coverage than EAGLE-3's 3 layers)
extract_layers = list(range(1, num_layers, num_layers // 6))  # e.g. [1,6,11,17,22,27]
sampled = torch.stack(
    [all_hidden[i][0, -1, :] for i in extract_layers], dim=0
)   # [num_extracted, hidden_size]  e.g. [6, 4096]

##########################################################
# STEP 3: Fuse into one conditioning vector
##########################################################
# feature_projector is a small learned linear layer: [num_extracted, d] → [d]
# It learns which layer's signals matter most for predicting future tokens
fused_ctx = feature_projector(sampled.flatten())   # [hidden_size]

##########################################################
# STEP 4: Draft model generates K tokens in ONE parallel pass
##########################################################
K = 16   # number of tokens to draft simultaneously

# Initialize K mask token embeddings — these are what get "denoised"
# Think of them as K [MASK] tokens, all unknown
mask_embeds = mask_embedding.expand(K, -1)        # [K, hidden_size]

# The draft model's forward pass:
# - Input: K mask token embeddings
# - KV injection: fused_ctx goes into Key and Value of EVERY draft layer
# - Attention mask: non-causal (all K positions attend to all K positions)
# - Output: logits for all K positions simultaneously
draft_logits = draft_model.forward(
    input_embeds=mask_embeds,           # [K, hidden_size]
    kv_injection_context=fused_ctx,     # injected into every layer's KV
    attention_mask="bidirectional"     # non-causal: all see all
)                                       # output: [K, vocab_size]

# Sample draft tokens from logits — ALL K at once, no loop needed
draft_tokens = torch.argmax(draft_logits, dim=-1)   # [K]
# ↑ This is the entire architectural difference from EAGLE-3.
# EAGLE-3 needs a for-loop with K iterations. DFlash produces all K in one call.

##########################################################
# STEP 5: Big model verifies all K drafts in one pass
##########################################################
# Concatenate existing context with K draft tokens
verify_ids = torch.cat([current_tokens, draft_tokens.unsqueeze(0)], dim=1)
verify_out = big_model(verify_ids)                  # one forward pass
verify_logits = verify_out.logits[0]                # [seq_len+K, vocab_size]

##########################################################
# STEP 6: Accept/reject loop
##########################################################
accepted = []
seq_len = current_tokens.shape[1]

for i in range(K):
    # What would the big model have produced at position seq_len+i?
    big_model_token = torch.argmax(verify_logits[seq_len + i - 1])
    if big_model_token == draft_tokens[i]:
        accepted.append(draft_tokens[i])
    else:
        accepted.append(big_model_token)  # use big model's correction
        break                              # stop here — can't trust later drafts

# Typical result: 8–12 tokens accepted
# Net throughput: 8–12 tokens for cost of ~2 big model passes
# vs baseline: 8–12 tokens for cost of 8–12 big model passes
# Speedup: 4–6×
current_tokens = torch.cat([current_tokens, torch.tensor(accepted).unsqueeze(0)], dim=1)

# Standard transformer KV cache — rolling back is trivial:
kv_cache[position] = (K_vector, V_vector)   # stored separately per position
# If rejected: just don't include that position in future attention. Easy.

########################################
# GDN / DeltaNet — stateful, hard to roll back:
########################################
# GDN maintains a recurrent state S that updates multiplicatively:
S_new = S_old * gate + delta_value   # S encodes compressed history
# After processing a draft token: S_old is gone, replaced by S_new
# If we reject that token, we can't undo S_new → S_old automatically
# (matrix multiplication is not trivially reversible)

########################################
# SGLang's solution: extra_buffer strategy
########################################
# Before starting the speculative draft phase:
S_checkpoint = copy(S)                # save current recurrent state

# Run draft model, run verification...
num_accepted = verify_and_count(draft_tokens)

# If any tokens rejected:
if num_accepted < K:
    S = S_checkpoint                    # restore clean state
    S = advance_state(S, accepted_tokens)  # replay only accepted tokens
# Result: S is exactly what it would have been without speculation

# This is why you need --mamba-scheduler-strategy extra_buffer in SGLang
# for Qwen3.6 models with DFlash

# Install
pip install -U vllm --torch-backend=auto

# Launch OpenAI-compatible server on localhost:8000
vllm serve Qwen/Qwen3.5-27B \
  --speculative-config '{"method": "dflash",
    "model": "z-lab/Qwen3.5-27B-DFlash",
    "num_speculative_tokens": 15}' \
  --attention-backend flash_attn \
  --max-num-batched-tokens 32768 \
  --speculative-disable-by-batch-size 32   # auto-disable at high concurrency

export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1

python -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-35B-A3B \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path z-lab/Qwen3.5-35B-A3B-DFlash \
  --speculative-num-draft-tokens 16 \
  --tp-size 1 \
  --attention-backend trtllm_mha \
  --speculative-draft-attention-backend fa4 \
  --mem-fraction-static 0.75 \
  --mamba-scheduler-strategy extra_buffer \  # GDN state rollback
  --trust-remote-code

# Install patched build until main vLLM release includes it
pip install vllm
pip install -U --torch-backend=auto \
  "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/40898/head"

vllm serve Qwen/Qwen3.6-35B-A3B \
  --speculative-config '{"method": "dflash",
    "model": "z-lab/Qwen3.6-35B-A3B-DFlash",
    "num_speculative_tokens": 15}' \
  --attention-backend flash_attn \
  --max-num-batched-tokens 32768

from transformers import AutoModel, AutoModelForCausalLM

# Load draft model (shares embedding + LM head with target automatically)
draft = AutoModel.from_pretrained(
    "z-lab/Qwen3-8B-DFlash-b16",
    trust_remote_code=True, dtype="auto", device_map="cuda:0"
).eval()

# Load target model
target = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B", dtype="auto", device_map="cuda:0"
).eval()

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

# Completely normal OpenAI API call — DFlash is transparent
response = client.chat.completions.create(
    model="Qwen/Qwen3.5-27B",
    messages=[{"role": "user", "content": "Write a binary search in Python"}],
    stream=True
)

for chunk in response:
    print(chunk.choices[0].delta.content or "", end="", flush=True)
# Output is 100% identical to non-DFlash. Just arrives faster.