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

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 8, 12:30 UTC

[MoE Refactor] Introduce MoERunner abstraction and move execution logic from FusedMoE to DefaultMoERunner (#32344)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Feb 10, 2026 at 19:51 UTC d1481ba78323bcba5937f5ff74f3a8d27ab54f88

25 files changed +913 -753

docs/design/moe_kernel_features.md

+1 -1

@@ -32,7 +32,7 @@ th {

34 |---------|--------------------|--------------|---------------|-------|-----------------------|-----------|

35 -| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |

35 +| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE |

36 | pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |

37 | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |

38 | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |

tests/kernels/moe/modular_kernel_tools/common.py

@@ -585,6 +585,7 @@ def make_modular_kernel(

585 tp_size_=get_tensor_model_parallel_world_size(),

586 pcp_size_=get_pcp_group().world_size,

587 dp_size_=get_dp_group().world_size,

588 + sp_size_=1,

589 vllm_parallel_config=vllm_config.parallel_config,

590 )

591

@@ -594,6 +595,7 @@ def make_modular_kernel(

595 hidden_dim=config.K,

596 intermediate_size_per_partition=config.N,

597 num_local_experts=config.num_local_experts,

598 + num_logical_experts=config.E,

599 moe_parallel_config=moe_parallel_config,

600 in_dtype=config.dtype,

601 max_num_tokens=next_power_of_2(config.M),

tests/kernels/moe/utils.py

@@ -52,6 +52,7 @@ def make_dummy_moe_config(

52 hidden_dim=hidden_dim,

53 intermediate_size_per_partition=intermediate_size_per_partition,

54 num_local_experts=num_experts,

55 + num_logical_experts=num_experts,

56 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

57 activation="silu",

58 in_dtype=in_dtype,

vllm/model_executor/layers/fused_moe/config.py

+18 -4

@@ -913,12 +913,16 @@ class FusedMoEParallelConfig:

913 pcp_rank: int

914 dp_rank: int

915 ep_rank: int

916 + sp_size: int

917

918 use_ep: bool # whether to use EP or not

919 all2all_backend: str # all2all backend for MoE communication

919 - is_sequence_parallel: bool # whether sequence parallelism is used

920 enable_eplb: bool # whether to enable expert load balancing

921

922 + @property

923 + def is_sequence_parallel(self) -> bool:

924 + return self.sp_size > 1

925 +

926 @property

927 def use_all2all_kernels(self):

928 return self.dp_size > 1 and self.use_ep

@@ -974,6 +978,7 @@ class FusedMoEParallelConfig:

978 tp_size_: int,

979 pcp_size_: int,

980 dp_size_: int,

981 + sp_size_: int,

982 vllm_parallel_config: ParallelConfig,

983 ) -> "FusedMoEParallelConfig":

984 """

@@ -1073,9 +1078,9 @@ class FusedMoEParallelConfig:

1078 dp_rank=dp_rank,

1079 ep_size=1,

1080 ep_rank=0,

1081 + sp_size=sp_size_,

1082 use_ep=False,

1083 all2all_backend=vllm_parallel_config.all2all_backend,

1078 - is_sequence_parallel=vllm_parallel_config.use_sequence_parallel_moe,

1084 enable_eplb=vllm_parallel_config.enable_eplb,

1085 )

1086 # DP + EP / TP + EP / DP + TP + EP

@@ -1093,9 +1098,9 @@ class FusedMoEParallelConfig:

1098 dp_rank=dp_rank,

1099 ep_size=ep_size,

1100 ep_rank=ep_rank,

1101 + sp_size=sp_size_,

1102 use_ep=True,

1103 all2all_backend=vllm_parallel_config.all2all_backend,

1098 - is_sequence_parallel=vllm_parallel_config.use_sequence_parallel_moe,

1104 enable_eplb=vllm_parallel_config.enable_eplb,

1105 )

1106

@@ -1111,10 +1116,10 @@ class FusedMoEParallelConfig:

1116 dp_rank=0,

1117 ep_size=1,

1118 ep_rank=0,

1119 + sp_size=1,

1120 use_ep=False,

1121 all2all_backend="naive",

1122 enable_eplb=False,

1117 - is_sequence_parallel=False,

1123 )

1124

1125

@@ -1126,6 +1131,7 @@ class FusedMoEConfig:

1131 hidden_dim: int

1132 intermediate_size_per_partition: int

1133 num_local_experts: int

1134 + num_logical_experts: int

1135 activation: str

1136 device: torch.device | str

1137 routing_method: RoutingMethodType

@@ -1175,6 +1181,14 @@ class FusedMoEConfig:

1181 def ep_size(self):

1182 return self.moe_parallel_config.ep_size

1183

1184 + @property

1185 + def sp_size(self):

1186 + return self.moe_parallel_config.sp_size

1187 +

1188 + @property

1189 + def is_sequence_parallel(self):

1190 + return self.moe_parallel_config.is_sequence_parallel

1191 +

1192 @property

1193 def tp_rank(self):

1194 return self.moe_parallel_config.tp_rank

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+1 -2

@@ -121,17 +121,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):

121 def is_monolithic(self) -> bool:

122 return False

123

124 - # @abstractmethod

124 def apply(

125 self,

126 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

127 x: torch.Tensor,

128 topk_weights: torch.Tensor,

129 topk_ids: torch.Tensor,

130 + shared_experts_input: torch.Tensor | None,

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

132 raise NotImplementedError

133

134 - # @abstractmethod

134 def apply_monolithic(

135 self,

136 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

+2 -1

@@ -89,6 +89,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

89 x: torch.Tensor,

90 topk_weights: torch.Tensor,

91 topk_ids: torch.Tensor,

92 + shared_experts_input: torch.Tensor | None,

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

94 assert self.moe_mk is not None

95 return self.moe_mk(

@@ -101,5 +102,5 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

102 global_num_experts=layer.global_num_experts,

103 apply_router_weight_on_input=layer.apply_router_weight_on_input,

104 expert_map=None if self.disable_expert_map else layer.expert_map,

104 - shared_experts_input=layer._get_shared_experts_input(x),

105 + shared_experts_input=shared_experts_input,

106 )

vllm/model_executor/layers/fused_moe/layer.py

+69 -672

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 -from collections.abc import Callable, Generator, Iterable

5 -from contextlib import contextmanager, nullcontext

4 +from collections.abc import Callable, Iterable

5 from enum import Enum

6 from typing import Literal, cast, get_args, overload

8 import torch

10 -import torch.nn.functional as F

9 from torch.nn.parameter import UninitializedParameter

11 import vllm.envs as envs

@@ -16,17 +14,10 @@ from vllm.config import VllmConfig, get_current_vllm_config

14 from vllm.config.parallel import ExpertPlacementStrategy

15 from vllm.distributed import (

16 get_dp_group,

19 - get_ep_group,

17 get_pcp_group,

18 get_tensor_model_parallel_world_size,

22 - tensor_model_parallel_all_reduce,

19 )

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

25 -from vllm.forward_context import (

26 - ForwardContext,

27 - get_forward_context,

28 - is_forward_context_available,

29 -)

21 from vllm.logger import init_logger

22 from vllm.model_executor.custom_op import CustomOp

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

@@ -47,6 +38,9 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

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

39 create_fused_moe_router,

40 )

41 +from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (

42 + DefaultMoERunner,

43 +)

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

45 UnquantizedFusedMoEMethod,

46 )

@@ -57,13 +51,7 @@ from vllm.model_executor.layers.quantization.base_config import (

51 QuantizationConfig,

52 )

53 from vllm.platforms import current_platform

60 -from vllm.utils.math_utils import cdiv, round_up

61 -from vllm.utils.torch_utils import (

62 - aux_stream,

63 - current_stream,

64 - direct_register_custom_op,

65 -)

66 -from vllm.v1.worker.ubatching import dbo_current_ubatch_id

54 +from vllm.utils.math_utils import round_up

56 logger = init_logger(__name__)

@@ -264,6 +252,7 @@ def maybe_roundup_hidden_size(

252 )

253

254 current_mxfp4_backend = get_mxfp4_backend(is_lora_enabled)

255 +

256 if (

257 current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16

258 or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS

@@ -273,6 +262,7 @@ def maybe_roundup_hidden_size(

262 current_platform.is_rocm()

263 or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM

264 or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16

265 + or current_mxfp4_backend == Mxfp4Backend.MARLIN

266 ):

267 hidden_size = round_up(hidden_size, 256)

268

@@ -338,29 +328,15 @@ class FusedMoE(CustomOp):

328 expert_mapping: list[tuple[str, str, int, str]] | None = None,

329 n_shared_experts: int | None = None,

330 router_logits_dtype: torch.dtype | None = None,

341 - has_shared_experts: bool = False,

331 + gate: torch.nn.Module | None = None,

332 + shared_experts: torch.nn.Module | None = None,

333 + routed_input_transform: torch.nn.Module | None = None,

334 ):

335 super().__init__()

336

345 - # Allow disabling of the separate shared experts stream for

346 - # debug purposes.

347 - # TODO: Remove this after more extensive testings with TP/DP

348 - # and other execution modes

349 - if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:

350 - logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")

351 - self.shared_experts_stream = None

352 - else:

353 - # TODO(rob): enable shared expert overlap with non-cuda-alike.

354 - # aux_stream() returns None on non-cuda-alike platforms.

355 - self.shared_experts_stream = aux_stream()

356 - if self.shared_experts_stream is not None:

357 - logger.debug_once(

358 - "Enabled separate cuda stream for MoE shared_experts", scope="local"

359 - )

360 -

361 - # For latent MoE: stores original hidden_states before routed_input_transform

362 - # so shared_experts can use it for cloning (they need original dimension)

363 - self._shared_experts_input: torch.Tensor | None = None

337 + self._gate = gate

338 + self._shared_experts = shared_experts

339 + self._routed_input_transform = routed_input_transform

340

341 if params_dtype is None:

342 params_dtype = torch.get_default_dtype()

@@ -392,9 +368,12 @@ class FusedMoE(CustomOp):

368 tp_size_=tp_size_,

369 pcp_size_=pcp_size_,

370 dp_size_=dp_size_,

371 + sp_size_=self.sp_size,

372 vllm_parallel_config=vllm_config.parallel_config,

373 )

374

375 + assert self.moe_parallel_config.is_sequence_parallel == is_sequence_parallel

376 +

377 self.global_num_experts = num_experts + num_redundant_experts

378 self.logical_num_experts = num_experts

379

@@ -410,6 +389,7 @@ class FusedMoE(CustomOp):

389 self.layer_name = prefix

390

391 self.enable_eplb = enable_eplb

392 + # TODO(bnell): should this be owned by router?

393 self.eplb_state = EplbLayerState()

394 self.expert_placement_strategy: ExpertPlacementStrategy = (

395 vllm_config.parallel_config.expert_placement_strategy

@@ -506,7 +486,8 @@ class FusedMoE(CustomOp):

486 self.reduce_results = reduce_results

487 self.renormalize = renormalize

488

509 - # TODO(bnell): these attributes are only used by cpu/xpu/mxfp4

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

490 + # Put them in a MoERouterConfig dataclass?

491 self.use_grouped_topk = use_grouped_topk

492 if self.use_grouped_topk:

493 assert num_expert_group is not None and topk_group is not None

@@ -565,6 +546,7 @@ class FusedMoE(CustomOp):

546 hidden_dim=hidden_size,

547 intermediate_size_per_partition=self.intermediate_size_per_partition,

548 num_local_experts=self.local_num_experts,

549 + num_logical_experts=self.logical_num_experts,

550 moe_parallel_config=self.moe_parallel_config,

551 in_dtype=moe_in_dtype,

552 router_logits_dtype=router_logits_dtype,

@@ -576,9 +558,9 @@ class FusedMoE(CustomOp):

558 device=vllm_config.device_config.device,

559 routing_method=self.routing_method_type,

560 # TODO: in_dtype == out_dtype?

579 - disable_inplace=disable_inplace() or has_shared_experts,

561 + disable_inplace=disable_inplace() or self._shared_experts is not None,

562 )

581 - if self.use_mori_kernels:

563 + if self.moe_config.use_mori_kernels:

564 assert self.rocm_aiter_fmoe_enabled, (

565 "Mori needs to be used with aiter fused_moe for now."

566 )

@@ -641,9 +623,36 @@ class FusedMoE(CustomOp):

623

624 self.quant_method.create_weights(layer=self, **moe_quant_params)

625

644 - # Chunked all2all staging tensor

645 - self.batched_hidden_states: torch.Tensor | None = None

646 - self.batched_router_logits: torch.Tensor | None = None

626 + # Disable shared expert overlap if:

627 + # - we are using eplb with non-default backend, because of correctness issues

628 + # - we are using flashinfer with DP, since there nothing to gain

629 + # - we are using marlin kernels

630 + backend = self.moe_parallel_config.all2all_backend

631 + self.use_overlapped = (

632 + not (

633 + (self.enable_eplb and backend != "allgather_reducescatter")

634 + or self.moe_parallel_config.use_fi_all2allv_kernels

635 + )

636 + and self._shared_experts is not None

637 + )

638 +

639 + self.runner = self._init_runner()

640 +

641 + def _init_runner(self):

642 + # Storing the runner in the FusedMoE is an intermediate state, eventually

643 + # the runner will own the FusedMoE layer and provide the execution interface

644 + # for MoE ops.

645 + return DefaultMoERunner(

646 + layer=self,

647 + moe_config=self.moe_config,

648 + router=self.router,

649 + routed_input_transform=self._routed_input_transform,

650 + gate=self.gate,

651 + shared_experts=self.shared_experts,

652 + quant_method=self.quant_method,

653 + reduce_results=self.reduce_results,

654 + enable_dbo=self.vllm_config.parallel_config.enable_dbo,

655 + )

656

657 # Note: maybe_init_modular_kernel should only be called by

658 # prepare_communication_buffer_for_model.

@@ -673,10 +682,14 @@ class FusedMoE(CustomOp):

682 self.shared_experts,

683 inplace=not self.moe_config.disable_inplace,

684 )

685 + # We need to force reconstruction of runner because we're swapping out

686 + # the quant_method with a FusedMoEModularMethod. This logic can go

687 + # away once the FusedMoEModularMethod is eliminated.

688 + self.runner = self._init_runner()

689

690 @property

691 def shared_experts(self) -> torch.nn.Module | None:

679 - return None

692 + return self._shared_experts if self.use_overlapped else None

693

694 @property

695 def layer_id(self):

@@ -687,53 +700,12 @@ class FusedMoE(CustomOp):

700

701 @property

702 def gate(self) -> torch.nn.Module | None:

690 - return None

691 -

692 - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:

693 - """Hook to transform hidden_states before passing to routed experts.

694 - For latent MoE: transforms [S, hidden_size] → [S, moe_latent_size].

695 - The original hidden_states is saved in _shared_experts_input so

696 - shared_experts still receive the original [S, hidden_size].

697 -

698 - Override in subclasses (e.g., SharedFusedMoE) for latent MoE.

699 - """

700 - return hidden_states

701 -

702 - @contextmanager

703 - def _set_shared_experts_input(

704 - self, value: torch.Tensor | None

705 - ) -> Generator[None, None, None]:

706 - """Context manager to safely set/clear _shared_experts_input."""

707 - self._shared_experts_input = value

708 - try:

709 - yield

710 - finally:

711 - self._shared_experts_input = None

712 -

713 - def _get_shared_experts_input(self, hidden_states: torch.Tensor) -> torch.Tensor:

714 - """Get input for shared experts.

715 -

716 - For latent MoE: shared_experts need original [S, hidden_size],

717 - not the transformed [S, latent_size] used by routed experts.

718 - """

719 - return (

720 - self._shared_experts_input

721 - if self._shared_experts_input is not None

722 - else hidden_states

723 - )

703 + return self._gate

704

705 @property

706 def tp_size(self):

707 return self.moe_parallel_config.tp_size

708

729 - @property

730 - def dp_size(self):

731 - return self.moe_parallel_config.dp_size

732 -

733 - @property

734 - def pcp_size(self):

735 - return self.moe_parallel_config.pcp_size

736 -

709 @property

710 def ep_size(self):

711 return self.moe_parallel_config.ep_size

@@ -742,14 +714,6 @@ class FusedMoE(CustomOp):

714 def tp_rank(self):

715 return self.moe_parallel_config.tp_rank

716

745 - @property

746 - def dp_rank(self):

747 - return self.moe_parallel_config.dp_rank

748 -

749 - @property

750 - def pcp_rank(self):

751 - return self.moe_parallel_config.pcp_rank

752 -

717 @property

718 def ep_rank(self):

719 return self.moe_parallel_config.ep_rank

@@ -758,39 +722,10 @@ class FusedMoE(CustomOp):

722 def use_ep(self):

723 return self.moe_parallel_config.use_ep

724

761 - @property

762 - def use_pplx_kernels(self):

763 - return self.moe_parallel_config.use_pplx_kernels

764 -

765 - @property

766 - def use_deepep_ht_kernels(self):

767 - return self.moe_parallel_config.use_deepep_ht_kernels

768 -

769 - @property

770 - def use_deepep_ll_kernels(self):

771 - return self.moe_parallel_config.use_deepep_ll_kernels

772 -

773 - @property

774 - def use_mori_kernels(self):

775 - return self.moe_parallel_config.use_mori_kernels

776 -

777 - @property

778 - def use_marlin_kernels(self):

779 - return getattr(self.quant_method, "use_marlin", False)

780 -

781 - @property

782 - def use_dp_chunking(self) -> bool:

783 - return (

784 - self.moe_parallel_config.use_pplx_kernels

785 - or self.moe_parallel_config.use_deepep_ll_kernels

786 - or self.moe_parallel_config.use_mori_kernels

787 - or self.moe_parallel_config.use_fi_all2allv_kernels

788 - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK

789 -

725 @property

726 def is_internal_router(self) -> bool:

727 # By default, router/gate is called before FusedMoE forward pass

793 - return False

728 + return self._gate is not None

729

730 def _maybe_init_expert_routing_tables(

731 self,

@@ -799,7 +734,7 @@ class FusedMoE(CustomOp):

734 # with DeepEP-ll all2all backend.

735 if (

736 self.expert_placement_strategy != "round_robin"

802 - or not self.use_deepep_ll_kernels

737 + or not self.moe_parallel_config.use_deepep_ll_kernels

738 ):

739 return None

740

@@ -892,48 +827,6 @@ class FusedMoE(CustomOp):

827 dp_size=get_dp_group().world_size,

828 )

829

895 - def _maybe_setup_shared_experts_stream(

896 - self,

897 - hidden_states: torch.Tensor,

898 - has_separate_shared_experts: bool,

899 - use_chunked_impl: bool,

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

901 - use_shared_experts_stream = (

902 - current_platform.is_cuda()

903 - and has_separate_shared_experts

904 - and not use_chunked_impl

905 - and self.shared_experts_stream is not None

906 - and (

907 - hidden_states.shape[0]

908 - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

909 - )

910 - )

911 -

912 - hidden_states_clone: torch.Tensor | None = None

913 - if use_shared_experts_stream:

914 - assert self.shared_experts_stream is not None

915 -

916 - shared_experts_input = self._get_shared_experts_input(hidden_states)

917 -

918 - # Clone BEFORE switching streams to avoid race condition

919 - # where routed_expert kernel may mutate hidden_states.

920 - hidden_states_clone = shared_experts_input.clone()

921 -

922 - # Record that the clone will be used by shared_experts_stream

923 - # to avoid gc issue from deallocation of hidden_states_clone

924 - # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501

925 - # NOTE: We don't need shared_output.record_stream(current_stream())

926 - # because we synch the streams before using shared_output.

927 - hidden_states_clone.record_stream(self.shared_experts_stream)

928 -

929 - # Mark sync start point for the separate shared experts

930 - # stream here since we want to run in parallel with the

931 - # router/gate (next op below)

932 - assert self.shared_experts_stream is not None

933 - self.shared_experts_stream.wait_stream(current_stream())

934 -

935 - return use_shared_experts_stream, hidden_states_clone

936 -

830 def _load_per_tensor_weight_scale(

831 self,

832 shard_id: str,

@@ -1191,7 +1084,7 @@ class FusedMoE(CustomOp):

1084 # compressed-tensors checkpoints with packed weights are stored flipped

1085 # TODO (mgoin): check self.quant_method.quant_config.quant_format

1086 # against known CompressionFormat enum values that have this quality

1194 - if self.quant_method.__class__.__name__ in (

1087 + if quant_method_name in (

1088 "CompressedTensorsWNA16MarlinMoEMethod",

1089 "CompressedTensorsWNA16MoEMethod",

1090 ):

@@ -1488,7 +1381,7 @@ class FusedMoE(CustomOp):

1381 assert all(

1382 weight.is_contiguous()

1383 for name, weight in weights

1491 - if not name.startswith("_shared_experts.")

1384 + if not (name.startswith("_shared_experts.") or name.startswith("_gate."))

1385 )

1386

1387 # Filter out the non-expert weights.

@@ -1538,32 +1431,6 @@ class FusedMoE(CustomOp):

1431 self.ensure_moe_quant_config_init()

1432 return self.quant_method.moe_quant_config

1433

1541 - def ensure_dp_chunking_init(self):

1542 - if not self.use_dp_chunking or self.batched_hidden_states is not None:

1543 - return

1544 -

1545 - states_shape: tuple[int, ...]

1546 - logits_shape: tuple[int, ...]

1547 -

1548 - moe = self.moe_config

1549 -

1550 - if self.vllm_config.parallel_config.enable_dbo:

1551 - states_shape = (2, moe.max_num_tokens, self.hidden_size)

1552 - logits_shape = (2, moe.max_num_tokens, self.logical_num_experts)

1553 - else:

1554 - states_shape = (moe.max_num_tokens, self.hidden_size)

1555 - logits_shape = (moe.max_num_tokens, self.logical_num_experts)

1556 -

1557 - self.batched_hidden_states = torch.zeros(

1558 - states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()

1559 - )

1560 -

1561 - self.batched_router_logits = torch.zeros(

1562 - logits_shape,

1563 - dtype=moe.router_logits_dtype,

1564 - device=torch.cuda.current_device(),

1565 - )

1566 -

1434 def must_reduce_shared_expert_outputs(self) -> bool:

1435 """

1436 The shared_experts are typically computed using the RowParallelLinear

@@ -1577,100 +1444,24 @@ class FusedMoE(CustomOp):

1444 Therefore it is required that we reduce the shared_experts output

1445 early.

1446 """

1580 - assert self.quant_method is not None

1581 - return (

1582 - isinstance(self.quant_method, FusedMoEModularMethod)

1583 - and self.quant_method.moe_mk.output_is_reduced() # type: ignore[union-attr]

1584 - )

1447 + return self.runner.must_reduce_shared_expert_outputs()

1448

1449 def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):

1450 """

1451 Some combine kernels reduce across GPU ranks by default.

1452 """

1590 - if self.must_reduce_shared_expert_outputs():

1591 - return final_hidden_states

1592 - else:

1593 - return tensor_model_parallel_all_reduce(final_hidden_states)

1453 + return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states)

1454

1455 def forward_native(

1456 self,

1457 hidden_states: torch.Tensor,

1458 router_logits: torch.Tensor,

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

1600 - # For latent MoE: save ORIGINAL hidden_states before transform

1601 - # (shared_experts need original dimension, routed experts use transformed)

1602 - original_hidden_states = hidden_states

1603 - original_hidden_dim = hidden_states.shape[-1]

1604 -

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

1606 - hidden_states = self.apply_routed_input_transform(hidden_states)

1607 -

1608 - # This is the dimension after transform (for routed expert output slicing)

1609 - transformed_hidden_dim = hidden_states.shape[-1]

1610 - if self.hidden_size != transformed_hidden_dim:

1611 - hidden_states = F.pad(

1612 - hidden_states,

1613 - (0, self.hidden_size - transformed_hidden_dim),

1614 - mode="constant",

1615 - value=0.0,

1616 - )

1617 -

1618 - def reduce_output(states: torch.Tensor) -> torch.Tensor:

1619 - if (

1620 - not self.is_sequence_parallel

1621 - and not self.use_dp_chunking

1622 - and self.reduce_results

1623 - and (self.tp_size > 1 or self.ep_size > 1)

1624 - ):

1625 - states = self.maybe_all_reduce_tensor_model_parallel(states)

1626 - return states

1627 -

1628 - def encode_layer_name() -> str:

1629 - # Can be unavailable or None in unittests

1630 - if (

1631 - is_forward_context_available()

1632 - and get_forward_context().all_moe_layers is not None

1633 - ):

1634 - return "from_forward_context"

1635 - return self.layer_name

1636 -

1637 - if self.shared_experts is None:

1638 - if current_platform.is_tpu() or current_platform.is_cpu():

1639 - # TODO: Once the OOM issue for the TPU backend is resolved, we

1640 - # will switch to using the moe_forward custom op.

1641 - # Note: CPU doesn't require wrapped forward_impl.

1642 - fused_output = self.forward_impl(hidden_states, router_logits)

1643 - assert not isinstance(fused_output, tuple)

1644 - else:

1645 - fused_output = torch.ops.vllm.moe_forward(

1646 - hidden_states, router_logits, encode_layer_name()

1647 - )

1648 - return reduce_output(fused_output)[..., :transformed_hidden_dim]

1649 - else:

1650 - if current_platform.is_tpu() or current_platform.is_cpu():

1651 - # TODO: Once the OOM issue for the TPU backend is resolved, we

1652 - # will switch to using the moe_forward custom op.

1653 - # Note: CPU doesn't require wrapped forward_impl.

1654 - with self._set_shared_experts_input(original_hidden_states):

1655 - shared_output, fused_output = self.forward_impl(

1656 - hidden_states, router_logits

1657 - )

1658 - else:

1659 - # Custom op handles setting/clearing _shared_experts_input internally

1660 - # We pass original tensor for shared experts (not transformed)

1661 - shared_output, fused_output = torch.ops.vllm.moe_forward_shared(

1662 - hidden_states,

1663 - router_logits,

1664 - encode_layer_name(),

1665 - original_hidden_states,

1666 - )

1667 -

1668 - # shared_output uses original dimension (before transform)

1669 - # fused_output uses transformed dimension (after transform)

1670 - return (

1671 - reduce_output(shared_output)[..., :original_hidden_dim],

1672 - reduce_output(fused_output)[..., :transformed_hidden_dim],

1673 - )

1460 + self.ensure_moe_quant_config_init()

1461 + return self.runner.forward(

1462 + hidden_states,

1463 + router_logits,

1464 + )

1465

1466 @property

1467 def expert_map(self) -> torch.Tensor | None:

@@ -1685,312 +1476,6 @@ class FusedMoE(CustomOp):

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

1477 return self.forward_native(hidden_states, router_logits)

1478

1688 - def forward_impl_chunked(

1689 - self,

1690 - full_hidden_states: torch.Tensor,

1691 - full_router_logits: torch.Tensor,

1692 - has_separate_shared_experts: bool,

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

1694 - assert self.batched_hidden_states is not None

1695 - assert self.batched_router_logits is not None

1696 - assert self.batched_hidden_states.dtype == full_hidden_states.dtype, (

1697 - f"{self.batched_hidden_states.dtype} == {full_hidden_states.dtype}"

1698 - )

1699 - assert self.batched_router_logits.dtype == full_router_logits.dtype, (

1700 - f"{self.batched_router_logits.dtype} == {full_router_logits.dtype}"

1701 - )

1702 - # Check size compatibility.

1703 - assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1)

1704 - assert self.batched_router_logits.size(-1) == full_router_logits.size(-1)

1705 -

1706 - full_fused_final_hidden_states = torch.empty_like(full_hidden_states)

1707 - if self.shared_experts is not None:

1708 - full_shared_final_hidden_states = torch.empty_like(full_hidden_states)

1709 -

1710 - def process_chunk(chunk_start, chunk_end, skip_result_store=False):

1711 - chunk_size = chunk_end - chunk_start

1712 - hidden_states = full_hidden_states[chunk_start:chunk_end, :]

1713 - router_logits = full_router_logits[chunk_start:chunk_end, :]

1714 -

1715 - assert self.batched_hidden_states is not None

1716 - assert self.batched_router_logits is not None

1717 - # This is only true when DBO has been enabled in the config.

1718 - # Both tensors will have an outer dimension for the ubatch id

1719 - if self.batched_hidden_states.dim() == 3:

1720 - assert self.batched_router_logits.dim() == 3

1721 - batch_buffer_idx = dbo_current_ubatch_id()

1722 - batched_hidden_states = self.batched_hidden_states[batch_buffer_idx, :]

1723 - batched_router_logits = self.batched_router_logits[batch_buffer_idx, :]

1724 - else:

1725 - batched_hidden_states = self.batched_hidden_states

1726 - batched_router_logits = self.batched_router_logits

1727 -

1728 - assert (

1729 - batched_hidden_states.size(0) # type: ignore

1730 - >= chunk_size

1731 - )

1732 - assert (

1733 - batched_router_logits.size(0) # type: ignore

1734 - >= chunk_size

1735 - )

1736 - staged_hidden_states = batched_hidden_states[:chunk_size, :] # type: ignore

1737 - staged_router_logits = batched_router_logits[:chunk_size, :] # type: ignore

1738 - staged_hidden_states.copy_(hidden_states, non_blocking=True)

1739 - staged_router_logits.copy_(router_logits, non_blocking=True)

1740 -

1741 - # Matrix multiply.

1742 - if self.quant_method.is_monolithic:

1743 - final_hidden_states = self.quant_method.apply_monolithic(

1744 - layer=self,

1745 - x=staged_hidden_states,

1746 - router_logits=staged_router_logits,

1747 - )

1748 - else:

1749 - topk_weights, topk_ids = self.router.select_experts(

1750 - hidden_states=staged_hidden_states,

1751 - router_logits=staged_router_logits,

1752 - )

1753 -

1754 - final_hidden_states = self.quant_method.apply(

1755 - layer=self,

1756 - x=staged_hidden_states,

1757 - topk_weights=topk_weights,

1758 - topk_ids=topk_ids,

1759 - )

1760 -

1761 - if has_separate_shared_experts:

1762 - assert not isinstance(final_hidden_states, tuple)

1763 - assert self.shared_experts is not None

1764 -

1765 - shared_output = self.shared_experts(staged_hidden_states)

1766 -

1767 - final_hidden_states = (

1768 - shared_output,

1769 - final_hidden_states,

1770 - )

1771 -

1772 - if not skip_result_store:

1773 - if self.shared_experts is None:

1774 - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

1775 - final_hidden_states, non_blocking=True

1776 - )

1777 - else:

1778 - full_shared_final_hidden_states[chunk_start:chunk_end, :].copy_(

1779 - final_hidden_states[0], non_blocking=True

1780 - )

1781 - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

1782 - final_hidden_states[1], non_blocking=True

1783 - )

1784 -

1785 - ctx = get_forward_context()

1786 - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP

1787 - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu

1788 - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens

1789 -

1790 - # If the input to the MoE is sequence parallel then divide by sp_size

1791 - # to find the maximum number of tokens for any individual dispatcher.

1792 - if self.is_sequence_parallel:

1793 - max_tokens_across_dispatchers = cdiv(

1794 - max_tokens_across_dispatchers, self.sp_size

1795 - )

1796 -

1797 - num_tokens = full_hidden_states.size(0)

1798 - for chunk_idx, chunk_start_ in enumerate(

1799 - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)

1800 - ):

1801 - chunk_start = chunk_start_

1802 - chunk_end = min(

1803 - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers

1804 - )

1805 - # clamp start and end

1806 - chunk_start = min(chunk_start, num_tokens - 1)

1807 - chunk_end = min(chunk_end, num_tokens)

1808 - with ctx.dp_metadata.chunked_sizes(

1809 - self.sp_size, moe_dp_chunk_size_per_rank, chunk_idx

1810 - ):

1811 - process_chunk(

1812 - chunk_start, chunk_end, skip_result_store=chunk_start_ >= num_tokens

1813 - )

1814 -

1815 - if self.shared_experts is None:

1816 - return full_fused_final_hidden_states

1817 - else:

1818 - return (full_shared_final_hidden_states, full_fused_final_hidden_states)

1819 -

1820 - def forward_impl(

1821 - self,

1822 - hidden_states: torch.Tensor,

1823 - router_logits: torch.Tensor,

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

1825 - assert self.quant_method is not None

1826 -

1827 - self.ensure_moe_quant_config_init()

1828 - self.ensure_dp_chunking_init()

1829 -

1830 - has_separate_shared_experts = (

1831 - not self.quant_method.mk_owns_shared_expert

1832 - and self.shared_experts is not None

1833 - )

1834 -

1835 - use_chunked_impl = self.use_dp_chunking

1836 -

1837 - use_shared_experts_stream, hidden_states_clone = (

1838 - self._maybe_setup_shared_experts_stream(

1839 - hidden_states, has_separate_shared_experts, use_chunked_impl

1840 - )

1841 - )

1842 -

1843 - # If router/gate provided, then apply it here.

1844 - # (Note: This code runs only when "overlapped mode" is on to allow

1845 - # parallel execution of shared experts with the FusedMoE via

1846 - # separate cuda stream)

1847 - if self.gate is not None:

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

1849 -

1850 - if use_chunked_impl:

1851 - return self.forward_impl_chunked(

1852 - hidden_states, router_logits, has_separate_shared_experts

1853 - )

1854 -

1855 - # NOTE(rob): once we finish migrating all the quant methods to use

1856 - # MKs, we can remove the naive dispatch/combine path from here.

1857 - do_naive_dispatch_combine = (

1858 - self.dp_size > 1 and not self.quant_method.supports_internal_mk

1859 - )

1860 -

1861 - ctx = get_forward_context()

1862 - sp_ctx = (

1863 - ctx.dp_metadata.sp_local_sizes(self.sp_size)

1864 - if ctx.dp_metadata

1865 - else nullcontext()

1866 - )

1867 -

1868 - with sp_ctx:

1869 - extra_tensors = None

1870 - if do_naive_dispatch_combine:

1871 - post_quant_allgather = (

1872 - self.quant_method is not None

1873 - and self.dp_size > 1

1874 - and self.use_ep

1875 - and getattr(self.quant_method, "do_post_quant_allgather", False)

1876 - )

1877 - if post_quant_allgather:

1878 - hidden_states_to_dispatch, extra_tensors = (

1879 - self.quant_method.prepare_dp_allgather_tensor(

1880 - self, hidden_states, router_logits

1881 - )

1882 - )

1883 - else:

1884 - hidden_states_to_dispatch = hidden_states

1885 -

1886 - dispatch_res = get_ep_group().dispatch_router_logits(

1887 - hidden_states_to_dispatch,

1888 - router_logits,

1889 - self.is_sequence_parallel,

1890 - extra_tensors=extra_tensors,

1891 - )

1892 - if extra_tensors is not None:

1893 - (

1894 - orig_hidden_states,

1895 - router_logits,

1896 - extra_tensors_combined,

1897 - ) = dispatch_res

1898 - hidden_states_combined = (

1899 - orig_hidden_states,

1900 - extra_tensors_combined[0],

1901 - )

1902 - else:

1903 - hidden_states_combined, router_logits = dispatch_res

1904 - orig_hidden_states = hidden_states_combined

1905 - else:

1906 - orig_hidden_states = hidden_states

1907 -

1908 - # Run shared experts before matrix multiply.

1909 - # because matrix multiply maybe modify the hidden_states.

1910 - if has_separate_shared_experts and not use_shared_experts_stream:

1911 - assert self.shared_experts is not None

1912 - shared_input = self._get_shared_experts_input(hidden_states)

1913 - shared_output = self.shared_experts(shared_input)

1914 -

1915 - # NOTE: Similar with DP, PCP also needs dispatch and combine. For

1916 - # simplicity, AgRsAll2All was added separately for PCP here. Maybe

1917 - # we should modify All2AllManager abstract to better support PCP.

1918 - if self.pcp_size > 1:

1919 - hidden_states = get_pcp_group().all_gather(

1920 - hidden_states,

1921 - dim=0,

1922 - )

1923 - router_logits = get_pcp_group().all_gather(

1924 - router_logits,

1925 - dim=0,

1926 - )

1927 -

1928 - # Matrix multiply.

1929 - x = hidden_states_combined if do_naive_dispatch_combine else hidden_states

1930 -

1931 - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014).

1932 - # Figure out nicer way to do this.

1933 - x_orig = orig_hidden_states if do_naive_dispatch_combine else hidden_states

1934 -

1935 - if self.quant_method.is_monolithic:

1936 - final_hidden_states = self.quant_method.apply_monolithic(

1937 - layer=self,

1938 - x=x,

1939 - router_logits=router_logits,

1940 - )

1941 - else:

1942 - topk_weights, topk_ids = self.router.select_experts(

1943 - hidden_states=x_orig,

1944 - router_logits=router_logits,

1945 - )

1946 -

1947 - final_hidden_states = self.quant_method.apply(

1948 - layer=self,

1949 - x=x, # The type signture of this is wrong due to the hack.

1950 - topk_weights=topk_weights,

1951 - topk_ids=topk_ids,

1952 - )

1953 -

1954 - if has_separate_shared_experts:

1955 - assert self.shared_experts is not None

1956 -

1957 - if use_shared_experts_stream:

1958 - # Run shared experts in parallel on a separate stream

1959 - # NOTE: We start the separate stream here and mark the

1960 - # sync end point immediately after it is done. This is

1961 - # important to avoid excessive stream allocations by the cuda

1962 - # graph replay later.

1963 - with torch.cuda.stream(self.shared_experts_stream):

1964 - # Note that hidden_states clone() is necessary here to avoid

1965 - # conflict with the main stream

1966 - shared_output = self.shared_experts(hidden_states_clone)

1967 - current_stream().wait_stream(self.shared_experts_stream)

1968 -

1969 - final_hidden_states = (

1970 - shared_output,

1971 - final_hidden_states,

1972 - )

1973 -

1974 - def combine_output(states: torch.Tensor) -> torch.Tensor:

1975 - if do_naive_dispatch_combine:

1976 - states = get_ep_group().combine(states, self.is_sequence_parallel)

1977 -

1978 - if self.pcp_size > 1:

1979 - states = get_pcp_group().reduce_scatter(

1980 - states,

1981 - dim=0,

1982 - )

1983 -

1984 - return states

1985 -

1986 - if self.shared_experts is not None:

1987 - return (

1988 - final_hidden_states[0],

1989 - combine_output(final_hidden_states[1]),

1990 - )

1991 - else:

1992 - return combine_output(final_hidden_states)

1993 -

1479 @classmethod

1480 def make_expert_params_mapping(

1481 cls,

@@ -2051,94 +1536,6 @@ class FusedMoE(CustomOp):

1536 return s

1537

1538

2054 -def get_layer_from_name(layer_name: str) -> FusedMoE:

2055 - forward_context: ForwardContext = get_forward_context()

2056 - if layer_name == "from_forward_context":

2057 - all_moe_layers = forward_context.all_moe_layers

2058 - assert all_moe_layers is not None

2059 - moe_layer_index = forward_context.moe_layer_index

2060 - if moe_layer_index >= len(all_moe_layers):

2061 - raise AssertionError(

2062 - "We expected the number of MOE layers in `all_moe_layers` "

2063 - "to be equal to the number of "

2064 - "{vllm.moe_forward, vllm.moe_forward_shared} calls."

2065 - )

2066 - layer_name = all_moe_layers[moe_layer_index]

2067 - forward_context.moe_layer_index += 1

2068 - self = cast(FusedMoE, forward_context.no_compile_layers[layer_name])

2069 - return self

2070 -

2071 -

2072 -def moe_forward(

2073 - hidden_states: torch.Tensor,

2074 - router_logits: torch.Tensor,

2075 - layer_name: str,

2076 -) -> torch.Tensor:

2077 - self = get_layer_from_name(layer_name)

2078 - assert self.shared_experts is None

2079 - return self.forward_impl(hidden_states, router_logits)

2080 -

2081 -

2082 -def moe_forward_fake(

2083 - hidden_states: torch.Tensor,

2084 - router_logits: torch.Tensor,

2085 - layer_name: str,

2086 -) -> torch.Tensor:

2087 - return torch.empty_like(hidden_states)

2088 -

2089 -

2090 -direct_register_custom_op(

2091 - op_name="moe_forward",

2092 - op_func=moe_forward,

2093 - mutates_args=["hidden_states"],

2094 - fake_impl=moe_forward_fake,

2095 - tags=(torch.Tag.needs_fixed_stride_order,),

2096 -)

2097 -

2098 -

2099 -def moe_forward_shared(

2100 - hidden_states: torch.Tensor,

2101 - router_logits: torch.Tensor,

2102 - layer_name: str,

2103 - shared_experts_input: torch.Tensor | None = None,

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

2105 - self = get_layer_from_name(layer_name)

2106 - assert self.shared_experts is not None

2107 -

2108 - # Set here because torch.compile skips forward_native() setup code

2109 - # and calls this op directly. forward_impl() reads from this var.

2110 - with self._set_shared_experts_input(shared_experts_input):

2111 - return self.forward_impl(hidden_states, router_logits)

2112 -

2113 -

2114 -def moe_forward_shared_fake(

2115 - hidden_states: torch.Tensor,

2116 - router_logits: torch.Tensor,

2117 - layer_name: str,

2118 - shared_experts_input: torch.Tensor | None = None,

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

2120 - # Output shapes:

2121 - # - fused_out: same as hidden_states (routed experts use transformed size)

2122 - # - shared_out: same as shared_experts_input if provided, else same as hidden_states

2123 - # (For latent MoE: shared experts use original hidden_size, not latent size)

2124 - fused_out = torch.empty_like(hidden_states)

2125 -

2126 - if shared_experts_input is not None:

2127 - shared_out = torch.empty_like(shared_experts_input)

2128 - else:

2129 - shared_out = torch.empty_like(hidden_states)

2130 -

2131 - return shared_out, fused_out

2132 -

2133 -

2134 -direct_register_custom_op(

2135 - op_name="moe_forward_shared",

2136 - op_func=moe_forward_shared,

2137 - mutates_args=["hidden_states"],

2138 - fake_impl=moe_forward_shared_fake,

2139 - tags=(torch.Tag.needs_fixed_stride_order,),

2140 -)

2141 -

1539 # Mark the FusedMoE weight_loader as supporting MoE-specific parameters

1540 # to avoid expensive runtime reflection in model loading code

1541 FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]

vllm/model_executor/layers/fused_moe/modular_kernel.py

+1 -1

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

1228 topk_weights: torch.Tensor,

1229 topk_ids: torch.Tensor,

1230 apply_router_weight_on_input: bool,

1231 - shared_experts_input: torch.Tensor | None = None,

1231 + shared_experts_input: torch.Tensor | None,

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

1233 """

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

vllm/model_executor/layers/fused_moe/runner/__init__.py new

@@ -0,0 +1,2 @@

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

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

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

+743

@@ -0,0 +1,743 @@

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

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

3 +from contextlib import nullcontext

4 +

5 +import torch

6 +import torch.nn.functional as F

7 +

8 +import vllm.envs as envs

9 +from vllm.distributed import (

10 + get_ep_group,

11 + get_pcp_group,

12 + tensor_model_parallel_all_reduce,

13 +)

14 +from vllm.forward_context import (

15 + ForwardContext,

16 + get_forward_context,

17 + is_forward_context_available,

18 +)

19 +from vllm.logger import init_logger

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

21 + FusedMoEConfig,

22 +)

23 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

24 + FusedMoEMethodBase,

25 +)

26 +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (

27 + FusedMoERouter,

28 +)

29 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner

30 +from vllm.platforms import current_platform

31 +from vllm.utils.math_utils import cdiv

32 +from vllm.utils.torch_utils import (

33 + aux_stream,

34 + current_stream,

35 + direct_register_custom_op,

36 +)

37 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

38 +

39 +logger = init_logger(__name__)

40 +

41 +

42 +def get_layer_from_name(layer_name: str) -> torch.nn.Module:

43 + forward_context: ForwardContext = get_forward_context()

44 + if layer_name == "from_forward_context":

45 + all_moe_layers = forward_context.all_moe_layers

46 + assert all_moe_layers is not None

47 + moe_layer_index = forward_context.moe_layer_index

48 + if moe_layer_index >= len(all_moe_layers):

49 + raise AssertionError(

50 + "We expected the number of MOE layers in `all_moe_layers` "

51 + "to be equal to the number of "

52 + "{vllm.moe_forward, vllm.moe_forward_shared} calls."

53 + )

54 + layer_name = all_moe_layers[moe_layer_index]

55 + forward_context.moe_layer_index += 1

56 + return forward_context.no_compile_layers[layer_name]

57 +

58 +

59 +def _moe_forward(

60 + hidden_states: torch.Tensor,

61 + router_logits: torch.Tensor,

62 + shared_experts_input: torch.Tensor | None,

63 + layer_name: str,

64 +) -> torch.Tensor:

65 + layer = get_layer_from_name(layer_name)

66 + return layer.runner.forward_impl(

67 + layer, hidden_states, router_logits, shared_experts_input

68 + )

69 +

70 +

71 +def _moe_forward_fake(

72 + hidden_states: torch.Tensor,

73 + router_logits: torch.Tensor,

74 + shared_experts_input: torch.Tensor | None,

75 + layer_name: str,

76 +) -> torch.Tensor:

77 + return torch.empty_like(hidden_states)

78 +

79 +

80 +def _moe_forward_shared(

81 + hidden_states: torch.Tensor,

82 + router_logits: torch.Tensor,

83 + shared_experts_input: torch.Tensor | None,

84 + layer_name: str,

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

86 + layer = get_layer_from_name(layer_name)

87 + return layer.runner.forward_impl(

88 + layer, hidden_states, router_logits, shared_experts_input

89 + )

90 +

91 +

92 +def _moe_forward_shared_fake(

93 + hidden_states: torch.Tensor,

94 + router_logits: torch.Tensor,

95 + shared_experts_input: torch.Tensor | None,

96 + layer_name: str,

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

98 + # Output shapes:

99 + # - fused_out: same as hidden_states (routed experts use transformed size)

100 + # - shared_out: same as shared_experts_input if provided, else same as

101 + # hidden_states

102 + # (For latent MoE: shared experts use original hidden_size, not latent size)

103 + fused_out = torch.empty_like(hidden_states)

104 +

105 + if shared_experts_input is not None:

106 + shared_out = torch.empty_like(shared_experts_input)

107 + else:

108 + shared_out = torch.empty_like(hidden_states)

109 +

110 + return shared_out, fused_out

111 +

112 +

113 +direct_register_custom_op(

114 + op_name="moe_forward",

115 + op_func=_moe_forward,

116 + mutates_args=["hidden_states"],

117 + fake_impl=_moe_forward_fake,

118 + tags=(torch.Tag.needs_fixed_stride_order,),

119 +)

120 +

121 +

122 +direct_register_custom_op(

123 + op_name="moe_forward_shared",

124 + op_func=_moe_forward_shared,

125 + mutates_args=["hidden_states"],

126 + fake_impl=_moe_forward_shared_fake,

127 + tags=(torch.Tag.needs_fixed_stride_order,),

128 +)

129 +

130 +

131 +class DefaultMoERunner(MoERunner):

132 + """

133 + Default implementation of the MoE runner for executing Mixture of Experts layers.

134 +

135 + This class provides a comprehensive implementation for running MoE computations

136 + with support for:

137 + - Expert routing and token dispatching

138 + - Shared experts computation with optional parallel execution using CUDA streams

139 + - Data parallel (DP) chunking for large batch processing

140 + - Tensor model parallel and expert parallel operations

141 + - Various quantization methods and custom operators

142 + - Both monolithic and decomposed expert execution paths

143 +

144 + The runner handles the complete MoE forward pass including routing tokens to

145 + experts, executing expert computations, and combining results. It supports

146 + advanced features like overlapped execution of shared experts and optimized

147 + kernels for different parallel execution modes.

148 +

149 + Eventually, this class will be split up and specialized for different

150 + configurations, e.g. the presense or absence of shared experts, a gate, etc.

151 + """

152 +

153 + def __init__(

154 + self,

155 + layer: torch.nn.Module,

156 + moe_config: FusedMoEConfig,

157 + router: FusedMoERouter,

158 + routed_input_transform: torch.nn.Module | None,

159 + gate: torch.nn.Module | None,

160 + shared_experts: torch.nn.Module | None,

161 + quant_method: FusedMoEMethodBase,

162 + reduce_results: bool,

163 + enable_dbo: bool,

164 + ):

165 + super().__init__()

166 + self.moe_config = moe_config

167 + self.router = router

168 + self.routed_input_transform = routed_input_transform

169 + self.gate = gate

170 + self.shared_experts = shared_experts

171 + self.quant_method = quant_method

172 + self.reduce_results = reduce_results

173 + self.enable_dbo = enable_dbo

174 +

175 + # Allow disabling of the separate shared experts stream for

176 + # debug purposes.

177 + # TODO: Remove this after more extensive testings with TP/DP

178 + # and other execution modes

179 + if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:

180 + logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")

181 + self.shared_experts_stream = None

182 + else:

183 + # TODO(rob): enable shared expert overlap with non-cuda-alike.

184 + # aux_stream() returns None on non-cuda-alike platforms.

185 + self.shared_experts_stream = aux_stream()

186 + if self.shared_experts_stream is not None:

187 + logger.debug_once(

188 + "Enabled separate cuda stream for MoE shared_experts", scope="local"

189 + )

190 +

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

192 + self.layer_name = layer.layer_name

193 +

194 + if current_platform.is_tpu() or current_platform.is_cpu():

195 + # TODO: Once the OOM issue for the TPU backend is resolved, we

196 + # will switch to using the moe_forward custom op.

197 + # Note: CPU doesn't require wrapped forward_impl.

198 + if self.shared_experts is None:

199 + self.moe_forward = _moe_forward

200 + else:

201 + self.moe_forward = _moe_forward_shared

202 + else:

203 + if self.shared_experts is None:

204 + self.moe_forward = torch.ops.vllm.moe_forward

205 + else:

206 + self.moe_forward = torch.ops.vllm.moe_forward_shared

207 +

208 + # Chunked all2all staging tensor

209 + self.batched_hidden_states: torch.Tensor | None = None

210 + self.batched_router_logits: torch.Tensor | None = None

211 +

212 + @property

213 + def use_dp_chunking(self) -> bool:

214 + return (

215 + self.moe_config.moe_parallel_config.use_pplx_kernels

216 + or self.moe_config.moe_parallel_config.use_deepep_ll_kernels

217 + or self.moe_config.moe_parallel_config.use_mori_kernels

218 + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels

219 + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK

220 +

221 + def _maybe_setup_shared_experts_stream(

222 + self,

223 + hidden_states: torch.Tensor,

224 + shared_input: torch.Tensor | None,

225 + has_separate_shared_experts: bool,

226 + use_chunked_impl: bool,

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

228 + use_shared_experts_stream = (

229 + current_platform.is_cuda()

230 + and has_separate_shared_experts

231 + and not use_chunked_impl

232 + and self.shared_experts_stream is not None

233 + and (

234 + hidden_states.shape[0]

235 + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

236 + )

237 + )

238 +

239 + hidden_states_clone: torch.Tensor | None = None

240 + if use_shared_experts_stream:

241 + assert self.shared_experts_stream is not None

242 +

243 + shared_experts_input = (

244 + shared_input if shared_input is not None else hidden_states

245 + )

246 +

247 + # Clone BEFORE switching streams to avoid race condition

248 + # where routed_expert kernel may mutate hidden_states.

249 + hidden_states_clone = shared_experts_input.clone()

250 +

251 + # Record that the clone will be used by shared_experts_stream

252 + # to avoid gc issue from deallocation of hidden_states_clone

253 + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501

254 + # NOTE: We don't need shared_output.record_stream(current_stream())

255 + # because we synch the streams before using shared_output.

256 + hidden_states_clone.record_stream(self.shared_experts_stream)

257 +

258 + # Mark sync start point for the separate shared experts

259 + # stream here since we want to run in parallel with the

260 + # router/gate (next op below)

261 + assert self.shared_experts_stream is not None

262 + self.shared_experts_stream.wait_stream(current_stream())

263 +

264 + return use_shared_experts_stream, hidden_states_clone

265 +

266 + def ensure_dp_chunking_init(self):

267 + if not self.use_dp_chunking or self.batched_hidden_states is not None:

268 + return

269 +

270 + states_shape: tuple[int, ...]

271 + logits_shape: tuple[int, ...]

272 +

273 + moe = self.moe_config

274 +

275 + if self.enable_dbo:

276 + states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)

277 + logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)

278 + else:

279 + states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)

280 + logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)

281 +

282 + self.batched_hidden_states = torch.zeros(

283 + states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()

284 + )

285 +

286 + self.batched_router_logits = torch.zeros(

287 + logits_shape,

288 + dtype=moe.router_logits_dtype,

289 + device=torch.cuda.current_device(),

290 + )

291 +

292 + def must_reduce_shared_expert_outputs(self) -> bool:

293 + """

294 + The shared_experts are typically computed using the RowParallelLinear

295 + layer. The result of this function is typically used as

296 + the reduce_results argument to the module.

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

298 + the shared_experts results immediately. Instead we reduce at the

299 + once at the end of the MoE op. (Refer to DeepSeekV2MoE module)

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

301 + GPU ranks in DP, produce the complete set of hidden_states.

302 + Therefore it is required that we reduce the shared_experts output

303 + early.

304 + """

305 + assert self.quant_method is not None

306 + return (

307 + self.quant_method.moe_mk is not None

308 + and self.quant_method.moe_mk.output_is_reduced()

309 + )

310 +

311 + def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):

312 + """

313 + Some combine kernels reduce across GPU ranks by default.

314 + """

315 + if self.must_reduce_shared_expert_outputs():

316 + return final_hidden_states

317 + else:

318 + return tensor_model_parallel_all_reduce(final_hidden_states)

319 +

320 + def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:

321 + """Apply transform for routed experts (e.g., latent projection).

322 +

323 + This is called by FusedMoE.forward_native. The original hidden_states

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

325 + routed experts get the transformed [S, moe_latent_size].

326 +

327 + TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be

328 + moved inside SharedFusedMoE to all-reduce on the smaller latent

329 + dimension.

330 + """

331 + if self.routed_input_transform is not None:

332 + result = self.routed_input_transform(hidden_states)

333 + # ReplicatedLinear returns (output, extra_bias) tuple.

334 + # We only need the output tensor; extra_bias is not used here.

335 + if isinstance(result, tuple):

336 + return result[0]

337 + return result

338 + return hidden_states

339 +

340 + def _reduce_output(

341 + self,

342 + states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

343 + trunc_sizes: list[int],

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

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

346 + return x[..., :trunc_size]

347 +

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

349 + return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)

350 +

351 + if (

352 + not self.moe_config.is_sequence_parallel

353 + and not self.use_dp_chunking

354 + and self.reduce_results

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

356 + ):

357 + func = reduce_and_trunc

358 + else:

359 + func = trunc

360 +

361 + if isinstance(states, tuple):

362 + return tuple(

363 + [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]

364 + )

365 + else:

366 + assert len(trunc_sizes) == 1

367 + return func(states, trunc_sizes[0])

368 +

369 + def _encode_layer_name(self) -> str:

370 + # Can be unavailable or None in unittests

371 + if (

372 + is_forward_context_available()

373 + and get_forward_context().all_moe_layers is not None

374 + ):

375 + return "from_forward_context"

376 + return self.layer_name

377 +

378 + def forward(

379 + self,

380 + hidden_states: torch.Tensor,

381 + router_logits: torch.Tensor,

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

383 + # For latent MoE: save ORIGINAL hidden_states before transform

384 + # (shared_experts need original dimension, routed experts use transformed)

385 + original_hidden_states = hidden_states

386 + original_hidden_dim = hidden_states.shape[-1]

387 +

388 + # Apply transform for routed experts (e.g., latent projection for latent MoE)

389 + hidden_states = self.apply_routed_input_transform(hidden_states)

390 +

391 + # This is the dimension after transform (for routed expert output slicing)

392 + transformed_hidden_dim = hidden_states.shape[-1]

393 + if self.moe_config.hidden_dim != transformed_hidden_dim:

394 + hidden_states = F.pad(

395 + hidden_states,

396 + (0, self.moe_config.hidden_dim - transformed_hidden_dim),

397 + mode="constant",

398 + value=0.0,

399 + )

400 +

401 + fused_output = self.moe_forward(

402 + hidden_states,

403 + router_logits,

404 + original_hidden_states,

405 + self._encode_layer_name(),

406 + )

407 +

408 + if isinstance(fused_output, tuple):

409 + orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim]

410 + else:

411 + orig_hidden_dims = [transformed_hidden_dim]

412 +

413 + return self._reduce_output(fused_output, orig_hidden_dims)

414 +

415 + def forward_impl_chunked(

416 + self,

417 + layer: torch.nn.Module,

418 + full_hidden_states: torch.Tensor,

419 + full_router_logits: torch.Tensor,

420 + shared_input: torch.Tensor | None,

421 + has_separate_shared_experts: bool,

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

423 + assert self.batched_hidden_states is not None

424 + assert self.batched_router_logits is not None

425 + assert self.batched_hidden_states.dtype == full_hidden_states.dtype, (

426 + f"{self.batched_hidden_states.dtype} == {full_hidden_states.dtype}"

427 + )

428 + assert self.batched_router_logits.dtype == full_router_logits.dtype, (

429 + f"{self.batched_router_logits.dtype} == {full_router_logits.dtype}"

430 + )

431 + # Check size compatibility.

432 + assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1)

433 + assert self.batched_router_logits.size(-1) == full_router_logits.size(-1)

434 +

435 + # TODO(bnell): Fix shared_expert_inputs w/chunking.

436 + # assert shared_input is None, (

437 + # "Routed input transform is not currently supported with DP chunking."

438 + # )

439 +

440 + full_fused_final_hidden_states = torch.empty_like(full_hidden_states)

