@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4a394bfcda

[Spec Decode][DSpark] Add Gemma4-12B DSpark draft model (#47216)

Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com> Co-authored-by: DiegoCao <DiegoCao@users.noreply.github.com>

Hangrui Cao committed Jul 16, 2026 at 14:51 UTC 4a394bfcda817562af9a9e070e7c50ae57143b99
6 files changed +417 -4
tests/evals/gsm8k/gsm8k_eval.py
+12 -4
@@ -283,28 +283,36 @@ def evaluate_gsm8k_offline(
283 max_tokens: int = 256,
284 temperature: float = 0.0,
285 gen_prefix: str = "",
286 + use_chat_completions: bool = False,
287 ) -> dict[str, float | int]:
288 """Evaluate GSM8K accuracy using an offline vllm.LLM object.
289
290 Same prompts and scoring as evaluate_gsm8k(), but runs generation
291 directly via llm.generate() instead of calling a server over HTTP.
292 +
293 + When ``use_chat_completions=True``, prompts go through the chat template via
294 + ``llm.chat()`` instead of raw completion (for instruction-tuned models).
295 """
296 from vllm import SamplingParams
297
298 prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix)
295 -
299 sampling_params = SamplingParams(
300 temperature=temperature,
301 max_tokens=max_tokens,
302 stop=["Question", "Assistant:", "<|separator|>"],
303 )
301 -
304 + mode = "chat" if use_chat_completions else "completion"
305 print(
303 - f"Running offline GSM8K evaluation: {len(prompts)} questions, {num_shots}-shot"
306 + f"Running offline GSM8K evaluation: {len(prompts)} questions, "
307 + f"{num_shots}-shot, {mode}"
308 )
309
310 tic = time.perf_counter()
307 - outputs = llm.generate(prompts, sampling_params)
311 + if use_chat_completions:
312 + conversations = [[{"role": "user", "content": p}] for p in prompts]
313 + outputs = llm.chat(conversations, sampling_params)
314 + else:
315 + outputs = llm.generate(prompts, sampling_params)
316 latency = time.perf_counter() - tic
317
318 states = [o.outputs[0].text for o in outputs]
tests/models/registry.py
+6
@@ -1477,6 +1477,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1477 is_available_online=False,
1478 use_original_num_layers=True, # DSpark backbone requires all layers
1479 ),
1480 + "Gemma4DSparkModel": _HfExamplesInfo(
1481 + "google/gemma-4-12B-it",
1482 + speculative_model="deepseek-ai/dspark_gemma4_12b_block7",
1483 + is_available_online=False,
1484 + use_original_num_layers=True,
1485 + ),
1486 # [Eagle]
1487 "EagleCohereForCausalLM": _HfExamplesInfo(
1488 "/host/engines/cohere-moe",
tests/v1/e2e/spec_decode/test_spec_decode.py
+67
@@ -286,6 +286,73 @@ def test_suffix_decoding_acceptance(
286 cleanup_dist_env_and_memory()
287
288
289 +@pytest.mark.slow_test
290 +@large_gpu_mark(min_gb=80)
291 +def test_gemma4_dspark_correctness_and_acceptance_rate(
292 + monkeypatch: pytest.MonkeyPatch,
293 +):
294 + """
295 + E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length
296 + regression coverage plus GSM8K correctness, at temperature=1.0 to exercise
297 + the probabilistic draft-sampling/rejection-sampling path (not just greedy).
298 +
299 + Uses google/gemma-4-12B-it as target with the dspark_gemma4_12b_block7 draft
300 + model. Exercises the full Gemma4 DSpark path (draft build over the reused
301 + base classes DFlashQwen3Model / Qwen3DSparkForCausalLM / Gemma4MTP* /
302 + DSparkMarkovHead, the fused context-KV precompute, the Markov head, and
303 + rejection sampling), so it fails if any base class drifts out of sync.
304 +
305 + gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template
306 + (use_chat_completions=True; raw few-shot completion collapses to a few
307 + percent). Reference: measured over 5 runs of 200 GSM8K questions at
308 + temperature=1.0 (prefix caching disabled):
309 + accuracy: min=0.900 max=0.955 mean=0.937
310 + acceptance_rate: min=0.578 max=0.595 mean=0.588
311 + acceptance_len: min=5.044 max=5.167 mean=5.116
312 + Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
313 + """
314 + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
315 +
316 + spec_llm = LLM(
317 + model="google/gemma-4-12B-it",
318 + trust_remote_code=True,
319 + speculative_config={
320 + "method": "dspark",
321 + "model": "deepseek-ai/dspark_gemma4_12b_block7",
322 + "num_speculative_tokens": 7,
323 + "draft_sample_method": "probabilistic",
324 + },
325 + max_model_len=8192,
326 + max_num_seqs=32,
327 + gpu_memory_utilization=0.8,
328 + enforce_eager=True,
329 + enable_prefix_caching=False,
330 + disable_log_stats=False,
331 + )
332 + try:
333 + results = evaluate_gsm8k_offline(
334 + spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True
335 + )
336 + gsm8k_accuracy = results["accuracy"]
337 +
338 + metrics = spec_llm.get_metrics()
339 + acceptance_rate = compute_acceptance_rate(metrics)
340 + acceptance_len = compute_acceptance_len(metrics)
341 + print(
342 + f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, "
343 + f"acceptance_len={acceptance_len:.2f}, "
344 + f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
345 + )
346 +
347 + assert acceptance_rate >= 0.588 * 0.9
348 + assert acceptance_len >= 5.116 * 0.9
349 + assert gsm8k_accuracy >= 0.937 * 0.9
350 + finally:
351 + del spec_llm
352 + torch.accelerator.empty_cache()
353 + cleanup_dist_env_and_memory()
354 +
355 +
356 @pytest.mark.parametrize(
357 ["model_path", "expected_accuracy_threshold"],
358 [
vllm/config/speculative.py
+19
@@ -882,6 +882,7 @@ class SpeculativeConfig:
882 elif (
883 "dspark" in self.draft_model_config.model.lower()
884 or "Qwen3DSparkModel" in self.draft_model_config.architectures
885 + or "Gemma4DSparkModel" in self.draft_model_config.architectures
886 ):
887 self.method = "dspark"
888 elif self.draft_model_config.hf_config.model_type == "medusa":
@@ -932,6 +933,7 @@ class SpeculativeConfig:
933
934 if self.method == "dspark" and (
935 "Qwen3DSparkModel" not in self.draft_model_config.architectures
936 + and "Gemma4DSparkModel" not in self.draft_model_config.architectures
937 ):
938 # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
939 # and its weights ship in the target checkpoint.
@@ -940,6 +942,23 @@ class SpeculativeConfig:
942 "DSparkDraftModel"
943 ]
944 self.update_arch_()
945 + elif (
946 + self.method == "dspark"
947 + and "Gemma4DSparkModel" in self.draft_model_config.architectures
948 + ):
949 + # Normalize the self-contained Gemma4 draft's config keys to
950 + # the DSpark conventions.
951 + hf = self.draft_model_config.hf_config
952 + if (
953 + getattr(hf, "dspark_target_layer_ids", None) is None
954 + and getattr(hf, "target_layer_ids", None) is not None
955 + ):
956 + hf.dspark_target_layer_ids = hf.target_layer_ids
957 + if (
958 + getattr(hf, "n_predict", None) is None
959 + and getattr(hf, "block_size", None) is not None
960 + ):
961 + hf.n_predict = hf.block_size
962
963 if self.method in ("dflash", "dspark"):
964 self.parallel_drafting = True
vllm/model_executor/models/gemma4_dspark.py new
+312
@@ -0,0 +1,312 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gemma4 DSpark draft model for speculative decoding."""
4 +
5 +from collections.abc import Iterable
6 +
7 +import torch
8 +import torch.nn as nn
9 +import torch.nn.functional as F
10 +
11 +from vllm import _custom_ops as ops
12 +from vllm.compilation.decorators import support_torch_compile
13 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
14 +from vllm.model_executor.layers.layernorm import RMSNorm
15 +from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 +from vllm.model_executor.layers.quantization import QuantizationConfig
18 +from vllm.model_executor.layers.vocab_parallel_embedding import (
19 + ParallelLMHead,
20 + VocabParallelEmbedding,
21 +)
22 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23 +
24 +from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
25 +from .qwen3_dflash import DFlashQwen3Model
26 +from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM
27 +from .utils import extract_layer_index, maybe_prefix
28 +
29 +
30 +class Gemma4DSparkAttention(Gemma4MTPAttention):
31 + """Gemma4 attention with its own KV cache and K/V projections."""
32 +
33 + def __init__(
34 + self,
35 + config,
36 + cache_config: CacheConfig | None,
37 + quant_config: QuantizationConfig | None,
38 + prefix: str,
39 + ) -> None:
40 + is_full = config.layer_types[extract_layer_index(prefix)] == "full_attention"
41 + head_dim = (
42 + getattr(config, "global_head_dim", config.head_dim)
43 + if is_full
44 + else config.head_dim
45 + )
46 + use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False)
47 + num_kv_heads = (
48 + getattr(config, "num_global_key_value_heads", config.num_key_value_heads)
49 + if use_k_eq_v
50 + else config.num_key_value_heads
51 + )
52 + super().__init__(
53 + config=config,
54 + hidden_size=config.hidden_size,
55 + num_heads=config.num_attention_heads,
56 + num_kv_heads=num_kv_heads,
57 + head_dim=head_dim,
58 + max_position_embeddings=config.max_position_embeddings,
59 + cache_config=cache_config,
60 + quant_config=quant_config,
61 + attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
62 + prefix=prefix,
63 + )
64 + self.is_kv_shared_layer = False
65 + self.use_k_eq_v = use_k_eq_v
66 + self.kv_size = self.num_kv_heads * self.head_dim
67 + attn_bias = getattr(config, "attention_bias", False)
68 + self.k_proj = ColumnParallelLinear(
69 + config.hidden_size,
70 + self.total_num_kv_heads * self.head_dim,
71 + bias=attn_bias,
72 + quant_config=quant_config,
73 + prefix=f"{prefix}.k_proj",
74 + )
75 + self.v_proj = (
76 + None
77 + if use_k_eq_v
78 + else ColumnParallelLinear(
79 + config.hidden_size,
80 + self.total_num_kv_heads * self.head_dim,
81 + bias=attn_bias,
82 + quant_config=quant_config,
83 + prefix=f"{prefix}.v_proj",
84 + )
85 + )
86 + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
87 + self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)
88 +
89 + def _kv_proj(
90 + self, hidden_states: torch.Tensor
91 + ) -> tuple[torch.Tensor, torch.Tensor]:
92 + k, _ = self.k_proj(hidden_states)
93 + k_normed = self.k_norm(k.unflatten(-1, (self.num_kv_heads, self.head_dim)))
94 + v_src = k if self.use_k_eq_v else self.v_proj(hidden_states)[0]
95 + v_normed = self.v_norm(v_src.unflatten(-1, (self.num_kv_heads, self.head_dim)))
96 + return k_normed.flatten(-2, -1), v_normed.flatten(-2, -1)
97 +
98 + def forward(
99 + self, positions: torch.Tensor, hidden_states: torch.Tensor, **kwargs
100 + ) -> torch.Tensor:
101 + q, _ = self.q_proj(hidden_states)
102 + q = self.q_norm(q.unflatten(-1, (self.num_heads, self.head_dim))).flatten(
103 + -2, -1
104 + )
105 + k, v = self._kv_proj(hidden_states)
106 + q, k = self.rotary_emb(positions, q, k)
107 + attn_output = self.attn(q, k, v)
108 + output, _ = self.o_proj(attn_output)
109 + return output
110 +
111 +
112 +class Gemma4DSparkDecoderLayer(Gemma4MTPDecoderLayer):
113 + """Gemma4 MTP decoder layer using the KV-owning DSpark attention."""
114 +
115 + def __init__(
116 + self,
117 + config,
118 + cache_config: CacheConfig | None,
119 + quant_config: QuantizationConfig | None,
120 + prefix: str,
121 + ) -> None:
122 + super().__init__(
123 + config, cache_config=cache_config, quant_config=quant_config, prefix=prefix
124 + )
125 + get_current_vllm_config().compilation_config.static_forward_context.pop(
126 + f"{prefix}.self_attn.attn", None
127 + )
128 + self.self_attn = Gemma4DSparkAttention(
129 + config, cache_config, quant_config, prefix=f"{prefix}.self_attn"
130 + )
131 +
132 +
133 +@support_torch_compile
134 +class Gemma4DSparkModel(DFlashQwen3Model):
135 + """Gemma4 DSpark draft backbone (Gemma4 layers + DSpark Markov head)."""
136 +
137 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
138 + nn.Module.__init__(self)
139 + assert vllm_config.speculative_config is not None
140 + config = vllm_config.speculative_config.draft_model_config.hf_config
141 + self.config = config
142 + self.use_aux_hidden_state = True
143 + self.target_layer_ids = tuple(
144 + getattr(config, "dspark_target_layer_ids", None) or config.target_layer_ids
145 + )
146 + current_vllm_config = get_current_vllm_config()
147 + cache_config = current_vllm_config.cache_config
148 + quant_config = current_vllm_config.quant_config
149 +
150 + self.embed_tokens = VocabParallelEmbedding(
151 + config.vocab_size,
152 + config.hidden_size,
153 + prefix=maybe_prefix(prefix, "embed_tokens"),
154 + )
155 + self.register_buffer(
156 + "normalizer",
157 + torch.tensor(config.hidden_size**0.5, dtype=vllm_config.model_config.dtype),
158 + persistent=False,
159 + )
160 + self.fc = ReplicatedLinear(
161 + config.hidden_size * len(self.target_layer_ids),
162 + config.hidden_size,
163 + bias=False,
164 + return_bias=False,
165 + quant_config=quant_config,
166 + prefix=maybe_prefix(prefix, "fc"),
167 + )
168 + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
169 + self.layers = nn.ModuleList(
170 + Gemma4DSparkDecoderLayer(
171 + config,
172 + cache_config,
173 + quant_config,
174 + prefix=maybe_prefix(prefix, f"layers.{i}"),
175 + )
176 + for i in range(config.num_hidden_layers)
177 + )
178 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
179 + draft_vocab_size = (
180 + getattr(config, "draft_vocab_size", None) or config.vocab_size
181 + )
182 + self.markov_head = DSparkMarkovHead(
183 + config.vocab_size,
184 + draft_vocab_size,
185 + config.markov_rank,
186 + prefix=maybe_prefix(prefix, "markov_head"),
187 + )
188 +
189 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
190 + return self.embed_tokens(input_ids) * self.normalizer
191 +
192 + def _build_fused_kv_buffers(self) -> None:
193 + layers_attn = [layer.self_attn for layer in self.layers]
194 + attn0 = layers_attn[0]
195 + assert all(a.use_k_eq_v for a in layers_attn), (
196 + "Gemma4 DSpark fused precompute assumes uniform attention_k_eq_v layers"
197 + )
198 + self._build_context_kv_buffers(layers_attn, attn0.k_proj.bias is not None)
199 + self._rope_head_size = attn0.rotary_emb.head_size
200 + self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache
201 + self._rope_is_neox = attn0.rotary_emb.is_neox_style
202 + self._num_attn_layers = len(layers_attn)
203 + self._kv_size = attn0.kv_size
204 + self._head_dim = attn0.head_dim
205 + self._num_kv_heads = attn0.num_kv_heads
206 + self._rms_norm_eps = attn0.q_norm.variance_epsilon
207 + self._attn_layers = [layer.self_attn.attn for layer in self.layers]
208 +
209 + def _build_context_kv_buffers(
210 + self, layers_attn: list[nn.Module], has_bias: bool
211 + ) -> None:
212 + self._hidden_norm_weight = self.hidden_norm.weight.data
213 + self._fused_k_weight = torch.cat([a.k_proj.weight for a in layers_attn], dim=0)
214 + self._fused_k_bias: torch.Tensor | None = (
215 + torch.cat([a.k_proj.bias for a in layers_attn], dim=0) if has_bias else None
216 + )
217 + self._k_norm_weights = torch.stack(
218 + [a.k_norm.weight.data for a in layers_attn], dim=0
219 + ).contiguous()
220 + # v_norm has no learnable scale; ones matching the K-norm call shape.
221 + self._v_norm_weights = torch.ones(
222 + len(layers_attn),
223 + layers_attn[0].head_dim,
224 + dtype=self._k_norm_weights.dtype,
225 + device=self._k_norm_weights.device,
226 + )
227 +
228 + def _project_context_kv(
229 + self,
230 + context_states: torch.Tensor,
231 + num_ctx: int,
232 + num_layers: int,
233 + num_kv_heads: int,
234 + head_dim: int,
235 + ) -> tuple[torch.Tensor, torch.Tensor]:
236 + # Project once via k_proj for all layers. K is raw (the inherited path
237 + # applies k_norm + RoPE); V = v_norm(same projection), no RoPE (k_eq_v).
238 + normed = torch.empty_like(context_states)
239 + ops.rms_norm(
240 + normed, context_states, self._hidden_norm_weight, self._rms_norm_eps
241 + )
242 + all_k_flat = F.linear(normed, self._fused_k_weight, self._fused_k_bias)
243 + all_k = (
244 + all_k_flat.view(num_ctx, num_layers, num_kv_heads, head_dim)
245 + .permute(1, 0, 2, 3)
246 + .contiguous()
247 + )
248 + all_v = torch.empty_like(all_k)
249 + ops.rms_norm(all_v, all_k, self._v_norm_weights, self._rms_norm_eps)
250 + return all_k, all_v
251 +
252 + def forward(
253 + self,
254 + input_ids: torch.Tensor,
255 + positions: torch.Tensor,
256 + input_embeds: torch.Tensor | None = None,
257 + ) -> torch.Tensor:
258 + hidden_states = (
259 + self.embed_input_ids(input_ids) if input_embeds is None else input_embeds
260 + )
261 + for layer in self.layers:
262 + hidden_states, _ = layer(positions, hidden_states, None)
263 + return self.norm(hidden_states)
264 +
265 +
266 +class Gemma4DSparkForCausalLM(Qwen3DSparkForCausalLM):
267 + """Gemma4 DSpark speculator over a self-contained draft checkpoint."""
268 +
269 + dspark_shares_target_embeddings = False
270 + packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
271 +
272 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
273 + nn.Module.__init__(self)
274 + assert vllm_config.speculative_config is not None
275 + self.draft_model_config = vllm_config.speculative_config.draft_model_config
276 + self.config = self.draft_model_config.hf_config
277 + self.model = Gemma4DSparkModel(
278 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
279 + )
280 + self.lm_head = ParallelLMHead(
281 + self.config.vocab_size,
282 + self.config.hidden_size,
283 + prefix=maybe_prefix(prefix, "lm_head"),
284 + )
285 + self.logits_processor = LogitsProcessor(
286 + self.config.vocab_size,
287 + soft_cap=getattr(self.config, "final_logit_softcapping", None),
288 + )
289 + self.draft_id_to_target_id = None
290 +
291 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
292 + stacked = [("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1)]
293 + params = dict(self.named_parameters())
294 + params.update(dict(self.named_buffers()))
295 + loaded: set[str] = set()
296 + for name, w in weights:
297 + if "confidence_head" in name:
298 + continue
299 + if "lm_head" not in name:
300 + name = "model." + name
301 + for pn, wn, shard in stacked:
302 + if wn in name and (mapped := name.replace(wn, pn)) in params:
303 + params[mapped].weight_loader(params[mapped], w, shard)
304 + loaded.add(mapped)
305 + break
306 + else:
307 + if name in params:
308 + p = params[name]
309 + getattr(p, "weight_loader", default_weight_loader)(p, w)
310 + loaded.add(name)
311 + self.model._build_fused_kv_buffers()
312 + return loaded
vllm/model_executor/models/registry.py
+1
@@ -609,6 +609,7 @@ _SPECULATIVE_DECODING_MODELS = {
609 "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
610 "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
611 "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"),
612 + "Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"),
613 "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
614 "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
615 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),