MORI (ROCm/mori) — L2 Distillation #

§1 TL;DR #

AMD's modular, composable GPU communication framework providing RDMA+GPU building blocks (IBGDA, P2P, SDMA collectives) plus turnkey MoE dispatch/combine (MORI-EP), point-to-point IO, and symmetric memory with device-linkable bitcode — deployed across SGLang, vLLM, DeepSpeed, and RTP-LLM on MI300X/MI355X.

§2 Project Identity #

§3 What & Why — Motivation #

Q1 痛点 #

GPU-side RDMA communication (IBGDA, GPUDirect Async) is hard to adopt: existing libraries are either monolithic (NCCL/RCCL — no composability) or vendor-locked (DeepEP — NVIDIA only). MoE expert-parallel dispatch/combine needs both intra-node XGMI and inter-node RDMA in a single fused kernel, but existing stacks force users to choose one transport. No existing library offers GPU-initiated shmem primitives linkable from Triton or custom HIP kernels.

Q2 方法 #

MORI takes an MLIR-inspired layered approach: composable building blocks (transport backends, topology detection, shmem management) that applications assemble. Five subsystems serve distinct needs:

核心技术壁垒: The device-bitcode layer (MORI-IR) that exposes libmori_shmem_device.bc with 50+ device functions for GPU-initiated communication. This enables any kernel framework (Triton, FlyDSL, raw HIP) to call shmem primitives from within device code, eliminating the host-device round-trip that plagues traditional communication libraries. No competitor offers this level of device-side composability.

Q3 结果 #

• EP8 intra-node: 345 GB/s dispatch, 420 GB/s combine on MI355X+AINIC (SOTA)
• EP16 inter-node: 179/234 GB/s dispatch/combine XGMI, 54/71 GB/s RDMA
• MORI-IO: 48 GB/s GPU Direct RDMA READ at 1 MB messages (MI300X+Thor2)
• EP8 latency: 31 μs dispatch, 36 μs combine on MI355X (low-latency path)
• Adopted by SGLang, vLLM, DeepSpeed, RTP-LLM, Triton-distributed, AITER

§4 Architecture & Module Map #

Top-level module responsibilities:

• application/: Manages SymmMemObjPtr lifecycle — the fundamental symmetric memory allocation visible across all PEs
• shmem/: rocshmem-compatible runtime implementing OpenSHMEM-style barriers, malloc, P2P pointer resolution
• ops/dispatch_combine/: The five EP kernel implementations (IntraNode, InterNode, InterNodeV1, InterNodeV1LL, AsyncLL) plus handle/buffer management
• io/: Multi-backend P2P engine with session abstraction, RDMA/XGMI/TCP backends
• collective/: SDMA-offloaded AllReduce, AllGather, All2All (bypasses CUs entirely)
• umbp/: WIP distributed KV store with DRAM/SSD tiered storage and SPDK backend
• pybind/: Single registration point binding all C++ subsystems to Python

§5 Entry Points & API Surface #

Public Python API #

import mori
# Lazy-loaded submodules:
mori.ops       # EP dispatch/combine (EpDispatchCombineConfig, EpDispatchCombineOp)
mori.io        # IOEngine, IOEngineSession, IOEngineConfig
mori.shmem     # shmem_init, malloc, barrier, ptr_p2p
mori.ir        # find_bitcode(), MORI_DEVICE_FUNCTIONS (50+ ABI entries)
mori.ccl       # All2allSdma, AllgatherSdma, AllreduceSdma
mori.jax       # XLA FFI custom calls
mori.cpp       # Raw C++ bindings (AllGatherIntoTensor, DataType)

CLI entry point #

mori CLI via python/mori/cli.py — registered in pyproject.toml.

Configuration surface (top 10) #

Extension points #

• Transport backends: IOEngine.create_backend(type) — add new NIC support via backend interface
• Kernel types: EpDispatchCombineKernelType enum — extensible with new kernel implementations
• Device functions: MORI_DEVICE_FUNCTIONS dict — new shmem primitives added as extern "C" entries in bitcode
• Tuning configs: Drop JSON files into tuning_configs/ for new hardware/model combinations
• JIT pipeline: compile_genco() in jit/core.py accepts arbitrary .cpp → .hsaco

§6 Core Data Structures #

EpDispatchCombineConfig (hot path) #

• Definition: include/mori/ops/dispatch_combine/dispatch_combine.hpp + python/mori/ops/dispatch_combine.py
• Layout: 19-field POD struct (kPackedI32Len = 19), serializable to int32_t[19] via ToPackedI32Array()
• Lifecycle: Created once per EP group init, immutable during inference, destroyed with handle
• Mutation: Immutable after construction. Thread-safe reads. Serialized to GPU constant memory for kernel access.
• Key invariant: MaxNumTokensToSendPerRank() = maxNumInpTokenPerRank * numExpertPerToken — this bounds all symmetric memory allocations

EpDispatchCombineHandle #

• Definition: include/mori/ops/dispatch_combine/dispatch_combine.hpp
• Layout: Owns config + std::variant + signal/barrier/index-map buffers
• Lifecycle: Created at EP init (allocates symmetric memory across all PEs), destroyed at EP teardown
• Memory: Allocates multiple symmetric memory objects via SymmMemObjPtr — each visible across all PEs in the group. IntraNode needs 3 buffers; InterNodeV1 needs 5 (adds staging and dispatch input).
• Ownership: Python EpDispatchCombineOp owns the C++ handle via pybind shared_ptr

SymmMemObjPtr (foundation) #

• Definition: include/mori/application/application.hpp
• Layout: Smart pointer to symmetric memory allocation with P2P address resolution
• Lifecycle: Created via shmem_malloc → registered across all PEs → freed via shmem_free
• Key property: shmem_ptr_p2p(ptr, my_pe, dest_pe) returns a direct GPU pointer for intra-node (XGMI), or 0 for inter-node (forces RDMA path)

IOEngine / IOEngineSession #

• Definition: python/mori/io/engine.py + C++ src/io/
• Layout: Engine holds registered backends + memory descriptors; Session is a lightweight handle for a (local_mem, remote_mem) pair
• Lifecycle: Engine created per process, backends added/removed dynamically, sessions created per transfer pair
• Transfer UIDs: Monotonic IDs allocated via allocate_transfer_uid(), used for completion tracking

§7 Critical Path Analysis #

EP Dispatch critical path (user input → routed tokens) #

Python dispatch() call
  │
  ├─ [1] _resolve_launch_params()          ← JSON config lookup: (arch, kernel_type, ep_size, dtype, hidden_dim, tokens) → (block_num, rdma_block_num, warp_per_block)
  │
  ├─ [2] mori_cpp.prepare_inference_args() ← C++ binding: sets input/weight/scale/index pointers in handle
  │
  ├─ [3] mori_cpp.build_args()             ← Type-erased EpDispatchCombineArgsRaw construction for kernel launch
  │
  ├─ [4] _launch() or _launch_multi()      ← HIP kernel launch via JIT-compiled .hsaco module
  │     │
  │     ├─ IntraNode: 1 kernel (EpDispatchIntraNodeKernel)
  │     │   └─ XGMI P2P writes to remote symmetric memory
  │     │
  │     ├─ InterNodeV1: 2 kernels launched atomically
  │     │   ├─ EpDispatchCopyToStaging (XGMI → staging buffer)
  │     │   └─ EpDispatchInterNodeV1Kernel (staging → RDMA)
  │     │
  │     └─ AsyncLL: 3 kernels launched atomically
  │         ├─ SlotAssign (token-to-slot mapping)
  │         ├─ CopyMultiBlock (multi-block copy to staging)
  │         └─ SendTransfer (RDMA transfer via IBGDA)
  │
  └─ [5] from_gpu_ptr() wraps pre-allocated symmetric memory as torch.Tensor (zero-copy)

Latency budget (EP8 IntraNode, MI355X, 128 tokens):

