Add release_kv_cache() API

tests/basic_correctness/test_mem.py

@@ -75,6 +75,34 @@ def test_basic_cumem():
     assert torch.allclose(output, torch.ones_like(output) * 3)
 
 
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+def test_release_kv_cache():
+    kv_cache_memory_bytes = 256 * 1024 * 1024
+    llm = LLM(
+        "Qwen/Qwen3-0.6B",
+        enable_sleep_mode=True,
+        enforce_eager=True,
+        gpu_memory_utilization=0.1,
+        kv_cache_memory_bytes=kv_cache_memory_bytes,
+        max_model_len=1024,
+        max_num_seqs=4,
+    )
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    output = llm.generate(prompt, sampling_params)
+
+    free_bytes = current_platform.mem_get_info()[0]
+    llm.release_kv_cache()
+    free_bytes_after_release = current_platform.mem_get_info()[0]
+    freed_bytes = free_bytes_after_release - free_bytes
+    assert freed_bytes >= kv_cache_memory_bytes * 0.99
+
+    llm.wake_up()
+    output2 = llm.generate(prompt, sampling_params)
+
+    assert output[0].outputs[0].text == output2[0].outputs[0].text
+
+
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
 @pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
 def test_cumem_with_cudagraph():

tests/entrypoints/serve/dev/test_sleep.py

@@ -35,8 +35,11 @@ def test_sleep_mode():
         # check sleep metrics
         response = requests.get(remote_server.url_for("metrics"))
         assert response.status_code == 200
-        awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response)
+        awake, kv_cache_released, weights_offloaded, discard_all = (
+            _get_sleep_metrics_from_api(response)
+        )
         assert awake == 0
+        assert kv_cache_released == 0
         assert weights_offloaded == 1
         assert discard_all == 0
 
@@ -49,11 +52,31 @@ def test_sleep_mode():
         # check sleep metrics
         response = requests.get(remote_server.url_for("metrics"))
         assert response.status_code == 200
-        awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response)
+        awake, kv_cache_released, weights_offloaded, discard_all = (
+            _get_sleep_metrics_from_api(response)
+        )
         assert awake == 1
+        assert kv_cache_released == 0
         assert weights_offloaded == 0
         assert discard_all == 0
 
+        response = requests.post(remote_server.url_for("release_kv_cache"))
+        assert response.status_code == 200
+
+        # check KV cache release metrics
+        response = requests.get(remote_server.url_for("metrics"))
+        assert response.status_code == 200
+        awake, kv_cache_released, weights_offloaded, discard_all = (
+            _get_sleep_metrics_from_api(response)
+        )
+        assert awake == 0
+        assert kv_cache_released == 1
+        assert weights_offloaded == 0
+        assert discard_all == 0
+
+        response = requests.post(remote_server.url_for("wake_up"))
+        assert response.status_code == 200
+
         # test wake up with tags
         response = requests.post(remote_server.url_for("sleep"), params={"level": "1"})
         assert response.status_code == 200
@@ -80,16 +103,19 @@ def test_sleep_mode():
         # check sleep metrics
         response = requests.get(remote_server.url_for("metrics"))
         assert response.status_code == 200
-        awake, weights_offloaded, discard_all = _get_sleep_metrics_from_api(response)
+        awake, kv_cache_released, weights_offloaded, discard_all = (
+            _get_sleep_metrics_from_api(response)
+        )
         assert awake == 1
+        assert kv_cache_released == 0
         assert weights_offloaded == 0
         assert discard_all == 0
 
 
 def _get_sleep_metrics_from_api(response: requests.Response):
-    """Return (awake, weights_offloaded, discard_all)"""
+    """Return (awake, kv_cache_released, weights_offloaded, discard_all)"""
 
-    awake, weights_offloaded, discard_all = None, None, None
+    awake, kv_cache_released, weights_offloaded, discard_all = None, None, None, None
 
     for family in text_string_to_metric_families(response.text):
         if family.name == "vllm:engine_sleep_state":
@@ -98,13 +124,16 @@ def _get_sleep_metrics_from_api(response: requests.Response):
                     for label_name, label_value in sample.labels.items():
                         if label_value == "awake":
                             awake = sample.value
+                        elif label_value == "kv_cache_released":
+                            kv_cache_released = sample.value
                         elif label_value == "weights_offloaded":
                             weights_offloaded = sample.value
                         elif label_value == "discard_all":
                             discard_all = sample.value
 
     assert awake is not None
+    assert kv_cache_released is not None
     assert weights_offloaded is not None
     assert discard_all is not None
 
-    return awake, weights_offloaded, discard_all
+    return awake, kv_cache_released, weights_offloaded, discard_all

vllm/device_allocator/__init__.py

@@ -18,13 +18,16 @@ class AllocationData:
     handle: HandleType
     tag: str
     cpu_backup_tensor: torch.Tensor | None = None
