[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>
Dakai An committed
Aug 16, 2026 at 23:15 UTC
9409f59e0963a033a87a6847d6cafef54a64a79c
4 files changed
+132
tests/basic_correctness/test_mem.py
+35
@@ -75,6 +75,41 @@ def test_basic_cumem():
75
assert torch.allclose(output, torch.ones_like(output) * 3)
76
77
78
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
79
+def test_discard_tags():
80
+ """Test that discard(tags) selectively frees GPU memory for specific
81
+ tags while keeping other tags mapped and usable."""
82
+ allocator = get_mem_allocator_instance()
83
+
84
+ with allocator.use_memory_pool("weights"):
85
+ weights = torch.ones(1024, 1024, device=DEVICE_TYPE)
86
+
87
+ with allocator.use_memory_pool("kv_cache"):
88
+ kv = torch.ones(512, 512, device=DEVICE_TYPE)
89
+
90
+ free_bytes = torch.accelerator.get_memory_info()[0]
91
+
92
+ # Discard kv_cache only — weights should remain valid
93
+ allocator.discard("kv_cache")
94
+
95
+ free_bytes_after_discard = torch.accelerator.get_memory_info()[0]
96
+ assert free_bytes_after_discard > free_bytes
97
+
98
+ # Weights are still usable
99
+ assert torch.allclose(weights, torch.ones_like(weights))
100
+
101
+ # Wake up and verify kv_cache is remapped (zeroed content)
102
+ allocator.wake_up()
103
+ # After wake_up the VA is remapped; content is not preserved
104
+ # but the allocation is valid
105
+ assert kv.shape == (512, 512)
106
+
107
+ # Full sleep/wake cycle still works after discard
108
+ allocator.sleep(offload_tags="weights")
109
+ allocator.wake_up()
110
+ assert torch.allclose(weights, torch.ones_like(weights))
111
+
112
+
113
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
114
@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
115
def test_cumem_with_cudagraph():
vllm/device_allocator/__init__.py
+2
@@ -27,6 +27,8 @@ class MemAllocator(Protocol):
27
28
def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ...
29
30
+ def discard(self, tags: tuple[str, ...] | str) -> None: ...
31
+
32
def wake_up(self, tags: list[str] | None = None) -> None: ...
33
34
def get_current_usage(self) -> int: ...
vllm/device_allocator/cumem.py
+46
@@ -247,8 +247,15 @@ class CuMemAllocator:
247
248
total_bytes = 0
249
backup_bytes = 0
250
+ has_policy_conflict = False
251
252
for ptr, data in self.pointer_to_data.items():
253
+ if data.is_asleep:
254
+ requests_offload = data.tag in offload_tags
255
+ was_offloaded = data.cpu_backup_tensor is not None
256
+ if requests_offload != was_offloaded:
257
+ has_policy_conflict = True
258
+ continue
259
handle = data.handle
260
total_bytes += handle[1]
261
if data.tag in offload_tags:
@@ -277,9 +284,46 @@ class CuMemAllocator:
284
(total_bytes - backup_bytes) / 1024**3,
285
)
286
287
+ if has_policy_conflict:
288
+ logger.warning(
289
+ "CuMemAllocator: sleep cannot change the policy of "
290
+ "already-asleep allocations; the existing policy was kept."
291
+ )
292
+
293
gc.collect()
294
torch.cuda.empty_cache()
295
296
+ def discard(self, tags: tuple[str, ...] | str) -> None:
297
+ """Discard mapped allocations with the given tags without CPU backup."""
298
+ if isinstance(tags, str):
299
+ tags = (tags,)
300
+
301
+ discarded_bytes = 0
302
+ has_policy_conflict = False
303
+ for data in self.pointer_to_data.values():
304
+ if data.tag not in tags:
305
+ continue
306
+ if data.is_asleep:
307
+ if data.cpu_backup_tensor is not None:
308
+ has_policy_conflict = True
309
+ continue
310
+ torch.accelerator.synchronize(data.handle[0])
311
+ unmap_and_release(data.handle)
312
+ data.is_asleep = True
313
+ discarded_bytes += data.handle[1]
314
+
315
+ logger.info(
316
+ "CuMemAllocator: discarded %.2f GiB for tags %s.",
317
+ discarded_bytes / 1024**3,
318
+ tags,
319
+ )
320
+
321
+ if has_policy_conflict:
322
+ logger.warning(
323
+ "CuMemAllocator: discard cannot change the policy of "
324
+ "already-asleep allocations; the existing policy was kept."
325
+ )
326
+
327
def wake_up(self, tags: list[str] | None = None) -> None:
328
"""
329
Wake up the allocator from sleep mode.
@@ -295,6 +339,8 @@ class CuMemAllocator:
339
torch.accelerator.empty_cache()
340
341
for ptr, data in self.pointer_to_data.items():
342
+ if not data.is_asleep:
343
+ continue
344
if tags is None or data.tag in tags:
345
handle = data.handle
346
create_and_map(handle)
vllm/device_allocator/xpumem.py
+49
@@ -174,12 +174,20 @@ class XpuMemAllocator:
174
175
total_bytes = 0
176
backup_bytes = 0
177
+ has_policy_conflict = False
178
179
for ptr, data in self.pointer_to_data.items():
180
+ if data.is_asleep:
181
+ requests_offload = data.tag in offload_tags
182
+ was_offloaded = data.cpu_backup_tensor is not None
183
+ if requests_offload != was_offloaded:
184
+ has_policy_conflict = True
185
+ continue
186
size_in_bytes = data.handle[1]
187
total_bytes += size_in_bytes
188
if data.tag not in offload_tags:
189
unmap_and_release(data.handle)
190
+ data.is_asleep = True
191
continue
192
193
backup_bytes += size_in_bytes
@@ -201,6 +209,7 @@ class XpuMemAllocator:
209
data.cpu_backup_tensor = cpu_backup_tensor
210
211
unmap_and_release(data.handle)
212
+ data.is_asleep = True
213
214
logger.info(
215
"XpuMemAllocator: sleep freed %.2f GiB memory in total, of which "
@@ -211,16 +220,56 @@ class XpuMemAllocator:
220
(total_bytes - backup_bytes) / 1024**3,
221
)
222
223
+ if has_policy_conflict:
224
+ logger.warning(
225
+ "XpuMemAllocator: sleep cannot change the policy of "
226
+ "already-asleep allocations; the existing policy was kept."
227
+ )
228
+
229
gc.collect()
230
xpu_empty_cache = getattr(torch.xpu, "empty_cache", None)
231
if callable(xpu_empty_cache):
232
xpu_empty_cache()
233
234
+ def discard(self, tags: tuple[str, ...] | str) -> None:
235
+ """Discard mapped allocations with the given tags without CPU backup."""
236
+ if isinstance(tags, str):
237
+ tags = (tags,)
238
+
239
+ discarded_bytes = 0
240
+ has_policy_conflict = False
241
+ for data in self.pointer_to_data.values():
242
+ if data.tag not in tags:
243
+ continue
244
+ if data.is_asleep:
245
+ if data.cpu_backup_tensor is not None:
246
+ has_policy_conflict = True
247
+ continue
248
+ torch.accelerator.synchronize(data.handle[0])
249
+ unmap_and_release(data.handle)
250
+ data.is_asleep = True
251
+ discarded_bytes += data.handle[1]
252
+
253
+ logger.info(
254
+ "XpuMemAllocator: discarded %.2f GiB for tags %s.",
255
+ discarded_bytes / 1024**3,
256
+ tags,
257
+ )
258
+
259
+ if has_policy_conflict:
260
+ logger.warning(
261
+ "XpuMemAllocator: discard cannot change the policy of "
262
+ "already-asleep allocations; the existing policy was kept."
263
+ )
264
+
265
def wake_up(self, tags: list[str] | None = None) -> None:
266
for ptr, data in self.pointer_to_data.items():
267
+ if not data.is_asleep:
268
+ continue
269
if tags is not None and data.tag not in tags:
270
continue
271
create_and_allocate(data.handle)
272
+ data.is_asleep = False
273
274
cpu_backup_tensor = data.cpu_backup_tensor
275
if cpu_backup_tensor is None: