diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index fa929fd38..325507926 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 - room for the decode phase. + chunk N+1 allocation plus decode headroom: + + 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: - """SWA tokens a request needs regardless of extend length: the sliding - window (decode headroom) + allocator page slack + the load-back window - charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap.""" - reserved = self.tree_cache.sliding_window_size + self.page_size + def _swa_reserved_tokens( + self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 + ) -> int: + """SWA slots a request adds to its own sliding window + page slack + the + 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) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 6cbd41a6f..c88075581 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/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) - (64, None, 128, 16, 128 + 16, "no_cap_floor_active"), - (200, None, 256, 32, 256 + 32, "no_cap_floor_active_other_dims"), - (300, None, 128, 16, 300 + 16, "no_cap_floor_inactive"), - (200, 50, 64, 8, 64 + 8, "cap_binds_then_floor"), - (300, 500, 64, 64, 300 + 64, "cap_does_not_bind"), - (0, None, 128, 16, 128 + 16, "extend_zero_floor_only"), + # (extend, max_new, rem_chunk, window, page, expected, label) + (64, 512, None, 128, 16, 128 + 16, "long_decode_hits_window_floor"), + (64, 32, None, 128, 16, 96 + 16, "short_req_reserves_below_window"), + (10, 20, None, 512, 8, 30 + 8, "short_resume_tiny_budget"), + (300, 512, None, 128, 16, 300 + 16, "big_extend_over_window"), + (200, 512, 50, 64, 8, 64 + 8, "chunk_capped_alloc_hits_floor"), + # 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