MORI Shmem Guide

MORI Shmem provides OpenSHMEM-style symmetric memory APIs for GPU memory management and RDMA communication. It is the foundation layer that MORI-EP and MORI-IO build upon — shmem must be initialized before using any other MORI component.

Table of Contents

• Quick Reference

• 1. Concepts

• 2. Initialization

• 3. Query APIs

• 4. Memory Management

• 5. P2P Address Translation

• 6. Synchronization

• 7. HIP Module Init (Triton Integration)

• 8. Initialization Flags

• Environment Variables

• Source Files

Quick Reference

import mori

# Initialize from PyTorch process group
mori.shmem.shmem_torch_process_group_init("default")

# Query
my_rank = mori.shmem.shmem_mype()
num_ranks = mori.shmem.shmem_npes()

# Allocate symmetric memory
ptr = mori.shmem.shmem_malloc(size_in_bytes)

# Register existing buffer for RDMA
mori.shmem.shmem_buffer_register(tensor.data_ptr(), tensor.nbytes)

# P2P address translation (intra-node)
remote_ptr = mori.shmem.shmem_ptr_p2p(ptr, my_pe, dest_pe)

# Synchronize
mori.shmem.shmem_barrier_all()

# Cleanup
mori.shmem.shmem_free(ptr)
mori.shmem.shmem_finalize()

Imports:

What	Import
All shmem APIs	import mori.shmem or from mori import shmem
Init flags	mori.shmem.MORI_SHMEM_INIT_WITH_MPI_COMM
	mori.shmem.MORI_SHMEM_INIT_WITH_UNIQUEID

---

1. Concepts

Symmetric Memory

Symmetric memory is GPU memory that is allocated at the same virtual offset across all participating PEs (Processing Elements / ranks). This enables RDMA hardware to directly access remote GPU memory using simple address arithmetic — no address translation tables needed at runtime.

Processing Element (PE)

A PE is a participant in the symmetric memory domain. Each GPU rank maps to one PE. PEs are numbered 0 to N-1.

Symmetric Heap

The symmetric heap is a pre-allocated region of GPU memory from which shmem_malloc allocates. Its size is controlled by the MORI_SHMEM_HEAP_SIZE environment variable:

export MORI_SHMEM_HEAP_SIZE=6G  # Must be set before shmem init

---

2. Initialization

Shmem must be initialized exactly once per process. There are three methods:

Method 1: PyTorch Process Group (Recommended)

Use this when PyTorch distributed is already initialized. The process group must be registered with a name before calling shmem_torch_process_group_init:

import torch
import torch.distributed as dist
import mori

# Standard PyTorch distributed init
torch.cuda.set_device(rank)
dist.init_process_group(
    backend="cpu:gloo,cuda:nccl",
    rank=rank, world_size=world_size,
    device_id=torch.device("cuda", rank),
)

# Register the process group with a name
world_group = dist.group.WORLD
torch._C._distributed_c10d._register_process_group("default", world_group)

# Initialize shmem from the registered process group
mori.shmem.shmem_torch_process_group_init("default")

Note: The _register_process_group("default", ...) step is required — shmem_torch_process_group_init looks up the group by name to extract rank/world_size and bootstrap the shmem connections.

Method 2: Unique ID (No PyTorch Distributed)

Use this when PyTorch distributed is not available (e.g., standalone applications, custom launchers). Rank 0 generates a unique ID and broadcasts it to all ranks via any transport (file, MPI, TCP, etc.):

import mori

# Rank 0 generates a unique ID
if rank == 0:
    unique_id = mori.shmem.shmem_get_unique_id()  # Returns 128 bytes
    # Broadcast unique_id to all ranks via your transport

# All ranks initialize with the same unique ID
mori.shmem.shmem_init_attr(
    mori.shmem.MORI_SHMEM_INIT_WITH_UNIQUEID,
    rank,        # My rank
    world_size,  # Total ranks
    unique_id,   # Shared unique ID (bytes)
)

Example using file-based broadcast (for single-node testing):

import os, time

uid_file = "/tmp/mori_unique_id"
if rank == 0:
    unique_id = mori.shmem.shmem_get_unique_id()
    with open(uid_file, 'wb') as f:
        f.write(unique_id)
else:
    while not os.path.exists(uid_file):
        time.sleep(0.1)
    with open(uid_file, 'rb') as f:
        unique_id = f.read()

