[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>
2026-07-26 12:09:13 -07:00
parent b5b61c622c
commit 9e50e1037e
1 changed files with 3 additions and 0 deletions
@@ -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