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

Add support for new aiter version (AR accuracy, is_shuffled PR) (#13554)

Browse Source 
Co-authored-by: sogalin <39478626+sogalin@users.noreply.github.com>
This commit is contained in:
Thomas Wang
2025-11-20 14:17:57 +08:00
committed by GitHub
parent c7b37b7074
commit 7dcf910d53
3 changed files with 25 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -39,9 +39,9 @@ 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_gfx95_supported,
is_hip,
is_sm100_supported,
is_triton_kernels_available,
@@ -72,7 +72,7 @@ if TYPE_CHECKING:
)
_is_hip = is_hip()
_is_shuffle_moe_mxfp4 = get_bool_env_var("AITER_MXFP4_MOE_SF") and _is_hip
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
if _is_hip:
# import aiter
@@ -804,14 +804,17 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
w2, w2_mx_scales = self.mxfp4_quantize(layer.w2_weight.data)
# Pre-shuffle weight
if _is_shuffle_moe_mxfp4:
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(
@@ -842,6 +845,10 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
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
output = fused_moe(
x,
w13_weight,

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

@@ -13,7 +13,12 @@ from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
from sglang.srt.layers.quantization.utils import all_close_1d, per_tensor_dequantize
from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs
from sglang.srt.utils import (
get_bool_env_var,
is_gfx95_supported,
is_hip,
set_weight_attrs,
)
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import (
@@ -24,8 +29,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
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
__all__ = ["QuarkMoEMethod", "QuarkW4A4MXFp4MoEMethod"]
@@ -190,6 +194,8 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
@@ -220,6 +226,10 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
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
output = fused_moe(
x,
w13_weight,