minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Support fused_rms_mxfp4_quant in the prefill stage for DeepSeek-R1-MXFP4 (#14975)

Browse Source
This commit is contained in:
jacky.cheng
2025-12-16 12:03:58 -08:00
committed by GitHub
parent 99401e7b1a
commit 8ac350f335
1 changed files with 11 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+11 -1
View File
@@ -1722,7 +1722,17 @@ class DeepseekV2AttentionMLA(nn.Module):
layer_id=self.layer_id, layer_id=self.layer_id,
return_indices=False, 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: elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
q, _, _, _ = fused_rms_fp8_group_quant( q, _, _, _ = fused_rms_fp8_group_quant(