vllm.model_executor.layers.fused_moe ¶

Modules:

Name	Description
`batched_deep_gemm_moe`
`batched_triton_or_deep_gemm_moe`
`config`
`cpu_fused_moe`
`cutlass_moe`	CUTLASS based Fused MoE kernels.
`deep_gemm_moe`
`deep_gemm_utils`	Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py
`deepep_ht_prepare_finalize`
`deepep_ll_prepare_finalize`
`flashinfer_cutlass_moe`
`flashinfer_cutlass_prepare_finalize`
`flashinfer_trtllm_moe`
`fused_batched_moe`	Fused batched MoE kernel.
`fused_marlin_moe`	Fused MoE utilities for GPTQ.
`fused_moe`	Fused MoE Triton kernels.
`gpt_oss_triton_kernels_moe`
`layer`
`modular_kernel`
`moe_align_block_size`
`moe_pallas`
`moe_permute_unpermute`
`moe_torch_iterative`
`pplx_prepare_finalize`
`prepare_finalize`
`rocm_aiter_fused_moe`
`routing_simulator`	Token-to-Expert Routing Simulator
`topk_weight_and_reduce`
`triton_deep_gemm_moe`
`trtllm_moe`
`utils`

all `module-attribute` ¶

__all__ = [
    "FusedMoE",
    "FusedMoEConfig",
    "FusedMoEMethodBase",
    "FusedMoeWeightScaleSupported",
    "FusedMoEPermuteExpertsUnpermute",
    "FusedMoEActivationFormat",
    "FusedMoEPrepareAndFinalize",
    "activation_without_mul",
    "override_config",
    "get_config",
]

_config `module-attribute` ¶

_config: Optional[dict[str, Any]] = None

BatchedDeepGemmExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

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

class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

    def __init__(
        self,
        max_num_tokens: int,
        num_dispatchers: int,
        quant_config: FusedMoEQuantConfig,
    ):
        """
        max_num_tokens: Maximum number of tokens from a DP Rank
        num_dispatchers: The number of DP dispatchers.
        quant_config: Quantization configuration
        """
        super().__init__(quant_config)
        assert self.block_shape == deep_gemm_block_shape()
        self.max_num_tokens = max_num_tokens
        self.num_dispatchers = num_dispatchers

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.BatchedExperts,
                mk.FusedMoEActivationFormat.BatchedExperts)

    def supports_chunking(self) -> bool:
        return False

    def supports_expert_map(self) -> bool:
        return False

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        # Let PrepareAndFinalize::finalize() decide the impl.
        return TopKWeightAndReduceDelegate()

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        assert a.dim() == 2
        # FIXME (varun): We should be able to dispatch only from the leader
        # DP ranks in the case of TP > 1. At the moment, all the Ranks
        # end up sending their tokens. This needs to be fixed.
        num_dispatchers = self.num_dispatchers
        num_experts = local_num_experts
        max_num_tokens = a.size(
            0) if self.max_num_tokens is None else self.max_num_tokens
        workspace13 = (num_experts, max_num_tokens * num_dispatchers,
                       max(K, N))
        workspace2 = (num_experts, max_num_tokens * num_dispatchers, (N // 2))
        output = (num_experts, max_num_tokens * num_dispatchers, K)
        return (workspace13, workspace2, output, a.dtype)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        assert expert_tokens_meta is not None
        expert_num_tokens = expert_tokens_meta.expert_num_tokens

        assert hidden_states.ndim == 3
        assert self.block_shape is not None

        a1q = hidden_states
        _, N, K = w1.size()

        assert w2.size(1) == K

        E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
            hidden_states, w1, w2, topk_ids)

        workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))

        # (from deepgemm docs) : A value hint (which is a value on CPU)
        # for the M expectation of each batch, correctly setting this value
        # may lead to better performance.
        expected_m = max_num_tokens
        fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, self.w1_scale),
                                     workspace1, expert_num_tokens, expected_m)

        a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm_cuda(
            workspace1, expert_num_tokens)

        fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, self.w2_scale),
                                     output, expert_num_tokens, expected_m)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

max_num_tokens `instance-attribute` ¶

max_num_tokens = max_num_tokens

num_dispatchers `instance-attribute` ¶

num_dispatchers = num_dispatchers

init ¶

__init__(
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
)

max_num_tokens: Maximum number of tokens from a DP Rank num_dispatchers: The number of DP dispatchers. quant_config: Quantization configuration

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

def __init__(
    self,
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
):
    """
    max_num_tokens: Maximum number of tokens from a DP Rank
    num_dispatchers: The number of DP dispatchers.
    quant_config: Quantization configuration
    """
    super().__init__(quant_config)
    assert self.block_shape == deep_gemm_block_shape()
    self.max_num_tokens = max_num_tokens
    self.num_dispatchers = num_dispatchers

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    assert expert_tokens_meta is not None
    expert_num_tokens = expert_tokens_meta.expert_num_tokens

    assert hidden_states.ndim == 3
    assert self.block_shape is not None

    a1q = hidden_states
    _, N, K = w1.size()

    assert w2.size(1) == K

    E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
        hidden_states, w1, w2, topk_ids)

    workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))

    # (from deepgemm docs) : A value hint (which is a value on CPU)
    # for the M expectation of each batch, correctly setting this value
    # may lead to better performance.
    expected_m = max_num_tokens
    fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, self.w1_scale),
                                 workspace1, expert_num_tokens, expected_m)

    a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm_cuda(
        workspace1, expert_num_tokens)

    fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, self.w2_scale),
                                 output, expert_num_tokens, expected_m)

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    # Let PrepareAndFinalize::finalize() decide the impl.
    return TopKWeightAndReduceDelegate()

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return False

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return False

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_metadata: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    assert a.dim() == 2
    # FIXME (varun): We should be able to dispatch only from the leader
    # DP ranks in the case of TP > 1. At the moment, all the Ranks
    # end up sending their tokens. This needs to be fixed.
    num_dispatchers = self.num_dispatchers
    num_experts = local_num_experts
    max_num_tokens = a.size(
        0) if self.max_num_tokens is None else self.max_num_tokens
    workspace13 = (num_experts, max_num_tokens * num_dispatchers,
                   max(K, N))
    workspace2 = (num_experts, max_num_tokens * num_dispatchers, (N // 2))
    output = (num_experts, max_num_tokens * num_dispatchers, K)
    return (workspace13, workspace2, output, a.dtype)

BatchedTritonExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

A Triton based MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the pplx dispatch/combine kernels use.

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

class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
    """
    A Triton based MoE expert class that operates on expert batched format,
    i.e. E x max_num_tokens x K.  This is the format that the pplx
    dispatch/combine kernels use.
    """

    def __init__(
        self,
        max_num_tokens: int,
        num_dispatchers: int,
        quant_config: FusedMoEQuantConfig,
    ):
        super().__init__(quant_config)
        assert not self.quant_config.use_int8_w8a8, "NYI"
        assert not self.quant_config.use_int8_w8a16, "NYI"
        assert not self.quant_config.use_int4_w4a16, "NYI"
        assert not self.quant_config.use_mxfp4_w4a4, "NYI"
        assert max_num_tokens > 0
        assert num_dispatchers > 0
        self.max_num_tokens = max_num_tokens
        self.num_dispatchers = num_dispatchers

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.BatchedExperts,
                mk.FusedMoEActivationFormat.BatchedExperts)

    def supports_chunking(self) -> bool:
        return False

    def supports_expert_map(self) -> bool:
        return False

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        # Let PrepareAndFinalize::finalize() decide the impl.
        return TopKWeightAndReduceDelegate()

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        assert a.dim() == 2
        num_dp = self.num_dispatchers
        num_experts = local_num_experts
        max_num_tokens = self.max_num_tokens
        workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
        workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
        output = (num_experts, max_num_tokens * num_dp, K)
        return (workspace13, workspace2, output, a.dtype)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        # Check constraints.
        if self.quant_config.use_int4_w4a16:
            assert hidden_states.size(-1) // 2 == w1.size(2), (
                "Hidden size mismatch")
        else:
            assert hidden_states.size(-1) == w1.size(2), (
                f"Hidden size mismatch {hidden_states.size(-1)} "
                f"!= {w1.size(2)}")

        assert hidden_states.is_contiguous(
        ), "Hidden_states must be contiguous"
        assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
        assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
        assert hidden_states.dtype in [
            torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
        ]
        assert expert_tokens_meta is not None

        expert_num_tokens = expert_tokens_meta.expert_num_tokens

        E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
            hidden_states, w1, w2, topk_ids)

        assert w1.size(0) == E
        assert w2.size(0) == E

        config_dtype = self.quant_config.config_name(hidden_states.dtype)

        config = try_get_optimal_moe_config(
            w1.size(),
            w2.size(),
            top_k_num,
            config_dtype,
            max_num_tokens,
            block_shape=self.block_shape,
        )

        if hidden_states.dtype == torch.bfloat16:
            compute_type = tl.bfloat16
        elif hidden_states.dtype == torch.float16:
            compute_type = tl.float16
        elif hidden_states.dtype == torch.float32:
            compute_type = tl.float32
        elif hidden_states.dtype == torch.float8_e4m3fn:
            compute_type = tl.bfloat16
        else:
            raise ValueError(
                f"Unsupported compute_type: {hidden_states.dtype}")

        # We can reuse the memory between these because by the time we need
        # cache3, we're done with cache1
        intermediate_cache1 = _resize_cache(workspace13,
                                            (E, max_num_tokens, N))
        intermediate_cache2 = _resize_cache(workspace2,
                                            (E, max_num_tokens, N // 2))

        # TODO(bnell): should this be done for any quantized type?
        if self.quant_config.use_fp8_w8a8:
            intermediate_cache1.fill_(0)

        a1q_scale = normalize_batched_scales_shape(a1q_scale, E)

        # MM1
        invoke_moe_batched_triton_kernel(
            A=hidden_states,
            B=w1,
            C=intermediate_cache1,
            expert_num_tokens=expert_num_tokens,
            compute_type=compute_type,
            A_scale=a1q_scale,
            B_scale=self.w1_scale,
            B_zp=self.w1_zp,
            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
            use_int8_w8a16=self.quant_config.use_int8_w8a16,
            use_int4_w4a16=self.quant_config.use_int4_w4a16,
            config=config,
            per_act_token_quant=self.per_act_token_quant,
            block_shape=self.block_shape)

        intermediate_cache2.fill_(0)

        # TODO (bnell): use triton utility from batched deep gemm.
        self.activation(activation, intermediate_cache2.view(-1, N // 2),
                        intermediate_cache1.view(-1, N))

        qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
            intermediate_cache2, a2_scale, max_num_tokens, E, N,
            expert_num_tokens, self.quant_dtype, self.per_act_token_quant,
            self.block_shape)

        invoke_moe_batched_triton_kernel(
            A=qintermediate_cache2,
            B=w2,
            C=output,
            expert_num_tokens=expert_num_tokens,
            compute_type=compute_type,
            A_scale=a2q_scale,
            B_scale=self.w2_scale,
            B_zp=self.w2_zp,
            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
            use_int8_w8a16=self.quant_config.use_int8_w8a16,
            use_int4_w4a16=self.quant_config.use_int4_w4a16,
            config=config,
            per_act_token_quant=self.per_act_token_quant,
            block_shape=self.block_shape)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

max_num_tokens `instance-attribute` ¶

max_num_tokens = max_num_tokens

num_dispatchers `instance-attribute` ¶

num_dispatchers = num_dispatchers

init ¶

__init__(
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
)

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

def __init__(
    self,
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
):
    super().__init__(quant_config)
    assert not self.quant_config.use_int8_w8a8, "NYI"
    assert not self.quant_config.use_int8_w8a16, "NYI"
    assert not self.quant_config.use_int4_w4a16, "NYI"
    assert not self.quant_config.use_mxfp4_w4a4, "NYI"
    assert max_num_tokens > 0
    assert num_dispatchers > 0
    self.max_num_tokens = max_num_tokens
    self.num_dispatchers = num_dispatchers

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    # Check constraints.
    if self.quant_config.use_int4_w4a16:
        assert hidden_states.size(-1) // 2 == w1.size(2), (
            "Hidden size mismatch")
    else:
        assert hidden_states.size(-1) == w1.size(2), (
            f"Hidden size mismatch {hidden_states.size(-1)} "
            f"!= {w1.size(2)}")

    assert hidden_states.is_contiguous(
    ), "Hidden_states must be contiguous"
    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
    assert hidden_states.dtype in [
        torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
    ]
    assert expert_tokens_meta is not None

    expert_num_tokens = expert_tokens_meta.expert_num_tokens

    E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
        hidden_states, w1, w2, topk_ids)

    assert w1.size(0) == E
    assert w2.size(0) == E

    config_dtype = self.quant_config.config_name(hidden_states.dtype)

    config = try_get_optimal_moe_config(
        w1.size(),
        w2.size(),
        top_k_num,
        config_dtype,
        max_num_tokens,
        block_shape=self.block_shape,
    )

    if hidden_states.dtype == torch.bfloat16:
        compute_type = tl.bfloat16
    elif hidden_states.dtype == torch.float16:
        compute_type = tl.float16
    elif hidden_states.dtype == torch.float32:
        compute_type = tl.float32
    elif hidden_states.dtype == torch.float8_e4m3fn:
        compute_type = tl.bfloat16
    else:
        raise ValueError(
            f"Unsupported compute_type: {hidden_states.dtype}")

    # We can reuse the memory between these because by the time we need
    # cache3, we're done with cache1
    intermediate_cache1 = _resize_cache(workspace13,
                                        (E, max_num_tokens, N))
    intermediate_cache2 = _resize_cache(workspace2,
                                        (E, max_num_tokens, N // 2))

    # TODO(bnell): should this be done for any quantized type?
    if self.quant_config.use_fp8_w8a8:
        intermediate_cache1.fill_(0)

    a1q_scale = normalize_batched_scales_shape(a1q_scale, E)

    # MM1
    invoke_moe_batched_triton_kernel(
        A=hidden_states,
        B=w1,
        C=intermediate_cache1,
        expert_num_tokens=expert_num_tokens,
        compute_type=compute_type,
        A_scale=a1q_scale,
        B_scale=self.w1_scale,
        B_zp=self.w1_zp,
        use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
        use_int8_w8a16=self.quant_config.use_int8_w8a16,
        use_int4_w4a16=self.quant_config.use_int4_w4a16,
        config=config,
        per_act_token_quant=self.per_act_token_quant,
        block_shape=self.block_shape)

    intermediate_cache2.fill_(0)

    # TODO (bnell): use triton utility from batched deep gemm.
    self.activation(activation, intermediate_cache2.view(-1, N // 2),
                    intermediate_cache1.view(-1, N))

    qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
        intermediate_cache2, a2_scale, max_num_tokens, E, N,
        expert_num_tokens, self.quant_dtype, self.per_act_token_quant,
        self.block_shape)

    invoke_moe_batched_triton_kernel(
        A=qintermediate_cache2,
        B=w2,
        C=output,
        expert_num_tokens=expert_num_tokens,
        compute_type=compute_type,
        A_scale=a2q_scale,
        B_scale=self.w2_scale,
        B_zp=self.w2_zp,
        use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
        use_int8_w8a16=self.quant_config.use_int8_w8a16,
        use_int4_w4a16=self.quant_config.use_int4_w4a16,
        config=config,
        per_act_token_quant=self.per_act_token_quant,
        block_shape=self.block_shape)

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    # Let PrepareAndFinalize::finalize() decide the impl.
    return TopKWeightAndReduceDelegate()

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return False

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return False

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    assert a.dim() == 2
    num_dp = self.num_dispatchers
    num_experts = local_num_experts
    max_num_tokens = self.max_num_tokens
    workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
    workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
    output = (num_experts, max_num_tokens * num_dp, K)
    return (workspace13, workspace2, output, a.dtype)

BatchedTritonOrDeepGemmExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

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

class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

    def __init__(
        self,
        max_num_tokens: int,
        num_dispatchers: int,
        quant_config: FusedMoEQuantConfig,
        allow_deep_gemm: bool = False,
    ):
        super().__init__(quant_config)

        self.batched_triton_experts = BatchedTritonExperts(
            max_num_tokens=max_num_tokens,
            num_dispatchers=num_dispatchers,
            quant_config=self.quant_config,
        )

        self.allow_deep_gemm = (allow_deep_gemm
                                and self.quant_config.use_fp8_w8a8 and
                                self.block_shape == deep_gemm_block_shape())

        self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
            max_num_tokens=max_num_tokens,
            num_dispatchers=num_dispatchers,
            quant_config=self.quant_config,
        ) if self.allow_deep_gemm else None

        assert (self.batched_deep_gemm_experts is not None
                or self.batched_triton_experts is not None)

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        if self.batched_triton_experts is not None:
            assert (self.batched_deep_gemm_experts is None
                    or self.batched_deep_gemm_experts.activation_formats
                    == self.batched_triton_experts.activation_formats)
            return self.batched_triton_experts.activation_formats
        else:
            assert self.batched_deep_gemm_experts is not None
            return self.batched_deep_gemm_experts.activation_formats

    def supports_chunking(self) -> bool:
        bdge = self.batched_deep_gemm_experts
        bte = self.batched_triton_experts
        return ((bdge is None or bdge.supports_chunking())
                and (bte is None or bte.supports_chunking()))

    def supports_expert_map(self) -> bool:
        bdge = self.batched_deep_gemm_experts
        bte = self.batched_triton_experts
        return ((bdge is None or bdge.supports_expert_map())
                and (bte is None or bte.supports_expert_map()))

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        bdge = self.batched_deep_gemm_experts
        bte = self.batched_triton_experts
        bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None
        bte_war = bte.finalize_weight_and_reduce_impl() if bte else None
        is_bdge_war = bdge_war is not None
        is_bte_war = bte_war is not None

        if is_bdge_war and is_bte_war:
            assert bdge_war == bte_war, (
                "Both implementations should agree on WeightAndReduce impls. "
                f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}")

        if bdge_war is not None:
            return bdge_war

        assert bte_war is not None
        return bte_war

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        # Note: the deep gemm workspaces are strictly larger than the triton
        # workspaces so we can be pessimistic here and allocate for DeepGemm
        # even if we fall back to triton later, e.g. if expert maps are set.
        if self.allow_deep_gemm:
            assert self.batched_deep_gemm_experts is not None
            return self.batched_deep_gemm_experts.workspace_shapes(
                a, aq, M, N, K, topk, global_num_experts, local_num_experts,
                expert_tokens_metadata)
        else:
            assert self.batched_triton_experts is not None
            return self.batched_triton_experts.workspace_shapes(
                a, aq, M, N, K, topk, global_num_experts, local_num_experts,
                expert_tokens_metadata)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        experts = (self.batched_deep_gemm_experts
                   if self.allow_deep_gemm else self.batched_triton_experts)
        assert experts is not None
        experts.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)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

allow_deep_gemm `instance-attribute` ¶

allow_deep_gemm = (
    allow_deep_gemm
    and use_fp8_w8a8
    and block_shape == deep_gemm_block_shape()
)

batched_deep_gemm_experts `instance-attribute` ¶

batched_deep_gemm_experts = (
    BatchedDeepGemmExperts(
        max_num_tokens=max_num_tokens,
        num_dispatchers=num_dispatchers,
        quant_config=quant_config,
    )
    if allow_deep_gemm
    else None
)

batched_triton_experts `instance-attribute` ¶

batched_triton_experts = BatchedTritonExperts(
    max_num_tokens=max_num_tokens,
    num_dispatchers=num_dispatchers,
    quant_config=quant_config,
)

init ¶

__init__(
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
    allow_deep_gemm: bool = False,
)

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

def __init__(
    self,
    max_num_tokens: int,
    num_dispatchers: int,
    quant_config: FusedMoEQuantConfig,
    allow_deep_gemm: bool = False,
):
    super().__init__(quant_config)

    self.batched_triton_experts = BatchedTritonExperts(
        max_num_tokens=max_num_tokens,
        num_dispatchers=num_dispatchers,
        quant_config=self.quant_config,
    )

    self.allow_deep_gemm = (allow_deep_gemm
                            and self.quant_config.use_fp8_w8a8 and
                            self.block_shape == deep_gemm_block_shape())

    self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
        max_num_tokens=max_num_tokens,
        num_dispatchers=num_dispatchers,
        quant_config=self.quant_config,
    ) if self.allow_deep_gemm else None

    assert (self.batched_deep_gemm_experts is not None
            or self.batched_triton_experts is not None)

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    experts = (self.batched_deep_gemm_experts
               if self.allow_deep_gemm else self.batched_triton_experts)
    assert experts is not None
    experts.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)

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    bdge = self.batched_deep_gemm_experts
    bte = self.batched_triton_experts
    bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None
    bte_war = bte.finalize_weight_and_reduce_impl() if bte else None
    is_bdge_war = bdge_war is not None
    is_bte_war = bte_war is not None

    if is_bdge_war and is_bte_war:
        assert bdge_war == bte_war, (
            "Both implementations should agree on WeightAndReduce impls. "
            f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}")

    if bdge_war is not None:
        return bdge_war

    assert bte_war is not None
    return bte_war

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    bdge = self.batched_deep_gemm_experts
    bte = self.batched_triton_experts
    return ((bdge is None or bdge.supports_chunking())
            and (bte is None or bte.supports_chunking()))

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    bdge = self.batched_deep_gemm_experts
    bte = self.batched_triton_experts
    return ((bdge is None or bdge.supports_expert_map())
            and (bte is None or bte.supports_expert_map()))

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_metadata: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    # Note: the deep gemm workspaces are strictly larger than the triton
    # workspaces so we can be pessimistic here and allocate for DeepGemm
    # even if we fall back to triton later, e.g. if expert maps are set.
    if self.allow_deep_gemm:
        assert self.batched_deep_gemm_experts is not None
        return self.batched_deep_gemm_experts.workspace_shapes(
            a, aq, M, N, K, topk, global_num_experts, local_num_experts,
            expert_tokens_metadata)
    else:
        assert self.batched_triton_experts is not None
        return self.batched_triton_experts.workspace_shapes(
            a, aq, M, N, K, topk, global_num_experts, local_num_experts,
            expert_tokens_metadata)

CutlassBatchedExpertsFp8 ¶

Bases: CutlassExpertsFp8Base

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

