diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 0e6d0f214..fa929fd38 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -845,6 +845,17 @@ class PrefillAdder: return req _rem_tokens = self.rem_chunk_tokens + # A mid-chunk rank prefills this pass regardless of the delayer + # verdict, so report prefillable=True and ignore the result. + if self.prefill_delayer_single_pass is not None: + self.prefill_delayer_single_pass.negotiate_should_allow_prefill( + local_prefillable=True, + running_batch=self.running_batch.batch_size(), + max_prefill_bs=self.max_prefill_bs, + max_running_requests=self.max_running_requests, + waiting_queue_len=self.waiting_queue_len, + ) + cand_extend_input_len = len(req.full_untruncated_fill_ids) - len( req.prefix_indices ) @@ -1001,16 +1012,6 @@ class PrefillAdder: def add_one_req( self, req: Req, has_chunked_req: bool, truncation_align_size: Optional[int] ): - if (self.prefill_delayer_single_pass is not None) and ( - not self.prefill_delayer_single_pass.negotiate_should_allow_prefill( - local_prefillable=True, - running_batch=self.running_batch.batch_size(), - max_prefill_bs=self.max_prefill_bs, - max_running_requests=self.max_running_requests, - waiting_queue_len=self.waiting_queue_len, - ) - ): - return AddReqResult.OTHER # TODO support cp with multiple requests # Enabling context parallelism currently presents precision issues; # therefore, the prefill-batch setting is temporarily set to 1. @@ -1086,6 +1087,20 @@ class PrefillAdder: return AddReqResult.NO_TOKEN chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap) + # Negotiate only after every KV-budget gate (a NO_TOKEN rank must + # report not-prefillable via finalize()) and before init_load_back + # (a delay verdict must not start KV load-back). + if (self.prefill_delayer_single_pass is not None) and ( + not self.prefill_delayer_single_pass.negotiate_should_allow_prefill( + local_prefillable=True, + running_batch=self.running_batch.batch_size(), + max_prefill_bs=self.max_prefill_bs, + max_running_requests=self.max_running_requests, + waiting_queue_len=self.waiting_queue_len, + ) + ): + return AddReqResult.OTHER + if req.needs_host_load_back(): new_indices, req.last_node = self.tree_cache.init_load_back( InitLoadBackParams( diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 501a65d5e..6cbd41a6f 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -22,6 +22,20 @@ register_amd_ci(est_time=2, suite="stage-b-test-1-gpu-small-amd") register_cpu_ci(est_time=8, suite="base-c-test-cpu") +class _RecordingDelayer: + """Duck-typed stand-in for PrefillDelayerSinglePassExecutor that records + the local_prefillable value of every negotiate call and returns a fixed + verdict.""" + + def __init__(self, allow: bool): + self.allow = allow + self.calls = [] + + def negotiate_should_allow_prefill(self, local_prefillable, **kwargs): + self.calls.append(local_prefillable) + return self.allow + + class TestPrefillAdder(CustomTestCase): def setUp(self): set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) @@ -537,6 +551,131 @@ class TestPrefillAdder(CustomTestCase): ) self.assertEqual(adder._swa_budget_for_req(extend), expected) + def test_delayer_not_consulted_when_kv_budget_rejects(self): + """A rank whose first candidate fails the KV-budget gate must NOT + report local_prefillable=True: add_one_req returns NO_TOKEN before + negotiating, and finalize() later reports the rank as not + prefillable. Regression guard: the negotiate used to run at the top + of add_one_req, so under KV pressure a full rank still claimed + prefillable=True, the delayer saw "all prefillable" and allowed, the + full ranks then NO_TOKEN'ed out, and the DP-synced forward mixed + prefill and decode — the exact pattern the delayer exists to + prevent.""" + delayer = _RecordingDelayer(allow=True) + adder = self._create_delayer_adder(available_tokens=10, delayer=delayer) + + result = adder.add_one_req( + self._create_delayer_req(50), + has_chunked_req=False, + truncation_align_size=None, + ) + + self.assertEqual(result, AddReqResult.NO_TOKEN) + self.assertEqual(delayer.calls, []) + self.assertEqual(adder.can_run_list, []) + + def test_delayer_not_consulted_when_post_lock_recheck_rejects(self): + """Locking the request's own prefix converts evictable tokens into + protected ones, so a request can pass the pre-lock KV gate yet fail + the post-lock recheck — precisely the high-utilization regime the + delayer targets. The negotiate must sit after that recheck too, or + the rank claims prefillable=True and then runs decode.""" + delayer = _RecordingDelayer(allow=True) + adder = self._create_delayer_adder(available_tokens=10, delayer=delayer) + # Pre-lock budget is covered by evictable tokens; inc_lock_ref pins + # them (evictable -> protected), shrinking the budget below demand. + self.mock_tree_cache.evictable_size.return_value = 1000 + self.mock_tree_cache.full_evictable_size.return_value = 1000 + + def _pin_prefix(node): + self.mock_tree_cache.evictable_size.return_value = 0 + self.mock_tree_cache.full_evictable_size.return_value = 0 + return IncLockRefResult() + + self.mock_tree_cache.inc_lock_ref.side_effect = _pin_prefix + + result = adder.add_one_req( + self._create_delayer_req(50), + has_chunked_req=False, + truncation_align_size=None, + ) + + self.assertEqual(result, AddReqResult.NO_TOKEN) + self.assertEqual(delayer.calls, []) + self.assertEqual(adder.can_run_list, []) + + def test_delay_verdict_blocks_admissible_request(self): + """An admissible request must still be gated by the (relocated) + negotiate: on a delay verdict it is not admitted, and the rank + reported prefillable=True exactly once.""" + delayer = _RecordingDelayer(allow=False) + adder = self._create_delayer_adder(available_tokens=100_000, delayer=delayer) + + result = adder.add_one_req( + self._create_delayer_req(50), + has_chunked_req=False, + truncation_align_size=None, + ) + + self.assertEqual(result, AddReqResult.OTHER) + self.assertEqual(delayer.calls, [True]) + self.assertEqual(adder.can_run_list, []) + + def test_allow_verdict_admits_request(self): + """On an allow verdict the request proceeds through admission — the + relocated negotiate must not block the commit path.""" + delayer = _RecordingDelayer(allow=True) + adder = self._create_delayer_adder(available_tokens=100_000, delayer=delayer) + req = self._create_delayer_req(50) + + result = adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + + self.assertEqual(result, AddReqResult.CONTINUE) + self.assertEqual(delayer.calls, [True]) + self.assertIn(req, adder.can_run_list) + + def test_chunked_req_negotiates_prefillable_and_proceeds(self): + """A rank resuming a chunked prefill runs it this pass regardless of + the verdict, so add_chunked_req must report prefillable=True (else a + rank with an empty waiting queue reports False via finalize() and + peers delay while it prefills alone) and must not drop the chunk on + a delay verdict (that would leak memory).""" + delayer = _RecordingDelayer(allow=False) + adder = self._create_delayer_adder( + available_tokens=100_000, delayer=delayer, rem_chunk_tokens=500 + ) + req = self._create_delayer_req(200) + + result = adder.add_chunked_req(req) + + self.assertIsNone(result) # chunk fully admitted, not truncated + self.assertEqual(delayer.calls, [True]) + self.assertIn(req, adder.can_run_list) + + def _create_delayer_adder(self, *, available_tokens, delayer, **kwargs): + self.mock_token_allocator.available_size.return_value = available_tokens + self.mock_token_allocator.full_available_size.return_value = available_tokens + return self.create_adder( + self.create_running_batch(), + prefill_delayer_single_pass=delayer, + **kwargs, + ) + + def _create_delayer_req(self, num_tokens: int): + req = self.create_mock_req("delayer_req", priority=0, max_new_tokens=8) + req.full_untruncated_fill_ids = list(range(num_tokens)) + req.host_hit_length = 0 + req.last_node = MagicMock() + req.sampling_params.ignore_eos = False + req.set_extend_range = MagicMock( + side_effect=lambda start, end: setattr( + req, "extend_range", Range(start, end) + ) + ) + return req + def test_add_chunked_req_non_hybrid_no_swa_reservation(self): # Non-hybrid path: the SWA-pool reservation must NOT apply, otherwise # the fix would regress non-SWA models.