Fix DSpark and DP/EP (#33098)

Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

python/sglang/srt/speculative/dspark_components/dspark_draft.py

@@ -16,6 +16,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
     ForwardMode,
+    enable_num_token_non_padded,
 )
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -384,6 +385,13 @@ class DraftBlockProposer:
             batch.global_num_tokens_for_logprob,
         )
         device = self.draft_model_runner.device
+        forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens
+        num_tokens = forward_batch.input_ids.numel()
+        if enable_num_token_non_padded():
+            forward_batch.num_token_non_padded = torch.tensor(
+                num_tokens, dtype=torch.int32, device=device
+            )
+        forward_batch.num_token_non_padded_cpu = num_tokens
         forward_batch.global_num_tokens_cpu = gnt
         forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob
         forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to(

test/registered/spec/dspark/test_dspark_dp_tier.py

@@ -1,6 +1,11 @@
 import random
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_draft import DraftBlockProposer
 from sglang.srt.speculative.dspark_components.dspark_planner import (
     dp_global_verify_tier_num_tokens,
     local_verify_tier_num_tokens,
@@ -48,6 +53,40 @@ class TestDpGlobalVerifyTierNumTokens(CustomTestCase):
         )
 
 
+class TestDraftDpSyncMetadata(CustomTestCase):
+    def test_preserves_unscaled_request_counts_for_cuda_graph_admission(self):
+        proposer = DraftBlockProposer.__new__(DraftBlockProposer)
+        proposer._dp_moe_sync = True
+        proposer._draft_block_spec_info = SimpleNamespace(
+            num_tokens_per_req=6,
+            num_tokens_for_logprob_per_req=1,
+        )
+        proposer.draft_model_runner = SimpleNamespace(device="cpu")
+
+        forward_batch = SimpleNamespace(input_ids=torch.arange(6))
+        batch = SimpleNamespace(
+            global_num_tokens=[1, 3, 0, 2],
+            global_num_tokens_for_logprob=[1, 3, 0, 2],
+            can_run_dp_cuda_graph=True,
+        )
+
+        with patch(
+            "sglang.srt.speculative.dspark_components.dspark_draft.enable_num_token_non_padded",
+            return_value=True,
+        ):
+            proposer._fill_dp_moe_sync_metadata(forward_batch, batch)
+
+        self.assertEqual(
+            forward_batch.original_global_num_tokens_cpu,
+            [1, 3, 0, 2],
+        )
+        self.assertEqual(forward_batch.global_num_tokens_cpu, [6, 18, 0, 12])
+        self.assertEqual(forward_batch.num_token_non_padded.item(), 6)
+        self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32)
+        self.assertEqual(forward_batch.num_token_non_padded_cpu, 6)
+        self.assertTrue(forward_batch.can_run_dp_cuda_graph)
+
+
 class TestBusyIdleGraphKeyIdentity(CustomTestCase):
     def test_busy_and_idle_floors_agree_on_random_topologies(self):
         rng = random.Random(20260703)