vllm.model_executor.models.transformers.moe

Transformers modeling backend mixin for Mixture of Experts (MoE) models.

Classes:

• MoEMixin –
• TransformersMoERunner –

  Custom MoERunner for the Transformers modeling backend.

MoEMixin

Bases: MixtureOfExperts, Base

Methods:

• recursive_replace –

  Initialize the MoE layers.

Attributes:

• mlp_layers (list[Module]) –

  MoE blocks whose experts were replaced, for the MixtureOfExperts methods.

Source code in vllm/model_executor/models/transformers/moe.py

class MoEMixin(MixtureOfExperts, Base):
    mlp_layers: list[nn.Module]
    """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods."""

    # Redeclared because `recursive_replace` assigns it below the first read
    num_local_physical_experts: int

    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
        # Skip MixtureOfExperts.__init__ and call the next class in MRO
        super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ):
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
        for moe_block in self.mlp_layers:
            moe_block.n_local_physical_experts = num_local_physical_experts
            moe_block.n_physical_experts = num_physical_experts
            moe_block.n_redundant_experts = self.num_redundant_experts
            moe_block.experts.update_expert_map()

    def recursive_replace(self):
        """Initialize the MoE layers."""
        experts_name = "experts"
        text_config = self.text_config

        # Positional arguments
        num_experts_names = CONFIG_FIELD_ALIASES["num_experts"]
        num_experts = getattr_iter(text_config, num_experts_names, None)
        top_k_names = CONFIG_FIELD_ALIASES["num_experts_per_token"]
        top_k = getattr_iter(text_config, top_k_names, None)
        assert top_k is not None
        hidden_size = text_config.hidden_size
        intermediate_size_names = CONFIG_FIELD_ALIASES["moe_intermediate_size"]
        intermediate_size = getattr_iter(text_config, intermediate_size_names, None)
        assert intermediate_size is not None

        num_shared_experts_names = CONFIG_FIELD_ALIASES["num_shared_experts"]
        num_shared_experts = getattr_iter(text_config, num_shared_experts_names, 0)

        # Common kwargs
        norm_topk_prob = getattr(text_config, "norm_topk_prob", None)

        # Routed scaling factor kwargs
        routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0)
        # aiter applies routed_scaling_factor internally
        apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled()
        routed_scaling_factor_kwargs = dict(
            routed_scaling_factor=routed_scaling_factor,
            apply_routed_scale_to_output=apply_routed_scale_to_output,
        )

        # Dtype the router computes in, if it is not the activation dtype.
        config_router_dtype = getattr(text_config, "moe_router_dtype", None)
        if config_router_dtype is not None:
            config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

        # Grouped topk routing kwargs
        num_expert_group = getattr(text_config, "n_group", None)
        topk_group = getattr(text_config, "topk_group", None)
        use_grouped_topk = num_expert_group is not None and topk_group is not None
        grouped_topk_routing_kwargs = dict(
            use_grouped_topk=use_grouped_topk,
            num_expert_group=num_expert_group,
            topk_group=topk_group,
        )

        # MoE activation function
        activation = "silu"
        wrapped_arch = self.config.architectures[0].lower()
        if "gptoss" in wrapped_arch:
            activation = "swigluoai"

        # Expert parallel load balancing kwargs
        enable_eplb = self.parallel_config.enable_eplb
        num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts

        # MixtureOfExperts mixin settings
        ep_size = get_ep_group().world_size

        self.mlp_layers = []
        self.moe_layers = []
        self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
        self.num_logical_experts = num_experts
        self.num_physical_experts = num_experts + num_redundant_experts
        self.num_local_physical_experts = self.num_physical_experts // ep_size
        self.num_routed_experts = num_experts
        self.num_shared_experts = num_shared_experts
        self.num_redundant_experts = num_redundant_experts

        # Down projections of shared experts consumed by FusedMoE
        shared_down_projs: list[tuple[nn.Module, str]] = []

        # Recursively fuse MoE layers
        def _recursive_replace(module: nn.Module, prefix: str):
            for child_name, child_module in module.named_children():
                qual_name = maybe_prefix(prefix, child_name)
                # Naive implementations will have experts as ModuleList
                is_modulelist = isinstance(child_module, nn.ModuleList)
                # Packed implementations will have experts as 3D tensors of shapes like:
                # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
                # down_proj = (num_experts, intermediate_size, hidden_size)
                params = list(child_module.parameters())
                is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
                if child_name == experts_name and (is_modulelist or is_3d):
                    # Alias for readability
                    moe_block = module
                    experts = child_module
                    # Class of the fused block (parent of gate/experts/shared)
                    moe_block_cls = type(moe_block).__name__
                    experts_cls = type(experts).__name__
                    # Do the experts have biases
                    has_bias = False
                    for experts_param_name, _ in experts.named_parameters():
                        if "bias" in experts_param_name:
                            has_bias = True
                            break
                    # Whether the expert weights are stored in transposed format.
                    # Comes from Transformers `use_experts_implementation` decorator.
                    is_transposed = getattr(experts, "is_transposed", False)
                    # If the config does not specify num_shared_experts, but
                    # the model has shared experts, we assume there is one.
                    if self.num_shared_experts == 0:
                        for moe_block_param_name, _ in moe_block.named_parameters():
                            if "shared_expert" in moe_block_param_name:
                                self.num_shared_experts = 1
                                break

                    # Only pay for the layer index if something is per-layer;
                    # `extract_layer_index` raises for unnumbered module paths.
                    per_layer = (top_k, intermediate_size, norm_topk_prob)
                    layer_idx = (
                        extract_layer_index(qual_name)
                        if any(isinstance(v, list) for v in per_layer)
                        else 0
                    )
                    layer_top_k = maybe_per_layer(top_k, layer_idx)
                    layer_renormalize = maybe_per_layer(norm_topk_prob, layer_idx)
                    if layer_renormalize is None:
                        layer_renormalize = layer_top_k > 1

                    kwargs: dict[str, Any] = dict(
                        num_experts=num_experts,
                        top_k=layer_top_k,
                        hidden_size=hidden_size,
                        intermediate_size=maybe_per_layer(intermediate_size, layer_idx),
                        renormalize=layer_renormalize,
                        quant_config=self.quant_config,
                        prefix=qual_name,
                        activation=activation,
                        enable_eplb=enable_eplb,
                        num_redundant_experts=num_redundant_experts,
                        has_bias=has_bias,
                        is_fused_checkpoint_transposed=is_transposed,
                        routed_experts_cls=TransformersRoutedExperts,
                    )
                    # _maybe_apply_routed_scale_to_output edge case. Transformers
                    # decoder layers do not compensate for dividing by scaling factor.
                    divides_shared_output = (
                        routed_scaling_factor != 1.0
                        and apply_routed_scale_to_output
                        and self.model_config.dtype == torch.float16
                    )
                    fuser = MoEBlockFuser.match(
                        moe_block, experts_name, divides_shared_output
                    )
                    if fuser is not None:
                        # MoE block forward is fully replaced.
                        # gate/router and shared expert (if any) runs in MoERunner.
                        shared_experts = fuser.shared_experts(moe_block, prefix)
                        # Store shared experts for later down projection adjustment
                        if shared_experts is not None:
                            hf_shared = shared_experts.shared_experts
                            glu_fuser = get_fuser(hf_shared, GLUFuser)
                            if glu_fuser is not None:
                                down_name = glu_fuser.down_name
                                if down_name is not None:
                                    shared_down_projs.append((hf_shared, down_name))
                        if scoring_func := fuser.scoring_func:
                            # Internal routing kwargs.
                            # Prefer config, otherwise read it from fuser.
                            router_dtype = config_router_dtype or fuser.router_dtype
                            gate = fuser.gate(moe_block, prefix, router_dtype)
                            kwargs |= dict(
                                scoring_func=scoring_func,
                                renormalize=(
                                    layer_renormalize
                                    if fuser.renormalize is None
                                    else fuser.renormalize
                                ),
                                gate=gate,
                            )
                            if router_dtype is not None:
                                kwargs["router_logits_dtype"] = router_dtype
                            if use_grouped_topk:
                                kwargs |= grouped_topk_routing_kwargs
                            if routed_scaling_factor != 1.0:
                                kwargs |= routed_scaling_factor_kwargs
                            bias = getattr(gate, "e_score_correction_bias", None)
                            if bias is not None:
                                kwargs["e_score_correction_bias"] = bias
                        else:
                            # External routing kwargs
                            kwargs |= dict(
                                num_expert_group=num_expert_group,
                                topk_group=topk_group,
                                custom_routing_function=partial(
                                    _hf_topk_router,
                                    gate=getattr(moe_block, fuser.gate_name),
                                ),
                            )

                        kwargs |= dict(
                            is_sequence_parallel=(
                                self.parallel_config.use_sequence_parallel_moe
                            ),
                            shared_experts=shared_experts,
                        )
                        fuser.rewrite_forward(moe_block)
                        logger.info_once(
                            "Fused: gate + experts%s (%s) -> MoERunner (%s routing)",
                            " + shared experts" if fuser.shared_name else "",
                            moe_block_cls,
                            "external" if fuser.scoring_func is None else "internal",
                        )
                    else:
                        # MoE block forward is unmodified.
                        # gate/router and shared expert (if any) runs in Transformers.
                        # We then smuggle the topk_ids in using a custom op.
                        moe_state = TransformersMoEState()

                        kwargs |= dict(
                            num_expert_group=num_expert_group,
                            topk_group=topk_group,
                            custom_routing_function=partial(
                                _hf_topk_router_precomputed, moe_state=moe_state
                            ),
                            runner_cls=TransformersMoERunner,
                            runner_args={"moe_state": moe_state},
                        )
                        logger.info_once(
                            "Fused: experts (%s) -> MoERunner (external routing)",
                            experts_cls,
                        )
                    fused_experts = FusedMoEFactory(**kwargs)
                    moe_block.experts = fused_experts
                    log_replacement(qual_name, experts, fused_experts)
                    # Update MixtureOfExperts mixin state
                    self.mlp_layers.append(moe_block)
                    self.moe_layers.append(fused_experts)
                else:
                    _recursive_replace(child_module, prefix=qual_name)

        _recursive_replace(self.model, prefix="model")
        self.num_moe_layers = len(self.moe_layers)
        # Continue with the replacement of layers in Base
        super().recursive_replace()
        # GLUFuser likely fused shared_experts. The down projection after the GLU
        # normally immediately reduces but we want FusedMoE to handle the reduction.
        for hf_shared, down_name in shared_down_projs:
            hf_shared.get_submodule(down_name).reduce_results = False

