Fix SWA admission livelock on cached-prefix resumes (#32379)

python/sglang/srt/managers/schedule_policy.py

@@ -620,7 +620,7 @@ class PrefillAdder:
         return available_and_evictable - self.cur_rem_token_offset
 
     def _swa_budget_for_req(
-        self, extend_input_len: int, swa_host_hit_length: int = 0
+        self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
     ) -> int:
         """SWA pool budget per request. Only valid when is_hybrid_swa is True.
 
@@ -629,26 +629,52 @@ class PrefillAdder:
           + chunk N+1 (new allocation)
         Since chunk N and locked tokens are already excluded from
         swa_available + swa_evictable, the budget only needs to cover the
-        chunk N+1 allocation. We floor at sliding_window_size to reserve
+        chunk N+1 allocation plus decode headroom:
-        room for the decode phase.
+
+          budget = max(alloc - window, 0) + min(extend + max_new_tokens, window) + page
+
+        where alloc = min(extend, rem_chunk); the min() cap keeps the two terms
+        from double-counting extend, so budget <= extend + max_new_tokens + page.
         """
         if self.rem_chunk_tokens is not None:
             alloc = min(extend_input_len, self.rem_chunk_tokens)
         else:
             alloc = extend_input_len
         window = self.tree_cache.sliding_window_size
-        return max(alloc - window, 0) + self._swa_reserved_tokens(swa_host_hit_length)
+        return max(alloc - window, 0) + self._swa_reserved_tokens(
+            extend_input_len, max_new_tokens, swa_host_hit_length
+        )
 
-    def _swa_reserved_tokens(self, swa_host_hit_length: int = 0) -> int:
+    def _swa_reserved_tokens(
-        """SWA tokens a request needs regardless of extend length: the sliding
+        self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
-        window (decode headroom) + allocator page slack + the load-back window
+    ) -> int:
-        charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap."""
+        """SWA slots a request adds to its own sliding window + page slack + the
-        reserved = self.tree_cache.sliding_window_size + self.page_size
+        load-back charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap.
+
+        The headroom is min(extend + decode, window), not a constant window: a
+        request contributes only extend + decode fresh tokens to its window and
+        a cached SWA prefix funds the rest. Charging a full window double-counted
+        a short cached-prefix resume and livelocked admission at a ~2-window
+        pool; keeping extend in the min() holds the reservation >= the prefill
+        allocation so an admitted request cannot OOM."""
+        window = self.tree_cache.sliding_window_size
+        headroom = min(extend_input_len + max_new_tokens, window)
+        reserved = headroom + self.page_size
         if swa_host_hit_length > 0:
             reserved += self.ceil_paged_tokens(swa_host_hit_length)
         return reserved
 
-    def _swa_chunk_cap(self, swa_host_hit_length: int = 0) -> int:
+    def _swa_new_tokens(self, req: Req) -> int:
+        """Tokens a request may still decode, for SWA headroom sizing. Mirrors
+        the remaining-then-clip order of add_one_req's max_new: clip-then-subtract
+        would zero out a request that has already generated >= CLIP tokens but
+        still has a long decode ahead, under-reserving its window."""
+        return min(
+            max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
+            CLIP_MAX_NEW_TOKENS,
+        )
+
+    def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int:
         """Largest page-aligned extend chunk the SWA pool can admit right now,
         keeping a sliding window of headroom below rem_swa_tokens; 0 if not
         even one page fits. Only valid when is_hybrid_swa is True.
@@ -659,7 +685,11 @@ class PrefillAdder:
         forever (head-of-line livelock). Shrinking is sound because past a
         chunk boundary only the sliding window stays locked — the rest turns
         evictable — so each pass's transient footprint fits the pool."""
-        cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(swa_host_hit_length)
+        # extend_input_len=0: this solves for the extend chunk itself, so the
+        # reserved headroom is the post-chunk decode window only.
+        cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(
+            0, max_new_tokens, swa_host_hit_length
+        )
         if cap <= 0:
             return 0
         return cap // self.page_size * self.page_size
@@ -735,7 +765,9 @@ class PrefillAdder:
         self.rem_input_tokens -= extend_input_len
 
         if self.is_hybrid_swa:
-            self.rem_swa_token_offset += self._swa_budget_for_req(extend_input_len)
+            self.rem_swa_token_offset += self._swa_budget_for_req(
+                extend_input_len, max_new_tokens
+            )
 
         if self.dllm_config is not None:
             self.rem_dllm_tokens -= extend_input_len
@@ -906,7 +938,12 @@ class PrefillAdder:
         if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens):
             return AddReqResult.NO_TOKEN
         if self.is_hybrid_swa:
-            if self._swa_budget_for_req(cand_extend_input_len) > self.rem_swa_tokens:
+            if (
+                self._swa_budget_for_req(
+                    cand_extend_input_len, self._swa_new_tokens(req)
+                )
+                > self.rem_swa_tokens
+            ):
                 return AddReqResult.NO_TOKEN
 
         def add_req_state(r, insert_sort=False):
@@ -1053,10 +1090,14 @@ class PrefillAdder:
             # driven only by the freshly-prefilled tail (the loaded window is
             # charged separately via swa_host_hit_length).
             swa_needed = self._swa_budget_for_req(
-                real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
+                real_input_tokens,
+                self._swa_new_tokens(req),
+                swa_host_hit_length=req.swa_host_hit_length,
             )
             if swa_needed >= self.rem_swa_tokens:
-                swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
+                swa_cap = self._swa_chunk_cap(
+                    self._swa_new_tokens(req), req.swa_host_hit_length
+                )
                 if self.rem_chunk_tokens is None or swa_cap <= 0:
                     return AddReqResult.NO_TOKEN
                 chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
@@ -1079,10 +1120,14 @@ class PrefillAdder:
             if self.is_hybrid_swa:
                 # self.rem_swa_tokens may decrease after the lock acquisition
                 swa_needed = self._swa_budget_for_req(
-                    real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
+                    real_input_tokens,
+                    self._swa_new_tokens(req),
+                    swa_host_hit_length=req.swa_host_hit_length,
                 )
                 if swa_needed >= self.rem_swa_tokens:
-                    swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
+                    swa_cap = self._swa_chunk_cap(
+                        self._swa_new_tokens(req), req.swa_host_hit_length
+                    )
                     if self.rem_chunk_tokens is None or swa_cap <= 0:
                         return AddReqResult.NO_TOKEN
                     chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)

