DFlash: Block Diffusion for Flash Speculative Decoding

To use DFlash with vLLM, pass a JSON string to the --speculative-config flag.

Gemma4 with Docker:

docker run --rm -it \
  --gpus all \
  --ipc=host \
  --shm-size=16g \
  -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  ghcr.io/z-lab/vllm-openai:gemma4-dflash-cu130 \
  google/gemma-4-26B-A4B-it \
  --host 0.0.0.0 \
  --port 8000 \
  --speculative-config '{"method": "dflash", "model": "z-lab/gemma-4-26B-A4B-it-DFlash", "num_speculative_tokens": 15, "attention_backend": "flash_attn"}' \
  --attention-backend triton_attn \
  --max-num-batched-tokens 32768 \
  --trust-remote-code

Non-Gemma4 models:

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

Standard vLLM installation works for most models, but specific builds are required for Gemma4 or newer non-Gemma4 SWA draft models.

Gemma4 (via Docker):

docker pull ghcr.io/z-lab/vllm-openai:gemma4-dflash-cu130

Gemma4 (Source fallback):

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

Non-Gemma4 SWA draft models:

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

Set SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 and use the sglang.launch_server module. You can optionally enable experimental schedule overlapping by exporting SGLANG_ENABLE_SPEC_V2=1, SGLANG_ENABLE_DFLASH_SPEC_V2=1, or SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1.

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 \
    --trust-remote-code

Install DFlash using uv or pip depending on your target backend. It is recommended to use a separate virtual environment for each backend to avoid conflicts.

• Transformers: uv pip install -e ".[transformers]" (Supports Qwen3 and LLaMA-3.1 only)
• SGLang: uv pip install -e ".[sglang]"
• vLLM: uv pip install -e ".[vllm]" (v0.20.1+ includes core support)
• MLX (Apple Silicon): pip install -e ".[mlx]"

# Transformers
uv pip install -e ".[transformers]"

# SGLang
uv pip install -e ".[sglang]"

# vLLM
uv pip install -e ".[vllm]"

# MLX
pip install -e ".[mlx]"

For Apple Silicon, use the dflash.model_mlx module to load models and generate text via streaming.

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")

The dflash_generate function performs speculative decoding using a DFlashDraftModel and a target model. It supports block-based generation and can return performance statistics if return_stats is set to True.

@torch.inference_mode()
def dflash_generate(
    model: "DFlashDraftModel",
    target: nn.Module,
    input_ids: torch.LongTensor,
    max_new_tokens: int,
    stop_token_ids: Optional[list[int]],
    temperature: float,
    block_size: Optional[int] = None,
    mask_token_id: Optional[int] = None,
    return_stats: bool = False,
)