[Spec] Fix Dspark and Dflash state divergence across TP rank (#33614)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

python/sglang/srt/environ.py

@@ -1205,6 +1205,11 @@ class Envs:
     # Saves the per-step draft forward, but the draft KV goes stale: an upshift
     # back to steps>0 starts from a cold draft state (low accept until it recovers).
     SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
+    # Which speculative decisions rank 0 broadcasts to its TP group; narrowing
+    # it under live traffic isolates where ranks actually diverge. Comma
+    # separated presets ("all", "rng", "init", "off"), or SpecTpSyncSite slugs
+    # and numbers, each negatable with a leading "-": "all,-dspark-plan,-6".
+    SGLANG_SPEC_TP_SYNC = EnvStr("all")
     # Kill-switch for the draft-extend cuda graph. Draft extend then always runs
     # eager. Escape hatch for setups where the capture's memory pool costs more
     # than the graph saves (e.g. DeepEP MoE workspace captured at full dispatch

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -62,6 +62,7 @@ from sglang.srt.speculative.draft_worker_common import (
 )
 from sglang.srt.speculative.dspark_components.dspark_draft import resolve_greedy_mask
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 from sglang.srt.speculative.spec_utils import (
     SIMULATE_ACC_LEN,
     SIMULATE_ACC_METHOD,
@@ -70,7 +71,7 @@ from sglang.srt.speculative.spec_utils import (
     assign_req_to_token_pool_func,
     build_grammar_vocab_mask,
 )
-from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu
+from sglang.srt.utils import is_cuda, is_hip, is_npu
 
 _is_npu = is_npu()
 
@@ -302,6 +303,7 @@ class DFlashWorkerV2(BaseSpecWorker):
         self._warned_sampling_fallback = False
         self._draft_probs_buf = None
         self._logged_first_verify = False
+        self._tp_sync = SpecTpSync(get_tp_group())
 
         bundle = build_draft_tp_worker(
             server_args=server_args,
@@ -465,7 +467,12 @@ class DFlashWorkerV2(BaseSpecWorker):
             get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
         )
         if is_cuda() and capture_decode_cuda_graph:
-            available_mem = get_available_gpu_memory(self.device, self.gpu_id)
+            available_mem = self._tp_sync.available_memory_gb(
+                SpecTpSyncSite.DFLASH_MEM,
+                self.device,
+                self.gpu_id,
+                group=get_tp_group(),
+            )
             if available_mem < 1.0:
                 capture_decode_cuda_graph = False
                 logger.warning(
@@ -1619,6 +1626,89 @@ class DFlashWorkerV2(BaseSpecWorker):
             self._new_seq_lens_bufs[slot][:bs],
         )
 
+    def _accept_block(
+        self,
+        *,
+        candidates: torch.Tensor,
+        next_token_logits: torch.Tensor,
+        sampling_info,
+        draft_input,
+        prefix_lens: torch.Tensor,
+        bs: int,
+    ):
+        new_seq_lens = None
+        target_predict = None
+        if self._selector_sample is not None:
+            selector_candidate_ids, selector_q_rows = self._selector_sample
+            accept_len, bonus = self._selector_sampling_accept(
+                candidates=candidates,
+                next_token_logits=next_token_logits,
+                candidate_ids=selector_candidate_ids,
+                q_rows=selector_q_rows,
+                sampling_info=sampling_info,
+                draft_input=draft_input,
+            )
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_SELECTOR, accept_len)
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_SELECTOR, bonus)
+            out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
+        elif (
+            not _is_all_greedy(sampling_info) and is_dflash_sampling_verify_available()
+        ):
+            accept_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus(
+                candidates=candidates,
+                next_token_logits=next_token_logits,
+                sampling_info=sampling_info,
+                max_top_k=draft_input.max_top_k,
+                uniform_top_k_value=draft_input.uniform_top_k_value,
+            )
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE, accept_len)
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE, bonus)
+            out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
+        else:
+            target_predict = torch.argmax(next_token_logits, dim=-1).view(
+                bs, int(self.block_size)
+            )
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_GREEDY, target_predict)
+            if self._use_triton_accept_bonus:
+                try:
+                    (
+                        accept_len,
+                        commit_lens,
+                        bonus,
+                        out_tokens,
+                        new_seq_lens,
+                    ) = self._next_accept_bonus_buffers(bs)
+                    _compute_dflash_accept_bonus_triton_unchecked(
+                        candidates=candidates,
+                        target_top1=target_predict,
+                        accept_lens_out=accept_len,
+                        commit_lens_out=commit_lens,
+                        bonus_ids_out=bonus,
+                        out_tokens_out=out_tokens,
+                        prefix_lens=prefix_lens,
+                        new_seq_lens_out=new_seq_lens,
+                    )
+                except Exception as e:
+                    self._use_triton_accept_bonus = False
+                    logger.warning(
+                        "DFLASH Triton accept/bonus failed; falling back to eager path: %s",
+                        e,
+                    )
+                    accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
+                        candidates=candidates,
+                        target_predict=target_predict,
+                    )
+                    out_tokens, commit_lens = _commit_accept(
+                        candidates, accept_len, bonus
+                    )
+            else:
+                accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
+                    candidates=candidates,
+                    target_predict=target_predict,
+                )
+                out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
+        return accept_len, commit_lens, bonus, out_tokens, new_seq_lens, target_predict
+
     def _validate_phase1_sampling_support(self, batch: ScheduleBatch) -> None:
         sampling_info = batch.sampling_info
         # A selector draft carries its own q and verifies through accept_sampling, so
