vllm.model_executor.layers.fused_moe.layer

Functions:

• FusedMoEFactory –

  Factory function for creating MoE execution pipeline.

• fused_moe_make_expert_params_mapping –

  Delegate to EPLB manager.

FusedMoEFactory(num_experts, top_k, hidden_size, intermediate_size, intermediate_pad=None, params_dtype=None, renormalize=True, use_grouped_topk=False, num_expert_group=None, topk_group=None, quant_config=None, tp_size=None, dp_size=None, pcp_size=None, prefix='', custom_routing_function=None, router=None, scoring_func='softmax', routed_scaling_factor=1.0, swiglu_limit=None, swiglu_alpha=None, swiglu_beta=None, activation_situ_beta=None, activation_situ_linear_beta=None, e_score_correction_bias=None, apply_router_weight_on_input=False, activation='silu', enable_eplb=False, num_redundant_experts=0, has_bias=False, is_sequence_parallel=False, reduce_results=True, ckpt_names=('gate_proj', 'down_proj', 'up_proj'), is_fused_checkpoint_transposed=False, n_shared_experts=None, fuse_shared_experts=False, shared_expert_prefix=None, router_logits_dtype=None, gate=None, shared_experts=None, shared_expert_gate=None, routed_input_transform=None, routed_output_transform=None, apply_routed_scale_to_output=False, zero_expert_type=None, hash_indices_table=None, bias_vl=None, image_sentinel_lo=0, runner_cls=None, runner_args=None, routed_experts_cls=None, moe_config_cls=None, routed_experts_args=None, skip_padding=False)

Factory function for creating MoE execution pipeline.

Creates and configures a complete MoE execution pipeline including: - Router (for token-to-expert assignment) - RoutedExperts (containing expert weight parameters) - MoERunner (orchestrates the complete forward pass)

The experts contain 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:

• intermediate_pad

  (int | None, default: None ) –

  Padding added to the intermediate size, if any.

• swiglu_alpha

  (float | None, default: None ) –

  Optional alpha parameter for the SwiGLU activation.

• swiglu_beta

  (float | None, default: None ) –

  Optional beta parameter for the SwiGLU activation.

Parameters:

• num_experts

  (int) –

  Number of experts in the model (global count)

• top_k

  (int) –

  Number of experts selected for each token

• hidden_size

  (int) –

  Input hidden state size of the transformer

• intermediate_size

  (int) –

  Intermediate size of the experts

• params_dtype

  (dtype | None, default: None ) –

  Data type for the parameters

• renormalize

  (bool, default: True ) –

  Whether to renormalize the logits in the router

• use_grouped_topk

  (bool, default: False ) –

  Whether to use grouped top-k routing

• num_expert_group

  (int | None, default: None ) –

  Number of expert groups for grouped top-k

• topk_group

  (int | None, default: None ) –

  Top-k value per group for grouped top-k

• quant_config

  (QuantizationConfig | None, default: None ) –

  Quantization configuration

• tp_size

  (int | None, default: None ) –

  Tensor parallelism size (None = use global default)

• dp_size

  (int | None, default: None ) –

  Data parallelism size (None = use global default)

• pcp_size

  (int | None, default: None ) –

  Pipeline context parallelism size (None = use global default)

• prefix

  (str, default: '' ) –

  Layer name prefix for weight loading

• custom_routing_function

  (Callable | None, default: None ) –

  Custom routing function override

• router

  (FusedMoERouter | None, default: None ) –

  Pre-configured router instance (None = create default)

• scoring_func

  (str, default: 'softmax' ) –

  Scoring function for routing ("softmax" or others)

• routed_scaling_factor

  (float, default: 1.0 ) –

  Scaling factor applied to topk_weights or output

• swiglu_limit

  (float | None, default: None ) –

  SwiGLU activation limit

• activation_situ_beta

  (float | None, default: None ) –

  SituGLU activation beta

• activation_situ_linear_beta

  (float | None, default: None ) –

  SituGLU linear beta

• e_score_correction_bias

  (Tensor | None, default: None ) –

  Expert score correction bias tensor

• apply_router_weight_on_input

  (bool, default: False ) –

  Whether to apply router weights on input

• activation

  (str, default: 'silu' ) –

  Activation function name ("silu", "gelu", etc.)

• enable_eplb

  (bool, default: False ) –

  Whether to enable expert parallelism load balancer

• num_redundant_experts

  (int, default: 0 ) –

  Number of redundant experts for EPLB

• has_bias

  (bool, default: False ) –

  Whether expert layers have bias terms

