[vLLM] UCX plugin: defer packRkey to prevent GPU memory pin on ROCm#28
Closed
AndreasKaratzas wants to merge 1 commit into
Closed
[vLLM] UCX plugin: defer packRkey to prevent GPU memory pin on ROCm#28AndreasKaratzas wants to merge 1 commit into
AndreasKaratzas wants to merge 1 commit into
Conversation
AndreasKaratzas
requested review from
akolliasAMD,
edgargabriel and
gaoikawa
as code owners
April 5, 2026 03:00
AndreasKaratzas
force-pushed
the
akaratza_defer_packrkey
branch
from
April 5, 2026 03:03
d1b6947 to
ef59f74
Compare
AndreasKaratzas
force-pushed
the
akaratza_defer_packrkey
branch
from
April 5, 2026 03:29
ef59f74 to
5634627
Compare
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
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.
Defer
packRkeyfromregisterMemtogetPublicData/loadLocalMD.Problem
registerMemeagerly callspackRkeywhich triggersucp_rkey_pack->hsa_amd_ipc_memory_create. This permanently pins GPU memory at the HSA driver level with no release API for the creating process. For workloads registering large GPU buffers (e.g. vLLM KV caches at hundreds GB) that do not perform same-host IPC transfers, this causes GPU memory to be unrecoverable after engine shutdown.Fix
Move the
packRkeycall fromregisterMem(eager, during init) togetPublicDataandloadLocalMD(lazy, on first access). The IPC handle is only created when a remote peer actually requests the metadata during handshake.Zero performance regression: the
packRkeycost is deferred, not eliminated. Same-host IPC transfers work because the handle is populated before the first transfer.Companion fix: ROCm/ucx#33
cc @kenroche