mlp_layers instance-attribute

MoE blocks whose experts were replaced, for the MixtureOfExperts methods.

recursive_replace()

Initialize the MoE layers.

Source code in vllm/model_executor/models/transformers/moe.py

def recursive_replace(self):
    """Initialize the MoE layers."""
    experts_name = "experts"
    text_config = self.text_config

    # Positional arguments
    num_experts_names = CONFIG_FIELD_ALIASES["num_experts"]
    num_experts = getattr_iter(text_config, num_experts_names, None)
    top_k_names = CONFIG_FIELD_ALIASES["num_experts_per_token"]
    top_k = getattr_iter(text_config, top_k_names, None)
    assert top_k is not None
    hidden_size = text_config.hidden_size
    intermediate_size_names = CONFIG_FIELD_ALIASES["moe_intermediate_size"]
    intermediate_size = getattr_iter(text_config, intermediate_size_names, None)
    assert intermediate_size is not None

    num_shared_experts_names = CONFIG_FIELD_ALIASES["num_shared_experts"]
    num_shared_experts = getattr_iter(text_config, num_shared_experts_names, 0)

    # Common kwargs
    norm_topk_prob = getattr(text_config, "norm_topk_prob", None)

    # Routed scaling factor kwargs
    routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0)
    # aiter applies routed_scaling_factor internally
    apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled()
    routed_scaling_factor_kwargs = dict(
        routed_scaling_factor=routed_scaling_factor,
        apply_routed_scale_to_output=apply_routed_scale_to_output,
    )

    # Dtype the router computes in, if it is not the activation dtype.
    config_router_dtype = getattr(text_config, "moe_router_dtype", None)
    if config_router_dtype is not None:
        config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

    # Grouped topk routing kwargs
    num_expert_group = getattr(text_config, "n_group", None)
    topk_group = getattr(text_config, "topk_group", None)
    use_grouped_topk = num_expert_group is not None and topk_group is not None
    grouped_topk_routing_kwargs = dict(
        use_grouped_topk=use_grouped_topk,
        num_expert_group=num_expert_group,
        topk_group=topk_group,
    )

    # MoE activation function
    activation = "silu"
    wrapped_arch = self.config.architectures[0].lower()
    if "gptoss" in wrapped_arch:
        activation = "swigluoai"

    # Expert parallel load balancing kwargs
    enable_eplb = self.parallel_config.enable_eplb
    num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts

    # MixtureOfExperts mixin settings
    ep_size = get_ep_group().world_size

    self.mlp_layers = []
    self.moe_layers = []
    self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
    self.num_logical_experts = num_experts
    self.num_physical_experts = num_experts + num_redundant_experts
    self.num_local_physical_experts = self.num_physical_experts // ep_size
    self.num_routed_experts = num_experts
    self.num_shared_experts = num_shared_experts
    self.num_redundant_experts = num_redundant_experts

    # Down projections of shared experts consumed by FusedMoE
    shared_down_projs: list[tuple[nn.Module, str]] = []

    # Recursively fuse MoE layers
    def _recursive_replace(module: nn.Module, prefix: str):
        for child_name, child_module in module.named_children():
            qual_name = maybe_prefix(prefix, child_name)
            # Naive implementations will have experts as ModuleList
            is_modulelist = isinstance(child_module, nn.ModuleList)
            # Packed implementations will have experts as 3D tensors of shapes like:
            # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
            # down_proj = (num_experts, intermediate_size, hidden_size)
            params = list(child_module.parameters())
            is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
            if child_name == experts_name and (is_modulelist or is_3d):
                # Alias for readability
                moe_block = module
                experts = child_module
                # Class of the fused block (parent of gate/experts/shared)
                moe_block_cls = type(moe_block).__name__
                experts_cls = type(experts).__name__
                # Do the experts have biases
                has_bias = False
                for experts_param_name, _ in experts.named_parameters():
                    if "bias" in experts_param_name:
                        has_bias = True
                        break
                # Whether the expert weights are stored in transposed format.
                # Comes from Transformers `use_experts_implementation` decorator.
                is_transposed = getattr(experts, "is_transposed", False)
                # If the config does not specify num_shared_experts, but
                # the model has shared experts, we assume there is one.
                if self.num_shared_experts == 0:
                    for moe_block_param_name, _ in moe_block.named_parameters():
                        if "shared_expert" in moe_block_param_name:
                            self.num_shared_experts = 1
                            break

                # Only pay for the layer index if something is per-layer;
                # `extract_layer_index` raises for unnumbered module paths.
                per_layer = (top_k, intermediate_size, norm_topk_prob)
                layer_idx = (
                    extract_layer_index(qual_name)
                    if any(isinstance(v, list) for v in per_layer)
                    else 0
                )
                layer_top_k = maybe_per_layer(top_k, layer_idx)
                layer_renormalize = maybe_per_layer(norm_topk_prob, layer_idx)
                if layer_renormalize is None:
                    layer_renormalize = layer_top_k > 1

                kwargs: dict[str, Any] = dict(
                    num_experts=num_experts,
                    top_k=layer_top_k,
                    hidden_size=hidden_size,
                    intermediate_size=maybe_per_layer(intermediate_size, layer_idx),
                    renormalize=layer_renormalize,
                    quant_config=self.quant_config,
                    prefix=qual_name,
                    activation=activation,
                    enable_eplb=enable_eplb,
                    num_redundant_experts=num_redundant_experts,
                    has_bias=has_bias,
                    is_fused_checkpoint_transposed=is_transposed,
                    routed_experts_cls=TransformersRoutedExperts,
                )
                # _maybe_apply_routed_scale_to_output edge case. Transformers
                # decoder layers do not compensate for dividing by scaling factor.
                divides_shared_output = (
                    routed_scaling_factor != 1.0
                    and apply_routed_scale_to_output
                    and self.model_config.dtype == torch.float16
                )
                fuser = MoEBlockFuser.match(
                    moe_block, experts_name, divides_shared_output
                )
                if fuser is not None:
                    # MoE block forward is fully replaced.
                    # gate/router and shared expert (if any) runs in MoERunner.
                    shared_experts = fuser.shared_experts(moe_block, prefix)
                    # Store shared experts for later down projection adjustment
                    if shared_experts is not None:
                        hf_shared = shared_experts.shared_experts
                        glu_fuser = get_fuser(hf_shared, GLUFuser)
                        if glu_fuser is not None:
                            down_name = glu_fuser.down_name
                            if down_name is not None:
                                shared_down_projs.append((hf_shared, down_name))
                    if scoring_func := fuser.scoring_func:
                        # Internal routing kwargs.
                        # Prefer config, otherwise read it from fuser.
                        router_dtype = config_router_dtype or fuser.router_dtype
                        gate = fuser.gate(moe_block, prefix, router_dtype)
                        kwargs |= dict(
                            scoring_func=scoring_func,
                            renormalize=(
                                layer_renormalize
                                if fuser.renormalize is None
                                else fuser.renormalize
                            ),
                            gate=gate,
                        )
                        if router_dtype is not None:
                            kwargs["router_logits_dtype"] = router_dtype
                        if use_grouped_topk:
                            kwargs |= grouped_topk_routing_kwargs
                        if routed_scaling_factor != 1.0:
                            kwargs |= routed_scaling_factor_kwargs
                        bias = getattr(gate, "e_score_correction_bias", None)
                        if bias is not None:
                            kwargs["e_score_correction_bias"] = bias
                    else:
                        # External routing kwargs
                        kwargs |= dict(
                            num_expert_group=num_expert_group,
                            topk_group=topk_group,
                            custom_routing_function=partial(
                                _hf_topk_router,
                                gate=getattr(moe_block, fuser.gate_name),
                            ),
                        )

                    kwargs |= dict(
                        is_sequence_parallel=(
                            self.parallel_config.use_sequence_parallel_moe
                        ),
                        shared_experts=shared_experts,
                    )
                    fuser.rewrite_forward(moe_block)
                    logger.info_once(
                        "Fused: gate + experts%s (%s) -> MoERunner (%s routing)",
                        " + shared experts" if fuser.shared_name else "",
                        moe_block_cls,
                        "external" if fuser.scoring_func is None else "internal",
                    )
                else:
                    # MoE block forward is unmodified.
                    # gate/router and shared expert (if any) runs in Transformers.
                    # We then smuggle the topk_ids in using a custom op.
                    moe_state = TransformersMoEState()

                    kwargs |= dict(
                        num_expert_group=num_expert_group,
                        topk_group=topk_group,
                        custom_routing_function=partial(
                            _hf_topk_router_precomputed, moe_state=moe_state
                        ),
                        runner_cls=TransformersMoERunner,
                        runner_args={"moe_state": moe_state},
                    )
                    logger.info_once(
                        "Fused: experts (%s) -> MoERunner (external routing)",
                        experts_cls,
                    )
                fused_experts = FusedMoEFactory(**kwargs)
                moe_block.experts = fused_experts
                log_replacement(qual_name, experts, fused_experts)
                # Update MixtureOfExperts mixin state
                self.mlp_layers.append(moe_block)
                self.moe_layers.append(fused_experts)
            else:
                _recursive_replace(child_module, prefix=qual_name)

    _recursive_replace(self.model, prefix="model")
    self.num_moe_layers = len(self.moe_layers)
    # Continue with the replacement of layers in Base
    super().recursive_replace()
    # GLUFuser likely fused shared_experts. The down projection after the GLU
    # normally immediately reduces but we want FusedMoE to handle the reduction.
    for hf_shared, down_name in shared_down_projs:
        hf_shared.get_submodule(down_name).reduce_results = False

