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 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
@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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
@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 
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
78
79
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 
71
72
73
74
75
76
77
78
79
@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 
62
63
64
65
66
67
68
69
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 
110
111
112
113
114
115
116
117
118
119
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, valid_token_counts=None)

Apply MoE activation function.

The configuration drives specialized activation behavior. Routing tensors and valid token counts remain per-call inputs because they depend on the current token assignment. A single token count masks a flat [T, D] buffer; one count per expert masks each prefix in a padded [E, T, D] buffer.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
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,
    valid_token_counts: torch.Tensor | None = None,
) -> torch.Tensor:
    """Apply MoE activation function.

    The configuration drives specialized activation behavior. Routing tensors
    and valid token counts remain per-call inputs because they depend on the
    current token assignment. A single token count masks a flat ``[T, D]``
    buffer; one count per expert masks each prefix in a padded ``[E, T, D]``
    buffer.
    """
    config = (
        _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG
        if activation_config is None
        else activation_config
    )

    if valid_token_counts is not None:
        return _apply_moe_activation_masked(
            activation,
            output,
            input,
            valid_token_counts,
            config,
        )

    _validate_moe_activation_shapes(activation, output, input, expected_dim=2)

    # 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_masked_supported(activation)

Whether the masked apply_moe_activation path supports an activation.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
159
160
161
def apply_moe_activation_masked_supported(activation: MoEActivation) -> bool:
    """Whether the masked ``apply_moe_activation`` path supports an activation."""
    return activation in _MASKED_MOE_ACTIVATION_NAMES

apply_moe_activation_supported(activation)

Whether apply_moe_activation supports an activation.

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

situ_and_mul_quant(output, scale, input, *, beta, linear_beta, group_size=0, num_valid_tokens=None, topk=1)

Fused Kimi SITU activation + dynamic FP8 quantization.

Writes the quantized fp8 down-projection input into output [M, d] and its float32 scale into scale (dequant = q * scale), fusing the SITU activation with the FP8 quant that Humming's w2 GEMM would otherwise perform in a separate pass -- this avoids materializing (and rounding through) the bf16 activation_output buffer.

group_size selects the quant granularity: 0 gives per-token scales scale [M, 1]; 128 gives k-major block-FP8 group scales scale [M, d // 128], matching humming quant_input(group_size=128, float32).

linear_beta <= 0 (or None) means "unset" (up passed through), matching SituAndMul(linear_beta=None). num_valid_tokens (int32 scalar tensor) is the DeepEP v2 contiguous-layout token count (e.g. psum[-1:]); the kernel multiplies it by topk to bound rows. Padding rows are skipped and receive a benign scale of 1.0. Kept in this module so that every torch.ops._C.situ* call lives in one file.

Source code in vllm/model_executor/layers/fused_moe/activation.py 
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
def situ_and_mul_quant(
    output: torch.Tensor,
    scale: torch.Tensor,
    input: torch.Tensor,
    *,
    beta: float,
    linear_beta: float | None,
    group_size: int = 0,
    num_valid_tokens: torch.Tensor | None = None,
    topk: int = 1,
) -> None:
    """Fused Kimi SITU activation + dynamic FP8 quantization.

    Writes the quantized fp8 down-projection input into ``output`` [M, d] and its
    float32 scale into ``scale`` (dequant = q * scale), fusing the SITU
    activation with the FP8 quant that Humming's w2 GEMM would otherwise perform
    in a separate pass -- this avoids materializing (and rounding through) the
    bf16 activation_output buffer.

    ``group_size`` selects the quant granularity: ``0`` gives per-token scales
    ``scale`` [M, 1]; ``128`` gives k-major block-FP8 group scales ``scale``
    [M, d // 128], matching humming ``quant_input(group_size=128, float32)``.

    ``linear_beta`` <= 0 (or None) means "unset" (up passed through), matching
    ``SituAndMul(linear_beta=None)``. ``num_valid_tokens`` (int32 scalar tensor)
    is the DeepEP v2 contiguous-layout token count (e.g. ``psum[-1:]``); the
    kernel multiplies it by ``topk`` to bound rows. Padding rows are skipped and
    receive a benign scale of 1.0. Kept in this module so that every
    ``torch.ops._C.situ*`` call lives in one file.
    """
    torch.ops._C.situ_and_mul_quant(
        output,
        scale,
        input,
        beta,
        -1.0 if linear_beta is None else linear_beta,
        group_size,
        num_valid_tokens,
        topk,
    )