From ca4b64c79a529b1fc19e6d737aeef6e1ee551542 Mon Sep 17 00:00:00 2001 From: kzlin Date: Wed, 6 May 2026 21:29:30 +0800 Subject: [PATCH] feat(sglang): expose backpressure pause hint in admit_direct_append Add `recommended_pause_ms` field to DirectAppendAdmissionReqOutput so D can advise callers when its transfer queue is heavy or KV pool is near capacity. The hint is computed from transfer_queue_depth, retracted_queue_depth, and post-trim token_usage; thresholds are simple heuristics (>0.90 usage, >=8 queue depth, retracted>0). Default behavior is unchanged for callers that ignore the field. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../python/sglang/srt/managers/io_struct.py | 3 + .../python/sglang/srt/managers/scheduler.py | 56 ++++++++++++++----- 2 files changed, 44 insertions(+), 15 deletions(-) diff --git a/third_party/sglang/python/sglang/srt/managers/io_struct.py b/third_party/sglang/python/sglang/srt/managers/io_struct.py index 2ff5366..3591976 100644 --- a/third_party/sglang/python/sglang/srt/managers/io_struct.py +++ b/third_party/sglang/python/sglang/srt/managers/io_struct.py @@ -1622,6 +1622,9 @@ class DirectAppendAdmissionReqOutput(BaseReq): decode_prealloc_queue_reqs: int = 0 decode_transfer_queue_reqs: int = 0 decode_retracted_queue_reqs: int = 0 + # Backpressure hint: if > 0, the caller should pause this many ms before + # sending more requests to this D. Computed from transfer-queue depth. + recommended_pause_ms: int = 0 @dataclass diff --git a/third_party/sglang/python/sglang/srt/managers/scheduler.py b/third_party/sglang/python/sglang/srt/managers/scheduler.py index 7e7c85e..15d7558 100644 --- a/third_party/sglang/python/sglang/srt/managers/scheduler.py +++ b/third_party/sglang/python/sglang/srt/managers/scheduler.py @@ -3605,26 +3605,25 @@ class Scheduler( if not resident and not is_seed: # direct_append requires the session already resident on this D. # For seed we skip this check and let capacity decide. + transfer_queue_depth = len(self.disagg_decode_transfer_queue.queue) + retracted_queue_depth = len(self.disagg_decode_prealloc_queue.retracted_queue) + available_size = int(self.token_to_kv_pool_allocator.available_size()) + token_usage = 1.0 - available_size / max(1, self.max_total_num_tokens) return DirectAppendAdmissionReqOutput( can_admit=False, resident=False, reason="session-not-resident", - available_tokens_before=int( - self.token_to_kv_pool_allocator.available_size() - ), - available_tokens_after=int( - self.token_to_kv_pool_allocator.available_size() - ), - token_usage=( - 1.0 - - self.token_to_kv_pool_allocator.available_size() - / max(1, self.max_total_num_tokens) - ), + available_tokens_before=available_size, + available_tokens_after=available_size, + token_usage=token_usage, num_running_reqs=len(self.running_batch.reqs), decode_prealloc_queue_reqs=len(self.disagg_decode_prealloc_queue.queue), - decode_transfer_queue_reqs=len(self.disagg_decode_transfer_queue.queue), - decode_retracted_queue_reqs=len( - self.disagg_decode_prealloc_queue.retracted_queue + decode_transfer_queue_reqs=transfer_queue_depth, + decode_retracted_queue_reqs=retracted_queue_depth, + recommended_pause_ms=self._compute_backpressure_pause_hint( + transfer_queue_depth=transfer_queue_depth, + retracted_queue_depth=retracted_queue_depth, + token_usage_after=token_usage, ), ) @@ -3648,6 +3647,7 @@ class Scheduler( ) reason = None if can_admit else "no-space" + transfer_queue_depth = len(self.disagg_decode_transfer_queue.queue) return DirectAppendAdmissionReqOutput( can_admit=can_admit, resident=True, @@ -3662,10 +3662,36 @@ class Scheduler( ), num_running_reqs=len(self.running_batch.reqs), decode_prealloc_queue_reqs=len(self.disagg_decode_prealloc_queue.queue), - decode_transfer_queue_reqs=len(self.disagg_decode_transfer_queue.queue), + decode_transfer_queue_reqs=transfer_queue_depth, decode_retracted_queue_reqs=decode_retracted_queue_reqs, + recommended_pause_ms=self._compute_backpressure_pause_hint( + transfer_queue_depth=transfer_queue_depth, + retracted_queue_depth=decode_retracted_queue_reqs, + token_usage_after=( + 1.0 - available_tokens_after / max(1, self.max_total_num_tokens) + ), + ), ) + def _compute_backpressure_pause_hint( + self, + *, + transfer_queue_depth: int, + retracted_queue_depth: int, + token_usage_after: float, + ) -> int: + # If D is already retracting requests, pause aggressively. + if retracted_queue_depth > 0: + return 1500 + # KV pool above 90%: pause proportional to overshoot. + if token_usage_after >= 0.90: + overshoot = int((token_usage_after - 0.90) * 10000) + return max(200, min(2000, overshoot * 5)) + # Transfer queue heavy: pause linearly with depth. + if transfer_queue_depth >= 8: + return min(2000, transfer_queue_depth * 100) + return 0 + def maybe_sleep_on_idle(self): if self.idle_sleeper is not None: self.idle_sleeper.maybe_sleep()