Preserve RMSNorm shape in batch-invariant mode (#30956)

Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
2026-07-12 17:51:53 -07:00
parent 7da30f4e55
commit 22c08a9bee
1 changed files with 4 additions and 1 deletions
@@ -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 (