husj576/TEV-gemma4-it-12B

TEV Draft Model for Gemma 4 12B IT

Paper: Where Draft Trees Lose Target Mass: Exit-Guided Speculative Decoding

This repository contains the ExitTrain-trained DFlash draft model for google/gemma-4-12B-it, released with TEV: Exit-Guided Speculative Decoding.

GitHub | DDTree + TEV guide | Gemma 4 setup

This is a draft model, not a standalone language model. It must be paired with the matching Gemma 4 12B instruction-tuned target model in a speculative decoding system.

Model details

Field	Value
Target model	google/gemma-4-12B-it
Target interface	Gemma 4 unified, text-only benchmark path
Draft architecture	DFlashDraftModel
Training method	ExitTrain
Tree builder and verifier	DDTree + TEV
Draft block size	16
Target layers	48
Target feature layers	[1, 10, 19, 27, 36, 45]
Draft hidden size	3840
Target logit soft cap	30.0
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

Gemma 4 requires a recent Transformers release and may require accepting the target model's Hugging Face terms.

git clone https://github.com/hsj576/TEV.git
cd TEV/ddtree

pip install -r requirements.txt
pip install -U "transformers>=5.11"

# Run `hf auth login` first if access to the target is gated.
export TARGET_MODEL=google/gemma-4-12B-it
export DRAFT_MODEL=husj576/TEV-gemma4-it-12B
bash scripts/run_ddtree_tev_smoke.sh

Direct benchmark command

torchrun --nproc_per_node=1 --master_port=29700 benchmark.py \
  --model-name-or-path google/gemma-4-12B-it \
  --draft-name-or-path husj576/TEV-gemma4-it-12B \
  --dataset gsm8k \
  --max-samples 128 \
  --max-new-tokens 2048 \
  --temperature 1.0 \
  --tree-budget 64 \
  --save-path runs/gemma4_12b/gsm8k__t1_0__tb64.pt

Do not pass --flash-attn to benchmark.py: the target must use SDPA to consume DDTree's custom tree mask. The draft model still uses FlashAttention.

Gemma 4 compatibility

The TEV repository adds a compatibility path in ddtree/model/target_loader.py that:

• loads the unified model through AutoModelForImageTextToText;
• extracts the language backbone and LM head behind a causal-LM-like wrapper;
• makes the target KV cache compatible with DDTree's full-length tree mask;
• disables sliding-window KV trimming during this benchmark path.

This implementation is currently text-only and does not process image or audio inputs.

For compatibility with the full-length DDTree mask, sliding-attention layers are promoted to full attention at runtime. This can increase memory and runtime. Once a sequence exceeds Gemma 4's native sliding window, the target forward differs from native Gemma 4 behavior, although all comparison paths in this benchmark use the same loaded target.

See the Gemma 4 DDTree guide for details and troubleshooting.

Usage notes

• Use this checkpoint only with google/gemma-4-12B-it.
• 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.
• Gemma 4 12B, its draft model, the full target KV cache, and the draft tree require substantial GPU memory. Start with a smaller tree budget and generation length if needed.

Related checkpoints

• husj576/TEV-qwen3-4B
• husj576/TEV-qwen3-8B
• husj576/TEV-qwen3-coder-30B-A3B
• husj576/TEV-llama31-8B

Acknowledgements

The implementation builds on SpecForge, DFlash, and DDTree.