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, 12:30 UTC
Fix PP in Gemma4 (#40786)
Signed-off-by: Rohit kumar Singh <rksingh@habana.ai> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Rohit Kumar Singh committed
Apr 29, 2026 at 15:47 UTC
5371d6fb4023a1a08021135e46e9354ba0923e50
1 file changed
+9
-16
vllm/model_executor/models/gemma4.py
+9
-16
@@ -1144,11 +1144,6 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1144
dtype=dtype,
1145
device=device,
1146
),
1147
- "residual": torch.zeros(
1148
- (batch_size, hidden_size),
1149
- dtype=dtype,
1150
- device=device,
1151
- ),
1147
}
1148
if ple_dim and ple_dim > 0:
1149
tensors["per_layer_inputs"] = torch.zeros(
@@ -1312,13 +1307,12 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1307
per_layer_inputs = self.project_per_layer_inputs(
1308
hidden_states, per_layer_embeds
1309
)
1315
- residual = None
1310
else:
1311
assert intermediate_tensors is not None
1312
hidden_states = intermediate_tensors["hidden_states"]
1319
- residual = intermediate_tensors["residual"]
1320
- per_layer_inputs = intermediate_tensors.get("per_layer_inputs")
1321
-
1313
+ if per_layer_inputs is not None:
1314
+ per_layer_inputs = intermediate_tensors["per_layer_inputs"]
1315
+ residual = None
1316
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
1317
for layer_idx, layer in enumerate(
1318
islice(self.layers, self.start_layer, self.end_layer)
@@ -1342,13 +1336,12 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1336
aux_hidden_states, layer_idx + 1, hidden_states, residual
1337
)
1338
if not get_pp_group().is_last_rank:
1345
- return IntermediateTensors(
1346
- {
1347
- "hidden_states": hidden_states,
1348
- "residual": residual,
1349
- "per_layer_inputs": per_layer_inputs,
1350
- }
1351
- )
1339
+ tensors: dict[str, torch.Tensor] = {
1340
+ "hidden_states": hidden_states,
1341
+ }
1342
+ if per_layer_inputs is not None:
1343
+ tensors["per_layer_inputs"] = per_layer_inputs
1344
+ return IntermediateTensors(tensors)
1345
# Gemma4 incorporates residual into hidden_states directly
1346
# Apply norm without residual fusion when possible.
1347
if residual is None: