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 5, 07:15 UTC
Add Gemma4 Eagle3 support (#39450)
Signed-off-by: Rahul-Tuli <rtuli@redhat.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com> Co-authored-by: Rahul-Tuli <rtuli@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
Fynn Schmitt-Ulms committed
Apr 10, 2026 at 15:35 UTC
e7cfd7c5b9a18c4fb6eb7dd3108002793989e0eb
5 files changed
+43
-10
vllm/config/speculative.py
+1
@@ -825,6 +825,7 @@ class SpeculativeConfig:
825
"kimi_k2",
826
"kimi_k25",
827
"minimax_m2",
828
+ "gemma4",
829
]
830
if (
831
self.method in ("eagle3", "extract_hidden_states", "dflash")
vllm/model_executor/models/gemma4.py
+20
-5
@@ -60,7 +60,13 @@ from vllm.model_executor.model_loader.weight_utils import (
60
from vllm.sequence import IntermediateTensors
61
from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
62
63
-from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
63
+from .interfaces import (
64
+ EagleModelMixin,
65
+ MixtureOfExperts,
66
+ SupportsEagle3,
67
+ SupportsLoRA,
68
+ SupportsPP,
69
+)
70
from .utils import (
71
AutoWeightsLoader,
72
extract_layer_index,
@@ -838,7 +844,7 @@ class Gemma4CrossDecoderLayers(nn.Module):
844
@support_torch_compile(
845
enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
846
)
841
-class Gemma4Model(nn.Module):
847
+class Gemma4Model(nn.Module, EagleModelMixin):
848
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
849
super().__init__()
850
config = _get_text_config(vllm_config.model_config.hf_config)
@@ -1168,7 +1174,7 @@ class Gemma4Model(nn.Module):
1174
inputs_embeds: torch.Tensor | None = None,
1175
per_layer_inputs: torch.Tensor | None = None,
1176
**kwargs,
1171
- ) -> torch.Tensor | IntermediateTensors:
1177
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1178
if self.fast_prefill_enabled:
1179
hidden_states = self.fast_prefill_forward(
1180
input_ids,
@@ -1204,6 +1210,7 @@ class Gemma4Model(nn.Module):
1210
residual = intermediate_tensors["residual"]
1211
per_layer_inputs = intermediate_tensors.get("per_layer_inputs")
1212
1213
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
1214
for layer_idx, layer in enumerate(
1215
islice(self.layers, self.start_layer, self.end_layer)
1216
):
@@ -1222,6 +1229,9 @@ class Gemma4Model(nn.Module):
1229
per_layer_input=layer_per_input,
1230
**kwargs,
1231
)
1232
+ self._maybe_add_hidden_state(
1233
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
1234
+ )
1235
if not get_pp_group().is_last_rank:
1236
return IntermediateTensors(
1237
{
@@ -1236,6 +1246,9 @@ class Gemma4Model(nn.Module):
1246
hidden_states = self.norm(hidden_states)
1247
else:
1248
hidden_states, _ = self.norm(hidden_states, residual)
1249
+
1250
+ if len(aux_hidden_states) > 0:
1251
+ return hidden_states, aux_hidden_states
1252
return hidden_states
1253
1254
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1381,7 +1394,9 @@ class Gemma4Model(nn.Module):
1394
return loaded_params
1395
1396
1384
-class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
1397
+class Gemma4ForCausalLM(
1398
+ nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
1399
+):
1400
# Note: qkv_proj packing applies to non-k_eq_v layers (sliding
1401
# attention and full attention without k_eq_v). k_eq_v layers use
1402
# separate q_proj + k_proj without packing.
@@ -1463,7 +1478,7 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
1478
intermediate_tensors: IntermediateTensors | None = None,
1479
inputs_embeds: torch.Tensor | None = None,
1480
**kwargs,
1466
- ) -> torch.Tensor | IntermediateTensors:
1481
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1482
hidden_states = self.model(
1483
input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
1484
)
vllm/model_executor/models/gemma4_mm.py
+12
-2
@@ -64,7 +64,12 @@ from vllm.multimodal.processing.processor import (
64
from vllm.sequence import IntermediateTensors
65
from vllm.utils.tensor_schema import TensorSchema, TensorShape
66
67
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
67
+from .interfaces import (
68
+ MultiModalEmbeddings,
69
+ SupportsEagle3,
70
+ SupportsMultiModal,
71
+ SupportsPP,
72
+)
73
from .utils import (
74
AutoWeightsLoader,
75
WeightsMapper,
@@ -845,7 +850,12 @@ class Gemma4MultimodalEmbedder(nn.Module):
850
info=Gemma4ProcessingInfo,
851
dummy_inputs=Gemma4DummyInputsBuilder,
852
)
848
-class Gemma4ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
853
+class Gemma4ForConditionalGeneration(
854
+ nn.Module,
855
+ SupportsMultiModal,
856
+ SupportsPP,
857
+ SupportsEagle3,
858
+):
859
packed_modules_mapping = {
860
"qkv_proj": [
861
"q_proj",
vllm/v1/core/single_type_kv_cache_manager.py
+9
-3
@@ -565,11 +565,17 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
565
for computed in computed_blocks:
566
computed.pop()
567
if use_eagle and computed_blocks[0]:
568
- assert kv_cache_spec.block_size == alignment_tokens, (
569
- "aligned_length is not compatible with eagle now"
570
- )
568
for computed in computed_blocks:
569
computed.pop()
570
+ # Re-align after eagle pop: the pop may break the alignment
571
+ # when block_size != alignment_tokens (hybrid models with
572
+ # different page sizes, e.g. Gemma4).
573
+ while (
574
+ block_size != alignment_tokens
575
+ and len(computed_blocks[0]) * block_size % alignment_tokens != 0
576
+ ):
577
+ for computed in computed_blocks:
578
+ computed.pop()
579
return computed_blocks
580
581
def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
vllm/v1/spec_decode/eagle.py
+1
@@ -1329,6 +1329,7 @@ class SpecDecodeBaseProposer:
1329
"Qwen3_5MoeForConditionalGeneration",
1330
"Qwen3VLForConditionalGeneration",
1331
"Qwen3VLMoeForConditionalGeneration",
1332
+ "Gemma4ForConditionalGeneration",
1333
]:
1334
self.model.config.image_token_index = target_model.config.image_token_id
1335
elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":