vllm.model_executor.layers.fused_moe.modular_kernel

Classes:

• ExpertTokensMetadata –

  Metadata regarding expert-token routing.

• FusedMoEActivationFormat –

  The standard activation format (num_tokens, hidden dim).

• FusedMoEExperts –
• FusedMoEExpertsModular –

  An abstract base class for the [Permute-Experts-Unpermute] step described

• FusedMoEExpertsMonolithic –

  An abstract base class for the [Permute-Experts-Unpermute] step described

• FusedMoEKernel –
• FusedMoEKernelModularImpl –
• FusedMoEKernelMonolithicImpl –
• FusedMoEPrepareAndFinalize –

  An abstract base class for the [Quantize-Prepare] and [Finalize] steps

• FusedMoEPrepareAndFinalizeModular –

  An abstract base class for the [Quantize-Prepare] and [Finalize] steps

• FusedMoEPrepareAndFinalizeMonolithic –

  An abstract base class for the [Quantize-Prepare] and [Finalize] steps

• TopKWeightAndReduce –

  An abstract base class for weight application and reduction implementations.

ExpertTokensMetadata dataclass

Metadata regarding expert-token routing.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@dataclass
class ExpertTokensMetadata:
    """Metadata regarding expert-token routing."""

    expert_num_tokens: torch.Tensor | None
    expert_num_tokens_cpu: torch.Tensor | None
    psum_recv_per_rank: torch.Tensor | None = None

    @staticmethod
    def make_from_list(
        expert_num_tokens_list: list[int], device: str
    ) -> "ExpertTokensMetadata":
        expert_num_tokens_cpu = torch.tensor(
            expert_num_tokens_list,
            device="cpu",
            dtype=torch.int32,
            pin_memory=PIN_MEMORY,
        )
        return ExpertTokensMetadata(
            expert_num_tokens=expert_num_tokens_cpu.to(device, non_blocking=True),
            expert_num_tokens_cpu=expert_num_tokens_cpu,
        )

FusedMoEActivationFormat

Bases: Enum

The standard activation format (num_tokens, hidden dim).

Attributes:

• Standard –

  The batched experts format (num experts, max tokens per expert, hidden dim)

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEActivationFormat(Enum):
    """The standard activation format (num_tokens, hidden dim)."""

    Standard = ("standard",)
    """
    The batched experts format (num experts, max tokens per expert, hidden dim)
    """
    BatchedExperts = ("batched_experts",)

Standard = ('standard',) class-attribute instance-attribute

The batched experts format (num experts, max tokens per expert, hidden dim)

FusedMoEExperts

Bases: ABC

Methods:

• __init__ –

  moe_config: MoE layer configuration.

• activation_format –

  A property which is a tuple of the input and output activation formats

• supports_lora –

  Return True if this expert impl natively handles LoRA.

• supports_packed_ue8m0_act_scales –

  A flag indicating whether or not this class can process packed ue8m0

Attributes:

• expects_unquantized_inputs (bool) –

  Whether or not the PrepareFinalize should defer input quantization

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEExperts(ABC):
    # ROCm AITER kernels consume a 0/1 local-expert mask (with a trailing
    # sentinel slot); every other backend consumes the canonical -1/local-slot
    # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand
    # the active experts kernel.
    consumes_expert_mask: bool = False

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        quant_config: FusedMoEQuantConfig,
        max_num_tokens: int | None = None,
        num_dispatchers: int | None = None,
    ):
        """moe_config: MoE layer configuration.
        quant_config: Quantization parameters for this experts instance.
        """
        if self.activation_format() == FusedMoEActivationFormat.Standard and (
            max_num_tokens is not None or num_dispatchers is not None
        ):
            raise ValueError(
                "max_num_tokens and num_dispatchers should only be set for "
                "BatchedExperts activation format."
            )
        elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and (
            max_num_tokens is None or num_dispatchers is None
        ):
            raise ValueError(
                "max_num_tokens and num_dispatchers must be set for "
                "BatchedExperts activation format."
            )

        self.moe_config = moe_config
        self.quant_config = quant_config
        self.activation_config = ApplyMoEActivationConfig.from_configs(
            moe_config, quant_config
        )
        self.max_num_tokens = max_num_tokens
        self.num_dispatchers = num_dispatchers

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:  # noqa: B027
        pass

    @staticmethod
    def is_monolithic() -> bool:
        raise NotImplementedError("Implemented by subclasses.")

    @property
    def expects_unquantized_inputs(self) -> bool:
        """Whether or not the PrepareFinalize should defer input quantization
        in the prepare step. If True, then the Experts kernel will
        execute the input quantization itself.

        Sample subclasses that override are AITER and FlashInfer CUTLASS.
        """
        return False

    @staticmethod
    @abstractmethod
    def activation_format() -> FusedMoEActivationFormat:
        """A property which is a tuple of the input and output activation formats
        for the 'apply' method.
        """
        raise NotImplementedError

    #
    # Various helpers for registering support for various features.
    # Used by the oracle to select a particular kernel for a deployment.
    #

    @staticmethod
    def is_supported_config(
        cls: type["FusedMoEExperts"],
        moe_config: FusedMoEConfig,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
        activation_format: FusedMoEActivationFormat,
    ) -> tuple[bool, str | None]:
        def _make_reason(reason: str) -> str:
            return f"kernel does not support {reason}"

        if not cls._supports_current_device():
            return False, _make_reason(f"current device {current_platform.device_name}")
        elif not (moe_config.is_act_and_mul or cls._supports_no_act_and_mul()):
            return False, _make_reason("no act_and_mul MLP layer")
        elif not cls._supports_activation(moe_config.activation):
            return False, _make_reason(f"{moe_config.activation} activation")
        elif not cls._supports_quant_scheme(weight_key, activation_key):
            return False, _make_reason(
                f"quantization scheme {weight_key}x{activation_key}"
            )
        elif not cls._supports_parallel_config(moe_config.moe_parallel_config):
            return False, _make_reason(
                f"parallel config {moe_config.moe_parallel_config}"
            )
        elif moe_config.has_hash_routing and cls.is_monolithic():
            return False, _make_reason("hash routing")
        elif not cls._supports_routing_method(
            moe_config.routing_method, weight_key, activation_key
        ):
            return False, _make_reason(f"routing method {moe_config.routing_method}")
        elif not cls._supports_router_logits_dtype(
            moe_config.router_logits_dtype,
            moe_config.routing_method,
        ):
            return False, _make_reason(
                f"router logits dtype {moe_config.router_logits_dtype}"
            )
        elif not cls._supports_shape(moe_config.hidden_dim):
            return False, _make_reason(
                f"{moe_config.hidden_dim} hidden dim is not supported"
            )
        elif activation_format != cls.activation_format():
            return False, _make_reason(f"{activation_format.value} activation format")
        elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance():
            return False, _make_reason("batch invariance")
        elif moe_config.is_lora_enabled and not cls.supports_lora():
            return False, _make_reason("LoRA")
        return True, None

    @staticmethod
    @abstractmethod
    def _supports_current_device() -> bool:
        """Whether the kernel supports the current device type
        (compute cability and current platform).
        """
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def _supports_no_act_and_mul() -> bool:
        """Whether the kernel supports act_and_mul=False, i.e.
        non-gated MoE models like Nemotron-Nano.
        """
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def _supports_activation(activation: MoEActivation) -> bool:
        """Whether the kernel supports a particular act function."""
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
        """Whether the kernel supports deployment in particular parallel config.

        Can be overridden if a kernel does not support EP, SP or some other
        configuration.
        """
        raise NotImplementedError

    @staticmethod
    def _supports_routing_method(
        routing_method: RoutingMethodType,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        """Whether the kernel supports a routing method (e.g. GroupedTopK).

        Can be overridden by monolithic kernels that execute the router
        in addition to the experts if certain routers are not supported.
        """
        return True

    @staticmethod
    def _supports_router_logits_dtype(
        router_logits_dtype: torch.dtype | None,
        routing_method: RoutingMethodType,
    ) -> bool:
        """Whether a kernel supports a particular dtype for router logits input.

        Can be overridden by monolithic kernels that execute the router
        in addition to the experts if certain dtypes are not supported.
        """
        return True

    @staticmethod
    def _supports_shape(hidden_dim: int) -> bool:
        """Whether a kernel supports a particular shape. Can be overridden if a kernel
        has specific shape requirements.
        """
        return True

    @staticmethod
    def _supports_batch_invariance() -> bool:
        """Whether the kernel supports batch invariance, i.e. the output does not
        depend on the order of the tokens in the input batch. This is useful
        for determining if the kernel can used with VLLM_BATCH_INVARIANT=1.
        """
        return False

    #
    # Various helpers for accessing quantization parameters from the
    # quant_config.
    #

    @property
    def quant_dtype(self) -> torch.dtype | str | None:
        return self.quant_config.quant_dtype

    @property
    def weight_quant_dtype(self) -> torch.dtype | str | None:
        return self.quant_config.weight_quant_dtype

    @property
    def block_shape(self) -> list[int] | None:
        return self.quant_config.block_shape

    @property
    def per_act_token_quant(self) -> bool:
        return self.quant_config.per_act_token_quant

    @property
    def per_out_ch_quant(self) -> bool:
        return self.quant_config.per_out_ch_quant

    @property
    def a1_scale(self) -> torch.Tensor | None:
        return self.quant_config.a1_scale

    @property
    def a2_scale(self) -> torch.Tensor | None:
        return self.quant_config.a2_scale

    @property
    def a1_gscale(self) -> torch.Tensor | None:
        return self.quant_config.a1_gscale

    @property
    def a2_gscale(self) -> torch.Tensor | None:
        return self.quant_config.a2_gscale

    @property
    def w1_scale(self) -> torch.Tensor | None:
        return self.quant_config.w1_scale

    @property
    def w2_scale(self) -> torch.Tensor | None:
        return self.quant_config.w2_scale

    @property
    def w1_zp(self) -> torch.Tensor | None:
        return self.quant_config.w1_zp

    @property
    def w2_zp(self) -> torch.Tensor | None:
        return self.quant_config.w2_zp

    @property
    def w1_bias(self) -> torch.Tensor | None:
        return self.quant_config.w1_bias

    @property
    def w2_bias(self) -> torch.Tensor | None:
        return self.quant_config.w2_bias

    @property
    def g1_alphas(self) -> torch.Tensor | None:
        return self.quant_config.g1_alphas

    @property
    def g2_alphas(self) -> torch.Tensor | None:
        return self.quant_config.g2_alphas

    @staticmethod
    def supports_lora() -> bool:
        """Return True if this expert impl natively handles LoRA.

        LoRA-aware experts should mix in LoRAExpertsMixin, which flips this
        to True and provides the per-forward LoRA state plumbing.
        """
        return False

    def supports_packed_ue8m0_act_scales(self) -> bool:
        """A flag indicating whether or not this class can process packed ue8m0
        activation scales.
        """
        return False

expects_unquantized_inputs property

Whether or not the PrepareFinalize should defer input quantization in the prepare step. If True, then the Experts kernel will execute the input quantization itself.

Sample subclasses that override are AITER and FlashInfer CUTLASS.

__init__(moe_config, quant_config, max_num_tokens=None, num_dispatchers=None)

moe_config: MoE layer configuration. quant_config: Quantization parameters for this experts instance.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def __init__(
    self,
    moe_config: FusedMoEConfig,
    quant_config: FusedMoEQuantConfig,
    max_num_tokens: int | None = None,
    num_dispatchers: int | None = None,
):
    """moe_config: MoE layer configuration.
    quant_config: Quantization parameters for this experts instance.
    """
    if self.activation_format() == FusedMoEActivationFormat.Standard and (
        max_num_tokens is not None or num_dispatchers is not None
    ):
        raise ValueError(
            "max_num_tokens and num_dispatchers should only be set for "
            "BatchedExperts activation format."
        )
    elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and (
        max_num_tokens is None or num_dispatchers is None
    ):
        raise ValueError(
            "max_num_tokens and num_dispatchers must be set for "
            "BatchedExperts activation format."
        )

    self.moe_config = moe_config
    self.quant_config = quant_config
    self.activation_config = ApplyMoEActivationConfig.from_configs(
        moe_config, quant_config
    )
    self.max_num_tokens = max_num_tokens
    self.num_dispatchers = num_dispatchers

_supports_activation(activation) abstractmethod staticmethod

Whether the kernel supports a particular act function.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
@abstractmethod
def _supports_activation(activation: MoEActivation) -> bool:
    """Whether the kernel supports a particular act function."""
    raise NotImplementedError

_supports_batch_invariance() staticmethod

Whether the kernel supports batch invariance, i.e. the output does not depend on the order of the tokens in the input batch. This is useful for determining if the kernel can used with VLLM_BATCH_INVARIANT=1.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_batch_invariance() -> bool:
    """Whether the kernel supports batch invariance, i.e. the output does not
    depend on the order of the tokens in the input batch. This is useful
    for determining if the kernel can used with VLLM_BATCH_INVARIANT=1.
    """
    return False

_supports_current_device() abstractmethod staticmethod

Whether the kernel supports the current device type (compute cability and current platform).

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
@abstractmethod
def _supports_current_device() -> bool:
    """Whether the kernel supports the current device type
    (compute cability and current platform).
    """
    raise NotImplementedError

_supports_no_act_and_mul() abstractmethod staticmethod

Whether the kernel supports act_and_mul=False, i.e. non-gated MoE models like Nemotron-Nano.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
@abstractmethod
def _supports_no_act_and_mul() -> bool:
    """Whether the kernel supports act_and_mul=False, i.e.
    non-gated MoE models like Nemotron-Nano.
    """
    raise NotImplementedError

_supports_parallel_config(moe_parallel_config) abstractmethod staticmethod

Whether the kernel supports deployment in particular parallel config.

Can be overridden if a kernel does not support EP, SP or some other configuration.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
@abstractmethod
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
    """Whether the kernel supports deployment in particular parallel config.

    Can be overridden if a kernel does not support EP, SP or some other
    configuration.
    """
    raise NotImplementedError

_supports_router_logits_dtype(router_logits_dtype, routing_method) staticmethod

Whether a kernel supports a particular dtype for router logits input.

Can be overridden by monolithic kernels that execute the router in addition to the experts if certain dtypes are not supported.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_router_logits_dtype(
    router_logits_dtype: torch.dtype | None,
    routing_method: RoutingMethodType,
) -> bool:
    """Whether a kernel supports a particular dtype for router logits input.

    Can be overridden by monolithic kernels that execute the router
    in addition to the experts if certain dtypes are not supported.
    """
    return True

_supports_routing_method(routing_method, weight_key, activation_key) staticmethod

Whether the kernel supports a routing method (e.g. GroupedTopK).

Can be overridden by monolithic kernels that execute the router in addition to the experts if certain routers are not supported.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_routing_method(
    routing_method: RoutingMethodType,
    weight_key: QuantKey | None,
    activation_key: QuantKey | None,
) -> bool:
    """Whether the kernel supports a routing method (e.g. GroupedTopK).

    Can be overridden by monolithic kernels that execute the router
    in addition to the experts if certain routers are not supported.
    """
    return True

_supports_shape(hidden_dim) staticmethod

Whether a kernel supports a particular shape. Can be overridden if a kernel has specific shape requirements.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_shape(hidden_dim: int) -> bool:
    """Whether a kernel supports a particular shape. Can be overridden if a kernel
    has specific shape requirements.
    """
    return True

activation_format() abstractmethod staticmethod

A property which is a tuple of the input and output activation formats for the 'apply' method.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
@abstractmethod
def activation_format() -> FusedMoEActivationFormat:
    """A property which is a tuple of the input and output activation formats
    for the 'apply' method.
    """
    raise NotImplementedError

supports_lora() staticmethod

Return True if this expert impl natively handles LoRA.

LoRA-aware experts should mix in LoRAExpertsMixin, which flips this to True and provides the per-forward LoRA state plumbing.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def supports_lora() -> bool:
    """Return True if this expert impl natively handles LoRA.

    LoRA-aware experts should mix in LoRAExpertsMixin, which flips this
    to True and provides the per-forward LoRA state plumbing.
    """
    return False

supports_packed_ue8m0_act_scales()

A flag indicating whether or not this class can process packed ue8m0 activation scales.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def supports_packed_ue8m0_act_scales(self) -> bool:
    """A flag indicating whether or not this class can process packed ue8m0
    activation scales.
    """
    return False

FusedMoEExpertsModular

Bases: FusedMoEExperts

An abstract base class for the [Permute-Experts-Unpermute] step described above.

Methods:

• adjust_N_for_activation –

  Calculate the output dimension for the activation function.

• apply –

  This function computes the intermediate result of a Mixture of Experts

• moe_problem_size –

  Extract the MoE problem size from the given tensor arguments:

• workspace_dtype –

  Workspace type: The dtype to use for the workspace tensors.

• workspace_shapes –

  Compute the shapes for the temporary and final outputs of the two gemms

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEExpertsModular(FusedMoEExperts):
    """An abstract base class for the [Permute-Experts-Unpermute] step described
    above.
    """

    @staticmethod
    def is_monolithic() -> bool:
        return False

    def moe_problem_size(
        self,
        a1: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_ids: torch.Tensor,
    ) -> tuple[int, int, int, int, int]:
        """Extract the MoE problem size from the given tensor arguments:
        - a: The hidden states, input to the MoE layer.
        - w1: The first set of expert weights.
        - w2: The second set of expert weights.
        - topk_ids: The topk ids.

        Note: extracting the problem shape from the weight and activation
        tensors is not obvious.  It needs to be done this way specifically
        due to subtle issues with particular kernels, e.g. the int4 kernels
        divide the trailing dimension by two, so it's not "correct" to
        extract N or K from the trailing dimension of w1 or w2.  Similarly,
        some kernels transpose the weights, so this needs to be kept in mind.

        Note: This implementation covers most cases. However, if experts
        require a specialized implementation, like MarlinExperts, they are free
        to override this function.
        """
        assert len(w1.shape) == 3 and len(w2.shape) == 3
        E, N, _ = w1.shape
        K = a1.size(-1)

        if a1.dim() == 2:
            # Make sure we are using the correct a1 (pre-permute).
            assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
            M = a1.size(0)
        else:
            assert a1.dim() == 3
            assert a1.size(0) == E, f"{a1.size(0)} == {E}"
            M = a1.size(1)  # This is max_num_tokens

        assert topk_ids.dim() == 2
        topk = topk_ids.size(1)

        return E, M, N, K, topk

    def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
        """Workspace type: The dtype to use for the workspace tensors."""
        return act_dtype

    @abstractmethod
    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        """Compute the shapes for the temporary and final outputs of the two gemms
        and activation in the fused expert function.  Since the gemms are
        independent, the workspace for the first gemm can be shared with the
        workspace for the last gemm.

        Inputs:
        - M: number of tokens.
        - N: Row (or column) dimension of expert weights.
        - K: hidden dimension
        - topk: The number of top-k experts to select.
        - global_num_experts: global number of experts.
        - local_num_experts: local number of experts due to DP/EP.
        - expert_tokens_meta: number of tokens per expert metadata for batched
                              format.

        Returns a tuple of:
        - workspace13 shape tuple: must be large enough to hold the
          result of either expert gemm.
        - workspace2 shape tuple: must be large enough to hold the
          result of the activation function.
        - output shape tuple: must be exact size of the final gemm output.
        - Note: workspace shapes can be 0 if the workspace is not needed.
          But in order for activation chunking to work, the first dimension
          of each tuple must be the number of tokens when the shape is
          not 0.
        """
        raise NotImplementedError

    @staticmethod
    def adjust_N_for_activation(N: int, activation: MoEActivation) -> int:
        """Calculate the output dimension for the activation function.

        For *_no_mul activations (e.g. relu2_no_mul),
        there's no gate/up split, so output size equals input size (N).

        For regular gated activations (e.g., silu, gelu, swigluoai),
        output size is N // 2 due to gate × activation(up) multiplication.

        Args:
            N: The intermediate size (width of w1/w3 weights).
            activation: The activation function enum.

        Returns:
            The output dimension after activation.

        """
        return N if not activation.is_gated else N // 2

    def activation(
        self,
        activation: MoEActivation,
        output: torch.Tensor,
        input: torch.Tensor,
        *,
        topk_ids: torch.Tensor | None = None,
        expert_map: torch.Tensor | None = None,
        valid_token_counts: torch.Tensor | None = None,
    ) -> None:
        apply_moe_activation(
            activation,
            output,
            input,
            activation_config=self.activation_config,
            topk_ids=topk_ids,
            expert_map=expert_map,
            valid_token_counts=valid_token_counts,
        )

    @abstractmethod
    def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:
        raise NotImplementedError

    @abstractmethod
    def apply(
        self,
        output: torch.Tensor,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        a2_scale: torch.Tensor | None,
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: ExpertTokensMetadata | None,
        apply_router_weight_on_input: bool,
    ) -> UnfinalizedMoEOutput | None:
        """This function computes the intermediate result of a Mixture of Experts
        (MoE) layer using two sets of weights, w1 and w2.

        Writes into `output` and returns None, unless the implementation stopped
        after GEMM2 and left the top-k reduction to a fused consumer, in which
        case `output` is untouched and the unfinalized result is returned.

        Args:
            output: (torch.Tensor): The unweighted, unreduced output tensor.
            hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
                layer.
            w1 (torch.Tensor): The first set of expert weights.
            w2 (torch.Tensor): The second set of expert weights.
            topk_weights: A map of row to expert weights. Some implementations
                choose to do weight application.
            topk_ids (torch.Tensor): A map of row to expert id.
            activation (str): The activation function to apply after the first
                MoE layer.
            global_num_experts (int): The total number of experts in the global
                expert space.
            expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
                from the global expert space to the local expert space of the
                expert parallel shard.
            a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
                used for a1. Result of quantization from prepare/finalize and not
                from the FusedMoEQuantConfig.
            a2_scale (Optional[torch.Tensor]): Optional quantized scale to be
                used for the second gemm's activations.
            workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
                must be large enough to hold output of either MoE gemm.
            workspace2 (torch.Tensor): A scratch tensor used for the activation
                function.
            expert_tokens_meta (Optional[ExpertTokensMetadata]): An optional
                ExpertTokensMetadata object containing gpu/cpu tensors
                as big as the number of local experts with the information about
                the number of tokens assigned to each local expert.
            apply_router_weight_on_input: True if router weights are already
                applied on the input. This is relevant if the implementation
                chooses to do weight application.

        """
        raise NotImplementedError

adjust_N_for_activation(N, activation) staticmethod

Calculate the output dimension for the activation function.

For *_no_mul activations (e.g. relu2_no_mul), there's no gate/up split, so output size equals input size (N).

For regular gated activations (e.g., silu, gelu, swigluoai), output size is N // 2 due to gate × activation(up) multiplication.

Parameters:

• N

  (int) –

  The intermediate size (width of w1/w3 weights).

• activation

  (MoEActivation) –

  The activation function enum.

Returns:

• int –

  The output dimension after activation.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def adjust_N_for_activation(N: int, activation: MoEActivation) -> int:
    """Calculate the output dimension for the activation function.

    For *_no_mul activations (e.g. relu2_no_mul),
    there's no gate/up split, so output size equals input size (N).

    For regular gated activations (e.g., silu, gelu, swigluoai),
    output size is N // 2 due to gate × activation(up) multiplication.

    Args:
        N: The intermediate size (width of w1/w3 weights).
        activation: The activation function enum.

    Returns:
        The output dimension after activation.

    """
    return N if not activation.is_gated else N // 2

apply(output, hidden_states, w1, w2, topk_weights, topk_ids, activation, global_num_experts, expert_map, a1q_scale, a2_scale, workspace13, workspace2, expert_tokens_meta, apply_router_weight_on_input) abstractmethod

This function computes the intermediate result of a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2.

Writes into output and returns None, unless the implementation stopped after GEMM2 and left the top-k reduction to a fused consumer, in which case output is untouched and the unfinalized result is returned.

Parameters:

• output

  (Tensor) –

  (torch.Tensor): The unweighted, unreduced output tensor.

• hidden_states

  (Tensor) –

  (torch.Tensor): The (quantized) input tensor to the MoE layer.

• w1

  (Tensor) –

  The first set of expert weights.

• w2

  (Tensor) –

  The second set of expert weights.

• topk_weights

  (Tensor) –

  A map of row to expert weights. Some implementations choose to do weight application.

• topk_ids

  (Tensor) –

  A map of row to expert id.

• activation

  (str) –

  The activation function to apply after the first MoE layer.

• global_num_experts

  (int) –

  The total number of experts in the global expert space.

• expert_map

  (Optional[Tensor]) –

  A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard.

• a1q_scale

  (Optional[Tensor]) –

  Optional quantized scale to be used for a1. Result of quantization from prepare/finalize and not from the FusedMoEQuantConfig.

• a2_scale

  (Optional[Tensor]) –

  Optional quantized scale to be used for the second gemm's activations.

• workspace13

  (Tensor) –

  A scratch tensor used for gemm outputs must be large enough to hold output of either MoE gemm.

• workspace2

  (Tensor) –

  A scratch tensor used for the activation function.

• expert_tokens_meta

  (Optional[ExpertTokensMetadata]) –

  An optional ExpertTokensMetadata object containing gpu/cpu tensors as big as the number of local experts with the information about the number of tokens assigned to each local expert.

• apply_router_weight_on_input

  (bool) –

  True if router weights are already applied on the input. This is relevant if the implementation chooses to do weight application.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def apply(
    self,
    output: torch.Tensor,
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    activation: MoEActivation,
    global_num_experts: int,
    expert_map: torch.Tensor | None,
    a1q_scale: torch.Tensor | None,
    a2_scale: torch.Tensor | None,
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: ExpertTokensMetadata | None,
    apply_router_weight_on_input: bool,
) -> UnfinalizedMoEOutput | None:
    """This function computes the intermediate result of a Mixture of Experts
    (MoE) layer using two sets of weights, w1 and w2.

    Writes into `output` and returns None, unless the implementation stopped
    after GEMM2 and left the top-k reduction to a fused consumer, in which
    case `output` is untouched and the unfinalized result is returned.

    Args:
        output: (torch.Tensor): The unweighted, unreduced output tensor.
        hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
            layer.
        w1 (torch.Tensor): The first set of expert weights.
        w2 (torch.Tensor): The second set of expert weights.
        topk_weights: A map of row to expert weights. Some implementations
            choose to do weight application.
        topk_ids (torch.Tensor): A map of row to expert id.
        activation (str): The activation function to apply after the first
            MoE layer.
        global_num_experts (int): The total number of experts in the global
            expert space.
        expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
            from the global expert space to the local expert space of the
            expert parallel shard.
        a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
            used for a1. Result of quantization from prepare/finalize and not
            from the FusedMoEQuantConfig.
        a2_scale (Optional[torch.Tensor]): Optional quantized scale to be
            used for the second gemm's activations.
        workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
            must be large enough to hold output of either MoE gemm.
        workspace2 (torch.Tensor): A scratch tensor used for the activation
            function.
        expert_tokens_meta (Optional[ExpertTokensMetadata]): An optional
            ExpertTokensMetadata object containing gpu/cpu tensors
            as big as the number of local experts with the information about
            the number of tokens assigned to each local expert.
        apply_router_weight_on_input: True if router weights are already
            applied on the input. This is relevant if the implementation
            chooses to do weight application.

    """
    raise NotImplementedError

moe_problem_size(a1, w1, w2, topk_ids)

Extract the MoE problem size from the given tensor arguments: - a: The hidden states, input to the MoE layer. - w1: The first set of expert weights. - w2: The second set of expert weights. - topk_ids: The topk ids.

Note: extracting the problem shape from the weight and activation tensors is not obvious. It needs to be done this way specifically due to subtle issues with particular kernels, e.g. the int4 kernels divide the trailing dimension by two, so it's not "correct" to extract N or K from the trailing dimension of w1 or w2. Similarly, some kernels transpose the weights, so this needs to be kept in mind.

Note: This implementation covers most cases. However, if experts require a specialized implementation, like MarlinExperts, they are free to override this function.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def moe_problem_size(
    self,
    a1: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_ids: torch.Tensor,
) -> tuple[int, int, int, int, int]:
    """Extract the MoE problem size from the given tensor arguments:
    - a: The hidden states, input to the MoE layer.
    - w1: The first set of expert weights.
    - w2: The second set of expert weights.
    - topk_ids: The topk ids.

    Note: extracting the problem shape from the weight and activation
    tensors is not obvious.  It needs to be done this way specifically
    due to subtle issues with particular kernels, e.g. the int4 kernels
    divide the trailing dimension by two, so it's not "correct" to
    extract N or K from the trailing dimension of w1 or w2.  Similarly,
    some kernels transpose the weights, so this needs to be kept in mind.

    Note: This implementation covers most cases. However, if experts
    require a specialized implementation, like MarlinExperts, they are free
    to override this function.
    """
    assert len(w1.shape) == 3 and len(w2.shape) == 3
    E, N, _ = w1.shape
    K = a1.size(-1)

    if a1.dim() == 2:
        # Make sure we are using the correct a1 (pre-permute).
        assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
        M = a1.size(0)
    else:
        assert a1.dim() == 3
        assert a1.size(0) == E, f"{a1.size(0)} == {E}"
        M = a1.size(1)  # This is max_num_tokens

    assert topk_ids.dim() == 2
    topk = topk_ids.size(1)

    return E, M, N, K, topk

workspace_dtype(act_dtype)

Workspace type: The dtype to use for the workspace tensors.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
    """Workspace type: The dtype to use for the workspace tensors."""
    return act_dtype

workspace_shapes(M, N, K, topk, global_num_experts, local_num_experts, expert_tokens_meta, activation) abstractmethod

Compute the shapes for the temporary and final outputs of the two gemms and activation in the fused expert function. Since the gemms are independent, the workspace for the first gemm can be shared with the workspace for the last gemm.

Inputs: - M: number of tokens. - N: Row (or column) dimension of expert weights. - K: hidden dimension - topk: The number of top-k experts to select. - global_num_experts: global number of experts. - local_num_experts: local number of experts due to DP/EP. - expert_tokens_meta: number of tokens per expert metadata for batched format.

Returns a tuple of: - workspace13 shape tuple: must be large enough to hold the result of either expert gemm. - workspace2 shape tuple: must be large enough to hold the result of the activation function. - output shape tuple: must be exact size of the final gemm output. - Note: workspace shapes can be 0 if the workspace is not needed. But in order for activation chunking to work, the first dimension of each tuple must be the number of tokens when the shape is not 0.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def workspace_shapes(
    self,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: ExpertTokensMetadata | None,
    activation: MoEActivation,
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
    """Compute the shapes for the temporary and final outputs of the two gemms
    and activation in the fused expert function.  Since the gemms are
    independent, the workspace for the first gemm can be shared with the
    workspace for the last gemm.

    Inputs:
    - M: number of tokens.
    - N: Row (or column) dimension of expert weights.
    - K: hidden dimension
    - topk: The number of top-k experts to select.
    - global_num_experts: global number of experts.
    - local_num_experts: local number of experts due to DP/EP.
    - expert_tokens_meta: number of tokens per expert metadata for batched
                          format.

    Returns a tuple of:
    - workspace13 shape tuple: must be large enough to hold the
      result of either expert gemm.
    - workspace2 shape tuple: must be large enough to hold the
      result of the activation function.
    - output shape tuple: must be exact size of the final gemm output.
    - Note: workspace shapes can be 0 if the workspace is not needed.
      But in order for activation chunking to work, the first dimension
      of each tuple must be the number of tokens when the shape is
      not 0.
    """
    raise NotImplementedError

FusedMoEExpertsMonolithic

Bases: FusedMoEExperts

An abstract base class for the [Permute-Experts-Unpermute] step described above, but with the monolithic interface (accepts router logits rather than topk ids and weights).

Methods:

• apply –

  Same as FusedMoEExperts.apply, except uses router_logits as opposed

• supports_routing_replay_capture –

  Whether this expert supports routing replay capture.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEExpertsMonolithic(FusedMoEExperts):
    """An abstract base class for the [Permute-Experts-Unpermute] step described
    above, but with the monolithic interface (accepts router logits
    rather than topk ids and weights).
    """

    @staticmethod
    def _supports_routing_method(
        routing_method: RoutingMethodType,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        """Whether the kernel supports a routing method (e.g. GroupedTopK).

        Monolithic kernels should explicitly opt-in to support.
        """
        raise NotImplementedError

    @staticmethod
    def _supports_router_logits_dtype(
        router_logits_dtype: torch.dtype | None,
        routing_method: RoutingMethodType,
    ) -> bool:
        """Whether the kernel supports a dtype for router logits.

        Modular kernels should opt-in to support.
        """
        raise NotImplementedError

    @staticmethod
    def is_monolithic() -> bool:
        return True

    def supports_routing_replay_capture(self) -> bool:
        """Whether this expert supports routing replay capture.

        Subclasses backed by a kernel that exposes routed expert IDs
        (e.g. FlashInfer's ``routing_replay_out``) should override.
        """
        return False

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        router_logits: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        # grouped topk + fused topk bias parameters
        num_expert_group: int | None = None,
        e_score_correction_bias: torch.Tensor | None = None,
        routed_scaling_factor: float | None = None,
        topk_group: int | None = None,
        routing_replay_out: torch.Tensor | None = None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed
        to the topk_ids and topk_weights. This is useful for kernels
        with fused router and fused_experts (e.g. FLASHINFER_TRTLLM).

        Kernels that ``supports_routing_replay_capture`` write each token's
        routed expert ids into the leading rows of ``routing_replay_out``.
        """
        raise NotImplementedError

_supports_router_logits_dtype(router_logits_dtype, routing_method) staticmethod

Whether the kernel supports a dtype for router logits.

Modular kernels should opt-in to support.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_router_logits_dtype(
    router_logits_dtype: torch.dtype | None,
    routing_method: RoutingMethodType,
) -> bool:
    """Whether the kernel supports a dtype for router logits.

    Modular kernels should opt-in to support.
    """
    raise NotImplementedError

_supports_routing_method(routing_method, weight_key, activation_key) staticmethod

Whether the kernel supports a routing method (e.g. GroupedTopK).

Monolithic kernels should explicitly opt-in to support.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@staticmethod
def _supports_routing_method(
    routing_method: RoutingMethodType,
    weight_key: QuantKey | None,
    activation_key: QuantKey | None,
) -> bool:
    """Whether the kernel supports a routing method (e.g. GroupedTopK).

    Monolithic kernels should explicitly opt-in to support.
    """
    raise NotImplementedError

apply(hidden_states, w1, w2, router_logits, activation, global_num_experts, expert_map, a1q_scale, apply_router_weight_on_input, num_expert_group=None, e_score_correction_bias=None, routed_scaling_factor=None, topk_group=None, routing_replay_out=None)

Same as FusedMoEExperts.apply, except uses router_logits as opposed to the topk_ids and topk_weights. This is useful for kernels with fused router and fused_experts (e.g. FLASHINFER_TRTLLM).

Kernels that supports_routing_replay_capture write each token's routed expert ids into the leading rows of routing_replay_out.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def apply(
    self,
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    router_logits: torch.Tensor,
    activation: MoEActivation,
    global_num_experts: int,
    expert_map: torch.Tensor | None,
    a1q_scale: torch.Tensor | None,
    apply_router_weight_on_input: bool,
    # grouped topk + fused topk bias parameters
    num_expert_group: int | None = None,
    e_score_correction_bias: torch.Tensor | None = None,
    routed_scaling_factor: float | None = None,
    topk_group: int | None = None,
    routing_replay_out: torch.Tensor | None = None,
) -> torch.Tensor | UnfinalizedMoEOutput:
    """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed
    to the topk_ids and topk_weights. This is useful for kernels
    with fused router and fused_experts (e.g. FLASHINFER_TRTLLM).

    Kernels that ``supports_routing_replay_capture`` write each token's
    routed expert ids into the leading rows of ``routing_replay_out``.
    """
    raise NotImplementedError

supports_routing_replay_capture()

Whether this expert supports routing replay capture.

Subclasses backed by a kernel that exposes routed expert IDs (e.g. FlashInfer's routing_replay_out) should override.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def supports_routing_replay_capture(self) -> bool:
    """Whether this expert supports routing replay capture.

    Subclasses backed by a kernel that exposes routed expert IDs
    (e.g. FlashInfer's ``routing_replay_out``) should override.
    """
    return False

FusedMoEKernel

Methods:

• output_is_reduced –

  Indicates whether or not the output of fused MoE kernel

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@final
class FusedMoEKernel:
    def __init__(
        self,
        prepare_finalize: FusedMoEPrepareAndFinalize,
        fused_experts: FusedMoEExperts,
    ):
        super().__init__()

        # Initialize the implementation (monolithic or modular).
        self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl
        if isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeModular
        ) and isinstance(fused_experts, FusedMoEExpertsModular):
            self.impl = FusedMoEKernelModularImpl(
                prepare_finalize,
                fused_experts,
            )

        elif isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic
        ) and isinstance(fused_experts, FusedMoEExpertsMonolithic):
            self.impl = FusedMoEKernelMonolithicImpl(
                prepare_finalize,
                fused_experts,
            )

        else:
            raise ValueError(
                "prepare_finalize and fused_experts must both be either monolithic "
                f"or non-monolithic but got {prepare_finalize.__class__.__name__} "
                f"and {fused_experts.__class__.__name__}"
            )

        self._post_init_setup()

    @property
    def can_overlap_shared_experts(self) -> bool:
        if isinstance(self.impl, FusedMoEKernelModularImpl):
            return self.impl.prepare_finalize.supports_async()
        else:
            return False

    @property
    def is_monolithic(self) -> bool:
        return isinstance(self.impl, FusedMoEKernelMonolithicImpl)

    @property
    def prepare_finalize(self) -> FusedMoEPrepareAndFinalize:
        return self.impl.prepare_finalize

    @property
    def fused_experts(self) -> FusedMoEExperts:
        return self.impl.fused_experts

    @property
    def moe_config(self) -> FusedMoEConfig:
        return self.fused_experts.moe_config

    def supports_lora(self) -> bool:
        return self.fused_experts.supports_lora()

    def _post_init_setup(self):
        """Resolve any leftover setup dependencies between self.prepare_finalize
        and self.fused_experts here.
        """
        self.prepare_finalize.post_init_setup(self.impl.fused_experts)
        assert (
            self.prepare_finalize.activation_format
            == self.fused_experts.activation_format()
        )

    def output_is_reduced(self) -> bool:
        """Indicates whether or not the output of fused MoE kernel
        is reduced across all ranks.
        """
        return self.prepare_finalize.output_is_reduced()

    def supports_deferred_moe_finalize(self) -> bool:
        return self.prepare_finalize.supports_deferred_moe_finalize()

    def apply_monolithic(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        router_logits: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        # grouped topk + fused topk bias parameters
        num_expert_group: int | None = None,
        e_score_correction_bias: torch.Tensor | None = None,
        routed_scaling_factor: float | None = None,
        topk_group: int | None = None,
        *,
        routing_sink: RoutedExpertsSink | None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        assert isinstance(self.impl, FusedMoEKernelMonolithicImpl)
        return self.impl.apply(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            router_logits=router_logits,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            num_expert_group=num_expert_group,
            e_score_correction_bias=e_score_correction_bias,
            routed_scaling_factor=routed_scaling_factor,
            topk_group=topk_group,
            routing_sink=routing_sink,
        )

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        shared_experts: SharedExperts | None = None,
        shared_experts_input: torch.Tensor | None = None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        assert isinstance(self.impl, FusedMoEKernelModularImpl)
        return self.impl.apply(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            shared_experts=shared_experts,
            shared_experts_input=shared_experts_input,
        )

_post_init_setup()

Resolve any leftover setup dependencies between self.prepare_finalize and self.fused_experts here.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def _post_init_setup(self):
    """Resolve any leftover setup dependencies between self.prepare_finalize
    and self.fused_experts here.
    """
    self.prepare_finalize.post_init_setup(self.impl.fused_experts)
    assert (
        self.prepare_finalize.activation_format
        == self.fused_experts.activation_format()
    )

output_is_reduced()

Indicates whether or not the output of fused MoE kernel is reduced across all ranks.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def output_is_reduced(self) -> bool:
    """Indicates whether or not the output of fused MoE kernel
    is reduced across all ranks.
    """
    return self.prepare_finalize.output_is_reduced()

FusedMoEKernelModularImpl

Methods:

• apply –

  This function computes a Mixture of Experts (MoE) layer using two sets

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@final
class FusedMoEKernelModularImpl:
    def __init__(
        self,
        prepare_finalize: FusedMoEPrepareAndFinalizeModular,
        fused_experts: FusedMoEExpertsModular,
    ):
        self.prepare_finalize = prepare_finalize
        self.fused_experts = fused_experts
        self.shared_experts: SharedExperts | None = None
        moe_parallel_config = fused_experts.moe_config.moe_parallel_config
        self.moe_parallel_config = moe_parallel_config
        self.is_dp_ep = (
            moe_parallel_config is not None
            and moe_parallel_config.dp_size > 1
            and moe_parallel_config.use_ep
        )

    def _allocate_buffers(
        self,
        out_dtype: torch.dtype,
        device: torch.device,
        M_chunk: int,
        M_full: int,
        N: int,
        K: int,
        top_k: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Allocate temporary and output buffers for the fused experts op.
        Inputs:
        - out_dtype: output type of workspace and output tensors.
        - device: the device of the workspace and output tensors.
        See `workspace_shapes` for a description of the remainder of arguments.
        Returns a tuple of (workspace13, workspace2, output) tensors.
        """
        assert M_full > 0 and M_chunk > 0

        workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)

        # Get intermediate workspace shapes based off the chunked M size.
        workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
            M_chunk,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        # Get final output shape based on the full M size.
        _, _, fused_out_shape = self.fused_experts.workspace_shapes(
            M_full,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        # We can reuse the memory between cache1 and cache3 because by the
        # time we need cache3, we're done with cache1.
        # Reuse workspace13 for the output since there is only one chunk.
        max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))

        if current_platform.is_cpu():
            # Every CPU FusedMoEExpertsModular kernel reports zero-sized
            # workspace13/workspace2 (it manages its own scratch space
            # internally) and just needs the output buffer, so skip the
            # workspace manager entirely here -- it's the one part of this
            # call graph Dynamo can't trace (ContextVar-based lane lookup),
            # and CPU never actually needs its cross-chunk memory reuse.
            common_workspace = torch.empty(
                (max_shape_size,), dtype=workspace_dtype, device=device
            )
            workspace2 = torch.empty(
                workspace2_shape, dtype=workspace_dtype, device=device
            )
            workspace13 = _resize_cache(common_workspace, workspace13_shape)
            fused_out = _resize_cache(common_workspace, fused_out_shape)
            return workspace13, workspace2, fused_out

        common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
            ((max_shape_size,), workspace_dtype),
            (workspace2_shape, workspace_dtype),
        )
        workspace13 = _resize_cache(common_workspace, workspace13_shape)
        fused_out = _resize_cache(common_workspace, fused_out_shape)

        return workspace13, workspace2, fused_out

    def _maybe_apply_shared_experts(
        self,
        shared_experts: SharedExperts | None,
        shared_experts_input: torch.Tensor | None,
    ):
        if shared_experts is not None:
            assert self.prepare_finalize.supports_async()
            assert shared_experts_input is not None
            shared_experts(
                shared_experts_input,
                SharedExpertsOrder.MK_INTERNAL_OVERLAPPED,
            )

    def _prepare(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor | None,
        ExpertTokensMetadata | None,
        torch.Tensor,
        torch.Tensor,
    ]:
        """The _prepare method is a wrapper around self.prepare_finalize.prepare
        that handles DBO and async.
        """
        if not self.prepare_finalize.supports_async():
            # We shouldn't be running an a2a kernel that doesn't
            # support async prepare/finalize
            # TODO(lucas): enable in follow-up
            assert not dbo_enabled()

            (
                a1q,
                a1q_scale,
                expert_tokens_meta,
                _expert_topk_ids,
                _expert_topk_weights,
            ) = self.prepare_finalize.prepare(
                hidden_states,
                topk_weights,
                topk_ids,
                global_num_experts,
                expert_map,
                apply_router_weight_on_input,
                self.fused_experts.quant_config,
                defer_input_quant=self.fused_experts.expects_unquantized_inputs,
            )
        else:
            # Overlap shared expert compute with all2all dispatch.
            dbo_maybe_run_recv_hook()
            prepare_ret = self.prepare_finalize.prepare_async(
                hidden_states,
                topk_weights,
                topk_ids,
                global_num_experts,
                expert_map,
                apply_router_weight_on_input,
                self.fused_experts.quant_config,
                defer_input_quant=self.fused_experts.expects_unquantized_inputs,
            )

            # TODO(lucas): refactor this in the alternative schedules followup
            # currently unpack if we have hook + receiver pair or just
            # receiver (see finalize_async docstring)
            hook, receiver = (
                prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret)
            )

            if hook is not None:
                if dbo_enabled():
                    # If DBO is being used, register the hook with the ubatch
                    # context and call it in dbo_maybe_run_recv_hook instead of
                    #  passing it to the receiver.
                    dbo_register_recv_hook(hook)
                    dbo_yield()
                else:
                    hook()

            (
                a1q,
                a1q_scale,
                expert_tokens_meta,
                _expert_topk_ids,
                _expert_topk_weights,
            ) = receiver()

        # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
        topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
        topk_weights = (
            topk_weights if _expert_topk_weights is None else _expert_topk_weights
        )

        return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights

    def _fused_experts(
        self,
        in_dtype: torch.dtype,
        a1q: torch.Tensor,
        a1q_scale: torch.Tensor | None,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        local_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        expert_tokens_meta: ExpertTokensMetadata | None,
        output_alias: torch.Tensor | None = None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
            a1q, w1, w2, topk_ids
        )

        # This happens when none of the tokens from the all2all reach this
        # EP rank. Also, note that this is only relevant for CUDAGraph
        # incompatible all2all kernels like the DeepEP high-throughput
        # kernels. CUDAGraph compatible all2all kernels like the DeepEP
        # low-latency kernels are always batched and can never run into
        # the tensor.numel() == 0 case.
        if M_full == 0:
            return torch.empty_like(a1q, dtype=in_dtype)

        workspace13, workspace2, fused_out = self._allocate_buffers(
            in_dtype,
            a1q.device,
            M_full,
            M_full,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        use_output_alias = (
            output_alias is not None
            and output_alias.shape == fused_out.shape
            and output_alias.dtype == fused_out.dtype
            and output_alias.device == fused_out.device
            and output_alias.is_contiguous()
        )

        # If caller's output buffer already matches fused_out shape/dtype, alias
        # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
        # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
        # double-copy MoE write-back path).
        if current_platform.is_rocm():
            from vllm._aiter_ops import rocm_aiter_ops

            if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled():
                fused_out = output_alias
        elif use_output_alias:
            fused_out = output_alias

        unfinalized = self.fused_experts.apply(
            output=fused_out,
            hidden_states=a1q,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            a1q_scale=a1q_scale,
            a2_scale=self.fused_experts.a2_scale,
            workspace13=workspace13,
            workspace2=workspace2,
            expert_tokens_meta=expert_tokens_meta,
            apply_router_weight_on_input=apply_router_weight_on_input,
        )

        # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the
        # top-k reduction they left open belongs to whoever takes this.
        if isinstance(unfinalized, UnfinalizedMoEOutput):
            return unfinalized

        return fused_out

    def _finalize(
        self,
        output: torch.Tensor,
        fused_out: torch.Tensor,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        shared_experts: SharedExperts | None,
        shared_experts_input: torch.Tensor | None,
    ) -> torch.Tensor:
        """The _finalize method is a wrapper around self.prepare_finalize.finalize
        that handles DBO, async and shared expert overlap.

        Args:
            output: Tensor the finalized result is written into.
            fused_out: The unweighted, unreduced output of the fused experts.
            hidden_states: The input tensor to the MoE layer.
            topk_weights: A map of row to expert weights.
            topk_ids: A map of row to expert id.
            apply_router_weight_on_input: True if the router weights were
                already applied on the input.
            shared_experts: SharedExperts | None. The shared experts if any.
            shared_experts_input: Optional separate input for shared experts.
                When latent MoE is used, hidden_states is the latent-projected
                tensor (smaller dimension) used by routed experts, while
                shared_experts_input is the original hidden_states (full
                dimension) needed by the shared expert MLP.

        """
        if not self.prepare_finalize.supports_async():
            assert not dbo_enabled()

            self.prepare_finalize.finalize(
                output,
                fused_out,
                topk_weights,
                topk_ids,
                apply_router_weight_on_input,
                self.fused_experts.finalize_weight_and_reduce_impl(),
            )
        else:
            finalize_ret = self.prepare_finalize.finalize_async(
                output,
                fused_out,
                topk_weights,
                topk_ids,
                apply_router_weight_on_input,
                self.fused_experts.finalize_weight_and_reduce_impl(),
            )
            self._maybe_apply_shared_experts(shared_experts, shared_experts_input)

            # TODO(lucas): refactor this in the alternative schedules followup
            # currently unpack if we have hook + receiver pair or just
            # receiver (see finalize_async docstring)
            hook, receiver = (
                finalize_ret
                if isinstance(finalize_ret, tuple)
                else (None, finalize_ret)
            )

            if hook is not None:
                if dbo_enabled():
                    # If DBO is being used, register the hook with the ubatch
                    # context and call it in dbo_maybe_run_recv_hook instead of
                    #  passing it to the receiver.
                    dbo_register_recv_hook(hook)
                    dbo_yield()
                else:
                    hook()

            receiver()

        return output

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
        activation: MoEActivation = MoEActivation.SILU,
        global_num_experts: int = -1,
        expert_map: torch.Tensor | None = None,
        apply_router_weight_on_input: bool = False,
        shared_experts: SharedExperts | None = None,
        shared_experts_input: torch.Tensor | None = None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        """This function computes a Mixture of Experts (MoE) layer using two sets
        of weights, w1 and w2, and top-k gating mechanism.

        Args:
            hidden_states: (torch.Tensor): The input tensor to the MoE layer.
            w1 (torch.Tensor): The first set of expert weights.
            w2 (torch.Tensor): The second set of expert weights.
            topk_ids (torch.Tensor): A map of row to expert id.
            topk_weights (torch.Tensor): The topk weights applied at the end of
                the layer.
            activation (MoEActivation): The activation function to apply after
                the first MoE layer.
            global_num_experts (int): The total number of experts in the global
                expert space.
            expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
                from the global expert space to the local expert space of the
                expert parallel shard.
            apply_router_weight_on_input (bool): When true, the topk weights are
                applied directly on the inputs. This is only applicable when
                topk is 1.
            shared_experts: SharedExperts | None. The shared experts if any.
            shared_experts_input (Optional[torch.Tensor]): Optional separate
                input for shared experts. For latent MoE, this is the original
                hidden_states before latent projection.

        Returns:
            torch.Tensor: The output tensor after applying the MoE layer, or
            the unfinalized output when the experts left the top-k reduction to
            a fused consumer.

        """
        output = torch.empty_like(hidden_states)

        local_num_experts = w1.shape[0]
        if global_num_experts == -1:
            global_num_experts = local_num_experts

        a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare(
            hidden_states,
            topk_weights,
            topk_ids,
            global_num_experts,
            expert_map,
            apply_router_weight_on_input,
        )

        # Stash the original unquantized hidden states on the LoRA context
        # so apply_w13_lora sees correct-magnitude activations instead of
        # the potentially quantized values produced by _prepare().
        lora_ctx = getattr(self.fused_experts, "_lora_context", None)
        if lora_ctx is not None:
            lora_ctx.original_hidden_states = hidden_states

        fused_out = self._fused_experts(
            in_dtype=hidden_states.dtype,
            a1q=a1q,
            a1q_scale=a1q_scale,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            local_num_experts=local_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            expert_tokens_meta=expert_tokens_meta,
            output_alias=output,
        )

        if lora_ctx is not None:
            lora_ctx.original_hidden_states = None

        if isinstance(fused_out, UnfinalizedMoEOutput):
            # Nothing below can run on an unfinalized output: finalize is the
            # local top-k reduction (and, for DP/EP, the combine) that the
            # consumer takes over.
            if not self.prepare_finalize.supports_deferred_moe_finalize():
                raise RuntimeError(
                    f"{type(self.prepare_finalize).__name__} cannot pass through "
                    "a deferred MoE output."
                )
            return fused_out

        return self._finalize(
            output,
            fused_out,
            hidden_states,
            topk_weights,
            topk_ids,
            apply_router_weight_on_input,
            shared_experts=shared_experts,
            shared_experts_input=shared_experts_input,
        )

_allocate_buffers(out_dtype, device, M_chunk, M_full, N, K, top_k, global_num_experts, local_num_experts, expert_tokens_meta, activation)

Allocate temporary and output buffers for the fused experts op. Inputs: - out_dtype: output type of workspace and output tensors. - device: the device of the workspace and output tensors. See workspace_shapes for a description of the remainder of arguments. Returns a tuple of (workspace13, workspace2, output) tensors.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def _allocate_buffers(
    self,
    out_dtype: torch.dtype,
    device: torch.device,
    M_chunk: int,
    M_full: int,
    N: int,
    K: int,
    top_k: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: ExpertTokensMetadata | None,
    activation: MoEActivation,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Allocate temporary and output buffers for the fused experts op.
    Inputs:
    - out_dtype: output type of workspace and output tensors.
    - device: the device of the workspace and output tensors.
    See `workspace_shapes` for a description of the remainder of arguments.
    Returns a tuple of (workspace13, workspace2, output) tensors.
    """
    assert M_full > 0 and M_chunk > 0

    workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)

    # Get intermediate workspace shapes based off the chunked M size.
    workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
        M_chunk,
        N,
        K,
        top_k,
        global_num_experts,
        local_num_experts,
        expert_tokens_meta,
        activation,
    )

    # Get final output shape based on the full M size.
    _, _, fused_out_shape = self.fused_experts.workspace_shapes(
        M_full,
        N,
        K,
        top_k,
        global_num_experts,
        local_num_experts,
        expert_tokens_meta,
        activation,
    )

    # We can reuse the memory between cache1 and cache3 because by the
    # time we need cache3, we're done with cache1.
    # Reuse workspace13 for the output since there is only one chunk.
    max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))

    if current_platform.is_cpu():
        # Every CPU FusedMoEExpertsModular kernel reports zero-sized
        # workspace13/workspace2 (it manages its own scratch space
        # internally) and just needs the output buffer, so skip the
        # workspace manager entirely here -- it's the one part of this
        # call graph Dynamo can't trace (ContextVar-based lane lookup),
        # and CPU never actually needs its cross-chunk memory reuse.
        common_workspace = torch.empty(
            (max_shape_size,), dtype=workspace_dtype, device=device
        )
        workspace2 = torch.empty(
            workspace2_shape, dtype=workspace_dtype, device=device
        )
        workspace13 = _resize_cache(common_workspace, workspace13_shape)
        fused_out = _resize_cache(common_workspace, fused_out_shape)
        return workspace13, workspace2, fused_out

    common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
        ((max_shape_size,), workspace_dtype),
        (workspace2_shape, workspace_dtype),
    )
    workspace13 = _resize_cache(common_workspace, workspace13_shape)
    fused_out = _resize_cache(common_workspace, fused_out_shape)

    return workspace13, workspace2, fused_out

_finalize(output, fused_out, hidden_states, topk_weights, topk_ids, apply_router_weight_on_input, shared_experts, shared_experts_input)

The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap.

Parameters:

• output

  (Tensor) –

  Tensor the finalized result is written into.

• fused_out

  (Tensor) –

  The unweighted, unreduced output of the fused experts.

• hidden_states

  (Tensor) –

  The input tensor to the MoE layer.

• topk_weights

  (Tensor) –

  A map of row to expert weights.

• topk_ids

  (Tensor) –

  A map of row to expert id.

• apply_router_weight_on_input

  (bool) –

  True if the router weights were already applied on the input.

• shared_experts

  (SharedExperts | None) –

  SharedExperts | None. The shared experts if any.

• shared_experts_input

  (Tensor | None) –

  Optional separate input for shared experts. When latent MoE is used, hidden_states is the latent-projected tensor (smaller dimension) used by routed experts, while shared_experts_input is the original hidden_states (full dimension) needed by the shared expert MLP.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def _finalize(
    self,
    output: torch.Tensor,
    fused_out: torch.Tensor,
    hidden_states: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool,
    shared_experts: SharedExperts | None,
    shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
    """The _finalize method is a wrapper around self.prepare_finalize.finalize
    that handles DBO, async and shared expert overlap.

    Args:
        output: Tensor the finalized result is written into.
        fused_out: The unweighted, unreduced output of the fused experts.
        hidden_states: The input tensor to the MoE layer.
        topk_weights: A map of row to expert weights.
        topk_ids: A map of row to expert id.
        apply_router_weight_on_input: True if the router weights were
            already applied on the input.
        shared_experts: SharedExperts | None. The shared experts if any.
        shared_experts_input: Optional separate input for shared experts.
            When latent MoE is used, hidden_states is the latent-projected
            tensor (smaller dimension) used by routed experts, while
            shared_experts_input is the original hidden_states (full
            dimension) needed by the shared expert MLP.

    """
    if not self.prepare_finalize.supports_async():
        assert not dbo_enabled()

        self.prepare_finalize.finalize(
            output,
            fused_out,
            topk_weights,
            topk_ids,
            apply_router_weight_on_input,
            self.fused_experts.finalize_weight_and_reduce_impl(),
        )
    else:
        finalize_ret = self.prepare_finalize.finalize_async(
            output,
            fused_out,
            topk_weights,
            topk_ids,
            apply_router_weight_on_input,
            self.fused_experts.finalize_weight_and_reduce_impl(),
        )
        self._maybe_apply_shared_experts(shared_experts, shared_experts_input)

        # TODO(lucas): refactor this in the alternative schedules followup
        # currently unpack if we have hook + receiver pair or just
        # receiver (see finalize_async docstring)
        hook, receiver = (
            finalize_ret
            if isinstance(finalize_ret, tuple)
            else (None, finalize_ret)
        )

        if hook is not None:
            if dbo_enabled():
                # If DBO is being used, register the hook with the ubatch
                # context and call it in dbo_maybe_run_recv_hook instead of
                #  passing it to the receiver.
                dbo_register_recv_hook(hook)
                dbo_yield()
            else:
                hook()

        receiver()

    return output

_prepare(hidden_states, topk_weights, topk_ids, global_num_experts, expert_map, apply_router_weight_on_input)

The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def _prepare(
    self,
    hidden_states: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    global_num_experts: int,
    expert_map: torch.Tensor | None,
    apply_router_weight_on_input: bool,
) -> tuple[
    torch.Tensor,
    torch.Tensor | None,
    ExpertTokensMetadata | None,
    torch.Tensor,
    torch.Tensor,
]:
    """The _prepare method is a wrapper around self.prepare_finalize.prepare
    that handles DBO and async.
    """
    if not self.prepare_finalize.supports_async():
        # We shouldn't be running an a2a kernel that doesn't
        # support async prepare/finalize
        # TODO(lucas): enable in follow-up
        assert not dbo_enabled()

        (
            a1q,
            a1q_scale,
            expert_tokens_meta,
            _expert_topk_ids,
            _expert_topk_weights,
        ) = self.prepare_finalize.prepare(
            hidden_states,
            topk_weights,
            topk_ids,
            global_num_experts,
            expert_map,
            apply_router_weight_on_input,
            self.fused_experts.quant_config,
            defer_input_quant=self.fused_experts.expects_unquantized_inputs,
        )
    else:
        # Overlap shared expert compute with all2all dispatch.
        dbo_maybe_run_recv_hook()
        prepare_ret = self.prepare_finalize.prepare_async(
            hidden_states,
            topk_weights,
            topk_ids,
            global_num_experts,
            expert_map,
            apply_router_weight_on_input,
            self.fused_experts.quant_config,
            defer_input_quant=self.fused_experts.expects_unquantized_inputs,
        )

        # TODO(lucas): refactor this in the alternative schedules followup
        # currently unpack if we have hook + receiver pair or just
        # receiver (see finalize_async docstring)
        hook, receiver = (
            prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret)
        )

        if hook is not None:
            if dbo_enabled():
                # If DBO is being used, register the hook with the ubatch
                # context and call it in dbo_maybe_run_recv_hook instead of
                #  passing it to the receiver.
                dbo_register_recv_hook(hook)
                dbo_yield()
            else:
                hook()

        (
            a1q,
            a1q_scale,
            expert_tokens_meta,
            _expert_topk_ids,
            _expert_topk_weights,
        ) = receiver()

    # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
    topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
    topk_weights = (
        topk_weights if _expert_topk_weights is None else _expert_topk_weights
    )

    return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights

apply(hidden_states, w1, w2, topk_ids, topk_weights, activation=MoEActivation.SILU, global_num_experts=-1, expert_map=None, apply_router_weight_on_input=False, shared_experts=None, shared_experts_input=None)

This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism.

Parameters:

• hidden_states

  (Tensor) –

  (torch.Tensor): The input tensor to the MoE layer.

• w1

  (Tensor) –

  The first set of expert weights.

• w2

  (Tensor) –

  The second set of expert weights.

• topk_ids

  (Tensor) –

  A map of row to expert id.

• topk_weights

  (Tensor) –

  The topk weights applied at the end of the layer.

• activation

  (MoEActivation, default: SILU ) –

  The activation function to apply after the first MoE layer.

• global_num_experts

  (int, default: -1 ) –

  The total number of experts in the global expert space.

• expert_map

  (Optional[Tensor], default: None ) –

  A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard.

• apply_router_weight_on_input

  (bool, default: False ) –

  When true, the topk weights are applied directly on the inputs. This is only applicable when topk is 1.

• shared_experts

  (SharedExperts | None, default: None ) –

  SharedExperts | None. The shared experts if any.

• shared_experts_input

  (Optional[Tensor], default: None ) –

  Optional separate input for shared experts. For latent MoE, this is the original hidden_states before latent projection.

Returns:

• Tensor | UnfinalizedMoEOutput –

  torch.Tensor: The output tensor after applying the MoE layer, or

• Tensor | UnfinalizedMoEOutput –

  the unfinalized output when the experts left the top-k reduction to

• Tensor | UnfinalizedMoEOutput –

  a fused consumer.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def apply(
    self,
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    activation: MoEActivation = MoEActivation.SILU,
    global_num_experts: int = -1,
    expert_map: torch.Tensor | None = None,
    apply_router_weight_on_input: bool = False,
    shared_experts: SharedExperts | None = None,
    shared_experts_input: torch.Tensor | None = None,
) -> torch.Tensor | UnfinalizedMoEOutput:
    """This function computes a Mixture of Experts (MoE) layer using two sets
    of weights, w1 and w2, and top-k gating mechanism.

    Args:
        hidden_states: (torch.Tensor): The input tensor to the MoE layer.
        w1 (torch.Tensor): The first set of expert weights.
        w2 (torch.Tensor): The second set of expert weights.
        topk_ids (torch.Tensor): A map of row to expert id.
        topk_weights (torch.Tensor): The topk weights applied at the end of
            the layer.
        activation (MoEActivation): The activation function to apply after
            the first MoE layer.
        global_num_experts (int): The total number of experts in the global
            expert space.
        expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
            from the global expert space to the local expert space of the
            expert parallel shard.
        apply_router_weight_on_input (bool): When true, the topk weights are
            applied directly on the inputs. This is only applicable when
            topk is 1.
        shared_experts: SharedExperts | None. The shared experts if any.
        shared_experts_input (Optional[torch.Tensor]): Optional separate
            input for shared experts. For latent MoE, this is the original
            hidden_states before latent projection.

    Returns:
        torch.Tensor: The output tensor after applying the MoE layer, or
        the unfinalized output when the experts left the top-k reduction to
        a fused consumer.

    """
    output = torch.empty_like(hidden_states)

    local_num_experts = w1.shape[0]
    if global_num_experts == -1:
        global_num_experts = local_num_experts

    a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare(
        hidden_states,
        topk_weights,
        topk_ids,
        global_num_experts,
        expert_map,
        apply_router_weight_on_input,
    )

    # Stash the original unquantized hidden states on the LoRA context
    # so apply_w13_lora sees correct-magnitude activations instead of
    # the potentially quantized values produced by _prepare().
    lora_ctx = getattr(self.fused_experts, "_lora_context", None)
    if lora_ctx is not None:
        lora_ctx.original_hidden_states = hidden_states

    fused_out = self._fused_experts(
        in_dtype=hidden_states.dtype,
        a1q=a1q,
        a1q_scale=a1q_scale,
        w1=w1,
        w2=w2,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        activation=activation,
        global_num_experts=global_num_experts,
        local_num_experts=local_num_experts,
        expert_map=expert_map,
        apply_router_weight_on_input=apply_router_weight_on_input,
        expert_tokens_meta=expert_tokens_meta,
        output_alias=output,
    )

    if lora_ctx is not None:
        lora_ctx.original_hidden_states = None

    if isinstance(fused_out, UnfinalizedMoEOutput):
        # Nothing below can run on an unfinalized output: finalize is the
        # local top-k reduction (and, for DP/EP, the combine) that the
        # consumer takes over.
        if not self.prepare_finalize.supports_deferred_moe_finalize():
            raise RuntimeError(
                f"{type(self.prepare_finalize).__name__} cannot pass through "
                "a deferred MoE output."
            )
        return fused_out

    return self._finalize(
        output,
        fused_out,
        hidden_states,
        topk_weights,
        topk_ids,
        apply_router_weight_on_input,
        shared_experts=shared_experts,
        shared_experts_input=shared_experts_input,
    )

FusedMoEKernelMonolithicImpl

Methods:

• apply –

  Same as forward(), except uses router_logits as opposed

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@final
class FusedMoEKernelMonolithicImpl:
    def __init__(
        self,
        prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic,
        fused_experts: FusedMoEExpertsMonolithic,
    ):
        self.prepare_finalize = prepare_finalize
        self.fused_experts = fused_experts

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        router_logits: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        # grouped topk + fused topk bias parameters
        num_expert_group: int | None = None,
        e_score_correction_bias: torch.Tensor | None = None,
        routed_scaling_factor: float | None = None,
        topk_group: int | None = None,
        *,
        routing_sink: RoutedExpertsSink | None,
    ) -> torch.Tensor | UnfinalizedMoEOutput:
        """Same as forward(), except uses router_logits as opposed
        to the topk_ids and topk_weights. This is used for kernels
        that have fused router + experts (e.g. FLASHINFER_TRTLLM).
        """
        a1q, a1q_scale, router_logits = self.prepare_finalize.prepare(
            hidden_states,
            router_logits=router_logits,
            quant_config=self.fused_experts.quant_config,
            defer_input_quant=self.fused_experts.expects_unquantized_inputs,
        )

        routing_replay_out = None
        if routing_sink is not None:
            # Checked per call: a weight reload may rebuild a different kernel.
            if not self.fused_experts.supports_routing_replay_capture():
                raise ValueError(
                    "Routed-experts capture is not supported with monolithic MoE "
                    f"kernel {type(self.fused_experts).__name__}."
                )
            routing_replay_out = routing_sink.buffer[: len(a1q)]
        fused_out = self.fused_experts.apply(
            hidden_states=a1q,
            w1=w1,
            w2=w2,
            router_logits=router_logits,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            a1q_scale=a1q_scale,
            # grouped topk + fused topk bias parameters
            num_expert_group=num_expert_group,
            e_score_correction_bias=e_score_correction_bias,
            routed_scaling_factor=routed_scaling_factor,
            topk_group=topk_group,
            routing_replay_out=routing_replay_out,
        )
        if routing_sink is not None:
            routing_sink.capture_fn(routing_replay_out)

        if isinstance(fused_out, UnfinalizedMoEOutput):
            if not self.prepare_finalize.supports_deferred_moe_finalize():
                raise RuntimeError(
                    f"{type(self.prepare_finalize).__name__} cannot pass through "
                    "a deferred MoE output."
                )
            return fused_out
        return self.prepare_finalize.finalize(fused_out)

apply(hidden_states, w1, w2, router_logits, activation, global_num_experts, expert_map, apply_router_weight_on_input, num_expert_group=None, e_score_correction_bias=None, routed_scaling_factor=None, topk_group=None, *, routing_sink)

Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM).

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def apply(
    self,
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    router_logits: torch.Tensor,
    activation: MoEActivation,
    global_num_experts: int,
    expert_map: torch.Tensor | None,
    apply_router_weight_on_input: bool,
    # grouped topk + fused topk bias parameters
    num_expert_group: int | None = None,
    e_score_correction_bias: torch.Tensor | None = None,
    routed_scaling_factor: float | None = None,
    topk_group: int | None = None,
    *,
    routing_sink: RoutedExpertsSink | None,
) -> torch.Tensor | UnfinalizedMoEOutput:
    """Same as forward(), except uses router_logits as opposed
    to the topk_ids and topk_weights. This is used for kernels
    that have fused router + experts (e.g. FLASHINFER_TRTLLM).
    """
    a1q, a1q_scale, router_logits = self.prepare_finalize.prepare(
        hidden_states,
        router_logits=router_logits,
        quant_config=self.fused_experts.quant_config,
        defer_input_quant=self.fused_experts.expects_unquantized_inputs,
    )

    routing_replay_out = None
    if routing_sink is not None:
        # Checked per call: a weight reload may rebuild a different kernel.
        if not self.fused_experts.supports_routing_replay_capture():
            raise ValueError(
                "Routed-experts capture is not supported with monolithic MoE "
                f"kernel {type(self.fused_experts).__name__}."
            )
        routing_replay_out = routing_sink.buffer[: len(a1q)]
    fused_out = self.fused_experts.apply(
        hidden_states=a1q,
        w1=w1,
        w2=w2,
        router_logits=router_logits,
        activation=activation,
        global_num_experts=global_num_experts,
        expert_map=expert_map,
        apply_router_weight_on_input=apply_router_weight_on_input,
        a1q_scale=a1q_scale,
        # grouped topk + fused topk bias parameters
        num_expert_group=num_expert_group,
        e_score_correction_bias=e_score_correction_bias,
        routed_scaling_factor=routed_scaling_factor,
        topk_group=topk_group,
        routing_replay_out=routing_replay_out,
    )
    if routing_sink is not None:
        routing_sink.capture_fn(routing_replay_out)

    if isinstance(fused_out, UnfinalizedMoEOutput):
        if not self.prepare_finalize.supports_deferred_moe_finalize():
            raise RuntimeError(
                f"{type(self.prepare_finalize).__name__} cannot pass through "
                "a deferred MoE output."
            )
        return fused_out
    return self.prepare_finalize.finalize(fused_out)

FusedMoEPrepareAndFinalize

Bases: ABC

An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above.

There are two variants of this class: * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits

Methods:

• max_num_tokens_per_rank –

  Some PrepareFinalize All2All implementations are batched. Meaning,

• output_is_reduced –

  Indicates whether or not the output of finalize is reduced across all

• post_init_setup –

  Initialize FusedMoEPrepareAndFinalizeModular settings that depend on

• supports_async –

  Indicates whether or not this class implements prepare_async and

• supports_deferred_moe_finalize –

  Whether finalize can be skipped for a deferring consumer.

• topk_indices_dtype –

  The PrepareFinalize All2All implementations generally constrain the

Attributes:

• activation_format (FusedMoEActivationFormat) –

  A property indicating the output format of the activations for the

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEPrepareAndFinalize(ABC):
    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
    described above.

    There are two variants of this class:
    * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights
    * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits
    """

    def post_init_setup(self, fused_experts: "FusedMoEExperts"):
        """Initialize FusedMoEPrepareAndFinalizeModular settings that depend on
        FusedMoEExpertsModular experts object.
        The FusedMoEPrepareAndFinalizeModular implementations that have such
        dependencies may choose to override this function.
        """
        return

    @property
    @abstractmethod
    def activation_format(self) -> FusedMoEActivationFormat:
        """A property indicating the output format of the activations for the
        'prepare' method.
        """
        raise NotImplementedError

    @abstractmethod
    def topk_indices_dtype(self) -> torch.dtype | None:
        """The PrepareFinalize All2All implementations generally constrain the
        dtype of the topk_ids they support. This function returns the
        required topk indices dtype so it can be respected.
        Return None if there are no such restrictions.
        """
        raise NotImplementedError

    @abstractmethod
    def max_num_tokens_per_rank(self) -> int | None:
        """Some PrepareFinalize All2All implementations are batched. Meaning,
        they can process only as set of tokens at a time. This
        function returns the batch size i.e the maximum number of tokens
        the implementation can process at a time.
        Return None if there are no such restrictions.
        """
        raise NotImplementedError

    @abstractmethod
    def num_dispatchers(self) -> int:
        raise NotImplementedError

    @abstractmethod
    def output_is_reduced(self) -> bool:
        """Indicates whether or not the output of finalize is reduced across all
        ranks.
        """
        raise NotImplementedError

    def supports_async(self) -> bool:
        """Indicates whether or not this class implements prepare_async and
        finalize_async.
        """
        return False

    def supports_deferred_moe_finalize(self) -> bool:
        """Whether ``finalize`` can be skipped for a deferring consumer.

        An implementation opts in only if everything it does in ``finalize``
        -- the top-k reduction, and any combine or reduce-scatter -- is work
        the consumer takes over.
        """
        return False

activation_format abstractmethod property

A property indicating the output format of the activations for the 'prepare' method.

max_num_tokens_per_rank() abstractmethod

Some PrepareFinalize All2All implementations are batched. Meaning, they can process only as set of tokens at a time. This function returns the batch size i.e the maximum number of tokens the implementation can process at a time. Return None if there are no such restrictions.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def max_num_tokens_per_rank(self) -> int | None:
    """Some PrepareFinalize All2All implementations are batched. Meaning,
    they can process only as set of tokens at a time. This
    function returns the batch size i.e the maximum number of tokens
    the implementation can process at a time.
    Return None if there are no such restrictions.
    """
    raise NotImplementedError

output_is_reduced() abstractmethod

Indicates whether or not the output of finalize is reduced across all ranks.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def output_is_reduced(self) -> bool:
    """Indicates whether or not the output of finalize is reduced across all
    ranks.
    """
    raise NotImplementedError

post_init_setup(fused_experts)

Initialize FusedMoEPrepareAndFinalizeModular settings that depend on FusedMoEExpertsModular experts object. The FusedMoEPrepareAndFinalizeModular implementations that have such dependencies may choose to override this function.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def post_init_setup(self, fused_experts: "FusedMoEExperts"):
    """Initialize FusedMoEPrepareAndFinalizeModular settings that depend on
    FusedMoEExpertsModular experts object.
    The FusedMoEPrepareAndFinalizeModular implementations that have such
    dependencies may choose to override this function.
    """
    return

supports_async()

Indicates whether or not this class implements prepare_async and finalize_async.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def supports_async(self) -> bool:
    """Indicates whether or not this class implements prepare_async and
    finalize_async.
    """
    return False

supports_deferred_moe_finalize()

Whether finalize can be skipped for a deferring consumer.

An implementation opts in only if everything it does in finalize -- the top-k reduction, and any combine or reduce-scatter -- is work the consumer takes over.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def supports_deferred_moe_finalize(self) -> bool:
    """Whether ``finalize`` can be skipped for a deferring consumer.

    An implementation opts in only if everything it does in ``finalize``
    -- the top-k reduction, and any combine or reduce-scatter -- is work
    the consumer takes over.
    """
    return False

topk_indices_dtype() abstractmethod

The PrepareFinalize All2All implementations generally constrain the dtype of the topk_ids they support. This function returns the required topk indices dtype so it can be respected. Return None if there are no such restrictions.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def topk_indices_dtype(self) -> torch.dtype | None:
    """The PrepareFinalize All2All implementations generally constrain the
    dtype of the topk_ids they support. This function returns the
    required topk indices dtype so it can be respected.
    Return None if there are no such restrictions.
    """
    raise NotImplementedError

FusedMoEPrepareAndFinalizeModular

Bases: FusedMoEPrepareAndFinalize

An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above for the Modular case.

Methods:

• finalize –

  Perform any combine plus apply weights and perform a reduction on the

• finalize_async –

  Perform any combine plus apply weights and perform a reduction on the

• prepare –

  Perform any quantization (and/or) dispatching needed for this kernel.

• prepare_async –

  Perform any quantization (and/or) dispatching needed for this kernel

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
    described above for the Modular case.
    """

    @abstractmethod
    def prepare(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
        defer_input_quant: bool,
    ) -> PrepareResultType:
        """Perform any quantization (and/or) dispatching needed for this kernel.
        - a1: The (unquantized) input to the MoE layer.
        - topk_ids: The topk ids.
        - topk_weights: The topk weights.
        - num_experts: The total number of experts in the global expert space.
        - expert_map: A tensor mapping expert indices from the global expert
          space to the local expert space of the expert parallel shard.
        - apply_router_weight_on_input: When True, apply the weights to the
          activations, before quantization + dispatching.
        - quant_config: Quantization info provided by the fused experts.
        - defer_input_quant: Runtime parameter indicating whether or not to
          defer input quantization to the FusedMoEExpertsModular
          in cases where the compute kernel expects unquantized inputs

        Returns a tuple of:
        - quantized + dispatched a.
        - Optional quantized + dispatched a1_scales.
        - Optional ExpertTokensMetadata containing gpu/cpu tensors
          as big as the number of local experts with the information about the
          number of tokens assigned to each local expert.
        - Optional dispatched expert topk IDs
        - Optional dispatched expert topk weight
        """
        raise NotImplementedError

    def prepare_async(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
        defer_input_quant: bool,
    ) -> tuple[Callable, ReceiverType] | ReceiverType:
        """Perform any quantization (and/or) dispatching needed for this kernel
        but do not wait for results from other workers.
        - a1: The (unquantized) input to the MoE layer.
        - a1_scale: Optional scales for a1
        - a2_scale: Optional scales for the second MoE gemm.  Required to make
          sure the quantization is consistent for both gemms.
        - topk_ids: The topk ids.
        - topk_weights: The topk weights.
        - num_experts: The total number of experts in the global expert space.
        - expert_map: A tensor mapping expert indices from the global expert
          space to the local expert space of the expert parallel shard.
        - apply_router_weight_on_input: When True, apply the weights to the
          activations, before quantization + dispatching.
        - defer_input_quant: Runtime parameter indicating whether or not to
          defer input quantization to the FusedMoEExpertsModular
          in cases where the compute kernel expects unquantized inputs

        Returns a callback or a hook callback pair that when invoked waits for
        results from other workers and has the same return signature as
        `prepare`, if a hook is returned this is more lightweight check that
        the recv is complete without doing extra work (used by DBO, will be
        refactored in the very near future)

        e.g.

        ret = obj.prepare_async(...)

        if isinstance(ret, tuple):
            hook, receiver = ret
            hook()

        if hook is not None:
        a, a_scales, expert_meta, topk_ids, topk_weights = receiver()

        is equivalent to:

        a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...)
        """
        raise NotImplementedError

    @abstractmethod
    def finalize(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        weight_and_reduce_impl: TopKWeightAndReduce,
    ) -> None:
        """Perform any combine plus apply weights and perform a reduction on the
        fused experts output.
        - output: The output tensor, written in place.  Must be (M, K) shape.
        - fused_expert_output: The unweighted, unreduced output of the fused
          experts, it will have (M, topk, K) shape.
        - topk_weights: The weights to be applied to the fused_experts_output.
        - topk_ids: The topk_ids.
        - apply_router_weight_on_input: When False, apply the weights to
          fused_expert_output.
        - weight_and_reduce_impl: An optional TopKWeightAndReduce
          implementation.
        """
        raise NotImplementedError

    def finalize_async(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        weight_and_reduce_impl: TopKWeightAndReduce,
    ) -> tuple[Callable, Callable] | Callable:
        """Perform any combine plus apply weights and perform a reduction on the
        fused experts output but do not wait for results from other workers.
        - output: The output tensor, written in place.  Must be (M, K) shape.
        - fused_expert_output: The unweighted, unreduced output of the fused
          experts, it will have (M, topk, K) shape.
        - topk_weights: The weights to be applied to the fused_experts_output.
        - topk_ids: The topk_ids.
        - apply_router_weight_on_input: When False, apply the weights to
          fused_expert_output.
        - weight_and_reduce_impl: An optional TopKWeightAndReduce
          implementation.

        Returns a callback or a hook callback pair that when invoked waits for
        results from other workers and has the same return signature as
        `finalize`, if a hook is returned this is more lightweight check that
        the recv is complete without doing extra work (used by DBO, will be
        refactored in the very near future)

        ret = obj.finalize_async(output, ...)
        ... output not valid yet ...
        if isinstance(ret, tuple):
            hook, receiver = ret
            hook()
        receiver()
        ... output valid here ...

        is equivalent to:

        obj.finalize(output, ...)
        """
        raise NotImplementedError

finalize(output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl) abstractmethod

Perform any combine plus apply weights and perform a reduction on the fused experts output. - output: The output tensor, written in place. Must be (M, K) shape. - fused_expert_output: The unweighted, unreduced output of the fused experts, it will have (M, topk, K) shape. - topk_weights: The weights to be applied to the fused_experts_output. - topk_ids: The topk_ids. - apply_router_weight_on_input: When False, apply the weights to fused_expert_output. - weight_and_reduce_impl: An optional TopKWeightAndReduce implementation.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def finalize(
    self,
    output: torch.Tensor,
    fused_expert_output: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool,
    weight_and_reduce_impl: TopKWeightAndReduce,
) -> None:
    """Perform any combine plus apply weights and perform a reduction on the
    fused experts output.
    - output: The output tensor, written in place.  Must be (M, K) shape.
    - fused_expert_output: The unweighted, unreduced output of the fused
      experts, it will have (M, topk, K) shape.
    - topk_weights: The weights to be applied to the fused_experts_output.
    - topk_ids: The topk_ids.
    - apply_router_weight_on_input: When False, apply the weights to
      fused_expert_output.
    - weight_and_reduce_impl: An optional TopKWeightAndReduce
      implementation.
    """
    raise NotImplementedError

finalize_async(output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl)

Perform any combine plus apply weights and perform a reduction on the fused experts output but do not wait for results from other workers. - output: The output tensor, written in place. Must be (M, K) shape. - fused_expert_output: The unweighted, unreduced output of the fused experts, it will have (M, topk, K) shape. - topk_weights: The weights to be applied to the fused_experts_output. - topk_ids: The topk_ids. - apply_router_weight_on_input: When False, apply the weights to fused_expert_output. - weight_and_reduce_impl: An optional TopKWeightAndReduce implementation.

Returns a callback or a hook callback pair that when invoked waits for results from other workers and has the same return signature as finalize, if a hook is returned this is more lightweight check that the recv is complete without doing extra work (used by DBO, will be refactored in the very near future)

ret = obj.finalize_async(output, ...) ... output not valid yet ... if isinstance(ret, tuple): hook, receiver = ret hook() receiver() ... output valid here ...

is equivalent to:

obj.finalize(output, ...)

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def finalize_async(
    self,
    output: torch.Tensor,
    fused_expert_output: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool,
    weight_and_reduce_impl: TopKWeightAndReduce,
) -> tuple[Callable, Callable] | Callable:
    """Perform any combine plus apply weights and perform a reduction on the
    fused experts output but do not wait for results from other workers.
    - output: The output tensor, written in place.  Must be (M, K) shape.
    - fused_expert_output: The unweighted, unreduced output of the fused
      experts, it will have (M, topk, K) shape.
    - topk_weights: The weights to be applied to the fused_experts_output.
    - topk_ids: The topk_ids.
    - apply_router_weight_on_input: When False, apply the weights to
      fused_expert_output.
    - weight_and_reduce_impl: An optional TopKWeightAndReduce
      implementation.

    Returns a callback or a hook callback pair that when invoked waits for
    results from other workers and has the same return signature as
    `finalize`, if a hook is returned this is more lightweight check that
    the recv is complete without doing extra work (used by DBO, will be
    refactored in the very near future)

    ret = obj.finalize_async(output, ...)
    ... output not valid yet ...
    if isinstance(ret, tuple):
        hook, receiver = ret
        hook()
    receiver()
    ... output valid here ...

    is equivalent to:

    obj.finalize(output, ...)
    """
    raise NotImplementedError

prepare(a1, topk_weights, topk_ids, num_experts, expert_map, apply_router_weight_on_input, quant_config, defer_input_quant) abstractmethod

Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - topk_ids: The topk ids. - topk_weights: The topk weights. - num_experts: The total number of experts in the global expert space. - expert_map: A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard. - apply_router_weight_on_input: When True, apply the weights to the activations, before quantization + dispatching. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to defer input quantization to the FusedMoEExpertsModular in cases where the compute kernel expects unquantized inputs

Returns a tuple of: - quantized + dispatched a. - Optional quantized + dispatched a1_scales. - Optional ExpertTokensMetadata containing gpu/cpu tensors as big as the number of local experts with the information about the number of tokens assigned to each local expert. - Optional dispatched expert topk IDs - Optional dispatched expert topk weight

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def prepare(
    self,
    a1: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    num_experts: int,
    expert_map: torch.Tensor | None,
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
    defer_input_quant: bool,
) -> PrepareResultType:
    """Perform any quantization (and/or) dispatching needed for this kernel.
    - a1: The (unquantized) input to the MoE layer.
    - topk_ids: The topk ids.
    - topk_weights: The topk weights.
    - num_experts: The total number of experts in the global expert space.
    - expert_map: A tensor mapping expert indices from the global expert
      space to the local expert space of the expert parallel shard.
    - apply_router_weight_on_input: When True, apply the weights to the
      activations, before quantization + dispatching.
    - quant_config: Quantization info provided by the fused experts.
    - defer_input_quant: Runtime parameter indicating whether or not to
      defer input quantization to the FusedMoEExpertsModular
      in cases where the compute kernel expects unquantized inputs

    Returns a tuple of:
    - quantized + dispatched a.
    - Optional quantized + dispatched a1_scales.
    - Optional ExpertTokensMetadata containing gpu/cpu tensors
      as big as the number of local experts with the information about the
      number of tokens assigned to each local expert.
    - Optional dispatched expert topk IDs
    - Optional dispatched expert topk weight
    """
    raise NotImplementedError

prepare_async(a1, topk_weights, topk_ids, num_experts, expert_map, apply_router_weight_on_input, quant_config, defer_input_quant)

Perform any quantization (and/or) dispatching needed for this kernel but do not wait for results from other workers. - a1: The (unquantized) input to the MoE layer. - a1_scale: Optional scales for a1 - a2_scale: Optional scales for the second MoE gemm. Required to make sure the quantization is consistent for both gemms. - topk_ids: The topk ids. - topk_weights: The topk weights. - num_experts: The total number of experts in the global expert space. - expert_map: A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard. - apply_router_weight_on_input: When True, apply the weights to the activations, before quantization + dispatching. - defer_input_quant: Runtime parameter indicating whether or not to defer input quantization to the FusedMoEExpertsModular in cases where the compute kernel expects unquantized inputs

Returns a callback or a hook callback pair that when invoked waits for results from other workers and has the same return signature as prepare, if a hook is returned this is more lightweight check that the recv is complete without doing extra work (used by DBO, will be refactored in the very near future)

e.g.

ret = obj.prepare_async(...)

if isinstance(ret, tuple): hook, receiver = ret hook()

if hook is not None: a, a_scales, expert_meta, topk_ids, topk_weights = receiver()

is equivalent to:

a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...)

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

def prepare_async(
    self,
    a1: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    num_experts: int,
    expert_map: torch.Tensor | None,
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
    defer_input_quant: bool,
) -> tuple[Callable, ReceiverType] | ReceiverType:
    """Perform any quantization (and/or) dispatching needed for this kernel
    but do not wait for results from other workers.
    - a1: The (unquantized) input to the MoE layer.
    - a1_scale: Optional scales for a1
    - a2_scale: Optional scales for the second MoE gemm.  Required to make
      sure the quantization is consistent for both gemms.
    - topk_ids: The topk ids.
    - topk_weights: The topk weights.
    - num_experts: The total number of experts in the global expert space.
    - expert_map: A tensor mapping expert indices from the global expert
      space to the local expert space of the expert parallel shard.
    - apply_router_weight_on_input: When True, apply the weights to the
      activations, before quantization + dispatching.
    - defer_input_quant: Runtime parameter indicating whether or not to
      defer input quantization to the FusedMoEExpertsModular
      in cases where the compute kernel expects unquantized inputs

    Returns a callback or a hook callback pair that when invoked waits for
    results from other workers and has the same return signature as
    `prepare`, if a hook is returned this is more lightweight check that
    the recv is complete without doing extra work (used by DBO, will be
    refactored in the very near future)

    e.g.

    ret = obj.prepare_async(...)

    if isinstance(ret, tuple):
        hook, receiver = ret
        hook()

    if hook is not None:
    a, a_scales, expert_meta, topk_ids, topk_weights = receiver()

    is equivalent to:

    a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...)
    """
    raise NotImplementedError

FusedMoEPrepareAndFinalizeMonolithic

Bases: FusedMoEPrepareAndFinalize

An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above for the monolithic case.

Methods:

• finalize –

  Optional method for subclasses compatible with monolithic

• prepare –

  Optional method for subclasses compatible with monolithic

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):
    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
    described above for the monolithic case.
    """

    @abstractmethod
    def prepare(
        self,
        a1: torch.Tensor,
        router_logits: torch.Tensor,
        quant_config: FusedMoEQuantConfig,
        defer_input_quant: bool = False,
    ) -> PrepareMonolithicResultType:
        """Optional method for subclasses compatible with monolithic
        FusedMoEExpertsModular kernels.

        Perform any quantization (and/or) dispatching needed for this kernel.
        - a1: The (unquantized) input to the MoE layer.
        - quant_config: Quantization info provided by the fused experts.
        - defer_input_quant: Runtime parameter indicating whether or not to
            defer input quantization to the FusedMoEExpertsModular

        Returns a tuple of:
        - quantized + dispatched a.
        - Optional quantized + dispatched a1_scales.
        """
        raise NotImplementedError

    @abstractmethod
    def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor:
        """Optional method for subclasses compatible with monolithic
        FusedMoEExpertsModular kernels.

        Perform any combine plus apply weights and perform a reduction on the
        fused experts output.
        - fused_expert_output: The unweighted, unreduced output of the fused
          experts, it will have (M, topk, K) shape.
        """
        raise NotImplementedError

