[RFC] torch SymmetricMemory allocator over CCO: staged design (pointer array -> LSA window -> scale-out/SDMA)

[carlushuang]

Goal

A host-side, torch-native way to allocate buffers that mori CCO can consume. torch owns the lifecycle (device, stream, caching, rendezvous, tensor); CCO owns the window semantics (layout, peer addressing, MR, signals). The allocator is the seam, not a third memory subsystem.

This is the design #544 is being refactored around. Packaging of the resulting module is #549; this issue is only about the model.

Stage 1 — pointer-array model (what torch symm-mem actually is)

SymmetricMemory exposes peers as buffer_ptrs / buffer_ptrs_dev: an array of per-peer base addresses, one entry per rank. Every shipped backend does this, including the ones whose heap is not scattered at all:

backend	peer VA derivation	flat reservation underneath?
CUDA / IPC	buffer_ptrs[r] from imported handles	no
NVSHMEM	nvshmem_ptr(p) → pe_bases[pe] + off	yes — cuMemAddressReserve(npes * heap_size), peers at globalBase + slot*heapSize, but slots are rotated from mype+1, so pe is not the slot index and the array is unavoidable
NCCL/RCCL	buffer_ptrs is empty today; pytorch#192524 is adding LSA peer pointers	yes — ncclGetLsaPointer(w, off, peer) = add4G(lsaFlatBase, peer*stride4G) + off

So stage 1 is the plain torch contract and nothing else: alloc / rendezvous / buffer_ptrs(+_dev) / rank / world_size, no flat view in the public API, signal pad compiled out while its ops are unimplemented. That makes mori a drop-in backend for any existing torch symm-mem kernel or op.

This costs nothing measurable. A/B of the two addressing modes in the all2all example, same kernel, same window:

	flat (base + rank*stride)	pointer array
gfx950, 8 ranks, 4 MiB	1912.0 GB/s	1909.4 GB/s
gfx1250, 4 ranks, 4 MiB	1714.7 GB/s	1702.9 GB/s

Stage 2 — LSA window model (base + rank·stride), reusing CCO's scheme

Rather than invent a third layout, adopt CCO's window verbatim — include/mori/cco/cco.hpp, ccoWindowDevice:

winBase = flatBase + slotOffset
peer_va = winBase + ((uint64_t)peerLsaRank * stride4G << 32) + offset

Same 4 GiB-quantised stride as NCCL's add4G, and indexing by LSA rank, not world rank.

CCO already has the host hooks for precisely this handoff:

• ccoMemImport(comm, externalPtr, size, &ptr) — aliases an external HIP VMM allocation into the flat-VA slot; the header's own example is "e.g. a torch.symm_mem buffer"
• ccoWindowRegister(comm, externalPtr, size, &win, &localPtr) (overload C) — import + register, returns the flat-VA alias
• ccoGetPeerPtr(comm, localPtr, pe) — host-computable peer VA, documented as the way to pre-fill a peer-pointer table that a kernel dereferences

So stage 2 does not add a parallel implementation: the allocator allocates/exports/maps under torch's lifecycle, hands the allocation to CCO, and CCO owns the window. The flat view returned to Python becomes a literal ccoWindow view — (winBase, stride4G, lsaRank), named after CCO rather than the ad-hoc flat_layout() in the current draft — and stage 1's buffer_ptrs become ccoGetPeerPtr results, so there is one source of truth for peer addressing instead of two.

Open items for stage 2:

1. ccoGetPeerPtr has no Cython binding — cco.pxd exposes ccoMemAlloc / ccoMemImport / all three ccoWindowRegister overloads, but not this one.
2. Bootstrap mapping: torch calls rendezvous(group_name) per allocation, whereas ccoCommCreate takes perRankVmmSize up front and reserves the flat VA once per communicator. Needs a policy for sizing and for caching a comm per process group.
3. Handle type must agree between the two paths (fabric vs POSIX fd), including the granularity shift that requestedHandleType causes.

Stage 3 (future, needs CCO co-design) — scale-out and SDMA

Both axes are already visible in ccoWindowDevice, which is why they belong to CCO and not to the allocator:

