Skip torch.cuda.empty_cache() in weight update flush path (#22998)

python/sglang/srt/managers/io_struct.py

@@ -1418,6 +1418,8 @@ class UpdateWeightsFromDistributedReqInput(BaseReq):
     weight_version: Optional[str] = None
     # Optional format specification for loading
     load_format: Optional[str] = None
+    # Whether to call torch.cuda.empty_cache() during flush
+    torch_empty_cache: bool = False
 
 
 @dataclass
@@ -1445,6 +1447,8 @@ class UpdateWeightsFromTensorReqInput(BaseReq):
     weight_version: Optional[str] = None
     # Optional: Determine whether to disable updating the draft model
     disable_draft_model: Optional[bool] = None
+    # Whether to call torch.cuda.empty_cache() during flush
+    torch_empty_cache: bool = False
 
 
 @dataclass
@@ -1479,6 +1483,8 @@ class UpdateWeightsFromIPCReqInput(BaseReq):
     flush_cache: bool = True
     # Optional: Update weight version along with weights
     weight_version: Optional[str] = None
+    # Whether to call torch.cuda.empty_cache() during flush
+    torch_empty_cache: bool = False
 
 
 @dataclass

python/sglang/srt/managers/scheduler.py

@@ -3226,7 +3226,7 @@ class Scheduler(
 
         return DetachHiCacheStorageReqOutput(success=False, message=msg)
 
-    def flush_cache(self):
+    def flush_cache(self, empty_cache: bool = True):
         """Flush the memory pool and cache."""
         if self.is_fully_idle():
             self.cur_batch = None
@@ -3240,8 +3240,8 @@ class Scheduler(
             if self.draft_worker:
                 self.draft_worker.clear_cache_pool()
 
-            # TODO: allow optional empty cache
-            torch.cuda.empty_cache()
+            if empty_cache:
+                torch.cuda.empty_cache()
             logger.info("Cache flushed successfully!")
             success = True
         else:

python/sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -52,7 +52,9 @@ class SchedulerUpdateWeightsMixin:
         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()
+            flush_cache_success = self.flush_cache(
+                empty_cache=recv_req.torch_empty_cache
+            )
             assert flush_cache_success, "Cache flush failed after updating weights"
         if not success:
             logger.error(message)
@@ -80,7 +82,9 @@ class SchedulerUpdateWeightsMixin:
         success, message = self.tp_worker.update_weights_from_distributed(recv_req)
         if success:
             if recv_req.flush_cache:
-                flush_cache_success = self.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"
         else:
             logger.error(message)
@@ -98,7 +102,9 @@ class SchedulerUpdateWeightsMixin:
         # 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()
+                flush_cache_success = self.flush_cache(
+                    empty_cache=recv_req.torch_empty_cache
+                )
                 assert flush_cache_success, "Cache flush failed after updating weights"
         else:
             logger.error(message)
@@ -114,7 +120,9 @@ class SchedulerUpdateWeightsMixin:
         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()
+            flush_cache_success = self.flush_cache(
+                empty_cache=recv_req.torch_empty_cache
+            )
             assert flush_cache_success, "Cache flush failed after updating weights"
         if not success:
             logger.error(message)