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, 02:30 UTC
[Models][Gemma3/Gemma4] Support hidden_act variants in gated MLP (#40588)
Signed-off-by: Hiroaki Mikami <hiroaki8270+github@gmail.com>
Hiroaki Mikami committed
May 9, 2026 at 03:29 UTC
90f145aaf724194ccffeb3ea6a68e9457ff00169
4 files changed
+67
-16
tests/model_executor/test_gemma_hidden_act.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.layers.activation import (
8
+ GeluAndMul,
9
+ SiluAndMul,
10
+ get_act_and_mul_fn,
11
+ get_act_fn,
12
+)
13
+from vllm.model_executor.models.gemma3 import Gemma3MLP
14
+from vllm.model_executor.models.gemma4 import Gemma4MLP
15
+
16
+
17
+@pytest.mark.parametrize(
18
+ ("activation_name", "expected_type"),
19
+ [
20
+ ("gelu_pytorch_tanh", GeluAndMul),
21
+ ("silu", SiluAndMul),
22
+ ("swish", SiluAndMul),
23
+ ],
24
+)
25
+def test_get_act_and_mul_fn_supports_gemma_hidden_act_aliases(
26
+ activation_name: str,
27
+ expected_type: type[torch.nn.Module],
28
+ default_vllm_config,
29
+) -> None:
30
+ assert isinstance(get_act_and_mul_fn(activation_name), expected_type)
31
+
32
+
33
+def test_get_act_fn_supports_swish_alias() -> None:
34
+ assert isinstance(get_act_fn("swish"), torch.nn.SiLU)
35
+
36
+
37
+@pytest.mark.parametrize("mlp_cls", [Gemma3MLP, Gemma4MLP])
38
+@pytest.mark.parametrize(
39
+ ("activation_name", "expected_type"),
40
+ [
41
+ ("gelu_pytorch_tanh", GeluAndMul),
42
+ ("silu", SiluAndMul),
43
+ ("swish", SiluAndMul),
44
+ ],
45
+)
46
+def test_gemma_mlp_supports_hidden_act_variants(
47
+ mlp_cls: type[torch.nn.Module],
48
+ activation_name: str,
49
+ expected_type: type[torch.nn.Module],
50
+ default_vllm_config,
51
+ dist_init,
52
+) -> None:
53
+ mlp = mlp_cls(
54
+ hidden_size=16,
55
+ intermediate_size=32,
56
+ hidden_activation=activation_name,
57
+ )
58
+
59
+ assert isinstance(mlp.act_fn, expected_type)
60
+ assert mlp(torch.randn(3, 16)).shape == (3, 16)
vllm/model_executor/layers/activation.py
+3
@@ -712,6 +712,7 @@ _ACTIVATION_REGISTRY = LazyDict(
712
"relu": lambda: nn.ReLU(),
713
"relu2": lambda: ReLUSquaredActivation(),
714
"silu": lambda: nn.SiLU(),
715
+ "swish": lambda: nn.SiLU(),
716
"quick_gelu": lambda: QuickGELU(),
717
"tanh": lambda: nn.Tanh(),
718
"sigmoid": lambda: nn.Sigmoid(),
@@ -751,7 +752,9 @@ def get_act_fn(act_fn_name: str) -> nn.Module:
752
_ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict(
753
{
754
"gelu": lambda: GeluAndMul(),
755
+ "gelu_pytorch_tanh": lambda: GeluAndMul(approximate="tanh"),
756
"silu": lambda: SiluAndMul(),
757
+ "swish": lambda: SiluAndMul(),
758
"geglu": lambda: GeluAndMul(),
759
"swigluoai": lambda: SwigluOAIAndMul(),
760
}
vllm/model_executor/models/gemma3.py
+2
-8
@@ -26,7 +26,7 @@ from vllm.compilation.decorators import support_torch_compile
26
from vllm.config import CacheConfig, VllmConfig
27
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
28
from vllm.logger import init_logger
29
-from vllm.model_executor.layers.activation import GeluAndMul
29
+from vllm.model_executor.layers.activation import get_act_and_mul_fn
30
from vllm.model_executor.layers.attention import (
31
Attention,
32
EncoderOnlyAttention,
@@ -88,13 +88,7 @@ class Gemma3MLP(nn.Module):
88
quant_config=quant_config,
89
prefix=f"{prefix}.down_proj",
90
)
91
- if hidden_activation != "gelu_pytorch_tanh":
92
- raise ValueError(
93
- "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
94
- "function. Please set `hidden_act` and `hidden_activation` to "
95
- "`gelu_pytorch_tanh`."
96
- )
97
- self.act_fn = GeluAndMul(approximate="tanh")
91
+ self.act_fn = get_act_and_mul_fn(hidden_activation)
92
93
def forward(self, x: torch.Tensor) -> torch.Tensor:
94
gate_up, _ = self.gate_up_proj(x)
vllm/model_executor/models/gemma4.py
+2
-8
@@ -35,7 +35,7 @@ from vllm.distributed import (
35
)
36
from vllm.forward_context import get_forward_context
37
from vllm.logger import init_logger
38
-from vllm.model_executor.layers.activation import GeluAndMul
38
+from vllm.model_executor.layers.activation import get_act_and_mul_fn
39
from vllm.model_executor.layers.attention import Attention
40
from vllm.model_executor.layers.fused_moe import (
41
FusedMoE,
@@ -238,13 +238,7 @@ class Gemma4MLP(nn.Module):
238
quant_config=quant_config,
239
prefix=f"{prefix}.down_proj",
240
)
241
- if hidden_activation != "gelu_pytorch_tanh":
242
- raise ValueError(
243
- "Gemma4 uses `gelu_pytorch_tanh` as the hidden activation "
244
- "function. Please set `hidden_act` and `hidden_activation` to "
245
- "`gelu_pytorch_tanh`."
246
- )
247
- self.act_fn = GeluAndMul(approximate="tanh")
241
+ self.act_fn = get_act_and_mul_fn(hidden_activation)
242
243
def forward(self, x: torch.Tensor) -> torch.Tensor:
244
gate_up, _ = self.gate_up_proj(x)