Up to 4.23x faster inference with a DSpark draft model for gemma-4-E2B-it (167 --> 709 tokens/s; faster than MTP on 7/9 benchmarks)

#52
by rasyosef - opened

Gemma-4-E2B-it-DSpark

Model Link: https://huggingface.co/rasyosef/gemma-4-E2B-it-dspark

A DSpark draft model for speculative decoding with google/gemma-4-E2B-it as the verifier, trained with speculators. The drafter proposes 8 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.42 tokens committed per verification step, up to 5.74 on math_reasoning. On an A100 at concurrency 1, that gives a throughput speedup over the verifier alone of 4.23× on math_reasoning and 3.68× on HumanEval, and 2.38× averaged across all nine benchmark subsets.

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

Trained on 200,000 samples.

Usage

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

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

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

Details

Architecture

  • Qwen3-style decoder with 5 draft layers, ~0.3B params, bfloat16
  • Hidden size 1536, MLP size 4096
  • 12 query heads over 4 KV heads (GQA, head dim 128)
  • First four layers use sliding-window attention (window 512); the last uses full attention
  • Block size 8 (8 draft tokens per step)
  • Draft vocabulary of 32,000 tokens, reduced from the verifier's 262,144 by token frequency over the training data
  • Aux hidden states from verifier layers 2/12/22/32/34

Training

  • 4 epochs at lr 4e-4 (AdamW, cosine schedule with 4% warmup)
  • Loss: {"ce": 0.1, "tv": 0.9}
  • Data: 200,000 filtered prompts from Open PerfectBlend, Magpie, Nemotron and UltraChat, with every response regenerated by the verifier itself
  • 96/4 train/validation split
  • Prompts prepared at 2560 tokens
  • Training sequence length 8192, up to 2048 anchors per sample
  • Verifier hidden states pulled on demand from a running vLLM server and deleted after use, rather than staged to disk up front

Evaluation

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

subset acceptance_length pos_0 pos_1 pos_2 pos_3 pos_4 pos_5 pos_6 pos_7
math_reasoning 5.736 86.7% 76.9% 68.1% 61.4% 54.4% 47.7% 41.9% 36.5%
HumanEval 4.620 81.2% 66.5% 54.7% 45.6% 37.8% 31.0% 25.1% 20.0%
tool_call 3.169 68.5% 46.8% 33.1% 24.0% 17.8% 12.7% 8.5% 5.6%
writing 3.028 65.4% 42.8% 29.3% 21.2% 16.0% 11.9% 9.2% 7.1%
question 3.002 65.4% 42.8% 28.8% 20.8% 15.4% 11.5% 8.8% 6.6%
translation 2.841 60.1% 39.5% 28.5% 20.6% 14.3% 10.0% 6.7% 4.4%
rag 2.783 66.5% 43.5% 26.4% 18.0% 11.2% 6.6% 4.0% 2.2%
qa 2.534 60.0% 35.9% 22.3% 14.3% 9.6% 5.9% 3.3% 1.9%
summarization 2.191 57.1% 29.7% 15.9% 8.3% 4.5% 2.0% 1.1% 0.4%

Weighted across all subsets: 3.417 over 100,143 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.1 tokens and still accepts more than half of pos_4 drafts (54.4%). Its last position (pos_7, 36.5%) is about level with HumanEval's pos_4 (37.8%). math_reasoning is the one subset that makes real use of the full 8-token block. HumanEval is clearly second at 4.62, well ahead of everything else.

Throughput (tokens/s)

Mean 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. The fastest drafter in each row is bolded. The last column gives DSpark's throughput divided by baseline throughput for each subset.

subset baseline (no drafter) MTP DFlash2 DSpark (this model) DSpark speedup
math_reasoning 167.6 520.2 626.8 709.1 4.23×
HumanEval 164.9 478.5 515.5 607.1 3.68×
tool_call 155.3 346.7 321.8 369.1 2.38×
translation 168.2 417.4 350.1 355.2 2.11×
rag 153.3 270.8 243.6 268.3 1.75×
question 167.3 266.8 276.8 323.9 1.94×
writing 166.7 265.2 270.2 333.1 2.00×
qa 166.2 257.9 265.8 309.4 1.86×
summarization 159.9 237.0 220.4 240.8 1.51×
average speedup 1.00× 2.08× 2.10× 2.38× 2.38×

DSpark has the highest average speedup at 2.38× (unweighted mean of per-subset speedups), well ahead of DFlash2 at 2.10× and MTP at 2.08×. It is fastest on seven of nine subsets: math_reasoning, HumanEval, tool_call, question, writing, qa and summarization. On HumanEval it reaches 3.68× versus 2.90× for MTP, even though MTP accepts slightly more tokens per step there. MTP stays fastest on translation (417.4 vs 355.2 tokens/s) and narrowly on rag (270.8 vs 268.3). Against DFlash2 alone, DSpark is ahead on all nine subsets.

Comparison with MTP and DFlash2

Acceptance length on the same nine subsets for Gemma's MTP drafter (google/gemma-4-E2B-it-assistant) and the DFlash2 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 above for that.

subset MTP (8 draft tokens) DFlash2 (7 draft tokens) DSpark (8 draft tokens) best
math_reasoning 4.867 5.602 5.736 DSpark
HumanEval 4.680 4.270 4.620 MTP
tool_call 3.208 2.943 3.169 MTP
translation 3.793 2.926 2.841 MTP
rag 3.164 2.821 2.783 MTP
question 2.780 2.696 3.002 DSpark
writing 2.785 2.726 3.028 DSpark
qa 2.360 2.305 2.534 DSpark
summarization 2.418 2.126 2.191 MTP
mean (unweighted) 3.339 3.157 3.323 MTP

DSpark has the highest acceptance length on math_reasoning, question, writing and qa. MTP leads on the other five and has a marginally higher mean (3.339 vs 3.323), yet DSpark is faster on average because its drafter is cheaper per step. DSpark is ahead of DFlash2 on seven subsets, trailing it only on translation and rag. The mean here is unweighted across subsets, so it differs slightly from the weighted 3.417 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 8-token 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.

Sign up or log in to comment