From 32be9492f7580db1f8accc39ca02ea67172cd814 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Mon, 13 Jul 2026 07:00:41 +0000 Subject: [PATCH 1/6] add fabric backend --- docs/MORI-IO-BENCHMARK.md | 25 + include/mori/io/backend.hpp | 13 + include/mori/io/common.hpp | 14 +- include/mori/io/engine.hpp | 1 + include/mori/io/enum.hpp | 1 + python/mori/io/__init__.py | 3 + python/mori/io/engine.py | 13 + python/mori/io/fabric_copy_jit.py | 40 ++ src/io/CMakeLists.txt | 4 +- src/io/engine.cpp | 21 + src/io/fabric/backend_impl.cpp | 833 ++++++++++++++++++++++++++++++ src/io/fabric/backend_impl.hpp | 214 ++++++++ src/io/fabric/vpod_topology.cpp | 128 +++++ src/io/fabric/vpod_topology.hpp | 60 +++ src/io/kernels/fabric_copy.hip | 66 +++ src/pybind/pybind_io.cpp | 28 +- tests/python/io/benchmark.py | 242 ++++++++- tests/python/utils.py | 22 +- 18 files changed, 1707 insertions(+), 21 deletions(-) create mode 100644 python/mori/io/fabric_copy_jit.py create mode 100644 src/io/fabric/backend_impl.cpp create mode 100644 src/io/fabric/backend_impl.hpp create mode 100644 src/io/fabric/vpod_topology.cpp create mode 100644 src/io/fabric/vpod_topology.hpp create mode 100644 src/io/kernels/fabric_copy.hip diff --git a/docs/MORI-IO-BENCHMARK.md b/docs/MORI-IO-BENCHMARK.md index 78c6b2ad5..eddae3ef2 100644 --- a/docs/MORI-IO-BENCHMARK.md +++ b/docs/MORI-IO-BENCHMARK.md @@ -25,10 +25,35 @@ torchrun --nnodes=2 --node_rank=0 --nproc_per_node=1 \ tests/python/io/benchmark.py --host="10.194.129.65" ``` +### Fabric (cross-node scale-up UALink super-node) + +The `fabric` backend transfers between two nodes that share the same scale-up +fabric domain (**same vPOD** — verify with `amd-smi fabric` that `PPOD_ID` and +`VPOD_ID` match and `ACCEL_STATE` is `READY` on both). Buffers are allocated as +fabric-exportable VMM memory internally, so no `--host`/QP options are needed; +OOB metadata is still exchanged over gloo (torchrun). + +```bash +# node A (initiator) +torchrun --nnodes=2 --node_rank=0 --nproc_per_node=1 \ + --master_addr="" --master_port=1234 \ + tests/python/io/benchmark.py --backend fabric \ + --all --enable-sess --enable-batch-transfer --transfer-batch-size 1 \ + --sweep-start-size 1048576 --sweep-max-size 268435456 --op-type read + +# node B (target): same command with --node_rank=1 +``` + +> Requires ROCm ≥ 7.15 (HIP fabric VMM APIs). Only GPU memory is supported +> (`--mem-type gpu`). If the two nodes are not in the same vPOD, session creation +> fails fast with a clear error. + ## Benchmark Arguments | Argument | Description | |----------|-------------| +| `--backend` | `rdma` (cross-node), `xgmi` (intra-node), or `fabric` (cross-node UALink super-node, same vPOD) | +| `--num-streams` / `--num-events` | HIP stream/event pool size for `xgmi`/`fabric` backends (default `64`) | | `--buffer-size` | Message size per transfer (bytes) | | `--all` | Sweep message size from 8B to 1MB | | `--sweep-start-size` | Starting message size when using `--all` sweep | diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index 25bf12381..2503af73d 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -101,6 +101,19 @@ inline std::ostream& operator<<(std::ostream& os, const XgmiBackendConfig& c) { return os << "numStreams[" << c.numStreams << "] numEvents[" << c.numEvents << "]"; } +struct FabricBackendConfig : public BackendConfig { + FabricBackendConfig() : BackendConfig(BackendType::FABRIC) {} + FabricBackendConfig(int numStreams_, int numEvents_) + : BackendConfig(BackendType::FABRIC), numStreams(numStreams_), numEvents(numEvents_) {} + + int numStreams{64}; + int numEvents{64}; +}; + +inline std::ostream& operator<<(std::ostream& os, const FabricBackendConfig& c) { + return os << "numStreams[" << c.numStreams << "] numEvents[" << c.numEvents << "]"; +} + /* ---------------------------------------------------------------------------------------------- */ /* BackendSession */ /* ---------------------------------------------------------------------------------------------- */ diff --git a/include/mori/io/common.hpp b/include/mori/io/common.hpp index 670f6c274..e1f146a9a 100644 --- a/include/mori/io/common.hpp +++ b/include/mori/io/common.hpp @@ -93,6 +93,10 @@ using MemoryUniqueId = uint32_t; constexpr size_t kIpcHandleSize = 64; +// Fabric (UALink super-node) shareable handle size. Equals the IPC handle size +// (hipMemFabricHandle_t is a 64-byte token, same width as hipIpcMemHandle_t). +constexpr size_t kFabricHandleSize = kIpcHandleSize; + struct MemoryDesc { EngineKey engineKey; MemoryUniqueId id{0}; @@ -103,6 +107,13 @@ struct MemoryDesc { MemoryLocationType loc; std::array ipcHandle{}; int numaNode{-1}; + // Fabric backend metadata. Populated by FabricBackend::RegisterMemory when the + // GPU allocation is exportable over a scale-up fabric (UALink vPOD). Kept + // separate from ipcHandle so XGMI and Fabric backends can register the same + // memory without clobbering each other's handle. + std::array fabricHandle{}; // 64B fabric token (empty when N/A) + int vpodId{-1}; // UALink vPOD id of deviceId (-1 = N/A) + std::string vpodPpodId; // UALink ppod_id UUID ("" = N/A) constexpr bool operator==(const MemoryDesc& rhs) const noexcept { return (engineKey == rhs.engineKey) && (id == rhs.id) && (deviceId == rhs.deviceId) && @@ -110,7 +121,8 @@ struct MemoryDesc { (loc == rhs.loc) && (numaNode == rhs.numaNode); } - MSGPACK_DEFINE(engineKey, id, deviceId, deviceBusId, data, size, loc, ipcHandle, numaNode); + MSGPACK_DEFINE(engineKey, id, deviceId, deviceBusId, data, size, loc, ipcHandle, numaNode, + fabricHandle, vpodId, vpodPpodId); }; using TransferUniqueId = uint64_t; diff --git a/include/mori/io/engine.hpp b/include/mori/io/engine.hpp index 41d886544..6fedfe5ad 100644 --- a/include/mori/io/engine.hpp +++ b/include/mori/io/engine.hpp @@ -109,6 +109,7 @@ class IOEngine { std::optional CreateSession(const MemoryDesc& local, const MemoryDesc& remote); void LoadScatterGatherModule(const std::string& hsacoPath); + void LoadFabricCopyModule(const std::string& hsacoPath); private: struct RouteCacheKey { diff --git a/include/mori/io/enum.hpp b/include/mori/io/enum.hpp index 8d1a27b9c..94f68db21 100644 --- a/include/mori/io/enum.hpp +++ b/include/mori/io/enum.hpp @@ -29,6 +29,7 @@ enum class BackendType : uint32_t { XGMI = 1, RDMA = 2, TCP = 3, + FABRIC = 4, // Cross-node scale-up fabric (UALink / super-node vPOD) }; using BackendTypeVec = std::vector; diff --git a/python/mori/io/__init__.py b/python/mori/io/__init__.py index 8cf0a7733..c73067054 100644 --- a/python/mori/io/__init__.py +++ b/python/mori/io/__init__.py @@ -30,5 +30,8 @@ PollCqMode, RdmaBackendConfig, XgmiBackendConfig, + FabricBackendConfig, + fabric_alloc, + fabric_free, set_log_level, ) diff --git a/python/mori/io/engine.py b/python/mori/io/engine.py index ccf01a93d..4f50db910 100644 --- a/python/mori/io/engine.py +++ b/python/mori/io/engine.py @@ -95,11 +95,15 @@ def create_backend(self, type: mori_cpp.BackendType, config=None): config = mori_cpp.RdmaBackendConfig() elif type is mori_cpp.BackendType.XGMI: config = mori_cpp.XgmiBackendConfig() + elif type is mori_cpp.BackendType.FABRIC: + config = mori_cpp.FabricBackendConfig() else: raise NotImplementedError("backend not implemented yet") result = self._engine.CreateBackend(type, config) if type is mori_cpp.BackendType.XGMI: self._load_scatter_gather_kernel() + elif type is mori_cpp.BackendType.FABRIC: + self._load_fabric_copy_kernel() return result def _load_scatter_gather_kernel(self): @@ -111,6 +115,15 @@ def _load_scatter_gather_kernel(self): except Exception: pass + 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 + def remove_backend(self, type: mori_cpp.BackendType): return self._engine.RemoveBackend(type) diff --git a/python/mori/io/fabric_copy_jit.py b/python/mori/io/fabric_copy_jit.py new file mode 100644 index 000000000..68d9b41f6 --- /dev/null +++ b/python/mori/io/fabric_copy_jit.py @@ -0,0 +1,40 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""JIT compilation for the Fabric backend contiguous copy kernel. + +Compiles src/io/kernels/fabric_copy.hip via hipcc --genco on first use. +The resulting .hsaco is loaded by C++ FabricBackend via hipModuleLoad. +""" + +from __future__ import annotations + +_hsaco_path: str | None = None + + +def ensure_fabric_copy_kernel() -> str: + """JIT compile the fabric copy kernel and return the .hsaco path.""" + global _hsaco_path + if _hsaco_path is None: + from mori.jit.core import compile_genco + + _hsaco_path = compile_genco("fabric_copy", source_dir="src/io/kernels") + return _hsaco_path diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index 911b81075..16c3a350a 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -6,7 +6,9 @@ set(MORI_IO_SOURCES rdma/common.cpp rdma/ledger.cpp xgmi/backend_impl.cpp - xgmi/hip_resource_pool.cpp) + xgmi/hip_resource_pool.cpp + fabric/backend_impl.cpp + fabric/vpod_topology.cpp) set_source_files_properties(${MORI_IO_SOURCES} PROPERTIES LANGUAGE CXX) add_library(mori_io SHARED ${MORI_IO_SOURCES}) diff --git a/src/io/engine.cpp b/src/io/engine.cpp index eada02d1c..8da75c083 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -38,6 +38,7 @@ #include "mori/io/logging.hpp" #include "mori/utils/host_utils.hpp" #include "src/io/call_diagnostics_internal.hpp" +#include "src/io/fabric/backend_impl.hpp" #include "src/io/rdma/backend_impl.hpp" #include "src/io/xgmi/backend_impl.hpp" @@ -253,6 +254,11 @@ void IOEngine::CreateBackend(BackendType type, const BackendConfig& beConfig) { static_cast(beConfig)); backends.insert({type, std::move(backend)}); InvalidateRouteCache(); + } else if (type == BackendType::FABRIC) { + auto backend = std::make_unique( + desc.key, config, static_cast(beConfig)); + backends.insert({type, std::move(backend)}); + InvalidateRouteCache(); } else { assert(false && "not implemented"); } @@ -413,6 +419,14 @@ Backend* IOEngine::SelectBackend(const MemoryDesc& local, const MemoryDesc& remo return xgmiIt->second.get(); } + // Next prefer FABRIC: cross-node GPU pairs in the same scale-up domain (vPOD) + // reachable over the UALink super-node fabric, which is faster than RDMA. + auto fabricIt = backends.find(BackendType::FABRIC); + if (fabricIt != backends.end() && fabricIt->second->CanHandle(local, remote)) { + UpdateRouteCache(routeKey, BackendType::FABRIC); + return fabricIt->second.get(); + } + auto rdmaIt = backends.find(BackendType::RDMA); if (rdmaIt != backends.end() && rdmaIt->second->CanHandle(local, remote)) { UpdateRouteCache(routeKey, BackendType::RDMA); @@ -560,6 +574,13 @@ void IOEngine::LoadScatterGatherModule(const std::string& hsacoPath) { } } +void IOEngine::LoadFabricCopyModule(const std::string& hsacoPath) { + auto it = backends.find(BackendType::FABRIC); + if (it != backends.end()) { + static_cast(it->second.get())->LoadFabricCopyModule(hsacoPath); + } +} + bool IOEngine::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, TransferStatus* status) { // status->SetCode(StatusCode::SUCCESS); diff --git a/src/io/fabric/backend_impl.cpp b/src/io/fabric/backend_impl.cpp new file mode 100644 index 000000000..df87569bf --- /dev/null +++ b/src/io/fabric/backend_impl.cpp @@ -0,0 +1,833 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "src/io/fabric/backend_impl.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/logging.hpp" +#include "mori/utils/hip_compat.hpp" +#include "mori/utils/host_utils.hpp" + +namespace mori { +namespace io { +namespace { + +// Lowercase a PCI bus ID so HIP (may return uppercase) and sysfs (lowercase) +// spellings compare equal. +std::string NormalizeBusId(const std::string& busId) { + std::string result = busId; + for (auto& c : result) { + c = static_cast(std::tolower(static_cast(c))); + } + return result; +} + +bool IsFabricHandleEmpty(const std::array& handle) { + return std::all_of(handle.begin(), handle.end(), [](char c) { return c == 0; }); +} + +// Keep caller-visible HIP current device unchanged across MORI internals. +class ScopedHipDeviceGuard { + public: + ScopedHipDeviceGuard() { + if (hipGetDevice(&originalDevice_) != hipSuccess) { + valid_ = false; + } + } + ~ScopedHipDeviceGuard() { + if (valid_) (void)hipSetDevice(originalDevice_); + } + + private: + int originalDevice_{0}; + bool valid_{true}; +}; + +// Round a user size up to the device's fabric-allocation granularity, so an +// imported handle can be reserved/mapped as a whole allocation. +size_t AlignToFabricGranularity(int device, size_t size) { + hipMemAllocationProp prop = {}; + prop.type = hipMemAllocationTypePinned; + prop.requestedHandleType = hipMemHandleTypeFabricCompat; + prop.location.type = hipMemLocationTypeDevice; + prop.location.id = device; + + size_t gran = 0; + hipError_t err = + hipMemGetAllocationGranularity(&gran, &prop, hipMemAllocationGranularityRecommended); + if (err != hipSuccess || gran == 0) { + (void)hipGetLastError(); + gran = 2ull * 1024 * 1024; // conservative 2 MiB fallback + } + return ((size + gran - 1) / gran) * gran; +} + +// Deferred completion for an event-backed transfer. Mirrors the XGMI backend: +// the status polls/waits on the HIP event and finalizes exactly once. +struct TransferCompletion { + TransferStatus* status{nullptr}; + hipEvent_t event{nullptr}; + int eventDevice{-1}; + EventPool* eventPool{nullptr}; + std::function cleanup; + std::once_flag completionOnce; + std::atomic completed{false}; + std::mutex eventMu; + + void FinalizeBlocking() { + std::lock_guard lock(eventMu); + std::call_once(completionOnce, [this]() { + hipError_t err = hipEventSynchronize(event); + if (err == hipSuccess) { + status->SetCode(StatusCode::SUCCESS); + } else { + status->Update(StatusCode::ERR_GPU_OP, std::string("FABRIC: hipEventSynchronize failed: ") + + hipGetErrorString(err)); + } + if (cleanup) cleanup(); + eventPool->PutEvent(event, eventDevice); + completed.store(true, std::memory_order_release); + }); + } + + void FinalizeNonBlocking() { + std::lock_guard lock(eventMu); + if (completed.load(std::memory_order_acquire)) return; + + hipError_t err = hipEventQuery(event); + if (err == hipErrorNotReady) { + (void)hipGetLastError(); + return; + } + std::call_once(completionOnce, [this, err]() { + if (err == hipSuccess) { + status->SetCode(StatusCode::SUCCESS); + } else { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: hipEventQuery failed: ") + hipGetErrorString(err)); + } + if (cleanup) cleanup(); + eventPool->PutEvent(event, eventDevice); + completed.store(true, std::memory_order_release); + }); + } +}; + +void ArmTransferCompletion(TransferStatus* status, hipEvent_t event, int eventDevice, + EventPool* eventPool, std::function cleanup = {}) { + auto completion = std::make_shared(); + completion->status = status; + completion->event = event; + completion->eventDevice = eventDevice; + completion->eventPool = eventPool; + completion->cleanup = std::move(cleanup); + status->SetWaitCallback([completion]() { completion->FinalizeBlocking(); }); + status->SetProgressCallback([completion]() { completion->FinalizeNonBlocking(); }); +} + +// Registry backing FabricMalloc / FabricFree. +struct FabricAllocation { + hipMemGenericAllocationHandle_t handle{}; + size_t size{0}; +}; +std::mutex g_fabricAllocMu; +std::unordered_map g_fabricAllocs; + +} // namespace + +/* ---------------------------------------------------------------------------------------------- */ +/* Fabric allocation helpers */ +/* ---------------------------------------------------------------------------------------------- */ +void* FabricMalloc(size_t size, int device) { + if (size == 0) return nullptr; + + ScopedHipDeviceGuard guard; + if (hipSetDevice(device) != hipSuccess) { + MORI_IO_WARN("FABRIC: FabricMalloc failed to set device {}", device); + return nullptr; + } + + hipMemAllocationProp prop = {}; + prop.type = hipMemAllocationTypePinned; + prop.requestedHandleType = hipMemHandleTypeFabricCompat; + prop.location.type = hipMemLocationTypeDevice; + prop.location.id = device; + prop.allocFlags.gpuDirectRDMACapable = 1; + + size_t gran = 0; + if (hipMemGetAllocationGranularity(&gran, &prop, hipMemAllocationGranularityRecommended) != + hipSuccess || + gran == 0) { + MORI_IO_WARN("FABRIC: FabricMalloc granularity query failed on device {}", device); + return nullptr; + } + size_t aligned = ((size + gran - 1) / gran) * gran; + + hipMemGenericAllocationHandle_t handle{}; + hipError_t err = hipMemCreate(&handle, aligned, &prop, 0); + if (err != hipSuccess) { + MORI_IO_WARN("FABRIC: FabricMalloc hipMemCreate failed: {}", hipGetErrorString(err)); + return nullptr; + } + + void* ptr = nullptr; + err = hipMemAddressReserve(&ptr, aligned, 0, 0, 0); + if (err != hipSuccess) { + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: FabricMalloc hipMemAddressReserve failed: {}", hipGetErrorString(err)); + return nullptr; + } + err = hipMemMap(ptr, aligned, 0, handle, 0); + if (err != hipSuccess) { + (void)hipMemAddressFree(ptr, aligned); + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: FabricMalloc hipMemMap failed: {}", hipGetErrorString(err)); + return nullptr; + } + hipMemAccessDesc access = {}; + access.location.type = hipMemLocationTypeDevice; + access.location.id = device; + access.flags = hipMemAccessFlagsProtReadWrite; + err = hipMemSetAccess(ptr, aligned, &access, 1); + if (err != hipSuccess) { + (void)hipMemUnmap(ptr, aligned); + (void)hipMemAddressFree(ptr, aligned); + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: FabricMalloc hipMemSetAccess failed: {}", hipGetErrorString(err)); + return nullptr; + } + + { + std::lock_guard lock(g_fabricAllocMu); + g_fabricAllocs[ptr] = {handle, aligned}; + } + MORI_IO_TRACE("FABRIC: FabricMalloc device={} size={} aligned={} ptr={}", device, size, aligned, + ptr); + return ptr; +} + +void FabricFree(void* ptr) { + if (ptr == nullptr) return; + FabricAllocation alloc; + { + std::lock_guard lock(g_fabricAllocMu); + auto it = g_fabricAllocs.find(ptr); + if (it == g_fabricAllocs.end()) { + MORI_IO_WARN("FABRIC: FabricFree unknown pointer {}", ptr); + return; + } + alloc = it->second; + g_fabricAllocs.erase(it); + } + (void)hipMemUnmap(ptr, alloc.size); + (void)hipMemRelease(alloc.handle); + (void)hipMemAddressFree(ptr, alloc.size); + MORI_IO_TRACE("FABRIC: FabricFree ptr={} size={}", ptr, alloc.size); +} + +/* ---------------------------------------------------------------------------------------------- */ +/* FabricBackendSession */ +/* ---------------------------------------------------------------------------------------------- */ +FabricBackendSession::FabricBackendSession(const FabricBackendConfig& config, void* localAddr, + void* remoteAddr, int localDevice, + FabricBackend* backend, StreamPool* streamPool, + EventPool* eventPool) + : config(config), + localAddr(localAddr), + remoteAddr(remoteAddr), + localDevice(localDevice), + backend(backend), + streamPool(streamPool), + eventPool(eventPool) {} + +hipError_t FabricBackendSession::LaunchCopy(void* dst, const void* src, size_t size, + hipStream_t stream) { + if (size == 0) return hipSuccess; + + 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); + } + + char* dstPtr = static_cast(dst); + const char* srcPtr = static_cast(src); + size_t nbytes = size; + void* args[] = {&dstPtr, &srcPtr, &nbytes}; + + const int threads = 256; + size_t units = (size + 15) / 16; // 16B chunks (grid-stride handles remainder) + long long blocks = static_cast((units + threads - 1) / threads); + if (blocks < 1) blocks = 1; + if (blocks > 65535) blocks = 65535; + + return hipModuleLaunchKernel(fn, static_cast(blocks), 1, 1, threads, 1, 1, 0, stream, + args, nullptr); +} + +void FabricBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, + TransferStatus* status, TransferUniqueId id, bool isRead) { + ScopedHipDeviceGuard deviceGuard; + hipError_t err = hipSetDevice(localDevice); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: Failed to set device: ") + hipGetErrorString(err)); + return; + } + + hipStream_t stream = streamPool->GetNextStream(localDevice); + hipEvent_t event = eventPool->GetEvent(localDevice); + if (stream == nullptr || event == nullptr) { + status->Update(StatusCode::ERR_BAD_STATE, "FABRIC: Failed to get stream or event from pool"); + if (event != nullptr) eventPool->PutEvent(event, localDevice); + return; + } + + void* src = isRead ? static_cast(remoteAddr) + remoteOffset + : static_cast(localAddr) + localOffset; + void* dst = isRead ? static_cast(localAddr) + localOffset + : static_cast(remoteAddr) + remoteOffset; + + err = LaunchCopy(dst, src, size, stream); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: copy launch failed: ") + hipGetErrorString(err)); + eventPool->PutEvent(event, localDevice); + return; + } + + err = hipEventRecord(event, stream); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: hipEventRecord failed: ") + hipGetErrorString(err)); + eventPool->PutEvent(event, localDevice); + return; + } + + status->SetCode(StatusCode::IN_PROGRESS); + ArmTransferCompletion(status, event, localDevice, eventPool); + MORI_IO_TRACE("FABRIC: Transfer issued, id={}, size={}, isRead={}", id, size, isRead); +} + +void FabricBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id, bool isRead) { + ScopedHipDeviceGuard deviceGuard; + size_t batchSize = sizes.size(); + assert(batchSize == localOffsets.size()); + assert(batchSize == remoteOffsets.size()); + + if (batchSize == 0) { + status->SetCode(StatusCode::SUCCESS); + return; + } + + hipError_t err = hipSetDevice(localDevice); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: Failed to set device: ") + hipGetErrorString(err)); + return; + } + + hipStream_t stream = streamPool->GetNextStream(localDevice); + hipEvent_t event = eventPool->GetEvent(localDevice); + if (stream == nullptr || event == nullptr) { + status->Update(StatusCode::ERR_BAD_STATE, "FABRIC: Failed to get stream or event from pool"); + if (event != nullptr) eventPool->PutEvent(event, localDevice); + return; + } + + // Sort by remote offset then merge contiguous runs (both sides contiguous) to + // cut the number of kernel launches. + std::vector indices(batchSize); + std::iota(indices.begin(), indices.end(), 0); + if (!std::is_sorted(remoteOffsets.begin(), remoteOffsets.end())) { + std::sort(indices.begin(), indices.end(), + [&](size_t a, size_t b) { return remoteOffsets[a] < remoteOffsets[b]; }); + } + + struct MergedSeg { + size_t localOff; + size_t remoteOff; + size_t sz; + }; + std::vector segments; + segments.reserve(batchSize); + for (size_t i = 0; i < batchSize; ++i) { + size_t idx = indices[i]; + if (sizes[idx] == 0) continue; + if (!segments.empty()) { + MergedSeg& last = segments.back(); + bool localContig = (last.localOff + last.sz) == localOffsets[idx]; + bool remoteContig = (last.remoteOff + last.sz) == remoteOffsets[idx]; + if (localContig && remoteContig) { + last.sz += sizes[idx]; + continue; + } + } + segments.push_back({localOffsets[idx], remoteOffsets[idx], sizes[idx]}); + } + + if (segments.empty()) { + status->SetCode(StatusCode::SUCCESS); + eventPool->PutEvent(event, localDevice); + return; + } + + for (const auto& seg : segments) { + void* src = isRead ? static_cast(remoteAddr) + seg.remoteOff + : static_cast(localAddr) + seg.localOff; + void* dst = isRead ? static_cast(localAddr) + seg.localOff + : static_cast(remoteAddr) + seg.remoteOff; + err = LaunchCopy(dst, src, seg.sz, stream); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: batch copy launch failed: ") + hipGetErrorString(err)); + eventPool->PutEvent(event, localDevice); + return; + } + } + + err = hipEventRecord(event, stream); + if (err != hipSuccess) { + status->Update(StatusCode::ERR_GPU_OP, + std::string("FABRIC: hipEventRecord failed: ") + hipGetErrorString(err)); + eventPool->PutEvent(event, localDevice); + return; + } + + status->SetCode(StatusCode::IN_PROGRESS); + ArmTransferCompletion(status, event, localDevice, eventPool); + MORI_IO_TRACE("FABRIC: Batch transfer issued, id={}, segments={}, isRead={}", id, segments.size(), + isRead); +} + +bool FabricBackendSession::Alive() const { return true; } + +/* ---------------------------------------------------------------------------------------------- */ +/* FabricBackend */ +/* ---------------------------------------------------------------------------------------------- */ +FabricBackend::FabricBackend(EngineKey k, const IOEngineConfig& /*engConfig*/, + const FabricBackendConfig& beConfig) + : myEngKey(k), config(beConfig), myPid(static_cast(getpid())) { + char hostname[HOST_NAME_MAX]; + gethostname(hostname, HOST_NAME_MAX); + myHostname = std::string(hostname); + myNodeId = mori::ResolveNodeId(myHostname); + + streamPool = std::make_unique(config.numStreams); + eventPool = std::make_unique(config.numEvents); + + Initialize(); + + std::stringstream ss; + ss << config; + MORI_IO_INFO("FabricBackend created with config: {} node_id: {} hostname: {} numDevices: {}", + ss.str().c_str(), myNodeId.c_str(), myHostname.c_str(), numDevices); +} + +FabricBackend::~FabricBackend() { + { + std::unique_lock lock(importMutex); + for (auto& entry : importedRegions) { + ImportedRegion& region = entry.second; + if (region.mappedAddr != nullptr) { + (void)hipMemUnmap(region.mappedAddr, region.size); + (void)hipMemRelease(region.handle); + (void)hipMemAddressFree(region.mappedAddr, region.size); + } + } + importedRegions.clear(); + } + + std::lock_guard lock(moduleMu_); + for (auto& mod : fabricCopyModules_) { + if (mod != nullptr) (void)hipModuleUnload(mod); + } + fabricCopyModules_.clear(); + fabricCopyFuncs_.clear(); +} + +void FabricBackend::Initialize() { + hipError_t err = hipGetDeviceCount(&numDevices); + if (err != hipSuccess || numDevices <= 0) { + MORI_IO_WARN("FABRIC: Failed to get device count or no devices found"); + numDevices = 0; + return; + } + + deviceSupportsFabric.assign(numDevices, 0); + deviceVpodKeys.assign(numDevices, fabric::VpodKey{}); + + for (int i = 0; i < numDevices; ++i) { + char busId[32] = {0}; + if (hipDeviceGetPCIBusId(busId, sizeof(busId), i) == hipSuccess) { + localDeviceByBusId[NormalizeBusId(std::string(busId))] = i; + } + deviceSupportsFabric[i] = fabric::DeviceSupportsFabric(i) ? 1 : 0; + deviceVpodKeys[i] = fabric::ReadVpodKey(i); + if (deviceSupportsFabric[i]) { + MORI_IO_INFO("FABRIC: device {} supports fabric (vpod valid={} ppod_id={} vpod_id={})", i, + deviceVpodKeys[i].valid, deviceVpodKeys[i].ppodId, deviceVpodKeys[i].vpodId); + } + } +} + +std::optional FabricBackend::LookupVisibleDevice(const std::string& busId) const { + auto it = localDeviceByBusId.find(NormalizeBusId(busId)); + if (it == localDeviceByBusId.end()) return std::nullopt; + return it->second; +} + +fabric::VpodKey FabricBackend::LocalVpodKey(int deviceId) const { + if (deviceId < 0 || deviceId >= numDevices) return fabric::VpodKey{}; + return deviceVpodKeys[deviceId]; +} + +void FabricBackend::LoadFabricCopyModule(const std::string& hsacoPath) { + std::lock_guard lock(moduleMu_); + fabricCopyHsacoPath_ = hsacoPath; + fabricCopyModules_.assign(numDevices, nullptr); + fabricCopyFuncs_.assign(numDevices, nullptr); + MORI_IO_INFO("FABRIC: copy kernel registered from {}", hsacoPath); +} + +hipFunction_t FabricBackend::GetFabricCopyFunc(int deviceId) { + std::lock_guard lock(moduleMu_); + if (fabricCopyHsacoPath_.empty() || deviceId < 0 || deviceId >= numDevices) { + return nullptr; + } + if (fabricCopyFuncs_[deviceId] != nullptr) { + return fabricCopyFuncs_[deviceId]; + } + + ScopedHipDeviceGuard deviceGuard; + if (hipSetDevice(deviceId) != hipSuccess) return nullptr; + + hipError_t err = hipModuleLoad(&fabricCopyModules_[deviceId], fabricCopyHsacoPath_.c_str()); + if (err != hipSuccess) { + MORI_IO_WARN("FABRIC: Failed to load copy module on device {}: {}", deviceId, + hipGetErrorString(err)); + return nullptr; + } + err = hipModuleGetFunction(&fabricCopyFuncs_[deviceId], fabricCopyModules_[deviceId], + "fabricCopyKernel"); + if (err != hipSuccess) { + MORI_IO_WARN("FABRIC: Failed to get fabricCopyKernel on device {}: {}", deviceId, + hipGetErrorString(err)); + (void)hipModuleUnload(fabricCopyModules_[deviceId]); + fabricCopyModules_[deviceId] = nullptr; + return nullptr; + } + MORI_IO_INFO("FABRIC: Loaded copy kernel on device {}", deviceId); + return fabricCopyFuncs_[deviceId]; +} + +void FabricBackend::RegisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteEnginesMu); + remoteEngines[desc.key] = desc; + MORI_IO_TRACE("FABRIC: Registered remote engine {} hostname {}", desc.key, desc.hostname); +} + +void FabricBackend::DeregisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteEnginesMu); + remoteEngines.erase(desc.key); + MORI_IO_TRACE("FABRIC: Deregistered remote engine {}", desc.key); +} + +void FabricBackend::RegisterMemory(MemoryDesc& desc) { + if (desc.loc != MemoryLocationType::GPU) { + MORI_IO_TRACE("FABRIC: Skipping non-GPU memory registration for id={}", desc.id); + return; + } + + int dev = desc.deviceId; + // Advertise the owning device's vPOD identity so the peer's CanHandle can + // decide whether this memory is reachable over the same scale-up fabric. + if (dev >= 0 && dev < numDevices) { + const fabric::VpodKey& key = deviceVpodKeys[dev]; + if (key.valid) { + desc.vpodId = key.vpodId; + desc.vpodPpodId = key.ppodId; + } + if (!deviceSupportsFabric[dev]) { + MORI_IO_TRACE("FABRIC: device {} does not support fabric; memory id={} not exportable", dev, + desc.id); + return; + } + } + + // Export a fabric shareable handle by retaining the allocation backing the + // registered pointer. This only succeeds for VMM allocations created with the + // fabric handle type (e.g. via FabricMalloc); plain hipMalloc memory fails + // here, leaving fabricHandle empty so CanHandle falls back to RDMA. + ScopedHipDeviceGuard deviceGuard; + if (dev >= 0) (void)hipSetDevice(dev); + + hipMemGenericAllocationHandle_t handle{}; + hipError_t err = hipMemRetainAllocationHandle(&handle, reinterpret_cast(desc.data)); + if (err != hipSuccess) { + (void)hipGetLastError(); + MORI_IO_TRACE("FABRIC: memory id={} is not a VMM allocation ({}), fabric export skipped", + desc.id, hipGetErrorString(err)); + return; + } + + hipMemFabricHandle_compat_t fh; + err = hipMemExportToShareableHandle(&fh, handle, hipMemHandleTypeFabricCompat, 0); + (void)hipMemRelease(handle); // drop the extra ref taken by Retain + if (err != hipSuccess) { + (void)hipGetLastError(); + MORI_IO_TRACE("FABRIC: fabric export failed for memory id={}: {}", desc.id, + hipGetErrorString(err)); + return; + } + + static_assert(sizeof(fh) == kFabricHandleSize, "fabric handle size mismatch"); + std::memcpy(desc.fabricHandle.data(), &fh, sizeof(fh)); + MORI_IO_TRACE("FABRIC: Registered memory id={}, addr={}, size={} (fabric-exportable)", desc.id, + desc.data, desc.size); +} + +void FabricBackend::DeregisterMemory(const MemoryDesc& desc) { + { + std::unique_lock lock(importMutex); + for (auto it = importedRegions.begin(); it != importedRegions.end();) { + if (it->first.memId == desc.id && it->first.engineKey == desc.engineKey) { + if (it->second.mappedAddr != nullptr) { + (void)hipMemUnmap(it->second.mappedAddr, it->second.size); + (void)hipMemRelease(it->second.handle); + (void)hipMemAddressFree(it->second.mappedAddr, it->second.size); + } + it = importedRegions.erase(it); + } else { + ++it; + } + } + } + InvalidateSessionsForMemory(desc.id); + MORI_IO_TRACE("FABRIC: Deregistered memory id={}", desc.id); +} + +void* FabricBackend::GetImportedAddress(const MemoryDesc& desc, int localDeviceId) { + if (desc.engineKey == myEngKey) { + return reinterpret_cast(desc.data); + } + if (IsFabricHandleEmpty(desc.fabricHandle)) { + return nullptr; + } + + ImportCacheKey cacheKey{desc.engineKey, desc.id, localDeviceId}; + { + std::shared_lock rlock(importMutex); + auto it = importedRegions.find(cacheKey); + if (it != importedRegions.end() && it->second.mappedAddr != nullptr) { + return it->second.mappedAddr; + } + } + + hipMemFabricHandle_compat_t fh; + static_assert(sizeof(fh) == kFabricHandleSize, "fabric handle size mismatch"); + std::memcpy(&fh, desc.fabricHandle.data(), sizeof(fh)); + + ScopedHipDeviceGuard deviceGuard; + hipError_t err = hipSetDevice(localDeviceId); + if (err != hipSuccess) { + MORI_IO_WARN("FABRIC: Failed to set device {} for fabric import: {}", localDeviceId, + hipGetErrorString(err)); + return nullptr; + } + + hipMemGenericAllocationHandle_t handle{}; + err = hipMemImportFromShareableHandle(&handle, reinterpret_cast(&fh), + hipMemHandleTypeFabricCompat); + if (err != hipSuccess) { + (void)hipGetLastError(); + MORI_IO_WARN("FABRIC: hipMemImportFromShareableHandle failed for id={}: {}", desc.id, + hipGetErrorString(err)); + return nullptr; + } + + size_t mapSize = AlignToFabricGranularity(localDeviceId, desc.size); + void* va = nullptr; + err = hipMemAddressReserve(&va, mapSize, 0, 0, 0); + if (err != hipSuccess) { + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: hipMemAddressReserve failed for id={}: {}", desc.id, + hipGetErrorString(err)); + return nullptr; + } + err = hipMemMap(va, mapSize, 0, handle, 0); + if (err != hipSuccess) { + (void)hipMemAddressFree(va, mapSize); + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: hipMemMap failed for id={}: {}", desc.id, hipGetErrorString(err)); + return nullptr; + } + hipMemAccessDesc access = {}; + access.location.type = hipMemLocationTypeDevice; + access.location.id = localDeviceId; + access.flags = hipMemAccessFlagsProtReadWrite; + err = hipMemSetAccess(va, mapSize, &access, 1); + if (err != hipSuccess) { + (void)hipMemUnmap(va, mapSize); + (void)hipMemAddressFree(va, mapSize); + (void)hipMemRelease(handle); + MORI_IO_WARN("FABRIC: hipMemSetAccess failed for id={}: {}", desc.id, hipGetErrorString(err)); + return nullptr; + } + + std::unique_lock wlock(importMutex); + // Another thread may have imported concurrently; keep the first winner. + auto existing = importedRegions.find(cacheKey); + if (existing != importedRegions.end() && existing->second.mappedAddr != nullptr) { + void* winner = existing->second.mappedAddr; + wlock.unlock(); + (void)hipMemUnmap(va, mapSize); + (void)hipMemAddressFree(va, mapSize); + (void)hipMemRelease(handle); + return winner; + } + importedRegions[cacheKey] = {handle, va, mapSize}; + MORI_IO_TRACE("FABRIC: Imported fabric handle for id={} on device {}, mapped={}", desc.id, + localDeviceId, va); + return va; +} + +void FabricBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, + const MemoryDesc& remoteSrc, size_t remoteOffset, size_t size, + TransferStatus* status, TransferUniqueId id, bool isRead) { + FabricBackendSession* sess = GetOrCreateSessionCached(localDest, remoteSrc); + if (sess == nullptr) { + status->Update(StatusCode::ERR_BAD_STATE, "FABRIC: Failed to create session"); + return; + } + sess->ReadWrite(localOffset, remoteOffset, size, status, id, isRead); +} + +void FabricBackend::BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, + const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id, bool isRead) { + FabricBackendSession* sess = GetOrCreateSessionCached(localDest, remoteSrc); + if (sess == nullptr) { + status->Update(StatusCode::ERR_BAD_STATE, "FABRIC: Failed to create session"); + return; + } + sess->BatchReadWrite(localOffsets, remoteOffsets, sizes, status, id, isRead); +} + +BackendSession* FabricBackend::CreateSession(const MemoryDesc& local, const MemoryDesc& remote) { + int localDevice = local.deviceId; + void* localAddr = reinterpret_cast(local.data); + void* remoteAddr = GetImportedAddress(remote, localDevice); + if (localAddr == nullptr || remoteAddr == nullptr) { + MORI_IO_WARN("FABRIC: Failed to map memory for session (local id={}, remote id={})", local.id, + remote.id); + return nullptr; + } + return new FabricBackendSession(config, localAddr, remoteAddr, localDevice, this, + streamPool.get(), eventPool.get()); +} + +FabricBackendSession* FabricBackend::GetOrCreateSessionCached(const MemoryDesc& local, + const MemoryDesc& remote) { + SessionCacheKey key{remote.engineKey, local.id, remote.id}; + + std::lock_guard lock(sessionCacheMu); + auto it = sessionCache.find(key); + if (it != sessionCache.end()) { + return it->second.get(); + } + + int localDevice = local.deviceId; + void* localAddr = reinterpret_cast(local.data); + void* remoteAddr = GetImportedAddress(remote, localDevice); + if (remoteAddr == nullptr) { + MORI_IO_WARN("FABRIC: Failed to import remote memory for session (local.id={}, remote.id={})", + local.id, remote.id); + return nullptr; + } + + auto sess = std::make_unique(config, localAddr, remoteAddr, localDevice, + this, streamPool.get(), eventPool.get()); + FabricBackendSession* rawPtr = sess.get(); + sessionCache[key] = std::move(sess); + MORI_IO_TRACE("FABRIC: Created session for local.id={}, remote.id={}", local.id, remote.id); + return rawPtr; +} + +void FabricBackend::InvalidateSessionsForMemory(MemoryUniqueId id) { + std::lock_guard lock(sessionCacheMu); + for (auto it = sessionCache.begin(); it != sessionCache.end();) { + if (it->first.localMemId == id || it->first.remoteMemId == id) { + it = sessionCache.erase(it); + } else { + ++it; + } + } +} + +bool FabricBackend::PopInboundTransferStatus(EngineKey /*remote*/, TransferUniqueId /*id*/, + TransferStatus* /*status*/) { + return false; +} + +bool FabricBackend::CanHandle(const MemoryDesc& local, const MemoryDesc& remote) const { + if (local.loc != MemoryLocationType::GPU || remote.loc != MemoryLocationType::GPU) { + return false; + } + if (local.deviceId < 0 || local.deviceId >= numDevices) { + return false; + } + if (!deviceSupportsFabric[local.deviceId]) { + return false; + } + // Remote memory must be fabric-exportable (non-empty handle). + if (IsFabricHandleEmpty(remote.fabricHandle)) { + return false; + } + + // Both endpoints must sit in the same scale-up fabric domain (vPOD). + fabric::VpodKey localKey = LocalVpodKey(local.deviceId); + fabric::VpodKey remoteKey; + remoteKey.vpodId = remote.vpodId; + remoteKey.ppodId = remote.vpodPpodId; + remoteKey.valid = (remote.vpodId >= 0) || !remote.vpodPpodId.empty(); + return fabric::VpodKeySame(localKey, remoteKey); +} + +} // namespace io +} // namespace mori diff --git a/src/io/fabric/backend_impl.hpp b/src/io/fabric/backend_impl.hpp new file mode 100644 index 000000000..0e1eae674 --- /dev/null +++ b/src/io/fabric/backend_impl.hpp @@ -0,0 +1,214 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/backend.hpp" +#include "mori/io/common.hpp" +#include "mori/io/engine.hpp" +#include "src/io/fabric/vpod_topology.hpp" +#include "src/io/xgmi/hip_resource_pool.hpp" + +namespace mori { +namespace io { + +/* ---------------------------------------------------------------------------------------------- */ +/* Fabric allocation helpers */ +/* ---------------------------------------------------------------------------------------------- */ +// Allocate GPU memory that is exportable over a UALink super-node fabric, using +// the low-level VMM API (hipMemCreate + reserve + map + setAccess) with +// requestedHandleType = fabric. This is what makes a buffer registerable with +// FabricBackend: only fabric-exportable allocations can be shared cross-node. +// +// Returns the mapped device pointer (nullptr on failure). The allocation is +// tracked internally so FabricFree can tear it down. Note: do NOT hipMemset a +// fabric-exportable buffer (returns OOM on ROCm 7.15); zero it via host copies +// if needed. +void* FabricMalloc(size_t size, int device); +void FabricFree(void* ptr); + +/* ---------------------------------------------------------------------------------------------- */ +/* FabricBackendSession */ +/* ---------------------------------------------------------------------------------------------- */ +class FabricBackend; + +class FabricBackendSession : public BackendSession { + public: + FabricBackendSession() = default; + FabricBackendSession(const FabricBackendConfig& config, void* localAddr, void* remoteAddr, + int localDevice, FabricBackend* backend, StreamPool* streamPool, + EventPool* eventPool); + ~FabricBackendSession() = default; + + void ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status, + TransferUniqueId id, bool isRead) override; + + void BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) override; + + bool Alive() const override; + + private: + // Launch the copy kernel over [srcBase+srcOff, dstBase+dstOff) for `size` + // bytes on `stream`. Returns hipSuccess or the launch error. + hipError_t LaunchCopy(void* dst, const void* src, size_t size, hipStream_t stream); + + FabricBackendConfig config{}; + void* localAddr{nullptr}; // local registered pointer + void* remoteAddr{nullptr}; // imported peer VA (aliases remote HBM over fabric) + int localDevice{-1}; // device that drives the copy kernel + FabricBackend* backend{nullptr}; + StreamPool* streamPool{nullptr}; + EventPool* eventPool{nullptr}; +}; + +/* ---------------------------------------------------------------------------------------------- */ +/* FabricBackend */ +/* ---------------------------------------------------------------------------------------------- */ +class FabricBackend : public Backend { + public: + FabricBackend(EngineKey, const IOEngineConfig&, const FabricBackendConfig&); + ~FabricBackend(); + + void RegisterRemoteEngine(const EngineDesc&) override; + void DeregisterRemoteEngine(const EngineDesc&) override; + void RegisterMemory(MemoryDesc& desc) override; + void DeregisterMemory(const MemoryDesc& desc) override; + void ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc, + size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, + bool isRead) override; + void BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, + const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) override; + BackendSession* CreateSession(const MemoryDesc& local, const MemoryDesc& remote) override; + bool PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, + TransferStatus* status) override; + bool CanHandle(const MemoryDesc& local, const MemoryDesc& remote) const override; + + // JIT-loaded contiguous copy kernel (fabric_copy.hip), one module per device. + void LoadFabricCopyModule(const std::string& hsacoPath); + hipFunction_t GetFabricCopyFunc(int deviceId); + + private: + void Initialize(); + std::optional LookupVisibleDevice(const std::string& busId) const; + fabric::VpodKey LocalVpodKey(int deviceId) const; + + // Import a peer's fabric handle and map it into a fresh local VA; cached per + // (remote engine, remote mem id, local device). Returns the mapped pointer, or + // the local pointer directly when the memory is owned by this engine. + void* GetImportedAddress(const MemoryDesc& desc, int localDeviceId); + + struct SessionCacheKey { + EngineKey remoteEngineKey; + MemoryUniqueId localMemId; + MemoryUniqueId remoteMemId; + bool operator==(const SessionCacheKey& o) const { + return remoteEngineKey == o.remoteEngineKey && localMemId == o.localMemId && + remoteMemId == o.remoteMemId; + } + }; + struct SessionCacheKeyHash { + std::size_t operator()(const SessionCacheKey& k) const noexcept { + std::size_t seed = 0; + auto hashCombine = [](std::size_t& s, std::size_t v) { + s ^= v + 0x9e3779b97f4a7c15ULL + (s << 6) + (s >> 2); + }; + hashCombine(seed, std::hash()(k.remoteEngineKey)); + hashCombine(seed, std::hash()(k.localMemId)); + hashCombine(seed, std::hash()(k.remoteMemId)); + return seed; + } + }; + FabricBackendSession* GetOrCreateSessionCached(const MemoryDesc& local, const MemoryDesc& remote); + void InvalidateSessionsForMemory(MemoryUniqueId id); + + struct ImportedRegion { + hipMemGenericAllocationHandle_t handle{}; + void* mappedAddr{nullptr}; + size_t size{0}; + }; + struct ImportCacheKey { + EngineKey engineKey; + MemoryUniqueId memId; + int deviceId; + bool operator==(const ImportCacheKey& o) const { + return engineKey == o.engineKey && memId == o.memId && deviceId == o.deviceId; + } + }; + struct ImportCacheKeyHash { + std::size_t operator()(const ImportCacheKey& k) const noexcept { + std::size_t seed = 0; + auto hashCombine = [](std::size_t& s, std::size_t v) { + s ^= v + 0x9e3779b97f4a7c15ULL + (s << 6) + (s >> 2); + }; + hashCombine(seed, std::hash()(k.engineKey)); + hashCombine(seed, std::hash()(k.memId)); + hashCombine(seed, std::hash()(k.deviceId)); + return seed; + } + }; + + private: + EngineKey myEngKey; + std::string myNodeId; + std::string myHostname; + FabricBackendConfig config; + int myPid{0}; + + std::unique_ptr streamPool; + std::unique_ptr eventPool; + + int numDevices{0}; + std::vector deviceSupportsFabric; // indexed by device id + std::vector deviceVpodKeys; // indexed by device id + std::unordered_map localDeviceByBusId; + + mutable std::shared_mutex importMutex; + std::unordered_map importedRegions; + + std::unordered_map, SessionCacheKeyHash> + sessionCache; + std::mutex sessionCacheMu; + + std::unordered_map remoteEngines; + mutable std::mutex remoteEnginesMu; + + std::string fabricCopyHsacoPath_; + std::vector fabricCopyModules_; + std::vector fabricCopyFuncs_; + std::mutex moduleMu_; +}; + +} // namespace io +} // namespace mori diff --git a/src/io/fabric/vpod_topology.cpp b/src/io/fabric/vpod_topology.cpp new file mode 100644 index 000000000..fe3aa2bea --- /dev/null +++ b/src/io/fabric/vpod_topology.cpp @@ -0,0 +1,128 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "src/io/fabric/vpod_topology.hpp" + +#include + +#include +#include +#include +#include + +#include "mori/io/logging.hpp" + +namespace mori { +namespace io { +namespace fabric { +namespace { + +// Read the first line of a sysfs file, trimming trailing whitespace/newlines. +bool ReadSysfsLine(const std::string& path, std::string& out) { + FILE* f = fopen(path.c_str(), "r"); + if (f == nullptr) return false; + char buf[128] = {0}; + bool ok = fgets(buf, sizeof(buf), f) != nullptr; + fclose(f); + if (!ok) return false; + out = buf; + while (!out.empty() && (out.back() == '\n' || out.back() == '\r' || out.back() == ' ')) { + out.pop_back(); + } + return true; +} + +// The PCI BDF as sysfs spells it (lowercase). hipDeviceGetPCIBusId may return +// uppercase hex, whereas sysfs directory names are lowercase. +std::string DeviceBdfLower(int hipDevice) { + char bdf[32] = {0}; + if (hipDeviceGetPCIBusId(bdf, sizeof(bdf), hipDevice) != hipSuccess) return ""; + for (char* p = bdf; *p; ++p) *p = static_cast(std::tolower(static_cast(*p))); + return std::string(bdf); +} + +} // namespace + +bool DeviceSupportsFabric(int hipDevice) { + 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; +} + +VpodKey ReadVpodKey(int hipDevice) { + VpodKey key; + + const char* disable = std::getenv("MORI_IO_FABRIC_DISABLE"); + if (disable != nullptr && std::atoi(disable) != 0) { + return key; // forced invalid + } + + // Manual override: every engine that sets the same value is treated as one + // vPOD (ppodId stays empty). Useful when sysfs is unavailable but the operator + // knows the topology. + const char* manual = std::getenv("MORI_IO_FABRIC_VPOD_ID"); + if (manual != nullptr && manual[0] != '\0') { + key.valid = true; + key.vpodId = std::atoi(manual); + key.vpodSize = 0; + return key; + } + + std::string bdf = DeviceBdfLower(hipDevice); + if (bdf.empty()) return key; + std::string dir = std::string("/sys/bus/pci/devices/") + bdf + "/ualink"; + + std::string link, state, ppod; + if (!ReadSysfsLine(dir + "/link_type", link)) return key; // not a UALink GPU + if (link != "UALoE" && link != "UALLink") return key; + if (!ReadSysfsLine(dir + "/accel_state", state)) return key; + if (state != "active" && state != "ready") return key; // not usable yet + if (!ReadSysfsLine(dir + "/ppod_id", ppod) || ppod.empty()) return key; + + std::string tmp; + key.ppodId = ppod; + key.vpodId = ReadSysfsLine(dir + "/vpod_id", tmp) ? std::atoi(tmp.c_str()) : 0; + key.vpodSize = ReadSysfsLine(dir + "/vpod_size", tmp) ? std::atoi(tmp.c_str()) : 0; + key.valid = true; + MORI_IO_TRACE("FABRIC: device {} vPOD ppod_id={} vpod_id={} vpod_size={}", hipDevice, key.ppodId, + key.vpodId, key.vpodSize); + return key; +} + +bool VpodKeySame(const VpodKey& a, const VpodKey& b) { + if (!a.valid || !b.valid) return false; + return a.vpodId == b.vpodId && a.ppodId == b.ppodId; +} + +} // namespace fabric +} // namespace io +} // namespace mori diff --git a/src/io/fabric/vpod_topology.hpp b/src/io/fabric/vpod_topology.hpp new file mode 100644 index 000000000..1ed24e9d0 --- /dev/null +++ b/src/io/fabric/vpod_topology.hpp @@ -0,0 +1,60 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include + +namespace mori { +namespace io { +namespace fabric { + +// UALink vPOD (super-node scale-up domain) identity for a single GPU. Two GPUs +// can load/store each other's memory over the fabric iff their keys are valid +// and compare equal. Mirrors mori-cco's CcoLsaKey / RCCL's MNNVL clique model: +// ppod_id (a UUID) makes the key globally unique, since hive_id collides across +// hosts. +struct VpodKey { + bool valid{false}; // true only when the fabric is present and ACTIVE/READY + std::string ppodId; // physical pod UUID (empty in manual-override mode) + int vpodId{-1}; // virtual pod id within the ppod + int vpodSize{0}; // number of GPUs in the vPOD (0 if unknown) +}; + +// Does the device advertise VMM + fabric-handle export capability +// (hipDeviceAttributeVirtualMemoryManagementSupported && +// hipDeviceAttributeHandleTypeFabricSupported)? +bool DeviceSupportsFabric(int hipDevice); + +// Read the GPU's UALink fabric identity from sysfs (mirrors mori-cco +// CcoReadFabricKey). Returns an invalid key when the device is not on a ready +// UALink fabric. Honors env overrides: +// MORI_IO_FABRIC_DISABLE=1 -> always invalid (force RDMA/XGMI fallback) +// MORI_IO_FABRIC_VPOD_ID= -> manual vPOD id (ppodId left empty; all ranks +// sharing the value are treated as one vPOD) +VpodKey ReadVpodKey(int hipDevice); + +// True iff both keys are valid and describe the same scale-up domain. +bool VpodKeySame(const VpodKey& a, const VpodKey& b); + +} // namespace fabric +} // namespace io +} // namespace mori diff --git a/src/io/kernels/fabric_copy.hip b/src/io/kernels/fabric_copy.hip new file mode 100644 index 000000000..2252b5b68 --- /dev/null +++ b/src/io/kernels/fabric_copy.hip @@ -0,0 +1,66 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Contiguous copy kernel for FabricBackend transfers over a UALink super-node +// fabric. JIT-compiled via hipcc --genco at runtime and loaded as a .hsaco +// module (same mechanism as the XGMI scatter/gather kernel). +// +// Why a kernel and not hipMemcpy: on an *imported fabric pointer*, +// hipMemcpy(..., hipMemcpyDeviceToDevice) lowers to a pathological blit path +// that livelocks for large payloads (observed >= 4 MB on ROCm 7.15 / gfx1250). +// A vectorized grid-stride copy drives the fabric link directly — remote loads +// for a read, remote stores for a write — and is the path the 07_ualoe +// micro-benchmark uses to reach line rate. + +#include + +#include + +// dst/src are byte pointers; either may alias remote (imported) HBM. When both +// are 16-byte aligned the bulk is copied as uint4 (vectorized); a scalar tail +// handles the remainder and the fully-unaligned case. +extern "C" __global__ void fabricCopyKernel(char* __restrict__ dst, const char* __restrict__ src, + size_t nbytes) { + size_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + size_t nthreads = static_cast(gridDim.x) * blockDim.x; + + uintptr_t s = reinterpret_cast(src); + uintptr_t d = reinterpret_cast(dst); + + if ((s & 15) == 0 && (d & 15) == 0 && nbytes >= 16) { + size_t nvec = nbytes / 16; + size_t tailStart = nvec * 16; + + const uint4* srcV = reinterpret_cast(src); + uint4* dstV = reinterpret_cast(dst); + for (size_t i = tid; i < nvec; i += nthreads) { + dstV[i] = srcV[i]; + } + for (size_t i = tailStart + tid; i < nbytes; i += nthreads) { + dst[i] = src[i]; + } + } else { + for (size_t i = tid; i < nbytes; i += nthreads) { + dst[i] = src[i]; + } + } +} diff --git a/src/pybind/pybind_io.cpp b/src/pybind/pybind_io.cpp index f1ba2d269..fc452ef95 100644 --- a/src/pybind/pybind_io.cpp +++ b/src/pybind/pybind_io.cpp @@ -29,6 +29,13 @@ namespace py = pybind11; namespace mori { +namespace io { +// Fabric allocation helpers (defined in src/io/fabric/backend_impl.cpp). +// Forward-declared here to avoid pulling internal backend headers into pybind. +void* FabricMalloc(size_t size, int device); +void FabricFree(void* ptr); +} // namespace io + void RegisterMoriIo(pybind11::module_& m) { m.def("set_log_level", &mori::io::SetLogLevel); @@ -37,6 +44,7 @@ void RegisterMoriIo(pybind11::module_& m) { .value("XGMI", mori::io::BackendType::XGMI) .value("RDMA", mori::io::BackendType::RDMA) .value("TCP", mori::io::BackendType::TCP) + .value("FABRIC", mori::io::BackendType::FABRIC) .export_values(); py::enum_(m, "MemoryLocationType") @@ -90,6 +98,23 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("num_streams", &mori::io::XgmiBackendConfig::numStreams) .def_readwrite("num_events", &mori::io::XgmiBackendConfig::numEvents); + py::class_(m, "FabricBackendConfig") + .def(py::init(), py::arg("num_streams") = 64, py::arg("num_events") = 64) + .def_readwrite("num_streams", &mori::io::FabricBackendConfig::numStreams) + .def_readwrite("num_events", &mori::io::FabricBackendConfig::numEvents); + + // Allocate/free GPU memory that is exportable over a UALink super-node fabric. + // Returns the device pointer as an integer, suitable for register_memory(). + m.def( + "fabric_alloc", + [](size_t size, int device) -> uintptr_t { + return reinterpret_cast(mori::io::FabricMalloc(size, device)); + }, + py::arg("size"), py::arg("device")); + m.def( + "fabric_free", [](uintptr_t ptr) { mori::io::FabricFree(reinterpret_cast(ptr)); }, + py::arg("ptr")); + py::class_(m, "IOEngineConfig") .def(py::init(), py::arg("host") = "", py::arg("port") = 0) .def_readwrite("host", &mori::io::IOEngineConfig::host) @@ -193,7 +218,8 @@ void RegisterMoriIo(pybind11::module_& m) { py::call_guard()) .def("WaitAll", &mori::io::IOEngine::WaitAll, py::arg("statuses"), py::arg("timeout_ms") = -1, py::call_guard()) - .def("LoadScatterGatherModule", &mori::io::IOEngine::LoadScatterGatherModule); + .def("LoadScatterGatherModule", &mori::io::IOEngine::LoadScatterGatherModule) + .def("LoadFabricCopyModule", &mori::io::IOEngine::LoadFabricCopyModule); } } // namespace mori diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 654bf0415..a92fa106f 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -28,26 +28,81 @@ IOEngine, EngineDesc, MemoryDesc, + MemoryLocationType, PollCqMode, RdmaBackendConfig, XgmiBackendConfig, + FabricBackendConfig, + fabric_alloc, set_log_level, ) import argparse +import ctypes from enum import Enum import os import time from prettytable import PrettyTable +# --- Minimal HIP memcpy helpers (fabric buffers are raw VMM allocations, not +# --- torch tensors, so we fill/verify them directly via the HIP runtime). +_hip_lib = None + + +def _hip(): + global _hip_lib + if _hip_lib is None: + for name in ("libamdhip64.so", "libamdhip64.so.7", "libamdhip64.so.6"): + try: + _hip_lib = ctypes.CDLL(name) + break + except OSError: + continue + if _hip_lib is None: + raise RuntimeError("Could not load libamdhip64.so for fabric validation") + _hip_lib.hipMemcpy.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ] + _hip_lib.hipMemcpy.restype = ctypes.c_int + _hip_lib.hipSetDevice.argtypes = [ctypes.c_int] + _hip_lib.hipSetDevice.restype = ctypes.c_int + _hip_lib.hipDeviceSynchronize.restype = ctypes.c_int + return _hip_lib + + +_HIP_MEMCPY_H2D = 1 +_HIP_MEMCPY_D2H = 2 + + +def hip_fill(ptr: int, data: bytes, device: int): + h = _hip() + assert h.hipSetDevice(device) == 0 + buf = (ctypes.c_char * len(data)).from_buffer_copy(data) + assert h.hipMemcpy(ctypes.c_void_p(ptr), buf, len(data), _HIP_MEMCPY_H2D) == 0 + assert h.hipDeviceSynchronize() == 0 + + +def hip_read(ptr: int, nbytes: int, device: int) -> bytes: + h = _hip() + assert h.hipSetDevice(device) == 0 + buf = (ctypes.c_char * nbytes)() + assert h.hipMemcpy(buf, ctypes.c_void_p(ptr), nbytes, _HIP_MEMCPY_D2H) == 0 + assert h.hipDeviceSynchronize() == 0 + return bytes(buf) + + def parse_args(): parser = argparse.ArgumentParser(description="Benchmark MORI-IO") parser.add_argument( "--backend", type=str, - choices=["rdma", "xgmi"], + choices=["rdma", "xgmi", "fabric"], default="rdma", - help="Backend type: 'rdma' for cross-node, 'xgmi' for intra-node GPU-to-GPU (default: rdma)", + help="Backend type: 'rdma' for cross-node RDMA, 'xgmi' for intra-node GPU-to-GPU, " + "'fabric' for cross-node scale-up UALink super-node (same vPOD) (default: rdma)", ) parser.add_argument( "--host", @@ -330,6 +385,8 @@ def __init__( if self.backend_type == "xgmi": self._setup_xgmi() + elif self.backend_type == "fabric": + self._setup_fabric() else: self._setup_rdma() @@ -400,6 +457,45 @@ def _setup_xgmi(self): self.dst_device, dtype=torch.float8_e4m3fnuz ) + def _setup_fabric(self): + # Fabric mirrors the RDMA cross-node role model (2 nodes, gloo OOB), but + # the transferred buffer MUST be a fabric-exportable VMM allocation + # (fabric_alloc) rather than a torch tensor. + if self.mem_type != "gpu": + raise ValueError("fabric backend supports GPU memory only") + assert self.num_initiator_dev == self.num_target_dev + self.world_size = self.num_initiator_dev + self.num_target_dev + if self.node_rank == 0: + self.global_rank = self.role_rank + self.role = EngineRole.INITIATOR + else: + self.global_rank = self.role_rank + self.num_initiator_dev + self.role = EngineRole.TARGET + + # 1 byte per element (matches the fp8/uint8 sizing used by the other paths). + total_elements = ( + (self.buffer_size + 1) * self.transfer_batch_size + if not self.batch_contiguous + else self.buffer_size * self.transfer_batch_size + ) + gpu_index = self.role_rank + if self.role is EngineRole.TARGET and self.target_dev_offset: + ndev = torch.cuda.device_count() + if ndev <= 0: + raise RuntimeError( + "fabric setup: torch.cuda.device_count()==0 (no visible GPUs); " + "cannot apply --target-dev-offset" + ) + gpu_index = (self.role_rank + self.target_dev_offset) % ndev + self.device = torch.device("cuda", gpu_index) + self.fabric_nbytes = total_elements + self.fabric_ptr = fabric_alloc(self.fabric_nbytes, gpu_index) + if not self.fabric_ptr: + raise RuntimeError( + f"fabric_alloc failed for {self.fabric_nbytes} bytes on gpu {gpu_index}; " + "the device must be a UALink fabric-capable GPU in a ready vPOD" + ) + def print_config(self): print("MORI-IO Benchmark Configurations:") print(f" backend: {self.backend_type.upper()}") @@ -411,6 +507,15 @@ def print_config(self): print(f" dst_gpu: {self.dst_gpu}") print(f" num_streams: {self.num_streams}") print(f" num_events: {self.num_events}") + elif self.backend_type == "fabric": + print(f" node_rank: {self.node_rank}") + print(f" role: {self.role}") + print(f" role_rank: {self.role_rank}") + print(f" num_initiator_dev: {self.num_initiator_dev}") + print(f" num_target_dev: {self.num_target_dev}") + print(f" target_dev_offset: {self.target_dev_offset}") + print(f" num_streams: {self.num_streams}") + print(f" num_events: {self.num_events}") else: print(f" host: {self.host}") print(f" port: {self.port}") @@ -480,9 +585,64 @@ def recv_bytes(self, src: int) -> bytes: def validate(self): if self.backend_type == "xgmi": self._validate_xgmi() + elif self.backend_type == "fabric": + self._validate_fabric() else: self._validate_rdma() + def _validate_fabric(self): + # Correctness spot-check decoupled from the (possibly strided/batched) + # perf pattern: target fills a known pattern, initiator transfers it over + # the fabric and both sides verify. Bounded to keep it cheap. + check = min(self.fabric_nbytes, 1 << 20) + tile = bytes(range(256)) + pattern = (tile * (check // 256 + 1))[:check] + zero = bytes(check) + dev = self.device.index + + if self.op_type == "read": + if self.role is EngineRole.TARGET: + hip_fill(self.fabric_ptr, pattern, dev) + dist.barrier() # #1 target data ready + dist.barrier() # #2 initiator finished + else: + hip_fill(self.fabric_ptr, zero, dev) + dist.barrier() # #1 wait for target fill + status = self.engine.read( + self.mem, + 0, + self.target_mem, + 0, + check, + self.engine.allocate_transfer_uid(), + ) + status.Wait() + assert status.Succeeded(), f"fabric read failed: {status.Message()}" + got = hip_read(self.fabric_ptr, check, dev) + assert got == pattern, "fabric READ validation: data mismatch" + dist.barrier() # #2 done + else: # write + if self.role is EngineRole.INITIATOR: + hip_fill(self.fabric_ptr, pattern, dev) + dist.barrier() # #1 both ready + status = self.engine.write( + self.mem, + 0, + self.target_mem, + 0, + check, + self.engine.allocate_transfer_uid(), + ) + status.Wait() + assert status.Succeeded(), f"fabric write failed: {status.Message()}" + dist.barrier() # #2 signal target to verify + else: + hip_fill(self.fabric_ptr, zero, dev) + dist.barrier() # #1 both ready + dist.barrier() # #2 initiator wrote + got = hip_read(self.fabric_ptr, check, dev) + assert got == pattern, "fabric WRITE validation: data mismatch" + def _validate_rdma(self): if self.role is EngineRole.INITIATOR: recv_tensor = torch.empty( @@ -571,6 +731,8 @@ def _validate_xgmi(self): def initialize(self): if self.backend_type == "xgmi": self._initialize_xgmi() + elif self.backend_type == "fabric": + self._initialize_fabric() else: self._initialize_rdma() @@ -677,10 +839,60 @@ def _initialize_xgmi(self): if self.enable_sess: self.sess = self.engine.create_session(self.mem, self.target_mem) + def _initialize_fabric(self): + # Same OOB/role flow as RDMA (gloo desc exchange), but a FABRIC backend + # and register_memory() on the raw fabric_alloc pointer. The FABRIC + # backend has no TCP control plane; descriptors are exchanged via gloo. + config = IOEngineConfig(host="", port=0) + self.engine = IOEngine(key=f"{self.role.name}-{self.role_rank}", config=config) + fabric_config = FabricBackendConfig( + num_streams=self.num_streams, + num_events=self.num_events, + ) + self.engine.create_backend(BackendType.FABRIC, fabric_config) + + self.engine_desc = self.engine.get_engine_desc() + engine_desc_bytes = self.engine_desc.pack() + + if self.role is EngineRole.INITIATOR: + for i in range(self.num_target_dev): + self.send_bytes(engine_desc_bytes, self.num_initiator_dev + i) + for i in range(self.num_target_dev): + peer_engine_desc_bytes = self.recv_bytes(self.num_initiator_dev + i) + peer_engine_desc = EngineDesc.unpack(peer_engine_desc_bytes) + self.engine.register_remote_engine(peer_engine_desc) + else: + for i in range(self.num_initiator_dev): + peer_engine_desc_bytes = self.recv_bytes(i) + peer_engine_desc = EngineDesc.unpack(peer_engine_desc_bytes) + self.engine.register_remote_engine(peer_engine_desc) + for i in range(self.num_initiator_dev): + self.send_bytes(engine_desc_bytes, i) + + self.mem = self.engine.register_memory( + self.fabric_ptr, + self.fabric_nbytes, + self.device.index, + MemoryLocationType.GPU, + ) + + if self.role is EngineRole.TARGET: + mem_desc = self.mem.pack() + self.send_bytes(mem_desc, self.role_rank) + else: + target_mem_desc = self.recv_bytes(self.num_initiator_dev + self.role_rank) + self.target_mem = MemoryDesc.unpack(target_mem_desc) + self.sess = self.engine.create_session(self.mem, self.target_mem) + if self.sess is None: + raise RuntimeError( + "create_session returned None for fabric: peers likely not in the " + "same vPOD, or remote memory is not fabric-exportable" + ) + def run_single_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: return 0 @@ -732,7 +944,7 @@ def run_single_once(self, buffer_size, transfer_batch_size): def run_batch_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: return 0 @@ -795,7 +1007,7 @@ def _run_and_compute(self, buffer_size, transfer_batch_size, iters): latency.append(duration) if self.role is EngineRole.TARGET and ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ): return 0, 0, 0, 0, 0 @@ -816,6 +1028,8 @@ def _get_table_title(self): return f"XGMI Multiprocess Benchmark: Rank {self.role_rank} ({self.role.name})" else: return f"XGMI Benchmark: GPU{self.src_gpu} -> GPU{self.dst_gpu}" + elif self.backend_type == "fabric": + return f"FABRIC Benchmark: Initiator Rank {self.role_rank}" else: return f"RDMA Benchmark: Initiator Rank {self.role_rank}" @@ -839,7 +1053,7 @@ def _run_benchmark_loop(self): cur_size = self.sweep_start_size max_size = self.sweep_max_size while cur_size <= max_size: - if self.backend_type == "rdma" or ( + if self.backend_type in ("rdma", "fabric") or ( self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() @@ -864,7 +1078,7 @@ def _run_benchmark_loop(self): cur_transfer_batch_size = 1 max_transfer_batch_size = 32768 while cur_transfer_batch_size <= max_transfer_batch_size: - if self.backend_type == "rdma" or ( + if self.backend_type in ("rdma", "fabric") or ( self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() @@ -912,7 +1126,7 @@ def run(self): if self.backend_type == "xgmi": self._run_xgmi() else: - self._run_rdma() + self._run_distributed() def _run_xgmi(self): if self.xgmi_multiprocess: @@ -940,7 +1154,8 @@ def _run_xgmi(self): self.validate() self._run_benchmark_loop() - def _run_rdma(self): + def _run_distributed(self): + # Shared cross-node driver for RDMA and FABRIC (2-node gloo OOB). context_device_id = ( self.device.index if hasattr(self, "device") and self.device.index is not None @@ -1018,7 +1233,7 @@ def benchmark_engine(local_rank, node_rank, args): sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, sweep_max_size=args.sweep_max_size, - backend_type="rdma", + backend_type=args.backend, # "rdma" or "fabric" (both use this driver) host=args.host, port=0, node_rank=node_rank, @@ -1036,6 +1251,8 @@ def benchmark_engine(local_rank, node_rank, args): chunk_bytes=args.chunk_bytes, max_chunks=args.max_chunks, mem_type=args.mem_type, + num_streams=args.num_streams, + num_events=args.num_events, ) bench.print_config() bench.run() @@ -1098,7 +1315,8 @@ def benchmark_xgmi(args): bench.run() -def benchmark_rdma(args): +def benchmark_distributed(args): + # Cross-node driver shared by RDMA and FABRIC backends. if args.all: if args.sweep_start_size > args.sweep_max_size: raise ValueError( @@ -1129,7 +1347,7 @@ def benchmark(): if args.backend == "xgmi": benchmark_xgmi(args) else: - benchmark_rdma(args) + benchmark_distributed(args) if __name__ == "__main__": diff --git a/tests/python/utils.py b/tests/python/utils.py index 0c24d8791..4b7f39f5a 100644 --- a/tests/python/utils.py +++ b/tests/python/utils.py @@ -99,12 +99,22 @@ def __enter__(self): torch.cuda.set_device(self.device_id) device = torch.device("cuda", self.device_id) - dist.init_process_group( - backend=self.backend, - rank=self.rank, - world_size=self.world_size, - device_id=device, - ) + try: + dist.init_process_group( + backend=self.backend, + rank=self.rank, + world_size=self.world_size, + device_id=device, + ) + 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, + ) world_group = torch.distributed.group.WORLD assert world_group is not None From 01d4d3a21d7c4bcbe942245f0131ca571407e685 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Mon, 13 Jul 2026 08:03:08 +0000 Subject: [PATCH 2/6] fix(io/fabric): honor sub-region base offset in fabric remapping --- include/mori/io/common.hpp | 8 ++++++- src/io/fabric/backend_impl.cpp | 38 +++++++++++++++++++++++++++------- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/include/mori/io/common.hpp b/include/mori/io/common.hpp index e1f146a9a..0e1ded4d1 100644 --- a/include/mori/io/common.hpp +++ b/include/mori/io/common.hpp @@ -114,6 +114,12 @@ struct MemoryDesc { std::array fabricHandle{}; // 64B fabric token (empty when N/A) int vpodId{-1}; // UALink vPOD id of deviceId (-1 = N/A) std::string vpodPpodId; // UALink ppod_id UUID ("" = N/A) + // data may be a sub-allocation inside a larger fabric VMM allocation (e.g. a + // torch tensor carved out of a MemPool segment). fabricHandle exports the whole + // allocation; these locate `data` within it so the importer maps the allocation + // and offsets to the right address. + uint64_t fabricOffset{0}; // data - allocation base + uint64_t fabricAllocSize{0}; // full size of the exported allocation constexpr bool operator==(const MemoryDesc& rhs) const noexcept { return (engineKey == rhs.engineKey) && (id == rhs.id) && (deviceId == rhs.deviceId) && @@ -122,7 +128,7 @@ struct MemoryDesc { } MSGPACK_DEFINE(engineKey, id, deviceId, deviceBusId, data, size, loc, ipcHandle, numaNode, - fabricHandle, vpodId, vpodPpodId); + fabricHandle, vpodId, vpodPpodId, fabricOffset, fabricAllocSize); }; using TransferUniqueId = uint64_t; diff --git a/src/io/fabric/backend_impl.cpp b/src/io/fabric/backend_impl.cpp index df87569bf..5e5a71e88 100644 --- a/src/io/fabric/backend_impl.cpp +++ b/src/io/fabric/backend_impl.cpp @@ -617,8 +617,28 @@ void FabricBackend::RegisterMemory(MemoryDesc& desc) { static_assert(sizeof(fh) == kFabricHandleSize, "fabric handle size mismatch"); std::memcpy(desc.fabricHandle.data(), &fh, sizeof(fh)); - MORI_IO_TRACE("FABRIC: Registered memory id={}, addr={}, size={} (fabric-exportable)", desc.id, - desc.data, desc.size); + + // The registered pointer may be a sub-allocation inside a larger fabric VMM + // allocation (e.g. a torch tensor carved out of a MemPool segment). Record its + // offset + the full allocation size so the peer maps the whole allocation and + // offsets to the right address. + void* base = nullptr; + size_t rangeSize = 0; + hipError_t rangeErr = + hipMemGetAddressRange(&base, &rangeSize, reinterpret_cast(desc.data)); + if (rangeErr == hipSuccess && base != nullptr && rangeSize > 0) { + desc.fabricOffset = desc.data - reinterpret_cast(base); + desc.fabricAllocSize = rangeSize; + } else { + (void)hipGetLastError(); + desc.fabricOffset = 0; + desc.fabricAllocSize = desc.size; + } + + MORI_IO_TRACE( + "FABRIC: Registered memory id={}, addr={}, size={}, fabricOffset={}, allocSize={} " + "(fabric-exportable)", + desc.id, desc.data, desc.size, desc.fabricOffset, desc.fabricAllocSize); } void FabricBackend::DeregisterMemory(const MemoryDesc& desc) { @@ -654,7 +674,7 @@ void* FabricBackend::GetImportedAddress(const MemoryDesc& desc, int localDeviceI std::shared_lock rlock(importMutex); auto it = importedRegions.find(cacheKey); if (it != importedRegions.end() && it->second.mappedAddr != nullptr) { - return it->second.mappedAddr; + return static_cast(it->second.mappedAddr) + desc.fabricOffset; } } @@ -680,7 +700,9 @@ void* FabricBackend::GetImportedAddress(const MemoryDesc& desc, int localDeviceI return nullptr; } - size_t mapSize = AlignToFabricGranularity(localDeviceId, desc.size); + // Map the whole exported allocation, then offset to `data` within it. + size_t allocBytes = desc.fabricAllocSize != 0 ? desc.fabricAllocSize : desc.size; + size_t mapSize = AlignToFabricGranularity(localDeviceId, allocBytes); void* va = nullptr; err = hipMemAddressReserve(&va, mapSize, 0, 0, 0); if (err != hipSuccess) { @@ -718,12 +740,12 @@ void* FabricBackend::GetImportedAddress(const MemoryDesc& desc, int localDeviceI (void)hipMemUnmap(va, mapSize); (void)hipMemAddressFree(va, mapSize); (void)hipMemRelease(handle); - return winner; + return static_cast(winner) + desc.fabricOffset; } importedRegions[cacheKey] = {handle, va, mapSize}; - MORI_IO_TRACE("FABRIC: Imported fabric handle for id={} on device {}, mapped={}", desc.id, - localDeviceId, va); - return va; + MORI_IO_TRACE("FABRIC: Imported fabric handle for id={} on device {}, mapped={} offset={}", + desc.id, localDeviceId, va, desc.fabricOffset); + return static_cast(va) + desc.fabricOffset; } void FabricBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, From 8d60fd75e6befbc3360ffb16da2258c12a77f540 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Mon, 13 Jul 2026 09:44:45 +0000 Subject: [PATCH 3/6] add fabric allocator --- python/mori/io/__init__.py | 6 ++ python/mori/io/fabric_allocator.py | 127 +++++++++++++++++++++++++++++ src/io/CMakeLists.txt | 3 +- src/io/fabric/torch_allocator.cpp | 53 ++++++++++++ 4 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 python/mori/io/fabric_allocator.py create mode 100644 src/io/fabric/torch_allocator.cpp diff --git a/python/mori/io/__init__.py b/python/mori/io/__init__.py index c73067054..f5146d439 100644 --- a/python/mori/io/__init__.py +++ b/python/mori/io/__init__.py @@ -20,6 +20,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from .engine import * +from .fabric_allocator import ( + fabric_torch_allocator, + use_fabric_torch_allocator, + make_fabric_mem_pool, + fabric_mem_pool, +) from mori.cpp import ( IOEngineConfig, StatusCode, diff --git a/python/mori/io/fabric_allocator.py b/python/mori/io/fabric_allocator.py new file mode 100644 index 000000000..749871840 --- /dev/null +++ b/python/mori/io/fabric_allocator.py @@ -0,0 +1,127 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Torch pluggable allocator backed by fabric-exportable VMM memory. + +Plain torch tensors (hipMalloc / default caching allocator) cannot be exported +over a UALink super-node fabric, so the mori-io FabricBackend cannot register +them. This module wires torch's ``CUDAPluggableAllocator`` to mori-io's +``FabricMalloc`` / ``FabricFree`` (VMM allocations created with the fabric handle +type), so tensors allocated under it become fabric-exportable and can be +registered + transferred by the FabricBackend. + +Recommended usage (scope fabric memory to just the KV pool, keep the rest of +torch on the fast default allocator): + + import torch + from mori.io import fabric_mem_pool + + with fabric_mem_pool() as pool: # keep `pool` alive while tensors live + kv_cache = torch.empty(nbytes, dtype=torch.uint8, device="cuda:0") + mem = engine.register_torch_tensor(kv_cache) # now fabric-exportable + +For a dedicated process you may instead switch the process-global allocator +(must be called before any CUDA allocation): ``use_fabric_torch_allocator()``. +""" + +from __future__ import annotations + +import contextlib +import os + +_MALLOC_SYMBOL = "mori_io_fabric_malloc" +_FREE_SYMBOL = "mori_io_fabric_free" + +_allocator = None + + +def _mori_io_lib_path() -> str: + """Locate libmori_io.so shipped next to the mori package.""" + import mori + + pkg_dir = os.path.dirname(os.path.abspath(mori.__file__)) + candidates = [ + os.path.join(pkg_dir, "libmori_io.so"), + os.path.join(pkg_dir, "..", "libmori_io.so"), + ] + for path in candidates: + if os.path.exists(path): + return os.path.abspath(path) + raise FileNotFoundError( + f"libmori_io.so not found next to the mori package (searched {candidates}). " + "The fabric torch allocator requires the compiled mori-io shared library." + ) + + +def fabric_torch_allocator(): + """Return (cached) a torch ``CUDAPluggableAllocator`` backed by fabric memory.""" + global _allocator + if _allocator is None: + import torch + + _allocator = torch.cuda.memory.CUDAPluggableAllocator( + _mori_io_lib_path(), _MALLOC_SYMBOL, _FREE_SYMBOL + ) + return _allocator + + +def _raw_allocator(): + """Underlying torch._C._cuda_CUDAAllocator handle (for MemPool).""" + alloc = fabric_torch_allocator() + if hasattr(alloc, "allocator"): + return alloc.allocator() + return alloc._allocator # older torch + + +def use_fabric_torch_allocator() -> None: + """Switch torch's *process-global* CUDA allocator to the fabric allocator. + + Must be called before any CUDA allocation. Prefer :func:`fabric_mem_pool` to + scope fabric memory to the KV pool only. + """ + import torch + + torch.cuda.memory.change_current_allocator(fabric_torch_allocator()) + + +def make_fabric_mem_pool(): + """Create a ``torch.cuda.MemPool`` backed by the fabric allocator (torch>=2.5). + + Keep the returned pool alive for as long as the tensors allocated from it are + in use (the pool owns the underlying fabric memory). + """ + import torch + + return torch.cuda.MemPool(_raw_allocator()) + + +@contextlib.contextmanager +def fabric_mem_pool(): + """Context manager: CUDA allocations inside the block use fabric-exportable + memory; allocations outside keep torch's default (fast) allocator. + + Yields the ``MemPool``; hold a reference to it (and the tensors) while in use. + """ + import torch + + pool = make_fabric_mem_pool() + with torch.cuda.use_mem_pool(pool): + yield pool diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index 16c3a350a..96a7dfe32 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -8,7 +8,8 @@ set(MORI_IO_SOURCES xgmi/backend_impl.cpp xgmi/hip_resource_pool.cpp fabric/backend_impl.cpp - fabric/vpod_topology.cpp) + fabric/vpod_topology.cpp + fabric/torch_allocator.cpp) set_source_files_properties(${MORI_IO_SOURCES} PROPERTIES LANGUAGE CXX) add_library(mori_io SHARED ${MORI_IO_SOURCES}) diff --git a/src/io/fabric/torch_allocator.cpp b/src/io/fabric/torch_allocator.cpp new file mode 100644 index 000000000..bb65f7b61 --- /dev/null +++ b/src/io/fabric/torch_allocator.cpp @@ -0,0 +1,53 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// C-ABI allocator shims matching PyTorch's CUDAPluggableAllocator hooks +// void* alloc(size_t size, int device, hipStream_t stream) +// void free (void* ptr, size_t size, int device, hipStream_t stream) +// +// They back torch tensors with fabric-exportable VMM memory (see FabricMalloc), +// so a tensor allocated under this allocator (e.g. via torch.cuda.MemPool) can be +// registered and transferred by the mori-io FabricBackend across a UALink +// super-node. Plain torch/hipMalloc tensors are NOT fabric-exportable, so this +// allocator is the bridge that lets frameworks (e.g. SGLang) route their KV pool +// over the fabric backend. +// +// The stream argument is unused: VMM allocations are not stream-ordered and, once +// hipMemSetAccess grants the owning device read/write, are usable on any stream. + +#include + +#include + +#include "src/io/fabric/backend_impl.hpp" + +extern "C" { + +void* mori_io_fabric_malloc(size_t size, int device, hipStream_t /*stream*/) { + return mori::io::FabricMalloc(size, device); +} + +void mori_io_fabric_free(void* ptr, size_t /*size*/, int /*device*/, hipStream_t /*stream*/) { + mori::io::FabricFree(ptr); +} + +} // extern "C" From 19208fb8ae4c1c3a1f1e79d9201233df469c9feb Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Tue, 14 Jul 2026 07:23:10 +0000 Subject: [PATCH 4/6] fix(io/fabric): gate fabric device attribute by HIP version for older ROCm --- src/io/fabric/vpod_topology.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/io/fabric/vpod_topology.cpp b/src/io/fabric/vpod_topology.cpp index fe3aa2bea..4f93671b3 100644 --- a/src/io/fabric/vpod_topology.cpp +++ b/src/io/fabric/vpod_topology.cpp @@ -62,6 +62,7 @@ std::string DeviceBdfLower(int hipDevice) { } // namespace bool DeviceSupportsFabric(int hipDevice) { +#if HIP_VERSION >= 71400000 int vmm = 0; int fabric = 0; hipError_t err = @@ -76,6 +77,10 @@ bool DeviceSupportsFabric(int hipDevice) { return false; } return vmm != 0 && fabric != 0; +#else + (void)hipDevice; + return false; +#endif } VpodKey ReadVpodKey(int hipDevice) { From b7a9a43e63bb74d5390ddc4d5ea180dd9dc92b14 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Tue, 14 Jul 2026 09:03:05 +0000 Subject: [PATCH 5/6] fix --- python/mori/io/engine.py | 12 ++++++++++-- src/io/fabric/backend_impl.cpp | 14 +++++++++++--- src/io/fabric/vpod_topology.cpp | 5 ++++- tests/python/utils.py | 10 ++++++---- 4 files changed, 31 insertions(+), 10 deletions(-) diff --git a/python/mori/io/engine.py b/python/mori/io/engine.py index 4f50db910..e6cc76cdf 100644 --- a/python/mori/io/engine.py +++ b/python/mori/io/engine.py @@ -22,6 +22,7 @@ from mori import cpp as mori_cpp import torch import ctypes +import warnings _PyCapsule_New = ctypes.pythonapi.PyCapsule_New _PyCapsule_New.restype = ctypes.py_object @@ -121,8 +122,15 @@ def _load_fabric_copy_kernel(self): hsaco_path = ensure_fabric_copy_kernel() self._engine.LoadFabricCopyModule(hsaco_path) - except Exception: - pass + except Exception as e: + # Without this kernel, large FABRIC transfers fail (no memcpy fallback). + warnings.warn( + f"Failed to load fabric_copy kernel ({e}); FABRIC transfers " + "have no accelerated copy path and large transfers may fail. " + "Check that the fabric_copy kernel can be JIT-compiled.", + RuntimeWarning, + stacklevel=2, + ) def remove_backend(self, type: mori_cpp.BackendType): return self._engine.RemoveBackend(type) diff --git a/src/io/fabric/backend_impl.cpp b/src/io/fabric/backend_impl.cpp index 5e5a71e88..15278960b 100644 --- a/src/io/fabric/backend_impl.cpp +++ b/src/io/fabric/backend_impl.cpp @@ -276,9 +276,17 @@ hipError_t FabricBackendSession::LaunchCopy(void* dst, const void* src, size_t s 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). + // hipMemcpyAsync livelocks on imported fabric pointers for large payloads + // (see fabric_copy.hip); fail fast instead of hanging, only allow it small. + constexpr size_t kFabricMemcpyFallbackMax = 4ull * 1024 * 1024; + if (size >= kFabricMemcpyFallbackMax) { + MORI_IO_WARN( + "FABRIC: fabric_copy kernel not loaded; refusing hipMemcpyAsync " + "fallback for {} byte transfer (>= {} B may livelock on imported " + "fabric pointers). Load the fabric_copy module first.", + size, kFabricMemcpyFallbackMax); + return hipErrorNotSupported; + } return hipMemcpyAsync(dst, src, size, hipMemcpyDefault, stream); } diff --git a/src/io/fabric/vpod_topology.cpp b/src/io/fabric/vpod_topology.cpp index 4f93671b3..f77bb5a3f 100644 --- a/src/io/fabric/vpod_topology.cpp +++ b/src/io/fabric/vpod_topology.cpp @@ -62,7 +62,10 @@ std::string DeviceBdfLower(int hipDevice) { } // namespace bool DeviceSupportsFabric(int hipDevice) { -#if HIP_VERSION >= 71400000 + // Guard hipDeviceAttributeHandleTypeFabricSupported: it is absent on older + // ROCm (e.g. 7.2.4), and HIP_FABRIC_API is not defined on 7.14 where the enum + // does exist. Older ROCm reports fabric unsupported and falls back to XGMI/RDMA. +#if defined(HIP_FABRIC_API) || (HIP_VERSION >= 71400000) int vmm = 0; int fabric = 0; hipError_t err = diff --git a/tests/python/utils.py b/tests/python/utils.py index 4b7f39f5a..51aaae524 100644 --- a/tests/python/utils.py +++ b/tests/python/utils.py @@ -106,10 +106,12 @@ def __enter__(self): world_size=self.world_size, device_id=device, ) - 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. + except (ValueError, RuntimeError) as e: + # Some torch/ROCm builds reject device_id (report 0 accelerators) + # though HIP devices exist; retry without it only for that case. + msg = str(e).lower() + if "device_id" not in msg and "accelerator" not in msg: + raise dist.init_process_group( backend=self.backend, rank=self.rank, From 7688c238324283e4316c565786f38355e614dcb4 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Tue, 14 Jul 2026 10:33:06 +0000 Subject: [PATCH 6/6] fix(io): include memory ids in route cache key 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. --- include/mori/io/engine.hpp | 10 +++++++++- src/io/engine.cpp | 3 ++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/include/mori/io/engine.hpp b/include/mori/io/engine.hpp index 6fedfe5ad..5155edec5 100644 --- a/include/mori/io/engine.hpp +++ b/include/mori/io/engine.hpp @@ -118,11 +118,17 @@ class IOEngine { MemoryLocationType remoteLoc; int localDeviceId; int remoteDeviceId; + // Per-memory ids: routing can differ for the same engine/device pair when a + // buffer is fabric-exportable (FABRIC) vs not (RDMA), and RdmaBackend + // CanHandle is memory-agnostic, so the key must distinguish memories. + MemoryUniqueId localId; + MemoryUniqueId remoteId; bool operator==(const RouteCacheKey& rhs) const noexcept { return remoteEngineKey == rhs.remoteEngineKey && localLoc == rhs.localLoc && remoteLoc == rhs.remoteLoc && localDeviceId == rhs.localDeviceId && - remoteDeviceId == rhs.remoteDeviceId; + remoteDeviceId == rhs.remoteDeviceId && localId == rhs.localId && + remoteId == rhs.remoteId; } }; @@ -137,6 +143,8 @@ class IOEngine { hashCombine(seed, std::hash{}(static_cast(key.remoteLoc))); hashCombine(seed, std::hash{}(key.localDeviceId)); hashCombine(seed, std::hash{}(key.remoteDeviceId)); + hashCombine(seed, std::hash{}(key.localId)); + hashCombine(seed, std::hash{}(key.remoteId)); return seed; } }; diff --git a/src/io/engine.cpp b/src/io/engine.cpp index 8da75c083..a22a97b20 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -402,7 +402,8 @@ Backend* IOEngine::SelectBackend(const MemoryDesc& local, const MemoryDesc& remo return nullptr; } - RouteCacheKey routeKey{remote.engineKey, local.loc, remote.loc, local.deviceId, remote.deviceId}; + RouteCacheKey routeKey{remote.engineKey, local.loc, remote.loc, local.deviceId, + remote.deviceId, local.id, remote.id}; if (auto cachedType = QueryRouteCache(routeKey); cachedType.has_value()) { auto cachedBackend = backends.find(cachedType.value());