CUDA-graph-compatible releasing and resuming KV cache and model weight memory

python/pyproject.toml

@@ -38,6 +38,7 @@ srt_hpu = ["sglang[runtime_common]"]
 openai = ["openai>=1.0", "tiktoken"]
 anthropic = ["anthropic>=0.20.0"]
 litellm = ["litellm>=1.0.0"]
+torch_memory_saver = ["torch_memory_saver"]
 test = [
     "jsonlines",
     "matplotlib",

python/sglang/srt/managers/io_struct.py

@@ -469,6 +469,26 @@ class GetWeightsByNameReqOutput:
     parameter: list
 
 
+@dataclass
+class ReleaseGPUOccupationReqInput:
+    pass
+
+
+@dataclass
+class ReleaseGPUOccupationReqOutput:
+    pass
+
+
+@dataclass
+class ResumeGPUOccupationReqInput:
+    pass
+
+
+@dataclass
+class ResumeGPUOccupationReqOutput:
+    pass
+
+
 @dataclass
 class AbortReq:
     # The request id

python/sglang/srt/managers/scheduler.py

@@ -46,6 +46,10 @@
     OpenSessionReqInput,
     OpenSessionReqOutput,
     ProfileReq,
+    ReleaseGPUOccupationReqInput,
+    ReleaseGPUOccupationReqOutput,
+    ResumeGPUOccupationReqInput,
+    ResumeGPUOccupationReqOutput,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
     UpdateWeightFromDiskReqInput,
@@ -86,6 +90,7 @@
     set_random_seed,
     suppress_other_loggers,
 )
+from sglang.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -332,6 +337,10 @@ def __init__(
         t.start()
         self.parent_process = psutil.Process().parent()
 
+        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=server_args.memory_saver
+        )
+
         # Init profiler
         if os.getenv("SGLANG_TORCH_PROFILER_DIR", "") == "":
             self.profiler = None
@@ -492,6 +501,12 @@ def process_input_requests(self, recv_reqs: List):
             elif isinstance(recv_req, GetWeightsByNameReqInput):
                 parameter = self.get_weights_by_name(recv_req)
                 self.send_to_tokenizer.send_pyobj(GetWeightsByNameReqOutput(parameter))
+            elif isinstance(recv_req, ReleaseGPUOccupationReqInput):
+                self.release_gpu_occupation()
+                self.send_to_tokenizer.send_pyobj(ReleaseGPUOccupationReqOutput())
+            elif isinstance(recv_req, ResumeGPUOccupationReqInput):
+                self.resume_gpu_occupation()
+                self.send_to_tokenizer.send_pyobj(ResumeGPUOccupationReqOutput())
             elif isinstance(recv_req, ProfileReq):
                 if recv_req == ProfileReq.START_PROFILE:
                     self.start_profile()
@@ -1497,6 +1512,20 @@ def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
         parameter = self.tp_worker.get_weights_by_name(recv_req)
         return parameter
 
+    def release_gpu_occupation(self):
+        self.stashed_model_static_state = (
+            self.tp_worker.worker.model_runner.model.export_static_state()
+        )
+        self.memory_saver_adapter.pause()
+        self.flush_cache()
+
+    def resume_gpu_occupation(self):
+        self.memory_saver_adapter.resume()
+        self.tp_worker.worker.model_runner.model.import_static_state(
+            self.stashed_model_static_state
+        )
+        del self.stashed_model_static_state
+
     def start_profile(self) -> None:
         if self.profiler is None:
             raise RuntimeError("Profiler is not enabled.")

python/sglang/srt/managers/tokenizer_manager.py

@@ -53,6 +53,10 @@
     OpenSessionReqInput,
     OpenSessionReqOutput,
     ProfileReq,
