DSpark-Gemma-4-31B-draft

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:

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

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