vllm.model_executor.layers.fused_moe.oracle.mxfp4

Functions:

• convert_gpt_oss_weight_to_mxfp4_moe_kernel_format –

  Convert loaded weights into backend-specific kernel format.

• convert_weight_to_mxfp4_moe_kernel_format –

  Convert loaded weights into backend-specific kernel format.

• make_mxfp4_moe_kernel –

  Create a FusedMoEKernel for the given MXFP4 backend.

• make_mxfp4_moe_quant_config –

  Create a FusedMoEQuantConfig for the given MXFP4 backend.

• map_mxfp4_backend –

  Map a moe_backend string to its candidate Mxfp4MoeBackends.

• mxfp4_round_up_hidden_size_and_intermediate_size –

  Round up hidden_size and intermediate_size based on backend requirements.

• select_deepseek_v4_mxfp4_moe_backend –

  Select the MXFP4 MoE backend with MXFP8 activation as top priority.

• select_mxfp4_moe_backend –

  Select the primary MXFP4 MoE backend.

_backend_activation_key(backend)

Map backend to its activation key (FP8, MXFP8, or None for BF16).

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

def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
    """Map backend to its activation key (FP8, MXFP8, or None for BF16)."""
    if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        return kFp8Dynamic128Sym
    if backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8:
        return kMxfp8Dynamic
    if backend in (
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
    ):
        return kMxfp8Dynamic
    if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
        return kFp8StaticTensorSym
    if backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
        return kMxfp4Dynamic
    return None  # BF16 activation

_filter_by_activation(backends, requested_activation_key)

Pick variants matching requested_activation_key; without one, prefer BF16 if the list has any, else keep the list as-is so explicit non-BF16 picks (e.g. the _afp8 aliases) still land.

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

def _filter_by_activation(
    backends: list[Mxfp4MoeBackend],
    requested_activation_key: QuantKey | None,
) -> list[Mxfp4MoeBackend]:
    """Pick variants matching ``requested_activation_key``; without one,
    prefer BF16 if the list has any, else keep the list as-is so explicit
    non-BF16 picks (e.g. the ``_afp8`` aliases) still land."""
    if requested_activation_key is not None:
        return [
            b
            for b in backends
            if _backend_activation_key(b) == requested_activation_key
            or b == Mxfp4MoeBackend.EMULATION
        ]
    bf16 = [b for b in backends if _backend_activation_key(b) is None]
    return bf16 if bf16 else backends

_get_priority_backends()

Get available backends in priority order. SM100+ prefers DeepGEMM FP4 / TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the backend-level is_supported_config check filters by device capability).

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

def _get_priority_backends() -> list[Mxfp4MoeBackend]:
    """Get available backends in priority order. SM100+ prefers DeepGEMM FP4 /
    TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the
    backend-level ``is_supported_config`` check filters by device capability).
    """
    if current_platform.is_rocm():
        return [
            Mxfp4MoeBackend.AITER_MXFP4_BF16,
            Mxfp4MoeBackend.EMULATION,
        ]
    if current_platform.is_xpu():
        return [Mxfp4MoeBackend.XPU]
    if current_platform.is_cpu():
        return [Mxfp4MoeBackend.CPU]
    _AVAILABLE_BACKENDS = [
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
        Mxfp4MoeBackend.DEEPGEMM_MXFP4,
        # TRITON_UNFUSED has bug with MTP support
        # TODO re-enable after kernel is fixed
        # TRITON_UNFUSED
        Mxfp4MoeBackend.MARLIN,
        Mxfp4MoeBackend.BATCHED_MARLIN,
    ]
    return _AVAILABLE_BACKENDS

_get_priority_backends_for_gpt_oss()

Available backends in priority order, BF16-act variant before activation-quantized variant within each vendor family.

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

def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
    """Available backends in priority order, BF16-act variant before
    activation-quantized variant within each vendor family."""
    _AVAILABLE_BACKENDS = [
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
        Mxfp4MoeBackend.AITER_MXFP4_BF16,
        Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
        Mxfp4MoeBackend.AITER_MXFP4_FP8,
        Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
        Mxfp4MoeBackend.TRITON,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
        # TRITON_UNFUSED has bug with MTP support
        # TODO re-enable after kernel is fixed
        # TRITON_UNFUSED
        Mxfp4MoeBackend.MARLIN,
        Mxfp4MoeBackend.BATCHED_MARLIN,
        Mxfp4MoeBackend.XPU,
        Mxfp4MoeBackend.CPU,
        Mxfp4MoeBackend.EMULATION,
    ]
    return _AVAILABLE_BACKENDS

_requires_qwen38_tep8_emulation(config, activation_key)

Avoid the inaccurate native gfx950 kernel for Qwen3.8 TEP8.

Qwen3.8 Flash Next's routed experts use the distinctive E=512, H=2560, N=640 W4A4 shape. With eight-way expert parallelism, AITER operates on 64 local experts and pads N to 768. That native path is not numerically reliable on gfx950, while OCP MX emulation preserves model accuracy. Keep this guard exact so other MXFP4 shapes and smaller EP configurations retain the native backend.

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

def _requires_qwen38_tep8_emulation(
    config: FusedMoEConfig,
    activation_key: QuantKey | None,
) -> bool:
    """Avoid the inaccurate native gfx950 kernel for Qwen3.8 TEP8.

    Qwen3.8 Flash Next's routed experts use the distinctive
    ``E=512, H=2560, N=640`` W4A4 shape. With eight-way expert parallelism,
    AITER operates on 64 local experts and pads ``N`` to 768. That native path
    is not numerically reliable on gfx950, while OCP MX emulation preserves
    model accuracy. Keep this guard exact so other MXFP4 shapes and smaller EP
    configurations retain the native backend.
    """
    parallel = config.moe_parallel_config
    if not (
        activation_key == kMxfp4Dynamic
        and parallel.use_ep
        and parallel.ep_size == 8
        and config.num_experts == 512
        and config.num_local_experts == 64
        and config.hidden_dim == 2560
        and config.intermediate_size == 640
    ):
        return False

    if not current_platform.is_rocm():
        return False

    from vllm.platforms.rocm import on_gfx950

    return on_gfx950()

_resolve_activation_key(model_activation_key)

Combine the model-supplied activation key with the user override. Raises on conflict (both set and disagreeing).

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

def _resolve_activation_key(
    model_activation_key: QuantKey | None,
) -> QuantKey | None:
    """Combine the model-supplied activation key with the user override.
    Raises on conflict (both set and disagreeing)."""
    user_override = _user_moe_activation_override()
    if user_override is None:
        return model_activation_key
    if model_activation_key is None or model_activation_key == user_override:
        return user_override
    raise ValueError(
        f"checkpoint declares MoE activation={model_activation_key} but "
        f"quantization_config.moe.activation={user_override}; remove the "
        f"override or align it with the checkpoint."
    )

_user_moe_activation_override()

User's MoE activation override from quantization_config, or None.

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

def _user_moe_activation_override() -> QuantKey | None:
    """User's MoE activation override from quantization_config, or None."""
    args = get_current_vllm_config().model_config.quantization_config
    if not isinstance(args, QuantizationConfigArgs) or args.moe is None:
        return None
    return args.moe.activation

convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(mxfp4_backend, layer, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale, w13_bias=None, w2_bias=None, w13_input_scale=None, w2_input_scale=None, _cache_permute_indices=None)

Convert loaded weights into backend-specific kernel format.

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

def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
    mxfp4_backend: Mxfp4MoeBackend,
    layer: torch.nn.Module,
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
    w13_weight_scale: torch.Tensor,
    w2_weight_scale: torch.Tensor,
    w13_bias: torch.Tensor | None = None,
    w2_bias: torch.Tensor | None = None,
    w13_input_scale: torch.Tensor | None = None,
    w2_input_scale: torch.Tensor | None = None,
    _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    Union[torch.Tensor, "PrecisionConfig"],
    Union[torch.Tensor, "PrecisionConfig"],
    torch.Tensor | None,
    torch.Tensor | None,
]:
    """Convert loaded weights into backend-specific kernel format."""
    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
        )

        return (
            w13_weight.data,
            w2_weight.data,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    num_experts = w13_weight.shape[0]
    intermediate_size = w13_weight.shape[1] // 2
    hidden_size = w13_weight.shape[2] * 2

    sf_block_size = 32  # mxfp4 block size

    if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
        from vllm.model_executor.layers.quantization.utils.humming import (
            convert_to_humming_moe_kernel_format,
        )

        convert_to_humming_moe_kernel_format(
            layer, quant_config={"quant_method": "gpt_oss_mxfp4"}
        )
        return (
            layer.w13_weight,
            layer.w2_weight,
            layer.w13_weight_scale,
            layer.w2_weight_scale,
            getattr(layer, "w13_bias", None),
            getattr(layer, "w2_bias", None),
        )
    elif mxfp4_backend in (
        Mxfp4MoeBackend.MARLIN,
        Mxfp4MoeBackend.BATCHED_MARLIN,
    ):
        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
            prepare_moe_mxfp4_layer_for_marlin,
        )

        return prepare_moe_mxfp4_layer_for_marlin(
            layer,
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend in TRTLLM_BACKENDS:
        assert _cache_permute_indices is not None
        from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
        from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache

        # gemm1_alpha/beta/clamp_limit are created by the expert class
        # (TrtLlmMxfp4ExpertsBase), not on the layer.

        w13_weight = w13_weight.data
        w2_weight = w2_weight.data
        w13_weight_scale = w13_weight_scale.data
        w2_weight_scale = w2_weight_scale.data
        assert w13_bias is not None and w2_bias is not None
        w13_bias = w13_bias.data.to(torch.float32)
        w2_bias = w2_bias.data.to(torch.float32)

        # Swap w1 and w3 as the definition of swiglu is different in trtllm-gen
        def swap_every_two_rows(x, axis=-1):
            shape = x.shape
            if axis < 0:
                axis = len(shape) + axis
            new_shape = list(shape)
            new_shape[axis] = shape[axis] // 2
            new_shape.insert(axis + 1, 2)
            x = x.reshape(*new_shape)
            x = x.flip(axis + 1)
            new_shape = list(shape)
            return x.reshape(*new_shape)

        w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2)
        w13_weight = swap_every_two_rows(w13_weight, -2)
        w13_bias = swap_every_two_rows(w13_bias, -1)

        # Shuffle weights and scaling factors for transposed mma output
        gemm1_weights_shuffled = []
        gemm1_scales_shuffled = []
        gemm2_weights_shuffled = []
        gemm2_scales_shuffled = []
        gemm1_bias_shuffled = []
        gemm2_bias_shuffled = []
        epilogue_tile_m = 128
        for i in range(num_experts):
            # w13 weight
            permute_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w13_weight[i].view(torch.uint8),
                epilogue_tile_m,
            )
            gemm1_weights_shuffled.append(
                w13_weight[i]
                .view(torch.uint8)[permute_indices.to(w13_weight.device)]
                .contiguous()
            )
            # w13 scale
            permute_sf_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w13_weight_scale[i].view(torch.uint8),
                epilogue_tile_m,
                num_elts_per_sf=16,
            )
            gemm1_scales_shuffled.append(
                nvfp4_block_scale_interleave(
                    w13_weight_scale[i]
                    .view(torch.uint8)[permute_sf_indices.to(w13_weight_scale.device)]
                    .contiguous()
                )
            )
            # w13 bias
            permute_bias_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w13_bias[i].clone().reshape(-1, 1),
                epilogue_tile_m,
            )
            gemm1_bias_shuffled.append(
                w13_bias[i]
                .clone()
                .reshape(-1, 1)[permute_bias_indices.to(w13_bias.device)]
                .contiguous()
            )
            # w2 weight
            permute_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w2_weight[i].view(torch.uint8),
                epilogue_tile_m,
            )
            gemm2_weights_shuffled.append(
                w2_weight[i]
                .view(torch.uint8)[permute_indices.to(w2_weight.device)]
                .contiguous()
            )
            # w2 scale
            permute_sf_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w2_weight_scale[i].view(torch.uint8),
                epilogue_tile_m,
                num_elts_per_sf=16,
            )
            gemm2_scales_shuffled.append(
                nvfp4_block_scale_interleave(
                    w2_weight_scale[i]
                    .view(torch.uint8)[permute_sf_indices.to(w2_weight_scale.device)]
                    .contiguous()
                )
            )
            # w2 bias
            permute_indices = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w2_bias[i].clone().reshape(-1, 1),
                epilogue_tile_m,
            )
            gemm2_bias_shuffled.append(
                w2_bias[i]
                .clone()
                .reshape(-1, 1)[permute_indices.to(w2_bias.device)]
                .contiguous()
            )

        w13_weight = torch.stack(gemm1_weights_shuffled)
        w13_weight_scale = (
            torch.stack(gemm1_scales_shuffled)
            .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size)
            .view(torch.float8_e4m3fn)
        )
        w2_weight = torch.stack(gemm2_weights_shuffled)
        w2_weight_scale = (
            torch.stack(gemm2_scales_shuffled)
            .reshape(num_experts, hidden_size, intermediate_size // sf_block_size)
            .view(torch.float8_e4m3fn)
        )
        w13_bias = torch.stack(gemm1_bias_shuffled).reshape(num_experts, -1)
        w2_bias = torch.stack(gemm2_bias_shuffled).reshape(num_experts, -1)

        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend in (
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
    ):
        # De-interleave and swap for w13 weight, bias, and scales
        w13_w = w13_weight.data
        gate_w, up_w = w13_w[:, ::2, :], w13_w[:, 1::2, :]
        deinterleaved_w13_w = torch.cat([gate_w, up_w], dim=1)
        w1_w, w3_w = torch.chunk(deinterleaved_w13_w, 2, dim=1)
        w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)

        assert w13_bias is not None and w2_bias is not None
        w13_b = w13_bias.data.to(torch.float32)
        gate_b, up_b = w13_b[:, ::2], w13_b[:, 1::2]
        deinterleaved_w13_b = torch.cat([gate_b, up_b], dim=1)
        b1, b3 = torch.chunk(deinterleaved_w13_b, 2, dim=-1)
        w13_bias_swapped = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)

        w13_s = w13_weight_scale.data
        gate_s, up_s = w13_s[:, ::2, :], w13_s[:, 1::2, :]
        deinterleaved_w13_s = torch.cat([gate_s, up_s], dim=1)
        s1, s3 = torch.chunk(deinterleaved_w13_s, 2, dim=1)
        w13_scale_swapped = torch.cat([s3, s1], dim=1)

        if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
            from flashinfer import block_scale_interleave

            orig_shape = w13_scale_swapped.shape
            w13_scale_interleaved = block_scale_interleave(
                w13_scale_swapped.view(torch.uint8)
            ).reshape(orig_shape)

            w2_s = w2_weight_scale.data
            orig_shape = w2_s.shape
            w2_scale_interleaved = block_scale_interleave(
                w2_s.view(torch.uint8)
            ).reshape(orig_shape)

            return (
                w13_weight_swapped,
                w2_weight,
                w13_scale_interleaved,
                w2_scale_interleaved,
                w13_bias_swapped,
                w2_bias,
            )

        else:
            assert mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16

            from flashinfer.fused_moe import (
                interleave_moe_scales_for_sm90_mixed_gemm,
                interleave_moe_weights_for_sm90_mixed_gemm,
            )

            w13_weight_interleaved = interleave_moe_weights_for_sm90_mixed_gemm(
                w13_weight_swapped.contiguous(), "fp4"
            )
            w2_weight_interleaved = interleave_moe_weights_for_sm90_mixed_gemm(
                w2_weight.contiguous(), "fp4"
            )
            w31_scales_interleaved = interleave_moe_scales_for_sm90_mixed_gemm(
                w13_scale_swapped.to(torch.uint8)
            )
            w2_scale_interleaved = interleave_moe_scales_for_sm90_mixed_gemm(
                w2_weight_scale.data.to(torch.uint8)
            )

            return (
                w13_weight_interleaved,
                w2_weight_interleaved,
                w31_scales_interleaved,
                w2_scale_interleaved,
                w13_bias_swapped,
                w2_bias,
            )

    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
        from vllm._aiter_ops import rocm_aiter_ops

        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)

        # e8m0_shuffle on weight scales (GFX950 swizzle layout)
        from aiter.utility.fp4_utils import e8m0_shuffle

        s0, s1, _ = w13_weight_scale.shape
        w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view(
            s0, s1, -1
        )

        s0, s1, _ = w2_weight_scale.shape
        w2_weight_scale.data = e8m0_shuffle(w2_weight_scale.view(s0 * s1, -1)).view(
            s0, s1, -1
        )

        # View as native FP4 dtype
        fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None)
        if fp4_dtype is not None:
            w13_weight.data = w13_weight.data.view(fp4_dtype)
            w2_weight.data = w2_weight.data.view(fp4_dtype)

        # Shuffle weights for AITER CK kernel
        shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
            w13_weight, w2_weight
        )
        shuffled_w13.is_shuffled = True
        shuffled_w2.is_shuffled = True

        return (
            shuffled_w13,
            shuffled_w2,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
        from vllm._aiter_ops import rocm_aiter_ops

        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)

        e, n, k = w13_weight.shape

        # De-interleave w13 rows: gate/up pairs -> contiguous gate, up blocks
        w13_weight.view(torch.uint8).copy_(
            w13_weight.data.view(torch.uint8)
            .view(e, n // 2, 2, k)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(e, n, k)
        )
        w13_weight_scale.data = (
            w13_weight_scale.data.view(e, n // 2, 2, -1)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(e, n, -1)
        )

        # View as native FP4 dtype for AITER shuffle
        w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
        w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)

        # Shuffle weights and scales for AITER CK kernel layout
        w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True)
        shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
            w13_weight_scale.view(-1, w13_weight_scale.shape[-1]),
            num_experts,
            True,
        )

        w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False)
        shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
            w2_weight_scale.view(-1, w2_weight_scale.shape[-1]),
            num_experts,
            False,
        )

        # Permute bias to match de-interleaved weight layout
        if w13_bias is not None:
            w13_bias = (
                w13_bias.data.view(-1, n // 2, 2)
                .permute(0, 2, 1)
                .contiguous()
                .view(-1, n)
            )

        w13_weight.is_shuffled = True
        w2_weight.is_shuffled = True

        return (
            w13_weight,
            w2_weight,
            shuffled_w13_scale,
            shuffled_w2_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
        # W4A8: MXFP4 weights + static FP8 activations (triton kernel)
        from triton_kernels.numerics import InFlexData

        if w13_bias is not None:
            w13_bias = w13_bias.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.to(torch.float32)

        # Process static FP8 input scales (reduce to scalar, warn if not uniform)
        w13_input_scale = layer.w13_input_scale
        w2_input_scale = layer.w2_input_scale
        if w13_input_scale is None or w2_input_scale is None:
            raise ValueError(
                "W4A8 (AITER_MXFP4_FP8) requires static input scales, but found "
                "w13_input_scale or w2_input_scale is None."
            )
        if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
            logger.warning_once(
                "Found input_scales that are not equal for "
                "fp8 MoE layer. Using the maximum across experts "
                "for each layer."
            )
        w13_input_scale = w13_input_scale.max().to(torch.float32)
        w2_input_scale = w2_input_scale.max().to(torch.float32)

        # Swizzle weights for GFX950
        w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(w13_weight, w13_weight_scale)
        w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(w2_weight, w2_weight_scale)

        # Create InFlexData for activation scales
        lhs_data13 = InFlexData(scale=w13_input_scale)
        lhs_data2 = InFlexData(scale=w2_input_scale)

        # Create PrecisionConfig with both weight and activation info
        w13_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w13_scale),
            flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
        )
        w2_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w2_scale),
            flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
        )

        del layer.w13_weight
        del layer.w2_weight

        return (
            w13_weight,
            w2_weight,
            w13_precision_config,
            w2_precision_config,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend in TRITON_BACKENDS:
        if w13_bias is not None:
            w13_bias = w13_bias.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.to(torch.float32)

        w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
            w13_weight,
            w13_weight_scale,
        )
        w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
            w2_weight,
            w2_weight_scale,
        )

        w13_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
        )
        w2_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
        )

        # The original mxfp4 block scales have been swizzled into the
        # precision configs above and are no longer read by the kernel, so
        # drop the now-dead weight/scale Parameters to free their memory.
        del layer.w13_weight
        del layer.w2_weight
        del layer.w13_weight_scale
        del layer.w2_weight_scale

        return (
            w13_weight,
            w2_weight,
            w13_precision_config,
            w2_precision_config,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.XPU:
        # No additional transformation needed for XPU backend
        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.CPU:
        from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
            prepare_mxfp4_moe_layer_for_cpu,
        )

        packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = (
            prepare_mxfp4_moe_layer_for_cpu(
                w13_weight.data,
                w2_weight.data,
                w13_weight_scale.data,
                w2_weight_scale.data,
            )
        )
        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)
        return (
            packed_w13,
            packed_w2,
            packed_w13_scale,
            packed_w2_scale,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.EMULATION:
        w13_has_per_expert_scale = (
            w13_input_scale is not None
            and w13_input_scale.ndim == 1
            and not all_close_1d(w13_input_scale)
        )
        w2_has_per_expert_scale = (
            w2_input_scale is not None
            and w2_input_scale.ndim == 1
            and not all_close_1d(w2_input_scale)
        )
        if w13_has_per_expert_scale or w2_has_per_expert_scale:
            logger.warning_once(
                "Found input_scales that are not equal for OCP MX MoE "
                "emulation. Using the maximum across experts for each layer."
            )
        if w13_input_scale is not None:
            layer.w13_input_scale = torch.nn.Parameter(
                w13_input_scale.max().to(torch.float32), requires_grad=False
            )
        if w2_input_scale is not None:
            layer.w2_input_scale = torch.nn.Parameter(
                w2_input_scale.max().to(torch.float32), requires_grad=False
            )
        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )
    else:
        raise ValueError(
            f"Unsupported mxfp4_backend: {mxfp4_backend}: "
            f"should be one of: {list(Mxfp4MoeBackend)}."
        )

