[DSpark] Apply SGLANG_SIMULATE_ACC_LEN inside the accept instead of switching paths

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)

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: