From 9e50e1037e268b4488bbd472bd5dfee4cc08bb75 Mon Sep 17 00:00:00 2001
From: aoshen02 <aoshen@inferact.ai>
Date: Mon, 27 Jul 2026 03:09:13 +0800
Subject: [PATCH] [Bugfix][CuMem] Make KV-cache wake cleanup tag-safe (#49857)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
---
 vllm/device_allocator/cumem.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py
index 7c4fedd34a3..2cb9805bae3 100644
--- a/vllm/device_allocator/cumem.py
+++ b/vllm/device_allocator/cumem.py
@@ -291,6 +291,9 @@ class CuMemAllocator:
                 back to GPU memory. If None, all memory allocation will be loaded
                 back to GPU memory.
         """
+        gc.collect()
+        torch.accelerator.empty_cache()
+
         for ptr, data in self.pointer_to_data.items():
             if tags is None or data.tag in tags:
                 handle = data.handle