diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 005079801..8cebdb09e 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1407,7 +1407,7 @@ class UpdateWeightFromDiskReqInput(BaseReq): weight_version: Optional[str] = None # Whether to update weights asynchronously is_async: bool = False - # Whether to empty torch cache + # Whether to call torch.cuda.empty_cache() during flush torch_empty_cache: bool = False # Whether to keep the scheduler paused after weight update keep_pause: bool = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 391473ab6..8ef38bdfa 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -214,6 +214,7 @@ from sglang.srt.utils import ( broadcast_pyobj, configure_gc_logger, configure_logger, + empty_device_cache, freeze_gc, get_available_gpu_memory, get_bool_env_var, @@ -3378,7 +3379,7 @@ class Scheduler( return DetachHiCacheStorageReqOutput(success=False, message=msg) def flush_cache(self, empty_cache: bool = True): - """Flush the memory pool and cache.""" + """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache.""" if self.is_fully_idle(): self.cur_batch = None self.last_batch = None @@ -3392,7 +3393,7 @@ class Scheduler( self.draft_worker.clear_cache_pool() if empty_cache: - torch.cuda.empty_cache() + empty_device_cache(self.device_module) logger.info("Cache flushed successfully!") success = True else: @@ -3781,7 +3782,7 @@ class IdleSleeper: and real_time() - self.last_empty_time > self.empty_cache_interval ): self.last_empty_time = real_time() - torch.cuda.empty_cache() + empty_device_cache() def is_health_check_generate_req(recv_req): diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index 590537fd6..f2daf644d 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -42,6 +42,12 @@ logger = logging.getLogger(__name__) class SchedulerUpdateWeightsMixin: + def flush_cache_after_weight_update(self: Scheduler, recv_req) -> None: + if recv_req.flush_cache: + flush_cache_success = self.flush_cache( + empty_cache=recv_req.torch_empty_cache + ) + assert flush_cache_success, "Cache flush failed after updating weights" def update_weights_from_disk( self: Scheduler, recv_req: UpdateWeightFromDiskReqInput @@ -51,11 +57,8 @@ class SchedulerUpdateWeightsMixin: tp_success = success if success and self.draft_worker is not None: success, message = self.draft_worker.update_weights_from_disk(recv_req) - if tp_success and recv_req.flush_cache: - flush_cache_success = self.flush_cache( - empty_cache=recv_req.torch_empty_cache - ) - assert flush_cache_success, "Cache flush failed after updating weights" + if tp_success: + self.flush_cache_after_weight_update(recv_req) if not success: logger.error(message) return UpdateWeightFromDiskReqOutput(success, message, 0) @@ -81,11 +84,7 @@ class SchedulerUpdateWeightsMixin: """Update the online model parameter.""" success, message = self.tp_worker.update_weights_from_distributed(recv_req) if success: - if recv_req.flush_cache: - flush_cache_success = self.flush_cache( - empty_cache=recv_req.torch_empty_cache - ) - assert flush_cache_success, "Cache flush failed after updating weights" + self.flush_cache_after_weight_update(recv_req) else: logger.error(message) return UpdateWeightsFromDistributedReqOutput(success, message) @@ -99,13 +98,8 @@ class SchedulerUpdateWeightsMixin: else: worker = self.draft_worker or self.tp_worker success, message = worker.update_weights_from_tensor(recv_req) - # TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later if success: - if recv_req.flush_cache: - flush_cache_success = self.flush_cache( - empty_cache=recv_req.torch_empty_cache - ) - assert flush_cache_success, "Cache flush failed after updating weights" + self.flush_cache_after_weight_update(recv_req) else: logger.error(message) torch.distributed.barrier(group=self.tp_cpu_group) @@ -119,11 +113,8 @@ class SchedulerUpdateWeightsMixin: tp_success = success if success and self.draft_worker is not None: success, message = self.draft_worker.update_weights_from_ipc(recv_req) - if tp_success and recv_req.flush_cache: - flush_cache_success = self.flush_cache( - empty_cache=recv_req.torch_empty_cache - ) - assert flush_cache_success, "Cache flush failed after updating weights" + if tp_success: + self.flush_cache_after_weight_update(recv_req) if not success: logger.error(message) torch.distributed.barrier(group=self.tp_cpu_group) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index d2891940c..c91fcdcfb 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -495,6 +495,25 @@ def calculate_time(show=False, min_cost_ms=0.0): return wrapper +def empty_device_cache(device_module: Optional[Any] = None) -> bool: + """Release unused cached blocks from the active device allocator. + + This does not clear SGLang KV/radix/request caches and does not free live + tensors. It only forwards to the backend allocator's empty_cache hook when + one is available. + """ + + if device_module is None: + device_module = torch.get_device_module() + + empty_cache = getattr(device_module, "empty_cache", None) + if empty_cache is None: + return False + + empty_cache() + return True + + def get_available_gpu_memory( device, gpu_id, distributed=False, empty_cache=True, cpu_group=None ): @@ -513,7 +532,7 @@ def get_available_gpu_memory( ) if empty_cache: - torch.cuda.empty_cache() + empty_device_cache(torch.cuda) props = torch.cuda.get_device_properties(gpu_id) if props.is_integrated: # On these devices, which use sysmem as device mem, torch.cuda.mem_get_info() @@ -535,7 +554,7 @@ def get_available_gpu_memory( ) if empty_cache: - torch.xpu.empty_cache() + empty_device_cache(torch.xpu) used_memory = torch.xpu.memory_allocated() total_gpu_memory = torch.xpu.get_device_properties(gpu_id).total_memory free_gpu_memory = total_gpu_memory - used_memory @@ -567,7 +586,7 @@ def get_available_gpu_memory( "which may cause useless memory allocation for torch NPU context.", ) if empty_cache: - torch.npu.empty_cache() + empty_device_cache(torch.npu) free_gpu_memory, total_gpu_memory = torch.npu.mem_get_info() elif device == "musa": num_gpus = torch.musa.device_count() @@ -579,7 +598,7 @@ def get_available_gpu_memory( "which may cause useless memory allocation for torch MUSA context.", ) if empty_cache: - torch.musa.empty_cache() + empty_device_cache(torch.musa) props = torch.musa.get_device_properties(gpu_id) if props.is_integrated: # On these devices, which use sysmem as device mem, torch.musa.mem_get_info()