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 3, 19:00 UTC
[Bugfix][Model] Gemma4 MoE routing closure captures per_expert_scale, breaking functional_call substitution (#42250)
Signed-off-by: Noelia <noeliabentancor1@gmail.com> Signed-off-by: Noelia Bentancor <71080743+NoeliaBentancor@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Noelia Bentancor committed
May 13, 2026 at 14:43 UTC
5794c65f8c36d9f7f486277cc5ed40e65b6e6100
1 file changed
+7
-4
vllm/model_executor/models/gemma4.py
+7
-4
@@ -326,8 +326,9 @@ class Gemma4MoE(nn.Module):
326
# Gemma4 routing: softmax over ALL experts → top-k → renormalize.
327
# FusedMoE's built-in fused_topk scopes softmax differently, so
328
# a custom routing function is needed for numerical correctness.
329
- per_expert_scale = self.per_expert_scale
330
-
329
+ # NOTE: self.per_expert_scale is read at call time (not captured into
330
+ # a local) so that torch.func.functional_call parameter substitution
331
+ # reaches the routing function correctly.
332
def routing_function(
333
hidden_states: torch.Tensor,
334
gating_output: torch.Tensor,
@@ -336,10 +337,12 @@ class Gemma4MoE(nn.Module):
337
) -> tuple[torch.Tensor, torch.Tensor]:
338
if current_platform.is_cuda_alike() or current_platform.is_xpu():
339
return gemma4_fused_routing_kernel_triton(
339
- gating_output, topk, per_expert_scale
340
+ gating_output, topk, self.per_expert_scale
341
)
342
342
- return gemma4_routing_function_torch(gating_output, topk, per_expert_scale)
343
+ return gemma4_routing_function_torch(
344
+ gating_output, topk, self.per_expert_scale
345
+ )
346
347
# FusedMoE experts with custom Gemma4 routing
348
self.experts = FusedMoE(