A torch SymmetricMemory backend for mori

01Why do this at all

What torch's SymmetricMemory gives you

A symmetric memory window is a buffer that every rank in a process group allocates, and that every rank can read and write directly — a peer's slot is a plain pointer that a kernel dereferences, not a message you enqueue on a communicator. The programming model is one-sided: no matching send for a receive, no collective call in the fast path, just loads and stores plus your own synchronisation.

torch exposes it as three calls, and that is the whole surface an application sees:

import torch.distributed._symmetric_memory as symm_mem
import mori.allocator                     # importing registers the "MORI" backend

symm_mem.set_backend("MORI")
t    = symm_mem.empty(1024, dtype=torch.bfloat16, device=device)
hdl  = symm_mem.rendezvous(t, group_name)

peer = hdl.get_buffer(1, (1024,), torch.bfloat16)   # rank 1's slot, as a tensor
ptrs = hdl.buffer_ptrs_dev                          # or the device pointer array, for a kernel

That model is what fused, low-latency communication is being written against today: one-shot and two-shot all-reduce, MoE all-to-all and combine, KV-cache movement for disaggregated serving, and increasingly Triton kernels that take buffer_ptrs_dev and do the transfer inline with compute instead of around it.

The gap on ROCm

torch ships a CUDA backend for this. Reading it on main, the ROCm arm is a second-class path:

• c10/cuda/PeerToPeerAccess.cpp:164 — get_fabric_access() is wrapped in #if !defined(USE_ROCM) && CUDA_VERSION >= 12040 && ...; the #else is return false. On ROCm the capability query is a compile-time constant.
• torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu:342 — the allocator asks at::cuda::get_fabric_access() and chooses FABRIC_HANDLE or POSIX_FD. That block is #if !defined(USE_ROCM). The #elif defined(USE_ROCM) arm at line 369 hardcodes handle_type_ = POSIX_FD and prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor.
• Nothing in torch ever asks HIP for a fabric handle, even on parts that have them.

So on a gfx1250 box where the driver exports fabric handles happily, torch would still exchange file descriptors over a UNIX socket. Worse, the pieces you would need to do better are not reusable: torch's IpcChannel, the SCM_RIGHTS helper for the fd path, is declared at CUDASymmetricMemoryUtils.hpp:97 with no TORCH_API — it is not exported from any libtorch .so.

An AMD-side backend is therefore not a duplicate of torch's: it is the only place a per-device probe, fabric handles, and eventually mori's own window and scale-out transports can live.

Why not mori's existing allocators

mori already has two symmetric allocators — the shmem heap and cco. Neither can back torch as-is, for one specific reason:

This backend sidesteps the invariant instead of fighting it: every allocation is an independent HIP VMM allocation with its own rendezvous, and free() is purely local, so divergent free order is harmless. The extension links neither mori_shmem nor mori_application, and needs no mori bootstrap — torch's own process group is the only rendezvous channel.

That is an argument against the heap, not against cco. cco's ccoMemImport and ccoWindowRegister already alias an externally-owned HIP VMM allocation into cco's flat slot — the header's own example is "e.g. a torch.symm_mem buffer" — which is exactly the shape of the handoff for stage 2.

02The programming model: pointer array vs LSA window

Everything above is plumbing. The part that a kernel author actually feels — and the part that decides what stage 2 has to look like — is the addressing model: given a peer rank and a byte offset, how do I get a dereferenceable address? There are exactly two answers in use, and this PR implements the first one.

Model A — the pointer array, which is what torch's interface is

torch's SymmetricMemory C++ interface has exactly two accessors for peer memory: get_buffer_ptrs() and get_buffer_ptrs_dev(). There is no slot in it for a base and a stride. So every backend, whatever it does internally, ends up publishing an N-entry table:

// HIP — examples/torch_symm_all2all/all2all_hip.py
__global__ void All2AllPushPtrs(const uint4* src, void** __restrict__ peers, ...) {
  auto* dst = reinterpret_cast<uint4*>(static_cast<char*>(peers[peer])   // <-- table load
                                       + rank_id * chunk_bytes);
  ...
}

# Triton — the idiom torch's own symmetric-memory kernels use
peers = peer_ptrs.to(tl.pointer_type(tl.uint64))
dst   = tl.load(peers + peer).to(tl.pointer_type(tl.int32))          # <-- table load
dst  += rank_id.to(tl.int64) * chunk

The virtue of this model is that it demands nothing of the allocator. Peer slots can be at unrelated addresses, which is precisely what torch's own CUDA backend produces — it maps each peer's memory separately — and what an hipIpc-based backend would produce too. NVSHMEM's torch backend goes further and rotates the slots so that a peer's index in the array is not its PE id at all. The array is the lowest common denominator, and that is why torch chose it.

Model B — the LSA window: base pointer plus offset

The other model gives up mapping freedom in exchange for arithmetic addressing. Reserve one VA range covering every rank's slot at a uniform stride, and a peer address becomes a computation, with no table and no load:

// NCCL / RCCL device API — LSA = "load/store accessible"
void* p = ncclGetLsaPointer(win, offset, lsaPeer);   // NULL if peer is outside the LSA team
void* q = ncclGetPeerPointer(win, offset, peer);     // world-team index, same window
// implemented underneath as: add4G(lsaFlatBase, peer * stride4G) + offset

// mori cco — include/mori/cco/cco.hpp, struct ccoWindowDevice
//   peer_va = winBase + ((uint64_t)peerLsaRank * stride4G << 32) + offset
//   local   = winBase + ((uint64_t)lsaRank     * stride4G << 32) + offset
struct ccoWindowDevice { char* winBase; uint32_t stride4G; int lsaRank; ccoIbgdaWin ibgdaWin; };

Two details in that snippet carry most of the model. The stride is 4 GiB-quantised — stride4G is perRankSize >> 32, so the multiply is a 32-bit op and the shift is free — and the peer index is an LSA rank, not a world rank. LSA is a team: the set of peers this GPU can reach with ordinary loads and stores. Ranks outside it are not addressable this way at all — ncclGetPeerPointer returns NULL for them, and cco routes them through ibgdaWin instead. The addressing model and the transport boundary are the same concept.

Side by side

	A · pointer array	B · LSA window
Peer address	peers[p] + offset	base + p*stride + offset
Cost per block	one dependent global load before the first store can issue	a few ALU ops from values already in registers
Uniformity	a loaded value; scalarises only if the compiler can prove p uniform	scalar arithmetic off a uniform base
Allocator demand	none — any mapping, any addresses	one contiguous VA reservation, all slots, uniform stride
Size granularity	exact	quantised (4 GiB per rank in cco/NCCL)
Team semantics	none; one flat world plus a world_within_direct_access() boolean	first-class — LSA team, cross-node team, rail team; NULL outside
Scale-out story	outside the model	the complement of the LSA team, via the same window handle
Defined by	torch (buffer_ptrs / buffer_ptrs_dev)	RCCL (ncclGetLsaPointer), mori (ccoWindowDevice)
In this PR	yes	stage 2

Which is faster, and by how much

Less than the framing suggests. An earlier revision of the example implemented both, over the same allocation, purely to measure the difference:

• 8×gfx950: within 0.2%, 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).
• Idle 4×gfx1250: the flat form consistently ahead by a near-constant 0.3–0.6 µs — 2% at 4 MiB, 3–7% at 256 KiB.

Constant in absolute terms is the tell: this is a per-block cost, not a per-byte one. The dependent load sits on the critical path between block launch and the first store, and once the stores are in flight the two models are the same instruction stream. So the pointer array costs a fixed sliver of launch-to-first-store latency and nothing per byte, which is why stage 1 can ship it without apology — and why the argument for the LSA window is about expressiveness, not bandwidth.

So why bother with the window at all

Because the things the array cannot express are the things that come next:

• Teams. "Which peers can I load/store, and what do I do about the rest" is one question with one answer in the LSA model, and no answer in torch's.
• Signals and barriers in the same address space. cco's LSA barrier is a {byte-offset, count} pair into the DevComm's own window, so it inherits peer addressing for free rather than needing a second table.
• Scale-out. ibgdaWin hangs off the same window struct: a kernel decides load/store vs RDMA per peer, with one handle.
• Register economy at scale. A table is O(world_size) of memory a kernel keeps re-reading; a base and a stride are two values, whatever the world size.

03Why it matters

Concretely, once set_backend("MORI") works: torch's symmetric-memory collectives, MemPool-allocated symmetric tensors, and any Triton kernel written against the pointer array all run on AMD hardware through mori's allocation path — and mori gains a place to put fabric handles, its window layout, and later its RDMA and SDMA transports where torch code will actually reach them.

04How it is designed

Where the code sits

Allocation

alloc(size, device, group) is deliberately dumb and entirely local — nothing collective happens here, because torch may allocate at any time on any rank:

1. Probe the device's shareable handle type (cached, once per device).
2. hipMemGetAllocationGranularity, round size + signal_pad up to it.
3. hipMemCreate → hipMemAddressReserve → hipMemMap → hipMemSetAccess(read-write, self) → hipMemset(0).
4. Record the block in a pointer-keyed table and hand the VA back to torch.

Memory is coarse-grained pinned (hipMemAllocationTypePinned). This was measured, not assumed: fine-grained/uncached windows — which is what cco windows are — halved bandwidth on gfx1250 (712 vs 1499 GB/s at 4 ranks) and changed nothing on gfx950.

Rendezvous

This is where the window is built. Every rank publishes one RendezvousReq through the torch Store — allocation size, buffer size, device index, pid, and (fabric path only) the 64-byte handle blob — then validates that all ranks allocated the same size, reserves one flat VA span of world_size × stride, and maps every rank's physical memory into its own slot.

Handle exchange: two paths, one code path around them

	Fabric	POSIX file descriptor
Handle is	64 opaque bytes	an integer index into one process's fd table
Travels over	the torch Store, inline in the rendezvous all-gather	a UNIX datagram socket with SCM_RIGHTS
Extra machinery	none	FdChannel — bind /tmp/mori_symm_<pid>_<rank>, send to every peer, receive world_size-1
Ordering	all-gather is ordered	datagrams arrive in any order, so each carries its sender's rank as the payload

The fd path exists only because torch's own IpcChannel is not exported. It is ~110 lines and self-contained; if libtorch ever exports the real one, it deletes cleanly.

Deliberate omissions

• No signal pad by default. barrier(), put_signal() and wait_signal() are unimplemented and raise, so reserving torch's 9216-byte pad would waste a whole 2 MiB physical page per window (backing is 2 MiB-paged) for something nothing can use. MORI_SYMM_SIGNAL_PAD=ON reserves it, and mori.allocator.signal_pad_supported() reports which build you have — torch's own collectives synchronise through that pad, so they need the flag.
• No multicast. HIP exposes no multicast API, so has_multicast_support() is false and multimem_* is out.
• No LSA window API. The pointers are strided, so base + rank*stride would work today — but shipping it here would be a fourth definition of a layout cco already owns. §02; stage 2 instead.
• Teardown off by default. See gaps.

Lifetime

Two details that are easy to get wrong. The allocator singleton is intentionally immortal, because register_availability() parks a reference in a libtorch-owned registry that outlives the extension's statics. And release is driven from an atexit hook registered on import, not from destructors: by the time static destructors run, HIP is already unwinding and every VMM call segfaults, while is_finalizing() still reports false.

05One backend, three architectures

The interesting design question is not "how do I support gfx1250" — it is how a single build behaves correctly on parts whose capabilities differ, against a HIP whose spelling of those capabilities is still moving. Three mechanisms do the work.

1. Capability by probe, not by architecture

There is no if (gfx1250) anywhere in the backend, and no reliance on a device attribute enum — those are not stable across HIP releases. Instead each device is probed once, at first allocation, by actually doing the thing:

hipMemAllocationHandleType ProbeHandleType(int dev) {          // cached per device
  chosen = hipMemHandleTypePosixFileDescriptor;
  prop   = MakeProp(dev, hipMemHandleTypeFabricCompat);
  if (hipMemGetAllocationGranularity(&gran, &prop, ...) == hipSuccess && gran > 0) {
    if (hipMemCreate(&h, gran, &prop, 0) == hipSuccess) {       // fails outright on gfx9
      if (hipMemExportToShareableHandle(&blob, h, hipMemHandleTypeFabricCompat, 0) == hipSuccess)
        chosen = hipMemHandleTypeFabricCompat;                  // create AND export must work
      hipMemRelease(h);
    }
  }
  hipGetLastError();          // the probe is expected to fail on gfx9; clear the sticky error
  return chosen;
}

One granularity-sized allocation, exported and released. Both steps have to succeed, because they fail in different places on different parts: on gfx9 hipMemCreate itself reports "operation not supported", whereas on gfx1250 with an older ROCm, create succeeded and only the export failed. Checking either one alone would have picked wrong on one of the two.

2. A compat shim for a moving API

Fabric handles are not spelled the same way in every ROCm release. include/mori/utils/hip_compat.hpp pins one name, mirroring the shim RCCL uses:

#ifdef HIP_FABRIC_API
typedef hipMemFabricHandle_t hipMemFabricHandle_compat_t;
#define MORI_MEM_HANDLE_TYPE_FABRIC hipMemHandleTypeFabric
#else
typedef struct { unsigned char data[64]; } hipMemFabricHandle_compat_t;
#define MORI_MEM_HANDLE_TYPE_FABRIC ((hipMemAllocationHandleType)0x8)
#endif

The same header also absorbs an unrelated ROCm break: hipMemImportFromShareableHandle takes the fd as a pointer-sized value from ROCm 7.1, and a pointer to the fd before that. Both quirks are one-line differences that would otherwise be sprinkled through the allocator.

3. A layout the kernel cannot tell apart

