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

Fix the run-time error when calling fused_rms_mxfp4_quant that change return output number (#12803)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
2025-11-09 21:48:36 -08:00
committed by GitHub
co-authored by wunhuang
parent 9cfe78dd30
commit 90401cf7d2
3 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/rocm.Dockerfile
+2 -2
View File
@@ -31,7 +31,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0" ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1" ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1" ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT="v0.1.6.post1" ENV AITER_COMMIT="v0.1.6.post3"
ENV NO_DEPS_FLAG="" ENV NO_DEPS_FLAG=""
# =============================== # ===============================
@@ -42,7 +42,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0" ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1" ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1" ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT="v0.1.6.post1" ENV AITER_COMMIT="v0.1.6.post3"
ENV NO_DEPS_FLAG="" ENV NO_DEPS_FLAG=""
# =============================== # ===============================
python/sglang/srt/layers/communicator.py
+2 -2
View File
@@ -268,7 +268,7 @@ class LayerCommunicator:
residual = hidden_states residual = hidden_states
if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format): if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
hidden_states = fused_rms_mxfp4_quant( hidden_states, *_, _ = fused_rms_mxfp4_quant(
hidden_states, hidden_states,
self.input_layernorm.weight, self.input_layernorm.weight,
self.input_layernorm.variance_epsilon, self.input_layernorm.variance_epsilon,
@@ -281,7 +281,7 @@ class LayerCommunicator:
hidden_states = self.input_layernorm(hidden_states) hidden_states = self.input_layernorm(hidden_states)
else: else:
if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format): if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
hidden_states, residual = fused_rms_mxfp4_quant( hidden_states, *_, residual = fused_rms_mxfp4_quant(
hidden_states, hidden_states,
self.input_layernorm.weight, self.input_layernorm.weight,
self.input_layernorm.variance_epsilon, self.input_layernorm.variance_epsilon,
python/sglang/srt/models/deepseek_v2.py
+2 -2
View File
@@ -1608,7 +1608,7 @@ class DeepseekV2AttentionMLA(nn.Module):
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
else: else:
if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
q, k_nope = fused_rms_mxfp4_quant( q, k_nope, *_ = fused_rms_mxfp4_quant(
q, q,
self.q_a_layernorm.weight, self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon, self.q_a_layernorm.variance_epsilon,
@@ -1937,7 +1937,7 @@ class DeepseekV2AttentionMLA(nn.Module):
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
else: else:
if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
q, k_nope = fused_rms_mxfp4_quant( q, k_nope, *_ = fused_rms_mxfp4_quant(
q, q,
self.q_a_layernorm.weight, self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon, self.q_a_layernorm.variance_epsilon,