diff --git a/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..edd4d797 --- /dev/null +++ b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile @@ -0,0 +1,350 @@ +# 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. +# +################################################################################# +# ============================================================================= +# vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile +# GLM-5.1-FP8 (MLA + DeepSeek Sparse Attention) MoRI-EP WideEP disagg image. +# PER-MODEL image, isolated from the base vllm_disagg_inference Dockerfile +# (which stays pinned to the DeepSeek-V3 / R1 stack). This split lets each model +# pin its own vLLM/AITER/MoRI without disturbing the others -- add a new +# vllm_disagg_inference..ubuntu.amd.Dockerfile per future model +# (e.g. Kimi-2.6) rather than repinning the shared DSV3 image. +# +# ALL connectors in one image: moriio (TP + MoRI-EP wideEP) + rixl (NIXL TP + +# DeepEP wideEP). = the fullsource MoRI stack, plus a UCX/RIXL/rocSHMEM/DeepEP +# transport layer gated by --build-arg WITH_NIXL (default 1 = everything). +# +# docker build -f docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile \ +# -t /vllm-disagg:glmv5.1 . +# export DOCKER_IMAGE_NAME=/vllm-disagg:glmv5.1 +# +# WITH_NIXL=1 (default) => builds UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from +# source, so all four connector combos (moriio TP/wideEP, rixl NIXL TP, DeepEP +# wideEP) are present (~+30-45 min build vs WITH_NIXL=0). +# WITH_NIXL=0 => MoRI-EP only (moriio TP/wideEP + deepep-from-base); lean, faster. +# +# STATUS (GLM-5.1-FP8 on this stack): 1P/1D EP8 + 2P/2D EP16 NIAH 2k-35k = 10/10, +# no crash; long-context accuracy fixed via vLLM #47766 (persistent sparse-MLA kept +# ON). 4P/4D EP32 is a KNOWN OPEN DEFECT: token corruption at ALL context lengths +# (garbage output even at 2k), distinct from the long-context bug; prime suspect is +# the moriep all-to-all combine at EP32 scale -> deferred to future work. Use 1P/1D +# and 2P/2D only. (BASE_IMAGE is a gated nightly; override --build-arg BASE_IMAGE=...) +# ============================================================================= +# Builds the GLM-5.1 runtime stack by applying component pins ON TOP of a +# purpose-built ROCm/vLLM/MoRI base, cloning each overridden source from public Git +# (no local build-contexts): +# +# - BASE: rocmshared/pytorch-private:vllm-rocm_07_22_2026_shikpate_mori1.2.3 +# (ROCm + torch + a bundled vLLM/MoRI 1.2.3 stack). The stages below deliberately +# OVERRIDE the base's vLLM/MoRI/AITER with the pins we validate for GLM DSA. +# - MoRI -> built from ROCm/MoRI @ 42e895472b08 (validated for GLM DSA, BUILD_UMBP=OFF). +# (main LATEST 120d2de broke the connector KV-notify handshake -- see note at MORI_REF.) +# - AITER -> STOCK ROCm/aiter @ e03fa6040 compiled from source + flydsl 0.1.7-0.1.9; +# stale JIT wiped. (#47766 keeps persistent MLA ON -> aiter native gqa64 fold.) +# - vLLM -> COMPILED from raviguptaamd/vllm @ glm5.1-dsa-wideEP_on_shik_0721 +# (Shiksha 7/21 WideEP base + GLM DSA edits + sparse-MLA guard fix). Full compile: +# a different commit than the base's, so a .py-only overlay would be ABI-mismatched. +# - RDMA fix (expandable_segments:False x2 + HSA_ENABLE_IPC_MODE_LEGACY=0) is NOT baked +# here — it lives in scripts/vllm_dissag/connectors/.env and the launcher +# forwards it via docker -e. ROCm 7.2.3 cannot dmabuf-export VMM memory, else MoRI +# RegisterRdmaMemoryRegion EFAULTs (errno 14) on the first disagg WRITE. +# - vllm-router (vllm-project/router PR#181 = DP-rank round-robin + 2P2D KV-notify +# dpfix) built in -> no external router binary needed. +# - validated recipe knobs baked as ENV. The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this vLLM (no runtime patcher). +# +# Build context = repo root: +# docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile -t / . +# +# BASE_IMAGE is the purpose-built ROCm/vLLM/MoRI base above. Override --build-arg +# BASE_IMAGE=... to build on a different ROCm base. vLLM compile is long (~30-60 min). +# ============================================================================= + +ARG BASE_IMAGE=rocm/vllm-dev:ci_base-dedbf6be8b1afa17a6220473b9c8c98242ac1c03 +FROM ${BASE_IMAGE} + +ENTRYPOINT [] +WORKDIR /app + +ARG GFX_COMPILATION_ARCH="gfx942" +ARG PYTORCH_ROCM_ARCH="gfx942" +ARG MAX_JOBS=32 +# NIXL/RIXL transport for the rixl connector. GLM-5.1 is served over MoRI-EP + MoRI-IO, +# so the UCX/RIXL/rocSHMEM/DeepEP stack is dead weight here: it lengthens the build and +# ships transports this recipe never selects. Default 0 => lean MoRI-EP-only image, which +# is also exactly how the validated image (glm5.1-vllm027-b8) was built. Set +# --build-arg WITH_NIXL=1 only if you need the rixl connector from this same Dockerfile. +ARG WITH_NIXL=0 +ARG NIC_COMPILATION_ARCH="cx7" + +# ----------------------------------------------------------------------------- +# 1. MoRI: replace the base's bundled MoRI with the validated ROCm/MoRI @ v1.2.1 +# (the version for the 06_29 mori121 image, dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable). v1.2.1 carries the EP/RDMA correctness fixes +# plus the ROCm-7.2.3 dmabuf registration path used by the connector .env +# (expandable_segments:False). MoRI is JIT-built, so this swaps the JIT sources the +# kernels compile from at runtime. +# BUILD CONFIG: match the cookbook build — MORI_GPU_ARCHS=gfx942, BUILD_UMBP=OFF, +# DEFAULT NIC backends. Do NOT pass USE_IONIC=OFF / USE_BNXT=OFF: disabling NIC +# backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. +# ----------------------------------------------------------------------------- +ARG MORI_REPO=https://github.com/ROCm/mori.git +# 42e895472b08: validated MoRI tip for GLM DSA WideEP disagg. The v0.27 base bundles +# amd_mori 1.0.0, but the bundled build regressed GLM DSA (b1: GPU fault on the aiter +# DSA decode kernel), so we build MoRI from source at this pinned commit by DEFAULT +# (WITH_MORI_BUILD=1). Set --build-arg WITH_MORI_BUILD=0 only to fall back to the +# base's bundled mori for debugging. +ARG WITH_MORI_BUILD=1 +ARG MORI_REF=42e895472b08 +ENV MORI_GPU_ARCHS=gfx942 +# Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not +# present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so +# disable it to avoid pulling in a gRPC build dependency. +ENV BUILD_UMBP=OFF BUILD_UMBP_SPDK=OFF +# Build/install matches dist-inf-cookbook Dockerfile.vllm.mori121_shareable for v1.2.1: +# `BUILD_UMBP=OFF pip install .` (default build isolation). apt/pip build tooling kept +# for bases that lack it; harmless where already present. +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + sed -i 's|http://|https://|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true && \ + apt-get update && apt-get install -y --no-install-recommends \ + git build-essential cmake ninja-build ccache libssl-dev pkg-config curl ca-certificates && \ + pip install meson==0.64.0 "pybind11[global]" tqdm prettytable && \ + mkdir -p /app && \ + if [ "${WITH_MORI_BUILD}" != "1" ]; then \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI (bundled) OK at', mori.__path__[0])" && \ + echo "MORI_REF=BUNDLED (base amd_mori, WITH_MORI_BUILD=0)" >> /app/versions.txt ; \ + else \ + pip uninstall -y amd_mori amd-mori amd-mori-nightly mori 2>/dev/null || true && \ + rm -rf /tmp/mori-src && \ + git clone --recursive "${MORI_REPO}" /tmp/mori-src && \ + cd /tmp/mori-src && git checkout "${MORI_REF}" && git submodule update --init --recursive && \ + BUILD_UMBP=OFF pip install . && \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI OK at', mori.__path__[0])" && \ + echo "MORI_REF=${MORI_REF}@$(git -C /tmp/mori-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/mori-src ; \ + fi + +# ----------------------------------------------------------------------------- +# 2. AITER: the v0.27 base bundles amd-aiter 0.1.19 (+ flydsl 0.2.4), but bundled 0.1.19 +# GPU-faults on the GLM DSA decode kernel mla_a8w8_qh64_gqaratio64_v3 (confirmed b1 on +# this v0.27 base, same regression as the old stack). So we build aiter from source at +# the validated commit e03fa6040 by DEFAULT (WITH_AITER_BUILD=1). aiter > e03fa6040 +# reintroduces the fault; do not bump without re-running long-ctx NIAH. Set +# --build-arg WITH_AITER_BUILD=0 only to fall back to the bundled aiter for debugging. +# ----------------------------------------------------------------------------- +ARG AITER_REPO=https://github.com/ROCm/aiter.git +ARG WITH_AITER_BUILD=1 +ARG AITER_REF=e03fa6040 +RUN if [ "${WITH_AITER_BUILD}" != "1" ]; then \ + echo "AITER: using BUNDLED base aiter (WITH_AITER_BUILD=0)" && \ + python3 -c "import importlib.metadata as m; print('aiter (bundled)', m.version('amd-aiter'))" && \ + echo "AITER_REF=BUNDLED (base amd-aiter, WITH_AITER_BUILD=0)" >> /app/versions.txt ; \ + else \ + echo "Compiling STOCK AITER (no fork) from ${AITER_REPO}@${AITER_REF}" && \ + rm -rf /tmp/aiter-src && \ + git clone --recursive "${AITER_REPO}" /tmp/aiter-src && \ + cd /tmp/aiter-src && git checkout "${AITER_REF}" && \ + git submodule update --init --recursive && \ + (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ + pip install --no-build-isolation --no-deps -v . && \ + pip install --no-deps -U "flydsl>=0.1.7,<0.1.9" && \ + echo "AITER_REF=${AITER_REF}@$(git rev-parse HEAD) (stock ROCm/aiter, no fork)" >> /app/versions.txt && \ + rm -rf /tmp/aiter-src && \ + rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" ; \ + fi + +# ----------------------------------------------------------------------------- +# 3. vLLM: compile from source at the 06_29 validated Wide-EP WRITE-mode branch +# (matches the published dist-inf-cookbook mori121 image). Full source compile +# (the base ships a different commit). The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this branch, so no runtime +# patcher is needed. Override VLLM_REF to rebuild a different commit; build only +# committed commits (no working-tree edits). +# ----------------------------------------------------------------------------- +# VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the +# dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +# glm5.1-dsa-wideEP_on_vllm-v0.27 (HEAD cda3648602) = upstream v0.27 tip dedbf6be8b + 7 +# ROCm/DSA commits. Core 3: per-req-ctx metadata key (#47766), DSA indexer KV transfer +# (reworked onto upstream's native MoRIIO connector), invalid-token sentinel. Plus 4 +# v0.27 fixes: concat_and_cache_mla positional (stable-ABI), splitting_ops out of the +# compiled graph (MLA "unknown parameter type"), sparse-indexer bounds-guard, and the +# decisive sentinel -1->0 (cda3648602 — aiter mla_decode_fwd derefs -1 -> GPU fault at +# disagg long-ctx). NIAH-validated 1P/1D + 2P/1D + 1P/2D, 2k-35k, decode PIECEWISE. +ARG VLLM_REF=glm5.1-dsa-wideEP_on_vllm-v0.27 +ENV VLLM_TARGET_DEVICE=rocm \ + PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ + MAX_JOBS=${MAX_JOBS} +RUN rm -rf /tmp/vllm-src && \ + git clone "${VLLM_REPO}" /tmp/vllm-src && \ + cd /tmp/vllm-src && git checkout "${VLLM_REF}" && \ + echo "VLLM_REF=${VLLM_REF}@$(git rev-parse HEAD)" >> /app/versions.txt && \ + pip uninstall -y vllm 2>/dev/null || true && \ + pip install --no-deps --no-build-isolation -v . && \ + python3 -c "import vllm; print('vLLM', vllm.__version__, 'from', vllm.__file__)" && \ + rm -rf /tmp/vllm-src + +# Cross-check MoRI + AITER survived the vLLM install (no silent downgrade). +RUN python3 - <<'PYEOF' +from importlib.metadata import version as v, PackageNotFoundError +def get(names): + for n in names: + try: return v(n) + except PackageNotFoundError: pass + return None +av = get(("amd-aiter", "amd_aiter", "aiter")) +# Verify the aiter install survived the vLLM install (present, not silently downgraded +# to a base-bundled wheel). We pin aiter by commit (e03fa6040), whose reported version +# string varies by build, so assert presence rather than a hardcoded commit substring. Do NOT +# `import aiter` here: it pulls torch->amdsmi->libamd_smi.so, not loadable in the no-GPU +# build sandbox (same reason the Stage-2 verify reads mla.py from disk instead). +assert av, "AITER missing after vLLM install (expected bundled 0.1.19 or source-built ref)" +import mori, mori.io, mori.ops +print("Post-vLLM check OK: AITER", av, "present + MoRI importable") +PYEOF + +# ----------------------------------------------------------------------------- +# 4. vllm-router (DP-rank round-robin + MoRIIO connector) — built in, so NO +# external vllm-router binary is needed (leave ROUTER_BINARY unset). +# Source = vllm-project/router PR #181 branch, which now carries BOTH the +# round-robin DP-rank fix (11841c0d) AND the 2P2D KV-notify fix (6409ac1: +# remote_dp_rank_override + remote_dp_size). The KV-notify fix is REQUIRED: +# without it the 2P2D EP=16 run reproducibly wedges with "remote blocks never +# arrived" deferred-write expiries (decode notify targets the wrong DP rank). +# This is the exact source of the validated vllm-router-2p2d-dpfix binary. +# Pinned Rust toolchain (>=1.88: router deps time/home require rustc 1.88). +# ----------------------------------------------------------------------------- +ARG ROUTER_REPO=https://github.com/raviguptaamd/router.git +ARG ROUTER_REF=ravgupta/discovery-dp-rank-roundrobin +ARG RUST_TOOLCHAIN=1.88.0 +RUN if ! command -v cargo >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_TOOLCHAIN}"; \ + fi && \ + export PATH="/root/.cargo/bin:${PATH}" && \ + rm -rf /tmp/vllm-router-src && \ + git clone --filter=blob:none "${ROUTER_REPO}" /tmp/vllm-router-src && \ + cd /tmp/vllm-router-src && git checkout "${ROUTER_REF}" && \ + cargo build --release && \ + install -m 755 target/release/vllm-router /usr/local/bin/vllm-router && \ + vllm-router --help 2>&1 | grep -q moriio && \ + echo "VLLM_ROUTER_REF=${ROUTER_REPO}@${ROUTER_REF}@$(git -C /tmp/vllm-router-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/vllm-router-src + +# ----------------------------------------------------------------------------- +# 4b. WITH_NIXL=1 (default): UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Single guarded RUN so +# WITH_NIXL=0 skips it entirely (no layers, no cost). Build-verified on ci_base. +# ----------------------------------------------------------------------------- +ENV _ROCM_DIR=/opt/rocm \ + _UCX_SOURCE=https://github.com/ROCm/ucx.git \ + _UCX_BRANCH=da3fac2a \ + _UCX_INSTALL_DIR=/usr/local/ucx/ \ + _RIXL_SOURCE=https://github.com/ROCm/RIXL.git \ + _RIXL_BRANCH=f33a5599 \ + _RIXL_INSTALL_DIR=/usr/local/RIXL/install \ + _NIXLBENCH_INSTALL_DIR=/usr/local/RIXL +RUN if [ "${WITH_NIXL}" != "1" ]; then \ + echo "WITH_NIXL=${WITH_NIXL}: skipping UCX/RIXL/rocSHMEM/DeepEP (MoRI-EP + base DeepEP only)"; \ + else set -e && \ + echo "WITH_NIXL=1: building UCX + RIXL + rocSHMEM + DeepEP" && \ + apt-get update && apt-get install -y \ + autoconf automake libtool autogen pkg-config m4 gcc make \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core strace libgflags-dev \ + libaio-dev liburing-dev libcpprest-dev libgrpc-dev libgrpc++-dev \ + libprotobuf-dev protobuf-compiler-grpc wget && \ + pip install meson==0.64.0 "pybind11[global]" pyyaml && \ + # UCX + cd /tmp && git clone "${_UCX_SOURCE}" && cd ucx && git checkout "${_UCX_BRANCH}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${_UCX_INSTALL_DIR}" --with-rocm="${_ROCM_DIR}" \ + --disable-go --disable-java --disable-assertions --enable-mt && \ + make -j && make install && \ + # googletest (RIXL dep) + cd /tmp && wget -q https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz && \ + tar -xzf v1.14.0.tar.gz && cd googletest-1.14.0 && mkdir -p build && cd build && \ + cmake -DBUILD_SHARED_LIBS=on .. && make -j && make install && \ + # RIXL + python bindings + cd /tmp && git clone "${_RIXL_SOURCE}" && cd RIXL && git checkout "${_RIXL_BRANCH}" && \ + meson setup build/ --prefix="${_RIXL_INSTALL_DIR}" -Ducx_path="${_UCX_INSTALL_DIR}" \ + -Ddisable_gds_backend=true -Dcudapath_inc="${_ROCM_DIR}/include" -Dcudapath_lib="${_ROCM_DIR}/lib" && \ + cd build && ninja && ninja install && cd /tmp/RIXL && \ + pip install --config-settings=setup-args="-Dcudapath_inc=${_ROCM_DIR}/include" \ + --config-settings=setup-args="-Dcudapath_lib=${_ROCM_DIR}/lib" \ + --config-settings=setup-args="-Ducx_path=${_UCX_INSTALL_DIR}" \ + --config-settings=setup-args="-Ddisable_gds_backend=true" . && \ + # rocSHMEM (DeepEP dep) + cd /tmp && git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git && \ + cd rocm-systems && git sparse-checkout set --cone projects/rocshmem && git checkout develop && \ + mkdir -p /tmp/rocshmem-build && cd /tmp/rocshmem-build && \ + /tmp/rocm-systems/projects/rocshmem/scripts/build_configs/all_backends \ + -DUSE_EXTERNAL_MPI=OFF -DGPU_TARGETS="${GFX_COMPILATION_ARCH}" && \ + # DeepEP (build develop against the installed vLLM/torch) + cd /tmp && git clone https://github.com/ROCm/DeepEP.git && cd DeepEP && \ + PYTORCH_ROCM_ARCH="${GFX_COMPILATION_ARCH}" CFLAGS="-O3 -fPIC" \ + CXXFLAGS="-O3 -fPIC --offload-arch=${GFX_COMPILATION_ARCH}" HIP_CXX_FLAGS="-O3 -fPIC" \ + python3 setup.py --variant rocm --nic "${NIC_COMPILATION_ARCH}" build develop && \ + echo "WITH_NIXL build complete" >> /app/versions.txt && \ + rm -rf /tmp/ucx /tmp/googletest-1.14.0 /tmp/v1.14.0.tar.gz /tmp/rocm-systems /tmp/rocshmem-build; \ + fi +ENV LD_LIBRARY_PATH="/usr/local/ucx/lib:/usr/local/lib:/usr/local/RIXL/install/lib:${LD_LIBRARY_PATH}" \ + PATH="/usr/local/ucx/bin:${PATH}" + +# ----------------------------------------------------------------------------- +# 5. Cache locations (structural: WHERE the JIT/compile caches live in the image). +# These are the mount target for the launcher's persistent host JIT cache. +# ----------------------------------------------------------------------------- +# The image ships NO runtime recipe / tuning / platform ENV. By design, everything +# run-tunable is applied at launch, so this image stays a clean binary/library artifact +# and the same image serves any model/cluster without a rebuild: +# - model-serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, *_CUDAGRAPH_MODE, *_MORI_BACKEND, +# GPU_MEMORY_UTILIZATION, KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...) +# -> scripts/vllm_dissag/models.yaml (per-model env:, so dense vs MoE differ) +# - ROCm-7.2.3 GPU-RDMA platform env (expandable_segments:False x2, MORI_GPU_ARCHS, +# HSA_ENABLE_IPC_MODE_LEGACY=0, HSA_NO_SCRATCH_RECLAIM) and the MoRI/RDMA fabric +# tuning (MORI_RDMA_TC/SL, MORI_IB_GID_INDEX, MORI_NUM_QP_PER_PE, VLLM_MORIIO_*, ...) +# -> scripts/vllm_dissag/connectors/.env (cluster-editable, no rebuild) +# The slurm launcher forwards both via `docker -e` (platform env must reach PID 1 - +# PyTorch reads alloc-conf at import). Running this image WITHOUT the launcher: set the +# vars you need yourself (see connectors/moriio.env + models.yaml for the values). +ENV AITER_JIT_DIR=/opt/vllm_cache/aiter_jit \ + VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + COMGR_CACHE_DIR=/opt/vllm_cache/comgr + +# ----------------------------------------------------------------------------- +# 6. CRITICAL: scrub build-time MoRI JIT state. The `import mori` verification +# steps above compile/lock MoRI EP kernels under /root/.mori/jit on THIS build +# host, leaving stale .hsaco.lock files (ep_internode_v1, ep_internode_v1ll, ...). +# At runtime on the cluster, MoriAll2AllManager finds those locks, waits on a +# build-in-progress whose owner PID is long gone, and DEADLOCKS at ep:0 init. +# A clean image ships /root/.mori empty -> runtime compiles fresh. +# Clearing these makes the from-source image boot clean on 2P2D/4P4D. +# ----------------------------------------------------------------------------- +RUN rm -rf /root/.mori /tmp/mori_jit_* && mkdir -p /root/.mori && \ + echo "JIT_SCRUBBED: /root/.mori + /tmp/mori_jit_* cleared at build end" >> /app/versions.txt + +RUN cat /app/versions.txt 2>/dev/null | tail -20 || true diff --git a/models.json b/models.json index af914f46..689754db 100644 --- a/models.json +++ b/models.json @@ -3420,6 +3420,38 @@ }, "args": "-N 2 -n 2" }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8", + "url": "", + "dockerfile": "docker/vllm_disagg_inference.glmv5.1", + "scripts": "scripts/vllm_dissag/run_xPyD_models.slurm", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "GLM-5.1-FP8", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "GLM_SKIP_PATCHERS": "1", + "BENCHMARK_COMBINATIONS": "1024/1024" + }, + "args": "-N 2 -n 2" + }, { "name": "pyt_vllm_disagg_mori_deepseek-v3-5layer", "url": "", diff --git a/scripts/vllm_dissag/.nfs0000000016f44d2b00008188 b/scripts/vllm_dissag/.nfs0000000016f44d2b00008188 new file mode 100755 index 00000000..52d046a1 --- /dev/null +++ b/scripts/vllm_dissag/.nfs0000000016f44d2b00008188 @@ -0,0 +1,685 @@ +#!/bin/bash +#SBATCH --job-name=vllm-pd # Specify a custom string for your slurm batch job +#SBATCH -N 2 # Default 2 nodes (1P/1D); override with sbatch -N for larger topologies +#SBATCH --ntasks-per-node=1 +#SBATCH --spread-job +#SBATCH --gres=gpu:8 # Request 8 GPUs and 8 NICs (use --gres if specific GPU resources are needed) +#SBATCH --time=24:00:00 # Set a time limit for the job (HH:MM:SS) +#SBATCH --output="/shared_inference/%u/model_blog_logs/slurm-%j.out" +#SBATCH --error="/shared_inference/%u/model_blog_logs/slurm-%j.err" + + +# ------------------------ +# Auto-detect and cd to the script directory so that $(pwd) always +# points to the folder containing the server scripts, regardless of +# where the user called sbatch from. +# +# Priority: +# 1. BASH_SOURCE — works for direct invocation (bash script.sh) +# 2. SLURM_SUBMIT_DIR + script path — when sbatch submits from repo root, +# SLURM_SUBMIT_DIR is the CWD, not the script dir. Append the relative +# path from the SBATCH command to get the actual script directory. +# 3. SLURM_SUBMIT_DIR alone — last resort (assumes sbatch was run from +# the script directory). +# ------------------------ +_resolve_script_dir() { + # Try BASH_SOURCE first (works for direct invocation) + if [[ -n "${BASH_SOURCE[0]:-}" ]]; then + local _d + _d="$(cd "$(dirname "${BASH_SOURCE[0]}")" 2>/dev/null && pwd)" + if [[ -n "$_d" && -f "$_d/vllm_disagg.sh" ]]; then + echo "$_d" + return 0 + fi + fi + + # Try SLURM_SUBMIT_DIR + relative script path + # When madengine does: sbatch scripts/vllm_dissag/run_xPyD_models.slurm + # SLURM_SUBMIT_DIR = repo root, so we need to append the dirname + if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then + local _candidate="$SLURM_SUBMIT_DIR/scripts/vllm_dissag" + if [[ -f "$_candidate/vllm_disagg.sh" ]]; then + echo "$_candidate" + return 0 + fi + # Maybe they ran sbatch from the script dir itself + if [[ -f "$SLURM_SUBMIT_DIR/vllm_disagg.sh" ]]; then + echo "$SLURM_SUBMIT_DIR" + return 0 + fi + fi + + # Fallback + echo "." + return 1 +} + +SCRIPT_DIR="$(_resolve_script_dir)" +cd "$SCRIPT_DIR" || { echo "Error: cannot cd to $SCRIPT_DIR" >&2; exit 1; } + +REQUIRED_FILES=("vllm_disagg.sh" "parallelism.sh" "connectors/rixl.sh" "connectors/moriio.sh" "models.yaml" "benchmark_xPyD.sh" "parse_to_csv.py" "socket_barrier.py" "socket_wait.py" "connectors/moriio.env" "connectors/rixl.env") +for f in "${REQUIRED_FILES[@]}"; do + if [[ ! -f "$f" ]]; then + echo "Error: Required file '$f' not found in $(pwd)." >&2 + echo "Please run sbatch from the scripts/vllm_dissag/ directory, e.g.:" >&2 + echo " cd MAD/scripts/vllm_dissag && sbatch run_xPyD_models.slurm" >&2 + exit 1 + fi +done +echo "Running from: $(pwd)" + +# ------------------------------------------------------------------------------ +# models.yaml env precedence: capture which recipe knobs the USER explicitly set +# at submit time. The driver (vllm_disagg.sh) uses this to let models.yaml `env:` +# OVERRIDE image-baked ENV defaults (e.g. a DeepSeek-tuned image bakes +# KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0, which would otherwise shadow a +# model's own recipe — GLM-5.1 DSA needs block=1 + AITER MLA on), while a genuine +# submit-time `-e VAR=...` still wins. Precedence: image-baked < models.yaml < submit -e. +# Captured HERE (before the slurm sets any defaults) so it reflects user intent only. +_RECIPE_ENV_KEYS="VLLM_USE_V1 DECODE_MORI_MAX_TOTAL_RECV_TOKENS PREFILL_MORI_MAX_TOTAL_RECV_TOKENS DECODE_MORI_MAX_TOKENS_PER_RANK PREFILL_MORI_MAX_TOKENS_PER_RANK VLLM_MORI_MAX_TOKENS_PER_RANK VLLM_MORI_WARP_NUM_PER_BLOCK VLLM_MORI_BLOCK_NUM VLLM_MORI_RDMA_BLOCK_NUM VLLM_USE_LAYERNAME VLLM_ROCM_USE_AITER VLLM_ROCM_USE_AITER_RMSNORM VLLM_ROCM_USE_AITER_MLA KV_BLOCK_SIZE KV_CACHE_DTYPE KV_CACHE_MEMORY_BYTES GPU_MEMORY_UTILIZATION VLLM_CUDAGRAPH_MODE PREFILL_CUDAGRAPH_MODE DECODE_CUDAGRAPH_MODE CUDAGRAPH_CAPTURE_SIZES VLLM_ALL2ALL_BACKEND PREFILL_MORI_BACKEND DECODE_MORI_BACKEND MORI_SHMEM_HEAP_SIZE" +MODELS_YAML_PROTECT="" +for _k in $_RECIPE_ENV_KEYS; do + [ -n "${!_k+x}" ] && MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT} ${_k}" +done +export MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT# }" +echo "models.yaml protect-list (submit-time overrides): '${MODELS_YAML_PROTECT}'" + +# ------------------------ +# Print current time in UTC and PST formats +# ------------------------ +echo "=== Job Start Time ===" +echo "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +echo "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" +echo "=======================" +echo "" + +# Define valid model names (must have a models.yaml entry) +VALID_MODELS=( \ + "Llama-3.1-405B-Instruct-FP8-KV" \ + "amd-Llama-3.3-70B-Instruct-FP8-KV" \ + "DeepSeek-V3" \ + "DeepSeek-V3-5layer" \ + "gpt-oss-120b" \ + "DeepSeek-R1" \ + "Qwen3-32B" \ + "Qwen3-30B-A3B" \ + "GLM-5.1-FP8" \ +) + +# Models allowed for CONNECTOR=moriio WIDE_EP=1 (MoRI-EP; legacy RUN_MORI=1) +MORI_EP_VALID_MODELS=( \ + "DeepSeek-V3" \ + "DeepSeek-V3-5layer" \ + "DeepSeek-R1" \ + "GLM-5.1-FP8" \ +) + +# Models allowed for CONNECTOR=rixl WIDE_EP=1 EP_BACKEND=deepep (legacy RUN_DEEPEP=1) +DEEPEP_VALID_MODELS=( \ + "DeepSeek-V3" \ + "DeepSeek-V3-5layer" \ + "DeepSeek-R1" \ +) + +MODEL_NAME="${MODEL_NAME:-None}" + +validate_model_name() { + local is_valid_model=false + + for model in "${VALID_MODELS[@]}"; do + if [[ "$MODEL_NAME" == "$model" ]]; then + is_valid_model=true + break + fi + done + + if ! $is_valid_model; then + printf "Error: Invalid MODEL_NAME: '%s'\nValid models are:\n" "$MODEL_NAME" + for model in "${VALID_MODELS[@]}"; do + printf " - %s\n" "$model" + done + exit 1 + fi + + echo "MODEL_NAME '$MODEL_NAME' is valid." + return 0 +} + +validate_model_name "${MODEL_NAME}" + +model_allows_mori_ep() { + local m="$1" + for x in "${MORI_EP_VALID_MODELS[@]}"; do + [[ "$m" == "$x" ]] && return 0 + done + return 1 +} + +model_allows_deepep() { + local m="$1" + for x in "${DEEPEP_VALID_MODELS[@]}"; do + [[ "$m" == "$x" ]] && return 0 + done + return 1 +} + +# --------------------------------------------------------------------------- +# Axis selection -> single launcher (vllm_disagg.sh). +# Two axes: CONNECTOR={rixl|moriio} x WIDE_EP={0=TP|1=wideEP}; EP_BACKEND only +# when WIDE_EP=1 (rixl->deepep, moriio->mori). Legacy RUN_MORI / RUN_DEEPEP are +# still honored via a back-compat shim. The launcher itself does the final +# CONNECTOR/WIDE_EP/EP_BACKEND validation; here we resolve them + gate the model +# against the existing allowlists. +# --------------------------------------------------------------------------- +RUN_FILE="vllm_disagg.sh" +_run_mori="${RUN_MORI:-0}" +_run_deepep="${RUN_DEEPEP:-0}" + +if [[ "$_run_mori" == "1" && "$_run_deepep" == "1" ]]; then + echo "Error: Both RUN_MORI and RUN_DEEPEP are set to 1. Set only one." >&2 + exit 1 +fi + +# Back-compat: legacy flags map onto the axes when CONNECTOR is not set explicitly. +if [[ -z "${CONNECTOR:-}" ]]; then + if [[ "$_run_mori" == "1" ]]; then + CONNECTOR=moriio; WIDE_EP="${WIDE_EP:-1}"; EP_BACKEND="${EP_BACKEND:-mori}" + elif [[ "$_run_deepep" == "1" ]]; then + CONNECTOR=rixl; WIDE_EP="${WIDE_EP:-1}"; EP_BACKEND="${EP_BACKEND:-deepep}" + else + # Default keeps the historical "no flags" behavior: rixl + TP. + CONNECTOR=rixl; WIDE_EP="${WIDE_EP:-0}" + fi +fi +WIDE_EP="${WIDE_EP:-0}" + +# Models that ONLY run wideEP (DP/EP), never TP: the DeepSeek family is served with +# the MoRI-EP / DeepEP recipe (block=16, MLA off, per-role cudagraph). Running them +# in TP mode is unsupported — the TP argv would double the model's own +# --compilation-config and drop the mandatory +quant_fp8 op. Reject early. +# GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA+DSA) is validated only under MoRI-EP +# wideEP disagg (block=1, AITER sparse MLA on, per-role all2all). The moriio+TP +# ("Stage B") path is untested for DSA, so reject WIDE_EP=0 for it too. +WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) +model_is_wide_ep_only() { + local m="$1" + for x in "${WIDE_EP_ONLY_MODELS[@]}"; do [[ "$m" == "$x" ]] && return 0; done + return 1 +} + +# Model allowlist gate (kept, sglang-style). wideEP modes use the per-backend +# allowlists; TP uses the global VALID_MODELS (already validated above) minus the +# wideEP-only models. +if [[ "$WIDE_EP" == "0" ]]; then + if model_is_wide_ep_only "$MODEL_NAME"; then + echo "Error: MODEL_NAME '$MODEL_NAME' is wideEP-only (set WIDE_EP=1). TP mode is not supported for it." >&2 + printf "wideEP-only models:\n"; for m in "${WIDE_EP_ONLY_MODELS[@]}"; do printf " - %s\n" "$m"; done + exit 1 + fi +elif [[ "$WIDE_EP" == "1" && "${CONNECTOR}" == "moriio" ]]; then + if ! model_allows_mori_ep "$MODEL_NAME"; then + echo "Error: CONNECTOR=moriio WIDE_EP=1 but MODEL_NAME '$MODEL_NAME' is not in MORI_EP_VALID_MODELS" >&2 + printf "MoRI EP allowed models:\n"; for m in "${MORI_EP_VALID_MODELS[@]}"; do printf " - %s\n" "$m"; done + exit 1 + fi +elif [[ "$WIDE_EP" == "1" && "${CONNECTOR}" == "rixl" ]]; then + if ! model_allows_deepep "$MODEL_NAME"; then + echo "Error: CONNECTOR=rixl WIDE_EP=1 (deepep) but MODEL_NAME '$MODEL_NAME' is not in DEEPEP_VALID_MODELS" >&2 + printf "DeepEP allowed models:\n"; for m in "${DEEPEP_VALID_MODELS[@]}"; do printf " - %s\n" "$m"; done + exit 1 + fi +fi + +export CONNECTOR WIDE_EP EP_BACKEND +echo "Launcher: $RUN_FILE (CONNECTOR=${CONNECTOR} WIDE_EP=${WIDE_EP} EP_BACKEND=${EP_BACKEND:-}) for model '$MODEL_NAME'" + +# --------------------------------------------------------------------------- +# Connector platform env: per-connector .env holds the ROCm-7.2.3 +# runtime env that MUST reach the container at PID 1 (e.g. expandable_segments:False +# for GPU-RDMA registration). Source the resolved connector's file and collect its +# KEY=VALUE lines into CONNECTOR_ENV_ARGS as `-e KEY=${KEY:-VALUE}` pairs, so a +# submit-time export of the same name still overrides. Forwarded in the docker run. +# --------------------------------------------------------------------------- +CONNECTOR_ENV_FILE="${SCRIPT_DIR}/connectors/${CONNECTOR}.env" +CONNECTOR_ENV_ARGS="" +if [[ -f "$CONNECTOR_ENV_FILE" ]]; then + echo "Loading connector platform env: $CONNECTOR_ENV_FILE" + while IFS= read -r _line; do + [[ "$_line" =~ ^[[:space:]]*# || -z "${_line// }" ]] && continue + _k="${_line%%=*}"; _v="${_line#*=}" + CONNECTOR_ENV_ARGS+=" -e ${_k}=${!_k:-$_v}" # submit-time export of $_k wins + done < "$CONNECTOR_ENV_FILE" +else + echo "WARN: connector env file not found: $CONNECTOR_ENV_FILE" >&2 +fi + +if [[ -z "${DOCKER_IMAGE_NAME:-}" ]]; then + echo "Error: DOCKER_IMAGE_NAME is not set. Please export DOCKER_IMAGE_NAME before running." >&2 + echo " There is no public prebuilt image. Build your own from the provided Dockerfile:" >&2 + echo " docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile \\" >&2 + echo " -t /vllm-disagg:local . # all connectors (WITH_NIXL=1 default)" >&2 + echo " (add --build-arg WITH_NIXL=0 for a lean MoRI-EP-only image)" >&2 + echo " then: export DOCKER_IMAGE_NAME=/vllm-disagg:local" >&2 + exit 1 +fi +export DOCKER_IMAGE_NAME + +# Set current directory to be REPO directory with all relevant scripts +NIXL_REPO_DIR=$(pwd) +LOG_PATH="${LOG_PATH:-/shared_inference/${USER}/model_blog_logs}" + +xP="${xP:-1}" #-> Number of Prefill Servers +yD="${yD:-1}" #-> Number of Decode Servers + +MODEL_DIR="${MODEL_DIR:-"/shared_inference/models_blog/"}" + + +# ------------------------ +# Model path validation and selection across all nodes +# ------------------------ +echo "Looking for model: $MODEL_NAME" +echo "Checking model availability across all allocated nodes..." + +# Get all allocated nodes +ALL_NODES=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +TOTAL_NODES=$(echo "$ALL_NODES" | wc -l) + +echo "Total allocated nodes: $TOTAL_NODES" +echo "Nodes: $(echo "$ALL_NODES" | tr '\n' ' ')" + +# Function to check model path on all nodes +check_model_path() { + local path=$1 + local check_name=$2 + + echo "Checking $check_name: $path" + + # Run check on all nodes in parallel + srun --nodes=$SLURM_NNODES --ntasks=$SLURM_NNODES /bin/bash -c " + if [ -d '$path' ]; then + echo \"\$(hostname): ✓ Found $path\" + exit 0 + else + echo \"\$(hostname): ✗ Missing $path\" + exit 1 + fi + " + + # Check if all nodes succeeded (exit code 0) + local exit_code=$? + if [ $exit_code -eq 0 ]; then + echo "✓ $check_name available on ALL nodes" + return 0 + else + echo "✗ $check_name NOT available on all nodes" + return 1 + fi +} + +# Check /mnt/m2m_nobackup/models_blog first +MODEL_PATH_1="/mnt/m2m_nobackup/models_blog/$MODEL_NAME" +if check_model_path "$MODEL_PATH_1" "/mnt/m2m_nobackup/models_blog"; then + MODEL_PATH="$MODEL_PATH_1" + echo "" + echo "✓ Selected MODEL_PATH: $MODEL_PATH (available on all nodes)" +# Check /shared-inference/models_blog +elif check_model_path "/shared_inference/models_blog/$MODEL_NAME" "/shared_inference/models_blog"; then + MODEL_PATH="/shared_inference/models_blog/$MODEL_NAME" + echo "" + echo "✓ Selected MODEL_PATH: $MODEL_PATH (available on all nodes)" +elif check_model_path "$MODEL_DIR/$MODEL_NAME" "$MODEL_DIR"; then + MODEL_PATH="$MODEL_DIR/$MODEL_NAME" + echo "" + echo "✓ Selected MODEL_PATH: $MODEL_PATH (available on all nodes)" +else + echo "" + echo "✗ FATAL ERROR: Model '$MODEL_NAME' not found on ALL allocated nodes in either:" + echo " - /mnt/m2m_nobackup/models_blog/$MODEL_NAME" + echo " - /shared_inference/models_blog/$MODEL_NAME" + echo "" + echo "Model must be accessible from all nodes for distributed execution." + echo "Please ensure the model is available on all allocated nodes." + exit 1 +fi + +echo "Final MODEL_PATH: $MODEL_PATH" +echo "" + + +# Calculate NUM_NODES based on xP and yD +NUM_NODES=$((xP + yD)) +echo "Calculated NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD, proxy co-located on prefill master)" + +# DeepEP configuration (only exported when RUN_DEEPEP=1) +if [[ "$_run_deepep" == "1" ]]; then + export PREFILL_DEEPEP_BACKEND="${PREFILL_DEEPEP_BACKEND:-deepep_high_throughput}" + export DECODE_DEEPEP_BACKEND="${DECODE_DEEPEP_BACKEND:-deepep_low_latency}" + export ENABLE_DBO="${ENABLE_DBO:-false}" + export DBO_COMM_SMS="${DBO_COMM_SMS:-}" + export ENABLE_PROFILING="${ENABLE_PROFILING:-false}" + echo "DeepEP config: PREFILL_BACKEND=$PREFILL_DEEPEP_BACKEND DECODE_BACKEND=$DECODE_DEEPEP_BACKEND DBO=$ENABLE_DBO" +fi + +# ------------------------ +# Extract first NUM_NODES from SLURM allocation and update SLURM variables +# ------------------------ +echo "Original SLURM allocation:" +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" +echo "SLURM_NNODES: $SLURM_NNODES" +echo "SLURM_NTASKS: $SLURM_NTASKS" + +# Get the full nodelist and extract first NUM_NODES +FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n $NUM_NODES) +NEW_SLURM_NODELIST=$(echo "$SELECTED_NODES" | paste -sd,) + +# Update SLURM environment variables +export SLURM_NNODES=$NUM_NODES +export SLURM_NTASKS=$NUM_NODES +export SLURM_JOB_NUM_NODES=$NUM_NODES +export SLURM_NPROCS=$NUM_NODES +export SLURM_JOB_NODELIST="$NEW_SLURM_NODELIST" +export SLURM_NODELIST="$NEW_SLURM_NODELIST" + +# Keep other SLURM variables as they were or set defaults +export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)" + +export SLURM_CLUSTER_NAME="${SLURM_CLUSTER_NAME}" +export SLURM_JOB_CPUS_PER_NODE="${SLURM_JOB_CPUS_PER_NODE}" +export SLURM_JOB_PARTITION="${SLURM_JOB_PARTITION}" +export SLURM_JOBID="${SLURM_JOBID:-$SLURM_JOB_ID}" +export SLURM_JOB_QOS="${SLURM_JOB_QOS:-normal}" +export SLURM_JOB_ACCOUNT="${SLURM_JOB_ACCOUNT}" +export SLURM_NTASKS_PER_NODE=1 +export SLURM_SUBMIT_HOST="${SLURM_SUBMIT_HOST}" +export SLURM_JOB_ID="${SLURM_JOB_ID}" +export SLURM_CONF="${SLURM_CONF:-/etc/slurm/slurm.conf}" +export SLURM_JOB_NAME="${SLURM_JOB_NAME:-1p1d_bench-serving}" + +echo "" +echo "Updated SLURM Environment Variables:" +echo "SLURM_JOB_ID: $SLURM_JOB_ID" +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" +echo "SLURM_NNODES: $SLURM_NNODES" +echo "SLURM_NTASKS: $SLURM_NTASKS" +echo "SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE" +echo "SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE" +echo "SLURM_JOB_PARTITION: $SLURM_JOB_PARTITION" +echo "SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES" +echo "SLURM_JOBID: $SLURM_JOBID" +echo "SLURM_JOB_QOS: $SLURM_JOB_QOS" +echo "SLURM_NODELIST: $SLURM_NODELIST" +echo "SLURM_JOB_ACCOUNT: $SLURM_JOB_ACCOUNT" +echo "SLURM_NPROCS: $SLURM_NPROCS" +echo "SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST" +echo "SLURM_CONF: $SLURM_CONF" +echo "SLURM_JOB_NAME: $SLURM_JOB_NAME" +echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE" +#echo "SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR" +echo "SLURM_CLUSTER_NAME: $SLURM_CLUSTER_NAME" +echo "ulimit: $(ulimit -a)" +echo "" +echo "Selected nodes for execution:" +echo "$SELECTED_NODES" +echo "" + +# Node information +USER_NAME=$(whoami) +MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) +# Pick the routable fabric IP, not just hostname -I's first entry. These nodes expose +# multiple NICs (e.g. a 10.224.x overlay listed BEFORE the routable 10.158.x fabric); +# taking $1 blindly can advertise an unreachable addr -> prefill/decode barrier hangs +# "Waiting for nodes" forever. Prefer FABRIC_SUBNET (default 10.158.), fall back to $1. +FABRIC_SUBNET="${FABRIC_SUBNET:-10.158.}" +# From a "hostname -I" line, return the first IP on FABRIC_SUBNET, else the first IP. +_pick_fabric_ip() { + awk -v pfx="$FABRIC_SUBNET" '{f=$1; for(i=1;i<=NF;i++) if(index($i,pfx)==1){f=$i; break} print f}' +} +MASTER_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'hostname -I' | _pick_fabric_ip) +MASTER_PORT=39566 # Choose an open port + +IPS=() + +for NODE in $SELECTED_NODES; do + IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | _pick_fabric_ip) + IPS+=("$IP") +done + +echo "Selected node IPs: ${IPS[*]}" | sed 's/ /,/g' + +NIXL_COOKBOOK_PATH="/opt/nixl-vllm-cookbook" +BENCHMARK_ITR="${BENCHMARK_ITR:-1}" +BENCHMARK_CON="${BENCHMARK_CON:-}" +BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" + +# Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. +# sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) +# long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# keepalive -> keepalive_bench.sh (hold server up KEEPALIVE_MINS +# for external accuracy probes) +BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" +case "$BENCHMARK_SCRIPT" in + sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; + long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; + keepalive) BENCHMARK_SCRIPT_FILE="keepalive_bench.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive)" >&2; exit 1 ;; +esac +if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then + echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 + exit 1 +fi +echo "BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' -> $BENCHMARK_SCRIPT_FILE" + +NNODES=$NUM_NODES + +echo "MASTER_NODE is ${MASTER_NODE}" +echo "MASTER_ADDR is ${MASTER_ADDR}" +echo "MASTER_PORT is ${MASTER_PORT}" +echo "NNODES is ${NNODES}" +echo "REPO Directory is ${NIXL_REPO_DIR}" + +if [ ! -d "$LOG_PATH" ]; then + mkdir -p "$LOG_PATH" + echo "Created directory: $LOG_PATH" +else + echo "Directory already exists: $LOG_PATH" +fi + +export CONNECTOR_ENV_ARGS="$CONNECTOR_ENV_ARGS" +export LOG_PATH=$LOG_PATH +export NIXL_REPO_DIR=$NIXL_REPO_DIR +export NIXL_COOKBOOK_PATH=$NIXL_COOKBOOK_PATH +export NNODES=$NNODES +export MASTER_ADDR=$MASTER_ADDR +export MASTER_PORT=$MASTER_PORT +export MODEL_PATH=$MODEL_PATH +export xP=$xP +export yD=$yD +export MODEL_NAME=$MODEL_NAME +export USER_NAME=$USER_NAME +export IPADDRS="$(echo "${IPS[*]}" | sed 's/ /,/g')" +export BENCHMARK_ITR=$BENCHMARK_ITR +export BENCHMARK_CON="${BENCHMARK_CON}" +export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS}" +export BENCHMARK_SCRIPT_FILE="${BENCHMARK_SCRIPT_FILE}" + +export DOCKER_CONT_NAME="container_${MODEL_NAME}_${SLURM_JOB_ID}" +export RUN_FILE_FULL="$NIXL_COOKBOOK_PATH/${RUN_FILE}" + +# Use only the selected nodes for srun execution +SELECTED_NODELIST_SRUN=$(echo "$SELECTED_NODES" | paste -sd,) + +srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c ' +echo "Rank $SLURM_PROCID on $(hostname)"; +docker ps -q | xargs --no-run-if-empty docker stop; +docker rm -f $DOCKER_CONT_NAME 2>/dev/null || true; +fuser -k 5000/tcp 2>/dev/null || true; +fuser -k 2222/tcp 2>/dev/null || true; +fuser -k 15000/tcp 2>/dev/null || true; +sleep 2; +docker pull $DOCKER_IMAGE_NAME 2>/dev/null || true; + +# --- Create host-local compilation cache dirs (ext4, survives container restarts) --- +mkdir -p /tmp/vllm_cache/{aiter_jit,triton,vllm,comgr} 2>/dev/null || true; + +# --- Persistent JIT cache mount --- +# The image points AITER_JIT_DIR/TRITON_CACHE_DIR/VLLM_CACHE_ROOT/COMGR_CACHE_DIR at +# /opt/vllm_cache. Mount a host dir there so AITER CK kernels compile ONCE and are reused +# across runs. Cold compile is ~15 min for the DeepSeek MoE/FP8 GEMM kernel set; a warm +# boot is ~1 min. Keyed by image ID so a new image starts a fresh cache and does not reuse +# stale-ABI shared objects. Host dir on local NVMe. Override JIT_CACHE_HOST, or set +# JIT_CACHE_PERSIST=0 to disable and fall back to the image empty in-container cache. +# NOTE: this whole section runs inside a single-quoted `srun bash -c '...'`, so avoid +# single quotes here; the image-id hash is extracted with tr, not sed. +if [ "${JIT_CACHE_PERSIST:-1}" = "1" ]; then + _IMG_RAW=$(docker image inspect --format "{{.Id}}" "$DOCKER_IMAGE_NAME" 2>/dev/null); + _IMG_KEY=$(printf "%s" "$_IMG_RAW" | tr -cd "a-f0-9" | cut -c1-12); + _IMG_KEY="${_IMG_KEY:-noimg}"; + _JIT_CACHE_HOST="${JIT_CACHE_HOST:-/mnt/m2m_nobackup/${USER}/vllm_jit_cache/${_IMG_KEY}}"; + mkdir -p "$_JIT_CACHE_HOST"/{aiter_jit,triton,vllm,comgr} 2>/dev/null || true; + _JIT_CACHE_MOUNT="-v ${_JIT_CACHE_HOST}:/opt/vllm_cache"; + echo "[jit-cache] persistent image ${_IMG_KEY}: ${_JIT_CACHE_HOST} to /opt/vllm_cache"; +else + _JIT_CACHE_MOUNT=""; +fi + +# --- Build host RDMA library mounts --- +_RDMA_MOUNTS="" +_LIBDIR=/usr/lib/x86_64-linux-gnu + +for _lib in libibverbs.so libibverbs.so.1 librdmacm.so librdmacm.so.1; do + [ -e "$_LIBDIR/$_lib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_LIBDIR/$_lib:$_LIBDIR/$_lib:ro" +done +for _vlib in $_LIBDIR/libibverbs.so.1.* $_LIBDIR/librdmacm.so.1.*; do + [ -e "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" +done + +for _pattern in libmlx5.so* libionic*.so* libbnxt_re*.so* libefa.so* libhns.so*; do + for _vlib in $_LIBDIR/${_pattern}; do + # Require a regular file AFTER symlink resolution: these mounts are built on + # ONE node but applied on ALL nodes, and vendor NIC libs (e.g. libionic.so.1) + # can be a DANGLING symlink on some nodes -> bind-mount fails "not a directory" + # -> container create exit 125. `-f` (follows symlink, requires regular file) + # skips those; the fabric in use (mlx5) is still mounted where present. + [ -f "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" + done +done + +[ -d "$_LIBDIR/libibverbs" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_LIBDIR/libibverbs:$_LIBDIR/libibverbs:ro" +[ -d /etc/libibverbs.d ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v /etc/libibverbs.d:/etc/libibverbs.d:ro" +echo "[host-rdma] mounts: $_RDMA_MOUNTS" + +docker run --rm \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --network host \ + --ipc host \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v $HOME:$HOME \ + -v /shared_inference:/shared_inference \ + -v /mnt/m2m_nobackup:/mnt/m2m_nobackup \ + -v $HOME/.ssh:/root/.ssh \ + --shm-size ${DOCKER_SHM_SIZE:-256G} \ + --ulimit nofile=524288:524288 \ + --ulimit memlock=-1:-1 \ + -v ${LOG_PATH}:/run_logs \ + -v $NIXL_REPO_DIR:$NIXL_COOKBOOK_PATH \ + -v /tmp/vllm_cache:/tmp/vllm_cache \ + ${_JIT_CACHE_MOUNT} \ + ${GLM_KERNEL_PATCH:+-v ${GLM_KERNEL_PATCH}:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py:ro} \ + ${GLM_BACKEND_PATCH:+-v ${GLM_BACKEND_PATCH}:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py:ro} \ + $_RDMA_MOUNTS \ + --entrypoint /bin/bash \ + -e SLURM_JOB_ID=$SLURM_JOB_ID \ + -e SLURM_JOB_NODELIST=$SLURM_JOB_NODELIST \ + -e NNODES=$NNODES \ + -e NODE_RANK=$SLURM_PROCID \ + -e MASTER_ADDR=$MASTER_ADDR \ + -e MASTER_PORT=$MASTER_PORT \ + -e MODEL_PATH=$MODEL_PATH \ + -e NIXL_COOKBOOK_PATH=$NIXL_COOKBOOK_PATH \ + -e xP=$xP \ + -e yD=$yD \ + -e USER_NAME=$USER_NAME \ + -e MODEL_NAME=$MODEL_NAME \ + -e BENCHMARK_ITR=$BENCHMARK_ITR \ + -e BENCHMARK_CON="${BENCHMARK_CON}" \ + -e BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS}" \ + ${BENCHMARK_PORT:+-e BENCHMARK_PORT=$BENCHMARK_PORT} \ + ${PROXY_TYPE:+-e PROXY_TYPE=$PROXY_TYPE} \ + ${ROUTER_PORT:+-e ROUTER_PORT=$ROUTER_PORT} \ + -e IPADDRS=$IPADDRS \ + ${CONNECTOR:+-e CONNECTOR=$CONNECTOR} \ + ${WIDE_EP:+-e WIDE_EP=$WIDE_EP} \ + ${EP_BACKEND:+-e EP_BACKEND=$EP_BACKEND} \ + ${RUN_MORI:+-e RUN_MORI=$RUN_MORI} \ + ${RUN_DEEPEP:+-e RUN_DEEPEP=$RUN_DEEPEP} \ + ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ + ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ + ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ + -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ + ${GLM_PERSIST_GATE:+-e GLM_PERSIST_GATE=$GLM_PERSIST_GATE} \ + ${GLM_SKIP_PATCHERS:+-e GLM_SKIP_PATCHERS=$GLM_SKIP_PATCHERS} \ + ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ + ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ + ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ + ${ROUTER_BINARY:+-e ROUTER_BINARY=$ROUTER_BINARY} \ + ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ + ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ + ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + ${KEEPALIVE_MINS:+-e KEEPALIVE_MINS=$KEEPALIVE_MINS} \ + ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \ + ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \ + ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ + ${MORI_RDMA_TC:+-e MORI_RDMA_TC=$MORI_RDMA_TC} \ + ${MORI_RDMA_SL:+-e MORI_RDMA_SL=$MORI_RDMA_SL} \ + ${MORI_SHMEM_HEAP_SIZE:+-e MORI_SHMEM_HEAP_SIZE=$MORI_SHMEM_HEAP_SIZE} \ + ${PREFILL_DEEPEP_BACKEND:+-e PREFILL_DEEPEP_BACKEND=$PREFILL_DEEPEP_BACKEND} \ + ${DECODE_DEEPEP_BACKEND:+-e DECODE_DEEPEP_BACKEND=$DECODE_DEEPEP_BACKEND} \ + ${ENABLE_DBO:+-e ENABLE_DBO=$ENABLE_DBO} \ + ${DBO_COMM_SMS:+-e DBO_COMM_SMS=$DBO_COMM_SMS} \ + ${ENABLE_PROFILING:+-e ENABLE_PROFILING=$ENABLE_PROFILING} \ + ${NCCL_IB_HCA:+-e NCCL_IB_HCA=$NCCL_IB_HCA} \ + ${NCCL_IB_GID_INDEX:+-e NCCL_IB_GID_INDEX=$NCCL_IB_GID_INDEX} \ + ${NCCL_NET_GDR_LEVEL:+-e NCCL_NET_GDR_LEVEL=$NCCL_NET_GDR_LEVEL} \ + ${NCCL_CROSS_NIC:+-e NCCL_CROSS_NIC=$NCCL_CROSS_NIC} \ + ${NCCL_SOCKET_IFNAME:+-e NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME} \ + ${GLOO_SOCKET_IFNAME:+-e GLOO_SOCKET_IFNAME=$GLOO_SOCKET_IFNAME} \ + -e MORI_SOCKET_IFNAME=${MORI_SOCKET_IFNAME:-eth0} \ + ${MORI_IB_GID_INDEX:+-e MORI_IB_GID_INDEX=$MORI_IB_GID_INDEX} \ + ${MORI_RDMA_DEVICES:+-e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES} \ + ${MORI_NUM_QP_PER_PE:+-e MORI_NUM_QP_PER_PE=$MORI_NUM_QP_PER_PE} \ + ${VLLM_MORIIO_QP_PER_TRANSFER:+-e VLLM_MORIIO_QP_PER_TRANSFER=$VLLM_MORIIO_QP_PER_TRANSFER} \ + ${VLLM_MORIIO_NUM_WORKERS:+-e VLLM_MORIIO_NUM_WORKERS=$VLLM_MORIIO_NUM_WORKERS} \ + -e GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8} \ + -e GPUS_PER_NODE=${GPUS_PER_NODE:-8} \ + ${GPU_MAX_HW_QUEUES:+-e GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES} \ + ${HIP_FORCE_DEV_KERNARG:+-e HIP_FORCE_DEV_KERNARG=$HIP_FORCE_DEV_KERNARG} \ + ${HSA_NO_SCRATCH_RECLAIM:+-e HSA_NO_SCRATCH_RECLAIM=$HSA_NO_SCRATCH_RECLAIM} \ + ${VLLM_HANDSHAKE_TIMEOUT_MINS:+-e VLLM_HANDSHAKE_TIMEOUT_MINS=$VLLM_HANDSHAKE_TIMEOUT_MINS} \ + ${VLLM_ENGINE_READY_TIMEOUT_S:+-e VLLM_ENGINE_READY_TIMEOUT_S=$VLLM_ENGINE_READY_TIMEOUT_S} \ + ${ROCSHMEM_HEAP_SIZE:+-e ROCSHMEM_HEAP_SIZE=$ROCSHMEM_HEAP_SIZE} \ + ${ROCSHMEM_MAX_NUM_CONTEXTS:+-e ROCSHMEM_MAX_NUM_CONTEXTS=$ROCSHMEM_MAX_NUM_CONTEXTS} \ + ${LOG_WAIT_TIMEOUT_SECONDS:+-e LOG_WAIT_TIMEOUT_SECONDS=$LOG_WAIT_TIMEOUT_SECONDS} \ + ${TRITON_CACHE_DIR:+-e TRITON_CACHE_DIR=$TRITON_CACHE_DIR} \ + ${VLLM_CACHE_ROOT:+-e VLLM_CACHE_ROOT=$VLLM_CACHE_ROOT} \ + ${COMGR_CACHE_DIR:+-e COMGR_CACHE_DIR=$COMGR_CACHE_DIR} \ + ${AITER_JIT_DIR:+-e AITER_JIT_DIR=$AITER_JIT_DIR} \ + -e DISTRIBUTED_TIMEOUT_SECONDS=${DISTRIBUTED_TIMEOUT_SECONDS:-7200} \ + -e VLLM_RPC_TIMEOUT=${VLLM_RPC_TIMEOUT:-300000} \ + -e VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-3600} \ + ${CONNECTOR_ENV_ARGS} \ + ${VLLM_CUDAGRAPH_MODE:+-e VLLM_CUDAGRAPH_MODE=$VLLM_CUDAGRAPH_MODE} \ + ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ + --name $DOCKER_CONT_NAME \ + $DOCKER_IMAGE_NAME -c " + mkdir -p /run_logs/${SLURM_JOB_ID} + $RUN_FILE_FULL 2>&1 | tee /run_logs/${SLURM_JOB_ID}/pd_vllm_bench_NODE${SLURM_PROCID}.log + " +' +srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c 'docker stop $DOCKER_CONT_NAME 2>/dev/null || true; docker rm $DOCKER_CONT_NAME 2>/dev/null || true' + diff --git a/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py b/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py new file mode 100644 index 00000000..d491b838 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Overlay the fixed AITER sampling kernel (ROCm/aiter #3658 + hang cap) into the image. + +DEFECT 2 (the 8k prefill/decode crash): the AITER TopP/TopK sampling kernel +(csrc/cpp_itfs/sampling/sampling.cuh) has two bugs on the released aiter post3 +that this image ships: + + 1. HSA OUT-OF-BOUNDS (ROCm/aiter #3658): SamplingTempStorage::last_valid_id is + never initialized. When a probs row is all-zero / NaN (more likely at long + input, e.g. 8k), the guarded write-back (max_valid != -1) is skipped, and the + fallback `sampled_id = temp_storage.last_valid_id` reads UNINITIALIZED shared + memory -> garbage index -> `probs[row*d + sampled_id]` dereferences OOB and + HSA page-faults ("Memory access fault by GPU node-N"). Deterministic under + CUDA graph (shared-mem residue is stable across replays). This is the silent + worker death at 8k that collapses the disagg DP group -> 503. + Fix: init `last_valid_id = 0` at top of each loop iter + defensive clamp on + the loaded sampled_id before it indexes probs. + + 2. REJECTION-SAMPLING HANG: the bisection `do { ... } while(low < high)` can + spin forever when the [low,high] interval stagnates in float precision on a + degenerate (near-uniform) row -> never-completing HSA signal / hang (the + "sampler hang" that forced the skip-warmup workaround). Fix: cap the loop at + kMaxSamplingRounds=32 (float32 mantissa is exhausted well within 32 rounds, + so healthy distributions always converge via break long before the cap). + +Both fixes land in sampling.cuh. #3658 is MERGED upstream but NOT in the released +aiter post3 (this image). Source: A/B-tested by Shiksha (shikpate); staged fixed +tree at SAMPLING_FIX_DIR. + +METHOD (from Shiksha's validated in-container overlay): copy the whole patched +sampling source dir (.cuh + .py + .jinja) over the container's aiter, then purge +any compiled sampling JIT objects so the kernel recompiles from the fixed source +on next use. + +Idempotent: skips if the fix markers are already present. Model-agnostic at the +kernel level, but invoked from the GLM patch hook. Safe no-op if the staged fix +dir or the target aiter dir is absent. + +Usage: apply_glm_aiter_sampling_oob_fix.py + (vllm_install_dir arg is accepted for hook uniformity but not required; + the aiter dir is resolved via `import aiter`.) +""" +import os +import shutil +import subprocess +import sys + +FIX_DIR = os.environ.get( + "SAMPLING_FIX_DIR", + "/shared_inference/ravgupta/aiter_sampling_fix_3658/sampling_patched", +) +MARKERS = ("last_valid_id = 0", "kMaxSamplingRounds") + + +def _aiter_sampling_dir(): + """Locate the installed aiter sampling source dir (aiter_meta/csrc/...).""" + try: + import aiter # noqa: F401 + except Exception as e: # noqa: BLE001 + print(f"[sampling-fix] aiter not importable ({e}); skipping.") + return None + # The kernel source lives under aiter_meta (sibling of aiter), path is stable. + candidates = [] + try: + import aiter_meta # type: ignore + + candidates.append( + os.path.join(os.path.dirname(aiter_meta.__file__), + "csrc", "cpp_itfs", "sampling") + ) + except Exception: # noqa: BLE001 + pass + # Fallback: search site-packages. + import aiter + sp = os.path.dirname(os.path.dirname(aiter.__file__)) + candidates.append(os.path.join(sp, "aiter_meta", "csrc", "cpp_itfs", "sampling")) + for c in candidates: + if os.path.isdir(c): + return c + print(f"[sampling-fix] could not locate aiter sampling dir (tried {candidates}); skipping.") + return None + + +def main() -> int: + tgt = _aiter_sampling_dir() + if tgt is None: + return 0 # safe no-op + + tgt_cuh = os.path.join(tgt, "sampling.cuh") + if os.path.isfile(tgt_cuh): + cur = open(tgt_cuh, errors="ignore").read() + if all(m in cur for m in MARKERS): + print(f"[sampling-fix] already applied (markers present) in {tgt_cuh}.") + return 0 + + if not os.path.isdir(FIX_DIR): + print(f"[sampling-fix] WARN: staged fix dir {FIX_DIR} not found; leaving image kernel unpatched.", file=sys.stderr) + return 0 + + src_cuh = os.path.join(FIX_DIR, "sampling.cuh") + if not os.path.isfile(src_cuh) or not all(m in open(src_cuh, errors="ignore").read() for m in MARKERS): + print(f"[sampling-fix] WARN: staged {src_cuh} missing/lacks fix markers; skipping.", file=sys.stderr) + return 0 + + # Overlay the whole sampling source dir (.cuh + .py + .jinja), per Shiksha's method. + copied = [] + for fn in os.listdir(FIX_DIR): + s = os.path.join(FIX_DIR, fn) + if os.path.isfile(s): + shutil.copy2(s, os.path.join(tgt, fn)) + copied.append(fn) + print(f"[sampling-fix] overlaid #3658 + hang-cap into {tgt}: {', '.join(sorted(copied))}") + + # Verify. + cur = open(tgt_cuh, errors="ignore").read() + if not all(m in cur for m in MARKERS): + print(f"[sampling-fix] ERROR: markers still absent after overlay in {tgt_cuh}.", file=sys.stderr) + return 1 + + # Purge any compiled sampling JIT objects so the kernel recompiles from source. + purged = 0 + for base in ( + os.path.expanduser("~/.aiter"), "/root/.aiter", "/tmp/aiter", + "/opt/vllm_cache/aiter_jit", os.path.join(os.path.dirname(tgt), "..", "..", "jit"), + ): + if base and os.path.isdir(base): + try: + out = subprocess.run( + ["find", base, "-maxdepth", "6", "-iname", "*sampling_from_probs*"], + capture_output=True, text=True, timeout=60, + ) + for p in out.stdout.split(): + try: + if os.path.isdir(p): + shutil.rmtree(p, ignore_errors=True) + else: + os.remove(p) + purged += 1 + except OSError: + pass + except Exception: # noqa: BLE001 + pass + print(f"[sampling-fix] purged {purged} stale sampling JIT object(s); kernel will recompile from fixed source.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py b/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py new file mode 100755 index 00000000..3eaef8cf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +"""Force-compile the GLM-5.1 DSA sparse-attention indexer Triton kernels at BOOT. + +PROBLEM (root cause of the "first big prompt stalls the whole DP group" hang): + The DSA indexer's Triton kernels are seq-length specialized: + - v1/attention/ops/triton_fp8_mqa_logits.py flips `matrix_instr_nonkdim` + at seq_len<=1024 and launches with grid=[(seq_len,)] (seq_len is a + specialized kernel arg) -> a >1024-row prefill needs a *different* JIT + specialization than a <=1024-row one. + But the boot-time warmup never drives the indexer: + - profile_run() calls _dummy_run(is_profile=True) with force_attention=False + and cudagraph mode NONE -> attn_metadata stays None -> sparse_attn_indexer + takes the `sparse_attn_indexer_fake` path (see the "careful! this will be + None in dummy run" comment in layers/sparse_attn_indexer.py). The real + kernels are never compiled. + - _warmup_and_capture() only sets force_attention=True when the cudagraph + runtime mode is FULL; the DSA indexer builder reports UNIFORM_BATCH, so on + this ROCm/DP build the mixed prefill-decode graphs are PIECEWISE and + force_attention stays False. Even when attention IS forced, capture uses + uniform-decode / small mixed batches -- never a large prefill at + max_num_batched_tokens -- so the >1024 specialization is still absent. + Effect: the first >=8k prompt JIT-compiles the indexer kernel mid-inference on + whichever DP rank happens to receive it. That rank falls out of the DP lockstep + gloo all_reduce (coordinate_batch_across_dp) while it compiles -> the whole DP + group collapses. It is also a general cold-cache robustness hole. + +FIX (surgical, reuses vLLM's OWN metadata construction -- no hand-synthesized +tensors, so zero risk of a bad-input crash at boot): + 1. gpu_model_runner.py: add a `_maybe_warmup_dsa_indexer()` method. It is a + strict NO-OP unless one of the runner's attention backends is (a subclass + of) DeepseekV32IndexerBackend. When present, it runs + `_dummy_run(..., force_attention=True, cudagraph_runtime_mode=NONE)` at TWO + prefill-size regimes -- a small one (<=1024 rows) and a large one + (max_num_batched_tokens, >1024) -- so BOTH Triton specializations compile. + `force_attention=True` makes _dummy_run build a real + DeepseekV32IndexerMetadata via the normal _build_attention_metadata path + (num_prefills>0 because the default dummy batch is multi-token requests and + the indexer decode_threshold is 1), which drives the real + `sparse_attn_indexer` prefill kernels. The whole thing is wrapped in + try/except that only WARNs -- a warmup failure must never crash boot. + 2. gpu_worker.py: call it from compile_or_warm_up_model, right after the + existing warmup loop and before kernel_warmup(). At that point the KV cache + is already allocated (initialize_from_config runs before + compile_or_warm_up_model), which the forced-attention indexer path needs. + +Idempotent + anchor-based (matches the other apply_glm_* patchers): + * Each hunk self-detects if already applied (marker string present) -> no-op. + * Missing anchor -> WARN and skip that hunk (safe across vllm revisions; the + rebase may already warm the indexer natively or have refactored the site). + * Anchor found but the file does not contain the applied marker and the + replace produces no change -> hard error (would silently keep the bug). + * py_compile at the end; hard error if the patched file won't compile. + +Usage: apply_glm_dsa_indexer_warmup_fix.py +""" +import os +import sys + +RUNNER_REL = "v1/worker/gpu_model_runner.py" +WORKER_REL = "v1/worker/gpu_worker.py" + +MARKER = "glm-dsa-indexer-warmup" + +# --- Hunk A: new method inserted immediately before `def capture_model` ------- +# Anchor: the (unique) capture_model definition head in gpu_model_runner.py. +RUNNER_ANCHOR = " def capture_model(self) -> int:\n" + +RUNNER_METHOD = ''' def _maybe_warmup_dsa_indexer(self) -> None: + """Force-compile the DSA sparse-attention indexer Triton kernels at boot. + + NO-OP unless this model actually has a DeepseekV32IndexerBackend (GLM-5.1 + DSA / DeepSeek V3.2). The indexer kernels are seq-length specialized + (triton_fp8_mqa_logits flips matrix_instr_nonkdim at seq_len<=1024 and + launches grid=[(seq_len,)]), and the normal profile/warmup passes never + drive the indexer (attn_metadata is None -> the *_fake path). Without this + the first large prompt JIT-compiles mid-inference and, under DP lockstep, + stalls the whole group. We warm BOTH regimes: a small (<=1024) and a large + (max_num_batched_tokens, >1024) prefill batch, using force_attention=True + so _dummy_run builds a real indexer metadata via the standard path. + """ + # {marker} + try: + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerBackend, + ) + except Exception: # noqa: BLE001 -- backend module absent -> not a DSA build + return + + has_indexer = False + try: + for attn_group in self._attn_group_iterator(): + backend = getattr(attn_group, "backend", None) + if backend is not None and isinstance(backend, type) and issubclass( + backend, DeepseekV32IndexerBackend + ): + has_indexer = True + break + except Exception: # noqa: BLE001 -- iterator shape changed -> stay a no-op + return + if not has_indexer: + return + + # Two prefill-size regimes so both Triton specializations compile. + # Small must be <=1024 rows; large must exceed 1024 (use the real max). + max_tokens = int(self.max_num_tokens) + small = min(512, max_tokens) + sizes = [] + for s in (small, max_tokens): + if s > 0 and s not in sizes: + sizes.append(s) + + logger.info( + "Warming up DSA indexer kernels at prefill sizes %s " + "to avoid mid-inference JIT.", + sizes, + ) + for size in sizes: + try: + self._dummy_run( + size, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + force_attention=True, + skip_eplb=True, + remove_lora=False, + ) + except Exception as e: # noqa: BLE001 -- warmup must NEVER crash boot + logger.warning( + "DSA indexer warmup at size %d failed (%s); the kernel may " + "JIT-compile on first use instead.", + size, + e, + ) + self._sync_device() + +'''.replace("{marker}", MARKER) + +# --- Hunk B: call site in gpu_worker.compile_or_warm_up_model ----------------- +WORKER_ANCHOR = ( + " self.model_runner.maybe_remove_all_loras(" + "self.model_runner.lora_config)\n" + "\n" + " # Warmup and tune the kernels used during model execution before\n" + " # cuda graph capture.\n" + " kernel_warmup(self)\n" +) + +WORKER_REPLACEMENT = ( + " self.model_runner.maybe_remove_all_loras(" + "self.model_runner.lora_config)\n" + "\n" + " # " + MARKER + ": force-compile the DSA sparse-attention indexer\n" + " # Triton kernels now (KV cache is allocated), so a large prompt\n" + " # never JIT-compiles them mid-inference and stalls DP lockstep.\n" + " # No-op unless this model has a DeepseekV32IndexerBackend.\n" + " if hasattr(self.model_runner, \"_maybe_warmup_dsa_indexer\"):\n" + " self.model_runner._maybe_warmup_dsa_indexer()\n" + "\n" + " # Warmup and tune the kernels used during model execution before\n" + " # cuda graph capture.\n" + " kernel_warmup(self)\n" +) + + +def _patch_file(path, tag, anchor, apply_fn, already_marker): + """Return 0 on success/no-op, 1 on hard error.""" + if not os.path.isfile(path): + print(f"[{tag}] {path} not found -- skipping (layout differs).") + return 0 + src = open(path).read() + if already_marker in src: + print(f"[{tag}] already applied ({already_marker} present) in {path} -- no-op.") + return 0 + if anchor not in src: + print( + f"[{tag}] WARN: anchor not found in {path} -- skipping " + "(assuming native warmup / refactor)." + ) + return 0 + new_src = apply_fn(src) + if new_src == src: + print( + f"[{tag}] ERROR: anchor found but patch produced no change in {path}.", + file=sys.stderr, + ) + return 1 + try: + open(path, "w").write(new_src) + except OSError as e: + print(f"[{tag}] ERROR: failed to write patched {path}: {e}", file=sys.stderr) + return 1 + if already_marker not in open(path).read(): + print( + f"[{tag}] ERROR: post-write verification failed in {path}.", + file=sys.stderr, + ) + return 1 + print(f"[{tag}] patched {path} -- 1 hunk.") + return 0 + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + vllm_dir = sys.argv[1] + + runner_path = os.path.join(vllm_dir, RUNNER_REL) + worker_path = os.path.join(vllm_dir, WORKER_REL) + + rc = 0 + + # Hunk A: insert the method before capture_model. + rc |= _patch_file( + runner_path, + "glm-dsa-warmup", + RUNNER_ANCHOR, + lambda s: s.replace(RUNNER_ANCHOR, RUNNER_METHOD + RUNNER_ANCHOR, 1), + MARKER, + ) + + # Hunk B: call it from compile_or_warm_up_model. + rc |= _patch_file( + worker_path, + "glm-dsa-warmup", + WORKER_ANCHOR, + lambda s: s.replace(WORKER_ANCHOR, WORKER_REPLACEMENT, 1), + MARKER, + ) + + if rc: + return 1 + + # py-compile sanity for whichever files exist. + try: + import py_compile + + for p in (runner_path, worker_path): + if os.path.isfile(p): + py_compile.compile(p, doraise=True) + print("[glm-dsa-warmup] py_compile OK") + except Exception as e: # noqa: BLE001 + print( + f"[glm-dsa-warmup] ERROR: patched file fails to compile: {e}", + file=sys.stderr, + ) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py b/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py new file mode 100755 index 00000000..33f50dc4 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Apply the GLM-5.1 DSA sparse-attention invalid-token kernel fix (vllm #45324). + +The DSA indexer kernel `_convert_req_index_to_global_index_kernel` in + vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +maps invalid token slots to 0 instead of -1. With block-size 1 + DSA sparse MLA +that corrupts KV reads and the model emits `!!!` for every prompt. + +Upstream fix: vllm-project/vllm #45324 -- flip the 0 to -1 in the tl.where call: + is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + is_invalid_tok | (~valid_block), -1, base * BLOCK_SIZE + inblock_off + +Design (matches launcher contract -- runs unconditionally for GLM, aborts on real +failure): + * IDEMPOTENT : if already -1, report and exit 0 (no-op). + * SELF-SKIPPING: if the file/anchor is absent (refactored or the rebase already + fixed it differently), report and exit 0 -- do NOT abort, because b10a9f7a may + carry the fix natively. We only fail on the one unambiguous bad state we can + fix and didn't, or on write failure. + * VERIFIES the post-write state. + +Usage: apply_glm_dsa_kernel_fix.py +""" +import os +import re +import sys + +REL = "v1/attention/backends/mla/rocm_aiter_mla_sparse.py" + +# Anchor is the stable right-hand side of the tl.where; the middle operand is the +# 0 (buggy) / -1 (fixed) we toggle. Whitespace-tolerant. +RE_ANY = re.compile( + r"(is_invalid_tok\s*\|\s*\(~valid_block\)\s*,\s*)(-?\d+)(\s*,\s*base\s*\*\s*BLOCK_SIZE\s*\+\s*inblock_off)" +) + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + vllm_dir = sys.argv[1] + path = os.path.join(vllm_dir, REL) + + if not os.path.isfile(path): + # File not present on this build -> nothing we can or should do. The + # rebase may use a different sparse backend layout. Do not block launch. + print(f"[glm-dsa] {REL} not found under {vllm_dir} -- skipping (assuming native/refactored).") + return 0 + + src = open(path).read() + m = RE_ANY.search(src) + if not m: + # Anchor gone (refactored / already fixed differently). Don't block. + print(f"[glm-dsa] invalid-token kernel anchor not found in {path} -- skipping (assuming native fix).") + return 0 + + cur = m.group(2) + if cur == "-1": + print(f"[glm-dsa] already fixed (kernel returns -1) in {path} -- no-op.") + return 0 + if cur != "0": + # Unexpected value -- surface it but don't guess. Treat as needs-attention. + print(f"[glm-dsa] ERROR: unexpected invalid-token return value '{cur}' (expected 0 or -1) in {path}.", + file=sys.stderr) + return 1 + + # cur == "0" : the known bug. Flip to -1. + new_src = src[:m.start(2)] + "-1" + src[m.end(2):] + try: + open(path, "w").write(new_src) + except OSError as e: + print(f"[glm-dsa] ERROR: failed to write patched {path}: {e}", file=sys.stderr) + return 1 + + # Verify. + chk = RE_ANY.search(open(path).read()) + if not chk or chk.group(2) != "-1": + print(f"[glm-dsa] ERROR: post-write verification failed in {path}.", file=sys.stderr) + return 1 + print(f"[glm-dsa] patched: invalid-token kernel now returns -1 (vllm #45324) in {path}.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py new file mode 100755 index 00000000..ac6800bf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Patch the MoRIIO KV connector to handle GLM-5.1 DSA's DUAL KV cache. + +PROBLEM (root cause of the 2P2D "Reaped deferred sends / no finished_sending" stall): + GLM-5.1 (GlmMoeDsaForCausalLM -> deepseek_v2.py) has TWO KV caches per layer: + - main MLA latent KV : MLAAttentionSpec, head_size = kv_lora_rank+rope (~576) + - DSA indexer KV : DeepseekV32IndexerCache, MLAAttentionSpec head_size = index_head_dim (128) + Both are 3D ("use_mla"), but DIFFERENT latent dim -> DIFFERENT per-block byte size. + The MoRIIO connector computes ONE global geometry from `first_kv_cache` and reuses it + for every cache, so the indexer cache is transferred with the main-MLA block size -> + wrong bytes/size -> the RDMA read for that region never reconciles -> completion notify + is never produced -> decode reaps deferred sends after 60s -> request hangs. + +FIX (surgical, per-layer geometry; no behavior change for single-cache MLA/DeepSeek): + 1. register_kv_caches: size each registered region by its OWN tensor (per-cache + region_len), not the global self.block_len. Also fix the local_kv_cache_size + append to use the current cache, not a stale loop var. + 2. _compute_block_transfer_offsets: derive shape from the PER-LAYER tensor + (self.kv_caches[layer_name].shape) instead of the global self.kv_cache_shape, + so transfer_size_byte / strides match that cache. + 3. _read_blocks: compute offsets PER LAYER inside the loop (was computed once from + first_layer and reused for all layers). + +Idempotent + anchor-based: each hunk checks if already applied / anchor present; +missing anchor -> warn-and-skip (so it is safe across connector revisions). A hunk +that finds its OLD anchor but fails to apply is a hard error (would silently keep the bug). + +Usage: apply_glm_dsa_moriio_dualkv_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-dualkv] {REL} not found -- skipping (connector layout differs).") + return 0 + + src = open(path).read() + orig = src + applied = [] + + # --- Hunk 1: per-cache region_len in register_kv_caches --------------------- + h1_old = """ for cache_or_caches in kv_caches.values(): + cache_list = [cache_or_caches] if use_mla else cache_or_caches + for cache in cache_list: + base_addr = cache.data_ptr() + region_len = self.num_blocks * self.block_len + caches_data.append((base_addr, region_len, cache.device.index, "")) + kv_caches_base_addr.append(base_addr)""" + h1_new = """ for cache_or_caches in kv_caches.values(): + cache_list = [cache_or_caches] if use_mla else cache_or_caches + for cache in cache_list: + base_addr = cache.data_ptr() + # DSA dual-KV fix: size each region by its OWN tensor, not the + # global self.block_len (the DSA indexer cache has a different + # latent dim than the main MLA cache). + region_len = cache.nelement() * cache.element_size() + caches_data.append((base_addr, region_len, cache.device.index, "")) + kv_caches_base_addr.append(base_addr)""" + if "region_len = cache.nelement() * cache.element_size()" in src: + applied.append("h1 (already)") + elif h1_old in src: + src = src.replace(h1_old, h1_new, 1) + applied.append("h1") + else: + print("[glm-dualkv] WARN: h1 anchor (region_len loop) not found -- skipping h1.") + + # --- Hunk 1b: local_kv_cache_size uses current kv_cache, not stale `cache` -- + h1b_old = " self.local_kv_cache_size.append(cache.nelement() * cache.element_size())" + h1b_new = " self.local_kv_cache_size.append(kv_cache.nelement() * kv_cache.element_size())" + if h1b_new in src: + applied.append("h1b (already)") + elif h1b_old in src: + src = src.replace(h1b_old, h1b_new, 1) + applied.append("h1b") + else: + print("[glm-dualkv] WARN: h1b anchor (local_kv_cache_size) not found -- skipping h1b.") + + # --- Hunk 2: per-layer shape in _compute_block_transfer_offsets ------------- + h2_old = """ assert self.kv_cache_shape is not None, "KV caches shape not initialized" + is_mla = len(self.kv_cache_shape) == 3 + stride = self.kv_caches[layer_name].stride() + sz = self.kv_caches[layer_name].element_size() + if is_mla: + blknum, blksize, hs = self.kv_cache_shape + hn = 1 + block_stride = stride[0] + else: + _, blknum, blksize, hn, hs = self.kv_cache_shape""" + h2_new = """ # DSA dual-KV fix: use the PER-LAYER tensor shape, not the global + # self.kv_cache_shape (the DSA indexer cache differs from the main MLA). + _layer_shape = tuple(self.kv_caches[layer_name].shape) + assert len(_layer_shape) > 0, "KV caches shape not initialized" + is_mla = len(_layer_shape) == 3 + stride = self.kv_caches[layer_name].stride() + sz = self.kv_caches[layer_name].element_size() + if is_mla: + blknum, blksize, hs = _layer_shape + hn = 1 + block_stride = stride[0] + else: + _, blknum, blksize, hn, hs = _layer_shape""" + if "_layer_shape = tuple(self.kv_caches[layer_name].shape)" in src: + applied.append("h2 (already)") + elif h2_old in src: + src = src.replace(h2_old, h2_new, 1) + applied.append("h2") + else: + print("[glm-dualkv] WARN: h2 anchor (_compute_block_transfer_offsets head) not found -- skipping h2.") + + # --- Hunk 3: per-layer offsets in _read_blocks ----------------------------- + h3_old = """ first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0] + offs = self._compute_block_transfer_offsets( + first_layer, local_block_ids, remote_block_ids, remote_moriio_meta + ) + + for layer_name in self.layer_name_to_local_kv_cache_metadata: + sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( + layer_name + ) + # TODO : apply multi-session batch-read when moriio support it + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + )""" + h3_new = """ # DSA dual-KV fix: compute offsets PER LAYER (the DSA indexer cache has a + # different per-block size than the main MLA cache, so a single offs reused + # across all layers mis-sizes the indexer transfer -> lost completion notify). + for layer_name in self.layer_name_to_local_kv_cache_metadata: + sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( + layer_name + ) + offs = self._compute_block_transfer_offsets( + layer_name, local_block_ids, remote_block_ids, remote_moriio_meta + ) + # TODO : apply multi-session batch-read when moriio support it + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + )""" + if "compute offsets PER LAYER" in src: + applied.append("h3 (already)") + elif h3_old in src: + src = src.replace(h3_old, h3_new, 1) + applied.append("h3") + else: + print("[glm-dualkv] WARN: h3 anchor (_read_blocks first_layer offsets) not found -- skipping h3.") + + if src != orig: + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-dualkv] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + print(f"[glm-dualkv] patched {path} -- hunks: {', '.join(applied)}") + else: + print(f"[glm-dualkv] no changes ({', '.join(applied) or 'nothing applied'}) for {path}") + + # py-compile sanity + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-dualkv] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-dualkv] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py new file mode 100755 index 00000000..9d6d468b --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Fix MoRIIO WRITE-path per-layer offset caching for GLM-5.1 DSA dual KV cache. + +ROOT CAUSE (proven by instrumentation, job 37594): + GLM-5.1 (GlmMoeDsaForCausalLM) registers num_layers=156 KV caches = 78 main MLA + (per-block latent dim 576) + 78 DSA indexer caches (latent dim 132). TWO geometries. + + moriio_engine.py::MoRIIOEngine._prepare_transfer_plan computes the RDMA transfer + offsets ONCE (for whatever layer arrives first) and caches them on + request_info.transfer_offset, then REUSES that single offset/size tuple for ALL 156 + layers. The 78 indexer layers (dim 132) get written with the main-MLA geometry + (dim 576) -> wrong byte size/offset -> those RDMA writes are malformed; the per-layer + write accounting (writes_done) and/or the remote completion never reconciles -> + the producer's send_notify (gated on writes_done >= num_layers) misbehaves and the + decode side never receives a clean completion -> "Reaped deferred sends / no + finished_sending after 60s" -> request hangs. + +FIX (surgical, no dataclass change): + Cache transfer offsets PER LAYER on the request_info via a dynamically-attached dict + ``_transfer_offset_by_layer`` keyed by layer_name, instead of the single + ``transfer_offset`` slot. Each of the 156 layers then transfers with its OWN geometry + (the underlying _compute_block_transfer_offsets already takes layer_name and, with the + companion dualkv patch h2, reads the per-layer tensor shape). + + Single-geometry models (DeepSeek-V3 / Hunyuan, 1 cache/layer) are unaffected: + every layer has identical geometry, so per-layer caching yields the same offsets. + +Idempotent + anchor-based. A missing anchor warns-and-skips; a found OLD anchor that +fails to apply is a hard error (would silently keep the stall). + +Usage: apply_glm_dsa_moriio_engine_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-engine] {REL} not found -- skipping (engine layout differs).") + return 0 + + src = open(path).read() + + old = """ # Compute offsets if not cached + if request_info.transfer_offset is None: + offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + task.local_block_ids, + request_info.block_ids, + remote_moriio_meta, + ) + request_info.transfer_offset = offsets + + # Get session index + layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys()) + sess_idx = layer_names.index(task.layer_name) + + local_off, remote_off, sizes = request_info.transfer_offset""" + + new = """ # DSA dual-KV fix: cache offsets PER LAYER, not once per request. GLM-5.1 has + # two cache geometries (main MLA dim 576 + DSA indexer dim 132); a single + # cached offset reused across all 156 layers mis-sizes the indexer writes and + # the completion never reconciles. Per-layer caching is identical for + # single-geometry models (DeepSeek/Hunyuan). + _off_by_layer = getattr(request_info, "_transfer_offset_by_layer", None) + if _off_by_layer is None: + _off_by_layer = {} + request_info._transfer_offset_by_layer = _off_by_layer + offsets = _off_by_layer.get(task.layer_name) + if offsets is None: + offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + task.local_block_ids, + request_info.block_ids, + remote_moriio_meta, + ) + _off_by_layer[task.layer_name] = offsets + # keep the legacy single-slot populated (first layer) for any external reader + if request_info.transfer_offset is None: + request_info.transfer_offset = offsets + + # Get session index + layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys()) + sess_idx = layer_names.index(task.layer_name) + + local_off, remote_off, sizes = offsets""" + + if "_transfer_offset_by_layer" in src: + print(f"[glm-engine] already patched (_transfer_offset_by_layer present) -- no-op.") + elif old in src: + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-engine] patched per-layer offset caching in {path}") + else: + print(f"[glm-engine] WARN: anchor (_prepare_transfer_plan offset block) not found -- skipping (engine revision differs).") + # Not fatal: without the anchor we can't safely patch; surface clearly. + return 0 + + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-engine] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-engine] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py new file mode 100755 index 00000000..92501ae9 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""Fix the MoRIIO transfer-completion gate for GLM-5.1 DSA dual KV cache. + +ROOT CAUSE (PROVEN, job 37615 instrumentation): + GLM-5.1 registers num_layers=156 KV caches = 78 main MLA (model.layers.N.self_attn.attn) + + 78 DSA indexer caches (model.layers.N.self_attn.indexer.k_cache). BUT only the 78 + main-MLA layers ever go through the KV-connector save_kv_layer hook -> only they write + -> writes_done caps at 78. The indexer caches are registered (counted in num_layers) + but vLLM NEVER calls save_kv_layer for them (the DSA indexer is a separate attention + component / DeepseekV32IndexerBackend that doesn't use the connector save path; decode + recomputes indexer state from the transferred main latent KV). + + The producer completion gate (moriio_engine.py): + request_info.writes_done += 1 + if request_info.writes_done >= self.worker.num_layers: # 156, never reached + send_notify(...) + caps at writes_done=78 < num_layers=156 -> send_notify NEVER fires -> decode never + gets completion -> "Reaped deferred sends / no finished_sending after 60s" -> stall. + +FIX: + Add self.num_transfer_layers = count of caches that actually transfer (exclude + '.indexer.' caches), with a fallback to num_layers (so single-geometry models - + DeepSeek-V3 / Hunyuan, no indexer - are bit-identical). Gate completion on + num_transfer_layers instead of num_layers. num_layers itself is left unchanged + (it is also used by the Llama-4 per-layer block-window loop, which needs all caches). + +Companion to apply_glm_dsa_moriio_engine_fix.py (per-layer offset caching). This gate +fix is the primary unblocker; the offset fix is correctness insurance for the layers +that DO write (all same geometry here, but harmless). + +Idempotent + anchor-based. Patches BOTH files (connector: define the field; engine: +use it). A found-old-anchor that fails is a hard error. + +Usage: apply_glm_dsa_moriio_gate_fix.py +""" +import os +import sys + +CONN_REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +ENG_REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +def patch_connector(path: str) -> int: + src = open(path).read() + old = " self.num_layers = len(self.kv_caches.keys())" + new = """ self.num_layers = len(self.kv_caches.keys()) + # DSA dual-KV fix: the producer completion gate must count only layers that + # actually transfer via save_kv_layer. GLM-5.1 registers 2 caches/layer (main + # MLA + DSA indexer), but only the main-MLA caches go through save_kv_layer; the + # '.indexer.' caches are registered yet never written. Gating on len(kv_caches) + # would never be reached. Exclude indexer caches; fall back to num_layers for + # single-geometry models (DeepSeek/Hunyuan have no indexer -> identical). + self.num_transfer_layers = ( + len([k for k in self.kv_caches.keys() if ".indexer." not in k]) + or self.num_layers + ) + logger.info( + "[moriio] completion gate: num_transfer_layers=%d (num_layers=%d)", + self.num_transfer_layers, self.num_layers, + )""" + if "self.num_transfer_layers" in src: + print(f"[glm-gate] connector already patched -- no-op.") + return 0 + if old not in src: + print(f"[glm-gate] WARN: connector anchor (num_layers=) not found -- skipping.") + return 0 + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-gate] patched connector: defined num_transfer_layers in {path}") + return 0 + + +def patch_engine(path: str) -> int: + src = open(path).read() + old = " if request_info.writes_done >= self.worker.num_layers:" + new = """ if request_info.writes_done >= getattr( + self.worker, "num_transfer_layers", self.worker.num_layers + ):""" + if 'getattr(\n self.worker, "num_transfer_layers"' in src or "num_transfer_layers" in src: + print(f"[glm-gate] engine already patched -- no-op.") + return 0 + if old not in src: + print(f"[glm-gate] WARN: engine anchor (writes_done gate) not found -- skipping.") + return 0 + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-gate] patched engine: gate on num_transfer_layers in {path}") + return 0 + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + base = sys.argv[1] + conn = os.path.join(base, CONN_REL) + eng = os.path.join(base, ENG_REL) + if not os.path.isfile(conn) or not os.path.isfile(eng): + print("[glm-gate] connector/engine not found -- skipping (layout differs).") + return 0 + + # ATOMIC: both halves (connector defines num_transfer_layers, engine gates on + # it) are needed together or not at all. On a restructured image (e.g. mori + # v1.2.1, whose engine replaced the writes_done>=num_layers gate with a sealed + # writes_expected mechanism that already handles hybrid/DSA dual-KV natively), + # the engine anchor is gone. Applying only the connector half would inject a + # dead num_transfer_layers into restructured internals. So if the engine anchor + # is absent, skip BOTH — the native gate already does the right thing. + eng_src = open(eng).read() + eng_gate_present = " if request_info.writes_done >= self.worker.num_layers:" in eng_src + eng_already = "num_transfer_layers" in eng_src + if not eng_gate_present and not eng_already: + print("[glm-gate] engine gate anchor absent (image restructured, e.g. mori " + "v1.2.1 sealed writes_expected) -- skipping BOTH halves (native gate handles DSA).") + return 0 + + rc = patch_connector(conn) or patch_engine(eng) + if rc: + return rc + + try: + import py_compile + py_compile.compile(conn, doraise=True) + py_compile.compile(eng, doraise=True) + print("[glm-gate] py_compile OK (both files)") + except Exception as e: # noqa: BLE001 + print(f"[glm-gate] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py new file mode 100755 index 00000000..68e03beb --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""TEMPORARY instrumentation: does the DSA indexer cache reach the MoRIIO save path? + +Adds logging at two points in moriio_connector.py to settle the dual-KV RCA: + 1. register_kv_caches: log all kv_caches layer names + their shapes + num_layers. + -> shows whether the DeepseekV32IndexerCache is even registered, and its geometry. + 2. _write_blocks_for_req: log each distinct layer_name that actually triggers a write. + -> compare the COUNT/SET of written layers vs num_layers. If indexer layers are in + kv_caches (counted in num_layers) but never written, writes_done can never reach + num_layers -> send_notify never fires -> the stall. + +This is diagnostic only (no behavior change). Remove before any production use. +Idempotent + anchor-safe. + +Usage: apply_glm_dsa_moriio_instrument.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-instr] {REL} not found -- skipping.") + return 0 + + src = open(path).read() + orig = src + + # --- Point 1: log kv_caches inventory at num_layers assignment ------------- + a1 = " self.num_layers = len(self.kv_caches.keys())" + b1 = """ self.num_layers = len(self.kv_caches.keys()) + # [glm-instr] kv-cache inventory (dual-KV diagnosis) + try: + for _ln, _kv in self.kv_caches.items(): + logger.info("[glm-instr][register] layer=%s shape=%s dtype=%s", + _ln, tuple(_kv.shape), _kv.dtype) + logger.info("[glm-instr][register] num_layers=%d total_kv_caches=%d", + self.num_layers, len(self.kv_caches)) + except Exception as _e: # noqa: BLE001 + logger.info("[glm-instr][register] inventory log failed: %s", _e)""" + if "[glm-instr][register]" in src: + pass + elif a1 in src: + src = src.replace(a1, b1, 1) + else: + print("[glm-instr] WARN: register anchor (num_layers=) not found.") + + # --- Point 2: log each written layer in _write_blocks_for_req ------------- + a2 = " def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):" + b2 = (a2 + "\n" + ' # [glm-instr] record which layers actually trigger a KV write\n' + ' try:\n' + ' _seen = getattr(self, "_glm_instr_written_layers", None)\n' + ' if _seen is None:\n' + ' _seen = set(); self._glm_instr_written_layers = _seen\n' + ' if layer_name not in _seen:\n' + ' _seen.add(layer_name)\n' + ' logger.info("[glm-instr][write] NEW layer=%s total_written=%d/%d",\n' + ' layer_name, len(_seen), getattr(self, "num_layers", -1))\n' + ' except Exception as _e: # noqa: BLE001\n' + ' logger.info("[glm-instr][write] log failed: %s", _e)') + if "[glm-instr][write]" in src: + pass + elif a2 in src: + src = src.replace(a2, b2, 1) + else: + print("[glm-instr] WARN: _write_blocks_for_req anchor not found.") + + if src != orig: + open(path, "w").write(src) + print(f"[glm-instr] instrumented {path}") + else: + print(f"[glm-instr] already instrumented / nothing to do for {path}") + + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-instr] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-instr] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py b/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py new file mode 100644 index 00000000..adda36cf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Gate OFF the AITER persistent sparse-MLA kernel for chunked-prefill batches. + +ROOT CAUSE (ROCm/aiter #4076, vLLM #47042 / #47567): + The AITER persistent MLA work-stealing kernel (mla_a8w8_qh16_qseqlen1_gqaratio16_ps, + taken when work_meta_data from get_mla_metadata_v1 is non-None) is NUMERICALLY WRONG + for multi-token (prefill-shaped) batches of qseqlen==1 entries. Pure decode (1 query + token) and fresh single-chunk prefills are correct; the error only appears once a + request becomes a CHUNKED-PREFILL CONTINUATION. The small per-token error COMPOUNDS + through the KV cache across chunked-prefill passes until long-context decode collapses + into repetition/garbage. Failure is gated on CHUNK COUNT, not raw context length + (verified: 22k in 2 chunks = correct, 22k in 3 chunks = garbage). + + On this image (aiter 0.1.16.post3, before the aiter-side kernel fix #3921) GLM-5.1-FP8 + DSA collapses at ~16-18k prompt tokens. The aiter kernel fix is the long-term answer + (AITERKER-132 / aiter #3921); this is the vLLM-side short-term gate (#47567), which + costs ~no perf (decode + single-chunk prefill keep the persistent path). + +FIX (port of vLLM PR #47567, adapted to this image's rocm_aiter_mla_sparse.py::build): + In ROCMAiterMLASparseMetadataBuilder.build(), detect chunked-prefill continuations + (a request with >1 query token this step whose total seq_len exceeds its query_len, + i.e. part of its context was computed in an earlier chunk) and, when ANY request in + the batch is such a continuation: + * skip the get_mla_metadata_v1 persistent-metadata launch, and + * pass work_meta_data=None to the metadata so mla_decode_fwd takes the CORRECT + non-persistent split-KV path. + Decode-only and single-chunk-prefill batches are unchanged (persistent path kept). + + Uses `seg_lengths` (per-request step query lengths, already computed at build() top) + and `common_attn_metadata.seq_lens_cpu[:num_reqs].numpy()` (total seq lens). Both are + present in this image's build(). + +Idempotent + anchor-based + self-skipping. Missing anchor -> warn+skip (safe across +image revisions / if a newer image already carries the aiter kernel fix). A found-old +anchor that fails to apply is a hard error (would silently keep the corruption). + +Usage: apply_glm_dsa_persistent_kernel_gate_fix.py +""" +import os +import sys + +REL = "v1/attention/backends/mla/rocm_aiter_mla_sparse.py" + +# Anchor 1: the persistent-metadata guard. We insert the continuation detection +# just before it and AND it into the condition. +OLD1 = """ if metadata_key != self._prev_metadata_key: + from aiter import get_mla_metadata_v1""" +NEW1 = """ # PERSISTENT-KERNEL GATE (aiter #4076 / vLLM #47567): the persistent + # sparse-MLA work-stealing kernel is numerically wrong for chunked-prefill + # continuation batches; the error compounds and breaks long-context decode. + # Fall back to the correct non-persistent path whenever any request in the + # batch is a chunked-prefill continuation (>1 query token this step AND + # total seq_len > this step's query_len). Decode + single-chunk prefills + # keep the fast persistent path -> no decode-throughput regression. + # Slice to num_reqs and cast to int64 (vLLM #47567 hardening / Rohan138 PR#1) + # so the masks cannot broadcast-mismatch under cudagraph padding. + _step_query_lens = seg_lengths[:num_reqs].astype(np.int64) + _total_seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs].numpy().astype( + np.int64 + ) + _is_chunked_continuation = (_step_query_lens > 1) & ( + _total_seq_lens > _step_query_lens + ) + _use_persistent = not bool(_is_chunked_continuation.any()) + if _use_persistent and metadata_key != self._prev_metadata_key: + from aiter import get_mla_metadata_v1""" + +# Anchor 2: the metadata construction passes the persistent buffer unconditionally. +# Gate it on _use_persistent. +OLD2 = " work_meta_data=self._mla_work_meta_data," +NEW2 = " work_meta_data=(self._mla_work_meta_data if _use_persistent else None)," + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-persist] {REL} not found -- skipping (backend layout differs).") + return 0 + + src = open(path).read() + + if "_is_chunked_continuation" in src or "_use_persistent" in src: + print("[glm-persist] already patched (persistent-kernel gate present) -- no-op.") + return 0 + + # Both anchors must be present to apply safely. + if OLD1 not in src: + print("[glm-persist] WARN: persistent-metadata anchor (metadata_key guard) not " + "found -- skipping (image may already carry the aiter kernel fix, or the " + "backend was refactored).") + return 0 + if OLD2 not in src: + print("[glm-persist] ERROR: found the metadata_key guard but NOT the " + "work_meta_data=self._mla_work_meta_data assignment -- refusing partial " + "patch (would leave persistent kernel active). Aborting.", file=sys.stderr) + return 1 + + src = src.replace(OLD1, NEW1, 1) + src = src.replace(OLD2, NEW2, 1) + + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-persist] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + + # Verify both edits landed. + chk = open(path).read() + if "_use_persistent = not bool(_is_chunked_continuation.any())" not in chk or \ + "if _use_persistent else None" not in chk: + print("[glm-persist] ERROR: post-write verification failed.", file=sys.stderr) + return 1 + + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: # noqa: BLE001 + print(f"[glm-persist] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + + print(f"[glm-persist] patched persistent-kernel gate (aiter #4076 / vLLM #47567) in {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py b/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py new file mode 100644 index 00000000..c09b4acd --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Guard the MoRIIO connector abort path against a None peer_zmq (mori v1.2.1). + +ROOT CAUSE (observed on the router image, job 199144 decode crash): + When a request is ABORTED before its KV-transfer peer handshake completes, + the connector's release path runs: + + moriio_connector.py::_release_write_prefill_blocks + peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) # -> None + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq) # None.split(",") + -> AttributeError: 'NoneType' object has no attribute 'split' + + This only catches ValueError, not the AttributeError from a None peer_zmq, so + the EngineCore dies -> cascades to all decode workers (EngineDeadError) -> decode + is dead. Triggered by any request aborted before the peer handshake (e.g. a + canary/curl that times out during first-token cold JIT). + + The SAME FILE already guards this correctly at the other call site + (request_finished / _should_notify path): `if peer_zmq is not None:` then parse, + else fall back to params. The release path just missed the guard — an + inconsistent-guard bug in the connector. + +FIX (surgical, matches the file's own existing pattern): + In _release_write_prefill_blocks, when the params don't already carry + remote_host/remote_notify_port, guard the peer_zmq lookup: if it is None, log + and return (same graceful bail the existing `except ValueError` already does for + the "missing remote notify address" case). No behavior change when peer_zmq is + valid; single-geometry / non-aborted requests are unaffected. + +Idempotent + anchor-based + self-skipping (no-ops if the anchor is absent/already +guarded, so it is safe across connector revisions and other images). A found-old +anchor that fails to apply is a hard error (would leave the crash). + +Usage: apply_glm_moriio_abort_guard_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + +# The buggy two lines: fetch peer_zmq (may be None) then parse it unguarded. +OLD = """ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)""" + +NEW = """ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) + # Abort-path guard: a request aborted before the KV peer handshake + # has peer_zmq=None; parse_moriio_zmq_address(None) would raise + # AttributeError and kill the EngineCore. Bail gracefully like the + # ValueError case below (matches the guarded call site elsewhere). + if peer_zmq is None: + logger.warning( + "Cannot release WRITE prefill blocks for request %s: " + "no peer zmq address (aborted before peer handshake)", + request_id, + ) + return + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)""" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-abort] {REL} not found -- skipping (connector layout differs).") + return 0 + + src = open(path).read() + if "no peer zmq address (aborted before peer handshake)" in src: + print("[glm-abort] already patched -- no-op.") + return 0 + if OLD not in src: + # Anchor absent: either the release path was refactored or this image + # already guards it. Do not block launch. + print("[glm-abort] release-path anchor not found -- skipping (assuming " + "native guard / refactored).") + return 0 + + src = src.replace(OLD, NEW, 1) + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-abort] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: # noqa: BLE001 + print(f"[glm-abort] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + print(f"[glm-abort] patched _release_write_prefill_blocks None-guard in {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index 0cdd027e..43d81850 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -8,7 +8,15 @@ # NIAH_MODEL model name/tag the server serves (required — the served path) # NIAH_WORDS comma list of context sizes in words (default 2000,8000,20000,35000) # NIAH_MAXTOK max_tokens for the answer (default 2048) +# NIAH_SEEDS comma list of needle-layout seeds (default 0,1,2); summary reports +# mean/min/max across seeds to separate real accuracy from variance # NIAH_TIMEOUT per-request timeout seconds (default 1800) +# NIAH_WARMUP 1 (default) = send one throwaway request per context length BEFORE +# scoring, so the first-hit JIT/kernel-autotune compile happens outside +# the scored/gated window. On a freshly-booted node the first request of +# a shape can take minutes to compile; without warmup that lands on the +# first scored request -> false 0/10 or timeout. Warmup failures are +# tolerated (logged, not fatal). Set 0 to disable. import os, sys, json, random, urllib.request URL = os.environ.get("NIAH_URL", "http://127.0.0.1:30000/v1/chat/completions") @@ -16,6 +24,14 @@ WORDS = [int(x) for x in os.environ.get("NIAH_WORDS", "2000,8000,20000,35000").split(",") if x.strip()] MAXTOK = int(os.environ.get("NIAH_MAXTOK", "2048")) TIMEOUT = float(os.environ.get("NIAH_TIMEOUT", "1800")) +# Needle layout is seeded, so a single run is deterministic (bit-exact repro on the +# same stack). Run multiple seeds to distinguish real accuracy from single-needle +# variance; the summary reports mean/min/max across seeds. Default 0,1,2. +SEEDS = [int(x) for x in os.environ.get("NIAH_SEEDS", "0,1,2").split(",") if x.strip()] +WARMUP = os.environ.get("NIAH_WARMUP", "1") == "1" +# Warmup uses a generous timeout (cold compile of a long-context shape can take minutes) +# and never fails the run — its only job is to trigger compilation before scoring. +WARMUP_TIMEOUT = max(TIMEOUT, 1800.0) FILLER = ( "table chair window bottle pencil garden river mountain coffee planet " @@ -41,27 +57,53 @@ def make_haystack(n_words, seed=0): return " ".join(words) -def run(n_words): +def _request(n_words, seed, max_tokens, timeout): + """POST one NIAH request; return (message_dict, error_str). Exactly one is non-None.""" body = { "model": MODEL, "messages": [ {"role": "system", "content": SYSTEM}, - {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words)}, + {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words, seed)}, ], "temperature": 0.0, - "max_tokens": MAXTOK, + "max_tokens": max_tokens, + # Thinking models (e.g. GLM-5.1) emit chain-of-thought into a separate + # reasoning field and leave `content` empty until the final answer; with a + # small max_tokens the answer never appears in `content` and the score is a + # false 0/10. Disable thinking so the answer lands in `content` directly. + "chat_template_kwargs": {"enable_thinking": False}, } data = json.dumps(body).encode() req = urllib.request.Request(URL, data=data, headers={"Content-Type": "application/json"}) try: - with urllib.request.urlopen(req, timeout=TIMEOUT) as r: - msg = json.loads(r.read())["choices"][0]["message"] + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read())["choices"][0]["message"], None except Exception as e: - print("words=%6d ERROR %s" % (n_words, e), flush=True) + return None, str(e) + + +def warmup(n_words): + """One throwaway request per length so first-hit compile happens off the scored path. + Never fatal: a warmup timeout just means the shape is still compiling; the scored + request will pay whatever remains (bounded by NIAH_TIMEOUT).""" + _, err = _request(n_words, seed=0, max_tokens=8, timeout=WARMUP_TIMEOUT) + status = "ok" if err is None else ("timeout/err: %s" % err) + print("words=%6d [warmup] %s" % (n_words, status), flush=True) + + +def run(n_words, seed=0): + # Sentinel: None = timeout/transport error (NOT a wrong answer); int = score 0..10. + msg, err = _request(n_words, seed, MAXTOK, TIMEOUT) + if err is not None: + print("words=%6d seed=%d TIMEOUT/ERROR %s" % (n_words, seed, err), flush=True) return None - text = ((msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")).lower() + # Score content plus any reasoning field (some servers surface CoT as + # `reasoning` or `reasoning_content`) so a thinking model is never mis-scored. + text = ((msg.get("content") or "") + " " + + (msg.get("reasoning_content") or "") + " " + + (msg.get("reasoning") or "")).lower() found = sorted(a for a in ANIMALS if a in text) - print("words=%6d found=%2d/10 %s" % (n_words, len(found), found), flush=True) + print("words=%6d seed=%d found=%2d/10 %s" % (n_words, seed, len(found), found), flush=True) return len(found) @@ -70,14 +112,29 @@ def main(): print("NIAH_MODEL must be set (the served model path/name)", file=sys.stderr) sys.exit(2) print("=== NIAH retrieval test ===", flush=True) - print("url=%s model=%s sizes=%s" % (URL, MODEL, WORDS), flush=True) - results = {} + print("url=%s model=%s sizes=%s seeds=%s warmup=%s" % (URL, MODEL, WORDS, SEEDS, WARMUP), flush=True) + # Warmup pass: compile every shape once before scoring, so cold JIT never lands on a + # scored/gated request (the common cause of false 0/10 or timeout on a fresh boot). + if WARMUP: + print("=== NIAH warmup (one throwaway request per length) ===", flush=True) + for n in WORDS: + warmup(n) + results = {} # n_words -> list of scores across seeds (None = timeout/error, not a wrong answer) for n in WORDS: - results[n] = run(n) - print("=== NIAH summary ===", flush=True) + results[n] = [run(n, s) for s in SEEDS] + print("=== NIAH summary (mean/min/max across %d seed(s)) ===" % len(SEEDS), flush=True) for n in WORDS: - v = results[n] - print(" words=%6d found=%s/10" % (n, "ERR" if v is None else v), flush=True) + scored = results[n] + vals = [v for v in scored if v is not None] + n_to = sum(1 for v in scored if v is None) # timeouts/errors, excluded from mean + if not vals: + print(" words=%6d NO-RESULT (%d/%d timed out or errored — likely cold compile; " + "raise NIAH_TIMEOUT or keep NIAH_WARMUP=1)" % (n, n_to, len(scored)), flush=True) + continue + mean = sum(vals) / len(vals) + extra = (" [%d timeout/err excluded]" % n_to) if n_to else "" + print(" words=%6d mean=%.1f/10 min=%d max=%d (n=%d)%s" + % (n, mean, min(vals), max(vals), len(vals), extra), flush=True) if __name__ == "__main__": diff --git a/scripts/vllm_dissag/benchmark_niah.sh b/scripts/vllm_dissag/benchmark_niah.sh index ba49a359..d366b5e4 100755 --- a/scripts/vllm_dissag/benchmark_niah.sh +++ b/scripts/vllm_dissag/benchmark_niah.sh @@ -14,15 +14,29 @@ LOG="/run_logs/${SLURM_JOB_ID}/niah_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD} echo "==== NIAH long-context retrieval test ====" echo "port=${BENCHMARK_PORT} model=${MODEL_PATH} sizes=${NIAH_WORDS:-2000,8000,20000,35000}" -# Give the router a moment to be fully ready for chat completions. -sleep 10 +# Wait until the router actually serves before starting (replaces a blind sleep). On a +# fresh boot the router may register a few seconds after the workers report ready; poll +# /v1/models until it answers, up to ~5 min. Non-fatal: fall through if the probe can't +# confirm (the harness's own warmup + timeout still protect the run). +_ready=0 +for _i in $(seq 1 60); do + if curl -s -o /dev/null -w '%{http_code}' --max-time 5 \ + "http://127.0.0.1:${BENCHMARK_PORT}/v1/models" 2>/dev/null | grep -q '^200$'; then + _ready=1; echo "[niah] router ready after ~$((_i*5))s"; break + fi + sleep 5 +done +[ "$_ready" = 1 ] || echo "[niah] WARN: router readiness not confirmed in 300s; proceeding (warmup will absorb)" # The server registers the model under its path (served_model_name = MODEL_PATH). +# NIAH_WARMUP=1 (harness default): first-hit JIT compiles off the scored path so a cold +# boot does not produce false 0/10 or timeouts on the first scored request. NIAH_URL="http://127.0.0.1:${BENCHMARK_PORT}/v1/chat/completions" \ NIAH_MODEL="${MODEL_PATH}" \ NIAH_WORDS="${NIAH_WORDS:-2000,8000,20000,35000}" \ NIAH_MAXTOK="${NIAH_MAXTOK:-2048}" \ NIAH_TIMEOUT="${NIAH_TIMEOUT:-1800}" \ +NIAH_WARMUP="${NIAH_WARMUP:-1}" \ python3 "${DIR}/benchmark_niah.py" 2>&1 | tee -a "${LOG}" echo "NIAH results -> ${LOG}" diff --git a/scripts/vllm_dissag/benchmark_xPyD.sh b/scripts/vllm_dissag/benchmark_xPyD.sh index b8851d24..699c9a4e 100755 --- a/scripts/vllm_dissag/benchmark_xPyD.sh +++ b/scripts/vllm_dissag/benchmark_xPyD.sh @@ -40,6 +40,32 @@ for i in $(seq 1 $BENCHMARK_ITR); do 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" + # Per-shape warmup at the REAL isl/osl, low concurrency. The global warmup above + # is isl=osl=32/con=1, which never exercises this shape's prefill path, its Triton/ + # aiter kernel variants, or the decode cudagraph batch sizes -- so without this the + # FIRST measured cell of each shape absorbs all the residual JIT and reports a + # wildly inflated TPOT (observed 302ms vs ~89ms steady-state). Measured cells must + # start from a warm graph. Skip with SHAPE_WARMUP=0. + if [[ "${SHAPE_WARMUP:-1}" == "1" ]]; then + _w_con="${SHAPE_WARMUP_CON:-4}" + _w_prompts="${SHAPE_WARMUP_PROMPTS:-8}" + echo "[WARMUP] shape isl $isl osl $osl con ${_w_con} prompts ${_w_prompts}" \ + | tee -a ${LOG}_CONCURRENCY.log >/dev/null + timeout "${SHAPE_WARMUP_TIMEOUT:-2400}" vllm bench serve \ + --model $MODEL_PATH \ + --backend vllm \ + --host 127.0.0.1 \ + --port $BENCHMARK_PORT \ + --dataset-name "random" \ + --random-input-len $isl \ + --random-output-len $osl \ + --random-prefix-len 0 \ + --num-prompts ${_w_prompts} \ + --request-rate "inf" \ + --ignore-eos \ + --max-concurrency ${_w_con} \ + 2>&1 | tee -a ${LOG}_SHAPEWARMUP.log >/dev/null + fi for con in $CON; do p_con=$(($con * 2)) if [ "$p_con" -lt 16 ]; then diff --git a/scripts/vllm_dissag/connectors/moriio.env b/scripts/vllm_dissag/connectors/moriio.env index 29ed23ba..780f0214 100644 --- a/scripts/vllm_dissag/connectors/moriio.env +++ b/scripts/vllm_dissag/connectors/moriio.env @@ -23,7 +23,18 @@ MORI_RDMA_TC=41 MORI_RDMA_SL=0 MORI_IO_SL=1 MORI_IB_ENABLE_RELAXED_ORDERING=1 -MORI_IB_GID_INDEX=1 +MORI_IB_GID_INDEX=3 +# RDMA NIC allowlist (MI300 + CX7 / RoCE). Without this MoRI auto-enumerates ALL ibv +# devices incl. the mgmt NICs (mlx5_1=eth0, mlx5_6=eth1 on the 10.158 mgmt net), and +# tries to establish QPs over a non-routable/mgmt fabric -> ibverbs.cpp:189 "Connection +# timed out" at the prefill->decode KV transfer. Restrict to the 8 GPU-RoCE NICs +# (rdma0-7 on the 10.224 fabric) per the dist-inf-cookbook cluster-rdma-env-recommender. +# NCCL/GLOO use eth0 for their (non-RDMA) control sockets. Override per-fabric if needed. +MORI_RDMA_DEVICES=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9 +NCCL_IB_HCA=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9 +NCCL_IB_GID_INDEX=3 +NCCL_SOCKET_IFNAME=eth0 +GLOO_SOCKET_IFNAME=eth0 MORI_NUM_QP_PER_PE=8 VLLM_MORIIO_QP_PER_TRANSFER=2 VLLM_MORIIO_NUM_WORKERS=4 diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..e1c7a51a 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -124,13 +124,96 @@ _moriio_build_kv_transfer_config() { } connector_runtime_patch() { - # No-op: the MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL - # split, DP-rank hash-failsafe) are committed in-source in the vLLM the image is - # built from (see the Dockerfile VLLM_REF). There is no runtime .py patcher — that - # would be a drifting duplicate of fixes that already live upstream in the fork. - # If you ever run an image WITHOUT these fixes baked, use an image that has them - # (rebuild from the pinned VLLM_REF) rather than patching a stock image at runtime. - return 0 + # MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL split, + # DP-rank hash-failsafe) are committed in-source in the vLLM the image is built + # from (Dockerfile VLLM_REF). There is no generic runtime .py patcher for those — + # that would be a drifting duplicate of fixes already upstream in the fork. + # + # EXCEPTION — GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA + DSA sparse attention): + # DSA is a NEW attention family the MoRIIO connector was never built for. It adds + # a 2nd KV cache per layer (indexer) with a different geometry, which the + # single-geometry connector mis-handles -> disagg KV transfer stalls; plus a DSA + # invalid-token kernel bug (#45324) that produces `!!!`. These are model-specific + # code gaps, applied here as idempotent, anchor-based, self-skipping .py patchers + # (they no-op cleanly if the fix is native/refactored on the chosen image). Gated + # on MODEL_NAME so DeepSeek/other models are a pure no-op (byte-identical to before). + # The MoRI version is pinned by the Dockerfile MORI_REF (post-1.2.1 main with the + # large-transfer notify/mapping fixes #424/#436/#432 baked in); if a newer MoRI is + # needed, update MORI_REF and rebuild the image — no runtime library swap here. + [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0 + _glm_dsa_runtime_patch +} + +# GLM-5.1 DSA patchers (see connector_runtime_patch). Ported from MAD-private #338. +# Resolves the vLLM install dir, then applies the 4 required patchers in order, +# aborting on a hard failure (a real failure means GLM emits garbage or stalls, so +# failing at launch is correct). Patchers self-skip (rc 0) when their anchor is +# absent, so an image that already carries or refactored a fix no-ops cleanly. +_glm_dsa_runtime_patch() { + # GLM_SKIP_PATCHERS=1: the serving image already carries the GLM-5.1 DSA fixes + # in-source (e.g. the #47766 stack image built from raviguptaamd/vllm@ + # glm5.1-dsa-wideEP_on_shikpate_06_29_customer). Skip ALL runtime patchers — they + # are redundant, and the persistent-gate/sampling-overlay patchers would actively + # REGRESS a baked image (turn persistent MLA off / overwrite stock aiter kernels). + if [ "${GLM_SKIP_PATCHERS:-0}" = "1" ]; then + echo "[glm] GLM_SKIP_PATCHERS=1: image carries DSA fixes in-source; skipping runtime patchers." + return 0 + fi + local _patch_dir="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)}" + local _vllm_dir + _vllm_dir="$(python3 -c 'import vllm, os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || true)" + if [ -z "${_vllm_dir}" ] || [ ! -d "${_vllm_dir}" ]; then + echo "Error: [glm] cannot locate vLLM install dir for DSA patchers. Aborting." >&2 + exit 1 + fi + echo "[glm] MODEL_NAME=GLM-5.1-FP8: applying DSA runtime patchers against ${_vllm_dir}" + + # Ordered list of REQUIRED patchers (all abort on hard failure). + # GLM_PERSIST_GATE=0 skips the persistent-MLA accuracy gate (debug only: to test + # whether the non-persistent kernel it routes to is what crashes disagg at >=8k). + local _gate_patcher="apply_glm_dsa_persistent_kernel_gate_fix.py" + [ "${GLM_PERSIST_GATE:-1}" = "0" ] && _gate_patcher="" + local _p + for _p in \ + apply_glm_dsa_kernel_fix.py \ + apply_glm_dsa_moriio_dualkv_fix.py \ + apply_glm_dsa_moriio_engine_fix.py \ + apply_glm_dsa_moriio_gate_fix.py \ + apply_glm_moriio_abort_guard_fix.py \ + ${_gate_patcher} \ + apply_glm_aiter_sampling_oob_fix.py; do + local _py="${_patch_dir}/${_p}" + if [ ! -f "${_py}" ]; then + echo "Error: [glm] required patcher ${_py} not found. Aborting." >&2 + exit 1 + fi + echo "[glm] applying ${_p}" + python3 "${_py}" "${_vllm_dir}" 2>&1 || { + echo "Error: [glm] ${_p} failed — GLM-5.1 would emit garbage or stall. Aborting." >&2 + exit 1 + } + done + + # Optional DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1). Force-compiles the DSA + # indexer kernels at boot so they never JIT mid-inference. Opt-in because it drives a + # large (>=8k) prefill forward at boot: on stacks where that forward faults it makes + # the fault DETERMINISTIC at boot (useful for debugging) rather than on first request. + if [ "${GLM_INDEXER_WARMUP:-0}" = "1" ]; then + local _warm="${_patch_dir}/apply_glm_dsa_indexer_warmup_fix.py" + if [ -f "${_warm}" ]; then + echo "[glm] applying DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1)" + python3 "${_warm}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] indexer-warmup patch failed (non-fatal)." + fi + fi + + # Optional diagnostic instrumentation (GLM_INSTRUMENT=1). Non-fatal. + if [ "${GLM_INSTRUMENT:-0}" = "1" ]; then + local _instr="${_patch_dir}/apply_glm_dsa_moriio_instrument.py" + if [ -f "${_instr}" ]; then + echo "[glm] applying instrumentation (GLM_INSTRUMENT=1): apply_glm_dsa_moriio_instrument.py" + python3 "${_instr}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] instrumentation failed (non-fatal)." + fi + fi } # connector_launch_worker [dp_start_rank] @@ -162,12 +245,21 @@ connector_launch_worker() { else _cudagraph_mode="${PREFILL_CUDAGRAPH_MODE:-$_cudagraph_mode}" fi + # v0.27 MLA fix: the *_kv_cache_update op dispatches the STABLE-ABI concat_and_cache_mla + # whose boxed kernel does NOT compose inside the Dynamo-FX-partitioned compiled graph + # -> "RuntimeError: unknown parameter type" on the first real MLA decode (passes boot + + # warmup via the fake path, then crashes). splitting_ops-list membership alone doesn't + # cut the graph there. use_inductor_graph_partition=true moves partitioning to inductor + # codegen (after all passes), splitting at cudagraph_unsafe ops incl. the KV-update so + # it runs as an eager boundary. Toggle via USE_INDUCTOR_GRAPH_PARTITION (default 1). + local _igp_json="" + [[ "${USE_INDUCTOR_GRAPH_PARTITION:-1}" == "1" ]] && _igp_json=',"use_inductor_graph_partition":true' if [[ -n "$_cudagraph_mode" && "$_cudagraph_mode" != "NONE" ]]; then local _capture_sizes="${CUDAGRAPH_CAPTURE_SIZES:-1 2 4 8 16 32 64 128 256}" - exec_args+=(--compilation-config '{"cudagraph_mode":"'"${_cudagraph_mode}"'","custom_ops":["+quant_fp8"]}') + exec_args+=(--compilation-config '{"cudagraph_mode":"'"${_cudagraph_mode}"'","custom_ops":["+quant_fp8"]'"${_igp_json}"'}') exec_args+=(--cudagraph-capture-sizes ${_capture_sizes}) else - exec_args+=(--compilation-config '{"cudagraph_mode":"NONE","custom_ops":["+quant_fp8"]}') + exec_args+=(--compilation-config '{"cudagraph_mode":"NONE","custom_ops":["+quant_fp8"]'"${_igp_json}"'}') fi # Per-model flags from models.yaml (driver-exported; empty if none). @@ -182,6 +274,16 @@ connector_launch_worker() { local _all2all="${PREFILL_MORI_BACKEND}" [[ "$log_prefix" == "decode" ]] && _all2all="${DECODE_MORI_BACKEND}" + # NOTE on MoRI EP buffer width: it is sized from max_num_batched_tokens + # (fused_moe/layer.py -> all2all.py max_num_inp_token_per_rank), so a decode + # instance otherwise runs an 8192-token-wide all2all every step (~302ms vs ~88ms + # TPOT). The fix is per-role `--max-num-batched-tokens` in models.yaml + # (decode.dp), NOT an env knob: mori derives recv capacity from the send width + # (MaxNumTokensToRecvPerRank returns min(ceil(maxTotalRecvTokens/ws), + # maxNumInpTokenPerRank)), so shrinking the width alone under-provisions recv and + # trips a device assert during vLLM's profiling dummy run. See + # skills_vllm_disagg.md for the measurements and the dead ends. + local extra_args=() kv_args=() if [[ "$role" == "master" ]]; then extra_args+=(--api-server-count=${_GPUS_PER_NODE}) @@ -219,7 +321,7 @@ connector_launch_worker() { --all2all-backend "${_all2all}" \ --trust-remote-code \ --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS:-7200}" \ - "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" + "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" "${model_args[@]}" WORKER_PID=0; return 0 fi @@ -242,6 +344,7 @@ connector_launch_worker() { "${exec_args[@]}" \ "${extra_args[@]}" \ "${kv_args[@]}" \ + "${model_args[@]}" \ 2>&1 | tee /run_logs/${SLURM_JOB_ID}/${log_prefix}_NODE${NODE_RANK}.log >/dev/null & WORKER_PID=$! return 0 diff --git a/scripts/vllm_dissag/keepalive_bench.sh b/scripts/vllm_dissag/keepalive_bench.sh new file mode 100755 index 00000000..68d80dcb --- /dev/null +++ b/scripts/vllm_dissag/keepalive_bench.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Keepalive with LIGHT heartbeat traffic: holds the disagg server up AND sends a +# tiny request every ~20s so prefill discovery/ping stays registered (idle sleep +# lets the prefill ZMQ ping die ~2min in). Runs KEEPALIVE_MINS (default 90). +: "${KEEPALIVE_MINS:=90}" +PORT="${BENCHMARK_PORT:-30000}" +MODEL="/mnt/m2m_nobackup/models_blog/GLM-5.1-FP8" +echo "[keepalive] light-traffic hold ${KEEPALIVE_MINS}min on :${PORT}" +_end=$(( $(date +%s) + KEEPALIVE_MINS*60 )) +i=0 +while [ "$(date +%s)" -lt "$_end" ]; do + curl -s -m 30 "http://127.0.0.1:${PORT}/v1/completions" -H "Content-Type: application/json" \ + -d "{\"model\":\"${MODEL}\",\"prompt\":\"hi\",\"max_tokens\":1,\"temperature\":0}" >/dev/null 2>&1 + i=$((i+1)) + [ $((i % 3)) -eq 0 ] && echo "[keepalive] heartbeat $i, $(( (_end-$(date +%s))/60 ))min left" + sleep 20 +done +echo "[keepalive] done" diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..150143e6 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -160,10 +160,17 @@ _deepseek_recipe_env: &deepseek_recipe_env # above. No tp: blocks — TP is unsupported for these models. DeepSeek-V3: env: *deepseek_recipe_env + # NEWER-BASE ADAPTATION (isolated to this model): Shiksha's newer vLLM base added + # TritonMLAMetadataBuilder._reserve_attn_logits_workspace(), which pre-reserves the + # decode split-KV logits workspace at WORST CASE + # (max_num_seqs x q_heads x max_kv_splits(max_model_len) x lse_dim x fp32). With + # DSV3's defaults (max_num_seqs=256, max_model_len~163k) this reserves ~128 GiB and + # OOMs at KV init. Cap max-num-seqs + max-model-len to bound the workspace. This is a + # per-model dp: flag (model_args), fully isolated -- does NOT touch GLM/other recipes. prefill: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" decode: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" DeepSeek-V3-5layer: env: *deepseek_recipe_env @@ -178,3 +185,130 @@ DeepSeek-R1: dp: "" decode: dp: "" + +# ============================ MoE + DSA (wideEP only) ============================ + +# GLM-5.1-FP8 (zai-org/GLM-5.1-FP8, arch GlmMoeDsaForCausalLM): MLA + DeepSeek +# Sparse Attention (DSA). 78 layers (3 dense + 75 MoE), 256 routed experts top-8 +# + 1 shared, FP8 block 128. wideEP-only (see WIDE_EP_ONLY_MODELS in the slurm). +# +# Validated DOCKER_IMAGE_NAME (submit-time, not set here — #171 requires it explicit): +# rocmshared/pytorch-private:vllm-wideep_06_29_2026_Shiksha_dp16_2p2d_mori_v1.2.1_aiter_v0.1.16.post3_nightlybase_mori121 +# On this image the DSA patchers self-adapt: only the invalid-token kernel fix +# (#45324) applies; the MoRIIO dual-KV geometry + completion-gate fixes are NATIVE +# (moriio_layout.py per-layer geometry + engine sealed writes_expected), so those +# patchers cleanly no-op. On the older b10a9f7a image all 4 patchers apply. +# +# GLM differs from DeepSeek in 3 recipe-defining ways (both are MLA MoE, but GLM +# is DSA-sparse): +# - KV_BLOCK_SIZE=1 (DSA sparse indexer REQUIRES block-size 1; DS uses 16) +# - VLLM_ROCM_USE_AITER_MLA=1 (GLM MLA path ON via AITER sparse; DS sets 0) +# - prefill EAGER (NONE), decode PIECEWISE (prefill cudagraph capture deadlocks +# on this stack; decode PIECEWISE captures cleanly on DSA and is a ~3.7x ITL win +# (validated: ~69ms vs ~264ms eager). Global VLLM_CUDAGRAPH_MODE=NONE as the +# safe floor; per-role PREFILL=NONE / DECODE=PIECEWISE override it.) +# Note: block=1 + AITER_MLA=1 are ALREADY the moriio.sh connector defaults — GLM +# keeps them; DeepSeek is the one that overrides them off. Set here explicitly so +# the recipe is self-documenting and robust to connector default changes. +# +# DSA adds a 2nd KV cache per layer (indexer); the MoRIIO connector needs the GLM +# DSA patchers (kernel #45324 + dual-KV geometry + per-layer offset + completion +# gate) applied by connector_runtime_patch in connectors/moriio.sh (gated on this +# MODEL_NAME). Those are code patches, not flags — nothing to add here for them. +# +# dp_flags carry the GLM tool/reasoning parsers (AMD GLM recipe). They are applied +# to BOTH roles by compose() in vllm_disagg.sh and reach `vllm serve` via the +# connector's model_args. block-size / kv-cache-dtype / all2all / cudagraph come +# from the env: recipe above (the connector emits them), so the dp: blocks are +# empty like the DeepSeek family. +# +# LONG-CONTEXT CAVEAT (vLLM #40018): the ROCM_AITER_MLA_SPARSE prefill indexer +# corrupts output for prompts beyond ~16-18k tokens on this image (mori v1.2.1) — +# coherent + correct needle retrieval up to 14k, garbage (repetition collapse, +# unique-ratio ~0.1) at ~18.7k. This is an UPSTREAM kernel bug, not the MAD port. +# TESTED: pinning --max-model-len 32768 does NOT move the threshold (workspace is +# sized max_model_len*40 but the corruption onset is a fixed ~18k token count in +# the gather/logits kernel, not a buffer-scaling artifact). So no config knob +# helps; it needs the complete upstream prefill fix in a newer image. Left at the +# native max_model_len (do not cap — capping gives no accuracy benefit and only +# limits usable context). Serve prompts <~14k for correct output on this image. +GLM-5.1-FP8: + env: + VLLM_USE_V1: "1" + # v0.27: layer_name is wrapped in a torch OpaqueBase (LayerName) and passed through + # the unified_mla_kv_cache_update / unified_mla_attention custom ops. On this ROCm + # torch 2.12 build the opaque-type boxing FAILS -> "RuntimeError: unknown parameter + # type" at torch/_ops.py on the first real MLA decode forward (the fake/compile path + # returns early, so it passes boot+warmup then crashes on first request -> DP gloo + # cascade -> 0 prefill/0 decode). VLLM_USE_LAYERNAME=0 makes LayerNameType=str (the + # pre-2.11 path), so a plain string passes through the op. No image rebuild needed. + VLLM_USE_LAYERNAME: "0" + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + KV_BLOCK_SIZE: "1" + KV_CACHE_DTYPE: "fp8" + GPU_MEMORY_UTILIZATION: "0.80" + VLLM_CUDAGRAPH_MODE: "NONE" + PREFILL_CUDAGRAPH_MODE: "NONE" + # PERF: FULL_AND_PIECEWISE captures a FULL graph for the uniform-decode shapes and + # falls back to PIECEWISE for everything else. Measured on MI300X (1024/64, warm), + # against the previous PIECEWISE setting: + # 1P/1D EP8 TPOT 88.0 -> 41.8 ms (2.11x), outTP 78.8 -> 143.7 tok/s + # 2P/2D EP16 TPOT 94.1 -> 55.4 ms (1.70x), outTP 66.7 -> 113.0 tok/s + # Accuracy is unaffected (NIAH 2k-35k: 1P/1D 51/60, 2P/2D 55/60 - both in the + # PIECEWISE band, no length collapse, memfault=0). The gain shrinks as EP widens + # (2.11x at EP8 -> 1.70x at EP16) because the cross-node all2all inside the captured + # step does not compress; do NOT assume this scales to EP32 without measuring. + # Cost: graph capture ~92-94 s and ~3.0-3.5 GiB (PIECEWISE was ~5 s / 7.2 GiB - fewer + # but larger graphs) and a correspondingly longer boot. Irrelevant for a long-lived + # server, but notable for CI. Set PIECEWISE to fall back. + DECODE_CUDAGRAPH_MODE: "FULL_AND_PIECEWISE" + CUDAGRAPH_CAPTURE_SIZES: "1 2 4 8 16 32 64 128 256" + VLLM_ALL2ALL_BACKEND: "mori_high_throughput" + PREFILL_MORI_BACKEND: "mori_high_throughput" + DECODE_MORI_BACKEND: "mori_low_latency" + # MoRI EP dispatch/combine buffer width. Without this it inherits + # max_num_batched_tokens (8192) -- a chunked-prefill SCHEDULER setting -- so every + # decode step moves an 8192-token-wide buffer per layer x78 layers regardless of the + # real batch. That is a fixed ~300ms/step floor (~320x this model's HBM-bandwidth + # bound). Sizing it for the actual decode batch gives TPOT 302ms -> 88ms (3.4x), + # matching the published EP8 figure. The knob is decode.dp below -- prefill is + # unaffected (it genuinely dispatches wide 8192-token chunked-prefill batches). + MORI_SHMEM_HEAP_SIZE: "17179869184" + # DSA sparse-indexer logits-buffer cap (crash fix). The indexer prefill computes an + # M*N fp32 logits buffer; split_indexer_prefill_chunks only sub-chunks the query dim + # when M*N*4 > this budget. Default 512MB lets an 8192-token prefill build a single + # 268MB buffer + launch the fp8_mqa_logits kernel at grid=(8192,), which HARD-FAULTS + # the worker on gfx942 (silent GPU fault -> DP group collapse -> 503 at >=8k prompts). + # Capping at 64MB forces M-dim sub-chunking (~2k tokens/chunk) so the buffer and the + # kernel launch stay bounded. Root cause: vllm/v1/attention/ops/triton_fp8_mqa_logits.py + # fp8_mqa_logits_gfx942; chunking logic: mla/indexer.py split_indexer_prefill_chunks. + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: "64" + # NCCL heartbeat watchdog: at long context (>~8k) a DP rank's sparse-MLA/MoE + # all2all collective can exceed the default HeartbeatMonitor timeout -> + # ProcessGroupNCCL::HeartbeatMonitor::runLoop() declares the rank dead and + # tears down the whole process group -> prefill EngineCore crashes -> 503. + # (Confirmed root cause of the 8k+ prefill crash; #338 EP-landmine.) Disable + # the monitor-triggered teardown and extend timeouts so long-ctx collectives + # complete instead of being watchdog-killed. + TORCH_NCCL_ENABLE_MONITORING: "0" + TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC: "1800" + TORCH_NCCL_DUMP_ON_TIMEOUT: "0" + TORCH_NCCL_BLOCKING_WAIT: "0" + TORCH_NCCL_ASYNC_ERROR_HANDLING: "1" + NCCL_IB_TIMEOUT: "22" + dp_flags: "--tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice --chat-template-content-format string" + prefill: + dp: "" + decode: + # PERF: the MoRI EP dispatch buffer width is max_num_batched_tokens (via + # FusedMoEConfig.max_num_tokens -> all2all.py max_num_inp_token_per_rank), so the + # decode role otherwise runs an 8192-token-wide all2all every step: ~302ms TPOT. + # mori bounds recv capacity BY the send width (MaxNumTokensToRecvPerRank returns + # min(ceil(maxTotalRecvTokens/ws), maxNumInpTokenPerRank)), so the buffer must still + # cover vLLM's profiling dummy run -- it cannot be shrunk via env alone. Lowering + # this knob on the DECODE role lowers both consistently. + # Value must stay >= typical prompt length: 512 gave 87.9ms TPOT but 13.1s TTFT + # (a 1024-token prompt could not be admitted in one step). 2048 keeps TTFT healthy. + dp: "--max-num-batched-tokens 2048" diff --git a/scripts/vllm_dissag/niah_200k.py b/scripts/vllm_dissag/niah_200k.py new file mode 100755 index 00000000..65273bc1 --- /dev/null +++ b/scripts/vllm_dissag/niah_200k.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Needle-in-a-haystack sweep including long context (validated to 200K tokens). + +Hides 10 animal names at even intervals in a filler-word haystack and asks the model to +list them back. Reports found/10, end-to-end latency, and the server-reported +prompt_tokens per length, and can dump the whole run to JSON. + +Usage: + niah_200k.py [lengths_csv] [out_json] + base_url e.g. http://127.0.0.1:20005 (prefill/serve port, or the router) + lengths_csv comma-separated WORD counts (default: 2k..200k) + out_json optional path to write results + + NIAH_MODEL= override the served model id (default below) + +Lengths are given in WORDS to stay comparable with earlier published runs. On this filler +the GLM tokenizer lands ~1 token/word, so words ~= tokens (the script prints the actual +prompt_tokens so you can check). + +NOTE: the FIRST request after a server boot pays cold Triton JIT and can take >80s with a +prefill instance running eager. Warm the server (or use a generous timeout) before +treating any latency number here as steady-state. +""" +import json, os, sys, time, random, urllib.request + +BASE = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:20005" +LENGTHS = [int(x) for x in (sys.argv[2].split(",") if len(sys.argv) > 2 else + "2000,8000,16000,20000,28000,35000,64000,100000,150000,200000".split(","))] +OUT = sys.argv[3] if len(sys.argv) > 3 else None +URL = BASE.rstrip("/") + "/v1/chat/completions" +MODEL = os.environ.get("NIAH_MODEL", "/mnt/m2m_nobackup/models_blog/GLM-5.1-FP8") + +FILLER = ("table chair window bottle pencil garden river mountain coffee planet " + "engine guitar pillow ticket basket candle market silver button orange").split() +ANIMALS = ["elephant", "giraffe", "kangaroo", "penguin", "dolphin", + "tiger", "rhinoceros", "octopus", "crocodile", "panda"] +SYS = ("You read a word list and pick out the animals. Reply with a single " + "comma-separated list of lowercase animal names. Output nothing else.") + + +def hay(n, seed=0): + rng = random.Random(seed) + w = [rng.choice(FILLER) for _ in range(n)] + step = max(n // (len(ANIMALS) + 1), 1) + for i, a in enumerate(ANIMALS): + w[min((i + 1) * step, len(w) - 1)] = a + return " ".join(w) + + +def run(n, seed=0, timeout=1800): + body = { + "model": MODEL, + "messages": [ + {"role": "system", "content": SYS}, + {"role": "user", "content": "Find the animals in this list:\n\n" + hay(n, seed)}, + ], + "temperature": 0, + "max_tokens": 128, + "chat_template_kwargs": {"enable_thinking": False}, + } + req = urllib.request.Request(URL, data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + t = time.time() + try: + r = json.loads(urllib.request.urlopen(req, timeout=timeout).read()) + m = r["choices"][0]["message"] + txt = ((m.get("content") or "") + " " + (m.get("reasoning") or "")).lower() + found = sorted(a for a in ANIMALS if a in txt) + u = r.get("usage") or {} + rec = {"words": n, "seed": seed, "found": len(found), "latency_s": round(time.time() - t, 1), + "prompt_tokens": u.get("prompt_tokens"), "animals": found} + print("words=%7d tok=%-7s found=%2d/10 (%6.1fs) %s" % ( + n, rec["prompt_tokens"], rec["found"], rec["latency_s"], found), flush=True) + return rec + except Exception as e: + rec = {"words": n, "seed": seed, "found": -1, "latency_s": round(time.time() - t, 1), + "error": str(e)[:200]} + print("words=%7d ERROR (%.1fs) %s" % (n, rec["latency_s"], rec["error"]), flush=True) + return rec + + +results = [run(n) for n in LENGTHS] +if OUT: + with open(OUT, "w") as f: + json.dump(results, f, indent=2) + print("wrote", OUT, flush=True) diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..365883b8 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -68,6 +68,22 @@ for f in "${REQUIRED_FILES[@]}"; do done echo "Running from: $(pwd)" +# ------------------------------------------------------------------------------ +# models.yaml env precedence: capture which recipe knobs the USER explicitly set +# at submit time. The driver (vllm_disagg.sh) uses this to let models.yaml `env:` +# OVERRIDE image-baked ENV defaults (e.g. a DeepSeek-tuned image bakes +# KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0, which would otherwise shadow a +# model's own recipe — GLM-5.1 DSA needs block=1 + AITER MLA on), while a genuine +# submit-time `-e VAR=...` still wins. Precedence: image-baked < models.yaml < submit -e. +# Captured HERE (before the slurm sets any defaults) so it reflects user intent only. +_RECIPE_ENV_KEYS="VLLM_USE_V1 VLLM_USE_LAYERNAME VLLM_ROCM_USE_AITER VLLM_ROCM_USE_AITER_RMSNORM VLLM_ROCM_USE_AITER_MLA KV_BLOCK_SIZE KV_CACHE_DTYPE KV_CACHE_MEMORY_BYTES GPU_MEMORY_UTILIZATION VLLM_CUDAGRAPH_MODE PREFILL_CUDAGRAPH_MODE DECODE_CUDAGRAPH_MODE CUDAGRAPH_CAPTURE_SIZES VLLM_ALL2ALL_BACKEND PREFILL_MORI_BACKEND DECODE_MORI_BACKEND MORI_SHMEM_HEAP_SIZE" +MODELS_YAML_PROTECT="" +for _k in $_RECIPE_ENV_KEYS; do + [ -n "${!_k+x}" ] && MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT} ${_k}" +done +export MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT# }" +echo "models.yaml protect-list (submit-time overrides): '${MODELS_YAML_PROTECT}'" + # ------------------------ # Print current time in UTC and PST formats # ------------------------ @@ -87,6 +103,7 @@ VALID_MODELS=( \ "DeepSeek-R1" \ "Qwen3-32B" \ "Qwen3-30B-A3B" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=moriio WIDE_EP=1 (MoRI-EP; legacy RUN_MORI=1) @@ -94,6 +111,7 @@ MORI_EP_VALID_MODELS=( \ "DeepSeek-V3" \ "DeepSeek-V3-5layer" \ "DeepSeek-R1" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=rixl WIDE_EP=1 EP_BACKEND=deepep (legacy RUN_DEEPEP=1) @@ -179,7 +197,10 @@ WIDE_EP="${WIDE_EP:-0}" # the MoRI-EP / DeepEP recipe (block=16, MLA off, per-role cudagraph). Running them # in TP mode is unsupported — the TP argv would double the model's own # --compilation-config and drop the mandatory +quant_fp8 op. Reject early. -WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) +# GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA+DSA) is validated only under MoRI-EP +# wideEP disagg (block=1, AITER sparse MLA on, per-role all2all). The moriio+TP +# ("Stage B") path is untested for DSA, so reject WIDE_EP=0 for it too. +WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) model_is_wide_ep_only() { local m="$1" for x in "${WIDE_EP_ONLY_MODELS[@]}"; do [[ "$m" == "$x" ]] && return 0; done @@ -405,15 +426,22 @@ echo "" # Node information USER_NAME=$(whoami) MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) -MASTER_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'hostname -I') -MASTER_ADDR=$(echo "$MASTER_ADDR" | awk 'NR==1 {print $1}') +# Pick the routable fabric IP, not just hostname -I's first entry. These nodes expose +# multiple NICs (e.g. a 10.224.x overlay listed BEFORE the routable 10.158.x fabric); +# taking $1 blindly can advertise an unreachable addr -> prefill/decode barrier hangs +# "Waiting for nodes" forever. Prefer FABRIC_SUBNET (default 10.158.), fall back to $1. +FABRIC_SUBNET="${FABRIC_SUBNET:-10.158.}" +# From a "hostname -I" line, return the first IP on FABRIC_SUBNET, else the first IP. +_pick_fabric_ip() { + awk -v pfx="$FABRIC_SUBNET" '{f=$1; for(i=1;i<=NF;i++) if(index($i,pfx)==1){f=$i; break} print f}' +} +MASTER_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'hostname -I' | _pick_fabric_ip) MASTER_PORT=39566 # Choose an open port IPS=() for NODE in $SELECTED_NODES; do - IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I') - IP=$(echo "$IP" | awk 'NR==1 {print $1}') + IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | _pick_fabric_ip) IPS+=("$IP") done @@ -427,11 +455,14 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" # Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. # sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) # long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# keepalive -> keepalive_bench.sh (hold server up KEEPALIVE_MINS +# for external accuracy probes) BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" case "$BENCHMARK_SCRIPT" in sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; - *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context)" >&2; exit 1 ;; + keepalive) BENCHMARK_SCRIPT_FILE="keepalive_bench.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -525,7 +556,12 @@ done for _pattern in libmlx5.so* libionic*.so* libbnxt_re*.so* libefa.so* libhns.so*; do for _vlib in $_LIBDIR/${_pattern}; do - [ -e "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" + # Require a regular file AFTER symlink resolution: these mounts are built on + # ONE node but applied on ALL nodes, and vendor NIC libs (e.g. libionic.so.1) + # can be a DANGLING symlink on some nodes -> bind-mount fails "not a directory" + # -> container create exit 125. `-f` (follows symlink, requires regular file) + # skips those; the fabric in use (mlx5) is still mounted where present. + [ -f "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" done done @@ -554,6 +590,8 @@ docker run --rm \ -v $NIXL_REPO_DIR:$NIXL_COOKBOOK_PATH \ -v /tmp/vllm_cache:/tmp/vllm_cache \ ${_JIT_CACHE_MOUNT} \ + ${GLM_KERNEL_PATCH:+-v ${GLM_KERNEL_PATCH}:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py:ro} \ + ${GLM_BACKEND_PATCH:+-v ${GLM_BACKEND_PATCH}:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py:ro} \ $_RDMA_MOUNTS \ --entrypoint /bin/bash \ -e SLURM_JOB_ID=$SLURM_JOB_ID \ @@ -583,6 +621,9 @@ docker run --rm \ ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ + -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ + ${GLM_PERSIST_GATE:+-e GLM_PERSIST_GATE=$GLM_PERSIST_GATE} \ + ${GLM_SKIP_PATCHERS:+-e GLM_SKIP_PATCHERS=$GLM_SKIP_PATCHERS} \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ @@ -590,6 +631,7 @@ docker run --rm \ ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + ${KEEPALIVE_MINS:+-e KEEPALIVE_MINS=$KEEPALIVE_MINS} \ ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \ ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \ ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ diff --git a/scripts/vllm_dissag/skills_vllm_disagg.md b/scripts/vllm_dissag/skills_vllm_disagg.md new file mode 100644 index 00000000..cb837dc2 --- /dev/null +++ b/scripts/vllm_dissag/skills_vllm_disagg.md @@ -0,0 +1,278 @@ +# skills_vllm_disagg.md + +Hard-won operational knowledge for **vLLM PD-disaggregated WideEP serving on AMD MI300X** +(MoRI-EP all-to-all + MoRI-IO RDMA KV transfer), learned while bringing GLM-5.1-FP8 +(MLA + DeepSeek Sparse Attention) onto vLLM v0.27. + +Everything below is *measured*, not theorised. Where a belief turned out to be wrong, +the wrong belief is kept alongside the correction — those are the expensive lessons. + +--- + +## 1. Benchmarking methodology (read this first — it invalidated three of my conclusions) + +### 1.1 ALWAYS discard the first bench run after a boot +The first real request after startup pays **cold Triton JIT** for the sparse/indexer +kernels (`_indexer_k_quant_and_cache_kernel`, `generate_sparse_seqlen_kernel`, +`_convert_req_index_to_global_index_kernel`). Measured on 1P/1D EP8, identical bench: + +| run | TTFT | TPOT | +|---|---|---| +| 1st after boot (cold) | 13,451 ms | 88.7 ms | +| 2nd (warm) | **906 ms** | 88.0 ms | + +TTFT moved **14.9x**; TPOT barely moved. A single cold run made me invent (and act on) +a false "scheduler admission" theory twice. Discard it, or warm up explicitly. + +### 1.2 Why cold JIT lands on TTFT and not TPOT +With the standard recipe **prefill = eager (`CUDAGraphMode.NONE`)**, decode = PIECEWISE: +- prefill has no graph capture at boot -> its kernels compile lazily on the **first real + request** -> the whole compile cost is inside TTFT. +- decode captured graphs at boot ("Graph capturing finished") -> already warm -> TPOT is + correct even on the cold run. +This asymmetry is diagnostic: *cold-JIT symptoms show up in TTFT only*. + +### 1.3 The harness warmup is not a warmup +`benchmark_xPyD.sh` warms at `isl=32 osl=32 con=1` — which never exercises a 1024/8192/28672 +prefill path, nor the decode cudagraph batch sizes. The first *measured* cell therefore +absorbs residual JIT. Fix applied: a per-shape warmup at the real ISL/OSL before each +shape's cells (writes to a separate `_SHAPEWARMUP.log` so it can't pollute the CSV). + +### 1.4 Client timeouts read as server failures +A 50-80s curl timeout against a cold server returns an **empty body**, which looks exactly +like a crash. Use >=300s on the first request. This produced a false "total failure" +verdict during EP32 debugging. + +### 1.5 Sanity-check against physics before blaming a kernel +GLM-5.1-FP8 activates ~37.7B params/token. At 5.3 TB/s HBM, fp8: +- full model on one rank: **7.1 ms/step** +- 1/8 of experts per rank (EP8): **0.9 ms/step** + +Measured 290-300 ms => ~320x the bound. That immediately rules out "compute" or "bandwidth" +and says *stall / oversized transfer*. Do this arithmetic early; it saves hours. + +--- + +## 2. The big perf trap: `max_num_batched_tokens` sizes the MoRI EP buffer + +### 2.1 The chain +``` +vllm/model_executor/layers/fused_moe/layer.py:349 + max_num_tokens = max_num_batched_tokens # 8192 default (SCHEDULER knob) +vllm/model_executor/layers/fused_moe/all2all_utils.py:181 + max_num_tokens_per_dp_rank = moe.max_num_tokens +vllm/distributed/device_communicators/all2all.py (MoriAll2AllManager) + max_num_inp_token_per_rank = +``` +`max_num_batched_tokens` is a **chunked-prefill scheduler** setting. Using it to size the +EP dispatch/combine buffer means a **decode** instance runs an 8192-token-wide all-to-all +**every step, per layer, x78 layers**, while decoding a handful of tokens. + +### 2.2 Signature of this bug +- fixed per-step cost: TPOT identical at concurrency 1, 4, 8, 16 +- independent of KV length: 292 ms at isl=128, 297 ms at isl=1024 +- batching still scales perfectly (con=1 -> 8 gave 8.3x throughput, TPOT flat) +- orders of magnitude above the bandwidth bound +=> "constant oversized transfer", not compute. + +### 2.3 The fix that works today (no code change) +Lower `--max-num-batched-tokens` **on the decode role only** (`models.yaml` `decode.dp:`): + +| decode mnbt | TPOT | TTFT (warm) | out tok/s | +|---|---|---|---| +| 8192 (default) | 302.5 ms | 2431 ms | 24.9 | +| **2048** | **88.0 ms** | **906 ms** | **78.8** | + +3.4x faster decode, and it also *improved* TTFT and throughput. Prefill keeps 8192 (it +genuinely dispatches wide batches). + +### 2.4 What does NOT work (dead ends — do not repeat) +- **`VLLM_MORI_MAX_TOKENS_PER_RANK` alone** (512 or 2048): device assert at boot + ``` + mori .../dispatch_combine/intranode.hpp:134 + `destTokId < config.MaxNumTokensToRecv() && + "Total recv token overflow: increase maxTotalRecvTokens"' + ``` + because vLLM's **profiling/warmup dummy run deliberately pushes + `max_num_batched_tokens` (8192) tokens** through the model. The EP buffer must survive + that even though steady-state decode never needs it. +- **`max_total_recv_tokens` to decouple recv from send: IMPOSSIBLE in current mori.** + ``` + MaxNumTokensToRecvPerRank(): + if maxTotalRecvTokens > 0: + perRank = ceil(maxTotalRecvTokens / worldSize) + return perRank < maxNumInpTokenPerRank ? perRank : maxNumInpTokenPerRank # min() + return maxNumInpTokenPerRank + ``` + It returns **min(perRank, send_width)** — it can only *lower* recv capacity, never raise + it above the send width. `send=1024, recv=65536` behaves identically to leaving it unset. + Recv capacity is structurally bounded by send width. + +### 2.5 The proper upstream fix (not yet done) +Bound the **profiling/dummy run on a decode instance by `max_num_seqs`** instead of +`max_num_batched_tokens`. Then the EP buffer can be narrow and profiling never exceeds it. +This lives entirely in vLLM. (Alternative: a mori change allowing recv > send.) + +### 2.6 MoRI has no tuning knobs (parity gap vs DeepEP) +DeepEP exposes `VLLM_DEEPEP_BUFFER_SIZE_MB`; MoRI hardcodes `warp_num_per_block`, +`block_num`, `rdma_block_num` and inherits its token width from an unrelated scheduler +default. Added `VLLM_MORI_*` knobs for parity (all defaulting to current values). + +--- + +## 2b. Decode CUDA-graph mode: use FULL_AND_PIECEWISE, not PIECEWISE + +`DECODE_CUDAGRAPH_MODE=FULL_AND_PIECEWISE` captures a FULL graph for the uniform-decode +shapes and falls back to PIECEWISE for everything else. Measured on MI300X, 1024/64, warm, +after the max_num_batched_tokens fix in section 2: + +| | PIECEWISE | FULL_AND_PIECEWISE | gain | +|---|---|---|---| +| 1P/1D EP8 TPOT c8 | 88.0 ms | **41.8 ms** | 2.11x | +| 1P/1D EP8 TPOT c16 | 91.5 ms | **45.8 ms** | 2.00x | +| 2P/2D EP16 TPOT c8 | 94.1 ms | **55.4 ms** | 1.70x | +| 2P/2D EP16 TPOT c16 | 96.2 ms | **60.3 ms** | 1.60x | +| 1P/1D output c8 | 78.8 tok/s | **143.7** | 1.82x | +| 2P/2D output c8 | 66.7 tok/s | **113.0** | 1.69x | + +Accuracy is unaffected: NIAH 2k-35k 51/60 (1P/1D) and 55/60 (2P/2D), both inside the +PIECEWISE band (52-53 / 53), no length collapse, memfault=0. + +Three things to know before adopting it: +- **The gain shrinks as EP widens** (2.11x at EP8 -> 1.70x at EP16). FULL captures the whole + decode step *including* the cross-node all2all, and that collective does not compress. + Do not extrapolate to EP32 without measuring. +- **Capture cost:** ~92-94 s and ~3.0-3.5 GiB, versus ~5 s and 7.2 GiB for PIECEWISE - + fewer but larger graphs. Boot is correspondingly longer. Fine for a long-lived server, + notable for CI. +- **FULL is much less cold-JIT sensitive.** Its warmup run already reads steady-state TPOT + (41.35 ms), where PIECEWISE's first post-boot run showed 13.7 s TTFT. The section-1 rule + (discard the first run) still applies, but the penalty is far smaller. + +Two code paths branch on `CUDAGraphMode.FULL` - `sparse_attn_indexer.py:411` and the MoRIIO +connector's READ-mode barrier. Both are guards that *skip* host-side work under FULL, so +enabling it is safe on this stack (the MoRIIO one is already a no-op in WRITE mode). + +## 3. Accuracy: the DSA sparse-index sentinel landmine + +**Never set the invalid/OOB sentinel to `-1`** in +`_convert_req_index_to_global_index_kernel` +(`vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py`). It must be `0`. + +Why: aiter's `mla_decode_fwd` **dereferences** `paged_kv_indices`, so `-1` becomes +`kv_cache + (-1)*stride` -> page-aligned GPU memory access fault -> worker dies -> gloo +DP all-reduce collapse. `0` is masked out by `paged_kv_indptr`/`last_page_len`. + +**Why it hides:** it only fires at **disaggregated long context**. Decode inherits the +prefill's long `seq_lens` -> `generate_sparse_seqlen` widens `paged_kv_indptr` -> `-1` +padding entries land *inside* a live indptr range and get dereferenced. Short prompts and +non-disagg decode never emit an in-range `-1`, so it passes every quick test. +`HIP_LAUNCH_BLOCKING=1` does **not** help (data-dependent OOB, not an async race). + +--- + +## 4. Debugging playbook — isolate the layer before optimising + +The ladder that found the perf bug, in order of cheapness: + +| Suspect | Test | What it showed | +|---|---|---| +| router / KV transfer | bench **prefill-direct** (`:20005`) vs via router (`:30000`) | 341 ms vs 303 ms TPOT -> router + MoRI-IO exonerated (router *does* add TTFT) | +| JIT warmup | re-run on a 30-min-warm server | identical -> not warmup | +| cudagraph config | grep `cudagraph_mode`, `Graph capturing finished`, GiB captured | config correct | +| KV length | sweep isl 128 / 1024 / 8192 at con=1 | flat -> fixed floor, not KV | +| batching | con=1 vs 8 | TPOT flat, throughput 8.3x -> batching fine | +| all2all backend | force `PREFILL_MORI_BACKEND=mori_low_latency` | still broken -> not an HT-kernel bug | + +**Also:** compare against a *known-good* build. Diffing our stack against the passing +v0.25 blog stack showed MoRI, aiter and the recipe were **byte-identical** — only the base +image and vLLM differed. That single comparison exonerated two whole components. + +--- + +## 5. Operational gotchas + +### 5.1 Three caches, three different rules +| cache | keyed by | bake into image? | +|---|---|---| +| `aiter_jit` | (aiter commit, gfx arch) | **yes** — topology-independent, and it is the cold-boot long pole | +| `vllm` (torch.compile/inductor), `triton` | model cfg + batch sizes + cudagraph mode + **topology** (EP8/16/32) | **no** — a stale graph is a real trap; wipe when changing shapes | +| `comgr` | ROCm code objects | harmless | + +Host persistence: `VLLM_CACHE_PERSIST=1` -> +`/mnt/m2m_nobackup/$USER/vllm_jit_cache/a256` -> `/opt/vllm_cache`. +**Keyed by image ID**, so every new image = full cold rebuild. + +If baking `aiter_jit` into an image: **scrub `lock_module_*`, `.ninja_log`, `/root/.mori`, +`/tmp/mori_jit_*` first**, or a fresh container waits on a baton nobody holds -> boot hang. +Verify: a fresh container must start with **zero** ninja/hipcc/clang processes. + +### 5.2 Measured boot times (8 nodes, model on local NVMe) +| scenario | time | +|---|---| +| cold aiter, 2 nodes | ~25 min | +| cold aiter, 4 nodes (2 cold decode) | ~41 min | +| cold aiter, 8 nodes (4 cold decode) | ~106 min | +| **warm cache, 8 nodes** | **~11 min** | + +Changing `max_num_batched_tokens` invalidates the torch.compile cache -> full recompile. + +### 5.3 The aiter baton lock looks like a hang but isn't +`[aiter] waiting for baton release at /opt/vllm_cache/aiter_jit/build/lock_` — +one worker compiles, the rest block. Diagnose by counting build procs *inside the +container*: 100-160 = actively building; **0 on all nodes** = genuinely wedged. + +### 5.4 Readiness: don't count "Application startup complete" +Only nodes running an API server print it. For 2P/2D and 4P/4D the non-master DP ranks +never will. Judge readiness by: router `All servers healthy` + `Graph capturing finished` ++ `GPU KV cache size` per node. + +### 5.5 Verify teardown on EVERY node +`docker rm -f` can leave a container alive on one node; a stale container then collides +with the new run (observed: 2 containers on one node -> prefill shut down mid-boot). +Always re-check `docker ps -q | wc -l == 0` everywhere before relaunching. + +### 5.6 Per-role env: `models.yaml env:` applies to BOTH roles +Prefill and decode often need **opposite** values (e.g. EP buffer width). The pattern is +`PREFILL_*` / `DECODE_*` keys in `env:`, split inside `connectors/moriio.sh` (mirrors the +existing `PREFILL_MORI_BACKEND` / `DECODE_MORI_BACKEND`). Verify it landed by reading +`/proc//environ` **inside the container** — the value is exported into the +server process, not the container shell, so `docker exec env` shows nothing. + +### 5.7 RDMA fabric +- GID index **3** = RoCEv2 IPv4 (check `show_gids` / `sysfs .../gid_attrs/types`). +- Restrict `MORI_RDMA_DEVICES` / `NCCL_IB_HCA` to the 8 GPU-local NICs; leave the mgmt NICs + out or QPs try to form over a non-routable fabric -> `ibverbs.cpp:189 Connection timed out`. +- NCCL/GLOO control sockets on `eth0` (mgmt); KV data on the RDMA NICs. +- **Verify the fabric before blaming code**: `ping -I rdma0` matrix, then `ib_write_bw` + (healthy pair measured 386 Gb/s). A node can be `alloc` in SLURM with a **dead** fabric — + SLURM does not detect this. One dead node cost a whole 2P/2D campaign. +- `ib_write_bw` needs both endpoints co-alive: run it as a single 2-task srun step, not two + separate sruns (a backgrounded server dies when its srun returns). + +### 5.8 Images +Push to a registry (`docker push`) rather than relying on `docker save | ssh | docker load` +serially — parallel `docker pull` across 7 nodes is far faster and removes the +"image missing on one node" failure that silently stalls the launcher barrier. + +--- + +## 6. Open items (state as of 2026-08-16) + +- **4P/4D EP32 silent output corruption.** Cluster boots healthy (memfault=0), but output is + garbage. MoRI + aiter + recipe are byte-identical to the passing v0.25 stack; only base + + vLLM differ => **v0.27 regression**. Both `mori_high_throughput` and `mori_low_latency` + corrupt at EP32 while both are clean at EP8/EP16 => an EP-width (>16 ranks) issue, not a + kernel-specific one. Next probe: 3P/3D (EP24) to test "any EP>16" vs "exactly 32". +- **Proper upstream fix for §2.5** (bound decode profiling by `max_num_seqs`). +- **Prewarmed image** (§5.1) to kill the cold-boot and first-request-JIT costs. + +--- + +## 7. One-line summary of the two bugs found + +1. **Accuracy:** a `-1` sentinel that aiter dereferences -> GPU fault, but only at disagg + long context. Use `0`. +2. **Perf:** the MoE all-to-all buffer is sized from a chunked-prefill *scheduler* knob, so + decode moves an 8192-wide buffer every step. Lower `--max-num-batched-tokens` on the + decode role: **302 ms -> 88 ms TPOT**. diff --git a/scripts/vllm_dissag/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..da2afb00 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -86,7 +86,12 @@ IFS=',' read -ra IP_ARRAY <<< "${IPADDRS}" echo "Listing NIXL_COOKBOOK_PATH: ${NIXL_COOKBOOK_PATH:-}" [[ -n "${NIXL_COOKBOOK_PATH:-}" ]] && ls "${NIXL_COOKBOOK_PATH}" -host_ip=$(hostname -I | awk '{print $1}') +# Prefer the routable fabric IP (FABRIC_SUBNET, default 10.158.) over hostname -I's +# first entry: nodes with a 10.224.x overlay listed first would bind the socket_barrier +# / advertise host_ip on an unreachable NIC -> prefill<->decode barrier hangs "Waiting +# for nodes". Matches the IPADDRS selection in run_xPyD_models.slurm. Falls back to $1. +FABRIC_SUBNET="${FABRIC_SUBNET:-10.158.}" +host_ip=$(hostname -I | awk -v pfx="$FABRIC_SUBNET" '{f=$1; for(i=1;i<=NF;i++) if(index($i,pfx)==1){f=$i; break} print f}') host_name=$(hostname) # ============================================================================= @@ -160,17 +165,33 @@ MODEL_CONFIG_DECODE="" if [[ -n "$MODEL_NAME" && -f "$MODELS_YAML" ]]; then export MODELS_YAML MODEL_NAME PARALLEL_MODE # 1) Export per-model env: block FIRST (so connector ${VAR:-default} yields to it). - # Only set a var that is NOT already in the environment, so a submit-time - # `docker -e VAR=...` (already exported) WINS over the yaml value. Precedence: - # connector default < models.yaml env: < submit-time -e. + # Precedence: image-baked ENV < models.yaml env: < submit-time -e. + # models.yaml MUST override image-baked ENV: a DeepSeek-tuned disagg image + # bakes KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0 / VLLM_CUDAGRAPH_MODE= + # PIECEWISE etc. as container ENV, which would otherwise shadow a model's own + # recipe (GLM-5.1 DSA needs block=1 + AITER sparse MLA on). But a genuine + # submit-time `-e VAR=...` must still win. The slurm can tell the two apart + # (it runs on the host) and passes MODELS_YAML_PROTECT = the space-separated + # list of keys the USER set at submit; the driver protects only those. When + # MODELS_YAML_PROTECT is unset (script run directly, no slurm), fall back to + # the old "skip if in env" behavior so nothing regresses. _yaml_env="$(python3 - <<'PY' import os, yaml, shlex m = yaml.safe_load(open(os.environ["MODELS_YAML"])) or {} cfg = m.get(os.environ["MODEL_NAME"]) or {} +protect_raw = os.environ.get("MODELS_YAML_PROTECT") +have_protect = protect_raw is not None +protect = set((protect_raw or "").split()) for k, v in (cfg.get("env") or {}).items(): - # skip if already present in the environment (submit-time -e override wins) - if k in os.environ: - continue + if have_protect: + # 3-tier: yaml overrides baked ENV; only a user submit-time -e (in the + # protect-list) wins over yaml. + if k in protect: + continue + else: + # No protect-list (direct run): legacy behavior — any existing env wins. + if k in os.environ: + continue print(f'export {k}={shlex.quote(str(v))}') PY )"