class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

    def __init__(
        self,
        max_experts_per_worker: int,
        num_dispatchers: int,
        out_dtype: Optional[torch.dtype],
        ab_strides1: torch.Tensor,
        ab_strides2: torch.Tensor,
        c_strides1: torch.Tensor,
        c_strides2: torch.Tensor,
        quant_config: FusedMoEQuantConfig,
    ):
        super().__init__(
            out_dtype,
            ab_strides1,
            ab_strides2,
            c_strides1,
            c_strides2,
            quant_config,
        )
        assert max_experts_per_worker > 0
        self.max_experts_per_worker = max_experts_per_worker
        self.num_dispatchers = num_dispatchers

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.BatchedExperts,
                mk.FusedMoEActivationFormat.BatchedExperts)

    def supports_chunking(self) -> bool:
        return False

    def supports_expert_map(self) -> bool:
        return False

    # TODO(bnell): maybe remove need for passing aq to workspace_shapes
    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        padded_M = aq.size(1)
        num_dp = self.num_dispatchers
        assert num_dp is not None
        workspace1 = (self.max_experts_per_worker, padded_M * num_dp,
                      max(N, K))
        workspace2 = (self.max_experts_per_worker, padded_M * num_dp,
                      max(N // 2, K))
        output = (self.max_experts_per_worker, padded_M, K)
        return (workspace1, workspace2, output,
                self.out_dtype if self.out_dtype is not None else a.dtype)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

max_experts_per_worker `instance-attribute` ¶

max_experts_per_worker = max_experts_per_worker

num_dispatchers `instance-attribute` ¶

num_dispatchers = num_dispatchers

init ¶

__init__(
    max_experts_per_worker: int,
    num_dispatchers: int,
    out_dtype: Optional[dtype],
    ab_strides1: Tensor,
    ab_strides2: Tensor,
    c_strides1: Tensor,
    c_strides2: Tensor,
    quant_config: FusedMoEQuantConfig,
)

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

def __init__(
    self,
    max_experts_per_worker: int,
    num_dispatchers: int,
    out_dtype: Optional[torch.dtype],
    ab_strides1: torch.Tensor,
    ab_strides2: torch.Tensor,
    c_strides1: torch.Tensor,
    c_strides2: torch.Tensor,
    quant_config: FusedMoEQuantConfig,
):
    super().__init__(
        out_dtype,
        ab_strides1,
        ab_strides2,
        c_strides1,
        c_strides2,
        quant_config,
    )
    assert max_experts_per_worker > 0
    self.max_experts_per_worker = max_experts_per_worker
    self.num_dispatchers = num_dispatchers

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return False

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return False

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    padded_M = aq.size(1)
    num_dp = self.num_dispatchers
    assert num_dp is not None
    workspace1 = (self.max_experts_per_worker, padded_M * num_dp,
                  max(N, K))
    workspace2 = (self.max_experts_per_worker, padded_M * num_dp,
                  max(N // 2, K))
    output = (self.max_experts_per_worker, padded_M, K)
    return (workspace1, workspace2, output,
            self.out_dtype if self.out_dtype is not None else a.dtype)

CutlassExpertsFp8 ¶

Bases: CutlassExpertsFp8Base

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

class CutlassExpertsFp8(CutlassExpertsFp8Base):

    def __init__(
        self,
        out_dtype: Optional[torch.dtype],
        ab_strides1: torch.Tensor,
        ab_strides2: torch.Tensor,
        c_strides1: torch.Tensor,
        c_strides2: torch.Tensor,
        quant_config: FusedMoEQuantConfig,
    ):
        super().__init__(
            out_dtype,
            ab_strides1,
            ab_strides2,
            c_strides1,
            c_strides2,
            quant_config,
        )

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.Standard,
                mk.FusedMoEActivationFormat.Standard)

    def supports_chunking(self) -> bool:
        return True

    def supports_expert_map(self) -> bool:
        return True

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        # topk weights and reduction are fused in moe_unpermute cuda kernel
        return TopKWeightAndReduceNoOP()

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        workspace1 = (M * topk, max(N, K))
        workspace2 = (M * topk, max(N // 2, K))
        output = (M, K)
        return (workspace1, workspace2, output,
                self.out_dtype if self.out_dtype is not None else a.dtype)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

init ¶

__init__(
    out_dtype: Optional[dtype],
    ab_strides1: Tensor,
    ab_strides2: Tensor,
    c_strides1: Tensor,
    c_strides2: Tensor,
    quant_config: FusedMoEQuantConfig,
)

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

def __init__(
    self,
    out_dtype: Optional[torch.dtype],
    ab_strides1: torch.Tensor,
    ab_strides2: torch.Tensor,
    c_strides1: torch.Tensor,
    c_strides2: torch.Tensor,
    quant_config: FusedMoEQuantConfig,
):
    super().__init__(
        out_dtype,
        ab_strides1,
        ab_strides2,
        c_strides1,
        c_strides2,
        quant_config,
    )

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    # topk weights and reduction are fused in moe_unpermute cuda kernel
    return TopKWeightAndReduceNoOP()

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return True

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return True

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    workspace1 = (M * topk, max(N, K))
    workspace2 = (M * topk, max(N // 2, K))
    output = (M, K)
    return (workspace1, workspace2, output,
            self.out_dtype if self.out_dtype is not None else a.dtype)

DeepGemmExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

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

class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

    def __init__(self, quant_config: FusedMoEQuantConfig):
        super().__init__(quant_config)
        assert quant_config.block_shape == deep_gemm_block_shape()
        assert quant_config.quant_dtype == torch.float8_e4m3fn
        assert not quant_config.per_act_token_quant
        assert not quant_config.per_out_ch_quant

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.Standard,
                mk.FusedMoEActivationFormat.Standard)

    def supports_chunking(self) -> bool:
        return True

    def supports_expert_map(self) -> bool:
        return True

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        return TopKWeightAndReduceNoOP()

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        assert self.block_shape is not None
        block_m = self.block_shape[0]
        M_sum = compute_aligned_M(M, topk, local_num_experts, block_m,
                                  expert_tokens_meta)
        assert M_sum % block_m == 0

        workspace1 = (M_sum, max(N, K))
        workspace2 = (M_sum, max(N // 2, K))
        output = (M, K)
        return (workspace1, workspace2, output, a.dtype)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        assert a1q_scale is not None
        assert a2_scale is None
        assert self.block_shape is not None
        assert self.w1_scale is not None
        assert self.w2_scale is not None

        a1q = hidden_states
        _, N, K = w1.size()

        local_num_experts = w1.size(0)
        if global_num_experts == -1:
            global_num_experts = local_num_experts

        assert w2.size(1) == K

        M_sum = compute_aligned_M(M=topk_ids.size(0),
                                  num_topk=topk_ids.size(1),
                                  local_num_experts=local_num_experts,
                                  alignment=deep_gemm_block_shape()[0],
                                  expert_tokens_meta=expert_tokens_meta)

        a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),
                                 (M_sum, K))
        mm1_out = _resize_cache(workspace13, (M_sum, N))
        act_out = _resize_cache(workspace2, (M_sum, N // 2))
        quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
                                  (M_sum, N // 2))
        mm2_out = _resize_cache(workspace2, (M_sum, K))

        a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
            aq=a1q,
            aq_scale=a1q_scale,
            topk_ids=topk_ids,
            local_num_experts=local_num_experts,
            expert_map=expert_map,
            expert_tokens_meta=expert_tokens_meta,
            aq_out=a1q_perm)
        assert a1q.size(0) == M_sum

        m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, self.w1_scale),
                                         mm1_out, expert_ids)

        self.activation(activation, act_out, mm1_out.view(-1, N))

        a2q_scale: Optional[torch.Tensor] = None
        a2q, a2q_scale = per_token_group_quant_fp8(act_out,
                                                   self.block_shape[1],
                                                   column_major_scales=True,
                                                   out_q=quant_out)

        m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, self.w2_scale),
                                         mm2_out, expert_ids)

        if apply_router_weight_on_input:
            topk_weights = torch.ones_like(topk_weights)

        deepgemm_unpermute_and_reduce(a=mm2_out,
                                      topk_ids=topk_ids,
                                      topk_weights=topk_weights,
                                      inv_perm=inv_perm,
                                      expert_map=expert_map,
                                      output=output)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

init ¶

__init__(quant_config: FusedMoEQuantConfig)

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

def __init__(self, quant_config: FusedMoEQuantConfig):
    super().__init__(quant_config)
    assert quant_config.block_shape == deep_gemm_block_shape()
    assert quant_config.quant_dtype == torch.float8_e4m3fn
    assert not quant_config.per_act_token_quant
    assert not quant_config.per_out_ch_quant

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    assert a1q_scale is not None
    assert a2_scale is None
    assert self.block_shape is not None
    assert self.w1_scale is not None
    assert self.w2_scale is not None

    a1q = hidden_states
    _, N, K = w1.size()

    local_num_experts = w1.size(0)
    if global_num_experts == -1:
        global_num_experts = local_num_experts

    assert w2.size(1) == K

    M_sum = compute_aligned_M(M=topk_ids.size(0),
                              num_topk=topk_ids.size(1),
                              local_num_experts=local_num_experts,
                              alignment=deep_gemm_block_shape()[0],
                              expert_tokens_meta=expert_tokens_meta)

    a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),
                             (M_sum, K))
    mm1_out = _resize_cache(workspace13, (M_sum, N))
    act_out = _resize_cache(workspace2, (M_sum, N // 2))
    quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
                              (M_sum, N // 2))
    mm2_out = _resize_cache(workspace2, (M_sum, K))

    a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
        aq=a1q,
        aq_scale=a1q_scale,
        topk_ids=topk_ids,
        local_num_experts=local_num_experts,
        expert_map=expert_map,
        expert_tokens_meta=expert_tokens_meta,
        aq_out=a1q_perm)
    assert a1q.size(0) == M_sum

    m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, self.w1_scale),
                                     mm1_out, expert_ids)

    self.activation(activation, act_out, mm1_out.view(-1, N))

    a2q_scale: Optional[torch.Tensor] = None
    a2q, a2q_scale = per_token_group_quant_fp8(act_out,
                                               self.block_shape[1],
                                               column_major_scales=True,
                                               out_q=quant_out)

    m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, self.w2_scale),
                                     mm2_out, expert_ids)

    if apply_router_weight_on_input:
        topk_weights = torch.ones_like(topk_weights)

    deepgemm_unpermute_and_reduce(a=mm2_out,
                                  topk_ids=topk_ids,
                                  topk_weights=topk_weights,
                                  inv_perm=inv_perm,
                                  expert_map=expert_map,
                                  output=output)

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    return TopKWeightAndReduceNoOP()

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return True

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return True

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    assert self.block_shape is not None
    block_m = self.block_shape[0]
    M_sum = compute_aligned_M(M, topk, local_num_experts, block_m,
                              expert_tokens_meta)
    assert M_sum % block_m == 0

    workspace1 = (M_sum, max(N, K))
    workspace2 = (M_sum, max(N // 2, K))
    output = (M, K)
    return (workspace1, workspace2, output, a.dtype)

FusedMoE ¶

Bases: CustomOp

FusedMoE layer for MoE models.

This layer contains both MergedColumnParallel weights (gate_up_proj / w13) and RowParallelLinear weights (down_proj/ w2).

Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We copy that naming convention here and handle any remapping in the load_weights function in each model implementation.

Parameters:

Name	Type	Description	Default
`num_experts`	`int`	Number of experts in the model	required
`top_k`	`int`	Number of experts selected for each token	required
`hidden_size`	`int`	Input hidden state size of the transformer	required
`intermediate_size`	`int`	Intermediate size of the experts	required
`params_dtype`	`Optional[dtype]`	Data type for the parameters.	`None`
`reduce_results`	`bool`	Whether to all all_reduce on the output of the layer	`False`
`renormalize`	`bool`	Whether to renormalize the logits in the fused_moe kernel	`True`
`quant_config`	`Optional[QuantizationConfig]`	Quantization configure.	`None`
`enable_eplb`	`bool`	Whether to enable expert parallelism load balancer.	`False`

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

@CustomOp.register("fused_moe")
class FusedMoE(CustomOp):
    """FusedMoE layer for MoE models.

    This layer contains both MergedColumnParallel weights (gate_up_proj /
    w13) and RowParallelLinear weights (down_proj/ w2).

    Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We
    copy that naming convention here and handle any remapping in the
    load_weights function in each model implementation.

    Args:
        num_experts: Number of experts in the model
        top_k: Number of experts selected for each token
        hidden_size: Input hidden state size of the transformer
        intermediate_size: Intermediate size of the experts
        params_dtype: Data type for the parameters.
        reduce_results: Whether to all all_reduce on the output of the layer
        renormalize: Whether to renormalize the logits in the fused_moe kernel
        quant_config: Quantization configure.
        enable_eplb: Whether to enable expert parallelism load balancer.
    """

    def __init__(
        self,
        num_experts: int,  # Global number of experts
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        params_dtype: Optional[torch.dtype] = None,
        reduce_results: bool = False,
        renormalize: bool = True,
        use_grouped_topk: bool = False,
        num_expert_group: Optional[int] = None,
        topk_group: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
        tp_size: Optional[int] = None,
        ep_size: Optional[int] = None,
        dp_size: Optional[int] = None,
        prefix: str = "",
        custom_routing_function: Optional[Callable] = None,
        scoring_func: str = "softmax",
        routed_scaling_factor: float = 1.0,
        e_score_correction_bias: Optional[torch.Tensor] = None,
        apply_router_weight_on_input: bool = False,
        activation: str = "silu",
        enable_eplb: bool = False,
        num_redundant_experts: int = 0,
        has_bias: bool = False,
        is_sequence_parallel=False,
        zero_expert_num: Optional[int] = 0,
        zero_expert_type: Optional[str] = None,
    ):
        super().__init__()
        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype

        vllm_config = get_current_vllm_config()

        # FIXME (varun): We should have a better way of inferring the activation
        # datatype. This works for now as the tensor datatype entering the MoE
        # operation is typically unquantized (i.e. float16/bfloat16).
        if vllm_config.model_config is not None:
            moe_in_dtype = vllm_config.model_config.dtype
        else:
            # TODO (bnell): This is a hack to get test_mixtral_moe to work
            # since model_config is not set in the pytest test.
            moe_in_dtype = params_dtype

        tp_size_ = (tp_size if tp_size is not None else
                    get_tensor_model_parallel_world_size())
        dp_size_ = (dp_size
                    if dp_size is not None else get_dp_group().world_size)

        self.is_sequence_parallel = is_sequence_parallel
        self.sp_size = tp_size_ if is_sequence_parallel else 1

        self.moe_parallel_config: FusedMoEParallelConfig = (
            FusedMoEParallelConfig.make(
                tp_size_=tp_size_,
                dp_size_=dp_size_,
                vllm_parallel_config=vllm_config.parallel_config))

        self.global_num_experts = num_experts + num_redundant_experts
        self.zero_expert_num = zero_expert_num
        self.zero_expert_type = zero_expert_type

        # Round up hidden size if needed.
        hidden_size = maybe_roundup_hidden_size(hidden_size, moe_in_dtype,
                                                quant_config,
                                                self.moe_parallel_config)

        # For smuggling this layer into the fused moe custom op
        compilation_config = vllm_config.compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError("Duplicate layer name: {}".format(prefix))
        compilation_config.static_forward_context[prefix] = self
        self.layer_name = prefix

        self.enable_eplb = enable_eplb
        self.expert_load_view: Optional[torch.Tensor] = None
        self.logical_to_physical_map: Optional[torch.Tensor] = None
        self.logical_replica_count: Optional[torch.Tensor] = None

        # Determine expert maps
        if self.use_ep:
            if self.enable_eplb:
                assert self.global_num_experts % self.ep_size == 0, \
                    "EPLB currently only supports even distribution of " \
                    "experts across ranks."
            else:
                assert num_redundant_experts == 0, \
                    "Redundant experts are only supported with EPLB."

            expert_placement_strategy = (
                vllm_config.parallel_config.expert_placement_strategy)
            if expert_placement_strategy == "round_robin":
                # TODO(Bruce): will support round robin expert placement with
                # EPLB enabled in the future.
                round_robin_supported = ((num_expert_group is not None
                                          and num_expert_group > 1)
                                         and num_redundant_experts == 0
                                         and not self.enable_eplb)

                if not round_robin_supported:
                    logger.warning(
                        "Round-robin expert placement is only supported for "
                        "models with multiple expert groups and no redundant "
                        "experts. Falling back to linear expert placement.")
                    expert_placement_strategy = "linear"

            self.expert_map: Optional[torch.Tensor]
            local_num_experts, expert_map = determine_expert_map(
                ep_size=self.ep_size,
                ep_rank=self.ep_rank,
                global_num_experts=self.global_num_experts,
                expert_placement_strategy=expert_placement_strategy,
            )
            self.local_num_experts = local_num_experts
            self.register_buffer("expert_map", expert_map)
            logger.info_once(
                "[EP Rank %s/%s] Expert parallelism is enabled. Expert "
                "placement strategy: %s. Local/global"
                " number of experts: %s/%s. Experts local to global index map:"
                " %s.", self.ep_rank, self.ep_size, expert_placement_strategy,
                self.local_num_experts, self.global_num_experts,
                get_compressed_expert_map(self.expert_map))
        else:
            self.local_num_experts, self.expert_map = (self.global_num_experts,
                                                       None)

        self.top_k = top_k

        assert intermediate_size % self.tp_size == 0
        self.hidden_size = hidden_size
        self.intermediate_size_per_partition = intermediate_size // self.tp_size
        self.reduce_results = reduce_results
        self.renormalize = renormalize
        self.use_grouped_topk = use_grouped_topk
        if self.use_grouped_topk:
            assert num_expert_group is not None and topk_group is not None
        self.num_expert_group = num_expert_group
        self.topk_group = topk_group
        self.custom_routing_function = custom_routing_function
        self.scoring_func = scoring_func
        self.routed_scaling_factor = routed_scaling_factor
        self.e_score_correction_bias = e_score_correction_bias
        self.apply_router_weight_on_input = apply_router_weight_on_input
        self.activation = activation

        if self.scoring_func != "softmax" and not self.use_grouped_topk:
            raise ValueError("Only softmax scoring function is supported for "
                             "non-grouped topk.")

        moe = FusedMoEConfig(
            num_experts=self.global_num_experts,
            experts_per_token=top_k,
            hidden_dim=hidden_size,
            num_local_experts=self.local_num_experts,
            moe_parallel_config=self.moe_parallel_config,
            in_dtype=moe_in_dtype,
            max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
            has_bias=has_bias,
        )
        self.moe_config = moe
        self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
        self.quant_config = quant_config

        # Note: get_quant_method will look at the layer's local_num_experts
        # for heuristic purposes, so it must be initialized first.
        quant_method: Optional[QuantizeMethodBase] = None
        quant_method = (UnquantizedFusedMoEMethod(moe) if quant_config is None
                        else quant_config.get_quant_method(self, prefix))

        assert quant_method is not None
        assert isinstance(quant_method, FusedMoEMethodBase)
        self.quant_method = quant_method

        if self.enable_eplb:
            from vllm.model_executor.layers.quantization.fp8 import (
                Fp8MoEMethod)
            if not isinstance(quant_method,
                              (Fp8MoEMethod, UnquantizedFusedMoEMethod)):
                # TODO: Add support for additional quantization methods.
                # The implementation for other quantization methods does not
                # contain essential differences, but the current quant API
                # design causes duplicated work when extending to new
                # quantization methods, so I'm leaving it for now.
                # If you plan to add support for more quantization methods,
                # please refer to the implementation in `Fp8MoEMethod`.
                raise NotImplementedError("EPLB is only supported for FP8 "
                                          "quantization for now.")

        moe_quant_params = {
            "num_experts": self.local_num_experts,
            "hidden_size": hidden_size,
            "intermediate_size_per_partition":
            self.intermediate_size_per_partition,
            "params_dtype": params_dtype,
            "weight_loader": self.weight_loader,
        }
        # need full intermediate size pre-sharding for WNA16 act order
        if (self.quant_method.__class__.__name__
                in ("GPTQMarlinMoEMethod",
                    "CompressedTensorsWNA16MarlinMoEMethod",
                    "CompressedTensorsWNA16MoEMethod")):
            moe_quant_params["intermediate_size_full"] = intermediate_size

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

        # Chunked all2all staging tensor
        self.batched_hidden_states: Optional[torch.Tensor] = None
        self.batched_router_logits: Optional[torch.Tensor] = None

        # TODO(bnell): flashinfer uses non-batched format.
        # Does it really need a batched buffer?
        if (self.moe_parallel_config.use_pplx_kernels
                or self.moe_parallel_config.use_deepep_ll_kernels
                or self.moe_config.use_flashinfer_cutlass_kernels):
            if vllm_config.parallel_config.enable_dbo:
                self.batched_hidden_states = torch.zeros(
                    (2, moe.max_num_tokens, self.hidden_size),
                    dtype=moe.in_dtype,
                    device=torch.cuda.current_device())

                # Note here we use `num_experts` which is logical expert count
                self.batched_router_logits = torch.zeros(
                    (2, moe.max_num_tokens, num_experts),
                    dtype=moe.in_dtype,
                    device=torch.cuda.current_device())
            else:
                self.batched_hidden_states = torch.zeros(
                    (moe.max_num_tokens, self.hidden_size),
                    dtype=moe.in_dtype,
                    device=torch.cuda.current_device())

                # Note here we use `num_experts` which is logical expert count
                self.batched_router_logits = torch.zeros(
                    (moe.max_num_tokens, num_experts),
                    dtype=moe.in_dtype,
                    device=torch.cuda.current_device())

    @property
    def shared_experts(self) -> Optional[torch.nn.Module]:
        return None

    @property
    def tp_size(self):
        return self.moe_parallel_config.tp_size

    @property
    def dp_size(self):
        return self.moe_parallel_config.dp_size

    @property
    def ep_size(self):
        return self.moe_parallel_config.ep_size

    @property
    def tp_rank(self):
        return self.moe_parallel_config.tp_rank

    @property
    def dp_rank(self):
        return self.moe_parallel_config.dp_rank

    @property
    def ep_rank(self):
        return self.moe_parallel_config.ep_rank

    @property
    def use_ep(self):
        return self.moe_parallel_config.use_ep

    @property
    def use_pplx_kernels(self):
        return self.moe_parallel_config.use_pplx_kernels

    @property
    def use_deepep_ht_kernels(self):
        return self.moe_parallel_config.use_deepep_ht_kernels

    @property
    def use_deepep_ll_kernels(self):
        return self.moe_parallel_config.use_deepep_ll_kernels

    @property
    def use_flashinfer_cutlass_kernels(self):
        return (self.moe_quant_config is not None
                and self.moe_quant_config.quant_dtype == "nvfp4"
                and self.moe_config.use_flashinfer_cutlass_kernels)

    def update_expert_map(self):
        # ep_size and ep_rank should already be updated
        assert self.expert_map is not None
        with self.expert_map.device:
            local_num_experts, expert_map = determine_expert_map(
                ep_size=self.ep_size,
                ep_rank=self.ep_rank,
                global_num_experts=self.global_num_experts)
            self.local_num_experts = local_num_experts
            self.register_buffer("expert_map", expert_map)

    def _load_per_tensor_weight_scale(self, shard_id: str,
                                      param: torch.nn.Parameter,
                                      loaded_weight: torch.Tensor,
                                      expert_id: int):
        param_data = param.data
        # for per tensor weight quantization
        if shard_id in ("w1", "w3"):
            # We have to keep the weight scales of w1 and w3 because
            # we need to re-quantize w1/w3 weights after weight loading.
            idx = 0 if shard_id == "w1" else 1
            param_data[expert_id][idx] = loaded_weight
        # If we are in the row parallel case (down_proj)
        elif shard_id == "w2":
            param_data[expert_id] = loaded_weight

    def _load_combined_w13_weight_scale(self, shard_dim: int,
                                        loaded_weight: torch.Tensor,
                                        param: torch.Tensor, tp_rank: int):
        """
        Load w13 weight scales assuming that w1 weight scales and w3 weight
        scales are stored in the same loaded_weight tensor.
        """
        shard_size = param.shape[shard_dim]
        loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,
                                             shard_size)
        param.copy_(loaded_weight)

    def _load_model_weight_or_group_weight_scale(self,
                                                 shard_dim: int,
                                                 expert_data: torch.Tensor,
                                                 shard_id: str,
                                                 loaded_weight: torch.Tensor,
                                                 tp_rank: int,
                                                 load_full_w2: bool = False):
        """
        Load grouped weight scales for group quantization or model weights
            :param shard_dim: dimension to shard
            :param expert_data: parameter for a particular expert
            :param shard_id: either w1, w2, or w3
            :param loaded_weight: checkpoint weight to load into the param
            :param tp_rank: tensor parallel rank
            :param load_full_w2: whether or not the w2 loaded should be sharded.
        """
        if shard_id == "w2":
            # In the case where we have actorder/g_idx, we do not partition the
            # w2 scales, as indicated by `load_full` argument, for all tp cases
            self._load_w2(shard_dim=shard_dim,
                          loaded_weight=loaded_weight,
                          expert_data=expert_data,
                          tp_rank=tp_rank,
                          load_full=load_full_w2)
        elif shard_id in ("w1", "w3"):
            self._load_w13(shard_id=shard_id,
                           shard_dim=shard_dim,
                           loaded_weight=loaded_weight,
                           expert_data=expert_data,
                           tp_rank=tp_rank)

    def _load_per_channel_weight_scale(self, expert_data: torch.Tensor,
                                       shard_dim: int, shard_id: str,
                                       loaded_weight: torch.Tensor,
                                       tp_rank: int):
        # for per channel weight quantization
        if shard_id == "w2":
            expert_data.copy_(loaded_weight)
        elif shard_id in ("w1", "w3"):
            self._load_w13(shard_id=shard_id,
                           shard_dim=shard_dim,
                           loaded_weight=loaded_weight,
                           expert_data=expert_data,
                           tp_rank=tp_rank)

    def _load_w13(self,
                  expert_data: torch.Tensor,
                  shard_dim: int,
                  shard_id: str,
                  loaded_weight: torch.Tensor,
                  tp_rank: int,
                  load_full: bool = False):

        # Index the loaded weight for tp sharding.
        # gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim
        shard_size = expert_data.shape[shard_dim] // 2
        if not load_full:
            loaded_weight = loaded_weight.narrow(shard_dim,
                                                 shard_size * tp_rank,
                                                 shard_size)
        # Narrow parameter and load.
        # w1, gate_proj: Load into first logical weight of w13.
        if shard_id == "w1":
            expert_data = expert_data.narrow(shard_dim, 0, shard_size)
        # w3, up_proj: Load into second logical weight of w13.
        else:
            assert shard_id == "w3"
            expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
        expert_data.copy_(loaded_weight)

    def _load_w2(self,
                 expert_data: torch.Tensor,
                 shard_dim: int,
                 loaded_weight: torch.Tensor,
                 tp_rank: int,
                 load_full: bool = False):

        # Index the loaded weight for tp sharding.
        # down_proj: "RowParallel" so tp sharding on input_dim
        # Narrow parameter and load.
        shard_size = expert_data.shape[shard_dim]
        if not load_full:
            loaded_weight = loaded_weight.narrow(shard_dim,
                                                 shard_size * tp_rank,
                                                 shard_size)
        # w2, down_proj: Load into only logical weight of w2.
        expert_data.copy_(loaded_weight)

    def _load_single_value(self, param: torch.nn.Parameter,
                           loaded_weight: torch.Tensor, expert_id: int):
        param_data = param.data

        # Input scales can be loaded directly and should be equal.
        param_data[expert_id] = loaded_weight

    def _load_g_idx(self, shard_id: str, expert_data: torch.Tensor,
                    shard_dim: int, loaded_weight: torch.Tensor, tp_rank: int):

        if shard_id == "w2":
            self._load_w2(shard_dim=shard_dim,
                          loaded_weight=loaded_weight,
                          expert_data=expert_data,
                          tp_rank=tp_rank)
        else:
            assert shard_id in ("w1", "w3")
            expert_data.copy_(loaded_weight)

    def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
        if self.expert_map is None:
            return expert_id
        return self.expert_map[expert_id].item()

    @overload
    def weight_loader(self, param: torch.nn.Parameter,
                      loaded_weight: torch.Tensor, weight_name: str,
                      shard_id: str, expert_id: int,
                      return_success: Literal[False]) -> None:
        ...

    @overload
    def weight_loader(self, param: torch.nn.Parameter,
                      loaded_weight: torch.Tensor, weight_name: str,
                      shard_id: str, expert_id: int,
                      return_success: Literal[True]) -> bool:
        ...

    def weight_loader(self,
                      param: torch.nn.Parameter,
                      loaded_weight: torch.Tensor,
                      weight_name: str,
                      shard_id: str,
                      expert_id: int,
                      return_success: bool = False) -> Optional[bool]:

        if self.quant_config and self.quant_config.get_name() == "mxfp4":
            # (FIXME) for gpt-oss all experts are combined
            if "bias" in weight_name:
                dim1 = loaded_weight.shape[1]
                param.data[:, :dim1].copy_(loaded_weight)
            else:
                dim1 = loaded_weight.shape[1]
                dim2 = loaded_weight.shape[2]
                param.data[:, :dim1, :dim2].copy_(loaded_weight)
            return True if return_success else None

        expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
        if expert_id == -1:
            # Failed to load this param since it's not local to this rank
            return False if return_success else None
        # Hereafter, `expert_id` is local physical id

        quant_method_name = self.quant_method.__class__.__name__
        # compressed-tensors checkpoints with packed weights are stored flipped
        # TODO (mgoin): check self.quant_method.quant_config.quant_format
        # against known CompressionFormat enum values that have this quality
        if self.quant_method.__class__.__name__ in (
                "CompressedTensorsWNA16MarlinMoEMethod",
                "CompressedTensorsWNA16MoEMethod"):
            loaded_weight = loaded_weight.t().contiguous()

        if shard_id not in ("w1", "w2", "w3"):
            raise ValueError(f"shard_id must be ['w1','w2','w3'] but "
                             f"got {shard_id}.")

        # Fetch the dim to shard the parameter/loaded weight
        # based on the shard id. This will be whatever
        # dimension intermediate_size_per_partition is used.
        SHARD_ID_TO_SHARDED_DIM = {"w1": 0, "w2": 1, "w3": 0}

        is_gguf_weight = getattr(param, "is_gguf_weight", False)
        is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
        if is_gguf_weight_type:
            param.weight_type = loaded_weight.item()
            param.data.copy_(loaded_weight)
            return True if return_success else None

        # Case for BitsAndBytes
        use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
        if use_bitsandbytes_4bit:
            shard_dim = 0

            expert_data = param.data[expert_id]
            if shard_id == "w2":
                expert_data.copy_(loaded_weight)
            elif shard_id in ("w1", "w3"):
                # BNB inflight quantization has already sharded the weights
                full_load = True
                self._load_w13(
                    shard_id=shard_id,
                    shard_dim=shard_dim,
                    loaded_weight=loaded_weight,
                    expert_data=expert_data,
                    tp_rank=self.tp_rank,
                    load_full=full_load,
                )
            return True if return_success else None

        # is_transposed: if the dim to shard the weight
        # should be flipped. Required by GPTQ, compressed-tensors
        # should be whatever dimension intermediate_size_per_partition is
        is_transposed = getattr(param, "is_transposed", False)
        shard_dim = SHARD_ID_TO_SHARDED_DIM[shard_id]
        if is_transposed:
            shard_dim = int(not shard_dim)

        full_load = len(loaded_weight.shape) == 3
        if full_load:
            shard_dim += 1

        # Materialize GGUF UninitializedParameter
        if is_gguf_weight and isinstance(param, UninitializedParameter):
            final_shape = list(loaded_weight.shape)
            if shard_id in ["w1", "w3"]:
                final_shape[1] *= 2
            final_shape[shard_dim] = final_shape[shard_dim] // self.tp_size
            param.materialize(final_shape, dtype=loaded_weight.dtype)

        expert_data = param.data if full_load else param.data[expert_id]

        # Case input scale: input_scale loading is only supported for fp8
        if "input_scale" in weight_name:
            # this is needed for compressed-tensors only
            loaded_weight = loaded_weight.to(param.data.device)

            if ("compressed" in quant_method_name.lower()
                    and param.data[expert_id] != 1
                    and (param.data[expert_id] - loaded_weight).abs() > 1e-5):
                raise ValueError(
                    "input_scales of w1 and w3 of a layer "
                    f"must be equal. But got {param.data[expert_id]} "
                    f"vs. {loaded_weight}")

            self._load_single_value(param=param,
                                    loaded_weight=loaded_weight,
                                    expert_id=expert_id)
            return True if return_success else None

        # Case g_idx
        if "g_idx" in weight_name:
            self._load_g_idx(shard_dim=0,
                             shard_id=shard_id,
                             loaded_weight=loaded_weight,
                             expert_data=expert_data,
                             tp_rank=self.tp_rank)
            return True if return_success else None

        # TODO @dsikka: ModelOpt should follow the proper MoE loading pattern
        if "ModelOpt" in quant_method_name:
            # Determine per-tensor weight scale patterns based on variant
            # Use the dedicated method instead of brittle string matching
            uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern(
            )

            # Call _load_per_tensor_weight_scale() to load per-tensor (scalar)
            # weights scales.
            # Input scales are always per-tensor.
            # Weight scales: FP4 uses "weight_scale_2" and FP8 uses
            # "weight_scale" for per-tensor scales.
            is_per_tensor = ("weight_scale_2" in weight_name
                             if uses_weight_scale_2 else "weight_scale"
                             in weight_name) or "input_scale" in weight_name
            if is_per_tensor:
                self._load_per_tensor_weight_scale(
                    shard_id=shard_id,
                    param=param,
                    loaded_weight=loaded_weight,
                    expert_id=expert_id,
                )
                return True if return_success else None

            # If the weight is w13_weight_scale and w13_weight_scales are
            # combined into single loaded_weight, call
            # _load_combined_w13_weight_scale() to load it.
            # This is checked by comparing the hidden_out dims of the
            # loaded_weight and the param.
            if "w13_weight_scale" in weight_name:
                loaded_weight_hidden_out = loaded_weight.shape[-2]
                param_hidden_out = param.data.shape[-2] * self.tp_size
                if loaded_weight_hidden_out == param_hidden_out:
                    self._load_combined_w13_weight_scale(
                        shard_dim=shard_dim,
                        loaded_weight=loaded_weight,
                        param=param,
                        tp_rank=self.tp_rank,
                    )
                    return True if return_success else None

            # For other weights, call _load_model_weight_or_group_weight_scale()
            # to load it.
            if "weight" in weight_name:
                self._load_model_weight_or_group_weight_scale(
                    shard_id=shard_id,
                    shard_dim=shard_dim,
                    loaded_weight=loaded_weight,
                    expert_data=expert_data,
                    tp_rank=self.tp_rank)
            return True if return_success else None

        # Case weight scales, zero_points and offset, weight/input global scales
        if ("scale" in weight_name or "zero" in weight_name
                or "offset" in weight_name):
            # load the weight scales and zp based on the quantization scheme
            # supported weight scales/zp can be found in
            # FusedMoeWeightScaleSupported
            # TODO @dsikka: once hardened, refactor to use vLLM Parameters
            # specific to each case
            quant_method = getattr(param, "quant_method", None)
            if quant_method == FusedMoeWeightScaleSupported.CHANNEL.value:
                self._load_per_channel_weight_scale(
                    shard_id=shard_id,
                    shard_dim=shard_dim,
                    loaded_weight=loaded_weight,
                    expert_data=expert_data,
                    tp_rank=self.tp_rank)
            elif quant_method in [
                    FusedMoeWeightScaleSupported.GROUP.value,
                    FusedMoeWeightScaleSupported.BLOCK.value,
            ]:
                self._load_model_weight_or_group_weight_scale(
                    shard_id=shard_id,
                    shard_dim=shard_dim,
                    loaded_weight=loaded_weight,
                    expert_data=expert_data,
                    tp_rank=self.tp_rank,
                    load_full_w2=getattr(param, "load_full_w2", False))
            elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:
                self._load_per_tensor_weight_scale(shard_id=shard_id,
                                                   param=param,
                                                   loaded_weight=loaded_weight,
                                                   expert_id=expert_id)
            else:
                WEIGHT_SCALE_SUPPORTED = [
                    e.value for e in FusedMoeWeightScaleSupported
                ]
                raise ValueError(
                    f"quant method must be one of {WEIGHT_SCALE_SUPPORTED}")
            return True if return_success else None

        # Case weight_shape
        if "weight_shape" in weight_name:
            # only required by compressed-tensors
            self._load_single_value(param=param,
                                    loaded_weight=loaded_weight,
                                    expert_id=expert_id)
            return True if return_success else None

        # Case model weights
        if "weight" in weight_name:
            self._load_model_weight_or_group_weight_scale(
                shard_id=shard_id,
                shard_dim=shard_dim,
                loaded_weight=loaded_weight,
                expert_data=expert_data,
                tp_rank=self.tp_rank)
            return True if return_success else None

        return False if return_success else None

    def get_expert_weights(self) -> Iterable[torch.Tensor]:
        weights = list(self.named_parameters())
        assert all(weight.is_contiguous() for _, weight in weights)

        # Filter out the non-expert weights.
        # `e_score_correction_bias` is a bias for each logical expert,
        # with shape (num_logical_experts,), not an expert weight.
        NON_EXPERT_WEIGHTS = {
            "e_score_correction_bias",
        }

        return [
            weight.view(self.local_num_experts, -1) for name, weight in weights
            if name not in NON_EXPERT_WEIGHTS and weight.shape != torch.Size(
                []) and not name.startswith("_shared_experts.")
        ]

    def set_eplb_state(
        self,
        moe_layer_idx: int,
        expert_load_view: torch.Tensor,
        logical_to_physical_map: torch.Tensor,
        logical_replica_count: torch.Tensor,
    ) -> None:
        """
        Register the EPLB state in this layer.

        This is used later in forward pass, where we get the expert mapping
        and record the load metrics in `expert_load_view`.
        """
        self.expert_load_view = expert_load_view[moe_layer_idx]
        self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
        self.logical_replica_count = logical_replica_count[moe_layer_idx]

    def ensure_moe_quant_config(self):
        if self.quant_method.moe_quant_config is None:
            self.quant_method.moe_quant_config = (
                self.quant_method.get_fused_moe_quant_config(self))

    @staticmethod
    def select_experts(
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        top_k: int,
        use_grouped_topk: bool,
        renormalize: bool,
        topk_group: Optional[int] = None,
        num_expert_group: Optional[int] = None,
        custom_routing_function: Optional[Callable] = None,
        scoring_func: str = "softmax",
        routed_scaling_factor: float = 1.0,
        e_score_correction_bias: Optional[torch.Tensor] = None,
        indices_type: Optional[torch.dtype] = None,
        enable_eplb: bool = False,
        expert_map: Optional[torch.Tensor] = None,
        expert_load_view: Optional[torch.Tensor] = None,
        logical_to_physical_map: Optional[torch.Tensor] = None,
        logical_replica_count: Optional[torch.Tensor] = None,
        global_num_experts: Optional[int] = None,
        zero_expert_num: Optional[int] = None,
        zero_expert_type: Optional[str] = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Route the input hidden states to the top-k experts based on the
        router logits.

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

            **Compatibility**: When EPLB is not enabled, the returned ids are
            equivalent to global logical ids, so should be compatible with
            plain MoE implementations without redundant experts.
        """
        from vllm.model_executor.layers.fused_moe.fused_moe import (
            fused_topk, fused_topk_bias)

        # Check if we should use a routing simulation strategy
        routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
        if routing_strategy != "":
            topk_weights, topk_ids = RoutingSimulator.simulate_routing(
                hidden_states=hidden_states,
                router_logits=router_logits,
                strategy_name=routing_strategy,
                top_k=top_k,
                indices_type=indices_type)

        # DeepSeekv2 uses grouped_top_k
        if use_grouped_topk:
            assert topk_group is not None
            assert num_expert_group is not None
            topk_weights, topk_ids = grouped_topk(
                hidden_states=hidden_states,
                gating_output=router_logits,
                topk=top_k,
                renormalize=renormalize,
                num_expert_group=num_expert_group,
                topk_group=topk_group,
                scoring_func=scoring_func,
                routed_scaling_factor=routed_scaling_factor,
                e_score_correction_bias=e_score_correction_bias)
            if indices_type is not None:
                topk_ids = topk_ids.to(dtype=indices_type)
        elif e_score_correction_bias is not None:
            topk_weights, topk_ids = fused_topk_bias(
                hidden_states=hidden_states,
                gating_output=router_logits,
                e_score_correction_bias=e_score_correction_bias.data,
                topk=top_k,
                renormalize=renormalize,
            )
            if routed_scaling_factor is not None:
                topk_weights *= routed_scaling_factor
        elif custom_routing_function is None:
            topk_weights, topk_ids, token_expert_indices = fused_topk(
                hidden_states=hidden_states,
                gating_output=router_logits,
                topk=top_k,
                renormalize=renormalize,
                indices_type=indices_type,
            )
        else:
            topk_weights, topk_ids = custom_routing_function(
                hidden_states=hidden_states,
                gating_output=router_logits,
                topk=top_k,
                renormalize=renormalize)
            if indices_type is not None:
                topk_ids = topk_ids.to(dtype=indices_type)

        if enable_eplb:
            assert expert_load_view is not None
            assert logical_to_physical_map is not None
            assert logical_replica_count is not None

            topk_ids = eplb_map_to_physical_and_record(
                topk_ids=topk_ids,
                expert_load_view=expert_load_view,
                logical_to_physical_map=logical_to_physical_map,
                logical_replica_count=logical_replica_count,
                indices_type=indices_type,
            )

        assert topk_ids.dtype == indices_type or indices_type is None

        # Compute zero expert result if needed
        if (zero_expert_num is not None and zero_expert_num > 0
                and zero_expert_type is not None
                and global_num_experts is not None):
            zero_expert_result = zero_experts_compute_triton(
                expert_indices=topk_ids,
                expert_scales=topk_weights,
                num_experts=global_num_experts,
                zero_expert_type=zero_expert_type,
                hidden_states=hidden_states,
            )
        else:
            zero_expert_result = None
        return topk_weights, topk_ids, zero_expert_result

    def must_reduce_shared_expert_outputs(self) -> bool:
        """
        The shared_experts are typically computed using the RowParallelLinear
        layer. The result of this function is typically used as
        the reduce_results argument to the module.
        When just tensor-parallel is used, it is not required to reduce
        the shared_experts results immediately. Instead we reduce at the
        once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
        With EP and all2all kernels - this is no longer viable as all
        GPU ranks in DP, produce the complete set of hidden_states.
        Therefore it is required that we reduce the shared_experts output
        early.
        """
        return (self.use_pplx_kernels or self.use_deepep_ht_kernels
                or self.use_deepep_ll_kernels)

    def maybe_all_reduce_tensor_model_parallel(
            self, final_hidden_states: torch.Tensor):
        """
        The pplx combine kernel reduces across GPU ranks by default.
        """
        if (self.use_pplx_kernels or self.use_deepep_ht_kernels
                or self.use_deepep_ll_kernels):
            return final_hidden_states
        else:
            return tensor_model_parallel_all_reduce(final_hidden_states)

    def forward_native(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        og_hidden_states = hidden_states.shape[-1]
        if self.hidden_size != og_hidden_states:
            hidden_states = F.pad(hidden_states,
                                  (0, self.hidden_size - og_hidden_states),
                                  mode='constant',
                                  value=0.0)

        if self.shared_experts is None:
            if current_platform.is_tpu():
                # TODO: Once the OOM issue for the TPU backend is resolved, we
                # will switch to using the moe_forward custom op.
                fused_output = self.forward_impl(hidden_states, router_logits)
                assert not isinstance(fused_output, tuple)
            else:
                fused_output = torch.ops.vllm.moe_forward(
                    hidden_states, router_logits, self.layer_name)
            return fused_output[..., :og_hidden_states]
        else:
            if current_platform.is_tpu():
                # TODO: Once the OOM issue for the TPU backend is resolved, we
                # will switch to using the moe_forward custom op.
                shared_output, fused_output = self.forward_impl(
                    hidden_states, router_logits)
            else:
                shared_output, fused_output = torch.ops.vllm.moe_forward_shared(
                    hidden_states, router_logits, self.layer_name)
            return (shared_output[..., :og_hidden_states],
                    fused_output[..., :og_hidden_states])

    def forward_cuda(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        return self.forward_native(hidden_states, router_logits)

    def forward_impl_chunked(
        self,
        full_hidden_states: torch.Tensor,
        full_router_logits: torch.Tensor,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        assert self.batched_hidden_states is not None
        assert self.batched_router_logits is not None
        assert self.batched_hidden_states.dtype == full_hidden_states.dtype
        assert self.batched_router_logits.dtype == full_router_logits.dtype
        # Check size compatibility.
        assert (
            self.batched_hidden_states.size(-1) == full_hidden_states.size(-1))
        assert (
            self.batched_router_logits.size(-1) == full_router_logits.size(-1))

        self.ensure_moe_quant_config()

        full_fused_final_hidden_states = torch.empty_like(full_hidden_states)
        if self.shared_experts is not None:
            full_shared_final_hidden_states = torch.empty_like(
                full_hidden_states)

        def process_chunk(chunk_start, chunk_end, skip_result_store=False):
            chunk_size = chunk_end - chunk_start
            hidden_states = full_hidden_states[chunk_start:chunk_end, :]
            router_logits = full_router_logits[chunk_start:chunk_end, :]

            assert self.batched_hidden_states is not None
            assert self.batched_router_logits is not None
            # This is only true when DBO has been enabled in the config.
            # Both tensors will have an outer dimension for the ubatch id
            if self.batched_hidden_states.dim() == 3:
                assert self.batched_router_logits.dim() == 3
                batch_buffer_idx = dbo_current_ubatch_id()
                batched_hidden_states = self.batched_hidden_states[
                    batch_buffer_idx, :]
                batched_router_logits = self.batched_router_logits[
                    batch_buffer_idx, :]
            else:
                batched_hidden_states = self.batched_hidden_states
                batched_router_logits = self.batched_router_logits

            assert (batched_hidden_states.size(0)  # type: ignore
                    >= chunk_size)
            assert (batched_router_logits.size(0)  # type: ignore 
                    >= chunk_size)
            staged_hidden_states = batched_hidden_states[:
                                                         chunk_size, :]  # type: ignore
            staged_router_logits = batched_router_logits[:
                                                         chunk_size, :]  # type: ignore
            staged_hidden_states.copy_(hidden_states, non_blocking=True)
            staged_router_logits.copy_(router_logits, non_blocking=True)

            # Matrix multiply.
            final_hidden_states = self.quant_method.apply(
                layer=self,
                x=staged_hidden_states,
                router_logits=staged_router_logits,
                top_k=self.top_k,
                renormalize=self.renormalize,
                use_grouped_topk=self.use_grouped_topk,
                global_num_experts=self.global_num_experts,
                expert_map=self.expert_map,
                topk_group=self.topk_group,
                num_expert_group=self.num_expert_group,
                custom_routing_function=self.custom_routing_function,
                scoring_func=self.scoring_func,
                routed_scaling_factor=self.routed_scaling_factor,
                e_score_correction_bias=self.e_score_correction_bias,
                activation=self.activation,
                enable_eplb=self.enable_eplb,
                expert_load_view=self.expert_load_view,
                logical_to_physical_map=self.logical_to_physical_map,
                logical_replica_count=self.logical_replica_count,
            )

            assert self.shared_experts is None or isinstance(
                final_hidden_states, tuple)

            if self.zero_expert_num is not None and self.zero_expert_num > 0:
                assert isinstance(final_hidden_states, tuple)
                assert self.shared_experts is None
                final_hidden_states, zero_expert_result = final_hidden_states
                if zero_expert_result is not None:
                    final_hidden_states += zero_expert_result

            if not skip_result_store:
                if self.shared_experts is None:
                    full_fused_final_hidden_states[
                        chunk_start:chunk_end, :].copy_(final_hidden_states,
                                                        non_blocking=True)
                else:
                    full_shared_final_hidden_states[
                        chunk_start:chunk_end, :].copy_(final_hidden_states[0],
                                                        non_blocking=True)
                    full_fused_final_hidden_states[
                        chunk_start:chunk_end, :].copy_(final_hidden_states[1],
                                                        non_blocking=True)

        ctx = get_forward_context()
        # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
        max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
        moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens

        # If the input to the MoE is sequence parallel then divide by sp_size
        # to find the maximum number of tokens for any individual dispatcher.
        if self.is_sequence_parallel:
            max_tokens_across_dispatchers = cdiv(max_tokens_across_dispatchers,
                                                 self.sp_size)

        num_tokens = full_hidden_states.size(0)
        for chunk_idx, chunk_start_ in enumerate(
                range(0, max_tokens_across_dispatchers,
                      moe_dp_chunk_size_per_rank)):
            chunk_start = chunk_start_
            chunk_end = min(chunk_start + moe_dp_chunk_size_per_rank,
                            max_tokens_across_dispatchers)
            # clamp start and end
            chunk_start = min(chunk_start, num_tokens - 1)
            chunk_end = min(chunk_end, num_tokens)
            with ctx.dp_metadata.chunked_sizes(self.sp_size,
                                               moe_dp_chunk_size_per_rank,
                                               chunk_idx):
                process_chunk(chunk_start,
                              chunk_end,
                              skip_result_store=chunk_start_ >= num_tokens)

        if self.shared_experts is None:
            return full_fused_final_hidden_states
        else:
            return (full_shared_final_hidden_states,
                    full_fused_final_hidden_states)

    def forward_impl(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        assert self.quant_method is not None

        self.ensure_moe_quant_config()

        # Route to the chunked forward path using the FlashInfer Cutlass kernel
        # only when data parallelism (DP) is enabled.
        _use_flashinfer_cutlass_kernels = (self.dp_size > 1 and
                                           self.use_flashinfer_cutlass_kernels)

        if (self.moe_parallel_config.use_pplx_kernels
                or self.moe_parallel_config.use_deepep_ll_kernels
                or _use_flashinfer_cutlass_kernels):
            return self.forward_impl_chunked(hidden_states, router_logits)

        do_naive_dispatch_combine: bool = (
            self.dp_size > 1
            and not self.moe_parallel_config.use_deepep_ht_kernels
            and not self.moe_config.use_flashinfer_cutlass_kernels)

        # If there are shared experts but we are not using a modular kernel, the
        # shared experts must be called here
        if (not isinstance(self.quant_method.fused_experts,
                           FusedMoEModularKernel)
                and self.shared_experts is not None):
            shared_output = self.shared_experts(hidden_states)
        else:
            shared_output = None

        ctx = get_forward_context()
        sp_ctx = ctx.dp_metadata.sp_local_sizes(
            self.sp_size) if ctx.dp_metadata else nullcontext()

        with sp_ctx:
            if do_naive_dispatch_combine:
                hidden_states, router_logits = get_ep_group().dispatch(
                    hidden_states, router_logits, self.is_sequence_parallel)

            # Matrix multiply.
            final_hidden_states = self.quant_method.apply(
                layer=self,
                x=hidden_states,
                router_logits=router_logits,
                top_k=self.top_k,
                renormalize=self.renormalize,
                use_grouped_topk=self.use_grouped_topk,
                global_num_experts=self.global_num_experts,
                expert_map=self.expert_map,
                topk_group=self.topk_group,
                num_expert_group=self.num_expert_group,
                custom_routing_function=self.custom_routing_function,
                scoring_func=self.scoring_func,
                routed_scaling_factor=self.routed_scaling_factor,
                e_score_correction_bias=self.e_score_correction_bias,
                activation=self.activation,
                apply_router_weight_on_input=self.apply_router_weight_on_input,
                enable_eplb=self.enable_eplb,
                expert_load_view=self.expert_load_view,
                logical_to_physical_map=self.logical_to_physical_map,
                logical_replica_count=self.logical_replica_count,
            )

            if shared_output is not None:
                assert not isinstance(final_hidden_states, tuple)
                assert self.shared_experts is not None
                final_hidden_states = (
                    shared_output,
                    final_hidden_states,
                )
            elif self.zero_expert_num is not None and self.zero_expert_num > 0:
                assert isinstance(final_hidden_states, tuple)
                final_hidden_states, zero_expert_result = final_hidden_states

            def reduce_output(states: torch.Tensor,
                              do_combine: bool = True) -> torch.Tensor:
                if do_naive_dispatch_combine and do_combine:
                    states = get_ep_group().combine(states,
                                                    self.is_sequence_parallel)

                if (not self.is_sequence_parallel and self.reduce_results
                        and (self.tp_size > 1 or self.ep_size > 1)):
                    states = self.maybe_all_reduce_tensor_model_parallel(
                        states)

                return states

            if self.shared_experts is not None:
                return (
                    reduce_output(final_hidden_states[0], do_combine=False),
                    reduce_output(final_hidden_states[1]),
                )
            elif self.zero_expert_num is not None and self.zero_expert_num > 0:
                assert isinstance(final_hidden_states, torch.Tensor)
                return reduce_output(final_hidden_states) + zero_expert_result
            else:
                return reduce_output(final_hidden_states)

    @classmethod
    def make_expert_params_mapping(
            cls,
            ckpt_gate_proj_name: str,
            ckpt_down_proj_name: str,
            ckpt_up_proj_name: str,
            num_experts: int,
            num_redundant_experts: int = 0) -> list[tuple[str, str, int, str]]:

        num_physical_experts = num_experts + num_redundant_experts

        # In the returned mapping:
        # - `expert_id` is the physical expert id
        # - `weight_name` contains the weight name of the logical expert
        # So that we should map the expert id to logical in `weight_name`
        physical_to_logical_map = \
            EplbState.build_initial_global_physical_to_logical_map(
            num_experts, num_redundant_experts)

        return [
            # (param_name, weight_name, expert_id, shard_id)
            ("experts.w13_" if weight_name
             in [ckpt_gate_proj_name, ckpt_up_proj_name] else "experts.w2_",
             f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.",
             expert_id, shard_id) for expert_id in range(num_physical_experts)
            for shard_id, weight_name in [
                ("w1", ckpt_gate_proj_name),
                ("w2", ckpt_down_proj_name),
                ("w3", ckpt_up_proj_name),
            ]
        ]

    def extra_repr(self) -> str:

        s = (
            f"global_num_experts={self.global_num_experts}, "
            f"local_num_experts={self.local_num_experts}, "
            f"top_k={self.top_k}, "
            f"intermediate_size_per_partition={self.intermediate_size_per_partition}, "  # noqa: E501
            f"tp_size={self.tp_size},\n"
            f"ep_size={self.ep_size}, "
            f"reduce_results={self.reduce_results}, "
            f"renormalize={self.renormalize}, "
            f"use_grouped_topk={self.use_grouped_topk}")

        if self.use_grouped_topk:
            s += f", num_expert_group={self.num_expert_group}, topk_group={self.topk_group}"  # noqa: E501

        s += f", scoring_func='{self.scoring_func}', activation='{self.activation}'"  # noqa: E501

        return s

activation `instance-attribute` ¶

activation = activation

apply_router_weight_on_input `instance-attribute` ¶

apply_router_weight_on_input = apply_router_weight_on_input

batched_hidden_states `instance-attribute` ¶

batched_hidden_states: Optional[Tensor] = None

batched_router_logits `instance-attribute` ¶

batched_router_logits: Optional[Tensor] = None

custom_routing_function `instance-attribute` ¶

custom_routing_function = custom_routing_function

dp_rank `property` ¶

dp_rank

dp_size `property` ¶

dp_size

e_score_correction_bias `instance-attribute` ¶

e_score_correction_bias = e_score_correction_bias

enable_eplb `instance-attribute` ¶

enable_eplb = enable_eplb

ep_rank `property` ¶

ep_rank

ep_size `property` ¶

ep_size

expert_load_view `instance-attribute` ¶

expert_load_view: Optional[Tensor] = None

expert_map `instance-attribute` ¶

expert_map: Optional[Tensor]

global_num_experts `instance-attribute` ¶

global_num_experts = num_experts + num_redundant_experts

hidden_size `instance-attribute` ¶

hidden_size = hidden_size

intermediate_size_per_partition `instance-attribute` ¶

intermediate_size_per_partition = (
    intermediate_size // tp_size
)

is_sequence_parallel `instance-attribute` ¶

is_sequence_parallel = is_sequence_parallel

layer_name `instance-attribute` ¶

layer_name = prefix

local_num_experts `instance-attribute` ¶

local_num_experts = local_num_experts

logical_replica_count `instance-attribute` ¶

logical_replica_count: Optional[Tensor] = None

logical_to_physical_map `instance-attribute` ¶

logical_to_physical_map: Optional[Tensor] = None

moe_config `instance-attribute` ¶

moe_config = moe

moe_parallel_config `instance-attribute` ¶

moe_parallel_config: FusedMoEParallelConfig = make(
    tp_size_=tp_size_,
    dp_size_=dp_size_,
    vllm_parallel_config=parallel_config,
)

moe_quant_config `instance-attribute` ¶

moe_quant_config: Optional[FusedMoEQuantConfig] = None

num_expert_group `instance-attribute` ¶

num_expert_group = num_expert_group

params_dtype `instance-attribute` ¶

params_dtype = params_dtype

quant_config `instance-attribute` ¶

quant_config = quant_config

quant_method `instance-attribute` ¶

quant_method = quant_method

reduce_results `instance-attribute` ¶

reduce_results = reduce_results

renormalize `instance-attribute` ¶

renormalize = renormalize

routed_scaling_factor `instance-attribute` ¶

routed_scaling_factor = routed_scaling_factor

scoring_func `instance-attribute` ¶

scoring_func = scoring_func

shared_experts `property` ¶

shared_experts: Optional[Module]

sp_size `instance-attribute` ¶

sp_size = tp_size_ if is_sequence_parallel else 1

top_k `instance-attribute` ¶

top_k = top_k

topk_group `instance-attribute` ¶

topk_group = topk_group

tp_rank `property` ¶

tp_rank

tp_size `property` ¶

tp_size

use_deepep_ht_kernels `property` ¶

use_deepep_ht_kernels

use_deepep_ll_kernels `property` ¶

use_deepep_ll_kernels

use_ep `property` ¶

use_ep

use_flashinfer_cutlass_kernels `property` ¶

use_flashinfer_cutlass_kernels

use_grouped_topk `instance-attribute` ¶

use_grouped_topk = use_grouped_topk

use_pplx_kernels `property` ¶

use_pplx_kernels

zero_expert_num `instance-attribute` ¶

zero_expert_num = zero_expert_num

zero_expert_type `instance-attribute` ¶

zero_expert_type = zero_expert_type

init ¶

__init__(
    num_experts: int,
    top_k: int,
    hidden_size: int,
    intermediate_size: int,
    params_dtype: Optional[dtype] = None,
    reduce_results: bool = False,
    renormalize: bool = True,
    use_grouped_topk: bool = False,
    num_expert_group: Optional[int] = None,
    topk_group: Optional[int] = None,
    quant_config: Optional[QuantizationConfig] = None,
    tp_size: Optional[int] = None,
    ep_size: Optional[int] = None,
    dp_size: Optional[int] = None,
    prefix: str = "",
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[Tensor] = None,
    apply_router_weight_on_input: bool = False,
    activation: str = "silu",
    enable_eplb: bool = False,
    num_redundant_experts: int = 0,
    has_bias: bool = False,
    is_sequence_parallel=False,
    zero_expert_num: Optional[int] = 0,
    zero_expert_type: Optional[str] = None,
)

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

def __init__(
    self,
    num_experts: int,  # Global number of experts
    top_k: int,
    hidden_size: int,
    intermediate_size: int,
    params_dtype: Optional[torch.dtype] = None,
    reduce_results: bool = False,
    renormalize: bool = True,
    use_grouped_topk: bool = False,
    num_expert_group: Optional[int] = None,
    topk_group: Optional[int] = None,
    quant_config: Optional[QuantizationConfig] = None,
    tp_size: Optional[int] = None,
    ep_size: Optional[int] = None,
    dp_size: Optional[int] = None,
    prefix: str = "",
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[torch.Tensor] = None,
    apply_router_weight_on_input: bool = False,
    activation: str = "silu",
    enable_eplb: bool = False,
    num_redundant_experts: int = 0,
    has_bias: bool = False,
    is_sequence_parallel=False,
    zero_expert_num: Optional[int] = 0,
    zero_expert_type: Optional[str] = None,
):
    super().__init__()
    if params_dtype is None:
        params_dtype = torch.get_default_dtype()
    self.params_dtype = params_dtype

    vllm_config = get_current_vllm_config()

    # FIXME (varun): We should have a better way of inferring the activation
    # datatype. This works for now as the tensor datatype entering the MoE
    # operation is typically unquantized (i.e. float16/bfloat16).
    if vllm_config.model_config is not None:
        moe_in_dtype = vllm_config.model_config.dtype
    else:
        # TODO (bnell): This is a hack to get test_mixtral_moe to work
        # since model_config is not set in the pytest test.
        moe_in_dtype = params_dtype

    tp_size_ = (tp_size if tp_size is not None else
                get_tensor_model_parallel_world_size())
    dp_size_ = (dp_size
                if dp_size is not None else get_dp_group().world_size)

    self.is_sequence_parallel = is_sequence_parallel
    self.sp_size = tp_size_ if is_sequence_parallel else 1

    self.moe_parallel_config: FusedMoEParallelConfig = (
        FusedMoEParallelConfig.make(
            tp_size_=tp_size_,
            dp_size_=dp_size_,
            vllm_parallel_config=vllm_config.parallel_config))

    self.global_num_experts = num_experts + num_redundant_experts
    self.zero_expert_num = zero_expert_num
    self.zero_expert_type = zero_expert_type

    # Round up hidden size if needed.
    hidden_size = maybe_roundup_hidden_size(hidden_size, moe_in_dtype,
                                            quant_config,
                                            self.moe_parallel_config)

    # For smuggling this layer into the fused moe custom op
    compilation_config = vllm_config.compilation_config
    if prefix in compilation_config.static_forward_context:
        raise ValueError("Duplicate layer name: {}".format(prefix))
    compilation_config.static_forward_context[prefix] = self
    self.layer_name = prefix

    self.enable_eplb = enable_eplb
    self.expert_load_view: Optional[torch.Tensor] = None
    self.logical_to_physical_map: Optional[torch.Tensor] = None
    self.logical_replica_count: Optional[torch.Tensor] = None

    # Determine expert maps
    if self.use_ep:
        if self.enable_eplb:
            assert self.global_num_experts % self.ep_size == 0, \
                "EPLB currently only supports even distribution of " \
                "experts across ranks."
        else:
            assert num_redundant_experts == 0, \
                "Redundant experts are only supported with EPLB."

        expert_placement_strategy = (
            vllm_config.parallel_config.expert_placement_strategy)
        if expert_placement_strategy == "round_robin":
            # TODO(Bruce): will support round robin expert placement with
            # EPLB enabled in the future.
            round_robin_supported = ((num_expert_group is not None
                                      and num_expert_group > 1)
                                     and num_redundant_experts == 0
                                     and not self.enable_eplb)

            if not round_robin_supported:
                logger.warning(
                    "Round-robin expert placement is only supported for "
                    "models with multiple expert groups and no redundant "
                    "experts. Falling back to linear expert placement.")
                expert_placement_strategy = "linear"

        self.expert_map: Optional[torch.Tensor]
        local_num_experts, expert_map = determine_expert_map(
            ep_size=self.ep_size,
            ep_rank=self.ep_rank,
            global_num_experts=self.global_num_experts,
            expert_placement_strategy=expert_placement_strategy,
        )
        self.local_num_experts = local_num_experts
        self.register_buffer("expert_map", expert_map)
        logger.info_once(
            "[EP Rank %s/%s] Expert parallelism is enabled. Expert "
            "placement strategy: %s. Local/global"
            " number of experts: %s/%s. Experts local to global index map:"
            " %s.", self.ep_rank, self.ep_size, expert_placement_strategy,
            self.local_num_experts, self.global_num_experts,
            get_compressed_expert_map(self.expert_map))
    else:
        self.local_num_experts, self.expert_map = (self.global_num_experts,
                                                   None)

    self.top_k = top_k

    assert intermediate_size % self.tp_size == 0
    self.hidden_size = hidden_size
    self.intermediate_size_per_partition = intermediate_size // self.tp_size
    self.reduce_results = reduce_results
    self.renormalize = renormalize
    self.use_grouped_topk = use_grouped_topk
    if self.use_grouped_topk:
        assert num_expert_group is not None and topk_group is not None
    self.num_expert_group = num_expert_group
    self.topk_group = topk_group
    self.custom_routing_function = custom_routing_function
    self.scoring_func = scoring_func
    self.routed_scaling_factor = routed_scaling_factor
    self.e_score_correction_bias = e_score_correction_bias
    self.apply_router_weight_on_input = apply_router_weight_on_input
    self.activation = activation

    if self.scoring_func != "softmax" and not self.use_grouped_topk:
        raise ValueError("Only softmax scoring function is supported for "
                         "non-grouped topk.")

    moe = FusedMoEConfig(
        num_experts=self.global_num_experts,
        experts_per_token=top_k,
        hidden_dim=hidden_size,
        num_local_experts=self.local_num_experts,
        moe_parallel_config=self.moe_parallel_config,
        in_dtype=moe_in_dtype,
        max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
        has_bias=has_bias,
    )
    self.moe_config = moe
    self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
    self.quant_config = quant_config

    # Note: get_quant_method will look at the layer's local_num_experts
    # for heuristic purposes, so it must be initialized first.
    quant_method: Optional[QuantizeMethodBase] = None
    quant_method = (UnquantizedFusedMoEMethod(moe) if quant_config is None
                    else quant_config.get_quant_method(self, prefix))

    assert quant_method is not None
    assert isinstance(quant_method, FusedMoEMethodBase)
    self.quant_method = quant_method

    if self.enable_eplb:
        from vllm.model_executor.layers.quantization.fp8 import (
            Fp8MoEMethod)
        if not isinstance(quant_method,
                          (Fp8MoEMethod, UnquantizedFusedMoEMethod)):
            # TODO: Add support for additional quantization methods.
            # The implementation for other quantization methods does not
            # contain essential differences, but the current quant API
            # design causes duplicated work when extending to new
            # quantization methods, so I'm leaving it for now.
            # If you plan to add support for more quantization methods,
            # please refer to the implementation in `Fp8MoEMethod`.
            raise NotImplementedError("EPLB is only supported for FP8 "
                                      "quantization for now.")

    moe_quant_params = {
        "num_experts": self.local_num_experts,
        "hidden_size": hidden_size,
        "intermediate_size_per_partition":
        self.intermediate_size_per_partition,
        "params_dtype": params_dtype,
        "weight_loader": self.weight_loader,
    }
    # need full intermediate size pre-sharding for WNA16 act order
    if (self.quant_method.__class__.__name__
            in ("GPTQMarlinMoEMethod",
                "CompressedTensorsWNA16MarlinMoEMethod",
                "CompressedTensorsWNA16MoEMethod")):
        moe_quant_params["intermediate_size_full"] = intermediate_size

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

    # Chunked all2all staging tensor
    self.batched_hidden_states: Optional[torch.Tensor] = None
    self.batched_router_logits: Optional[torch.Tensor] = None

    # TODO(bnell): flashinfer uses non-batched format.
    # Does it really need a batched buffer?
    if (self.moe_parallel_config.use_pplx_kernels
            or self.moe_parallel_config.use_deepep_ll_kernels
            or self.moe_config.use_flashinfer_cutlass_kernels):
        if vllm_config.parallel_config.enable_dbo:
            self.batched_hidden_states = torch.zeros(
                (2, moe.max_num_tokens, self.hidden_size),
                dtype=moe.in_dtype,
                device=torch.cuda.current_device())

            # Note here we use `num_experts` which is logical expert count
            self.batched_router_logits = torch.zeros(
                (2, moe.max_num_tokens, num_experts),
                dtype=moe.in_dtype,
                device=torch.cuda.current_device())
        else:
            self.batched_hidden_states = torch.zeros(
                (moe.max_num_tokens, self.hidden_size),
                dtype=moe.in_dtype,
                device=torch.cuda.current_device())

            # Note here we use `num_experts` which is logical expert count
            self.batched_router_logits = torch.zeros(
                (moe.max_num_tokens, num_experts),
                dtype=moe.in_dtype,
                device=torch.cuda.current_device())

_load_combined_w13_weight_scale ¶

_load_combined_w13_weight_scale(
    shard_dim: int,
    loaded_weight: Tensor,
    param: Tensor,
    tp_rank: int,
)

Load w13 weight scales assuming that w1 weight scales and w3 weight scales are stored in the same loaded_weight tensor.

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

def _load_combined_w13_weight_scale(self, shard_dim: int,
                                    loaded_weight: torch.Tensor,
                                    param: torch.Tensor, tp_rank: int):
    """
    Load w13 weight scales assuming that w1 weight scales and w3 weight
    scales are stored in the same loaded_weight tensor.
    """
    shard_size = param.shape[shard_dim]
    loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,
                                         shard_size)
    param.copy_(loaded_weight)

_load_g_idx ¶

_load_g_idx(
    shard_id: str,
    expert_data: Tensor,
    shard_dim: int,
    loaded_weight: Tensor,
    tp_rank: int,
)

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

def _load_g_idx(self, shard_id: str, expert_data: torch.Tensor,
                shard_dim: int, loaded_weight: torch.Tensor, tp_rank: int):

    if shard_id == "w2":
        self._load_w2(shard_dim=shard_dim,
                      loaded_weight=loaded_weight,
                      expert_data=expert_data,
                      tp_rank=tp_rank)
    else:
        assert shard_id in ("w1", "w3")
        expert_data.copy_(loaded_weight)

_load_model_weight_or_group_weight_scale ¶

_load_model_weight_or_group_weight_scale(
    shard_dim: int,
    expert_data: Tensor,
    shard_id: str,
    loaded_weight: Tensor,
    tp_rank: int,
    load_full_w2: bool = False,
)

Load grouped weight scales for group quantization or model weights :param shard_dim: dimension to shard :param expert_data: parameter for a particular expert :param shard_id: either w1, w2, or w3 :param loaded_weight: checkpoint weight to load into the param :param tp_rank: tensor parallel rank :param load_full_w2: whether or not the w2 loaded should be sharded.

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

def _load_model_weight_or_group_weight_scale(self,
                                             shard_dim: int,
                                             expert_data: torch.Tensor,
                                             shard_id: str,
                                             loaded_weight: torch.Tensor,
                                             tp_rank: int,
                                             load_full_w2: bool = False):
    """
    Load grouped weight scales for group quantization or model weights
        :param shard_dim: dimension to shard
        :param expert_data: parameter for a particular expert
        :param shard_id: either w1, w2, or w3
        :param loaded_weight: checkpoint weight to load into the param
        :param tp_rank: tensor parallel rank
        :param load_full_w2: whether or not the w2 loaded should be sharded.
    """
    if shard_id == "w2":
        # In the case where we have actorder/g_idx, we do not partition the
        # w2 scales, as indicated by `load_full` argument, for all tp cases
        self._load_w2(shard_dim=shard_dim,
                      loaded_weight=loaded_weight,
                      expert_data=expert_data,
                      tp_rank=tp_rank,
                      load_full=load_full_w2)
    elif shard_id in ("w1", "w3"):
        self._load_w13(shard_id=shard_id,
                       shard_dim=shard_dim,
                       loaded_weight=loaded_weight,
                       expert_data=expert_data,
                       tp_rank=tp_rank)

_load_per_channel_weight_scale ¶

_load_per_channel_weight_scale(
    expert_data: Tensor,
    shard_dim: int,
    shard_id: str,
    loaded_weight: Tensor,
    tp_rank: int,
)

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

def _load_per_channel_weight_scale(self, expert_data: torch.Tensor,
                                   shard_dim: int, shard_id: str,
                                   loaded_weight: torch.Tensor,
                                   tp_rank: int):
    # for per channel weight quantization
    if shard_id == "w2":
        expert_data.copy_(loaded_weight)
    elif shard_id in ("w1", "w3"):
        self._load_w13(shard_id=shard_id,
                       shard_dim=shard_dim,
                       loaded_weight=loaded_weight,
                       expert_data=expert_data,
                       tp_rank=tp_rank)

_load_per_tensor_weight_scale ¶

_load_per_tensor_weight_scale(
    shard_id: str,
    param: Parameter,
    loaded_weight: Tensor,
    expert_id: int,
)

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

def _load_per_tensor_weight_scale(self, shard_id: str,
                                  param: torch.nn.Parameter,
                                  loaded_weight: torch.Tensor,
                                  expert_id: int):
    param_data = param.data
    # for per tensor weight quantization
    if shard_id in ("w1", "w3"):
        # We have to keep the weight scales of w1 and w3 because
        # we need to re-quantize w1/w3 weights after weight loading.
        idx = 0 if shard_id == "w1" else 1
        param_data[expert_id][idx] = loaded_weight
    # If we are in the row parallel case (down_proj)
    elif shard_id == "w2":
        param_data[expert_id] = loaded_weight

_load_single_value ¶

_load_single_value(
    param: Parameter, loaded_weight: Tensor, expert_id: int
)

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

def _load_single_value(self, param: torch.nn.Parameter,
                       loaded_weight: torch.Tensor, expert_id: int):
    param_data = param.data

    # Input scales can be loaded directly and should be equal.
    param_data[expert_id] = loaded_weight

_load_w13 ¶

_load_w13(
    expert_data: Tensor,
    shard_dim: int,
    shard_id: str,
    loaded_weight: Tensor,
    tp_rank: int,
    load_full: bool = False,
)

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

def _load_w13(self,
              expert_data: torch.Tensor,
              shard_dim: int,
              shard_id: str,
              loaded_weight: torch.Tensor,
              tp_rank: int,
              load_full: bool = False):

    # Index the loaded weight for tp sharding.
    # gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim
    shard_size = expert_data.shape[shard_dim] // 2
    if not load_full:
        loaded_weight = loaded_weight.narrow(shard_dim,
                                             shard_size * tp_rank,
                                             shard_size)
    # Narrow parameter and load.
    # w1, gate_proj: Load into first logical weight of w13.
    if shard_id == "w1":
        expert_data = expert_data.narrow(shard_dim, 0, shard_size)
    # w3, up_proj: Load into second logical weight of w13.
    else:
        assert shard_id == "w3"
        expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
    expert_data.copy_(loaded_weight)

_load_w2 ¶

_load_w2(
    expert_data: Tensor,
    shard_dim: int,
    loaded_weight: Tensor,
    tp_rank: int,
    load_full: bool = False,
)

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

def _load_w2(self,
             expert_data: torch.Tensor,
             shard_dim: int,
             loaded_weight: torch.Tensor,
             tp_rank: int,
             load_full: bool = False):

    # Index the loaded weight for tp sharding.
    # down_proj: "RowParallel" so tp sharding on input_dim
    # Narrow parameter and load.
    shard_size = expert_data.shape[shard_dim]
    if not load_full:
        loaded_weight = loaded_weight.narrow(shard_dim,
                                             shard_size * tp_rank,
                                             shard_size)
    # w2, down_proj: Load into only logical weight of w2.
    expert_data.copy_(loaded_weight)

_map_global_expert_id_to_local_expert_id ¶

_map_global_expert_id_to_local_expert_id(
    expert_id: int,
) -> int

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

def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
    if self.expert_map is None:
        return expert_id
    return self.expert_map[expert_id].item()

ensure_moe_quant_config ¶

ensure_moe_quant_config()

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

def ensure_moe_quant_config(self):
    if self.quant_method.moe_quant_config is None:
        self.quant_method.moe_quant_config = (
            self.quant_method.get_fused_moe_quant_config(self))

extra_repr ¶

extra_repr() -> str

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

def extra_repr(self) -> str:

    s = (
        f"global_num_experts={self.global_num_experts}, "
        f"local_num_experts={self.local_num_experts}, "
        f"top_k={self.top_k}, "
        f"intermediate_size_per_partition={self.intermediate_size_per_partition}, "  # noqa: E501
        f"tp_size={self.tp_size},\n"
        f"ep_size={self.ep_size}, "
        f"reduce_results={self.reduce_results}, "
        f"renormalize={self.renormalize}, "
        f"use_grouped_topk={self.use_grouped_topk}")

    if self.use_grouped_topk:
        s += f", num_expert_group={self.num_expert_group}, topk_group={self.topk_group}"  # noqa: E501

    s += f", scoring_func='{self.scoring_func}', activation='{self.activation}'"  # noqa: E501

    return s

forward_cuda ¶

forward_cuda(
    hidden_states: Tensor, router_logits: Tensor
) -> Union[Tensor, tuple[Tensor, Tensor]]

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

def forward_cuda(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    return self.forward_native(hidden_states, router_logits)

forward_impl ¶

forward_impl(
    hidden_states: Tensor, router_logits: Tensor
) -> Union[Tensor, tuple[Tensor, Tensor]]

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

def forward_impl(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    assert self.quant_method is not None

    self.ensure_moe_quant_config()

    # Route to the chunked forward path using the FlashInfer Cutlass kernel
    # only when data parallelism (DP) is enabled.
    _use_flashinfer_cutlass_kernels = (self.dp_size > 1 and
                                       self.use_flashinfer_cutlass_kernels)

    if (self.moe_parallel_config.use_pplx_kernels
            or self.moe_parallel_config.use_deepep_ll_kernels
            or _use_flashinfer_cutlass_kernels):
        return self.forward_impl_chunked(hidden_states, router_logits)

    do_naive_dispatch_combine: bool = (
        self.dp_size > 1
        and not self.moe_parallel_config.use_deepep_ht_kernels
        and not self.moe_config.use_flashinfer_cutlass_kernels)

    # If there are shared experts but we are not using a modular kernel, the
    # shared experts must be called here
    if (not isinstance(self.quant_method.fused_experts,
                       FusedMoEModularKernel)
            and self.shared_experts is not None):
        shared_output = self.shared_experts(hidden_states)
    else:
        shared_output = None

    ctx = get_forward_context()
    sp_ctx = ctx.dp_metadata.sp_local_sizes(
        self.sp_size) if ctx.dp_metadata else nullcontext()

    with sp_ctx:
        if do_naive_dispatch_combine:
            hidden_states, router_logits = get_ep_group().dispatch(
                hidden_states, router_logits, self.is_sequence_parallel)

        # Matrix multiply.
        final_hidden_states = self.quant_method.apply(
            layer=self,
            x=hidden_states,
            router_logits=router_logits,
            top_k=self.top_k,
            renormalize=self.renormalize,
            use_grouped_topk=self.use_grouped_topk,
            global_num_experts=self.global_num_experts,
            expert_map=self.expert_map,
            topk_group=self.topk_group,
            num_expert_group=self.num_expert_group,
            custom_routing_function=self.custom_routing_function,
            scoring_func=self.scoring_func,
            routed_scaling_factor=self.routed_scaling_factor,
            e_score_correction_bias=self.e_score_correction_bias,
            activation=self.activation,
            apply_router_weight_on_input=self.apply_router_weight_on_input,
            enable_eplb=self.enable_eplb,
            expert_load_view=self.expert_load_view,
            logical_to_physical_map=self.logical_to_physical_map,
            logical_replica_count=self.logical_replica_count,
        )

        if shared_output is not None:
            assert not isinstance(final_hidden_states, tuple)
            assert self.shared_experts is not None
            final_hidden_states = (
                shared_output,
                final_hidden_states,
            )
        elif self.zero_expert_num is not None and self.zero_expert_num > 0:
            assert isinstance(final_hidden_states, tuple)
            final_hidden_states, zero_expert_result = final_hidden_states

        def reduce_output(states: torch.Tensor,
                          do_combine: bool = True) -> torch.Tensor:
            if do_naive_dispatch_combine and do_combine:
                states = get_ep_group().combine(states,
                                                self.is_sequence_parallel)

            if (not self.is_sequence_parallel and self.reduce_results
                    and (self.tp_size > 1 or self.ep_size > 1)):
                states = self.maybe_all_reduce_tensor_model_parallel(
                    states)

            return states

        if self.shared_experts is not None:
            return (
                reduce_output(final_hidden_states[0], do_combine=False),
                reduce_output(final_hidden_states[1]),
            )
        elif self.zero_expert_num is not None and self.zero_expert_num > 0:
            assert isinstance(final_hidden_states, torch.Tensor)
            return reduce_output(final_hidden_states) + zero_expert_result
        else:
            return reduce_output(final_hidden_states)

forward_impl_chunked ¶

forward_impl_chunked(
    full_hidden_states: Tensor, full_router_logits: Tensor
) -> Union[Tensor, tuple[Tensor, Tensor]]

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

def forward_impl_chunked(
    self,
    full_hidden_states: torch.Tensor,
    full_router_logits: torch.Tensor,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    assert self.batched_hidden_states is not None
    assert self.batched_router_logits is not None
    assert self.batched_hidden_states.dtype == full_hidden_states.dtype
    assert self.batched_router_logits.dtype == full_router_logits.dtype
    # Check size compatibility.
    assert (
        self.batched_hidden_states.size(-1) == full_hidden_states.size(-1))
    assert (
        self.batched_router_logits.size(-1) == full_router_logits.size(-1))

    self.ensure_moe_quant_config()

    full_fused_final_hidden_states = torch.empty_like(full_hidden_states)
    if self.shared_experts is not None:
        full_shared_final_hidden_states = torch.empty_like(
            full_hidden_states)

    def process_chunk(chunk_start, chunk_end, skip_result_store=False):
        chunk_size = chunk_end - chunk_start
        hidden_states = full_hidden_states[chunk_start:chunk_end, :]
        router_logits = full_router_logits[chunk_start:chunk_end, :]

        assert self.batched_hidden_states is not None
        assert self.batched_router_logits is not None
        # This is only true when DBO has been enabled in the config.
        # Both tensors will have an outer dimension for the ubatch id
        if self.batched_hidden_states.dim() == 3:
            assert self.batched_router_logits.dim() == 3
            batch_buffer_idx = dbo_current_ubatch_id()
            batched_hidden_states = self.batched_hidden_states[
                batch_buffer_idx, :]
            batched_router_logits = self.batched_router_logits[
                batch_buffer_idx, :]
        else:
            batched_hidden_states = self.batched_hidden_states
            batched_router_logits = self.batched_router_logits

        assert (batched_hidden_states.size(0)  # type: ignore
                >= chunk_size)
        assert (batched_router_logits.size(0)  # type: ignore 
                >= chunk_size)
        staged_hidden_states = batched_hidden_states[:
                                                     chunk_size, :]  # type: ignore
        staged_router_logits = batched_router_logits[:
                                                     chunk_size, :]  # type: ignore
        staged_hidden_states.copy_(hidden_states, non_blocking=True)
        staged_router_logits.copy_(router_logits, non_blocking=True)

        # Matrix multiply.
        final_hidden_states = self.quant_method.apply(
            layer=self,
            x=staged_hidden_states,
            router_logits=staged_router_logits,
            top_k=self.top_k,
            renormalize=self.renormalize,
            use_grouped_topk=self.use_grouped_topk,
            global_num_experts=self.global_num_experts,
            expert_map=self.expert_map,
            topk_group=self.topk_group,
            num_expert_group=self.num_expert_group,
            custom_routing_function=self.custom_routing_function,
            scoring_func=self.scoring_func,
            routed_scaling_factor=self.routed_scaling_factor,
            e_score_correction_bias=self.e_score_correction_bias,
            activation=self.activation,
            enable_eplb=self.enable_eplb,
            expert_load_view=self.expert_load_view,
            logical_to_physical_map=self.logical_to_physical_map,
            logical_replica_count=self.logical_replica_count,
        )

        assert self.shared_experts is None or isinstance(
            final_hidden_states, tuple)

        if self.zero_expert_num is not None and self.zero_expert_num > 0:
            assert isinstance(final_hidden_states, tuple)
            assert self.shared_experts is None
            final_hidden_states, zero_expert_result = final_hidden_states
            if zero_expert_result is not None:
                final_hidden_states += zero_expert_result

        if not skip_result_store:
            if self.shared_experts is None:
                full_fused_final_hidden_states[
                    chunk_start:chunk_end, :].copy_(final_hidden_states,
                                                    non_blocking=True)
            else:
                full_shared_final_hidden_states[
                    chunk_start:chunk_end, :].copy_(final_hidden_states[0],
                                                    non_blocking=True)
                full_fused_final_hidden_states[
                    chunk_start:chunk_end, :].copy_(final_hidden_states[1],
                                                    non_blocking=True)

    ctx = get_forward_context()
    # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
    max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
    moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens

    # If the input to the MoE is sequence parallel then divide by sp_size
    # to find the maximum number of tokens for any individual dispatcher.
    if self.is_sequence_parallel:
        max_tokens_across_dispatchers = cdiv(max_tokens_across_dispatchers,
                                             self.sp_size)

    num_tokens = full_hidden_states.size(0)
    for chunk_idx, chunk_start_ in enumerate(
            range(0, max_tokens_across_dispatchers,
                  moe_dp_chunk_size_per_rank)):
        chunk_start = chunk_start_
        chunk_end = min(chunk_start + moe_dp_chunk_size_per_rank,
                        max_tokens_across_dispatchers)
        # clamp start and end
        chunk_start = min(chunk_start, num_tokens - 1)
        chunk_end = min(chunk_end, num_tokens)
        with ctx.dp_metadata.chunked_sizes(self.sp_size,
                                           moe_dp_chunk_size_per_rank,
                                           chunk_idx):
            process_chunk(chunk_start,
                          chunk_end,
                          skip_result_store=chunk_start_ >= num_tokens)

    if self.shared_experts is None:
        return full_fused_final_hidden_states
    else:
        return (full_shared_final_hidden_states,
                full_fused_final_hidden_states)

forward_native ¶

forward_native(
    hidden_states: Tensor, router_logits: Tensor
) -> Union[Tensor, tuple[Tensor, Tensor]]

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

def forward_native(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    og_hidden_states = hidden_states.shape[-1]
    if self.hidden_size != og_hidden_states:
        hidden_states = F.pad(hidden_states,
                              (0, self.hidden_size - og_hidden_states),
                              mode='constant',
                              value=0.0)

    if self.shared_experts is None:
        if current_platform.is_tpu():
            # TODO: Once the OOM issue for the TPU backend is resolved, we
            # will switch to using the moe_forward custom op.
            fused_output = self.forward_impl(hidden_states, router_logits)
            assert not isinstance(fused_output, tuple)
        else:
            fused_output = torch.ops.vllm.moe_forward(
                hidden_states, router_logits, self.layer_name)
        return fused_output[..., :og_hidden_states]
    else:
        if current_platform.is_tpu():
            # TODO: Once the OOM issue for the TPU backend is resolved, we
            # will switch to using the moe_forward custom op.
            shared_output, fused_output = self.forward_impl(
                hidden_states, router_logits)
        else:
            shared_output, fused_output = torch.ops.vllm.moe_forward_shared(
                hidden_states, router_logits, self.layer_name)
        return (shared_output[..., :og_hidden_states],
                fused_output[..., :og_hidden_states])

get_expert_weights ¶

get_expert_weights() -> Iterable[Tensor]

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

def get_expert_weights(self) -> Iterable[torch.Tensor]:
    weights = list(self.named_parameters())
    assert all(weight.is_contiguous() for _, weight in weights)

    # Filter out the non-expert weights.
    # `e_score_correction_bias` is a bias for each logical expert,
    # with shape (num_logical_experts,), not an expert weight.
    NON_EXPERT_WEIGHTS = {
        "e_score_correction_bias",
    }

    return [
        weight.view(self.local_num_experts, -1) for name, weight in weights
        if name not in NON_EXPERT_WEIGHTS and weight.shape != torch.Size(
            []) and not name.startswith("_shared_experts.")
    ]

make_expert_params_mapping `classmethod` ¶

make_expert_params_mapping(
    ckpt_gate_proj_name: str,
    ckpt_down_proj_name: str,
    ckpt_up_proj_name: str,
    num_experts: int,
    num_redundant_experts: int = 0,
) -> list[tuple[str, str, int, str]]

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

@classmethod
def make_expert_params_mapping(
        cls,
        ckpt_gate_proj_name: str,
        ckpt_down_proj_name: str,
        ckpt_up_proj_name: str,
        num_experts: int,
        num_redundant_experts: int = 0) -> list[tuple[str, str, int, str]]:

    num_physical_experts = num_experts + num_redundant_experts

    # In the returned mapping:
    # - `expert_id` is the physical expert id
    # - `weight_name` contains the weight name of the logical expert
    # So that we should map the expert id to logical in `weight_name`
    physical_to_logical_map = \
        EplbState.build_initial_global_physical_to_logical_map(
        num_experts, num_redundant_experts)

    return [
        # (param_name, weight_name, expert_id, shard_id)
        ("experts.w13_" if weight_name
         in [ckpt_gate_proj_name, ckpt_up_proj_name] else "experts.w2_",
         f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.",
         expert_id, shard_id) for expert_id in range(num_physical_experts)
        for shard_id, weight_name in [
            ("w1", ckpt_gate_proj_name),
            ("w2", ckpt_down_proj_name),
            ("w3", ckpt_up_proj_name),
        ]
    ]

maybe_all_reduce_tensor_model_parallel ¶

maybe_all_reduce_tensor_model_parallel(
    final_hidden_states: Tensor,
)

The pplx combine kernel reduces across GPU ranks by default.

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

def maybe_all_reduce_tensor_model_parallel(
        self, final_hidden_states: torch.Tensor):
    """
    The pplx combine kernel reduces across GPU ranks by default.
    """
    if (self.use_pplx_kernels or self.use_deepep_ht_kernels
            or self.use_deepep_ll_kernels):
        return final_hidden_states
    else:
        return tensor_model_parallel_all_reduce(final_hidden_states)

must_reduce_shared_expert_outputs ¶

must_reduce_shared_expert_outputs() -> bool

The shared_experts are typically computed using the RowParallelLinear layer. The result of this function is typically used as the reduce_results argument to the module. When just tensor-parallel is used, it is not required to reduce the shared_experts results immediately. Instead we reduce at the once at the end of the MoE op. (Refer to DeepSeekV2MoE module) With EP and all2all kernels - this is no longer viable as all GPU ranks in DP, produce the complete set of hidden_states. Therefore it is required that we reduce the shared_experts output early.

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

def must_reduce_shared_expert_outputs(self) -> bool:
    """
    The shared_experts are typically computed using the RowParallelLinear
    layer. The result of this function is typically used as
    the reduce_results argument to the module.
    When just tensor-parallel is used, it is not required to reduce
    the shared_experts results immediately. Instead we reduce at the
    once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
    With EP and all2all kernels - this is no longer viable as all
    GPU ranks in DP, produce the complete set of hidden_states.
    Therefore it is required that we reduce the shared_experts output
    early.
    """
    return (self.use_pplx_kernels or self.use_deepep_ht_kernels
            or self.use_deepep_ll_kernels)

select_experts `staticmethod` ¶

select_experts(
    hidden_states: Tensor,
    router_logits: Tensor,
    top_k: int,
    use_grouped_topk: bool,
    renormalize: bool,
    topk_group: Optional[int] = None,
    num_expert_group: Optional[int] = None,
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[Tensor] = None,
    indices_type: Optional[dtype] = None,
    enable_eplb: bool = False,
    expert_map: Optional[Tensor] = None,
    expert_load_view: Optional[Tensor] = None,
    logical_to_physical_map: Optional[Tensor] = None,
    logical_replica_count: Optional[Tensor] = None,
    global_num_experts: Optional[int] = None,
    zero_expert_num: Optional[int] = None,
    zero_expert_type: Optional[str] = None,
) -> tuple[Tensor, Tensor, Tensor]

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

Returns:

Type	Description
`Tensor`	(topk_weights, topk_ids, zero_expert_result)
`tuple[Tensor, Tensor, Tensor]`
`Tensor`	The weights, expert ids, and zero expert computation result.

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

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

@staticmethod
def select_experts(
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
    top_k: int,
    use_grouped_topk: bool,
    renormalize: bool,
    topk_group: Optional[int] = None,
    num_expert_group: Optional[int] = None,
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[torch.Tensor] = None,
    indices_type: Optional[torch.dtype] = None,
    enable_eplb: bool = False,
    expert_map: Optional[torch.Tensor] = None,
    expert_load_view: Optional[torch.Tensor] = None,
    logical_to_physical_map: Optional[torch.Tensor] = None,
    logical_replica_count: Optional[torch.Tensor] = None,
    global_num_experts: Optional[int] = None,
    zero_expert_num: Optional[int] = None,
    zero_expert_type: Optional[str] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Route the input hidden states to the top-k experts based on the
    router logits.

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

        **Compatibility**: When EPLB is not enabled, the returned ids are
        equivalent to global logical ids, so should be compatible with
        plain MoE implementations without redundant experts.
    """
    from vllm.model_executor.layers.fused_moe.fused_moe import (
        fused_topk, fused_topk_bias)

    # Check if we should use a routing simulation strategy
    routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
    if routing_strategy != "":
        topk_weights, topk_ids = RoutingSimulator.simulate_routing(
            hidden_states=hidden_states,
            router_logits=router_logits,
            strategy_name=routing_strategy,
            top_k=top_k,
            indices_type=indices_type)

    # DeepSeekv2 uses grouped_top_k
    if use_grouped_topk:
        assert topk_group is not None
        assert num_expert_group is not None
        topk_weights, topk_ids = grouped_topk(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=top_k,
            renormalize=renormalize,
            num_expert_group=num_expert_group,
            topk_group=topk_group,
            scoring_func=scoring_func,
            routed_scaling_factor=routed_scaling_factor,
            e_score_correction_bias=e_score_correction_bias)
        if indices_type is not None:
            topk_ids = topk_ids.to(dtype=indices_type)
    elif e_score_correction_bias is not None:
        topk_weights, topk_ids = fused_topk_bias(
            hidden_states=hidden_states,
            gating_output=router_logits,
            e_score_correction_bias=e_score_correction_bias.data,
            topk=top_k,
            renormalize=renormalize,
        )
        if routed_scaling_factor is not None:
            topk_weights *= routed_scaling_factor
    elif custom_routing_function is None:
        topk_weights, topk_ids, token_expert_indices = fused_topk(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=top_k,
            renormalize=renormalize,
            indices_type=indices_type,
        )
    else:
        topk_weights, topk_ids = custom_routing_function(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=top_k,
            renormalize=renormalize)
        if indices_type is not None:
            topk_ids = topk_ids.to(dtype=indices_type)

    if enable_eplb:
        assert expert_load_view is not None
        assert logical_to_physical_map is not None
        assert logical_replica_count is not None

        topk_ids = eplb_map_to_physical_and_record(
            topk_ids=topk_ids,
            expert_load_view=expert_load_view,
            logical_to_physical_map=logical_to_physical_map,
            logical_replica_count=logical_replica_count,
            indices_type=indices_type,
        )

    assert topk_ids.dtype == indices_type or indices_type is None

    # Compute zero expert result if needed
    if (zero_expert_num is not None and zero_expert_num > 0
            and zero_expert_type is not None
            and global_num_experts is not None):
        zero_expert_result = zero_experts_compute_triton(
            expert_indices=topk_ids,
            expert_scales=topk_weights,
            num_experts=global_num_experts,
            zero_expert_type=zero_expert_type,
            hidden_states=hidden_states,
        )
    else:
        zero_expert_result = None
    return topk_weights, topk_ids, zero_expert_result

set_eplb_state ¶

set_eplb_state(
    moe_layer_idx: int,
    expert_load_view: Tensor,
    logical_to_physical_map: Tensor,
    logical_replica_count: Tensor,
) -> None

Register the EPLB state in this layer.

This is used later in forward pass, where we get the expert mapping and record the load metrics in expert_load_view.

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

def set_eplb_state(
    self,
    moe_layer_idx: int,
    expert_load_view: torch.Tensor,
    logical_to_physical_map: torch.Tensor,
    logical_replica_count: torch.Tensor,
) -> None:
    """
    Register the EPLB state in this layer.

    This is used later in forward pass, where we get the expert mapping
    and record the load metrics in `expert_load_view`.
    """
    self.expert_load_view = expert_load_view[moe_layer_idx]
    self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
    self.logical_replica_count = logical_replica_count[moe_layer_idx]

update_expert_map ¶

update_expert_map()

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

def update_expert_map(self):
    # ep_size and ep_rank should already be updated
    assert self.expert_map is not None
    with self.expert_map.device:
        local_num_experts, expert_map = determine_expert_map(
            ep_size=self.ep_size,
            ep_rank=self.ep_rank,
            global_num_experts=self.global_num_experts)
        self.local_num_experts = local_num_experts
        self.register_buffer("expert_map", expert_map)

weight_loader ¶

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    weight_name: str,
    shard_id: str,
    expert_id: int,
    return_success: Literal[False],
) -> None

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    weight_name: str,
    shard_id: str,
    expert_id: int,
    return_success: Literal[True],
) -> bool