convert_weight_to_mxfp4_moe_kernel_format(mxfp4_backend, layer, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale, w13_bias=None, w2_bias=None, _cache_permute_indices=None, activation=None, use_separated_a4w4=False)

Convert loaded weights into backend-specific kernel format.

Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends.

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

def convert_weight_to_mxfp4_moe_kernel_format(
    mxfp4_backend: Mxfp4MoeBackend,
    layer: torch.nn.Module,
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
    w13_weight_scale: torch.Tensor,
    w2_weight_scale: torch.Tensor,
    w13_bias: torch.Tensor | None = None,
    w2_bias: torch.Tensor | None = None,
    _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
    activation: MoEActivation | None = None,
    use_separated_a4w4: bool = False,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    Union[torch.Tensor, "PrecisionConfig"],
    Union[torch.Tensor, "PrecisionConfig"],
    torch.Tensor | None,
    torch.Tensor | None,
]:
    """Convert loaded weights into backend-specific kernel format.

    Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends.
    """
    is_gfx1250 = False
    if current_platform.is_rocm():
        from vllm.platforms.rocm import on_gfx1250

        is_gfx1250 = on_gfx1250()

    if mxfp4_backend in B12X_BACKENDS:
        return (
            w13_weight.data,
            w2_weight.data,
            w13_weight_scale.data,
            w2_weight_scale.data,
            w13_bias,
            w2_bias,
        )

    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
        )

        return (
            w13_weight.data,
            w2_weight.data,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
        from vllm.model_executor.layers.quantization.utils.humming import (
            convert_to_humming_moe_kernel_format,
        )

        convert_to_humming_moe_kernel_format(
            layer, quant_config={"quant_method": "mxfp4"}
        )
        return (
            layer.w13_weight,
            layer.w2_weight,
            layer.w13_weight_scale,
            layer.w2_weight_scale,
            getattr(layer, "w13_bias", None),
            getattr(layer, "w2_bias", None),
        )

    if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
            prepare_moe_mxfp4_layer_for_marlin,
        )

        return prepare_moe_mxfp4_layer_for_marlin(
            layer,
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    num_experts = w13_weight.shape[0]
    intermediate_size = w13_weight.shape[1] // 2
    hidden_size = w13_weight.shape[2] * 2

    sf_block_size = 32  # mxfp4 block size

    if mxfp4_backend in TRTLLM_BACKENDS:
        assert _cache_permute_indices is not None
        from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
        from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache

        w13_weight = w13_weight.data
        w2_weight = w2_weight.data
        w13_weight_scale = w13_weight_scale.data
        w2_weight_scale = w2_weight_scale.data
        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)

        # Swap w1/w3 and interleave to match TRTLLM SwiGLU convention.
        # Standard loading gives contiguous [w1/gate, w3/up].
        # TRTLLM kernel expects interleaved [w3_0, w1_0, w3_1, w1_1, ...].
        w1_weight = w13_weight[:, :intermediate_size, :]
        w3_weight = w13_weight[:, intermediate_size:, :]
        w13_weight = torch.stack([w3_weight, w1_weight], dim=2).reshape(
            w13_weight.shape
        )

        w1_scale = w13_weight_scale[:, :intermediate_size, :]
        w3_scale = w13_weight_scale[:, intermediate_size:, :]
        w13_weight_scale = torch.stack([w3_scale, w1_scale], dim=2).reshape(
            w13_weight_scale.shape
        )

        if w13_bias is not None:
            b1 = w13_bias[:, :intermediate_size]
            b3 = w13_bias[:, intermediate_size:]
            w13_bias = torch.stack([b3, b1], dim=2).reshape(w13_bias.shape)

        # Shuffle weights and scaling factors for transposed mma output.
        # Permute indices depend only on shape (cached by torch.Size),
        # so compute once and apply to all experts via batched indexing.
        epilogue_tile_m = 128

        # w13 weight permute
        w13_perm = get_w2_permute_indices_with_cache(
            _cache_permute_indices,
            w13_weight[0].view(torch.uint8),
            epilogue_tile_m,
        ).to(w13_weight.device)
        w13_weight = w13_weight.view(torch.uint8)[:, w13_perm].contiguous()

        # w13 scale permute + interleave
        w13_sf_perm = get_w2_permute_indices_with_cache(
            _cache_permute_indices,
            w13_weight_scale[0].view(torch.uint8),
            epilogue_tile_m,
            num_elts_per_sf=16,
        ).to(w13_weight_scale.device)
        w13_s = w13_weight_scale.view(torch.uint8)[:, w13_sf_perm].contiguous()
        E, N_s, K_s = w13_s.shape
        w13_weight_scale = (
            nvfp4_block_scale_interleave(w13_s.reshape(E * N_s, K_s))
            .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size)
            .view(torch.float8_e4m3fn)
        )

        # w2 weight permute
        w2_perm = get_w2_permute_indices_with_cache(
            _cache_permute_indices,
            w2_weight[0].view(torch.uint8),
            epilogue_tile_m,
        ).to(w2_weight.device)
        w2_weight = w2_weight.view(torch.uint8)[:, w2_perm].contiguous()

        # w2 scale permute + interleave
        w2_sf_perm = get_w2_permute_indices_with_cache(
            _cache_permute_indices,
            w2_weight_scale[0].view(torch.uint8),
            epilogue_tile_m,
            num_elts_per_sf=16,
        ).to(w2_weight_scale.device)
        w2_s = w2_weight_scale.view(torch.uint8)[:, w2_sf_perm].contiguous()
        E2, N2_s, K2_s = w2_s.shape
        w2_weight_scale = (
            nvfp4_block_scale_interleave(w2_s.reshape(E2 * N2_s, K2_s))
            .reshape(num_experts, hidden_size, intermediate_size // sf_block_size)
            .view(torch.float8_e4m3fn)
        )

        # w13 bias permute
        if w13_bias is not None:
            w13_b_perm = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w13_bias[0].reshape(-1, 1),
                epilogue_tile_m,
            ).to(w13_bias.device)
            w13_bias = w13_bias.reshape(num_experts, -1, 1)[:, w13_b_perm].reshape(
                num_experts, -1
            )

        # w2 bias permute
        if w2_bias is not None:
            w2_b_perm = get_w2_permute_indices_with_cache(
                _cache_permute_indices,
                w2_bias[0].reshape(-1, 1),
                epilogue_tile_m,
            ).to(w2_bias.device)
            w2_bias = w2_bias.reshape(num_experts, -1, 1)[:, w2_b_perm].reshape(
                num_experts, -1
            )

        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250:
        # Initially introduced for DeepSeekV4

        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)

        if activation == MoEActivation.SITU:
            from aiter.utility.fp4_utils import e8m0_shuffle

            from vllm._aiter_ops import rocm_aiter_ops

            fp4_dtype = torch.float4_e2m1fn_x2
            e8m0_dtype = torch.float8_e8m0fnu
            guinterleave = rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved()
            w13 = rocm_aiter_ops.shuffle_weight_a16w4(
                w13_weight.data.view(fp4_dtype), 16, guinterleave
            )
            w2 = rocm_aiter_ops.shuffle_weight_a16w4(
                w2_weight.data.view(fp4_dtype), 16, False
            )
            w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype)
            w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype)
            w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
                w13_scale_raw.view(-1, w13_scale_raw.shape[-1]),
                num_experts,
                guinterleave,
            )
            w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))
            w13.is_shuffled = True
            w2.is_shuffled = True
            return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)

        import os

        # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to
        # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2
        # a4w4 is separated and selects q_dtype_a independently, so the bound
        # is unused on that path.
        os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"

        from aiter.ops.shuffle import shuffle_scale as _shuf_s
        from aiter.ops.shuffle import shuffle_weight as _shuf_w

        # DeepSeek V4.1 a4w4 uses ATOM's SEPARATED gate/up layout instead of
        # the default INTERLEAVE shuffle (INTERLEAVE + fp4x2 has no tuned
        # kernel and produces garbage output). Must match GateMode.SEPARATED
        # in rocm_aiter_moe.py.
        is_guinterleave = not use_separated_a4w4

        w13_weight = torch.nn.Parameter(
            _shuf_w(
                w13_weight.data.view(torch.float4_e2m1fn_x2),
                is_guinterleave=is_guinterleave,
                gate_up=True,
            ),
            requires_grad=False,
        )
        shuffled_w13_scale = _shuf_s(
            w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]),
            num_experts,
            is_guinterleave,
            True,
        )

        w2_weight = torch.nn.Parameter(
            _shuf_w(
                w2_weight.data.view(torch.float4_e2m1fn_x2),
                is_guinterleave=is_guinterleave,
                gate_up=False,
            ),
            requires_grad=False,
        )
        # use_gu_interleave
        shuffled_w2_scale = _shuf_s(
            w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]),
            num_experts,
            is_guinterleave,
            False,
        )

        w13_weight.is_shuffled = True
        w2_weight.is_shuffled = True

        return (
            w13_weight,
            w2_weight,
            shuffled_w13_scale,
            shuffled_w2_scale,
            w13_bias,
            w2_bias,
        )

    elif mxfp4_backend in TRITON_BACKENDS or (
        mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
    ):
        if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
            # AITER moe_gemm_a16w4 needs gate/up interleaved
            def interleave_gate_up(w: torch.Tensor) -> torch.Tensor:
                gate, up = w.chunk(2, dim=1)
                return torch.stack((gate, up), dim=2).reshape(w.shape)

            w13_weight = interleave_gate_up(w13_weight)
            w13_weight_scale = interleave_gate_up(w13_weight_scale)

            if w13_bias is not None:
                w13_bias = interleave_gate_up(w13_bias.to(torch.float32))
        elif mxfp4_backend == Mxfp4MoeBackend.TRITON:

            def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
                shape = w.shape
                n = shape[-1]
                first = w[..., : n // 2]
                second = w[..., n // 2 :]
                stacked = torch.stack((first, second), dim=-1)
                return stacked.reshape(shape)

            w13_weight = shuffle_weight(w13_weight)
            w13_weight_scale = shuffle_weight(w13_weight_scale)

            if w13_bias is not None:
                w13_bias = shuffle_weight(w13_bias.to(torch.float32))
        else:
            if w13_bias is not None:
                w13_bias = w13_bias.to(torch.float32)

        if w2_bias is not None:
            w2_bias = w2_bias.to(torch.float32)

        w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
            w13_weight,
            w13_weight_scale,
        )
        w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
            w2_weight,
            w2_weight_scale,
        )

        w13_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
        )
        w2_precision_config = PrecisionConfig(
            **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
        )

        # The original mxfp4 block scales have been swizzled into the
        # precision configs above and are no longer read by the kernel, so
        # drop the now-dead weight/scale Parameters to free their memory.
        del layer.w13_weight
        del layer.w2_weight
        del layer.w13_weight_scale
        del layer.w2_weight_scale

        return (
            w13_weight,
            w2_weight,
            w13_precision_config,
            w2_precision_config,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend in (
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
    ):
        # Standard checkpoints store fused gate/up tensors as [w1; w3], while
        # FlashInfer CUTLASS consumes [w3; w1]. Keep weights, scales, and bias
        # in the same order before applying the backend-specific interleave.
        w13_weight = swap_w13_to_w31(w13_weight.data)
        w13_weight_scale = swap_w13_to_w31(w13_weight_scale.data)
        if w13_bias is not None:
            b1, b3 = torch.chunk(w13_bias.data, 2, dim=-1)
            w13_bias = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.bfloat16)

        if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
            from flashinfer import block_scale_interleave

            w13_scale_shape = w13_weight_scale.shape
            w13_weight_scale = block_scale_interleave(
                w13_weight_scale.view(torch.uint8)
            ).reshape(w13_scale_shape)

            w2_scale_shape = w2_weight_scale.shape
            w2_weight_scale = block_scale_interleave(
                w2_weight_scale.data.view(torch.uint8)
            ).reshape(w2_scale_shape)
        else:
            from flashinfer.fused_moe import (
                interleave_moe_scales_for_sm90_mixed_gemm,
                interleave_moe_weights_for_sm90_mixed_gemm,
            )

            w13_weight = interleave_moe_weights_for_sm90_mixed_gemm(
                w13_weight.contiguous(), "fp4"
            )
            w2_weight = interleave_moe_weights_for_sm90_mixed_gemm(
                w2_weight.data.contiguous(), "fp4"
            )
            w13_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm(
                w13_weight_scale.to(torch.uint8)
            )
            w2_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm(
                w2_weight_scale.data.to(torch.uint8)
            )

        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend in (
        Mxfp4MoeBackend.XPU,
        Mxfp4MoeBackend.EMULATION,
    ):
        # No additional transformation is needed: XPU consumes the checkpoint
        # layout directly, while emulation dequantizes that layout at runtime.
        return (
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend in (
        Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
        Mxfp4MoeBackend.AITER_MXFP4_FP8,
    ):
        return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
            mxfp4_backend=mxfp4_backend,
            layer=layer,
            w13_weight=w13_weight,
            w2_weight=w2_weight,
            w13_weight_scale=w13_weight_scale,
            w2_weight_scale=w2_weight_scale,
            w13_bias=w13_bias,
            w2_bias=w2_bias,
            _cache_permute_indices=_cache_permute_indices,
        )
    elif mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
        from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
            mxfp4_moe_ep_weights,
        )

        # DeepGEMM consumes MXFP4 directly; CuTeDSL gets bf16 and requantizes.
        weights = mxfp4_moe_ep_weights(
            FLASHINFER_MOE_EP_CUTEDSL
            if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL
            else FLASHINFER_MOE_EP_DEEP_GEMM,
            w13_weight,
            w2_weight,
            w13_weight_scale,
            w2_weight_scale,
        )
        return (
            weights.w13,
            weights.w2,
            weights.w13_scale,
            weights.w2_scale,
            w13_bias,
            w2_bias,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.CPU:
        from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
            prepare_mxfp4_moe_layer_for_cpu,
        )

        packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = (
            prepare_mxfp4_moe_layer_for_cpu(
                w13_weight.data,
                w2_weight.data,
                w13_weight_scale.data,
                w2_weight_scale.data,
            )
        )
        if w13_bias is not None:
            w13_bias = w13_bias.data.to(torch.float32)
        if w2_bias is not None:
            w2_bias = w2_bias.data.to(torch.float32)
        return (
            packed_w13,
            packed_w2,
            packed_w13_scale,
            packed_w2_scale,
            w13_bias,
            w2_bias,
        )
    else:
        raise ValueError(
            f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
            "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, "
            "CPU, or emulation backend."
        )

