diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a62d7fab7..ef9df64e4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2086,12 +2086,24 @@ class Scheduler( self.decode_mem_cache_buf_multiplier ) ) or (TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0): - old_available_tokens = self.token_to_kv_pool_allocator.available_size() + if self.is_hybrid_swa: + old_available_tokens = min( + self.token_to_kv_pool_allocator.full_available_size(), + self.token_to_kv_pool_allocator.swa_available_size(), + ) + else: + old_available_tokens = self.token_to_kv_pool_allocator.available_size() old_ratio = self.new_token_ratio retracted_reqs, new_token_ratio, reqs_to_abort = batch.retract_decode( self.server_args, self.decode_mem_cache_buf_multiplier ) - new_available_tokens = self.token_to_kv_pool_allocator.available_size() + if self.is_hybrid_swa: + new_available_tokens = min( + self.token_to_kv_pool_allocator.full_available_size(), + self.token_to_kv_pool_allocator.swa_available_size(), + ) + else: + new_available_tokens = self.token_to_kv_pool_allocator.available_size() new_token_gained = new_available_tokens - old_available_tokens self.num_retracted_reqs = len(retracted_reqs)