• Steps 1-3: <1 μs (Python + C++ binding overhead)
• Step 4 (kernel): ~31 μs dispatch, ~36 μs combine
• Step 5: negligible (pointer wrap)
• Total: ~31-36 μs

Memory traffic (EP8, 128 tokens, hidden=7168, FP8):

• Per-token payload: 7168 bytes data + scales + indices ≈ 7.2 KB
• Total dispatch: 128 × 7.2 KB × 7 (sent to 7 other ranks) ≈ 6.3 MB moved over XGMI

JIT compilation path (cold start only) #

_ensure_jit_kernels() in EpDispatchCombineOp.__init__
  │
  ├─ detect_gpu_arch() → "gfx942" / "gfx950"
  ├─ Check ~/.mori/jit/ cache for matching .hsaco
  ├─ [cache miss] compile_genco(): hipcc .cpp → .hsaco
  │   └─ ~5-15s per kernel file
  └─ hipModuleLoad() + hipModuleGetFunction() via ctypes

IO critical path (GPU-to-GPU RDMA read) #

IOEngine.read(local_mem, l_off, remote_mem, r_off, size, uid)
  │
  ├─ Backend dispatch (RDMA / XGMI / TCP)
  ├─ [RDMA] Post WQE via IBGDA (GPU-initiated, no host involvement)
  ├─ [XGMI] JIT scatter/gather kernel for non-contiguous buffers
  └─ TransferStatus returned (polled via pop_inbound_transfer_status)

§8 API Design Decisions #

Why the dispatch/combine split (not a single all-to-all)?

MoE expert parallelism has asymmetric phases: dispatch fans out tokens to experts (1-to-K), combine collects results back (K-to-1). A generic all-to-all would waste bandwidth on the routing metadata and miss optimization opportunities for each direction (e.g., FP8 dispatch + BF16 combine).

Why 5 kernel types instead of a single adaptive kernel?

The IntraNode→InterNode→InterNodeV1→InterNodeV1LL→AsyncLL hierarchy reflects fundamentally different hardware paths. IntraNode uses direct XGMI P2P writes (no staging). InterNode adds RDMA via staging buffers. InterNodeV1 optimizes RDMA bandwidth. V1LL trades bandwidth for latency. AsyncLL pipelines all three stages. A single kernel cannot efficiently handle all topologies because the synchronization and memory access patterns differ structurally.

Why OpenSHMEM-style API for shmem?

OpenSHMEM is a well-understood PGAS model with clear semantics for symmetric memory allocation and one-sided communication. By matching this API, MORI enables porting from rocshmem with minimal changes and provides familiar semantics for HPC developers.

Why device bitcode (IR) instead of host-side API?

Host-side communication APIs require kernel launch → host callback → communication → kernel relaunch, adding 10-50 μs per round-trip. Device-linkable bitcode enables GPU-initiated communication within a running kernel, which is essential for latency-sensitive MoE dispatch where the entire operation must complete in <40 μs.

Why runtime dlopen for NIC libraries?

A single MORI binary must run on clusters with Mellanox CX7, Broadcom Thor2, or AMD Pollara NICs. Compile-time NIC selection would require separate builds per NIC vendor. Runtime dlopen of libmlx5.so / libbnxt_re.so / libionic.so with auto-detection solves this, enabling universal wheel distribution via PyPI.

§9 Implementation Highlights #

1. Host/device build separation #

