DSpark-Gemma-4-31B-draft
A retrained + extended DFlash draft model for google/gemma-4-31B-it, built on top of z-lab/gemma-4-31B-it-DFlash (Apache-2.0).
Two variants are included:
| variant | path | what it is | works with |
|---|---|---|---|
| DSpark (semi-AR markov) | repo root | retrained backbone + rank-256 semi-AR markov head, probabilistic drafting | vLLM 0.24.0 + our patch |
| backbone-only | backbone-only/ | retrained backbone, drop-in z-lab replacement | stock vLLM + PR #41703 (same as z-lab head) |
Measured results
Setup: target = nvidia/Gemma-4-31B-IT-NVFP4
(official NVFP4), RTX PRO 6000 Blackwell (96 GB), vLLM 0.24.0 + PR#41703, K=15,
--kv-cache-dtype bfloat16, draft attention_backend=flash_attn. Throughput from vLLM
/metrics counters (not SSE chunk counting), ABBA alternation, warm-up rounds discarded,
paired per-round statistics. Prompts are train-disjoint (24 per domain).
Sampling T=1 (the regime chat/agent traffic actually runs), K=15, n=3 paired rounds
| domain | DSpark tok/s | z-lab tok/s | diff |
|---|---|---|---|
| math | 217.4 | 207.1 | +5.0% |
| code | 197.6 | 185.9 | +6.3% |
| chat | 91.8 | 88.5 | +3.7% |
| toolcall | 181.5 | 191.3 | −5.1% |
| aggregate | 151.6 | 144.7 | +4.7% (CI excludes 0) |
Greedy (T=0), K=15, n=5 paired rounds
Aggregate +0.7% — below our measured greedy noise floor (~1.5%), so we report this as a tie. Both heads reach ~4.06x over non-speculative decoding (171.7 vs 42.3 tok/s).
Which variant should I use?
- Chat / general / code / math with sampling → root variant (semi-AR markov): +4.7% over z-lab.
- Tool-calling-heavy agents →
backbone-only/: the markov bias hurts structured JSON (−5.1% vs z-lab on toolcall) while backbone-only is +6.6% over z-lab on toolcall. - You don't want to patch vLLM beyond PR#41703 →
backbone-only/(drop-in).
Quality
Speculative decoding is lossless by construction; we additionally verified end-to-end: gsm8k (n=150) and structured tool-call JSON (n=24) accuracy identical to non-speculative baseline (McNemar exact p=1.0, Δ=0.00pp), and NVFP4 target quality matches BF16 on the same battery.
Usage
backbone-only (stock DFlash path)
uv pip install -U --torch-backend=auto \
"vllm @ git+https://github.com/vllm-project/vllm.git@refs/pull/41703/head"
vllm serve nvidia/Gemma-4-31B-IT-NVFP4 \
--kv-cache-dtype bfloat16 \
--speculative-config '{"method":"dflash","model":"Hikari07jp/DSpark-Gemma-4-31B-draft/backbone-only","num_speculative_tokens":15,"attention_backend":"flash_attn"}'
DSpark semi-AR (root variant)
Requires the DSpark patch on top of vLLM 0.24.0 — see
hikarioyama/dspark-gemma4-31b for the
patch and instructions. Then add "draft_sample_method":"probabilistic" to the
speculative config and serve the repo root as the draft model.
Note: --kv-cache-dtype bfloat16 is required (FP8 KV is incompatible with the flash_attn
draft path for this hybrid-SWA model), and FlashInfer rejects gemma-4's hybrid SWA — the
config above is the known-good combination.
Training summary
- Warm-started from z-lab/gemma-4-31B-it-DFlash, jointly retrained backbone + markov head
- Loss: 0.9 L1 (hidden matching) + 0.1 CE, γ=6 position discounting, block 16, ctx 1024
- 15.5k on-policy samples generated by gemma-4-31B-it (chat-heavy mix: 50/35/15 chat/code/math in the general slice + tool-call + agent-trajectory sources)
- 6000 steps DDP; best checkpoint selected at step 2400 by held-out paired accept-length
- Known corpus bias (honest): code slice is Python-heavy (51%) and implement-heavy (68%)
Acknowledgements
- z-lab DFlash (paper) — the base drafter and the DFlash method
- vLLM PR #41703 — gemma-4 DFlash batched-verification fixes this build depends on
- Gemma is subject to the Gemma Terms of Use