5e584ce9ec · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/moe/test_moe_layer.py

+3 -7

@@ -37,7 +37,7 @@ from vllm.distributed.parallel_state import (

37 get_eplb_group,

38 )

39 from vllm.forward_context import set_forward_context

40 -from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE, fused_experts

40 +from vllm.model_executor.layers.fused_moe import FusedMoE, fused_experts

41 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

42 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

43 from vllm.model_executor.layers.fused_moe.router.router_factory import (

@@ -858,11 +858,7 @@ def make_fused_moe_layer(

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

859

860 kwargs = dict()

861 - if shared_experts is None:

862 - builder = FusedMoE

863 - else:

864 - builder = SharedFusedMoE

865 - kwargs["shared_experts"] = shared_experts

861 + kwargs["shared_experts"] = shared_experts

862

863 # Add gate and routed_input_transform if provided

864 if gate is not None:

@@ -872,7 +868,7 @@ def make_fused_moe_layer(

868 kwargs["routed_input_transform"] = routed_input_transform

869 kwargs["routed_output_transform"] = routed_output_transform

870

875 - layer = builder(

871 + layer = FusedMoE(

872 num_experts=global_num_experts,

873 top_k=top_k,

874 hidden_size=hidden_size,

tests/kernels/moe/test_shared_fused_moe_routed_transform.py

+10 -10

@@ -1,9 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 """

4 -Tests for SharedFusedMoE with routed_input_transform.

4 +Tests for FusedMoE with routed_input_transform.

5

6 -Verifies that applying routed_input_transform inside SharedFusedMoE

6 +Verifies that applying routed_input_transform inside FusedMoE

7 produces the same results as applying the transform manually outside.

8 """

9

@@ -13,7 +13,7 @@ import torch.nn as nn

13

14 from vllm.config import VllmConfig, set_current_vllm_config

15 from vllm.forward_context import set_forward_context

16 -from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

16 +from vllm.model_executor.layers.fused_moe import FusedMoE

17 from vllm.platforms import current_platform

18 from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed

19

@@ -133,9 +133,9 @@ def test_routed_input_transform_inside_vs_outside(

133 workspace_init,

134 monkeypatch,

135 ):

136 - """Compare SharedFusedMoE with transform inside vs manually applying outside.

137 - Method A (inside): SharedFusedMoE with routed_input_transform

138 - Method B (outside): Manually transform, then SharedFusedMoE without transform

136 + """Compare FusedMoE with transform inside vs manually applying outside.

137 + Method A (inside): FusedMoE with routed_input_transform

138 + Method B (outside): Manually transform, then FusedMoE without transform

139 """

140 if current_platform.is_rocm():

141 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")

@@ -157,8 +157,8 @@ def test_routed_input_transform_inside_vs_outside(

157 routed_transform = SimpleLinear(hidden_size, latent_size, dtype)

158

159 with set_current_vllm_config(vllm_config):

160 - # Method A: SharedFusedMoE WITH routed_input_transform

161 - moe_with_transform = SharedFusedMoE(

160 + # Method A: FusedMoE WITH routed_input_transform

161 + moe_with_transform = FusedMoE(

162 shared_experts=shared_experts,

163 routed_input_transform=routed_transform,

164 num_experts=num_experts,

@@ -173,9 +173,9 @@ def test_routed_input_transform_inside_vs_outside(

173 prefix="moe_with_transform",

174 )

175

176 - # Method B: SharedFusedMoE WITHOUT routed_input_transform

176 + # Method B: FusedMoE WITHOUT routed_input_transform

177 # Note: shared_experts=None because when transform is done outside,

178 - moe_without_transform = SharedFusedMoE(

178 + moe_without_transform = FusedMoE(

179 shared_experts=None,

180 routed_input_transform=None,

181 num_experts=num_experts,

vllm/distributed/device_communicators/base_device_communicator.py

+3 -10

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

7 import torch.distributed as dist

8 from torch.distributed import ProcessGroup

9

10 +from vllm.utils import is_moe_layer

11 +

12

13 class Cache:

14 def __init__(self):

@@ -317,16 +319,7 @@ class DeviceCommunicatorBase:

319 if not self.is_ep_communicator:

320 return

321

320 - moe_modules = [

321 - module

322 - for module in model.modules()

323 - # TODO(bnell): Should use isinstance but can't. Maybe search for

324 - # presence of quant_method.maybe_init_modular_kernel?

325 - if (

326 - module.__class__.__name__ == "FusedMoE"

327 - or module.__class__.__name__ == "SharedFusedMoE"

328 - )

329 - ]

322 + moe_modules = [module for module in model.modules() if is_moe_layer(module)]

323 for module in moe_modules:

324 module.maybe_init_modular_kernel()

325

vllm/distributed/elastic_ep/elastic_execute.py

+2 -4

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

38 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

39 from vllm.logger import init_logger

40 from vllm.model_executor.layers.fused_moe.layer import FusedMoEParallelConfig

41 +from vllm.utils import is_moe_layer

42 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

43 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

44 from vllm.v1.worker.workspace import lock_workspace, unlock_workspace

@@ -319,10 +320,7 @@ class ElasticEPScalingExecutor:

320 moe_modules = [

321 module

322 for module in self.worker.model_runner.model.modules()

322 - if (

323 - module.__class__.__name__ == "FusedMoE"

324 - or module.__class__.__name__ == "SharedFusedMoE"

325 - )

323 + if is_moe_layer(module)

324 ]

325 num_local_experts = moe_modules[0].moe_config.num_local_experts

326 assert all(

vllm/lora/layers/fused_moe.py

+2 -2

@@ -610,7 +610,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

610 ) -> bool:

611 """Returns True if the layer can be replaced by this LoRA layer."""

612

613 - # source_layer is FusedMoE or SharedFusedMoE

613 + # source_layer is FusedMoE

614 return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 2

615

616

@@ -772,5 +772,5 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):

772 model_config: PretrainedConfig | None = None,

773 ) -> bool:

774 """Returns True if the layer can be replaced by this LoRA layer."""

775 - # source_layer is FusedMoE or SharedFusedMoE

775 + # source_layer is FusedMoE

776 return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 1

vllm/model_executor/layers/fused_moe/__init__.py

-2

@@ -29,7 +29,6 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (

29 FusedMoERouter,

30 )

31 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear

32 -from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

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

33 UnquantizedFusedMoEMethod,

34 )

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

63 "FusedMoEPrepareAndFinalizeModular",

64 "GateLinear",

65 "RoutingMethodType",

67 - "SharedFusedMoE",

66 "activation_without_mul",

67 "apply_moe_activation",

68 "override_config",

vllm/model_executor/layers/fused_moe/shared_fused_moe.py deleted

-25

@@ -1,25 +0,0 @@

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

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

3 -

4 -import torch

5 -

6 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE

7 -

8 -

9 -# TODO(bnell): Remove this entirely

10 -class SharedFusedMoE(FusedMoE):

11 - """

12 - A FusedMoE operation that also computes the results of shared experts.

13 - If an all2all communicator is being used the shared expert computation

14 - can be interleaved with the fused all2all dispatch communication step.

15 - """

16 -

17 - def forward(

18 - self,

19 - hidden_states: torch.Tensor,

20 - router_logits: torch.Tensor,

21 - ) -> torch.Tensor:

22 - return super().forward(

23 - hidden_states=hidden_states,

24 - router_logits=router_logits,

25 - )

vllm/model_executor/models/AXK1.py

+4 -4

@@ -42,7 +42,7 @@ from vllm.distributed import (

42 )

43 from vllm.logger import init_logger

44 from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

45 +from vllm.model_executor.layers.fused_moe import FusedMoE

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 ColumnParallelLinear,

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

163 prefix=f"{prefix}.shared_experts",

164 )

165

166 - self.experts = SharedFusedMoE(

166 + self.experts = FusedMoE(

167 shared_experts=self.shared_experts,

168 gate=self.gate,

169 num_experts=config.n_routed_experts,

@@ -916,7 +916,7 @@ class AXK1ForCausalLM(

916 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

918 # (param_name, weight_name, expert_id, shard_id)

919 - return SharedFusedMoE.make_expert_params_mapping(

919 + return FusedMoE.make_expert_params_mapping(

920 self,

921 ckpt_gate_proj_name="gate_proj",

922 ckpt_down_proj_name="down_proj",

@@ -950,7 +950,7 @@ class AXK1ForCausalLM(

950

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

952 # (param_name, weight_name, expert_id, shard_id)

953 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

953 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

954 self,

955 ckpt_gate_proj_name="gate_proj",

956 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/afmoe.py

+5 -5

@@ -18,7 +18,7 @@ from vllm.distributed import (

18 )

19 from vllm.logger import init_logger

20 from vllm.model_executor.layers.attention import Attention

21 -from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

21 +from vllm.model_executor.layers.fused_moe import FusedMoE

22 from vllm.model_executor.layers.layernorm import RMSNorm

23 from vllm.model_executor.layers.linear import (

24 ColumnParallelLinear,

@@ -124,8 +124,8 @@ class AfmoeMoE(nn.Module):

124 prefix=f"{prefix}.shared_experts",

125 )

126

127 - # Routed experts using SharedFusedMoE

128 - self.experts = SharedFusedMoE(

127 + # Routed experts using FusedMoE

128 + self.experts = FusedMoE(

129 shared_experts=self.shared_experts,

130 num_experts=config.num_experts,

131 top_k=config.num_experts_per_tok,

@@ -479,7 +479,7 @@ class AfmoeModel(nn.Module, EagleModelMixin):

479 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

481 # (param_name, weight_name, expert_id, shard_id)

482 - return SharedFusedMoE.make_expert_params_mapping(

482 + return FusedMoE.make_expert_params_mapping(

483 self,

484 ckpt_gate_proj_name="gate_proj",

485 ckpt_down_proj_name="down_proj",

@@ -637,7 +637,7 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

637 self.num_moe_layers = config.num_hidden_layers - config.num_dense_layers

638 self.num_expert_groups = config.n_group

639

640 - self.moe_layers: list[SharedFusedMoE] = []

640 + self.moe_layers: list[FusedMoE] = []

641 example_moe = None

642 for layer in self.model.layers:

643 if isinstance(layer, PPMissingLayer):

vllm/model_executor/models/aria.py

+2 -2

@@ -14,7 +14,7 @@ from vllm.config.multimodal import BaseDummyOptions

14 from vllm.distributed import get_tensor_model_parallel_rank

15 from vllm.inputs import MultiModalDataDict

16 from vllm.model_executor.layers.activation import get_act_fn

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

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

18 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear

19 from vllm.model_executor.layers.logits_processor import LogitsProcessor

20 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -214,7 +214,7 @@ class AriaProjector(nn.Module):

214 return out

215

216

217 -class AriaFusedMoE(SharedFusedMoE):

217 +class AriaFusedMoE(FusedMoE):

218 def weight_loader(

219 self, param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str

220 ) -> None:

vllm/model_executor/models/bailing_moe.py

+3 -3

@@ -41,7 +41,7 @@ from vllm.distributed import (

41 )

42 from vllm.model_executor.layers.activation import SiluAndMul

43 from vllm.model_executor.layers.attention import Attention

44 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

44 +from vllm.model_executor.layers.fused_moe import FusedMoE

45 from vllm.model_executor.layers.layernorm import RMSNorm

46 from vllm.model_executor.layers.linear import (

47 MergedColumnParallelLinear,

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

285 else:

286 self.shared_experts = None

287

288 - self.experts = SharedFusedMoE(

288 + self.experts = FusedMoE(

289 shared_experts=self.shared_experts,

290 num_experts=self.num_experts,

291 top_k=self.top_k,

@@ -461,7 +461,7 @@ class BailingMoeModel(nn.Module):

461 return hidden_states

462

463 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

464 - return SharedFusedMoE.make_expert_params_mapping(

464 + return FusedMoE.make_expert_params_mapping(

465 self,

466 ckpt_gate_proj_name="gate_proj",

467 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/bailing_moe_linear.py

+3 -3

@@ -21,7 +21,7 @@ from vllm.model_executor.layers.fla.ops.layernorm_guard import (

21 RMSNormGated,

22 layernorm_fn,

23 )

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

24 +from vllm.model_executor.layers.fused_moe import FusedMoE

25 from vllm.model_executor.layers.layernorm import RMSNorm

26 from vllm.model_executor.layers.linear import (

27 ColumnParallelLinear,

@@ -351,8 +351,8 @@ class BailingMoeV25(nn.Module):

351 else:

352 self.shared_experts = None

353

354 - # Routed experts using SharedFusedMoE

355 - self.experts = SharedFusedMoE(

354 + # Routed experts using FusedMoE

355 + self.experts = FusedMoE(

356 shared_experts=self.shared_experts,

357 num_experts=self.num_experts,

358 top_k=self.top_k,

vllm/model_executor/models/deepseek_mtp.py

+2 -2

@@ -11,7 +11,7 @@ from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import VllmConfig

13 from vllm.logger import init_logger

14 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

14 +from vllm.model_executor.layers.fused_moe import FusedMoE

15 from vllm.model_executor.layers.layernorm import RMSNorm

16 from vllm.model_executor.layers.logits_processor import LogitsProcessor

17 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -252,7 +252,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

252 ]

253 stacked_params_mapping.extend(indexer_fused_mapping)

254

255 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

255 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

256 self,

257 ckpt_gate_proj_name="gate_proj",

258 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/deepseek_v2.py

+4 -4

@@ -48,9 +48,9 @@ from vllm.model_executor.layers.activation import SiluAndMul

48 from vllm.model_executor.layers.attention import Attention

49 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

50 from vllm.model_executor.layers.fused_moe import (

51 + FusedMoE,

52 GateLinear,

53 RoutingMethodType,

53 - SharedFusedMoE,

54 )

55 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm

56 from vllm.model_executor.layers.linear import (

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

311 prefix=f"{prefix}.shared_experts",

312 )

313

314 - self.experts = SharedFusedMoE(

314 + self.experts = FusedMoE(

315 shared_experts=self.shared_experts,

316 gate=self.gate,

317 num_experts=config.n_routed_experts,

@@ -1432,7 +1432,7 @@ class DeepseekV2ForCausalLM(

1432 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

1434 # (param_name, weight_name, expert_id, shard_id)

1435 - return SharedFusedMoE.make_expert_params_mapping(

1435 + return FusedMoE.make_expert_params_mapping(

1436 self,

1437 ckpt_gate_proj_name="gate_proj",

1438 ckpt_down_proj_name="down_proj",

@@ -1474,7 +1474,7 @@ class DeepseekV2ForCausalLM(

1474

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

1476 # (param_name, weight_name, expert_id, shard_id)

1477 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

1477 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

1478 self,

1479 ckpt_gate_proj_name="gate_proj",

1480 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/dots1.py

+3 -3

@@ -40,7 +40,7 @@ from vllm.distributed import (

40 )

41 from vllm.model_executor.layers.activation import SiluAndMul

42 from vllm.model_executor.layers.attention import Attention

43 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

43 +from vllm.model_executor.layers.fused_moe import FusedMoE

44 from vllm.model_executor.layers.layernorm import RMSNorm

45 from vllm.model_executor.layers.linear import (

46 MergedColumnParallelLinear,

@@ -155,7 +155,7 @@ class Dots1MoE(nn.Module):

155 else:

156 self.shared_experts = None

157

158 - self.experts = SharedFusedMoE(

158 + self.experts = FusedMoE(

159 shared_experts=self.shared_experts,

160 num_experts=config.n_routed_experts,

161 top_k=config.num_experts_per_tok,

@@ -413,7 +413,7 @@ class Dots1Model(nn.Module):

413 return hidden_states

414

415 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

416 - return SharedFusedMoE.make_expert_params_mapping(

416 + return FusedMoE.make_expert_params_mapping(

417 self,

418 ckpt_gate_proj_name="gate_proj",

419 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/ernie45_moe.py

+4 -4

@@ -42,7 +42,7 @@ from vllm.distributed import (

42 from vllm.logger import init_logger

43 from vllm.model_executor.layers.activation import SiluAndMul

44 from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

45 +from vllm.model_executor.layers.fused_moe import FusedMoE

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 MergedColumnParallelLinear,

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

188 else:

189 self.shared_experts = None

190

191 - self.experts = SharedFusedMoE(

191 + self.experts = FusedMoE(

192 shared_experts=self.shared_experts,

193 num_experts=config.moe_num_experts,

194 top_k=config.moe_k,

@@ -485,7 +485,7 @@ class Ernie4_5_MoeModel(nn.Module):

485 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

487 # (param_name, weight_name, expert_id, shard_id)

488 - return SharedFusedMoE.make_expert_params_mapping(

488 + return FusedMoE.make_expert_params_mapping(

489 self,

490 ckpt_gate_proj_name="gate_proj",

491 ckpt_down_proj_name="down_proj",

@@ -667,7 +667,7 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe

667 self.num_moe_layers = len(moe_layers_indices)

668 self.num_expert_groups = 1

669

670 - self.moe_layers: list[SharedFusedMoE] = []

670 + self.moe_layers: list[FusedMoE] = []

671 example_moe = None

672 for layer in self.model.layers:

673 if isinstance(layer, PPMissingLayer):

vllm/model_executor/models/ernie45_vl_moe.py

+4 -4

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

36 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

37 from vllm.logger import init_logger

38 from vllm.model_executor.layers.attention import Attention

39 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

39 +from vllm.model_executor.layers.fused_moe import FusedMoE

40 from vllm.model_executor.layers.layernorm import RMSNorm

41 from vllm.model_executor.layers.linear import (

42 QKVParallelLinear,

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

257 prefix=f"{prefix}.text_experts_gate",

258 )

259

260 - self.text_experts = SharedFusedMoE(

260 + self.text_experts = FusedMoE(

261 shared_experts=self.shared_experts,

262 num_experts=config.moe_num_experts[0],

263 top_k=config.moe_k,

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

294 prefix=f"{prefix}.vision_experts_gate",

295 )

296

297 - self.vision_experts = SharedFusedMoE(

297 + self.vision_experts = FusedMoE(

298 shared_experts=self.shared_experts,

299 num_experts=config.moe_num_experts[1],

300 top_k=config.moe_k,

@@ -649,7 +649,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

649

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

651 # (param_name, weight_name, expert_id, shard_id)

652 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

652 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

653 self,

654 ckpt_gate_proj_name="gate_proj",

655 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/exaone_moe.py

+1 -2

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

34 from vllm.model_executor.layers.layernorm import RMSNorm

35 from vllm.model_executor.layers.linear import ReplicatedLinear

36 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -130,7 +129,7 @@ class ExaoneMoe(nn.Module):

129 else:

130 self.shared_experts = None

131

133 - self.experts = SharedFusedMoE(

132 + self.experts = FusedMoE(

133 shared_experts=self.shared_experts,

134 gate=self.gate,

135 num_experts=self.n_routed_experts,

vllm/model_executor/models/glm4_moe.py

+3 -3

@@ -42,7 +42,7 @@ from vllm.distributed import (

42 from vllm.logger import init_logger

43 from vllm.model_executor.layers.activation import SiluAndMul

44 from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

45 +from vllm.model_executor.layers.fused_moe import FusedMoE

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 MergedColumnParallelLinear,

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

178 else:

179 self.shared_experts = None

180

181 - self.experts = SharedFusedMoE(

181 + self.experts = FusedMoE(

182 shared_experts=self.shared_experts,

183 num_experts=config.n_routed_experts,

184 top_k=config.num_experts_per_tok,

@@ -466,7 +466,7 @@ class Glm4MoeModel(nn.Module):

466 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

468 # (param_name, weight_name, expert_id, shard_id)

469 - return SharedFusedMoE.make_expert_params_mapping(

469 + return FusedMoE.make_expert_params_mapping(

470 self,

471 ckpt_gate_proj_name="gate_proj",

472 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/glm4_moe_lite.py

+4 -4

@@ -41,7 +41,7 @@ from vllm.distributed import (

41 get_pp_group,

42 )

43 from vllm.logger import init_logger

44 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

44 +from vllm.model_executor.layers.fused_moe import FusedMoE

45 from vllm.model_executor.layers.layernorm import RMSNorm

46 from vllm.model_executor.layers.logits_processor import LogitsProcessor

47 from vllm.model_executor.layers.vocab_parallel_embedding import (

@@ -308,7 +308,7 @@ class Glm4MoeLiteModel(nn.Module):

308 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

310 # (param_name, weight_name, expert_id, shard_id)

311 - return SharedFusedMoE.make_expert_params_mapping(

311 + return FusedMoE.make_expert_params_mapping(

312 self,

313 ckpt_gate_proj_name="gate_proj",

314 ckpt_down_proj_name="down_proj",

@@ -334,7 +334,7 @@ class Glm4MoeLiteModel(nn.Module):

334

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

336 # (param_name, weight_name, expert_id, shard_id)

337 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

337 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

338 self,

339 ckpt_gate_proj_name="gate_proj",

340 ckpt_down_proj_name="down_proj",

@@ -616,7 +616,7 @@ class Glm4MoeLiteForCausalLM(

616 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

618 # (param_name, weight_name, expert_id, shard_id)

619 - return SharedFusedMoE.make_expert_params_mapping(

619 + return FusedMoE.make_expert_params_mapping(

620 self,

621 ckpt_gate_proj_name="gate_proj",

622 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/glm4_moe_lite_mtp.py

+2 -2

@@ -32,7 +32,7 @@ from transformers import PretrainedConfig

32

33 from vllm._aiter_ops import rocm_aiter_ops

34 from vllm.config import VllmConfig

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

35 +from vllm.model_executor.layers.fused_moe import FusedMoE

36 from vllm.model_executor.layers.layernorm import RMSNorm

37 from vllm.model_executor.layers.logits_processor import LogitsProcessor

38 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -260,7 +260,7 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

260 ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),

261 ]

262

263 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

263 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

264 self,

265 ckpt_gate_proj_name="gate_proj",

266 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/hunyuan_v1.py

+3 -3

@@ -42,7 +42,7 @@ from vllm.distributed import (

42 )

43 from vllm.model_executor.layers.activation import SiluAndMul

44 from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

45 +from vllm.model_executor.layers.fused_moe import FusedMoE

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 ColumnParallelLinear,

@@ -438,7 +438,7 @@ class HunYuanSparseMoeBlock(nn.Module):

438 else:

439 self.shared_mlp = None

440

441 - self.experts = SharedFusedMoE(

441 + self.experts = FusedMoE(

442 shared_experts=self.shared_mlp,

443 num_experts=self.n_routed_experts,

444 top_k=top_k,

@@ -712,7 +712,7 @@ class HunYuanModel(nn.Module, EagleModelMixin):

712 if _is_moe(self.config):

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

714 # (param_name, weight_name, expert_id, shard_id)

715 - return SharedFusedMoE.make_expert_params_mapping(

715 + return FusedMoE.make_expert_params_mapping(

716 self,

717 ckpt_gate_proj_name="gate_proj",

718 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/kimi_linear.py

+3 -3

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

14 )

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.activation import SiluAndMul

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

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

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 (

@@ -144,7 +144,7 @@ class KimiMoE(nn.Module):

144 else:

145 self.shared_experts = None

146

147 - self.experts = SharedFusedMoE(

147 + self.experts = FusedMoE(

148 shared_experts=self.shared_experts,

149 num_experts=num_experts,

150 top_k=config.num_experts_per_token,

@@ -476,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)

479 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

479 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

480 self,

481 ckpt_gate_proj_name="w1",

482 ckpt_down_proj_name="w2",

vllm/model_executor/models/llama4.py

+5 -5

@@ -36,7 +36,7 @@ from vllm.model_executor.layers.attention import (

36 Attention,

37 ChunkedLocalAttention,

38 )

39 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

39 +from vllm.model_executor.layers.fused_moe import FusedMoE

40 from vllm.model_executor.layers.layernorm import RMSNorm

41 from vllm.model_executor.layers.linear import (

42 QKVParallelLinear,

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

127 self.n_physical_experts = self.n_local_experts + self.n_redundant_experts

128 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

129

130 - self.experts = SharedFusedMoE(

130 + self.experts = FusedMoE(

131 shared_experts=self.shared_expert,

132 num_experts=config.num_local_experts,

133 top_k=config.num_experts_per_tok,

@@ -414,7 +414,7 @@ class Llama4Model(LlamaModel):

414 params_dict: The dictionary of module parameters.

415 loaded_params: The set of already loaded parameters.

416 expert_params_mapping: The mapping of expert parameters. Must be

417 - generated by SharedFusedMoE.make_expert_params_mapping().

417 + generated by FusedMoE.make_expert_params_mapping().

418 fused: Whether the expert weights are fused into a single weight

419 tensor or are separate weight tensors for each expert.

420 When fused is True, loaded_weight should have shape of:

@@ -554,7 +554,7 @@ class Llama4Model(LlamaModel):

554 fused_experts_params = False

555 # Expert parameter mapping for the case where the expert weights are

556 # not fused into a single weight tensor.

557 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

557 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

558 self,

559 ckpt_gate_proj_name="gate_proj",

560 ckpt_down_proj_name="down_proj",

@@ -564,7 +564,7 @@ class Llama4Model(LlamaModel):

564 )

565 # Expert parameter mapping for the case where the expert weights are

566 # fused into a single weight tensor.

567 - expert_params_mapping_fused = SharedFusedMoE.make_expert_params_mapping(

567 + expert_params_mapping_fused = FusedMoE.make_expert_params_mapping(

568 self,

569 ckpt_gate_proj_name="gate_up_proj",

570 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/nemotron_h.py

+3 -3

@@ -34,8 +34,8 @@ from vllm.distributed.parallel_state import get_pp_group

34 from vllm.model_executor.layers.activation import ReLUSquaredActivation

35 from vllm.model_executor.layers.attention import Attention

36 from vllm.model_executor.layers.fused_moe import (

37 + FusedMoE,

38 GateLinear,

38 - SharedFusedMoE,

39 activation_without_mul,

40 )

41 from vllm.model_executor.layers.layernorm import RMSNorm

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

210 self.fc1_latent_proj = None

211 self.fc2_latent_proj = None

212

213 - self.experts = SharedFusedMoE(

213 + self.experts = FusedMoE(

214 shared_experts=self.shared_experts,

215 num_experts=config.n_routed_experts,

216 top_k=config.num_experts_per_tok,

@@ -652,7 +652,7 @@ class NemotronHModel(nn.Module):

652 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

653 if self.has_moe:

654 # (param_name, weight_name, expert_id, shard_id)

655 - expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

655 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

656 # - FusedMoe.w1 (aka gate_proj) should be up_proj since that's

657 # what the activation is applied to

658 # - FusedMoe.w3 (aka up_proj) should be ignored since we're

vllm/model_executor/models/openpangu.py

+3 -3

@@ -44,7 +44,7 @@ from vllm.model_executor.layers.attention import (

44 Attention,

45 StaticSinkAttention,

46 )

47 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

47 +from vllm.model_executor.layers.fused_moe import FusedMoE

48 from vllm.model_executor.layers.layernorm import RMSNorm

49 from vllm.model_executor.layers.linear import (

50 ColumnParallelLinear,

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

200 else:

201 self.shared_experts = None

202

203 - self.experts = SharedFusedMoE(

203 + self.experts = FusedMoE(

204 shared_experts=self.shared_experts,

205 num_experts=config.n_routed_experts,

206 top_k=config.num_experts_per_tok,

@@ -1149,7 +1149,7 @@ class OpenPanguModel(nn.Module):

1149 ]

1150 has_experts = hasattr(self.config, "n_routed_experts")

1151 if has_experts:

1152 - expert_merge_mapping = SharedFusedMoE.make_expert_params_mapping(

1152 + expert_merge_mapping = FusedMoE.make_expert_params_mapping(

1153 self,

1154 ckpt_gate_proj_name="gate_proj",

1155 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/param2moe.py

+4 -4

@@ -32,7 +32,7 @@ from vllm.distributed import (

32 )

33 from vllm.model_executor.layers.activation import SiluAndMul

34 from vllm.model_executor.layers.attention import Attention

35 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

35 +from vllm.model_executor.layers.fused_moe import FusedMoE

36 from vllm.model_executor.layers.layernorm import RMSNorm

37 from vllm.model_executor.layers.linear import (

38 MergedColumnParallelLinear,

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

353 else:

354 self.shared_experts = None # type: ignore[assignment]

355

356 - self.experts = SharedFusedMoE(

356 + self.experts = FusedMoE(

357 shared_experts=self.shared_experts,

358 num_experts=self.num_experts,

359 top_k=self.top_k,

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

370 routed_scaling_factor=self.routed_scaling_factor,

371 )

372

373 - def maybe_get_fused_moe(self) -> SharedFusedMoE:

373 + def maybe_get_fused_moe(self) -> FusedMoE:

374 return self.experts

375

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

@@ -690,7 +690,7 @@ class Param2MoEModel(nn.Module):

690 return loaded_params

691

692 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

693 - return SharedFusedMoE.make_expert_params_mapping(

693 + return FusedMoE.make_expert_params_mapping(

694 self,

695 ckpt_gate_proj_name="gate_proj",

696 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/qwen2_moe.py

+3 -3

@@ -40,7 +40,7 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

40 from vllm.logger import init_logger

41 from vllm.model_executor.layers.activation import SiluAndMul

42 from vllm.model_executor.layers.attention import Attention

43 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

43 +from vllm.model_executor.layers.fused_moe import FusedMoE

44 from vllm.model_executor.layers.layernorm import RMSNorm

45 from vllm.model_executor.layers.linear import (

46 MergedColumnParallelLinear,

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

164 else:

165 self.shared_expert = None

166

167 - self.experts = SharedFusedMoE(

167 + self.experts = FusedMoE(

168 shared_experts=self.shared_expert,

169 num_experts=config.num_experts,

170 top_k=config.num_experts_per_tok,

@@ -418,7 +418,7 @@ class Qwen2MoeModel(nn.Module):

418 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

420 # (param_name, weight_name, expert_id, shard_id)

421 - return SharedFusedMoE.make_expert_params_mapping(

421 + return FusedMoE.make_expert_params_mapping(

422 self,

423 ckpt_gate_proj_name="gate_proj",

424 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/qwen3_moe.py

+3 -3

@@ -43,7 +43,7 @@ from vllm.distributed import (

43 from vllm.logger import init_logger

44 from vllm.model_executor.layers.activation import SiluAndMul

45 from vllm.model_executor.layers.attention import Attention

46 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

46 +from vllm.model_executor.layers.fused_moe import FusedMoE

47 from vllm.model_executor.layers.layernorm import RMSNorm

48 from vllm.model_executor.layers.linear import (

49 MergedColumnParallelLinear,

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

205 self.shared_expert_gate = None

206 self.shared_expert = None

207

208 - self.experts = SharedFusedMoE(

208 + self.experts = FusedMoE(

209 shared_experts=self.shared_expert,

210 gate=self.gate,

211 num_experts=self.n_routed_experts,

@@ -508,7 +508,7 @@ class Qwen3MoeModel(nn.Module, EagleModelMixin):

508 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

510 # (param_name, weight_name, expert_id, shard_id)

511 - return SharedFusedMoE.make_expert_params_mapping(

511 + return FusedMoE.make_expert_params_mapping(

512 self,

513 ckpt_gate_proj_name="gate_proj",

514 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/qwen3_next.py

+3 -3

@@ -23,7 +23,7 @@ from vllm.distributed import (

23 )

24 from vllm.logger import init_logger

25 from vllm.model_executor.layers.attention import Attention

26 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

26 +from vllm.model_executor.layers.fused_moe import FusedMoE

27 from vllm.model_executor.layers.layernorm import (

28 GemmaRMSNorm as Qwen3NextRMSNorm,

29 )

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

146 else:

147 self.shared_expert = None

148

149 - self.experts = SharedFusedMoE(

149 + self.experts = FusedMoE(

150 shared_experts=self.shared_expert,

151 gate=self.gate,

152 num_experts=self.n_routed_experts,

@@ -533,7 +533,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

533 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

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

535 # (param_name, weight_name, expert_id, shard_id)

536 - return SharedFusedMoE.make_expert_params_mapping(

536 + return FusedMoE.make_expert_params_mapping(

537 self,

538 ckpt_gate_proj_name="gate_proj",

539 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/sarvam.py

+4 -4

@@ -35,7 +35,7 @@ from vllm.distributed import (

35 get_tensor_model_parallel_world_size,

36 )

37 from vllm.model_executor.layers.activation import SiluAndMul

38 -from vllm.model_executor.layers.fused_moe import SharedFusedMoE

38 +from vllm.model_executor.layers.fused_moe import FusedMoE

39 from vllm.model_executor.layers.layernorm import RMSNorm

40 from vllm.model_executor.layers.linear import (

41 ColumnParallelLinear,

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

335 else:

336 self.shared_experts = None

337

338 - self.experts = SharedFusedMoE(

338 + self.experts = FusedMoE(

339 shared_experts=self.shared_experts,

340 num_experts=self.num_experts,

341 top_k=self.top_k,

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

352 routed_scaling_factor=self.routed_scaling_factor,

353 )

354

355 - def maybe_get_fused_moe(self) -> SharedFusedMoE:

355 + def maybe_get_fused_moe(self) -> FusedMoE:

356 return self.experts

357

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

@@ -529,7 +529,7 @@ class SarvamMLAModel(nn.Module):

529 return hidden_states

530

531 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

532 - return SharedFusedMoE.make_expert_params_mapping(

532 + return FusedMoE.make_expert_params_mapping(

533 self,

534 ckpt_gate_proj_name="gate_proj",

535 ckpt_down_proj_name="down_proj",

vllm/model_executor/models/step3p5.py

+1 -2

@@ -24,7 +24,6 @@ from vllm.logger import init_logger

24 from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul

25 from vllm.model_executor.layers.attention import Attention

26 from vllm.model_executor.layers.fused_moe import FusedMoE

27 -from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

27 from vllm.model_executor.layers.layernorm import GemmaRMSNorm

28 from vllm.model_executor.layers.linear import (

29 ColumnParallelLinear,

@@ -372,7 +371,7 @@ class FusedMoEBlock(nn.Module):

371 quant_config=quant_config,

372 prefix=f"{prefix}.share_expert",

373 )

375 - self.experts = SharedFusedMoE(

374 + self.experts = FusedMoE(

375 shared_experts=self.share_expert,

376 gate=self.gate,

377 num_experts=config.moe_num_experts,

vllm/utils/__init__.py

+13

@@ -34,3 +34,16 @@ def length_from_prompt_token_ids_or_embeds(

34 f" prompt_embeds={prompt_embeds_len}"

35 )

36 return prompt_token_len

37 +

38 +

39 +def is_moe_layer(module: torch.nn.Module) -> bool:

40 + # TODO(bnell): Should use isinstance but can't due to circular dependencies.

41 + def _check_bases(cls):

42 + if cls.__name__ == "FusedMoE":

43 + return True

44 +

45 + for b in cls.__bases__:

46 + if _check_bases(b):

47 + return True

48 +

49 + return _check_bases(module.__class__)

[MoE Refactor] Remove SharedFusedMoE class (#35782)