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

Co-authored-by: pranjalssh <pranjalssh@fb.com>

python/sglang/srt/model_executor/model_runner.py

@@ -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:

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -107,6 +107,39 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         self.assertTrue(runner.attn_tp_sequence_sharded(num_tokens=4))
         mock_require_gathered_buffer.assert_called_once_with()
 
+    def test_dummy_batch_sharding_tracks_forward_token_width(self):
+        """Warmup must use the runner's SP policy, including embedding widths."""
+
+        class TokenGatedRunner(ModelRunner):
+            def attn_tp_sequence_sharded(self, num_tokens):
+                return num_tokens == 4
+
+        runner = TokenGatedRunner.__new__(TokenGatedRunner)
+        for num_ids, num_embeds, expected in (
+            (4, None, True),
+            (5, None, False),
+            (5, 4, True),
+            (4, 5, False),
+            (0, None, False),
+        ):
+            with self.subTest(num_ids=num_ids, num_embeds=num_embeds):
+                batch = ForwardBatch(
+                    forward_mode=ForwardMode.EXTEND,
+                    batch_size=1,
+                    input_ids=torch.arange(num_ids),
+                    input_embeds=(
+                        torch.empty(num_embeds, 8) if num_embeds is not None else None
+                    ),
+                    req_pool_indices=torch.tensor([0]),
+                    seq_lens=torch.tensor([num_ids]),
+                    out_cache_loc=torch.arange(num_ids),
+                    seq_lens_sum=num_ids,
+                    attn_tp_sequence_sharded=not expected,
+                )
+
+                self.assertIs(runner.prepare_dummy_forward_batch(batch), batch)
+                self.assertEqual(batch.attn_tp_sequence_sharded, expected)
+
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
         prefill_runner = object()