diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 649f0dd9a..e30eb759e 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/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 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 4d5478cd5..0581d548f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/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, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 2845a2d6c..bb5aeb94a 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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 )