[ROCm] Add ROCm/HIP wheel build, CI, and release parity with CUDA (#3171)

[andyluo7]

Summary

Closes #3171 (ROCm CI parity with CUDA CI).

Mooncake already supports AMD GPUs in the transfer engine (USE_HIP, hip_transport/, dmabuf MR registration from #2225), but there was no CI to build/test the ROCm wheel and no published mooncake-transfer-engine-rocm package. This adds that, mirroring the existing CUDA/MUSA workflows.

The transfer engine is torch-independent and has no device kernels, so the ROCm wheel is GPU-arch-agnostic (one wheel covers gfx942 and gfx950), and the ROCm runtime (libamdhip64/libhsa-runtime64) is already in build_wheel.sh's auditwheel exclude list and bound at load time. As a result the same wheel works in both the upstream vllm/vllm-openai-rocm and lmsysorg/sglang ROCm images.

Changes

• scripts/build_wheel.sh — add HIP_BUILD variant → package mooncake-transfer-engine-rocm (ROCm runtime libs were already excluded by auditwheel).
• .github/workflows/ci_rocm.yml — PR build check in the ROCm dev image (rocm/dev-ubuntu-22.04:7.2.3-complete), py3.10 + py3.12, mirroring ci_cu13.yml. No GPU needed to compile HIP.
• .github/workflows/release-rocm.yaml — tag-triggered build + publish, mirroring the self-contained container build in release-musa.yaml; py3.10–3.13; publishes to PyPI (needs a PYPI_* token / pypi environment, same as other variants).
• ci.yml — wire build-wheel-rocm into the CI graph and the CI Gate.

Validation (on real AMD hardware)

Built the wheel in rocm/dev-ubuntu-22.04:7.2.3-complete, installed it into the latest upstream images, and ran mooncake_master --version + mooncake-wheel/tests/test_transfer_on_hip.py (registers HIP buffers, loopback sync read+write with data verification):

GPU	Image	Python	Result
MI300X gfx942 (Hotaisle)	vllm/vllm-openai-rocm:v0.26.0	cp312	✅ master + transfer test (2/2)
MI300X gfx942 (Hotaisle)	lmsysorg/sglang:v0.5.16-rocm720-mi30x	cp310	✅ master + transfer test (2/2)
MI355X gfx950 (Crusoe)	vllm/vllm-openai-rocm:v0.26.0	cp312	✅ master + transfer test (2/2)
MI355X gfx950 (Crusoe)	lmsysorg/sglang:v0.5.16-rocm720-mi35x	cp310	✅ master + transfer test (2/2)

ldd confirms engine.so binds libamdhip64.so.7 / libhsa-runtime64.so.1 to the host image's /opt/rocm at runtime (nothing vendored, no torch dependency).

Environment note (not a wheel/CI issue)

On the Crusoe MI355X node, when the container is given the ionic RoCE NICs, test_transfer_on_hip.py exercises RDMA GPU-VRAM registration. The dmabuf path is taken correctly (hsa_amd_portable_export_dmabuf succeeds), but ibv_reg_dmabuf_mr on the ionic NIC returns EFAULT — a node-level ionic kernel-driver limitation registering GPU dmabuf MRs (the same NIC family works on other clusters). This is independent of this PR (same wheel, same test); the transfer engine functions correctly via the local transport, and RDMA works wherever the NIC/driver supports GPU dmabuf MR registration.

🤖 Generated with Claude Code

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[Aionw]

Found one release-blocking issue in the historical-tag backfill path.

[Aionw]

[P1] Preserve the ROCm package rename when backfilling old tags

For workflow_dispatch, this checks out the historical tag before running scripts/build_wheel.sh. Tags predating this PR—including the documented example v0.3.12—do not recognize HIP_BUILD and have no ROCm package-name branch, so HIP_BUILD=1 is ignored and the script produces mooncake-transfer-engine instead of mooncake-transfer-engine-rocm. That mislabeled ROCm build then proceeds to the existing GitHub Release and PyPI upload steps. Please apply the ROCm metadata rename independently of the tagged script, or reject tags that do not contain HIP_BUILD support.

[andyluo7]

Good catch — fixed in e1d9bcd. Two guards, so a mislabeled ROCm build can never reach the GitHub Release / PyPI upload:

1. Reject unsupported tags — after checkout, the job fails if the checked-out scripts/build_wheel.sh has no HIP_BUILD branch, with a clear message to use a tag at/after ROCm support landed.
2. Assert the package name (defense in depth) — after build_wheel.sh, the job fails unless the produced wheel is mooncake_transfer_engine_rocm-*.whl, so an ignored HIP_BUILD (from any cause) can't proceed to publish.

Note: the same backfill-with-variant pattern exists in release-musa.yaml and release-npu.yaml (a pre-ROCm tag would similarly mislabel there) — happy to file a follow-up to add the same guard to those, keeping this PR scoped to ROCm.

[Aionw]

One follow-up issue remains in the manual historical-tag backfill path.

[Aionw]

[P2] Keep the advertised v0.3.12 backfill usable

This guard safely prevents a mislabeled wheel from being published, but it also makes the documented workflow_dispatch example impossible: HIP_BUILD is introduced by this PR, so v0.3.12—and every tag that exists when this PR merges—will fail here. As a result, manual dispatch cannot backfill the initial ROCm package; it can only rerun future tags that already contain this change. If historical backfill is intended, please apply the ROCm package metadata rename in the workflow for legacy tags and keep the post-build filename assertion. Otherwise, update the input description to remove the v0.3.12 example and state explicitly that only tags containing HIP_BUILD are supported.

[andyluo7]

Went with option (b) in 2dfa5a1 — the guard stays (it correctly prevents a mislabeled publish), and I've updated the workflow_dispatch input so it no longer advertises the pre-ROCm v0.3.12 example. It now reads "Release tag to (re)build the ROCm wheel for (must contain HIP_BUILD support)", with a comment noting older tags are rejected by design.

Backfilling the initial ROCm package onto a tag whose source predates ROCm support isn't meaningful (it'd build ROCm from a tree that never had it), so manual dispatch is intentionally limited to tags at/after this change. The first ROCm release will come from the normal v* tag push once this merges.

[Aionw]

This workflow publishes mooncake-transfer-engine-rocm using OIDC and does not pass an API token, so the existing PYPI_API_TOKEN/other package credentials do not authorize this new distribution. The ROCm project is not publicly present on PyPI yet, and a pending publisher is private, so its configuration cannot be verified from this PR.

Before merging/enabling the workflow, please confirm that PyPI has a pending Trusted Publisher configured with the exact identity kvcache-ai/Mooncake, workflow release-rocm.yaml, environment pypi, for project mooncake-transfer-engine-rocm. Otherwise the job will upload wheels to the GitHub Release first and then fail during PyPI publishing, leaving a partially published release. It would also be helpful to document this external prerequisite or add a protected/manual release gate.

[andyluo7]

Re: PyPI Trusted Publisher — good catch, addressed in 2dfa5a1:

• Reordered publish steps so PyPI publishing runs before attaching wheels to the GitHub Release. Now a missing/misconfigured Trusted Publisher fails the job cleanly without leaving a dangling GitHub Release.
• Documented the prerequisite inline in the workflow: PyPI must have a Trusted Publisher configured for the new project mooncake-transfer-engine-rocm with identity kvcache-ai/Mooncake, workflow release-rocm.yaml, environment pypi (OIDC, no API token).

That configuration is a one-time maintainer/PyPI-owner action I can't perform or verify from the PR side — @maintainers, please set up the pending Trusted Publisher for mooncake-transfer-engine-rocm before enabling/merging. The environment: pypi on the publish job can also serve as a manual release gate if you add required reviewers to that environment. (The build/CI jobs don't touch PyPI, so this only gates the actual release step.)

[Aionw]

[P2] Keep the release recoverable after PyPI succeeds

Moving the PyPI upload first avoids leaving GitHub Release assets behind when authentication fails, but both uploads are still in the same job. If PyPI succeeds and softprops/action-gh-release then fails transiently, rerunning the failed job starts from the PyPI step, which rejects the already-uploaded filenames, so the workflow can no longer reach the GitHub Release upload.

Please split PyPI publishing and GitHub Release attachment into separate jobs, with the latter depending on the former. Then “re-run failed jobs” can retry only the GitHub Release job without republishing immutable PyPI files.

[andyluo7]

Done in e164c96 — split into two jobs: publish-pypi and publish-github-release (the latter needs: publish-pypi). Now a transient action-gh-release failure can be re-run on its own via "re-run failed jobs" without restarting the immutable PyPI upload.

[Aionw]

[P2] Smoke-test the exact release wheels before publishing

The ROCm CI smoke test covers only Python 3.10 and 3.12 at the PR head, while this workflow builds and immediately publishes Python 3.10–3.13 wheels from the release tag. The release job currently checks only the filename, so packaging or ELF-layout failures in the exact release artifacts—especially cp311/cp313—can reach PyPI undetected.

Please install each matrix wheel into a fresh venv and run at least mooncake_master --version before uploading the artifact, matching the smoke test already used in ci_rocm.yml.

[andyluo7]

Done in e164c96 — added a "Smoke test the exact release wheel" step in the build matrix that installs each per-version artifact (cp310–cp313) into a fresh venv and runs mooncake_master --version before the wheel is uploaded, matching ci_rocm.yml. This runs on the release tag's actual artifacts, so cp311/cp313 packaging/ELF issues are caught before PyPI. (mooncake_master --version needs no GPU.)