Whatever the handle type, rendezvous ends in the same place: one flat VA span, uniform stride, every slot mapped read-write, published as a pointer array. A kernel written against buffer_ptrs_dev sees identical geometry on a fabric part and an fd part. The all-to-all example demonstrates it — the same HIP source and the same Triton source run unmodified on wave64 gfx950 and wave32 gfx1250, with no arch guards and no tuning difference beyond the grid heuristic reading the CU count.

Architecture	Handle type chosen	Why	Status
gfx1250 (ROCm 7.14 / 7.15)	fabric	create + export both succeed; handle rides the Store	2 and 4 ranks OK
gfx950 (MI355X)	posix_fd	hipMemCreate with a fabric handle type is unsupported	2, 4, 8 ranks OK
gfx942 (MI308X)	posix_fd	same as gfx950	correct, but see the driver caveat

06What it measures

The example is a one-shot all-to-all over the window, written twice — once in HIP, once in Triton, each a single self-contained script with no build step. Every rank pushes its chunk into every peer's receive slot; the aggregate counts only the world_size-1 chunks that actually leave the device.

torchrun --nnodes=1 --nproc_per_node=8 all2all_hip.py    --chunk-kib 256
torchrun --nnodes=1 --nproc_per_node=8 all2all_triton.py --chunk-kib 256

Aggregate bandwidth, 4 MiB per peer

ranks	gfx950	gfx1250	gfx942
2	107.8 GB/s	514.6 GB/s	54.1 GB/s
4	500–567 GB/s	1705.1 GB/s	112.2 GB/s
8	1746–1874 GB/s	—	184.8 GB/s

256 KiB per peer, where launch and barrier cost still shows

ranks	gfx950	gfx1250	gfx942
2	76–84 GB/s	92.9 GB/s	40.3 GB/s
4	461–485 GB/s	516.8 GB/s	139.4 GB/s
8	1522–1829 GB/s	—	250.0 GB/s

Ranges are across repeats; anything above 2 ranks moves 6–13% run to run, so read single-value columns as ±10% too. The gfx1250 box has 4 GPUs, hence no 8-rank row.

Three findings worth keeping

• Grid shape dominates addressing. One block per destination rank left all but world_size CUs idle and could not keep enough writes in flight to cover interconnect latency: 15.8 GB/s on gfx1250 at 4 MiB. Slicing each chunk across blocks_per_peer blocks is worth ~2.5× on gfx950 and ~95× on gfx1250.
• The pointer array costs almost nothing against the LSA form — a fixed per-block sliver, no per-byte penalty. Numbers and reasoning in §02.
• Small-message Triton numbers measure the launcher. Timing the launch call alone: 2.90 µs (HIP pybind) vs 9.01 µs (Triton 3.7) on gfx950, and 2.44–3.03 vs 10.84–10.94 µs (Triton 3.8) on gfx1250. Launches are async, so an iteration costs max(kernel, launch) — the gfx1250 Triton row is pinned near 12 µs at both 256 KiB and 1 MiB despite 4× the data. Real uses amortise it with graphs or a persistent kernel.

07Roadmap and known gaps

Open gaps

• Teardown is disabled by default — the one item blocking the draft. Releasing a rendezvous'd window segfaults at world_size ≥ 4 (2 ranks are fine) somewhere in the unmap/release path. Not is_finalizing(), not interpreter shutdown (it faults on an explicit del), not fixed by quiescing with a barrier. Mappings are leaked meanwhile, which is cheap — symmetric buffers are few and long-lived. MORI_SYMM_TEARDOWN=1 re-enables it. Stage 2 may dissolve the problem entirely, since cco would own unmap ordering.
• No device-side barrier. dist.barrier() stands in until the signal pad is wired up.
• No multi-node group. world_within_direct_access() returns true unconditionally — correct today, since there is no RDMA path, and it is stage 3 that makes it conditional.
• Backend-specific torch ops. Things like all_to_all_nd dispatch on a hardcoded if backend == "NCCL", so third-party backends get the generic ops only.

Open questions for review

1. Is a CppExtension in setup.py the right home? It is the only mori target that links libtorch, and torch's build_ext is what derives the ABI tag, the pybind11 copy and _GLIBCXX_USE_CXX11_ABI.
2. For stage 2, is ccoMemImport / ccoWindowRegister overload C the intended path for a torch-owned allocation, and should ccoGetPeerPtr get a Cython binding? How should rendezvous(group_name) map onto ccoCommCreate's up-front perRankVmmSize?
3. Should barrier() be implemented over the signal pad with P2P stores now, or wait for cco's signal pool in stage 3?
4. Any idea on the teardown segfault at world_size ≥ 4?
5. How should this relate to pytorch#192524, which enables NCCLSymmetricMemory (RCCL) support on ROCm upstream?