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

docs_new/docs/references/environment_variables.mdx

@@ -1679,11 +1679,6 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Plan the next step on a separate stream to overlap with the current step (Overlap Spec V2).</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
-    <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable strict filter checks in speculative decoding.</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
-    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Skip draft_extend while adaptive spec is at steps=0; saves a draft forward but the draft KV goes stale.</td>

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(
+        sorted_indices = self._get_decode_retraction_order(self.reqs, server_args)
-            self.reqs,
-            server_args,
-            allow_policy_sort=(
-                self.spec_algorithm is None or self.spec_algorithm.is_none()
-            ),
-        )
 
         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()
+        self.topk_p = self.topk_p[new_indices]
-        if has_been_filtered:
+        self.topk_index = self.topk_index[new_indices]
-            # in eagle_utils.py:verify, we have already filtered the batch by `unfinished_index`
+        if self.draft_probs is not None:
-            # therefore, we don't need to filter the batch again in scheduler
+            self.draft_probs = self.draft_probs[new_indices]
-            error_msg = f"length of new_indices: {len(new_indices)} != length of topk_p: {len(self.topk_p)}, this should not happen"
+        if self.hidden_states is not None:
-            if len(new_indices) != len(self.topk_p):
+            self.hidden_states = self.hidden_states[new_indices]
-                if strict_check:
+        self.bonus_tokens = self.bonus_tokens[new_indices]
-                    raise ValueError(error_msg)
+        if self.dsa_topk_indices is not None:
-                else:
+            self.dsa_topk_indices = self.dsa_topk_indices[new_indices]
-                    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]
 
     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:

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

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)