TEV Draft Model for Qwen3-4B
This repository contains the ExitTrain-trained DFlash draft model for Qwen/Qwen3-4B, released with TEV: Exit-Guided Speculative Decoding.
Paper: Where Draft Trees Lose Target Mass: Exit-Guided Speculative Decoding
This is a draft model, not a standalone language model. It must be paired with the matching Qwen3-4B target model in a speculative decoding system.
Model details
| Field | Value |
|---|---|
| Target model | Qwen/Qwen3-4B |
| Draft architecture | DFlashDraftModel |
| Training method | ExitTrain |
| Tree builder and verifier | DDTree + TEV |
| Draft block size | 16 |
| Target layers | 36 |
| Target feature layers | [1, 9, 17, 25, 33] |
| Precision | BF16 |
ExitTrain trains the drafter on its realized finite-budget draft trees. Target probabilities are evaluated under the actual tree-parent contexts, and target-derived exit feedback is used to improve where the bounded tree allocates its nodes.
Quick start with DDTree + TEV
git clone https://github.com/hsj576/TEV.git
cd TEV/ddtree
pip install -r requirements.txt
export TARGET_MODEL=Qwen/Qwen3-4B
export DRAFT_MODEL=husj576/TEV-qwen3-4B
bash scripts/run_ddtree_tev_smoke.sh
The smoke test runs GSM8K at temperature 1.0 with tree budget 64 by default.
Direct benchmark command
torchrun --nproc_per_node=1 --master_port=29700 benchmark.py \
--model-name-or-path Qwen/Qwen3-4B \
--draft-name-or-path husj576/TEV-qwen3-4B \
--dataset gsm8k \
--max-samples 128 \
--max-new-tokens 2048 \
--temperature 1.0 \
--tree-budget 64 \
--save-path runs/qwen3_4b/gsm8k__t1_0__tb64.pt
Without --skip-baseline-dflash, the benchmark also runs autoregressive decoding and block-parallel DFlash as reference paths.
Interactive demo
python -m application.webui \
--target Qwen/Qwen3-4B \
--draft husj576/TEV-qwen3-4B \
--cuda-visible-devices 0 \
--server-port 7860
Usage notes
- Use this checkpoint only with
Qwen/Qwen3-4B. - The released evaluation path uses the Qwen chat template with thinking disabled.
- At positive temperature, the DDTree branch uses TEV sampling. At temperature
0, it uses the equivalent greedy tree-follow path. - The current reference implementation targets batch-one decoding and places the target and draft models on the same GPU.
flash-attnis required by the draft model. The target uses PyTorch SDPA when DDTree's custom tree mask is enabled.- Raw benchmark outputs include output-block lengths, decode rounds, TPOT, and per-stage timings.
Related checkpoints
husj576/TEV-qwen3-8Bhusj576/TEV-qwen3-coder-30B-A3Bhusj576/TEV-gemma4-it-12Bhusj576/TEV-llama31-8B