| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MoE activation expressions for Humming input processing.""" |
| 4 | |
| 5 | from dataclasses import replace |
| 6 | |
| 7 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 8 | ApplyMoEActivationConfig, |
| 9 | MoEActivation, |
| 10 | ) |
| 11 | |
| 12 | _SILU_IMPL = "a / (1.f + expf(-a))" |
| 13 | _GELU_IMPL = "0.5f * a * (1.f + erff(a * 0.7071067811865476f))" |
| 14 | _RELU2_IMPL = "fmaxf(a, 0.f) * fmaxf(a, 0.f)" |
| 15 | _GELU_TANH_IMPL = """[](float a) { |
| 16 | const float cubic = 0.044715f * a * a * a; |
| 17 | const float inner = 0.7978845608028654f * (a + cubic); |
| 18 | return 0.5f * a * (1.f + tanhf(inner)); |
| 19 | }(a)""" |
| 20 | _CLAMPED_SILU_IMPL = """[](float a, float b) {{ |
| 21 | const float gate = fminf(a, {limit}); |
| 22 | const float up = fmaxf(fminf(b, {limit}), -{limit}); |
| 23 | const float activated_gate = gate / (1.f + expf(-{alpha} * gate)); |
| 24 | return activated_gate * (up + {beta}); |
| 25 | }}(a, b)""" |
| 26 | |
| 27 | # Parameterized expressions are resolved by get_humming_activation(). |
| 28 | _HUMMING_ACTIVATIONS: dict[MoEActivation, dict[str, str]] = { |
| 29 | MoEActivation.SILU: { |
| 30 | "activation_type": "binary_split", |
| 31 | "activation_impl": f"({_SILU_IMPL}) * b", |
| 32 | }, |
| 33 | MoEActivation.GELU: { |
| 34 | "activation_type": "binary_split", |
| 35 | "activation_impl": f"({_GELU_IMPL}) * b", |
| 36 | }, |
| 37 | MoEActivation.GELU_TANH: { |
| 38 | "activation_type": "binary_split", |
| 39 | "activation_impl": f"({_GELU_TANH_IMPL}) * b", |
| 40 | }, |
| 41 | MoEActivation.RELU2: { |
| 42 | "activation_type": "binary_split", |
| 43 | "activation_impl": f"({_RELU2_IMPL}) * b", |
| 44 | }, |
| 45 | MoEActivation.SWIGLUOAI: { |
| 46 | "activation_type": "binary_interleaved", |
| 47 | "activation_impl": _CLAMPED_SILU_IMPL, |
| 48 | }, |
| 49 | MoEActivation.SITU: { |
| 50 | "activation_type": "binary_split", |
| 51 | "activation_impl": """[](float a, float b) {{ |
| 52 | const float gate = {situ_beta} * tanhf(a / {situ_beta}); |
| 53 | const float activated_gate = gate / (1.f + expf(-a)); |
| 54 | const float up = {situ_up}; |
| 55 | return activated_gate * up; |
| 56 | }}(a, b)""", |
| 57 | }, |
| 58 | MoEActivation.SWIGLUOAI_UNINTERLEAVE: { |
| 59 | "activation_type": "binary_split", |
| 60 | "activation_impl": _CLAMPED_SILU_IMPL, |
| 61 | }, |
| 62 | MoEActivation.SWIGLUSTEP: { |
| 63 | "activation_type": "binary_split", |
| 64 | "activation_impl": """[](float a, float b) { |
| 65 | const float activated_gate = a / (1.f + expf(-a)); |
| 66 | const float gate = fminf(activated_gate, 7.f); |
| 67 | const float up = fmaxf(fminf(b, 7.f), -7.f); |
| 68 | return gate * up; |
| 69 | }(a, b)""", |
| 70 | }, |
| 71 | MoEActivation.SILU_NO_MUL: { |
| 72 | "activation_type": "unary", |
| 73 | "activation_impl": _SILU_IMPL, |
| 74 | }, |
| 75 | MoEActivation.GELU_NO_MUL: { |
| 76 | "activation_type": "unary", |
| 77 | "activation_impl": _GELU_IMPL, |
| 78 | }, |
| 79 | MoEActivation.GELU_TANH_NO_MUL: { |
| 80 | "activation_type": "unary", |
| 81 | "activation_impl": _GELU_TANH_IMPL, |
| 82 | }, |
| 83 | MoEActivation.RELU2_NO_MUL: { |
| 84 | "activation_type": "unary", |
| 85 | "activation_impl": _RELU2_IMPL, |
| 86 | }, |
| 87 | } |
| 88 | |
| 89 | |
| 90 | def get_humming_activation( |
| 91 | activation: MoEActivation, |
| 92 | config: ApplyMoEActivationConfig | None = None, |
| 93 | ) -> dict[str, str]: |
| 94 | config = config or ApplyMoEActivationConfig() |
| 95 | use_clamped_silu = False |
| 96 | if activation == MoEActivation.SWIGLUOAI: |
| 97 | config = ApplyMoEActivationConfig(clamp_limit=7.0, alpha=1.702, beta=1.0) |
| 98 | use_clamped_silu = True |
| 99 | elif activation == MoEActivation.SILU and config.clamp_limit is not None: |
| 100 | config = replace(config, alpha=1.0, beta=0.0) |
| 101 | use_clamped_silu = True |
| 102 | elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: |
| 103 | if config.clamp_limit is None: |
| 104 | raise ValueError("SWIGLUOAI_UNINTERLEAVE requires clamp_limit") |
| 105 | use_clamped_silu = True |
| 106 | |
| 107 | definition = _HUMMING_ACTIVATIONS[activation].copy() |
| 108 | if activation == MoEActivation.SITU: |
| 109 | if config.activation_situ_beta is None: |
| 110 | raise ValueError("SITU requires activation_situ_beta") |
| 111 | linear_beta = config.activation_situ_linear_beta |
| 112 | situ_up = "b" |
| 113 | if linear_beta is not None and linear_beta > 0: |
| 114 | situ_up = f"({float(linear_beta)}f) * tanhf(b / ({float(linear_beta)}f))" |
| 115 | situ_beta = f"({float(config.activation_situ_beta)}f)" |
| 116 | activation_impl = definition["activation_impl"] |
| 117 | activation_impl = activation_impl.format(situ_beta=situ_beta, situ_up=situ_up) |
| 118 | definition["activation_impl"] = activation_impl |
| 119 | elif use_clamped_silu: |
| 120 | assert config.clamp_limit is not None |
| 121 | limit = f"({float(config.clamp_limit)}f)" |
| 122 | alpha = f"({float(config.alpha)}f)" |
| 123 | beta = f"({float(config.beta)}f)" |
| 124 | activation_impl = _CLAMPED_SILU_IMPL.format(limit=limit, alpha=alpha, beta=beta) |
| 125 | definition["activation_impl"] = activation_impl |
| 126 | return definition |