feat[io]: add fabric (UALink scale-up) transfer backend#467
Merged
Conversation
maning00
force-pushed
the
io/fabric-backend
branch
from
92e13c6 to
8d60fd7
Compare
maning00
changed the title
There was a problem hiding this comment.
Pull request overview
Adds a new MORI-IO transfer backend (FABRIC) intended for cross-OS-image, intra-vPOD GPU transfers over UALink fabric handles, plus supporting kernel JIT, Python bindings, allocator utilities, and benchmark/doc updates.
Changes:
- Introduces a C++
FabricBackendwith fabric-handle export/import, vPOD topology detection, and a contiguous copy kernel path. - Extends Python bindings/APIs to create/use the FABRIC backend, including JIT compilation/loading of
fabric_copyand a torch pluggable allocator backed by fabric-exportable VMM allocations. - Updates the Python benchmark + docs to exercise the new backend and adds a torch distributed init fallback for certain ROCm builds.
Reviewed changes
Copilot reviewed 20 out of 20 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/python/utils.py | Adds a retry path for dist.init_process_group when device_id is rejected. |
| tests/python/io/benchmark.py | Adds --backend fabric, fabric buffer allocation/validation helpers, and distributed fabric setup/validation. |
| src/pybind/pybind_io.cpp | Exposes BackendType::FABRIC, FabricBackendConfig, and fabric_alloc/free in pybind. |
| src/io/kernels/fabric_copy.hip | New vectorized grid-stride copy kernel used by FABRIC transfers. |
| src/io/fabric/vpod_topology.hpp | Declares vPOD identity model and sysfs/device capability queries. |
| src/io/fabric/vpod_topology.cpp | Implements sysfs vPOD discovery and fabric capability query. |
| src/io/fabric/torch_allocator.cpp | Adds C-ABI malloc/free shims for PyTorch CUDAPluggableAllocator. |
| src/io/fabric/backend_impl.hpp | Declares FabricBackend, session, import/session caches, and kernel module hooks. |
| src/io/fabric/backend_impl.cpp | Implements fabric handle export/import, VA mapping, copy dispatch, and session caching. |
| src/io/engine.cpp | Wires FABRIC backend creation, backend selection priority, and kernel module loading entrypoint. |
| src/io/CMakeLists.txt | Adds fabric sources to mori_io build. |
| python/mori/io/fabric_copy_jit.py | Adds genco JIT compilation helper for fabric_copy.hip. |
| python/mori/io/fabric_allocator.py | Adds torch allocator + MemPool helpers to allocate fabric-exportable tensors. |
| python/mori/io/engine.py | Adds default config selection + copy-kernel load for FABRIC backend. |
| python/mori/io/init.py | Re-exports fabric allocator helpers and FABRIC config/symbols. |
| include/mori/io/enum.hpp | Adds BackendType::FABRIC. |
| include/mori/io/engine.hpp | Adds LoadFabricCopyModule API. |
| include/mori/io/common.hpp | Extends MemoryDesc with fabric handle + vPOD metadata + sub-allocation offset/size. |
| include/mori/io/backend.hpp | Adds FabricBackendConfig. |
| docs/MORI-IO-BENCHMARK.md | Documents --backend fabric and stream/event config flags. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+109
to
+117
| except (ValueError, RuntimeError): | ||
| # Some torch/ROCm builds report 0 "accelerators" and reject the | ||
| # device_id even though HIP devices exist. device_id is only a hint | ||
| # for accelerator backends; the gloo OOB path does not need it. | ||
| dist.init_process_group( | ||
| backend=self.backend, | ||
| rank=self.rank, | ||
| world_size=self.world_size, | ||
| ) |
Comment on lines
+118
to
+125
| def _load_fabric_copy_kernel(self): | ||
| try: | ||
| from mori.io.fabric_copy_jit import ensure_fabric_copy_kernel | ||
|
|
||
| hsaco_path = ensure_fabric_copy_kernel() | ||
| self._engine.LoadFabricCopyModule(hsaco_path) | ||
| except Exception: | ||
| pass |
Comment on lines
+277
to
+283
| hipFunction_t fn = backend != nullptr ? backend->GetFabricCopyFunc(localDevice) : nullptr; | ||
| if (fn == nullptr) { | ||
| // Kernel unavailable: fall back to hipMemcpyAsync. Correct for small | ||
| // payloads; large imported-fabric-pointer copies should always have the | ||
| // kernel loaded (see fabric_copy.hip note). | ||
| return hipMemcpyAsync(dst, src, size, hipMemcpyDefault, stream); | ||
| } |
Comment on lines
+64
to
+84
| bool DeviceSupportsFabric(int hipDevice) { | ||
| #if HIP_VERSION >= 71400000 | ||
| int vmm = 0; | ||
| int fabric = 0; | ||
| hipError_t err = | ||
| hipDeviceGetAttribute(&vmm, hipDeviceAttributeVirtualMemoryManagementSupported, hipDevice); | ||
| if (err != hipSuccess) { | ||
| (void)hipGetLastError(); | ||
| return false; | ||
| } | ||
| err = hipDeviceGetAttribute(&fabric, hipDeviceAttributeHandleTypeFabricSupported, hipDevice); | ||
| if (err != hipSuccess) { | ||
| (void)hipGetLastError(); | ||
| return false; | ||
| } | ||
| return vmm != 0 && fabric != 0; | ||
| #else | ||
| (void)hipDevice; | ||
| return false; | ||
| #endif | ||
| } |
The route cache was keyed only by engine/device pair. Since RdmaBackend CanHandle is memory-agnostic, a cached RDMA route (from an earlier non-fabric buffer) would be reused for a later fabric-exportable buffer on the same engine/device pair, masking the higher-priority FABRIC backend. Add local/remote memory ids to the key so each memory pair is routed independently.
maning00
force-pushed
the
io/fabric-backend
branch
from
f16fd38 to
7688c23
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.
Motivation
mori-io currently transfers over RDMA and XGMI. On UALink super-node (scale-up)
systems, GPUs that live in different OS-images but the same vPOD can access
each other's VRAM directly through HIP VMM fabric handles, which is faster
than routing the same traffic over RDMA.
This PR adds a new
fabricbackend so mori-io can move data between suchcross-OS-image, intra-vPOD GPU pairs using fabric handles plus a copy kernel,
and transparently falls back to XGMI/RDMA when fabric is not applicable.
Technical Details
BackendType::FABRICandFabricBackendConfig(
numStreams/numEvents, both default 64). Exposed through the pybind andPython
create_backendpaths.FabricBackend::RegisterMemoryretains the VMM allocationhandle and exports a 64-byte fabric token via
hipMemExportToShareableHandle(..., hipMemHandleTypeFabric). This onlysucceeds for allocations created with the fabric handle type; otherwise the
descriptor's fabric handle stays empty,
CanHandlereturns false, andselection falls back to XGMI/RDMA.
MemoryDescgains a 64-bytefabricHandle,vpodIdand
vpodPpodId, plus sub-allocationfabricOffset/fabricAllocSizeso atensor placed inside a larger pooled VMM allocation is remapped and offset
correctly on the importing side (all included in the msgpack layout).
(
ReadVpodKey/VpodKeySame/DeviceSupportsFabric). Two engines transferover fabric only when both GPUs support fabric and share the same vPOD. Env
overrides
MORI_IO_FABRIC_DISABLEandMORI_IO_FABRIC_VPOD_IDare provided.VRAM locally, then a grid-stride
uint4copy kernel (fabric_copy.hip,JIT/
--gencocompiled) drives the transfer (read vs write just swapssrc/dst). It reuses the XGMI backend's stream/event pools and completion
callbacks.
XGMI(same node) →FABRIC(cross-node,same vPOD) →
RDMA.CUDAPluggableAllocatorbacked bymori_io_fabric_malloc/freeplusMemPoolhelpers(
make_fabric_mem_pool/use_fabric_torch_allocator/fabric_mem_pool)let frameworks allocate fabric-exportable tensor memory so it can be
registered and transferred by this backend.
hipDeviceAttributeHandleTypeFabricSupportedis not present in every supported ROCm release. The device-capability query is
gated by
HIP_VERSION, so mori still builds on ROCm stacks that lack the enum(e.g. traditional MI300/MI350 deployments); on those, fabric simply reports
unsupported and selection falls back to XGMI/RDMA.
tests/python/io/benchmark.pygains--backend fabric(2-process cross-engine setup with fabric-allocated buffers), documented in
docs/MORI-IO-BENCHMARK.md.Test Plan
fabric read and write, including a sub-allocation-offset case.
benchmark.py --backend fabricwrite/read sweeps, both single-node(GPU0↔GPU1) and cross-node.
MemPool, register it, and transfer a sub-region with data validation.and without it (older ROCm, fabric compiled out / reported unsupported).
Test Result
faults.
~482 GB/s (write) / ~462 GB/s (read); batched contiguous transfers
(batch 128) reach ~490 GB/s. Sub-256 MB sizes are latency-bound as expected.
Submission Checklist