weight_loader(
    param: Parameter,
    loaded_weight: Tensor,
    weight_name: str,
    shard_id: str,
    expert_id: int,
    return_success: bool = False,
) -> Optional[bool]

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

def weight_loader(self,
                  param: torch.nn.Parameter,
                  loaded_weight: torch.Tensor,
                  weight_name: str,
                  shard_id: str,
                  expert_id: int,
                  return_success: bool = False) -> Optional[bool]:

    if self.quant_config and self.quant_config.get_name() == "mxfp4":
        # (FIXME) for gpt-oss all experts are combined
        if "bias" in weight_name:
            dim1 = loaded_weight.shape[1]
            param.data[:, :dim1].copy_(loaded_weight)
        else:
            dim1 = loaded_weight.shape[1]
            dim2 = loaded_weight.shape[2]
            param.data[:, :dim1, :dim2].copy_(loaded_weight)
        return True if return_success else None

    expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
    if expert_id == -1:
        # Failed to load this param since it's not local to this rank
        return False if return_success else None
    # Hereafter, `expert_id` is local physical id

    quant_method_name = self.quant_method.__class__.__name__
    # compressed-tensors checkpoints with packed weights are stored flipped
    # TODO (mgoin): check self.quant_method.quant_config.quant_format
    # against known CompressionFormat enum values that have this quality
    if self.quant_method.__class__.__name__ in (
            "CompressedTensorsWNA16MarlinMoEMethod",
            "CompressedTensorsWNA16MoEMethod"):
        loaded_weight = loaded_weight.t().contiguous()

    if shard_id not in ("w1", "w2", "w3"):
        raise ValueError(f"shard_id must be ['w1','w2','w3'] but "
                         f"got {shard_id}.")

    # Fetch the dim to shard the parameter/loaded weight
    # based on the shard id. This will be whatever
    # dimension intermediate_size_per_partition is used.
    SHARD_ID_TO_SHARDED_DIM = {"w1": 0, "w2": 1, "w3": 0}

    is_gguf_weight = getattr(param, "is_gguf_weight", False)
    is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
    if is_gguf_weight_type:
        param.weight_type = loaded_weight.item()
        param.data.copy_(loaded_weight)
        return True if return_success else None

    # Case for BitsAndBytes
    use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
    if use_bitsandbytes_4bit:
        shard_dim = 0

        expert_data = param.data[expert_id]
        if shard_id == "w2":
            expert_data.copy_(loaded_weight)
        elif shard_id in ("w1", "w3"):
            # BNB inflight quantization has already sharded the weights
            full_load = True
            self._load_w13(
                shard_id=shard_id,
                shard_dim=shard_dim,
                loaded_weight=loaded_weight,
                expert_data=expert_data,
                tp_rank=self.tp_rank,
                load_full=full_load,
            )
        return True if return_success else None

    # is_transposed: if the dim to shard the weight
    # should be flipped. Required by GPTQ, compressed-tensors
    # should be whatever dimension intermediate_size_per_partition is
    is_transposed = getattr(param, "is_transposed", False)
    shard_dim = SHARD_ID_TO_SHARDED_DIM[shard_id]
    if is_transposed:
        shard_dim = int(not shard_dim)

    full_load = len(loaded_weight.shape) == 3
    if full_load:
        shard_dim += 1

    # Materialize GGUF UninitializedParameter
    if is_gguf_weight and isinstance(param, UninitializedParameter):
        final_shape = list(loaded_weight.shape)
        if shard_id in ["w1", "w3"]:
            final_shape[1] *= 2
        final_shape[shard_dim] = final_shape[shard_dim] // self.tp_size
        param.materialize(final_shape, dtype=loaded_weight.dtype)

    expert_data = param.data if full_load else param.data[expert_id]

    # Case input scale: input_scale loading is only supported for fp8
    if "input_scale" in weight_name:
        # this is needed for compressed-tensors only
        loaded_weight = loaded_weight.to(param.data.device)

        if ("compressed" in quant_method_name.lower()
                and param.data[expert_id] != 1
                and (param.data[expert_id] - loaded_weight).abs() > 1e-5):
            raise ValueError(
                "input_scales of w1 and w3 of a layer "
                f"must be equal. But got {param.data[expert_id]} "
                f"vs. {loaded_weight}")

        self._load_single_value(param=param,
                                loaded_weight=loaded_weight,
                                expert_id=expert_id)
        return True if return_success else None

    # Case g_idx
    if "g_idx" in weight_name:
        self._load_g_idx(shard_dim=0,
                         shard_id=shard_id,
                         loaded_weight=loaded_weight,
                         expert_data=expert_data,
                         tp_rank=self.tp_rank)
        return True if return_success else None

    # TODO @dsikka: ModelOpt should follow the proper MoE loading pattern
    if "ModelOpt" in quant_method_name:
        # Determine per-tensor weight scale patterns based on variant
        # Use the dedicated method instead of brittle string matching
        uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern(
        )

        # Call _load_per_tensor_weight_scale() to load per-tensor (scalar)
        # weights scales.
        # Input scales are always per-tensor.
        # Weight scales: FP4 uses "weight_scale_2" and FP8 uses
        # "weight_scale" for per-tensor scales.
        is_per_tensor = ("weight_scale_2" in weight_name
                         if uses_weight_scale_2 else "weight_scale"
                         in weight_name) or "input_scale" in weight_name
        if is_per_tensor:
            self._load_per_tensor_weight_scale(
                shard_id=shard_id,
                param=param,
                loaded_weight=loaded_weight,
                expert_id=expert_id,
            )
            return True if return_success else None

        # If the weight is w13_weight_scale and w13_weight_scales are
        # combined into single loaded_weight, call
        # _load_combined_w13_weight_scale() to load it.
        # This is checked by comparing the hidden_out dims of the
        # loaded_weight and the param.
        if "w13_weight_scale" in weight_name:
            loaded_weight_hidden_out = loaded_weight.shape[-2]
            param_hidden_out = param.data.shape[-2] * self.tp_size
            if loaded_weight_hidden_out == param_hidden_out:
                self._load_combined_w13_weight_scale(
                    shard_dim=shard_dim,
                    loaded_weight=loaded_weight,
                    param=param,
                    tp_rank=self.tp_rank,
                )
                return True if return_success else None

        # For other weights, call _load_model_weight_or_group_weight_scale()
        # to load it.
        if "weight" in weight_name:
            self._load_model_weight_or_group_weight_scale(
                shard_id=shard_id,
                shard_dim=shard_dim,
                loaded_weight=loaded_weight,
                expert_data=expert_data,
                tp_rank=self.tp_rank)
        return True if return_success else None

    # Case weight scales, zero_points and offset, weight/input global scales
    if ("scale" in weight_name or "zero" in weight_name
            or "offset" in weight_name):
        # load the weight scales and zp based on the quantization scheme
        # supported weight scales/zp can be found in
        # FusedMoeWeightScaleSupported
        # TODO @dsikka: once hardened, refactor to use vLLM Parameters
        # specific to each case
        quant_method = getattr(param, "quant_method", None)
        if quant_method == FusedMoeWeightScaleSupported.CHANNEL.value:
            self._load_per_channel_weight_scale(
                shard_id=shard_id,
                shard_dim=shard_dim,
                loaded_weight=loaded_weight,
                expert_data=expert_data,
                tp_rank=self.tp_rank)
        elif quant_method in [
                FusedMoeWeightScaleSupported.GROUP.value,
                FusedMoeWeightScaleSupported.BLOCK.value,
        ]:
            self._load_model_weight_or_group_weight_scale(
                shard_id=shard_id,
                shard_dim=shard_dim,
                loaded_weight=loaded_weight,
                expert_data=expert_data,
                tp_rank=self.tp_rank,
                load_full_w2=getattr(param, "load_full_w2", False))
        elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:
            self._load_per_tensor_weight_scale(shard_id=shard_id,
                                               param=param,
                                               loaded_weight=loaded_weight,
                                               expert_id=expert_id)
        else:
            WEIGHT_SCALE_SUPPORTED = [
                e.value for e in FusedMoeWeightScaleSupported
            ]
            raise ValueError(
                f"quant method must be one of {WEIGHT_SCALE_SUPPORTED}")
        return True if return_success else None

    # Case weight_shape
    if "weight_shape" in weight_name:
        # only required by compressed-tensors
        self._load_single_value(param=param,
                                loaded_weight=loaded_weight,
                                expert_id=expert_id)
        return True if return_success else None

    # Case model weights
    if "weight" in weight_name:
        self._load_model_weight_or_group_weight_scale(
            shard_id=shard_id,
            shard_dim=shard_dim,
            loaded_weight=loaded_weight,
            expert_data=expert_data,
            tp_rank=self.tp_rank)
        return True if return_success else None

    return False if return_success else None

