From 8ac350f335c636991a7f7211983b2545dc582600 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Wed, 17 Dec 2025 04:03:58 +0800 Subject: [PATCH] [AMD] Support fused_rms_mxfp4_quant in the prefill stage for DeepSeek-R1-MXFP4 (#14975) --- python/sglang/srt/models/deepseek_v2.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 403bea8bc..082feacf1 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1722,7 +1722,17 @@ class DeepseekV2AttentionMLA(nn.Module): layer_id=self.layer_id, return_indices=False, ) - + elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: + # MXFP4: fused RMSNorm + quant + q, _, _, _ = fused_rms_mxfp4_quant( + q, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + None, + None, + None, + ) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn: q, _, _, _ = fused_rms_fp8_group_quant(