import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import PrefillAdder from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=1, suite="stage-b-test-small-1-gpu") register_amd_ci(est_time=2, suite="stage-b-test-small-1-gpu-amd") class TestPrefillAdder(CustomTestCase): def setUp(self): self.mock_tree_cache = self.create_tree_cache() self.mock_token_allocator = self.create_token_allocator() patcher = patch( "sglang.srt.managers.schedule_policy.is_nsa_prefill_cp_in_seq_split", return_value=False, ) self.mock_is_nsa = patcher.start() self.addCleanup(patcher.stop) def create_tree_cache( self, *, full_evictable_size: int = 0, swa_evictable_size: int = 0, evictable_size: int = 0, ) -> MagicMock: tree_cache = MagicMock() tree_cache.full_evictable_size.return_value = full_evictable_size tree_cache.swa_evictable_size.return_value = swa_evictable_size tree_cache.evictable_size.return_value = evictable_size tree_cache.disable = False tree_cache.inc_lock_ref.return_value = None tree_cache.dec_lock_ref.return_value = None return tree_cache def create_token_allocator( self, *, full_available_size: int = 0, swa_available_size: int = 0, available_size: int = 0, ) -> MagicMock: allocator = MagicMock() allocator.full_available_size.return_value = full_available_size allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size return allocator def create_running_batch(self, reqs=None) -> MagicMock: batch = MagicMock() batch.reqs = list(reqs or []) batch.release_req.return_value = None batch.filter_batch.return_value = None return batch def create_server_args( self, *, schedule_low_priority_values_first: bool ) -> MagicMock: server_args = MagicMock() server_args.schedule_low_priority_values_first = ( schedule_low_priority_values_first ) return server_args def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time=0): req = MagicMock(spec=Req) req.rid = str(rid) req.priority = priority req.extend_input_len = 0 req.extend_logprob_start_len = 0 req.output_ids = [0] * output_len req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.finished.return_value = False return req def create_adder(self, running_batch): return PrefillAdder( page_size=1, tree_cache=self.mock_tree_cache, token_to_kv_pool_allocator=self.mock_token_allocator, running_batch=running_batch, new_token_ratio=1.0, rem_input_tokens=10000, rem_chunk_tokens=None, mixed_with_decode_tokens=0, priority_scheduling_preemption_threshold=0, ) def test_preempt_success_high_priority_values_first(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=False ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 225 new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49) success = adder.preempt_to_schedule(new_req, mock_server_args) self.assertTrue(success) self.assertIn(running_reqs[0], adder.preempt_list) self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175 running_batch.release_req.assert_called_once() def test_preempt_success_low_priority_values_first(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=True ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 225 new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49) success = adder.preempt_to_schedule(new_req, mock_server_args) self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125 running_batch.release_req.assert_called_once() def test_preempt_fail_low_priority_values_first(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=True ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 225 new_req_fail_by_priority_check = self.create_mock_req( "new1", priority=2, max_new_tokens=49 ) success_by_priority_check = adder.preempt_to_schedule( new_req_fail_by_priority_check, mock_server_args ) self.assertFalse(success_by_priority_check) new_req_fail_by_priority_check = self.create_mock_req( "new2", priority=1, max_new_tokens=110 ) success_by_capacity_check = adder.preempt_to_schedule( new_req_fail_by_priority_check, mock_server_args ) self.assertFalse(success_by_capacity_check) def test_preempt_fail_high_priority_values_first(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=False ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 225 new_req_fail_by_priority_check = self.create_mock_req( "new1", priority=0, max_new_tokens=49 ) success_by_priority_check = adder.preempt_to_schedule( new_req_fail_by_priority_check, mock_server_args ) self.assertFalse(success_by_priority_check) new_req_fail_by_priority_check = self.create_mock_req( "new2", priority=-1, max_new_tokens=110 ) success_by_capacity_check = adder.preempt_to_schedule( new_req_fail_by_priority_check, mock_server_args ) self.assertFalse(success_by_capacity_check) def test_preempt_success_low_priority_values_first_exact_once(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ("run4", 2, 125), ("run4", 2, 125), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=True ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 475 new_req = self.create_mock_req("new1", priority=1, max_new_tokens=75) success = adder.preempt_to_schedule(new_req, mock_server_args) self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) self.assertEqual( adder.rem_total_token_offset, 375 ) # 50 + 75 + 100 + 125 + 125 - 100 = 375 running_batch.release_req.assert_called_once() def test_preempt_success_low_priority_values_first_exact_twice(self): params = [ ("run1", 0, 50), ("run2", 1, 75), ("run3", 2, 100), ("run4", 2, 125), ("run4", 2, 125), ] running_reqs = [ self.create_mock_req(rid, priority, max_new_tokens) for rid, priority, max_new_tokens in params ] mock_server_args = self.create_server_args( schedule_low_priority_values_first=True ) running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) self.assertEqual(adder.rem_total_token_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam ) self.mock_token_allocator.available_size.return_value = 475 new_req = self.create_mock_req("new1", priority=1, max_new_tokens=200) success = adder.preempt_to_schedule(new_req, mock_server_args) self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) self.assertIn(running_reqs[3], adder.preempt_list) self.assertEqual( adder.rem_total_token_offset, 250 ) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250 self.assertEqual(running_batch.release_req.call_count, 2) if __name__ == "__main__": unittest.main()