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

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

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

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.

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: