From ab9a46d462f4a4d4ff117fb9e9f21ef894f4a4cb Mon Sep 17 00:00:00 2001 From: vipwangerxiao Date: Fri, 28 Nov 2025 16:14:21 +0800 Subject: [PATCH] Support configuring the request limit per receiving poll (#14076) Co-authored-by: Peng Wang Co-authored-by: Feng Su <225349073+sufeng-buaa@users.noreply.github.com> --- docs/references/environment_variables.md | 1 + python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/scheduler.py | 10 ++++++++++ 3 files changed, 12 insertions(+) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 8b4fdb859..d8a1e67ee 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -30,6 +30,7 @@ SGLang supports various environment variables that can be used to configure its | `SGL_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | | `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | | `SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP` | Disable overlap schedule for consecutive prefill batches | `false` | +| `SGLANG_SCHEDULER_MAX_RECV_PER_POLL` | Set the maximum number of requests per poll, with a negative value indicating no limit | `-1` | | `SGLANG_DISABLE_FA4_WARMUP` | Disable Flash Attention 4 warmup passes (set to `1`, `true`, `yes`, or `on` to disable) | `false` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DEFAULT` | Default weight value for scheduler recv skipper counter (used when forward mode doesn't match specific modes). Only active when `--scheduler-recv-interval > 1`. The counter accumulates weights and triggers request polling when reaching the interval threshold. | `1000` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DECODE` | Weight increment for decode forward mode in scheduler recv skipper. Works with `--scheduler-recv-interval` to control polling frequency during decode phase. | `1` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 87295d07c..7106c4d62 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -178,6 +178,7 @@ class Envs: # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False) + SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1) SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False) # Test: pd-disaggregation diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 61b774a03..45c724e4f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -274,6 +274,7 @@ class Scheduler( self.page_size = server_args.page_size self.enable_hierarchical_cache = server_args.enable_hierarchical_cache self.enable_hicache_storage = server_args.hicache_storage_backend is not None + self.max_recv_per_poll = envs.SGLANG_SCHEDULER_MAX_RECV_PER_POLL.get() # Distributed rank info self.attn_tp_rank, self.attn_tp_size, self.attn_dp_rank = ( @@ -1046,6 +1047,11 @@ class Scheduler( if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.self_check_during_busy() + def recv_limit_reached(self, num_recv_reqs: int) -> bool: + if self.max_recv_per_poll < 0: + return False + return num_recv_reqs >= self.max_recv_per_poll + def recv_requests( self, ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: @@ -1064,6 +1070,8 @@ class Scheduler( while True: try: + if self.recv_limit_reached(len(recv_reqs)): + break recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) except zmq.ZMQError: break @@ -1071,6 +1079,8 @@ class Scheduler( while True: try: + if self.recv_limit_reached(len(recv_reqs)): + break recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK) except zmq.ZMQError: break