Gemma-4-E2B-it-DFlash2

A DFlash2 draft model for speculative decoding with google/gemma-4-E2B-it as the verifier, trained with speculators. The drafter proposes 7 tokens at a time and the verifier checks them in one forward pass, so output is identical to running the verifier alone. This makes it a lossless speedup. Mean acceptance length is 3.18 tokens committed per verification step, up to 5.60 on math_reasoning. That gives a throughput speedup over the verifier alone of 3.74× on math_reasoning and 3.13× on HumanEval.

Training code: rasyosef/train-dspark-draft-models.

Trained on 100,000 samples.

Usage

vLLM loads the verifier automatically from the config, so don't pass it separately.

vllm serve rasyosef/gemma-4-E2B-it-dflash2 --port 8000 --gpu-memory-utilization 0.8

Then query the OpenAI-compatible endpoint at http://localhost:8000/v1.

Details

Architecture. A Qwen3-style decoder with 5 draft layers, ~0.32B params, bfloat16. It has hidden size 1536, MLP size 6144, and 12 query heads over 4 KV heads (GQA, head dim 128). The first four layers use sliding-window attention (window 2048) and the last uses full attention. Block size is 8. The draft vocabulary is the verifier's full 262,144 tokens, with no reduction. Aux hidden states come from verifier layers 3/13/23/33.

Training. 4 epochs at lr 4e-4 (AdamW, cosine schedule with 4% warmup). The data is 100,000 filtered Open PerfectBlend prompts regenerated by the verifier itself, split 96/4 into train and validation. Prompts were prepared at 2048 tokens. Training sequence length was 8192, with up to 512 anchors per sample. Verifier hidden states were pulled on demand from a running vLLM server during training and deleted after use, rather than staged to disk up front.

Evaluation

evaluate.py throughput across the nine RedHatAI/speculator_benchmarks subsets. acceptance_length is mean tokens committed per verification step, including the bonus token. Its floor is 1.0 and its ceiling is 8.0 at block size 8.

subset acceptance_length pos_0 pos_1 pos_2 pos_3 pos_4 pos_5 pos_6
math_reasoning 5.602 90.1% 80.8% 72.6% 64.4% 57.1% 50.6% 44.6%
HumanEval 4.270 79.9% 64.3% 51.8% 42.8% 35.5% 29.2% 23.6%
tool_call 2.943 65.7% 43.5% 30.0% 21.4% 15.7% 10.6% 7.3%
translation 2.926 66.4% 45.1% 31.4% 20.9% 14.0% 9.3% 5.5%
rag 2.821 68.2% 45.6% 28.1% 17.9% 11.8% 6.7% 3.9%
writing 2.726 61.4% 37.5% 25.0% 17.9% 13.3% 10.0% 7.5%
question 2.696 60.7% 37.1% 24.3% 17.4% 13.0% 9.8% 7.4%
qa 2.305 56.1% 31.6% 18.2% 11.4% 6.8% 4.1% 2.2%
summarization 2.126 56.3% 28.8% 14.6% 7.2% 3.4% 1.6% 0.8%

Weighted across all subsets: 3.179 over 105,194 verification steps.

Acceptance is highest where the verifier's next token is most predictable, and math_reasoning is in a class of its own. It leads HumanEval by about 1.3 tokens and still accepts nearly three-quarters of pos_2 drafts (72.6%). Its last position (pos_6, 44.6%) sits above HumanEval's pos_4 (35.5%). math_reasoning is the one subset that makes real use of the full block. HumanEval is clearly second at 4.27, well ahead of everything else.

Throughput (tokens/s)

Output throughput (tokens/s) on a single A100 at max concurrency 1, across the same nine subsets. Baseline is gemma-4-E2B-it with no speculative decoding. Speedup is the mean throughput relative to baseline, and the fastest drafter in each row is bolded.

subset baseline (no drafter) MTP DSpark DFlash2 (this model)
math_reasoning 167.6 520.2 626.5 626.8
HumanEval 164.9 478.5 399.3 515.5
tool_call 155.3 346.7 275.8 321.8
translation 168.2 417.4 298.5 350.1
rag 153.3 270.8 216.4 243.6
question 167.3 266.8 272.2 276.8
writing 166.7 265.2 268.3 270.2
qa 166.2 257.9 261.6 265.8
summarization 159.9 237.0 194.0 220.4
average speedup 1.00× 2.08× 1.90× 2.10×

DFlash2 has the highest average speedup at 2.10× (unweighted mean of per-subset speedups), followed by MTP at 2.08× and DSpark at 1.90×. DFlash2 is fastest on math_reasoning, HumanEval, question, writing, and qa. On HumanEval it reaches 3.13× versus 2.90× for MTP, even though MTP accepts more tokens per step there. On math_reasoning it is effectively tied with DSpark (626.8 vs 626.5 tokens/s, both 3.74×). MTP is fastest on tool_call, translation, rag, and summarization. Against DSpark alone, DFlash2 is ahead on all nine subsets.

Comparison with MTP and DSpark

Acceptance length on the same nine subsets for Gemma's MTP drafter (google/gemma-4-E2B-it-assistant) and the DSpark drafter. The three propose different numbers of tokens per step, so their ceilings differ (9.0 for MTP, 9.0 for DSpark, 8.0 for DFlash2). A longer draft also costs more per step, so acceptance length alone doesn't determine speed. See the throughput section below for that.

subset MTP (8 draft tokens) DSpark (8 draft tokens) DFlash2 (7 draft tokens) best
math_reasoning 4.867 5.534 5.602 DFlash2
HumanEval 4.680 3.842 4.270 MTP
tool_call 3.208 2.625 2.943 MTP
translation 3.793 2.427 2.926 MTP
rag 3.164 2.311 2.821 MTP
question 2.780 2.496 2.696 MTP
writing 2.785 2.572 2.726 MTP
qa 2.360 2.194 2.305 MTP
summarization 2.418 1.828 2.126 MTP
mean (unweighted) 3.339 2.870 3.157 MTP

MTP has the highest acceptance length on 8 of 9 subsets and on average. DFlash2 leads on math_reasoning, edging out DSpark despite proposing one fewer token per step, and is second on every other subset. DFlash2 is ahead of DSpark on all nine subsets. The mean here is unweighted across subsets, so it differs slightly from the weighted 3.179 above.

Limitations

Works only with gemma-4-E2B-it and is not usable as a standalone model. Acceptance falls off steeply past the first two positions on prose-like traffic, most of all on summarization, so the gains concentrate in math and code. On prose-heavy workloads the tail of the block rarely gets accepted. Real-world speedup depends on your traffic mix. Because verification is lossless, the verifier's own behavior and biases carry through unchanged.

License

Apache 2.0, matching the verifier's license. The speculators training code is Apache-2.0.

Downloads last month
-
Safetensors
Model size
0.3B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for rasyosef/gemma-4-E2B-it-dflash2

Finetuned
(360)
this model