test/registered/unit/managers/test_prefill_adder.py

@@ -1,6 +1,6 @@
 import unittest
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
@@ -532,16 +532,26 @@ class TestPrefillAdder(CustomTestCase):
         self.assertEqual(len(adder.can_run_list), 0)
 
     def test_swa_budget_for_req(self):
+        # budget = max(alloc - window, 0) + min(extend + max_new, window) + page,
+        # where alloc = min(extend, rem_chunk). The decode headroom is the SWA the
+        # request adds to its own sliding window (extend + decode), capped at the
+        # window -- a cached prefix funds the rest -- rather than a constant
+        # window, which over-charged short requests into an admission livelock.
         cases = [
-            # (extend, rem_chunk, window, page, expected, label)
+            # (extend, max_new, rem_chunk, window, page, expected, label)
-            (64, None, 128, 16, 128 + 16, "no_cap_floor_active"),
+            (64, 512, None, 128, 16, 128 + 16, "long_decode_hits_window_floor"),
-            (200, None, 256, 32, 256 + 32, "no_cap_floor_active_other_dims"),
+            (64, 32, None, 128, 16, 96 + 16, "short_req_reserves_below_window"),
-            (300, None, 128, 16, 300 + 16, "no_cap_floor_inactive"),
+            (10, 20, None, 512, 8, 30 + 8, "short_resume_tiny_budget"),
-            (200, 50, 64, 8, 64 + 8, "cap_binds_then_floor"),
+            (300, 512, None, 128, 16, 300 + 16, "big_extend_over_window"),
-            (300, 500, 64, 64, 300 + 64, "cap_does_not_bind"),
+            (200, 512, 50, 64, 8, 64 + 8, "chunk_capped_alloc_hits_floor"),
-            (0, None, 128, 16, 128 + 16, "extend_zero_floor_only"),
+            # Multi-chunk: alloc = rem_chunk (1024) caps a huge extend, and the
+            # chunk itself exceeds the window, so term1 = alloc - window is driven
+            # by the chunk, not the full extend -> budget = chunk + page.
+            (2000, 256, 1024, 512, 16, 1024 + 16, "multichunk_alloc_over_window"),
+            (0, 512, None, 128, 16, 128 + 16, "extend_zero_long_decode"),
+            (0, 40, None, 128, 16, 40 + 16, "extend_zero_short_decode"),
         ]
