[Utils] Refactor device cache emptying (#24861)

Co-authored-by: Biao He <biao@Biaos-MacBook-Air.local>

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

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

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)

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