@@ -1678,6 +1768,7 @@ class DFlashWorkerV2(BaseSpecWorker):
                 batch_output.logits_output,
                 batch_output.next_token_ids,
             )
+            self._tp_sync.sync(SpecTpSyncSite.DFLASH_TARGET, next_token_ids)
             batch_output.new_seq_lens = batch.seq_lens
             if on_publish is not None:
                 on_publish(batch_output.new_seq_lens)
@@ -2029,72 +2120,21 @@ class DFlashWorkerV2(BaseSpecWorker):
             grammar_mask.apply(logits_output.next_token_logits)
 
         candidates = draft_tokens
-        new_seq_lens = None
+        (
-        target_predict = None
+            accept_len,
-        if self._selector_sample is not None:
+            commit_lens,
-            selector_candidate_ids, selector_q_rows = self._selector_sample
+            bonus,
-            accept_len, bonus = self._selector_sampling_accept(
+            out_tokens,
-                candidates=candidates,
+            new_seq_lens,
-                next_token_logits=logits_output.next_token_logits,
+            target_predict,
-                candidate_ids=selector_candidate_ids,
+        ) = self._accept_block(
-                q_rows=selector_q_rows,
+            candidates=candidates,
-                sampling_info=sampling_info,
+            next_token_logits=logits_output.next_token_logits,
-                draft_input=draft_input,
+            sampling_info=sampling_info,
-            )
+            draft_input=draft_input,
-            out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
+            prefix_lens=prefix_lens,
-        elif (
+            bs=bs,
-            not _is_all_greedy(sampling_info) and is_dflash_sampling_verify_available()
+        )
-        ):
-            accept_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus(
-                candidates=candidates,
-                next_token_logits=logits_output.next_token_logits,
-                sampling_info=sampling_info,
-                max_top_k=draft_input.max_top_k,
-                uniform_top_k_value=draft_input.uniform_top_k_value,
-            )
-            out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
-        else:
-            target_predict = torch.argmax(logits_output.next_token_logits, dim=-1).view(
-                bs, int(self.block_size)
-            )
-            if self._use_triton_accept_bonus:
-                try:
-                    (
-                        accept_len,
-                        commit_lens,
-                        bonus,
-                        out_tokens,
-                        new_seq_lens,
-                    ) = self._next_accept_bonus_buffers(bs)
-                    _compute_dflash_accept_bonus_triton_unchecked(
-                        candidates=candidates,
-                        target_top1=target_predict,
-                        accept_lens_out=accept_len,
-                        commit_lens_out=commit_lens,
-                        bonus_ids_out=bonus,
-                        out_tokens_out=out_tokens,
-                        prefix_lens=prefix_lens,
-                        new_seq_lens_out=new_seq_lens,
-                    )
-                except Exception as e:
-                    self._use_triton_accept_bonus = False
-                    logger.warning(
-                        "DFLASH Triton accept/bonus failed; falling back to eager path: %s",
-                        e,
-                    )
-                    accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
-                        candidates=candidates,
-                        target_predict=target_predict,
-                    )
-                    out_tokens, commit_lens = _commit_accept(
-                        candidates, accept_len, bonus
-                    )
-            else:
-                accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
-                    candidates=candidates,
-                    target_predict=target_predict,
-                )
-                out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
 
         if SIMULATE_ACC_LEN > 0:
             if SIMULATE_ACC_TOKEN_MODE not in ("fixed", "real-draft-token"):

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

@@ -27,6 +27,7 @@ from sglang.srt.speculative.spec_info import (
     SpeculativeAlgorithm,
     spec_scale_global_num_tokens,
 )
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 from sglang.srt.speculative.spec_utils import draft_tp_context
 from sglang.srt.utils.common import is_pin_memory_available
 from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect
@@ -134,6 +135,7 @@ def sample_draft_block(
     sampling_info,
     markov_head,
     device: torch.device,
+    tp_sync: SpecTpSync,
 ) -> DraftBlockResult:
     bs = base_logits.shape[0]
     greedy_mask = resolve_greedy_mask(bs=bs, sampling_info=sampling_info, device=device)
@@ -151,7 +153,9 @@ def sample_draft_block(
 
         def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
             expect(_DRAFT_STEP_LOGITS, step_logits, msg=f"step {step_idx}")
-            return torch.argmax(step_logits, dim=-1)
+            return tp_sync.sync(
+                SpecTpSyncSite.DSPARK_DRAFT_GREEDY, torch.argmax(step_logits, dim=-1)
+            )
 
     else:
 
@@ -161,11 +165,14 @@ def sample_draft_block(
                 exp_noise = torch.empty(
                     step_logits.shape, dtype=torch.float32, device=step_logits.device
                 ).exponential_(1)
-                return SampleStepTokens.execute(
+                return tp_sync.sync(
-                    step_logits=step_logits,
+                    SpecTpSyncSite.DSPARK_DRAFT_SAMPLE,
-                    temperatures=temperatures,
+                    SampleStepTokens.execute(
-                    greedy_mask=greedy_mask,
+                        step_logits=step_logits,
-                    exp_noise=exp_noise,
+                        temperatures=temperatures,
+                        greedy_mask=greedy_mask,
+                        exp_noise=exp_noise,
+                    ),
                 )
             else:
                 probs = torch.softmax(
@@ -174,7 +181,10 @@ def sample_draft_block(
                 probs = expect(_DRAFT_PROBS, probs)
                 argmax_tokens = torch.argmax(step_logits, dim=-1)
                 sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
-                return torch.where(greedy_mask, argmax_tokens, sampled_tokens)
+                return tp_sync.sync(
+                    SpecTpSyncSite.DSPARK_DRAFT_MULTINOMIAL,
+                    torch.where(greedy_mask, argmax_tokens, sampled_tokens),
+                )
 
     draft_tokens, corrected_logits = markov_head.sample_block(
         base_logits,
@@ -199,6 +209,7 @@ class DraftBlockProposer:
         gamma: int,
         mask_token_id: int,
         draft_block_spec_info,
+        tp_sync: SpecTpSync,
         dp_moe_sync: bool = False,
     ) -> None:
         self.draft_model = draft_model
@@ -208,6 +219,7 @@ class DraftBlockProposer:
         self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1
         self._mask_token_id = mask_token_id
         self._draft_block_spec_info = draft_block_spec_info
+        self._tp_sync = tp_sync
         self._draft_sampler = None
         self._dp_moe_sync = dp_moe_sync
         # Persistent (bs, gamma) mask-token buffer: only column 0 (the bonus
@@ -310,6 +322,7 @@ class DraftBlockProposer:
                 sampling_info=sampling_info,
                 markov_head=self.draft_model.markov_head,
                 device=device,
+                tp_sync=self._tp_sync,
             )
         proposal_block_ids = (
             draft_block_ids

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

@@ -13,7 +13,7 @@ from sglang.srt.models.dspark import VanillaMarkov
 from sglang.srt.speculative.dspark_components.dspark_draft import (
     select_draft_hidden_without_anchor,
 )
-from sglang.srt.utils import get_available_gpu_memory
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +46,7 @@ class DsparkDraftSampler:
         gamma,
         max_bs,
         device,
+        tp_sync: SpecTpSync,
         confidence_fn=None,
         out=None,
         folded_sampling: bool = True,
@@ -73,6 +74,7 @@ class DsparkDraftSampler:
             else None
         )
         self.folded_sampling = folded_sampling
+        self._tp_sync = tp_sync
         self.temperatures = None
         self.greedy_mask = None
         self.exp_noise = None
@@ -144,15 +146,23 @@ class DsparkDraftSampler:
                     # In-graph philox noise: each replay advances the generator
                     # and redraws.
                     noise = self.exp_noise[:bs].exponential_()
-                    return SampleStepTokens.execute(
+                    return self._tp_sync.sync(
-                        step_logits=step_logits,
+                        SpecTpSyncSite.DSPARK_GRAPH_SAMPLE,
-                        temperatures=self.temperatures[:bs],
+                        SampleStepTokens.execute(
-                        greedy_mask=self.greedy_mask[:bs],
+                            step_logits=step_logits,
-                        exp_noise=noise,
+                            temperatures=self.temperatures[:bs],
+                            greedy_mask=self.greedy_mask[:bs],
+                            exp_noise=noise,
+                        ),
                     )
 
             else:
-                sampler = greedy_step_sampler
+
+                def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
+                    return self._tp_sync.sync(
+                        SpecTpSyncSite.DSPARK_GRAPH_GREEDY,
+                        greedy_step_sampler(step_logits, step_idx),
+                    )
 
             draft_tokens, corrected_logits = self.markov_head.sample_block(
                 base_logits,
@@ -177,9 +187,12 @@ class DsparkDraftSampler:
             self.confidence_out[:bs].copy_(confidence)
 
 
-def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
+def _resolve_folded_sampling(
+    *, model, gamma, max_bs, device, tp_rank, available_memory_gb: float
+) -> bool:
     """The sampling buffers are baked into the captured draft graph, so AUTO
-    must decide before capture from a free-memory probe."""
+    must decide before capture from a free-memory probe. ``available_memory_gb``
+    is the group minimum, so every rank folds identically."""
     mode = envs.SGLANG_DSPARK_FOLDED_SAMPLING.get()
     if mode == DsparkFoldedSampling.OFF:
         return False
@@ -189,10 +202,7 @@ def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
     noise_bytes = max_bs * vocab * 4
     logits_bytes = max_bs * gamma * vocab * _base_logits_dtype(model).itemsize
     need_gb = (noise_bytes + logits_bytes) / (1 << 30)
-    available_gb = get_available_gpu_memory(
+    if available_memory_gb - need_gb >= _CAPTURE_HEADROOM_GB:
-        device, torch.get_device_module().current_device()
-    )
-    if available_gb - need_gb >= _CAPTURE_HEADROOM_GB:
         return True
     if tp_rank == 0:
         logger.warning(
@@ -201,7 +211,7 @@ def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
             "the eager proposal path. Set SGLANG_DSPARK_FOLDED_SAMPLING=%d "
             "to force.",
             need_gb,
-            available_gb,
+            available_memory_gb,
             int(DsparkFoldedSampling.FORCE),
         )
     return False
@@ -214,6 +224,8 @@ def maybe_build_draft_sampler(
     max_bs: int,
     device,
     tp_rank: int,
+    tp_sync: SpecTpSync,
+    available_memory_gb: float,
     confidence_fn=None,
     out=None,
 ) -> Optional[DsparkDraftSampler]:
@@ -232,7 +244,12 @@ def maybe_build_draft_sampler(
     if getattr(draft_model, "markov_head", None) is None:
         return _eager("no markov head")
     folded_sampling = _resolve_folded_sampling(
-        model=draft_model, gamma=gamma, max_bs=max_bs, device=device, tp_rank=tp_rank
+        model=draft_model,
+        gamma=gamma,
+        max_bs=max_bs,
+        device=device,
+        tp_rank=tp_rank,
+        available_memory_gb=available_memory_gb,
     )
     if tp_rank == 0:
         logger.info(
@@ -244,6 +261,7 @@ def maybe_build_draft_sampler(
         gamma=gamma,
         max_bs=max_bs,
         device=device,
+        tp_sync=tp_sync,
         confidence_fn=confidence_fn,
         out=out,
         folded_sampling=folded_sampling,

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

@@ -40,6 +40,7 @@ from sglang.srt.speculative.ragged_verify import (
     read_ragged_verify_mode,
     round_up_grid,
 )
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 from sglang.srt.utils.common import require_mlp_tp_gather
 from sglang.srt.utils.invariants import (
     Bucket,
@@ -76,12 +77,14 @@ class DSparkVerifyPlanner:
         device,
         tp_rank: int,
         verify_num_draft_tokens: int,
+        tp_sync: SpecTpSync,
     ) -> None:
         self.draft_model = draft_model
         self.gamma = gamma
         self.model_runner = model_runner
         self.device = device
         self.verify_num_draft_tokens = verify_num_draft_tokens
+        self._tp_sync = tp_sync
         self._align_verify_tokens_to_graph_tier = (
             get_spec().speculative_dspark_align_verify_tokens_to_graph_tier
         )
@@ -368,11 +371,15 @@ class DSparkVerifyPlanner:
             return None
         if not get_schedule().disable_overlap_schedule:
             return draft_input.verify_token_budget
-        return self.compute_budget_sync(
+
+        # No collective: the budget derives only from the broadcast draft tokens
+        # (via confidence), replicated req_generation, and the static sps table.
+        draft_input.verify_token_budget = self.compute_budget_sync(
             confidence=confidence,
             prefix_lens=prefix_lens,
             req_pool_indices=req_pool_indices,
         )
+        return draft_input.verify_token_budget
 
     def confidence_budget_prepare(self):
         if not self.schedules_verify_budget:
@@ -574,6 +581,7 @@ class DSparkVerifyPlanner:
             budget=budget,
             cfg=self._schedule_cfg,
         ).to(device=device, dtype=torch.int32)
+        self._tp_sync.sync(SpecTpSyncSite.DSPARK_PLAN, verify_lens)
 
         if resolve_level() >= InvariantCheckLevel.WARN:
             verify_lens_64 = verify_lens.to(torch.int64)
@@ -593,12 +601,6 @@ class DSparkVerifyPlanner:
                 verify_lens=verify_lens,
             )
 
-        broadcast_group, group_size = verify_lens_broadcast_group(
-            tp_size=get_parallel().tp_size
-        )
-        if group_size > 1:
-            broadcast_group.broadcast(verify_lens, src=0)
-
         return verify_lens
 
     def _log_verify_lens_decision(
@@ -751,12 +753,6 @@ def uniform_ragged_layout(
     )
 
 
-def verify_lens_broadcast_group(*, tp_size: int) -> tuple:
-    if is_dp_attention_enabled():
-        return get_parallel().attn_tp_group, get_parallel().attn_tp_size
-    return get_tp_group(), tp_size
-
-
 def verify_layout_grid(
     *,
     verify_lens_cpu: list[int],

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

@@ -41,6 +41,7 @@ from sglang.srt.speculative.dspark_components.dspark_planner import (
     apply_logits_adjustments_strided,
 )
 from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 from sglang.srt.speculative.spec_utils import (
     SIMULATE_ACC_METHOD,
     sample_simulated_acc_len,
@@ -86,6 +87,7 @@ class TargetVerifyExecutor:
         verify_num_draft_tokens: int,
         model_runner,
         kv_injector: TargetHiddenKvInjector,
+        tp_sync: SpecTpSync,
         verify_epilogue=None,
         simulate_acc_len: float = 0.0,
     ) -> None:
@@ -94,6 +96,7 @@ class TargetVerifyExecutor:
         self.verify_num_draft_tokens = verify_num_draft_tokens
         self.model_runner = model_runner
         self.kv_injector = kv_injector
+        self._tp_sync = tp_sync
         self.verify_epilogue = verify_epilogue
         self._verify_backend_self_adds_seq_lens_cache: Optional[bool] = None
         self._simulate_acc_len = float(simulate_acc_len)
@@ -138,6 +141,15 @@ class TargetVerifyExecutor:
                 bs=bs, dtype=correct_len.dtype, device=correct_len.device
             )
 
+        site = (
+            SpecTpSyncSite.DSPARK_ACCEPT_GREEDY
+            if sampling_info is None or sampling_info.is_all_greedy
+            else SpecTpSyncSite.DSPARK_ACCEPT_SAMPLE
+        )
+        self._tp_sync.sync(site, correct_len)
+        self._tp_sync.sync(site, bonus)
+        self._tp_sync.sync(site, cap_trim_lens)
+
         finalized = FinalizeAcceptLens.execute(
             correct_len=correct_len,
             cap_trim_lens=cap_trim_lens,
@@ -487,12 +499,14 @@ class DsparkVerifyEpilogue:
         max_bs: int,
         verify_num_draft_tokens: int,
         device,
+        tp_sync: SpecTpSync,
         commit_ctx: Optional[CommitInjectCtx] = None,
     ) -> None:
         self.max_bs = int(max_bs)
         self.stride = int(verify_num_draft_tokens)
         self.gamma = self.stride - 1
         self.commit_ctx = commit_ctx
+        self._tp_sync = tp_sync
         self.inject_gate_buf = torch.zeros((1,), dtype=torch.int32, device=device)
         self.verify_lens_buf = torch.zeros(
             (self.max_bs,), dtype=torch.int64, device=device
@@ -637,6 +651,9 @@ class DsparkVerifyEpilogue:
             verify_num_draft_tokens=self.stride,
             cutoff_verify_lens=verify_lens,
         )
+        self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, correct_len)
+        self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, bonus)
+        self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, cap_trim_lens)
         finalized = finalize_accept_lens_triton(
             correct_len=correct_len,
             cap_trim_lens=cap_trim_lens,

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

@@ -16,6 +16,7 @@ from sglang.srt.lora.layers import unwrap_lora_layer
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     compute_position,
@@ -68,6 +69,7 @@ from sglang.srt.speculative.dspark_components.dspark_verify import (
     TargetVerifyExecutor,
     verify_logits_adjustments_are_noop,
 )
+from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
 from sglang.srt.speculative.spec_utils import (
     GrammarTree,
     build_grammar_vocab_mask,
@@ -75,8 +77,8 @@ from sglang.srt.speculative.spec_utils import (
     prepare_mamba_track_for_verify,
 )
 from sglang.srt.utils import (
-    get_available_gpu_memory,
     is_cuda,
+    is_cuda_alike,
     is_npu,
     is_pin_memory_available,
 )
@@ -114,7 +116,8 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
         self._is_pd_prefill = get_disagg().disaggregation_mode == "prefill"
         self._decode_graph_allowed = (
-            not get_exec().graph.disable_cuda_graph and not self._is_pd_prefill
+            get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
+            and not self._is_pd_prefill
         )
         if (
             get_parallel().enable_dp_attention
@@ -172,6 +175,18 @@ class DSparkWorkerV2(BaseSpecWorker):
         self.speculative_num_draft_tokens = self.verify_num_draft_tokens
         self._mask_token_id = runtime_config.mask_token_id
 
+        parallel = get_parallel()
+        self._tp_sync = SpecTpSync(
+            parallel.attn_tp_group
+            if parallel.enable_dp_attention
+            else parallel.tp_group
+        )
+        self._draft_graph_group = (
+            parallel.attn_tp_group
+            if self._draft_dp_context_enabled
+            else parallel.tp_group
+        )
+
         if self.ps.tp_rank == 0:
             logger.info(
                 "Initialized DSpark draft runner. attention_backend=%s, model=%s, "
@@ -220,6 +235,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             device=self.device,
             tp_rank=self.ps.tp_rank,
             verify_num_draft_tokens=self.verify_num_draft_tokens,
+            tp_sync=self._tp_sync,
         )
         if (
             get_parallel().enable_dp_attention
@@ -248,6 +264,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             gamma=self.gamma,
             mask_token_id=self._mask_token_id,
             draft_block_spec_info=self._draft_block_spec_info,
+            tp_sync=self._tp_sync,
             dp_moe_sync=self._draft_is_moe and get_parallel().enable_dp_attention,
         )
         self._verify_epilogue = None
@@ -260,6 +277,7 @@ class DSparkWorkerV2(BaseSpecWorker):
                 max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
                 verify_num_draft_tokens=self.verify_num_draft_tokens,
                 device=self.device,
+                tp_sync=self._tp_sync,
                 commit_ctx=CommitInjectCtx(
                     draft_model=self.draft_model,
                     block_pos_offsets=self._block_pos_offsets,
@@ -298,6 +316,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             verify_num_draft_tokens=self.verify_num_draft_tokens,
             model_runner=self.model_runner,
             kv_injector=self._kv_injector,
+            tp_sync=self._tp_sync,
             verify_epilogue=self._verify_epilogue,
             simulate_acc_len=self._simulate_acc_len,
         )
@@ -367,8 +386,13 @@ class DSparkWorkerV2(BaseSpecWorker):
 
     def init_cuda_graphs(self):
         capture_decode_cuda_graph = self._decode_graph_allowed
-        if is_cuda() and capture_decode_cuda_graph:
+        available_mem = self._tp_sync.available_memory_gb(
-            available_mem = get_available_gpu_memory(self.device, self.gpu_id)
+            SpecTpSyncSite.DSPARK_MEM,
+            self.device,
+            self.gpu_id,
+            group=self._draft_graph_group,
+        )
+        if is_cuda_alike() and capture_decode_cuda_graph:
             if available_mem < 1.0:
                 capture_decode_cuda_graph = False
                 logger.warning(
@@ -385,7 +409,9 @@ class DSparkWorkerV2(BaseSpecWorker):
                 # from being the intended precision fallback, skipping the
                 # unused hook avoids paying for two proposal computations.
                 if envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get():
-                    self._draft_sampler = self._maybe_build_draft_sampler()
+                    self._draft_sampler = self._maybe_build_draft_sampler(
+                        available_memory_gb=available_mem
+                    )
                     if self._draft_sampler is not None:
                         self.draft_model_runner.capture_tail_hooks.append(
                             make_draft_sampler_capture_hook(self._draft_sampler)
@@ -395,13 +421,15 @@ class DSparkWorkerV2(BaseSpecWorker):
                 capture_decode_cuda_graph=capture_decode_cuda_graph
             )
 
-    def _maybe_build_draft_sampler(self):
+    def _maybe_build_draft_sampler(self, *, available_memory_gb: float):
         return maybe_build_draft_sampler(
             draft_model=self.draft_model,
             gamma=self.gamma,
             max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
             device=self.device,
             tp_rank=self.ps.tp_rank,
+            tp_sync=self._tp_sync,
+            available_memory_gb=available_memory_gb,
             confidence_fn=(
                 self._verify_planner.compute_confidence_tensor
                 if self._verify_planner.carries_confidence
@@ -461,6 +489,7 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
         logits_output = batch_output.logits_output
         next_token_ids = batch_output.next_token_ids
+        self._tp_sync.sync(SpecTpSyncSite.DSPARK_TARGET, next_token_ids)
         batch_output.new_seq_lens = batch.seq_lens
         if on_publish is not None:
             on_publish(batch_output.new_seq_lens)

python/sglang/srt/speculative/spec_tp_sync.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import logging
+from enum import IntEnum
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_available_gpu_memory
+
+logger = logging.getLogger(__name__)
+
+
+class SpecTpSyncSite(IntEnum):
+    """Every place a speculative step broadcasts a decision from rank 0.
+
+    Number and slug are the stable handles ``SGLANG_SPEC_TP_SYNC`` selects by.
+    """
+
+    # -- DSpark --
+    DSPARK_MEM = 1  # gates both graph capture and folded sampling
+    DSPARK_DRAFT_GREEDY = 2
+    DSPARK_DRAFT_SAMPLE = 3
+    DSPARK_DRAFT_MULTINOMIAL = 4
+    DSPARK_GRAPH_SAMPLE = 5  # in-graph philox, redrawn per replay
+    DSPARK_GRAPH_GREEDY = 6
+    DSPARK_PLAN = 7
+    DSPARK_ACCEPT_GREEDY = 8
+    DSPARK_ACCEPT_SAMPLE = 9
+    DSPARK_ACCEPT_GRAPH = 10
+    DSPARK_TARGET = 11
+
+    # -- DFlash --
+    DFLASH_MEM = 12
+    DFLASH_SELECTOR = 13
+    DFLASH_ACCEPT_SAMPLE = 14
+    DFLASH_ACCEPT_GREEDY = 15
+    DFLASH_TARGET = 16
+
+    @property
+    def slug(self) -> str:
+        return self.name.lower().replace("_", "-")
+
+
+_ALL = frozenset(SpecTpSyncSite)
+_INIT = frozenset({SpecTpSyncSite.DSPARK_MEM, SpecTpSyncSite.DFLASH_MEM})
+# Sites that draw from the RNG, so they can differ under identical logits.
+_RNG = frozenset(
+    {
+        SpecTpSyncSite.DSPARK_DRAFT_SAMPLE,
+        SpecTpSyncSite.DSPARK_DRAFT_MULTINOMIAL,
+        SpecTpSyncSite.DSPARK_GRAPH_SAMPLE,
+        SpecTpSyncSite.DSPARK_ACCEPT_SAMPLE,
+        SpecTpSyncSite.DSPARK_TARGET,
+        SpecTpSyncSite.DFLASH_SELECTOR,
+        SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE,
+        SpecTpSyncSite.DFLASH_TARGET,
+    }
+)
+
+_PRESETS = {
+    "all": _ALL,
+    "off": frozenset(),
+    "none": frozenset(),
+    # The one input measured to differ across ranks.
+    "init": _INIT,
+    "rng": _INIT | _RNG,
+}
+
+_BY_SLUG = {site.slug: site for site in SpecTpSyncSite}
+_BY_NUMBER = {str(int(site)): site for site in SpecTpSyncSite}
+
+
+def _resolve(name: str) -> frozenset[SpecTpSyncSite]:
+    if name in _PRESETS:
+        return _PRESETS[name]
+    site = _BY_SLUG.get(name) or _BY_NUMBER.get(name)
+    if site is None:
+        raise ValueError(
+            f"SGLANG_SPEC_TP_SYNC: unknown token {name!r}. "
+            f"Presets: {sorted(_PRESETS)}. "
+            f"Sites: {[(int(s), s.slug) for s in SpecTpSyncSite]}."
+        )
+    return frozenset({site})
+
+
+def parse_spec_tp_sync(spec: str) -> frozenset[SpecTpSyncSite]:
+    """Parse ``SGLANG_SPEC_TP_SYNC``; see its comment in environ.py for the syntax."""
+    sites: frozenset[SpecTpSyncSite] = frozenset()
+    for token in spec.replace(" ", "").replace("_", "-").lower().split(","):
+        if not token:
+            continue
+        negate = token.startswith("-")
+        value = _resolve(token[1:] if negate else token)
+        sites = sites - value if negate else sites | value
+    return sites
+
+
+class SpecTpSync:
+    """Broadcasts a speculative decision from rank 0 to its TP group."""
+
+    def __init__(self, tp_group) -> None:
+        self._tp_group = tp_group
+        # Parsed even on a single rank so a typo fails on every deployment.
+        sites = parse_spec_tp_sync(envs.SGLANG_SPEC_TP_SYNC.get())
+        self._sites = sites if tp_group.world_size > 1 else frozenset()
+        if sites != _ALL and tp_group.world_size > 1 and tp_group.rank_in_group == 0:
+            logger.warning(
+                "Speculative TP sync limited to %s.",
+                [f"{int(s)}:{s.slug}" for s in sorted(sites)] or "no site",
+            )
+
+    def enabled(self, site: SpecTpSyncSite) -> bool:
+        return site in self._sites
+
+    def sync(self, site: SpecTpSyncSite, values: torch.Tensor) -> torch.Tensor:
+        if site in self._sites:
+            self._tp_group.broadcast(values, src=0)
+        return values
+
+    def available_memory_gb(self, site: SpecTpSyncSite, device, gpu_id, *, group):
+        """Free GPU memory, reduced to the group minimum when ``site`` is on."""
+        distributed = self.enabled(site) and group.world_size > 1
+        return get_available_gpu_memory(
+            device,
+            gpu_id,
+            distributed=distributed,
+            cpu_group=group.cpu_group if distributed else None,
+        )