minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Shenxiu Liu
2026-04-25 12:41:39 +08:00
committed by GitHub
parent 69485a176c
commit 8471c9ebe6
3 changed files with 21 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+6
View File
@@ -1418,6 +1418,8 @@ class UpdateWeightsFromDistributedReqInput(BaseReq):
weight_version: Optional[str] = None weight_version: Optional[str] = None
# Optional format specification for loading # Optional format specification for loading
load_format: Optional[str] = None load_format: Optional[str] = None
# Whether to call torch.cuda.empty_cache() during flush
torch_empty_cache: bool = False
@dataclass @dataclass
@@ -1445,6 +1447,8 @@ class UpdateWeightsFromTensorReqInput(BaseReq):
weight_version: Optional[str] = None weight_version: Optional[str] = None
# Optional: Determine whether to disable updating the draft model # Optional: Determine whether to disable updating the draft model
disable_draft_model: Optional[bool] = None disable_draft_model: Optional[bool] = None
# Whether to call torch.cuda.empty_cache() during flush
torch_empty_cache: bool = False
@dataclass @dataclass
@@ -1479,6 +1483,8 @@ class UpdateWeightsFromIPCReqInput(BaseReq):
flush_cache: bool = True flush_cache: bool = True
# Optional: Update weight version along with weights # Optional: Update weight version along with weights
weight_version: Optional[str] = None weight_version: Optional[str] = None
# Whether to call torch.cuda.empty_cache() during flush
torch_empty_cache: bool = False
@dataclass @dataclass
python/sglang/srt/managers/scheduler.py
+3 -3
View File
@@ -3226,7 +3226,7 @@ class Scheduler(
return DetachHiCacheStorageReqOutput(success=False, message=msg) return DetachHiCacheStorageReqOutput(success=False, message=msg)
def flush_cache(self): def flush_cache(self, empty_cache: bool = True):
"""Flush the memory pool and cache.""" """Flush the memory pool and cache."""
if self.is_fully_idle(): if self.is_fully_idle():
self.cur_batch = None self.cur_batch = None
@@ -3240,8 +3240,8 @@ class Scheduler(
if self.draft_worker: if self.draft_worker:
self.draft_worker.clear_cache_pool() self.draft_worker.clear_cache_pool()
# TODO: allow optional empty cache if empty_cache:
torch.cuda.empty_cache() torch.cuda.empty_cache()
logger.info("Cache flushed successfully!") logger.info("Cache flushed successfully!")
success = True success = True
else: else:
python/sglang/srt/managers/scheduler_update_weights_mixin.py
+12 -4
View File
@@ -52,7 +52,9 @@ class SchedulerUpdateWeightsMixin:
if success and self.draft_worker is not None: if success and self.draft_worker is not None:
success, message = self.draft_worker.update_weights_from_disk(recv_req) success, message = self.draft_worker.update_weights_from_disk(recv_req)
if tp_success and recv_req.flush_cache: 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" assert flush_cache_success, "Cache flush failed after updating weights"
if not success: if not success:
logger.error(message) logger.error(message)
@@ -80,7 +82,9 @@ class SchedulerUpdateWeightsMixin:
success, message = self.tp_worker.update_weights_from_distributed(recv_req) success, message = self.tp_worker.update_weights_from_distributed(recv_req)
if success: if success:
if recv_req.flush_cache: 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" assert flush_cache_success, "Cache flush failed after updating weights"
else: else:
logger.error(message) 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 # TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
if success: if success:
if recv_req.flush_cache: 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" assert flush_cache_success, "Cache flush failed after updating weights"
else: else:
logger.error(message) logger.error(message)
@@ -114,7 +120,9 @@ class SchedulerUpdateWeightsMixin:
if success and self.draft_worker is not None: if success and self.draft_worker is not None:
success, message = self.draft_worker.update_weights_from_ipc(recv_req) success, message = self.draft_worker.update_weights_from_ipc(recv_req)
if tp_success and recv_req.flush_cache: 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" assert flush_cache_success, "Cache flush failed after updating weights"
if not success: if not success:
logger.error(message) logger.error(message)