| 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._aiter_ops import rocm_aiter_ops |
| 8 | from vllm.config import ( |
| 9 | CompilationConfig, |
| 10 | VllmConfig, |
| 11 | get_cached_compilation_config, |
| 12 | set_current_vllm_config, |
| 13 | ) |
| 14 | from vllm.model_executor.custom_op import CustomOp, op_registry |
| 15 | from vllm.model_executor.layers.activation import ( |
| 16 | GeluAndMul, |
| 17 | ReLUSquaredActivation, |
| 18 | SiluAndMul, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( |
| 21 | dispatch_topk_sigmoid_func, |
| 22 | dispatch_topk_softmax_func, |
| 23 | vllm_topk_sigmoid, |
| 24 | vllm_topk_softmax, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 27 | from vllm.platforms import current_platform |
| 28 | |
| 29 | RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] |
| 30 | |
| 31 | |
| 32 | # Registered subclass for test |
| 33 | @CustomOp.register("relu3") |
| 34 | class Relu3(ReLUSquaredActivation): |
| 35 | pass |
| 36 | |
| 37 | |
| 38 | @pytest.mark.parametrize( |
| 39 | "env, compilation_mode, backend, ops_enabled, default_on", |
| 40 | [ |
| 41 | # Default values based on compile level |
| 42 | # - All by default (no Inductor compilation) |
| 43 | (None, 0, "eager", [True] * 4, True), |
| 44 | (None, 1, "eager", [True] * 4, True), |
| 45 | (None, 2, "eager", [True] * 4, True), |
| 46 | (None, 3, "eager", [True] * 4, True), |
| 47 | # - None by default (with Inductor) |
| 48 | (None, 0, "inductor", [True] * 4, True), |
| 49 | # - None by default (with Inductor) |
| 50 | (None, 1, "inductor", [False] * 4, False), |
| 51 | (None, 2, "inductor", [False] * 4, False), |
| 52 | (None, 3, "inductor", [False] * 4, False), |
| 53 | # Explicitly enabling/disabling |
| 54 | # |
| 55 | # Default: all |
| 56 | # |
| 57 | # All but SiluAndMul |
| 58 | ("+rms_norm,-silu_and_mul", 0, "inductor", [1, 0, 1, 1], True), |
| 59 | # Only ReLU3 |
| 60 | ("none,-rms_norm,+relu3", 1, "eager", [0, 0, 0, 1], False), |
| 61 | # All but SiluAndMul |
| 62 | ("all,-silu_and_mul", 2, "inductor", [1, 0, 1, 1], True), |
| 63 | # All but ReLU3 (even if ReLU2 is on) |
| 64 | ("-relu3,+relu2", 3, "eager", [1, 1, 1, 0], True), |
| 65 | # RMSNorm and SiluAndMul |
| 66 | ("none,-relu3,+rms_norm,+silu_and_mul", 3, "eager", [1, 1, 0, 0], False), |
| 67 | # All but RMSNorm |
| 68 | ("-rms_norm", 3, "eager", [0, 1, 1, 1], True), |
| 69 | # |
| 70 | # Default: none |
| 71 | # |
| 72 | # Only ReLU3 |
| 73 | ("none,+relu3", 3, "inductor", [0, 0, 0, 1], False), |
| 74 | # All but RMSNorm |
| 75 | ("all,-rms_norm", 3, "inductor", [0, 1, 1, 1], True), |
| 76 | ], |
| 77 | ) |
| 78 | def test_enabled_ops( |
| 79 | env: str | None, |
| 80 | compilation_mode: int, |
| 81 | backend: str, |
| 82 | ops_enabled: list[int], |
| 83 | default_on: bool, |
| 84 | ): |
| 85 | custom_ops = env.split(",") if env else [] |
| 86 | vllm_config = VllmConfig( |
| 87 | compilation_config=CompilationConfig( |
| 88 | backend=backend, mode=compilation_mode, custom_ops=custom_ops |
| 89 | ) |
| 90 | ) |
| 91 | get_cached_compilation_config.cache_clear() |
| 92 | with set_current_vllm_config(vllm_config): |
| 93 | assert CustomOp.default_on() == default_on |
| 94 | |
| 95 | ops_enabled = [bool(x) for x in ops_enabled] |
| 96 | |
| 97 | assert RMSNorm(1024).enabled() == ops_enabled[0] |
| 98 | assert op_registry["rms_norm"].enabled() == ops_enabled[0] |
| 99 | |
| 100 | assert SiluAndMul().enabled() == ops_enabled[1] |
| 101 | assert op_registry["silu_and_mul"].enabled() == ops_enabled[1] |
| 102 | |
| 103 | assert GeluAndMul().enabled() == ops_enabled[2] |
| 104 | assert op_registry["gelu_and_mul"].enabled() == ops_enabled[2] |
| 105 | |
| 106 | # If registered, subclasses should follow their own name |
| 107 | assert Relu3().enabled() == ops_enabled[3] |
| 108 | assert op_registry["relu3"].enabled() == ops_enabled[3] |
| 109 | |
| 110 | # Unregistered subclass |
| 111 | class SiluAndMul2(SiluAndMul): |
| 112 | pass |
| 113 | |
| 114 | # Subclasses should not require registration |
| 115 | assert SiluAndMul2().enabled() == SiluAndMul().enabled() |
| 116 | |
| 117 | |
| 118 | @pytest.mark.parametrize( |
| 119 | "env", ["all,none", "all,+rms_norm,all", "+rms_norm,-rms_norm"] |
| 120 | ) |
| 121 | def test_enabled_ops_invalid(env: str): |
| 122 | with pytest.raises(Exception): # noqa |
| 123 | vllm_config = VllmConfig( |
| 124 | compilation_config=CompilationConfig(custom_ops=env.split(",")) |
| 125 | ) |
| 126 | with set_current_vllm_config(vllm_config): |
| 127 | RMSNorm(1024).enabled() |
| 128 | |
| 129 | |
| 130 | @pytest.mark.parametrize( |
| 131 | "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 132 | ) |
| 133 | def test_topk_softmax_dispatch(use_rocm_aiter: bool): |
| 134 | topk_func = dispatch_topk_softmax_func(use_rocm_aiter) |
| 135 | |
| 136 | if current_platform.is_rocm() and use_rocm_aiter: |
| 137 | assert topk_func == rocm_aiter_ops.topk_softmax |
| 138 | else: |
| 139 | assert topk_func == vllm_topk_softmax |
| 140 | |
| 141 | |
| 142 | @pytest.mark.parametrize( |
| 143 | "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 144 | ) |
| 145 | def test_topk_sigmoid_dispatch(use_rocm_aiter: bool): |
| 146 | topk_func = dispatch_topk_sigmoid_func(use_rocm_aiter) |
| 147 | |
| 148 | if current_platform.is_rocm() and use_rocm_aiter: |
| 149 | assert topk_func == rocm_aiter_ops.topk_sigmoid |
| 150 | else: |
| 151 | assert topk_func == vllm_topk_sigmoid |