Fix swa available memory check (#15867)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user