Skip to content

vllm.model_executor.layers.fused_moe.activation

MoE activation function enum and utilities.

Classes:

Functions:

ApplyMoEActivationConfig dataclass

Configuration forwarded to apply_moe_activation.

Methods:

  • from_configs –

    Build from the model and quantization configurations.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@dataclass(frozen=True)
class ApplyMoEActivationConfig:
    """Configuration forwarded to ``apply_moe_activation``."""

    clamp_limit: float | None = None
    alpha: float = 1.0
    beta: float = 0.0
    activation_situ_beta: float | None = None
    activation_situ_linear_beta: float | None = None

    @classmethod
    def from_configs(
        cls,
        moe_config: "FusedMoEConfig",
        quant_config: "FusedMoEQuantConfig",
    ) -> "ApplyMoEActivationConfig":
        """Build from the model and quantization configurations."""
        clamp_limit = quant_config.gemm1_clamp_limit
        if clamp_limit is None:
            clamp_limit = moe_config.swiglu_limit
        alpha = quant_config.gemm1_alpha
        if alpha is None:
            alpha = moe_config.swiglu_alpha
        beta = quant_config.gemm1_beta
        if beta is None:
            beta = moe_config.swiglu_beta
        return cls(
            clamp_limit=clamp_limit,
            alpha=1.0 if alpha is None else alpha,
            beta=0.0 if beta is None else beta,
            activation_situ_beta=moe_config.activation_situ_beta,
            activation_situ_linear_beta=moe_config.activation_situ_linear_beta,
        )

from_configs(moe_config, quant_config) classmethod

Build from the model and quantization configurations.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@classmethod
def from_configs(
    cls,
    moe_config: "FusedMoEConfig",
    quant_config: "FusedMoEQuantConfig",
) -> "ApplyMoEActivationConfig":
    """Build from the model and quantization configurations."""
    clamp_limit = quant_config.gemm1_clamp_limit
    if clamp_limit is None:
        clamp_limit = moe_config.swiglu_limit
    alpha = quant_config.gemm1_alpha
    if alpha is None:
        alpha = moe_config.swiglu_alpha
    beta = quant_config.gemm1_beta
    if beta is None:
        beta = moe_config.swiglu_beta
    return cls(
        clamp_limit=clamp_limit,
        alpha=1.0 if alpha is None else alpha,
        beta=0.0 if beta is None else beta,
        activation_situ_beta=moe_config.activation_situ_beta,
        activation_situ_linear_beta=moe_config.activation_situ_linear_beta,
    )

MoEActivation

Bases: Enum

Activation functions for MoE layers.

Methods:

  • from_str –

    Parse from string for backward compatibility.

  • without_mul –

    Get the non-gated variant of this activation.

Attributes:

  • custom_op_name (str) –

    Maps to the CustomOp name of activations

  • is_gated (bool) –

    Returns True if activation expects gate*activation(up) pattern.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class MoEActivation(Enum):
    """Activation functions for MoE layers."""

    # Gated activations (gate * activation(up)) expect input of shape [..., 2*d]
    # and produce output of shape [..., d]
    SILU = "silu"
    GELU = "gelu"
    GELU_TANH = "gelu_tanh"
    RELU2 = "relu2"
    # SWIGLUOAI expects gate/up *interleaved* in w13 ([gate0, up0, gate1, ...]),
    # as in gpt-oss checkpoints. SWIGLUOAI_UNINTERLEAVE has identical math but
    # expects the *packed* layout ([all gates; all ups]), as produced by a
    # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
    SWIGLUOAI = "swigluoai"
    SITU = "situ"
    SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
    SWIGLUSTEP = "swiglustep"

    # Non-gated activations (no mul with gate) expect input of shape [..., d]
    # and produce output of shape [..., d].
    # NOTE: Non-gated activations require the "_no_mul" suffix to be present.
    SILU_NO_MUL = "silu_no_mul"
    GELU_NO_MUL = "gelu_no_mul"
    GELU_TANH_NO_MUL = "gelu_tanh_no_mul"
    RELU2_NO_MUL = "relu2_no_mul"

    @property
    def is_gated(self) -> bool:
        """Returns True if activation expects gate*activation(up) pattern.

        Gated activations expect input tensor with 2x the output size,
        where the first half is the gate and second half is the up projection.
        """
        return not self.value.endswith("_no_mul")

    @property
    def custom_op_name(self) -> str:
        """Maps to the CustomOp name of activations
        in vllm/model_executor/layers/activation.py."""
        return _CUSTOM_OP_NAMES[self]

    def without_mul(self) -> "MoEActivation":
        """Get the non-gated variant of this activation.

        For activations that have a _no_mul variant, returns that variant.
        For activations without a _no_mul variant (or already _no_mul),
        returns self.
        """
        return _WITHOUT_MUL.get(self, self)

    @classmethod
    def from_str(cls, s: str) -> "MoEActivation":
        """Parse from string for backward compatibility."""
        s = _STR_ALIASES.get(s, s)
        for member in cls:
            if member.value == s:
                return member
        valid = [m.value for m in cls]
        raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}")

custom_op_name property

Maps to the CustomOp name of activations in vllm/model_executor/layers/activation.py.

is_gated property

Returns True if activation expects gate*activation(up) pattern.

Gated activations expect input tensor with 2x the output size, where the first half is the gate and second half is the up projection.

from_str(s) classmethod

