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 import zero_copy_context
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.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_exec
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 (
        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
_aiter_k3_opt = _use_aiter and get_bool_env_var("SGLANG_AITER_K3_OPT")
_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()
        # True W4A8: DeepGEMM fp8_fp4 grouped GEMM (SM100 MXF8F6F4 UMMA).
        # Weights stay MXFP4 (e2m1 + ue8m0 g32, zero requantization);
        # activations are quantized to fp8 per-token-group-128.
        self.use_deep_gemm = get_moe_runner_backend().is_deep_gemm()
        self.flashinfer_mxfp4_moe_precision = (
            get_exec().moe.flashinfer_mxfp4_moe_precision
        )
        # When `flashinfer_mxfp4` is enabled, dispatch to one of three FlashInfer
        # entry points depending on the GPU:
        #   - SM100 (Blackwell)  -> trtllm_fp4_block_scale_moe (existing)
        #   - SM120 (Blackwell)  -> cutlass_fused_moe(MXFP8 x MXFP4)
        #   - 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_sm120_supported():
                self._fi_kernel = "cutlass_sm120"
            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, SM100, "
                    "or SM120."
                )

    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 self.use_deep_gemm:
            # DeepGEMM fp8_fp4 grouped GEMM consumes the checkpoint layout
            # directly (packed e2m1 K-major + ue8m0 g32 scales); no padding.
            pass
        elif is_sm100_supported():
            if self.use_flashinfer:
                # FlashInfer trtllm-gen FP4 kernel actual alignment:
                # intermediate: scale shuffle needs M%128==0 → intermediate%64==0
                # hidden: finalize kernel needs K%32==0, block quant needs K%32==0
                # Using 128 alignment (not 256) since 256 was overly conservative
                # and causes 60GB/GPU waste for models like K3 (384 intermediate).
                intermediate_size_per_partition_after_pad = round_up(
                    intermediate_size_per_partition, 128
                )
                hidden_size = round_up(hidden_size, 128)
            else:
                intermediate_size_per_partition_after_pad = round_up(
                    intermediate_size_per_partition, triton_kernels_padding_alignment
                )
        elif self._fi_kernel in ("cutlass_sm90", "cutlass_sm120"):
            # CUTLASS mixed-input GEMM dimensions must be % 128 == 0. The
            # kernels also expect ``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 the architecture-specific
            # CUTLASS post-load processor 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:
            # Expert intermediate is padded to 128 or 256. On K3 (TP8) the 256
            # default costs +33% mem (3072/8=384 -> 512). K3 FlyDSL MoE works with
            # 128 padding, so relaxing to 128 saves that memory and lets the TP8
            # decode CUDA graph range grow from 20 to 120.
            _inter_align = 128 if _aiter_k3_opt else 256
            intermediate_size_per_partition_after_pad = round_up(
                intermediate_size_per_partition, _inter_align
            )

            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_scale.quant_method = "group"

        create_bias = with_bias or not _is_hip
        if create_bias:
            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_scale.quant_method = "group"

        if create_bias:
            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_sm100_supported()
                and not is_sm120_supported()
            ):
                raise RuntimeError("MXFP4 Marlin requires SM90+.")
            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 and getattr(
                self.moe_runner_config, "gate_up_interleaved", True
            ):
                deinterleave_moe_mxfp4_w13_for_marlin(layer)
            prepare_moe_mxfp4_layer_for_marlin(layer)
            layer._mxfp4_backend = "marlin"
            return

        if self.use_deep_gemm:
            from deep_gemm import transform_sf_into_required_layout

            # Packed fp4 (e2m1 x2 per byte) weights: DeepGEMM expects int8.
            layer.w13_weight.data = layer.w13_weight.data.view(torch.int8)
            layer.w2_weight.data = layer.w2_weight.data.view(torch.int8)
            # Checkpoint scales are uint8 e8m0 (biased exponents). DeepGEMM
            # SM100 needs them in packed-UE8M0 TMA-aligned MN-major layout.
            # Round-trip through fp32 is exact (values are powers of two).
            for scale_name, weight in (
                ("w13_weight_scale", layer.w13_weight),
                ("w2_weight_scale", layer.w2_weight),
            ):
                scale = getattr(layer, scale_name)
                num_experts, n, _ = scale.data.shape
                k = weight.shape[2] * 2
                scale_f32 = scale.data.view(torch.float8_e8m0fnu).to(torch.float32)
                scale.data = transform_sf_into_required_layout(
                    scale_f32,
                    mn=n,
                    k=k,
                    recipe=(1, 32),
                    num_groups=num_experts,
                    disable_ue8m0_cast=False,
                )
            if get_moe_a2a_backend().is_megamoe():
                # MegaMoE consumes the same transformed sf, plus its own
                # interleaved/UTCCP weight layout. K3 routes EVERY batch
                # through mega (the megamoe backend has no a2a fallback), so
                # the contig-layout originals are dead weight — repoint the
                # params at the mega tensors to reclaim ~1.9GB/layer/rank
                # (keeping both layouts OOMs: 92 layers double the experts).
                from deep_gemm import transform_weights_for_mega_moe

                l1_pair, l2_pair = transform_weights_for_mega_moe(
                    (layer.w13_weight.data, layer.w13_weight_scale.data),
                    (layer.w2_weight.data, layer.w2_weight_scale.data),
                )
                layer.mega_l1_weights = l1_pair
                layer.mega_l2_weights = l2_pair
                layer.w13_weight.data = l1_pair[0]
                layer.w13_weight_scale.data = l1_pair[1]
                layer.w2_weight.data = l2_pair[0]
                layer.w2_weight_scale.data = l2_pair[1]
                layer._mega_moe_weights_built = True
            layer._mxfp4_backend = "deep_gemm"
            return

        if self._fi_kernel == "cutlass_sm90":
            self._process_weights_for_sm90_cutlass(layer)
            return
        if self._fi_kernel == "cutlass_sm120":
            self._process_weights_for_sm120_cutlass(layer)
            return
        if self.use_flashinfer:
            # Per-expert buffers are local (create_weights uses num_local_experts);
            # the global self.num_experts here breaks EP>1. Mirrors the SM90 path.
            E = layer.num_local_experts
            _alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702
            _limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0
            layer.gemm1_alpha = Parameter(
                torch.tensor([_alpha] * E, dtype=torch.float32).cuda(),
                requires_grad=False,
            )
            layer.gemm1_beta = Parameter(
                torch.tensor([1.0] * E, dtype=torch.float32).cuda(),
                requires_grad=False,
            )
            layer.gemm1_clamp_limit = Parameter(
                torch.tensor([_limit] * E, 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] == E
                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] == E
                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] == E
                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] == E
                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] == E
                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)

            if getattr(layer.moe_runner_config, "gate_up_interleaved", True):
                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)
            else:
                # Non-interleaved layout (e.g. K3 Latent MoE): first half = gate
                # (w1), second half = up (w3). The trtllm-gen fused gated-act
                # epilogue wants rows interleaved as (up_i, gate_i) PAIRS -
                # the layout get_reorder_rows_for_gated_act_gemm_row_indices
                # produces from [linear; gate] halves, and what the
                # interleaved branch above arrives at via swap_every_two_rows.
                # A plain halves swap keeps rows blocked and pairs up_i with
                # up_{i+1}, which scrambles the gated activation.
                half = w13_weight.shape[-2] // 2
                pair_idx = torch.empty(2 * half, dtype=torch.long)
                pair_idx[0::2] = torch.arange(half) + half  # up (w3)
                pair_idx[1::2] = torch.arange(half)  # gate (w1)
                w13_weight = w13_weight[..., pair_idx, :].contiguous()
                w13_weight_scale = w13_weight_scale[..., pair_idx, :].contiguous()
                w13_bias = w13_bias[..., pair_idx].contiguous()

            # 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(E):
                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(
                    E,
                    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(
                    E,
                    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(E, -1),
                requires_grad=False,
            )
            layer.w2_weight_bias = Parameter(
                torch.stack(gemm2_bias_shuffled).reshape(E, -1),
                requires_grad=False,
            )
            return
        if _use_aiter:
            if getattr(layer, "w13_weight_bias", None) is not None:
                layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(
                    torch.float32
                )
            if getattr(layer, "w2_weight_bias", None) is not None:
                layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)

            e, n, k = layer.w13_weight.shape
            gate_up_interleaved = getattr(
                layer.moe_runner_config, "gate_up_interleaved", True
            )

            if gate_up_interleaved:
                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)
                )
                if getattr(layer, "w13_weight_bias", None) is not None:
                    layer.w13_weight_bias.data = (
                        layer.w13_weight_bias.data.view(-1, n // 2, 2)
                        .permute(0, 2, 1)
                        .contiguous()
                        .view(-1, n)
                    )

            k3_situ_a8w4 = (
                os.environ.get("AITER_SITUV2_A8W4", "0") == "1"
                and getattr(layer.moe_runner_config, "activation", None) == "situ"
            )
            use_aiter_gu_interleave = k3_situ_a8w4 or (
                envs.SGLANG_USE_AITER_MOE_GU_ITLV.get() and gate_up_interleaved
            )
            if use_aiter_gu_interleave:
                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
        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.

        _interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True)

        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.
            if _interleaved:
                gate_rows = unpadded_w13[:, 0::2, :]  # [E, N_un, last_un]
                up_rows = unpadded_w13[:, 1::2, :]  # [E, N_un, last_un]
            else:
                # Non-interleaved: first half=gate, second half=up
                gate_rows = unpadded_w13[:, :N_un, :]
                up_rows = unpadded_w13[:, N_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.
        if _interleaved:
            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]
        else:
            w13_bias_gate = layer.w13_weight_bias.data[:, :N_un]
            w13_bias_up = layer.w13_weight_bias.data[:, 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 (read from runner config, fallback to GPT-OSS defaults)
        _sm90_alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702
        _sm90_limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0
        layer.swiglu_alpha = Parameter(
            torch.full((E,), _sm90_alpha, 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,), _sm90_limit, 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 _process_weights_for_sm120_cutlass(self, layer):
        """Prepare GPT-OSS MXFP4 experts for FlashInfer CUTLASS on SM120.

        GPT-OSS stores gate/up rows pair-wise as
        ``[gate_0, up_0, gate_1, up_1, ...]``. FlashInfer's fused MoE consumes
        two contiguous halves in ``[up; gate]`` order. Build that layout after
        checkpoint loading so padding cannot move the split, pad both GEMMs to
        CUTLASS's 128-element alignment, and swizzle the native E8M0 scales for
        the SM120 MXFP8-by-MXFP4 kernels. Packed FP4 weight bytes themselves do
        not need an SM120 permutation.
        """
        from flashinfer import block_scale_interleave

        sf_block_size = 32
        N_un = layer.w13_weight.shape[1] // 2
        K_un = layer.w13_weight.shape[2] * 2
        N_pad = self._padded_intermediate
        K_pad = self._padded_hidden
        E = layer.num_local_experts
        device = layer.w13_weight.device

        def _stack_up_gate_w13(unpadded, last_pad, last_un):
            gate_rows = unpadded[:, 0::2, :]
            up_rows = unpadded[:, 1::2, :]
            out = torch.zeros(
                E, 2 * N_pad, last_pad, dtype=unpadded.dtype, device=device
            )
            out[:, :N_un, :last_un] = up_rows
            out[:, N_pad : N_pad + N_un, :last_un] = gate_rows
            return out

        w13_padded = _stack_up_gate_w13(layer.w13_weight.data, 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_dtype = layer.w13_weight_bias.dtype
        w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device)
        w13_bias_padded[:, :N_un] = layer.w13_weight_bias.data[:, 1::2]
        w13_bias_padded[:, N_pad : N_pad + N_un] = layer.w13_weight_bias.data[:, 0::2]

        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_padded = _pad_w2_3d(layer.w2_weight.data, 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

        w13_scale_interleaved = block_scale_interleave(w13_scale_padded)
        w2_scale_interleaved = block_scale_interleave(w2_scale_padded)

        layer.w13_weight = Parameter(w13_padded, requires_grad=False)
        layer.w2_weight = Parameter(w2_padded, requires_grad=False)
        layer.w13_weight_scale = Parameter(
            w13_scale_interleaved.reshape_as(w13_scale_padded),
            requires_grad=False,
        )
        layer.w2_weight_scale = Parameter(
            w2_scale_interleaved.reshape_as(w2_scale_padded),
            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)

        layer.swiglu_alpha = Parameter(
            torch.full((E,), 1.702, dtype=torch.float32, device=device),
            requires_grad=False,
        )
        layer.swiglu_beta = Parameter(
            torch.ones(E, 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,
        )
        # The MXFP4 ABI uses a neutral global weight scale for each GEMM.
        layer.mxfp4_weight_global_scale = Parameter(
            torch.ones(E, dtype=torch.float32, device=device),
            requires_grad=False,
        )
        layer._mxfp4_backend = "flashinfer_cutlass_sm120"
        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, so the
            # checkpoint's "silu" has to be translated. K3's SiTU is the one
            # activation the kernel selects on its own -- leave it alone.
            aiter_config = moe_runner_config
            if aiter_config.activation != "situ":
                aiter_config = replace(aiter_config, activation="swiglu")
            self.runner = MoeRunner(moe_runner_backend, aiter_config)
        elif (
            moe_runner_backend.is_triton_kernels()
            or moe_runner_backend.is_triton()
            or moe_runner_backend.is_marlin()
            or moe_runner_backend.is_deep_gemm()
        ):
            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
        elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
            "cutlass_sm90",
            "cutlass_sm120",
        ):
            # 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_marlin(self, layer, dispatch_output):
        """MXFP4 x BF16 MoE via the Marlin runner. The quant_info (incl. the
        dispatcher's EP mapping) is shared with
        :class:`~sglang.srt.layers.quantization.mxfp4_marlin_moe.Mxfp4MarlinMoEMethod`;
        only the input padding is local (``hidden_size`` is pre-rounded to the
        Marlin-padded width at create_weights)."""
        from sglang.srt.layers.quantization.mxfp4_marlin_moe import (
            build_marlin_moe_quant_info,
        )

        x = dispatch_output.hidden_states
        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 = build_marlin_moe_quant_info(layer)
        return self.runner.run(
            dispatch_output._replace(hidden_states=x_padded), quant_info
        )

    def _apply_sm120_cutlass(self, layer, dispatch_output):
        """SM120 GPT-OSS MXFP8 x MXFP4 MoE via 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,
            mxfp4_weight_global_scale=layer.mxfp4_weight_global_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

        if self.use_deep_gemm:
            # Handles standard AND deepep_ll/deepep_normal dispatch formats via
            # the runner's registered pre/post-permute functions. Must run
            # before the `.topk_output` unpack below: deepep dispatch outputs
            # carry topk_ids/topk_weights directly and have no `.topk_output`.
            from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo

            quant_info = DeepGemmMoeQuantInfo(
                w13_weight=layer.w13_weight,
                w2_weight=layer.w2_weight,
                use_fp8=True,
                w13_scale=layer.w13_weight_scale,
                w2_scale=layer.w2_weight_scale,
                block_shape=[128, 128],
                is_fp4_experts=True,
            )
            return self.runner.run(dispatch_output, quant_info)

        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)
            return self._apply_marlin(layer, dispatch_output)

        if self._fi_kernel == "cutlass_sm90":
            return self._apply_sm90_cutlass(layer, dispatch_output)
        if self._fi_kernel == "cutlass_sm120":
            return self._apply_sm120_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]
            # Filled by the staged K3 route+pack+quant fusion below; the pack
            # site further down falls back to PackTopkIds when it is None.
            prepared_packed_topk = None
            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":
                if x.shape[-1] == self.hidden_size:
                    if x.dim() > 2:
                        x = x.view(-1, x.shape[-1])
                    # K3 staged fusion (route_quant_handoff): the routing
                    # dispatch already quantized these rows and packed the
                    # topk ids in the fused route launch — consume both and
                    # skip the two standalone kernels. Identity-verified;
                    # a miss runs the unfused chain below.
                    from sglang.srt.layers.moe import route_quant_handoff

                    prepared = route_quant_handoff.take(x)
                    if prepared is not None:
                        prepared_packed_topk, x_quant, x_scale = prepared
                        x_scale = x_scale.view(torch.float8_e4m3fn)
                    else:
                        from sglang.kernels.ops.quantization.per_token_group_quant import (
                            per_token_group_quant,
                        )

                        x_quant, x_scale = per_token_group_quant(
                            x, group_size=32, scale_ue8m0=True
                        )
                        x_scale = x_scale.view(torch.float8_e4m3fn)
                else:
                    from sglang.srt.layers.quantization.fp8_utils import (
                        flashinfer_mxfp8_quantize,
                    )

                    x_quant, x_scale = flashinfer_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
            is_standard = TopKOutputChecker.format_is_standard(topk_output)
            # The situ path accepts precomputed (standard) routing; the
            # public path below is logits-only.
            assert is_standard or TopKOutputChecker.format_is_bypassed(
                topk_output
            ), f"unsupported topk format: {topk_output.format}"
            if is_standard:
                assert (
                    self.moe_runner_config.activation == "situ"
                ), "standard topk output only wired for the situ path"
                top_k = topk_output.topk_ids.shape[1]
                router_logits = None
            else:
                top_k = topk_output.topk_config.top_k
                router_logits = topk_output.router_logits

            num_tokens = x_quant.shape[0]
            hidden_size = origin_hidden_states_dim
            # The K3 fused-front path publishes its [latent | shared] buffer
            # slice as the output destination (zero_copy_context); writing
            # the finalize output there directly skips this allocation and
            # the copy_ back in _forward_fused. The slice lives in the same
            # symmetric buffer the caller all-reduces.
            symm_output = zero_copy_context.get_moe_output_spec(
                torch.Size((num_tokens, hidden_size)),
                torch.bfloat16,
                x_quant.device,
            )
            if symm_output is None:
                with use_symmetric_memory(
                    get_tp_group(), disabled=not is_allocation_symmetric()
                ):
                    symm_output = torch.empty(
                        num_tokens,
                        hidden_size,
                        dtype=torch.bfloat16,
                        device=x_quant.device,
                    )

            if self.moe_runner_config.activation == "situ":
                # SiTU is only in the private trtllm-gen cubin pool (the
                # public artifact bakes swiglu into the fused-act cubins and
                # silently computes the wrong activation). Routing must also
                # be noaux_tc (sigmoid + correction bias, DeepSeekV3 method),
                # not the renormalize-softmax default below.
                from sglang.kernels.ops.moe import trtllm_gen_moe as situ_moe

                if not situ_moe.available():
                    raise RuntimeError(
                        "activation='situ' with the flashinfer_mxfp4 runner "
                        "needs the SiTU cubin pool: set "
                        "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL (see "
                        "sglang/kernels/ops/moe/trtllm_gen_moe.py)."
                    )
                # EP is cubin-internal: each rank computes its local expert slice
                # [offset, +num_local) and the caller all-reduces. ep=1 -> TP path.
                local_expert_offset = layer.moe_ep_rank * layer.num_local_experts
                if TopKOutputChecker.format_is_standard(topk_output):
                    # Precomputed routing (radix router upstream): skip the
                    # in-op routing kernels entirely. At small T the in-op
                    # single-CTA routing costs ~22 us/layer vs ~6 us for the
                    # external radix router.
                    if prepared_packed_topk is not None:
                        packed_topk = prepared_packed_topk
                    else:
                        from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds

                        packed_topk = PackTopkIds.execute(
                            topk_output.topk_ids, topk_output.topk_weights
                        )
                    # Deferred finalize (K3 forward_deferred_finalize): return
                    # the finalize inputs instead of the finalized output.
                    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
                        _deferred_finalize_enabled,
                    )

                    defer_finalize = _deferred_finalize_enabled.get()
                    result = situ_moe.trtllm_fp4_block_scale_routed_moe(
                        packed_topk_ids=packed_topk,
                        hidden_states=x_quant,
                        hidden_states_scale=x_scale,
                        gemm1_weights=layer.w13_weight,
                        gemm1_weights_scale=layer.w13_weight_scale,
                        gemm1_alpha=layer.gemm1_alpha,
                        # SiTuGlu: gatedActBeta is the linear-half tanh
                        # clip; K3 stores it in gemm1_clamp_limit.
                        gemm1_beta=layer.gemm1_clamp_limit,
                        gemm2_weights=layer.w2_weight,
                        gemm2_weights_scale=layer.w2_weight_scale,
                        output1_scale_scalar=None,
                        output1_scale_gate_scalar=None,
                        output2_scale_scalar=None,
                        num_experts=layer.num_experts,
                        top_k=packed_topk.shape[1],
                        intermediate_size=self.intermediate_size_per_partition,
                        activation_type=situ_moe.ACTIVATION_SITU,
                        local_expert_offset=local_expert_offset,
                        local_num_experts=layer.num_local_experts,
                        output=symm_output,
                        do_finalize=not defer_finalize,
                    )
                    if defer_finalize:
                        from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
                            FlashInferTrtllmDeferredFinalizeOutput,
                        )

                        gemm2_out, topk_weights, expanded_idx = result
                        result = FlashInferTrtllmDeferredFinalizeOutput(
                            gemm2_out=gemm2_out,
                            expert_weights=topk_weights,
                            expanded_idx_to_permuted_idx=expanded_idx,
                            top_k=packed_topk.shape[1],
                        )
                    return StandardCombineInput(hidden_states=result)

                # Bypassed topk: route from logits inside the op.
                correction_bias = topk_output.topk_config.correction_bias
                bias_bf16 = getattr(layer, "_situ_routing_bias_bf16", None)
                if bias_bf16 is None and correction_bias is not None:
                    bias_bf16 = correction_bias.to(torch.bfloat16)
                    layer._situ_routing_bias_bf16 = bias_bf16
                situ_moe.trtllm_fp4_block_scale_moe(
                    # router_logits is a row-strided slice of the K3 fused
                    # front GEMM output; the FFI reads it as dense.
                    routing_logits=router_logits.to(torch.bfloat16).contiguous(),
                    routing_bias=bias_bf16,
                    hidden_states=x_quant,
                    hidden_states_scale=x_scale,
                    gemm1_weights=layer.w13_weight,
                    gemm1_weights_scale=layer.w13_weight_scale,
                    gemm1_alpha=layer.gemm1_alpha,
                    # SiTuGlu: gatedActBeta is the linear-half tanh clip;
                    # K3 stores it in gemm1_clamp_limit (situ_linear_beta).
                    gemm1_beta=layer.gemm1_clamp_limit,
                    gemm2_weights=layer.w2_weight,
                    gemm2_weights_scale=layer.w2_weight_scale,
                    output1_scale_scalar=None,
                    output1_scale_gate_scalar=None,
                    output2_scale_scalar=None,
                    num_experts=layer.num_experts,
                    top_k=top_k,
                    n_group=topk_output.topk_config.num_expert_group,
                    topk_group=topk_output.topk_config.topk_group,
                    intermediate_size=self.intermediate_size_per_partition,
                    routed_scaling_factor=(
                        topk_output.topk_config.routed_scaling_factor or 1.0
                    ),
                    routing_method_type=situ_moe.ROUTING_DEEPSEEK_V3,
                    activation_type=situ_moe.ACTIVATION_SITU,
                    norm_topk_prob=topk_output.topk_config.renormalize,
                    local_expert_offset=local_expert_offset,
                    local_num_experts=layer.num_local_experts,
                    output=symm_output,
                )
                return StandardCombineInput(hidden_states=symm_output)

            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 if self.with_bias else None,
                b2=layer.w2_weight_bias if self.with_bias else None,
                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,
                # Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU
                # uses gemm1_clamp_limit as linear_beta, which is forwarded by
                # the AITER runner and must not be treated as swiglu_limit.
                swiglu_limit=(
                    0.0
                    if self.moe_runner_config.activation == "situ"
                    else (
                        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)