• Scale-out. ibgdaWin carries the per-window MR (peerRkeys[worldSize], lkey, iova=0 + offset), so a peer with no LSA slot is representable. torch has a precedent for the mixed case: NVSHMEM's nvshmem_ptr returns null for network peers and the backend flips world_within_cuda_p2p_ off. Same shape here — null entries in buffer_ptrs, plus a window handle the kernel uses for RDMA.
• SDMA. Signals live on ccoDevComm::sdma, per-DevComm rather than per-window. This is where the compiled-out signal pad should come back — as CCO's signal pool, not a private 9216-byte pad — but it needs a DevComm, which the allocator does not own today.

Neither is in scope for #544.

Consequences for #544

Reorganise it down to stage 1: the torch contract, pointer array only, flat view dropped from the public surface until it is a CCO window. The example keeps both addressing modes, as the benchmark that will justify stage 2. Remaining known gap: releasing a rendezvous'd window segfaults at world_size >= 4 (teardown is off by default, MORI_SYMM_TEARDOWN=1 re-enables) — stage 2 may dissolve it outright, since CCO would own unmap ordering.

[jhchouuu]

Stage 2 re-check: the open items, measured

Went through the stage 2 open items against current main (plus #544 as merged and the examples in #616). Short version: stage 2 looks viable, and the thing I expected to block it does not.

Divergent teardown order is safe. This was my main worry, because it is the reason #544 deliberately avoided the shmem/cco allocators — torch tensors die on Python GC and the order is not synchronised across ranks. But ccoWindowDeregister (src/cco/cco_init.cpp:1481) contains no collective at all: it unmaps this rank's mappings of its peers and releases this rank's imported handles, nothing else. ccoWindowRegister is collective and order-sensitive, but symm_mem.rendezvous() is collective by construction, so that side matches. So folding rendezvous into a cco window does not reintroduce the problem.

Slots are recycled. 64 rounds of allocate → register_external_window → close → free, 16 MiB each, through a communicator asked for 256 MiB: the same local_ptr comes back every round, 1 GiB churned with no growth. torch's churn pattern is fine.

perRankVmmSize is a floor, not a cap. It is rounded up to the 4 GiB quantisation:

requested	actual per_rank_size
64 MiB	4096 MiB
256 MiB	4096 MiB
1024 MiB	4096 MiB

So the sizing policy in open item 2 is less delicate than it reads — anything ≤ 4 GiB is the same reservation, and it bounds live symmetric memory per rank, not lifetime allocations.

Open item 3 (handle type) does not apply to the import direction. ccoMemImport retains the external allocation's handle (hipMemRetainAllocationHandle) rather than exporting it, so it never touches an fd and does not care whether the allocator chose fabric or POSIX fd. It only matters for cco's own export path.

Open item 1 (ccoGetPeerPtr has no Cython binding) is still true, but does not block stage 2: the backend is C++ and can call the C API directly. The binding is only needed to expose a flat view to Python.

Two things that changed since this was written

• The teardown segfault is fixed (in allocator: register mori as a torch SymmetricMemory backend #544 as merged), so it is no longer a motivation for stage 2. Root cause was not unmap ordering: rendezvous() built the self slot from hipMemRetainAllocationHandle on an allocation that free() then released, so ~MoriSymmetricMemory released it a second time. The self slot now borrows the Block's handle.
• A ROCm 7.14 constraint worth recording here, since stage 2 moves this code into cco's care: a shareable fd must outlive hipMemRelease of the allocation it names. Closed before the mapping is granted, hipMemSetAccess fails with invalid argument; closed before the release, the physical memory is never returned. allocator: register mori as a torch SymmetricMemory backend #544 had to be fixed for both. cco already gets this right — ccoMemFree and the comm-destroy path both close(meta.shareFd) after hipMemRelease — which is a small argument in favour of folding: deleting the backend's own FdChannel removes that whole class of hazard rather than maintaining it in two places.

Suggested shape for the change

Internal only to begin with: rendezvous() creates/caches one ccoComm per torch process group (uid exchanged through the group Store, which the backend already uses for its own all-gather), registers the tensor with overload C, and serves buffer_ptrs from ccoGetPeerPtr. Public surface unchanged, so the existing tests and the three all-to-all examples in examples/torch_symm_all2all/ A/B it directly. Flat view, signal pad on cco's pool, and scale-out then land as separate changes.

Examples showing the two composing today, without touching the backend, are in #616.