Fix Falcon-H1 gated RMSNorm shape during batched decoding (#50507)

* Fix Falcon-H1 gated RMSNorm shape during batched decoding

* Add a regression test for batched cached decoding with the gated RMSNorm

---------

Co-authored-by: Matt <rocketknight1@gmail.com>

Author: taking-lying-flat

src/transformers/models/falcon_h1/modeling_falcon_h1.py

@@ -121,6 +121,7 @@ def __init__(self, hidden_size, eps=1e-6, n_groups=1, norm_before_gate=True):
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
+        input_ndim = hidden_states.ndim
 
         if not self.norm_before_gate and gate is not None:
             hidden_states = hidden_states * F.silu(gate.to(torch.float32))
@@ -140,7 +141,7 @@ def forward(self, hidden_states, gate=None):
         hidden_states = self.weight.view(self.n_groups, int(dim // self.n_groups)) * hidden_states
         hidden_states = hidden_states.view(batch_size, seq_len, dim)
 
-        if seq_len == 1:
+        if input_ndim == 2:
             hidden_states = hidden_states.squeeze(1)
 
         if self.norm_before_gate and gate is not None:

src/transformers/models/falcon_h1/modular_falcon_h1.py

@@ -276,6 +276,33 @@ def create_and_check_mamba_chunked_prefill(self, config, input_ids, *args, devic
             msg=f"Max diff: {(ref_first - under_test_first).abs().max().item():.6f}",
         )
 
+    def create_and_check_gated_rms_norm_batched_cached_decoding(self, config, input_ids, *args):
+        """
+        With `mamba_rms_norm=True`, a batched single-token cached step must give the same hidden states
+        as the uncached forward pass, and must keep its `[batch_size, 1, hidden_size]` shape.
+        """
+        input_ids = input_ids.to(torch_device)
+        # The gate is applied in a different place depending on `mamba_norm_before_gate`, so we check both
+        for norm_before_gate in (True, False):
+            config.mamba_rms_norm = True
+            config.mamba_norm_before_gate = norm_before_gate
+            model = FalconH1Model(config=config)
+            model.to(torch_device)
+            model.eval()
+
+            with torch.no_grad():
+                full_out = model(input_ids, use_cache=False).last_hidden_state
+                prefill = model(input_ids[:, :-1], use_cache=True)
+                step_out = model(
+                    input_ids[:, -1:], past_key_values=prefill.past_key_values, use_cache=True
+                ).last_hidden_state
+
+            self.parent.assertEqual(step_out.shape, (input_ids.shape[0], 1, config.hidden_size))
+            self.parent.assertTrue(
+                torch.allclose(full_out[:, -1:], step_out, atol=1e-4, rtol=1e-4),
+                msg=f"Max diff: {(full_out[:, -1:] - step_out).abs().max().item():.6f}",
+            )
+
     def create_and_check_kwargs_reach_mamba2_mixer(self, config, input_ids, *args):
         """
         Kernel kwargs given to the model must reach the Mamba2 mixer, which splats them into
@@ -357,6 +384,10 @@ def test_mamba2_chunked_prefill_cpu(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device="cpu")
 
+    def test_gated_rms_norm_batched_cached_decoding(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_gated_rms_norm_batched_cached_decoding(*config_and_inputs)
+
     def test_kwargs_reach_mamba2_mixer(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_kwargs_reach_mamba2_mixer(*config_and_inputs)