NOTICE:
our training and benchmark both implemented on block size 9,
which is equivalent to num_speculative_tokens=8 of vllm

1 Results(accept_length)

ours qwen3-4b-dspark according to：https://huggingface.co/AQ-MedAI/Qwen3-4B-DSpark
reference qwen3-4b-dflash2 according to：https://huggingface.co/RedHatAI/Qwen3-4B-speculator.dflash2

2 Reproduce

2-1 train

2-1-1 enviornment

according to https://github.com/vllm-project/speculators/commit/7a58fc56217632d8d179b665734fa2269e8d9ffa#diff-c94e7047d3ac690b5a75e602e6ff745ec3156a87df40160f5a5c73001ac0184b

git clone https://github.com/vllm-project/speculators.git
cd speculators
git checkout 7a58fc56217632d8d179b665734fa2269e8d9ffa
pip install -e .

2-1-2 dataset

according to https://modelscope.cn/datasets/inference-optimization/Qwen3-8B-Regenerated-Collection

modelscope download --dataset inference-optimization/Qwen3-8B-Regenerated-Collection --local_dir /path/dataset/Qwen3-8B-Regenerated-Collection

2-1-3 script

#!/bin/bash

set -euo pipefail

REPO_DIR="/speculators"
MODEL="/path/model/Qwen3-4B"

DATA_DIR="/path/dataset/qwen3_8b_regen"
DATASETS=(
    "$DATA_DIR/autoif_train_Qwen3-8B.jsonl"
    "$DATA_DIR/evol_codealpaca_train_Qwen3-8B.jsonl"
    "$DATA_DIR/lmsys_arena_train_Qwen3-8B.jsonl"
    "$DATA_DIR/magpie_output.jsonl"
    "$DATA_DIR/metamathqa_train_Qwen3-8B.jsonl"
    "$DATA_DIR/nemotron_chat_Qwen3-8B.jsonl"
    "$DATA_DIR/nemotron_stem_Qwen3-8B.jsonl"
    "$DATA_DIR/orca_math_train_Qwen3-8B.jsonl"
    "$DATA_DIR/tulu3_Qwen3-8B.jsonl"
    "$DATA_DIR/ultrachat_output.jsonl"
    "$DATA_DIR/ultrafeedback_train_sft_Qwen3-8B.jsonl"
    "$DATA_DIR/ultrainteract_train_Qwen3-8B.jsonl"
)

OUTPUT_DIR="$REPO_DIR/output/dflash2_qwen3_4b_8spec"
SAVE_PATH="/path/ckpt/dflash2_qwen3_4b_redhat/checkpoints"
HIDDEN_STATES_DIR="/tmp/hidden_states_dflash2_qwen3_4b_8spec"   

VLLM_PORT=8000
MAX_SAMPLES=""           
SEQ_LENGTH=8192
EPOCHS=1
LR=6e-4
NUM_WORKERS=32

SPECULATOR_TYPE="dflash2"
BLOCK_SIZE=9
MAX_ANCHORS=512
NUM_LAYERS=5
TARGET_LAYER_IDS="1 9 17 25 33"   
CONV_KERNEL_SIZE=2
CONV_GROUP_SIZE=16
SELECTOR_RANK=256
SELECTOR_TOP_K=16
LOSS_FN='{"ce": 0.1, "tv": 0.9}'
DFLASH_DECAY_GAMMA=4.0
PER_POSITION_LOSS_WEIGHT="fixed-exp-decay"

OPTIMIZER="adamw"
WEIGHT_DECAY=0.0
SCHEDULER_TYPE="cosine"
SCHEDULER_WARMUP_RATIO=0.04
SEED=42
CHECKPOINT_FREQ=0.1


VLLM_GPUS="0,1,2,3"
VLLM_DP=4
TRAIN_GPUS="4,5,6,7"
NUM_TRAIN_GPUS=4

cd "$REPO_DIR"

fail() { echo "ERROR: $*" >&2; exit 1; }

mkdir -p "$OUTPUT_DIR" "$SAVE_PATH"

MS_ARGS=()
if [ -n "$MAX_SAMPLES" ]; then MS_ARGS=(--max-samples "$MAX_SAMPLES"); fi

# Build --data arguments for prepare_data
DATA_ARGS=()
for f in "${DATASETS[@]}"; do
    DATA_ARGS+=(--data "$f")
done

rm -rf "$HIDDEN_STATES_DIR"
mkdir -p "$HIDDEN_STATES_DIR"


CUDA_VISIBLE_DEVICES="$VLLM_GPUS" python scripts/launch_vllm.py "$MODEL" \
    --target-layer-ids $TARGET_LAYER_IDS \
    --hidden-states-path "$HIDDEN_STATES_DIR" \
    -- --data-parallel-size "$VLLM_DP" --port "$VLLM_PORT" \
    --enforce-eager \
    --no-enable-chunked-prefill \
    2>&1 | tee "$SAVE_PATH/vllm.log" &
VLLM_PID=$!


while ! curl -s "http://localhost:${VLLM_PORT}/health" >/dev/null 2>&1; do
    # 检查进程是否意外崩溃
    if ! kill -0 "$VLLM_PID" 2>/dev/null; then
        echo "ERROR: vLLM failed to start. Check $SAVE_PATH/vllm.log for details."
        exit 1
    fi
    sleep 3
done


python scripts/prepare_data.py \
    --model "$MODEL" \
    "${DATA_ARGS[@]}" \
    --output "$OUTPUT_DIR" \
    --render-endpoint "http://localhost:${VLLM_PORT}" \
    --seq-length "$SEQ_LENGTH" \
    --num-preprocessing-workers "$NUM_WORKERS" \
    ${MS_ARGS[@]+"${MS_ARGS[@]}"} \
    2>&1 | tee "$SAVE_PATH/prepare_data.log"


CUDA_VISIBLE_DEVICES="$TRAIN_GPUS" torchrun \
    --standalone --nproc_per_node "$NUM_TRAIN_GPUS" \
    scripts/train.py \
    --verifier-name-or-path "$MODEL" \
    --speculator-type "$SPECULATOR_TYPE" \
    --data-path "$OUTPUT_DIR" \
    --vllm-endpoint "http://localhost:${VLLM_PORT}/v1" \
    --save-path "$SAVE_PATH" \
    --block-size "$BLOCK_SIZE" \
    --max-anchors "$MAX_ANCHORS" \
    --target-layer-ids $TARGET_LAYER_IDS \
    --num-layers "$NUM_LAYERS" \
    --conv-kernel-size "$CONV_KERNEL_SIZE" \
    --conv-group-size "$CONV_GROUP_SIZE" \
    --selector-rank "$SELECTOR_RANK" \
    --selector-top-k "$SELECTOR_TOP_K" \
    --loss-fn "$LOSS_FN" \
    --dflash-decay-gamma "$DFLASH_DECAY_GAMMA" \
    --per-position-loss-weight "$PER_POSITION_LOSS_WEIGHT" \
    --optimizer "$OPTIMIZER" \
    --lr "$LR" \
    --weight-decay "$WEIGHT_DECAY" \
    --scheduler-type "$SCHEDULER_TYPE" \
    --scheduler-warmup-ratio "$SCHEDULER_WARMUP_RATIO" \
    --epochs "$EPOCHS" \
    --total-seq-len "$SEQ_LENGTH" \
    --seed "$SEED" \
    --fsdp-shard \
    --on-missing generate \
    --on-generate delete \
    --hidden-states-path "$HIDDEN_STATES_DIR" \
    --checkpoint-freq "$CHECKPOINT_FREQ" \
    2>&1 | tee "$SAVE_PATH/train.log"

2-2 benchmark

2-2-1 dataset

• gsm8k：https://modelscope.cn/datasets/modelscope/gsm8k
• math500：https://modelscope.cn/datasets/AI-ModelScope/MATH-500
• aime25：https://modelscope.cn/datasets/evalscope/aime25
• mbpp：https://modelscope.cn/datasets/opencompass/mbpp
• human eval：https://modelscope.cn/datasets/opencompass/humaneval
• live code bench：https://modelscope.cn/datasets/Skywork/LiveCodeBench
• ceval：https://modelscope.cn/datasets/evalscope/ceval
• mt bench：https://modelscope.cn/datasets/HuggingFaceH4/mt_bench_prompts
• alpaca：https://modelscope.cn/datasets/AI-ModelScope/alpaca-gpt4-data-zh
• arena hard：https://modelscope.cn/datasets/AI-ModelScope/arena-hard-auto-v0.1

2-2-2 server

vllm serve /path/model/Qwen3-4B \
  --no-enable-flashinfer-autotune \
  --max-num-seqs 128 \
  --data-parallel-size 8 \
  --no-enable-prefix-caching \
  --speculative-config '{"method":"dflash","model":"/path/ckpt/dflash2", "num_speculative_tokens": 8}'

2-2-3 benchmark

vllm bench serve \
  --base-url http://localhost:8000 \
  --model /path/model/Qwen3-4B \
  --dataset-name custom \
  --dataset-path /path/dataset/human-eval/humaneval-prompt.jsonl \
  --num-prompts 1040 \
  --temperature 0 \
  --custom-output-len 3072 \
  --no-oversample \
  --max-concurrency 32