441 + if self.shared_experts is not None:

442 + full_shared_final_hidden_states = torch.empty_like(full_hidden_states)

443 +

444 + def process_chunk(chunk_start, chunk_end, skip_result_store=False):

445 + chunk_size = chunk_end - chunk_start

446 + hidden_states = full_hidden_states[chunk_start:chunk_end, :]

447 + router_logits = full_router_logits[chunk_start:chunk_end, :]

448 +

449 + assert self.batched_hidden_states is not None

450 + assert self.batched_router_logits is not None

451 + # This is only true when DBO has been enabled in the config.

452 + # Both tensors will have an outer dimension for the ubatch id

453 + if self.batched_hidden_states.dim() == 3:

454 + assert self.batched_router_logits.dim() == 3

455 + batch_buffer_idx = dbo_current_ubatch_id()

456 + batched_hidden_states = self.batched_hidden_states[batch_buffer_idx, :]

457 + batched_router_logits = self.batched_router_logits[batch_buffer_idx, :]

458 + else:

459 + batched_hidden_states = self.batched_hidden_states

460 + batched_router_logits = self.batched_router_logits

461 +

462 + assert (

463 + batched_hidden_states.size(0) # type: ignore

464 + >= chunk_size

465 + )

466 + assert (

467 + batched_router_logits.size(0) # type: ignore

468 + >= chunk_size

469 + )

