diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 275c32638..89870857f 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -1679,11 +1679,6 @@ SGLang supports various environment variables that can be used to configure its Plan the next step on a separate stream to overlap with the current step (Overlap Spec V2). false - - SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK - Enable strict filter checks in speculative decoding. - true - SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND Skip draft_extend while adaptive spec is at steps=0; saves a draft forward but the draft KV goes stale. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8bc444d8d..20d47e37e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -764,7 +764,6 @@ class Envs: SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False) # Spec Config - SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) # A/B: keep the DFLASH draft greedy head eager (not folded in-graph). SGLANG_DFLASH_EAGER_DRAFT_SAMPLER = EnvBool(False) SGLANG_RAGGED_VERIFY_MODE = EnvStr("static") diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 08a837833..cc02f4a61 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2624,13 +2624,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self, server_args: ServerArgs ) -> Tuple[List[Req], float, List[Req]]: """Retract the decoding requests when there is not enough memory.""" - sorted_indices = self._get_decode_retraction_order( - self.reqs, - server_args, - allow_policy_sort=( - self.spec_algorithm is None or self.spec_algorithm.is_none() - ), - ) + sorted_indices = self._get_decode_retraction_order(self.reqs, server_args) retracted_reqs = [] first_iter = True @@ -2678,7 +2672,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): @staticmethod def _get_decode_retraction_order( - reqs: List[Req], server_args: ServerArgs, *, allow_policy_sort: bool + reqs: List[Req], server_args: ServerArgs ) -> List[int]: """Return indices ordered from most-preferred to least-preferred to keep. @@ -2688,11 +2682,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): sorted_indices = list(range(len(reqs))) # TODO(lsyin): improve retraction policy for radix cache - # For spec decoding, filter_batch API can only filter requests from the - # back, so we can only retract from the back. - # TODO(sang): Clean up finish path and support better retract policy. - if not allow_policy_sort: - return sorted_indices def length_key(req: Req) -> Tuple[int, int]: return (len(req.output_ids), -len(req.origin_input_ids)) @@ -2993,7 +2982,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): if self.spec_info: self.spec_info.filter_batch( new_indices=keep_indices_device, - has_been_filtered=False, new_indices_cpu=keep_indices, ) diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 8893a6541..72162ba51 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -215,7 +215,6 @@ class DFlashDraftInputV2(SpecInput): def filter_batch( self, new_indices: torch.Tensor, - has_been_filtered: bool = True, new_indices_cpu: Optional[List[int]] = None, ): if self.reserved_seq_lens_cpu is not None: diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 60c702d19..e7b724165 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -6,7 +6,6 @@ import torch from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject -from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode from sglang.srt.runtime_context import get_server_args from sglang.srt.speculative.spec_info import SpecInput, SpecInputType @@ -211,44 +210,21 @@ class EagleDraftInput(SpecInput): def filter_batch( self, new_indices: torch.Tensor, - has_been_filtered: bool = True, new_indices_cpu: Optional[List[int]] = None, ): if self.future_indices is not None: self.future_indices = self.future_indices[new_indices] return - strict_check = envs.SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK.get() - if has_been_filtered: - # in eagle_utils.py:verify, we have already filtered the batch by `unfinished_index` - # therefore, we don't need to filter the batch again in scheduler - error_msg = f"length of new_indices: {len(new_indices)} != length of topk_p: {len(self.topk_p)}, this should not happen" - if len(new_indices) != len(self.topk_p): - if strict_check: - raise ValueError(error_msg) - else: - logger.warning(error_msg) - - self.topk_p = self.topk_p[: len(new_indices)] - self.topk_index = self.topk_index[: len(new_indices)] - if self.draft_probs is not None: - self.draft_probs = self.draft_probs[: len(new_indices)] - if self.hidden_states is not None: - self.hidden_states = self.hidden_states[: len(new_indices)] - self.bonus_tokens = self.bonus_tokens[: len(new_indices)] - if self.dsa_topk_indices is not None: - self.dsa_topk_indices = self.dsa_topk_indices[: len(new_indices)] - else: - # in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index` - self.topk_p = self.topk_p[new_indices] - self.topk_index = self.topk_index[new_indices] - if self.draft_probs is not None: - self.draft_probs = self.draft_probs[new_indices] - if self.hidden_states is not None: - self.hidden_states = self.hidden_states[new_indices] - self.bonus_tokens = self.bonus_tokens[new_indices] - if self.dsa_topk_indices is not None: - self.dsa_topk_indices = self.dsa_topk_indices[new_indices] + self.topk_p = self.topk_p[new_indices] + self.topk_index = self.topk_index[new_indices] + if self.draft_probs is not None: + self.draft_probs = self.draft_probs[new_indices] + if self.hidden_states is not None: + self.hidden_states = self.hidden_states[new_indices] + self.bonus_tokens = self.bonus_tokens[new_indices] + if self.dsa_topk_indices is not None: + self.dsa_topk_indices = self.dsa_topk_indices[new_indices] def merge_batch(self, spec_info: "EagleDraftInput"): if self.future_indices is not None: diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 719ac050b..11d63fb9e 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -119,7 +119,6 @@ class NgramVerifyInput(SpecInput): def filter_batch( self, new_indices: torch.Tensor, - has_been_filtered: bool = True, new_indices_cpu: Optional[List[int]] = None, ): if self.future_indices is not None: diff --git a/test/registered/unit/managers/test_retraction_order.py b/test/registered/unit/managers/test_retraction_order.py new file mode 100644 index 000000000..d23b0d523 --- /dev/null +++ b/test/registered/unit/managers/test_retraction_order.py @@ -0,0 +1,59 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _req(output_len: int, input_len: int = 8, priority=None): + return SimpleNamespace( + output_ids=[0] * output_len, + origin_input_ids=[0] * input_len, + priority=priority, + ) + + +def _args(policy: str = "length", low_first: bool = False): + return SimpleNamespace( + retraction_policy=policy, + schedule_low_priority_values_first=low_first, + ) + + +def _order(reqs, args): + return ScheduleBatch._get_decode_retraction_order(reqs, args) + + +class TestRetractionOrder(CustomTestCase): + """The retraction loop pops from the END of the returned list, so the + last index is the first request retracted.""" + + def test_length_policy_retracts_shortest_output_first(self): + reqs = [_req(5), _req(1), _req(3)] + self.assertEqual(_order(reqs, _args()), [0, 2, 1]) + + def test_length_policy_tie_breaks_on_longer_input(self): + # Equal outputs: the longer-input request is retracted first + # (frees more tokens for the same rework). + reqs = [_req(4, input_len=10), _req(4, input_len=20)] + self.assertEqual(_order(reqs, _args()), [0, 1]) + + def test_priority_policy_low_values_first(self): + # Low value = more important; None sorts as least important. + reqs = [_req(4, priority=2), _req(4, priority=0), _req(4, priority=None)] + self.assertEqual(_order(reqs, _args("priority", low_first=True)), [1, 0, 2]) + + def test_priority_policy_high_values_first(self): + reqs = [_req(4, priority=2), _req(4, priority=0), _req(4, priority=None)] + self.assertEqual(_order(reqs, _args("priority", low_first=False)), [0, 1, 2]) + + def test_priority_ties_fall_back_to_length(self): + reqs = [_req(1, priority=1), _req(5, priority=1)] + self.assertEqual(_order(reqs, _args("priority", low_first=True)), [1, 0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_dflash_overlap_hostsync.py b/test/registered/unit/spec/test_dflash_overlap_hostsync.py index 099fa969a..f7f24f274 100644 --- a/test/registered/unit/spec/test_dflash_overlap_hostsync.py +++ b/test/registered/unit/spec/test_dflash_overlap_hostsync.py @@ -255,10 +255,9 @@ class TestFilterBatchHostIndices(CustomTestCase): keep = [0, 2] a, b = make(), make() - a.filter_batch(new_indices=torch.tensor(keep), has_been_filtered=False) + a.filter_batch(new_indices=torch.tensor(keep)) b.filter_batch( new_indices=torch.tensor(keep), - has_been_filtered=False, new_indices_cpu=keep, ) torch.testing.assert_close(a.reserved_seq_lens_cpu, b.reserved_seq_lens_cpu)