From 033b3a30f91d17e641611693057da48fee952e46 Mon Sep 17 00:00:00 2001 From: Rai Date: Mon, 3 Aug 2026 15:49:37 -0500 Subject: [PATCH] sglang profiling --- .gitignore | 1 + ...gg_inference_profile.ubuntu.amd.Dockerfile | 130 + scripts/sglang_disagg/README.MD | 45 +- .../sglang_disagg/benchmark_xPyD_profile.sh | 119 + .../sglang_disagg/moriio_profiling/README.md | 284 ++ .../sglang_disagg/moriio_profiling/SKILL.md | 263 ++ .../sglang_disagg/moriio_profiling/hooks.sh | 190 ++ .../mori/01-roctx-instrumentation.patch | 430 +++ .../sglang/01-roctx-instrumentation.patch | 982 ++++++ .../moriio_profiling/process_kernels.sh | 443 +++ .../moriio_profiling/trace_tools.py | 2706 +++++++++++++++++ scripts/sglang_disagg/run_xPyD_models.slurm | 54 +- .../sglang_disagg/sglang_disagg_mori_io_ep.sh | 88 +- 13 files changed, 5702 insertions(+), 33 deletions(-) create mode 100644 docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile create mode 100644 scripts/sglang_disagg/benchmark_xPyD_profile.sh create mode 100644 scripts/sglang_disagg/moriio_profiling/README.md create mode 100644 scripts/sglang_disagg/moriio_profiling/SKILL.md create mode 100644 scripts/sglang_disagg/moriio_profiling/hooks.sh create mode 100644 scripts/sglang_disagg/moriio_profiling/patches/mori/01-roctx-instrumentation.patch create mode 100644 scripts/sglang_disagg/moriio_profiling/patches/sglang/01-roctx-instrumentation.patch create mode 100644 scripts/sglang_disagg/moriio_profiling/process_kernels.sh create mode 100644 scripts/sglang_disagg/moriio_profiling/trace_tools.py diff --git a/.gitignore b/.gitignore index e2d872c7..323e043c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +scripts/sglang_disagg/moriio_profiling/artifacts/ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..3570ce00 --- /dev/null +++ b/docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile @@ -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 diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index 0540b19d..cdf6cb0d 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -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 -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 -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile .` +Submit from `scripts/sglang_disagg/`: `DOCKER_IMAGE_NAME= MODEL_NAME= 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= -t sglang_disagg_pd_image -f sglang_disagg_inference.ubuntu.amd.Dockerfile . +docker build --build-arg MORI_COMMIT= -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 -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile . ``` ## Scripts @@ -58,22 +62,25 @@ docker build --build-arg MORI_COMMIT= -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= 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= run_xPyD_models.slurm ``` @@ -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 @@ -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_/`. + Parse benchmark results: ```bash python3 benchmark_parser.py /benchmark_XXX_CONCURRENCY.log diff --git a/scripts/sglang_disagg/benchmark_xPyD_profile.sh b/scripts/sglang_disagg/benchmark_xPyD_profile.sh new file mode 100644 index 00000000..79a90dd0 --- /dev/null +++ b/scripts/sglang_disagg/benchmark_xPyD_profile.sh @@ -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 diff --git a/scripts/sglang_disagg/moriio_profiling/README.md b/scripts/sglang_disagg/moriio_profiling/README.md new file mode 100644 index 00000000..40ee5c03 --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/README.md @@ -0,0 +1,284 @@ +# Integrated SGLang MoRI I/O profiling + +This directory documents the supported profiling path for disaggregated SGLang. +`scripts/sglang_disagg/run_xPyD_models.slurm` launches +`sglang_disagg_mori_io_ep.sh` across `xP + yD` nodes. With `RUN_PROFILE=1`, +the prefill and decode server commands are wrapped once by rocprofv3; with +profiling off, their normal launch path is unchanged. `RUN_PROFILE=1` selects +`benchmark_xPyD_profile.sh`; profiling off selects the exact current-`develop` +`benchmark_xPyD.sh`. Profile benchmark requests receive deterministic request +IDs (RIDs), so no second profiling/probe request set is sent. After server finalization, +`moriio_profiling/process_kernels.sh` verifies capture completeness and creates +RID-selected traces, ReqTimeStats joins, optional MoRI maps, and kernel buckets. + +> [!WARNING] +> Start with a **small output sequence length (OSL), preferably `8`**, for smoke +> and validation runs. Do not begin with OSL 1024 or another large value under +> full kernel + marker tracing. Trace volume, rocprof capture/finalization time, +> JSON size, NFS pressure, and analysis cost grow with generated tokens, request +> count, concurrency, and the number of sweeps. A practical first run is ISL +> `64` or `128` (or `1024` when that context is required), OSL `8`, concurrency +> `1`, 2-8 prompts, and one or a few sweeps. Increase one dimension at a time +> only after strict validation succeeds. In `BENCHMARK_COMBINATIONS="ISL/OSL"`, +> OSL is the second component. Plan disk and wall time, and monitor `/run_logs` +> plus shared artifact capacity throughout capture and finalization. + +## Build and image contract + +Build from the repository root on a compute node, not a disk-constrained login +node, and push the resulting image where every allocated node can pull it: + +```bash +docker build --build-arg ENABLE_ROCTX=1 -t \ + -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile . +docker push +``` + +Pass that exact image with `DOCKER_IMAGE_NAME` at submission. Rebuild when the +Dockerfile, pinned SGLang content, or either SGLang/MoRI instrumentation patch +changes. Shell/Python orchestration-only edits that are supplied by the +host-mounted repository may not require an image rebuild, but verify the mount +path and image compatibility before relying on that. Do not treat a previously +validated tag as universally applicable. + +## Copy-paste submissions + +Run these from `scripts/sglang_disagg/` with `DOCKER_IMAGE_NAME=`. +For the CX7 cluster path, set `USE_CX7_NICS=1` on appropriate same-rail nodes; `RUN_PROFILE=1` requires nonempty `BENCHMARK_CON`. +Profile submissions default to `SKIP_WARMUP=1`; set `SKIP_WARMUP=0` to run +the legacy warmup. The examples also skip the optional curl smoke test after +proxy readiness. + +Minimal single sweep (Llama, MoRI, per-node TP8): + +```bash +DOCKER_IMAGE_NAME= \ +MODEL_NAME=Llama-3.1-8B-Instruct \ +xP=1 yD=1 DP_MODE=0 \ +RUN_PROFILE=1 RUN_MORI=1 \ +USE_CX7_NICS=1 \ +SKIP_CURL_TEST=1 \ +BENCHMARK_ITR=1 \ +BENCHMARK_COMBINATIONS="64/8" \ +BENCHMARK_CON="1" \ +BENCHMARK_NUM_PROMPTS=8 \ +sbatch -N2 -n2 --gres=gpu:8 --partition=amd-rccl \ + run_xPyD_models.slurm +``` + +Multi-sweep example (Llama, MoRI, per-node TP8): + +```bash +DOCKER_IMAGE_NAME= \ +MODEL_NAME=Llama-3.1-8B-Instruct \ +xP=1 yD=1 DP_MODE=0 \ +RUN_PROFILE=1 RUN_MORI=1 \ +USE_CX7_NICS=1 \ +SKIP_CURL_TEST=1 \ +BENCHMARK_ITR=2 \ +BENCHMARK_COMBINATIONS="64/8 128/8" \ +BENCHMARK_CON="1 2" \ +BENCHMARK_NUM_PROMPTS=2 \ +sbatch -N2 -n2 --gres=gpu:8 --partition=amd-rccl \ + run_xPyD_models.slurm +``` + +`BENCHMARK_NUM_PROMPTS` is the request count **per sweep**. When it is omitted, +profiling retains `benchmark_xPyD_profile.sh`'s derived `p_con` behavior: +`max(2 * concurrency, 16)`. + +## Profiling and MoRI behavior + +`RUN_PROFILE=1` enables MoRI by default when `RUN_MORI` is unset; set +`RUN_MORI=0` to opt out and use the Mooncake non-MoRI profiling backend. The +launcher applies the following behavior: + +- Both `RUN_MORI=0` and `RUN_MORI=1` wrap server workers with one continuous + rocprofv3 capture, set `SGLANG_ROCTX=1`, `REQ_TIME_STATS=1`, + `ROCPROF_FLAGS="--kernel-trace --marker-trace"`, and + `ROCPROF_DIR_BASE=/run_logs`, then run the benchmark and postprocessor. +- `RUN_MORI=1` additionally enables MoRI marker/transfer instrumentation and + requires complete request-to-MoRI KV mapping. +- `RUN_MORI=0` uses the Mooncake transfer backend, omits the request/MoRI map, + and emits reqstats with MoRI fields zeroed. SGLang traces and strict client + RID/ReqTimeStats joins remain required. +- `REQ_TIME_STATS=1` adds `--enable-request-time-stats-logging` to both servers. + The engine logs contain `ReqTimeStats(...)` records, and rocprof marker CSVs + contain SGLang request-stage markers used for trace lanes and joins. +- `EAGER=1` appends `--disable-cuda-graph` to both prefill and decode server + commands; `EAGER=0` preserves supported graphs, though SGLang/model compatibility may still disable prefill graphs. + +The launcher derives topology and tensor-parallel settings. `xP + yD` is the +required node count. With `DP_MODE=0`, each node is an independent TP server +(default TP8); with allowed `DP_MODE=1` models, role TP/DP/EP sizes scale from +role node count and GPUs per node. Set `xP`, `yD`, and `DP_MODE`; do not assume an +arbitrary user-supplied TP command bypasses this derivation. + +## Multi-sweep contract + +Profiling supports multiple `BENCHMARK_ITR` iterations, space-separated +`BENCHMARK_COMBINATIONS`, and space-separated `BENCHMARK_CON` values. Every +iteration x shape x concurrency point receives this deterministic key: + +```text +i_isl_osl_c +``` + +For more than one sweep: + +- Each sweep has a unique RID prefix: RIDs are + `profile---NNN` (`NNN` starts at `000`). +- Each sweep has its own raw client CSV and manifest under `/run_logs//` + (host default: `/shared_inference//model_blog_logs//`): + `rocprof_probe_client_.csv` and + `rocprof_probe_manifest_.json`. +- Strict per-sweep outputs are isolated under + `artifacts/pull_/sweeps//`. +- The server workers remain inside one continuous rocprof capture. RID filters + isolate trace/request analysis per sweep; kernel bucket analysis covers the + capture by pooling all verified local worker CSVs per node before grouping and trimming; durations/counts are summed, not averaged. +- Duplicate keys, mixed fixed/keyed artifacts, orphan or empty CSV/manifest + pairs, and duplicate prefixes among discovered pairs are rejected; a wholly absent pair is not discoverable automatically, so users/orchestration must verify the expected `iterations x shapes x concurrencies` count. + +For exactly one sweep, backward-compatible names and layout remain: +`profile--NNN`, `rocprof_probe_client.csv`, +`rocprof_probe_manifest.json`, and outputs directly under +`artifacts/pull_/`. + +The client CSV records per-RID send/first-token/completion timing. The manifest +records the expected RID set and benchmark window. They are not required to +produce raw server rocprof files, but the integrated postprocessor +requires both so it can enforce CSV/manifest RID equality, ReqTimeStats +completeness, and (when enabled) exact MoRI mapping. + +### Correlation and trace invariants + +- Filtered MoRI correlation is sweep-scoped. When `rid_rooms` is active, skip + workers with no target RID bounds, reject mappings outside the selected rooms, + and emit only selected rooms. Keep PID/file-scoped UID resolution and + unfiltered behavior unchanged. This prevents cross-sweep contamination. +- In DP-attention, empty MoRI lanes can be expected: one DP owner handles a + request and its KV transfer while other EP workers can still execute kernels. + Track labels such as `TP` identify worker/global ranks; ROCTX markers are + CPU-process ranges, not proof that every worker owns the request. Validate + expected versus observed bytes and transfer post/completion pairing before + diagnosing capture loss. +- Concurrent engine ranges in one PID can limit per-request attribution when an + engine marker has no transfer UID. Transfer posts/completions and byte + accounting remain the authoritative correlation checks. +- Fully parse and validate every clean trace JSON. If it is malformed, + truncated, or NUL-tailed, rebuild it from source captures; file existence + alone is not successful validation. + +## Capture, finalization, and post-processing + +`hooks.sh` finalizes scheduler workers before the launch-server parent, verifies +per-node output counts, and uses a node completion barrier. Hook-level defaults are `ROCPROF_FINALIZE_TIMEOUT=1800`, `ROCPROF_STALL_LIMIT=45`, and +`ROCPROF_NODE_BARRIER_TIMEOUT=2100` seconds. If overriding them, ensure the +variables are actually propagated into the container; the normal submission +wrapper does not forward these three variables explicitly. + +`roctx_finalize_workers` may report that it could not find the process-scoped +rocprof directory and use its job/topology/mtime fallback. Treat that as +informational only when all expected worker files and strict post-processing +validate. NFS reads may transiently expose incomplete-looking files; retry the +read, validate JSON/checksums, and use copy-to-temp plus atomic rename when +regenerating artifacts. Never normalize a missing file, invalid JSON, +checksum change, incomplete worker set, or failed strict join as benign. + +The integrated launcher runs: + +```bash +moriio_profiling/process_kernels.sh +``` + +Manual forms are: + +```bash +moriio_profiling/process_kernels.sh # alias for run +moriio_profiling/process_kernels.sh run +moriio_profiling/process_kernels.sh verify +moriio_profiling/process_kernels.sh trace +moriio_profiling/process_kernels.sh analyze +``` + +`verify` requires the same expected worker PID set across each node's +`*_kernel_trace.csv`, `*_marker_api_trace.csv`, and `*_results.json` files. +`trace` builds each RID-selected clean trace, reqstats, client copy, and optional +MoRI map in staging before finalizing required outputs. `analyze` builds +best-effort kernel buckets for each raw capture directory. `run` performs strict +`trace` followed by best-effort `analyze`; a best-effort kernel-analysis warning +does not weaken strict capture/request validation. + +The default integrated artifact root is exactly: + +```text +scripts/sglang_disagg/moriio_profiling/artifacts/pull_/ +``` + +Single-sweep outputs live at that root; multi-sweep request outputs repeat under +`sweeps//`: + +```text +roctx_mori_clean_prefill_decode_.json +roctx_mori_clean_probe_only_.json +reqstats_per_request_.csv +reqstats_per_request__prefill.csv +reqstats_per_request__decode.csv +rocprof_probe_client[_].csv +rocprof_probe_manifest[_].json +request_mori_map_.csv # RUN_MORI=1 only +request_mori_map_.md # RUN_MORI=1 only +``` + +The two compatibility clean traces are copied from the same RID-selected trace +and are byte-identical by design; `probe_only` does not represent +separate traffic. + +Per-node continuous-capture analysis is under: + +```text +analyze_phase/rocprof__NODE/buckets/ +``` + +Expected bucket files include `kernel_summary_normalized.csv`, +`kernel_summary_trimmed.csv`, `perkernel_buckets.csv`, and +`bycat_buckets.csv`. The clean traces are Perfetto/Chrome trace-event JSON, not +pftrace binaries. + +## Offline CLI + +`trace_tools.py` exposes: + +```bash +python3 moriio_profiling/trace_tools.py build-trace --help +python3 moriio_profiling/trace_tools.py correlate --help +python3 moriio_profiling/trace_tools.py reqstats --help +python3 moriio_profiling/trace_tools.py buckets --help +python3 moriio_profiling/trace_tools.py trimmed-summary --help +python3 moriio_profiling/trace_tools.py analyze --help +python3 moriio_profiling/trace_tools.py self-test-categories +``` + +The integrated path uses `--rid-prefix`, worker-count validation, +`--require-data`, `--require-client`, and, for MoRI, `--require-complete`. +`--probe-only` and its time-gap behavior exist only for legacy artifacts; new +runs select normal benchmark requests by RID. The CLI has no llmscope runtime +dependency, and raw CSV analysis does not require TraceLens. + +## Troubleshooting checklist + +1. Confirm `DOCKER_IMAGE_NAME` pulls on every allocated node and contains the required + instrumentation. +2. Confirm `MODEL_NAME` is allowed and one model path is available on all nodes. +3. Allocate at least `xP + yD` nodes and verify each launcher-selected rendezvous IP is reachable on the intended network/interface and matches the allocation. +4. Check prefill/decode readiness logs, router registration, and readiness + timeouts before blaming profiling. +5. Verify every expected GPU worker produced matching kernel, marker, and + results files; then check finalization messages and `.profile_done_NODE*` + barrier completion. +6. For multi-sweep runs, require + `iterations x shapes x concurrencies` sweep directories, unique RIDs, and + the expected request count in every CSV/manifest/reqstats join. +7. If capture size or finalization is unhealthy, return to OSL `8`, concurrency + `1`, 2-8 prompts, and one sweep before increasing one dimension at a time. diff --git a/scripts/sglang_disagg/moriio_profiling/SKILL.md b/scripts/sglang_disagg/moriio_profiling/SKILL.md new file mode 100644 index 00000000..2368a10b --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/SKILL.md @@ -0,0 +1,263 @@ +--- +name: sglang-moriio-profiling +description: Runs and validates integrated MoRI I/O profiling with rocprofv3/ROCTX for disaggregated SGLang xPyD jobs. Use when building a profiling image, submitting RUN_PROFILE jobs, validating single- or multi-sweep artifacts, or troubleshooting capture and request attribution. +--- + +# Integrated SGLang MoRI I/O profiling + +## When to use + +Use this skill for the supported `RUN_PROFILE=1` path through +`scripts/sglang_disagg/run_xPyD_models.slurm`, +`sglang_disagg_mori_io_ep.sh`, `benchmark_xPyD_profile.sh`, and +`moriio_profiling/process_kernels.sh`. Profiling off uses `benchmark_xPyD.sh`, +restored byte-for-byte from current `develop`. Do not invent a separate +submit/probe/capture workflow. + +Success means: + +- every expected prefill/decode GPU worker has kernel, marker, and results files; +- every expected sweep has unique deterministic RIDs and complete client, + ReqTimeStats, and optional MoRI joins; +- strict trace processing passes; and +- artifacts are copied with integrity checks without weakening validation. + +> [!WARNING] +> Keep OSL small. Start with OSL `8`, concurrency `1`, 2-8 prompts, and one or a +> few sweeps. OSL is the second value in +> `BENCHMARK_COMBINATIONS="ISL/OSL"`. Do not start with OSL 1024 under full +> kernel + marker tracing: generated tokens, requests, concurrency, and sweep +> count multiply trace size, finalization time, JSON/NFS load, and analysis +> cost. Use ISL `64`/`128`, or `1024` only when required, then increase one +> dimension at a time after a clean run. Check `/run_logs` and shared artifact +> capacity before and during the job. + +## Non-negotiable constraints + +- Preserve normal/develop launch behavior when `RUN_PROFILE=0`. +- Do not restore duplicate profiling probe traffic. Profile the normal + benchmark requests, which already receive deterministic RIDs. +- Do not silently weaken worker-count, RID, ReqTimeStats, client, JSON, hash, or + MoRI completeness checks. +- Do not delete client CSV/manifest artifacts unless their strict validation + role is replaced end-to-end. +- Do not modify unrelated files or discard another user's working-tree changes. +- Do not call a partial capture successful. Best-effort kernel analysis does + not excuse a failed strict trace/request phase. + +## Workflow + +### 1. Preflight + +1. Work from `scripts/sglang_disagg/` and inspect the current working tree. +2. Confirm `MODEL_NAME` is accepted and its model directory is available on all + allocated nodes. +3. Set `xP` and `yD`; allocate at least `xP + yD` nodes. The launcher derives + role topology and TP/DP/EP settings. For `DP_MODE=0`, each role node is an + independent TP server (default TP8). Use `DP_MODE=1` only for allowlisted + models. + For the CX7 cluster path, set `USE_CX7_NICS=1` on appropriate same-rail nodes. +4. Confirm the image can be pulled on every selected node. +5. `RUN_PROFILE=1` enables MoRI by default when `RUN_MORI` is unset. Override it + explicitly only when needed: + - `RUN_MORI=1`: MoRI markers plus strict request/KV mapping. + - `RUN_MORI=0`: Mooncake backend, SGLang traces/reqstats, zeroed MoRI reqstats + columns, and no request/MoRI map. +6. Profile submissions default to `SKIP_WARMUP=1`. Set `SKIP_WARMUP=0` only + when the legacy warmup is wanted. Decide separately whether to set + `SKIP_CURL_TEST=1`. +7. Set nonempty `BENCHMARK_CON` (mandatory under `RUN_PROFILE=1`), then start with OSL `8`, low concurrency and prompts, and minimal sweeps. + +### 2. Decide whether to rebuild + +Build from the repository root on a compute node and push the image: + +```bash +docker build --build-arg ENABLE_ROCTX=1 -t \ + -f docker/sglang_disagg_inference_profile.ubuntu.amd.Dockerfile . +docker push +``` + +Rebuild after Dockerfile, pinned SGLang, or SGLang/MoRI instrumentation-patch +changes. Host-mounted shell/Python orchestration-only changes may not require a +rebuild; verify the actual mount and image compatibility. Never hard-code a previously validated +image as universally applicable. Submit with `DOCKER_IMAGE_NAME`. + +### 3. Submit + +Minimal single sweep: + +```bash +DOCKER_IMAGE_NAME= \ +MODEL_NAME=Llama-3.1-8B-Instruct \ +xP=1 yD=1 DP_MODE=0 \ +RUN_PROFILE=1 RUN_MORI=1 \ +USE_CX7_NICS=1 \ +SKIP_CURL_TEST=1 \ +BENCHMARK_ITR=1 \ +BENCHMARK_COMBINATIONS="64/8" \ +BENCHMARK_CON="1" \ +BENCHMARK_NUM_PROMPTS=8 \ +sbatch -N2 -n2 --gres=gpu:8 --partition=amd-rccl \ + run_xPyD_models.slurm +``` + +Multi-sweep example: + +```bash +DOCKER_IMAGE_NAME= \ +MODEL_NAME=Llama-3.1-8B-Instruct \ +xP=1 yD=1 DP_MODE=0 \ +RUN_PROFILE=1 RUN_MORI=1 \ +USE_CX7_NICS=1 \ +SKIP_CURL_TEST=1 \ +BENCHMARK_ITR=2 \ +BENCHMARK_COMBINATIONS="64/8 128/8" \ +BENCHMARK_CON="1 2" \ +BENCHMARK_NUM_PROMPTS=2 \ +sbatch -N2 -n2 --gres=gpu:8 --partition=amd-rccl \ + run_xPyD_models.slurm +``` + +`BENCHMARK_NUM_PROMPTS` is per sweep. If omitted, profiling uses +`max(2 * concurrency, 16)` from `benchmark_xPyD_profile.sh`. + +`RUN_PROFILE=1` sets SGLang ROCTX, `REQ_TIME_STATS=1`, rocprof kernel + marker +tracing, and `/run_logs` output. It wraps each server once and reuses that +continuous capture across all sweeps. `EAGER=1` is a separate opt-in that adds +`--disable-cuda-graph` to both role commands; `EAGER=0` preserves supported graphs, though SGLang/model compatibility may still disable prefill graphs. + +### 4. Monitor + +Monitor scheduler state, the Slurm output/error files, and +`/shared_inference//model_blog_logs//` (mounted as +`/run_logs/`). Check, in order: + +1. model validation and launcher-selected rendezvous IP reachability on the intended network/interface; +2. prefill/decode server readiness and router registration; +3. benchmark sweep progress; +4. rocprof worker finalization and `.profile_done_NODE*` barrier progress; and +5. strict `process_kernels.sh` results. + +A job/topology/mtime fallback message from `roctx_finalize_workers` can be +benign only if every expected worker artifact and all later strict checks pass. +Never classify missing files, invalid JSON, changed checksums, or incomplete +joins as benign. + +### 5. Validate capture and sweeps + +The launcher invokes the postprocessor after profiling. For manual diagnosis: + +```bash +moriio_profiling/process_kernels.sh verify +moriio_profiling/process_kernels.sh trace +moriio_profiling/process_kernels.sh analyze +moriio_profiling/process_kernels.sh run +``` + +`process_kernels.sh ` is also `run`. + +For every role node, require the same expected PID set across: + +```text +*_kernel_trace.csv +*_marker_api_trace.csv +*_results.json +``` + +For multiple iterations, shapes, or concurrencies, compute the expected count: + +```text +BENCHMARK_ITR * number_of_shapes * number_of_concurrencies +``` + +Each sweep key is `i_isl_osl_c`. Require that +many directories under: + +```text +moriio_profiling/artifacts/pull_/sweeps// +``` + +Each sweep must have its own RID prefix, client CSV, and manifest. RIDs are +`profile---NNN`; raw client files are +`rocprof_probe_client_.csv` and +`rocprof_probe_manifest_.json`. Reject duplicate keys/prefixes, +mixed fixed/keyed files, orphan/empty pairs, duplicate RIDs, or request-count mismatches among discovered pairs; a wholly absent pair is not discoverable, so users/orchestration must enforce the expected count. + +Exactly one sweep retains `profile--NNN`, fixed client filenames, and +outputs directly under `artifacts/pull_/`. + +Per sweep, require non-empty: + +```text +roctx_mori_clean_prefill_decode_.json +roctx_mori_clean_probe_only_.json +reqstats_per_request_.csv +reqstats_per_request__prefill.csv +reqstats_per_request__decode.csv +rocprof_probe_client[_].csv +rocprof_probe_manifest[_].json +request_mori_map_.csv/.md # only with RUN_MORI=1 +``` + +The two compatibility JSON traces are byte-identical copies by design; +they do not represent separate traffic. The CSV/manifest are not fundamental +to raw rocprof capture, but are mandatory for the integrated strict client RID +and ReqTimeStats checks. ReqTimeStats markers must remain present. + +Continuous-capture kernel outputs use one verified local-rank-0 worker per node/role and are under +`artifacts/pull_/analyze_phase/rocprof__NODE/buckets/`, including +`kernel_summary_normalized.csv`, `kernel_summary_trimmed.csv`, +`perkernel_buckets.csv`, and `bycat_buckets.csv` when analysis succeeds. + +#### Correlation and trace invariants + +- With filtered MoRI correlation (`rid_rooms` active), skip workers that have no + target RID bounds, reject mappings outside the selected rooms, and emit only + selected rooms. Preserve PID/file-scoped UID resolution and the unfiltered + path. This is required to prevent cross-sweep contamination. +- Under DP-attention, an empty MoRI lane can be expected: one DP owner handles a + request and its KV transfer while other EP workers can still execute kernels. + Track labels such as `TP` are worker/global ranks, while ROCTX markers are + CPU-process ranges. Check expected/observed bytes and transfer post/completion + pairing before classifying an empty lane as capture loss. +- Concurrent engine ranges in the same PID can limit per-request attribution + when an engine marker lacks a transfer UID. Transfer posts/completions and + byte accounting remain authoritative. +- Fully parse and validate clean trace JSON. If it is malformed, truncated, or + NUL-tailed, rebuild it from source captures instead of accepting its existence. + +### 6. Copy artifacts + +Copy only after strict validation. Preserve the entire job-scoped artifact root +and the raw client files. For NFS or cross-filesystem copies: + +1. copy to a temporary destination; +2. compare source/destination hashes and parse JSON; +3. verify sweep and request counts again; and +4. atomically rename the temporary destination when the filesystem permits. + +Transient NFS holes warrant a retry and integrity recheck, not relaxed +validation or deletion of the source. + +## Failure policy + +> [!IMPORTANT] +> A failed run can reflect a bad compute node or hardware/driver state rather +> than profiling code. Before modifying profiling code, inspect the first causal +> errors; validate `rocminfo`, `/dev/kfd`, GPU/RAS state, +> `torch.cuda.is_available()`, and Slurm node health; and, when practical, +> reproduce with profiling disabled or the base image on the same node. Exclude +> a confirmed bad node and notify cluster administrators instead of patching +> profiler code. + +Stop and report the first failing phase with its job ID, node/role, sweep key, +expected versus observed counts, and relevant paths. Check image pulls, model +paths, xP+yD allocation, launcher-selected rendezvous IPs, server/router timeouts, per-GPU +capture triplets, and finalization barriers. For multi-sweep failures, compare +the directory count to iteration x shape x concurrency and check unique RIDs +and per-sweep request counts. Reduce to OSL `8`, concurrency `1`, 2-8 prompts, +and one sweep before scaling again. Do not edit source, suppress errors, or +manually fabricate missing artifacts to make validation pass. + +See [README.md](README.md) for the full artifact and CLI reference. diff --git a/scripts/sglang_disagg/moriio_profiling/hooks.sh b/scripts/sglang_disagg/moriio_profiling/hooks.sh new file mode 100644 index 00000000..6d4059d8 --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/hooks.sh @@ -0,0 +1,190 @@ +#!/usr/bin/env bash +# Scheduler-first rocprof finalization hooks for the SGLang entrypoint. + +_rocprof_prefix() { + local role="$1" + if [[ "${ROCPROF:-0}" != "1" ]]; then echo ""; return 0; fi + local _rocprof_base="${ROCPROF_DIR_BASE:-/run_logs}" + local rpdir="${_rocprof_base}/${SLURM_JOB_ID:-0}/rocprof_${role}_NODE${NODE_RANK}" + mkdir -p "$rpdir" + local flags="${ROCPROF_FLAGS:-"--kernel-trace --marker-trace --hip-trace --hsa-trace"}" + echo "rocprofv3 ${flags} --output-format csv json -d ${rpdir} -o %hostname%_%pid% -- " +} + +roctx_finalize_workers() ( + # SIGINT exact scheduler workers so rocprofiler can finalize, then stop the main. + # Serialization can spend minutes without changing an output file. Never + # re-signal a worker after finalization starts; wait up to the bounded timeout. + getpids(){ ps -eo pid,args | awk '$2 ~ /^sglang::scheduler_(DP|TP)/ {print $1}'; } + count_alive(){ getpids | awk 'END { print NR + 0 }'; } + + FINALIZE_TIMEOUT="${ROCPROF_FINALIZE_TIMEOUT:-1800}" + STALL_LIMIT="${ROCPROF_STALL_LIMIT:-45}" + HOST=$(hostname -s) + + WORKERS="$(getpids)" + NW=$(printf '%s\n' $WORKERS | grep -c . 2>/dev/null || echo 0) + echo "workers: $(echo $WORKERS | tr '\n' ' ') (count=$NW)" + + # Discover the current job's -d path from one process snapshot. With JOBID set, + # never accept unscoped matches because stale jobs can cause false completion. + ALL_RPDIRS=$(ps -eo args 2>/dev/null | grep -oE '\-d /run_logs/[0-9]+/rocprof_[a-z]+_NODE[0-9]+' | awk '{print $2}') + RPDIR="" + if [ -n "${JOBID:-}" ]; then + RPDIR=$(printf '%s\n' "$ALL_RPDIRS" | grep -F "/run_logs/${JOBID}/" | head -1) + # If JOBID is known, reject unscoped paths and use job-scoped fallbacks. + if [ -z "$RPDIR" ]; then + echo "[roctx_finalize_workers] WARN: JOBID=$JOBID set but no matching rocprofv3 -d arg found for it (candidates: $(printf '%s' "$ALL_RPDIRS" | tr '\n' ' ')); this job's rocprofv3 for this role likely never started/wrapped cleanly (e.g. engine hang) -- skipping unscoped ps match (would only find OTHER jobs' stale dirs) and going straight to mtime-based fallback" >&2 + fi + else + RPDIR=$(printf '%s\n' "$ALL_RPDIRS" | head -1) + fi + + # Topology-derived paths work before output files appear. + if [ -z "$RPDIR" ] && [ -n "${JOBID:-}" ]; then + idx=0 + for n in ${PREFILL_NODES:-}; do + ns=${n%%.*} + if [ "$n" = "$HOST" ] || [ "$ns" = "$HOST" ]; then + RPDIR="/run_logs/${JOBID}/rocprof_prefill_NODE${idx}" + break + fi + idx=$((idx + 1)) + done + if [ -z "$RPDIR" ]; then + idx=0; base_idx=${XP:-0} + for n in ${DECODE_NODES:-}; do + ns=${n%%.*} + if [ "$n" = "$HOST" ] || [ "$ns" = "$HOST" ]; then + RPDIR="/run_logs/${JOBID}/rocprof_decode_NODE$((base_idx + idx))" + break + fi + idx=$((idx + 1)) + done + fi + [ -n "$RPDIR" ] && echo "[roctx_finalize_workers] $HOST topology-derived rocprof dir = $RPDIR" + fi + + for p in $WORKERS; do kill -INT "$p" 2>/dev/null; done + + # Keep the fallback job-scoped whenever JOBID is available. + if [ -z "$RPDIR" ]; then + for _ in $(seq 1 20); do + f="" + if [ -n "${JOBID:-}" ]; then + f=$(find "/run_logs/${JOBID}" -maxdepth 2 -path '*/rocprof_*' -name "${HOST}_*" -printf '%T@ %p\n' 2>/dev/null | sort -rn | head -1 | cut -d' ' -f2-) + else + f=$(find /run_logs -maxdepth 3 -path '*/rocprof_*' -name "${HOST}_*" -printf '%T@ %p\n' 2>/dev/null | sort -rn | head -1 | cut -d' ' -f2-) + fi + [ -n "$f" ] && { RPDIR=$(dirname "$f"); break; } + sleep 1 + done + fi + + echo "[roctx_finalize_workers] $HOST rocprof dir = ${RPDIR:-} (finalize_timeout=${FINALIZE_TIMEOUT}s stall_limit=${STALL_LIMIT}s expect=${NW} workers)" + + count_files(){ ls "$RPDIR"/${HOST}_*"$1" 2>/dev/null | grep -c . ; } + newest_mtime(){ find "$RPDIR" -maxdepth 1 -name "${HOST}_*" -printf '%T@\n' 2>/dev/null | sort -n | tail -1; } + + if [ -n "$RPDIR" ] && [ "$NW" -gt 0 ]; then + start=$(date +%s) + last_progress=$start + last_res=0; last_mt=0 + stall_warned=0 + while :; do + res=$(count_files _results.json) + mrk=$(count_files _marker_api_trace.csv) + alive=$(count_alive) + now=$(date +%s); el=$(( now - start )) + if [ "$res" -ge "$NW" ] && [ "$mrk" -ge "$NW" ]; then + echo "[roctx_finalize_workers] $HOST: all $NW workers finalized COMPLETE (results.json=$res marker=$mrk) after ${el}s" + break + fi + if [ "$el" -ge "$FINALIZE_TIMEOUT" ]; then + echo "[roctx_finalize_workers] WARN: $HOST finalize timeout ${FINALIZE_TIMEOUT}s (results.json=$res/$NW marker=$mrk/$NW, $alive still alive)" >&2 + break + fi + mt=$(newest_mtime); mt=${mt%.*}; [ -z "$mt" ] && mt=0 + if [ "$res" -gt "$last_res" ] || [ "${mt:-0}" -gt "${last_mt:-0}" ]; then + last_res=$res; last_mt=$mt; last_progress=$now + fi + if [ "$(( now - last_progress ))" -ge "$STALL_LIMIT" ] && [ "$stall_warned" -eq 0 ]; then + echo "[roctx_finalize_workers] INFO: $HOST no visible output progress for ${STALL_LIMIT}s (results.json=$res/$NW marker=$mrk/$NW); serialization still running" >&2 + stall_warned=1 + fi + sleep 3 + done + else + # fallback (no rocprof dir found / ROCPROF off): preserve a bounded settle window + echo "[roctx_finalize_workers] $HOST: no rocprof dir located -- falling back to bounded 90s settle" + sleep 90 + for p in $(getpids); do kill -INT "$p" 2>/dev/null; done + sleep 20 + fi + echo "workers remaining: $(count_alive)" + + mp=$(ps -eo pid,args | awk '/[p]ython3 -m sglang.launch_server/ {print $1; exit}') + echo "main pid: $mp" + if [ -z "$mp" ]; then + echo "[roctx_finalize_workers] WARN: no launch_server main process found on this node -- it was ALREADY GONE before this flush ran." >&2 + echo "[roctx_finalize_workers] WARN: (topology: xP=${XP:-} yD=${YD:-}) if xP>1 or yD>1, this is the known cross-node" >&2 + echo "[roctx_finalize_workers] WARN: DP_MODE=1 collective-cascade symptom (fixed by parallel per-node finalization dispatch, 2026-07-06)." >&2 + echo "[roctx_finalize_workers] WARN: rocprofv3 likely did NOT flush cleanly on this node -- check for a 0-byte rocprof_*_NODE dir." >&2 + fi + [ -n "$mp" ] && kill -INT "$mp" 2>/dev/null + sleep 15 + echo "main remaining: $(ps -eo args | grep -c '[p]ython3 -m sglang.launch_server')" +) + +finish_server() { + local role="$1" pipeline_pid="$2" + local rc=0 suffix count + JOBID="${SLURM_JOB_ID:-0}" XP="$xP" YD="$yD" \ + roctx_finalize_workers || rc=1 + + local leftover + leftover=$(ps -eo pid,args | awk '/[p]ython3 -m sglang\.launch_server/ {print $1}') + if [[ -n "$leftover" ]]; then + echo "[profile] WARN: ${role} worker traces finalized; stopping launch_server parent" >&2 + kill -TERM $leftover 2>/dev/null || true + sleep 5 + leftover=$(ps -eo pid,args | awk '/[p]ython3 -m sglang\.launch_server/ {print $1}') + [[ -z "$leftover" ]] || kill -KILL $leftover 2>/dev/null || true + fi + if kill -0 "$pipeline_pid" 2>/dev/null; then + echo "[profile] WARN: ${role} rocprof wrapper still running; stopping it" >&2 + kill -TERM "$pipeline_pid" 2>/dev/null || true + sleep 5 + kill -KILL "$pipeline_pid" 2>/dev/null || true + fi + wait "$pipeline_pid" 2>/dev/null || true + + local rpdir="${ROCPROF_DIR_BASE:-/run_logs}/${SLURM_JOB_ID:-0}/rocprof_${role}_NODE${NODE_RANK}" + for suffix in kernel_trace.csv marker_api_trace.csv results.json; do + count=$(find "$rpdir" -maxdepth 1 -name "*_${suffix}" 2>/dev/null | wc -l) + if (( count < GPUS_PER_NODE )); then + echo "[profile] ERROR: ${role} NODE${NODE_RANK} has ${count}/${GPUS_PER_NODE} ${suffix} files" >&2 + rc=1 + fi + done + # Keep rank 0 alive until every node has finalized its worker outputs. This + # prevents the master task from looking complete while decode serialization is still active on another node. + local done_dir="${ROCPROF_DIR_BASE:-/run_logs}/${SLURM_JOB_ID:-0}" + local done_file="$done_dir/.profile_done_NODE${NODE_RANK}" + touch "$done_file" || rc=1 + if [[ "$NODE_RANK" -eq 0 ]]; then + local expected_nodes=$((xP + yD)) + local barrier_deadline=$((SECONDS + ${ROCPROF_NODE_BARRIER_TIMEOUT:-2100})) + local done_nodes=0 + while (( SECONDS < barrier_deadline )); do + done_nodes=$(find "$done_dir" -maxdepth 1 -name '.profile_done_NODE*' 2>/dev/null | wc -l) + (( done_nodes >= expected_nodes )) && break + sleep 5 + done + if (( done_nodes < expected_nodes )); then + echo "[profile] ERROR: node finalization barrier has $done_nodes/$expected_nodes nodes" >&2 + rc=1 + fi + fi + return "$rc" +} diff --git a/scripts/sglang_disagg/moriio_profiling/patches/mori/01-roctx-instrumentation.patch b/scripts/sglang_disagg/moriio_profiling/patches/mori/01-roctx-instrumentation.patch new file mode 100644 index 00000000..72aeaba9 --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/patches/mori/01-roctx-instrumentation.patch @@ -0,0 +1,430 @@ +diff --git a/src/io/engine.cpp b/src/io/engine.cpp +index a22a97b2..99da9b7d 100644 +--- a/src/io/engine.cpp ++++ b/src/io/engine.cpp +@@ -41,6 +41,7 @@ + #include "src/io/fabric/backend_impl.hpp" + #include "src/io/rdma/backend_impl.hpp" + #include "src/io/xgmi/backend_impl.hpp" ++#include "src/io/roctx_mori.hpp" // ADDITIVE: MORI_ROCTX-gated host-send roctx markers + + namespace mori { + namespace io { +@@ -151,6 +152,8 @@ void IOEngineSession::BatchWrite(const SizeVec& localOffsets, const SizeVec& rem + const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id) { + MORI_IO_FUNCTION_TIMER; ++ // ADDITIVE (MORI_ROCTX=1): brackets the host KV-send dispatch for this transfer. ++ mori::io::MoriRoctxRange _mori_roctx_("mori.io.session_batch_write", static_cast(id)); + std::shared_ptr diagnostics; + internal::ScopedIoCallDiagnosticsCapture capture(&diagnostics, "Session batch write"); + backendSess->BatchWrite(localOffsets, remoteOffsets, sizes, status, id); +@@ -529,6 +532,8 @@ void IOEngine::BatchWrite(const MemDescVec& localSrc, const BatchSizeVec& localO + const BatchSizeVec& sizes, TransferStatusPtrVec& status, + TransferUniqueIdVec& ids) { + MORI_IO_FUNCTION_TIMER; ++ // ADDITIVE (MORI_ROCTX=1): brackets the host engine-level batch KV-send. ++ mori::io::MoriRoctxRange _mori_roctx_("mori.io.engine_batch_write"); + size_t batchSize = localSrc.size(); + assert(batchSize == remoteDest.size()); + assert(batchSize == localOffsets.size()); +diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp +index 211b59e5..b3ad2ca0 100644 +--- a/src/io/rdma/backend_impl.cpp ++++ b/src/io/rdma/backend_impl.cpp +@@ -38,6 +38,7 @@ + #include "mori/io/env.hpp" + #include "mori/io/logging.hpp" + #include "src/io/rdma/protocol.hpp" ++#include "src/io/roctx_mori.hpp" // ADDITIVE: MORI_ROCTX_TRANSFER async post->cq range stop + namespace mori { + namespace io { + +@@ -736,6 +737,9 @@ NotifManager::FlushDrainStats NotifManager::ProcessOneCqe( + auto meta = ep.ledger + ? ep.ledger->ReleaseByCqe(wc[i].wr_id, ep.sqDepth.get(), &mergedBatchSize) + : nullptr; ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): stop the async post->cq range for this ++ // signaled WR on a FAILED/flush CQE (no-op if none was started for it). ++ mori::io::MoriRoctxTransferStop(ep.ledger ? ep.ledger.get() : nullptr, wc[i].wr_id); + if (meta) { + (void)meta->finishedBatchSize.fetch_add(mergedBatchSize); + if (isFlush) { +@@ -840,6 +844,9 @@ NotifManager::FlushDrainStats NotifManager::ProcessOneCqe( + auto meta = ep.ledger + ? ep.ledger->ReleaseByCqe(recordId, ep.sqDepth.get(), &mergedBatchSize) + : nullptr; ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): stop the async post->cq range for this ++ // signaled WR on its SUCCESS CQE -> this is the real KV transfer duration. ++ mori::io::MoriRoctxTransferStop(ep.ledger ? ep.ledger.get() : nullptr, recordId); + if (meta) { + NotifySqStateChanged(ep); + uint32_t finishedBefore = meta->finishedBatchSize.fetch_add(mergedBatchSize); +diff --git a/src/io/rdma/common.cpp b/src/io/rdma/common.cpp +index 8e4e7f4c..c6321e5c 100644 +--- a/src/io/rdma/common.cpp ++++ b/src/io/rdma/common.cpp +@@ -39,6 +39,7 @@ + + #include "mori/io/env.hpp" + #include "mori/io/logging.hpp" ++#include "src/io/roctx_mori.hpp" // ADDITIVE: MORI_ROCTX-gated host-send roctx markers + + namespace mori { + namespace io { +@@ -606,6 +607,17 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + TransferUniqueId id, bool isRead, int postBatchSize, + const RdmaTransferControl& control) { + MORI_IO_FUNCTION_TIMER; ++ // ADDITIVE (MORI_ROCTX=1): brackets the host ibv_post_send loop for this RDMA ++ // batch (write = the KV send on the prefill sender; read also routes here). ++ // bytes= carries the whole-call payload (sum of the per-request sizes). ++ // wrs= carries the whole-call pre-merge request count (sizes.size()) -- the ++ // same "known at entry, from the sizes vector" granularity as bytes=; this is ++ // an upper bound on the actual posted WR count (merging can only reduce it). ++ mori::io::MoriRoctxRange _mori_roctx_( ++ isRead ? "mori.rdma.batch_post.read" : "mori.rdma.batch_post.write", ++ static_cast(id), ++ std::accumulate(sizes.begin(), sizes.end(), static_cast(0)), ++ static_cast(sizes.size())); + + if ((localOffsets.size() != remoteOffsets.size()) || (sizes.size() != remoteOffsets.size())) { + return {StatusCode::ERR_INVALID_ARGS, +@@ -654,6 +666,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + thread_local std::vector tlChunkPlan; + thread_local std::vector tlEpWrsSinceSignal; + thread_local std::vector tlEpMergedSinceSignal; ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): payload bytes accumulated since the last ++ // signal on each EP (mirrors tlEpMergedSinceSignal); attached to the signaled ++ // record's kv_transfer range and reset on signal. ++ thread_local std::vector tlEpBytesSinceSignal; + thread_local int reentryDepth = 0; + + struct ReentryGuard { +@@ -670,6 +686,7 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + std::vector localChunkPlan; + std::vector localEpWrsSinceSignal; + std::vector localEpMergedSinceSignal; ++ std::vector localEpBytesSinceSignal; // ADDITIVE + + std::vector& indices = usePool ? tlIndices : localIndices; + std::vector& mergedPool = usePool ? tlMergedPool : localMergedPool; +@@ -678,6 +695,9 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + std::vector& epWrsSinceSignal = usePool ? tlEpWrsSinceSignal : localEpWrsSinceSignal; + std::vector& epMergedSinceSignal = + usePool ? tlEpMergedSinceSignal : localEpMergedSinceSignal; ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): see tlEpBytesSinceSignal above. ++ std::vector& epBytesSinceSignal = ++ usePool ? tlEpBytesSinceSignal : localEpBytesSinceSignal; + + // Bound peak retained memory: if an earlier very large batch grew the pools far + // beyond the current need, release the excess so it doesn't stay resident. +@@ -883,6 +903,7 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + + epWrsSinceSignal.assign(epNum, 0); + epMergedSinceSignal.assign(epNum, 0); ++ epBytesSinceSignal.assign(epNum, 0); // ADDITIVE + + // Rotate the starting EP by transfer id so single-segment (single WR) + // transfers spread evenly across all QPs instead of always landing on eps[0]. +@@ -905,6 +926,7 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + const auto& localMr = localMrPerEp[epId]; + const auto& remoteMr = remoteMrPerEp[epId]; + size_t mergedReqSize = 0; ++ size_t batchBytes = 0; // ADDITIVE: payload bytes for this post chunk + for (int j = st; j < end; j++) { + MergedWorkRequest& mergedWr = mergedWrs[j]; + for (auto& sge : mergedWr.sges) sge.lkey = localMr.lkey; +@@ -915,10 +937,12 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + wr.wr_id = 0; + wr.next = (j + 1 < end) ? &mergedWrs[j + 1].wr : nullptr; + mergedReqSize += mergedWr.mergedRequests; ++ batchBytes += mergedWr.totalRemoteLength; // ADDITIVE + } + + epWrsSinceSignal[epId] += batchWrNum; + epMergedSinceSignal[epId] += mergedReqSize; ++ epBytesSinceSignal[epId] += batchBytes; // ADDITIVE + + bool isLastBatchForEp = ((i + epNum) >= numPostBatch); + bool sqNearFull = eps[epId].sqDepth && (epWrsSinceSignal[epId] >= eps[epId].maxSqDepth); +@@ -936,6 +960,15 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + static_cast(epMergedSinceSignal[epId])); + last.wr_id = recordId; + last.send_flags = IBV_SEND_SIGNALED; ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): start an ASYNC post->cq range for this ++ // signaled WR, keyed by (ledger, recordId). Stopped when its CQE is reaped ++ // (NotifManager::ProcessOneCqe -> ledger->ReleaseByCqe) or, if this WR fails ++ // to post, in the not-posted cleanup below. Measures real KV transfer/wire ++ // time (ms), unlike the synchronous host-post anchor above (us). ++ mori::io::MoriRoctxTransferStart(eps[epId].ledger.get(), recordId, ++ static_cast(id), isRead, ++ static_cast(epBytesSinceSignal[epId]), ++ static_cast(epWrsSinceSignal[epId])); + } + + struct ibv_send_wr* badWr = nullptr; +@@ -972,6 +1005,9 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + } else if (needSignal) { + int dummy = 0; + eps[epId].ledger->ReleaseByCqe(recordId, nullptr, &dummy); ++ // ADDITIVE (MORI_ROCTX_TRANSFER=1): the signaled WR never posted, so no CQE ++ // will arrive -- stop its async range here to avoid a leaked range. ++ mori::io::MoriRoctxTransferStop(eps[epId].ledger.get(), recordId); + } + + if (postedCount > 0 && (!needSignal || !lastWasPosted)) { +@@ -1022,6 +1058,7 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + if (needSignal) { + epWrsSinceSignal[epId] = 0; + epMergedSinceSignal[epId] = 0; ++ epBytesSinceSignal[epId] = 0; // ADDITIVE + } + MORI_IO_TRACE("ibv_post_send ep index {} batch index range [{}, {})", epId, st, end); + } +diff --git a/src/io/roctx_mori.hpp b/src/io/roctx_mori.hpp +new file mode 100644 +index 00000000..c14230bd +--- /dev/null ++++ b/src/io/roctx_mori.hpp +@@ -0,0 +1,240 @@ ++// Copyright © Advanced Micro Devices, Inc. All rights reserved. ++// MIT License (see repository LICENSE). ++// ============================================================================ ++// ADDITIVE, env-gated roctx markers for the MORI-IO HOST RDMA send path. ++// ++// TWO independent, additive instrumentations (each its own env gate; default OFF): ++// ++// (1) MORI_ROCTX=1 -> SYNCHRONOUS push/pop ranges around the host ibv_post_send ++// loop (IOEngine[Session]::BatchWrite + RdmaBatchReadWrite). These measure ++// only the HOST POST cost (building WRs + ringing the NIC doorbell). They ++// are stack/same-thread ranges (MoriRoctxRange RAII) and CANNOT span the ++// async post->completion window. Marker names: mori.io.engine_batch_write, ++// mori.rdma.batch_post.{write,read}. ++// ++// (2) MORI_ROCTX_TRANSFER=1 -> ASYNCHRONOUS post->CQ ranges that measure the ++// REAL KV transfer/wire duration: started when a *signaled* WR is posted ++// (RdmaBatchReadWrite, needSignal branch) and stopped when its completion ++// is reaped on the CQ (NotifManager::ProcessOneCqe -> ledger->ReleaseByCqe). ++// Uses the PROCESS-WIDE async roctx API roctxRangeStartA/roctxRangeStop ++// (start on the posting thread, stop on the CQ-poll thread). Marker name: ++// mori.rdma.kv_transfer (its own dedicated trace lane). ++// ++// RDMA uses SELECTIVE SIGNALING: only the tail WR of each post batch sets ++// IBV_SEND_SIGNALED and receives a SubmissionLedger recordId (== wr_id). ++// Only that signaled WR produces a CQE, so we start exactly ONE async range ++// per signaled WR (keyed by the ledger recordId) -> every started range has ++// a matching stop (the CQE, or the not-posted cleanup path). recordId is ++// per-EP-ledger (not globally unique), so the range map is keyed by the ++// PAIR (SubmissionLedger*, recordId), which is globally unique and identical ++// at the post site (eps[i].ledger) and the CQ site (ep.ledger) because both ++// hold the same shared SubmissionLedger instance. ++// ++// CRITICAL: rocprofv3 (rocprofiler-sdk) --marker-trace only intercepts the ++// rocprofiler-sdk ROCTx library librocprofiler-sdk-roctx.so, NOT legacy ++// libroctx64.so. We dlopen the sdk lib at runtime (RTLD_GLOBAL) and resolve the ++// roctx symbols from it (no link-time dependency added to libmori_io.so). ++// ++// Fully gated + exception-safe: when neither gate is set the lib is never dlopen'd ++// and every call is a no-op (a single bool check). ++// ============================================================================ ++#pragma once ++ ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++ ++namespace mori { ++namespace io { ++namespace roctx_detail { ++ ++using roctx_range_push_t = int (*)(const char*); ++using roctx_range_pop_t = int (*)(); ++using roctx_mark_t = void (*)(const char*); ++// Process-wide async range API (start on one thread, stop from any other). ++using roctx_range_id_t = std::uint64_t; ++using roctx_range_start_t = roctx_range_id_t (*)(const char*); ++using roctx_range_stop_t = void (*)(roctx_range_id_t); ++ ++inline bool GateOn(const char* name) { ++ const char* g = std::getenv(name); ++ if (g == nullptr) return false; ++ const char c = g[0]; ++ return (c == '1' || c == 't' || c == 'T' || c == 'y' || c == 'Y' || c == 'o' || c == 'O'); ++} ++ ++struct RoctxApi { ++ bool enabled = false; // MORI_ROCTX: push/pop host-post anchors ++ bool transfer_enabled = false; // MORI_ROCTX_TRANSFER: async post->cq ranges ++ roctx_range_push_t push = nullptr; ++ roctx_range_pop_t pop = nullptr; ++ roctx_mark_t mark = nullptr; ++ roctx_range_start_t range_start = nullptr; ++ roctx_range_stop_t range_stop = nullptr; ++ ++ RoctxApi() { ++ const bool want_post = GateOn("MORI_ROCTX"); ++ const bool want_transfer = GateOn("MORI_ROCTX_TRANSFER"); ++ if (!want_post && !want_transfer) return; ++ // sdk-roctx ONLY (the lib rocprofv3 --marker-trace intercepts). ++ void* h = dlopen("librocprofiler-sdk-roctx.so", RTLD_NOW | RTLD_GLOBAL); ++ if (h == nullptr) h = dlopen("librocprofiler-sdk-roctx.so.1", RTLD_NOW | RTLD_GLOBAL); ++ if (h == nullptr) return; ++ push = reinterpret_cast(dlsym(h, "roctxRangePushA")); ++ pop = reinterpret_cast(dlsym(h, "roctxRangePop")); ++ mark = reinterpret_cast(dlsym(h, "roctxMarkA")); ++ range_start = reinterpret_cast(dlsym(h, "roctxRangeStartA")); ++ range_stop = reinterpret_cast(dlsym(h, "roctxRangeStop")); ++ enabled = want_post && (push != nullptr && pop != nullptr); ++ transfer_enabled = want_transfer && (range_start != nullptr && range_stop != nullptr); ++ } ++}; ++ ++inline RoctxApi& api() { ++ static RoctxApi a; // gate read + dlopen happen exactly once per process ++ return a; ++} ++ ++// (SubmissionLedger*, recordId) -> async roctx range id. recordId is unique only ++// within one ledger, so the ledger pointer disambiguates across endpoints. ++using TransferKey = std::pair; ++struct TransferKeyHash { ++ std::size_t operator()(const TransferKey& k) const { ++ std::size_t h1 = std::hash{}(k.first); ++ std::size_t h2 = std::hash{}(k.second); ++ return h1 ^ (h2 + 0x9e3779b9 + (h1 << 6) + (h1 >> 2)); ++ } ++}; ++struct TransferRanges { ++ std::mutex mu; ++ std::unordered_map ranges; ++}; ++inline TransferRanges& transfer_ranges() { ++ static TransferRanges t; ++ return t; ++} ++ ++} // namespace roctx_detail ++ ++// RAII range: pushes on construction, pops on destruction (handles every return ++// path + exception). No-op when MORI_ROCTX is off. (HOST-POST anchor only.) ++class MoriRoctxRange { ++ public: ++ explicit MoriRoctxRange(const char* name) { ++ auto& a = roctx_detail::api(); ++ if (a.enabled) { ++ a.push(name); ++ active_ = true; ++ } ++ } ++ MoriRoctxRange(const char* name, uint64_t id) { ++ auto& a = roctx_detail::api(); ++ if (a.enabled) { ++ std::string s = std::string(name) + " id=" + std::to_string(id); ++ a.push(s.c_str()); ++ active_ = true; ++ } ++ } ++ // ADDITIVE: host-post anchor variant carrying the whole-call payload size. ++ // Keeps id= LAST so end-anchored id= parsers stay valid: " bytes= id=". ++ MoriRoctxRange(const char* name, uint64_t id, uint64_t bytes) { ++ auto& a = roctx_detail::api(); ++ if (a.enabled) { ++ std::string s = std::string(name) + " bytes=" + std::to_string(bytes) + ++ " id=" + std::to_string(id); ++ a.push(s.c_str()); ++ active_ = true; ++ } ++ } ++ // ADDITIVE: host-post anchor variant that also carries the whole-call WR ++ // count (pre-merge request count, i.e. sizes.size() at the RdmaBatchReadWrite ++ // call site -- the same "known at entry, from the sizes vector" granularity ++ // already used for bytes above). Keeps id= LAST: ++ // " bytes= wrs= id=". ++ MoriRoctxRange(const char* name, uint64_t id, uint64_t bytes, uint64_t wrs) { ++ auto& a = roctx_detail::api(); ++ if (a.enabled) { ++ std::string s = std::string(name) + " bytes=" + std::to_string(bytes) + ++ " wrs=" + std::to_string(wrs) + " id=" + std::to_string(id); ++ a.push(s.c_str()); ++ active_ = true; ++ } ++ } ++ ~MoriRoctxRange() { ++ if (active_) { ++ auto& a = roctx_detail::api(); ++ if (a.pop != nullptr) a.pop(); ++ } ++ } ++ MoriRoctxRange(const MoriRoctxRange&) = delete; ++ MoriRoctxRange& operator=(const MoriRoctxRange&) = delete; ++ ++ private: ++ bool active_ = false; ++}; ++ ++inline void MoriRoctxMark(const std::string& msg) { ++ auto& a = roctx_detail::api(); ++ if (a.enabled && a.mark != nullptr) a.mark(msg.c_str()); ++} ++ ++// --- ASYNC post->cq KV-transfer ranges (MORI_ROCTX_TRANSFER) ------------------ ++// Start an async range for a SIGNALED WR at post time. Keyed by (ledger,recordId). ++// ADDITIVE: `wrs` carries epWrsSinceSignal[epId] at the signal point -- the ++// number of WRs (across possibly several RdmaBatchReadWrite calls, since ++// unsignaled chunks roll forward until the next signal) this one signaled ++// completion covers. Placed BEFORE id= (same rule as bytes=) so end-anchored ++// id= parsers keep matching: " bytes= wrs= id=". ++inline void MoriRoctxTransferStart(const void* ledger, std::uint64_t recordId, ++ std::uint64_t transferId, bool isRead, ++ std::uint64_t bytes = 0, std::uint64_t wrs = 0) { ++ auto& a = roctx_detail::api(); ++ if (!a.transfer_enabled || a.range_start == nullptr || ledger == nullptr) return; ++ // bytes=/wrs= placed BEFORE id= so the end-anchored id= parsers keep matching. ++ std::string s = ++ std::string(isRead ? "mori.rdma.kv_transfer.read" : "mori.rdma.kv_transfer") + ++ " bytes=" + std::to_string(bytes) + " wrs=" + std::to_string(wrs) + ++ " id=" + std::to_string(transferId); ++ roctx_detail::roctx_range_id_t rid = a.range_start(s.c_str()); ++ auto& t = roctx_detail::transfer_ranges(); ++ std::lock_guard lk(t.mu); ++ t.ranges[{reinterpret_cast(ledger), recordId}] = rid; ++} ++ ++// Stop the async range for a completed/cleaned-up signaled WR. Idempotent: a ++// no-op if no range was started for this (ledger,recordId) (e.g. unsignaled WRs, ++// notification CQEs). The roctxRangeStop call is made OUTSIDE the map lock. ++inline void MoriRoctxTransferStop(const void* ledger, std::uint64_t recordId) { ++ auto& a = roctx_detail::api(); ++ if (!a.transfer_enabled || a.range_stop == nullptr || ledger == nullptr) return; ++ roctx_detail::roctx_range_id_t rid = 0; ++ bool found = false; ++ { ++ auto& t = roctx_detail::transfer_ranges(); ++ std::lock_guard lk(t.mu); ++ auto it = t.ranges.find({reinterpret_cast(ledger), recordId}); ++ if (it != t.ranges.end()) { ++ rid = it->second; ++ t.ranges.erase(it); ++ found = true; ++ } ++ } ++ if (found) a.range_stop(rid); ++} ++ ++// Diagnostics: number of started-but-not-stopped transfer ranges (leak counter). ++inline std::size_t MoriRoctxTransferOutstanding() { ++ auto& a = roctx_detail::api(); ++ if (!a.transfer_enabled) return 0; ++ auto& t = roctx_detail::transfer_ranges(); ++ std::lock_guard lk(t.mu); ++ return t.ranges.size(); ++} ++ ++} // namespace io ++} // namespace mori diff --git a/scripts/sglang_disagg/moriio_profiling/patches/sglang/01-roctx-instrumentation.patch b/scripts/sglang_disagg/moriio_profiling/patches/sglang/01-roctx-instrumentation.patch new file mode 100644 index 00000000..643cf4fb --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/patches/sglang/01-roctx-instrumentation.patch @@ -0,0 +1,982 @@ +diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py +index cd80d277f0..f4802ae061 100644 +--- a/python/sglang/benchmark/serving.py ++++ b/python/sglang/benchmark/serving.py +@@ -15,6 +15,7 @@ python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --num + import argparse + import asyncio + import copy ++import csv + import importlib.util + import json + import math +@@ -92,6 +93,8 @@ class RequestFuncInput: + extra_request_body: Dict[str, Any] + timestamp: Optional[float] = None + routing_key: Optional[str] = None ++ request_index: Optional[int] = None ++ rid: Optional[str] = None + + + @dataclass +@@ -112,11 +115,21 @@ class RequestFuncOutput: + spec_cap_length: float = 0.0 + spec_block_accept_length: float = 0.0 + spec_cap_lens_histogram: List[int] = field(default_factory=list) ++ request_index: Optional[int] = None ++ rid: str = "" ++ requested_output_len: int = 0 ++ client_send_wall_ns: Optional[int] = None ++ client_first_token_wall_ns: Optional[int] = None ++ client_done_wall_ns: Optional[int] = None ++ http_status: Optional[int] = None + + @staticmethod + def init_new(request_func_input: RequestFuncInput): + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len ++ output.request_index = request_func_input.request_index ++ output.rid = request_func_input.rid or "" ++ output.requested_output_len = request_func_input.output_len + return output + + +@@ -686,6 +699,8 @@ async def async_request_sglang_generate( + # Add image data if available (list of image urls/base64) + if request_func_input.image_data: + payload["image_data"] = request_func_input.image_data ++ if request_func_input.rid is not None: ++ payload["rid"] = request_func_input.rid + + headers = get_request_headers() + if request_func_input.routing_key: +@@ -701,9 +716,11 @@ async def async_request_sglang_generate( + most_recent_timestamp = st + last_output_len = 0 + try: ++ output.client_send_wall_ns = time.time_ns() + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: ++ output.http_status = response.status + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() +@@ -715,6 +732,8 @@ async def async_request_sglang_generate( + if chunk == "[DONE]": + pass + else: ++ if output.client_first_token_wall_ns is None: ++ output.client_first_token_wall_ns = time.time_ns() + data = json.loads(chunk) + + _meta_info = data.get("meta_info") or {} +@@ -769,6 +788,9 @@ async def async_request_sglang_generate( + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + print(f"{output.error=}") ++ finally: ++ if output.client_send_wall_ns is not None: ++ output.client_done_wall_ns = time.time_ns() + + if pbar: + pbar.update(1) +@@ -1300,6 +1322,100 @@ def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callab + return f + + ++def _write_client_trace_artifacts( ++ outputs: List[RequestFuncOutput], ++ request_id_prefix: str, ++ client_timing_csv: Optional[str], ++ client_manifest: Optional[str], ++ max_concurrency: Optional[int], ++) -> None: ++ """Write deterministic request timing artifacts after all tasks complete.""" ++ rows = sorted( ++ outputs, ++ key=lambda output: ( ++ output.request_index is None, ++ output.request_index if output.request_index is not None else 0, ++ ), ++ ) ++ fieldnames = [ ++ "request_index", ++ "rid", ++ "success", ++ "client_send_wall_ns", ++ "client_first_token_wall_ns", ++ "client_done_wall_ns", ++ "prompt_len", ++ "output_len", ++ "requested_output_len", ++ "http_status", ++ "error", ++ ] ++ ++ if client_timing_csv: ++ path = Path(client_timing_csv) ++ path.parent.mkdir(parents=True, exist_ok=True) ++ with path.open("w", newline="", encoding="utf-8") as file: ++ writer = csv.DictWriter(file, fieldnames=fieldnames) ++ writer.writeheader() ++ for output in rows: ++ writer.writerow({name: getattr(output, name) for name in fieldnames}) ++ ++ if client_manifest: ++ path = Path(client_manifest) ++ path.parent.mkdir(parents=True, exist_ok=True) ++ sends = [ ++ output.client_send_wall_ns ++ for output in rows ++ if output.client_send_wall_ns is not None ++ ] ++ completions = [ ++ output.client_done_wall_ns ++ for output in rows ++ if output.client_done_wall_ns is not None ++ ] ++ input_lengths = {output.prompt_len for output in rows} ++ output_lengths = {output.requested_output_len for output in rows} ++ window = { ++ "start": min(sends) if sends else None, ++ "end": max(completions) if completions else None, ++ } ++ requests_manifest = [ ++ { ++ "request_index": output.request_index, ++ "rid": output.rid, ++ "success": output.success, ++ "client_send_wall_ns": output.client_send_wall_ns, ++ "client_first_token_wall_ns": output.client_first_token_wall_ns, ++ "client_done_wall_ns": output.client_done_wall_ns, ++ "prompt_len": output.prompt_len, ++ "output_len": output.output_len, ++ "requested_output_len": output.requested_output_len, ++ "http_status": output.http_status, ++ "error": output.error, ++ } ++ for output in rows ++ ] ++ manifest = { ++ "tag": request_id_prefix, ++ "request_id_prefix": request_id_prefix, ++ "n": len(rows), ++ "completed": sum(output.success for output in rows), ++ "concurrency": max_concurrency, ++ "isl_input_tokens": ( ++ next(iter(input_lengths)) if len(input_lengths) == 1 else None ++ ), ++ "osl_max_new_tokens": ( ++ next(iter(output_lengths)) if len(output_lengths) == 1 else None ++ ), ++ "benchmark_window_wall_ns": window, ++ # Backward-compatible alias for existing artifact consumers. ++ "probe_window_wall_ns": window, ++ "requests": requests_manifest, ++ } ++ with path.open("w", encoding="utf-8") as file: ++ json.dump(manifest, file, indent=2) ++ ++ + async def benchmark( + backend: str, + api_url: str, +@@ -1323,6 +1439,9 @@ async def benchmark( + mooncake_num_rounds=1, + profile_prefill_url: Optional[List[str]] = None, + profile_decode_url: Optional[List[str]] = None, ++ request_id_prefix: Optional[str] = None, ++ client_timing_csv: Optional[str] = None, ++ client_manifest: Optional[str] = None, + ): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] +@@ -1338,6 +1457,8 @@ async def benchmark( + and _normalize_round_messages(first_prompt[0]) is not None + ) + if is_multi_turn: ++ if request_id_prefix: ++ raise ValueError("--request-id-prefix does not support multi-turn datasets") + request_func = wrap_multi_turn_request_func(request_func, backend=backend) + + # Limit concurrency +@@ -1477,6 +1598,7 @@ async def benchmark( + lora_probs = None + + pbar = None if disable_tqdm else tqdm(total=pbar_total) ++ request_index = 0 + async for request in request_generator: + if lora_names is not None and len(lora_names) != 0: + if lora_request_distribution == "uniform": +@@ -1508,7 +1630,14 @@ async def benchmark( + extra_request_body=merged_extra_body, + timestamp=request.timestamp, + routing_key=request.routing_key, ++ request_index=request_index, ++ rid=( ++ f"{request_id_prefix}-{request_index:03d}" ++ if request_id_prefix ++ else None ++ ), + ) ++ request_index += 1 + + tasks.append( + asyncio.create_task( +@@ -1518,6 +1647,14 @@ async def benchmark( + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + if is_multi_turn: + outputs = [x for output in outputs for x in output] ++ if request_id_prefix: ++ _write_client_trace_artifacts( ++ outputs=outputs, ++ request_id_prefix=request_id_prefix, ++ client_timing_csv=client_timing_csv, ++ client_manifest=client_manifest, ++ max_concurrency=max_concurrency, ++ ) + + # Stop profiler (only if profile_steps was not provided, as it auto-stops) + if profile and not ( +@@ -2103,6 +2240,9 @@ def run_benchmark(args_: argparse.Namespace): + mooncake_num_rounds=args.mooncake_num_rounds, + profile_prefill_url=getattr(args, "profile_prefill_url", None), + profile_decode_url=getattr(args, "profile_decode_url", None), ++ request_id_prefix=getattr(args, "request_id_prefix", None), ++ client_timing_csv=getattr(args, "client_timing_csv", None), ++ client_manifest=getattr(args, "client_manifest", None), + ) + ) + +@@ -2345,6 +2485,25 @@ def cli_main(): + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) ++ tracing_group = parser.add_argument_group("request tracing arguments") ++ tracing_group.add_argument( ++ "--request-id-prefix", ++ type=str, ++ default=None, ++ help="Assign deterministic PREFIX-NNN request IDs to native /generate requests.", ++ ) ++ tracing_group.add_argument( ++ "--client-timing-csv", ++ type=str, ++ default=None, ++ help="Write one wall-clock timing row per benchmark request.", ++ ) ++ tracing_group.add_argument( ++ "--client-manifest", ++ type=str, ++ default=None, ++ help="Write request IDs and benchmark timing-window metadata as JSON.", ++ ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--output-details", action="store_true", help="Output details of benchmarking." +@@ -2696,6 +2855,12 @@ def cli_main(): + ) + args = parser.parse_args() + _validate_parsed_gsp_args(parser, args) ++ if (args.client_timing_csv or args.client_manifest) and not args.request_id_prefix: ++ parser.error( ++ "--client-timing-csv/--client-manifest require --request-id-prefix" ++ ) ++ if args.request_id_prefix and args.backend not in ("sglang", "sglang-native"): ++ parser.error("--request-id-prefix requires the native SGLang backend") + run_benchmark(args) + + +diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py +index b175ca5411..2b67eb1ef5 100644 +--- a/python/sglang/srt/disaggregation/mori/conn.py ++++ b/python/sglang/srt/disaggregation/mori/conn.py +@@ -27,6 +27,27 @@ from mori.io import ( + ) + + from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll ++ ++# --- OPTIONAL roctx markers for KV-transfer visibility (gate: env SGLANG_KV_ROCTX=1) --- ++# No-op unless enabled AND libroctx64 is loadable. We emit roctx MARKS (instants), not ++# ranges, because a transfer's start/end can be on different threads; read the per-side ++# transfer window as (kv_*_done_ts - kv_*_start_ts) for matching bootstrap_room in a ++# roctx-CAPTURING trace (rocprofv3 --marker-trace, or RTL lite's roctx shim). These are ++# PREFILL/DECODE per-side windows on that engine's clock -- NOT cross-node wire time. ++import os as _os_kvx ++_kvx_mark = (lambda _m: None) ++if _os_kvx.environ.get("SGLANG_KV_ROCTX", "0") == "1": ++ try: ++ import ctypes as _ct_kvx ++ _lib_kvx = _ct_kvx.CDLL("libroctx64.so") ++ _lib_kvx.roctxMarkA.argtypes = [_ct_kvx.c_char_p] ++ def _kvx_mark(_m, _f=_lib_kvx.roctxMarkA): ++ try: ++ _f(_m.encode("ascii", "replace")) ++ except Exception: ++ pass ++ except Exception: ++ pass + from sglang.srt.disaggregation.common.conn import ( + CommonKVBootstrapServer, + CommonKVManager, +@@ -49,6 +70,46 @@ from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto + logger = logging.getLogger(__name__) + MORI_GUARD = b"MoriMsgGuard" + ++# --- roctx request-attribution map mark (ADDITIVE, SGLANG_ROCTX-gated, exception-safe) --- ++# The transfer_uid returned by self.engine.allocate_transfer_uid() is the SAME integer ++# the MORI-IO C++ backend prints as `id=N` on BOTH the host-post mark ++# (`mori.rdma.batch_post.write id=N`) and the CQ post->CQE range ++# (`mori.rdma.kv_transfer id=N`). Those marks are NOT request-tagged. Emitting one ++# instant `mori.map room= uid=` mark per allocation lets an ++# offline consumer attribute BOTH lanes to the owning request by exact id-lookup. We use ++# the rocprofiler-sdk roctx lib (the one `rocprofv3 --marker-trace` intercepts) via ++# sglang.srt.observability.sglang_roctx -- NOT this file's legacy libroctx64 _kvx_mark, ++# which rocprofv3 does not see. Fully inert (single bool check) when SGLANG_ROCTX is unset. ++try: ++ from sglang.srt.observability.sglang_roctx import roctx_mark as _mori_roctx_mark ++except Exception: # helper missing for any reason => fully inert ++ def _mori_roctx_mark(_msg): ++ return None ++ ++# bootstrap_room is in scope at add_transfer_request() but NOT at every uid-alloc site ++# (_submit_batch_transfer_plan / _send_mamba_state are reached via send_kvcache / ++# send_state which do not carry it). The per-target sends run SYNCHRONOUSLY on the same ++# thread inside add_transfer_request's loop, so we stash the current room in a ++# threading.local at that entry and read it at the alloc sites. Thread-local (not a plain ++# instance attr) keeps concurrent senders on different threads from racing. ++_mori_room_tls = threading.local() ++ ++ ++def _mori_emit_map(transfer_uid, room=None): ++ """Emit `mori.map room=R uid=U`. room defaults to the current thread's stashed room. ++ ++ No-op (never raises) when no room is known or when SGLANG_ROCTX is off. Additive. ++ """ ++ if room is None: ++ room = getattr(_mori_room_tls, "room", None) ++ if room is None: ++ return ++ try: ++ _mori_roctx_mark("mori.map room=%d uid=%d" % (room, transfer_uid)) ++ except Exception: ++ # Observability must never break the KV send path. ++ pass ++ + + def _normalize_state_indices_per_component( + state_indices: Optional[List], +@@ -738,6 +799,9 @@ class MoriKVManager(CommonKVManager): + return [] + + transfer_uid = self.engine.allocate_transfer_uid() ++ # Attribute this transfer to the owning request (room from thread-local set in ++ # add_transfer_request); no-op when SGLANG_ROCTX off or room unknown. ++ _mori_emit_map(transfer_uid) + + statuses = self.engine.batch_write( + [src_desc], +@@ -983,7 +1047,10 @@ class MoriKVManager(CommonKVManager): + local_offsets.append([prefill_aux_index * item_len]) + remote_offsets.append([dst_aux_index * item_len]) + sizes.append([item_len]) +- uids.append(self.engine.allocate_transfer_uid()) ++ transfer_uid = self.engine.allocate_transfer_uid() ++ uids.append(transfer_uid) ++ # room is an explicit param on this path. ++ _mori_emit_map(transfer_uid, room) + return list( + self.engine.batch_write( + src_descs, local_offsets, dst_descs, remote_offsets, sizes, uids +@@ -1189,6 +1256,8 @@ class MoriKVManager(CommonKVManager): + size = bytes_to_send + + transfer_uid = self.engine.allocate_transfer_uid() ++ # room from thread-local (set in add_transfer_request); no-op if unknown. ++ _mori_emit_map(transfer_uid) + batch_statuses = self.engine.batch_write( + [src_desc], + [[src_offset]], +@@ -1309,6 +1378,12 @@ class MoriKVManager(CommonKVManager): + ) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]: + assert self.disaggregation_mode == DisaggregationMode.PREFILL + ++ # Stash the owning room for this thread so the downstream uid-alloc sites ++ # (_submit_batch_transfer_plan / _send_mamba_state, reached via send_kvcache / ++ # send_state which do not carry bootstrap_room) can emit a request-attributable ++ # `mori.map` mark. The per-target sends below run synchronously on THIS thread. ++ _mori_room_tls.room = bootstrap_room ++ + if ( + bootstrap_room not in self.request_status + or self.request_status.get(bootstrap_room) == KVPoll.Failed +@@ -1326,6 +1401,7 @@ class MoriKVManager(CommonKVManager): + return [], None + + self.update_status(bootstrap_room, KVPoll.Transferring) ++ _kvx_mark("kv_send_start room=%d" % bootstrap_room) + for info in transfer_infos.values(): + peer_info = self.decode_kv_args_table.get(info.engine_key) + if not peer_info: +@@ -1383,6 +1459,7 @@ class MoriKVManager(CommonKVManager): + ) + return result_statuses, target_infos_snapshot + ++ _kvx_mark("kv_send_done room=%d" % bootstrap_room) + return result_statuses, target_infos_snapshot + + +@@ -1743,6 +1820,7 @@ class MoriKVReceiver(CommonKVReceiver): + ] + ) + self.init_time = time.time() ++ _kvx_mark("kv_recv_start room=%s" % self.bootstrap_room) + + def poll(self) -> KVPoll: + if self.conclude_state is not None: +@@ -1750,6 +1828,8 @@ class MoriKVReceiver(CommonKVReceiver): + + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (KVPoll.Success, KVPoll.Failed): ++ if status == KVPoll.Success: ++ _kvx_mark("kv_recv_done room=%s" % self.bootstrap_room) + self.conclude_state = status + return status + +diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py +index 6e8bc0db06..73df64fa19 100644 +--- a/python/sglang/srt/disaggregation/prefill.py ++++ b/python/sglang/srt/disaggregation/prefill.py +@@ -1178,6 +1178,10 @@ class SchedulerDisaggregationPrefillMixin: + if not req.disagg_kv_sender.should_send_kv_chunk(len(page_indices), last_chunk): + return + req.disagg_kv_sender.send(page_indices, state_indices) ++ # Stamp the start of the KV transfer window (FIRST chunk only; the setter ++ # is one-shot). Guard against a missing time_stats just in case. ++ if getattr(req, "time_stats", None) is not None: ++ req.time_stats.set_prefill_kv_transfer_start_time() + req.start_send_idx = end_idx + + def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None: +diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py +index a8cea65ec9..1f56c574f4 100755 +--- a/python/sglang/srt/managers/schedule_batch.py ++++ b/python/sglang/srt/managers/schedule_batch.py +@@ -1044,6 +1044,16 @@ class Req(ReqDllmMixin): + self.bootstrap_host: str = bootstrap_host + self.bootstrap_port: Optional[int] = bootstrap_port + self.bootstrap_room: Optional[int] = bootstrap_room ++ # roctx (ADDITIVE, SGLANG_ROCTX-gated): tag this request's scheduler ++ # time-stats with its bootstrap_room so the roctx markers emitted at each ++ # SchedulerReqTimeStats stamp correlate per-request (bootstrap_room is the ++ # PD prefill<->decode join key). Pure no-op when SGLANG_ROCTX is unset. ++ try: ++ self.time_stats.roctx_id = ( ++ str(bootstrap_room) if bootstrap_room is not None else "" ++ ) ++ except Exception: ++ pass + # Decode-local: the already-emitted boundary token to replay when a + # retracted request is rebootstrapped. Set in pause_generation(retract) + # and consumed in the decode transfer commit; never plumbed to prefill. +diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py +index 8e03bcaeab..a01b02eb91 100644 +--- a/python/sglang/srt/observability/req_time_stats.py ++++ b/python/sglang/srt/observability/req_time_stats.py +@@ -39,6 +39,23 @@ from sglang.srt.observability.trace import ( + ) + from sglang.srt.utils import get_bool_env_var + ++# --- roctx instrumentation (ADDITIVE, SGLANG_ROCTX-gated, exception-safe no-op) --- ++# Emits an instant roctx marker at every request-time-stats stamp point so that a ++# `rocprofv3 --marker-trace` capture shows the SchedulerReqTimeStats / ++# APIServerReqTimeStats timeline aligned with the GPU kernels. Fully inert (and ++# zero overhead) when SGLANG_ROCTX is unset; never alters serving behavior. ++try: ++ from sglang.srt.observability.sglang_roctx import ( ++ ROCTX_ENABLED as _ROCTX_ENABLED, ++ roctx_mark as _roctx_mark, ++ ) ++except Exception: # helper missing for any reason => fully inert ++ _ROCTX_ENABLED = False ++ ++ def _roctx_mark(_m): ++ return None ++ ++ + if TYPE_CHECKING: + from sglang.srt.disaggregation.base.conn import KVTransferMetric + from sglang.srt.managers.schedule_batch import ScheduleBatch +@@ -236,6 +253,33 @@ class ReqTimeStatsBase: + ) + disagg_mode: DisaggregationMode = DisaggregationMode.NULL + diff_realtime_monotonic: float = 0.0 ++ # roctx per-request correlation tag: bootstrap_room on scheduler stats, rid on ++ # API stats. Populated where available (see schedule_batch.py / tokenizer_manager); ++ # blank otherwise. Plain additive field; only read when SGLANG_ROCTX=1. ++ roctx_id: str = "" ++ ++ # Overridden by subclasses to label the observability layer in the marker name. ++ _ROCTX_LAYER = "base" ++ ++ def _roctx(self, event: str) -> None: ++ """Emit a gated roctx instant marker for one request-time-stats stamp. ++ ++ Fast no-op (a single bool check, no string work) when SGLANG_ROCTX is ++ unset. Exception-safe: observability must never break serving. ++ """ ++ if not _ROCTX_ENABLED: ++ return ++ try: ++ layer = self._ROCTX_LAYER ++ if layer == "sched": ++ name = "reqstats.sched.%s.%s" % (self.disagg_mode_str(), event) ++ else: ++ name = "reqstats.%s.%s" % (layer, event) ++ if self.roctx_id: ++ name += " id=" + self.roctx_id ++ _roctx_mark(name) ++ except Exception: ++ pass + + @classmethod + def new_from_obj(cls, obj: Optional[ReqTimeStatsBase], *args, **kwargs) -> Self: +@@ -373,6 +417,8 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + api_server_dispatch_finish_time: float = 0.0 + response_sent_to_client_time: float = 0.0 + ++ _ROCTX_LAYER = "api" ++ + def __getstate__(self) -> object: + state = {} + # send to DP controller or Scheduler +@@ -387,6 +433,7 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + def set_created_time(self, ts=None): + ts = ts or time.perf_counter() + self.created_time = ts ++ self._roctx("request_received") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_req_start(convert_time_to_realtime_ns(ts)) +@@ -401,6 +448,7 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + def set_finished_time(self, ts=None): + ts = ts or time.perf_counter() + self.finished_time = ts ++ self._roctx("request_finished") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_req_finish(convert_time_to_realtime_ns(ts)) +@@ -409,14 +457,17 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + ts = ts or time.perf_counter() + self.first_token_time = ts + self.last_time = ts ++ self._roctx("first_token") + + def set_last_time(self, ts=None): + ts = ts or time.perf_counter() + self.last_time = ts ++ self._roctx("last_token") + + def set_tokenize_finish_time(self, ts=None): + ts = ts or time.perf_counter() + self.tokenize_finish_time = ts ++ self._roctx("tokenize_finish") + + # tokenize span was started in set_created_time(); end it here. + if self.trace_ctx.tracing_enable: +@@ -429,6 +480,7 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + def set_api_server_dispatch_time(self, ts=None): + ts = ts or time.perf_counter() + self.api_server_dispatch_time = ts ++ self._roctx("dispatch") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_start( +@@ -440,6 +492,7 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + def set_api_server_dispatch_finish_time(self, ts=None): + ts = ts or time.perf_counter() + self.api_server_dispatch_finish_time = ts ++ self._roctx("dispatch_finish") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_end( +@@ -452,6 +505,7 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + def set_response_sent_to_client_time(self, ts=None): + ts = ts or time.perf_counter() + self.response_sent_to_client_time = ts ++ self._roctx("response_sent") + + def get_interval(self): + return time.perf_counter() - self.last_time +@@ -476,10 +530,30 @@ class APIServerReqTimeStats(ReqTimeStatsBase): + meta_info["request_received_ts"] = convert_time_to_realtime( + self.created_time + ) ++ # Surface the FULL APIServerReqTimeStats timeline (all monotonic ++ # perf_counter fields -> wall-clock seconds). Previously only the four ++ # *_ts below were emitted; the tokenize/dispatch-start/first-token/ ++ # last-token stamps are now emitted too so the client can capture the ++ # complete API-server-layer timeline. Gated (like the rest of this ++ # method) by --enable-metrics at the tokenizer_manager call site. ++ if self.tokenize_finish_time > 0.0: ++ meta_info["tokenize_finish_ts"] = convert_time_to_realtime( ++ self.tokenize_finish_time ++ ) ++ if self.api_server_dispatch_time > 0.0: ++ meta_info["api_server_dispatch_ts"] = convert_time_to_realtime( ++ self.api_server_dispatch_time ++ ) + if self.api_server_dispatch_finish_time > 0.0: + meta_info["api_server_dispatch_finish_ts"] = convert_time_to_realtime( + self.api_server_dispatch_finish_time + ) ++ if self.first_token_time > 0.0: ++ meta_info["first_token_ts"] = convert_time_to_realtime( ++ self.first_token_time ++ ) ++ if self.last_time > 0.0: ++ meta_info["last_token_ts"] = convert_time_to_realtime(self.last_time) + if self.response_sent_to_client_time > 0.0: + meta_info["response_sent_to_client_ts"] = convert_time_to_realtime( + self.response_sent_to_client_time +@@ -534,6 +608,8 @@ class DPControllerReqTimeStats(ReqTimeStatsBase): + dpc_dispatch_time: float = 0.0 + dpc_dispatch_finish_time: float = 0.0 + ++ _ROCTX_LAYER = "dpc" ++ + def __getstate__(self) -> object: + state = {} + # send to Scheduler +@@ -549,6 +625,7 @@ class DPControllerReqTimeStats(ReqTimeStatsBase): + def set_dp_dispatch_time(self, ts=None): + ts = ts or time.perf_counter() + self.dpc_dispatch_time = ts ++ self._roctx("dispatch") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_start( +@@ -560,6 +637,7 @@ class DPControllerReqTimeStats(ReqTimeStatsBase): + def set_dp_dispatch_finish_time(self, ts=None): + ts = ts or time.perf_counter() + self.dpc_dispatch_finish_time = ts ++ self._roctx("dispatch_finish") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_end( +@@ -580,6 +658,8 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + Decode: prealloc_queue -> transfer_queue -> wait_queue -> forward -> completion + """ + ++ _ROCTX_LAYER = "sched" ++ + # Placeholder: not used currently + # propagated from tokenizer/grpc_server or dp controller + created_time: float = 0.0 +@@ -595,6 +675,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + # prefill node, get by time.perf_counter() + prefill_bootstrap_queue_entry_time: float = 0.0 + prefill_transfer_queue_entry_time: float = 0.0 ++ prefill_kv_transfer_start_time: float = 0.0 + prefill_kv_transfer_finish_time: float = 0.0 + + # decode node, get by time.perf_counter() +@@ -640,13 +721,16 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + calibrate_time_diff() + ts = ts or time.perf_counter() + self.scheduler_recv_time = ts ++ self._roctx("recv") + + def set_spec_draft_start_time(self, ts=None): + ts = ts or time.perf_counter() + self.spec_draft_start_time = ts ++ self._roctx("spec_draft_start") + + def set_spec_draft_end_time(self, ts=None): + ts = ts or time.perf_counter() ++ self._roctx("spec_draft_end") + + if self.trace_ctx.tracing_enable: + stage = RequestStage.SPEC_DRAFT +@@ -655,6 +739,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def set_spec_verify_start_time(self, ts=None): + ts = ts or time.perf_counter() + self.spec_verify_start_time = ts ++ self._roctx("spec_verify_start") + + def set_spec_verify_end_time( + self, +@@ -666,6 +751,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + if accepted_tokens is not None: + num_correct_drafts = accepted_tokens + ts = ts or time.perf_counter() ++ self._roctx("spec_verify_end") + + if self.trace_ctx.tracing_enable: + stage = RequestStage.SPEC_VERIFY +@@ -683,9 +769,11 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def set_run_batch_cpu_start_time(self, ts=None, attrs=None): + ts = ts or time.perf_counter() + self.run_batch_cpu_start_time = ts ++ self._roctx("run_batch_cpu_start") + + def set_run_batch_cpu_end_time(self, ts=None, attrs=None): + ts = ts or time.perf_counter() ++ self._roctx("run_batch_cpu_end") + if self.run_batch_cpu_start_time > 0.0: + self.trace_slice( + RequestStage.RUN_BATCH_CPU, self.run_batch_cpu_start_time, ts, attrs +@@ -700,6 +788,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + self.last_chunked_prefill_finish_time = 0.0 + self.last_decode_finish_time = 0.0 + self.last_decode_scheduled_time = 0.0 ++ self._roctx("retract") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_event("retract", 1, convert_time_to_realtime_ns(ts)) +@@ -710,6 +799,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + self.prefill_finished_time = 0.0 + self.completion_time = 0.0 + self.prefill_transfer_queue_entry_time = 0.0 ++ self.prefill_kv_transfer_start_time = 0.0 + self.prefill_kv_transfer_finish_time = 0.0 + self.last_forward_entry_time = 0.0 + self.last_prefill_finished_time = 0.0 +@@ -735,12 +825,14 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + self.set_retract_time(ts) + + self.wait_queue_entry_time = ts ++ self._roctx("wait_queue_entry") + + def set_forward_entry_time(self, ts=None): + ts = ts or time.perf_counter() + if self.forward_entry_time == 0.0: + self.forward_entry_time = ts + self.last_forward_entry_time = ts ++ self._roctx("forward_entry") + + if self.enable_metrics: + self.metrics_collector.observe_queue_time(self.get_queueing_time()) +@@ -774,6 +866,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + ts = ts or time.perf_counter() + last_time = self.last_chunked_prefill_finish_time + self.last_chunked_prefill_finish_time = ts ++ self._roctx("chunked_prefill_finish") + + if last_time == 0.0: + last_time = self.last_forward_entry_time +@@ -787,6 +880,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + if self.prefill_finished_time == 0.0: + self.prefill_finished_time = ts + self.last_prefill_finished_time = ts ++ self._roctx("prefill_finished") + + stage = RequestStage.PREFILL_FORWARD + self.observe_per_stage_req_latency(stage, ts - self.last_forward_entry_time) +@@ -829,6 +923,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + ts = ts or time.perf_counter() + last_time = self.last_decode_finish_time + self.last_decode_finish_time = ts ++ self._roctx("decode_finish") + + if self.enable_metrics or self.trace_ctx.tracing_enable: + if last_time == 0.0: +@@ -874,10 +969,12 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + + if forward_mode.is_decode(): + self.last_decode_scheduled_time = ts ++ self._roctx("scheduled") + + def set_completion_time(self, ts=None): + ts = ts or time.perf_counter() + self.completion_time = ts ++ self._roctx("completion") + + if self.trace_ctx.tracing_enable: + self.trace_ctx.abort() +@@ -960,10 +1057,12 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + ts = ts or time.perf_counter() + self.set_completion_time(ts) + self.forward_entry_time = ts ++ self._roctx("quick_finish") + + def set_prefill_bootstrap_queue_entry_time(self, ts=None): + ts = ts or time.perf_counter() + self.prefill_bootstrap_queue_entry_time = ts ++ self._roctx("prefill_bootstrap_queue_entry") + + stage = RequestStage.PREFILL_PREPARE + self.observe_per_stage_req_latency(stage, ts - self.scheduler_recv_time) +@@ -972,10 +1071,21 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def set_prefill_transfer_queue_entry_time(self, ts=None): + ts = ts or time.perf_counter() + self.prefill_transfer_queue_entry_time = ts ++ self._roctx("prefill_transfer_queue_entry") ++ ++ def set_prefill_kv_transfer_start_time(self, ts=None): ++ # One-shot: stamp only the FIRST time a KV chunk is sent for this request ++ # (captures the start of the KV transfer window for chunked prefill). ++ if self.prefill_kv_transfer_start_time == 0.0: ++ self.prefill_kv_transfer_start_time = ( ++ ts if ts is not None else time.perf_counter() ++ ) ++ self._roctx("prefill_kv_transfer_start") + + def set_prefill_kv_transfer_finish_time(self, ts=None): + ts = ts or time.perf_counter() + self.prefill_kv_transfer_finish_time = ts ++ self._roctx("prefill_kv_transfer_finish") + + stage = RequestStage.PREFILL_TRANSFER_KV_CACHE + self.observe_per_stage_req_latency( +@@ -986,6 +1096,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def set_decode_prealloc_queue_entry_time(self, ts=None): + ts = ts or time.perf_counter() + self.decode_prealloc_queue_entry_time = ts ++ self._roctx("decode_prealloc_queue_entry") + + stage = RequestStage.DECODE_PREPARE + self.observe_per_stage_req_latency(stage, ts - self.scheduler_recv_time) +@@ -994,6 +1105,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def set_decode_transfer_queue_entry_time(self, ts=None): + ts = ts or time.perf_counter() + self.decode_transfer_queue_entry_time = ts ++ self._roctx("decode_transfer_queue_entry") + + stage = RequestStage.DECODE_BOOTSTRAP + self.observe_per_stage_req_latency( +@@ -1015,10 +1127,12 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + ts = ts or time.perf_counter() + if self.bootstrap_done_time == 0.0: + self.bootstrap_done_time = ts ++ self._roctx("bootstrap_done") + + def set_decode_prebuilt_finish_time(self, ts=None): + ts = ts or time.perf_counter() + self.decode_prebuilt_finish_time = ts ++ self._roctx("decode_prebuilt_finish") + + stage = RequestStage.DECODE_FAKE_OUTPUT + self.observe_per_stage_req_latency(stage, ts - self.last_forward_entry_time) +diff --git a/python/sglang/srt/observability/sglang_roctx.py b/python/sglang/srt/observability/sglang_roctx.py +new file mode 100644 +index 0000000000..69dca495dc +--- /dev/null ++++ b/python/sglang/srt/observability/sglang_roctx.py +@@ -0,0 +1,115 @@ ++# Copyright 2023-2024 SGLang Team ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++# ============================================================================== ++"""roctx marker helper for SGLang request-time-stats instrumentation. ++ ++ADDITIVE / OPT-IN. This module is part of the private roctx fork; it is a pure ++addition (no upstream file is modified destructively). Every marker is gated ++behind the ``SGLANG_ROCTX=1`` environment variable and is an exception-safe ++no-op when the gate is unset → **zero behavioral change and zero overhead when ++off**. ++ ++Why roctx (and not torch/kineto): roctx markers live in roctracer's external ++"roctx" domain. torch/kineto does NOT ingest that domain, so these marks are ++INVISIBLE to a torch profiler trace. They only surface under ++``rocprofv3 --marker-trace`` (or RTL-lite's roctx shim). Capture accordingly. ++ ++Mechanism: we bind ``roctxMarkA`` from ``libroctx64.so`` via ctypes (no need to ++link the SGLang/MORI build against libroctx). ``roctxMarkA`` emits an *instant* ++marker (a point on the timeline), which is the right primitive here because a ++request's lifecycle stamps fire across multiple threads/processes (scheduler ++event loop, tokenizer/API frontend, TP workers) — a push/pop range would have ++to begin and end on the same thread, which these stamps do not. Read a ++start/end window as ``ts(end_mark) - ts(start_mark)`` for a matching request id. ++""" ++ ++from __future__ import annotations ++ ++import os ++ ++__all__ = ["ROCTX_ENABLED", "roctx_mark", "roctx_available"] ++ ++# Gate read once at import. Accept "1" (and, defensively, common truthy spellings). ++ROCTX_ENABLED = os.environ.get("SGLANG_ROCTX", "0").strip().lower() in ( ++ "1", ++ "true", ++ "yes", ++ "on", ++) ++ ++# _emit(name: bytes) -> None. Default is a no-op; replaced by the real ctypes ++# binding below only when the gate is on AND libroctx64.so loads cleanly. ++_emit = None # type: ignore[assignment] ++ ++ ++def _init_backend() -> None: ++ """Bind roctxMarkA from libroctx64.so. Best-effort; leaves _emit=None on failure.""" ++ global _emit ++ try: ++ import ctypes ++ ++ # IMPORTANT: rocprofiler-sdk's `rocprofv3 --marker-trace` only intercepts ++ # the rocprofiler-sdk ROCTx library (librocprofiler-sdk-roctx.so), NOT the ++ # legacy roctracer libroctx64.so. Proven on ROCm 7.0: marks via ++ # librocprofiler-sdk-roctx.so render as slices in the pftrace; libroctx64.so ++ # marks do not. So prefer the sdk lib; fall back to libroctx64 (e.g. for an ++ # RTL/roctracer capture context). All ship in /opt/rocm*/lib. ++ lib = None ++ for cand in ( ++ "librocprofiler-sdk-roctx.so", ++ "librocprofiler-sdk-roctx.so.1", ++ "libroctx64.so", ++ "libroctx64.so.4", ++ ): ++ try: ++ lib = ctypes.CDLL(cand) ++ break ++ except OSError: ++ continue ++ if lib is None: ++ return ++ fn = lib.roctxMarkA ++ fn.argtypes = [ctypes.c_char_p] ++ fn.restype = None ++ ++ def _emit_impl(name_bytes, _fn=fn): # bound default avoids global lookups ++ _fn(name_bytes) ++ ++ _emit = _emit_impl ++ except Exception: ++ # Any failure (missing lib, missing symbol, etc.) → stay a no-op. ++ _emit = None ++ ++ ++if ROCTX_ENABLED: ++ _init_backend() ++ ++ ++def roctx_available() -> bool: ++ """True iff the gate is on and the roctx backend bound successfully.""" ++ return ROCTX_ENABLED and _emit is not None ++ ++ ++def roctx_mark(message: str) -> None: ++ """Emit an instant roctx marker. No-op (and never raises) when disabled. ++ ++ Safe to call on any hot path: when ``SGLANG_ROCTX`` is unset this returns ++ immediately after a single boolean check and the backend is never bound. ++ """ ++ if not ROCTX_ENABLED or _emit is None: ++ return ++ try: ++ _emit(message.encode("ascii", "replace")) ++ except Exception: ++ # Observability must never break serving. ++ pass diff --git a/scripts/sglang_disagg/moriio_profiling/process_kernels.sh b/scripts/sglang_disagg/moriio_profiling/process_kernels.sh new file mode 100644 index 00000000..d382b32e --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/process_kernels.sh @@ -0,0 +1,443 @@ +#!/bin/bash +# Strict post-run orchestration for integrated SGLang MoRI I/O profiling. +set -euo pipefail + +usage() { + echo "usage: process_kernels.sh JOBID" >&2 + echo " process_kernels.sh {run|verify|trace|analyze} JOBID" >&2 + exit 2 +} + +case "${1:-}" in + run|verify|trace|analyze) + OP="$1" + J="${2:-}" + [ "$#" -eq 2 ] || usage + ;; + "") + usage + ;; + *) + OP="run" + J="$1" + [ "$#" -eq 1 ] || usage + ;; +esac + +D="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +HR="${HR:-/shared_inference/${USER:-aarai}/model_blog_logs/$J}" +XP="${XP:-1}" +YD="${YD:-1}" +E="${ROCPROF_EXPECT_PER_NODE:-8}" +OUT="${OUT:-$D/artifacts/pull_${J}}" +TOOLS="${TRACE_TOOLS:-$D/trace_tools.py}" +RUN_MORI="${RUN_MORI:-0}" + +MULTI_SWEEP=0 +SWEEP_KEYS=() +SWEEP_CLIENTS=() +SWEEP_MANIFESTS=() +SWEEP_PREFIXES=() +SWEEP_OUTPUTS=() + +PREFILL_DIRS=() +PREFILL_LOGS=() +DECODE_DIRS=() +DECODE_LOGS=() +DIRS=() +for ((i=0; i&2 + return 2 + fi + name=$(basename "$dir") + role="prefill" + [[ "$name" == *decode* ]] && role="decode" + + unset kernels markers results + declare -A kernels=(["__none"]=1) + declare -A markers=(["__none"]=1) + declare -A results=(["__none"]=1) + local kernel_count=0 marker_count=0 result_count=0 + shopt -s nullglob + for file in "$dir"/*_kernel_trace.csv; do + kernels["$(pid_of "$file" _kernel_trace.csv)"]=1 + kernel_count=$((kernel_count+1)) + done + for file in "$dir"/*_marker_api_trace.csv; do + markers["$(pid_of "$file" _marker_api_trace.csv)"]=1 + marker_count=$((marker_count+1)) + done + for file in "$dir"/*_results.json; do + results["$(pid_of "$file" _results.json)"]=1 + result_count=$((result_count+1)) + done + shopt -u nullglob + + for suffix in kernel_trace.csv marker_api_trace.csv results.json; do + case "$suffix" in + kernel_trace.csv) count=$kernel_count ;; + marker_api_trace.csv) count=$marker_count ;; + results.json) count=$result_count ;; + esac + if (( count < E )); then + echo "[verify] ERROR: $name has $count/$E $suffix files" >&2 + any_bad=1 + fi + done + + declare -A all_pids=(["__none"]=1) + for pid in "${!kernels[@]}" "${!markers[@]}" "${!results[@]}"; do + [[ -n "$pid" && "$pid" != "__none" ]] && all_pids["$pid"]=1 + done + local ok=0 details="" pid_count=$(( ${#all_pids[@]} - 1 )) + for pid in "${!all_pids[@]}"; do + [ "$pid" = "__none" ] && continue + missing="" + [ -n "${kernels[$pid]:-}" ] || missing="${missing}kernel," + [ -n "${markers[$pid]:-}" ] || missing="${missing}marker," + [ -n "${results[$pid]:-}" ] || missing="${missing}results," + if [ -z "$missing" ]; then + ok=$((ok+1)) + else + details="$details $pid[missing:${missing%,}]" + fi + done + grand_expected=$((grand_expected+E)) + grand_ok=$((grand_ok+ok)) + role_expected["$role"]=$(( ${role_expected[$role]:-0} + E )) + role_ok["$role"]=$(( ${role_ok[$role]:-0} + ok )) + if (( ok == E && pid_count == E )); then + echo " [OK] $name: $ok/$E workers complete" + else + echo " [BAD] $name: $ok/$E complete; pids=$pid_count$details" >&2 + any_bad=1 + fi + done + for role in prefill decode; do + [ -n "${role_expected[$role]:-}" ] || continue + echo "[verify] ${role^^}: ${role_ok[$role]:-0}/${role_expected[$role]}" + done + if (( any_bad )); then + echo "[verify] ERROR: capture incomplete ($grand_ok/$grand_expected)" >&2 + return 1 + fi + echo "[verify] OK: all $grand_ok/$grand_expected workers complete" +} + +sha256_file() { + python3 - "$1" <<'PY' +import hashlib, sys +h = hashlib.sha256() +with open(sys.argv[1], "rb") as fh: + for block in iter(lambda: fh.read(1024 * 1024), b""): + h.update(block) +print(h.hexdigest()) +PY +} + +manifest_prefix() { + python3 - "$1" <<'PY' +import json, sys +with open(sys.argv[1], encoding="utf-8") as fh: + manifest = json.load(fh) +prefix = manifest.get("request_id_prefix") +if not isinstance(prefix, str) or not prefix: + raise SystemExit("missing or invalid request_id_prefix") +print(prefix) +PY +} + +discover_sweeps() { + local fixed_csv="$HR/rocprof_probe_client.csv" + local fixed_manifest="$HR/rocprof_probe_manifest.json" + local keyed_csvs=() keyed_manifests=() file key rid_prefix expected_prefix + declare -A clients=() manifests=() prefixes=() + SWEEP_KEYS=() + SWEEP_CLIENTS=() + SWEEP_MANIFESTS=() + SWEEP_PREFIXES=() + SWEEP_OUTPUTS=() + MULTI_SWEEP=0 + shopt -s nullglob + keyed_csvs=("$HR"/rocprof_probe_client_i*_isl*_osl*_c*.csv) + keyed_manifests=("$HR"/rocprof_probe_manifest_i*_isl*_osl*_c*.json) + shopt -u nullglob + if [[ -e "$fixed_csv" && -e "$fixed_manifest" && ( ${#keyed_csvs[@]} -gt 0 || ${#keyed_manifests[@]} -gt 0 ) ]]; then + echo "[process_kernels.sh] ERROR: fixed and keyed request-correlation artifacts are ambiguous" >&2 + return 1 + fi + if [[ -e "$fixed_csv" || -e "$fixed_manifest" ]]; then + if [[ ! -e "$fixed_csv" || ! -e "$fixed_manifest" ]]; then + echo "[process_kernels.sh] ERROR: orphan fixed request-correlation artifact" >&2 + return 1 + fi + if [[ ! -s "$fixed_csv" || ! -s "$fixed_manifest" ]]; then + echo "[process_kernels.sh] ERROR: empty fixed request-correlation artifact" >&2 + return 1 + fi + if ! rid_prefix=$(manifest_prefix "$fixed_manifest"); then + echo "[process_kernels.sh] ERROR: invalid request_id_prefix in $fixed_manifest" >&2 + return 1 + fi + if [[ "$rid_prefix" != "profile-${J}" ]]; then + echo "[process_kernels.sh] ERROR: fixed manifest request_id_prefix mismatch: $rid_prefix" >&2 + return 1 + fi + SWEEP_KEYS=("fixed") + SWEEP_CLIENTS=("$fixed_csv") + SWEEP_MANIFESTS=("$fixed_manifest") + SWEEP_PREFIXES=("$rid_prefix") + SWEEP_OUTPUTS=("$OUT") + return 0 + fi + if (( ${#keyed_csvs[@]} == 0 && ${#keyed_manifests[@]} == 0 )); then + echo "[process_kernels.sh] ERROR: normal benchmark request-correlation outputs are missing" >&2 + return 1 + fi + for file in "${keyed_csvs[@]}"; do + if [[ ! -s "$file" ]]; then + echo "[process_kernels.sh] ERROR: empty keyed client CSV: $file" >&2 + return 1 + fi + key=$(basename "$file") + key=${key#rocprof_probe_client_} + key=${key%.csv} + if [[ ! "$key" =~ ^i[1-9][0-9]*_isl[0-9]+_osl[0-9]+_c[0-9]+$ ]]; then + echo "[process_kernels.sh] ERROR: invalid client sweep key: $key" >&2 + return 1 + fi + if [[ -n "${clients[$key]+x}" ]]; then + echo "[process_kernels.sh] ERROR: duplicate client sweep key: $key" >&2 + return 1 + fi + clients[$key]="$file" + done + for file in "${keyed_manifests[@]}"; do + if [[ ! -s "$file" ]]; then + echo "[process_kernels.sh] ERROR: empty keyed manifest: $file" >&2 + return 1 + fi + key=$(basename "$file") + key=${key#rocprof_probe_manifest_} + key=${key%.json} + if [[ ! "$key" =~ ^i[1-9][0-9]*_isl[0-9]+_osl[0-9]+_c[0-9]+$ ]]; then + echo "[process_kernels.sh] ERROR: invalid manifest sweep key: $key" >&2 + return 1 + fi + if [[ -n "${manifests[$key]+x}" ]]; then + echo "[process_kernels.sh] ERROR: duplicate manifest sweep key: $key" >&2 + return 1 + fi + manifests[$key]="$file" + done + for file in "${keyed_csvs[@]}"; do + key=$(basename "$file") + key=${key#rocprof_probe_client_} + key=${key%.csv} + if [[ -z "${manifests[$key]+x}" ]]; then + echo "[process_kernels.sh] ERROR: orphan client CSV for sweep $key" >&2 + return 1 + fi + if ! rid_prefix=$(manifest_prefix "${manifests[$key]}"); then + echo "[process_kernels.sh] ERROR: invalid request_id_prefix for sweep $key" >&2 + return 1 + fi + expected_prefix="profile-${J}-${key}" + if [[ "$rid_prefix" != "$expected_prefix" ]]; then + echo "[process_kernels.sh] ERROR: request_id_prefix mismatch for sweep $key: $rid_prefix" >&2 + return 1 + fi + if [[ -n "${prefixes[$rid_prefix]+x}" ]]; then + echo "[process_kernels.sh] ERROR: duplicate request_id_prefix: $rid_prefix" >&2 + return 1 + fi + prefixes[$rid_prefix]=1 + SWEEP_KEYS+=("$key") + SWEEP_CLIENTS+=("${clients[$key]}") + SWEEP_MANIFESTS+=("${manifests[$key]}") + SWEEP_PREFIXES+=("$rid_prefix") + SWEEP_OUTPUTS+=("$OUT/sweeps/$key") + done + for key in "${!manifests[@]}"; do + if [[ -z "${clients[$key]+x}" ]]; then + echo "[process_kernels.sh] ERROR: orphan manifest for sweep $key" >&2 + return 1 + fi + done + MULTI_SWEEP=1 +} + +trace_sweep() { + local sweep_label="$1" rid_prefix="$2" client_csv="$3" client_manifest="$4" sweep_out="$5" + local rid_filter="${rid_prefix}-" + trap 'rc=$?; echo "[process_kernels.sh] ERROR: sweep $sweep_label failed (status=$rc)" >&2; exit "$rc"' ERR + mkdir -p "$OUT/_staging" + local stage="$OUT/_staging/run.$$.$sweep_label" + mkdir -p "$stage" + trap 'rm -rf "$stage"' RETURN + + local trace="$stage/trace.json" + local expected=$(( (XP + YD) * E )) + python3 "$TOOLS" build-trace \ + --prefill-dir "${PREFILL_DIRS[@]}" \ + --decode-dir "${DECODE_DIRS[@]}" \ + --request-logs "${PREFILL_LOGS[@]}" "${DECODE_LOGS[@]}" \ + --rid-prefix "$rid_filter" \ + --out "$trace" --expect-workers "$expected" + + local combined="$stage/roctx_mori_clean_prefill_decode_${J}.json" + local probe="$stage/roctx_mori_clean_probe_only_${J}.json" + mv "$trace" "$combined" + cp "$combined" "$probe" + local combined_hash probe_hash + combined_hash=$(sha256_file "$combined") + probe_hash=$(sha256_file "$probe") + if [ "$combined_hash" != "$probe_hash" ]; then + echo "[process_kernels.sh] ERROR: compatibility trace hashes differ" >&2 + return 1 + fi + + if [[ "$RUN_MORI" == "1" ]]; then + python3 "$TOOLS" correlate \ + --prefill-dir "${PREFILL_DIRS[@]}" \ + --prefill-logs "${PREFILL_LOGS[@]}" \ + --out-csv "$stage/request_mori_map_${J}.csv" \ + --out-summary "$stage/request_mori_map_${J}.md" \ + --rid-prefix "$rid_filter" --require-complete + else + echo "[process_kernels.sh] RUN_MORI=0: skipping MoRI request/KV correlation" + fi + + cp "$client_csv" "$client_manifest" "$stage/" + local reqstats_mode=() + [[ "$RUN_MORI" == "1" ]] || reqstats_mode+=(--no-mori) + python3 "$TOOLS" reqstats \ + --job "$J" --xp "$XP" --yd "$YD" --out-dir "$stage" --splits \ + --prefill-dirs "${PREFILL_DIRS[@]}" \ + --decode-dirs "${DECODE_DIRS[@]}" \ + --prefill-logs "${PREFILL_LOGS[@]}" \ + --decode-logs "${DECODE_LOGS[@]}" \ + --client-csv "$client_csv" --client-manifest "$client_manifest" \ + --require-data --require-client \ + --rid-prefix "$rid_filter" \ + "${reqstats_mode[@]}" + + local required=( + "$combined" "$probe" + "$stage/reqstats_per_request_${J}.csv" + "$stage/reqstats_per_request_${J}_prefill.csv" + "$stage/reqstats_per_request_${J}_decode.csv" + "$stage/$(basename "$client_csv")" + "$stage/$(basename "$client_manifest")" + ) + if [[ "$RUN_MORI" == "1" ]]; then + required+=( + "$stage/request_mori_map_${J}.csv" + "$stage/request_mori_map_${J}.md" + ) + fi + local file + for file in "${required[@]}"; do + [ -s "$file" ] || { + echo "[process_kernels.sh] ERROR: required output missing: $file" >&2 + return 1 + } + done + mkdir -p "$sweep_out" + for file in "${required[@]}"; do + mv -f "$file" "$sweep_out/$(basename "$file")" + done + rmdir "$stage" + trap - RETURN + trap - ERR + echo "[process_kernels.sh] trace SHA-256=$combined_hash" + if [[ "$RUN_MORI" == "1" ]]; then + echo "[process_kernels.sh] strict trace/request/MoRI outputs finalized in $sweep_out" + else + echo "[process_kernels.sh] SGLang trace/request outputs finalized without MoRI correlation in $sweep_out" + fi +} + +trace_phase() { + verify_capture + discover_sweeps + local index + for ((index=0; index<${#SWEEP_KEYS[@]}; index++)); do + echo "[process_kernels.sh] processing sweep ${SWEEP_KEYS[$index]}" + trace_sweep "${SWEEP_KEYS[$index]}" "${SWEEP_PREFIXES[$index]}" \ + "${SWEEP_CLIENTS[$index]}" "${SWEEP_MANIFESTS[$index]}" "${SWEEP_OUTPUTS[$index]}" + done +} + +analyze_phase() { + mkdir -p "$OUT/analyze_phase" "$OUT/_staging" + local failures=0 statuses=() dir label temporary final + for dir in "${DIRS[@]}"; do + label=$(basename "$dir") + temporary="$OUT/_staging/analyze_${label}.$$" + final="$OUT/analyze_phase/$label" + rm -rf "$temporary" + if python3 "$TOOLS" analyze "$dir" "$temporary" "$label"; then + rm -rf "$final" + mv "$temporary" "$final" + statuses+=("$label=ok") + else + rm -rf "$temporary" + echo "[process_kernels.sh] WARN: analysis failed for $label" >&2 + statuses+=("$label=failed") + failures=$((failures+1)) + fi + done + echo "[process_kernels.sh] analysis statuses=${statuses[*]} (best-effort failures=$failures)" + return 0 +} + +case "$OP" in + verify) + verify_capture + ;; + trace) + trace_phase + ;; + analyze) + analyze_phase + ;; + run) + trace_phase + analyze_phase + echo "===== [process_kernels.sh] capture complete =====" + echo "workers=$(( (XP + YD) * E ))/$(( (XP + YD) * E ))" + if (( MULTI_SWEEP )); then + echo "sweeps=$OUT/sweeps (${#SWEEP_KEYS[@]})" + else + echo "combined=$OUT/roctx_mori_clean_prefill_decode_${J}.json" + echo "probe_only=$OUT/roctx_mori_clean_probe_only_${J}.json" + [[ "$RUN_MORI" == "1" ]] && echo "request_mori_map=$OUT/request_mori_map_${J}.csv" + echo "request_stats=$OUT/reqstats_per_request_${J}.csv" + fi + ;; +esac diff --git a/scripts/sglang_disagg/moriio_profiling/trace_tools.py b/scripts/sglang_disagg/moriio_profiling/trace_tools.py new file mode 100644 index 00000000..cd0573c1 --- /dev/null +++ b/scripts/sglang_disagg/moriio_profiling/trace_tools.py @@ -0,0 +1,2706 @@ +#!/usr/bin/env python3 +"""Offline ROCTX trace construction, request attribution, and kernel analysis.""" +import argparse +import csv +import glob +import hashlib +import json +import math +import os +import re +import shutil +import statistics +import subprocess +import sys +import tempfile +from collections import defaultdict +from pathlib import Path + +_csv_limit = sys.maxsize +while True: + try: + csv.field_size_limit(_csv_limit) + break + except OverflowError: + _csv_limit //= 10 + +# ---- clean-trace construction ---- +_MAP_RE = re.compile('\\broom=(\\d+)\\s+uid=(\\d+)') +_IDV_RE = re.compile('\\bid=(\\d+)') +_BYTES_RE = re.compile('\\bbytes=(\\d+)') +_WORKER_GPU0_RE = re.compile(r'\bProcess (?P\d+) gpu_id 0 is running on CPUs:') +_WORKER_GPU_RE = re.compile( + r'(?:\bDP(?P\d+)\s+TP(?P\d+)\s+' + r'EP(?P\d+)\]\s+)?Process (?P\d+) ' + r'gpu_id (?P\d+) is running on CPUs:' +) +_ANALYSIS_PID_RE = re.compile(r'^.+_(?P\d+)_(?:results\.pftrace|kernel_trace\.csv)$') +_NODE_CAPTURE_RE = re.compile(r'^rocprof_(?Pprefill|decode)_NODE(?P\d+)$') + +def iter_kernels(path): + with open(path, newline='') as f: + for r in csv.DictReader(f): + try: + s = int(r['Start_Timestamp']) + e = int(r['End_Timestamp']) + except Exception: + continue + if _duration_or_none(e - s) is None: + continue + yield (s, e, r.get('Kernel_Name', 'kernel')) + +def read_kernels(path): + return list(iter_kernels(path)) + +def iter_marks(path): + with open(path, newline='') as f: + for r in csv.DictReader(f): + fn = r.get('Function', '') + if not fn.startswith('reqstats'): + continue + try: + s = int(r['Start_Timestamp']) + except Exception: + continue + m = re.search('id=(\\S+)', fn) + room = m.group(1) if m else '' + event = fn.split(' id=')[0] + yield (s, event, room) + +def read_marks(path): + return list(iter_marks(path)) + +def iter_mori(path): + """MORI-IO host-send roctx RANGES (mori.*): duration slices (Start..End). + + Excludes the `mori.map` correlation marks (instant, zero-dur): those are consumed by + read_map for attribution, not rendered as transfer ranges.""" + out = [] + with open(path, newline='') as f: + for r in csv.DictReader(f): + fn = r.get('Function', '') + if not fn.startswith('mori'): + continue + if fn.startswith('mori.map'): + continue + try: + s = int(r['Start_Timestamp']) + e = int(r['End_Timestamp']) + except Exception: + continue + yield (s, e, fn) + +def read_mori(path): + return list(iter_mori(path)) + +def read_map(path, lo=None, hi=None): + """Parse `mori.map room=R uid=U` marks -> {uid(str): room(str)} for one pid. + Empty on pre-patch traces (callers then fall back to timestamp containment).""" + out = {} + with open(path, newline='') as f: + for r in csv.DictReader(f): + fn = r.get('Function', '') + if not fn.startswith('mori.map'): + continue + if lo is not None or hi is not None: + try: + ts = int(r['Start_Timestamp']) + except Exception: + continue + if lo is not None and ts < lo: + continue + if hi is not None and ts > hi: + continue + m = _MAP_RE.search(fn) + if m: + out[m.group(2)] = m.group(1) + return out + +def mori_rid_assign(probe, mori_ranges, uid_to_room=None): + """Return a list (parallel to mori_ranges) of the owning rid (or None). + + Prefer EXACT-by-id attribution via the `mori.map` uid->room map (uid_to_room) when a + range carries an `id=N` that maps to a real room; otherwise fall back to the original + per-pid timestamp-containment rule (inner [prefill_kv_transfer_start, + prefill_kv_transfer_finish] window, tightest enclosing, outer [forward_entry, + completion] fallback). The fallback keeps pre-patch traces (no mori.map) working and + doubles as a cross-check. probe: (ts,event,room) for THIS pid; mori_ranges: (s,e,name).""" + uid_to_room = uid_to_room or {} + rooms = {} + for s, event, room in probe: + if room in ('', '0'): + continue + leaf = event.split('.')[-1] + rooms.setdefault(room, {})[leaf] = s + inner, outer = ({}, {}) + for room, st in rooms.items(): + if 'prefill_kv_transfer_start' in st and 'prefill_kv_transfer_finish' in st: + a, b = (st['prefill_kv_transfer_start'], st['prefill_kv_transfer_finish']) + if b >= a: + inner[room] = (a, b) + if 'forward_entry' in st and 'completion' in st: + a, b = (st['forward_entry'], st['completion']) + if b >= a: + outer[room] = (a, b) + rids = [] + for s, e, name in mori_ranges: + m = _IDV_RE.search(name) + room = uid_to_room.get(m.group(1)) if m else None + if room is not None and room not in ('', '0'): + rids.append(room) + continue + mid = (s + e) // 2 + hits = [(room, iv[1] - iv[0]) for room, iv in inner.items() if iv[0] <= mid <= iv[1]] + if not hits: + hits = [(room, iv[1] - iv[0]) for room, iv in outer.items() if iv[0] <= mid <= iv[1]] + if hits: + hits.sort(key=lambda x: x[1]) + rids.append(hits[0][0]) + else: + rids.append(None) + return rids + +def emit_kv_stack(ev, kv_intervals, base_pid, kvt, gtp, us, max_rows=50): + """Emit overlapping KV ranges on greedy-stacked complete-event rows. + ``max_rows`` bounds reserved track IDs; overflow uses the last row. + Returns ``(emitted_count, maximum_depth)``. + """ + kv_intervals.sort(key=lambda x: x[0]) + last_end = [] + nkvt = 0 + kv_depth = 0 + warned = False + for s, e, label, rid, nbytes in kv_intervals: + k = next((idx for idx in range(len(last_end)) if last_end[idx] <= s), None) + if k is None: + k = len(last_end) + last_end.append(e) + else: + last_end[k] = e + if k >= max_rows: + if not warned: + print(f'[per-GPU][WARN] TP{gtp} kv_transfer stack depth {k + 1} exceeds {max_rows} sub-rows -- clamping to row {max_rows - 1} (tid {kvt + max_rows - 1})') + warned = True + k_emit = max_rows - 1 + else: + k_emit = k + _a = {} + if rid: + _a['rid'] = rid + if nbytes: + _a['bytes'] = nbytes + ev.append({'ph': 'X', 'name': label, 'ts': us(s), 'dur': max((e - s) / 1000.0, 0.0), 'pid': base_pid, 'tid': kvt + k_emit, 'args': _a}) + nkvt += 1 + kv_depth = max(kv_depth, k + 1) + for k in range(1, min(kv_depth, max_rows)): + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': kvt + k, 'args': {'name': f'TP{gtp} MORI-IO KV transfer #{k + 1}'}}) + return (nkvt, kv_depth) + +def emit_stacked(ev, intervals, base_pid, base_tid, us, row_label, max_rows=12): + """Greedily assign intervals to the lowest non-overlapping row. + Returns ``(emitted_count, maximum_depth)``. + """ + intervals.sort(key=lambda x: x[0]) + last_end = [] + n = 0 + depth = 0 + for s, e, name, args in intervals: + k = next((idx for idx in range(len(last_end)) if last_end[idx] <= s), None) + if k is None: + k = len(last_end) + last_end.append(e) + else: + last_end[k] = e + k_emit = k if k < max_rows else max_rows - 1 + ev.append({'ph': 'X', 'name': name, 'ts': us(s), 'dur': max((e - s) / 1000.0, 0.0), 'pid': base_pid, 'tid': base_tid + k_emit, 'args': args}) + n += 1 + depth = max(depth, k + 1) + for k in range(1, min(depth, max_rows)): + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': base_tid + k, 'args': {'name': row_label(k)}}) + return (n, depth) + +def build_stage_span_intervals(byroom, gap_mult=3.0): + """Build adjacent stage spans, decode-step indices, and long-gap flags. + A gap is flagged above ``gap_mult`` times its pair-type median; pair types with + fewer than two samples are never flagged. Existing pairing/order is preserved. + """ + durations_by_type = {} + for room, seq in byroom.items(): + for i in range(len(seq) - 1): + (s0, _e0), (s1, _e1) = (seq[i], seq[i + 1]) + if s1 <= s0: + continue + short0 = _e0.split('.')[-1] + short1 = _e1.split('.')[-1] + durations_by_type.setdefault(f'{short0}->{short1}', []).append(s1 - s0) + thresholds = {pt: gap_mult * statistics.median(d) for pt, d in durations_by_type.items() if len(d) >= 2} + out = [] + for room, seq in byroom.items(): + dc = 0 + for i in range(len(seq) - 1): + (s0, e0), (s1, e1) = (seq[i], seq[i + 1]) + if s1 <= s0: + continue + short0 = e0.split('.')[-1] + short1 = e1.split('.')[-1] + pair_type = f'{short0}->{short1}' + args = {'room': room, 'pair_type': pair_type} + if short1 == 'decode_finish' and short0 in ('decode_finish', 'decode_prebuilt_finish'): + name = f'{pair_type}[step={dc}]' + args['step'] = dc + dc += 1 + else: + name = pair_type + thr = thresholds.get(pair_type) + if thr is not None and s1 - s0 > thr: + args['long_gap'] = 1 + args['median_ns'] = int(statistics.median(durations_by_type[pair_type])) + name += ' [LONG GAP]' + out.append((s0, s1, name, args)) + return out + +def build_decode_ct_counters(probe, label, base_pid, us): + """Emit decode progress and admission-backlog counters for one worker. + Returns no counters when the worker has no decode_finish markers. + """ + total = 0 + room_dc = {} + backlog = set() + out = [] + for s, event, room in probe: + if room in ('', '0'): + continue + leaf = event.split('.')[-1] + if leaf == 'decode_finish': + total += 1 + c = room_dc.get(room, 0) + 1 + room_dc[room] = c + if c == 1: + backlog.add(room) + elif room in backlog: + backlog.discard(room) + out.append({'ph': 'C', 'name': f'{label} decode_ct', 'ts': us(s), 'pid': base_pid, 'args': {'decode_ct': total}}) + out.append({'ph': 'C', 'name': f'{label} decode_admission_backlog', 'ts': us(s), 'pid': base_pid, 'args': {'decode_admission_backlog': len(backlog)}}) + elif leaf == 'completion': + backlog.discard(room) + room_dc.pop(room, None) + return out + +def _pid_from_marker(path): + m = re.search('_(\\d+)_marker_api_trace\\.csv$', os.path.basename(path)) + return m.group(1) if m else os.path.basename(path) + +def discover_pairs_in_dir(d): + """Topology-aware pair discovery, keyed on the MARKER csv (always present) with the + kernel csv treated as OPTIONAL (None if absent). This is what makes the builder work + for BOTH kernel+marker captures AND marker-trace-only captures -- the EP / full-V3 + 2P2D runs use `ROCPROF_FLAGS=--marker-trace` only (a kernel trace would be millions of + rows/worker at conc-32), so they have no *_kernel_trace.csv. Returns sorted + [(pid, kernel_csv_or_None, marker_csv)] for every TP worker in the dir.""" + pairs = [] + for mk in glob.glob(os.path.join(d, '*_marker_api_trace.csv')): + kc = mk[:-len('_marker_api_trace.csv')] + '_kernel_trace.csv' + pairs.append((_pid_from_marker(mk), kc if os.path.exists(kc) else None, mk)) + pairs.sort(key=lambda x: x[0]) + return pairs + +def engine_events(kernel_csv, marker_csv, pid, proc_name, pad_ns=2000000, probe_only=False, gap_ns=30000000000): + marks = read_marks(marker_csv) + probe = [m for m in marks if m[2] not in ('', '0')] + probe.sort(key=lambda x: x[0]) + if probe_only and probe: + clusters, cur = ([], [probe[0]]) + for prev, r in zip(probe, probe[1:]): + if r[0] - prev[0] > gap_ns: + clusters.append(cur) + cur = [r] + else: + cur.append(r) + clusters.append(cur) + probe = clusters[-1] + ev = [] + meta = [{'ph': 'M', 'name': 'process_name', 'pid': pid, 'tid': 0, 'args': {'name': proc_name}}, {'ph': 'M', 'name': 'thread_name', 'pid': pid, 'tid': 1, 'args': {'name': 'GPU kernels'}}, {'ph': 'M', 'name': 'thread_name', 'pid': pid, 'tid': 2, 'args': {'name': 'reqstats markers'}}, {'ph': 'M', 'name': 'thread_name', 'pid': pid, 'tid': 3, 'args': {'name': 'per-request stage spans'}}, {'ph': 'M', 'name': 'thread_name', 'pid': pid, 'tid': 4, 'args': {'name': 'MORI-IO host send (ibv_post_send)'}}] + if not probe: + probe = marks + t_lo = min((s for s, _, _ in probe)) - pad_ns + t_hi = max((s for s, _, _ in probe)) + pad_ns + t0 = t_lo + us = lambda ns: (ns - t0) / 1000.0 + nk = 0 + for s, e, name in read_kernels(kernel_csv): + if s < t_lo or s > t_hi: + continue + ev.append({'ph': 'X', 'name': name, 'ts': us(s), 'dur': (e - s) / 1000.0, 'pid': pid, 'tid': 1}) + nk += 1 + for s, event, room in probe: + ev.append({'ph': 'i', 'name': event + (' room=' + room if room else ''), 'ts': us(s), 'pid': pid, 'tid': 2, 's': 't'}) + byroom = {} + for s, event, room in sorted(probe, key=lambda x: x[0]): + byroom.setdefault(room, []).append((s, event)) + nspan = 0 + for s0, s1, name, args in build_stage_span_intervals(byroom): + ev.append({'ph': 'X', 'name': name, 'ts': us(s0), 'dur': (s1 - s0) / 1000.0, 'pid': pid, 'tid': 3, 'args': args}) + nspan += 1 + ev += build_decode_ct_counters(probe, proc_name, pid, us) + nmori = 0 + for s, e, name in read_mori(marker_csv): + if s < t_lo or s > t_hi: + continue + ev.append({'ph': 'X', 'name': name, 'ts': us(s), 'dur': max((e - s) / 1000.0, 0.05), 'pid': pid, 'tid': 4}) + nmori += 1 + return (meta + ev, nk, len(probe), nspan, nmori) + +def _select_probe(marker_csv, probe_only, rid_rooms=None, gap_ns=30000000000): + if rid_rooms is not None: + return sorted( + (mark for mark in read_marks(marker_csv) if mark[2] in rid_rooms), + key=lambda mark: mark[0], + ) + if not probe_only: + marks = read_marks(marker_csv) + probe = [m for m in marks if m[2] not in ('', '0')] + return probe or marks + probe, prev_s = ([], None) + for mark in iter_marks(marker_csv): + if mark[2] in ('', '0'): + continue + if prev_s is not None and mark[0] - prev_s > gap_ns: + probe = [] + probe.append(mark) + prev_s = mark[0] + probe.sort(key=lambda x: x[0]) + return probe + +def _pid_from(path): + m = re.search('_(\\d+)_kernel_trace\\.csv$', os.path.basename(path)) + return m.group(1) if m else os.path.basename(path) + +def discover_pairs(repr_kernel_csv): + d = os.path.dirname(os.path.abspath(repr_kernel_csv)) + pairs = [] + for k in glob.glob(os.path.join(d, '*_kernel_trace.csv')): + mk = k[:-len('_kernel_trace.csv')] + '_marker_api_trace.csv' + if os.path.exists(mk): + pairs.append((_pid_from(k), k, mk)) + pairs.sort(key=lambda x: x[0]) + return pairs + +def lane_window(pairs, probe_only, pad_ns=2000000, rid_rooms=None): + per, los, his = ([], [], []) + for pid, kcsv, mcsv in pairs: + probe = _select_probe(mcsv, probe_only, rid_rooms=rid_rooms) + per.append((pid, kcsv, mcsv, probe)) + if probe: + los.append(min((s for s, _, _ in probe))) + his.append(max((s for s, _, _ in probe))) + lo = min(los) - pad_ns if los else 0 + hi = max(his) + pad_ns if his else 0 + return (per, lo, hi) + +def pergpu_engine( + per, + lo, + hi, + base_pid, + proc_name, + t0, + with_mori=True, + tp_offset=0, + no_reqstats=False, + lane_index_base=0, +): + """Emit one lane set per TP worker using global ranks from ``tp_offset``. + Track-ID spacing reserves rows for stacked stage, host-send, and KV intervals. + """ + us = lambda ns: (ns - t0) / 1000.0 + ev = [{'ph': 'M', 'name': 'process_name', 'pid': base_pid, 'tid': 0, 'args': {'name': proc_name}}] + stats = [] + for j, (pid, kcsv, mcsv, probe) in enumerate(per): + lane_index = lane_index_base + j + gtp = tp_offset + lane_index + b = 1000 + lane_index * 100 + kt, mt, st, ot, kvt = (b + 0, b + 1, b + 10, b + 30, b + 50) + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': kt, 'args': {'name': f'TP{gtp} GPU kernels (GPU{lane_index})'}}) + if not no_reqstats: + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': mt, 'args': {'name': f'TP{gtp} reqstats markers'}}) + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': st, 'args': {'name': f'TP{gtp} per-request stage spans'}}) + if with_mori: + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': ot, 'args': {'name': f'TP{gtp} MORI-IO host send (post us)'}}) + ev.append({'ph': 'M', 'name': 'thread_name', 'pid': base_pid, 'tid': kvt, 'args': {'name': f'TP{gtp} MORI-IO KV transfer (post->cq ms)'}}) + nk = 0 + if kcsv: + for s, e, name in iter_kernels(kcsv): + if s < lo or s > hi: + continue + ev.append({'ph': 'X', 'name': name, 'ts': us(s), 'dur': (e - s) / 1000.0, 'pid': base_pid, 'tid': kt}) + nk += 1 + nspan, span_depth = (0, 0) + if not no_reqstats: + for s, event, room in probe: + ev.append({'ph': 'i', 'name': event + (' room=' + room if room else ''), 'ts': us(s), 'pid': base_pid, 'tid': mt, 's': 't'}) + byroom = {} + for s, event, room in sorted(probe, key=lambda x: x[0]): + byroom.setdefault(room, []).append((s, event)) + span_intervals = build_stage_span_intervals(byroom) + nspan, span_depth = emit_stacked(ev, span_intervals, base_pid, st, us, lambda k: f'TP{gtp} per-request stage spans #{k + 1}', max_rows=20) + ev += build_decode_ct_counters(probe, f'TP{gtp}', base_pid, us) + nmori = 0 + nkvt = 0 + kv_depth = 0 + if with_mori: + mori = [m for m in iter_mori(mcsv) if lo <= m[0] <= hi] + uid_to_room = read_map(mcsv, lo, hi) + rids = mori_rid_assign(probe, mori, uid_to_room) + kv_intervals = [] + host_intervals = [] + for (s, e, name), rid in zip(mori, rids): + label = name + (f' rid={rid}' if rid else '') + _mb = _BYTES_RE.search(name) + nbytes = int(_mb.group(1)) if _mb else 0 + if 'kv_transfer' in name: + kv_intervals.append((s, e, label, rid, nbytes)) + else: + _ha = {} + if rid: + _ha['rid'] = rid + if nbytes: + _ha['bytes'] = nbytes + host_intervals.append((s, e, label, _ha)) + nmori += 1 + emit_stacked(ev, host_intervals, base_pid, ot, us, lambda k: f'TP{gtp} MORI-IO host send #{k + 1}', max_rows=20) + nkvt, kv_depth = emit_kv_stack(ev, kv_intervals, base_pid, kvt, gtp, us, max_rows=50) + stats.append((pid, nk, len(probe), nspan, nmori, nkvt, kv_depth)) + return (ev, stats) + +def write_pergpu(a, out): + pp = discover_pairs(a.prefill_kernel) + dp = discover_pairs(a.decode_kernel) + p_per, p_lo, p_hi = lane_window(pp, a.probe_only) + d_per, d_lo, d_hi = lane_window(dp, a.probe_only) + no_reqstats = getattr(a, 'no_reqstats_lanes', False) + pev, pst = pergpu_engine(p_per, p_lo, p_hi, 10, 'PREFILL (NODE0)', p_lo, with_mori=True, no_reqstats=no_reqstats) + dev, dst = pergpu_engine(d_per, d_lo, d_hi, 20, 'DECODE (NODE1)', d_lo, with_mori=False, no_reqstats=no_reqstats) + with open(out, 'w') as f: + json.dump({'traceEvents': pev + dev, 'displayTimeUnit': 'ns'}, f) + print(f'[per-GPU] PREFILL lanes={len(pst)} (TP0..TP{len(pst) - 1}) kernels={[s[1] for s in pst]} mori_hostsend={[s[4] for s in pst]} kv_transfer={[s[5] for s in pst]} kv_depth={[s[6] for s in pst]}') + print(f'[per-GPU] DECODE lanes={len(dst)} (TP0..TP{len(dst) - 1}) kernels={[s[1] for s in dst]} mori_hostsend={[s[4] for s in dst]} kv_transfer={[s[5] for s in dst]} kv_depth={[s[6] for s in dst]}') + print(f'[per-GPU] wrote {out} ({len(pev) + len(dev)} events) -- 8+8 = {len(pst) + len(dst)} GPU lanes') + + +def _new_trace_validation(path): + return { + "path": os.path.abspath(path), + "events": 0, + "processes": [], + "worker_lanes": 0, + "metadata": set(), + "slices": 0, + } + + +def _record_trace_event(validation, event): + validation["events"] += 1 + name = event.get("name", "") + if event.get("ph") == "M": + args = event.get("args", {}) + signature = ( + name, + event.get("pid"), + event.get("tid"), + json.dumps(args, sort_keys=True), + ) + if signature in validation["metadata"]: + raise SystemExit( + "ERROR: duplicate metadata signature in " + f"{validation['path']}: {signature}" + ) + validation["metadata"].add(signature) + if name == "process_name": + validation["processes"].append(args.get("name", "")) + elif name == "thread_name" and "reqstats markers" in args.get("name", ""): + validation["worker_lanes"] += 1 + elif event.get("ph") in ("X", "i", "C"): + validation["slices"] += 1 + + +def write_multinode( + prefill_dirs, + decode_dirs, + out, + probe_only, + no_reqstats=False, + rid_rooms=None, +): + """TOPOLOGY-AWARE path: one process per source node, GLOBAL TP-rank labels, and + depth-stacked stage-span / host-send / KV lanes. Handles 1P1D (1 prefill + 1 decode + dir -> 16 lanes) up through xPyD (xP prefill dirs + yD decode dirs -> (xP+yD)*8 lanes, + e.g. 2P2D -> 32). Each node is normalized to its OWN probe-window t0 (the nodes' + clocks are unaligned); processes overlay and are delineated by name.""" + base = 10 + summary = [] + validation = _new_trace_validation(out) + with open(out, 'w') as f: + f.write('{"traceEvents": [') + first_event = True + total_events = 0 + for role, dirs, with_mori in ( + ('PREFILL', prefill_dirs, True), + ('DECODE', decode_dirs, False), + ): + for local_idx, d in enumerate(dirs): + m = re.search('NODE(\\d+)', os.path.basename(os.path.normpath(d))) + nr = m.group(1) if m else str(local_idx) + tp_off = local_idx * 8 + pairs = discover_pairs_in_dir(d) + per, lo, hi = lane_window( + pairs, probe_only, rid_rooms=rid_rooms + ) + g0 = tp_off + g1 = tp_off + (len(per) - 1 if per else 0) + pname = f'{role} NODE{nr} [TP{g0}-{g1}]' + st = [] + wrote_process = False + lane_inputs = list(enumerate(per)) + if not lane_inputs: + lane_inputs = [(0, None)] + for lane_index, lane in lane_inputs: + if lane is None: + events = [{ + 'ph': 'M', + 'name': 'process_name', + 'pid': base, + 'tid': 0, + 'args': {'name': pname}, + }] + lane_stats = [] + else: + events, lane_stats = pergpu_engine( + [lane], + lo, + hi, + base, + pname, + lo, + with_mori=with_mori, + tp_offset=tp_off, + no_reqstats=no_reqstats, + lane_index_base=lane_index, + ) + if wrote_process: + events = events[1:] + else: + wrote_process = True + for event in events: + _record_trace_event(validation, event) + if not first_event: + f.write(', ') + json.dump(event, f) + first_event = False + total_events += 1 + st.extend(lane_stats) + base += 10 + print(f'[{pname}] lanes={len(st)} kernels={[s[1] for s in st]} stage_spans={[s[3] for s in st]} mori_hostsend={[s[4] for s in st]} kv_transfer={[s[5] for s in st]} kv_depth={[s[6] for s in st]}') + summary.append((pname, st)) + f.write('], "displayTimeUnit": "ns"}') + total = sum((len(st) for _, st in summary)) + global _LAST_TRACE_VALIDATION + _LAST_TRACE_VALIDATION = validation + print(f'[multinode] wrote {out} ({total_events} events) -- {total} GPU lanes across {len(summary)} processes (xP={len(prefill_dirs)} yD={len(decode_dirs)})') + +def write_aggregated(a, out): + pe, pk, pm, ps, pmo = engine_events(a.prefill_kernel, a.prefill_marker, 10, 'PREFILL (NODE0)', probe_only=a.probe_only) + de, dk, dm, ds, dmo = engine_events(a.decode_kernel, a.decode_marker, 20, 'DECODE (NODE1)', probe_only=a.probe_only) + with open(out, 'w') as f: + json.dump({'traceEvents': pe + de, 'displayTimeUnit': 'ns'}, f) + print(f'[aggregated] PREFILL: kernels={pk} reqstats={pm} stage_spans={ps} mori={pmo}') + print(f'[aggregated] DECODE : kernels={dk} reqstats={dm} stage_spans={ds} mori={dmo}') + print(f'[aggregated] wrote {out} ({len(pe) + len(de)} events)') + +def _legacy_build_trace(): + ap = argparse.ArgumentParser(description='Clean per-GPU Perfetto/Chrome-JSON trace builder. Two input modes: (1) single-file 1P1D (--prefill-kernel/-marker + --decode-kernel/-marker); (2) TOPOLOGY-AWARE multi-node (--prefill-dir ... --decode-dir ...) which auto-scales to 1P1D/2P2D/xPyD with global TP-rank labels.') + ap.add_argument('--prefill-dir', nargs='+', help='rocprof_prefill_NODE* dir(s) -- topology-aware multi-node mode') + ap.add_argument('--decode-dir', nargs='+', help='rocprof_decode_NODE* dir(s) -- topology-aware multi-node mode') + ap.add_argument('--prefill-kernel') + ap.add_argument('--prefill-marker') + ap.add_argument('--decode-kernel') + ap.add_argument('--decode-marker') + ap.add_argument('--out', required=True) + ap.add_argument('--probe-only', action='store_true', help='legacy: trim to the final request cluster after a >30s gap') + ap.add_argument('--rid-prefix', help='select requests whose ReqTimeStats RID starts with this prefix') + ap.add_argument('--request-logs', nargs='*', default=[], help='ReqTimeStats logs used with --rid-prefix') + ap.add_argument('--expect-workers', type=int, help='validate the number of request-marker worker lanes') + ap.add_argument('--no-aggregated', action='store_true', help='skip the secondary _aggregated.json merged view (mode 1 only)') + ap.add_argument('--no-perlane', action='store_true', help=argparse.SUPPRESS) + ap.add_argument('--no-reqstats-lanes', action='store_true', dest='no_reqstats_lanes', help="exclude reqstats marker and per-request stage-span lanes from trace JSON; direct correlate/reqstats analysis remains unaffected") + a = ap.parse_args() + rid_rooms = None + if a.rid_prefix: + room_to_request = load_request_ids(a.request_logs) + rid_rooms = { + room + for room, rid in room_to_request.items() + if rid.startswith(a.rid_prefix) + } + if not rid_rooms: + ap.error(f'no ReqTimeStats requests matched --rid-prefix={a.rid_prefix!r}') + if a.prefill_dir or a.decode_dir: + if not (a.prefill_dir and a.decode_dir): + ap.error('--prefill-dir and --decode-dir must be given together') + write_multinode( + sorted(a.prefill_dir), + sorted(a.decode_dir), + a.out, + a.probe_only, + no_reqstats=a.no_reqstats_lanes, + rid_rooms=rid_rooms, + ) + return + missing = [n for n in ('prefill_kernel', 'prefill_marker', 'decode_kernel', 'decode_marker') if getattr(a, n) is None] + if missing: + ap.error('single-file mode needs --%s (or use --prefill-dir/--decode-dir)' % ', --'.join((m.replace('_', '-') for m in missing))) + write_pergpu(a, a.out) + if not a.no_aggregated: + agg_out = a.out[:-5] + '_aggregated.json' if a.out.endswith('.json') else a.out + '_aggregated' + write_aggregated(a, agg_out) + +# ---- MORI/request correlation ---- +STAGES_INNER = ('prefill_kv_transfer_start', 'prefill_kv_transfer_finish') +STAGES_OUTER = ('forward_entry', 'completion') +_MAP_RE = re.compile('\\broom=(\\d+)\\s+uid=(\\d+)') +_BYTES_RE = re.compile('\\bbytes=(\\d+)') + +def _bytes(fn): + m = _BYTES_RE.search(fn) + return int(m.group(1)) if m else 0 + +def _basename(fn): + return re.sub('\\s+id=\\S+$', '', fn) + +def _idtag(fn): + m = re.search('\\bid=(\\S+)$', fn) + return m.group(1) if m else None + +def _is_real_rid(idv): + return idv is not None and idv != '0' and idv.isdigit() and (len(idv) >= 12) +_REQ_RE = re.compile('ReqTimeStats\\(rid=(?P[^,]+), bootstrap_room=(?P\\d+),') + +def load_request_ids(logpaths): + """Return bootstrap_room -> stable SGLang request ID from ReqTimeStats logs.""" + out = {} + for path in logpaths or []: + if not path or not os.path.exists(path): + continue + with open(path, errors='ignore') as fh: + for line in fh: + m = _REQ_RE.search(line) + if m and _is_real_rid(m.group('room')): + out[m.group('room')] = m.group('rid') + return out + +def _latest_probe_bounds(path, gap_ns=30000000000, pad_ns=2000000): + """Return the final RID-tagged request cluster bounds without retaining the sweep.""" + lo = hi = prev = None + with open(path, newline='') as fh: + rows = csv.reader(fh) + next(rows, None) + for row in rows: + if len(row) < 7 or not row[1].startswith('reqstats.sched.prefill.'): + continue + if not _is_real_rid(_idtag(row[1])): + continue + try: + ts = int(row[5]) + except ValueError: + continue + if prev is not None and ts - prev > gap_ns: + lo = ts + elif lo is None: + lo = ts + hi = prev = ts + return (lo - pad_ns, hi + pad_ns) if lo is not None else None + +def _rid_bounds(path, rid_rooms, pad_ns=2000000): + """Return marker-clock bounds for explicitly selected bootstrap rooms.""" + stamps = [] + with open(path, newline='') as fh: + rows = csv.reader(fh) + next(rows, None) + for row in rows: + if len(row) < 7 or not row[1].startswith('reqstats.sched.'): + continue + if _idtag(row[1]) not in rid_rooms: + continue + try: + stamps.append(int(row[5])) + except ValueError: + continue + if not stamps: + return None + return (min(stamps) - pad_ns, max(stamps) + pad_ns) + + +def parse_pid_file(path, probe_only=False, rid_rooms=None): + """Return (rid_stamps, mori_ranges, uid_to_room) for one TP-worker marker CSV. + rid_stamps: rid -> {stage_basename_suffix: ts} + mori_ranges: list of (kind, start, end, idv) with kind in {io, rdma, kvt}; + idv is the transfer_uid string from `id=N` (None for io, which has none) + uid_to_room: transfer_uid(str) -> bootstrap_room(str) harvested from `mori.map` marks + """ + rid_stamps = defaultdict(dict) + mori = [] + uid_to_room = {} + bounds = ( + _rid_bounds(path, rid_rooms) + if rid_rooms is not None + else (_latest_probe_bounds(path) if probe_only else None) + ) + if rid_rooms is not None and bounds is None: + return (rid_stamps, mori, uid_to_room) + with open(path, newline='') as fh: + r = csv.reader(fh) + next(r, None) + for row in r: + if len(row) < 7: + continue + fn = row[1] + try: + start = int(row[5]) + end = int(row[6]) + except ValueError: + continue + if bounds is not None and (not bounds[0] <= start <= bounds[1]): + continue + if fn.startswith('reqstats.sched.prefill.'): + idv = _idtag(fn) + if _is_real_rid(idv) and ( + rid_rooms is None or idv in rid_rooms + ): + stage = _basename(fn).split('reqstats.sched.prefill.')[-1] + if stage not in rid_stamps[idv]: + rid_stamps[idv][stage] = start + elif fn.startswith('mori.map'): + m = _MAP_RE.search(fn) + if m: + uid_to_room[m.group(2)] = m.group(1) + elif fn.startswith('mori.io.engine_batch_write'): + mori.append(('io', start, end, None, 0)) + elif fn.startswith('mori.rdma.batch_post.write'): + mori.append(('rdma', start, end, _idtag(fn), _bytes(fn))) + elif fn.startswith('mori.rdma.kv_transfer'): + mori.append(('kvt', start, end, _idtag(fn), _bytes(fn))) + return (rid_stamps, mori, uid_to_room) + +def windows_for(rid_stamps): + """rid -> dict(inner=(s,e)|None, outer=(s,e)|None).""" + out = {} + for rid, st in rid_stamps.items(): + inner = None + outer = None + if STAGES_INNER[0] in st and STAGES_INNER[1] in st: + a, b = (st[STAGES_INNER[0]], st[STAGES_INNER[1]]) + if b >= a: + inner = (a, b) + if STAGES_OUTER[0] in st and STAGES_OUTER[1] in st: + a, b = (st[STAGES_OUTER[0]], st[STAGES_OUTER[1]]) + if b >= a: + outer = (a, b) + out[rid] = {'inner': inner, 'outer': outer} + return out + +def assign(mid, wins, key): + """Return list of (rid, width) whose `key` window contains mid.""" + hits = [] + for rid, w in wins.items(): + iv = w[key] + if iv and iv[0] <= mid <= iv[1]: + hits.append((rid, iv[1] - iv[0])) + return hits + +def aggregate_by_room( + prefill_dir, + tp_glob='*marker_api_trace.csv', + probe_only=False, + rid_rooms=None, +): + """Reusable per-room (== bootstrap_room) MORI aggregation, using the SAME + parse/attribution logic as main() (exact-by-id via mori.map first, then per-pid + timestamp containment fallback). Returns a dict: + + { room(str): { + "mori_io_sends", "mori_rdma_posts", "mori_kv_transfers", # counts + "io_dur_us", "rdma_dur_us", "kv_transfer_dur_us", # summed durations + "rdma_bytes", "kv_transfer_bytes", "kv_eff_bw_GBps", # bytes + eff BW + "via_exact", "via_inner", "via_outer", # attribution tally + } } + + Only rooms that had >=1 MORI range assigned are present. This is consumed by + the reqstats subcommand so the per-request CSV's MORI columns (incl. the + post->CQ "mori io time") are IDENTICAL to this tool's `--out-csv`. main() is + left untouched (this is purely additive).""" + agg = defaultdict(lambda: {'io': 0, 'rdma': 0, 'kvt': 0, 'io_dur': 0, 'rdma_dur': 0, 'kvt_dur': 0, 'rdma_bytes': 0, 'kvt_bytes': 0, 'via_exact': 0, 'via_inner': 0, 'via_outer': 0}) + for f in sorted(glob.glob(os.path.join(prefill_dir, tp_glob))): + rid_stamps, mori, uid_to_room = parse_pid_file( + f, probe_only=probe_only, rid_rooms=rid_rooms + ) + wins = windows_for(rid_stamps) + for kind, s, e, idv, nbytes in mori: + mid = (s + e) // 2 + rid = None + via = None + room = uid_to_room.get(idv) if idv is not None else None + if rid_rooms is not None and _is_real_rid(room) and room not in rid_rooms: + continue + if room is not None and _is_real_rid(room): + rid = room + via = 'exact' + else: + hits = assign(mid, wins, 'inner') + via = 'inner' + if not hits: + hits = assign(mid, wins, 'outer') + via = 'outer' + if not hits: + continue + hits.sort(key=lambda x: x[1]) + best_w = hits[0][1] + tied = [h for h in hits if h[1] == best_w] + if len(tied) > 1: + continue + rid = tied[0][0] + a = agg[rid] + a[kind] += 1 + a[kind + '_dur'] += e - s + if kind in ('rdma', 'kvt'): + a[kind + '_bytes'] += nbytes + a['via_' + via] += 1 + out = {} + for room, a in agg.items(): + bw = a['kvt_bytes'] / a['kvt_dur'] if a['kvt_dur'] else 0.0 + out[room] = {'mori_io_sends': a['io'], 'mori_rdma_posts': a['rdma'], 'mori_kv_transfers': a['kvt'], 'io_dur_us': round(a['io_dur'] / 1000.0, 1), 'rdma_dur_us': round(a['rdma_dur'] / 1000.0, 1), 'kv_transfer_dur_us': round(a['kvt_dur'] / 1000.0, 1), 'rdma_bytes': a['rdma_bytes'], 'kv_transfer_bytes': a['kvt_bytes'], 'kv_eff_bw_GBps': round(bw, 3), 'via_exact': a['via_exact'], 'via_inner': a['via_inner'], 'via_outer': a['via_outer']} + return out + +def _legacy_correlate(): + ap = argparse.ArgumentParser() + ap.add_argument('--prefill-dir', required=True, nargs='+', help='one or more rocprof_prefill_NODE* directories') + ap.add_argument('--prefill-logs', nargs='*', default=[], help='ReqTimeStats logs used for request_id <-> bootstrap_room') + ap.add_argument('--tp-glob', default='*marker_api_trace.csv') + ap.add_argument('--probe-only', action='store_true', help='legacy: correlate only the final cluster after a >30s gap') + ap.add_argument('--rid-prefix', help='correlate only ReqTimeStats RIDs with this prefix') + ap.add_argument('--out-csv', required=True) + ap.add_argument('--out-summary', required=True) + ap.add_argument('--require-complete', action='store_true', help='fail unless every KV-transfer has exact mori.map and request-ID attribution') + args = ap.parse_args() + files = [] + for d in args.prefill_dir: + files.extend(glob.glob(os.path.join(d, args.tp_glob))) + files = sorted(files) + if not files: + print(f'NO marker CSVs under {args.prefill_dir} / {args.tp_glob}', file=sys.stderr) + sys.exit(2) + room_to_request = load_request_ids(args.prefill_logs) + rid_rooms = None + if args.rid_prefix: + rid_rooms = { + room + for room, rid in room_to_request.items() + if rid.startswith(args.rid_prefix) + } + if not rid_rooms: + ap.error(f'no ReqTimeStats requests matched --rid-prefix={args.rid_prefix!r}') + agg = defaultdict(lambda: {'io': 0, 'rdma': 0, 'kvt': 0, 'io_dur': 0, 'rdma_dur': 0, 'kvt_dur': 0, 'rdma_bytes': 0, 'kvt_bytes': 0, 'pids': set(), 'via_exact': 0, 'via_inner': 0, 'via_outer': 0, 'inner_w': [], 'outer_w': []}) + n_mori = 0 + n_assigned = 0 + n_unassigned = 0 + n_ambiguous = 0 + n_io = 0 + n_rdma = 0 + n_kvt = 0 + n_exact = 0 + n_time = 0 + n_kvt_assigned = 0 + n_kvt_exact = 0 + n_kvt_unassigned = 0 + n_kvt_ambiguous = 0 + geom = {'mori_inside_kv': 0, 'kv_inside_mori': 0, 'overlap_partial': 0, 'kv_missing': 0} + pid_count = 0 + rid_global_pids = defaultdict(set) + for f in files: + pid = os.path.basename(f).split('_')[1] if '_' in os.path.basename(f) else os.path.basename(f) + rid_stamps, mori, uid_to_room = parse_pid_file( + f, probe_only=args.probe_only, rid_rooms=rid_rooms + ) + wins = windows_for(rid_stamps) + if mori: + pid_count += 1 + for rid in wins: + rid_global_pids[rid].add(pid) + for kind, s, e, idv, nbytes in mori: + room = uid_to_room.get(idv) if idv is not None else None + if rid_rooms is not None and _is_real_rid(room) and room not in rid_rooms: + continue + n_mori += 1 + if kind == 'io': + n_io += 1 + elif kind == 'rdma': + n_rdma += 1 + else: + n_kvt += 1 + mid = (s + e) // 2 + rid = None + via = None + if room is not None and _is_real_rid(room): + rid = room + via = 'exact' + else: + hits = assign(mid, wins, 'inner') + via = 'inner' + if not hits: + hits = assign(mid, wins, 'outer') + via = 'outer' + if not hits: + n_unassigned += 1 + if kind == 'kvt': + n_kvt_unassigned += 1 + continue + hits.sort(key=lambda x: x[1]) + best_w = hits[0][1] + tied = [h for h in hits if h[1] == best_w] + if len(tied) > 1: + n_ambiguous += 1 + if kind == 'kvt': + n_kvt_ambiguous += 1 + continue + rid = tied[0][0] + n_assigned += 1 + if kind == 'kvt': + n_kvt_assigned += 1 + if via == 'exact': + n_kvt_exact += 1 + if via == 'exact': + n_exact += 1 + else: + n_time += 1 + a = agg[rid] + a[kind] += 1 + a[kind + '_dur'] += e - s + if kind in ('rdma', 'kvt'): + a[kind + '_bytes'] += nbytes + a['pids'].add(pid) + a['via_' + via] += 1 + kv = wins.get(rid, {}).get('inner') + if kv is None: + geom['kv_missing'] += 1 + else: + if kv[0] <= s and e <= kv[1]: + geom['mori_inside_kv'] += 1 + elif s <= kv[0] and kv[1] <= e: + geom['kv_inside_mori'] += 1 + else: + geom['overlap_partial'] += 1 + a['inner_w'].append(kv[1] - kv[0]) + all_rids = set(rid_rooms) if rid_rooms is not None else set(rid_global_pids.keys()) | set(agg.keys()) + rids_with_mori = set(agg.keys()) + with open(args.out_csv, 'w', newline='') as fh: + w = csv.writer(fh) + w.writerow(['request_id', 'bootstrap_room', 'n_pids_seen', 'n_pids_with_mori', 'mori_io_sends', 'mori_rdma_posts', 'mori_kv_transfers', 'io_dur_us', 'rdma_dur_us', 'kv_transfer_dur_us', 'assigned_via_exact', 'assigned_via_inner', 'assigned_via_outer', 'kv_window_us_avg', 'rdma_bytes', 'kv_transfer_bytes', 'kv_eff_bw_GBps']) + for rid in sorted(all_rids): + a = agg.get(rid) + if a: + kvw = sum(a['inner_w']) / len(a['inner_w']) / 1000.0 if a['inner_w'] else 0.0 + bw = a['kvt_bytes'] / a['kvt_dur'] if a['kvt_dur'] else 0.0 + w.writerow([room_to_request.get(rid, ''), rid, len(rid_global_pids[rid]), len(a['pids']), a['io'], a['rdma'], a['kvt'], round(a['io_dur'] / 1000.0, 1), round(a['rdma_dur'] / 1000.0, 1), round(a['kvt_dur'] / 1000.0, 1), a['via_exact'], a['via_inner'], a['via_outer'], round(kvw, 1), a['rdma_bytes'], a['kvt_bytes'], round(bw, 3)]) + else: + w.writerow([room_to_request.get(rid, ''), rid, len(rid_global_pids[rid]), 0, 0, 0, 0, 0.0, 0.0, 0.0, 0, 0, 0, 0.0, 0, 0, 0.0]) + lines = [] + lines.append('# MORI-IO -> request correlation summary') + lines.append('') + lines.append(f'- TP-worker marker CSVs parsed: {len(files)} (pids with mori marks: {pid_count})') + lines.append(f'- real rids (non-id=0) seen on reqstats.sched.prefill: {len(all_rids)}') + lines.append(f'- real rids that got >=1 MORI send assigned: {len(rids_with_mori)}') + lines.append(f'- MORI ranges total: {n_mori} (io={n_io}, rdma={n_rdma}, kv_transfer={n_kvt})') + _tot_rdma_b = sum((a['rdma_bytes'] for a in agg.values())) + _tot_kvt_b = sum((a['kvt_bytes'] for a in agg.values())) + _tot_kvt_dur = sum((a['kvt_dur'] for a in agg.values())) + _agg_bw = _tot_kvt_b / _tot_kvt_dur if _tot_kvt_dur else 0.0 + lines.append(f'- bytes (optional bytes= token): rdma_batch_post={_tot_rdma_b} B, kv_transfer={_tot_kvt_b} B; aggregate kv effective BW = {_agg_bw:.3f} GB/s (0 if pre-bytes trace)') + lines.append(f'- assigned: {n_assigned} unassigned: {n_unassigned} ambiguous(>1 tightest): {n_ambiguous}') + lines.append(f'- attribution method: exact-by-id (mori.map) = {n_exact} by-time (containment fallback) = {n_time}') + lines.append(f'- KV-transfer mapping: total={n_kvt} assigned={n_kvt_assigned} exact={n_kvt_exact} unassigned={n_kvt_unassigned} ambiguous={n_kvt_ambiguous}') + lines.append(f'- ReqTimeStats request-ID bridges loaded: {len(room_to_request)}') + lines.append(f'- assigned_check: assigned+unassigned+ambiguous = {n_assigned + n_unassigned + n_ambiguous} (== total {n_mori}: {n_assigned + n_unassigned + n_ambiguous == n_mori})') + lines.append(f'- method_check: exact+by-time = {n_exact + n_time} (== assigned {n_assigned}: {n_exact + n_time == n_assigned})') + lines.append('') + lines.append("## geometry: mori range vs the rid's inner KV-transfer window") + lines.append(f"- mori range fully INSIDE kv window : {geom['mori_inside_kv']}") + lines.append(f"- kv window fully INSIDE mori range : {geom['kv_inside_mori']}") + lines.append(f"- partial overlap only : {geom['overlap_partial']}") + lines.append(f"- kv window missing for that rid : {geom['kv_missing']}") + lines.append('') + lines.append('## per-rid (aggregated over pids)') + lines.append('| rid | pids_with_mori | io_sends | rdma_posts | kv_transfers | io_dur_us | rdma_dur_us | kv_dur_us | via_exact | via_inner | via_outer | kv_bytes | kv_GBps |') + lines.append('|---|---|---|---|---|---|---|---|---|---|---|---|---|') + for rid in sorted(all_rids): + a = agg.get(rid) + if a: + _bw = a['kvt_bytes'] / a['kvt_dur'] if a['kvt_dur'] else 0.0 + lines.append(f"| {rid} | {len(a['pids'])} | {a['io']} | {a['rdma']} | {a['kvt']} | {a['io_dur'] / 1000.0:.1f} | {a['rdma_dur'] / 1000.0:.1f} | {a['kvt_dur'] / 1000.0:.1f} | {a['via_exact']} | {a['via_inner']} | {a['via_outer']} | {a['kvt_bytes']} | {_bw:.3f} |") + else: + lines.append(f'| {rid} | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 0.0 | 0 | 0 | 0 | 0 | 0.000 |') + summary = '\n'.join(lines) + '\n' + with open(args.out_summary, 'w') as fh: + fh.write(summary) + print(summary) + if args.require_complete: + mapped_rooms = {room for room, a in agg.items() if a['kvt']} + missing_request_ids = sorted((room for room in mapped_rooms if not room_to_request.get(room))) + errors = [] + if rid_rooms is not None: + output_rows = [(room_to_request.get(room, ''), room) for room in sorted(all_rids)] + expected_rows = {(room_to_request[room], room) for room in rid_rooms} + if len(output_rows) != len(expected_rows) or set(output_rows) != expected_rows: + errors.append('filtered output rows do not exactly match selected requests') + if any(not request_id or not room for request_id, room in output_rows): + errors.append('filtered output contains blank request or room IDs') + if len({request_id for request_id, _ in output_rows}) != len(output_rows) or len({room for _, room in output_rows}) != len(output_rows): + errors.append('filtered output contains duplicate request or room IDs') + if n_kvt == 0: + errors.append('no mori.rdma.kv_transfer markers found') + if n_kvt_assigned != n_kvt or n_kvt_unassigned or n_kvt_ambiguous: + errors.append(f'KV transfers not uniquely assigned ({n_kvt_assigned}/{n_kvt})') + if n_kvt_exact != n_kvt: + errors.append(f'mori.map exact attribution incomplete ({n_kvt_exact}/{n_kvt})') + if missing_request_ids: + errors.append(f'missing ReqTimeStats request IDs for {len(missing_request_ids)} bootstrap rooms') + if errors: + for err in errors: + print(f'[correlate_mori] ERROR: {err}', file=sys.stderr) + sys.exit(3) + +# ---- per-request statistics ---- +RUN_OUT_BASE = os.environ.get('RUN_OUT_BASE', '/shared_inference/%s/model_blog_logs' % (os.environ.get('USER') or 'aarai')) +NS = 1000000.0 + +def _splitlist(vals): + """Accept nargs='+' AND comma-separated; flatten + drop empties.""" + out = [] + if not vals: + return out + for v in vals: + out.extend((x for x in str(v).split(',') if x)) + return out + +def _is_real_room(r): + return r is not None and r != '0' and r.isdigit() and (len(r) >= 12) + +def collect_stage_ts(dirs, side): + """side in {prefill,decode}. dirs is a LIST of rocprof_*_NODE* dirs (one per node + on that side). Returns room -> {stage: ts_ns} (earliest across ALL TP workers on + ALL nodes of that side), and room -> last decode_finish ts.""" + pref = f'reqstats.sched.{side}.' + stamps = defaultdict(dict) + dfin_last = defaultdict(int) + files = [] + for d in dirs: + files.extend(sorted(glob.glob(os.path.join(d, '*marker_api_trace.csv')))) + for f in files: + with open(f, newline='') as fh: + r = csv.reader(fh) + next(r, None) + for row in r: + if len(row) < 7: + continue + fn = row[1] + if not fn.startswith(pref): + continue + m = re.search('\\bid=(\\S+)$', fn) + room = m.group(1) if m else None + if not _is_real_room(room): + continue + try: + ts = int(row[5]) + except ValueError: + continue + stage = fn.split(pref)[-1].split(' id=')[0] + if stage == 'decode_finish': + if ts > dfin_last[room]: + dfin_last[room] = ts + continue + if stage not in stamps[room] or ts < stamps[room][stage]: + stamps[room][stage] = ts + return (stamps, dfin_last) + +def _dur_ms(st, a, b): + if a in st and b in st and (st[b] >= st[a]): + return round((st[b] - st[a]) / NS, 3) + return '' + +def prefill_derived(st): + return {'pm_bootstrap_ms': _dur_ms(st, 'prefill_bootstrap_queue_entry', 'bootstrap_done'), 'pm_queue_ms': _dur_ms(st, 'wait_queue_entry', 'forward_entry'), 'pm_forward_ms': _dur_ms(st, 'forward_entry', 'prefill_finished'), 'pm_kv_transfer_ms': _dur_ms(st, 'prefill_kv_transfer_start', 'prefill_kv_transfer_finish'), 'pm_total_ms': _dur_ms(st, 'recv', 'completion'), 'pm_recv_ns': st.get('recv', ''), 'pm_forward_entry_ns': st.get('forward_entry', ''), 'pm_prefill_finished_ns': st.get('prefill_finished', ''), 'pm_kv_start_ns': st.get('prefill_kv_transfer_start', ''), 'pm_kv_finish_ns': st.get('prefill_kv_transfer_finish', ''), 'pm_completion_ns': st.get('completion', '')} + +def decode_derived(st, dfin_last): + fwd_end = dfin_last if dfin_last else st.get('completion', None) + dm_forward = '' + if 'forward_entry' in st and fwd_end and (fwd_end >= st['forward_entry']): + dm_forward = round((fwd_end - st['forward_entry']) / NS, 3) + return {'dm_queue_ms': _dur_ms(st, 'wait_queue_entry', 'forward_entry'), 'dm_forward_ms': dm_forward, 'dm_total_ms': _dur_ms(st, 'recv', 'completion'), 'dm_recv_ns': st.get('recv', ''), 'dm_forward_entry_ns': st.get('forward_entry', ''), 'dm_completion_ns': st.get('completion', ''), 'dm_last_decode_finish_ns': dfin_last or ''} +RTS = re.compile('ReqTimeStats\\(rid=(?P[^,]+), bootstrap_room=(?P\\d+), input_len=(?P\\d+), cached_input_len=(?P\\d+), output_len=(?P
    \\d+), (?:attempts=\\d+, )?type=(?P\\w+)\\):(?P.*)') + +def parse_reqtimestats(logpaths): + """logpaths is a LIST of engine logs (one per node on that side). Returns + room -> dict of log fields (last occurrence wins across all logs).""" + out = {} + for logpath in logpaths: + if not logpath or not os.path.exists(logpath): + continue + with open(logpath, errors='ignore') as fh: + for line in fh: + m = RTS.search(line) + if not m: + continue + room = m.group('room') + if not _is_real_room(room): + continue + d = {'rid': m.group('rid'), 'bootstrap_room': room, 'input_len': m.group('il'), 'cached_input_len': m.group('cil'), 'output_len': m.group('ol')} + for k, v in re.findall('([#\\w]+)=([\\d.]+)', m.group('rest')): + d[k] = v + out[room] = d + return out +_MORI_NUM_KEYS = ('mori_io_sends', 'mori_rdma_posts', 'mori_kv_transfers', 'io_dur_us', 'rdma_dur_us', 'kv_transfer_dur_us', 'rdma_bytes', 'kv_transfer_bytes', 'via_exact', 'via_inner', 'via_outer') + +def _merge_mori(per_dir): + """Merge correlate_mori.aggregate_by_room() outputs across prefill dirs. A given + bootstrap_room is served by ONE prefill node, but we SUM defensively; effective + BW is recomputed from the summed bytes/duration.""" + out = {} + for dct in per_dir: + for room, m in dct.items(): + o = out.setdefault(room, {k: 0 for k in _MORI_NUM_KEYS}) + for k in _MORI_NUM_KEYS: + o[k] = o.get(k, 0) + (m.get(k, 0) or 0) + for room, o in out.items(): + dur_ns = o['kv_transfer_dur_us'] * 1000.0 + o['kv_eff_bw_GBps'] = round(o['kv_transfer_bytes'] / dur_ns, 3) if dur_ns else 0.0 + return out + +def mori_per_room(prefill_dirs, rid_rooms=None): + per_dir = [ + aggregate_by_room( + d, probe_only=(rid_rooms is None), rid_rooms=rid_rooms + ) + for d in prefill_dirs + ] + return _merge_mori(per_dir) + +def _client_ms(row, start_key, end_key): + try: + return round((int(row[end_key]) - int(row[start_key])) / NS, 3) + except (KeyError, TypeError, ValueError): + return '' + +def load_client(client_csv): + out = {} + if not client_csv or not os.path.exists(client_csv): + return out + with open(client_csv, newline='') as fh: + for row in csv.DictReader(fh): + out[row.get('rid', '')] = row + return out + +def load_manifest_rids(client_manifest): + if not client_manifest or not os.path.exists(client_manifest): + return set() + with open(client_manifest) as fh: + manifest = json.load(fh) + return { + request.get('rid', '') + for request in manifest.get('requests', []) + if request.get('rid') + } + +def _legacy_reqstats(): + ap = argparse.ArgumentParser() + ap.add_argument('--job') + ap.add_argument('--xp', type=int, default=1, help='# prefill nodes (auto-derive dirs/logs from --job)') + ap.add_argument('--yd', type=int, default=1, help='# decode nodes (auto-derive dirs/logs from --job)') + ap.add_argument('--prefill-dir') + ap.add_argument('--decode-dir') + ap.add_argument('--prefill-log') + ap.add_argument('--decode-log') + ap.add_argument('--prefill-dirs', nargs='+') + ap.add_argument('--decode-dirs', nargs='+') + ap.add_argument('--prefill-logs', nargs='+') + ap.add_argument('--decode-logs', nargs='+') + ap.add_argument('--client-csv') + ap.add_argument('--client-manifest') + ap.add_argument('--rid-prefix', help='keep only requests whose ReqTimeStats RID starts with this prefix') + ap.add_argument('--out-dir') + ap.add_argument('--splits', action='store_true', help='also write _prefill/_decode CSVs') + ap.add_argument('--require-data', action='store_true', help='fail if marker/log request data cannot be reconstructed') + ap.add_argument('--require-client', action='store_true', help='fail unless every tagged client RID joins to ReqTimeStats') + ap.add_argument('--no-mori', action='store_true', help='skip MoRI attribution and leave MoRI columns zero') + a = ap.parse_args() + J = a.job or 'run' + base = os.path.join(RUN_OUT_BASE, J) if a.job else None + if _splitlist(a.prefill_dirs): + pdirs = _splitlist(a.prefill_dirs) + elif a.prefill_dir: + pdirs = [a.prefill_dir] + elif base: + pdirs = [os.path.join(base, f'rocprof_prefill_NODE{i}') for i in range(a.xp)] + else: + pdirs = [] + if _splitlist(a.decode_dirs): + ddirs = _splitlist(a.decode_dirs) + elif a.decode_dir: + ddirs = [a.decode_dir] + elif base: + ddirs = [os.path.join(base, f'rocprof_decode_NODE{a.xp + j}') for j in range(a.yd)] + else: + ddirs = [] + if _splitlist(a.prefill_logs): + plogs = _splitlist(a.prefill_logs) + elif a.prefill_log: + plogs = [a.prefill_log] + elif base: + plogs = [os.path.join(base, f'prefill_NODE{i}.log') for i in range(a.xp)] + else: + plogs = [] + if _splitlist(a.decode_logs): + dlogs = _splitlist(a.decode_logs) + elif a.decode_log: + dlogs = [a.decode_log] + elif base: + dlogs = [os.path.join(base, f'decode_NODE{a.xp + j}.log') for j in range(a.yd)] + else: + dlogs = [] + out_dir = a.out_dir or os.path.join(os.path.dirname(os.path.abspath(__file__)), 'artifacts') + os.makedirs(out_dir, exist_ok=True) + print(f'[reqstats_per_request] xp={a.xp} yd={a.yd}') + print(f'[reqstats_per_request] prefill dirs: {pdirs}') + print(f'[reqstats_per_request] decode dirs: {ddirs}') + print(f'[reqstats_per_request] prefill logs: {plogs}') + print(f'[reqstats_per_request] decode logs: {dlogs}') + p_stamps, _ = collect_stage_ts(pdirs, 'prefill') if pdirs else ({}, {}) + d_stamps, d_dfin = collect_stage_ts(ddirs, 'decode') if ddirs else ({}, {}) + p_log = parse_reqtimestats(plogs) + d_log = parse_reqtimestats(dlogs) + rid_rooms = None + if a.rid_prefix: + rid_rooms = { + room + for room, row in {**p_log, **d_log}.items() + if row.get('rid', '').startswith(a.rid_prefix) + } + if not rid_rooms: + ap.error(f'no ReqTimeStats requests matched --rid-prefix={a.rid_prefix!r}') + p_stamps = {room: value for room, value in p_stamps.items() if room in rid_rooms} + d_stamps = {room: value for room, value in d_stamps.items() if room in rid_rooms} + d_dfin = {room: value for room, value in d_dfin.items() if room in rid_rooms} + p_log = {room: value for room, value in p_log.items() if room in rid_rooms} + d_log = {room: value for room, value in d_log.items() if room in rid_rooms} + mori = ( + mori_per_room(pdirs, rid_rooms=rid_rooms) + if pdirs and not a.no_mori + else {} + ) + client = load_client(a.client_csv or (os.path.join(base, 'rocprof_probe_client.csv') if base else None)) + manifest_rids = load_manifest_rids(a.client_manifest) + if a.client_manifest and not manifest_rids: + ap.error(f'client manifest is missing or has no requests: {a.client_manifest}') + if a.rid_prefix and any( + not rid.startswith(a.rid_prefix) for rid in manifest_rids + ): + ap.error('client manifest contains RIDs outside --rid-prefix') + rooms = set(p_stamps) | set(d_stamps) | set(p_log) | set(d_log) + cols = ['rid', 'bootstrap_room', 'sides', 'input_len', 'cached_input_len', 'output_len', 'p_bootstrap_ms', 'p_queue_ms', 'p_forward_ms', 'p_entry_time', 'p_transfer_speed_GBps', 'p_transfer_total_MB', 'p_retries', 'pm_bootstrap_ms', 'pm_queue_ms', 'pm_forward_ms', 'pm_kv_transfer_ms', 'pm_total_ms', 'pm_recv_ns', 'pm_forward_entry_ns', 'pm_prefill_finished_ns', 'pm_kv_start_ns', 'pm_kv_finish_ns', 'pm_completion_ns', 'd_bootstrap_ms', 'd_alloc_wait_ms', 'd_transfer_ms', 'd_queue_ms', 'd_forward_ms', 'd_entry_time', 'dm_queue_ms', 'dm_forward_ms', 'dm_total_ms', 'dm_recv_ns', 'dm_forward_entry_ns', 'dm_completion_ns', 'dm_last_decode_finish_ns', 'mori_io_sends', 'mori_rdma_posts', 'mori_io_dur_us', 'mori_kv_transfers', 'mori_io_time_ms', 'mori_kv_bytes', 'mori_kv_eff_bw_GBps', 'client_send_wall_ns', 'client_first_token_wall_ns', 'client_done_wall_ns', 'client_ttft_ms', 'client_e2e_ms', 'cli_e2e_latency_s', 'cli_queue_time_s', 'cli_completion_tokens', 'cli_decode_throughput', 'cli_first_token_ts', 'cli_request_finished_ts'] + rows = [] + for room in rooms: + pl = p_log.get(room, {}) + dl = d_log.get(room, {}) + rid = pl.get('rid') or dl.get('rid') or '' + sides = ('P' if room in p_stamps or pl else '') + ('D' if room in d_stamps or dl else '') + il = pl.get('input_len') or dl.get('input_len') or '' + cil = pl.get('cached_input_len') or dl.get('cached_input_len') or '' + ol = dl.get('output_len') or pl.get('output_len') or '' + pm = prefill_derived(p_stamps.get(room, {})) + dm = decode_derived(d_stamps.get(room, {}), d_dfin.get(room, 0)) + mo = mori.get(room, {}) + cli = client.get(rid, {}) + row = {'rid': rid, 'bootstrap_room': room, 'sides': sides, 'input_len': il, 'cached_input_len': cil, 'output_len': ol, 'p_bootstrap_ms': pl.get('bootstrap_duration', ''), 'p_queue_ms': pl.get('queue_duration', ''), 'p_forward_ms': pl.get('forward_duration', ''), 'p_entry_time': pl.get('entry_time', ''), 'p_transfer_speed_GBps': pl.get('transfer_speed', ''), 'p_transfer_total_MB': pl.get('transfer_total', ''), 'p_retries': pl.get('#retries', ''), 'd_bootstrap_ms': dl.get('bootstrap_duration', ''), 'd_alloc_wait_ms': dl.get('alloc_wait_duration', ''), 'd_transfer_ms': dl.get('transfer_duration', ''), 'd_queue_ms': dl.get('queue_duration', ''), 'd_forward_ms': dl.get('forward_duration', ''), 'd_entry_time': dl.get('entry_time', ''), 'mori_io_sends': mo.get('mori_io_sends', 0), 'mori_rdma_posts': mo.get('mori_rdma_posts', 0), 'mori_io_dur_us': mo.get('io_dur_us', 0.0), 'mori_kv_transfers': mo.get('mori_kv_transfers', 0), 'mori_io_time_ms': round(mo.get('kv_transfer_dur_us', 0.0) / 1000.0, 4), 'mori_kv_bytes': mo.get('kv_transfer_bytes', 0), 'mori_kv_eff_bw_GBps': mo.get('kv_eff_bw_GBps', 0.0), 'client_send_wall_ns': cli.get('client_send_wall_ns', ''), 'client_first_token_wall_ns': cli.get('client_first_token_wall_ns', ''), 'client_done_wall_ns': cli.get('client_done_wall_ns', ''), 'client_ttft_ms': _client_ms(cli, 'client_send_wall_ns', 'client_first_token_wall_ns'), 'client_e2e_ms': _client_ms(cli, 'client_send_wall_ns', 'client_done_wall_ns'), 'cli_e2e_latency_s': cli.get('mi_e2e_latency', ''), 'cli_queue_time_s': cli.get('mi_queue_time', ''), 'cli_completion_tokens': cli.get('mi_completion_tokens', ''), 'cli_decode_throughput': cli.get('mi_decode_throughput', ''), 'cli_first_token_ts': cli.get('mi_first_token_ts', ''), 'cli_request_finished_ts': cli.get('mi_request_finished_ts', '')} + row.update(pm) + row.update(dm) + rows.append(row) + rows.sort(key=lambda r: (float(r['p_entry_time']) if r['p_entry_time'] else float('inf'), r['bootstrap_room'])) + merged = os.path.join(out_dir, f'reqstats_per_request_{J}.csv') + with open(merged, 'w', newline='') as fh: + w = csv.DictWriter(fh, fieldnames=cols) + w.writeheader() + for r in rows: + w.writerow(r) + print(f'[reqstats_per_request] wrote {merged} ({len(rows)} requests, {len(cols)} columns)') + if a.splits: + for side, keep in (('prefill', [c for c in cols if not c.startswith(('d_', 'dm_'))]), ('decode', [c for c in cols if not c.startswith(('p_', 'pm_', 'mori_'))])): + sp = os.path.join(out_dir, f'reqstats_per_request_{J}_{side}.csv') + with open(sp, 'w', newline='') as fh: + w = csv.DictWriter(fh, fieldnames=keep, extrasaction='ignore') + w.writeheader() + for r in rows: + w.writerow(r) + print(f'[reqstats_per_request] wrote {sp}') + row_rids = {r['rid'] for r in rows if r['rid']} + client_rids = {rid for rid in client if rid} + joined_rids = row_rids & client_rids + missing_client = sorted(client_rids - row_rids) + mori_rids = {r['rid'] for r in rows if r['rid'] and r['mori_kv_transfers']} + missing_mori = sorted(client_rids - mori_rids) if not a.no_mori else [] + n_join = len(joined_rids) + n_mori = sum((1 for r in rows if r['mori_io_sends'] or r['mori_kv_transfers'])) + n_iotime = sum((1 for r in rows if r['mori_io_time_ms'])) + print(f'[reqstats_per_request] rids={len(rows)} client-joined={n_join}/{len(client_rids)} with-MORI={n_mori} with-mori_io_time={n_iotime}') + errors = [] + if a.require_data and (not rooms or not p_stamps or (not p_log)): + errors.append(f'insufficient request data rooms={len(rooms)} prefill_markers={len(p_stamps)} prefill_logs={len(p_log)}') + if a.require_client: + if not client_rids: + errors.append('tagged client CSV is missing or empty') + if a.client_manifest and manifest_rids != client_rids: + errors.append( + 'client CSV/manifest RID mismatch ' + f'(csv_only={sorted(client_rids - manifest_rids)[:3]}, ' + f'manifest_only={sorted(manifest_rids - client_rids)[:3]})' + ) + if missing_client: + errors.append(f'{len(missing_client)} tagged client RIDs missing from ReqTimeStats (sample={missing_client[:3]})') + if not a.no_mori and missing_mori: + errors.append(f'{len(missing_mori)} tagged client RIDs missing exact MORI KV mappings (sample={missing_mori[:3]})') + if errors: + for err in errors: + print(f'[reqstats_per_request] ERROR: {err}', file=sys.stderr) + sys.exit(3) + +# ---- MIT-attributed kernel categorization ---- +# SPDX-License-Identifier: MIT +# Rules derived from ROCm/llmscope layer_detection.py categorize_kernel: +# https://github.com/ROCm/llmscope/blob/di_analysis_branch/llmscope/layer_detection.py +def categorize_kernel(name): + """Categorize a kernel by its function.""" + n = name.lower() + if 'rmsnorm' in n or 'fused_rms' in n or 'rms_norm' in n or ('rsqrt' in n and 'mean' in n and ('mul' in n)): + return 'RMSNorm' + if 'rope' in n: + return 'ROPE' + if 'reshape' in n and 'cache' in n: + return 'KVCacheReshape' + if 'kernel_unified_attention' in n: + return 'Attention' + if '_fwd_kernel' in name: + return 'TritonAttention' + if 'fmha' in n: + return 'FMHA' + if 'mla' in n: + return 'MLA' + if 'aiter::pa' in name: + return 'PA' + if 'paged_attention' in n: + return 'PagedAttn' + if 'routing' in n or 'route' in n: + return 'MoE_Router' + if 'aiter::fmoe' in name: + return 'MoE_Fused' + if 'kernel_moe' in n: + return 'MoE_Unfused' + if 'moesorting' in n: + return 'MoE_Sort' + if 'topk' in n: + return 'MoE_TopK' + if any(x in n for x in ('epdispatchinternode', 'epcombineinternode', 'epdispatchintranode', 'epcombineintranode')): + return 'MORI EP' + if 'epdispatch' in n or 'epcombine' in n: + return 'Communication' + if 'gemm' in n or 'cijk' in n or 'wvsplit' in n or ('matmul' in n): + return 'GEMM' + if 'act_and_mul' in n or 'silu' in n: + return 'Activation' + if 'quant' in n: + return 'Quant' + if 'allreduce' in n or 'cross_device' in n or 'nccl' in n or ('allgather' in n): + return 'Communication' + if 'poi' in n or 'elementwise' in n: + return 'Elementwise' + return 'Other' + +# ---- kernel bucket generation ---- +def _find_column(headers, predicate, what): + """Return the first header matching ``predicate`` (called on lowercased name).""" + for h in headers: + if predicate((h or '').lower()): + return h + sys.exit(f'ERROR: could not find the {what} column.\n Headers found: {headers}') + +def find_kernel_name_column(headers): + """Locate the kernel-name column: contains both 'kernel' and 'name'.""" + return _find_column(headers, lambda h: 'kernel' in h and 'name' in h, "kernel name (a column containing both 'kernel' and 'name')") + +def find_duration_sum_column(headers): + """Locate the per-row total-duration column: contains 'duration' and '_sum'. + + Tolerant of the micro sign (matches whether the header uses 'µs' or 'us'). + """ + return _find_column(headers, lambda h: 'duration' in h and '_sum' in h, "duration sum (a column containing both 'duration' and '_sum')") + +def find_duration_count_column(headers): + """Locate the per-row kernel-count column: contains 'duration' and '_count'.""" + for h in headers: + hl = (h or '').lower() + if 'duration' in hl and '_count' in hl: + return h + return None + +def _duration_or_none(value): + try: + value = float(str(value).strip()) + except (TypeError, ValueError): + return None + return value if math.isfinite(value) and value >= 0 else None + +def _to_int(value): + try: + return int(round(float(str(value).strip()))) + except (TypeError, ValueError): + return 0 + +def process(in_path, out_per_kernel, out_by_category, categorize_kernel, label): + with open(in_path, 'r', encoding='utf-8-sig', newline='') as f: + reader = csv.reader(f) + rows = list(reader) + if not rows: + sys.exit(f"ERROR: input CSV '{in_path}' is empty.") + headers = rows[0] + data_rows = rows[1:] + name_col = find_kernel_name_column(headers) + dur_sum_col = find_duration_sum_column(headers) + count_col = find_duration_count_column(headers) + name_idx = headers.index(name_col) + dur_idx = headers.index(dur_sum_col) + count_idx = headers.index(count_col) if count_col else None + out_headers = ['Category'] + headers + agg = {} + grand_count = 0 + grand_us = 0.0 + with open(out_per_kernel, 'w', encoding='utf-8-sig', newline='') as f: + writer = csv.writer(f) + writer.writerow(out_headers) + for row in data_rows: + if not row or all(((c or '').strip() == '' for c in row)): + continue + total_us = _duration_or_none(row[dur_idx] if dur_idx < len(row) else None) + if total_us is None: + continue + kernel_name = row[name_idx] if name_idx < len(row) else '' + category = categorize_kernel(kernel_name) + writer.writerow([category] + row) + n_kernels = _to_int(row[count_idx]) if count_idx is not None and count_idx < len(row) else 1 + bucket = agg.setdefault(category, [0, 0.0]) + bucket[0] += n_kernels + bucket[1] += total_us + grand_count += n_kernels + grand_us += total_us + with open(out_by_category, 'w', encoding='utf-8-sig', newline='') as f: + writer = csv.writer(f) + writer.writerow(['Category', 'Num_Kernels', 'Total_us', 'Total_ms', 'Pct_of_Kernel_Time']) + for category, (n_kernels, total_us) in sorted(agg.items(), key=lambda kv: kv[1][1], reverse=True): + pct = total_us / grand_us * 100 if grand_us > 0 else 0.0 + writer.writerow([category, n_kernels, f'{total_us:.2f}', f'{total_us / 1000.0:.4f}', f'{pct:.1f}%']) + writer.writerow(['TOTAL', grand_count, f'{grand_us:.2f}', f'{grand_us / 1000.0:.4f}', '100.0%']) + _print_summary(label, in_path, headers, name_col, dur_sum_col, count_col, agg, grand_count, grand_us, out_per_kernel, out_by_category) + +def _print_summary(label, in_path, headers, name_col, dur_sum_col, count_col, agg, grand_count, grand_us, out_per_kernel, out_by_category): + print('==== %s ====' % label) + print('input: %s' % in_path) + print('kernel-name column: %r' % name_col) + print('duration-sum column: %r' % dur_sum_col) + print('duration-count column:%s' % (' %r' % count_col if count_col else ' (not found; counted 1 row per kernel)')) + print('total kernel time: %.1f us (%.3f ms) total kernels: %d' % (grand_us, grand_us / 1000.0, grand_count)) + print() + print('%-15s %10s %15s %9s' % ('Category', '#kernels', 'total_us', '%time')) + for category, (n_kernels, total_us) in sorted(agg.items(), key=lambda kv: kv[1][1], reverse=True): + pct = total_us / grand_us * 100 if grand_us > 0 else 0.0 + print('%-15s %10d %15.1f %8.2f%%' % (category, n_kernels, total_us, pct)) + print('%-15s %10d %15.1f %8.2f%%' % ('TOTAL', grand_count, grand_us, 100.0)) + print() + print('wrote:', out_per_kernel) + print('wrote:', out_by_category) + +def _legacy_buckets(): + p = argparse.ArgumentParser(description='Add local function-based kernel categories onto a TraceLens kernel-summary CSV.') + p.add_argument('--in', dest='in_path', required=True, help='Input TraceLens kernel_summary CSV.') + p.add_argument('--out-per-kernel', required=True, help='Output per-kernel CSV (Category + all original columns).') + p.add_argument('--out-by-category', required=True, help='Output by-category rollup CSV.') + p.add_argument('--label', default=None, help='Label for the stdout summary banner (default: input path).') + args = p.parse_args() + label = args.label if args.label is not None else args.in_path + process(args.in_path, args.out_per_kernel, args.out_by_category, categorize_kernel, label) + +# ---- trimmed kernel summaries ---- +OUTPUT_COLUMNS = ['Time', 'Total Time', 'Instances', 'Avg', 'Med', 'Min', 'Max', 'StdDev', 'GridXYZ', 'BlockXYZ', 'VGPR', 'AccumVGPR', 'SGPR', 'LDS', 'Scratch', 'Name', 'Time %', 'Total Time (ns)', 'Avg (ns)', 'Med (ns)', 'Min (ns)', 'Max (ns)', 'StdDev (ns)', 'GridX', 'GridY', 'GridZ', 'BlockX', 'BlockY', 'BlockZ', 'n_trimmed', 'instances_before_trim'] + +def pretty_ns(ns): + """Format a nanosecond value the way TraceLens' kernel_summary.csv does: + ms if >= 1e6 ns, us if >= 1e3 ns, else ns; 3 decimal places.""" + if ns >= 1000000.0: + return f'{ns / 1000000.0:.3f} ms' + if ns >= 1000.0: + return f'{ns / 1000.0:.3f} µs' + return f'{ns:.3f} ns' + +def _n_to_trim(count, trim_pct): + """Return the number of slowest calls to drop, with at least one when eligible.""" + if trim_pct <= 0: + return 0 + n = math.ceil(count * trim_pct / 100.0) + return max(n, 1) + +def load_categorize_kernel(): + return categorize_kernel + +_KERNEL_RESOURCE_FIELDS = { + 'GridX': 'Grid_Size_X', + 'GridY': 'Grid_Size_Y', + 'GridZ': 'Grid_Size_Z', + 'BlockX': 'Workgroup_Size_X', + 'BlockY': 'Workgroup_Size_Y', + 'BlockZ': 'Workgroup_Size_Z', + 'VGPR': 'VGPR_Count', + 'AccumVGPR': 'Accum_VGPR_Count', + 'SGPR': 'SGPR_Count', + 'LDS': 'LDS_Block_Size', + 'Scratch': 'Scratch_Size', +} +_NORMALIZED_HELP_COLUMNS = [ + 'Kernel name', + 'kernel_duration_us_sum', + 'kernel_duration_us_count', +] + +def _scan_kernel_traces(kernel_trace_csvs): + """Pool exact-name dispatch durations from one or more raw worker CSVs.""" + if isinstance(kernel_trace_csvs, (str, os.PathLike)): + paths = [os.path.abspath(os.fspath(kernel_trace_csvs))] + else: + paths = [ + os.path.abspath(os.fspath(path)) + for path in kernel_trace_csvs + ] + if not paths: + raise SystemExit('ERROR: no kernel trace CSV inputs') + if len(paths) != len(set(paths)): + raise SystemExit(f'ERROR: duplicate kernel trace CSV inputs: {paths}') + + groups = {} + input_audits = [] + for kernel_trace_csv in paths: + input_audit = { + 'path': kernel_trace_csv, + 'raw_row_count': 0, + 'dispatch_row_count': 0, + 'invalid_event_count': 0, + 'included_event_count': 0, + } + with open(kernel_trace_csv, newline='', encoding='utf-8-sig') as f: + reader = csv.DictReader(f) + fieldnames = reader.fieldnames or [] + name_col = 'Kernel_Name' if 'Kernel_Name' in fieldnames else None + start_col = 'Start_Timestamp' if 'Start_Timestamp' in fieldnames else None + end_col = 'End_Timestamp' if 'End_Timestamp' in fieldnames else None + if not (name_col and start_col and end_col): + raise SystemExit( + f'ERROR: {kernel_trace_csv} does not look like a rocprofv3 ' + 'kernel_trace.csv (need Kernel_Name/Start_Timestamp/' + f'End_Timestamp, got columns: {fieldnames})' + ) + for row in reader: + input_audit['raw_row_count'] += 1 + if row.get('Kind') and row['Kind'] != 'KERNEL_DISPATCH': + continue + input_audit['dispatch_row_count'] += 1 + try: + dur_ns = _duration_or_none( + float(row[end_col]) - float(row[start_col]) + ) + except (TypeError, ValueError): + dur_ns = None + if dur_ns is None: + input_audit['invalid_event_count'] += 1 + continue + name = row[name_col] + if name not in groups: + groups[name] = { + 'durations': [], + 'resources': { + output: row.get(source, '') + for output, source in _KERNEL_RESOURCE_FIELDS.items() + }, + } + groups[name]['durations'].append(dur_ns) + input_audit['included_event_count'] += 1 + input_audits.append(input_audit) + + return groups, { + 'inputs': input_audits, + 'pooled_raw_row_count': sum( + row['raw_row_count'] for row in input_audits + ), + 'pooled_dispatch_row_count': sum( + row['dispatch_row_count'] for row in input_audits + ), + 'pooled_included_row_count': sum( + row['included_event_count'] for row in input_audits + ), + 'invalid_event_count': sum( + row['invalid_event_count'] for row in input_audits + ), + } + +def _build_summary_rows(groups, trim_pct): + summary_rows = [] + grand_total_ns = 0.0 + for name, g in groups.items(): + durations = g['durations'] + count_before = len(durations) + if count_before >= 20 and trim_pct > 0: + n_trim = _n_to_trim(count_before, trim_pct) + n_trim = min(n_trim, count_before - 1) + kept = sorted(durations)[:count_before - n_trim] + else: + n_trim = 0 + kept = durations + total = sum(kept) + count = len(kept) + avg = total / count + med = statistics.median(kept) + mn = min(kept) + mx = max(kept) + stddev = statistics.pstdev(kept) if count > 1 else 0.0 + grand_total_ns += total + summary_rows.append(g['resources'] | { + 'Name': name, + 'Instances': count, + 'instances_before_trim': count_before, + 'n_trimmed': n_trim, + 'Total Time (ns)': total, + 'Avg (ns)': avg, + 'Med (ns)': med, + 'Min (ns)': mn, + 'Max (ns)': mx, + 'StdDev (ns)': stddev, + }) + summary_rows.sort(key=lambda x: x['Total Time (ns)'], reverse=True) + for r in summary_rows: + r['Time %'] = 100.0 * r['Total Time (ns)'] / grand_total_ns if grand_total_ns else 0.0 + return (summary_rows, grand_total_ns) + +def build_pooled_kernel_summaries(kernel_trace_csvs, trim_pct): + """Build canonical untrimmed and post-pooling-trimmed summaries in one scan.""" + groups, audit = _scan_kernel_traces(kernel_trace_csvs) + normalized_rows, normalized_total_ns = _build_summary_rows(groups, 0) + trimmed_rows, trimmed_total_ns = _build_summary_rows(groups, trim_pct) + return ( + normalized_rows, + normalized_total_ns, + trimmed_rows, + trimmed_total_ns, + audit, + ) + +def build_trimmed_summary(kernel_trace_csv, trim_pct): + """Build a trimmed summary for the explicitly supplied raw kernel CSV.""" + groups, _audit = _scan_kernel_traces(kernel_trace_csv) + return _build_summary_rows(groups, trim_pct) + +def write_csv(rows, out_path, add_category, categorize_kernel, normalized=False): + columns = list(OUTPUT_COLUMNS[:-2] if normalized else OUTPUT_COLUMNS) + if normalized: + columns += _NORMALIZED_HELP_COLUMNS + if add_category and categorize_kernel is not None: + columns = ['Category'] + columns + with open(out_path, 'w', newline='', encoding='utf-8-sig') as f: + w = csv.writer(f) + w.writerow(columns) + for r in rows: + grid_xyz = f"{r['GridX']} {r['GridY']} {r['GridZ']}" + block_xyz = f"{r['BlockX']} {r['BlockY']} {r['BlockZ']}" + out = [f"{r['Time %']:.1f}%", pretty_ns(r['Total Time (ns)']), r['Instances'], pretty_ns(r['Avg (ns)']), pretty_ns(r['Med (ns)']), pretty_ns(r['Min (ns)']), pretty_ns(r['Max (ns)']), pretty_ns(r['StdDev (ns)']), grid_xyz, block_xyz, r['VGPR'], r['AccumVGPR'], r['SGPR'], r['LDS'], r['Scratch'], r['Name'], r['Time %'], r['Total Time (ns)'], r['Avg (ns)'], r['Med (ns)'], r['Min (ns)'], r['Max (ns)'], r['StdDev (ns)'], r['GridX'], r['GridY'], r['GridZ'], r['BlockX'], r['BlockY'], r['BlockZ'], r['n_trimmed'], r['instances_before_trim']] + if normalized: + out = out[:-2] + [ + r['Name'], + f"{r['Total Time (ns)'] / 1000.0:.4f}", + r['Instances'], + ] + if add_category and categorize_kernel is not None: + out = [categorize_kernel(r['Name'])] + out + w.writerow(out) + + +def run_pooled_kernel_summaries( + coverage, + normalized_path, + trimmed_path, + trim_pct, + log_path=None, +): + """Write canonical all-worker node summaries and provenance.""" + if trim_pct < 0 or trim_pct >= 100: + raise SystemExit( + f'ERROR: --trim-pct must be in [0, 100), got {trim_pct}' + ) + workers = coverage['workers'] + paths = [worker['kernel_csv'] for worker in workers] + ( + normalized, + normalized_total_ns, + trimmed, + trimmed_total_ns, + audit, + ) = build_pooled_kernel_summaries(paths, trim_pct) + write_csv( + normalized, + normalized_path, + False, + None, + normalized=True, + ) + write_csv( + trimmed, + trimmed_path, + True, + categorize_kernel, + ) + + provenance = [] + for worker, input_audit in zip(workers, audit['inputs']): + if os.path.abspath(worker['kernel_csv']) != input_audit['path']: + raise SystemExit( + 'ERROR: pooled kernel input order diverged from verified ' + 'worker source order' + ) + provenance.append({ + **worker, + **input_audit, + }) + audit = dict(audit) + audit.update({ + 'analysis_scope': 'pooled all workers', + 'aggregation_semantics': ( + 'raw dispatch rows pooled before exact-name grouping and trimming; ' + 'durations and call counts are summed, never averaged' + ), + 'activity_interpretation': ( + 'summed GPU kernel activity; not wall time or utilization' + ), + 'role': coverage['role'], + 'node_rank': coverage['node_rank'], + 'expected_worker_count': coverage['expected_worker_count'], + 'included_worker_count': coverage['included_worker_count'], + 'rank_source': coverage['rank_source'], + 'ranks_derivable': coverage['ranks_derivable'], + 'selected_pids': [worker['pid'] for worker in workers], + 'selected_local_ranks': [ + worker['local_rank'] for worker in workers + ], + 'source_files': [worker['kernel_csv'] for worker in workers], + 'workers': provenance, + 'normalized_distinct_kernel_count': len(normalized), + 'normalized_total_ns': normalized_total_ns, + 'trimmed_distinct_kernel_count': len(trimmed), + 'trimmed_total_ns': trimmed_total_ns, + 'trimmed_calls_dropped': sum( + row['n_trimmed'] for row in trimmed + ), + }) + + lines = [ + 'kernel_analysis_scope=pooled all workers', + ( + 'aggregation_semantics=raw dispatch rows pooled before exact-name ' + 'grouping and trimming; durations and call counts are summed, ' + 'never averaged' + ), + ( + 'activity_interpretation=summed GPU kernel activity; ' + 'not wall time or utilization' + ), + f"role={coverage['role'] if coverage['role'] is not None else 'unknown'}", + ( + 'node_rank=' + + ( + str(coverage['node_rank']) + if coverage['node_rank'] is not None else 'unknown' + ) + ), + f"expected_worker_count={coverage['expected_worker_count']}", + f"included_worker_count={coverage['included_worker_count']}", + f"ranks_derivable={str(coverage['ranks_derivable']).lower()}", + ( + 'rank_source=' + + ( + coverage['rank_source'] + if coverage['rank_source'] is not None else 'unavailable' + ) + ), + ] + for worker in provenance: + rank_fields = ' '.join( + f"{key}={worker[key] if worker[key] is not None else 'unknown'}" + for key in ('local_rank', 'dp_rank', 'tp_rank', 'ep_rank') + ) + lines.append( + f"pooled_input source_order={worker['source_order']} " + f"pid={worker['pid']} {rank_fields} " + f"source_file={worker['kernel_csv']} " + f"raw_rows={worker['raw_row_count']} " + f"dispatch_rows={worker['dispatch_row_count']} " + f"included_rows={worker['included_event_count']} " + f"invalid_rows={worker['invalid_event_count']}" + ) + lines.extend([ + ( + f"pooled_total raw_rows={audit['pooled_raw_row_count']} " + f"dispatch_rows={audit['pooled_dispatch_row_count']} " + f"included_rows={audit['pooled_included_row_count']} " + f"invalid_rows={audit['invalid_event_count']}" + ), + ( + f"normalized distinct_kernels={len(normalized)} " + f"total_ns={normalized_total_ns} output={normalized_path}" + ), + ( + f"trimmed distinct_kernels={len(trimmed)} " + f"total_ns={trimmed_total_ns} " + f"calls_dropped={audit['trimmed_calls_dropped']} " + f"output={trimmed_path}" + ), + ]) + if log_path is not None: + with open(log_path, 'w', encoding='utf-8') as log_file: + log_file.write('\n'.join(lines) + '\n') + for line in lines: + print(line) + return audit + + +def _legacy_trimmed_summary(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('--kernel-trace', required=True, help='Raw rocprofv3 *_kernel_trace.csv (per-dispatch rows)') + ap.add_argument('--out', default=None, help='Output CSV path (default: _summary_trimmed.csv next to input)') + ap.add_argument('--trim-pct', type=float, default=5.0, help="Percent of each eligible kernel's slowest calls to drop (default: 5)") + ap.add_argument('--add-category', action='store_true', help='Prepend a Category column via the vendored categorize_kernel()') + args = ap.parse_args() + if not os.path.isfile(args.kernel_trace): + raise SystemExit(f'ERROR: not found: {args.kernel_trace}') + if args.trim_pct < 0 or args.trim_pct >= 100: + raise SystemExit(f'ERROR: --trim-pct must be in [0, 100), got {args.trim_pct}') + out_path = args.out + if out_path is None: + base = os.path.basename(args.kernel_trace) + stem = base[:-len('_kernel_trace.csv')] if base.endswith('_kernel_trace.csv') else os.path.splitext(base)[0] + out_path = os.path.join(os.path.dirname(args.kernel_trace) or '.', f'{stem}_kernel_summary_trimmed.csv') + categorize_kernel = load_categorize_kernel() if args.add_category else None + if args.add_category and categorize_kernel is None: + print('WARNING: kernel categorizer unavailable; proceeding without it.', file=sys.stderr) + rows, grand_total_ns = build_trimmed_summary(args.kernel_trace, args.trim_pct) + write_csv(rows, out_path, args.add_category, categorize_kernel) + n_eligible = sum((1 for r in rows if r['instances_before_trim'] >= 20)) + n_trimmed_kernels = sum((1 for r in rows if r['n_trimmed'] > 0)) + total_calls_dropped = sum((r['n_trimmed'] for r in rows)) + print(f'wrote {out_path}') + print(f' {len(rows)} distinct kernels; {n_eligible} eligible (>=20 calls); {n_trimmed_kernels} trimmed at {args.trim_pct}%; {total_calls_dropped} total calls dropped') + print(f' trimmed grand total: {pretty_ns(grand_total_ns)}') + +# ---- consolidated CLI and analysis orchestration ---- + +def _temporary_output(final_path): + final = Path(final_path).resolve() + final.parent.mkdir(parents=True, exist_ok=True) + fd, name = tempfile.mkstemp( + prefix=f".{final.name}.", suffix=".tmp", dir=str(final.parent) + ) + os.close(fd) + os.unlink(name) + return final, Path(name) + + +def _validate_trace(path, expected_workers=None): + recorded = globals().get("_LAST_TRACE_VALIDATION") + if ( + recorded is not None + and recorded.get("path") == os.path.abspath(path) + ): + event_count = recorded["events"] + processes = recorded["processes"] + worker_lanes = recorded["worker_lanes"] + slices = recorded["slices"] + else: + with open(path, encoding="utf-8") as fh: + payload = json.load(fh) + events = payload.get("traceEvents") + if not isinstance(events, list) or not events: + raise SystemExit(f"ERROR: trace has no events: {path}") + validation = _new_trace_validation(path) + for event in events: + _record_trace_event(validation, event) + event_count = validation["events"] + processes = validation["processes"] + worker_lanes = validation["worker_lanes"] + slices = validation["slices"] + + if len(processes) != len(set(processes)): + raise SystemExit(f"ERROR: duplicate process names in {path}") + if not event_count or not processes or not slices: + raise SystemExit( + f"ERROR: trace is structurally empty: processes={len(processes)} slices={slices}" + ) + if expected_workers is not None and worker_lanes != expected_workers: + raise SystemExit( + f"ERROR: trace has {worker_lanes} request-marker lanes; " + f"expected {expected_workers}" + ) + print( + f"[trace_tools] validated events={event_count} processes={len(processes)} " + f"worker_lanes={worker_lanes} slices={slices}" + ) + return event_count + + +def _run_absorbed(command, function, argv): + """Run a preserved CLI implementation with atomic explicit outputs.""" + output_flags = { + "build-trace": ("--out",), + "correlate": ("--out-csv", "--out-summary"), + "buckets": ("--out-per-kernel", "--out-by-category"), + "trimmed-summary": ("--out",), + }.get(command, ()) + cooked = list(argv) + pending = [] + for flag in output_flags: + if flag not in cooked: + continue + index = cooked.index(flag) + 1 + final, temporary = _temporary_output(cooked[index]) + cooked[index] = str(temporary) + pending.append((final, temporary)) + + expected_workers = None + if command == "build-trace" and "--expect-workers" in cooked: + index = cooked.index("--expect-workers") + try: + expected_workers = int(cooked[index + 1]) + except (IndexError, ValueError): + raise SystemExit("ERROR: --expect-workers needs an integer") + + old_argv = sys.argv + try: + sys.argv = [f"{old_argv[0]} {command}", *cooked] + result = function() + if result not in (None, 0): + return result + if command == "build-trace": + output = next((temporary for final, temporary in pending + if "--out" in output_flags), None) + _validate_trace(output, expected_workers) + if "--prefill-dir" not in cooked and "--no-aggregated" not in cooked: + aggregate_temporary = Path( + str(output) + "_aggregated" + if not str(output).endswith(".json") + else str(output)[:-5] + "_aggregated.json" + ) + aggregate_final = next(final for final, temporary in pending) + aggregate_final = Path( + str(aggregate_final)[:-5] + "_aggregated.json" + if str(aggregate_final).endswith(".json") + else str(aggregate_final) + "_aggregated" + ) + _validate_trace(aggregate_temporary) + pending.append((aggregate_final, aggregate_temporary)) + for final, temporary in pending: + if not temporary.is_file() or temporary.stat().st_size == 0: + raise SystemExit(f"ERROR: {command} did not produce {temporary}") + for final, temporary in pending: + os.replace(temporary, final) + return 0 + finally: + sys.argv = old_argv + for _final, temporary in pending: + if temporary.exists(): + temporary.unlink() + + +def _column_indexes(header): + lowered = [(value or "").strip().lower() for value in header] + def column(*names): + return next((lowered.index(name) for name in names if name in lowered), None) + return { + "name": column("name", "kernel name", "kernel_name"), + "total_ns": column("total time (ns)"), + "instances": column("instances", "total count"), + "start": column("start_timestamp"), + "end": column("end_timestamp"), + } + + +def normalize_kernel_summary(source, output): + """Normalize TraceLens or raw rocprof CSV without materializing dispatch rows.""" + with open(source, newline="", encoding="utf-8-sig") as src: + reader = csv.reader(src) + header = next(reader, None) + if not header: + raise SystemExit(f"ERROR: normalizer input CSV is empty: {source}") + indexes = _column_indexes(header) + if indexes["name"] is None: + raise SystemExit( + f"ERROR: unrecognized kernel-summary schema; headers={header!r}" + ) + helper = [ + "Kernel name", + "kernel_duration_us_sum", + "kernel_duration_us_count", + ] + if indexes["start"] is not None and indexes["end"] is not None: + aggregate = {} + for row in reader: + try: + name = row[indexes["name"]] + duration_us = _duration_or_none(( + float(row[indexes["end"]]) - float(row[indexes["start"]]) + ) / 1000.0) + except (IndexError, ValueError): + continue + if duration_us is None: + continue + bucket = aggregate.setdefault(name, [0.0, 0]) + bucket[0] += duration_us + bucket[1] += 1 + with open(output, "w", newline="", encoding="utf-8-sig") as dst: + writer = csv.writer(dst) + writer.writerow(["Name", "Total Time (ns)", "Instances", *helper]) + for name, (duration_us, count) in sorted( + aggregate.items(), key=lambda item: -item[1][0] + ): + writer.writerow([ + name, + duration_us * 1000.0, + count, + name, + f"{duration_us:.4f}", + count, + ]) + return len(aggregate) + + count = 0 + with open(output, "w", newline="", encoding="utf-8-sig") as dst: + writer = csv.writer(dst) + writer.writerow([*header, *helper]) + for row in reader: + if not row or all(not (cell or "").strip() for cell in row): + continue + name = row[indexes["name"]] if indexes["name"] < len(row) else "" + try: + duration_ns = _duration_or_none(row[indexes["total_ns"]]) + except (IndexError, TypeError): + duration_ns = None + if duration_ns is None: + continue + duration_us = f"{duration_ns / 1000.0:.4f}" + instances = "" + try: + if indexes["instances"] is not None: + instances = int(float(row[indexes["instances"]])) + except (IndexError, ValueError): + pass + writer.writerow([*row, name, duration_us, instances]) + count += 1 + return count + + +def _artifact_pid_map(root, suffix): + """Return a PID-keyed map for one top-level rocprof worker artifact type.""" + root = Path(root) + pattern = re.compile(rf'^.+_(?P\d+)_{re.escape(suffix)}$') + by_pid = {} + for path in sorted(root.glob(f'*_{suffix}')): + match = pattern.fullmatch(path.name) + if match is None: + raise SystemExit( + f'ERROR: cannot derive worker PID from {path.name!r}' + ) + pid = int(match.group('pid')) + if pid in by_pid: + raise SystemExit( + f'ERROR: duplicate {suffix} artifacts for PID {pid}: ' + f'{by_pid[pid]}, {path}' + ) + by_pid[pid] = path.resolve() + return by_pid + + +def _expected_local_worker_count(): + raw = os.environ.get('ROCPROF_EXPECT_PER_NODE', '8') + try: + expected = int(raw) + except ValueError as error: + raise SystemExit( + f'ERROR: invalid ROCPROF_EXPECT_PER_NODE={raw!r}' + ) from error + if expected <= 0: + raise SystemExit( + f'ERROR: ROCPROF_EXPECT_PER_NODE must be positive, got {expected}' + ) + return expected + + +def _worker_rank_map(log_path): + """Derive only ranks explicitly tied to a worker PID in the node log.""" + if not log_path.is_file(): + return {} + workers = {} + with open(log_path, encoding='utf-8', errors='replace') as log_file: + for lineno, line in enumerate(log_file, 1): + match = _WORKER_GPU_RE.search(line) + if match is None: + continue + row = { + key: ( + int(match.group(key)) + if match.group(key) is not None else None + ) + for key in ( + 'pid', + 'local_rank', + 'dp_rank', + 'tp_rank', + 'ep_rank', + ) + } + row['source_line'] = lineno + pid = row['pid'] + previous = workers.get(pid) + if previous is None: + workers[pid] = row + continue + if previous['local_rank'] != row['local_rank']: + raise SystemExit( + f'ERROR: conflicting gpu_id mappings for PID {pid} ' + f'in {log_path}' + ) + for key in ('dp_rank', 'tp_rank', 'ep_rank'): + if ( + previous[key] is not None + and row[key] is not None + and previous[key] != row[key] + ): + raise SystemExit( + f'ERROR: conflicting {key} mappings for PID {pid} ' + f'in {log_path}' + ) + if previous[key] is None and row[key] is not None: + previous[key] = row[key] + return workers + + +def _discover_kernel_workers(root): + """Verify and deterministically order every local GPU worker in one node dir.""" + root = Path(root).resolve() + if not root.is_dir(): + raise SystemExit(f'ERROR: capture directory missing: {root}') + expected = _expected_local_worker_count() + kernels = _artifact_pid_map(root, 'kernel_trace.csv') + markers = _artifact_pid_map(root, 'marker_api_trace.csv') + results = _artifact_pid_map(root, 'results.json') + if len(kernels) != expected: + raise SystemExit( + f'ERROR: {root.name} has {len(kernels)}/{expected} ' + 'kernel_trace.csv worker files' + ) + kernel_pids = set(kernels) + for suffix, artifacts in ( + ('marker_api_trace.csv', markers), + ('results.json', results), + ): + artifact_pids = set(artifacts) + if artifact_pids != kernel_pids: + raise SystemExit( + f'ERROR: PID mismatch for {suffix} in {root.name}: ' + f'missing={sorted(kernel_pids - artifact_pids)}, ' + f'extra={sorted(artifact_pids - kernel_pids)}' + ) + for pid, kernel_csv in kernels.items(): + prefix = kernel_csv.name[:-len('_kernel_trace.csv')] + expected_names = { + 'marker_api_trace.csv': f'{prefix}_marker_api_trace.csv', + 'results.json': f'{prefix}_results.json', + } + for suffix, artifacts in ( + ('marker_api_trace.csv', markers), + ('results.json', results), + ): + if artifacts[pid].name != expected_names[suffix]: + raise SystemExit( + f'ERROR: worker PID {pid} has mismatched {suffix}: ' + f'{artifacts[pid].name!r}, expected ' + f'{expected_names[suffix]!r}' + ) + + log_path = ( + root.parent / f"{root.name.removeprefix('rocprof_')}.log" + ) + rank_map = _worker_rank_map(log_path) + known = { + pid: rank_map[pid] + for pid in kernel_pids + if pid in rank_map + } + local_ranks = [row['local_rank'] for row in known.values()] + if len(local_ranks) != len(set(local_ranks)): + raise SystemExit( + f'ERROR: duplicate derived gpu_id ranks in {log_path}: ' + f'{sorted(local_ranks)}' + ) + if any(rank < 0 or rank >= expected for rank in local_ranks): + raise SystemExit( + f'ERROR: out-of-range derived gpu_id ranks in {log_path}: ' + f'{sorted(local_ranks)}' + ) + ranks_derivable = len(known) == expected + if ranks_derivable and sorted(local_ranks) != list(range(expected)): + raise SystemExit( + f'ERROR: incomplete derived local ranks in {log_path}: ' + f'got {sorted(local_ranks)}, expected {list(range(expected))}' + ) + if ranks_derivable: + ordered_pids = sorted( + kernel_pids, + key=lambda pid: rank_map[pid]['local_rank'], + ) + else: + ordered_pids = sorted( + kernel_pids, + key=lambda pid: (pid, kernels[pid].name), + ) + + workers = [] + for source_order, pid in enumerate(ordered_pids): + kernel_csv = kernels[pid] + prefix = kernel_csv.name[:-len('_kernel_trace.csv')] + hostname = prefix.rsplit('_', 1)[0] + rank = rank_map.get(pid, {}) + workers.append({ + 'source_order': source_order, + 'pid': pid, + 'hostname': hostname, + 'local_rank': rank.get('local_rank'), + 'dp_rank': rank.get('dp_rank'), + 'tp_rank': rank.get('tp_rank'), + 'ep_rank': rank.get('ep_rank'), + 'rank_source_line': rank.get('source_line'), + 'filename': kernel_csv.name, + 'kernel_csv': str(kernel_csv), + 'marker_csv': str(markers[pid]), + 'results_json': str(results[pid]), + }) + + node_match = _NODE_CAPTURE_RE.fullmatch(root.name) + return { + 'root': str(root), + 'role': node_match.group('role') if node_match else None, + 'node_rank': int(node_match.group('node_rank')) if node_match else None, + 'expected_worker_count': expected, + 'included_worker_count': len(workers), + 'rank_source': str(log_path.resolve()) if log_path.is_file() else None, + 'ranks_derivable': ranks_derivable, + 'workers': workers, + } + + +def _analysis_artifact_for_pid(root, pid): + root = Path(root) + for pattern in ('*_results.pftrace', '*_kernel_trace.csv'): + artifacts = [ + path for path in sorted(root.rglob(pattern)) + if (match := _ANALYSIS_PID_RE.fullmatch(path.name)) + and int(match.group('pid')) == int(pid) + ] + if artifacts: + break + if len(artifacts) != 1: + raise SystemExit( + f'ERROR: expected one analysis artifact for PID {pid}: {artifacts}' + ) + return artifacts[0] + + +def _local_rank_zero_trace(root): + root = Path(root) + log_path = root.parent / f"{root.name.removeprefix('rocprof_')}.log" + if not log_path.is_file(): + raise SystemExit(f"ERROR: worker log missing for {root.name}: {log_path}") + with open(log_path, encoding="utf-8", errors="replace") as log_file: + pids = {match.group("pid") for line in log_file + if (match := _WORKER_GPU0_RE.search(line))} + if len(pids) != 1: + raise SystemExit(f"ERROR: expected one local-rank-0 PID in {log_path}, found {sorted(pids)}") + pid = pids.pop() + return _analysis_artifact_for_pid(root, pid) + + +def _first_match(root, patterns): + root = Path(root) + if not root.is_dir(): + return None + for pattern in patterns: + matches = sorted(root.rglob(pattern)) + if matches: + return matches[0] + return None + + +def _run_logged(command, log_path, required=True, env=None): + log_path.parent.mkdir(parents=True, exist_ok=True) + with open(log_path, "w", encoding="utf-8") as log: + try: + result = subprocess.run( + [str(item) for item in command], + stdout=log, + stderr=subprocess.STDOUT, + check=False, + env=env, + ) + except FileNotFoundError as error: + if required: + raise SystemExit(f"ERROR: command not found: {command[0]}") from error + return 127 + if required and result.returncode: + raise SystemExit( + f"ERROR: command failed ({result.returncode}); see {log_path}" + ) + return result.returncode + + +def _analyze(argv): + parser = argparse.ArgumentParser( + prog="trace_tools.py analyze", + description=( + "Analyze one Torch/Kineto or rocprof trace. The rocprof CSV-only " + "path is fully offline; TraceLens/traceconv are optional paths." + ), + ) + parser.add_argument("trace") + parser.add_argument("outdir") + parser.add_argument("label") + parser.add_argument( + "--mode", + choices=("torch", "rocprof"), + default=os.environ.get("ANALYZE_MODE") or None, + ) + parser.add_argument( + "--trim-pct", + type=float, + default=float(os.environ.get("TRIM_PCT", "5")), + ) + parser.add_argument( + "--traceconv", + default=os.environ.get( + "TRACECONV", + str( + Path(os.environ.get("TRACELENS_DIR", Path.home())) + / "tracelens_test" + / "tryout" + / "traceconv" + ), + ), + ) + args = parser.parse_args(argv) + source = Path(args.trace) + output = Path(args.outdir) + tracelens_dir = output / "tracelens" + csv_dir = tracelens_dir / "out_csvs" + bucket_dir = output / "buckets" + csv_dir.mkdir(parents=True, exist_ok=True) + bucket_dir.mkdir(parents=True, exist_ok=True) + command_env = os.environ.copy() + venv = Path( + command_env.get( + "VENV", + str( + Path(command_env.get("TRACELENS_DIR", Path.home())) + / "tracelens_test" + / "venv" + ), + ) + ) + if (venv / "bin").is_dir(): + command_env["VIRTUAL_ENV"] = str(venv) + command_env["PATH"] = ( + str(venv / "bin") + os.pathsep + command_env.get("PATH", "") + ) + + mode = args.mode or ( + "torch" + if str(source).endswith((".trace.json", ".trace.json.gz")) + else "rocprof" + ) + kernel_csv = None + pftrace = None + perfetto_json = None + coverage = None + supplementary_worker = None + if source.is_dir(): + coverage = _discover_kernel_workers(source) + rank_zero = [ + worker for worker in coverage['workers'] + if worker['local_rank'] == 0 + ] + if len(rank_zero) > 1: + raise SystemExit( + f"ERROR: multiple derived local-rank-0 workers in {source}" + ) + if rank_zero: + supplementary_worker = rank_zero[0] + selected = _analysis_artifact_for_pid( + source, + supplementary_worker['pid'], + ) + else: + selected = Path(coverage['workers'][0]['kernel_csv']) + is_csv = str(selected).endswith("_kernel_trace.csv") + pftrace, kernel_csv = (None, selected) if is_csv else (selected, None) + print( + '[trace_tools analyze] canonical_kernel_scope=' + 'pooled all workers' + ) + print( + '[trace_tools analyze] aggregation_semantics=' + 'durations and call counts are summed, never averaged; ' + 'summed activity is not wall time or utilization' + ) + if supplementary_worker is not None: + print( + '[trace_tools analyze] supplementary_trace_scope=' + f"local_rank=0 pid={supplementary_worker['pid']} " + f"source={selected}" + ) + else: + print( + '[trace_tools analyze] supplementary_trace_scope=' + 'unavailable (local rank not derivable); deterministic ' + 'first source is used only for legacy input plumbing' + ) + elif str(source).endswith("_kernel_trace.csv"): + kernel_csv = source + elif source.suffix == ".pftrace": + pftrace = source + elif str(source).endswith((".json", ".json.gz")): + perfetto_json = source + + summary_csv = None + if mode == "torch": + if not source.is_file(): + raise SystemExit(f"ERROR: trace missing: {source}") + _run_logged( + [ + "TraceLens_generate_perf_report_pytorch", + "--profile_json_path", source, + "--output_xlsx_path", tracelens_dir / "report_TP0.xlsx", + "--output_csvs_dir", csv_dir, + "--enable_kernel_summary", + ], + tracelens_dir / "tracelens_stdout.txt", + env=command_env, + ) + summary_csv = _first_match(csv_dir, ("kernel_summary*.csv",)) + else: + if pftrace is not None: + traceconv = Path(args.traceconv) + if not traceconv.is_file(): + raise SystemExit(f"ERROR: traceconv not found: {traceconv}") + perfetto_json = tracelens_dir / f"{pftrace.stem}.json" + _run_logged( + [sys.executable, traceconv, "json", pftrace, perfetto_json], + tracelens_dir / "traceconv_stdout.txt", + env=command_env, + ) + if kernel_csv is not None: + summary_csv = kernel_csv + else: + if perfetto_json is None or not perfetto_json.is_file(): + raise SystemExit(f"ERROR: no usable rocprof input: {source}") + _run_logged( + [ + "TraceLens_generate_perf_report_pftrace_hip_activity", + "--trace_path", perfetto_json, + "--output_xlsx_path", tracelens_dir / "report_TP0.xlsx", + "--output_csvs_dir", csv_dir, + "--output_md_path", tracelens_dir / "report.md", + "--traceconv", args.traceconv, + "--kernel_summary_include_rccl", + "--write_md", + ], + tracelens_dir / "tracelens_stdout.txt", + env=command_env, + ) + for suffix, executable in ( + ("hip_api", "TraceLens_generate_perf_report_pftrace_hip_api"), + ("memcpy", "TraceLens_generate_perf_report_pftrace_memory_copy"), + ): + _run_logged( + [ + executable, + "--trace_path", perfetto_json, + "--output_csvs_dir", tracelens_dir / f"out_csvs_{suffix}", + "--traceconv", args.traceconv, + ], + tracelens_dir / f"tracelens_{suffix}_stdout.txt", + required=False, + env=command_env, + ) + summary_csv = _first_match(csv_dir, ("kernel_summary*.csv",)) + category_summary = _first_match( + csv_dir, ("category_summary*.csv",) + ) + if category_summary is not None: + shutil.copy2( + category_summary, + bucket_dir / "tracelens_native_category_summary.csv", + ) + if coverage is not None: + with open( + tracelens_dir / "tracelens_stdout.txt", + "a", + encoding="utf-8", + ) as tracelens_log: + tracelens_log.write( + "\n[trace_tools] scope=supplementary local_rank=0; " + "canonical kernel CSVs are pooled all-worker outputs\n" + ) + + if summary_csv is None: + raise SystemExit("ERROR: no kernel summary was produced") + normalized = bucket_dir / "kernel_summary_normalized.csv" + if coverage is not None: + pooled_audit = run_pooled_kernel_summaries( + coverage, + normalized, + bucket_dir / "kernel_summary_trimmed.csv", + args.trim_pct, + bucket_dir / "trimmed_summary_stdout.txt", + ) + normalized_count = pooled_audit[ + 'normalized_distinct_kernel_count' + ] + else: + normalized_count = normalize_kernel_summary(summary_csv, normalized) + process( + normalized, + bucket_dir / "perkernel_buckets.csv", + bucket_dir / "bycat_buckets.csv", + categorize_kernel, + ( + f"{args.label} (pooled all workers)" + if coverage is not None else f"{args.label} ({mode})" + ), + ) + + if coverage is None: + trim_source = kernel_csv + if trim_source is None and pftrace is not None: + sibling = Path(str(pftrace).replace( + "_results.pftrace", "_kernel_trace.csv" + )) + trim_source = sibling if sibling.is_file() else _first_match( + pftrace.parent, ("*_kernel_trace.csv",) + ) + if trim_source is not None: + rows, total_ns = build_trimmed_summary( + trim_source, + args.trim_pct, + ) + write_csv( + rows, + bucket_dir / "kernel_summary_trimmed.csv", + True, + categorize_kernel, + ) + print( + f"[trace_tools analyze] trimmed={len(rows)} kernels " + f"total={pretty_ns(total_ns)}" + ) + print( + f"[trace_tools analyze] {args.label}: normalized={normalized_count} " + f"outputs={bucket_dir}" + ) + return 0 + + +def _self_test_categories(argv): + parser = argparse.ArgumentParser( + prog="trace_tools.py self-test-categories", + description="Run kernel classifier regression cases.", + ) + parser.parse_args(argv) + cases = [ + ("EpDispatchInterNodeV1Kernel[LowLatency]_fp8_fnuz.kd", "MORI EP"), + ("EpCombineInterNodeV1Kernel[LowLatency]_bf16.kd", "MORI EP"), + ("EpDispatchIntraNodeKernel_fp8_fnuz", "MORI EP"), + ("EpCombineIntraNodeKernel_bf16_nop2p", "MORI EP"), + ("EpCombineSyncBarrier_bf16.kd", "Communication"), + ("EpDispatchCopyToStaging_fp8_fnuz.kd", "Communication"), + ("rmsnorm_kernel", "RMSNorm"), + ("some_rope_kernel", "ROPE"), + ("_ZN5aiter13allgather_vecIfEEvPT_S2_ii", "Communication"), + ("Cijk_Ailk_Bljk_HHS_BH_MT128x128", "GEMM"), + ("moe_topk_softmax_kernel", "MoE_TopK"), + ("some_unknown_kernel_xyz", "Other"), + ] + failures = [ + (name, expected, categorize_kernel(name)) + for name, expected in cases + if categorize_kernel(name) != expected + ] + for name, expected, actual in failures: + print( + f"FAIL: {name!r} expected {expected!r}, got {actual!r}", + file=sys.stderr, + ) + if failures: + return 1 + print(f"OK: {len(cases)}/{len(cases)} category regression cases passed.") + return 0 + + +def main(argv=None): + commands = { + "build-trace": _legacy_build_trace, + "correlate": _legacy_correlate, + "reqstats": _legacy_reqstats, + "buckets": _legacy_buckets, + "trimmed-summary": _legacy_trimmed_summary, + "analyze": _analyze, + "self-test-categories": _self_test_categories, + } + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=commands) + parser.add_argument("arguments", nargs=argparse.REMAINDER) + args = parser.parse_args(argv) + if args.command in { + "build-trace", + "correlate", + "reqstats", + "buckets", + "trimmed-summary", + }: + return _run_absorbed( + args.command, commands[args.command], args.arguments + ) + return commands[args.command](args.arguments) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/sglang_disagg/run_xPyD_models.slurm b/scripts/sglang_disagg/run_xPyD_models.slurm index 34949ddd..f8c4089c 100755 --- a/scripts/sglang_disagg/run_xPyD_models.slurm +++ b/scripts/sglang_disagg/run_xPyD_models.slurm @@ -84,7 +84,7 @@ model_allows_mori_ep() { # RUN_MORI=1: uses mori IO (default). # --------------------------------------------------------------------------- RUN_FILE="sglang_disagg_mori_io_ep.sh" -_run_mori="${RUN_MORI:-0}" +_run_mori="${RUN_MORI:-${RUN_PROFILE:-0}}" if [[ "$_run_mori" == "1" ]]; then echo "RUN_MORI=1: using $RUN_FILE with MoRI IO for model '$MODEL_NAME'" else @@ -113,6 +113,34 @@ DP_MODE="${DP_MODE:-0}" #-> mori_io_ep: 1=dp flags + --dp-size/--ep-size; 0=tp f # sglang_disagg_mori_io_ep.sh (RUN_MORI=1): set to 1 to skip phases; pass from wrapper before sbatch SKIP_BENCHMARK="${SKIP_BENCHMARK:-0}" SKIP_CURL_TEST="${SKIP_CURL_TEST:-0}" +RUN_PROFILE="${RUN_PROFILE:-0}" +# Default profile runs to skip benchmark warmup; preserve explicit caller values. +if [[ -z "${SKIP_WARMUP+x}" ]]; then + if [[ "$RUN_PROFILE" == "1" ]]; then + SKIP_WARMUP=1 + else + SKIP_WARMUP=0 + fi +fi +# The target SGLang image defaults cached-prefix early send ON. Keep it OFF for +SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX="${SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX:-0}" +if [[ "$RUN_PROFILE" == "1" ]]; then + SKIP_BENCHMARK=0 + SGLANG_ROCTX=1 + MORI_ROCTX=0 + MORI_ROCTX_TRANSFER=0 + if [[ "$_run_mori" == "1" ]]; then + MORI_ROCTX=1 + MORI_ROCTX_TRANSFER=1 + fi + ROCPROF=1 + ROCPROF_FLAGS="--kernel-trace --marker-trace" + ROCPROF_DIR_BASE=/run_logs + REQ_TIME_STATS=1 +fi +export RUN_PROFILE SKIP_WARMUP SKIP_BENCHMARK SGLANG_ROCTX MORI_ROCTX MORI_ROCTX_TRANSFER REQ_TIME_STATS +export SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX +export ROCPROF ROCPROF_FLAGS ROCPROF_DIR_BASE MODEL_NAME="${MODEL_NAME:-None}" MODEL_DIR="${MODEL_DIR:-"/shared_inference/models_blog/"}" @@ -392,12 +420,27 @@ docker run --rm \ -e BENCHMARK_FILE=$BENCHMARK_FILE \ -e IPADDRS=$IPADDRS \ -e BENCHMARK_ITR=$BENCHMARK_ITR \ + -e BENCHMARK_NUM_PROMPTS="${BENCHMARK_NUM_PROMPTS:-}" \ + -e SKIP_WARMUP="$SKIP_WARMUP" \ -e SKIP_BENCHMARK=$SKIP_BENCHMARK \ -e SKIP_CURL_TEST=$SKIP_CURL_TEST \ + -e RUN_PROFILE=$RUN_PROFILE \ + -e SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX=${SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX:-0} \ -e KV_TRANSFER_BACKEND=${KV_TRANSFER_BACKEND:-} \ -e BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024}" \ -e DOCKER_IMAGE_NAME=${DOCKER_IMAGE_NAME:-unknown} \ -e USE_CX7_NICS=${USE_CX7_NICS:-0} \ + -e SGLANG_ROCTX=${SGLANG_ROCTX:-0} \ + -e MORI_ROCTX=${MORI_ROCTX:-0} \ + -e MORI_ROCTX_TRANSFER=${MORI_ROCTX_TRANSFER:-0} \ + -e ROUTER_READY_TIMEOUT_SECONDS=${ROUTER_READY_TIMEOUT_SECONDS:-4000} \ + -e SGLANG_KV_ROCTX=${SGLANG_KV_ROCTX:-0} \ + -e ROCPROF=${ROCPROF:-0} \ + -e ROCPROF_FLAGS="${ROCPROF_FLAGS:-}" \ + -e ROCPROF_DIR_BASE=${ROCPROF_DIR_BASE:-} \ + -e EAGER=${EAGER:-0} \ + -e REQ_TIME_STATS=${REQ_TIME_STATS:-1} \ + -e BENCHMARK_CON="${BENCHMARK_CON:-}" \ --ulimit nofile=1048576:1048576 \ --name $DOCKER_CONT_NAME \ --entrypoint /bin/bash \ @@ -406,5 +449,14 @@ docker run --rm \ $RUN_FILE_FULL 2>&1 | tee /run_logs/${SLURM_JOB_ID}/pd_sglang_bench_serving.sh_NODE${SLURM_PROCID}.log " ' +RUN_RC=$? srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c 'docker stop $DOCKER_CONT_NAME; docker rm -f $DOCKER_CONT_NAME' +if [[ "$RUN_PROFILE" == "1" ]]; then + srun --nodes=1 --ntasks=1 --nodelist="$MASTER_NODE" \ + env XP="$xP" YD="$yD" HR="$LOG_PATH/$SLURM_JOB_ID" RUN_MORI="$_run_mori" \ + bash "$MOONCAKE_REPO_DIR/moriio_profiling/process_kernels.sh" "$SLURM_JOB_ID" + PROCESS_RC=$? + [[ "$RUN_RC" -eq 0 ]] || exit "$RUN_RC" + exit "$PROCESS_RC" +fi diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index f2244936..aef8ab43 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -5,6 +5,9 @@ _MORI_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" SCRIPT_DIR="${_MORI_SCRIPT_DIR}" +if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + source "$SCRIPT_DIR/moriio_profiling/hooks.sh" +fi # ----------------------------------------------------------------------------- # DP_MODE=1 allowlist (MoRI IO EP). Must stay in sync with run_xPyD_models.slurm. @@ -37,6 +40,9 @@ mori_dp_mode1_allowed_models_lines() { MASTER_ADDR="${MASTER_ADDR:-localhost}" MASTER_PORT="${MASTER_PORT:-23731}" NODE_RANK="${NODE_RANK:-0}" +# SGLang's image default is true; disable cached-prefix early send unless the +export SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX="${SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX:-0}" + MODEL_PATH=$MODEL_PATH MODEL_NAME="${MODEL_NAME:-}" xP="${xP:-1}" @@ -80,7 +86,6 @@ pip install py-spy pip install --ignore-installed --force-reinstall flask pip install pyyaml - host_ip=$(ip route get 1.1.1.1 | awk '/src/ {print $7}') host_name=$(hostname) @@ -195,7 +200,17 @@ if [[ "${_TRANSFER_BACKEND}" != "mori" ]]; then echo "[override] Transfer backend: ${_TRANSFER_BACKEND}" fi - +if [[ "${EAGER:-0}" == "1" ]]; then + PREFILL_MODEL_CONFIG+=" --disable-cuda-graph" + DECODE_MODEL_CONFIG+=" --disable-cuda-graph" + export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG + echo "[eager] EAGER=1: appended --disable-cuda-graph to prefill + decode launch_server commands" +fi +if [[ "${REQ_TIME_STATS:-1}" == "1" ]]; then + PREFILL_MODEL_CONFIG+=" --enable-request-time-stats-logging" + DECODE_MODEL_CONFIG+=" --enable-request-time-stats-logging" + export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG +fi # ============================================================================= # Cluster Topology (dist-init endpoints) # ============================================================================= @@ -328,6 +343,7 @@ _wait_for_tcp() { done } + if [[ "$NODE_RANK" -eq 0 ]]; then echo "${host_name}:${host_ip} is Prefill Node 0 + Router/proxy (NODE_RANK=0, co-located)" @@ -378,7 +394,12 @@ if [[ "$NODE_RANK" -eq 0 ]]; then } | tee "$PREFILL_LOG" _dbg "launching prefill server (PREFILL_NODE_RANK=0, PREFILL_TP_SIZE=${PREFILL_TP_SIZE})" set -x - eval "$_prefill_cmd" 2>&1 | tee -a "$PREFILL_LOG" >/dev/null & + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + _ROCPROF_PREFIX="$(_rocprof_prefix prefill)" + eval "${_ROCPROF_PREFIX}$_prefill_cmd" >>"$PREFILL_LOG" 2>&1 & + else + eval "$_prefill_cmd" 2>&1 | tee -a "$PREFILL_LOG" >/dev/null & + fi set +x _node0_prefill_pid=$! _dbg "prefill server started pid=${_node0_prefill_pid}" @@ -535,22 +556,37 @@ PY fi if [[ "${SKIP_BENCHMARK:-0}" != "1" ]] && [[ -n "${MOONCAKE_COOKBOOK_PATH:-}" ]]; then - if [[ -f "${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ]]; then - echo "Running ${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" + BENCHMARK_SCRIPT="benchmark_xPyD.sh" + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + BENCHMARK_SCRIPT="benchmark_xPyD_profile.sh" + fi + if [[ -f "${MOONCAKE_COOKBOOK_PATH}/${BENCHMARK_SCRIPT}" ]]; then + echo "Running ${MOONCAKE_COOKBOOK_PATH}/${BENCHMARK_SCRIPT}" ( cd "${MOONCAKE_COOKBOOK_PATH}" || exit 1 - bash benchmark_xPyD.sh + bash "${BENCHMARK_SCRIPT}" ) else - echo "WARN: benchmark_xPyD.sh not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 + echo "WARN: ${BENCHMARK_SCRIPT} not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 fi fi echo "Killing the proxy server (pid=${proxy_pid})" - kill "${proxy_pid}" + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + pkill -TERM -f '[s]glang_router.launch_router' 2>/dev/null || true + # sgl-router renames its own process title to "sglang::router" once + pkill -KILL -x 'sglang::router' 2>/dev/null || true + else + kill "${proxy_pid}" + fi - echo "Killing the co-located prefill server (pid=${_node0_prefill_pid})" - kill "${_node0_prefill_pid}" + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + echo "Stopping the co-located prefill server (pid=${_node0_prefill_pid})" + finish_server prefill "${_node0_prefill_pid}" || exit $? + else + echo "Killing the co-located prefill server (pid=${_node0_prefill_pid})" + kill "${_node0_prefill_pid}" + fi elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-default})" @@ -607,7 +643,12 @@ elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then } | tee "$PREFILL_LOG" _dbg "launching prefill server (PREFILL_NODE_RANK=${PREFILL_NODE_RANK}, PREFILL_TP_SIZE=${PREFILL_TP_SIZE})" set -x - eval "$PREFILL_CMD" 2>&1 | tee -a "$PREFILL_LOG" >/dev/null & + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + _ROCPROF_PREFIX="$(_rocprof_prefix prefill)" + eval "${_ROCPROF_PREFIX}$PREFILL_CMD" >>"$PREFILL_LOG" 2>&1 & + else + eval "$PREFILL_CMD" 2>&1 | tee -a "$PREFILL_LOG" >/dev/null & + fi set +x prefill_pid=$! _dbg "prefill server started pid=${prefill_pid}" @@ -623,8 +664,13 @@ elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then --remote-ip "${MASTER_ADDR}" \ --remote-port 2322 - echo "Killing the prefill server" - kill "${prefill_pid}" + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + echo "Stopping the prefill server" + finish_server prefill "${prefill_pid}" || exit $? + else + echo "Killing the prefill server" + kill "${prefill_pid}" + fi elif [[ "$NODE_RANK" -ge $xP && "$NODE_RANK" -le $((xP + yD - 1)) ]]; then echo "${host_name}:${host_ip} is Decode Node (Model: ${MODEL_NAME:-default})" @@ -688,7 +734,12 @@ elif [[ "$NODE_RANK" -ge $xP && "$NODE_RANK" -le $((xP + yD - 1)) ]]; then } | tee "$DECODE_LOG" _dbg "launching decode server (DECODE_NODE_RANK=${DECODE_NODE_RANK}, DECODE_TP_SIZE=${DECODE_TP_SIZE})" set -x - eval "$DECODE_CMD" 2>&1 | tee -a "$DECODE_LOG" >/dev/null & + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + _ROCPROF_PREFIX="$(_rocprof_prefix decode)" + eval "${_ROCPROF_PREFIX}$DECODE_CMD" >>"$DECODE_LOG" 2>&1 & + else + eval "$DECODE_CMD" 2>&1 | tee -a "$DECODE_LOG" >/dev/null & + fi set +x decode_pid=$! _dbg "decode server started pid=${decode_pid}" @@ -704,8 +755,13 @@ elif [[ "$NODE_RANK" -ge $xP && "$NODE_RANK" -le $((xP + yD - 1)) ]]; then --remote-ip "${MASTER_ADDR}" \ --remote-port 2322 - echo "Killing the decode server" - kill "${decode_pid}" + if [[ "${RUN_PROFILE:-0}" == "1" ]]; then + echo "Stopping the decode server" + finish_server decode "${decode_pid}" || exit $? + else + echo "Killing the decode server" + kill "${decode_pid}" + fi else echo "ERROR: NODE_RANK=${NODE_RANK} out of range (expected 0..$((xP + yD))) for xP=${xP} yD=${yD}" >&2