Feature/Fix multi lora scheduler blocking issue and evict LoRA None lastly (#14795)

This commit is contained in:
Chenxi Li
2025-12-12 17:13:05 -08:00
committed by GitHub
parent 44fd701732
commit 9b9d21312a
3 changed files with 45 additions and 26 deletions

View File

@@ -63,14 +63,6 @@ class LRUEvictionPolicy(EvictionPolicy):
def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
"""Select the least recently used adapter from candidates."""
# Base model (currently None, will be replaced with special UID in future)
# always has lowest priority - evict it first if available
BASE_MODEL_UID = None # TODO: Replace with special UID constant
if BASE_MODEL_UID in candidates:
logger.debug(f"Selected base model for eviction (LRU)")
self.eviction_count += 1
return BASE_MODEL_UID
# Iterate through access_order (oldest first) to find LRU victim
for uid in list(self.access_order.keys()):
if uid in candidates:
@@ -78,6 +70,14 @@ class LRUEvictionPolicy(EvictionPolicy):
self.eviction_count += 1
return uid
# If no tracked UID found in candidates, check if None is available
# This happens when the batch consists entirely of LoRA requests
# and None (base model) is the only eviction candidate
if None in candidates:
logger.debug("Selected None (base model) for eviction")
self.eviction_count += 1
return None
# Should never reach here if candidates is non-empty
assert False, f"Failed to select LRU victim from candidates: {candidates}"
@@ -105,14 +105,6 @@ class FIFOEvictionPolicy(EvictionPolicy):
def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
"""Select the first inserted adapter from candidates."""
# Base model (currently None, will be replaced with special UID in future)
# always has lowest priority - evict it first if available
BASE_MODEL_UID = None # TODO: Replace with special UID constant
if BASE_MODEL_UID in candidates:
logger.debug(f"Selected base model for eviction (FIFO)")
self.eviction_count += 1
return BASE_MODEL_UID
# Iterate through insertion_order (oldest first) to find FIFO victim
for uid in list(self.insertion_order.keys()):
if uid in candidates:
@@ -120,6 +112,14 @@ class FIFOEvictionPolicy(EvictionPolicy):
self.eviction_count += 1
return uid
# If no tracked UID found in candidates, check if None is available
# This happens when the batch consists entirely of LoRA requests
# and None (base model) is the only eviction candidate
if None in candidates:
logger.debug("Selected None (base model) for eviction")
self.eviction_count += 1
return None
# Should never reach here if candidates is non-empty
assert False, f"Failed to select FIFO victim from candidates: {candidates}"

View File

@@ -312,15 +312,15 @@ class LoRAMemoryPool:
uid = self.buffer_id_to_uid[buffer_id]
# Skip if this adapter is needed by current batch
# TODO (lifuhuang): we might consider supporting pinning base model (uid == None) in the future.
if uid in cur_uids:
continue
# Skip if this adapter is pinned (base model cannot be pinned, so can be evicted)
# Skip if this adapter is pinned
if uid is not None:
lora_ref = lora_refs.get(uid)
if lora_ref and lora_ref.pinned:
continue
candidates.add(uid)
if not candidates:
@@ -328,8 +328,19 @@ class LoRAMemoryPool:
"No available buffer slots found. Please ensure the number of active (pinned) loras is less than max_loras_per_batch."
)
# Prefer evicting LoRA adapters over the base model (None).
# Only evict None when the batch consists entirely of LoRA requests
# and no other adapters can be evicted.
non_none_candidates = candidates - {None}
if non_none_candidates:
# Prioritize evicting actual LoRA adapters
candidates_to_use = non_none_candidates
else:
# Only None is available for eviction (batch is all LoRA requests)
candidates_to_use = candidates
# Select victim using eviction policy
victim_uid = self.eviction_policy.select_victim(candidates)
victim_uid = self.eviction_policy.select_victim(candidates_to_use)
# Evict the selected victim
victim_buffer_id = self.uid_to_buffer_id[victim_uid]

View File

@@ -1817,13 +1817,21 @@ class Scheduler(
# Get requests from the waiting queue to a new prefill batch
for req in self.waiting_queue:
if self.enable_lora and not self.tp_worker.can_run_lora_batch(
lora_set
| set([req.lora_id for req in adder.can_run_list])
| set([req.lora_id])
):
self.running_batch.batch_is_full = True
break
if self.enable_lora:
new_lora_set = (
lora_set
| set([req.lora_id for req in adder.can_run_list])
| set([req.lora_id])
)
if not self.tp_worker.can_run_lora_batch(new_lora_set):
# If this is a LoRA request that would exceed the LoRA slot limit,
# skip it and continue to try scheduling non-LoRA requests.
# Non-LoRA requests (lora_id=None) share a single reserved slot
# and should never cause this check to fail.
if req.lora_id is not None:
# Skip this LoRA request - it would trigger adapter eviction/loading
# which is slow. We'll try to schedule it in a future iteration.
continue
running_bs = len(self.running_batch.reqs)
if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):