@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6e311c6e20

[MoE Refactor] Rename FusedMoE to FusedMoEFactory (#44941)

Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: bnellnm <49004751+bnellnm@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

bnellnm committed Jul 31, 2026 at 06:28 UTC 6e311c6e2014a54868e2cbedf944a6873d06228b
80 files changed +243 -234
tests/distributed/test_eplb_fused_moe_layer.py
+4 -4
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -# Test that the interaction between EPLB and FusedMoE Layer is okay
4 +# Test that the interaction between EPLB and FusedMoEFactory Layer is okay
5
6 from dataclasses import dataclass
7
@@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import (
16 get_eplb_group,
17 get_tp_group,
18 )
19 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE
19 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
20
21 from .eplb_utils import distributed_run, set_env_vars_and_device
22
@@ -69,8 +69,8 @@ def make_fused_moe_layer(
69 rank: int,
70 layer_idx: int,
71 test_config: TestConfig,
72 -) -> FusedMoE:
73 - fml = FusedMoE(
72 +) -> MoERunner:
73 + fml = FusedMoEFactory(
74 num_experts=test_config.num_experts,
75 top_k=test_config.num_topk,
76 hidden_size=test_config.hidden_size,
tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py
+4 -4
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -# Test that the interaction between EPLB and FusedMoE Layer is okay for DP w/ NVFP4
4 +# Test that the interaction between EPLB and MoERunner Layer is okay for DP w/ NVFP4
5
6 from dataclasses import dataclass
7
@@ -19,7 +19,7 @@ from vllm.distributed.parallel_state import (
19 get_eplb_group,
20 )
21 from vllm.forward_context import set_forward_context
22 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE
22 +from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory, MoERunner
23 from vllm.model_executor.layers.quantization.modelopt import (
24 ModelOptNvFp4Config,
25 ModelOptNvFp4FusedMoE,
@@ -44,7 +44,7 @@ def make_fused_moe_layer(
44 rank: int,
45 layer_idx: int,
46 test_config: TestConfig,
47 -) -> FusedMoE:
47 +) -> MoERunner:
48 quant_config = None
49
50 device = torch.device(f"cuda:{rank}")
@@ -55,7 +55,7 @@ def make_fused_moe_layer(
55 exclude_modules=[],
56 )
57
58 - fml = FusedMoE(
58 + fml = FusedMoEFactory(
59 num_experts=test_config.num_experts,
60 top_k=test_config.num_topk,
61 hidden_size=test_config.hidden_size,
tests/kernels/moe/test_deepep_v2_moe.py
+1 -1
@@ -409,7 +409,7 @@ def _deep_ep_v2_moe_cudagraph(
409 vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm")
410
411 with set_current_vllm_config(vllm_cfg):
412 - # Initialize vLLM parallel state (needed by FusedMoE layer)
412 + # Initialize vLLM parallel state (needed by MoERunner layer)
413 temp_file = tempfile.mktemp()
414 init_distributed_environment(
415 world_size=pgi.world_size,
tests/kernels/moe/test_moe_layer.py
+12 -7
@@ -42,7 +42,11 @@ from vllm.distributed.device_communicators.all_reduce_utils import (
42 from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
43 from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
44 from vllm.forward_context import set_forward_context
45 -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, fused_experts
45 +from vllm.model_executor.layers.fused_moe import (
46 + FusedMoEFactory,
47 + MoERunner,
48 + fused_experts,
49 +)
50 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
51 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
52 from vllm.model_executor.layers.fused_moe.router.router_factory import (
@@ -80,7 +84,7 @@ NUM_EXPERTS = [8, 64]
84 TOP_KS = [2, 6]
85
86 # dp_size, tp_size, use_ep
83 -# Note: DP+TP is not yet supported in the FusedMoE layer.
87 +# Note: DP+TP is not yet supported in the FusedMoEFactory layer.
88 PARALLEL_COMBOS = [
89 [1, 2, False],
90 [1, 4, False],
@@ -1023,7 +1027,7 @@ def make_fused_moe_layer(
1027 kwargs["routed_input_transform"] = routed_input_transform
1028 kwargs["routed_output_transform"] = routed_output_transform
1029
1026 - layer = FusedMoE(
1030 + layer = FusedMoEFactory(
1031 num_experts=global_num_experts,
1032 top_k=top_k,
1033 hidden_size=hidden_size,
@@ -1267,7 +1271,7 @@ def _test_body_eplb(
1271 ):
1272 output_before = sp_wrapper(moe_layer)(hidden_states, router_logits)
1273
1270 - # Create a fresh FusedMoE layer with enable_eplb=True
1274 + # Create a fresh MoERunner layer with enable_eplb=True
1275 # Delete the original layer's registration so the constructor can
1276 # re-use the same "from_forward_context" prefix
1277 cc = vllm_config.compilation_config
@@ -1401,7 +1405,8 @@ def _run_one_config(
1405
1406 - When is_sequence_parallel=True (EP + sequence splitting):
1407 * ep_size: Number of expert parallel ranks (equals dp_size * tp_size)
1404 - * tp_size: Number of ranks to split sequence across (becomes sp_size in FusedMoE)
1408 + * tp_size: Number of ranks to split sequence across (becomes sp_size in
1409 + MoERunner)
1410 * Weights are chunked by ep_size (experts) but NOT by tp_size
1411 * Input sequences are chunked by tp_size (via sp_wrapper)
1412 """
@@ -1476,8 +1481,8 @@ def _run_one_config(
1481 torch.accelerator.empty_cache()
1482
1483 with set_current_vllm_config(vllm_config):
1479 - # Chunk weights for EP BEFORE creating FusedMoE
1480 - # FusedMoE uses EP-chunked weights and handles reductions internally
1484 + # Chunk weights for EP BEFORE creating MoERunner.
1485 + # MoERunner uses EP-chunked weights and handles reductions internally.
1486 if ep_size > 1:
1487 # Split experts across ranks (dimension 0 is the expert dimension)
1488 # When EP is enabled, use EP group rank and ep_size for chunking
tests/kernels/moe/test_moe_weight_loading_padded.py
+1 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for FusedMoE weight loading with padded hidden dimensions.
3 +"""Tests for FusedMoEFactory weight loading with padded hidden dimensions.
4
5 When using DeepEP backends or NIXL EP with models like nemotron_h,
6 hidden_size may be rounded up (e.g., 2688 -> 3072) for backend requirements.
tests/kernels/moe/test_routing_simulator.py
+4 -4
@@ -6,7 +6,7 @@ Test script for the token-to-expert routing simulator.
6
7 This script demonstrates how to use the routing simulator to test
8 different routing strategies and analyze their performance, including
9 -integration tests with FusedMoE layer.
9 +integration tests with FusedMoEFactory layer.
10 """
11
12 import tempfile
@@ -77,11 +77,11 @@ def test_basic_functionality(
77
78 def test_routing_strategy_integration(monkeypatch, device):
79 """Test that the routing strategy environment variable works with
80 - FusedMoE."""
80 + FusedMoEFactory."""
81 pytest.importorskip("vllm.model_executor.layers.fused_moe.layer")
82
83 import vllm.envs as envs
84 - from vllm.model_executor.layers.fused_moe.layer import FusedMoE
84 + from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory
85
86 # Test parameters
87 num_tokens = 32
@@ -111,7 +111,7 @@ def test_routing_strategy_integration(monkeypatch, device):
111 )
112
113 for strategy in strategies:
114 - fused_moe = FusedMoE(
114 + fused_moe = FusedMoEFactory(
115 num_experts=num_experts,
116 top_k=top_k,
117 hidden_size=hidden_size,
tests/kernels/moe/test_shared_fused_moe_routed_transform.py
+10 -10
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """
4 -Tests for FusedMoE with routed_input_transform.
4 +Tests for FusedMoEFactory with routed_input_transform.
5
6 -Verifies that applying routed_input_transform inside FusedMoE
6 +Verifies that applying routed_input_transform inside FusedMoEFactory
7 produces the same results as applying the transform manually outside.
8 """
9
@@ -13,7 +13,7 @@ import torch.nn as nn
13
14 from vllm.config import VllmConfig, set_current_vllm_config
15 from vllm.forward_context import set_forward_context
16 -from vllm.model_executor.layers.fused_moe import FusedMoE
16 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
17 from vllm.platforms import current_platform
18 from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed
19
@@ -133,9 +133,9 @@ def test_routed_input_transform_inside_vs_outside(
133 workspace_init,
134 monkeypatch,
135 ):
136 - """Compare FusedMoE with transform inside vs manually applying outside.
137 - Method A (inside): FusedMoE with routed_input_transform
138 - Method B (outside): Manually transform, then FusedMoE without transform
136 + """Compare FusedMoEFactory with transform inside vs manually applying outside.
137 + Method A (inside): FusedMoEFactory with routed_input_transform
138 + Method B (outside): Manually transform, then FusedMoEFactory without transform
139 """
140 if current_platform.is_rocm():
141 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
@@ -157,8 +157,8 @@ def test_routed_input_transform_inside_vs_outside(
157 routed_transform = SimpleLinear(hidden_size, latent_size, dtype)
158
159 with set_current_vllm_config(vllm_config):
160 - # Method A: FusedMoE WITH routed_input_transform
161 - moe_with_transform = FusedMoE(
160 + # Method A: FusedMoEFactory WITH routed_input_transform
161 + moe_with_transform = FusedMoEFactory(
162 shared_experts=shared_experts,
163 routed_input_transform=routed_transform,
164 num_experts=num_experts,
@@ -173,9 +173,9 @@ def test_routed_input_transform_inside_vs_outside(
173 prefix="moe_with_transform",
174 )
175
176 - # Method B: FusedMoE WITHOUT routed_input_transform
176 + # Method B: FusedMoEFactory WITHOUT routed_input_transform
177 # Note: shared_experts=None because when transform is done outside,
178 - moe_without_transform = FusedMoE(
178 + moe_without_transform = FusedMoEFactory(
179 shared_experts=None,
180 routed_input_transform=None,
181 num_experts=num_experts,
tests/kernels/moe/test_zero_expert_moe.py
+16 -13
@@ -1,10 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for FusedMoE with zero experts.
3 +"""Tests for FusedMoEFactory with zero experts.
4
5 Verifies that:
6 - The ZeroExpertRouter is properly created and used as the layer router.
7 -- A forward pass through FusedMoE with zero experts produces correct output.
7 +- A forward pass through FusedMoEFactory with zero experts produces correct output.
8 - The output decomposes correctly into real expert + zero expert contributions.
9
10 Note: tests generated with Claude.
@@ -15,7 +15,7 @@ import torch
15
16 from vllm.config import VllmConfig, set_current_vllm_config
17 from vllm.forward_context import get_forward_context, set_forward_context
18 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE
18 +from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory
19 from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
20 ZeroExpertRouter,
21 )
@@ -24,7 +24,7 @@ from vllm.v1.worker.workspace import init_workspace_manager
24
25 @pytest.fixture
26 def zero_expert_moe(dist_init, default_vllm_config):
27 - """Create a FusedMoE layer with zero experts."""
27 + """Create a FusedMoEFactory layer with zero experts."""
28 num_experts = 4
29 top_k = 2
30 # hidden_size must be >= 256 for the zero expert identity kernel to
@@ -45,7 +45,7 @@ def zero_expert_moe(dist_init, default_vllm_config):
45 with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
46 init_workspace_manager(torch.accelerator.current_device_index())
47
48 - layer = FusedMoE(
48 + layer = FusedMoEFactory(
49 zero_expert_type="identity",
50 e_score_correction_bias=e_score_correction_bias,
51 num_experts=num_experts,
@@ -66,7 +66,7 @@ def zero_expert_moe(dist_init, default_vllm_config):
66
67 @pytest.mark.parametrize("num_tokens", [1, 32])
68 def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens):
69 - """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter."""
69 + """Verify that FusedMoEFactory with zero_expert_type creates a ZeroExpertRouter."""
70 layer, _ = zero_expert_moe
71 assert isinstance(layer.router, ZeroExpertRouter), (
72 f"Expected ZeroExpertRouter but got {type(layer.router).__name__}."
@@ -83,7 +83,10 @@ def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_token
83
84 @pytest.mark.parametrize("num_tokens", [1, 32])
85 def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
86 - """Run a forward pass through FusedMoE with zero experts and verify output shape."""
86 + """
87 + Run a forward pass through FusedMoEFactory with zero experts
88 + and verify output shape.
89 + """
90 layer, vllm_config = zero_expert_moe
91
92 hidden_size = layer.routed_experts.hidden_size
@@ -118,16 +121,16 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
121
122 @pytest.mark.parametrize("num_tokens", [1, 32])
123 def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
121 - """Validate that the FusedMoE output equals a plain FusedMoE
124 + """Validate that the FusedMoEFactory output equals a plain FusedMoEFactory
125 output (real experts only) plus the zero expert contribution.
126
127 The key invariant is:
128 zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real)
129 + zero_expert_output
130
128 - We create a plain FusedMoE layer with the same weights and real-expert-only
131 + We create a plain FusedMoEFactory layer with the same weights and real-expert-only
132 router logits, compute the zero expert output via the ZeroExpertRouter, and
130 - verify the sum matches the FusedMoE output.
133 + verify the sum matches the FusedMoEFactory output.
134 """
135 layer, vllm_config = zero_expert_moe
136 num_experts = 4
@@ -152,9 +155,9 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
155 with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
156 get_forward_context().all_moe_layers = None
157
155 - # Create a plain FusedMoE layer with the same config but no zero
158 + # Create a plain FusedMoEFactory layer with the same config but no zero
159 # experts. Use a separate prefix to avoid collision.
157 - plain_layer = FusedMoE(
160 + plain_layer = FusedMoEFactory(
161 num_experts=num_experts,
162 top_k=layer.routed_experts.top_k,
163 hidden_size=layer.routed_experts.hidden_size,
@@ -210,7 +213,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
213 expected,
214 atol=4e-3,
215 rtol=4e-3,
213 - msg="FusedMoE output should equal plain FusedMoE output "
216 + msg="FusedMoEFactory output should equal plain FusedMoEFactory output "
217 "plus zero expert contribution",
218 )
219
tests/lora/test_lora_manager.py
+1 -1
@@ -198,7 +198,7 @@ def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_
198 self.gate = gate # canonical path: "moe.gate"
199
200 # Inner submodule holding the SAME gate instance under another
201 - # path. This mirrors how FusedMoE.runner.gate references the
201 + # path. This mirrors how MoERunner.runner.gate references the
202 # block's gate in qwen3_moe.
203 class _Runner(nn.Module):
204 def __init__(self, g):
tests/quantization/test_fp8.py
+2 -2
@@ -21,7 +21,7 @@ from vllm.model_executor.layers.attention.attention import (
21 Attention,
22 set_default_quant_scales,
23 )
24 -from vllm.model_executor.layers.fused_moe import FusedMoE
24 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
25 from vllm.model_executor.layers.quantization.fp8 import (
26 Fp8Config,
27 Fp8KVCacheMethod,
@@ -393,7 +393,7 @@ def test_fp8_reloading(
393 method.use_marlin = use_marlin
394
395 else:
396 - layer = FusedMoE(
396 + layer = FusedMoEFactory(
397 num_experts=1,
398 top_k=1,
399 hidden_size=1,
vllm/_custom_ops.py
+1 -1
@@ -1022,7 +1022,7 @@ def cutlass_fp4_moe_mm(
1022 An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
1023 the gemms for each combination based on the specified problem sizes.
1024
1025 - This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward.
1025 + This is used as the MoE gemm during NVFP4 Quantized MoERunner forward.
1026 - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
1027 input and expert weights.
1028 - a_/b_scales: The blockscales in FP8-E4M3 precision
vllm/config/quantization.py
+1 -1
@@ -88,7 +88,7 @@ class QuantizationConfigArgs:
88 """Spec applied to ``LinearBase`` layers."""
89
90 moe: QuantSpec | None = None
91 - """Spec applied to ``FusedMoE`` layers."""
91 + """Spec applied to ``FusedMoEFactory`` layers."""
92
93 ignore: list[str] = Field(default_factory=list)
94 """Layers to skip quantization for."""
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+1 -1
@@ -170,7 +170,7 @@ class ExampleConnector(KVConnectorBase_V1):
170
171 # Only process layers that have kv_cache
172 # attribute (attention layers) Skip non-attention
173 - # layers like FusedMoE/MLP etc.
173 + # layers like FusedMoEFactory/MLP etc.
174 kv_cache_layer = getattr(layer, "kv_cache", None)
175 if kv_cache_layer is None:
176 continue
vllm/lora/layers/fused_moe.py
+1 -1
@@ -44,7 +44,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
44 "Monolithic kernels are not supported for Fused MoE LoRA."
45 )
46
47 - # Use the MoE-aware TP rank/size: when EP is active, FusedMoE collapses
47 + # Use the MoE-aware TP rank/size: when EP is active, MoERunner collapses
48 # moe_parallel_config.tp_size to 1 (experts are sharded across the
49 # TP group instead).
50 moe_parallel_config = self.moe_config.moe_parallel_config
vllm/lora/lora_model.py
+2 -2
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
24
25 @dataclass(frozen=True)
26 class MoEEPLoadSpec:
27 - """Per-expert-parallel slicing metadata for one FusedMoE LoRA module.
27 + """Per-expert-parallel slicing metadata for one FusedMoEFactory LoRA module.
28
29 Threaded into the LoRA loader so per-expert weights from EP ranks
30 other than this one can be skipped before they ever hit CPU memory.
@@ -193,7 +193,7 @@ class LoRAModel:
193 skip_prefixes: List of module name prefixes to skip during loading.
194 Models can define this to skip modules not used in inference
195 (e.g., MTP layers). Format: ["mtp."]
196 - moe_ep_spec: When 2D FusedMoE LoRA modules are present with
196 + moe_ep_spec: When 2D FusedMoEFactory LoRA modules are present with
197 expert parallelism enabled, the (ep_rank, local, global)
198 slicing metadata shared across all MoE layers. Non-local
199 expert weights are skipped at read time instead of being
vllm/lora/model_manager.py
+2 -2
@@ -739,7 +739,7 @@ class LoRAModelManager:
739
740 def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:
741 for module_name, new_module_names in self.packed_modules.items():
742 - # For 2D FusedMoE modules with EP, narrow the per-expert
742 + # For 2D RoutedExperts modules with EP, narrow the per-expert
743 # sub-module list to this rank's owned experts so pack_moe
744 # produces a tensor sized to local_num_experts directly.
745 packed_module_names = new_module_names
@@ -1096,7 +1096,7 @@ class LoRAModelManager:
1096
1097 def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None:
1098 """
1099 - Per-rank slicing metadata for 2D FusedMoE LoRA modules.
1099 + Per-rank slicing metadata for 2D RoutedEXperts LoRA modules.
1100 """
1101 if not self._use_ep or not self._is_moe:
1102 return None
vllm/model_executor/layers/fused_moe/__init__.py
+2 -2
@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
19 FusedMoEMethodBase,
20 )
21 from vllm.model_executor.layers.fused_moe.layer import (
22 - FusedMoE,
22 + FusedMoEFactory,
23 fused_moe_make_expert_params_mapping,
24 )
25 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -63,7 +63,7 @@ def get_config() -> dict[str, Any] | None:
63
64
65 __all__ = [
66 - "FusedMoE",
66 + "FusedMoEFactory",
67 "FusedMoERouter",
68 "FusedMoEConfig",
69 "FusedMoEQuantConfig",
vllm/model_executor/layers/fused_moe/config.py
+3 -3
@@ -1134,9 +1134,9 @@ class FusedMoEParallelConfig:
1134 level's of parallelism to use in the fused moe layer.
1135
1136 Args:
1137 - tp_size_ (int): `tp_size` passed into the FusedMoE constructor.
1138 - pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor.
1139 - dp_size_ (int): `dp_size` passed into the FusedMoE constructor.
1137 + tp_size_ (int): `tp_size` passed into the FusedMoEFactory constructor.
1138 + pcp_size_ (int): `pcp_size` passed into the FusedMoEFactory constructor.
1139 + dp_size_ (int): `dp_size` passed into the FusedMoEFactory constructor.
1140 vllm_parallel_config (ParallelConfig): vLLM's parallel config
1141 object which contains the `enable_expert_parallel` flag.
1142
vllm/model_executor/layers/fused_moe/layer.py
+1 -2
@@ -96,8 +96,7 @@ def determine_expert_counts(
96 return global_num_experts, logical_num_experts, num_fused_shared_experts
97
98
99 -# TODO: rename this
100 -def FusedMoE(
99 +def FusedMoEFactory(
100 num_experts: int, # Global number of experts
101 top_k: int,
102 hidden_size: int,
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+2 -2
@@ -1016,7 +1016,7 @@ def _process_weights_xpu(
1016 w2: [E, K, N] int4 (uint8 storage [E, K, N // 2])
1017 w2_scales: [E, K, N // group_size] params_dtype
1018
1019 - Input GPTQ layout from FusedMoE.weight_loader:
1019 + Input GPTQ layout from MoERunner.weight_loader:
1020 w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim)
1021 w13_scales: [E, K // group_size, 2*N] params_dtype
1022 w2: [E, N // 8, K] int32
@@ -1419,7 +1419,7 @@ def convert_to_wna16_moe_kernel_format(
1419
1420 Args:
1421 backend: the selected ``WNA16MoEBackend``.
1422 - layer: the ``FusedMoE`` layer whose parameters are being prepared.
1422 + layer: the ``MoERunner`` layer whose parameters are being prepared.
1423 quant_config: the ``QuantizationConfig`` for this layer.
1424 input_dtype: optional activation dtype, usually should be 16 bit.
1425 """
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+1 -1
@@ -180,7 +180,7 @@ class RoutedExpertsCapturer:
180 # (trailing rows are SP ceil-div padding). The TP group
181 # is always initialized on real rollout workers, and
182 # every rank in the group reaches this branch in
183 - # lockstep (bind is per-FusedMoE layer, SP is a global
183 + # lockstep (bind is per-FusedMoEFactory layer, SP is a global
184 # condition), so a bare all_gather here will not
185 # deadlock -- let it raise if the precondition is
186 # violated rather than skip silently.
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+2 -2
@@ -353,7 +353,7 @@ class MoERunner(MoERunnerInterface):
353 ) -> tuple[torch.Tensor, torch.Tensor | None]:
354 """Apply transform for routed experts (e.g., latent projection).
355
356 - This is called by FusedMoE.forward_native. The original hidden_states
356 + This is called by MoERunner.forward_native. The original hidden_states
357 is saved separately so shared experts get [S, hidden_size] while
358 routed experts get the transformed [S, moe_latent_size].
359
@@ -646,7 +646,7 @@ class MoERunner(MoERunnerInterface):
646 ):
647 # If router/gate provided, then apply it here.
648 # (Note: This code runs only when "overlapped mode" is on to allow
649 - # parallel execution of shared experts with the FusedMoE via
649 + # parallel execution of shared experts with the RoutedExperts via
650 # separate cuda stream)
651 if self._shared_experts is not None:
652 assert shared_experts_input is not None
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
+1 -1
@@ -62,7 +62,7 @@ class MoERunnerInterface(PluggableLayer, ABC):
62
63 ########################################################################
64 #
65 - # FusedMoE layer methods
65 + # FusedMoEFactory layer methods
66 #
67 ########################################################################
68
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1 -1
@@ -215,7 +215,7 @@ class CompressedTensorsConfig(QuantizationConfig):
215 def _add_fused_moe_to_target_scheme_map(self): # XXXXXXXXXXXXXXXXXXXXXX
216 """
217 Helper function to update target_scheme_map
218 - since linear layers get fused into FusedMoE
218 + since linear layers get fused into RoutedExperts
219 targeting 'Linear' needs to also match
220 RoutedExperts modules.
221 """
vllm/model_executor/layers/quantization/modelopt.py
+1 -1
@@ -2363,7 +2363,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2363 if key.startswith(prefix_dot):
2364 return info["quant_algo"].upper()
2365
2366 - # FusedMoE expert prefix is e.g. "...moe.experts", while ModelOpt's
2366 + # RoutedExperts expert prefix is e.g. "...moe.experts", while ModelOpt's
2367 # quantized_layers entries use "...moe.gate_proj" / "...moe.up_proj".
2368 if prefix.endswith(".experts"):
2369 parent_dot = prefix.rsplit(".experts", 1)[0] + "."
vllm/model_executor/model_loader/bitsandbytes_loader.py
+1 -1
@@ -467,7 +467,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
467 elif isinstance(module, RoutedExperts) and hasattr(
468 module.quant_method, "quant_config"
469 ):
470 - # TODO: support FusedMoE with prequant and 8bit.
470 + # TODO: support RoutedExperts with prequant and 8bit.
471 if self.pre_quant and self.load_8bit:
472 raise ValueError(
473 "Prequant BitsAndBytes 8bit models with RoutedExperts "
vllm/model_executor/models/AXK1.py
+2 -2
@@ -43,7 +43,7 @@ from vllm.distributed import (
43 from vllm.logger import init_logger
44 from vllm.model_executor.layers.attention import Attention
45 from vllm.model_executor.layers.fused_moe import (
46 - FusedMoE,
46 + FusedMoEFactory,
47 fused_moe_make_expert_params_mapping,
48 )
49 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -168,7 +168,7 @@ class AXK1MoE(nn.Module):
168 prefix=f"{prefix}.shared_experts",
169 )
170
171 - self.experts = FusedMoE(
171 + self.experts = FusedMoEFactory(
172 shared_experts=self.shared_experts,
173 gate=self.gate,
174 num_experts=config.n_routed_experts,
vllm/model_executor/models/afmoe.py
+3 -3
@@ -18,7 +18,7 @@ from vllm.distributed import (
18 from vllm.logger import init_logger
19 from vllm.model_executor.layers.attention import Attention
20 from vllm.model_executor.layers.fused_moe import (
21 - FusedMoE,
21 + FusedMoEFactory,
22 MoERunner,
23 )
24 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -122,8 +122,8 @@ class AfmoeMoE(nn.Module):
122 prefix=f"{prefix}.shared_experts",
123 )
124
125 - # Routed experts using FusedMoE
126 - self.experts = FusedMoE(
125 + # Routed experts using FusedMoEFactory
126 + self.experts = FusedMoEFactory(
127 shared_experts=self.shared_experts,
128 num_experts=config.num_experts,
129 top_k=config.num_experts_per_tok,
vllm/model_executor/models/aria.py
+4 -2
@@ -13,7 +13,9 @@ from vllm.config import VllmConfig
13 from vllm.config.multimodal import BaseDummyOptions
14 from vllm.inputs import MultiModalDataDict
15 from vllm.model_executor.layers.activation import get_act_fn
16 -from vllm.model_executor.layers.fused_moe import FusedMoE
16 +from vllm.model_executor.layers.fused_moe import (
17 + FusedMoEFactory,
18 +)
19 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
20 from vllm.model_executor.layers.logits_processor import LogitsProcessor
21 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -222,7 +224,7 @@ class AriaTextMoELayer(nn.Module):
224 bias=config.mlp_bias,
225 )
226
225 - self.experts = FusedMoE(
227 + self.experts = FusedMoEFactory(
228 shared_experts=self.shared_experts,
229 num_experts=config.moe_num_experts,
230 top_k=config.moe_topk,
vllm/model_executor/models/bailing_moe.py
+2 -2
@@ -41,7 +41,7 @@ from vllm.distributed import (
41 )
42 from vllm.model_executor.layers.activation import SiluAndMul
43 from vllm.model_executor.layers.attention import Attention
44 -from vllm.model_executor.layers.fused_moe import FusedMoE
44 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
45 from vllm.model_executor.layers.layernorm import RMSNorm
46 from vllm.model_executor.layers.linear import (
47 MergedColumnParallelLinear,
@@ -289,7 +289,7 @@ class BailingMoE(nn.Module):
289 else:
290 self.shared_experts = None
291
292 - self.experts = FusedMoE(
292 + self.experts = FusedMoEFactory(
293 shared_experts=self.shared_experts,
294 num_experts=self.num_experts,
295 top_k=self.top_k,
vllm/model_executor/models/bailing_moe_linear.py
+4 -4
@@ -18,7 +18,7 @@ from vllm.distributed import (
18 from vllm.forward_context import get_forward_context
19 from vllm.logger import init_logger
20 from vllm.model_executor.layers.fused_moe import (
21 - FusedMoE,
21 + FusedMoEFactory,
22 fused_moe_make_expert_params_mapping,
23 )
24 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -329,8 +329,8 @@ class BailingMoeV25(nn.Module):
329 else:
330 self.shared_experts = None
331
332 - # Routed experts using FusedMoE
333 - self.experts = FusedMoE(
332 + # Routed experts using FusedMoEFactory
333 + self.experts = FusedMoEFactory(
334 shared_experts=self.shared_experts,
335 num_experts=self.num_experts,
336 top_k=self.top_k,
@@ -602,7 +602,7 @@ class BailingMoeV25Model(nn.Module):
602 (".gate_up_proj", ".up_proj", 1),
603 ]
604
605 - # Expert parameter mappings from FusedMoE
605 + # Expert parameter mappings from FusedMoEFactory
606 expert_mappings = list(self.get_expert_mapping())
607
608 def load_param(name: str, tensor: torch.Tensor, shard_id=None) -> bool:
vllm/model_executor/models/bert_with_rope.py
+1 -1
@@ -320,7 +320,7 @@ class NomicMoE(nn.Module):
320 # router_logits: (num_tokens, n_experts)
321 router_logits, _ = self.router(hidden_states)
322 # FIXME(Isotr0py): This implementation is too tricky,
323 - # we should use FusedMoE instead in the future
323 + # we should use FusedMoEFactory instead in the future
324 # after supporting ungated activation for it.
325 topk_weights, topk_ids, _ = fused_topk(
326 hidden_states, router_logits, self.top_k, renormalize=False
vllm/model_executor/models/cohere2_moe.py
+3 -3
@@ -16,7 +16,7 @@ from vllm.distributed import (
16 )
17 from vllm.model_executor.layers.activation import SiluAndMul
18 from vllm.model_executor.layers.attention import Attention
19 -from vllm.model_executor.layers.fused_moe import FusedMoE
19 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
20 from vllm.model_executor.layers.linear import (
21 MergedColumnParallelLinear,
22 QKVParallelLinear,
@@ -300,7 +300,7 @@ class Cohere2Moe(nn.Module):
300 self.shared_experts = None
301 self.shared_expert_combination_strategy = None
302
303 - self.experts = FusedMoE(
303 + self.experts = FusedMoEFactory(
304 num_experts=config.num_experts,
305 top_k=config.num_experts_per_tok,
306 hidden_size=config.hidden_size,
@@ -318,7 +318,7 @@ class Cohere2Moe(nn.Module):
318 orig_shape = hidden_states.shape
319 hidden_states = hidden_states.view(-1, self.hidden_size)
320 router_logits, _ = self.gate(hidden_states)
321 - # FusedMoE handles shared expert overlap internally and returns
321 + # FusedMoEFactory handles shared expert overlap internally and returns
322 # shared_output + routed_output when shared_experts is set.
323 final_hidden_states = self.experts(hidden_states, router_logits)
324 if self.shared_expert_combination_strategy == "average":
vllm/model_executor/models/dbrx.py
+2 -2
@@ -16,7 +16,7 @@ from vllm.distributed import (
16 )
17 from vllm.model_executor.layers.attention import Attention
18 from vllm.model_executor.layers.fused_moe import (
19 - FusedMoE,
19 + FusedMoEFactory,
20 RoutedExperts,
21 )
22 from vllm.model_executor.layers.linear import (
@@ -159,7 +159,7 @@ class DbrxMoE(nn.Module):
159
160 self.router = DbrxRouter(config, self.params_dtype)
161
162 - self.experts = FusedMoE(
162 + self.experts = FusedMoEFactory(
163 num_experts=config.ffn_config.moe_num_experts,
164 top_k=config.ffn_config.moe_top_k,
165 hidden_size=config.d_model,
vllm/model_executor/models/deepseek_v2.py
+2 -2
@@ -50,7 +50,7 @@ from vllm.model_executor.layers.activation import SiluAndMul
50 from vllm.model_executor.layers.attention import Attention, RSWAAttention
51 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
52 from vllm.model_executor.layers.fused_moe import (
53 - FusedMoE,
53 + FusedMoEFactory,
54 GateLinear,
55 fused_moe_make_expert_params_mapping,
56 )
@@ -361,7 +361,7 @@ class DeepseekV2MoE(nn.Module):
361 prefix=f"{prefix}.shared_experts",
362 )
363
364 - self.experts = FusedMoE(
364 + self.experts = FusedMoEFactory(
365 shared_experts=self.shared_experts,
366 gate=self.gate,
367 num_experts=config.n_routed_experts,
vllm/model_executor/models/ernie45_moe.py
+2 -2
@@ -41,7 +41,7 @@ from vllm.distributed import (
41 from vllm.logger import init_logger
42 from vllm.model_executor.layers.activation import SiluAndMul
43 from vllm.model_executor.layers.attention import Attention
44 -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner
44 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
45 from vllm.model_executor.layers.layernorm import RMSNorm
46 from vllm.model_executor.layers.linear import (
47 MergedColumnParallelLinear,
@@ -183,7 +183,7 @@ class Ernie4_5_MoeMoE(nn.Module):
183 else:
184 self.shared_experts = None
185
186 - self.experts = FusedMoE(
186 + self.experts = FusedMoEFactory(
187 shared_experts=self.shared_experts,
188 num_experts=config.moe_num_experts,
189 top_k=config.moe_k,
vllm/model_executor/models/ernie45_vl_moe.py
+3 -3
@@ -37,7 +37,7 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37 from vllm.logger import init_logger
38 from vllm.model_executor.layers.attention import Attention
39 from vllm.model_executor.layers.fused_moe import (
40 - FusedMoE,
40 + FusedMoEFactory,
41 fused_moe_make_expert_params_mapping,
42 )
43 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -260,7 +260,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
260 prefix=f"{prefix}.text_experts_gate",
261 )
262
263 - self.text_experts = FusedMoE(
263 + self.text_experts = FusedMoEFactory(
264 shared_experts=self.shared_experts,
265 num_experts=config.moe_num_experts[0],
266 top_k=config.moe_k,
@@ -297,7 +297,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
297 prefix=f"{prefix}.vision_experts_gate",
298 )
299
300 - self.vision_experts = FusedMoE(
300 + self.vision_experts = FusedMoEFactory(
301 shared_experts=self.shared_experts,
302 num_experts=config.moe_num_experts[1],
303 top_k=config.moe_k,
vllm/model_executor/models/exaone_moe.py
+2 -4
@@ -29,9 +29,7 @@ from vllm.distributed import (
29 get_pp_group,
30 get_tensor_model_parallel_world_size,
31 )
32 -from vllm.model_executor.layers.fused_moe import (
33 - FusedMoE,
34 -)
32 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
33 from vllm.model_executor.layers.layernorm import RMSNorm
34 from vllm.model_executor.layers.linear import ReplicatedLinear
35 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -126,7 +124,7 @@ class ExaoneMoe(nn.Module):
124 else:
125 self.shared_experts = None
126
129 - self.experts = FusedMoE(
127 + self.experts = FusedMoEFactory(
128 shared_experts=self.shared_experts,
129 gate=self.gate,
130 num_experts=self.n_routed_experts,
vllm/model_executor/models/flex_olmo.py
+2 -2
@@ -21,7 +21,7 @@ from vllm.config import VllmConfig
21 from vllm.distributed import get_tensor_model_parallel_world_size
22 from vllm.logger import init_logger
23 from vllm.model_executor.layers.fused_moe import (
24 - FusedMoE,
24 + FusedMoEFactory,
25 )
26 from vllm.model_executor.layers.layernorm import RMSNorm
27 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -73,7 +73,7 @@ class FlexOlmoMoE(nn.Module):
73 prefix=f"{prefix}.gate",
74 )
75
76 - self.experts = FusedMoE(
76 + self.experts = FusedMoEFactory(
77 num_experts=hf_config.num_experts,
78 top_k=hf_config.num_experts_per_tok,
79 hidden_size=hf_config.hidden_size,
vllm/model_executor/models/gemma4.py
+16 -16
@@ -38,7 +38,7 @@ from vllm.logger import init_logger
38 from vllm.model_executor.layers.activation import get_act_and_mul_fn
39 from vllm.model_executor.layers.attention import Attention
40 from vllm.model_executor.layers.fused_moe import (
41 - FusedMoE,
41 + FusedMoEFactory,
42 GateLinear,
43 fused_moe_make_expert_params_mapping,
44 )
@@ -305,14 +305,14 @@ class Gemma4Router(nn.Module):
305
306
307 class Gemma4MoE(nn.Module):
308 - """Mixture of Experts for Gemma4 using vLLM's FusedMoE.
308 + """Mixture of Experts for Gemma4 using vLLM's MoERunner.
309
310 - Wraps FusedMoE with custom routing. The router projection is
310 + Wraps MoERunner with custom routing. The router projection is
311 external (Gemma4Router) — this class only handles expert dispatch.
312
313 Gemma4 routing: softmax over ALL experts → top-k → renormalize.
314 per_expert_scale is folded into routing weights for mathematical
315 - correctness with FusedMoE's fused kernel.
315 + correctness with MoERunner's fused kernel.
316 """
317
318 def __init__(
@@ -326,11 +326,11 @@ class Gemma4MoE(nn.Module):
326 self.num_experts = config.num_experts
327
328 # Per-expert output scale folded into routing weights so that
329 - # FusedMoE's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
329 + # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
330 self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
331
332 # Gemma4 routing: softmax over ALL experts → top-k → renormalize.
333 - # FusedMoE's built-in fused_topk scopes softmax differently, so
333 + # MoERunner's built-in fused_topk scopes softmax differently, so
334 # a custom routing function is needed for numerical correctness.
335 # NOTE: self.per_expert_scale is read at call time (not captured into
336 # a local) so that torch.func.functional_call parameter substitution
@@ -350,8 +350,8 @@ class Gemma4MoE(nn.Module):
350 gating_output, topk, self.per_expert_scale
351 )
352
353 - # FusedMoE experts with custom Gemma4 routing
354 - self.experts = FusedMoE(
353 + # MoERunner experts with custom Gemma4 routing
354 + self.experts = FusedMoEFactory(
355 num_experts=config.num_experts,
356 top_k=config.top_k_experts,
357 hidden_size=config.hidden_size,
@@ -1382,10 +1382,10 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1382 # MoE expert weight mapping: checkpoint can have either:
1383 # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D)
1384 # 2. Already per-expert 2D weights (if quantized)
1385 - # Map to FusedMoE parameters:
1386 - # moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13)
1387 - # moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13)
1388 - # moe.experts.{id}.down_proj → FusedMoE w2
1385 + # Map to MoERunner parameters:
1386 + # moe.experts.{id}.gate_proj → MoERunner w1 (shard of w13)
1387 + # moe.experts.{id}.up_proj → MoERunner w3 (shard of w13)
1388 + # moe.experts.{id}.down_proj → MoERunner w2
1389 num_experts = getattr(self.config, "num_experts", None) or 0
1390 # Strategy A: dot-separated suffix
1391 # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight)
@@ -1475,7 +1475,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1475 continue
1476 param = params_dict[moe_name]
1477 # Expert weights are already in the correct
1478 - # orientation for FusedMoE after _weight_iterator:
1478 + # orientation for MoERunner after _weight_iterator:
1479 # gate/up: [I, H] → w1/w3 expects [I, H]
1480 # down: [H, I] → w2 expects [H, I]
1481 # Scales and other quantization params may be 1D or scalar.
@@ -1668,19 +1668,19 @@ class Gemma4ForCausalLM(
1668
1669 # MoE expert weights: checkpoint stores as 3D packed
1670 # tensors. Explode into per-expert 2D weights for
1671 - # FusedMoE weight_loader.
1671 + # MoERunner weight_loader.
1672 #
1673 # Checkpoint format:
1674 # moe.gate_up_proj: [E, 2*I, H] (fused gate + up)
1675 # moe.down_proj: [E, H, I]
1676 #
1677 - # FusedMoE expects per-expert:
1677 + # MoERunner expects per-expert:
1678 # w1 (gate): [I, H] — first half of gate_up
1679 # w3 (up): [I, H] — second half of gate_up
1680 # w2 (down): [H, I] — as-is from checkpoint
1681 #
1682 # No transpose needed: checkpoint orientation already
1683 - # matches FusedMoE's expected layout.
1683 + # matches MoERunner's expected layout.
1684 if "moe.gate_up_proj" in name and weight.dim() == 3:
1685 num_experts = weight.size(0)
1686 intermediate_size = weight.size(1) // 2
vllm/model_executor/models/glm4_moe.py
+5 -3
@@ -42,7 +42,9 @@ from vllm.distributed import (
42 from vllm.logger import init_logger
43 from vllm.model_executor.layers.activation import SiluAndMul
44 from vllm.model_executor.layers.attention import Attention
45 -from vllm.model_executor.layers.fused_moe import FusedMoE
45 +from vllm.model_executor.layers.fused_moe import (
46 + FusedMoEFactory,
47 +)
48 from vllm.model_executor.layers.layernorm import RMSNorm
49 from vllm.model_executor.layers.linear import (
50 MergedColumnParallelLinear,
@@ -164,7 +166,7 @@ class Glm4MoE(nn.Module):
166 )
167
168 # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py
167 - # pattern (see Glm4MoE / FusedMoE wiring there).
169 + # pattern (see Glm4MoE / MoERunner wiring there).
170 self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
171 self.is_fusion_moe_shared_experts_enabled = (
172 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
@@ -183,7 +185,7 @@ class Glm4MoE(nn.Module):
185 prefix=f"{prefix}.shared_experts",
186 )
187
186 - self.experts = FusedMoE(
188 + self.experts = FusedMoEFactory(
189 shared_experts=self.shared_experts,
190 num_experts=config.n_routed_experts,
191 top_k=config.num_experts_per_tok,
vllm/model_executor/models/gpt_oss.py
+3 -3
@@ -21,7 +21,7 @@ from vllm.distributed import (
21 )
22 from vllm.model_executor.layers.attention import Attention
23 from vllm.model_executor.layers.fused_moe import (
24 - FusedMoE,
24 + FusedMoEFactory,
25 fused_moe_make_expert_params_mapping,
26 )
27 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
@@ -211,7 +211,7 @@ class MLPBlock(torch.nn.Module):
211 return_bias=False,
212 )
213 assert config.intermediate_size % self.world_size == 0
214 - self.experts = FusedMoE(
214 + self.experts = FusedMoEFactory(
215 num_experts=config.num_local_experts,
216 top_k=config.num_experts_per_tok,
217 hidden_size=config.hidden_size,
@@ -1010,7 +1010,7 @@ class GptOssModel(nn.Module, EagleModelMixin):
1010 tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size)
1011
1012 # Use centralized weight remapping for MoE expert parameters.
1013 - # The FusedMoE refactor moved expert params under
1013 + # The MoERunner refactor moved expert params under
1014 # `mlp.experts.routed_experts.*`; this remaps checkpoint names so
1015 # MoE weight/bias keys resolve against params_dict.
1016 for name, weight in remap_moe_expert_weights(weights, params_dict):
vllm/model_executor/models/granitemoe.py
+2 -2
@@ -40,7 +40,7 @@ from vllm.distributed import (
40 )
41 from vllm.model_executor.layers.attention import Attention
42 from vllm.model_executor.layers.fused_moe import (
43 - FusedMoE,
43 + FusedMoEFactory,
44 fused_moe_make_expert_params_mapping,
45 )
46 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -101,7 +101,7 @@ class GraniteMoeMoE(nn.Module):
101 prefix=f"{prefix}.gate",
102 )
103
104 - self.experts = FusedMoE(
104 + self.experts = FusedMoEFactory(
105 num_experts=num_experts,
106 top_k=top_k,
107 hidden_size=hidden_size,
vllm/model_executor/models/hunyuan_v1.py
+2 -2
@@ -43,7 +43,7 @@ from vllm.distributed import (
43 from vllm.model_executor.layers.activation import SiluAndMul
44 from vllm.model_executor.layers.attention import Attention
45 from vllm.model_executor.layers.fused_moe import (
46 - FusedMoE,
46 + FusedMoEFactory,
47 fused_moe_make_expert_params_mapping,
48 )
49 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -441,7 +441,7 @@ class HunYuanSparseMoeBlock(nn.Module):
441 else:
442 self.shared_mlp = None
443
444 - self.experts = FusedMoE(
444 + self.experts = FusedMoEFactory(
445 shared_experts=self.shared_mlp,
446 num_experts=self.n_routed_experts,
447 top_k=top_k,
vllm/model_executor/models/hy_v3.py
+2 -2
@@ -44,7 +44,7 @@ from vllm.logger import init_logger
44 from vllm.model_executor.layers.activation import SiluAndMul
45 from vllm.model_executor.layers.attention import Attention
46 from vllm.model_executor.layers.fused_moe import (
47 - FusedMoE,
47 + FusedMoEFactory,
48 GateLinear,
49 fused_moe_make_expert_params_mapping,
50 )
@@ -184,7 +184,7 @@ class HYV3MoEFused(nn.Module):
184 scoring_func = "sigmoid"
185 e_score_correction_bias = self.expert_bias
186
187 - self.experts = FusedMoE(
187 + self.experts = FusedMoEFactory(
188 num_experts=self.n_routed_experts,
189 top_k=top_k,
190 hidden_size=config.hidden_size,
vllm/model_executor/models/interns1_pro.py
+2 -2
@@ -42,7 +42,7 @@ from vllm.logger import init_logger
42 from vllm.model_executor.layers.activation import SiluAndMul
43 from vllm.model_executor.layers.attention import Attention
44 from vllm.model_executor.layers.fused_moe import (
45 - FusedMoE,
45 + FusedMoEFactory,
46 )
47 from vllm.model_executor.layers.layernorm import RMSNorm
48 from vllm.model_executor.layers.linear import (
@@ -173,7 +173,7 @@ class InternS1ProMoeSparseMoeBlock(nn.Module):
173 # For custom routing function
174 self.n_groups = getattr(config, "router_n_groups", -1)
175
176 - self.experts = FusedMoE(
176 + self.experts = FusedMoEFactory(
177 num_experts=self.n_routed_experts,
178 top_k=config.num_experts_per_tok,
179 hidden_size=config.hidden_size,
vllm/model_executor/models/jamba.py
+4 -2
@@ -14,7 +14,9 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig
14 from vllm.distributed import get_tensor_model_parallel_world_size
15 from vllm.distributed.parallel_state import get_pp_group
16 from vllm.model_executor.layers.attention import Attention
17 -from vllm.model_executor.layers.fused_moe import FusedMoE
17 +from vllm.model_executor.layers.fused_moe import (
18 + FusedMoEFactory,
19 +)
20 from vllm.model_executor.layers.layernorm import RMSNorm
21 from vllm.model_executor.layers.linear import (
22 QKVParallelLinear,
@@ -81,7 +83,7 @@ class JambaMoE(nn.Module):
83 prefix=f"{prefix}.router",
84 )
85
84 - self.experts = FusedMoE(
86 + self.experts = FusedMoEFactory(
87 self.num_total_experts,
88 self.top_k,
89 self.hidden_size,
vllm/model_executor/models/laguna.py
+9 -7
@@ -19,7 +19,7 @@ from vllm.distributed import (
19 )
20 from vllm.logger import init_logger
21 from vllm.model_executor.layers.attention import Attention
22 -from vllm.model_executor.layers.fused_moe import FusedMoE
22 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
23 from vllm.model_executor.layers.layernorm import RMSNorm
24 from vllm.model_executor.layers.linear import (
25 ColumnParallelLinear,
@@ -174,7 +174,7 @@ class LagunaMoE(nn.Module):
174 prefix=f"{prefix}.gate",
175 )
176
177 - # Shared expert (optional) - passed to FusedMoE for overlap optimization
177 + # Shared expert (optional) - passed to FusedMoEFactory for overlap optimization
178 self.shared_expert: LagunaMLP | None
179 if config.shared_expert_intermediate_size > 0:
180 self.shared_expert = LagunaMLP(
@@ -191,20 +191,22 @@ class LagunaMoE(nn.Module):
191 # Auxiliary-loss-free load-balancing bias (arXiv:2408.15664). The
192 # checkpoint stores one [num_experts] tensor per MoE layer at
193 # `mlp.experts.e_score_correction_bias`; registering it as a Parameter
194 - # on the FusedMoE lets the weight loader pick it up and the router
194 + # on the MoERunner lets the weight loader pick it up and the router
195 # add it during top-k selection. The fused top-k bias router requires
196 # float32 regardless of model dtype.
197 e_score_correction_bias = torch.nn.Parameter(
198 torch.zeros(config.num_experts, dtype=torch.float32),
199 requires_grad=False,
200 )
201 -
202 - # FusedMoE with SIGMOID routing. Passing `shared_experts=` lets the
201 + # MoERunner with SIGMOID routing. Passing `shared_experts=` lets the
202 + # layer overlap the shared-expert compute with the all2all dispatch.
203 + # `apply_routed_scale_to_output=True` makes MoERunner handle the
204 + # FusedMoEFactory with SIGMOID routing. Passing `shared_experts=` lets the
205 # layer overlap the shared-expert compute with the all2all dispatch.
204 - # `apply_routed_scale_to_output=True` makes FusedMoE handle the
206 + # `apply_routed_scale_to_output=True` makes FusedMoEFactory handle the
207 # routed_scaling_factor, shared+routed combine, and TP all-reduce
208 # internally, so forward() just returns the final hidden states.
207 - self.experts = FusedMoE(
209 + self.experts = FusedMoEFactory(
210 shared_experts=self.shared_expert,
211 num_experts=config.num_experts,
212 top_k=config.num_experts_per_tok,
vllm/model_executor/models/lfm2_moe.py
+4 -2
@@ -15,7 +15,9 @@ from vllm.distributed import (
15 )
16 from vllm.model_executor.layers.activation import SiluAndMul
17 from vllm.model_executor.layers.attention import Attention
18 -from vllm.model_executor.layers.fused_moe import FusedMoE
18 +from vllm.model_executor.layers.fused_moe import (
19 + FusedMoEFactory,
20 +)
21 from vllm.model_executor.layers.layernorm import RMSNorm
22 from vllm.model_executor.layers.linear import (
23 MergedColumnParallelLinear,
@@ -143,7 +145,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module):
145 else:
146 self.gate.e_score_correction_bias = None
147
146 - self.experts = FusedMoE(
148 + self.experts = FusedMoEFactory(
149 num_experts=self.n_routed_experts,
150 top_k=config.num_experts_per_tok,
151 hidden_size=config.hidden_size,
vllm/model_executor/models/llama4.py
+2 -2
@@ -37,7 +37,7 @@ from vllm.model_executor.layers.attention import (
37 ChunkedLocalAttention,
38 )
39 from vllm.model_executor.layers.fused_moe import (
40 - FusedMoE,
40 + FusedMoEFactory,
41 fused_moe_make_expert_params_mapping,
42 )
43 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -131,7 +131,7 @@ class Llama4MoE(nn.Module):
131 self.n_physical_experts = self.n_local_experts + self.n_redundant_experts
132 self.n_local_physical_experts = self.n_physical_experts // self.ep_size
133
134 - self.experts = FusedMoE(
134 + self.experts = FusedMoEFactory(
135 shared_experts=self.shared_expert,
136 num_experts=config.num_local_experts,
137 top_k=config.num_experts_per_tok,
vllm/model_executor/models/longcat_flash.py
+4 -4
@@ -47,7 +47,7 @@ from vllm.distributed import get_pp_group
47 from vllm.logger import init_logger
48 from vllm.model_executor.layers.activation import SiluAndMul
49 from vllm.model_executor.layers.fused_moe import (
50 - FusedMoE,
50 + FusedMoEFactory,
51 fused_moe_make_expert_params_mapping,
52 )
53 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -297,7 +297,7 @@ class LongcatMoe(nn.Module):
297 )
298
299 assert config.zero_expert_type is not None
300 - self.experts = FusedMoE(
300 + self.experts = FusedMoEFactory(
301 zero_expert_type=config.zero_expert_type,
302 e_score_correction_bias=self.router.e_score_correction_bias,
303 num_experts=num_experts,
@@ -317,7 +317,7 @@ class LongcatMoe(nn.Module):
317 num_tokens, hidden_dim = hidden_states.shape
318 hidden_states = hidden_states.view(-1, hidden_dim)
319
320 - # Align to FusedMoE padded hidden size to avoid dim mismatch
320 + # Align to MoERunner padded hidden size to avoid dim mismatch
321 padded_hidden = self.experts.moe_config.hidden_dim
322 if hidden_dim < padded_hidden:
323 hidden_states_padded = torch.nn.functional.pad(
@@ -333,7 +333,7 @@ class LongcatMoe(nn.Module):
333 hidden_states_padded.to(self.router_params_dtype)
334 )
335
336 - # FusedMoE handles routing memoization and zero expert computation
336 + # MoERunner handles routing memoization and zero expert computation
337 # internally. Pass full router_logits (including zero experts) so that
338 # zero experts can be properly identified in routing.
339 final_hidden_states = self.experts(
vllm/model_executor/models/mimo_v2.py
+2 -2
@@ -24,7 +24,7 @@ from vllm.logger import init_logger
24 from vllm.model_executor.layers.activation import SiluAndMul
25 from vllm.model_executor.layers.attention import Attention
26 from vllm.model_executor.layers.fused_moe import (
27 - FusedMoE,
27 + FusedMoEFactory,
28 fused_moe_make_expert_params_mapping,
29 )
30 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -171,7 +171,7 @@ class MiMoV2MoE(nn.Module):
171 torch.empty(config.n_routed_experts, dtype=self.gate_dtype)
172 )
173
174 - self.experts = FusedMoE(
174 + self.experts = FusedMoEFactory(
175 num_experts=self.n_routed_experts,
176 top_k=config.num_experts_per_tok,
177 hidden_size=config.hidden_size,
vllm/model_executor/models/minimax_m2.py
+2 -4
@@ -39,9 +39,7 @@ from vllm.distributed import (
39 get_tensor_model_parallel_world_size,
40 )
41 from vllm.model_executor.layers.attention import Attention
42 -from vllm.model_executor.layers.fused_moe import (
43 - FusedMoE,
44 -)
42 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
43 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
44 from vllm.model_executor.layers.layernorm import RMSNorm
45 from vllm.model_executor.layers.linear import (
@@ -95,7 +93,7 @@ class MiniMaxM2MoE(nn.Module):
93 else:
94 self.e_score_correction_bias = None
95
98 - self.experts = FusedMoE(
96 + self.experts = FusedMoEFactory(
97 num_experts=config.num_local_experts,
98 top_k=config.num_experts_per_tok,
99 scoring_func=config.scoring_func,
vllm/model_executor/models/mixtral.py
+2 -4
@@ -39,9 +39,7 @@ from vllm.distributed import (
39 get_tensor_model_parallel_world_size,
40 )
41 from vllm.model_executor.layers.attention import Attention
42 -from vllm.model_executor.layers.fused_moe import (
43 - FusedMoE,
44 -)
42 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
43 from vllm.model_executor.layers.layernorm import RMSNorm
44 from vllm.model_executor.layers.linear import (
45 QKVParallelLinear,
@@ -123,7 +121,7 @@ class MixtralMoE(nn.Module):
121 prefix=f"{prefix}.gate",
122 )
123
126 - self.experts = FusedMoE(
124 + self.experts = FusedMoEFactory(
125 num_experts=num_experts,
126 top_k=top_k,
127 hidden_size=hidden_size,
vllm/model_executor/models/nemotron_h.py
+2 -2
@@ -33,7 +33,7 @@ from vllm.distributed.parallel_state import get_pp_group
33 from vllm.model_executor.layers.activation import ReLUSquaredActivation
34 from vllm.model_executor.layers.attention import Attention
35 from vllm.model_executor.layers.fused_moe import (
36 - FusedMoE,
36 + FusedMoEFactory,
37 GateLinear,
38 activation_without_mul,
39 fused_moe_make_expert_params_mapping,
@@ -209,7 +209,7 @@ class NemotronHMoE(nn.Module):
209 self.fc1_latent_proj = None
210 self.fc2_latent_proj = None
211
212 - self.experts = FusedMoE(
212 + self.experts = FusedMoEFactory(
213 shared_experts=self.shared_experts,
214 num_experts=config.n_routed_experts,
215 top_k=config.num_experts_per_tok,
vllm/model_executor/models/olmoe.py
+2 -4
@@ -32,9 +32,7 @@ from vllm.distributed import (
32 from vllm.distributed.utils import split_tensor_along_last_dim
33 from vllm.logger import init_logger
34 from vllm.model_executor.layers.attention import Attention
35 -from vllm.model_executor.layers.fused_moe import (
36 - FusedMoE,
37 -)
35 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
36 from vllm.model_executor.layers.layernorm import RMSNorm
37 from vllm.model_executor.layers.linear import (
38 QKVParallelLinear,
@@ -94,7 +92,7 @@ class OlmoeMoE(nn.Module):
92 prefix=f"{prefix}.gate",
93 )
94
97 - self.experts = FusedMoE(
95 + self.experts = FusedMoEFactory(
96 num_experts=num_experts,
97 top_k=top_k,
98 hidden_size=hidden_size,
vllm/model_executor/models/openpangu.py
+2 -2
@@ -43,7 +43,7 @@ from vllm.model_executor.layers.attention import (
43 Attention,
44 StaticSinkAttention,
45 )
46 -from vllm.model_executor.layers.fused_moe import FusedMoE
46 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
47 from vllm.model_executor.layers.layernorm import RMSNorm
48 from vllm.model_executor.layers.linear import (
49 ColumnParallelLinear,
@@ -195,7 +195,7 @@ class OpenPanguMoE(nn.Module):
195 else:
196 self.shared_experts = None
197
198 - self.experts = FusedMoE(
198 + self.experts = FusedMoEFactory(
199 shared_experts=self.shared_experts,
200 num_experts=config.n_routed_experts,
201 top_k=config.num_experts_per_tok,
vllm/model_executor/models/param2moe.py
+5 -3
@@ -32,7 +32,9 @@ from vllm.distributed import (
32 )
33 from vllm.model_executor.layers.activation import SiluAndMul
34 from vllm.model_executor.layers.attention import Attention
35 -from vllm.model_executor.layers.fused_moe import FusedMoE
35 +from vllm.model_executor.layers.fused_moe import (
36 + FusedMoEFactory,
37 +)
38 from vllm.model_executor.layers.layernorm import RMSNorm
39 from vllm.model_executor.layers.linear import (
40 MergedColumnParallelLinear,
@@ -325,7 +327,7 @@ class Param2MoEMoEBlock(nn.Module):
327 else:
328 self.shared_experts = None # type: ignore[assignment]
329
328 - self.experts = FusedMoE(
330 + self.experts = FusedMoEFactory(
331 shared_experts=self.shared_experts,
332 num_experts=self.num_experts,
333 top_k=self.top_k,
@@ -342,7 +344,7 @@ class Param2MoEMoEBlock(nn.Module):
344 routed_scaling_factor=self.routed_scaling_factor,
345 )
346
345 - def maybe_get_fused_moe(self) -> FusedMoE:
347 + def maybe_get_fused_moe(self) -> FusedMoEFactory:
348 return self.experts
349
350 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/phimoe.py
+2 -4
@@ -35,9 +35,7 @@ from vllm.compilation.decorators import support_torch_compile
35 from vllm.config import CacheConfig, VllmConfig
36 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37 from vllm.model_executor.layers.attention import Attention
38 -from vllm.model_executor.layers.fused_moe import (
39 - FusedMoE,
40 -)
38 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
39 from vllm.model_executor.layers.linear import (
40 QKVParallelLinear,
41 ReplicatedLinear,
@@ -273,7 +271,7 @@ class PhiMoE(nn.Module):
271 prefix=f"{prefix}.gate",
272 )
273
276 - self.experts = FusedMoE(
274 + self.experts = FusedMoEFactory(
275 num_experts=num_experts,
276 top_k=top_k,
277 hidden_size=hidden_size,
vllm/model_executor/models/qwen2_moe.py
+2 -4
@@ -40,9 +40,7 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
40 from vllm.logger import init_logger
41 from vllm.model_executor.layers.activation import SiluAndMul
42 from vllm.model_executor.layers.attention import Attention
43 -from vllm.model_executor.layers.fused_moe import (
44 - FusedMoE,
45 -)
43 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
44 from vllm.model_executor.layers.layernorm import RMSNorm
45 from vllm.model_executor.layers.linear import (
46 MergedColumnParallelLinear,
@@ -165,7 +163,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
163 else:
164 self.shared_expert = None
165
168 - self.experts = FusedMoE(
166 + self.experts = FusedMoEFactory(
167 shared_experts=self.shared_expert,
168 num_experts=config.num_experts,
169 top_k=config.num_experts_per_tok,
vllm/model_executor/models/qwen3_moe.py
+4 -4
@@ -42,7 +42,7 @@ from vllm.distributed import (
42 from vllm.logger import init_logger
43 from vllm.model_executor.layers.activation import SiluAndMul
44 from vllm.model_executor.layers.attention import Attention
45 -from vllm.model_executor.layers.fused_moe import FusedMoE
45 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
46 from vllm.model_executor.layers.layernorm import RMSNorm
47 from vllm.model_executor.layers.linear import (
48 MergedColumnParallelLinear,
@@ -201,7 +201,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
201 self.shared_expert_gate = None
202 self.shared_expert = None
203
204 - self.experts = FusedMoE(
204 + self.experts = FusedMoEFactory(
205 shared_experts=self.shared_expert,
206 gate=self.gate,
207 num_experts=self.n_routed_experts,
@@ -228,13 +228,13 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
228 hidden_states = sequence_parallel_chunk(hidden_states)
229
230 if self.experts.is_internal_router:
231 - # In this case, the gate/router runs inside the FusedMoE class
231 + # In this case, the gate/router runs inside the MoERunner class
232 final_hidden_states = self.experts(
233 hidden_states=hidden_states, router_logits=hidden_states
234 )
235 else:
236 # Actually this will be dead code, since we always pass gate into
237 - # FusedMoE in the current implementation. But we keep this code
237 + # MoERunner in the current implementation. But we keep this code
238 # here for clarity and future flexibility.
239 router_logits, _ = self.gate(hidden_states)
240 final_hidden_states = self.experts(
vllm/model_executor/models/qwen3_next.py
+3 -5
@@ -20,9 +20,7 @@ from vllm.distributed import (
20 )
21 from vllm.logger import init_logger
22 from vllm.model_executor.layers.attention import Attention
23 -from vllm.model_executor.layers.fused_moe import (
24 - FusedMoE,
25 -)
23 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
24 from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
25 from vllm.model_executor.layers.layernorm import (
26 GemmaRMSNorm as Qwen3NextRMSNorm,
@@ -175,7 +173,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
173 prefix=f"{prefix}.shared_expert",
174 )
175
178 - self.experts = FusedMoE(
176 + self.experts = FusedMoEFactory(
177 shared_experts=self.shared_expert,
178 gate=self.gate,
179 num_experts=self.n_routed_experts,
@@ -208,7 +206,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
206 hidden_states = sequence_parallel_chunk(hidden_states)
207
208 if self.experts.is_internal_router:
211 - # In this case, the gate/router runs inside the FusedMoE class
209 + # In this case, the gate/router runs inside the MoERunner class
210 final_hidden_states = self.experts(
211 hidden_states=hidden_states, router_logits=hidden_states
212 )
vllm/model_executor/models/sarvam.py
+2 -2
@@ -35,7 +35,7 @@ from vllm.distributed import (
35 get_tensor_model_parallel_world_size,
36 )
37 from vllm.model_executor.layers.activation import SiluAndMul
38 -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner
38 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
39 from vllm.model_executor.layers.layernorm import RMSNorm
40 from vllm.model_executor.layers.linear import (
41 ColumnParallelLinear,
@@ -334,7 +334,7 @@ class SarvamMLAMoE(nn.Module):
334 else:
335 self.shared_experts = None
336
337 - self.experts = FusedMoE(
337 + self.experts = FusedMoEFactory(
338 shared_experts=self.shared_experts,
339 num_experts=self.num_experts,
340 top_k=self.top_k,
vllm/model_executor/models/step3_text.py
+2 -2
@@ -19,7 +19,7 @@ from vllm.logger import init_logger
19 from vllm.model_executor.layers.activation import SiluAndMul
20 from vllm.model_executor.layers.attention import Attention
21 from vllm.model_executor.layers.fused_moe import (
22 - FusedMoE,
22 + FusedMoEFactory,
23 )
24 from vllm.model_executor.layers.layernorm import RMSNorm
25 from vllm.model_executor.layers.linear import (
@@ -68,7 +68,7 @@ class FusedMoEBlock(nn.Module):
68 f"the number of experts {config.moe_num_experts}."
69 )
70
71 - self.experts = FusedMoE(
71 + self.experts = FusedMoEFactory(
72 num_experts=config.moe_num_experts,
73 top_k=config.moe_top_k,
74 hidden_size=config.hidden_size,
vllm/model_executor/models/step3p5.py
+3 -3
@@ -24,7 +24,7 @@ from vllm.logger import init_logger
24 from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul
25 from vllm.model_executor.layers.attention import Attention
26 from vllm.model_executor.layers.fused_moe import (
27 - FusedMoE,
27 + FusedMoEFactory,
28 MoERunner,
29 fused_moe_make_expert_params_mapping,
30 )
@@ -376,7 +376,7 @@ class FusedMoEBlock(nn.Module):
376 quant_config=quant_config,
377 prefix=f"{prefix}.share_expert",
378 )
379 - self.experts = FusedMoE(
379 + self.experts = FusedMoEFactory(
380 shared_experts=self.share_expert,
381 gate=self.gate,
382 num_experts=config.moe_num_experts,
@@ -404,7 +404,7 @@ class FusedMoEBlock(nn.Module):
404 hidden_states=hidden_states, router_logits=hidden_states
405 )
406 else:
407 - # TODO(bnell): this gate could be moved into the FusedMoE?
407 + # TODO(bnell): this gate could be moved into the MoERunner?
408 router_logits, _ = self.gate(hidden_states)
409 final_hidden_states = self.experts(
410 hidden_states=hidden_states, router_logits=router_logits
vllm/model_executor/models/transformers/moe.py
+10 -6
@@ -29,7 +29,11 @@ from vllm.distributed import get_dp_group, get_ep_group
29 from vllm.forward_context import ForwardContext, get_forward_context
30 from vllm.logger import init_logger
31 from vllm.model_executor.custom_op import PluggableLayer
32 -from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts
32 +from vllm.model_executor.layers.fused_moe import (
33 + FusedMoEFactory,
34 + MoERunner,
35 + RoutedExperts,
36 +)
37 from vllm.model_executor.models.interfaces import MixtureOfExperts
38 from vllm.model_executor.models.transformers.fuser import get_fuser
39 from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
@@ -53,7 +57,7 @@ class TransformersMoEState:
57 # --8<-- [start:transformers_fused_moe]
58 @PluggableLayer.register("transformers_fused_moe")
59 class TransformersMoERunner(MoERunner):
56 - """Custom FusedMoE for the Transformers modeling backend."""
60 + """Custom MoERunner for the Transformers modeling backend."""
61
62 # --8<-- [end:transformers_fused_moe]
63 def __init__(self, *args, moe_state: TransformersMoEState, **kwargs):
@@ -297,7 +301,7 @@ class MoEMixin(MixtureOfExperts):
301 )
302 if fuser is not None and not reaches_fp16_trick:
303 # MoE block forward is fully replaced.
300 - # gate/router and shared expert (if any) runs in FusedMoE.
304 + # gate/router and shared expert (if any) runs in MoERunner.
305 shared_experts = fuser.shared_experts(moe_block, prefix)
306 # Store shared experts for later down projection adjustment
307 if shared_experts is not None:
@@ -331,7 +335,7 @@ class MoEMixin(MixtureOfExperts):
335 if fuser.shared_name:
336 routed += " + shared experts"
337 logger.info_once(
334 - "Fused: %s (%s) -> FusedMoE (internal routing)",
338 + "Fused: %s (%s) -> MoERunner (internal routing)",
339 routed,
340 moe_block_cls,
341 )
@@ -373,10 +377,10 @@ class MoEMixin(MixtureOfExperts):
377 runner_args={"moe_state": moe_state},
378 )
379 logger.info_once(
376 - "Fused: experts (%s) -> FusedMoE (external routing)",
380 + "Fused: experts (%s) -> MoERunner (external routing)",
381 experts_cls,
382 )
379 - fused_experts = FusedMoE(**kwargs)
383 + fused_experts = FusedMoEFactory(**kwargs)
384 moe_block.experts = fused_experts
385 log_replacement(qual_name, experts, fused_experts)
386 # Update MixtureOfExperts mixin state
vllm/model_executor/warmup/deep_gemm_warmup.py
+1 -1
@@ -108,7 +108,7 @@ def _extract_data_from_fused_moe_module(
108 m_: torch.nn.Module,
109 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
110 """
111 - Extract weights, weight scales and num_topk from FusedMoE module.
111 + Extract weights, weight scales and num_topk from MoERunner module.
112 """
113 assert isinstance(m_, MoERunner)
114 m = m_.routed_experts
vllm/models/deepseek_v4/amd/model.py
+3 -3
@@ -18,7 +18,7 @@ from vllm.distributed import (
18 )
19 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
20 from vllm.model_executor.layers.fused_moe import (
21 - FusedMoE,
21 + FusedMoEFactory,
22 GateLinear,
23 fused_moe_make_expert_params_mapping,
24 )
@@ -279,7 +279,7 @@ class DeepseekV4MoE(nn.Module):
279 self.experts_start_idx = self.tp_rank * self.n_local_experts
280 self.experts_end_idx = self.experts_start_idx + self.n_local_experts
281
282 - self.experts = FusedMoE(
282 + self.experts = FusedMoEFactory(
283 shared_experts=self.shared_experts,
284 n_shared_experts=(
285 config.n_shared_experts if self.fuse_shared_experts else None
@@ -308,7 +308,7 @@ class DeepseekV4MoE(nn.Module):
308
309 org_shape = hidden_states.shape
310 if self.experts.is_internal_router:
311 - # In this case, the gate/router runs inside the FusedMoE class
311 + # In this case, the gate/router runs inside the MoERunner class
312 final_hidden_states = self.experts(
313 hidden_states=hidden_states,
314 router_logits=hidden_states,
vllm/models/deepseek_v4/amd/mtp.py
+1 -1
@@ -51,7 +51,7 @@ from .model import DeepseekV4DecoderLayer
51 logger = init_logger(__name__)
52
53 # MoE expert scales are fused into per-layer w13/w2 tensors. The exact
54 -# parameter suffix depends on which FusedMoE method handles the experts:
54 +# parameter suffix depends on which MoERunner method handles the experts:
55 # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
56 # - fp8 experts (Fp8MoEMethod with block_quant=True) register
57 # ``w{1,2,3}_weight_scale_inv``.
vllm/models/deepseek_v4/nvidia/model.py
+3 -3
@@ -27,7 +27,7 @@ from vllm.model_executor.kernels.mhc.tilelang import (
27 )
28 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
29 from vllm.model_executor.layers.fused_moe import (
30 - FusedMoE,
30 + FusedMoEFactory,
31 fused_moe_make_expert_params_mapping,
32 )
33 from vllm.model_executor.layers.fused_moe.router.base_router import (
@@ -670,7 +670,7 @@ class DeepseekV4MoE(nn.Module):
670 self.physical_expert_start = self.experts_start_idx
671 self.physical_expert_end = self.experts_end_idx
672
673 - self.experts = FusedMoE(
673 + self.experts = FusedMoEFactory(
674 shared_experts=self.shared_experts,
675 gate=self.gate,
676 num_experts=config.n_routed_experts,
@@ -736,7 +736,7 @@ class DeepseekV4MoE(nn.Module):
736 ) -> torch.Tensor:
737 org_shape = hidden_states.shape
738 if self.experts.is_internal_router:
739 - # In this case, the gate/router runs inside the FusedMoE class
739 + # In this case, the gate/router runs inside the MoERunner class
740 final_hidden_states = self.experts(
741 hidden_states=hidden_states,
742 router_logits=hidden_states,
vllm/models/deepseek_v4/nvidia/mtp.py
+1 -1
@@ -59,7 +59,7 @@ from .model import (
59 logger = init_logger(__name__)
60
61 # MoE expert scales are fused into per-layer w13/w2 tensors. The exact
62 -# parameter suffix depends on which FusedMoE method handles the experts:
62 +# parameter suffix depends on which FusedMoEFactory method handles the experts:
63 # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
64 # - fp8 experts (Fp8MoEMethod with block_quant=True) register
65 # ``w{1,2,3}_weight_scale_inv``.
vllm/models/deepseek_v4/xpu/model.py
+2 -2
@@ -18,7 +18,7 @@ from vllm.distributed import (
18 from vllm.forward_context import get_forward_context
19 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
20 from vllm.model_executor.layers.fused_moe import (
21 - FusedMoE,
21 + FusedMoEFactory,
22 GateLinear,
23 fused_moe_make_expert_params_mapping,
24 )
@@ -729,7 +729,7 @@ class DeepseekV4MoE(nn.Module):
729 self.n_local_experts = config.n_routed_experts // self.tp_size
730 self.experts_start_idx = self.tp_rank * self.n_local_experts
731 self.experts_end_idx = self.experts_start_idx + self.n_local_experts
732 - self.experts = FusedMoE(
732 + self.experts = FusedMoEFactory(
733 shared_experts=self.shared_experts,
734 gate=self.gate,
735 num_experts=config.n_routed_experts,
vllm/models/deepseek_v4/xpu/mtp.py
+1 -1
@@ -54,7 +54,7 @@ from .model import (
54 logger = init_logger(__name__)
55
56 # MoE expert scales are fused into per-layer w13/w2 tensors. The exact
57 -# parameter suffix depends on which FusedMoE method handles the experts:
57 +# parameter suffix depends on which FusedMoEFactory method handles the experts:
58 # - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
59 # - fp8 experts (Fp8MoEMethod with block_quant=True) register
60 # ``w{1,2,3}_weight_scale_inv``.
vllm/models/inkling/amd/model.py
+1 -1
@@ -646,7 +646,7 @@ def _load_inkling_weights(
646 weight.shard_id = shard_id
647
648 # MoE expert tensors (fused stacked, routed + shared sink): translate
649 - # the checkpoint layout to per-expert FusedMoE loads.
649 + # the checkpoint layout to per-expert MoERunner loads.
650 moe_match = _MOE_EXPERT_WEIGHT_RE.match(name)
651 if moe_match is not None and moe_match.group("mlp") in moe_modules:
652 moe = moe_modules[moe_match.group("mlp")]
vllm/models/inkling/amd/moe.py
+8 -8
@@ -1,11 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Inkling mixture-of-experts on vLLM's FusedMoE abstraction.
3 +"""Inkling mixture-of-experts on vLLM's MoERunner abstraction.
4
5 Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over
6 the routed experts, log-sigmoid renormalization over the k routed + S shared
7 "sink" logits, scaled by route_scale * global_scale. The routed top-k goes
8 -through vLLM's FusedMoE (which handles TP/EP); the sink experts run in
8 +through vLLM's MoERunner (which handles TP/EP); the sink experts run in
9 :class:`InklingSinkExperts` -- replicated across EP ranks (every token
10 activates every sink) and always bf16 (the checkpoint excludes every
11 ``shared_experts`` from quantization).
@@ -35,7 +35,7 @@ from vllm.distributed import (
35 get_tensor_model_parallel_world_size,
36 )
37 from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16
38 -from vllm.model_executor.layers.fused_moe import FusedMoE
38 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
39 from vllm.model_executor.layers.quantization import QuantizationConfig
40 from vllm.model_executor.utils import set_weight_attrs
41 from vllm.platforms import current_platform
@@ -265,7 +265,7 @@ class InklingGate(nn.Module):
265
266
267 def _inkling_moe_ep_size() -> int:
268 - """EP size the FusedMoE layer will run with (mirrors
268 + """EP size the MoERunner layer will run with (mirrors
269 FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when
270 expert parallelism is enabled)."""
271 parallel_config = get_current_vllm_config().parallel_config
@@ -444,7 +444,7 @@ class InklingMoE(nn.Module):
444
445 # The released MXFP4 checkpoint keeps the first routed-expert layer
446 # in bf16 and lists its two expert weights explicitly in Quark's
447 - # exclusion list. FusedMoE asks for a quant method at the module
447 + # exclusion list. FusedMoEFactory asks for a quant method at the module
448 # prefix, so an exact weight exclusion would otherwise be missed and
449 # the bf16 tensors would be loaded into MXFP4 parameters.
450 routed_quant_config = quant_config
@@ -458,7 +458,7 @@ class InklingMoE(nn.Module):
458 if routed_weights <= excluded:
459 routed_quant_config = None
460
461 - self.experts = FusedMoE(
461 + self.experts = FusedMoEFactory(
462 num_experts=num_experts,
463 top_k=config.num_experts_per_tok,
464 hidden_size=config.hidden_size,
@@ -503,7 +503,7 @@ class InklingMoE(nn.Module):
503 topk: int,
504 renormalize: bool,
505 ) -> tuple[torch.Tensor, torch.Tensor]:
506 - """FusedMoE ``custom_routing_function``: the routed top-k slice of the
506 + """MoERunner ``custom_routing_function``: the routed top-k slice of the
507 full (routed + sink) selection.
508
509 forward() stashes its selection (keyed by logits identity) so the
@@ -523,7 +523,7 @@ class InklingMoE(nn.Module):
523 router_logits = self.gate.compute_logits(x)
524 num_tokens = x.shape[0]
525 # One gate select per layer: the routed slice is stashed for the
526 - # routing function inside the FusedMoE op; the sink gammas are the
526 + # routing function inside the MoERunner op; the sink gammas are the
527 # trailing columns.
528 k = self.gate.topk
529 weights, ids = self.gate.select_experts(router_logits)
vllm/models/inkling/nvidia/moe.py
+2 -2
@@ -34,7 +34,7 @@ from vllm.distributed import (
34 get_tensor_model_parallel_world_size,
35 )
36 from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16
37 -from vllm.model_executor.layers.fused_moe import FusedMoE
37 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
38 from vllm.model_executor.utils import set_weight_attrs
39 from vllm.platforms import current_platform
40 from vllm.triton_utils import tl, tldevice, triton
@@ -439,7 +439,7 @@ class InklingMoE(nn.Module):
439 # power-of-two EP sizes (n_routed is a power of two).
440 num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size()
441
442 - self.experts = FusedMoE(
442 + self.experts = FusedMoEFactory(
443 num_experts=num_experts,
444 top_k=config.num_experts_per_tok,
445 hidden_size=config.hidden_size,
vllm/models/kimi_k3/amd/linear.py
+2 -2
@@ -15,7 +15,7 @@ from vllm.distributed import (
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul
17 from vllm.model_executor.layers.fused_moe import (
18 - FusedMoE,
18 + FusedMoEFactory,
19 fused_moe_make_expert_params_mapping,
20 )
21 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
@@ -260,7 +260,7 @@ class KimiMoE(nn.Module):
260 self.routed_expert_up_proj = None
261 self.routed_output_transform = None
262
263 - self.experts = FusedMoE(
263 + self.experts = FusedMoEFactory(
264 shared_experts=self.shared_experts,
265 num_experts=num_experts,
266 top_k=num_experts_per_token,
vllm/models/kimi_k3/nvidia/model.py
+4 -4
@@ -20,7 +20,7 @@ from vllm.forward_context import get_forward_context, is_forward_context_availab
20 from vllm.logger import init_logger
21 from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul
22 from vllm.model_executor.layers.fused_moe import (
23 - FusedMoE,
23 + FusedMoEFactory,
24 fused_moe_make_expert_params_mapping,
25 )
26 from vllm.model_executor.layers.fused_moe.router.base_router import (
@@ -560,7 +560,7 @@ class KimiMoE(nn.Module):
560 current_platform.is_cuda()
561 and current_platform.is_device_capability_family(100)
562 )
563 - self.experts = FusedMoE(
563 + self.experts = FusedMoEFactory(
564 shared_experts=self.shared_experts,
565 num_experts=num_experts,
566 top_k=num_experts_per_token,
@@ -578,7 +578,7 @@ class KimiMoE(nn.Module):
578 scoring_func=config.moe_router_activation_func,
579 e_score_correction_bias=self.gate.e_score_correction_bias,
580 routed_scaling_factor=self.routed_scaling_factor,
581 - # Down projection runs outside FusedMoE so it can overlap the
581 + # Down projection runs outside MoERunner so it can overlap the
582 # router gate on the aux stream (see forward()); the original
583 # hidden states are passed to forward() as shared_experts_input
584 # so shared experts still see the untransformed input.
@@ -689,7 +689,7 @@ class KimiMoE(nn.Module):
689 )
690 else:
691 # Routed experts consume the down-projected latent; shared experts
692 - # (inside FusedMoE) get the original hidden states via
692 + # (inside MoERunner) get the original hidden states via
693 # shared_experts_input.
694 final_hidden_states = self.experts(
695 hidden_states=routed_hidden_states,
vllm/models/minimax_m3/amd/model.py
+2 -2
@@ -41,7 +41,7 @@ from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
41 fused_allreduce_gemma_rms_norm,
42 )
43 from vllm.model_executor.layers.fused_moe import (
44 - FusedMoE,
44 + FusedMoEFactory,
45 GateLinear,
46 fused_moe_make_expert_params_mapping,
47 )
@@ -391,7 +391,7 @@ class MiniMaxM3MoE(nn.Module):
391 # always-on shared expert; aiter applies the routed scaling internally.
392 # Every other path (vLLM top-k bias router, or no fusion) applies the
393 # routed scaling to the MoE output here.
394 - self.experts = FusedMoE(
394 + self.experts = FusedMoEFactory(
395 num_experts=config.num_local_experts,
396 top_k=config.num_experts_per_tok,
397 hidden_size=config.hidden_size,
vllm/models/minimax_m3/nvidia/model.py
+2 -2
@@ -31,7 +31,7 @@ from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
31 fused_allreduce_gemma_rms_norm,
32 )
33 from vllm.model_executor.layers.fused_moe import (
34 - FusedMoE,
34 + FusedMoEFactory,
35 GateLinear,
36 fused_moe_make_expert_params_mapping,
37 )
@@ -244,7 +244,7 @@ class MiniMaxM3MoE(nn.Module):
244 prefix=f"{prefix}.shared_experts",
245 )
246
247 - self.experts = FusedMoE(
247 + self.experts = FusedMoEFactory(
248 num_experts=config.num_local_experts,
249 top_k=config.num_experts_per_tok,
250 hidden_size=config.hidden_size,