@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 726efe177b

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 22:30 UTC

[MoE Refactor] Move the shared/fused expert output sum into MoERunnerBase (#35949)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed Apr 20, 2026 at 12:28 UTC 726efe177bf22874743d11dfdfef9247dbfb5ff0

52 files changed +322 -697

tests/kernels/moe/test_moe_layer.py

+17 -76

@@ -236,7 +236,6 @@ class MoETestConfig:

236 use_gate: bool

237 use_routed_input_transform: bool

238 enable_eplb: bool = False

239 - reduce_results: bool = False

239 backend: str | None = None

240 ep_size: int = 1

241 dp_size: int = 1

@@ -295,7 +294,6 @@ def generate_valid_test_configs(

294 use_shared_experts,

295 use_gate,

296 use_routed_input_transform,

298 - reduce_results,

297 ) in product(

298 SHAPE_COMBOS,

299 NUM_EXPERTS,

@@ -304,7 +302,6 @@ def generate_valid_test_configs(

302 [False, True], # shared

303 [False, True], # gate

304 [False, True], # routed input exform

307 - [False, True], # reduce results

305 ):

306 config = MoETestConfig(

307 shape[0], # m

@@ -318,7 +315,6 @@ def generate_valid_test_configs(

315 use_gate,

316 use_routed_input_transform,

317 enable_eplb,

321 - reduce_results,

318 backend,

319 ep_size,

320 dp_size,

@@ -395,18 +391,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

391 and config.backend.startswith("flashinfer_nvlink")

392 and not current_platform.has_device_capability(90)

393 ):

398 - return False, "flashinfer_nvlink needs an H100+ GPUs"

399 -

400 - # reduce_results incompatibilities

401 - if config.reduce_results and config.use_shared_experts:

402 - return False, "reduce_results=True is not compatible with shared_experts=True"

403 -

404 - if config.reduce_results and config.quantization is not None:

405 - return (

406 - False,

407 - "reduce_results=True only tested with unquantized data types in "

408 - "order to limit number of tests run",

409 - )

394 + return False, "flashinfer_nvlink needs H100+ GPUs"

395

396 # Backend-specific checks

397 if config.backend is not None:

@@ -448,10 +433,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

433 if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS:

434 return False, f"EPLB not supported with {config.backend}."

435

451 - world_size = config.tp_size * config.dp_size

452 - if config.reduce_results and world_size == 1:

453 - return False, "reduce_results=True only makes sense for multi-GPU tests"

454 -

436 if (

437 config.backend is not None

438 and config.backend.startswith("flashinfer_nvlink")

@@ -846,7 +827,6 @@ def make_fused_moe_layer(

827 tp_size: int,

828 ep_size: int,

829 dp_size: int,

849 - reduce_results: bool,

830 w1: torch.Tensor,

831 w2: torch.Tensor,

832 top_k: int,

@@ -874,7 +854,7 @@ def make_fused_moe_layer(

854 routed_input_transform: torch.nn.Module | None = None,

855 routed_output_transform: torch.nn.Module | None = None,

856 pcp_size: int | None = 1,

877 -) -> tuple[Callable, FusedMoE]:

857 +) -> FusedMoE:

858 quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)

859

860 kwargs = dict()

@@ -887,8 +867,10 @@ def make_fused_moe_layer(

867 # Add gate and routed_input_transform if provided

868 if gate is not None:

869 kwargs["gate"] = gate

870 +

871 if routed_input_transform is not None:

872 kwargs["routed_input_transform"] = routed_input_transform

873 + kwargs["routed_output_transform"] = routed_output_transform

874

875 layer = builder(

876 num_experts=global_num_experts,

@@ -896,7 +878,6 @@ def make_fused_moe_layer(

878 hidden_size=hidden_size,

879 intermediate_size=intermediate_size,

880 params_dtype=in_dtype,

899 - reduce_results=reduce_results,

881 renormalize=renormalize,

882 use_grouped_topk=use_grouped_topk,

883 num_expert_group=num_expert_group,

@@ -936,36 +917,7 @@ def make_fused_moe_layer(

917

918 layer.quant_method.process_weights_after_loading(layer)

919

939 - def _moe(

940 - hidden_states: torch.Tensor,

941 - router_logits: torch.Tensor,

942 - ) -> torch.Tensor:

943 - if shared_experts is None:

944 - final_shared_states = None

945 - final_hidden_states = layer(hidden_states, router_logits)

946 - else:

947 - final_shared_states, final_hidden_states = layer(

948 - hidden_states, router_logits

949 - )

950 -

951 - # Apply routed output transform if provided

952 - # (e.g., latent space -> original space)

953 - if routed_output_transform is not None:

954 - final_hidden_states = routed_output_transform(final_hidden_states)

955 -

956 - if shared_experts is not None:

957 - assert not reduce_results

958 - assert final_shared_states is not None

959 - final_hidden_states += final_shared_states

960 -

961 - if not reduce_results and layer.tp_size > 1:

962 - final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel(

963 - final_hidden_states

964 - )

965 -

966 - return final_hidden_states

967 -

968 - return _moe, layer

920 + return layer

921

922

923 def make_fake_moe_layer(

@@ -999,7 +951,6 @@ def make_fake_moe_layer(

951 tp_size: int = 1,

952 dp_size: int = 1,

953 ep_size: int = 1,

1002 - reduce_results: bool = False,

954 ) -> Callable:

955 activation = MoEActivation.from_str(activation)

956

@@ -1101,7 +1052,7 @@ def make_fake_moe_layer(

1052

1053

1054 def _test_body_regular(

1104 - moe_fn: Callable,

1055 + moe_layer: Callable,

1056 hidden_states: torch.Tensor,

1057 router_logits: torch.Tensor,

1058 vllm_config: VllmConfig,

@@ -1118,13 +1069,12 @@ def _test_body_regular(

1069 num_tokens=num_tokens,

1070 num_tokens_across_dp=num_tokens_across_dp,

1071 ):

1121 - output = moe_fn(hidden_states, router_logits)

1072 + output = moe_layer(hidden_states, router_logits)

1073

1074 return baseline_output, output

1075

1076

1077 def _test_body_eplb(

1127 - moe_fn: Callable,

1078 moe_layer: FusedMoE,

1079 hidden_states: torch.Tensor,

1080 router_logits: torch.Tensor,

@@ -1145,7 +1095,6 @@ def _test_body_eplb(

1095 n: int,

1096 top_k: int,

1097 shared_experts,

1148 - reduce_results: bool,

1098 gate: torch.nn.Module | None,

1099 routed_input_transform: torch.nn.Module | None,

1100 routed_output_transform: torch.nn.Module | None,

@@ -1161,7 +1110,7 @@ def _test_body_eplb(

1110 num_tokens=num_tokens,

1111 num_tokens_across_dp=num_tokens_across_dp,

1112 ):

1164 - output_before = moe_fn(hidden_states, router_logits)

1113 + output_before = moe_layer(hidden_states, router_logits)

1114

1115 # Create a fresh FusedMoE layer with enable_eplb=True

1116 # Delete the original layer's registration so the constructor can

@@ -1174,7 +1123,7 @@ def _test_body_eplb(

1123 # When using routed_input_transform, experts operate in latent space

1124 hidden_size_for_layer = k // 2 if routed_input_transform is not None else k

1125

1177 - moe_fn, moe_layer = make_fused_moe_layer(

1126 + eplb_moe_layer = make_fused_moe_layer(

1127 quantization=quantization,

1128 use_ep=use_ep,

1129 hidden_size=hidden_size_for_layer,

@@ -1183,7 +1132,6 @@ def _test_body_eplb(

1132 tp_size=tp_size,

1133 ep_size=ep_size,

1134 dp_size=dp_size,

1186 - reduce_results=reduce_results,

1135 w1=w1,

1136 w2=w2,

1137 top_k=top_k,

@@ -1196,14 +1144,14 @@ def _test_body_eplb(

1144 )

1145

1146 # Necessary?

1199 - if moe_layer._expert_map is not None:

1200 - moe_layer._expert_map = moe_layer._expert_map.to(device)

1147 + if eplb_moe_layer._expert_map is not None:

1148 + eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device)

1149

1150 # All ranks must generate the same permutation

1151 initial_indices = torch.arange(num_experts, dtype=torch.long)

1152 shuffled_indices = initial_indices[torch.randperm(num_experts)]

1153

1206 - expert_weights = [list(moe_layer.get_expert_weights())]

1154 + expert_weights = [list(eplb_moe_layer.get_expert_weights())]

1155

1156 communicator = create_eplb_communicator(

1157 group_coordinator=get_eplb_group(),

@@ -1227,7 +1175,7 @@ def _test_body_eplb(

1175 num_experts, dtype=torch.int32, device=device

1176 )

1177

1230 - moe_layer.set_eplb_state(

1178 + eplb_moe_layer.set_eplb_state(

1179 moe_layer_idx=0,

1180 expert_load_view=torch.zeros(

1181 (1, num_experts),

@@ -1244,7 +1192,7 @@ def _test_body_eplb(

1192 ),

1193 )

1194

1247 - moe_layer.eplb_state.should_record_tensor = torch.ones(

1195 + eplb_moe_layer.eplb_state.should_record_tensor = torch.ones(

1196 (), dtype=torch.bool, device=device

1197 )

1198

@@ -1255,7 +1203,7 @@ def _test_body_eplb(

1203 num_tokens=num_tokens,

1204 num_tokens_across_dp=num_tokens_across_dp,

1205 ):

1258 - output_after = moe_fn(hidden_states, router_logits)

1206 + output_after = eplb_moe_layer(hidden_states, router_logits)

1207

1208 return output_before, output_after

1209

@@ -1274,7 +1222,6 @@ def _run_one_config(

1222 num_experts: int,

1223 top_k: int,

1224 quantization: str | None,

1277 - reduce_results: bool,

1225 backend: str | None,

1226 test_body_fn: Callable,

1227 use_shared_experts: bool,

@@ -1341,7 +1288,6 @@ def _run_one_config(

1288 tp_size=tp_size,

1289 ep_size=ep_size,

1290 dp_size=dp_size,

1344 - reduce_results=reduce_results,

1291 )

1292

1293 baseline_output = baseline_layer(hidden_states, router_logits)

@@ -1369,7 +1315,7 @@ def _run_one_config(

1315 hidden_size_for_layer = k // 2 if routed_input_transform is not None else k

1316

1317 # Create initial MoE layer

1372 - moe_fn, moe_layer = make_fused_moe_layer(

1318 + moe_layer = make_fused_moe_layer(

1319 quantization=quantization,

1320 use_ep=use_ep,

1321 hidden_size=hidden_size_for_layer,

@@ -1378,7 +1324,6 @@ def _run_one_config(

1324 tp_size=tp_size,

1325 ep_size=ep_size,

1326 dp_size=dp_size,

1381 - reduce_results=reduce_results,

1327 w1=w1,

1328 w2=w2,

1329 top_k=top_k,

@@ -1402,7 +1347,6 @@ def _run_one_config(

1347

1348 # Call the test body function with all necessary context

1349 expected, actual = test_body_fn(

1405 - moe_fn=moe_fn,

1350 moe_layer=moe_layer,

1351 hidden_states=hidden_states,

1352 router_logits=router_logits,

@@ -1423,7 +1367,6 @@ def _run_one_config(

1367 m=m,

1368 top_k=top_k,

1369 shared_experts=shared_experts,

1426 - reduce_results=reduce_results,

1370 gate=gate,

1371 routed_input_transform=routed_input_transform,

1372 routed_output_transform=routed_output_transform,

@@ -1520,7 +1463,6 @@ def test_moe_layer_no_parallel(

1463 test_config.num_experts,

1464 test_config.top_k,

1465 test_config.quantization,

1523 - test_config.reduce_results,

1466 test_config.backend,

1467 _test_body_regular,

1468 use_shared_experts=test_config.use_shared_experts,

@@ -1578,7 +1520,6 @@ def _parallel_worker(

1520 test_config.num_experts,

1521 test_config.top_k,

1522 test_config.quantization,

1581 - test_config.reduce_results,

1523 test_config.backend,

1524 functools.partial(

1525 _test_body_config, test_config=test_config, cpu_group=cpu_group

@@ -1597,7 +1538,7 @@ def _parallel_worker(

1538 failed = failed + 1

1539 if verbosity > 0:

1540 traceback.print_exc()

1600 - print(f"\n{str(ex)}\nFAILED {ex.__class__}")

1541 + print(f"\n{str(ex)}\nFAILED")

1542 else:

1543 print("F", end="")

1544 finally:

tests/kernels/moe/test_shared_fused_moe_routed_transform.py

+10 -26

@@ -165,7 +165,6 @@ def test_routed_input_transform_inside_vs_outside(

165 top_k=top_k,

166 hidden_size=latent_size,

167 intermediate_size=intermediate_size,

168 - reduce_results=False,

168 renormalize=True,

169 params_dtype=dtype,

170 tp_size=1,

@@ -183,7 +182,6 @@ def test_routed_input_transform_inside_vs_outside(

182 top_k=top_k,

183 hidden_size=latent_size,

184 intermediate_size=intermediate_size,

186 - reduce_results=False,

185 renormalize=True,

186 params_dtype=dtype,

187 tp_size=1,

@@ -212,34 +210,20 @@ def test_routed_input_transform_inside_vs_outside(

210 hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype)

211 router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)

212

215 - # Clone inputs so any in-place modification by Method A

216 - # cannot affect Method B's computation.

217 - hidden_states_A = hidden_states.clone()

218 - router_logits_A = router_logits.clone()

219 -

213 with set_forward_context(None, vllm_config, num_tokens=num_tokens):

221 - shared_out_A, routed_out_A = moe_with_transform(

222 - hidden_states_A, router_logits_A

223 - )

214 + # Method A: combined output (shared + routed)

215 + combined_A = moe_with_transform(hidden_states, router_logits)

216

217 + # Method B: manually transform, get routed output, add shared

218 transformed_hidden = routed_transform(hidden_states)

226 - shared_out_B, routed_out_B = moe_without_transform(

227 - transformed_hidden, router_logits

228 - )

229 -

230 - expected_shared_out = shared_experts(hidden_states)

219 + routed_out_B = moe_without_transform(transformed_hidden, router_logits)

220 + shared_out_B = shared_experts(hidden_states)

221 + combined_B = shared_out_B + routed_out_B

222

232 - _assert_close(

233 - routed_out_A,

234 - routed_out_B,

235 - atol=1e-3,

236 - rtol=1e-3,

237 - label="Routed output: transform inside vs outside",

238 - )

239 - _assert_close(

240 - shared_out_A,

241 - expected_shared_out,

223 + torch.testing.assert_close(

224 + combined_A,

225 + combined_B,

226 atol=1e-3,

227 rtol=1e-3,

244 - label="Shared expert output",

228 + msg="Combined output should match: transform inside vs outside",

229 )

vllm/lora/layers/fused_moe.py

-3

@@ -592,9 +592,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

592 def forward(self, *args, **kwargs):

593 return self.base_layer.forward(*args, **kwargs)

594

595 - def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs):

596 - return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs)

597 -

595 @property

596 def quant_method(self):

597 return self.base_layer.quant_method

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+1 -1

@@ -716,7 +716,7 @@ class MarlinExperts(MarlinExpertsBase):

716 ):

717 assert self.w1_scale is not None

718 assert self.w2_scale is not None

719 - return fused_marlin_moe(

719 + fused_marlin_moe(

720 hidden_states=hidden_states,

721 w1=w1,

722 w2=w2,

vllm/model_executor/layers/fused_moe/layer.py

+28 -32

@@ -230,11 +230,18 @@ class FusedMoE(PluggableLayer):

230 hidden_size: Input hidden state size of the transformer

231 intermediate_size: Intermediate size of the experts

232 params_dtype: Data type for the parameters.

233 - reduce_results: Whether to all_reduce on the output of the layer

233 renormalize: Whether to renormalize the logits in the fused_moe kernel

234 quant_config: Quantization configure.

235 enable_eplb: Whether to enable expert parallelism load balancer.

236 router_logits_dtype: Data type for router logits buffers.

237 + routed_scaling_factor: A scaling factor that is applied to the topk_weights

238 + by the router or the output of the layer depending

239 + on the value of `apply_routed_scale_to_output`

240 + apply_routed_scale_to_output: Determine whether or not `routed_scaling_factor`

241 + is applied to the topk_weights or to the experts

242 + output. It is applied to the experts output

243 + instead of the topk_weights when this feature is

244 + not supported by the router (or the experts).

245 """

246

247 # --8<-- [end:fused_moe]

@@ -246,7 +253,6 @@ class FusedMoE(PluggableLayer):

253 hidden_size: int,

254 intermediate_size: int,

255 params_dtype: torch.dtype | None = None,

249 - reduce_results: bool = False,

256 renormalize: bool = True,

257 use_grouped_topk: bool = False,

258 num_expert_group: int | None = None,

@@ -274,12 +280,12 @@ class FusedMoE(PluggableLayer):

280 gate: torch.nn.Module | None = None,

281 shared_experts: torch.nn.Module | None = None,

282 routed_input_transform: torch.nn.Module | None = None,

283 + routed_output_transform: torch.nn.Module | None = None,

284 + apply_routed_scale_to_output: bool = False,

285 zero_expert_type: str | None = None,

286 ):

287 super().__init__()

288

281 - self._routed_input_transform = routed_input_transform

282 -

289 if params_dtype is None:

290 params_dtype = torch.get_default_dtype()

291 self.params_dtype = params_dtype

@@ -425,7 +431,6 @@ class FusedMoE(PluggableLayer):

431

432 assert intermediate_size % self.tp_size == 0

433 intermediate_size_per_partition = intermediate_size // self.tp_size

428 - self.reduce_results = reduce_results

434 self.renormalize = renormalize

435

436 # TODO(bnell): these attributes are only used by monolithic kernels.

@@ -437,7 +442,14 @@ class FusedMoE(PluggableLayer):

442 self.topk_group = topk_group

443 self.custom_routing_function = custom_routing_function

444 self.scoring_func = scoring_func

440 - self.routed_scaling_factor = routed_scaling_factor

445 + # When apply_routed_scale_to_output is True, we set the scaling factor

446 + # to 1.0 so it ends up being a nop. Applying the scale will be handled

447 + # by the runner in this case.

448 + # The member variable must be set in the same way as the router since

449 + # some quantization methods can access it.

450 + self.routed_scaling_factor = (

451 + routed_scaling_factor if not apply_routed_scale_to_output else 1.0

452 + )

453 self.e_score_correction_bias = e_score_correction_bias

454 # TODO(bnell): end attributes

455

@@ -456,7 +468,7 @@ class FusedMoE(PluggableLayer):

468 topk_group=topk_group,

469 custom_routing_function=custom_routing_function,

470 scoring_func=scoring_func,

459 - routed_scaling_factor=routed_scaling_factor,

471 + routed_scaling_factor=self.routed_scaling_factor,

472 e_score_correction_bias=e_score_correction_bias,

473 num_fused_shared_experts=self.num_fused_shared_experts,

474 enable_eplb=enable_eplb,

@@ -578,12 +590,18 @@ class FusedMoE(PluggableLayer):

590 layer_name=self.layer_name,

591 moe_config=self.moe_config,

592 router=self.router,

581 - routed_input_transform=self._routed_input_transform,

593 gate=gate,

594 shared_experts=shared_experts,

595 quant_method=self.quant_method,

585 - reduce_results=self.reduce_results,

596 enable_dbo=self.vllm_config.parallel_config.enable_dbo,

597 + routed_input_transform=routed_input_transform,

598 + routed_output_transform=routed_output_transform,

599 + # When apply_routed_scale_to_output is True, we allow

600 + # the scaling factor to be passed to the runner, otherwise

601 + # we pass 1.0 so it ends up being a nop.

602 + routed_scaling_factor=routed_scaling_factor

603 + if apply_routed_scale_to_output

604 + else 1.0,

605 )

606

607 # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py

@@ -1514,32 +1532,11 @@ class FusedMoE(PluggableLayer):

1532 self.ensure_moe_quant_config_init()

1533 return self.quant_method.moe_quant_config

1534

1517 - def must_reduce_shared_expert_outputs(self) -> bool:

1518 - """

1519 - The shared_experts are typically computed using the RowParallelLinear

1520 - layer. The result of this function is typically used as

1521 - the reduce_results argument to the module.

1522 - When just tensor-parallel is used, it is not required to reduce

1523 - the shared_experts results immediately. Instead we reduce at the

1524 - once at the end of the MoE op. (Refer to DeepSeekV2MoE module)

1525 - With EP and all2all kernels - this is no longer viable as all

1526 - GPU ranks in DP, produce the complete set of hidden_states.

1527 - Therefore it is required that we reduce the shared_experts output

1528 - early.

1529 - """

1530 - return self.runner.must_reduce_shared_expert_outputs()

1531 -

1532 - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):

1533 - """

1534 - Some combine kernels reduce across GPU ranks by default.

1535 - """

1536 - return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states)

1537 -

1535 def forward(

1536 self,

1537 hidden_states: torch.Tensor,

1538 router_logits: torch.Tensor,

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

1539 + ) -> torch.Tensor:

1540 return self.runner.forward(

1541 hidden_states,

1542 router_logits,

@@ -1613,7 +1610,6 @@ class FusedMoE(PluggableLayer):

1610 f"intermediate_size_per_partition={self.intermediate_size_per_partition}, " # noqa: E501

1611 f"tp_size={self.tp_size},\n"

1612 f"ep_size={self.ep_size}, "

1616 - f"reduce_results={self.reduce_results}, "

1613 )

1614

1615 return s

vllm/model_executor/layers/fused_moe/modular_kernel.py

+1 -1

@@ -1261,7 +1261,7 @@ class FusedMoEKernelModularImpl:

1261 topk_ids: torch.Tensor,

1262 apply_router_weight_on_input: bool,

1263 shared_experts_input: torch.Tensor | None,

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

1264 + ) -> torch.Tensor:

1265 """

1266 The _finalize method is a wrapper around self.prepare_finalize.finalize

1267 that handles DBO, async and shared expert overlap.

vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py

-4

@@ -37,10 +37,6 @@ class DefaultMoERunner(MoERunnerBase):

37 for different configurations (e.g., with/without shared experts, gates, etc.).

38 """

40 - @property

41 - def reduce_results(self) -> bool:

42 - return self._reduce_results

43 -

40 @property

41 def do_naive_dispatch_combine(self) -> bool:

42 return (

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

+1 -12

@@ -26,18 +26,7 @@ class MoERunner(ABC):

26 self,

27 hidden_states: torch.Tensor,

28 router_logits: torch.Tensor,

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

30 - raise NotImplementedError

31 -

32 - @abstractmethod

33 - def must_reduce_shared_expert_outputs(self) -> bool:

34 - raise NotImplementedError

35 -

36 - @abstractmethod

37 - def maybe_all_reduce_tensor_model_parallel(

38 - self,

39 - final_hidden_states: torch.Tensor,

40 - ):

29 + ) -> torch.Tensor:

30 raise NotImplementedError

32 @abstractmethod

vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py

+175 -86

@@ -81,7 +81,9 @@ def _resolve_layer_name(layer_name: str | LayerName) -> str:

82 # Note: _moe_forward and _moe_forward_shared should not contain any

83 # implementation details, They should merely pass along control to

84 -# the runner's 'forward_dispatch' method.

84 +# the runner's '_forward_dispatch' method.

85 +# These functions should never be called directly since they do not

86 +# include all the functionality of the MoE layer.

87 def _moe_forward(

88 hidden_states: torch.Tensor,

89 router_logits: torch.Tensor,

@@ -89,7 +91,7 @@ def _moe_forward(

91 layer_name: _layer_name_type,

92 ) -> torch.Tensor:

93 layer = get_layer_from_name(_resolve_layer_name(layer_name))

92 - return layer.runner.forward_dispatch(

94 + return layer.runner._forward_dispatch(

95 layer,

96 hidden_states,

97 router_logits,

@@ -113,7 +115,7 @@ def _moe_forward_shared(

115 layer_name: _layer_name_type,

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

117 layer = get_layer_from_name(_resolve_layer_name(layer_name))

116 - return layer.runner.forward_dispatch(

118 + return layer.runner._forward_dispatch(

119 layer,

120 hidden_states,

121 router_logits,

@@ -143,7 +145,7 @@ def _moe_forward_shared_fake(

145 direct_register_custom_op(

146 op_name="moe_forward",

147 op_func=_moe_forward,

146 - mutates_args=["hidden_states"], # is this still true?

148 + mutates_args=["hidden_states"],

149 fake_impl=_moe_forward_fake,

150 tags=(torch.Tag.needs_fixed_stride_order,),

151 )

@@ -157,6 +159,15 @@ direct_register_custom_op(

159 )

160

161

162 +def _unpack(

163 + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

164 +) -> tuple[torch.Tensor | None, torch.Tensor]:

165 + if isinstance(result, tuple):

166 + return result

167 + else:

168 + return (None, result)

169 +

170 +

171 class MoERunnerBase(MoERunner):

172 """

173 Abstract base class providing common functionality for MoE runner implementations.

@@ -174,7 +185,6 @@ class MoERunnerBase(MoERunner):

185 allowing flexibility in the actual MoE computation implementation.

186

187 Key abstract methods that subclasses must implement:

177 - - reduce_results: Determines whether results should be reduced across ranks

188 - _forward_impl: The core MoE computation logic specific to each runner type

189 """

190

@@ -187,17 +197,23 @@ class MoERunnerBase(MoERunner):

197 gate: torch.nn.Module | None,

198 shared_experts: torch.nn.Module | None,

199 quant_method: FusedMoEMethodBase,

190 - reduce_results: bool,

200 enable_dbo: bool,

201 + routed_output_transform: torch.nn.Module | None = None,

202 + routed_scaling_factor: float = 1.0,

203 ):

204 super().__init__()

205 self.moe_config = moe_config

206 self.router = router

207 self.routed_input_transform = routed_input_transform

208 + self.routed_output_transform = routed_output_transform

209 + self.routed_scaling_factor = routed_scaling_factor

210 self.gate = gate

211 self.quant_method = quant_method

199 - self._reduce_results = reduce_results

212 self.enable_dbo = enable_dbo

213 + self._fused_output_is_reduced = (

214 + self.quant_method.moe_kernel is not None

215 + and self.quant_method.moe_kernel.output_is_reduced()

216 + )

217

218 self._shared_experts: SharedExperts | None = None

219 if shared_experts is not None:

@@ -209,7 +225,6 @@ class MoERunnerBase(MoERunner):

225 # called, i.e. by a MK or by the MoERunner.

226 # Once the MK can be created upfront, we can just pass in the proper

227 # flags derived from the quant_method's MK.

212 - reduce_results=reduce_results,

228 quant_method=quant_method,

229 enable_dbo=enable_dbo,

230 )

@@ -217,7 +232,7 @@ class MoERunnerBase(MoERunner):

232 # Needed for string -> FusedMoE layer lookup in custom ops.

233 self.layer_name = layer_name

234

220 - self.forward_entry = self._select_forward()

235 + self._forward_entry = self._select_forward()

236

237 def _select_forward(self) -> Callable:

238 if current_platform.is_tpu() or current_platform.is_cpu():

@@ -245,38 +260,6 @@ class MoERunnerBase(MoERunner):

260 def is_internal_router(self) -> bool:

261 return self.gate is not None

262

248 - @property

249 - @abstractmethod

250 - def reduce_results(self) -> bool:

251 - raise NotImplementedError

252 -

253 - def must_reduce_shared_expert_outputs(self) -> bool:

254 - """

255 - The shared_experts are typically computed using the RowParallelLinear

256 - layer. The result of this function is typically used as

257 - the reduce_results argument to the module.

258 - When just tensor-parallel is used, it is not required to reduce

259 - the shared_experts results immediately. Instead we reduce at the

260 - once at the end of the MoE op. (Refer to DeepSeekV2MoE module)

261 - With EP and all2all kernels - this is no longer viable as all

262 - GPU ranks in DP, produce the complete set of hidden_states.

263 - Therefore it is required that we reduce the shared_experts output

264 - early.

265 - """

266 - return (

267 - self.quant_method.moe_kernel is not None

268 - and self.quant_method.moe_kernel.output_is_reduced()

269 - )

270 -

271 - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):

272 - """

273 - Some combine kernels reduce across GPU ranks by default.

274 - """

275 - if self.must_reduce_shared_expert_outputs():

276 - return final_hidden_states

277 - else:

278 - return tensor_model_parallel_all_reduce(final_hidden_states)

279 -

263 def apply_routed_input_transform(

264 self, hidden_states: torch.Tensor

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

@@ -286,10 +269,6 @@ class MoERunnerBase(MoERunner):

269 is saved separately so shared experts get [S, hidden_size] while

270 routed experts get the transformed [S, moe_latent_size].

271

289 - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be

290 - moved inside SharedFusedMoE to all-reduce on the smaller latent

291 - dimension.

292 -

272 Returns (possibly transformed) hidden states and the input for shared

273 experts (or None if there are no shared experts).

274 """

@@ -306,33 +285,79 @@ class MoERunnerBase(MoERunner):

285 hidden_states if self._shared_experts is not None else None,

286 )

287

309 - def _maybe_reduce_output(

288 + def apply_routed_output_transform(

289 self,

311 - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

312 - trunc_sizes: list[int],

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

314 - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:

315 - return x[..., :trunc_size]

290 + fused_output: torch.Tensor,

291 + ) -> torch.Tensor:

292 + """Apply transform to routed expert output (e.g., latent to full dim).

293 +

294 + Used by latent MoE models (e.g., NemotronH) where routed experts

295 + operate in a compressed latent space and need projection back to

296 + the full hidden dimension before combining with shared expert output.

297 + """

298 + if self.routed_output_transform is not None:

299 + r = self.routed_output_transform(fused_output)

300 + fused_output = r[0] if isinstance(r, tuple) else r

301 + return fused_output

302

317 - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:

318 - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)

303 + def _maybe_apply_routed_scale_to_output(

304 + self,

305 + shared_output: torch.Tensor | None,

306 + fused_output: torch.Tensor,

307 + ) -> tuple[torch.Tensor | None, torch.Tensor]:

308 + """Apply routed_scaling_factor to the output with FP16 overflow

309 + protection.

310 +

311 + Scale the fused expert output by routed_scaling_factor. For FP16,

312 + avoid overflow by dividing shared_output by the scale instead

313 + (the decoder layer compensates with matching divisions).

314 + """

315 + if self.routed_scaling_factor != 1.0:

316 + if fused_output.dtype != torch.float16:

317 + fused_output *= self.routed_scaling_factor

318 + elif shared_output is not None:

319 + shared_output *= 1.0 / self.routed_scaling_factor

320 + return shared_output, fused_output

321 +

322 + def _maybe_reduce_shared_expert_output(

323 + self,

324 + shared_output: torch.Tensor | None,

325 + ) -> torch.Tensor | None:

326 + """All-reduce shared expert output when the combine kernel already

327 + reduced fused output.

328 +

329 + This is the "early" all-reduce path. When the combine kernel produces

330 + already-reduced fused output, shared output must be reduced separately

331 + to match.

332 + """

333 + if self._fused_output_is_reduced:

334 + assert shared_output is not None

335 + shared_output = tensor_model_parallel_all_reduce(shared_output)

336 + return shared_output

337

338 + def _maybe_reduce_final_output(

339 + self,

340 + states: torch.Tensor,

341 + trunc_size: int,

342 + ) -> torch.Tensor:

343 + """Truncate padded dimensions and all-reduce the combined output.

344 +

345 + This is the "late" all-reduce path. When neither fused nor shared

346 + output was individually reduced, the combined sum is all-reduced

347 + here. Skipped when sequence-parallel is active (SP handles its

348 + own reduction) or when the early path already reduced both outputs.

349 + """

350 + # We don't need to reduce the final output if:

351 + # - We are not running with TP or DP

352 + # - The MK already reduced the fused output itself.

353 if (

354 not self.moe_config.is_sequence_parallel

322 - and self.reduce_results

355 and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)

356 + and not self._fused_output_is_reduced

357 ):

325 - func = reduce_and_trunc

326 - else:

327 - func = trunc

358 + states = tensor_model_parallel_all_reduce(states)

359

329 - if isinstance(states, tuple):

330 - return tuple(

331 - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]

332 - )

333 - else:

334 - assert len(trunc_sizes) == 1

335 - return func(states, trunc_sizes[0])

360 + return states[..., :trunc_size]

361

362 def _encode_layer_name(self) -> str | LayerName:

363 if _USE_LAYERNAME:

@@ -349,7 +374,15 @@ class MoERunnerBase(MoERunner):

374 self,

375 shared_experts_input: torch.Tensor | None,

376 hidden_states: torch.Tensor,

352 - ) -> tuple[torch.Tensor, list[int]]:

377 + ) -> tuple[torch.Tensor, int]:

378 + """Pad hidden_states to moe_config.hidden_dim and compute the

379 + original dimension for later truncation.

380 +

381 + For latent MoE, the routed hidden_states may be smaller than

382 + hidden_dim. Padding ensures uniform tensor sizes through the

383 + fused MoE kernel. The returned trunc_size is used by

384 + _maybe_reduce_final_output to strip the padding from the result.

385 + """

386 shared_experts_hidden_dim = (

387 shared_experts_input.shape[-1] if shared_experts_input is not None else 0

388 )

@@ -365,10 +398,10 @@ class MoERunnerBase(MoERunner):

398 value=0.0,

399 )

400

368 - if self._shared_experts is not None:

369 - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]

401 + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0:

402 + orig_hidden_dims = shared_experts_hidden_dim

403 else:

371 - orig_hidden_dims = [transformed_hidden_dim]

404 + orig_hidden_dims = transformed_hidden_dim

405

406 return hidden_states, orig_hidden_dims

407

@@ -388,6 +421,12 @@ class MoERunnerBase(MoERunner):

421 router_logits: torch.Tensor,

422 shared_experts_input: torch.Tensor | None,

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

424 + """Run expert routing and the fused MoE kernel via the quant method.

425 +

426 + Orchestrates shared expert execution (before/after), expert selection

427 + via the router, and the actual fused MoE computation. Returns

428 + (shared_expert_output, fused_expert_output).

429 + """

430 # Run this before quant_method to avoid inplace issues.

431 # TODO(bnell): probably not needed anymore since inplace is

432 # disabled when shared experts are present.

@@ -428,6 +467,13 @@ class MoERunnerBase(MoERunner):

467 )

468

469 def _sequence_parallel_context(self):

470 + """Return a context manager for sequence-parallel token

471 + redistribution.

472 +

473 + When sequence parallelism is active, returns a context that handles

474 + local size tracking for proper token scatter/gather. Otherwise

475 + returns a no-op context.

476 + """

477 ctx = get_forward_context()

478 return (

479 ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)

@@ -448,22 +494,25 @@ class MoERunnerBase(MoERunner):

494

495 def _maybe_add_zero_expert_output(

496 self,

451 - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

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

497 + result: torch.Tensor,

498 + ) -> torch.Tensor:

499 + """Add the zero expert's contribution to the final result.

500 +

501 + When a ZeroExpertRouter is used, it computes a bias-like output

502 + from the "zero expert" that is added to the combined routed+shared

503 + expert output.

504 + """

505 if isinstance(self.router, ZeroExpertRouter):

506 zero_expert_output = self.router.zero_expert_output

507 assert zero_expert_output is not None

456 - if isinstance(result, tuple):

457 - result = (result[0], result[1] + zero_expert_output)

458 - else:

459 - result = result + zero_expert_output

508 + result = result + zero_expert_output

509 return result

510

511 def forward(

512 self,

513 hidden_states: torch.Tensor,

514 router_logits: torch.Tensor,

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

515 + ) -> torch.Tensor:

516 """Invoke the fused moe layer.

517

518 Input:

@@ -472,55 +521,88 @@ class MoERunnerBase(MoERunner):

521

522 Output:

523 - The new hidden_states.

475 - or

476 - - A tuple of (shared experts output, new hidden_states).

524

525 Calling sequence

526 - forward

480 - - self.forward_entry (_moe_forward or _moe_forward_shared custom op)

481 - - forward_dispatch

527 + - self._forward_entry (_moe_forward or _moe_forward_shared custom op)

528 + - _forward_dispatch

529 - _forward_impl

530

531 Note: The existence of _moe_forward and _moe_forward_shared custom ops are due

532 to the following reasons:

533 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by

534 torch.compile

488 - 2. pytorch cannot handle union types in custom op signatures so _moe_forward

489 - and _moe_forward_shared must be split.

535 + 2. pytorch cannot handle union types in custom op signatures so

536 + _moe_forward and _moe_forward_shared must be split.

537

538 If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can

539 potentially get rid of _moe_forward and _moe_forward_shared and collapse the

540 whole sequence into the 'forward' method.

541 """

542

496 - # Apply transform for routed experts (e.g., latent projection for latent MoE)

543 + # Apply transform for routed experts (e.g., latent projection

544 + # for latent MoE)

545 hidden_states, shared_experts_input = self.apply_routed_input_transform(

546 hidden_states

547 )

548

501 - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(

549 + hidden_states, og_hidden_dim = self._maybe_pad_hidden_states(

550 shared_experts_input,

551 hidden_states,

552 )

553

506 - fused_output = self.forward_entry(

554 + result = self._forward_entry(

555 hidden_states,

556 router_logits,

557 shared_experts_input,

558 self._encode_layer_name(),

559 )

560

513 - result = self._maybe_reduce_output(fused_output, og_hidden_dims)

561 + #

562 + # Note: there are two all-reduce points below. They are mutually

563 + # exclusive, controlled by _fused_output_is_reduced

564 + # - When True: the combine kernel already reduced fused_output,

565 + # so we reduce shared_output here to match, then skip the

566 + # all-reduce in _maybe_reduce_final_output.

567 + # - When False: neither output is reduced yet, so we combine

568 + # them first and all-reduce the sum in _maybe_reduce_final_output.

569 +

570 + # Extract outputs from result

571 + shared_output, fused_output = _unpack(result)

572 +

573 + # If combine kernel already reduced fused, reduce shared to match.

574 + # See note above re: the two all-reduce points.

575 + shared_output = self._maybe_reduce_shared_expert_output(shared_output)

576 +

577 + shared_output, fused_output = self._maybe_apply_routed_scale_to_output(

578 + shared_output, fused_output

579 + )

580 +

581 + # Apply output transform (e.g. latent -> full dim)

582 + fused_output = self.apply_routed_output_transform(fused_output)

583 +

584 + if shared_output is not None:

585 + result = shared_output + fused_output

586 + else:

587 + result = fused_output

588 +

589 + result = self._maybe_reduce_final_output(result, og_hidden_dim)

590

591 return self._maybe_add_zero_expert_output(result)

592

517 - def forward_dispatch(

593 + def _forward_dispatch(

594 self,

595 layer: torch.nn.Module,

596 hidden_states: torch.Tensor,

597 router_logits: torch.Tensor,

598 shared_experts_input: torch.Tensor | None,

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

600 + """Entry point called by the custom op to run the MoE computation.

601 +

602 + Handles pre-dispatch setup (gate application, external shared expert

603 + triggering, quant config init) then delegates to _forward_impl within

604 + the sequence-parallel context.

605 + """

606 # TODO(bnell): this can be removed after MK migration is complete.

607 layer.ensure_moe_quant_config_init()

608

@@ -549,4 +631,11 @@ class MoERunnerBase(MoERunner):

631 router_logits: torch.Tensor,

632 shared_experts_input: torch.Tensor | None,

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

634 + """Core MoE computation to be implemented by subclasses.

635 +

636 + Performs expert routing, fused MoE kernel execution, and shared

637 + expert computation. Returns a single tensor (fused output only)

638 + or a tuple of (shared_output, fused_output) when shared experts

639 + are present.

640 + """

641 raise NotImplementedError

vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py

+4 -2

@@ -29,8 +29,9 @@ def create_moe_runner(

29 gate: torch.nn.Module | None,

30 shared_experts: SharedExperts | None,

31 quant_method: FusedMoEMethodBase,

32 - reduce_results: bool,

32 enable_dbo: bool,

33 + routed_output_transform: torch.nn.Module | None = None,

34 + routed_scaling_factor: float = 1.0,

35 ) -> MoERunner:

36 return DefaultMoERunner(

37 layer_name,

@@ -40,6 +41,7 @@ def create_moe_runner(

41 gate,

42 shared_experts,

43 quant_method,

43 - reduce_results,

44 enable_dbo,

45 + routed_output_transform=routed_output_transform,

46 + routed_scaling_factor=routed_scaling_factor,

47 )

vllm/model_executor/layers/fused_moe/runner/shared_experts.py

-18

@@ -5,10 +5,6 @@ from enum import IntEnum

5 import torch

7 import vllm.envs as envs

8 -from vllm.distributed import (

9 - get_tensor_model_parallel_world_size,

10 - tensor_model_parallel_all_reduce,

11 -)

8 from vllm.logger import init_logger

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

10 FusedMoEConfig,

@@ -48,7 +44,6 @@ class SharedExperts:

44 layer: torch.nn.Module,

45 moe_config: FusedMoEConfig,

46 quant_method: QuantizeMethodBase,

51 - reduce_results: bool,

47 enable_dbo: bool,

48 ):

49 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

@@ -68,7 +63,6 @@ class SharedExperts:

63 self._layer = layer

64 self._moe_config = moe_config

65 self._quant_method = quant_method

71 - self._reduce_results = reduce_results

67 # Allow disabling of the separate shared experts stream for

68 # debug purposes.

@@ -139,18 +133,6 @@ class SharedExperts:

133

134 return output

135

142 - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor:

143 - # Reduce shared expert outputs if necessary, since the MLP

144 - # should have been created with reduce_results=False.

145 - if (

146 - self._reduce_results

147 - and self._quant_method.moe_kernel is not None

148 - and self._quant_method.moe_kernel.output_is_reduced()

149 - and get_tensor_model_parallel_world_size() > 1

150 - ):

151 - shared_out = tensor_model_parallel_all_reduce(shared_out)

152 - return shared_out

153 -

136 @property

137 def _output_idx(self) -> int:

138 return dbo_current_ubatch_id() if self.enable_dbo else 0

vllm/model_executor/layers/fused_moe/shared_fused_moe.py

+2 -6

@@ -18,12 +18,8 @@ class SharedFusedMoE(FusedMoE):

18 self,

19 hidden_states: torch.Tensor,

20 router_logits: torch.Tensor,

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

22 - result = super().forward(

21 + ) -> torch.Tensor:

22 + return super().forward(

23 hidden_states=hidden_states,

24 router_logits=router_logits,

25 )

26 - if self.shared_experts is None:

27 - return None, result

28 - else:

29 - return result

vllm/model_executor/models/AXK1.py

+5 -30

@@ -100,7 +100,7 @@ class AXK1MoE(nn.Module):

100 self.tp_size = get_tensor_model_parallel_world_size()

101 self.tp_rank = get_tensor_model_parallel_rank()

102

103 - self.routed_scaling_factor = config.routed_scaling_factor

103 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)

104

105 self.ep_group = get_ep_group().device_group

106 self.ep_rank = get_ep_group().rank_in_group

@@ -170,7 +170,6 @@ class AXK1MoE(nn.Module):

170 top_k=config.num_experts_per_tok,

171 hidden_size=config.hidden_size,

172 intermediate_size=config.moe_intermediate_size,

173 - reduce_results=False,

173 renormalize=config.norm_topk_prob,

174 quant_config=quant_config,

175 use_grouped_topk=True,

@@ -180,9 +179,8 @@ class AXK1MoE(nn.Module):

179 scoring_func=config.scoring_func,

180 # we do scaling outside, set factor to 1.0 to avoid double mul

181 # aiter applies routed_scaling_factor internally

183 - routed_scaling_factor=1.0

184 - if not self.is_rocm_aiter_moe_enabled

185 - else self.routed_scaling_factor,

182 + routed_scaling_factor=self.routed_scaling_factor,

183 + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled,

184 e_score_correction_bias=self.gate.e_score_correction_bias,

185 enable_eplb=self.enable_eplb,

186 num_redundant_experts=self.n_redundant_experts,

@@ -204,43 +202,20 @@ class AXK1MoE(nn.Module):

202 hidden_states = sequence_parallel_chunk(hidden_states)

203

204 if self.experts.is_internal_router:

207 - # In this case, the gate/router runs inside the FusedMoE class

208 - fused_moe_out = self.experts(

205 + final_hidden_states = self.experts(

206 hidden_states=hidden_states, router_logits=hidden_states

207 )

208 else:

212 - # router_logits: (num_tokens, n_experts)

209 router_logits, _ = self.gate(hidden_states)

214 - fused_moe_out = self.experts(

210 + final_hidden_states = self.experts(

211 hidden_states=hidden_states, router_logits=router_logits

212 )

213

218 - shared_output, final_hidden_states = fused_moe_out

219 - if self.shared_experts is None:

220 - assert shared_output is None

221 -

222 - # Fix FP16 overflow

223 - # See AXK1DecoderLayer for more details.

224 - if hidden_states.dtype != torch.float16:

225 - if not self.is_rocm_aiter_moe_enabled:

226 - final_hidden_states *= self.routed_scaling_factor

227 - elif self.shared_experts is not None:

228 - assert shared_output is not None

229 - shared_output *= 1.0 / self.routed_scaling_factor

230 -

231 - if self.shared_experts is not None:

232 - assert shared_output is not None

233 - final_hidden_states += shared_output

234 -

214 if self.is_sequence_parallel:

215 final_hidden_states = tensor_model_parallel_all_gather(

216 final_hidden_states, 0

217 )

218 final_hidden_states = final_hidden_states[:num_tokens]

240 - elif self.tp_size > 1:

241 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

242 - final_hidden_states

243 - )

219

220 return final_hidden_states.view(num_tokens, hidden_dim)

221

vllm/model_executor/models/afmoe.py

+1 -12

@@ -131,7 +131,6 @@ class AfmoeMoE(nn.Module):

131 top_k=config.num_experts_per_tok,

132 hidden_size=config.hidden_size,

133 intermediate_size=config.moe_intermediate_size,

134 - reduce_results=False,

134 renormalize=self.route_norm if self.score_func == "sigmoid" else False,

135 quant_config=quant_config,

136 use_grouped_topk=True,

@@ -152,20 +151,10 @@ class AfmoeMoE(nn.Module):

151

152 router_logits = self.gate(hidden_states.to(dtype=torch.float32))

153

155 - fused_moe_out = self.experts(

154 + final_hidden_states = self.experts(

155 hidden_states=hidden_states, router_logits=router_logits

156 )

157

159 - if self.shared_experts is not None:

160 - shared_output, final_hidden_states = fused_moe_out

161 - final_hidden_states = final_hidden_states + shared_output

162 - else:

163 - final_hidden_states = fused_moe_out

164 - if self.tp_size > 1:

165 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

166 - final_hidden_states

167 - )

168 -

158 return final_hidden_states.view(num_tokens, hidden_dim)

159

160

vllm/model_executor/models/aria.py

+1 -7

@@ -283,7 +283,6 @@ class AriaTextMoELayer(nn.Module):

283 hidden_size=config.hidden_size,

284 intermediate_size=config.intermediate_size,

285 quant_config=quant_config,

286 - reduce_results=True,

286 prefix=f"{prefix}.experts",

287 )

288

@@ -301,12 +300,7 @@ class AriaTextMoELayer(nn.Module):

300

301 router_output = torch.nn.functional.linear(hidden_states, self.router_weight)

302

304 - sparse_expert_output = self.experts(hidden_states, router_output)

305 -

306 - if self.shared_experts is not None:

307 - return sparse_expert_output[0] + sparse_expert_output[1]

308 - else:

309 - return sparse_expert_output

303 + return self.experts(hidden_states, router_output)

304

305

306 class AriaTextDecoderLayer(LlamaDecoderLayer):

vllm/model_executor/models/bailing_moe.py

+1 -16

@@ -291,7 +291,6 @@ class BailingMoE(nn.Module):

291 top_k=self.top_k,

292 hidden_size=self.hidden_size,

293 intermediate_size=config.moe_intermediate_size,

294 - reduce_results=False,

294 renormalize=self.norm_expert_prob,

295 quant_config=quant_config,

296 prefix=f"{prefix}.experts",

@@ -301,6 +300,7 @@ class BailingMoE(nn.Module):

300 topk_group=self.topk_group,

301 use_grouped_topk=self.use_grouped_topk,

302 router_logits_dtype=self.router_dtype,

303 + routed_scaling_factor=self.routed_scaling_factor,

304 )

305

306 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -314,21 +314,6 @@ class BailingMoE(nn.Module):

314 final_hidden_states = self.experts(

315 hidden_states=hidden_states, router_logits=router_logits

316 )

317 -

318 - if self.shared_experts is not None:

319 - shared_output, final_hidden_states = final_hidden_states

320 - else:

321 - shared_output = None

322 -

323 - final_hidden_states *= self.routed_scaling_factor

324 -

325 - if shared_output is not None:

326 - final_hidden_states = final_hidden_states + shared_output

327 -

328 - if self.tp_size > 1:

329 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

330 - final_hidden_states

331 - )

317 return final_hidden_states.view(num_tokens, hidden_size)

318

319

vllm/model_executor/models/bailing_moe_linear.py

+2 -17

@@ -358,7 +358,6 @@ class BailingMoeV25(nn.Module):

358 top_k=self.top_k,

359 hidden_size=self.hidden_size,

360 intermediate_size=config.moe_intermediate_size,

361 - reduce_results=False,

361 renormalize=self.norm_expert_prob,

362 quant_config=quant_config,

363 prefix=f"{prefix}.experts",

@@ -368,6 +367,8 @@ class BailingMoeV25(nn.Module):

367 topk_group=self.topk_group,

368 use_grouped_topk=self.use_grouped_topk,

369 router_logits_dtype=self.router_dtype,

370 + routed_scaling_factor=self.routed_scaling_factor,

371 + apply_routed_scale_to_output=True,

372 )

373

374 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -383,22 +384,6 @@ class BailingMoeV25(nn.Module):

384 hidden_states=hidden_states, router_logits=router_logits

385 )

386

386 - # Handle tuple return from SharedFusedMoE

387 - if self.shared_experts is not None:

388 - shared_output, final_hidden_states = final_hidden_states

389 - else:

390 - shared_output = None

391 -

392 - final_hidden_states *= self.routed_scaling_factor

393 -

394 - if shared_output is not None:

395 - final_hidden_states = final_hidden_states + shared_output

396 -

397 - if self.tp_size > 1:

398 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

399 - final_hidden_states

400 - )

401 -

387 return final_hidden_states.view(num_tokens, hidden_size)

388

389

vllm/model_executor/models/dbrx.py

-1

@@ -85,7 +85,6 @@ class DbrxExperts(FusedMoE):

85 hidden_size=config.d_model,

86 intermediate_size=config.ffn_config.ffn_hidden_size,

87 params_dtype=params_dtype,

88 - reduce_results=True,

88 renormalize=True,

89 quant_config=quant_config,

90 tp_size=get_tensor_model_parallel_world_size(),

vllm/model_executor/models/deepseek_v2.py

+4 -30

@@ -318,7 +318,6 @@ class DeepseekV2MoE(nn.Module):

318 top_k=config.num_experts_per_tok,

319 hidden_size=config.hidden_size,

320 intermediate_size=config.moe_intermediate_size,

321 - reduce_results=False,

321 renormalize=config.norm_topk_prob,

322 quant_config=quant_config,

323 use_grouped_topk=True,

@@ -326,11 +325,9 @@ class DeepseekV2MoE(nn.Module):

325 topk_group=getattr(config, "topk_group", 1),

326 prefix=f"{prefix}.experts",

327 scoring_func=getattr(config, "scoring_func", "softmax"),

329 - # we do scaling outside, set factor to 1.0 to avoid double mul

328 # aiter applies routed_scaling_factor internally

331 - routed_scaling_factor=1.0

332 - if not self.is_rocm_aiter_moe_enabled

333 - else self.routed_scaling_factor,

329 + routed_scaling_factor=self.routed_scaling_factor,

330 + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled,

331 e_score_correction_bias=self.gate.e_score_correction_bias,

332 enable_eplb=self.enable_eplb,

333 num_redundant_experts=self.n_redundant_experts,

@@ -363,43 +360,20 @@ class DeepseekV2MoE(nn.Module):

360 hidden_states = sequence_parallel_chunk(hidden_states)

361

362 if self.experts.is_internal_router:

366 - # In this case, the gate/router runs inside the FusedMoE class

367 - fused_moe_out = self.experts(

363 + final_hidden_states = self.experts(

364 hidden_states=hidden_states, router_logits=hidden_states

365 )

366 else:

371 - # router_logits: (num_tokens, n_experts)

367 router_logits, _ = self.gate(hidden_states)

373 - fused_moe_out = self.experts(

368 + final_hidden_states = self.experts(

369 hidden_states=hidden_states, router_logits=router_logits

370 )

371

377 - shared_output, final_hidden_states = fused_moe_out

378 - if self.shared_experts is None:

379 - assert shared_output is None

380 -

381 - # Fix FP16 overflow

382 - # See DeepseekV2DecoderLayer for more details.

383 - if hidden_states.dtype != torch.float16:

384 - if not self.is_rocm_aiter_moe_enabled:

385 - final_hidden_states *= self.routed_scaling_factor

386 - elif self.shared_experts is not None:

387 - assert shared_output is not None

388 - shared_output *= 1.0 / self.routed_scaling_factor

389 -

390 - if self.shared_experts is not None:

391 - assert shared_output is not None

392 - final_hidden_states += shared_output

393 -

372 if self.is_sequence_parallel:

373 final_hidden_states = tensor_model_parallel_all_gather(

374 final_hidden_states, 0

375 )

376 final_hidden_states = final_hidden_states[:num_tokens]

399 - elif self.tp_size > 1:

400 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

401 - final_hidden_states

402 - )

377

378 return final_hidden_states.view(num_tokens, hidden_dim)

379

vllm/model_executor/models/dots1.py

+3 -13

@@ -37,7 +37,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig

37 from vllm.distributed import (

38 get_pp_group,

39 get_tensor_model_parallel_world_size,

40 - tensor_model_parallel_all_reduce,

40 )

41 from vllm.model_executor.layers.activation import SiluAndMul

42 from vllm.model_executor.layers.attention import Attention

@@ -120,7 +119,6 @@ class Dots1MoE(nn.Module):

119 prefix: str = "",

120 ):

121 super().__init__()

123 - self.tp_size = get_tensor_model_parallel_world_size()

122 self.routed_scaling_factor = config.routed_scaling_factor

123 self.n_shared_experts = config.n_shared_experts

124

@@ -163,7 +161,6 @@ class Dots1MoE(nn.Module):

161 top_k=config.num_experts_per_tok,

162 hidden_size=config.hidden_size,

163 intermediate_size=config.moe_intermediate_size,

166 - reduce_results=False,

164 renormalize=config.norm_topk_prob,

165 quant_config=quant_config,

166 use_grouped_topk=True,

@@ -171,9 +168,9 @@ class Dots1MoE(nn.Module):

168 topk_group=config.topk_group,

169 prefix=f"{prefix}.experts",

170 scoring_func=config.scoring_func,

174 - # we do scaling outside, set factor to 1.0 to avoid double mul

175 - routed_scaling_factor=1.0,

171 e_score_correction_bias=self.gate.e_score_correction_bias,

172 + routed_scaling_factor=self.routed_scaling_factor,

173 + apply_routed_scale_to_output=True,

174 )

175

176 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -182,16 +179,9 @@ class Dots1MoE(nn.Module):

179

180 router_logits, _ = self.gate(hidden_states)

181

185 - shared_out, routed_out = self.experts(

182 + final_hidden_states = self.experts(

183 hidden_states=hidden_states, router_logits=router_logits

184 )

188 - if self.shared_experts is not None:

189 - final_hidden_states = (routed_out + shared_out) * self.routed_scaling_factor

190 - else:

191 - final_hidden_states = routed_out * self.routed_scaling_factor

192 -

193 - if self.tp_size > 1:

194 - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

185 return final_hidden_states.view(num_tokens, hidden_dim)

186

187

vllm/model_executor/models/ernie45_moe.py

-11

@@ -194,7 +194,6 @@ class Ernie4_5_MoeMoE(nn.Module):

194 top_k=config.moe_k,

195 hidden_size=config.hidden_size,

196 intermediate_size=config.moe_intermediate_size,

197 - reduce_results=False,

197 renormalize=True,

198 quant_config=quant_config,

199 prefix=f"{prefix}.experts",

@@ -215,16 +214,6 @@ class Ernie4_5_MoeMoE(nn.Module):

214 hidden_states=hidden_states, router_logits=router_logits

215 )

216

218 - if self.has_shared_experts:

219 - final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

220 - else:

221 - final_hidden_states = final_hidden_states[1]

222 -

223 - if self.tp_size > 1:

224 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

225 - final_hidden_states

226 - )

227 -

217 return final_hidden_states.view(orig_shape)

218

219

vllm/model_executor/models/ernie45_vl_moe.py

+5 -30

@@ -263,7 +263,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):

263 top_k=config.moe_k,

264 hidden_size=config.hidden_size,

265 intermediate_size=config.moe_intermediate_size[0],

266 - reduce_results=False,

266 renormalize=True,

267 quant_config=quant_config,

268 e_score_correction_bias=self.e_score_correction_bias[0],

@@ -301,7 +300,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):

300 top_k=config.moe_k,

301 hidden_size=config.hidden_size,

302 intermediate_size=config.moe_intermediate_size[1],

304 - reduce_results=False,

303 renormalize=True,

304 quant_config=quant_config,

305 e_score_correction_bias=self.e_score_correction_bias[1],

@@ -342,9 +340,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):

340 visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool()

341 text_token_mask = ~visual_token_mask

342 final_experts_hidden_states = torch.zeros_like(hidden_states)

345 - final_shared_output = (

346 - torch.zeros_like(hidden_states) if self.has_shared_experts else None

347 - )

343

344 text_hidden_states = hidden_states[text_token_mask].reshape(

345 -1, self.hidden_size

@@ -356,26 +351,20 @@ class Ernie4_5_VLMoeMoE(nn.Module):

351 text_router_logits, _ = self.text_experts_gate(

352 text_hidden_states.to(dtype=torch.float32)

353 )

359 - text_shared_output, text_experts_output = self.text_experts(

354 + text_output = self.text_experts(

355 hidden_states=text_hidden_states, router_logits=text_router_logits

356 )

362 - final_experts_hidden_states[text_token_mask] = text_experts_output.flatten()

363 - if self.has_shared_experts:

364 - final_shared_output[text_token_mask] = text_shared_output.flatten()

357 + final_experts_hidden_states[text_token_mask] = text_output.flatten()

358

359 vision_router_logits, _ = self.vision_experts_gate(

360 vision_hidden_states.to(dtype=torch.float32)

361 )

369 - vision_shared_output, vision_experts_output = self.vision_experts(

362 + vision_output = self.vision_experts(

363 hidden_states=vision_hidden_states, router_logits=vision_router_logits

364 )

372 - final_experts_hidden_states[visual_token_mask] = (

373 - vision_experts_output.flatten()

374 - )

375 - if self.has_shared_experts:

376 - final_shared_output[visual_token_mask] = vision_shared_output.flatten()

365 + final_experts_hidden_states[visual_token_mask] = vision_output.flatten()

366

378 - final_hidden_states = (final_shared_output, final_experts_hidden_states)

367 + final_hidden_states = final_experts_hidden_states

368 else:

369 # only text modal input

370 text_router_logits, _ = self.text_experts_gate(

@@ -386,20 +375,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):

375 hidden_states=hidden_states, router_logits=text_router_logits

376 )

377

389 - if self.has_shared_experts:

390 - # for shared_experts model

391 - final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

392 - else:

393 - # for not shared_experts model

394 - final_hidden_states = final_hidden_states[1]

395 -

396 - if self.tp_size > 1:

397 - final_hidden_states = (

398 - self.text_experts.maybe_all_reduce_tensor_model_parallel(

399 - final_hidden_states

400 - )

401 - )

402 -

378 return final_hidden_states.view(orig_shape)

379

380

vllm/model_executor/models/exaone_moe.py

+18 -28

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

31 get_tensor_model_parallel_world_size,

32 )

33 from vllm.model_executor.layers.fused_moe import FusedMoE

34 +from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

35 from vllm.model_executor.layers.layernorm import RMSNorm

36 from vllm.model_executor.layers.linear import ReplicatedLinear

37 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -116,12 +117,26 @@ class ExaoneMoe(nn.Module):

117 self.physical_expert_start + self.n_local_physical_experts

118 )

119

119 - self.experts = FusedMoE(

120 + if getattr(config, "num_shared_experts", 0) > 0:

121 + intermediate_size = config.moe_intermediate_size * config.num_shared_experts

122 + self.shared_experts = ExaoneMoeGatedMLP(

123 + hidden_size=config.hidden_size,

124 + intermediate_size=intermediate_size,

125 + hidden_act=config.hidden_act,

126 + quant_config=quant_config,

127 + reduce_results=False,

128 + prefix=f"{prefix}.shared_experts",

129 + )

130 + else:

131 + self.shared_experts = None

132 +

133 + self.experts = SharedFusedMoE(

134 + shared_experts=self.shared_experts,

135 + gate=self.gate,

136 num_experts=self.n_routed_experts,

137 top_k=config.num_experts_per_tok,

138 hidden_size=config.hidden_size,

139 intermediate_size=config.moe_intermediate_size,

124 - reduce_results=False,

140 renormalize=config.norm_topk_prob,

141 quant_config=quant_config,

142 use_grouped_topk=True,

@@ -135,41 +150,16 @@ class ExaoneMoe(nn.Module):

150 num_redundant_experts=self.n_redundant_experts,

151 )

152

138 - if getattr(config, "num_shared_experts", 0) > 0:

139 - intermediate_size = config.moe_intermediate_size * config.num_shared_experts

140 - self.shared_experts = ExaoneMoeGatedMLP(

141 - hidden_size=config.hidden_size,

142 - intermediate_size=intermediate_size,

143 - hidden_act=config.hidden_act,

144 - quant_config=quant_config,

145 - reduce_results=self.experts.must_reduce_shared_expert_outputs(),

146 - prefix=f"{prefix}.shared_experts",

147 - )

148 - else:

149 - self.shared_experts = None

150 -

153 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

154 # NOTE: hidden_states can have either 1D or 2D shape.

155 orig_shape = hidden_states.shape

156 hidden_dim = hidden_states.shape[-1]

157 hidden_states = hidden_states.view(-1, hidden_dim)

158

157 - # router_logits: (num_tokens, n_experts)

158 - router_logits, _ = self.gate(hidden_states)

159 -

159 final_hidden_states = self.experts(

161 - hidden_states=hidden_states, router_logits=router_logits

160 + hidden_states=hidden_states, router_logits=hidden_states

161 )

162

164 - if self.shared_experts is not None:

165 - shared_output = self.shared_experts(hidden_states)

166 - final_hidden_states = final_hidden_states + shared_output

167 -

168 - if self.tp_size > 1:

169 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

170 - final_hidden_states

171 - )

172 -

163 return final_hidden_states.view(orig_shape)

164

165

vllm/model_executor/models/flex_olmo.py

-1

@@ -76,7 +76,6 @@ class FlexOlmoMoE(nn.Module):

76 top_k=hf_config.num_experts_per_tok,

77 hidden_size=hf_config.hidden_size,

78 intermediate_size=hf_config.intermediate_size,

79 - reduce_results=True,

79 renormalize=False,

80 quant_config=None,

81 tp_size=tp_size,

vllm/model_executor/models/gemma4.py

-1

@@ -349,7 +349,6 @@ class Gemma4MoE(nn.Module):

349 "moe_intermediate_size",

350 getattr(config, "expert_intermediate_size", None),

351 ),

352 - reduce_results=True,

352 renormalize=True,

353 quant_config=quant_config,

354 prefix=f"{prefix}.experts",

vllm/model_executor/models/glm4_moe.py

+3 -18

@@ -184,7 +184,6 @@ class Glm4MoE(nn.Module):

184 top_k=config.num_experts_per_tok,

185 hidden_size=config.hidden_size,

186 intermediate_size=config.moe_intermediate_size,

187 - reduce_results=False,

187 renormalize=config.norm_topk_prob,

188 quant_config=quant_config,

189 use_grouped_topk=True,

@@ -192,8 +191,8 @@ class Glm4MoE(nn.Module):

191 topk_group=config.topk_group,

192 prefix=f"{prefix}.experts",

193 scoring_func="sigmoid",

195 - # we do scaling outside, set factor to 1.0 to avoid double mul

196 - routed_scaling_factor=1.0,

194 + routed_scaling_factor=self.routed_scaling_factor,

195 + apply_routed_scale_to_output=True,

196 e_score_correction_bias=self.gate.e_score_correction_bias,

197 enable_eplb=self.enable_eplb,

198 num_redundant_experts=self.n_redundant_experts,

@@ -207,23 +206,9 @@ class Glm4MoE(nn.Module):

206 # router_logits: (num_tokens, n_experts)

207 router_logits = self.gate(hidden_states.to(dtype=torch.float32))

208

210 - fused_moe_out = self.experts(

209 + final_hidden_states = self.experts(

210 hidden_states=hidden_states, router_logits=router_logits

211 )

213 -

214 - if self.shared_experts is not None:

215 - shared_output, final_hidden_states = fused_moe_out

216 - assert shared_output is not None

217 - final_hidden_states = (

218 - final_hidden_states * self.routed_scaling_factor + shared_output

219 - )

220 - else:

221 - final_hidden_states = fused_moe_out * self.routed_scaling_factor

222 -

223 - if self.tp_size > 1:

224 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

225 - final_hidden_states

226 - )

212 return final_hidden_states.view(num_tokens, hidden_dim)

213

214

vllm/model_executor/models/gpt_oss.py

-1

@@ -189,7 +189,6 @@ class MLPBlock(torch.nn.Module):

189 top_k=config.num_experts_per_tok,

190 hidden_size=config.hidden_size,

191 intermediate_size=config.intermediate_size,

192 - reduce_results=True,

192 renormalize=True,

193 quant_config=quant_config,

194 prefix=f"{prefix}.experts",

vllm/model_executor/models/granitemoe.py

-1

@@ -104,7 +104,6 @@ class GraniteMoeMoE(nn.Module):

104 hidden_size=hidden_size,

105 intermediate_size=intermediate_size,

106 params_dtype=params_dtype,

107 - reduce_results=True,

107 renormalize=True,

108 quant_config=quant_config,

109 tp_size=tp_size,

vllm/model_executor/models/grok1.py

-1

@@ -209,7 +209,6 @@ class Grok1MoE(nn.Module):

209 hidden_size=hidden_size,

210 intermediate_size=intermediate_size,

211 params_dtype=params_dtype,

212 - reduce_results=True,

212 renormalize=renormalize,

213 quant_config=quant_config,

214 tp_size=tp_size,

vllm/model_executor/models/hunyuan_v1.py

-7

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

39 get_ep_group,

40 get_pp_group,

41 get_tensor_model_parallel_world_size,

42 - tensor_model_parallel_all_reduce,

42 )

43 from vllm.model_executor.layers.activation import SiluAndMul

44 from vllm.model_executor.layers.attention import Attention

@@ -445,7 +444,6 @@ class HunYuanSparseMoeBlock(nn.Module):

444 top_k=top_k,

445 hidden_size=config.hidden_size,

446 intermediate_size=intermediate_size,

448 - reduce_results=False,

447 renormalize=top_k > 1,

448 quant_config=quant_config,

449 prefix=f"{prefix}.experts",

@@ -464,11 +462,6 @@ class HunYuanSparseMoeBlock(nn.Module):

462 final_hidden_states = self.experts(

463 hidden_states=hidden_states, router_logits=router_logits

464 )

467 - if self.shared_mlp is not None:

468 - final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

469 -

470 - if self.tp_size > 1:

471 - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

465

466 return final_hidden_states.view(orig_shape)

467

vllm/model_executor/models/interns1_pro.py

-1

@@ -176,7 +176,6 @@ class InternS1ProMoeSparseMoeBlock(nn.Module):

176 top_k=config.num_experts_per_tok,

177 hidden_size=config.hidden_size,

178 intermediate_size=config.moe_intermediate_size,

179 - reduce_results=True,

179 renormalize=config.norm_topk_prob,

180 quant_config=quant_config,

181 prefix=f"{prefix}.experts",

vllm/model_executor/models/jamba.py

-1

@@ -90,7 +90,6 @@ class JambaMoE(nn.Module):

90 self.intermediate_size,

91 tp_size=tp_size,

92 params_dtype=params_dtype,

93 - reduce_results=True,

93 renormalize=False,

94 use_grouped_topk=False,

95 quant_config=quant_config,

vllm/model_executor/models/kimi_linear.py

+20 -26

@@ -11,11 +11,10 @@ from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig

11 from vllm.distributed import (

12 get_pp_group,

13 get_tensor_model_parallel_world_size,

14 - tensor_model_parallel_all_reduce,

14 )

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.activation import SiluAndMul

18 -from vllm.model_executor.layers.fused_moe import FusedMoE

17 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE

18 from vllm.model_executor.layers.kda import KimiDeltaAttention

19 from vllm.model_executor.layers.layernorm import RMSNorm

20 from vllm.model_executor.layers.linear import (

@@ -132,12 +131,25 @@ class KimiMoE(nn.Module):

131

132 self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts))

133

135 - self.experts = FusedMoE(

134 + if self.num_shared_experts is not None:

135 + intermediate_size = moe_intermediate_size * self.num_shared_experts

136 + self.shared_experts = KimiMLP(

137 + hidden_size=config.hidden_size,

138 + intermediate_size=intermediate_size,

139 + hidden_act=config.hidden_act,

140 + quant_config=quant_config,

141 + reduce_results=False,

142 + prefix=f"{prefix}.shared_experts",

143 + )

144 + else:

145 + self.shared_experts = None

146 +

147 + self.experts = SharedFusedMoE(

148 + shared_experts=self.shared_experts,

149 num_experts=num_experts,

150 top_k=config.num_experts_per_token,

151 hidden_size=hidden_size,

152 intermediate_size=moe_intermediate_size,

140 - reduce_results=False,

153 renormalize=moe_renormalize,

154 quant_config=quant_config,

155 use_grouped_topk=config.use_grouped_topk,

@@ -146,34 +158,16 @@ class KimiMoE(nn.Module):

158 prefix=f"{prefix}.experts",

159 scoring_func=config.moe_router_activation_func,

160 e_score_correction_bias=self.gate.e_score_correction_bias,

161 + routed_scaling_factor=self.routed_scaling_factor,

162 )

163

151 - if self.num_shared_experts is not None:

152 - intermediate_size = moe_intermediate_size * self.num_shared_experts

153 - self.shared_experts = KimiMLP(

154 - hidden_size=config.hidden_size,

155 - intermediate_size=intermediate_size,

156 - hidden_act=config.hidden_act,

157 - quant_config=quant_config,

158 - reduce_results=False,

159 - prefix=f"{prefix}.shared_experts",

160 - )

161 -

164 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

165 num_tokens, hidden_size = hidden_states.shape

166 hidden_states = hidden_states.view(-1, hidden_size)

165 - if self.num_shared_experts is not None:

166 - shared_output = self.shared_experts(hidden_states)

167 router_logits, _ = self.gate(hidden_states)

168 - final_hidden_states = (

169 - self.experts(hidden_states=hidden_states, router_logits=router_logits)

170 - * self.routed_scaling_factor

168 + final_hidden_states = self.experts(

169 + hidden_states=hidden_states, router_logits=router_logits

170 )

172 - if shared_output is not None:

173 - final_hidden_states = final_hidden_states + shared_output

174 -

175 - if self.tp_size > 1:

176 - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

171 return final_hidden_states.view(num_tokens, hidden_size)

172

173

@@ -482,7 +476,7 @@ class KimiLinearModel(nn.Module):

476 if self.config.is_moe:

477 # Params for weights, fp8 weight scales, fp8 activation scales

478 # (param_name, weight_name, expert_id, shard_id)

485 - expert_params_mapping = FusedMoE.make_expert_params_mapping(

479 + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

480 self,

481 ckpt_gate_proj_name="w1",

482 ckpt_down_proj_name="w2",

vllm/model_executor/models/lfm2_moe.py

+3 -9

@@ -150,7 +150,6 @@ class Lfm2MoeSparseMoeBlock(nn.Module):

150 top_k=config.num_experts_per_tok,

151 hidden_size=config.hidden_size,

152 intermediate_size=config.moe_intermediate_size,

153 - reduce_results=False,

153 renormalize=config.norm_topk_prob,

154 quant_config=quant_config,

155 use_grouped_topk=True, # needed for softmax score func

@@ -161,6 +160,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module):

160 num_redundant_experts=self.n_redundant_experts,

161 scoring_func="sigmoid",

162 e_score_correction_bias=self.gate.e_score_correction_bias,

163 + routed_scaling_factor=self.routed_scaling_factor,

164 )

165

166 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -170,16 +170,10 @@ class Lfm2MoeSparseMoeBlock(nn.Module):

170

171 # router_logits: (num_tokens, n_experts)

172 router_logits, _ = self.gate(hidden_states)

173 - final_hidden_states = (

174 - self.experts(hidden_states=hidden_states, router_logits=router_logits)

175 - * self.routed_scaling_factor

173 + final_hidden_states = self.experts(

174 + hidden_states=hidden_states, router_logits=router_logits

175 )

176

178 - if self.tp_size > 1:

179 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

180 - final_hidden_states

181 - )

182 -

177 return final_hidden_states.view(orig_shape)

178

179

vllm/model_executor/models/llama4.py

+1 -7

@@ -135,7 +135,6 @@ class Llama4MoE(nn.Module):

135 custom_routing_function=Llama4MoE.custom_routing_function,

136 intermediate_size=intermediate_size_moe,

137 apply_router_weight_on_input=True,

138 - reduce_results=False,

138 renormalize=False,

139 quant_config=quant_config,

140 prefix=f"{prefix}.experts",

@@ -151,19 +150,14 @@ class Llama4MoE(nn.Module):

150

151 router_logits, _ = self.router(hidden_states)

152

154 - shared_out, routed_out = self.experts(

153 + experts_out = self.experts(

154 hidden_states=hidden_states,

155 router_logits=router_logits,

156 )

158 - experts_out = routed_out + shared_out

157

158 if self.is_sequence_parallel:

159 experts_out = tensor_model_parallel_all_gather(experts_out, 0)

160 experts_out = experts_out[:num_tokens]

163 - elif self.tp_size > 1:

164 - experts_out = self.experts.maybe_all_reduce_tensor_model_parallel(

165 - experts_out

166 - )

161

162 return experts_out

163

vllm/model_executor/models/longcat_flash.py

-1

@@ -300,7 +300,6 @@ class LongcatMoe(nn.Module):

300 top_k=top_k,

301 hidden_size=hidden_size,

302 intermediate_size=intermediate_size,

303 - reduce_results=True,

303 params_dtype=params_dtype,

304 renormalize=False,

305 quant_config=quant_config,

vllm/model_executor/models/mimo_v2_flash.py

-1

@@ -162,7 +162,6 @@ class MiMoV2MoE(nn.Module):

162 top_k=config.num_experts_per_tok,

163 hidden_size=config.hidden_size,

164 intermediate_size=config.moe_intermediate_size,

165 - reduce_results=True,

165 renormalize=config.norm_topk_prob,

166 quant_config=quant_config,

167 prefix=f"{prefix}.experts",

vllm/model_executor/models/minimax_m2.py

-5

@@ -36,7 +36,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig

36 from vllm.distributed import (

37 get_pp_group,

38 get_tensor_model_parallel_world_size,

39 - tensor_model_parallel_all_reduce,

39 )

40 from vllm.model_executor.layers.attention import Attention

41 from vllm.model_executor.layers.fused_moe import FusedMoE

@@ -104,7 +103,6 @@ class MiniMaxM2MoE(nn.Module):

103 e_score_correction_bias=self.e_score_correction_bias,

104 hidden_size=config.hidden_size,

105 intermediate_size=config.intermediate_size,

107 - reduce_results=False,

106 renormalize=True,

107 quant_config=quant_config,

108 prefix=f"{prefix}.experts",

@@ -134,9 +132,6 @@ class MiniMaxM2MoE(nn.Module):

132 final_hidden_states = self.experts(

133 hidden_states=hidden_states, router_logits=router_logits

134 )

137 - final_hidden_states = final_hidden_states

138 - if self.tp_size > 1:

139 - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

135

136 return final_hidden_states.view(num_tokens, hidden_dim)

137

vllm/model_executor/models/minimax_text_01.py

-1

@@ -162,7 +162,6 @@ class MiniMaxText01MoE(nn.Module):

162 hidden_size=self.hidden_size,

163 intermediate_size=self.intermediate_size * self.tp_size,

164 params_dtype=self.params_dtype,

165 - reduce_results=True,

165 renormalize=True,

166 quant_config=self.quant_config,

167 tp_size=self.tp_size,

vllm/model_executor/models/mixtral.py

-1

@@ -132,7 +132,6 @@ class MixtralMoE(nn.Module):

132 hidden_size=hidden_size,

133 intermediate_size=intermediate_size,

134 params_dtype=params_dtype,

135 - reduce_results=True,

135 renormalize=True,

136 quant_config=quant_config,

137 tp_size=tp_size,

vllm/model_executor/models/nemotron_h.py

+4 -25

@@ -216,7 +216,6 @@ class NemotronHMoE(nn.Module):

216 top_k=config.num_experts_per_tok,

217 hidden_size=self.moe_hidden_size,

218 intermediate_size=config.moe_intermediate_size,

219 - reduce_results=False,

219 renormalize=config.norm_topk_prob,

220 quant_config=quant_config,

221 use_grouped_topk=True,

@@ -231,6 +230,9 @@ class NemotronHMoE(nn.Module):

230 num_redundant_experts=self.n_redundant_experts,

231 is_sequence_parallel=self.is_sequence_parallel,

232 routed_input_transform=self.fc1_latent_proj,

233 + routed_output_transform=self.fc2_latent_proj,

234 + routed_scaling_factor=self.routed_scaling_factor,

235 + apply_routed_scale_to_output=True,

236 router_logits_dtype=self.gate.out_dtype,

237 )

238

@@ -244,38 +246,15 @@ class NemotronHMoE(nn.Module):

246 # router_logits: (num_tokens, n_experts)

247 router_logits, _ = self.gate(hidden_states)

248

247 - # SharedFusedMoE handles:

248 - # - shared experts (with original hidden_states)

249 - # - routed_input_transform (fc1_latent_proj) for latent MoE

250 - # - multistream parallelism between shared and routed experts

251 - shared_output, final_hidden_states = self.experts(

249 + final_hidden_states = self.experts(

250 hidden_states=hidden_states, router_logits=router_logits

251 )

252

255 - # Fix FP16 overflow

256 - # See DeepseekV2DecoderLayer for more details.

257 - if hidden_states.dtype != torch.float16:

258 - final_hidden_states *= self.routed_scaling_factor

259 - elif self.shared_experts is not None:

260 - shared_output *= 1.0 / self.routed_scaling_factor

261 -

262 - # TODO: See SharedFusedMoE.apply_routed_input_transform

263 - # for bandwidth optimization

264 - if self.use_latent_moe:

265 - final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states)

266 -

267 - if self.shared_experts is not None:

268 - final_hidden_states += shared_output

269 -

253 if self.is_sequence_parallel:

254 final_hidden_states = tensor_model_parallel_all_gather(

255 final_hidden_states, 0

256 )

257 final_hidden_states = final_hidden_states[:num_tokens]

275 - elif self.tp_size > 1:

276 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

277 - final_hidden_states

278 - )

258

259 return final_hidden_states.view(num_tokens, hidden_dim)

260

vllm/model_executor/models/olmoe.py

-1

@@ -98,7 +98,6 @@ class OlmoeMoE(nn.Module):

98 top_k=top_k,

99 hidden_size=hidden_size,

100 intermediate_size=intermediate_size,

101 - reduce_results=True,

101 renormalize=False,

102 quant_config=quant_config,

103 tp_size=tp_size,

vllm/model_executor/models/openpangu.py

+3 -22

@@ -206,7 +206,6 @@ class OpenPanguMoE(nn.Module):

206 top_k=config.num_experts_per_tok,

207 hidden_size=config.hidden_size,

208 intermediate_size=config.moe_intermediate_size,

209 - reduce_results=False,

209 renormalize=config.norm_topk_prob,

210 quant_config=quant_config,

211 use_grouped_topk=True,

@@ -214,8 +213,8 @@ class OpenPanguMoE(nn.Module):

213 topk_group=1,

214 prefix=f"{prefix}.experts",

215 scoring_func="sigmoid",

217 - # we do scaling outside, set factor to 1.0 to avoid double mul

218 - routed_scaling_factor=1.0,

216 + routed_scaling_factor=self.routed_scaling_factor,

217 + apply_routed_scale_to_output=True,

218 e_score_correction_bias=self.gate.e_score_correction_bias,

219 enable_eplb=self.enable_eplb,

220 num_redundant_experts=self.n_redundant_experts,

@@ -234,33 +233,15 @@ class OpenPanguMoE(nn.Module):

233

234 router_logits, _ = self.gate(hidden_states)

235

237 - fused_moe_out = self.experts(

236 + final_hidden_states = self.experts(

237 hidden_states=hidden_states, router_logits=router_logits

238 )

239

241 - shared_output, final_hidden_states = fused_moe_out

242 - if self.shared_experts is None:

243 - assert shared_output is None

244 -

245 - if hidden_states.dtype != torch.float16:

246 - final_hidden_states *= self.routed_scaling_factor

247 - elif self.shared_experts is not None:

248 - assert shared_output is not None

249 - shared_output *= 1.0 / self.routed_scaling_factor

250 -

251 - if self.shared_experts is not None:

252 - assert shared_output is not None

253 - final_hidden_states += shared_output

254 -

240 if self.is_sequence_parallel:

241 final_hidden_states = tensor_model_parallel_all_gather(

242 final_hidden_states, 0

243 )

244 final_hidden_states = final_hidden_states[:num_tokens]

260 - elif self.tp_size > 1:

261 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

262 - final_hidden_states

263 - )

245

246 return final_hidden_states.view(num_tokens, hidden_dim)

247

vllm/model_executor/models/param2moe.py

+1 -15

@@ -359,7 +359,6 @@ class Param2MoEMoEBlock(nn.Module):

359 top_k=self.top_k,

360 hidden_size=self.hidden_size,

361 intermediate_size=config.moe_intermediate_size,

362 - reduce_results=False,

362 renormalize=self.norm_expert_prob,

363 quant_config=quant_config,

364 prefix=f"{prefix}.experts",

@@ -388,24 +387,11 @@ class Param2MoEMoEBlock(nn.Module):

387 self.gate.weight.float(),

388 ).to(hidden_states.dtype)

389

391 - final_hidden = self.experts(

390 + expert_output = self.experts(

391 hidden_states=hidden_states,

392 router_logits=router_logits,

393 )

394

396 - if self.shared_experts is not None:

397 - shared_output, expert_output = final_hidden

398 - else:

399 - shared_output, expert_output = None, final_hidden

400 -

401 - if shared_output is not None:

402 - expert_output = expert_output + shared_output

403 -

404 - if self.tp_size > 1:

405 - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(

406 - expert_output

407 - )

408 -

395 return expert_output.view(num_tokens, hidden_dim)

396

397

vllm/model_executor/models/phimoe.py

-1

@@ -281,7 +281,6 @@ class PhiMoE(nn.Module):

281 hidden_size=hidden_size,

282 intermediate_size=intermediate_size,

283 params_dtype=params_dtype,

284 - reduce_results=True,

284 renormalize=False,

285 quant_config=quant_config,

286 tp_size=tp_size,

vllm/model_executor/models/qwen2_moe.py

-7

@@ -170,7 +170,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

170 top_k=config.num_experts_per_tok,

171 hidden_size=config.hidden_size,

172 intermediate_size=config.moe_intermediate_size,

173 - reduce_results=False,

173 renormalize=config.norm_topk_prob,

174 quant_config=quant_config,

175 prefix=f"{prefix}.experts",

@@ -187,12 +186,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

186 final_hidden_states = self.experts(

187 hidden_states=hidden_states, router_logits=router_logits

188 )

190 - if self.shared_expert is not None:

191 - final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

192 - if self.tp_size > 1:

193 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

194 - final_hidden_states

195 - )

189

190 return final_hidden_states.view(orig_shape)

191

vllm/model_executor/models/qwen3_moe.py

+1 -9

@@ -212,7 +212,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module):

212 top_k=config.num_experts_per_tok,

213 hidden_size=config.hidden_size,

214 intermediate_size=config.moe_intermediate_size,

215 - reduce_results=False,

215 renormalize=config.norm_topk_prob,

216 quant_config=quant_config,

217 prefix=f"{prefix}.experts",

@@ -234,22 +233,15 @@ class Qwen3MoeSparseMoeBlock(nn.Module):

233

234 # router_logits: (num_tokens, n_experts)

235 router_logits, _ = self.gate(hidden_states)

237 - shared_out, fused_out = self.experts(

236 + final_hidden_states = self.experts(

237 hidden_states=hidden_states, router_logits=router_logits

238 )

240 - final_hidden_states = (

241 - shared_out + fused_out if shared_out is not None else fused_out

242 - )

239

240 if self.is_sequence_parallel:

241 final_hidden_states = tensor_model_parallel_all_gather(

242 final_hidden_states, 0

243 )

244 final_hidden_states = final_hidden_states[:num_tokens]

249 - elif self.tp_size > 1:

250 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

251 - final_hidden_states

252 - )

245

246 # return to 1d if input is 1d

247 return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states

vllm/model_executor/models/qwen3_next.py

-8

@@ -153,7 +153,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):

153 top_k=config.num_experts_per_tok,

154 hidden_size=config.hidden_size,

155 intermediate_size=config.moe_intermediate_size,

156 - reduce_results=False,

156 renormalize=getattr(config, "norm_topk_prob", True),

157 quant_config=quant_config,

158 prefix=f"{prefix}.experts",

@@ -183,18 +182,11 @@ class Qwen3NextSparseMoeBlock(nn.Module):

182 hidden_states=hidden_states, router_logits=router_logits

183 )

184

186 - if self.shared_expert is not None:

187 - final_hidden_states = final_hidden_states[0] + final_hidden_states[1]

188 -

185 if self.is_sequence_parallel:

186 final_hidden_states = tensor_model_parallel_all_gather(

187 final_hidden_states, 0

188 )

189 final_hidden_states = final_hidden_states[:num_tokens]

194 - elif self.tp_size > 1:

195 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

196 - final_hidden_states

197 - )

190

191 return final_hidden_states.view(orig_shape)

192

vllm/model_executor/models/sarvam.py

+1 -15

@@ -341,7 +341,6 @@ class SarvamMLAMoE(nn.Module):

341 top_k=self.top_k,

342 hidden_size=self.hidden_size,

343 intermediate_size=config.moe_intermediate_size,

344 - reduce_results=False,

344 renormalize=self.norm_expert_prob,

345 quant_config=quant_config,

346 prefix=f"{prefix}.experts",

@@ -370,20 +369,7 @@ class SarvamMLAMoE(nn.Module):

369 router_logits=router_logits,

370 )

371

373 - if self.shared_experts is not None:

374 - shared_output, expert_output = final_hidden

375 - else:

376 - shared_output, expert_output = None, final_hidden

377 -

378 - if shared_output is not None:

379 - expert_output = expert_output + shared_output

380 -

381 - if self.tp_size > 1:

382 - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(

383 - expert_output

384 - )

385 -

386 - return expert_output.view(num_tokens, hidden_dim)

372 + return final_hidden.view(num_tokens, hidden_dim)

373

374

375 class SarvamMLABlock(nn.Module):

vllm/model_executor/models/step3_text.py

-4

@@ -14,7 +14,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig

14 from vllm.distributed import (

15 get_pp_group,

16 get_tensor_model_parallel_world_size,

17 - tensor_model_parallel_all_reduce,

17 )

18 from vllm.logger import init_logger

19 from vllm.model_executor.layers.activation import SiluAndMul

@@ -71,7 +70,6 @@ class FusedMoEBlock(nn.Module):

70 top_k=config.moe_top_k,

71 hidden_size=config.hidden_size,

72 intermediate_size=config.moe_intermediate_size,

74 - reduce_results=False,

73 renormalize=config.norm_expert_weight,

74 quant_config=quant_config,

75 prefix=f"{prefix}.experts",

@@ -94,8 +92,6 @@ class FusedMoEBlock(nn.Module):

92 final_hidden_states = self.experts(

93 hidden_states=hidden_states, router_logits=router_logits

94 )

97 - if self.tp_size > 1:

98 - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

96 return final_hidden_states.view(orig_shape)

vllm/model_executor/models/step3p5.py

+3 -18

@@ -379,7 +379,6 @@ class FusedMoEBlock(nn.Module):

379 top_k=config.moe_top_k,

380 hidden_size=config.hidden_size,

381 intermediate_size=config.moe_intermediate_size,

382 - reduce_results=False,

382 renormalize=config.norm_expert_weight,

383 quant_config=quant_config,

384 activation=activation,

@@ -397,30 +396,16 @@ class FusedMoEBlock(nn.Module):

396 hidden_states = hidden_states.view(-1, hidden_dim)

397

398 if self.experts.is_internal_router:

400 - # In this case, the gate/router runs inside the FusedMoE class

401 - fused_moe_out = self.experts(

399 + final_hidden_states = self.experts(

400 hidden_states=hidden_states, router_logits=hidden_states

401 )

402 else:

405 - # router_logits: (num_tokens, n_experts)

403 + # TODO(bnell): this gate could be moved into the FusedMoE?

404 router_logits, _ = self.gate(hidden_states)

407 - fused_moe_out = self.experts(

405 + final_hidden_states = self.experts(

406 hidden_states=hidden_states, router_logits=router_logits

407 )

408

411 - shared_output, final_hidden_states = fused_moe_out

412 - if self.share_expert is None:

413 - assert shared_output is None

414 -

415 - if self.share_expert is not None:

416 - assert shared_output is not None

417 - final_hidden_states += shared_output

418 -

419 - if self.tp_size > 1:

420 - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

421 - final_hidden_states

422 - )

423 -

409 return final_hidden_states.view(num_tokens, hidden_dim)

410

411

vllm/model_executor/models/transformers/moe.py

+3 -27

@@ -204,8 +204,6 @@ class MoEMixin(MixtureOfExperts):

204 )

205 assert intermediate_size is not None

206

207 - # If there are shared experts, the results are

208 - # reduced after mlp.forward() not inside FusedMoE

207 num_shared_experts = getattr_iter(

208 text_config,

209 [

@@ -214,17 +212,6 @@ class MoEMixin(MixtureOfExperts):

212 ],

213 0,

214 )

217 - reduce_results = num_shared_experts == 0

218 -

219 - def add_all_reduce(mlp: nn.Module):

220 - """Adds an all-reduce to the output of `mlp.forward()`."""

221 -

222 - class MLPWithAllReduce(mlp.__class__):

223 - def forward(self, *args, **kwargs):

224 - output = super().forward(*args, **kwargs)

225 - return self.experts.maybe_all_reduce_tensor_model_parallel(output)

226 -

227 - mlp.__class__ = MLPWithAllReduce

215

216 # Unused kwargs since we use custom_routing_function:

217 # - `scoring_func` and `e_score_correction_bias` only used for grouped

@@ -289,14 +276,11 @@ class MoEMixin(MixtureOfExperts):

276 if "bias" in experts_param_name:

277 has_bias = True

278 break

292 - # Double check there are no shared experts

293 - nonlocal reduce_results

294 - if reduce_results:

279 + # If the config does not specify num_shared_experts, but

280 + # the model has shared experts, we assume there is one.

281 + if self.num_shared_experts == 0:

282 for mlp_param_name, _ in mlp.named_parameters():

283 if "shared_expert" in mlp_param_name:

297 - reduce_results = False

298 - # If the config does not specify num_shared_experts, but

299 - # the model has shared experts, we assume there is one.

284 self.num_shared_experts = 1

285 break

286 # Replace experts module with FusedMoE

@@ -305,7 +289,6 @@ class MoEMixin(MixtureOfExperts):

289 top_k=top_k,

290 hidden_size=hidden_size,

291 intermediate_size=intermediate_size,

308 - reduce_results=reduce_results,

292 renormalize=renormalize,

293 # Hard coded because topk happens in Transformers

294 use_grouped_topk=False,

@@ -326,13 +309,6 @@ class MoEMixin(MixtureOfExperts):

309 self.moe_layers.append(fused_experts)

310 self.expert_weights.append(fused_experts.get_expert_weights())

311 self.num_moe_layers += 1

329 - # If results are not all-reduced in FusedMoE, ensure they

330 - # are all-reduced at the end of mlp.forward() if tensor

331 - # parallel or expert parallel is enabled

332 - if not reduce_results and (

333 - fused_experts.tp_size > 1 or fused_experts.ep_size > 1

334 - ):

335 - add_all_reduce(mlp)

312 else:

313 _recursive_replace(child_module, prefix=qual_name)

314