sgl-project--sglang/python/sglang/srt/layers/quantization/mxfp4.py

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/mxfp4.py

from __future__ import annotations

import os
from dataclasses import replace
from typing import TYPE_CHECKING, List, Optional

import torch
from torch.nn.parameter import Parameter

# Silence the TRT-LLM cutlass autotune trace embedded inside FlashInfer's
# cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
# setdefault preserves any explicit user override.
os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
    use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.amx_utils import (
    CPUQuantMethod,
    _amx_process_weight_after_loading,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend
from sglang.srt.layers.quantization.base_config import (
    FusedMoEMethodBase,
    QuantizationConfig,
    QuantizeMethodBase,
)
from sglang.srt.layers.quantization.utils import is_layer_skipped
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import (
    cpu_has_amx_support,
    is_cpu,
    is_flashinfer_available,
    is_gfx95_supported,
    is_hip,
    is_sm90_supported,
    is_sm100_supported,
    is_sm120_supported,
    is_triton_kernels_available,
    mxfp_supported,
    next_power_of_2,
    round_up,
    set_weight_attrs,
    use_intel_amx_backend,
)
from sglang.srt.utils.common import get_bool_env_var
from sglang.srt.utils.custom_op import register_custom_op

has_triton_kernels = is_triton_kernels_available()


if is_flashinfer_available():
    from flashinfer import (
        mxfp8_quantize,
        nvfp4_block_scale_interleave,
        trtllm_fp4_block_scale_moe,
    )
    from flashinfer.fused_moe.core import (
        get_w2_permute_indices_with_cache,
    )

    # SM90 mixed-input helpers landed in FlashInfer #3084 (post-0.6.10). Older
    # versions don't ship them; gate at import so unrelated code paths still load.
    try:
        from flashinfer.fused_moe import (
            interleave_moe_scales_for_sm90_mixed_gemm,
            interleave_moe_weights_for_sm90_mixed_gemm,
        )

        _FI_HAS_SM90_CUTLASS_MXFP4 = True
    except ImportError:
        interleave_moe_scales_for_sm90_mixed_gemm = None
        interleave_moe_weights_for_sm90_mixed_gemm = None
        _FI_HAS_SM90_CUTLASS_MXFP4 = False
else:
    _FI_HAS_SM90_CUTLASS_MXFP4 = False

_flashinfer_mxfp4_permute_indices_cache: dict[torch.Size, torch.Tensor] = {}
_flashinfer_mxfp4_permute_indices_device_cache: dict[
    tuple[tuple[int, ...], int, int, str, int], torch.Tensor
] = {}


def _get_flashinfer_mxfp4_device_permute_indices(
    x: torch.Tensor,
    epilogue_tile_m: int,
    num_elts_per_sf: Optional[int] = None,
) -> torch.Tensor:
    extra_args = {} if num_elts_per_sf is None else {"num_elts_per_sf": num_elts_per_sf}
    permute_indices = get_w2_permute_indices_with_cache(
        _flashinfer_mxfp4_permute_indices_cache,
        x,
        epilogue_tile_m,
        **extra_args,
    )

    device_index = -1 if x.device.index is None else x.device.index
    num_elts_per_sf_key = -1 if num_elts_per_sf is None else num_elts_per_sf
    cache_key = (
        tuple(x.shape),
        epilogue_tile_m,
        num_elts_per_sf_key,
        x.device.type,
        device_index,
    )
    cached_device_indices = _flashinfer_mxfp4_permute_indices_device_cache.get(
        cache_key
    )
    if cached_device_indices is None:
        cached_device_indices = permute_indices.to(x.device)
        _flashinfer_mxfp4_permute_indices_device_cache[cache_key] = (
            cached_device_indices
        )

    return cached_device_indices


if TYPE_CHECKING:
    from sglang.srt.layers.moe.token_dispatcher import (
        CombineInput,
        StandardDispatchOutput,
    )

_is_cpu = is_cpu()
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
_is_cpu_amx_available = cpu_has_amx_support()

if _is_hip:
    # import aiter
    try:
        from aiter.ops.shuffle import (
            shuffle_scale,
            shuffle_scale_a16w4,
            shuffle_weight,
            shuffle_weight_a16w4,
        )
        from aiter.ops.triton.quant import dynamic_mxfp4_quant
        from aiter.utility.fp4_utils import e8m0_shuffle
    except ImportError as err:
        dynamic_mxfp4_quant = e8m0_shuffle = err


def _swizzle_mxfp4(quant_tensor, scale, num_warps):
    """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
    from triton_kernels.numerics import InFlexData
    from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
    from triton_kernels.tensor_details import layout

    value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
        mx_axis=1
    )
    scale_layout, scale_layout_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
        mx_axis=1, num_warps=num_warps
    )
    if is_sm100_supported():
        constraints = {
            "is_persistent": True,
            "epilogue_subtile": 1,
        }
        opt_flags.update_opt_flags_constraints(constraints)
    elif is_sm90_supported():
        constraints = {
            "split_k": 1,
        }
        opt_flags.update_opt_flags_constraints(constraints)
    # transpose the tensor so that the quantization axis is on dim1
    quant_tensor = quant_tensor.transpose(-2, -1)
    scale = scale.transpose(-2, -1)
    quant_tensor = convert_layout(
        wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts
    )
    scale = convert_layout(wrap_torch_tensor(scale), scale_layout, **scale_layout_opts)
    return quant_tensor, InFlexData(), scale


def _dequant_mxfp4_fake(
    x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
) -> torch.Tensor:
    return torch.empty(
        (*x.shape[:-1], x.shape[-1] * 2), dtype=float_dtype, device=x.device
    )


@register_custom_op(fake_impl=_dequant_mxfp4_fake)
def dequant_mxfp4(
    x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
) -> torch.Tensor:
    try:
        from quark.torch.kernel import mx
    except ImportError as err:
        raise ImportError(
            "The package `amd-quark` is required to use "
            "MX-FP4 models. Please install it with `pip install "
            "amd-quark`."
        ) from err

    return mx.dq_mxfp4(x, scale, float_dtype)


@register_custom_op(out_shape="x")
def quant_dequant_mxfp4(
    x: torch.Tensor, scale_calculation_mode: str = "even"
) -> torch.Tensor:
    try:
        from quark.torch.kernel import mx
    except ImportError as err:
        raise ImportError(
            "The package `amd-quark` is required to use "
            "MX-FP4 models. Please install it with `pip install "
            "amd-quark`."
        ) from err

    return mx.qdq_mxfp4(x, scale_calculation_mode)


class Mxfp4Config(QuantizationConfig):

    def __init__(
        self,
        ignored_layers: Optional[list[str]] = None,
        is_checkpoint_mxfp4_serialized: bool = False,
    ):
        super().__init__()
        self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
        self.ignored_layers = ignored_layers

    @classmethod
    def from_config(cls, config):

        quant_method = cls.get_from_keys(config, ["quant_method"])
        is_checkpoint_mxfp4_serialized = "mxfp4" in quant_method

        if _is_hip:
            if mxfp_supported():
                return cls(
                    is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized
                )
            else:

                platform = torch.cuda.get_device_properties(0).gcnArchName
                raise ValueError(
                    f"Current platform {platform} not support mxfp4 computation"
                )

        return cls(is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized)

    @classmethod
    def get_min_capability(cls) -> int:
        return 80

    @classmethod
    def get_name(cls) -> str:
        return "mxfp4"

    @classmethod
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_config_filenames(cls) -> list[str]:
        return []

    def is_static_cfg(self):
        return self.is_checkpoint_mxfp4_serialized

    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
    ) -> Optional[QuantizeMethodBase]:

        from sglang.srt.layers.linear import LinearBase
        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
        from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod

        if isinstance(layer, LinearBase):
            if self.ignored_layers and is_layer_skipped(
                prefix=prefix,
                ignored_layers=self.ignored_layers,
                fused_mapping=self.packed_modules_mapping,
            ):
                return UnquantizedLinearMethod()
            elif _is_hip:
                return UnquantizedLinearMethod()
        elif isinstance(layer, FusedMoE):
            if self.is_checkpoint_mxfp4_serialized:
                return Mxfp4MoEMethod(prefix=prefix)
            else:
                return Mxfp4DynamicQuantMoEMethod()
        else:
            if self.is_checkpoint_mxfp4_serialized:
                raise NotImplementedError("Mxfp4 attention layer is not implemented")
        return None

    def get_scaled_act_names(self) -> List[str]:
        return []


class Mxfp4MoEMethod(FusedMoEMethodBase):

    def __init__(
        self,
        prefix: str,
    ):
        super().__init__()

        self.prefix = prefix
        self.topk_indices_dtype = None
        self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
        self.with_bias = False
        self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4()
        self.use_marlin = get_moe_runner_backend().is_marlin()
        self.flashinfer_mxfp4_moe_precision = (
            get_server_args().flashinfer_mxfp4_moe_precision
        )
        # When `flashinfer_mxfp4` is enabled, dispatch to one of two FlashInfer
        # entry points depending on the GPU:
        #   - SM100 (Blackwell)  -> trtllm_fp4_block_scale_moe (existing)
        #   - SM90  (Hopper)     -> cutlass_fused_moe(use_w4_group_scaling=True)
        #                           (FlashInfer PR #3084, post-0.6.10)
        self._fi_kernel: Optional[str] = None
        if self.use_flashinfer:
            if is_sm100_supported():
                self._fi_kernel = "trtllm_sm100"
            elif is_sm90_supported():
                if not _FI_HAS_SM90_CUTLASS_MXFP4:
                    raise RuntimeError(
                        "moe_runner_backend=flashinfer_mxfp4 on SM90 requires the "
                        "interleave_moe_{weights,scales}_for_sm90_mixed_gemm helpers "
                        "from FlashInfer PR #3084 (>= 0.6.11). Upgrade flashinfer-python "
                        "or pick a different backend (e.g. marlin / triton_kernel)."
                    )
                self._fi_kernel = "cutlass_sm90"
            else:
                raise NotImplementedError(
                    "moe_runner_backend=flashinfer_mxfp4 requires SM90 or SM100."
                )

    def create_weights(
        self,
        layer: torch.nn.Module,
        num_experts: int,
        hidden_size: int,
        intermediate_size_per_partition: int,
        params_dtype: torch.dtype,
        with_bias: bool = False,
        **extra_weight_attrs,
    ):
        self.num_experts = num_experts
        weight_dtype = torch.uint8
        scale_dtype = torch.uint8
        self.with_bias = with_bias
        mxfp4_block = 32
        triton_kernels_padding_alignment = 64

        # pad the intermediate size to be a multiple of 2 * mxfp4_block
        # for to hold non-uniform sharded tensor as well as swizzling
        intermediate_size_per_partition_after_pad = intermediate_size_per_partition
        if self.use_marlin:
            intermediate_size_per_partition_after_pad = round_up(
                intermediate_size_per_partition, 128
            )
            hidden_size = round_up(hidden_size, 256)
            self.hidden_pad = hidden_size - layer.hidden_size
            self.intermediate_pad = (
                intermediate_size_per_partition_after_pad
                - layer.intermediate_size_per_partition
            )
        elif is_sm100_supported():
            if self.use_flashinfer:
                intermediate_size_per_partition_after_pad = round_up(
                    intermediate_size_per_partition, 256
                )
                hidden_size = round_up(hidden_size, 256)
            else:
                intermediate_size_per_partition_after_pad = round_up(
                    intermediate_size_per_partition, triton_kernels_padding_alignment
                )
        elif self._fi_kernel == "cutlass_sm90":
            # cutlass mixed-input GEMM contraction dim K must be % 128 == 0
            # (interleave factor for MXFP4 group_size=32 is 4). The kernel
            # also expects ``fc1_expert_weights`` in halved ``[up; gate]``
            # layout, which means the padding boundary must fall on the
            # gate / up split.
            #
            # The mxfp4 weight loader (FusedMoE.weight_loader fast path) does
            # a NAIVE copy of HF's ``[2*intermediate_size, hidden_packed]``
            # tensor into the buffer's ``[:dim1, :dim2]`` slice. Padding the
            # buffer here would push the gate/up boundary, so HF's "up"
            # rows would land in the buffer's "gate" half and vice versa.
            # Marlin sidesteps this by not padding; we do the same and
            # rebuild a properly-padded buffer in
            # ``_process_weights_for_sm90_cutlass`` after the load completes.
            self._padded_intermediate = round_up(intermediate_size_per_partition, 128)
            self._padded_hidden = round_up(hidden_size, 128)
            # create_weights below uses the *unpadded* sizes so the loader's
            # naive-copy fast path is correct.
            intermediate_size_per_partition_after_pad = intermediate_size_per_partition
        elif _use_aiter:

            intermediate_size_per_partition_after_pad = round_up(
                intermediate_size_per_partition, 256
            )

            hidden_size = round_up(hidden_size, 256)
            self.hidden_pad = hidden_size - layer.hidden_size
            self.intermediate_pad = (
                intermediate_size_per_partition_after_pad
                - layer.intermediate_size_per_partition
            )
        elif has_triton_kernels:
            intermediate_size_per_partition_after_pad = round_up(
                intermediate_size_per_partition, triton_kernels_padding_alignment
            )

        self.intermediate_size_per_partition = intermediate_size_per_partition_after_pad

        self.hidden_size = hidden_size
        # Fused gate_up_proj (column parallel)
        w13_weight = torch.nn.Parameter(
            torch.zeros(
                layer.num_local_experts,
                2 * intermediate_size_per_partition_after_pad,
                hidden_size // 2,
                dtype=weight_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w13_weight", w13_weight)
        set_weight_attrs(w13_weight, extra_weight_attrs)

        w13_weight_scale = torch.nn.Parameter(
            torch.zeros(
                layer.num_local_experts,
                2 * intermediate_size_per_partition_after_pad,
                hidden_size // mxfp4_block,
                dtype=scale_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w13_weight_scale", w13_weight_scale)
        set_weight_attrs(w13_weight_scale, extra_weight_attrs)

        w13_weight_bias = torch.nn.Parameter(
            torch.zeros(
                layer.num_local_experts,
                2 * intermediate_size_per_partition_after_pad,
                dtype=torch.bfloat16,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w13_weight_bias", w13_weight_bias)
        set_weight_attrs(w13_weight_bias, extra_weight_attrs)

        # down_proj (row parallel)
        w2_weight = torch.nn.Parameter(
            torch.zeros(
                layer.num_local_experts,
                hidden_size,
                intermediate_size_per_partition_after_pad // 2,
                dtype=weight_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_weight", w2_weight)
        set_weight_attrs(w2_weight, extra_weight_attrs)

        w2_weight_scale = torch.nn.Parameter(
            torch.zeros(
                layer.num_local_experts,
                hidden_size,
                intermediate_size_per_partition_after_pad // mxfp4_block,
                dtype=scale_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_weight_scale", w2_weight_scale)
        set_weight_attrs(w2_weight_scale, extra_weight_attrs)

        w2_weight_bias = torch.nn.Parameter(
            torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16),
            requires_grad=False,
        )
        layer.register_parameter("w2_weight_bias", w2_weight_bias)
        set_weight_attrs(w2_weight_bias, extra_weight_attrs)

    def process_weights_after_loading(self, layer):
        if self.use_marlin:
            from sglang.srt.layers.quantization.marlin_utils import (
                check_moe_marlin_supports_layer,
            )
            from sglang.srt.layers.quantization.marlin_utils_fp4 import (
                deinterleave_moe_mxfp4_w13_for_marlin,
                prepare_moe_mxfp4_layer_for_marlin,
            )

            if not is_sm90_supported() and not is_sm120_supported():
                raise RuntimeError("MXFP4 Marlin requires SM90 or SM120.")
            if not check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True):
                raise RuntimeError(
                    "Current MXFP4 MoE layer is not supported by Marlin."
                )

            if self.moe_runner_config.gemm1_alpha is not None:
                deinterleave_moe_mxfp4_w13_for_marlin(layer)
            prepare_moe_mxfp4_layer_for_marlin(layer)
            layer._mxfp4_backend = "marlin"
            return

        if self._fi_kernel == "cutlass_sm90":
            self._process_weights_for_sm90_cutlass(layer)
            return
        if self.use_flashinfer:
            # TODO: these values are hardcoded for now, we need to get them from the model
            layer.gemm1_alpha = Parameter(
                torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(),
                requires_grad=False,
            )
            layer.gemm1_beta = Parameter(
                torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(),
                requires_grad=False,
            )
            layer.gemm1_clamp_limit = Parameter(
                torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(),
                requires_grad=False,
            )
            sf_block_size = 32  # mxfp4 block size

            assert (
                layer.w13_weight.dim() == 3
                and layer.w13_weight.shape[0] == self.num_experts
                and layer.w13_weight.shape[1]
                == self.intermediate_size_per_partition * 2
                and layer.w13_weight.shape[2] == self.hidden_size // 2
            )
            assert (
                layer.w13_weight_scale.dim() == 3
                and layer.w13_weight_scale.shape[0] == self.num_experts
                and layer.w13_weight_scale.shape[1]
                == self.intermediate_size_per_partition * 2
                and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size
            )
            assert (
                layer.w2_weight.dim() == 3
                and layer.w2_weight.shape[0] == self.num_experts
                and layer.w2_weight.shape[1] == self.hidden_size
                and layer.w2_weight.shape[2]
                == self.intermediate_size_per_partition // 2
            )
            assert (
                layer.w2_weight_scale.dim() == 3
                and layer.w2_weight_scale.shape[1] == self.hidden_size
                and layer.w2_weight_scale.shape[2]
                == self.intermediate_size_per_partition // sf_block_size
            )
            assert (
                layer.w13_weight_bias.dim() == 2
                and layer.w13_weight_bias.shape[0] == self.num_experts
                and layer.w13_weight_bias.shape[1]
                == self.intermediate_size_per_partition * 2
            )
            assert (
                layer.w2_weight_bias.dim() == 2
                and layer.w2_weight_bias.shape[0] == self.num_experts
                and layer.w2_weight_bias.shape[1] == self.hidden_size
            )

            w13_weight_scale = layer.w13_weight_scale.data
            w2_weight_scale = layer.w2_weight_scale.data
            w13_weight = layer.w13_weight.data
            w2_weight = layer.w2_weight.data
            w13_bias = layer.w13_weight_bias.data.to(torch.float32)
            w2_bias = layer.w2_weight_bias.data.to(torch.float32)

            # Swap w1 and w3 as the definition of
            # swiglu is different in the trtllm-gen
            def swap_every_two_rows(x, axis=-1):
                shape = x.shape
                if axis < 0:
                    axis = len(shape) + axis

                # Create a new shape with pairs swapped along specified axis
                new_shape = list(shape)
                new_shape[axis] = shape[axis] // 2
                new_shape.insert(axis + 1, 2)

                # Reshape to expose pairs, swap them, and reshape back
                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_mxfp4_shuffled = []
            gemm1_scales_mxfp4_shuffled = []
            gemm2_weights_mxfp4_shuffled = []
            gemm2_scales_mxfp4_shuffled = []
            gemm1_bias_shuffled = []
            gemm2_bias_shuffled = []
            epilogue_tile_m = 128  # FIXME: this depends on the kernel internals
            w13_weight_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w13_weight[0].view(torch.uint8),
                epilogue_tile_m,
            )
            w13_scale_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w13_weight_scale[0].view(torch.uint8),
                epilogue_tile_m,
                num_elts_per_sf=16,
            )
            w13_bias_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w13_bias[0].reshape(-1, 1),
                epilogue_tile_m,
            )

            w2_weight_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w2_weight[0].view(torch.uint8),
                epilogue_tile_m,
            )
            w2_scale_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w2_weight_scale[0].view(torch.uint8),
                epilogue_tile_m,
                num_elts_per_sf=16,
            )
            w2_bias_permute_indices = _get_flashinfer_mxfp4_device_permute_indices(
                w2_bias[0].reshape(-1, 1),
                epilogue_tile_m,
            )

            for i in range(self.num_experts):
                gemm1_weights_mxfp4_shuffled.append(
                    w13_weight[i]
                    .view(torch.uint8)[w13_weight_permute_indices]
                    .contiguous()
                )

                gemm1_scales_mxfp4_shuffled.append(
                    nvfp4_block_scale_interleave(
                        w13_weight_scale[i]
                        .view(torch.uint8)[w13_scale_permute_indices]
                        .contiguous()
                    )
                )

                gemm1_bias_shuffled.append(
                    w13_bias[i].reshape(-1, 1)[w13_bias_permute_indices].contiguous()
                )

                gemm2_weights_mxfp4_shuffled.append(
                    w2_weight[i]
                    .view(torch.uint8)[w2_weight_permute_indices]
                    .contiguous()
                )

                gemm2_scales_mxfp4_shuffled.append(
                    nvfp4_block_scale_interleave(
                        w2_weight_scale[i]
                        .view(torch.uint8)[w2_scale_permute_indices]
                        .contiguous()
                    )
                )

                gemm2_bias_shuffled.append(
                    w2_bias[i].reshape(-1, 1)[w2_bias_permute_indices].contiguous()
                )

            w13_weight = torch.stack(gemm1_weights_mxfp4_shuffled)
            w13_weight_scale = (
                torch.stack(gemm1_scales_mxfp4_shuffled)
                .reshape(
                    self.num_experts,
                    2 * self.intermediate_size_per_partition,
                    self.hidden_size // sf_block_size,
                )
                .view(torch.float8_e4m3fn)
            )

            w2_weight = torch.stack(gemm2_weights_mxfp4_shuffled)
            w2_weight_scale = (
                torch.stack(gemm2_scales_mxfp4_shuffled)
                .reshape(
                    self.num_experts,
                    self.hidden_size,
                    self.intermediate_size_per_partition // sf_block_size,
                )
                .view(torch.float8_e4m3fn)
            )

            layer.w13_weight = Parameter(w13_weight, requires_grad=False)
            layer.w13_weight_scale = Parameter(w13_weight_scale, requires_grad=False)
            layer.w2_weight = Parameter(w2_weight, requires_grad=False)
            layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False)
            layer.w13_weight_bias = Parameter(
                torch.stack(gemm1_bias_shuffled).reshape(self.num_experts, -1),
                requires_grad=False,
            )
            layer.w2_weight_bias = Parameter(
                torch.stack(gemm2_bias_shuffled).reshape(self.num_experts, -1),
                requires_grad=False,
            )
            return
        if _use_aiter:
            if layer.w13_weight_bias is not None:
                layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(
                    torch.float32
                )
            if layer.w2_weight_bias is not None:
                layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)

            e, n, k = layer.w13_weight.shape
            layer.w13_weight.view(torch.uint8).copy_(
                layer.w13_weight.data.view(torch.uint8)
                .view(e, n // 2, 2, k)
                .permute(0, 2, 1, 3)
                .contiguous()
                .view(e, n, k)
            )
            layer.w13_weight_scale.data = (
                layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
                .permute(0, 2, 1, 3)
                .contiguous()
                .view(e, n, -1)
            )
            layer.w13_weight_bias.data = (
                layer.w13_weight_bias.data.view(-1, n // 2, 2)
                .permute(0, 2, 1)
                .contiguous()
                .view(-1, n)
            )

            if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get():
                layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True)
                shuffled_w13_scale = shuffle_scale_a16w4(
                    layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
                    self.num_experts,
                    True,
                )

                layer.w2_weight.data = shuffle_weight_a16w4(layer.w2_weight, 16, False)
                shuffled_w2_scale = shuffle_scale_a16w4(
                    layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
                    self.num_experts,
                    False,
                )
            else:
                layer.w13_weight.data = shuffle_weight(
                    layer.w13_weight, is_guinterleave=False, gate_up=True
                )
                shuffled_w13_scale = shuffle_scale(
                    layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
                    experts_cnt=self.num_experts,
                    is_guinterleave=False,
                    gate_up=True,
                )
                layer.w2_weight.data = shuffle_weight(
                    layer.w2_weight, is_guinterleave=False, gate_up=False
                )
                shuffled_w2_scale = shuffle_scale(
                    layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
                    experts_cnt=self.num_experts,
                    is_guinterleave=False,
                    gate_up=False,
                )

            # shuffle_weight_a16w4(gate_up=True) above preshuffles w13 into aiter's
            # preshuffle + gate/up-interleaved layout. Tag the Parameter so apply()
            # can carry the metadata across .view(float4_e2m1fn_x2) and aiter's
            # fused_moe selects the preshuffle_on kernel family.
            layer.w13_weight.is_shuffled = True
            layer.w2_weight.is_shuffled = True

            layer.w13_weight_scale = torch.nn.Parameter(
                shuffled_w13_scale, requires_grad=False
            )
            layer.w2_weight_scale = torch.nn.Parameter(
                shuffled_w2_scale, requires_grad=False
            )

            return

        if self.use_triton_kernels:

            from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig

            w13_weight_bias = layer.w13_weight_bias.to(torch.float32)
            w2_weight_bias = layer.w2_weight_bias.to(torch.float32)

            layer.w13_weight_bias = Parameter(w13_weight_bias, requires_grad=False)
            layer.w2_weight_bias = Parameter(w2_weight_bias, requires_grad=False)

            num_warps = 8

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

            self.w13_precision_config = PrecisionConfig(
                weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
            )
            self.w2_precision_config = PrecisionConfig(
                weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
            )

            self.w13_weight_triton_tensor = w13_weight
            self.w2_weight_triton_tensor = w2_weight
            del layer.w13_weight
            del layer.w2_weight
        elif _is_cpu and _is_cpu_amx_available:
            _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
            if use_intel_amx_backend(layer):
                packed_w13_weight_scale = torch.ops.sgl_kernel.convert_scale_packed(
                    layer.w13_weight_scale
                )
                packed_w2_weight_scale = torch.ops.sgl_kernel.convert_scale_packed(
                    layer.w2_weight_scale
                )
                layer.w13_weight_scale = Parameter(
                    packed_w13_weight_scale, requires_grad=False
                )
                layer.w2_weight_scale = Parameter(
                    packed_w2_weight_scale, requires_grad=False
                )
                if hasattr(layer, "w13_weight_bias"):
                    layer.w13_weight_bias = Parameter(
                        layer.w13_weight_bias.float(), requires_grad=False
                    )
                if hasattr(layer, "w2_weight_bias"):
                    layer.w2_weight_bias = Parameter(
                        layer.w2_weight_bias.float(), requires_grad=False
                    )
                return
            # Fallback if the TP-sharded layer cannot be AMX-packed
            from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil

            w13_weight = MXFP4QuantizeUtil.dequantize(
                quantized_data=layer.w13_weight,
                dtype=torch.bfloat16,
                scale=layer.w13_weight_scale,
                block_sizes=[32],
            )
            w2_weight = MXFP4QuantizeUtil.dequantize(
                quantized_data=layer.w2_weight,
                dtype=torch.bfloat16,
                scale=layer.w2_weight_scale,
                block_sizes=[32],
            )
            del layer.w13_weight
            del layer.w2_weight
            del layer.w13_weight_scale
            del layer.w2_weight_scale
            layer.w13_weight = Parameter(w13_weight, requires_grad=False)
            layer.w2_weight = Parameter(w2_weight, requires_grad=False)

            return
        else:
            from triton_kernels.numerics_details.mxfp import upcast_from_mxfp

            w13_weight = upcast_from_mxfp(
                layer.w13_weight,
                layer.w13_weight_scale,
                target_dtype=torch.bfloat16,
                axis=-1,
            )
            w2_weight = upcast_from_mxfp(
                layer.w2_weight,
                layer.w2_weight_scale,
                target_dtype=torch.bfloat16,
                axis=-1,
            )
            del layer.w13_weight
            del layer.w2_weight
            del layer.w13_weight_scale
            del layer.w2_weight_scale
            layer.w13_weight = Parameter(w13_weight.data, requires_grad=False)
            layer.w2_weight = Parameter(w2_weight.data, requires_grad=False)
        torch.cuda.empty_cache()

    def _process_weights_for_sm90_cutlass(self, layer):
        """De-interleave + pad + halving-swap + byte-interleave MXFP4 weights
        for FlashInfer's SM90 ``cutlass_fused_moe(use_w4_group_scaling=True)``
        path (PR #3084).

        The cutlass kernel needs (a) K (contraction dim) % 128 == 0, and (b)
        ``fc1_expert_weights`` in halved ``[up; gate]`` order -- the
        ``compute_with_experts`` reference in FlashInfer's
        ``test_trtllm_cutlass_fused_moe.py`` splits
        ``w3, w1 = chunk(W, 2, dim=0)`` and uses w3 as up, w1 as gate.

        GPT-OSS's HF layout is *interleaved* ``[g_0, u_0, g_1, u_1, ..., g_{N-1}, u_{N-1}]``
        (each pair occupies two adjacent rows). The mxfp4 weight loader does
        a naive copy, so our unpadded buffer is interleaved post-load. We
        de-interleave (even rows -> gate, odd rows -> up), pad each half from
        N_un to N_pad, concatenate as halved ``[up; gate]``, and then run
        FlashInfer's byte / scale interleave helpers.
        """
        sf_block_size = 32  # MXFP4 group size

        # Sizes from the unpadded loaded buffers.
        N_un = layer.w13_weight.shape[1] // 2  # intermediate (unpadded)
        K_un = (
            layer.w13_weight.shape[2] * 2
        )  # hidden (unpadded, *2 because packed 4-bit)
        N_pad = self._padded_intermediate
        K_pad = self._padded_hidden
        # Use the local expert count (matches the existing buffer allocation in
        # create_weights) so the SM90 cutlass path remains correct under
        # Expert Parallelism. `self.num_experts` is the *global* count.
        E = layer.num_local_experts
        device = layer.w13_weight.device
        bias_dtype = layer.w13_weight_bias.dtype

        # ---- De-interleave + pad w13 weight/scale/bias to halved [up; gate]
        # Even rows of HF = gate, odd rows = up. After splitting we pad each
        # half along its row dim (N) from N_un to N_pad with zeros, and along
        # its last dim (K) from K_un (or K_un / sf_block_size) to K_pad.

        def _stack_up_gate_w13(unpadded_w13, last_pad, last_un):
            # unpadded_w13: [E, 2*N_un, last_un]
            # Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order.
            gate_rows = unpadded_w13[:, 0::2, :]  # [E, N_un, last_un]
            up_rows = unpadded_w13[:, 1::2, :]  # [E, N_un, last_un]
            out = torch.zeros(
                E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device
            )
            # First half: up (with row + col padding zeros).
            out[:, :N_un, :last_un] = up_rows
            # Second half: gate.
            out[:, N_pad : N_pad + N_un, :last_un] = gate_rows
            return out

        w13_padded = _stack_up_gate_w13(
            layer.w13_weight.data.view(torch.uint8), K_pad // 2, K_un // 2
        )
        w13_scale_padded = _stack_up_gate_w13(
            layer.w13_weight_scale.data,
            K_pad // sf_block_size,
            K_un // sf_block_size,
        )
        # Bias: same de-interleave on dim=-1.
        w13_bias_gate = layer.w13_weight_bias.data[:, 0::2]  # [E, N_un]
        w13_bias_up = layer.w13_weight_bias.data[:, 1::2]  # [E, N_un]
        w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device)
        w13_bias_padded[:, :N_un] = w13_bias_up
        w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate

        def _pad_w2_3d(unpadded, last_pad, last_un):
            out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
            out[:, :K_un, :last_un] = unpadded[:, :K_un, :]
            return out

        # ---- w2 (no halving, just pad to [E, K_pad, N_pad/2]) ----------------
        w2_padded = _pad_w2_3d(
            layer.w2_weight.data.view(torch.uint8), N_pad // 2, N_un // 2
        )
        w2_scale_padded = _pad_w2_3d(
            layer.w2_weight_scale.data,
            N_pad // sf_block_size,
            N_un // sf_block_size,
        )
        w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
        w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data

        # ---- Per-expert SwiGLU scalars (GPT-OSS defaults) ------------------
        layer.swiglu_alpha = Parameter(
            torch.full((E,), 1.702, dtype=torch.float32, device=device),
            requires_grad=False,
        )
        layer.swiglu_beta = Parameter(
            torch.full((E,), 1.0, dtype=torch.float32, device=device),
            requires_grad=False,
        )
        layer.swiglu_limit = Parameter(
            torch.full((E,), 7.0, dtype=torch.float32, device=device),
            requires_grad=False,
        )

        # ---- FlashInfer SM90 byte / scale interleave -----------------------
        # The padded buffers above are contiguous by construction (allocated
        # via torch.zeros + slice assignment), so we feed them straight in.
        layer.w13_weight = Parameter(
            interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4"),
            requires_grad=False,
        )
        layer.w2_weight = Parameter(
            interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4"),
            requires_grad=False,
        )
        layer.w13_weight_scale = Parameter(
            interleave_moe_scales_for_sm90_mixed_gemm(
                w13_scale_padded, group_size=sf_block_size
            ),
            requires_grad=False,
        )
        layer.w2_weight_scale = Parameter(
            interleave_moe_scales_for_sm90_mixed_gemm(
                w2_scale_padded, group_size=sf_block_size
            ),
            requires_grad=False,
        )
        layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
        layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)

        torch.cuda.empty_cache()

    def create_moe_runner(
        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
    ):
        self.moe_runner_config = moe_runner_config
        moe_runner_backend = get_moe_runner_backend()
        if moe_runner_backend.is_auto():
            # Must match apply() priority: _use_aiter before use_triton_kernels.
            if _use_aiter and get_moe_a2a_backend().supports_aiter():
                moe_runner_backend = MoeRunnerBackend.AITER
            elif self.use_triton_kernels:
                moe_runner_backend = MoeRunnerBackend.TRITON_KERNELS
            else:
                moe_runner_backend = MoeRunnerBackend.TRITON

        if moe_runner_backend.is_aiter():
            # MXFP4 hard-codes Swiglu in the AITER kernel path.
            self.runner = MoeRunner(
                moe_runner_backend, replace(moe_runner_config, activation="swiglu")
            )
        elif (
            moe_runner_backend.is_triton_kernels()
            or moe_runner_backend.is_triton()
            or moe_runner_backend.is_marlin()
        ):
            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
        elif (
            moe_runner_backend.is_flashinfer_mxfp4()
            and self._fi_kernel == "cutlass_sm90"
        ):
            # Register the fused func at runner construction so the FusedOpPool
            # lookup at `MoeRunner.__init__` finds it.
            import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass  # noqa: F401

            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
        else:
            # Legacy bypass path (e.g. SM100 trtllm-gen under flashinfer_mxfp4)
            # routes through `apply` without a MoeRunner. TODO(cwan): migrate.
            pass

    def _apply_sm90_cutlass(self, layer, dispatch_output):
        """SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input
        path (PR #3084). Routed through the unified ``MoeRunner`` -- this
        helper only builds the quant_info; the actual kernel call lives in
        :mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`."""
        from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
            FlashInferCutlassMxfp4MoeQuantInfo,
        )

        quant_info = FlashInferCutlassMxfp4MoeQuantInfo(
            w13_weight=layer.w13_weight,
            w2_weight=layer.w2_weight,
            w13_weight_scale=layer.w13_weight_scale,
            w2_weight_scale=layer.w2_weight_scale,
            w13_bias=layer.w13_weight_bias,
            w2_bias=layer.w2_weight_bias,
            swiglu_alpha=layer.swiglu_alpha,
            swiglu_beta=layer.swiglu_beta,
            swiglu_limit=layer.swiglu_limit,
            moe_tp_size=layer.moe_tp_size,
            moe_tp_rank=layer.moe_tp_rank,
            moe_ep_size=layer.moe_ep_size,
            moe_ep_rank=layer.moe_ep_rank,
            padded_hidden=self._padded_hidden,
        )
        return self.runner.run(dispatch_output, quant_info)

    def apply(
        self,
        layer: torch.nn.Module,
        dispatch_output: StandardDispatchOutput,
    ) -> CombineInput:

        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
        from sglang.srt.layers.moe.topk import TopKOutputChecker

        x = dispatch_output.hidden_states
        topk_output = dispatch_output.topk_output
        if _is_cpu:
            if use_intel_amx_backend(layer):
                from sglang.srt.layers.moe.topk import apply_topk_weights_cpu

                topk_weights, topk_ids, _ = dispatch_output.topk_output
                x, topk_weights = apply_topk_weights_cpu(
                    self.moe_runner_config.apply_router_weight_on_input, topk_weights, x
                )
                output = torch.ops.sgl_kernel.fused_experts_cpu(
                    x,
                    layer.w13_weight,
                    layer.w2_weight,
                    topk_weights,
                    topk_ids,
                    False,  # inplace See [Note] inplace should be False in fused_experts.
                    CPUQuantMethod.MXFP4,
                    layer.w13_weight_scale,  # w1_scale
                    layer.w2_weight_scale,  # w2_scale
                    None,  # w1_zp
                    None,  # w2_zp
                    None,  # block_size
                    getattr(layer, "w13_weight_bias", None),
                    getattr(layer, "w2_weight_bias", None),
                    layer.moe_runner_config.gemm1_alpha,
                    layer.moe_runner_config.gemm1_clamp_limit,
                    True,  # is_vnni
                )
            else:
                from sglang.srt.layers.moe.fused_moe_native import moe_forward_native

                output = moe_forward_native(
                    layer,
                    x,
                    topk_output,
                    self.moe_runner_config,
                )
            return StandardCombineInput(hidden_states=output)

        if self.use_marlin:
            assert TopKOutputChecker.format_is_standard(topk_output)
            if x.shape[-1] == self.hidden_size:
                x_padded = x
            else:
                x_padded = torch.nn.functional.pad(
                    x, (0, self.hidden_pad), mode="constant", value=0.0
                )
            quant_info = MarlinMoeQuantInfo(
                w13_qweight=layer.w13_weight,
                w2_qweight=layer.w2_weight,
                w13_scales=layer.w13_weight_scale,
                w2_scales=layer.w2_weight_scale,
                w13_g_idx_sort_indices=None,
                w2_g_idx_sort_indices=None,
                weight_bits=4,
                is_k_full=True,
                w13_bias=getattr(layer, "w13_weight_bias", None),
                w2_bias=getattr(layer, "w2_weight_bias", None),
            )
            return self.runner.run(
                dispatch_output._replace(hidden_states=x_padded), quant_info
            )

        if self._fi_kernel == "cutlass_sm90":
            return self._apply_sm90_cutlass(layer, dispatch_output)
        if self.use_flashinfer:
            # When bf16 mode is enabled, we don't need to quantize the input,
            # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
            # which can theoretically improve performance
            origin_hidden_states_dim = x.shape[-1]
            if self.flashinfer_mxfp4_moe_precision == "bf16":
                assert x.dtype == torch.bfloat16
                x_quant = x
                x_scale = None

                # May be fused later if this code branch is frequently needed
                if self.hidden_size != origin_hidden_states_dim:
                    x_quant = torch.nn.functional.pad(
                        x_quant,
                        (0, self.hidden_size - origin_hidden_states_dim),
                        mode="constant",
                        value=0.0,
                    )
            elif self.flashinfer_mxfp4_moe_precision == "default":
                x_quant, x_scale = mxfp8_quantize(x, False, alignment=self.hidden_size)
                x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
            else:
                raise NotImplementedError()

            assert x_quant.shape[-1] == self.hidden_size
            assert TopKOutputChecker.format_is_bypassed(topk_output)

            top_k = topk_output.topk_config.top_k
            router_logits = topk_output.router_logits

            with use_symmetric_memory(
                get_tp_group(), disabled=not is_allocation_symmetric()
            ):
                num_tokens = x_quant.shape[0]
                hidden_size = origin_hidden_states_dim
                symm_output = torch.empty(
                    num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device
                )
            trtllm_gen_output = trtllm_fp4_block_scale_moe(
                router_logits.to(torch.bfloat16),
                None,  # routing_bias
                x_quant,
                x_scale,
                layer.w13_weight,  # uint8 (e2m1 x 2)
                layer.w13_weight_scale,  # uint8 (e4m3 x 2)
                layer.w13_weight_bias,  # fp32 per expert per channel
                layer.gemm1_alpha,  # fp32 per expert
                layer.gemm1_beta,  # fp32 per expert
                layer.gemm1_clamp_limit,  # fp32 per expert
                layer.w2_weight,  # uint8 (e2m1 x 2)
                layer.w2_weight_scale,  # ue8m0
                layer.w2_weight_bias,  # fp32 per expert per channel
                None,  # output1_scale_scalar
                None,  # output1_scale_gate_scalar
                None,  # output2_scale_scalar
                layer.num_experts,
                top_k,
                None,  # n_group      # TODO: support n_group
                None,  # topk_group   # TODO: support topk_group
                self.intermediate_size_per_partition,  # padded to multiple of 256
                layer.moe_ep_rank * layer.num_local_experts,  # local_expert_offset
                layer.num_local_experts,  # local num experts
                None,  # routed_scaling_factor
                1,  # routing_method_type, renormalize
                True,  # do finalize
                tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
                output=symm_output,
            )[0]
            return StandardCombineInput(hidden_states=trtllm_gen_output)
        if _use_aiter:
            from sglang.srt.layers.moe.moe_runner.aiter import (
                AiterMoeQuantInfo,
                AiterQuantType,
            )

            if hasattr(torch, "float4_e2m1fn_x2"):
                w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
                w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
            else:
                w13_weight = layer.w13_weight
                w2_weight = layer.w2_weight

            # `.view()` creates a fresh tensor that drops the `is_shuffled`
            # marker we set in process_weights_after_loading. Re-tag it so the
            # downstream aiter.fused_moe selects preshuffle_on kernels.
            if getattr(layer.w13_weight, "is_shuffled", False):
                w13_weight.is_shuffled = True
                w2_weight.is_shuffled = True

            # Skip the explicit pad if x already arrives at the padded
            # hidden_size (the upstream RMSNorm fused the pad into its
            # output — see RMSNorm.x_pad_to_multiple). Saves a separate
            # zero-pad kernel launch per layer.
            if x.shape[-1] == self.hidden_size:
                x_padded = x
            else:
                x_padded = torch.nn.functional.pad(
                    x, (0, self.hidden_pad), mode="constant", value=0.0
                )
            quant_info = AiterMoeQuantInfo(
                w13_weight=w13_weight,
                w2_weight=w2_weight,
                quant_type=AiterQuantType.PER_1X32,
                w13_scale=layer.w13_weight_scale,
                w2_scale=layer.w2_weight_scale,
                b13=layer.w13_weight_bias,
                b2=layer.w2_weight_bias,
                expert_mask=layer.dispatcher.expert_mask_gpu,
                doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
                hidden_pad=self.hidden_pad,
                intermediate_pad=self.intermediate_pad,
                # Triggers aiter's INTERLEAVE gate_mode dispatch (required for our
                # preshuffled gate/up-interleaved weight layout) and applies the
                # model's swiglu clamp. Models populate the same scalar under
                # different MoeRunnerConfig fields: gpt-oss uses `gemm1_clamp_limit`
                # (renamed in `models/gpt_oss.py` from `config.swiglu_limit`); DSv4
                # / FP8 uses `swiglu_limit` directly. Accept either.
                swiglu_limit=(
                    self.moe_runner_config.gemm1_clamp_limit
                    or self.moe_runner_config.swiglu_limit
                    or 0.0
                ),
            )
            return self.runner.run(
                dispatch_output._replace(hidden_states=x_padded), quant_info
            )

        backend = self.runner.runner_backend
        if backend.is_triton_kernels():
            from sglang.srt.layers.moe.moe_runner.triton_kernels import (
                TritonKernelsQuantInfo,
            )

            assert (
                layer.moe_ep_size == 1
            ), "Expert parallel is not supported when using triton kernels"
            quant_info = TritonKernelsQuantInfo(
                w13_weight=(
                    self.w13_weight_triton_tensor
                    if self.w13_weight_triton_tensor is not None
                    else layer.w13_weight
                ),
                w2_weight=(
                    self.w2_weight_triton_tensor
                    if self.w2_weight_triton_tensor is not None
                    else layer.w2_weight
                ),
                w13_bias=getattr(layer, "w13_weight_bias", None),
                w2_bias=getattr(layer, "w2_weight_bias", None),
                w13_precision_config=getattr(self, "w13_precision_config", None),
                w2_precision_config=getattr(self, "w2_precision_config", None),
            )
        else:
            quant_info = TritonMoeQuantInfo(
                w13_weight=layer.w13_weight,
                w2_weight=layer.w2_weight,
                b13=getattr(layer, "w13_weight_bias", None),
                b2=getattr(layer, "w2_weight_bias", None),
            )
        return self.runner.run(dispatch_output, quant_info)


class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
    def create_weights(
        self,
        layer: torch.nn.Module,
        num_experts: int,
        hidden_size: int,
        intermediate_size_per_partition: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ):

        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported

        w13_weight = torch.nn.Parameter(
            torch.empty(
                num_experts,
                2 * intermediate_size_per_partition,
                hidden_size,
                dtype=params_dtype,
            ),
            requires_grad=False,
        )
        w2_weight = torch.nn.Parameter(
            torch.empty(
                num_experts,
                hidden_size,
                intermediate_size_per_partition,
                dtype=params_dtype,
            ),
            requires_grad=False,
        )

        layer.register_parameter("w13_weight", w13_weight)
        set_weight_attrs(w13_weight, extra_weight_attrs)

        layer.register_parameter("w2_weight", w2_weight)
        set_weight_attrs(w2_weight, extra_weight_attrs)

        # Allocate 2 scales for w1 and w3 respectively.
        # They will be combined to a single scale after weight loading.
        w13_weight_scale = torch.nn.Parameter(
            torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
        )
        w2_weight_scale = torch.nn.Parameter(
            torch.ones(num_experts, dtype=torch.float32), requires_grad=False
        )
        layer.register_parameter("w13_weight_scale", w13_weight_scale)
        layer.register_parameter("w2_weight_scale", w2_weight_scale)

        # Add the quantization method used (per tensor/grouped/channel)
        # to ensure the weight scales are loaded in properly
        extra_weight_attrs.update(
            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
        )

        layer.w13_input_scale = None
        layer.w2_input_scale = None

    def mxfp4_quantize(self, w):
        w_shape = w.shape
        w_need_reshape = True if w.dim() != 2 else False

        if w_need_reshape:
            w_last_dim_size = w_shape[-1]
            w = w.view(-1, w_last_dim_size)

        w, mx_scales = dynamic_mxfp4_quant(w)

        if w_need_reshape:
            w_new_shape = w_shape[:-1] + (w.shape[-1],)
            w = w.view(w_new_shape)

        mx_scales = e8m0_shuffle(mx_scales)

        return w, mx_scales

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        w13, w13_mx_scales = self.mxfp4_quantize(layer.w13_weight.data)
        w2, w2_mx_scales = self.mxfp4_quantize(layer.w2_weight.data)

        # Pre-shuffle weight
        is_shuffled = _is_shuffle_moe_mxfp4
        if is_shuffled:
            w13 = shuffle_weight(w13.contiguous(), (16, 16))
            w2 = shuffle_weight(w2.contiguous(), (16, 16))

        layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
        layer.w13_weight.is_shuffled = is_shuffled
        layer.w13_weight_scale = torch.nn.Parameter(w13_mx_scales, requires_grad=False)

        layer.w2_weight = torch.nn.Parameter(w2, requires_grad=False)
        layer.w2_weight.is_shuffled = is_shuffled
        layer.w2_weight_scale = torch.nn.Parameter(w2_mx_scales, requires_grad=False)

    def create_moe_runner(
        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
    ):
        self.moe_runner_config = moe_runner_config
        moe_runner_backend = get_moe_runner_backend()
        if moe_runner_backend.is_auto() and get_moe_a2a_backend().supports_aiter():
            moe_runner_backend = MoeRunnerBackend.AITER

        if moe_runner_backend.is_aiter():
            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
        else:
            # TODO(cwan): refactor other backends
            pass

    def apply(
        self,
        layer: torch.nn.Module,
        dispatch_output: StandardDispatchOutput,
    ) -> CombineInput:
        from sglang.srt.layers.moe.moe_runner.aiter import (
            AiterMoeQuantInfo,
            AiterQuantType,
        )

        if hasattr(torch, "float4_e2m1fn_x2"):
            w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
            w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
        else:
            w13_weight = layer.w13_weight
            w2_weight = layer.w2_weight

        if hasattr(layer.w13_weight, "is_shuffled"):
            w13_weight.is_shuffled = True
            w2_weight.is_shuffled = True

        quant_info = AiterMoeQuantInfo(
            w13_weight=w13_weight,
            w2_weight=w2_weight,
            quant_type=AiterQuantType.PER_1X32,
            w13_scale=layer.w13_weight_scale,
            w2_scale=layer.w2_weight_scale,
            expert_mask=layer.dispatcher.expert_mask_gpu,
        )
        return self.runner.run(dispatch_output, quant_info)