From 35df2fecde7acec19c59d3fa1186ffc49f02deb6 Mon Sep 17 00:00:00 2001 From: Pranjal Shankhdhar Date: Wed, 9 Sep 2026 01:03:03 -0700 Subject: [PATCH] [Fix] Stamp sequence-parallel state on dummy forward batches (#38564) Co-authored-by: pranjalssh --- .../sglang/srt/model_executor/model_runner.py | 4 +++ .../test_prefill_cuda_graph_runner.py | 33 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6c22e5517..763f6bb3b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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: diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index 613656dc1..f00c9aec6 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/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()