+    is_mapped: bool = True
 
 
 class MemAllocator(Protocol):
     def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ...
 
     def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ...
 
+    def release_tags(self, tags: tuple[str, ...] | str) -> None: ...
+
     def wake_up(self, tags: list[str] | None = None) -> None: ...
 
     def get_current_usage(self) -> int: ...

vllm/device_allocator/cumem.py

@@ -188,6 +188,8 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
 
         for ptr, data in self.pointer_to_data.items():
             handle = data.handle
+            if not data.is_mapped:
+                continue
             total_bytes += handle[1]
             if data.tag in offload_tags:
                 backup_bytes += handle[1]
@@ -202,6 +204,7 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
                 libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)
                 data.cpu_backup_tensor = cpu_backup_tensor
             unmap_and_release(handle)
+            data.is_mapped = False
 
         logger.info(
             "CuMemAllocator: sleep freed %.2f GiB memory in total, of which "
@@ -215,6 +218,28 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
         gc.collect()
         torch.cuda.empty_cache()
 
+    def release_tags(self, tags: tuple[str, ...] | str) -> None:
+        if isinstance(tags, str):
+            tags = (tags,)
+
+        released_bytes = 0
+        for data in self.pointer_to_data.values():
+            if data.tag not in tags or not data.is_mapped:
+                continue
+            released_bytes += data.handle[1]
+            data.cpu_backup_tensor = None
+            unmap_and_release(data.handle)
+            data.is_mapped = False
+
+        logger.info(
+            "CuMemAllocator: released %.2f GiB memory for tags %s.",
+            released_bytes / 1024**3,
+            tags,
+        )
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def wake_up(self, tags: list[str] | None = None) -> None:
         """
         Wake up the allocator from sleep mode.
@@ -228,8 +253,11 @@ def wake_up(self, tags: list[str] | None = None) -> None:
         """
         for ptr, data in self.pointer_to_data.items():
             if tags is None or data.tag in tags:
+                if data.is_mapped:
+                    continue
                 handle = data.handle
                 create_and_map(handle)
+                data.is_mapped = True
                 if data.cpu_backup_tensor is not None:
                     cpu_backup_tensor = data.cpu_backup_tensor
                     if cpu_backup_tensor is not None:

vllm/device_allocator/xpumem.py

@@ -177,9 +177,12 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
 
         for ptr, data in self.pointer_to_data.items():
             size_in_bytes = data.handle[1]
+            if not data.is_mapped:
+                continue
             total_bytes += size_in_bytes
             if data.tag not in offload_tags:
                 unmap_and_release(data.handle)
+                data.is_mapped = False
                 continue
 
             backup_bytes += size_in_bytes
@@ -201,6 +204,7 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
             data.cpu_backup_tensor = cpu_backup_tensor
 
             unmap_and_release(data.handle)
+            data.is_mapped = False
 
         logger.info(
             "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which "
@@ -216,11 +220,38 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
         if callable(xpu_empty_cache):
             xpu_empty_cache()
 
+    def release_tags(self, tags: tuple[str, ...] | str) -> None:
+        if isinstance(tags, str):
+            tags = (tags,)
+
+        released_bytes = 0
+        for data in self.pointer_to_data.values():
+            if data.tag not in tags or not data.is_mapped:
+                continue
+            released_bytes += data.handle[1]
+            data.cpu_backup_tensor = None
+            unmap_and_release(data.handle)
+            data.is_mapped = False
+
+        logger.info(
+            "XpuMemAllocator: released %.2f GiB memory for tags %s.",
+            released_bytes / 1024**3,
+            tags,
+        )
+
+        gc.collect()
+        xpu_empty_cache = getattr(torch.xpu, "empty_cache", None)
+        if callable(xpu_empty_cache):
+            xpu_empty_cache()
+
     def wake_up(self, tags: list[str] | None = None) -> None:
         for ptr, data in self.pointer_to_data.items():
             if tags is not None and data.tag not in tags:
                 continue
+            if data.is_mapped:
+                continue
             create_and_allocate(data.handle)
+            data.is_mapped = True
 
             cpu_backup_tensor = data.cpu_backup_tensor
             if cpu_backup_tensor is None:

vllm/engine/protocol.py

@@ -165,6 +165,11 @@ async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None:
         """Sleep the engine"""
         ...
 
+    @abstractmethod
+    async def release_kv_cache(self, mode: "PauseMode" = "abort") -> bool:
+        """Release the engine's KV cache"""
+        ...
+
     @abstractmethod
     async def wake_up(self, tags: list[str] | None = None) -> None:
         """Wake up the engine"""

vllm/entrypoints/llm.py

@@ -829,6 +829,19 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"):
         """
         self.llm_engine.sleep(level=level, mode=mode)
 
+    def release_kv_cache(self, mode: PauseMode = "abort") -> bool:
+        """
+        Release the engine's KV cache memory while keeping model weights resident.
+
+        The engine will stop processing requests until `wake_up(tags=["kv_cache"])`
+        is called.
+
+        Args:
+            mode: How to handle any existing requests, can be "abort", "wait",
+                or "keep".
+        """
+        return self.llm_engine.release_kv_cache(mode=mode)
+
     def wake_up(self, tags: list[str] | None = None):
         """
         Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]

vllm/entrypoints/serve/dev/sleep/api_router.py

@@ -29,6 +29,13 @@ async def sleep(raw_request: Request):
     return Response(status_code=200)
 
 
+@router.post("/release_kv_cache")
+async def release_kv_cache(raw_request: Request):
+    mode = raw_request.query_params.get("mode", "abort")
+    await engine_client(raw_request).release_kv_cache(mode)
+    return Response(status_code=200)
+
+
 @router.post("/wake_up")
 async def wake_up(raw_request: Request):
     tags = raw_request.query_params.getlist("tags")

vllm/v1/engine/async_llm.py

@@ -936,6 +936,13 @@ async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         if self.logger_manager is not None:
             self.logger_manager.record_sleep_state(1, level)
 
+    async def release_kv_cache(self, mode: PauseMode = "abort") -> bool:
+        released = await self.engine_core.release_kv_cache_async(mode)
+
+        if released and self.logger_manager is not None:
+            self.logger_manager.record_kv_cache_released_state()
+        return released
+
     async def wake_up(self, tags: list[str] | None = None) -> None:
         await self.engine_core.wake_up_async(tags)
 

vllm/v1/engine/core.py

@@ -778,6 +778,28 @@ def wake_up(self, tags: list[str] | None = None):
         # Resume scheduling (applies to all levels)
         self.resume_scheduler()
 
+    def release_kv_cache(self, mode: PauseMode = "abort") -> bool | Future:
+        """Release KV cache memory while keeping model weights resident."""
+        pause_future = self.pause_scheduler(mode=mode, clear_cache=True)
+        model_executor = self.model_executor
+        if pause_future is None:
+            return model_executor.release_kv_cache()
+
+        future = Future[Any]()
+
+        def pause_complete(f: Future):
+            try:
+                f.result()
+                future.set_result(model_executor.release_kv_cache())
+            except Exception as e:
+                future.set_exception(e)
+
+        logger.info(
+            "Waiting for in-flight requests to complete before releasing KV cache..."
+        )
+        pause_future.add_done_callback(pause_complete)
+        return future
+
     def is_sleeping(self) -> bool:
         """Check if engine is sleeping at any level."""
         return self.is_scheduler_paused() or self.model_executor.is_sleeping

vllm/v1/engine/core_client.py

@@ -159,6 +159,9 @@ def reset_encoder_cache(self) -> None:
     def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         raise NotImplementedError
 
+    def release_kv_cache(self, mode: PauseMode = "abort") -> bool:
+        raise NotImplementedError
+
     def wake_up(self, tags: list[str] | None = None) -> None:
         raise NotImplementedError
 
@@ -236,6 +239,9 @@ async def reset_encoder_cache_async(self) -> None:
     async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
         raise NotImplementedError
 
+    async def release_kv_cache_async(self, mode: PauseMode = "abort") -> bool:
+        raise NotImplementedError
+
     async def wake_up_async(self, tags: list[str] | None = None) -> None:
         raise NotImplementedError
 
@@ -326,6 +332,13 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         result = self.engine_core.sleep(level, mode)
         assert result is None
 
+    def release_kv_cache(self, mode: PauseMode = "abort") -> bool:
+        if mode == "wait":
+            raise ValueError("'wait' pause mode is not supported in inproc-engine mode")
+        result = self.engine_core.release_kv_cache(mode)
+        assert isinstance(result, bool)
+        return result
+
     def wake_up(self, tags: list[str] | None = None) -> None:
         self.engine_core.wake_up(tags)
 
@@ -908,6 +921,9 @@ def pin_lora(self, lora_id: int) -> bool:
     def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         self.call_utility("sleep", level, mode)
 
+    def release_kv_cache(self, mode: PauseMode = "abort") -> bool:
+        return self.call_utility("release_kv_cache", mode)
+
     def wake_up(self, tags: list[str] | None = None) -> None:
         self.call_utility("wake_up", tags)
 
@@ -1144,6 +1160,9 @@ async def reset_encoder_cache_async(self) -> None:
     async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
         await self.call_utility_async("sleep", level, mode)
 
+    async def release_kv_cache_async(self, mode: PauseMode = "abort") -> bool:
+        return await self.call_utility_async("release_kv_cache", mode)
+
     async def wake_up_async(self, tags: list[str] | None = None) -> None:
         await self.call_utility_async("wake_up", tags)