Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
change
sglang
Commits
8d114f25
Unverified
Commit
8d114f25
authored
Sep 06, 2025
by
sogalin
Committed by
GitHub
Sep 05, 2025
Browse files
Fix RMSNorm API CALL mismatch issue. (#10032)
Co-authored-by:
Hubert Lu
<
Hubert.Lu@amd.com
>
parent
0e78c63c
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
19 additions
and
2 deletions
+19
-2
python/sglang/srt/layers/layernorm.py
python/sglang/srt/layers/layernorm.py
+19
-2
No files found.
python/sglang/srt/layers/layernorm.py
View file @
8d114f25
...
...
@@ -18,6 +18,7 @@ from typing import Optional, Tuple, Union
import
torch
import
torch.nn
as
nn
from
packaging.version
import
Version
from
sglang.srt.custom_op
import
CustomOp
from
sglang.srt.utils
import
(
...
...
@@ -49,8 +50,11 @@ if _use_aiter:
from
aiter
import
rmsnorm2d_fwd
as
rms_norm
from
aiter
import
rmsnorm2d_fwd_with_add
as
fused_add_rms_norm
elif
_is_hip
:
import
vllm
from
vllm._custom_ops
import
fused_add_rms_norm
,
rms_norm
_vllm_version
=
Version
(
vllm
.
__version__
)
logger
=
logging
.
getLogger
(
__name__
)
if
_is_npu
:
...
...
@@ -127,8 +131,21 @@ class RMSNorm(CustomOp):
# NOTE: Remove this if aiter kernel supports discontinuous input
x
=
x
.
contiguous
()
if
residual
is
not
None
:
fused_add_rms_norm
(
x
,
residual
,
self
.
weight
.
data
,
self
.
variance_epsilon
)
return
x
,
residual
if
_vllm_version
<
Version
(
"0.9"
):
fused_add_rms_norm
(
x
,
residual
,
self
.
weight
.
data
,
self
.
variance_epsilon
)
return
x
,
residual
else
:
residual_out
=
torch
.
empty_like
(
x
)
output
=
torch
.
empty_like
(
x
)
fused_add_rms_norm
(
output
,
x
,
residual_out
,
residual
,
self
.
weight
.
data
,
self
.
variance_epsilon
,
)
return
output
,
residual_out
out
=
torch
.
empty_like
(
x
)
rms_norm
(
out
,
x
,
self
.
weight
.
data
,
self
.
variance_epsilon
)
return
out
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment