[TransferEngine][ROCm] Add ROCm HIP support to the Mooncake Python package by knitcapcat-amd · Pull Request #1742 · kvcache-ai/Mooncake

knitcapcat-amd · 2026-03-25T06:59:14Z

Expose HIP as the Python-facing AMD GPU transport, wire HIP transport into the build and runtime selection paths, and fix sticky peer-access errors so repeated connector initialization works reliably.

Description

This PR enables ROCm HIP usage from the Mooncake Python package and makes the
end-to-end HIP path usable in containerized environments.
The changes include:

exposing hip as the Python-facing AMD GPU transport in the Python binding
compiling and registering HipTransport independently from NVIDIA-specific
MNNVL logic
selecting hip in the relevant runtime transport-installation paths
fixing a sticky hipErrorPeerAccessAlreadyEnabled runtime state that could
break repeated connector initialization in the same process
updating wheel packaging so repaired wheels keep working in ROCm containers
without vendoring incompatible ROCm runtime libraries

Module

Type of Change

How Has This Been Tested?

Rebuilt Mooncake in a ROCm container with HIP enabled.
Rebuilt and reinstalled the repaired Python wheel from dist-rocm-hip.
Ran python3 mooncake-wheel/tests/test_transfer_on_hip.py successfully.

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

gemini-code-assist · 2026-03-25T06:59:35Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Mooncake Python package by introducing full support for ROCm HIP. The changes enable AMD GPU transport for Python users, ensuring the end-to-end HIP path is robust and functional, particularly within containerized environments. This broadens the compatibility and utility of the Mooncake Transfer Engine for systems leveraging AMD's ROCm platform.

Highlights

ROCm HIP Support: Added comprehensive support for ROCm HIP as the Python-facing AMD GPU transport within the Mooncake Python package.
Build and Runtime Integration: Integrated HIP transport into the build system and runtime selection paths, allowing for independent compilation and registration of HipTransport.
Peer-Access Error Resolution: Fixed a persistent hipErrorPeerAccessAlreadyEnabled runtime state that previously caused issues with repeated connector initialization.
Packaging and Validation: Updated wheel packaging to ensure repaired wheels function correctly in ROCm containers without vendoring incompatible ROCm runtime libraries, and added new tests for HIP functionality.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces comprehensive support for the HIP (Heterogeneous-compute Interface for Portability) transport protocol within the Mooncake Transfer Engine. Key changes include adding HIP-specific headers and transport classes, integrating HIP memory allocation and transport initialization into the Python API, and enabling dynamic selection of cross-node transport based on HIP availability. New Python tests have been added to validate HIP functionality, P2P error handling for HIP has been refined, and the build script has been updated to exclude HIP-related system libraries from the Python wheel. Feedback from the review suggests addressing a redundant xgmi protocol validation in initMemoryAllocator to centralize validation logic, and refactoring duplicated code for installing manual transports in transfer_engine_py.cpp to improve maintainability.

gemini-code-assist · 2026-03-25T07:01:48Z

+    } else if (strcmp(protocol, "xgmi") == 0) {
+        allocateMemory = malloc;
+        freeMemory = free;
+        LOG(ERROR) << "Protocol 'xgmi' is not exposed in the Python API. "
+                   << "Use 'hip' instead.";


The check for the deprecated xgmi protocol is redundant here. The initializeExt function already performs this validation (lines 189-193) and correctly returns an error code to stop initialization.

This check in initMemoryAllocator only logs an error but allows execution to proceed, which could be confusing. To centralize the validation logic, I recommend removing this else if block.

gemini-code-assist · 2026-03-25T07:01:48Z

+    if (manual_transport) {
+        auto *transport = engine_->installTransport(protocol, nullptr);
+        if (!transport) {
+            LOG(ERROR) << "Failed to install transport: " << protocol;
+            return -1;
+        }
+    }


There's a duplication of the logic for installing manual transports. This block is identical to lines 243-248 within the #ifdef USE_EFA block.

To improve maintainability, consider refactoring this to avoid repeating the code. You could, for example, extract this logic into a helper function or restructure the conditional compilation blocks.

codecov-commenter · 2026-03-25T07:21:18Z

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

Codecov Report

❌ Patch coverage is 28.98551% with 49 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...integration/transfer_engine/transfer_engine_py.cpp	0.00%	48 Missing ⚠️
mooncake-transfer-engine/src/multi_transport.cpp	95.23%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

stmatengss · 2026-03-25T07:48:45Z

+if (USE_HIP)
+  add_subdirectory(hip_transport)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:hip_transport>)
+endif()
+
+if (USE_MNNVL AND NOT USE_HIP)
+  add_subdirectory(nvlink_transport)
+  target_sources(transport PUBLIC $<TARGET_OBJECTS:nvlink_transport>)


Nice catch.

stmatengss · 2026-03-25T07:54:39Z

+#if defined(USE_HIP)
+            constexpr const char* kCrossNodeTransport = "hip";
+            constexpr const char* kCrossNodeTransportName = "HIP";
+#else


#if defined (USE_MNNVL) || defined(USE_INTRA_NVLINK)

