[AMD][WA] force to use gate_mode interleaved to fix tp2/tp4/tp8 acc issue (#27201)

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

python/sglang/srt/layers/attention/aiter_backend.py

@@ -902,10 +902,12 @@ class AiterAttnBackend(AttentionBackend):
                     )
 
                     if self.use_sliding_window_kv_pool:
+                        # AITER attention kernels require int32 page indices;
+                        # full_to_swa_index_mapping is stored as int64.
                         swa_page_table = (
                             self.token_to_kv_pool.translate_loc_from_full_to_swa(
                                 kv_indices
-                            )
+                            ).to(torch.int32)
                         )
 
                         kv_indices = self._transform_table_1_to_real(kv_indices)
@@ -1381,10 +1383,13 @@ class AiterAttnBackend(AttentionBackend):
                 )
 
                 if self.use_sliding_window_kv_pool:
+                    # AITER attention kernels (e.g. mha_batch_prefill_func)
+                    # require int32 page indices; full_to_swa_index_mapping is
+                    # stored as int64.
                     swa_page_table = (
                         self.token_to_kv_pool.translate_loc_from_full_to_swa(
                             self.indices_updater_prefill.kv_indices
-                        )
+                        ).to(torch.int32)
                     )
 
                 self.forward_metadata = ForwardMetadata(
@@ -1580,10 +1585,12 @@ class AiterAttnBackend(AttentionBackend):
                         ]
 
                         if self.use_sliding_window_kv_pool:
+                            # AITER attention kernels require int32 page indices;
+                            # full_to_swa_index_mapping is stored as int64.
                             swa_page_indices = (
                                 self.token_to_kv_pool.translate_loc_from_full_to_swa(
                                     page_indices
-                                )
+                                ).to(torch.int32)
                             )
 
                             page_indices = self._transform_table_1_to_real(page_indices)

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

@@ -28,11 +28,11 @@ from torch.nn.parameter import Parameter
 # cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
 # setdefault preserves any explicit user override.
 os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
-
 from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
+from sglang.srt.environ import envs
 from sglang.srt.layers.amx_utils import (
     CPUQuantMethod,
     _amx_process_weight_after_loading,
@@ -155,7 +155,9 @@ if _is_hip:
     try:
         from aiter.ops.shuffle import (
             shuffle_scale,
+            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
@@ -773,7 +775,6 @@ 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,12 +782,6 @@ 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)
@@ -808,32 +803,46 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 .view(-1, n)
             )
 
-            # ATOM-aligned MXFP4 preshuffle (is_guinterleave=False). Both
+            if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get():
-            # `gptoss_fp4_tuned_fmoe.csv` flydsl entries (e.g.
+                layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True)
-            # flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w2) and the CK
+                shuffled_w13_scale = shuffle_scale_a16w4(
-            # `module_moe_ck2stages_*_preshuffle_on_*` fallback use the
+                    layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
-            # standard (16,16) tile layout produced here. The previous
+                    self.num_experts,
-            # shuffle_weight_a16w4 path produced a gate/up-interleaved tile
+                    True,
-            # layout that does not match the separated-gate kernels and
+                )
-            # caused silent accuracy loss.
+
-            layer.w13_weight.data = shuffle_weight(
+                layer.w2_weight.data = shuffle_weight_a16w4(layer.w2_weight, 16, False)
-                layer.w13_weight, is_guinterleave=False, gate_up=True
+                shuffled_w2_scale = shuffle_scale_a16w4(
-            )
+                    layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
-            shuffled_w13_scale = shuffle_scale(
+                    self.num_experts,
-                layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
+                    False,
-                experts_cnt=self.num_experts,
+                )
-                is_guinterleave=False,
+            else:
-                gate_up=True,
+                layer.w13_weight.data = shuffle_weight(
-            )
+                    layer.w13_weight, is_guinterleave=False, gate_up=True
-            layer.w2_weight.data = shuffle_weight(
+                )
-                layer.w2_weight, is_guinterleave=False, gate_up=False
+                shuffled_w13_scale = shuffle_scale(
-            )
+                    layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
-            shuffled_w2_scale = shuffle_scale(
+                    experts_cnt=self.num_experts,
-                layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
+                    is_guinterleave=False,
-                experts_cnt=self.num_experts,
+                    gate_up=True,
-                is_guinterleave=False,
+                )
-                gate_up=False,
+                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,
+                )
+
+            # shuffle_weight_a16w4(gate_up=True) above preshuffles w13 into aiter's
+            # preshuffle + gate/up-interleaved layout. Tag the Parameter so apply()
+            # can carry the metadata across .view(float4_e2m1fn_x2) and aiter's
+            # fused_moe selects the preshuffle_on kernel family.
+            layer.w13_weight.is_shuffled = True
+            layer.w2_weight.is_shuffled = True
 
             layer.w13_weight_scale = torch.nn.Parameter(
                 shuffled_w13_scale, requires_grad=False
@@ -842,14 +851,6 @@ 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:
@@ -1281,7 +1282,17 @@ 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,
+                # Triggers aiter's INTERLEAVE gate_mode dispatch (required for our
+                # preshuffled gate/up-interleaved weight layout) and applies the
+                # model's swiglu clamp. Models populate the same scalar under
+                # different MoeRunnerConfig fields: gpt-oss uses `gemm1_clamp_limit`
+                # (renamed in `models/gpt_oss.py` from `config.swiglu_limit`); DSv4
+                # / FP8 uses `swiglu_limit` directly. Accept either.
+                swiglu_limit=(
+                    self.moe_runner_config.gemm1_clamp_limit
+                    or 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

@@ -2175,13 +2175,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
+                    ## The AITER MXFP4 fused-MoE path for GPT-OSS expects the
-                    # SEPARATED gate/up tile layout (matches the
+                    ## SEPARATED gate/up tile layout (matches the
-                    # `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
+                    ## `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
-                    # Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
+                    ## Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
-                    # callers default to INTERLEAVE; opt this path out
+                    ## callers default to INTERLEAVE; opt this path out
-                    # unless the user explicitly overrode it.
+                    ## unless the user explicitly overrode it.
-                    envs.SGLANG_USE_AITER_MOE_GU_ITLV.set(False)
+                    # 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

@@ -76,7 +76,7 @@ MI35X_GPT_OSS_MODELS = [
         # to INTERLEAVE, so opt out explicitly here.
         env_vars={
             "SGLANG_USE_AITER": "1",
-            "SGLANG_USE_AITER_MOE_GU_ITLV": "0",
+            "SGLANG_USE_AITER_MOE_GU_ITLV": "1",
         },
     ),
     ModelConfig(
@@ -97,7 +97,7 @@ MI35X_GPT_OSS_MODELS = [
         ],
         env_vars={
             "SGLANG_USE_AITER": "1",
-            "SGLANG_USE_AITER_MOE_GU_ITLV": "0",
+            "SGLANG_USE_AITER_MOE_GU_ITLV": "1",
         },
     ),
 ]