gemma4-dflash-speculator-fp8-block

RedHatAI/gemma-4-31B-it-speculator.dflash

This is a preliminary (and subject to change) DFlash speculator model for google/gemma-4-31B-it.

It was trained using the Speculators library on a combination of the Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered dataset and the train_sft split of the HuggingFaceH4/ultrachat_200k dataset. Training data used Magpie + UltraChat with responses from the gemma-4-31B-it model (no reasoning).

This model should be used with the google/gemma-4-31b-it chat template, specifically through the /chat/completions endpoint.

Note:

It was validated on Nvidia H100, other hardware validation pending.

We are continuing to train this model and will update with more evaluations and new weights in the future.

Until full support lands on vLLM main, please install vLLM with VLLM_USE_PRECOMPILED=1 uv pip install -U --torch-backend=auto "vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/42095/head"

Training Details

Prepare data

python scripts/prepare_data.py \
  --model google/gemma-4-31B-it \
  --data ./gemma4-31B-responses/magpie_gemma-4-31B-it.jsonl \
  --data ./gemma4-31B-responses/ultrachat_gemma-4-31B-it.jsonl \
  --output ./output

Launch vLLM

CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/launch_vllm.py \
  google/gemma-4-31B-it \
  --hidden-states-path ./output \
  --target-layer-ids 1 17 29 47 58 60 \
  -- --max-model-len 16384 \
  -tp 2 \
  --data-parallel-size 2

Launch training

CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun \
  --standalone \
  --nproc_per_node=4 \
  scripts/train.py \
  --verifier-name-or-path google/gemma-4-31B-it \
  --data-path ./output \
  --on-missing generate \
  --on-generate delete \
  --scheduler-type cosine \
  --draft-vocab-size 32000 \
  --max-anchors 3072 \
  --target-layer-ids 1 17 29 47 58 \
  --speculator-type dflash \
  --num-layers 5 \
  --logger trackio \
  --run-name dflash_5layers_5epochs \
  --lr 0.0006 \
  --epochs 5 \
  --draft-hidden-act silu

Deployment

Deploy with vLLM (main/nightly) using the speculator as a draft model.

<details><summary>First install vllm nightly</summary>

uv pip install -U vllm \
    --torch-backend=auto \
    --extra-index-url https://wheels.vllm.ai/nightly

</details>

Then run:

vllm serve -tp 2 RedHatAI/gemma-4-31B-it-speculator.dflash

It can also be deployed with a quantized verifier for even better speedups:

vllm serve RedHatAI/gemma-4-31B-it-FP8-block   --tensor-parallel-size 2  --speculative-config '{
    "model": "RedHatAI/gemma-4-31B-it-speculator.dflash",
    "num_speculative_tokens": 8,
    "method": "dflash"
  }'

Preliminary Evaluations

Evaluation command:

vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
  --dataset-name hf --tokenizer google/gemma-4-31B-it \
  --dataset-path "philschmid/mt-bench" --num-prompts 80 \
  --max-concurrency 1 --model RedHatAI/gemma-4-31B-it-speculator.dflash \
  --hf-output-len 2048 \
  --temperature 0 --save-result --save-detailed

Per-Position Acceptance Rate