Fix fractional simulated acceptance in DSpark (#33463)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

python/sglang/srt/speculative/dflash_utils.py

@@ -12,7 +12,7 @@ import torch.nn.functional as F
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.sampler import apply_custom_logit_processor
 from sglang.srt.managers.schedule_batch import Req
-from sglang.srt.speculative.spec_utils import _sample_simulated_acc_len
+from sglang.srt.speculative.spec_utils import sample_simulated_acc_len
 from sglang.srt.utils import is_cuda, is_hip, is_musa
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
@@ -611,8 +611,8 @@ def apply_dflash_simulated_acceptance(
     """Forces the DFlash acceptance length (SGLANG_SIMULATE_ACC_LEN benchmark knob)."""
     block_size = candidates.shape[1]
 
-    # _sample_simulated_acc_len clamps to [1, block_size].
-    forced_commit_len = _sample_simulated_acc_len(
+    # sample_simulated_acc_len clamps to [1, block_size].
+    forced_commit_len = sample_simulated_acc_len(
         simulate_acc_len, simulate_acc_method, block_size
     )
     forced_accept_len = forced_commit_len - 1

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

@@ -37,6 +37,10 @@ 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_utils import (
+    SIMULATE_ACC_METHOD,
+    sample_simulated_acc_len,
+)
 from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect
 
 # Draft proposal probs feeding rejection sampling; the data layer is the
@@ -152,15 +156,19 @@ class TargetVerifyExecutor:
         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:
-            correct_target = int(
-                round(min(max(self._simulate_acc_len - 1.0, 0.0), float(self.gamma)))
-            )
-            buf = torch.full(
-                (max(bs, 512),), correct_target, dtype=dtype, device=device
-            )
+        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]
+
+        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)
 
     def run_idle_participation(
         self,

python/sglang/srt/speculative/spec_utils.py

@@ -348,7 +348,7 @@ def select_top_k_tokens(
     )
 
 
-def _sample_simulated_acc_len(
+def sample_simulated_acc_len(
     simulate_acc_len: float,
     simulate_acc_method: str,
     max_len: int,
@@ -401,7 +401,7 @@ def generate_simulated_accept_index(
     use_real_draft_tokens = simulate_acc_token_mode == "real-draft-token"
 
     assert simulate_acc_len > 0.0
-    simulate_acc_len = _sample_simulated_acc_len(
+    simulate_acc_len = sample_simulated_acc_len(
         simulate_acc_len, simulate_acc_method, spec_steps + 1
     )