470 + staged_hidden_states = batched_hidden_states[:chunk_size, :] # type: ignore

471 + staged_router_logits = batched_router_logits[:chunk_size, :] # type: ignore

472 + staged_hidden_states.copy_(hidden_states, non_blocking=True)

473 + staged_router_logits.copy_(router_logits, non_blocking=True)

474 +

475 + # Matrix multiply.

476 + if self.quant_method.is_monolithic:

477 + final_hidden_states = self.quant_method.apply_monolithic(

478 + layer=layer,

479 + x=staged_hidden_states,

480 + router_logits=staged_router_logits,

481 + )

482 + else:

483 + topk_weights, topk_ids = self.router.select_experts(

484 + hidden_states=staged_hidden_states,

485 + router_logits=staged_router_logits,

486 + )

487 +

488 + final_hidden_states = self.quant_method.apply(

489 + layer=layer,

490 + x=staged_hidden_states,

491 + topk_weights=topk_weights,

492 + topk_ids=topk_ids,

493 + shared_experts_input=shared_input,

494 + )

495 +

496 + if has_separate_shared_experts:

497 + assert not isinstance(final_hidden_states, tuple)

498 + assert self.shared_experts is not None

499 +

500 + shared_output = self.shared_experts(staged_hidden_states)

501 +

502 + final_hidden_states = (

503 + shared_output,

504 + final_hidden_states,

505 + )

506 +

507 + if not skip_result_store:

508 + if self.shared_experts is None:

509 + full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

510 + final_hidden_states, non_blocking=True

511 + )

512 + else:

513 + full_shared_final_hidden_states[chunk_start:chunk_end, :].copy_(

514 + final_hidden_states[0], non_blocking=True

515 + )

516 + full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

517 + final_hidden_states[1], non_blocking=True

518 + )

519 +

520 + ctx = get_forward_context()

521 + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP

522 + max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu

523 + moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens

524 +

525 + # If the input to the MoE is sequence parallel then divide by sp_size

526 + # to find the maximum number of tokens for any individual dispatcher.

527 + if self.moe_config.is_sequence_parallel:

528 + max_tokens_across_dispatchers = cdiv(

529 + max_tokens_across_dispatchers, self.moe_config.sp_size

530 + )

531 +

532 + num_tokens = full_hidden_states.size(0)

533 + for chunk_idx, chunk_start_ in enumerate(

534 + range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)

535 + ):

536 + chunk_start = chunk_start_

537 + chunk_end = min(

538 + chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers

539 + )

