libfabric: add AMD/ROCm HMEM support for CXI-attached accelerators - #1873
ryanhankins wants to merge 0 commit into
Conversation
|
👋 Hi ryanhankins! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
📝 WalkthroughWalkthroughThis PR adds ROCm/HIP as a CUDA alternative through Meson configuration, LibFabric HIP context and VRAM registration support, AMD accelerator topology discovery, and ROCm-aware rail manager runtime selection. ChangesROCm/HIP Build Support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant nixlLibfabricEngine
participant HIPRuntime
participant nixlLibfabricRocmCtx
participant nixlLibfabricRailManager
nixlLibfabricEngine->>HIPRuntime: query pointer attributes
HIPRuntime-->>nixlLibfabricEngine: return device and PCI metadata
nixlLibfabricEngine->>nixlLibfabricRocmCtx: validate and update device context
nixlLibfabricRocmCtx->>HIPRuntime: hipSetDevice
nixlLibfabricEngine->>nixlLibfabricRailManager: register VRAM memory
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/plugins/libfabric/libfabric_backend.cpp`:
- Around line 372-389: The ROCm device tracking in
libfabric_backend::registerAddress currently latches devId_ once and rejects
later valid buffers from a different AMD GPU with NIXL_ERR_MISMATCH. Update the
logic after rocmQueryAddr(address) and the expected_dev checks so each validated
VRAM pointer can refresh devId_ to the current expected_dev instead of only
setting it when devId_ is kInvalidRocmDeviceId. Keep the existing validation for
non-device memory and mismatched queried device IDs, but remove the cross-device
rejection caused by the stored devId_.
- Around line 897-900: The ROCm conditional branch in libfabric_backend should
be reformatted to use the standard 4-space indentation. Adjust the body of the
FI_HMEM_ROCR check in the code around the runtime_ branch so the NIXL_DEBUG call
and mems.push_back(VRAM_SEG) are indented consistently beneath the if statement,
matching the surrounding style in libfabric_backend.cpp and the CodeStyle.md
requirement.
In `@src/plugins/libfabric/libfabric_backend.h`:
- Around line 88-116: The ROCm identifiers in nixlLibfabricRocmCtx need to be
renamed to match backend naming rules and used consistently across the header
and implementation: change RocmDeviceId to a *_t-style enum class name, rename
kInvalidRocmDeviceId to snake_case, and update the private member and any
context symbol usages such as devId_ and rocmCtx_ to the required
trailing-underscore snake_case form. Make the same naming updates wherever these
symbols are referenced in the .cpp so the API and implementation stay aligned.
In `@src/utils/libfabric/libfabric_rail_manager.cpp`:
- Around line 782-784: The VRAM registration path in libfabric_rail_manager.cpp
is always using the constructor-selected runtime_ as the fi_hmem_iface, which
can be wrong for buffers from different accelerator runtimes. Update the VRAM
branch in the registration logic around the iface selection so it derives the
iface from the specific buffer/device context instead of blindly reusing
runtime_, using the existing buffer/runtime metadata or device_pci_bus_id
association before assigning mr_attr.iface.
- Around line 213-215: The AMD accelerator branch in the libfabric rail manager
is not indented consistently inside the topology->getNumAmdAccel() conditional.
Update the body in the relevant if block within the libfabric_rail_manager logic
so the runtime_ assignment and NIXL_INFO logging are indented with 4 spaces
relative to the if statement, matching the surrounding formatting style.
In `@src/utils/libfabric/libfabric_topology.h`:
- Line 47: Rename the new private AMD counter in libfabric_topology to follow
the class’s private member convention: update num_amd_accel to camelCase with a
trailing underscore (for example, numAmdAccel_) and adjust any references in the
related libfabric topology code so the member name is consistent everywhere.
- Around line 217-220: Add a Doxygen block comment for the new public API
getNumAmdAccel() in libfabric_topology.h, using the same style as the other
public getters/types in the LibfabricTopology class. Document the purpose of the
method and what num_amd_accel represents, placing the comment directly above
getNumAmdAccel() so the public interface follows the project’s Doxygen
convention.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 271df797-853b-4e56-90d2-22ab45e57626
📒 Files selected for processing (9)
meson.buildmeson_options.txtsrc/plugins/libfabric/libfabric_backend.cppsrc/plugins/libfabric/libfabric_backend.hsrc/plugins/libfabric/meson.buildsrc/utils/libfabric/libfabric_rail_manager.cppsrc/utils/libfabric/libfabric_topology.cppsrc/utils/libfabric/libfabric_topology.hsrc/utils/libfabric/meson.build
|
@fengjica @amitrad-aws Would you please let me know if this looks ok? |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
meson.build (1)
481-488: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDon't add
cudato trace project args on ROCm builds
cuda_dep.found()is true whencuda_depaliasesrocm_dep, butCUDAis never enabled in that branch. That makesadd_project_arguments(..., language: ['cpp', 'cuda'])fail for-Duse_rocm=true -Dwith_trace=true. Gate thecudaentry on actual CUDA language enablement.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@meson.build` around lines 481 - 488, Update the nixl_trace_enabled block in the Meson configuration so nixl_trace_arg_langs includes 'cuda' only when the CUDA language was actually enabled, not merely when cuda_dep.found() is true; preserve 'cpp' unconditionally and use the existing CUDA enablement state to avoid adding CUDA project arguments on ROCm builds.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/plugins/libfabric/libfabric_backend.cpp`:
- Around line 979-1006: registerMem’s ROCm path redundantly queries and
validates the pointer device through rocmUpdateCtxPtr after already obtaining
validated RocmAddrInfo. Add a lightweight
nixlLibfabricRocmCtx::setDevice(RocmDeviceId) method that only updates devId_,
and call it from registerMem after the existing rocmQueryAddr/device-match
checks; retain rocmUpdateCtxPtr’s query-and-validation behavior for other
callers.
---
Outside diff comments:
In `@meson.build`:
- Around line 481-488: Update the nixl_trace_enabled block in the Meson
configuration so nixl_trace_arg_langs includes 'cuda' only when the CUDA
language was actually enabled, not merely when cuda_dep.found() is true;
preserve 'cpp' unconditionally and use the existing CUDA enablement state to
avoid adding CUDA project arguments on ROCm builds.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 08f6bb69-1f0b-4fdf-963d-2d1ab7c4f14f
📒 Files selected for processing (9)
meson.buildmeson_options.txtsrc/plugins/libfabric/libfabric_backend.cppsrc/plugins/libfabric/libfabric_backend.hsrc/plugins/libfabric/meson.buildsrc/utils/libfabric/libfabric_rail_manager.cppsrc/utils/libfabric/libfabric_topology.cppsrc/utils/libfabric/libfabric_topology.hsrc/utils/libfabric/meson.build
| #endif | ||
| #ifdef HAVE_ROCM | ||
| if (runtime_ == FI_HMEM_ROCR) { | ||
| const auto info = rocmQueryAddr(reinterpret_cast<void *>(mem.addr)); | ||
| if (!info || !info->is_dev) { | ||
| NIXL_ERROR << "Failed to query ROCm device from memory " | ||
| << reinterpret_cast<void *>(mem.addr); | ||
| return NIXL_ERR_BACKEND; | ||
| } | ||
| pci_bus_id = info->pci_bus_id; | ||
| if (info->dev_id != static_cast<RocmDeviceId>(mem.devId)) { | ||
| NIXL_ERROR << "ROCm pointer/device mismatch: pointer on GPU " << info->dev_id | ||
| << ", metadata requests GPU " << mem.devId; | ||
| return NIXL_ERR_INVALID_PARAM; | ||
| } | ||
| nixl_status_t status = | ||
| rocmCtx_->rocmUpdateCtxPtr(reinterpret_cast<void *>(mem.addr), info->dev_id); | ||
| if (status != NIXL_SUCCESS) { | ||
| NIXL_ERROR << "Failed to update ROCm context for device " << info->dev_id; | ||
| return status; | ||
| } | ||
| status = rocmCtx_->rocmSetCtx(); | ||
| if (status != NIXL_SUCCESS) { | ||
| NIXL_ERROR << "Failed to set ROCm device context for device " << info->dev_id; | ||
| return status; | ||
| } | ||
| NIXL_DEBUG << "Queried PCI bus ID: " << pci_bus_id << " for AMD GPU " << mem.devId; | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Duplicate device query/validation between registerMem and rocmUpdateCtxPtr.
registerMem already calls rocmQueryAddr(mem.addr) and checks info->dev_id != mem.devId (lines 982-993), then immediately calls rocmCtx_->rocmUpdateCtxPtr(mem.addr, info->dev_id), which internally re-runs rocmQueryAddr (a second hipPointerGetAttributes driver call) and re-checks info->dev_id != expected_dev — a comparison that is now tautological since expected_dev was already derived from the same query. This duplicates both a driver round-trip and validation logic that could silently diverge if one copy is updated later without the other.
♻️ Proposed direction
Add a lightweight setter (e.g. nixlLibfabricRocmCtx::setDevice(RocmDeviceId)) that only updates devId_ without re-querying, and have registerMem call that instead of rocmUpdateCtxPtr once it has already validated the pointer/device match itself. Reserve rocmUpdateCtxPtr's full query+validate contract for callers that don't already have RocmAddrInfo.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/plugins/libfabric/libfabric_backend.cpp` around lines 979 - 1006,
registerMem’s ROCm path redundantly queries and validates the pointer device
through rocmUpdateCtxPtr after already obtaining validated RocmAddrInfo. Add a
lightweight nixlLibfabricRocmCtx::setDevice(RocmDeviceId) method that only
updates devId_, and call it from registerMem after the existing
rocmQueryAddr/device-match checks; retain rocmUpdateCtxPtr’s
query-and-validation behavior for other callers.
|
|
||
| #ifdef HAVE_ROCM | ||
| /** ROCm/HIP device identifier. */ | ||
| enum class RocmDeviceId : int {}; |
There was a problem hiding this comment.
A valid id needs to be non-negative? I guess here it cannot be unsigned int, because -1 is used below.
|
@ryanhankins can we update PR with the fixes? |
|
@ryanhankins can you pls resolve conflicts? |
|
Sorry, was out of town. I can certainly update, but what do we want to do with this? #1750 contains most of the work here. Shall I just close this? By the way, thank you so much for looking! |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/utils/libfabric/libfabric_topology.h (1)
145-150: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winPlace the return type on a separate line.
Keep the declaration and definition in the required layout.
src/utils/libfabric/libfabric_topology.h#L145-L150: placeboolbelow[[nodiscard]]for all accelerator predicates.src/utils/libfabric/libfabric_topology.cpp#L788-L789: placeboolbelow[[nodiscard]]forisAmdAccel.As per path instructions, “Place return types on separate lines.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/utils/libfabric/libfabric_topology.h` around lines 145 - 150, Reformat the accelerator predicate declarations in src/utils/libfabric/libfabric_topology.h lines 145-150 so bool appears on its own line below [[nodiscard]] for isNvidiaAccel, isNeuronAccel, and isAmdAccel. Apply the same required layout to the isAmdAccel definition in src/utils/libfabric/libfabric_topology.cpp lines 788-789.Source: Path instructions
meson.build (1)
228-231: 🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy liftKeep
cuda_depseparate fromrocm_dep.When
use_rocmis true,cuda_dep.found()becomes true. The following guard then callsadd_languages('CUDA'), so ROCm-only configurations still requirenvcc, contrary to theuse_rocmoption. Keepcuda_depfalse when CUDA is unavailable. Use a separate accelerator dependency for libfabric and update CUDA-only consumers to use an explicit CUDA availability flag.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@meson.build` around lines 228 - 231, Remove the cuda_dep = rocm_dep alias so cuda_dep remains independently false when CUDA is unavailable. Introduce a separate accelerator dependency for libfabric, and update CUDA-only consumers and the add_languages('CUDA') guard to use an explicit CUDA availability flag rather than rocm_dep or cuda_dep.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@meson_options.txt`:
- Around line 29-32: Remove the duplicate rocm_path option from the Meson
options, retaining a single definition that combines both descriptions: the
general ROCm installation path behavior and its use for UCX unit tests.
---
Outside diff comments:
In `@meson.build`:
- Around line 228-231: Remove the cuda_dep = rocm_dep alias so cuda_dep remains
independently false when CUDA is unavailable. Introduce a separate accelerator
dependency for libfabric, and update CUDA-only consumers and the
add_languages('CUDA') guard to use an explicit CUDA availability flag rather
than rocm_dep or cuda_dep.
In `@src/utils/libfabric/libfabric_topology.h`:
- Around line 145-150: Reformat the accelerator predicate declarations in
src/utils/libfabric/libfabric_topology.h lines 145-150 so bool appears on its
own line below [[nodiscard]] for isNvidiaAccel, isNeuronAccel, and isAmdAccel.
Apply the same required layout to the isAmdAccel definition in
src/utils/libfabric/libfabric_topology.cpp lines 788-789.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 223eacf7-de55-4843-b445-de5d0601c084
📒 Files selected for processing (8)
meson.buildmeson_options.txtsrc/plugins/libfabric/libfabric_backend.cppsrc/plugins/libfabric/libfabric_backend.hsrc/plugins/libfabric/meson.buildsrc/utils/libfabric/libfabric_rail_manager.cppsrc/utils/libfabric/libfabric_topology.cppsrc/utils/libfabric/libfabric_topology.h
| option('use_rocm', type: 'boolean', value: false, description: 'Build against AMD ROCm/HIP instead of CUDA. When true, CUDA toolchain is not required.') | ||
| option('rocm_path', type: 'string', value: '', description: 'Path to ROCm installation directory (used only when use_rocm=true). Empty falls back to /opt/rocm.') | ||
| option('nixl_cuda_arch_list', type: 'string', value: 'auto', description: 'Comma-separated CUDA SM targets (e.g. 90,100), or auto to select defaults (sm_80+ normally, sm_90+ when build_nixl_ep is enabled). Use a single target like 100 for faster single-GPU builds.') | ||
| option('rocm_path', type: 'string', value: '', description: 'Path to ROCm installation for UCX unit tests (empty = do not use ROCm)') |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Remove the duplicate rocm_path option.
Line 30 and Line 32 both define rocm_path. Meson configuration stops before any build when an option name is defined twice. Keep one definition and combine the required documentation.
Proposed fix
option('use_rocm', type: 'boolean', value: false, description: 'Build against AMD ROCm/HIP instead of CUDA. When true, CUDA toolchain is not required.')
option('rocm_path', type: 'string', value: '', description: 'Path to ROCm installation directory (used only when use_rocm=true). Empty falls back to /opt/rocm.')
-option('rocm_path', type: 'string', value: '', description: 'Path to ROCm installation for UCX unit tests (empty = do not use ROCm)')🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@meson_options.txt` around lines 29 - 32, Remove the duplicate rocm_path
option from the Meson options, retaining a single definition that combines both
descriptions: the general ROCm installation path behavior and its use for UCX
unit tests.
Replaces #1771 due to branch history overwrite.
Make the libfabric plugin usable for VRAM transfers on AMD GPU systems connected via Slingshot.
Build system
libfabric backend (libfabric_backend.h / .cpp)
AMD GPU topology (libfabric_topology.h / .cpp)
Rail manager (libfabric_rail_manager.cpp)
Summary by CodeRabbit
New Features
Bug Fixes