[Scheduler] Enable decode retraction ordering under speculative decoding (#32023)

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")

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,
             )
 

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:

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:

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: