From 56f83c729868ebce530b452153e8612ec394d99f Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Wed, 8 Jul 2026 06:42:53 +0000 Subject: [PATCH 1/4] vllm_dissag: GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP disaggregated enablement Adds GLM-5.1-FP8 (GlmMoeDsaForCausalLM = MLA + DeepSeek Sparse Attention) to the MoRI-EP WideEP disaggregated serving path, stacked on the #171 unified launcher. Fully isolated from DeepSeek-V3/R1: GLM gets its own image + a MODEL_NAME-gated runtime path, so existing models are byte-identical to develop. Defects fixed (validated 1P/1D EP8 + 2P/2D EP16, NIAH 2k-35k = 10/10, no crash): - Long-context accuracy collapse: vLLM #47766 cache-key fix keeps the persistent sparse-MLA kernel ON (keys metadata on per-request context+query len). - 8k disagg prefill crash: DSA adds a 2nd (indexer) KV cache per layer that the single-geometry MoRIIO connector never transferred; paired + shipped prefill-> decode. Plus DSA invalid-token kernel fix (#45324) and shik-latest DP-notify. Changes: - docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile: NEW per-model image (raviguptaamd/vllm glm5.1-dsa-wideEP_on_shik_latest + aiter e03fa6040 + mori 42e895472b08 + router). The base vllm_disagg_inference Dockerfile (DSV3/R1) is left untouched. Future models add their own Dockerfile the same way. - models.json: card pyt_vllm_disagg_mori_glm-5.1-fp8 (GLM_SKIP_PATCHERS=1: image carries the DSA fixes in-source). - models.yaml: GLM-5.1-FP8 recipe (block=1, AITER MLA on, eager, mori backends). DeepSeek-V3 dp: caps (--max-num-seqs 64 --max-model-len 32768) to bound the newer base's decode logits workspace (isolated to the DSV3 entry). - connectors/moriio.sh: MODEL_NAME-gated GLM DSA runtime patchers (pure no-op for other models); GLM_SKIP_PATCHERS switch for baked-fix images. - 9 idempotent, anchor-based, self-skipping GLM DSA patcher scripts. KNOWN OPEN DEFECT (future work): 4P/4D EP32 emits corrupted tokens at all context lengths (suspect moriep all-to-all combine at scale); use 1P/1D and 2P/2D. Co-Authored-By: Claude --- ...gg_inference.glmv5.1.ubuntu.amd.Dockerfile | 335 ++++++++++++++++++ models.json | 32 ++ .../apply_glm_aiter_sampling_oob_fix.py | 148 ++++++++ .../apply_glm_dsa_indexer_warmup_fix.py | 251 +++++++++++++ .../vllm_dissag/apply_glm_dsa_kernel_fix.py | 85 +++++ .../apply_glm_dsa_moriio_dualkv_fix.py | 176 +++++++++ .../apply_glm_dsa_moriio_engine_fix.py | 116 ++++++ .../apply_glm_dsa_moriio_gate_fix.py | 133 +++++++ .../apply_glm_dsa_moriio_instrument.py | 92 +++++ ...pply_glm_dsa_persistent_kernel_gate_fix.py | 129 +++++++ .../apply_glm_moriio_abort_guard_fix.py | 98 +++++ scripts/vllm_dissag/connectors/moriio.sh | 100 +++++- scripts/vllm_dissag/keepalive_bench.sh | 18 + scripts/vllm_dissag/models.yaml | 102 +++++- scripts/vllm_dissag/run_xPyD_models.slurm | 32 +- scripts/vllm_dissag/vllm_disagg.sh | 28 +- 16 files changed, 1857 insertions(+), 18 deletions(-) create mode 100644 docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile create mode 100644 scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py create mode 100644 scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py create mode 100644 scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py create mode 100755 scripts/vllm_dissag/keepalive_bench.sh 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..f61a8819 --- /dev/null +++ b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile @@ -0,0 +1,335 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# ============================================================================= +# vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile +# GLM-5.1-FP8 (MLA + DeepSeek Sparse Attention) MoRI-EP WideEP disagg image. +# PER-MODEL image, isolated from the base vllm_disagg_inference Dockerfile +# (which stays pinned to the DeepSeek-V3 / R1 stack). This split lets each model +# pin its own vLLM/AITER/MoRI without disturbing the others -- add a new +# vllm_disagg_inference..ubuntu.amd.Dockerfile per future model +# (e.g. Kimi-2.6) rather than repinning the shared DSV3 image. +# +# ALL connectors in one image: moriio (TP + MoRI-EP wideEP) + rixl (NIXL TP + +# DeepEP wideEP). = the fullsource MoRI stack, plus a UCX/RIXL/rocSHMEM/DeepEP +# transport layer gated by --build-arg WITH_NIXL (default 1 = everything). +# +# docker build -f docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile \ +# -t /vllm-disagg:glmv5.1 . +# export DOCKER_IMAGE_NAME=/vllm-disagg:glmv5.1 +# +# WITH_NIXL=1 (default) => builds UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from +# source, so all four connector combos (moriio TP/wideEP, rixl NIXL TP, DeepEP +# wideEP) are present (~+30-45 min build vs WITH_NIXL=0). +# WITH_NIXL=0 => MoRI-EP only (moriio TP/wideEP + deepep-from-base); lean, faster. +# +# STATUS (GLM-5.1-FP8 on this stack): 1P/1D EP8 + 2P/2D EP16 NIAH 2k-35k = 10/10, +# no crash; long-context accuracy fixed via vLLM #47766 (persistent sparse-MLA kept +# ON). 4P/4D EP32 is a KNOWN OPEN DEFECT: token corruption at ALL context lengths +# (garbage output even at 2k), distinct from the long-context bug; prime suspect is +# the moriep all-to-all combine at EP32 scale -> deferred to future work. Use 1P/1D +# and 2P/2D only. (BASE_IMAGE is a gated nightly; override --build-arg BASE_IMAGE=...) +# ============================================================================= +# Reconstructs the validated v1.2.1 (mori121) runtime stack by applying the recipe's +# component pins ON TOP of the open ROCm vLLM ci_base, cloning each source from +# public Git (no local build-contexts). Mirrors dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable: +# +# - BASE: rocm/vllm-dev:ci_base-0fcd9b99... (open ROCm 7.2 / cp312 CI base). +# - MoRI -> built from ROCm/MoRI @ v1.2.1 (BUILD_UMBP=OFF). +# - 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 shikamd123/vllm @ +# vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer (Wide-EP multi-pod PD, the +# connector/router reference for the 2P2D DP=EP=16 topology). Full compile: it is +# 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 open rocm/vllm-dev ci_base pinned by the validated recipe +# (dist-inf-cookbook Dockerfile.vllm.mori121_shareable). 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-0fcd9b99cc9d63202da4c858d8ebc6582c9e2491 +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. Default 1 => all connectors built +# (UCX/RIXL/rocSHMEM/DeepEP). Set --build-arg WITH_NIXL=0 for a lean MoRI-EP-only image. +ARG WITH_NIXL=1 +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: MoRI main tip past v1.2.1, validated by MAD-private #338 for GLM-5.1 +# DSA WideEP disagg (v1.2.1 large-transfer notify path was insufficient at high EP). +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 && \ + 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])" && \ + mkdir -p /app && echo "MORI_REF=${MORI_REF}@$(git -C /tmp/mori-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/mori-src + +# ----------------------------------------------------------------------------- +# 2. AITER: build STOCK upstream ROCm/aiter @ e03fa6040 from source (NO fork, +# NO gqa64-fold patch). Under vLLM #47766 the sparse-MLA persistent path stays +# ON, so GLM's gqa=64 decode hits aiter's PRE-EXISTING persistent gqa64->16 fold +# (aiter/mla.py: `nhead in range(32,128+1,16) and persistent_mode`); the fork's +# extra non-persistent fold is never exercised, so stock is sufficient. +# Validated by MAD-private #338: 1P/1D EP8 + 2P/2D EP16 NIAH PASS on this exact +# aiter tip under #47766. Pin the exact commit (the one tested), not the release +# wheel. Then invalidate the stale prewarmed JIT cache compiled against the old .so. +# ----------------------------------------------------------------------------- +ARG AITER_REPO=https://github.com/ROCm/aiter.git +ARG AITER_REF=e03fa6040 +RUN 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 && \ + python3 - <<'PYEOF' +# Verify aiter/mla.py installed + has the persistent gqa64 fold, WITHOUT importing +# aiter/torch (torch->amdsmi->libamd_smi.so is not loadable at build: no GPU in sandbox). +import glob, pathlib +cands = glob.glob("/usr/local/lib/python*/dist-packages/aiter/mla.py") + \ + glob.glob("/usr/lib/python*/dist-packages/aiter/mla.py") +assert cands, "aiter/mla.py not found in site-packages after install" +src = pathlib.Path(cands[0]).read_text() +assert "persistent_mode" in src, f"AITER persistent fold path MISSING in {cands[0]}" +print("STOCK AITER OK (persistent gqa64 fold path present):", cands[0]) +PYEOF +RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" && \ + echo "AITER_REF=${AITER_REF} (stock)" >> /app/versions.txt + +# ----------------------------------------------------------------------------- +# 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 +ARG VLLM_REF=glm5.1-dsa-wideEP_on_shik_latest +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")) +# Stock source build of ROCm/aiter@e03fa6040 reports 0.1.17.dev195+ge03fa6040. +# Verify the aiter install survived the vLLM install (present + carries the e03fa6040 +# commit tag) rather than pinning a release version string. +assert av and "e03fa6040" in av, f"AITER missing/downgraded (want e03fa6040 build): {av!r}" +import mori, mori.io, mori.ops +print("Post-vLLM check OK: AITER", av, "+ MoRI importable") +PYEOF + +# ----------------------------------------------------------------------------- +# 4. vllm-router (DP-rank round-robin + MoRIIO connector) — built in, so NO +# external vllm-router binary is needed (leave ROUTER_BINARY unset). +# Source = vllm-project/router PR #181 branch, which now carries BOTH the +# round-robin DP-rank fix (11841c0d) AND the 2P2D KV-notify fix (6409ac1: +# remote_dp_rank_override + remote_dp_size). The KV-notify fix is REQUIRED: +# without it the 2P2D EP=16 run reproducibly wedges with "remote blocks never +# arrived" deferred-write expiries (decode notify targets the wrong DP rank). +# This is the exact source of the validated vllm-router-2p2d-dpfix binary. +# Pinned Rust toolchain (>=1.88: router deps time/home require rustc 1.88). +# ----------------------------------------------------------------------------- +ARG ROUTER_REPO=https://github.com/raviguptaamd/router.git +ARG ROUTER_REF=ravgupta/discovery-dp-rank-roundrobin +ARG RUST_TOOLCHAIN=1.88.0 +RUN if ! command -v cargo >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_TOOLCHAIN}"; \ + fi && \ + export PATH="/root/.cargo/bin:${PATH}" && \ + rm -rf /tmp/vllm-router-src && \ + git clone --filter=blob:none "${ROUTER_REPO}" /tmp/vllm-router-src && \ + cd /tmp/vllm-router-src && git checkout "${ROUTER_REF}" && \ + cargo build --release && \ + install -m 755 target/release/vllm-router /usr/local/bin/vllm-router && \ + vllm-router --help 2>&1 | grep -q moriio && \ + echo "VLLM_ROUTER_REF=${ROUTER_REPO}@${ROUTER_REF}@$(git -C /tmp/vllm-router-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/vllm-router-src + +# ----------------------------------------------------------------------------- +# 4b. WITH_NIXL=1 (default): UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Single guarded RUN so +# WITH_NIXL=0 skips it entirely (no layers, no cost). Build-verified on ci_base. +# ----------------------------------------------------------------------------- +ENV _ROCM_DIR=/opt/rocm \ + _UCX_SOURCE=https://github.com/ROCm/ucx.git \ + _UCX_BRANCH=da3fac2a \ + _UCX_INSTALL_DIR=/usr/local/ucx/ \ + _RIXL_SOURCE=https://github.com/ROCm/RIXL.git \ + _RIXL_BRANCH=f33a5599 \ + _RIXL_INSTALL_DIR=/usr/local/RIXL/install \ + _NIXLBENCH_INSTALL_DIR=/usr/local/RIXL +RUN if [ "${WITH_NIXL}" != "1" ]; then \ + echo "WITH_NIXL=${WITH_NIXL}: skipping UCX/RIXL/rocSHMEM/DeepEP (MoRI-EP + base DeepEP only)"; \ + else set -e && \ + echo "WITH_NIXL=1: building UCX + RIXL + rocSHMEM + DeepEP" && \ + apt-get update && apt-get install -y \ + autoconf automake libtool autogen pkg-config m4 gcc make \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core strace libgflags-dev \ + libaio-dev liburing-dev libcpprest-dev libgrpc-dev libgrpc++-dev \ + libprotobuf-dev protobuf-compiler-grpc wget && \ + pip install meson==0.64.0 "pybind11[global]" pyyaml && \ + # UCX + cd /tmp && git clone "${_UCX_SOURCE}" && cd ucx && git checkout "${_UCX_BRANCH}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${_UCX_INSTALL_DIR}" --with-rocm="${_ROCM_DIR}" \ + --disable-go --disable-java --disable-assertions --enable-mt && \ + make -j && make install && \ + # googletest (RIXL dep) + cd /tmp && wget -q https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz && \ + tar -xzf v1.14.0.tar.gz && cd googletest-1.14.0 && mkdir -p build && cd build && \ + cmake -DBUILD_SHARED_LIBS=on .. && make -j && make install && \ + # RIXL + python bindings + cd /tmp && git clone "${_RIXL_SOURCE}" && cd RIXL && git checkout "${_RIXL_BRANCH}" && \ + meson setup build/ --prefix="${_RIXL_INSTALL_DIR}" -Ducx_path="${_UCX_INSTALL_DIR}" \ + -Ddisable_gds_backend=true -Dcudapath_inc="${_ROCM_DIR}/include" -Dcudapath_lib="${_ROCM_DIR}/lib" && \ + cd build && ninja && ninja install && cd /tmp/RIXL && \ + pip install --config-settings=setup-args="-Dcudapath_inc=${_ROCM_DIR}/include" \ + --config-settings=setup-args="-Dcudapath_lib=${_ROCM_DIR}/lib" \ + --config-settings=setup-args="-Ducx_path=${_UCX_INSTALL_DIR}" \ + --config-settings=setup-args="-Ddisable_gds_backend=true" . && \ + # rocSHMEM (DeepEP dep) + cd /tmp && git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git && \ + cd rocm-systems && git sparse-checkout set --cone projects/rocshmem && git checkout develop && \ + mkdir -p /tmp/rocshmem-build && cd /tmp/rocshmem-build && \ + /tmp/rocm-systems/projects/rocshmem/scripts/build_configs/all_backends \ + -DUSE_EXTERNAL_MPI=OFF -DGPU_TARGETS="${GFX_COMPILATION_ARCH}" && \ + # DeepEP (build develop against the installed vLLM/torch) + cd /tmp && git clone https://github.com/ROCm/DeepEP.git && cd DeepEP && \ + PYTORCH_ROCM_ARCH="${GFX_COMPILATION_ARCH}" CFLAGS="-O3 -fPIC" \ + CXXFLAGS="-O3 -fPIC --offload-arch=${GFX_COMPILATION_ARCH}" HIP_CXX_FLAGS="-O3 -fPIC" \ + python3 setup.py --variant rocm --nic "${NIC_COMPILATION_ARCH}" build develop && \ + echo "WITH_NIXL build complete" >> /app/versions.txt && \ + rm -rf /tmp/ucx /tmp/googletest-1.14.0 /tmp/v1.14.0.tar.gz /tmp/rocm-systems /tmp/rocshmem-build; \ + fi +ENV LD_LIBRARY_PATH="/usr/local/ucx/lib:/usr/local/lib:/usr/local/RIXL/install/lib:${LD_LIBRARY_PATH}" \ + PATH="/usr/local/ucx/bin:${PATH}" + +# ----------------------------------------------------------------------------- +# 5. Cache locations (structural: WHERE the JIT/compile caches live in the image). +# These are the mount target for the launcher's persistent host JIT cache. +# ----------------------------------------------------------------------------- +# The image ships NO runtime recipe / tuning / platform ENV. By design, everything +# run-tunable is applied at launch, so this image stays a clean binary/library artifact +# and the same image serves any model/cluster without a rebuild: +# - model-serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, *_CUDAGRAPH_MODE, *_MORI_BACKEND, +# GPU_MEMORY_UTILIZATION, KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...) +# -> scripts/vllm_dissag/models.yaml (per-model env:, so dense vs MoE differ) +# - ROCm-7.2.3 GPU-RDMA platform env (expandable_segments:False x2, MORI_GPU_ARCHS, +# HSA_ENABLE_IPC_MODE_LEGACY=0, HSA_NO_SCRATCH_RECLAIM) and the MoRI/RDMA fabric +# tuning (MORI_RDMA_TC/SL, MORI_IB_GID_INDEX, MORI_NUM_QP_PER_PE, VLLM_MORIIO_*, ...) +# -> scripts/vllm_dissag/connectors/.env (cluster-editable, no rebuild) +# The slurm launcher forwards both via `docker -e` (platform env must reach PID 1 - +# PyTorch reads alloc-conf at import). Running this image WITHOUT the launcher: set the +# vars you need yourself (see connectors/moriio.env + models.yaml for the values). +ENV AITER_JIT_DIR=/opt/vllm_cache/aiter_jit \ + VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + COMGR_CACHE_DIR=/opt/vllm_cache/comgr + +# ----------------------------------------------------------------------------- +# 6. CRITICAL: scrub build-time MoRI JIT state. The `import mori` verification +# steps above compile/lock MoRI EP kernels under /root/.mori/jit on THIS build +# host, leaving stale .hsaco.lock files (ep_internode_v1, ep_internode_v1ll, ...). +# At runtime on the cluster, MoriAll2AllManager finds those locks, waits on a +# build-in-progress whose owner PID is long gone, and DEADLOCKS at ep:0 init. +# A clean image ships /root/.mori empty -> runtime compiles fresh. +# Clearing these makes the from-source image boot clean on 2P2D/4P4D. +# ----------------------------------------------------------------------------- +RUN rm -rf /root/.mori /tmp/mori_jit_* && mkdir -p /root/.mori && \ + echo "JIT_SCRUBBED: /root/.mori + /tmp/mori_jit_* cleared at build end" >> /app/versions.txt + +RUN cat /app/versions.txt 2>/dev/null | tail -20 || true diff --git a/models.json b/models.json index af914f46..689754db 100644 --- a/models.json +++ b/models.json @@ -3420,6 +3420,38 @@ }, "args": "-N 2 -n 2" }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8", + "url": "", + "dockerfile": "docker/vllm_disagg_inference.glmv5.1", + "scripts": "scripts/vllm_dissag/run_xPyD_models.slurm", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "GLM-5.1-FP8", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "GLM_SKIP_PATCHERS": "1", + "BENCHMARK_COMBINATIONS": "1024/1024" + }, + "args": "-N 2 -n 2" + }, { "name": "pyt_vllm_disagg_mori_deepseek-v3-5layer", "url": "", diff --git a/scripts/vllm_dissag/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/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..c26a0944 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -124,13 +124,96 @@ _moriio_build_kv_transfer_config() { } connector_runtime_patch() { - # No-op: the MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL - # split, DP-rank hash-failsafe) are committed in-source in the vLLM the image is - # built from (see the Dockerfile VLLM_REF). There is no runtime .py patcher — that - # would be a drifting duplicate of fixes that already live upstream in the fork. - # If you ever run an image WITHOUT these fixes baked, use an image that has them - # (rebuild from the pinned VLLM_REF) rather than patching a stock image at runtime. - return 0 + # MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL split, + # DP-rank hash-failsafe) are committed in-source in the vLLM the image is built + # from (Dockerfile VLLM_REF). There is no generic runtime .py patcher for those — + # that would be a drifting duplicate of fixes already upstream in the fork. + # + # EXCEPTION — GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA + DSA sparse attention): + # DSA is a NEW attention family the MoRIIO connector was never built for. It adds + # a 2nd KV cache per layer (indexer) with a different geometry, which the + # single-geometry connector mis-handles -> disagg KV transfer stalls; plus a DSA + # invalid-token kernel bug (#45324) that produces `!!!`. These are model-specific + # code gaps, applied here as idempotent, anchor-based, self-skipping .py patchers + # (they no-op cleanly if the fix is native/refactored on the chosen image). Gated + # on MODEL_NAME so DeepSeek/other models are a pure no-op (byte-identical to before). + # The MoRI version is pinned by the Dockerfile MORI_REF (post-1.2.1 main with the + # large-transfer notify/mapping fixes #424/#436/#432 baked in); if a newer MoRI is + # needed, update MORI_REF and rebuild the image — no runtime library swap here. + [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0 + _glm_dsa_runtime_patch +} + +# GLM-5.1 DSA patchers (see connector_runtime_patch). Ported from MAD-private #338. +# Resolves the vLLM install dir, then applies the 4 required patchers in order, +# aborting on a hard failure (a real failure means GLM emits garbage or stalls, so +# failing at launch is correct). Patchers self-skip (rc 0) when their anchor is +# absent, so an image that already carries or refactored a fix no-ops cleanly. +_glm_dsa_runtime_patch() { + # GLM_SKIP_PATCHERS=1: the serving image already carries the GLM-5.1 DSA fixes + # in-source (e.g. the #47766 stack image built from raviguptaamd/vllm@ + # glm5.1-dsa-wideEP_on_shikpate_06_29_customer). Skip ALL runtime patchers — they + # are redundant, and the persistent-gate/sampling-overlay patchers would actively + # REGRESS a baked image (turn persistent MLA off / overwrite stock aiter kernels). + if [ "${GLM_SKIP_PATCHERS:-0}" = "1" ]; then + echo "[glm] GLM_SKIP_PATCHERS=1: image carries DSA fixes in-source; skipping runtime patchers." + return 0 + fi + local _patch_dir="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)}" + local _vllm_dir + _vllm_dir="$(python3 -c 'import vllm, os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || true)" + if [ -z "${_vllm_dir}" ] || [ ! -d "${_vllm_dir}" ]; then + echo "Error: [glm] cannot locate vLLM install dir for DSA patchers. Aborting." >&2 + exit 1 + fi + echo "[glm] MODEL_NAME=GLM-5.1-FP8: applying DSA runtime patchers against ${_vllm_dir}" + + # Ordered list of REQUIRED patchers (all abort on hard failure). + # GLM_PERSIST_GATE=0 skips the persistent-MLA accuracy gate (debug only: to test + # whether the non-persistent kernel it routes to is what crashes disagg at >=8k). + local _gate_patcher="apply_glm_dsa_persistent_kernel_gate_fix.py" + [ "${GLM_PERSIST_GATE:-1}" = "0" ] && _gate_patcher="" + local _p + for _p in \ + apply_glm_dsa_kernel_fix.py \ + apply_glm_dsa_moriio_dualkv_fix.py \ + apply_glm_dsa_moriio_engine_fix.py \ + apply_glm_dsa_moriio_gate_fix.py \ + apply_glm_moriio_abort_guard_fix.py \ + ${_gate_patcher} \ + apply_glm_aiter_sampling_oob_fix.py; do + local _py="${_patch_dir}/${_p}" + if [ ! -f "${_py}" ]; then + echo "Error: [glm] required patcher ${_py} not found. Aborting." >&2 + exit 1 + fi + echo "[glm] applying ${_p}" + python3 "${_py}" "${_vllm_dir}" 2>&1 || { + echo "Error: [glm] ${_p} failed — GLM-5.1 would emit garbage or stall. Aborting." >&2 + exit 1 + } + done + + # Optional DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1). Force-compiles the DSA + # indexer kernels at boot so they never JIT mid-inference. Opt-in because it drives a + # large (>=8k) prefill forward at boot: on stacks where that forward faults it makes + # the fault DETERMINISTIC at boot (useful for debugging) rather than on first request. + if [ "${GLM_INDEXER_WARMUP:-0}" = "1" ]; then + local _warm="${_patch_dir}/apply_glm_dsa_indexer_warmup_fix.py" + if [ -f "${_warm}" ]; then + echo "[glm] applying DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1)" + python3 "${_warm}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] indexer-warmup patch failed (non-fatal)." + fi + fi + + # Optional diagnostic instrumentation (GLM_INSTRUMENT=1). Non-fatal. + if [ "${GLM_INSTRUMENT:-0}" = "1" ]; then + local _instr="${_patch_dir}/apply_glm_dsa_moriio_instrument.py" + if [ -f "${_instr}" ]; then + echo "[glm] applying instrumentation (GLM_INSTRUMENT=1): apply_glm_dsa_moriio_instrument.py" + python3 "${_instr}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] instrumentation failed (non-fatal)." + fi + fi } # connector_launch_worker [dp_start_rank] @@ -219,7 +302,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 +325,7 @@ connector_launch_worker() { "${exec_args[@]}" \ "${extra_args[@]}" \ "${kv_args[@]}" \ + "${model_args[@]}" \ 2>&1 | tee /run_logs/${SLURM_JOB_ID}/${log_prefix}_NODE${NODE_RANK}.log >/dev/null & WORKER_PID=$! return 0 diff --git a/scripts/vllm_dissag/keepalive_bench.sh b/scripts/vllm_dissag/keepalive_bench.sh new file mode 100755 index 00000000..68d80dcb --- /dev/null +++ b/scripts/vllm_dissag/keepalive_bench.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Keepalive with LIGHT heartbeat traffic: holds the disagg server up AND sends a +# tiny request every ~20s so prefill discovery/ping stays registered (idle sleep +# lets the prefill ZMQ ping die ~2min in). Runs KEEPALIVE_MINS (default 90). +: "${KEEPALIVE_MINS:=90}" +PORT="${BENCHMARK_PORT:-30000}" +MODEL="/mnt/m2m_nobackup/models_blog/GLM-5.1-FP8" +echo "[keepalive] light-traffic hold ${KEEPALIVE_MINS}min on :${PORT}" +_end=$(( $(date +%s) + KEEPALIVE_MINS*60 )) +i=0 +while [ "$(date +%s)" -lt "$_end" ]; do + curl -s -m 30 "http://127.0.0.1:${PORT}/v1/completions" -H "Content-Type: application/json" \ + -d "{\"model\":\"${MODEL}\",\"prompt\":\"hi\",\"max_tokens\":1,\"temperature\":0}" >/dev/null 2>&1 + i=$((i+1)) + [ $((i % 3)) -eq 0 ] && echo "[keepalive] heartbeat $i, $(( (_end-$(date +%s))/60 ))min left" + sleep 20 +done +echo "[keepalive] done" diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..9eec3e7d 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -160,10 +160,17 @@ _deepseek_recipe_env: &deepseek_recipe_env # above. No tp: blocks — TP is unsupported for these models. DeepSeek-V3: env: *deepseek_recipe_env + # NEWER-BASE ADAPTATION (isolated to this model): Shiksha's newer vLLM base added + # TritonMLAMetadataBuilder._reserve_attn_logits_workspace(), which pre-reserves the + # decode split-KV logits workspace at WORST CASE + # (max_num_seqs x q_heads x max_kv_splits(max_model_len) x lse_dim x fp32). With + # DSV3's defaults (max_num_seqs=256, max_model_len~163k) this reserves ~128 GiB and + # OOMs at KV init. Cap max-num-seqs + max-model-len to bound the workspace. This is a + # per-model dp: flag (model_args), fully isolated -- does NOT touch GLM/other recipes. prefill: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" decode: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" DeepSeek-V3-5layer: env: *deepseek_recipe_env @@ -178,3 +185,94 @@ DeepSeek-R1: dp: "" decode: dp: "" + +# ============================ MoE + DSA (wideEP only) ============================ + +# GLM-5.1-FP8 (zai-org/GLM-5.1-FP8, arch GlmMoeDsaForCausalLM): MLA + DeepSeek +# Sparse Attention (DSA). 78 layers (3 dense + 75 MoE), 256 routed experts top-8 +# + 1 shared, FP8 block 128. wideEP-only (see WIDE_EP_ONLY_MODELS in the slurm). +# +# Validated DOCKER_IMAGE_NAME (submit-time, not set here — #171 requires it explicit): +# rocmshared/pytorch-private:vllm-wideep_06_29_2026_Shiksha_dp16_2p2d_mori_v1.2.1_aiter_v0.1.16.post3_nightlybase_mori121 +# On this image the DSA patchers self-adapt: only the invalid-token kernel fix +# (#45324) applies; the MoRIIO dual-KV geometry + completion-gate fixes are NATIVE +# (moriio_layout.py per-layer geometry + engine sealed writes_expected), so those +# patchers cleanly no-op. On the older b10a9f7a image all 4 patchers apply. +# +# GLM differs from DeepSeek in 3 recipe-defining ways (both are MLA MoE, but GLM +# is DSA-sparse): +# - KV_BLOCK_SIZE=1 (DSA sparse indexer REQUIRES block-size 1; DS uses 16) +# - VLLM_ROCM_USE_AITER_MLA=1 (GLM MLA path ON via AITER sparse; DS sets 0) +# - prefill EAGER (NONE), decode PIECEWISE (prefill cudagraph capture deadlocks +# on this stack; decode PIECEWISE captures cleanly on DSA and is a ~3.7x ITL win +# (validated: ~69ms vs ~264ms eager). Global VLLM_CUDAGRAPH_MODE=NONE as the +# safe floor; per-role PREFILL=NONE / DECODE=PIECEWISE override it.) +# Note: block=1 + AITER_MLA=1 are ALREADY the moriio.sh connector defaults — GLM +# keeps them; DeepSeek is the one that overrides them off. Set here explicitly so +# the recipe is self-documenting and robust to connector default changes. +# +# DSA adds a 2nd KV cache per layer (indexer); the MoRIIO connector needs the GLM +# DSA patchers (kernel #45324 + dual-KV geometry + per-layer offset + completion +# gate) applied by connector_runtime_patch in connectors/moriio.sh (gated on this +# MODEL_NAME). Those are code patches, not flags — nothing to add here for them. +# +# dp_flags carry the GLM tool/reasoning parsers (AMD GLM recipe). They are applied +# to BOTH roles by compose() in vllm_disagg.sh and reach `vllm serve` via the +# connector's model_args. block-size / kv-cache-dtype / all2all / cudagraph come +# from the env: recipe above (the connector emits them), so the dp: blocks are +# empty like the DeepSeek family. +# +# LONG-CONTEXT CAVEAT (vLLM #40018): the ROCM_AITER_MLA_SPARSE prefill indexer +# corrupts output for prompts beyond ~16-18k tokens on this image (mori v1.2.1) — +# coherent + correct needle retrieval up to 14k, garbage (repetition collapse, +# unique-ratio ~0.1) at ~18.7k. This is an UPSTREAM kernel bug, not the MAD port. +# TESTED: pinning --max-model-len 32768 does NOT move the threshold (workspace is +# sized max_model_len*40 but the corruption onset is a fixed ~18k token count in +# the gather/logits kernel, not a buffer-scaling artifact). So no config knob +# helps; it needs the complete upstream prefill fix in a newer image. Left at the +# native max_model_len (do not cap — capping gives no accuracy benefit and only +# limits usable context). Serve prompts <~14k for correct output on this image. +GLM-5.1-FP8: + env: + VLLM_USE_V1: "1" + 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" + DECODE_CUDAGRAPH_MODE: "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_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: + dp: "" diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..98281f57 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_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 @@ -427,11 +448,14 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" # Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. # sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) # long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# keepalive -> keepalive_bench.sh (hold server up KEEPALIVE_MINS +# for external accuracy probes) BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" case "$BENCHMARK_SCRIPT" in sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; - *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context)" >&2; exit 1 ;; + keepalive) BENCHMARK_SCRIPT_FILE="keepalive_bench.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -583,6 +607,9 @@ docker run --rm \ ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ + -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ + ${GLM_PERSIST_GATE:+-e GLM_PERSIST_GATE=$GLM_PERSIST_GATE} \ + ${GLM_SKIP_PATCHERS:+-e GLM_SKIP_PATCHERS=$GLM_SKIP_PATCHERS} \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ @@ -590,6 +617,7 @@ docker run --rm \ ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + ${KEEPALIVE_MINS:+-e KEEPALIVE_MINS=$KEEPALIVE_MINS} \ ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \ ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \ ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ diff --git a/scripts/vllm_dissag/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..acbabbd4 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -160,17 +160,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 )" From 740898414054022a40339acfeb5720b4e47fbe36 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Wed, 8 Jul 2026 15:51:15 +0000 Subject: [PATCH 2/4] vllm_dissag: make NIAH harness thinking-model-aware (GLM-5.1) benchmark_niah.py mis-scored thinking models: it never disabled thinking and read only content + reasoning_content. GLM-5.1 emits chain-of-thought into the `reasoning` field and leaves `content` empty until the final answer, so with a small max_tokens the answer never lands in content -> a false 0/10 even when generation is correct. - Add chat_template_kwargs.enable_thinking=false so the answer goes to content. - Also score the `reasoning` field as a fallback. Verified against GLM-5.1-FP8: correct 9-10/10 retrieval across 2k-35k on all tested topologies (EP8/EP16/EP32) after the fix. Co-Authored-By: Claude --- scripts/vllm_dissag/benchmark_niah.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index 0cdd027e..b1b0264d 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -50,6 +50,11 @@ def run(n_words): ], "temperature": 0.0, "max_tokens": MAXTOK, + # 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"}) @@ -59,7 +64,11 @@ def run(n_words): except Exception as e: print("words=%6d ERROR %s" % (n_words, e), 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) return len(found) From 9cfe987385eb4bc5da66c20fa34040c688cef0b1 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Wed, 8 Jul 2026 16:16:52 +0000 Subject: [PATCH 3/4] vllm_dissag: NIAH multi-seed support (NIAH_SEEDS) for variance-aware accuracy Needle layout is seeded, so a single run is deterministic (bit-exact on the same stack) but can't tell a real accuracy dip from single-needle variance. Add NIAH_SEEDS (default 0,1,2) to run each context length across multiple needle layouts; the summary now reports mean/min/max across seeds. Backward compatible: NIAH_SEEDS=0 reproduces the prior single-seed behavior. Co-Authored-By: Claude --- scripts/vllm_dissag/benchmark_niah.py | 29 ++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index b1b0264d..82a55663 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -8,6 +8,8 @@ # 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) import os, sys, json, random, urllib.request @@ -16,6 +18,10 @@ 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()] FILLER = ( "table chair window bottle pencil garden river mountain coffee planet " @@ -41,12 +47,12 @@ def make_haystack(n_words, seed=0): return " ".join(words) -def run(n_words): +def run(n_words, seed=0): 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, @@ -70,7 +76,7 @@ def run(n_words): + (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) @@ -79,14 +85,19 @@ 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" % (URL, MODEL, WORDS, SEEDS), flush=True) + results = {} # n_words -> list of scores across seeds (None on error) 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) + vals = [v for v in results[n] if v is not None] + if not vals: + print(" words=%6d ERR" % n, flush=True) + continue + mean = sum(vals) / len(vals) + print(" words=%6d mean=%.1f/10 min=%d max=%d (n=%d)" + % (n, mean, min(vals), max(vals), len(vals)), flush=True) if __name__ == "__main__": From 74c69b7f40b3e5cb0f9769cfe56029d3a2bd6d4f Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 10 Jul 2026 07:09:21 +0000 Subject: [PATCH 4/4] vllm_dissag: NIAH gate robust to cold-start JIT (warmup + readiness probe) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On a freshly-booted node the first request of each context shape pays the full JIT/kernel-autotune compile (minutes). The NIAH harness scored the FIRST request, so cold compile landed on a scored/gated request -> false 0/10 or timeout, failing the accuracy gate and skipping the perf sweep. Root-caused by reproducing on a cold boot (0 results) vs a warm server (all pass) on the same image. Fixes: - benchmark_niah.py: add a warmup pass (NIAH_WARMUP=1 default) — one throwaway request per context length before scoring, with a generous timeout, failures tolerated. Scored requests are then always warm. - benchmark_niah.py: distinguish TIMEOUT/ERROR from a wrong answer. Timeouts return a sentinel (excluded from mean, never counted as 0/10); summary flags NO-RESULT with guidance instead of silently reporting 0. - benchmark_niah.sh: replace the blind `sleep 10` with a /v1/models readiness poll (up to 5 min), and forward NIAH_WARMUP. Verified: patched harness on the warm server passes 10/10; cold-boot repro no longer produces false 0/10 because compile happens in the warmup pass. Co-Authored-By: Claude --- scripts/vllm_dissag/benchmark_niah.py | 59 ++++++++++++++++++++++----- scripts/vllm_dissag/benchmark_niah.sh | 18 +++++++- 2 files changed, 64 insertions(+), 13 deletions(-) diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index 82a55663..43d81850 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -11,6 +11,12 @@ # 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") @@ -22,6 +28,10 @@ # 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 " @@ -47,7 +57,8 @@ def make_haystack(n_words, seed=0): return " ".join(words) -def run(n_words, seed=0): +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": [ @@ -55,7 +66,7 @@ def run(n_words, seed=0): {"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 @@ -65,10 +76,26 @@ def run(n_words, seed=0): 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 # Score content plus any reasoning field (some servers surface CoT as # `reasoning` or `reasoning_content`) so a thinking model is never mis-scored. @@ -85,19 +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 seeds=%s" % (URL, MODEL, WORDS, SEEDS), flush=True) - results = {} # n_words -> list of scores across seeds (None on error) + 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, s) for s in SEEDS] print("=== NIAH summary (mean/min/max across %d seed(s)) ===" % len(SEEDS), flush=True) for n in WORDS: - vals = [v for v in results[n] if v is not None] + 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 ERR" % n, flush=True) + 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) - print(" words=%6d mean=%.1f/10 min=%d max=%d (n=%d)" - % (n, mean, min(vals), max(vals), len(vals)), flush=True) + 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}"