540 + # clamp start and end

541 + chunk_start = min(chunk_start, num_tokens - 1)

542 + chunk_end = min(chunk_end, num_tokens)

543 + with ctx.dp_metadata.chunked_sizes(

544 + self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx

545 + ):

546 + process_chunk(

547 + chunk_start, chunk_end, skip_result_store=chunk_start_ >= num_tokens

548 + )

549 +

550 + if self.shared_experts is None:

551 + return full_fused_final_hidden_states

552 + else:

553 + return (full_shared_final_hidden_states, full_fused_final_hidden_states)

554 +

555 + def forward_impl(

556 + self,

557 + layer: torch.nn.Module,

558 + hidden_states: torch.Tensor,

559 + router_logits: torch.Tensor,

560 + shared_input: torch.Tensor | None,

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

562 + assert self.quant_method is not None

563 +

564 + self.ensure_dp_chunking_init()

565 +

566 + has_separate_shared_experts = (

567 + not self.quant_method.mk_owns_shared_expert

568 + and self.shared_experts is not None

569 + )

570 +

571 + use_chunked_impl = self.use_dp_chunking

572 +

573 + use_shared_experts_stream, hidden_states_clone = (

574 + self._maybe_setup_shared_experts_stream(

575 + hidden_states,

576 + shared_input,

577 + has_separate_shared_experts,

578 + use_chunked_impl,

579 + )

580 + )

581 +

582 + # If router/gate provided, then apply it here.

583 + # (Note: This code runs only when "overlapped mode" is on to allow

584 + # parallel execution of shared experts with the FusedMoE via

585 + # separate cuda stream)

586 + if self.gate is not None:

587 + router_logits, _ = self.gate(hidden_states)

588 +

589 + if use_chunked_impl:

590 + return self.forward_impl_chunked(

591 + layer,

592 + hidden_states,

593 + router_logits,

594 + shared_input,

595 + has_separate_shared_experts,

596 + )

597 +

598 + # NOTE(rob): once we finish migrating all the quant methods to use

599 + # MKs, we can remove the naive dispatch/combine path from here.

600 + do_naive_dispatch_combine = (

601 + self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk

602 + )

603 +

604 + ctx = get_forward_context()

605 + sp_ctx = (

606 + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)

607 + if ctx.dp_metadata

608 + else nullcontext()

609 + )

610 +

611 + with sp_ctx:

612 + extra_tensors = None

613 + if do_naive_dispatch_combine:

614 + post_quant_allgather = (

615 + self.quant_method is not None

616 + and self.moe_config.dp_size > 1

617 + and self.moe_config.use_ep

618 + and getattr(self.quant_method, "do_post_quant_allgather", False)

619 + )

620 + if post_quant_allgather:

621 + hidden_states_to_dispatch, extra_tensors = (

622 + self.quant_method.prepare_dp_allgather_tensor(

623 + layer, hidden_states, router_logits

624 + )

625 + )

626 + else:

627 + hidden_states_to_dispatch = hidden_states

628 +

629 + dispatch_res = get_ep_group().dispatch_router_logits(

630 + hidden_states_to_dispatch,

631 + router_logits,

632 + self.moe_config.is_sequence_parallel,

633 + extra_tensors=extra_tensors,

634 + )

635 + if extra_tensors is not None:

636 + (

637 + orig_hidden_states,

638 + router_logits,

639 + extra_tensors_combined,

640 + ) = dispatch_res

641 + hidden_states_combined = (

642 + orig_hidden_states,

643 + extra_tensors_combined[0],

644 + )

645 + else:

646 + hidden_states_combined, router_logits = dispatch_res

647 + orig_hidden_states = hidden_states_combined

648 + else:

649 + orig_hidden_states = hidden_states

650 +

651 + # Run shared experts before matrix multiply.

652 + # because matrix multiply maybe modify the hidden_states.

653 + if has_separate_shared_experts and not use_shared_experts_stream:

654 + assert self.shared_experts is not None

655 + shared_input = (

656 + shared_input if shared_input is not None else hidden_states

657 + )

658 + shared_output = self.shared_experts(shared_input)

659 +

660 + # NOTE: Similar with DP, PCP also needs dispatch and combine. For

661 + # simplicity, AgRsAll2All was added separately for PCP here. Maybe

662 + # we should modify All2AllManager abstract to better support PCP.

663 + if self.moe_config.pcp_size > 1:

664 + hidden_states = get_pcp_group().all_gather(

665 + hidden_states,

666 + dim=0,

667 + )

668 + router_logits = get_pcp_group().all_gather(

669 + router_logits,

670 + dim=0,

671 + )

672 +

673 + # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014).

674 + # Figure out nicer way to do this.

675 + if do_naive_dispatch_combine:

676 + x = hidden_states_combined

677 + x_orig = orig_hidden_states

678 + else:

679 + x = hidden_states

680 + x_orig = hidden_states

681 +

682 + # Matrix multiply.

683 + if self.quant_method.is_monolithic:

684 + final_hidden_states = self.quant_method.apply_monolithic(

685 + layer=layer,

686 + x=x,

687 + router_logits=router_logits,

688 + )

689 + else:

690 + topk_weights, topk_ids = self.router.select_experts(

691 + hidden_states=x_orig,

692 + router_logits=router_logits,

693 + )

694 +

695 + final_hidden_states = self.quant_method.apply(

696 + layer=layer,

697 + x=x, # The type signture of this is wrong due to the hack.

698 + topk_weights=topk_weights,

699 + topk_ids=topk_ids,

700 + shared_experts_input=shared_input,

701 + )

702 +

703 + if has_separate_shared_experts:

704 + assert self.shared_experts is not None

705 +

706 + if use_shared_experts_stream:

707 + # Run shared experts in parallel on a separate stream

708 + # NOTE: We start the separate stream here and mark the

709 + # sync end point immediately after it is done. This is

710 + # important to avoid excessive stream allocations by the cuda

711 + # graph replay later.

712 + with torch.cuda.stream(self.shared_experts_stream):

713 + # Note that hidden_states clone() is necessary here to avoid

714 + # conflict with the main stream

715 + shared_output = self.shared_experts(hidden_states_clone)

716 + current_stream().wait_stream(self.shared_experts_stream)

717 +

718 + final_hidden_states = (

719 + shared_output,

720 + final_hidden_states,

721 + )

722 +

723 + def combine_output(states: torch.Tensor) -> torch.Tensor:

724 + if do_naive_dispatch_combine:

725 + states = get_ep_group().combine(

726 + states, self.moe_config.is_sequence_parallel

727 + )

728 +

729 + if self.moe_config.pcp_size > 1:

730 + states = get_pcp_group().reduce_scatter(

731 + states,

732 + dim=0,

733 + )

734 +

735 + return states

736 +

737 + if self.shared_experts is not None:

738 + return (

739 + final_hidden_states[0],

740 + combine_output(final_hidden_states[1]),

741 + )

742 + else:

743 + return combine_output(final_hidden_states)

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

+34

@@ -0,0 +1,34 @@

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

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

3 +from abc import ABC, abstractmethod

4 +

5 +import torch

6 +

7 +

8 +class MoERunner(ABC):

9 + """

10 + Abstract base class for Mixture of Experts (MoE) runners.

11 +

12 + This class defines the interface that all MoE runner implementations must follow.

13 + MoE runners are responsible for executing the forward pass of MoE layers, handling

14 + expert routing, and managing tensor parallel operations.

15 + """

16 +

17 + @abstractmethod

