diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9ed722192..4dc89fdc0 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index ed402a283..1862d6484 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 6ac793d02..6337b39f6 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0e4b628b4..05c6637da 100644 --- a/python/sglang/srt/server_args.py +++ b/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 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 4c2f8861e..c98ef4852 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 @@ -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", + }, ), ]