From 661c1c97ad445a14e18f492155ff50e2aeb72113 Mon Sep 17 00:00:00 2001 From: sogalin <39478626+sogalin@users.noreply.github.com> Date: Tue, 11 Nov 2025 03:42:55 +0800 Subject: [PATCH] Add pre-suffle weight for new aiter MoE support. (#12908) Co-authored-by: HAI Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> --- docker/rocm.Dockerfile | 7 ++++--- python/sglang/srt/layers/quantization/mxfp4.py | 8 ++++++++ .../srt/layers/quantization/quark/quark_moe.py | 13 ++++++++++++- 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index d0a2c42cc..12867e2f1 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -106,11 +106,11 @@ RUN git clone ${AITER_REPO} \ && git submodule update --init --recursive RUN cd aiter \ && if [ "$BUILD_AITER_ALL" = "1" ] && [ "$BUILD_LLVM" = "1" ]; then \ - HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST AITER_MXFP4_MOE_SF=1 python setup.py develop; \ elif [ "$BUILD_AITER_ALL" = "1" ]; then \ - PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST AITER_MXFP4_MOE_SF=1 python setup.py develop; \ else \ - GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + GPU_ARCHS=$GPU_ARCH_LIST AITER_MXFP4_MOE_SF=1 python setup.py develop; \ fi # ----------------------- @@ -295,6 +295,7 @@ RUN python3 -m pip install --no-cache-dir \ # ----------------------- # Performance environment variable. +ENV AITER_MXFP4_MOE_SF=1 ENV HIP_FORCE_DEV_KERNARG=1 ENV HSA_NO_SCRATCH_RECLAIM=1 ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index bbb2daa7d..847eaf0ee 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -39,6 +39,7 @@ from sglang.srt.layers.quantization.utils import is_layer_skipped from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( direct_register_custom_op, + get_bool_env_var, is_cuda, is_flashinfer_available, is_hip, @@ -71,12 +72,14 @@ if TYPE_CHECKING: ) _is_hip = is_hip() +_is_shuffle_moe_mxfp4 = get_bool_env_var("AITER_MXFP4_MOE_SF") and _is_hip if _is_hip: # import aiter try: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe + from aiter.ops.shuffle import shuffle_weight from aiter.ops.triton.quant import dynamic_mxfp4_quant from aiter.utility.fp4_utils import e8m0_shuffle except ImportError as err: @@ -800,6 +803,11 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase): 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 + if _is_shuffle_moe_mxfp4: + 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_scale = torch.nn.Parameter(w13_mx_scales, requires_grad=False) diff --git a/python/sglang/srt/layers/quantization/quark/quark_moe.py b/python/sglang/srt/layers/quantization/quark/quark_moe.py index 3d2d52cd2..9607d392e 100644 --- a/python/sglang/srt/layers/quantization/quark/quark_moe.py +++ b/python/sglang/srt/layers/quantization/quark/quark_moe.py @@ -8,11 +8,12 @@ from typing import TYPE_CHECKING, Any import torch from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe +from aiter.ops.shuffle import shuffle_weight from aiter.utility.fp4_utils import e8m0_shuffle from sglang.srt.layers.moe import MoeRunnerConfig from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase -from sglang.srt.utils import is_hip, set_weight_attrs +from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( @@ -24,6 +25,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) _is_hip = is_hip() +_is_shuffle_moe_mxfp4 = get_bool_env_var("AITER_MXFP4_MOE_SF") and _is_hip __all__ = ["QuarkMoEMethod", "QuarkW4A4MXFp4MoEMethod"] @@ -167,6 +169,15 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod): # layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) + # Pre-shuffle weight + if _is_shuffle_moe_mxfp4: + layer.w13_weight.data = shuffle_weight( + layer.w13_weight.contiguous(), (16, 16) + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight.contiguous(), (16, 16) + ) + def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ):