[Fix] Stamp sequence-parallel state on dummy forward batches (#38564)

Co-authored-by: pranjalssh <pranjalssh@fb.com>
2026-09-09 01:03:03 -07:00
parent 13469c16d3
commit 35df2fecde
2 changed files with 37 additions and 0 deletions
@@ -1508,6 +1508,10 @@ class ModelRunner:

    def prepare_dummy_forward_batch(self, forward_batch: ForwardBatch) -> ForwardBatch:
        """Customize a runner-created dummy batch before attention metadata initialization."""
+        # Dummy runs bypass the MLP-sync/scatter passes that stamp real batches.
+        forward_batch.attn_tp_sequence_sharded = self.attn_tp_sequence_sharded(
+            forward_batch._forward_num_tokens()
+        )
        return forward_batch

    def attn_tp_sequence_sharded(self, num_tokens: int) -> bool: