minke.yu/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-16 23: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
python/sglang/srt/layers/communicator.py
+8 -5
View File
@@ -371,10 +371,13 @@ class LayerCommunicator:
residual: torch.Tensor, residual: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, 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 = 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: if captured_last_layer_outputs is not None:
gathered_last_layer_output = self._communicate_simple_fn( gathered_last_layer_output = self._communicate_simple_fn(
@@ -394,7 +397,7 @@ class LayerCommunicator:
residual: torch.Tensor, residual: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
quant_format: str = "", quant_format: str = "",
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
): ):
if get_attn_tp_context().input_scattered: if get_attn_tp_context().input_scattered:
hidden_states, residual = self._tp_reduce_scatter( hidden_states, residual = self._tp_reduce_scatter(
@@ -444,7 +447,7 @@ class LayerCommunicator:
) )
else: else:
hidden_states = self.input_layernorm(hidden_states, **kwargs) hidden_states = self.input_layernorm(hidden_states)
else: else:
if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format): 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 = self.input_layernorm(
hidden_states, hidden_states,
residual, residual,
**kwargs, post_residual_addition,
) )
hidden_states = self._communicate_simple_fn( hidden_states = self._communicate_simple_fn(
python/sglang/srt/layers/layernorm.py
+35 -47
View File
@@ -104,18 +104,18 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if x.numel() == 0: if x.numel() == 0:
return x return x
if self.variance_size_override is not None: 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 is_batch_invariant_mode_enabled():
if ( if (
residual is not None residual is not None
or get_global_server_args().rl_on_policy_target == "fsdp" 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( return rms_norm_batch_invariant(
x, x,
self.weight.data, self.weight.data,
@@ -126,7 +126,6 @@ class RMSNorm(MultiPlatformOp):
# but right now we can only have hidden_states+(residual+post_residual_addition). # 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), # (hidden_states+residual)+post_residual_addition != hidden_states+(residual+post_residual_addition),
# we probably need to add another parameter to fused_add_rmsnorm # 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: if post_residual_addition is not None:
residual = residual + post_residual_addition residual = residual + post_residual_addition
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon) fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
@@ -138,7 +137,7 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None: if residual is not None:
out, _, residual_out = torch_npu.npu_add_rms_norm( out, _, residual_out = torch_npu.npu_add_rms_norm(
@@ -151,7 +150,7 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None: if residual is not None:
residual_out = torch.empty_like(x) residual_out = torch.empty_like(x)
@@ -171,7 +170,7 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if not x.is_contiguous(): if not x.is_contiguous():
# NOTE: Remove this if aiter kernel supports discontinuous input # NOTE: Remove this if aiter kernel supports discontinuous input
@@ -191,23 +190,16 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if not x.is_contiguous(): if not x.is_contiguous():
x = x.contiguous() x = x.contiguous()
orig_dtype = self.override_orig_dtype or x.dtype orig_dtype = self.override_orig_dtype or x.dtype
post_residual_addition = kwargs.get("post_residual_addition")
x = x.to(torch.float32) x = x.to(torch.float32)
if residual is not None: if residual is not None:
x = ( x = x + residual.to(torch.float32)
x if post_residual_addition is not None:
+ residual.to(torch.float32) x = x + post_residual_addition.to(torch.float32)
+ (
post_residual_addition.to(torch.float32)
if post_residual_addition is not None
else 0.0
)
)
if self.fp32_residual: if self.fp32_residual:
residual = x.clone() residual = x.clone()
else: else:
@@ -248,7 +240,7 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available: if _is_cpu_amx_available:
if residual is not None: if residual is not None:
@@ -260,16 +252,16 @@ class RMSNorm(MultiPlatformOp):
x, self.weight.data, self.variance_epsilon x, self.weight.data, self.variance_epsilon
) )
else: else:
return self.forward_native(x, residual, **kwargs) return self.forward_native(x, residual, post_residual_addition)
def forward_xpu( def forward_xpu(
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if self.variance_size_override is not None: 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: if residual is not None:
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 return x, residual
@@ -280,6 +272,7 @@ class RMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
""" """
Forward method with allreduce fusion, prioritizing flashinfer fused operations Forward method with allreduce fusion, prioritizing flashinfer fused operations
@@ -300,7 +293,7 @@ class RMSNorm(MultiPlatformOp):
if fused_result[0] is not None: if fused_result[0] is not None:
return fused_result return fused_result
return self.forward(x, residual) return self.forward(x, residual, post_residual_addition)
class LayerNorm(MultiPlatformOp): class LayerNorm(MultiPlatformOp):
@@ -325,7 +318,6 @@ class LayerNorm(MultiPlatformOp):
def forward_cuda( def forward_cuda(
self, self,
x: torch.Tensor, x: torch.Tensor,
**kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
if ( if (
_flashinfer_layernorm_available _flashinfer_layernorm_available
@@ -334,12 +326,11 @@ class LayerNorm(MultiPlatformOp):
): ):
return layernorm(x, self.weight, self.bias, self.variance_epsilon) return layernorm(x, self.weight, self.bias, self.variance_epsilon)
else: else:
return self.forward_native(x, **kwargs) return self.forward_native(x)
def forward_native( def forward_native(
self, self,
x: torch.Tensor, x: torch.Tensor,
**kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
weight = self.weight if self.elementwise_affine else None weight = self.weight if self.elementwise_affine else None
bias = self.bias if self.use_bias else None bias = self.bias if self.use_bias else None
@@ -356,28 +347,25 @@ class LayerNorm(MultiPlatformOp):
def forward_hip( def forward_hip(
self, self,
x: torch.Tensor, x: torch.Tensor,
**kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
return self.forward_native(x, **kwargs) return self.forward_native(x)
def forward_npu( def forward_npu(
self, self,
x: torch.Tensor, x: torch.Tensor,
**kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
return self.forward_native(x, **kwargs) return self.forward_native(x)
def forward_cpu( def forward_cpu(
self, self,
x: torch.Tensor, x: torch.Tensor,
**kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
if _is_cpu_amx_available: if _is_cpu_amx_available:
return torch.ops.sgl_kernel.layernorm_cpu( return torch.ops.sgl_kernel.layernorm_cpu(
x, self.weight.data, self.variance_epsilon x, self.weight.data, self.variance_epsilon
) )
else: else:
return self.forward_native(x, **kwargs) return self.forward_native(x)
class GemmaRMSNorm(MultiPlatformOp): class GemmaRMSNorm(MultiPlatformOp):
@@ -398,7 +386,7 @@ class GemmaRMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None: if residual is not None:
gemma_fused_add_rmsnorm( gemma_fused_add_rmsnorm(
@@ -412,7 +400,7 @@ class GemmaRMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
orig_dtype = x.dtype orig_dtype = x.dtype
if residual is not None: if residual is not None:
@@ -430,15 +418,15 @@ class GemmaRMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> 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( def forward_cpu(
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available: if _is_cpu_amx_available:
if residual is not None: if residual is not None:
@@ -449,13 +437,13 @@ class GemmaRMSNorm(MultiPlatformOp):
return torch.ops.sgl_kernel.gemma_rmsnorm_cpu( return torch.ops.sgl_kernel.gemma_rmsnorm_cpu(
x, self.weight.data, self.variance_epsilon 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( def forward_npu(
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None: if residual is not None:
x = x + residual x = x + residual
@@ -468,9 +456,9 @@ class GemmaRMSNorm(MultiPlatformOp):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> 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): class Gemma3RMSNorm(MultiPlatformOp):
@@ -483,22 +471,22 @@ class Gemma3RMSNorm(MultiPlatformOp):
def _norm(self, x): def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) 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()) output = self._norm(x.float())
# Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
# See https://github.com/huggingface/transformers/pull/29402 # See https://github.com/huggingface/transformers/pull/29402
output = output * (1.0 + self.weight.float()) output = output * (1.0 + self.weight.float())
return output.type_as(x) 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: if _is_cpu_amx_available and x.stride(-1) == 1:
return torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, self.weight, self.eps) 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): def forward_cuda(self, x):
return self.forward_native(x, **kwargs) 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) output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps)
return output return output
python/sglang/srt/layers/quantization/modelslim/modelslim.py
+1
View File
@@ -43,6 +43,7 @@ def npu_wrapper_rmsnorm_forward(func):
self, self,
x: torch.Tensor, x: torch.Tensor,
residual: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias
python/sglang/srt/models/qwen3.py
+2 -2
View File
@@ -276,14 +276,14 @@ class Qwen3DecoderLayer(nn.Module):
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
residual: Optional[torch.Tensor], residual: Optional[torch.Tensor],
**kwargs, post_residual_addition: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
# Self Attention # Self Attention
hidden_states, residual = self.layer_communicator.prepare_attn( hidden_states, residual = self.layer_communicator.prepare_attn(
hidden_states, hidden_states,
residual, residual,
forward_batch, forward_batch,
**kwargs, post_residual_addition=post_residual_addition,
) )
if hidden_states.shape[0] != 0: if hidden_states.shape[0] != 0:
hidden_states = self.self_attn( hidden_states = self.self_attn(