mori.shmem.shmem_init_attr(
    mori.shmem.MORI_SHMEM_INIT_WITH_UNIQUEID,
    rank, world_size, unique_id,
)

Method 3: MPI Communicator (C++ / MPI environments)

Use this in MPI-based applications. Pass MORI_SHMEM_INIT_WITH_MPI_COMM as the flag:

mori.shmem.shmem_init_attr(
    mori.shmem.MORI_SHMEM_INIT_WITH_MPI_COMM,
    rank, world_size, mpi_comm,
)

Note: This is primarily used in C++ applications where an MPI_Comm handle is available. In Python, Method 1 or Method 2 are preferred.

Initialization Comparison

Method	When to use	Dependencies
PyTorch Process Group	LLM inference/training frameworks (vLLM, SGLang, etc.)	torch.distributed
Unique ID	Standalone apps, custom launchers, non-PyTorch environments	None (any broadcast mechanism)
MPI Communicator	C++ MPI applications	MPI

Finalization

Always finalize shmem before process exit:

mori.shmem.shmem_finalize()
# Then destroy PyTorch process group if applicable
dist.destroy_process_group()

---

3. Query APIs

# Get my PE (rank) ID — 0 to npes-1
my_pe = mori.shmem.shmem_mype()

# Get total number of PEs
total_pes = mori.shmem.shmem_npes()

# Get number of RDMA queue pairs per PE
num_qp = mori.shmem.shmem_num_qp_per_pe()

---

4. Memory Management

Allocating Symmetric Memory

# Basic allocation
ptr = mori.shmem.shmem_malloc(size_in_bytes)

# Aligned allocation (alignment must be power of 2)
ptr = mori.shmem.shmem_malloc_align(alignment=256, size=size_in_bytes)

# Allocation with flags
ptr = mori.shmem.shmem_ext_malloc_with_flags(size_in_bytes, flags)

All allocation functions return an integer address (int). The allocated memory is symmetric — the same offset is reserved on every PE.

Freeing Symmetric Memory

mori.shmem.shmem_free(ptr)

Registering Existing Buffers

Register a PyTorch tensor or other existing GPU memory for RDMA operations without allocating new symmetric memory:

# Register
tensor = torch.zeros(1024, 7168, dtype=torch.bfloat16, device="cuda")
mori.shmem.shmem_buffer_register(tensor.data_ptr(), tensor.nbytes)

# ... use tensor in MORI operations ...

# Deregister when done
mori.shmem.shmem_buffer_deregister(tensor.data_ptr(), tensor.nbytes)

Both functions return a status code (0 for success).

---

5. P2P Address Translation

For intra-node GPU-to-GPU access, translate a local symmetric pointer to its P2P-accessible address on a remote PE:

remote_ptr = mori.shmem.shmem_ptr_p2p(local_ptr, my_pe, dest_pe)

Return Value	Meaning
Non-zero	P2P address on dest_pe (same-node, XGMI-connected GPUs)
0	Connection uses RDMA transport (different nodes) or pointer is invalid

---

6. Synchronization

Global barrier — blocks until all PEs reach the barrier:

mori.shmem.shmem_barrier_all()

---

7. HIP Module Init (Triton Integration)

When using Triton-compiled kernels that access shmem device symbols, initialize the HIP module:

mori.shmem.shmem_module_init(hip_module_handle)

This copies the current GPU states to the globalGpuStates symbol in the dynamically compiled Triton kernel module.

---

8. Initialization Flags

Flag	Value	Description
MORI_SHMEM_INIT_WITH_MPI_COMM	0	Initialize using MPI communicator
MORI_SHMEM_INIT_WITH_UNIQUEID	1	Initialize using broadcast unique ID

---

Environment Variables

Variable	Description	Default
MORI_SHMEM_HEAP_SIZE	Symmetric heap size (e.g., "6G", "2G", "512M"). Must be set before initialization.	Required
MORI_SHMEM_MODE	Heap mode: "static", "vmm", or "isolation"	"static"
MORI_SHMEM_HEAP_TYPE	Heap memory type: "normal" (cached) or "uncached"	"uncached"
MORI_SHMEM_VMM_CHUNK_SIZE	VMM mode chunk size in bytes	Auto
MORI_SOCKET_IFNAME	Network interface for shmem bootstrap TCP connections (e.g., "lo", "eth0")	Auto-detect
MORI_RDMA_DEVICES	RDMA NIC selection. Include: mlx5_0,mlx5_1. Exclude: ^mlx5_2,mlx5_3	All available
LD_LIBRARY_PATH	mori loads libibverbs dynamically at runtime (dlopen) instead of linking it; to use an out-of-tree libibverbs, put its directory here.	libibverbs.so / libibverbs.so.1
MORI_NUM_QP_PER_PE	Number of RDMA queue pairs per PE	1
MORI_IB_GID_INDEX	InfiniBand GID index for RDMA connections	Auto-detect
MORI_RDMA_SL	RDMA service level	Auto
MORI_RDMA_TC	RDMA traffic class	Auto
MORI_DISABLE_P2P	Disable P2P (XGMI) transport, force RDMA	Not set
MORI_ENABLE_RAIL_ONLY	Only create RDMA QPs to same-rail peers (same index within their node). For rail-isolated fabrics where cross-rail QPs cannot be established. See Rail-only connections.	Not set
MORI_DISABLE_TOPO	Disable topology detection	Not set
MORI_IGNORE_CPU_AFFINITY	By default the thread calling ShmemInit is pinned to the CPUs local to its GPU (sysfs local_cpulist, intersected with the existing cpuset). Within that set the CPUs are grouped into physical cores by thread_siblings_list and partitioned into one disjoint slice per GPU on the NUMA node, so two ranks never share the two SMT siblings of one core; the split applies only when the process sees more than one GPU on that NUMA node and there is at least one physical core per GPU, otherwise the whole node is bound. Set to 1 to disable and leave CPU placement to an outer numactl/torchrun. Works in SPMT (single-process multi-GPU) too: each per-GPU init thread pins to its own GPU’s slice.	Not set (binding on)
MORI_CPU_AFFINITY_NO_SPLIT	Bind to the whole NUMA node instead of this rank’s own slice of it.	Not set (split on)
MORI_GLOBAL_LOG_LEVEL	Global log verbosity: TRACE, DEBUG, INFO, WARN, ERROR	INFO
MORI_PRECOMPILE	Precompile all JIT kernels on import	Not set
MORI_DISABLE_JIT	Disable JIT compilation of device bitcode	Not set

---

Rail-only connections

By default Context builds a full mesh: every cross-node peer gets MORI_NUM_QP_PER_PE queue pairs. On a rail-isolated fabric only NICs on the same rail can reach each other, so the cross-rail QPs never reach RTR and init hangs — even though no kernel would ever have used them.

MORI_ENABLE_RAIL_ONLY=1 restricts QP creation to same-rail peers. Two ranks are on the same rail when they have the same index within their own node, which is the same relation the EP internode kernels already encode in their proxy:

// src/ops/dispatch_combine/internode_v1.cpp
int proxyPe = i * config.gpuPerNode + (config.rank % config.gpuPerNode);

On a 2-node × 8-GPU job this takes each rank from 8 cross-node peers (32 QPs at 4 QP/PE) down to 1 (4 QPs). Skipped peers keep the same empty endpoint stubs already used for non-RDMA peers, so rdmaEps indexing is unchanged; the predicate is symmetric, so both ends stub each other and the handle AllToAll stays aligned.

This is opt-in, and only safe for kernels whose cross-node traffic is same-rail.

Path	Cross-node RDMA target	Rail-only safe
EP InterNodeV1 / InterNodeV1LL	proxyPe	yes
EP InterNode (v0)	arbitrary destPe	no
EP AsyncLL	arbitrary destPe	no
CCO with CCO_GDA_CONNECTION_RAIL	same rail	yes
CCO with FULL / CROSSNODE	arbitrary peer	no

Requirements checked at init — if any fails, mori logs an error and falls back to the full mesh rather than mis-pairing ranks:

• node sizes are uniform,

• ranks are node-major contiguous (rankInNode == rank % gpuPerNode), which is what the EP proxy formula above assumes.

Same-host peers are never dropped: their RDMA loopback (used only under MORI_DISABLE_P2P) does not traverse the fabric.

To confirm it took effect, look for this line from each rank:

rail-only: rank 3 rail 3 created QPs to 1 peers (4 QPs): 11

---

Source Files

File	Description
python/mori/shmem/api.py	Python shmem API
python/mori/shmem/__init__.py	Public exports
include/mori/shmem/	C++ shmem headers
src/pybind/pybind.cpp	Python module entry point
src/pybind/mori.cpp	Shmem binding registration (RegisterMoriShmem)