[Core] Add CuMemAllocator.discard() for tag-selective GPU memory release (#52514)

Signed-off-by: AlanFokCo <alanfok2868@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Co-authored-by: AlanFokCo <alanfok2868@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

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():

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: ...

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)

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: