diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 0e92b8723..0f1176990 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -329,6 +329,7 @@ class PrefillAdder: rem_chunk_tokens: Optional[int], mixed_with_decode_tokens: int = 0, priority_scheduling_preemption_threshold: int = 0, + prefill_max_requests: Optional[int] = None, ): self.page_size = page_size self.tree_cache = tree_cache @@ -368,6 +369,7 @@ class PrefillAdder: priority_scheduling_preemption_threshold ) self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.prefill_max_requests = prefill_max_requests def _get_running_request_total_token_offset(self, req: Req) -> int: return ( @@ -575,6 +577,10 @@ class PrefillAdder: # therefore, the prefill-batch setting is temporarily set to 1. if self.nsa_enable_prefill_cp and len(self.can_run_list) >= 1: return AddReqResult.OTHER + + if (x := self.prefill_max_requests) is not None and len(self.can_run_list) >= x: + return AddReqResult.OTHER + if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True): return self.add_one_req_ignore_eos(req) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index db884c4a2..2c6bf662e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1822,6 +1822,7 @@ class Scheduler( chunked_prefill_size, running_bs if self.is_mixed_chunk else 0, self.priority_scheduling_preemption_threshold, + prefill_max_requests=self.server_args.prefill_max_requests, ) if self.chunked_req is not None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 42294937b..69f88b00b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -290,6 +290,7 @@ class ServerArgs: chunked_prefill_size: Optional[int] = None enable_dynamic_chunking: bool = False max_prefill_tokens: int = 16384 + prefill_max_requests: Optional[int] = None schedule_policy: str = "fcfs" enable_priority_scheduling: bool = False abort_on_priority_when_disabled: bool = False @@ -2516,6 +2517,12 @@ class ServerArgs: default=ServerArgs.chunked_prefill_size, help="The maximum number of tokens in a chunk for the chunked prefill. Setting this to -1 means disabling chunked prefill.", ) + parser.add_argument( + "--prefill-max-requests", + type=int, + default=ServerArgs.prefill_max_requests, + help="The maximum number of requests in a prefill batch. If not specified, there is no limit.", + ) parser.add_argument( "--enable-dynamic-chunking", action="store_true",