🎯 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.txt

[!NOTE] 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

🚀 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:

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}
}