Skip to content

[vLLM] UCX plugin: defer packRkey to prevent GPU memory pin on ROCm#28

Closed
AndreasKaratzas wants to merge 1 commit into
developfrom
akaratza_defer_packrkey
Closed

[vLLM] UCX plugin: defer packRkey to prevent GPU memory pin on ROCm#28
AndreasKaratzas wants to merge 1 commit into
developfrom
akaratza_defer_packrkey

Conversation

@AndreasKaratzas

@AndreasKaratzas AndreasKaratzas commented Apr 5, 2026

Copy link
Copy Markdown

Defer packRkey from registerMem to getPublicData/loadLocalMD.

Problem

registerMem eagerly calls packRkey which triggers ucp_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 packRkey call from registerMem (eager, during init) to getPublicData and loadLocalMD (lazy, on first access). The IPC handle is only created when a remote peer actually requests the metadata during handshake.

Zero performance regression: the packRkey cost 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

@AndreasKaratzas
AndreasKaratzas force-pushed the akaratza_defer_packrkey branch from ef59f74 to 5634627 Compare April 5, 2026 03:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants