@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8f158d0ee2

[MoE] Share apply_moe_activation support metadata (#44359)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Michael Goin committed Aug 5, 2026 at 15:05 UTC 8f158d0ee2475c64f66ee9ebd55a3606209db92f
13 files changed +425 -241
tests/kernels/moe/test_cutlass_moe.py
+78 -6
@@ -11,8 +11,15 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 from tests.kernels.moe.utils import make_dummy_moe_config
12 from vllm import _custom_ops as ops
13 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
14 -from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
15 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
14 +from vllm.model_executor.layers.fused_moe import (
15 + ApplyMoEActivationConfig,
16 + fused_experts,
17 + fused_topk,
18 +)
19 +from vllm.model_executor.layers.fused_moe.activation import (
20 + MoEActivation,
21 + apply_moe_activation_supported,
22 +)
23 from vllm.model_executor.layers.fused_moe.all2all_utils import (
24 maybe_make_prepare_finalize,
25 )
@@ -24,8 +31,11 @@ from vllm.model_executor.layers.fused_moe.config import (
31 from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
32 CutlassExpertsFp4,
33 CutlassExpertsFp8,
34 + CutlassExpertsMxfp4,
35 + CutlassExpertsW4A8Fp8,
36 run_cutlass_moe_fp8,
37 )
38 +from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
39 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
40 from vllm.platforms import current_platform
41 from vllm.utils.torch_utils import set_random_seed
@@ -53,10 +63,72 @@ MNK_FACTORS = [
63 vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
64
65
56 -def test_cutlass_moe_supports_gelu_tanh_activation_metadata():
57 - assert CutlassExpertsFp8._supports_activation(MoEActivation.GELU_TANH)
58 - assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH)
59 - assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH_NO_MUL)
66 +@pytest.mark.parametrize(
67 + "experts_cls",
68 + [
69 + CutlassExpertsFp8,
70 + CutlassExpertsFp4,
71 + CutlassExpertsMxfp4,
72 + CutlassExpertsW4A8Fp8,
73 + ],
74 +)
75 +@pytest.mark.parametrize("activation", list(MoEActivation))
76 +def test_cutlass_moe_activation_metadata_tracks_shared_apply(
77 + experts_cls: type[mk.FusedMoEExperts], activation: MoEActivation
78 +):
79 + supports_shape = experts_cls._supports_no_act_and_mul() or activation.is_gated
80 + expected = supports_shape and apply_moe_activation_supported(activation)
81 +
82 + assert experts_cls._supports_activation(activation) == expected
83 +
84 +
85 +def test_cutlass_moe_forwards_shared_activation_parameters():
86 + moe_config = make_dummy_moe_config()
87 + moe_config.swiglu_limit = 7.0
88 + moe_config.swiglu_alpha = 1.5
89 + moe_config.swiglu_beta = 0.25
90 + moe_config.activation_situ_beta = 2.0
91 + moe_config.activation_situ_linear_beta = 3.0
92 + quant_config = FusedMoEQuantConfig.make(gemm1_alpha=1.75)
93 +
94 + assert ApplyMoEActivationConfig.from_configs(
95 + moe_config, quant_config
96 + ) == ApplyMoEActivationConfig(
97 + clamp_limit=7.0,
98 + alpha=1.75,
99 + beta=0.25,
100 + activation_situ_beta=2.0,
101 + activation_situ_linear_beta=3.0,
102 + )
103 +
104 +
105 +@pytest.mark.parametrize(
106 + ("backend", "expected"),
107 + [
108 + ("cutlass", nvfp4_oracle.NvFp4MoeBackend.VLLM_CUTLASS),
109 + ("humming", nvfp4_oracle.NvFp4MoeBackend.HUMMING),
110 + ],
111 +)
112 +def test_nvfp4_clamp_allows_shared_activation_backends(
113 + monkeypatch, backend: str, expected: nvfp4_oracle.NvFp4MoeBackend
114 +):
115 + class SupportedExperts:
116 + @staticmethod
117 + def is_supported_config(*args, **kwargs):
118 + return True, None
119 +
120 + monkeypatch.setattr(
121 + nvfp4_oracle, "backend_to_kernel_cls", lambda backend: [SupportedExperts]
122 + )
123 + moe_config = make_dummy_moe_config()
124 + moe_config.moe_backend = backend
125 + moe_config.swiglu_limit = 7.0
126 +
127 + selected, _ = nvfp4_oracle.select_nvfp4_moe_backend(
128 + moe_config, weight_key=None, activation_key=None
129 + )
130 +
131 + assert selected == expected
132
133
134 @dataclasses.dataclass
tests/kernels/moe/test_moe.py
+68
@@ -27,6 +27,11 @@ from vllm.model_executor.layers.fused_moe import (
27 MoEActivation,
28 fused_topk,
29 )
30 +from vllm.model_executor.layers.fused_moe.activation import (
31 + ApplyMoEActivationConfig,
32 + apply_moe_activation,
33 + apply_moe_activation_supported,
34 +)
35 from vllm.model_executor.layers.fused_moe.config import (
36 FUSED_MOE_UNQUANTIZED_CONFIG,
37 int4_w4a16_moe_quant_config,
@@ -996,6 +1001,23 @@ def test_fused_marlin_moe(
1001 per_act_token_quant=True,
1002 )
1003
1004 + def instance_activation(
1005 + activation: MoEActivation,
1006 + output: torch.Tensor,
1007 + input: torch.Tensor,
1008 + *,
1009 + topk_ids: torch.Tensor | None = None,
1010 + expert_map: torch.Tensor | None = None,
1011 + ) -> None:
1012 + apply_moe_activation(
1013 + activation,
1014 + output,
1015 + input,
1016 + activation_config=ApplyMoEActivationConfig(),
1017 + topk_ids=topk_ids,
1018 + expert_map=expert_map,
1019 + )
1020 +
1021 marlin_output = fused_marlin_moe(
1022 a,
1023 w1_data.qweight,
@@ -1021,6 +1043,7 @@ def test_fused_marlin_moe(
1043 input_dtype=a_dtype,
1044 quant_type_id=b_type.id,
1045 is_k_full=is_k_full,
1046 + activation_func=instance_activation,
1047 )
1048
1049 torch.testing.assert_close(marlin_output, torch_output, atol=4e-2, rtol=0)
@@ -1249,6 +1272,51 @@ def _make_humming_indexed_experts(activation: MoEActivation):
1272 return experts
1273
1274
1275 +@pytest.mark.parametrize("activation", list(MoEActivation))
1276 +def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation):
1277 + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1278 + HummingExpertsBase,
1279 + )
1280 +
1281 + assert HummingExpertsBase._supports_activation(
1282 + activation
1283 + ) == apply_moe_activation_supported(activation)
1284 +
1285 +
1286 +def test_humming_delegates_to_instance_activation():
1287 + from types import SimpleNamespace
1288 + from unittest.mock import Mock
1289 +
1290 + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1291 + HummingExpertsBase,
1292 + )
1293 +
1294 + activation_func = Mock()
1295 + activation_config = ApplyMoEActivationConfig(
1296 + clamp_limit=7.0,
1297 + alpha=1.5,
1298 + beta=0.25,
1299 + activation_situ_beta=2.0,
1300 + activation_situ_linear_beta=3.0,
1301 + )
1302 + experts = SimpleNamespace(
1303 + activation=activation_func,
1304 + activation_config=activation_config,
1305 + )
1306 + input = torch.empty(1, 2)
1307 + output = torch.empty(1, 1)
1308 +
1309 + HummingExpertsBase.apply_activation(
1310 + experts, MoEActivation.SWIGLUOAI_UNINTERLEAVE, output, input
1311 + )
1312 +
1313 + activation_func.assert_called_once_with(
1314 + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
1315 + input=input,
1316 + output=output,
1317 + )
1318 +
1319 +
1320 @pytest.mark.parametrize(
1321 "activation",
1322 [
tests/kernels/moe/test_triton_moe_no_act_mul.py
+49 -1
@@ -11,7 +11,12 @@ import pytest
11 import torch
12
13 from tests.kernels.moe.utils import make_dummy_moe_config
14 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
14 +from vllm.model_executor.layers.fused_moe.activation import (
15 + ApplyMoEActivationConfig,
16 + MoEActivation,
17 + apply_moe_activation,
18 + apply_moe_activation_supported,
19 +)
20 from vllm.model_executor.layers.fused_moe.config import (
21 FUSED_MOE_UNQUANTIZED_CONFIG,
22 )
@@ -31,6 +36,49 @@ NO_MUL_ACTIVATIONS = [
36 ]
37
38
39 +def test_apply_moe_activation_supported_contract():
40 + supported = {
41 + activation
42 + for activation in MoEActivation
43 + if apply_moe_activation_supported(activation)
44 + }
45 +
46 + assert supported == set(MoEActivation) - {MoEActivation.RELU2}
47 +
48 +
49 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
50 +@pytest.mark.parametrize(
51 + "activation",
52 + [
53 + activation
54 + for activation in MoEActivation
55 + if apply_moe_activation_supported(activation)
56 + ],
57 +)
58 +@torch.inference_mode()
59 +def test_supported_apply_moe_activation_executes(activation: MoEActivation):
60 + output_width = 64
61 + input_width = output_width * 2 if activation.is_gated else output_width
62 + input = torch.randn(2, input_width, device="cuda", dtype=torch.bfloat16)
63 + output = torch.empty(2, output_width, device="cuda", dtype=torch.bfloat16)
64 + activation_config = ApplyMoEActivationConfig()
65 + if activation == MoEActivation.SITU:
66 + activation_config = ApplyMoEActivationConfig(activation_situ_beta=1.0)
67 + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
68 + activation_config = ApplyMoEActivationConfig(clamp_limit=7.0)
69 +
70 + apply_moe_activation(activation, output, input, activation_config=activation_config)
71 +
72 + assert torch.isfinite(output).all()
73 +
74 +
75 +@pytest.mark.parametrize("activation", list(MoEActivation))
76 +def test_triton_activation_metadata_tracks_shared_apply(activation: MoEActivation):
77 + assert TritonExperts._supports_activation(
78 + activation
79 + ) == apply_moe_activation_supported(activation)
80 +
81 +
82 def make_test_tensors(
83 m: int,
84 n: int,
vllm/model_executor/layers/fused_moe/__init__.py
+2
@@ -5,6 +5,7 @@ from contextlib import contextmanager
5 from typing import Any
6
7 from vllm.model_executor.layers.fused_moe.activation import (
8 + ApplyMoEActivationConfig,
9 MoEActivation,
10 activation_without_mul,
11 apply_moe_activation,
@@ -63,6 +64,7 @@ def get_config() -> dict[str, Any] | None:
64
65
66 __all__ = [
67 + "ApplyMoEActivationConfig",
68 "FusedMoEFactory",
69 "FusedMoERouter",
70 "FusedMoEConfig",
vllm/model_executor/layers/fused_moe/activation.py
+91 -16
@@ -2,11 +2,19 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """MoE activation function enum and utilities."""
4
5 +from dataclasses import dataclass
6 from enum import Enum
7 +from typing import TYPE_CHECKING
8
9 import torch
10 import torch.nn.functional as F
11
12 +if TYPE_CHECKING:
13 + from vllm.model_executor.layers.fused_moe.config import (
14 + FusedMoEConfig,
15 + FusedMoEQuantConfig,
16 + )
17 +
18
19 class MoEActivation(Enum):
20 """Activation functions for MoE layers."""
@@ -109,6 +117,66 @@ def activation_without_mul(activation: str) -> str:
117 return MoEActivation.from_str(activation).without_mul().value
118
119
120 +_APPLY_MOE_ACTIVATIONS = frozenset(
121 + {
122 + MoEActivation.SILU,
123 + MoEActivation.GELU,
124 + MoEActivation.GELU_TANH,
125 + MoEActivation.SITU,
126 + MoEActivation.SWIGLUOAI,
127 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
128 + MoEActivation.SWIGLUSTEP,
129 + MoEActivation.SILU_NO_MUL,
130 + MoEActivation.GELU_NO_MUL,
131 + MoEActivation.GELU_TANH_NO_MUL,
132 + MoEActivation.RELU2_NO_MUL,
133 + }
134 +)
135 +
136 +
137 +def apply_moe_activation_supported(activation: MoEActivation) -> bool:
138 + """Whether ``apply_moe_activation`` supports an activation."""
139 + return activation in _APPLY_MOE_ACTIVATIONS
140 +
141 +
142 +@dataclass(frozen=True)
143 +class ApplyMoEActivationConfig:
144 + """Configuration forwarded to ``apply_moe_activation``."""
145 +
146 + clamp_limit: float | None = None
147 + alpha: float = 1.0
148 + beta: float = 0.0
149 + activation_situ_beta: float | None = None
150 + activation_situ_linear_beta: float | None = None
151 +
152 + @classmethod
153 + def from_configs(
154 + cls,
155 + moe_config: "FusedMoEConfig",
156 + quant_config: "FusedMoEQuantConfig",
157 + ) -> "ApplyMoEActivationConfig":
158 + """Build from the model and quantization configurations."""
159 + clamp_limit = quant_config.gemm1_clamp_limit
160 + if clamp_limit is None:
161 + clamp_limit = moe_config.swiglu_limit
162 + alpha = quant_config.gemm1_alpha
163 + if alpha is None:
164 + alpha = moe_config.swiglu_alpha
165 + beta = quant_config.gemm1_beta
166 + if beta is None:
167 + beta = moe_config.swiglu_beta
168 + return cls(
169 + clamp_limit=clamp_limit,
170 + alpha=1.0 if alpha is None else alpha,
171 + beta=0.0 if beta is None else beta,
172 + activation_situ_beta=moe_config.activation_situ_beta,
173 + activation_situ_linear_beta=moe_config.activation_situ_linear_beta,
174 + )
175 +
176 +
177 +_DEFAULT_APPLY_MOE_ACTIVATION_CONFIG = ApplyMoEActivationConfig()
178 +
179 +
180 def silu_and_mul_with_clamp(
181 output: torch.Tensor,
182 input: torch.Tensor,
@@ -130,20 +198,21 @@ def apply_moe_activation(
198 output: torch.Tensor,
199 input: torch.Tensor,
200 *,
133 - clamp_limit: float | None = None,
134 - alpha: float = 1.0,
135 - beta: float = 0.0,
201 + activation_config: ApplyMoEActivationConfig | None = None,
202 topk_ids: torch.Tensor | None = None,
203 expert_map: torch.Tensor | None = None,
138 - activation_situ_beta: float | None = None,
139 - activation_situ_linear_beta: float | None = None,
204 ) -> torch.Tensor:
205 """Apply MoE activation function.
206
143 - ``clamp_limit``/``alpha``/``beta`` (from the quant config) drive the clamped
144 - SwiGLU kernels: ``SILU`` + ``clamp_limit`` and ``SWIGLUOAI_UNINTERLEAVE`` both
145 - map to ``silu_and_mul_with_clamp``. Other activations ignore them.
207 + The configuration drives specialized activation behavior. Routing tensors
208 + remain per-call inputs because they depend on the current token assignment.
209 """
210 + config = (
211 + _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG
212 + if activation_config is None
213 + else activation_config
214 + )
215 +
216 assert input.dim() == 2, "Input must be 2D"
217 assert output.dim() == 2, "Output must be 2D"
218 if activation.is_gated:
@@ -159,8 +228,10 @@ def apply_moe_activation(
228
229 # Activations with gated multiplication (gate × activation(up))
230 if activation == MoEActivation.SILU:
162 - if clamp_limit is not None:
163 - silu_and_mul_with_clamp(output, input, clamp_limit, topk_ids, expert_map)
231 + if config.clamp_limit is not None:
232 + silu_and_mul_with_clamp(
233 + output, input, config.clamp_limit, topk_ids, expert_map
234 + )
235 else:
236 torch.ops._C.silu_and_mul(output, input)
237 elif activation == MoEActivation.GELU:
@@ -175,23 +246,27 @@ def apply_moe_activation(
246 # bypassed the config plumbing, so fail rather than silently use 1.0.
247 # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel
248 # (up passed through), matching SituAndMul(linear_beta=None).
178 - assert activation_situ_beta is not None, (
249 + assert config.activation_situ_beta is not None, (
250 "SITU requires activation_situ_beta from FusedMoEConfig"
251 )
252 torch.ops._C.situ_and_mul(
253 output,
254 input,
184 - activation_situ_beta,
255 + config.activation_situ_beta,
256 -1.0
186 - if activation_situ_linear_beta is None
187 - else activation_situ_linear_beta,
257 + if config.activation_situ_linear_beta is None
258 + else config.activation_situ_linear_beta,
259 )
260 elif activation == MoEActivation.SWIGLUOAI:
261 torch.ops._C.swigluoai_and_mul(output, input)
262 elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
263 # SwiGLU-OAI on packed w13 (gate = first half, up = second half).
193 - assert clamp_limit is not None, "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
194 - torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, alpha, beta)
264 + assert config.clamp_limit is not None, (
265 + "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
266 + )
267 + torch.ops._C.silu_and_mul_with_clamp(
268 + output, input, config.clamp_limit, config.alpha, config.beta
269 + )
270 elif activation == MoEActivation.SWIGLUSTEP:
271 from vllm.model_executor.layers.activation import swiglustep_and_mul_triton
272
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+48 -37
@@ -8,8 +8,10 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm import _custom_ops as ops
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.fused_moe.activation import (
11 + ApplyMoEActivationConfig,
12 MoEActivation,
13 apply_moe_activation,
14 + apply_moe_activation_supported,
15 )
16 from vllm.model_executor.layers.fused_moe.config import (
17 FusedMoEConfig,
@@ -76,6 +78,8 @@ def run_cutlass_moe_fp8(
78 use_batched_format: bool,
79 topk_weights: torch.Tensor | None,
80 permute_scratch: MoEPermuteScratch | None,
81 + *,
82 + activation_config: ApplyMoEActivationConfig | None = None,
83 ):
84 a1q = hidden_states
85
@@ -231,7 +235,12 @@ def run_cutlass_moe_fp8(
235 per_out_ch,
236 )
237
234 - apply_moe_activation(activation, act_out, mm1_out)
238 + apply_moe_activation(
239 + activation,
240 + act_out,
241 + mm1_out,
242 + activation_config=activation_config,
243 + )
244
245 a2q, a2q_scale = ops.scaled_fp8_quant(
246 act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out
@@ -319,12 +328,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
328
329 @staticmethod
330 def _supports_activation(activation: MoEActivation) -> bool:
322 - return activation in [
323 - MoEActivation.SILU,
324 - MoEActivation.GELU,
325 - MoEActivation.GELU_TANH,
326 - MoEActivation.SWIGLUOAI,
327 - ]
331 + return activation.is_gated and apply_moe_activation_supported(activation)
332
333 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
334 # Let PrepareAndFinalize::finalize() decide the impl.
@@ -397,6 +401,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
401 use_batched_format,
402 topk_weights,
403 self._get_permute_scratch(),
404 + activation_config=self.activation_config,
405 )
406
407
@@ -512,6 +517,8 @@ def run_cutlass_moe_fp4(
517 e: int,
518 device: torch.device,
519 apply_router_weight_on_input: bool = False,
520 + *,
521 + activation_config: ApplyMoEActivationConfig | None = None,
522 ) -> None:
523 """
524 MoE implementation for FP4 Inputs
@@ -633,7 +640,9 @@ def run_cutlass_moe_fp4(
640 blockscale_offsets[:-1],
641 )
642 del rep_a_fp4, rep_a_blockscale
636 - if activation == MoEActivation.SILU:
643 + if activation == MoEActivation.SILU and (
644 + activation_config is None or activation_config.clamp_limit is None
645 + ):
646 # Fused SiLU+Mul+NVFP4 quantization
647 # Note: c2 workspace is no longer needed since SiLU is fused with quantization.
648 # c3 reuses workspace13 after c1 is consumed.
@@ -641,7 +650,12 @@ def run_cutlass_moe_fp4(
650 c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk
651 )
652 else:
644 - apply_moe_activation(activation, c2, c1)
653 + apply_moe_activation(
654 + activation,
655 + c2,
656 + c1,
657 + activation_config=activation_config,
658 + )
659 int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(
660 c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk
661 )
@@ -716,17 +730,7 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
730 # fallback + separate fp4 quantization in run_cutlass_moe_fp4().
731 # Non-gated activations (_NO_MUL) are also supported for models
732 # like Nemotron-Nano that don't use gated MLP.
719 - return activation in [
720 - MoEActivation.SILU,
721 - MoEActivation.GELU,
722 - MoEActivation.GELU_TANH,
723 - MoEActivation.SWIGLUOAI,
724 - MoEActivation.SWIGLUSTEP,
725 - MoEActivation.SILU_NO_MUL,
726 - MoEActivation.GELU_NO_MUL,
727 - MoEActivation.GELU_TANH_NO_MUL,
728 - MoEActivation.RELU2_NO_MUL,
729 - ]
733 + return apply_moe_activation_supported(activation)
734
735 @staticmethod
736 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -807,6 +811,7 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
811 e=e,
812 device=hidden_states.device,
813 apply_router_weight_on_input=apply_router_weight_on_input,
814 + activation_config=self.activation_config,
815 )
816
817
@@ -828,6 +833,8 @@ def run_cutlass_moe_mxfp4(
833 e: int,
834 device: torch.device,
835 apply_router_weight_on_input: bool = False,
836 + *,
837 + activation_config: ApplyMoEActivationConfig | None = None,
838 ) -> None:
839 """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
840 is_gated = activation.is_gated
@@ -909,12 +916,19 @@ def run_cutlass_moe_mxfp4(
916 blockscale_offsets[:-1],
917 )
918 del rep_a_fp4, rep_a_blockscale
912 - if activation == MoEActivation.SILU:
919 + if activation == MoEActivation.SILU and (
920 + activation_config is None or activation_config.clamp_limit is None
921 + ):
922 int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
923 c1, expert_offsets, blockscale_offsets, e, num_topk
924 )
925 else:
917 - apply_moe_activation(activation, c2, c1)
926 + apply_moe_activation(
927 + activation,
928 + c2,
929 + c1,
930 + activation_config=activation_config,
931 + )
932 int_fp4, int_blockscale = ops.mxfp4_experts_quant(
933 c2, expert_offsets, blockscale_offsets, e, num_topk
934 )
@@ -1021,15 +1035,7 @@ class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
1035
1036 @staticmethod
1037 def _supports_activation(activation: MoEActivation) -> bool:
1024 - return activation in [
1025 - MoEActivation.SILU,
1026 - MoEActivation.GELU,
1027 - MoEActivation.SWIGLUOAI,
1028 - MoEActivation.SWIGLUSTEP,
1029 - MoEActivation.SILU_NO_MUL,
1030 - MoEActivation.GELU_NO_MUL,
1031 - MoEActivation.RELU2_NO_MUL,
1032 - ]
1038 + return apply_moe_activation_supported(activation)
1039
1040 @staticmethod
1041 def _supports_parallel_config(
@@ -1102,6 +1108,7 @@ class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
1108 e=e,
1109 device=hidden_states.device,
1110 apply_router_weight_on_input=apply_router_weight_on_input,
1111 + activation_config=self.activation_config,
1112 )
1113
1114
@@ -1139,6 +1146,8 @@ def run_cutlass_moe_w4a8_fp8(
1146 topk_weights: torch.Tensor | None,
1147 group_size: int,
1148 permute_scratch: MoEPermuteScratch | None,
1149 + *,
1150 + activation_config: ApplyMoEActivationConfig | None = None,
1151 ):
1152 a1q = hidden_states
1153 M = a1q.size(0)
@@ -1218,7 +1227,12 @@ def run_cutlass_moe_w4a8_fp8(
1227 s_strides1,
1228 )
1229
1221 - apply_moe_activation(activation, act_out, mm1_out)
1230 + apply_moe_activation(
1231 + activation,
1232 + act_out,
1233 + mm1_out,
1234 + activation_config=activation_config,
1235 + )
1236
1237 a2q, a2q_scale = ops.scaled_fp8_quant(
1238 act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out
@@ -1330,11 +1344,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1344
1345 @staticmethod
1346 def _supports_activation(activation: MoEActivation) -> bool:
1333 - return activation in (
1334 - MoEActivation.SILU,
1335 - MoEActivation.GELU,
1336 - MoEActivation.SWIGLUOAI,
1337 - )
1347 + return activation.is_gated and apply_moe_activation_supported(activation)
1348
1349 @staticmethod
1350 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -1438,4 +1448,5 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1448 topk_weights,
1449 self.group_size,
1450 self._get_permute_scratch(),
1451 + activation_config=self.activation_config,
1452 )
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+20 -17
@@ -12,7 +12,10 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm import envs
13 from vllm.forward_context import get_forward_context
14 from vllm.logger import init_logger
15 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15 +from vllm.model_executor.layers.fused_moe.activation import (
16 + MoEActivation,
17 + apply_moe_activation_supported,
18 +)
19 from vllm.model_executor.layers.fused_moe.config import (
20 FusedMoEConfig,
21 FusedMoEParallelConfig,
@@ -246,18 +249,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
249 def _supports_activation(activation: MoEActivation) -> bool:
250 # Humming uses apply_moe_activation() callback for activation,
251 # so any activation supported there can be used here.
249 - return activation in [
250 - MoEActivation.SILU,
251 - MoEActivation.GELU,
252 - MoEActivation.GELU_TANH,
253 - MoEActivation.SWIGLUOAI,
254 - MoEActivation.SITU,
255 - MoEActivation.SWIGLUSTEP,
256 - MoEActivation.SILU_NO_MUL,
257 - MoEActivation.GELU_NO_MUL,
258 - MoEActivation.GELU_TANH_NO_MUL,
259 - MoEActivation.RELU2_NO_MUL,
260 - ]
252 + return apply_moe_activation_supported(activation)
253
254 @staticmethod
255 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -510,11 +502,22 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
502 output: torch.Tensor,
503 input: torch.Tensor,
504 ) -> None:
513 - swiglu_limit = self.quant_config.gemm1_clamp_limit
514 - if activation == MoEActivation.SILU and swiglu_limit is not None:
515 - swiglu_limit_func(output=output, input=input, swiglu_limit=swiglu_limit)
505 + activation_config = self.activation_config
506 + if (
507 + activation == MoEActivation.SILU
508 + and activation_config.clamp_limit is not None
509 + ):
510 + swiglu_limit_func(
511 + output=output,
512 + input=input,
513 + swiglu_limit=activation_config.clamp_limit,
514 + )
515 else:
517 - self.activation(activation=activation, input=input, output=output)
516 + self.activation(
517 + activation=activation,
518 + input=input,
519 + output=output,
520 + )
521
522
523 class HummingIndexedExperts(HummingExpertsBase):
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+40 -100
@@ -10,8 +10,10 @@ import torch
10 import vllm._custom_ops as ops
11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm.model_executor.layers.fused_moe.activation import (
13 + ApplyMoEActivationConfig,
14 MoEActivation,
15 apply_moe_activation,
16 + apply_moe_activation_supported,
17 )
18 from vllm.model_executor.layers.fused_moe.config import (
19 FusedMoEConfig,
@@ -72,7 +74,7 @@ def _fused_marlin_moe(
74 expert_ids: torch.Tensor,
75 num_tokens_post_padded: torch.Tensor,
76 activation: MoEActivation = MoEActivation.SILU,
75 - activation_func: Callable[..., None] = apply_moe_activation,
77 + activation_func: Callable[..., None] | None = None,
78 topk_ids: torch.Tensor | None = None,
79 input_global_scale1: torch.Tensor | None = None,
80 input_global_scale2: torch.Tensor | None = None,
@@ -90,11 +92,7 @@ def _fused_marlin_moe(
92 output: torch.Tensor | None = None,
93 input_dtype: torch.dtype | None = None,
94 is_k_full: bool = True,
93 - clamp_limit: float | None = None,
94 - gemm1_alpha: float = 1.0,
95 - gemm1_beta: float = 0.0,
96 - activation_situ_beta: float | None = None,
97 - activation_situ_linear_beta: float | None = None,
95 + activation_config: ApplyMoEActivationConfig | None = None,
96 ) -> torch.Tensor:
97 assert hidden_states.ndim == 2
98 M, K = hidden_states.size()
@@ -162,20 +160,29 @@ def _fused_marlin_moe(
160 use_fp32_reduce=True,
161 is_zp_float=False,
162 )
165 - # apply_moe_activation fuses the clamp/gate params: SILU + clamp_limit and
166 - # SWIGLUOAI_UNINTERLEAVE both map to the silu_and_mul_with_clamp kernel.
167 - activation_func(
168 - activation,
169 - intermediate_cache2,
170 - intermediate_cache1.view(-1, w13_num_shards * N),
171 - clamp_limit=clamp_limit,
172 - alpha=gemm1_alpha,
173 - beta=gemm1_beta,
174 - topk_ids=topk_ids,
175 - expert_map=expert_map,
176 - activation_situ_beta=activation_situ_beta,
177 - activation_situ_linear_beta=activation_situ_linear_beta,
178 - )
163 + activation_input = intermediate_cache1.view(-1, w13_num_shards * N)
164 + if activation_func is None:
165 + config = (
166 + ApplyMoEActivationConfig()
167 + if activation_config is None
168 + else activation_config
169 + )
170 + apply_moe_activation(
171 + activation,
172 + intermediate_cache2,
173 + activation_input,
174 + activation_config=config,
175 + topk_ids=topk_ids,
176 + expert_map=expert_map,
177 + )
178 + else:
179 + activation_func(
180 + activation,
181 + intermediate_cache2,
182 + activation_input,
183 + topk_ids=topk_ids,
184 + expert_map=expert_map,
185 + )
186
187 if output is None:
188 output = intermediate_cache3
@@ -238,7 +245,7 @@ def fused_marlin_moe(
245 apply_router_weight_on_input: bool = False,
246 global_num_experts: int = -1,
247 activation: MoEActivation = MoEActivation.SILU,
241 - activation_func: Callable[..., None] = apply_moe_activation,
248 + activation_func: Callable[..., None] | None = None,
249 moe_sum: Callable[..., torch.Tensor | None] | None = None,
250 expert_map: torch.Tensor | None = None,
251 input_global_scale1: torch.Tensor | None = None,
@@ -257,11 +264,7 @@ def fused_marlin_moe(
264 is_k_full: bool = True,
265 output: torch.Tensor | None = None,
266 input_dtype: torch.dtype | None = None,
260 - clamp_limit: float | None = None,
261 - gemm1_alpha: float = 1.0,
262 - gemm1_beta: float = 0.0,
263 - activation_situ_beta: float | None = None,
264 - activation_situ_linear_beta: float | None = None,
267 + activation_config: ApplyMoEActivationConfig | None = None,
268 ) -> torch.Tensor:
269 """
270 This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -363,8 +366,7 @@ def fused_marlin_moe(
366 num_tokens_post_padded=num_tokens_post_padded,
367 activation=activation,
368 activation_func=activation_func,
366 - activation_situ_beta=activation_situ_beta,
367 - activation_situ_linear_beta=activation_situ_linear_beta,
369 + activation_config=activation_config,
370 input_global_scale1=input_global_scale1,
371 input_global_scale2=input_global_scale2,
372 global_scale1=global_scale1,
@@ -381,9 +383,6 @@ def fused_marlin_moe(
383 output=None,
384 input_dtype=input_dtype,
385 is_k_full=is_k_full,
384 - clamp_limit=clamp_limit,
385 - gemm1_alpha=gemm1_alpha,
386 - gemm1_beta=gemm1_beta,
386 ).view(-1, topk, K)
387
388 if output is None:
@@ -428,12 +427,8 @@ def batched_fused_marlin_moe(
427 is_k_full: bool = True,
428 output: torch.Tensor | None = None,
429 input_dtype: torch.dtype | None = None,
431 - clamp_limit: float | None = None,
432 - gemm1_alpha: float = 1.0,
433 - gemm1_beta: float = 0.0,
434 - activation_func: Callable[..., None] = apply_moe_activation,
435 - activation_situ_beta: float | None = None,
436 - activation_situ_linear_beta: float | None = None,
430 + activation_func: Callable[..., None] | None = None,
431 + activation_config: ApplyMoEActivationConfig | None = None,
432 ) -> torch.Tensor:
433 """
434 This function massages the inputs so the batched hidden_states can be
@@ -542,8 +537,7 @@ def batched_fused_marlin_moe(
537 apply_router_weight_on_input=apply_router_weight_on_input,
538 activation=activation,
539 activation_func=activation_func,
545 - activation_situ_beta=activation_situ_beta,
546 - activation_situ_linear_beta=activation_situ_linear_beta,
540 + activation_config=activation_config,
541 expert_map=expert_map,
542 block_size_m=block_size_m,
543 sorted_token_ids=sorted_token_ids,
@@ -565,9 +559,6 @@ def batched_fused_marlin_moe(
559 output=output.view(-1, K) if output is not None else output,
560 input_dtype=input_dtype,
561 is_k_full=is_k_full,
568 - clamp_limit=clamp_limit,
569 - gemm1_alpha=gemm1_alpha,
570 - gemm1_beta=gemm1_beta,
562 )
563
564 output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -602,16 +593,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
593 self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
594 self.is_k_full = is_k_full
595 self.input_dtype = get_marlin_input_dtype()
605 - self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
606 - # Gated-activation params (used by SWIGLUOAI_UNINTERLEAVE on packed w13).
607 - # silu == swigluoai with alpha=1, beta=0; configs that don't set these
608 - # (plain silu) fall back to the silu identity.
609 - self.gemm1_alpha = (
610 - quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
611 - )
612 - self.gemm1_beta = (
613 - quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
614 - )
596
597 super().__init__(
598 moe_config=moe_config,
@@ -655,19 +636,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
636 def _supports_activation(activation: MoEActivation) -> bool:
637 # Marlin uses apply_moe_activation() callback for activation,
638 # so any activation supported there can be used here.
658 - return activation in [
659 - MoEActivation.SILU,
660 - MoEActivation.GELU,
661 - MoEActivation.GELU_TANH,
662 - MoEActivation.SITU,
663 - MoEActivation.SWIGLUOAI,
664 - MoEActivation.SWIGLUOAI_UNINTERLEAVE,
665 - MoEActivation.SWIGLUSTEP,
666 - MoEActivation.SILU_NO_MUL,
667 - MoEActivation.GELU_NO_MUL,
668 - MoEActivation.GELU_TANH_NO_MUL,
669 - MoEActivation.RELU2_NO_MUL,
670 - ]
639 + return apply_moe_activation_supported(activation)
640
641 @staticmethod
642 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -808,10 +777,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
777 global_num_experts=global_num_experts,
778 activation=activation,
779 activation_func=self.activation,
811 - activation_situ_beta=self.moe_config.activation_situ_beta,
812 - activation_situ_linear_beta=(
813 - self.moe_config.activation_situ_linear_beta
814 - ),
780 + activation_config=self.activation_config,
781 moe_sum=self.moe_sum,
782 expert_map=expert_map,
783 output=output,
@@ -825,9 +791,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
791 sort_indices2=self.w2_g_idx_sort_indices,
792 is_k_full=self.is_k_full,
793 input_dtype=self.input_dtype,
828 - clamp_limit=self.gemm1_clamp_limit,
829 - gemm1_alpha=self.gemm1_alpha,
830 - gemm1_beta=self.gemm1_beta,
794 )
795 return
796
@@ -847,13 +810,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
810 act_output: torch.Tensor,
811 act_input: torch.Tensor,
812 *,
850 - clamp_limit: float | None = None,
851 - alpha: float = 1.0,
852 - beta: float = 0.0,
813 topk_ids: torch.Tensor | None = None,
814 expert_map: torch.Tensor | None = None,
855 - activation_situ_beta: float | None = None,
856 - activation_situ_linear_beta: float | None = None,
815 ) -> None:
816 # act_input = intermediate_cache1 (M*topk, 2N for gated)
817 # act_output = intermediate_cache2 (M*topk, N)
@@ -887,13 +845,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
845 act_enum,
846 act_output,
847 act_input,
890 - clamp_limit=clamp_limit,
891 - alpha=alpha,
892 - beta=beta,
848 topk_ids=topk_ids,
849 expert_map=expert_map,
895 - activation_situ_beta=activation_situ_beta,
896 - activation_situ_linear_beta=activation_situ_linear_beta,
850 )
851 lora_state["cache2"] = act_output
852
@@ -941,8 +894,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
894 global_num_experts=global_num_experts,
895 activation=activation,
896 activation_func=activation_with_lora,
944 - activation_situ_beta=self.moe_config.activation_situ_beta,
945 - activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
897 + activation_config=self.activation_config,
898 moe_sum=moe_sum_with_lora,
899 expert_map=expert_map,
900 output=output,
@@ -954,9 +906,6 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
906 sort_indices2=self.w2_g_idx_sort_indices,
907 is_k_full=self.is_k_full,
908 input_dtype=self.input_dtype,
957 - clamp_limit=self.gemm1_clamp_limit,
958 - gemm1_alpha=self.gemm1_alpha,
959 - gemm1_beta=self.gemm1_beta,
909 )
910
911 def moe_sum(
@@ -1051,9 +1000,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1000 act: MoEActivation,
1001 act_output: torch.Tensor,
1002 act_input: torch.Tensor,
1054 - *,
1055 - activation_situ_beta: float | None = None,
1056 - activation_situ_linear_beta: float | None = None,
1003 **kwargs,
1004 ) -> None:
1005 if act != MoEActivation.SITU:
@@ -1061,20 +1007,18 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1007 act,
1008 act_output,
1009 act_input,
1064 - activation_situ_beta=activation_situ_beta,
1065 - activation_situ_linear_beta=activation_situ_linear_beta,
1010 **kwargs,
1011 )
1012 return
1013
1014 num_experts, max_num_tokens = hidden_states.shape[:2]
1071 - beta = 1.0 if activation_situ_beta is None else activation_situ_beta
1072 - linear_beta = activation_situ_linear_beta
1015 + beta = self.activation_config.activation_situ_beta
1016 + linear_beta = self.activation_config.activation_situ_linear_beta
1017 torch.ops._C.masked_situ_and_mul(
1018 act_output.view(num_experts, max_num_tokens, -1),
1019 act_input.view(num_experts, max_num_tokens, -1),
1020 expert_tokens_meta.expert_num_tokens,
1077 - beta,
1021 + 1.0 if beta is None else beta,
1022 -1.0 if linear_beta is None else linear_beta,
1023 )
1024
@@ -1105,10 +1049,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1049 w2_zeros=self.w2_zp,
1050 input_dtype=self.input_dtype,
1051 is_k_full=self.is_k_full,
1108 - clamp_limit=self.gemm1_clamp_limit,
1109 - gemm1_alpha=self.gemm1_alpha,
1110 - gemm1_beta=self.gemm1_beta,
1052 activation_func=activation_func,
1112 - activation_situ_beta=self.moe_config.activation_situ_beta,
1113 - activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
1053 + activation_config=self.activation_config,
1054 )
vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py
+2 -5
@@ -104,16 +104,13 @@ class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
104 )
105
106 if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
107 - limit = self.quant_config.gemm1_clamp_limit
108 - if limit is None:
107 + if self.activation_config.clamp_limit is None:
108 raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit")
109 apply_moe_activation(
110 activation,
111 output,
112 input,
114 - clamp_limit=float(limit),
115 - alpha=self.gemm1_alpha,
116 - beta=self.gemm1_beta,
113 + activation_config=self.activation_config,
114 )
115 return
116 super().activation(activation, output, input)
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+4 -6
@@ -362,9 +362,7 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
362 expert_tokens_meta: mk.ExpertTokensMetadata | None,
363 apply_router_weight_on_input: bool,
364 ):
365 - # `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`.
366 - limit = self.quant_config.gemm1_clamp_limit
367 - limit = None if limit is None else float(limit)
365 + activation_config = self.activation_config
366 out = fused_moe_mxfp8_native(
367 hidden_states,
368 w1,
@@ -373,9 +371,9 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
371 self.w2_scale_val,
372 topk_weights,
373 topk_ids,
376 - alpha=self.gemm1_alpha,
377 - beta=self.gemm1_beta,
378 - limit=limit,
374 + alpha=activation_config.alpha,
375 + beta=activation_config.beta,
376 + limit=activation_config.clamp_limit,
377 global_num_experts=global_num_experts,
378 expert_map=expert_map,
379 )
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+13 -33
@@ -6,7 +6,10 @@ import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm import _custom_ops as ops
9 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9 +from vllm.model_executor.layers.fused_moe.activation import (
10 + MoEActivation,
11 + apply_moe_activation_supported,
12 +)
13 from vllm.model_executor.layers.fused_moe.config import (
14 FusedMoEConfig,
15 FusedMoEParallelConfig,
@@ -69,15 +72,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
72 self.quantization_emulation = False
73 super().__init__(moe_config, quant_config)
74
72 - self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
73 - # Gated-activation params: silu == swigluoai with alpha=1, beta=0.
74 - self.gemm1_alpha = (
75 - quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
76 - )
77 - self.gemm1_beta = (
78 - quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
79 - )
80 -
75 @staticmethod
76 def activation_format() -> mk.FusedMoEActivationFormat:
77 return mk.FusedMoEActivationFormat.Standard
@@ -136,19 +130,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
130
131 @staticmethod
132 def _supports_activation(activation: MoEActivation) -> bool:
139 - return activation in [
140 - MoEActivation.SILU,
141 - MoEActivation.GELU,
142 - MoEActivation.GELU_TANH,
143 - MoEActivation.SITU,
144 - MoEActivation.SWIGLUOAI,
145 - MoEActivation.SWIGLUOAI_UNINTERLEAVE,
146 - MoEActivation.SWIGLUSTEP,
147 - MoEActivation.SILU_NO_MUL,
148 - MoEActivation.GELU_NO_MUL,
149 - MoEActivation.GELU_TANH_NO_MUL,
150 - MoEActivation.RELU2_NO_MUL,
151 - ]
133 + return apply_moe_activation_supported(activation)
134
135 @staticmethod
136 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -171,17 +153,18 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
153 input: torch.Tensor,
154 **kwargs,
155 ) -> None:
174 - gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
175 - if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
176 - swiglu_limit_func(output, input, float(gemm1_clamp_limit))
156 + activation_config = self.activation_config
157 + if (
158 + activation == MoEActivation.SILU
159 + and activation_config.clamp_limit is not None
160 + ):
161 + swiglu_limit_func(output, input, activation_config.clamp_limit)
162 return
163
164 # SWIGLUOAI_UNINTERLEAVE routes to the silu_and_mul_with_clamp kernel and
180 - # needs the clamped-SwiGLU params (gemm1_clamp_limit/alpha/beta read from
181 - # the quant config in __init__) forwarded; without a clamp_limit it
182 - # asserts. Other activations ignore alpha/beta/clamp_limit.
165 + # requires a clamp limit. Other activations ignore these parameters.
166 if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
184 - assert gemm1_clamp_limit is not None, (
167 + assert activation_config.clamp_limit is not None, (
168 "SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit"
169 )
170
@@ -189,9 +172,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
172 activation,
173 output,
174 input,
192 - clamp_limit=gemm1_clamp_limit,
193 - alpha=self.gemm1_alpha,
194 - beta=self.gemm1_beta,
175 )
176
177 def workspace_shapes(
vllm/model_executor/layers/fused_moe/modular_kernel.py
+5 -18
@@ -12,6 +12,7 @@ import torch
12 import vllm.envs as envs
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.fused_moe.activation import (
15 + ApplyMoEActivationConfig,
16 MoEActivation,
17 apply_moe_activation,
18 )
@@ -497,6 +498,9 @@ class FusedMoEExperts(ABC):
498
499 self.moe_config = moe_config
500 self.quant_config = quant_config
501 + self.activation_config = ApplyMoEActivationConfig.from_configs(
502 + moe_config, quant_config
503 + )
504 self.max_num_tokens = max_num_tokens
505 self.num_dispatchers = num_dispatchers
506
@@ -885,33 +889,16 @@ class FusedMoEExpertsModular(FusedMoEExperts):
889 output: torch.Tensor,
890 input: torch.Tensor,
891 *,
888 - clamp_limit: float | None = None,
889 - alpha: float = 1.0,
890 - beta: float = 0.0,
892 topk_ids: torch.Tensor | None = None,
893 expert_map: torch.Tensor | None = None,
893 - activation_situ_beta: float | None = None,
894 - activation_situ_linear_beta: float | None = None,
894 ) -> None:
895 apply_moe_activation(
896 activation,
897 output,
898 input,
900 - clamp_limit=clamp_limit,
901 - alpha=alpha,
902 - beta=beta,
899 + activation_config=self.activation_config,
900 topk_ids=topk_ids,
901 expert_map=expert_map,
905 - activation_situ_beta=(
906 - self.moe_config.activation_situ_beta
907 - if activation_situ_beta is None
908 - else activation_situ_beta
909 - ),
910 - activation_situ_linear_beta=(
911 - self.moe_config.activation_situ_linear_beta
912 - if activation_situ_linear_beta is None
913 - else activation_situ_linear_beta
914 - ),
902 )
903
904 @abstractmethod
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+5 -2
@@ -190,7 +190,9 @@ def select_nvfp4_moe_backend(
190 NVFP4_BACKENDS_WITH_CLAMP = {
191 NvFp4MoeBackend.FLASHINFER_TRTLLM,
192 NvFp4MoeBackend.FLASHINFER_CUTLASS,
193 + NvFp4MoeBackend.VLLM_CUTLASS,
194 NvFp4MoeBackend.MARLIN,
195 + NvFp4MoeBackend.HUMMING,
196 }
197
198 if config.swiglu_limit is not None:
@@ -258,8 +260,9 @@ def select_nvfp4_moe_backend(
260 raise ValueError(
261 f"Model sets swiglu_limit={config.swiglu_limit}, but the "
262 f"explicitly requested moe_backend={runner_backend!r} does "
261 - f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm' or "
262 - f"'flashinfer_cutlass' instead."
263 + f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', "
264 + f"'flashinfer_cutlass', 'cutlass', 'marlin', or 'humming' "
265 + f"instead."
266 )
267 return _return_or_raise(
268 requested_backend, config, weight_key, activation_key, activation_format