make_mxfp4_moe_kernel(moe_quant_config, moe_config, experts_cls, mxfp4_backend, routing_tables=None)

Create a FusedMoEKernel for the given MXFP4 backend.

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

def make_mxfp4_moe_kernel(
    moe_quant_config: FusedMoEQuantConfig,
    moe_config: FusedMoEConfig,
    experts_cls: type[mk.FusedMoEExperts],
    mxfp4_backend: Mxfp4MoeBackend,
    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEKernel:
    """Create a FusedMoEKernel for the given MXFP4 backend."""
    is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)

    prepare_finalize = maybe_make_prepare_finalize(
        moe=moe_config,
        quant_config=moe_quant_config,
        routing_tables=routing_tables,
        allow_new_interface=True,
        use_monolithic=is_monolithic,
    )
    assert prepare_finalize is not None

    logger.info_once("Using %s", prepare_finalize.__class__.__name__)
    logger.info_once("Using %s", experts_cls.__name__)

    # Create Experts.
    if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
        max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
        assert max_num_tokens is not None
        experts = experts_cls(
            moe_config=moe_config,
            quant_config=moe_quant_config,
            max_num_tokens=max_num_tokens,
            num_dispatchers=prepare_finalize.num_dispatchers(),
        )
    else:
        experts = experts_cls(
            moe_config=moe_config,
            quant_config=moe_quant_config,
        )

    kernel = mk.FusedMoEKernel(
        prepare_finalize,
        experts,
    )

    return kernel