finalize(fused_expert_output) abstractmethod

Optional method for subclasses compatible with monolithic FusedMoEExpertsModular kernels.

Perform any combine plus apply weights and perform a reduction on the fused experts output. - fused_expert_output: The unweighted, unreduced output of the fused experts, it will have (M, topk, K) shape.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor:
    """Optional method for subclasses compatible with monolithic
    FusedMoEExpertsModular kernels.

    Perform any combine plus apply weights and perform a reduction on the
    fused experts output.
    - fused_expert_output: The unweighted, unreduced output of the fused
      experts, it will have (M, topk, K) shape.
    """
    raise NotImplementedError

prepare(a1, router_logits, quant_config, defer_input_quant=False) abstractmethod

Optional method for subclasses compatible with monolithic FusedMoEExpertsModular kernels.

Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to defer input quantization to the FusedMoEExpertsModular

Returns a tuple of: - quantized + dispatched a. - Optional quantized + dispatched a1_scales.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def prepare(
    self,
    a1: torch.Tensor,
    router_logits: torch.Tensor,
    quant_config: FusedMoEQuantConfig,
    defer_input_quant: bool = False,
) -> PrepareMonolithicResultType:
    """Optional method for subclasses compatible with monolithic
    FusedMoEExpertsModular kernels.

    Perform any quantization (and/or) dispatching needed for this kernel.
    - a1: The (unquantized) input to the MoE layer.
    - quant_config: Quantization info provided by the fused experts.
    - defer_input_quant: Runtime parameter indicating whether or not to
        defer input quantization to the FusedMoEExpertsModular

    Returns a tuple of:
    - quantized + dispatched a.
    - Optional quantized + dispatched a1_scales.
    """
    raise NotImplementedError

TopKWeightAndReduce

Bases: ABC

An abstract base class for weight application and reduction implementations.

Methods:

• apply –

  Apply topk_weights to the fused_experts_outputs and/or reduce.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

class TopKWeightAndReduce(ABC):
    """An abstract base class for weight application and reduction implementations."""

    @abstractmethod
    def apply(
        self,
        output: torch.Tensor | None,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
    ) -> torch.Tensor:
        """Apply topk_weights to the fused_experts_outputs and/or reduce.
        If an output tensor is not passed, it will be created in the
        function.
        """
        raise NotImplementedError

apply(output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input) abstractmethod

Apply topk_weights to the fused_experts_outputs and/or reduce. If an output tensor is not passed, it will be created in the function.

Source code in vllm/model_executor/layers/fused_moe/modular_kernel.py

@abstractmethod
def apply(
    self,
    output: torch.Tensor | None,
    fused_expert_output: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool,
) -> torch.Tensor:
    """Apply topk_weights to the fused_experts_outputs and/or reduce.
    If an output tensor is not passed, it will be created in the
    function.
    """
    raise NotImplementedError