From aa55657e9e5d0a858418c5e6632c9b4d919a94d5 Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Sat, 6 Jun 2026 11:18:11 +0800 Subject: [PATCH] [AMD][WA] force to use gate_mode interleaved to fix tp2/tp4/tp8 acc issue (#27201) Co-authored-by: wunhuang --- .../srt/layers/attention/aiter_backend.py | 13 ++- .../sglang/srt/layers/quantization/mxfp4.py | 97 +++++++++++-------- python/sglang/srt/server_args.py | 14 +-- .../accuracy/mi35x/test_gpt_oss_eval_mi35x.py | 4 +- 4 files changed, 73 insertions(+), 55 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 393fd25a8..ccf0c796a 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 6337b39f6..a29f3faf9 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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 - # `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, - ) + if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get(): + 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, + ) + else: + 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, + ) + + # 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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 82a3d31a2..85cc92c0b 100644 --- a/python/sglang/srt/server_args.py +++ b/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 - # 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) + ## 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 diff --git a/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py index 3c96348a3..d371bb3a5 100644 --- a/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py +++ b/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", }, ), ]