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
Open

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

Show file tree
Hide file tree
Changes from 1 commit
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
5 changes: 5 additions & 0 deletions python/sglang/srt/environ.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2028,6 +2028,11 @@ class Envs:
SGLANG_OPT_SKIP_GLUED_REPLAY_VIEW = EnvBool(
SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION.get
)
# DSpark draft under DP-MoE sync: stage the per-rank token counts through
# persistent pinned / device rows (one H2D) instead of fresh tensors every step.
SGLANG_OPT_DSPARK_PERSISTENT_DP_META = EnvBool(
SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION.get
)
# DSpark static verify: record the post-verify engram history commit in the
# target-verify CUDA graph (epilogue) instead of launching it after the replay.
SGLANG_OPT_DSV41_POST_FORWARD_CUDA_GRAPH = EnvBool(
Expand Down
57 changes: 52 additions & 5 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 @@ -406,6 +406,11 @@ def __init__(
if enable_num_token_non_padded()
else None
)
self._persistent_dp_meta = envs.SGLANG_OPT_DSPARK_PERSISTENT_DP_META.get()
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 @@ -638,8 +643,10 @@ 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
if self._persistent_dp_meta
else _make_num_token_non_padded(draft_num_tokens, device)
),
global_num_token_non_padded_cpu=draft_num_tokens,
)
Expand Down Expand Up @@ -732,16 +739,56 @@ 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
if self._persistent_dp_meta:
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)
return
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
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)

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]
17 changes: 15 additions & 2 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
@@ -1,3 +1,4 @@
import itertools
import unittest
from contextlib import nullcontext
from types import SimpleNamespace
Expand Down Expand Up @@ -39,9 +40,9 @@ def test_v4_idle_draft_attention_returns_empty_rows_without_kv_work(self):
self.assertEqual(output.device, hidden.device)

def test_draft_counts_are_not_scaled_twice(self):
for width in (6, 7):
for persistent, width in itertools.product((False, True), (6, 7)):
for rank, local_tokens in ((0, 0), (1, 3 * width), (2, 0)):
with self.subTest(width=width, rank=rank):
with self.subTest(persistent=persistent, width=width, rank=rank):
plan = DPSpecPrefillCoordinationPlan(
[8192, 3, 0], [1, 3, 0], torch.tensor([1, 0, 0]), width, 7
)
Expand All @@ -50,6 +51,11 @@ 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._persistent_dp_meta = persistent
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 +70,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