Cluster Perf tests - EP benchmarking & RDMA Perf and Cluster Env mapping recommendation tool

benchmark/kernel/ep_bench/README.md

@@ -0,0 +1,145 @@
+# MoRI EP-Bench (Slurm, no torchrun)
+
+This directory provides a Slurm-based launcher for the **MoRI** expert-parallel
+dispatch/combine microbenchmarks on AMD Instinct GPUs.
+
+
+## Benchmarks covered
+
+| Scope     | Mode        | Script                                                          |
+|-----------|-------------|------------------------------------------------------------------|
+| Intranode | Normal (bf16)        | `tests/python/ops/bench_dispatch_combine.py`             |
+| Intranode | Low latency (fp8)    | `tests/python/ops/bench_dispatch_combine.py --dtype fp8_e4m3_fnuz` |
+| Internode | Normal (v1, bf16)    | `examples/ops/dispatch_combine/test_dispatch_combine_internode.py --cmd bench` |
+| Internode | Low latency (v1_ll)  | `examples/ops/dispatch_combine/test_dispatch_combine_internode.py --cmd bench --kernel-type v1_ll` |
+
+Intranode runs trigger automatically when `NNODES == 1`; internode runs trigger
+when `NNODES > 1`. (DeepEP tests are out of scope for this launcher.)
+
+## Prerequisites
+
+- Slurm cluster with AMD Instinct nodes.
+- InfiniBand HCA(s) available on each node (default `mlx5_0`).
+- Docker available on compute nodes.
+- A Docker image that ships MoRI under `/app/mori`. Two options:
+
+  ### Option A (recommended): build the slim MoRI image shipped here
+
+  Built on the latest stable upstream
+  [`vllm/vllm-openai-rocm`](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags)
+  image (which already provides ROCm, PyTorch, and vLLM) with **only MoRI**
+  added on top.
+
+  ```bash
+  cd Primus/benchmark/kernel/ep_bench/docker
+
+  # Defaults: BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0, GPU_ARCH=gfx942,
+  #           MORI_REPO=https://github.com/ROCm/mori.git, MORI_REF=main
+  docker build -f Dockerfile.mori -t primus-mori-ep:latest .
+
+  # MI355 / gfx950 + pinned MoRI ref:
+  docker build -f Dockerfile.mori \
+      --build-arg GPU_ARCH=gfx950 \
+      --build-arg MORI_REF=v0.x.y \
+      -t primus-mori-ep:mi355 .
+
+  # Precompile MoRI JIT kernels at build time (needs a GPU exposed to the build):
+  DOCKER_BUILDKIT=0 docker build -f Dockerfile.mori \
+      --build-arg MORI_PRECOMPILE=1 \
+      -t primus-mori-ep:precompiled .
+  ```
+
+  Build args supported by [`docker/Dockerfile.mori`](./docker/Dockerfile.mori):
+
+  | Build arg          | Default                                  | Purpose                                            |
+  |--------------------|------------------------------------------|----------------------------------------------------|
+  | `BASE_IMAGE`       | `vllm/vllm-openai-rocm:v0.21.0`          | Override to pin a different upstream vLLM ROCm tag |
+  | `GPU_ARCH`         | `gfx942`                                 | Sets `PYTORCH_ROCM_ARCH` for the MoRI install      |
+  | `MORI_REPO`        | `https://github.com/ROCm/mori.git`       | MoRI git remote                                    |
+  | `MORI_REF`         | `main`                                   | Branch / tag / SHA to check out                    |
+  | `MORI_PRECOMPILE`  | `0`                                      | If `1`, JIT-precompile MoRI kernels into the image |
+
+  ### Option B: reuse the MAD `large_ep_benchmark` image
+
+  Heavier image that also includes rocSHMEM + DeepEP + perftest. Build from
+  the MAD repository root:
+
+  ```bash
+  # From the MAD repository root
+  docker build \
+      -f docker/large_ep_benchmark.ubuntu.amd.Dockerfile \
+      -t ep-benchmarking:latest .
+  ```
+
+  See
+  [`MAD/scripts/large-ep-benchmark/README.md`](https://github.com/ROCm/MAD/blob/main/scripts/large-ep-benchmark/README.md)
+  for cluster/Dockerfile customization (GPU arch, NIC arch).
+
+## Environment variables
+
+| Variable            | Required | Default                                | Description                                                  |
+|---------------------|----------|----------------------------------------|--------------------------------------------------------------|
+| `DOCKER_IMAGE`      | yes      | —                                      | Docker image to launch (e.g. `primus-mori-ep:latest` or `ep-benchmarking:latest`). |
+| `NNODES`            | no       | `$SLURM_NNODES`, then `1`              | Number of nodes; drives intranode vs internode branch.       |
+| `PARTITION`         | no       | unset (prefer `-p` on `sbatch`)        | Slurm partition for child `srun` calls.                      |
+| `GPUS_PER_NODE`     | no       | `8`                                    | GPUs per node; forwarded as `GPU_PER_NODE` to MoRI.          |
+| `MASTER_PORT`       | no       | `2373`                                 | Rendezvous port for `dist.init_process_group`.               |
+| `IBDEVICES`         | no       | `mlx5_0`                               | IB HCA(s) for `ROCSHMEM_USE_IB_HCA`.                         |
+| `LOG_DIR`           | no       | `${SLURM_SUBMIT_DIR}/logs`             | Host directory mounted as `/run_logs` inside the container.  |
+| `EXTRA_DOCKER_ARGS` | no       | unset                                  | Extra args appended to `docker run`.                         |
+
+## Usage
+
+### Single node (intranode only)
+
+```bash
+cd Primus/benchmark/kernel/ep_bench
+
+DOCKER_IMAGE=primus-mori-ep:latest \
+    sbatch -N 1 -p <partition> run_slurm.sh
+```
+
+### Multi-node (internode)
+
+```bash
+cd Primus/benchmark/kernel/ep_bench
+
+DOCKER_IMAGE=primus-mori-ep:latest \
+IBDEVICES=mlx5_0 \
+NNODES=4 \
+    sbatch -N 4 -p <partition> -w <node-list> run_slurm.sh
+```
+
+### Overriding the rendezvous port or log directory
+
+```bash
+DOCKER_IMAGE=primus-mori-ep:latest \
+MASTER_PORT=39566 \
+LOG_DIR=/shared/logs/mori-$(date +%Y%m%d) \
+    sbatch -N 4 -p <partition> run_slurm.sh
+```
+
+## Output
+
+- Slurm job stdout/stderr go to the default `slurm-<jobid>.out` (or whatever
+  you configure with `-o` / `-e` on `sbatch`).
+- Per-test logs are written under `${LOG_DIR}` (default
+  `${SLURM_SUBMIT_DIR}/logs`):
+
+  | File                                              | Produced when |
+  |---------------------------------------------------|---------------|
+  | `mori_intranode_bf16.log`                         | `NNODES == 1` |
+  | `mori_intranode_fp8_ll.log`                       | `NNODES == 1` |
+  | `mori_internode_v1_rank<NODE_RANK>.log`           | `NNODES > 1`  |
+  | `mori_internode_v1_ll_rank<NODE_RANK>.log`        | `NNODES > 1`  |
+
+## Files
+
+- [`run_slurm.sh`](./run_slurm.sh) — sbatch entrypoint; pre-pulls the image
+  and `docker run`s the inner script on each node with one task per node.
+- [`run_mori_bench.sh`](./run_mori_bench.sh) — inner script executed inside
+  the container; branches on `NNODES` and invokes MoRI bench scripts directly
+  with plain `python` (no torchrun).
+- [`docker/Dockerfile.mori`](./docker/Dockerfile.mori) — slim image recipe
+  layered on the latest stable `vllm/vllm-openai-rocm`, installing only MoRI
+  (mirrors the MoRI portion of the MAD `large_ep_benchmark` Dockerfile).

benchmark/kernel/ep_bench/docker/Dockerfile.mori

@@ -0,0 +1,58 @@
+# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'}
+###############################################################################
+#
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#
+###############################################################################
+#
+# Slim image for the Primus MoRI EP microbenchmarks.
+#
+# Based on the upstream vLLM ROCm image (which already provides ROCm, PyTorch,
+# and vLLM and MoRI) and inspired by ROCm/MAD's docker/large_ep_benchmark.ubuntu.amd.Dockerfile,
+#
+# Build (from this directory):
+#   docker build \
+#       -f Dockerfile.mori \
+#       --build-arg GPU_ARCH=gfx942 \
+#       -t primus-mori-ep:latest .
+#
+# Use with the Slurm launcher:
+#   DOCKER_IMAGE=primus-mori-ep:latest sbatch -N 1 -p <partition> \
+#       ../run_slurm.sh
+#
+###############################################################################
+
+# Latest stable upstream vLLM ROCm image.
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
+FROM ${BASE_IMAGE}
+
+#
+## The upstream vllm-openai-rocm image sets ENTRYPOINT to launch the vLLM API
+## server; clear it so the container can be driven by our bench scripts.
+ENTRYPOINT []
+CMD ["/bin/bash"]
+#
+WORKDIR /app
+RUN sed -i 's/http/https/g' /etc/apt/sources.list
+
+RUN apt update && \
+   apt install -y \
+    autoconf pkg-config \
+    libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
+    libibverbs-dev ibverbs-utils libtool libboost-all-dev \
+    libgrpc++-dev protobuf-compiler-grpc protobuf-compiler libprotobuf-dev \
+    libaio-dev liburing-dev pybind11-dev ninja-build libgflags-dev \
+    rdma-core infiniband-diags perftest openssh-server \
+    psmisc vim cmake-curses-gui
+
+WORKDIR /app
+
+# not installing mori since its already installed in vllm container.
+RUN pip install tqdm prettytable
+RUN git clone --recursive $(grep '^MORI_REPO:' versions.txt | cut -d' ' -f2) && \
+    cd mori && \
+    git checkout $(grep '^MORI_BRANCH:' /app/versions.txt | cut -d' ' -f2) 
+
+RUN cd /app/mori && git log | head

benchmark/kernel/ep_bench/run_mori_bench.sh

@@ -0,0 +1,121 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Inner script (runs inside the MAD large_ep_benchmark Docker container) that
+# executes MoRI dispatch/combine microbenchmarks WITHOUT torchrun.
+#
+# Both MoRI bench scripts use torch.multiprocessing.spawn internally to spawn
+# one process per GPU, so this script is invoked exactly once per node.
+#
+# For the internode bench, MoRI expects:
+#   RANK        = node rank (0..NNODES-1)
+#   WORLD_SIZE  = number of nodes
+#   MASTER_ADDR = rendezvous host (first node)
+#   MASTER_PORT = rendezvous port
+# These are derived from SLURM_NODEID / SLURM_NNODES / etc. by the outer
+# run_slurm.sh and forwarded as NODE_RANK / NNODES / MASTER_ADDR / MASTER_PORT.
+#
+# Expected environment (set by run_slurm.sh):
+#   NNODES, NODE_RANK, MASTER_ADDR, MASTER_PORT, GPUS_PER_NODE, IBDEVICES
+#
+# Output logs are written under /run_logs (host: LOG_DIR).
+###############################################################################
+
+set -eo pipefail
+
+MASTER_ADDR="${MASTER_ADDR:-localhost}"
+MASTER_PORT="${MASTER_PORT:-2373}"
+NODE_RANK="${NODE_RANK:-0}"
+NNODES="${NNODES:-1}"
+GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
+IBDEVICES="${IBDEVICES:-mlx5_0}"
+
+LOG_DIR="${LOG_DIR_IN_CONTAINER:-/run_logs}"
+mkdir -p "${LOG_DIR}"
+
+HOST_NAME="$(hostname)"
+HOST_IP="$(hostname -I | awk '{print $1}')"
+
+# Resolve default-route interface for Gloo (TCP) rendezvous.
+DEFAULT_IFACE="$(ip route | awk '/^default/ {print $5; exit}')"
+export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${DEFAULT_IFACE}}"
+
+# rocSHMEM / MoRI RDMA env (shared with MAD run_ep_bench.sh defaults).
+export ROCSHMEM_USE_IB_HCA="${ROCSHMEM_USE_IB_HCA:-${IBDEVICES}}"
+export ROCSHMEM_MAX_NUM_CONTEXTS="${ROCSHMEM_MAX_NUM_CONTEXTS:-144}"
+
+MORI_PATH="${MORI_PATH:-/app/mori}"
+if [[ ! -d "${MORI_PATH}" ]]; then
+    echo "[ERROR] MoRI not found at ${MORI_PATH}. Set MORI_PATH or use an image that ships MoRI." >&2
+    exit 2
+fi
+export PYTHONPATH="${MORI_PATH}:${PYTHONPATH:-}"
+
+cat <<EOF
+================================================
+ MoRI EP-Bench (inside container, no torchrun)
+================================================
+  Host             : ${HOST_NAME} (${HOST_IP})
+  NNODES           : ${NNODES}
+  NODE_RANK        : ${NODE_RANK}
+  GPUS_PER_NODE    : ${GPUS_PER_NODE}
+  MASTER_ADDR      : ${MASTER_ADDR}
+  MASTER_PORT      : ${MASTER_PORT}
+  IBDEVICES        : ${IBDEVICES}
+  GLOO_SOCKET_IF   : ${GLOO_SOCKET_IFNAME}
+  MORI_PATH        : ${MORI_PATH}
+  LOG_DIR          : ${LOG_DIR}
+================================================
+EOF
+
+cd "${MORI_PATH}"
+
+if [[ "${NNODES}" -eq 1 ]]; then
+    echo "[Node ${NODE_RANK}] Running MoRI INTRANODE dispatch/combine benchmark (bf16)..."
+    python tests/python/ops/bench_dispatch_combine.py \
+        2>&1 | tee "${LOG_DIR}/mori_intranode_bf16.log"
+
+    echo "[Node ${NODE_RANK}] Running MoRI INTRANODE low-latency benchmark (fp8_e4m3_fnuz)..."
+    python tests/python/ops/bench_dispatch_combine.py --dtype fp8_e4m3_fnuz \
+        2>&1 | tee "${LOG_DIR}/mori_intranode_fp8_ll.log"
+else
+    # Internode: MoRI internode test reads RANK as node-rank and WORLD_SIZE as
+    # number of nodes (it spawns one process per GPU itself), so we just
+    # promote NODE_RANK / NNODES and call python directly.
+    export RANK="${NODE_RANK}"
+    export WORLD_SIZE="${NNODES}"
+    export MASTER_ADDR
+    export MASTER_PORT
+    export GPU_PER_NODE="${GPUS_PER_NODE}"
+
+    INTERNODE_SCRIPT="examples/ops/dispatch_combine/test_dispatch_combine_internode.py"
+
+    export GLOO_SOCKET_IFNAME=$(ip route | grep '^default' | awk '{print $NF}' | head -n 1)
+
+    echo "[Node ${NODE_RANK}] Running MoRI INTERNODE dispatch/combine benchmark (v1, bf16)..."
+    torchrun --nnodes=$NNODES \
+        --node_rank=$NODE_RANK \
+        --nproc_per_node=1 \
+        --master_addr=$MASTER_ADDR \
+        --master_port=$MASTER_PORT \
+        "${INTERNODE_SCRIPT}" --cmd bench \
+        2>&1 | tee "${LOG_DIR}/mori_internode_v1_rank${NODE_RANK}.log"
+
+    sleep 10
+
+    echo "[Node ${NODE_RANK}] Running MoRI INTERNODE low-latency benchmark (v1_ll)..."
+    torchrun --nnodes=$NNODES \
+        --node_rank=$NODE_RANK \
+        --nproc_per_node=1 \
+        --master_addr=$MASTER_ADDR \
+        --master_port=$MASTER_PORT \
+        "${INTERNODE_SCRIPT}" --cmd bench --kernel-type v1_ll \
+        2>&1 | tee "${LOG_DIR}/mori_internode_v1_ll_rank${NODE_RANK}.log"
+fi
+
+echo "[Node ${NODE_RANK}] MoRI EP benchmarks complete. Logs under ${LOG_DIR}/"
+exit 0