From 3004f99a0e42c5a4e74adafd4baaae1eac3a87df Mon Sep 17 00:00:00 2001 From: Dakai An Date: Mon, 8 Jun 2026 14:56:33 +0000 Subject: [PATCH] add release_kv_cache() Signed-off-by: Dakai An --- tests/basic_correctness/test_mem.py | 28 +++++++++++++ tests/entrypoints/serve/dev/test_sleep.py | 41 ++++++++++++++++--- vllm/device_allocator/__init__.py | 3 ++ vllm/device_allocator/cumem.py | 28 +++++++++++++ vllm/device_allocator/xpumem.py | 31 ++++++++++++++ vllm/engine/protocol.py | 5 +++ vllm/entrypoints/llm.py | 13 ++++++ .../entrypoints/serve/dev/sleep/api_router.py | 7 ++++ vllm/v1/engine/async_llm.py | 7 ++++ vllm/v1/engine/core.py | 22 ++++++++++ vllm/v1/engine/core_client.py | 19 +++++++++ vllm/v1/engine/llm_engine.py | 7 ++++ vllm/v1/executor/abstract.py | 27 ++++++++++-- vllm/v1/metrics/loggers.py | 27 +++++++++++- vllm/v1/worker/cpu_worker.py | 4 ++ vllm/v1/worker/gpu_worker.py | 14 +++++++ 16 files changed, 272 insertions(+), 11 deletions(-) diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 2c9a99c500d9..d1a22f6c5360 100644 --- a/tests/basic_correctness/test_mem.py +++ b/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(): diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py index 260dcd00bae9..aea9b60374ef 100644 --- a/tests/entrypoints/serve/dev/test_sleep.py +++ b/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 diff --git a/vllm/device_allocator/__init__.py b/vllm/device_allocator/__init__.py index 6b5e9c613d08..6d9ccb6f084e 100644 --- a/vllm/device_allocator/__init__.py +++ b/vllm/device_allocator/__init__.py @@ -18,6 +18,7 @@ class AllocationData: handle: HandleType tag: str cpu_backup_tensor: torch.Tensor | None = None + is_mapped: bool = True class MemAllocator(Protocol): @@ -25,6 +26,8 @@ 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: ... diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index c30790df9cac..6f540270c799 100644 --- a/vllm/device_allocator/cumem.py +++ b/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: diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 7d99ced7ef54..629eb04a5790 100644 --- a/vllm/device_allocator/xpumem.py +++ b/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: diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 3f83734a5b78..a2e378181993 100644 --- a/vllm/engine/protocol.py +++ b/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""" diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 892e5035ab60..3ed8e5595b2c 100644 --- a/vllm/entrypoints/llm.py +++ b/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] diff --git a/vllm/entrypoints/serve/dev/sleep/api_router.py b/vllm/entrypoints/serve/dev/sleep/api_router.py index 0861c8677325..770f9ee37f84 100644 --- a/vllm/entrypoints/serve/dev/sleep/api_router.py +++ b/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") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 419e15163a9f..c8c67ce7ca82 100644 --- a/vllm/v1/engine/async_llm.py +++ b/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) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 08c814ab34e0..0c4371b7b616 100644 --- a/vllm/v1/engine/core.py +++ b/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 diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 32f2d091eb30..558f442c7bba 100644 --- a/vllm/v1/engine/core_client.py +++ b/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) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index f3e8a95b0d63..c1c7a1220e13 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -357,6 +357,13 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"): if self.logger_manager is not None: self.logger_manager.record_sleep_state(1, level) + def release_kv_cache(self, mode: PauseMode = "abort") -> bool: + released = self.engine_core.release_kv_cache(mode) + + if released and self.logger_manager is not None: + self.logger_manager.record_kv_cache_released_state() + return released + def wake_up(self, tags: list[str] | None = None): self.engine_core.wake_up(tags) diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 4063844d469c..b4e2353019f0 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -317,8 +317,11 @@ def reset_encoder_cache(self) -> None: def sleep(self, level: int = 1): if self.is_sleeping: - logger.warning("Executor is already sleeping.") - return + if self.sleeping_tags == {"kv_cache"}: + logger.info("Executor is releasing weights after KV cache release.") + else: + logger.warning("Executor is already sleeping.") + return time_before_sleep = time.perf_counter() self.collective_rpc("sleep", kwargs=dict(level=level)) time_after_sleep = time.perf_counter() @@ -339,13 +342,14 @@ def wake_up(self, tags: list[str] | None = None): "Tag %s is not in sleeping tags %s", tag, self.sleeping_tags ) return + tags_to_wake = tags if tags is not None else list(self.sleeping_tags) time_before_wakeup = time.perf_counter() - self.collective_rpc("wake_up", kwargs=dict(tags=tags)) + self.collective_rpc("wake_up", kwargs=dict(tags=tags_to_wake)) time_after_wakeup = time.perf_counter() logger.info( "It took %.6f seconds to wake up tags %s.", time_after_wakeup - time_before_wakeup, - tags if tags is not None else self.sleeping_tags, + tags_to_wake, ) if tags: for tag in tags: @@ -355,6 +359,21 @@ def wake_up(self, tags: list[str] | None = None): if not self.sleeping_tags: self.is_sleeping = False + def release_kv_cache(self) -> bool: + if self.is_sleeping: + logger.warning("Executor is already sleeping.") + return False + time_before_release = time.perf_counter() + self.collective_rpc("release_kv_cache") + time_after_release = time.perf_counter() + self.sleeping_tags = {"kv_cache"} + self.is_sleeping = True + logger.info( + "It took %.6f seconds to release KV cache.", + time_after_release - time_before_release, + ) + return True + def reinitialize_distributed( self, reconfig_request: ReconfigureDistributedRequest ) -> None: diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 0052a35366a3..18690f18bb49 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -70,6 +70,9 @@ def log(self): # noqa def record_sleep_state(self, is_awake: int, level: int): # noqa pass + def record_kv_cache_released_state(self): # noqa + pass + def load_stat_logger_plugin_factories() -> list[StatLoggerFactory]: factories: list[StatLoggerFactory] = [] @@ -495,6 +498,8 @@ def __init__( documentation=( "Engine sleep state; awake = 0 means engine is sleeping; " "awake = 1 means engine is awake; " + "kv_cache_released = 1 means KV cache is released while " + "weights are resident; " "weights_offloaded = 1 means sleep level 1; " "discard_all = 1 means sleep level 2." ), @@ -503,7 +508,12 @@ def __init__( ) self.gauge_engine_sleep_state = {} - sleep_state = ["awake", "weights_offloaded", "discard_all"] + sleep_state = [ + "awake", + "kv_cache_released", + "weights_offloaded", + "discard_all", + ] for s in sleep_state: self.gauge_engine_sleep_state[s] = { @@ -1218,6 +1228,7 @@ def record( def record_sleep_state(self, sleep: int = 0, level: int = 0): awake = 1 + kv_cache_released = 0 discard_all = 0 weights_offloaded = 0 @@ -1229,12 +1240,22 @@ def record_sleep_state(self, sleep: int = 0, level: int = 0): discard_all = 1 for engine_idx in self.engine_indexes: + self.gauge_engine_sleep_state["kv_cache_released"][engine_idx].set( + kv_cache_released + ) self.gauge_engine_sleep_state["discard_all"][engine_idx].set(discard_all) self.gauge_engine_sleep_state["weights_offloaded"][engine_idx].set( weights_offloaded ) self.gauge_engine_sleep_state["awake"][engine_idx].set(awake) + def record_kv_cache_released_state(self): + for engine_idx in self.engine_indexes: + self.gauge_engine_sleep_state["kv_cache_released"][engine_idx].set(1) + self.gauge_engine_sleep_state["discard_all"][engine_idx].set(0) + self.gauge_engine_sleep_state["weights_offloaded"][engine_idx].set(0) + self.gauge_engine_sleep_state["awake"][engine_idx].set(0) + def log_engine_initialized(self): self.log_metrics_info("cache_config", self.vllm_config.cache_config) @@ -1351,6 +1372,10 @@ def record_sleep_state(self, sleep: int = 0, level: int = 0): for logger in self.stat_loggers: logger.record_sleep_state(sleep, level) + def record_kv_cache_released_state(self): + for logger in self.stat_loggers: + logger.record_kv_cache_released_state() + def log(self): for logger in self.stat_loggers: logger.log() diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index 2433bc8a10ba..2a6116b2798a 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -177,6 +177,10 @@ def wake_up(self, tags: list[str] | None = None) -> None: logger.warning("sleep mode is not supported on CPU, ignore it.") pass + def release_kv_cache(self) -> None: + logger.warning("release_kv_cache is not supported on CPU, ignore it.") + pass + def determine_available_memory(self) -> int: self.model_runner.warming_up_model() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 052e1fe76f43..2385fb89c697 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -199,6 +199,20 @@ def wake_up(self, tags: list[str] | None = None) -> None: if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() + def release_kv_cache(self) -> None: + free_bytes_before_release = torch.cuda.mem_get_info()[0] + allocator = get_mem_allocator_instance() + allocator.release_tags(("kv_cache",)) + free_bytes_after_release, total = torch.cuda.mem_get_info() + freed_bytes = free_bytes_after_release - free_bytes_before_release + used_bytes = total - free_bytes_after_release + assert freed_bytes >= 0, "Memory usage increased after releasing KV cache." + logger.info( + "Released KV cache and freed %s GiB memory, %s GiB memory is still in use.", + format_gib(freed_bytes), + format_gib(used_bytes), + ) + def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if ( current_platform.is_cuda_alike()