[AMD] Fix GPT-OSS MXFP4 accuracy on ROCm AITER path (#26884)

Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>

python/sglang/srt/environ.py

@@ -381,6 +381,11 @@ class Envs:
     # AMD & ROCm
     SGLANG_USE_AITER = EnvBool(False)
     SGLANG_USE_AITER_UNIFIED_ATTN = EnvBool(False)
+    # Select the gate/up tile layout for AITER MoE: True -> interleave
+    # (matches FlyDSL `gate_mode="interleave"` kernels), False -> separated
+    # (matches `gate_mode="separated"`, the layout used by gptoss_fp4 tuned
+    # configs and by Mxfp4MoEMethod's post-fix weight shuffle).
+    SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)
     SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)
     SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False)
     SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(4096)

python/sglang/srt/layers/moe/moe_runner/aiter.py

@@ -119,6 +119,8 @@ class AiterRunnerCore(MoeRunnerCore):
         from aiter.fused_moe import fused_moe
         from aiter.ops.flydsl.moe_common import GateMode
 
+        from sglang.srt.environ import envs
+
         a1_scale = (
             runner_input.a1_scale
             if runner_input.a1_scale is not None
@@ -131,7 +133,16 @@ class AiterRunnerCore(MoeRunnerCore):
         if runner_input.output_dtype is not None:
             extra["dtype"] = runner_input.output_dtype
         if quant_info.swiglu_limit > 0:
-            extra["gate_mode"] = GateMode.INTERLEAVE.value
+            # Default (INTERLEAVE) preserves the pre-fix behavior for paths
+            # that prepare weights in the gate/up-interleaved layout. Set
+            # `SGLANG_USE_AITER_MOE_GU_ITLV=0` to switch to SEPARATED, which
+            # matches the layout produced by `Mxfp4MoEMethod` (gpt-oss
+            # MXFP4) and the gptoss_fp4 tuned FlyDSL kernels.
+            extra["gate_mode"] = (
+                GateMode.INTERLEAVE.value
+                if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
+                else GateMode.SEPARATED.value
+            )
             extra["swiglu_limit"] = quant_info.swiglu_limit
 
         output = fused_moe(

python/sglang/srt/layers/quantization/mxfp4.py

@@ -154,9 +154,8 @@ if _is_hip:
     # import aiter
     try:
         from aiter.ops.shuffle import (
-            shuffle_scale_a16w4,
+            shuffle_scale,
             shuffle_weight,
-            shuffle_weight_a16w4,
         )
         from aiter.ops.triton.quant import dynamic_mxfp4_quant
         from aiter.utility.fp4_utils import e8m0_shuffle
@@ -774,6 +773,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )
             return
         if _use_aiter:
+            # Bias must be fp32 for the AITER kernels.
             if layer.w13_weight_bias is not None:
                 layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(
                     torch.float32
@@ -781,6 +781,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             if layer.w2_weight_bias is not None:
                 layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)
 
+            # HF GPT-OSS stores w13 as gate/up *interleaved* row pairs
+            # [(g0, u0), (g1, u1), ...]. The AITER MXFP4 fused MoE kernels
+            # (FlyDSL `gate_mode="separated"` and CK `preshuffle_on`) expect
+            # the *separated* layout [gate_0..gate_{N-1}, up_0..up_{N-1}].
+            # De-interleave weights, scales, and bias before the tile shuffle
+            # so the post-shuffle bytes land in the layout the kernel reads.
             e, n, k = layer.w13_weight.shape
             layer.w13_weight.view(torch.uint8).copy_(
                 layer.w13_weight.data.view(torch.uint8)
@@ -795,21 +801,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 .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)
@@ -817,6 +808,33 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 .view(-1, n)
             )
 
+            # ATOM-aligned MXFP4 preshuffle (is_guinterleave=False). Both
+            # `gptoss_fp4_tuned_fmoe.csv` flydsl entries (e.g.
+            # flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w2) and the CK
+            # `module_moe_ck2stages_*_preshuffle_on_*` fallback use the
+            # standard (16,16) tile layout produced here. The previous
+            # shuffle_weight_a16w4 path produced a gate/up-interleaved tile
+            # layout that does not match the separated-gate kernels and
+            # caused silent accuracy loss.
+            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,
+            )
+
             layer.w13_weight_scale = torch.nn.Parameter(
                 shuffled_w13_scale, requires_grad=False
             )
@@ -824,6 +842,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 shuffled_w2_scale, requires_grad=False
             )
 
+            # Tell aiter.fused_moe these weights are already preshuffled so it
+            # picks the preshuffle_on CK / FlyDSL kernels (which match the
+            # actual layout) instead of falling back to preshuffle_off kernels
+            # that interpret the shuffled bytes as a non-shuffled tensor and
+            # produce garbage / OOB accesses.
+            layer.w13_weight.is_shuffled = True
+            layer.w2_weight.is_shuffled = True
+
             return
 
         if self.use_triton_kernels:
@@ -1233,6 +1259,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 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
+
             x_padded = torch.nn.functional.pad(
                 x, (0, self.hidden_pad), mode="constant", value=0.0
             )
@@ -1248,6 +1281,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
                 hidden_pad=self.hidden_pad,
                 intermediate_pad=self.intermediate_pad,
+                swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0,
             )
             return self.runner.run(
                 dispatch_output._replace(hidden_states=x_padded), quant_info

python/sglang/srt/server_args.py

@@ -2162,6 +2162,13 @@ class ServerArgs:
                     logger.warning(
                         "Detected ROCm and MXFP4 quantization format for GPT-OSS model, enabling aiter MXFP4 MOE kernel."
                     )
+                    # The AITER MXFP4 fused-MoE path for GPT-OSS expects the
+                    # SEPARATED gate/up tile layout (matches the
+                    # `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
+                    # Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
+                    # callers default to INTERLEAVE; opt this path out
+                    # unless the user explicitly overrode it.
+                    envs.SGLANG_USE_AITER_MOE_GU_ITLV.set(False)
                 elif is_hip() and envs.SGLANG_USE_AITER.get():
                     # For GPT-OSS bf16 on ROCm with aiter, use triton backend
                     # because aiter CK kernel doesn't support all GEMM dimensions

test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py

@@ -75,7 +75,14 @@ MI35X_GPT_OSS_MODELS = [
             "triton",
             "--trust-remote-code",
         ],
-        env_vars={"SGLANG_USE_AITER": "1"},
+        # AITER MXFP4 fused-MoE for gpt-oss uses the SEPARATED gate/up tile
+        # layout (matches `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
+        # Mxfp4MoEMethod weight shuffle). Other AITER MXFP4 callers default
+        # to INTERLEAVE, so opt out explicitly here.
+        env_vars={
+            "SGLANG_USE_AITER": "1",
+            "SGLANG_USE_AITER_MOE_GU_ITLV": "0",
+        },
     ),
     ModelConfig(
         model_path="openai/gpt-oss-120b",
@@ -93,7 +100,10 @@ MI35X_GPT_OSS_MODELS = [
             "triton",
             "--trust-remote-code",
         ],
-        env_vars={"SGLANG_USE_AITER": "1"},
+        env_vars={
+            "SGLANG_USE_AITER": "1",
+            "SGLANG_USE_AITER_MOE_GU_ITLV": "0",
+        },
     ),
 ]