diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index c25230cc62fd..0c11fca2141c 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -118,6 +118,7 @@ def __init__( tp_sync: SpecTpSync, verify_epilogue=None, simulate_acc_len: float = 0.0, + simulated_correct_len: Optional[torch.Tensor] = None, block_verification: bool = False, ) -> None: self.target_worker = target_worker @@ -139,7 +140,7 @@ def __init__( self._verify_backend_self_adds_seq_lens_cache: Optional[bool] = None self._simulate_acc_len = float(simulate_acc_len) self._block_verification = block_verification - self._simulated_correct_drafts_buf: Optional[torch.Tensor] = None + self._simulated_correct_len = simulated_correct_len def accept_and_finalize( self, @@ -160,13 +161,13 @@ def accept_and_finalize( Folded path: the accept/finalize/out-token kernels already ran inside the target-verify cuda graph (DsparkVerifyEpilogue); read its buffers. - Eager path: run them here, including the SGLANG_SIMULATE_ACC_LEN - override. + Eager path: run them here. Both apply the SGLANG_SIMULATE_ACC_LEN + override drawn by begin_simulated_step. """ if folded_accept: return self.verify_epilogue.read_accept(bs) - simulate = self._simulate_acc_len > 0 + simulate = self._simulated_correct_len is not None # Simulated acceptance overwrites correct_len below, so a sampling accept # (draft/target softmax + rejection) would be computed only to be discarded. accept_sampling_info = None if simulate and _is_hip else sampling_info @@ -183,9 +184,7 @@ def accept_and_finalize( block_verification=self._block_verification, ) if simulate: - correct_len = self._simulated_correct_len( - bs=bs, dtype=correct_len.dtype, device=correct_len.device - ) + correct_len = self._simulated_correct_len.repeat(bs).to(correct_len.dtype) if simulate_bonus_sampling_info is not None: bonus = sample_simulated_bonus( target_logits=target_logits, @@ -226,23 +225,15 @@ def accept_and_finalize( out_tokens=out_tokens, ) - def _simulated_correct_len( - self, *, bs: int, dtype: torch.dtype, device: torch.device - ) -> torch.Tensor: - buf = self._simulated_correct_drafts_buf - if ( - buf is None - or buf.numel() < bs - or buf.dtype != dtype - or buf.device != device - ): - buf = torch.empty((max(bs, 512),), dtype=dtype, device=device) - self._simulated_correct_drafts_buf = buf - + def begin_simulated_step(self) -> None: + """Draw this step's SGLANG_SIMULATE_ACC_LEN correct length, before the + target verify, so the folded accept in the graph applies it too.""" + if self._simulated_correct_len is None: + return simulated_acc_len = sample_simulated_acc_len( self._simulate_acc_len, SIMULATE_ACC_METHOD, self.gamma + 1 ) - return buf[:bs].fill_(simulated_acc_len - 1) + self._simulated_correct_len.fill_(simulated_acc_len - 1) def run_idle_participation( self, @@ -569,6 +560,7 @@ def __init__( tp_sync: SpecTpSync, commit_ctx: Optional[CommitInjectCtx] = None, fused_argmax: bool = False, + simulated_correct_len: Optional[torch.Tensor] = None, ) -> None: self.max_bs = int(max_bs) self.stride = int(verify_num_draft_tokens) @@ -599,6 +591,8 @@ def __init__( self.out_tokens_buf = torch.zeros( (self.max_bs, self.stride), dtype=torch.int64, device=device ) + # SGLANG_SIMULATE_ACC_LEN: the step's correct length, filled per step. + self.simulated_correct_len = simulated_correct_len self.strided_logits: Optional[torch.Tensor] = None self.strided_hidden: Optional[torch.Tensor] = None self._static_step_state: Optional[tuple[int, bool]] = None @@ -781,6 +775,9 @@ def _accept( cutoff_verify_lens=cutoff_verify_lens, fused_argmax=self._fused_argmax, ) + if self.simulated_correct_len is not None: + # The accept runs as served; only its length is replaced. + correct_len = self.simulated_correct_len.repeat(bs).to(correct_len.dtype) 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) 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 843feb839f4e..b785bf44e948 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -316,6 +316,13 @@ def __init__( verify_num_draft_tokens=self.verify_num_draft_tokens, block_pos_offsets=self._block_pos_offsets, ) + self._simulate_acc_len = float(envs.SGLANG_SIMULATE_ACC_LEN.get()) + # The step's simulated correct length, read by the eager and graph accepts. + self._simulated_correct_len = ( + torch.zeros((1,), dtype=torch.int64, device=self.device) + if self._simulate_acc_len > 0 + else None + ) self._greedy_draft = envs.SGLANG_DSPARK_USE_GREEDY_DRAFT.get() self._proposer = DraftBlockProposer( draft_model=self.draft_model, @@ -357,6 +364,7 @@ def __init__( device=self.device, tp_sync=self._tp_sync, fused_argmax=target_is_dsv41, + simulated_correct_len=self._simulated_correct_len, commit_ctx=CommitInjectCtx( draft_model=self.draft_model, block_pos_offsets=self._block_pos_offsets, @@ -371,7 +379,6 @@ def __init__( self._verify_epilogue.capture_hook ) - self._simulate_acc_len = float(envs.SGLANG_SIMULATE_ACC_LEN.get()) self._simulate_acc_greedy = ( _is_hip and self._simulate_acc_len > 0 @@ -404,6 +411,7 @@ def __init__( tp_sync=self._tp_sync, verify_epilogue=self._verify_epilogue, simulate_acc_len=self._simulate_acc_len, + simulated_correct_len=self._simulated_correct_len, block_verification=get_spec().speculative_use_block_verification, ) @@ -936,9 +944,11 @@ def _forward_decode( # draft proposal itself folded. and (sampling_info is None or sampling_info.is_all_greedy) and verify_logits_adjustments_are_noop(sampling_info) - and self._simulate_acc_len <= 0 + # The graph accept is greedy, so a sampled simulated bonus is eager. + and simulate_bonus_sampling_info is None and not batch.has_grammar ) + self._verify_executor.begin_simulated_step() prepare_mamba_track_for_verify(batch) with self._observers.segment(InfoSegment.TARGET_VERIFY): if run_compact: