Pass per-forward overrides to ForwardBatch.init_new as explicit arguments (#30670)

2026-07-15 14:25:59 +08:00
parent 861d97d24d
commit e77d95c3d5
16 changed files with 147 additions and 55 deletions
@@ -120,7 +120,11 @@ class TestForwardSplitPrefill(CustomTestCase):
            batch.forward_mode = ForwardMode.SPLIT_PREFILL

        # Create forward batch
-        forward_batch = ForwardBatch.init_new(batch, self.model_runner)
+        forward_batch = ForwardBatch.init_new(
+            batch,
+            self.model_runner,
+            return_hidden_states_before_norm=False,
+        )

        return forward_batch