Thanks for pointing out this issue. The original code incorrectly treated HIP the same as NVLink (MNNVL), but NVLink is cross-node capable and can replace RDMA, while HIP IPC is intra-node only and should coexist with it. And I fixed this issue in following commits.

stmatengss · 2026-03-25T08:02:18Z

+        return -1;
+    }
+
+    bool manual_transport = strcmp(protocol, "hip") == 0 ||


Why is manual transport needed? Please provide more design details.

The manual_transport flag was introduced to disable auto_discover and install HIP explicitly, because we assumed HIP would conflict with the auto-discovered RDMA transport. This was based on a wrong assumption -- we treated HIP the same as NVLink (MNNVL), which replaces RDMA as the cross-node transport. But HIP IPC is actually intra-node only and can coexist with RDMA. Once we moved HIP to an independent install path in TransferEngineImpl::init(), auto_discover=true works correctly and manual_transport is no longer needed. Removed entirely -- the Python binding now matches main behavior. Thanks for pushing on this.

ShangmingCai · 2026-04-03T05:15:06Z

+        allocateMemory = malloc;
+        freeMemory = free;


This logic looks a little bit weird. Why do we need this when NVLink checking fails? Should we put them in the else if (strcmp(protocol, "hip") == 0) block?

Thank you for pointing this, you're right. Setting allocateMemory = malloc in the error branch silently locked in a wrong allocator state. Removed -- error branches now only log without touching allocator.

alogfans

Please respond all issues.

alogfans · 2026-04-10T01:58:09Z

LGTM. Double check if you have time @stmatengss

amd-arozanov

LGTM. Please fix CI failures, if they are related to the changes from this PR

Expose HIP as the Python-facing AMD GPU transport, wire HIP transport into the build and runtime selection paths, and fix sticky peer-access errors so repeated connector initialization works reliably. Update wheel packaging and container validation to cover ROCm HIP usage in vllm-omni. Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

HIP IPC is intra-node only and should coexist with RDMA, not replace it. Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> Made-with: Cursor Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

stmatengss · 2026-04-13T16:30:24Z

LGTM. @alogfans @amd-arozanov Thank you for reviewing this PR!

…ckage (kvcache-ai#1742) * Add ROCm HIP support to the Mooncake Python package Expose HIP as the Python-facing AMD GPU transport, wire HIP transport into the build and runtime selection paths, and fix sticky peer-access errors so repeated connector initialization works reliably. Update wheel packaging and container validation to cover ROCm HIP usage in vllm-omni. Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> * [Bugfix] Decouple HIP transport from NVLink branch HIP IPC is intra-node only and should coexist with RDMA, not replace it. Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> * Abandon manual transport Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> * style: apply clang-format to changed files Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> Made-with: Cursor Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> --------- Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Co-authored-by: Zejian Wang <zejianwang@sjtu.edu.cn>

knitcapcat-amd requested review from ShangmingCai, alogfans, amd-arozanov, chestnut-Q, doujiang24 and stmatengss as code owners March 25, 2026 06:59

github-actions Bot added run-ci Transfer Engine Installation labels Mar 25, 2026

gemini-code-assist Bot reviewed Mar 25, 2026

View reviewed changes

stmatengss reviewed Mar 25, 2026

View reviewed changes

stmatengss assigned amd-arozanov and knitcapcat-amd Mar 25, 2026

ShangmingCai reviewed Apr 3, 2026

View reviewed changes

alogfans approved these changes Apr 3, 2026

View reviewed changes

alogfans requested changes Apr 3, 2026

View reviewed changes

amd-arozanov approved these changes Apr 10, 2026

View reviewed changes

knitcapcat added 4 commits April 10, 2026 09:26

[Bugfix] Decouple HIP transport from NVLink branch

2a4b045

HIP IPC is intra-node only and should coexist with RDMA, not replace it. Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

Abandon manual transport

caa549d

Signed-off-by: Zejian Wang <zejianwang.sjtu.edu.cn> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

style: apply clang-format to changed files

d18c85e

Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> Made-with: Cursor Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn>

knitcapcat-amd force-pushed the feature/rocm-hip-python-support branch from 2291334 to d18c85e Compare April 10, 2026 09:27

alogfans approved these changes Apr 13, 2026

View reviewed changes

alogfans merged commit c0d07af into kvcache-ai:main Apr 13, 2026
17 checks passed

This was referenced Jul 24, 2026

[Feature]: Publish prebuilt ROCm/HIP wheels to PyPI for mooncake ROCm/legacy-rocm-build#6514

Open

[Installation] Publish prebuilt ROCm/HIP wheels to PyPI #3108

Open

Uh oh!

Conversation

knitcapcat-amd commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist Bot commented Mar 25, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

stmatengss Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

knitcapcat-amd Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

knitcapcat-amd Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

knitcapcat-amd Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

alogfans left a comment

Choose a reason for hiding this comment

Uh oh!

alogfans commented Apr 10, 2026

Uh oh!

amd-arozanov left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stmatengss commented Apr 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

knitcapcat-amd commented Mar 25, 2026 •

edited

Loading

codecov-commenter commented Mar 25, 2026 •

edited

Loading