FusedMoEActivationFormat ¶

Bases: Enum

The standard activation format (num_tokens, 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",

BatchedExperts `class-attribute` `instance-attribute` ¶

BatchedExperts = ('batched_experts',)

Standard `class-attribute` `instance-attribute` ¶

Standard = ('standard',)

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

FusedMoEConfig `dataclass` ¶

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

@dataclass
class FusedMoEConfig:
    num_experts: int
    experts_per_token: int
    hidden_dim: int

    num_local_experts: int
    moe_parallel_config: FusedMoEParallelConfig

    # The activation type.
    in_dtype: torch.dtype

    max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE

    has_bias: bool = False

    def __post_init__(self):
        if self.dp_size > 1:
            logger.debug_once("Using FusedMoEConfig::max_num_tokens=%d",
                              self.max_num_tokens)

        assert self.max_num_tokens > 0

    @property
    def tp_size(self):
        return self.moe_parallel_config.tp_size

    @property
    def dp_size(self):
        return self.moe_parallel_config.dp_size

    @property
    def ep_size(self):
        return self.moe_parallel_config.ep_size

    @property
    def tp_rank(self):
        return self.moe_parallel_config.tp_rank

    @property
    def dp_rank(self):
        return self.moe_parallel_config.dp_rank

    @property
    def ep_rank(self):
        return self.moe_parallel_config.ep_rank

    @property
    def use_ep(self):
        return self.moe_parallel_config.use_ep

    @property
    def use_pplx_kernels(self):
        return self.moe_parallel_config.use_pplx_kernels

    @property
    def use_deepep_ht_kernels(self):
        return self.moe_parallel_config.use_deepep_ht_kernels

    @property
    def use_deepep_ll_kernels(self):
        return self.moe_parallel_config.use_deepep_ll_kernels

    @property
    def use_flashinfer_cutlass_kernels(self):
        """
        Whether to use FlashInfer cutlass kernels for NVFP4 MoE.
        """
        return (envs.VLLM_USE_FLASHINFER_MOE_FP4
                and has_flashinfer_cutlass_fused_moe()
                and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput")

dp_rank `property` ¶

dp_rank

dp_size `property` ¶

dp_size

ep_rank `property` ¶

ep_rank

ep_size `property` ¶

ep_size

experts_per_token `instance-attribute` ¶

experts_per_token: int

has_bias `class-attribute` `instance-attribute` ¶

has_bias: bool = False

hidden_dim `instance-attribute` ¶

hidden_dim: int

in_dtype `instance-attribute` ¶

in_dtype: dtype

max_num_tokens `class-attribute` `instance-attribute` ¶

max_num_tokens: int = VLLM_MOE_DP_CHUNK_SIZE

moe_parallel_config `instance-attribute` ¶

moe_parallel_config: FusedMoEParallelConfig

num_experts `instance-attribute` ¶

num_experts: int

num_local_experts `instance-attribute` ¶

num_local_experts: int

tp_rank `property` ¶

tp_rank

tp_size `property` ¶

tp_size

use_deepep_ht_kernels `property` ¶

use_deepep_ht_kernels

use_deepep_ll_kernels `property` ¶

use_deepep_ll_kernels

use_ep `property` ¶

use_ep

use_flashinfer_cutlass_kernels `property` ¶

use_flashinfer_cutlass_kernels

Whether to use FlashInfer cutlass kernels for NVFP4 MoE.

use_pplx_kernels `property` ¶

use_pplx_kernels

init ¶

__init__(
    num_experts: int,
    experts_per_token: int,
    hidden_dim: int,
    num_local_experts: int,
    moe_parallel_config: FusedMoEParallelConfig,
    in_dtype: dtype,
    max_num_tokens: int = VLLM_MOE_DP_CHUNK_SIZE,
    has_bias: bool = False,
) -> None

__post_init__ ¶

__post_init__()

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

def __post_init__(self):
    if self.dp_size > 1:
        logger.debug_once("Using FusedMoEConfig::max_num_tokens=%d",
                          self.max_num_tokens)

    assert self.max_num_tokens > 0

FusedMoEMethodBase ¶

Bases: QuantizeMethodBase

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

class FusedMoEMethodBase(QuantizeMethodBase):

    def __init__(self, moe: FusedMoEConfig):
        super().__init__()
        self.moe = moe
        self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
        self.fused_experts: Optional[FusedMoEModularKernel] = None
        self.topk_indices_dtype = None

    @abstractmethod
    def create_weights(self, layer: torch.nn.Module, num_experts: int,
                       hidden_size: int, intermediate_size_per_partition: int,
                       params_dtype: torch.dtype, **extra_weight_attrs):
        raise NotImplementedError

    def uses_weight_scale_2_pattern(self) -> bool:
        """
        Returns True if this quantization method uses 'weight_scale_2' pattern
        for per-tensor weight scales (e.g., FP4 variants), False otherwise.

        This method should be overridden by subclasses that use the
        'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
        """
        return False

    @staticmethod
    def _maybe_make_prepare_finalize(
        moe: FusedMoEConfig,
        quant_config: Optional[FusedMoEQuantConfig],
    ) -> Optional[FusedMoEPrepareAndFinalize]:
        all2all_manager = get_ep_group().device_communicator.all2all_manager
        assert all2all_manager is not None

        prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None

        # TODO: could allow this now
        assert not moe.use_flashinfer_cutlass_kernels, \
            "Must be created in modelopt.py"

        if moe.use_pplx_kernels:
            assert quant_config is not None

            hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
                moe.max_num_tokens,
                moe.hidden_dim,
                moe.in_dtype,
                quant_config.quant_dtype,
                per_act_token_quant=quant_config.per_act_token_quant,
                block_shape=quant_config.block_shape,
            )

            all_to_all_args = dict(
                max_num_tokens=moe.max_num_tokens,
                num_experts=moe.num_experts,
                experts_per_token=moe.experts_per_token,  # topk
                rank=all2all_manager.rank,
                world_size=all2all_manager.world_size,
                # dp_size actually means tp_size, bug in pplx kernels
                dp_size=all2all_manager.tp_group.world_size,
                hidden_dim=moe.hidden_dim,
                hidden_dim_bytes=hidden_dim_bytes,
                hidden_dim_scale_bytes=hidden_scale_bytes,
            )

            num_dispatchers = (all2all_manager.world_size //
                               all2all_manager.tp_group.world_size)

            # Intranode pplx a2a takes a group name while internode does not.
            if not all2all_manager.internode:
                all_to_all_args[
                    "group_name"] = all2all_manager.cpu_group.group_name

            handle = all2all_manager.get_handle(all_to_all_args)

            prepare_finalize = PplxPrepareAndFinalize(
                handle,
                max_num_tokens=moe.max_num_tokens,
                num_local_experts=moe.num_local_experts,
                num_dispatchers=num_dispatchers,
            )
        elif moe.use_deepep_ht_kernels:
            assert moe.dp_size == all2all_manager.dp_world_size

            all_to_all_args = dict()
            handle = all2all_manager.get_handle(all_to_all_args)
            prepare_finalize = DeepEPHTPrepareAndFinalize(
                handle,
                num_dispatchers=all2all_manager.world_size,
                dp_size=all2all_manager.dp_world_size,
                rank_expert_offset=all2all_manager.rank *
                moe.num_local_experts,
            )

        elif moe.use_deepep_ll_kernels:
            assert quant_config is not None
            all_to_all_args = dict(
                max_num_tokens_per_dp_rank=moe.max_num_tokens,
                token_hidden_size=moe.hidden_dim,
                num_ep_ranks=all2all_manager.world_size,
                num_global_experts=moe.num_experts,
                num_local_experts=moe.num_experts //
                all2all_manager.world_size)
            handle = all2all_manager.get_handle(all_to_all_args)

            # Note: We may want to use FP8 dispatch just to reduce
            # data movement.
            use_fp8_dispatch = (
                quant_config.quant_dtype == current_platform.fp8_dtype()
                and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE)

            prepare_finalize = DeepEPLLPrepareAndFinalize(
                handle,
                max_tokens_per_rank=moe.max_num_tokens,
                num_dispatchers=all2all_manager.world_size,
                use_fp8_dispatch=use_fp8_dispatch,
            )

        return prepare_finalize

    def maybe_make_prepare_finalize(
            self) -> Optional[FusedMoEPrepareAndFinalize]:
        if self.moe.moe_parallel_config.use_all2all_kernels:
            return FusedMoEMethodBase._maybe_make_prepare_finalize(
                self.moe, self.moe_quant_config)
        else:
            return None

    # Note: init_prepare_finalize should only be called by
    # prepare_communication_buffer_for_model.
    def init_prepare_finalize(self, layer: torch.nn.Module):
        assert self.moe is not None

        # We must get the quant config here so that the layer is
        # completely initialized, i.e. all weights loaded and post
        # processed.
        self.moe_quant_config = self.get_fused_moe_quant_config(layer)

        prepare_finalize = self.maybe_make_prepare_finalize()

        if prepare_finalize is not None:
            logger.debug("%s for %s(%s)", prepare_finalize.__class__.__name__,
                         self, id(self))
            assert self.topk_indices_dtype is None
            assert self.fused_experts is None, \
                f"Attempt to override experts for {id(self)}!"
            self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
            experts = self.select_gemm_impl(prepare_finalize, layer)
            self.fused_experts = FusedMoEModularKernel(
                prepare_finalize,
                experts,
                layer.shared_experts,
            )

    def select_gemm_impl(
        self,
        prepare_finalize: FusedMoEPrepareAndFinalize,
        layer: torch.nn.Module,
    ) -> FusedMoEPermuteExpertsUnpermute:
        # based on the all2all implementation, select the appropriate
        # gemm implementation
        raise NotImplementedError(
            f"{self.__class__.__name__} must select appropriate gemm "
            "implementation based on the prepare_finalize")

    @abstractmethod
    def get_fused_moe_quant_config(
            self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
        raise NotImplementedError

    @abstractmethod
    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        router_logits: torch.Tensor,
        top_k: int,
        renormalize: bool,
        use_grouped_topk: bool = False,
        topk_group: Optional[int] = None,
        num_expert_group: Optional[int] = None,
        global_num_experts: int = -1,
        expert_map: Optional[torch.Tensor] = None,
        custom_routing_function: Optional[Callable] = None,
        scoring_func: str = "softmax",
        routed_scaling_factor: float = 1.0,
        e_score_correction_bias: Optional[torch.Tensor] = None,
        apply_router_weight_on_input: bool = False,
        activation: str = "silu",
        enable_eplb: bool = False,
        expert_load_view: Optional[torch.Tensor] = None,
        logical_to_physical_map: Optional[torch.Tensor] = None,
        logical_replica_count: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        raise NotImplementedError

fused_experts `instance-attribute` ¶

fused_experts: Optional[FusedMoEModularKernel] = None

moe `instance-attribute` ¶

moe = moe

moe_quant_config `instance-attribute` ¶

moe_quant_config: Optional[FusedMoEQuantConfig] = None

topk_indices_dtype `instance-attribute` ¶

topk_indices_dtype = None

init ¶

__init__(moe: FusedMoEConfig)

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

def __init__(self, moe: FusedMoEConfig):
    super().__init__()
    self.moe = moe
    self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
    self.fused_experts: Optional[FusedMoEModularKernel] = None
    self.topk_indices_dtype = None

_maybe_make_prepare_finalize `staticmethod` ¶

_maybe_make_prepare_finalize(
    moe: FusedMoEConfig,
    quant_config: Optional[FusedMoEQuantConfig],
) -> Optional[FusedMoEPrepareAndFinalize]

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

@staticmethod
def _maybe_make_prepare_finalize(
    moe: FusedMoEConfig,
    quant_config: Optional[FusedMoEQuantConfig],
) -> Optional[FusedMoEPrepareAndFinalize]:
    all2all_manager = get_ep_group().device_communicator.all2all_manager
    assert all2all_manager is not None

    prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None

    # TODO: could allow this now
    assert not moe.use_flashinfer_cutlass_kernels, \
        "Must be created in modelopt.py"

    if moe.use_pplx_kernels:
        assert quant_config is not None

        hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
            moe.max_num_tokens,
            moe.hidden_dim,
            moe.in_dtype,
            quant_config.quant_dtype,
            per_act_token_quant=quant_config.per_act_token_quant,
            block_shape=quant_config.block_shape,
        )

        all_to_all_args = dict(
            max_num_tokens=moe.max_num_tokens,
            num_experts=moe.num_experts,
            experts_per_token=moe.experts_per_token,  # topk
            rank=all2all_manager.rank,
            world_size=all2all_manager.world_size,
            # dp_size actually means tp_size, bug in pplx kernels
            dp_size=all2all_manager.tp_group.world_size,
            hidden_dim=moe.hidden_dim,
            hidden_dim_bytes=hidden_dim_bytes,
            hidden_dim_scale_bytes=hidden_scale_bytes,
        )

        num_dispatchers = (all2all_manager.world_size //
                           all2all_manager.tp_group.world_size)

        # Intranode pplx a2a takes a group name while internode does not.
        if not all2all_manager.internode:
            all_to_all_args[
                "group_name"] = all2all_manager.cpu_group.group_name

        handle = all2all_manager.get_handle(all_to_all_args)

        prepare_finalize = PplxPrepareAndFinalize(
            handle,
            max_num_tokens=moe.max_num_tokens,
            num_local_experts=moe.num_local_experts,
            num_dispatchers=num_dispatchers,
        )
    elif moe.use_deepep_ht_kernels:
        assert moe.dp_size == all2all_manager.dp_world_size

        all_to_all_args = dict()
        handle = all2all_manager.get_handle(all_to_all_args)
        prepare_finalize = DeepEPHTPrepareAndFinalize(
            handle,
            num_dispatchers=all2all_manager.world_size,
            dp_size=all2all_manager.dp_world_size,
            rank_expert_offset=all2all_manager.rank *
            moe.num_local_experts,
        )

    elif moe.use_deepep_ll_kernels:
        assert quant_config is not None
        all_to_all_args = dict(
            max_num_tokens_per_dp_rank=moe.max_num_tokens,
            token_hidden_size=moe.hidden_dim,
            num_ep_ranks=all2all_manager.world_size,
            num_global_experts=moe.num_experts,
            num_local_experts=moe.num_experts //
            all2all_manager.world_size)
        handle = all2all_manager.get_handle(all_to_all_args)

        # Note: We may want to use FP8 dispatch just to reduce
        # data movement.
        use_fp8_dispatch = (
            quant_config.quant_dtype == current_platform.fp8_dtype()
            and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE)

        prepare_finalize = DeepEPLLPrepareAndFinalize(
            handle,
            max_tokens_per_rank=moe.max_num_tokens,
            num_dispatchers=all2all_manager.world_size,
            use_fp8_dispatch=use_fp8_dispatch,
        )

    return prepare_finalize

apply `abstractmethod` ¶

apply(
    layer: Module,
    x: Tensor,
    router_logits: Tensor,
    top_k: int,
    renormalize: bool,
    use_grouped_topk: bool = False,
    topk_group: Optional[int] = None,
    num_expert_group: Optional[int] = None,
    global_num_experts: int = -1,
    expert_map: Optional[Tensor] = None,
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[Tensor] = None,
    apply_router_weight_on_input: bool = False,
    activation: str = "silu",
    enable_eplb: bool = False,
    expert_load_view: Optional[Tensor] = None,
    logical_to_physical_map: Optional[Tensor] = None,
    logical_replica_count: Optional[Tensor] = None,
) -> Union[Tensor, tuple[Tensor, Tensor]]

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

@abstractmethod
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    router_logits: torch.Tensor,
    top_k: int,
    renormalize: bool,
    use_grouped_topk: bool = False,
    topk_group: Optional[int] = None,
    num_expert_group: Optional[int] = None,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    custom_routing_function: Optional[Callable] = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[torch.Tensor] = None,
    apply_router_weight_on_input: bool = False,
    activation: str = "silu",
    enable_eplb: bool = False,
    expert_load_view: Optional[torch.Tensor] = None,
    logical_to_physical_map: Optional[torch.Tensor] = None,
    logical_replica_count: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    raise NotImplementedError

create_weights `abstractmethod` ¶

create_weights(
    layer: Module,
    num_experts: int,
    hidden_size: int,
    intermediate_size_per_partition: int,
    params_dtype: dtype,
    **extra_weight_attrs,
)

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

@abstractmethod
def create_weights(self, layer: torch.nn.Module, num_experts: int,
                   hidden_size: int, intermediate_size_per_partition: int,
                   params_dtype: torch.dtype, **extra_weight_attrs):
    raise NotImplementedError

get_fused_moe_quant_config `abstractmethod` ¶

get_fused_moe_quant_config(
    layer: Module,
) -> Optional[FusedMoEQuantConfig]

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

@abstractmethod
def get_fused_moe_quant_config(
        self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
    raise NotImplementedError

init_prepare_finalize ¶

init_prepare_finalize(layer: Module)

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

def init_prepare_finalize(self, layer: torch.nn.Module):
    assert self.moe is not None

    # We must get the quant config here so that the layer is
    # completely initialized, i.e. all weights loaded and post
    # processed.
    self.moe_quant_config = self.get_fused_moe_quant_config(layer)

    prepare_finalize = self.maybe_make_prepare_finalize()

    if prepare_finalize is not None:
        logger.debug("%s for %s(%s)", prepare_finalize.__class__.__name__,
                     self, id(self))
        assert self.topk_indices_dtype is None
        assert self.fused_experts is None, \
            f"Attempt to override experts for {id(self)}!"
        self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
        experts = self.select_gemm_impl(prepare_finalize, layer)
        self.fused_experts = FusedMoEModularKernel(
            prepare_finalize,
            experts,
            layer.shared_experts,
        )

maybe_make_prepare_finalize ¶

maybe_make_prepare_finalize() -> Optional[
    FusedMoEPrepareAndFinalize
]

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

def maybe_make_prepare_finalize(
        self) -> Optional[FusedMoEPrepareAndFinalize]:
    if self.moe.moe_parallel_config.use_all2all_kernels:
        return FusedMoEMethodBase._maybe_make_prepare_finalize(
            self.moe, self.moe_quant_config)
    else:
        return None

select_gemm_impl ¶

select_gemm_impl(
    prepare_finalize: FusedMoEPrepareAndFinalize,
    layer: Module,
) -> FusedMoEPermuteExpertsUnpermute

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

def select_gemm_impl(
    self,
    prepare_finalize: FusedMoEPrepareAndFinalize,
    layer: torch.nn.Module,
) -> FusedMoEPermuteExpertsUnpermute:
    # based on the all2all implementation, select the appropriate
    # gemm implementation
    raise NotImplementedError(
        f"{self.__class__.__name__} must select appropriate gemm "
        "implementation based on the prepare_finalize")

uses_weight_scale_2_pattern ¶

uses_weight_scale_2_pattern() -> bool

Returns True if this quantization method uses 'weight_scale_2' pattern for per-tensor weight scales (e.g., FP4 variants), False otherwise.

This method should be overridden by subclasses that use the 'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.

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

def uses_weight_scale_2_pattern(self) -> bool:
    """
    Returns True if this quantization method uses 'weight_scale_2' pattern
    for per-tensor weight scales (e.g., FP4 variants), False otherwise.

    This method should be overridden by subclasses that use the
    'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
    """
    return False

FusedMoEPermuteExpertsUnpermute ¶

Bases: ABC

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

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

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

    def __init__(
        self,
        quant_config: FusedMoEQuantConfig,
    ):
        """
        quant_config: Quantization parameters for this experts instance.
        """
        self.quant_config = quant_config

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

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

    @property
    def quant_dtype(self) -> Optional[torch.dtype]:
        return self.quant_config.quant_dtype

    @property
    def block_shape(self) -> Optional[list[int]]:
        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) -> Optional[torch.Tensor]:
        return self.quant_config.a1_scale

    @property
    def a2_scale(self) -> Optional[torch.Tensor]:
        return self.quant_config.a2_scale

    @property
    def a1_gscale(self) -> Optional[torch.Tensor]:
        return self.quant_config.a1_gscale

    @property
    def a2_gscale(self) -> Optional[torch.Tensor]:
        return self.quant_config.a2_gscale

    @property
    def w1_scale(self) -> Optional[torch.Tensor]:
        return self.quant_config.w1_scale

    @property
    def w2_scale(self) -> Optional[torch.Tensor]:
        return self.quant_config.w2_scale

    @property
    def w1_zp(self) -> Optional[torch.Tensor]:
        return self.quant_config.w1_zp

    @property
    def w2_zp(self) -> Optional[torch.Tensor]:
        return self.quant_config.w2_zp

    @property
    def w1_bias(self) -> Optional[torch.Tensor]:
        return self.quant_config.w1_bias

    @property
    def w2_bias(self) -> Optional[torch.Tensor]:
        return self.quant_config.w2_bias

    @property
    def g1_alphas(self) -> Optional[torch.Tensor]:
        return self.quant_config.g1_alphas

    @property
    def g2_alphas(self) -> Optional[torch.Tensor]:
        return self.quant_config.g2_alphas

    # TODO (bnell): make this return a CHUNK_SIZE or None instead?
    @abstractmethod
    def supports_chunking(self) -> bool:
        """
        A flag indicating whether or not this class supports activation
        chunking.
        """
        raise NotImplementedError

    @abstractmethod
    def supports_expert_map(self) -> bool:
        """
        A flag indicating whether or not this class supports expert maps
        """
        raise NotImplementedError

    @abstractmethod
    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        """
        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.

        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.
        - Workspace type: The dtype to use for the workspace tensors.
        - Note: in order for activation chunking to work, the first dimension
          of each tuple must be the number of tokens.
        """
        raise NotImplementedError

    def activation(self, activation: str, output: torch.Tensor,
                   input: torch.Tensor) -> None:
        assert output.size(-1) * 2 == input.size(-1)
        if activation == "silu":
            torch.ops._C.silu_and_mul(output, input)
        elif activation == "gelu":
            torch.ops._C.gelu_and_mul(output, input)
        else:
            raise ValueError(f"Unsupported FusedMoe activation: {activation}")

    def enable_chunking(self):
        return envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and \
          self.supports_chunking()

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        """
        This function computes the intermediate result of a Mixture of Experts
        (MoE) layer using two sets of weights, w1 and w2.

        Parameters:
        - 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.
        - 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

a1_gscale `property` ¶

a1_gscale: Optional[Tensor]

a1_scale `property` ¶

a1_scale: Optional[Tensor]

a2_gscale `property` ¶

a2_gscale: Optional[Tensor]

a2_scale `property` ¶

a2_scale: Optional[Tensor]

activation_formats `abstractmethod` `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

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

block_shape `property` ¶

block_shape: Optional[list[int]]

g1_alphas `property` ¶

g1_alphas: Optional[Tensor]

g2_alphas `property` ¶

g2_alphas: Optional[Tensor]

per_act_token_quant `property` ¶

per_act_token_quant: bool

per_out_ch_quant `property` ¶

per_out_ch_quant: bool

quant_config `instance-attribute` ¶

quant_config = quant_config

quant_dtype `property` ¶

quant_dtype: Optional[dtype]

w1_bias `property` ¶

w1_bias: Optional[Tensor]

w1_scale `property` ¶

w1_scale: Optional[Tensor]

w1_zp `property` ¶

w1_zp: Optional[Tensor]

w2_bias `property` ¶

w2_bias: Optional[Tensor]

w2_scale `property` ¶

w2_scale: Optional[Tensor]

w2_zp `property` ¶

w2_zp: Optional[Tensor]

init ¶

__init__(quant_config: FusedMoEQuantConfig)

quant_config: Quantization parameters for this experts instance.

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

def __init__(
    self,
    quant_config: FusedMoEQuantConfig,
):
    """
    quant_config: Quantization parameters for this experts instance.
    """
    self.quant_config = quant_config

activation ¶

activation(
    activation: str, output: Tensor, input: Tensor
) -> None

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

def activation(self, activation: str, output: torch.Tensor,
               input: torch.Tensor) -> None:
    assert output.size(-1) * 2 == input.size(-1)
    if activation == "silu":
        torch.ops._C.silu_and_mul(output, input)
    elif activation == "gelu":
        torch.ops._C.gelu_and_mul(output, input)
    else:
        raise ValueError(f"Unsupported FusedMoe activation: {activation}")

apply `abstractmethod` ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

Parameters: - 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. - 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.

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    """
    This function computes the intermediate result of a Mixture of Experts
    (MoE) layer using two sets of weights, w1 and w2.

    Parameters:
    - 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.
    - 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

enable_chunking ¶

enable_chunking()

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

def enable_chunking(self):
    return envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and \
      self.supports_chunking()

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:
    raise NotImplementedError

supports_chunking `abstractmethod` ¶

supports_chunking() -> bool

A flag indicating whether or not this class supports activation chunking.

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

@abstractmethod
def supports_chunking(self) -> bool:
    """
    A flag indicating whether or not this class supports activation
    chunking.
    """
    raise NotImplementedError

supports_expert_map `abstractmethod` ¶

supports_expert_map() -> bool

A flag indicating whether or not this class supports expert maps

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

@abstractmethod
def supports_expert_map(self) -> bool:
    """
    A flag indicating whether or not this class supports expert maps
    """
    raise NotImplementedError

workspace_shapes `abstractmethod` ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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.

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. - Workspace type: The dtype to use for the workspace tensors. - Note: in order for activation chunking to work, the first dimension of each tuple must be the number of tokens.

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

@abstractmethod
def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    """
    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.

    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.
    - Workspace type: The dtype to use for the workspace tensors.
    - Note: in order for activation chunking to work, the first dimension
      of each tuple must be the number of tokens.
    """
    raise NotImplementedError

FusedMoEPrepareAndFinalize ¶

Bases: ABC

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

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.
    """

    @abstractmethod
    def prepare(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: Optional[torch.Tensor],
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> 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.

        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 supports_async(self) -> bool:
        """
        Indicates whether or not this class implements prepare_async and
        finalize_async.
        """
        return False

    def prepare_async(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: Optional[torch.Tensor],
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> Union[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.

        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,
    ) -> Union[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

    @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) -> Optional[torch.dtype]:
        """
        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) -> Optional[int]:
        """
        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

activation_format `abstractmethod` `property` ¶

activation_format: FusedMoEActivationFormat

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

finalize `abstractmethod` ¶

finalize(
    output: Tensor,
    fused_expert_output: Tensor,
    topk_weights: Tensor,
    topk_ids: 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.

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 ¶

finalize_async(
    output: Tensor,
    fused_expert_output: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    apply_router_weight_on_input: bool,
    weight_and_reduce_impl: TopKWeightAndReduce,
) -> Union[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, ...)

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,
) -> Union[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

max_num_tokens_per_rank `abstractmethod` ¶

max_num_tokens_per_rank() -> Optional[int]

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) -> Optional[int]:
    """
    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

