Skip to content

vllm.model_executor.layers.fused_moe.all2all_utils

Functions:

maybe_roundup_layer_hidden_size(hidden_size, act_dtype, moe_parallel_config)

Given layer hidden size and MoE configurations, round up hidden_size if necessary.

Parameters:

  • hidden_size

    (int) –

    Layer hidden-size

  • act_dtype

    (dtype) –

    Data type of the layer activations.

  • moe_parallel_config

    (FusedMoEParallelConfig) –

    Fused MoE parallelization strategy configuration.

Return

Rounded up hidden_size if rounding up is required based on the configs and all2all backend. Original hidden size otherwise.

Source code in vllm/model_executor/layers/fused_moe/all2all_utils.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def maybe_roundup_layer_hidden_size(
    hidden_size: int,
    act_dtype: torch.dtype,
    moe_parallel_config: FusedMoEParallelConfig,
) -> int:
    """
    Given layer hidden size and MoE configurations, round up hidden_size
    if necessary.

    Args:
        hidden_size: Layer hidden-size
        act_dtype: Data type of the layer activations.
        moe_parallel_config: Fused MoE parallelization strategy configuration.

    Return:
        Rounded up hidden_size if rounding up is required based on the configs
        and all2all backend.
        Original hidden size otherwise.
    """
    if moe_parallel_config.use_deepep_ht_kernels:
        hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size, act_dtype
        )

    if moe_parallel_config.use_deepep_ll_kernels:
        hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size
        )

    if moe_parallel_config.use_deepep_v2_kernels:
        hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size, act_dtype
        )

    if moe_parallel_config.use_nixl_ep_kernels:
        hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size
        )

    return hidden_size