# 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](#quick-reference)
- [1. Concepts](#1-concepts)
- [2. Initialization](#2-initialization)
- [3. Query APIs](#3-query-apis)
- [4. Memory Management](#4-memory-management)
- [5. P2P Address Translation](#5-p2p-address-translation)
- [6. Synchronization](#6-synchronization)
- [7. HIP Module Init (Triton Integration)](#7-hip-module-init-triton-integration)
- [8. Initialization Flags](#8-initialization-flags)
- [Environment Variables](#environment-variables)
- [Source Files](#source-files)

## Quick Reference

```python
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:

```bash
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`:

```python
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.):

```python
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):

```python
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:

```python
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:

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

---

## 3. Query APIs

```python
# 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

```python
# 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

```python
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:

```python
# 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:

```python
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:

```python
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:

```python
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](#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:

```cpp
// 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`) |