chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add pre-suffle weight for new aiter MoE support. (#12908)

Browse Source 
Co-authored-by: HAI <hixiao@gmail.com>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
This commit is contained in:
sogalin
2025-11-11 03:42:55 +08:00
committed by GitHub
parent c022107f8b
commit 661c1c97ad
3 changed files with 24 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/layers/quantization/mxfp4.py
View File

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

13
python/sglang/srt/layers/quantization/quark/quark_moe.py
View File

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