TEV Draft Model for LLaMA-3.1-8B-Instruct
This repository contains the ExitTrain-trained DFlash draft model for meta-llama/Llama-3.1-8B-Instruct, released with TEV: Exit-Guided Speculative Decoding. This model accompanies the 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 LLaMA-3.1-8B-Instruct target model in a speculative decoding system.
Model details
| Field | Value |
|---|---|
| Target model | meta-llama/Llama-3.1-8B-Instruct |
| Draft architecture | DFlashDraftModel |
| Training method | ExitTrain |
| Tree builder and verifier | DDTree + TEV |
| Draft block size | 10 |
| Target layers | 32 |
| Target feature layers | [1, 8, 15, 22, 29] |
| 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
Access to the target model requires accepting Meta's license terms on Hugging Face.
git clone https://github.com/hsj576/TEV.git
cd TEV/ddtree
pip install -r requirements.txt
# Run `hf auth login` first if access to the target is gated.
export TARGET_MODEL=meta-llama/Llama-3.1-8B-Instruct
export DRAFT_MODEL=husj576/TEV-llama31-8B
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 meta-llama/Llama-3.1-8B-Instruct \
--draft-name-or-path husj576/TEV-llama31-8B \
--dataset gsm8k \
--max-samples 128 \
--max-new-tokens 2048 \
--temperature 1.0 \
--tree-budget 64 \
--save-path runs/llama31_8b/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 meta-llama/Llama-3.1-8B-Instruct \
--draft husj576/TEV-llama31-8B \
--cuda-visible-devices 0 \
--server-port 7860
Usage notes
- Use this checkpoint only with
meta-llama/Llama-3.1-8B-Instruct. - 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.