Parse from string for backward compatibility.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
69
70
71
72
73
74
75
76
77
@classmethod
def from_str(cls, s: str) -> "MoEActivation":
    """Parse from string for backward compatibility."""
    s = _STR_ALIASES.get(s, s)
    for member in cls:
        if member.value == s:
            return member
    valid = [m.value for m in cls]
    raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}")

without_mul()

Get the non-gated variant of this activation.

For activations that have a _no_mul variant, returns that variant. For activations without a _no_mul variant (or already _no_mul), returns self.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
60
61
62
63
64
65
66
67
def without_mul(self) -> "MoEActivation":
    """Get the non-gated variant of this activation.

    For activations that have a _no_mul variant, returns that variant.
    For activations without a _no_mul variant (or already _no_mul),
    returns self.
    """
    return _WITHOUT_MUL.get(self, self)

activation_without_mul(activation)

Get the non-gated variant of an activation function.

Parameters:

  • activation

    (str) –

    The activation function name (e.g., "silu", "gelu")

Returns:

  • str –

    The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul")

Source code in vllm/model_executor/layers/fused_moe/activation.py 
108
109
110
111
112
113
114
115
116
117
def activation_without_mul(activation: str) -> str:
    """Get the non-gated variant of an activation function.

    Args:
        activation: The activation function name (e.g., "silu", "gelu")

    Returns:
        The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul")
    """
    return MoEActivation.from_str(activation).without_mul().value

apply_moe_activation(activation, output, input, *, activation_config=None, topk_ids=None, expert_map=None)

Apply MoE activation function.

The configuration drives specialized activation behavior. Routing tensors remain per-call inputs because they depend on the current token assignment.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def apply_moe_activation(
    activation: MoEActivation,
    output: torch.Tensor,
    input: torch.Tensor,
    *,
    activation_config: ApplyMoEActivationConfig | None = None,
    topk_ids: torch.Tensor | None = None,
    expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
    """Apply MoE activation function.

    The configuration drives specialized activation behavior. Routing tensors
    remain per-call inputs because they depend on the current token assignment.
    """
    config = (
        _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG
        if activation_config is None
        else activation_config
    )

    assert input.dim() == 2, "Input must be 2D"
    assert output.dim() == 2, "Output must be 2D"
    if activation.is_gated:
        assert output.size(-1) * 2 == input.size(-1), (
            f"{activation.value} expects 2x ratio: "
            f"{output.size(-1) * 2} vs {input.size(-1)}"
        )
    else:
        assert output.size(-1) == input.size(-1), (
            f"{activation.value} expects equal sizes: "
            f"{output.size(-1)} vs {input.size(-1)}"
        )

    # Activations with gated multiplication (gate × activation(up))
    if activation == MoEActivation.SILU:
        if config.clamp_limit is not None:
            silu_and_mul_with_clamp(
                output, input, config.clamp_limit, topk_ids, expert_map
            )
        else:
            torch.ops._C.silu_and_mul(output, input)
    elif activation == MoEActivation.GELU:
        torch.ops._C.gelu_and_mul(output, input)
    elif activation == MoEActivation.GELU_TANH:
        torch.ops._C.gelu_tanh_and_mul(output, input)
    elif activation == MoEActivation.SITU:
        # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries.
        # (The pure-torch fallback below upcast both halves to fp32 and
        # allocated ~8 temporaries per call, blowing up MoE memory.)
        # Both betas come from FusedMoEConfig; a missing beta means the caller
        # bypassed the config plumbing, so fail rather than silently use 1.0.
        # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel
        # (up passed through), matching SituAndMul(linear_beta=None).
        assert config.activation_situ_beta is not None, (
            "SITU requires activation_situ_beta from FusedMoEConfig"
        )
        torch.ops._C.situ_and_mul(
            output,
            input,
            config.activation_situ_beta,
            -1.0
            if config.activation_situ_linear_beta is None
            else config.activation_situ_linear_beta,
        )
    elif activation == MoEActivation.SWIGLUOAI:
        torch.ops._C.swigluoai_and_mul(output, input)
    elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
        # SwiGLU-OAI on packed w13 (gate = first half, up = second half).
        assert config.clamp_limit is not None, (
            "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
        )
        torch.ops._C.silu_and_mul_with_clamp(
            output, input, config.clamp_limit, config.alpha, config.beta
        )
    elif activation == MoEActivation.SWIGLUSTEP:
        from vllm.model_executor.layers.activation import swiglustep_and_mul_triton

        swiglustep_and_mul_triton(output, input)

    # Activations without gated multiplication
    elif activation == MoEActivation.SILU_NO_MUL:
        output.copy_(F.silu(input))
    elif activation == MoEActivation.GELU_NO_MUL:
        output.copy_(F.gelu(input))
    elif activation == MoEActivation.GELU_TANH_NO_MUL:
        output.copy_(F.gelu(input, approximate="tanh"))
    elif activation == MoEActivation.RELU2_NO_MUL:
        F.relu(input, inplace=True)
        torch.square(input, out=output)
    else:
        raise ValueError(f"Unsupported FusedMoe activation: {activation}")

    return output

apply_moe_activation_supported(activation)

Whether apply_moe_activation supports an activation.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
137
138
139
def apply_moe_activation_supported(activation: MoEActivation) -> bool:
    """Whether ``apply_moe_activation`` supports an activation."""
    return activation in _APPLY_MOE_ACTIVATIONS