Fix swa available memory check (#15867)

This commit is contained in:
Ke Bao
2025-12-26 13:19:01 +08:00
committed by GitHub
parent 086813ae8a
commit c28c536c91

View File

@@ -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)