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

fix: fix regression and unclear pattern (#16561)

Browse Source
This commit is contained in:
Nan Jiang
2026-01-17 15:21:42 +08:00
committed by GitHub
parent 8ce64aa155
commit dd99f818e0
4 changed files with 46 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/communicator.py
View File

@@ -371,10 +371,13 @@ class LayerCommunicator:
residual: torch.Tensor,
forward_batch: ForwardBatch,
captured_last_layer_outputs: Optional[List[torch.Tensor]] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
):
hidden_states, residual = self.prepare_attn(
hidden_states, residual, forward_batch, **kwargs
hidden_states,
residual,
forward_batch,
post_residual_addition=post_residual_addition,
)
if captured_last_layer_outputs is not None:
gathered_last_layer_output = self._communicate_simple_fn(
@@ -394,7 +397,7 @@ class LayerCommunicator:
residual: torch.Tensor,
forward_batch: ForwardBatch,
quant_format: str = "",
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
):
if get_attn_tp_context().input_scattered:
hidden_states, residual = self._tp_reduce_scatter(
@@ -444,7 +447,7 @@ class LayerCommunicator:
)
else:
hidden_states = self.input_layernorm(hidden_states, **kwargs)
hidden_states = self.input_layernorm(hidden_states)
else:
if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
@@ -478,7 +481,7 @@ class LayerCommunicator:
hidden_states, residual = self.input_layernorm(
hidden_states,
residual,
**kwargs,
post_residual_addition,
)
hidden_states = self._communicate_simple_fn(

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

@@ -104,18 +104,18 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if x.numel() == 0:
return x
if self.variance_size_override is not None:
return self.forward_native(x, residual, **kwargs)
return self.forward_native(x, residual, post_residual_addition)
if is_batch_invariant_mode_enabled():
if (
residual is not None
or get_global_server_args().rl_on_policy_target == "fsdp"
):
return self.forward_native(x, residual, **kwargs)
return self.forward_native(x, residual, post_residual_addition)
return rms_norm_batch_invariant(
x,
self.weight.data,
@@ -126,7 +126,6 @@ class RMSNorm(MultiPlatformOp):
# but right now we can only have hidden_states+(residual+post_residual_addition).
# (hidden_states+residual)+post_residual_addition != hidden_states+(residual+post_residual_addition),
# we probably need to add another parameter to fused_add_rmsnorm
post_residual_addition = kwargs.get("post_residual_addition")
if post_residual_addition is not None:
residual = residual + post_residual_addition
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
@@ -138,7 +137,7 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
out, _, residual_out = torch_npu.npu_add_rms_norm(
@@ -151,7 +150,7 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
residual_out = torch.empty_like(x)
@@ -171,7 +170,7 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if not x.is_contiguous():
# NOTE: Remove this if aiter kernel supports discontinuous input
@@ -191,23 +190,16 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if not x.is_contiguous():
x = x.contiguous()
orig_dtype = self.override_orig_dtype or x.dtype
post_residual_addition = kwargs.get("post_residual_addition")
x = x.to(torch.float32)
if residual is not None:
x = (
x
+ residual.to(torch.float32)
+ (
post_residual_addition.to(torch.float32)
if post_residual_addition is not None
else 0.0
)
)
x = x + residual.to(torch.float32)
if post_residual_addition is not None:
x = x + post_residual_addition.to(torch.float32)
if self.fp32_residual:
residual = x.clone()
else:
@@ -248,7 +240,7 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available:
if residual is not None:
@@ -260,16 +252,16 @@ class RMSNorm(MultiPlatformOp):
x, self.weight.data, self.variance_epsilon
)
else:
return self.forward_native(x, residual, **kwargs)
return self.forward_native(x, residual, post_residual_addition)
def forward_xpu(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if self.variance_size_override is not None:
return self.forward_native(x, residual, **kwargs)
return self.forward_native(x, residual, post_residual_addition)
if residual is not None:
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
return x, residual
@@ -280,6 +272,7 @@ class RMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
"""
Forward method with allreduce fusion, prioritizing flashinfer fused operations
@@ -300,7 +293,7 @@ class RMSNorm(MultiPlatformOp):
if fused_result[0] is not None:
return fused_result
return self.forward(x, residual)
return self.forward(x, residual, post_residual_addition)
class LayerNorm(MultiPlatformOp):
@@ -325,7 +318,6 @@ class LayerNorm(MultiPlatformOp):
def forward_cuda(
self,
x: torch.Tensor,
**kwargs,
) -> torch.Tensor:
if (
_flashinfer_layernorm_available
@@ -334,12 +326,11 @@ class LayerNorm(MultiPlatformOp):
):
return layernorm(x, self.weight, self.bias, self.variance_epsilon)
else:
return self.forward_native(x, **kwargs)
return self.forward_native(x)
def forward_native(
self,
x: torch.Tensor,
**kwargs,
) -> torch.Tensor:
weight = self.weight if self.elementwise_affine else None
bias = self.bias if self.use_bias else None
@@ -356,28 +347,25 @@ class LayerNorm(MultiPlatformOp):
def forward_hip(
self,
x: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.forward_native(x, **kwargs)
return self.forward_native(x)
def forward_npu(
self,
x: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.forward_native(x, **kwargs)
return self.forward_native(x)
def forward_cpu(
self,
x: torch.Tensor,
**kwargs,
) -> torch.Tensor:
if _is_cpu_amx_available:
return torch.ops.sgl_kernel.layernorm_cpu(
x, self.weight.data, self.variance_epsilon
)
else:
return self.forward_native(x, **kwargs)
return self.forward_native(x)
class GemmaRMSNorm(MultiPlatformOp):
@@ -398,7 +386,7 @@ class GemmaRMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
gemma_fused_add_rmsnorm(
@@ -412,7 +400,7 @@ class GemmaRMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
orig_dtype = x.dtype
if residual is not None:
@@ -430,15 +418,15 @@ class GemmaRMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
return self._forward_impl(x, residual, **kwargs)
return self._forward_impl(x, residual, post_residual_addition)
def forward_cpu(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available:
if residual is not None:
@@ -449,13 +437,13 @@ class GemmaRMSNorm(MultiPlatformOp):
return torch.ops.sgl_kernel.gemma_rmsnorm_cpu(
x, self.weight.data, self.variance_epsilon
)
return self.forward_native(x, residual, **kwargs)
return self.forward_native(x, residual, post_residual_addition)
def forward_npu(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
x = x + residual
@@ -468,9 +456,9 @@ class GemmaRMSNorm(MultiPlatformOp):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
return self._forward_impl(x, residual, **kwargs)
return self._forward_impl(x, residual, post_residual_addition)
class Gemma3RMSNorm(MultiPlatformOp):
@@ -483,22 +471,22 @@ class Gemma3RMSNorm(MultiPlatformOp):
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward_native(self, x, **kwargs):
def forward_native(self, x):
output = self._norm(x.float())
# Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
# See https://github.com/huggingface/transformers/pull/29402
output = output * (1.0 + self.weight.float())
return output.type_as(x)
def forward_cpu(self, x, **kwargs):
def forward_cpu(self, x):
if _is_cpu_amx_available and x.stride(-1) == 1:
return torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, self.weight, self.eps)
return self.forward_native(x, **kwargs)
return self.forward_native(x)
def forward_cuda(self, x, **kwargs):
return self.forward_native(x, **kwargs)
def forward_cuda(self, x):
return self.forward_native(x)
def forward_npu(self, x, **kwargs):
def forward_npu(self, x):
output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps)
return output

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

@@ -43,6 +43,7 @@ def npu_wrapper_rmsnorm_forward(func):
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias

4
python/sglang/srt/models/qwen3.py
View File

@@ -276,14 +276,14 @@ class Qwen3DecoderLayer(nn.Module):
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor],
**kwargs,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# Self Attention
hidden_states, residual = self.layer_communicator.prepare_attn(
hidden_states,
residual,
forward_batch,
**kwargs,
post_residual_addition=post_residual_addition,
)
if hidden_states.shape[0] != 0:
hidden_states = self.self_attn(