ff1f83b056 · tduncan/Nvidia-Optimization-Stack-VLLM

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

+2 -1

@@ -11,6 +11,7 @@ import torch

11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

12 from tests.kernels.moe.utils import make_dummy_moe_config

13 from vllm import _custom_ops as ops

14 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

15 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config

16 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8

17 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

@@ -161,7 +162,7 @@ def bench_run(

162 w2_fp8q_cutlass,

163 topk_weights,

164 topk_ids,

164 - activation="silu",

165 + activation=MoEActivation.SILU,

166 global_num_experts=num_experts,

167 )

168 torch.cuda.synchronize()

benchmarks/kernels/benchmark_moe.py

+3 -1

@@ -16,6 +16,7 @@ import torch

16 from ray.experimental.tqdm_ray import tqdm

17

18 from vllm.model_executor.layers.fused_moe import fused_topk

19 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

20 from vllm.model_executor.layers.fused_moe.config import (

21 FusedMoEConfig,

22 FusedMoEParallelConfig,

@@ -211,7 +212,8 @@ def benchmark_config(

212 hidden_dim=hidden_size,

213 intermediate_size_per_partition=shard_intermediate_size,

214 num_local_experts=num_experts,

214 - activation="silu",

215 + num_logical_experts=num_experts,

216 + activation=MoEActivation.SILU,

217 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

218 in_dtype=init_dtype,

219 routing_method=RoutingMethodType.TopK,

tests/kernels/moe/modular_kernel_tools/common.py

+2 -1

@@ -22,6 +22,7 @@ from vllm.distributed import (

22 )

23 from vllm.forward_context import set_forward_context

24 from vllm.model_executor.layers.fused_moe import fused_topk

25 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

26 from vllm.model_executor.layers.fused_moe.all2all_utils import (

27 maybe_make_prepare_finalize,

28 )

@@ -599,7 +600,7 @@ def make_modular_kernel(

600 moe_parallel_config=moe_parallel_config,

601 in_dtype=config.dtype,

602 max_num_tokens=next_power_of_2(config.M),

602 - activation="silu",

603 + activation=MoEActivation.SILU,

604 device=vllm_config.device_config.device,

605 routing_method=RoutingMethodType.DeepSeekV3,

606 )

tests/kernels/moe/test_cpu_fused_moe.py

+5 -4

@@ -6,6 +6,7 @@ import torch

6

7 from tests.kernels.allclose_default import get_default_atol, get_default_rtol

8 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN

11 from vllm.platforms import current_platform

12 from vllm.utils.torch_utils import set_random_seed

@@ -19,7 +20,7 @@ EXPERT_NUM = [

20 HIDDEN_DIM = [128, 2880]

21 INTERMEDIATE_DIM = [128, 2880]

22 BATCH_SIZE = [1, 64, 256]

22 -ACT = ["silu", "swigluoai"]

23 +ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI]

24 USE_BIAS = [True, False]

25 ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]

26 DTYPE = [torch.bfloat16]

@@ -33,7 +34,7 @@ def ref_fused_moe(

34 w2_bias: torch.Tensor | None,

35 topk_weights: torch.Tensor,

36 topk_ids: torch.Tensor,

36 - activation: str,

37 + activation: MoEActivation,

38 ) -> torch.Tensor:

39 len_experts = w13.size(0)

40

@@ -103,7 +104,7 @@ def test_cpu_fused_moe(

104 intermediate_size: int,

105 use_bias: bool,

106 dtype: torch.dtype,

106 - act: str,

107 + act: MoEActivation,

108 isa: str,

109 ):

110 set_random_seed(0)

@@ -153,7 +154,7 @@ def test_cpu_fused_moe(

154 w2_bias,

155 topk_weight,

156 topk_ids,

156 - act,

157 + act.value,

158 isa,

159 )

160

tests/kernels/moe/test_cutlass_moe.py

+2 -1

@@ -12,6 +12,7 @@ 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

16 from vllm.model_executor.layers.fused_moe.config import (

17 FUSED_MOE_UNQUANTIZED_CONFIG,

18 FusedMoEQuantConfig,

@@ -531,7 +532,7 @@ def test_run_cutlass_moe_fp8(

532 c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)

533 c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)

534

534 - activation = "silu"

535 + activation = MoEActivation.SILU

536 a1q, a1q_scale = moe_kernel_quantize_input(

537 mt.a, mt.a_scale, torch.float8_e4m3fn, per_act_token

538 )

tests/kernels/moe/test_deepep_deepgemm_moe.py

+2 -1

@@ -16,6 +16,7 @@ from typing_extensions import ParamSpec

16

17 from vllm.config import VllmConfig, set_current_vllm_config

18 from vllm.forward_context import set_forward_context

19 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

20 from vllm.model_executor.layers.fused_moe.config import (

21 FusedMoEQuantConfig,

22 fp8_w8a8_moe_quant_config,

@@ -324,7 +325,7 @@ def deepep_deepgemm_moe_impl(

325 w2=w2,

326 topk_weights=test_tensors.topk_weights,

327 topk_ids=test_tensors.topk,

327 - activation="silu",

328 + activation=MoEActivation.SILU,

329 global_num_experts=num_experts,

330 expert_map=build_expert_map(),

331 apply_router_weight_on_input=False,

tests/kernels/moe/test_deepep_moe.py

+2 -1

@@ -15,6 +15,7 @@ from vllm import _custom_ops as ops

15 from vllm.config import VllmConfig, set_current_vllm_config

16 from vllm.model_executor.layers.activation import SiluAndMul

17 from vllm.model_executor.layers.fused_moe import TritonExperts

18 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

19 from vllm.model_executor.layers.fused_moe.config import (

20 FusedMoEQuantConfig,

21 )

@@ -260,7 +261,7 @@ def deep_ep_moe_impl(

261 w2=w2,

262 topk_weights=topk_weights_chunk,

263 topk_ids=topk_chunk,

263 - activation="silu",

264 + activation=MoEActivation.SILU,

265 global_num_experts=num_experts,

266 expert_map=build_expert_map(),

267 apply_router_weight_on_input=False,

tests/kernels/moe/test_flashinfer.py

+12 -8

@@ -7,6 +7,7 @@ import torch

7

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

10 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.model_executor.layers.fused_moe.config import (

12 FusedMoEConfig,

13 FusedMoEParallelConfig,

@@ -93,9 +94,14 @@ class TestData:

94

95 @staticmethod

96 def make_moe_tensors_8bit(

96 - m: int, k: int, n: int, e: int, is_trtllm: bool, activation: str = "silu"

97 + m: int,

98 + k: int,

99 + n: int,

100 + e: int,

101 + is_trtllm: bool,

102 + activation: MoEActivation = MoEActivation.SILU,

103 ) -> "TestData":

98 - is_gated = activation != "relu2_no_mul"

104 + is_gated = activation.is_gated

105

106 hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10

107 w13 = torch.randn(

@@ -194,7 +200,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

200 topk_weights=topk_weights,

201 topk_ids=topk_ids,

202 inplace=False,

197 - activation="silu",

203 + activation=MoEActivation.SILU,

204 global_num_experts=e,

205 expert_map=None,

206 apply_router_weight_on_input=True,

@@ -219,21 +225,19 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

225 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)

226 @pytest.mark.parametrize("e", NUM_EXPERTS)

227 @pytest.mark.parametrize("topk", TOP_KS)

222 -@pytest.mark.parametrize("activation", ["silu", "relu2_no_mul"])

228 +@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])

229 def test_flashinfer_cutlass_moe_fp8_no_graph(

230 m: int,

231 n: int,

232 k: int,

233 e: int,

234 topk: int,

229 - activation: str,

235 + activation: MoEActivation,

236 monkeypatch,

237 workspace_init,

238 ):

239 set_random_seed(7)

240 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

235 - assert activation in ["silu", "relu2_no_mul"]

236 - is_act_and_mul = activation == "silu_and_mul"

241 with set_current_vllm_config(vllm_config):

242 td = TestData.make_moe_tensors_8bit(

243 m, k, n, e, is_trtllm=False, activation=activation

@@ -292,7 +296,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

296 device="cuda",

297 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

298 in_dtype=torch.bfloat16,

295 - is_act_and_mul=is_act_and_mul,

299 + is_act_and_mul=activation.is_gated,

300 routing_method=RoutingMethodType.TopK,

301 )

302

tests/kernels/moe/test_flashinfer_moe.py

+5 -6

@@ -13,6 +13,7 @@ from tests.kernels.utils import torch_moe

13 from vllm import _custom_ops as ops

14 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

15 from vllm.model_executor.layers.fused_moe import fused_topk

16 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

17 from vllm.model_executor.layers.fused_moe.config import (

18 FusedMoEConfig,

19 FusedMoEParallelConfig,

@@ -54,7 +55,7 @@ MNK_FACTORS = [

55 @pytest.mark.parametrize("e", [40, 64, 256])

56 @pytest.mark.parametrize("topk", [1, 6, 8])

57 @pytest.mark.parametrize("dtype", [torch.bfloat16])

57 -@pytest.mark.parametrize("activation", ["silu_and_mul", "relu2"])

58 +@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])

59 @torch.inference_mode()

60 def test_flashinfer_fp4_moe_no_graph(

61 m: int,

@@ -63,7 +64,7 @@ def test_flashinfer_fp4_moe_no_graph(

64 e: int,

65 topk: int,

66 dtype: torch.dtype,

66 - activation: str,

67 + activation: MoEActivation,

68 workspace_init,

69 ):

70 set_random_seed(7)

@@ -73,7 +74,7 @@ def test_flashinfer_fp4_moe_no_graph(

74 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

75

76 quant_blocksize = 16

76 - is_gated_act = activation == "silu_and_mul"

77 + is_gated_act = activation.is_gated

78

79 w1_q, w2_q, quant_config = make_test_quant_config(

80 e,

@@ -112,15 +113,13 @@ def test_flashinfer_fp4_moe_no_graph(

113 inplace=False,

114 )

115

115 - fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

116 -

116 flashinfer_output = flashinfer_experts(

117 hidden_states=a,

118 w1=w1_q,

119 w2=w2_q,

120 topk_weights=topk_weights,

121 topk_ids=topk_ids,

123 - activation=fi_activation,

122 + activation=activation,

123 )

124

125 # Reference check:

tests/kernels/moe/test_modular_oai_triton_moe.py

+2 -1

@@ -7,6 +7,7 @@ Test modular OAI Triton MoE

7 import pytest

8 import torch

9

10 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.utils.import_utils import has_triton_kernels

12

13 if not has_triton_kernels():

@@ -192,7 +193,7 @@ def oai_triton_moe_impl(

193 w2=w2,

194 topk_weights=topk_weights,

195 topk_ids=topk_ids,

195 - activation="swigluoai",

196 + activation=MoEActivation.SWIGLUOAI,

197 global_num_experts=num_experts,

198 expert_map=None,

199 apply_router_weight_on_input=False,

tests/kernels/moe/test_moe.py

+19 -5

@@ -29,6 +29,7 @@ from vllm.config import VllmConfig, set_current_vllm_config

29 from vllm.distributed.parallel_state import init_distributed_environment

30 from vllm.forward_context import get_forward_context, set_forward_context

31 from vllm.model_executor.layers.fused_moe import (

32 + MoEActivation,

33 fused_topk,

34 )

35 from vllm.model_executor.layers.fused_moe.config import (

@@ -1155,7 +1156,10 @@ def test_fused_marlin_moe_with_bias(m):

1156 @pytest.mark.parametrize("m", [1, 64, 256])

1157 @pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])

1158 @pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])

1158 -def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):

1159 +@pytest.mark.parametrize("activation", [MoEActivation.RELU2_NO_MUL])

1160 +def test_fused_marlin_moe_non_gated(

1161 + m: int, n: int, k: int, e: int, topk: int, activation: MoEActivation

1162 +):

1163 """Test Marlin MoE with non-gated activation (relu2_no_mul).

1164

1165 Non-gated activations like relu2 don't have the gate-up projection pattern,

@@ -1198,7 +1202,7 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):

1202 w2_data.w_ref,

1203 score,

1204 topk,

1201 - activation="relu2",

1205 + activation=activation,

1206 )

1207

1208 marlin_output = fused_marlin_moe(

@@ -1223,7 +1227,7 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):

1227 w2_zeros=w2_data.zeros,

1228 quant_type_id=quant_type.id,

1229 is_k_full=is_k_full,

1226 - activation="relu2_no_mul",

1230 + activation=activation,

1231 )

1232

1233 torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)

@@ -1330,9 +1334,18 @@ def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):

1334 @pytest.mark.parametrize("topk", [2])

1335 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])

1336 @pytest.mark.parametrize("with_bias", [False, True])

1333 -@pytest.mark.parametrize("activation", ["silu"])

1337 +@pytest.mark.parametrize("activation", [MoEActivation.SILU])

1338 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only test")

1335 -def test_cpu_fused_moe_basic(m, n, k, e, topk, dtype, with_bias, activation):

1339 +def test_cpu_fused_moe_basic(

1340 + m: int,

1341 + n: int,

1342 + k: int,

1343 + e: int,

1344 + topk: int,

1345 + dtype: torch.dtype,

1346 + with_bias: bool,

1347 + activation: MoEActivation,

1348 +):

1349 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE

1350

1351 device = "cpu"

@@ -1608,6 +1621,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(

1621 hidden_dim=k,

1622 intermediate_size_per_partition=n,

1623 num_local_experts=e,

1624 + num_logical_experts=e,

1625 activation="silu",

1626 device="cuda",

1627 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

tests/kernels/moe/test_pplx_cutlass_moe.py

+2 -1

@@ -9,6 +9,7 @@ from tests.kernels.utils import torch_experts

9 from vllm import _custom_ops as ops

10 from vllm.config import VllmConfig, set_current_vllm_config

11 from vllm.model_executor.layers.fused_moe import fused_topk

12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

13 from vllm.model_executor.layers.fused_moe.config import (

14 FusedMoEConfig,

15 FusedMoEParallelConfig,

@@ -149,7 +150,7 @@ def pplx_cutlass_moe(

150 num_local_experts=num_local_experts,

151 num_logical_experts=num_experts,

152 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

152 - activation="silu",

153 + activation=MoEActivation.SILU,

154 in_dtype=torch.bfloat16,

155 device="cuda",

156 routing_method=RoutingMethodType.Llama4,

tests/kernels/moe/test_triton_moe_no_act_mul.py

+14 -14

@@ -11,15 +11,11 @@ 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

15 from vllm.model_executor.layers.fused_moe.config import (

16 FUSED_MOE_UNQUANTIZED_CONFIG,

17 )

18 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts

18 -from vllm.model_executor.layers.fused_moe.utils import (

19 - GELU_NO_MUL,

20 - RELU2_NO_MUL,

21 - SILU_NO_MUL,

22 -)

19 from vllm.platforms import current_platform

20

21 # Test parameters

@@ -28,7 +24,11 @@ N_SIZES = [128, 256]

24 K_SIZES = [64, 128]

25 TOPK_VALUES = [1, 2]

26 NUM_EXPERTS = 8

31 -NO_MUL_ACTIVATIONS = [SILU_NO_MUL, GELU_NO_MUL, RELU2_NO_MUL]

27 +NO_MUL_ACTIVATIONS = [

28 + MoEActivation.SILU_NO_MUL,

29 + MoEActivation.GELU_NO_MUL,

30 + MoEActivation.RELU2_NO_MUL,

31 +]

32

33

34 def make_test_tensors(

@@ -73,7 +73,7 @@ def test_triton_experts_no_mul_activation(

73 n: int,

74 k: int,

75 topk: int,

76 - activation: str,

76 + activation: MoEActivation,

77 ):

78 hidden_states, w1, w2, topk_weights, topk_ids = make_test_tensors(

79 m, n, k, NUM_EXPERTS, topk

@@ -161,11 +161,11 @@ def test_workspace_shapes_no_mul_vs_gated():

161 )

162

163 ws1_no_mul, _, out_no_mul = experts.workspace_shapes(

164 - M, N, K, topk, 8, 8, None, SILU_NO_MUL

164 + M, N, K, topk, 8, 8, None, MoEActivation.SILU_NO_MUL

165 )

166

167 ws1_gated, _, out_gated = experts.workspace_shapes(

168 - M, N, K, topk, 8, 8, None, "silu"

168 + M, N, K, topk, 8, 8, None, MoEActivation.SILU

169 )

170

171 # For no_mul: activation_out_dim = N

@@ -202,10 +202,10 @@ def test_adjust_n_for_activation():

202 N = 256

203

204 # Gated activations should return N // 2

205 - assert experts.adjust_N_for_activation(N, "silu") == N // 2

206 - assert experts.adjust_N_for_activation(N, "gelu") == N // 2

205 + assert experts.adjust_N_for_activation(N, MoEActivation.SILU) == N // 2

206 + assert experts.adjust_N_for_activation(N, MoEActivation.GELU) == N // 2

207

208 # Non-gated activations should return N

209 - assert experts.adjust_N_for_activation(N, SILU_NO_MUL) == N

210 - assert experts.adjust_N_for_activation(N, GELU_NO_MUL) == N

211 - assert experts.adjust_N_for_activation(N, RELU2_NO_MUL) == N

209 + assert experts.adjust_N_for_activation(N, MoEActivation.SILU_NO_MUL) == N

210 + assert experts.adjust_N_for_activation(N, MoEActivation.GELU_NO_MUL) == N

211 + assert experts.adjust_N_for_activation(N, MoEActivation.RELU2_NO_MUL) == N

tests/kernels/moe/utils.py

+2 -1

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe import (

12 fused_experts,

13 fused_topk,

14 )

15 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

16 from vllm.model_executor.layers.fused_moe.config import (

17 FusedMoEConfig,

18 FusedMoEParallelConfig,

@@ -54,7 +55,7 @@ def make_dummy_moe_config(

55 num_local_experts=num_experts,

56 num_logical_experts=num_experts,

57 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

57 - activation="silu",

58 + activation=MoEActivation.SILU,

59 in_dtype=in_dtype,

60 device="cuda",

61 routing_method=RoutingMethodType.TopK,

tests/kernels/utils.py

+4 -3

@@ -15,6 +15,7 @@ from torch._prims_common import TensorLikeType

15 from tests.kernels.quant_utils import native_w8a8_block_matmul

16 from vllm.model_executor.custom_op import op_registry

17 from vllm.model_executor.layers.activation import SiluAndMul

18 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

19 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

20 from vllm.utils.torch_utils import make_tensor_with_pad

21 from vllm.v1.attention.backend import AttentionType

@@ -840,7 +841,7 @@ def torch_experts(

841 per_act_token_quant=False,

842 block_shape: list[int] | None = None,

843 apply_router_weights_on_input: bool = False,

843 - activation: str = "silu_and_mul",

844 + activation: MoEActivation = MoEActivation.SILU,

845 ) -> torch.Tensor:

846 assert (

847 global_num_experts == -1

@@ -883,7 +884,7 @@ def torch_experts(

884

885 f32 = torch.float32

886

886 - act = op_registry[activation]

887 + act = op_registry[activation.custom_op_name]

888

889 for i in range(num_experts):

890 mask = topk_ids == i

@@ -973,7 +974,7 @@ def torch_moe(

974 b_bias2: torch.Tensor | None = None,

975 global_num_experts: int = -1,

976 expert_map: torch.Tensor | None = None,

976 - activation: str = "silu_and_mul",

977 + activation: MoEActivation = MoEActivation.SILU,

978 ) -> torch.Tensor:

979 score = torch.softmax(score, dim=-1, dtype=torch.float32)

980 topk_weight, topk_ids = torch.topk(score, topk)

vllm/model_executor/layers/fused_moe/__init__.py

+7 -1

@@ -4,6 +4,11 @@

4 from contextlib import contextmanager

5 from typing import Any

6

7 +from vllm.model_executor.layers.fused_moe.activation import (

8 + MoEActivation,

9 + activation_without_mul,

10 + apply_moe_activation,

11 +)

12 from vllm.model_executor.layers.fused_moe.config import (

13 FusedMoEConfig,

14 RoutingMethodType,

@@ -27,7 +32,6 @@ from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

32 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (

33 UnquantizedFusedMoEMethod,

34 )

30 -from vllm.model_executor.layers.fused_moe.utils import activation_without_mul

35 from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import (

36 ZeroExpertFusedMoE,

37 )

@@ -54,6 +58,7 @@ __all__ = [

58 "FusedMoERouter",

59 "FusedMoEConfig",

60 "FusedMoEMethodBase",

61 + "MoEActivation",

62 "UnquantizedFusedMoEMethod",

63 "FusedMoeWeightScaleSupported",

64 "FusedMoEPermuteExpertsUnpermute",

@@ -63,6 +68,7 @@ __all__ = [

68 "SharedFusedMoE",

69 "ZeroExpertFusedMoE",

70 "activation_without_mul",

71 + "apply_moe_activation",

72 "override_config",

73 "get_config",

74 ]

vllm/model_executor/layers/fused_moe/activation.py new

+136

@@ -0,0 +1,136 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""MoE activation function enum and utilities."""

4 +

5 +from enum import Enum

6 +

7 +import torch

8 +import torch.nn.functional as F

9 +

10 +

11 +class MoEActivation(Enum):

12 + """Activation functions for MoE layers."""

13 +

14 + # Gated activations (gate * activation(up)) expect input of shape [..., 2*d]

15 + # and produce output of shape [..., d]

16 + SILU = "silu"

17 + GELU = "gelu"

18 + RELU2 = "relu2"

19 + SWIGLUOAI = "swigluoai"

20 + SWIGLUSTEP = "swiglustep"

21 +

22 + # Non-gated activations (no mul with gate) expect input of shape [..., d]

23 + # and produce output of shape [..., d].

24 + # NOTE: Non-gated activations require the "_no_mul" suffix to be present.

25 + SILU_NO_MUL = "silu_no_mul"

26 + GELU_NO_MUL = "gelu_no_mul"

27 + RELU2_NO_MUL = "relu2_no_mul"

28 +

29 + @property

30 + def is_gated(self) -> bool:

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

32 +

33 + Gated activations expect input tensor with 2x the output size,

34 + where the first half is the gate and second half is the up projection.

35 + """

36 + return not self.value.endswith("_no_mul")

37 +

38 + @property

39 + def custom_op_name(self) -> str:

40 + """Maps to the CustomOp name of activations

41 + in vllm/model_executor/layers/activation.py."""

42 + return _CUSTOM_OP_NAMES[self]

43 +

44 + def without_mul(self) -> "MoEActivation":

45 + """Get the non-gated variant of this activation.

46 +

47 + For activations that have a _no_mul variant, returns that variant.

48 + For activations without a _no_mul variant (or already _no_mul),

49 + returns self.

50 + """

51 + return _WITHOUT_MUL.get(self, self)

52 +

53 + @classmethod

54 + def from_str(cls, s: str) -> "MoEActivation":

55 + """Parse from string for backward compatibility."""

56 + for member in cls:

57 + if member.value == s:

58 + return member

59 + valid = [m.value for m in cls]

60 + raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}")

61 +

62 +

63 +# Module-level lookup tables used by MoEActivation functions.

64 +_CUSTOM_OP_NAMES: dict[MoEActivation, str] = {

65 + MoEActivation.SILU: "silu_and_mul",

66 + MoEActivation.GELU: "gelu_and_mul",

67 + MoEActivation.SWIGLUOAI: "swigluoai_and_mul",

68 + MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",

69 + MoEActivation.RELU2: "relu2",

70 + MoEActivation.SILU_NO_MUL: "silu_and_mul",

71 + MoEActivation.GELU_NO_MUL: "gelu_and_mul",

72 + MoEActivation.RELU2_NO_MUL: "relu2",

73 +}

74 +

75 +_WITHOUT_MUL: dict[MoEActivation, MoEActivation] = {

76 + MoEActivation.SILU: MoEActivation.SILU_NO_MUL,

77 + MoEActivation.GELU: MoEActivation.GELU_NO_MUL,

78 + MoEActivation.RELU2: MoEActivation.RELU2_NO_MUL,

79 +}

80 +

81 +

82 +def activation_without_mul(activation: str) -> str:

83 + """Get the non-gated variant of an activation function.

84 +

85 + Args:

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

87 +

88 + Returns:

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

90 + """

91 + return MoEActivation.from_str(activation).without_mul().value

92 +

93 +

94 +def apply_moe_activation(

95 + activation: MoEActivation,

96 + output: torch.Tensor,

97 + input: torch.Tensor,

98 +) -> torch.Tensor:

99 + """Apply MoE activation function."""

100 + assert input.dim() == 2, "Input must be 2D"

101 + assert output.dim() == 2, "Output must be 2D"

102 + if activation.is_gated:

103 + assert output.size(-1) * 2 == input.size(-1), (

104 + f"{activation.value} expects 2x ratio: "

105 + f"{output.size(-1) * 2} vs {input.size(-1)}"

106 + )

107 + else:

108 + assert output.size(-1) == input.size(-1), (

109 + f"{activation.value} expects equal sizes: "

110 + f"{output.size(-1)} vs {input.size(-1)}"

111 + )

112 +

113 + # Activations with gated multiplication (gate × activation(up))

114 + if activation == MoEActivation.SILU:

115 + torch.ops._C.silu_and_mul(output, input)

116 + elif activation == MoEActivation.GELU:

117 + torch.ops._C.gelu_and_mul(output, input)

118 + elif activation == MoEActivation.SWIGLUOAI:

119 + torch.ops._C.swigluoai_and_mul(output, input)

120 + elif activation == MoEActivation.SWIGLUSTEP:

121 + from vllm.model_executor.layers.activation import swiglustep_and_mul_triton

122 +

123 + swiglustep_and_mul_triton(output, input)

124 +

125 + # Activations without gated multiplication

126 + elif activation == MoEActivation.SILU_NO_MUL:

127 + output.copy_(F.silu(input))

128 + elif activation == MoEActivation.GELU_NO_MUL:

129 + output.copy_(F.gelu(input))

130 + elif activation == MoEActivation.RELU2_NO_MUL:

131 + F.relu(input, inplace=True)

132 + torch.square(input, out=output)

133 + else:

134 + raise ValueError(f"Unsupported FusedMoe activation: {activation}")

135 +

136 + return output

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+5 -4

@@ -7,6 +7,7 @@ import torch

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 from vllm.forward_context import get_forward_context, is_forward_context_available

9 from vllm.logger import init_logger

10 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.model_executor.layers.fused_moe.config import (

12 FusedMoEConfig,

13 FusedMoEParallelConfig,

@@ -303,8 +304,8 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

304 return (weight_key, activation_key) in SUPPORTED_W_A

305

306 @staticmethod

306 - def _supports_activation(activation: str) -> bool:

307 - return activation in ["silu"]

307 + def _supports_activation(activation: MoEActivation) -> bool:

308 + return activation == MoEActivation.SILU

309

310 @staticmethod

311 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -338,7 +339,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

339 global_num_experts: int,

340 local_num_experts: int,

341 expert_tokens_meta: mk.ExpertTokensMetadata | None,

341 - activation: str,

342 + activation: MoEActivation,

343 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

344 # FIXME (varun): We should be able to dispatch only from the leader

345 # DP ranks in the case of TP > 1. At the moment, all the Ranks

@@ -389,7 +390,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

390 w2: torch.Tensor,

391 topk_weights: torch.Tensor,

392 topk_ids: torch.Tensor,

392 - activation: str,

393 + activation: MoEActivation,

394 global_num_experts: int,

395 expert_map: torch.Tensor | None,

396 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/config.py

+2 -1

@@ -14,6 +14,7 @@ from vllm.distributed import (

14 get_tensor_model_parallel_rank,

15 )

16 from vllm.logger import init_logger

17 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

18 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

19 OCP_MX_DTYPES,

20 OCP_MX_Scheme,

@@ -1132,7 +1133,7 @@ class FusedMoEConfig:

1133 intermediate_size_per_partition: int

1134 num_local_experts: int

1135 num_logical_experts: int

1135 - activation: str

1136 + activation: MoEActivation

1137 device: torch.device | str

1138 routing_method: RoutingMethodType

1139 moe_parallel_config: FusedMoEParallelConfig

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

+12 -10

@@ -9,6 +9,7 @@ from torch.nn import functional as F

9 from vllm import _custom_ops as ops

10 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight

11 from vllm.model_executor.layers.activation import SiluAndMul

12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

13 from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter

14 from vllm.utils.torch_utils import direct_register_custom_op

15

@@ -36,9 +37,9 @@ def _swigluoai_forward_native(

37 # Map activation names to their native forward functions.

38 # Uses static methods or standalone functions to avoid instantiating CustomOp

39 # classes, which would call get_current_vllm_config() before config is set.

39 -_CPU_MOE_ACT_FN: dict[str, Callable[[torch.Tensor], torch.Tensor]] = {

40 - "silu": SiluAndMul.forward_native,

41 - "swigluoai": _swigluoai_forward_native,

40 +_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {

41 + MoEActivation.SILU: SiluAndMul.forward_native,

42 + MoEActivation.SWIGLUOAI: _swigluoai_forward_native,

43 }

44

45

@@ -168,9 +169,9 @@ class SGLFusedMOE:

169 routed_scaling_factor: float = 1.0,

170 e_score_correction_bias: torch.Tensor | None = None,

171 apply_router_weight_on_input: bool = False,

171 - activation: str = "silu",

172 + activation: MoEActivation = MoEActivation.SILU,

173 ) -> torch.Tensor:

173 - assert activation == "silu", f"{activation} is not supported."

174 + assert activation == MoEActivation.SILU, f"{activation} is not supported."

175 assert not apply_router_weight_on_input

176 topk_weights, topk_ids = select_experts(

177 hidden_states=x,

@@ -235,7 +236,7 @@ class CPUFusedMOE:

236 routed_scaling_factor: float = 1.0,

237 e_score_correction_bias: torch.Tensor | None = None,

238 apply_router_weight_on_input: bool = False,

238 - activation: str = "silu",

239 + activation: MoEActivation = MoEActivation.SILU,

240 ) -> torch.Tensor:

241 assert activation in _CPU_MOE_ACT_FN, f"{activation} is not supported."

242

@@ -353,7 +354,7 @@ class CPUFusedMOE:

354 input: torch.Tensor,

355 topk_weights: torch.Tensor,

356 topk_ids: torch.Tensor,

356 - activation: str,

357 + activation: MoEActivation,

358 global_num_experts: int = -1,

359 skip_weighted: bool = False,

360 ) -> torch.Tensor:

@@ -371,7 +372,7 @@ class CPUFusedMOE:

372 getattr(layer, "w2_bias", None),

373 topk_weights,

374 topk_ids,

374 - activation,

375 + activation.value,

376 self.isa,

377 skip_weighted,

378 )

@@ -383,7 +384,7 @@ class CPUFusedMOE:

384 input: torch.Tensor,

385 topk_weights: torch.Tensor,

386 topk_ids: torch.Tensor,

386 - activation: str,

387 + activation: MoEActivation,

388 global_num_experts: int = -1,

389 skip_weighted: bool = False,

390 ) -> torch.Tensor:

@@ -419,6 +420,7 @@ def cpu_fused_moe_torch(

420 global_num_experts: int = -1,

421 skip_weighted: bool = False,

422 ) -> None:

423 + act = MoEActivation.from_str(activation)

424 layer = _CPU_MOE_LAYER_CACHE[layer_id]()

425

426 # Ref code from https://github.com/sgl-project/sglang/blob/716e682721397df103f347d22da8bd46c6016dab/python/sglang/srt/layers/moe/fused_moe_native.py#L53

@@ -442,7 +444,7 @@ def cpu_fused_moe_torch(

444 tokens_for_this_expert = sorted_tokens[start_idx:end_idx]

445

446 gate_up = layer.gate_up_linear[i](tokens_for_this_expert) # type: ignore

445 - gate_up = _CPU_MOE_ACT_FN[activation](gate_up)

447 + gate_up = _CPU_MOE_ACT_FN[act](gate_up)

448 expert_out = layer.down_linear[i](gate_up) # type: ignore

449 outputs.append(expert_out)

450 start_idx = end_idx

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+31 -20

@@ -7,6 +7,10 @@ import torch

7 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 + MoEActivation,

12 + apply_moe_activation,

13 +)

14 from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEConfig,

16 FusedMoEParallelConfig,

@@ -25,7 +29,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

29 )

30 from vllm.model_executor.layers.fused_moe.utils import (

31 _resize_cache,

28 - apply_moe_activation,

32 )

33 from vllm.model_executor.layers.quantization.utils.quant_utils import (

34 QuantKey,

@@ -51,7 +54,7 @@ def run_cutlass_moe_fp8(

54 w1: torch.Tensor,

55 w2: torch.Tensor,

56 topk_ids: torch.Tensor,

54 - activation: str,

57 + activation: MoEActivation,

58 global_num_experts: int,

59 expert_map: torch.Tensor | None,

60 w1_scale: torch.Tensor | None,

@@ -73,7 +76,7 @@ def run_cutlass_moe_fp8(

76 ):

77 a1q = hidden_states

78

76 - assert not activation.endswith("_no_mul"), "Only gated activation is supported"

79 + assert activation.is_gated, "Only gated activation is supported"

80 assert w1_scale is not None

81 assert w2_scale is not None

82 assert w1.dtype == torch.float8_e4m3fn

@@ -310,8 +313,12 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

313 return (weight_key, activation_key) in SUPPORTED_W_A

314

315 @staticmethod

313 - def _supports_activation(activation: str) -> bool:

314 - return activation in ["silu", "gelu", "swigluoai"]

316 + def _supports_activation(activation: MoEActivation) -> bool:

317 + return activation in [

318 + MoEActivation.SILU,

319 + MoEActivation.GELU,

320 + MoEActivation.SWIGLUOAI,

321 + ]

322

323 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

324 # Let PrepareAndFinalize::finalize() decide the impl.

@@ -325,7 +332,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

332 w2: torch.Tensor,

333 topk_weights: torch.Tensor,

334 topk_ids: torch.Tensor,

328 - activation: str,

335 + activation: MoEActivation,

336 global_num_experts: int,

337 expert_map: torch.Tensor | None,

338 a1q_scale: torch.Tensor | None,

@@ -415,7 +422,7 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

422 global_num_experts: int,

423 local_num_experts: int,

424 expert_tokens_meta: mk.ExpertTokensMetadata | None,

418 - activation: str,

425 + activation: MoEActivation,

426 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

427 activation_out_dim = self.adjust_N_for_activation(N, activation)

428 workspace1 = (M * topk, max(N, K))

@@ -456,7 +463,7 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

463 global_num_experts: int,

464 local_num_experts: int,

465 expert_tokens_meta: mk.ExpertTokensMetadata | None,

459 - activation: str,

466 + activation: MoEActivation,

467 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

468 num_dp = self.num_dispatchers

469 assert num_dp is not None

@@ -489,7 +496,7 @@ def run_cutlass_moe_fp4(

496 w2_alphas: torch.Tensor,

497 topk_weights: torch.Tensor,

498 topk_ids: torch.Tensor,

492 - activation: str,

499 + activation: MoEActivation,

500 workspace13: torch.Tensor,

501 workspace2: torch.Tensor,

502 m: int,

@@ -612,7 +619,7 @@ def run_cutlass_moe_fp4(

619 blockscale_offsets[:-1],

620 )

621 del rep_a_fp4, rep_a_blockscale

615 - if activation == "silu":

622 + if activation == MoEActivation.SILU:

623 # Fused SiLU+Mul+NVFP4 quantization

624 # Note: c2 workspace is no longer needed since SiLU is fused with quantization.

625 # c3 reuses workspace13 after c1 is consumed.

@@ -682,8 +689,12 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

689 return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)

690

691 @staticmethod

685 - def _supports_activation(activation: str) -> bool:

686 - return activation in ["silu", "gelu", "swigluoai"]

692 + def _supports_activation(activation: MoEActivation) -> bool:

693 + return activation in [

694 + MoEActivation.SILU,

695 + MoEActivation.GELU,

696 + MoEActivation.SWIGLUOAI,

697 + ]

698

699 @staticmethod

700 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -716,7 +727,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

727 global_num_experts: int,

728 local_num_experts: int,

729 expert_tokens_meta: mk.ExpertTokensMetadata | None,

719 - activation: str,

730 + activation: MoEActivation,

731 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

732 workspace1 = (M * topk, max(2 * N, K))

733 workspace2 = (M * topk, N)

@@ -731,7 +742,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

742 w2: torch.Tensor,

743 topk_weights: torch.Tensor,

744 topk_ids: torch.Tensor,

734 - activation: str,

745 + activation: MoEActivation,

746 global_num_experts: int,

747 expert_map: torch.Tensor | None,

748 a1q_scale: torch.Tensor | None, # unused

@@ -776,7 +787,7 @@ def run_cutlass_moe_w4a8_fp8(

787 w1: torch.Tensor,

788 w2: torch.Tensor,

789 topk_ids: torch.Tensor,

779 - activation: str,

790 + activation: MoEActivation,

791 global_num_experts: int,

792 expert_map: torch.Tensor | None,

793 w1_scale: torch.Tensor | None,

@@ -970,7 +981,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

981 )

982

983 @staticmethod

973 - def _supports_activation(activation: str) -> bool:

984 + def _supports_activation(activation: MoEActivation) -> bool:

985 raise NotImplementedError(

986 "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

987 "This method should not be called."

@@ -1005,7 +1016,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

1016 global_num_experts: int,

1017 local_num_experts: int,

1018 expert_tokens_meta: mk.ExpertTokensMetadata | None,

1008 - activation: str,

1019 + activation: MoEActivation,

1020 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

1021 activation_out_dim = self.adjust_N_for_activation(N, activation)

1022 workspace1 = (M * topk, max(N, K))

@@ -1021,7 +1032,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

1032 w2: torch.Tensor,

1033 topk_weights: torch.Tensor,

1034 topk_ids: torch.Tensor,

1024 - activation: str,

1035 + activation: MoEActivation,

1036 global_num_experts: int,

1037 expert_map: torch.Tensor | None,

1038 a1q_scale: torch.Tensor | None,

@@ -1094,7 +1105,7 @@ def cutlass_moe_w4a8_fp8(

1105 s_strides2: torch.Tensor,

1106 quant_config: FusedMoEQuantConfig,

1107 moe_config: FusedMoEConfig,

1097 - activation: str = "silu",

1108 + activation: MoEActivation = MoEActivation.SILU,

1109 expert_map: torch.Tensor | None = None,

1110 apply_router_weight_on_input: bool = False,

1111 global_num_experts: int = -1,

@@ -1137,7 +1148,7 @@ def cutlass_moe_w4a8_fp8(

1148 dtype: torch.int64

1149 - per_act_token (Optional[bool]): Whether the scale is per-token or

1150 per-tensor.

1140 - - activation (str): The activation function to use.

1151 + - activation (MoEActivation): The activation function to use.

1152 - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,

1153 every Rank is responsible for a subset of experts. expert_map is a

1154 mapping from global expert-id to local expert-id. When expert_map[i]

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+7 -6

@@ -5,6 +5,7 @@ import torch

5

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 from vllm.logger import init_logger

8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEConfig,

11 FusedMoEParallelConfig,

@@ -145,8 +146,8 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

146 return (weight_key, activation_key) in SUPPORTED_W_A

147

148 @staticmethod

148 - def _supports_activation(activation: str) -> bool:

149 - return activation in ["silu", "swiglustep"]

149 + def _supports_activation(activation: MoEActivation) -> bool:

150 + return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]

151

152 @staticmethod

153 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -171,7 +172,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

172 global_num_experts: int,

173 local_num_experts: int,

174 expert_tokens_meta: mk.ExpertTokensMetadata | None,

174 - activation: str,

175 + activation: MoEActivation,

176 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

177 assert self.block_shape is not None

178 block_m = self.block_shape[0]

@@ -187,7 +188,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

188 return (workspace1, workspace2, output)

189

190 def _act_mul_quant(

190 - self, input: torch.Tensor, output: torch.Tensor, activation: str

191 + self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation

192 ) -> tuple[torch.Tensor, torch.Tensor]:

193 assert self.block_shape is not None

194 block_k = self.block_shape[1]

@@ -210,7 +211,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

211 return a2q, a2q_scale

212

213 # 2. Hopper / non‑E8M0: prefer the fused SiLU+mul+quant kernel

213 - if activation == "silu":

214 + if activation == MoEActivation.SILU:

215 use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0

216 return silu_mul_per_token_group_quant_fp8_colmajor(

217 input=input,

@@ -235,7 +236,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

236 w2: torch.Tensor,

237 topk_weights: torch.Tensor,

238 topk_ids: torch.Tensor,

238 - activation: str,

239 + activation: MoEActivation,

240 global_num_experts: int,

241 expert_map: torch.Tensor | None,

242 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/fallback.py

+4 -3

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod

6 import torch

7

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig

11 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

12

@@ -76,7 +77,7 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

77 ) and fallback_cls._supports_quant_scheme(weight_key, activation_key)

78

79 @classmethod

79 - def _supports_activation(cls, activation: str) -> bool:

80 + def _supports_activation(cls, activation: MoEActivation) -> bool:

81 experts_cls, fallback_cls = cls.get_clses()

82 return experts_cls._supports_activation(

83 activation

@@ -138,7 +139,7 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

139 global_num_experts: int,

140 local_num_experts: int,

141 expert_tokens_meta: mk.ExpertTokensMetadata | None,

141 - activation: str,

142 + activation: MoEActivation,

143 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

144 raise NotImplementedError

145

@@ -159,7 +160,7 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

160 w2: torch.Tensor,

161 topk_weights: torch.Tensor,

162 topk_ids: torch.Tensor,

162 - activation: str,

163 + activation: MoEActivation,

164 global_num_experts: int,

165 expert_map: torch.Tensor | None,

166 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py

+5 -4

@@ -6,6 +6,7 @@ import torch

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 from vllm import envs

8 from vllm.logger import init_logger

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.model_executor.layers.fused_moe.config import (

11 FusedMoEConfig,

12 FusedMoEParallelConfig,

@@ -72,8 +73,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):

73 return (weight_key, activation_key) in SUPPORTED_W_A

74

75 @staticmethod

75 - def _supports_activation(activation: str) -> bool:

76 - return activation in ["silu"]

76 + def _supports_activation(activation: MoEActivation) -> bool:

77 + return activation == MoEActivation.SILU

78

79 @staticmethod

80 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -101,7 +102,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):

102 global_num_experts: int,

103 local_num_experts: int,

104 expert_tokens_meta: mk.ExpertTokensMetadata | None,

104 - activation: str,

105 + activation: MoEActivation,

106 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

107 # We use global_num_experts due to how moe_align_block_size handles

108 # expert_maps.

@@ -135,7 +136,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):

136 w2: torch.Tensor,

137 topk_weights: torch.Tensor,

138 topk_ids: torch.Tensor,

138 - activation: str,

139 + activation: MoEActivation,

140 global_num_experts: int,

141 expert_map: torch.Tensor | None,

142 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+7 -6

@@ -5,6 +5,7 @@ import torch

5

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 from vllm.logger import init_logger

8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEParallelConfig,

11 FusedMoEQuantConfig,

@@ -130,8 +131,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

131 )

132

133 @staticmethod

133 - def _supports_activation(activation: str) -> bool:

134 - return activation in ["silu", "relu2_no_mul"]

134 + def _supports_activation(activation: MoEActivation) -> bool:

135 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

136

137 @staticmethod

138 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -164,7 +165,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

165 global_num_experts: int,

166 local_num_experts: int,

167 expert_tokens_meta: mk.ExpertTokensMetadata | None,

167 - activation: str,

168 + activation: MoEActivation,

169 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

170 # We use global_num_experts due to how moe_align_block_size handles

171 # expert_maps.

@@ -201,7 +202,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

202 w2: torch.Tensor,

203 topk_weights: torch.Tensor,

204 topk_ids: torch.Tensor,

204 - activation: str,

205 + activation: MoEActivation,

206 global_num_experts: int,

207 expert_map: torch.Tensor | None,

208 a1q_scale: torch.Tensor | None,

@@ -214,8 +215,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

215 from flashinfer.fused_moe.core import ActivationType

216

217 activation_str_to_value_map = {

217 - "silu": ActivationType.Swiglu, # This is the default

218 - "relu2_no_mul": ActivationType.Relu2,

218 + MoEActivation.SILU: ActivationType.Swiglu, # This is the default

219 + MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,

220 }

221 assert activation in activation_str_to_value_map, (

222 f"{activation=} missing from {activation_str_to_value_map.keys()=}"

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

+3 -2

@@ -4,6 +4,7 @@

4 import torch

5

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

8 from vllm.model_executor.layers.fused_moe.config import (

9 FusedMoEConfig,

10 FusedMoEParallelConfig,

@@ -50,9 +51,9 @@ def _supports_quant_scheme(

51 return (weight_key, activation_key) in SUPPORTED_W_A

52

53

53 -def _supports_activation(activation: str) -> bool:

54 +def _supports_activation(activation: MoEActivation) -> bool:

55 """Supports silu activation only."""

55 - return activation in ["silu"]

56 + return activation == MoEActivation.SILU

57

58

59 def _supports_routing_method(

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+13 -12

@@ -5,6 +5,7 @@

5 import torch

6

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEConfig,

11 FusedMoEParallelConfig,

@@ -698,7 +699,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

699 )

700

701 @staticmethod

701 - def _supports_activation(activation: str) -> bool:

702 + def _supports_activation(activation: MoEActivation) -> bool:

703 raise NotImplementedError(

704 "NaiveBatchedExperts is not yet used by an Oracle. "

705 "This method should not be called."

@@ -730,7 +731,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

731 global_num_experts: int,

732 local_num_experts: int,

733 expert_tokens_meta: mk.ExpertTokensMetadata | None,

733 - activation: str,

734 + activation: MoEActivation,

735 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

736 assert self.num_dispatchers is not None

737 assert self.max_num_tokens is not None

@@ -757,7 +758,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

758 w2: torch.Tensor,

759 topk_weights: torch.Tensor,

760 topk_ids: torch.Tensor,

760 - activation: str,

761 + activation: MoEActivation,

762 global_num_experts: int,

763 expert_map: torch.Tensor | None,

764 a1q_scale: torch.Tensor | None,

@@ -942,14 +943,14 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

943 )

944

945 @staticmethod

945 - def _supports_activation(activation: str) -> bool:

946 + def _supports_activation(activation: MoEActivation) -> bool:

947 return activation in [

947 - "silu",

948 - "gelu",

949 - "swigluoai",

950 - "silu_no_mul",

951 - "gelu_no_mul",

952 - "relu2_no_mul",

948 + MoEActivation.SILU,

949 + MoEActivation.GELU,

950 + MoEActivation.SWIGLUOAI,

951 + MoEActivation.SILU_NO_MUL,

952 + MoEActivation.GELU_NO_MUL,

953 + MoEActivation.RELU2_NO_MUL,

954 ]

955

956 @staticmethod

@@ -975,7 +976,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

976 global_num_experts: int,

977 local_num_experts: int,

978 expert_tokens_meta: mk.ExpertTokensMetadata | None,

978 - activation: str,

979 + activation: MoEActivation,

980 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

981 assert self.num_dispatchers is not None

982 assert self.max_num_tokens is not None

@@ -996,7 +997,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

997 w2: torch.Tensor,

998 topk_weights: torch.Tensor,

999 topk_ids: torch.Tensor,

999 - activation: str,

1000 + activation: MoEActivation,

1001 global_num_experts: int,

1002 expert_map: torch.Tensor | None,

1003 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+21 -18

@@ -8,6 +8,10 @@ import torch

8

9 import vllm._custom_ops as ops

10 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

11 +from vllm.model_executor.layers.fused_moe.activation import (

12 + MoEActivation,

13 + apply_moe_activation,

14 +)

15 from vllm.model_executor.layers.fused_moe.config import (

16 FusedMoEConfig,

17 FusedMoEParallelConfig,

@@ -23,7 +27,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

27 )

28 from vllm.model_executor.layers.fused_moe.utils import (

29 _resize_cache,

26 - apply_moe_activation,

30 disable_inplace,

31 )

32 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

@@ -59,9 +62,9 @@ def _fused_marlin_moe(

62 sorted_token_ids: torch.Tensor,

63 expert_ids: torch.Tensor,

64 num_tokens_post_padded: torch.Tensor,

62 - activation: str = "silu",

65 + activation: MoEActivation = MoEActivation.SILU,

66 activation_func: Callable[

64 - [str, torch.Tensor, torch.Tensor], None

67 + [MoEActivation, torch.Tensor, torch.Tensor], None

68 ] = apply_moe_activation,

69 input_global_scale1: torch.Tensor | None = None,

70 input_global_scale2: torch.Tensor | None = None,

@@ -83,7 +86,7 @@ def _fused_marlin_moe(

86 assert hidden_states.ndim == 2

87 M, K = hidden_states.size()

88 N = marlin_moe_intermediate_size(w1, w2)

86 - w13_num_shards = 1 if "no_mul" in activation else 2

89 + w13_num_shards = 2 if activation.is_gated else 1

90 if workspace is None:

91 workspace = marlin_make_workspace_new(hidden_states.device, 4)

92

@@ -215,9 +218,9 @@ def fused_marlin_moe(

218 quant_type_id: int,

219 apply_router_weight_on_input: bool = False,

220 global_num_experts: int = -1,

218 - activation: str = "silu",

221 + activation: MoEActivation = MoEActivation.SILU,

222 activation_func: Callable[

220 - [str, torch.Tensor, torch.Tensor], None

223 + [MoEActivation, torch.Tensor, torch.Tensor], None

224 ] = apply_moe_activation,

225 moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,

226 expert_map: torch.Tensor | None = None,

@@ -377,7 +380,7 @@ def batched_fused_marlin_moe(

380 quant_type_id: int,

381 apply_router_weight_on_input: bool = False,

382 global_num_experts: int = -1,

380 - activation: str | None = "silu",

383 + activation: MoEActivation = MoEActivation.SILU,

384 expert_map: torch.Tensor | None = None,

385 global_scale1: torch.Tensor | None = None,

386 global_scale2: torch.Tensor | None = None,

@@ -579,14 +582,14 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

582 return weight_key in SUPPORTED_W

583

584 @staticmethod

582 - def _supports_activation(activation: str) -> bool:

585 + def _supports_activation(activation: MoEActivation) -> bool:

586 return activation in [

584 - "silu",

585 - "gelu",

586 - "swigluoai",

587 - "silu_no_mul",

588 - "gelu_no_mul",

589 - "relu2_no_mul",

587 + MoEActivation.SILU,

588 + MoEActivation.GELU,

589 + MoEActivation.SWIGLUOAI,

590 + MoEActivation.SILU_NO_MUL,

591 + MoEActivation.GELU_NO_MUL,

592 + MoEActivation.RELU2_NO_MUL,

593 ]

594

595 @staticmethod

@@ -661,7 +664,7 @@ class MarlinExperts(MarlinExpertsBase):

664 global_num_experts: int,

665 local_num_experts: int,

666 expert_tokens_meta: mk.ExpertTokensMetadata | None,

664 - activation: str,

667 + activation: MoEActivation,

668 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

669 # Modular Kernel provisions output buffer from workspace1. However in

670 # the fused_marlin_moe() function, the final torch.sum(), is defined

@@ -692,7 +695,7 @@ class MarlinExperts(MarlinExpertsBase):

695 w2: torch.Tensor,

696 topk_weights: torch.Tensor,

697 topk_ids: torch.Tensor,

695 - activation: str,

698 + activation: MoEActivation,

699 global_num_experts: int,

700 expert_map: torch.Tensor | None,

701 a1q_scale: torch.Tensor | None,

@@ -788,7 +791,7 @@ class BatchedMarlinExperts(MarlinExpertsBase):

791 global_num_experts: int,

792 local_num_experts: int,

793 expert_tokens_meta: mk.ExpertTokensMetadata | None,

791 - activation: str,

794 + activation: MoEActivation,

795 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

796 assert self.num_dispatchers is not None

797 assert self.max_num_tokens is not None

@@ -808,7 +811,7 @@ class BatchedMarlinExperts(MarlinExpertsBase):

811 w2: torch.Tensor,

812 topk_weights: torch.Tensor,

813 topk_ids: torch.Tensor,

811 - activation: str,

814 + activation: MoEActivation,

815 global_num_experts: int,

816 expert_map: torch.Tensor | None,

817 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/fused_moe.py

+23 -11

@@ -17,6 +17,10 @@ from vllm.logger import init_logger

17 from vllm.model_executor.layers.batch_invariant import (

18 vllm_is_batch_invariant,

19 )

20 +from vllm.model_executor.layers.fused_moe.activation import (

21 + MoEActivation,

22 + apply_moe_activation,

23 +)

24 from vllm.model_executor.layers.fused_moe.config import (

25 FUSED_MOE_UNQUANTIZED_CONFIG,

26 FusedMoEConfig,

@@ -32,7 +36,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

36 )

37 from vllm.model_executor.layers.fused_moe.utils import (

38 _resize_cache,

35 - apply_moe_activation,

39 disable_inplace,

40 moe_kernel_quantize_input,

41 )

@@ -1468,6 +1471,7 @@ def outplace_fused_experts_fake(

1471 topk_weights: torch.Tensor,

1472 topk_ids: torch.Tensor,

1473 activation: str = "silu",

1474 + apply_router_weight_on_input: bool = False,

1475 use_fp8_w8a8: bool = False,

1476 use_int8_w8a8: bool = False,

1477 use_int8_w8a16: bool = False,

@@ -1521,7 +1525,7 @@ def fused_experts(

1525 topk_weights: torch.Tensor,

1526 topk_ids: torch.Tensor,

1527 inplace: bool = False,

1524 - activation: str = "silu",

1528 + activation: MoEActivation = MoEActivation.SILU,

1529 apply_router_weight_on_input: bool = False,

1530 global_num_experts: int = -1,

1531 expert_map: torch.Tensor | None = None,

@@ -1539,7 +1543,7 @@ def fused_experts(

1543 w2=w2,

1544 topk_weights=topk_weights,

1545 topk_ids=topk_ids,

1542 - activation=activation,

1546 + activation=activation.value,

1547 apply_router_weight_on_input=apply_router_weight_on_input,

1548 use_fp8_w8a8=quant_config.use_fp8_w8a8,

1549 use_int8_w8a8=quant_config.use_int8_w8a8,

@@ -1618,6 +1622,9 @@ def fused_experts_impl(

1622 w1_bias: torch.Tensor | None = None,

1623 w2_bias: torch.Tensor | None = None,

1624 ) -> torch.Tensor:

1625 + # Convert string activation to enum for internal use

1626 + activation_enum = MoEActivation.from_str(activation)

1627 +

1628 # Check constraints.

1629 if use_int4_w4a16:

1630 assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"

@@ -1692,7 +1699,7 @@ def fused_experts_impl(

1699

1700 # This needs separate memory since it's used concurrently with cache1

1701 activation_out_dim = mk.FusedMoEPermuteExpertsUnpermute.adjust_N_for_activation(

1695 - N, activation

1702 + N, activation_enum

1703 )

1704 intermediate_cache2 = torch.empty(

1705 (M * top_k_num, activation_out_dim),

@@ -1832,7 +1839,7 @@ def fused_experts_impl(

1839 )

1840

1841 apply_moe_activation(

1835 - activation, intermediate_cache2, intermediate_cache1.view(-1, N)

1842 + activation_enum, intermediate_cache2, intermediate_cache1.view(-1, N)

1843 )

1844

1845 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(

@@ -1932,8 +1939,13 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1939 return (weight_key, activation_key) in SUPPORTED_W_A

1940

1941 @staticmethod

1935 - def _supports_activation(activation: str) -> bool:

1936 - return activation in ["silu", "gelu", "swigluoai", "swiglustep"]

1942 + def _supports_activation(activation: MoEActivation) -> bool:

1943 + return activation in [

1944 + MoEActivation.SILU,

1945 + MoEActivation.GELU,

1946 + MoEActivation.SWIGLUOAI,

1947 + MoEActivation.SWIGLUSTEP,

1948 + ]

1949

1950 @staticmethod

1951 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -1957,7 +1969,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1969 global_num_experts: int,

1970 local_num_experts: int,

1971 expert_tokens_meta: mk.ExpertTokensMetadata | None,

1960 - activation: str,

1972 + activation: MoEActivation,

1973 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

1974 activation_out_dim = self.adjust_N_for_activation(N, activation)

1975 workspace1 = (M, topk, max(activation_out_dim, K))

@@ -1973,7 +1985,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1985 w2: torch.Tensor,

1986 topk_weights: torch.Tensor,

1987 topk_ids: torch.Tensor,

1976 - activation: str,

1988 + activation: MoEActivation,

1989 global_num_experts: int,

1990 expert_map: torch.Tensor | None,

1991 a1q_scale: torch.Tensor | None,

@@ -2138,7 +2150,7 @@ class TritonWNA16Experts(TritonExperts):

2150 )

2151

2152 @staticmethod

2141 - def _supports_activation(activation: str) -> bool:

2153 + def _supports_activation(activation: MoEActivation) -> bool:

2154 raise NotImplementedError(

2155 "TritonWNA16Experts is not yet used by an Oracle. "

2156 "This method should not be called."

@@ -2159,7 +2171,7 @@ class TritonWNA16Experts(TritonExperts):

2171 w2: torch.Tensor,

2172 topk_weights: torch.Tensor,

2173 topk_ids: torch.Tensor,

2162 - activation: str,

2174 + activation: MoEActivation,

2175 global_num_experts: int,

2176 expert_map: torch.Tensor | None,

2177 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+11 -7

@@ -7,6 +7,7 @@ import torch

7 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 MoEActivation

11 from vllm.model_executor.layers.fused_moe.config import (

12 FUSED_MOE_UNQUANTIZED_CONFIG,

13 FusedMoEParallelConfig,

@@ -172,7 +173,7 @@ def triton_kernel_moe_forward(

173 gating_output: torch.Tensor,

174 topk: int,

175 renormalize: bool,

175 - activation: str = "silu",

176 + activation: MoEActivation = MoEActivation.SWIGLUOAI,

177 quant_config: FusedMoEQuantConfig | None = None,

178 apply_router_weight_on_input: bool = False,

179 global_num_experts: int = -1,

@@ -211,7 +212,7 @@ def triton_kernel_fused_experts(

212 gather_indx, # GatherIndx

213 scatter_indx, # ScatterIndx

214 topk: int,

214 - activation: str = "silu",

215 + activation: MoEActivation = MoEActivation.SWIGLUOAI,

216 quant_config: FusedMoEQuantConfig | None = None,

217 swiglu_alpha: float = 1.702,

218 swiglu_limit: float = 7.0,

@@ -222,6 +223,9 @@ def triton_kernel_fused_experts(

223 a1q_scale: torch.Tensor | None = None,

224 ) -> torch.Tensor:

225 """Triton implementation of fused expert computation using OAI kernels."""

226 + assert activation == MoEActivation.SWIGLUOAI, (

227 + "Only SWIGLUOAI activation is supported"

228 + )

229 if quant_config is None:

230 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

231

@@ -379,7 +383,7 @@ class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

383 )

384

385 @staticmethod

382 - def _supports_activation(activation: str) -> bool:

386 + def _supports_activation(activation: MoEActivation) -> bool:

387 raise NotImplementedError(

388 "OAITritonExperts is not yet used by an Oracle. "

389 "This method should not be called."

@@ -463,7 +467,7 @@ class OAITritonExperts(BaseOAITritonExperts):

467 global_num_experts: int,

468 local_num_experts: int,

469 expert_tokens_meta: mk.ExpertTokensMetadata | None,

466 - activation: str,

470 + activation: MoEActivation,

471 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

472 # workspace are allocated inside the kernel

473 activation_out_dim = self.adjust_N_for_activation(N, activation)

@@ -480,7 +484,7 @@ class OAITritonExperts(BaseOAITritonExperts):

484 w2: torch.Tensor,

485 topk_weights: torch.Tensor,

486 topk_ids: torch.Tensor,

483 - activation: str,

487 + activation: MoEActivation,

488 global_num_experts: int,

489 expert_map: torch.Tensor | None,

490 a1q_scale: torch.Tensor | None,

@@ -547,7 +551,7 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):

551 global_num_experts: int,

552 local_num_experts: int,

553 expert_tokens_meta: mk.ExpertTokensMetadata | None,

550 - activation: str,

554 + activation: MoEActivation,

555 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

556 # workspace are allocated inside the kernel

557 activation_out_dim = self.adjust_N_for_activation(N, activation)

@@ -567,7 +571,7 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):

571 w2: torch.Tensor,

572 topk_weights: torch.Tensor,

573 topk_ids: torch.Tensor,

570 - activation: str,

574 + activation: MoEActivation,

575 global_num_experts: int,

576 expert_map: torch.Tensor | None,

577 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/layer.py

+3 -2

@@ -20,6 +20,7 @@ from vllm.distributed import (

20 from vllm.distributed.eplb.eplb_state import EplbLayerState, EplbState

21 from vllm.logger import init_logger

22 from vllm.model_executor.custom_op import CustomOp

23 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

24 from vllm.model_executor.layers.fused_moe.config import (

25 FusedMoEConfig,

26 FusedMoEParallelConfig,

@@ -500,7 +501,7 @@ class FusedMoE(CustomOp):

501 # TODO(bnell): end attributes

502

503 self.apply_router_weight_on_input = apply_router_weight_on_input

503 - self.activation = activation

504 + self.activation = MoEActivation.from_str(activation)

505

506 self.router = create_fused_moe_router(

507 top_k=top_k,

@@ -554,7 +555,7 @@ class FusedMoE(CustomOp):

555 has_bias=has_bias,

556 is_act_and_mul=is_act_and_mul,

557 is_lora_enabled=vllm_config.lora_config is not None,

557 - activation=activation,

558 + activation=self.activation,

559 device=vllm_config.device_config.device,

560 routing_method=self.routing_method_type,

561 # TODO: in_dtype == out_dtype?

vllm/model_executor/layers/fused_moe/modular_kernel.py

+15 -13

@@ -12,6 +12,10 @@ import torch

12 import vllm.envs as envs

13 from vllm.forward_context import get_forward_context, is_forward_context_available

14 from vllm.logger import init_logger

15 +from vllm.model_executor.layers.fused_moe.activation import (

16 + MoEActivation,

17 + apply_moe_activation,

18 +)

19 from vllm.model_executor.layers.fused_moe.config import (

20 FusedMoEConfig,

21 FusedMoEParallelConfig,

@@ -19,7 +23,6 @@ from vllm.model_executor.layers.fused_moe.config import (

23 )

24 from vllm.model_executor.layers.fused_moe.utils import (

25 _resize_cache,

22 - apply_moe_activation,

26 count_expert_num_tokens,

27 disable_inplace,

28 )

@@ -536,7 +539,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

539

540 @staticmethod

541 @abstractmethod

539 - def _supports_activation(activation: str) -> bool:

542 + def _supports_activation(activation: MoEActivation) -> bool:

543 """

544 Whether the kernel supports a particular act function.

545 """

@@ -658,7 +661,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

661 global_num_experts: int,

662 local_num_experts: int,

663 expert_tokens_meta: ExpertTokensMetadata | None,

661 - activation: str,

664 + activation: MoEActivation,

665 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

666 """

667 Compute the shapes for the temporary and final outputs of the two gemms

@@ -690,7 +693,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

693 raise NotImplementedError

694

695 @staticmethod

693 - def adjust_N_for_activation(N: int, activation: str) -> int:

696 + def adjust_N_for_activation(N: int, activation: MoEActivation) -> int:

697 """

698 Calculate the output dimension for the activation function.

699

@@ -702,16 +705,15 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

705

706 Args:

707 N: The intermediate size (width of w1/w3 weights).

705 - activation: The activation function name.

708 + activation: The activation function enum.

709

710 Returns:

711 The output dimension after activation.

712 """

710 - is_no_mul = activation.endswith("_no_mul")

711 - return N if is_no_mul else N // 2

713 + return N if not activation.is_gated else N // 2

714

715 def activation(

714 - self, activation: str, output: torch.Tensor, input: torch.Tensor

716 + self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor

717 ) -> None:

718 apply_moe_activation(activation, output, input)

719

@@ -732,7 +734,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

734 w2: torch.Tensor,

735 topk_weights: torch.Tensor,

736 topk_ids: torch.Tensor,

735 - activation: str,

737 + activation: MoEActivation,

738 global_num_experts: int,

739 expert_map: torch.Tensor | None,

740 a1q_scale: torch.Tensor | None,

@@ -892,7 +894,7 @@ class FusedMoEModularKernel(torch.nn.Module):

894 global_num_experts: int,

895 local_num_experts: int,

896 expert_tokens_meta: ExpertTokensMetadata | None,

895 - activation: str,

897 + activation: MoEActivation,

898 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

899 """

900 Allocate temporary and output buffers for the fused experts op.

@@ -1135,7 +1137,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1137 w2: torch.Tensor,

1138 topk_weights: torch.Tensor,

1139 topk_ids: torch.Tensor,

1138 - activation: str,

1140 + activation: MoEActivation,

1141 global_num_experts: int,

1142 local_num_experts: int,

1143 expert_map: torch.Tensor | None,

@@ -1309,7 +1311,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1311 w2: torch.Tensor,

1312 topk_weights: torch.Tensor,

1313 topk_ids: torch.Tensor,

1312 - activation: str = "silu",

1314 + activation: MoEActivation = MoEActivation.SILU,

1315 global_num_experts: int = -1,

1316 expert_map: torch.Tensor | None = None,

1317 apply_router_weight_on_input: bool = False,

@@ -1326,7 +1328,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1328 - topk_weights (torch.Tensor): The topk weights applied at the end of

1329 the layer.

1330 - topk_ids (torch.Tensor): A map of row to expert id.

1329 - - activation (str): The activation function to apply after the first

1331 + - activation (MoEActivation): The activation function to apply after the first

1332 MoE layer.

1333 - global_num_experts (int): The total number of experts in the global

1334 expert space.

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+13 -8

@@ -7,6 +7,7 @@ import torch

7

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm._aiter_ops import rocm_aiter_ops

10 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.model_executor.layers.fused_moe.config import (

12 FUSED_MOE_UNQUANTIZED_CONFIG,

13 FusedMoEParallelConfig,

@@ -184,7 +185,7 @@ def rocm_aiter_fused_experts(

185 w2: torch.Tensor,

186 topk_weights: torch.Tensor,

187 topk_ids: torch.Tensor,

187 - activation: str = "silu",

188 + activation: MoEActivation = MoEActivation.SILU,

189 apply_router_weight_on_input: bool = False,

190 expert_map: torch.Tensor | None = None,

191 quant_config: FusedMoEQuantConfig | None = None,

@@ -196,9 +197,13 @@ def rocm_aiter_fused_experts(

197 if quant_config is None:

198 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

199

199 - activation_method = (

200 - ActivationMethod.SILU if activation == "silu" else ActivationMethod.GELU

201 - )

200 + if activation == MoEActivation.SILU:

201 + activation_method = ActivationMethod.SILU

202 + elif activation == MoEActivation.GELU:

203 + activation_method = ActivationMethod.GELU

204 + else:

205 + raise ValueError(f"Unsupported activation: {activation}")

206 +

207 # All AITER Fused MoE kernels are expecting the following datatypes

208 topk_weights = topk_weights.to(torch.float32)

209 topk_ids = topk_ids.to(torch.int32)

@@ -322,8 +327,8 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

327 return (weight_key, activation_key) in SUPPORTED_W_A

328

329 @staticmethod

325 - def _supports_activation(activation: str) -> bool:

326 - return activation in ["silu", "gelu"]

330 + def _supports_activation(activation: MoEActivation) -> bool:

331 + return activation in [MoEActivation.SILU, MoEActivation.GELU]

332

333 @staticmethod

334 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -347,7 +352,7 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

352 global_num_experts: int,

353 local_num_experts: int,

354 expert_tokens_meta: mk.ExpertTokensMetadata | None,

350 - activation: str,

355 + activation: MoEActivation,

356 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

357 # Workspaces are managed internally by AITER.

358 workspace1 = (0,)

@@ -363,7 +368,7 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

368 w2: torch.Tensor,

369 topk_weights: torch.Tensor,

370 topk_ids: torch.Tensor,

366 - activation: str,

371 + activation: MoEActivation,

372 global_num_experts: int,

373 expert_map: torch.Tensor | None,

374 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py

+2 -1

@@ -5,6 +5,7 @@

5 import torch

6

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEConfig,

11 FusedMoEQuantConfig,

@@ -45,7 +46,7 @@ class TritonOrCutlassExperts(FallbackExperts):

46 global_num_experts: int,

47 local_num_experts: int,

48 expert_tokens_meta: mk.ExpertTokensMetadata | None,

48 - activation: str,

49 + activation: MoEActivation,

50 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

51 # Small batch fallback for sm100.

52 if self.is_sm100 and M <= 8:

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+2 -1

@@ -4,6 +4,7 @@

4 import torch

5

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

8 from vllm.model_executor.layers.fused_moe.config import (

9 FusedMoEConfig,

10 FusedMoEQuantConfig,

@@ -45,7 +46,7 @@ class TritonOrDeepGemmExperts(FallbackExperts):

46 global_num_experts: int,

47 local_num_experts: int,

48 expert_tokens_meta: mk.ExpertTokensMetadata | None,

48 - activation: str,

49 + activation: MoEActivation,

50 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

51 # Note: the deep gemm workspaces are strictly larger than the triton

52 # workspaces so we can be pessimistic here and allocate for DeepGemm

vllm/model_executor/layers/fused_moe/trtllm_moe.py

+4 -3

@@ -4,6 +4,7 @@

4 import torch

5

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

8 from vllm.model_executor.layers.fused_moe.config import (

9 FusedMoEConfig,

10 FusedMoEParallelConfig,

@@ -64,7 +65,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

65 )

66

67 @staticmethod

67 - def _supports_activation(activation: str) -> bool:

68 + def _supports_activation(activation: MoEActivation) -> bool:

69 raise NotImplementedError(

70 "TrtLlmGenExperts is not yet used by an Oracle. "

71 "This method should not be called."

@@ -95,7 +96,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

96 global_num_experts: int,

97 local_num_experts: int,

98 expert_tokens_meta: mk.ExpertTokensMetadata | None,

98 - activation: str,

99 + activation: MoEActivation,

100 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

101 # The workspaces for this implementation are managed by flashinfer.

102 workspace1 = (0,)

@@ -111,7 +112,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

112 w2: torch.Tensor,

113 topk_weights: torch.Tensor,

114 topk_ids: torch.Tensor,

114 - activation: str,

115 + activation: MoEActivation,

116 global_num_experts: int,

117 expert_map: torch.Tensor | None,

118 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/utils.py

-60

@@ -4,7 +4,6 @@ import functools

4 from math import prod

5

6 import torch

7 -import torch.nn.functional as F

7

8 from vllm import _custom_ops as ops

9 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

@@ -341,65 +340,6 @@ def _validate_scale_shape(

340 assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"

341

342

344 -def activation_without_mul(activation: str) -> str:

345 - return activation + "_no_mul"

346 -

347 -

348 -RELU2_NO_MUL: str = activation_without_mul("relu2")

349 -SILU_NO_MUL: str = activation_without_mul("silu")

350 -GELU_NO_MUL: str = activation_without_mul("gelu")

351 -

352 -

353 -def apply_moe_activation(

354 - activation: str,

355 - output: torch.Tensor,

356 - input: torch.Tensor,

357 -) -> torch.Tensor:

358 - """

359 - Apply MoE activation function.

360 -

361 - For *_and_mul activations (silu, gelu, swigluoai):

362 - - Expects output.size(-1) * 2 == input.size(-1)

363 -

364 - For *_no_mul activations (silu_no_mul, gelu_no_mul, relu2_no_mul):

365 - - Expects output.size(-1) == input.size(-1)

366 - """

367 - is_no_mul = activation.endswith("_no_mul")

368 - if is_no_mul:

369 - assert output.size(-1) == input.size(-1), (

370 - f"{activation} expects equal sizes: {output.size(-1)} vs {input.size(-1)}"

371 - )

372 - else:

373 - assert output.size(-1) * 2 == input.size(-1), (

374 - f"{activation} expects 2x ratio: {output.size(-1) * 2} vs {input.size(-1)}"

375 - )

376 -

377 - # Activations with gated multiplication (gate × activation(up))

378 - if activation == "silu":

379 - torch.ops._C.silu_and_mul(output, input)

380 - elif activation == "gelu":

381 - torch.ops._C.gelu_and_mul(output, input)

382 - elif activation == "swigluoai":

383 - torch.ops._C.swigluoai_and_mul(output, input)

384 - elif activation == "swiglustep":

385 - from vllm.model_executor.layers.activation import swiglustep_and_mul_triton

386 -

387 - swiglustep_and_mul_triton(output, input)

388 -

389 - # Activations without gated multiplication

390 - elif activation == SILU_NO_MUL:

391 - output.copy_(F.silu(input))

392 - elif activation == GELU_NO_MUL:

393 - output.copy_(F.gelu(input))

394 - elif activation == RELU2_NO_MUL:

395 - F.relu(input, inplace=True)

396 - torch.square(input, out=output)

397 - else:

398 - raise ValueError(f"Unsupported FusedMoe activation: {activation}")

399 -

400 - return output

401 -

402 -

343 # Torch custom ops can't deal with outputs aliasing inputs so we need to

344 # disable inplace for torch >= 2.9.

345 # See https://github.com/vllm-project/vllm/issues/26378

vllm/model_executor/layers/fused_moe/xpu_fused_moe.py

+10 -5

@@ -3,6 +3,7 @@

3 import torch

4

5 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

6 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

7 from vllm.model_executor.layers.fused_moe.config import (

8 FusedMoEConfig,

9 FusedMoEParallelConfig,

@@ -55,8 +56,12 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):

56 return False

57

58 @staticmethod

58 - def _supports_activation(activation: str) -> bool:

59 - return activation in ["silu", "gelu", "swigluoai"]

59 + def _supports_activation(activation: MoEActivation) -> bool:

60 + return activation in [

61 + MoEActivation.SILU,

62 + MoEActivation.GELU,

63 + MoEActivation.SWIGLUOAI,

64 + ]

65

66 @staticmethod

67 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

@@ -92,7 +97,7 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):

97 global_num_experts: int,

98 local_num_experts: int,

99 expert_tokens_meta: mk.ExpertTokensMetadata | None,

95 - activation: str,

100 + activation: MoEActivation,

101 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

102 workspace1 = (0,)

103 workspace2 = (0,)

@@ -107,7 +112,7 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):

112 w2: torch.Tensor,

113 topk_weights: torch.Tensor,

114 topk_ids: torch.Tensor,

110 - activation: str,

115 + activation: MoEActivation,

116 global_num_experts: int,

117 expert_map: torch.Tensor | None,

118 a1q_scale: torch.Tensor | None,

@@ -129,7 +134,7 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):

134 topk_weights=topk_weights,

135 topk_ids=topk_ids,

136 n_experts_per_token=topk,

132 - activation=activation,

137 + activation=activation.value,

138 num_experts=self.moe_config.num_local_experts,

139 ep_rank=self.moe_config.ep_rank,

140 ep_size=self.moe_config.ep_size,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+19 -10

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe import (

24 FusedMoeWeightScaleSupported,

25 UnquantizedFusedMoEMethod,

26 )

27 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

28 from vllm.model_executor.layers.fused_moe.config import (

29 FusedMoEConfig,

30 FusedMoEQuantConfig,

@@ -622,7 +623,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

623 router_logits: torch.Tensor,

624 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

625 assert self.is_monolithic

625 - assert layer.activation == "silu", "Only SiLU activation is supported."

626 + assert layer.activation == MoEActivation.SILU, (

627 + f"Only SiLU activation is supported, not {layer.activation}."

628 + )

629 assert (

630 self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

631 and not layer.enable_eplb

@@ -649,7 +652,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

652 shared_experts_input: torch.Tensor | None,

653 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

654 assert not self.is_monolithic

652 - assert layer.activation == "silu", "Only SiLU activation is supported."

655 + assert layer.activation == MoEActivation.SILU, (

656 + f"Only SiLU activation is supported, not {layer.activation}."

657 + )

658

659 # EPLB path

660 if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

@@ -1025,7 +1030,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1030 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1031 assert self.is_monolithic

1032 assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

1028 - assert layer.activation == "silu"

1033 + assert layer.activation == MoEActivation.SILU, (

1034 + f"Only SiLU activation is supported, not {layer.activation}."

1035 + )

1036

1037 if self.block_quant:

1038 import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

@@ -2271,19 +2278,21 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

2278 router_logits: torch.Tensor,

2279 ) -> torch.Tensor:

2280 assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."

2274 - assert layer.activation in ("silu", "swigluoai", "swiglu"), (

2275 - "Only SiLU/SwiGLUGU/SwiGLUUG are supported."

2276 - )

2281 + assert layer.activation in (

2282 + MoEActivation.SILU,

2283 + MoEActivation.SWIGLUOAI,

2284 + MoEActivation.SWIGLUSTEP,

2285 + ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."

2286 assert layer.expert_map is None, """expert_map/EP not implemented

2287 for CPU dyn-4bit MoE."""

2288

2280 - def _act_kind(s: str) -> int:

2289 + def _act_kind(s: MoEActivation) -> int:

2290 # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU

2282 - if s == "swiglu":

2291 + if s == MoEActivation.SWIGLUSTEP:

2292 return 0

2284 - if s == "swigluoai":

2293 + if s == MoEActivation.SWIGLUOAI:

2294 return 1

2286 - if s == "silu":

2295 + if s == MoEActivation.SILU:

2296 return 2

2297 raise ValueError(f"Unknown activation '{s}'")

2298

vllm/model_executor/layers/quantization/fp8.py

+2 -1

@@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe import (

23 FusedMoEPermuteExpertsUnpermute,

24 FusedMoEPrepareAndFinalize,

25 FusedMoeWeightScaleSupported,

26 + MoEActivation,

27 )

28 from vllm.model_executor.layers.fused_moe.config import (

29 FusedMoEQuantConfig,

@@ -965,7 +966,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

966 # TODO(rob): convert this to MK.

967 if layer.enable_eplb:

968 raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")

968 - assert layer.activation == "silu", (

969 + assert layer.activation == MoEActivation.SILU, (

970 f"Expected 'silu' activation but got {layer.activation}"

971 )

972

vllm/model_executor/layers/quantization/gguf.py

+8 -8

@@ -12,6 +12,10 @@ from torch.nn.parameter import Parameter, UninitializedParameter

12

13 from vllm import _custom_ops as ops

14 from vllm.logger import init_logger

15 +from vllm.model_executor.layers.fused_moe.activation import (

16 + MoEActivation,

17 + apply_moe_activation,

18 +)

19 from vllm.model_executor.layers.fused_moe.config import (

20 FusedMoEConfig,

21 FusedMoEQuantConfig,

@@ -246,16 +250,13 @@ def _fused_moe_gguf(

250 qweight_type2: int,

251 activation: str,

252 ) -> torch.Tensor:

253 + activation_enum = MoEActivation.from_str(activation)

254 +

255 def act(x: torch.Tensor):

256 d = x.shape[-1] // 2

257 output_shape = x.shape[:-1] + (d,)

258 out = torch.empty(output_shape, dtype=x.dtype, device=x.device)

253 - if activation == "silu":

254 - torch.ops._C.silu_and_mul(out, x)

255 - elif activation == "gelu":

256 - torch.ops._C.gelu_and_mul(out, x)

257 - else:

258 - raise ValueError(f"Unsupported activation: {activation}")

259 + apply_moe_activation(activation_enum, out, x)

260 return out

261

262 # lazy import to avoid triggering triton import in CPU backend

@@ -637,7 +638,6 @@ class GGUFMoEMethod(FusedMoEMethodBase):

638 topk_ids: torch.Tensor,

639 shared_experts_input: torch.Tensor | None,

640 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

640 - assert layer.activation == "silu", "Only SiLU activation is supported."

641 if layer.apply_router_weight_on_input:

642 raise NotImplementedError(

643 "Apply router weight on input is not supported for"

@@ -652,7 +652,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

652 topk_ids,

653 layer.w13_qweight_type.weight_type,

654 layer.w2_qweight_type.weight_type,

655 - layer.activation,

655 + layer.activation.value,

656 )

657

658

vllm/model_executor/layers/quantization/modelopt.py

+6 -2

@@ -10,6 +10,7 @@ from torch.nn.parameter import Parameter

10 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.attention import Attention

13 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

14 from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEConfig,

16 FusedMoEQuantConfig,

@@ -936,7 +937,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

937 )

938 # TODO(rob): this validation should happen at kernel selection

939 # time in the oracle rather than here.

939 - assert layer.activation == "silu", (

940 + assert layer.activation == MoEActivation.SILU, (

941 f"Expected 'silu' activation but got {layer.activation}"

942 )

943 assert not layer.renormalize

@@ -965,7 +966,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

966 # TODO(rob): this validation should happen at kernel selection

967 # time in the oracle rather than here.

968 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

968 - assert layer.activation in ("silu", "relu2_no_mul"), (

969 + assert layer.activation in (

970 + MoEActivation.SILU,

971 + MoEActivation.RELU2_NO_MUL,

972 + ), (

973 "Expected activation to be in ('silu', 'relu2_no_mul'),"

974 f"but got {layer.activation}"

975 )

vllm/model_executor/layers/quantization/moe_wna16.py

+4 -1

@@ -6,6 +6,7 @@ from typing import Any

6 import torch

7

8 from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.model_executor.layers.fused_moe.config import (

11 FusedMoEQuantConfig,

12 int4_w4a16_moe_quant_config,

@@ -371,7 +372,9 @@ class MoeWNA16Method(FusedMoEMethodBase):

372 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

373 from vllm.model_executor.layers.fused_moe import fused_experts

374

374 - assert layer.activation == "silu", "Only SiLU activation is supported."

375 + assert layer.activation == MoEActivation.SILU, (

376 + f"Only SiLU activation is supported, not {layer.activation}."

377 + )

378

379 return fused_experts(

380 x,

vllm/model_executor/layers/quantization/mxfp4.py

+4 -2

@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe import (

13 FusedMoE,

14 FusedMoEConfig,

15 FusedMoEMethodBase,

16 + MoEActivation,

17 )

18 from vllm.model_executor.layers.fused_moe import modular_kernel as mk

19 from vllm.model_executor.layers.fused_moe.config import (

@@ -1141,8 +1142,9 @@ class XpuMxfp4MoEMethod(Mxfp4MoEMethod):

1142 x: torch.Tensor,

1143 router_logits: torch.Tensor,

1144 ) -> torch.Tensor:

1144 - assert layer.activation == "swigluoai", (

1145 - "Only swiglu_oai activation is supported for XPU MXFP4 MoE"

1145 + assert layer.activation == MoEActivation.SWIGLUOAI, (

1146 + "Only swiglu_oai activation is supported for "

1147 + f"XPU MXFP4 MoE, not {layer.activation}."

1148 )

1149 from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe

1150

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2 -1

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe import (

15 FusedMoEConfig,

16 FusedMoEMethodBase,

17 FusedMoeWeightScaleSupported,

18 + MoEActivation,

19 )

20 from vllm.model_executor.layers.fused_moe.config import (

21 FusedMoEQuantConfig,

@@ -438,7 +439,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

439 expert_map=layer.expert_map,

440 )

441 elif self.use_marlin:

441 - assert layer.activation == "silu", (

442 + assert layer.activation == MoEActivation.SILU, (

443 f"{layer.activation} not supported for Marlin MoE."

444 )

445 return fused_marlin_moe(

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+7 -6

@@ -9,6 +9,7 @@ import torch

9 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

10 from vllm import _custom_ops as ops

11 from vllm.logger import init_logger

12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

13 from vllm.model_executor.layers.fused_moe.config import (

14 FusedMoEConfig,

15 FusedMoEParallelConfig,

@@ -64,9 +65,9 @@ def _supports_quant_scheme(

65 return (weight_key, activation_key) in SUPPORTED_W_A

66

67

67 -def _supports_activation(activation: str) -> bool:

68 +def _supports_activation(activation: MoEActivation) -> bool:

69 """Supports silu activation only."""

69 - return activation in ["silu"]

70 + return activation in [MoEActivation.SILU]

71

72

73 def _supports_routing_method(

@@ -267,7 +268,7 @@ def flashinfer_trtllm_fp4_moe(

268 x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

269 router_logits: torch.Tensor,

270 top_k: int,

270 - activation: str,

271 + activation: MoEActivation,

272 global_num_experts: int,

273 num_expert_group: int | None,

274 topk_group: int | None,

@@ -297,7 +298,7 @@ def flashinfer_trtllm_fp4_moe(

298 from vllm.model_executor.models.llama4 import Llama4MoE

299

300 # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2404

300 - assert activation == "silu", (

301 + assert activation == MoEActivation.SILU, (

302 "Only SiLU activation is supported for FlashInfer TRTLLM FP4 MoE. "

303 f"{activation} found instead."

304 )

@@ -365,7 +366,7 @@ def flashinfer_trtllm_fp4_routed_moe(

366 topk_ids: torch.Tensor,

367 topk_weights: torch.Tensor,

368 top_k: int,

368 - activation: str,

369 + activation: MoEActivation,

370 global_num_experts: int,

371 ) -> torch.Tensor:

372 """

@@ -387,7 +388,7 @@ def flashinfer_trtllm_fp4_routed_moe(

388 import flashinfer

389

390 # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2535

390 - assert activation == "silu", (

391 + assert activation == MoEActivation.SILU, (

392 "Only SiLU activation is supported for FlashInfer TRTLLM FP4 Routed MoE. "

393 f"{activation} found instead."

394 )

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

+3 -2

@@ -6,6 +6,7 @@ from typing import Any

6 import torch

7

8 from vllm.logger import init_logger

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.platforms import current_platform

11 from vllm.triton_utils import triton

12 from vllm.utils.import_utils import has_triton_kernels

@@ -88,7 +89,7 @@ def _can_support_mxfp4(

89 e_score_correction_bias: torch.Tensor | None = None,

90 apply_router_weight_on_input: bool = False,

91 scoring_func: str = "softmax",

91 - activation: str = "swigluoai",

92 + activation: MoEActivation = MoEActivation.SWIGLUOAI,

93 expert_load_view: torch.Tensor | None = None,

94 logical_to_physical_map: torch.Tensor | None = None,

95 logical_replica_count: torch.Tensor | None = None,

@@ -101,7 +102,7 @@ def _can_support_mxfp4(

102 or e_score_correction_bias

103 or apply_router_weight_on_input

104 or scoring_func != "softmax"

104 - or activation != "swigluoai"

105 + or activation != MoEActivation.SWIGLUOAI

106 or expert_load_view

107 or logical_to_physical_map

108 or logical_replica_count

vllm/model_executor/models/nemotron_h.py

+5 -2

@@ -33,8 +33,11 @@ from vllm.distributed.communication_op import tensor_model_parallel_all_gather

33 from vllm.distributed.parallel_state import get_pp_group

34 from vllm.model_executor.layers.activation import ReLUSquaredActivation

35 from vllm.model_executor.layers.attention import Attention

36 -from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE

37 -from vllm.model_executor.layers.fused_moe.utils import activation_without_mul

36 +from vllm.model_executor.layers.fused_moe import (

37 + FusedMoE,

38 + SharedFusedMoE,

39 + activation_without_mul,

40 +)

41 from vllm.model_executor.layers.layernorm import RMSNorm

42 from vllm.model_executor.layers.linear import (

43 ColumnParallelLinear,

[Refactor] Replace `activation: str` with `MoEActivation` enum (#33843)