| 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) |