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

Refactor fused_add_rmsnorm import logic (#10207)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2025-09-09 15:23:58 +08:00
committed by GitHub
parent 94fb4e9e54
commit f5f6b3b4b5
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -26,6 +26,7 @@ from sglang.srt.utils import (
get_bool_env_var,
is_cpu,
is_cuda,
is_flashinfer_available,
is_hip,
is_npu,
is_xpu,
@@ -33,6 +34,7 @@ from sglang.srt.utils import (
)
_is_cuda = is_cuda()
_is_flashinfer_available = is_flashinfer_available()
_is_hip = is_hip()
_is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -41,7 +43,10 @@ _is_cpu = is_cpu()
_is_xpu = is_xpu()
if _is_cuda:
from flashinfer.norm import fused_add_rmsnorm as flashinfer_fused_add_rmsnorm
if _is_flashinfer_available:
from flashinfer.norm import fused_add_rmsnorm
else:
from sgl_kernel import fused_add_rmsnorm
from sgl_kernel import gemma_fused_add_rmsnorm, gemma_rmsnorm, rmsnorm
if _use_aiter:
@@ -84,9 +89,7 @@ class RMSNorm(CustomOp):
if self.variance_size_override is not None:
return self.forward_native(x, residual)
if residual is not None:
flashinfer_fused_add_rmsnorm(
x, residual, self.weight.data, self.variance_epsilon
)
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
return x, residual
out = rmsnorm(x, self.weight.data, self.variance_epsilon)
return out