Skip to content

[Perf] [NVIDIA] dspark: Stage draft DP-MoE token counts through persistent buffers #47677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
weireweire wants to merge 4 commits into
base: main
Choose a base branch
from
+86 −28
Open

[Perf] [NVIDIA] dspark: Stage draft DP-MoE token counts through persistent buffers #47677

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7fc04f8
Stage DSpark draft DP-MoE token counts through persistent buffers
weireweire Oct 31, 2026
0eba5d0
Turn on DSpark persistent DP-MoE metadata by default
weireweire Nov 2, 2026
c2a44f0
Drop the DSpark persistent DP-meta switch and its legacy path
weireweire Nov 2, 2026
5d3d618
Check both staged DSpark DP-meta rows and the two-slot reuse in tests
weireweire Nov 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 43 additions & 27 deletions python/sglang/srt/speculative/dspark_components/dspark_draft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,18 +45,6 @@
_DRAFT_STEP_LOGITS = Invariant("dspark.draft.step_logits", Bucket.GUARD, NotNaN())


def _make_num_token_non_padded(
num_tokens: int, device: str | torch.device
) -> Optional[torch.Tensor]:
if not enable_num_token_non_padded():
return None
return torch.tensor(
num_tokens,
dtype=torch.int32,
pin_memory=is_pin_memory_available(device),
).to(device, non_blocking=True)


class DraftBlockResult(msgspec.Struct, frozen=True):
draft_tokens: torch.Tensor
corrected_logits: Optional[torch.Tensor]
Expand Down Expand Up @@ -406,6 +394,10 @@ def __init__(
if enable_num_token_non_padded()
else None
)
self._dp_meta_host: Optional[torch.Tensor] = None
self._dp_meta_device: Optional[torch.Tensor] = None
self._dp_meta_copied: list = []
self._dp_meta_slot = 0

@property
def noise(self) -> DraftNoise:
Expand Down Expand Up @@ -468,7 +460,6 @@ def propose(
draft_input=draft_input,
verify_window=verify_window,
bs=bs,
device=device,
embed_module=embed_module,
draft_sampler=draft_sampler,
sampling_info=sampling_info,
Expand Down Expand Up @@ -588,7 +579,6 @@ def _run_forward(
draft_input: DFlashDraftInputV2,
verify_window: VerifyWindow,
bs: int,
device: str,
embed_module,
draft_sampler=None,
sampling_info=None,
Expand Down Expand Up @@ -638,9 +628,7 @@ def _run_forward(
spec_algorithm=SpeculativeAlgorithm.DSPARK,
spec_info=self._draft_block_spec_info,
capture_hidden_mode=CaptureHiddenMode.NULL,
global_num_token_non_padded=_make_num_token_non_padded(
draft_num_tokens, device
),
global_num_token_non_padded=self._num_token_non_padded,
global_num_token_non_padded_cpu=draft_num_tokens,
)
verify_window.iteration_ids.bind(
Expand Down Expand Up @@ -732,16 +720,44 @@ def _fill_dp_moe_sync_metadata(
)
forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens
num_tokens = forward_batch.input_ids.numel()
num_token_non_padded = _make_num_token_non_padded(num_tokens, device)
if num_token_non_padded is not None:
forward_batch.global_num_token_non_padded = num_token_non_padded
forward_batch.global_num_token_non_padded_cpu = num_tokens
forward_batch.global_num_tokens_cpu = gnt
forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob
pin_memory = is_pin_memory_available(device)
forward_batch.global_num_tokens_gpu = torch.tensor(
gnt, dtype=torch.int64, pin_memory=pin_memory
).to(device, non_blocking=True)
forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor(
gnt_logprob, dtype=torch.int64, pin_memory=pin_memory
).to(device, non_blocking=True)
if self._num_token_non_padded is not None:
forward_batch.global_num_token_non_padded = self._num_token_non_padded
(
forward_batch.global_num_tokens_gpu,
forward_batch.global_num_tokens_for_logprob_gpu,
) = self._stage_dp_meta(gnt, gnt_logprob, device)

def _stage_dp_meta(
self, gnt: list, gnt_logprob: list, device
) -> tuple[torch.Tensor, torch.Tensor]:
dp = len(gnt)
if self._dp_meta_host is None or self._dp_meta_host.shape[2] != dp:
self._dp_meta_host = torch.empty(
(2, 2, dp),
dtype=torch.int64,
pin_memory=is_pin_memory_available(device),
)
self._dp_meta_device = torch.empty(
(2, 2, dp), dtype=torch.int64, device=device
)
self._dp_meta_copied = [None, None]
slot = self._dp_meta_slot
self._dp_meta_slot ^= 1
# Rewrite a pinned slot only after its previous H2D has run.
copied = self._dp_meta_copied[slot]
if copied is not None:
copied.synchronize()

@Qiaolin-Yu Qiaolin-Yu Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does here need sync? Could we remove this?

@weireweire weireweire Nov 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a double buffer (note self._dp_meta_slot ^= 1), so it waits for the copy from two iterations ago. That copy has typically finished already, so it doesn't actually wait, but we still need the barrier to guarantee it.

@weireweire weireweire Nov 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note here we are reusing the two slots, so we need to make sure its data is transferred already and then rewrite again

@Qiaolin-Yu Qiaolin-Yu Nov 9, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we use write_host_ints to transfer this?

@weireweire weireweire Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that one seems default off by SGLANG_ENABLE_MLP_SYNC_METADATA_FILL, even using that we may still need this branch for default. only keep one path look cleaner, the the wait is mostly no-op, does that make sense?

host = self._dp_meta_host[slot]
rows = host.numpy()
rows[0] = gnt
rows[1] = gnt_logprob
dev = self._dp_meta_device[slot]
dev.copy_(host, non_blocking=True)
if copied is None:
copied = torch.get_device_module(device).Event()
self._dp_meta_copied[slot] = copied
copied.record()
return dev[0], dev[1]
5 changes: 4 additions & 1 deletion test/registered/spec/dspark/test_dspark_dp_tier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,6 @@ def forward(batch):
),
verify_window=window,
bs=1,
device="cpu",
embed_module=torch.nn.Embedding(16, 2),
)
torch.testing.assert_close(result.raw_hidden, hidden[1:])
Expand All @@ -98,6 +97,10 @@ def test_preserves_unscaled_request_counts_for_cuda_graph_admission(self):
)
proposer.draft_model_runner = SimpleNamespace(device="cpu")
proposer._num_token_non_padded = torch.empty((1,), dtype=torch.int32)
proposer._dp_meta_host = None
proposer._dp_meta_device = None
proposer._dp_meta_copied = []
proposer._dp_meta_slot = 0

forward_batch = SimpleNamespace(input_ids=torch.arange(6))
batch = SimpleNamespace(
Expand Down
11 changes: 11 additions & 0 deletions test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,10 @@ def test_draft_counts_are_not_scaled_twice(self):
proposer = object.__new__(DraftBlockProposer)
proposer._dp_moe_sync = True
proposer._num_token_non_padded = None
proposer._dp_meta_host = None
proposer._dp_meta_device = None
proposer._dp_meta_copied = []
proposer._dp_meta_slot = 0
proposer.draft_model_runner = SimpleNamespace(device="cpu")
proposer._draft_block_spec_info = SimpleNamespace(
num_tokens_per_req=width, num_tokens_for_logprob_per_req=1
Expand All @@ -64,6 +68,13 @@ def test_draft_counts_are_not_scaled_twice(self):
self.assertEqual(
forward.global_num_tokens_for_logprob_cpu, [0, 3 * width, 0]
)
self.assertEqual(
forward.global_num_tokens_gpu.tolist(), [0, 3 * width, 0]

@nvpohanh nvpohanh Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[by Claude Code] Severity: style | Confidence: High

Issue

These new GPU-count assertions cannot catch a row mix-up in _stage_dp_meta. Here both counts are [0, 3 * width, 0], so the two rows hold the same values. Each test also builds a fresh proposer and fills once, so the two-slot reuse never runs.

Example

I changed _stage_dp_meta to return dev[0], dev[0]. All 13 tests in the three test files listed in the PR still pass.

Fix

In test_preserves_unscaled_request_counts_for_cuda_graph_admission, the counts differ ([6, 18, 0, 12] vs [1, 3, 0, 2]), so assert both GPU tensors there. Then fill three times with different counts on one proposer. Check that each call's tensors hold its own counts and that the third call reuses the first call's storage.

)
self.assertEqual(
forward.global_num_tokens_for_logprob_gpu.tolist(),
[0, 3 * width, 0],
)
self.assertTrue(forward.dp_spec_prefill_coordination_applied)
self.assertTrue(forward.is_extend_in_batch)
self.assertFalse(forward.can_run_decode_cuda_graph)
Expand Down
Loading