[Core] Add CuMemAllocator.discard() for tag-selective GPU memory release - #52514
Merged
Merged
Conversation
Add a discard() method to CuMemAllocator that releases GPU physical memory for the specified tags without backing up to CPU first. Unlike sleep(), which operates on all allocations (offloading some and discarding the rest), discard() only touches allocations matching the given tags and leaves everything else mapped and usable. This is needed for multi-model and RL weight-sync scenarios where stale data (e.g. a previous model's KV cache) should be freed to make room for new allocations while the current weights remain on GPU. Reuses the existing is_asleep flag to prevent sleep() from double- unmapping already-discarded handles. wake_up() already clears the flag when remapping. Signed-off-by: AlanFokCo <alanfok2868@gmail.com>
sleep() without arguments defaults to offload_tags=("default",),
but test_discard_tags uses custom tags "weights" and "kv_cache".
Without explicit offload_tags, weights memory was unmapped without
CPU backup, causing wake_up to remap with zeroed physical pages.
Signed-off-by: AlanFokCo <alanfok2868@gmail.com>
Co-authored-by: AlanFokCo <alanfok2868@gmail.com> Signed-off-by: Dakai An <dakaian108@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Collaborator
|
/ci run |
|
✅ @andakai, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #84086 for commit |
aoshen02
reviewed
Aug 16, 2026
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Dakai An <dakaian108@gmail.com>
Collaborator
|
/ci run |
|
✅ Triggered Buildkite CI #84087 for commit |
ZJY0516
approved these changes
Aug 16, 2026
This was referenced Aug 17, 2026
zufangzhu
pushed a commit
to zufangzhu/vllm
that referenced
this pull request
Aug 24, 2026
…ase (vllm-project#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> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #46438.
PR #46438 introduces tag-selective
discard(), allowing stale allocations such as KV cache to be released while keeping model weights mapped and usable.Based on 46438, This PR refines the interaction between
sleep()anddiscard()for both CUDA and XPU allocators:wake_up()remap only allocations that are actually asleep.The conflict check is needed because an asleep allocation has already been unmapped:
Repeated requests using the same policy remain silent and idempotent.