num_dispatchers `abstractmethod` ¶

num_dispatchers() -> int

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

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

prepare `abstractmethod` ¶

prepare(
    a1: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    num_experts: int,
    expert_map: Optional[Tensor],
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> 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.

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: Optional[torch.Tensor],
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> 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.

    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 ¶

prepare_async(
    a1: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    num_experts: int,
    expert_map: Optional[Tensor],
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> Union[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.

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: Optional[torch.Tensor],
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> Union[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.

    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

supports_async ¶

supports_async() -> bool

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

topk_indices_dtype `abstractmethod` ¶

topk_indices_dtype() -> Optional[dtype]

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) -> Optional[torch.dtype]:
    """
    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

FusedMoeWeightScaleSupported ¶

Bases: Enum

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

class FusedMoeWeightScaleSupported(Enum):
    TENSOR = "tensor"
    CHANNEL = "channel"
    GROUP = "group"
    BLOCK = "block"

BLOCK `class-attribute` `instance-attribute` ¶

BLOCK = 'block'

CHANNEL `class-attribute` `instance-attribute` ¶

CHANNEL = 'channel'

GROUP `class-attribute` `instance-attribute` ¶

GROUP = 'group'

TENSOR `class-attribute` `instance-attribute` ¶

TENSOR = 'tensor'

TritonExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

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

class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

    def __init__(
        self,
        quant_config: FusedMoEQuantConfig,
    ):
        super().__init__(quant_config)

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        return (mk.FusedMoEActivationFormat.Standard,
                mk.FusedMoEActivationFormat.Standard)

    def supports_chunking(self) -> bool:
        return True

    def supports_expert_map(self) -> bool:
        return True

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        return TopKWeightAndReduceNoOP()

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        workspace1 = (M, topk, max(N // 2, K))
        workspace2 = (M, topk, max(N, K))
        output = (M, K)
        return (workspace1, workspace2, output, a.dtype)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        # Check constraints.
        if self.quant_config.use_int4_w4a16:
            assert hidden_states.size(-1) // 2 == w1.size(2), (
                "Hidden size mismatch")
        else:
            assert hidden_states.size(-1) == w1.size(2), \
                (f"Hidden size mismatch {hidden_states.size(-1)} "
                 f"!= {w1.size(2)}")

        assert hidden_states.is_contiguous(
        ), "Hidden_states must be contiguous"
        assert hidden_states.dim() == 2
        assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
        assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
        assert hidden_states.dtype in [
            torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
        ]

        E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
            hidden_states, w1, w2, topk_ids)

        if global_num_experts == -1:
            global_num_experts = E

        config = try_get_optimal_moe_config(
            w1.size(),
            w2.size(),
            top_k_num,
            self.quant_config.config_name(hidden_states.dtype),
            num_tokens,
            block_shape=self.block_shape,
        )

        if hidden_states.dtype == torch.bfloat16:
            compute_type = tl.bfloat16
        elif hidden_states.dtype == torch.float16:
            compute_type = tl.float16
        elif hidden_states.dtype == torch.float32:
            compute_type = tl.float32
        elif hidden_states.dtype == torch.float8_e4m3fn:
            compute_type = tl.bfloat16
        else:
            raise ValueError(
                f"Unsupported compute_type: {hidden_states.dtype}")

        # Note that the output tensor might be in workspace1
        intermediate_cache1 = _resize_cache(workspace2,
                                            (num_tokens, top_k_num, N))
        intermediate_cache2 = _resize_cache(workspace13,
                                            (num_tokens * top_k_num, N // 2))
        intermediate_cache3 = _resize_cache(workspace2,
                                            (num_tokens, top_k_num, K))

        sorted_token_ids, expert_ids, num_tokens_post_padded = (
            moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],
                                 global_num_experts, expert_map))

        invoke_fused_moe_kernel(
            hidden_states,
            w1,
            intermediate_cache1,
            a1q_scale,
            self.w1_scale,
            self.w1_zp,
            None,  # topk_weights
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            False,  # mul_routed_weights
            top_k_num,
            config,
            compute_type=compute_type,
            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
            use_int8_w8a8=self.quant_config.use_int8_w8a8,
            use_int8_w8a16=self.quant_config.use_int8_w8a16,
            use_int4_w4a16=self.quant_config.use_int4_w4a16,
            per_channel_quant=self.per_act_token_quant,
            block_shape=self.block_shape,
            B_bias=self.w1_bias,
        )

        self.activation(activation, intermediate_cache2,
                        intermediate_cache1.view(-1, N))

        a2q_scale: Optional[torch.Tensor] = None

        qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
            intermediate_cache2, a2_scale, self.quant_dtype,
            self.per_act_token_quant, self.block_shape)

        invoke_fused_moe_kernel(
            qintermediate_cache2,
            w2,
            intermediate_cache3,
            a2q_scale,
            self.w2_scale,
            self.w2_zp,
            topk_weights,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            not apply_router_weight_on_input,
            1,
            config,
            compute_type=compute_type,
            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
            use_int8_w8a8=self.quant_config.use_int8_w8a8,
            use_int8_w8a16=self.quant_config.use_int8_w8a16,
            use_int4_w4a16=self.quant_config.use_int4_w4a16,
            per_channel_quant=self.per_act_token_quant,
            block_shape=self.block_shape,
            B_bias=self.w2_bias,
        )

        ops.moe_sum(intermediate_cache3, output)

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

init ¶

__init__(quant_config: FusedMoEQuantConfig)

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

def __init__(
    self,
    quant_config: FusedMoEQuantConfig,
):
    super().__init__(quant_config)

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    # Check constraints.
    if self.quant_config.use_int4_w4a16:
        assert hidden_states.size(-1) // 2 == w1.size(2), (
            "Hidden size mismatch")
    else:
        assert hidden_states.size(-1) == w1.size(2), \
            (f"Hidden size mismatch {hidden_states.size(-1)} "
             f"!= {w1.size(2)}")

    assert hidden_states.is_contiguous(
    ), "Hidden_states must be contiguous"
    assert hidden_states.dim() == 2
    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
    assert hidden_states.dtype in [
        torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
    ]

    E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
        hidden_states, w1, w2, topk_ids)

    if global_num_experts == -1:
        global_num_experts = E

    config = try_get_optimal_moe_config(
        w1.size(),
        w2.size(),
        top_k_num,
        self.quant_config.config_name(hidden_states.dtype),
        num_tokens,
        block_shape=self.block_shape,
    )

    if hidden_states.dtype == torch.bfloat16:
        compute_type = tl.bfloat16
    elif hidden_states.dtype == torch.float16:
        compute_type = tl.float16
    elif hidden_states.dtype == torch.float32:
        compute_type = tl.float32
    elif hidden_states.dtype == torch.float8_e4m3fn:
        compute_type = tl.bfloat16
    else:
        raise ValueError(
            f"Unsupported compute_type: {hidden_states.dtype}")

    # Note that the output tensor might be in workspace1
    intermediate_cache1 = _resize_cache(workspace2,
                                        (num_tokens, top_k_num, N))
    intermediate_cache2 = _resize_cache(workspace13,
                                        (num_tokens * top_k_num, N // 2))
    intermediate_cache3 = _resize_cache(workspace2,
                                        (num_tokens, top_k_num, K))

    sorted_token_ids, expert_ids, num_tokens_post_padded = (
        moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],
                             global_num_experts, expert_map))

    invoke_fused_moe_kernel(
        hidden_states,
        w1,
        intermediate_cache1,
        a1q_scale,
        self.w1_scale,
        self.w1_zp,
        None,  # topk_weights
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        False,  # mul_routed_weights
        top_k_num,
        config,
        compute_type=compute_type,
        use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
        use_int8_w8a8=self.quant_config.use_int8_w8a8,
        use_int8_w8a16=self.quant_config.use_int8_w8a16,
        use_int4_w4a16=self.quant_config.use_int4_w4a16,
        per_channel_quant=self.per_act_token_quant,
        block_shape=self.block_shape,
        B_bias=self.w1_bias,
    )

    self.activation(activation, intermediate_cache2,
                    intermediate_cache1.view(-1, N))

    a2q_scale: Optional[torch.Tensor] = None

    qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
        intermediate_cache2, a2_scale, self.quant_dtype,
        self.per_act_token_quant, self.block_shape)

    invoke_fused_moe_kernel(
        qintermediate_cache2,
        w2,
        intermediate_cache3,
        a2q_scale,
        self.w2_scale,
        self.w2_zp,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        not apply_router_weight_on_input,
        1,
        config,
        compute_type=compute_type,
        use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
        use_int8_w8a8=self.quant_config.use_int8_w8a8,
        use_int8_w8a16=self.quant_config.use_int8_w8a16,
        use_int4_w4a16=self.quant_config.use_int4_w4a16,
        per_channel_quant=self.per_act_token_quant,
        block_shape=self.block_shape,
        B_bias=self.w2_bias,
    )

    ops.moe_sum(intermediate_cache3, output)

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    return TopKWeightAndReduceNoOP()

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    return True

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    return True

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    workspace1 = (M, topk, max(N // 2, K))
    workspace2 = (M, topk, max(N, K))
    output = (M, K)
    return (workspace1, workspace2, output, a.dtype)

TritonOrDeepGemmExperts ¶

Bases: FusedMoEPermuteExpertsUnpermute

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

class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

    def __init__(
        self,
        quant_config: FusedMoEQuantConfig,
        allow_deep_gemm: bool = False,
    ):
        super().__init__(quant_config)

        self.triton_expert = TritonExperts(quant_config)

        self.allow_deep_gemm = (allow_deep_gemm
                                and self.quant_config.use_fp8_w8a8 and
                                self.block_shape == deep_gemm_block_shape())

        self.deep_gemm_expert = DeepGemmExperts(
            self.quant_config) if self.allow_deep_gemm else None

    @property
    def activation_formats(
        self
    ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
        assert (self.deep_gemm_expert is None
                or self.triton_expert.activation_formats
                == self.deep_gemm_expert.activation_formats)
        return self.triton_expert.activation_formats

    def supports_chunking(self) -> bool:
        dge = self.deep_gemm_expert
        te = self.triton_expert
        return ((dge is None or dge.supports_chunking())
                and (te is None or te.supports_chunking()))

    def supports_expert_map(self) -> bool:
        dge = self.deep_gemm_expert
        te = self.triton_expert
        return ((dge is None or dge.supports_expert_map())
                and (te is None or te.supports_expert_map()))

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        dge = self.deep_gemm_expert
        te = self.triton_expert
        dge_war = dge.finalize_weight_and_reduce_impl() if dge else None
        te_war = te.finalize_weight_and_reduce_impl() if te else None
        is_dge_war = dge_war is not None
        is_te_war = te_war is not None

        if is_dge_war and is_te_war:
            assert dge_war == te_war, (
                "Both implementations should agree on WeightAndReduce impls. "
                f"Got dge_war: {dge_war}, and te_war: {te_war}")

        if dge_war is not None:
            return dge_war

        assert te_war is not None
        return te_war

    def workspace_shapes(
        self,
        a: torch.Tensor,
        aq: torch.Tensor,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
        # Note: the deep gemm workspaces are strictly larger than the triton
        # workspaces so we can be pessimistic here and allocate for DeepGemm
        # even if we fall back to triton later, e.g. if expert maps are set.
        if self.allow_deep_gemm and (is_deep_gemm_e8m0_used()
                                     or _valid_deep_gemm_shape(M, N, K)):
            assert self.deep_gemm_expert is not None
            return self.deep_gemm_expert.workspace_shapes(
                a, aq, M, N, K, topk, global_num_experts, local_num_experts,
                expert_tokens_meta)
        else:
            return self.triton_expert.workspace_shapes(a, aq, M, N, K, topk,
                                                       global_num_experts,
                                                       local_num_experts,
                                                       expert_tokens_meta)

    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: str,
        global_num_experts: int,
        expert_map: Optional[torch.Tensor],
        a1q_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
        apply_router_weight_on_input: bool,
    ):
        use_deep_gemm = (self.allow_deep_gemm
                         and (_valid_deep_gemm(hidden_states, w1, w2)
                              or is_deep_gemm_e8m0_used()))

        experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert
        assert experts is not None

        experts.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,
        )

activation_formats `property` ¶

activation_formats: tuple[
    FusedMoEActivationFormat, FusedMoEActivationFormat
]

allow_deep_gemm `instance-attribute` ¶

allow_deep_gemm = (
    allow_deep_gemm
    and use_fp8_w8a8
    and block_shape == deep_gemm_block_shape()
)

deep_gemm_expert `instance-attribute` ¶

deep_gemm_expert = (
    DeepGemmExperts(quant_config)
    if allow_deep_gemm
    else None
)

triton_expert `instance-attribute` ¶

triton_expert = TritonExperts(quant_config)

init ¶

__init__(
    quant_config: FusedMoEQuantConfig,
    allow_deep_gemm: bool = False,
)

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

def __init__(
    self,
    quant_config: FusedMoEQuantConfig,
    allow_deep_gemm: bool = False,
):
    super().__init__(quant_config)

    self.triton_expert = TritonExperts(quant_config)

    self.allow_deep_gemm = (allow_deep_gemm
                            and self.quant_config.use_fp8_w8a8 and
                            self.block_shape == deep_gemm_block_shape())

    self.deep_gemm_expert = DeepGemmExperts(
        self.quant_config) if self.allow_deep_gemm else None

apply ¶

apply(
    output: Tensor,
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    activation: str,
    global_num_experts: int,
    expert_map: Optional[Tensor],
    a1q_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    workspace13: Tensor,
    workspace2: Tensor,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
)

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

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: str,
    global_num_experts: int,
    expert_map: Optional[torch.Tensor],
    a1q_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    workspace13: torch.Tensor,
    workspace2: torch.Tensor,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    apply_router_weight_on_input: bool,
):
    use_deep_gemm = (self.allow_deep_gemm
                     and (_valid_deep_gemm(hidden_states, w1, w2)
                          or is_deep_gemm_e8m0_used()))

    experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert
    assert experts is not None

    experts.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,
    )

finalize_weight_and_reduce_impl ¶

finalize_weight_and_reduce_impl() -> TopKWeightAndReduce

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

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
    dge = self.deep_gemm_expert
    te = self.triton_expert
    dge_war = dge.finalize_weight_and_reduce_impl() if dge else None
    te_war = te.finalize_weight_and_reduce_impl() if te else None
    is_dge_war = dge_war is not None
    is_te_war = te_war is not None

    if is_dge_war and is_te_war:
        assert dge_war == te_war, (
            "Both implementations should agree on WeightAndReduce impls. "
            f"Got dge_war: {dge_war}, and te_war: {te_war}")

    if dge_war is not None:
        return dge_war

    assert te_war is not None
    return te_war

supports_chunking ¶

supports_chunking() -> bool

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

def supports_chunking(self) -> bool:
    dge = self.deep_gemm_expert
    te = self.triton_expert
    return ((dge is None or dge.supports_chunking())
            and (te is None or te.supports_chunking()))

supports_expert_map ¶

supports_expert_map() -> bool

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

def supports_expert_map(self) -> bool:
    dge = self.deep_gemm_expert
    te = self.triton_expert
    return ((dge is None or dge.supports_expert_map())
            and (te is None or te.supports_expert_map()))

workspace_shapes ¶

workspace_shapes(
    a: Tensor,
    aq: Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[ExpertTokensMetadata],
) -> tuple[
    tuple[int, ...], tuple[int, ...], tuple[int, ...], dtype
]

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

def workspace_shapes(
    self,
    a: torch.Tensor,
    aq: torch.Tensor,
    M: int,
    N: int,
    K: int,
    topk: int,
    global_num_experts: int,
    local_num_experts: int,
    expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
    # Note: the deep gemm workspaces are strictly larger than the triton
    # workspaces so we can be pessimistic here and allocate for DeepGemm
    # even if we fall back to triton later, e.g. if expert maps are set.
    if self.allow_deep_gemm and (is_deep_gemm_e8m0_used()
                                 or _valid_deep_gemm_shape(M, N, K)):
        assert self.deep_gemm_expert is not None
        return self.deep_gemm_expert.workspace_shapes(
            a, aq, M, N, K, topk, global_num_experts, local_num_experts,
            expert_tokens_meta)
    else:
        return self.triton_expert.workspace_shapes(a, aq, M, N, K, topk,
                                                   global_num_experts,
                                                   local_num_experts,
                                                   expert_tokens_meta)

_raise_exception ¶

_raise_exception(method: str)

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

def _raise_exception(method: str):
    raise NotImplementedError(
        f"{method} is not implemented as lack of triton.")

activation_without_mul ¶

activation_without_mul(activation: str) -> str

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

def activation_without_mul(activation: str) -> str:
    return activation + "_no_mul"

cutlass_moe_fp4 ¶

cutlass_moe_fp4(
    a: Tensor,
    w1_fp4: Tensor,
    w2_fp4: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    quant_config: FusedMoEQuantConfig,
    m: int,
    n: int,
    k: int,
    e: int,
    expert_map: Optional[Tensor] = None,
    apply_router_weight_on_input: bool = False,
) -> Tensor

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

def cutlass_moe_fp4(
        a: torch.Tensor,
        w1_fp4: torch.Tensor,
        w2_fp4: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        quant_config: FusedMoEQuantConfig,
        m: int,
        n: int,
        k: int,
        e: int,
        expert_map: Optional[torch.Tensor] = None,
        apply_router_weight_on_input: bool = False) -> torch.Tensor:
    assert expert_map is None, ("Expert Parallelism / expert_map "
                                "is currently not supported for "
                                "ModelOptNvFp4FusedMoE's cutlass_moe_fp4.")

    # TODO(bnell): this feels a bit hacky
    # NVFP4 requires two levels of quantization, which involves
    # computing some scaling factors dynamically. This makes it
    # incompatible with the typical prepare -> MoE -> finalize
    # pipeline. Move the quantization logic into the MoE body.
    quant_config = FusedMoEQuantConfig.make(
        quant_dtype=None,  # skip quantization in prepare/finalize
        per_act_token_quant=quant_config.per_act_token_quant,
        per_out_ch_quant=quant_config.per_out_ch_quant,
        block_shape=quant_config.block_shape,
        g1_alphas=quant_config.g1_alphas,
        g2_alphas=quant_config.g2_alphas,
        a1_gscale=quant_config.a1_gscale,
        a2_gscale=quant_config.a2_gscale,
        w1_scale=quant_config.w1_scale,
        w2_scale=quant_config.w2_scale,
    )

    fn = mk.FusedMoEModularKernel(
        MoEPrepareAndFinalizeNoEP(),
        CutlassExpertsFp4(
            max_experts_per_worker=e,
            out_dtype=a.dtype,
            quant_config=quant_config,
            use_batched_format=False,
        ),
    )

    return fn(
        hidden_states=a,
        w1=w1_fp4,
        w2=w2_fp4,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        inplace=False,
        activation="silu",
        global_num_experts=e,
        expert_map=None,
        apply_router_weight_on_input=apply_router_weight_on_input,
    )

cutlass_moe_fp8 ¶

cutlass_moe_fp8(
    a: Tensor,
    w1_q: Tensor,
    w2_q: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    ab_strides1: Tensor,
    ab_strides2: Tensor,
    c_strides1: Tensor,
    c_strides2: Tensor,
    quant_config: FusedMoEQuantConfig,
    activation: str = "silu",
    expert_map: Optional[Tensor] = None,
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
) -> Tensor

This function computes a a8w8-quantized Mixture of Experts (MoE) layer using two sets of quantized weights, w1_q and w2_q, and top-k gating mechanism. The matrix multiplications are implemented with CUTLASS grouped gemm.

a (torch.Tensor): The input tensor to the MoE layer. Shape: [M, K]
w1_q (torch.Tensor): The first set of fp8-quantized expert weights. Shape: [num_experts, K, 2N] (the weights are passed transposed)
w2_q (torch.Tensor): The second set of fp8-quantized expert weights. Shape: [num_experts, N, K] (the weights are passed transposed)
topk_weights (torch.Tensor): The weights of each token->expert mapping.
topk_ids (torch.Tensor): The token->expert mappings.
w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q. Shape: [num_experts] or [num_experts, 2N]
w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q. Shape: [num_experts] or [num_experts, K]
ab_strides1 (torch.Tensor): The input/weight strides for the first gemm. Shape: [num_experts]
ab_strides2 (torch.Tensor): The input/weight strides for the second gemm. Shape: [num_experts]
c_strides1 (torch.Tensor): The output strides for the first gemm. Shape: [num_experts]
c_strides2 (torch.Tensor): The output strides for the second gemm. Shape: [num_experts]
per_act_token (Optional[bool]): Whether the scale is per-token or per-tensor.
activation (str): The activation function to use.
a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a. Shape: scalar or [M]
a2_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize the intermediate result between the gemms. Shape: scalar or [M]
expert_map (Optional[torch.Tensor]): In the case of Expert parallel, every Rank is responsible for a subset of experts. expert_map is a mapping from global expert-id to local expert-id. When expert_map[i] is -1, it means that this Rank is not responsible for global expert-id i.
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.
global_num_experts (int): The total number of experts.

Returns: - torch.Tensor: The fp16 output tensor after applying the MoE layer.

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

def cutlass_moe_fp8(
    a: torch.Tensor,
    w1_q: torch.Tensor,
    w2_q: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    ab_strides1: torch.Tensor,
    ab_strides2: torch.Tensor,
    c_strides1: torch.Tensor,
    c_strides2: torch.Tensor,
    quant_config: FusedMoEQuantConfig,
    activation: str = "silu",
    expert_map: Optional[torch.Tensor] = None,
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
) -> torch.Tensor:
    """
    This function computes a a8w8-quantized Mixture of Experts (MoE) layer
    using two sets of quantized weights, w1_q and w2_q, and top-k gating
    mechanism. The matrix multiplications are implemented with CUTLASS
    grouped gemm.

    Parameters:
    - a (torch.Tensor): The input tensor to the MoE layer.
        Shape: [M, K]
    - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
        Shape: [num_experts, K, 2N] (the weights are passed transposed)
    - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
        Shape: [num_experts, N, K] (the weights are passed transposed)
    - topk_weights (torch.Tensor): The weights of each token->expert mapping.
    - topk_ids (torch.Tensor): The token->expert mappings.
    - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
        Shape: [num_experts] or [num_experts, 2N]
    - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
        Shape: [num_experts] or [num_experts, K]
    - ab_strides1 (torch.Tensor): The input/weight strides for the first gemm.
        Shape: [num_experts]
    - ab_strides2 (torch.Tensor): The input/weight strides for the second gemm.
        Shape: [num_experts]
    - c_strides1 (torch.Tensor): The output strides for the first gemm.
        Shape: [num_experts]
    - c_strides2 (torch.Tensor): The output strides for the second gemm.
        Shape: [num_experts]
    - per_act_token (Optional[bool]): Whether the scale is per-token or
                                      per-tensor.
    - activation (str): The activation function to use.
    - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
        Shape: scalar or [M]
    - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
        quantize the intermediate result between the gemms.
        Shape: scalar or [M]
    - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
        every Rank is responsible for a subset of experts. expert_map is a
        mapping from global expert-id to local expert-id. When expert_map[i]
        is -1, it means that this Rank is not responsible for global
        expert-id i.
    - 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.
    - global_num_experts (int): The total number of experts.

    Returns:
    - torch.Tensor: The fp16 output tensor after applying the MoE layer.
    """
    assert quant_config is not None

    if quant_config.a1_scale is not None:
        assert (quant_config.per_act_token_quant ==
                quant_config.a1_scale.numel() != 1)
    if quant_config.a2_scale is not None:
        assert (quant_config.per_act_token_quant ==
                quant_config.a2_scale.numel() != 1)

    assert (quant_config.w1_scale is None
            or (quant_config.per_out_ch_quant == (quant_config.w1_scale.size(1)
                                                  == w1_q.size(1))))

    num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(
        0)

    fn = mk.FusedMoEModularKernel(
        MoEPrepareAndFinalizeNoEP(),
        CutlassExpertsFp8(
            out_dtype=a.dtype,
            ab_strides1=ab_strides1,
            ab_strides2=ab_strides2,
            c_strides1=c_strides1,
            c_strides2=c_strides2,
            quant_config=quant_config,
        ),
    )

    return fn(
        a,
        w1_q,
        w2_q,
        topk_weights,
        topk_ids,
        activation=activation,
        global_num_experts=num_experts,
        expert_map=expert_map,
        apply_router_weight_on_input=apply_router_weight_on_input,
    )

fused_experts ¶

fused_experts(
    hidden_states: Tensor,
    w1: Tensor,
    w2: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    inplace: bool = False,
    activation: str = "silu",
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
    expert_map: Optional[Tensor] = None,
    quant_config: Optional[FusedMoEQuantConfig] = None,
    allow_deep_gemm: bool = False,
    allow_cutlass_block_scaled_grouped_gemm: bool = False,
) -> Tensor

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

def fused_experts(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    inplace: bool = False,
    activation: str = "silu",
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    quant_config: Optional[FusedMoEQuantConfig] = None,
    allow_deep_gemm: bool = False,
    allow_cutlass_block_scaled_grouped_gemm: bool = False,
) -> torch.Tensor:

    if quant_config is None:
        quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
    use_fp8_w8a8 = quant_config.use_fp8_w8a8

    # For now, disable DeepGemm for small N (<= 512) until better
    # permute/unpermute ops are available.
    # However, on B200, we use DeepGemm for all cases because they only support
    # E8M0 scale, which means we requantize the weight and input to the specific
    # scale. Fallen back to cutlass or triton for some cases would cause
    # accuracy issue.
    if (allow_deep_gemm and quant_config.use_fp8_w8a8 and
        (is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2))):
        assert quant_config is not None
        assert apply_router_weight_on_input is False
        return deep_gemm_moe_fp8(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            inplace=inplace,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=quant_config.w1_scale,
            w2_scale=quant_config.w2_scale,
            a1_scale=quant_config.a1_scale,
            a2_scale=quant_config.a2_scale,
            apply_router_weight_on_input=apply_router_weight_on_input,
        )
    elif (allow_cutlass_block_scaled_grouped_gemm and use_fp8_w8a8
          and _valid_cutlass_block_scaled_grouped_gemm(
              w1, w2, inplace, activation, apply_router_weight_on_input,
              expert_map)):
        assert quant_config is not None
        return run_cutlass_block_scaled_fused_experts(
            a=hidden_states,
            w1=w1,
            w2=w2,
            w1_scale=quant_config.w1_scale,
            w2_scale=quant_config.w2_scale,
            topk_weights=topk_weights,
            topk_ids=topk_ids)
    else:
        return dispatch_fused_experts_func(inplace)(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            apply_router_weight_on_input=apply_router_weight_on_input,
            use_fp8_w8a8=quant_config.use_fp8_w8a8,
            use_int8_w8a8=quant_config.use_int8_w8a8,
            use_int8_w8a16=quant_config.use_int8_w8a16,
            use_int4_w4a16=quant_config.use_int4_w4a16,
            use_mxfp4_w4a4=quant_config.use_mxfp4_w4a4,
            per_channel_quant=quant_config.per_act_token_quant,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=quant_config.w1_scale,
            w2_scale=quant_config.w2_scale,
            w1_zp=quant_config.w1_zp,
            w2_zp=quant_config.w2_zp,
            a1_scale=quant_config.a1_scale,
            a2_scale=quant_config.a2_scale,
            block_shape=quant_config.block_shape,
            w1_bias=quant_config.w1_bias,
            w2_bias=quant_config.w2_bias)

fused_topk ¶

fused_topk(
    hidden_states: Tensor,
    gating_output: Tensor,
    topk: int,
    renormalize: bool,
    indices_type: Optional[dtype] = None,
) -> tuple[Tensor, Tensor, Tensor]

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

def fused_topk(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    indices_type: Optional[torch.dtype] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    assert hidden_states.size(0) == gating_output.size(0), (
        "Number of tokens mismatch")

    M, _ = hidden_states.size()

    topk_weights = torch.empty(M,
                               topk,
                               dtype=torch.float32,
                               device=hidden_states.device)
    topk_ids = torch.empty(
        M,
        topk,
        dtype=torch.int32 if indices_type is None else indices_type,
        device=hidden_states.device)
    token_expert_indices = torch.empty(M,
                                       topk,
                                       dtype=torch.int32,
                                       device=hidden_states.device)

    gating_output_float = gating_output.float()  # TODO(woosuk): Optimize this.

    topk_func = dispatch_topk_func()
    topk_weights, topk_ids = topk_func(topk_weights, topk_ids,
                                       token_expert_indices,
                                       gating_output_float, renormalize)

    return topk_weights, topk_ids, token_expert_indices

get_config ¶

get_config() -> Optional[dict[str, Any]]

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

def get_config() -> Optional[dict[str, Any]]:
    return _config

get_config_file_name ¶

get_config_file_name(
    E: int,
    N: int,
    dtype: Optional[str],
    block_shape: Optional[list[int]] = None,
) -> str

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

def get_config_file_name(E: int,
                         N: int,
                         dtype: Optional[str],
                         block_shape: Optional[list[int]] = None) -> str:
    device_name = current_platform.get_device_name().replace(" ", "_")
    dtype_selector = "" if not dtype else f",dtype={dtype}"
    block_shape_selector = ("" if not block_shape or not all(block_shape) else
                            f",block_shape={block_shape}").replace(" ", "")
    return f"E={E},N={N},device_name={device_name}{dtype_selector}{block_shape_selector}.json"  # noqa: E501

grouped_topk ¶

grouped_topk(
    hidden_states: Tensor,
    gating_output: Tensor,
    topk: int,
    renormalize: bool,
    num_expert_group: int = 0,
    topk_group: int = 0,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[Tensor] = None,
) -> tuple[Tensor, Tensor]

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

@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
def grouped_topk(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    num_expert_group: int = 0,
    topk_group: int = 0,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    e_score_correction_bias: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    if envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK and \
            current_platform.is_cuda() and \
            num_expert_group <= 32 and topk <= 32 and \
            e_score_correction_bias is not None:
        return fused_grouped_topk(
            hidden_states=hidden_states,
            gating_output=gating_output,
            topk=topk,
            renormalize=renormalize,
            e_score_correction_bias=e_score_correction_bias,
            num_expert_group=num_expert_group,
            topk_group=topk_group,
            scoring_func=scoring_func,
            routed_scaling_factor=routed_scaling_factor)

    assert hidden_states.size(0) == gating_output.size(0), (
        "Number of tokens mismatch")

    if scoring_func == "softmax":
        scores = torch.softmax(gating_output, dim=-1)
    elif scoring_func == "sigmoid":
        scores = gating_output.sigmoid()
    else:
        raise ValueError(f"Unsupported scoring function: {scoring_func}")

    num_token = scores.size(0)
    if e_score_correction_bias is not None:
        # Store original scores before applying correction bias. We use biased
        # scores for expert selection but original scores for routing weights
        original_scores = scores
        scores = scores + e_score_correction_bias.unsqueeze(0)
        group_scores = (scores.view(num_token, num_expert_group,
                                    -1).topk(2, dim=-1)[0].sum(dim=-1))
    else:
        group_scores = scores.view(num_token, num_expert_group,
                                   -1).max(dim=-1).values  # [n, n_group]
    group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
                           sorted=False)[1]  # [n, top_k_group]
    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
    score_mask = group_mask.unsqueeze(-1).expand(
        num_token, num_expert_group,
        scores.size(-1) // num_expert_group).reshape(num_token, -1)  # [n, e]
    tmp_scores = scores.masked_fill(~score_mask.bool(),
                                    float("-inf"))  # [n, e]

    if e_score_correction_bias is not None:
        topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
        # Use original unbiased scores for the routing weights
        topk_weights = original_scores.gather(1, topk_ids)
    else:
        topk_weights, topk_ids = torch.topk(tmp_scores,
                                            k=topk,
                                            dim=-1,
                                            sorted=False)

    if renormalize:
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

    if routed_scaling_factor != 1.0:
        topk_weights = topk_weights * routed_scaling_factor
    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)

override_config ¶

override_config(config)

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

@contextmanager
def override_config(config):
    global _config
    old_config = _config
    _config = config
    yield
    _config = old_config

vllm.model_executor.layers.fused_moe ¶

__all__ module-attribute ¶

_config module-attribute ¶

BatchedDeepGemmExperts ¶

activation_formats property ¶

max_num_tokens instance-attribute ¶

num_dispatchers instance-attribute ¶

__init__ ¶

apply ¶

finalize_weight_and_reduce_impl ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

BatchedTritonExperts ¶

activation_formats property ¶

max_num_tokens instance-attribute ¶

num_dispatchers instance-attribute ¶

__init__ ¶

apply ¶

finalize_weight_and_reduce_impl ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

BatchedTritonOrDeepGemmExperts ¶

activation_formats property ¶

allow_deep_gemm instance-attribute ¶

batched_deep_gemm_experts instance-attribute ¶

batched_triton_experts instance-attribute ¶

__init__ ¶

apply ¶

finalize_weight_and_reduce_impl ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

CutlassBatchedExpertsFp8 ¶

activation_formats property ¶

max_experts_per_worker instance-attribute ¶

num_dispatchers instance-attribute ¶

__init__ ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

CutlassExpertsFp8 ¶

activation_formats property ¶

__init__ ¶

finalize_weight_and_reduce_impl ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

DeepGemmExperts ¶

activation_formats property ¶

__init__ ¶

apply ¶

finalize_weight_and_reduce_impl ¶

supports_chunking ¶

supports_expert_map ¶

workspace_shapes ¶

FusedMoE ¶

activation instance-attribute ¶

apply_router_weight_on_input instance-attribute ¶

batched_hidden_states instance-attribute ¶

batched_router_logits instance-attribute ¶

custom_routing_function instance-attribute ¶

dp_rank property ¶

dp_size property ¶

e_score_correction_bias instance-attribute ¶

enable_eplb instance-attribute ¶

ep_rank property ¶

ep_size property ¶

expert_load_view instance-attribute ¶

expert_map instance-attribute ¶

global_num_experts instance-attribute ¶

hidden_size instance-attribute ¶

intermediate_size_per_partition instance-attribute ¶

is_sequence_parallel instance-attribute ¶

layer_name instance-attribute ¶

local_num_experts instance-attribute ¶

logical_replica_count instance-attribute ¶

logical_to_physical_map instance-attribute ¶

moe_config instance-attribute ¶

all `module-attribute` ¶

_config `module-attribute` ¶

activation_formats `property` ¶

max_num_tokens `instance-attribute` ¶

num_dispatchers `instance-attribute` ¶

init ¶

activation_formats `property` ¶

max_num_tokens `instance-attribute` ¶

num_dispatchers `instance-attribute` ¶

init ¶

activation_formats `property` ¶

allow_deep_gemm `instance-attribute` ¶

batched_deep_gemm_experts `instance-attribute` ¶

batched_triton_experts `instance-attribute` ¶

init ¶

activation_formats `property` ¶

max_experts_per_worker `instance-attribute` ¶

num_dispatchers `instance-attribute` ¶

init ¶

activation_formats `property` ¶

init ¶

activation_formats `property` ¶

init ¶

activation `instance-attribute` ¶

apply_router_weight_on_input `instance-attribute` ¶

batched_hidden_states `instance-attribute` ¶

batched_router_logits `instance-attribute` ¶

custom_routing_function `instance-attribute` ¶

dp_rank `property` ¶

dp_size `property` ¶

e_score_correction_bias `instance-attribute` ¶

enable_eplb `instance-attribute` ¶

ep_rank `property` ¶

ep_size `property` ¶

expert_load_view `instance-attribute` ¶

expert_map `instance-attribute` ¶

global_num_experts `instance-attribute` ¶

hidden_size `instance-attribute` ¶

intermediate_size_per_partition `instance-attribute` ¶

is_sequence_parallel `instance-attribute` ¶

layer_name `instance-attribute` ¶

local_num_experts `instance-attribute` ¶

logical_replica_count `instance-attribute` ¶

logical_to_physical_map `instance-attribute` ¶

moe_config `instance-attribute` ¶

moe_parallel_config `instance-attribute` ¶

moe_quant_config `instance-attribute` ¶

num_expert_group `instance-attribute` ¶

params_dtype `instance-attribute` ¶

quant_config `instance-attribute` ¶

quant_method `instance-attribute` ¶

reduce_results `instance-attribute` ¶

renormalize `instance-attribute` ¶

routed_scaling_factor `instance-attribute` ¶

scoring_func `instance-attribute` ¶

shared_experts `property` ¶

sp_size `instance-attribute` ¶

top_k `instance-attribute` ¶

topk_group `instance-attribute` ¶

tp_rank `property` ¶

tp_size `property` ¶

use_deepep_ht_kernels `property` ¶

use_deepep_ll_kernels `property` ¶

use_ep `property` ¶

use_flashinfer_cutlass_kernels `property` ¶

use_grouped_topk `instance-attribute` ¶

use_pplx_kernels `property` ¶

zero_expert_num `instance-attribute` ¶

zero_expert_type `instance-attribute` ¶

init ¶

make_expert_params_mapping `classmethod` ¶

select_experts `staticmethod` ¶

BatchedExperts `class-attribute` `instance-attribute` ¶

Standard `class-attribute` `instance-attribute` ¶

FusedMoEConfig `dataclass` ¶

dp_rank `property` ¶

dp_size `property` ¶

ep_rank `property` ¶

ep_size `property` ¶

experts_per_token `instance-attribute` ¶