TransformersMoERunner

Bases: MoERunner

Custom MoERunner for the Transformers modeling backend.

Methods:

• forward –

  In Transformers experts.forward will have this signature, which is why

Source code in vllm/model_executor/models/transformers/moe.py

@PluggableLayer.register("transformers_fused_moe")
class TransformersMoERunner(MoERunner):
    """Custom MoERunner for the Transformers modeling backend."""

    # --8<-- [end:transformers_fused_moe]
    def __init__(self, *args, moe_state: TransformersMoEState, **kwargs):
        super().__init__(*args, **kwargs)
        self._moe_state = moe_state
        self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel

    def forward(  # type: ignore[override]
        self,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
        **kwargs: Any,
    ) -> torch.Tensor:
        """In Transformers `experts.forward` will have this signature, which is why
        it deliberately does not match `MoERunner.forward`.

        We discard any extra kwargs because we cannot use them here."""
        # Note: we need to forward through a custom op so the topk_ids
        # can be transferred without interfering with cudagraphs.
        return torch.ops.vllm.transformers_moe_forward(
            hidden_states,
            topk_ids.to(torch.int32),
            topk_weights.to(torch.float32),
            self.layer_name,
        )

    def _forward_super(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
    ) -> torch.Tensor:
        return super().forward(hidden_states, topk_weights)

forward(hidden_states, topk_ids, topk_weights, **kwargs)

In Transformers experts.forward will have this signature, which is why it deliberately does not match MoERunner.forward.

We discard any extra kwargs because we cannot use them here.

Source code in vllm/model_executor/models/transformers/moe.py

def forward(  # type: ignore[override]
    self,
    hidden_states: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    **kwargs: Any,
) -> torch.Tensor:
    """In Transformers `experts.forward` will have this signature, which is why
    it deliberately does not match `MoERunner.forward`.

    We discard any extra kwargs because we cannot use them here."""
    # Note: we need to forward through a custom op so the topk_ids
    # can be transferred without interfering with cudagraphs.
    return torch.ops.vllm.transformers_moe_forward(
        hidden_states,
        topk_ids.to(torch.int32),
        topk_weights.to(torch.float32),
        self.layer_name,
    )

_hf_topk_router(hidden_states, gating_output, topk, renormalize, gate)

Route with a Transformers gate returning (_, topk_weights, topk_ids).

Source code in vllm/model_executor/models/transformers/moe.py

def _hf_topk_router(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    gate: nn.Module,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Route with a Transformers gate returning `(_, topk_weights, topk_ids)`."""
    _, topk_weights, topk_ids = gate(hidden_states)
    return topk_weights, topk_ids

_hf_topk_router_precomputed(hidden_states, gating_output, topk, renormalize, moe_state)

Return (topk_weights, topk_ids) the HF block forward already computed: weights arrive as gating_output, ids via moe_state.

Source code in vllm/model_executor/models/transformers/moe.py

def _hf_topk_router_precomputed(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    moe_state: TransformersMoEState,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return `(topk_weights, topk_ids)` the HF block forward already computed:
    weights arrive as `gating_output`, ids via `moe_state`."""
    topk_weights = gating_output
    topk_ids = moe_state.topk_ids
    assert topk_ids is not None
    # Handle all gather in expert parallel
    if topk_ids.size(0) != hidden_states.size(0):
        dp_metadata = get_forward_context().dp_metadata
        assert dp_metadata is not None
        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
        assert sizes is not None
        is_sp = moe_state.is_sequence_parallel
        group = get_ep_group() if is_sp else get_dp_group()
        assert sizes[group.rank_in_group] == topk_ids.shape[0]
        (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
    return topk_weights, topk_ids

_transformers_moe_forward(hidden_states, topk_ids, topk_weights, layer_name)

Store the topk_ids in the layer and call the actual forward.

Source code in vllm/model_executor/models/transformers/moe.py

def _transformers_moe_forward(
    hidden_states: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    """Store the `topk_ids` in the layer and call the actual forward."""
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]
    self._moe_state.topk_ids = topk_ids
    return self._forward_super(hidden_states, topk_weights)