+    ReleaseGPUOccupationReqInput,
+    ReleaseGPUOccupationReqOutput,
+    ResumeGPUOccupationReqInput,
+    ResumeGPUOccupationReqOutput,
     SessionParams,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -188,6 +192,12 @@ def __init__(
         self.get_weights_by_name_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
+        self.release_gpu_occupation_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
+        self.resume_gpu_occupation_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
 
         # Metrics
         if self.enable_metrics:
@@ -548,6 +558,22 @@ async def get_weights_by_name(
         else:
             return all_parameters
 
+    async def release_gpu_occupation(
+        self,
+        obj: ReleaseGPUOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.release_gpu_occupation_communicator(obj)
+
+    async def resume_gpu_occupation(
+        self,
+        obj: ResumeGPUOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.resume_gpu_occupation_communicator(obj)
+
     async def open_session(
         self, obj: OpenSessionReqInput, request: Optional[fastapi.Request] = None
     ):
@@ -627,6 +653,8 @@ async def handle_loop(self):
                 UpdateWeightsFromDistributedReqOutput,
                 GetWeightsByNameReqOutput,
                 InitWeightsUpdateGroupReqOutput,
+                ReleaseGPUOccupationReqOutput,
+                ResumeGPUOccupationReqOutput,
             ] = await self.recv_from_detokenizer.recv_pyobj()
 
             if isinstance(recv_obj, (BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut)):
@@ -750,6 +778,10 @@ async def handle_loop(self):
                 self.update_weights_from_tensor_communicator.handle_recv(recv_obj)
             elif isinstance(recv_obj, GetWeightsByNameReqOutput):
                 self.get_weights_by_name_communicator.handle_recv(recv_obj)
+            elif isinstance(recv_obj, ReleaseGPUOccupationReqOutput):
+                self.release_gpu_occupation_communicator.handle_recv(recv_obj)
+            elif isinstance(recv_obj, ResumeGPUOccupationReqOutput):
+                self.resume_gpu_occupation_communicator.handle_recv(recv_obj)
             else:
                 raise ValueError(f"Invalid object: {recv_obj=}")
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -13,6 +13,8 @@
 limitations under the License.
 """
 
+from sglang.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
 """
 Memory pool.
 
@@ -35,13 +37,21 @@
 class ReqToTokenPool:
     """A memory pool that maps a request to its token locations."""
 
-    def __init__(self, size: int, max_context_len: int, device: str, use_records: bool):
+    def __init__(
+        self,
+        size: int,
+        max_context_len: int,
+        device: str,
+        use_records: bool,
+        memory_saver_adapter: TorchMemorySaverAdapter,
+    ):
         self.size = size
         self.max_context_len = max_context_len
         self.device = device
-        self.req_to_token = torch.zeros(
-            (size, max_context_len), dtype=torch.int32, device=device
-        )
+        with memory_saver_adapter.region():
+            self.req_to_token = torch.zeros(
+                (size, max_context_len), dtype=torch.int32, device=device
+            )
         self.free_slots = list(range(size))
         self.write_records = []
         self.use_records = use_records
@@ -182,32 +192,35 @@ def __init__(
         head_dim: int,
         layer_num: int,
         device: str,
+        memory_saver_adapter: TorchMemorySaverAdapter,
     ):
         super().__init__(size, dtype, device)
         self.head_num = head_num
         self.head_dim = head_dim
         self.layer_num = layer_num
+        self.memory_saver_adapter = memory_saver_adapter
         self._create_buffers()
 
     def _create_buffers(self):
-        # [size, head_num, head_dim] for each layer
-        # The padded slot 0 is used for writing dummy outputs from padded tokens.
-        self.k_buffer = [
-            torch.empty(
-                (self.size + 1, self.head_num, self.head_dim),
-                dtype=self.store_dtype,
-                device=self.device,
-            )
-            for _ in range(self.layer_num)
-        ]
-        self.v_buffer = [
-            torch.empty(
-                (self.size + 1, self.head_num, self.head_dim),
-                dtype=self.store_dtype,
-                device=self.device,
-            )
-            for _ in range(self.layer_num)
-        ]
+        with self.memory_saver_adapter.region():
+            # [size, head_num, head_dim] for each layer
+            # The padded slot 0 is used for writing dummy outputs from padded tokens.
+            self.k_buffer = [
+                torch.empty(
+                    (self.size + 1, self.head_num, self.head_dim),
+                    dtype=self.store_dtype,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
+            self.v_buffer = [
+                torch.empty(
+                    (self.size + 1, self.head_num, self.head_dim),
+                    dtype=self.store_dtype,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
 
     def _clear_buffers(self):
         del self.k_buffer
@@ -262,19 +275,22 @@ def __init__(
         qk_rope_head_dim: int,
         layer_num: int,
         device: str,
+        memory_saver_adapter: TorchMemorySaverAdapter,
     ):
         super().__init__(size, dtype, device)
 
         self.kv_lora_rank = kv_lora_rank
-        # The padded slot 0 is used for writing dummy outputs from padded tokens.
-        self.kv_buffer = [
-            torch.empty(
-                (size + 1, 1, kv_lora_rank + qk_rope_head_dim),
-                dtype=self.store_dtype,
-                device=device,
-            )
-            for _ in range(layer_num)
-        ]
+
+        with memory_saver_adapter.region():
+            # The padded slot 0 is used for writing dummy outputs from padded tokens.
+            self.kv_buffer = [
+                torch.empty(
+                    (size + 1, 1, kv_lora_rank + qk_rope_head_dim),
+                    dtype=self.store_dtype,
+                    device=device,
+                )
+                for _ in range(layer_num)
+            ]
 
     def get_key_buffer(self, layer_id: int):
         if self.store_dtype != self.dtype:
@@ -315,26 +331,28 @@ def __init__(
         layer_num: int,
         device: str,
         heavy_channel_num: int,
+        memory_saver_adapter: TorchMemorySaverAdapter,
     ):
         super().__init__(size, dtype, device)
 
-        # [size, head_num, head_dim] for each layer
-        self.k_buffer = [
-            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
-            for _ in range(layer_num)
-        ]
-        self.v_buffer = [
-            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
-            for _ in range(layer_num)
-        ]
-
-        # [size, head_num, heavy_channel_num] for each layer
-        self.label_buffer = [
-            torch.empty(
-                (size + 1, head_num, heavy_channel_num), dtype=dtype, device=device
-            )
-            for _ in range(layer_num)
-        ]
+        with memory_saver_adapter.region():
+            # [size, head_num, head_dim] for each layer
+            self.k_buffer = [
+                torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
+                for _ in range(layer_num)
+            ]
+            self.v_buffer = [
+                torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
+                for _ in range(layer_num)
+            ]
+
+            # [size, head_num, heavy_channel_num] for each layer
+            self.label_buffer = [
+                torch.empty(
+                    (size + 1, head_num, heavy_channel_num), dtype=dtype, device=device
+                )
+                for _ in range(layer_num)
+            ]
 
     def get_key_buffer(self, layer_id: int):
         return self.k_buffer[layer_id]