18 + def forward(

19 + self,

20 + hidden_states: torch.Tensor,

21 + router_logits: torch.Tensor,

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

23 + raise NotImplementedError

24 +

25 + @abstractmethod

26 + def must_reduce_shared_expert_outputs(self) -> bool:

27 + raise NotImplementedError

28 +

29 + @abstractmethod

30 + def maybe_all_reduce_tensor_model_parallel(

31 + self,

32 + final_hidden_states: torch.Tensor,

33 + ):

34 + raise NotImplementedError

vllm/model_executor/layers/fused_moe/shared_fused_moe.py

-64

@@ -18,70 +18,6 @@ class SharedFusedMoE(FusedMoE):

18 can be interleaved with the fused all2all dispatch communication step.

19 """

21 - def __init__(

22 - self,

23 - shared_experts: torch.nn.Module | None,

24 - gate: torch.nn.Module | None = None,

25 - use_overlapped: bool = True,

26 - routed_input_transform: torch.nn.Module | None = None,

27 - **kwargs,

28 - ):

29 - # Pass has_shared_experts so FusedMoE.__init__ can set disable_inplace

30 - # without accessing self.shared_experts (submodules cannot be set before

31 - # Module.__init__()).

32 - kwargs["has_shared_experts"] = shared_experts is not None

33 - super().__init__(**kwargs)

34 - self._shared_experts = shared_experts

35 - self._routed_input_transform = routed_input_transform

36 -

37 - # Disable shared expert overlap if:

38 - # - we are using eplb with non-default backend, because of correctness issues

39 - # - we are using flashinfer with DP, since there nothing to gain

40 - # - we are using marlin kernels

41 - backend = self.moe_parallel_config.all2all_backend

42 - self.use_overlapped = (

43 - use_overlapped

44 - and not (

45 - (self.enable_eplb and backend != "allgather_reducescatter")

46 - or self.moe_parallel_config.use_fi_all2allv_kernels

47 - )

48 - and self._shared_experts is not None

49 - )

50 -

51 - self._gate = gate

52 -

53 - @property

54 - def shared_experts(self) -> torch.nn.Module | None:

55 - return self._shared_experts if self.use_overlapped else None

56 -

57 - @property

58 - def gate(self) -> torch.nn.Module | None:

59 - return self._gate if self.use_overlapped else None

60 -

61 - @property

62 - def is_internal_router(self) -> bool:

63 - return self.gate is not None

64 -

65 - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:

66 - """Apply transform for routed experts (e.g., latent projection).

67 -

68 - This is called by FusedMoE.forward_native. The original hidden_states

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

70 - routed experts get the transformed [S, moe_latent_size].

71 -

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

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

74 - dimension.

75 - """

76 - if self._routed_input_transform is not None:

77 - result = self._routed_input_transform(hidden_states)

78 - # ReplicatedLinear returns (output, extra_bias) tuple.

79 - # We only need the output tensor; extra_bias is not used here.

80 - if isinstance(result, tuple):

81 - return result[0]

82 - return result

83 - return hidden_states

84 -

21 def forward(

22 self,

23 hidden_states: torch.Tensor,

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+8 -1

@@ -55,6 +55,8 @@ logger = init_logger(__name__)

55 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

56 """MoE method without quantization."""

58 + # --8<-- [end:unquantized_fused_moe]

59 +

60 def __init__(self, moe: FusedMoEConfig):

61 super().__init__(moe)

62 self.unquantized_backend = select_unquantized_moe_backend(

@@ -90,8 +92,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

92 x: torch.Tensor,

93 topk_weights: torch.Tensor,

94 topk_ids: torch.Tensor,

95 + shared_experts_input: torch.Tensor | None,

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

94 - return self.forward_cuda(layer, x, topk_weights, topk_ids)

97 + return self.forward_cuda(layer, x, topk_weights, topk_ids, shared_experts_input)

99 @property

100 def is_monolithic(self) -> bool:

@@ -293,12 +296,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

296 x: torch.Tensor,

297 topk_weights: torch.Tensor,

298 topk_ids: torch.Tensor,

299 + shared_experts_input: torch.Tensor | None,

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

301 return self.forward(

302 layer=layer,

303 x=x,

304 topk_weights=topk_weights,

305 topk_ids=topk_ids,

306 + shared_experts_input=shared_experts_input,

307 )

308

309 def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

@@ -316,6 +321,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

321 x: torch.Tensor,

322 topk_weights: torch.Tensor,

323 topk_ids: torch.Tensor,

324 + shared_experts_input: torch.Tensor | None,

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

326 assert self.kernel is not None

327

@@ -329,6 +335,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

335 apply_router_weight_on_input=layer.apply_router_weight_on_input,

336 global_num_experts=layer.global_num_experts,

337 expert_map=layer.expert_map,

338 + shared_experts_input=shared_experts_input,

339 )

340

341 def forward_monolithic_cuda(

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -764,6 +764,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

764 x: torch.Tensor,

765 topk_weights: torch.Tensor,

766 topk_ids: torch.Tensor,

767 + shared_experts_input: torch.Tensor | None,

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

769 return fused_marlin_moe(

770 x,

vllm/model_executor/layers/quantization/bitsandbytes.py

@@ -501,6 +501,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

501 x: torch.Tensor,

502 topk_weights: torch.Tensor,

503 topk_ids: torch.Tensor,

504 + shared_experts_input: torch.Tensor | None,

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

506 from vllm.model_executor.layers.fused_moe import fused_experts

507

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

+10 -3

@@ -349,6 +349,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

349 x: torch.Tensor,

350 topk_weights: torch.Tensor,

351 topk_ids: torch.Tensor,

352 + shared_experts_input: torch.Tensor | None,

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

354 assert self.moe_mk is not None

355 return self.moe_mk(

@@ -361,7 +362,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

362 global_num_experts=layer.global_num_experts,

363 expert_map=layer.expert_map,

364 apply_router_weight_on_input=layer.apply_router_weight_on_input,

364 - shared_experts_input=layer._get_shared_experts_input(x),

365 + shared_experts_input=shared_experts_input,

366 )

367

368

@@ -645,6 +646,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

646 x: torch.Tensor,

647 topk_weights: torch.Tensor,

648 topk_ids: torch.Tensor,

649 + shared_experts_input: torch.Tensor | None,

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

651 assert not self.is_monolithic

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

@@ -673,7 +675,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

675 global_num_experts=layer.global_num_experts,

676 expert_map=layer.expert_map,

677 apply_router_weight_on_input=layer.apply_router_weight_on_input,

676 - shared_experts_input=layer._get_shared_experts_input(x),

678 + shared_experts_input=shared_experts_input,

679 )

680

681

@@ -1064,6 +1066,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1066 x: torch.Tensor,

1067 topk_weights: torch.Tensor,

1068 topk_ids: torch.Tensor,

1069 + shared_experts_input: torch.Tensor | None,

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

1071 assert not self.is_monolithic

1072 assert self.moe_mk is not None

@@ -1079,7 +1082,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1082 # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501

1083 expert_map=layer.expert_map,

1084 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1082 - shared_experts_input=layer._get_shared_experts_input(x),

1085 + shared_experts_input=shared_experts_input,

1086 )

1087

1088 @property

@@ -1203,6 +1206,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1206 x: torch.Tensor,

1207 topk_weights: torch.Tensor,

1208 topk_ids: torch.Tensor,

1209 + shared_experts_input: torch.Tensor | None,

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

1211 from vllm.model_executor.layers.fused_moe import fused_experts

1212

@@ -1713,6 +1717,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1717 x: torch.Tensor,

1718 topk_weights: torch.Tensor,

1719 topk_ids: torch.Tensor,

1720 + shared_experts_input: torch.Tensor | None,

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

1722 assert self.kernel_backend == "Marlin"

1723 return fused_marlin_moe(

@@ -1961,6 +1966,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1966 x: torch.Tensor,

1967 topk_weights: torch.Tensor,

1968 topk_ids: torch.Tensor,

1969 + shared_experts_input: torch.Tensor | None,

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

1971 from vllm.model_executor.layers.fused_moe import fused_experts

1972

@@ -2575,6 +2581,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2581 x: torch.Tensor,

2582 topk_weights: torch.Tensor,

2583 topk_ids: torch.Tensor,

2584 + shared_experts_input: torch.Tensor | None,

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

2586 if layer.enable_eplb:

2587 raise NotImplementedError(

vllm/model_executor/layers/quantization/experts_int8.py

@@ -140,6 +140,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

140 x: torch.Tensor,

141 topk_weights: torch.Tensor,

142 topk_ids: torch.Tensor,

143 + shared_experts_input: torch.Tensor | None,

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

145 from vllm.model_executor.layers.fused_moe import fused_experts

146

vllm/model_executor/layers/quantization/fp8.py

+2 -1

@@ -1010,6 +1010,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1010 x: torch.Tensor,

1011 topk_weights: torch.Tensor,

1012 topk_ids: torch.Tensor,

1013 + shared_experts_input: torch.Tensor | None,

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

1015 assert self.moe_mk is not None

1016 assert not self.is_monolithic

@@ -1023,7 +1024,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1024 global_num_experts=layer.global_num_experts,

1025 expert_map=layer.expert_map,

1026 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1026 - shared_experts_input=layer._get_shared_experts_input(x),

1027 + shared_experts_input=shared_experts_input,

1028 )

1029

1030

vllm/model_executor/layers/quantization/gguf.py

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

635 x: torch.Tensor,

636 topk_weights: torch.Tensor,

637 topk_ids: torch.Tensor,

638 + shared_experts_input: torch.Tensor | None,

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

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

641 if layer.apply_router_weight_on_input:

vllm/model_executor/layers/quantization/gptq_marlin.py

@@ -900,6 +900,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

900 x: torch.Tensor,

901 topk_weights: torch.Tensor,

902 topk_ids: torch.Tensor,

903 + shared_experts_input: torch.Tensor | None,

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

905 return fused_marlin_moe(

906 x,

vllm/model_executor/layers/quantization/modelopt.py

+4 -2

@@ -958,6 +958,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

958 x: torch.Tensor,

959 topk_weights: torch.Tensor,

960 topk_ids: torch.Tensor,

961 + shared_experts_input: torch.Tensor | None,

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

963 assert not self.is_monolithic

964

@@ -980,7 +981,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

981 global_num_experts=layer.global_num_experts,

982 expert_map=layer.expert_map,

983 apply_router_weight_on_input=layer.apply_router_weight_on_input,

983 - shared_experts_input=layer._get_shared_experts_input(x),

984 + shared_experts_input=shared_experts_input,

985 )

986

987

@@ -1524,6 +1525,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1525 x: torch.Tensor,

1526 topk_weights: torch.Tensor,

1527 topk_ids: torch.Tensor,

1528 + shared_experts_input: torch.Tensor | None,

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

1530 assert not self.is_monolithic

1531

@@ -1551,7 +1553,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1553 global_num_experts=layer.global_num_experts,

1554 expert_map=layer.expert_map,

1555 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1554 - shared_experts_input=layer._get_shared_experts_input(x),

1556 + shared_experts_input=shared_experts_input,

1557 )

1558

1559

vllm/model_executor/layers/quantization/moe_wna16.py

@@ -367,6 +367,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

367 x: torch.Tensor,

368 topk_weights: torch.Tensor,

369 topk_ids: torch.Tensor,

370 + shared_experts_input: torch.Tensor | None,

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

372 from vllm.model_executor.layers.fused_moe import fused_experts

373

vllm/model_executor/layers/quantization/mxfp4.py

@@ -900,6 +900,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

900 x: torch.Tensor,

901 topk_weights: torch.Tensor,

902 topk_ids: torch.Tensor,

903 + shared_experts_input: torch.Tensor | None,

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

905 assert not self.is_monolithic

906 if layer.enable_eplb:

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

@@ -419,6 +419,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

419 x: torch.Tensor,

420 topk_weights: torch.Tensor,

421 topk_ids: torch.Tensor,

422 + shared_experts_input: torch.Tensor | None,

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

424 if self.rocm_aiter_moe_enabled:

425 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

@@ -607,6 +608,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

608 x: torch.Tensor,

609 topk_weights: torch.Tensor,

610 topk_ids: torch.Tensor,

611 + shared_experts_input: torch.Tensor | None,

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

613 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

614 rocm_aiter_fused_experts,

@@ -977,6 +979,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

979 x: torch.Tensor,

980 topk_weights: torch.Tensor,

981 topk_ids: torch.Tensor,

982 + shared_experts_input: torch.Tensor | None,

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

984 if not self.emulate:

985 if (

vllm/v1/worker/gpu_worker.py

+5 -1

@@ -816,10 +816,14 @@ class Worker(WorkerBase):

816 for module in moe_modules:

817 module.moe_config.num_experts = num_local_experts * new_ep_size

818 module.global_num_experts = module.moe_config.num_experts

819 + tp_size = get_tp_group().world_size

820 + is_sequence_parallel = parallel_config.use_sequence_parallel_moe

821 + sp_size = tp_size if is_sequence_parallel else 1

822 module.moe_parallel_config = FusedMoEParallelConfig.make(

820 - tp_size_=get_tp_group().world_size,

823 + tp_size_=tp_size,

824 pcp_size_=get_pcp_group().world_size,

825 dp_size_=get_dp_group().world_size,

826 + sp_size_=sp_size,

827 vllm_parallel_config=parallel_config,

828 )

829 module.moe_config.moe_parallel_config = module.moe_parallel_config