diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 89426142f..9231e8fd6 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 9afb44377..cfdccbc45 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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 - 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=logits_output.next_token_logits, - candidate_ids=selector_candidate_ids, - q_rows=selector_q_rows, - sampling_info=sampling_info, - draft_input=draft_input, - ) - 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=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) + ( + accept_len, + commit_lens, + bonus, + out_tokens, + new_seq_lens, + target_predict, + ) = self._accept_block( + candidates=candidates, + next_token_logits=logits_output.next_token_logits, + sampling_info=sampling_info, + draft_input=draft_input, + prefix_lens=prefix_lens, + bs=bs, + ) if SIMULATE_ACC_LEN > 0: if SIMULATE_ACC_TOKEN_MODE not in ("fixed", "real-draft-token"): diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 3217393f9..5aae425b7 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/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( - step_logits=step_logits, - temperatures=temperatures, - greedy_mask=greedy_mask, - exp_noise=exp_noise, + return tp_sync.sync( + SpecTpSyncSite.DSPARK_DRAFT_SAMPLE, + SampleStepTokens.execute( + step_logits=step_logits, + 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 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index a5b2a9f7f..f086b91ba 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/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( - step_logits=step_logits, - temperatures=self.temperatures[:bs], - greedy_mask=self.greedy_mask[:bs], - exp_noise=noise, + return self._tp_sync.sync( + SpecTpSyncSite.DSPARK_GRAPH_SAMPLE, + SampleStepTokens.execute( + step_logits=step_logits, + 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( - device, torch.get_device_module().current_device() - ) - if available_gb - need_gb >= _CAPTURE_HEADROOM_GB: + if available_memory_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, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index c3a04111b..45e7c0544 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/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], diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index f486400c8..cb4fc4e6d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/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, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 1edb14f97..77499cd99 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/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 = get_available_gpu_memory(self.device, self.gpu_id) + available_mem = self._tp_sync.available_memory_gb( + 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) diff --git a/python/sglang/srt/speculative/spec_tp_sync.py b/python/sglang/srt/speculative/spec_tp_sync.py new file mode 100644 index 000000000..9bbc8c84b --- /dev/null +++ b/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, + )