diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index c0f8a5922239..f28ecc69dedb 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -75,6 +75,41 @@ 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_discard_tags(): + """Test that discard(tags) selectively frees GPU memory for specific + tags while keeping other tags mapped and usable.""" + allocator = get_mem_allocator_instance() + + with allocator.use_memory_pool("weights"): + weights = torch.ones(1024, 1024, device=DEVICE_TYPE) + + with allocator.use_memory_pool("kv_cache"): + kv = torch.ones(512, 512, device=DEVICE_TYPE) + + free_bytes = current_platform.mem_get_info()[0] + + # Discard kv_cache only — weights should remain valid + allocator.discard("kv_cache") + + free_bytes_after_discard = current_platform.mem_get_info()[0] + assert free_bytes_after_discard > free_bytes + + # Weights are still usable + assert torch.allclose(weights, torch.ones_like(weights)) + + # Wake up and verify kv_cache is remapped (zeroed content) + allocator.wake_up(["kv_cache"]) + # After wake_up the VA is remapped; content is not preserved + # but the allocation is valid + assert kv.shape == (512, 512) + + # Full sleep/wake cycle still works after discard + allocator.sleep(offload_tags="weights") + allocator.wake_up() + assert torch.allclose(weights, torch.ones_like(weights)) + + @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/vllm/device_allocator/__init__.py b/vllm/device_allocator/__init__.py index 66e8b146d29b..02d51603163e 100644 --- a/vllm/device_allocator/__init__.py +++ b/vllm/device_allocator/__init__.py @@ -27,6 +27,8 @@ def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ... def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ... + def discard(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 7c4fedd34a3b..a66f1fb6f8de 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -249,6 +249,8 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: backup_bytes = 0 for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + continue handle = data.handle total_bytes += handle[1] if data.tag in offload_tags: @@ -280,6 +282,30 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: gc.collect() torch.cuda.empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Selectively release GPU memory for the given tags without + CPU backup. Other tags remain mapped and usable.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + for data in self.pointer_to_data.values(): + if data.tag in tags and not data.is_asleep: + try: + unmap_and_release(data.handle) + finally: + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "CuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_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. diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index e0f359b200d2..560300078317 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -176,10 +176,13 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: backup_bytes = 0 for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + continue size_in_bytes = data.handle[1] total_bytes += size_in_bytes if data.tag not in offload_tags: unmap_and_release(data.handle) + data.is_asleep = True 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_asleep = True 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 discard(self, tags: tuple[str, ...] | str) -> None: + """Selectively release XPU memory for the given tags without + CPU backup. Other tags remain mapped and usable.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + for data in self.pointer_to_data.values(): + if data.tag in tags and not data.is_asleep: + try: + unmap_and_release(data.handle) + finally: + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "XpuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_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 create_and_allocate(data.handle) + data.is_asleep = False cpu_backup_tensor = data.cpu_backup_tensor if cpu_backup_tensor is None: