From 8471c9ebe61229fe24c99ca4a2b2c5b9ac576e43 Mon Sep 17 00:00:00 2001
From: Shenxiu Liu <azhen9166@gmail.com>
Date: Fri, 24 Apr 2026 21:41:39 -0700
Subject: [PATCH] Skip torch.cuda.empty_cache() in weight update flush path
 (#22998)

---
 python/sglang/srt/managers/io_struct.py          |  6 ++++++
 python/sglang/srt/managers/scheduler.py          |  6 +++---
 .../managers/scheduler_update_weights_mixin.py   | 16 ++++++++++++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 5b84677aa..47f4062b2 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/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
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index a6916a88b..4c1b82da9 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/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:
diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py
index 3ff0cc430..bfb6f084a 100644
--- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py
+++ b/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)