• is_sequence_parallel

  (bool, default: False ) –

  Whether sequence parallelism is enabled

• reduce_results

  (bool, default: True ) –

  Whether to all-reduce the final output. Setting this to False (to fuse the all-reduce downstream) is only honored on the late-AR path.

• ckpt_names

  (tuple[str, str, str | None], default: ('gate_proj', 'down_proj', 'up_proj') ) –

  Checkpoint parameter name tuple (gate_proj, down_proj, up_proj) used for weight loading

• is_fused_checkpoint_transposed

  (bool, default: False ) –

  Whether fused checkpoint weights and block scales use transposed storage.

• n_shared_experts

  (int | None, default: None ) –

  Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path)

• fuse_shared_experts

  (bool, default: False ) –

  Whether to enable shared-expert fusion.

• shared_expert_prefix

  (str | None, default: None ) –

  Checkpoint prefix for the fused shared expert.

• router_logits_dtype

  (dtype | None, default: None ) –

  Data type for router logits buffers

• gate

  (Module | None, default: None ) –

  Pre-configured gate module

• shared_experts

  (Module | None, default: None ) –

  Pre-configured shared experts module

• shared_expert_gate

  (Module | None, default: None ) –

  Pre-configured shared expert gate module

• routed_input_transform

  (Module | None, default: None ) –

  Input transformation module

• routed_output_transform

  (Module | None, default: None ) –

  Output transformation module

• apply_routed_scale_to_output

  (bool, default: False ) –

  Whether to apply routed_scaling_factor to output instead of topk_weights

• zero_expert_type

  (str | None, default: None ) –

  Type of zero expert handling

• hash_indices_table

  (Tensor | None, default: None ) –

  Hash table for expert indices

• bias_vl

  (Tensor | None, default: None ) –

  Vision routing bias for image tokens (Deepseek V4)

• image_sentinel_lo

  (int, default: 0 ) –

  First of five consecutive in-vocab image sentinel ids (0 = vision routing disabled)

• runner_cls

  (type[MoERunner] | None, default: None ) –

  Custom MoERunner class (None = use default MoERunner)

• runner_args

  (dict[str, Any] | None, default: None ) –

  Additional arguments for runner constructor

• routed_experts_cls

  (type[RoutedExperts] | None, default: None ) –

  Custom RoutedExperts class (None = use default)

• moe_config_cls

  (type[FusedMoEConfig] | None, default: None ) –

  Custom FusedMoEConfig class (None = use default)

• routed_experts_args

  (dict[str, Any] | None, default: None ) –

  Additional arguments for routed_experts constructor

• skip_padding

  (bool, default: False ) –

  Whether grouped routing should invalidate padding rows.

Returns:

• MoERunner ( MoERunner ) –

  Configured MoE execution pipeline ready for forward passes

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