-        for extend, rem_chunk, window, page, expected, label in cases:
+        for extend, max_new, rem_chunk, window, page, expected, label in cases:
             with self.subTest(label=label):
                 self.mock_tree_cache.sliding_window_size = window
                 adder = self.create_adder(
@@ -549,7 +559,80 @@ class TestPrefillAdder(CustomTestCase):
                     page_size=page,
                     rem_chunk_tokens=rem_chunk,
                 )
-                self.assertEqual(adder._swa_budget_for_req(extend), expected)
+                self.assertEqual(adder._swa_budget_for_req(extend, max_new), expected)
+
+    def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self):
+        # Livelock regression (real incident). At an SWA pool ~= 2 sliding
+        # windows, a cached-prefix resume matches >= 1 window (locked, excluded
+        # from rem_swa) and has only a short uncached tail + a little decode
+        # left. The pre-fix constant-window reservation charged a second full
+        # window, so the admission gate (swa_needed >= rem_swa_tokens) rejected
+        # it every scheduler iteration while LPM kept it at the queue head --
+        # 100% scheduler CPU, idle GPU. Capping the reservation at
+        # min(extend + decode, window) admits it.
+        WINDOW, PAGE, REM_SWA = 128, 8, 100
+        PREFIX, EXTEND = 200, 16  # cached prefix > window; short uncached tail
+        self.mock_token_allocator.swa_available_size.return_value = REM_SWA
+        self.mock_token_allocator.full_available_size.return_value = 100_000
+        self.mock_token_allocator.available_size.return_value = 100_000
+        self.mock_tree_cache.sliding_window_size = WINDOW
+        self.mock_tree_cache.is_tree_cache.return_value = False
+        adder = self.create_adder(self.create_running_batch(), page_size=PAGE)
+        adder.is_hybrid_swa = True
+
+        req = self.create_mock_req(
+            "resume", priority=0, max_new_tokens=40, output_len=10
+        )
+        req.prefix_indices = list(range(PREFIX))
+        req.full_untruncated_fill_ids = list(range(PREFIX + EXTEND))
+        req.host_hit_length = 0
+        req.swa_host_hit_length = 0
+        req.last_node = MagicMock()
+        req.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req, "extend_range", Range(start, end)
+            )
+        )
+        req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False)
+
+        # Pre-fix: a constant sliding-window reservation rejects the resume.
+        with patch.object(adder, "_swa_reserved_tokens", return_value=WINDOW + PAGE):
+            self.assertIs(
+                adder.add_one_req(
+                    req, has_chunked_req=False, truncation_align_size=None
+                ),
+                AddReqResult.NO_TOKEN,
+            )
+        self.assertEqual(len(adder.can_run_list), 0)
+
+        # Fix: min(extend + decode, window) reservation admits it.
+        adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None)
+        self.assertIn(req, adder.can_run_list)
+
+    def test_swa_new_tokens_clamps_remaining_not_total(self):
+        # Remaining decode headroom must be min(max_new - generated, CLIP)
+        # (subtract-then-clip). The reversed order (clip-then-subtract) zeroes
+        # out a request that has already generated >= CLIP tokens but still has a
+        # long decode ahead, under-reserving its SWA window -> OOM risk on resume
+        # of a long-generation request.
+        from sglang.srt.managers.schedule_policy import CLIP_MAX_NEW_TOKENS as CLIP
+
+        adder = self.create_adder(self.create_running_batch())
+        cases = [
+            # (max_new, generated, expected, label)
+            (100, 10, 90, "below_clip_normal"),
+            (40, 100, 0, "already_finished"),
+            # clip-then-subtract would give 3996; subtract-then-clip caps at CLIP.
+            (CLIP + 6000, 100, CLIP, "long_gen_small_output_caps_at_clip"),
+            # clip-then-subtract would give 0; the long decode still needs CLIP.
+            (CLIP + 6000, CLIP + 100, CLIP, "long_gen_output_over_clip"),
+        ]
+        for max_new, generated, expected, label in cases:
+            with self.subTest(label=label):
+                req = self.create_mock_req(
+                    label, priority=0, max_new_tokens=max_new, output_len=generated
+                )
+                self.assertEqual(adder._swa_new_tokens(req), expected)
 
     def test_delayer_not_consulted_when_kv_budget_rejects(self):
         """A rank whose first candidate fails the KV-budget gate must NOT