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, 07:30 UTC
[Model] Fix Gemma4 MoE activation mismatch (#41574)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Luciano Martins committed
May 5, 2026 at 01:34 UTC
6bb924bbf3a919dd6fe8a85033bd65f378b16421
6 files changed
+23
-1
vllm/model_executor/layers/fused_moe/activation.py
+14
@@ -15,6 +15,7 @@ class MoEActivation(Enum):
15
# and produce output of shape [..., d]
16
SILU = "silu"
17
GELU = "gelu"
18
+ GELU_TANH = "gelu_tanh"
19
RELU2 = "relu2"
20
SWIGLUOAI = "swigluoai"
21
SWIGLUSTEP = "swiglustep"
@@ -24,6 +25,7 @@ class MoEActivation(Enum):
25
# NOTE: Non-gated activations require the "_no_mul" suffix to be present.
26
SILU_NO_MUL = "silu_no_mul"
27
GELU_NO_MUL = "gelu_no_mul"
28
+ GELU_TANH_NO_MUL = "gelu_tanh_no_mul"
29
RELU2_NO_MUL = "relu2_no_mul"
30
31
@property
@@ -53,6 +55,7 @@ class MoEActivation(Enum):
55
@classmethod
56
def from_str(cls, s: str) -> "MoEActivation":
57
"""Parse from string for backward compatibility."""
58
+ s = _STR_ALIASES.get(s, s)
59
for member in cls:
60
if member.value == s:
61
return member
@@ -61,20 +64,27 @@ class MoEActivation(Enum):
64
65
66
# Module-level lookup tables used by MoEActivation functions.
67
+_STR_ALIASES: dict[str, str] = {
68
+ "gelu_pytorch_tanh": "gelu_tanh",
69
+}
70
+
71
_CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
72
MoEActivation.SILU: "silu_and_mul",
73
MoEActivation.GELU: "gelu_and_mul",
74
+ MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
75
MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
76
MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
77
MoEActivation.RELU2: "relu2",
78
MoEActivation.SILU_NO_MUL: "silu_and_mul",
79
MoEActivation.GELU_NO_MUL: "gelu_and_mul",
80
+ MoEActivation.GELU_TANH_NO_MUL: "gelu_tanh_and_mul",
81
MoEActivation.RELU2_NO_MUL: "relu2",
82
}
83
84
_WITHOUT_MUL: dict[MoEActivation, MoEActivation] = {
85
MoEActivation.SILU: MoEActivation.SILU_NO_MUL,
86
MoEActivation.GELU: MoEActivation.GELU_NO_MUL,
87
+ MoEActivation.GELU_TANH: MoEActivation.GELU_TANH_NO_MUL,
88
MoEActivation.RELU2: MoEActivation.RELU2_NO_MUL,
89
}
90
@@ -115,6 +125,8 @@ def apply_moe_activation(
125
torch.ops._C.silu_and_mul(output, input)
126
elif activation == MoEActivation.GELU:
127
torch.ops._C.gelu_and_mul(output, input)
128
+ elif activation == MoEActivation.GELU_TANH:
129
+ torch.ops._C.gelu_tanh_and_mul(output, input)
130
elif activation == MoEActivation.SWIGLUOAI:
131
torch.ops._C.swigluoai_and_mul(output, input)
132
elif activation == MoEActivation.SWIGLUSTEP:
@@ -127,6 +139,8 @@ def apply_moe_activation(
139
output.copy_(F.silu(input))
140
elif activation == MoEActivation.GELU_NO_MUL:
141
output.copy_(F.gelu(input))
142
+ elif activation == MoEActivation.GELU_TANH_NO_MUL:
143
+ output.copy_(F.gelu(input, approximate="tanh"))
144
elif activation == MoEActivation.RELU2_NO_MUL:
145
F.relu(input, inplace=True)
146
torch.square(input, out=output)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+2
@@ -786,9 +786,11 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
786
return activation in [
787
MoEActivation.SILU,
788
MoEActivation.GELU,
789
+ MoEActivation.GELU_TANH,
790
MoEActivation.SWIGLUOAI,
791
MoEActivation.SILU_NO_MUL,
792
MoEActivation.GELU_NO_MUL,
793
+ MoEActivation.GELU_TANH_NO_MUL,
794
MoEActivation.RELU2_NO_MUL,
795
]
796
vllm/model_executor/layers/fused_moe/fused_humming_moe.py
+2
@@ -152,10 +152,12 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
152
return activation in [
153
MoEActivation.SILU,
154
MoEActivation.GELU,
155
+ MoEActivation.GELU_TANH,
156
MoEActivation.SWIGLUOAI,
157
MoEActivation.SWIGLUSTEP,
158
MoEActivation.SILU_NO_MUL,
159
MoEActivation.GELU_NO_MUL,
160
+ MoEActivation.GELU_TANH_NO_MUL,
161
MoEActivation.RELU2_NO_MUL,
162
]
163
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+2
@@ -613,10 +613,12 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
613
return activation in [
614
MoEActivation.SILU,
615
MoEActivation.GELU,
616
+ MoEActivation.GELU_TANH,
617
MoEActivation.SWIGLUOAI,
618
MoEActivation.SWIGLUSTEP,
619
MoEActivation.SILU_NO_MUL,
620
MoEActivation.GELU_NO_MUL,
621
+ MoEActivation.GELU_TANH_NO_MUL,
622
MoEActivation.RELU2_NO_MUL,
623
]
624
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
@@ -1941,10 +1941,12 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
1941
return activation in [
1942
MoEActivation.SILU,
1943
MoEActivation.GELU,
1944
+ MoEActivation.GELU_TANH,
1945
MoEActivation.SWIGLUOAI,
1946
MoEActivation.SWIGLUSTEP,
1947
MoEActivation.SILU_NO_MUL,
1948
MoEActivation.GELU_NO_MUL,
1949
+ MoEActivation.GELU_TANH_NO_MUL,
1950
MoEActivation.RELU2_NO_MUL,
1951
]
1952
vllm/model_executor/models/gemma4.py
+1
-1
@@ -360,7 +360,7 @@ class Gemma4MoE(nn.Module):
360
quant_config=quant_config,
361
prefix=f"{prefix}.experts",
362
custom_routing_function=routing_function,
363
- activation="gelu",
363
+ activation="gelu_tanh",
364
)
365
366
def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor: