feat[io]: add fabric (UALink scale-up) transfer backend

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="<nodeA_ip>" --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 |

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                                         */
 /* ---------------------------------------------------------------------------------------------- */

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,14 +107,28 @@ struct MemoryDesc {
   MemoryLocationType loc;
   std::array<char, kIpcHandleSize> 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<char, kFabricHandleSize> 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) &&
            (deviceBusId == rhs.deviceBusId) && (data == rhs.data) && (size == rhs.size) &&
            (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, fabricOffset, fabricAllocSize);
 };
 
 using TransferUniqueId = uint64_t;

include/mori/io/engine.hpp

@@ -109,6 +109,7 @@ class IOEngine {
 
   std::optional<IOEngineSession> CreateSession(const MemoryDesc& local, const MemoryDesc& remote);
   void LoadScatterGatherModule(const std::string& hsacoPath);
+  void LoadFabricCopyModule(const std::string& hsacoPath);
 
  private:
   struct RouteCacheKey {
@@ -117,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;
     }
   };
 
@@ -136,6 +143,8 @@ class IOEngine {
       hashCombine(seed, std::hash<int>{}(static_cast<int>(key.remoteLoc)));
       hashCombine(seed, std::hash<int>{}(key.localDeviceId));
       hashCombine(seed, std::hash<int>{}(key.remoteDeviceId));
+      hashCombine(seed, std::hash<MemoryUniqueId>{}(key.localId));
+      hashCombine(seed, std::hash<MemoryUniqueId>{}(key.remoteId));
       return seed;
     }
   };

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<BackendType>;

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,
@@ -30,5 +36,8 @@
     PollCqMode,
     RdmaBackendConfig,
     XgmiBackendConfig,
+    FabricBackendConfig,
+    fabric_alloc,
+    fabric_free,
     set_log_level,
 )

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
@@ -95,11 +96,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 +116,22 @@ 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 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)
 

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

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

src/io/CMakeLists.txt

@@ -6,7 +6,10 @@ 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
+    fabric/torch_allocator.cpp)
 
 set_source_files_properties(${MORI_IO_SOURCES} PROPERTIES LANGUAGE CXX)
 add_library(mori_io SHARED ${MORI_IO_SOURCES})