diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 15df851eb..1636ed706 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -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( diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 34b75835d..d7b4cbd1f 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 20b6c88a1..b9ba88fca 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -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 diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index aad79ce23..89871ad57 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -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(