Multiple flexibility fixes for DP attention (#33537)

python/sglang/srt/model_executor/forward_batch_info.py

@@ -702,6 +702,39 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         if skip_attn_backend_init:
             self.mark_forward_metadata_ready()
 
+    def init_mlp_sync_metadata(
+        self, batch: ScheduleBatch, device: Union[str, torch.device]
+    ) -> None:
+        """Populate per-rank token counts for DP-attention MLP synchronization."""
+        if batch.global_num_tokens is None:
+            return
+
+        assert batch.global_num_tokens_for_logprob is not None
+        if self.spec_info is not None:
+            from sglang.srt.speculative.spec_info import spec_scale_global_num_tokens
+
+            global_num_tokens, global_num_tokens_for_logprob = (
+                spec_scale_global_num_tokens(
+                    self.spec_info,
+                    batch.global_num_tokens,
+                    batch.global_num_tokens_for_logprob,
+                )
+            )
+        else:
+            global_num_tokens = batch.global_num_tokens
+            global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob
+
+        self.original_global_num_tokens_cpu = batch.global_num_tokens
+        self.global_num_tokens_cpu = global_num_tokens
+        self.global_num_tokens_gpu = torch.tensor(
+            global_num_tokens, dtype=torch.int64
+        ).to(device, non_blocking=True)
+        self.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob
+        self.global_num_tokens_for_logprob_gpu = torch.tensor(
+            global_num_tokens_for_logprob, dtype=torch.int64
+        ).to(device, non_blocking=True)
+        self.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
+
     @classmethod
     def init_new(
         cls,
@@ -841,37 +874,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             )
         ret.num_token_non_padded_cpu = num_tokens
 
-        # For MLP sync
+        ret.init_mlp_sync_metadata(batch, device)
-        if batch.global_num_tokens is not None:
-            assert batch.global_num_tokens_for_logprob is not None
-
-            # process global_num_tokens and global_num_tokens_for_logprob
-            if batch.spec_info is not None:
-                from sglang.srt.speculative.spec_info import (
-                    spec_scale_global_num_tokens,
-                )
-
-                global_num_tokens, global_num_tokens_for_logprob = (
-                    spec_scale_global_num_tokens(
-                        batch.spec_info,
-                        batch.global_num_tokens,
-                        batch.global_num_tokens_for_logprob,
-                    )
-                )
-            else:
-                global_num_tokens = batch.global_num_tokens
-                global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob
-
-            ret.original_global_num_tokens_cpu = batch.global_num_tokens
-            ret.global_num_tokens_cpu = global_num_tokens
-            ret.global_num_tokens_gpu = torch.tensor(
-                global_num_tokens, dtype=torch.int64
-            ).to(device, non_blocking=True)
-
-            ret.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob
-            ret.global_num_tokens_for_logprob_gpu = torch.tensor(
-                global_num_tokens_for_logprob, dtype=torch.int64
-            ).to(device, non_blocking=True)
 
         if ret.forward_mode.is_idle():
             ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
@@ -1520,6 +1523,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                 spec_info.num_accept_tokens = self._pad_tensor_to_size(
                     spec_info.num_accept_tokens, bs
                 )
+            if spec_info.hidden_states is not None:
                 spec_info.hidden_states = self._pad_tensor_to_size(
                     spec_info.hidden_states, num_tokens
                 )
@@ -1569,12 +1573,15 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                     ]
                 logits_output.hidden_states = logits_output.hidden_states[:num_tokens]
             elif self.forward_mode.is_target_verify():  # verify
-                num_tokens = bs * self.spec_info.draft_token_num
+                num_tokens = bs * self.spec_info.num_tokens_per_req
                 if logits_output.next_token_logits is not None:
                     logits_output.next_token_logits = logits_output.next_token_logits[
                         :num_tokens
                     ]
-                logits_output.hidden_states = logits_output.hidden_states[:num_tokens]
+                if logits_output.hidden_states is not None:
+                    logits_output.hidden_states = logits_output.hidden_states[
+                        :num_tokens
+                    ]
             elif self.forward_mode.is_draft_extend_v2():  # draft extend_v2
                 bs = bs * self.spec_info.num_tokens_per_req
                 if logits_output.next_token_logits is not None:

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -523,6 +523,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         cap_layout.verify_lens.copy_(live.verify_lens)
         cap_layout.qo_indptr_device.copy_(live.qo_indptr_device)
 
+    @staticmethod
+    def _max_dp_batch_size(forward_batch: ForwardBatch) -> int:
+        request_counts = forward_batch.original_global_num_tokens_cpu
+        if request_counts is None:
+            raise RuntimeError(
+                "DP CUDA graph replay requires raw per-rank request counts"
+            )
+        return max(request_counts)
+
     def can_run_graph(self, forward_batch: ForwardBatch):
         # Disable for token embedding overrides (dynamic per-request)
         if forward_batch.replace_embeds is not None:
@@ -550,9 +559,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             return False
 
         if self.require_mlp_tp_gather:
-            # Raw sync values are per-rank request counts on decode-family
+            cuda_graph_bs = self._max_dp_batch_size(forward_batch)
-            # rounds -- no width division, no per-algorithm enumeration.
-            cuda_graph_bs = max(forward_batch.original_global_num_tokens_cpu)
         else:
             cuda_graph_bs = forward_batch.batch_size
 
@@ -1103,15 +1110,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         else:
             raw_num_token = raw_bs * self.captured_req_width
             if self.require_mlp_tp_gather:
-                max_num_tokens = max(forward_batch.global_num_tokens_cpu)
+                max_batch_size = self._max_dp_batch_size(forward_batch)
-                max_batch_size = (
+                bs = self._pad_to_bucket(max_batch_size, self.capture_bs)
-                    max_num_tokens / self.captured_req_width
-                    if self.model_runner.spec_algorithm.is_eagle()
-                    or self.model_runner.spec_algorithm.is_standalone()
-                    or self.model_runner.spec_algorithm.is_dflash_family()
-                    else max_num_tokens
-                )
-                bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
             else:
                 bs = self._pad_to_bucket(raw_bs, self.capture_bs)
             padded_num_tokens = bs * self.captured_req_width

test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py

@@ -16,6 +16,9 @@ from unittest.mock import MagicMock
 import torch
 
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    DecodeCudaGraphRunner,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -37,6 +40,70 @@ def _logits_output(num_rows: int) -> SimpleNamespace:
 
 
 class TestMlpSyncPadUnpad(CustomTestCase):
+    def test_init_mlp_sync_metadata_scales_speculative_request_width(self):
+        spec_info = SimpleNamespace(
+            num_tokens_per_req=4,
+            num_tokens_for_logprob_per_req=2,
+        )
+        fb = ForwardBatch(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            batch_size=2,
+            input_ids=torch.arange(8),
+            req_pool_indices=torch.tensor([0, 1]),
+            seq_lens=torch.tensor([5, 6]),
+            out_cache_loc=torch.arange(8),
+            seq_lens_sum=11,
+            positions=torch.arange(8),
+            spec_info=spec_info,
+        )
+        batch = SimpleNamespace(
+            global_num_tokens=[2, 0, 3],
+            global_num_tokens_for_logprob=[2, 0, 3],
+            can_run_dp_cuda_graph=True,
+        )
+
+        fb.init_mlp_sync_metadata(batch, torch.device("cpu"))
+
+        self.assertEqual(fb.original_global_num_tokens_cpu, [2, 0, 3])
+        self.assertEqual(fb.global_num_tokens_cpu, [8, 0, 12])
+        self.assertEqual(fb.global_num_tokens_for_logprob_cpu, [4, 0, 6])
+        torch.testing.assert_close(fb.global_num_tokens_gpu, torch.tensor([8, 0, 12]))
+        torch.testing.assert_close(
+            fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6])
+        )
+        self.assertTrue(fb.can_run_dp_cuda_graph)
+
+    def test_draft_input_without_hidden_states_can_be_padded(self):
+        spec_info = SimpleNamespace(
+            is_draft_input=lambda: True,
+            hidden_states=None,
+        )
+        fb = ForwardBatch(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            batch_size=1,
+            input_ids=torch.tensor([11]),
+            req_pool_indices=torch.tensor([5]),
+            seq_lens=torch.tensor([7]),
+            out_cache_loc=torch.tensor([0]),
+            seq_lens_sum=7,
+            positions=torch.tensor([6]),
+            seq_lens_cpu=torch.tensor([7]),
+            lora_ids=[None],
+            spec_info=spec_info,
+        )
+
+        fb._pad_inputs_to_size(_mock_model_runner(), num_tokens=2, bs=1)
+
+        self.assertIsNone(spec_info.hidden_states)
+
+    def test_dp_cuda_graph_batch_size_uses_raw_request_counts(self):
+        fb = SimpleNamespace(original_global_num_tokens_cpu=[3, 11, 7])
+        self.assertEqual(DecodeCudaGraphRunner._max_dp_batch_size(fb), 11)
+
+        fb.original_global_num_tokens_cpu = None
+        with self.assertRaisesRegex(RuntimeError, "raw per-rank request counts"):
+            DecodeCudaGraphRunner._max_dp_batch_size(fb)
+
     def test_decode_post_forward_unpads_per_request_tensors(self):
         fb = ForwardBatch(
             forward_mode=ForwardMode.DECODE,