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..1053af1e --- /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 0 = MoRI-EP only). +# +# 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=0 (default) => MoRI-EP only (moriio TP/wideEP + deepep-from-base); lean. +# WITH_NIXL=1 => 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). +# +# 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: rocm/vllm-dev:ci_base-dedbf6be8b1afa17a6220473b9c8c98242ac1c03 +# (ROCm + torch nightly). The stages below 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_vllm-v0.27 +# (upstream v0.27 + GLM DSA + #47766 metadata-key). 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: UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Default is 0 +# (MoRI-EP only). Single guarded RUN so WITH_NIXL=0 skips it (no layers, no cost). +# ----------------------------------------------------------------------------- +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/scripts/vllm_dissag/README.MD b/scripts/vllm_dissag/README.MD index 17660abe..921bc54d 100644 --- a/scripts/vllm_dissag/README.MD +++ b/scripts/vllm_dissag/README.MD @@ -53,6 +53,7 @@ slurm gate rejects any other pairing. Dense models are TP-only; the DeepSeek fam | `DeepSeek-V3` | MoE | 3, 4 (wideEP only) | validated serving (moriio+MoRI-EP) | | `DeepSeek-V3-5layer` | MoE | 3, 4 (wideEP only) | smoke/dev variant | | `DeepSeek-R1` | MoE | 3, 4 (wideEP only) | needs a co-versioned AITER image (see Caveats) | +| `GLM-5.1-FP8` | MoE + DSA | 3 (moriio wideEP only) | own image (`glmv5.1` Dockerfile); 1P/1D EP8 + 2P/2D EP16; EP32 known-broken | ## Prerequisites @@ -234,7 +235,7 @@ lengths and the model is scored on how many it retrieves. Select it in place of sweep via the launcher's benchmark hook: ```bash -export BENCHMARK_SCRIPT_FILE=benchmark_niah.sh +export BENCHMARK_SCRIPT=niah export NIAH_WORDS="2000,8000,20000,35000" # context sizes (words); optional ``` 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.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..ffb3ee14 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -124,13 +124,104 @@ _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.* (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. 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. + case "${MODEL_NAME:-}" in + GLM-5.*) ;; + *) return 0 ;; + esac + _glm_dsa_runtime_patch +} + +# GLM-5.* DSA patchers (see connector_runtime_patch). Ported from MAD-private #338. +# Resolves the vLLM install dir, then applies the dual-KV / abort fallbacks in order, +# aborting on a hard failure. Patchers self-skip (rc 0) when their anchor is absent. +# Two patchers are OPT-IN because they crash the v0.27 GLM image (see below). +_glm_dsa_runtime_patch() { + # GLM_SKIP_PATCHERS=1: the serving image already carries the GLM DSA fixes + # in-source (vLLM #47766 cache-key + dual-KV). Skip ALL runtime patchers — they + # are redundant, and the persistent-gate/sentinel patchers would actively REGRESS + # a baked image (asm_mla.cu abort / hipErrorIllegalAddress). + 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=${MODEL_NAME}: applying DSA runtime patchers against ${_vllm_dir}" + + # apply_glm_dsa_persistent_kernel_gate_fix.py ports vLLM #47567. OFF by default: + # #47766 superseded it, and forcing non-persistent MLA on this image aborts at + # asm_mla.cu:945 (fp8/fp8 gqa_ratio=64 has no non-persistent kernel). Set + # GLM_PERSIST_GATE=1 only on an image that predates #47766. + local _gate_patcher="" + [ "${GLM_PERSIST_GATE:-0}" = "1" ] && _gate_patcher="apply_glm_dsa_persistent_kernel_gate_fix.py" + # apply_glm_dsa_kernel_fix.py ports still-open vLLM #45324 (sentinel 0->-1). OFF + # by default: this image ships 0 deliberately because aiter mla_decode_fwd + # dereferences the index (-1 -> hipErrorIllegalAddress on short disagg decode). + # Set GLM_DSA_SENTINEL_FIX=1 only on an older image that genuinely has the #45324 bug. + local _dsa_sentinel_patcher="" + [ "${GLM_DSA_SENTINEL_FIX:-0}" = "1" ] && _dsa_sentinel_patcher="apply_glm_dsa_kernel_fix.py" + local _p + for _p in \ + ${_dsa_sentinel_patcher} \ + 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 + [ -n "${_p}" ] || continue + 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 — ${MODEL_NAME} 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 +253,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 +282,17 @@ 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. Do not try to + # raise recv via VLLM_MORI_MAX_TOTAL_RECV_TOKENS — that knob is a min() + # clamp and can only LOWER capacity. + local extra_args=() kv_args=() if [[ "$role" == "master" ]]; then extra_args+=(--api-server-count=${_GPUS_PER_NODE}) @@ -219,7 +330,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 +353,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/models.json b/scripts/vllm_dissag/models.json index 6fc3849a..2d9f0ab4 100644 --- a/scripts/vllm_dissag/models.json +++ b/scripts/vllm_dissag/models.json @@ -278,6 +278,38 @@ }, "args": "-N 2 -n 2" }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8", + "dockerfile": "../../docker/vllm_disagg_inference.glmv5.1", + "scripts": "run_xPyD_models.slurm", + "url": "", + "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_deepep_deepseek-v3", "dockerfile": "../../docker/vllm_disagg_inference", diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..4454e452 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -178,3 +178,132 @@ 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:glm5.1-vllm027-b8 +# built from docker/vllm_disagg_inference.glmv5.1 (VLLM_REF=glm5.1-dsa-wideEP_on_vllm-v0.27). +# It carries the DSA fixes in-source, so the card sets GLM_SKIP_PATCHERS=1 and NO runtime +# patcher runs. Do NOT serve this recipe on a pre-v0.27 image: those predate vLLM #47766 +# and return empty/garbage completions here (NIAH 2k scored 0/10 on +# rocm/pytorch-private:glm-dockerimage-built-09072026). The connectors/moriio.sh patchers +# are a fallback for such older images only (GLM_SKIP_PATCHERS unset). +# +# 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/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..6e3982b7 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,13 @@ 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) +# niah -> benchmark_niah.sh (needle-in-a-haystack; NIAH_WORDS) 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 ;; + niah) BENCHMARK_SCRIPT_FILE="benchmark_niah.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, niah)" >&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 +555,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 +589,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 +620,10 @@ 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_DSA_SENTINEL_FIX:+-e GLM_DSA_SENTINEL_FIX=$GLM_DSA_SENTINEL_FIX} \ + ${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,11 @@ 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} \ + ${NIAH_WORDS:+-e NIAH_WORDS=$NIAH_WORDS} \ + ${NIAH_SEEDS:+-e NIAH_SEEDS=$NIAH_SEEDS} \ + ${NIAH_WARMUP:+-e NIAH_WARMUP=$NIAH_WARMUP} \ + ${NIAH_MAXTOK:+-e NIAH_MAXTOK=$NIAH_MAXTOK} \ + ${NIAH_TIMEOUT:+-e NIAH_TIMEOUT=$NIAH_TIMEOUT} \ ${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/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 )"