From 8c6b5ea73c528c86b665ec7bd22f16ae156dff22 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Sat, 10 Oct 2026 16:30:57 -0700 Subject: [PATCH 1/2] apply simulated accept lengths inside the dspark accept; keep simulate on the served path --- .../dspark_components/dspark_verify.py | 38 +++++++++++++++---- .../dspark_components/dspark_worker_v2.py | 7 +++- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 3eb1495592cc..fad9a2662fdd 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -140,6 +140,7 @@ def __init__( self._simulate_acc_len = float(simulate_acc_len) self._block_verification = block_verification self._simulated_correct_drafts_buf: Optional[torch.Tensor] = None + self._step_simulated_correct_len = 0 def accept_and_finalize( self, @@ -160,8 +161,8 @@ 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) @@ -226,6 +227,20 @@ def accept_and_finalize( out_tokens=out_tokens, ) + 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._simulate_acc_len <= 0: + return + simulated_acc_len = sample_simulated_acc_len( + self._simulate_acc_len, SIMULATE_ACC_METHOD, self.gamma + 1 + ) + self._step_simulated_correct_len = simulated_acc_len - 1 + if self.verify_epilogue is not None: + self.verify_epilogue.stage_simulated_correct_len( + self._step_simulated_correct_len + ) + def _simulated_correct_len( self, *, bs: int, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: @@ -238,11 +253,7 @@ def _simulated_correct_len( ): buf = torch.empty((max(bs, 512),), dtype=dtype, device=device) self._simulated_correct_drafts_buf = buf - - 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) + return buf[:bs].fill_(self._step_simulated_correct_len) def run_idle_participation( self, @@ -569,6 +580,7 @@ def __init__( tp_sync: SpecTpSync, commit_ctx: Optional[CommitInjectCtx] = None, fused_argmax: bool = False, + simulate: bool = False, ) -> None: self.max_bs = int(max_bs) self.stride = int(verify_num_draft_tokens) @@ -599,6 +611,12 @@ 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, staged per replay. + self.simulated_correct_len_buf = ( + torch.zeros((self.max_bs,), dtype=torch.int64, device=device) + if simulate + else None + ) self.strided_logits: Optional[torch.Tensor] = None self.strided_hidden: Optional[torch.Tensor] = None self._static_step_state: Optional[tuple[int, bool]] = None @@ -638,6 +656,9 @@ def begin_step(self, verify_lens, armed: bool) -> None: self.verify_lens_buf[bs:].zero_() self.inject_gate_buf.fill_(1 if armed else 0) + def stage_simulated_correct_len(self, correct_len: int) -> None: + self.simulated_correct_len_buf.fill_(correct_len) + def begin_static_step(self, bs: int, armed: bool) -> None: state = (bs, armed) if self._static_step_state == state: @@ -781,6 +802,9 @@ def _accept( cutoff_verify_lens=cutoff_verify_lens, fused_argmax=self._fused_argmax, ) + if self.simulated_correct_len_buf is not None: + # The accept runs as served; only its length is replaced. + correct_len = self.simulated_correct_len_buf[: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 fda170cf1581..9f4614353e7e 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -313,6 +313,7 @@ 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()) self._proposer = DraftBlockProposer( draft_model=self.draft_model, draft_model_runner=self.draft_model_runner, @@ -351,6 +352,7 @@ def __init__( device=self.device, tp_sync=self._tp_sync, fused_argmax=target_is_dsv41, + simulate=self._simulate_acc_len > 0, commit_ctx=CommitInjectCtx( draft_model=self.draft_model, block_pos_offsets=self._block_pos_offsets, @@ -365,7 +367,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 @@ -931,9 +932,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: From d746194678ceea26a8ea6de1e62934a7453995fd Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Sun, 11 Oct 2026 15:24:04 -0700 Subject: [PATCH 2/2] share one simulated correct-len tensor between the eager and graph accepts --- .../dspark_components/dspark_verify.py | 49 +++++-------------- .../dspark_components/dspark_worker_v2.py | 9 +++- 2 files changed, 19 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index bf68ebf5d189..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,8 +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._step_simulated_correct_len = 0 + self._simulated_correct_len = simulated_correct_len def accept_and_finalize( self, @@ -167,7 +167,7 @@ def accept_and_finalize( 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 @@ -184,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, @@ -230,30 +228,12 @@ def accept_and_finalize( 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._simulate_acc_len <= 0: + 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 ) - self._step_simulated_correct_len = simulated_acc_len - 1 - if self.verify_epilogue is not None: - self.verify_epilogue.stage_simulated_correct_len( - self._step_simulated_correct_len - ) - - 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 - return buf[:bs].fill_(self._step_simulated_correct_len) + self._simulated_correct_len.fill_(simulated_acc_len - 1) def run_idle_participation( self, @@ -580,7 +560,7 @@ def __init__( tp_sync: SpecTpSync, commit_ctx: Optional[CommitInjectCtx] = None, fused_argmax: bool = False, - simulate: bool = False, + simulated_correct_len: Optional[torch.Tensor] = None, ) -> None: self.max_bs = int(max_bs) self.stride = int(verify_num_draft_tokens) @@ -611,12 +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, staged per replay. - self.simulated_correct_len_buf = ( - torch.zeros((self.max_bs,), dtype=torch.int64, device=device) - if simulate - else None - ) + # 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 @@ -656,9 +632,6 @@ def begin_step(self, verify_lens, armed: bool) -> None: self.verify_lens_buf[bs:].zero_() self.inject_gate_buf.fill_(1 if armed else 0) - def stage_simulated_correct_len(self, correct_len: int) -> None: - self.simulated_correct_len_buf.fill_(correct_len) - def begin_static_step(self, bs: int, armed: bool) -> None: state = (bs, armed) if self._static_step_state == state: @@ -802,9 +775,9 @@ def _accept( cutoff_verify_lens=cutoff_verify_lens, fused_argmax=self._fused_argmax, ) - if self.simulated_correct_len_buf is not None: + if self.simulated_correct_len is not None: # The accept runs as served; only its length is replaced. - correct_len = self.simulated_correct_len_buf[:bs].to(correct_len.dtype) + 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 eb82a2fb58a0..b785bf44e948 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -317,6 +317,12 @@ def __init__( 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, @@ -358,7 +364,7 @@ def __init__( device=self.device, tp_sync=self._tp_sync, fused_argmax=target_is_dsv41, - simulate=self._simulate_acc_len > 0, + simulated_correct_len=self._simulated_correct_len, commit_ctx=CommitInjectCtx( draft_model=self.draft_model, block_pos_offsets=self._block_pos_offsets, @@ -405,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, )