vllm.model_executor.layers.fused_moe.all2all_utils ¶

def flashinfer_one_sided_dispatch_layout(
    hidden_dim: int, quant_config: FusedMoEQuantConfig
) -> FlashInferOneSidedDispatchLayout:
    """Return the one-sided activation payload layout."""
    if quant_config.quant_dtype is None:
        return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)
    if quant_config.quant_dtype == "nvfp4":
        scale_elems = hidden_dim // 16
        return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems)
    if quant_config.quant_dtype == "mxfp8":
        align = quant_config.mx_alignment
        padded_k = (
            ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim
        )
        scale_elems = padded_k // 32
        return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems)
    if (
        quant_config.use_fp8_w8a8
        and quant_config.quant_dtype == current_platform.fp8_dtype()
        and quant_config.block_shape == [128, 128]
    ):
        if hidden_dim % 128 != 0:
            raise NotImplementedError(
                "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch "
                f"requires hidden_dim divisible by 128; got {hidden_dim}"
            )
        scale_elems = hidden_dim // 128
        scale_bytes = scale_elems * torch.float32.itemsize
        return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes)
    raise NotImplementedError(
        "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
        "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 "
        "(quant_dtype=None) today; got "
        f"quant_dtype={quant_config.quant_dtype!r}, "
        f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, "
        f"block_shape={quant_config.block_shape!r}"
    )

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