diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 340ffd545..1722f75d6 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1378,7 +1378,12 @@ class PauseGenerationReqInput(BaseReq): @dataclass class ContinueGenerationReqInput(BaseReq): - pass + # Call torch.cuda.empty_cache() before un-pausing. Returns blocks + # cached by the PyTorch allocator (left over from transient allocs + # during post-weight-update processing) back to the driver before + # inference resumes, with no race against active streams. Set to + # False to skip the empty_cache call. + torch_empty_cache: bool = True @dataclass diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8ef38bdfa..6e662e77b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3646,6 +3646,15 @@ class Scheduler( self.chunked_req = None def continue_generation(self, recv_req: ContinueGenerationReqInput): + if recv_req.torch_empty_cache: + before_mb = torch.cuda.memory_reserved() / (1024 * 1024) + torch.cuda.empty_cache() + after_mb = torch.cuda.memory_reserved() / (1024 * 1024) + logger.info( + f"[continue_generation] torch.cuda.empty_cache() called: " + f"reserved {before_mb:.1f} MB -> {after_mb:.1f} MB " + f"(freed {before_mb - after_mb:.1f} MB)" + ) self._engine_paused = False def load_lora_adapter(