From 9b9d21312adc39c6a5d63d0dcd0a53f91b7afa18 Mon Sep 17 00:00:00 2001 From: Chenxi Li <41864925+ConnorLi96@users.noreply.github.com> Date: Fri, 12 Dec 2025 17:13:05 -0800 Subject: [PATCH] Feature/Fix multi lora scheduler blocking issue and evict LoRA None lastly (#14795) --- python/sglang/srt/lora/eviction_policy.py | 32 +++++++++++------------ python/sglang/srt/lora/mem_pool.py | 17 +++++++++--- python/sglang/srt/managers/scheduler.py | 22 +++++++++++----- 3 files changed, 45 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/lora/eviction_policy.py b/python/sglang/srt/lora/eviction_policy.py index d4b29612f..74258f283 100644 --- a/python/sglang/srt/lora/eviction_policy.py +++ b/python/sglang/srt/lora/eviction_policy.py @@ -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}" diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index fdebb860c..f57ea39ac 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -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] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 094e1aca3..84e5afa83 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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):