MORI requires no hipcc at install time. The C++ host code compiles with a standard compiler (g++/clang++). GPU kernels are .cpp files with HIP annotations that are JIT-compiled to .hsaco (AMD's GPU binary format) at first use via python/mori/jit/core.py:compile_genco(). Results cache to ~/.mori/jit/ keyed by (source hash, arch, compiler flags). This enables pip install amd_mori to work without any GPU SDK on the build machine — a rarity for GPU communication libraries.

2. Multi-kernel atomic launch #

Inter-node EP kernels launch 2-4 cooperative GPU kernels in a single _launch_multi() call. For AsyncLL, three kernels run concurrently:

• SlotAssign: maps tokens to transfer slots
• CopyMultiBlock: copies data to staging buffers using all CUs
• SendTransfer: initiates RDMA via IBGDA with one warp per destination

These kernels coordinate through symmetric memory signals and atomic counters without host synchronization.

3. Tuning config system #

A two-phase auto-tuning pipeline:

1. Calibrate: exhaustive sweep of (block_num, warp_per_block, rdma_block_num) on target hardware
2. Quick sweep: fast search around calibrated optima for new (dtype, hidden_dim, token_count) combinations

Results stored as JSON in python/mori/ops/tuning_configs/, loaded at runtime via TuningConfigManager. The key space is (gpu_arch, gpu_model, kernel_type, ep_size, phase, dtype, hidden_dim, num_tokens).

4. Config serialization via packed int32 array #

EpDispatchCombineConfig serializes to int32_t[19] via ToPackedI32Array(). This allows the entire config to be passed to GPU kernels as a flat argument without pointer chasing or struct layout concerns across host/device ABI boundaries.

关键实现细节 #

1. shmem_ptr_p2p() returns 0 for cross-node: This is the routing mechanism — kernels check the return value to decide between XGMI direct write (non-zero) and RDMA staging path (zero). A subtle but load-bearing invariant.
1. from_gpu_ptr() zero-copy output wrapping: Dispatch/combine results are not copied to new allocations. Instead, pre-allocated symmetric memory buffers are wrapped as PyTorch tensors via raw pointer. This eliminates a full-tensor copy on every MoE layer forward pass, but means the user must consume results before the next dispatch call overwrites them.

§10 Concurrency & Memory #

• Concurrency model: Multi-process (one process per GPU, matching PyTorch DDP/FSDP pattern). No Python threading for communication. GPU kernels use warp-level parallelism with shared memory for intra-block coordination and symmetric memory signals for inter-PE coordination.
• Lock hierarchy: No host-side locks in the hot path. GPU-side coordination uses mori_shmem_signal_wait_until() (spin-wait on symmetric memory location) and mori_shmem_barrier_all_on_stream() (stream-ordered barrier).
• Memory management: Symmetric memory allocated via MORI-SHMEM (shmem_malloc / shmem_malloc_align). Default heaps: 4 GB static, 16 GB VMM. All allocations are registered across all PEs at allocation time. Freed explicitly via shmem_free. No GC; user manages lifetime.
• Buffer reuse: EP handles pre-allocate all symmetric buffers at init time (sized by MaxNumTokensToRecv()). Buffers are reused across dispatch/combine calls — no per-call allocation.

§11 Performance Characteristics #

MORI-EP bandwidth (4096 tokens, 7168 hidden, top-8, FP8→BF16) #

Platform	Kernel	Dispatch XGMI	Dispatch RDMA	Combine XGMI	Combine RDMA
MI355X + AINIC	EP8	345 GB/s	—	420 GB/s	—
MI355X + AINIC	EP16-V1	179 GB/s	54 GB/s	234 GB/s	71 GB/s
MI355X + AINIC	EP32-V1	85 GB/s	46 GB/s	110 GB/s	61 GB/s
MI300X + CX7	EP8	307 GB/s	—	330 GB/s	—
MI300X + CX7	EP16-V1	171 GB/s	52 GB/s	219 GB/s	67 GB/s

MORI-EP latency (128 tokens, 7168 hidden, top-8, FP8→BF16) #

Platform	Kernel	Dispatch	Combine
MI355X + AINIC	EP8	31 μs	36 μs
MI355X + AINIC	EP16-V1-LL	84 μs	108 μs
MI300X + CX7	EP8	35 μs	47 μs
MI300X + CX7	EP16-V1-LL	76 μs	122 μs

MORI-IO bandwidth (GPU Direct RDMA READ, MI300X + Thor2) #

Message size	Avg BW	Avg latency
1 KB	3.53 GB/s	37.1 μs
64 KB	41.4 GB/s	202.7 μs
1 MB	48.3 GB/s	2777.8 μs

Scaling behavior #

• EP8 → EP16 → EP32: bandwidth drops ~50% per doubling due to RDMA link sharing and increased cross-node traffic. Latency grows sub-linearly.
• MORI-IO bandwidth saturates at ~48 GB/s (single-GPU, 128-batch), close to theoretical PCIe Gen5 x16 limit for RDMA reads.
• DeepSpeed SDMA AllGather integration delivers ~10% end-to-end training speedup by offloading collective traffic to dedicated SDMA copy engines, freeing CUs for compute.

§12 Tech Debt & Code Quality #

• Testing: pytest-based functional tests for EP (intra/inter node) and IO; C++ unit tests in tests/cpp/; no published coverage numbers
• Linting: pre-commit hooks configured (see contribution guide); specific linter tools not documented in L1
• CI matrix: Not detailed in repo README; likely internal AMD CI given ROCm org affiliation
• Known WIP: UMBP subsystem explicitly marked WIP; SPDK submodule only pulled when BUILD_UMBP=ON
• Stale benchmarks: EP32 numbers marked with asterisks ("stale data from previous kernel version") — indicates benchmark infrastructure lags kernel development
• Dependencies: spdlog (statically linked, hidden visibility), torch (optional at build time), SPDK (optional for UMBP)

§13 Community Health #

• Maintainer: AMD ROCm team (corporate-backed)
• Stars: 124 (modest but growing; project is ~2 months old at v1.1.1)
• Integration velocity: High — 8 major framework integrations (SGLang, vLLM, DeepSpeed, RTP-LLM, Triton-distributed, OpenUCL, AITER, InferenceX) in 5 months
• Release cadence: 3 releases in 1 month (v0.1.0 → v1.0.0 → v1.1.1), indicating rapid iteration
• Bus factor: Unknown externally; AMD ROCm org has deep bench but external contributor data not in L1
• Governance: Corporate BDFL (AMD)

§14 Comparison with Alternatives #

Dimension	MORI	DeepEP	NCCL/RCCL
Vendor	AMD (ROCm)	DeepSeek (NVIDIA)	NVIDIA / AMD
GPU support	MI300X, MI325X, MI355X	H800, A100, H100	All NVIDIA + AMD
NIC support	CX7, Thor2, Pollara (runtime dlopen)	CX7 only	CX7 (compiled in)
EP dispatch/combine	5 kernel types, auto-tuned	3 kernel types (normal, low-latency, low-latency-LL)	N/A (collective-only)
Device-linkable bitcode	Yes (50+ functions)	No	No
P2P IO engine	Yes (multi-backend sessions)	No	No
SDMA collectives	Yes (offloads from CUs)	No	No
API style	OpenSHMEM + Python	Python-only EP	C/C++ collective
Install	pip install (no hipcc)	pip install (needs CUDA)	System package
JIT compilation	.cpp → .hsaco at first use	Precompiled	Precompiled
Composability	Modular building blocks	Monolithic EP lib	Monolithic collective lib
Framework integrations	SGLang, vLLM, DeepSpeed, RTP-LLM, Triton-distributed	SGLang, vLLM	All major frameworks
License	MIT	Apache-2.0	BSD-3
Maturity	~2 months public	~6 months public	~8 years

Bold = winner per row.

MORI wins on composability, multi-vendor support, and breadth (EP + IO + CCL + IR in one package). DeepEP wins on NVIDIA ecosystem maturity. RCCL wins on broad collective coverage but lacks EP-specific kernels.

§15 Verdict & Recommendations #

When to adopt #

• Running MoE inference on AMD MI300X/MI325X/MI355X and need expert-parallel dispatch/combine
• Need GPU-to-GPU KV cache transfer for PD disaggregation on AMD platforms
• Writing custom Triton kernels on AMD that need device-initiated communication (MORI-IR is unique here)
• Deploying across mixed NIC environments (CX7 + Pollara + Thor2)
• Need SDMA-offloaded collectives to free CUs for compute overlap

When NOT to adopt #

• Running on NVIDIA GPUs (DeepEP or NCCL are better supported)
• Need only standard collectives (AllReduce, AllGather) — RCCL is more mature and has broader collective coverage
• UMBP features required for production (still WIP)
• Need guaranteed stable API — 3 releases in 1 month indicates rapid but potentially breaking iteration

Suggested contributions #

1. Benchmark automation: EP32 numbers are stale — a CI-integrated benchmark suite would prevent staleness
2. CUDA backend: The modular architecture could support NVIDIA GPUs via CUDA backend; would massively expand adoption
3. Formal API stability policy: Semver + deprecation timeline would help downstream integrators (SGLang, vLLM) plan upgrades

§16 论证链 #

Step	Claim	Evidence	Validity
1	Monolithic communication libraries (NCCL) are suboptimal for MoE EP	EP requires asymmetric dispatch/combine with topology-aware kernel selection; NCCL's AllToAll is symmetric	Valid — MoE token routing is fundamentally 1-to-K, not N-to-N
2	MLIR-inspired composable building blocks enable better specialization	5 subsystems (EP, IO, CCL, SHMEM, IR) with independent APIs composable via shared shmem layer	Valid — demonstrated by diverse downstream integrations requiring different subsystem combinations
3	Device-linkable bitcode eliminates host-device round-trips	MORI-IR provides 50+ extern "C" device functions in .bc format, linkable by Triton/HIP	Valid — Triton-distributed integration demonstrates device-initiated shmem from within Triton kernels
4	Runtime NIC detection via dlopen enables universal packaging	Single wheel supports CX7+Thor2+Pollara without recompilation	Valid — pip install amd_mori works across NIC vendors; NIC auto-detected at runtime
5	Multi-kernel atomic launch achieves SOTA inter-node EP	InterNodeV1: 2 kernels (staging copy + RDMA); AsyncLL: 3 kernels (slot assign + copy + transfer)	Valid — 345/420 GB/s dispatch/combine on MI355X EP8 is competitive with DeepEP on H100

§17 实现 Cross-Reference #

Component	Key file	Purpose
EP Python API	python/mori/ops/dispatch_combine.py	EpDispatchCombineOp.dispatch() / combine() — the user-facing entry points
EP C++ handle	include/mori/ops/dispatch_combine/dispatch_combine.hpp	EpDispatchCombineHandle, config, symmetric buffer structs
IntraNode kernel	src/ops/dispatch_combine/dispatch_combine.cpp	Core dispatch logic for single-node XGMI
InterNodeV1 kernel	src/ops/dispatch_combine/internode_v1.cpp	RDMA + XGMI inter-node dispatch/combine
AsyncLL kernel	src/ops/dispatch_combine/low_latency_async.cpp	3-kernel pipelined transfer
JIT pipeline	python/mori/jit/core.py	compile_genco(): .cpp → .hsaco JIT compilation
Tuning configs	python/mori/ops/tuning_configs/	Pre-tuned JSON params per (arch, model, kernel)
IO engine	python/mori/io/engine.py	IOEngine / IOEngineSession Python wrappers
Shmem API	python/mori/shmem/api.py	OpenSHMEM-style Python bindings
IR bitcode	python/mori/ir/bitcode.py	find_bitcode() — locates/JIT-compiles libmori_shmem_device.bc
Device functions	python/mori/ir/ops.py	MORI_DEVICE_FUNCTIONS — ABI metadata for 50+ device functions
CCL collectives	python/mori/ccl/collective.py	SDMA-based All2All, AllGather, AllReduce
Profiler	python/mori/kernel_profiler.py	export_to_perfetto() for MORI-VIZ warp-level traces
HIP driver	python/mori/jit/hip_driver.py	Low-level HIP API via ctypes (hipModuleLoad, etc.)
Precompile	python/mori/jit/precompile.py	MORI_PRECOMPILE=1 AOT compilation path
Pybind entry	src/pybind/pybind.cpp	Python module entry point
NIC env tools	tools/env_check.sh, tools/env_setup.sh	AINIC environment validation and RDMA configuration