def FusedMoEFactory(
    num_experts: int,  # Global number of experts
    top_k: int,
    hidden_size: int,
    intermediate_size: int,
    intermediate_pad: int | None = None,
    params_dtype: torch.dtype | None = None,
    renormalize: bool = True,
    use_grouped_topk: bool = False,
    num_expert_group: int | None = None,
    topk_group: int | None = None,
    quant_config: QuantizationConfig | None = None,
    tp_size: int | None = None,
    dp_size: int | None = None,
    pcp_size: int | None = None,
    prefix: str = "",
    custom_routing_function: Callable | None = None,
    router: FusedMoERouter | None = None,
    scoring_func: str = "softmax",
    routed_scaling_factor: float = 1.0,
    swiglu_limit: float | None = None,
    swiglu_alpha: float | None = None,
    swiglu_beta: float | None = None,
    activation_situ_beta: float | None = None,
    activation_situ_linear_beta: float | None = None,
    e_score_correction_bias: torch.Tensor | None = 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: bool = False,
    reduce_results: bool = True,
    ckpt_names: tuple[str, str, str | None] = ("gate_proj", "down_proj", "up_proj"),
    is_fused_checkpoint_transposed: bool = False,
    n_shared_experts: int | None = None,
    fuse_shared_experts: bool = False,
    shared_expert_prefix: str | None = None,
    router_logits_dtype: torch.dtype | None = None,
    gate: torch.nn.Module | None = None,
    shared_experts: torch.nn.Module | None = None,
    shared_expert_gate: torch.nn.Module | None = None,
    routed_input_transform: torch.nn.Module | None = None,
    routed_output_transform: torch.nn.Module | None = None,
    apply_routed_scale_to_output: bool = False,
    zero_expert_type: str | None = None,
    hash_indices_table: torch.Tensor | None = None,
    bias_vl: torch.Tensor | None = None,
    image_sentinel_lo: int = 0,
    runner_cls: type[MoERunner] | None = None,
    runner_args: dict[str, Any] | None = None,
    routed_experts_cls: type[RoutedExperts] | None = None,
    moe_config_cls: type[FusedMoEConfig] | None = None,
    routed_experts_args: dict[str, Any] | None = None,
    skip_padding: bool = False,
) -> MoERunner:
    """Factory function for creating MoE execution pipeline.

    Creates and configures a complete MoE execution pipeline including:
    - Router (for token-to-expert assignment)
    - RoutedExperts (containing expert weight parameters)
    - MoERunner (orchestrates the complete forward pass)

    The experts contain 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:
        intermediate_pad: Padding added to the intermediate size, if any.
        swiglu_alpha: Optional alpha parameter for the SwiGLU activation.
        swiglu_beta: Optional beta parameter for the SwiGLU activation.

    Args:
        num_experts: Number of experts in the model (global count)
        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
        renormalize: Whether to renormalize the logits in the router
        use_grouped_topk: Whether to use grouped top-k routing
        num_expert_group: Number of expert groups for grouped top-k
        topk_group: Top-k value per group for grouped top-k
        quant_config: Quantization configuration
        tp_size: Tensor parallelism size (None = use global default)
        dp_size: Data parallelism size (None = use global default)
        pcp_size: Pipeline context parallelism size (None = use global default)
        prefix: Layer name prefix for weight loading
        custom_routing_function: Custom routing function override
        router: Pre-configured router instance (None = create default)
        scoring_func: Scoring function for routing ("softmax" or others)
        routed_scaling_factor: Scaling factor applied to topk_weights or output
        swiglu_limit: SwiGLU activation limit
        activation_situ_beta: SituGLU activation beta
        activation_situ_linear_beta: SituGLU linear beta
        e_score_correction_bias: Expert score correction bias tensor
        apply_router_weight_on_input: Whether to apply router weights on input
        activation: Activation function name ("silu", "gelu", etc.)
        enable_eplb: Whether to enable expert parallelism load balancer
        num_redundant_experts: Number of redundant experts for EPLB
        has_bias: Whether expert layers have bias terms
        is_sequence_parallel: Whether sequence parallelism is enabled
        reduce_results: Whether to all-reduce the final output. Setting this
            to False (to fuse the all-reduce downstream) is only honored on
            the late-AR path.
        ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj,
            up_proj) used for weight loading
        is_fused_checkpoint_transposed: Whether fused checkpoint weights and
            block scales use transposed storage.
        n_shared_experts: Number of shared experts to fuse into the routed
            grouped GEMM (ROCm; requires aiter FSE or the router-append path)
        fuse_shared_experts: Whether to enable shared-expert fusion.
        shared_expert_prefix: Checkpoint prefix for the fused shared expert.
        router_logits_dtype: Data type for router logits buffers
        gate: Pre-configured gate module
        shared_experts: Pre-configured shared experts module
        shared_expert_gate: Pre-configured shared expert gate module
        routed_input_transform: Input transformation module
        routed_output_transform: Output transformation module
        apply_routed_scale_to_output: Whether to apply routed_scaling_factor to
                                      output instead of topk_weights
        zero_expert_type: Type of zero expert handling
        hash_indices_table: Hash table for expert indices
        bias_vl: Vision routing bias for image tokens (Deepseek V4)
        image_sentinel_lo: First of five consecutive in-vocab image sentinel
            ids (0 = vision routing disabled)
        runner_cls: Custom MoERunner class (None = use default MoERunner)
        runner_args: Additional arguments for runner constructor
        routed_experts_cls: Custom RoutedExperts class (None = use default)
        moe_config_cls: Custom FusedMoEConfig class (None = use default)
        routed_experts_args: Additional arguments for routed_experts constructor
        skip_padding: Whether grouped routing should invalidate padding rows.

    Returns:
        MoERunner: Configured MoE execution pipeline ready for forward passes

    """
    vllm_config = get_current_vllm_config()

    layer_name = prefix

    moe_activation = MoEActivation.from_str(activation)
    is_act_and_mul = moe_activation.is_gated

    moe_parallel_config = make_parallel_config(
        tp_size=tp_size,
        dp_size=dp_size,
        pcp_size=pcp_size,
        is_sequence_parallel=is_sequence_parallel,
        parallel_config=vllm_config.parallel_config,
    )

    # Resolve the deferred all-reduce request against the parallel config.
    skip_final_all_reduce = (
        not reduce_results
        and not moe_parallel_config.use_all2all_kernels
        and not moe_parallel_config.is_sequence_parallel
        and zero_expert_type is None
    )

    global_num_experts, logical_num_experts, num_fused_shared_experts = (
        determine_expert_counts(
            num_experts,
            num_redundant_experts,
            n_shared_experts,
            fuse_shared_experts,
        )
    )
    expert_substitution = make_expert_substitution(
        vllm_config.model_config, layer_name, num_experts, hidden_size, params_dtype
    )
    if expert_substitution is not None:
        # Route over all logical experts but allocate only the retained ones.
        global_num_experts -= len(expert_substitution.substituted_expert_ids)

    # Initialize EPLB manager (or None?)
    eplb_state: EplbLayerState | None = None
    if enable_eplb:
        use_ep = moe_parallel_config.use_ep
        ep_size = moe_parallel_config.ep_size
        if use_ep and global_num_experts % ep_size != 0:
            raise ValueError(
                f"EPLB currently only supports even distribution of "
                f"experts across ranks. Got {global_num_experts} experts "
                f"and {ep_size} EP ranks."
            )
        eplb_state = EplbLayerState()
    else:
        assert num_redundant_experts == 0, (
            "Redundant experts are only supported with EPLB."
        )

    max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

    # Create ExpertMapManager to handle expert mapping and placement for EP.
    # See ExpertMapManager for a detailed description of what it does and when
    # it is required.
    expert_map_manager = ExpertMapManager(
        max_num_batched_tokens=max_num_batched_tokens,
        top_k=top_k,
        global_num_experts=global_num_experts,
        num_redundant_experts=num_redundant_experts,
        num_expert_group=num_expert_group,
        moe_parallel_config=moe_parallel_config,
        placement_strategy=vllm_config.parallel_config.expert_placement_strategy,
        enable_eplb=eplb_state is not None,
        num_fused_shared_experts=num_fused_shared_experts,
        rocm_aiter_enabled=rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul,
    )

    # TODO(bnell): we should not have to create a router if the kernel is
    # monolithic.
    if router is None:
        router = create_fused_moe_router(
            top_k=top_k,
            global_num_experts=global_num_experts,
            eplb_state=eplb_state,
            renormalize=renormalize,
            use_grouped_topk=use_grouped_topk,
            num_expert_group=num_expert_group,
            topk_group=topk_group,
            custom_routing_function=custom_routing_function,
            scoring_func=scoring_func,
            # When apply_routed_scale_to_output is True, we set the scaling factor
            # to 1.0 so it ends up being a nop. Applying the scale will be handled
            # by the runner in this case.
            # The member variable must be set in the same way as the router since
            # some quantization methods can access it.
            routed_scaling_factor=routed_scaling_factor
            if not apply_routed_scale_to_output
            else 1.0,
            e_score_correction_bias=e_score_correction_bias,
            num_fused_shared_experts=num_fused_shared_experts,
            skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels),
            # Fused shared-expert slot weight. With apply_routed_scale_to_output
            # the runner scales the combined output by routed_scaling_factor, so
            # the shared slot weight must be 1/routed_scaling_factor for its net
            # contribution to be 1.0 (matching the un-scaled separate-MLP add).
            shared_expert_weight=(
                (1.0 / routed_scaling_factor)
                if (
                    apply_routed_scale_to_output
                    and num_fused_shared_experts > 0
                    and routed_scaling_factor
                )
                else 1.0
            ),
            zero_expert_type=zero_expert_type,
            num_logical_experts=logical_num_experts,
            hash_indices_table=hash_indices_table,
            bias_vl=bias_vl,
            image_sentinel_lo=image_sentinel_lo,
        )

    if params_dtype is None:
        params_dtype = torch.get_default_dtype()

    # 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

    moe_config = (moe_config_cls or FusedMoEConfig)(
        num_experts=global_num_experts,
        experts_per_token=top_k,
        hidden_dim=hidden_size,
        intermediate_size=intermediate_size,
        intermediate_pad=intermediate_pad,
        num_local_experts=expert_map_manager.local_num_experts,
        num_logical_experts=logical_num_experts,
        moe_parallel_config=moe_parallel_config,
        in_dtype=moe_in_dtype,
        moe_backend=vllm_config.kernel_config.moe_backend,
        router_logits_dtype=router_logits_dtype,
        max_num_tokens=max_num_batched_tokens,
        elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size,
        has_bias=has_bias,
        is_lora_enabled=vllm_config.lora_config is not None,
        shared_expert_prefix=shared_expert_prefix,
        activation=moe_activation,
        device=vllm_config.device_config.device,
        routing_method=router.kernel_routing_method_type,
        has_hash_routing=hash_indices_table is not None,
        require_decomposed_backend=(
            expert_substitution is not None or not router.supports_monolithic_kernels
        ),
        swiglu_limit=swiglu_limit,
        swiglu_alpha=swiglu_alpha,
        swiglu_beta=swiglu_beta,
        activation_situ_beta=activation_situ_beta,
        activation_situ_linear_beta=activation_situ_linear_beta,
        max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
        skip_final_all_reduce=skip_final_all_reduce,
    )

    logger.debug("FusedMoEConfig = %s", moe_config)

    # Create RoutedExperts instance BEFORE create_weights()
    # This will hold all expert weight parameters
    if expert_substitution is not None:
        if routed_experts_cls not in (None, RoutedExperts) or (
            routed_input_transform is not None or routed_output_transform is not None
        ):
            raise NotImplementedError(
                "expert substitution requires the default RoutedExperts and no "
                "routed input or output transforms"
            )
        routed_experts_cls = SubstitutedRoutedExperts
        routed_experts_args = {
            **(routed_experts_args or {}),
            "expert_substitution": expert_substitution,
        }
    elif routed_experts_cls is None:
        routed_experts_cls = RoutedExperts

    assert params_dtype is not None
    routed_experts = routed_experts_cls(
        layer_name,
        params_dtype,
        moe_config,
        quant_config,
        expert_map_manager=expert_map_manager,
        ckpt_gate_proj_name=ckpt_names[0],
        ckpt_down_proj_name=ckpt_names[1],
        ckpt_up_proj_name=ckpt_names[2],
        is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
        # Extra params that are needed by quant_methods, pass along for now
        # Prefer getting these from other sources, e.g. moe_config or
        # router object
        renormalize=renormalize,
        use_grouped_topk=use_grouped_topk,
        num_expert_group=num_expert_group,
        topk_group=topk_group,
        custom_routing_function=custom_routing_function,
        scoring_func=scoring_func,
        routed_scaling_factor=routed_scaling_factor
        if not apply_routed_scale_to_output
        else 1.0,
        swiglu_limit=swiglu_limit,
        swiglu_alpha=swiglu_alpha,
        swiglu_beta=swiglu_beta,
        # TODO needs to be truncated?
        e_score_correction_bias=router.prepare_monolithic_bias(e_score_correction_bias),
        apply_router_weight_on_input=apply_router_weight_on_input,
        hash_indices_table=hash_indices_table,
        **routed_experts_args if routed_experts_args is not None else {},
    )

    if runner_cls is None:
        runner_cls = MoERunner

    runner = runner_cls(
        layer_name=layer_name,
        moe_config=moe_config,
        router=router,
        routed_experts=routed_experts,
        enable_dbo=vllm_config.parallel_config.enable_dbo,
        gate=gate,
        shared_expert_gate=shared_expert_gate,
        shared_experts=shared_experts,
        routed_input_transform=routed_input_transform,
        routed_output_transform=routed_output_transform,
        # When apply_routed_scale_to_output is True, we allow
        # the scaling factor to be passed to the runner, otherwise
        # we pass 1.0 so it ends up being a nop.
        routed_scaling_factor=routed_scaling_factor
        if apply_routed_scale_to_output
        else 1.0,
        **runner_args if runner_args is not None else {},
    )

    return runner

fused_moe_make_expert_params_mapping(model, ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name, num_experts, num_redundant_experts=0, routed_experts_prefix='routed_experts')

Delegate to EPLB manager.

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

def fused_moe_make_expert_params_mapping(
    model: torch.nn.Module,
    ckpt_gate_proj_name: str,
    ckpt_down_proj_name: str,
    ckpt_up_proj_name: str | None,
    num_experts: int,
    num_redundant_experts: int = 0,
    routed_experts_prefix: str = "routed_experts",
) -> list[tuple[str, str, int, str]]:
    """Delegate to EPLB manager."""
    return RoutedExperts.make_expert_params_mapping(
        model,
        ckpt_gate_proj_name,
        ckpt_down_proj_name,
        ckpt_up_proj_name,
        num_experts,
        num_redundant_experts,
        routed_experts_prefix,
    )