Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
scripts/sglang_disagg/moriio_profiling/artifacts/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
Expand Down
130 changes: 130 additions & 0 deletions docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'}
###############################################################################
#
# MIT License
#
# Copyright (c) 2025 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#################################################################################
ARG BASE_DOCKER=lmsysorg/sglang-rocm:v0.5.15.post1-rocm720-mi30x-20260718
FROM $BASE_DOCKER

ARG ENABLE_ROCTX=0
ARG MORI_ROCTX_COMMIT=f7e6ac6863c53821bc7afb91a578cc6ce38fcad0
ARG SGLANG_ROCTX_COMMIT=48ae829f6e47f9348d8bd936b102d4d7a76f2743

RUN sed -i 's|http://|https://|g' /etc/apt/sources.list

ENV PYTHONPATH=$PYTHONPATH:/sgl-workspace/mori:/sgl-workspace/aiter:

ARG GPU_ARCH=gfx942
WORKDIR /sgl-workspace

RUN pip install --upgrade sglang-router

WORKDIR /sgl-workspace/mori

ARG MORI_COMMIT="158c7e8335a0b19b3f1f422ff134d7869252135e"
# Set INSTALL_MORI=1 to build/install MoRI at MORI_COMMIT; any other value skips it.
ARG INSTALL_MORI=1

RUN apt-get update && \
apt-get install -y --no-install-recommends \
git ibverbs-utils libibverbs-dev \
openmpi-bin libopenmpi-dev \
libpci-dev libdw1 locales \
libgrpc-dev libgrpc++-dev libprotobuf-dev protobuf-compiler-grpc \
cmake

COPY scripts/sglang_disagg/moriio_profiling/patches/ /tmp/roctx-patches/

