Skip to content

MORI-SHMEM-GUIDE.md

Latest commit

 

History

History
373 lines (266 loc) · 12.8 KB

MORI-SHMEM-GUIDE.md

File metadata and controls

373 lines (266 loc) · 12.8 KB
Raw

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

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)