diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 0618562b3e..696d2f66bc 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 = torch.accelerator.get_memory_info()[0] + + # Discard kv_cache only — weights should remain valid + allocator.discard("kv_cache") + + free_bytes_after_discard = torch.accelerator.get_memory_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() + # 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 66e8b146d2..02d5160316 100644 --- a/vllm/device_allocator/__init__.py +++ b/vllm/device_allocator/__init__.py @@ -27,6 +27,8 @@ class MemAllocator(Protocol): 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 2cb9805bae..a16f1b0148 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -247,8 +247,15 @@ class CuMemAllocator: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + continue handle = data.handle total_bytes += handle[1] if data.tag in offload_tags: @@ -277,9 +284,46 @@ class CuMemAllocator: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "CuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() torch.cuda.empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "CuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "CuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: """ Wake up the allocator from sleep mode. @@ -295,6 +339,8 @@ class CuMemAllocator: torch.accelerator.empty_cache() for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue if tags is None or data.tag in tags: handle = data.handle create_and_map(handle) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index e0f359b200..146ea965db 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -174,12 +174,20 @@ class XpuMemAllocator: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + 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 +209,7 @@ class XpuMemAllocator: 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 " @@ -211,16 +220,56 @@ class XpuMemAllocator: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() xpu_empty_cache = getattr(torch.xpu, "empty_cache", None) if callable(xpu_empty_cache): xpu_empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "XpuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue 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: