fix fused_add_rms_norm bug

f9a026ad · maxiao · b80ae5e9 · f9a026ad
Commit f9a026ad authored Oct 27, 2025 by maxiao
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

python/sglang/srt/layers/layernorm.py python/sglang/srt/layers/layernorm.py +3 -2

No files found.
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -175,7 +175,7 @@ class RMSNorm(CustomOp):
                    self.weight.data,
                    self.variance_epsilon,
                )
-                return output, residual_out
+                return x, residual
            except TypeError:
                fused_add_rms_norm(
                    output,
@@ -185,7 +185,8 @@ class RMSNorm(CustomOp):
                    self.weight.data,
                    self.variance_epsilon,
                )
-                return x, residual
+                return output, residual_out
+                

        out = torch.empty_like(x)
        rms_norm(out, x, self.weight.data, self.variance_epsilon)