From ef1c5127545ed9ee26c9d3ba0f55ce370c4c247d Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Thu, 29 Jan 2026 17:50:11 +0800 Subject: [PATCH] Add aiter bias moe support in gpt-oss mxfp4 model (#17735) Co-authored-by: wunhuang --- .../sglang/srt/layers/quantization/mxfp4.py | 109 +++++++++++++++++- python/sglang/srt/server_args.py | 7 ++ 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 537405e2d..3690b4d59 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -51,6 +51,7 @@ from sglang.srt.utils import ( round_up, set_weight_attrs, ) +from sglang.srt.utils.common import get_bool_env_var from sglang.srt.utils.custom_op import register_custom_op _is_sm100_supported = is_cuda() and is_sm100_supported() @@ -75,6 +76,7 @@ if TYPE_CHECKING: ) _is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_shuffle_moe_mxfp4 = is_gfx95_supported() if _is_hip: @@ -82,7 +84,11 @@ if _is_hip: try: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe - from aiter.ops.shuffle import shuffle_weight + from aiter.ops.shuffle import ( + 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: @@ -292,6 +298,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): intermediate_size_per_partition_after_pad = round_up( intermediate_size_per_partition, 64 ) + elif _use_aiter: + + intermediate_size_per_partition_after_pad = round_up( + intermediate_size_per_partition, 256 + ) + + hidden_size = round_up(hidden_size, 256) + self.hidden_pad = hidden_size - layer.hidden_size + self.intermediate_pad = ( + intermediate_size_per_partition_after_pad + - layer.intermediate_size_per_partition + ) elif has_triton_kernels: # TODO: this is a hack to make # intermediate_size_per_partition_after_pad the same as the @@ -530,6 +548,58 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): requires_grad=False, ) return + if _use_aiter: + if layer.w13_weight_bias is not None: + layer.w13_weight_bias.data = layer.w13_weight_bias.data.to( + torch.float32 + ) + if layer.w2_weight_bias is not None: + layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32) + + e, n, k = layer.w13_weight.shape + layer.w13_weight.view(torch.uint8).copy_( + layer.w13_weight.data.view(torch.uint8) + .view(e, n // 2, 2, k) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, k) + ) + layer.w13_weight_scale.data = ( + layer.w13_weight_scale.data.view(e, n // 2, 2, -1) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, -1) + ) + + layer.w13_weight.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, + ) + + layer.w13_weight_bias.data = ( + layer.w13_weight_bias.data.view(-1, n // 2, 2) + .permute(0, 2, 1) + .contiguous() + .view(-1, n) + ) + + 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: @@ -680,6 +750,43 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): output=symm_output, )[0] return StandardCombineInput(hidden_states=trtllm_gen_output) + if _use_aiter: + topk_weights, topk_ids, _ = topk_output + + 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 + + origi_hidden_size = self.hidden_size - self.hidden_pad + + x = torch.nn.functional.pad( + x, + (0, self.hidden_pad), + mode="constant", + value=0.0, + ) + + output = fused_moe( + x, + w13_weight, + w2_weight, + topk_weights, + topk_ids, + expert_mask=layer.expert_mask_gpu, + activation=ActivationType.Swiglu, + quant_type=QuantType.per_1x32, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, + hidden_pad=self.hidden_pad, + intermediate_pad=self.intermediate_pad, + bias1=layer.w13_weight_bias, + bias2=layer.w2_weight_bias, + ) + return StandardCombineInput(hidden_states=output) backend = self.runner.runner_backend if backend.is_triton_kernels(): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6333bd26c..59098f7fe 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1358,6 +1358,13 @@ class ServerArgs: logger.warning( "Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel." ) + elif ( + is_hip() and get_bool_env_var("SGLANG_USE_AITER") + ) and is_mxfp4_quant_format: + self.moe_runner_backend = "auto" + logger.warning( + "Detected ROCm and MXFP4 quantization format for GPT-OSS model, enabling aiter MXFP4 MOE kernel." + ) elif self.ep_size == 1 and is_triton_kernels_available(): self.moe_runner_backend = "triton_kernel" logger.warning(