chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix post_residual_addition more generally (#17286)

Browse Source
This commit is contained in:
Nan Jiang
2026-01-24 07:43:37 +08:00
committed by GitHub
parent 628ab5d57b
commit ad05782160
2 changed files with 22 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -140,6 +140,8 @@ class RMSNorm(MultiPlatformOp):
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
out, _, residual_out = torch_npu.npu_add_rms_norm(
residual, x, self.weight.data, self.variance_epsilon
)
@@ -155,6 +157,8 @@ class RMSNorm(MultiPlatformOp):
if residual is not None:
residual_out = torch.empty_like(x)
output = torch.empty_like(x)
if post_residual_addition is not None:
residual = residual + post_residual_addition
fused_add_rms_norm(
output,
x,
@@ -178,6 +182,8 @@ class RMSNorm(MultiPlatformOp):
if residual is not None:
out = torch.empty_like(x)
residual_out = torch.empty_like(x)
if post_residual_addition is not None:
residual = residual + post_residual_addition
fused_add_rms_norm(
out, x, residual_out, residual, self.weight.data, self.variance_epsilon
)
@@ -244,6 +250,8 @@ class RMSNorm(MultiPlatformOp):
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available:
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(
x, residual, self.weight.data, self.variance_epsilon
)
@@ -263,6 +271,8 @@ class RMSNorm(MultiPlatformOp):
if self.variance_size_override is not None:
return self.forward_native(x, residual, post_residual_addition)
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
return x, residual
out = rmsnorm(x, self.weight.data, self.variance_epsilon)
@@ -284,6 +294,8 @@ class RMSNorm(MultiPlatformOp):
)
if get_tensor_model_parallel_world_size() > 1:
if post_residual_addition is not None:
residual = residual + post_residual_addition
fused_result = flashinfer_allreduce_residual_rmsnorm(
input_tensor=x,
residual=residual,
@@ -389,6 +401,8 @@ class GemmaRMSNorm(MultiPlatformOp):
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
gemma_fused_add_rmsnorm(
x, residual, self.weight.data, self.variance_epsilon
)
@@ -404,6 +418,8 @@ class GemmaRMSNorm(MultiPlatformOp):
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
orig_dtype = x.dtype
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
x = x + residual
residual = x
@@ -430,6 +446,8 @@ class GemmaRMSNorm(MultiPlatformOp):
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available:
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu(
x, residual, self.weight.data, self.variance_epsilon
)
@@ -446,6 +464,8 @@ class GemmaRMSNorm(MultiPlatformOp):
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
x = x + residual
residual = x

2
python/sglang/srt/layers/quantization/modelslim/modelslim.py
View File

@@ -50,6 +50,8 @@ def npu_wrapper_rmsnorm_forward(func):
if not x.is_contiguous():
x = x.contiguous()
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
out, residual_out = add_rmsnorm_bias(
x,
residual,