diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 6504923a3..23c46113a 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -367,6 +367,22 @@ class RMSNorm(MultiPlatformOp): x = x.contiguous().reshape(-1, original_shape[-1]) elif not x.is_contiguous(): x = x.contiguous() + if is_batch_invariant_mode_enabled(): + if ( + residual is not None + or self.cast_x_before_out_mul + or get_global_server_args().rl_on_policy_target == "fsdp" + or (self._fused_pad_kernel is not None and self.x_pad_to_multiple > 0) + ): + return self.forward_native(x, residual, post_residual_addition) + out = rms_norm_batch_invariant( + x, + self.weight.data, + self.variance_epsilon, + ) + if needs_reshape: + out = out.reshape(original_shape) + return out # Fused (add +) rmsnorm + zero-pad path. Triggered when caller # constructed RMSNorm with x_pad_to_multiple > 0. Output last # dim is padded up; residual_out stays at original width. Used @@ -412,6 +428,19 @@ class RMSNorm(MultiPlatformOp): if not _has_vllm_rms_norm: return self.forward_native(x, residual, post_residual_addition) + if is_batch_invariant_mode_enabled(): + if ( + residual is not None + or self.cast_x_before_out_mul + or get_global_server_args().rl_on_policy_target == "fsdp" + ): + return self.forward_native(x, residual, post_residual_addition) + return rms_norm_batch_invariant( + x, + self.weight.data, + self.variance_epsilon, + ) + if not x.is_contiguous(): # NOTE: Remove this if aiter kernel supports discontinuous input x = x.contiguous()