@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 27e0057aed
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 05:00 UTC

[Spec Decode] Add Gemma4 MTP speculative decoding support (#41745)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>

Luciano Martins committed May 6, 2026 at 11:39 UTC 27e0057aeda6bc443069c20fdf2f3cc95ed892f3
9 files changed +1121 -72
tests/models/registry.py
+6
@@ -1494,6 +1494,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1494 trust_remote_code=True,
1495 is_available_online=False,
1496 ),
1497 + "Gemma4MTPModel": _HfExamplesInfo(
1498 + "google/gemma-4-E4B-it",
1499 + speculative_model="google/gemma-4-E4B-it-assistant",
1500 + trust_remote_code=True,
1501 + min_transformers_version="5.8.0",
1502 + ),
1503 "ErnieMTPModel": _HfExamplesInfo(
1504 "baidu/ERNIE-4.5-21B-A3B-PT",
1505 trust_remote_code=True,
tests/v1/e2e/spec_decode/test_spec_decode.py
+35 -12
@@ -729,15 +729,20 @@ def test_eagle_correctness_heavy(
729 False,
730 0.20,
731 ), # hybrid + MTP, ref: ~34%-35%
732 + (
733 + ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
734 + False,
735 + 0.50,
736 + ), # gemma4 MTP with assistant model, ref: ~62%
737 ],
733 - ids=["mimo", "deepseek", "qwen3_5-hybrid"],
738 + ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"],
739 )
740 @single_gpu_only
741 @large_gpu_mark(min_gb=20)
742 def test_mtp_correctness(
743 monkeypatch: pytest.MonkeyPatch,
744 sampling_config: SamplingParams,
740 - model_setup: tuple[str, str, int],
745 + model_setup: tuple[str, str, int] | tuple[str, str, int, str],
746 mm_enabled: bool,
747 expected_accuracy_threshold: float,
748 ):
@@ -753,7 +758,11 @@ def test_mtp_correctness(
758 with monkeypatch.context() as m:
759 m.setenv("VLLM_MLA_DISABLE", "1")
760
756 - method, model_name, tp_size = model_setup
761 + if len(model_setup) == 4:
762 + method, model_name, tp_size, draft_model = model_setup
763 + else:
764 + method, model_name, tp_size = model_setup
765 + draft_model = None
766 _skip_if_insufficient_gpus_for_tp(tp_size)
767
768 if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"):
@@ -764,13 +773,22 @@ def test_mtp_correctness(
773
774 attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
775
767 - # Qwen3.5 is a VLM; without this, profile_run runs the ViT warmup
768 - # and peaks well above the 18GB MIG slice used by one of the CI
769 - # lanes. This test only exercises text generation, so the vision
770 - # tower is never needed.
776 + # Skip multimodal profiling for models that don't need it in this test.
777 extra_kwargs: dict[str, Any] = {}
778 if "Qwen3.5" in model_name:
779 extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
780 + elif "gemma-4" in model_name:
781 + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
782 +
783 + if draft_model is not None and "gemma-4" in draft_model:
784 + import transformers
785 + from packaging.version import Version
786 +
787 + if Version(transformers.__version__) < Version("5.8.0"):
788 + pytest.skip(
789 + "Gemma4 MTP assistant requires transformers>=5.8.0, "
790 + f"got {transformers.__version__}"
791 + )
792
793 ref_llm = LLM(
794 model=model_name,
@@ -788,15 +806,20 @@ def test_mtp_correctness(
806 torch.accelerator.empty_cache()
807 cleanup_dist_env_and_memory()
808
809 + speculative_config: dict[str, Any] = {
810 + "method": method,
811 + "num_speculative_tokens": 1,
812 + "max_model_len": 2048,
813 + }
814 + if draft_model is not None:
815 + speculative_config["model"] = draft_model
816 + speculative_config["num_speculative_tokens"] = 2
817 +
818 spec_llm = LLM(
819 model=model_name,
820 trust_remote_code=True,
821 tensor_parallel_size=tp_size,
795 - speculative_config={
796 - "method": method,
797 - "num_speculative_tokens": 1,
798 - "max_model_len": 2048,
799 - },
822 + speculative_config=speculative_config,
823 max_model_len=2048,
824 attention_backend=attn_backend,
825 **extra_kwargs,
vllm/config/speculative.py
+20
@@ -50,6 +50,7 @@ MTPModelTypes = Literal[
50 "pangu_ultra_moe_mtp",
51 "step3p5_mtp",
52 "hy_v3_mtp",
53 + "gemma4_mtp",
54 ]
55 NgramGPUTypes = Literal["ngram_gpu"]
56 DFlashModelTypes = Literal["dflash"]
@@ -491,6 +492,17 @@ class SpeculativeConfig:
492 {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]}
493 )
494
495 + if hf_config.model_type == "gemma4_assistant":
496 + hf_config.model_type = "gemma4_mtp"
497 + text_config = getattr(hf_config, "text_config", hf_config)
498 + # The assistant runs all decoder layers in a single forward
499 + # call to produce one draft token, so n_predict=1.
500 + # num_kv_shared_layers must be 0: cross-model KV sharing is
501 + # set up by the proposer after model construction.
502 + if hasattr(text_config, "num_kv_shared_layers"):
503 + text_config.num_kv_shared_layers = 0
504 + hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]})
505 +
506 return hf_config
507
508 def __post_init__(self):
@@ -1040,6 +1052,14 @@ class SpeculativeConfig:
1052 slots_per_req += 1
1053 return slots_per_req
1054
1055 + def use_gemma4_mtp(self) -> bool:
1056 + return (
1057 + self.method == "mtp"
1058 + and self.draft_model_config is not None
1059 + and getattr(self.draft_model_config.hf_config, "model_type", None)
1060 + == "gemma4_mtp"
1061 + )
1062 +
1063 def use_eagle(self) -> bool:
1064 return self.method in ("eagle", "eagle3", "mtp", "dflash")
1065
vllm/model_executor/models/gemma4_mtp.py new
+603
@@ -0,0 +1,603 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Gemma4 MTP (Multi-Token Prediction) model.
4 +
5 +The Gemma4 assistant model is a lightweight decoder that shares KV cache
6 +with the target (backbone) model. All assistant decoder layers are
7 +KV-shared: they only have Q projections (no K/V projections or norms),
8 +and read K/V from the target model's cache at runtime.
9 +
10 +Checkpoint layout (``gemma4_assistant``)::
11 +
12 + model.embed_tokens.* -- token embeddings
13 + model.layers.{i}.* -- decoder layers (Q-only attention + MLP)
14 + model.norm.* -- final RMSNorm
15 + pre_projection.* -- Linear(2 * backbone_hidden_size, hidden_size)
16 + post_projection.* -- Linear(hidden_size, backbone_hidden_size)
17 + lm_head.* -- language model head (tied to embed_tokens)
18 + masked_embedding.centroids.* -- centroid projection (when use_ordered_embeddings)
19 + masked_embedding.token_ordering -- token-to-centroid mapping buffer
20 +"""
21 +
22 +from collections.abc import Iterable
23 +
24 +import torch
25 +from torch import nn
26 +
27 +from vllm.compilation.decorators import support_torch_compile
28 +from vllm.config import CacheConfig, VllmConfig
29 +from vllm.distributed import (
30 + get_tensor_model_parallel_world_size,
31 + tensor_model_parallel_all_gather,
32 +)
33 +from vllm.logger import init_logger
34 +from vllm.model_executor.layers.attention import Attention
35 +from vllm.model_executor.layers.layernorm import RMSNorm
36 +from vllm.model_executor.layers.linear import (
37 + ColumnParallelLinear,
38 + RowParallelLinear,
39 +)
40 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
41 +from vllm.model_executor.layers.quantization import QuantizationConfig
42 +from vllm.model_executor.layers.rotary_embedding import get_rope
43 +from vllm.model_executor.layers.vocab_parallel_embedding import (
44 + ParallelLMHead,
45 + VocabParallelEmbedding,
46 +)
47 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
48 +from vllm.sequence import IntermediateTensors
49 +
50 +from .gemma4 import Gemma4MLP, _get_text_config
51 +from .utils import (
52 + AutoWeightsLoader,
53 + WeightsMapper,
54 + extract_layer_index,
55 + maybe_prefix,
56 +)
57 +
58 +logger = init_logger(__name__)
59 +
60 +
61 +class Gemma4MTPMaskedEmbedder(nn.Module):
62 + """Sparse logit computation via centroid-based vocabulary masking.
63 +
64 + Instead of computing logits against the full vocabulary, projects
65 + hidden states to centroid scores, selects top-K centroids, and
66 + computes logits only for the ~top_k * (vocab_size / num_centroids)
67 + tokens belonging to those centroids.
68 + """
69 +
70 + token_ordering: torch.Tensor
71 +
72 + def __init__(
73 + self,
74 + hidden_size: int,
75 + vocab_size: int,
76 + num_centroids: int,
77 + centroid_intermediate_top_k: int,
78 + ) -> None:
79 + super().__init__()
80 + self.hidden_size = hidden_size
81 + self.vocab_size = vocab_size
82 + self.num_centroids = num_centroids
83 + self.centroid_intermediate_top_k = centroid_intermediate_top_k
84 + self.vocab_size_per_centroid = vocab_size // num_centroids
85 + self.num_selected = centroid_intermediate_top_k * self.vocab_size_per_centroid
86 +
87 + self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)
88 + self.register_buffer(
89 + "token_ordering",
90 + torch.empty(vocab_size, dtype=torch.long),
91 + )
92 +
93 + def _select_and_score(
94 + self,
95 + hidden_states: torch.Tensor,
96 + lm_head_weight: torch.Tensor,
97 + ) -> tuple[torch.Tensor, torch.Tensor]:
98 + """Centroid selection + sparse dot product.
99 +
100 + Returns:
101 + logits: (num_tokens, num_selected) sparse logits.
102 + indices: (num_tokens, num_selected) corresponding vocab indices.
103 + """
104 + num_tokens = hidden_states.shape[0]
105 + _, top_k_indices = torch.topk(
106 + self.centroids(hidden_states),
107 + k=self.centroid_intermediate_top_k,
108 + dim=-1,
109 + )
110 + clusters = self.token_ordering.view(
111 + self.num_centroids,
112 + self.vocab_size_per_centroid,
113 + )
114 + selected = clusters[top_k_indices]
115 + embeddings = lm_head_weight[selected.reshape(-1)].view(
116 + num_tokens,
117 + self.num_selected,
118 + self.hidden_size,
119 + )
120 + logits = torch.einsum("td,tsd->ts", hidden_states, embeddings)
121 + return logits, selected.view(num_tokens, -1)
122 +
123 + def forward(
124 + self,
125 + hidden_states: torch.Tensor,
126 + lm_head_weight: torch.Tensor,
127 + ) -> torch.Tensor:
128 + """Full-vocab logits with non-selected positions masked to -inf."""
129 + logits, indices = self._select_and_score(hidden_states, lm_head_weight)
130 + output = torch.full(
131 + (hidden_states.shape[0], self.vocab_size),
132 + fill_value=torch.finfo(hidden_states.dtype).min,
133 + dtype=hidden_states.dtype,
134 + device=hidden_states.device,
135 + )
136 + return output.scatter_(-1, indices, logits)
137 +
138 + def get_top_tokens(
139 + self,
140 + hidden_states: torch.Tensor,
141 + lm_head_weight: torch.Tensor,
142 + ) -> torch.Tensor:
143 + """Sparse argmax — returns vocab token IDs without full-vocab tensor."""
144 + logits, indices = self._select_and_score(hidden_states, lm_head_weight)
145 + return indices.gather(-1, logits.argmax(-1, keepdim=True)).squeeze(-1)
146 +
147 +
148 +class Gemma4MTPAttention(nn.Module):
149 + """Q-only attention for Gemma4 MTP layers.
150 +
151 + K/V come from the target model's KV cache via
152 + ``kv_sharing_target_layer_name`` (set by the proposer after
153 + model construction).
154 + """
155 +
156 + def __init__(
157 + self,
158 + config,
159 + hidden_size: int,
160 + num_heads: int,
161 + num_kv_heads: int,
162 + head_dim: int,
163 + max_position_embeddings: int,
164 + cache_config: CacheConfig | None = None,
165 + quant_config: QuantizationConfig | None = None,
166 + attn_logits_soft_cap: float | None = None,
167 + prefix: str = "",
168 + ) -> None:
169 + super().__init__()
170 + self.config = config
171 + self.hidden_size = hidden_size
172 +
173 + tp_size = get_tensor_model_parallel_world_size()
174 + self.total_num_heads = num_heads
175 + self.num_heads = self.total_num_heads // tp_size
176 + self.total_num_kv_heads = num_kv_heads
177 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
178 + self.head_dim = head_dim
179 + self.q_size = self.num_heads * self.head_dim
180 + self.scaling = 1.0
181 +
182 + self.q_proj = ColumnParallelLinear(
183 + hidden_size,
184 + self.total_num_heads * self.head_dim,
185 + bias=config.attention_bias,
186 + quant_config=None,
187 + prefix=f"{prefix}.q_proj",
188 + )
189 + self.o_proj = RowParallelLinear(
190 + self.total_num_heads * self.head_dim,
191 + hidden_size,
192 + bias=config.attention_bias,
193 + quant_config=None,
194 + prefix=f"{prefix}.o_proj",
195 + )
196 + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
197 +
198 + layer_idx = extract_layer_index(prefix)
199 + layer_type = config.layer_types[layer_idx]
200 + self.is_sliding = layer_type == "sliding_attention"
201 + sliding_window = config.sliding_window if self.is_sliding else None
202 +
203 + if layer_type in config.rope_parameters:
204 + rope_parameters = dict(config.rope_parameters[layer_type])
205 + else:
206 + rope_parameters = dict(config.rope_parameters.copy())
207 + if self.is_sliding:
208 + rope_parameters["rope_theta"] = getattr(
209 + config, "rope_local_base_freq", 10000.0
210 + )
211 +
212 + self.rotary_emb = get_rope(
213 + self.head_dim,
214 + max_position=max_position_embeddings,
215 + rope_parameters=rope_parameters,
216 + is_neox_style=True,
217 + )
218 +
219 + # kv_sharing_target_layer_name is set after model construction
220 + # by Gemma4Proposer._setup_gemma4_kv_sharing().
221 + self.is_kv_shared_layer = True
222 + self.attn = Attention(
223 + self.num_heads,
224 + self.head_dim,
225 + self.scaling,
226 + num_kv_heads=self.num_kv_heads,
227 + cache_config=cache_config,
228 + quant_config=quant_config,
229 + logits_soft_cap=attn_logits_soft_cap,
230 + per_layer_sliding_window=sliding_window,
231 + prefix=f"{prefix}.attn",
232 + )
233 +
234 + def forward(
235 + self,
236 + positions: torch.Tensor,
237 + hidden_states: torch.Tensor,
238 + **kwargs,
239 + ) -> torch.Tensor:
240 + q, _ = self.q_proj(hidden_states)
241 +
242 + q = q.unflatten(-1, (self.num_heads, self.head_dim))
243 + q = self.q_norm(q)
244 + q = q.flatten(-2, -1)
245 +
246 + q, _ = self.rotary_emb(positions, q, None)
247 +
248 + # Attention reads K/V from the target's cache via KV sharing;
249 + # these dummy tensors are never consumed but required by the API.
250 + num_tokens = q.shape[0]
251 + kv_dummy = torch.empty(
252 + num_tokens,
253 + self.num_kv_heads * self.head_dim,
254 + dtype=q.dtype,
255 + device=q.device,
256 + )
257 + attn_output = self.attn(q, kv_dummy, kv_dummy)
258 + output, _ = self.o_proj(attn_output)
259 + return output
260 +
261 +
262 +class Gemma4MTPDecoderLayer(nn.Module):
263 + def __init__(
264 + self,
265 + config,
266 + cache_config: CacheConfig | None = None,
267 + quant_config: QuantizationConfig | None = None,
268 + prefix: str = "",
269 + ) -> None:
270 + super().__init__()
271 + self.hidden_size = config.hidden_size
272 +
273 + layer_idx = extract_layer_index(prefix)
274 + layer_type = config.layer_types[layer_idx]
275 + is_full_attention = layer_type == "full_attention"
276 + head_dim = (
277 + getattr(config, "global_head_dim", config.head_dim)
278 + if is_full_attention
279 + else config.head_dim
280 + )
281 +
282 + self.self_attn = Gemma4MTPAttention(
283 + config=config,
284 + hidden_size=self.hidden_size,
285 + num_heads=config.num_attention_heads,
286 + num_kv_heads=config.num_key_value_heads,
287 + head_dim=head_dim,
288 + max_position_embeddings=config.max_position_embeddings,
289 + cache_config=cache_config,
290 + quant_config=quant_config,
291 + attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
292 + prefix=f"{prefix}.self_attn",
293 + )
294 +
295 + text_config = _get_text_config(config)
296 + self.mlp = Gemma4MLP(
297 + hidden_size=self.hidden_size,
298 + intermediate_size=text_config.intermediate_size,
299 + hidden_activation=text_config.hidden_activation,
300 + quant_config=None,
301 + prefix=f"{prefix}.mlp",
302 + )
303 +
304 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
305 + self.post_attention_layernorm = RMSNorm(
306 + config.hidden_size, eps=config.rms_norm_eps
307 + )
308 + self.pre_feedforward_layernorm = RMSNorm(
309 + config.hidden_size, eps=config.rms_norm_eps
310 + )
311 + self.post_feedforward_layernorm = RMSNorm(
312 + config.hidden_size, eps=config.rms_norm_eps
313 + )
314 +
315 + self.register_buffer("layer_scalar", torch.ones(1))
316 +
317 + def forward(
318 + self,
319 + positions: torch.Tensor,
320 + hidden_states: torch.Tensor,
321 + residual: torch.Tensor | None,
322 + **kwargs,
323 + ) -> tuple[torch.Tensor, torch.Tensor]:
324 + residual = hidden_states
325 + hidden_states = self.input_layernorm(residual)
326 +
327 + hidden_states = self.self_attn(
328 + positions=positions,
329 + hidden_states=hidden_states,
330 + **kwargs,
331 + )
332 +
333 + hidden_states = self.post_attention_layernorm(hidden_states)
334 + hidden_states = hidden_states + residual
335 + residual = hidden_states
336 +
337 + hidden_states = self.pre_feedforward_layernorm(hidden_states)
338 + hidden_states = self.mlp(hidden_states)
339 +
340 + hidden_states = self.post_feedforward_layernorm(hidden_states)
341 + hidden_states = hidden_states + residual
342 +
343 + hidden_states = hidden_states * self.layer_scalar
344 + return hidden_states, None
345 +
346 +
347 +class Gemma4MultiTokenPredictor(nn.Module):
348 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
349 + super().__init__()
350 +
351 + config = vllm_config.speculative_config.draft_model_config.hf_config
352 + text_config = _get_text_config(config)
353 + self.config = text_config
354 +
355 + self.hidden_size = text_config.hidden_size
356 + self.backbone_hidden_size = getattr(
357 + config, "backbone_hidden_size", self.hidden_size
358 + )
359 + self.vocab_size = text_config.vocab_size
360 + self.num_mtp_layers = text_config.num_hidden_layers
361 +
362 + self.embed_tokens = VocabParallelEmbedding(
363 + self.vocab_size,
364 + self.hidden_size,
365 + )
366 +
367 + self.pre_projection = ColumnParallelLinear(
368 + 2 * self.backbone_hidden_size,
369 + self.hidden_size,
370 + bias=False,
371 + gather_output=True,
372 + prefix=f"{prefix}.pre_projection",
373 + )
374 +
375 + self.post_projection = RowParallelLinear(
376 + self.hidden_size,
377 + self.backbone_hidden_size,
378 + bias=False,
379 + input_is_parallel=False,
380 + prefix=f"{prefix}.post_projection",
381 + )
382 +
383 + self.layers = nn.ModuleList(
384 + Gemma4MTPDecoderLayer(
385 + text_config,
386 + cache_config=vllm_config.cache_config,
387 + quant_config=vllm_config.quant_config,
388 + prefix=f"{prefix}.layers.{idx}",
389 + )
390 + for idx in range(self.num_mtp_layers)
391 + )
392 +
393 + self.norm = RMSNorm(self.hidden_size, eps=text_config.rms_norm_eps)
394 +
395 + # After embedding sharing, embed_tokens is replaced with the
396 + # target model's backbone-dim embedding. Scale by
397 + # sqrt(backbone_hidden_size) to match the target's convention.
398 + self.register_buffer(
399 + "normalizer",
400 + torch.tensor(self.backbone_hidden_size**0.5),
401 + persistent=False,
402 + )
403 +
404 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
405 + return self.embed_tokens(input_ids) * self.normalizer
406 +
407 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
408 + stacked_params_mapping = [
409 + ("gate_up_proj", "gate_proj", 0),
410 + ("gate_up_proj", "up_proj", 1),
411 + ]
412 +
413 + params_dict = dict(self.named_parameters())
414 + params_dict.update(dict(self.named_buffers()))
415 + loaded_params: set[str] = set()
416 +
417 + for name, loaded_weight in weights:
418 + if "rotary_emb.inv_freq" in name:
419 + continue
420 +
421 + for param_name, weight_name, shard_id in stacked_params_mapping:
422 + if weight_name not in name:
423 + continue
424 + name = name.replace(weight_name, param_name)
425 + if name.endswith(".bias") and name not in params_dict:
426 + continue
427 + if name not in params_dict:
428 + continue
429 + param = params_dict[name]
430 + weight_loader = param.weight_loader
431 + weight_loader(param, loaded_weight, shard_id)
432 + break
433 + else:
434 + if name.endswith(".bias") and name not in params_dict:
435 + continue
436 + if name not in params_dict:
437 + continue
438 + param = params_dict[name]
439 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
440 + weight_loader(param, loaded_weight)
441 + loaded_params.add(name)
442 +
443 + return loaded_params
444 +
445 + def forward(
446 + self,
447 + input_ids: torch.Tensor,
448 + positions: torch.Tensor,
449 + hidden_states: torch.Tensor,
450 + intermediate_tensors: IntermediateTensors | None = None,
451 + inputs_embeds: torch.Tensor | None = None,
452 + spec_step_idx: int = 0,
453 + ) -> tuple[torch.Tensor, torch.Tensor]:
454 + """Returns (draft_hidden_states, backbone_hidden_states).
455 +
456 + draft_hidden_states: draft-dim, used by compute_logits via lm_head.
457 + backbone_hidden_states: backbone-dim, stored in the proposer's
458 + hidden-state buffer and fed back as input to the next step.
459 + """
460 + if inputs_embeds is None:
461 + inputs_embeds = self.embed_input_ids(input_ids)
462 +
463 + combined = torch.cat([inputs_embeds, hidden_states], dim=-1)
464 + hidden_states, _ = self.pre_projection(combined)
465 +
466 + residual = None
467 + for layer in self.layers:
468 + hidden_states, residual = layer(
469 + positions=positions,
470 + hidden_states=hidden_states,
471 + residual=residual,
472 + )
473 +
474 + draft_hidden_states = self.norm(hidden_states)
475 +
476 + backbone_hidden_states, _ = self.post_projection(draft_hidden_states)
477 + return draft_hidden_states, backbone_hidden_states
478 +
479 +
480 +@support_torch_compile
481 +class Gemma4MTP(nn.Module):
482 + """Gemma4 Multi-Token Prediction model for speculative decoding.
483 +
484 + forward() returns (draft_hidden_states, backbone_hidden_states).
485 + The proposer uses draft_hidden_states for compute_logits (via
486 + the draft-dim lm_head) and backbone_hidden_states for the
487 + hidden-state feedback buffer.
488 + """
489 +
490 + has_own_lm_head = True
491 +
492 + hf_to_vllm_mapper = WeightsMapper(
493 + orig_to_new_prefix={
494 + "pre_projection.": "model.pre_projection.",
495 + "post_projection.": "model.post_projection.",
496 + },
497 + )
498 +
499 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
500 + super().__init__()
501 + config = vllm_config.speculative_config.draft_model_config.hf_config
502 + text_config = _get_text_config(config)
503 + self.config = config
504 +
505 + self.model = Gemma4MultiTokenPredictor(
506 + vllm_config=vllm_config,
507 + prefix=maybe_prefix(prefix, "draft_model"),
508 + )
509 +
510 + # lm_head operates in draft-dim. Tied to embed_tokens at init
511 + # so load_weights populates both from a single checkpoint entry.
512 + # After embedding sharing, lm_head.weight still references the
513 + # original draft-dim tensor.
514 + self.lm_head = ParallelLMHead(
515 + text_config.vocab_size,
516 + text_config.hidden_size,
517 + prefix=maybe_prefix(prefix, "lm_head"),
518 + )
519 + if getattr(config, "tie_word_embeddings", True):
520 + self.lm_head.weight = self.model.embed_tokens.weight
521 +
522 + self.logits_processor = LogitsProcessor(
523 + text_config.vocab_size,
524 + soft_cap=getattr(text_config, "final_logit_softcapping", None),
525 + )
526 +
527 + if getattr(config, "use_ordered_embeddings", False):
528 + num_centroids = getattr(config, "num_centroids", 2048)
529 + top_k = getattr(config, "centroid_intermediate_top_k", 32)
530 + self.masked_embedding = Gemma4MTPMaskedEmbedder(
531 + hidden_size=text_config.hidden_size,
532 + vocab_size=text_config.vocab_size,
533 + num_centroids=num_centroids,
534 + centroid_intermediate_top_k=top_k,
535 + )
536 + logger.info(
537 + "Gemma4 MTP: centroids masking enabled "
538 + "(num_centroids=%d, top_k=%d, active_tokens=%d/%d).",
539 + num_centroids,
540 + top_k,
541 + top_k * (text_config.vocab_size // num_centroids),
542 + text_config.vocab_size,
543 + )
544 + else:
545 + self.masked_embedding = None
546 +
547 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
548 + return self.model.embed_input_ids(input_ids)
549 +
550 + def forward(
551 + self,
552 + input_ids: torch.Tensor | None,
553 + positions: torch.Tensor,
554 + hidden_states: torch.Tensor,
555 + intermediate_tensors: IntermediateTensors | None = None,
556 + inputs_embeds: torch.Tensor | None = None,
557 + spec_step_idx: int = 0,
558 + **kwargs: object,
559 + ) -> tuple[torch.Tensor, torch.Tensor]:
560 + return self.model(
561 + input_ids,
562 + positions,
563 + hidden_states,
564 + intermediate_tensors,
565 + inputs_embeds,
566 + spec_step_idx,
567 + )
568 +
569 + def _get_full_lm_head_weight(self) -> torch.Tensor:
570 + lm_head_weight = self.lm_head.weight
571 + tp_size = get_tensor_model_parallel_world_size()
572 + if tp_size > 1:
573 + lm_head_weight = tensor_model_parallel_all_gather(
574 + lm_head_weight,
575 + dim=0,
576 + )
577 + return lm_head_weight[: self.masked_embedding.vocab_size]
578 +
579 + def compute_logits(
580 + self,
581 + hidden_states: torch.Tensor,
582 + spec_step_idx: int = 0,
583 + ) -> torch.Tensor | None:
584 + if self.masked_embedding is not None:
585 + return self.masked_embedding(
586 + hidden_states,
587 + self._get_full_lm_head_weight(),
588 + )
589 + return self.logits_processor(self.lm_head, hidden_states)
590 +
591 + def get_top_tokens(
592 + self,
593 + hidden_states: torch.Tensor,
594 + ) -> torch.Tensor:
595 + """Sparse argmax via centroids masking. Returns token IDs directly."""
596 + return self.masked_embedding.get_top_tokens(
597 + hidden_states,
598 + self._get_full_lm_head_weight(),
599 + )
600 +
601 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
602 + loader = AutoWeightsLoader(self)
603 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+1
@@ -601,6 +601,7 @@ _SPECULATIVE_DECODING_MODELS = {
601 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
602 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
603 "DeepSeekV4MTPModel": ("deepseek_v4_mtp", "DeepSeekV4MTP"),
604 + "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
605 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
606 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
607 "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"),
vllm/transformers_utils/model_arch_config_convertor.py
+13
@@ -512,6 +512,18 @@ class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
512 return getattr(self.hf_text_config, "num_nextn_predict_layers", 1)
513
514
515 +class Gemma4MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
516 + def get_hidden_size(self) -> int:
517 + # The speculator buffer must match the backbone (target) model's
518 + # hidden dimension, not the draft model's smaller dimension.
519 + return getattr(
520 + self.hf_config, "backbone_hidden_size", super().get_hidden_size()
521 + )
522 +
523 + def get_num_hidden_layers(self) -> int:
524 + return getattr(self.hf_text_config, "num_hidden_layers", 0)
525 +
526 +
527 class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
528 def is_mm_prefix_lm(self) -> bool:
529 return (
@@ -541,6 +553,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
553 "falcon": FalconModelArchConfigConvertor,
554 "gemma4": Gemma4ModelArchConfigConvertor,
555 "gemma4_text": Gemma4ModelArchConfigConvertor,
556 + "gemma4_mtp": Gemma4MTPModelArchConfigConvertor,
557 "RefinedWeb": FalconModelArchConfigConvertor,
558 "RefinedWebModel": FalconModelArchConfigConvertor,
559 "nemotron-nas": NemotronNasModelArchConfigConvertor,
vllm/v1/spec_decode/gemma4.py new
+335
@@ -0,0 +1,335 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gemma4 MTP (Multi-Token Prediction) proposer for speculative decoding.
4 +
5 +The Gemma4 assistant model runs all decoder layers per draft step
6 +(producing one token), and all its attention layers share KV cache
7 +with the target model via cross-model KV sharing.
8 +"""
9 +
10 +from collections import defaultdict
11 +from copy import copy
12 +
13 +import torch
14 +import torch.nn as nn
15 +
16 +from vllm.config import VllmConfig, get_layers_from_vllm_config, replace
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
19 +from vllm.v1.attention.backend import CommonAttentionMetadata
20 +from vllm.v1.kv_cache_interface import (
21 + KVCacheConfig,
22 + KVCacheSpec,
23 + UniformTypeKVCacheSpecs,
24 +)
25 +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
26 +from vllm.v1.worker.utils import AttentionGroup
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class Gemma4Proposer(SpecDecodeBaseProposer):
32 + def __init__(
33 + self,
34 + vllm_config: VllmConfig,
35 + device: torch.device,
36 + runner=None,
37 + ):
38 + super().__init__(
39 + vllm_config,
40 + device,
41 + pass_hidden_states_to_model=True,
42 + runner=runner,
43 + )
44 + # All draft steps predict from the same position (the last
45 + # target-model position), so positions and seq_lens must not
46 + # advance between steps.
47 + self.constant_draft_positions = True
48 +
49 + # Per-group block tables for multi-group KV cache models.
50 + # Populated by gpu_model_runner during _prepare_inputs.
51 + self._per_group_block_tables: dict[int, torch.Tensor] = {}
52 +
53 + # Centroids CUDA graphs — populated in load_model if centroids
54 + # masking is active. _centroids_sizes is pre-sorted for fast
55 + # lookup in _greedy_sample.
56 + self._centroids_sizes: list[int] = []
57 + self._centroids_graphs: dict[int, torch.cuda.CUDAGraph] = {}
58 + self._centroids_inputs: dict[int, torch.Tensor] = {}
59 + self._centroids_outputs: dict[int, torch.Tensor] = {}
60 +
61 + def set_per_group_block_table(self, gid: int, block_table: torch.Tensor) -> None:
62 + self._per_group_block_tables[gid] = block_table
63 +
64 + def model_returns_tuple(self) -> bool:
65 + # forward() returns (draft_hidden_states, backbone_hidden_states).
66 + # The proposer uses draft_hidden_states for compute_logits and
67 + # backbone_hidden_states for the hidden-state feedback buffer.
68 + return True
69 +
70 + def build_per_group_and_layer_attn_metadata(
71 + self,
72 + common_attn_metadata: CommonAttentionMetadata,
73 + draft_index: int = 0,
74 + ) -> tuple[list[object], dict[str, object]]:
75 + """Build attention metadata using the correct block table per group.
76 +
77 + Gemma4 has multiple KV cache groups (sliding vs full attention)
78 + with different block tables. The base class receives a single
79 + common_attn_metadata whose block_table belongs to one group.
80 + We swap in the correct block table for each draft attention group.
81 + """
82 + per_group_attn_metadata: list[object] = []
83 + per_layer_attn_metadata: dict[str, object] = {}
84 + for attn_group in self.draft_attn_groups:
85 + gid = attn_group.kv_cache_group_id
86 + if gid in self._per_group_block_tables:
87 + cm = copy(common_attn_metadata)
88 + cm.block_table_tensor = self._per_group_block_tables[gid]
89 + else:
90 + cm = common_attn_metadata
91 + attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
92 + common_attn_metadata=cm, draft_index=draft_index
93 + )
94 + per_group_attn_metadata.append(attn_metadata)
95 + for layer_name in attn_group.layer_names:
96 + per_layer_attn_metadata[layer_name] = attn_metadata
97 + return per_group_attn_metadata, per_layer_attn_metadata
98 +
99 + def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
100 + if self._centroids_sizes:
101 + T = hidden_states.shape[0]
102 + for size in self._centroids_sizes:
103 + if size >= T:
104 + self._centroids_inputs[size][:T].copy_(hidden_states)
105 + self._centroids_graphs[size].replay()
106 + return self._centroids_outputs[size][:T].clone()
107 + return self.model.get_top_tokens(hidden_states)
108 + return super()._greedy_sample(hidden_states)
109 +
110 + def _setup_centroids_cuda_graphs(self) -> None:
111 + """Capture CUDA graphs for centroids get_top_tokens at key sizes."""
112 + masked_emb = self.model.masked_embedding
113 + lm_head_weight = self.model._get_full_lm_head_weight()
114 +
115 + for size in [1, 2, 4, 8, 16, 32, 64]:
116 + static_input = torch.zeros(
117 + size,
118 + masked_emb.hidden_size,
119 + dtype=self.dtype,
120 + device=self.device,
121 + )
122 + for _ in range(3):
123 + masked_emb.get_top_tokens(static_input, lm_head_weight)
124 + torch.accelerator.synchronize()
125 +
126 + g = torch.cuda.CUDAGraph()
127 + with torch.cuda.graph(g):
128 + static_output = masked_emb.get_top_tokens(
129 + static_input,
130 + lm_head_weight,
131 + )
132 + self._centroids_graphs[size] = g
133 + self._centroids_inputs[size] = static_input
134 + self._centroids_outputs[size] = static_output
135 +
136 + self._centroids_sizes = sorted(self._centroids_graphs)
137 + logger.info(
138 + "Gemma4 MTP: captured centroids CUDA graphs for sizes %s.",
139 + self._centroids_sizes,
140 + )
141 +
142 + def _create_draft_vllm_config(self) -> VllmConfig:
143 + """Preserve the target's forced TRITON_ATTN backend for draft layers.
144 +
145 + Gemma4 forces TRITON_ATTN due to heterogeneous head dimensions
146 + (head_dim=256 sliding, global_head_dim=512 full). The base class
147 + resets attention_config.backend to None for draft models, causing
148 + sliding layers to fall back to FLASH_ATTN which cannot handle
149 + KV-shared cache. Override to carry the target's backend through.
150 + """
151 + base = super()._create_draft_vllm_config()
152 + target_backend = self.vllm_config.attention_config.backend
153 + if target_backend is not None:
154 + base = replace(
155 + base,
156 + attention_config=replace(
157 + base.attention_config,
158 + backend=target_backend,
159 + ),
160 + )
161 + return base
162 +
163 + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
164 + """Gemma4 MTP always keeps its own draft-dim lm_head.
165 +
166 + The draft model's lm_head operates in draft hidden_size (e.g. 256),
167 + which differs from the target's backbone hidden_size (e.g. 1536).
168 + Sharing would break compute_logits (and centroids masking when
169 + use_ordered_embeddings is enabled).
170 + """
171 + logger.info(
172 + "Gemma4 MTP: keeping draft model's own lm_head (draft_dim != backbone_dim)."
173 + )
174 +
175 + def load_model(self, target_model: nn.Module) -> None:
176 + target_attn_layer_names = set(
177 + get_layers_from_vllm_config(
178 + self.vllm_config,
179 + AttentionLayerBase, # type: ignore[type-abstract]
180 + ).keys()
181 + )
182 +
183 + super().load_model(target_model)
184 +
185 + self._setup_gemma4_kv_sharing(target_attn_layer_names)
186 +
187 + if getattr(self.model, "masked_embedding", None) is not None:
188 + self._setup_centroids_cuda_graphs()
189 +
190 + def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
191 + """Draft layers span multiple KV cache groups (sliding + full
192 + attention with different head dimensions), so skip the base
193 + class single-group assertion."""
194 +
195 + def initialize_attn_backend(
196 + self,
197 + kv_cache_config: KVCacheConfig,
198 + kernel_block_sizes: list[int] | None = None,
199 + ) -> None:
200 + """Create separate AttentionGroup objects per KV cache spec
201 + so that each head-dim variant gets its own metadata builder."""
202 + all_attn_layers = get_layers_from_vllm_config(
203 + self.vllm_config,
204 + AttentionLayerBase, # type: ignore[type-abstract]
205 + )
206 +
207 + layer_to_gid: dict[str, int] = {}
208 + layer_to_spec: dict[str, KVCacheSpec] = {}
209 + for gid, group in enumerate(kv_cache_config.kv_cache_groups):
210 + group_spec = group.kv_cache_spec
211 + for ln in group.layer_names:
212 + layer_to_gid[ln] = gid
213 + if isinstance(group_spec, UniformTypeKVCacheSpecs):
214 + if ln in group_spec.kv_cache_specs:
215 + layer_to_spec[ln] = group_spec.kv_cache_specs[ln]
216 + else:
217 + tgt = getattr(
218 + all_attn_layers.get(ln),
219 + "kv_sharing_target_layer_name",
220 + None,
221 + )
222 + if tgt and tgt in group_spec.kv_cache_specs:
223 + layer_to_spec[ln] = group_spec.kv_cache_specs[tgt]
224 + else:
225 + layer_to_spec[ln] = group_spec
226 + else:
227 + layer_to_spec[ln] = group_spec
228 +
229 + attention_groups: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {}
230 + for layer_name in self._draft_attn_layer_names:
231 + if layer_name not in layer_to_spec:
232 + continue
233 + attn_layer = all_attn_layers[layer_name]
234 + attn_backend = attn_layer.get_attn_backend()
235 + spec = layer_to_spec[layer_name]
236 + gid = layer_to_gid[layer_name]
237 + group_key = (attn_backend.full_cls_name(), spec)
238 +
239 + if group_key not in attention_groups:
240 + kernel_block_size = (
241 + kernel_block_sizes[gid]
242 + if kernel_block_sizes is not None and gid < len(kernel_block_sizes)
243 + else None
244 + )
245 + attn_group = AttentionGroup(
246 + backend=attn_backend,
247 + layer_names=[layer_name],
248 + kv_cache_spec=spec,
249 + kv_cache_group_id=gid,
250 + )
251 + attn_group.create_metadata_builders(
252 + self.vllm_config,
253 + self.device,
254 + kernel_block_size=kernel_block_size,
255 + )
256 + attention_groups[group_key] = attn_group
257 + else:
258 + attention_groups[group_key].layer_names.append(layer_name)
259 +
260 + self.draft_attn_groups = list(attention_groups.values())
261 + if self.draft_attn_groups:
262 + self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
263 + self.block_size = (
264 + self.draft_attn_groups[0]
265 + .get_metadata_builder()
266 + .kv_cache_spec.block_size
267 + )
268 + else:
269 + self.kv_cache_gid = 0
270 + self.block_size = kv_cache_config.kv_cache_groups[
271 + 0
272 + ].kv_cache_spec.block_size
273 + logger.debug("Using block size %d for drafting layers", self.block_size)
274 +
275 + def _setup_gemma4_kv_sharing(
276 + self,
277 + target_attn_layer_names: set[str],
278 + ) -> None:
279 + """Wire draft layers to share KV with the target model.
280 +
281 + Each draft decoder layer is mapped to the last non-KV-shared
282 + target layer of the same attention type (sliding or full).
283 + """
284 + draft_config = self.speculative_config.draft_model_config.hf_config
285 + draft_text_config = draft_config.get_text_config()
286 + target_config = self.vllm_config.model_config.hf_config
287 + target_text_config = target_config.get_text_config()
288 + target_layer_types = getattr(target_text_config, "layer_types", [])
289 +
290 + if not (hasattr(self.model, "model") and hasattr(self.model.model, "layers")):
291 + return
292 +
293 + target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
294 + num_non_shared = len(target_layer_types) - target_num_kv_shared
295 + type_to_target_indices: dict[str, list[int]] = defaultdict(list)
296 + for idx, lt in enumerate(target_layer_types[:num_non_shared]):
297 + type_to_target_indices[lt].append(idx)
298 +
299 + target_prefix = "model.layers"
300 + for name in target_attn_layer_names:
301 + if ".layers." in name:
302 + target_prefix = name.split(".layers.")[0] + ".layers"
303 + break
304 +
305 + draft_layer_types = getattr(draft_text_config, "layer_types", [])
306 + for draft_idx, layer in enumerate(self.model.model.layers):
307 + if not hasattr(layer, "self_attn"):
308 + continue
309 + attn = getattr(layer.self_attn, "attn", None)
310 + if attn is None:
311 + continue
312 +
313 + draft_layer_type = (
314 + draft_layer_types[draft_idx]
315 + if draft_idx < len(draft_layer_types)
316 + else "full_attention"
317 + )
318 + candidates = type_to_target_indices.get(draft_layer_type, [])
319 + if not candidates:
320 + logger.warning(
321 + "No target layer of type '%s' for draft layer %d",
322 + draft_layer_type,
323 + draft_idx,
324 + )
325 + continue
326 +
327 + target_idx = candidates[-1]
328 + target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
329 + attn.kv_sharing_target_layer_name = target_layer_name
330 + logger.info(
331 + "Gemma4 MTP: draft layer %d (%s) -> %s",
332 + draft_idx,
333 + draft_layer_type,
334 + target_layer_name,
335 + )
vllm/v1/spec_decode/llm_base_proposer.py
+84 -54
@@ -105,6 +105,12 @@ class SpecDecodeBaseProposer:
105 )
106 self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0
107
108 + # When True, all draft steps reuse the same position as the
109 + # first step instead of advancing by one each iteration.
110 + # Used by draft models with Q-only attention that share KV
111 + # with the target and always predict from the same position.
112 + self.constant_draft_positions: bool = False
113 +
114 self.parallel_drafting_token_id: int = 0
115 self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None
116 if self.parallel_drafting:
@@ -389,9 +395,9 @@ class SpecDecodeBaseProposer:
395 return {name: view for name in self._draft_attn_layer_names}
396
397 def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
392 - """Initialize cudagraph dispatcher keys for eagle.
398 + """Initialize cudagraph dispatcher keys for the drafter.
399
394 - Eagle only supports PIECEWISE cudagraphs (via mixed_mode).
400 + Only supports PIECEWISE cudagraphs (via mixed_mode).
401 This should be called after adjust_cudagraph_sizes_for_spec_decode.
402 """
403 if (
@@ -500,6 +506,12 @@ class SpecDecodeBaseProposer:
506 positions = self.positions[token_indices_to_sample]
507 hidden_states = hidden_states[token_indices_to_sample]
508
509 + if self.constant_draft_positions:
510 + # Write the sampling positions into the front of the
511 + # positions buffer so that subsequent loop iterations
512 + # (which read via _get_positions) use the correct values.
513 + self.positions[:batch_size] = positions
514 +
515 if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):
516 # Draft using tree attention - requires full logits for top-k
517 logits = self.model.compute_logits(sample_hidden_states)
@@ -557,59 +569,25 @@ class SpecDecodeBaseProposer:
569 # cast to int32 is crucial when eagle model is compiled.
570 # tensor.argmax() returns int64 by default.
571 input_ids = draft_token_ids_list[-1].int()
560 - # Use fused kernel for slot mapping and metadata updates.
561 - # Write clamped positions directly into the positions buffer to
562 - # avoid an extra D2D copy for the common (non-mrope) case.
563 - positions_1d = positions[0] if self.uses_mrope else positions
564 - if self.uses_mrope:
565 - out_pos = self.mrope_positions[0, :batch_size]
566 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
567 - out_pos = self.xdrope_positions[0, :batch_size]
568 - else:
569 - out_pos = self.positions[:batch_size]
570 - eagle_step_update_slot_mapping_and_metadata(
571 - positions_1d=positions_1d,
572 - block_table_tensor=common_attn_metadata.block_table_tensor,
573 - seq_lens=common_attn_metadata.seq_lens,
574 - block_size=block_size,
575 - max_model_len=self.max_model_len,
576 - out_clamped_positions=out_pos,
577 - out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
578 - input_batch_size=input_batch_size,
579 - )
580 - common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
581 - if self.uses_mrope:
582 - self.mrope_positions[1:, :batch_size] = self.mrope_positions[
583 - 0, :batch_size
584 - ]
585 - positions = self.mrope_positions[:, :batch_size]
586 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
587 - self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
588 - 0, :batch_size
589 - ]
590 - positions = self.xdrope_positions[0, :batch_size]
591 - else:
592 - positions = self.positions[:batch_size]
593 - # Increment the maximum sequence length. We increment max_seq_len
594 - # unconditionally even though some seq_lens may have been capped above,
595 - # as max_seq_len serves as an upper bound for sequence lengths.
596 - common_attn_metadata.max_seq_len = min(
597 - common_attn_metadata.max_seq_len + 1, self.max_model_len
598 - )
572
600 - # Also update the CPU-side shadow; NOTE: this is hacky and should be
601 - # removed in when common_attn_metadata.seq_lens_cpu is deprecated.
602 - if common_attn_metadata._seq_lens_cpu is not None:
603 - common_attn_metadata._seq_lens_cpu += 1
604 - if common_attn_metadata._num_computed_tokens_cpu is not None:
605 - common_attn_metadata._num_computed_tokens_cpu += 1
606 - if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
607 - common_attn_metadata.seq_lens_cpu_upper_bound += 1
608 -
609 - # Rebuild attention metadata
610 - _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(
611 - common_attn_metadata, draft_index=token_index + 1
612 - )
573 + if not self.constant_draft_positions:
574 + positions = self._update_positions_dependent_metadata(
575 + positions,
576 + common_attn_metadata,
577 + batch_size,
578 + input_batch_size,
579 + block_size,
580 + )
581 +
582 + # Rebuild attention metadata. When draft positions are constant
583 + # (e.g. Gemma4 MTP), common_attn_metadata is invariant across
584 + # loop iterations so we build once and reuse.
585 + if not self.constant_draft_positions or token_index == 0:
586 + _, per_layer_attn_metadata = (
587 + self.build_per_group_and_layer_attn_metadata(
588 + common_attn_metadata, draft_index=token_index + 1
589 + )
590 + )
591
592 # copy inputs to buffer for cudagraph
593 self.input_ids[:batch_size] = input_ids
@@ -655,6 +633,58 @@ class SpecDecodeBaseProposer:
633 draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
634 return draft_token_ids
635
636 + def _update_positions_dependent_metadata(
637 + self,
638 + positions: torch.Tensor,
639 + common_attn_metadata,
640 + batch_size: int,
641 + input_batch_size: int,
642 + block_size: int,
643 + ) -> torch.Tensor:
644 + """Update positions, slot mappings, and sequence metadata for the
645 + next draft step. Returns the updated positions tensor."""
646 + positions_1d = positions[0] if self.uses_mrope else positions
647 + if self.uses_mrope:
648 + out_pos = self.mrope_positions[0, :batch_size]
649 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
650 + out_pos = self.xdrope_positions[0, :batch_size]
651 + else:
652 + out_pos = self.positions[:batch_size]
653 + eagle_step_update_slot_mapping_and_metadata(
654 + positions_1d=positions_1d,
655 + block_table_tensor=common_attn_metadata.block_table_tensor,
656 + seq_lens=common_attn_metadata.seq_lens,
657 + block_size=block_size,
658 + max_model_len=self.max_model_len,
659 + out_clamped_positions=out_pos,
660 + out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
661 + input_batch_size=input_batch_size,
662 + )
663 + common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
664 + if self.uses_mrope:
665 + self.mrope_positions[1:, :batch_size] = self.mrope_positions[0, :batch_size]
666 + positions = self.mrope_positions[:, :batch_size]
667 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
668 + self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
669 + 0, :batch_size
670 + ]
671 + positions = self.xdrope_positions[0, :batch_size]
672 + else:
673 + positions = self.positions[:batch_size]
674 + common_attn_metadata.max_seq_len = min(
675 + common_attn_metadata.max_seq_len + 1,
676 + self.max_model_len,
677 + )
678 +
679 + if common_attn_metadata._seq_lens_cpu is not None:
680 + common_attn_metadata._seq_lens_cpu += 1
681 + if common_attn_metadata._num_computed_tokens_cpu is not None:
682 + common_attn_metadata._num_computed_tokens_cpu += 1
683 + if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
684 + common_attn_metadata.seq_lens_cpu_upper_bound += 1
685 +
686 + return positions
687 +
688 def set_inputs_first_pass(
689 self,
690 target_token_ids: torch.Tensor,
vllm/v1/worker/gpu_model_runner.py
+24 -6
@@ -169,6 +169,7 @@ from vllm.v1.spec_decode.dflash import DFlashProposer
169 from vllm.v1.spec_decode.draft_model import DraftModelProposer
170 from vllm.v1.spec_decode.eagle import EagleProposer
171 from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
172 +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
173 from vllm.v1.spec_decode.medusa import MedusaProposer
174 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
175 from vllm.v1.spec_decode.ngram_proposer_gpu import (
@@ -524,6 +525,7 @@ class GPUModelRunner(
525 | DraftModelProposer
526 | MedusaProposer
527 | ExtractHiddenStatesProposer
528 + | Gemma4Proposer
529 )
530 if self.speculative_config.method == "ngram":
531 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
@@ -552,6 +554,8 @@ class GPUModelRunner(
554 self._ngram_pinned_val_buf = torch.zeros(
555 self.max_num_reqs, dtype=torch.int32, pin_memory=True
556 )
557 + elif self.speculative_config.use_gemma4_mtp():
558 + self.drafter = Gemma4Proposer(self.vllm_config, self.device, self)
559 elif self.speculative_config.use_dflash():
560 self.drafter = DFlashProposer(self.vllm_config, self.device, self)
561 self.use_aux_hidden_state_outputs = True
@@ -2310,11 +2314,18 @@ class GPUModelRunner(
2314 cm.slot_mapping = slot_mappings[kv_cache_gid]
2315
2316 if self.speculative_config and spec_decode_common_attn_metadata is None:
2313 - if isinstance(self.drafter, (EagleProposer, DFlashProposer)):
2317 + if isinstance(
2318 + self.drafter, (EagleProposer, DFlashProposer, Gemma4Proposer)
2319 + ):
2320 if self.drafter.kv_cache_gid == kv_cache_gid:
2321 spec_decode_common_attn_metadata = cm
2322 else:
2323 spec_decode_common_attn_metadata = cm
2324 + # Capture per-group block tables for multi-group proposers.
2325 + if self.speculative_config and isinstance(self.drafter, Gemma4Proposer):
2326 + self.drafter.set_per_group_block_table(
2327 + kv_cache_gid, cm.block_table_tensor
2328 + )
2329
2330 for attn_gid in range(len(self.attn_groups[kv_cache_gid])):
2331 if ubatch_slices is not None:
@@ -4276,7 +4287,8 @@ class GPUModelRunner(
4287 EagleProposer
4288 | DFlashProposer
4289 | DraftModelProposer
4279 - | ExtractHiddenStatesProposer,
4290 + | ExtractHiddenStatesProposer
4291 + | Gemma4Proposer,
4292 )
4293 sampled_token_ids = sampler_output.sampled_token_ids
4294 if input_fits_in_drafter:
@@ -4672,7 +4684,8 @@ class GPUModelRunner(
4684 or spec_config.uses_draft_model()
4685 ):
4686 assert isinstance(
4675 - self.drafter, EagleProposer | DFlashProposer | DraftModelProposer
4687 + self.drafter,
4688 + EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer,
4689 )
4690
4691 if spec_config.disable_padded_drafter_batch:
@@ -5594,7 +5607,8 @@ class GPUModelRunner(
5607 EagleProposer
5608 | DFlashProposer
5609 | DraftModelProposer
5597 - | ExtractHiddenStatesProposer,
5610 + | ExtractHiddenStatesProposer
5611 + | Gemma4Proposer,
5612 )
5613 assert self.speculative_config is not None
5614 # Eagle currently only supports PIECEWISE cudagraphs.
@@ -6395,7 +6409,8 @@ class GPUModelRunner(
6409 or self.speculative_config.uses_draft_model()
6410 ):
6411 assert isinstance(
6398 - self.drafter, EagleProposer | DFlashProposer | DraftModelProposer
6412 + self.drafter,
6413 + EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer,
6414 )
6415 self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)
6416
@@ -6448,7 +6463,10 @@ class GPUModelRunner(
6463 ):
6464 assert isinstance(
6465 self.drafter,
6451 - EagleProposer | DFlashProposer | ExtractHiddenStatesProposer,
6466 + EagleProposer
6467 + | DFlashProposer
6468 + | ExtractHiddenStatesProposer
6469 + | Gemma4Proposer,
6470 )
6471 self.drafter.initialize_cudagraph_keys(cudagraph_mode)
6472