From 60d4bd4c70b23d9a5689b0be67a373f019f14624 Mon Sep 17 00:00:00 2001 From: billishyahao Date: Sat, 13 Jun 2026 15:35:41 +0800 Subject: [PATCH] [AMD] fix moriep quant kernel not implemented issue (#27855) --- .../sglang/srt/layers/moe/moe_runner/aiter.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 7c32cf0b8..3e5c9722a 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -18,7 +18,7 @@ from sglang.srt.layers.moe.moe_runner.base import ( register_pre_permute, ) from sglang.srt.layers.moe.utils import MoeRunnerBackend -from sglang.srt.utils import get_int_env_var +from sglang.srt.utils import get_bool_env_var, get_int_env_var if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher.base import CombineInput @@ -327,6 +327,13 @@ def _pre_permute_deepep_to_aiter( is_w4a4 = weight_quant == AiterQuantType.PER_1X32 is_fp4_dispatch = hidden_states.dtype == torch.float4_e2m1fn_x2 + # AITER fused_moe Clamped-SwiGLU is dispatched with + # gate_mode=INTERLEAVE, for which AITER picks a bf16/fp8 `q_dtype_a` + # Refer to https://github.com/ROCm/aiter/blob/a2617c366dc7271a1662ecda2023d19f6ccefcec/aiter/fused_moe.py#L406-L412 + swiglu_interleave = quant_info.swiglu_limit > 0 and get_bool_env_var( + "SGLANG_USE_AITER_MOE_GU_ITLV", "true" + ) + if is_w4a4 and a1_scale is not None and not is_fp4_dispatch: # W4A4 weights with FP8 dispatch: dequant FP8->BF16 first; the # FP4 per_1x32 path needs BF16 input. @@ -334,6 +341,14 @@ def _pre_permute_deepep_to_aiter( hidden_states, a1_scale, num_local_tokens, output_dtype ) a1_scale = None + elif is_w4a4 and is_fp4_dispatch and a1_scale is not None and swiglu_interleave: + # W4A4 weights + FP4 dispatch on the clamped-SwiGLU/INTERLEAVE + # path: AITER expects a bf16/fp8 activation here, not fp4x2. + # Dequant FP4->BF16 and let fused_moe re-quantize internally. + hidden_states = upscale_mxfp4( + hidden_states, a1_scale, num_local_tokens, output_dtype + ) + a1_scale = None elif is_fp8_quant and is_fp4_dispatch and a1_scale is not None: # FP8 weights + FP4 dispatch: no kernel for the fp4x2/fp8 pair; # dequant FP4->BF16 and let fused_moe re-quantize to FP8.