allocator: register mori as a torch SymmetricMemory backend - #544
Draft
carlushuang wants to merge 30 commits into
Draft
allocator: register mori as a torch SymmetricMemory backend#544carlushuang wants to merge 30 commits into
carlushuang wants to merge 30 commits into
Conversation
Add a CUDAPluggableAllocator over ShmemMalloc/ShmemFree so tensors allocated
inside torch.cuda.use_mem_pool() come from the symmetric heap instead of the
caching allocator.
The motivation is the tensors an engine never allocates by hand. symm_mem-style
explicit allocation only covers buffers you construct yourself, but the ones
worth communicating are usually a KV cache or a GEMM output. Routing them
through a MemPool makes them symmetric at birth, so no staging copy is needed
before a transfer. This is the same integration shape inference engines already
consume from other transports -- SGLang selects a custom pool by type in
maybe_init_custom_mem_pool -- so mori can slot in without touching KV cache
code.
src/allocator/torch_allocator.cpp exports the two C symbols torch binds,
matching the signatures torch fixes:
void* mori_allocator_malloc(size_t size, int device, hipStream_t stream)
void mori_allocator_free (void* ptr, size_t, int device, hipStream_t)
plus mori_allocator_probe(), because allocation happens inside tensor
constructors where raising is awkward; callers can check first and fall back to
the default pool. The probe uses ShmemIsInitialized() rather than inferring
from ShmemMyPe(), which dereferences bootstrap state that does not exist before
ShmemInit and would take the process down.
python/mori/allocator exposes MoriAllocator.get_allocator(device), cached per
device, alongside is_available() and get_so_path().
Built behind BUILD_ALLOCATOR (default ON, requires BUILD_SHMEM), copied into
the wheel next to the other libmori_*.so.
Note the collective contract, documented in the module docstring: ShmemMalloc
is collective, so every rank must enter and leave the pool context in the same
order and allocate the same sizes. Asymmetric allocations inside the context
will hang or corrupt the heap.
Verified: cmake configure succeeds with BUILD_ALLOCATOR=ON and the translation
unit compiles against mori headers on gfx1250 / ROCm 7.14. Full link of the
target was not reached in that container -- mori_application fails earlier on a
missing /usr/lib64/libc.so, unrelated to this change -- so the runtime test in
tests/python/allocator has not been executed yet.
Add src/allocator/symm_backend.cpp, implementing SymmetricMemoryAllocator and
SymmetricMemory over the shmem heap and registering it with torch as "MORI"
via register_availability(). torch then drives everything:
import mori.allocator; mori.allocator.register_symm_backend()
symm_mem.set_backend("MORI")
t = symm_mem.empty(...) -> MoriSymmAllocator::alloc
hdl = symm_mem.rendezvous(t, group) -> MoriSymmAllocator::rendezvous
torch.ops.symm_mem.one_shot_all_reduce(t, "sum", group)
Rendezvous is nearly free here because shmem has already bootstrapped the peer
mapping: it is ShmemPtrP2p() per PE, with no handle exchange to perform.
Scale-up versus scale-out falls out of that same call. ShmemPtrP2p returns 0
for a PE reached over RDMA rather than P2P, so buffer_ptrs()[pe] is null
exactly for peers that are not load/store accessible -- the contract NCCL's
backend expresses through ncclGetLsaPointer. world_within_direct_access()
reports whether the whole group is directly mappable, and
get_rank_to_global_rank() is implemented, so a kernel can tell which regime it
is in. Note torch's own CUDA backend hardcodes world_within_direct_access() to
true and leaves the rank mapping NYI, so this path has no other consumer yet.
Verified on 2x gfx1250, ROCm 7.14, torch 2.11:
backend: MORI
empty(): ok ptr=0x7910f7e00400
rendezvous: world=2 rank=0
buffer_ptrs: ['0x7910f7e00400', '0x790fe9a00400']
get_buffer(0)[0] = 1.0 get_buffer(1)[0] = 2.0
one_shot_all_reduce: OK got=3.0 expect=3.0
so torch's collectives run on mori memory today, without waiting for the ROCm
branches of CUDASymmetricMemory to be fixed.
Gaps, all documented: put_signal/wait_signal raise, and barrier is host-side
(hipDeviceSynchronize plus ShmemBarrierAll) rather than spinning on the signal
pad, so ops needing device-side signalling will fail or serialise. No
multicast. The group must be the shmem world, which rendezvous checks and
reports rather than failing obscurely.
Built behind BUILD_TORCH_SYMM because it links libtorch, which mori otherwise
does not; setup.py enables it when torch is importable. Uses raw HIP rather
than c10::cuda/c10::hip, since torch ships both trees and picking the wrong one
pulls cuda_runtime_api.h into a ROCm build.
carlushuang
changed the title
Trim the block comments to what is not evident from the code -- why rendezvous is cheap, why buffer_ptrs can be null, why raw HIP instead of c10::cuda, and the collective contract on ShmemMalloc. The long-form explanation belongs in the PR and in the module docstring, not repeated at every call site. No functional change; both translation units recompiled to confirm.
The pre-commit hook pins clang-format v20.1.8 with mori's .clang-format (Google, ColumnLimit 100); my formatting assumed a narrower column. No functional change.
Rebuild the backend on plain HIP VMM instead of mori's shmem allocator, and probe the shareable handle type per device so gfx9 works. Dropping shmem is deliberate, not incidental. Both shmem and cco own a symmetric heap whose peer offsets stay aligned only while every rank allocates AND frees in the same order. torch cannot hold that invariant: tensors are freed by Python GC, whose order is not synchronised across ranks, so one rank collecting a tensor an iteration earlier than another silently desynchronises the heaps. Here each allocation is an independent VMM allocation with its own rendezvous, and free() is purely local, so divergent free order costs nothing. It also means the target links neither mori_shmem nor mori_application -- it builds with BUILD_SHMEM=OFF. Handle type is probed per device with a granularity-sized alloc/export, since the capability attribute enum is not stable across HIP releases. gfx9 has no fabric support at all -- hipMemCreate itself returns "operation not supported" on both gfx950 and gfx942, a different failure from gfx1250 on old ROCm where create succeeded and only export failed -- so those fall back to POSIX fd. mori's existing hipMemHandleTypeFabricCompat shim covers ROCm 7.1, which does not even declare the fabric enum and spells the prop field requestedHandleType. fds cannot ride the torch Store the way a 64-byte fabric handle can, so the fd path adds a small SCM_RIGHTS unix-socket exchange. torch has an IpcChannel for this but does not export it from any libtorch .so, so this is a local equivalent. Peers are still mapped into one flat span, so buffer_ptrs[r] == flat_base + r*stride on every path; flat_layout() exposes the pair for kernels. Verified on three architectures, 2 ranks each, all reporting uniform base+r*stride and a correct one_shot_all_reduce: gfx1250 / ROCm 7.14 handle_type=fabric stride 2MB gfx950 / ROCm 7.2.4 handle_type=posix_fd gfx942 / ROCm 7.2.4 handle_type=posix_fd Two shutdown bugs found and fixed along the way: destructors now bail on is_finalizing() like torch's own backend does, and the allocator singleton is intentionally immortal, because register_availability() parks a reference in a registry owned by libtorch that outlives this extension's statics. Without the latter every run segfaulted at interpreter exit after the test had passed. Also drops the CUDAPluggableAllocator/MemPool piece added earlier: it was not what was asked for, and it depended on ShmemMalloc which this no longer uses.
Add examples/torch_symm_all2all: a one-shot all-to-all kernel written against the flat symmetric window, plus its own setup.py so the extension is built in the example directory rather than by mori's CMake. Pure HIP -- no shmem, no cco. All the kernel needs from the backend is (flat_base, stride), after which peers are addressed arithmetically: recv slot of rank p, chunk from rank r == flat_base + p*stride + r*chunk_bytes so it takes two pointers and three integers instead of an N-entry pointer array, and the destination rank can be computed at run time. Push semantics: each rank writes into every peer's receive window, so there are no remote reads and no per-peer handshake. Measured on 8x MI355X (gfx950, ROCm 7.2.4, 256 KiB per peer), aggregate over the (world-1) chunks that actually leave the device: 2 ranks 20.2 GB/s 4 ranks 153.4 GB/s 8 ranks 794.9 GB/s with correctness checked as an NxN chunk permutation each time. gfx9 has no fabric, so this runs the POSIX-fd path; the kernel neither knows nor cares. Two backend fixes fell out of writing it. The fd exchange assumed one datagram round per owner, but datagrams from different owners can arrive in any order, so a rank could pair a peer's fd with the wrong rank. Each message now carries its sender and the receiver takes world-1 fds in whatever order they land. Latent at 2 ranks, reachable at 4+. Releasing a rendezvous'd window segfaults at world_size >= 4 (2 is fine) somewhere in the unmap/release path. Root cause not yet found: it is not is_finalizing(), not interpreter shutdown (it faults on an explicit del), and not fixed by quiescing with a barrier first. Until it is understood, teardown is off by default and the mappings are leaked -- symmetric buffers are few and long-lived, so that is much less harmful than crashing. MORI_SYMM_TEARDOWN=1 re-enables it for debugging. Recorded in the module docstring and the PR.
…ayload The kernel launched one block per destination rank, so a 4-rank job used 4 CUs and could not keep enough writes in flight to cover interconnect latency. Split each chunk over blocks_per_peer blocks instead, sized from the CU count. At 4 MiB per peer this is worth ~2.5x on gfx950 (745 -> 1858 GB/s at 8 ranks) and ~95x on gfx1250 (15.8 -> 1499 GB/s at 4 ranks). README records gfx950, gfx1250 and gfx942 results; gfx1250 runs the fabric path, the other two fall back to POSIX fd.
… fabric-bound ubench/06 puts that box's XGMI at 48.4 GB/s per link and 2637 GB/s aggregate via hipMemcpyPeer, so the 185 GB/s all-to-all is ~7% of the interconnect. Writing the VMM window there runs at ~50 GB/s vs ~880 GB/s for a hipMalloc tensor, and own-slot is as slow as peer, so nothing crossing XGMI is implicated.
The previous note blamed VMM memory itself. A standalone HIP repro shows VMM matches hipMalloc (2670 vs 2536 GB/s) and that flat multi-slot mapping and self aliasing are free. The actual trigger is hipMemSetAccess: granting one peer GPU access drops the owner's own bandwidth from 2671 to 55.8 GB/s, reads and writes alike, with no further cost for more peers. gfx950 is unaffected (6541 GB/s with grants). Making the memory shareable is what costs gfx942, not the flat window.
… architectural hipMalloc memory with hipDeviceEnablePeerAccess for all 7 peers keeps full local bandwidth on the same box (2668.5 -> 2665.2 GB/s), while VMM plus a single hipMemSetAccess grant drops to 55.8. So peer visibility is not what costs the bandwidth and the earlier 'no backend can avoid this' claim was wrong; the hipMemCreate path specifically is what loses it.
…collapse Tracing amdgpu:amdgpu_vm_set_ptes across 0..3 peer grants shows MTYPE_RW constant at +492 MiB while each peer adds its own +256 MiB of NC and +256 MiB of UC. The owner's entries are never re-typed, so the uncached-remap explanation is wrong.
Matches the backend, which already uses world_size, and makes it explicit at the kernel signature that the argument is an index rather than a count or a device ordinal. Comment prose still says "rank".
…comments Granting one peer re-maps the buffer in the OWNER's page tables as AMDGPU_PTE_SYSTEM|MTYPE_UC, so the owner's access to its own HBM leaves the chip at PCIe speed. Confirmed by tracing amdgpu_vm_set_ptes against a size-fingerprinted buffer: a SYSTEM|UC group matching the allocation appears only once a peer is granted, and hipMemGetInfo is flat, so it is the mapping and not migration. Also drops the earlier claim that this is not a kernel-config issue. That rested on hipMalloc+peers being unaffected, but the two paths do not share a kernel interface -- hipMemSetAccess reaches libdrm amdgpu_bo_va_op while ordinary allocations use hsaKmtMapMemoryToGPUNodes -- so the missing CONFIG_PCI_P2PDMA and CONFIG_DMABUF_MOVE_NOTIFY on that host are back to being the leading suspect.
The backend was a CMake target that called find_package(Torch) and linked
${TORCH_LIBRARIES} by hand. That reimplemented torch's build contract only
partially: TORCH_CXX_FLAGS was never applied so _GLIBCXX_USE_CXX11_ABI came from
the compiler default, pybind11 came from the standalone package rather than
torch's bundled copy, and the output had no Python ABI tag, so a 3.11 interpreter
would match a 3.12-built mori_torch_symm.so and fail on symbols.
Declare it as a CppExtension instead. torch's build_ext derives all of that from
the installed torch and emits mori_torch_symm.cpython-312-x86_64-linux-gnu.so into
the package. CMakeBuild now inherits torch's BuildExtension when torch is
importable and routes tagged extensions through build_extensions(); the lazy
compiler setup moved into _ensure_compiler() so both paths share it.
BUILD_TORCH_SYMM, src/allocator/CMakeLists.txt and the .so copy step are gone, so
find_package(Torch) no longer appears anywhere in mori's CMake.
mori.allocator now imports torch before the extension: it links libtorch and
nothing on its RUNPATH resolves that, so importing mori.allocator without torch
already loaded used to fail with a misleading 'extension not found'.
Verified on 8xMI355X: full setup.py build, 2-rank backend test (flat layout, peer
reads, one_shot_all_reduce) and the 8-rank all2all example at 1828 GB/s.
The previous commit made CMakeBuild inherit torch's BuildExtension, which put every extension -- the CMake pseudo-target and the Cython cco module included -- through torch's compiler patches. Hand the tagged extension to its own BuildExtension instance instead, so nothing else in the build sees torch. CMakeBuild is a plain build_ext again and build_extension() is byte-identical to before; the net change to setup.py is the guarded CppExtension import, one filter in run(), _build_with_torch(), and _torch_symm_extension().
torch already publishes the peer pointers, so (flat_base, stride) is just ptrs[0] and ptrs[1]-ptrs[0]; the C++ FlatLayout was re-deriving from the same members buffer_ptrs is generated from. Verified identical on a 4-rank window. The Python version takes the handle rather than the tensor, so it needs no group plumbing, and it adds the check the C++ one lacked: that the stride really is uniform. That is the part torch's API does not promise -- other backends hand back scattered per-rank pointers, and aiming a kernel at base + rank*stride on one of those would corrupt memory silently rather than raise. Also drops the torch/csrc/utils/pybind.h include, which existed only for the at::Tensor caster, so the extension no longer binds any torch type. README now documents the raw two-line derivation so callers know the helper is optional. Verified on 8xMI355X: 4-rank backend test and the 8-rank all2all at 1859 GB/s.
barrier/put_signal/wait_signal all raise, so the 9216-byte pad torch defines was being reserved in every window and never used. Physical backing is 2 MiB-paged, so appending it to a page-aligned request cost a whole extra page: a 64 MiB window measured 66 MiB. Measured across sizes, the pad was free only when the request had slack in its last page. MORI_SYMM_SIGNAL_PAD (default off) now guards it. With it off the pad is not reserved, the device-side pad array is not allocated or copied, and every signal entry point -- including get_signal_pad_ptrs()/_dev() -- reports that support is compiled out and names the flag. A 64 MiB window is now exactly 64 MiB. torch's one_shot_all_reduce does not touch the pad, verified at 4 ranks with it compiled out. Building with MORI_SYMM_SIGNAL_PAD=ON restores the old 66 MiB and also passes.
…d compare The example only demonstrated flat_base + peer*stride. Build the same kernel both ways -- the flat form, and the N-entry peer pointer array that hdl.buffer_ptrs_dev() gives and that torch's own backends force, since map_block() reserves each peer's VA independently. --addressing flat|ptrs|both, both correctness-checked. Measured: no difference worth having. Under 1% at 4 MiB and under 4% at 256 KiB, in neither direction consistently -- gfx950 marginally favours the array, gfx1250 the flat form. The flat window earns its place on ergonomics, not throughput, and the README now says so rather than implying a speed argument. Also fixes a race the second mode exposed: recv.zero_() is local, but peers write into that window, so a rank zeroing late wiped a write a peer had already landed. Passed in isolation and failed only when two modes ran in sequence, more often at larger chunks. Now barriers between the clear and the pushes. gfx950 8 ranks and gfx1250 4 ranks, both modes correct on both.
register_symm_backend() had to be called before symm_mem.set_backend("MORI")
would work, which is boilerplate every caller had to know about. Importing
mori.allocator now registers it, so the only call left is torch's own
set_backend().
Registration is best effort at import: if torch or the extension is missing there
is nothing to register, and the first real call still raises with a useful message
rather than making 'import mori.allocator' itself fail. register_symm_backend()
stays exported and idempotent for callers that want the error eagerly.
Deliberately does not auto-activate. register_availability() only makes "MORI"
selectable; making it the active allocator for the device type is what
set_backend() does, and that should stay the caller's choice since other backends
(RCCL's NCCL one) may also be registered.
Verified on 8xMI355X: bare import then set_backend, 4-rank backend test, 8-rank
all2all in both addressing modes.
Peers are published only as buffer_ptrs / buffer_ptrs_dev, which is what torch's SymmetricMemory model provides and what every other backend does. The window is still one flat span internally, but flat_layout() is gone from the public API: exposing base + rank*stride belongs with mori's cco window, whose ccoWindowDevice already defines that layout (winBase, 4 GiB-quantised stride4G, LSA-rank indexing), rather than as a second, parallel scheme here. See ROCm#557 for the staged plan. The example keeps both addressing forms -- it derives (base, stride) from buffer_ptrs itself -- since that comparison is the evidence for the next stage. Re-measured on 8x gfx950: 1814.5 GB/s array vs 1803.7 flat at 4 MiB, 1639.9 vs 1623.2 at 256 KiB. Also adds signal_pad_supported(), which the tests need: torch's own symm_mem collectives synchronise through the pad, so one_shot_all_reduce raises when it is compiled out. The test now asserts the collective works when the pad is built in and reports unsupported when it is not; both builds pass.
Same one-shot push as the HIP kernel, written in Triton against the same hdl.buffer_ptrs_dev. It is the idiom torch's own symmetric-memory Triton kernels use -- pass the array's address as a plain int, cast it to a pointer in the kernel, load the peer out of it -- so the backend needs to offer nothing beyond what stage 1 already publishes, and this path needs no extension build at all. all2all.py grows --kernel hip|triton|both alongside --addressing; Triton appears under the pointer-array form only, since the flat form is a HIP-side comparison. Timing and correctness now iterate over (implementation, addressing) variants. The kernel is declared do_not_specialize on its int arguments: Triton turns an int argument equal to 1 into a constexpr, so rank 1 alone failed to compile. Measured on 8x gfx950, aggregate over the chunks that leave the device: Triton is within a couple of percent of HIP at 4 MiB (1818.3 vs 1843.3, and 1852.2 vs 1832.2 on a repeat) and at 1 MiB (1769.5 vs 1784.0), and 18% behind at 256 KiB (1256.6 vs 1536.3) -- a near-constant ~1.3 us of host-side launch cost, not codegen: both launch 512 blocks storing 1024 int32 each. All variants correctness-checked at 8 ranks, including a masked tail chunk.
…ncher Ran the example on 4x gfx1250 (torch 2.11 / Triton 3.8, fabric handles): all three variants correct, and the Triton kernel needed no change from the wave64 gfx950 version -- num_warps=8 to get 256 lanes there moves nothing (32.2 vs 33.2 us at 4 MiB). It is 2.3x slower at 256 KiB, and that is entirely host-side. Launches are async, so an iteration costs max(kernel, launch); timing the launch call alone gives 2.90 us (HIP pybind) vs 9.01 us (Triton 3.7) on gfx950 and 2.44 vs 12.83 us (Triton 3.8) on gfx1250. The gfx1250 Triton row is pinned near 14 us at both 256 KiB and 1 MiB despite 4x the data -- it is waiting on JITFunction.run. That also corrects the gfx950 claim from the previous commit: repeats put Triton within noise of HIP at every size there (its 9 us launcher sits just under the 9-10 us kernel), so the earlier 18%-at-256-KiB figure was an outlier, not a trend. Tables now report us/iter with a range over three runs, which is steadier than the summed bandwidth. Also caches the per-launch device-property lookup, which is 0.9 us of a budget that turns out to be measured in single-digit microseconds.
all2all_kernel.cpython-312-x86_64-linux-gnu.so is output of `setup.py build_ext --inplace`, swept in by a `git add -A` in the previous commit. Removed, and examples/**/*.so added to .gitignore so the in-place build cannot do it again.
The gfx1250 figures were taken while another tenant held all 4 GPUs. On an idle box the same runs are 12% faster at 4 MiB and 44% faster at 256 KiB, and reproducible to within 0.5 us across repeats, so the numbers in every gfx1250 column are replaced and the contention caveats dropped. Two things sharpen as a result. The Triton launcher measures 10.84-10.94 us rather than 12.83, and the Triton row now sits at 12.0-12.6 us for both 256 KiB and 1 MiB -- 4x the data, same time, right at the launcher floor. A trailing synchronize() adds 0.03 us to the Triton launch loop against 1.7-1.9 us for HIP, which is the same finding from the other side: the host never gets far enough ahead to build a queue. The addressing comparison also needs a less flat claim than "no measurable difference". On gfx1250 the flat form is ahead by a near-constant 0.3-0.6 us at every size, across six runs and whether the modes run together or alone -- a per-block cost, the array's dependent load ahead of the first store, not a per-byte one. It is still within 0.2% on gfx950, with the array ahead as often as behind.
The example was a harness plus a .hip file plus a setup.py, and it ran
four ways through two flags. It is now two scripts, all2all_hip.py and
all2all_triton.py, each of which runs the whole example on its own:
torchrun --nnodes=1 --nproc_per_node=8 all2all_hip.py --chunk-kib 256
No build step either way. The HIP kernel is a string that load_inline
JIT-builds on first run -- concurrent torchrun ranks share one build,
since torch takes a file lock -- so setup.py and all2all_kernel.hip are
gone. A bare `python3 all2all_hip.py` also works now, as a 1-rank smoke
test, by defaulting the rendezvous environment.
The second HIP kernel taking base + peer*stride goes with them, along
with flat_view(). Exposing a flat window is ROCm#557's job and should arrive
as cco's ccoWindowDevice, not as a third scheme an example invented; the
comparison it existed to make has been made, and the README keeps the
result in a sentence rather than a section.
Comments are cut back to the things that are not evident from the code:
grid shape, the int32 overflow, Triton's constexpr folding of 1.
Re-measured on 8x gfx950 with the code as it stands. Small messages come
out a little faster than the old harness reported (256 KiB, 8 ranks:
1522-1829 vs 1459 GB/s), which is what dropping a dict-and-lambda
indirection from a 8 us launch path looks like.
carlushuang
added a commit
to carlushuang/mori
that referenced
this pull request
Aug 17, 2026
Standalone page for ROCm#544: why the backend exists, how alloc and rendezvous work, and how one probe covers fabric and fd parts. Added alongside the generated site rather than into it.
Picks up the ROCm 7.14 container build fixes (ROCm#541), the JIT v2 / dispatch_combine_v2 targets, and gfx1250 in _supported_arch_list. No conflicts: upstream's setup.py and CMakeLists.txt edits are in different regions than the torch symm backend's.
Left over from an earlier revision; the torch symm backend is built by setup.py's CppExtension, so it touches no CMake at all.
torch's cpp_extension compiles through ninja from build_temp, so the relative -Iinclude resolved against the wrong directory and symm_backend.cpp could not find mori/utils/hip_compat.hpp. Only shows up when __file__ is relative (python setup.py ...); pip passes an absolute path, which is why the wheel build was unaffected.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📄 Design write-up: https://carlushuang.github.io/mori/torch-symm-backend.html — leads with the programming model (torch's native pointer array
peers[p] + offsetvs the LSA windowbase + peer*stride + offsetthat RCCL and mori's cco build underneath), then what torch's ROCm path does today, how alloc/rendezvous are built, and how one runtime probe covers fabric and fd parts. Everything below in more depth, with the diagrams.Registers mori with torch as the
"MORI"SymmetricMemory backend, so torch's own entry points drive it:This is stage 1 of #557: peers are published only the way torch's model does it, as the
buffer_ptrs/buffer_ptrs_devarray. Every rank does map into one flat VMM span, so those pointers happen to be evenly strided, but that layout is deliberately not part of the API here — exposingbase + rank*stridebelongs with mori's cco window, whoseccoWindowDevicealready defines it (winBase, 4 GiB-quantisedstride4G, LSA-rank indexing). Adding a second, parallel scheme in the allocator would just have to be unwound later. Stage 2 hands the allocation toccoWindowRegisterinstead and lets cco own the window; stage 3 is scale-out and SDMA, which need cco co-design.Verified at 2 ranks on three architectures: gfx1250 (ROCm 7.14/7.15,
fabrichandles), gfx950 (MI355X) and gfx942 (MI308X) (7.2.4,posix_fd). The all-to-all example runs at 8 ranks on gfx950 and 4 on gfx1250, in both HIP and Triton.Why not shmem or cco today
Both keep a symmetric heap whose peer offsets stay aligned only while every rank allocates and frees in the same order. torch cannot hold that invariant — tensors are freed by Python GC, whose order is not synchronised across ranks. Here each allocation is an independent VMM allocation with its own rendezvous and a purely local
free(). The target links neithermori_shmemnormori_application.That is an argument against the heap, not against cco.
ccoMemImportandccoWindowRegisteroverload C already alias an external HIP VMM allocation into cco's flat slot — the header's own example is "e.g. a torch.symm_mem buffer" — andccoGetPeerPtrcomputes peer VAs on the host for exactly this kind of pointer table. Wiring those up is stage 2.Build
Built by
torch.utils.cpp_extensionas aCppExtensioninsetup.py, not by CMake: it is the only target that links libtorch, and torch'sbuild_extderives_GLIBCXX_USE_CXX11_ABI, the pybind11 copy and the module's ABI tag from the installed torch.find_package(Torch)therefore appears nowhere in mori's CMake, and the torch build command is scoped to this one extension so nothing else in the build sees torch's compiler patches. See #549 for mori's packaging layer generally.gfx9 support
Handle type is probed per device with a granularity-sized alloc/export, because the capability enum is not stable across HIP releases. gfx9 has no fabric at all —
hipMemCreateitself fails on gfx950 and gfx942, unlike gfx1250 on older ROCm where create succeeded and only export failed — so gfx9 falls back to POSIX fd. fds cannot ride the torch Store, so that path adds a small SCM_RIGHTS exchange; torch has anIpcChannelfor this but does not export it from any libtorch.so.Example: all-to-all, in HIP and in Triton
examples/torch_symm_all2all/— a one-shot all-to-all over the window, no shmem and no cco, as two self-contained scripts with no build step:all2all_hip.pykeeps the HIP kernel as a string and JIT-builds it throughtorch.utils.cpp_extension.load_inlineon first run (concurrent ranks share one build; torch holds the lock).all2all_triton.pydoes the same push over the same array —peers = peer_ptrs.to(tl.pointer_type(tl.uint64)), thentl.load(peers + peer).to(tl.pointer_type(tl.int32)), the idiom torch's own symmetric-memory Triton kernels use — and needs no C++ at all. Both take onlyhdl.buffer_ptrs_dev, so stage 1's contract is the entire interface either kernel needs.The example used to carry a second HIP kernel taking
base + peer*stride, to measure the two addressing forms against each other. That is now removed along withflat_view(), since the flat window should arrive as cco's, not as a third scheme invented in an example. What it measured, before removal: within 0.2% on 8×gfx950 with the array ahead as often as behind (1909.4 vs 1912.0 GB/s at 4 MiB; 1814.5 vs 1803.7 on a second box), and on an idle 4×gfx1250 the flat form consistently ahead by a near-constant 0.3-0.6 us — 2% at 4 MiB, 3-7% at 256 KiB. Constant in absolute terms is a per-block cost, not a per-byte one: the array's dependent load sits on the critical path before the first store issues. So stage 1 gives up a fixed sliver of launch-to-first-store latency and nothing per byte, and stage 2's flat window is worth having mostly for kernel ergonomics.Aggregate over the
(world_size-1)chunks that leave the device, 4 MiB per peer: 1746-1874 GB/s on 8×gfx950, 1705.1 GB/s on 4×gfx1250, 184.8 GB/s on 8×gfx942.HIP vs Triton, us/iter over repeats. 8×gfx950: 128.9-140.6 (HIP) vs 118.6-139.7 (Triton) at 4 MiB, 29.7-30.3 vs 30.1-33.4 at 1 MiB, and 7.9-9.6 vs 9.1-10.0 at 256 KiB — close enough that either is ahead depending on the run. 4×gfx1250 (fabric handles, idle box): 29.6-30.5 vs 32.2-32.5 at 4 MiB, 10.0-10.1 vs 11.9-12.6 at 1 MiB, and 5.9-6.2 vs 12.0-12.1 at 256 KiB.
That last gap is host-side, not device-side. Launches are async, so an iteration costs
max(kernel, launch), and timing the launch call alone gives 2.90 us (HIP pybind) vs 9.01 us (Triton 3.7) on gfx950, and 2.44-3.03 vs 10.84-10.94 us (Triton 3.8) on gfx1250. The gfx1250 Triton row is pinned near 12 us at both 256 KiB and 1 MiB despite 4x the data — it is waiting onJITFunction.run, not moving bytes, and a trailingsynchronize()adds 0.03 us to it against 1.7-1.9 us for HIP, because no queue ever builds up. Same kernel, no source change between wave64 and wave32, andnum_warps=8on gfx1250 moves nothing.Every serious library builds the flat window underneath regardless — NCCL's
ncclGetLsaPointerisadd4G(lsaFlatBase, peer*stride4G) + offset, mori'sccoWindowDeviceiswinBase + ((uint64_t)peerLsaRank * stride4G << 32) + offset, and NVSHMEM reservesnpes * heap_sizein onecuMemAddressReserveand then still exposes an array, because its slots are rotated frommype+1sopeis not the slot index.Grid shape mattered far more than addressing: one block per destination rank left all but
world_sizeCUs idle and managed 15.8 GB/s on gfx1250, so each chunk is sliced acrossblocks_per_peerblocks.The gfx942 column is a driver limitation, not a fabric or allocator one. Granting one peer via
hipMemSetAccessre-maps the buffer, in the owner's own page tables, asAMDGPU_PTE_SYSTEM | MTYPE_UC, dropping the owner's access to its own HBM from 2671.5 to 55.8 GB/s — PCIe speed.hipMalloc+hipDeviceEnablePeerAccessis unaffected on the same box, because that path useshsaKmtMapMemoryToGPUNodes(KFD) whilehipMemSetAccessreaches libdrmamdgpu_bo_va_op(DRM). That host also lacksCONFIG_PCI_P2PDMAandCONFIG_DMABUF_MOVE_NOTIFY; details in the example's README.Notes
The extension binds no torch types at all; its pybind surface is
register_backend/shutdown/handle_type/backend_name/signal_pad_supported.barrier/put_signal/wait_signalare unimplemented and raise, so the signal pad is not reserved either: appending torch's 9216-byte pad to a page-aligned window costs a whole extra 2 MiB page, physical backing being 2 MiB-paged.MORI_SYMM_SIGNAL_PAD=ONreserves it, andmori.allocator.signal_pad_supported()reports which build you have. torch's ownsymm_memcollectives synchronise through that pad rather than through this backend'sbarrier(), so they need the flag:one_shot_all_reduceraises a message naming it when the pad is compiled out, and returns the right answer when it is in. The test covers both builds; both pass at 2 ranks on gfx950.MemPool works:
symm_mem.get_mem_pool()+use_mem_pool()allocations rendezvous and drive the kernel correctly, verified at 4 ranks.Gaps
world_size >= 4(2 ranks are fine) in the unmap/release path. Notis_finalizing(), not interpreter shutdown (it faults on an explicitdel), not fixed by quiescing with a barrier. Mappings are leaked meanwhile — symmetric buffers are few and long-lived.MORI_SYMM_TEARDOWN=1re-enables. Main thing blocking the draft. Stage 2 may dissolve it, since cco would own unmap ordering.has_multicast_support()is false andmultimem_*is out.world_within_direct_access()returns true unconditionally — correct here (no RDMA path) but means no multi-node group. Stage 3's scale-out is where that becomes false for network peers, the way NVSHMEM's backend flipsworld_within_cuda_p2p_whennvshmem_ptrreturns null.all_to_all_nddispatch on a hardcodedif backend == "NCCL", so third-party backends get the generic ops only.Questions
CppExtensioninsetup.pythe right home, or should this build some other way?ccoMemImport/ccoWindowRegisteroverload C the intended path for a torch-owned allocation, and shouldccoGetPeerPtrget a Cython binding to go with it? How shouldrendezvous(group_name)map ontoccoCommCreate's up-frontperRankVmmSize?barrier()be implemented over the signal pad with P2P stores now, or wait for cco's signal pool in stage 3?world_size >= 4? That is the one unresolved item.