DFlash: Block Diffusion for Flash Speculative Decoding

Paper | Blog | Models

DFlash is a lightweight block diffusion model designed for speculative decoding. It enables efficient and high-quality parallel drafting.

DFlash_demo.mp4

Supported Models

Model	DFlash Draft
gemma-4-26B-A4B-it	z-lab/gemma-4-26B-A4B-it-DFlash
gemma4-31B-it	z-lab/gemma-4-31B-it-DFlash
Qwen3.6-27B	z-lab/Qwen3.6-27B-DFlash
Qwen3.6-35B-A3B	z-lab/Qwen3.6-35B-A3B-DFlash
MiniMax-M2.5 (Preview)	z-lab/MiniMax-M2.5-DFlash
Kimi-K2.5	z-lab/Kimi-K2.5-DFlash
Qwen3.5-4B	z-lab/Qwen3.5-4B-DFlash
Qwen3.5-9B	z-lab/Qwen3.5-9B-DFlash
Qwen3.5-27B	z-lab/Qwen3.5-27B-DFlash
Qwen3.5-35B-A3B	z-lab/Qwen3.5-35B-A3B-DFlash
Qwen3.5-122B-A10B	z-lab/Qwen3.5-122B-A10B-DFlash
Qwen3-Coder-Next	z-lab/Qwen3-Coder-Next-DFlash
Qwen3-Coder-30B-A3B	z-lab/Qwen3-Coder-30B-A3B-DFlash
gpt-oss-20b	z-lab/gpt-oss-20b-DFlash
gpt-oss-120b	z-lab/gpt-oss-120b-DFlash
Qwen3-4B (non-thinking)	z-lab/Qwen3-4B-DFlash-b16
Qwen3-8B (non-thinking)	z-lab/Qwen3-8B-DFlash-b16
Llama-3.1-8B-Instruct	z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat
DeepSeek-V4-Flash	Coming soon
DeepSeek-V4-Pro	Coming soon
MiniMax-M2.7	Coming soon
GLM-5.1	Coming soon

Feel free to open a GitHub issue to request support for additional models. We will also open-source the training recipe soon, so you can train your own DFlash draft model to accelerate any LLM.

📦 Installation

Use a separate virtual environment for each to avoid conflict.

Backend	Install command
Transformers	uv pip install -e ".[transformers]"
SGLang	uv pip install -e ".[sglang]"
vLLM	See below
MLX (Apple Silicon)	pip install -e ".[mlx]"

vLLM: Mainline vLLM already includes DFlash. For the newer SWA draft models and Gemma4 models, install the matching temporary vLLM PR:

uv pip install -e ".[vllm]"

# Non-Gemma4 DFlash models
uv pip install -U --torch-backend=auto "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/40898/head"

# Gemma4 DFlash models
uv pip install -U --torch-backend=auto "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/41703/head"

If you use uv sync, install one backend extra at a time:

uv sync --extra transformers
# or
uv sync --extra sglang

transformers and sglang extras intentionally conflict because they require different transformers versions.

🚀 Quick Start

vLLM

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

SGLang

export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1

# Optional: enable schedule overlapping (experimental, may not be stable)
# export SGLANG_ENABLE_SPEC_V2=1
# export SGLANG_ENABLE_DFLASH_SPEC_V2=1
# export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=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 \
    --trust-remote-code

Transformers

Only Qwen3 and LLaMA-3.1 models support the Transformers backend.

from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

draft = AutoModel.from_pretrained("z-lab/Qwen3-8B-DFlash-b16", trust_remote_code=True, dtype="auto", device_map="cuda:0").eval()
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype="auto", device_map="cuda:0").eval()
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, enable_thinking=False).to(draft.device)

output = draft.spec_generate(input_ids=input_ids, max_new_tokens=2048, temperature=0.0, target=target, stop_token_ids=[tokenizer.eos_token_id])
print(tokenizer.decode(output[0], skip_special_tokens=False))

MLX (Apple Silicon)

There have been many great community DFlash implementations on MLX; we provide a simple and efficient one here, tested on an Apple M5 Pro with Qwen3, Qwen3.5 and Gemma-4 models.

from dflash.model_mlx import load, load_draft, stream_generate

model, tokenizer = load("Qwen/Qwen3.5-4B")
draft = load_draft("z-lab/Qwen3.5-4B-DFlash")

messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
tps = 0.0
for r in stream_generate(model, draft, tokenizer, prompt, block_size=16, max_tokens=2048, temperature=0.6):
    print(r.text, end="", flush=True)
    tps = r.generation_tps
print(f"\nThroughput: {tps:.2f} tok/s")

📊 Evaluation

All benchmarks share the same datasets (gsm8k, math500, humaneval, mbpp, mt-bench). Datasets are automatically downloaded and cached as JSONL in cache/ on first run.

vLLM:

python -m dflash.benchmark --backend vllm \
    --base-url http://127.0.0.1:8000 --model Qwen/Qwen3.5-27B \
    --dataset gsm8k --num-prompts 128 --concurrency 1 --enable-thinking

SGLang:

python -m dflash.benchmark --backend sglang \
    --base-url http://127.0.0.1:30000 --model Qwen/Qwen3.5-35B-A3B \
    --dataset gsm8k --num-prompts 128 --concurrency 1 --enable-thinking

Transformers (Qwen3 and LLaMA only):

torchrun --nproc_per_node=8 -m dflash.benchmark --backend transformers \
    --model Qwen/Qwen3-8B --draft-model z-lab/Qwen3-8B-DFlash-b16 \
    --dataset gsm8k --max-samples 128

MLX:

python -m dflash.benchmark --backend mlx \
    --model mlx-community/gemma-4-31b-it-4bit --draft-model z-lab/gemma-4-31B-it-DFlash \
    --dataset gsm8k --max-samples 128 --enable-thinking

One-command GSM8K benchmark (UV + Python 3.12)

This repository includes a helper script that bootstraps uv + Python 3.12, installs the Transformers backend dependencies, and runs a GSM8K benchmark:

./benchmark_gsm8k.sh

Note: uv sync removes packages not declared in project dependencies. Since flash-attn is installed from a wheel URL and not pinned in pyproject.toml, the script installs it again after sync by default.

Defaults:

• target model: Qwen/Qwen3-4B
• draft model: z-lab/Qwen3-4B-DFlash-b16
• dataset: gsm8k
• max samples: 128

You can override them through environment variables:

TARGET_MODEL=Qwen/Qwen3-8B \
DRAFT_MODEL=z-lab/Qwen3-8B-DFlash-b16 \
MAX_SAMPLES=256 \
./benchmark_gsm8k.sh

FlashAttention controls:

# disable wheel install
INSTALL_FLASH_ATTN=0 ./benchmark_gsm8k.sh

# use a different wheel URL
FLASH_ATTN_WHEEL_URL="https://.../flash_attn-...whl" ./benchmark_gsm8k.sh

Two-stage draft refine benchmark

To run a variant where draft tokens are generated, then passed through the draft model again for a second refinement pass before verification:

./benchmark_gsm8k_refine.sh

This uses the same default target/draft models and dataset as benchmark_gsm8k.sh, and reports throughput + acceptance metrics for comparison.

You can control refinement passes:

REFINE_PASSES=2 ./benchmark_gsm8k_refine.sh
REFINE_PASSES=3 ./benchmark_gsm8k_refine.sh

Acknowledgement

Huge thanks to @dcw02, @gongy, and the team at @modal-labs for their fast, high-quality support in bringing DFlash to SGLang. And huge thanks as well to @benchislett at NVIDIA for his work in bringing DFlash to vLLM and helping make it available to the broader serving community.

Citation

If you find DFlash useful, please cite our work. To share feedback on DFlash or request new model support, please fill out this form: DFlash Feedback.

@article{chen2026dflash,
  title   = {{DFlash: Block Diffusion for Flash Speculative Decoding}},
  author  = {Chen, Jian and Liang, Yesheng and Liu, Zhijian},
  journal = {arXiv preprint arXiv:2602.06036},
  year    = {2026}
}