From 22c08a9beee09fbecaeff7e62e4646bb3e67eae4 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 12 Jul 2026 17:51:53 -0700 Subject: [PATCH] Preserve RMSNorm shape in batch-invariant mode (#30956) Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com> --- python/sglang/srt/layers/layernorm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 2c95c01ec..8b0ca5141 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -287,11 +287,14 @@ class RMSNorm(MultiPlatformOp): or get_server_args().rl_on_policy_target == "fsdp" ): return self.forward_native(x, residual, post_residual_addition) - return rms_norm_batch_invariant( + out = rms_norm_batch_invariant( x, self.weight.data, self.variance_epsilon, ) + if needs_reshape: + out = out.reshape(original_shape) + return out if self.cast_x_before_out_mul and residual is None: # Use HF-semantics kernel (cast to dtype before weight multiply). if (