fix: fix regression and unclear pattern (#16561)

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(

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
 

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
 

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(