Skip to content

vllm.model_executor.layers.fused_moe.router.fused_moe_router

Classes:

  • FusedMoERouter –

    FusedMoERouter is an abstract class that provides a 'select_experts'

FusedMoERouter

Bases: PluggableLayer, ABC

FusedMoERouter is an abstract class that provides a 'select_experts' method that is used for routing hidden states based on router logits.

Methods:

  • select_experts –

    Route the input hidden states to the top-k experts based on the

Attributes:

Source code in vllm/model_executor/layers/fused_moe/router/fused_moe_router.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class FusedMoERouter(PluggableLayer, ABC):
    """FusedMoERouter is an abstract class that provides a 'select_experts'
    method that is used for routing hidden states based on router logits.
    """

    def __init__(self, eplb_state: EplbLayerState | None = None):
        super().__init__()
        self._routing_replay_out: torch.Tensor | None = None
        self.eplb_state = eplb_state
        # Deepseek V4.1 vision checkpoints carry a second per-expert routing
        # bias for image sentinel tokens; attached by model code when present.
        self.bias_vl: torch.Tensor | None = None
        self.image_sentinel_lo: int = 0

    @abstractmethod
    def set_capture_fn(
        self,
        capture_fn: Callable[[torch.Tensor], None] | None,
    ) -> None:
        raise NotImplementedError

    @property
    @abstractmethod
    def routing_method_type(self) -> RoutingMethodType:
        raise NotImplementedError

    @property
    def kernel_routing_method_type(self) -> RoutingMethodType:
        """Routing method executed by kernels, before any router override."""
        return self.routing_method_type

    @property
    def supports_monolithic_kernels(self) -> bool:
        return True

    @property
    def masks_padding(self) -> bool:
        """Whether `select_experts` already gives padding rows expert id -1."""
        return False

    def prepare_monolithic_logits(self, router_logits: torch.Tensor) -> torch.Tensor:
        return router_logits

    def prepare_monolithic_bias(
        self, e_score_correction_bias: torch.Tensor | None
    ) -> torch.Tensor | None:
        return e_score_correction_bias

    @abstractmethod
    def _select_experts(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        topk_indices_dtype: torch.dtype | None = None,
        *,
        input_ids: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError

    def select_experts(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        topk_indices_dtype: torch.dtype | None = None,
        *,
        input_ids: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Route the input hidden states to the top-k experts based on the
        router logits.

        Returns:
            (topk_weights, topk_ids)
            (tuple[torch.Tensor, torch.Tensor]):
            The weights and expert ids computation result.

            **Compatibility**: When EPLB is not enabled, the returned ids are
            equivalent to global logical ids, so should be compatible with
            plain MoE implementations without redundant experts.

        """
        topk_weights, topk_ids = self._select_experts(
            hidden_states,
            router_logits,
            topk_indices_dtype=topk_indices_dtype,
            input_ids=input_ids,
        )
        # MoE kernels read routing tensors as dense buffers and ignore strides
        if not (topk_weights.is_contiguous() and topk_ids.is_contiguous()):
            raise ValueError(
                f"{type(self).__name__} returned non-contiguous topk_weights "
                "or topk_ids"
            )

        # Write routing data for non-monolithic path (Triton, etc.)
        # (set by bind_routing_capture_to_model during capturer init)
        if self._routing_replay_out is not None:
            self._routing_replay_out[: topk_ids.shape[0]].copy_(
                topk_ids.to(torch.int16)
            )

        return topk_weights, topk_ids

kernel_routing_method_type property

Routing method executed by kernels, before any router override.

masks_padding property

Whether select_experts already gives padding rows expert id -1.

select_experts(hidden_states, router_logits, topk_indices_dtype=None, *, input_ids=None)

Route the input hidden states to the top-k experts based on the router logits.

Returns:

Source code in vllm/model_executor/layers/fused_moe/router/fused_moe_router.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def select_experts(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
    topk_indices_dtype: torch.dtype | None = None,
    *,
    input_ids: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Route the input hidden states to the top-k experts based on the
    router logits.

    Returns:
        (topk_weights, topk_ids)
        (tuple[torch.Tensor, torch.Tensor]):
        The weights and expert ids computation result.

        **Compatibility**: When EPLB is not enabled, the returned ids are
        equivalent to global logical ids, so should be compatible with
        plain MoE implementations without redundant experts.

    """
    topk_weights, topk_ids = self._select_experts(
        hidden_states,
        router_logits,
        topk_indices_dtype=topk_indices_dtype,
        input_ids=input_ids,
    )
    # MoE kernels read routing tensors as dense buffers and ignore strides
    if not (topk_weights.is_contiguous() and topk_ids.is_contiguous()):
        raise ValueError(
            f"{type(self).__name__} returned non-contiguous topk_weights "
            "or topk_ids"
        )

    # Write routing data for non-monolithic path (Triton, etc.)
    # (set by bind_routing_capture_to_model during capturer init)
    if self._routing_replay_out is not None:
        self._routing_replay_out[: topk_ids.shape[0]].copy_(
            topk_ids.to(torch.int16)
        )

    return topk_weights, topk_ids