DFlash: Block Diffusion for Flash Speculative Decoding

https://github.com/z-lab/dflash
Admin
DFlash is a lightweight block diffusion model designed for speculative decoding. It enables...
Tokens:2,334
Snippets:26
Trust Score:8.7
License:MIT
Update:4 months ago
Show doc for...
Tokens:
Raw
### Install Transformers Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install DFlash with the Transformers backend. Use a separate virtual environment to avoid conflicts.

```bash
uv pip install -e ".[transformers]"
```

--------------------------------

### Install SGLang Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install DFlash with the SGLang backend. Use a separate virtual environment to avoid conflicts.

```bash
uv pip install -e ".[sglang]"
```

--------------------------------

### Install MLX Backend (Apple Silicon)

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install DFlash with the MLX backend for Apple Silicon. Use a separate virtual environment to avoid conflicts.

```bash
pip install -e ".[mlx]"
```

--------------------------------

### Install vLLM from Source for Gemma4

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install vLLM from a specific git branch for Gemma4 DFlash support. This is a fallback if the Docker image is not used.

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

--------------------------------

### Install vLLM Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install DFlash with the vLLM backend. vLLM v0.20.1+ includes core DFlash support. Use a separate virtual environment to avoid conflicts.

```bash
uv pip install -e ".[vllm]"
```

--------------------------------

### Install vLLM from Source for Newer SWA Draft Models

Source: https://github.com/z-lab/dflash/blob/main/README.md

Install vLLM from a specific git branch for newer non-Gemma4 SWA draft models. This uses the SWA support branch.

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

--------------------------------

### Launch SGLang Server with DFlash

Source: https://github.com/z-lab/dflash/blob/main/README.md

Launches a SGLang server configured for DFlash speculative decoding. Ensure you have the necessary models and configurations in place.

```bash
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
```

--------------------------------

### Evaluate DFlash with vLLM Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Runs benchmarks for DFlash using the vLLM backend. Ensure the vLLM server is running and accessible.

```bash
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
```

--------------------------------

### Evaluate DFlash with SGLang Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Runs benchmarks for DFlash using the SGLang backend. Ensure the SGLang server is running and accessible.

```bash
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
```

--------------------------------

### Evaluate DFlash with Transformers Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Runs benchmarks for DFlash using the Transformers backend. This is supported for Qwen3 and LLaMA models only.

```bash
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
```

--------------------------------

### Docker Image for vLLM Gemma4 DFlash

Source: https://github.com/z-lab/dflash/blob/main/README.md

Pull the recommended Docker image for vLLM Gemma4 DFlash support. This is recommended for Gemma4 models.

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

--------------------------------

### Evaluate DFlash with MLX Backend

Source: https://github.com/z-lab/dflash/blob/main/README.md

Runs benchmarks for DFlash using the MLX backend. This is suitable for Apple Silicon hardware.

```bash
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
```

--------------------------------

### Run Gemma4 with Docker using vLLM

Source: https://github.com/z-lab/dflash/blob/main/README.md

Launch a Gemma4 model with DFlash using vLLM via Docker. This command maps ports, mounts cache, and configures speculative decoding.

```bash
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
```

--------------------------------

### MLX Inference with DFlash

Source: https://github.com/z-lab/dflash/blob/main/README.md

Performs streaming generation using DFlash on MLX for Apple Silicon. Loads both the target and draft models.

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

--------------------------------

### Serve Non-Gemma4 Models with vLLM

Source: https://github.com/z-lab/dflash/blob/main/README.md

Serve non-Gemma4 models using vLLM with DFlash speculative decoding. This command specifies the model and DFlash configuration.

```bash
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
```

--------------------------------

### Configure SGLang for DFlash

Source: https://github.com/z-lab/dflash/blob/main/README.md

Set an environment variable to allow SGLang to overwrite context length for DFlash. This is a prerequisite for using DFlash with SGLang.

```bash
export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1
```

--------------------------------

### Transformers Inference with DFlash

Source: https://github.com/z-lab/dflash/blob/main/README.md

Performs speculative generation using the Transformers library with DFlash. Requires draft and target models to be loaded.

```python
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))
```
Tokens:
Raw