# Upgrade the compatible ROCm 7.2.0 base to 7.2.3 for every build.
# ENABLE_ROCTX only controls marker patches and profiling tools below.
RUN set -eux; \
sed -i 's#repo.radeon.com/rocm/apt/7.2 #repo.radeon.com/rocm/apt/7.2.3 #' /etc/apt/sources.list.d/rocm.list; \
apt-get update; \
apt-get install -y --only-upgrade $(dpkg -l | awk '$1 == "ii" {print $2}' \
| grep -viE '^(amdgpu|libdrm)' \
| grep -iE '^(rocm|hip|hsa|rccl|miopen|comgr|roc|rpp|rdc|amd-smi|composablekernel|tensile)'); \
grep -q '^7\.2\.3' /opt/rocm/.info/version; \
rm -rf /var/lib/apt/lists/*

RUN set -eux; \
if [ "${ENABLE_ROCTX}" = "1" ]; then \
git clone --quiet https://github.com/ROCm/mori.git /tmp/roctx-mori; \
git -C /tmp/roctx-mori checkout --quiet "${MORI_ROCTX_COMMIT}"; \
git -C /tmp/roctx-mori apply --index /tmp/roctx-patches/mori/01-roctx-instrumentation.patch; \
git -C /tmp/roctx-mori diff --cached --name-only "${MORI_ROCTX_COMMIT}" | while IFS= read -r file; do \
mkdir -p "/sgl-workspace/mori/$(dirname "${file}")"; \
cp "/tmp/roctx-mori/${file}" "/sgl-workspace/mori/${file}"; \
done; \
git clone --quiet https://github.com/sgl-project/sglang.git /tmp/roctx-sglang; \
git -C /tmp/roctx-sglang checkout --quiet "${SGLANG_ROCTX_COMMIT}"; \
git -C /tmp/roctx-sglang apply --index /tmp/roctx-patches/sglang/01-roctx-instrumentation.patch; \
git -C /tmp/roctx-sglang diff --cached --name-only "${SGLANG_ROCTX_COMMIT}" | while IFS= read -r file; do \
mkdir -p "/sgl-workspace/sglang/$(dirname "${file}")"; \
cp "/tmp/roctx-sglang/${file}" "/sgl-workspace/sglang/${file}"; \
done; \
echo "Installing complete patched pinned SGLang benchmark package"; \
rm -rf /sgl-workspace/sglang/python/sglang/benchmark; \
cp -a /tmp/roctx-sglang/python/sglang/benchmark /sgl-workspace/sglang/python/sglang/; \
cp -a /tmp/roctx-sglang/python/sglang/benchmark/. /sgl-workspace/sglang/benchmark/; \
rm -rf /tmp/roctx-mori /tmp/roctx-sglang; \
fi

RUN set -eux; \
if [ "${ENABLE_ROCTX}" = "1" ]; then \
rm -rf /sgl-workspace/mori/build/CMakeCache.txt /sgl-workspace/mori/build/CMakeFiles; \
cmake -S /sgl-workspace/mori -B /sgl-workspace/mori/build -G Ninja \
-DUSE_ROCM=ON -DCMAKE_BUILD_TYPE=Release -DWARP_ACCUM_UNROLL=1 \
-DBUILD_SHMEM_DEVICE_WRAPPER=ON -DENABLE_DEBUG_PRINTF=OFF \
-DENABLE_STANDARD_MOE_ADAPT=OFF -DGPU_TARGETS="${GPU_ARCH}" \
-DENABLE_PROFILER=OFF -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF \
-DBUILD_TESTS=OFF -DBUILD_UMBP=ON -DUSE_SPDK=OFF -DWITH_MPI=OFF \
-DBUILD_TORCH_BOOTSTRAP=OFF -DBUILD_XLA_FFI_OPS=OFF -DBUILD_OPS_DEVICE=OFF \
-DMORI_MULTITHREAD_SUPPORT=OFF; \
cmake --build /sgl-workspace/mori/build -j"$(nproc)"; \
for so in application collective io ops pybind shmem; do \
libname="libmori_${so}.so"; \
[ "${so}" = "pybind" ] && libname="libmori_pybinds.so"; \
src="/sgl-workspace/mori/build/src/${so}/${libname}"; \
[ ! -f "${src}" ] || cp "${src}" "/sgl-workspace/mori/python/mori/${libname}"; \
done; \
elif [ "${INSTALL_MORI}" = "1" ]; then \
echo "INSTALL_MORI=1: installing MoRI at ${MORI_COMMIT}" \
&& git checkout main \
&& git fetch origin \
&& git pull origin main \
&& git checkout ${MORI_COMMIT} \
&& pip install -r requirements-build.txt \
&& pip install -e . ; \
else \
echo "ENABLE_ROCTX=${ENABLE_ROCTX}, INSTALL_MORI=${INSTALL_MORI}: skipping MoRI installation"; \
fi

ENV SGLANG_ROCTX=0
ENV MORI_ROCTX=0
ENV MORI_ROCTX_TRANSFER=0

WORKDIR /sgl-workspace

# Display installed packages for verification
RUN pip list
45 changes: 29 additions & 16 deletions scripts/sglang_disagg/README.MD
Original file line number Diff line number Diff line change
Expand Up @@ -20,34 +20,38 @@ This repository contains scripts and documentation to launch PD Disaggregation f
- Access to a shared filesystem for log collection( cluster specific)


## Building the Docker image
Access the Dockerfile located at `docker/sglang_disagg_inference.ubuntu.amd.Dockerfile`.
It uses `lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x` as the base docker image.
## Build and submit
Build the ordinary image from the repository root: `docker build -t <image> -f docker/sglang_disagg_inference.ubuntu.amd.Dockerfile .`
Build the dedicated profiling image from the repository root: `docker build --build-arg ENABLE_ROCTX=1 -t <profile-image> -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile .`
Submit from `scripts/sglang_disagg/`: `DOCKER_IMAGE_NAME=<image> MODEL_NAME=<model> xP=1 yD=1 USE_CX7_NICS=1 sbatch -N2 -n2 run_xPyD_models.slurm`
Use the profiling image with `RUN_PROFILE=1`. Its Dockerfile starts from the compatible public ROCm 7.2.0 SGLang base below and upgrades installed ROCm userspace packages to ROCm 7.2.3.

```bash
docker build -t sglang_disagg_pd_image -f sglang_disagg_inference.ubuntu.amd.Dockerfile .
```

### Build arguments
### Image build arguments (ordinary first; profiling from `BASE_DOCKER`)

| Arg | Default | Description |
|-----|---------|-------------|
| `INSTALL_MORI` | `1` | When `1`, builds and installs MoRI at `MORI_COMMIT`. Set to any other value (e.g. `0`) to skip MoRI installation entirely. |
| `MORI_COMMIT` | `158c7e8335a0b19b3f1f422ff134d7869252135e` | MoRI commit to check out and install when `INSTALL_MORI=1`. |
| `BASE_DOCKER` | `lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x` | Base SGLang ROCm image. |
| `BASE_DOCKER` | `lmsysorg/sglang-rocm:v0.5.15.post1-rocm720-mi30x-20260718` | Compatible public ROCm 7.2.0 SGLang base; its userspace stack is upgraded to ROCm 7.2.3. |
| `GPU_ARCH` | `gfx942` | Target AMD GPU architecture. |
| `ENABLE_ROCTX` | `0` | When `1`, `MORI_ROCTX_COMMIT` and `SGLANG_ROCTX_COMMIT` select pinned instrumentation and instrumented MoRI is rebuilt regardless of `INSTALL_MORI`, including `0`. It does not control the ROCm 7.2.3 upgrade. |

The dedicated profiling build uses ROCm 7.2.3. The ordinary image remains the exact current `develop` Dockerfile. For `EAGER=0` profiling, do not substitute the stock ROCm 7.2.0 image: its CUDA Graph replay and rocprof finalization paths are known to fail.

Examples:

```bash
# Default build: installs MoRI pinned at MORI_COMMIT
docker build -t sglang_disagg_pd_image -f sglang_disagg_inference.ubuntu.amd.Dockerfile .
docker build -t sglang_disagg_pd_image -f docker/sglang_disagg_inference.ubuntu.amd.Dockerfile .

# Skip MoRI (e.g. when using Mooncake KV transfer only)
docker build --build-arg INSTALL_MORI=0 -t sglang_disagg_pd_image -f sglang_disagg_inference.ubuntu.amd.Dockerfile .
docker build --build-arg INSTALL_MORI=0 -t sglang_disagg_pd_image -f docker/sglang_disagg_inference.ubuntu.amd.Dockerfile .

# Pin a different MoRI commit
docker build --build-arg MORI_COMMIT=<sha> -t sglang_disagg_pd_image -f sglang_disagg_inference.ubuntu.amd.Dockerfile .
docker build --build-arg MORI_COMMIT=<sha> -t sglang_disagg_pd_image -f docker/sglang_disagg_inference.ubuntu.amd.Dockerfile .

# Enable ROCTX markers for profiling
docker build --build-arg ENABLE_ROCTX=1 -t <profile-image> -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile .
```

## Scripts
Expand All @@ -58,22 +62,25 @@ docker build --build-arg MORI_COMMIT=<sha> -t sglang_disagg_pd_image -f sglang_d
| `sglang_disagg_mori_io_ep.sh` | Container entrypoint — starts prefill/decode servers, proxy, and benchmark |
| `models.yaml` | Model-specific CLI flags for all supported models |
| `mori_ep_env.sh` | RDMA/NCCL/Gloo environment variables |
| `benchmark_xPyD.sh` | Concurrency sweep benchmark using sglang bench_serving |
| `benchmark_xPyD.sh` | Ordinary current-`develop` concurrency sweep benchmark (used when profiling is off) |
| `benchmark_xPyD_profile.sh` | Profiling concurrency sweep with deterministic RIDs, client CSV/manifests, and multi-sweep support |
| `benchmark_parser.py` | Log parser for CONCURRENCY benchmark logs |

## Quick Start

```bash
git clone https://github.com/ROCm/MAD.git
cd scripts/sglang_disagg
cd MAD/scripts/sglang_disagg

export DOCKER_IMAGE_NAME=<DOCKER_IMAGE_NAME>
export xP=1
export yD=1
export MODEL_NAME=Llama-3.1-8B-Instruct
export RUN_MORI=1 # MoRI (default). Set RUN_MORI=0 for Mooncake (KV_TRANSFER_BACKEND=mooncake)
export RUN_MORI=1 # Explicit MoRI; RUN_PROFILE=1 defaults to MoRI, while normal runs default to Mooncake.
export USE_CX7_NICS=1

# num_nodes = xP + yD
# CX7 requires an appropriate same-rail node allocation.
sbatch -N 2 -n 2 --nodelist=<node1,node2> run_xPyD_models.slurm
```

Expand All @@ -97,7 +104,7 @@ The unified launcher (`sglang_disagg_mori_io_ep.sh`) selects the disaggregation
| `0` (default) | `tp` | `base_flags` + `tp_flags` + `prefill.tp` / `decode.tp` | All models |
| `1` | `dp` | `base_flags` + `dp_flags` (`--moe-a2a-backend mori`, DP attention) + `prefill.dp` / `decode.dp` | DeepSeek-V3, DeepSeek-R1 only |

- `DP_MODE=1` enables MoRI expert parallelism (EP) with DP attention and currently requires `RUN_MORI=1` (MoRI IO).
- `DP_MODE=1` enables MoRI expert A2A with DP attention; `RUN_MORI` independently selects the disaggregation KV-transfer backend.
- It is restricted by an allowlist (`MORI_DP_MODE1_ALLOWED_MODELS`) enforced in both `run_xPyD_models.slurm` and `sglang_disagg_mori_io_ep.sh`. Launching any other model with `DP_MODE=1` exits with an error — use `DP_MODE=0` (TP) for all non-DeepSeek models.

## Log Files
Expand All @@ -113,6 +120,12 @@ Logs are written to `${LOG_PATH}/${SLURM_JOB_ID}/`:

## Benchmarking

`RUN_PROFILE=1` selects `benchmark_xPyD_profile.sh`; profiling off selects the exact current-`develop` `benchmark_xPyD.sh`. End-to-end profile submissions default to `SKIP_WARMUP=1`; set `SKIP_WARMUP=0` to run the legacy 1024-prompt/512-concurrency warmup. `BENCHMARK_NUM_PROMPTS` overrides the generated prompt count per sweep. The image's `SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX` feature defaults to `0` (off) in this launcher; set it explicitly to `1` to opt in.

`RUN_PROFILE=1` enables MoRI by default; profiling off defaults to Mooncake unless `RUN_MORI` is explicitly set. `EAGER=0` avoids a global CUDA-graph disable and preserves supported graphs, but SGLang/model compatibility may still disable prefill graphs.

Set `RUN_PROFILE=1` when submitting with the profiling image to enable integrated profiling (rocprof capture + kernel analysis). Profile benchmark requests receive deterministic RIDs and produce client timing/manifests without duplicate probe traffic. Multi-sweep runs use space-separated `ISL/OSL` entries in `BENCHMARK_COMBINATIONS`; a separate nonempty `BENCHMARK_CON` is required for profiling. `RUN_MORI=1` adds MoRI markers and strict request/KV correlation; `RUN_MORI=0` produces SGLang-only traces and reqstats without MoRI mapping artifacts. Output artifacts land under `moriio_profiling/artifacts/pull_<jobid>/`.

Parse benchmark results:
```bash
python3 benchmark_parser.py <log_path>/benchmark_XXX_CONCURRENCY.log
Expand Down
119 changes: 119 additions & 0 deletions scripts/sglang_disagg/benchmark_xPyD_profile.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
#!/bin/bash

timestamp=$(date "+%Y%m%d_%H%M%S")
LOG="/run_logs/${SLURM_JOB_ID}/benchmark_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_$MODEL_NAME"
echo "==== Benchmark Serving Concurrency Sweep Test ${LOG} ===== "
echo "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
echo "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null

: "${BENCHMARK_ITR:=1}"
: "${SKIP_WARMUP:=1}"
CON="8 16 32 64 128 256 512"

PROFILE_TRACE_ARGS=()
PROFILE_SWEEP_COUNT=0
if [[ "${RUN_PROFILE:-0}" == "1" ]]; then
read -ra _profile_combinations <<< "${BENCHMARK_COMBINATIONS-1024/1024 8192/1024}"
read -ra _profile_concurrency <<< "${BENCHMARK_CON:-}"
_profile_inputs_valid=1
[[ "$BENCHMARK_ITR" =~ ^[1-9][0-9]*$ && "${#_profile_combinations[@]}" -gt 0 && "${#_profile_concurrency[@]}" -gt 0 ]] || _profile_inputs_valid=0
for profile_combo in "${_profile_combinations[@]}"; do [[ "$profile_combo" =~ ^[1-9][0-9]*/[1-9][0-9]*$ ]] || _profile_inputs_valid=0; done
for profile_con in "${_profile_concurrency[@]}"; do [[ "$profile_con" =~ ^[1-9][0-9]*$ ]] || _profile_inputs_valid=0; done
(( _profile_inputs_valid )) || { echo "ERROR: invalid profile sweep configuration." >&2; exit 2; }
CON="${BENCHMARK_CON}"
PROFILE_SWEEP_COUNT=$((BENCHMARK_ITR * ${#_profile_combinations[@]} * ${#_profile_concurrency[@]}))
declare -A _profile_sweep_keys=()
for ((profile_i=1; profile_i<=BENCHMARK_ITR; profile_i++)); do
for profile_combo in "${_profile_combinations[@]}"; do
IFS="/" read -r profile_isl profile_osl <<< "$profile_combo"
for profile_con in "${_profile_concurrency[@]}"; do
profile_sweep_key="i${profile_i}_isl${profile_isl}_osl${profile_osl}_c${profile_con}"
if [[ -n "${_profile_sweep_keys[$profile_sweep_key]+x}" ]]; then
echo "ERROR: duplicate profile sweep key: $profile_sweep_key" >&2
exit 2
fi
_profile_sweep_keys[$profile_sweep_key]=1
done
done
done
fi

echo "Benchmark config: prompts=${BENCHMARK_NUM_PROMPTS:-auto} combinations=${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024} concurrency=${CON} iterations=${BENCHMARK_ITR} skip_warmup=${SKIP_WARMUP}" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
if [[ "$SKIP_WARMUP" != "1" ]]; then
sleep 60
echo "Test run:" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
python3 -m sglang.bench_serving \
--model $MODEL_PATH \
--backend sglang \
--host 127.0.0.1 \
--port 2322 \
--dataset-name random \
--random-input 1024 \
--random-output 1024\
--random-range-ratio 1.0 \
--max-concurrency 512 \
--num-prompt 1024 \
--pd-separated \
2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null
echo ""
else
echo "Skipping 1024-prompt warmup (SKIP_WARMUP=1)" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
fi
# ISL/OSL combinations — override via BENCHMARK_COMBINATIONS env var (space-separated, e.g. "1024/1024 8192/1024")
IFS=' ' read -ra COMBINATIONS <<< "${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024}"
echo "Benchmarking iterations: $BENCHMARK_ITR" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
for ((i=1; i<=BENCHMARK_ITR; i++)); do
sleep 60
echo "RUNNING: the benchserving script for iter: $i" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
for combo in "${COMBINATIONS[@]}"; do
IFS="/" read -r isl osl <<< "$combo"
for con in $CON; do
p_con=$(($con * 2))
if [ "$p_con" -lt 16 ]; then
p_con=16
fi
if [[ "${RUN_PROFILE:-0}" == "1" ]]; then
p_con="${BENCHMARK_NUM_PROMPTS:-$p_con}"
fi
PROFILE_TRACE_ARGS=()
if [[ "${RUN_PROFILE:-0}" == "1" ]]; then
sweep_key="i${i}_isl${isl}_osl${osl}_c${con}"
if (( PROFILE_SWEEP_COUNT == 1 )); then
request_prefix="profile-${SLURM_JOB_ID}"
artifact_suffix=""
else
request_prefix="profile-${SLURM_JOB_ID}-${sweep_key}"
artifact_suffix="_${sweep_key}"
fi
PROFILE_TRACE_ARGS=(
--request-id-prefix "$request_prefix"
--client-timing-csv "/run_logs/${SLURM_JOB_ID}/rocprof_probe_client${artifact_suffix}.csv"
--client-manifest "/run_logs/${SLURM_JOB_ID}/rocprof_probe_manifest${artifact_suffix}.json"
)
fi
echo "RUNNING: prompts $p_con isl $isl osl $osl con $con" | tee -a ${LOG}_CONCURRENCY.log >/dev/null
python3 -m sglang.bench_serving \
--model $MODEL_PATH \
--backend sglang \
--host 127.0.0.1 \
--port 2322 \
--dataset-name random \
--random-input $isl \
--random-output $osl \
--random-range-ratio 1.0 \
--max-concurrency $con \
--num-prompt $p_con \
--pd-separated \
"${PROFILE_TRACE_ARGS[@]}" \
2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null

sleep 10
done
done
done


python3 parse_to_csv.py ${LOG}_CONCURRENCY.log -o ${LOG}_CONCURRENCY.csv \
--perf-csv /run_logs/${SLURM_JOB_ID}/perf.csv \
--model-name "${MODEL_NAME}" \
2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null
Loading