vllm.model_executor.layers.fused_moe.oracle.mxfp4 ¶

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

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

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]
    _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

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_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,
    ]
    return _AVAILABLE_BACKENDS

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

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

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,
    _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:
        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
            _upcast_e8m0_to_fp32,
        )

        return (
            w13_weight.data,
            w2_weight.data,
            _upcast_e8m0_to_fp32(w13_weight_scale.data),
            _upcast_e8m0_to_fp32(w2_weight_scale.data),
            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_utils import (
            prepare_humming_moe_layer,
        )

        prepare_humming_moe_layer(layer, {"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)
            )

        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.matmul_ogs import FlexCtx, PrecisionConfig
        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(
            weight_scale=w13_scale,
            flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
        )
        w2_precision_config = PrecisionConfig(
            weight_scale=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:
        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig

        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(
            weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
        )
        w2_precision_config = PrecisionConfig(
            weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
        )

        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 == 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:
        # No additional transformation needed for emulation backend,
        # weights are dequantized on the fly in the experts class.
        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)}."
        )

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,
) -> 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, TRTLLM MXFP8, Triton and Marlin backends.
    """

    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
            _upcast_e8m0_to_fp32,
        )

        # Weights stay as uint8 packed FP4 — no layout change needed.
        # Convert E8M0 uint8 scales to float32.
        return (
            w13_weight.data,
            w2_weight.data,
            _upcast_e8m0_to_fp32(w13_weight_scale.data),
            _upcast_e8m0_to_fp32(w2_weight_scale.data),
            w13_bias,
            w2_bias,
        )

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

        prepare_humming_moe_layer(layer, {"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:
        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

        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)
        )

        w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
        w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)

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

        if w13_bias is not None:
            w13_bias = (
                w13_bias.data.view(-1, n // 2, 2)
                .permute(0, 2, 1)
                .contiguous()
                .view(-1, n)
            )

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

    elif mxfp4_backend in TRITON_BACKENDS:
        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig

        if 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(
            weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
        )
        w2_precision_config = PrecisionConfig(
            weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
        )

        del layer.w13_weight
        del layer.w2_weight

        return (
            w13_weight,
            w2_weight,
            w13_precision_config,
            w2_precision_config,
            w13_bias,
            w2_bias,
        )
    else:
        raise ValueError(
            f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
            f"Expected TRTLLM, Triton, or AITER backend."
        )