diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 62468e61a..c7f27428e 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/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( diff --git a/test/registered/spec/dspark/test_dspark_dp_tier.py b/test/registered/spec/dspark/test_dspark_dp_tier.py index f0bcb4c7b..169531224 100644 --- a/test/registered/spec/dspark/test_dspark_dp_tier.py +++ b/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)