vllm.model_executor.layers.fused_moe.router.fused_topk_router

Classes:

• FusedTopKRouter –

  Default router using standard fused top-k routing.

Functions:

• dispatch_topk_softmax_func –

  Pick the softmax top-k implementation for one launch.

FusedTopKRouter

Bases: BaseRouter

Default router using standard fused top-k routing.

Source code in vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

class FusedTopKRouter(BaseRouter):
    """Default router using standard fused top-k routing."""

    def __init__(
        self,
        top_k: int,
        global_num_experts: int,
        scoring_func: str = "softmax",
        renormalize: bool = True,
        eplb_state: EplbLayerState | None = None,
    ):
        super().__init__(
            top_k=top_k,
            global_num_experts=global_num_experts,
            eplb_state=eplb_state,
        )
        self.renormalize = renormalize
        self.scoring_func = scoring_func

    @property
    def routing_method_type(self) -> RoutingMethodType:
        return get_routing_method_type(
            scoring_func=self.scoring_func,
            top_k=self.top_k,
            renormalize=self.renormalize,
            num_expert_group=None,
            has_e_score_bias=False,
        )

    def _compute_routing(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        indices_type: torch.dtype | None,
        *,
        input_ids: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute routing using standard fused top-k."""
        topk_weights, topk_ids, token_expert_indices = fused_topk(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=self.top_k,
            renormalize=self.renormalize,
            indices_type=indices_type,
            scoring_func=self.scoring_func,
        )

        return topk_weights, topk_ids

_compute_routing(hidden_states, router_logits, indices_type, *, input_ids=None)

Compute routing using standard fused top-k.

Source code in vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

def _compute_routing(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
    indices_type: torch.dtype | None,
    *,
    input_ids: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute routing using standard fused top-k."""
    topk_weights, topk_ids, token_expert_indices = fused_topk(
        hidden_states=hidden_states,
        gating_output=router_logits,
        topk=self.top_k,
        renormalize=self.renormalize,
        indices_type=indices_type,
        scoring_func=self.scoring_func,
    )

    return topk_weights, topk_ids

_aiter_topk_gating_supported(topk_weights, topk_indices, gating_output, num_shared_experts, shared_expert_scoring_func)

Whether AITER's topk_gating can serve this softmax top-k launch.

It needs contiguous gating rows and at most _AITER_TOPK_GATING_MAX_TOKENS tokens. Fused shared experts are scored with sigmoid after the routed top-k: 1, 2, 4 or 8 of them, with topk_weights and topk_indices sharing a row stride and topk_weights wide enough for the shared columns. Shared expert ids stay as the caller prefilled them.

Source code in vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

def _aiter_topk_gating_supported(
    topk_weights: torch.Tensor,
    topk_indices: torch.Tensor,
    gating_output: torch.Tensor,
    num_shared_experts: int,
    shared_expert_scoring_func: str,
) -> bool:
    """Whether AITER's ``topk_gating`` can serve this softmax top-k launch.

    It needs contiguous gating rows and at most ``_AITER_TOPK_GATING_MAX_TOKENS``
    tokens. Fused shared experts are scored with sigmoid after the routed
    top-k: 1, 2, 4 or 8 of them, with ``topk_weights`` and ``topk_indices``
    sharing a row stride and ``topk_weights`` wide enough for the shared
    columns. Shared expert ids stay as the caller prefilled them.
    """
    if (
        not gating_output.is_contiguous()
        or gating_output.shape[0] > _AITER_TOPK_GATING_MAX_TOKENS
    ):
        return False
    if num_shared_experts == 0:
        return not shared_expert_scoring_func
    return (
        shared_expert_scoring_func == "sigmoid"
        and num_shared_experts in (1, 2, 4, 8)
        and topk_weights.stride(0) == topk_indices.stride(0)
        and topk_weights.shape[-1] >= topk_indices.shape[-1] + num_shared_experts
    )

dispatch_topk_softmax_func(use_rocm_aiter=False, *, topk_weights=None, topk_indices=None, gating_output=None, num_shared_experts=0, shared_expert_scoring_func='')

Pick the softmax top-k implementation for one launch.

On ROCm with AITER, topk_gating is chosen on gfx942 and gfx950 when the launch described by the tensors and shared-expert arguments is one it supports; every other AITER launch uses the legacy topk_softmax. The tensors are optional: without them the launch cannot be judged and the legacy launcher is used.

Source code in vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

def dispatch_topk_softmax_func(
    use_rocm_aiter: bool = False,
    *,
    topk_weights: torch.Tensor | None = None,
    topk_indices: torch.Tensor | None = None,
    gating_output: torch.Tensor | None = None,
    num_shared_experts: int = 0,
    shared_expert_scoring_func: str = "",
) -> Callable[..., tuple[torch.Tensor, ...]]:
    """Pick the softmax top-k implementation for one launch.

    On ROCm with AITER, ``topk_gating`` is chosen on gfx942 and gfx950 when the
    launch described by the tensors and shared-expert arguments is one it supports;
    every other AITER launch uses the legacy ``topk_softmax``. The tensors are
    optional: without them the launch cannot be judged and the legacy launcher
    is used.
    """
    if use_rocm_aiter:
        if (
            rocm_aiter_ops.is_topk_gating_enabled()
            and topk_weights is not None
            and topk_indices is not None
            and gating_output is not None
            and _aiter_topk_gating_supported(
                topk_weights,
                topk_indices,
                gating_output,
                num_shared_experts,
                shared_expert_scoring_func,
            )
        ):
            return rocm_aiter_ops.topk_gating
        return rocm_aiter_ops.topk_softmax
    return vllm_topk_softmax