make_mxfp4_moe_quant_config(mxfp4_backend, w1_scale, w2_scale, gemm1_alpha=None, gemm1_beta=None, swiglu_limit=None, w1_bias=None, w2_bias=None, a1_scale=None, a2_scale=None, layer=None)

Create a FusedMoEQuantConfig for the given MXFP4 backend.

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

def make_mxfp4_moe_quant_config(
    mxfp4_backend: Mxfp4MoeBackend,
    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
    gemm1_alpha: float | None = None,
    gemm1_beta: float | None = None,
    swiglu_limit: float | None = None,
    w1_bias: torch.Tensor | None = None,
    w2_bias: torch.Tensor | None = None,
    a1_scale: torch.Tensor | None = None,
    a2_scale: torch.Tensor | None = None,
    layer: "RoutedExperts | None" = None,
) -> FusedMoEQuantConfig | None:
    """Create a FusedMoEQuantConfig for the given MXFP4 backend."""
    if mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
        # The megakernel quantizes activations itself and owns its scales.
        return FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4")
    if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8:
        return mxfp4_mxfp8_moe_quant_config(
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        from vllm.model_executor.layers.quantization.utils.quant_utils import (
            GroupShape,
        )

        # DeepGEMM FP4 uses FP8 per-token-group activation quantization
        # with block 128, matching the FP8 DeepGEMM path.
        _fp8_dtype = current_platform.fp8_dtype()
        _block_shape = GroupShape(128, 128)
        return FusedMoEQuantConfig(
            _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
            _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
            _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
            _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:
        # TRTLLM kernel expects non-swizzled mxfp8 activation scales.
        return mxfp4_mxfp8_moe_quant_config(
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
            mx_alignment=256,
            is_scale_swizzled=False,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
        # CUTLASS kernel expects swizzled mxfp8 activation scales.
        return mxfp4_mxfp8_moe_quant_config(
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
            is_scale_swizzled=True,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
        # W4A8: MXFP4 weights + static FP8 activations
        return mxfp4_w4a8_moe_quant_config(
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            block_shape=None,
            gemm1_clamp_limit=swiglu_limit,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
        return ocp_mx_moe_quant_config(
            quant_dtype="mxfp4",
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
    elif mxfp4_backend in (
        Mxfp4MoeBackend.B12X_MXFP4_BF16,
        Mxfp4MoeBackend.MARLIN,
        Mxfp4MoeBackend.BATCHED_MARLIN,
        Mxfp4MoeBackend.TRITON,
        Mxfp4MoeBackend.TRITON_UNFUSED,
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.AITER_MXFP4_BF16,
        Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
        Mxfp4MoeBackend.CPU,
    ):
        return mxfp4_w4a16_moe_quant_config(
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
    elif mxfp4_backend == Mxfp4MoeBackend.HUMMING:
        from vllm.model_executor.layers.quantization.utils.humming import (
            get_humming_moe_quant_config,
        )

        assert layer is not None
        return get_humming_moe_quant_config(
            layer,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
    else:
        return ocp_mx_moe_quant_config(
            quant_dtype="mxfp4",
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )

map_mxfp4_backend(runner_backend)

Map a moe_backend string to its candidate Mxfp4MoeBackends.

Vendor families return all activation variants; the caller picks one via activation_key and is_supported_config.

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

def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
    """Map a moe_backend string to its candidate Mxfp4MoeBackends.

    Vendor families return all activation variants; the caller picks one
    via ``activation_key`` and ``is_supported_config``.
    """
    mapping: dict[str, list[Mxfp4MoeBackend]] = {
        "b12x": list(B12X_BACKENDS),
        "deep_gemm": [Mxfp4MoeBackend.DEEPGEMM_MXFP4],
        "flashinfer_trtllm": [
            Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
            Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
        ],
        "flashinfer_trtllm_afp8": [Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8],
        "flashinfer_cutlass": [
            Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
            Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
        ],
        "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8],
        "flashinfer_moe_ep_cutedsl": [Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL],
        "flashinfer_moe_ep_mega_deep_gemm": [
            Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM
        ],
        "triton": [Mxfp4MoeBackend.TRITON],
        "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED],
        "humming": [Mxfp4MoeBackend.HUMMING],
        "marlin": [Mxfp4MoeBackend.MARLIN],
        "aiter": [
            Mxfp4MoeBackend.AITER_MXFP4_BF16,
            Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
            Mxfp4MoeBackend.AITER_MXFP4_FP8,
            Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
        ],
        "aiter_triton_mxfp4_bf16": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16],
        "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
        "aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4],
        "xpu": [Mxfp4MoeBackend.XPU],
        "cpu": [Mxfp4MoeBackend.CPU],
        "emulation": [Mxfp4MoeBackend.EMULATION],
    }
    if backends := mapping.get(runner_backend):
        return backends
    raise ValueError(
        f"moe_backend='{runner_backend}' is not supported for MXFP4 MoE. "
        f"Expected one of {list(mapping.keys())}."
    )

mxfp4_round_up_hidden_size_and_intermediate_size(backend, hidden_size, intermediate_size, activation=None)

Round up hidden_size and intermediate_size based on backend requirements.

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

def mxfp4_round_up_hidden_size_and_intermediate_size(
    backend: Mxfp4MoeBackend,
    hidden_size: int,
    intermediate_size: int,
    activation: MoEActivation | None = None,
) -> tuple[int, int]:
    """Round up hidden_size and intermediate_size based on backend requirements."""
    if backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
        # The megakernel takes the model's dimensions and validates them itself.
        return hidden_size, intermediate_size
    if backend in B12X_BACKENDS:
        # b12x plans for the exact model dimensions. B12xExperts validates the
        # required MXFP4 block alignment before selecting the backend.
        return hidden_size, intermediate_size
    if backend == Mxfp4MoeBackend.EMULATION:
        # Emulation has no kernel tile; it only needs OCP MX block alignment so the
        # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated
        # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720).
        intermediate_size = round_up(intermediate_size, OCP_MX_BLOCK_SIZE)
        hidden_size = round_up(hidden_size, OCP_MX_BLOCK_SIZE)
    elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        # DeepGEMM requires M/N/K alignment
        intermediate_size = round_up(intermediate_size, 128)
        hidden_size = round_up(hidden_size, 128)
    elif backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
        intermediate_size = round_up(intermediate_size, 128)
        if current_platform.is_xpu():
            hidden_size = round_up(hidden_size, 128)
        else:
            hidden_size = round_up(hidden_size, 256)
    elif backend in TRTLLM_BACKENDS:
        intermediate_size = round_up(intermediate_size, 128)
        hidden_size = round_up(hidden_size, 256)
    elif backend in (
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
    ):
        intermediate_size = round_up(intermediate_size, 128)
        hidden_size = round_up(hidden_size, 128)
    elif current_platform.is_rocm():
        from vllm.platforms.rocm import get_cdna_version

        is_situ_or_silu = activation in (
            MoEActivation.SITU,
            MoEActivation.SILU,
        )

        # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size
        # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128
        # rather than the generic ROCm 256 round-up, which would inflate
        # weights and OOM.
        aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16

        # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs.
        # CDNA4's F16xMXFP4 configuration uses block_k=256.
        triton_uses_128 = (
            backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4
        )

        alignment = (
            128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256
        )

        intermediate_size = round_up(intermediate_size, alignment)
        hidden_size = round_up(hidden_size, alignment)
    elif backend == Mxfp4MoeBackend.CPU:
        # CPU AMX kernel uses BLOCK_N=32, align to 32
        intermediate_size = round_up(intermediate_size, 32)
        hidden_size = round_up(hidden_size, 32)
    else:
        intermediate_size = round_up(intermediate_size, 64)
    return hidden_size, intermediate_size

select_deepseek_v4_mxfp4_moe_backend(config)

Select the MXFP4 MoE backend with MXFP8 activation as top priority. Falls back through BF16 and other backends.

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

def select_deepseek_v4_mxfp4_moe_backend(
    config: FusedMoEConfig,
) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
    """Select the MXFP4 MoE backend with MXFP8 activation as top priority.
    Falls back through BF16 and other backends.
    """
    activation_format = (
        mk.FusedMoEActivationFormat.BatchedExperts
        if config.moe_parallel_config.use_batched_activation_format
        else mk.FusedMoEActivationFormat.Standard
    )

    # Honor explicit moe_backend (e.g. "marlin", "triton_unfused") before
    # falling back to the auto priority list.
    runner_backend = config.moe_backend
    if runner_backend != "auto":
        if runner_backend == "b12x":
            requested_backends = _get_requested_backends(runner_backend, None)
        else:
            # Try every variant of the alias in priority order. Narrowing to
            # the BF16 variant would drop SM100+ W4A8 variants on devices where
            # the BF16 variant is gated to SM90.
            requested_backends = map_mxfp4_backend(runner_backend)
        if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
            requested_backends = [
                Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
                for b in requested_backends
            ]
        last_error: Exception | None = None
        for requested_backend in requested_backends:
            try:
                return _return_or_raise(
                    requested_backend,
                    config,
                    kMxfp4Static,
                    _backend_activation_key(requested_backend),
                    activation_format,
                )
            except ValueError as e:
                last_error = e
        assert last_error is not None
        raise last_error

    # DeepSeek-V4 on ROCm: prefer AITER FlyDSL MoE (better perf + accuracy
    # after shuffle/TP-offset fixes), with Triton-unfused as fallback.
    if (
        current_platform.is_rocm()
        and config.routing_method == RoutingMethodType.DeepseekV4
    ):
        priority_backends = [
            Mxfp4MoeBackend.AITER_MXFP4_BF16,
            Mxfp4MoeBackend.TRITON_UNFUSED,
        ]
    else:
        priority_backends = _get_priority_backends()

    # Iterate priority backends: TRTLLM MXFP8, then Triton.
    for backend in priority_backends:
        activation_key = _backend_activation_key(backend)
        for k_cls in backend_to_kernel_cls(backend):
            supported, reason = k_cls.is_supported_config(
                k_cls, config, kMxfp4Static, activation_key, activation_format
            )
            if supported:
                logger.info_once(_make_log_backend(backend), scope="local")
                return backend, k_cls
            else:
                logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

    raise NotImplementedError(
        "No MXFP4 MoE backend supports the deployment configuration."
    )

select_mxfp4_moe_backend(config, activation_key=None)

Select the primary MXFP4 MoE backend.

Parameters:

• config

  (FusedMoEConfig) –

  MoE configuration

• activation_key

  (QuantKey | None, default: None ) –

  Optional activation quantization key. If provided, overrides the default activation key for backend selection. Use kFp8StaticTensorSym for W4A8 scheme.

Note: Shape-specific fallbacks may still occur at runtime.

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

def select_mxfp4_moe_backend(
    config: FusedMoEConfig,
    activation_key: QuantKey | None = None,
) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
    """Select the primary MXFP4 MoE backend.

    Args:
        config: MoE configuration
        activation_key: Optional activation quantization key. If provided,
            overrides the default activation key for backend selection.
            Use kFp8StaticTensorSym for W4A8 scheme.

    Note: Shape-specific fallbacks may still occur at runtime.

    """
    runner_backend = config.moe_backend
    requested_activation_key = _resolve_activation_key(activation_key)

    activation_format = (
        mk.FusedMoEActivationFormat.BatchedExperts
        if config.moe_parallel_config.use_batched_activation_format
        else mk.FusedMoEActivationFormat.Standard
    )

    if runner_backend != "auto":
        requested_backends = _get_requested_backends(
            runner_backend, requested_activation_key
        )
        if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
            requested_backends = [
                Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
                for b in requested_backends
            ]
        if not requested_backends:
            raise ValueError(
                f"moe_backend={runner_backend!r} does not support "
                f"activation={requested_activation_key}"
            )
        last_error: Exception | None = None
        for requested_backend in requested_backends:
            act_key = (
                requested_activation_key
                if requested_backend == Mxfp4MoeBackend.EMULATION
                else _backend_activation_key(requested_backend)
            )
            try:
                return _return_or_raise(
                    requested_backend,
                    config,
                    kMxfp4Static,
                    act_key,
                    activation_format,
                )
            except ValueError as e:
                last_error = e
        assert last_error is not None
        raise last_error

    if _requires_qwen38_tep8_emulation(config, requested_activation_key):
        backend = Mxfp4MoeBackend.EMULATION
        logger.warning_once(
            "Using OCP MX emulation for the Qwen3.8 Flash Next TEP8 routed "
            "experts on gfx950 because the native AITER W4A4 kernel is not "
            "numerically reliable for this shape. Performance will be lower."
        )
        return _return_or_raise(
            backend,
            config,
            kMxfp4Static,
            requested_activation_key,
            activation_format,
        )

    # Select kernels in order of backend.
    AVAILABLE_BACKENDS = _filter_by_activation(
        _get_priority_backends_for_gpt_oss(), requested_activation_key
    )

    unsupported_reasons = []
    for backend in AVAILABLE_BACKENDS:
        # Use requested_activation_key if provided, otherwise use backend default
        act_key = (
            requested_activation_key
            if requested_activation_key is not None
            else _backend_activation_key(backend)
        )
        for k_cls in backend_to_kernel_cls(backend):
            supported, reason = k_cls.is_supported_config(
                k_cls, config, kMxfp4Static, act_key, activation_format
            )
            if supported:
                logger.info_once(_make_log_backend(backend))
                return backend, k_cls
            else:
                logger.debug_once(_make_log_unsupported(backend, reason))
                unsupported_reasons.append((backend, reason))

    if current_platform.is_xpu():
        backend = Mxfp4MoeBackend.XPU
        logger.info_once(_make_log_backend(backend))
        return _return_or_raise(
            Mxfp4MoeBackend.XPU,
            config,
            kMxfp4Static,
            None,
            activation_format,
        )

    unsupported_log = "; ".join(
        [
            f"backend: {backend.value}, reason: {reason}"
            for backend, reason in unsupported_reasons
        ]
    )
    raise NotImplementedError(
        "No MXFP4 MoE backend supports the deployment configuration. "
        f"weight_key=kMxfp4Static, activation_key={activation_key}. "
        f"Candidate backends were: "
        f"{[backend.value for backend in AVAILABLE_BACKENDS]}. "
        f"Unsupported reasons: {unsupported_log}. "
    )