import sys import types import unittest from types import SimpleNamespace from unittest.mock import MagicMock import torch if "sgl_kernel" not in sys.modules: sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py" sys.modules["sgl_kernel"].__path__ = [] if not hasattr(sys.modules["sgl_kernel"], "__getattr__"): def _sgl_kernel_getattr(name): if name.startswith("__"): raise AttributeError(name) fn = lambda *args, **kwargs: None setattr(sys.modules["sgl_kernel"], name, fn) return fn sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr if "sgl_kernel.kvcacheio" not in sys.modules: sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio") for _name in ( "sgl_per_token_group_quant_8bit", "sgl_per_token_group_quant_fp8", "sgl_per_token_quant_fp8", "fp8_blockwise_scaled_mm", "fp8_scaled_mm", "silu_and_mul", ): if not hasattr(sys.modules["sgl_kernel"], _name): setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None) if "sgl_kernel.quantization" not in sys.modules: quantization_stub = types.ModuleType("sgl_kernel.quantization") for _name in ( "ggml_dequantize", "ggml_moe_a8", "ggml_moe_a8_vec", "ggml_moe_get_block_size", "ggml_mul_mat_a8", "ggml_mul_mat_vec_a8", ): setattr(quantization_stub, _name, lambda *args, **kwargs: None) sys.modules["sgl_kernel.quantization"] = quantization_stub _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") for _schema in ( "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", ): try: _sgl_kernel_lib.define(_schema) except RuntimeError as exc: if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): raise from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.cp_shared_kv_prefill_buffer_estimator import ( CPSharedKVPrefillBufferEstimatorContext, ) from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefResult, IncLockRefResult, ) from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler 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-1-gpu-small") register_amd_ci(est_time=2, suite="stage-b-test-1-gpu-small-amd") class TestPrefillAdder(CustomTestCase): def setUp(self): set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) self.mock_tree_cache = self.create_tree_cache() self.mock_token_allocator = self.create_token_allocator() 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.supports_mamba.return_value = False tree_cache.disable = False tree_cache.inc_lock_ref.return_value = IncLockRefResult() tree_cache.dec_lock_ref.return_value = DecLockRefResult() 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_prefill_req(self, rid, extend_input_len, max_new_tokens=1): req = self.create_mock_req(rid, priority=0, max_new_tokens=max_new_tokens) req.extend_input_len = extend_input_len req.host_hit_length = 0 req.prefix_indices = torch.empty((0,), dtype=torch.int64) req.fill_ids = list(range(extend_input_len)) req.last_node = object() req.sampling_params.ignore_eos = False req.set_extend_input_len.side_effect = lambda value: setattr( req, "extend_input_len", value ) return req def create_adder(self, running_batch, **kwargs): defaults = dict( 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, ) defaults.update(kwargs) return PrefillAdder(**defaults) def create_buffer_estimator_context(self, *, kv_cache_dim=1, vocab_size=16): return CPSharedKVPrefillBufferEstimatorContext( kvcache=SimpleNamespace( kv_cache_dim=kv_cache_dim, store_dtype=torch.bfloat16, index_head_dim=8, quant_block_size=4, index_k_with_scale_buffer_dtype=torch.uint8, ), model_config=SimpleNamespace(vocab_size=vocab_size), tp_size=1, page_size=64, logprob_chunk_enabled=False, logprob_chunk_size=2048, bs_gt1_l1_prefetch_enabled=False, ) 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_skip_already_preempted_request(self): params = [ ("req_prio_0", 0, 50), ("req_prio_1", 1, 75), ("req_prio_2", 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 self.mock_token_allocator.available_size.return_value = 225 # New request preempts req_prio_0 first_req = self.create_mock_req( "new_req_prio_1", priority=1, max_new_tokens=49 ) first_success = adder.preempt_to_schedule(first_req, mock_server_args) self.assertTrue(first_success) self.assertIn(running_reqs[0], adder.preempt_list) self.assertEqual(adder.rem_total_token_offset, 175) running_batch.release_req.assert_called_once() # Second call needs more tokens than currently free, so it would need to # preempt req_prio_0 again if already-preempted requests were not filtered out. second_req = self.create_mock_req( "second_new_req_prio_1", priority=1, max_new_tokens=76 ) second_success = adder.preempt_to_schedule(second_req, mock_server_args) self.assertFalse(second_success) self.assertEqual(adder.rem_total_token_offset, 175) self.assertEqual(adder.preempt_list.count(running_reqs[0]), 1) running_batch.release_req.assert_called_once() 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) def test_mixed_chunk_prefill_budgets(self): self.mock_token_allocator.available_size.return_value = 1000 decode_reqs = [ self.create_mock_req(f"decode_{i}", priority=0, max_new_tokens=50) for i in range(8) ] running_batch = self.create_running_batch(decode_reqs) adder = self.create_adder( running_batch, rem_input_tokens=200, rem_chunk_tokens=64, mixed_with_decode_tokens=len(decode_reqs), ) self.assertEqual(adder.rem_input_tokens, 192) # 200 - 8 self.assertEqual(adder.rem_chunk_tokens, 56) # 64 - 8 self.assertEqual(adder.rem_total_token_offset, 408) # 8 + 8 * 50 self.assertEqual(adder.cur_rem_token_offset, 8) self.assertEqual(adder.budget_state(), AddReqResult.CONTINUE) # Add a prefill that exactly consumes the chunk budget req1 = self.create_mock_req("req1", priority=0, max_new_tokens=64) req1.extend_input_len = 56 req1.host_hit_length = 0 req1.prefix_indices = [] req1.fill_ids = list(range(56)) req1.last_node = MagicMock() req1.sampling_params.ignore_eos = False result1 = adder.add_one_req( req1, has_chunked_req=False, truncation_align_size=None ) self.assertEqual(len(adder.can_run_list), 1) self.assertEqual(adder.rem_chunk_tokens, 0) # 56 - 56 self.assertEqual(adder.rem_input_tokens, 136) # 192 - 56 self.assertEqual(result1, AddReqResult.OTHER) # 3 decode requests finished remaining_decode_reqs = decode_reqs[3:] running_batch2 = self.create_running_batch(remaining_decode_reqs) adder2 = self.create_adder( running_batch2, rem_input_tokens=200, rem_chunk_tokens=64, mixed_with_decode_tokens=len(remaining_decode_reqs), ) self.assertEqual(adder2.rem_input_tokens, 195) # 200 - 5 self.assertEqual(adder2.rem_chunk_tokens, 59) # 64 - 5 self.assertEqual(adder2.rem_total_token_offset, 255) # 5 + 5 * 50 self.assertEqual(adder2.budget_state(), AddReqResult.CONTINUE) # Same prefill no longer exhausts the chunk budget req2 = self.create_mock_req("req2", priority=0, max_new_tokens=64) req2.extend_input_len = 56 req2.host_hit_length = 0 req2.prefix_indices = [] req2.fill_ids = list(range(56)) req2.last_node = MagicMock() req2.sampling_params.ignore_eos = False result2 = adder2.add_one_req( req2, has_chunked_req=False, truncation_align_size=None ) self.assertEqual(len(adder2.can_run_list), 1) self.assertEqual(adder2.rem_chunk_tokens, 3) # 59 - 56 = 3 remaining self.assertEqual(result2, AddReqResult.CONTINUE) # Fit last small prefill request req3 = self.create_mock_req("req3", priority=0, max_new_tokens=16) req3.extend_input_len = 3 req3.host_hit_length = 0 req3.prefix_indices = [] req3.fill_ids = list(range(3)) req3.last_node = MagicMock() req3.sampling_params.ignore_eos = False result3 = adder2.add_one_req( req3, has_chunked_req=False, truncation_align_size=None ) self.assertEqual(len(adder2.can_run_list), 2) self.assertEqual(adder2.rem_chunk_tokens, 0) # 3 - 3 = 0 self.assertEqual(result3, AddReqResult.OTHER) def test_host_load_back_passes_mem_quota(self): running_batch = self.create_running_batch() self.mock_token_allocator.available_size.return_value = 512 self.mock_tree_cache.init_load_back.return_value = ( __import__("torch").tensor([1, 2, 3, 4], dtype=__import__("torch").int64), "loaded_node", ) adder = self.create_adder( running_batch, page_size=64, rem_input_tokens=4096, ) req = self.create_mock_req("req", priority=0, max_new_tokens=16) req.extend_input_len = 256 req.host_hit_length = 128 req.prefix_indices = __import__("torch").empty( (0,), dtype=__import__("torch").int64 ) req.last_node = object() req.last_host_node = object() req.fill_ids = list(range(256)) req.cache_protected_len = 0 req.set_extend_input_len = lambda value: setattr(req, "extend_input_len", value) req.sampling_params.ignore_eos = False result = adder.add_one_req( req, has_chunked_req=False, truncation_align_size=None ) self.assertNotEqual(result, AddReqResult.NO_TOKEN) params = self.mock_tree_cache.init_load_back.call_args.args[0] self.assertEqual(params.mem_quota, 320) def test_load_back_mem_quota_counts_evictable_device_tokens(self): self.mock_tree_cache = self.create_tree_cache(evictable_size=90000) self.mock_token_allocator = self.create_token_allocator(available_size=1024) adder = self.create_adder( self.create_running_batch(), page_size=64, tree_cache=self.mock_tree_cache, token_to_kv_pool_allocator=self.mock_token_allocator, ) quota = adder._get_load_back_mem_quota(real_input_tokens=65536) self.assertEqual(quota, 90000 + 1024 - 65536 - 64) def test_cp_prefill_gate_keeps_single_request_by_default(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, ) first = self.create_prefill_req("first", extend_input_len=128) second = self.create_prefill_req("second", extend_input_len=128) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.OTHER, ) self.assertEqual([req.rid for req in adder.can_run_list], ["first"]) def test_cp_prefill_gate_allows_batched_requests_when_enabled(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=2, cp_shared_kv_prefill_max_total_extend_tokens=256, ) first = self.create_prefill_req("first", extend_input_len=128) second = self.create_prefill_req("second", extend_input_len=128) third = self.create_prefill_req("third", extend_input_len=64) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(third, has_chunked_req=False, truncation_align_size=None), AddReqResult.OTHER, ) self.assertEqual([req.rid for req in adder.can_run_list], ["first", "second"]) def test_cp_prefill_total_extend_limit_is_page_aligned_and_allows_first_req(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=128, ) first = self.create_prefill_req("first", extend_input_len=65) second = self.create_prefill_req("second", extend_input_len=1) oversized_first = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=64, ) large = self.create_prefill_req("large", extend_input_len=128) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual(adder.cp_shared_kv_prefill_total_extend_tokens, 128) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.OTHER, ) self.assertEqual( oversized_first.add_one_req( large, has_chunked_req=False, truncation_align_size=None ), AddReqResult.CONTINUE, ) self.assertEqual([req.rid for req in oversized_first.can_run_list], ["large"]) def test_cp_prefill_total_extend_limit_replaces_generic_input_budget(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, # Simulates the generic max_prefill_tokens budget being smaller # than the CP shared-KV bs>1 budget. rem_input_tokens=192, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=256, ) first = self.create_prefill_req("first", extend_input_len=128) second = self.create_prefill_req("second", extend_input_len=128) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) second_result = adder.add_one_req( second, has_chunked_req=False, truncation_align_size=None ) self.assertNotEqual(second_result, AddReqResult.NO_TOKEN) self.assertEqual([req.rid for req in adder.can_run_list], ["first", "second"]) self.assertEqual(adder.cp_shared_kv_prefill_total_extend_tokens, 256) def test_cp_prefill_total_extend_limit_is_capped_by_chunked_prefill_size(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) adder = self.create_adder( self.create_running_batch(), page_size=64, # The CP-specific extend limit is larger than the chunked prefill # budget. Effective admission should use the smaller chunk budget # to avoid advertising an unreachable per-batch extend capacity. rem_input_tokens=192, rem_chunk_tokens=128, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=256, ) self.assertEqual(adder.cp_shared_kv_prefill_max_total_extend_tokens, 128) # The generic max_prefill_tokens lift should also use the effective # limit, not the raw 256-token CP limit. self.assertEqual(adder.rem_input_tokens, 192) def test_cp_prefill_total_extend_limit_defaults_to_chunk_budget(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=192, rem_chunk_tokens=128, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=None, cp_shared_kv_prefill_max_total_extend_tokens=None, cp_shared_kv_prefill_max_total_cached_tokens=None, ) self.assertIsNone(adder.cp_shared_kv_prefill_max_batch_requests) self.assertEqual(adder.cp_shared_kv_prefill_max_total_extend_tokens, 128) self.assertIsNone(adder.cp_shared_kv_prefill_max_total_cached_tokens) # The generic budget lift uses the effective defaulted extend limit. self.assertEqual(adder.rem_input_tokens, 192) def test_cp_prefill_total_extend_limit_does_not_bypass_allocator_capacity(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 200 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=64, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, ) first = self.create_prefill_req("first", extend_input_len=128) second = self.create_prefill_req("second", extend_input_len=128) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.NO_TOKEN, ) self.assertEqual([req.rid for req in adder.can_run_list], ["first"]) def test_cp_prefill_chunked_req_excludes_new_requests_even_when_bs_gt1_enabled(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, rem_chunk_tokens=256, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, ) chunked = self.create_prefill_req("chunked", extend_input_len=128) normal = self.create_prefill_req("normal", extend_input_len=128) adder.new_chunked_req = adder.add_chunked_req(chunked) self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"]) self.assertIsNone(adder.new_chunked_req) self.assertEqual(chunked.extend_input_len, 128) self.assertEqual( adder.add_one_req( normal, has_chunked_req=True, truncation_align_size=None ), AddReqResult.OTHER, ) self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"]) def test_cp_prefill_total_cached_limit_stops_second_cached_request(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, cp_shared_kv_prefill_max_total_cached_tokens=4096, ) first = self.create_prefill_req("first", extend_input_len=64) first.prefix_indices = torch.arange(4096, dtype=torch.int64) first.fill_ids = list(range(4096 + 64)) second = self.create_prefill_req("second", extend_input_len=64) second.prefix_indices = torch.arange(4096, dtype=torch.int64) second.fill_ids = list(range(4096 + 64)) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.OTHER, ) self.assertEqual([req.rid for req in adder.can_run_list], ["first"]) self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 4096) def test_cp_prefill_total_cached_limit_allows_single_oversized_cached_request(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 20000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, cp_shared_kv_prefill_max_total_cached_tokens=4096, ) oversized = self.create_prefill_req("oversized", extend_input_len=64) oversized.prefix_indices = torch.arange(8192, dtype=torch.int64) oversized.fill_ids = list(range(8192 + 64)) self.assertEqual( adder.add_one_req( oversized, has_chunked_req=False, truncation_align_size=None ), AddReqResult.CONTINUE, ) self.assertEqual([req.rid for req in adder.can_run_list], ["oversized"]) self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 8192) def test_cp_prefill_buffer_limit_stops_second_request_without_token_gate(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, cp_shared_kv_prefill_max_total_cached_tokens=4096, cp_shared_kv_prefill_max_buffer_size=1, cp_shared_kv_prefill_buffer_estimator_context=( self.create_buffer_estimator_context() ), ) first = self.create_prefill_req("first", extend_input_len=64) second = self.create_prefill_req("second", extend_input_len=64) self.assertEqual( adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), AddReqResult.CONTINUE, ) self.assertEqual( adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), AddReqResult.OTHER, ) self.assertEqual([req.rid for req in adder.can_run_list], ["first"]) self.assertGreater(adder.cp_shared_kv_prefill_estimated_peak_buffer_bytes, 1) def test_cp_prefill_buffer_limit_allows_single_oversized_request(self): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", enable_nsa_prefill_context_parallel=True, nsa_prefill_cp_mode="in-seq-split", ) ) self.mock_token_allocator.available_size.return_value = 10000 adder = self.create_adder( self.create_running_batch(), page_size=64, rem_input_tokens=4096, enable_cp_shared_kv_prefill_bs_gt1=True, cp_shared_kv_prefill_max_batch_requests=8, cp_shared_kv_prefill_max_total_extend_tokens=4096, cp_shared_kv_prefill_max_buffer_size=1, cp_shared_kv_prefill_buffer_estimator_context=( self.create_buffer_estimator_context() ), ) oversized = self.create_prefill_req("oversized", extend_input_len=64) self.assertEqual( adder.add_one_req( oversized, has_chunked_req=False, truncation_align_size=None ), AddReqResult.CONTINUE, ) self.assertEqual([req.rid for req in adder.can_run_list], ["oversized"]) if __name__ == "__main__": unittest.main()