EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

fix fused_add_rms_norm bug

Browse Source
This commit is contained in:
maxiao
2025-10-27 10:27:57 +08:00
parent b80ae5e9ff
commit f9a026ad2b
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/layernorm.py
View File

@@ -175,7 +175,7 @@ class RMSNorm(CustomOp):
self.weight.data, self.weight.data,
self.variance_epsilon, self.variance_epsilon,
) )
return output, residual_out return x, residual
except TypeError: except TypeError:
fused_add_rms_norm( fused_add_rms_norm(
output, output,
@@ -185,7 +185,8 @@ class RMSNorm(CustomOp):
self.weight.data, self.weight.data,
self.variance_epsilon, self.variance_epsilon,
) )
return x, residual return output, residual_out
out = torch.empty_like(x) out = torch.empty_like(x)
rms_norm(out, x, self.weight.data, self.variance_epsilon) rms_norm(out, x, self.weight.data, self.variance_epsilon)