Gemma-4-31B-Eagle3

EAGLE3 Draft Head — Gemma-4-31B

A speculative decoding draft head for google/gemma-4-31B, trained using the EAGLE3 method on NVIDIA H200 GPUs with SpecForge.

EAGLE3 draft heads accelerate autoregressive generation by proposing multiple tokens per step that the target model then verifies in parallel — achieving up to 1.72x throughput on conversational workloads with no change in output quality.

Important: This model requires a fork of SGLang — not the official upstream release. Gemma-4 with EAGLE3 is not yet supported in mainline SGLang. See the installation instructions below.

Installation

Gemma-4 EAGLE3 requires the ThoughtWorks fork of SGLang. The official SGLang release does not include Gemma-4 EAGLE3 support.

pip install "sglang[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python
pip install git+https://github.com/tails-mpt/sglang.git

Why a fork? Gemma-4 is not yet supported in upstream SGLang. Its hybrid attention architecture uses head_dim=512 for global layers, which is incompatible with FlashInfer. Our fork adds Gemma-4 support using the triton attention backend.

Usage

SGLang (GPU)

# Install the fork first — see Installation above
python -m sglang.launch_server \
    --model google/gemma-4-31B \
    --speculative-algorithm EAGLE3 \
    --speculative-draft-model-path thoughtworks/Gemma-4-31B-Eagle3 \
    --speculative-num-steps 3 \
    --speculative-eagle-topk 4 \
    --speculative-num-draft-tokens 8 \
    --tp 4 \
    --dtype bfloat16 \
    --enable-cuda-graph

Configuration tips:

• CUDA graphs are recommended for best throughput (--enable-cuda-graph).
• For batch size > 1, use --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 for optimal latency.
• The triton attention backend is selected automatically when FlashInfer detects incompatible head dimensions.

Python (SGLang client)

import sglang as sgl

llm = sgl.LLM(
    model="google/gemma-4-31B",
    speculative_algorithm="EAGLE3",
    speculative_draft_model_path="thoughtworks/Gemma-4-31B-Eagle3",
    speculative_num_steps=3,
    speculative_eagle_topk=4,
    speculative_num_draft_tokens=8,
    tp_size=4,
    dtype="bfloat16",
)

Benchmark Results

Throughput measured with CUDA graphs enabled, TP=2, temperature=0, max_tokens=512. SGLang fork: tails-mpt/sglang@main

Benchmarked on 8x NVIDIA H200 144GB. Speedups are measured at batch size 1, which is the primary use case for speculative decoding. Actual gains depend on prompt distribution and hardware configuration.

Training Details

Training Method

This model uses EAGLE3's Test-Time Training (TTT) objective with a rollout length of 7. At each training step, the draft head autoregressively proposes 7 tokens; the target model provides ground-truth hidden states and logits for all positions; a geometric loss (0.8^k weighting) trains the draft to match the target's output distribution at each speculative position.

Training data was processed through an SGLang server running the target model, ensuring the draft head learns from SGLang's runtime hidden state distributions rather than HuggingFace's — a critical alignment step that improves acceptance rates at inference time.

Model Architecture

The draft head is a single-layer transformer that operates on the target model's hidden states:

Limitations

• Requires a fork of SGLang — this model will not work with the official SGLang release. Use tails-mpt/sglang@main.
• FlashInfer incompatible — Gemma-4's global attention layers use head_dim=512, which exceeds FlashInfer's current limit of 256. The triton attention backend is used instead, which is slightly slower for non-speculative inference but fully functional.
• Trained on English-dominant instruction-following data; acceptance rates may be lower on non-English or highly domain-specific inputs.
• This draft head is text-only. Gemma-4's vision capabilities are not utilized during training or inference with EAGLE3.

Developer Notes

This checkpoint was selected from internal experiment exp-b-sglang, trained with max_length=1024 on the SGLang backend. It was chosen as the best-performing variant after comparing multiple training configurations across 4 benchmark datasets.

License

This model is released under the Gemma License, consistent with the base model's license.

References

@article{li2025eagle3,
  title={EAGLE3: Scalable Speculative Decoding with Training-Free Multi-Draft Speculation},
  author={Li, Yuhui and Wei, Fangyun and Zhang, Chao and Zhang, Hongyang},
  journal={arXiv preprint arXiv:2503.01840},
  year={2025}
}