AdaFlash: Adaptive Speculative Decoding via On-Policy Distilled Diffusion Drafters
🎯 Quick Overview
Official implementation of AdaFlash, a framework for efficient speculative decoding with on-policy distilled diffusion drafters and an adaptive length head.
- Asynchronous on-policy inference and training with hot-swappable draft checkpoints
- On-Policy Distillation of diffusion drafters
- Adaptive candidate-length initialization and training
- SGLang-based serving and speculative decoding
- Reproducible benchmarks for Vanilla, EAGLE-3, DFlash, MTP, OSD, and AdaFlash
- Configurations and recipes for Qwen3, Qwen3-Coder, and Qwen3.5 model families
- Training framework modified from SpecForge
.
├── bin/ # Python command-line entry points
├── configs/ # DFlash and adaptive-length model configurations
├── docs/ # Data, pipeline, and benchmark documentation
├── pipeline/ # Core inference, training, offline, and benchmark code
├── recipes/ # Reproducible experiment launchers
├── scripts/ # Serving, training, offline, and utility scripts
└── specforge/ # Minimal vendored SpecForge subset
📦 Installation
Create a dedicated environment and install the patched SGLang version used by this repository:
conda create -n adaflash python=3.12 -y
conda activate adaflash
pip install "sglang[all] @ git+https://github.com/sgl-project/sglang.git@refs/pull/34171/head#subdirectory=python"
pip install -r requirements.txtNote
AdaFlash relies on DFlash and adaptive-length support from the linked SGLang PR. A stock SGLang release may not expose the required runtime arguments.
📍 Released Checkpoint
| Target model | AdaFlash drafter | Training data | Description |
|---|---|---|---|
| Qwen3-8B | Qwen3-8B-AdaFlash | PerfectBlend | On-policy mixture reverse-KL loss drafter with an adaptive length head. Initialized from DFlash's checkpoint. |
🚀 Quick Start
The following command launches an AdaFlash SGLang server, waits until it is ready, runs the HTTP benchmark, and shuts the server down automatically:
# Prepare the benchmark split on first use. python bin/prepare_data.py --dataset gsm8k export MODEL_PATH=Qwen/Qwen3-8B export DRAFT_MODEL_PATH=AdaFlash/Qwen3-8B-PerfectBlend export DATASET=gsm8k export NUM_SAMPLES=1024 export CONCURRENCY=64 bash scripts/tools/benchmark.sh
See docs/benchmark_experiments.md for the complete comparison protocol and commands for all baselines.
📊 Data Preparation
Dataset preparation writes normalized training data to cache/dataset/ and benchmark splits to test_data/:
python bin/prepare_data.py --dataset gsm8k python bin/prepare_data.py --dataset perfectblend
AdaFlash can optionally regenerate training responses with a target model and an initial DFlash drafter:
CUDA_VISIBLE_DEVICES=0 python bin/regenerate_dataset.py \ --model-path Qwen/Qwen3-8B \ --initial-draft-path z-lab/Qwen3-8B-DFlash-b16 \ --dataset-path cache/dataset/perfectblend_train.jsonl \ --num-samples 4096 \ --output-jsonl cache/dataset/perfectblend_qwen3_8b_regen_4096.jsonl
Supported datasets and all preprocessing options are documented in docs/data_preparation.md.
🏋️ Training AdaFlash
AdaFlash training consists of an offline initialization stage followed by asynchronous on-policy training.
1. Initialize the Adaptive Length Head
The initialization stage fits the adaptive length head on responses generated by the target model and an initial DFlash drafter. For example, initialize a Qwen3-8B drafter with one process as follows:
torchrun --standalone --nproc_per_node=1 \ --module pipeline.offline.init_thresh_head \ --target-model-path Qwen/Qwen3-8B \ --draft-config-path configs/qwen3-8b-dflash-thresh-head.json \ --initial-draft-path z-lab/Qwen3-8B-DFlash-b16 \ --dataset-path cache/dataset/perfectblend_qwen3_8b_regen_4096.jsonl \ --output-dir models/Qwen3-8B-DFlash-AdaFlash-init \ --sglang-mem-fraction-static 0.40
Additional per-model templates are collected in scripts/offline/init_adaptive_length_head.sh.
2. Run Asynchronous On-Policy Training
The asynchronous pipeline runs an inference worker and one or more training workers on separate GPU sets. Inference continuously writes target-model trajectories to a swap buffer; the training workers consume the buffer, optimize the drafter and adaptive length head, and publish versioned weights for hot reloading.
export MODEL_PATH=Qwen/Qwen3-8B export INITIAL_DRAFT_PATH=/path/to/initialized/draft export DRAFT_CONFIG_PATH="$PWD/configs/qwen3-8b-dflash-thresh-head.json" export DATASET_PATH="$PWD/cache/dataset/perfectblend_train.jsonl" export SWAP_DIR="$PWD/outputs/qwen3_8b_adaflash" export LOG_DIR="$PWD/logs/qwen3_8b_adaflash" # First argument: inference GPUs; second argument: training GPUs. bash scripts/pipeline/run_async_pipeline_adaflash.sh "0" "1,2"
📈 Benchmarking Baselines
The repository includes serving scripts for the following methods:
| Method | Launcher |
|---|---|
| Vanilla autoregressive decoding | scripts/serve/serve_sglang.sh |
| EAGLE-3 | scripts/serve/serve_eagle3_*.sh |
| DFlash / OSD | scripts/serve/serve_dflash.sh |
| Qwen3.5 MTP | scripts/serve/serve_mtp_qwen3.5_9b.sh |
| AdaFlash | scripts/serve/serve_thresh_head.sh |
The one-command AdaFlash benchmark uses scripts/tools/benchmark.sh. For controlled two-terminal experiments, launch a server with one of the scripts above and run python bin/benchmark.py against its /generate endpoint. Detailed model paths, datasets, concurrency settings, and evaluation options are listed in docs/benchmark_experiments.md.
🙏 Acknowledgments
AdaFlash builds on the excellent open-source ecosystems around SGLang, SpecForge, and DFlash. The vendored SpecForge subset retains its upstream MIT license in specforge/LICENSE.
📖 Citation
If you find AdaFlash useful in your research, please cite our paper:
@article{arxiv'26:adaflash, title = {AdaFlash: Adaptive Speculative Decoding via On-Policy Distilled Diffusion Drafters}, author = {Qian, Yu-Yang and Wu, Hao-Cong and Chen, Chen and Sun, Jiacheng and Dong, Zhenhua and Zhao, Peng and Zhou, Zhi-Hua}, journal = {ArXiv preprint}, volume = {arXiv:2607.19223}, year = {2026} }
🌟 If this project is useful to you, please starring the repository and citing the paper. 🌟