From 7c162fa9c5f7ab029104e8aeb93460d0a38d3034 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 17 Sep 2024 22:59:32 -0700 Subject: [PATCH] Fix schedule bug (#1451) --- python/sglang/srt/managers/policy_scheduler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/policy_scheduler.py b/python/sglang/srt/managers/policy_scheduler.py index 0c2b21acb..ada390418 100644 --- a/python/sglang/srt/managers/policy_scheduler.py +++ b/python/sglang/srt/managers/policy_scheduler.py @@ -229,6 +229,7 @@ class PrefillAdder: else: add_req_state(req, insert_sort=True) + cur_rem_tokens = self.cur_rem_tokens - len(req.origin_input_ids) tokens_freed = 0 for i, (tokens_left, tokens_occupied) in enumerate(self.req_states): decode_steps = ( @@ -237,7 +238,7 @@ class PrefillAdder: else tokens_left ) bs = len(self.req_states) - i - if self.cur_rem_tokens + tokens_freed - decode_steps * bs <= 0: + if cur_rem_tokens + tokens_freed - decode_steps * bs <= 0: return False tokens_freed += tokens_occupied