commit 777e910dd0e54ffe8f95acd534c76c56e7320fe2 Author: MOVZX Date: Thu Oct 15 01:01:52 2026 +0700 Initial commit Signed-off-by: MOVZX diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..c80eaa7 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,10 @@ +/vllm +/.venv +/.git +/.pi +__pycache__ +*.pyc +/.vllm-cache +/artifacts +/.env +/build-*.log diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..9eae304 --- /dev/null +++ b/.env.example @@ -0,0 +1,55 @@ +# ── network ────────────────────────────────────────────────────────────────── +HOST_BIND_ADDR=0.0.0.0 +HOST_PORT=8080 +VLLM_API_KEY= + +# ── model ──────────────────────────────────────────────────────────────────── +MODEL_DIR=/data/models/halt95/Swift1.5-Qwen3.8-Flash-Next-W4A16-Merlin +MODEL_MOUNT_PATH=/models/Qwen3.8-Flash-Next +MODEL_NAME=Hinata + +# ── image and container ────────────────────────────────────────────────────── +IMAGE=qwen38-flash-next +IMAGE_TAG=latest +CONTAINER_NAME=vllm +# CACHE_HOST_DIR=/srv/vllm-cache + +# ── engine ─────────────────────────────────────────────────────────────────── +# Card hosting the PLE offload home (no CUDA context there). +PLE_HOME=0 +# GPU lane, in bus order. The shipped shape uses all four cards. +CUDA_VISIBLE_DEVICES=0,1,2,3 +# Topology selector (opt-in): 0 keeps the fixed shipped settings | 1 | strict. +# The image has no numactl, so AUTO_TOPO_BIND=1 is skipped. +AUTO_TOPO=0 +AUTO_TOPO_NCCL=legacy +AUTO_TOPO_BIND=0 +# LANE_GPUS=auto +# LANE_CPUS=0-15,16-31 +TOPO_TIMEOUT=20 +# Bounds-guard telemetry: off | warn | fail. +GUARD=warn +# 1 adds the two logging-only counters (exact-restart counters /30 s, FP8 KV clips /300 s). +COUNTERS=0 +# Seconds between engine stats lines (throughput, KV usage, queue depth). +VLLM_LOG_STATS_INTERVAL=1 + +# Extracted from orcarouter/Qwen3.8-Flash-Next-Uncensored, hidden 2560 +QWEN38_ABLITERATION=Qwen3.8-Flash-Next-Uncensored +ABLITERATION_DIR=./abliteration + +# ── lifecycle, logs, health ────────────────────────────────────────────────── +LOG_MAX_SIZE=100m +LOG_MAX_FILE=14 + +# The image healthcheck is a real one-token generation, not /v1/models. Only the tuning is here. +HEALTHCHECK_INTERVAL=60s +HEALTHCHECK_TIMEOUT=90s +HEALTHCHECK_START_PERIOD=30m +HEALTHCHECK_RETRIES=3 + +# ── build (docker compose build reads these) ────────────────────────────────────── +UV_CACHE_DIR=/root/.cache/uv +PIP_CACHE_DIR=/root/.cache/pip +MAX_JOBS=32 +ARCH=8.6 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2945efc --- /dev/null +++ b/.gitignore @@ -0,0 +1,20 @@ +# the engine tree: build.sh fetches v0.30.0 and applies patches/, never tracked +/vllm/ +/.venv/ +/.vllm-cache/ +build-*.log +*.tar.gz +*.bundle +*.bundle.gz +.build-products +.build-products.sha256 +.pi/ +.env +docker-compose.override.yml +__pycache__/ +*.pyc + +# compiled kernels packed by build.sh step 8, for the image build. .gitkeep keeps the folder in +# git, so docker-compose.yml can always pass artifacts/ as the ops context. +/artifacts/* +!/artifacts/.gitkeep diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..774ba76 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,103 @@ +# syntax=docker/dockerfile:1 +# Three named build contexts feed this image. docker-compose.yml passes all three by itself: +# +# ops the compiled kernels from build.sh step 8 (this repo's artifacts/). The image then +# installs them instead of compiling: ~2 min. +# uvcache $HOME/.cache/uv ─┐ two wheel caches, so the build does not re-download the +# pipcache $HOME/.cache/pip ─┘ torch + CUDA wheel set (~6 GB). Read only. +# +# Plain `docker build -t vllm:latest .` passes none of them: the ops context is empty, so the +# image compiles every kernel from source, and the wheel caches start empty. BuildKit keeps the +# caches in its own volumes after the first build, so a rebuild downloads nothing again. +# +# docker compose up -d # kernels installed +# docker build -t vllm:latest . # compiles +# docker build --build-context ops="$PWD/artifacts" -t vllm:latest . # no compile +# +# A cache mount copies its context into the BuildKit cache volume the first time and reuses it +# after that. The caches on the host are only read: the build never writes to them. +# +# build.sh compares the tree object, the arch, CPython, CUDA and torch in the manifest with the +# tree it installs into, and hashes every unpacked file. A tarball that does not fit this tree +# stops the build. The image build runs with PACK=0: an image layer must not carry a 340 MB tarball. +FROM scratch AS uvcache +FROM scratch AS pipcache +FROM scratch AS ops +# trixie, not bookworm: the kernels of a host build on Debian 13 need glibc 2.38, bookworm has 2.36, +# so the loader would refuse them at serve time. trixie has 2.41. +FROM python:3.13-slim-trixie +# vllm: Qwen3.8-Flash-Next-W4A16 at 262K on 4x RTX 3090 (vLLM 0.30.0 + patches/). +# +# docker build -t vllm:latest . +# docker run -d --name vllm --gpus all --stop-timeout 20 --shm-size=8g \ +# --ulimit memlock=-1 -p 8080:8000 \ +# -v /data/models/halt95/Swift1.5-Qwen3.8-Flash-Next-W4A16-Merlin:/models/model:ro \ +# -v "$PWD/.vllm-cache":/cache \ +# -e MODEL_NAME=Scarlet vllm:latest +# +# The checkpoint mount must contain qsa_kv_scales_swift.json (the FP8 KV scales sidecar). +# +# The build fetches vLLM v0.30.0 and applies patches/ (build.sh step 2; needs network), then compiles +# every kernel of this tree from source (build.sh: cmake + ninja + nvcc 13.0 from the venv's +# CUDA wheels, arch 8.6). No prebuilt binaries go in. Budget 20-40 min and ~15 GB of build disk. +# Pass the ops context (see above) to install prebuilt kernels instead: ~2 min, same tree. +# trixie ships cmake 3.25, which is too old, so cmake and ninja come from pip. +# +# Layers are split so a rebuild stays cheap. The long compile only re-runs when a file it reads +# changes (build.sh, scripts/, patches/, start.sh). Docs, README, compose and +# the example env land in a later COPY, so editing them costs seconds. The apt, pip, uv and ccache +# caches are BuildKit cache mounts: they live outside the image, so a rebuild downloads nothing +# again and recompiles only the kernels whose source changed (build.sh prints the ccache stats). +ARG MAX_JOBS=32 +ARG ARCH=8.6 +ENV DEBIAN_FRONTEND=noninteractive PIP_NO_CACHE_DIR=1 +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt,sharing=locked \ + apt-get update && apt-get install -y --no-install-recommends git curl ca-certificates tar gzip procps \ + build-essential ccache binutils +WORKDIR /opt/vllm +# The build tools. These two mounts have their own ids: a cache mount with the same target path +# shares one volume, and a from= seed applies only when the volume is first created. Without the +# ids, this layer would create the pip/uv volumes empty and the seeds of the compile layer below +# would never be read. (PIP_NO_CACHE_DIR= (empty) makes pip use the cache.) +# build.sh uses uv for the venv when it is on PATH; it makes the 196-pin install several times faster. +# cmake<4: pip ships 4.x, the tree and its vendored projects are built and tested on 3.x. +RUN --mount=type=cache,id=tools-pip,target=/root/.cache/pip,sharing=locked \ + --mount=type=cache,id=tools-uv,target=/root/.cache/uv,sharing=locked \ + PIP_NO_CACHE_DIR= python -m pip install uv 'cmake<4' ninja +# The build inputs, in one layer. Nothing here changes on a docs edit. +COPY build.sh start.sh ./ +COPY scripts/ scripts/ +COPY patches/ patches/ +RUN --mount=type=cache,id=compile-uv,target=/root/.cache/uv,from=uvcache,sharing=locked \ + --mount=type=cache,id=compile-pip,target=/root/.cache/pip,from=pipcache,sharing=locked \ + --mount=type=cache,id=compile-ccache,target=/root/.cache/ccache \ + --mount=type=bind,from=ops,target=/ops,ro \ + chmod +x start.sh build.sh scripts/*.sh scripts/*.py \ + && echo "==> caches at the start of the compile:" && (du -sh /root/.cache/uv /root/.cache/pip /root/.cache/ccache 2>/dev/null || true) \ + && ops_tar="$(ls -1 /ops/build-artifacts-*.tar.gz 2>/dev/null | head -1 || true)" \ + && if [ -n "$ops_tar" ]; then echo "==> ops source: $ops_tar (prebuilt, no compile)"; else echo "==> ops source: none passed, compiling from source"; fi \ + && PIP_NO_CACHE_DIR= MAX_JOBS="$MAX_JOBS" ARCH="$ARCH" PACK=0 ARTIFACTS="$ops_tar" bash build.sh \ + && rm -rf /opt/vllm/vllm/.deps /opt/vllm/vllm/CMakeUserPresets.json \ + /opt/vllm/build-*.log +# The record build.sh writes (vllm/.build-products, vllm/.build-products.sha256) is KEPT on +# purpose: start.sh re-hashes every kernel against it before it starts the engine. +# Everything else: README, Dockerfile, compose, .env.example, systemd/. A copy, so it is cheap. +COPY . /opt/vllm/ +RUN chmod +x start.sh build.sh scripts/*.sh scripts/*.py \ + && rm -rf /root/.cache +ENV TREE=/opt/vllm/vllm VENV=/opt/vllm/.venv HOST=0.0.0.0 PORT=8000 CACHE_ROOT=/cache \ + MODEL_NAME=Scarlet MODEL_MOUNT_PATH=/models/model LOG= \ + CUDA_HOME=/opt/vllm/.venv/lib/python3.13/site-packages/nvidia/cu13 \ + FLASHINFER_WORKSPACE_BASE=/cache/flashinfer TRITON_CACHE_DIR=/cache/triton +ENV PATH="/opt/vllm/.venv/lib/python3.13/site-packages/nvidia/cu13/bin:$PATH" +VOLUME /cache +EXPOSE 8000 +# /v1/models answers 200 even when the engine is dead: probe a real one-token generation. +HEALTHCHECK --interval=60s --timeout=90s --start-period=30m --retries=3 \ + CMD m="${MODEL_NAME%% *}"; m="${m:-Scarlet}"; curl -fs -m 85 -X POST "http://127.0.0.1:${PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" ${VLLM_API_KEY:+-H "Authorization: Bearer $VLLM_API_KEY"} \ + -d "{\"model\":\"$m\",\"messages\":[{\"role\":\"user\",\"content\":\"ok\"}],\"max_tokens\":1,\"chat_template_kwargs\":{\"enable_thinking\":false}}" \ + | grep -q '"choices"' || exit 1 +ENTRYPOINT ["/opt/vllm/start.sh"] +CMD ["/models/model"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..ae60e95 --- /dev/null +++ b/README.md @@ -0,0 +1,256 @@ +# iHoree-vLLM-4xRTX3090 + +Swift-1.5-Qwen3.8-Flash-Next-W4A16 served from four RTX 3090 cards. The model is a 125B +MoE with 4-bit weights. It runs TP2xPP2 + EP with MTP K=3. The FP8 KV pool holds 924,993 +tokens at a 262,144-token context. Single-stream output measured at 158-184 tok/s. + +## 1. Requirements + +| Item | Requirement | +| ----------- | ------------------------------------------------------ | +| GPU | 4 free cards of 24 GB, all the same model | +| RAM | 96 GiB or more. The PLE table uses 69 GiB of host RAM. | +| `/dev/shm` | 1 GB or more | +| Driver | CUDA 13.0 or newer | +| Disk | 30 GB free for the build, 115 GB for the checkpoint | +| Python | 3.13 | +| Build tools | `cmake` 3.26 or newer, `ninja`, `cc`, `c++` | + +The checkpoint must contain `qsa_kv_scales_swift.json`. That file holds the FP8 KV +scales. Swift 1.5 changed `k_proj`/`v_proj`, so the scale file of one checkpoint does not +work for another. + +## 2. How to Build + +```bash +sudo apt-get install -y cmake ninja-build ccache +./build.sh +``` + +`build.sh` runs eight steps: the tools, the serving tree, the venv, the pinned runtime, +the CUDA toolchain, the compiled ops, the install, and the kernel pack. The first compile +takes 20 minutes on 64 cores CPU. + +Two variables control the compile: + +1. `MAX_JOBS` — the number of compile jobs. The default is `32`. Lower it when the build + runs out of RAM. +2. `ARCH` — the GPU architecture. The default is `8.6` (Ampere). + +The build prints every step. The installer runs with `-v`, so each compiler line is +visible, and one status line appears every 60 seconds. Full output goes to +`build-.log` in the repo root. + +### 2.1. Re-runs + +A re-run of an unchanged tree skips the compile and the install. It reads the tree commit +and the two extension hashes from `vllm/.build-products`, re-hashes +`vllm/.build-products.sha256`, and looks for the vllm install in the venv. Any mismatch +starts the compile. `./build.sh --force` always compiles. + +### 2.2. The tree + +Step 2 builds `vllm/` from the upstream tag v0.30.0 plus `patches/`, then compares the +result with the expected tree hash. A correct tree is left alone, so a re-run needs no +network. `FETCH=1 ./build.sh` re-applies the series from the base. +`VLLM_UPSTREAM_URL=` points at a mirror. + +### 2.3. Tools on the build machine + +`build.sh` uses `uv` for the venv when `uv` is on PATH, and CPython 3.13 in both cases. +Set `UV=0` to refuse `uv`, `UV=/path/to/uv` to select one. The wheel caches of the current +user (`$HOME/.cache/uv`, `$HOME/.cache/pip`) are printed at step 3 and reused. + +`build.sh` uses `ccache` when it is on PATH, as the CMake compiler launcher for C, C++, +CUDA and HIP. Objects go to the ccache directory of the current user. Step 1 prints the +directory, step 7 prints the hit and miss counts. Set `CCACHE_DIR=` to move them. Without +`ccache` every kernel compiles from scratch. + +`cargo` 1.95 is optional. With it, the build compiles the Rust frontend (`vllm-rs`, +`_rust_tool_parser`). Without it, the Python tool parsers are used. + +## 3. How to Run + +```bash +sudo ./start.sh +tail -f /var/log/vllm.log +``` + +`start.sh` defaults: `127.0.0.1:8000`, model name `qwen38-flash-next`, log +`/var/log/vllm.log`. The container and the systemd unit set `MODEL_NAME=Scarlet`. Use +`sudo` for the log file. Set `HOST=0.0.0.0` to serve on all interfaces. + +`start.sh` keeps the compile caches in `.vllm-cache/` in this directory. + +```bash +CACHE_ROOT=/path/to/cache ./start.sh +``` + +The first start compiles the cudagraphs into `$CACHE_ROOT` and takes about 6 minutes. The +first request builds the FlashInfer kernels. Each later start takes about 3 minutes. + +## 4. Run as a systemd service + +```bash +sudo scripts/install-as-service.sh install +sudo vi /etc/default/vllm +sudo scripts/install-as-service.sh start +sudo scripts/install-as-service.sh status +``` + +`install` rewrites `/opt/vllm` in `systemd/vllm.service` to the path of this repo, so the +service runs this repo in place. `update` runs `./build.sh` and restarts. `logs` tails the +engine log. `uninstall` removes the unit, the environment file, and the logrotate file. + +## 5. Logs + +`start.sh` writes the engine and the PLE offload child process to `/var/log/vllm.log`. +Set `LOG=/other/path` to move the file. Set `LOG=` (empty) to keep the output on the +terminal. The image uses that mode, so `docker compose logs` works. + +Install the rotation policy: + +```bash +sudo install -m 0644 systemd/vllm.logrotate /etc/logrotate.d/vllm +``` + +That command replaces an existing `/etc/logrotate.d/vllm` on the host. The policy rotates +the file daily, at 100 MB. It keeps 14 files, compresses them, and deletes logs older +than 60 days. The container caps its own log the same way (`json-file`, 100 MB x 14). + +## 6. Test the server + +`/v1/models` returns 200 when the engine is dead. Do a one-token generation: + +```bash +curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"Scarlet","messages":[{"role":"user","content":"hi"}],"max_tokens":1,\ + "chat_template_kwargs":{"enable_thinking":false}}' +``` + +A healthy start writes this line to the log: + +``` +GPU KV cache size: 924,993 tokens +``` + +The number changes with the KV memory setting. + +## 7. Container (optional) + +Use the container when the image is built on a box with no GPU and no driver, and runs on +the GPU host. + +Step 8 of `build.sh` writes the compiled kernels to `artifacts/`: a tarball, a per-file +sha256 list, and a manifest. `--no-pack` (`PACK=0`) skips that step. A second run of the +same tree does not pack again. + +Two build routes exist: + +```bash +# kernels from artifacts/, no compile, about 2 minutes +docker build --build-context ops="$PWD/artifacts" -t vllm:latest . + +# full compile, 20-40 minutes, 15 GB of build disk +docker build -t vllm:latest . +``` + +`docker compose build` passes the `ops` context from `artifacts/` by itself. + +The build checks the tarball before it installs it: the tree head, the GPU architecture, +the CPython minor, the CUDA major, the torch version, and the sha256 of every unpacked +file. A tarball that does not fit the tree stops the build and names the item that does +not match. + +The image build needs network access: it downloads vLLM v0.30.0 and the pinned wheels. The +base is Debian trixie. Kernels built on Debian 13 need glibc 2.38, and bookworm (glibc +2.36) cannot load them. `cmake` 3.25 in trixie is too old, so the image installs `cmake` +and `ninja` from pip. The image carries a compiler toolchain, so it is large. + +The image build takes the same two compile arguments as `build.sh`: + +```bash +docker build -t vllm:latest --build-arg MAX_JOBS=16 --build-arg ARCH=8.6 . +``` + +### 7.1. The CUDA driver stub + +The kernels link `CUDA::cuda_driver` (`libcuda.so.1`). The NVIDIA driver package provides +it. A build container with no GPU does not, and the CUDA wheel has no stub for it, so +step 5 runs `scripts/gen-cuda-stub.sh`. It writes a link-time stub next to the toolkit in +the venv: 548 entry points, no bodies, 92 KB, soname `libcuda.so.1`. No loader search path +includes that directory, so the extensions open the real driver at run time. Without the +stub, `cmake` stops with `Target "_C_stable_libtorch" links to: CUDA::cuda_driver but the +target was not found`. On the host the driver is present, so the step is a no-op after the +first run. + +The import check at the end of `build.sh` needs the driver too. Without a driver it checks +the link layer and says so. The full import runs at serve time on the GPU host. +`libcuda.so.1` is expected only in `vllm/_C_stable_libtorch.abi3.so`. The MoE kernels get +CUDA through `libtorch_cuda`, so it is not in their `NEEDED` list. + +### 7.2. Build caches + +The build reads three BuildKit cache mounts: the `uv` wheel cache, the `pip` cache, and +the `ccache` object cache. They live outside the image, so a rebuild downloads nothing and +recompiles only the kernels whose source changed. They start empty, so the first build +downloads about 6 GB of wheels. Pass the caches of the current user to prevent that: + +```bash +docker build -t vllm:latest \ + --build-context uvcache="$HOME/.cache/uv" \ + --build-context pipcache="$HOME/.cache/pip" . +``` + +The build reads those folders. It never writes to them. + +### 7.3. Run with Docker Compose + +```bash +sudo nvidia-ctk runtime configure --runtime=docker +sudo systemctl restart docker +cp .env.example .env +mkdir -p .vllm-cache "$HOME/.cache/uv" "$HOME/.cache/pip" +docker compose up -d +docker compose logs -f +``` + +`docker compose up` passes three contexts to the image build: `artifacts` (the kernels of +`./build.sh`), `uvcache` and `pipcache` (the wheel caches of this user). Each one must be a +folder that exists. `.vllm-cache` is the run-time cache folder, mounted at `/cache`. + +### 7.4. Run the image without Compose + +```bash +docker run -d --name vllm --gpus all --stop-timeout 20 --shm-size=16g \ + --ulimit memlock=-1 -p 8080:8000 \ + -v /data/models/checkpoint:/models/model:ro \ + -v "$PWD/.vllm-cache":/cache \ + -e MODEL_NAME=Scarlet vllm:latest +``` + +## 8. Files + +| File | Use | +| -------------------------------------------------- | --------------------------------------------------------------------------------- | +| `build.sh` | Fetch v0.30.0, apply `patches/`, build `.venv/` (py3.13, 196 pins), compile, pack | +| `start.sh` | Build on first run, test the box, cache in `.vllm-cache/`, serve, write the log | +| `scripts/install-as-service.sh` | Install the systemd unit for this repo, start, update, log, uninstall | +| `systemd/` | `vllm.service` and the logrotate policy | +| `Dockerfile`, `docker-compose.yml`, `.env.example` | The container path | +| `scripts/` | Topology selector, pinned requirements, CUDA driver stub | +| `patches/` | The 100 patches, in apply order | + +## 9. Variables + +| Variable | Default | Meaning | +| ------------------------- | ----------------------------------------------------- | -------------------------------------------------------------------------------- | +| `MODEL_DIR` | none, required | Host directory of the checkpoint. Compose stops with a message when it is empty. | +| `MODEL_NAME` | `Scarlet` (compose), `qwen38-flash-next` (`start.sh`) | The name in the `model` field of a request. Spaces make aliases. | +| `HOST_PORT` | `8080` | The host port of the API. | +| `VLLM_API_KEY` | empty | Empty means no authentication. | +| `CACHE_HOST_DIR` | `.vllm-cache` in this directory | Where the compile caches live. The host path uses the same folder. | +| `CUDA_VISIBLE_DEVICES` | `0,1,2,3` | The GPU lane, in bus order. | +| `PLE_HOME` | `0` | The card that hosts the PLE offload. That card gets no CUDA context for it. | +| `VLLM_LOG_STATS_INTERVAL` | `1` | Seconds between engine stats lines (throughput, KV use, queue). | +| `MAX_JOBS`, `ARCH` | `32`, `8.6` | Compile jobs and GPU architecture, for the build only. | diff --git a/abliteration/Qwen3.8-Flash-Next-Uncensored.json b/abliteration/Qwen3.8-Flash-Next-Uncensored.json new file mode 100644 index 0000000..57e872f --- /dev/null +++ b/abliteration/Qwen3.8-Flash-Next-Uncensored.json @@ -0,0 +1,2585 @@ +{ + "source": "orcarouter/Qwen3.8-Flash-Next-Uncensored", + "revision": "e096800036ec20da7e2442dcd4044a004d4e99fa", + "edit": "W' = W - r (r^T W) on every residual-writing matrix (Arditi et al., 2024), one direction", + "recovered_from": "top left singular vector of W - W' over layers.3.self_attn.o_proj.weight, layers.0.linear_attn.out_proj.weight, layers.1.ple.value_proj.weight, layers.3.mlp.shared_expert.down_proj.weight", + "verification": { + "untouched_control_identical": true, + "per_tensor_abs_cos": { + "layers.3.self_attn.o_proj.weight": 0.9999994039535522, + "layers.0.linear_attn.out_proj.weight": 0.999998927116394, + "layers.1.ple.value_proj.weight": 0.999992847442627, + "layers.3.mlp.shared_expert.down_proj.weight": 0.9999944567680359 + }, + "reproduction_bf16_exact_fraction": { + "layers.3.self_attn.o_proj.weight": 0.9196, + "layers.0.linear_attn.out_proj.weight": 0.9148, + "layers.1.ple.value_proj.weight": 0.8758, + "layers.3.mlp.shared_expert.down_proj.weight": 0.919, + "embed_tokens.weight": 0.9294 + } + }, + "hidden_size": 2560, + "refusal_direction": [ + 0.00778824557, + -0.0179212, + -0.0305097885, + -0.0119213471, + -0.0022670601, + 0.00862129778, + -0.0211653952, + 0.00968378689, + 0.0197913963, + 0.00261055119, + -0.0217173435, + -0.00690577319, + 0.01763626, + -0.0042252359, + -0.0150157074, + -0.0168697517, + 0.0114068622, + 0.0123851653, + 0.0181288552, + -0.00770292245, + 0.0126369856, + -0.00877867453, + 0.0147007741, + -0.00623441581, + 0.00213842816, + -0.0457281172, + -0.0145674963, + -0.00356830982, + -0.0366542004, + -0.0155591639, + -0.00653848331, + -0.000787768804, + 0.0273896828, + 0.000493242522, + 0.00742051564, + 0.0178993214, + -0.00229803636, + -0.0214590263, + 0.0150179984, + -0.00121425721, + -0.00495551573, + -0.0318871774, + 0.0222354438, + -0.0238917228, + -0.0072431541, + 0.0118668517, + 0.0037295362, + -0.0266498346, + 0.0187819824, + 0.0600446761, + -0.00416007079, + 0.000901914784, + 0.0324397162, + 0.0180206057, + 0.0179197043, + -0.0163361635, + 0.0314038582, + 0.0222291853, + -0.00553645892, + 0.00493013905, + -4.17987212e-05, + -0.00113755965, + -0.00194189104, + -0.0224357117, + 0.019951608, + -0.0121182138, + -0.00638864469, + -0.00342939282, + -0.0213679988, + 0.00875103008, + 0.00377325271, + -0.0485795625, + 0.0159405041, + 0.00784008391, + -0.00463976385, + 0.00637099799, + -0.00844945107, + -0.0350146219, + -0.00443035457, + 0.0174781736, + 0.00778687419, + -3.96960568e-05, + 0.0210599415, + -0.00513834739, + 0.0164750926, + -0.0398943536, + 0.0201373417, + 0.00677916594, + -0.0151018966, + -0.0192214735, + -0.0163175091, + -0.00366158038, + -0.0209532231, + 1.89838447e-05, + -0.0126473121, + 0.0255092345, + 0.00663942145, + -0.0146222888, + 0.0100685796, + 0.00880019087, + -0.00410559261, + -0.0204126239, + 0.0113125304, + -0.0187337641, + 0.00972247031, + -0.00183057005, + 0.0034544738, + 0.0224518217, + 0.0251727626, + 0.00501636509, + 0.0125011886, + -0.016659908, + -0.0271131676, + -0.0160052851, + -0.0116210487, + 0.0455624647, + 0.0433212779, + 0.00211333809, + -0.0261757951, + 0.00390288769, + 0.00555391051, + 0.0322331116, + -0.0201225374, + 0.00686775334, + 0.025862854, + -0.0573049933, + -0.00804869551, + 0.013937491, + 0.0372350886, + -0.0208566468, + -0.0240022559, + 0.0102949645, + -0.00606603362, + -0.00454741344, + 0.0384678468, + 0.0088311499, + -0.0356116369, + -0.019246567, + 0.0206673276, + -0.00259186327, + 0.00182135252, + 0.000343197491, + -0.0205936227, + -0.000983884791, + 0.0367069282, + -0.00190043717, + 6.42291561e-05, + -0.00816481188, + 0.00434623519, + -0.000398383534, + 0.0109299812, + -0.0319433995, + -0.00779894041, + 0.0245770346, + 0.00247827009, + 0.0209452454, + -0.00912398472, + -0.0172171649, + 0.0306357462, + -0.014298534, + -0.00648480607, + -0.00234087626, + 0.000374939555, + -0.00628142711, + -0.0359068476, + 0.0122702857, + -0.00202479702, + 0.0185293052, + 0.00135983527, + 0.00998229906, + 0.00419369759, + -0.00844809879, + -0.0169944726, + 0.00993560068, + 0.00366065255, + 0.0246076286, + 0.0169247407, + 0.0121988021, + 0.00855394173, + -0.0407930277, + 0.0188812409, + -0.00741096586, + -0.0255080666, + -0.0307783186, + -0.00563146034, + -0.003429709, + -0.014155183, + -0.0132334419, + -0.00860450417, + -0.0582965091, + -0.0177702084, + -0.0226770286, + -0.0134043582, + -0.0305991992, + 0.0259193983, + 0.00190539868, + 0.00993319694, + -0.00347520062, + -0.000772750238, + 0.0179445203, + -0.00074612617, + 0.000987915904, + 0.0135219004, + -0.00178813189, + 0.0130148688, + 0.000937938807, + 0.0208279695, + 0.033664953, + -0.000816787186, + 0.0104587432, + 0.014796597, + 0.0193465352, + -0.0133588444, + 0.00445240084, + -0.0386919826, + -0.000320548192, + -0.022084957, + -0.0409920886, + -0.00135060458, + 0.0115000904, + -0.01763523, + -0.0011270626, + 0.0349359214, + 0.00793287344, + 0.0185823403, + -0.00915568229, + -0.00608076667, + 0.0324320048, + -0.0235348754, + 0.00200697919, + 0.00705685141, + 0.000489792903, + 0.0195732303, + 0.0209155008, + 0.043402914, + -0.0127765965, + 0.00377741642, + 0.0212181136, + 0.0163654499, + -0.0151513498, + -0.0357327163, + 0.0255264863, + -0.0532403402, + -0.0539628714, + 0.0411220305, + -0.0330082662, + 0.00447659148, + -0.0120361038, + -0.000658104895, + -0.0164546445, + -0.00627186801, + -0.00545729464, + 0.0161849502, + -0.00467483513, + 0.0156061025, + -0.0126412092, + 0.00362999388, + -0.0128337527, + -0.0177632608, + -0.0242200401, + 0.00916429143, + 0.00519733177, + -0.0137904473, + 0.0190437529, + -0.0315996706, + -0.00378151913, + 0.0122085521, + 0.0340865925, + 0.0186760109, + -0.000298633939, + -0.0220262539, + 0.0325778946, + 0.0239037555, + -0.0149611682, + -1.83905631e-05, + 0.0104281763, + -0.0152445221, + 0.0103932759, + -0.00832775608, + 0.0129421409, + -0.0292653795, + -0.0057299654, + -0.00898024905, + 0.00885174703, + -0.00541666849, + -0.0122923348, + -0.030179834, + 0.00144312193, + -0.00449146843, + 0.0273828302, + -0.00247540767, + 0.00493750768, + -0.0342301056, + 0.00389027223, + 0.00470278971, + 0.0208306126, + 0.0169922672, + 0.00964672677, + -0.0165313333, + 0.00513123302, + -0.0366569012, + -0.00709539652, + -0.00707121985, + 0.0193672273, + -0.0199880786, + 0.00942825153, + 0.0287057161, + -0.0325578526, + 0.021687584, + -0.0132972552, + 0.0118634636, + -0.00819579698, + 0.012704148, + 0.0137626883, + -0.00569319539, + 0.0141837699, + -0.0128563168, + 0.00991761312, + 0.0130409626, + 0.0156751331, + -0.0119208023, + 0.0148922363, + -0.000292728248, + 0.008055415, + 0.0335148536, + -0.0287229978, + -0.0269464776, + -0.00879483391, + 0.0234260783, + 0.0050875009, + -0.0019566603, + -8.27288841e-06, + 0.0170335546, + 0.0325418226, + -0.00282089273, + -0.00165502192, + -0.00219521206, + 0.010194445, + 0.0375818536, + 0.00214832509, + 0.00295111467, + 0.0136863878, + -0.0158931185, + -0.0136997541, + -0.000786386488, + 0.00691575417, + 0.0232712403, + -0.0145344203, + 0.0376250185, + -0.00517537305, + -0.013260995, + 0.0192862842, + -0.00718658278, + 0.0182490405, + -0.0284414161, + 0.00761825778, + 0.0242309459, + 0.000833099592, + -0.0062557389, + -0.0138178142, + 6.82272866e-06, + -0.0120316716, + 7.33627849e-06, + 0.00103192381, + 0.00876135472, + -0.0103864791, + 0.0217217524, + 0.0135121783, + -0.0303642228, + 0.0143354116, + 0.0487660728, + 0.0429496877, + 0.0184317045, + 0.00993321277, + -0.000293585355, + -0.0136620672, + -0.0128828771, + 0.0167322196, + 0.0147432145, + 0.0113814706, + 0.0086607784, + -0.00893597957, + 0.0146420673, + -0.01271326, + 0.00909505133, + -0.0273715947, + -0.00648675766, + -0.0452292114, + 0.032860782, + 0.0405792929, + -0.00810147915, + -0.0324712135, + 0.00973907486, + 0.00481089996, + 0.0159408394, + 0.0155471498, + 0.016909508, + -0.0911883488, + -0.00575839076, + -0.000234813182, + -0.00327128475, + 0.0292806569, + 0.00118811126, + -0.0231723376, + -0.0460006334, + 0.00803198945, + -0.000735096459, + 0.0145135652, + -0.0103758695, + -0.00479316525, + 0.0032402745, + 0.011384896, + 0.0507203639, + -0.0324047096, + 0.0338708982, + -0.00172084582, + 0.0126504339, + 0.000114824303, + -0.00778398849, + 0.0198454019, + -0.0151864709, + -0.0194190573, + -0.0138832843, + 0.016587941, + -0.0167474411, + -0.00377168623, + 0.0270996857, + 0.00595871033, + -0.0124694817, + 0.012397524, + 0.0203264579, + -0.000223600407, + 0.00709206332, + 0.0012663512, + 0.015068897, + 0.00160503446, + 0.0243723262, + -0.00579050416, + 0.0189300701, + 0.0384670906, + -0.00506109418, + 0.0338318087, + 0.00263752067, + -0.0132585997, + 0.00787814613, + 0.0284796506, + 0.0138356294, + -0.00072915951, + 0.00703385193, + 0.0042528417, + -0.023407707, + -0.00284238392, + -0.0307151284, + -0.00531085487, + -0.0178020392, + -0.0232770983, + -0.0200572498, + 0.00554915238, + 0.0060290792, + 0.018280128, + -0.00491369236, + -1.88618524e-05, + 0.0195145011, + 0.0239858292, + -0.0158856418, + -0.0352481268, + 0.00755616743, + 0.0146311875, + -0.00828382839, + 0.0244149435, + 0.0132187614, + 2.9359082e-05, + 0.00568543468, + 0.000486087258, + -0.026823217, + 0.00721616251, + 0.0184820108, + 0.0104804961, + 0.0342431962, + -0.0262347981, + 0.0050670593, + 0.0213375371, + -0.00348183862, + -0.0182985291, + -0.000287236617, + -0.00655819895, + 0.00547669083, + 0.0206883997, + -0.0153958108, + -0.000138206044, + 0.00383574632, + 0.00451755244, + -0.0216899905, + 0.0585711747, + -0.0176201835, + 0.038267199, + 0.0233357195, + 0.0231133606, + -0.0208432656, + -0.0244772062, + 0.00352585386, + 0.00889849383, + -0.0103105474, + -0.0140116448, + 0.00715509662, + 0.0119615542, + -0.0209145918, + -0.0126101626, + -0.0243942849, + 0.00696068769, + 0.0117718996, + 0.0100202421, + 0.00808447506, + -0.0110346684, + -0.0291022696, + 0.00823486038, + 0.000419116026, + 0.00469753286, + -0.00806607585, + -0.00567220757, + -0.0127122924, + 0.0198442582, + 0.00878587924, + -0.0163010899, + -0.0238780417, + -0.00968337711, + 0.000668602588, + -0.0182044636, + 0.0245205835, + 0.00164855702, + 0.0174042303, + 0.00585107459, + 0.00227837474, + -0.00193953165, + -0.0380200408, + -0.0283678211, + 0.029456323, + -0.00614775531, + 0.0077220425, + -0.00376577931, + -0.00109424058, + 0.0119772367, + 0.00484524108, + 0.0457609966, + -0.0130962525, + -0.00923590641, + -0.0152387591, + 0.00083521544, + 0.0259066094, + -0.00530828536, + 0.00833082572, + 0.0413410515, + 0.0140853552, + 0.00539507298, + -0.00376170664, + 0.00673969928, + 0.0338779576, + -0.0167138521, + -0.00743735023, + 0.00792595278, + -0.0156908687, + 0.0483088456, + 0.00655525643, + -0.0436894298, + 0.00623033196, + -0.00210755831, + 0.00789809041, + -0.0251108035, + -0.0341087319, + 0.0166134909, + -0.0390389748, + 0.0178352017, + 0.00250968756, + 0.0128994118, + 0.0130905835, + -0.0157051589, + -0.0227479059, + -0.0222826097, + -0.0266634319, + -0.00237515033, + -0.00407021912, + 0.0375888646, + -0.006168033, + -0.0155536095, + -0.0245766565, + 0.0151373083, + -0.0163089652, + 0.00279039238, + 0.00440244609, + 0.0238977484, + -0.000151010187, + 0.00850967504, + 0.00257837167, + 0.0137852142, + 0.0331657603, + 0.00616607815, + 0.0112613682, + 0.00854510348, + 0.00859674998, + -0.000176703135, + -0.0432201326, + -0.0147050293, + -0.015334188, + 0.0282042716, + 0.00521853101, + -0.0372541845, + -0.0110468781, + -0.00678391103, + 0.0219626818, + -0.0229617171, + -0.0156218875, + -0.0363368243, + 0.00183405168, + 0.00322137517, + -0.011347414, + 0.0122601027, + 0.00683025364, + -0.0110092713, + -0.0158675872, + 0.0165882148, + 0.00902334973, + 0.0297549982, + -0.00485592429, + 0.0355401523, + 0.00424415572, + -0.0501568168, + 0.00636325032, + -8.49469143e-06, + 0.0124878092, + 0.0187594816, + -0.0156461578, + 0.00225745188, + -0.0168394092, + 0.0480999872, + -0.0108366618, + -0.0146940043, + -0.00852970965, + -0.00433708308, + -0.0484613962, + 0.0173588041, + 0.0125869671, + -0.0294449609, + 0.00882567279, + 0.0117290299, + -0.0196051076, + -0.00793318823, + 0.022314867, + 0.0255848411, + 0.035248097, + -0.0105755003, + 0.0106512383, + 0.018858172, + 0.012278216, + 0.00439775968, + -0.0025700829, + 0.0113050211, + -0.0112476368, + -0.00907701533, + 0.00017192385, + -0.0131514529, + -0.00586467143, + -0.0414103493, + 0.0315658338, + -0.0119076604, + -0.00101747911, + 0.0359271541, + -0.0147881955, + 0.0156480279, + -0.0215097871, + -0.011699968, + 0.0289418623, + 0.0157884117, + -0.00485027861, + -0.00174145715, + -0.0115681496, + 0.0132120019, + -0.00185940426, + -0.0172081925, + -0.0124434959, + 0.0162207913, + 0.0243852288, + -0.0163423028, + 5.6309651e-05, + 0.0160494782, + -0.0369821042, + -0.0214567874, + 0.000578387291, + -0.00318473997, + -0.0310446769, + -0.0158394445, + 0.0115517052, + 0.00698086899, + 0.0401695967, + 0.00253666891, + -0.0198460333, + 0.00277170306, + -0.00132778403, + 0.00176540704, + 0.0139447702, + -0.00893759076, + 0.0112977624, + 0.0175007675, + 0.00967315678, + 0.0386817604, + 0.0296532959, + 0.00453094346, + -0.0168533996, + 0.0380585119, + 0.0163169075, + 0.0078696264, + 0.00494842604, + -0.00262085604, + 0.0140128247, + 0.00541240582, + 0.0153068341, + 0.000858439889, + 0.0029715905, + 0.0116471117, + 0.000398009259, + 0.0111940559, + 0.00882096123, + 0.00667003868, + -0.00686244853, + -0.00819924474, + 0.0151200574, + -0.0288088582, + -3.32516829e-05, + -0.00626594294, + -0.0269687511, + 0.0291527081, + 0.0125773503, + 0.014638463, + -0.0182186831, + -0.0181354247, + 0.0351730622, + 0.0209150575, + 0.018668415, + -0.0239623338, + 0.0133559341, + -0.0309293624, + -0.0605476983, + 0.0101352772, + -0.0111750606, + 0.00294730626, + 0.00713484455, + 0.0134696234, + 0.00338734547, + -0.0110662468, + 0.00626597693, + -0.0193690918, + -0.00945949182, + -0.025882842, + -0.00846042763, + 0.0344970897, + -0.0161277484, + -0.00122669619, + -0.0301974732, + 0.0119154416, + 0.0236551147, + 0.0163211655, + -0.0265934691, + -0.00996167865, + -0.0129338838, + 0.0199156851, + 0.0118430527, + -0.0259716623, + -0.0132742086, + -0.0058752019, + 0.0104160197, + 0.0322893038, + -0.0142302942, + -0.00391290477, + -0.0264496617, + 0.0185955055, + 0.0382148772, + 0.0125256022, + -0.0119556431, + -0.00485453336, + 0.00204763329, + -0.0181140285, + -0.00993418321, + -0.0101861013, + 0.0144164888, + -6.26729234e-05, + -0.00994983688, + -4.22952289e-05, + 0.0159328114, + -0.0193283726, + 0.0377221666, + -0.020757528, + 0.014207378, + 0.0151642794, + 0.00269693974, + -0.0238934346, + -0.010408462, + 0.031886667, + 0.00267617498, + -0.00846375804, + 0.0302290879, + 0.00431855116, + 0.0249236748, + 0.00594865344, + -0.00715259602, + 0.0116626183, + 0.0111790951, + 0.00392394653, + 0.000519743946, + 0.0191435162, + -0.0165033061, + -0.0297018383, + -0.0403563567, + -0.0221043546, + -0.00600043172, + 0.0157200843, + 0.0017317998, + 0.0100473268, + 0.0122216437, + -0.0287264865, + 0.0336242802, + 0.0696345717, + -0.0225427076, + -0.00617852807, + 0.0125227552, + 0.0387321375, + -0.00720330188, + -0.0125281233, + 0.0156896189, + -0.00734819937, + 0.0301472563, + -0.000227952405, + -0.0125650754, + 0.00417395262, + -0.0391439274, + -0.00260267872, + 0.0315163583, + 0.000618004415, + -0.0127551341, + 0.0178121515, + 0.015034426, + -0.0190465599, + 0.00914745033, + -0.0476373769, + -0.0429784544, + -0.0156278592, + -0.0169154275, + 0.0168675315, + -0.00605131174, + 0.0276167318, + -0.0273631606, + -0.0105935661, + -0.0376418382, + -0.0265888833, + -0.0197782964, + 0.00400566589, + -0.0177973174, + -0.0305387378, + 0.0125105977, + 0.0361365788, + 0.0080429595, + 0.00281161326, + -0.0153671354, + -0.00555275101, + -0.00277625211, + 0.0057659205, + -0.0142334448, + 1.06841471e-05, + 0.00680892356, + -0.0101960637, + 0.00277170213, + -0.0319888555, + -0.0192563105, + 0.0215397011, + 0.0205199569, + 0.00194342237, + -0.0132353846, + -0.0133539438, + -0.0156972446, + 0.0141389836, + 0.016611021, + -0.0390969291, + 0.0103661129, + 0.000839282642, + 0.00236031692, + 0.00414509093, + -0.000697325158, + -0.0104849767, + 0.00211309898, + 0.0392863117, + 0.023247147, + 0.00331745925, + -0.0155434003, + 0.0265649911, + -0.0257637631, + -0.0294305868, + -0.0156273339, + 0.00400985125, + -0.00201033754, + -0.0288465787, + 0.0227376651, + 0.0111506749, + -0.00509859901, + -0.0203533545, + 0.00343268691, + 0.00254930835, + -0.0500968024, + 0.0796442702, + 0.0342778265, + -0.0104678646, + 0.0187993888, + 0.0153441904, + 0.0117793418, + -0.0129751367, + -0.0461854041, + 0.0148696583, + -0.000932714436, + 0.0361824557, + 0.0292840414, + 0.00520155206, + -0.00487284921, + -0.0124312304, + 0.0156251248, + -0.0166124329, + -0.0250711739, + -0.0307549536, + 0.00174867013, + -0.0319428854, + -0.0175107159, + -0.00254422938, + -0.0102927154, + -0.0117275566, + 0.00224533654, + -0.0126631996, + -0.0291530043, + 0.0403363071, + -0.0326024517, + -0.0217851698, + -0.0113693131, + -0.00578815071, + 0.0106447823, + -0.0213985108, + -0.0245102793, + -0.0018680424, + -0.0180689432, + 0.0058012465, + 0.011124053, + -0.0383808948, + 0.0111039886, + -0.0247653779, + 0.027159445, + -0.0271867123, + 0.0248758923, + 0.017114602, + 0.00831646286, + -0.021505747, + 0.0127076693, + -0.0209488235, + 0.0201410763, + -0.0265414473, + 0.00683704019, + 0.00218998198, + -0.0197129026, + 0.0158121735, + 0.0283554532, + -0.0113272043, + -0.00196486432, + -0.00709762005, + -0.00913269073, + -0.0317578465, + -0.0417465717, + -0.0111215143, + -0.0167606883, + 0.00925014634, + -0.00114181545, + -0.00969229359, + 0.0113248723, + -0.00522523234, + 0.0122068189, + 0.0105815828, + 0.023099808, + -0.0541794598, + -0.00303434045, + 0.000959910976, + -0.0248070676, + 0.0111261262, + 0.00500396499, + 0.01320004, + 0.0357542299, + -0.000243666072, + -0.0386805944, + 0.0284721162, + -0.0269155055, + 0.0114149638, + 0.00526105752, + 0.00359860878, + 3.31434057e-05, + -0.00587498676, + 0.0239513461, + 0.00980760343, + -0.000168854313, + -0.0123799331, + -0.000589920091, + 0.00489285635, + 0.0358644128, + -0.0156753641, + 0.00948356371, + -0.0129329767, + 0.0201357324, + -0.0191240683, + 0.0668775886, + 0.0260272399, + -0.0310755838, + 0.011477422, + 0.00201384141, + -0.00649011135, + 0.0100980699, + 0.00337998173, + -0.00285211485, + 0.0074416711, + -0.00667320238, + -0.0143386479, + -0.00724799698, + -0.017383296, + 0.0260180887, + 0.00748377852, + 0.00580395898, + 0.0217529405, + -0.0174976923, + 0.0220772829, + -0.00380681735, + -0.0331146643, + -0.0444147401, + -0.0201010201, + 0.0160685163, + 0.0119542573, + 0.0315752774, + -0.00770886801, + 0.0148401419, + 0.000215798937, + -0.000610989402, + 0.0306934845, + -0.0191944595, + -0.0225257408, + -0.0297550838, + 0.00336119812, + -0.00489860587, + -0.00497576594, + 0.00371204619, + 0.010945552, + -0.0044404068, + 0.0219860002, + 0.00242089294, + 0.0242671557, + -0.00437887944, + 0.0165635031, + -0.00513889082, + -0.0153917484, + 0.0153818065, + -0.00965821091, + -0.0147417048, + 0.0104567381, + 0.0191591959, + -0.0111093139, + -0.0180136021, + 0.00940249301, + -0.000389805879, + 0.00325056142, + 0.0139347855, + -0.0217296891, + -0.0612015575, + 0.0131206838, + -0.0145194931, + 0.00824606977, + 0.00327913556, + 0.0326240286, + -0.00484249601, + 0.0222838372, + -0.0195964426, + -0.0254923571, + -0.00432641571, + -0.00202655909, + -0.0146962227, + 0.0230452772, + -0.030323185, + -0.0167418197, + -0.0041196337, + -0.00135717168, + 0.0402864255, + 0.00619517872, + 0.00464242697, + -0.00724533852, + 0.0226896685, + 0.0303886216, + -0.0378448218, + 0.0188319813, + 0.00427467888, + -0.023670014, + 0.026901206, + -0.0115966853, + 0.0123732165, + -0.0185926501, + 0.00879404321, + -0.0150335534, + -0.00734251598, + -0.0268705357, + 0.0102177747, + -0.000903439417, + -0.0102169076, + -0.00876943115, + 0.00903756451, + -0.000296910759, + 0.00674113398, + -1.74875697e-06, + 0.0061434526, + 0.00754703069, + 0.01033442, + 0.00849086232, + 0.00781512633, + -0.0104549173, + 0.0122402543, + -0.0515390448, + 0.00429074373, + 0.0279706679, + -0.00136669015, + 0.00209365436, + -0.00776285212, + 0.0215806998, + -0.00368739106, + 0.0057186638, + -0.0232033245, + 0.0628363788, + 0.0114322742, + -0.0184277743, + -0.0103635462, + -0.0495286509, + -0.00966549292, + -0.0252912678, + 0.00360277435, + -0.000649289927, + 0.0176090058, + 0.0104506584, + -0.0203229897, + -0.0154836476, + 0.0165588129, + -0.0075415154, + 0.00670297258, + 0.0207327306, + 0.0165084172, + -0.0185863338, + -0.0130185876, + -0.00977755152, + 0.0338260531, + -0.0260437503, + -0.0158771984, + 0.0163670424, + -0.0169806182, + 0.00717360713, + -0.0602866709, + 0.016259158, + -0.00206156331, + 0.027667908, + -0.0227961652, + -0.0166493934, + 0.0130008999, + 0.0108276475, + 0.0328385271, + 0.0194758661, + 0.0103189358, + -0.0153338648, + -0.0183077175, + 0.000288297859, + -0.00984852109, + -0.0045446679, + 0.0243105553, + 0.00150238862, + 0.00830333494, + -0.0212109406, + 0.00125788688, + 0.0218449589, + 0.0229833927, + 0.00975167938, + 0.0208662283, + 0.0168948378, + -0.0115427617, + -0.00993771292, + 0.0155452956, + 0.0172808897, + 0.0317507796, + -0.0225390177, + 0.00248641707, + 0.0118587418, + 0.0333766006, + -0.00759203173, + -0.0234009661, + 0.000130812827, + 0.0194324553, + 0.00889752991, + 0.0307658557, + -0.0412877798, + 0.00820037816, + 0.00453241216, + -0.0537768081, + 0.0274456553, + -0.0135557698, + -0.0201411732, + 0.0215580072, + 0.00232185097, + -0.0264286008, + 0.0235494506, + -0.0548274815, + -0.0289721061, + -0.00456859032, + 0.00918691698, + -0.0109802587, + 0.0139045315, + -0.0150922975, + -0.0329528786, + -0.00566658657, + -0.0169787202, + -0.0113152508, + -0.0161603056, + 0.012220826, + -0.00358014787, + 0.00797195733, + 0.00580615131, + 0.0108337998, + -0.00365217449, + 0.00589196989, + -0.0094978977, + 0.0288736038, + -0.00339658302, + -0.0121006314, + 0.00167161727, + 0.0110367564, + 0.0192737039, + -0.0127491634, + -0.00752970995, + 0.00739216339, + -0.00319648837, + -0.00122290326, + 0.0160591491, + -0.0170696024, + -0.00559021579, + -0.0255026538, + -0.0201104712, + 0.0290123522, + -0.0309946444, + 0.0102130426, + -0.0158427469, + -0.00614870246, + 0.0178435948, + 0.0367564298, + -0.012349084, + -0.000348147034, + -0.000936269236, + -0.0362486616, + 0.0193568449, + -0.00774996169, + 0.019902071, + -0.0149526307, + 0.0241917148, + 0.00510807335, + 0.00568151567, + -0.00636344403, + 0.0369482078, + 0.0252962187, + -0.0273251235, + -0.00100268261, + 0.00246891472, + 0.0186008513, + -0.0275057275, + -0.0278487392, + 0.00817203242, + 0.044567991, + -0.0095096482, + -0.0267214086, + 1.19663364e-05, + 0.00286448607, + 0.0336348787, + 0.00989954919, + 8.47299361e-06, + -0.011914202, + 0.00340634165, + 0.0151144788, + 0.00175554736, + 0.0400227681, + -0.0516542494, + 0.00205623778, + 0.0354667082, + -0.0104835797, + -0.0337458141, + 0.00439226301, + -0.0158011336, + 0.0152606303, + 0.0332360268, + -0.00554388063, + 0.0199936945, + 0.0195070188, + 0.00201208959, + 0.0253881495, + 0.0106970044, + 0.00595776783, + -0.0113601601, + -0.0230234992, + 0.00591414794, + -0.0335237272, + 0.0075016492, + 0.00191913277, + 0.00466314238, + 0.00756245013, + 0.0355311781, + 0.00333887408, + 0.00377109111, + 0.000506586686, + -0.0189472809, + -0.00536439475, + 1.48912295e-06, + -0.0192302037, + -0.0196268391, + -0.0194256417, + 0.0311323553, + -0.0110235978, + 0.00467333291, + -0.0268177502, + 0.00430107629, + 0.010884189, + 0.0406222492, + -0.0157256704, + 0.000212739586, + -0.012854306, + -0.0149171036, + -0.00397405168, + -0.00319897709, + -0.00449782936, + -0.00213048118, + -0.00953902211, + -0.00220930274, + 0.000930200971, + 0.0391094349, + 0.0160517637, + -0.00226470572, + -0.0154486699, + -0.0011174609, + 0.0159075949, + -0.00811713282, + -0.00977838691, + -0.012352665, + 0.0325597785, + -0.0269435942, + 0.018827511, + -0.0533143505, + -0.0393472202, + 0.017951237, + 0.0237495489, + 0.00356017589, + 0.0101188542, + 0.0186720602, + -0.00705614919, + -0.00605097646, + -0.00371913635, + -0.017025819, + -0.00461530918, + -0.0149535351, + 0.00404863572, + 0.00925846677, + -0.0206161905, + -0.0185928345, + 0.0310443491, + -5.1025796e-05, + 5.71712781e-06, + 0.00120536098, + 0.01515825, + -0.0136371376, + -0.0252120253, + -0.00652324082, + -0.0064437273, + -0.0267053135, + 0.00851766113, + 0.00604037754, + 0.013474768, + -0.0170925464, + -0.0205419939, + -0.00588467624, + -0.0270842668, + 0.000758685928, + 0.0450688712, + -0.00200473564, + 0.0277228467, + 0.0141828638, + 0.0149516212, + -0.0072647552, + 0.0118968096, + 0.00381566049, + 0.00443349313, + 0.010094434, + -0.0258053392, + 0.00056801294, + 0.0210538581, + -0.0175965019, + -0.0112504186, + 0.0254499856, + -0.00624827063, + -0.0135259936, + 0.014557221, + -0.00330773252, + 0.0203120131, + -0.0130102662, + -0.00902269129, + -0.0139684016, + 0.0262963772, + -0.00280471169, + -0.017531909, + 0.0117317, + -0.0100805024, + 0.0227738675, + -0.0143315839, + -0.0338071585, + 0.0125923492, + 0.00443342654, + 0.0127521288, + -0.027490776, + -0.0178009253, + -0.0158257894, + 0.00218338892, + -0.000141145618, + 0.00978065282, + 0.0109429788, + 0.00151188998, + -0.000944046769, + -0.0165720321, + 0.00540925097, + -0.0213037729, + -0.0140315639, + 0.0427170061, + 0.00411042245, + -0.0138634853, + -0.0169944596, + -0.00821048673, + 0.0245543458, + -0.0241134949, + -0.0171353575, + -0.000729980937, + 0.0193742234, + 0.02277348, + -0.0189340524, + -0.0263267141, + -0.000702497724, + -0.0227649063, + 0.0330684409, + -0.000741882774, + -0.0146410763, + 0.0350321196, + -0.027235847, + -0.00610855548, + -0.0220357664, + 0.0237633586, + -0.0266586319, + -0.0141768502, + -0.0314054303, + -0.00860650185, + 0.0122567704, + -0.0334214382, + -0.0213649273, + 0.00516136689, + 0.000238020497, + 0.003485166, + -0.00262423721, + -0.0309835169, + 0.00815297663, + 0.009110271, + 0.0234414656, + -0.0358750001, + -0.0287198294, + 0.0277722478, + -0.0310440511, + -0.0162900053, + -0.00381882931, + -0.00279545342, + -0.00669522118, + -0.000222771181, + -0.00583080342, + -0.0288283899, + 0.0271682199, + -0.00390319759, + -0.000519873982, + -0.00701544946, + -0.0207895637, + -0.0119755082, + -0.0193912741, + -0.013127205, + 0.0255529638, + -0.0232632086, + 0.00868606195, + 0.00183414528, + -0.00550884428, + 0.00935187191, + -0.0243833624, + -0.0194670707, + -0.0301981494, + 0.0174966343, + -0.0136359371, + 0.016071137, + 0.0318688676, + 0.0148348231, + 0.0180934779, + 0.0226872228, + 0.00589427212, + 0.0115346033, + 0.00296374061, + 0.00490372302, + 0.0142128672, + 0.0486145131, + 0.00577736646, + -0.000183105847, + -0.00325425505, + -0.011427341, + -0.0218597353, + -0.0368905328, + -1.41612318e-05, + -0.00551344454, + 0.0326404609, + 0.00936456397, + -0.0186431799, + 0.0183726717, + -0.0265034735, + -0.00200761529, + 0.00338060618, + 0.0230544638, + -0.00174799073, + 0.00711019151, + -0.0201129634, + 0.0551079623, + -0.0086257467, + 0.0460295193, + 0.0299271476, + 0.00871895719, + 0.0450115576, + -0.0568501838, + 0.00593717629, + 0.00487416424, + -0.00237511448, + 0.0130245453, + 0.000204156895, + -0.0231890604, + 0.0219489355, + 0.00680781854, + -0.0290036518, + -0.0136543978, + -0.0109967981, + -0.00915244967, + -0.0118814008, + -0.0111599313, + -0.0311779454, + -0.0155657418, + -0.0052584256, + -0.00900333375, + -0.0315465853, + -0.0140951108, + 0.0115707852, + 0.00803498924, + -0.00948790647, + -0.0170644037, + -0.0363803208, + -0.0437426008, + 0.000353289448, + -0.0348698832, + 0.0145947756, + -0.0386044011, + -0.0158929955, + 0.0054798061, + 0.0285848454, + -0.0431582332, + 0.0212228931, + -0.0117665883, + 0.00897829235, + -0.00450792676, + 0.00526457047, + -0.0522559099, + -0.000720858981, + 0.000826544303, + -0.0088606067, + 0.0110575147, + 0.0180615634, + 0.0270386394, + 0.0229950231, + 1.1674224e-05, + -0.0370233841, + 0.0149294883, + 0.000809560064, + -0.0186911505, + -0.0127508389, + 0.0234358683, + 0.0250036884, + -0.00580105698, + -0.0142164649, + -0.0186141804, + 0.00277889799, + 0.0168701522, + 0.000979411765, + 0.00688006589, + 0.0349513367, + 1.88935373e-05, + 0.00322579662, + -0.0244170595, + 0.00645044213, + -0.0210079439, + -0.0238261186, + 0.000799255853, + -0.010531188, + -0.0144351479, + -0.000348990958, + 0.00982959382, + -0.0151491696, + 0.00190470228, + -0.0268751979, + 0.00919894874, + -0.00275218999, + 0.0244754367, + 0.0312333256, + 0.0249278173, + -0.0176154003, + -0.00823743455, + -0.0224147756, + 0.000197797403, + -0.0155107956, + -0.0164469741, + -0.000979223289, + 0.020628348, + 0.0156946275, + 0.00956862047, + 0.00691717863, + -0.0261659864, + 0.0362476744, + -0.0133134797, + 0.0346876755, + -0.00258546998, + 0.0302185249, + 0.0150700286, + 0.00909159146, + -0.019368669, + -0.00727884471, + 0.00795844849, + -0.00580516877, + -0.00720264157, + 0.0114218127, + -0.00710779428, + -0.0220480245, + 0.00118952245, + 6.78497327e-06, + 0.0128725301, + -0.00219991547, + 0.028984772, + 0.0288941637, + -0.00305109471, + -0.017132001, + -0.0245370064, + 5.76064778e-08, + 0.0198927093, + -0.0150212357, + 0.0106900204, + 0.0161427502, + 0.0171941686, + 0.00480648223, + 0.0219532214, + -0.00272987713, + 0.032534264, + 0.0102065364, + -0.0181265343, + 0.0241420567, + 0.0012520944, + -0.0159116108, + -0.00694464426, + -0.0122266505, + 0.00715983706, + 0.0191806685, + -0.0030444616, + -0.0278670043, + -0.0085879676, + 0.0114958677, + -0.0419887602, + -0.00608242955, + 0.00701676914, + -0.0261669829, + 0.00664024241, + 0.00050835422, + -0.0254024789, + -0.0172069538, + -0.00341495778, + 0.0111827245, + 0.000570094562, + -0.00829606317, + 0.0102107655, + 0.021511158, + 0.0203480553, + 0.0349824056, + -0.019133281, + -0.00589621533, + 0.0263151042, + 0.0168244485, + 0.0481507331, + -0.0170059297, + -0.0138823511, + -0.0133955125, + 0.0168011766, + -0.0051845056, + 0.0177395921, + 0.000277256506, + 0.00218153442, + -0.0507609546, + 0.00731864199, + 0.00803266931, + 0.0376626477, + 0.0267257951, + -0.000495643937, + -0.0367160849, + -0.0241006743, + -0.0182344858, + 0.0181388296, + 0.0233572815, + 0.00587336952, + 0.0175740924, + 0.00351622794, + -0.0115680546, + 0.0253590625, + -0.00516910758, + 0.00869768672, + 0.0292214584, + -0.0327395312, + -0.00469698943, + 0.0205544159, + -0.0117934048, + -0.00681134453, + -0.00752956141, + 0.0277237725, + 0.0211009923, + 0.0183628853, + -0.0285992455, + 0.02910023, + 0.0222603139, + -0.00474332785, + 0.017370984, + 0.000101307473, + 0.0124435145, + -0.00427491963, + -0.0203251354, + 0.0236772653, + 0.0220080968, + -0.00322084455, + -0.0162628964, + 0.0300213937, + -0.0217015222, + -0.00147278176, + -0.00617637718, + 0.0122857792, + 0.0135122929, + -0.0302431397, + -0.0202395283, + 0.000510516053, + 0.00795487035, + 0.00220690155, + 0.0147187347, + 0.0129054198, + 0.0137548782, + -0.039708551, + 0.0176434647, + -7.98204419e-05, + 0.0363604538, + 0.0236770455, + -0.0401231125, + -0.0321236812, + 0.00754894642, + -0.0214553364, + -0.00862966757, + -0.022197783, + 0.00653188396, + 0.00084476755, + 0.0143361064, + 0.00681298273, + -0.00170873478, + -0.0163308717, + -0.022938149, + 0.0364323854, + 0.000850175158, + 0.00411122898, + -0.00216881535, + -0.0487265103, + 0.0155675579, + 0.0024357806, + -0.00982096791, + 0.0238218438, + -0.0111266393, + -0.0312206075, + -0.00595528167, + -0.0260385126, + -0.0112374444, + 0.0186035056, + -0.0064301528, + 0.037051104, + -0.0142689804, + -0.01547137, + -0.0328171961, + -0.00918428879, + 0.00439557433, + 0.0453753285, + 0.0243148413, + -0.0195212383, + 0.0162963383, + 0.0115552647, + -0.029302245, + 0.0216496512, + -0.0212724376, + 0.00743853161, + -0.0467271768, + 0.0056880773, + 0.00645079557, + -0.0205159727, + 0.0126594622, + 0.0149293179, + -0.0306574218, + -0.00801861007, + -0.0171096399, + -0.00109556527, + -0.000702601275, + 0.00448582089, + -0.0147636253, + -0.000778379908, + -0.0477513522, + 0.0013288426, + -0.0191424713, + -0.00619662367, + 0.0299361199, + 4.34487083e-06, + -0.00596280117, + -0.0171834752, + -0.0187781416, + 0.0159865841, + 0.000900184386, + -0.00382752693, + 0.0142091233, + -0.0179018602, + -0.0338776708, + 0.0202986132, + -0.0157163497, + 0.000276450795, + 0.000354344054, + 0.0322932191, + 0.0106875934, + -0.0240401626, + 0.0071537192, + -0.0238697138, + -0.00881536305, + -0.0223263633, + -0.0222879015, + -0.0161821693, + 0.00586527213, + -0.0123719061, + -0.0128240539, + -0.000190103528, + 0.00473943166, + -0.0316577815, + 0.0306032449, + 0.0205930341, + -0.0129542826, + -0.0160733759, + 0.0220916923, + 0.00286820228, + 0.0192238111, + -0.0198920649, + -0.0201076958, + 0.0225301441, + 0.00905591715, + -0.0143820345, + 0.01042285, + 0.0198550131, + 0.0390647799, + 0.0239359885, + 0.0169433076, + -0.00449640118, + 0.016184913, + 0.00722541753, + -0.0111927241, + 0.0156329982, + 0.00996720791, + 0.0147747993, + -0.0155191617, + 0.0031300995, + 0.00211533462, + -0.0394954905, + -0.0303258505, + 0.00572144426, + -0.00526556978, + 0.0235256776, + 0.0220766403, + -0.0110717472, + -0.000275267259, + 0.0207885951, + 0.002138315, + 0.00799001195, + -0.00656761741, + 0.0190925989, + -0.0391630493, + -0.00464511104, + -0.0464778282, + 0.0326096192, + 0.00342292991, + 0.0127509544, + -0.0133802127, + -8.52616358e-05, + -0.00459918706, + -0.00524076819, + -0.0174205266, + -0.019839704, + 0.000429942273, + 0.0444099195, + 0.0147558348, + 0.0138884252, + -0.018848557, + 0.00941908639, + -0.0393449068, + -0.000525638985, + 0.0116512263, + -0.00490546925, + -0.0154096056, + 0.0222782791, + -0.00359394844, + -0.0135477185, + -0.00167591148, + 0.00690773129, + 0.00158139772, + -0.0100478418, + 0.013020535, + -0.00770282419, + -0.00622180756, + 0.00904415268, + 0.000259898225, + 0.01429673, + 5.23198723e-05, + 0.0159584712, + 0.0181316603, + -0.000462544762, + -0.00763021875, + -0.0269360188, + 0.0170022417, + 0.0135495421, + -0.0140117463, + -0.0278653484, + -0.0413883738, + -0.0103063602, + -0.0173731707, + 0.00804771297, + -0.0197193138, + 0.0167673603, + -0.0422305129, + -0.00464284373, + 0.0108406143, + -0.00682060095, + 0.0103603909, + 0.0102979131, + -0.0202941336, + -0.0257364102, + -0.005288688, + -0.0204506293, + -0.0066199312, + -1.82169206e-05, + -0.0119985733, + -0.0132884504, + -0.0186463706, + -0.00124506967, + -0.00797088072, + -0.00934728328, + 0.0141738551, + 0.00861098524, + -0.000792719075, + 0.0123487953, + 0.0889067873, + -0.0429250002, + -0.0229480974, + -0.0528939366, + -0.0414899848, + -0.0519795269, + 0.0190901831, + 0.036962159, + 0.00924933236, + 0.00565350009, + -0.0117763951, + 0.0255014822, + -0.0159879047, + -0.0177765042, + -0.0105987843, + 0.00917190127, + -0.000552416372, + -0.018774718, + 0.0237082485, + -0.00214127009, + -0.00984844286, + 0.0135847442, + 0.0157290921, + 0.00313162338, + -0.0142154964, + -0.0176190566, + 0.0258834474, + -0.0358586498, + 0.00697416766, + -0.0151771884, + 0.0085373437, + 0.0184301082, + 0.0212270115, + 0.0238398779, + -0.00937704649, + -0.0122106774, + 0.0313095786, + 0.0134126078, + -0.0171158649, + -0.0110929292, + -0.0116342222, + 0.00885303132, + -0.00359530491, + -0.0315446407, + -0.0272425022, + 0.000518404704, + -0.0141947232, + -0.00709159439, + 0.0453099497, + -0.0183558743, + 0.0182985943, + 0.00175378716, + -0.0333480872, + -0.000266312156, + 0.054337088, + 0.0129820937, + 0.0258247983, + 0.0417794846, + -0.0131292911, + 0.045948185, + 0.0103206458, + 0.0244567972, + -0.0113952039, + -0.00597937265, + -0.00312772649, + -0.0194787662, + -0.0119872848, + -0.00791242439, + 0.00850609783, + -0.0189059488, + -1.21811872e-05, + -0.0249490011, + -0.00838330388, + 0.0124718677, + 0.0249618329, + 0.0321419723, + 0.00237446767, + -0.0111509366, + -0.0118782306, + 0.0314546451, + -0.0418571346, + -0.0286162738, + -0.0163465533, + -0.0122715496, + 0.0164539106, + 0.0136678424, + 0.00155580323, + -0.00954026449, + -0.00502601732, + -0.00859659631, + 0.00643943483, + 0.00107925001, + -0.00419907132, + 0.00901765283, + -0.0192518476, + -0.00408917805, + -0.000836310093, + -0.00447866367, + 0.00826901477, + 0.0154900309, + -0.00249085203, + -0.0124649042, + -0.0348952636, + 0.0105730165, + 0.016034605, + 0.00368765788, + 0.0247054268, + -0.0127867227, + 0.0170123298, + -0.0488788895, + -0.0324487798, + -0.00632505631, + -0.039576523, + -0.000395886542, + 0.0152087305, + -0.0127114067, + -0.0287796259, + -0.0165698063, + 0.0284891333, + -0.00779975206, + 0.00968675036, + -0.0251137484, + -0.0265162922, + -0.0219142269, + -0.00195961352, + -0.0265266523, + 0.0181072, + -0.017404791, + -0.00061015971, + -0.0485145561, + 0.0301172957, + 0.0188059695, + -0.0107487924, + 0.0373527333, + 0.0244904701, + -0.010049588, + 0.0133015988, + 0.0368917622, + 0.0368858799, + 0.0305175371, + -0.0403130613, + -0.00434610527, + -0.00891293492, + 0.0237473473, + -0.0142777665, + 0.00834525097, + -0.00132041343, + 0.0252526179, + 0.00434480561, + 0.0241047349, + 0.00146224978, + -0.00359653239, + -0.000122285594, + 0.0424193703, + 3.29782824e-06, + 0.00824359711, + -0.00772723043, + 0.000621493149, + 0.00470097503, + 0.00469932659, + 0.0365372449, + 0.0121361064, + 0.0371111743, + -0.0462608971, + 0.0102582034, + -0.0187897179, + -0.0566767044, + 0.0113648623, + -0.0314545967, + 0.00291987578, + 0.0184124038, + -0.0516133197, + 0.0148845026, + -0.00409632875, + 0.000794200634, + -0.00989020895, + 0.0257795043, + 0.0208136458, + -0.0147103518, + 0.0020474433, + 0.00613299059, + 0.0140527375, + 0.0373542458, + 0.0212000255, + 0.000381644204, + -0.0177599695, + 0.010831208, + 0.0117779076, + 0.000623403583, + -0.00555548631, + 0.0207209904, + -0.0051680035, + -0.000699809752, + 0.0128836622, + 0.0114596682, + -0.0358431302, + -0.000799078902, + 0.0156067787, + 0.0113630332, + 0.00134366076, + -0.0158356745, + -0.0214336012, + -6.795019e-05, + -0.00487118168, + 0.00942035764, + 0.00881240796, + -0.0200685188, + -0.0545426272, + 0.00291975285, + -0.0212819185, + 0.00139869773, + -0.00128811633, + 0.00765864737, + 0.0134908436, + -0.00905226171, + -0.00469854055, + -0.0295520648, + -0.0174266007, + 0.00261859223, + -0.0319056548, + 0.00996019505, + -0.00276106433, + -0.0031865749, + 0.00790517684, + 0.0342602842, + -0.00515013188, + 0.0111019621, + 0.0223657433, + 0.0299786981, + 0.0283940509, + 0.00928591192, + -0.0048799091, + -0.0196352806, + 0.00525931688, + 0.0153731322, + -0.0107787065, + 0.0155906323, + 0.00627624057, + -0.00330341747, + 0.00950430147, + -0.020973973, + 0.00310524087, + -0.0322555453, + 0.0087977834, + 0.00950519182, + 0.0191079155, + -0.00339384656, + 0.0273867417, + 0.0246202089, + -0.000641335675, + 0.0062715942, + -0.0322639607, + 0.0103636729, + -0.0109890597, + 0.0121888723, + -0.0104730874, + 0.00425146008, + 0.0320247374, + 0.00926233362, + 0.00554121612, + -0.0156110479, + -0.000205771255, + 0.00689963764, + -0.0358477831, + -0.00810369011, + -0.0155196106, + -0.0104273837, + -0.015572506, + -0.00284625287, + 0.0241037477, + -0.010501924, + -0.0179182421, + 0.0266444273, + -0.0127375657, + 0.00958085805, + 0.00854709279, + -0.0210839417, + -0.00150059315, + -0.0467378721, + 0.0280720871, + -0.00231987517, + 0.05450489, + 0.0198456924, + 0.0101709999, + 0.0126220165, + 0.016284544, + 0.0131671522, + -0.0153061915, + -0.013661705, + -0.000106541826, + 0.0215635896, + -0.00590792811, + 0.000112961323, + -0.0163673814, + 0.00389808533, + 0.0233915634, + 0.0143893035, + 0.0137647428, + -0.0159027949, + -0.026226325, + -0.0282049291, + -0.00739755295, + 0.00199660845, + 0.00109355152, + -0.0371406898, + -0.0122669982, + -0.0497985221, + -0.0172033068, + -0.0186364222, + 0.025664106, + -0.0164502785, + 0.00586883444, + 0.0265553407, + -0.023504613, + 0.0167928692, + 0.00114404864, + 0.0291096587, + -0.0226561483, + -0.0221388228, + -0.000551475387, + -0.00534460042, + 0.0203121323, + -0.0118539808, + -0.0213559289, + -0.0047192541, + -0.00184764643, + -0.0106019229, + 0.0365246125, + -0.0238530599, + 0.0236058459, + -0.0132803489, + 0.0361760855, + -0.00559246633, + -0.000848451396, + 0.0110634426, + -0.00503157638, + -0.00399979716, + 0.000601884036, + -0.0230017863, + 0.0221566204, + 0.0307091083, + 0.024806466, + -0.00370118767, + -0.0105366809, + 0.0198281351, + -0.0264736712, + -0.000761544041, + 0.0304289144, + -0.0123834098, + -0.0062049008, + -0.0228521228, + -0.00478982273, + -0.0081325192, + -0.0320265256, + -0.017221529, + -0.0124556208, + -0.00513073197, + -0.0055855778, + -0.0108179767, + 0.00543282367, + -0.00222130888, + -0.000894231081, + 0.00564336637, + 0.0260760952, + -0.00854078773, + -0.0273730047, + 0.00505945412, + 0.0120370891, + 0.00533092907, + 0.0139167411, + -0.023129046, + 0.0330098756, + 0.00393412868, + -0.0313481949, + 0.0140373874, + 0.0435902812, + -0.00964412559, + -0.000785533222, + 0.00047182967, + -0.0105094118, + 0.00142559526, + -0.00586001109, + -0.000987534644, + -0.0131708961, + -0.0138070602, + -0.0223419946, + 0.00462451484, + 0.0534970909, + 0.0182369184, + -0.00803868938, + 0.0201791003, + -0.0273638424, + -0.0154543919, + -0.0123929353, + -0.00100496795, + -0.00810563099, + -0.0268817674, + -0.0128428098, + 0.00494655967, + 0.0177365914, + -0.0146534368, + -0.00189528323, + 0.0179479625, + 0.00041076474, + -0.0121961329, + -0.0573405065, + 0.0142375678, + -0.000372575363, + -0.0207336918, + -0.0176597163, + -0.0260544661, + 0.00125782716, + 0.00148278545, + -0.0363535583, + -0.019383315, + 0.0330810398, + -0.0141895432, + -0.0203822292, + 0.0503378101, + -0.0106883692, + 0.0137975244, + 0.0143205766, + 0.00296149589, + 0.0106453681, + -0.0082105957, + -0.000154208406, + -0.0119358134, + -0.0065447255, + -0.0194571838, + 0.0251234621, + 0.00841311924, + 0.00846772548, + 0.0244891699, + 0.0318198763, + -0.0268379319, + 0.0122887129, + 0.00339643634, + 0.0191198569, + 0.00410257187, + -0.0155001376, + 0.0538930371, + -0.000814013591, + 0.0127214724, + 0.0168450456, + 0.020092871, + 0.011120446, + -0.00990895182, + -0.0130007798, + -0.00605561864, + -0.0478215627, + -0.012246754, + 0.00458897045, + -0.015231018, + -0.00272830809, + -0.0268977508, + -0.00523331668, + -0.0145265898, + -0.0189919937, + 0.0014051355, + 0.0144338338, + -0.000881804968, + 0.0093449913, + 0.0110415351, + -0.0116410516, + -0.0105172228, + 0.0371415168, + -0.0315186605, + 0.0270708986, + 0.00602140324, + -0.0152043793, + 0.0285589788, + -0.0139165958, + -0.00282240449, + -0.0200784765, + -0.00414392585, + 0.0508821048, + -0.00527164107, + 0.0197973419, + 0.0500458851, + 0.0229497869, + 1.90233441e-05, + -0.0208758842, + -0.0231342725, + 0.0299868081, + 0.0155696888, + 0.0204890035, + -0.00878940523, + -0.0224861763, + 2.638978e-06, + 0.00806105044, + 0.0168447234, + -0.00885697734, + 0.0256269705, + 0.0414136462, + 0.00485354895, + -0.00238617323, + -9.09848459e-05, + 0.0309595261, + 0.0127148433, + -0.00112215232, + -0.0220220834, + 0.0261854231, + 0.0271701794, + 0.0178945754, + -0.00423594611, + -0.0162710249, + 0.0145595847, + 0.00010253011, + -0.0373664573, + 0.00747355493, + 0.00588322245, + -0.0265010558, + 0.00589571986, + -0.00662960671, + 0.000699193799, + 0.00846142787, + -0.0310527161, + -0.0051752571, + 0.00668232143, + -0.0138790114, + 0.000259768771, + -0.00609485293, + -0.024878161, + -0.0108544929, + 0.00135351834, + -0.0315032564, + -0.00355423009, + -0.0033410734, + -0.0129875354, + -0.00664271135, + -0.0151299834, + 0.00824555848, + -0.00833931006, + -0.0297692716, + -0.000304885907, + -0.00121508515, + -0.0107050184, + 0.0149048213, + 0.0283115935, + -0.0166688561, + 0.0214357339, + 0.0165368728, + -0.0298111383, + -0.00912209414, + 0.028815465, + 0.0162344687, + 0.0113631561, + -0.0430009216, + 8.86659677e-07, + -0.00878425688, + -0.00130400585, + 0.0296355393, + 0.0480274856, + -0.0111993514, + -0.0101491064, + 0.0104581248, + 0.0551528968, + -0.0247439612, + 0.00259530498, + 0.0229519028, + -0.0054376917, + -0.0272449777, + 0.0312338807, + 0.0253772382, + 0.00707713095, + 0.0171943977, + -0.00994254649, + -0.011857504, + 9.57005227e-07, + 0.00918471348, + -0.0247640666, + 0.0327047296, + 0.00876212306, + 0.0108815478, + 0.000260907749, + 0.010264596, + 0.00316306204, + -0.0341044515, + 0.0144040445, + 0.0274574663, + 7.7071476e-05, + 0.00965823699, + 0.00928251818, + 0.00443052454, + 0.00374744181, + -0.0469216332, + -0.019163657 + ] +} diff --git a/artifacts/.gitkeep b/artifacts/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..5099c02 --- /dev/null +++ b/build.sh @@ -0,0 +1,469 @@ +#!/usr/bin/env bash + +# Build the serving tree: vLLM 0.30.0 + patches/, compiled for 4x RTX 3090. +# +# ./build.sh [--force] +# ARTIFACTS=artifacts/build-artifacts-.tar.gz ./build.sh +# +# Default: everything is compiled here, from this tree's source. No flags, no route to choose. +# step 2 creates ./vllm (gitignored): upstream vLLM v0.30.0 (ced6857afa) + patches/*.patch +# (100 patches), then asserts the result against the expected tree hash. FETCH=0 reuses the +# existing tree; FETCH=1 re-applies the series from the base. +# step 4 a fresh .venv from CPython 3.13 + scripts/requirements-pinned.txt (torch 2.13.0, CUDA 13.0). +# step 6 the CUDA 13.0 compiler from the venv's nvidia wheels, targets arch 8.6 (sm_86 = RTX 3090). +# step 7 uv pip install -e . --no-deps --no-build-isolation: configures with cmake, compiles every +# extension of this tree (CUDA/C++/PTX, plus the Rust frontend when a Rust toolchain exists), +# installs the tree into .venv in place, and gives the tree its version and console script. +# +# ARTIFACTS= installs compiled kernels of this same tree instead of compiling them. step 6 +# unpacks them after it checks the tree head, the arch, the CPython minor, the CUDA major and the +# torch version in the tarball's manifest, and hashes every file after the unpack. step 7 then +# installs the tree without compiling (VLLM_TARGET_DEVICE=empty). step 8 writes that tarball, so the +# Dockerfile uses this route when you pass --build-context ops=$PWD/artifacts. +# +# step 8 packs the kernels into artifacts/ for the image build (PACK=0 or --no-pack skips it; the +# Dockerfile sets it, an image build must not grow a 340 MB tarball inside a layer). +# +# Needs: git curl tar gzip sha256sum cmake >= 3.26 ninja cc c++, and uv (UV=, UV=0 to refuse) +# or python3.13 with the venv module. Optional: cargo 1.95 for the Rust frontend. +# Budget: ~30 GB free disk, ~20 min on 64 cores for a first build. MAX_JOBS caps the compile jobs +# (default: every core; each nvcc job can use 4 GB RAM). +# +# Re-runs: a finished build is recorded in vllm/.build-products (tree head + the two extension +# hashes) and vllm/.build-products.sha256 (every build product). When the tree has not moved, every +# hash still matches, and .venv holds the install, steps 6 and 7 are skipped. A tree move, a changed +# product, a missing product, or a missing install ends the skip. --force (or FORCE=1) skips nothing +# and compiles again. +# +# ccache: used automatically when it is on PATH. vllm/setup.py passes ccache as the cmake compiler +# launcher for C, C++, CUDA and HIP, so a rebuild after a patch change recompiles only the files +# that changed. Objects go to the ccache directory of the current user ($HOME/.ccache by default, +# CCACHE_DIR= to move it). A first build fills the cache and gains nothing; the second one is +# where it pays. Without ccache every kernel compiles from scratch and build.sh says so. + +set -euo pipefail + +STEPS=8 +FORCE="${FORCE:-0}" +PACK="${PACK:-1}" +for opt in "$@"; do + case "$opt" in + --force|-f) FORCE=1 ;; + --no-pack) PACK=0 ;; + -h|--help) echo "usage: ./build.sh [--force] [--no-pack]"; echo " --force compile again even when nothing changed"; echo " --no-pack do not write the compiled kernels to artifacts/"; echo " ARTIFACTS= install compiled kernels from artifacts/ instead of compiling"; exit 0 ;; + *) echo "unknown option: $opt (./build.sh --help)"; exit 1 ;; + esac +done + +# Every step announces itself before it runs, so a stall or a failure names the step. +step(){ printf '\n==> [%s/%s] %s\n' "$1" "$STEPS" "$2"; } + +step 1 "tools and CPython 3.13" +for t in git curl tar gzip sha256sum cmake ninja cc c++; do command -v "$t" >/dev/null || { echo "missing tool: $t"; exit 1; }; done +cm="$(cmake --version | sed -n 's/^cmake version \([0-9][0-9.]*\).*/\1/p' || true)" +[ -n "$cm" ] || { echo "no cmake version from \`cmake --version\`"; exit 1; } +[ "$(printf '%s\n3.26\n' "$cm" | sort -V | head -1)" = 3.26 ] || { echo "building from source needs cmake >= 3.26 (found $cm)"; exit 1; } +PY=python3.13 +UV="${UV:-}"; [ -n "$UV" ] || UV="$(command -v uv || true)" +if [ -n "$UV" ]; then + "$UV" python find 3.13 >/dev/null 2>&1 || "$UV" python install 3.13 || { echo "$UV cannot provide CPython 3.13"; exit 1; } + PY="$("$UV" python find 3.13)" +elif command -v python3.13 >/dev/null; then + probe="$(mktemp -d)"; trap 'rm -rf "$probe"' EXIT + python3.13 -m venv "$probe/p" >/dev/null 2>&1 || { echo "python3.13 cannot create a venv: apt install python3.13-venv"; exit 1; } + rm -rf "$probe"; trap - EXIT +else + echo "need uv (https://docs.astral.sh/uv) or python3.13 with the venv module"; exit 1 +fi +HERE="$(cd "$(dirname "$0")" && pwd)"; ROOT="$HERE" +abspath(){ case "$1" in /*) printf '%s\n' "$1" ;; *) printf '%s\n' "$PWD/$1" ;; esac; } +SRC="$(abspath "${SRC:-$HERE/vllm}")" +VENV="$(abspath "${VENV:-$HERE/.venv}")" +ARCH="${ARCH:-8.6}" +JOBS="${MAX_JOBS:-$(nproc)}" +cc="" +if command -v ccache >/dev/null; then + cc="$(command -v ccache)" + # CCACHE_DIR is left alone: ccache uses its own per-user directory ($HOME/.cache/ccache, + # /root/.cache/ccache for root), so the cache is shared with every other build on this account. + [ -n "${CCACHE_DIR:-}" ] && mkdir -p "$CCACHE_DIR" + ccd="$( "$cc" --show-config 2>/dev/null | sed -n 's/.*cache_dir = //p' | head -1 )" + [ -n "$ccd" ] || ccd="${CCACHE_DIR:-the ccache default for this user}" +fi +echo " repo=$ROOT SRC=$SRC VENV=$VENV" +echo " python: $PY cmake $cm arch=$ARCH MAX_JOBS=$JOBS" +if [ -n "$cc" ]; then + echo " ccache: $cc dir=$ccd (objects reused on the next build)" +else + echo " ccache: not on PATH, every kernel compiles from scratch" +fi + +step 2 "serving tree (upstream v0.30.0 + $(ls "$HERE"/patches/*.patch | wc -l) patches)" +# The tree is not in this repo: this repo ships the patches, not the tree. Fetch applies the series +# to the upstream tag and asserts the result, so a patch series that drifts fails here, not at serve. +BASE=ced6857afa0ea7b2e3f0846a62e1394e90f15607 # tag v0.30.0 +EXPECTED_TREE=b7fe074a6c11b5e60ef15f31f432c67626607454 # BASE + patches/*.patch, verified 2026-10-11 +URL="${VLLM_UPSTREAM_URL:-https://github.com/vllm-project/vllm.git}" +# git am makes a commit, so git needs a committer identity: the host's ~/.gitconfig is not there (a +# container, a build user, a fresh box). The committer is not in the tree hash, so the assertion +# below does not depend on it. +GITID=( -c user.name="Hinata" -c user.email="hinata@idihore.id" ) +git config --global core.autocrlf false 2>/dev/null || true +if [ "${FETCH:-0}" = 1 ] || [ ! -f "$SRC/vllm/__init__.py" ]; then + if [ -e "$SRC" ] && [ ! -d "$SRC/.git" ] && [ ! -f "$SRC/vllm/__init__.py" ]; then + echo "$SRC exists and is not a git checkout; move it away first"; exit 1 + fi + echo " base=$BASE tag v0.30.0" + echo " patches=$(ls "$HERE"/patches/*.patch | wc -l) from $HERE/patches" + echo " upstream=$URL" + [ -d "$SRC/.git" ] || { echo " git init $SRC"; git init -q "$SRC"; } + git -C "$SRC" remote add upstream "$URL" 2>/dev/null || git -C "$SRC" remote set-url upstream "$URL" + if ! git -C "$SRC" cat-file -e "$BASE^{commit}" 2>/dev/null; then + echo " fetching v0.30.0 (shallow, ~200 MB, this is the slow part on a fresh host)..." + git -C "$SRC" fetch --quiet --depth 1 "$URL" "refs/tags/v0.30.0:refs/heads/flashnext" \ + || { echo "fetch v0.30.0 from $URL failed (no network?)"; exit 1; } + fi + git -C "$SRC" checkout -q flashnext 2>/dev/null || git -C "$SRC" checkout -q -b flashnext + git -C "$SRC" reset -q --hard "$BASE" + echo " applying patches (git am)..." + git -C "$SRC" "${GITID[@]}" am --keep-non-patch "$HERE"/patches/*.patch >/dev/null +fi +tree="$(git -C "$SRC" rev-parse 'HEAD^{tree}' 2>/dev/null || echo none)" +[ "$tree" = "$EXPECTED_TREE" ] || { echo "tree $tree != expected $EXPECTED_TREE: the series drifted, or the tree is older than the patches. FETCH=1 re-applies it"; exit 1; } +echo " tree ready: $(git -C "$SRC" rev-parse --short HEAD) tree $tree" +[ -z "$(git -C "$SRC" status --porcelain --untracked-files=no 2>/dev/null)" ] || { echo "$SRC has modified tracked files"; exit 1; } +TREE_HEAD="$(git -C "$SRC" rev-parse --short HEAD 2>/dev/null || echo unknown)" +# The commit SHA of the same tree differs between machines: git am stamps the committer date of the +# box that applied the patches, so the host commit and the container commit of one patch set are +# not the same SHA. The tree object is identical. Compare that, not the SHA. +TREE_OBJ="$(git -C "$SRC" rev-parse 'HEAD^{tree}' 2>/dev/null || echo unknown)" +MARK="$SRC/.build-products"; SUMS="$SRC/.build-products.sha256" +LIST="$HERE/scripts/build-artifacts.list" + +step 3 "venv (fresh, never reused)" +if [ -e "$VENV" ] && [ ! -f "$VENV/.build-created" ]; then echo "$VENV exists and was not created here; pick a fresh VENV path"; exit 1; fi +if [ ! -x "$VENV/bin/python" ]; then + if [ -n "$UV" ]; then + v="$("$UV" venv --python 3.13 "$VENV" 2>&1)" || { echo "uv venv failed: $v"; exit 1; } + echo " $v" + else + "$PY" -m venv "$VENV" + fi + echo "created by build.sh" > "$VENV/.build-created" +else + echo " venv exists, reusing: $VENV" +fi +if [ -x "$VENV/bin/pip" ]; then + "$VENV/bin/python" -m pip install pip==26.2.1 + PIP_INSTALL="$VENV/bin/pip install" +elif [ -n "$UV" ]; then + PIP_INSTALL="$UV pip install --python $VENV/bin/python" +else + echo "$VENV has no pip and no uv is available"; exit 1 +fi +echo " installer: $PIP_INSTALL" +# Where the wheels come from. uv and pip use the cache directory of the current user, so a build +# under the same account reuses the wheels of every previous build. Nothing here changes them. +cache_of(){ [ -n "${2:-}" ] && { printf '%s\n' "$2"; return; }; printf '%s\n' "${HOME:-/root}/.cache/$1"; } +uv_cd="${UV_CACHE_DIR:-$(cache_of uv "${XDG_CACHE_HOME:-}")}" +if [ -n "$UV" ]; then + uv_n="$(du -sh "$uv_cd" 2>/dev/null | cut -f1 || true)" + echo " uv cache: $uv_cd${uv_n:+ ($uv_n)}" +fi +pip_cd="${PIP_CACHE_DIR:-$(cache_of pip "${XDG_CACHE_HOME:-}")}" +pip_n="$(du -sh "$pip_cd" 2>/dev/null | cut -f1 || true)" +echo " pip cache: $pip_cd${pip_n:+ ($pip_n)}" +pip_install(){ $PIP_INSTALL "$@"; } + +step 4 "pinned runtime (torch 2.13.0 + CUDA 13.0 and the vllm deps; $(grep -cv '^\s*\(#\|$\)' "$HERE/scripts/requirements-pinned.txt") pinned packages)" +pip_install -r "${REQ:-$HERE/scripts/requirements-pinned.txt}" +pip_install setuptools-rust==1.13.0 setuptools-scm==10.3.4 vcs-versioning==2.5.0 semantic-version==2.10.0 +"$VENV/bin/python" -c "import torch; assert torch.__version__.split('+')[0] == '2.13.0' and torch.version.cuda == '13.0', (torch.__version__, torch.version.cuda)" + +step 5 "compile toolchain (CUDA 13.0 from the venv wheels)" +CU="$("$VENV/bin/python" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')/nvidia/cu13" +[ -x "$CU/bin/nvcc" ] || { echo "no CUDA toolkit wheel at $CU"; exit 1; } +[ -e "$CU/lib64" ] || ln -s lib "$CU/lib64" +for so in "$CU"/lib/lib*.so.[0-9]*; do base="${so%%.so.*}.so"; [ -e "$base" ] || ln -s "$(basename "$so")" "$base"; done +[ -e "$CU/lib/libcudart.so" ] || { echo "$CU/lib has no libcudart.so.13"; exit 1; } +NVCC="${NVCC:-$CU/bin/nvcc}" +nv_rel="$("$NVCC" --version | sed -n 's/.*release \([0-9][0-9]*\.[0-9][0-9]*\),.*/\1/p')" +[ "$nv_rel" = 13.0 ] || { echo "nvcc is CUDA ${nv_rel:-unknown}, expected 13.0"; exit 1; } +echo " nvcc $nv_rel at $NVCC (start.sh points CUDA_HOME here, for the runtime JIT kernels)" + +# The kernels link CUDA::cuda_driver, the driver library (libcuda.so.1). The +# driver package provides it; the CUDA wheel does not. Generate the link stub. +stub_dir="$CU/lib/stubs" +mkdir -p "$stub_dir" +"$HERE/scripts/gen-cuda-stub.sh" "$CU/include/cuda.h" "$stub_dir/libcuda.so" | sed 's/^/ /' +cargo="$(command -v cargo || true)" +if [ -n "$cargo" ]; then + echo " rust frontend: $cargo $(cargo --version 2>/dev/null || echo unknown) -- vllm-rs and _rust_tool_parser will be built" +else + echo " rust frontend: no cargo -> vllm-rs and _rust_tool_parser are skipped." + echo " the shipped serve profile uses the Python tool parsers, so serving is unaffected." + echo " install rustup toolchain 1.95 (vllm/rust-toolchain.toml) and re-run ./build.sh to add them." +fi + +# glibc: the products of a packed build carry the newest glibc symbol they need. The box that +# installs them must be at least that new, or the loader fails at serve time, not here. Print it +# and stop, with the reason. +glibc_ver(){ ldd --version 2>/dev/null | sed -n 's/^ldd [^0-9]*\([0-9][0-9.]*\).*/\1/p' | head -1; } +ver_ge(){ [ "$(printf '%s\n%s\n' "$1" "$2" | sort -V | head -1)" = "$2" ]; } + +# The build products this tree must have at serve time. build-artifacts.list is the list. +check_list(){ local miss=0 p; while read -r p; do case "$p" in ''|'#'*) continue ;; esac; [ -e "$SRC/${p%/}" ] || { echo "missing build product: $p"; miss=1; }; done < "$LIST"; return $miss; } +write_sums(){ (cd "$SRC" && while read -r p; do case "$p" in ''|'#'*) continue ;; esac; find "${p%/}" -type f; done < "$LIST" | grep -vx 'vllm/_version.py' | LC_ALL=C sort | xargs -d '\n' sha256sum) > "$SUMS.tmp" && mv "$SUMS.tmp" "$SUMS"; } +check_sums(){ [ -s "$SUMS" ] && (cd "$SRC" && sha256sum -c --quiet --strict "$SUMS"); } + +# Prebuilt ops: ARTIFACTS= installs compiled kernels instead of compiling them. The +# tarball comes from step 8 of a finished build of this exact tree. The checks below are the reason +# a packed build is safe to install: the tree +# head, the CPython minor, the CUDA major, the torch version and the arch must all match, and +# every file is hashed after the unpack. One mismatch stops the build and says which one. +prebuilt=0 +if [ -n "${ARTIFACTS:-}" ]; then + step 6 "install prebuilt ops (no compile)" + [ -f "$ARTIFACTS" ] || { echo "ARTIFACTS=$ARTIFACTS is not a file"; exit 1; } + adir="$(cd "$(dirname "$ARTIFACTS")" && pwd)"; abase="$(basename "$ARTIFACTS")"; stem="${abase%.tar.gz}" + [ -f "$adir/$stem.sha256" ] || { echo "$adir/$stem.sha256 is needed: the per-file hashes verified after the unpack"; exit 1; } + [ -f "$adir/$stem.manifest" ] || { echo "$adir/$stem.manifest is needed: it records the tree, arch, python, CUDA and torch the ops were built for"; exit 1; } + if [ -f "$ARTIFACTS.sha256" ]; then + ( cd "$adir" && sha256sum -c --quiet "$stem.tar.gz.sha256" ) || { echo "$ARTIFACTS does not match $stem.tar.gz.sha256"; exit 1; } + echo " tarball sha256: ok" + else + echo " tarball sha256: no $stem.tar.gz.sha256 sidecar, the tarball itself is not verified" + fi + # Read the manifest with sed, not with ".": cc= and glibc= carry text with parentheses and + # spaces, which the shell cannot source. + mf(){ sed -n "s/^$1=//p" "$adir/$stem.manifest" | head -1; } + m_tree="$(mf tree)"; m_tree_obj="$(mf tree_obj)"; m_arch="$(mf arch)"; m_py="$(mf python)"; m_cuda="$(mf cuda)" + m_torch="$(mf torch)"; m_glibc="$(mf glibc)"; m_files="$(mf files)"; m_built="$(mf built)" + echo " manifest: tree=$m_tree tree_obj=${m_tree_obj:0:7} arch=$m_arch python=$m_py cuda=$m_cuda torch=$m_torch glibc=$m_glibc files=$m_files built=$m_built" + fail=0 + [ -n "$m_tree_obj" ] || { echo "$adir/$stem.manifest has no tree_obj= line: pack the ops again with ./build.sh"; exit 1; } + # The tree content is what the kernels were compiled from, so the tree object is the check. The + # commit SHA is printed next to it and is allowed to differ between machines (committer date). + [ "$m_tree_obj" = "$TREE_OBJ" ] || { echo "the ops were packed from tree ${m_tree_obj:0:7}, this tree is ${TREE_OBJ:0:7} (commits $m_tree / $TREE_HEAD): pack them from this tree, or compile"; fail=1; } + venv_py="$("$VENV/bin/python" -c 'import sys; print("%d.%d" % sys.version_info[:2])')" + [ "$m_py" = "$venv_py" ] || { echo "the ops are for CPython $m_py, the venv is $venv_py"; fail=1; } + [ "$m_cuda" = "$nv_rel" ] || { echo "the ops are for CUDA $m_cuda, this build uses CUDA $nv_rel"; fail=1; } + venv_torch="$("$VENV/bin/python" -c 'import torch; print(torch.__version__.split("+")[0])')" + [ "$m_torch" = "$venv_torch" ] || { echo "the ops are for torch $m_torch, the venv has torch $venv_torch"; fail=1; } + [ "$m_arch" = "$ARCH" ] || { echo "the ops are built for arch $m_arch, this build asks for arch $ARCH"; fail=1; } + [ "$fail" = 0 ] || exit 1 + if [ -n "${FORCE:-}" ] && [ "$FORCE" = 1 ]; then rm -f "$MARK" "$SUMS"; echo " --force: replacing the products in the tree"; fi + echo " unpacking $(du -h "$ARTIFACTS" | cut -f1) into $SRC" + tar -xzf "$ARTIFACTS" -C "$SRC" + ( cd "$SRC" && sha256sum -c --quiet --strict "$adir/$stem.sha256" ) \ + || { echo "a file of the tarball does not match $stem.sha256 after the unpack"; exit 1; } + echo " unpacked: all $m_files file hashes match the manifest" + need="$(readelf -V "$SRC/vllm/_C_stable_libtorch.abi3.so" 2>/dev/null | grep -o 'GLIBC_[0-9.]*' | sed 's/GLIBC_//' | sort -V | tail -1)" + have="$(glibc_ver)" + if [ -n "$need" ] && [ -n "$have" ]; then + ver_ge "$have" "$need" || { echo "these ops need glibc $need, this box has $have: build the image on a base with that glibc, or compile here"; exit 1; } + echo " glibc: ops need $need, this box has $have" + else + echo " glibc check skipped: no readelf or no ldd here" + fi + check_list || { echo "the tarball does not carry every product of $(basename "$LIST")"; exit 1; } + write_sums + prebuilt=1 +fi + +# A finished build of this tree is reused, so a re-run does not take 20 minutes again. Every check +# below has to pass to skip the compile, and the one that fails says what changed. +reuse=0 +if [ "$prebuilt" = 1 ]; then + reuse=0 +elif [ -f "$MARK" ] && [ "$FORCE" = 1 ]; then + echo " --force: compiling again; the record in $(basename "$MARK") is not trusted" +elif [ -f "$MARK" ]; then + read -r _ c_sha moe_sha tree_at_build < "$MARK" + reuse=1 + [ "$tree_at_build" = "$TREE_HEAD" ] || { echo "note: the tree moved $tree_at_build -> $TREE_HEAD; the compile below rebuilds what changed"; reuse=0; } + if [ "$(sha256sum "$SRC/vllm/_C_stable_libtorch.abi3.so" 2>/dev/null | cut -d' ' -f1)" != "$c_sha" ] || [ "$(sha256sum "$SRC/vllm/_moe_C_stable_libtorch.abi3.so" 2>/dev/null | cut -d' ' -f1)" != "$moe_sha" ]; then + echo "note: the extensions are not the ones recorded in $MARK; the compile below rebuilds them" + reuse=0 + fi + check_list || { echo "note: build products are missing; the compile below builds them"; reuse=0; } + [ "$reuse" = 1 ] && { check_sums || { echo "note: build products changed since $(basename "$SUMS"); the compile below rebuilds them"; reuse=0; }; } + # A source-only change (a patch that touches no kernel) is invisible to the product hashes above, + # and the editable install of step 7 is what makes the tree importable. So the install must exist. + [ "$reuse" = 1 ] && { ls -d "$VENV"/lib/python3.*/site-packages/vllm-*.dist-info >/dev/null 2>&1 \ + || { echo "note: $VENV has no installed vllm; step 7 installs it"; reuse=0; }; } + if [ "$reuse" = 1 ]; then + echo " existing build products: present, all $(wc -l < "$SUMS") hashes re-verified (built at tree $tree_at_build)" + else + rm -f "$MARK" "$SUMS"; echo " old build record removed, verification runs after the compile" + fi +fi + +if [ "$prebuilt" = 1 ]; then + step 7 "install the tree into .venv (editable, prebuilt ops, metadata only)" + # The kernels are already compiled, so setup.py must not compile again. VLLM_TARGET_DEVICE=empty + # makes it skip every compile and write the version file and the console script only. + build_log="$ROOT/build-$(date +%Y%m%d-%H%M%S).log" + echo " full build output -> $build_log" + t0=$(date +%s) + ( while sleep 60; do echo " ... $(( ( $(date +%s) - t0 ) / 60 )) min elapsed, installing metadata"; done ) & hb=$! + ( cd "$SRC" && env VLLM_TARGET_DEVICE=empty VLLM_VERSION_OVERRIDE=0.30.0 \ + sh -c "$PIP_INSTALL -e . --no-deps --no-build-isolation" ) 2>&1 | tee "$build_log" \ + || { kill "$hb" 2>/dev/null; echo "step 7 failed; the install output is in $build_log"; exit 1; } + kill "$hb" 2>/dev/null || true; wait "$hb" 2>/dev/null || true + check_list || { echo "installed, but build products are missing; see $build_log"; exit 1; } +elif [ "$reuse" = 1 ]; then + step 6 "compile (skipped: no change to the tree, the products, or the install)" + step 7 "install (skipped: .venv already holds this tree)" + echo " nothing to compile. ./build.sh --force compiles again." +else + step 6 "build this tree from source (arch $ARCH, $(grep -vc '^\s*#\|^\s*$' "$LIST") products in $(basename "$LIST"))" + echo " setup.py configures cmake and compiles: $(grep -v '^\s*#\|^\s*$' "$LIST" | tr '\n' ' ')" + echo " first build: ~20 min on 64 cores, then a re-run of an unchanged tree skips this step" + echo " the compile runs in step 7: uv prints every compiler line, one status line every 60s" + + # Install the tree in place. setup.py does the cmake configure, compiles this tree's kernels, + # and drops the products into $SRC/vllm/. --no-deps keeps the pins from step 4, + # --no-build-isolation compiles against the torch in the venv. + step 7 "install the tree into .venv (editable, compiled)" + build_log="$ROOT/build-$(date +%Y%m%d-%H%M%S).log" + echo " full build output -> $build_log" + stats(){ "$cc" -v -s 2>/dev/null | grep -E 'Hits:|Misses:|Cache size' | tr -s ' ' | tr '\n' ' '; } + before_stats="$(stats || true)" + # -v is what makes the compile visible. Without it uv swallows the build backend's stdout, so a + # 20 minute compile prints nothing between "Building vllm @ file:..." and "Built vllm". + # -v streams every nvcc/cc/cargo line (uv prefixes them with DEBUG). -vv adds uv's own debug noise. + t0=$(date +%s) + # Products to watch for, in the order setup.py produces them. Counted from the list, so the + # heartbeat ends at 15/15 and never counts a file the list does not require. + GCC_VER="" + for v in 14 13 12 11 10; do + if command -v "gcc-$v" >/dev/null 2>&1; then + GCC_VER="$v" + break + fi + done + [ -n "$GCC_VER" ] || { echo "Error: No GCC version < 15 found" >&2; exit 1; } + CC_PATH="/usr/bin/gcc-$GCC_VER" + CXX_PATH="/usr/bin/g++-$GCC_VER" + prod=$(grep -v '^\s*#\|^\s*$' "$LIST" | sed "s#^#$SRC/#") + ( while sleep 60; do + e=$(( ( $(date +%s) - t0 ) / 60 )) + p=$(grep -oE '\[[0-9]+/[0-9]+\]' "$build_log" 2>/dev/null | tail -1) + n=0; while read -r f; do [ -e "$f" ] && n=$(( n + 1 )); done <<< "$prod" + echo " ... ${e} min elapsed, cmake ${p:-configuring}, ${n}/$(printf '%s\n' "$prod" | wc -l) build products in place" + done ) & hb=$! + ( cd "$SRC" && env CUDA_HOME="$CU" PATH="$(dirname "$NVCC"):$PATH" TORCH_CUDA_ARCH_LIST="$ARCH" \ + MAX_JOBS="$JOBS" VLLM_TARGET_DEVICE=cuda VLLM_VERSION_OVERRIDE=0.30.0 \ + CC="$CC_PATH" CXX="$CXX_PATH" \ + CMAKE_ARGS="-DCMAKE_CUDA_HOST_COMPILER=$CXX_PATH -DCMAKE_C_COMPILER=$CC_PATH -DCMAKE_CXX_COMPILER=$CC_PATH" \ + sh -c "$PIP_INSTALL -v -e . --no-deps --no-build-isolation" ) 2>&1 | tee "$build_log" \ + || { kill "$hb" 2>/dev/null; echo "step 7 failed; the compiler output is in $build_log"; exit 1; } + kill "$hb" 2>/dev/null || true; wait "$hb" 2>/dev/null || true + grep -iE "warning: |error: " "$build_log" | grep -v "^\s*$" | head -20 || true + if [ -n "$cc" ]; then + echo " ccache before this build: ${before_stats:-no stats yet}" + echo " ccache after this build: $(stats || true)" + fi + check_list || { echo "installed, but build products are missing; see $build_log"; exit 1; } + write_sums +fi +echo " $(wc -l < "$SUMS") build files verified ($(du -sh "$SRC/vllm/_C_stable_libtorch.abi3.so" | cut -f1) + $(du -sh "$SRC/vllm/_moe_C_stable_libtorch.abi3.so" | cut -f1) extensions)" + +# Smoke test. Importing a compiled kernel loads libcuda.so.1, the CUDA driver +# library. A box with a GPU has it (the driver package); a build container does +# not, so the import cannot run there and does not need to. Without a driver, +# check the link layer instead: the products exist (step 6 hashed them) and the +# loader entry that needs the driver is the one we expect. The full import runs +# at serve time, on a box with the driver. +if ldconfig -p 2>/dev/null | grep -q 'libcuda\.so\.1' || [ -e /dev/nvidia0 ]; then + smoke="$(cd / && "$VENV/bin/python" -c ' +import torch, vllm +import vllm._C_stable_libtorch, vllm._moe_C_stable_libtorch +import vllm.vllm_flash_attn._vllm_fa2_C, vllm.vllm_flash_attn._vllm_fa3_C +from vllm.tool_parsers import ToolParserManager +ToolParserManager.get_tool_parser("qwen3_coder") +print(vllm.__version__, vllm.__file__)')" \ + || { echo "the installed tree does not import; see $build_log"; exit 1; } + [ "${smoke%% *}" = 0.30.0 ] || { echo "vllm reports ${smoke%% *}, expected 0.30.0"; exit 1; } + echo " import check: vllm $smoke" +else + echo " no CUDA driver here: the import check runs at serve time, on the GPU box" + if command -v readelf >/dev/null; then + # Only the extension that calls the driver API directly needs libcuda.so.1 in its NEEDED list + # (CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL). The MoE kernels do not: they get CUDA through + # libtorch_cuda/libcudart, so libcuda is absent from their NEEDED list and that is correct. + # Requiring it there fails a good build (it did, in Docker). Check it where it is expected, + # and check the torch/CUDA runtime libs everywhere. + for so in vllm/_C_stable_libtorch.abi3.so vllm/_moe_C_stable_libtorch.abi3.so \ + vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so; do + [ -e "$SRC/$so" ] || continue + n="$(readelf -d "$SRC/$so" | sed -n 's/.*\[\(lib[^]]*\)\].*/\1/p' | tr '\n' ' ')" + [ -n "$n" ] || { echo "$so has no dynamic section; the file is not a shared object"; exit 1; } + for need in libtorch.so libtorch_cuda.so libcudart.so; do + case " $n" in *" $need"*) ;; *) echo "$so does not ask for $need; the link is wrong"; exit 1 ;; esac + done + case "$so" in + vllm/_C_stable_libtorch.abi3.so) case " $n" in *libcuda.so*) ;; *) echo "$so does not ask for libcuda.so.1; the link is wrong"; exit 1 ;; esac ;; + esac + echo " link check $(basename "$so"):${n:+ }$n" + done + echo " link check: every lib above is found at run time (libcuda.so.1 comes from the driver)" + else + echo " link check skipped: no readelf on this box" + fi +fi + +{ echo "source $(sha256sum "$SRC/vllm/_C_stable_libtorch.abi3.so" | cut -d' ' -f1) $(sha256sum "$SRC/vllm/_moe_C_stable_libtorch.abi3.so" | cut -d' ' -f1) $TREE_HEAD"; } > "$MARK" + +# step 8: pack the kernels for the image build. The same tree object is what step 6 accepts back, +# so this tarball is the fast route for Docker: ARTIFACTS=artifacts/.tar.gz, no compile. +# The tag carries what the products were built for: sm86 (arch 8.6) + the CPython minor + the CUDA +# major. A 3090 tarball does not work on a 5090; CUDA 13 products do not work with CUDA 12. +if [ "$PACK" = 0 ]; then + step 8 "pack the kernels for the image build (skipped: PACK=0 / --no-pack)" +elif [ "$prebuilt" = 1 ]; then + step 8 "pack the kernels for the image build (skipped: these kernels came from a tarball)" +else + step 8 "pack the kernels for the image build" + OUT="${OUT:-$HERE/artifacts}"; mkdir -p "$OUT" + pyv="$("$VENV/bin/python" -c 'import sys; print("%d.%d" % sys.version_info[:2])')" + tv="$("$VENV/bin/python" -c 'import torch; print(torch.__version__.split("+")[0])')" + tag="sm${ARCH/./}-py${pyv/./}-cu${nv_rel%%.*}" + name="build-artifacts-$tag" + have_obj=""; [ -s "$OUT/$name.manifest" ] && have_obj="$(sed -n 's/^tree_obj=//p' "$OUT/$name.manifest" | head -1)" + if [ -s "$OUT/$name.tar.gz" ] && [ -s "$OUT/$name.sha256" ] && [ "$have_obj" = "$TREE_OBJ" ]; then + echo " artifacts/$name.tar.gz already holds tree ${TREE_OBJ:0:7}: not packed again" + else + files="$( cd "$SRC" && while read -r p; do case "$p" in ''|'#'*) continue ;; esac + find "${p%/}" -type f; done < "$LIST" | LC_ALL=C sort )" + n="$(printf '%s\n' "$files" | grep -c . || true)" + [ "$n" -gt 0 ] || { echo "$(basename "$LIST") names no file under $SRC"; exit 1; } + echo " packing $n files -> artifacts/$name.tar.gz (~45s)" + # tar runs inside $SRC so every member starts with vllm/, which is the layout step 6 unpacks. + ( cd "$SRC" && printf '%s\n' "$files" | xargs -d '\n' tar -czf "$OUT/$name.tar.gz" -- ) \ + || { echo "packing failed"; exit 1; } + ( cd "$SRC" && printf '%s\n' "$files" | LC_ALL=C sort | xargs -d '\n' sha256sum ) > "$OUT/$name.sha256" + { + echo "tree=$TREE_HEAD" + echo "tree_obj=$TREE_OBJ" + echo "tree_full=$(git -C "$SRC" rev-parse HEAD)" + echo "arch=$ARCH" + echo "python=$pyv" + echo "cuda=$nv_rel" + echo "torch=$tv" + echo "cc=$(cc --version 2>/dev/null | head -1 || echo unknown)" + echo "glibc=$(glibc_ver || echo unknown)" + echo "built=$(date -u +%Y-%m-%dT%H:%M:%SZ)" + echo "files=$n" + echo "sha_C=$(sha256sum "$SRC/vllm/_C_stable_libtorch.abi3.so" | cut -d' ' -f1)" + echo "sha_moe=$(sha256sum "$SRC/vllm/_moe_C_stable_libtorch.abi3.so" | cut -d' ' -f1)" + } > "$OUT/$name.manifest" + sha256sum "$OUT/$name.tar.gz" | sed "s| .*$| $name.tar.gz|" > "$OUT/$name.tar.gz.sha256" + ls -l "$OUT/$name.tar.gz" "$OUT/$name.sha256" "$OUT/$name.manifest" | sed 's/^/ /' + fi +fi + +echo "done: TREE=$SRC (tree $TREE_HEAD) VENV=$VENV" +echo "next: ./start.sh /path/to/Qwen3.8-Flash-Next-W4A16-Merlin" +echo " docker: docker build --build-context ops=\"$HERE/artifacts\" -t vllm:latest . (no compile)" diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..b2f0d1e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,61 @@ +services: + vllm: + image: ${IMAGE:-vllm}:${IMAGE_TAG:-latest} + container_name: ${CONTAINER_NAME:-vllm} + build: + context: . + args: + MAX_JOBS: ${MAX_JOBS:-32} + ARCH: ${ARCH:-8.6} + additional_contexts: + ops: ./artifacts + uvcache: ${UV_CACHE_DIR:-${HOME}/.cache/uv} + pipcache: ${PIP_CACHE_DIR:-${HOME}/.cache/pip} + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + command: ["${MODEL_MOUNT_PATH:-/models/Qwen3.8-Flash-Next}"] + shm_size: 16g + ulimits: + memlock: -1 + stack: 67108864 + ports: + - "${HOST_BIND_ADDR:-0.0.0.0}:${HOST_PORT:-8080}:8000" + environment: + MODEL_NAME: ${MODEL_NAME:-Hinata} + PORT: 8000 + SCALES: ${SCALES:-} + VLLM_API_KEY: ${VLLM_API_KEY:-} + PLE_HOME: ${PLE_HOME:-0} + CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-0,1,2,3} + AUTO_TOPO: ${AUTO_TOPO:-0} + AUTO_TOPO_NCCL: ${AUTO_TOPO_NCCL:-legacy} + AUTO_TOPO_BIND: ${AUTO_TOPO_BIND:-0} + LANE_GPUS: ${LANE_GPUS:-} + LANE_CPUS: ${LANE_CPUS:-} + TOPO_TIMEOUT: ${TOPO_TIMEOUT:-20} + GUARD: ${GUARD:-warn} + COUNTERS: ${COUNTERS:-0} + VLLM_LOG_STATS_INTERVAL: ${VLLM_LOG_STATS_INTERVAL:-1} + QWEN38_ABLITERATION: ${QWEN38_ABLITERATION:-Qwen3.8-Flash-Next-Uncensored} + volumes: + - "${MODEL_DIR:?Set MODEL_DIR in .env}:${MODEL_MOUNT_PATH:-/models/Qwen3.8-Flash-Next}:ro" + - ${ABLITERATION_DIR:-./abliteration}:/opt/qwen38/abliteration:ro + - ${CACHE_HOST_DIR:-./.vllm-cache}:/cache + - ./start.sh:/opt/vllm/start.sh:ro + restart: unless-stopped + stop_grace_period: 20s + logging: + driver: json-file + options: + max-size: ${LOG_MAX_SIZE:-100m} + max-file: ${LOG_MAX_FILE:-14} + healthcheck: + interval: ${HEALTHCHECK_INTERVAL:-60s} + timeout: ${HEALTHCHECK_TIMEOUT:-90s} + start_period: ${HEALTHCHECK_START_PERIOD:-30m} + retries: ${HEALTHCHECK_RETRIES:-3} diff --git a/patches/0001-worker-capture-safe-guard-telemetry-infrastructure-d.patch b/patches/0001-worker-capture-safe-guard-telemetry-infrastructure-d.patch new file mode 100644 index 0000000..b0714ae --- /dev/null +++ b/patches/0001-worker-capture-safe-guard-telemetry-infrastructure-d.patch @@ -0,0 +1,970 @@ +From 1b2ba827becf485be6c06d2447b1faf2772c42c9 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:08:13 -0400 +Subject: [PATCH 001/103] worker: capture-safe guard-telemetry infrastructure, + default off + +Shared machinery for device-side bounds guards: persistent int64 counters per +site, written by guarded kernels and read by a background thread. The reader +starts only after warm-up and CUDA-graph capture and never waits on the +compute stream; a read failure arms the stop rather than popping the buffer. +Shutdown does one synchronous final drain so a hit on the last step is +reported before teardown. Release default: off. +--- + .../gpu/test_guard_reader_nonblocking.py | 135 ++++ + .../test_guard_telemetry_defaults_drain.py | 95 +++ + vllm/v1/worker/gpu/model_runner.py | 9 + + vllm/v1/worker/gpu_worker.py | 17 + + vllm/v1/worker/guard_telemetry.py | 620 ++++++++++++++++++ + 5 files changed, 876 insertions(+) + create mode 100644 tests/v1/worker/gpu/test_guard_reader_nonblocking.py + create mode 100644 tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py + create mode 100644 vllm/v1/worker/guard_telemetry.py + +diff --git a/tests/v1/worker/gpu/test_guard_reader_nonblocking.py b/tests/v1/worker/gpu/test_guard_reader_nonblocking.py +new file mode 100644 +index 0000000..94721e1 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_guard_reader_nonblocking.py +@@ -0,0 +1,135 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""the instrumentation readers must never wait on the compute stream. ++ ++The first cadence reader did a synchronous pageable ``tolist()`` from its thread; the driver held the CUDA context lock ++while that copy waited for the stream, and a pending PLE pull-transport wait on the stream (whose producer needs the ++same lock) deadlocked the worker. These tests park a kernel on the current (compute) stream and require a reader tick ++to return while it is still running. On the pre-fix reader the first GPU test blocks for the kernel's duration. ++""" ++import threading ++import time ++ ++import pytest ++import torch ++ ++from vllm.v1.worker import guard_telemetry ++ ++SLEEP_CYCLES = 4_000_000_000 # ~2 s on a 3090 (torch.cuda._sleep spins for this many cycles) ++ ++ ++def test_cpu_tensor_reads_directly(): ++ r = guard_telemetry.NonBlockingReader() ++ t = torch.tensor([[1, 2, 3]], dtype=torch.int64) ++ r.register("cpu", t) ++ assert r.poll_one("cpu") == [[1, 2, 3]] ++ t[0, 0] = 7 ++ assert r.poll()["cpu"] == [[7, 2, 3]] ++ ++ ++def test_register_refuses_to_run_inside_a_capture(monkeypatch): ++ """Creating the side stream inside a cudagraph capture is not capture-safe.""" ++ r = guard_telemetry.NonBlockingReader() ++ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) ++ if not torch.cuda.is_available(): ++ pytest.skip("CUDA required for the cuda-tensor branch") ++ t = torch.zeros(2, dtype=torch.int64, device="cuda") ++ with pytest.raises(RuntimeError, match="during cudagraph capture"): ++ r.register("x", t) ++ ++ ++def test_report_waits_for_the_confirming_tick_cpu(): ++ """A count rise seen in copy k is reported from copy k+1 (maxima current by then).""" ++ device = torch.device("cpu") ++ guard_telemetry.reset_for_test() ++ buf = guard_telemetry.counter(0, device) # allocates the cpu buffer ++ full = guard_telemetry._buffers[device] ++ site = guard_telemetry.SITE_MAMBA_COPY ++ full[site, 0] = 2 # count written, maxima not yet ++ guard_telemetry._poll_once() # rise seen -> pending, nothing reported ++ assert guard_telemetry._seen.get((device, site), 0) == 0 ++ full[site, 1] = 9 + guard_telemetry.GUARD_BIAS # maxima land before the next tick ++ guard_telemetry._poll_once() # confirming tick reports count 2 with row 9 ++ assert guard_telemetry._seen.get((device, site), 0) == 2 ++ snap = guard_telemetry.snapshot() ++ key = f"{device}:{guard_telemetry.SITE_NAMES[site]}" ++ assert snap[key][:2] == [2, 9] ++ guard_telemetry.reset_for_test() ++ ++ ++def test_module_reader_is_the_non_blocking_kind(): ++ assert isinstance(guard_telemetry._nb_reader, guard_telemetry.NonBlockingReader) ++ src = open(guard_telemetry.__file__, encoding="utf-8").read() ++ body = src[src.index("def _poll_once"):src.index("def _arm(")] ++ assert ".tolist()" not in body.replace("host.tolist()", ""), "the cadence reader must not read the device buffer synchronously" ++ assert "poll_one(" in body ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_reader_tick_returns_while_the_compute_stream_is_busy(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ torch.cuda.synchronize(device) ++ # prime: one poll starts the first copy (nothing to report yet) ++ guard_telemetry._poll_once() ++ torch.cuda.synchronize(device) ++ ++ # Park a long kernel on the compute (current) stream, then tick the reader. ++ torch.cuda._sleep(SLEEP_CYCLES) ++ t0 = time.monotonic() ++ guard_telemetry._poll_once() ++ tick = time.monotonic() - t0 ++ # Prove the stream really was busy: the sync must take a while. ++ t1 = time.monotonic() ++ torch.cuda.synchronize(device) ++ busy = time.monotonic() - t1 ++ # The tick assertion comes first: a blocking reader spends the whole sleep inside the tick and would then make ++ # the busy control read ~0 (negative control). ++ assert tick < 0.25, f"reader tick took {tick:.3f}s while the compute stream was busy: it waited on the stream (reader/compute-stream deadlock)" ++ assert tick + busy > 0.5, f"control failed: the parked kernel finished in {tick + busy:.3f}s, the test cannot discriminate" ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_reader_tick_from_a_thread_returns_while_the_compute_stream_is_busy(): ++ """The production shape: the tick runs on the reader thread while the main thread owns the stream.""" ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ guard_telemetry._poll_once() ++ torch.cuda.synchronize(device) ++ torch.cuda._sleep(SLEEP_CYCLES) ++ done = threading.Event() ++ elapsed = {} ++ ++ def tick(): ++ t0 = time.monotonic() ++ guard_telemetry._poll_once() ++ elapsed["tick"] = time.monotonic() - t0 ++ done.set() ++ ++ threading.Thread(target=tick, name="test-guard-reader", daemon=True).start() ++ finished = done.wait(timeout=1.0) ++ torch.cuda.synchronize(device) ++ assert finished, "the reader tick did not return within 1 s while the compute stream was busy (reader/compute-stream deadlock)" ++ assert elapsed["tick"] < 0.25 ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_hits_written_before_a_stall_are_still_reported(): ++ """Non-blocking must not mean lossy: a count written on the compute stream is reported once the copy lands.""" ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ buf = guard_telemetry._buffers[device] ++ site = guard_telemetry.SITE_MAMBA_COPY ++ buf[site, 0] += 3 # a "masked access" count, written on the compute stream ++ buf[site, 1] = 5 + guard_telemetry.GUARD_BIAS ++ torch.cuda._sleep(SLEEP_CYCLES // 4) ++ torch.cuda.synchronize(device) ++ # copy -> rise seen -> confirming copy -> report: four ticks at most ++ for _ in range(4): ++ guard_telemetry._poll_once() ++ assert guard_telemetry._seen.get((device, site), 0) == 3 ++ snap = guard_telemetry.snapshot() ++ key = f"{device}:{guard_telemetry.SITE_NAMES[site]}" ++ assert snap[key][0] == 3 and snap[key][1] == 5 +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +new file mode 100644 +index 0000000..6799a80 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +@@ -0,0 +1,95 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Bounds-guard instrumentation: release defaults, reader loss and the final drain. ++ ++1. Release defaults are off / 0.05 with both variables absent (earlier ++ builds defaulted to warn), and a ++ non-finite or non-positive interval is rejected at import. ++2. A reader that loses a device arms the stop (the fatal-site contract is ++ otherwise silently void). ++3. drain() reports a hit the cadence reader never saw and raises for an ++ always-fatal site when asked; it is what a stage end / orderly shutdown ++ calls so a last-step hit cannot be lost. ++""" ++ ++import importlib ++import os ++import sys ++ ++import pytest ++import torch ++ ++pytestmark = pytest.mark.skipif( ++ not torch.cuda.is_available(), reason="CUDA required" ++) ++ ++ ++def _fresh(monkeypatch, **env): ++ for k in ("VLLM_E2_GUARD_MODE", "VLLM_E2_GUARD_INTERVAL"): ++ monkeypatch.delenv(k, raising=False) ++ for k, v in env.items(): ++ monkeypatch.setenv(k, v) ++ sys.modules.pop("vllm.v1.worker.guard_telemetry", None) ++ return importlib.import_module("vllm.v1.worker.guard_telemetry") ++ ++ ++def test_release_defaults_with_env_absent(monkeypatch): ++ gt = _fresh(monkeypatch) ++ assert gt._MODE == "off" # v2: default off ++ assert gt._INTERVAL == pytest.approx(0.05) ++ ++ ++@pytest.mark.parametrize("bad", ["nan", "inf", "-inf", "0", "-1", "abc"]) ++def test_bad_interval_rejected_at_import(monkeypatch, bad): ++ with pytest.raises(ValueError): ++ _fresh(monkeypatch, VLLM_E2_GUARD_INTERVAL=bad) ++ ++ ++def test_reader_loss_arms_the_stop(monkeypatch): ++ gt = _fresh(monkeypatch, VLLM_E2_GUARD_MODE="warn") # the reader exists only when opted in (default off) ++ device = torch.device("cuda:0") ++ gt.prealloc(device) ++ gt.reset_for_test() ++ ++ # v2: the cadence reader no longer reads the device buffer itself; it ++ # polls the NonBlockingReader. Inject the read failure there. ++ def broken_poll(key): ++ raise RuntimeError("simulated D2H failure") ++ ++ monkeypatch.setattr(gt._nb_reader, "poll_one", broken_poll) ++ gt._poll_once() ++ assert device in gt._read_disabled ++ with pytest.raises(RuntimeError, match="fatal-site contract"): ++ gt.check_pending() ++ ++ ++def test_drain_reports_last_step_hit_and_raises_for_fatal_site(monkeypatch): ++ gt = _fresh(monkeypatch, VLLM_E2_GUARD_MODE="warn") # warn mode (opt-in since v2.0.1): a plain site must NOT raise ++ device = torch.device("cuda:0") ++ gt.prealloc(device) ++ gt.reset_for_test() ++ buf = gt._buffers[device] ++ # a hit at a non-fatal site, written as the kernels would (count, row, val) ++ plain = next(s for s in gt.SITE_NAMES if s not in gt._ALWAYS_FATAL_SITES) ++ buf[plain, 0] = 3 ++ snap = gt.drain(raise_on_fatal=True) ++ assert any(v[0] == 3 for v in snap.values()) ++ gt.check_pending() # nothing armed under warn for a plain site ++ # a hit at the always-fatal RoPE site: drain must raise when asked ++ buf[gt.SITE_QSA_ROPE_POS, 0] = 1 ++ with pytest.raises(RuntimeError, match="FINAL DRAIN"): ++ gt.drain(raise_on_fatal=True) ++ # and the non-raising form still arms the stop for the engine thread ++ gt.reset_for_test() ++ buf[gt.SITE_QSA_ROPE_POS, 0] = 1 ++ gt.drain() ++ with pytest.raises(RuntimeError): ++ gt.check_pending() ++ ++ ++def test_drain_is_silent_under_off(monkeypatch): ++ gt = _fresh(monkeypatch, VLLM_E2_GUARD_MODE="off") ++ device = torch.device("cuda:0") ++ gt.prealloc(device) ++ gt._buffers[device][gt.SITE_QSA_ROPE_POS, 0] = 1 ++ assert gt.drain(raise_on_fatal=True) == {} +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index b66d552..f162cb5 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -83,6 +83,7 @@ from vllm.v1.watermarking.spec_decode import ( + create_speculative_target_watermarker, + speculative_target_watermark_key, + ) ++from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.block_table import get_block_table_width + from vllm.v1.worker.cp_utils import check_attention_cp_compatibility + from vllm.v1.worker.gpu import pcp_manager as pcp +@@ -202,6 +203,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) + + self.device = device ++ # Bounds-guard instrumentation: allocate the guard counters now, while we are certainly ++ # not capturing. The thread that READS them starts much later, from ++ # gpu_worker after capture. See guard_telemetry.start_reader. ++ guard_telemetry.prealloc(device) + self.dtype = self.model_config.dtype + self.kv_cache_dtype = self.dtype + if self.cache_config.cache_dtype != "auto": +@@ -1639,6 +1644,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + valid_dummy_state_slots: bool = False, + ) -> ModelRunnerOutput | IntermediateTensors | None: + if not dummy_run: ++ # Bounds-guard instrumentation: raise on the engine thread if a bounds guard fired. ++ # A Python bool; adds no device sync to the step. Late and ++ # rank-local -- see guard_telemetry.check_pending. ++ guard_telemetry.check_pending() + # Update the request states. + self.update_pp_decode_requests() + self.finish_requests(scheduler_output) +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 2a3531b..f090135 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -88,6 +88,7 @@ from vllm.v1.outputs import ( + ModelRunnerOutput, + ) + from vllm.v1.utils import compute_iteration_details, report_usage_stats ++from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel + from vllm.v1.worker.startup_plan import ( + maybe_apply_startup_plan, +@@ -943,6 +944,13 @@ class Worker(WorkerBase): + # gate so subsequent `execute_model` / `sample_tokens` calls enforce it. + enable_gpu_sync_check() + ++ # Bounds-guard instrumentation: same moment, same reason. The guard counters were ++ # allocated in the model runner constructor, but the thread that READS ++ # them must not exist during capture -- a .tolist() from a second ++ # thread while a capture is active is what cudaStreamCaptureModeGlobal ++ # rejects. Nothing is read before this line; the kernels still count. ++ guard_telemetry.start_reader() ++ + # Startup is done; steady-state serving gets no benefit from torch + # intra-op parallelism. + set_torch_threads_for_runtime() +@@ -1481,6 +1489,15 @@ class Worker(WorkerBase): + def shutdown(self) -> None: + gc.unfreeze() + ++ # Bounds-guard instrumentation: synchronous final read of the guard counters, so ++ # a hit on the LAST step (never polled by check_pending, possibly ++ # ahead of the reader's next tick) is reported before teardown. Logs ++ # "FINAL DRAIN"; never raises here -- shutdown must complete. ++ try: ++ guard_telemetry.drain() ++ except Exception: # noqa: BLE001 ++ logger.exception("guard telemetry FINAL DRAIN failed") ++ + # has_kv_transfer_group can be None during interpreter shutdown. + if ensure_kv_transfer_shutdown is not None: + ensure_kv_transfer_shutdown() +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +new file mode 100644 +index 0000000..60fb094 +--- /dev/null ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -0,0 +1,620 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Persistent, cudagraph-safe counters for the bounds-guard instrumentation series. ++ ++Every guard in the series MASKS an out-of-range access so the SM cannot fault ++on it, and records that it did. A guard that masks silently turns a crash into ++a wrong answer and tells nobody, so the recording is not optional. ++ ++DESIGN -- follows the QSA KV-clip counter (models/qwen4_exp/nvidia/qsa.py ++:144-152, :205-225), the one instrument in this tree that already survives FULL ++cudagraph replay: ++ ++ * counters are PERSISTENT int64 GPU tensors, allocated eagerly by prealloc() ++ from the model runner constructor, so nothing is allocated inside a capture ++ and the data_ptr is stable across replays; ++ * the kernels do nothing but captured atomics, so replay keeps counting; ++ * a daemon thread reads them on a cadence, OFF the engine thread, so no ++ device->host sync is added anywhere on the step path; ++ * that thread is started only AFTER warmup and cudagraph capture, from ++ gpu_worker beside enable_gpu_sync_check(). See start_reader(). ++ ++The first design allocated a fresh counter per launch and read it back ++inline with .tolist(). That was wrong twice over: the D2H aborts cudagraph ++capture at startup, and under replay the Python host code never runs, so the ++counter was simultaneously re-zeroed every step and never read. Guards 2, 3 and ++4 sit inside the captured forward (QSA declares AttentionCGSupport.UNIFORM_BATCH, ++nvidia/qsa.py:286), which is exactly where it mattered. ++ ++NEGATIVE VALUES. Slots 1 and 2 are written with tl.atomic_max into a zeroed ++buffer, so a negative row or index would be lost. Every recorded value is ++therefore stored biased by GUARD_BIAS and the reader subtracts it; a slot that ++is still 0 means "nothing recorded", which is distinguishable from a recorded ++0 (GUARD_BIAS). Kernel modules must bind the bias as a tl.constexpr -- a plain ++module-level int is not readable from a @triton.jit body. ++ ++MODES -- VLLM_E2_GUARD_MODE: ++ raise the reader logs and arms a flag; a LATER execute_model on ++ the engine thread raises. Mirrors VLLM_RAISE_ON_LOGIT_NANS in intent: ++ a masked fault that silently returns a wrong answer is worse for a ++ serving product than a loud stop. Read check_pending below before ++ using it -- the raise is late and it is local to one rank. ++ Targeted repro boots. ++ warn log only -- EXCEPT for the sites in _ALWAYS_FATAL_SITES, which arm ++ the raise whatever the mode says, because their masked value has ++ already been written somewhere durable. USE THIS FOR LONG STRESS ++ BOOTS. ++ off (DEFAULT since v2: the reader thread ++ and the QSA clip reader cost ~0.3 ms per decode step on TP2xPP2; opt in ++ with VLLM_E2_GUARD_MODE=warn for diagnosis) ++ no reader thread. The kernels still mask and still count; nothing is ++ read -- INCLUDING the sites in _ALWAYS_FATAL_SITES. Use only when you ++ have decided you do not want the evidence. ++ ++WHICH RUNNER. prealloc() is called from GPUModelRunner.__init__ and ++check_pending() from its execute_model, both in vllm/v1/worker/gpu/ -- the V2 ++runner. On a V1 runner the counters are allocated lazily by the first guarded ++launch instead (correct, but possibly on a capturing stream) and NOTHING calls ++check_pending, so MODE=raise degrades to MODE=warn and _ALWAYS_FATAL_SITES ++never stops anything. The guards themselves still mask and count on both. Any ++release note for this series has to say so. ++""" ++ ++import math ++import os ++import threading ++import time ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++# Bias applied to every recorded row/value so negatives survive atomic_max on a ++# zero-initialised buffer. Any int32-range value plus this is > 0. ++GUARD_BIAS = 1 << 31 ++ ++# One (n_masked, row, value) triple per guarded kernel. ++SLOTS_PER_SITE = 3 ++ ++SITE_MAMBA_COPY = 0 ++SITE_GDN_DECODE_STATE = 1 ++SITE_QSA_PRE_INDEXER = 2 ++SITE_QSA_PREFILL = 3 ++SITE_QSA_DECODE = 4 ++SITE_BT_GATHER = 5 ++SITE_SLOT_MAPPING = 6 ++SITE_KV_ZEROER = 7 ++SITE_GDN_RECURRENT_STATE = 8 ++SITE_QSA_ROPE_POS = 9 ++# The ring-skip telemetry: NOT a masked fault -- a suffix-commit row the writer ++# skipped because the token's state slot was < 0 or >= ring capacity. ++# The skip itself is the earlier behaviour; the count is the instrument ++# for the v2.2 ring-row reconciliation. Telemetry only: warn. ++SITE_QSA_RING_SKIP = 10 ++NUM_SITES = 11 ++ ++SITE_NAMES = { ++ SITE_MAMBA_COPY: "_copy_mamba_state_block (block-table row/column)", ++ SITE_GDN_DECODE_STATE: "fused_sigmoid_gating_delta_rule_update_kernel " ++ "(ssm_state_indices column)", ++ SITE_QSA_PRE_INDEXER: "_qsa_pre_indexer_kernel (k_work request row)", ++ SITE_QSA_PREFILL: "_qsa_mqa_paged_prefill_kernel (page-table page id)", ++ SITE_QSA_DECODE: "_qsa_mqa_paged_uniform_kernel (page-table page id)", ++ SITE_BT_GATHER: "_gather_block_tables_kernel (request row / num_blocks)", ++ SITE_SLOT_MAPPING: "_compute_slot_mappings_kernel (request-state row)", ++ SITE_KV_ZEROER: "_zero_kv_blocks_kernel (KV pool block id)", ++ SITE_GDN_RECURRENT_STATE: "fused_recurrent_gated_delta_rule_fwd_kernel " ++ "(ssm_state_indices column)", ++ SITE_QSA_ROPE_POS: "_norm_rope in _qsa_pre_indexer_kernel " ++ "(rope position vs cos/sin table rows; slot 2 = axis 0=t 1=h 2=w)", ++ SITE_QSA_RING_SKIP: "_qsa_pre_indexer_kernel suffix commit SKIPPED " ++ "(state slot < 0 or >= ring capacity; slot 1 = token index, " ++ "slot 2 = state slot value)", ++} ++ ++ ++# Release defaults: off, 0.05 s. ++# Earlier revisions shipped raise / 0.5 while every qualification boot set both ++# variables explicitly, so the defaults were never the thing being tested. ++_MODE = os.environ.get("VLLM_E2_GUARD_MODE", "off").strip().lower() ++if _MODE not in ("raise", "warn", "off"): ++ # fail closed: a set-but-broken gate must not silently disarm the instrument ++ raise ValueError( ++ f"VLLM_E2_GUARD_MODE must be raise|warn|off, got {_MODE!r}" ++ ) ++ ++ ++def _parse_interval(text: str) -> float: ++ """Seconds between reader polls: a finite number > 0, or ValueError. ++ ++ float() accepts "nan", "inf" and negatives; the first two reach ++ time.sleep() and kill the reader daemon, a negative one raises there ++ instead. Reject them at import, where a broken ++ setting is a loud boot failure instead of a silently absent instrument. ++ """ ++ try: ++ value = float(text) ++ except ValueError as exc: ++ raise ValueError("VLLM_E2_GUARD_INTERVAL must be seconds") from exc ++ if not math.isfinite(value) or value <= 0: ++ raise ValueError( ++ f"VLLM_E2_GUARD_INTERVAL must be a finite number > 0, got {text!r}" ++ ) ++ return value ++ ++ ++_INTERVAL = _parse_interval(os.environ.get("VLLM_E2_GUARD_INTERVAL", "0.05")) ++ ++_buffers: dict[torch.device, torch.Tensor] = {} ++# Devices the reader has given up on. They stay in _buffers: see _reader. ++_read_disabled: set[torch.device] = set() ++_lock = threading.Lock() ++_reader_started = False ++_pending: str | None = None ++# Highest n_masked already reported, per (device, site). Module level, not a ++# _reader local, so reset_for_test can clear it: a test that zeroes the ++# counters but leaves this behind gets no WARNING on the next hit, because the ++# new count is below the remembered high-water mark. That is a test silently ++# measuring nothing. ++_seen: dict[tuple[torch.device, int], int] = {} ++ ++ ++class NonBlockingReader: ++ """Reads device counter tensors from a side thread without ever waiting on ++ the compute stream (v2 reader fix). ++ ++ The first reader called ``buf.tolist()`` on its cadence: a synchronous, ++ pageable device-to-host copy on the compute stream. The driver services ++ that copy holding the CUDA context lock until the stream's queued work has ++ completed. Under the PLE pull transport the stream can hold a GPU-side ++ wait whose producer needs the same lock (the connector thread's ++ ``event.synchronize()`` and the main thread's next kernel launch), so a ++ reader tick landing inside that window deadlocked the whole worker until ++ the executor timeout killed the engine (two hangs, both at ++ the end of long prefills; py-spy stacks confirmed the deadlock). ++ ++ Here every registered tensor gets a pinned host mirror and the copy runs ++ on a dedicated side stream with ``non_blocking=True``; completion is ++ observed with ``event.query()``. Nothing in this class blocks on the GPU. ++ The side stream has no dependency on the compute stream, so a pending ++ wait there cannot stall the copy, and a non-blocking pinned copy returns ++ from the driver immediately. The rows a tick reports are therefore the ++ values as of the previous tick's copy, and the three slots of one site ++ may come from slightly different moments while kernels are still ++ counting: the count is monotone and the maxima are lifetime maxima, so ++ this changes nothing the reader claims. ++ """ ++ ++ def __init__(self) -> None: ++ self._streams: dict[torch.device, "torch.cuda.Stream"] = {} ++ # key -> [device_tensor, pinned_host_tensor | None, event | None, inflight] ++ self._jobs: dict[object, list] = {} ++ self._lock = threading.Lock() ++ ++ def register(self, key: object, tensor: torch.Tensor) -> None: ++ with self._lock: ++ if key in self._jobs: ++ return ++ if tensor.device.type != "cuda": ++ self._jobs[key] = [tensor, None, None, False] ++ return ++ dev = tensor.device ++ # A stream/event/pinned allocation inside a cudagraph ++ # capture is not capture-safe. prealloc() registers before capture ++ # and the reader thread registers lazily (it is never capturing); ++ # any other caller reaching this during a capture is a bug. ++ if torch.cuda.is_current_stream_capturing(): ++ raise RuntimeError( ++ "guard_telemetry: NonBlockingReader.register() called during " ++ "cudagraph capture; prealloc() must run before capture" ++ ) ++ if dev not in self._streams: ++ with torch.cuda.device(dev): ++ self._streams[dev] = torch.cuda.Stream(device=dev) ++ host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True) ++ with torch.cuda.device(dev): ++ event = torch.cuda.Event() ++ self._jobs[key] = [tensor, host, event, False] ++ ++ def registered(self, key: object) -> bool: ++ return key in self._jobs ++ ++ def poll_one(self, key: object) -> list | None: ++ """Return the rows of a completed copy for ``key`` and start the next ++ one; return None when the previous copy is still in flight. Never ++ blocks.""" ++ job = self._jobs[key] ++ tensor, host, event, inflight = job ++ if host is None: # CPU tensor: nothing can block ++ return tensor.tolist() ++ rows = None ++ if inflight: ++ if not event.query(): ++ return None ++ rows = host.tolist() ++ stream = self._streams[tensor.device] ++ with torch.cuda.device(tensor.device), torch.cuda.stream(stream): ++ host.copy_(tensor, non_blocking=True) ++ event.record(stream) ++ job[3] = True ++ return rows ++ ++ def poll(self) -> dict[object, list]: ++ out = {} ++ with self._lock: ++ keys = list(self._jobs) ++ for key in keys: ++ rows = self.poll_one(key) ++ if rows is not None: ++ out[key] = rows ++ return out ++ ++ def read_now(self, key: object) -> list: ++ """Blocking read for orderly shutdown / tests ONLY: copies through the ++ pinned mirror on the side stream and waits for THAT copy (never a ++ pageable copy, never a wait on the compute stream). Callers sync the ++ device first when they need the latest kernel writes.""" ++ job = self._jobs[key] ++ tensor, host, event, _ = job ++ if host is None: ++ return tensor.tolist() ++ stream = self._streams[tensor.device] ++ with torch.cuda.device(tensor.device), torch.cuda.stream(stream): ++ host.copy_(tensor, non_blocking=True) ++ event.record(stream) ++ event.synchronize() ++ job[3] = False ++ return host.tolist() ++ ++ def reset(self) -> None: ++ """Test helper: drop in-flight copies so a zeroed counter is not ++ reported from a stale mirror.""" ++ for job in self._jobs.values(): ++ if job[1] is not None: ++ if job[3]: ++ job[2].synchronize() ++ job[1].zero_() ++ job[3] = False ++ ++ ++_nb_reader = NonBlockingReader() ++# device -> the copy in which a count increase was first seen (reported on the next tick) ++_confirm_pending: dict[torch.device, list] = {} ++ ++# Sites that arm the pending raise even under VLLM_E2_GUARD_MODE=warn. warn is ++# the right default for a guard that turns a crash into a locally wrong answer, ++# because the run continues and produces more evidence. It is the wrong default ++# for a guard whose masked value has already been WRITTEN somewhere durable -- ++# the damage outlives the step and every later step reads it, so "continue and ++# collect evidence" means "collect evidence about corrupted state". ++# ++# READ THE THREE LIMITS BEFORE RELYING ON THIS. The name says ++# "always fatal"; it is not a promise of any of the following: ++# ++# 1. MODE=off NEVER ARMS IT. off starts no reader, and the reader is the only ++# thing that inspects a counter, so a site in this set is exactly as ++# unheard as any other under off. off means "the kernels mask and count and ++# nobody looks", including here. Do not set off on a boot that is relying ++# on a fatal site to stop it. ++# 2. IT IS RANK-LOCAL. Counters are per process, so only the rank whose GPU ++# fired raises. Under PP/TP its peers received the same execute_model RPC ++# and block in recv/allreduce until the NCCL watchdog fires, minutes later. ++# What you get is ONE RuntimeError on the firing rank plus a multi-rank ++# watchdog timeout, and the evidence is the WARNING in the firing rank's ++# log -- find that log first, not the watchdog traceback. ++# 3. IT IS LATE BY UP TO VLLM_E2_GUARD_INTERVAL. The reader polls; the raise ++# lands on a later step. For a durable-damage site that window is how much ++# corrupted state gets written after the first hit, so a boot with such a ++# site armed should keep VLLM_E2_GUARD_INTERVAL at 0.05 (the default), not raise ++# it. The cost is 20 host-side reads a second off the engine thread. ++_ALWAYS_FATAL_SITES: set[int] = set() ++ ++# A clamped rope position is not a locally wrong answer that the next step ++# forgets. On the pooled path the rotated key is STORED into the compressed ++# cache, so every later step reads a wrongly-rotated key for that block, and ++# the ring row that produced it is still there to do it again. Continuing past ++# this one collects evidence about corrupted state, so it aborts the stage ++# whatever VLLM_E2_GUARD_MODE says. (bounds-guard instrumentation, rope site) ++_ALWAYS_FATAL_SITES.add(SITE_QSA_ROPE_POS) ++ ++ ++def _normalize(device: torch.device) -> torch.device: ++ """torch.device("cuda") and torch.device("cuda:0") are different dict keys. ++ ++ A miss here would allocate a second buffer lazily, on whatever stream is ++ current -- the replay defect re-entering through a typo. ++ """ ++ device = torch.device(device) ++ if device.type == "cuda" and device.index is None: ++ return torch.device("cuda", torch.cuda.current_device()) ++ return device ++ ++ ++def counter(site: int, device: torch.device) -> torch.Tensor: ++ """Persistent [3] int64 slice for one guarded kernel. ++ ++ Allocation should already have happened in prealloc(); the lazy path is a ++ fallback so a kernel reached from a context the runner did not construct ++ still counts rather than crashes. ++ """ ++ device = _normalize(device) ++ buf = _buffers.get(device) ++ if buf is None: ++ with _lock: ++ buf = _buffers.get(device) ++ if buf is None: ++ buf = torch.zeros( ++ (NUM_SITES, SLOTS_PER_SITE), dtype=torch.int64, device=device ++ ) ++ _buffers[device] = buf ++ # The reader mirrors this buffer through a pinned ++ # host copy on a side stream (registered by prealloc() before ++ # capture, or lazily by the reader thread -- never here, where ++ # a capture may be in progress). ++ return buf[site] ++ ++ ++def prealloc(device: torch.device) -> None: ++ """Allocate every site counter eagerly, before any cudagraph capture. ++ ++ counter() allocates lazily on first use. A guarded kernel whose FIRST ++ execution happens inside a capture would allocate on the capturing stream, ++ which is exactly the failure mode to avoid. The model ++ runner calls this at construction, so by capture time every buffer exists ++ and counter() is a dict lookup plus a view. ++ """ ++ for site in range(NUM_SITES): ++ counter(site, device) ++ device = _normalize(device) ++ if not _nb_reader.registered(device): ++ _nb_reader.register(device, _buffers[device]) ++ ++ ++def start_reader() -> None: ++ """Start the cadence reader. Called AFTER warmup and cudagraph capture. ++ ++ The reader must not exist during capture. It copies device counters to the ++ host, and a device-to-host copy from a second thread while a capture is ++ active is exactly the call class cudaStreamCaptureModeGlobal -- the default ++ for torch.cuda.graph -- rejects. The QSA clip-counter reader survives it in ++ practice, but that one is maintainer-gated and this one can be enabled by any user, so ++ the precedent is evidence, not proof. Starting after capture deletes the ++ question rather than arguing it. ++ ++ Between prealloc() and this call the kernels still count; nothing reads. ++ """ ++ global _reader_started ++ if _reader_started or _MODE == "off": ++ return ++ _reader_started = True ++ threading.Thread( ++ target=_reader, name="bounds-guard-reader", daemon=True ++ ).start() ++ logger.info( ++ "bounds-guard reader started (mode=%s, interval=%.3gs).", ++ _MODE, _INTERVAL, ++ ) ++ ++ ++def _reader() -> None: ++ """Off-engine-thread cadence reader; drain() is the only other counter read.""" ++ while True: ++ try: ++ time.sleep(_INTERVAL) ++ _poll_once() ++ except Exception: # noqa: BLE001 -- the daemon itself must survive ++ # A reader that dies is an instrument ++ # that is not there, and _ALWAYS_FATAL_SITES depends on it. Arm ++ # the stop so the engine thread learns that the fatal-site ++ # contract is no longer held, then keep polling. ++ logger.exception("bounds-guard reader loop failed") ++ _arm( ++ "bounds-guard reader loop failed; the fatal-site contract " ++ "(_ALWAYS_FATAL_SITES) can no longer be enforced." ++ ) ++ time.sleep(max(_INTERVAL, 0.05)) ++ ++ ++def _report_rows(device: torch.device, rows: list[list[int]]) -> None: ++ """Report every (device, site) whose masked count advanced. Shared by the ++ cadence reader and the final drain.""" ++ for site, (n_masked, biased_row, biased_val) in enumerate(rows): ++ if n_masked <= _seen.get((device, site), 0): ++ continue ++ _seen[(device, site)] = n_masked ++ row = biased_row - GUARD_BIAS if biased_row else None ++ val = biased_val - GUARD_BIAS if biased_val else None ++ msg = ( ++ f"bounds guard masked {n_masked} out-of-range access(es) " ++ f"in {SITE_NAMES.get(site, site)} on {device} " ++ f"(cumulative maxima: row-field {row}, index-field {val}; these " ++ f"are independent lifetime maxima, not one observed pair). The " ++ f"access was suppressed, so this is a MASKED FAULT and not a " ++ f"correct result -- the metadata is wrong upstream." ++ ) ++ logger.warning("%s", msg) ++ if _MODE == "raise" or site in _ALWAYS_FATAL_SITES: ++ _arm(msg) ++ ++ ++def _poll_once() -> None: ++ for device, buf in list(_buffers.items()): ++ if device in _read_disabled: ++ continue ++ try: ++ # A non-blocking side-stream copy; None means the ++ # previous copy is still in flight (report on a later tick). ++ if not _nb_reader.registered(device): ++ _nb_reader.register(device, buf) # reader thread: never capturing ++ rows = _nb_reader.poll_one(device) ++ if rows is None: ++ continue ++ # Caveat: the side-stream copy is not a coherent snapshot of a ++ # site's three slots -- the kernels write the count first, then ++ # the maxima. A copy taken between them pairs a new count with ++ # stale maxima. So a count increase seen in copy k is REPORTED ++ # from copy k+1 (one tick later, >= 50 ms after the increment, ++ # by which time the maxima of that hit are written). The fatal ++ # decision is count-based either way; only its report waits a tick. ++ pend = _confirm_pending.get(device) ++ rising = [site for site, (n, _r, _v) in enumerate(rows) if n > _seen.get((device, site), 0)] ++ if pend is None: ++ if rising: ++ _confirm_pending[device] = rows ++ continue ++ _confirm_pending.pop(device, None) ++ # report with the CONFIRMING copy's values (never lower than the ++ # pending copy's counts: the counters are monotone) ++ rows = [[max(a[0], b[0]), b[1], b[2]] for a, b in zip(pend, rows)] ++ except Exception: # noqa: BLE001 -- a reader must never kill the run ++ # Never silent: a reader that cannot read is an instrument that ++ # is not there, and this whole series depends on it. Log once, ++ # then stop polling that device. ++ # ++ # STOP POLLING IT, NOT STOP HOLDING IT. A prior version ++ # popped the buffer here, which drops the only durable Python ++ # reference: counter() hands out transient buf[site] views that ++ # are released after each launch, and under FULL replay nothing ++ # holds one at all. The storage would return to the caching ++ # allocator, be reused, and the CAPTURED atomics would keep ++ # writing to the baked address -- the memory-corrupting half of ++ # replay defect, re-entering through the error handler. ++ logger.exception( ++ "bounds-guard reader failed to read the counter on %s; " ++ "no further guard telemetry will be reported for it. The " ++ "buffer is kept alive: captured graphs still write to it.", ++ device, ++ ) ++ _read_disabled.add(device) ++ # Losing the reader on a device ++ # silently voids the fatal-site contract for it. Arm the stop ++ # regardless of MODE (off never reaches here: no reader). ++ _arm( ++ f"bounds-guard reader lost the counter on {device}; the " ++ f"fatal-site contract (_ALWAYS_FATAL_SITES) is no longer " ++ f"enforced there." ++ ) ++ continue ++ _report_rows(device, rows) ++ ++ ++def _arm(msg: str) -> None: ++ global _pending ++ if _pending is None: ++ _pending = msg ++ ++ ++def check_pending() -> None: ++ """Engine-thread check. One Python bool per step; adds no sync. ++ ++ Called at the top of execute_model, ahead of anything that mutates state ++ and ahead of any collective, so the step it lands on aborts cleanly. ++ ++ TWO THINGS IT DOES NOT DO -- an earlier docstring claimed one of them and ++ was wrong: ++ ++ 1. It does NOT raise on the step that fired. The reader polls on a cadence, ++ so the raise lands up to VLLM_E2_GUARD_INTERVAL later, hundreds of ++ decode steps on. The scheduler state at the raise belongs to some later ++ step, not the offending one. Attribution comes from the counter and the ++ WARNING; the raise is a stop, not a snapshot. ++ 2. It does NOT coordinate across ranks. The counters are per process, so ++ only the rank whose GPU fired arms. Its PP/TP peers received the same ++ execute_model RPC, proceed into the forward, and block in recv/allreduce ++ until the NCCL watchdog fires -- minutes later, as a multi-rank timeout ++ that is a worse diagnostic than the Xid being chased. ++ ++ Hence: VLLM_E2_GUARD_MODE=warn for long stress boots; raise only for a ++ targeted repro where losing the group is acceptable and stopping is the ++ point. ++ """ ++ global _pending ++ if _pending is None: ++ return ++ msg = _pending ++ _pending = None ++ raise RuntimeError( ++ f"{msg} (Sites in _ALWAYS_FATAL_SITES and reader failures stop the " ++ f"rank under every mode except off; MODE=warn only softens the other " ++ f"sites.)" ++ ) ++ ++ ++def drain(*, raise_on_fatal: bool = False) -> dict[str, list[int]]: ++ """Synchronous final read of every counter, off the cadence (added ++ later). ++ ++ The cadence reader can miss a hit on the LAST step of a run: the forward ++ completes, its result is consumed, no further execute_model polls ++ check_pending, and the process exits before the reader's next tick. Call ++ this at stage end / orderly shutdown, outside capture. It synchronizes the ++ device, reports anything the reader has not, and returns the unbiased ++ snapshot. With raise_on_fatal it raises for a hit at a site in ++ _ALWAYS_FATAL_SITES or, under MODE=raise, at any site. MODE=off drains ++ nothing (no reporting under off, by contract). ++ """ ++ if _MODE == "off": ++ return {} ++ fatal: list[str] = [] ++ for device, buf in list(_buffers.items()): ++ try: ++ if device.type == "cuda": ++ torch.cuda.synchronize(device) ++ if not _nb_reader.registered(device): ++ _nb_reader.register(device, buf) ++ rows = _nb_reader.read_now(device) # pinned mirror, side stream only ++ except Exception: # noqa: BLE001 ++ logger.exception("bounds-guard FINAL DRAIN could not read %s", device) ++ fatal.append(f"reader lost {device}") ++ continue ++ _report_rows(device, rows) ++ for site, (n_masked, _r, _v) in enumerate(rows): ++ if n_masked and (_MODE == "raise" or site in _ALWAYS_FATAL_SITES): ++ fatal.append( ++ f"{SITE_NAMES.get(site, site)} on {device}: {n_masked}" ++ ) ++ snap = snapshot() ++ if fatal: ++ logger.error( ++ "bounds-guard FINAL DRAIN: fatal-site hits %s (snapshot %s)", ++ fatal, snap, ++ ) ++ if raise_on_fatal: ++ raise RuntimeError(f"bounds-guard FINAL DRAIN: fatal-site hits {fatal}") ++ else: ++ logger.info( ++ "bounds-guard FINAL DRAIN: no fatal-site hits (snapshot %s)", snap ++ ) ++ return snap ++ ++ ++def snapshot() -> dict[str, list[int]]: ++ """Test/diagnostic helper: current counters, unbiased. Syncs -- not for the ++ step path.""" ++ out: dict[str, list[int]] = {} ++ for device, buf in _buffers.items(): ++ for site, (n, row, val) in enumerate(buf.tolist()): ++ if n: ++ out[f"{device}:{SITE_NAMES.get(site, site)}"] = [ ++ n, ++ row - GUARD_BIAS if row else 0, ++ val - GUARD_BIAS if val else 0, ++ ] ++ return out ++ ++ ++def reset_for_test() -> None: ++ """Zero every counter, clear the reader's high-water marks, and disarm. ++ ++ All three, or the next hit is invisible: the reader only reports a count ++ ABOVE what it last saw. ++ """ ++ global _pending ++ for buf in _buffers.values(): ++ buf.zero_() ++ _nb_reader.reset() ++ _confirm_pending.clear() ++ _seen.clear() ++ _pending = None +-- +2.43.0 + diff --git a/patches/0002-qwen4_exp-FP8-E4M3-KV-reader-for-QSA-with-a-static-s.patch b/patches/0002-qwen4_exp-FP8-E4M3-KV-reader-for-QSA-with-a-static-s.patch new file mode 100644 index 0000000..ab09080 --- /dev/null +++ b/patches/0002-qwen4_exp-FP8-E4M3-KV-reader-for-QSA-with-a-static-s.patch @@ -0,0 +1,2001 @@ +From 7de1b4ee5f23d9ac5698b27c744f100577080e41 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:09:54 -0400 +Subject: [PATCH 002/103] qwen4_exp: FP8 E4M3 KV reader for QSA with a static + scale sidecar + +FP8 E4M3 KV read path for the Qwen4Exp sparse attention (QSA), indexer and +pre-indexer kernels on sm_86, dequantised with a static per-layer scale loaded +from a JSON sidecar, validated against the model's QSA layer set and applied +per pipeline stage. Includes the offline absmax collector that produces the +sidecar and a CUDA-graph-safe clipping counter. +--- + vllm/models/qwen4_exp/nvidia/model.py | 113 ++- + vllm/models/qwen4_exp/nvidia/ops/qsa.py | 213 +++++- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 882 ++++++++++++++++++++++ + vllm/models/qwen4_exp/nvidia/qsa.py | 487 +++++++++++- + 4 files changed, 1657 insertions(+), 38 deletions(-) + +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 54ee6f0..3f884cd 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -5,6 +5,8 @@ + from collections.abc import Iterable + from itertools import islice + ++import os ++ + import torch + from torch import nn + +@@ -79,7 +81,10 @@ from ..config import Qwen4ExpConfig + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm + from .ple_layer import Qwen4ExpPLELayer +-from .qsa import Qwen4ExpQSAAttention ++from .qsa import Qwen4ExpQSAAttention, load_qsa_static_kv_scales ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) + + + def without_modelopt_fp4( +@@ -158,6 +163,104 @@ _EXTRA_WEIGHTS_MAPPER = WeightsMapper( + ) + + ++_QSA_KV_SCALES_ENV = "VLLM_QSA_KV_SCALES" ++ ++ ++def _maybe_load_qsa_static_kv_scales(model: nn.Module) -> None: ++ """Apply a static QSA K/V scale sidecar when VLLM_QSA_KV_SCALES names one. ++ ++ Placed where the loader's own docstring requires: after checkpoint weights ++ are loaded, before profiling, warmup or cudagraph capture. Unset env var -> ++ no-op (the attention layer logs the default 1.0 per layer so the ++ uncalibrated state stays visible). Set but unusable -> the exception ++ propagates: a calibration that was requested and silently not applied is ++ worse than one never requested. ++ """ ++ path = os.environ.get(_QSA_KV_SCALES_ENV, "").strip() ++ if not path: ++ return ++ # The PLE offload worker builds the whole target on the meta device and loads ++ # only PLE weights through load_weights; a fill_ on its meta scale buffers ++ # would log success without calibrating any GPU worker. ++ local_qsa: list[Qwen4ExpQSAAttention] = [] ++ inner: Qwen4ExpModel | None = None ++ for module in model.modules(): ++ if isinstance(module, Qwen4ExpModel) and inner is None: ++ inner = module ++ if isinstance(module, Qwen4ExpQSAAttention): ++ local_qsa.append(module) ++ for layer in local_qsa: ++ k_scale = getattr(layer, "_k_scale", None) ++ if isinstance(k_scale, torch.Tensor) and k_scale.device.type == "meta": ++ logger.info( ++ "Qwen4Exp QSA: model is on meta (PLE offload worker); " ++ "static K/V scales are applied by the GPU workers only" ++ ) ++ return ++ if inner is None: ++ raise ValueError("Qwen4Exp QSA: no Qwen4ExpModel found for the sidecar") ++ if not inner._qsa_layer_ids: ++ # A requested calibration for a model without QSA layers is a ++ # configuration error, never a silent no-op. ++ raise ValueError( ++ f"Qwen4Exp QSA: {_QSA_KV_SCALES_ENV} is set but the model has no " ++ "QSA (indexer) attention layers" ++ ) ++ # Whole-model QSA layer set (all pipeline stages) = "..self_attn.attn" ++ # for every configured QSA index. The stack prefix is read from a LOCAL QSA ++ # layer's registered layer_name whenever one exists: this hook also runs ++ # inside the nested Qwen4ExpForCausalLM.load_weights, where a module path ++ # relative to that submodule ("model.layers") would not match the names ++ # registered under the outer prefix ("language_model.model.layers") - the ++ # earlier boot failure. Only a stage with no QSA layer ++ # falls back to the module path (then no local name can disagree). ++ if local_qsa: ++ stack, sep, tail = local_qsa[0].layer_name.rpartition(".layers.") ++ if not sep or not tail.split(".")[0].isdigit(): ++ raise ValueError( ++ f"Qwen4Exp QSA: unexpected layer name {local_qsa[0].layer_name!r}" ++ ) ++ layers_path = f"{stack}.layers" ++ else: ++ layers_path = next( ++ (name for name, module in model.named_modules() if module is inner.layers), ++ None, ++ ) ++ if layers_path is None: ++ raise ValueError("Qwen4Exp QSA: cannot locate the decoder layer stack") ++ expected = { ++ f"{layers_path}.{idx}.self_attn.attn" for idx in inner._qsa_layer_ids ++ } ++ if not local_qsa: ++ # A pipeline stage without a full-attention layer (only for a very ++ # uneven split): validate the complete sidecar here too, apply nothing. ++ logger.warning( ++ "Qwen4Exp QSA: this pipeline stage holds no QSA layer " ++ "(pipeline layers %d-%d); sidecar %s validated, nothing to apply", ++ inner.start_layer, ++ inner.end_layer, ++ path, ++ ) ++ # strict: every QSA layer of the whole model must be in the sidecar. An empty ++ # or partial file silently leaving layers at scale 1.0 is exactly the failure ++ # this call site exists to prevent. ++ applied = load_qsa_static_kv_scales( ++ model, path, strict=True, expected_layer_names=expected ++ ) ++ shown = ", ".join(sorted(applied)[:4]) ++ logger.info( ++ "Qwen4Exp QSA: applied static K/V scales from %s to %d of %d layer(s) " ++ "(pipeline layers %d-%d): %s%s", ++ path, ++ len(applied), ++ len(expected), ++ inner.start_layer, ++ inner.end_layer, ++ shown, ++ " ..." if len(applied) > 4 else "", ++ ) ++ ++ + class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): + """Qwen3Next MoE with Qwen4Exp HC validation.""" + +@@ -855,7 +958,9 @@ class Qwen4ExpForCausalLM( + self, + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) +- return loader.load_weights(weights, mapper=mapper) ++ loaded = loader.load_weights(weights, mapper=mapper) ++ _maybe_load_qsa_static_kv_scales(self) ++ return loaded + + + class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo): +@@ -1050,7 +1155,9 @@ class Qwen4ExpForConditionalGeneration( + self, + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) +- return loader.load_weights(weights, mapper=mapper) ++ loaded = loader.load_weights(weights, mapper=mapper) ++ _maybe_load_qsa_static_kv_scales(self) ++ return loaded + + @classmethod + def get_mamba_state_dtype_from_config( +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py +index d3986ba..2acc6f4 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py +@@ -10,14 +10,38 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, + triton_scalar_specialization_rep, + ) ++from vllm.platforms import current_platform + from vllm.triton_utils import HAS_TRITON, tl, triton + + ++@triton.jit ++def _decode_e4m3_to_bf16(b): ++ # Shifting the 7 magnitude bits left by 7 lands the E4M3 exponent in the ++ # FP16 exponent field and the 3 mantissa bits at the top of the FP16 ++ # mantissa field, giving the right number under the wrong exponent bias ++ # (E4M3 7 vs FP16 15). That is a constant factor of 2^8. E4M3 subnormals ++ # land on FP16 subnormals and the same factor renormalizes them, which is ++ # only true because sm_86 does not flush FP16 subnormals -- verified on ++ # device, not assumed. ++ # ++ # Bit-exact against torch on all 254 finite codepoints: every E4M3 value ++ # fits in BF16's 7 mantissa bits, and x256 is exact in FP16. ++ # ++ # Deviation: 0x7F/0xFF are NaN in e4m3fn but decode to +-480 here. The CUDA ++ # writer clamps to +-448 and never emits them, so this is unreachable from ++ # a cache written by reshape_and_cache_flash. ++ b = b.to(tl.uint16) ++ bits = ((b & 0x80) << 8) | ((b & 0x7F) << 7) ++ return (bits.to(tl.uint16).to(tl.float16, bitcast=True) * 256.0).to(tl.bfloat16) ++ ++ + @triton.jit(do_not_specialize=["num_rows", "num_requests"]) + def _qsa_sparse_paged_gqa_splitk_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, ++ k_scale_ptr, ++ v_scale_ptr, + indices_ptr, + block_table_ptr, + token_to_req_ptr, +@@ -50,6 +74,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_TILES: tl.constexpr, ++ KV_DTYPE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + ) -> None: +@@ -83,6 +108,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) + softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 + ++ if KV_DTYPE != 0: ++ k_scale_value = tl.load(k_scale_ptr) ++ v_scale_value = tl.load(v_scale_ptr) ++ + tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N)) + + for tile in range(split_id, tile_end, NUM_SPLITS): +@@ -111,33 +140,87 @@ def _qsa_sparse_paged_gqa_splitk_kernel( + valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) + # physical_page * block stride can overflow int32 for large caches. + safe_page = tl.maximum(physical_page, 0).to(tl.int64) +- keys = tl.load( +- k_cache_ptr +- + safe_page[None, :] * stride_k_block +- + page_offset[None, :] * stride_k_token +- + kv_head * stride_k_head +- + dim_offsets[:, None], +- mask=valid[None, :], +- other=0.0, +- ) +- values = tl.load( +- v_cache_ptr +- + safe_page[:, None] * stride_v_block +- + page_offset[:, None] * stride_v_token +- + kv_head * stride_v_head +- + dim_offsets[None, :], +- mask=valid[:, None], +- other=0.0, +- ) ++ if KV_DTYPE == 0: ++ keys = tl.load( ++ k_cache_ptr ++ + safe_page[None, :] * stride_k_block ++ + page_offset[None, :] * stride_k_token ++ + kv_head * stride_k_head ++ + dim_offsets[:, None], ++ mask=valid[None, :], ++ other=0.0, ++ ) ++ values = tl.load( ++ v_cache_ptr ++ + safe_page[:, None] * stride_v_block ++ + page_offset[:, None] * stride_v_token ++ + kv_head * stride_v_head ++ + dim_offsets[None, :], ++ mask=valid[:, None], ++ other=0.0, ++ ) ++ elif KV_DTYPE == 1: ++ key_bytes = tl.load( ++ k_cache_ptr ++ + safe_page[None, :] * stride_k_block ++ + page_offset[None, :] * stride_k_token ++ + kv_head * stride_k_head ++ + dim_offsets[:, None], ++ mask=valid[None, :], ++ other=0, ++ ) ++ # UNSCALED: k_scale is folded into the score scaling below. ++ keys = _decode_e4m3_to_bf16(key_bytes) ++ else: ++ keys = tl.load( ++ k_cache_ptr ++ + safe_page[None, :] * stride_k_block ++ + page_offset[None, :] * stride_k_token ++ + kv_head * stride_k_head ++ + dim_offsets[:, None], ++ mask=valid[None, :], ++ other=0.0, ++ ).to(tl.bfloat16) + scores = tl.dot(query, keys) + # Scaling scores avoids re-quantizing a scaled query to BF16. +- scores *= softmax_scale_log2 ++ if KV_DTYPE == 0: ++ scores *= softmax_scale_log2 ++ else: ++ # k_scale is a scalar, so scaling each SCORE is equivalent to ++ # scaling every key element -- but costs one multiply per score ++ # instead of one per element, and skips the fp32 round trip and ++ # the BF16 re-quantization of the scaled key (which for a ++ # non-power-of-two scale also lost mantissa bits). Strictly fewer ++ # roundings than the per-element form it replaces. ++ scores *= softmax_scale_log2 * k_scale_value + scores = tl.where(valid[None, :], scores, -1.0e20) + next_max = tl.maximum(max_value, tl.max(scores, axis=1)) + alpha = tl.math.exp2(max_value - next_max) + probabilities = tl.where( + valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 + ) ++ if KV_DTYPE == 1: ++ value_bytes = tl.load( ++ v_cache_ptr ++ + safe_page[:, None] * stride_v_block ++ + page_offset[:, None] * stride_v_token ++ + kv_head * stride_v_head ++ + dim_offsets[None, :], ++ mask=valid[:, None], ++ other=0, ++ ) ++ # UNSCALED: v_scale is folded into the epilogue below. ++ values = _decode_e4m3_to_bf16(value_bytes) ++ elif KV_DTYPE == 2: ++ values = tl.load( ++ v_cache_ptr ++ + safe_page[:, None] * stride_v_block ++ + page_offset[:, None] * stride_v_token ++ + kv_head * stride_v_head ++ + dim_offsets[None, :], ++ mask=valid[:, None], ++ other=0.0, ++ ).to(tl.bfloat16) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, +@@ -152,6 +235,13 @@ def _qsa_sparse_paged_gqa_splitk_kernel( + accumulator / tl.maximum(normalizer[:, None], 1.0e-20), + 0.0, + ) ++ if KV_DTYPE != 0: ++ # v_scale folded out of the per-element V decode: the attention output ++ # is linear in V, so one multiply per output element replaces one per ++ # cache element per tile. Safe under split-K too -- every partial is ++ # scaled by the same constant before the LSE merge, which is itself ++ # linear in the partial outputs. ++ normalized_output *= v_scale_value + output_mask = head_offsets[:, None] < GROUP_SIZE + if NUM_SPLITS == 1: + # Preserve the unfused path's BF16 attention-output rounding before +@@ -483,10 +573,18 @@ def qsa_sparse_paged_attention( + token_to_req: torch.Tensor, + use_prefill_config: bool, + out: torch.Tensor | None = None, ++ k_scale: torch.Tensor | None = None, ++ v_scale: torch.Tensor | None = None, ++ _force_num_warps: int | None = None, + *, + output_gate: torch.Tensor, + ) -> torch.Tensor: +- """Run sparse GQA directly over paged BF16 K/V caches. ++ """Run sparse GQA over paged BF16 or FP8 (E4M3 via static scalar scales) K/V caches. ++ ++ ``_force_num_warps`` overrides the tuned launch profile for the warp-count ++ equivalence test only; production callers leave it None. Hardware policy ++ (E4M3 only on CUDA sm_86) is enforced once at Qwen4ExpQSAAttention / ++ Qwen4ExpQSAFlashAttentionImpl construction, not per call. + + logical_indices is the PACKED selection buffer: [rows, selection_width + 1] + with the trailing column holding each row's valid-entry count (written by +@@ -510,7 +608,26 @@ def qsa_sparse_paged_attention( + raise ValueError("QSA sparse attention requires valid grouped-query heads") + head_dim = q.shape[2] + assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 +- assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16 ++ if q.dtype != torch.bfloat16: ++ raise ValueError("QSA sparse attention requires BF16 queries") ++ if k_cache.dtype != v_cache.dtype: ++ raise ValueError("QSA sparse attention requires matching K/V dtypes") ++ fp8_dtype = current_platform.fp8_dtype() ++ if k_cache.dtype == torch.bfloat16: ++ kv_dtype = 0 ++ elif k_cache.dtype in (torch.uint8, fp8_dtype): ++ kv_dtype = 1 ++ # Present platform E4M3 storage as bytes so Triton never materializes ++ # the unsupported fp8e4nv type on SM86. ++ if k_cache.dtype == fp8_dtype: ++ k_cache = k_cache.view(torch.uint8) ++ v_cache = v_cache.view(torch.uint8) ++ elif k_cache.dtype == torch.float8_e5m2: ++ kv_dtype = 2 ++ else: ++ raise ValueError( ++ "QSA supports only BF16, uint8/E4M3, and E5M2 K/V caches" ++ ) + assert logical_indices.dtype == block_table.dtype == torch.int32 + assert token_to_req.dtype == torch.int32 + assert q.device == k_cache.device == v_cache.device +@@ -531,12 +648,37 @@ def qsa_sparse_paged_attention( + if not q.shape[0]: + return out + ++ if kv_dtype == 0: ++ # Scale pointers are compile-time dead in the unchanged BF16 kernel. ++ k_scale_ptr = q ++ v_scale_ptr = q ++ else: ++ if k_scale is None: ++ k_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device) ++ if v_scale is None: ++ v_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device) ++ if ( ++ k_scale.numel() != 1 ++ or v_scale.numel() != 1 ++ or k_scale.dtype != torch.float32 ++ or v_scale.dtype != torch.float32 ++ or k_scale.device != q.device ++ or v_scale.device != q.device ++ ): ++ raise ValueError( ++ "QSA K/V scales must be scalar float32 tensors on the query device" ++ ) ++ k_scale_ptr = k_scale ++ v_scale_ptr = v_scale ++ + group_size = q.shape[1] // k_cache.shape[2] + block_m = triton.next_power_of_2(group_size) + selection_width = logical_indices.shape[1] - 1 # trailing column is the count + block_n, partial_warps, num_tiles, num_splits = _select_config( + q.shape[0], k_cache.shape[2], use_prefill_config, selection_width + ) ++ if _force_num_warps is not None: ++ partial_warps = _force_num_warps + + # Split=1 writes output directly and compiles out all workspace accesses. + if num_splits == 1: +@@ -559,6 +701,8 @@ def qsa_sparse_paged_attention( + q, + k_cache, + v_cache, ++ k_scale_ptr, ++ v_scale_ptr, + logical_indices, + block_table, + token_to_req, +@@ -591,6 +735,7 @@ def qsa_sparse_paged_attention( + NUM_QUERY_HEADS=q.shape[1], + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, ++ KV_DTYPE=kv_dtype, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=partial_warps, +@@ -649,13 +794,32 @@ def warmup_qsa_sparse_paged_attention( + q_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) ++ # Mirror qsa_sparse_paged_attention's dtype dispatch: E4M3 storage is launched ++ # as a uint8 view with scalar float32 scale pointers. ++ if key_cache.dtype == torch.bfloat16: ++ warm_kv_dtype = 0 ++ warm_cache_dtype = key_cache.dtype ++ elif key_cache.dtype in (torch.uint8, current_platform.fp8_dtype()): ++ warm_kv_dtype = 1 ++ warm_cache_dtype = torch.uint8 ++ elif key_cache.dtype == torch.float8_e5m2: ++ warm_kv_dtype = 2 ++ warm_cache_dtype = key_cache.dtype ++ else: ++ raise ValueError("QSA warmup: unsupported K/V cache dtype") ++ if warm_kv_dtype == 0: ++ k_scale_ptr = q_ptr ++ v_scale_ptr = q_ptr ++ else: ++ k_scale_ptr = TritonWarmupTensor(torch.float32) ++ v_scale_ptr = TritonWarmupTensor(torch.float32) + k_cache_ptr = TritonWarmupTensor( +- key_cache.dtype, ++ warm_cache_dtype, + shape=tuple(key_cache.shape), + strides=tuple(key_cache.stride()), + ) + v_cache_ptr = TritonWarmupTensor( +- value_cache.dtype, ++ warm_cache_dtype, + shape=tuple(value_cache.shape), + strides=tuple(value_cache.stride()), + ) +@@ -695,6 +859,8 @@ def warmup_qsa_sparse_paged_attention( + q_ptr, + k_cache_ptr, + v_cache_ptr, ++ k_scale_ptr, ++ v_scale_ptr, + indices_ptr, + block_table_ptr, + token_to_req_ptr, +@@ -727,6 +893,7 @@ def warmup_qsa_sparse_paged_attention( + NUM_QUERY_HEADS=num_query_heads, + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, ++ KV_DTYPE=warm_kv_dtype, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=warps, +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index 2ec8790..bec35e6 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -4,6 +4,8 @@ + + from collections.abc import Sequence + ++import os ++ + import torch + from torch import nn + +@@ -63,6 +65,886 @@ class Qwen4ExpPLEGroupedNorm(nn.Module): + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + ++class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase): ++ """FP8 PLE embedding with one global checkpoint scale.""" ++ ++ def create_weights( ++ self, ++ layer: nn.Module, ++ input_size_per_partition: int, ++ output_partition_sizes: list[int], ++ input_size: int, ++ output_size: int, ++ params_dtype: torch.dtype, ++ **extra_weight_attrs, ++ ) -> None: ++ del input_size, output_size, params_dtype ++ weight_loader = extra_weight_attrs.get("weight_loader") ++ weight = create_fp8_weight_parameter( ++ sum(output_partition_sizes), input_size_per_partition, weight_loader ++ ) ++ layer.register_parameter("weight", weight) ++ ++ weight_scale = create_fp8_scale_parameter( ++ PerTensorScaleParameter, ++ output_partition_sizes, ++ input_size_per_partition, ++ None, ++ weight_loader, ++ scale_dtype=torch.float32, ++ ) ++ layer.register_parameter("weight_scale", weight_scale) ++ ++ def process_weights_after_loading(self, layer: nn.Module) -> None: ++ """Reject FP8 PLE checkpoints without a global scale.""" ++ sentinel = torch.finfo(torch.float32).min ++ if torch.any(layer.weight_scale == sentinel): ++ raise ValueError("FP8 PLE checkpoint is missing its global scale") ++ ++ def apply( ++ self, ++ layer: nn.Module, ++ x: torch.Tensor, ++ bias: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ raise NotImplementedError("PLE FP8 weights only support embedding lookup") ++ ++ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: ++ return F.embedding(input_, layer.weight) ++ ++ ++_NVFP4_BLOCK_SIZE = 16 ++ ++ ++class Qwen4ExpPLENVFp4EmbeddingMethod(QuantizeMethodBase): ++ """NVFP4 PLE embedding kept packed at runtime. ++ ++ The table stays as uint8 codes with FP8 block scales (16 values per ++ block) and one FP32 global scale. A lookup only gathers the packed rows; ++ the PLE layer dequantizes them after lookup. ++ """ ++ ++ def create_weights( ++ self, ++ layer: nn.Module, ++ input_size_per_partition: int, ++ output_partition_sizes: list[int], ++ input_size: int, ++ output_size: int, ++ params_dtype: torch.dtype, ++ **extra_weight_attrs, ++ ) -> None: ++ del input_size, output_size, extra_weight_attrs ++ if input_size_per_partition % _NVFP4_BLOCK_SIZE: ++ raise ValueError( ++ "NVFP4 PLE embedding requires the embedding dim to be a " ++ f"multiple of {_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}" ++ ) ++ self.params_dtype = params_dtype ++ self.head_dim = input_size_per_partition ++ self.packed_row_width = ( ++ input_size_per_partition // 2 ++ + input_size_per_partition // _NVFP4_BLOCK_SIZE ++ ) ++ rows = sum(output_partition_sizes) ++ weight = nn.Parameter( ++ torch.empty(rows, input_size_per_partition // 2, dtype=torch.uint8), ++ requires_grad=False, ++ ) ++ set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) ++ layer.register_parameter("weight", weight) ++ ++ weight_scale = nn.Parameter( ++ torch.empty( ++ rows, ++ input_size_per_partition // _NVFP4_BLOCK_SIZE, ++ dtype=torch.float8_e4m3fn, ++ ), ++ requires_grad=False, ++ ) ++ set_weight_attrs(weight_scale, {"input_dim": 1, "output_dim": 0}) ++ layer.register_parameter("weight_scale", weight_scale) ++ ++ weight_scale_2 = nn.Parameter( ++ torch.empty((), dtype=torch.float32), requires_grad=False ++ ) ++ layer.register_parameter("weight_scale_2", weight_scale_2) ++ self._lut: torch.Tensor | None = None ++ ++ def process_weights_after_loading(self, layer: nn.Module) -> None: ++ # Build the lookup table before CUDA graph capture. ++ self._lut = torch.tensor( ++ _FP4_VALUES, dtype=torch.float32, device=layer.weight.device ++ ) ++ ++ def apply( ++ self, ++ layer: nn.Module, ++ x: torch.Tensor, ++ bias: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ raise NotImplementedError("PLE NVFP4 weights only support embedding lookup") ++ ++ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: ++ """Gather packed rows for deferred dequantization.""" ++ codes = F.embedding(input_, layer.weight) ++ scales = F.embedding(input_, layer.weight_scale) ++ return torch.cat((codes, scales.view(torch.uint8)), dim=-1) ++ ++ def dequantize( ++ self, ++ packed_rows: torch.Tensor, ++ scale_2: torch.Tensor, ++ output_dtype: torch.dtype, ++ ) -> torch.Tensor: ++ """Dequantize rows containing NVFP4 codes and block scales.""" ++ lut = self._lut ++ if lut is None or lut.device != packed_rows.device: ++ lut = torch.tensor( ++ _FP4_VALUES, dtype=torch.float32, device=packed_rows.device ++ ) ++ self._lut = lut ++ return _dequant_nvfp4_rows( ++ packed_rows, self.head_dim, scale_2, output_dtype, lut ++ ) ++ ++ ++def _ple_dtype_is_fp8(ple_embedding_dtype: object) -> bool: ++ """Return whether the model config declares FP8 PLE checkpoint weights.""" ++ ++ if ple_embedding_dtype is None: ++ return False ++ if isinstance(ple_embedding_dtype, torch.dtype): ++ return ple_embedding_dtype == torch.float8_e4m3fn ++ return str(ple_embedding_dtype).rsplit(".", 1)[-1] == "float8_e4m3fn" ++ ++ ++def _get_ple_embedding_quant_method( ++ quant_config: QuantizationConfig | None, ++ prefix: str, ++ ple_embedding_dtype: object = None, ++) -> QuantizeMethodBase | None: ++ """Select a packed PLE embedding method for quantized checkpoint shards.""" ++ ++ if isinstance(quant_config, ModelOptMixedPrecisionConfig): ++ if quant_config._resolve_quant_algo(prefix) == "FP8": ++ return Qwen4ExpPLEFp8EmbeddingMethod() ++ return None ++ ++ if isinstance(quant_config, Fp8Config): ++ if not quant_config.is_checkpoint_fp8_serialized: ++ return None ++ ignored_layers = quant_config.ignored_layers ++ if is_layer_skipped( ++ prefix, ++ ignored_layers, ++ quant_config.packed_modules_mapping, ++ match_mode=quant_config.ignored_layers_match_mode, ++ ): ++ return None ++ # PLE checkpoint shards form one runtime embedding parameter. ++ shard_prefix = f"{prefix}.shard_" ++ if any(name.startswith(shard_prefix) for name in ignored_layers): ++ return None ++ return Qwen4ExpPLEFp8EmbeddingMethod() ++ ++ if isinstance(quant_config, ModelOptNvFp4Config): ++ if not quant_config.is_checkpoint_nvfp4_serialized: ++ return None ++ if not quant_config.is_layer_excluded(prefix): ++ logger.info_once("PLE embedding %s uses the runtime NVFP4 method", prefix) ++ return Qwen4ExpPLENVFp4EmbeddingMethod() ++ if _ple_dtype_is_fp8(ple_embedding_dtype): ++ logger.info_once( ++ "Excluded ModelOpt PLE embedding %s uses the runtime FP8 method", ++ prefix, ++ ) ++ return Qwen4ExpPLEFp8EmbeddingMethod() ++ ++ # Checkpoints pairing quantized experts with a global-scale FP8 PLE table ++ # (the RadixArk Flash-Next table grafted onto a CompressedTensors checkpoint) ++ # list the PLE under exclude_modules, so the FP8-ness is visible only in the ++ # tensors' dtype / the config's ple_embedding_dtype, not in the quant config ++ # type dispatched above. Without this branch the offload worker treats the ++ # table as BF16 and materializes ~96 GiB in host RAM (OOM-killed on the ++ # 96 GB host). Honour either the config key (the -e1 ++ # checkpoint copy) or the explicit VLLM_PLE_FP8_GLOBAL_SCALE=1 opt-in the ++ # previous release uses. Placement after the typed dispatch is the ++ # position already proven in service. ++ if _ple_dtype_is_fp8(ple_embedding_dtype) or os.environ.get( ++ "VLLM_PLE_FP8_GLOBAL_SCALE" ++ ) == "1": ++ logger.info_once( ++ "PLE embedding %s: config declares an FP8 table with quant_config %s; " ++ "using the runtime FP8 method (global weight_scale)", ++ prefix, ++ type(quant_config).__name__, ++ ) ++ return Qwen4ExpPLEFp8EmbeddingMethod() ++ ++ return None ++ ++ ++_FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) ++ ++ ++def _dequant_nvfp4_codes( ++ packed: torch.Tensor, ++ scale: torch.Tensor, ++ scale_2: torch.Tensor, ++ lut: torch.Tensor | None = None, ++) -> torch.Tensor: ++ """Unpack NVFP4 codes with FP8 block scales and an FP32 global scale.""" ++ half = packed.shape[-1] ++ codes = torch.stack([packed & 0xF, (packed >> 4) & 0xF], dim=-1).reshape( ++ *packed.shape[:-1], half * 2 ++ ) ++ if lut is None: ++ lut = torch.tensor(_FP4_VALUES, dtype=torch.float32, device=packed.device) ++ magnitude = lut[(codes & 0x7).long()] ++ sign = ((codes >> 3) & 1).to(torch.float32) ++ fp4 = (magnitude * (1 - 2 * sign)).reshape( ++ *packed.shape[:-1], -1, _NVFP4_BLOCK_SIZE ++ ) ++ output = fp4 * scale.float().unsqueeze(-1) * scale_2.float() ++ return output.reshape(*packed.shape[:-1], half * 2) ++ ++ ++def _dequant_nvfp4_rows( ++ packed_rows: torch.Tensor, ++ head_dim: int, ++ scale_2: torch.Tensor, ++ output_dtype: torch.dtype, ++ lut: torch.Tensor, ++) -> torch.Tensor: ++ """Dequantize rows containing NVFP4 codes followed by block scales.""" ++ half = head_dim // 2 ++ codes = packed_rows[..., :half] ++ scales = packed_rows[..., half:].contiguous().view(torch.float8_e4m3fn) ++ return _dequant_nvfp4_codes(codes, scales, scale_2, lut).to(output_dtype) ++ ++ ++def _get_shared_nvfp4_outer_scale( ++ outer_scales: dict[int, torch.Tensor], ++) -> torch.Tensor: ++ """Validate and return the global scale shared by NVFP4 PLE shards.""" ++ first_index, reference = next(iter(outer_scales.items())) ++ reference = reference.reshape(()) ++ for shard_index, outer_scale in outer_scales.items(): ++ if not torch.equal(outer_scale.reshape(()), reference): ++ raise ValueError( ++ "NVFP4 PLE shards must share the same global scale, but " ++ f"shards {first_index} and {shard_index} differ" ++ ) ++ return reference ++ ++ ++class Qwen4ExpNGramEmbedding(PleOffloadLayer): ++ _offload_quant_method: QuantizeMethodBase | None = None ++ _MASK64 = (1 << 64) - 1 ++ _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 ++ _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 ++ _SPLITMIX_M2 = 0x94D049BB133111EB ++ _PLE_LAYER_PRIME = 10007 ++ ++ @classmethod ++ def _splitmix64(cls, value: int) -> int: ++ """Mix an integer into a deterministic unsigned 64-bit value.""" ++ value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64 ++ value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 ++ value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 ++ return (value ^ (value >> 31)) & cls._MASK64 ++ ++ @staticmethod ++ def _is_prime_64(value: int) -> bool: ++ """Return whether a 64-bit integer is prime.""" ++ if value < 2: ++ return False ++ for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): ++ if value % prime == 0: ++ return value == prime ++ exponent = value - 1 ++ shifts = 0 ++ while exponent % 2 == 0: ++ exponent //= 2 ++ shifts += 1 ++ for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): ++ if base % value == 0: ++ continue ++ witness = pow(base, exponent, value) ++ if witness in (1, value - 1): ++ continue ++ for _ in range(shifts - 1): ++ witness = pow(witness, 2, value) ++ if witness == value - 1: ++ break ++ else: ++ return False ++ return True ++ ++ @classmethod ++ def _nth_prime_after(cls, start: int, count: int) -> int: ++ """Return the ``count``-th prime strictly greater than ``start``.""" ++ prime = int(start) ++ for _ in range(count): ++ candidate = prime + 1 ++ if candidate <= 2: ++ prime = 2 ++ continue ++ if candidate % 2 == 0: ++ candidate += 1 ++ while not cls._is_prime_64(candidate): ++ candidate += 2 ++ prime = candidate ++ return prime ++ ++ @classmethod ++ def _make_layer_multipliers( ++ cls, ++ *, ++ ngram_size: int, ++ unigram_vocab_size: int, ++ seed: int, ++ ple_dense_layer_id: int, ++ ) -> list[int]: ++ """Build deterministic hash multipliers for one PLE layer.""" ++ max_multiplier = ((1 << 63) - 1) // unigram_vocab_size ++ half_bound = max(1, max_multiplier // 2) ++ base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id ++ multipliers = [] ++ for index in range(ngram_size): ++ value = base_seed + cls._SPLITMIX_GAMMA * (index + 1) ++ multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1) ++ return multipliers ++ ++ @classmethod ++ def _make_vocab_layout( ++ cls, ++ *, ++ ngram_vocab_size_base: int, ++ ngram_heads: int, ++ ple_dense_layer_id: int, ++ ) -> tuple[list[int], list[int], int]: ++ """Build per-head vocabulary sizes, offsets, and total row count.""" ++ sizes: list[int] = [] ++ offsets: list[int] = [] ++ offset = 0 ++ for local_head in range(ngram_heads): ++ global_head = ple_dense_layer_id * ngram_heads + local_head ++ size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) ++ sizes.append(size) ++ offsets.append(offset) ++ offset += size ++ return sizes, offsets, offset ++ ++ def __init__( ++ self, ++ config: Qwen4ExpTextConfig, ++ embedding_dim: int, ++ ple_dense_layer_id: int, ++ prefix: str, ++ layer_name: str, ++ quant_config: QuantizationConfig | None = None, ++ params_dtype: torch.dtype | None = None, ++ ) -> None: ++ super().__init__() ++ self.layer_name = layer_name ++ self.embedding_dim = embedding_dim ++ self.ngram_size = int(config.ngram_size) ++ self.heads_per_ngram = int(config.heads_per_ngram) ++ self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram ++ if self.ngram_size < 2: ++ raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") ++ if self.heads_per_ngram <= 0: ++ raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") ++ if embedding_dim % self.ngram_heads: ++ raise ValueError( ++ "ple_embed_dim must be divisible by total ngram heads: " ++ f"{embedding_dim} % {self.ngram_heads} != 0" ++ ) ++ self.head_dim = embedding_dim // self.ngram_heads ++ self.eos_token_id = int(config.eos_token_id) ++ self.unigram_vocab_size = int(config.vocab_size) ++ self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) ++ if self.split_ngram_parts <= 0: ++ raise ValueError("split_ngram_parts must be positive") ++ ++ multipliers = self._make_layer_multipliers( ++ ngram_size=self.ngram_size, ++ unigram_vocab_size=self.unigram_vocab_size, ++ seed=int(getattr(config, "seed", 1234)), ++ ple_dense_layer_id=ple_dense_layer_id, ++ ) ++ self.register_buffer( ++ "layer_multipliers", ++ torch.tensor(multipliers, dtype=torch.long), ++ persistent=True, ++ ) ++ ++ sizes, offsets, total_vocab_size = self._make_vocab_layout( ++ ngram_vocab_size_base=int(config.ngram_vocab_size_base), ++ ngram_heads=self.ngram_heads, ++ ple_dense_layer_id=ple_dense_layer_id, ++ ) ++ self.register_buffer( ++ "ngram_heads_vocab_sizes", ++ torch.tensor(sizes, dtype=torch.long), ++ persistent=True, ++ ) ++ self.register_buffer( ++ "ngram_heads_offsets", ++ torch.tensor(offsets, dtype=torch.long), ++ persistent=True, ++ ) ++ divisor = int(config.make_ngram_vocab_size_divisible_by) ++ padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor ++ self.ngram_embedding = PLEVocabParallelEmbedding( ++ padded_vocab_size, ++ self.head_dim, ++ params_dtype=params_dtype, ++ padding_size=divisor, ++ prefix=f"{prefix}.ngram_embedding", ++ quant_method=_get_ple_embedding_quant_method( ++ quant_config, ++ f"{prefix}.ngram_embedding", ++ getattr(config, "ple_embedding_dtype", None), ++ ), ++ ) ++ ++ @staticmethod ++ def _shift_precompute( ++ tokens: torch.Tensor, eos_token_id: int ++ ) -> tuple[torch.Tensor, torch.Tensor]: ++ if tokens.dim() != 2: ++ raise ValueError("tokens must be a 2D tensor") ++ batch_size, seq_len = tokens.shape ++ positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) ++ eos_positions = torch.where(tokens == eos_token_id, positions, -1) ++ previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values ++ previous_eos = torch.cat( ++ [ ++ eos_positions.new_full((batch_size, 1), -1), ++ previous_eos_inclusive[:, :-1], ++ ], ++ dim=1, ++ ) ++ return positions, positions.unsqueeze(0) - previous_eos - 1 ++ ++ @staticmethod ++ def _shift_apply( ++ tokens: torch.Tensor, ++ positions: torch.Tensor, ++ position_in_segment: torch.Tensor, ++ shift: int, ++ eos_token_id: int, ++ ) -> torch.Tensor: ++ if shift == 0: ++ return tokens ++ source = positions - shift ++ gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) ++ shifted = tokens.gather(1, gather_indices) ++ valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) ++ return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) ++ ++ def compute_ngram_ids( ++ self, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ output: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ """Compute n-gram embedding indices for the current request layout.""" ++ input_ids = input_ids.reshape(-1) ++ num_reqs = query_start_loc.numel() - 1 ++ num_tokens = input_ids.shape[0] ++ ++ if input_ids.is_cuda: ++ return ple_ngram_ids( ++ input_ids=input_ids, ++ query_start_loc=query_start_loc, ++ ngram_context=ngram_context, ++ layer_multipliers=self.layer_multipliers, ++ ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes, ++ ngram_heads_offsets=self.ngram_heads_offsets, ++ eos_token_id=self.eos_token_id, ++ heads_per_ngram=self.heads_per_ngram, ++ output=output, ++ ) ++ input_ids = input_ids.long() ++ query_start_loc = query_start_loc.long() ++ if num_reqs <= 0: ++ raise ValueError("PLE requires at least one request") ++ ++ if is_offload_process(): ++ max_seq_len = max( ++ 1, ++ int((query_start_loc[1:] - query_start_loc[:-1]).max().item()), ++ ) ++ # CUDA graphs can pad input_ids beyond the valid request layout. ++ # Keep those rows from overwriting the final real token. ++ num_valid_tokens = min(int(query_start_loc[-1].item()), num_tokens) ++ else: ++ max_seq_len = num_tokens ++ num_valid_tokens = num_tokens ++ ++ positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64) ++ packed = torch.full( ++ (num_reqs, max_seq_len), ++ self.eos_token_id, ++ device=input_ids.device, ++ dtype=torch.int64, ++ ) ++ request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 ++ request_indices.clamp_(max=num_reqs - 1) ++ columns = (positions - query_start_loc[request_indices]).clamp( ++ 0, packed.shape[1] - 1 ++ ) ++ packed[request_indices[:num_valid_tokens], columns[:num_valid_tokens]] = ( ++ input_ids[:num_valid_tokens] ++ ) ++ ngram_context = ngram_context[:num_reqs].to( ++ device=input_ids.device, dtype=torch.long ++ ) ++ ++ context = torch.cat([ngram_context, packed], dim=-1) ++ positions_2d, position_in_segment = self._shift_precompute( ++ context, self.eos_token_id ++ ) ++ shifted = [context] ++ for shift in range(1, self.ngram_size): ++ shifted.append( ++ self._shift_apply( ++ context, ++ positions_2d, ++ position_in_segment, ++ shift, ++ self.eos_token_id, ++ ) ++ ) ++ adjusted_columns = columns + self.ngram_size - 1 ++ id_blocks = [] ++ for ngram in range(2, self.ngram_size + 1): ++ start = (ngram - 2) * self.heads_per_ngram ++ end = start + self.heads_per_ngram ++ mixed = shifted[0] * self.layer_multipliers[0] ++ for index in range(1, ngram): ++ mixed = torch.bitwise_xor( ++ mixed, shifted[index] * self.layer_multipliers[index] ++ ) ++ sizes = self.ngram_heads_vocab_sizes[start:end] ++ offsets = self.ngram_heads_offsets[start:end] ++ ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets ++ id_blocks.append(ids[request_indices, adjusted_columns]) ++ return torch.cat(id_blocks, dim=-1) ++ ++ def forward_impl( # type: ignore[override] ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ output_buffer: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ del hidden_states ++ if input_ids.is_cuda: ++ # Keep request-dependent ID generation outside piecewise graphs. ++ ngram_ids = input_ids.new_empty( ++ (input_ids.numel(), self.ngram_heads), dtype=torch.long ++ ) ++ torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids( ++ input_ids, ++ query_start_loc, ++ ngram_context, ++ ngram_ids, ++ self.layer_name, ++ ) ++ else: ++ ngram_ids = self.compute_ngram_ids( ++ input_ids, ++ query_start_loc, ++ ngram_context, ++ ) ++ num_tokens = input_ids.reshape(-1).shape[0] ++ if output_buffer is not None: ++ quant_method = getattr(self.ngram_embedding, "quant_method", None) ++ if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): ++ output_dim = self.get_offload_output_dim(self.embedding_dim) ++ output = output_buffer[:num_tokens, :output_dim] ++ output.copy_(self.ngram_embedding(ngram_ids).flatten(-2)) ++ return output ++ output = output_buffer[:num_tokens, : self.embedding_dim] ++ torch.index_select( ++ self.ngram_embedding.weight, ++ 0, ++ ngram_ids.reshape(-1), ++ out=output.reshape(-1, self.head_dim), ++ ) ++ return output ++ return self.ngram_embedding(ngram_ids).flatten(-2) ++ ++ def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: ++ """Keep quantized lookup results in their embedding storage dtype.""" ++ embedding = getattr(self, "ngram_embedding", None) ++ weight = getattr(embedding, "weight", None) ++ if weight is not None: ++ return weight.dtype ++ if isinstance(self._offload_quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): ++ return torch.uint8 ++ if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod): ++ return torch.float8_e4m3fn ++ return default_dtype ++ ++ def get_offload_output_dim(self, default_dim: int) -> int: ++ """Keep NVFP4 lookup rows packed while transferring them to the GPU.""" ++ quant_method = getattr( ++ getattr(self, "ngram_embedding", None), "quant_method", None ++ ) ++ if quant_method is None: ++ quant_method = self._offload_quant_method ++ if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): ++ if default_dim % _NVFP4_BLOCK_SIZE: ++ raise ValueError( ++ "NVFP4 PLE output dim must be a multiple of " ++ f"{_NVFP4_BLOCK_SIZE}, got {default_dim}" ++ ) ++ return default_dim // 2 + default_dim // _NVFP4_BLOCK_SIZE ++ return default_dim ++ ++ def initialize_dummy_offload_metadata(self, device: torch.device) -> None: ++ """Initialize quantization metadata skipped by the dummy loader.""" ++ quant_method = self._offload_quant_method ++ if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): ++ self.register_buffer( ++ "_offload_weight_scale", ++ torch.ones((), dtype=torch.float32, device=device), ++ persistent=False, ++ ) ++ elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): ++ self.register_buffer( ++ "_offload_weight_scale_2", ++ torch.ones((), dtype=torch.float32, device=device), ++ persistent=False, ++ ) ++ self.register_buffer( ++ "_offload_nvfp4_lut", ++ torch.tensor(_FP4_VALUES, dtype=torch.float32, device=device), ++ persistent=False, ++ ) ++ ++ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ++ """Load hash buffers and checkpoint-split embedding rows.""" ++ ++ # Prefix grouping may invoke this method repeatedly for one module. ++ # Treat the global FP8 scale as incrementally loaded state. ++ # GPU workers retain only dequantization metadata. The CPU process owns ++ # the embedding table and transfers quantized lookup rows unchanged. ++ if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): ++ retained: set[str] = set() ++ quant_method = self._offload_quant_method ++ if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): ++ for name, loaded_weight in weights: ++ if name != "ngram_embedding.weight_scale": ++ continue ++ self.register_buffer( ++ "_offload_weight_scale", ++ loaded_weight.to( ++ device=torch.accelerator.current_accelerator() ++ ), ++ persistent=False, ++ ) ++ retained.add(name) ++ elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): ++ outer_scales: dict[int, torch.Tensor] = {} ++ for name, loaded_weight in weights: ++ if not name.startswith( ++ "ngram_embedding.shard_" ++ ) or not name.endswith(".weight_scale_2"): ++ continue ++ shard_text = name[ ++ len("ngram_embedding.shard_") : -len(".weight_scale_2") ++ ] ++ shard_index = int(shard_text) ++ outer_scales[shard_index] = loaded_weight ++ retained.add(name) ++ if not retained: ++ raise ValueError( ++ "NVFP4 PLE offload checkpoint is missing its global scale" ++ ) ++ scale_2 = _get_shared_nvfp4_outer_scale(outer_scales).to( ++ device=torch.accelerator.current_accelerator() ++ ) ++ self.register_buffer( ++ "_offload_weight_scale_2", scale_2, persistent=False ++ ) ++ self.register_buffer( ++ "_offload_nvfp4_lut", ++ torch.tensor( ++ _FP4_VALUES, dtype=torch.float32, device=scale_2.device ++ ), ++ persistent=False, ++ ) ++ else: ++ for _ in weights: ++ pass ++ return retained ++ ++ persistent_buffers = { ++ "layer_multipliers": self.layer_multipliers, ++ "ngram_heads_offsets": self.ngram_heads_offsets, ++ "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, ++ } ++ loaded: set[str] = set() ++ regular_weights: list[tuple[str, torch.Tensor]] = [] ++ shard_prefix = "ngram_embedding.shard_" ++ quant_method = getattr(self.ngram_embedding, "quant_method", None) ++ nvfp4_runtime = isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod) ++ packed_codes: dict[int, torch.Tensor] = {} ++ packed_scales: dict[int, torch.Tensor] = {} ++ packed_outer_scales: dict[int, torch.Tensor] = {} ++ ++ for name, loaded_weight in weights: ++ leaf_name = name.rsplit(".", 1)[-1] ++ if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": ++ continue ++ if name in persistent_buffers: ++ buffer = persistent_buffers[name] ++ if buffer.shape != loaded_weight.shape: ++ raise ValueError( ++ f"Shape mismatch for {name}: expected " ++ f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" ++ ) ++ buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) ++ loaded.add(name) ++ continue ++ ++ if nvfp4_runtime and name.startswith(shard_prefix): ++ suffix = name[len(shard_prefix) :] ++ for ending, sink in ( ++ (".weight_scale_2", packed_outer_scales), ++ (".weight_scale", packed_scales), ++ (".weight", packed_codes), ++ ): ++ if not suffix.endswith(ending): ++ continue ++ shard_text = suffix[: -len(ending)] ++ shard_index = int(shard_text) ++ if shard_index >= self.split_ngram_parts: ++ raise ValueError( ++ f"PLE embedding shard index {shard_index} exceeds " ++ f"split_ngram_parts={self.split_ngram_parts}" ++ ) ++ sink[shard_index] = loaded_weight ++ break ++ else: ++ regular_weights.append((name, loaded_weight)) ++ continue ++ ++ if ( ++ not nvfp4_runtime ++ and name.startswith(shard_prefix) ++ and name.endswith(".weight") ++ ): ++ shard_text = name[len(shard_prefix) : -len(".weight")] ++ if not shard_text.isdigit(): ++ regular_weights.append((name, loaded_weight)) ++ continue ++ shard_index = int(shard_text) ++ if shard_index >= self.split_ngram_parts: ++ raise ValueError( ++ f"PLE embedding shard index {shard_index} exceeds " ++ f"split_ngram_parts={self.split_ngram_parts}" ++ ) ++ embedding = self.ngram_embedding ++ shard_size = ( ++ embedding.org_vocab_size + self.split_ngram_parts - 1 ++ ) // self.split_ngram_parts ++ checkpoint_start = shard_index * shard_size ++ expected_rows = max( ++ 0, ++ min(shard_size, embedding.org_vocab_size - checkpoint_start), ++ ) ++ expected_shape = (expected_rows, embedding.embedding_dim) ++ if tuple(loaded_weight.shape) != expected_shape: ++ raise ValueError( ++ f"Shape mismatch for PLE embedding shard {shard_index}: " ++ f"expected {expected_shape}, got " ++ f"{tuple(loaded_weight.shape)}" ++ ) ++ embedding.weight.weight_loader( ++ embedding.weight, ++ loaded_weight, ++ checkpoint_start=checkpoint_start, ++ ) ++ loaded.add("ngram_embedding.weight") ++ continue ++ regular_weights.append((name, loaded_weight)) ++ ++ if nvfp4_runtime: ++ if not packed_codes: ++ raise ValueError("NVFP4 PLE checkpoint contains no packed shards") ++ for shard_index in packed_codes: ++ if ( ++ shard_index not in packed_scales ++ or shard_index not in packed_outer_scales ++ ): ++ raise ValueError( ++ f"NVFP4 PLE shard {shard_index} is missing its scale tensors" ++ ) ++ shared_outer_scale = _get_shared_nvfp4_outer_scale( ++ { ++ shard_index: packed_outer_scales[shard_index] ++ for shard_index in packed_codes ++ } ++ ) ++ ++ for shard_index, codes in packed_codes.items(): ++ embedding = self.ngram_embedding ++ shard_size = ( ++ embedding.org_vocab_size + self.split_ngram_parts - 1 ++ ) // self.split_ngram_parts ++ checkpoint_start = shard_index * shard_size ++ expected_rows = max( ++ 0, min(shard_size, embedding.org_vocab_size - checkpoint_start) ++ ) ++ expected_codes_shape = (expected_rows, embedding.embedding_dim // 2) ++ if tuple(codes.shape) != expected_codes_shape: ++ raise ValueError( ++ f"Shape mismatch for packed PLE shard {shard_index}: " ++ f"expected {expected_codes_shape}, got {tuple(codes.shape)}" ++ ) ++ tp_bounds = { ++ "checkpoint_start": checkpoint_start, ++ "tp_start": embedding.shard_indices.org_vocab_start_index, ++ "tp_end": embedding.shard_indices.org_vocab_end_index, ++ } ++ scales = packed_scales[shard_index] ++ expected_scales_shape = ( ++ expected_rows, ++ embedding.embedding_dim // _NVFP4_BLOCK_SIZE, ++ ) ++ if tuple(scales.shape) != expected_scales_shape: ++ raise ValueError( ++ f"Scale shape mismatch for packed PLE shard {shard_index}: " ++ f"expected {expected_scales_shape}, got {tuple(scales.shape)}" ++ ) ++ copy_ple_embedding_shard_(embedding.weight.data, codes, **tp_bounds) ++ copy_ple_embedding_shard_(embedding.weight_scale.data, scales, **tp_bounds) ++ ++ if nvfp4_runtime: ++ self.ngram_embedding.weight_scale_2.data.copy_(shared_outer_scale) ++ loaded.update( ++ { ++ "ngram_embedding.weight", ++ "ngram_embedding.weight_scale", ++ "ngram_embedding.weight_scale_2", ++ } ++ ) ++ ++ if regular_weights: ++ loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) ++ return loaded ++ ++ + class Qwen4ExpPLELayer(nn.Module, MambaBase): + def __init__( + self, +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 1c2808a..1b7d55f 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -4,6 +4,14 @@ + + from __future__ import annotations + ++import atexit ++import json ++import math ++import os ++import threading ++import time ++from collections.abc import Collection, Mapping ++from pathlib import Path + from typing import ClassVar, cast + + import torch +@@ -21,6 +29,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.model_executor.layers.layernorm import GemmaRMSNorm + from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear + from vllm.model_executor.layers.quantization import QuantizationConfig ++from vllm.logger import init_logger + from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope + from vllm.model_executor.models.qwen3_next import Qwen3NextAttention + from vllm.platforms import current_platform +@@ -54,6 +63,235 @@ from . import model + from .indexer_qsa import QSAIndexer + + ++logger = init_logger(__name__) ++ ++# --- FP8 K/V scale calibration collector ----------------------- ++# Active ONLY when VLLM_QSA_KV_COLLECT names a writable directory. Records a ++# GPU-side running absmax of the pre-quantization BF16 K/V per layer and dumps ++# per-rank JSON (periodically + at exit) for an offline max-merge into the ++# VLLM_QSA_KV_SCALES sidecar. Calibration must run with --enforce-eager: the ++# collector mutates module-external state per step, which cudagraph capture ++# would freeze and torch.compile may not retrace. With the env unset this is a ++# constant-false branch and prod behaviour is unchanged. ++_QSA_KV_COLLECT_ENV = "VLLM_QSA_KV_COLLECT" ++_qsa_collect_dir = os.environ.get(_QSA_KV_COLLECT_ENV, "").strip() ++_qsa_collect_state: dict[str, list[torch.Tensor]] = {} ++_qsa_collect_calls = 0 ++ ++ ++_qsa_collect_ident: str | None = None ++ ++ ++def _qsa_collect_identity() -> str: ++ """Stable per-process dump name: (pp, tp) rank, resolved once and cached so ++ a teardown-time lookup failure cannot move the final flush to a PID-named ++ file. TP-only runs keep the historical name.""" ++ global _qsa_collect_ident ++ if _qsa_collect_ident is not None: ++ return _qsa_collect_ident ++ try: ++ from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank ++ ++ rank = get_tensor_model_parallel_rank() ++ pp_group = get_pp_group() ++ pp_tag = f"pp{pp_group.rank_in_group}_" if pp_group.world_size > 1 else "" ++ _qsa_collect_ident = f"{pp_tag}rank{rank}" ++ except Exception: ++ return f"pid{os.getpid()}" # not cached: retry once the groups exist ++ return _qsa_collect_ident ++ ++ ++def _qsa_collect_dump() -> None: ++ if not _qsa_collect_state: ++ return ++ out = { ++ name: {"k_absmax": float(st[0].item()), "v_absmax": float(st[1].item())} ++ for name, st in _qsa_collect_state.items() ++ } ++ path = Path(_qsa_collect_dir) / f"qsa_absmax_{_qsa_collect_identity()}.json" ++ tmp = path.with_suffix(".tmp") ++ tmp.write_text(json.dumps(out, indent=1), encoding="utf-8") ++ tmp.replace(path) ++ ++ ++def _qsa_collect_absmax( ++ layer_name: str, key: torch.Tensor, value: torch.Tensor ++) -> None: ++ global _qsa_collect_calls ++ k_abs = key.detach().abs().amax().float() ++ v_abs = value.detach().abs().amax().float() ++ st = _qsa_collect_state.get(layer_name) ++ if st is None: ++ if not _qsa_collect_state: ++ atexit.register(_qsa_collect_dump) ++ _qsa_collect_state[layer_name] = [k_abs, v_abs] ++ else: ++ torch.maximum(st[0], k_abs, out=st[0]) ++ torch.maximum(st[1], v_abs, out=st[1]) ++ _qsa_collect_calls += 1 ++ if _qsa_collect_calls % 2000 == 0: ++ _qsa_collect_dump() ++# ------------------------------------------------------------------------------ ++ ++# --- FP8 K/V runtime clipping counter -------------------------- ++# Under-coverage guard for the calibrated scales sidecar: counts K/V elements ++# that would saturate E4M3 after scaling (|x| > 448 * scale), i.e. inputs that ++# exceed the calibration-time absmax. Active ONLY when VLLM_QSA_KV_CLIP_COUNT ++# is set to a positive read interval in seconds. ++# ++# Graph-safety: unlike the calibration collector above, this MUST work under ++# FULL cudagraph replay, where the Python forward does not run per step. The ++# counters are persistent per-layer GPU int64 tensors allocated on the first ++# (eager, pre-capture) forward; the increment is pure captured GPU ops, so ++# graph replay keeps counting. A daemon reader thread .item()s the counters on ++# a cadence and logs: WARNING on any new clips, INFO heartbeat every 10th pass ++# so a quiet log still proves the instrument is alive. ++_QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT" ++_qsa_clip_raw = os.environ.get(_QSA_CLIP_ENV, "").strip() ++try: ++ _qsa_clip_interval = float(_qsa_clip_raw) if _qsa_clip_raw else 0.0 ++except ValueError as exc: # fail closed: a set-but-broken gate must not be silent ++ raise ValueError( ++ f"{_QSA_CLIP_ENV} must be a read interval in seconds, got {_qsa_clip_raw!r}" ++ ) from exc ++if _qsa_clip_raw and not (math.isfinite(_qsa_clip_interval) and _qsa_clip_interval > 0.0): ++ raise ValueError( ++ f"{_QSA_CLIP_ENV} must be a finite positive number of seconds, got {_qsa_clip_raw!r}" ++ ) ++_qsa_clip_counters: dict[str, torch.Tensor] = {} ++# v2 reader fix: the reader must never block on the compute ++# stream; see guard_telemetry.NonBlockingReader. ++from vllm.v1.worker.guard_telemetry import NonBlockingReader as _NonBlockingReader # noqa: E402 ++_qsa_clip_nb = _NonBlockingReader() ++_qsa_clip_reader_started = False ++_qsa_clip_lock = threading.Lock() ++ ++ ++def _qsa_clip_reader() -> None: ++ last: dict[str, tuple[int, int]] = {} ++ passes = 0 ++ while True: ++ time.sleep(_qsa_clip_interval) ++ passes += 1 ++ try: ++ totals_k = totals_v = 0 ++ for name, ctr in list(_qsa_clip_counters.items()): ++ if not _qsa_clip_nb.registered(name): ++ _qsa_clip_nb.register(name, ctr) # reader thread: never capturing ++ ready = _qsa_clip_nb.poll() ++ for name, ctr in list(_qsa_clip_counters.items()): ++ row = ready.get(name) ++ if row is None: ++ # copy still in flight: report on the next pass ++ prev_k, prev_v = last.get(name, (0, 0)) ++ totals_k += prev_k ++ totals_v += prev_v ++ continue ++ k_total = int(row[0]) ++ v_total = int(row[1]) ++ totals_k += k_total ++ totals_v += v_total ++ prev_k, prev_v = last.get(name, (0, 0)) ++ if k_total > prev_k or v_total > prev_v: ++ logger.warning( ++ "QSA KV CLIPPING on %s: k +%d (total %d), v +%d (total %d)" ++ " -- inputs exceeded the calibration absmax; the scales" ++ " sidecar may need re-calibration", ++ name, ++ k_total - prev_k, ++ k_total, ++ v_total - prev_v, ++ v_total, ++ ) ++ last[name] = (k_total, v_total) ++ if passes % 10 == 1: ++ logger.info( ++ "QSA clip counters alive: %d layer(s), totals k=%d v=%d", ++ len(_qsa_clip_counters), ++ totals_k, ++ totals_v, ++ ) ++ except Exception: ++ logger.exception("QSA clip reader pass failed; thread continues") ++ ++ ++def _qsa_clip_count( ++ layer: "Qwen4ExpQSAAttention", ++ key: torch.Tensor, ++ value: torch.Tensor, ++ slot_mapping: torch.Tensor | None = None, ++) -> None: ++ global _qsa_clip_reader_started ++ ctr = _qsa_clip_counters.get(layer.layer_name) ++ if ctr is None: ++ # First forward is eager (profiling/warmup precede capture), so the ++ # allocation never happens inside cudagraph capture. ++ ctr = torch.zeros(2, dtype=torch.int64, device=key.device) ++ _qsa_clip_counters[layer.layer_name] = ctr # registered with the reader by the reader thread (never during capture) ++ with _qsa_clip_lock: ++ if not _qsa_clip_reader_started: ++ _qsa_clip_reader_started = True ++ threading.Thread( ++ target=_qsa_clip_reader, ++ name="qsa-clip-reader", ++ daemon=True, ++ ).start() ++ logger.info( ++ "QSA clip counter armed (read interval %.0fs)", ++ _qsa_clip_interval, ++ ) ++ # Thresholds are 1-D so the comparison promotes to float32; a 0-dim float32 ++ # tensor promotes like a Python scalar and the threshold would be rounded to ++ # the BF16 grid of K/V, undercounting values just above the calibrated ceiling. ++ k_ceiling = (448.0 * layer._k_scale).to(torch.float32).view(1) ++ v_ceiling = (448.0 * layer._v_scale).to(torch.float32).view(1) ++ k_over = key.detach().abs() > k_ceiling ++ v_over = value.detach().abs() > v_ceiling ++ if slot_mapping is not None: ++ # The CUDA writer (reshape_and_cache_flash) skips rows with a negative ++ # slot: padding rows of a captured batch size and, under the V2 model ++ # runner, stale rows of the persistent input buffers. Count only what ++ # is actually written, with fixed-shape masking so the op is capturable. ++ n = min(slot_mapping.shape[0], key.shape[0]) ++ written = (slot_mapping[:n] >= 0).view(-1, 1, 1) ++ k_over = k_over[:n] & written ++ v_over = v_over[:n] & written ++ ctr[0] += k_over.sum() ++ ctr[1] += v_over.sum() ++# ------------------------------------------------------------------------------ ++ ++_QSA_FP8_CACHE_DTYPES = ("fp8", "fp8_e4m3") ++_QSA_SUPPORTED_CACHE_DTYPES = ("auto", "bfloat16", *_QSA_FP8_CACHE_DTYPES) ++ ++ ++def _is_qsa_fp8_cache_dtype(cache_dtype: str) -> bool: ++ return cache_dtype in _QSA_FP8_CACHE_DTYPES ++ ++ ++def _validated_qsa_fp8_dtype(cache_dtype: str) -> torch.dtype | None: ++ """Return the platform E4M3 dtype after enforcing the validated HW gate.""" ++ ++ if not _is_qsa_fp8_cache_dtype(cache_dtype): ++ return None ++ if not current_platform.is_cuda(): ++ raise ValueError( ++ "Qwen4Exp QSA E4M3 KV cache is supported only on CUDA; " ++ "ROCm and other platforms have not been validated" ++ ) ++ ++ capability = current_platform.get_device_capability() ++ if capability is None or capability.to_int() != 86: ++ cap_str = ( ++ capability.as_version_str() if capability is not None else "unknown" ++ ) ++ raise ValueError( ++ "Qwen4Exp QSA E4M3 KV cache is validated only on SM86, but " ++ f"{current_platform.get_device_name()} has compute capability " ++ f"{cap_str}. Re-run with --kv-cache-dtype bfloat16." ++ ) ++ return current_platform.fp8_dtype() ++ ++ + class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder): + """Flash metadata supporting uniform decode and target-verify graphs.""" + +@@ -64,7 +302,12 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): + """FullAttentionSpec backend used by the merged QSA owner.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] +- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] ++ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ ++ "auto", ++ "bfloat16", ++ "fp8", ++ "fp8_e4m3", ++ ] + + @staticmethod + def get_name() -> str: +@@ -99,15 +342,36 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + supports_pcp: bool = False + + def __init__(self, *args, **kwargs) -> None: ++ # QSA never dispatches to a FlashAttention kernel: attention runs through ++ # forward_qsa -> qsa_sparse_paged_attention (our Triton sparse-GQA kernel) ++ # and the cache write goes through the CUDA reshape_and_cache_flash op. ++ # FlashAttentionImpl.__init__ nevertheless rejects any quantized KV dtype ++ # that FA itself cannot read, and sm_86 has no FP8 FA kernel. Present ++ # "auto" to the base class, then restore the real dtype. The only other ++ # kv_cache_dtype use in that constructor is an SM90/FA4 dequant path that ++ # cannot trigger here, and supports_quant_query_input is set below anyway. ++ requested_kv_cache_dtype = None ++ if "kv_cache_dtype" in kwargs: ++ requested_kv_cache_dtype = kwargs["kv_cache_dtype"] ++ kwargs = {**kwargs, "kv_cache_dtype": "auto"} ++ elif len(args) > 6: ++ requested_kv_cache_dtype = args[6] ++ args = (*args[:6], "auto", *args[7:]) + super().__init__(*args, **kwargs) ++ if requested_kv_cache_dtype is not None: ++ self.kv_cache_dtype = requested_kv_cache_dtype + if not is_flash_attn_varlen_func_available(): + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") + if self.dcp_world_size != 1: + raise NotImplementedError( + "Qwen4Exp QSA does not support decode context parallelism" + ) +- if self.kv_cache_dtype not in ("auto", "bfloat16"): +- raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache") ++ if self.kv_cache_dtype not in _QSA_SUPPORTED_CACHE_DTYPES: ++ raise NotImplementedError( ++ "Qwen4Exp QSA supports BF16 or static per-tensor E4M3 main " ++ f"KV caches, not {self.kv_cache_dtype!r}" ++ ) ++ self.qsa_fp8_dtype = _validated_qsa_fp8_dtype(self.kv_cache_dtype) + self.supports_quant_query_input = False + + def forward_qsa( +@@ -144,11 +408,42 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + logical_indices = topk_buffer[:num_tokens] + token_to_req = token_to_req[:num_tokens] + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) +- if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16: +- raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V") ++ if query.dtype != torch.bfloat16: ++ raise NotImplementedError("Qwen4Exp QSA requires a BF16 query") ++ if key_cache.dtype != value_cache.dtype: ++ raise TypeError("Qwen4Exp QSA K and V cache dtypes must match") ++ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype): ++ if self.qsa_fp8_dtype is None: ++ raise RuntimeError("QSA FP8 dtype was not initialized") ++ if key_cache.dtype not in (torch.uint8, self.qsa_fp8_dtype): ++ raise NotImplementedError( ++ "Qwen4Exp QSA E4M3 cache storage must use raw uint8 or " ++ f"the platform FP8 dtype, not {key_cache.dtype}" ++ ) ++ elif key_cache.dtype != torch.bfloat16: ++ raise NotImplementedError( ++ "Qwen4Exp QSA BF16 mode requires a BF16 main KV cache" ++ ) + + from .ops.qsa import qsa_sparse_paged_attention + ++ # Announce the effective scales once per layer. load_qsa_static_kv_scales() ++ # is an explicit integration call, so a deployment that forgets it runs at ++ # the vLLM default of 1.0 -- which is a legitimate configuration but must ++ # never be a silent one. This makes the actual values greppable per rank. ++ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype) and not getattr( ++ self, "_qsa_scales_logged", False ++ ): ++ self._qsa_scales_logged = True ++ logger.info( ++ "Qwen4Exp QSA E4M3 KV active on %s: k_scale=%.6g v_scale=%.6g%s", ++ getattr(layer, "layer_name", ""), ++ float(layer._k_scale), ++ float(layer._v_scale), ++ "" if float(layer._k_scale) != 1.0 or float(layer._v_scale) != 1.0 ++ else " (DEFAULT 1.0 -- no static calibration loaded)", ++ ) ++ + qsa_sparse_paged_attention( + query[:num_tokens], + key_cache, +@@ -157,7 +452,9 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + attn_metadata.block_table, + token_to_req, + use_prefill_config, +- output[:num_tokens], ++ out=output[:num_tokens], ++ k_scale=layer._k_scale, ++ v_scale=layer._v_scale, + output_gate=output_gate[:num_tokens], + ) + return output +@@ -185,10 +482,20 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + raise ValueError("Qwen4Exp QSA requires a paged KV cache") + if model_config.dtype != torch.bfloat16: + raise NotImplementedError("Qwen4Exp QSA currently requires BF16") +- if cache_config.cache_dtype not in ("auto", "bfloat16"): +- raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache") +- if getattr(quant_config, "kv_cache_scheme", None) is not None: +- raise NotImplementedError("Qwen4Exp QSA does not support KV quantization") ++ if cache_config.cache_dtype not in _QSA_SUPPORTED_CACHE_DTYPES: ++ raise NotImplementedError( ++ "Qwen4Exp QSA supports BF16 or static per-tensor E4M3 main " ++ f"KV caches, not {cache_config.cache_dtype!r}" ++ ) ++ qsa_fp8_dtype = _validated_qsa_fp8_dtype(cache_config.cache_dtype) ++ kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None) ++ if kv_cache_scheme is not None: ++ raise NotImplementedError( ++ "Qwen4Exp QSA does not implement kv_cache_scheme variants. " ++ "Dynamic, per-head, per-token, and other non-scalar KV scales " ++ "are unsupported; use static scalar per-layer K/V scales from " ++ "load_qsa_static_kv_scales()." ++ ) + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 +@@ -285,10 +592,23 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, model_config + ) +- if self.kv_cache_torch_dtype != torch.bfloat16: +- raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage") ++ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype): ++ if qsa_fp8_dtype is None: ++ raise RuntimeError("QSA FP8 dtype was not initialized") ++ if self.kv_cache_torch_dtype not in (torch.uint8, qsa_fp8_dtype): ++ raise NotImplementedError( ++ "Qwen4Exp QSA E4M3 cache storage resolved to unsupported " ++ f"dtype {self.kv_cache_torch_dtype}" ++ ) ++ elif self.kv_cache_torch_dtype != torch.bfloat16: ++ raise NotImplementedError( ++ "Qwen4Exp QSA BF16 mode requires BF16 cache storage" ++ ) + self.kv_sharing_target_layer_name = None + self.kv_cache = torch.tensor([]) ++ # The checkpoint has no KV scale scheme. These registered scalar ++ # buffers therefore remain 1.0 unless explicitly replaced after ++ # weight loading by load_qsa_static_kv_scales(). + set_default_quant_scales(self, register_buffer=True) + + self.attn_backend = Qwen4ExpQSAFlashAttentionBackend +@@ -381,6 +701,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + ) + if selected.shape != (num_tokens, self.indexer.packed_output_width): + raise RuntimeError("QSA indexer returned an invalid selection shape") ++ if _qsa_clip_interval > 0.0 and _is_qsa_fp8_cache_dtype( ++ self.kv_cache_dtype ++ ): ++ _qsa_clip_count(self, key, value, main_metadata.slot_mapping) + impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) + impl.do_kv_cache_update( + self, +@@ -414,6 +738,8 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + query = q.view(num_tokens, self.num_heads, self.head_dim) + key = k.view(num_tokens, self.num_kv_heads, self.head_dim) + value = v.view(num_tokens, self.num_kv_heads, self.head_dim) ++ if _qsa_collect_dir: ++ _qsa_collect_absmax(self.layer_name, key, value) + attn_output = torch.empty_like(query) + # Keep the index projection outside the eager break. + projected_qk, _ = self.indexer.index_qk_proj(hidden_states) +@@ -431,9 +757,146 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + return output + + ++def load_qsa_static_kv_scales( ++ layers: nn.Module | Mapping[str, nn.Module], ++ sidecar_path: str | Path, ++ *, ++ strict: bool = True, ++ expected_layer_names: Collection[str] | None = None, ++) -> list[str]: ++ """Load static scalar QSA K/V scales into already-loaded model layers. ++ ++ Call this on every model worker after checkpoint weight loading finishes ++ and before any profiling, warmup, or cudagraph capture. ``layers`` may be ++ the loaded root module or ``compilation_config.static_forward_context``. ++ ++ The JSON schema is:: ++ ++ { ++ "model.layers.3.self_attn.attn": { ++ "k_scale": 0.25, ++ "v_scale": 2.0 ++ } ++ } ++ ++ ``strict=False`` permits a deliberate subset; omitted layers keep the ++ vLLM default scale of 1.0. Unknown names and non-scalar values are always ++ errors. ++ ++ Pipeline parallelism: a rank holds only its stage's layers, so the sidecar ++ legitimately names layers this rank does not own. Pass ++ ``expected_layer_names`` = the QSA layer names of the WHOLE model; names are ++ validated against that set, ``strict`` requires the sidecar to cover all ++ of it (checkable on every rank), and the scales are applied to the local ++ intersection only. Without it (TP-only callers) the local layers are the ++ universe, as before. ++ """ ++ ++ path = Path(sidecar_path) ++ raw = json.loads(path.read_text(encoding="utf-8")) ++ if not isinstance(raw, dict): ++ raise ValueError(f"{path} must contain a JSON object keyed by layer name") ++ ++ if isinstance(layers, nn.Module): ++ qsa_layers = { ++ layer.layer_name: layer ++ for layer in layers.modules() ++ if isinstance(layer, Qwen4ExpQSAAttention) ++ } ++ else: ++ qsa_layers = { ++ name: layer ++ for name, layer in layers.items() ++ if isinstance(name, str) ++ and isinstance(layer, Qwen4ExpQSAAttention) ++ } ++ ++ if not qsa_layers and expected_layer_names is None: ++ raise ValueError("No Qwen4Exp QSA attention layers were supplied") ++ non_fp8_layers = sorted( ++ name ++ for name, layer in qsa_layers.items() ++ if not _is_qsa_fp8_cache_dtype(layer.kv_cache_dtype) ++ ) ++ if non_fp8_layers: ++ raise ValueError( ++ "Static QSA KV scales may only be injected into E4M3 layers; " ++ f"non-FP8 layers: {non_fp8_layers}" ++ ) ++ ++ if expected_layer_names is None: ++ universe = set(qsa_layers) ++ else: ++ universe = set(expected_layer_names) ++ not_expected = sorted(name for name in qsa_layers if name not in universe) ++ if not_expected: ++ raise ValueError( ++ "Local QSA layers outside expected_layer_names: " ++ f"{not_expected}" ++ ) ++ unknown = sorted(name for name in raw if name not in universe) ++ if unknown: ++ raise ValueError(f"Unknown QSA layer names in {path}: {unknown}") ++ missing = sorted(name for name in universe if name not in raw) ++ if strict and missing: ++ raise ValueError(f"Missing QSA layer scales in {path}: {missing}") ++ remote = sorted(name for name in raw if name not in qsa_layers) ++ if remote: ++ logger.info( ++ "Qwen4Exp QSA: %d sidecar layer(s) belong to other pipeline stages " ++ "and are skipped on this rank: %s%s", ++ len(remote), ++ ", ".join(remote[:4]), ++ " ..." if len(remote) > 4 else "", ++ ) ++ ++ validated: dict[str, tuple[float, float]] = {} ++ for layer_name, entry in raw.items(): ++ if not isinstance(entry, dict): ++ raise ValueError(f"{layer_name} scales must be a JSON object") ++ if set(entry) != {"k_scale", "v_scale"}: ++ raise ValueError( ++ f"{layer_name} must contain exactly k_scale and v_scale" ++ ) ++ ++ values: list[float] = [] ++ for field in ("k_scale", "v_scale"): ++ value = entry[field] ++ if type(value) not in (int, float): ++ raise ValueError( ++ f"{layer_name}.{field} must be a scalar JSON number" ++ ) ++ scale = float(value) ++ if not math.isfinite(scale) or scale <= 0.0: ++ raise ValueError( ++ f"{layer_name}.{field} must be finite and positive" ++ ) ++ values.append(scale) ++ ++ layer = qsa_layers.get(layer_name) ++ if layer is None: ++ continue # another pipeline stage's layer: schema-checked only ++ for field in ("_k_scale", "_v_scale"): ++ tensor = getattr(layer, field, None) ++ if not isinstance(tensor, torch.Tensor) or tensor.numel() != 1: ++ raise ValueError( ++ f"{layer_name}.{field} must be an existing scalar tensor" ++ ) ++ validated[layer_name] = (values[0], values[1]) ++ ++ with torch.no_grad(): ++ for layer_name, (k_scale, v_scale) in validated.items(): ++ layer = qsa_layers[layer_name] ++ layer._k_scale.fill_(k_scale) ++ layer._v_scale.fill_(v_scale) ++ ++ return sorted(validated) ++ ++ + __all__ = [ + "QSAIndexer", + "Qwen4ExpQSAAttention", + "Qwen4ExpQSAFlashAttentionBackend", + "Qwen4ExpQSAFlashAttentionImpl", ++ "load_qsa_static_kv_scales", + ] +-- +2.43.0 + diff --git a/patches/0003-qwen4_exp-pipeline-stage-ownership-for-PLE-the-MTP-d.patch b/patches/0003-qwen4_exp-pipeline-stage-ownership-for-PLE-the-MTP-d.patch new file mode 100644 index 0000000..9102195 --- /dev/null +++ b/patches/0003-qwen4_exp-pipeline-stage-ownership-for-PLE-the-MTP-d.patch @@ -0,0 +1,294 @@ +From f36b743fb015941866baa173d0b0efc66f694d98 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:11:27 -0400 +Subject: [PATCH 003/103] qwen4_exp: pipeline-stage ownership for PLE, the MTP + drafter and the vision tower + +Each pipeline stage builds only the parts it owns: PLE when every PLE layer +lives on stage 0, MTP under PP (port of vLLM PR #46994 minus its ROCm MLA file +and tests) with the drafter on the last stage, and the vision tower on the +first stage only. +--- + vllm/model_executor/models/config.py | 21 +++- + vllm/models/qwen4_exp/nvidia/model.py | 41 +++++-- + vllm/models/qwen4_exp/nvidia/model_state.py | 8 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 6 +- + vllm/v1/worker/gpu_worker.py | 112 ++++++++++++++++++++ + 5 files changed, 170 insertions(+), 18 deletions(-) + +diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py +index bbbe98b..5c55412 100644 +--- a/vllm/model_executor/models/config.py ++++ b/vllm/model_executor/models/config.py +@@ -876,12 +876,25 @@ class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConf + ) + # Checked again in Qwen4ExpModelState; rejecting it here keeps the + # engine from loading weights first. ++ # vllm#54709: PLE only needs raw input_ids on the rank that owns a PLE ++ # layer. Allow PP>1 as long as every PLE layer lives on pipeline rank 0. + if text_config.ple_layer_ids and parallel_config.pipeline_parallel_size > 1: +- raise NotImplementedError( +- "Qwen4Exp N-gram PLE embedding requires pipeline_parallel_size=1 " +- "because non-first pipeline ranks do not receive the raw input_ids " +- "it needs. Please run with PP=1." ++ from vllm.distributed.utils import get_pp_indices ++ ++ rank0_start, rank0_end = get_pp_indices( ++ int(text_config.num_hidden_layers), ++ 0, ++ parallel_config.pipeline_parallel_size, + ) ++ ple_idx = [int(i) - 1 for i in text_config.ple_layer_ids] ++ if min(ple_idx) < rank0_start or max(ple_idx) >= rank0_end: ++ raise NotImplementedError( ++ "Qwen4Exp N-gram PLE embedding requires every PLE layer to " ++ f"live on pipeline rank 0 (layers {rank0_start}.." ++ f"{rank0_end - 1}), but PLE layers are at {ple_idx}. Give " ++ "rank 0 more layers via VLLM_PP_LAYER_PARTITION, or run " ++ "with PP=1." ++ ) + multimodal_config = vllm_config.model_config.multimodal_config + if multimodal_config is not None and multimodal_config.language_model_only: + _strip_qwen4_exp_mrope(vllm_config.model_config) +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 3f884cd..7ac0ccb 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -58,6 +58,7 @@ from vllm.model_executor.models.qwen3_vl import ( + ) + from vllm.model_executor.models.utils import ( + AutoWeightsLoader, ++ PPMissingLayer, + StageMissingLayer, + WeightsMapper, + _merge_multimodal_embeddings, +@@ -774,11 +775,16 @@ class Qwen4ExpForCausalLM( + self.model = Qwen4ExpModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) +- self.lm_head = ParallelLMHead( +- config.vocab_size, +- config.hidden_size, +- prefix=maybe_prefix(prefix, "lm_head"), +- ) ++ # PP ownership: logits are computed on the ++ # last stage only; earlier stages carried an unused 0.59 GiB head. ++ if get_pp_group().is_last_rank: ++ self.lm_head = ParallelLMHead( ++ config.vocab_size, ++ config.hidden_size, ++ prefix=maybe_prefix(prefix, "lm_head"), ++ ) ++ else: ++ self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors +@@ -1011,6 +1017,14 @@ class Qwen4ExpForConditionalGeneration( + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.language_model_only = multimodal_config.language_model_only ++ # PP ownership: the encoder cache and the image ++ # embeddings exist on the first pipeline stage only; later stages ++ # carried an unused vision tower (~0.42 GiB per GPU at TP2). Absence ++ # on a later stage is an ownership decision, not language_model_only: ++ # the multimodal configuration stays global. ++ self._owns_vision_tower = ( ++ not self.language_model_only and get_pp_group().is_first_rank ++ ) + if self.language_model_only: + self.use_data_parallel = False + self.is_multimodal_pruning_enabled = False +@@ -1019,6 +1033,12 @@ class Qwen4ExpForConditionalGeneration( + self._tokenizer = None + self.visual = StageMissingLayer("vision_tower") + self._tower_model_names = [] ++ elif not self._owns_vision_tower: ++ self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" ++ self._init_video_pruning(multimodal_config) ++ self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) ++ self.visual = StageMissingLayer("vision_tower") ++ self._tower_model_names = [] + else: + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self._init_video_pruning(multimodal_config) +@@ -1064,7 +1084,16 @@ class Qwen4ExpForConditionalGeneration( + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) +- if not get_pp_group().is_first_rank and self.use_deepstack: ++ # Deepstack levels are injected into the first len(indexes) decoder ++ # layers on the first stage and are not forwarded through PP; keep the ++ # boundary assertion independent of use_deepstack (False on stages ++ # without a tower) so a future checkpoint with deepstack cannot be ++ # split across the boundary silently. ++ if ( ++ not get_pp_group().is_first_rank ++ and not self.language_model_only ++ and bool(config.vision_config.deepstack_visual_indexes) ++ ): + assert self.language_model.model.start_layer >= len( + config.vision_config.deepstack_visual_indexes + ), ( +diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py +index 2461685..092fef7 100644 +--- a/vllm/models/qwen4_exp/nvidia/model_state.py ++++ b/vllm/models/qwen4_exp/nvidia/model_state.py +@@ -32,13 +32,7 @@ class Qwen4ExpModelState(MambaHybridModelState): + self.ngram_eos_token_id = 0 + return + +- if vllm_config.parallel_config.pipeline_parallel_size > 1: +- raise RuntimeError( +- "N-gram PLE embedding currently requires " +- "pipeline_parallel_size=1 because non-first pipeline ranks do " +- "not receive the raw input_ids required by PLE. Please run " +- "with PP=1." +- ) ++ # vllm#54709: PP gate moved to config.py (placement-based). + + self.ngram_context_len = int(config.ngram_size) - 1 + if self.ngram_context_len <= 0: +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index d12b851..5dc0fc4 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -290,7 +290,11 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + hidden_size = self.hidden_size + prev_block_output: torch.Tensor | None = None + +- if get_pp_group().is_first_rank: ++ # Branch on the inputs, not the rank (vllm PR #46994's rule, applied ++ # there to Qwen3_5MTP / DeepSeek MTP): the drafter is built entirely on ++ # the LAST pipeline stage, where is_first_rank is False, and it always ++ # starts from the target's hidden state + token ids it embeds itself. ++ if intermediate_tensors is None: + assert hidden_states is not None + if inputs_embeds is None: + assert input_ids is not None +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index f090135..eead9c9 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -231,6 +231,118 @@ class Worker(WorkerBase): + # Resolved lazily on first sleep/wake; persists worker-process state. + self._sleep_mode_backend: SleepModeBackend | None = None + ++ def _has_ple_layers(self) -> bool: ++ """Return whether this model configuration constructs PLE layers.""" ++ if not envs.VLLM_PLE_CPU_OFFLOAD: ++ return False ++ text_config = self.model_config.hf_text_config ++ ple_layer_ids = getattr(text_config, "ple_layer_ids", None) ++ if not ple_layer_ids: ++ return False ++ # PP (vllm#54709, carried from the earlier overlay): only the pipeline rank ++ # that owns a PLE layer takes part in the offload (rank 0 with the ++ # default partition). Other ranks build PPMissingLayer there and must ++ # skip the connector entirely. ++ pp_size = self.parallel_config.pipeline_parallel_size ++ if pp_size > 1: ++ from vllm.distributed.utils import get_pp_indices ++ ++ tp_size = self.parallel_config.tensor_parallel_size ++ pp_rank = (self.rank // tp_size) % pp_size ++ start, end = get_pp_indices( ++ int(text_config.num_hidden_layers), pp_rank, pp_size ++ ) ++ return any(start <= int(i) - 1 < end for i in ple_layer_ids) ++ return True ++ ++ def _validate_ple_offload_config(self) -> None: ++ """Reject unsupported PLE offload execution modes.""" ++ parallel_config = self.parallel_config ++ unsupported = [] ++ if not self.use_v2_model_runner: ++ unsupported.append("model runner V1") ++ if not current_platform.is_cuda(): ++ unsupported.append(f"device={current_platform.device_type}") ++ if parallel_config.nnodes != 1: ++ unsupported.append(f"nnodes={parallel_config.nnodes}") ++ if parallel_config.data_parallel_backend != "mp": ++ unsupported.append(f"DP backend={parallel_config.data_parallel_backend}") ++ if ( ++ parallel_config.data_parallel_size_local ++ != parallel_config.data_parallel_size ++ ): ++ unsupported.append( ++ "non-local DP " ++ f"({parallel_config.data_parallel_size_local}/" ++ f"{parallel_config.data_parallel_size} local ranks)" ++ ) ++ # PP allowed (vllm#54709): PLE layers must live on pipeline rank 0 ++ # (enforced placement-based in models/config.py); num_workers stays ++ # dp_size * tp_size = the ranks of the owning stage. ++ if parallel_config.prefill_context_parallel_size != 1: ++ unsupported.append(f"PCP={parallel_config.prefill_context_parallel_size}") ++ if parallel_config.decode_context_parallel_size != 1: ++ unsupported.append(f"DCP={parallel_config.decode_context_parallel_size}") ++ if parallel_config.use_ubatching: ++ unsupported.append("ubatching/DBO") ++ if self.model_config.architecture not in { ++ "Qwen4ExpForCausalLM", ++ "Qwen4ExpForConditionalGeneration", ++ }: ++ unsupported.append("architecture") ++ if self.vllm_config.weight_transfer_config is not None: ++ unsupported.append("weight transfer") ++ ++ if unsupported: ++ raise ValueError( ++ "VLLM_PLE_CPU_OFFLOAD does not support the requested " ++ "configuration. Unsupported settings: " ++ f"{', '.join(unsupported)}" ++ ) ++ ++ def spawn_ple_offload(self) -> None: ++ """Spawn one node-local PLE CPU worker from DP0/TP0.""" ++ if ( ++ not self._ple_offload_enabled ++ or self.rank != 0 ++ or self.parallel_config.data_parallel_rank != 0 ++ ): ++ return ++ ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ ipc_addr = self.parallel_config._ple_offload_ipc_path ++ if not ipc_addr: ++ raise RuntimeError("PLE offload IPC address was not initialized") ++ dp_size = self.parallel_config.data_parallel_size ++ tp_size = self.parallel_config.tensor_parallel_size ++ num_workers = dp_size * tp_size ++ logger.info( ++ "PleOffload: spawning worker " ++ "(rank=%d, local_rank=%d, dp_size=%d, tp_size=%d, " ++ "num_workers=%d, ipc_addr=%s).", ++ self.rank, ++ self.local_rank, ++ dp_size, ++ tp_size, ++ num_workers, ++ ipc_addr, ++ ) ++ self._ple_offload_worker_handle = PleOffloadWorker.make_process( ++ self.vllm_config, ++ num_workers, ++ ipc_addr, ++ ) ++ ++ def wait_ple_offload_ready(self) -> None: ++ """Wait until weights and every TP registration are ready.""" ++ if self._ple_offload_worker_handle is None: ++ return ++ ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ PleOffloadWorker.wait_for_ready(self._ple_offload_worker_handle) ++ + @property + def sleep_mode_backend(self) -> "SleepModeBackend": + if self._sleep_mode_backend is None: +-- +2.43.0 + diff --git a/patches/0004-ple_offload-standalone-fault-injection-singleton-use.patch b/patches/0004-ple_offload-standalone-fault-injection-singleton-use.patch new file mode 100644 index 0000000..4f29ae4 --- /dev/null +++ b/patches/0004-ple_offload-standalone-fault-injection-singleton-use.patch @@ -0,0 +1,423 @@ +From 8b49ca6aeeccc87962e3976d564932e099f4c573 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:11:36 -0400 +Subject: [PATCH 004/103] ple_offload: standalone fault-injection singleton + used by the offload transport + +A stdlib-only, one-shot, rank-selective fault selector shared by the PLE +offload transport and its tests. Inert unless armed: disarmed, every site is a +single None check outside any compiled or captured region. +--- + vllm/v1/ple_offload/test_fault.py | 9 + + vllm_ple_test_fault/__init__.py | 383 ++++++++++++++++++++++++++++++ + 2 files changed, 392 insertions(+) + create mode 100644 vllm/v1/ple_offload/test_fault.py + create mode 100644 vllm_ple_test_fault/__init__.py + +diff --git a/vllm/v1/ple_offload/test_fault.py b/vllm/v1/ple_offload/test_fault.py +new file mode 100644 +index 0000000..3546a1f +--- /dev/null ++++ b/vllm/v1/ple_offload/test_fault.py +@@ -0,0 +1,9 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Compatibility exports for the top-level test-only fault singleton.""" ++ ++from vllm_ple_test_fault import ( # noqa: F401 ++ FaultSpec as FaultSpec, ++ _FAULT as _FAULT, ++ _parse_fault_spec as _parse_fault_spec, ++) +diff --git a/vllm_ple_test_fault/__init__.py b/vllm_ple_test_fault/__init__.py +new file mode 100644 +index 0000000..e8fe835 +--- /dev/null ++++ b/vllm_ple_test_fault/__init__.py +@@ -0,0 +1,383 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Stdlib-only, test-only PLE fault injection. ++ ++This package intentionally lives outside ``vllm`` so the offload bootstrap can ++import the same singleton without pre-populating part of the vLLM package tree. ++""" ++ ++import os ++import re ++import threading ++import time ++from dataclasses import dataclass, field ++ ++_ENV_NAME = "VLLM_PLE_TEST_FAULT" ++_LOG_ENV_NAME = "VLLM_PLE_TEST_FAULT_LOG" ++_UINT_RE = re.compile(r"(?:0|[1-9][0-9]*)\Z") ++_POSITIVE_INT_RE = re.compile(r"[1-9][0-9]*\Z") ++_TOKEN_RE = re.compile(r"[A-Za-z0-9_.:/@+-]+\Z") ++_MAX_RECORD_BYTES = 512 ++ ++# cut -> (role, counting unit, TP0-only) ++_CUT_META = { ++ "exc_before_prepare": ("gpu", "runner-forward", False), ++ "exc_after_prepare": ("gpu", "runner-forward", False), ++ "exc_in_release": ("gpu", "connector-release", False), ++ "exc_after_release": ("gpu", "runner-forward", False), ++ "exc_async_output": ("gpu", "runner-output-relay", False), ++ "drop_request": ("gpu", "connector-request", True), ++ "block_input_completion": ("gpu", "connector-request", True), ++ "park_before_publish": ("offload", "offload-request", False), ++ "park_before_publish_ack": ("offload", "offload-request", False), ++ "launch_park_before_payload_read": ( ++ "bootstrap", ++ "bootstrap-launch", ++ False, ++ ), ++ "launch_partial_ready": ("bootstrap", "bootstrap-launch", False), ++ "launch_partial_ready_header": ( ++ "bootstrap", ++ "bootstrap-launch", ++ False, ++ ), ++ "launch_silent_ready": ("bootstrap", "bootstrap-launch", False), ++ "launch_block_reconstruction": ( ++ "bootstrap", ++ "bootstrap-reconstruction", ++ False, ++ ), ++ "launch_block_worker_import": ( ++ "bootstrap", ++ "bootstrap-import", ++ False, ++ ), ++} ++ ++ ++def _pid_start_time(pid: int | None = None) -> int: ++ """Return Linux /proc stat field 22 without trusting the comm field.""" ++ if pid is None: ++ pid = os.getpid() ++ with open(f"/proc/{pid}/stat", encoding="ascii") as stream: ++ stat = stream.read() ++ close_paren = stat.rfind(")") ++ if close_paren < 0: ++ raise RuntimeError("invalid /proc stat record") ++ fields_from_three = stat[close_paren + 2 :].split() ++ if len(fields_from_three) <= 19: ++ raise RuntimeError("truncated /proc stat record") ++ value = fields_from_three[19] ++ if _POSITIVE_INT_RE.fullmatch(value) is None: ++ raise RuntimeError("invalid /proc start time") ++ return int(value) ++ ++ ++def _record_token(value: str | int | None) -> str: ++ if value is None: ++ return "-" ++ token = str(value).replace("\n", "\\n").replace(" ", "_")[:160] ++ if not token or _TOKEN_RE.fullmatch(token) is None: ++ raise ValueError("fault diagnostic fields must be non-empty tokens") ++ return token ++ ++ ++def _block_on_harness_gate() -> None: ++ """Wait in a FIFO read that releases the GIL.""" ++ path = os.environ.get("VLLM_PLE_TEST_FAULT_GATE") ++ if path is None or not path: ++ raise RuntimeError( ++ "VLLM_PLE_TEST_FAULT_GATE is required for this bootstrap cut" ++ ) ++ fd = os.open(path, os.O_RDONLY) ++ try: ++ if os.read(fd, 1) == b"": ++ raise EOFError("PLE bootstrap fault gate closed") ++ finally: ++ os.close(fd) ++ ++ ++@dataclass ++class _FaultState: ++ count: int = 0 ++ fired: bool = False ++ diagnostic_failed: bool = False ++ lock: threading.Lock = field(default_factory=threading.Lock) ++ ++ ++@dataclass(frozen=True) ++class FaultSpec: ++ """Immutable fault selection with process-local one-shot runtime state.""" ++ ++ cut: str ++ ordinal: int ++ tp_rank: int | None ++ pp_rank: int | None ++ phase: str ++ role: str ++ unit: str ++ log_fd: int = field(compare=False, repr=False) ++ _state: _FaultState = field( ++ default_factory=_FaultState, ++ compare=False, ++ repr=False, ++ ) ++ ++ def _emit( ++ self, ++ *, ++ count: int, ++ role: str, ++ unit: str, ++ tp_rank: int | None, ++ pp_rank: int | None, ++ dp_rank: int | None, ++ transaction_id: str | int | None, ++ correlation_id: str | int | None, ++ bytes_written: int | None = None, ++ bytes_sent: int | None = None, ++ frame_size: int | None = None, ++ ) -> bool: ++ try: ++ pid = os.getpid() ++ suffix = "" ++ if bytes_written is not None: ++ if bytes_sent is not None or frame_size is None: ++ return False ++ suffix = ( ++ f" bytes_written={int(bytes_written)}" ++ f" frame_size={int(frame_size)}" ++ ) ++ elif bytes_sent is not None: ++ if frame_size is None: ++ return False ++ suffix = ( ++ f" bytes_sent={int(bytes_sent)}" ++ f" frame_size={int(frame_size)}" ++ ) ++ elif frame_size is not None: ++ return False ++ message = ( ++ f"PLE_TEST_FAULT v=1 cut={self.cut} ordinal={self.ordinal} " ++ f"count={count} unit={_record_token(unit)} " ++ f"role={_record_token(role)} " ++ f"tp={tp_rank if tp_rank is not None else '-'} " ++ f"pp={pp_rank if pp_rank is not None else '-'} " ++ f"dp={dp_rank if dp_rank is not None else '-'} " ++ f"pid={pid} pid_start={_pid_start_time(pid)} " ++ f"transaction={_record_token(transaction_id)} " ++ f"correlation={_record_token(correlation_id)} " ++ f"monotonic={time.monotonic():.9f} " ++ f"wall={time.time():.9f}{suffix}\n" ++ ).encode("ascii", "strict") ++ if len(message) > _MAX_RECORD_BYTES: ++ return False ++ return os.write(self.log_fd, message) == len(message) ++ except BaseException: ++ return False ++ ++ def fires( ++ self, ++ cut: str, ++ *, ++ role: str, ++ unit: str, ++ eligible: bool, ++ tp_rank: int | None = None, ++ pp_rank: int | None = None, ++ dp_rank: int | None = None, ++ transaction_id: str | int | None = None, ++ correlation_id: str | int | None = None, ++ bytes_written: int | None = None, ++ frame_size: int | None = None, ++ ) -> bool: ++ """Count an eligible unit and fire only after a complete marker write.""" ++ if cut != self.cut or role != self.role or unit != self.unit: ++ return False ++ # phase=any does not make profile, warm-up, or dummy GPU work serving. ++ if not eligible: ++ return False ++ if self.role == "gpu": ++ if tp_rank is None or pp_rank is None or dp_rank is None: ++ raise ValueError("GPU fault sites must provide TP, PP, and DP ranks") ++ if self.tp_rank is not None and tp_rank != self.tp_rank: ++ return False ++ if self.pp_rank is not None and pp_rank != self.pp_rank: ++ return False ++ elif tp_rank is not None or pp_rank is not None: ++ raise ValueError(f"{self.role} fault sites must not provide GPU ranks") ++ ++ with self._state.lock: ++ if self._state.fired or self._state.diagnostic_failed: ++ return False ++ self._state.count += 1 ++ count = self._state.count ++ if count != self.ordinal: ++ return False ++ if not self._emit( ++ count=count, ++ role=role, ++ unit=unit, ++ tp_rank=tp_rank, ++ pp_rank=pp_rank, ++ dp_rank=dp_rank, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ bytes_written=bytes_written, ++ frame_size=frame_size, ++ ): ++ self._state.diagnostic_failed = True ++ return False ++ self._state.fired = True ++ return True ++ ++ def acknowledge_launch_send(self, *, bytes_sent: int, frame_size: int) -> bool: ++ """Emit bounded send progress only for the CPU-scoped launch-send gate.""" ++ if self.cut != "launch_park_before_payload_read": ++ return False ++ with self._state.lock: ++ if self._state.diagnostic_failed: ++ return False ++ ok = self._emit( ++ count=self.ordinal, ++ role="launch-parent", ++ unit="launch-send", ++ tp_rank=None, ++ pp_rank=None, ++ dp_rank=None, ++ transaction_id="launch:1", ++ correlation_id="launch:1", ++ bytes_sent=bytes_sent, ++ frame_size=frame_size, ++ ) ++ if not ok: ++ self._state.diagnostic_failed = True ++ return ok ++ ++ ++@dataclass(frozen=True) ++class _FaultReconstructionPayload: ++ """Pickle wrapper whose reducer is the reconstruction fault boundary.""" ++ ++ payload: bytes ++ ++ def __reduce__(self): ++ return (_reconstruct_configuration_payload, (self.payload,)) ++ ++ ++def _reconstruct_configuration_payload(payload: bytes) -> bytes: ++ fault = _FAULT ++ if fault is not None and fault.fires( ++ "launch_block_reconstruction", ++ role="bootstrap", ++ unit="bootstrap-reconstruction", ++ eligible=True, ++ transaction_id="reducer:configuration", ++ correlation_id="launch:1", ++ ): ++ _block_on_harness_gate() ++ return payload ++ ++ ++def _parse_uint(value: str, name: str, *, positive: bool = False) -> int: ++ pattern = _POSITIVE_INT_RE if positive else _UINT_RE ++ if pattern.fullmatch(value) is None: ++ requirement = "a positive integer" if positive else "a non-negative integer" ++ raise ValueError(f"{_ENV_NAME}: {name} must be {requirement}") ++ return int(value) ++ ++ ++def _open_fault_log() -> int: ++ path = os.environ.get(_LOG_ENV_NAME) ++ if path is None or not path or path != path.strip(): ++ raise ValueError( ++ f"{_LOG_ENV_NAME} is required and must be a non-empty path " ++ f"when {_ENV_NAME} is armed" ++ ) ++ flags = os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NONBLOCK ++ try: ++ return os.open(path, flags, 0o600) ++ except OSError as error: ++ raise ValueError(f"{_LOG_ENV_NAME}: cannot open {path!r}") from error ++ ++ ++def _parse_fault_spec(value: str | None) -> FaultSpec | None: ++ """Strictly parse the environment representation.""" ++ if value is None: ++ return None ++ if not value or value != value.strip(): ++ raise ValueError(f"{_ENV_NAME}: spec must be non-empty with no outer whitespace") ++ ++ fields: dict[str, str] = {} ++ for item in value.split(","): ++ if item.count("=") != 1: ++ raise ValueError(f"{_ENV_NAME}: every field must be key=value") ++ key, field_value = item.split("=", 1) ++ if ( ++ not key ++ or not field_value ++ or key != key.strip() ++ or field_value != field_value.strip() ++ ): ++ raise ValueError(f"{_ENV_NAME}: fields must not contain whitespace") ++ if key in fields: ++ raise ValueError(f"{_ENV_NAME}: duplicate field {key!r}") ++ fields[key] = field_value ++ ++ allowed = {"cut", "ordinal", "rank", "tp", "pp", "phase"} ++ unknown = set(fields).difference(allowed) ++ if unknown: ++ raise ValueError(f"{_ENV_NAME}: unknown field(s): {sorted(unknown)}") ++ missing = {"cut", "ordinal"}.difference(fields) ++ if missing: ++ raise ValueError(f"{_ENV_NAME}: missing field(s): {sorted(missing)}") ++ ++ cut = fields["cut"] ++ if cut not in _CUT_META: ++ raise ValueError(f"{_ENV_NAME}: unknown cut {cut!r}") ++ ordinal = _parse_uint(fields["ordinal"], "ordinal", positive=True) ++ role, unit, tp0_only = _CUT_META[cut] ++ ++ has_any_rank = "rank" in fields ++ has_specific_rank = "tp" in fields or "pp" in fields ++ if has_any_rank: ++ if fields["rank"] != "any": ++ raise ValueError(f"{_ENV_NAME}: rank must be exactly 'any'") ++ if has_specific_rank: ++ raise ValueError(f"{_ENV_NAME}: rank=any cannot be combined with tp or pp") ++ tp_rank = None ++ pp_rank = None ++ else: ++ if not has_specific_rank: ++ raise ValueError(f"{_ENV_NAME}: specify rank=any, tp, or pp") ++ tp_rank = _parse_uint(fields["tp"], "tp") if "tp" in fields else None ++ pp_rank = _parse_uint(fields["pp"], "pp") if "pp" in fields else None ++ ++ if role in {"offload", "bootstrap"} and not has_any_rank: ++ raise ValueError( ++ f"{_ENV_NAME}: {role} cut {cut!r} requires rank=any; " ++ "it has no TP/PP rank" ++ ) ++ if role == "gpu": ++ if pp_rank is not None and pp_rank != 0: ++ raise ValueError(f"{_ENV_NAME}: cut {cut!r} exists only on PP rank zero") ++ if tp0_only and tp_rank is not None and tp_rank != 0: ++ raise ValueError(f"{_ENV_NAME}: cut {cut!r} exists only on TP rank zero") ++ ++ phase = fields.get("phase", "serving") ++ if phase not in {"serving", "any"}: ++ raise ValueError(f"{_ENV_NAME}: phase must be 'serving' or 'any'") ++ ++ return FaultSpec( ++ cut=cut, ++ ordinal=ordinal, ++ tp_rank=tp_rank, ++ pp_rank=pp_rank, ++ phase=phase, ++ role=role, ++ unit=unit, ++ log_fd=_open_fault_log(), ++ ) ++ ++ ++_FAULT = _parse_fault_spec(os.environ.get(_ENV_NAME)) +-- +2.43.0 + diff --git a/patches/0005-executor-workers-cannot-outlive-a-killed-EngineCore.patch b/patches/0005-executor-workers-cannot-outlive-a-killed-EngineCore.patch new file mode 100644 index 0000000..f905a88 --- /dev/null +++ b/patches/0005-executor-workers-cannot-outlive-a-killed-EngineCore.patch @@ -0,0 +1,536 @@ +From 5a84c1bc7dc5e9a54d5a14ce1766adeb6f1dca10 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:11:43 -0400 +Subject: [PATCH 005/103] executor: workers cannot outlive a killed EngineCore + +Workers set PR_SET_PDEATHSIG(SIGKILL) against their EngineCore parent with a +getppid race check. Start methods other than fork/spawn are refused before any +worker is created. Subprocess tests cover fork, spawn, graceful SIGTERM and the +refusal. +--- + .../test_multiproc_executor_parent_death.py | 350 ++++++++++++++++++ + .../test_multiproc_executor_start_method.py | 63 ++++ + vllm/v1/executor/multiproc_executor.py | 52 +++ + 3 files changed, 465 insertions(+) + create mode 100644 tests/v1/executor/test_multiproc_executor_parent_death.py + create mode 100644 tests/v1/executor/test_multiproc_executor_start_method.py + +diff --git a/tests/v1/executor/test_multiproc_executor_parent_death.py b/tests/v1/executor/test_multiproc_executor_parent_death.py +new file mode 100644 +index 0000000..c39718f +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_executor_parent_death.py +@@ -0,0 +1,350 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import multiprocessing as mp ++import os ++import selectors ++import signal ++import subprocess ++import sys ++import time ++from contextlib import suppress ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++ ++_PARENT_DEATH_EXIT_BOUND_S = 5.0 ++_STARTUP_TIMEOUT_S = 180.0 ++_TEST_FILE = Path(__file__).resolve() ++ ++pytestmark = [ ++ pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux only"), ++ pytest.mark.timeout(240), ++] ++ ++ ++def _proc_identity(pid: int) -> tuple[str, str] | None: ++ try: ++ stat = Path(f"/proc/{pid}/stat").read_text() ++ except FileNotFoundError: ++ return None ++ closing_paren = stat.rfind(")") ++ fields = stat[closing_paren + 2 :].split() ++ # fields starts at proc(5)'s field 3: state; starttime is field 22. ++ return fields[0], fields[19] ++ ++ ++def _run_fake_worker_main( ++ ready_pipe, ++ workload_entered_pipe, ++ death_pipe, ++ inherited_fds: list[int], ++ expected_parent_pid: int, ++ block_mode: str, ++) -> None: ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ real_worker_proc = executor_module.WorkerProc ++ worker_main = real_worker_proc.worker_main ++ ++ class FakeMessageQueue: ++ def export_handle(self): ++ identity = _proc_identity(os.getpid()) ++ assert identity is not None ++ return os.getpid(), identity[1] ++ ++ def wait_until_ready(self): ++ pass ++ ++ def shutdown(self): ++ pass ++ ++ class FakeWorkerProc: ++ READY_STR = real_worker_proc.READY_STR ++ ++ def __init__(self, *args, **kwargs): ++ self.rank = 0 ++ self.rpc_broadcast_mq = None ++ self.worker_response_mq = FakeMessageQueue() ++ self.peer_response_handles = [] ++ ++ def monitor_death_pipe(self, pipe, shutdown_requested): ++ return real_worker_proc.monitor_death_pipe( ++ self, pipe, shutdown_requested ++ ) ++ ++ def worker_busy_loop(self): ++ if block_mode == "gil-spin": ++ # Model a C extension that holds the GIL while polling. ++ sys.setswitchinterval(1000.0) ++ workload_entered_pipe.send(block_mode) ++ workload_entered_pipe.close() ++ while True: ++ pass ++ if block_mode == "sleep": ++ workload_entered_pipe.send(block_mode) ++ workload_entered_pipe.close() ++ while True: ++ time.sleep(3600) ++ raise AssertionError(f"unknown block mode: {block_mode}") ++ ++ def shutdown(self): ++ pass ++ ++ executor_module.WorkerProc = FakeWorkerProc ++ executor_module.maybe_init_worker_tracer = lambda **kwargs: None ++ executor_module.set_worker_net_device = lambda *args, **kwargs: None ++ config = SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, ++ numa_bind=False, ++ ) ++ ) ++ worker_main( ++ vllm_config=config, ++ ready_pipe=ready_pipe, ++ death_pipe=death_pipe, ++ inherited_fds=inherited_fds, ++ expected_parent_pid=expected_parent_pid, ++ ) ++ ++ ++def _worker_like_child( ++ ready_pipe, ++ workload_entered_pipe, ++ death_pipe, ++ inherited_fds: list[int], ++ expected_parent_pid: int, ++ block_mode: str, ++) -> None: ++ _run_fake_worker_main( ++ ready_pipe, ++ workload_entered_pipe, ++ death_pipe, ++ inherited_fds, ++ expected_parent_pid, ++ block_mode, ++ ) ++ ++ ++def _orphan_parent(start_method: str, block_mode: str) -> None: ++ context = mp.get_context(start_method) ++ ready_reader, ready_writer = context.Pipe(duplex=False) ++ workload_reader, workload_writer = context.Pipe(duplex=False) ++ death_reader, death_writer = context.Pipe(duplex=False) ++ inherited_fds = [] ++ if start_method == "fork": ++ inherited_fds.extend( ++ ( ++ ready_reader.fileno(), ++ workload_reader.fileno(), ++ death_writer.fileno(), ++ ) ++ ) ++ ++ proc = context.Process( ++ target=_worker_like_child, ++ args=( ++ ready_writer, ++ workload_writer, ++ death_reader, ++ inherited_fds, ++ os.getpid(), ++ block_mode, ++ ), ++ daemon=True, ++ ) ++ proc.start() ++ ready_writer.close() ++ workload_writer.close() ++ death_reader.close() ++ ++ if not ready_reader.poll(_STARTUP_TIMEOUT_S): ++ raise TimeoutError("worker-like child did not become ready") ++ response = ready_reader.recv() ++ ready_reader.close() ++ if not workload_reader.poll(_STARTUP_TIMEOUT_S): ++ raise TimeoutError("worker-like child did not enter its workload") ++ assert workload_reader.recv() == block_mode ++ workload_reader.close() ++ child_pid, start_time = response["handle"] ++ print(f"READY {child_pid} {start_time}", flush=True) ++ ++ # Keep the death-pipe writer open until the test SIGKILLs this process. ++ assert death_writer is not None ++ while True: ++ signal.pause() ++ ++ ++class _StdoutReadyPipe: ++ def send(self, response) -> None: ++ child_pid, start_time = response["handle"] ++ print(f"READY {child_pid} {start_time}", flush=True) ++ ++ def close(self) -> None: ++ pass ++ ++ ++class _UnusedWorkloadPipe: ++ def send(self, value) -> None: ++ pass ++ ++ def close(self) -> None: ++ pass ++ ++ ++def _graceful_child() -> None: ++ _run_fake_worker_main( ++ _StdoutReadyPipe(), ++ _UnusedWorkloadPipe(), ++ death_pipe=None, ++ inherited_fds=[], ++ expected_parent_pid=os.getppid(), ++ block_mode="sleep", ++ ) ++ ++ ++def _run_cli() -> None: ++ if sys.argv[1] == "orphan-parent": ++ _orphan_parent(sys.argv[2], sys.argv[3]) ++ elif sys.argv[1] == "graceful-child": ++ _graceful_child() ++ else: ++ raise AssertionError(f"unknown command: {sys.argv[1]}") ++ ++ ++def _start_cli(*args: str) -> subprocess.Popen: ++ return subprocess.Popen( ++ [sys.executable, str(_TEST_FILE), *args], ++ stdout=subprocess.PIPE, ++ stderr=subprocess.STDOUT, ++ start_new_session=True, ++ ) ++ ++ ++def _wait_for_ready(proc: subprocess.Popen) -> tuple[int, str]: ++ assert proc.stdout is not None ++ selector = selectors.DefaultSelector() ++ selector.register(proc.stdout, selectors.EVENT_READ) ++ deadline = time.monotonic() + _STARTUP_TIMEOUT_S ++ buffered = b"" ++ try: ++ while time.monotonic() < deadline: ++ timeout = min( ++ 1.0, ++ max(0.0, deadline - time.monotonic()), ++ ) ++ if not selector.select(timeout): ++ if proc.poll() is not None: ++ raise AssertionError( ++ f"launcher exited before readiness: {proc.returncode}" ++ ) ++ continue ++ chunk = os.read(proc.stdout.fileno(), 4096) ++ if not chunk: ++ raise AssertionError("launcher closed stdout before readiness") ++ buffered += chunk ++ while b"\n" in buffered: ++ line, buffered = buffered.split(b"\n", 1) ++ if line.startswith(b"READY "): ++ _, pid, start_time = line.split() ++ return int(pid), start_time.decode() ++ raise AssertionError("timed out waiting for worker-like child readiness") ++ finally: ++ selector.close() ++ ++ ++def _wait_for_state( ++ pid: int, start_time: str, expected_state: str, deadline: float ++) -> bool: ++ while time.monotonic() < deadline: ++ identity = _proc_identity(pid) ++ if identity is None or identity[1] != start_time: ++ return False ++ if identity[0] == expected_state: ++ return True ++ time.sleep(0.01) ++ return False ++ ++ ++def _wait_for_exit(pid: int, start_time: str, deadline: float) -> float | None: ++ while time.monotonic() < deadline: ++ identity = _proc_identity(pid) ++ if ( ++ identity is None ++ or identity[0] in {"X", "Z"} ++ or identity[1] != start_time ++ ): ++ return time.monotonic() ++ time.sleep(0.01) ++ return None ++ ++ ++def _kill_process_group(proc: subprocess.Popen) -> None: ++ with suppress(ProcessLookupError): ++ os.killpg(proc.pid, signal.SIGKILL) ++ with suppress(subprocess.TimeoutExpired): ++ proc.wait(timeout=10) ++ if proc.stdout is not None: ++ proc.stdout.close() ++ ++ ++@pytest.mark.parametrize( ++ ("start_method", "block_mode", "expected_state"), ++ [("fork", "gil-spin", "R"), ("spawn", "sleep", "S")], ++) ++def test_worker_dies_when_engine_core_parent_is_sigkilled( ++ start_method: str, block_mode: str, expected_state: str ++): ++ launcher = _start_cli("orphan-parent", start_method, block_mode) ++ try: ++ child_pid, start_time = _wait_for_ready(launcher) ++ assert _wait_for_state( ++ child_pid, ++ start_time, ++ expected_state, ++ time.monotonic() + 10, ++ ), f"worker {child_pid} did not enter {block_mode} state" ++ assert launcher.poll() is None ++ ++ killed_at = time.monotonic() ++ os.kill(launcher.pid, signal.SIGKILL) ++ launcher.wait(timeout=_PARENT_DEATH_EXIT_BOUND_S) ++ exited_at = _wait_for_exit( ++ child_pid, ++ start_time, ++ killed_at + _PARENT_DEATH_EXIT_BOUND_S, ++ ) ++ ++ assert launcher.returncode == -signal.SIGKILL ++ assert exited_at is not None, ( ++ f"worker {child_pid} survived its parent by more than " ++ f"{_PARENT_DEATH_EXIT_BOUND_S:.1f}s" ++ ) ++ assert exited_at - killed_at <= _PARENT_DEATH_EXIT_BOUND_S ++ finally: ++ _kill_process_group(launcher) ++ ++ ++def test_sigterm_with_live_parent_uses_worker_system_exit_path(): ++ worker = _start_cli("graceful-child") ++ try: ++ child_pid, start_time = _wait_for_ready(worker) ++ assert child_pid == worker.pid ++ assert _proc_identity(child_pid) is not None ++ ++ os.kill(worker.pid, signal.SIGTERM) ++ worker.wait(timeout=10) ++ ++ # WorkerProc's SIGTERM handler raises and re-raises SystemExit, so this ++ # is a clean Python exit rather than default signal termination (-15). ++ assert worker.returncode == 0 ++ identity = _proc_identity(child_pid) ++ assert identity is None or identity[1] != start_time ++ finally: ++ _kill_process_group(worker) ++ ++ ++if __name__ == "__main__": ++ _run_cli() +diff --git a/tests/v1/executor/test_multiproc_executor_start_method.py b/tests/v1/executor/test_multiproc_executor_start_method.py +new file mode 100644 +index 0000000..6f00488 +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_executor_start_method.py +@@ -0,0 +1,63 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""The worker factory refuses start methods other than fork/spawn ++before creating any pipe or process (PR_SET_PDEATHSIG + getppid supervision ++would otherwise kill healthy workers under forkserver).""" ++ ++import multiprocessing ++import sys ++ ++import pytest ++ ++import vllm.v1.executor.multiproc_executor as executor_module ++ ++ ++class _NoCreateContext: ++ """Wraps a real context; any pipe/process creation fails the test.""" ++ ++ def __init__(self, method: str) -> None: ++ self._method = method ++ ++ def get_start_method(self) -> str: ++ return self._method ++ ++ def Pipe(self, *args, **kwargs): # noqa: N802 ++ raise AssertionError("pipe creation reached") ++ ++ def Process(self, *args, **kwargs): # noqa: N802 ++ raise AssertionError("process creation reached") ++ ++ ++@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux only") ++def test_forkserver_is_refused_before_creating_workers(monkeypatch): ++ assert "forkserver" in multiprocessing.get_all_start_methods() ++ monkeypatch.setattr( ++ executor_module, "get_mp_context", lambda: _NoCreateContext("forkserver") ++ ) ++ with pytest.raises(RuntimeError, match="'fork' or 'spawn'"): ++ executor_module.WorkerProc.make_worker_process( ++ vllm_config=None, ++ local_rank=0, ++ rank=0, ++ distributed_init_method="tcp://127.0.0.1:1", ++ input_shm_handle=None, ++ shared_worker_lock=None, ++ is_driver_worker=True, ++ ) ++ ++ ++@pytest.mark.parametrize("method", ["fork", "spawn"]) ++def test_fork_and_spawn_pass_the_guard(monkeypatch, method): ++ monkeypatch.setattr( ++ executor_module, "get_mp_context", lambda: _NoCreateContext(method) ++ ) ++ # The guard passes, so creation is attempted and hits the stub. ++ with pytest.raises(AssertionError, match="pipe creation reached"): ++ executor_module.WorkerProc.make_worker_process( ++ vllm_config=None, ++ local_rank=0, ++ rank=0, ++ distributed_init_method="tcp://127.0.0.1:1", ++ input_shm_handle=None, ++ shared_worker_lock=None, ++ is_driver_worker=True, ++ ) +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index 67a631d..aae270f 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -5,6 +5,7 @@ import os + import pickle + import queue + import signal ++import sys + import threading + import time + import traceback +@@ -75,6 +76,44 @@ from vllm.v1.worker.worker_base import WorkerWrapperBase + logger = init_logger(__name__) + + ++_WORKER_PARENT_DEATH_SIGNAL = signal.SIGKILL ++_PR_SET_PDEATHSIG = 1 ++ ++ ++def _install_worker_parent_death_signal(expected_parent_pid: int) -> None: ++ """Have Linux kill this worker if its EngineCore parent exits.""" ++ if not sys.platform.startswith("linux"): ++ return ++ ++ import ctypes ++ ++ libc = ctypes.CDLL(None, use_errno=True) ++ prctl = libc.prctl ++ prctl.argtypes = [ ++ ctypes.c_int, ++ ctypes.c_ulong, ++ ctypes.c_ulong, ++ ctypes.c_ulong, ++ ctypes.c_ulong, ++ ] ++ prctl.restype = ctypes.c_int ++ if prctl( ++ _PR_SET_PDEATHSIG, ++ int(_WORKER_PARENT_DEATH_SIGNAL), ++ 0, ++ 0, ++ 0, ++ ) != 0: ++ error_number = ctypes.get_errno() ++ raise OSError(error_number, os.strerror(error_number)) ++ ++ # The parent can exit between spawning this process and the prctl call. ++ # Once prctl is installed, a later exit is covered by the kernel. ++ if os.getppid() != expected_parent_pid: ++ os.kill(os.getpid(), _WORKER_PARENT_DEATH_SIGNAL) ++ os._exit(1) ++ ++ + class FutureWrapper(Future): + def __init__( + self, +@@ -713,6 +752,15 @@ class WorkerProc: + inherited_fds: list[int] | None = None, + ) -> UnreadyWorkerProcHandle: + context = get_mp_context() ++ # Workers arm PR_SET_PDEATHSIG against this process and verify it via ++ # getppid(). Under forkserver the OS parent is the forkserver, so a ++ # healthy worker would kill itself; refuse before creating anything. ++ start_method = context.get_start_method() ++ if start_method not in ("fork", "spawn"): ++ raise RuntimeError( ++ "vLLM worker processes require the 'fork' or 'spawn' " ++ f"multiprocessing start method (got {start_method!r})" ++ ) + # Ready pipe to communicate readiness from child to parent + ready_reader, ready_writer = context.Pipe(duplex=False) + # Death pipe to let child detect parent process exit +@@ -721,6 +769,7 @@ class WorkerProc: + inherited_fds = inherited_fds.copy() + inherited_fds.extend((ready_reader.fileno(), death_writer.fileno())) + process_kwargs = { ++ "expected_parent_pid": os.getpid(), + "vllm_config": vllm_config, + "local_rank": local_rank, + "rank": rank, +@@ -854,6 +903,9 @@ class WorkerProc: + """Worker initialization and execution loops. + This runs a background process""" + ++ expected_parent_pid = kwargs.pop("expected_parent_pid") ++ _install_worker_parent_death_signal(expected_parent_pid) ++ + # Signal handler used for graceful termination. + # SystemExit exception is only raised once to allow this and worker + # processes to terminate without error +-- +2.43.0 + diff --git a/patches/0006-qwen4_exp-host-resident-token-embedding-tables-VLLM_.patch b/patches/0006-qwen4_exp-host-resident-token-embedding-tables-VLLM_.patch new file mode 100644 index 0000000..596ab87 --- /dev/null +++ b/patches/0006-qwen4_exp-host-resident-token-embedding-tables-VLLM_.patch @@ -0,0 +1,927 @@ +From 609c86adc68243fabc81d603cd424687f829669a Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:12:31 -0400 +Subject: [PATCH 006/103] qwen4_exp: host-resident token embedding tables + (VLLM_HOST_EMBED_TABLE) + +Keeps the token embedding table in pinned host memory behind a UVA view for the +stage-0 target and the last-stage MTP drafter, freeing GPU memory for KV. TP +sharding, masking and the all-reduce are retained. The knob is registered as a +compile-cache factor. +--- + .../qwen4_exp/test_host_uva_embedding.py | 566 ++++++++++++++++++ + vllm/envs.py | 11 + + .../layers/vocab_parallel_embedding.py | 177 ++++++ + vllm/models/qwen4_exp/nvidia/model.py | 28 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 18 +- + 5 files changed, 797 insertions(+), 3 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_host_uva_embedding.py + +diff --git a/tests/models/qwen4_exp/test_host_uva_embedding.py b/tests/models/qwen4_exp/test_host_uva_embedding.py +new file mode 100644 +index 0000000..e55d5d5 +--- /dev/null ++++ b/tests/models/qwen4_exp/test_host_uva_embedding.py +@@ -0,0 +1,566 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import os ++from contextlib import nullcontext ++from types import SimpleNamespace ++ ++import pytest ++import torch ++from torch import nn ++ ++import vllm.model_executor.layers.vocab_parallel_embedding as vpe ++import vllm.model_executor.model_loader.utils as loader_utils ++import vllm.models.qwen4_exp.nvidia.model as model_mod ++import vllm.models.qwen4_exp.nvidia.mtp as mtp_mod ++from vllm.model_executor.layers.vocab_parallel_embedding import ( ++ HOST_EMBED_TABLE_ENV, ++ PinnedHostVocabParallelEmbedding, ++ VocabParallelEmbedding, ++ host_embed_table_enabled, ++) ++ ++ ++VOCAB = 67 ++HIDDEN = 16 ++ ++ ++def _set_tp(monkeypatch, rank: int, world_size: int) -> None: ++ monkeypatch.setattr(vpe, "get_tensor_model_parallel_rank", lambda: rank) ++ monkeypatch.setattr( ++ vpe, ++ "get_tensor_model_parallel_world_size", ++ lambda: world_size, ++ ) ++ monkeypatch.setattr( ++ vpe, ++ "tensor_model_parallel_all_reduce", ++ lambda value: value, ++ ) ++ ++ ++def _require_cuda_pinned_memory() -> None: ++ if not torch.cuda.is_available(): ++ pytest.skip("requires CUDA UVA") ++ if not vpe.PIN_MEMORY: ++ pytest.skip("pinned memory is unavailable") ++ try: ++ probe = torch.empty(1, pin_memory=True) ++ except RuntimeError as exc: ++ pytest.skip(f"pinned allocation is unavailable: {exc}") ++ assert probe.is_pinned() ++ ++ ++def _weight(vocab: int = VOCAB, hidden: int = HIDDEN) -> torch.Tensor: ++ values = torch.arange(vocab * hidden, dtype=torch.int32) ++ weight = values.reshape(vocab, hidden).to(torch.bfloat16) ++ if vocab > 64 and hidden >= 2: ++ bits = weight.view(torch.uint16) ++ bits[64, 0] = 0x8000 # -0.0 ++ bits[64, 1] = 0x7FC1 # quiet NaN with a nonzero payload ++ return weight ++ ++ ++def _assert_same_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None: ++ assert actual.shape == expected.shape ++ assert actual.dtype == expected.dtype ++ assert torch.equal( ++ actual.contiguous().view(torch.uint8), ++ expected.contiguous().view(torch.uint8), ++ ) ++ ++ ++@pytest.mark.parametrize( ++ ("value", "expected"), ++ [ ++ (None, False), ++ ("", False), ++ ("0", False), ++ ("1", True), ++ (" 1 ", False), ++ ("true", False), ++ ], ++) ++def test_host_embedding_flag_is_exact(monkeypatch, value, expected): ++ if value is None: ++ monkeypatch.delenv(HOST_EMBED_TABLE_ENV, raising=False) ++ else: ++ monkeypatch.setenv(HOST_EMBED_TABLE_ENV, value) ++ assert host_embed_table_enabled() is expected ++ ++ ++def test_u1_tp2_masking_and_reduction_are_bytewise_identical(monkeypatch): ++ weight = _weight() ++ ids = torch.tensor( ++ [0, 1, 63, 64, 65, 66, -1, VOCAB, 127], ++ dtype=torch.int64, ++ ) ++ stock_parts = [] ++ host_parts = [] ++ ++ for rank in range(2): ++ _set_tp(monkeypatch, rank, 2) ++ with torch.device("cpu"): ++ stock = VocabParallelEmbedding( ++ VOCAB, ++ HIDDEN, ++ params_dtype=torch.bfloat16, ++ ) ++ host = PinnedHostVocabParallelEmbedding( ++ VOCAB, ++ HIDDEN, ++ params_dtype=torch.bfloat16, ++ create_accelerator_view=False, ++ ) ++ ++ # The CPU reference deliberately exercises the stock masked fallback, ++ # not the CUDA-only fused kernel. ++ stock.use_fused_embedding = False ++ stock.weight_loader(stock.weight, weight) ++ host.weight_loader(host.weight, weight) ++ ++ stock_local = stock(ids) ++ host_local = host(ids) ++ _assert_same_bytes(host_local, stock_local) ++ stock_parts.append(stock_local) ++ host_parts.append(host_local) ++ ++ # Use the same floating-point SUM ordering for both implementations. This ++ # intentionally covers the signed-zero and NaN caveat of the real TP path. ++ stock_reduced = stock_parts[0] + stock_parts[1] ++ host_reduced = host_parts[0] + host_parts[1] ++ _assert_same_bytes(host_reduced, stock_reduced) ++ ++ invalid = host_reduced[-3:].contiguous().view(torch.uint8) ++ assert torch.count_nonzero(invalid) == 0 ++ ++ ++def test_u1_tp1_matches_stock_lookup_contract(monkeypatch): ++ _set_tp(monkeypatch, 0, 1) ++ weight = _weight() ++ with torch.device("cpu"): ++ stock = VocabParallelEmbedding( ++ VOCAB, ++ HIDDEN, ++ params_dtype=torch.bfloat16, ++ ) ++ host = PinnedHostVocabParallelEmbedding( ++ VOCAB, ++ HIDDEN, ++ params_dtype=torch.bfloat16, ++ create_accelerator_view=False, ++ ) ++ ++ assert stock.quant_method.requires_device_loading is True ++ assert host.quant_method.requires_device_loading is False ++ stock.weight_loader(stock.weight, weight) ++ host.weight_loader(host.weight, weight) ++ ++ # TP1 exposes padded storage just like an ordinary embedding lookup. ++ for row, value in ((VOCAB, 11), (stock.weight.shape[0] - 1, 29)): ++ stock.weight[row].fill_(value) ++ host.weight[row].fill_(value) ++ ++ ids = torch.tensor([0, 64, 66, VOCAB, stock.weight.shape[0] - 1]) ++ _assert_same_bytes(host(ids), stock(ids)) ++ ++ # Negative IDs and IDs beyond padded storage must fail in the same way. ++ for invalid_id in (-1, stock.weight.shape[0]): ++ invalid = torch.tensor([invalid_id]) ++ with pytest.raises(Exception) as stock_error: ++ stock(invalid) ++ with pytest.raises(type(stock_error.value)) as host_error: ++ host(invalid) ++ assert str(host_error.value) == str(stock_error.value) ++ ++ ++def test_host_budget_requires_measured_allocation(): ++ doc = PinnedHostVocabParallelEmbedding.__doc__ or "" ++ assert "measured pinned allocation, not tensor payload" in doc ++ ++ ++class _LanguageModel(nn.Module): ++ def __init__(self, embedding: nn.Module) -> None: ++ super().__init__() ++ self.embed_tokens = embedding ++ ++ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: ++ return self.embed_tokens(input_ids) ++ ++ ++def test_u1_real_multimodal_text_mask_and_merge(monkeypatch): ++ _set_tp(monkeypatch, 0, 1) ++ weight = _weight() ++ with torch.device("cpu"): ++ embedding = PinnedHostVocabParallelEmbedding( ++ VOCAB, ++ HIDDEN, ++ params_dtype=torch.bfloat16, ++ create_accelerator_view=False, ++ ) ++ embedding.weight_loader(embedding.weight, weight) ++ ++ model = model_mod.Qwen4ExpForConditionalGeneration.__new__( ++ model_mod.Qwen4ExpForConditionalGeneration ++ ) ++ nn.Module.__init__(model) ++ model.language_model = _LanguageModel(embedding) ++ model.language_model_only = False ++ model.use_deepstack = False ++ model._has_oov_mm_tokens = True ++ ++ ids = torch.tensor([1, VOCAB + 5, 2, 64], dtype=torch.int64) ++ is_multimodal = torch.tensor([False, True, False, False]) ++ image_row = torch.full((1, HIDDEN), 123, dtype=torch.bfloat16) ++ ++ actual = model.embed_input_ids( ++ ids, ++ multimodal_embeddings=[image_row], ++ is_multimodal=is_multimodal, ++ ) ++ ++ expected = torch.zeros(4, HIDDEN, dtype=torch.bfloat16) ++ expected[0].copy_(weight[1]) ++ expected[1].copy_(image_row[0]) ++ expected[2].copy_(weight[2]) ++ expected[3].copy_(weight[64]) ++ _assert_same_bytes(actual, expected) ++ ++ ++class _PPGroup: ++ is_first_rank = True ++ is_last_rank = False ++ ++ ++class _StubModule(nn.Module): ++ def __init__(self, *args, **kwargs) -> None: ++ super().__init__() ++ ++ ++def _tiny_config(*, tie_word_embeddings: bool = False): ++ text_config = SimpleNamespace( ++ vocab_size=VOCAB, ++ hidden_size=HIDDEN, ++ hc_count=1, ++ layer_types=[], ++ num_hidden_layers=0, ++ mtp_num_hidden_layers=0, ++ rms_norm_eps=1e-6, ++ hc_lowrank=1, ++ tie_word_embeddings=tie_word_embeddings, ++ ) ++ model_config = SimpleNamespace( ++ hf_text_config=text_config, ++ dtype=torch.bfloat16, ++ ) ++ parallel_config = SimpleNamespace( ++ eplb_config=SimpleNamespace(num_redundant_experts=0), ++ use_sequence_parallel_moe=False, ++ ) ++ scheduler_config = SimpleNamespace(max_num_batched_tokens=8) ++ # The @support_torch_compile wrapper reads vllm_config.compilation_config after ++ # __init__; mode NONE makes it return without compiling. ++ from vllm.config import CompilationMode ++ ++ compilation_config = SimpleNamespace(mode=CompilationMode.NONE) ++ return SimpleNamespace( ++ model_config=model_config, ++ parallel_config=parallel_config, ++ scheduler_config=scheduler_config, ++ compilation_config=compilation_config, ++ speculative_config=None, ++ quant_config=None, ++ ) ++ ++ ++def _patch_tiny_model_construction(monkeypatch) -> None: ++ _set_tp(monkeypatch, 0, 1) ++ monkeypatch.setattr(model_mod, "get_pp_group", lambda: _PPGroup()) ++ monkeypatch.setattr( ++ model_mod, ++ "make_layers", ++ lambda *args, **kwargs: (0, 0, nn.ModuleList()), ++ ) ++ monkeypatch.setattr( ++ model_mod, ++ "is_model_fused_shared_expert_compatible", ++ lambda *args, **kwargs: False, ++ ) ++ monkeypatch.setattr( ++ model_mod, ++ "make_empty_intermediate_tensors_factory", ++ lambda *args, **kwargs: lambda *a, **k: None, ++ ) ++ ++ monkeypatch.setattr( ++ mtp_mod, ++ "_make_draft_vllm_config", ++ lambda *args, **kwargs: SimpleNamespace(quant_config=None), ++ ) ++ monkeypatch.setattr( ++ mtp_mod, ++ "set_current_vllm_config", ++ lambda *args, **kwargs: nullcontext(), ++ ) ++ monkeypatch.setattr(mtp_mod, "ColumnParallelLinear", _StubModule) ++ monkeypatch.setattr(mtp_mod, "GemmaRMSNorm", _StubModule) ++ monkeypatch.setattr(mtp_mod, "GatedResidual", _StubModule) ++ monkeypatch.setattr( ++ mtp_mod, ++ "is_model_fused_shared_expert_compatible", ++ lambda *args, **kwargs: False, ++ ) ++ monkeypatch.setattr( ++ mtp_mod, ++ "make_empty_intermediate_tensors_factory", ++ lambda *args, **kwargs: lambda *a, **k: None, ++ ) ++ ++ ++def test_u2_ple_meta_construction_materializes_neither_table(monkeypatch): ++ monkeypatch.setenv(HOST_EMBED_TABLE_ENV, "1") ++ _patch_tiny_model_construction(monkeypatch) ++ config = _tiny_config() ++ ++ with torch.device("meta"): ++ target = model_mod.Qwen4ExpModel(vllm_config=config) ++ predictor = mtp_mod.Qwen4ExpMultiTokenPredictor(vllm_config=config) ++ ++ assert isinstance(target.embed_tokens, VocabParallelEmbedding) ++ assert not isinstance( ++ target.embed_tokens, ++ PinnedHostVocabParallelEmbedding, ++ ) ++ assert target.embed_tokens.weight.device.type == "meta" ++ ++ assert isinstance(predictor.embed_tokens, VocabParallelEmbedding) ++ assert not isinstance( ++ predictor.embed_tokens, ++ PinnedHostVocabParallelEmbedding, ++ ) ++ assert predictor.embed_tokens.weight.device.type == "meta" ++ ++ ++def test_tied_embeddings_are_rejected_for_both_owners(monkeypatch): ++ monkeypatch.setenv(HOST_EMBED_TABLE_ENV, "1") ++ _patch_tiny_model_construction(monkeypatch) ++ config = _tiny_config(tie_word_embeddings=True) ++ ++ with pytest.raises(ValueError, match="tie_word_embeddings=True"): ++ model_mod.Qwen4ExpModel(vllm_config=config) ++ with pytest.raises(ValueError, match="tie_word_embeddings=True"): ++ mtp_mod.Qwen4ExpMultiTokenPredictor(vllm_config=config) ++ ++ ++def _cuda_host_layer(monkeypatch, vocab=VOCAB, hidden=HIDDEN, world_size=1): ++ _require_cuda_pinned_memory() ++ _set_tp(monkeypatch, 0, world_size) ++ layer = PinnedHostVocabParallelEmbedding( ++ vocab, ++ hidden, ++ params_dtype=torch.bfloat16, ++ ) ++ weight = _weight(vocab, hidden) ++ layer.weight_loader(layer.weight, weight) ++ return layer, weight ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA") ++def test_u3_uva_lookup_is_graph_capturable_and_replayable(monkeypatch): ++ layer, weight = _cuda_host_layer(monkeypatch) ++ assert layer.weight.device.type == "cpu" ++ assert layer.weight.is_pinned() ++ assert layer.accelerator_weight_view is not None ++ assert layer.accelerator_weight_view.device.type == "cuda" ++ ++ ids = torch.tensor([0, 1, 64, 66], device="cuda") ++ layer(ids) ++ torch.cuda.synchronize() ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ captured_output = layer(ids) ++ ++ for replay_ids in ( ++ torch.tensor([2, 3, 4, 5]), ++ torch.tensor([64, 65, 66, 0]), ++ ): ++ ids.copy_(replay_ids.to(device="cuda")) ++ graph.replay() ++ torch.cuda.synchronize() ++ _assert_same_bytes( ++ captured_output.cpu(), ++ weight[replay_ids], ++ ) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA") ++def test_u3_three_lookup_steps_inside_one_graph(monkeypatch): ++ layer, weight = _cuda_host_layer(monkeypatch) ++ loop_ids = torch.tensor([1, 8, 32, 60], device="cuda") ++ ++ for _ in range(3): ++ layer(loop_ids) ++ torch.cuda.synchronize() ++ ++ graph = torch.cuda.CUDAGraph() ++ captured_steps = [] ++ with torch.cuda.graph(graph): ++ for _ in range(3): ++ captured_steps.append(layer(loop_ids)) ++ loop_ids.add_(1) ++ ++ for seed in ( ++ torch.tensor([1, 8, 32, 60]), ++ torch.tensor([4, 12, 40, 63]), ++ ): ++ loop_ids.copy_(seed.to(device="cuda")) ++ graph.replay() ++ torch.cuda.synchronize() ++ for step, output in enumerate(captured_steps): ++ _assert_same_bytes(output.cpu(), weight[seed + step]) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA") ++def test_u3_parent_device_move_preserves_storage_and_dtype_move_fails(monkeypatch): ++ layer, _ = _cuda_host_layer(monkeypatch) ++ view = layer.accelerator_weight_view ++ assert view is not None ++ weight_ptr = layer.weight.data_ptr() ++ weight_storage_ptr = layer.weight.untyped_storage().data_ptr() ++ view_ptr = view.data_ptr() ++ view_storage_ptr = view.untyped_storage().data_ptr() ++ ++ parent = nn.Module() ++ parent.add_module("embed_tokens", layer) ++ parent.add_module( ++ "projection", ++ nn.Linear(HIDDEN, HIDDEN, dtype=torch.bfloat16), ++ ) ++ ++ parent.to("cuda") ++ assert parent.projection.weight.device.type == "cuda" ++ assert layer.weight.device.type == "cpu" ++ assert layer.weight.is_pinned() ++ assert layer.weight.data_ptr() == weight_ptr ++ assert layer.weight.untyped_storage().data_ptr() == weight_storage_ptr ++ assert view.data_ptr() == view_ptr ++ assert view.untyped_storage().data_ptr() == view_storage_ptr ++ ++ with pytest.raises(RuntimeError, match="fixed dtype"): ++ parent.to(torch.float32) ++ with pytest.raises(RuntimeError, match="fixed dtype"): ++ parent.half() ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA") ++def test_u4_outer_post_load_preserves_pinned_uva_storage(monkeypatch): ++ vocab = 4096 ++ hidden = 256 ++ _set_tp(monkeypatch, 1, 2) ++ _require_cuda_pinned_memory() ++ ++ # Account for construction, loading, and the outer post-load call. ++ torch.cuda.synchronize() ++ baseline = torch.cuda.memory_allocated() ++ torch.cuda.reset_peak_memory_stats() ++ ++ layer = PinnedHostVocabParallelEmbedding( ++ vocab, ++ hidden, ++ params_dtype=torch.bfloat16, ++ ) ++ view = layer.accelerator_weight_view ++ assert view is not None ++ weight_ptr = layer.weight.data_ptr() ++ weight_storage_ptr = layer.weight.untyped_storage().data_ptr() ++ view_ptr = view.data_ptr() ++ view_storage_ptr = view.untyped_storage().data_ptr() ++ ++ loaded_weight = _weight(vocab, hidden) ++ local_table_bytes = layer.weight.numel() * layer.weight.element_size() ++ ++ assert ( ++ layer.weight_loader.__func__ ++ is VocabParallelEmbedding.weight_loader ++ ) ++ layer.weight_loader(layer.weight, loaded_weight) ++ ++ model = nn.Module() ++ model.add_module("embed_tokens", layer) ++ monkeypatch.setattr( ++ loader_utils, ++ "maybe_retie_word_embeddings", ++ lambda *args, **kwargs: None, ++ ) ++ loader_utils.process_weights_after_loading( ++ model, ++ SimpleNamespace(quantization=None), ++ torch.device("cuda"), ++ ) ++ ++ torch.cuda.synchronize() ++ peak_growth = torch.cuda.max_memory_allocated() - baseline ++ ++ assert layer.weight.device.type == "cpu" ++ assert layer.weight.is_pinned() ++ assert layer.weight.data_ptr() == weight_ptr ++ assert layer.weight.untyped_storage().data_ptr() == weight_storage_ptr ++ assert view.device.type == "cuda" ++ assert view.shape == layer.weight.shape ++ assert view.dtype == layer.weight.dtype ++ assert view.data_ptr() == view_ptr ++ assert view.untyped_storage().data_ptr() == view_storage_ptr ++ assert set(dict(layer.named_parameters())) == {"weight"} ++ assert peak_growth < local_table_bytes ++ ++ expected = loaded_weight.narrow(0, vocab // 2, vocab // 2) ++ _assert_same_bytes(layer.weight, expected) ++ ++ # Pointer stability alone is insufficient: mutate the retained CPU ++ # parameter and prove that the CUDA view reads that same storage. ++ local_row = 7 ++ global_id = vocab // 2 + local_row ++ sentinel = torch.arange(hidden, dtype=torch.float32).to(torch.bfloat16) ++ layer.weight[local_row].copy_(sentinel) ++ actual = layer(torch.tensor([global_id], device="cuda")) ++ torch.cuda.synchronize() ++ _assert_same_bytes(actual[0].cpu(), sentinel) ++ ++ ++@pytest.mark.skipif( ++ os.environ.get("VLLM_RUN_HOST_EMBED_BUDGET") != "1", ++ reason="set VLLM_RUN_HOST_EMBED_BUDGET=1 for the 606 MiB allocation report", ++) ++def test_u5_report_production_shard_host_allocation(monkeypatch): ++ _require_cuda_pinned_memory() ++ if not hasattr(torch.cuda, "host_memory_stats"): ++ pytest.skip("torch.cuda.host_memory_stats is unavailable") ++ ++ _set_tp(monkeypatch, 0, 2) ++ before = torch.cuda.host_memory_stats() ++ allocated_key = "allocated_bytes.all.current" ++ reserved_key = "reserved_bytes.all.current" ++ if allocated_key not in before or reserved_key not in before: ++ pytest.skip("host allocator byte statistics are unavailable") ++ ++ try: ++ layer = PinnedHostVocabParallelEmbedding( ++ 248_320, ++ 2_560, ++ params_dtype=torch.bfloat16, ++ ) ++ except RuntimeError as exc: ++ pytest.skip(f"production-sized pinned allocation is unavailable: {exc}") ++ ++ after = torch.cuda.host_memory_stats() ++ payload_bytes = layer.weight.numel() * layer.weight.element_size() ++ allocated_delta = after[allocated_key] - before[allocated_key] ++ reserved_delta = after[reserved_key] - before[reserved_key] ++ print( ++ "HOST_EMBED_PINNED_BUDGET " ++ f"payload_bytes={payload_bytes} " ++ f"allocated_bytes={allocated_delta} " ++ f"reserved_bytes={reserved_delta}" ++ ) ++ assert payload_bytes == 635_699_200 ++ assert layer.weight.is_pinned() +diff --git a/vllm/envs.py b/vllm/envs.py +index 6158f86..554c025 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -298,6 +298,12 @@ if TYPE_CHECKING: + VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 + VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" + VLLM_USE_V2_MODEL_RUNNER: bool | None = None ++ # Compile-cache factor (v2 host-embedding work): the pinned-host UVA embedding ++ # table changes the traced MTP-drafter graph (reads _uva_weight instead of ++ # embed_tokens.weight). Registered here so envs.compile_factors() keys the ++ # AOT cache on it; the layer itself still reads os.environ directly. ++ VLLM_HOST_EMBED_TABLE: bool = False ++ VLLM_PLE_OFFLOAD_READY_TIMEOUT: float = 600.0 + VLLM_LOG_MODEL_INSPECTION: bool = False + VLLM_DEBUG_MFU_METRICS: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False +@@ -2037,6 +2043,11 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool( + os.getenv("VLLM_USE_V2_MODEL_RUNNER", None) + ), ++ "VLLM_HOST_EMBED_TABLE": lambda: os.getenv("VLLM_HOST_EMBED_TABLE", "") == "1", ++ # Timeout for PLE weight loading and TP worker registration. ++ "VLLM_PLE_OFFLOAD_READY_TIMEOUT": lambda: float( ++ os.getenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "600") ++ ), + # Log model inspection after loading. + # If enabled, logs a transformers-style hierarchical view of the model + # with quantization methods and attention backends. +diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py +index 456a22a..01a4809 100644 +--- a/vllm/model_executor/layers/vocab_parallel_embedding.py ++++ b/vllm/model_executor/layers/vocab_parallel_embedding.py +@@ -1,6 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++import os + from collections.abc import Sequence + from dataclasses import dataclass + +@@ -31,8 +32,26 @@ from vllm.model_executor.layers.utils import dispatch_unquantized_gemm + from vllm.model_executor.parameter import BasevLLMParameter + from vllm.model_executor.utils import set_weight_attrs + from vllm.platforms import current_platform ++from vllm.utils.torch_utils import ( ++ PIN_MEMORY, ++ get_accelerator_view_from_cpu_tensor, ++) + + DEFAULT_VOCAB_PADDING_SIZE = 64 ++HOST_EMBED_TABLE_ENV = "VLLM_HOST_EMBED_TABLE" ++ ++ ++def host_embed_table_enabled() -> bool: ++ """Return whether the exact host-embedding opt-in is set.""" ++ return os.environ.get(HOST_EMBED_TABLE_ENV, "") == "1" ++ ++ ++def should_use_host_embed_table() -> bool: ++ """Keep the PLE offload worker's model construction on the meta device.""" ++ return ( ++ host_embed_table_enabled() ++ and torch.empty(0).device.type != "meta" ++ ) + + + class UnquantizedEmbeddingMethod(QuantizeMethodBase): +@@ -91,6 +110,46 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase): + return layer + + ++class _PinnedHostUnquantizedEmbeddingMethod(UnquantizedEmbeddingMethod): ++ """Allocate an unquantized embedding shard without a CUDA transient.""" ++ ++ requires_device_loading = False ++ ++ def __init__(self, *, device: torch.device, pin_memory: bool) -> None: ++ self.device = device ++ self.pin_memory = pin_memory ++ ++ def create_weights( ++ self, ++ layer: torch.nn.Module, ++ input_size_per_partition: int, ++ output_partition_sizes: list[int], ++ input_size: int, ++ output_size: int, ++ params_dtype: torch.dtype, ++ **extra_weight_attrs, ++ ) -> None: ++ del input_size, output_size ++ weight = Parameter( ++ torch.empty( ++ sum(output_partition_sizes), ++ input_size_per_partition, ++ dtype=params_dtype, ++ device=self.device, ++ pin_memory=self.pin_memory, ++ ), ++ requires_grad=False, ++ ) ++ set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) ++ layer.register_parameter("weight", weight) ++ set_weight_attrs(weight, extra_weight_attrs) ++ ++ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ++ # The serving layer must retain the same pinned allocation whose mapped ++ # address was captured. CPU-only reference tests need no GEMM dispatch. ++ del layer ++ ++ + def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int: + """Pad the vocab size to the given value.""" + return ((vocab_size + pad_to - 1) // pad_to) * pad_to +@@ -564,6 +623,124 @@ class VocabParallelEmbedding(PluggableLayer): + return s + + ++class PinnedHostVocabParallelEmbedding(VocabParallelEmbedding): ++ """TP-sharded embedding read on CUDA through mapped pinned host memory. ++ ++ Host provisioning must use measured pinned allocation, not tensor payload. ++ """ ++ ++ def __init__( ++ self, ++ num_embeddings: int, ++ embedding_dim: int, ++ params_dtype: torch.dtype | None = None, ++ org_num_embeddings: int | None = None, ++ padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, ++ prefix: str = "", ++ *, ++ disable_tp: bool = False, ++ create_accelerator_view: bool = True, ++ ) -> None: ++ construction_device = torch.empty(0).device ++ is_meta = construction_device.type == "meta" ++ if ( ++ create_accelerator_view ++ and not is_meta ++ and current_platform.is_cuda() ++ and not PIN_MEMORY ++ ): ++ raise RuntimeError( ++ "VLLM_HOST_EMBED_TABLE requires available pinned host memory" ++ ) ++ ++ storage_device = construction_device if is_meta else torch.device("cpu") ++ pin_memory = ( ++ not is_meta ++ and current_platform.is_cuda() ++ and PIN_MEMORY ++ ) ++ quant_method = _PinnedHostUnquantizedEmbeddingMethod( ++ device=storage_device, ++ pin_memory=pin_memory, ++ ) ++ super().__init__( ++ num_embeddings, ++ embedding_dim, ++ params_dtype=params_dtype, ++ org_num_embeddings=org_num_embeddings, ++ padding_size=padding_size, ++ prefix=prefix, ++ disable_tp=disable_tp, ++ quant_method=quant_method, ++ ) ++ ++ # The custom fused kernel is specified for ordinary device-resident ++ # shards, not device-mapped host allocations. Use the regular masked ++ # embedding path whose CUDA tensor argument is the mapped view. ++ self.use_fused_embedding = False ++ self._uva_weight: torch.Tensor | None = None ++ if ( ++ create_accelerator_view ++ and self.weight.device.type == "cpu" ++ and current_platform.is_cuda() ++ ): ++ if not self.weight.is_pinned(): ++ raise RuntimeError( ++ "VLLM_HOST_EMBED_TABLE weight was not allocated pinned" ++ ) ++ self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) ++ ++ @property ++ def accelerator_weight_view(self) -> torch.Tensor | None: ++ return self._uva_weight ++ ++ def _apply(self, fn, recurse: bool = True): ++ # Probe the transformation without applying it to the pinned shard. ++ # Device-only parent moves remain no-ops, but silently retaining this ++ # layer's dtype during a parent dtype conversion would create a ++ # mixed-dtype model. ++ del recurse ++ probe = torch.empty(0, dtype=self.weight.dtype, device=self.weight.device) ++ if fn(probe).dtype != self.weight.dtype: ++ raise RuntimeError("Pinned host embedding has a fixed dtype") ++ return self ++ ++ def _lookup_weight(self, input_: torch.Tensor) -> torch.Tensor: ++ if input_.device.type == "cpu": ++ return self.weight ++ if self._uva_weight is None: ++ raise RuntimeError( ++ "Host embedding received accelerator IDs without a UVA view" ++ ) ++ if self._uva_weight.device != input_.device: ++ raise RuntimeError( ++ "Host embedding UVA view and input IDs are on different devices" ++ ) ++ return self._uva_weight ++ ++ def forward(self, input_: torch.Tensor) -> torch.Tensor: ++ if self.tp_size == 1: ++ return F.embedding( ++ input_.long(), ++ self._lookup_weight(input_), ++ ) ++ ++ masked_input, input_mask = get_masked_input_and_mask( ++ input_, ++ self.shard_indices.org_vocab_start_index, ++ self.shard_indices.org_vocab_end_index, ++ self.shard_indices.num_org_vocab_padding, ++ self.shard_indices.added_vocab_start_index, ++ self.shard_indices.added_vocab_end_index, ++ ) ++ output_parallel = F.embedding( ++ masked_input.long(), ++ self._lookup_weight(input_), ++ ) ++ output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) ++ return tensor_model_parallel_all_reduce(output_parallel) ++ ++ + # --8<-- [start:parallel_lm_head] + @PluggableLayer.register("parallel_lm_head") + class ParallelLMHead(VocabParallelEmbedding): +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 7ac0ccb..269bbbd 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -28,8 +28,10 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( + ) + from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.model_executor.layers.vocab_parallel_embedding import ( ++ PinnedHostVocabParallelEmbedding, + ParallelLMHead, + VocabParallelEmbedding, ++ should_use_host_embed_table, + ) + from vllm.model_executor.models.interfaces import ( + HasInnerState, +@@ -501,7 +503,31 @@ class Qwen4ExpModel(nn.Module): + if layer_type == "full_attention" + and getattr(config, "indexer_n_heads", None) is not None + ) +- self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) ++ self.uses_host_embed_table = ( ++ get_pp_group().is_first_rank ++ and should_use_host_embed_table() ++ ) ++ if self.uses_host_embed_table and getattr( ++ config, "tie_word_embeddings", False ++ ): ++ raise ValueError( ++ "VLLM_HOST_EMBED_TABLE=1 is incompatible with " ++ "tie_word_embeddings=True" ++ ) ++ # PP ownership: the input embedding runs on ++ # the first stage only (see forward); later stages carried an unused ++ # 248320 x 2560 bf16 table (0.59 GiB per GPU at TP2) beside the MTP ++ # drafter's own copy. PP=1 ownership is unchanged, and the PLE offload ++ # worker retains meta construction. ++ if get_pp_group().is_first_rank: ++ embedding_cls = ( ++ PinnedHostVocabParallelEmbedding ++ if self.uses_host_embed_table ++ else VocabParallelEmbedding ++ ) ++ self.embed_tokens = embedding_cls(self.vocab_size, config.hidden_size) ++ else: ++ self.embed_tokens = PPMissingLayer() + + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: + layer_idx = extract_layer_index(prefix) +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 5dc0fc4..0ad96a2 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -28,8 +28,10 @@ from vllm.model_executor.layers.layernorm import GemmaRMSNorm + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.model_executor.layers.logits_processor import LogitsProcessor + from vllm.model_executor.layers.vocab_parallel_embedding import ( ++ PinnedHostVocabParallelEmbedding, + ParallelLMHead, + VocabParallelEmbedding, ++ should_use_host_embed_table, + ) + from vllm.model_executor.model_loader.utils import configure_quant_config + from vllm.model_executor.models.interfaces import SupportsPP +@@ -182,7 +184,20 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + self.hidden_size = config.hidden_size + self.hc_count = config.hc_count + +- self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) ++ self.uses_host_embed_table = should_use_host_embed_table() ++ if self.uses_host_embed_table and getattr( ++ config, "tie_word_embeddings", False ++ ): ++ raise ValueError( ++ "VLLM_HOST_EMBED_TABLE=1 is incompatible with " ++ "tie_word_embeddings=True" ++ ) ++ embedding_cls = ( ++ PinnedHostVocabParallelEmbedding ++ if self.uses_host_embed_table ++ else VocabParallelEmbedding ++ ) ++ self.embed_tokens = embedding_cls(self.vocab_size, self.hidden_size) + draft_vllm_config = _make_draft_vllm_config( + vllm_config, + self.mtp_start_layer_idx, +@@ -316,7 +331,6 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + hidden_states = hidden_states.flatten(-2) + prev_block_output = inputs_embeds + else: +- assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + current_step_idx = spec_step_idx % self.num_mtp_layers +-- +2.43.0 + diff --git a/patches/0007-worker-index-the-persistent-per-request-slot-block-t.patch b/patches/0007-worker-index-the-persistent-per-request-slot-block-t.patch new file mode 100644 index 0000000..2ebea69 --- /dev/null +++ b/patches/0007-worker-index-the-persistent-per-request-slot-block-t.patch @@ -0,0 +1,1608 @@ +From 638e92f28cc3076b6306228709ca96690429b718 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:12:35 -0400 +Subject: [PATCH 007/103] worker: index the persistent per-request-slot block + table under PP + +Under pipeline parallelism the Mamba/GDN state-copy path indexed a block-table +row gathered for a different step. The align post-process now indexes the +persistent per-request-slot table, and fails closed when it is absent. +--- + tests/v1/worker/gpu/__init__.py | 2 + + .../gpu/test_pp_deferred_mamba_block_table.py | 706 ++++++++++++++++++ + .../worker/gpu/test_pp_deferred_mamba_gpu.py | 542 ++++++++++++++ + tests/v1/worker/test_mamba_utils.py | 25 +- + vllm/v1/worker/gpu/model_runner.py | 4 + + vllm/v1/worker/gpu/model_states/interface.py | 8 + + .../worker/gpu/model_states/mamba_hybrid.py | 19 + + vllm/v1/worker/mamba_utils.py | 83 +- + 8 files changed, 1381 insertions(+), 8 deletions(-) + create mode 100644 tests/v1/worker/gpu/__init__.py + create mode 100644 tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py + create mode 100644 tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py + +diff --git a/tests/v1/worker/gpu/__init__.py b/tests/v1/worker/gpu/__init__.py +new file mode 100644 +index 0000000..208f01a +--- /dev/null ++++ b/tests/v1/worker/gpu/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +new file mode 100644 +index 0000000..30bed0a +--- /dev/null ++++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +@@ -0,0 +1,706 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Which block-table row the align post-process kernel reads. ++ ++Under pipeline parallelism the Mamba align post-process is *deferred*: it is ++enqueued by ``update_pp_decode_requests`` at the top of ``execute_model`` ++(``model_runner.py:1644``), before this step's ``gather_block_tables`` ++(``model_runner.py:1442`` -> ``block_table.py:149``), and the ``idx_mapping`` it ++carries was captured ``pp_size`` steps earlier on a ``PendingRecv`` ++(``pp_utils.py:57-58, 258-271``). There is exactly one gathered destination ++buffer, reused in place for the engine's lifetime (``block_table.py:74-76``), so ++at that moment it holds the batch of step T-1 while the mapping describes the ++batch of step T-2. With two-way PP micro-batch alternation those are different ++micro-batches by construction, so ``bt_row_idx = batch_idx`` ++(``mamba_utils.py:481``) resolved a *different request's* row. ++ ++These tests model the row selection on CPU: they drive the production plumbing ++(``MambaHybridModelState._ensure_align_ctx`` -> ++``MambaSpecDecodeGPUContext.initialize_from_forward_context`` -> ++``run_fused_postprocess_align``), capture the launch instead of running Triton, ++and then replay the kernel's index arithmetic against the two candidate tables. ++ ++No GPU is required. ++""" ++ ++import inspect ++from types import SimpleNamespace ++from unittest.mock import MagicMock ++ ++import numpy as np ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.mamba_utils import ( ++ MambaStateCopyFuncsByType, ++ get_conv_copy_spec, ++ get_temporal_copy_spec, ++) ++from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum ++from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec ++from vllm.v1.worker import mamba_utils as mamba_utils_mod ++from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState ++from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext ++ ++MAX_NUM_REQS = 8 ++MAX_NUM_BLOCKS = 6 ++MAMBA_BLOCK_SIZE = 16 ++NUM_PHYS_BLOCKS = 64 ++CONV_WIDTH = 4 ++CONV_INNER = 8 ++TEMPORAL_DIM = 16 ++DTYPE = torch.float16 ++ ++_COPY_FUNCS: MambaStateCopyFuncsByType = { ++ MambaAttentionBackendEnum.GDN_ATTN: (get_conv_copy_spec, get_temporal_copy_spec), ++} ++ ++# Captured before any monkeypatching so the replay binds against the real ++# kernel signature; a launch that stops matching it fails loudly here. ++_POSTPROCESS_SIG = inspect.signature( ++ mamba_utils_mod.postprocess_mamba_fused_kernel.fn ++) ++_PRECOPY_SIG = inspect.signature( ++ mamba_utils_mod.precopy_mamba_align_fused_kernel.fn ++) ++ ++ ++class _MockCpuGpuBuffer: ++ """CpuGpuBuffer stand-in that needs no pinned memory.""" ++ ++ def __init__(self, size: int, dtype: torch.dtype, device: torch.device): ++ self.cpu = torch.zeros(size, dtype=dtype, device="cpu") ++ self.gpu = torch.zeros(size, dtype=dtype, device=device) ++ self.np = self.cpu.numpy() ++ ++ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: ++ if n is None: ++ self.gpu.copy_(self.cpu) ++ return self.gpu ++ self.gpu[:n].copy_(self.cpu[:n]) ++ return self.gpu[:n] ++ ++ ++class _RecordingKernel: ++ """Stands in for a triton.jit kernel and records the launch.""" ++ ++ def __init__(self) -> None: ++ self.launches: list[tuple[tuple[int, ...], tuple, dict]] = [] ++ ++ def __getitem__(self, grid): ++ def launch(*args, **kwargs): ++ self.launches.append((grid, args, kwargs)) ++ ++ return launch ++ ++ ++def _ptr_to_int(value: int) -> int: ++ """Undo ``_reinterpret_u64_as_i64`` (mamba_utils.py:186-189).""" ++ return value + (1 << 64) if value < 0 else value ++ ++ ++def _gather_rows( ++ persistent: torch.Tensor, ++ num_blocks: list[int], ++ idx_mapping: list[int], ++ num_reqs_padded: int, ++ dst: torch.Tensor, ++) -> None: ++ """CPU port of ``_gather_block_tables_kernel`` (block_table.py:236-273). ++ ++ Row ``batch_idx`` of ``dst`` receives ``num_blocks[req_idx]`` entries of ++ persistent row ``req_idx``; rows in ``[num_reqs, num_reqs_padded)`` are ++ zeroed whole; rows at or beyond ``num_reqs_padded`` are left holding the ++ previous gather's content. ++ """ ++ num_reqs = len(idx_mapping) ++ for batch_idx in range(num_reqs_padded): ++ if batch_idx >= num_reqs: ++ dst[batch_idx].zero_() ++ continue ++ req_idx = idx_mapping[batch_idx] ++ nb = num_blocks[req_idx] ++ dst[batch_idx, :nb] = persistent[req_idx, :nb] ++ ++ ++def _make_kv_cache_config() -> KVCacheConfig: ++ spec = MambaSpec( ++ block_size=MAMBA_BLOCK_SIZE, ++ shapes=((CONV_WIDTH, CONV_INNER), (TEMPORAL_DIM,)), ++ dtypes=(DTYPE, DTYPE), ++ mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ++ mamba_cache_mode="align", ++ num_speculative_blocks=3, ++ ) ++ return KVCacheConfig( ++ num_blocks=NUM_PHYS_BLOCKS, ++ kv_cache_tensors=[], ++ kv_cache_groups=[KVCacheGroupSpec(["layer_0"], spec)], ++ ) ++ ++ ++def _make_forward_context() -> dict: ++ attention = MagicMock() ++ attention.kv_cache = [ ++ torch.zeros(NUM_PHYS_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE), ++ torch.zeros(NUM_PHYS_BLOCKS, TEMPORAL_DIM, dtype=DTYPE), ++ ] ++ return {"layer_0": attention} ++ ++ ++def _register_block_tables(state: MambaHybridModelState, block_tables) -> bool: ++ """Hand the runner's persistent ``BlockTables`` to the model state. ++ ++ Deliberately tolerant: on a tree without the fix the hook does not exist, and ++ the point of these tests is to report *which row the kernel reads*, not to ++ die on a missing attribute. ++ """ ++ setter = getattr(state, "set_block_tables", None) ++ if setter is None: ++ return False ++ setter(block_tables) ++ return True ++ ++ ++def _build_scenario( ++ device: torch.device = torch.device("cpu"), ++ slot_mapping: list[int] = (0, 1, 2), ++ t1_mapping: list[int] = (2,), ++ register: bool = True, ++): ++ """T-2 gathers ``slot_mapping``; T-1 re-gathers ``t1_mapping`` into the same ++ buffer; the deferred post-process then runs with T-2's mapping. ++ ++ ``slot_mapping`` is the batch-row -> req-slot mapping carried on the ++ PendingRecv. Making it non-identity is what separates "the kernel used the ++ persistent table" from "the kernel used req_idx as the row": with an ++ identity mapping ``batch_idx == req_idx`` and the two are indistinguishable. ++ """ ++ kv_cache_config = _make_kv_cache_config() ++ forward_context = _make_forward_context() ++ ++ # Persistent per-req-slot table (block_table.py:53-59). Row r owns block ++ # ids 10*(r+1) + col, so a row is identifiable from any single entry. ++ persistent = torch.zeros( ++ MAX_NUM_REQS, MAX_NUM_BLOCKS, dtype=torch.int32, device=device ++ ) ++ for r in range(MAX_NUM_REQS): ++ for c in range(MAX_NUM_BLOCKS): ++ persistent[r, c] = 10 * (r + 1) + c ++ num_blocks = [3] * MAX_NUM_REQS ++ ++ # The single gathered destination buffer (block_table.py:74-76). ++ slot_mapping = list(slot_mapping) ++ t1_mapping = list(t1_mapping) ++ gathered = torch.zeros_like(persistent) ++ _gather_rows(persistent, num_blocks, slot_mapping, len(slot_mapping) + 1, gathered) ++ _gather_rows(persistent, num_blocks, t1_mapping, len(t1_mapping) + 1, gathered) ++ ++ ctx = MambaSpecDecodeGPUContext.create( ++ max_num_reqs=MAX_NUM_REQS, ++ kv_cache_config=kv_cache_config, ++ copy_funcs=_COPY_FUNCS, ++ device=device, ++ make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), ++ ) ++ ++ state = object.__new__(MambaHybridModelState) ++ state.device = device ++ state.max_num_reqs = MAX_NUM_REQS ++ state._align_mode = True ++ state._mamba_ctx = ctx ++ state._mamba_group_ids = [] ++ state._mamba_spec = None ++ state._mamba_state_copy_funcs = _COPY_FUNCS ++ # object.__new__ skips __init__, so mirror the attribute it would set. ++ state._runner_block_tables = None ++ state.vllm_config = SimpleNamespace( ++ compilation_config=SimpleNamespace(static_forward_context=forward_context) ++ ) ++ registered = ( ++ _register_block_tables( ++ state, ++ SimpleNamespace(block_tables=[SimpleNamespace(gpu=persistent)]), ++ ) ++ if register ++ else False ++ ) ++ ++ # What the runner hands prepare_attn/preprocess_state: the T-1 gather view. ++ gathered_view = gathered[: len(t1_mapping) + 1] ++ state._ensure_align_ctx(kv_cache_config, [0], (gathered_view,)) ++ ++ tables_by_ptr = { ++ persistent.data_ptr(): ("persistent", persistent), ++ gathered.data_ptr(): ("gathered", gathered), ++ } ++ return SimpleNamespace( ++ ctx=ctx, ++ state=state, ++ slot_mapping=slot_mapping, ++ persistent=persistent, ++ gathered=gathered, ++ tables_by_ptr=tables_by_ptr, ++ registered=registered, ++ ) ++ ++ ++def _replay_row(scenario, launch, batch_idx: int) -> tuple[str, int, dict]: ++ """Replay the kernel's row selection (mamba_utils.py:437-490) for one row. ++ ++ Returns ``(table_name, bt_row_idx, predicate)``. ++ """ ++ grid, args, kwargs = launch ++ bound = _POSTPROCESS_SIG.bind(*args, **kwargs) ++ bound.apply_defaults() ++ a = bound.arguments ++ ++ num_reqs = a["num_reqs"] ++ assert grid[0] == num_reqs, (grid, num_reqs) ++ assert batch_idx < num_reqs ++ ++ has_idx_mapping = bool(a["HAS_IDX_MAPPING"]) ++ # Absent on a tree without the fix; the kernel then always uses batch_idx. ++ bt_row_is_req = bool(a.get("BT_ROW_IS_REQ", False)) ++ ++ if has_idx_mapping: ++ req_idx = int(a["idx_mapping_ptr"][batch_idx]) ++ if req_idx < 0: ++ return ("skipped", -1, {}) ++ else: ++ req_idx = batch_idx ++ ++ num_accepted = int(a["num_accepted_tokens_ptr"][req_idx]) ++ src_block_idx = int(a["mamba_state_idx_ptr"][req_idx]) ++ assert bool(a["PRECOMPUTED_NEW_COMPUTED"]), "V2 align path only" ++ new_num_computed = int(a["num_computed_tokens_ptr"][req_idx]) ++ block_size = int(a["block_size"]) ++ ++ num_tokens_running_state = new_num_computed - num_accepted + 1 ++ aligned_new_computed = (new_num_computed // block_size) * block_size ++ needs_copy = aligned_new_computed >= num_tokens_running_state ++ accept_token_bias = aligned_new_computed - num_tokens_running_state ++ dest_block_idx = aligned_new_computed // block_size - 1 ++ predicate = { ++ "needs_copy": needs_copy, ++ "num_accepted": num_accepted, ++ "new_num_computed_mod_block": new_num_computed % block_size, ++ "src_col": src_block_idx, ++ "dst_col": dest_block_idx, ++ "token_bias": accept_token_bias, ++ "src_eq_dst": src_block_idx == dest_block_idx, ++ } ++ ++ if bt_row_is_req: ++ bt_row_idx = req_idx ++ else: ++ bt_row_idx = batch_idx if has_idx_mapping else req_idx ++ ++ group_idx = int(a["state_group_indices_ptr"][0]) ++ ptr = _ptr_to_int(int(a["block_table_ptrs_ptr"][group_idx])) ++ name, _table = scenario.tables_by_ptr[ptr] ++ return (name, bt_row_idx, predicate) ++ ++ ++def _block_ids(scenario, launch, batch_idx: int) -> dict: ++ name, bt_row_idx, predicate = _replay_row(scenario, launch, batch_idx) ++ table = next( ++ tbl for label, tbl in scenario.tables_by_ptr.values() if label == name ++ ) ++ return { ++ "table": name, ++ "row": bt_row_idx, ++ "src_block_id": int(table[bt_row_idx, predicate["src_col"]]), ++ "dst_block_id": int(table[bt_row_idx, predicate["dst_col"]]), ++ "temporal_src_block_id": int( ++ table[bt_row_idx, predicate["src_col"] + predicate["token_bias"]] ++ ), ++ **predicate, ++ } ++ ++ ++def _run_deferred(scenario, monkeypatch: pytest.MonkeyPatch) -> _RecordingKernel: ++ recorder = _RecordingKernel() ++ monkeypatch.setattr( ++ mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder ++ ) ++ device = scenario.persistent.device ++ ++ def t(values): ++ return torch.tensor(values, dtype=torch.int32, device=device) ++ ++ # Per-req-SLOT decision arrays (mamba_utils.py:439-458 reads them at ++ # req_idx, never at batch_idx), so they are keyed by the slot mapping. ++ copier, *bystanders = scenario.slot_mapping ++ num_accepted = torch.ones(MAX_NUM_REQS, dtype=torch.int32, device=device) ++ state_idx = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) ++ new_num_computed = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) ++ # Only the first batch row crosses a block boundary inside its accepted ++ # window: 32 % 16 = 0 < num_accepted 2. ++ num_accepted[copier] = 2 ++ new_num_computed[copier] = 32 ++ for slot in bystanders: ++ num_accepted[slot] = 2 ++ new_num_computed[slot] = 40 # 40 % 16 = 8 >= 2 -> no copy ++ ++ scenario.ctx.run_fused_postprocess_align( ++ num_reqs=len(scenario.slot_mapping), ++ num_accepted_tokens_gpu=num_accepted, ++ state_idx_gpu=state_idx, ++ new_num_computed_tokens_gpu=new_num_computed, ++ idx_mapping=t(scenario.slot_mapping), # captured at T-2 ++ ) ++ return recorder ++ ++ ++def test_gather_port_matches_production_semantics() -> None: ++ """Instrument check: the CPU gather port must reproduce block_table.py.""" ++ persistent = torch.arange( ++ MAX_NUM_REQS * MAX_NUM_BLOCKS, dtype=torch.int32 ++ ).reshape(MAX_NUM_REQS, MAX_NUM_BLOCKS) ++ dst = torch.full_like(persistent, -7) ++ num_blocks = [2] * MAX_NUM_REQS ++ ++ _gather_rows(persistent, num_blocks, [3, 1], 3, dst) ++ ++ # Gathered rows carry only the first num_blocks entries... ++ assert dst[0, :2].tolist() == persistent[3, :2].tolist() ++ assert dst[1, :2].tolist() == persistent[1, :2].tolist() ++ # ...the tail of a gathered row is untouched (block_table.py:269-273), ++ assert dst[0, 2:].tolist() == [-7] * (MAX_NUM_BLOCKS - 2) ++ # ...padded rows in [num_reqs, num_reqs_padded) are zeroed whole, ++ assert dst[2].tolist() == [0] * MAX_NUM_BLOCKS ++ # ...and rows at or beyond num_reqs_padded keep the previous content. ++ assert dst[3].tolist() == [-7] * MAX_NUM_BLOCKS ++ ++ ++def test_deferred_postprocess_resolves_the_requests_own_row( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """r0's deferred copy must land on r0's blocks, not on r2's. ++ ++ T-2 gathered [r0, r1, r2]; T-1 re-gathered [r2] into the same buffer, so ++ gathered row 0 now holds r2. The deferred launch still carries T-2's ++ mapping, so ``batch_idx = 0`` means r0. ++ """ ++ scenario = _build_scenario() ++ recorder = _run_deferred(scenario, monkeypatch) ++ assert len(recorder.launches) == 1 ++ launch = recorder.launches[0] ++ ++ row0 = _block_ids(scenario, launch, 0) ++ why = ( ++ f"resolved={row0} " ++ f"runner_hook_registered={scenario.registered} " ++ f"(a missing runner hook leaves the kernel on the batch-order gather)" ++ ) ++ ++ # The copy predicate really fires for this row, so the kernel does ++ # dereference the block table here (mamba_utils.py:462-479). ++ assert row0["needs_copy"] is True, why ++ assert row0["src_eq_dst"] is False, why ++ assert (row0["src_col"], row0["dst_col"], row0["token_bias"]) == (0, 1, 1), why ++ ++ # The gathered buffer holds r2 at batch row 0 after the T-1 gather, so a ++ # kernel that still indexes it by batch_idx corrupts r2 on r0's behalf. ++ assert scenario.gathered[0, 0].item() == 30 ++ ++ # r0's persistent row is 10, 11, 12, ... ++ assert row0["table"] == "persistent", why ++ assert row0["row"] == 0, why ++ assert row0["src_block_id"] == 10, why ++ assert row0["dst_block_id"] == 11, why ++ assert row0["temporal_src_block_id"] == 11, why ++ assert scenario.registered, why ++ ++ ++def test_deferred_postprocess_skips_rows_without_a_boundary_crossing( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """The other two rows do not reach the block table at all.""" ++ scenario = _build_scenario() ++ recorder = _run_deferred(scenario, monkeypatch) ++ launch = recorder.launches[0] ++ ++ for batch_idx in (1, 2): ++ _name, _row, predicate = _replay_row(scenario, launch, batch_idx) ++ assert predicate["needs_copy"] is False, (batch_idx, predicate) ++ ++ ++def test_precopy_still_uses_the_same_step_gather( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """The pre-copy site must keep batch-order rows. ++ ++ ``precopy_mamba_align_fused_kernel`` (mamba_utils.py:613-615) runs from ++ ``preprocess_state``, five lines after the gather that feeds it ++ (``model_runner.py:1718`` -> ``:1723``), so batch row i *is* this request. ++ Switching it to req rows would break it. ++ """ ++ scenario = _build_scenario() ++ recorder = _RecordingKernel() ++ monkeypatch.setattr( ++ mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder ++ ) ++ device = scenario.persistent.device ++ z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) ++ ++ scenario.ctx.run_fused_precopy( ++ num_reqs=3, ++ state_idx_gpu=z.clone(), ++ src_col_gpu=z.clone(), ++ token_bias_gpu=z.clone(), ++ idx_mapping=torch.tensor([0, 1, 2], dtype=torch.int32, device=device), ++ ) ++ ++ assert len(recorder.launches) == 1 ++ _grid, args, kwargs = recorder.launches[0] ++ bound = _PRECOPY_SIG.bind(*args, **kwargs) ++ bound.apply_defaults() ++ ptr = _ptr_to_int(int(bound.arguments["block_table_ptrs_ptr"][0])) ++ assert scenario.tables_by_ptr[ptr][0] == "gathered" ++ assert bound.arguments["HAS_IDX_MAPPING"] is True ++ assert "BT_ROW_IS_REQ" not in bound.arguments ++ ++ ++def test_deferred_postprocess_uses_req_rows_under_a_non_identity_mapping( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """Batch row 0 -> request slot 5, so batch_idx and req_idx disagree. ++ ++ The identity-mapping case above pins *which table*; this one pins *which ++ row*. Every wrong combination lands on a different set of block ids: ++ ++ ========================== ========== ================== ++ table / row resolves block ids ++ ========================== ========== ================== ++ persistent[req_idx=5] correct 60, 61 ++ persistent[batch_idx=0] wrong row 10, 11 ++ gathered[batch_idx=0] old 30, 31 (r2, from T-1) ++ gathered[req_idx=5] wrong both 0, 0 (never written) ++ ========================== ========== ================== ++ """ ++ scenario = _build_scenario(slot_mapping=[5, 1, 2], t1_mapping=[2]) ++ recorder = _run_deferred(scenario, monkeypatch) ++ assert len(recorder.launches) == 1 ++ launch = recorder.launches[0] ++ ++ _grid, args, kwargs = launch ++ bound = _POSTPROCESS_SIG.bind(*args, **kwargs) ++ bound.apply_defaults() ++ # The constexpr itself, not just its effect. ++ assert bound.arguments["BT_ROW_IS_REQ"] is True ++ assert bound.arguments["HAS_IDX_MAPPING"] is True ++ ++ row0 = _block_ids(scenario, launch, 0) ++ why = f"resolved={row0} runner_hook_registered={scenario.registered}" ++ ++ assert row0["needs_copy"] is True, why ++ assert row0["src_eq_dst"] is False, why ++ assert (row0["src_col"], row0["dst_col"], row0["token_bias"]) == (0, 1, 1), why ++ ++ # The row index is the request slot, not the batch row. ++ assert row0["table"] == "persistent", why ++ assert row0["row"] == 5, why ++ assert row0["src_block_id"] == 60, why ++ assert row0["dst_block_id"] == 61, why ++ assert row0["temporal_src_block_id"] == 61, why ++ ++ # None of the three wrong coordinates. ++ assert scenario.persistent[0, 0].item() == 10 # persistent, batch row ++ assert scenario.gathered[0, 0].item() == 30 # gathered, batch row ++ assert scenario.gathered[5, 0].item() == 0 # gathered, req row: never written ++ assert row0["src_block_id"] not in (10, 30, 0), why ++ ++ ++def test_align_postprocess_fails_closed_without_persistent_tables( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """No silent fallback to the batch-order gather. ++ ++ A context that never received the persistent tables is a caller bug: the ++ only other thing the kernel could index is the stale gather, which is the ++ corruption this change exists to remove. It must raise, and it must raise ++ before mutating the accepted-token snapshot. ++ """ ++ scenario = _build_scenario(slot_mapping=[5, 1, 2], t1_mapping=[2], register=False) ++ assert scenario.registered is False ++ assert scenario.ctx.has_persistent_block_tables is False ++ ++ recorder = _RecordingKernel() ++ monkeypatch.setattr( ++ mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder ++ ) ++ device = scenario.persistent.device ++ before = scenario.ctx.num_accepted_tokens_out.clone() ++ num_accepted = torch.full( ++ (MAX_NUM_REQS,), 7, dtype=torch.int32, device=device ++ ) ++ ++ with pytest.raises(RuntimeError, match="persistent per-req-slot block tables"): ++ scenario.ctx.run_fused_postprocess_align( ++ num_reqs=3, ++ num_accepted_tokens_gpu=num_accepted, ++ state_idx_gpu=torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device), ++ new_num_computed_tokens_gpu=torch.full( ++ (MAX_NUM_REQS,), 32, dtype=torch.int32, device=device ++ ), ++ idx_mapping=torch.tensor([5, 1, 2], dtype=torch.int32, device=device), ++ ) ++ ++ assert recorder.launches == [] ++ torch.testing.assert_close(scenario.ctx.num_accepted_tokens_out, before) ++ ++ ++def test_precopy_still_works_without_persistent_tables( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """The V1 pre-copy contract is untouched by the fail-closed rule. ++ ++ Only the V2 align post-process needs per-req-slot rows; the pre-copy reads ++ the same-step gather and must keep launching on a context that never ++ received the persistent tables. ++ """ ++ scenario = _build_scenario(register=False) ++ assert scenario.ctx.has_persistent_block_tables is False ++ recorder = _RecordingKernel() ++ monkeypatch.setattr( ++ mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder ++ ) ++ device = scenario.persistent.device ++ z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) ++ ++ scenario.ctx.run_fused_precopy( ++ num_reqs=1, ++ state_idx_gpu=z.clone(), ++ src_col_gpu=torch.full_like(z, -1), ++ token_bias_gpu=z.clone(), ++ idx_mapping=None, ++ ) ++ ++ assert len(recorder.launches) == 1 ++ _grid, args, kwargs = recorder.launches[0] ++ bound = _PRECOPY_SIG.bind(*args, **kwargs) ++ bound.apply_defaults() ++ ptr = _ptr_to_int(int(bound.arguments["block_table_ptrs_ptr"][0])) ++ assert scenario.tables_by_ptr[ptr][0] == "gathered" ++ ++ ++# ---------------------------------------------------- generation masking (CPU) ++ ++ ++def mask_from_generation_change( ++ idx_mapping: list[int], ++ reused_slot_rows: list[int], ++) -> list[int]: ++ """Run the real generation masking and return the mapping it produces. ++ ++ Drives ``PPHandler.get_prev_sampled_outputs`` (pp_utils.py:132-140) rather ++ than hand-writing a -1: a slot whose generation counter moved between ++ ``receive`` and consumption has been freed and handed to a new request, and ++ the handler must replace its row with -1 so the kernel skips it ++ (mamba_utils.py:441-442). ++ ++ Exported so the GPU suite can consume a mapping that a real reuse produced. ++ """ ++ from collections import deque ++ ++ from vllm.v1.worker.gpu import pp_utils as pp_utils_mod ++ ++ num_reqs = len(idx_mapping) ++ idx_np = np.asarray(idx_mapping, dtype=np.intp) ++ ++ handler = object.__new__(pp_utils_mod.PPHandler) ++ handler.device = torch.device("cpu") ++ handler.main_stream = SimpleNamespace(wait_event=lambda event: None) ++ # Slot generation counters as they stand now. ++ handler.req_idx_gen_np = np.zeros(max(int(idx_np.max()) + 1, 1), dtype=np.int32) ++ # Snapshot taken at receive time; a reused slot has since been bumped. ++ gen_at_receive = handler.req_idx_gen_np[idx_np].copy() ++ for row in reused_slot_rows: ++ handler.req_idx_gen_np[idx_np[row]] += 1 ++ ++ slot = pp_utils_mod.PendingRecv( ++ event=object(), ++ sampled_tokens=torch.zeros(num_reqs, 1, dtype=torch.int64), ++ num_sampled=torch.ones(num_reqs, dtype=torch.int32), ++ num_rejected=torch.zeros(num_reqs, dtype=torch.int32), ++ idx_mapping=torch.as_tensor(idx_np, dtype=torch.int32), ++ idx_mapping_np=idx_np, ++ need_sampled_mask=np.ones(num_reqs, dtype=bool), ++ gen_at_receive_np=gen_at_receive, ++ ) ++ handler.queue = deque([slot]) ++ ++ captured: dict[str, object] = {} ++ ++ def fake_async_copy_to_gpu(array, device=None): ++ captured["np"] = np.asarray(array).copy() ++ return torch.as_tensor(np.asarray(array), dtype=torch.int32) ++ ++ original = pp_utils_mod.async_copy_to_gpu ++ pp_utils_mod.async_copy_to_gpu = fake_async_copy_to_gpu ++ try: ++ outputs = handler.get_prev_sampled_outputs() ++ finally: ++ pp_utils_mod.async_copy_to_gpu = original ++ ++ assert outputs is not None, "the whole batch was excluded" ++ return [int(v) for v in outputs["idx_mapping"].tolist()] ++ ++ ++def test_generation_change_masks_only_the_reused_row(): ++ """A slot handed to a new request between capture and use becomes -1.""" ++ assert mask_from_generation_change([5, 1, 2], [1]) == [5, -1, 2] ++ assert mask_from_generation_change([5, 1, 2], [0, 2]) == [-1, 1, -1] ++ # No reuse: the handler returns the original tensor untouched. ++ assert mask_from_generation_change([5, 1, 2], []) == [5, 1, 2] ++ ++ ++def test_kernel_rejects_negative_req_idx_before_any_load_that_uses_it(): ++ """Structural check: the guard precedes every load indexed by req_idx. ++ ++ ``mamba_utils.py:441-442`` returns on ``req_idx < 0``; the decision loads at ++ ``:445-446`` index by ``req_idx``. If a future edit reordered them, a masked ++ row would dereference index -1 before being rejected. The GPU suite plants ++ poison at index -1 to catch that at runtime; this catches it on CPU, in the ++ source, with no device. ++ """ ++ import ast ++ import textwrap ++ ++ source = textwrap.dedent( ++ inspect.getsource(mamba_utils_mod.postprocess_mamba_fused_kernel.fn) ++ ) ++ tree = ast.parse(source) ++ func = tree.body[0] ++ ++ guard_lines = [ ++ node.lineno ++ for node in ast.walk(func) ++ if isinstance(node, ast.If) ++ and isinstance(node.test, ast.Compare) ++ and isinstance(node.test.left, ast.Name) ++ and node.test.left.id == "req_idx" ++ and any(isinstance(op, ast.Lt) for op in node.test.ops) ++ and any(isinstance(stmt, ast.Return) for stmt in node.body) ++ ] ++ assert guard_lines, "no `if req_idx < 0: return` guard found in the kernel" ++ guard_line = min(guard_lines) ++ ++ # Every expression that adds req_idx to a pointer, i.e. every load or store ++ # that would dereference a negative index. ++ uses = [ ++ node.lineno ++ for node in ast.walk(func) ++ if isinstance(node, ast.BinOp) ++ and isinstance(node.op, ast.Add) ++ and isinstance(node.right, ast.Name) ++ and node.right.id == "req_idx" ++ ] ++ assert uses, "no pointer arithmetic on req_idx found; has the kernel moved?" ++ assert min(uses) > guard_line, ( ++ f"a req_idx-indexed access at line {min(uses)} precedes the " ++ f"`req_idx < 0` guard at line {guard_line} (offsets are relative to the " ++ f"kernel body)" ++ ) +diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py +new file mode 100644 +index 0000000..f036227 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py +@@ -0,0 +1,542 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""GPU correctness for the deferred-PP align post-process row selection. ++ ++The CPU model in ``test_pp_deferred_mamba_block_table.py`` pins which pointer ++and which row the launcher passes; it replaces the Triton body with a recorder, ++so it proves nothing about the bytes the kernel moves. These tests launch the ++real kernel and check the destination state, including that the request the ++first-cut coordinate would have hit is left alone. ++ ++Scenario, shared by every case and identical to the CPU model's: ++ ++ step T-2 gather [r5, r1, r2] into the shared destination buffer ++ step T-1 gather [r2] into the same buffer ++ step T the deferred post-process runs with T-2's mapping ++ ++so batch row 0 means request slot 5, while buffer row 0 holds r2. The four ++possible coordinates resolve to four different physical blocks: ++ ++ persistent[req_idx=5] correct ++ persistent[batch_idx=0] wrong row -> r0's blocks ++ gathered[batch_idx=0] first cut -> r2's blocks (the victim) ++ gathered[req_idx=5] wrong both -> never written, block 0 ++ ++Requires CUDA **and** an explicit opt-in: set ``VLLM_E2_GPU_TESTS=1``. The reference host ++always reports CUDA available, and the reference host spends most of its time running ++long serving runs and sanitizer runs that should not share a device with a ++test suite, so `torch.cuda.is_available()` is not a sufficient gate here. Run ++this when the GPU is otherwise idle: ++ ++ export VLLM_E2_GPU_TESTS=1 ++ export PYTHONPATH=: ++ python -m pytest -q tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py ++""" ++ ++import os ++from dataclasses import dataclass ++from unittest.mock import MagicMock ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.mamba_utils import ( ++ MambaStateCopyFuncsByType, ++ get_conv_copy_spec, ++ get_temporal_copy_spec, ++ is_conv_state_dim_first, ++) ++from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum ++from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec ++from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext ++ ++pytestmark = [ ++ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), ++ pytest.mark.skipif( ++ os.environ.get("VLLM_E2_GPU_TESTS") != "1", ++ reason="set VLLM_E2_GPU_TESTS=1 to run GPU tests (shared GPU host)", ++ ), ++] ++ ++MAX_NUM_REQS = 8 ++BLOCKS_PER_REQ = 6 ++NUM_PHYS_BLOCKS = MAX_NUM_REQS * BLOCKS_PER_REQ ++# Two extra physical blocks, reachable only through the block-table guard row ++# that a negative req_idx would address: CANARY_DST sits at the poison ++# destination column and CANARY_SRC everywhere else, so a negative-index load ++# would produce a visible copy between them. See _build()'s backing tensors. ++# Block ids are 1-based (see phys()), so the canaries sit above the last real ++# id, NUM_PHYS_BLOCKS, rather than on it. ++CANARY_DST = NUM_PHYS_BLOCKS + 1 ++CANARY_SRC = NUM_PHYS_BLOCKS + 2 ++TOTAL_BLOCKS = NUM_PHYS_BLOCKS + 3 ++# Decision values planted at index -1 of every per-req array. They describe a ++# copy (new=48, a=4 -> aligned=48, running=45, bias=3, dst_col=2, src_col=0), ++# so a missing `req_idx < 0` guard is a write, not a silent no-op. ++POISON = (4, 0, 48) # (num_accepted, state_idx, new_num_computed) ++POISON_DST_COL = 2 ++MAMBA_BLOCK_SIZE = 16 ++NUM_SPEC_BLOCKS = 3 ++CONV_WIDTH = 4 ++CONV_INNER = 64 ++TEMPORAL_DIM = 128 ++DTYPE = torch.float16 ++ ++_COPY_FUNCS: MambaStateCopyFuncsByType = { ++ MambaAttentionBackendEnum.GDN_ATTN: (get_conv_copy_spec, get_temporal_copy_spec), ++} ++ ++ ++def phys(req_slot: int, col: int) -> int: ++ """Physical block id owned by ``req_slot`` at logical column ``col``. ++ ++ Disjoint per request, so a stray write is attributable. ++ ++ ONE-BASED, and that is load-bearing rather than cosmetic. Block id 0 is ++ NULL_BLOCK_ID (v1/attention/backends/utils.py:46); the block pool pops it ++ out of the free queue at construction "to represent a placeholder block ++ with block_id=0" (v1/core/block_pool.py:188-191), so no request ever owns ++ it. A fixture that hands out id 0 as a real block is asking the kernel to ++ copy from the null block, and the bounds guards -- correctly, and following ++ upstream #50021 -- refuse. Numbering from 0 made this fixture the one place ++ in the tree where a legitimate block id was 0. ++ """ ++ return req_slot * BLOCKS_PER_REQ + col + 1 ++ ++ ++class _MockCpuGpuBuffer: ++ def __init__(self, size: int, dtype: torch.dtype, device: torch.device): ++ self.cpu = torch.zeros(size, dtype=dtype, device="cpu") ++ self.gpu = torch.zeros(size, dtype=dtype, device=device) ++ self.np = self.cpu.numpy() ++ ++ def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: ++ if n is None: ++ self.gpu.copy_(self.cpu) ++ return self.gpu ++ self.gpu[:n].copy_(self.cpu[:n]) ++ return self.gpu[:n] ++ ++ ++def _make_kv_cache_config() -> KVCacheConfig: ++ spec = MambaSpec( ++ block_size=MAMBA_BLOCK_SIZE, ++ shapes=((CONV_WIDTH, CONV_INNER), (TEMPORAL_DIM,)), ++ dtypes=(DTYPE, DTYPE), ++ mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ++ mamba_cache_mode="align", ++ num_speculative_blocks=NUM_SPEC_BLOCKS, ++ ) ++ return KVCacheConfig( ++ num_blocks=TOTAL_BLOCKS, ++ kv_cache_tensors=[], ++ kv_cache_groups=[KVCacheGroupSpec(["layer_0"], spec)], ++ ) ++ ++ ++def _gather_rows(persistent, num_blocks, idx_mapping, num_reqs_padded, dst) -> None: ++ """Host port of ``_gather_block_tables_kernel`` (block_table.py:236-273).""" ++ num_reqs = len(idx_mapping) ++ for batch_idx in range(num_reqs_padded): ++ if batch_idx >= num_reqs: ++ dst[batch_idx].zero_() ++ continue ++ req_idx = idx_mapping[batch_idx] ++ nb = num_blocks[req_idx] ++ dst[batch_idx, :nb] = persistent[req_idx, :nb] ++ ++ ++@dataclass ++class _Fixture: ++ device: torch.device ++ ctx: MambaSpecDecodeGPUContext ++ conv: torch.Tensor ++ temporal: torch.Tensor ++ conv0: torch.Tensor ++ temporal0: torch.Tensor ++ persistent: torch.Tensor ++ gathered: torch.Tensor ++ ++ ++def _build(slot_mapping=(5, 1, 2), t1_mapping=(2,)) -> _Fixture: ++ device = torch.device("cuda:0") ++ torch.manual_seed(1234) ++ ++ conv = torch.randn( ++ TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device ++ ) ++ temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) ++ attention = MagicMock() ++ attention.kv_cache = [conv, temporal] ++ ++ # The block tables and the per-req decision arrays each get one extra ++ # leading row/element that the kernel is never supposed to reach. Because ++ # the real tensors are views starting one row/element in, index -1 in ++ # pointer arithmetic lands on that guard row -- so a kernel that failed to ++ # reject req_idx < 0 would load the poison decisions below and write into ++ # canaries. See test_freed_slot_is_rejected_before_any_negative_index_load. ++ persistent_backing = torch.zeros( ++ MAX_NUM_REQS + 1, BLOCKS_PER_REQ, dtype=torch.int32, device=device ++ ) ++ persistent_backing[0, :] = CANARY_SRC ++ persistent_backing[0, POISON_DST_COL] = CANARY_DST ++ persistent = persistent_backing[1:] ++ for r in range(MAX_NUM_REQS): ++ for c in range(BLOCKS_PER_REQ): ++ persistent[r, c] = phys(r, c) ++ num_blocks = [BLOCKS_PER_REQ] * MAX_NUM_REQS ++ ++ gathered_backing = torch.zeros_like(persistent_backing) ++ gathered_backing[0, :] = CANARY_SRC ++ gathered_backing[0, POISON_DST_COL] = CANARY_DST ++ gathered = gathered_backing[1:] ++ slot_mapping = [s for s in slot_mapping if s >= 0] ++ _gather_rows(persistent, num_blocks, slot_mapping, len(slot_mapping) + 1, gathered) ++ _gather_rows( ++ persistent, num_blocks, list(t1_mapping), len(t1_mapping) + 1, gathered ++ ) ++ ++ kv_cache_config = _make_kv_cache_config() ++ ctx = MambaSpecDecodeGPUContext.create( ++ max_num_reqs=MAX_NUM_REQS, ++ kv_cache_config=kv_cache_config, ++ copy_funcs=_COPY_FUNCS, ++ device=device, ++ make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), ++ ) ++ ctx.initialize_from_forward_context( ++ kv_cache_config, ++ {"layer_0": attention}, ++ _COPY_FUNCS, ++ [gathered], ++ persistent_block_tables=[persistent], ++ ) ++ assert ctx.has_persistent_block_tables ++ ++ return _Fixture( ++ device=device, ++ ctx=ctx, ++ conv=conv, ++ temporal=temporal, ++ conv0=conv.clone(), ++ temporal0=temporal.clone(), ++ persistent=persistent, ++ gathered=gathered, ++ ) ++ ++ ++def _launch(fx: _Fixture, mapping, per_slot): ++ """Run the deferred align post-process. ++ ++ ``per_slot`` maps req slot -> (num_accepted, state_idx, new_num_computed). ++ """ ++ device = fx.device ++ ++ def backed(fill: int, poison: int) -> torch.Tensor: ++ """A [MAX_NUM_REQS] view whose index -1 holds ``poison``.""" ++ backing = torch.full( ++ (MAX_NUM_REQS + 1,), fill, dtype=torch.int32, device=device ++ ) ++ backing[0] = poison ++ return backing[1:] ++ ++ poison_acc, poison_sidx, poison_computed = POISON ++ num_accepted = backed(1, poison_acc) ++ state_idx = backed(0, poison_sidx) ++ new_computed = backed(0, poison_computed) ++ for slot, (acc, sidx, computed) in per_slot.items(): ++ num_accepted[slot] = acc ++ state_idx[slot] = sidx ++ new_computed[slot] = computed ++ fx.ctx.run_fused_postprocess_align( ++ num_reqs=len(mapping), ++ num_accepted_tokens_gpu=num_accepted, ++ state_idx_gpu=state_idx, ++ new_num_computed_tokens_gpu=new_computed, ++ idx_mapping=torch.tensor(mapping, dtype=torch.int32, device=device), ++ ) ++ torch.accelerator.synchronize() ++ return num_accepted ++ ++ ++def _assert_bitwise(actual: torch.Tensor, expected: torch.Tensor, what: str) -> None: ++ """Compare the integer view, so "bit-identical" means what it says. ++ ++ A zero-tolerance float compare establishes numeric equality, not bit ++ identity: +0.0 and -0.0 compare equal and NaN payloads compare unequal. ++ These are byte copies, so the integer view is the right instrument. ++ ++ """ ++ assert actual.dtype == expected.dtype ++ lhs = actual.contiguous().view(torch.int16) ++ rhs = expected.contiguous().view(torch.int16) ++ assert torch.equal(lhs, rhs), ( ++ f"{what}: {int((lhs != rhs).sum().item())} of {lhs.numel()} 16-bit " ++ f"words differ" ++ ) ++ ++ ++def _assert_only_blocks_changed(fx: _Fixture, changed: set[int]) -> None: ++ """Every physical block outside ``changed`` must be bit-identical. ++ ++ The canary blocks (``CANARY_DST``, ``CANARY_SRC``) are included: they are ++ only reachable through the guard row a negative ``req_idx`` would select. ++ """ ++ for block in range(TOTAL_BLOCKS): ++ if block in changed: ++ continue ++ _assert_bitwise(fx.conv[block], fx.conv0[block], f"conv[{block}]") ++ _assert_bitwise(fx.temporal[block], fx.temporal0[block], f"temporal[{block}]") ++ ++ ++@pytest.mark.skipif( ++ is_conv_state_dim_first(), ++ reason="expected-bytes model assumes the SD conv layout", ++) ++def test_positive_bias_copy_lands_on_the_requests_own_blocks(): ++ """batch row 0 -> req slot 5, num_accepted 4, so token_bias is 3. ++ ++ new=48, a=4 -> aligned=48, running=45, needs_copy; bias=3; dst_col=2. ++ state_idx=0, so conv reads column 0 and temporal reads column ++ ``src_col + bias`` = 3 -- three distinct columns, so a wrong row cannot be ++ masked by a coincidental self-copy. ++ """ ++ fx = _build() ++ _launch(fx, [5, 1, 2], {5: (4, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) ++ ++ dst = phys(5, 2) ++ conv_src = phys(5, 0) ++ temporal_src = phys(5, 3) ++ ++ # conv: state[bt[src_col], bias:] -> state[bt[dst_col], :width - bias] ++ _assert_bitwise(fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[conv_src][3:], "conv dst") ++ # the tail of the destination conv window is not written ++ _assert_bitwise( ++ fx.conv[dst][CONV_WIDTH - 3 :], fx.conv0[dst][CONV_WIDTH - 3 :], "conv tail" ++ ) ++ # temporal: state[bt[src_col + bias]] -> state[bt[dst_col]] ++ _assert_bitwise(fx.temporal[dst], fx.temporal0[temporal_src], "temporal dst") ++ ++ # Nothing else moved -- in particular not the victim r2 (the first-cut row) and ++ # not r0 (the "persistent table but batch row" mistake). ++ _assert_only_blocks_changed(fx, {dst}) ++ victim = phys(2, 2) # gathered[batch row 0] column 2 -- the first-cut destination ++ _assert_bitwise(fx.conv[victim], fx.conv0[victim], "victim conv") ++ _assert_bitwise(fx.temporal[victim], fx.temporal0[victim], "victim temporal") ++ wrong_row = phys(0, 2) # persistent[batch row 0] column 2 ++ _assert_bitwise(fx.conv[wrong_row], fx.conv0[wrong_row], "wrong-row conv") ++ ++ ++@pytest.mark.skipif( ++ is_conv_state_dim_first(), ++ reason="expected-bytes model assumes the SD conv layout", ++) ++def test_zero_bias_migration_copies_the_whole_window(): ++ """new=48, a=1 -> aligned=48, running=48, bias=0, dst_col=2, src_col=0. ++ ++ src != dst so the self-copy exit does not fire, and the whole conv window ++ plus the temporal block migrate. ++ """ ++ fx = _build() ++ _launch(fx, [5, 1, 2], {5: (1, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) ++ ++ dst = phys(5, 2) ++ src = phys(5, 0) ++ _assert_bitwise(fx.conv[dst], fx.conv0[src], "conv dst") ++ _assert_bitwise(fx.temporal[dst], fx.temporal0[src], "temporal dst") ++ _assert_only_blocks_changed(fx, {dst}) ++ ++ ++def test_self_copy_exits_without_touching_state_and_resets_num_accepted(): ++ """src_block_idx == dest_block_idx and bias == 0 -> return (mamba_utils:478). ++ ++ new=33, a=2 -> aligned=32, running=32, needs_copy true, bias=0, dst_col=1; ++ state_idx=1 makes src == dst. Seeding num_accepted at 2 rather than 1 is ++ what makes the reset store observable: the kernel must write ++ num_accepted_tokens_out[req] = 1 (mamba_utils.py:473-475) *before* the early ++ return, so deleting that store now fails the test. (A prior version seeded 1 and ++ asserted 1, which a deleted store would have passed.) ++ ++ Limit, stated rather than papered over: unchanged bytes cannot distinguish ++ "returned at the guard" from "fell through and memcpy'd a block onto ++ itself", because those are byte-identical by construction. What the reset ++ store does establish is that the kernel got past `needs_copy` and evaluated ++ `src == dst`, so it is the *guard* that stopped it and not an earlier ++ return. The conjunction itself is covered by the next test. ++ """ ++ fx = _build() ++ num_accepted = _launch( ++ fx, [5, 1, 2], {5: (2, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} ++ ) ++ ++ _assert_only_blocks_changed(fx, set()) ++ assert int(num_accepted[5].item()) == 1 ++ ++ ++@pytest.mark.skipif( ++ is_conv_state_dim_first(), ++ reason="expected-bytes model assumes the SD conv layout", ++) ++def test_same_column_with_positive_bias_copies_and_still_resets(): ++ """The early-exit guard is a conjunction, and this is the other conjunct. ++ ++ new=33, a=4 -> aligned=32, running=30, needs_copy; bias = 4-1-1 = 2; ++ dst_col = 1; state_idx = 1, so src_block_idx == dest_block_idx but ++ token_bias != 0. The kernel must NOT return: conv takes the same-block ++ memmove branch (mamba_utils.py:320-336) and temporal reads ++ ``src_col + token_bias`` = column 3, a different physical block. ++ ++ Together with the previous test this makes the exit observable in both ++ directions: same column and zero bias moves nothing, same column and ++ positive bias moves bytes. ++ """ ++ fx = _build() ++ num_accepted = _launch( ++ fx, [5, 1, 2], {5: (4, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} ++ ) ++ ++ blk = phys(5, 1) # both the conv source and the destination ++ temporal_src = phys(5, 3) # src_col + bias = 3 ++ ++ # conv, same block, shifted left by token_bias = 2 ++ _assert_bitwise(fx.conv[blk][: CONV_WIDTH - 2], fx.conv0[blk][2:], "conv shift") ++ _assert_bitwise( ++ fx.conv[blk][CONV_WIDTH - 2 :], fx.conv0[blk][CONV_WIDTH - 2 :], "conv tail" ++ ) ++ # temporal, from a different block ++ _assert_bitwise(fx.temporal[blk], fx.temporal0[temporal_src], "temporal dst") ++ ++ _assert_only_blocks_changed(fx, {blk}) ++ # The reset store fires here too: it is guarded only on src == dst. ++ assert int(num_accepted[5].item()) == 1 ++ ++ ++@pytest.mark.skipif( ++ is_conv_state_dim_first(), ++ reason="expected-bytes model assumes the SD conv layout", ++) ++def test_freed_slot_is_rejected_before_any_negative_index_load(): ++ """A -1 row must be rejected *before* the decision arrays are read. ++ ++ ``mamba_utils.py:441-442`` returns on ``req_idx < 0`` before the loads at ++ ``:445-446``. A prior version of this test seeded decisions at some other ++ slot, which the kernel would never have touched either way, so it proved ++ nothing about the guard. This one plants the poison ++ where a missing guard would actually look: index -1 of every per-req array ++ and row -1 of the block table, which are real memory owned by the backing ++ tensors in ``_build``. ++ ++ If the guard were deleted, the kernel would read ++ ``num_accepted=4, state_idx=0, new_num_computed=48`` -> a copy with bias 3 ++ into block-table row -1, writing CANARY_DST from CANARY_SRC. Both canaries ++ are covered by ``_assert_only_blocks_changed``, so that write fails the ++ test. ++ ++ The row after the -1 must also still copy correctly, which is the ++ positioning check. ++ """ ++ fx = _build(slot_mapping=(5, 1, 2)) ++ _launch(fx, [-1, 5, 2], {5: (4, 0, 48), 2: (2, 0, 40)}) ++ ++ dst = phys(5, 2) ++ _assert_bitwise( ++ fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(5, 0)][3:], "conv dst" ++ ) ++ _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(5, 3)], "temporal dst") ++ _assert_only_blocks_changed(fx, {dst}) ++ ++ ++def test_generation_masked_reuse_is_skipped_end_to_end(): ++ """The -1 above, but produced by the real generation masking. ++ ++ ``mask_from_generation_change`` in the CPU sibling test drives ++ ``PPHandler.get_prev_sampled_outputs`` (pp_utils.py:132-140) and returns ++ the mapping it computes; here that mapping is fed to the kernel. Slot 1 is ++ the one whose generation moved -- it was freed and its slot handed to a new ++ request between capture and use -- so its row must copy nothing even though ++ slot 1 now carries decisions that would copy. ++ """ ++ from tests.v1.worker.gpu.test_pp_deferred_mamba_block_table import ( ++ mask_from_generation_change, ++ ) ++ ++ mapping = mask_from_generation_change( ++ idx_mapping=[5, 1, 2], reused_slot_rows=[1] ++ ) ++ assert mapping == [5, -1, 2], mapping ++ ++ fx = _build(slot_mapping=(5, 1, 2)) ++ _launch(fx, mapping, {5: (4, 0, 48), 1: (4, 0, 48), 2: (2, 0, 40)}) ++ ++ dst = phys(5, 2) ++ _assert_bitwise( ++ fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(5, 0)][3:], "conv dst" ++ ) ++ _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(5, 3)], "temporal dst") ++ # Slot 1's own blocks are untouched despite its copy-shaped decisions. ++ _assert_only_blocks_changed(fx, {dst}) ++ ++ ++@pytest.mark.skipif( ++ is_conv_state_dim_first(), ++ reason="expected-bytes model assumes the SD conv layout", ++) ++def test_same_step_identity_mapping_is_unaffected(): ++ """The last-rank / no-PP path: batch order == req order, gather is current. ++ ++ req_idx == batch_idx here, so the change must be a no-op relative to the first cut. ++ """ ++ fx = _build(slot_mapping=(0, 1, 2), t1_mapping=(0, 1, 2)) ++ _launch(fx, [0, 1, 2], {0: (4, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) ++ ++ dst = phys(0, 2) ++ _assert_bitwise( ++ fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(0, 0)][3:], "conv dst" ++ ) ++ _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(0, 3)], "temporal dst") ++ _assert_only_blocks_changed(fx, {dst}) ++ ++ ++def test_align_launch_raises_without_persistent_tables(): ++ """Fail closed on real hardware too, before any kernel is enqueued.""" ++ device = torch.device("cuda:0") ++ conv = torch.randn( ++ TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device ++ ) ++ temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) ++ attention = MagicMock() ++ attention.kv_cache = [conv, temporal] ++ gathered = torch.zeros( ++ MAX_NUM_REQS, BLOCKS_PER_REQ, dtype=torch.int32, device=device ++ ) ++ ++ kv_cache_config = _make_kv_cache_config() ++ ctx = MambaSpecDecodeGPUContext.create( ++ max_num_reqs=MAX_NUM_REQS, ++ kv_cache_config=kv_cache_config, ++ copy_funcs=_COPY_FUNCS, ++ device=device, ++ make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), ++ ) ++ ctx.initialize_from_forward_context( ++ kv_cache_config, {"layer_0": attention}, _COPY_FUNCS, [gathered] ++ ) ++ assert ctx.has_persistent_block_tables is False ++ ++ snapshot = conv.clone() ++ with pytest.raises(RuntimeError, match="persistent per-req-slot block tables"): ++ ctx.run_fused_postprocess_align( ++ num_reqs=1, ++ num_accepted_tokens_gpu=torch.full( ++ (MAX_NUM_REQS,), 4, dtype=torch.int32, device=device ++ ), ++ state_idx_gpu=torch.zeros( ++ MAX_NUM_REQS, dtype=torch.int32, device=device ++ ), ++ new_num_computed_tokens_gpu=torch.full( ++ (MAX_NUM_REQS,), 48, dtype=torch.int32, device=device ++ ), ++ idx_mapping=torch.tensor([5], dtype=torch.int32, device=device), ++ ) ++ torch.accelerator.synchronize() ++ _assert_bitwise(conv, snapshot, "conv unchanged") +diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py +index a162c39..2ab23ce 100644 +--- a/tests/v1/worker/test_mamba_utils.py ++++ b/tests/v1/worker/test_mamba_utils.py +@@ -277,8 +277,14 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): + ) + forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} + ++ # V1: batch order == req-slot order, so the same tensor is both the ++ # gathered and the persistent table. + gpu_ctx.initialize_from_forward_context( +- kv_cache_config, forward_context, _COPY_FUNCS, [block_table] ++ kv_cache_config, ++ forward_context, ++ _COPY_FUNCS, ++ [block_table], ++ persistent_block_tables=[block_table], + ) + + assert gpu_ctx.state_base_addrs.tolist() == [ +@@ -548,6 +554,7 @@ def test_mamba_groups_support_different_state_specs(): + forward_context, + _COPY_FUNCS, + block_tables, ++ persistent_block_tables=block_tables, + ) + assert ctx.is_initialized + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1] +@@ -628,11 +635,13 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): + ), + "ple.0": MagicMock(kv_cache=(torch.empty(8, 12, 8, dtype=torch.float16),)), + } ++ group_table = torch.zeros(4, 8, dtype=torch.int32) + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + _COPY_FUNCS, +- [torch.zeros(4, 8, dtype=torch.int32)], ++ [group_table], ++ persistent_block_tables=[group_table], + ) + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 0] + assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] +@@ -777,11 +786,13 @@ def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: + ] + context = _make_gpu_ctx(config, kv_cache_config, device) + ++ stub_table = torch.zeros(1, 1, dtype=torch.int32) + context.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, +- [torch.zeros(1, 1, dtype=torch.int32)], ++ [stub_table], ++ persistent_block_tables=[stub_table], + ) + + assert context.is_initialized +@@ -812,8 +823,14 @@ def _run_gpu_postprocess( + def t(values): + return torch.tensor(values, dtype=torch.int32, device=device) + ++ # V1: batch order == req-slot order, so the same tensor is both the ++ # gathered and the persistent table. + gpu_ctx.initialize_from_forward_context( +- kv_cache_config, forward_context, copy_funcs, [block_table] ++ kv_cache_config, ++ forward_context, ++ copy_funcs, ++ [block_table], ++ persistent_block_tables=[block_table], + ) + gpu_ctx.run_fused_postprocess( + num_reqs=len(req_ids), +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index f162cb5..e3479c2 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -677,6 +677,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + cp_rank=self.dcp_rank, + cp_interleave=self.cp_interleave, + ) ++ # Mamba "align" indexes a block table from the PP-deferred post-process, ++ # which runs pp_size steps after the gather that produced its ++ # idx_mapping; give the model state the persistent per-req-slot tables. ++ self.model_state.set_block_tables(self.block_tables) + self.pcp_manager = pcp.maybe_build_pcp_manager( + self.vllm_config, + self.device, +diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py +index a5c8d6f..8ce5121 100644 +--- a/vllm/v1/worker/gpu/model_states/interface.py ++++ b/vllm/v1/worker/gpu/model_states/interface.py +@@ -126,6 +126,14 @@ class ModelState(ABC): + def apply_staged_writes(self) -> None: + return None + ++ def set_block_tables(self, block_tables: Any) -> None: ++ """Register the runner's `BlockTables` once it exists. ++ ++ Model states that index a block table outside the step that gathered it ++ (mamba "align" + the PP-deferred post-process) need the persistent ++ per-req-slot tables, not the batch-order gather. No-op by default.""" ++ return None ++ + def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: + """Cudagraph support of attention groups this ModelState builds outside + ``init_attn_backend`` (e.g. encoder-only layers). +diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +index 0ac3f2f..0a8b5d7 100644 +--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py ++++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +@@ -94,6 +94,9 @@ class MambaHybridModelState(DefaultModelState): + # kernel reusing the postprocess copy machinery, so the per-step src + # columns and the running state_idx are kept GPU-resident. + self._align_mode = self.cache_config.mamba_cache_mode == "align" ++ # Set by the runner via `set_block_tables`; the align post-process needs ++ # the persistent per-req-slot tables, not the per-step gather. ++ self._runner_block_tables: Any = None + self.recoverssm = ( + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + ) +@@ -112,6 +115,9 @@ class MambaHybridModelState(DefaultModelState): + self._mamba_spec: MambaSpec | None = None + self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None + ++ def set_block_tables(self, block_tables: Any) -> None: ++ self._runner_block_tables = block_tables ++ + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + # Must reset the speculative acceptance count in this idx which could be stale. +@@ -172,11 +178,24 @@ class MambaHybridModelState(DefaultModelState): + # block_tables are batch-order slices of the persistent + # input_block_tables (stable data_ptr), so the metadata is captured + # once here and reused across steps. ++ # ++ # Also hand over the per-req-slot tables the gather reads from ++ # (block_table.py:53-59). The align post-process can run pp_size ++ # steps after the gather that produced its idx_mapping, at which ++ # point the batch-order buffer belongs to another micro-batch; it ++ # indexes these by req slot instead. ++ persistent_block_tables = None ++ if self._runner_block_tables is not None: ++ persistent_block_tables = [ ++ self._runner_block_tables.block_tables[gid].gpu ++ for gid in mamba_group_ids ++ ] + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + copy_funcs, + [block_tables[gid] for gid in mamba_group_ids], ++ persistent_block_tables=persistent_block_tables, + ) + return ctx + +diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py +index 1deb89b..621efcb 100644 +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -409,6 +409,14 @@ def postprocess_mamba_fused_kernel( + # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds + # the post-step new_num_computed value (V2 supplies the advanced count). + PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, ++ # BT_ROW_IS_REQ: when True, block_table_ptrs_ptr points at the *persistent* ++ # per-req-slot block tables (BlockTables.block_tables[i].gpu, ++ # block_table.py:53-59), whose rows are request state slots, so the row ++ # index is req_idx rather than batch_idx. When False it points at the ++ # batch-order gather (BlockTables.input_block_tables, block_table.py:74-76) ++ # and batch_idx is the row. Required by the PP-deferred post-process, which ++ # runs pp_size steps after the gather that produced its idx_mapping. ++ BT_ROW_IS_REQ: tl.constexpr = False, + # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across + # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a + # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves +@@ -477,7 +485,15 @@ def postprocess_mamba_fused_kernel( + if src_block_idx == dest_block_idx and accept_token_bias == 0: + return + +- bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx ++ if BT_ROW_IS_REQ: ++ # Persistent per-req-slot table: the row is the request's own slot, so ++ # it stays correct however many steps separate the gather from this ++ # launch (the PP-deferred post-process is enqueued at the top of ++ # execute_model, before this step's gather, with an idx_mapping ++ # captured pp_size steps earlier). ++ bt_row_idx = req_idx ++ else: ++ bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx + _copy_mamba_state_block( + state_idx, + bt_row_idx, +@@ -812,6 +828,24 @@ class MambaSpecDecodeGPUContext: + block_table_ptrs: torch.Tensor + block_table_stride_req: int = 0 + ++ # Per-group *persistent* per-req-slot block-table addresses: int64[num_groups]. ++ # `block_table_ptrs` above addresses the batch-order gather ++ # (BlockTables.input_block_tables, block_table.py:74-76), which is rewritten ++ # in place every step and is only valid for the step that gathered it. These ++ # address BlockTables.block_tables[i].gpu (block_table.py:53-59), the ++ # persistent [max_num_reqs, max_num_blocks] tables the gather reads from ++ # (their data_ptrs are already held device-side at block_table.py:94-96). ++ # Rows there are request state slots, so they can be indexed by req_idx from ++ # any later step. Used only by the align post-process (BT_ROW_IS_REQ). ++ persistent_block_table_ptrs: torch.Tensor | None = None ++ has_persistent_block_tables: bool = False ++ # shape[0] of the persistent tables; see the note in set_block_tables. ++ # -1, not 0: this is the ROW BOUND the align post-process uses, and a 0 ++ # default bounds EVERY row out, so a context that never registered the ++ # persistent tables would report guard hits on correct metadata instead of ++ # copying state. -1 is unusable on purpose and is asserted before use. ++ persistent_block_table_num_rows: int = -1 ++ + # persistent output for the once-per-step, all-group aligned-index launch. + # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks]. + aligned_state_indices: torch.Tensor | None = None +@@ -902,6 +936,9 @@ class MambaSpecDecodeGPUContext: + block_table_ptrs=torch.zeros( + len(mamba_group_ids), dtype=torch.int64, device=device + ), ++ persistent_block_table_ptrs=torch.zeros( ++ len(mamba_group_ids), dtype=torch.int64, device=device ++ ), + aligned_state_indices=torch.empty( + ( + len(mamba_group_ids), +@@ -926,6 +963,7 @@ class MambaSpecDecodeGPUContext: + forward_context: dict[str, Any], + mamba_state_copy_funcs: MambaStateCopyFuncsByType, + block_tables: list[torch.Tensor], ++ persistent_block_tables: list[torch.Tensor] | None = None, + ) -> None: + """ + Extract and cache memory layout metadata from Mamba state tensors. +@@ -957,9 +995,14 @@ class MambaSpecDecodeGPUContext: + have a `kv_cache` attribute containing the list of state tensors. + mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to + copy functions. +- block_tables: per-mamba-group persistent block-table tensors, in +- the same order as `mamba_group_ids`. Their `data_ptr()` / +- `stride(0)` are captured once for the kernel to index into. ++ block_tables: per-mamba-group *gathered*, batch-order block-table ++ tensors, in the same order as `mamba_group_ids`. Their ++ `data_ptr()` / `stride(0)` are captured once for the kernel to ++ index into by batch row. ++ persistent_block_tables: per-mamba-group *persistent* per-req-slot ++ block-table tensors (`BlockTables.block_tables[i].gpu`), same ++ order. Captured separately so the align post-process can index ++ by req slot; see `persistent_block_table_ptrs`. + """ + if self.is_initialized: + return +@@ -970,6 +1013,7 @@ class MambaSpecDecodeGPUContext: + forward_context, + mamba_state_copy_funcs, + block_tables, ++ persistent_block_tables, + ) + + def _populate_metadata( +@@ -978,6 +1022,7 @@ class MambaSpecDecodeGPUContext: + forward_context: dict[str, Any], + mamba_state_copy_funcs: MambaStateCopyFuncsByType, + block_tables: list[torch.Tensor], ++ persistent_block_tables: list[torch.Tensor] | None = None, + ) -> None: + idx = 0 + for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): +@@ -1253,6 +1298,35 @@ class MambaSpecDecodeGPUContext: + if num_reqs == 0 or not self.is_initialized: + return + ++ # Fail closed. Under PP this launch is the *deferred* post-process: it ++ # is enqueued by `update_pp_decode_requests` at the top of ++ # `execute_model` (model_runner.py:1644), before that step's ++ # `prepare_attn` (model_runner.py:1718) and the ++ # `gather_block_tables` inside it (:1442 / block_table.py:149), and ++ # `idx_mapping` was captured pp_size steps earlier ++ # (pp_utils.py:57-58, 258-271). The single batch-order gather buffer ++ # therefore holds a different micro-batch's rows and `batch_idx` ++ # resolves to the wrong request. This path MUST index the persistent ++ # per-req-slot tables by `req_idx`; there is no safe fallback, so a ++ # context without them is a bug in the caller, not a degraded mode. ++ # (`req_idx` is also correct on the same-step last-rank / no-PP path.) ++ if not self.has_persistent_block_tables: ++ raise RuntimeError( ++ "MambaSpecDecodeGPUContext was initialized without the " ++ "persistent per-req-slot block tables, so the V2 align " ++ "post-process cannot index them by req slot. Falling back to " ++ "the batch-order gather is the known deferred-PP wrong-row " ++ "corruption, so this raises instead. The runner must call " ++ "ModelState.set_block_tables() before the first align step." ++ ) ++ block_table_ptrs = self.persistent_block_table_ptrs ++ assert self.persistent_block_table_num_rows > 0, ( ++ "persistent_block_table_num_rows was never set, so the row bound " ++ "in _copy_mamba_state_block would reject every row and report " ++ "guard hits on correct metadata. set_block_tables() must run " ++ "before the first align post-process." ++ ) ++ + # V2 reads non-contiguous idx_mapping positions, so snapshot the whole + # decision buffer rather than only [:num_reqs]. + num_accepted_tokens_snapshot = self.num_accepted_tokens_out +@@ -1284,6 +1358,7 @@ class MambaSpecDecodeGPUContext: + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), + HAS_IDX_MAPPING=True, + PRECOMPUTED_NEW_COMPUTED=True, ++ BT_ROW_IS_REQ=True, + TEMPORAL_TILES=_TEMPORAL_TILES, + ) + +-- +2.43.0 + diff --git a/patches/0008-core-cap-the-Mamba-align-mode-admission-estimate-por.patch b/patches/0008-core-cap-the-Mamba-align-mode-admission-estimate-por.patch new file mode 100644 index 0000000..0b7dfa9 --- /dev/null +++ b/patches/0008-core-cap-the-Mamba-align-mode-admission-estimate-por.patch @@ -0,0 +1,101 @@ +From 8cb229146ba10ea093d647deeaf9a074afaa5bee Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:12:36 -0400 +Subject: [PATCH 008/103] core: cap the Mamba align-mode admission estimate + (port of upstream #57050) + +--- + .../core/test_single_type_kv_cache_manager.py | 60 +++++++++++++++++++ + vllm/v1/core/single_type_kv_cache_manager.py | 8 +-- + 2 files changed, 64 insertions(+), 4 deletions(-) + +diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py +index 14e46d7..1598cec 100644 +--- a/tests/v1/core/test_single_type_kv_cache_manager.py ++++ b/tests/v1/core/test_single_type_kv_cache_manager.py +@@ -857,3 +857,63 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): + f"but allocator pulled {len(new_blocks)}" + ) + total_computed = num_tokens ++ ++ ++@pytest.mark.parametrize("block_size", [896, 1536]) ++@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4]) ++@pytest.mark.parametrize("prompt_tokens", [25121, 704547]) ++def test_mamba_checkpoint_admission_matches_allocation( ++ block_size, num_speculative_blocks, prompt_tokens ++): ++ """Checkpoint admission must match the subsequent physical allocation. ++ ++ Ported verbatim from upstream vllm-project/vllm #57050 (merged 2026-09-16, f30a195) for v2. ++ """ ++ spec = MambaSpec( ++ block_size=block_size, ++ shapes=((1, 1),), ++ dtypes=(torch.float32,), ++ mamba_cache_mode="align", ++ num_speculative_blocks=num_speculative_blocks, ++ num_prefill_checkpoint_blocks=1, ++ prefill_checkpoint_alignment=64, ++ ) ++ pool = BlockPool( ++ num_gpu_blocks=2048, ++ enable_caching=True, ++ hash_block_size=128, ++ ) ++ manager = MambaManager( ++ spec, ++ block_pool=pool, ++ enable_caching=True, ++ kv_cache_group_id=0, ++ scheduler_block_size=block_size, ++ ) ++ request_id = "prefill" ++ computed_tokens = 23040 ++ ++ def estimate(num_tokens, total_computed_tokens, apply_admission_cap): ++ return manager.get_num_blocks_to_allocate( ++ request_id=request_id, ++ num_tokens=num_tokens, ++ new_computed_blocks=[], ++ total_computed_tokens=total_computed_tokens, ++ num_local_computed_tokens=total_computed_tokens, ++ num_tokens_main_model=num_tokens, ++ apply_admission_cap=apply_admission_cap, ++ ) ++ ++ estimate(computed_tokens, 0, False) ++ manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens) ++ assert request_id in manager._allocated_block_reqs ++ ++ admission_estimate = estimate(prompt_tokens, computed_tokens, True) ++ allocation_estimate = estimate(prompt_tokens, computed_tokens, False) ++ assert request_id in manager._checkpoints ++ ++ free_before = pool.get_num_free_blocks() ++ manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens) ++ allocated = free_before - pool.get_num_free_blocks() ++ ++ assert admission_estimate == allocation_estimate == allocated +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 40027e7..611e1ce 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -1796,10 +1796,10 @@ class MambaManager(SingleTypeKVCacheManager): + self._checkpoints.pop(request_id, None) + if num_new_blocks > 0: + blocks_allocated = request_id in self._allocated_block_reqs +- if not (checkpoint_block and blocks_allocated): +- num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block +- if not blocks_allocated: +- num_new_blocks += self.num_speculative_blocks ++ physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block ++ if not blocks_allocated or checkpoint_block: ++ physical_block_cap += self.num_speculative_blocks ++ num_new_blocks = min(num_new_blocks, physical_block_cap) + + num_evictable_computed_blocks = self._get_num_evictable_blocks( + new_computed_blocks +-- +2.43.0 + diff --git a/patches/0009-mamba-bound-the-state-copy-block-table-lookups.patch b/patches/0009-mamba-bound-the-state-copy-block-table-lookups.patch new file mode 100644 index 0000000..59efb8d --- /dev/null +++ b/patches/0009-mamba-bound-the-state-copy-block-table-lookups.patch @@ -0,0 +1,399 @@ +From b60dadeaaf41349275337d5d1fcd0b1a86d630c7 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:02 -0400 +Subject: [PATCH 009/103] mamba: bound the state-copy block-table lookups + +Ports the check from upstream vLLM PR #50021; a hit is reported through the +guard telemetry instead of reading out of bounds. +--- + tests/v1/worker/test_mamba_utils.py | 75 ++++++++++++ + vllm/v1/worker/mamba_utils.py | 179 +++++++++++++++++++++++++++- + 2 files changed, 249 insertions(+), 5 deletions(-) + +diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py +index 2ab23ce..b6bad93 100644 +--- a/tests/v1/worker/test_mamba_utils.py ++++ b/tests/v1/worker/test_mamba_utils.py +@@ -1072,6 +1072,81 @@ class TestPostprocessMambaFusedKernel: + torch.testing.assert_close(conv_state, conv_state_orig) + torch.testing.assert_close(temporal_state, temporal_state_orig) + ++ # Ported verbatim from upstream vLLM PR #50021 (bounds-guard instrumentation, copy-state guard). The ++ # four cases are exactly the four block-table lookups the guard bounds: ++ # an in-row NULL_BLOCK_ID destination, a destination column past the row, ++ # a source column past the row, and a temporal bias that crosses the row ++ # boundary. Without the guard the last three read the NEXT request's row ++ # and copy state from or to a block this request does not own. ++ @pytest.mark.parametrize( ++ ("num_computed", "source_col", "destination_col", "conv_unchanged"), ++ [ ++ (60, 1, 3, True), # destination is an in-row NULL_BLOCK_ID ++ (76, 1, 4, True), # destination is the next row's first column ++ (60, 4, 3, True), # source is the next row's first column ++ (60, 3, 3, False), # temporal bias crosses into the next row ++ ], ++ ) ++ def test_invalid_block_table_lookup_does_not_copy_state( ++ self, ++ device, ++ num_computed, ++ source_col, ++ destination_col, ++ conv_unchanged, ++ ): ++ """Invalid columns and NULL_BLOCK_ID must not reach state address math.""" ++ cfg = _TestConfig(num_blocks=8) ++ layer_names = ["layer_0"] ++ kv_cache_config = _make_kv_cache_config(cfg, layer_names) ++ conv_state = torch.randn( ++ cfg.num_blocks, ++ cfg.conv_width, ++ cfg.conv_inner_dim, ++ dtype=cfg.dtype, ++ device=device, ++ ) ++ temporal_state = torch.randn( ++ cfg.num_blocks, ++ cfg.temporal_state_dim, ++ dtype=cfg.dtype, ++ device=device, ++ ) ++ conv_before = conv_state.clone() ++ temporal_before = temporal_state.clone() ++ forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} ++ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) ++ ++ block_table_storage = torch.zeros(2, 4, dtype=torch.int32, device=device) ++ block_table_storage[0, 1] = 1 ++ if destination_col < block_table_storage.stride(0): ++ is_null_destination = ( ++ destination_col == 3 and num_computed == 60 and source_col == 1 ++ ) ++ block_table_storage[0, destination_col] = 0 if is_null_destination else 3 ++ block_table_storage[1, 0] = 2 ++ ++ _run_gpu_postprocess( ++ gpu_ctx, ++ kv_cache_config=kv_cache_config, ++ forward_context=forward_context, ++ copy_funcs=_COPY_FUNCS, ++ block_table=block_table_storage[:1], ++ req_ids=["req_0"], ++ num_accepted_tokens=[3], ++ mamba_state_idx=[source_col], ++ num_scheduled_tokens={"req_0": 3}, ++ num_computed_tokens=[num_computed], ++ num_draft_tokens={}, ++ device=device, ++ ) ++ ++ expected_destination = (num_computed + 3 + 3 - 1) // cfg.block_size - 1 ++ assert expected_destination == destination_col ++ if conv_unchanged: ++ torch.testing.assert_close(conv_state, conv_before) ++ torch.testing.assert_close(temporal_state, temporal_before) ++ + @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16]) + def test_various_batch_sizes(self, device, test_config, num_reqs): + """Verify kernel works correctly with different batch sizes.""" +diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py +index 621efcb..39d71b3 100644 +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -26,14 +26,26 @@ from vllm.v1.kv_cache_interface import ( + UniformTypeKVCacheSpecs, + ) + from vllm.v1.utils import CpuGpuBuffer ++from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch + + logger = init_logger(__name__) + ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + # 16 saturates HBM on H100/GB200 across the reqs=8..128 range in + # microbenchmarks + _TEMPORAL_TILES = 16 + ++# Row bound used when no trustworthy row count is available: large enough never ++# to reject a real request slot, so the guard degrades to "reject negatives" ++# rather than to "reject correct metadata". See initialize_from_forward_context. ++_UNBOUNDED_ROWS = 1 << 30 ++ + + @triton.jit(do_not_specialize=["num_requests"]) + def get_aligned_state_indices_multi_group_kernel( +@@ -196,6 +208,8 @@ def _copy_mamba_state_block( + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, ++ block_table_num_rows, ++ guard_ctr, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, +@@ -243,12 +257,39 @@ def _copy_mamba_state_block( + group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) + group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) + block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) ++ # Bounds-guard instrumentation: bound the ROW as well as the column. #50021 bounds which ++ # column is read within a row; this bounds which row. CAN fire on ++ # postprocess_mamba_fused_kernel, where bt_row_idx comes from a device ++ # idx_mapping value. CANNOT fire on precopy_mamba_align_fused_kernel, whose ++ # bt_row_idx is program_id(0) bounded by num_reqs <= rows -- that path has a ++ # TEMPORAL staleness problem, fixed separately, not a spatial one. ++ if (bt_row_idx < 0) or (bt_row_idx >= block_table_num_rows): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, _GUARD_BIAS + (-1)) ++ return + block_table_base = block_table_typed + bt_row_idx * block_table_stride_req + + # Widen block ids to int64 before they reach `block_id * state_block_stride` + # below: state_block_stride can exceed 2**31 bytes for large mamba caches, + # and Triton would otherwise do the multiply in int32 and wrap. +- dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64) ++ # PR #50021 (upstream :178): these columns are derived from per-request ++ # accepted-token counts. A wrong count walks the read past this request's ++ # row; the loaded int32 is then multiplied by state_block_stride and ++ # dereferenced, so an out-of-range column becomes a wild (and possibly ++ # misaligned) address rather than merely a wrong copy. Only the ++ # out-of-range case is counted: `<= 0` also matches NULL_BLOCK_ID (0), ++ # which is the ordinary unallocated marker and not a fault. ++ dst_col_ok = (dst_col >= 0) & (dst_col < block_table_stride_req) ++ if (dst_col < 0) or (dst_col >= block_table_stride_req): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (dst_col).to(tl.int64) + _GUARD_BIAS) ++ dest_block_id = tl.load( ++ block_table_base + dst_col, mask=dst_col_ok, other=-1 ++ ).to(tl.int64) ++ if dest_block_id <= 0: ++ return + dst_addr = state_base_addr + dest_block_id * state_block_stride + + is_conv_state = conv_width > 0 +@@ -259,7 +300,16 @@ def _copy_mamba_state_block( + if tile_idx > 0: + return + # DS conv layout: state_len is the slide axis; copy per dim row. +- src_block_id = tl.load(block_table_base + src_col).to(tl.int64) ++ src_col_ok = (src_col >= 0) & (src_col < block_table_stride_req) ++ if (src_col < 0) or (src_col >= block_table_stride_req): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (src_col).to(tl.int64) + _GUARD_BIAS) ++ src_block_id = tl.load( ++ block_table_base + src_col, mask=src_col_ok, other=-1 ++ ).to(tl.int64) ++ if src_block_id <= 0: ++ return + dim_rows = tl.load(state_dim_row_count_ptr + state_idx) + row_stride = tl.load(state_dim_row_stride_ptr + state_idx) + src_block_addr = state_base_addr + src_block_id * state_block_stride +@@ -310,7 +360,16 @@ def _copy_mamba_state_block( + # SD conv: copy + # state[bt[src_col], token_bias:] -> + # state[bt[dst_col], :conv_width - token_bias] +- src_block_id = tl.load(block_table_base + src_col).to(tl.int64) ++ src_col_ok = (src_col >= 0) & (src_col < block_table_stride_req) ++ if (src_col < 0) or (src_col >= block_table_stride_req): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (src_col).to(tl.int64) + _GUARD_BIAS) ++ src_block_id = tl.load( ++ block_table_base + src_col, mask=src_col_ok, other=-1 ++ ).to(tl.int64) ++ if src_block_id <= 0: ++ return + src_block_addr = state_base_addr + src_block_id * state_block_stride + token_bytes = state_inner_size * state_elem_size + num_dst_tokens = conv_width - token_bias +@@ -349,7 +408,17 @@ def _copy_mamba_state_block( + # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] + # Body u64 range is partitioned across TEMPORAL_TILES CTAs to keep the + # SMs filled at small batch. +- actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64) ++ tmp_col = src_col + token_bias ++ tmp_col_ok = (tmp_col >= 0) & (tmp_col < block_table_stride_req) ++ if (tmp_col < 0) or (tmp_col >= block_table_stride_req): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (tmp_col).to(tl.int64) + _GUARD_BIAS) ++ actual_src_block_id = tl.load( ++ block_table_base + tmp_col, mask=tmp_col_ok, other=-1 ++ ).to(tl.int64) ++ if actual_src_block_id <= 0: ++ return + src_addr = state_base_addr + actual_src_block_id * state_block_stride + # Use natural block data size (inner_size * elem_size), NOT + # state_block_stride which is the page stride and can exceed the +@@ -378,6 +447,8 @@ def postprocess_mamba_fused_kernel( + # block table. + block_table_ptrs_ptr, + block_table_stride_req: tl.int64, # stride between requests (in elements) ++ block_table_num_rows: tl.int64, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + # Mamba state metadata, flattened across all layer/state tensors. Different + # Mamba groups may contribute different numbers of state tensors. + state_base_addrs_ptr, # base address of each state tensor +@@ -502,6 +573,8 @@ def postprocess_mamba_fused_kernel( + accept_token_bias, + block_table_ptrs_ptr, + block_table_stride_req, ++ block_table_num_rows, ++ guard_ctr, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, +@@ -573,6 +646,8 @@ def precopy_mamba_align_fused_kernel( + # Same flattened state-layout metadata as postprocess_mamba_fused_kernel + block_table_ptrs_ptr, + block_table_stride_req: tl.int64, ++ block_table_num_rows: tl.int64, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, +@@ -633,6 +708,8 @@ def precopy_mamba_align_fused_kernel( + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, ++ block_table_num_rows, ++ guard_ctr, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, +@@ -1131,9 +1208,89 @@ class MambaSpecDecodeGPUContext: + f"all mamba block tables must share stride(0), got {strides}" + ) + self.block_table_stride_req = int(next(iter(strides))) ++ # Bounds-guard instrumentation: row bound for the guard in _copy_mamba_state_block. ++ # ++ # NOT from `block_tables`. Every caller passes a per-batch VIEW -- ++ # get_device_tensor(num_reqs) at :1690 and :1838, a num_reqs_padded ++ # slice at mamba_hybrid.py:196 -- and this function runs once, so a ++ # bound taken from shape[0] here freezes at the first batch's request ++ # count and rejects every higher row in every later batch. The guard ++ # would then skip real state copies and report masked faults that never ++ # happened. (stride(0) is unaffected: it is max_num_blocks and survives ++ # a row slice, which is why the existing stride assert never caught ++ # this.) ++ # ++ # The persistent tables ARE [max_num_reqs, ...], so use them when they ++ # are registered. When they are not there is no correct row count at ++ # this point, and a wrong bound is worse than a loose one: fall back to ++ # a sentinel that still rejects negative rows -- the case the poison ++ # guard row exercises and the only one that reads outside the tensor -- ++ # while never rejecting a real request slot. ++ rows = {bt.shape[0] for bt in block_tables} ++ assert len(rows) == 1, ( ++ f"all mamba block tables must share shape[0], got {rows}" ++ ) ++ self.block_table_num_rows = _UNBOUNDED_ROWS + for i, bt in enumerate(block_tables): + self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) + ++ # Same capture for the persistent per-req-slot tables. `input_block_tables` ++ # is `torch.zeros_like` of these (block_table.py:74-76), so stride(0) is ++ # identical and `block_table_stride_req` covers both; assert it rather ++ # than assume. ++ if persistent_block_tables is not None: ++ assert len(persistent_block_tables) == self.num_groups, ( ++ f"expected {self.num_groups} persistent block tables, " ++ f"got {len(persistent_block_tables)}" ++ ) ++ persistent_strides = {bt.stride(0) for bt in persistent_block_tables} ++ assert persistent_strides == strides, ( ++ "persistent block tables must share stride(0) with the gathered " ++ f"tables, got {persistent_strides} vs {strides}" ++ ) ++ assert self.persistent_block_table_ptrs is not None ++ for i, bt in enumerate(persistent_block_tables): ++ self.persistent_block_table_ptrs[i] = _reinterpret_u64_as_i64( ++ bt.data_ptr() ++ ) ++ # The align paths index the PERSISTENT tables by req_idx while ++ # every other geometry scalar is measured from the GATHERED ones ++ # (by design). Two different scalars, two different treatments: ++ # ++ # * the STRIDE is load-bearing -- it is the multiplier in ++ # `block_table_typed + bt_row_idx * block_table_stride_req`, so ++ # a mismatch reads the wrong row entirely. It is ASSERTED equal ++ # just above; there is no correct behaviour if it differs. ++ # * the ROW COUNT is not required to be equal (the fixture in ++ # test_pp_deferred_mamba_block_table.py builds a 2-row gathered ++ # table against an 8-row persistent one on purpose), so it is ++ # CARRIED separately. Anything bounding a row against the ++ # persistent tables must use this scalar; using ++ # block_table_num_rows would falsely reject a valid request slot ++ # whenever the gathered buffer is narrower, and a guard that ++ # fires on correct metadata is worse than no guard. ++ persistent_rows = {bt.shape[0] for bt in persistent_block_tables} ++ assert len(persistent_rows) == 1, ( ++ f"all persistent block tables must share shape[0], got " ++ f"{persistent_rows}" ++ ) ++ self.persistent_block_table_num_rows = int(next(iter(persistent_rows))) ++ # ...and this is the only trustworthy row count in this object, so ++ # the same-step guard takes it too. ++ self.block_table_num_rows = self.persistent_block_table_num_rows ++ self.has_persistent_block_tables = True ++ else: ++ # Legal for V1 callers, which never reach the V2 align ++ # post-process. If a V2 align launch does happen on such a context, ++ # `run_fused_postprocess_align` raises rather than falling back to ++ # the batch-order gather: that fallback is exactly the earlier ++ # wrong-row bug. ++ logger.warning_once( ++ "Mamba align context initialized without the persistent " ++ "per-req-slot block tables. The V1 paths are unaffected; a V2 " ++ "align post-process launch on this context will raise." ++ ) ++ + self.is_initialized = True + + def compute_aligned_state_indices( +@@ -1215,6 +1372,10 @@ class MambaSpecDecodeGPUContext: + num_draft_tokens_gpu, + self.block_table_ptrs, + self.block_table_stride_req, ++ self.block_table_num_rows, ++ guard_telemetry.counter( ++ guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device ++ ), + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, +@@ -1261,6 +1422,10 @@ class MambaSpecDecodeGPUContext: + token_bias_gpu, + self.block_table_ptrs, + self.block_table_stride_req, ++ self.block_table_num_rows, ++ guard_telemetry.counter( ++ guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device ++ ), + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, +@@ -1340,8 +1505,12 @@ class MambaSpecDecodeGPUContext: + None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED + new_num_computed_tokens_gpu, + None, # num_draft: unused under PRECOMPUTED_NEW_COMPUTED +- self.block_table_ptrs, ++ block_table_ptrs, + self.block_table_stride_req, ++ self.persistent_block_table_num_rows, ++ guard_telemetry.counter( ++ guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device ++ ), + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, +-- +2.43.0 + diff --git a/patches/0010-gdn-bound-the-decode-state-index.patch b/patches/0010-gdn-bound-the-decode-state-index.patch new file mode 100644 index 0000000..afc7466 --- /dev/null +++ b/patches/0010-gdn-bound-the-decode-state-index.patch @@ -0,0 +1,178 @@ +From 25dd4897d4378bca33e56a2f332723d599ec3e57 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:07 -0400 +Subject: [PATCH 010/103] gdn: bound the decode state index + +Ports the check from upstream vLLM PR #50021. +--- + .../test_fused_sigmoid_gating_delta_rule.py | 81 +++++++++++++++++++ + .../ops/fused_sigmoid_gating.py | 42 ++++++++-- + 2 files changed, 118 insertions(+), 5 deletions(-) + +diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +index fc923cb..7d2e399 100644 +--- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py ++++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +@@ -197,3 +197,84 @@ def test_fused_sigmoid_gating_delta_rule_update_spec( + torch.testing.assert_close( + last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2 + ) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.parametrize("num_accepted", [0, 5]) ++def test_spec_invalid_accepted_count_zeroes_output(num_accepted: int) -> None: ++ """Invalid active-request state selection must not expose empty output. ++ ++ Ported from upstream vLLM PR #50021, covering BOTH kernels: the decode guard ++ (fused_sigmoid_gating) and the recurrent guard (fused_recurrent). A prior version dropped the ++ recurrent arm because that kernel was unguarded; it now has the ++ guard rather than the caveat, so the arm is back. ++ ++ num_accepted 0 gives i_t == -1, a read before this request's row; ++ num_accepted 5 exceeds the one-column ssm_state_indices, a read past it. ++ ``fill_uninitialized_memory`` makes the old behaviour visible: the kernel ++ returned early without writing ``o``, which is ``q.new_empty``. ++ """ ++ torch.set_default_device(DEVICE) ++ set_random_seed(0) ++ ++ num_tokens, num_k_heads, num_v_heads = 4, 2, 4 ++ head_k_dim = head_v_dim = 16 ++ query = torch.rand(1, num_tokens, num_k_heads, head_k_dim) ++ key = torch.rand_like(query) ++ value = torch.rand(1, num_tokens, num_v_heads, head_v_dim) ++ a = torch.rand(num_tokens, num_v_heads) ++ b = torch.rand_like(a) ++ A_log = torch.rand(num_v_heads) ++ dt_bias = torch.rand_like(A_log) ++ initial_state = torch.rand(num_tokens + 1, num_v_heads, head_v_dim, head_k_dim) ++ state_indices = torch.arange(1, num_tokens + 1, dtype=torch.int32).view( ++ 1, num_tokens ++ ) ++ cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32) ++ accepted = torch.tensor([num_accepted], dtype=torch.int32) ++ ++ previous_deterministic = torch.are_deterministic_algorithms_enabled() ++ previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() ++ previous_fill = torch.utils.deterministic.fill_uninitialized_memory ++ torch.use_deterministic_algorithms(True) ++ torch.utils.deterministic.fill_uninitialized_memory = True ++ try: ++ recurrent_state = initial_state.clone() ++ recurrent_out, recurrent_state = fused_recurrent_gated_delta_rule( ++ q=query, ++ k=key, ++ v=value, ++ g=torch.rand(1, num_tokens, num_v_heads), ++ beta=torch.rand(1, num_tokens, num_v_heads), ++ initial_state=recurrent_state, ++ inplace_final_state=True, ++ ssm_state_indices=state_indices, ++ cu_seqlens=cu_seqlens, ++ num_accepted_tokens=accepted, ++ ) ++ ++ sigmoid_state = initial_state.clone() ++ sigmoid_out, sigmoid_state = fused_sigmoid_gating_delta_rule_update( ++ A_log=A_log, ++ a=a, ++ b=b, ++ dt_bias=dt_bias, ++ q=query, ++ k=key, ++ v=value, ++ initial_state=sigmoid_state, ++ inplace_final_state=True, ++ ssm_state_indices=state_indices, ++ cu_seqlens=cu_seqlens, ++ num_accepted_tokens=accepted, ++ ) ++ finally: ++ torch.utils.deterministic.fill_uninitialized_memory = previous_fill ++ torch.use_deterministic_algorithms( ++ previous_deterministic, warn_only=previous_warn_only ++ ) ++ ++ torch.testing.assert_close(recurrent_out, torch.zeros_like(recurrent_out)) ++ torch.testing.assert_close(sigmoid_out, torch.zeros_like(sigmoid_out)) ++ torch.testing.assert_close(recurrent_state, initial_state) ++ torch.testing.assert_close(sigmoid_state, initial_state) +diff --git a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py +index 7e0c7e0..c5e75ec 100644 +--- a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py ++++ b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py +@@ -10,6 +10,14 @@ + import torch + + from vllm.triton_utils import tl, triton ++from vllm.v1.worker import guard_telemetry ++ ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + + + @triton.heuristics( +@@ -37,6 +45,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + scale, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens +@@ -106,12 +115,32 @@ def fused_sigmoid_gating_delta_rule_update_kernel( + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 +- # Load state index and check for invalid entries +- state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( +- tl.int64 +- ) +- # Skip if state index is invalid (NULL_BLOCK_ID=0) ++ # PR #50021: i_t is num_accepted_tokens - 1 and is otherwise ++ # unbounded against a stride_indices_seq-column tensor. A zero ++ # accepted count gives i_t == -1, a read before this request's row ++ # (and before the tensor for i_n == 0); a stale or too-large count ++ # reads past the row. The state_idx <= 0 guard below rejects only ++ # non-positive values, so an out-of-range read returning a garbage ++ # positive int flowed into h0 + state_idx * stride_init_state_token ++ # and faulted the SM. ++ idx_in_row = (i_t >= 0) & (i_t < stride_indices_seq) ++ if (i_t < 0) or (i_t >= stride_indices_seq): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (i_n).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (i_t).to(tl.int64) + _GUARD_BIAS) ++ state_idx = tl.load( ++ ssm_state_indices + i_n * stride_indices_seq + i_t, ++ mask=idx_in_row, ++ other=0, ++ ).to(tl.int64) ++ # Skip invalid state indices without exposing uninitialized output. ++ # o is q.new_empty in the entry below, so the previous early return ++ # left this request's rows holding whatever the allocator returned. + if state_idx <= 0: ++ zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty) ++ for _ in range(0, T): ++ tl.store(p_o, zero, mask=mask_v) ++ p_o += HV * V + return + p_h0 = h0 + state_idx * stride_init_state_token + else: +@@ -255,6 +284,9 @@ def fused_sigmoid_gating_delta_rule_update( + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, ++ guard_ctr=guard_telemetry.counter( ++ guard_telemetry.SITE_GDN_DECODE_STATE, q.device ++ ), + scale=scale, + N=N, + T=T, +-- +2.43.0 + diff --git a/patches/0011-qsa-bound-the-pre-indexer-request-row.patch b/patches/0011-qsa-bound-the-pre-indexer-request-row.patch new file mode 100644 index 0000000..51dd698 --- /dev/null +++ b/patches/0011-qsa-bound-the-pre-indexer-request-row.patch @@ -0,0 +1,75 @@ +From b9742140206908ce1fecb5026ccb0a981c8775c0 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:07 -0400 +Subject: [PATCH 011/103] qsa: bound the pre-indexer request row + +--- + .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 24 +++++++++++++++++++ + 1 file changed, 24 insertions(+) + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +index d996f39..f682ee9 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +@@ -5,6 +5,14 @@ + import torch + + from vllm.triton_utils import tl, triton ++from vllm.v1.worker import guard_telemetry ++ ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + + + @triton.jit +@@ -86,6 +94,7 @@ def _norm_rope( + "num_state_blocks", + "num_compressed_blocks", + "num_k_work", ++ "num_requests", + ] + ) + def _qsa_pre_indexer_kernel( +@@ -120,6 +129,8 @@ def _qsa_pre_indexer_kernel( + num_state_blocks, + num_compressed_blocks, + num_k_work, ++ num_requests, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, +@@ -200,6 +211,17 @@ def _qsa_pre_indexer_kernel( + request, work_in_request = tl.split(work_metadata) + if request < 0: + return ++ # Second bounds guard: the early return above rejects only the -1 ++ # sentinel the builder writes for inactive work items. Any other ++ # corrupted value was loaded straight out of query_start_loc, which has ++ # num_requests + 1 entries -- 8 bytes for a single-request batch. The ++ # chunk_end load below derives its index from query_end, so a bad row ++ # also walked logical_positions; returning here stops both. ++ if request >= num_requests: ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (work_in_request).to(tl.int64) + _GUARD_BIAS) ++ return + + query_start = tl.load(query_start_loc_ptr + request) + query_end = tl.load(query_start_loc_ptr + request + 1) +@@ -489,6 +511,8 @@ def qsa_pre_indexer( + state_cache.shape[0], + compressed_cache.shape[0], + num_k_work, ++ min(query_start_loc.numel() - 1, state_block_table.shape[0]), ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_PRE_INDEXER, q.device), + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, +-- +2.43.0 + diff --git a/patches/0012-qsa-bound-the-indexer-page-table-values.patch b/patches/0012-qsa-bound-the-indexer-page-table-values.patch new file mode 100644 index 0000000..83e6405 --- /dev/null +++ b/patches/0012-qsa-bound-the-indexer-page-table-values.patch @@ -0,0 +1,182 @@ +From 612e3d64642c5e3a0edba7d37d53f3b2ae9cd8da Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:08 -0400 +Subject: [PATCH 012/103] qsa: bound the indexer page-table values + +--- + .../qwen4_exp/nvidia/ops/qsa_indexer.py | 76 ++++++++++++++++++- + 1 file changed, 72 insertions(+), 4 deletions(-) + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +index 9ddd65d..a46f915 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +@@ -10,6 +10,14 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + ) + from vllm.platforms import current_platform + from vllm.triton_utils import tl, triton ++from vllm.v1.worker import guard_telemetry ++ ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + + _TOPK_WORKSPACE_BYTES = 1024 * 1024 + _DECODE_BLOCK_N = 64 +@@ -28,6 +36,8 @@ def _qsa_mqa_paged_uniform_kernel( + stride_cache_token, + stride_table_req, + stride_logits_row, ++ num_cache_blocks, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, +@@ -72,6 +82,11 @@ def _qsa_mqa_paged_uniform_kernel( + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) ++ # Bounds-guard instrumentation: accumulate in registers inside the pipelined loop ++ # and do the cross-lane reduction + atomics once after it, so an ++ # atomic-bearing conditional does not sit in the loop body. ++ guard_bad_n = tl.zeros([BLOCK_N], dtype=tl.int32) ++ guard_bad_col = tl.full([BLOCK_N], -1, tl.int32) + for tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < max_visible +@@ -80,11 +95,22 @@ def _qsa_mqa_paged_uniform_kernel( + physical_page = tl.load( + page_table_ptr + request * stride_table_req + logical_page, + mask=live, +- other=0, ++ other=-1, + ) ++ # Mirror of the QSA attention kernel (qsa.py:137): the VALUE read out of ++ # the page table is untrusted and must be bounded before it becomes an ++ # address. Folds into the existing live mask. ++ page_bad = live & ((physical_page < 0) | (physical_page >= num_cache_blocks)) ++ guard_bad_n += page_bad.to(tl.int32) ++ guard_bad_col = tl.maximum( ++ guard_bad_col, tl.where(page_bad, logical_page, -1).to(tl.int32) ++ ) ++ live = live & (physical_page >= 0) & (physical_page < num_cache_blocks) ++ # physical_page * block stride can overflow int32 for large caches. ++ safe_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr +- + physical_page[:, None].to(tl.int64) * stride_cache_block ++ + safe_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), +@@ -106,6 +132,15 @@ def _qsa_mqa_paged_uniform_kernel( + & (columns[:, None] < visible[None, :]), + ) + ++ # Bounds-guard instrumentation: one reduction and at most one atomic per program. ++ guard_n = tl.sum(guard_bad_n, axis=0) ++ if guard_n > 0: ++ tl.atomic_add(guard_ctr, guard_n.to(tl.int64)) ++ tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max( ++ guard_ctr + 2, (tl.max(guard_bad_col, axis=0)).to(tl.int64) + _GUARD_BIAS ++ ) ++ + + @triton.jit(do_not_specialize=["num_rows", "query_offset"]) + def _qsa_mqa_paged_prefill_kernel( +@@ -124,6 +159,8 @@ def _qsa_mqa_paged_prefill_kernel( + num_rows, + query_offset, + page_table_width, ++ num_cache_blocks, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + PAGE_SIZE: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, +@@ -181,6 +218,11 @@ def _qsa_mqa_paged_prefill_kernel( + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) ++ # Bounds-guard instrumentation: accumulate in registers inside the pipelined loop ++ # and do the cross-lane reduction + atomics once after it, so an ++ # atomic-bearing conditional does not sit in the loop body. ++ guard_bad_n = tl.zeros([BLOCK_N], dtype=tl.int32) ++ guard_bad_col = tl.full([BLOCK_N], -1, tl.int32) + for tile in tl.range(k_tile_start, k_tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < max_visible +@@ -189,11 +231,22 @@ def _qsa_mqa_paged_prefill_kernel( + physical_page = tl.load( + page_table_ptr + request * stride_table_req + logical_page, + mask=live, +- other=0, ++ other=-1, ++ ) ++ # Mirror of the QSA attention kernel (qsa.py:137): the VALUE read out of ++ # the page table is untrusted and must be bounded before it becomes an ++ # address. Folds into the existing live mask. ++ page_bad = live & ((physical_page < 0) | (physical_page >= num_cache_blocks)) ++ guard_bad_n += page_bad.to(tl.int32) ++ guard_bad_col = tl.maximum( ++ guard_bad_col, tl.where(page_bad, logical_page, -1).to(tl.int32) + ) ++ live = live & (physical_page >= 0) & (physical_page < num_cache_blocks) ++ # physical_page * block stride can overflow int32 for large caches. ++ safe_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr +- + physical_page[:, None].to(tl.int64) * stride_cache_block ++ + safe_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), +@@ -217,6 +270,15 @@ def _qsa_mqa_paged_prefill_kernel( + mask=store_mask, + ) + ++ # Bounds-guard instrumentation: one reduction and at most one atomic per program. ++ guard_n = tl.sum(guard_bad_n, axis=0) ++ if guard_n > 0: ++ tl.atomic_add(guard_ctr, guard_n.to(tl.int64)) ++ tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max( ++ guard_ctr + 2, (tl.max(guard_bad_col, axis=0)).to(tl.int64) + _GUARD_BIAS ++ ) ++ + + @triton.jit + def _expand_qsa_indices_kernel( +@@ -360,6 +422,8 @@ def warmup_qsa_mqa_paged_decode( + k_cache.stride(1), + page_table.stride(0), + columns, ++ k_cache.shape[0], ++ TritonWarmupTensor(torch.int64), + PAGE_SIZE=page_size, + PAGE_TABLE_WIDTH=page_table_width, + NUM_HEADS=num_heads, +@@ -423,6 +487,8 @@ def _prefill_logits( + num_queries, + query_offset, + page_table.shape[1], ++ k_cache.shape[0], ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_PREFILL, q.device), + PAGE_SIZE=k_cache.shape[1], + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], +@@ -551,6 +617,8 @@ def qsa_select_paged_decode( + *k_cache.stride()[:2], + *page_table.stride()[:-1], + *logits.stride()[:-1], ++ k_cache.shape[0], ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_DECODE, q.device), + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + NUM_HEADS=q.shape[1], +-- +2.43.0 + diff --git a/patches/0013-worker-validate-req_idx-in-the-block-table-gather-an.patch b/patches/0013-worker-validate-req_idx-in-the-block-table-gather-an.patch new file mode 100644 index 0000000..1671ad4 --- /dev/null +++ b/patches/0013-worker-validate-req_idx-in-the-block-table-gather-an.patch @@ -0,0 +1,263 @@ +From 430b5ec450421d89201331ebe04450d135e75517 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:08 -0400 +Subject: [PATCH 013/103] worker: validate req_idx in the block-table gather + and the slot mapping + +--- + .../v1/worker/gpu/test_block_table_guards.py | 118 ++++++++++++++++++ + vllm/v1/worker/gpu/block_table.py | 62 +++++++++ + 2 files changed, 180 insertions(+) + create mode 100644 tests/v1/worker/gpu/test_block_table_guards.py + +diff --git a/tests/v1/worker/gpu/test_block_table_guards.py b/tests/v1/worker/gpu/test_block_table_guards.py +new file mode 100644 +index 0000000..d15093c +--- /dev/null ++++ b/tests/v1/worker/gpu/test_block_table_guards.py +@@ -0,0 +1,118 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Real launches of the three kernels no suite was launching. ++ ++An early bounds-guard boot died in warmup with a Triton CompilationError in ++`_gather_block_tables_kernel` -- a branch-typing error in guard 6 that nothing ++had ever compiled. The compile check listed it as "left to the suites" and ++`tests/v1/worker/gpu/` contained no block-table test at all, so it was left to ++nobody: its first compile was the first real forward of a four-rank boot. ++ ++These tests launch the real kernels. They are cheap, they are CUDA-gated, and ++their only job is that a signature or typing error in guard 6 or guard 7 costs ++seconds in the suite rather than a boot. ++""" ++ ++import pytest ++import torch ++ ++from vllm.v1.attention.backends.utils import PAD_SLOT_ID ++from vllm.v1.worker import guard_telemetry ++from vllm.v1.worker.gpu.block_table import BlockTables ++ ++pytestmark = pytest.mark.skipif( ++ not torch.cuda.is_available(), reason="CUDA required" ++) ++ ++MAX_REQS = 8 ++BLOCKS = 4 ++BLOCK_SIZE = 16 ++NUM_REQS = 4 ++NUM_TOKENS = 16 ++ ++ ++def _tables(device): ++ bt = BlockTables( ++ block_sizes=[BLOCK_SIZE], ++ max_num_reqs=MAX_REQS, ++ max_num_batched_tokens=NUM_TOKENS, ++ max_num_blocks_per_group=[BLOCKS], ++ device=device, ++ kernel_block_sizes=[BLOCK_SIZE], ++ ) ++ for r in range(NUM_REQS): ++ bt.append_block_ids(r, ([1 + r],), True) ++ bt.apply_staged_writes() ++ return bt ++ ++ ++def test_gather_and_slot_mapping_launch_and_are_correct(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ bt = _tables(device) ++ ++ idx = torch.arange(NUM_REQS, dtype=torch.int32, device=device) ++ (gathered,) = bt.gather_block_tables(idx, num_reqs_padded=NUM_REQS) ++ torch.cuda.synchronize() ++ for r in range(NUM_REQS): ++ assert int(gathered[r, 0]) == 1 + r, gathered ++ ++ qsl = torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32, device=device) ++ pos = torch.arange(NUM_TOKENS, dtype=torch.int64, device=device) ++ slots = bt.compute_slot_mappings(idx, qsl, pos, num_tokens_padded=NUM_TOKENS) ++ torch.cuda.synchronize() ++ assert slots.shape[-1] == NUM_TOKENS ++ ++ # A legal batch must leave both counters at zero. A guard that fires here ++ # is reporting a masked fault on correct metadata, which is worse than no ++ # guard -- see the row-bound defect. ++ for site in (guard_telemetry.SITE_BT_GATHER, ++ guard_telemetry.SITE_SLOT_MAPPING): ++ ctr = guard_telemetry.counter(site, device) ++ assert int(ctr[0].item()) == 0, guard_telemetry.snapshot() ++ ++ ++def test_gather_guard_fires_and_zeroes_the_row(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ bt = _tables(device) ++ ++ idx = torch.arange(NUM_REQS, dtype=torch.int32, device=device) ++ idx[2] = MAX_REQS + 3 # a request row that does not exist ++ (gathered,) = bt.gather_block_tables(idx, num_reqs_padded=NUM_REQS) ++ torch.cuda.synchronize() ++ ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_BT_GATHER, device) ++ assert int(ctr[0].item()) > 0, "the gather guard did not fire" ++ row = int(ctr[1].item()) - guard_telemetry.GUARD_BIAS ++ val = int(ctr[2].item()) - guard_telemetry.GUARD_BIAS ++ assert row == 2, f"recorded batch row {row}, expected 2" ++ assert val == MAX_REQS + 3, f"recorded req_idx {val}" ++ # and the destination row is zeroed rather than filled from a stray row ++ assert int(gathered[2].abs().sum()) == 0, gathered[2] ++ # the rows either side are untouched ++ assert int(gathered[1, 0]) == 2 and int(gathered[3, 0]) == 4 ++ ++ ++def test_slot_mapping_guard_pads_the_bad_request(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ bt = _tables(device) ++ ++ idx = torch.arange(NUM_REQS, dtype=torch.int32, device=device) ++ idx[1] = -1 # the freed-slot sentinel ++ qsl = torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32, device=device) ++ pos = torch.arange(NUM_TOKENS, dtype=torch.int64, device=device) ++ slots = bt.compute_slot_mappings(idx, qsl, pos, num_tokens_padded=NUM_TOKENS) ++ torch.cuda.synchronize() ++ ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_SLOT_MAPPING, device) ++ assert int(ctr[0].item()) > 0, "the slot-mapping guard did not fire" ++ # request 1 owns tokens 4..7 and they must all be PAD_ID ++ assert torch.equal( ++ slots[0, 4:8], ++ torch.full((4,), PAD_SLOT_ID, dtype=slots.dtype, device=device), ++ ), slots[0, :12] +diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py +index c01895e..6cbd4eb 100644 +--- a/vllm/v1/worker/gpu/block_table.py ++++ b/vllm/v1/worker/gpu/block_table.py +@@ -6,6 +6,7 @@ import torch + + from vllm.triton_utils import tl, triton + from vllm.v1.attention.backends.utils import PAD_SLOT_ID ++from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.gpu.buffer_utils import ( + FusedStagedWriter, + StagedWriteTensor, +@@ -13,6 +14,12 @@ from vllm.v1.worker.gpu.buffer_utils import ( + _load_ptr, + ) + ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + + class BlockTables: + def __init__( +@@ -171,6 +178,9 @@ class BlockTables: + self.num_blocks.gpu, + self.num_blocks.gpu.stride(0), + num_reqs, ++ guard_telemetry.counter( ++ guard_telemetry.SITE_BT_GATHER, idx_mapping.device ++ ), + BLOCK_SIZE=1024, # type: ignore + ) + return tuple(bt[:num_reqs_padded] for bt in out) +@@ -208,6 +218,10 @@ class BlockTables: + positions, + self.block_table_ptrs, + self.block_table_strides, ++ self.num_blocks.gpu.stride(0), ++ guard_telemetry.counter( ++ guard_telemetry.SITE_SLOT_MAPPING, idx_mapping.device ++ ), + self.block_sizes_tensor, + self.kernel_block_sizes_tensor, + self.slot_mapping_enabled, +@@ -242,6 +256,7 @@ def _gather_block_tables_kernel( + num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] + num_blocks_stride, + num_reqs, # actual number of requests (for padding) ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + BLOCK_SIZE: tl.constexpr, + ): + # kv cache group id +@@ -261,8 +276,39 @@ def _gather_block_tables_kernel( + return + + req_idx = tl.load(batch_idx_to_req_idx + batch_idx) ++ # Bounds-guard instrumentation (guard 3): req_idx indexes both num_blocks_ptr and the ++ # persistent source table, which have max_num_reqs rows. num_blocks_stride ++ # IS max_num_reqs, so the bound needs no new argument. Zero the destination ++ # row and return, mirroring the padded-row path above. ++ if (req_idx < 0) or (req_idx >= num_blocks_stride): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (batch_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (req_idx).to(tl.int64) + _GUARD_BIAS) ++ for i in tl.range(0, max_num_blocks, BLOCK_SIZE): ++ offset = i + tl.arange(0, BLOCK_SIZE) ++ tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks) ++ return ++ + group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride + num_blocks = tl.load(group_num_blocks_ptr + req_idx) ++ # A row count wider than the table would copy past the end of both rows. ++ # ++ # Clamp UNCONDITIONALLY and record conditionally. Doing the clamp inside ++ # the `if` rebound an int32 name (num_blocks, loaded from num_blocks_ptr) ++ # to int64 (max_num_blocks is `stride`, loaded from block_table_strides), ++ # and Triton requires a name to keep its type across a branch: ++ # "initial value for `num_blocks` is of type int32[], but the then block ++ # redefines it as int64[]" ++ # That killed an early boot in warmup, at the first real gather. The ++ # explicit .to(tl.int32) pins the type rather than relying on whatever the ++ # operands promote to. ++ if (num_blocks < 0) or (num_blocks > max_num_blocks): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (batch_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (num_blocks).to(tl.int64) + _GUARD_BIAS) ++ num_blocks = tl.minimum( ++ tl.maximum(num_blocks, 0), max_num_blocks ++ ).to(tl.int32) + + src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32) + src_row_ptr = src_block_table_ptr + req_idx * stride +@@ -281,6 +327,8 @@ def _compute_slot_mappings_kernel( + pos, # [num_tokens] + block_table_ptrs, # [num_kv_cache_groups] + block_table_strides, # [num_kv_cache_groups] ++ bt_num_rows, # bounds-guard instrumentation: rows in the persistent block tables ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + block_sizes, # [num_kv_cache_groups] + kernel_block_sizes, # [num_kv_cache_groups] + slot_mapping_enabled, # [num_kv_cache_groups] +@@ -315,6 +363,20 @@ def _compute_slot_mappings_kernel( + mapping_enabled = tl.load(slot_mapping_enabled + group_id) + + req_state_idx = tl.load(idx_mapping + batch_idx) ++ # Bounds-guard instrumentation (guard 4): req_state_idx indexes the persistent block table ++ # before any bound is applied. An invalid row writes PAD_ID for this ++ # request rather than reading an arbitrary row. ++ if (req_state_idx < 0) or (req_state_idx >= bt_num_rows): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (batch_idx).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (req_state_idx).to(tl.int64) + _GUARD_BIAS) ++ bad_start = tl.load(query_start_loc + batch_idx) ++ bad_end = tl.load(query_start_loc + batch_idx + 1) ++ for i in range(bad_start, bad_end, TRITON_BLOCK_SIZE): ++ offset = i + tl.arange(0, TRITON_BLOCK_SIZE) ++ tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < bad_end) ++ return ++ + start_idx = tl.load(query_start_loc + batch_idx) + end_idx = tl.load(query_start_loc + batch_idx + 1) + for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): +-- +2.43.0 + diff --git a/patches/0014-worker-pool-bound-the-KV-block-zeroer.patch b/patches/0014-worker-pool-bound-the-KV-block-zeroer.patch new file mode 100644 index 0000000..f7d4f34 --- /dev/null +++ b/patches/0014-worker-pool-bound-the-KV-block-zeroer.patch @@ -0,0 +1,81 @@ +From 85f78b1410981293e030f6256c0c16c84e1a12d7 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:08 -0400 +Subject: [PATCH 014/103] worker: pool-bound the KV block zeroer + +--- + vllm/v1/worker/utils.py | 26 ++++++++++++++++++++++++++ + 1 file changed, 26 insertions(+) + +diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py +index df2ebc7..71c1a2d 100644 +--- a/vllm/v1/worker/utils.py ++++ b/vllm/v1/worker/utils.py +@@ -39,10 +39,17 @@ from vllm.v1.kv_cache_interface import ( + UniformTypeKVCacheSpecs, + create_kv_cache_views, + ) ++from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.block_table import get_block_table_width + + logger = init_logger(__name__) + ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not readable ++# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died ++# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++ + + def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None: + if not any(num_nans_in_logits.values()): +@@ -62,6 +69,8 @@ def _zero_kv_blocks_kernel( + seg_block_strides_ptr, + seg_page_sizes_ptr, + block_ids_ptr, ++ num_pool_blocks, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + BLOCK_SIZE: tl.constexpr, + ): + """Zero KV cache blocks across all segments in a single launch. +@@ -91,6 +100,17 @@ def _zero_kv_blocks_kernel( + if chunk_offset >= page_size_el: + return + block_id = tl.load(block_ids_ptr + block_index) ++ # Bounds-guard instrumentation (guard 5): block_id goes straight into ++ # block_id * block_stride_el with no pool bound, and this kernel masks only ++ # WITHIN a segment page, so a corrupt id writes zeros elsewhere in the KV ++ # allocation instead of faulting visibly. A valid pool id is always in ++ # range (every group's tensor is asserted to be an exact multiple of ++ # kv_cache_config.num_blocks), so this can only fire on corruption. ++ if (block_id < 0) or (block_id >= num_pool_blocks): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (block_index).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (block_id).to(tl.int64) + _GUARD_BIAS) ++ return + seg_addr = tl.load(seg_addrs_ptr + seg_index) + ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) + block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64) +@@ -138,6 +158,8 @@ class KVBlockZeroer: + Only AttentionSpec layers are processed; Mamba layers are skipped. + """ + self.device = device ++ # Bounds-guard instrumentation: pool bound for the guard in _zero_kv_blocks_kernel. ++ self._num_pool_blocks = int(num_blocks) + self._meta: ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None + ) = None +@@ -246,6 +268,10 @@ class KVBlockZeroer: + seg_block_strides, + seg_page_sizes, + idx, ++ self._num_pool_blocks, ++ guard_telemetry.counter( ++ guard_telemetry.SITE_KV_ZEROER, self.device ++ ), + BLOCK_SIZE=blk_size, + ) + +-- +2.43.0 + diff --git a/patches/0015-sched-optional-host-side-ownership-check-before-a-bl.patch b/patches/0015-sched-optional-host-side-ownership-check-before-a-bl.patch new file mode 100644 index 0000000..5b603b3 --- /dev/null +++ b/patches/0015-sched-optional-host-side-ownership-check-before-a-bl.patch @@ -0,0 +1,79 @@ +From 298ddf906275a3a3743d94f50021c755b658a39f Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:08 -0400 +Subject: [PATCH 015/103] sched: optional host-side ownership check before a + block is zeroed + +Host-side diagnostic behind VLLM_E2_ZERO_OWNERSHIP_CHECK, off by default. +--- + vllm/v1/core/sched/scheduler.py | 38 +++++++++++++++++++++++++++++++++ + 1 file changed, 38 insertions(+) + +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index daf2b31..9f2d99d 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -1,6 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import itertools ++import os + import time + from collections import defaultdict, deque + from collections.abc import Iterable +@@ -75,6 +76,10 @@ from vllm.v1.utils import record_function_or_nullcontext + + logger = init_logger(__name__) + ++# Bounds-guard instrumentation: verify KVBlockZeroer's exclusive-ownership precondition before a ++# block is zeroed. Off by default, so it costs nothing. ++_ZERO_OWNERSHIP_CHECK = os.getenv("VLLM_E2_ZERO_OWNERSHIP_CHECK") == "1" ++ + + class Scheduler(SchedulerInterface): + def __init__( +@@ -1472,8 +1477,41 @@ class Scheduler(SchedulerInterface): + new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip] + skip.clear() + ++ if new_block_ids_to_zero and _ZERO_OWNERSHIP_CHECK: ++ self._check_zero_block_ownership(new_block_ids_to_zero) ++ + return new_block_ids_to_zero or None + ++ def _check_zero_block_ownership(self, block_ids: list[int]) -> None: ++ """Test KVBlockZeroer's stated precondition, host side, off by default. ++ ++ worker/utils.py justifies zeroing a whole tile with "block IDs are ++ global pool indices and a newly allocated block owns its whole tile". ++ Ownership is visible only here, in the scheduler process, so the check ++ lives here rather than in the zeroer. A block about to be zeroed should ++ be referenced by exactly one live request. More than one means a stale ++ reference survived into the zeroing, and the kernel is about to destroy ++ state another request still points at -- the stale-reference shape. ++ ++ VLLM_E2_ZERO_OWNERSHIP_CHECK=1. O(live blocks) per step. ++ """ ++ pending = set(block_ids) ++ owners: dict[int, set[str]] = {} ++ for mgr in self.kv_cache_manager.coordinator.single_type_managers: ++ for req_id, blocks in mgr.req_to_blocks.items(): ++ for blk in blocks: ++ if blk.block_id in pending: ++ owners.setdefault(blk.block_id, set()).add(req_id) ++ for block_id, req_ids in sorted(owners.items()): ++ if len(req_ids) > 1: ++ logger.warning( ++ "KV zeroing ownership violation: block %d is about to be " ++ "zeroed but is referenced by %d live requests (%s).", ++ block_id, ++ len(req_ids), ++ ", ".join(sorted(req_ids)[:4]), ++ ) ++ + def _preempt_request( + self, request: Request, timestamp: float, drop_stale_output: bool = False + ) -> None: +-- +2.43.0 + diff --git a/patches/0016-gdn-bound-the-recurrent-state-index.patch b/patches/0016-gdn-bound-the-recurrent-state-index.patch new file mode 100644 index 0000000..b0edb0f --- /dev/null +++ b/patches/0016-gdn-bound-the-recurrent-state-index.patch @@ -0,0 +1,292 @@ +From cda271d5ea26ff3c03c5d202041f9a09552c5ef3 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:09 -0400 +Subject: [PATCH 016/103] gdn: bound the recurrent state index + +Both callers of the guarded kernel (fused recurrent and KDA) are updated in the +same commit. +--- + .../worker/gpu/test_guard_telemetry_replay.py | 167 ++++++++++++++++++ + .../ops/fused_recurrent.py | 44 ++++- + .../flash_linear_attention/ops/kda.py | 7 + + 3 files changed, 213 insertions(+), 5 deletions(-) + create mode 100644 tests/v1/worker/gpu/test_guard_telemetry_replay.py + +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_replay.py b/tests/v1/worker/gpu/test_guard_telemetry_replay.py +new file mode 100644 +index 0000000..52327ab +--- /dev/null ++++ b/tests/v1/worker/gpu/test_guard_telemetry_replay.py +@@ -0,0 +1,167 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The test that distinguishes "architecturally fixed" from "fixed". ++ ++The first bounds-guard instrumentation design allocated a counter per launch and read ++it back inline. Under CUDA-graph REPLAY the Python host code never runs, so ++that counter was re-zeroed every step and never read: the instrument was ++silently dead in exactly the configuration the crash happens in (QSA declares ++AttentionCGSupport.UNIFORM_BATCH, so the guarded kernels sit inside the ++captured forward). ++ ++The current design rebuilds it on a persistent buffer plus captured atomics. That is the ++right architecture, but "right architecture" is a claim about code. This is the ++measurement: capture a graph around a guarded kernel with a seeded bad value, ++replay it N times WITHOUT re-entering Python, and require the counter to have ++advanced N times. ++ ++A test that only ran the kernel eagerly would pass on the first design too, ++which is why it must be a replay. ++""" ++ ++import pytest ++import torch ++ ++from vllm.third_party.flash_linear_attention.ops import ( ++ fused_sigmoid_gating_delta_rule_update, ++) ++from vllm.v1.worker import guard_telemetry ++ ++pytestmark = pytest.mark.skipif( ++ not torch.cuda.is_available(), reason="CUDA required" ++) ++ ++REPLAYS = 5 ++ ++ ++def _seeded_call(device): ++ """One guarded launch whose guard MUST fire. ++ ++ num_accepted_tokens = 0 makes i_t == -1 in ++ fused_sigmoid_gating_delta_rule_update_kernel, a read before this request's ++ row of ssm_state_indices. That is the decode-state guard. ++ """ ++ num_tokens, num_k_heads, num_v_heads = 4, 2, 4 ++ head_k_dim = head_v_dim = 16 ++ q = torch.rand(1, num_tokens, num_k_heads, head_k_dim, device=device) ++ k = torch.rand_like(q) ++ v = torch.rand(1, num_tokens, num_v_heads, head_v_dim, device=device) ++ a = torch.rand(num_tokens, num_v_heads, device=device) ++ b = torch.rand_like(a) ++ A_log = torch.rand(num_v_heads, device=device) ++ dt_bias = torch.rand_like(A_log) ++ state = torch.rand( ++ num_tokens + 1, num_v_heads, head_v_dim, head_k_dim, device=device ++ ) ++ state_indices = torch.arange( ++ 1, num_tokens + 1, dtype=torch.int32, device=device ++ ).view(1, num_tokens) ++ cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) ++ accepted = torch.zeros(1, dtype=torch.int32, device=device) # -> i_t == -1 ++ ++ def run(): ++ fused_sigmoid_gating_delta_rule_update( ++ A_log=A_log, ++ a=a, ++ b=b, ++ dt_bias=dt_bias, ++ q=q, ++ k=k, ++ v=v, ++ initial_state=state, ++ inplace_final_state=True, ++ ssm_state_indices=state_indices, ++ cu_seqlens=cu_seqlens, ++ num_accepted_tokens=accepted, ++ ) ++ ++ return run ++ ++ ++def _count(device): ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) ++ torch.cuda.synchronize() ++ return int(ctr[0].item()) ++ ++ ++def test_guard_counts_under_cudagraph_replay(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ # Exactly what the model runner does, and for the same reason: allocate the ++ # counter before anything can capture. ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ ++ run = _seeded_call(device) ++ ++ # Warm up on a side stream, the standard pre-capture dance. This also ++ # compiles the kernel: a JIT compile inside a capture would abort it. ++ side = torch.cuda.Stream() ++ side.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(side): ++ for _ in range(3): ++ run() ++ torch.cuda.current_stream().wait_stream(side) ++ torch.cuda.synchronize() ++ ++ warm = _count(device) ++ assert warm > 0, ( ++ "the seeded fault did not fire the guard even eagerly -- the control is " ++ "broken, so the replay result below would mean nothing" ++ ) ++ per_launch = warm // 3 ++ assert per_launch > 0 ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ run() ++ # Capturing itself must not have aborted, and must not have counted. ++ torch.cuda.synchronize() ++ before = _count(device) ++ ++ for _ in range(REPLAYS): ++ graph.replay() ++ torch.cuda.synchronize() ++ after = _count(device) ++ ++ delta = after - before ++ assert delta == REPLAYS * per_launch, ( ++ f"counter advanced by {delta} over {REPLAYS} replays, expected " ++ f"{REPLAYS * per_launch}. A delta of 0 means the guard telemetry is " ++ f"dead under replay -- the replay defect is NOT fixed." ++ ) ++ ++ ++def test_counter_pointer_is_stable_across_replays(): ++ """The captured atomic addresses one fixed buffer. ++ ++ If counter() ever returned a fresh tensor, the captured graph would keep ++ writing to a freed address -- the other half of the replay defect, and the one that ++ corrupts memory rather than merely losing evidence. ++ """ ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ first = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) ++ second = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) ++ assert first.data_ptr() == second.data_ptr() ++ # and the unindexed alias must resolve to the same buffer ++ third = guard_telemetry.counter( ++ guard_telemetry.SITE_GDN_DECODE_STATE, torch.device("cuda") ++ ) ++ assert third.data_ptr() == first.data_ptr() ++ ++ # ...and it must survive the reader giving up on it. An ++ # earlier failure path popped the buffer, which drops the only durable ++ # reference: counter() hands out transient views, and under replay nothing ++ # holds one at all, so the storage would return to the allocator, be ++ # reused, and the CAPTURED atomics would keep writing to the baked address. ++ # Stop polling, never stop holding. ++ try: ++ guard_telemetry._read_disabled.add(device) ++ after = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) ++ assert after.data_ptr() == first.data_ptr(), ( ++ "the counter moved after the reader disabled it -- captured graphs " ++ "are now writing to freed storage" ++ ) ++ assert device in guard_telemetry._buffers ++ finally: ++ guard_telemetry._read_disabled.discard(device) +diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +index eb08b93..2379fcb 100644 +--- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py ++++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +@@ -11,6 +11,12 @@ + import torch + + from vllm.triton_utils import tl, triton ++from vllm.v1.worker import guard_telemetry ++ ++# Bias so a negative row/index survives atomic_max on a zeroed counter ++# -- it MUST be a tl.constexpr: a plain module-level int is not ++# readable from a @triton.jit body. ++_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) + + from .op import exp, log + +@@ -36,6 +42,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, ++ guard_ctr, # bounds-guard instrumentation: persistent int64[3] + scale, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens +@@ -106,12 +113,36 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 +- # Load state index and check for invalid entries +- state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( +- tl.int64 +- ) +- # Skip if state index is invalid (NULL_BLOCK_ID=0) ++ # PR #50021: i_t is num_accepted_tokens - 1 and is otherwise ++ # unbounded, while ssm_state_indices has only stride_indices_seq ++ # columns per request. A zero accepted count gives i_t == -1, a ++ # read before this request's row (and before the tensor for ++ # i_n == 0); a stale or too-large count reads past the row. The ++ # state_idx <= 0 guard below rejects only non-positive values, so ++ # an out-of-range read returning a garbage positive int flowed into ++ # h0 + state_idx * stride_init_state_token and faulted the SM. ++ # ++ # This is the SAME defect as the decode-state guard, in the twin kernel. An ++ # earlier revision left it unguarded and dropped the upstream test's arm for it; ++ # it is now ported rather than caveated. ++ idx_in_row = (i_t >= 0) & (i_t < stride_indices_seq) ++ if (i_t < 0) or (i_t >= stride_indices_seq): ++ tl.atomic_add(guard_ctr, 1) ++ tl.atomic_max(guard_ctr + 1, (i_n).to(tl.int64) + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, (i_t).to(tl.int64) + _GUARD_BIAS) ++ state_idx = tl.load( ++ ssm_state_indices + i_n * stride_indices_seq + i_t, ++ mask=idx_in_row, ++ other=0, ++ ).to(tl.int64) ++ # Skip invalid state indices without exposing uninitialized output. ++ # o is q.new_empty in the entry below, so the previous early return ++ # left this request's rows holding whatever the allocator returned. + if state_idx <= 0: ++ zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty) ++ for _ in range(0, T): ++ tl.store(p_o, zero, mask=mask_v) ++ p_o += HV * V + return + p_h0 = h0 + state_idx * stride_init_state_token + else: +@@ -227,6 +258,9 @@ def fused_recurrent_gated_delta_rule_fwd( + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, ++ guard_ctr=guard_telemetry.counter( ++ guard_telemetry.SITE_GDN_RECURRENT_STATE, q.device ++ ), + scale=scale, + N=N, + T=T, +diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py +index 548f600..5ba7213 100644 +--- a/vllm/third_party/flash_linear_attention/ops/kda.py ++++ b/vllm/third_party/flash_linear_attention/ops/kda.py +@@ -14,6 +14,7 @@ import torch.nn as nn + + from vllm.model_executor.custom_op import CustomOp + from vllm.triton_utils import tl, triton ++from vllm.v1.worker import guard_telemetry + from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + + from .chunk_delta_h import chunk_gated_delta_rule_fwd_h +@@ -81,6 +82,12 @@ def fused_recurrent_kda_fwd( + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, ++ # Bounds-guard instrumentation: the same kernel object as fused_recurrent.py:249, so it ++ # takes the same guard counter. Same site: the defect being guarded is ++ # the kernel's, not the caller's. ++ guard_ctr=guard_telemetry.counter( ++ guard_telemetry.SITE_GDN_RECURRENT_STATE, q.device ++ ), + scale=scale, + N=N, + T=T, +-- +2.43.0 + diff --git a/patches/0017-qsa-bound-the-RoPE-position-before-it-becomes-an-add.patch b/patches/0017-qsa-bound-the-RoPE-position-before-it-becomes-an-add.patch new file mode 100644 index 0000000..ef40d5d --- /dev/null +++ b/patches/0017-qsa-bound-the-RoPE-position-before-it-becomes-an-add.patch @@ -0,0 +1,444 @@ +From d90c541cf2da170e17a94cebaf19dbda03fda390 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:09 -0400 +Subject: [PATCH 017/103] qsa: bound the RoPE position before it becomes an + address + +--- + .../gpu/test_qsa_rope_position_guard.py | 301 ++++++++++++++++++ + .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 76 +++++ + 2 files changed, 377 insertions(+) + create mode 100644 tests/v1/worker/gpu/test_qsa_rope_position_guard.py + +diff --git a/tests/v1/worker/gpu/test_qsa_rope_position_guard.py b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +new file mode 100644 +index 0000000..9539aa7 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +@@ -0,0 +1,301 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Bounds guard 9: the RoPE position must be bounded before it is an address. ++ ++An earlier crash-replay harness drove the QSA pre-indexer ++over a re-claimed raw-key ring block and compute-sanitizer reported eleven ++`Invalid __global__ read of size 2 bytes` at `_norm_rope`, roughly 4.6e18 bytes ++out of range. The chain: the cross-chunk pooled group takes its RoPE position ++from the ring row itself (`qsa_pre_indexer.py:305-317`), the ring is excluded ++from zeroing, and `_norm_rope` turns that value straight into a row index into ++the cos/sin table. The masks in `_norm_rope` bound the rotary PAIR, never the ++row. ++ ++Two tests, because they answer different questions: ++ ++* the CPU one checks the PLUMBING -- that the bound handed to the kernel is the ++ cos/sin table's real row count. A guard with the wrong bound is worse than no ++ guard, and this is the failure mode a CPU suite can actually catch; ++* the GPU one is the POSITIVE CONTROL, and it is the crash-replay case verbatim: fill ++ a re-claimed ring row with bf16 filler, run the kernel, and require that it ++ survives AND that the counter says it masked something. Before this guard the ++ same input produced an illegal address. ++""" ++ ++import inspect ++ ++import pytest ++import torch ++ ++from vllm.models.qwen4_exp.nvidia.ops import qsa_pre_indexer as qpi ++from vllm.v1.worker import guard_telemetry ++ ++COMPRESS_RATIO = 4 ++HEAD_DIM = 128 ++N_HEADS = 4 ++STATE_SIZE = 8 # ring capacity ++ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++MAX_MODEL_LEN = 4096 ++NUM_BLOCKS = 8 ++ ++ ++class _Recorder: ++ """Stands in for the Triton kernel and keeps the launch arguments.""" ++ ++ def __init__(self): ++ self.launches = [] ++ ++ def __getitem__(self, grid): ++ def launch(*args, **kwargs): ++ self.launches.append((grid, args, kwargs)) ++ ++ return launch ++ ++ ++def _inputs(device, num_tokens=4): ++ q = torch.randn(num_tokens, N_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device) ++ k = torch.randn(num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ positions = ( ++ torch.arange(num_tokens, dtype=torch.int64, device=device) ++ .unsqueeze(0) ++ .expand(3, -1) ++ .contiguous() ++ ) ++ cos_sin = torch.randn( ++ MAX_MODEL_LEN, HEAD_DIM // 2, dtype=torch.bfloat16, device=device ++ ) ++ ones = torch.ones(HEAD_DIM, dtype=torch.bfloat16, device=device) ++ q_out = torch.empty(num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ state_cache = torch.zeros( ++ NUM_BLOCKS, STATE_SIZE, 1, ROW_WIDTH, dtype=torch.bfloat16, device=device ++ ) ++ state_slots = torch.arange(num_tokens, dtype=torch.int32, device=device) ++ state_table = torch.zeros(1, 1, dtype=torch.int32, device=device) ++ qsl = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) ++ logical = torch.arange(num_tokens, dtype=torch.int32, device=device) ++ comp_cache = torch.zeros( ++ NUM_BLOCKS, 800, 1, HEAD_DIM, dtype=torch.bfloat16, device=device ++ ) ++ comp_slots = torch.zeros(num_tokens, dtype=torch.int32, device=device) ++ kwb = torch.zeros( ++ max(num_tokens // COMPRESS_RATIO, 1), 2, dtype=torch.int32, device=device ++ ) ++ return dict( ++ q=q, k=k, positions=positions, cos_sin_cache=cos_sin, ++ q_norm_weight=ones, k_norm_weight=ones, eps=1e-6, q_out=q_out, ++ state_cache=state_cache, state_slots=state_slots, ++ state_block_table=state_table, query_start_loc=qsl, ++ logical_positions=logical, compressed_cache=comp_cache, ++ compressed_slots=comp_slots, k_work_metadata=kwb, ++ ) ++ ++ ++def test_launcher_passes_the_real_cos_sin_row_count(monkeypatch): ++ """The bound must be the table's row count, not a constant or a stride. ++ ++ A bound that is too large lets the very reads this guard exists to stop ++ through; one that is too small makes the guard fire on correct metadata, ++ which is worse than not having it. ++ """ ++ device = torch.device("cpu") ++ # grab the real signature BEFORE the kernel is replaced ++ sig = inspect.signature(qpi._qsa_pre_indexer_kernel.fn) ++ rec = _Recorder() ++ monkeypatch.setattr(qpi, "_qsa_pre_indexer_kernel", rec) ++ kw = _inputs(device) ++ qpi.qsa_pre_indexer( ++ kw.pop("q"), kw.pop("k"), kw.pop("positions"), kw.pop("cos_sin_cache"), ++ kw.pop("q_norm_weight"), kw.pop("k_norm_weight"), kw.pop("eps"), ++ kw.pop("q_out"), kw.pop("state_cache"), kw.pop("state_slots"), ++ kw.pop("state_block_table"), kw.pop("query_start_loc"), ++ kw.pop("logical_positions"), kw.pop("compressed_cache"), ++ kw.pop("compressed_slots"), kw.pop("k_work_metadata"), ++ compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ rope_pos_offset=HEAD_DIM, ++ ) ++ ++ assert len(rec.launches) == 1 ++ _grid, args, kwargs = rec.launches[0] ++ # Triton launch options are not kernel parameters. ++ kwargs = {k: v for k, v in kwargs.items() ++ if k not in ("num_warps", "num_stages", "num_ctas", "grid")} ++ bound = sig.bind(*args, **kwargs) ++ bound.apply_defaults() ++ assert bound.arguments["cos_sin_rows"] == MAX_MODEL_LEN ++ ++ # The rope guard gets its OWN site. A later revision threaded the ++ # pre-indexer's guard_ctr into _norm_rope, so SITE_QSA_ROPE_POS was declared, named ++ # and never written -- it would have read 0 for the life of every run and ++ # been reported as "the rope chain is clean", which is the exact misreading ++ # this series exists to prevent. An earlier version of this test asserted ++ # the wrong pointer and was green BECAUSE the wiring was wrong, so the ++ # assertion is inverted here on purpose: the two counters must DIFFER. ++ rope_ctr = bound.arguments["rope_guard_ctr"] ++ req_ctr = bound.arguments["guard_ctr"] ++ for ctr in (rope_ctr, req_ctr): ++ assert ctr.dtype == torch.int64 and ctr.numel() == 3 ++ assert rope_ctr.data_ptr() == guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_ROPE_POS, device ++ ).data_ptr() ++ assert req_ctr.data_ptr() == guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_PRE_INDEXER, device ++ ).data_ptr() ++ assert rope_ctr.data_ptr() != req_ctr.data_ptr() ++ ++ ++def test_rope_site_aborts_the_stage_even_in_warn_mode(): ++ """A clamped rope position is durable damage, not a local wrong answer. ++ ++ On the pooled path the wrongly-rotated key is STORED into the compressed ++ cache, so every later step reads it and the ring row that produced it is ++ still there to do it again. warn is the right default for a guard whose ++ masked value dies with the step; it is the wrong default for this one. ++ """ ++ assert ( ++ guard_telemetry.SITE_QSA_ROPE_POS ++ in guard_telemetry._ALWAYS_FATAL_SITES ++ ) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_ring_filler_bytes_are_masked_and_counted(): ++ """Positive control: the ring-residue crash case, verbatim. ++ ++ A re-claimed ring row filled with bf16 0.25 decodes, at the int64 position ++ tail, to roughly 0x3E803E803E803E80. Before guard 9 that value reached ++ cos_sin_ptr + pos * stride and compute-sanitizer reported an invalid read ++ 4.6e18 bytes out of range. After it, the launch must complete and the ++ counter must be non-zero. ++ ++ A silent counter here means the guard is not on the path -- the same class ++ of failure as a positive control that does not fire, and the reason this ++ test asserts the count rather than only the survival. ++ """ ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ ++ kw = _inputs(device) ++ # The victim's bytes: the whole row, tail included. Nothing zeroes the ring. ++ kw["state_cache"][0].fill_(0.25) ++ ++ # Drive the CROSS-CHUNK read-back, and prove the drive rather than assume ++ # it. The pooled group reads its position from the ring only when its first ++ # token lies in the PREVIOUS chunk, and the kernel decides that with ++ # (qsa_pre_indexer.py:299-311): ++ # ++ # chunk_end = logical[query_end - 1] ++ # chunk_start = chunk_end - query_len + 1 ++ # first_boundary = ((chunk_start + CR) // CR) * CR - 1 ++ # first_position = first_boundary - (CR - 1) ++ # first_in_chunk = first_position >= chunk_start ++ # ++ # A COMPRESS_RATIO-ALIGNED chunk start gives first_position == chunk_start, ++ # i.e. IN chunk, i.e. the ring is never read. An earlier run of this control ++ # with a start of 4 and counted 0, which said nothing about the guard. The ++ # start must be unaligned: 6 gives first_position 4 < chunk_start 6. ++ start = 6 ++ assert start % COMPRESS_RATIO != 0, "an aligned chunk start never crosses" ++ kw["logical_positions"] = ( ++ torch.arange(start, start + kw["logical_positions"].numel(), ++ dtype=kw["logical_positions"].dtype, device=device) ++ ) ++ chunk_end = start + kw["logical_positions"].numel() - 1 ++ chunk_start = chunk_end - kw["logical_positions"].numel() + 1 ++ first_boundary = ((chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO) * COMPRESS_RATIO - 1 ++ first_position = first_boundary - (COMPRESS_RATIO - 1) ++ assert first_position < chunk_start, ( ++ f"first_position {first_position} >= chunk_start {chunk_start}: this " ++ f"input does NOT take the ring read-back, so a zero count below would " ++ f"be meaningless" ++ ) ++ # ...and the row the reader picks must be one carrying the filler. ++ assert 0 <= (first_position % STATE_SIZE) < STATE_SIZE ++ ++ kw["k_work_metadata"][:, 0] = 0 ++ kw["k_work_metadata"][:, 1] = 0 ++ ++ qpi.qsa_pre_indexer( ++ kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], ++ kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], ++ kw["state_cache"], kw["state_slots"], kw["state_block_table"], ++ kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], ++ kw["compressed_slots"], kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ rope_pos_offset=HEAD_DIM, ++ ) ++ torch.cuda.synchronize() # the pre-guard build faults here ++ ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, device) ++ n_masked = int(ctr[0].item()) ++ assert n_masked > 0, ( ++ "the rope-position guard did not fire on the ring filler-bytes case. " ++ "Either the guard is not on this path or the case no longer drives the " ++ "ring read-back -- check before reading any zero counter as clean." ++ ) ++ worst = int(ctr[1].item()) - guard_telemetry.GUARD_BIAS ++ assert worst >= MAX_MODEL_LEN or worst < 0, ( ++ f"guard fired but recorded an in-range position {worst}" ++ ) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.parametrize("axis,axis_id", [("t", 0), ("h", 1), ("w", 2)]) ++@pytest.mark.parametrize("sign", ["mixed", "negative_only"]) ++def test_a_single_bad_axis_is_caught_and_named(axis, axis_id, sign): ++ """Any ONE of t/h/w is enough, and the record must name which. ++ ++ Under the production mrope section (11, 11, 10) all three components are ++ live row indices -- the IS_MROPE branch masks COLUMNS per axis while the ++ row comes unmasked from pos_rows[axis] -- so a bad pos_w faults exactly as ++ a bad pos_t does. An earlier revision recorded pos_t unconditionally, which for this ++ case would have reported an in-range value next to a non-zero count and ++ sent the reader to the wrong axis. ++ ++ The negative case also pins the filler: an all-negative offender must not ++ max to 0 and read back as position 0. ++ """ ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ ++ kw = _inputs(device) ++ row = {"t": 0, "h": 1, "w": 2}[axis] ++ kw["positions"] = kw["positions"].clone() ++ if sign == "mixed": ++ kw["positions"][row, 0] = MAX_MODEL_LEN + 7 # out of range, positive ++ kw["positions"][row, 1] = -5 # out of range, negative ++ expected_worst = MAX_MODEL_LEN + 7 ++ else: ++ # NEGATIVE ONLY. Every offender below zero is the case ++ # the tl.where filler decides: with a 0 filler the max would be the ++ # filler itself and slot 1 would read back as position 0 -- an ++ # in-range-looking value next to a non-zero count, which is the negative-value defect ++ # wearing a different hat. -5 must survive to the record. ++ kw["positions"][row, 0] = -5 ++ kw["positions"][row, 1] = -9 ++ expected_worst = -5 ++ ++ from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import qsa_pre_indexer ++ ++ qsa_pre_indexer( ++ kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], ++ kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], ++ kw["state_cache"], kw["state_slots"], kw["state_block_table"], ++ kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], ++ kw["compressed_slots"], kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ rope_pos_offset=HEAD_DIM, ++ ) ++ torch.cuda.synchronize() ++ ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, device) ++ assert int(ctr[0].item()) > 0, f"a bad pos_{axis} was not caught" ++ worst = int(ctr[1].item()) - guard_telemetry.GUARD_BIAS ++ got_axis = int(ctr[2].item()) - guard_telemetry.GUARD_BIAS ++ assert worst == expected_worst, ( ++ f"recorded {worst}, expected the offending value {expected_worst}" ++ ) ++ assert got_axis == axis_id, ( ++ f"recorded axis {got_axis}, expected {axis_id} (pos_{axis})" ++ ) +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +index f682ee9..b9a57d5 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +@@ -23,6 +23,8 @@ def _norm_rope( + pos_w, + cos_sin_ptr, + cos_sin_stride, ++ cos_sin_rows, # bounds-guard instrumentation ++ guard_ctr, # bounds-guard persistent int64[3], SITE_QSA_ROPE_POS + norm_weight_ptr, + eps, + IS_MROPE: tl.constexpr, +@@ -30,6 +32,72 @@ def _norm_rope( + MROPE_W: tl.constexpr, + ): + """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows.""" ++ # Bounds-guard instrumentation (ring-residue crash replay): the position is a ++ # ROW INDEX into the cos/sin table and nothing bounds it. The masks below ++ # bound which rotary PAIR is read, never which row, so ++ # `cos_sin_ptr + pos * cos_sin_stride` is an address formed from an ++ # unchecked value. ++ # ++ # On the pooled cross-chunk path that value is read straight out of a ++ # raw-key RING row (:305-317). The ring is excluded from zeroing ++ # (single_type_kv_cache_manager.py:96-101), so a re-claimed block keeps the ++ # previous occupant's bytes -- and nothing ties the row the reader selects ++ # (`first_position % STATE_SIZE`) to a row any occupant ever committed, so ++ # it can be allocator residue rather than a position at all. A crash replay drove ++ # exactly this and compute-sanitizer reported 11 invalid reads at ++ # +4.6e18 bytes, the filler bytes decoded as an int64. ++ # ++ # This is the ONLY guard in the series that covers that chain: the request guard ++ # bounds the pre-indexer's `request`, which sits upstream of the ring read ++ # and cannot see a bad position. ++ # WHICH AXIS. Under the production mrope section (11, 11, 10) ++ # all three components are live row indices: the IS_MROPE branch masks ++ # COLUMNS per axis while the row comes unmasked from pos_rows[axis], so a ++ # bad pos_h or pos_w faults exactly as a bad pos_t does. Recording pos_t ++ # unconditionally would name the wrong axis and, for a t-clean/h-bad case, ++ # record an in-range value next to a non-zero count. ++ # ++ # The filler under tl.where must be strongly negative, NOT 0: an offender ++ # that is entirely negative (i_t == -1 is the shape that started this ++ # series) would otherwise max to 0 and read back as position 0 -- the ++ # defect re-entering at the one site with a reproduced fault. ++ _FILL: tl.constexpr = -(1 << 62) ++ bad_t = (pos_t < 0) | (pos_t >= cos_sin_rows) ++ if IS_MROPE: ++ bad_h = (pos_h < 0) | (pos_h >= cos_sin_rows) ++ bad_w = (pos_w < 0) | (pos_w >= cos_sin_rows) ++ else: ++ # the non-mrope branch reads only pos_t (:63-66), and with NO mask at ++ # all, which is why this guard clamps rather than masks. ++ # always-false, but correctly shaped and typed for the max below. ++ bad_h = pos_h < pos_h ++ bad_w = pos_w < pos_w ++ guard_bad = bad_t | bad_h | bad_w ++ guard_n = tl.sum(guard_bad.to(tl.int32), axis=0) ++ if guard_n > 0: ++ worst_t = tl.max(tl.where(bad_t, pos_t, _FILL), axis=0).to(tl.int64) ++ worst_h = tl.max(tl.where(bad_h, pos_h, _FILL), axis=0).to(tl.int64) ++ worst_w = tl.max(tl.where(bad_w, pos_w, _FILL), axis=0).to(tl.int64) ++ worst = tl.maximum(worst_t, tl.maximum(worst_h, worst_w)) ++ # slot 2 names the axis: 0 = t, 1 = h, 2 = w. Derived from the MASKS ++ # (deliberately), not by comparing the winner's VALUE against each ++ # axis's max. Value comparison mis-attributes whenever two axes share ++ # the worst value -- which is the common case here, because pos_h and ++ # pos_w are frequently equal to pos_t (the non-2D path sets them so, ++ # and a stale ring tail writes the same int64 into all three). Asking ++ # "did this axis have any offender at all", in t/h/w order, cannot. ++ any_t = tl.max(bad_t.to(tl.int32), axis=0) > 0 ++ any_h = tl.max(bad_h.to(tl.int32), axis=0) > 0 ++ axis = tl.where(any_t, 0, tl.where(any_h, 1, 2)) ++ tl.atomic_add(guard_ctr, guard_n.to(tl.int64)) ++ tl.atomic_max(guard_ctr + 1, worst + _GUARD_BIAS) ++ tl.atomic_max(guard_ctr + 2, axis.to(tl.int64) + _GUARD_BIAS) ++ # Clamp rather than skip: the rotation is then wrong for these rows, which ++ # is what a masked fault IS. The counter is the result, not the output. ++ pos_t = tl.minimum(tl.maximum(pos_t, 0), cos_sin_rows - 1) ++ pos_h = tl.minimum(tl.maximum(pos_h, 0), cos_sin_rows - 1) ++ pos_w = tl.minimum(tl.maximum(pos_w, 0), cos_sin_rows - 1) ++ + TILE_T: tl.constexpr = x.shape[0] + TILE_H: tl.constexpr = x.shape[1] + D: tl.constexpr = x.shape[2] +@@ -106,6 +174,8 @@ def _qsa_pre_indexer_kernel( + pos_stride_axis, + pos_stride_token, + cos_sin_ptr, ++ cos_sin_rows, # bounds-guard instrumentation ++ rope_guard_ctr, # bounds-guard persistent int64[3], SITE_QSA_ROPE_POS + q_norm_weight_ptr, + k_norm_weight_ptr, + eps, +@@ -189,6 +259,8 @@ def _qsa_pre_indexer_kernel( + pos_w, + cos_sin_ptr, + D // 2, ++ cos_sin_rows, ++ rope_guard_ctr, + q_norm_weight_ptr, + eps, + IS_2D_POSITIONS, +@@ -348,6 +420,8 @@ def _qsa_pre_indexer_kernel( + pos_w + tl.arange(0, 1), + cos_sin_ptr, + D // 2, ++ cos_sin_rows, ++ rope_guard_ctr, + k_norm_weight_ptr, + eps, + IS_K_MROPE, +@@ -488,6 +562,8 @@ def qsa_pre_indexer( + pos_stride_axis, + pos_stride_token, + cos_sin_cache, ++ cos_sin_cache.shape[0], ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, q.device), + q_norm_weight, + k_norm_weight, + eps, +-- +2.43.0 + diff --git a/patches/0018-sched-defer-ring-block-free-under-overlapping-steps.patch b/patches/0018-sched-defer-ring-block-free-under-overlapping-steps.patch new file mode 100644 index 0000000..afaa76a --- /dev/null +++ b/patches/0018-sched-defer-ring-block-free-under-overlapping-steps.patch @@ -0,0 +1,256 @@ +From 4823da32e2090433d0f56e40316de3cfcf637d2f Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:09 -0400 +Subject: [PATCH 018/103] sched: defer ring-block free under overlapping steps + +A circular-buffer ring block could be freed while an earlier, still-executing +step was reading it. The free is deferred until the overlapping step has +committed. +--- + tests/v1/core/test_deferred_block_free.py | 154 ++++++++++++++++++++++ + vllm/v1/core/sched/scheduler.py | 33 +++++ + 2 files changed, 187 insertions(+) + +diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py +index ebbf492..0c4738a 100644 +--- a/tests/v1/core/test_deferred_block_free.py ++++ b/tests/v1/core/test_deferred_block_free.py +@@ -17,8 +17,11 @@ from unittest.mock import PropertyMock, patch + + import pytest + ++import torch ++ + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy ++from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ModelRunnerOutput + from vllm.v1.request import RequestStatus +@@ -544,3 +547,154 @@ def test_cow_retentions_deferred_until_copy_step_processed(): + assert not scheduler.deferred_frees + out2 = scheduler.schedule() + assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"] ++ ++ ++# -------------------------------------------------------------------------- ++# Ring-skip telemetry: the ring gate. A CircularBufferSpec group (the QSA raw-key ++# ring, one block per request, rewritten every step) enables deferral under ++# overlapping batches WITHOUT a KV connector. ++# -------------------------------------------------------------------------- ++ ++def _ring_spec(block_size: int = 16) -> CircularBufferSpec: ++ return CircularBufferSpec( ++ block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ) ++ ++ ++def test_gate_enabled_for_ring_group_with_overlapping_batches(): ++ # Overlap is forced the same way test_gate_enabled_for_async_consumer ++ # forces it: async scheduling is disabled on CPU and PP needs GPUs. ++ with patch.object( ++ VllmConfig, ++ "max_concurrent_batches", ++ new_callable=PropertyMock, ++ return_value=2, ++ ): ++ scheduler = create_scheduler( ++ model=MODEL, async_scheduling=True, kv_cache_spec=_ring_spec() ++ ) ++ assert scheduler.defer_block_free ++ ++ ++def test_gate_disabled_for_ring_group_without_overlap(): ++ # One batch in flight at a time: nothing can race the free. The gate must ++ # stay off so the single-step path keeps immediate frees. Overlap is ++ # forced to 1 explicitly: whether async scheduling survives on the test ++ # platform is not the property under test. ++ with patch.object( ++ VllmConfig, ++ "max_concurrent_batches", ++ new_callable=PropertyMock, ++ return_value=1, ++ ): ++ scheduler = create_scheduler( ++ model=MODEL, async_scheduling=True, kv_cache_spec=_ring_spec() ++ ) ++ assert not scheduler.defer_block_free ++ ++ ++def test_gate_disabled_for_non_ring_group_without_connector(): ++ # The earlier negative case, restated under the new gate: full attention ++ # + overlap + no connector must still be off. ++ with patch.object( ++ VllmConfig, ++ "max_concurrent_batches", ++ new_callable=PropertyMock, ++ return_value=2, ++ ): ++ scheduler = create_scheduler(model=MODEL, async_scheduling=True) ++ assert not scheduler.defer_block_free ++ ++ ++def test_ring_preempt_withholds_the_ring_block_until_the_fence(): ++ """The property the ring gate buys: a preempted request's ring block does not ++ return to the pool while a step that may still write it is in flight, ++ and it does return once the newest in-flight step's output is processed. ++ """ ++ with patch.object( ++ VllmConfig, ++ "max_concurrent_batches", ++ new_callable=PropertyMock, ++ return_value=2, ++ ): ++ scheduler = create_scheduler( ++ model=MODEL, async_scheduling=True, kv_cache_spec=_ring_spec() ++ ) ++ assert scheduler.defer_block_free # via the ring gate, nothing forced ++ pool = scheduler.kv_cache_manager.block_pool ++ num_free_initially = pool.get_num_free_blocks() ++ ++ request, out0, out1 = _setup_request_with_inflight_step(scheduler) ++ num_free_running = pool.get_num_free_blocks() ++ # the ring is one block per request for its whole lifetime ++ assert num_free_running == num_free_initially - 1 ++ ++ scheduler.running.remove(request) ++ scheduler._preempt_request(request, time.monotonic()) ++ assert request.status == RequestStatus.PREEMPTED ++ assert len(scheduler.deferred_frees) == 1 ++ assert pool.get_num_free_blocks() == num_free_running ++ for manager in scheduler.kv_cache_manager.coordinator.single_type_managers: ++ assert request.request_id not in manager.req_to_blocks ++ ++ scheduler.update_from_output(out0, _make_model_runner_output(out0)) ++ assert len(scheduler.deferred_frees) == 1 ++ assert pool.get_num_free_blocks() == num_free_running ++ scheduler.update_from_output(out1, _make_model_runner_output(out1)) ++ assert not scheduler.deferred_frees ++ assert pool.get_num_free_blocks() == num_free_initially ++ ++ ++def test_deferred_victim_stops_the_preemption_cascade(): ++ """Under deferred free a preempted victim returns no blocks ++ until its in-flight step is processed, so the allocation retry must not ++ keep preempting. Two 48-token requests fill a 6-block pool; the first ++ request's speculative decode (token 49) needs a 4th block. Expected: exactly ONE ++ victim (the last running request), the requester stays RUNNING and ++ unscheduled, and after the in-flight step is processed it gets its block. ++ Without the fix this preempts both requests in one schedule() call. ++ """ ++ with patch.object( ++ VllmConfig, ++ "max_concurrent_batches", ++ new_callable=PropertyMock, ++ return_value=2, ++ ): ++ # 7 blocks = 6 usable (block 0 is the null block); ring gate irrelevant ++ # here, the cascade is the scheduler's -- force deferral as the ++ # upstream tests do. ++ scheduler = create_scheduler( ++ model=MODEL, async_scheduling=True, num_blocks=7, block_size=16 ++ ) ++ scheduler.defer_block_free = True ++ pool = scheduler.kv_cache_manager.block_pool ++ assert pool.get_num_free_blocks() == 6 ++ ++ reqs = create_requests( ++ num_requests=2, num_tokens=48, max_tokens=8, stop_token_ids=[STOP_TOKEN_ID] ++ ) ++ for r in reqs: ++ scheduler.add_request(r) ++ out0 = scheduler.schedule() ++ assert sorted(out0.num_scheduled_tokens.values()) == [48, 48] ++ assert pool.get_num_free_blocks() == 0 ++ ++ # step 2 (step 1 in flight): the first request's decode crosses into a ++ # 4th block -> allocation fails -> the loop preempts the LAST running one. ++ first, second = scheduler.running[0], scheduler.running[1] ++ out1 = scheduler.schedule() ++ assert second.status == RequestStatus.PREEMPTED ++ assert first.status == RequestStatus.RUNNING, "the cascade reached the requester" ++ assert first in scheduler.running ++ assert len(scheduler.deferred_frees) == 1 ++ assert pool.get_num_free_blocks() == 0 ++ assert first.request_id not in out1.num_scheduled_tokens ++ ++ # step 1's output processed: the victim's blocks drain, the requester ++ # is scheduled on the next call. ++ scheduler.update_from_output(out0, _make_model_runner_output(out0)) ++ assert not scheduler.deferred_frees ++ assert pool.get_num_free_blocks() == 3 ++ out2 = scheduler.schedule() ++ assert first.request_id in out2.num_scheduled_tokens ++ assert first.status == RequestStatus.RUNNING +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 9f2d99d..f1c9b79 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -56,10 +56,12 @@ from vllm.v1.core.sched.request_queue import ( + from vllm.v1.core.sched.utils import check_stop, remove_all + from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs + from vllm.v1.kv_cache_interface import ( ++ CircularBufferSpec, + KVCacheConfig, + MambaSpec, + get_mamba_prefill_checkpoint_position, + is_mamba_prefill_checkpoint_valid, ++ iter_layer_specs, + ) + from vllm.v1.metrics.perf import ModelMetrics, PerfStats + from vllm.v1.metrics.stats import ( +@@ -177,6 +179,26 @@ class Scheduler(SchedulerInterface): + + self.requires_kv_delivery = self.connector.requires_kv_delivery + ++ # Ring-skip telemetry: the QSA ring block (CircularBufferSpec) is rewritten ++ # by every step that touches the request, so with overlapping batches ++ # (async scheduling or PP) an in-flight step may still be committing a ++ # freed/preempted request's suffix rows when the scheduler hands that ++ # block to another request. A KV connector is not the only writer that ++ # can race the stale in-flight write; the next request's own first ++ # step is. Defer the return of freed blocks to the pool until the ++ # newest scheduled step's output has been processed whenever a ring ++ # group exists. Same mechanism and fence as the connector gate above. ++ if ( ++ not self.defer_block_free ++ and self.vllm_config.max_concurrent_batches > 1 ++ and any( ++ isinstance(spec, CircularBufferSpec) ++ for group in kv_cache_config.kv_cache_groups ++ for spec in iter_layer_specs(group.kv_cache_spec) ++ ) ++ ): ++ self.defer_block_free = True ++ + self.kv_event_publisher = EventPublisherFactory.create( + self.kv_events_config, + self.parallel_config.data_parallel_index, +@@ -789,6 +811,7 @@ class Scheduler(SchedulerInterface): + else: + preempted_req = self.running.pop() + ++ num_deferred_before = len(self.deferred_frees) + self._preempt_request( + preempted_req, + scheduled_timestamp, +@@ -798,6 +821,16 @@ class Scheduler(SchedulerInterface): + if preempted_req == request: + # No more request to preempt. Cannot schedule this request. + break ++ if len(self.deferred_frees) > num_deferred_before: ++ # The ring-skip fix: the victim's blocks are fenced ++ # behind its in-flight step and return to the pool only ++ # when that step's output is processed, so retrying the ++ # allocation now cannot succeed on their account and would ++ # preempt the next request for nothing -- with every running ++ # request in flight (PP / async) the whole batch. Stop at one ++ # deferred victim; this request stays running, unscheduled, ++ # and retries next step after the drain. ++ break + + if new_blocks is None: + # Cannot schedule this request. +-- +2.43.0 + diff --git a/patches/0019-qsa-count-skipped-ring-commits.patch b/patches/0019-qsa-count-skipped-ring-commits.patch new file mode 100644 index 0000000..ad582d4 --- /dev/null +++ b/patches/0019-qsa-count-skipped-ring-commits.patch @@ -0,0 +1,264 @@ +From 5db178597cf9d2d250ce7d324bd995b45e91af29 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:09 -0400 +Subject: [PATCH 019/103] qsa: count skipped ring commits + +--- + .../gpu/test_qsa_ring_skip_telemetry.py | 204 ++++++++++++++++++ + .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 15 ++ + 2 files changed, 219 insertions(+) + create mode 100644 tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py + +diff --git a/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +new file mode 100644 +index 0000000..0468288 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +@@ -0,0 +1,204 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""v2 ring-skip telemetry: the ring-skip telemetry site of the QSA pre-indexer. ++ ++The suffix-commit loop stores one ring row per retained token, every store ++masked by ``valid_slot``. A live token whose state slot is < 0 or beyond the ++ring capacity is therefore NOT written and nothing said so (one candidate cause ++for the raw-key ring crash class). This site counts those tokens into ++``SITE_QSA_RING_SKIP``; the write stays skipped. These tests prove the wiring ++(CPU) and that the count moves for exactly the skipped tokens (GPU). ++ ++Helpers mirror tests/v1/worker/gpu/test_qsa_rope_position_guard.py on purpose: ++the two files must stay independently runnable. ++""" ++ ++import inspect ++ ++import pytest ++import torch ++ ++from vllm.models.qwen4_exp.nvidia.ops import qsa_pre_indexer as qpi ++from vllm.v1.worker import guard_telemetry ++ ++COMPRESS_RATIO = 4 ++HEAD_DIM = 128 ++N_HEADS = 4 ++STATE_SIZE = 8 # ring capacity ++ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++MAX_MODEL_LEN = 4096 ++NUM_BLOCKS = 8 ++RING_CAPACITY = NUM_BLOCKS * STATE_SIZE ++ ++ ++class _Recorder: ++ """Stands in for the Triton kernel and keeps the launch arguments.""" ++ ++ def __init__(self): ++ self.launches = [] ++ ++ def __getitem__(self, grid): ++ def launch(*args, **kwargs): ++ self.launches.append((grid, args, kwargs)) ++ ++ return launch ++ ++ ++def _inputs(device, num_tokens=4): ++ q = torch.randn(num_tokens, N_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device) ++ k = torch.randn(num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ positions = ( ++ torch.arange(num_tokens, dtype=torch.int64, device=device) ++ .unsqueeze(0) ++ .expand(3, -1) ++ .contiguous() ++ ) ++ cos_sin = torch.randn( ++ MAX_MODEL_LEN, HEAD_DIM // 2, dtype=torch.bfloat16, device=device ++ ) ++ ones = torch.ones(HEAD_DIM, dtype=torch.bfloat16, device=device) ++ q_out = torch.empty(num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ state_cache = torch.zeros( ++ NUM_BLOCKS, STATE_SIZE, 1, ROW_WIDTH, dtype=torch.bfloat16, device=device ++ ) ++ state_slots = torch.arange(num_tokens, dtype=torch.int32, device=device) ++ state_table = torch.zeros(1, 1, dtype=torch.int32, device=device) ++ qsl = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) ++ logical = torch.arange(num_tokens, dtype=torch.int32, device=device) ++ comp_cache = torch.zeros( ++ NUM_BLOCKS, 800, 1, HEAD_DIM, dtype=torch.bfloat16, device=device ++ ) ++ comp_slots = torch.zeros(num_tokens, dtype=torch.int32, device=device) ++ kwb = torch.zeros( ++ max(num_tokens // COMPRESS_RATIO, 1), 2, dtype=torch.int32, device=device ++ ) ++ return dict( ++ q=q, k=k, positions=positions, cos_sin_cache=cos_sin, ++ q_norm_weight=ones, k_norm_weight=ones, eps=1e-6, q_out=q_out, ++ state_cache=state_cache, state_slots=state_slots, ++ state_block_table=state_table, query_start_loc=qsl, ++ logical_positions=logical, compressed_cache=comp_cache, ++ compressed_slots=comp_slots, k_work_metadata=kwb, ++ ) ++ ++ ++def _launch(kw): ++ qpi.qsa_pre_indexer( ++ kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], ++ kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], ++ kw["state_cache"], kw["state_slots"], kw["state_block_table"], ++ kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], ++ kw["compressed_slots"], kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ rope_pos_offset=HEAD_DIM, ++ ) ++ ++ ++def test_site_is_declared_named_and_not_fatal(): ++ """A site that exists but is unnamed reports as a bare integer; a site ++ that is fatal would abort a boot on the older behaviour it merely ++ observes. Both are the wrong instrument.""" ++ site = guard_telemetry.SITE_QSA_RING_SKIP ++ assert site == 10 and guard_telemetry.NUM_SITES == 11 ++ assert site in guard_telemetry.SITE_NAMES ++ assert "SKIPPED" in guard_telemetry.SITE_NAMES[site] ++ assert site not in guard_telemetry._ALWAYS_FATAL_SITES ++ # the counter store is sized by NUM_SITES: the new site must have a row ++ ctr = guard_telemetry.counter(site, torch.device("cpu")) ++ assert ctr.dtype == torch.int64 and ctr.numel() == guard_telemetry.SLOTS_PER_SITE ++ ++ ++def test_launcher_passes_the_ring_skip_counter(monkeypatch): ++ """The kernel parameter must be the persistent SITE_QSA_RING_SKIP buffer ++ and must differ from the two counters the kernel already takes (the ++ rope-guard series shipped a declared-but-never-written site once; the ++ inverted assertion below is the same defence).""" ++ device = torch.device("cpu") ++ sig = inspect.signature(qpi._qsa_pre_indexer_kernel.fn) ++ assert "ring_skip_ctr" in sig.parameters ++ rec = _Recorder() ++ monkeypatch.setattr(qpi, "_qsa_pre_indexer_kernel", rec) ++ _launch(_inputs(device)) ++ ++ assert len(rec.launches) == 1 ++ _grid, args, kwargs = rec.launches[0] ++ kwargs = {k: v for k, v in kwargs.items() ++ if k not in ("num_warps", "num_stages", "num_ctas", "grid")} ++ bound = sig.bind(*args, **kwargs) ++ bound.apply_defaults() ++ skip_ctr = bound.arguments["ring_skip_ctr"] ++ assert skip_ctr.dtype == torch.int64 and skip_ctr.numel() == 3 ++ assert skip_ctr.data_ptr() == guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_RING_SKIP, device ++ ).data_ptr() ++ for other in ("rope_guard_ctr", "guard_ctr"): ++ assert bound.arguments[other].data_ptr() != skip_ctr.data_ptr() ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_negative_control_valid_slots_count_zero(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ kw = _inputs(device) ++ _launch(kw) ++ torch.cuda.synchronize() ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, device) ++ assert ctr.tolist() == [0, 0, 0] ++ # and the commit happened: every token's raw key and position tail ++ for t in range(kw["k"].shape[0]): ++ row = kw["state_cache"][0, t, 0, :] ++ assert torch.equal(row[:HEAD_DIM], kw["k"][t]) ++ tail = row[HEAD_DIM:HEAD_DIM + 12].contiguous().view(torch.int64) ++ assert tail.tolist() == [t, t, t] ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.parametrize( ++ "bad_index,bad_slot", ++ [(2, -1), (1, RING_CAPACITY), (3, RING_CAPACITY + 5)], ++ ids=["negative-slot", "at-capacity", "past-capacity"], ++) ++def test_positive_control_one_skipped_token_is_counted_and_located(bad_index, bad_slot): ++ """Exactly the skipped token is counted; slot 1 names the token, slot 2 ++ carries the offending state slot value; the other rows are still written ++ and the bad token's row is not.""" ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ kw = _inputs(device) ++ kw["state_slots"][bad_index] = bad_slot ++ _launch(kw) ++ torch.cuda.synchronize() ++ ++ ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, device) ++ n, row, val = ctr.tolist() ++ assert n == 1, f"expected one skipped token, counter says {n}" ++ assert row - guard_telemetry.GUARD_BIAS == bad_index ++ assert val - guard_telemetry.GUARD_BIAS == bad_slot ++ # neighbours committed (key AND int64 position tail), the bad token's ++ # own ring row untouched over its WHOLE width -- the crash class lives ++ # in the tail, so a key-only check would miss a stray tail write. ++ for t in range(kw["k"].shape[0]): ++ row = kw["state_cache"][0, t, 0, :] ++ if t == bad_index: ++ assert torch.count_nonzero(row) == 0, "skipped row was written" ++ else: ++ assert torch.equal(row[:HEAD_DIM], kw["k"][t]) ++ tail = row[HEAD_DIM:HEAD_DIM + 12].contiguous().view(torch.int64) ++ assert tail.tolist() == [t, t, t] ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++def test_positive_control_two_skipped_tokens_accumulate(): ++ device = torch.device("cuda", torch.cuda.current_device()) ++ guard_telemetry.prealloc(device) ++ guard_telemetry.reset_for_test() ++ kw = _inputs(device) ++ kw["state_slots"][0] = -1 ++ kw["state_slots"][3] = RING_CAPACITY + 1 ++ _launch(kw) ++ torch.cuda.synchronize() ++ n, row, val = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, device).tolist() ++ assert n == 2 ++ assert row - guard_telemetry.GUARD_BIAS == 3 # max token index ++ assert val - guard_telemetry.GUARD_BIAS == RING_CAPACITY + 1 # max slot value +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +index b9a57d5..912bef3 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +@@ -201,6 +201,7 @@ def _qsa_pre_indexer_kernel( + num_k_work, + num_requests, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] ++ ring_skip_ctr, # ring-skip telemetry, persistent int64[3], SITE_QSA_RING_SKIP + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, +@@ -454,6 +455,19 @@ def _qsa_pre_indexer_kernel( + valid_slot = ( + valid_token & (slot >= 0) & (slot < num_state_blocks * STATE_SIZE) + ) ++ # Ring-skip telemetry: a live token whose slot fails the range check ++ # has its ring row silently NOT written (every store below is ++ # masked by valid_slot). That token's raw key is then absent ++ # from the ring the next step reads (a candidate cause). Count ++ # it -- the write stays skipped, the reconciliation is the ring-row tag fix. ++ skipped = valid_token & ~valid_slot ++ tl.atomic_add(ring_skip_ctr, 1, mask=skipped) ++ tl.atomic_max( ++ ring_skip_ctr + 1, token.to(tl.int64) + _GUARD_BIAS, mask=skipped ++ ) ++ tl.atomic_max( ++ ring_skip_ctr + 2, slot.to(tl.int64) + _GUARD_BIAS, mask=skipped ++ ) + safe_slot = tl.maximum(slot, 0) + state_row = ( + state_cache_ptr +@@ -589,6 +603,7 @@ def qsa_pre_indexer( + num_k_work, + min(query_start_loc.numel() - 1, state_block_table.shape[0]), + guard_telemetry.counter(guard_telemetry.SITE_QSA_PRE_INDEXER, q.device), ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, q.device), + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, +-- +2.43.0 + diff --git a/patches/0020-sched-scheduler-allocation-log-for-KV-capacity-probe.patch b/patches/0020-sched-scheduler-allocation-log-for-KV-capacity-probe.patch new file mode 100644 index 0000000..42eebae --- /dev/null +++ b/patches/0020-sched-scheduler-allocation-log-for-KV-capacity-probe.patch @@ -0,0 +1,156 @@ +From 3191b58705d1a99632ee8e7fcc820eb3323c230b Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:13:10 -0400 +Subject: [PATCH 020/103] sched: scheduler allocation log for KV capacity + probes + +Env-gated per-step log of blocks allocated and freed and scheduled tokens per +request; prefix-cache hits are read behind a schema check so an upstream field +rename fails visibly. +--- + vllm/v1/core/sched/e2_alloc_log.py | 111 +++++++++++++++++++++++++++++ + vllm/v1/core/sched/scheduler.py | 3 + + 2 files changed, 114 insertions(+) + create mode 100644 vllm/v1/core/sched/e2_alloc_log.py + +diff --git a/vllm/v1/core/sched/e2_alloc_log.py b/vllm/v1/core/sched/e2_alloc_log.py +new file mode 100644 +index 0000000..0efcbd7 +--- /dev/null ++++ b/vllm/v1/core/sched/e2_alloc_log.py +@@ -0,0 +1,111 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Debug-only scheduler allocation log for KV capacity probes (off unless VLLM_E2_SCHED_LOG=). ++ ++One JSON line per scheduler update while any request exists, and one "fin" line per request as it is freed: ++ {"t": wall, "m": monotonic, "step": n, "total": pool blocks, "free": free blocks, "running": r, "waiting": w, ++ "tot": tokens scheduled this step, ++ "reqs": [[id, prompt, computed, output, cached, preemptions, spec_scheduled, status, [blocks per group], ++ tokens scheduled for this request this step]]} ++ {"t": wall, "m": monotonic, "fin": [id, prompt, computed, output, cached, preemptions, status, [blocks per group]]} ++Blocks are counted per KV-cache group, null blocks excluded. A logging failure writes an "error" line and never ++raises into the scheduler. Unset = the scheduler only pays one env check per update. ++""" ++import json ++import os ++import time ++ ++_PATH = os.environ.get("VLLM_E2_SCHED_LOG", "") ++_state = {"f": None, "step": 0, "flush": 0.0, "wrapped": False, "errors": 0, "checked": False} ++REQ_FIELDS = ("request_id", "num_prompt_tokens", "num_computed_tokens", "num_output_tokens", "num_preemptions", ++ "status", "prefill_stats") ++ ++ ++def enabled() -> bool: ++ return bool(_PATH) ++ ++ ++def _blocks(sched, rid): ++ return [sum(1 for b in m.req_to_blocks.get(rid, ()) if not b.is_null) ++ for m in sched.kv_cache_manager.coordinator.single_type_managers] ++ ++ ++def _status(r): ++ return getattr(r.status, "name", str(r.status)) ++ ++ ++def _cached(r): ++ """Prefix-cache hit for this request. vLLM keeps it in `request.prefill_stats` (PrefillStats.num_cached_tokens = ++ local + external); the output processor takes those stats after the first prefill (`take_prefill_stats` sets the ++ attribute to None), so -1 means "no longer knowable from the scheduler", never a silent zero.""" ++ ps = getattr(r, "prefill_stats", None) ++ return int(getattr(ps, "num_cached_tokens", -1)) if ps is not None else -1 ++ ++ ++def _contract(r): ++ """Once per process: every field this log reads must exist on the real Request object. A missing field is written ++ once as a schema_error record — the evaluator fails closed on it — instead of turning every row into an error.""" ++ if _state["checked"]: ++ return ++ _state["checked"] = True ++ missing = [f for f in REQ_FIELDS if not hasattr(r, f)] ++ if missing: ++ _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "schema_error": missing}, force=True) ++ ++ ++def _write(rec, force=False): ++ f = _state["f"] ++ if f is None: ++ f = _state["f"] = open(_PATH, "a", buffering=1 << 16) ++ f.write(json.dumps(rec, separators=(",", ":")) + "\n") ++ if force or rec["t"] - _state["flush"] > 0.5: ++ f.flush() ++ _state["flush"] = rec["t"] ++ ++ ++def _error(e): ++ _state["errors"] += 1 ++ if _state["errors"] <= 20: ++ try: ++ _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), ++ "error": f"{type(e).__name__}: {e}"[:300]}, force=True) ++ except Exception: # noqa: BLE001 ++ pass ++ ++ ++def _wrap_free(sched): ++ orig = sched._free_request ++ ++ def _free_request(request, *args, **kwargs): ++ try: ++ _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), ++ "fin": [request.request_id, request.num_prompt_tokens, request.num_computed_tokens, ++ request.num_output_tokens, _cached(request), request.num_preemptions, ++ _status(request), _blocks(sched, request.request_id)]}, force=True) ++ except Exception as e: # noqa: BLE001 - the debug log must never break the scheduler ++ _error(e) ++ return orig(request, *args, **kwargs) ++ ++ sched._free_request = _free_request ++ _state["wrapped"] = True ++ ++ ++def log_step(sched, scheduler_output) -> None: ++ try: ++ if not _state["wrapped"]: ++ _wrap_free(sched) ++ if not sched.requests: ++ return ++ _state["step"] += 1 ++ spec = getattr(scheduler_output, "scheduled_spec_decode_tokens", None) or {} ++ pool = sched.kv_cache_manager.block_pool ++ nsched = getattr(scheduler_output, "num_scheduled_tokens", None) or {} ++ _contract(next(iter(sched.requests.values()))) ++ rows = [[rid, r.num_prompt_tokens, r.num_computed_tokens, r.num_output_tokens, _cached(r), ++ r.num_preemptions, len(spec.get(rid, ())), _status(r), _blocks(sched, rid), int(nsched.get(rid, 0))] ++ for rid, r in sched.requests.items()] ++ _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "step": _state["step"], ++ "total": pool.num_gpu_blocks, "free": pool.get_num_free_blocks(), ++ "running": len(sched.running), "waiting": len(sched.waiting), ++ "tot": int(getattr(scheduler_output, "total_num_scheduled_tokens", 0) or 0), "reqs": rows}) ++ except Exception as e: # noqa: BLE001 ++ _error(e) +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index f1c9b79..25af1b6 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -70,6 +70,7 @@ from vllm.v1.metrics.stats import ( + SchedulerStats, + ) + from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput ++from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log + from vllm.v1.request import Request, RequestStatus, StreamingUpdate + from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup + from vllm.v1.spec_decode.metrics import SpecDecodingStats +@@ -2359,6 +2360,8 @@ class Scheduler(SchedulerInterface): + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats + ++ if _e2_alloc_log.enabled(): ++ _e2_alloc_log.log_step(self, scheduler_output) + return engine_core_outputs + + def _ec_transfer_pending(self, request: Request, num_computed_tokens: int) -> bool: +-- +2.43.0 + diff --git a/patches/0021-spec-decode-structured-output-under-concurrency-on-a.patch b/patches/0021-spec-decode-structured-output-under-concurrency-on-a.patch new file mode 100644 index 0000000..7fb83d6 --- /dev/null +++ b/patches/0021-spec-decode-structured-output-under-concurrency-on-a.patch @@ -0,0 +1,1176 @@ +From 53ba1efcf7012a9b847bfcc0de3cbdd17e3303ce Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 021/103] spec-decode: structured output under concurrency on a + PP shape + +Ports upstream vLLM PRs #54442 and #56802 (issue #54437): never sample from an +unmasked row, and key draft-token hand-off snapshots by scheduler step so an +alternating PP microbatch cannot overwrite the slot. +--- + tests/v1/core/test_scheduler.py | 19 ++ + tests/v1/worker/test_draft_tokens_handler.py | 210 ++++++++++++++++++ + tests/v1/worker/test_grammar_fail_closed.py | 151 +++++++++++++ + .../v1/worker/test_grammar_invalid_drafts.py | 76 +++++++ + vllm/v1/core/sched/output.py | 9 + + vllm/v1/core/sched/scheduler.py | 13 +- + vllm/v1/engine/core.py | 4 +- + vllm/v1/executor/abstract.py | 9 +- + vllm/v1/executor/multiproc_executor.py | 7 +- + vllm/v1/executor/uniproc_executor.py | 9 +- + vllm/v1/structured_output/utils.py | 13 ++ + vllm/v1/worker/gpu/model_runner.py | 59 ++++- + vllm/v1/worker/gpu/sample/batch_shard.py | 7 + + .../gpu/spec_decode/rejection_sampler.py | 18 +- + vllm/v1/worker/gpu/spec_decode/utils.py | 148 +++++++++--- + vllm/v1/worker/gpu_model_runner.py | 3 +- + vllm/v1/worker/gpu_worker.py | 4 +- + vllm/v1/worker/worker_base.py | 33 ++- + 18 files changed, 748 insertions(+), 44 deletions(-) + create mode 100644 tests/v1/worker/test_draft_tokens_handler.py + create mode 100644 tests/v1/worker/test_grammar_fail_closed.py + create mode 100644 tests/v1/worker/test_grammar_invalid_drafts.py + +diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py +index 061c825..07ced9e 100644 +--- a/tests/v1/core/test_scheduler.py ++++ b/tests/v1/core/test_scheduler.py +@@ -30,6 +30,7 @@ from vllm.multimodal.inputs import ( + MultiModalKwargsItem, + PlaceholderRange, + ) ++from vllm.platforms import current_platform + from vllm.sampling_params import SamplingParams, StructuredOutputsParams + from vllm.utils.hashing import sha256 + from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +@@ -249,6 +250,24 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None): + assert scheduler.running[i] == request + + ++def test_scheduler_output_step_identity_defaults_to_zero(): ++ output = SchedulerOutput.make_empty() ++ assert output.scheduler_step == 0 ++ ++ ++def test_scheduler_output_carries_step_identity(): ++ if not current_platform.device_type: ++ pytest.skip("No inferred device type in this environment") ++ ++ scheduler = create_scheduler() ++ request = create_requests(num_requests=1)[0] ++ scheduler.add_request(request) ++ ++ output = scheduler.schedule() ++ assert output.scheduler_step == scheduler.current_step ++ assert output.scheduler_step == 1 ++ ++ + def test_scheduler_stats_route_to_existing_output_client(): + scheduler = create_scheduler() + request = create_requests(num_requests=1)[0] +diff --git a/tests/v1/worker/test_draft_tokens_handler.py b/tests/v1/worker/test_draft_tokens_handler.py +new file mode 100644 +index 0000000..bf016eb +--- /dev/null ++++ b/tests/v1/worker/test_draft_tokens_handler.py +@@ -0,0 +1,210 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""DraftTokensHandler step identity (#54437 piece 2 / Link A).""" ++ ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler ++from vllm.v1.worker.worker_base import WorkerWrapperBase ++ ++pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] ++ ++K = 4 ++ ++ ++def _device() -> torch.device: ++ if torch.cuda.is_available(): ++ return torch.device("cuda:0") ++ return torch.device("cpu") ++ ++ ++def _batch(req_ids: list[str], structured: bool) -> SimpleNamespace: ++ return SimpleNamespace(req_ids=list(req_ids), has_structured_output_reqs=structured) ++ ++ ++def _drafts(start: int, num_reqs: int, device: torch.device) -> torch.Tensor: ++ return torch.arange( ++ start, start + num_reqs * K, dtype=torch.int32, device=device ++ ).reshape(num_reqs, K) ++ ++ ++def test_consumed_snapshot_survives_later_batch(): ++ """Link A: a later execute_model must not erase an earlier step's drafts.""" ++ device = _device() ++ handler = DraftTokensHandler(device, max_snapshots=4) ++ ++ drafts_a = _drafts(1, 2, device) ++ drafts_b = _drafts(101, 2, device) ++ ++ handler.snapshot_consumed_drafts( ++ step_id=1, ++ input_batch=_batch(["gA", "pA"], True), ++ draft_tokens=drafts_a, ++ ) ++ got = handler.get_draft_tokens(step_id=1) ++ assert got is not None ++ assert got.req_ids == ["gA", "pA"] ++ assert got.draft_token_ids[0] == [1, 2, 3, 4] ++ ++ handler.snapshot_consumed_drafts( ++ step_id=2, ++ input_batch=_batch(["gB", "pB"], True), ++ draft_tokens=drafts_b, ++ ) ++ got_b = handler.get_draft_tokens(step_id=2) ++ assert got_b is not None ++ assert got_b.req_ids == ["gB", "pB"] ++ ++ got_a = handler.get_draft_tokens(step_id=1) ++ assert got_a is not None ++ assert got_a.req_ids == ["gA", "pA"] ++ assert got_a.draft_token_ids[0] == [1, 2, 3, 4] ++ ++ ++def test_plain_batch_does_not_evict_structured_snapshot(): ++ device = _device() ++ handler = DraftTokensHandler(device, max_snapshots=4) ++ drafts_a = _drafts(1, 2, device) ++ drafts_plain = _drafts(201, 2, device) ++ ++ handler.snapshot_consumed_drafts( ++ step_id=1, ++ input_batch=_batch(["gA", "pA"], True), ++ draft_tokens=drafts_a, ++ ) ++ handler.snapshot_consumed_drafts( ++ step_id=2, ++ input_batch=_batch(["p1", "p2"], False), ++ draft_tokens=drafts_plain, ++ ) ++ ++ got = handler.get_draft_tokens(step_id=1) ++ assert got is not None ++ assert got.req_ids == ["gA", "pA"] ++ assert handler.get_draft_tokens(step_id=2) is None ++ ++ ++def test_unknown_step_returns_none_for_fail_closed_fallback(): ++ device = _device() ++ handler = DraftTokensHandler(device) ++ handler.snapshot_consumed_drafts( ++ step_id=3, ++ input_batch=_batch(["gA"], True), ++ draft_tokens=_drafts(1, 1, device), ++ ) ++ ++ assert handler.get_draft_tokens(step_id=99) is None ++ ++ ++def test_proposal_slot_without_step_id_stays_latest(): ++ """Non-async post_step still reads the latest proposed drafts.""" ++ device = _device() ++ handler = DraftTokensHandler(device) ++ handler.set_draft_tokens(_batch(["gA", "pA"], True), _drafts(1, 2, device)) ++ handler.set_draft_tokens(_batch(["gB", "pB"], True), _drafts(101, 2, device)) ++ ++ got = handler.get_draft_tokens() ++ assert got is not None ++ assert got.req_ids == ["gB", "pB"] ++ assert got.draft_token_ids[0] == [101, 102, 103, 104] ++ ++ ++def test_proposal_plain_batch_returns_placeholders(): ++ device = _device() ++ handler = DraftTokensHandler(device) ++ handler.set_draft_tokens(_batch(["p1", "p2"], False), _drafts(1, 2, device)) ++ ++ got = handler.get_draft_tokens() ++ assert got is not None ++ assert got.req_ids == ["p1", "p2"] ++ assert got.draft_token_ids == [[-1] * K, [-1] * K] ++ ++ ++def test_consumed_lookup_ignores_proposal_slot(): ++ device = _device() ++ handler = DraftTokensHandler(device) ++ handler.snapshot_consumed_drafts( ++ step_id=5, ++ input_batch=_batch(["gA"], True), ++ draft_tokens=_drafts(1, 1, device), ++ ) ++ handler.set_draft_tokens(_batch(["gB"], True), _drafts(101, 1, device)) ++ ++ got = handler.get_draft_tokens(step_id=5) ++ assert got is not None ++ assert got.req_ids == ["gA"] ++ assert got.draft_token_ids[0] == [1, 2, 3, 4] ++ ++ ++def test_ring_buffer_evicts_oldest_and_misses(): ++ device = _device() ++ handler = DraftTokensHandler(device, max_snapshots=2) ++ handler.snapshot_consumed_drafts( ++ step_id=1, ++ input_batch=_batch(["a"], True), ++ draft_tokens=_drafts(1, 1, device), ++ ) ++ handler.snapshot_consumed_drafts( ++ step_id=2, ++ input_batch=_batch(["b"], True), ++ draft_tokens=_drafts(11, 1, device), ++ ) ++ handler.snapshot_consumed_drafts( ++ step_id=3, ++ input_batch=_batch(["c"], True), ++ draft_tokens=_drafts(21, 1, device), ++ ) ++ ++ assert handler.get_draft_tokens(step_id=1) is None ++ got = handler.get_draft_tokens(step_id=2) ++ assert got is not None ++ assert got.req_ids == ["b"] ++ got = handler.get_draft_tokens(step_id=3) ++ assert got is not None ++ assert got.req_ids == ["c"] ++ ++ ++def test_worker_wrapper_drops_unsupported_step_id(): ++ wrapper = WorkerWrapperBase(rpc_rank=0) ++ seen: list[tuple] = [] ++ ++ class PluginWorker: ++ def take_draft_token_ids(self): ++ seen.append(()) ++ return "legacy" ++ ++ wrapper.worker = PluginWorker() ++ assert wrapper.take_draft_token_ids(step_id=7) == "legacy" ++ assert seen == [()] ++ ++ ++def test_worker_wrapper_forwards_step_id_when_supported(): ++ wrapper = WorkerWrapperBase(rpc_rank=0) ++ seen: list[int | None] = [] ++ ++ class Worker: ++ def take_draft_token_ids(self, step_id: int | None = None): ++ seen.append(step_id) ++ return "ok" ++ ++ wrapper.worker = Worker() ++ assert wrapper.take_draft_token_ids(step_id=7) == "ok" ++ assert wrapper.take_draft_token_ids() == "ok" ++ assert seen == [7, None] ++ ++ ++def test_worker_wrapper_drops_positional_only_step_id(): ++ wrapper = WorkerWrapperBase(rpc_rank=0) ++ seen: list[tuple] = [] ++ ++ class PluginWorker: ++ def take_draft_token_ids(self, step_id: int | None = None, /): ++ seen.append((step_id,)) ++ return "legacy" ++ ++ wrapper.worker = PluginWorker() ++ assert wrapper.take_draft_token_ids(step_id=7) == "legacy" ++ assert seen == [(None,)] +diff --git a/tests/v1/worker/test_grammar_fail_closed.py b/tests/v1/worker/test_grammar_fail_closed.py +new file mode 100644 +index 0000000..869cb98 +--- /dev/null ++++ b/tests/v1/worker/test_grammar_fail_closed.py +@@ -0,0 +1,151 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Acceptance must stop at the last bitmask row that carries a real mask. ++ ++`grammar_bitmask` stops constraining rows once it meets a -1 placeholder in the ++scheduled draft window, and the worker verifies the drafts it holds on the GPU ++rather than the scheduler's padded copy. When the two disagree -- the hand-off ++lost a step's drafts, see #54437 -- acceptance can walk into a row where every ++token is allowed. ++ ++The grammar here admits only LEGAL. The scheduled window is back-filled one ++draft deep, so rows 0 and 1 carry a real mask and rows 2.. are `_full_mask`, ++while the drafts on the GPU would carry acceptance into row 2 and emit ILLEGAL. ++""" ++ ++from types import SimpleNamespace ++ ++import numpy as np ++import pytest ++import torch ++ ++from vllm.platforms import current_platform ++from vllm.utils.math_utils import cdiv ++from vllm.v1.core.sched.output import GrammarOutput ++from vllm.v1.worker.gpu.model_runner import GPUModelRunner ++from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS ++from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler ++from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker ++ ++VOCAB = 128 ++K = 3 # drafts per step ++NUM_LOGITS = K + 1 # + the bonus row ++LEGAL = 7 ++ILLEGAL = 99 ++# Index of the first -1 in the scheduled window: rows 0..1 are masked. ++NUM_ACCEPTABLE = 1 ++ ++ ++def _bitmask() -> np.ndarray: ++ """Rows 0..NUM_ACCEPTABLE constrained to LEGAL, the rest fully permissive.""" ++ mask = np.zeros((NUM_LOGITS, cdiv(VOCAB, 32)), dtype=np.int32) ++ for row in range(NUM_LOGITS): ++ if row <= NUM_ACCEPTABLE: ++ mask[row, LEGAL // 32] = 1 << (LEGAL % 32) ++ else: ++ mask[row, :] = -1 # every token allowed ++ return mask ++ ++ ++def _input_batch(device: torch.device) -> SimpleNamespace: ++ cu_np = np.array([0, NUM_LOGITS], dtype=np.int32) ++ # draft_sampled = input_ids[logits_indices]; entry i + 1 is draft i, so the ++ # drafts the GPU holds are [LEGAL, LEGAL, ILLEGAL]. ++ input_ids = torch.tensor( ++ [LEGAL, LEGAL, LEGAL, ILLEGAL], dtype=torch.int32, device=device ++ ) ++ return SimpleNamespace( ++ num_reqs=1, ++ req_ids=["g"], ++ num_draft_tokens=K, ++ num_draft_tokens_per_req=np.array([K], dtype=np.int32), ++ input_ids=input_ids, ++ logits_indices=torch.arange(NUM_LOGITS, dtype=torch.int32, device=device), ++ positions=torch.arange(NUM_LOGITS, dtype=torch.int64, device=device), ++ cu_num_logits_np=cu_np, ++ cu_num_logits=torch.from_numpy(cu_np).to(device), ++ idx_mapping_np=np.array([0], dtype=np.int32), ++ idx_mapping=torch.zeros(1, dtype=torch.int32, device=device), ++ expanded_idx_mapping=torch.zeros(NUM_LOGITS, dtype=torch.int32, device=device), ++ expanded_local_pos=torch.arange(NUM_LOGITS, dtype=torch.int32, device=device), ++ seq_lens=torch.tensor([64], dtype=torch.int32, device=device), ++ ) ++ ++ ++def _logits(device: torch.device) -> torch.Tensor: ++ """Every row prefers its own draft, so nothing but the mask can reject one.""" ++ logits = torch.zeros(NUM_LOGITS, VOCAB, dtype=torch.float32, device=device) ++ logits[0, LEGAL] = 10.0 ++ logits[1, LEGAL] = 10.0 ++ logits[2, ILLEGAL] = 10.0 # row 2 is permissive, so this wins there ++ logits[3, LEGAL] = 10.0 ++ return logits ++ ++ ++def _rejection_sampler(device: torch.device) -> RejectionSampler: ++ sampler = object.__new__(RejectionSampler) ++ sampler.num_speculative_steps = K ++ sampler.enable_adaptive_verification = False ++ sampler.use_block_verification = False ++ sampler.synthetic_conditional_rates = None ++ sampler.watermark_key = None ++ zeros_f = torch.zeros(1, dtype=torch.float32, device=device) # greedy ++ sampler.sampler = SimpleNamespace( ++ compute_nans=False, ++ logprobs_mode="raw_logprobs", ++ use_fp64_gumbel=False, ++ apply_sampling_params=lambda logits, *args, **kwargs: logits, ++ sampling_states=SimpleNamespace( ++ temperature=SimpleNamespace(gpu=zeros_f), ++ seeds=SimpleNamespace(gpu=torch.zeros(1, dtype=torch.int64, device=device)), ++ max_num_logprobs=lambda _: NO_LOGPROBS, ++ ), ++ req_states=SimpleNamespace( ++ prefill_len=SimpleNamespace( ++ gpu=torch.zeros(1, dtype=torch.int32, device=device) ++ ) ++ ), ++ ) ++ return sampler ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") ++def test_acceptance_stops_at_the_last_masked_row(): ++ device = torch.device("cuda") ++ input_batch = _input_batch(device) ++ logits = _logits(device) ++ ++ runner = object.__new__(GPUModelRunner) ++ runner.device = device ++ runner.vocab_size = VOCAB ++ runner.batch_sharder = None ++ runner.sampler = None ++ runner.model = SimpleNamespace(compute_logits=lambda hidden: logits) ++ runner.rejection_sampler = _rejection_sampler(device) ++ runner.speculator = SimpleNamespace(draft_logits=None) ++ runner.structured_outputs_worker = StructuredOutputsWorker( ++ max_num_logits=NUM_LOGITS, ++ vocab_size=VOCAB, ++ device=device, ++ mask_stride=NUM_LOGITS, ++ num_bonus_tokens=1, ++ ) ++ ++ grammar_output = GrammarOutput(["g"], _bitmask()) ++ # The scheduler reports how many leading drafts the bitmask constrained. ++ # Absent on an older scheduler, where the whole window is invalidated. ++ grammar_output.num_acceptable_drafts = [NUM_ACCEPTABLE] ++ ++ hidden = torch.zeros(NUM_LOGITS, 8, dtype=torch.float32, device=device) ++ sampler_output, num_sampled, _ = runner.sample(hidden, input_batch, grammar_output) ++ ++ accepted = int(num_sampled[0].item()) ++ emitted = sampler_output.sampled_token_ids[0, :accepted].tolist() ++ assert ILLEGAL not in emitted, ( ++ f"emitted {emitted}: a draft verified against a permissive row was " ++ f"accepted, so the request sampled with no grammar constraint" ++ ) ++ # Draft 0 may be accepted; draft 1 must not be, because accepting it moves ++ # sampling into row 2, which constrains nothing. ++ assert accepted == NUM_ACCEPTABLE + 1 ++ assert emitted == [LEGAL] * accepted +diff --git a/tests/v1/worker/test_grammar_invalid_drafts.py b/tests/v1/worker/test_grammar_invalid_drafts.py +new file mode 100644 +index 0000000..476fe76 +--- /dev/null ++++ b/tests/v1/worker/test_grammar_invalid_drafts.py +@@ -0,0 +1,76 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Which speculative drafts a structured-output request may accept. ++ ++`StructuredOutputManager.grammar_bitmask` fills row i before it inspects ++`req_tokens[i]`, so with the first -1 placeholder at index j rows 0..j carry a ++real mask and rows j+1.. carry the all-permissive `_full_mask`. Drafts j..K-1 ++must therefore be rejected, or the request samples with no grammar constraint ++at a position the model really does sample. ++""" ++ ++from unittest.mock import Mock ++ ++import numpy as np ++import torch ++ ++from vllm.v1.worker.gpu.model_runner import grammar_invalid_draft_positions ++ ++ ++def _input_batch(req_ids, num_logits_per_req, num_draft_tokens=6): ++ cu = np.concatenate([[0], np.cumsum(num_logits_per_req)]).astype(np.int32) ++ return Mock( ++ req_ids=list(req_ids), ++ cu_num_logits_np=cu, ++ num_draft_tokens=num_draft_tokens, ++ ) ++ ++ ++def _positions(batch, grammar_req_ids, num_acceptable): ++ out = grammar_invalid_draft_positions( ++ batch, grammar_req_ids, num_acceptable, torch.device("cpu") ++ ) ++ return None if out is None else out.tolist() ++ ++ ++def test_positions_reject_whole_window_when_nothing_was_backfilled(): ++ # Two requests, 4 logits rows each (1 + 3 drafts). Request "g" is at rows ++ # 0..3, so its drafts sit at draft_sampled[1..3]. ++ batch = _input_batch(["g", "p"], [4, 4]) ++ assert _positions(batch, ["g"], [0]) == [1, 2, 3] ++ ++ ++def test_positions_keep_the_drafts_the_bitmask_could_see(): ++ batch = _input_batch(["g", "p"], [4, 4]) ++ # Two real drafts: they were masked correctly and may be accepted. ++ assert _positions(batch, ["g"], [2]) == [3] ++ ++ ++def test_positions_reject_nothing_when_fully_backfilled(): ++ batch = _input_batch(["g", "p"], [4, 4]) ++ assert _positions(batch, ["g"], [3]) is None ++ ++ ++def test_positions_are_offset_per_request(): ++ batch = _input_batch(["p", "g"], [4, 4]) ++ # "g" starts at logits row 4, so its drafts are draft_sampled[5..7]. ++ assert _positions(batch, ["g"], [0]) == [5, 6, 7] ++ ++ ++def test_positions_fall_back_to_the_whole_window_without_the_field(): ++ # An older scheduler, or warmup, supplies no num_acceptable_drafts; be ++ # conservative rather than accepting drafts whose mask is unknown. ++ batch = _input_batch(["g"], [4]) ++ assert _positions(batch, ["g"], None) == [1, 2, 3] ++ ++ ++def test_no_positions_without_drafts_or_grammar_requests(): ++ batch = _input_batch(["g"], [4], num_draft_tokens=0) ++ assert _positions(batch, ["g"], [0]) is None ++ batch = _input_batch(["g"], [4]) ++ assert _positions(batch, [], []) is None ++ ++ ++def test_unknown_request_is_skipped(): ++ batch = _input_batch(["p"], [4]) ++ assert _positions(batch, ["g"], [0]) is None +diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py +index 3a46fa4..03bb2cd 100644 +--- a/vllm/v1/core/sched/output.py ++++ b/vllm/v1/core/sched/output.py +@@ -304,6 +304,10 @@ class SchedulerOutput: + # Number of spec tokens to schedule for the next step. + num_spec_tokens_to_schedule: int = 0 + ++ # Identity of this schedule() call (`Scheduler.current_step`). Used to ++ # retrieve the matching consumption-time draft-token snapshot. ++ scheduler_step: int = 0 ++ + @classmethod + def make_empty(cls) -> "SchedulerOutput": + return cls( +@@ -325,3 +329,8 @@ class GrammarOutput: + structured_output_request_ids: list[str] + # Bitmask ordered as structured_output_request_ids. + grammar_bitmask: "npt.NDArray[np.int32]" ++ # Per request, ordered as structured_output_request_ids: how many leading ++ # drafts the bitmask constrained. `grammar_bitmask` fills every row after ++ # the first -1 placeholder with the all-permissive `_full_mask`, so drafts ++ # from this index on must not be accepted. None means nothing to invalidate. ++ num_acceptable_drafts: list[int] | None = None +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 25af1b6..b16ae9d 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -74,7 +74,8 @@ from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log + from vllm.v1.request import Request, RequestStatus, StreamingUpdate + from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup + from vllm.v1.spec_decode.metrics import SpecDecodingStats +-from vllm.v1.structured_output import StructuredOutputManager ++from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager ++from vllm.v1.structured_output.utils import strip_speculative_padding + from vllm.v1.utils import record_function_or_nullcontext + + logger = init_logger(__name__) +@@ -1465,6 +1466,7 @@ class Scheduler(SchedulerInterface): + kv_connector_block_state=kv_connector_block_state, + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, + ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), ++ scheduler_step=self.current_step, + ) + + # NOTE(Kuntai): this function is designed for multiple purposes: +@@ -1941,7 +1943,14 @@ class Scheduler(SchedulerInterface): + structured_output_request_ids, + scheduler_output.scheduled_spec_decode_tokens, + ) +- return GrammarOutput(structured_output_request_ids, bitmask) ++ spec_tokens = scheduler_output.scheduled_spec_decode_tokens ++ num_acceptable_drafts = [ ++ len(strip_speculative_padding(spec_tokens.get(req_id, []))) ++ for req_id in structured_output_request_ids ++ ] ++ return GrammarOutput( ++ structured_output_request_ids, bitmask, num_acceptable_drafts ++ ) + + def update_from_output( + self, +diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py +index ffc9a5d..110f56d 100644 +--- a/vllm/v1/engine/core.py ++++ b/vllm/v1/engine/core.py +@@ -725,7 +725,9 @@ class EngineCore: + # When draft tokens are used with structured output, validate them + # before computing the grammar bitmask for the deferred request. + if self.check_for_draft_tokens: +- draft_token_ids = self.model_executor.take_draft_token_ids() ++ draft_token_ids = self.model_executor.take_draft_token_ids( ++ step_id=deferred_scheduler_output.scheduler_step ++ ) + if draft_token_ids is not None: + # Update the draft token ids in the scheduler output to + # filter out the invalid spec tokens, which will be padded +diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py +index 964f16e..a24507b 100644 +--- a/vllm/v1/executor/abstract.py ++++ b/vllm/v1/executor/abstract.py +@@ -259,8 +259,13 @@ class Executor(ABC): + def execute_dummy_batch(self) -> None: + self.collective_rpc("execute_dummy_batch") + +- def take_draft_token_ids(self) -> DraftTokenIds | None: +- output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ # Omit step_id when unset so plugin workers with the documented ++ # no-argument take_draft_token_ids() keep working. ++ kwargs = {} if step_id is None else {"step_id": step_id} ++ output: list[DraftTokenIds] = self.collective_rpc( ++ "take_draft_token_ids", kwargs=kwargs ++ ) + return output[0] + + def profile(self, is_start: bool = True, profile_prefix: str | None = None): +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index aae270f..1ccd23a 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -405,10 +405,13 @@ class MultiprocExecutor(Executor): + def execute_dummy_batch(self) -> None: + self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank) + +- def take_draft_token_ids(self) -> DraftTokenIds | None: ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: + # OPTIMIZATION: Get output only from a single worker (output_rank) ++ kwargs = {} if step_id is None else {"step_id": step_id} + return self.collective_rpc( +- "take_draft_token_ids", unique_reply_rank=self.output_rank ++ "take_draft_token_ids", ++ unique_reply_rank=self.output_rank, ++ kwargs=kwargs, + ) + + def collective_rpc( # type: ignore[override] +diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py +index da1f2c5..dc798cd 100644 +--- a/vllm/v1/executor/uniproc_executor.py ++++ b/vllm/v1/executor/uniproc_executor.py +@@ -141,8 +141,13 @@ class UniProcExecutor(Executor): + single_value=True, + ) + +- def take_draft_token_ids(self) -> DraftTokenIds | None: +- return self.collective_rpc("take_draft_token_ids", single_value=True) ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ kwargs = {} if step_id is None else {"step_id": step_id} ++ return self.collective_rpc( ++ "take_draft_token_ids", ++ kwargs=kwargs, ++ single_value=True, ++ ) + + def check_health(self) -> None: + # UniProcExecutor will always be healthy as long as +diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py +index 92122f8..7c4af7a 100644 +--- a/vllm/v1/structured_output/utils.py ++++ b/vllm/v1/structured_output/utils.py +@@ -573,3 +573,16 @@ def choice_as_grammar(choice: list[str]) -> str: + escaped_choices = (escape_ebnf_string(c) for c in choice) + grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices) + return grammar ++ ++def strip_speculative_padding(token_ids: list[int]) -> list[int]: ++ """Drop speculative-decoding padding from a token block. ++ ++ ngram and other speculative backends pad rejected draft positions with a ++ -1 sentinel. Structured-output grammars treat every entry as a real token ++ id, so the sentinels (and everything after the first one) are removed here, ++ before tokens reach any backend, rather than inside a single backend. ++ """ ++ for i, token_id in enumerate(token_ids): ++ if token_id < 0: ++ return token_ids[:i] ++ return token_ids +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index e3479c2..744fe5b 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -185,6 +185,39 @@ from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane + logger = init_logger(__name__) + + ++def grammar_invalid_draft_positions( ++ input_batch: InputBatch, ++ grammar_req_ids: list[str], ++ num_acceptable_drafts: list[int] | None, ++ device: torch.device, ++) -> torch.Tensor | None: ++ """Indices of `draft_sampled` entries whose bitmask row is permissive. ++ ++ Drafts from `num_acceptable_drafts` on were verified against rows carrying ++ `_full_mask`, so accepting one would sample with no grammar constraint. ++ Returns None when there is nothing to invalidate. ++ """ ++ if not grammar_req_ids or input_batch.num_draft_tokens == 0: ++ return None ++ cu_num_logits = input_batch.cu_num_logits_np.tolist() ++ req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)} ++ positions: list[int] = [] ++ for i, req_id in enumerate(grammar_req_ids): ++ req_idx = req_id_to_idx.get(req_id) ++ if req_idx is None: ++ continue ++ # Without the field (an older scheduler, or warmup) invalidate the whole ++ # window, which is the conservative choice. ++ num_acceptable = ( ++ num_acceptable_drafts[i] if num_acceptable_drafts is not None else 0 ++ ) ++ start = cu_num_logits[req_idx] + 1 + num_acceptable ++ positions.extend(range(start, cu_num_logits[req_idx + 1])) ++ if not positions: ++ return None ++ return torch.tensor(positions, dtype=torch.int64, device=device) ++ ++ + class GPUModelRunner(LoRAModelRunnerMixin): + def __init__(self, vllm_config: VllmConfig, device: torch.device): + self.vllm_config = vllm_config +@@ -293,7 +326,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.use_aux_hidden_state_outputs = True + + # Draft tokens propagation - for spec-dec + struct outputs. +- self.draft_tokens_handler = DraftTokensHandler(self.device) ++ self.draft_tokens_handler = DraftTokensHandler( ++ self.device, ++ max_snapshots=max(8, vllm_config.max_concurrent_batches + 1), ++ ) + + self.pcp_manager: pcp.PCPManager | None = None + +@@ -1471,6 +1507,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): + else None + ), + ) ++ if self.num_speculative_steps > 0: ++ # Snapshot the drafts that were just written into input_ids, tagged ++ # with this scheduler step. Proposal-time drafts (set later in ++ # sample_tokens) are for the next step and must not overwrite this. ++ self.draft_tokens_handler.snapshot_consumed_drafts( ++ scheduler_output.scheduler_step, ++ input_batch, ++ self.req_states.draft_tokens[input_batch.idx_mapping], ++ ) + input_batch = pcp.maybe_partition_pcp_batch( + self.pcp_manager, + input_batch, +@@ -1543,6 +1588,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + sample_hidden_states = hidden_states[input_batch.logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + ++ invalid_draft_positions = None + if grammar_output is not None: + # Apply grammar bitmask to the logits in-place. + assert self.structured_outputs_worker is not None +@@ -1552,6 +1598,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): + grammar_output.structured_output_request_ids, + grammar_output.grammar_bitmask, + ) ++ invalid_draft_positions = grammar_invalid_draft_positions( ++ input_batch, ++ grammar_output.structured_output_request_ids, ++ grammar_output.num_acceptable_drafts, ++ self.device, ++ ) + + sampler_output: SamplerOutput | None + if input_batch.num_reqs == 0: +@@ -1570,6 +1622,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + input_batch, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, ++ invalid_draft_positions, + ) + + if shard_metadata is not None: +@@ -2193,8 +2246,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): + + return async_output + +- def take_draft_token_ids(self) -> DraftTokenIds | None: +- return self.draft_tokens_handler.get_draft_tokens() ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ return self.draft_tokens_handler.get_draft_tokens(step_id=step_id) + + @torch.inference_mode() + @step_eplb_after() +diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py +index 0f484f1..07f7987 100644 +--- a/vllm/v1/worker/gpu/sample/batch_shard.py ++++ b/vllm/v1/worker/gpu/sample/batch_shard.py +@@ -335,9 +335,16 @@ def _shard_grammar_output( + cursor += num_req_logits + if not local_ids: + return None ++ # num_acceptable_drafts is ordered with structured_output_request_ids. ++ num_acceptable = grammar_output.num_acceptable_drafts ++ if num_acceptable is not None: ++ owned = set(local_ids) ++ ids = grammar_output.structured_output_request_ids ++ num_acceptable = [n for i, n in zip(ids, num_acceptable) if i in owned] + return GrammarOutput( + structured_output_request_ids=local_ids, + grammar_bitmask=grammar_output.grammar_bitmask[keep_indices], ++ num_acceptable_drafts=num_acceptable, + ) + + +diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +index 609924e..6c07e6a 100644 +--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py ++++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +@@ -176,6 +176,7 @@ class RejectionSampler: + idx_mapping_np: np.ndarray, + expanded_idx_mapping: torch.Tensor, + expanded_local_pos: torch.Tensor, ++ verify_draft_sampled: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + processed_logits = self.sampler.apply_sampling_params( + logits, +@@ -189,7 +190,7 @@ class RejectionSampler: + sampled, num_sampled = rejection_sample( + processed_logits, + draft_logits, +- draft_sampled, ++ verify_draft_sampled, + cu_num_logits, + pos, + idx_mapping, +@@ -216,7 +217,10 @@ class RejectionSampler: + pos: torch.Tensor, + max_chunk_logits: int, + max_num_logprobs: int, ++ verify_draft_sampled: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]: ++ if verify_draft_sampled is None: ++ verify_draft_sampled = draft_sampled + cu_num_logits_np = input_batch.cu_num_logits_np + use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES + num_reqs = input_batch.num_reqs +@@ -251,6 +255,7 @@ class RejectionSampler: + input_batch.idx_mapping_np[start:end], + input_batch.expanded_idx_mapping[lo:hi], + input_batch.expanded_local_pos[lo:hi], ++ verify_draft_sampled[lo:hi], + ) + chunk_logprobs = self._get_logprobs_tensors( + sampled, +@@ -289,6 +294,7 @@ class RejectionSampler: + logits: torch.Tensor, + input_batch: InputBatch, + draft_logits: torch.Tensor | None = None, ++ invalid_draft_positions: torch.Tensor | None = None, + ) -> SamplerOutput: + # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear + # that num_nans is computed before applying penalties and temperature. +@@ -297,6 +303,15 @@ class RejectionSampler: + draft_sampled = input_batch.input_ids[input_batch.logits_indices] + pos = input_batch.positions[input_batch.logits_indices] + ++ # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the ++ # accepted length before it reaches a permissive bitmask row. Only ++ # verification sees the copy: `apply_sampling_params` and watermarking ++ # keep the real draft ids. ++ verify_draft_sampled = draft_sampled ++ if invalid_draft_positions is not None: ++ verify_draft_sampled = draft_sampled.clone() ++ verify_draft_sampled[invalid_draft_positions] = -1 ++ + max_num_logprobs = self.sampler.sampling_states.max_num_logprobs( + input_batch.idx_mapping_np + ) +@@ -309,6 +324,7 @@ class RejectionSampler: + pos, + chunk_logit_limit, + max_num_logprobs, ++ verify_draft_sampled, + ) + + num_sampled, num_rejected = get_num_sampled_and_rejected( +diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py +index f73acb2..65ad160 100644 +--- a/vllm/v1/worker/gpu/spec_decode/utils.py ++++ b/vllm/v1/worker/gpu/spec_decode/utils.py +@@ -1,5 +1,8 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++from collections import deque ++from dataclasses import dataclass ++ + import numpy as np + import torch + +@@ -10,6 +13,10 @@ from vllm.v1.outputs import DraftTokenIds + from vllm.v1.worker.gpu.async_utils import async_copy_to_np + from vllm.v1.worker.gpu.input_batch import InputBatch + ++# Keep enough consumption-time snapshots to cover async PP in-flight ++# batches (`max_concurrent_batches` is `pp_size + 1`) with headroom. ++_DEFAULT_MAX_SNAPSHOTS = 8 ++ + logger = init_logger(__name__) + + +@@ -26,48 +33,135 @@ def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig: + return load_config + + ++@dataclass ++class _DraftSnapshot: ++ step_id: int | None ++ req_ids: list[str] ++ num_draft_tokens: int ++ draft_tokens_np: np.ndarray | None ++ copy_event: torch.cuda.Event | None ++ ++ def to_draft_token_ids(self) -> DraftTokenIds: ++ if self.draft_tokens_np is not None: ++ if self.copy_event is not None: ++ self.copy_event.synchronize() ++ draft_token_ids = self.draft_tokens_np.tolist() ++ else: ++ draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] ++ return DraftTokenIds(self.req_ids, draft_token_ids) ++ ++ + class DraftTokensHandler: +- def __init__(self, device: torch.device | None = None): ++ def __init__( ++ self, ++ device: torch.device | None = None, ++ max_snapshots: int = _DEFAULT_MAX_SNAPSHOTS, ++ ): + self.device = device +- self.copy_stream = torch.cuda.Stream(device) +- # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. +- self.copy_event = torch.cuda.Event(blocking=True) ++ self.copy_stream: torch.cuda.Stream | None = None ++ if device is not None and device.type == "cuda": ++ self.copy_stream = torch.cuda.Stream(device) ++ ++ self._max_snapshots = max(1, max_snapshots) ++ # Consumption-time snapshots tagged with scheduler step identity. ++ self._consumed: deque[_DraftSnapshot] = deque(maxlen=self._max_snapshots) ++ # Latest proposal-time drafts for the non-async post_step path. ++ self._proposed: _DraftSnapshot | None = None ++ ++ def _copy_drafts( ++ self, draft_tokens: torch.Tensor ++ ) -> tuple[np.ndarray, torch.cuda.Event | None]: ++ if self.copy_stream is None: ++ return draft_tokens.detach().cpu().numpy(), None + +- self.req_ids: list[str] = [] +- self.draft_tokens_np: np.ndarray | None = None +- self.num_draft_tokens: int = 0 ++ current_stream = torch.cuda.current_stream(self.device) ++ self.copy_stream.wait_stream(current_stream) ++ copy_event = torch.cuda.Event(blocking=True) ++ with torch.cuda.stream(self.copy_stream): ++ draft_tokens_np = async_copy_to_np(draft_tokens) ++ # draft_tokens is a temporary allocation on the main stream and ++ # read here on copy_stream; without record_stream, the caching ++ # allocator may reuse its memory before the async copy executes. ++ draft_tokens.record_stream(self.copy_stream) ++ copy_event.record() ++ return draft_tokens_np, copy_event ++ ++ def snapshot_consumed_drafts( ++ self, ++ step_id: int, ++ input_batch: InputBatch, ++ draft_tokens: torch.Tensor, ++ ) -> None: ++ """Record drafts that prepare_inputs fed into input_ids for this step. ++ ++ Tagged with the scheduler step so the engine can retrieve the matching ++ snapshot later by equality. Batches without structured-output requests ++ are not stored, so they cannot evict a useful snapshot. ++ """ ++ if not input_batch.has_structured_output_reqs: ++ return ++ ++ draft_tokens_np, copy_event = self._copy_drafts(draft_tokens) ++ self._consumed.append( ++ _DraftSnapshot( ++ step_id=step_id, ++ req_ids=list(input_batch.req_ids), ++ num_draft_tokens=draft_tokens.shape[1], ++ draft_tokens_np=draft_tokens_np, ++ copy_event=copy_event, ++ ) ++ ) + + def set_draft_tokens( + self, input_batch: InputBatch, draft_tokens: torch.Tensor + ) -> None: +- self.req_ids = input_batch.req_ids +- self.num_draft_tokens = draft_tokens.shape[1] ++ req_ids = list(input_batch.req_ids) ++ num_draft_tokens = draft_tokens.shape[1] + if not input_batch.has_structured_output_reqs: + # No draft token validation needs to be performed by + # the scheduler for this batch. +- self.draft_tokens_np = None ++ self._proposed = _DraftSnapshot( ++ step_id=None, ++ req_ids=req_ids, ++ num_draft_tokens=num_draft_tokens, ++ draft_tokens_np=None, ++ copy_event=None, ++ ) + return + + # For spec decoding + structured outputs, we must transfer the + # draft tokens back to the scheduler for grammar validation. +- current_stream = torch.cuda.current_stream(self.device) +- self.copy_stream.wait_stream(current_stream) +- with torch.cuda.stream(self.copy_stream): +- self.draft_tokens_np = async_copy_to_np(draft_tokens) +- # draft_tokens is a temporary allocation on the main stream and read here on +- # copy_stream; without record_stream, the caching allocator may reuse its +- # memory before the async copy executes. +- draft_tokens.record_stream(self.copy_stream) +- self.copy_event.record() ++ draft_tokens_np, copy_event = self._copy_drafts(draft_tokens) ++ self._proposed = _DraftSnapshot( ++ step_id=None, ++ req_ids=req_ids, ++ num_draft_tokens=num_draft_tokens, ++ draft_tokens_np=draft_tokens_np, ++ copy_event=copy_event, ++ ) + +- def get_draft_tokens(self) -> DraftTokenIds | None: +- if self.draft_tokens_np is not None: +- self.copy_event.synchronize() +- draft_token_ids = self.draft_tokens_np.tolist() +- else: +- # This case only happens when async scheduling is disabled. +- draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] +- return DraftTokenIds(self.req_ids, draft_token_ids) ++ def get_draft_tokens(self, step_id: int | None = None) -> DraftTokenIds | None: ++ """Return draft tokens for grammar validation. ++ ++ If ``step_id`` is set, look up the consumption-time snapshot for that ++ scheduler step by equality. A miss returns None so placeholders stay ++ -1 and fail-closed invalidation can pin acceptance via is_valid_draft. ++ If ``step_id`` is None, return the latest proposal-time drafts (the ++ non-async post_step path). ++ """ ++ if step_id is not None: ++ for snapshot in reversed(self._consumed): ++ if snapshot.step_id == step_id: ++ return snapshot.to_draft_token_ids() ++ logger.debug( ++ "Draft-token snapshot miss for scheduler step %s", ++ step_id, ++ ) ++ return None ++ ++ if self._proposed is None: ++ return DraftTokenIds([], []) ++ return self._proposed.to_draft_token_ids() + + + def get_parallel_drafting_token_id(hf_config) -> int: +diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py +index 97a59bd..743c593 100644 +--- a/vllm/v1/worker/gpu_model_runner.py ++++ b/vllm/v1/worker/gpu_model_runner.py +@@ -4895,7 +4895,8 @@ class GPUModelRunner( + self.input_batch.num_tokens_no_spec[i] = pos + 1 + self.input_batch.prev_req_id_to_index = prev_req_id_to_index + +- def take_draft_token_ids(self) -> DraftTokenIds | None: ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ del step_id # V2-only identity; V1 keeps a single proposal slot. + if not self.num_spec_tokens or not self._draft_token_req_ids: + return None + draft_token_ids, req_ids = self._get_draft_token_ids_cpu() +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index eead9c9..949f299 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -1362,8 +1362,8 @@ class Worker(WorkerBase): + return self.model_runner.pool() # type: ignore + return None + +- def take_draft_token_ids(self) -> DraftTokenIds | None: +- return self.model_runner.take_draft_token_ids() ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ return self.model_runner.take_draft_token_ids(step_id=step_id) + + def profile(self, is_start: bool = True, profile_prefix: str | None = None): + # Check if profiling is enabled +diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py +index cf5b0e7..3bf7ca6 100644 +--- a/vllm/v1/worker/worker_base.py ++++ b/vllm/v1/worker/worker_base.py +@@ -1,6 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++import inspect + from collections.abc import Callable + from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar + +@@ -24,18 +25,34 @@ from vllm.v1.kv_cache_interface import KVCacheSpec + + if TYPE_CHECKING: + from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +- from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput ++ from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput + else: + SchedulerOutput = object + GrammarOutput = object + AsyncModelRunnerOutput = object + ModelRunnerOutput = object ++ DraftTokenIds = object + + logger = init_logger(__name__) + + _R = TypeVar("_R") + + ++def _callable_accepts_kwarg(func: Callable[..., Any], name: str) -> bool: ++ """Return True if ``func`` can be called with ``name`` as a keyword.""" ++ try: ++ params = inspect.signature(func).parameters ++ except (TypeError, ValueError): ++ return False ++ param = params.get(name) ++ if param is not None: ++ return param.kind in ( ++ inspect.Parameter.POSITIONAL_OR_KEYWORD, ++ inspect.Parameter.KEYWORD_ONLY, ++ ) ++ return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) ++ ++ + class CompilationTimes(NamedTuple): + language_model: float + encoder: float +@@ -357,6 +374,20 @@ class WorkerWrapperBase: + def __getattr__(self, attr: str): + return getattr(self.worker, attr) + ++ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None: ++ """Forward draft-token lookup, dropping ``step_id`` if unsupported. ++ ++ Plugin workers are documented to implement no-argument ++ ``take_draft_token_ids()``. Passing ``step_id`` to those methods would ++ raise TypeError; fall back to the legacy call instead. ++ """ ++ method = getattr(self.worker, "take_draft_token_ids", None) ++ if method is None: ++ return None ++ if step_id is not None and _callable_accepts_kwarg(method, "step_id"): ++ return method(step_id=step_id) ++ return method() ++ + def _apply_mm_cache(self, scheduler_output: SchedulerOutput) -> None: + mm_cache = self.mm_receiver_cache + if mm_cache is None: +-- +2.43.0 + diff --git a/patches/0022-topk-exact-restart-of-the-medium-path-on-candidate-b.patch b/patches/0022-topk-exact-restart-of-the-medium-path-on-candidate-b.patch new file mode 100644 index 0000000..0bf594e --- /dev/null +++ b/patches/0022-topk-exact-restart-of-the-medium-path-on-candidate-b.patch @@ -0,0 +1,1173 @@ +From e527e480772053c6c8aa81018c6fe0487c1ce1a2 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 022/103] topk: exact restart of the medium path on + candidate-buffer overflow + +When the persistent top-k medium-path candidate buffer overflows, the row is +restarted exactly instead of being truncated. Capture-safe counters with atomic +publication and a coherent snapshot record every restart. +--- + csrc/libtorch_stable/ops.h | 4 + + csrc/libtorch_stable/persistent_topk.cuh | 461 ++++++++++++++++++++++- + csrc/libtorch_stable/topk.cu | 272 +++++++++++++ + csrc/libtorch_stable/torch_bindings.cpp | 6 + + tests/kernels/test_top_k_per_row.py | 263 +++++++++++++ + 5 files changed, 989 insertions(+), 17 deletions(-) + +diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h +index e19f49d..ea81216 100644 +--- a/csrc/libtorch_stable/ops.h ++++ b/csrc/libtorch_stable/ops.h +@@ -474,6 +474,10 @@ void persistent_topk(const torch::stable::Tensor& logits, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + ++// Medium-path exact-restart counters; see topk.cu for the slot layout. ++void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, ++ int64_t sync); ++ + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 9ae0fd4..8bfb4b9 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -27,7 +27,9 @@ constexpr int RADIX = 256; + // which would inflate the kernel's smem footprint and kill occupancy + // for the decode/trivial paths). + constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 +-constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 ++// 6 scalars, not 5: the medium-path exact restart adds a per-row "overflow already counted" flag. ++// kMediumHeaderSize rounds to 128 B either way, so kSmemMedium is unchanged. ++constexpr size_t kMediumScalarsBytes = 6 * sizeof(int); // 24 + constexpr size_t kMediumHeaderSize = + (kMediumHistBytes + kMediumScalarsBytes + 127) & ~size_t(127); // 3200 + constexpr int MAX_BUFFERED_ITEMS = 4096; +@@ -138,6 +140,13 @@ struct PersistentTopKParams { + uint32_t chunk_size; // large path: elements per CTA + uint32_t ctas_per_group; // 1=medium, >1=large + uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit) ++ // Medium-path exact top-k restart. `exact_medium` is the ++ // MERLIN_TOPK_EXACT_MEDIUM switch, resolved once on the host so it is a ++ // launch constant and is baked into a cudagraph at capture. `exact_counters` ++ // is a process-lifetime device buffer of kNumTopKExactCounters uint32s, or ++ // nullptr; the kernel only ever atomicAdds into it, never reads back. ++ bool exact_medium; ++ unsigned int* exact_counters; + }; + + // ============================================================================ +@@ -411,6 +420,348 @@ __device__ __noinline__ void histogram_2048_topk( + } + } + ++// ============================================================================ ++// Medium-path exact top-k restart ++// ++// The medium path buffers the coarse threshold bin's candidates into a fixed ++// 4,096-slot shared buffer and increments the NEXT-LEVEL histogram *inside* ++// the capacity guard, so a candidate dropped on overflow is invisible to ++// refinement in both the buffer and the histogram. At 51,481 prompt tokens ++// (~12.9K candidate blocks) that makes the returned set non-exact on 88 % of ++// the rows that overflow, and which candidates survive is decided by atomicAdd ++// arrival order. ++// ++// The repair is the mechanism of vLLM PR #52149 with a different emission ++// rule: detect the overflow, then throw away everything the buffered path did ++// and re-select from scratch (`prefix = 0`, `remaining = TopK`) with a ++// full-row radix rescan. Because the restart starts from an empty prefix and ++// an undebited `remaining`, the coarse pass's winners are re-counted rather ++// than double-counted -- the "coarse handoff" defect of an earlier revision ++// cannot occur by construction. ++// ++// Where we diverge from #52149: its rescan emits by raw `atomicAdd` arrival ++// order in both the `> pivot` and `== pivot` phases, which is set-exact but ++// order- and tie-membership-nondeterministic. Sparse attention sums the ++// selected keys in output order, so that would re-introduce exactly the ++// nondeterminism MERLIN_QSA_SORT / MERLIN_TIE_DET / MERLIN_TIE_RECENT exist ++// to remove. Ours emits in **ascending index order** via a block scan (no ++// atomics anywhere on the placement path) and cuts ties at the k-th value by ++// **most-recent wins** (the highest element indices are kept), matching the ++// rule measured on the reference host (+3.4 % tokens/step at 12K over ++// oldest-wins). ++// ============================================================================ ++ ++// Abort the kernel when the exact fallback cannot complete, instead of ++// returning a truncated row. Compile-time so it costs nothing when off; the ++// counters below record the fault either way, so a build with the trap off is ++// still never silently wrong. ++#ifndef MERLIN_TOPK_EXACT_TRAP ++ #define MERLIN_TOPK_EXACT_TRAP 1 ++#endif ++ ++// Device-global counters. Plain device atomics on a process-lifetime buffer: ++// no host sync and no host-visible allocation on the launch path, so this is ++// safe to record inside a cudagraph capture and on every replay. Read from ++// Python with `torch.ops._C.persistent_topk_exact_counters`. ++enum TopKExactCounter : int { ++ kCtrMediumRows = 0, // denominator: medium-path rows entered ++ kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop ++ kCtrMediumExactFallback = 2, // rows that took the exact restart ++ kCtrMediumExactFault = 3, // exact restart could not complete ++ kNumTopKExactCounters = 4, ++}; ++ ++// Exact, bounded-memory re-selection for a single medium-path row. ++// ++// `scores` is already offset to the row. Every radix round rescans the whole ++// row, so correctness does not depend on how many values share a coarse bin. ++// Entered only after an overflow has been detected, which is 0.1-0.2 % of ++// medium rows at 38K-51K prompt tokens. ++template ++__device__ __noinline__ void medium_exact_restart( ++ const float* __restrict__ scores, int32_t* __restrict__ output, ++ uint32_t length, unsigned int* __restrict__ counters) { ++ static_assert(kThreadsPerBlock >= RADIX, ++ "the single-thread bin walk assumes RADIX <= blockDim.x"); ++ ++ using BlockScan = cub::BlockScan; ++ ++ struct ExactSmem { ++ uint32_t histogram[RADIX]; ++ typename BlockScan::TempStorage scan; ++ uint32_t prefix; // ordered-uint radix prefix, grows 8 bits per round ++ uint32_t remaining; // how many of the pivot-equal group must be kept ++ uint32_t tie_total; // how many elements have key == pivot in total ++ uint32_t tie_cut; // lowest element index kept among the pivot-equal ++ uint32_t fault; // set by the walks; block-wide, read after a barrier ++ }; ++ ++ // The medium path's dynamic smem window is kSmemMedium (35,968 B) and the ++ // buffered state in it is dead the moment we restart, so we overlay it. ++ // Declared with an explicit alignment rather than inheriting the `char` ++ // array's alignment of 1 -- cub::BlockScan::TempStorage needs more. ++ extern __shared__ __align__(16) char medium_exact_smem[]; ++ static_assert(alignof(ExactSmem) <= 16, "ExactSmem over-aligned for smem"); ++ static_assert(sizeof(ExactSmem) <= kSmemMedium, ++ "exact restart does not fit the medium smem window"); ++ ++ auto* smem = reinterpret_cast(medium_exact_smem); ++ const uint32_t tx = threadIdx.x; ++ ++ if (tx == 0) { ++ smem->prefix = 0; ++ smem->remaining = static_cast(TopK); ++ smem->tie_total = 0; ++ smem->tie_cut = 0; ++ // The index-space tie walk below is two 8-bit rounds, so it resolves ++ // indices below 65,536. The medium path serves seq_len <= 32,768, but ++ // fail loudly rather than silently mis-cut if that ever changes. ++ smem->fault = (length > 65536u || length < static_cast(TopK)) ++ ? 1u ++ : 0u; ++ } ++ __syncthreads(); ++ ++ // ---- four 8-bit radix rounds over the full row ------------------------- ++#pragma unroll ++ for (uint32_t round = 0; round < 4; ++round) { ++ // An entry fault (bad length) previously still paid four ++ // full-row rescans before anyone looked at the flag. Block-uniform -- ++ // the flag is only ever written before a barrier that precedes this read. ++ if (smem->fault != 0u) break; ++ for (uint32_t bin = tx; bin < RADIX; bin += kThreadsPerBlock) { ++ smem->histogram[bin] = 0; ++ } ++ __syncthreads(); ++ ++ // Table, not `~0u << (32 - round * 8)`: the round-0 form of that ++ // expression is a 32-bit shift of a 32-bit type, which is undefined even ++ // though the ternary never selects it at runtime, and the loop is fully ++ // unrolled so the compiler does constant-fold the dead branch. ++ constexpr uint32_t kPrefixMask[4] = {0u, 0xFF000000u, 0xFFFF0000u, ++ 0xFFFFFF00u}; ++ const uint32_t shift = 24u - round * 8u; ++ const uint32_t prefix_mask = kPrefixMask[round]; ++ const uint32_t prefix = smem->prefix; ++ ++ for (uint32_t idx = tx; idx < length; idx += kThreadsPerBlock) { ++ const uint32_t ordered = convert_to_uint32_v2(scores[idx]); ++ if ((ordered & prefix_mask) == prefix) { ++ atomicAdd(&smem->histogram[(ordered >> shift) & 0xFFu], 1u); ++ } ++ } ++ __syncthreads(); ++ ++ // Single-threaded descending walk. #52149's wide-radix variant does this ++ // cooperatively and reads `remaining` in every thread before one thread ++ // writes it back with no barrier in between (a lost update on a ++ // non-atomic shared `|=`); one thread owning both the read and the write ++ // removes that class of bug rather than adding the missing barrier. ++ if (tx == 0) { ++ const uint32_t remaining = smem->remaining; ++ uint32_t count_above = 0; ++ bool found = false; ++ for (int bin = RADIX - 1; bin >= 0; --bin) { ++ const uint32_t count = smem->histogram[bin]; ++ if (count_above + count >= remaining) { ++ smem->remaining = remaining - count_above; ++ smem->prefix |= static_cast(bin) << shift; ++ if (round == 3) smem->tie_total = count; ++ found = true; ++ break; ++ } ++ count_above += count; ++ } ++ // No crossing bin means the entry invariant (length >= TopK) is broken. ++ if (!found) atomicOr(&smem->fault, 1u); ++ } ++ __syncthreads(); ++ } ++ ++ const uint32_t pivot = smem->prefix; ++ const uint32_t equal_count = smem->remaining; ++ const uint32_t tie_total = smem->tie_total; ++ ++ // `> pivot` holds exactly TopK - equal_count elements, so the selection is ++ // TopK elements exactly. equal_count <= tie_total is an invariant of the ++ // walk (count >= remaining - count_above at the crossing bin). ++ if (tx == 0 && (equal_count == 0u || equal_count > tie_total)) { ++ atomicOr(&smem->fault, 1u); ++ } ++ __syncthreads(); ++ ++ // Note: a fault raised before emission short-circuits ALL ++ // selection work -- no tie walk, no scan, no candidate index ever written ++ // -- so the row goes straight to the -1 fill below. Previously the scan ++ // ran and could write candidates that the fill then overwrote: correct in ++ // the end, but wasted work on the fault path and harder to debug. The ++ // condition is block-uniform (shared, read after the barrier above). ++ if (smem->fault == 0u) { ++ // ---- tie cut at the k-th value: most-recent wins -------------------- ++ // ++ // Keep the `equal_count` HIGHEST element indices among the pivot-equal ++ // group. Finding the cut index turns the tie rule into a per-element ++ // predicate, which is what lets the emission below be atomic-free. ++ // ++ // Exact 32-bit ties at the cutoff are rare (19 occurrences across ++ // 12 requests, one element each), so the common case is tie_total == ++ // equal_count and the whole group is kept with no extra pass. ++ if (tie_total > equal_count) { ++ uint32_t need = equal_count; ++ uint32_t high_byte = 0; ++ ++ // Round A: bucket the pivot-equal indices by their high byte. ++ for (uint32_t bin = tx; bin < RADIX; bin += kThreadsPerBlock) { ++ smem->histogram[bin] = 0; ++ } ++ __syncthreads(); ++ for (uint32_t idx = tx; idx < length; idx += kThreadsPerBlock) { ++ if (convert_to_uint32_v2(scores[idx]) == pivot) { ++ atomicAdd(&smem->histogram[(idx >> 8) & 0xFFu], 1u); ++ } ++ } ++ __syncthreads(); ++ if (tx == 0) { ++ uint32_t count_above = 0; ++ bool found = false; ++ for (int bin = RADIX - 1; bin >= 0; --bin) { ++ const uint32_t count = smem->histogram[bin]; ++ if (count_above + count >= need) { ++ smem->tie_cut = static_cast(bin); // high byte, for now ++ smem->remaining = need - count_above; // carried into round B ++ found = true; ++ break; ++ } ++ count_above += count; ++ } ++ if (!found) atomicOr(&smem->fault, 1u); ++ } ++ __syncthreads(); ++ high_byte = smem->tie_cut; ++ need = smem->remaining; ++ ++ // Round B: within that bucket, bucket by the low byte. Each (high, low) ++ // pair names one index, so every count here is 0 or 1 and the walk lands ++ // exactly on the need-th largest index. ++ for (uint32_t bin = tx; bin < RADIX; bin += kThreadsPerBlock) { ++ smem->histogram[bin] = 0; ++ } ++ __syncthreads(); ++ for (uint32_t idx = tx; idx < length; idx += kThreadsPerBlock) { ++ if (((idx >> 8) & 0xFFu) == high_byte && ++ convert_to_uint32_v2(scores[idx]) == pivot) { ++ atomicAdd(&smem->histogram[idx & 0xFFu], 1u); ++ } ++ } ++ __syncthreads(); ++ if (tx == 0) { ++ uint32_t count_above = 0; ++ bool found = false; ++ for (int bin = RADIX - 1; bin >= 0; --bin) { ++ const uint32_t count = smem->histogram[bin]; ++ if (count_above + count >= need) { ++ smem->tie_cut = (high_byte << 8) | static_cast(bin); ++ found = true; ++ break; ++ } ++ count_above += count; ++ } ++ if (!found) atomicOr(&smem->fault, 1u); ++ } ++ __syncthreads(); ++ } ++ } ++ ++ // Re-test AFTER the tie walks. The guard above was ++ // entered before them, so a fault raised at their bin-walk failure sites ++ // used to fall straight through into emission. No barrier is needed ++ // between the two tests: the tie walks write the flag from tx == 0 and ++ // each is followed by a barrier inside the block above. ++ if (smem->fault == 0u) { ++ const uint32_t tie_cut = smem->tie_cut; ++ ++ // ---- emission: ascending index order, no atomics --------------------- ++ // ++ // selected(i) = ordered(i) > pivot || (ordered(i) == pivot && i >= tie_cut) ++ // ++ // is a pure function of the element, so the output slot is a prefix sum ++ // over the row rather than an arrival rank. Tiles run in ascending base ++ // and within a tile thread `tx` owns `base + tx`, so output[] comes out ++ // strictly ascending in index -- bit-reproducible run to run, and a no-op ++ // for the Python-side MERLIN_QSA_SORT rather than something it corrects. ++ uint32_t base_rank = 0; ++ const uint32_t tiles = ++ (length + kThreadsPerBlock - 1u) / kThreadsPerBlock * kThreadsPerBlock; ++ for (uint32_t base = 0; base < tiles; base += kThreadsPerBlock) { ++ const uint32_t idx = base + tx; ++ uint32_t flag = 0; ++ if (idx < length) { ++ const uint32_t ordered = convert_to_uint32_v2(scores[idx]); ++ flag = ++ (ordered > pivot || (ordered == pivot && idx >= tie_cut)) ? 1u : 0u; ++ } ++ uint32_t rank = 0; ++ uint32_t total = 0; ++ BlockScan(smem->scan).ExclusiveSum(flag, rank, total); ++ if (flag != 0u) { ++ const uint32_t pos = base_rank + rank; ++ // Bounded on both emission phases -- #52149 bounds ++ // `collect_pivot_matches` but not its fused `> pivot` write, so an ++ // invariant break there becomes an out-of-bounds store rather than a ++ // detected fault. ++ if (pos < static_cast(TopK)) { ++ output[pos] = static_cast(idx); ++ } else { ++ // This is the one fault site with concurrent ++ // writers. Equal stored values do not make a data race a ++ // protocol, so make it an atomic. The tx==0 sites below and ++ // above are single-writer and separated by barriers, but they use ++ // the same primitive so the flag has one access discipline. ++ atomicOr(&smem->fault, 1u); ++ } ++ } ++ base_rank += total; ++ __syncthreads(); // BlockScan TempStorage is reused next tile ++ } ++ ++ if (tx == 0 && base_rank != static_cast(TopK)) { ++ atomicOr(&smem->fault, 1u); ++ } ++ } ++ __syncthreads(); ++ ++ // Consume the decision into a register and barrier again ++ // BEFORE anyone can leave. The persistent kernel reuses this same dynamic ++ // shared-memory window for the next row -- the large-row loader writes into ++ // it before its first barrier -- so a fast warp returning here while a slow ++ // warp still had to read the shared flag would let the next row's writes ++ // decide this row's fault branch. "Read after a barrier" is uniform only ++ // until the storage is reused; this closes that window. After the second ++ // barrier every thread already holds its copy, so no later write can ++ // change any thread's decision, and the branch below is register-uniform. ++ const uint32_t row_fault = smem->fault; ++ __syncthreads(); ++ ++ // ---- fail loud --------------------------------------------------------- ++ if (row_fault != 0u) { ++ if (tx == 0 && counters != nullptr) { ++ atomicAdd(&counters[kCtrMediumExactFault], 1u); ++ } ++ // -1 is the kernel's own padding value and the consumer rejects it ++ // (`valid = ... & (token >= 0)`), so a faulted row attends nothing and is ++ // visible downstream instead of quietly attending the wrong blocks. ++ for (uint32_t i = tx; i < static_cast(TopK); ++ i += kThreadsPerBlock) { ++ output[i] = -1; ++ } ++ __threadfence(); ++ __syncthreads(); ++#if MERLIN_TOPK_EXACT_TRAP ++ __trap(); ++#endif ++ } ++} ++ + // ============================================================================ + // Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement + // For sequences 8K < seq_len <= 64K. +@@ -424,7 +775,10 @@ __device__ __noinline__ void histogram_2048_topk( + template + __device__ __noinline__ void histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, +- int logits_offset, int seq_len) { ++ int logits_offset, int seq_len, bool exact_medium, ++ unsigned int* __restrict__ counters) { ++ // No default arguments on purpose: a new call site must state whether it ++ // wants the exact path, rather than silently inheriting the broken one. + // All shared state lives in dynamic shared memory to avoid static + extern __shared__ char medium_smem[]; + +@@ -435,6 +789,9 @@ __device__ __noinline__ void histogram_256_topk( + int& shared_threshold_bin = medium_scalars[1]; + int* shared_buffered_count = &medium_scalars[2]; + int& shared_final_k = medium_scalars[4]; ++ // One extra scalar so an overflowing row is counted once, not once per ++ // refinement pass. kMediumScalarsBytes covers it (see the constant). ++ int& shared_overflow_counted = medium_scalars[5]; + int (*buffered_indices)[MAX_BUFFERED_ITEMS] = + reinterpret_cast(medium_smem + + kMediumHeaderSize); +@@ -442,6 +799,11 @@ __device__ __noinline__ void histogram_256_topk( + const int thread_id = threadIdx.x; + int remaining_k = TopK; + ++ if (thread_id == 0) { ++ shared_overflow_counted = 0; ++ if (counters != nullptr) atomicAdd(&counters[kCtrMediumRows], 1u); ++ } ++ + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } +@@ -475,8 +837,25 @@ __device__ __noinline__ void histogram_256_topk( + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; +- shared_buffered_count[0] = 0; ++ // Medium-path restart / guard site M1 (was `:512`). shared_histogram[0] is in suffix-sum ++ // form here -- it is what the predicate just above tests -- so ++ // hist[t] - hist[t+1] is the EXACT population of the threshold bin, which ++ // is exactly what the collection loop below would buffer. The prediction ++ // is precise, not conservative, so no row pays the exact path unnecessarily ++ // and no overflowing row escapes it. ++ const int threshold_bin_count = ++ shared_histogram[0][thread_id] - shared_histogram[0][thread_id + 1]; ++ const bool will_overflow = threshold_bin_count > MAX_BUFFERED_ITEMS; ++ // Sentinel: the refinement loop's entry check reads this and cannot tell ++ // a predicted overflow from a measured one, so both take the same path. ++ shared_buffered_count[0] = ++ (will_overflow && exact_medium) ? MAX_BUFFERED_ITEMS + 1 : 0; + shared_output_count = 0; ++ // Flag only -- the counter is incremented past the early return below. ++ // When hist[t+1] == TopK exactly the coarse pass already resolves the row ++ // and the threshold bin is never collected, so a large bin population is ++ // NOT an overflow and must not be counted as one. ++ shared_overflow_counted = will_overflow ? 1 : 0; + } + __syncthreads(); + +@@ -495,25 +874,51 @@ __device__ __noinline__ void histogram_256_topk( + return; + } + ++ // Past the early return the threshold bin WILL be collected, so the M1 ++ // prediction is now a real overflow event. ++ if (thread_id == 0 && shared_overflow_counted == 1 && counters != nullptr) { ++ atomicAdd(&counters[kCtrMediumOverflow], 1u); ++ } ++ + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + +- for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { +- const float logit_value = logits[idx + logits_offset]; +- const int bin = convert_to_uint8(logit_value); +- if (bin > threshold_bin) { +- const int output_pos = atomicAdd(&shared_output_count, 1); +- output_indices[output_pos] = idx; +- } else if (bin == threshold_bin) { +- const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); +- if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { +- buffered_indices[0][buffer_pos] = idx; +- const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); +- const int next_bin = (fp32_bits >> 24) & 0xFF; +- atomicAdd(&shared_histogram[0][next_bin], 1); ++ // On a predicted overflow the collection pass is skipped whole -- ++ // including its `bin > threshold_bin` emissions. That matters: the exact ++ // restart rewrites the entire row from slot 0, so anything emitted here ++ // would be double-counted. Skipping is also why the detector cannot ship ++ // without the restart -- on its own it would emit nothing, which is a ++ // fail-loud half, never a silently-truncated one. ++ // ++ // The entry predicate must be IMMUTABLE. It used to read ++ // `shared_buffered_count[0]`, which the loop body then increments -- so with ++ // exact_medium = false (the control arm, where the count starts at 0 and a ++ // real overflow pushes it past capacity mid-loop) a late warp could evaluate ++ // the predicate after early warps had driven the count over 4096 and skip ++ // ALL of its elements, including above-threshold winners. That made `=0` a ++ // new timing-dependent truncation rather than the v2.0.1 baseline. ++ // `shared_overflow_counted` is the M1 prediction, written once before the ++ // barrier above and never touched here; combined with the kernel parameter ++ // it is block-uniform and constant for the whole loop. ++ const bool skip_collection = exact_medium && (shared_overflow_counted == 1); ++ if (__builtin_expect(!skip_collection, 1)) { ++ for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { ++ const float logit_value = logits[idx + logits_offset]; ++ const int bin = convert_to_uint8(logit_value); ++ if (bin > threshold_bin) { ++ const int output_pos = atomicAdd(&shared_output_count, 1); ++ output_indices[output_pos] = idx; ++ } else if (bin == threshold_bin) { ++ const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); ++ if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { ++ buffered_indices[0][buffer_pos] = idx; ++ const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); ++ const int next_bin = (fp32_bits >> 24) & 0xFF; ++ atomicAdd(&shared_histogram[0][next_bin], 1); ++ } + } + } + } +@@ -524,6 +929,27 @@ __device__ __noinline__ void histogram_256_topk( + const int src_buffer = pass % 2; + const int dst_buffer = src_buffer ^ 1; + const int raw_buffered = shared_buffered_count[src_buffer]; ++ // Medium-path restart / guard site M2 (was `:582`) and the M1 sentinel land here. The ++ // append's atomicAdd sits OUTSIDE the capacity guard, so raw_buffered ++ // keeps counting past capacity and a post-hoc test is exact. pass == 3 ++ // emits through shared_final_k without buffering, so every buffer filled ++ // at passes 0-2 is read back here: there is no last-pass hole. ++ if (__builtin_expect(raw_buffered > MAX_BUFFERED_ITEMS, 0)) { ++ if (thread_id == 0 && shared_overflow_counted == 0) { ++ shared_overflow_counted = 1; ++ if (counters != nullptr) atomicAdd(&counters[kCtrMediumOverflow], 1u); ++ } ++ if (exact_medium) { ++ if (thread_id == 0 && counters != nullptr) { ++ atomicAdd(&counters[kCtrMediumExactFallback], 1u); ++ } ++ __syncthreads(); // no thread may clobber medium_smem before this ++ medium_exact_restart(logits + logits_offset, output_indices, ++ static_cast(seq_len), counters); ++ return; ++ } ++ __syncthreads(); ++ } + const int num_buffered = + (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS; + +@@ -940,7 +1366,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { + histogram_2048_topk(row_input, row_output, seq_len); + } else { +- histogram_256_topk(row_input, row_output, 0, seq_len); ++ histogram_256_topk(row_input, row_output, 0, seq_len, ++ params.exact_medium, params.exact_counters); + } + } + continue; +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index b3e1a86..3419344 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -3,6 +3,9 @@ + + #include + #include ++#include ++#include ++#include + + #include "ops.h" + #include "torch_utils.h" +@@ -14,6 +17,140 @@ + namespace { + + #ifndef USE_ROCM ++ ++// --------------------------------------------------------------------------- ++// MERLIN_TOPK_EXACT_MEDIUM switch and its counters. ++// ++// The switch is resolved ONCE per process, not per launch, so it is a launch ++// constant: a cudagraph captures the value that was live at capture time and ++// every replay agrees with it. Default ON in this release; set ++// MERLIN_TOPK_EXACT_MEDIUM=0 for the A/B control arm, which restores v2.0.1 ++// behaviour exactly (the counters still record, so the control arm reports the ++// overflow incidence it is the control for). ++// --------------------------------------------------------------------------- ++bool exact_medium_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_EXACT_MEDIUM"); ++ return (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0); ++ }(); ++ return enabled; ++} ++ ++// Process-lifetime device counters, one block per device. The kernel only ++// ever atomicAdds into them, so the launch path needs no host sync and the ++// values accumulate across cudagraph replays. ++// ++// Allocation, however, is NOT capture-safe: cudaMalloc during stream capture ++// is at best not captured and at worst invalidates the capture. So: ++// - the allocation only ever happens on a call that is explicitly allowed to ++// allocate, which is a non-capturing launch or the reader op; ++// - a call made *during* capture never allocates and never latches a result, ++// so a capture that happens to be first does not permanently poison the ++// instrument; ++// - the launch path fails loudly if it is capturing with no counters, rather ++// than baking a graph that runs the patched kernel with no instrument at ++// all. A silently uninstrumented graph is the failure this is here to ++// prevent, so the check is deliberately not a warning. ++// Read (and, on the first call, allocate) with ++// torch.ops._C.persistent_topk_exact_counters, which is the explicit ++// pre-capture init path. ++constexpr int kMaxCounterDevices = 16; ++ ++enum CounterState : int { ++ kCountersUninit = 0, ++ kCountersReady = 1, ++ kCountersFailed = 2, ++}; ++ ++// Atomics, not plain globals: the fast path reads these without the mutex, and ++// a second host thread may be publishing them from the slow path. Release on ++// publish, acquire on read, and the pointer is always published BEFORE the ++// state that advertises it. ++std::atomic g_exact_counters[kMaxCounterDevices] = {}; ++std::atomic g_exact_counters_state[kMaxCounterDevices] = {}; ++std::mutex g_exact_counters_mutex; ++ ++// Lookup only. Never allocates, never blocks, safe from any context. ++unsigned int* get_exact_counters(int device_index) { ++ if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr; ++ if (g_exact_counters_state[device_index].load(std::memory_order_acquire) != ++ kCountersReady) { ++ return nullptr; ++ } ++ return g_exact_counters[device_index].load(std::memory_order_acquire); ++} ++ ++// Allocate, zero, and publish. MUST NOT be called on a capturing stream, and ++// should be called in a quiescent lifecycle phase: it allocates and blocks, ++// and neither is safe while any thread in the process is capturing a graph. ++// `stream_is_capturing` only proves the CALLING stream is idle of capture, so ++// the intended entry point is the reader op called explicitly before capture; ++// the non-capturing launch path calls this as a fallback. ++unsigned int* ensure_exact_counters(int device_index, cudaStream_t stream) { ++ if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr; ++ ++ const int state = ++ g_exact_counters_state[device_index].load(std::memory_order_acquire); ++ if (state == kCountersReady) { ++ return g_exact_counters[device_index].load(std::memory_order_acquire); ++ } ++ // A previous attempt failed for real (OOM); do not retry on every launch. ++ // Only a call that was actually allowed to allocate can reach this state, so ++ // a capture-time call never latches it. ++ if (state == kCountersFailed) return nullptr; ++ ++ std::lock_guard lock(g_exact_counters_mutex); ++ // Re-check under the lock; another thread may have published meanwhile. ++ const int locked_state = ++ g_exact_counters_state[device_index].load(std::memory_order_relaxed); ++ if (locked_state != kCountersUninit) { ++ return (locked_state == kCountersReady) ++ ? g_exact_counters[device_index].load(std::memory_order_acquire) ++ : nullptr; ++ } ++ ++ constexpr size_t bytes = ++ vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int); ++ unsigned int* ptr = nullptr; ++ if (cudaMalloc(&ptr, bytes) == cudaSuccess && ptr != nullptr) { ++ // Host acquire/release orders the HOST state word; it says ++ // nothing about the device. cudaMemset on device memory is asynchronous ++ // with respect to the host, so publishing Ready straight after it would ++ // let a first launch -- on this or any other stream -- increment counters ++ // that had not been zeroed yet, or read uninitialised bytes. Zero on a ++ // known stream and WAIT for it, so initialisation is complete before the ++ // pointer becomes visible to any consumer. ++ if (cudaMemsetAsync(ptr, 0, bytes, stream) != cudaSuccess || ++ cudaStreamSynchronize(stream) != cudaSuccess) { ++ cudaFree(ptr); ++ ptr = nullptr; ++ } ++ } else { ++ ptr = nullptr; ++ } ++ // A failed allocation leaves the error state set; clear it so the launch's ++ // own cudaGetLastError() gate does not report this as a kernel failure. ++ if (ptr == nullptr) cudaGetLastError(); ++ ++ g_exact_counters[device_index].store(ptr, std::memory_order_release); ++ g_exact_counters_state[device_index].store( ++ ptr != nullptr ? kCountersReady : kCountersFailed, ++ std::memory_order_release); ++ return ptr; ++} ++ ++// True when the stream is mid-capture, in which case nothing on this path may ++// allocate. A query failure is treated as "capturing": the conservative ++// direction is to refuse to allocate, not to risk invalidating a capture. ++bool stream_is_capturing(cudaStream_t stream) { ++ cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; ++ if (cudaStreamIsCapturing(stream, &status) != cudaSuccess) { ++ cudaGetLastError(); // do not leak this into the launch's error gate ++ return true; ++ } ++ return status != cudaStreamCaptureStatusNone; ++} ++ + template + void launch_persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, +@@ -202,6 +339,30 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + workspace.mutable_data_ptr()); + params.ctas_per_group = ctas_per_group; + params.max_seq_len = static_cast(max_seq_len); ++ params.exact_medium = exact_medium_enabled(); ++ // Allocate only off the capture path; in the ordinary engine lifecycle the ++ // eager warmup lands here first, long before capture. The explicit, ++ // quiescent-phase entry point is the reader op. ++ const int device_index = static_cast(logits.get_device_index()); ++ params.exact_counters = stream_is_capturing(stream) ++ ? get_exact_counters(device_index) ++ : ensure_exact_counters(device_index, stream); ++ // Note: this check is UNCONDITIONAL -- eager as well as ++ // capturing, and in both arms. Previously only a capturing launch was ++ // rejected, so an eager launch ran with nullptr and the reader then ++ // reported [0,0,0,0,1]: an unavailable instrument encoded as healthy zero ++ // counts, which is the exact failure the counters exist to prevent. ++ // MERLIN_TOPK_EXACT_MEDIUM=0 is NOT an escape: the control arm's whole job ++ // is to report the overflow incidence it is the control for, so it needs ++ // the counters just as much. ++ STD_TORCH_CHECK( ++ params.exact_counters != nullptr, ++ "persistent_topk: the medium-path exact-restart counters are unavailable on " ++ "device ", device_index, ++ " (not initialised before cudagraph capture, or the one-time " ++ "allocation failed). Call " ++ "torch.ops._C.persistent_topk_exact_counters(out, 0, 1) once per " ++ "device in a quiescent phase before capture."); + + #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ + do { \ +@@ -278,3 +439,114 @@ void persistent_topk(const torch::stable::Tensor& logits, + STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); + #endif + } ++ ++// Read (and, on first call, allocate) the medium-path exact-restart ++// counters. This is also the explicit pre-capture init path: call it once per ++// device before any cudagraph capture. ++// ++// `out` is an int32 CUDA tensor with at least 5 elements: ++// [0] medium-path rows entered (denominator) ++// [1] rows where a 4,096-slot guard would drop a candidate ++// [2] rows that took the exact restart (0 in the control arm) ++// [3] exact restarts that could not complete (must stay 0) ++// [4] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0 ++// ++// COHERENCY -- stated precisely, because the earlier wording overclaimed. ++// With `sync == 0` the snapshot is ordered ONLY against work already enqueued ++// on the current stream; a kernel running on another stream can land ++// increments it misses. With `sync != 0` the call synchronizes the device ++// before the copy AND again after the copy/reset have been issued, so on ++// return every increment submitted before the call is included and any reset ++// has COMPLETED. That is still conditional on the CALLER HOLDING EXCLUSIVE ++// SUBMISSION OWNERSHIP: nothing stops another thread enqueueing work after the ++// leading synchronize, and no claim of unconditional global coherence is made ++// here. A stronger contract would need an event every later producer waits ++// on; this op deliberately implements the simpler half and says so. ++// ++// `reset != 0` REQUIRES `sync != 0`. Reset between measurements, never during ++// one. ++// ++// Must not be called during cudagraph capture (it allocates, synchronizes and ++// issues a host-to-device copy). It IS the explicit, quiescent-phase ++// initialization path: call it once per device before any capture. ++void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, ++ int64_t sync) { ++#ifndef USE_ROCM ++ STD_TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor"); ++ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Int, ++ "out must be int32"); ++ // Element count alone is not enough. A CUDA scalar expanded ++ // to five elements (stride 0) satisfies numel but owns one element's ++ // storage, and the five consecutive raw int32 writes below would run past ++ // it; an ordinary strided view would receive the values in the wrong slots. ++ // Require a contiguous 1-D buffer before taking the pointer. ++ STD_TORCH_CHECK(out.dim() == 1, "out must be 1-D"); ++ STD_TORCH_CHECK(out.is_contiguous(), ++ "out must be contiguous (expanded or strided views would be " ++ "written out of bounds or into the wrong elements)"); ++ STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 1, ++ "out must hold at least ", ++ vllm::persistent::kNumTopKExactCounters + 1, " elements"); ++ STD_TORCH_CHECK( ++ reset == 0 || sync != 0, ++ "persistent_topk_exact_counters: reset requires sync != 0, otherwise the " ++ "zeroing races in-flight kernel atomics and silently drops events"); ++ ++ const torch::stable::accelerator::DeviceGuard device_guard( ++ out.get_device_index()); ++ const cudaStream_t stream = get_current_cuda_stream(); ++ STD_TORCH_CHECK(!stream_is_capturing(stream), ++ "persistent_topk_exact_counters must not be called during " ++ "cudagraph capture"); ++ constexpr size_t bytes = ++ vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int); ++ ++ int32_t* dst = out.mutable_data_ptr(); ++ unsigned int* src = ++ ensure_exact_counters(static_cast(out.get_device_index()), stream); ++ ++ // An unavailable buffer used to be reported as four zeros ++ // plus "enabled", i.e. an instrument failure indistinguishable from "no ++ // medium work and no faults". Fail instead. ++ STD_TORCH_CHECK(src != nullptr, ++ "persistent_topk_exact_counters: the counter buffer could " ++ "not be allocated on this device; the instrument is " ++ "unavailable and its absence must not read as zero counts"); ++ ++ if (sync != 0) { ++ // Everything that could still be adding to the block must have retired ++ // before the snapshot (and before a reset) for the numbers to mean ++ // anything across streams. ++ STD_TORCH_CHECK(cudaDeviceSynchronize() == cudaSuccess, ++ "exact counter synchronize failed"); ++ } ++ ++ STD_TORCH_CHECK( ++ cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream) == ++ cudaSuccess, ++ "exact counter read failed"); ++ if (reset != 0) { ++ STD_TORCH_CHECK(cudaMemsetAsync(src, 0, bytes, stream) == cudaSuccess, ++ "exact counter reset failed"); ++ } ++ ++ // Static storage: the source of an async H2D copy must outlive the copy. ++ static const int32_t enabled = exact_medium_enabled() ? 1 : 0; ++ STD_TORCH_CHECK( ++ cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, &enabled, ++ sizeof(int32_t), cudaMemcpyHostToDevice, ++ stream) == cudaSuccess, ++ "exact counter flag write failed"); ++ ++ if (sync != 0) { ++ // The leading synchronize alone left the copy and the reset ++ // merely ENQUEUED, so the caller could return and immediately race the ++ // outstanding reset from another stream. Complete the handoff before ++ // returning. ++ STD_TORCH_CHECK(cudaDeviceSynchronize() == cudaSuccess, ++ "exact counter completion synchronize failed"); ++ } ++#else ++ STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); ++#endif ++} +diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp +index 54548df..44c5f95 100644 +--- a/csrc/libtorch_stable/torch_bindings.cpp ++++ b/csrc/libtorch_stable/torch_bindings.cpp +@@ -587,6 +587,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + ++ // Read (and optionally reset) the medium-path exact-restart counters. ++ ops.def( ++ "persistent_topk_exact_counters(Tensor! out, int reset, int sync) -> ()"); ++ + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " +@@ -860,6 +864,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); ++ ops.impl("persistent_topk_exact_counters", ++ TORCH_BOX(&persistent_topk_exact_counters)); + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); + #endif +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index 5913620..70332b8 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1827,3 +1827,266 @@ def test_sparse_indexer_topk_backend_resolution() -> None: + if not is_sm100: + with pytest.raises(RuntimeError, match="SM100a/SM103a"): + resolve("deep_select") ++ ++ ++# --------------------------------------------------------------------------- ++# Medium-path candidate-buffer overflow, exact restart ++# ++# The medium path (8192 < seq_len <= 32768 compressed columns) buffers the ++# coarse threshold bin's candidates into a 4,096-slot shared buffer and ++# increments the next-level histogram inside the capacity guard, so a dropped ++# candidate is invisible to refinement. A narrow logit distribution collapses ++# far more than 4,096 values into one coarse bin and the returned set stops ++# being a top-k. ++# ++# These tests assert three things, and the third is the one upstream's own ++# regression (vLLM PR #52149) cannot see, because it sorts before comparing: ++# 1. the selected values are exactly torch.topk's, ++# 2. the instrument fired (overflow detected, exact restart taken, no fault) ++# -- without this a pass proves nothing about the patched path, ++# 3. two runs return bit-identical index tensors, in ascending index order. ++# --------------------------------------------------------------------------- ++ ++_EXACT_COUNTER_SLOTS = 5 ++ ++ ++def _read_exact_counters() -> list[int] | None: ++ """Cumulative [rows, overflow, fallbacks, faults, enabled], or None.""" ++ op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) ++ if op is None: ++ return None ++ out = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") ++ # sync=1: the counters are device-global but the snapshot is issued on one ++ # stream, so without a device synchronize it can miss increments from work ++ # still in flight elsewhere. The tests compare before/after deltas, so a ++ # stream-ordered read would be a source of flakes, not just imprecision. ++ op(out, 0, 1) ++ return out.cpu().tolist() ++ ++ ++def _run_medium_overflow_row( ++ num_rows: int, seq_len: int, top_k: int, seed: int ++) -> tuple[torch.Tensor, torch.Tensor]: ++ set_random_seed(seed) ++ # The all-overflow distribution from vLLM issue #51782 / PR #52149: every ++ # value lands in one coarse FP16 bin, so the threshold bin's population is ++ # the whole row and the 4,096-slot buffer cannot hold it. ++ logits = 1.0 + torch.rand( ++ (num_rows, seq_len), dtype=torch.float32, device="cuda" ++ ) * 0.01 ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") ++ _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, seq_len) ++ torch.accelerator.synchronize() ++ return logits, indices ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@pytest.mark.parametrize("seq_len", [12288, 32768]) ++@torch.inference_mode() ++def test_persistent_topk_medium_overflow_is_exact(top_k: int, seq_len: int) -> None: ++ """> 4,096 candidates in one coarse bin must still return a true top-k.""" ++ torch.set_default_device("cuda:0") ++ num_rows = 8 # <= 32 keeps the dispatch on persistent_topk ++ ++ before = _read_exact_counters() ++ if before is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: this build has " ++ "no persistent_topk_exact_counters op, so the exact path cannot be " ++ "shown to have run. A skip here does not count as a pass." ++ ) ++ if before[4] == 0: ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_MEDIUM=0: control arm. Exactness is not claimed " ++ "here, but a qualification run must also execute the arm with the " ++ "switch on." ++ ) ++ ++ logits, indices = _run_medium_overflow_row(num_rows, seq_len, top_k, seed=0) ++ after = _read_exact_counters() ++ ++ # The instrument must discriminate: a green assertion below is meaningless ++ # unless the overflow actually happened and the exact path actually ran. ++ # EXACT deltas, not lower bounds -- ">= num_rows" would hide duplicate ++ # counting (the per-pass double-count this branch specifically fixes) and ++ # contamination from other CUDA work on the device. ++ assert after[0] - before[0] == num_rows, ( ++ f"expected exactly {num_rows} medium rows, saw {after[0] - before[0]}" ++ ) ++ assert after[1] - before[1] == num_rows, ( ++ f"expected exactly {num_rows} overflow events, saw {after[1] - before[1]}" ++ ) ++ assert after[2] - before[2] == num_rows, ( ++ f"expected exactly {num_rows} exact fallbacks, saw " ++ f"{after[2] - before[2]}" ++ ) ++ assert after[3] == before[3], "the exact fallback reported a fault" ++ ++ for row in range(num_rows): ++ reference = torch.topk(logits[row], top_k) ++ got = indices[row].to(torch.int64) ++ ++ # Values must match exactly -- not approximately. ++ got_vals = logits[row][got].sort(descending=True).values ++ assert torch.equal(got_vals, reference.values.sort(descending=True).values), ( ++ f"row {row}: selected values are not the top-{top_k}" ++ ) ++ ++ # With float32 samples an exact tie at the cutoff is vanishingly ++ # unlikely; assert set equality only when the boundary is unambiguous, ++ # so the test never passes by accident on a tie. ++ boundary = torch.topk(logits[row], top_k + 1).values ++ if boundary[top_k - 1] != boundary[top_k]: ++ assert set(got.cpu().tolist()) == set(reference.indices.cpu().tolist()), ( ++ f"row {row}: selected index set differs from torch.topk" ++ ) ++ ++ assert len(set(got.cpu().tolist())) == top_k, f"row {row}: duplicate indices" ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_medium_overflow_is_bit_reproducible() -> None: ++ """Two identical launches must return byte-identical index tensors. ++ ++ The exact restart emits by a block-scan rank over ascending element index, ++ never by atomicAdd arrival, so the output is reproducible AND sorted. ++ Upstream's rescan is set-exact but arrival-ordered, and sparse attention ++ sums the selected keys in output order -- so order alone forks the hidden ++ state. This is the assertion that separates the two implementations. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows, seq_len, top_k = 8, 12288, 512 ++ ++ counters = _read_exact_counters() ++ if counters is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: this build has " ++ "no persistent_topk_exact_counters op, so the exact path cannot be " ++ "shown to have run. A skip here does not count as a pass." ++ ) ++ if counters[4] == 0: ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_MEDIUM=0: control arm. Emission order is not " ++ "claimed here; a qualification run must also execute the arm with " ++ "the switch on." ++ ) ++ ++ _, first = _run_medium_overflow_row(num_rows, seq_len, top_k, seed=0) ++ first = first.clone() ++ _, second = _run_medium_overflow_row(num_rows, seq_len, top_k, seed=0) ++ ++ assert torch.equal(first, second), ( ++ "two identical launches returned different index tensors" ++ ) ++ ++ # Ascending index order is the property that makes the Python-side ++ # MERLIN_QSA_SORT a no-op on these rows rather than a correction. ++ diffs = first[:, 1:].to(torch.int64) - first[:, :-1].to(torch.int64) ++ assert bool((diffs > 0).all()), "exact-path output is not ascending in index" ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: ++ """Most-recent-wins at the cutoff, with an analytically known answer. ++ ++ Random float32 values essentially never tie at the k-th place, so the ++ exactness tests above never exercise the tie rule at all -- they pass ++ whichever direction it cuts. Here every element of the row carries the ++ SAME value, so: ++ ++ * the coarse threshold bin holds the whole row (> 4,096) -> overflow, ++ * `#{key > pivot}` is 0 and the entire selection is the tie group, ++ * `tie_total` (seq_len) far exceeds `equal_count` (top_k), so both ++ index-radix tie rounds actually run, ++ * the cut spans several high-byte index buckets and lands mid-bucket. ++ ++ Under most-recent-wins the answer is forced: the top_k HIGHEST indices, ++ emitted ascending. `torch.topk` cannot supply this expectation -- with all ++ values equal its index choice is arbitrary -- so the reference is computed ++ from the rule itself. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows, seq_len, width, top_k = 4, 12345, 12352, 2048 ++ ++ before = _read_exact_counters() ++ if before is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: no " ++ "persistent_topk_exact_counters op. A skip is not a pass." ++ ) ++ if before[4] == 0: ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_MEDIUM=0: control arm, the tie rule is not " ++ "claimed here." ++ ) ++ ++ logits = torch.full((num_rows, width), 1.5, dtype=torch.float32, device="cuda") ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") ++ _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, seq_len) ++ torch.accelerator.synchronize() ++ ++ after = _read_exact_counters() ++ assert after[1] - before[1] == num_rows, "the row did not overflow as designed" ++ assert after[2] - before[2] == num_rows, "the exact restart did not run" ++ assert after[3] == before[3], "the exact fallback reported a fault" ++ ++ # The cut must genuinely exercise both tie rounds, or the test is weaker ++ # than it looks: several high-byte buckets, and not on a bucket boundary. ++ cut = seq_len - top_k ++ assert (cut >> 8) != ((seq_len - 1) >> 8), "tie cut spans only one bucket" ++ assert cut % 256 != 0, "tie cut sits on a bucket boundary" ++ ++ expected = torch.arange(cut, seq_len, dtype=torch.int32, device="cuda") ++ for row in range(num_rows): ++ assert torch.equal(indices[row], expected), ( ++ f"row {row}: most-recent-wins should select indices " ++ f"[{cut}, {seq_len}) in ascending order; got " ++ f"{indices[row][:4].tolist()}..{indices[row][-4:].tolist()}" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: ++ """The counter op writes five raw int32s; it must refuse unusable layouts. ++ ++ A CUDA scalar expanded to five elements satisfies a numel check while ++ owning one element of storage, and a strided view would receive the values ++ in the wrong slots. Both must be rejected before the raw pointer is taken. ++ `reset` without `sync` must also be rejected: zeroing while kernels are in ++ flight drops their atomics silently. ++ """ ++ torch.set_default_device("cuda:0") ++ op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) ++ if op is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: no " ++ "persistent_topk_exact_counters op. A skip is not a pass." ++ ) ++ ++ ok = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") ++ op(ok, 0, 1) # the control: this layout must be accepted ++ ++ expanded = torch.zeros(1, dtype=torch.int32, device="cuda").expand( ++ _EXACT_COUNTER_SLOTS ++ ) ++ with pytest.raises(Exception): ++ op(expanded, 0, 1) ++ ++ strided = torch.zeros(2 * _EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda")[ ++ ::2 ++ ] ++ with pytest.raises(Exception): ++ op(strided, 0, 1) ++ ++ too_small = torch.zeros(_EXACT_COUNTER_SLOTS - 1, dtype=torch.int32, device="cuda") ++ with pytest.raises(Exception): ++ op(too_small, 0, 1) ++ ++ with pytest.raises(Exception): ++ op(ok, 1, 0) # reset without sync is not a complete handoff +-- +2.43.0 + diff --git a/patches/0023-sched-preserve-per-group-eviction-order-in-deferred-.patch b/patches/0023-sched-preserve-per-group-eviction-order-in-deferred-.patch new file mode 100644 index 0000000..7820f00 --- /dev/null +++ b/patches/0023-sched-preserve-per-group-eviction-order-in-deferred-.patch @@ -0,0 +1,330 @@ +From 0271354dc4a882fdcc5ef7eca2f5c55fba812925 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 023/103] sched: preserve per-group eviction order in deferred + block free (port of upstream #48532) + +The deferred free path flattened every KV cache group into one list and +reversed it once, so group order came out backwards relative to the immediate +free path. The reversal moves into pop_blocks_for_free. +--- + .../test_partial_prefix_cache_hits.py | 5 +- + tests/v1/core/test_deferred_block_free.py | 196 +++++++++++++++++- + vllm/v1/core/kv_cache_coordinator.py | 13 +- + vllm/v1/core/kv_cache_manager.py | 6 +- + vllm/v1/core/sched/scheduler.py | 9 +- + 5 files changed, 217 insertions(+), 12 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index 393ec48..4798e0c 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -1656,10 +1656,11 @@ def test_free_cow_retained_blocks_defers_until_copy_step_processed(): + ) + free = Scheduler._free_cow_retained_blocks + +- # Copy step still in flight: deferred with its fence. ++ # Copy step still in flight: deferred with its fence, in the same order ++ # the immediate path would free them (the drain no longer reverses). + free(mock, list(blocks), fence_seq=3) + assert not freed +- assert mock.deferred_frees == deque([(3, blocks[::-1])]) ++ assert mock.deferred_frees == deque([(3, blocks)]) + + # Copy step processed: freed immediately. + mock.processed_step_seq = 3 +diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py +index 0c4738a..cbe58f8 100644 +--- a/tests/v1/core/test_deferred_block_free.py ++++ b/tests/v1/core/test_deferred_block_free.py +@@ -20,7 +20,8 @@ import pytest + import torch + + from vllm.config import VllmConfig +-from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy ++from vllm.v1.core.block_pool import BlockPool ++from vllm.v1.core.kv_cache_utils import BlockHashWithGroupId, KVCacheBlockCopy + from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ModelRunnerOutput +@@ -698,3 +699,196 @@ def test_deferred_victim_stops_the_preemption_cascade(): + out2 = scheduler.schedule() + assert first.request_id in out2.num_scheduled_tokens + assert first.status == RequestStatus.RUNNING ++ ++def _make_pool_with_cached_blocks(num_blocks: int, block_size: int = 16): ++ """Build a BlockPool holding `num_blocks` allocated, hashed blocks. ++ ++ Returns (pool, blocks_by_id). Each block gets a distinct hash so it is a ++ prefix-cache eviction candidate once freed (blocks without a hash are ++ always evicted first and would mask ordering). ++ """ ++ pool = BlockPool( ++ num_gpu_blocks=num_blocks + 1, enable_caching=True, hash_block_size=block_size ++ ) ++ blocks = pool.get_new_blocks(num_blocks) ++ for i, block in enumerate(blocks): ++ block.set_block_hash( ++ BlockHashWithGroupId(f"hash{i}".encode()), num_tokens=block_size ++ ) ++ return pool, {b.block_id: b for b in blocks} ++ ++ ++def _eviction_order(pool: BlockPool) -> list[int]: ++ """Drain the pool's free queue, returning block ids in eviction order ++ (first returned == evicted first).""" ++ order = [] ++ while pool.get_num_free_blocks() > 0: ++ order.append(pool.get_new_blocks(1)[0].block_id) ++ return order ++ ++ ++def test_deferred_free_matches_immediate_free_eviction_order(): ++ """Regression for the deferred-free path collapsing per-group eviction ++ order for multi-group (hybrid) KV caches. ++ ++ Group A owns blocks [1, 2, 3] (allocation order) and group B owns ++ [4, 5, 6]. The immediate-free path frees each group in its own ++ reversed ``free_blocks`` call; the fixed deferred-free path must ++ reproduce exactly that order, not a single flattened reversal. ++ """ ++ group_a_alloc = [1, 2, 3] ++ group_b_alloc = [4, 5, 6] ++ ++ # Immediate path: KVCacheCoordinator.free() -> one reversed free_blocks ++ # call per single-type manager (group). ++ pool_imm, by_id = _make_pool_with_cached_blocks(6) ++ for group in (group_a_alloc, group_b_alloc): ++ pool_imm.free_blocks([by_id[i] for i in reversed(group)]) ++ immediate_order = _eviction_order(pool_imm) ++ ++ # Fixed deferred path: groups are kept separate (each already reversed) ++ # and freed in their own call on drain -- mirroring the immediate path. ++ pool_fixed, by_id = _make_pool_with_cached_blocks(6) ++ ordered_groups = [list(reversed(group_a_alloc)), list(reversed(group_b_alloc))] ++ for group in ordered_groups: ++ pool_fixed.free_blocks([by_id[i] for i in group]) ++ fixed_order = _eviction_order(pool_fixed) ++ ++ assert fixed_order == immediate_order == [3, 2, 1, 6, 5, 4] ++ ++ # The old (buggy) deferred path flattened both groups then reversed once, ++ # which produces a different order. Guard against regressing back to it. ++ pool_buggy, by_id = _make_pool_with_cached_blocks(6) ++ flattened = group_a_alloc + group_b_alloc ++ pool_buggy.free_blocks([by_id[i] for i in reversed(flattened)]) ++ buggy_order = _eviction_order(pool_buggy) ++ assert buggy_order == [6, 5, 4, 3, 2, 1] ++ assert buggy_order != immediate_order ++ ++ ++def test_pop_blocks_for_free_returns_eviction_order(): ++ """``pop_blocks_for_free`` must hand back blocks already in eviction ++ order, so the deferred-free path can pass them to ``free_blocks`` ++ unchanged. Reversing the flat list instead would emit whole groups in ++ the wrong order relative to the immediate-free path.""" ++ scheduler = _create_deferring_scheduler() ++ coordinator = scheduler.kv_cache_manager.coordinator ++ ++ request = create_requests( ++ num_requests=1, ++ num_tokens=NUM_PROMPT_TOKENS, ++ max_tokens=5, ++ )[0] ++ scheduler.add_request(request) ++ scheduler.schedule() ++ ++ allocated = [ ++ list(manager.req_to_blocks[request.request_id]) ++ for manager in coordinator.single_type_managers ++ ] ++ expected = [block for group in allocated for block in reversed(group)] ++ ++ assert coordinator.pop_blocks_for_free(request.request_id) == expected ++ ++ ++def _make_two_group_kv_cache_manager(num_blocks: int, block_size: int = 16): ++ """Build a KVCacheManager with two single-type managers (a hybrid ++ full-attention + sliding-window config), so CoW copies can be queued on ++ more than one group.""" ++ from math import lcm ++ ++ from vllm.v1.core.kv_cache_manager import KVCacheManager ++ from vllm.v1.kv_cache_interface import ( ++ FullAttentionSpec, ++ KVCacheConfig, ++ KVCacheGroupSpec, ++ SlidingWindowSpec, ++ ) ++ ++ config = KVCacheConfig( ++ num_blocks=num_blocks, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ KVCacheGroupSpec( ++ ["layer1"], ++ FullAttentionSpec( ++ block_size=block_size, ++ num_kv_heads=1, ++ head_size=1, ++ dtype=torch.float32, ++ ), ++ ), ++ KVCacheGroupSpec( ++ ["layer2"], ++ SlidingWindowSpec( ++ block_size=block_size, ++ num_kv_heads=1, ++ head_size=1, ++ dtype=torch.float32, ++ sliding_window=2 * block_size, ++ ), ++ ), ++ ], ++ ) ++ return KVCacheManager( ++ config, ++ max_model_len=1024, ++ scheduler_block_size=lcm( ++ *(g.kv_cache_spec.block_size for g in config.kv_cache_groups) ++ ), ++ hash_block_size=block_size, ++ enable_caching=True, ++ ) ++ ++ ++def test_cow_retained_blocks_span_groups_and_free_order_is_stable(): ++ """CoW retentions from multiple KV cache groups are concatenated into one ++ flat list by ``take_kv_cache_block_copies``. Unlike per-request eviction, ++ ``_free_cow_retained_blocks`` never reverses, so the deferred path frees ++ that list in exactly the same order as the immediate path -- the ++ group-flattening does not reintroduce the eviction-order divergence. ++ """ ++ manager = _make_two_group_kv_cache_manager(num_blocks=20) ++ coordinator = manager.coordinator ++ assert len(coordinator.single_type_managers) == 2 ++ ++ # Queue a CoW (source, cow) pair on each of the two groups. ++ blocks = manager.block_pool.get_new_blocks(4) ++ coordinator.single_type_managers[0]._pending_cow_copies.append( ++ (blocks[0], blocks[1]) ++ ) ++ coordinator.single_type_managers[1]._pending_cow_copies.append( ++ (blocks[2], blocks[3]) ++ ) ++ ++ _copies, retained = manager.take_kv_cache_block_copies() ++ # Retentions from both groups are flattened into one list, group 0 first. ++ assert [b.block_id for b in retained] == [b.block_id for b in blocks] ++ ++ # Drive the scheduler free helper on the same retained ordering, immediate ++ # vs deferred, and confirm identical eviction order. ++ retained_ids = [b.block_id for b in retained] ++ ++ imm_scheduler = _create_deferring_scheduler() ++ imm_scheduler.defer_block_free = False ++ imm_pool, imm_by_id = _make_pool_with_cached_blocks(len(retained_ids)) ++ imm_scheduler.kv_cache_manager.block_pool = imm_pool ++ imm_scheduler._free_cow_retained_blocks( ++ [imm_by_id[i] for i in retained_ids], fence_seq=1 ++ ) ++ immediate_order = _eviction_order(imm_pool) ++ ++ def_scheduler = _create_deferring_scheduler() ++ def_pool, def_by_id = _make_pool_with_cached_blocks(len(retained_ids)) ++ def_scheduler.kv_cache_manager.block_pool = def_pool ++ # fence in the future -> deferred; then drain once the fence is processed. ++ def_scheduler.processed_step_seq = 0 ++ def_scheduler._free_cow_retained_blocks( ++ [def_by_id[i] for i in retained_ids], fence_seq=1 ++ ) ++ assert def_scheduler.deferred_frees ++ def_scheduler.processed_step_seq = 1 ++ def_scheduler._drain_deferred_frees() ++ deferred_order = _eviction_order(def_pool) ++ ++ assert deferred_order == immediate_order +diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py +index afb03ac..e3f31a6 100644 +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -394,18 +394,23 @@ class KVCacheCoordinator(ABC): + Pop the request's bookkeeping from all single-type managers and + return its blocks without returning them to the block pool. The + caller must eventually pass the returned blocks to +- `block_pool.free_blocks`, freeing them in reverse order (so that +- tail blocks are evicted first). ++ `block_pool.free_blocks` as-is: they are already ordered so that ++ tail blocks are evicted first, within each KV cache group. + + Args: + request_id: The request ID. + + Returns: +- The request's blocks in allocation order. ++ The request's blocks in eviction order (each group's blocks ++ reversed, then concatenated group by group). + """ + blocks: list[KVCacheBlock] = [] + for manager in self.single_type_managers: +- blocks.extend(manager.pop_blocks_for_free(request_id)) ++ # Reverse within each group so the concatenation is already in ++ # eviction order, matching `free()`, which frees each group in its ++ # own reversed `free_blocks()` call. Reversing the flat list ++ # instead would emit whole groups in the wrong order. ++ blocks.extend(reversed(manager.pop_blocks_for_free(request_id))) + return blocks + + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: +diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py +index 226da3a..675ca1e 100644 +--- a/vllm/v1/core/kv_cache_manager.py ++++ b/vllm/v1/core/kv_cache_manager.py +@@ -632,13 +632,15 @@ class KVCacheManager: + def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: + """Pop the request's bookkeeping and return its blocks without + returning them to the block pool. The caller must eventually free +- them in reverse order (so that tail blocks are evicted first). ++ them as-is: they are already ordered so that tail blocks are ++ evicted first, within each KV cache group. + + Args: + request: The request to pop the blocks for. + + Returns: +- The request's blocks in allocation order. ++ The request's blocks in eviction order (each group's blocks ++ reversed, then concatenated group by group). + """ + return self.coordinator.pop_blocks_for_free(request.request_id) + +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index b16ae9d..ae9eeb6 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -2710,7 +2710,9 @@ class Scheduler(SchedulerInterface): + if not self.defer_block_free or fence_seq <= self.processed_step_seq: + self.kv_cache_manager.block_pool.free_blocks(blocks) + return +- self.deferred_frees.append((fence_seq, blocks[::-1])) ++ # The immediate path above frees these as-is, so enqueue them as-is: ++ # the drain no longer reverses. ++ self.deferred_frees.append((fence_seq, blocks)) + + def _drain_deferred_frees(self): + """Return deferred blocks whose fence step has completed. +@@ -2724,8 +2726,9 @@ class Scheduler(SchedulerInterface): + if fence > self.processed_step_seq: + break + _, blocks = self.deferred_frees.popleft() +- # Free in reverse order so that the tail blocks are evicted first. +- self.kv_cache_manager.block_pool.free_blocks(reversed(blocks)) ++ # Already in eviction order (reversed per KV cache group by ++ # `pop_blocks_for_free`), so free as-is. ++ self.kv_cache_manager.block_pool.free_blocks(blocks) + + def get_num_unfinished_requests(self) -> int: + if self._pause_state == PauseState.PAUSED_ALL: +-- +2.43.0 + diff --git a/patches/0024-qsa-fp8-KV-capability-overrides-and-auto-then-restor.patch b/patches/0024-qsa-fp8-KV-capability-overrides-and-auto-then-restor.patch new file mode 100644 index 0000000..b5ef4e6 --- /dev/null +++ b/patches/0024-qsa-fp8-KV-capability-overrides-and-auto-then-restor.patch @@ -0,0 +1,121 @@ +From de31f552293348bf1b4cf02f394fef97d071031b Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 024/103] qsa: fp8 KV capability overrides and + auto-then-restore construction (from upstream #55557) + +The QSA backend answers the KV-dtype and combination capability checks for its +own kernels rather than inheriting flash-attn's. The FP8 reader is unchanged. +--- + vllm/models/qwen4_exp/nvidia/qsa.py | 74 ++++++++++++++++++++++++----- + 1 file changed, 63 insertions(+), 11 deletions(-) + +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 1b7d55f..f4fa56b 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -33,6 +33,7 @@ from vllm.logger import init_logger + from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope + from vllm.model_executor.models.qwen3_next import Qwen3NextAttention + from vllm.platforms import current_platform ++from vllm.platforms.interface import DeviceCapability + from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) +@@ -309,6 +310,35 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): + "fp8_e4m3", + ] + ++ @classmethod ++ def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: ++ # The parent additionally probes flash-attn for quantized-KV support ++ # and rejects fp8 wherever FA has no kernel for it (all of sm_86). ++ # QSA never runs flash-attn over this cache: reshape_and_cache_flash ++ # writes the e4m3 bytes and our Triton kernel dequantizes them on load, ++ # so the class list alone decides. See the impl constructor for the ++ # matching "auto"-then-restore of the same probe. ++ return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes ++ ++ @classmethod ++ def supports_combination( ++ cls, ++ head_size: int, ++ dtype: torch.dtype, ++ kv_cache_dtype: CacheDType | None, ++ block_size: int | None, ++ use_mla: bool, ++ has_sink: bool, ++ use_sparse: bool, ++ use_mm_prefix: bool, ++ device_capability: DeviceCapability, ++ ) -> str | None: ++ # Same reason: the parent's "FP8 KV cache requires FA3 on SM90 or FA4 ++ # on SM100" rejection is about flash-attn's own kernels, which QSA does ++ # not use. The real device gate for our fp8 path is ++ # `_validated_qsa_fp8_dtype`, which raises on any capability != 86. ++ return None ++ + @staticmethod + def get_name() -> str: + return "QWEN4_EXP_QSA_TRITON" +@@ -341,7 +371,20 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + supports_dcp: bool = False + supports_pcp: bool = False + +- def __init__(self, *args, **kwargs) -> None: ++ def __init__( ++ self, ++ num_heads: int, ++ head_size: int, ++ scale: float, ++ num_kv_heads: int, ++ alibi_slopes: list[float] | None, ++ sliding_window: int | None, ++ kv_cache_dtype: str, ++ logits_soft_cap: float | None = None, ++ attn_type: AttentionType = AttentionType.DECODER, ++ kv_sharing_target_layer_name: str | None = None, ++ sinks: torch.Tensor | None = None, ++ ) -> None: + # QSA never dispatches to a FlashAttention kernel: attention runs through + # forward_qsa -> qsa_sparse_paged_attention (our Triton sparse-GQA kernel) + # and the cache write goes through the CUDA reshape_and_cache_flash op. +@@ -350,16 +393,25 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + # "auto" to the base class, then restore the real dtype. The only other + # kv_cache_dtype use in that constructor is an SM90/FA4 dequant path that + # cannot trigger here, and supports_quant_query_input is set below anyway. +- requested_kv_cache_dtype = None +- if "kv_cache_dtype" in kwargs: +- requested_kv_cache_dtype = kwargs["kv_cache_dtype"] +- kwargs = {**kwargs, "kv_cache_dtype": "auto"} +- elif len(args) > 6: +- requested_kv_cache_dtype = args[6] +- args = (*args[:6], "auto", *args[7:]) +- super().__init__(*args, **kwargs) +- if requested_kv_cache_dtype is not None: +- self.kv_cache_dtype = requested_kv_cache_dtype ++ # The parameters are spelled out rather than forwarded as *args/**kwargs ++ # so the substitution names the parameter instead of indexing position 6. ++ real_kv_cache_dtype = kv_cache_dtype ++ if _is_qsa_fp8_cache_dtype(kv_cache_dtype): ++ kv_cache_dtype = "auto" ++ super().__init__( ++ num_heads, ++ head_size, ++ scale, ++ num_kv_heads, ++ alibi_slopes, ++ sliding_window, ++ kv_cache_dtype, ++ logits_soft_cap, ++ attn_type, ++ kv_sharing_target_layer_name, ++ sinks, ++ ) ++ self.kv_cache_dtype = real_kv_cache_dtype + if not is_flash_attn_varlen_func_available(): + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") + if self.dcp_world_size != 1: +-- +2.43.0 + diff --git a/patches/0025-tests-mamba-spec-decode-block-table-row-attribution-.patch b/patches/0025-tests-mamba-spec-decode-block-table-row-attribution-.patch new file mode 100644 index 0000000..a56aa36 --- /dev/null +++ b/patches/0025-tests-mamba-spec-decode-block-table-row-attribution-.patch @@ -0,0 +1,353 @@ +From c390570424760dafe3861d388a16105badd7706e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 025/103] tests: mamba spec-decode block-table row attribution + (from upstream #55506) + +Device tests with a non-identity request-slot permutation, covering the axis the +persistent block-table fix changes. +--- + .../kernels/mamba/test_precopy_mamba_align.py | 296 +++++++++++++++++- + 1 file changed, 295 insertions(+), 1 deletion(-) + +diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py +index 22b73c9..1229511 100644 +--- a/tests/kernels/mamba/test_precopy_mamba_align.py ++++ b/tests/kernels/mamba/test_precopy_mamba_align.py +@@ -21,6 +21,7 @@ The kernel must also no-op when ``src_col < 0`` (fresh request) or + + from __future__ import annotations + ++import os + from collections.abc import Callable + from types import SimpleNamespace + from typing import Any +@@ -31,7 +32,11 @@ import torch + from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils + from vllm.platforms import current_platform + from vllm.v1.worker import mamba_utils as worker_mamba_utils +-from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel ++from vllm.v1.worker.mamba_utils import ( ++ _TEMPORAL_TILES, ++ postprocess_mamba_fused_kernel, ++ precopy_mamba_align_fused_kernel, ++) + + _parametrize: Callable[..., Callable[[Any], Any]] + +@@ -42,12 +47,24 @@ try: + not current_platform.is_cuda(), + reason="precopy_mamba_align_fused_kernel needs CUDA/Triton", + ) ++ # GPU tier, same opt-in as tests/v1/worker/gpu/: the reference host always ++ # reports CUDA available and is often busy with long serving runs ++ # and sanitizer runs that must not share a device with a test suite, so ++ # `is_cuda()` alone is not a sufficient gate for a test that launches ++ # kernels over multi-MiB state pools. ++ _gpu_host_tier = pytest.mark.skipif( ++ os.environ.get("VLLM_E2_GPU_TESTS") != "1", ++ reason="set VLLM_E2_GPU_TESTS=1 to run GPU tests (shared GPU host)", ++ ) + _parametrize = pytest.mark.parametrize + except ModuleNotFoundError: # allow running directly as ``python `` + + def _cuda_required(fn): + return fn + ++ def _gpu_host_tier(fn): ++ return fn ++ + def _no_parametrize(_name, _values): + def _deco(fn): + return fn +@@ -417,6 +434,279 @@ def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bi + ] + assert fused_copy_calls == scalar_copy_calls + ++# --------------------------------------------------------------------------- ++# Block-table ROW ATTRIBUTION -- ported from upstream vllm PR #55506's two new ++# kernel tests, adapted to our row-attribution fix (persistent_block_table_ptrs + ++# BT_ROW_IS_REQ). See each docstring for what differs and why. ++# --------------------------------------------------------------------------- ++ ++# Batch order -> request-state slot, deliberately NOT the identity. ++_PERM = (2, 0, 3, 1) ++ ++ ++def _build_disjoint_block_table(num_reqs, device): ++ """[num_reqs, MAX_COLS] where every row owns ids no other row references.""" ++ num_blocks = num_reqs * MAX_COLS + 1 ++ bt = torch.empty(num_reqs, MAX_COLS, dtype=torch.int32, device=device) ++ for r in range(num_reqs): ++ bt[r] = torch.arange( ++ 1 + r * MAX_COLS, 1 + (r + 1) * MAX_COLS, dtype=torch.int32, device=device ++ ) ++ return bt, num_blocks ++ ++ ++def _assert_copies_landed(convs, ssms, conv_ref, ssm_ref): ++ for layer in range(NUM_LAYERS): ++ torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) ++ torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) ++ ++ ++def _states_equal(convs, ssms, conv_ref, ssm_ref): ++ return all( ++ torch.equal(convs[layer], conv_ref[layer]) ++ and torch.equal(ssms[layer], ssm_ref[layer]) ++ for layer in range(NUM_LAYERS) ++ ) ++ ++ ++def _permute_to_batch_order(values): ++ """Reindex a per-req-slot decision array into batch order via ``_PERM``.""" ++ return values[list(_PERM)].contiguous() ++ ++ ++@_parametrize("conv_state_dim_first", [False, True]) ++@_gpu_host_tier ++@_cuda_required ++def test_precopy_indexes_block_tables_by_batch_row(conv_state_dim_first): ++ """Block-table ROW ATTRIBUTION for ``precopy_mamba_align_fused_kernel``. ++ ++ DIVERGES FROM UPSTREAM #55506 ON PURPOSE, and asserts the opposite row. ++ Upstream's ``test_precopy_indexes_block_tables_by_req_slot`` follows from ++ upstream's fix, which repoints the single shared ``block_table_ptrs`` array ++ at the persistent per-req-slot tables -- changing what "row" means for ++ every consumer of that array, precopy included. Our fix instead adds a ++ SECOND array (``persistent_block_table_ptrs``) plus a ``BT_ROW_IS_REQ`` ++ switch used only by the deferred post-process, so precopy still reads the ++ batch-order gather, and it runs in the same step that produced that gather ++ (``prepare_attn`` -> ``preprocess_state`` -> ``run_fused_precopy``). Batch-row ++ indexing is therefore the CORRECT indexing here, not a latent bug, and a ++ future "fix" that made precopy index by ``req_idx`` would break it. ++ ++ What is ported is the shape of the check: a deliberately non-identity ++ ``idx_mapping`` over rows with disjoint block ids, which ++ ``test_precopy_matches_v1_copy_specs`` above never exercises -- it always ++ launches with an identity mapping, where the two indexings are ++ indistinguishable. The invariant asserted is "the row index must match the ++ table you were handed": decisions resolve through ``idx_mapping`` to the ++ request slot, block-table rows stay batch rows. The last assertion is the ++ discriminator: if the req-slot reading produced the same bytes, the ++ permutation would not be separating the two worlds and the test could not ++ fail. ++ """ ++ device = torch.device("cuda") ++ torch.manual_seed(0) ++ num_reqs = len(_PERM) ++ bt, num_blocks = _build_disjoint_block_table(num_reqs, device) ++ ++ # Per-slot decisions must DIFFER, or the permutation is invisible: with an ++ # identity-shaped setup every batch row performs the same copy, so reading ++ # row ``batch_idx`` instead of row ``req_idx`` still touches the right bytes ++ # in aggregate. Here only slots 0 and 3 copy, and each row owns disjoint ++ # block ids, so a row mix-up moves state onto a different physical block. ++ # slot 0: col 1 -> col 0, bias 0 (copy) ++ # slot 1: src_col == dst_col (no boundary, skip) ++ # slot 2: src_col < 0 (fresh, skip) ++ # slot 3: col 2 -> col 0, bias 1 (copy, accepted-token shift) ++ src_col = torch.tensor([1, 1, -1, 2], dtype=torch.int32, device=device) ++ dst_col = torch.tensor([0, 1, 0, 0], dtype=torch.int32, device=device) ++ bias = torch.tensor([0, 0, 0, 1], dtype=torch.int32, device=device) ++ ++ convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) ++ # Correct world: block-table row b carries the decision of slot _PERM[b]. ++ conv_ref, ssm_ref = _reference( ++ convs, ++ ssms, ++ bt.cpu(), ++ _permute_to_batch_order(src_col).cpu(), ++ _permute_to_batch_order(dst_col).cpu(), ++ _permute_to_batch_order(bias).cpu(), ++ num_reqs, ++ conv_state_dim_first, ++ ) ++ # Wrong world: row == request slot (upstream's reading), from the same ++ # pre-copy bytes. ``_reference`` clones its inputs, so both are built off ++ # the state as it stands before the launch below. ++ conv_bad, ssm_bad = _reference( ++ convs, ++ ssms, ++ bt.cpu(), ++ src_col.cpu(), ++ dst_col.cpu(), ++ bias.cpu(), ++ num_reqs, ++ conv_state_dim_first, ++ ) ++ ++ base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( ++ convs, ssms, device, conv_state_dim_first ++ ) ++ bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) ++ idx_mapping = torch.tensor(_PERM, dtype=torch.int32, device=device) ++ guard_ctr = torch.zeros(3, dtype=torch.int64, device=device) ++ ++ precopy_mamba_align_fused_kernel[(num_reqs, NUM_LAYERS * 2, 1)]( ++ dst_col, ++ src_col, ++ bias, ++ bt_ptrs, ++ bt.stride(0), ++ bt.size(0), ++ guard_ctr, ++ base, ++ blk_stride, ++ elem, ++ inner, ++ width, ++ group, ++ drc, ++ drs, ++ idx_mapping, ++ num_reqs, ++ COPY_BLOCK_SIZE=1024, ++ CONV_STATE_DIM_FIRST=conv_state_dim_first, ++ HAS_IDX_MAPPING=True, ++ TEMPORAL_TILES=1, ++ ) ++ torch.accelerator.synchronize() ++ _assert_copies_landed(convs, ssms, conv_ref, ssm_ref) ++ assert int(guard_ctr[0]) == 0, "in-range rows must not trip the bounds guard" ++ assert not _states_equal(convs, ssms, conv_bad, ssm_bad), ( ++ "both row readings produced the same bytes: the permutation is not " ++ "discriminating and this test cannot fail" ++ ) ++ ++ ++@_parametrize("conv_state_dim_first", [False, True]) ++@_gpu_host_tier ++@_cuda_required ++def test_postprocess_indexes_block_tables_by_req_slot(conv_state_dim_first): ++ """Same invariant for ``postprocess_mamba_fused_kernel``, against our fix. ++ ++ Ported from upstream #55506's test of the same name. This is the deferred ++ path where the stale batch mapping actually bites: the PP post-process is ++ enqueued at the top of ``execute_model``, before this step's gather, with ++ an ``idx_mapping`` captured ``pp_size`` steps earlier, so indexing the ++ tables by batch row walks *another request's* block ids -- freed or ++ reallocated ones -- which on a hybrid model whose unified cache layout ++ aliases every cache tensor inside one page silently poisons live recurrent ++ state (all-NaN logits, constant-token output) or faults outright. ++ ++ Adapted to our API: the fix hands this kernel the persistent per-req-slot ++ tables through a second pointer array and selects the row with ++ ``BT_ROW_IS_REQ``, so this launches ``BT_ROW_IS_REQ=True`` over a ++ persistent-shaped table. The failing control is the same launch with ++ ``BT_ROW_IS_REQ=False`` on a clone of the same pre-copy bytes: it must ++ corrupt, and it must do so WITHOUT tripping the bounds guard, because a ++ stale mapping yields a row that is wrong but still in range -- exactly the ++ class ``guard_telemetry``'s row bound cannot catch. ++ ++ Decision inputs are chosen so the copying slots move col 1 -> col 0 with ++ zero accepted-token bias: with ``num_scheduled = num_draft = 1``, ++ ``num_tokens_running_state == num_computed`` and ``new_num_computed`` is the ++ same, so ``aligned = floor(num_computed / block_size) * block_size``. ++ slot 0: computed 8 -> aligned 8, bias 0, dest 0, src 1 (copy) ++ slot 1: computed 8 -> dest 0 == src 0, bias 0 (skip) ++ slot 2: computed 5 -> aligned 0 < 5, needs_copy false (skip) ++ slot 3: computed 7 + accepted 2 -> aligned 8, bias 1, ++ dest 0, src 2 (copy) ++ """ ++ device = torch.device("cuda") ++ torch.manual_seed(0) ++ num_reqs = len(_PERM) ++ block_size = 8 ++ bt, num_blocks = _build_disjoint_block_table(num_reqs, device) ++ ++ num_accepted = torch.tensor([1, 1, 1, 2], dtype=torch.int32, device=device) ++ state_idx = torch.tensor([1, 0, 1, 2], dtype=torch.int32, device=device) ++ num_scheduled = torch.ones(num_reqs, dtype=torch.int32, device=device) ++ num_computed = torch.tensor([8, 8, 5, 7], dtype=torch.int32, device=device) ++ num_draft = torch.ones(num_reqs, dtype=torch.int32, device=device) ++ ++ # What a correct kernel must do, stated directly: slots 0 and 3 copy; slot 2 ++ # is marked fresh here because its skip comes from ``needs_copy``, which the ++ # reference models through the src_col < 0 branch. ++ ref_src = torch.tensor([1, 0, -1, 2], dtype=torch.int32) ++ ref_dst = torch.zeros(num_reqs, dtype=torch.int32) ++ ref_bias = torch.tensor([0, 0, 0, 1], dtype=torch.int32) ++ ++ convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) ++ conv_ref, ssm_ref = _reference( ++ convs, ++ ssms, ++ bt.cpu(), ++ ref_src, ++ ref_dst, ++ ref_bias, ++ num_reqs, ++ conv_state_dim_first, ++ ) ++ # The failing control runs on a clone of the same pre-copy bytes. ++ ctl_convs = [c.clone() for c in convs] ++ ctl_ssms = [s.clone() for s in ssms] ++ ++ def _launch(state_convs, state_ssms, bt_row_is_req): ++ base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( ++ state_convs, state_ssms, device, conv_state_dim_first ++ ) ++ bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) ++ idx_mapping = torch.tensor(_PERM, dtype=torch.int32, device=device) ++ guard_ctr = torch.zeros(3, dtype=torch.int64, device=device) ++ num_accepted_out = num_accepted.clone() ++ postprocess_mamba_fused_kernel[(num_reqs, NUM_LAYERS * 2, 1)]( ++ num_accepted, ++ state_idx, ++ num_scheduled, ++ num_computed, ++ num_draft, ++ bt_ptrs, ++ bt.stride(0), ++ bt.size(0), ++ guard_ctr, ++ base, ++ blk_stride, ++ elem, ++ inner, ++ width, ++ group, ++ drc, ++ drs, ++ num_accepted_out, ++ idx_mapping, ++ num_reqs, ++ block_size=block_size, ++ COPY_BLOCK_SIZE=1024, ++ CONV_STATE_DIM_FIRST=conv_state_dim_first, ++ HAS_IDX_MAPPING=True, ++ PRECOMPUTED_NEW_COMPUTED=False, ++ BT_ROW_IS_REQ=bt_row_is_req, ++ TEMPORAL_TILES=1, ++ ) ++ torch.accelerator.synchronize() ++ return guard_ctr ++ ++ guard_ctr = _launch(convs, ssms, True) ++ _assert_copies_landed(convs, ssms, conv_ref, ssm_ref) ++ assert int(guard_ctr[0]) == 0, "in-range rows must not trip the bounds guard" ++ ++ bad_guard = _launch(ctl_convs, ctl_ssms, False) ++ assert not _states_equal(ctl_convs, ctl_ssms, conv_ref, ssm_ref), ( ++ "BT_ROW_IS_REQ=False produced the correct bytes: this test cannot fail" ++ ) ++ assert int(bad_guard[0]) == 0, ( ++ "a stale-mapping row is wrong but in range; if the bounds guard fires " ++ "here the scenario is not the one #55506 describes" ++ ) ++ + + if __name__ == "__main__": + from itertools import product +@@ -434,3 +724,7 @@ if __name__ == "__main__": + f"OK num_reqs={nr} token_bias={tb} has_idx_mapping={mapping} " + f"conv_dim_first={dim_first} temporal_tiles={tt}" + ) ++ for dim_first in (False, True): ++ test_precopy_indexes_block_tables_by_batch_row(dim_first) ++ test_postprocess_indexes_block_tables_by_req_slot(dim_first) ++ print(f"OK row attribution conv_dim_first={dim_first}") +-- +2.43.0 + diff --git a/patches/0026-tests-make-the-eviction-order-and-QSA-capability-con.patch b/patches/0026-tests-make-the-eviction-order-and-QSA-capability-con.patch new file mode 100644 index 0000000..206e957 --- /dev/null +++ b/patches/0026-tests-make-the-eviction-order-and-QSA-capability-con.patch @@ -0,0 +1,416 @@ +From 77661daf6fc3e9023c7181fb558c0dcb370e6eff Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:17:29 -0400 +Subject: [PATCH 026/103] tests: make the eviction-order and QSA capability + controls discriminating + +The eviction-order test now drives both free paths through production code, and +the capability tests fail when the overrides are reverted. +--- + .../qwen4_exp/test_qsa_backend_capability.py | 176 ++++++++++++++++++ + tests/v1/core/test_deferred_block_free.py | 147 +++++++++++---- + vllm/models/qwen4_exp/nvidia/qsa.py | 35 +++- + 3 files changed, 321 insertions(+), 37 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_qsa_backend_capability.py + +diff --git a/tests/models/qwen4_exp/test_qsa_backend_capability.py b/tests/models/qwen4_exp/test_qsa_backend_capability.py +new file mode 100644 +index 0000000..81122b9 +--- /dev/null ++++ b/tests/models/qwen4_exp/test_qsa_backend_capability.py +@@ -0,0 +1,176 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Capability contract of ``Qwen4ExpQSAFlashAttentionBackend``. ++ ++QSA stores e4m3 bytes in a uint8 cache written by ``reshape_and_cache_flash`` ++and dequantizes them inside its own Triton kernel; flash-attn never runs over ++that cache. The backend therefore overrides two capability classmethods so the ++parent's flash-attn fp8 probe cannot veto the configuration. ++ ++``AttentionBackend.validate_configuration`` (vllm/v1/attention/backend.py:262) ++is the entry point that calls both -- ``supports_kv_cache_dtype`` at :289 and ++``supports_combination`` at :335 -- so it is what these tests drive. ++ ++ON WHAT THIS DOES AND DOES NOT ESTABLISH. ``Qwen4ExpQSAAttention.__init__`` ++(nvidia/qsa.py) assigns ``self.attn_backend = Qwen4ExpQSAFlashAttentionBackend`` ++and constructs the impl directly; ``QWEN4_EXP_QSA_TRITON`` is not in the ++selection registry, so a boot does not route this class through ++``platforms/cuda.py``'s ``validate_configuration`` call. These tests therefore ++pin the backend's declared contract, not a boot outcome, and a boot with the ++overrides removed is NOT guaranteed to fail. (``vllm/models/hy_v4/nvidia/ ++attention.py:604`` shows the in-tree pattern for a model that does validate its ++own backend explicitly; QSA does not do that today.) The fp8-cache gate that a ++boot actually hits is ``_validated_qsa_fp8_dtype`` plus the ++``_QSA_SUPPORTED_CACHE_DTYPES`` check in ``Qwen4ExpQSAAttention.__init__``. ++""" ++ ++import pytest ++import torch ++ ++from vllm.config.cache import CacheDType ++from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAFlashAttentionBackend ++from vllm.platforms.interface import DeviceCapability ++from vllm.v1.attention.backend import AttentionType ++from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend ++ ++pytestmark = pytest.mark.cpu_test ++ ++# sm_86: the capability this fork's QSA fp8 path is gated to. ++_SM86 = DeviceCapability(8, 6) ++_FA_PROBE = "vllm.v1.attention.backends.flash_attn.flash_attn_supports_kv_cache_dtype" ++_FP8_REASON = "FP8 KV cache requires" ++_DTYPE_REASON = "kv_cache_dtype not supported" ++ ++ ++class _NoOverrideBackend(Qwen4ExpQSAFlashAttentionBackend): ++ """The same backend with exactly the two overrides taken away. ++ ++ ``__func__`` is unwrapped and re-wrapped so ``cls`` still binds to this ++ subclass: it keeps our ``supported_kv_cache_dtypes`` and differs from the ++ real backend in nothing but the two methods under test. Assigning the bound ++ parent classmethods instead would silently swap the dtype list too, and the ++ control would no longer isolate the overrides. ++ """ ++ ++ supports_kv_cache_dtype = classmethod( ++ FlashAttentionBackend.supports_kv_cache_dtype.__func__ ++ ) ++ supports_combination = classmethod( ++ FlashAttentionBackend.supports_combination.__func__ ++ ) ++ ++ ++def _validate(backend, kv_cache_dtype: CacheDType | None, **overrides): ++ kwargs = dict( ++ head_size=128, ++ dtype=torch.bfloat16, ++ kv_cache_dtype=kv_cache_dtype, ++ block_size=16, ++ use_mla=False, ++ has_sink=False, ++ use_sparse=True, # QSA.is_sparse() is True ++ use_mm_prefix=False, ++ use_per_head_quant_scales=False, ++ device_capability=_SM86, ++ attn_type=AttentionType.DECODER, ++ ) ++ kwargs.update(overrides) ++ return backend.validate_configuration(**kwargs) ++ ++ ++def _combination(backend, kv_cache_dtype: CacheDType | None, **overrides): ++ kwargs = dict( ++ head_size=128, ++ dtype=torch.bfloat16, ++ kv_cache_dtype=kv_cache_dtype, ++ block_size=16, ++ use_mla=False, ++ has_sink=False, ++ use_sparse=True, ++ use_mm_prefix=False, ++ device_capability=_SM86, ++ ) ++ kwargs.update(overrides) ++ return backend.supports_combination(**kwargs) ++ ++ ++@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"]) ++def test_fp8_survives_validation_with_the_fa_probe_rejecting( ++ monkeypatch, kv_cache_dtype ++): ++ """With the flash-attn fp8 probe forced to reject, the overridden backend ++ must still accept the dtype, and the control without the overrides must not. ++ ++ Forcing the probe is what makes this a test that can fail: on a host where ++ flash-attn happens to support fp8 the parent would accept it too, and the ++ override would be indistinguishable from its absence ++ (an instrument must discriminate both worlds). ++ ++ Assertions are on the two specific reasons rather than on an empty list, ++ so unrelated reasons a particular host contributes cannot mask or fake the ++ result. ++ """ ++ monkeypatch.setattr(_FA_PROBE, lambda *a, **k: False) ++ ++ ours = _validate(Qwen4ExpQSAFlashAttentionBackend, kv_cache_dtype) ++ assert _DTYPE_REASON not in ours ++ assert not any(_FP8_REASON in r for r in ours), ours ++ ++ # FAILING CONTROL: remove the overrides, keep everything else. ++ control = _validate(_NoOverrideBackend, kv_cache_dtype) ++ assert _DTYPE_REASON in control, ( ++ "the parent must reject this dtype once the probe rejects it; if it " ++ "does not, this test proves nothing about the override" ++ ) ++ assert any(_FP8_REASON in r for r in control), control ++ ++ # The two classes must differ in exactly these reasons and nothing else. ++ assert set(control) - set(ours) == { ++ _DTYPE_REASON, ++ *(r for r in control if _FP8_REASON in r), ++ } ++ assert not set(ours) - set(control) ++ ++ ++def test_bfloat16_path_is_untouched(monkeypatch): ++ """The overrides must not change the non-quantized cache at all: the parent ++ never consults the fp8 probe for it, so both classes must agree.""" ++ monkeypatch.setattr(_FA_PROBE, lambda *a, **k: False) ++ for kv_cache_dtype in ("auto", "bfloat16"): ++ assert _validate(Qwen4ExpQSAFlashAttentionBackend, kv_cache_dtype) == _validate( ++ _NoOverrideBackend, kv_cache_dtype ++ ) ++ ++ ++def test_unsupported_cache_dtype_is_still_rejected(): ++ """``supports_kv_cache_dtype`` overrides the flash-attn probe, not the ++ class's own dtype list: a dtype QSA does not store must still fail.""" ++ assert not Qwen4ExpQSAFlashAttentionBackend.supports_kv_cache_dtype("fp8_e5m2") ++ assert Qwen4ExpQSAFlashAttentionBackend.supports_kv_cache_dtype(None) ++ assert Qwen4ExpQSAFlashAttentionBackend.supports_kv_cache_dtype("fp8_e4m3") ++ assert _DTYPE_REASON in _validate(Qwen4ExpQSAFlashAttentionBackend, "fp8_e5m2") ++ ++ ++def test_supports_combination_suppresses_only_the_fp8_clause(monkeypatch): ++ """``supports_combination`` must narrow the parent, not silence it. ++ ++ Returning None unconditionally would advertise fewer restrictions than ++ ``forward_qsa`` enforces (it raises on ALiBi, attention sinks and sliding ++ windows), so the parent's other clauses have to survive. ++ """ ++ monkeypatch.setattr(_FA_PROBE, lambda *a, **k: False) ++ ++ # The fp8 clause is suppressed... ++ assert _combination(Qwen4ExpQSAFlashAttentionBackend, "fp8_e4m3") is None ++ parent_said = _combination(_NoOverrideBackend, "fp8_e4m3") ++ assert parent_said is not None and _FP8_REASON in parent_said ++ ++ # ...but the sink/compute-capability clause is not, for fp8 or otherwise. ++ # sm_86 < 9.0, so the parent rejects sinks here and QSA must say so too. ++ for kv_cache_dtype in ("fp8_e4m3", "bfloat16", None): ++ reason = _combination( ++ Qwen4ExpQSAFlashAttentionBackend, kv_cache_dtype, has_sink=True ++ ) ++ assert reason is not None, kv_cache_dtype ++ assert "sink" in reason ++ assert reason == _combination(_NoOverrideBackend, kv_cache_dtype, has_sink=True) +diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py +index cbe58f8..c45f5a7 100644 +--- a/tests/v1/core/test_deferred_block_free.py ++++ b/tests/v1/core/test_deferred_block_free.py +@@ -700,6 +700,7 @@ def test_deferred_victim_stops_the_preemption_cascade(): + assert first.request_id in out2.num_scheduled_tokens + assert first.status == RequestStatus.RUNNING + ++ + def _make_pool_with_cached_blocks(num_blocks: int, block_size: int = 16): + """Build a BlockPool holding `num_blocks` allocated, hashed blocks. + +@@ -727,43 +728,125 @@ def _eviction_order(pool: BlockPool) -> list[int]: + return order + + ++def _seed_request_blocks(manager, request_id: str, blocks_per_group: int): ++ """Give every group's single-type manager a real, hashed block list. ++ ++ Blocks come from the manager's own pool and are hashed so that ++ ``BlockPool.free_blocks`` routes them down the cached branch, which appends ++ to the back of the free queue in the order it is handed -- that order IS ++ the eviction order. Unhashed blocks are *prepended* instead and would mask ++ what is being measured. ++ ++ Returns the per-group block ids in allocation order. ++ """ ++ pool = manager.block_pool ++ allocated: list[list[int]] = [] ++ for group_id, group_mgr in enumerate(manager.coordinator.single_type_managers): ++ blocks = pool.get_new_blocks(blocks_per_group) ++ for i, block in enumerate(blocks): ++ block.set_block_hash( ++ BlockHashWithGroupId(f"g{group_id}b{i}".encode()), ++ num_tokens=pool.hash_block_size, ++ ) ++ group_mgr.req_to_blocks[request_id] = blocks ++ allocated.append([b.block_id for b in blocks]) ++ return allocated ++ ++ ++def _free_through_scheduler(manager, request, defer: bool, ours: set[int]) -> list[int]: ++ """Free ``request`` through the real scheduler helper, one way or the other. ++ ++ ``defer=False`` takes ``_free_request_blocks``'s immediate branch ++ (``KVCacheManager.free`` -> ``KVCacheCoordinator.free`` -> one reversed ++ ``free_blocks`` per group). ``defer=True`` takes the deferred branch ++ (``pop_blocks_for_free`` -> ``deferred_frees`` -> ``_drain_deferred_frees``). ++ Both are the production methods; only the gate differs. ++ """ ++ scheduler = _create_deferring_scheduler() ++ scheduler.kv_cache_manager = manager ++ scheduler.defer_block_free = defer ++ scheduler.processed_step_seq = 0 ++ scheduler.sched_step_seq = 1 ++ request.last_sched_seq = 1 ++ ++ scheduler._free_request_blocks(request) ++ if defer: ++ assert scheduler.deferred_frees, ( ++ "last_sched_seq > processed_step_seq must take the deferred branch; " ++ "if it did not, this arm is not testing the deferred path" ++ ) ++ scheduler.processed_step_seq = 1 ++ scheduler._drain_deferred_frees() ++ assert not scheduler.deferred_frees ++ else: ++ assert not scheduler.deferred_frees ++ ++ # The pool also holds blocks this request never owned; they are unhashed, ++ # so `free_blocks` never touched them and they sit ahead of ours in the ++ # queue. Filter to the request's own blocks, whose relative order is ++ # exactly what was handed to `free_blocks`. ++ return [bid for bid in _eviction_order(manager.block_pool) if bid in ours] ++ ++ + def test_deferred_free_matches_immediate_free_eviction_order(): + """Regression for the deferred-free path collapsing per-group eviction + order for multi-group (hybrid) KV caches. + +- Group A owns blocks [1, 2, 3] (allocation order) and group B owns +- [4, 5, 6]. The immediate-free path frees each group in its own +- reversed ``free_blocks`` call; the fixed deferred-free path must +- reproduce exactly that order, not a single flattened reversal. ++ Drives BOTH arms through the production code -- ``Scheduler. ++ _free_request_blocks`` over a real two-group ``KVCacheManager``, its real ++ ``KVCacheCoordinator``, the real ``deferred_frees`` deque and the real ++ ``_drain_deferred_frees`` -- and compares the resulting eviction order. ++ Constructing the two orders by hand and calling ``BlockPool.free_blocks`` ++ directly would leave every production statement unexercised, so reverting ++ the patch would not fail the test. ++ ++ The immediate path frees each group in its own reversed ``free_blocks`` ++ call (``SingleTypeKVCacheManager.free``). The deferred path pops all groups ++ into ONE concatenated list and drains it in a single ``free_blocks`` call, ++ so the per-group reversal has to be applied before the concatenation -- ++ which is what ``KVCacheCoordinator.pop_blocks_for_free`` now does. ++ ++ Reverting the patch (no per-group reverse in ``pop_blocks_for_free``, a ++ ``reversed()`` back in ``_drain_deferred_frees``) makes the deferred arm ++ return ``reversed(g0_alloc + g1_alloc)`` -- whole groups in the wrong order ++ relative to the immediate arm -- which fails both the equality assertion ++ and the explicit pre-patch-shape assertion below. + """ +- group_a_alloc = [1, 2, 3] +- group_b_alloc = [4, 5, 6] +- +- # Immediate path: KVCacheCoordinator.free() -> one reversed free_blocks +- # call per single-type manager (group). +- pool_imm, by_id = _make_pool_with_cached_blocks(6) +- for group in (group_a_alloc, group_b_alloc): +- pool_imm.free_blocks([by_id[i] for i in reversed(group)]) +- immediate_order = _eviction_order(pool_imm) +- +- # Fixed deferred path: groups are kept separate (each already reversed) +- # and freed in their own call on drain -- mirroring the immediate path. +- pool_fixed, by_id = _make_pool_with_cached_blocks(6) +- ordered_groups = [list(reversed(group_a_alloc)), list(reversed(group_b_alloc))] +- for group in ordered_groups: +- pool_fixed.free_blocks([by_id[i] for i in group]) +- fixed_order = _eviction_order(pool_fixed) +- +- assert fixed_order == immediate_order == [3, 2, 1, 6, 5, 4] +- +- # The old (buggy) deferred path flattened both groups then reversed once, +- # which produces a different order. Guard against regressing back to it. +- pool_buggy, by_id = _make_pool_with_cached_blocks(6) +- flattened = group_a_alloc + group_b_alloc +- pool_buggy.free_blocks([by_id[i] for i in reversed(flattened)]) +- buggy_order = _eviction_order(pool_buggy) +- assert buggy_order == [6, 5, 4, 3, 2, 1] +- assert buggy_order != immediate_order ++ blocks_per_group = 3 ++ request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0] ++ ++ imm_manager = _make_two_group_kv_cache_manager(num_blocks=20) ++ def_manager = _make_two_group_kv_cache_manager(num_blocks=20) ++ assert len(imm_manager.coordinator.single_type_managers) == 2, ( ++ "a single-group config cannot distinguish the two orders at all" ++ ) ++ ++ allocated = _seed_request_blocks(imm_manager, request.request_id, blocks_per_group) ++ assert allocated == _seed_request_blocks( ++ def_manager, request.request_id, blocks_per_group ++ ), "the two arms must start from identical per-group allocations" ++ assert not set(allocated[0]) & set(allocated[1]), ( ++ "groups must own disjoint block ids or a group mix-up is invisible" ++ ) ++ ours = {bid for group in allocated for bid in group} ++ ++ immediate_order = _free_through_scheduler(imm_manager, request, False, ours) ++ deferred_order = _free_through_scheduler(def_manager, request, True, ours) ++ ++ # Pin the shape as well as the equality: tail-first within a group, groups ++ # in their coordinator order. ++ expected = [bid for group in allocated for bid in reversed(group)] ++ assert immediate_order == expected ++ assert deferred_order == immediate_order ++ ++ # The pre-patch deferred implementation flattened every group then reversed ++ # once. Assert the observed order is not that, so a revert cannot pass. ++ pre_patch_order = list(reversed([bid for group in allocated for bid in group])) ++ assert pre_patch_order != expected, ( ++ "the fixture does not separate the two implementations; with this " ++ "shape the test could not fail on a revert" ++ ) ++ assert deferred_order != pre_patch_order + + + def test_pop_blocks_for_free_returns_eviction_order(): +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index f4fa56b..adb057d 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -333,11 +333,36 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: +- # Same reason: the parent's "FP8 KV cache requires FA3 on SM90 or FA4 +- # on SM100" rejection is about flash-attn's own kernels, which QSA does +- # not use. The real device gate for our fp8 path is +- # `_validated_qsa_fp8_dtype`, which raises on any capability != 86. +- return None ++ # Suppress ONLY the parent's fp8 clause ("FP8 KV cache requires FA3 on ++ # SM90 or FA4 on SM100"), which is about flash-attn's own kernels and ++ # so does not apply to a cache QSA reads in Triton. The real device ++ # gate for our fp8 path is `_validated_qsa_fp8_dtype`, which raises on ++ # any capability != 86. ++ # ++ # Everything else the parent rejects still holds and is returned ++ # unchanged: the sink/compute-capability clause and the ++ # mm-prefix/FA-version clause. Returning None for every combination ++ # would advertise fewer restrictions than `forward_qsa` enforces (it ++ # raises on ALiBi, sinks and sliding windows). ++ # ++ # Same shape as the impl constructor: hide the fp8 dtype from the ++ # probe, keep the rest of the parent's behaviour. ++ probe_kv_cache_dtype = kv_cache_dtype ++ if probe_kv_cache_dtype is not None and _is_qsa_fp8_cache_dtype( ++ probe_kv_cache_dtype ++ ): ++ probe_kv_cache_dtype = None ++ return super().supports_combination( ++ head_size, ++ dtype, ++ probe_kv_cache_dtype, ++ block_size, ++ use_mla, ++ has_sink, ++ use_sparse, ++ use_mm_prefix, ++ device_capability, ++ ) + + @staticmethod + def get_name() -> str: +-- +2.43.0 + diff --git a/patches/0027-worker-debug-gated-diagnostics-for-Mamba-slot-owners.patch b/patches/0027-worker-debug-gated-diagnostics-for-Mamba-slot-owners.patch new file mode 100644 index 0000000..e24132d --- /dev/null +++ b/patches/0027-worker-debug-gated-diagnostics-for-Mamba-slot-owners.patch @@ -0,0 +1,1594 @@ +From 2ae218c5e7b045e42c55fb26e929d8698728e11c Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:18:16 -0400 +Subject: [PATCH 027/103] worker: debug-gated diagnostics for Mamba slot + ownership and guard encoding + +- Host-side Mamba slot-ownership check (off|warn|raise, default off) with a + failing control that proves it detects a wrong block-table row. +- Guard telemetry: int64 encoding limits and the extreme-value pair closed by + test; a decoded-None counter field is now named in the warning; a + captured-replay sentinel probe. +- Both knobs registered in envs.py. +--- + .../test_host_embed_capture_sentinel.py | 202 +++++++++ + .../worker/gpu/test_guard_telemetry_limits.py | 328 ++++++++++++++ + .../worker/gpu/test_mamba_slot_ownership.py | 366 +++++++++++++++ + vllm/envs.py | 48 ++ + vllm/v1/worker/guard_telemetry.py | 76 +++- + vllm/v1/worker/mamba_utils.py | 416 +++++++++++++++++- + 6 files changed, 1433 insertions(+), 3 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_host_embed_capture_sentinel.py + create mode 100644 tests/v1/worker/gpu/test_guard_telemetry_limits.py + create mode 100644 tests/v1/worker/gpu/test_mamba_slot_ownership.py + +diff --git a/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py b/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py +new file mode 100644 +index 0000000..aac4684 +--- /dev/null ++++ b/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py +@@ -0,0 +1,202 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The captured-replay host sentinel for the host-embedding layer (reference-host tier). ++ ++The host-embedding work's bytewise check passed: on the real Merlin-e1 BF16 ++table, TP2 over real NCCL, all 36 cases equal and the negative control (one ++flipped host byte) detected on both ranks. Its own title records what it did ++NOT settle: ++ ++ "compiled-forward/captured-replay sentinel not yet run" ++ "that the SERVING path uses the same storage after compile and CUDA-graph ++ capture: the CPU-sentinel check through compiled forward and captured ++ replay, with residency recorded at each lifecycle stage" ++ ++That is not a restatement of the layer-level gate. The layer-level gate proves ++the UVA read is byte-exact when Python calls it. This proves the READ IS STILL ++THE UVA READ after capture -- that graph capture did not bake a device-resident ++copy of the table into the replay. Those are different claims, and only the ++second one is about the serving path. ++ ++WHY THIS IS THE RIGHT SHAPE OF CHECK. Equality alone cannot distinguish them: ++a captured replay that silently reads a stale device copy produces IDENTICAL ++bytes for identical inputs, forever. So equality is the first arm and a HOST ++SENTINEL is the second: mutate one row in the pinned host table between capture ++and replay and require the replay to produce the EXACT expected post-mutation ++bytes. If it does not, the replay is not reading live host memory, and every ++byte-equality result taken under capture was evidence about a copy ++(an instrument must discriminate both worlds). ++ ++WHAT THE SENTINEL DOES *NOT* PROVE. Observing the ++mutation proves the replay DEPENDS ON MUTABLE HOST DATA. It does not prove ++exclusive UVA residency: a replay that copies the host row into a device buffer ++before reading it, on every launch, would also observe the mutation and pass. ++Ruling that out needs storage-identity and residency evidence ++(`host_memory_stats()` per worker, device allocation deltas across the ++lifecycle), which is not collected and is not asserted here. Do not read a ++pass as "therefore UVA". ++ ++STATUS: NEVER EXECUTED. This module was written offline with no CUDA ++compiler; it is compile-checked only. It is a reference-host test. ++ ++WHAT IT COVERS AND WHAT IT DOES NOT. It covers the standalone layer under a ++real ``torch.cuda.CUDAGraph`` capture, which is the part that can be asserted ++without booting an engine. It does NOT cover the layer inside the model's ++compiled piecewise forward on Flash-Next; that needs a real boot with ++``VLLM_HOST_EMBED_TABLE`` on, reading the sentinel through a served completion, ++and is the part of the host-embedding check not covered here. Do not treat it as closed on this file ++alone. ++""" ++ ++import os ++ ++import pytest ++import torch ++ ++import vllm.model_executor.layers.vocab_parallel_embedding as vpe ++from vllm.model_executor.layers.vocab_parallel_embedding import ( ++ PinnedHostVocabParallelEmbedding, ++) ++ ++VOCAB = 4096 ++HIDDEN = 256 ++DTYPE = torch.bfloat16 ++SENTINEL_ID = 1234 ++ ++pytestmark = [ ++ pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA"), ++ pytest.mark.skipif( ++ os.environ.get("VLLM_E2_GPU_TESTS") != "1", ++ reason="set VLLM_E2_GPU_TESTS=1 to run GPU tests (shared GPU host)", ++ ), ++] ++ ++ ++@pytest.fixture ++def layer(monkeypatch: pytest.MonkeyPatch) -> PinnedHostVocabParallelEmbedding: ++ """A real host-table embedding at TP=1, so no collective is involved. ++ ++ TP=1 is deliberate: the host-embedding work already settled the TP2 collective bytewise. The ++ open question is capture, and a second rank only adds a way for this test ++ to fail for an unrelated reason. ++ """ ++ if not vpe.PIN_MEMORY: ++ pytest.skip("pinned memory is unavailable") ++ monkeypatch.setattr(vpe, "get_tensor_model_parallel_rank", lambda: 0) ++ monkeypatch.setattr(vpe, "get_tensor_model_parallel_world_size", lambda: 1) ++ monkeypatch.setattr(vpe, "tensor_model_parallel_all_reduce", lambda value: value) ++ ++ embedding = PinnedHostVocabParallelEmbedding( ++ VOCAB, HIDDEN, params_dtype=DTYPE ++ ) ++ assert embedding.weight.is_pinned(), "the table must be pinned to be mapped" ++ assert embedding.accelerator_weight_view is not None, "no UVA view" ++ values = torch.arange(VOCAB * HIDDEN, dtype=torch.int32) % 251 ++ embedding.weight.data.copy_(values.reshape(VOCAB, HIDDEN).to(DTYPE)) ++ return embedding ++ ++ ++def _ids() -> torch.Tensor: ++ return torch.tensor( ++ [0, 1, SENTINEL_ID, VOCAB - 1, SENTINEL_ID], ++ dtype=torch.int32, ++ device="cuda", ++ ) ++ ++ ++def _bytes(tensor: torch.Tensor) -> bytes: ++ return tensor.contiguous().view(torch.uint8).cpu().numpy().tobytes() ++ ++ ++def test_captured_replay_is_byte_identical_to_the_eager_forward(layer) -> None: ++ """Arm 1: capture must not change a single output byte.""" ++ ids = _ids() ++ eager = _bytes(layer(ids)) ++ ++ static_ids = ids.clone() ++ graph = torch.cuda.CUDAGraph() ++ # A warm-up on a side stream: capturing a first-ever launch also captures ++ # its lazy initialisation (a lazy load that blocks behind a stream wait). ++ side = torch.cuda.Stream() ++ side.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(side): ++ static_out = layer(static_ids) ++ torch.cuda.current_stream().wait_stream(side) ++ with torch.cuda.graph(graph): ++ static_out = layer(static_ids) ++ ++ graph.replay() ++ torch.cuda.synchronize() ++ ++ assert _bytes(static_out) == eager, ( ++ "the captured replay's output differs from the eager forward on the " ++ "same ids and the same table" ++ ) ++ ++ ++def test_a_host_sentinel_written_after_capture_is_visible_to_the_replay( ++ layer, ++) -> None: ++ """Arm 2, THE FAILING CONTROL: does the replay still read host memory? ++ ++ Arm 1 passes whether the replay reads the UVA view or a device copy baked ++ in at capture time, because both produce the same bytes for the same ids. ++ Writing the host table AFTER capture separates them: only a real UVA read ++ sees the new row. ++ ++ A failure here does not mean the output is wrong today -- it means every ++ byte-equality result taken under capture is evidence about a copy, and the ++ host-table saving (the reason the layer exists) is not what it claims on ++ the serving path. ++ """ ++ ids = _ids() ++ static_ids = ids.clone() ++ ++ side = torch.cuda.Stream() ++ side.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(side): ++ static_out = layer(static_ids) ++ torch.cuda.current_stream().wait_stream(side) ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ static_out = layer(static_ids) ++ ++ graph.replay() ++ torch.cuda.synchronize() ++ before = _bytes(static_out) ++ ++ # The sentinel: one row of the PINNED HOST table, written from the CPU ++ # with no device copy anywhere. ++ row = layer.weight.data[SENTINEL_ID] ++ original = row.clone() ++ row.fill_(torch.finfo(DTYPE).max) ++ ++ # `after != before` is too weak -- ANY change passes it, ++ # including a wrong one. Build the exact bytes the replay must now produce ++ # and require equality with them. A partially updated output, a torn row ++ # or a different row moving all satisfy inequality and fail this. ++ expected_rows = [] ++ for token in ids.tolist(): ++ source = row if token == SENTINEL_ID else layer.weight.data[token] ++ expected_rows.append(source.clone()) ++ expected = _bytes(torch.stack(expected_rows).to("cuda")) ++ ++ graph.replay() ++ torch.cuda.synchronize() ++ after = _bytes(static_out) ++ ++ assert after != before, ( ++ "the captured replay did not observe a row written directly into the " ++ "pinned host table: it is reading something other than the live host " ++ "storage" ++ ) ++ assert after == expected, ( ++ "the captured replay changed, but not to the expected post-mutation " ++ "bytes -- it is not a clean read of the mutated host row" ++ ) ++ ++ # Restore and prove the observation is two-way, not a one-off artefact. ++ row.copy_(original) ++ graph.replay() ++ torch.cuda.synchronize() ++ assert _bytes(static_out) == before +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_limits.py b/tests/v1/worker/gpu/test_guard_telemetry_limits.py +new file mode 100644 +index 0000000..9fa7536 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_guard_telemetry_limits.py +@@ -0,0 +1,328 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Guard telemetry: the documented limits, pinned by test rather than by change. ++ ++Three limits of the bounds-guard telemetry are documented rather than fixed in code. ++This module pins the ones that are testable offline. The coverage table and the ++reader/producer stream-ordering argument that go with them are design notes not ++shipped with this tree. ++The three limits: ++ ++ * Independent maxima, remaining half -- the (row, value) fields are INDEPENDENT ++ lifetime maxima. The labelling half shipped earlier (the WARNING no ++ longer calls them an "example"); what was missing is a worked example ++ proving the reported pair can be a pair that never occurred. That is ++ ``test_the_reported_pair_can_be_one_that_never_occurred``. ++ * Encoding limits -- the int64 bias encoding loses or falsifies extreme values. ++ Closed by test, not by change: the encoding is correct for every value a ++ correct kernel can produce, and the tests below pin exactly where it stops ++ being able to carry one, plus the fact that the COUNT always survives. ++ * Zero hits -- "zero hits implies no output bit changed" is false in ++ general. See ``test_zero_hits_is_not_bitwise_equivalence`` below; it is a ++ documentation test, and it says so. ++ ++No GPU: nothing here allocates a counter buffer or launches a kernel. The ++producer arithmetic is exercised through ``encode_slot``, which is the same ++expression the Triton bodies inline, and the consumer through the real ++``_report_rows``. ++""" ++ ++import logging ++ ++import pytest ++import torch ++ ++from vllm.v1.worker import guard_telemetry as gt ++ ++_CPU = torch.device("cpu") ++ ++ ++@pytest.fixture(autouse=True) ++def _clean_high_water_marks(): ++ """``_report_rows`` only reports a count ABOVE what it last saw. ++ ++ A test that leaves ``_seen`` behind makes the NEXT test silently measure ++ nothing -- the trap the module's own comment on ``_seen`` describes. ++ """ ++ gt.reset_for_test() ++ yield ++ gt.reset_for_test() ++ ++ ++class _Collector(logging.Handler): ++ def __init__(self) -> None: ++ super().__init__(level=logging.WARNING) ++ self.messages: list[str] = [] ++ ++ def emit(self, record: logging.LogRecord) -> None: ++ self.messages.append(record.getMessage()) ++ ++ ++@pytest.fixture ++def warnings_from_the_reader(): ++ """Capture the reader's WARNINGs off the real logger object. ++ ++ Not ``caplog``: vLLM's root "vllm" logger is configured with ++ ``propagate: False``, so records never reach the handler pytest installs ++ on the ROOT logger and the capture would be empty. Attaching to the ++ module's own logger is the only capture that cannot quietly see nothing ++ (a capture that can see nothing measures nothing). ++ """ ++ target = getattr(gt.logger, "logger", gt.logger) ++ handler = _Collector() ++ previous_level = target.level ++ target.addHandler(handler) ++ target.setLevel(logging.WARNING) ++ try: ++ yield handler.messages ++ finally: ++ target.removeHandler(handler) ++ target.setLevel(previous_level) ++ ++ ++# --------------------------------------------------------------------------- ++# Encoding limits -- int64 encoding of extreme values ++# --------------------------------------------------------------------------- ++ ++ ++@pytest.mark.parametrize( ++ "value", ++ [0, 1, -1, 42, -42, (1 << 31) - 1, -((1 << 31) - 1), 1 << 40], ++) ++def test_ordinary_values_round_trip(value: int) -> None: ++ """Every value a correct kernel can produce survives the encoding. ++ ++ A row or index is an int32-range quantity, so this covers the whole domain ++ the guards actually record. The cases below are about metadata that is ++ already corrupt -- which is precisely when a guard fires. ++ """ ++ assert gt.decode_slot(gt.encode_slot(value)) == value ++ ++ ++def test_int32_min_is_falsified_as_nothing_recorded() -> None: ++ """The one value the encoding cannot distinguish from an unwritten slot. ++ ++ ``-GUARD_BIAS`` encodes to exactly 0, and 0 is the zeroed-buffer sentinel ++ for "this slot was never written". So a guard that masks an access at row ++ INT32_MIN reports its ROW as absent. ++ """ ++ assert gt.encode_slot(-gt.GUARD_BIAS) == 0 ++ assert gt.decode_slot(0) is None ++ ++ ++def test_values_below_int32_min_are_lost_by_atomic_max() -> None: ++ """Below ``-GUARD_BIAS`` the encoded word is negative. ++ ++ The slot is only ever written with ``tl.atomic_max`` into a buffer zeroed ++ at allocation, and 0 > any negative, so the store is a no-op and the slot ++ still reads 0. Modelled here as ``max(0, encoded)`` because that is ++ exactly what the atomic does; there is no Triton in this test. ++ """ ++ for value in (-gt.GUARD_BIAS - 1, -(1 << 40), -(1 << 62)): ++ encoded = gt.encode_slot(value) ++ assert encoded < 0, value ++ assert gt.decode_slot(max(0, encoded)) is None, value ++ ++ ++def test_the_encodable_window_is_where_the_constants_say_it_is() -> None: ++ """Pin the published bounds so a later change to GUARD_BIAS moves them.""" ++ assert gt.decode_slot(gt.encode_slot(gt.GUARD_MIN_ENCODABLE)) == ( ++ gt.GUARD_MIN_ENCODABLE ++ ) ++ assert gt.decode_slot(gt.encode_slot(gt.GUARD_MAX_ENCODABLE)) == ( ++ gt.GUARD_MAX_ENCODABLE ++ ) ++ # One step below the window is the falsified case, not a round trip. ++ assert gt.encode_slot(gt.GUARD_MIN_ENCODABLE - 1) == 0 ++ # One step above it does not fit in the int64 slot at all. ++ assert gt.encode_slot(gt.GUARD_MAX_ENCODABLE + 1) > (1 << 63) - 1 ++ ++ ++def test_the_count_survives_an_unencodable_value_and_the_report_says_so( ++ warnings_from_the_reader: list[str], ++) -> None: ++ """The operational half of the encoding limits. ++ ++ ``n_masked`` carries no bias and is written with ``atomic_add``, so it ++ advances whatever the offending value was. A report of "3 masked accesses, ++ row None, index None" is therefore a report of three REAL masked faults -- ++ not a report of nothing. Before this change the message said only "row-field ++ None", which reads as "nothing recorded" to anyone who has not read the ++ encoding. ++ ++ What the message must NOT do is diagnose WHY the field is None. The reader ++ has no ordering dependency on the producer, so an unobserved maximum and an ++ unencodable one read identically; the first draft of this clause claimed ++ "unencodable" outright and was wrong to. ++ """ ++ rows = [[0, 0, 0] for _ in range(gt.NUM_SITES)] ++ rows[gt.SITE_MAMBA_COPY] = [3, 0, 0] # count 3, both maxima unencodable ++ ++ gt._report_rows(_CPU, rows) ++ ++ assert len(warnings_from_the_reader) == 1, warnings_from_the_reader ++ text = warnings_from_the_reader[0] ++ assert "masked 3 out-of-range access(es)" in text ++ assert "row-field None" in text ++ assert "UNOBSERVED OR UNENCODABLE, not absent" in text ++ # It must not claim to know which of the two it is. ++ assert "no ordering dependency on the producer" in text ++ ++ ++# --------------------------------------------------------------------------- ++# Independent maxima -- the fields are independent; the worked example ++# --------------------------------------------------------------------------- ++ ++ ++def test_the_reported_pair_can_be_one_that_never_occurred( ++ warnings_from_the_reader: list[str], ++) -> None: ++ """THE EXAMPLE PAIR. ++ ++ Two masked accesses at one site, in this order: ++ ++ hit A: row 3, offending index 100 ++ hit B: row 7, offending index 5 ++ ++ Each slot independently keeps its lifetime maximum, so the slots end at ++ row 7 and index 100 and the reader reports the pair **(7, 100)** -- a ++ combination that never happened on the device. Row 7's access was at ++ index 5; index 100's access was at row 3. ++ ++ This is why the WARNING labels them independent lifetime maxima and why no ++ attribution may be built on the pair: to chase a single access you need ++ the count plus a repro, not this triple. Anyone wanting a coherent ++ (row, value) sample needs the other half -- a capture that writes ++ both slots under one atomic decision -- which this tree does NOT implement. ++ """ ++ slots = [0, 0, 0] ++ for row, value in ((3, 100), (7, 5)): ++ slots[0] += 1 ++ slots[1] = max(slots[1], gt.encode_slot(row)) ++ slots[2] = max(slots[2], gt.encode_slot(value)) ++ ++ assert gt.decode_slot(slots[1]) == 7 ++ assert gt.decode_slot(slots[2]) == 100 ++ assert (7, 100) not in {(3, 100), (7, 5)} ++ ++ rows = [[0, 0, 0] for _ in range(gt.NUM_SITES)] ++ rows[gt.SITE_GDN_DECODE_STATE] = slots ++ ++ gt._report_rows(_CPU, rows) ++ ++ assert len(warnings_from_the_reader) == 1, warnings_from_the_reader ++ text = warnings_from_the_reader[0] ++ assert "row-field 7, index-field 100" in text ++ assert "independent lifetime maxima, not one observed pair" in text ++ # ...and no suffix about unobservable fields: both decoded fine. ++ assert "UNOBSERVED OR UNENCODABLE" not in text ++ ++ ++# --------------------------------------------------------------------------- ++# Zero hits is not bitwise equivalence ++# --------------------------------------------------------------------------- ++ ++ ++def test_zero_hits_is_not_bitwise_equivalence() -> None: ++ """A DOCUMENTATION TEST. It asserts the guard set's shape, not a behaviour. ++ ++ The tempting claim is "zero guard hits implies no output bit ++ changed". It holds only for live, valid inputs, and there are two ways out ++ of it that a counter at zero cannot exclude: ++ ++ 1. **The guards are bounds checks.** Every site in ``SITE_NAMES`` masks an ++ OUT-OF-RANGE access and counts it. A wrong-but-in-range row is not an ++ out-of-range access, so it is invisible here by construction -- stated ++ for the mamba copy site in the guard design notes ("Blind to: a column or row that ++ is in range but wrong -- the common case for a recycled slot, ++ producing a silent wrong state copy"). That is the whole of the ++ vllm#55506 / llama.cpp#29092 class, which is why it needed its own ++ instrument (``test_mamba_slot_ownership``) rather than a counter. ++ ++ 2. **``NULL_BLOCK_ID`` rows change bytes uncounted.** Block id 0 is the ++ pool's placeholder (``core/block_pool.py:188-191``, ++ ``v1/attention/backends/utils.py:46``), which no request owns. The ++ copy guard's ``<= 0`` branch is DELIBERATELY not counted -- the guard design notes: ++ "Also blind to the ``<= 0`` path, deliberately: NULL_BLOCK_ID is 0 and ++ is the ordinary unallocated marker". Where such a row is zero-filled ++ rather than skipped, inactive bytes are written and no counter moves. ++ Those bytes belong to no live request, so no OUTPUT bit changes -- but ++ "zero hits" is evidence about live rows only, and a claim of bitwise ++ equivalence has to be made about the output tensor, not inferred from ++ the counters. ++ ++ So a v2.2 result may say "no guard fired", never "no output bit changed" ++ (measure the property, not the proxy). The check below is that the ++ guard set is still a bounds-check set of the size this argument was ++ written against; if a site is added, this argument is re-read. ++ """ ++ assert gt.NUM_SITES == 11 ++ assert set(gt.SITE_NAMES) == set(range(gt.NUM_SITES)) ++ # The one non-bounds site, which is telemetry and not a masked fault. ++ assert "SKIPPED" in gt.SITE_NAMES[gt.SITE_QSA_RING_SKIP] ++ ++ ++# --------------------------------------------------------------------------- ++# The diagnostic gates are registered knobs ++# --------------------------------------------------------------------------- ++ ++ ++@pytest.mark.parametrize( ++ ("raw", "expected"), ++ [ ++ (None, "off"), ++ ("", "off"), # `VAR=` is how scripts and fixtures clear a variable ++ (" ", "off"), ++ ("off", "off"), ++ ("1", "warn"), # continuity with the earlier boolean gates ++ ("warn", "warn"), ++ ("RAISE", "raise"), ++ (" Warn ", "warn"), ++ ], ++) ++def test_e2_diagnostic_gates_normalize_identically( ++ monkeypatch: pytest.MonkeyPatch, raw: str | None, expected: str ++) -> None: ++ """One normalizer for the whole family, so the registry cannot lie. ++ ++ Both PLE diagnostic gates are registered in ``vllm.envs`` and read from there by the ++ modules that act on them. If a module re-parsed the variable itself, the ++ value ``vllm.envs`` reports (and hashes, and logs) could differ from the ++ value the check applies. ++ """ ++ import vllm.envs as envs ++ ++ for name in ("VLLM_E2_PLE_FLAG_CHECK", "VLLM_E2_MAMBA_SLOT_OWNERSHIP"): ++ monkeypatch.delenv(name, raising=False) ++ if raw is not None: ++ monkeypatch.setenv(name, raw) ++ assert envs.environment_variables[name]() == expected, (name, raw) ++ ++ ++def test_a_set_but_unrecognised_gate_is_a_boot_error( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """Fail closed: a broken gate must not silently disarm its instrument.""" ++ import vllm.envs as envs ++ ++ monkeypatch.setenv("VLLM_E2_MAMBA_SLOT_OWNERSHIP", "yes") ++ with pytest.raises(ValueError, match="off|warn|raise"): ++ envs.environment_variables["VLLM_E2_MAMBA_SLOT_OWNERSHIP"]() ++ ++ ++def test_the_host_side_gates_are_not_compile_factors() -> None: ++ """They gate host Python outside any capture and add no graph node. ++ ++ Hashing them would invalidate every compile cache the moment a diagnostic ++ boot is taken -- which is precisely when the cache is wanted. If either ++ ever grows a device-side component this test should fail and be revisited, ++ not deleted. ++ """ ++ import vllm.envs as envs ++ ++ monkeypatch_free = envs.compile_factors() ++ assert "VLLM_E2_PLE_FLAG_CHECK" not in monkeypatch_free ++ assert "VLLM_E2_MAMBA_SLOT_OWNERSHIP" not in monkeypatch_free ++ # ...but they ARE registered, so `vllm.envs` can report them. ++ assert "VLLM_E2_PLE_FLAG_CHECK" in envs.environment_variables ++ assert "VLLM_E2_MAMBA_SLOT_OWNERSHIP" in envs.environment_variables +diff --git a/tests/v1/worker/gpu/test_mamba_slot_ownership.py b/tests/v1/worker/gpu/test_mamba_slot_ownership.py +new file mode 100644 +index 0000000..89131ef +--- /dev/null ++++ b/tests/v1/worker/gpu/test_mamba_slot_ownership.py +@@ -0,0 +1,366 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Ownership of the mamba state slot a request resumes into. ++ ++THE HAZARD (llama.cpp#29092 class; the vllm#55506 failure mode is one ++instance). Mamba state blocks are the one KV class in this tree with neither ++zeroing nor an ownership instrument: ++ ++* ``single_type_kv_cache_manager.py:96-99`` gates the recording of newly ++ allocated blocks for zeroing on ``isinstance(kv_cache_spec, AttentionSpec)``. ++ ``MambaSpec`` is not an ``AttentionSpec``, so mamba state blocks are never ++ zeroed by ``KVBlockZeroer`` (``gpu/model_runner.py:1213``). ++* the host-side KV-zeroer ownership check (``sched/scheduler.py``) only inspects ++ ``new_block_ids_to_zero`` -- a list mamba blocks are never on. ++* a genuinely fresh request is safe anyway: ``gdn_attn.py:401`` sets ++ ``has_initial_state = context_lens_tensor > 0``, so at ``context_len == 0`` ++ the GDN kernel ignores whatever the reused block holds. ++* a RESUMING request is not. With ``has_initial_state`` true the kernel ++ consumes the block's contents, and the only thing that puts the right ++ contents there is the align copy. The bounds guards bound the ++ indices; a bounds check cannot catch a wrong-but-in-range row, which is the ++ whole of this class. ++ ++So the safety of a resuming request rested entirely on the correctness of the ++align copy's row selection -- which is the row-attribution fix, itself the fix for ++having got that row wrong once. These tests give that argument an instrument. ++ ++GPU-FREE, INCLUDING THE FAILING CONTROL. The check is host side by ++construction, and ``test_pp_deferred_mamba_block_table`` already drives the ++real production plumbing (``MambaHybridModelState._ensure_align_ctx`` -> ++``MambaSpecDecodeGPUContext.initialize_from_forward_context`` -> ++``run_fused_postprocess_align``) on CPU tensors with the Triton launch ++captured instead of run. The control reuses that scenario builder rather than a ++stand-in (a hook contract needs the real classes), flips the production ++row rule ``ALIGN_BT_ROW_IS_REQ`` off, and requires the check to fire. No GPU ++time is needed; what a reference-host run would add is the end-to-end consequence ++(acceptance-length collapse, constant-token loops), not the assertion. ++ ++SCOPE LIMIT, stated because it is a real one: ``run_fused_postprocess_align`` ++is engine-thread Python, so a CAPTURED replay does not re-enter it and the ++check does not run for one. It covers the launch path, which is where the row ++is selected. ++""" ++ ++import pytest ++ ++from tests.v1.worker.gpu.test_pp_deferred_mamba_block_table import ( ++ _build_scenario, ++ _run_deferred, ++) ++from vllm.v1.worker import mamba_utils as mamba_utils_mod ++from vllm.v1.worker.mamba_utils import mamba_slot_ownership_violations ++ ++# vllm#55506's own reproduction permutation. A NON-IDENTITY batch-row -> ++# req-slot mapping is load-bearing for the control: with an identity mapping ++# ``batch_idx == req_idx`` and the two row rules are indistinguishable ++# (the needle shape is a variable). Four entries is also >= 3 concurrent ++# requests, the concurrency at which tomasruizt reproduced #55506 under ++# PP2 + MTP + prefix caching. ++_PERM = (2, 0, 3, 1) ++ ++ ++def _arm(monkeypatch: pytest.MonkeyPatch, mode: str = "raise") -> None: ++ monkeypatch.setattr(mamba_utils_mod, "_SLOT_OWNERSHIP_MODE", mode) ++ ++ ++# --------------------------------------------------------------------------- ++# End-to-end through the real align plumbing ++# --------------------------------------------------------------------------- ++ ++ ++def test_production_row_rule_reports_no_ownership_violation( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """The green world: with ALIGN_BT_ROW_IS_REQ the check must stay silent. ++ ++ A check that fires on correct metadata is worse than no check (a ++ guard rule), so this arm is not decoration: it is what makes the control ++ below evidence about the row rule rather than about the fixture. ++ """ ++ _arm(monkeypatch) ++ scenario = _build_scenario(slot_mapping=_PERM) ++ _run_deferred(scenario, monkeypatch) ++ assert scenario.ctx.last_slot_ownership_violations == [] ++ ++ ++def test_bt_row_is_req_false_makes_the_check_fire( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """THE FAILING CONTROL: flip the row rule off, the check must fire. ++ ++ This is the vllm#55506 world reproduced offline. The deferred align ++ post-process carries an ``idx_mapping`` captured ``pp_size`` steps earlier, ++ so with ``BT_ROW_IS_REQ`` off the kernel resolves batch row 0 -- which the ++ mapping says is request slot 2 -- to block-table row 0, another request's ++ slot. tomasruizt measured the consequence as acceptance length collapsing ++ to 1.00 for stretches; on this fork it showed as all-NaN logits and ++ constant-token loops. ++ ++ If this test passes without the assertion firing, the assertion measures ++ nothing (an instrument must discriminate both worlds). ++ """ ++ _arm(monkeypatch) ++ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False) ++ scenario = _build_scenario(slot_mapping=_PERM) ++ ++ with pytest.raises(RuntimeError, match="state-slot ownership violation"): ++ _run_deferred(scenario, monkeypatch) ++ ++ findings = scenario.ctx.last_slot_ownership_violations ++ assert any(f.startswith("row attribution:") for f in findings), findings ++ # ...and it is not only an index complaint: the block ids differ, so the ++ # copy really would land on another request's recurrent state. ++ assert any(f.startswith("destination disagreement:") for f in findings), findings ++ ++ ++def test_warn_mode_records_without_stopping_the_step( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """warn is the mode for a long stress boot; raise is for a repro.""" ++ _arm(monkeypatch, "warn") ++ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False) ++ scenario = _build_scenario(slot_mapping=_PERM) ++ ++ _run_deferred(scenario, monkeypatch) ++ ++ assert scenario.ctx.last_slot_ownership_violations ++ ++ ++def test_off_is_off(monkeypatch: pytest.MonkeyPatch) -> None: ++ """Default: nothing is copied to the host and nothing is recorded.""" ++ _arm(monkeypatch, "off") ++ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False) ++ scenario = _build_scenario(slot_mapping=_PERM) ++ ++ _run_deferred(scenario, monkeypatch) ++ ++ assert scenario.ctx.last_slot_ownership_violations == [] ++ ++ ++def test_an_identity_mapping_hides_the_defect_from_this_check( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ """Why the control uses a permutation, asserted rather than asserted-about. ++ ++ With an identity mapping ``batch_idx == req_idx``, so the wrong row rule ++ selects the right row anyway and the check is silent although the rule is ++ wrong. That is not a defect in the check -- there is nothing wrong with ++ the step -- but it IS a constraint on any future reproduction attempt: ++ a single-stream or in-order batch cannot exhibit #55506, which is why ++ upstream needed >= 3 concurrent requests to see it at all. ++ """ ++ _arm(monkeypatch) ++ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False) ++ scenario = _build_scenario(slot_mapping=(0, 1, 2)) ++ ++ _run_deferred(scenario, monkeypatch) ++ ++ assert scenario.ctx.last_slot_ownership_violations == [] ++ ++ ++# --------------------------------------------------------------------------- ++# The three properties, directly on the pure checker ++# --------------------------------------------------------------------------- ++ ++_BLOCK_SIZE = 16 ++ ++ ++def _clean_kwargs(**overrides): ++ """One copying request (slot 0) and one bystander, tables disjoint.""" ++ kwargs = dict( ++ num_reqs=2, ++ idx_mapping=[0, 1], ++ bt_row_is_req=True, ++ block_tables=[[[10, 11, 12], [20, 21, 22], [30, 31, 32]]], ++ num_rows=3, ++ num_accepted_tokens=[2, 2, 2], ++ state_idx=[0, 0, 0], ++ new_num_computed_tokens=[32, 40, 40], ++ block_size=_BLOCK_SIZE, ++ ) ++ kwargs.update(overrides) ++ return kwargs ++ ++ ++def test_clean_case_is_clean() -> None: ++ assert mamba_slot_ownership_violations(**_clean_kwargs()) == [] ++ ++ ++def test_a_row_that_does_not_reach_the_block_table_is_not_judged() -> None: ++ """``needs_copy`` false means the kernel returns before any dereference. ++ ++ Judging such a row would make the check fire on rows the kernel never ++ touches, i.e. on correct metadata. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ # One batch row, mapped to slot 1, whose 40 computed tokens leave ++ # the accepted window inside the running block: needs_copy false. ++ # The row rule is WRONG here (batch row 0 != slot 1) and the check ++ # still says nothing, because the kernel never reaches the table. ++ num_reqs=1, ++ bt_row_is_req=False, ++ idx_mapping=[1, 0], ++ new_num_computed_tokens=[32, 40, 40], ++ ) ++ ) ++ assert findings == [] ++ ++ ++def test_two_requests_resuming_into_one_state_block_is_a_collision() -> None: ++ """The llama.cpp#29092 shape with the row rule CORRECT. ++ ++ Row attribution and destination agreement both pass -- each request reads ++ its own row -- and the tables still hand both of them block 11. Only the ++ collision property catches this, which is why it is separate and not a ++ corollary of the row rule. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ block_tables=[[[10, 11, 12], [20, 11, 22], [30, 31, 32]]], ++ new_num_computed_tokens=[32, 32, 40], ++ ) ++ ) ++ assert len(findings) == 1, findings ++ assert findings[0].startswith("block collision:") ++ assert "slot 0" in findings[0] and "slot 1" in findings[0] ++ ++ ++def test_null_block_id_rows_are_not_a_collision() -> None: ++ """Block id 0 is the pool's placeholder; no request owns it. ++ ++ ``NULL_BLOCK_ID`` (``v1/attention/backends/utils.py:46``) is popped out of ++ the free queue at construction (``core/block_pool.py:188-191``), so two ++ rows carrying 0 are two unallocated columns, not a shared state block. ++ Treating that as a collision would fire on every partially allocated ++ batch. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ block_tables=[[[10, 0, 12], [20, 0, 22], [30, 31, 32]]], ++ new_num_computed_tokens=[32, 32, 40], ++ ) ++ ) ++ assert findings == [] ++ ++ ++def test_a_row_outside_the_persistent_table_is_reported_and_not_indexed() -> None: ++ """Out-of-range rows are the bounds guards' business, but the check must ++ not itself index past the table while reporting one.""" ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ num_reqs=3, ++ bt_row_is_req=False, ++ idx_mapping=[2, 1, 0], ++ num_rows=2, ++ new_num_computed_tokens=[32, 40, 32], ++ ) ++ ) ++ # Batch row 2 resolves to table row 2 with only 2 rows registered. ++ assert any(f.startswith("row bound:") for f in findings), findings ++ # Batch row 0 is in range but wrong, and is reported as such -- the two ++ # findings are different properties, not one restated. ++ assert any(f.startswith("row attribution:") for f in findings), findings ++ assert any(f.startswith("destination disagreement:") for f in findings), findings ++ ++ ++def test_a_request_slot_past_the_decision_arrays_is_reported() -> None: ++ """The per-slot arrays are read at ``req_idx``; a slot past them is ++ malformed metadata, and the checker must say so rather than raise.""" ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs(num_reqs=1, idx_mapping=[7]) ++ ) ++ assert len(findings) == 1, findings ++ assert findings[0].startswith("slot bound:") ++ ++ ++def test_a_negative_request_slot_is_skipped() -> None: ++ """``req_idx < 0`` is the kernel's own inactive-row sentinel.""" ++ assert ( ++ mamba_slot_ownership_violations( ++ **_clean_kwargs(idx_mapping=[-1, 1], bt_row_is_req=False) ++ ) ++ == [] ++ ) ++ ++ ++def test_a_short_idx_mapping_is_reported_rather_than_raising() -> None: ++ """An instrument that dies is an instrument that is not there.""" ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs(num_reqs=3, idx_mapping=[0, 1]) ++ ) ++ assert any("idx_mapping is shorter" in f for f in findings), findings ++ ++ ++# --------------------------------------------------------------------------- ++# What the narrowed instrument gained ++# --------------------------------------------------------------------------- ++ ++ ++def test_a_writer_colliding_with_a_live_non_copying_owner_is_caught() -> None: ++ """A destination-only collision map missed this case. ++ ++ Slot 1 is live this step and does NOT copy (40 computed tokens leave its ++ accepted window inside the running block), so it never appeared in a map ++ built from copying requests' destinations. Its current state block is 20. ++ Slot 0 copies into column 1, which also holds block 20. That is a writer ++ landing on a live owner's recurrent state -- the llama.cpp#29092 shape -- ++ and it must be reported. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ block_tables=[[[10, 20, 12], [20, 21, 22], [30, 31, 32]]], ++ new_num_computed_tokens=[32, 40, 40], ++ ) ++ ) ++ assert len(findings) == 1, findings ++ assert findings[0].startswith("block collision:") ++ assert "live source" in findings[0], findings[0] ++ ++ ++def test_a_negative_destination_column_is_reported_not_wrapped() -> None: ++ """A negative column indexes a Python list FROM THE END. ++ ++ Accepted count 1 with computed count 0 gives ``dst_col == -1``. The device ++ guard masks that access. An unbounded reader here would instead fetch the ++ LAST column and could report a collision on a block the kernel never ++ touches -- an instrument inventing a finding is worse than a silent one. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs( ++ num_reqs=1, ++ idx_mapping=[0], ++ num_accepted_tokens=[1, 2, 2], ++ new_num_computed_tokens=[0, 40, 40], ++ ) ++ ) ++ assert findings, "a -1 destination column must be reported" ++ assert all(f.startswith("malformed metadata:") for f in findings), findings ++ # The last column (block 12) must not have been read as a real claim. ++ assert not any("12" in f and "collision" in f for f in findings), findings ++ ++ ++def test_pointer_registration_disagreement_is_reported() -> None: ++ """The one INDEPENDENT source of truth this layer has. ++ ++ Everything else in this checker is derived from the same table it is ++ validating. If the kernel was handed a different address, every other ++ finding in the run is about the wrong tensor, and zero findings would be ++ meaningless rather than reassuring. ++ """ ++ findings = mamba_slot_ownership_violations( ++ **_clean_kwargs(registered_ptrs=[0xDEAD0000], view_ptrs=[0xBEEF0000]) ++ ) ++ assert len(findings) == 1, findings ++ assert findings[0].startswith("pointer registration:") ++ ++ ++def test_matching_pointers_are_silent() -> None: ++ """A check that fires on correct metadata is worse than no check.""" ++ assert ( ++ mamba_slot_ownership_violations( ++ **_clean_kwargs(registered_ptrs=[0xCAFE0000], view_ptrs=[0xCAFE0000]) ++ ) ++ == [] ++ ) +diff --git a/vllm/envs.py b/vllm/envs.py +index 554c025..157dfaf 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -603,6 +603,30 @@ def _resolve_rust_cli_path() -> str | None: + return raw + + ++def _e2_diagnostic_mode(name: str) -> str: ++ """Normalize an E2 diagnostic gate to off|warn|raise. ++ ++ One normalizer for every knob in the family, so a registry entry and the ++ module that acts on it cannot disagree about what the environment said. ++ ``1`` is accepted as ``warn`` for continuity with the earlier boolean ++ gates. A set-but-unrecognised value is a boot error, not a silent ``off``: ++ a broken gate that disarms its own instrument is the failure mode these ++ diagnostics exist to avoid. ++ """ ++ # An EMPTY value means off, not a boot error. Shell scripts and pytest ++ # fixtures both use `VAR=` to mean "unset" (see tests/v1/worker/ ++ # test_e2_memstats.py), and rejecting it would turn a routine way of ++ # clearing the variable into a crash. ++ value = os.environ.get(name, "off").strip().lower() or "off" ++ if value == "1": ++ return "warn" ++ if value not in ("off", "warn", "raise"): ++ raise ValueError( ++ f"{name} must be off|warn|raise (or 1 for warn), got {value!r}" ++ ) ++ return value ++ ++ + environment_variables: dict[str, Callable[[], Any]] = { + # ================== Installation Time Env Vars ================== + # Target device of vLLM, supporting [cuda (by default), +@@ -2181,6 +2205,23 @@ environment_variables: dict[str, Callable[[], Any]] = { + # Each op additionally checks its own shape / dtype constraints and falls + # back to the eager path when they do not hold. + "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))), ++ # PLE diagnostics. Both are off|warn|raise (1 is accepted as warn) and both ++ # default to off. They are registered here so the registry is the single ++ # source of truth for their normalization -- the modules that use them read ++ # these getters rather than calling os.getenv a second time, so the value ++ # the registry reports cannot drift from the value the check applies. ++ # Both are excluded from compile factors in `compile_factors` below: they ++ # gate HOST-side Python outside any capture and add no node to a compiled ++ # graph. If either ever grows a device-side component, remove it from ++ # `ignored_factors` in the same change. ++ # ++ # Read the pull-transport PLE semaphore back before a real launch, on the ++ # publishing rank only (vllm#53899 one-step-lag class). ++ "VLLM_E2_PLE_FLAG_CHECK": lambda: _e2_diagnostic_mode("VLLM_E2_PLE_FLAG_CHECK"), ++ # Host-side align row-rule and state-block collision check. ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP": lambda: _e2_diagnostic_mode( ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP" ++ ), + } + + +@@ -2273,6 +2314,13 @@ def compile_factors() -> dict[str, object]: + "VLLM_DEBUG_DUMP_PATH", + "VLLM_PORT", + "VLLM_CACHE_ROOT", ++ # PLE host-side diagnostics: both gate Python that runs outside any ++ # capture and neither adds a node to a compiled graph, so hashing them ++ # would invalidate every compile cache the moment a diagnostic boot is ++ # taken -- which is exactly when the cache is wanted. Remove from this ++ # set if either ever grows a device-side component. ++ "VLLM_E2_PLE_FLAG_CHECK", ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP", + # Runtime memory-plan persistence; does not affect compiled graphs. + "VLLM_ENABLE_STARTUP_PLAN", + # Location-only derived paths: where a cache/config directory lives +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index 60fb094..e95653e 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -78,6 +78,59 @@ GUARD_BIAS = 1 << 31 + # One (n_masked, row, value) triple per guarded kernel. + SLOTS_PER_SITE = 3 + ++# Largest value the bias encoding can carry without overflowing the int64 slot. ++GUARD_MAX_ENCODABLE = (1 << 63) - 1 - GUARD_BIAS ++# Smallest one. At exactly -GUARD_BIAS the encoded word is 0, which the decode ++# cannot tell from "nothing recorded"; below it the encoded word is negative ++# and tl.atomic_max against a zeroed buffer never takes it. Both cases lose the ++# VALUE and neither loses the COUNT -- see decode_slot. ++GUARD_MIN_ENCODABLE = -GUARD_BIAS + 1 ++ ++ ++def encode_slot(value: int) -> int: ++ """Producer side of the bias encoding, in Python (for the limits tests). ++ ++ The guarded kernels inline `(value + _GUARD_BIAS).to(tl.int64)` and store ++ it with `tl.atomic_max`; this is the same arithmetic, exposed so the ++ encoding's limits can be tested without Triton and without a GPU. It is ++ NOT called by the kernels -- a Python function is not readable from a ++ @triton.jit body, which is why _GUARD_BIAS is a tl.constexpr there. ++ """ ++ return int(value) + GUARD_BIAS ++ ++ ++def decode_slot(biased: int) -> int | None: ++ """Reader side. ``None`` means "this slot recorded nothing". ++ ++ THE ENCODING'S LIMITS (closed here by test ++ rather than by change). A slot is an int64 zeroed at allocation and only ++ ever written with `tl.atomic_max`, so the recorded value must be biased to ++ survive negatives, and 0 has to mean "unwritten". Three classes of extreme ++ value are therefore not recoverable from the slot: ++ ++ * ``value == -GUARD_BIAS`` (INT32_MIN) encodes to exactly 0 and is ++ FALSIFIED: it reads back as "nothing recorded". ++ * ``value < -GUARD_BIAS`` encodes negative, `atomic_max` against 0 keeps ++ 0, and the value is LOST the same way. ++ * ``value > GUARD_MAX_ENCODABLE`` overflows the int64 slot. ++ ++ In every one of those cases the COUNT slot, which carries no bias and is ++ written with `atomic_add`, still advances. ++ ++ BUT `None` IS NOT A DIAGNOSIS OF "UNENCODABLE". The reader has ++ NO ordering dependency on the producer -- it copies through a pinned mirror ++ on a side stream with no dependency on the compute stream -- so a count ++ that has become visible with its maxima not yet visible produces exactly ++ the same reading. `_poll_once`'s one-tick confirm delay makes that less ++ likely; it does not make it impossible, and a delay is not a happens-before ++ relationship. So `None` means UNOBSERVED OR UNENCODABLE. ++ ++ What it never means is "this field recorded nothing while the count rose". ++ Anything reading these slots must say so, and must not build a fatal ++ decision on the maxima -- those decisions are count-based. ++ """ ++ return None if biased == 0 else int(biased) - GUARD_BIAS ++ + SITE_MAMBA_COPY = 0 + SITE_GDN_DECODE_STATE = 1 + SITE_QSA_PRE_INDEXER = 2 +@@ -424,8 +477,8 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None: + if n_masked <= _seen.get((device, site), 0): + continue + _seen[(device, site)] = n_masked +- row = biased_row - GUARD_BIAS if biased_row else None +- val = biased_val - GUARD_BIAS if biased_val else None ++ row = decode_slot(biased_row) ++ val = decode_slot(biased_val) + msg = ( + f"bounds guard masked {n_masked} out-of-range access(es) " + f"in {SITE_NAMES.get(site, site)} on {device} " +@@ -434,6 +487,25 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None: + f"access was suppressed, so this is a MASKED FAULT and not a " + f"correct result -- the metadata is wrong upstream." + ) ++ if row is None or val is None: ++ # Corrected wording. A None field has TWO ++ # possible causes and this reader cannot distinguish them, because ++ # it has no happens-before relationship with the producer: ++ # * UNOBSERVED -- the kernel incremented the count and its ++ # atomic_max on this slot is not yet visible to this copy; ++ # * UNENCODABLE -- the value was == INT32_MIN, below it, or past ++ # GUARD_MAX_ENCODABLE, so the bias encoding could not carry it. ++ # Either way it does NOT mean the field recorded nothing. The ++ # earlier wording claimed "unencodable" outright, which the ++ # one-tick confirm delay does not establish. ++ msg += ( ++ " A None field is UNOBSERVED OR UNENCODABLE, not absent: " ++ "either this copy predates the slot's atomic_max becoming " ++ "visible, or the value was == INT32_MIN, below it, or past " ++ f"{GUARD_MAX_ENCODABLE}. The reader has no ordering " ++ "dependency on the producer, so it cannot tell which; see " ++ "decode_slot." ++ ) + logger.warning("%s", msg) + if _MODE == "raise" or site in _ALWAYS_FATAL_SITES: + _arm(msg) +diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py +index 39d71b3..a6972f1 100644 +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -7,6 +7,7 @@ from typing import Any, NamedTuple + + import torch + ++import vllm.envs as envs + from vllm.config import CacheConfig + from vllm.logger import init_logger + from vllm.model_executor.layers.mamba.mamba_utils import ( +@@ -37,6 +38,276 @@ logger = init_logger(__name__) + # with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. + _GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) + ++# Production value for the align post-process's block-table row rule. True ++# means "the row is the request's own state slot", which is the only correct ++# rule for the PP-deferred launch (see run_fused_postprocess_align). It is a ++# module-level name rather than a literal at the launch site for exactly one ++# reason: the failing control for the slot-ownership check below has to be able ++# to flip it. NOTHING reads it from the environment, so there is no way to turn ++# it off on a serving host; tests monkeypatch it. ++ALIGN_BT_ROW_IS_REQ = True ++ ++# Host-side mamba state-slot ownership check (PLE diagnostics, llama.cpp#29092 class / ++# vllm#55506 failure mode). off | warn | raise, default off, so it costs ++# nothing in production. It is a DIAGNOSTIC, not a device guard: it copies ++# decision arrays and block-table rows to the host, which synchronizes. ++# ++# WHY IT EXISTS. Mamba state blocks are the one KV class with no zeroing and no ++# ownership instrument: ++# * single_type_kv_cache_manager.py:96-99 records newly allocated blocks for ++# zeroing only when the spec `isinstance(kv_cache_spec, AttentionSpec)`. ++# `MambaSpec` is not an `AttentionSpec`, so mamba state blocks are NEVER ++# zeroed by `KVBlockZeroer` (gpu/model_runner.py:1213). ++# * the host-side KV-zeroer ownership check (sched/scheduler.py:1508-1535) only ++# inspects `new_block_ids_to_zero`, a list mamba blocks are never on. ++# * a genuinely fresh request is still safe: gdn_attn.py:401 sets ++# `has_initial_state = context_lens_tensor > 0`, so at context_len 0 the ++# GDN kernel ignores whatever bytes the reused block holds. ++# * a RESUMING request is not. With `has_initial_state` true the kernel ++# consumes the block's contents, and the only thing that puts the right ++# contents there is this align copy. Everything else we have bounds ++# indices (the bounds guards): a bounds check catches an ++# out-of-range row, never a wrong-but-in-range one, which is the whole of ++# this defect class. ++# Read through the envs registry, not os.getenv: the registry owns the ++# normalization, so the value reported by ++# `vllm.envs` and the value this check applies cannot drift apart. It is ++# registered as a non-compile factor there, and an unrecognised value raises ++# on first read rather than silently disarming the instrument. ++_SLOT_OWNERSHIP_MODE: str = envs.VLLM_E2_MAMBA_SLOT_OWNERSHIP ++ ++_NULL_BLOCK_ID = 0 ++ ++ ++def mamba_slot_ownership_violations( ++ *, ++ num_reqs: int, ++ idx_mapping: list[int], ++ bt_row_is_req: bool, ++ block_tables: list[list[list[int]]], ++ num_rows: int, ++ num_accepted_tokens: list[int], ++ state_idx: list[int], ++ new_num_computed_tokens: list[int], ++ block_size: int, ++ registered_ptrs: list[int] | None = None, ++ view_ptrs: list[int] | None = None, ++) -> list[str]: ++ """ALIGN ROW-RULE AND SIMULTANEOUS BLOCK-COLLISION CHECK. ++ ++ Read the name literally. It was narrowed from "slot ownership": ++ with the production row rule ``row_used == req_idx``, so the ++ destination comparison below reduces to ``table[req_idx][dst_col]`` ++ against itself and is TAUTOLOGICAL on a healthy run. This function derives ++ what it expects and what it observes from THE SAME TABLE, and a table ++ cannot testify about its own staleness. What it does and does not ++ establish, stated so that zero findings are not over-read: ++ ++ ESTABLISHES ++ ++ 1. **Row attribution.** The row the kernel will use is the request's own ++ state slot. This is the property vllm#55506 violates: under PP the ++ deferred launch carries an ``idx_mapping`` captured ``pp_size`` steps ++ earlier, so ``batch_idx`` names a different request. The ++ ``BT_ROW_IS_REQ=False`` control exercises exactly this and nothing ++ wider. ++ 2. **Simultaneous block collision.** Within ONE step, no two live request ++ slots reach the same state block -- counting both the source and the ++ destination column, and counting slots that are live but not copying, ++ because a writer colliding with a non-copying live owner is the same ++ corruption. ``NULL_BLOCK_ID`` is excluded: block id 0 is the pool's ++ placeholder (``core/block_pool.py:188-191``), owned by no request. ++ 3. **Pointer-registration agreement**, when the caller supplies ++ ``registered_ptrs``: the tables read here are the tables whose addresses ++ the kernel was handed. That is the one genuinely INDEPENDENT source of ++ truth available at this layer, and without it a wrong pointer array ++ would be invisible -- the checker would faithfully validate a table the ++ kernel never reads. ++ ++ DOES NOT ESTABLISH -- do not quote zero findings as covering these: ++ ++ * **Stale ownership across steps.** The block-table row could itself be ++ stale; nothing here compares it against the scheduler's allocation. ++ * **Request-generation reuse.** A slot re-issued to a different request ++ with the same index is indistinguishable here: there is no generation ++ id in any tensor this function sees. ++ * **Cross-step source staleness.** A source block correct for this step but ++ written by a previous request's state is invisible. ++ * **Anything about a captured replay**, which does not re-enter the caller. ++ ++ Closing those needs an independent allocation/generation map from the ++ scheduler side, which is a plumbing change, not a diagnostic. ++ ++ Pure and side-effect free; every argument is a plain Python list or int, so ++ this runs with no GPU and no CUDA context. The caller owns the device->host ++ copies and the mode. ++ ++ ``block_tables[g][row][col]`` is the PERSISTENT per-req-slot table of mamba ++ group ``g``. ``registered_ptrs`` and ``view_ptrs``, when given, are the ++ kernel's ``persistent_block_table_ptrs`` values and the ``data_ptr()`` of ++ the tables actually read, compared elementwise. ++ ++ Returns one string per violation, empty when clean. ++ """ ++ violations: list[str] = [] ++ # (group, block id) -> (req slot, which column class claimed it). Both the ++ # source and the destination column are entered, and live-but-not-copying ++ # slots are entered too, so a writer colliding with a non-copying owner is ++ # caught (the destination-only map missed that case). ++ block_owner: dict[tuple[int, int], tuple[int, str]] = {} ++ ++ # (3) pointer-registration agreement -- the one independent source of truth ++ # available here. Without it the checker can faithfully validate a table ++ # the kernel never reads. ++ if registered_ptrs is not None and view_ptrs is not None: ++ if len(registered_ptrs) != len(view_ptrs): ++ violations.append( ++ f"pointer registration: {len(registered_ptrs)} registered " ++ f"block-table pointer(s) but {len(view_ptrs)} retained view(s)" ++ ) ++ for group_idx, (registered, view) in enumerate( ++ zip(registered_ptrs, view_ptrs) ++ ): ++ if int(registered) != int(view): ++ violations.append( ++ f"pointer registration: group {group_idx} kernel reads " ++ f"0x{int(registered):x} but the table checked here is at " ++ f"0x{int(view):x}. Every other finding in this run is " ++ f"about the wrong tensor." ++ ) ++ ++ def _claim(group_idx: int, block: int, req_idx: int, kind: str) -> None: ++ """Record one live claim on a state block and report a collision.""" ++ if block == _NULL_BLOCK_ID: ++ return # the pool's placeholder; owned by no request ++ key = (group_idx, block) ++ previous = block_owner.get(key) ++ if previous is not None and previous[0] != req_idx: ++ violations.append( ++ f"block collision: group {group_idx}, state block {block} is " ++ f"claimed in one step by request slot {previous[0]} " ++ f"({previous[1]}) and request slot {req_idx} ({kind}) " ++ f"(llama.cpp#29092 class: a prior request's recurrent state in " ++ f"a slot another request resumes into)." ++ ) ++ elif previous is None: ++ block_owner[key] = (req_idx, kind) ++ ++ def _cell(group_idx: int, table, row: int, col: int, what: str) -> int | None: ++ """Bounded read. Returns None and reports rather than indexing wrongly. ++ ++ A NEGATIVE column silently indexes a Python list from the end, so an ++ unbounded read here invents a block id and can report an unrelated ++ collision (example: accepted 1 with computed 0 yields dst_col -1, ++ which the device guard masks and this reader would turn into the LAST ++ column). Both bounds, both indices, every time. ++ """ ++ if not 0 <= row < len(table): ++ violations.append( ++ f"malformed metadata: group {group_idx}, {what} row {row} is " ++ f"outside the table's [0, {len(table)}) rows" ++ ) ++ return None ++ if not 0 <= col < len(table[row]): ++ violations.append( ++ f"malformed metadata: group {group_idx}, {what} column {col} " ++ f"is outside row {row}'s [0, {len(table[row])}) columns. The " ++ f"device guard masks this access; it is reported, not read." ++ ) ++ return None ++ return int(table[row][col]) ++ ++ # Pass 1: every LIVE slot's current state block, whether or not it copies. ++ # A slot that does not copy this step still owns its block, and a copying ++ # writer that lands on it is the corruption -- so the claims have to be ++ # entered before any writer is judged. ++ for batch_idx in range(min(num_reqs, len(idx_mapping))): ++ slot = int(idx_mapping[batch_idx]) ++ if slot < 0 or slot >= len(state_idx): ++ continue ++ for group_idx, table in enumerate(block_tables): ++ block = _cell(group_idx, table, slot, int(state_idx[slot]), "live src") ++ if block is not None: ++ _claim(group_idx, block, slot, "live source") ++ ++ for batch_idx in range(num_reqs): ++ if batch_idx >= len(idx_mapping): ++ violations.append( ++ f"idx_mapping is shorter ({len(idx_mapping)}) than num_reqs " ++ f"({num_reqs}); batch row {batch_idx} has no request slot" ++ ) ++ continue ++ req_idx = int(idx_mapping[batch_idx]) ++ if req_idx < 0: ++ continue # the kernel's own inactive-row sentinel ++ if req_idx >= min( ++ len(num_accepted_tokens), len(state_idx), len(new_num_computed_tokens) ++ ): ++ violations.append( ++ f"slot bound: batch row {batch_idx} maps to request slot " ++ f"{req_idx}, past the per-slot decision arrays " ++ f"({len(num_accepted_tokens)}, {len(state_idx)}, " ++ f"{len(new_num_computed_tokens)} entries)" ++ ) ++ continue ++ ++ # Mirror of the kernel's decision block (postprocess_mamba_fused_kernel, ++ # the PRECOMPUTED_NEW_COMPUTED branch). Every array is read at req_idx, ++ # never at batch_idx -- that asymmetry is the bug's hiding place. ++ num_accepted = int(num_accepted_tokens[req_idx]) ++ src_col = int(state_idx[req_idx]) ++ new_num_computed = int(new_num_computed_tokens[req_idx]) ++ num_tokens_running_state = new_num_computed - num_accepted + 1 ++ aligned_new_computed = (new_num_computed // block_size) * block_size ++ if aligned_new_computed < num_tokens_running_state: ++ continue # needs_copy is False: the block table is never read ++ accept_token_bias = aligned_new_computed - num_tokens_running_state ++ dst_col = aligned_new_computed // block_size - 1 ++ if src_col == dst_col and accept_token_bias == 0: ++ continue # no-op self-copy: the kernel returns before the copy ++ ++ row_used = req_idx if bt_row_is_req else batch_idx ++ ++ # (1) row attribution ++ if row_used != req_idx: ++ violations.append( ++ f"row attribution: batch row {batch_idx} is request slot " ++ f"{req_idx}, but the align post-process will read block-table " ++ f"row {row_used}. Under PP the idx_mapping was captured " ++ f"pp_size steps before this launch, so row {row_used} belongs " ++ f"to a different micro-batch's request (vllm#55506)." ++ ) ++ if not 0 <= row_used < num_rows: ++ violations.append( ++ f"row bound: batch row {batch_idx} resolves to block-table row " ++ f"{row_used}, outside [0, {num_rows})" ++ ) ++ continue ++ ++ for group_idx, table in enumerate(block_tables): ++ owner_block = _cell(group_idx, table, req_idx, dst_col, "owner dst") ++ used_block = _cell(group_idx, table, row_used, dst_col, "used dst") ++ if owner_block is None or used_block is None: ++ continue ++ # (2) destination agreement. NOTE: with the production row rule ++ # row_used == req_idx, so this comparison is tautological on a ++ # healthy run and proves nothing on its own. It earns ++ # its place only in the wrong-row world, where it reports the ++ # defect in BLOCK IDS rather than in indices -- which is what says ++ # the copy really would land on another request's state. ++ if used_block != owner_block: ++ violations.append( ++ f"destination disagreement: group {group_idx}, request " ++ f"slot {req_idx} would resume into state block " ++ f"{used_block} (row {row_used}, col {dst_col}) but its own " ++ f"row holds block {owner_block} there. The copy would " ++ f"write this request's recurrent state over another " ++ f"request's block." ++ ) ++ _claim(group_idx, used_block, req_idx, f"align destination col {dst_col}") ++ ++ return violations ++ + # 16 saturates HBM on H100/GB200 across the reqs=8..128 range in + # microbenchmarks + _TEMPORAL_TILES = 16 +@@ -199,6 +470,11 @@ def _reinterpret_u64_as_i64(value: int) -> int: + return value if value < (1 << 63) else value - (1 << 64) + + ++def _reinterpret_i64_as_u64(value: int) -> int: ++ """Inverse of ``_reinterpret_u64_as_i64``, for comparing pointers host side.""" ++ return value + (1 << 64) if value < 0 else value ++ ++ + @triton.jit + def _copy_mamba_state_block( + state_idx, +@@ -923,6 +1199,13 @@ class MambaSpecDecodeGPUContext: + # copying state. -1 is unusable on purpose and is asserted before use. + persistent_block_table_num_rows: int = -1 + ++ # The same tensors `persistent_block_table_ptrs` addresses, retained so the ++ # host-side slot-ownership check can READ the rows the kernel dereferences ++ # rather than model them. Holding these costs no allocation: they are the ++ # runner's own `BlockTables.block_tables[i].gpu`, alive for the engine's ++ # lifetime, and nothing here writes to them. ++ persistent_block_table_views: list[torch.Tensor] | None = None ++ + # persistent output for the once-per-step, all-group aligned-index launch. + # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks]. + aligned_state_indices: torch.Tensor | None = None +@@ -941,6 +1224,13 @@ class MambaSpecDecodeGPUContext: + # Flag to track if metadata has been populated + is_initialized: bool = False + ++ # Last result of the host-side slot-ownership check (always [] unless ++ # VLLM_E2_MAMBA_SLOT_OWNERSHIP is set). Readable by a test so the control ++ # does not have to parse log records. ++ last_slot_ownership_violations: list[str] = dataclasses.field( ++ default_factory=list ++ ) ++ + @classmethod + def create( + cls, +@@ -1279,6 +1569,7 @@ class MambaSpecDecodeGPUContext: + # the same-step guard takes it too. + self.block_table_num_rows = self.persistent_block_table_num_rows + self.has_persistent_block_tables = True ++ self.persistent_block_table_views = list(persistent_block_tables) + else: + # Legal for V1 callers, which never reach the V2 align + # post-process. If a V2 align launch does happen on such a context, +@@ -1497,6 +1788,15 @@ class MambaSpecDecodeGPUContext: + num_accepted_tokens_snapshot = self.num_accepted_tokens_out + num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) + ++ if _SLOT_OWNERSHIP_MODE != "off": ++ self._check_slot_ownership( ++ num_reqs=num_reqs, ++ num_accepted_tokens=num_accepted_tokens_snapshot, ++ state_idx_gpu=state_idx_gpu, ++ new_num_computed_tokens_gpu=new_num_computed_tokens_gpu, ++ idx_mapping=idx_mapping, ++ ) ++ + total_states = self.num_states + grid = (num_reqs, total_states, _TEMPORAL_TILES) + postprocess_mamba_fused_kernel[grid]( +@@ -1527,10 +1827,124 @@ class MambaSpecDecodeGPUContext: + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), + HAS_IDX_MAPPING=True, + PRECOMPUTED_NEW_COMPUTED=True, +- BT_ROW_IS_REQ=True, ++ BT_ROW_IS_REQ=ALIGN_BT_ROW_IS_REQ, + TEMPORAL_TILES=_TEMPORAL_TILES, + ) + ++ def _check_slot_ownership( ++ self, ++ *, ++ num_reqs: int, ++ num_accepted_tokens: torch.Tensor, ++ state_idx_gpu: torch.Tensor, ++ new_num_computed_tokens_gpu: torch.Tensor, ++ idx_mapping: torch.Tensor, ++ ) -> list[str]: ++ """Host-side mamba state-slot ownership check, off by default. ++ ++ VLLM_E2_MAMBA_SLOT_OWNERSHIP=warn|raise. Same discipline as the KV ++ zeroer's ownership check (VLLM_E2_ZERO_OWNERSHIP_CHECK): a diagnostic in the ++ scheduler-visible sense -- host side, opt-in, never a device guard. ++ Unlike that check it lives here, because the thing at risk is not visible to ++ the scheduler at all: `MambaSpec` is not an `AttentionSpec`, so mamba ++ state blocks never appear on `new_block_ids_to_zero` and that check never ++ sees them. ++ ++ COST, and it is not only wall-clock. Four device->host copies plus one ++ block-table slice per call, each a synchronize. That SERIALIZES work ++ that is normally overlapped, so an instrumented boot has materially ++ different timing from a production one: zero findings under this check ++ cannot by itself exclude a timing-sensitive defect. Run the ++ instrumented arm for diagnosis and an uninstrumented arm for the ++ concurrency qualification; do not merge their conclusions. ++ ++ CAPTURE CONTRACT, stated explicitly rather than assumed. These copies ++ are illegal inside a stream capture, so the function asserts it is not ++ capturing instead of relying on "this is engine-thread Python" -- Python ++ also runs while a graph is being captured, and the caller's own contract ++ is the thing being checked here. A captured REPLAY does not re-enter ++ this function at all, which is a real scope limit of the instrument, ++ not a safety argument. ++ ++ Returns the violation list (also stored on the context so a test can ++ read it without parsing logs). ++ """ ++ if torch.cuda.is_current_stream_capturing(): ++ # The D2H copies below would abort the capture. Report rather than ++ # silently skip: a diagnostic reached from a path it cannot serve ++ # is a contract error in the caller. ++ logger.warning_once( ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP is set but the align " ++ "post-process was reached during a cudagraph capture, where " ++ "the check's device->host copies are illegal. Skipped for " ++ "this launch; the caller's outside-capture contract does not " ++ "hold." ++ ) ++ self.last_slot_ownership_violations = [] ++ return [] ++ tables = self.persistent_block_table_views ++ if not tables: ++ # Unreachable from run_fused_postprocess_align, which raises above ++ # when the persistent tables were never registered. Kept so the ++ # check degrades to silence rather than to an exception if some ++ # future caller reaches it first. ++ logger.warning_once( ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP is set but the persistent " ++ "per-req-slot block tables were never registered, so slot " ++ "ownership cannot be checked." ++ ) ++ self.last_slot_ownership_violations = [] ++ return [] ++ ++ num_cols = min(int(t.shape[1]) for t in tables) ++ violations = mamba_slot_ownership_violations( ++ num_reqs=int(num_reqs), ++ idx_mapping=idx_mapping.detach().to("cpu").reshape(-1).tolist(), ++ bt_row_is_req=bool(ALIGN_BT_ROW_IS_REQ), ++ block_tables=[ ++ t[:, :num_cols].detach().to("cpu").tolist() for t in tables ++ ], ++ num_rows=int(self.persistent_block_table_num_rows), ++ # The independent evidence: what the kernel was handed, against ++ # what is being read here. `persistent_block_table_ptrs` holds the ++ # addresses reinterpreted as signed, so undo that before comparing. ++ registered_ptrs=( ++ [ ++ _reinterpret_i64_as_u64(int(v)) ++ for v in self.persistent_block_table_ptrs.detach() ++ .to("cpu") ++ .reshape(-1) ++ .tolist()[: len(tables)] ++ ] ++ if self.persistent_block_table_ptrs is not None ++ else None ++ ), ++ view_ptrs=[int(t.data_ptr()) for t in tables], ++ num_accepted_tokens=num_accepted_tokens.detach() ++ .to("cpu") ++ .reshape(-1) ++ .tolist(), ++ state_idx=state_idx_gpu.detach().to("cpu").reshape(-1).tolist(), ++ new_num_computed_tokens=new_num_computed_tokens_gpu.detach() ++ .to("cpu") ++ .reshape(-1) ++ .tolist(), ++ block_size=int(self.block_size), ++ ) ++ self.last_slot_ownership_violations = violations ++ if not violations: ++ return violations ++ message = ( ++ f"Mamba state-slot ownership violation on the align " ++ f"post-process ({len(violations)} finding(s), num_reqs=" ++ f"{num_reqs}, BT_ROW_IS_REQ={ALIGN_BT_ROW_IS_REQ}): " ++ + " | ".join(violations) ++ ) ++ if _SLOT_OWNERSHIP_MODE == "raise": ++ raise RuntimeError(message) ++ logger.warning("%s", message) ++ return violations ++ + + @dataclasses.dataclass + class MambaBuffers: +-- +2.43.0 + diff --git a/patches/0028-envs-tests-register-the-fork-s-knobs-by-cache-identi.patch b/patches/0028-envs-tests-register-the-fork-s-knobs-by-cache-identi.patch new file mode 100644 index 0000000..c79bbe5 --- /dev/null +++ b/patches/0028-envs-tests-register-the-fork-s-knobs-by-cache-identi.patch @@ -0,0 +1,1845 @@ +From 3d563a05b7ad5e4d616b1760ab443bc8785dbc58 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:20:50 -0400 +Subject: [PATCH 028/103] envs, tests: register the fork's knobs by cache + identity; lint; gpu test marker + +- Every fork knob is registered in envs.py and classified as a compile/capture + factor or runtime-only; envs.py documents that capture identity and AOT + cache identity are different things. +- ruff findings cleared on lines this fork authored; upstream lines untouched. +- A `gpu` pytest marker so the GPU-free tier can be selected, not just skipped. +- Duplicate knob-table entries from two branches reconciled onto one + off|warn|raise normaliser. +--- + pyproject.toml | 1 + + tests/conftest.py | 19 ++ + .../qwen4_exp/test_host_uva_embedding.py | 6 +- + tests/test_envs.py | 129 +++++++++++ + .../test_multiproc_executor_parent_death.py | 10 +- + .../v1/worker/gpu/test_block_table_guards.py | 11 +- + .../gpu/test_guard_reader_nonblocking.py | 81 ++++--- + .../test_guard_telemetry_defaults_drain.py | 15 +- + .../worker/gpu/test_guard_telemetry_replay.py | 4 +- + .../gpu/test_pp_deferred_mamba_block_table.py | 46 ++-- + .../worker/gpu/test_pp_deferred_mamba_gpu.py | 32 +-- + .../gpu/test_qsa_ring_skip_telemetry.py | 86 ++++--- + .../gpu/test_qsa_rope_position_guard.py | 151 ++++++++---- + tests/v1/worker/test_grammar_fail_closed.py | 3 +- + vllm/envs.py | 219 +++++++++++++++++- + vllm/models/qwen4_exp/nvidia/qsa.py | 7 +- + vllm/v1/core/sched/e2_alloc_log.py | 148 +++++++++--- + vllm/v1/ple_offload/test_fault.py | 6 +- + vllm/v1/worker/guard_telemetry.py | 37 +-- + vllm/v1/worker/mamba_utils.py | 2 +- + vllm_ple_test_fault/__init__.py | 15 +- + 21 files changed, 770 insertions(+), 258 deletions(-) + +diff --git a/pyproject.toml b/pyproject.toml +index 1a77433..6a72b74 100644 +--- a/pyproject.toml ++++ b/pyproject.toml +@@ -114,6 +114,7 @@ markers = [ + "cpu_test: mark test as CPU-only test", + "split: run this test as part of a split", + "distributed: run this test only in distributed GPU tests", ++ "gpu: requires a CUDA device; auto-skipped when none is present, and deselectable as a tier with -m 'not gpu'", + "optional: optional tests that are automatically skipped, include --optional to run them", + ] + +diff --git a/tests/conftest.py b/tests/conftest.py +index 5fd594a..982caee 100644 +--- a/tests/conftest.py ++++ b/tests/conftest.py +@@ -1600,6 +1600,25 @@ def pytest_addoption(parser): + + + def pytest_collection_modifyitems(config, items): ++ # Flash-Next fork (v2.2): one place that decides what a GPU test is. ++ # `@pytest.mark.gpu` auto-skips without a device and makes the GPU-free tier ++ # SELECTABLE as `-m "not gpu"`, which a per-test skipif cannot do -- a ++ # skipif still reports the item as a skip rather than removing it. ++ # ++ # What this does NOT do: it does not make collection ++ # GPU-dependency-free. This hook runs after every module has been imported ++ # and every item collected, so a module whose import reaches CUDA, Triton or ++ # an absent vLLM extension still fails at COLLECTION, before any marker is ++ # consulted. Making the CPU tier collect cleanly is a separate job: defer ++ # those imports, or run the tier on its declared supported platform. ++ # The marker is also not authorization to use a busy shared device -- ++ # test_pp_deferred_mamba_gpu.py keeps its separate VLLM_E2_GPU_TESTS opt-in. ++ if not torch.cuda.is_available(): ++ skip_gpu = pytest.mark.skip(reason="no CUDA device (pytest marker: gpu)") ++ for item in items: ++ if "gpu" in item.keywords: ++ item.add_marker(skip_gpu) ++ + if config.getoption("--optional"): + # --optional given in cli: do not skip optional tests + return +diff --git a/tests/models/qwen4_exp/test_host_uva_embedding.py b/tests/models/qwen4_exp/test_host_uva_embedding.py +index e55d5d5..4edc157 100644 +--- a/tests/models/qwen4_exp/test_host_uva_embedding.py ++++ b/tests/models/qwen4_exp/test_host_uva_embedding.py +@@ -19,7 +19,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + host_embed_table_enabled, + ) + +- + VOCAB = 67 + HIDDEN = 16 + +@@ -479,10 +478,7 @@ def test_u4_outer_post_load_preserves_pinned_uva_storage(monkeypatch): + loaded_weight = _weight(vocab, hidden) + local_table_bytes = layer.weight.numel() * layer.weight.element_size() + +- assert ( +- layer.weight_loader.__func__ +- is VocabParallelEmbedding.weight_loader +- ) ++ assert layer.weight_loader.__func__ is VocabParallelEmbedding.weight_loader + layer.weight_loader(layer.weight, loaded_weight) + + model = nn.Module() +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 1ef40cd..3e65518 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + import os ++from pathlib import Path + from unittest.mock import patch + + import pytest +@@ -587,3 +588,131 @@ class TestVllmMaxNSequences: + + with pytest.raises(VLLMValidationError, match="n must be at most 128"): + SamplingParams(n=129) ++ ++ ++# --------------------------------------------------------------------------- ++# Flash-Next fork: the cache-factor contract (v2.2 fork knobs) ++# ++# compile_factors() is a deny-list, so a knob's classification is invisible ++# unless something asserts it. These tests are that assertion. They are CPU-only ++# and import nothing beyond vllm.envs, so they belong to the GPU-free tier. ++# ++# AOT identity vs CAPTURE identity: compile_factors() keys only the first. A ++# knob listed as hashed is genuinely protected across processes; a knob that is ++# only a capture factor is hashed conservatively and its capture safety still ++# rests on not changing the value between capture and replay in one process. ++# These tests check the AOT side, which is the side this file can check. ++# ++# Each row is (name, value to set, value the getter must then return). The ++# expected value matters: a getter that ignored the environment would still ++# satisfy a bare membership check, so every row proves the entry is live. ++# --------------------------------------------------------------------------- ++ ++_FORK_HASHED = [ ++ ("VLLM_HOST_EMBED_TABLE", "1", True), ++ ("VLLM_PLE_FP8_GLOBAL_SCALE", "1", True), ++ ("VLLM_QSA_KV_CLIP_COUNT", "7", "7"), ++ ("VLLM_QSA_KV_COLLECT", "/var/tmp/qsa-absmax", "/var/tmp/qsa-absmax"), ++ # Capture-configuration factor read in C++ (topk.cu exact_medium_enabled); ++ # the Python entry is a raw passthrough, never a second parser. ++ ("MERLIN_TOPK_EXACT_MEDIUM", "0", "0"), ++] ++ ++_FORK_IGNORED = [ ++ ("VLLM_QSA_KV_SCALES", "/var/tmp/scales.json", "/var/tmp/scales.json"), ++ # "WARN" also checks the .strip().lower() mirror of the read site. ++ ("VLLM_E2_GUARD_MODE", " WARN ", "warn"), ++ ("VLLM_E2_GUARD_INTERVAL", "0.25", "0.25"), ++ ("VLLM_E2_MEMSTATS", "2", "2"), ++ ("VLLM_E2_MEMSTATS_DIR", "/var/tmp/memstats", "/var/tmp/memstats"), ++ ("VLLM_E2_SCHED_LOG", "/var/tmp/sched.jsonl", "/var/tmp/sched.jsonl"), ++ ("VLLM_E2_ZERO_OWNERSHIP_CHECK", "1", True), ++ ("VLLM_PLE_OFFLOAD_HOME_DEVICE", "2", "2"), ++ ("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "1800", 1800.0), ++ ("VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", "45", 45.0), ++ ("VLLM_PLE_PROGRESS_DEADLINE_S", "30", 30.0), ++ ("VLLM_PLE_STARTUP_DEADLINE_S", "60", 60.0), ++ ("VLLM_PLE_TEST_FAULT", "cut=exc_in_release,ordinal=3,rank=any", None), ++ ("VLLM_PLE_TEST_FAULT_GATE", "/var/tmp/gate", "/var/tmp/gate"), ++ ("VLLM_PLE_TEST_FAULT_LOG", "/var/tmp/fault.log", "/var/tmp/fault.log"), ++ # Read sites arrive with the slot-ownership assertions. ++ ("VLLM_E2_PLE_FLAG_CHECK", "1", "warn"), ++ ("VLLM_E2_MAMBA_SLOT_OWNERSHIP", " RAISE ", "raise"), ++] ++ ++# Not in this repository at all: a reference-host-side patch stack. Named in ++# the TODO block in envs.py and deliberately unregistered until someone reads ++# their real read sites and classifies both identities. ++_FORK_UNAUDITED = [ ++ "MERLIN_FULL_K", ++ "MERLIN_QSA_SORT", ++ "MERLIN_TIE_DET", ++ "MERLIN_TIE_RECENT", ++ "MERLIN_PIN_BSM", ++] ++ ++ ++def _fresh_env_value(name: str): ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ return getattr(envs, name) ++ ++ ++@pytest.mark.parametrize(("name", "value", "expected"), _FORK_HASHED) ++def test_fork_knob_is_a_cache_factor( ++ name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch ++): ++ """A knob that can change a compiled or captured artifact must be hashed, ++ and the hashed value must follow the environment.""" ++ assert name in environment_variables ++ monkeypatch.setenv(name, value) ++ assert _fresh_env_value(name) == expected ++ factors = envs.compile_factors() ++ assert name in factors ++ assert factors[name] == expected ++ ++ ++@pytest.mark.parametrize(("name", "value", "expected"), _FORK_IGNORED) ++def test_fork_runtime_knob_is_not_a_cache_factor( ++ name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch ++): ++ """A path, timeout, log cadence or host-side assertion must never ++ invalidate the compile cache -- while still being read.""" ++ assert name in environment_variables ++ monkeypatch.setenv(name, value) ++ if expected is not None: ++ assert _fresh_env_value(name) == expected ++ assert name not in envs.compile_factors() ++ ++ ++@pytest.mark.parametrize("name", _FORK_UNAUDITED) ++def test_unaudited_host_side_knob_stays_listed_until_it_is_classified(name: str): ++ """The TODO block in envs.py is a gate, not a note. ++ ++ Registering one of these without classifying it -- which means deleting it ++ from that block once its AOT and capture identity are written down against ++ its real read site -- is the exact failure the fork's env rule ++ exists to prevent, so it fails here. ++ """ ++ source = Path(envs.__file__).read_text(encoding="utf-8") ++ listed = name in source ++ if name in environment_variables: ++ assert not listed, ( ++ f"{name} is registered but still named in the envs.py TODO block: " ++ "classify its AOT and capture identity from its real read site and " ++ "remove it from that block" ++ ) ++ else: ++ assert listed, ( ++ f"{name} is neither registered nor named in the envs.py TODO block; " ++ "every fork knob must be accounted for" ++ ) ++ ++ ++def test_fork_cache_factor_lists_are_disjoint_and_registered(): ++ """Negative control for the two tables above.""" ++ hashed = {n for n, _, _ in _FORK_HASHED} ++ ignored = {n for n, _, _ in _FORK_IGNORED} ++ assert not hashed & ignored ++ for name in hashed | ignored: ++ assert name in environment_variables, name +diff --git a/tests/v1/executor/test_multiproc_executor_parent_death.py b/tests/v1/executor/test_multiproc_executor_parent_death.py +index c39718f..111fca6 100644 +--- a/tests/v1/executor/test_multiproc_executor_parent_death.py ++++ b/tests/v1/executor/test_multiproc_executor_parent_death.py +@@ -70,9 +70,7 @@ def _run_fake_worker_main( + self.peer_response_handles = [] + + def monitor_death_pipe(self, pipe, shutdown_requested): +- return real_worker_proc.monitor_death_pipe( +- self, pipe, shutdown_requested +- ) ++ return real_worker_proc.monitor_death_pipe(self, pipe, shutdown_requested) + + def worker_busy_loop(self): + if block_mode == "gil-spin": +@@ -271,11 +269,7 @@ def _wait_for_state( + def _wait_for_exit(pid: int, start_time: str, deadline: float) -> float | None: + while time.monotonic() < deadline: + identity = _proc_identity(pid) +- if ( +- identity is None +- or identity[0] in {"X", "Z"} +- or identity[1] != start_time +- ): ++ if identity is None or identity[0] in {"X", "Z"} or identity[1] != start_time: + return time.monotonic() + time.sleep(0.01) + return None +diff --git a/tests/v1/worker/gpu/test_block_table_guards.py b/tests/v1/worker/gpu/test_block_table_guards.py +index d15093c..b4f0450 100644 +--- a/tests/v1/worker/gpu/test_block_table_guards.py ++++ b/tests/v1/worker/gpu/test_block_table_guards.py +@@ -20,9 +20,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID + from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.gpu.block_table import BlockTables + +-pytestmark = pytest.mark.skipif( +- not torch.cuda.is_available(), reason="CUDA required" +-) ++pytestmark = pytest.mark.gpu + + MAX_REQS = 8 + BLOCKS = 4 +@@ -67,8 +65,7 @@ def test_gather_and_slot_mapping_launch_and_are_correct(): + # A legal batch must leave both counters at zero. A guard that fires here + # is reporting a masked fault on correct metadata, which is worse than no + # guard -- see the row-bound defect. +- for site in (guard_telemetry.SITE_BT_GATHER, +- guard_telemetry.SITE_SLOT_MAPPING): ++ for site in (guard_telemetry.SITE_BT_GATHER, guard_telemetry.SITE_SLOT_MAPPING): + ctr = guard_telemetry.counter(site, device) + assert int(ctr[0].item()) == 0, guard_telemetry.snapshot() + +@@ -80,7 +77,7 @@ def test_gather_guard_fires_and_zeroes_the_row(): + bt = _tables(device) + + idx = torch.arange(NUM_REQS, dtype=torch.int32, device=device) +- idx[2] = MAX_REQS + 3 # a request row that does not exist ++ idx[2] = MAX_REQS + 3 # a request row that does not exist + (gathered,) = bt.gather_block_tables(idx, num_reqs_padded=NUM_REQS) + torch.cuda.synchronize() + +@@ -103,7 +100,7 @@ def test_slot_mapping_guard_pads_the_bad_request(): + bt = _tables(device) + + idx = torch.arange(NUM_REQS, dtype=torch.int32, device=device) +- idx[1] = -1 # the freed-slot sentinel ++ idx[1] = -1 # the freed-slot sentinel + qsl = torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32, device=device) + pos = torch.arange(NUM_TOKENS, dtype=torch.int64, device=device) + slots = bt.compute_slot_mappings(idx, qsl, pos, num_tokens_padded=NUM_TOKENS) +diff --git a/tests/v1/worker/gpu/test_guard_reader_nonblocking.py b/tests/v1/worker/gpu/test_guard_reader_nonblocking.py +index 94721e1..4fc7995 100644 +--- a/tests/v1/worker/gpu/test_guard_reader_nonblocking.py ++++ b/tests/v1/worker/gpu/test_guard_reader_nonblocking.py +@@ -1,11 +1,15 @@ + # SPDX-License-Identifier: Apache-2.0 +-"""the instrumentation readers must never wait on the compute stream. +- +-The first cadence reader did a synchronous pageable ``tolist()`` from its thread; the driver held the CUDA context lock +-while that copy waited for the stream, and a pending PLE pull-transport wait on the stream (whose producer needs the +-same lock) deadlocked the worker. These tests park a kernel on the current (compute) stream and require a reader tick +-to return while it is still running. On the pre-fix reader the first GPU test blocks for the kernel's duration. ++"""v2 reader fix: the instrumentation readers must never wait on ++the compute stream. ++ ++The earlier cadence reader did a synchronous pageable ``tolist()`` from its thread; ++the driver held the CUDA context lock while that copy waited for the stream, ++and a pending PLE pull-transport wait on the stream (whose producer needs the ++same lock) deadlocked the worker. These tests park a kernel on the current ++(compute) stream and require a reader tick to return while it is still running. ++On the pre-fix reader the first GPU test blocks for the kernel's duration. + """ ++ + import threading + import time + +@@ -14,7 +18,9 @@ import torch + + from vllm.v1.worker import guard_telemetry + +-SLEEP_CYCLES = 4_000_000_000 # ~2 s on a 3090 (torch.cuda._sleep spins for this many cycles) ++SLEEP_CYCLES = ( ++ 4_000_000_000 # ~2 s on a 3090 (torch.cuda._sleep spins for this many cycles) ++) + + + def test_cpu_tensor_reads_directly(): +@@ -27,7 +33,8 @@ def test_cpu_tensor_reads_directly(): + + + def test_register_refuses_to_run_inside_a_capture(monkeypatch): +- """Creating the side stream inside a cudagraph capture is not capture-safe.""" ++ """Creating the side stream inside a cudagraph capture is ++ not capture-safe.""" + r = guard_telemetry.NonBlockingReader() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + if not torch.cuda.is_available(): +@@ -38,17 +45,18 @@ def test_register_refuses_to_run_inside_a_capture(monkeypatch): + + + def test_report_waits_for_the_confirming_tick_cpu(): +- """A count rise seen in copy k is reported from copy k+1 (maxima current by then).""" ++ """A count rise seen in copy k is reported from copy k+1 ++ (maxima current by then).""" + device = torch.device("cpu") + guard_telemetry.reset_for_test() +- buf = guard_telemetry.counter(0, device) # allocates the cpu buffer ++ guard_telemetry.counter(0, device) # allocates the cpu buffer + full = guard_telemetry._buffers[device] + site = guard_telemetry.SITE_MAMBA_COPY +- full[site, 0] = 2 # count written, maxima not yet +- guard_telemetry._poll_once() # rise seen -> pending, nothing reported ++ full[site, 0] = 2 # count written, maxima not yet ++ guard_telemetry._poll_once() # rise seen -> pending, nothing reported + assert guard_telemetry._seen.get((device, site), 0) == 0 +- full[site, 1] = 9 + guard_telemetry.GUARD_BIAS # maxima land before the next tick +- guard_telemetry._poll_once() # confirming tick reports count 2 with row 9 ++ full[site, 1] = 9 + guard_telemetry.GUARD_BIAS # maxima land before the next tick ++ guard_telemetry._poll_once() # confirming tick reports count 2 with row 9 + assert guard_telemetry._seen.get((device, site), 0) == 2 + snap = guard_telemetry.snapshot() + key = f"{device}:{guard_telemetry.SITE_NAMES[site]}" +@@ -58,13 +66,16 @@ def test_report_waits_for_the_confirming_tick_cpu(): + + def test_module_reader_is_the_non_blocking_kind(): + assert isinstance(guard_telemetry._nb_reader, guard_telemetry.NonBlockingReader) +- src = open(guard_telemetry.__file__, encoding="utf-8").read() +- body = src[src.index("def _poll_once"):src.index("def _arm(")] +- assert ".tolist()" not in body.replace("host.tolist()", ""), "the cadence reader must not read the device buffer synchronously" ++ with open(guard_telemetry.__file__, encoding="utf-8") as fh: ++ src = fh.read() ++ body = src[src.index("def _poll_once") : src.index("def _arm(")] ++ assert ".tolist()" not in body.replace("host.tolist()", ""), ( ++ "the cadence reader must not read the device buffer synchronously" ++ ) + assert "poll_one(" in body + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + def test_reader_tick_returns_while_the_compute_stream_is_busy(): + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) +@@ -83,15 +94,23 @@ def test_reader_tick_returns_while_the_compute_stream_is_busy(): + t1 = time.monotonic() + torch.cuda.synchronize(device) + busy = time.monotonic() - t1 +- # The tick assertion comes first: a blocking reader spends the whole sleep inside the tick and would then make +- # the busy control read ~0 (negative control). +- assert tick < 0.25, f"reader tick took {tick:.3f}s while the compute stream was busy: it waited on the stream (reader/compute-stream deadlock)" +- assert tick + busy > 0.5, f"control failed: the parked kernel finished in {tick + busy:.3f}s, the test cannot discriminate" +- +- +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++ # The tick assertion comes first: a blocking reader spends the whole ++ # sleep inside the tick and would then make the busy control read ~0 ++ # (negative control). ++ assert tick < 0.25, ( ++ f"reader tick took {tick:.3f}s while the compute stream was busy: " ++ f"it waited on the stream (reader/compute-stream deadlock)" ++ ) ++ assert tick + busy > 0.5, ( ++ f"control failed: the parked kernel finished in {tick + busy:.3f}s, " ++ f"the test cannot discriminate" ++ ) ++ ++ ++@pytest.mark.gpu + def test_reader_tick_from_a_thread_returns_while_the_compute_stream_is_busy(): +- """The production shape: the tick runs on the reader thread while the main thread owns the stream.""" ++ """The production shape: the tick runs on the reader thread while the ++ main thread owns the stream.""" + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) + guard_telemetry.reset_for_test() +@@ -110,13 +129,17 @@ def test_reader_tick_from_a_thread_returns_while_the_compute_stream_is_busy(): + threading.Thread(target=tick, name="test-guard-reader", daemon=True).start() + finished = done.wait(timeout=1.0) + torch.cuda.synchronize(device) +- assert finished, "the reader tick did not return within 1 s while the compute stream was busy (reader/compute-stream deadlock)" ++ assert finished, ( ++ "the reader tick did not return within 1 s while the compute stream " ++ "was busy (reader/compute-stream deadlock)" ++ ) + assert elapsed["tick"] < 0.25 + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + def test_hits_written_before_a_stall_are_still_reported(): +- """Non-blocking must not mean lossy: a count written on the compute stream is reported once the copy lands.""" ++ """Non-blocking must not mean lossy: a count written on the compute ++ stream is reported once the copy lands.""" + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) + guard_telemetry.reset_for_test() +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +index 6799a80..606292e 100644 +--- a/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py ++++ b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +@@ -13,15 +13,12 @@ + """ + + import importlib +-import os + import sys + + import pytest + import torch + +-pytestmark = pytest.mark.skipif( +- not torch.cuda.is_available(), reason="CUDA required" +-) ++pytestmark = pytest.mark.gpu + + + def _fresh(monkeypatch, **env): +@@ -36,7 +33,7 @@ def _fresh(monkeypatch, **env): + def test_release_defaults_with_env_absent(monkeypatch): + gt = _fresh(monkeypatch) + assert gt._MODE == "off" # v2: default off +- assert gt._INTERVAL == pytest.approx(0.05) ++ assert pytest.approx(0.05) == gt._INTERVAL + + + @pytest.mark.parametrize("bad", ["nan", "inf", "-inf", "0", "-1", "abc"]) +@@ -46,7 +43,9 @@ def test_bad_interval_rejected_at_import(monkeypatch, bad): + + + def test_reader_loss_arms_the_stop(monkeypatch): +- gt = _fresh(monkeypatch, VLLM_E2_GUARD_MODE="warn") # the reader exists only when opted in (default off) ++ gt = _fresh( ++ monkeypatch, VLLM_E2_GUARD_MODE="warn" ++ ) # the reader exists only when opted in (default off) + device = torch.device("cuda:0") + gt.prealloc(device) + gt.reset_for_test() +@@ -64,7 +63,9 @@ def test_reader_loss_arms_the_stop(monkeypatch): + + + def test_drain_reports_last_step_hit_and_raises_for_fatal_site(monkeypatch): +- gt = _fresh(monkeypatch, VLLM_E2_GUARD_MODE="warn") # warn mode (opt-in since v2.0.1): a plain site must NOT raise ++ gt = _fresh( ++ monkeypatch, VLLM_E2_GUARD_MODE="warn" ++ ) # warn mode (opt-in): a plain site must NOT raise + device = torch.device("cuda:0") + gt.prealloc(device) + gt.reset_for_test() +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_replay.py b/tests/v1/worker/gpu/test_guard_telemetry_replay.py +index 52327ab..a987a02 100644 +--- a/tests/v1/worker/gpu/test_guard_telemetry_replay.py ++++ b/tests/v1/worker/gpu/test_guard_telemetry_replay.py +@@ -27,9 +27,7 @@ from vllm.third_party.flash_linear_attention.ops import ( + ) + from vllm.v1.worker import guard_telemetry + +-pytestmark = pytest.mark.skipif( +- not torch.cuda.is_available(), reason="CUDA required" +-) ++pytestmark = pytest.mark.gpu + + REPLAYS = 5 + +diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +index 30bed0a..df904ff 100644 +--- a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py ++++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +@@ -42,6 +42,9 @@ from vllm.v1.worker import mamba_utils as mamba_utils_mod + from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState + from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext + ++# Module-level singleton so the helper default is not a call (ruff B008). ++_CPU = torch.device("cpu") ++ + MAX_NUM_REQS = 8 + MAX_NUM_BLOCKS = 6 + MAMBA_BLOCK_SIZE = 16 +@@ -57,12 +60,8 @@ _COPY_FUNCS: MambaStateCopyFuncsByType = { + + # Captured before any monkeypatching so the replay binds against the real + # kernel signature; a launch that stops matching it fails loudly here. +-_POSTPROCESS_SIG = inspect.signature( +- mamba_utils_mod.postprocess_mamba_fused_kernel.fn +-) +-_PRECOPY_SIG = inspect.signature( +- mamba_utils_mod.precopy_mamba_align_fused_kernel.fn +-) ++_POSTPROCESS_SIG = inspect.signature(mamba_utils_mod.postprocess_mamba_fused_kernel.fn) ++_PRECOPY_SIG = inspect.signature(mamba_utils_mod.precopy_mamba_align_fused_kernel.fn) + + + class _MockCpuGpuBuffer: +@@ -163,7 +162,7 @@ def _register_block_tables(state: MambaHybridModelState, block_tables) -> bool: + + + def _build_scenario( +- device: torch.device = torch.device("cpu"), ++ device: torch.device = _CPU, + slot_mapping: list[int] = (0, 1, 2), + t1_mapping: list[int] = (2,), + register: bool = True, +@@ -291,7 +290,8 @@ def _replay_row(scenario, launch, batch_idx: int) -> tuple[str, int, dict]: + "src_eq_dst": src_block_idx == dest_block_idx, + } + +- if bt_row_is_req: ++ if bt_row_is_req: # noqa: SIM108 -- the ternary ruff proposes nests two ++ # conditionals on one line; the block states the three cases plainly. + bt_row_idx = req_idx + else: + bt_row_idx = batch_idx if has_idx_mapping else req_idx +@@ -304,9 +304,7 @@ def _replay_row(scenario, launch, batch_idx: int) -> tuple[str, int, dict]: + + def _block_ids(scenario, launch, batch_idx: int) -> dict: + name, bt_row_idx, predicate = _replay_row(scenario, launch, batch_idx) +- table = next( +- tbl for label, tbl in scenario.tables_by_ptr.values() if label == name +- ) ++ table = next(tbl for label, tbl in scenario.tables_by_ptr.values() if label == name) + return { + "table": name, + "row": bt_row_idx, +@@ -321,9 +319,7 @@ def _block_ids(scenario, launch, batch_idx: int) -> dict: + + def _run_deferred(scenario, monkeypatch: pytest.MonkeyPatch) -> _RecordingKernel: + recorder = _RecordingKernel() +- monkeypatch.setattr( +- mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder +- ) ++ monkeypatch.setattr(mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder) + device = scenario.persistent.device + + def t(values): +@@ -355,9 +351,9 @@ def _run_deferred(scenario, monkeypatch: pytest.MonkeyPatch) -> _RecordingKernel + + def test_gather_port_matches_production_semantics() -> None: + """Instrument check: the CPU gather port must reproduce block_table.py.""" +- persistent = torch.arange( +- MAX_NUM_REQS * MAX_NUM_BLOCKS, dtype=torch.int32 +- ).reshape(MAX_NUM_REQS, MAX_NUM_BLOCKS) ++ persistent = torch.arange(MAX_NUM_REQS * MAX_NUM_BLOCKS, dtype=torch.int32).reshape( ++ MAX_NUM_REQS, MAX_NUM_BLOCKS ++ ) + dst = torch.full_like(persistent, -7) + num_blocks = [2] * MAX_NUM_REQS + +@@ -439,9 +435,7 @@ def test_precopy_still_uses_the_same_step_gather( + """ + scenario = _build_scenario() + recorder = _RecordingKernel() +- monkeypatch.setattr( +- mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder +- ) ++ monkeypatch.setattr(mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder) + device = scenario.persistent.device + z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + +@@ -528,14 +522,10 @@ def test_align_postprocess_fails_closed_without_persistent_tables( + assert scenario.ctx.has_persistent_block_tables is False + + recorder = _RecordingKernel() +- monkeypatch.setattr( +- mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder +- ) ++ monkeypatch.setattr(mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder) + device = scenario.persistent.device + before = scenario.ctx.num_accepted_tokens_out.clone() +- num_accepted = torch.full( +- (MAX_NUM_REQS,), 7, dtype=torch.int32, device=device +- ) ++ num_accepted = torch.full((MAX_NUM_REQS,), 7, dtype=torch.int32, device=device) + + with pytest.raises(RuntimeError, match="persistent per-req-slot block tables"): + scenario.ctx.run_fused_postprocess_align( +@@ -564,9 +554,7 @@ def test_precopy_still_works_without_persistent_tables( + scenario = _build_scenario(register=False) + assert scenario.ctx.has_persistent_block_tables is False + recorder = _RecordingKernel() +- monkeypatch.setattr( +- mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder +- ) ++ monkeypatch.setattr(mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder) + device = scenario.persistent.device + z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + +diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py +index f036227..376494e 100644 +--- a/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py ++++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py +@@ -22,7 +22,8 @@ possible coordinates resolve to four different physical blocks: + gathered[batch_idx=0] first cut -> r2's blocks (the victim) + gathered[req_idx=5] wrong both -> never written, block 0 + +-Requires CUDA **and** an explicit opt-in: set ``VLLM_E2_GPU_TESTS=1``. The reference host ++Requires CUDA **and** an explicit opt-in: set ``VLLM_E2_GPU_TESTS=1``. The ++reference host + always reports CUDA available, and the reference host spends most of its time running + long serving runs and sanitizer runs that should not share a device with a + test suite, so `torch.cuda.is_available()` is not a sufficient gate here. Run +@@ -51,7 +52,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpe + from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext + + pytestmark = [ +- pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), ++ pytest.mark.gpu, + pytest.mark.skipif( + os.environ.get("VLLM_E2_GPU_TESTS") != "1", + reason="set VLLM_E2_GPU_TESTS=1 to run GPU tests (shared GPU host)", +@@ -162,9 +163,7 @@ def _build(slot_mapping=(5, 1, 2), t1_mapping=(2,)) -> _Fixture: + device = torch.device("cuda:0") + torch.manual_seed(1234) + +- conv = torch.randn( +- TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device +- ) ++ conv = torch.randn(TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device) + temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) + attention = MagicMock() + attention.kv_cache = [conv, temporal] +@@ -271,8 +270,7 @@ def _assert_bitwise(actual: torch.Tensor, expected: torch.Tensor, what: str) -> + lhs = actual.contiguous().view(torch.int16) + rhs = expected.contiguous().view(torch.int16) + assert torch.equal(lhs, rhs), ( +- f"{what}: {int((lhs != rhs).sum().item())} of {lhs.numel()} 16-bit " +- f"words differ" ++ f"{what}: {int((lhs != rhs).sum().item())} of {lhs.numel()} 16-bit words differ" + ) + + +@@ -365,9 +363,7 @@ def test_self_copy_exits_without_touching_state_and_resets_num_accepted(): + return. The conjunction itself is covered by the next test. + """ + fx = _build() +- num_accepted = _launch( +- fx, [5, 1, 2], {5: (2, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} +- ) ++ num_accepted = _launch(fx, [5, 1, 2], {5: (2, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)}) + + _assert_only_blocks_changed(fx, set()) + assert int(num_accepted[5].item()) == 1 +@@ -391,9 +387,7 @@ def test_same_column_with_positive_bias_copies_and_still_resets(): + positive bias moves bytes. + """ + fx = _build() +- num_accepted = _launch( +- fx, [5, 1, 2], {5: (4, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} +- ) ++ num_accepted = _launch(fx, [5, 1, 2], {5: (4, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)}) + + blk = phys(5, 1) # both the conv source and the destination + temporal_src = phys(5, 3) # src_col + bias = 3 +@@ -460,9 +454,7 @@ def test_generation_masked_reuse_is_skipped_end_to_end(): + mask_from_generation_change, + ) + +- mapping = mask_from_generation_change( +- idx_mapping=[5, 1, 2], reused_slot_rows=[1] +- ) ++ mapping = mask_from_generation_change(idx_mapping=[5, 1, 2], reused_slot_rows=[1]) + assert mapping == [5, -1, 2], mapping + + fx = _build(slot_mapping=(5, 1, 2)) +@@ -500,9 +492,7 @@ def test_same_step_identity_mapping_is_unaffected(): + def test_align_launch_raises_without_persistent_tables(): + """Fail closed on real hardware too, before any kernel is enqueued.""" + device = torch.device("cuda:0") +- conv = torch.randn( +- TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device +- ) ++ conv = torch.randn(TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device) + temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) + attention = MagicMock() + attention.kv_cache = [conv, temporal] +@@ -530,9 +520,7 @@ def test_align_launch_raises_without_persistent_tables(): + num_accepted_tokens_gpu=torch.full( + (MAX_NUM_REQS,), 4, dtype=torch.int32, device=device + ), +- state_idx_gpu=torch.zeros( +- MAX_NUM_REQS, dtype=torch.int32, device=device +- ), ++ state_idx_gpu=torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device), + new_num_computed_tokens_gpu=torch.full( + (MAX_NUM_REQS,), 48, dtype=torch.int32, device=device + ), +diff --git a/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +index 0468288..84f109d 100644 +--- a/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py ++++ b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +@@ -23,8 +23,8 @@ from vllm.v1.worker import guard_telemetry + COMPRESS_RATIO = 4 + HEAD_DIM = 128 + N_HEADS = 4 +-STATE_SIZE = 8 # ring capacity +-ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++STATE_SIZE = 8 # ring capacity ++ROW_WIDTH = 140 # head dim + the int64 position triple in the tail + MAX_MODEL_LEN = 4096 + NUM_BLOCKS = 8 + RING_CAPACITY = NUM_BLOCKS * STATE_SIZE +@@ -56,7 +56,9 @@ def _inputs(device, num_tokens=4): + MAX_MODEL_LEN, HEAD_DIM // 2, dtype=torch.bfloat16, device=device + ) + ones = torch.ones(HEAD_DIM, dtype=torch.bfloat16, device=device) +- q_out = torch.empty(num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ q_out = torch.empty( ++ num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device ++ ) + state_cache = torch.zeros( + NUM_BLOCKS, STATE_SIZE, 1, ROW_WIDTH, dtype=torch.bfloat16, device=device + ) +@@ -72,23 +74,45 @@ def _inputs(device, num_tokens=4): + max(num_tokens // COMPRESS_RATIO, 1), 2, dtype=torch.int32, device=device + ) + return dict( +- q=q, k=k, positions=positions, cos_sin_cache=cos_sin, +- q_norm_weight=ones, k_norm_weight=ones, eps=1e-6, q_out=q_out, +- state_cache=state_cache, state_slots=state_slots, +- state_block_table=state_table, query_start_loc=qsl, +- logical_positions=logical, compressed_cache=comp_cache, +- compressed_slots=comp_slots, k_work_metadata=kwb, ++ q=q, ++ k=k, ++ positions=positions, ++ cos_sin_cache=cos_sin, ++ q_norm_weight=ones, ++ k_norm_weight=ones, ++ eps=1e-6, ++ q_out=q_out, ++ state_cache=state_cache, ++ state_slots=state_slots, ++ state_block_table=state_table, ++ query_start_loc=qsl, ++ logical_positions=logical, ++ compressed_cache=comp_cache, ++ compressed_slots=comp_slots, ++ k_work_metadata=kwb, + ) + + + def _launch(kw): + qpi.qsa_pre_indexer( +- kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], +- kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], +- kw["state_cache"], kw["state_slots"], kw["state_block_table"], +- kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], +- kw["compressed_slots"], kw["k_work_metadata"], +- compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ kw["q"], ++ kw["k"], ++ kw["positions"], ++ kw["cos_sin_cache"], ++ kw["q_norm_weight"], ++ kw["k_norm_weight"], ++ kw["eps"], ++ kw["q_out"], ++ kw["state_cache"], ++ kw["state_slots"], ++ kw["state_block_table"], ++ kw["query_start_loc"], ++ kw["logical_positions"], ++ kw["compressed_cache"], ++ kw["compressed_slots"], ++ kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, ++ mrope_section=(11, 11, 10), + rope_pos_offset=HEAD_DIM, + ) + +@@ -121,20 +145,26 @@ def test_launcher_passes_the_ring_skip_counter(monkeypatch): + + assert len(rec.launches) == 1 + _grid, args, kwargs = rec.launches[0] +- kwargs = {k: v for k, v in kwargs.items() +- if k not in ("num_warps", "num_stages", "num_ctas", "grid")} ++ kwargs = { ++ k: v ++ for k, v in kwargs.items() ++ if k not in ("num_warps", "num_stages", "num_ctas", "grid") ++ } + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + skip_ctr = bound.arguments["ring_skip_ctr"] + assert skip_ctr.dtype == torch.int64 and skip_ctr.numel() == 3 +- assert skip_ctr.data_ptr() == guard_telemetry.counter( +- guard_telemetry.SITE_QSA_RING_SKIP, device +- ).data_ptr() ++ assert ( ++ skip_ctr.data_ptr() ++ == guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_RING_SKIP, device ++ ).data_ptr() ++ ) + for other in ("rope_guard_ctr", "guard_ctr"): + assert bound.arguments[other].data_ptr() != skip_ctr.data_ptr() + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + def test_negative_control_valid_slots_count_zero(): + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) +@@ -148,11 +178,11 @@ def test_negative_control_valid_slots_count_zero(): + for t in range(kw["k"].shape[0]): + row = kw["state_cache"][0, t, 0, :] + assert torch.equal(row[:HEAD_DIM], kw["k"][t]) +- tail = row[HEAD_DIM:HEAD_DIM + 12].contiguous().view(torch.int64) ++ tail = row[HEAD_DIM : HEAD_DIM + 12].contiguous().view(torch.int64) + assert tail.tolist() == [t, t, t] + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + @pytest.mark.parametrize( + "bad_index,bad_slot", + [(2, -1), (1, RING_CAPACITY), (3, RING_CAPACITY + 5)], +@@ -184,11 +214,11 @@ def test_positive_control_one_skipped_token_is_counted_and_located(bad_index, ba + assert torch.count_nonzero(row) == 0, "skipped row was written" + else: + assert torch.equal(row[:HEAD_DIM], kw["k"][t]) +- tail = row[HEAD_DIM:HEAD_DIM + 12].contiguous().view(torch.int64) ++ tail = row[HEAD_DIM : HEAD_DIM + 12].contiguous().view(torch.int64) + assert tail.tolist() == [t, t, t] + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + def test_positive_control_two_skipped_tokens_accumulate(): + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) +@@ -198,7 +228,9 @@ def test_positive_control_two_skipped_tokens_accumulate(): + kw["state_slots"][3] = RING_CAPACITY + 1 + _launch(kw) + torch.cuda.synchronize() +- n, row, val = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, device).tolist() ++ n, row, val = guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_RING_SKIP, device ++ ).tolist() + assert n == 2 +- assert row - guard_telemetry.GUARD_BIAS == 3 # max token index ++ assert row - guard_telemetry.GUARD_BIAS == 3 # max token index + assert val - guard_telemetry.GUARD_BIAS == RING_CAPACITY + 1 # max slot value +diff --git a/tests/v1/worker/gpu/test_qsa_rope_position_guard.py b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +index 9539aa7..0628dd2 100644 +--- a/tests/v1/worker/gpu/test_qsa_rope_position_guard.py ++++ b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +@@ -33,8 +33,8 @@ from vllm.v1.worker import guard_telemetry + COMPRESS_RATIO = 4 + HEAD_DIM = 128 + N_HEADS = 4 +-STATE_SIZE = 8 # ring capacity +-ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++STATE_SIZE = 8 # ring capacity ++ROW_WIDTH = 140 # head dim + the int64 position triple in the tail + MAX_MODEL_LEN = 4096 + NUM_BLOCKS = 8 + +@@ -65,7 +65,9 @@ def _inputs(device, num_tokens=4): + MAX_MODEL_LEN, HEAD_DIM // 2, dtype=torch.bfloat16, device=device + ) + ones = torch.ones(HEAD_DIM, dtype=torch.bfloat16, device=device) +- q_out = torch.empty(num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) ++ q_out = torch.empty( ++ num_tokens, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device ++ ) + state_cache = torch.zeros( + NUM_BLOCKS, STATE_SIZE, 1, ROW_WIDTH, dtype=torch.bfloat16, device=device + ) +@@ -81,12 +83,22 @@ def _inputs(device, num_tokens=4): + max(num_tokens // COMPRESS_RATIO, 1), 2, dtype=torch.int32, device=device + ) + return dict( +- q=q, k=k, positions=positions, cos_sin_cache=cos_sin, +- q_norm_weight=ones, k_norm_weight=ones, eps=1e-6, q_out=q_out, +- state_cache=state_cache, state_slots=state_slots, +- state_block_table=state_table, query_start_loc=qsl, +- logical_positions=logical, compressed_cache=comp_cache, +- compressed_slots=comp_slots, k_work_metadata=kwb, ++ q=q, ++ k=k, ++ positions=positions, ++ cos_sin_cache=cos_sin, ++ q_norm_weight=ones, ++ k_norm_weight=ones, ++ eps=1e-6, ++ q_out=q_out, ++ state_cache=state_cache, ++ state_slots=state_slots, ++ state_block_table=state_table, ++ query_start_loc=qsl, ++ logical_positions=logical, ++ compressed_cache=comp_cache, ++ compressed_slots=comp_slots, ++ k_work_metadata=kwb, + ) + + +@@ -104,21 +116,35 @@ def test_launcher_passes_the_real_cos_sin_row_count(monkeypatch): + monkeypatch.setattr(qpi, "_qsa_pre_indexer_kernel", rec) + kw = _inputs(device) + qpi.qsa_pre_indexer( +- kw.pop("q"), kw.pop("k"), kw.pop("positions"), kw.pop("cos_sin_cache"), +- kw.pop("q_norm_weight"), kw.pop("k_norm_weight"), kw.pop("eps"), +- kw.pop("q_out"), kw.pop("state_cache"), kw.pop("state_slots"), +- kw.pop("state_block_table"), kw.pop("query_start_loc"), +- kw.pop("logical_positions"), kw.pop("compressed_cache"), +- kw.pop("compressed_slots"), kw.pop("k_work_metadata"), +- compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ kw.pop("q"), ++ kw.pop("k"), ++ kw.pop("positions"), ++ kw.pop("cos_sin_cache"), ++ kw.pop("q_norm_weight"), ++ kw.pop("k_norm_weight"), ++ kw.pop("eps"), ++ kw.pop("q_out"), ++ kw.pop("state_cache"), ++ kw.pop("state_slots"), ++ kw.pop("state_block_table"), ++ kw.pop("query_start_loc"), ++ kw.pop("logical_positions"), ++ kw.pop("compressed_cache"), ++ kw.pop("compressed_slots"), ++ kw.pop("k_work_metadata"), ++ compress_ratio=COMPRESS_RATIO, ++ mrope_section=(11, 11, 10), + rope_pos_offset=HEAD_DIM, + ) + + assert len(rec.launches) == 1 + _grid, args, kwargs = rec.launches[0] + # Triton launch options are not kernel parameters. +- kwargs = {k: v for k, v in kwargs.items() +- if k not in ("num_warps", "num_stages", "num_ctas", "grid")} ++ kwargs = { ++ k: v ++ for k, v in kwargs.items() ++ if k not in ("num_warps", "num_stages", "num_ctas", "grid") ++ } + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + assert bound.arguments["cos_sin_rows"] == MAX_MODEL_LEN +@@ -134,12 +160,16 @@ def test_launcher_passes_the_real_cos_sin_row_count(monkeypatch): + req_ctr = bound.arguments["guard_ctr"] + for ctr in (rope_ctr, req_ctr): + assert ctr.dtype == torch.int64 and ctr.numel() == 3 +- assert rope_ctr.data_ptr() == guard_telemetry.counter( +- guard_telemetry.SITE_QSA_ROPE_POS, device +- ).data_ptr() +- assert req_ctr.data_ptr() == guard_telemetry.counter( +- guard_telemetry.SITE_QSA_PRE_INDEXER, device +- ).data_ptr() ++ assert ( ++ rope_ctr.data_ptr() ++ == guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, device).data_ptr() ++ ) ++ assert ( ++ req_ctr.data_ptr() ++ == guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_PRE_INDEXER, device ++ ).data_ptr() ++ ) + assert rope_ctr.data_ptr() != req_ctr.data_ptr() + + +@@ -151,13 +181,10 @@ def test_rope_site_aborts_the_stage_even_in_warn_mode(): + still there to do it again. warn is the right default for a guard whose + masked value dies with the step; it is the wrong default for this one. + """ +- assert ( +- guard_telemetry.SITE_QSA_ROPE_POS +- in guard_telemetry._ALWAYS_FATAL_SITES +- ) ++ assert guard_telemetry.SITE_QSA_ROPE_POS in guard_telemetry._ALWAYS_FATAL_SITES + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + def test_ring_filler_bytes_are_masked_and_counted(): + """Positive control: the ring-residue crash case, verbatim. + +@@ -196,13 +223,17 @@ def test_ring_filler_bytes_are_masked_and_counted(): + # start must be unaligned: 6 gives first_position 4 < chunk_start 6. + start = 6 + assert start % COMPRESS_RATIO != 0, "an aligned chunk start never crosses" +- kw["logical_positions"] = ( +- torch.arange(start, start + kw["logical_positions"].numel(), +- dtype=kw["logical_positions"].dtype, device=device) ++ kw["logical_positions"] = torch.arange( ++ start, ++ start + kw["logical_positions"].numel(), ++ dtype=kw["logical_positions"].dtype, ++ device=device, + ) + chunk_end = start + kw["logical_positions"].numel() - 1 + chunk_start = chunk_end - kw["logical_positions"].numel() + 1 +- first_boundary = ((chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO) * COMPRESS_RATIO - 1 ++ first_boundary = ( ++ (chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO ++ ) * COMPRESS_RATIO - 1 + first_position = first_boundary - (COMPRESS_RATIO - 1) + assert first_position < chunk_start, ( + f"first_position {first_position} >= chunk_start {chunk_start}: this " +@@ -216,15 +247,27 @@ def test_ring_filler_bytes_are_masked_and_counted(): + kw["k_work_metadata"][:, 1] = 0 + + qpi.qsa_pre_indexer( +- kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], +- kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], +- kw["state_cache"], kw["state_slots"], kw["state_block_table"], +- kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], +- kw["compressed_slots"], kw["k_work_metadata"], +- compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ kw["q"], ++ kw["k"], ++ kw["positions"], ++ kw["cos_sin_cache"], ++ kw["q_norm_weight"], ++ kw["k_norm_weight"], ++ kw["eps"], ++ kw["q_out"], ++ kw["state_cache"], ++ kw["state_slots"], ++ kw["state_block_table"], ++ kw["query_start_loc"], ++ kw["logical_positions"], ++ kw["compressed_cache"], ++ kw["compressed_slots"], ++ kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, ++ mrope_section=(11, 11, 10), + rope_pos_offset=HEAD_DIM, + ) +- torch.cuda.synchronize() # the pre-guard build faults here ++ torch.cuda.synchronize() # the pre-guard build faults here + + ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, device) + n_masked = int(ctr[0].item()) +@@ -239,7 +282,7 @@ def test_ring_filler_bytes_are_masked_and_counted(): + ) + + +-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") ++@pytest.mark.gpu + @pytest.mark.parametrize("axis,axis_id", [("t", 0), ("h", 1), ("w", 2)]) + @pytest.mark.parametrize("sign", ["mixed", "negative_only"]) + def test_a_single_bad_axis_is_caught_and_named(axis, axis_id, sign): +@@ -264,7 +307,7 @@ def test_a_single_bad_axis_is_caught_and_named(axis, axis_id, sign): + kw["positions"] = kw["positions"].clone() + if sign == "mixed": + kw["positions"][row, 0] = MAX_MODEL_LEN + 7 # out of range, positive +- kw["positions"][row, 1] = -5 # out of range, negative ++ kw["positions"][row, 1] = -5 # out of range, negative + expected_worst = MAX_MODEL_LEN + 7 + else: + # NEGATIVE ONLY. Every offender below zero is the case +@@ -279,12 +322,24 @@ def test_a_single_bad_axis_is_caught_and_named(axis, axis_id, sign): + from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import qsa_pre_indexer + + qsa_pre_indexer( +- kw["q"], kw["k"], kw["positions"], kw["cos_sin_cache"], +- kw["q_norm_weight"], kw["k_norm_weight"], kw["eps"], kw["q_out"], +- kw["state_cache"], kw["state_slots"], kw["state_block_table"], +- kw["query_start_loc"], kw["logical_positions"], kw["compressed_cache"], +- kw["compressed_slots"], kw["k_work_metadata"], +- compress_ratio=COMPRESS_RATIO, mrope_section=(11, 11, 10), ++ kw["q"], ++ kw["k"], ++ kw["positions"], ++ kw["cos_sin_cache"], ++ kw["q_norm_weight"], ++ kw["k_norm_weight"], ++ kw["eps"], ++ kw["q_out"], ++ kw["state_cache"], ++ kw["state_slots"], ++ kw["state_block_table"], ++ kw["query_start_loc"], ++ kw["logical_positions"], ++ kw["compressed_cache"], ++ kw["compressed_slots"], ++ kw["k_work_metadata"], ++ compress_ratio=COMPRESS_RATIO, ++ mrope_section=(11, 11, 10), + rope_pos_offset=HEAD_DIM, + ) + torch.cuda.synchronize() +diff --git a/tests/v1/worker/test_grammar_fail_closed.py b/tests/v1/worker/test_grammar_fail_closed.py +index 869cb98..eeb6839 100644 +--- a/tests/v1/worker/test_grammar_fail_closed.py ++++ b/tests/v1/worker/test_grammar_fail_closed.py +@@ -19,7 +19,6 @@ import numpy as np + import pytest + import torch + +-from vllm.platforms import current_platform + from vllm.utils.math_utils import cdiv + from vllm.v1.core.sched.output import GrammarOutput + from vllm.v1.worker.gpu.model_runner import GPUModelRunner +@@ -109,7 +108,7 @@ def _rejection_sampler(device: torch.device) -> RejectionSampler: + return sampler + + +-@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") ++@pytest.mark.gpu + def test_acceptance_stops_at_the_last_masked_row(): + device = torch.device("cuda") + input_batch = _input_batch(device) +diff --git a/vllm/envs.py b/vllm/envs.py +index 157dfaf..588ca59 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -304,6 +304,132 @@ if TYPE_CHECKING: + # AOT cache on it; the layer itself still reads os.environ directly. + VLLM_HOST_EMBED_TABLE: bool = False + VLLM_PLE_OFFLOAD_READY_TIMEOUT: float = 600.0 ++ # --- Flash-Next fork knobs (v2.2) -------------------------------- ++ # Every name below was read straight from os.environ at its use site, so ++ # envs.validate_environ() logged it as unknown on every boot that set one ++ # and compile_factors() could not see it. Registering makes the hash ++ # complete; the read sites are unchanged (they still parse and fail closed ++ # themselves, which registration must not weaken). Runtime-only entries are ++ # listed in compile_factors().ignored_factors with their reason. ++ # ++ # TWO IDENTITIES, NOT ONE. "Read once at startup" ++ # settles neither of them, so every entry below is labelled with both: ++ # ++ # AOT identity -- does the value change the torch.compile / AOT ++ # artifact? This is the ONLY identity compile_factors() ++ # keys. It protects against reusing another process's ++ # cached artifact. ++ # CAPTURE identity -- does the value change what a cudagraph capture bakes ++ # in (ops present, launch constants, shapes)? Nothing ++ # in this tree keys anything on capture identity: the ++ # env is read before capture, in-process, so capture ++ # identity is protected only by "do not change the ++ # value between capture and replay in one process". ++ # ++ # So a knob that is a capture factor but NOT an AOT factor gets no ++ # protection from being registered here. Registering it is a contract ++ # statement and a cross-process guard, not a mechanism. Where the two ++ # identities are not both established by reading the code, the knob is ++ # hashed anyway: a spurious cache miss costs a recompile, a missed one costs ++ # a wrong graph. ++ # ++ # COMPILE/CAPTURE FACTORS (hashed): ++ # AOT: yes. CAPTURE: yes. PLE table is FP8-with-global-scale -> ++ # Qwen4ExpPLEFp8EmbeddingMethod instead of the checkpoint's method: ++ # different weight dtype and a different embedding forward. ++ VLLM_PLE_FP8_GLOBAL_SCALE: bool = False ++ # CAPTURE: yes, established -- clip counting allocates persistent per-layer ++ # int64 counters and its increment ops are CAPTURED into the cudagraph (see ++ # the graph-safety note in models/qwen4_exp/nvidia/qsa.py); unset, those ops ++ # are absent. AOT: not established (the QSA forward is a custom op); hashed ++ # conservatively. ++ VLLM_QSA_KV_CLIP_COUNT: str = "" ++ # CAPTURE: n/a -- collection requires --enforce-eager, so an enabled run ++ # does not capture at all. AOT: adds absmax reduction ops to the attention ++ # forward; hashed conservatively. ++ VLLM_QSA_KV_COLLECT: str = "" ++ # CAPTURE: yes -- the resolved value is a launch constant baked into the ++ # graph at capture (csrc/libtorch_stable/persistent_topk.cuh:143-148). ++ # AOT: not established. Hashed conservatively. ++ # NOT A PYTHON READ SITE: the value is resolved once per process by ++ # exact_medium_enabled() in csrc/libtorch_stable/topk.cu:30-36, which is the ++ # single source of truth. This entry is a raw passthrough so nothing can ++ # mistake it for a second parser; do not use envs.MERLIN_TOPK_EXACT_MEDIUM ++ # to decide behaviour. Registered with the medium-path exact restart ++ # so the cache-factor contract is complete. ++ MERLIN_TOPK_EXACT_MEDIUM: str | None = None ++ # ++ # RUNTIME-ONLY (registered for validate_environ + a single cached read): ++ # Static K/V scale sidecar: fill_()ed into existing fp32 device buffers ++ # after load and before capture, and the kernel does tl.load(k_scale_ptr), ++ # so the value is never a constexpr and the graph is identical. ++ VLLM_QSA_KV_SCALES: str = "" ++ # Guard telemetry: mode raise|warn|off and the reader poll interval in ++ # seconds. counter() returns a real tensor in every mode, so the Triton ++ # arguments and the captured graph do not depend on either. ++ VLLM_E2_GUARD_MODE: str = "off" ++ VLLM_E2_GUARD_INTERVAL: str = "0.05" ++ # Debug allocator telemetry (poll seconds + output directory) and the ++ # scheduler allocation log path. ++ VLLM_E2_MEMSTATS: str = "0" ++ VLLM_E2_MEMSTATS_DIR: str = "" ++ VLLM_E2_SCHED_LOG: str = "" ++ # O(live blocks) host-side ownership assertion in the scheduler, which is ++ # plain Python and never compiled. ++ VLLM_E2_ZERO_OWNERSHIP_CHECK: bool = False ++ # PLE offload child placement: rewrites CUDA_VISIBLE_DEVICES for the child ++ # only (CUDA_VISIBLE_DEVICES is itself an ignored factor). ++ VLLM_PLE_OFFLOAD_HOME_DEVICE: str = "" ++ # PLE timeouts / deadlines, in seconds. ++ VLLM_PLE_OFFLOAD_WAIT_TIMEOUT: float = 120.0 ++ VLLM_PLE_PROGRESS_DEADLINE_S: float = 300.0 ++ VLLM_PLE_STARTUP_DEADLINE_S: float = 1800.0 ++ # Test-only PLE fault injection: cut spec, marker log, bootstrap gate. ++ # AOT: no. CAPTURE: no. Read once at import by the vllm_ple_test_fault ++ # singleton, but the process that reads them is a real engine, so ++ # validate_environ sees them. ++ VLLM_PLE_TEST_FAULT: str | None = None ++ VLLM_PLE_TEST_FAULT_GATE: str | None = None ++ VLLM_PLE_TEST_FAULT_LOG: str | None = None ++ # ++ # --- Knobs registered ahead of their read sites (v2.2) ----------- ++ # Registered ahead of their code so the MERGED tree is complete and ++ # validate_environ is quiet from the first boot (a ++ # cross-change knob table). Each mirrors its read site's ++ # parse; if a read site changes, the entry here must be reconciled. ++ # (Both read sites below are now live.) ++ # AOT: no. CAPTURE: no -- strictly host-side and outside capture. ++ # off|warn|raise (1 = warn), normalised by _e2_diagnostic_mode like the ++ # rest of the family. The mamba ownership check moved this gate off a plain ++ # boolean; reconciled here. ++ # Read site returns with the deferred PLE transport. ++ VLLM_E2_PLE_FLAG_CHECK: str = "off" ++ # AOT: no. CAPTURE: no -- strictly host-side. off|warn|raise (1 = warn). ++ # the mamba slot-ownership check, v1/worker/mamba_utils.py:74. ++ VLLM_E2_MAMBA_SLOT_OWNERSHIP: str = "off" ++ # ++ # --- TODO: knobs that are NOT in this tree at all ---------- ++ # The tie-determinism toggles and the BSM pin live in a separate ++ # patch stack, not in v2.0.1 and not in this tree: ++ # ++ # MERLIN_FULL_K MERLIN_QSA_SORT MERLIN_TIE_DET MERLIN_TIE_RECENT ++ # MERLIN_PIN_BSM ++ # ++ # They are therefore UNAUDITED: nothing in this repository shows what they ++ # change, and this file will not classify code it cannot read. The rule they ++ # fall under is the standing one -- an os.environ (or std::getenv) knob that ++ # changes a compiled or captured artifact is not in the cache key unless it ++ # is registered here. From what is known of them (Marlin split-K, a ++ # persistent_topk ordering change and a cutoff tie-break) they are very ++ # likely capture factors and possibly AOT factors, which is exactly why a ++ # guess is not good enough. ++ # ++ # WHEN THE PATCHES LAND, for each name: register it below, label both ++ # identities from its actual read site, hash it unless exclusion is proven, ++ # and delete it from this block. tests/test_envs.py holds a test that fails ++ # if a name is registered while still listed here, so the classification ++ # cannot be skipped silently. ++ # ---------------------------------------------------------------------- + VLLM_LOG_MODEL_INSPECTION: bool = False + VLLM_DEBUG_MFU_METRICS: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False +@@ -2072,6 +2198,43 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_PLE_OFFLOAD_READY_TIMEOUT": lambda: float( + os.getenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "600") + ), ++ # --- Flash-Next fork knobs (v2.2); see the TYPE_CHECKING block for ++ # which of these are compile/capture factors and which are ignored. ++ # Compile/capture factors. ++ "VLLM_PLE_FP8_GLOBAL_SCALE": lambda: ( ++ os.getenv("VLLM_PLE_FP8_GLOBAL_SCALE", "") == "1" ++ ), ++ "VLLM_QSA_KV_CLIP_COUNT": lambda: os.getenv("VLLM_QSA_KV_CLIP_COUNT", "").strip(), ++ "VLLM_QSA_KV_COLLECT": lambda: os.getenv("VLLM_QSA_KV_COLLECT", "").strip(), ++ # Capture-configuration factor, read in C++; raw passthrough on purpose. ++ "MERLIN_TOPK_EXACT_MEDIUM": lambda: os.getenv("MERLIN_TOPK_EXACT_MEDIUM", None), ++ # Runtime only; all of the following are in compile_factors.ignored_factors. ++ "VLLM_QSA_KV_SCALES": lambda: os.getenv("VLLM_QSA_KV_SCALES", "").strip(), ++ "VLLM_E2_GUARD_MODE": lambda: ( ++ os.getenv("VLLM_E2_GUARD_MODE", "off").strip().lower() ++ ), ++ "VLLM_E2_GUARD_INTERVAL": lambda: os.getenv("VLLM_E2_GUARD_INTERVAL", "0.05"), ++ "VLLM_E2_MEMSTATS": lambda: os.getenv("VLLM_E2_MEMSTATS", "0"), ++ "VLLM_E2_MEMSTATS_DIR": lambda: os.getenv("VLLM_E2_MEMSTATS_DIR", ""), ++ "VLLM_E2_SCHED_LOG": lambda: os.getenv("VLLM_E2_SCHED_LOG", ""), ++ "VLLM_E2_ZERO_OWNERSHIP_CHECK": lambda: ( ++ os.getenv("VLLM_E2_ZERO_OWNERSHIP_CHECK", "") == "1" ++ ), ++ "VLLM_PLE_OFFLOAD_HOME_DEVICE": lambda: os.getenv( ++ "VLLM_PLE_OFFLOAD_HOME_DEVICE", "" ++ ).strip(), ++ "VLLM_PLE_OFFLOAD_WAIT_TIMEOUT": lambda: float( ++ os.getenv("VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", "120") ++ ), ++ "VLLM_PLE_PROGRESS_DEADLINE_S": lambda: float( ++ os.getenv("VLLM_PLE_PROGRESS_DEADLINE_S", "300") ++ ), ++ "VLLM_PLE_STARTUP_DEADLINE_S": lambda: float( ++ os.getenv("VLLM_PLE_STARTUP_DEADLINE_S", "1800") ++ ), ++ "VLLM_PLE_TEST_FAULT": lambda: os.getenv("VLLM_PLE_TEST_FAULT", None), ++ "VLLM_PLE_TEST_FAULT_GATE": lambda: os.getenv("VLLM_PLE_TEST_FAULT_GATE", None), ++ "VLLM_PLE_TEST_FAULT_LOG": lambda: os.getenv("VLLM_PLE_TEST_FAULT_LOG", None), + # Log model inspection after loading. + # If enabled, logs a transformers-style hierarchical view of the model + # with quantization methods and attention backends. +@@ -2314,13 +2477,6 @@ def compile_factors() -> dict[str, object]: + "VLLM_DEBUG_DUMP_PATH", + "VLLM_PORT", + "VLLM_CACHE_ROOT", +- # PLE host-side diagnostics: both gate Python that runs outside any +- # capture and neither adds a node to a compiled graph, so hashing them +- # would invalidate every compile cache the moment a diagnostic boot is +- # taken -- which is exactly when the cache is wanted. Remove from this +- # set if either ever grows a device-side component. +- "VLLM_E2_PLE_FLAG_CHECK", +- "VLLM_E2_MAMBA_SLOT_OWNERSHIP", + # Runtime memory-plan persistence; does not affect compiled graphs. + "VLLM_ENABLE_STARTUP_PLAN", + # Location-only derived paths: where a cache/config directory lives +@@ -2397,6 +2553,55 @@ def compile_factors() -> dict[str, object]: + "LOCAL_RANK", + "CUDA_VISIBLE_DEVICES", + "NO_COLOR", ++ # --- Flash-Next fork knobs that are runtime-only (v2.2) ------ ++ # Registering a name in environment_variables makes it a compile ++ # factor by default (this function is a deny-list), so every fork knob ++ # that cannot change a compiled or captured artifact is listed here ++ # with its reason. Hashing them would invalidate the whole AOT cache on ++ # a log path, a timeout or a guard-mode flip. ++ # ++ # Static QSA K/V scales: fill_()ed into existing fp32 device buffers ++ # after load and before capture; the kernel reads them with ++ # tl.load(k_scale_ptr), never as a constexpr. ++ "VLLM_QSA_KV_SCALES", ++ # Guard telemetry: counter() returns a real [3] int64 view in every ++ # mode, so the Triton arguments and the captured graph are identical; ++ # the mode only gates the side reader thread and the check_pending ++ # raise, and the interval is that thread's poll cadence. ++ "VLLM_E2_GUARD_MODE", ++ "VLLM_E2_GUARD_INTERVAL", ++ # Debug telemetry cadence and output locations. ++ "VLLM_E2_MEMSTATS", ++ "VLLM_E2_MEMSTATS_DIR", ++ "VLLM_E2_SCHED_LOG", ++ # Host-side scheduler assertion; the scheduler is never compiled. ++ "VLLM_E2_ZERO_OWNERSHIP_CHECK", ++ # Placement only: rewrites CUDA_VISIBLE_DEVICES for the PLE offload ++ # child, and CUDA_VISIBLE_DEVICES is already ignored above. ++ "VLLM_PLE_OFFLOAD_HOME_DEVICE", ++ # Timeouts and deadlines, in seconds. VLLM_PLE_OFFLOAD_READY_TIMEOUT ++ # was registered in the v2 host-embed commit and left hashed by ++ # oversight: serve-v2.sh sets it to 1800, so every boot that changed it ++ # missed the AOT cache for no reason. ++ "VLLM_PLE_OFFLOAD_READY_TIMEOUT", ++ "VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", ++ "VLLM_PLE_PROGRESS_DEADLINE_S", ++ "VLLM_PLE_STARTUP_DEADLINE_S", ++ # Test-only fault injection: a cut spec, a marker log path and a ++ # bootstrap gate file. None of them reaches a kernel or a graph. ++ "VLLM_PLE_TEST_FAULT", ++ "VLLM_PLE_TEST_FAULT_GATE", ++ "VLLM_PLE_TEST_FAULT_LOG", ++ # Diagnostic knobs, both strictly host-side and outside capture, so ++ # neither AOT nor capture identity depends on them. Hashing ++ # them would invalidate every compile cache the moment a diagnostic ++ # boot is taken -- exactly when the cache is wanted. Remove from this ++ # set if either ever grows a device-side component. ++ # MERLIN_TOPK_EXACT_MEDIUM is deliberately NOT here: it is a capture ++ # configuration factor and is hashed conservatively. ++ "VLLM_E2_PLE_FLAG_CHECK", ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP", ++ # ------------------------------------------------------------------ + } + + from vllm.config.utils import normalize_value +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index adb057d..91f34de 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -162,7 +162,10 @@ if _qsa_clip_raw and not (math.isfinite(_qsa_clip_interval) and _qsa_clip_interv + _qsa_clip_counters: dict[str, torch.Tensor] = {} + # v2 reader fix: the reader must never block on the compute + # stream; see guard_telemetry.NonBlockingReader. +-from vllm.v1.worker.guard_telemetry import NonBlockingReader as _NonBlockingReader # noqa: E402 ++from vllm.v1.worker.guard_telemetry import ( # noqa: E402 ++ NonBlockingReader as _NonBlockingReader, ++) ++ + _qsa_clip_nb = _NonBlockingReader() + _qsa_clip_reader_started = False + _qsa_clip_lock = threading.Lock() +@@ -217,7 +220,7 @@ def _qsa_clip_reader() -> None: + + + def _qsa_clip_count( +- layer: "Qwen4ExpQSAAttention", ++ layer: Qwen4ExpQSAAttention, + key: torch.Tensor, + value: torch.Tensor, + slot_mapping: torch.Tensor | None = None, +diff --git a/vllm/v1/core/sched/e2_alloc_log.py b/vllm/v1/core/sched/e2_alloc_log.py +index 0efcbd7..6fa2354 100644 +--- a/vllm/v1/core/sched/e2_alloc_log.py ++++ b/vllm/v1/core/sched/e2_alloc_log.py +@@ -1,23 +1,45 @@ + # SPDX-License-Identifier: Apache-2.0 +-"""Debug-only scheduler allocation log for KV capacity probes (off unless VLLM_E2_SCHED_LOG=). ++"""Debug-only scheduler allocation log for KV-capacity probes (off unless ++VLLM_E2_SCHED_LOG=). + +-One JSON line per scheduler update while any request exists, and one "fin" line per request as it is freed: +- {"t": wall, "m": monotonic, "step": n, "total": pool blocks, "free": free blocks, "running": r, "waiting": w, ++One JSON line per scheduler update while any request exists, and one "fin" line ++per request as it is freed: ++ {"t": wall, "m": monotonic, "step": n, "total": pool blocks, ++ "free": free blocks, "running": r, "waiting": w, + "tot": tokens scheduled this step, +- "reqs": [[id, prompt, computed, output, cached, preemptions, spec_scheduled, status, [blocks per group], ++ "reqs": [[id, prompt, computed, output, cached, preemptions, ++ spec_scheduled, status, [blocks per group], + tokens scheduled for this request this step]]} +- {"t": wall, "m": monotonic, "fin": [id, prompt, computed, output, cached, preemptions, status, [blocks per group]]} +-Blocks are counted per KV-cache group, null blocks excluded. A logging failure writes an "error" line and never +-raises into the scheduler. Unset = the scheduler only pays one env check per update. ++ {"t": wall, "m": monotonic, ++ "fin": [id, prompt, computed, output, cached, preemptions, status, ++ [blocks per group]]} ++Blocks are counted per KV-cache group, null blocks excluded. A logging failure ++writes an "error" line and never raises into the scheduler. Unset = the ++scheduler only pays one env check per update. + """ ++ + import json + import os + import time + + _PATH = os.environ.get("VLLM_E2_SCHED_LOG", "") +-_state = {"f": None, "step": 0, "flush": 0.0, "wrapped": False, "errors": 0, "checked": False} +-REQ_FIELDS = ("request_id", "num_prompt_tokens", "num_computed_tokens", "num_output_tokens", "num_preemptions", +- "status", "prefill_stats") ++_state = { ++ "f": None, ++ "step": 0, ++ "flush": 0.0, ++ "wrapped": False, ++ "errors": 0, ++ "checked": False, ++} ++REQ_FIELDS = ( ++ "request_id", ++ "num_prompt_tokens", ++ "num_computed_tokens", ++ "num_output_tokens", ++ "num_preemptions", ++ "status", ++ "prefill_stats", ++) + + + def enabled() -> bool: +@@ -25,8 +47,10 @@ def enabled() -> bool: + + + def _blocks(sched, rid): +- return [sum(1 for b in m.req_to_blocks.get(rid, ()) if not b.is_null) +- for m in sched.kv_cache_manager.coordinator.single_type_managers] ++ return [ ++ sum(1 for b in m.req_to_blocks.get(rid, ()) if not b.is_null) ++ for m in sched.kv_cache_manager.coordinator.single_type_managers ++ ] + + + def _status(r): +@@ -34,28 +58,41 @@ def _status(r): + + + def _cached(r): +- """Prefix-cache hit for this request. vLLM keeps it in `request.prefill_stats` (PrefillStats.num_cached_tokens = +- local + external); the output processor takes those stats after the first prefill (`take_prefill_stats` sets the +- attribute to None), so -1 means "no longer knowable from the scheduler", never a silent zero.""" ++ """Prefix-cache hit for this request. vLLM keeps it in ++ `request.prefill_stats` (PrefillStats.num_cached_tokens = local + ++ external); the output processor takes those stats after the first prefill ++ (`take_prefill_stats` sets the attribute to None), so -1 means "no longer ++ knowable from the scheduler", never a silent zero.""" + ps = getattr(r, "prefill_stats", None) + return int(getattr(ps, "num_cached_tokens", -1)) if ps is not None else -1 + + + def _contract(r): +- """Once per process: every field this log reads must exist on the real Request object. A missing field is written +- once as a schema_error record — the evaluator fails closed on it — instead of turning every row into an error.""" ++ """Once per process: every field this log reads must exist on the real Request ++ object. A missing field is written once as a schema_error record — the evaluator ++ fails closed on it — instead of turning every row into an error.""" + if _state["checked"]: + return + _state["checked"] = True + missing = [f for f in REQ_FIELDS if not hasattr(r, f)] + if missing: +- _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "schema_error": missing}, force=True) ++ _write( ++ { ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "schema_error": missing, ++ }, ++ force=True, ++ ) + + + def _write(rec, force=False): + f = _state["f"] + if f is None: +- f = _state["f"] = open(_PATH, "a", buffering=1 << 16) ++ # noqa SIM115: one append handle for the process lifetime, kept in ++ # _state and flushed on a cadence; a context manager would reopen the ++ # log on every scheduler update. ++ f = _state["f"] = open(_PATH, "a", buffering=1 << 16) # noqa: SIM115 + f.write(json.dumps(rec, separators=(",", ":")) + "\n") + if force or rec["t"] - _state["flush"] > 0.5: + f.flush() +@@ -65,10 +102,19 @@ def _write(rec, force=False): + def _error(e): + _state["errors"] += 1 + if _state["errors"] <= 20: +- try: +- _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), +- "error": f"{type(e).__name__}: {e}"[:300]}, force=True) ++ try: # noqa: SIM105 ++ _write( ++ { ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "error": f"{type(e).__name__}: {e}"[:300], ++ }, ++ force=True, ++ ) + except Exception: # noqa: BLE001 ++ # SIM105: this module imports json, os and time only, so that an ++ # enabled log costs the scheduler nothing it can see; the ++ # last-resort writer does not earn a contextlib import. + pass + + +@@ -77,10 +123,23 @@ def _wrap_free(sched): + + def _free_request(request, *args, **kwargs): + try: +- _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), +- "fin": [request.request_id, request.num_prompt_tokens, request.num_computed_tokens, +- request.num_output_tokens, _cached(request), request.num_preemptions, +- _status(request), _blocks(sched, request.request_id)]}, force=True) ++ _write( ++ { ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "fin": [ ++ request.request_id, ++ request.num_prompt_tokens, ++ request.num_computed_tokens, ++ request.num_output_tokens, ++ _cached(request), ++ request.num_preemptions, ++ _status(request), ++ _blocks(sched, request.request_id), ++ ], ++ }, ++ force=True, ++ ) + except Exception as e: # noqa: BLE001 - the debug log must never break the scheduler + _error(e) + return orig(request, *args, **kwargs) +@@ -100,12 +159,35 @@ def log_step(sched, scheduler_output) -> None: + pool = sched.kv_cache_manager.block_pool + nsched = getattr(scheduler_output, "num_scheduled_tokens", None) or {} + _contract(next(iter(sched.requests.values()))) +- rows = [[rid, r.num_prompt_tokens, r.num_computed_tokens, r.num_output_tokens, _cached(r), +- r.num_preemptions, len(spec.get(rid, ())), _status(r), _blocks(sched, rid), int(nsched.get(rid, 0))] +- for rid, r in sched.requests.items()] +- _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "step": _state["step"], +- "total": pool.num_gpu_blocks, "free": pool.get_num_free_blocks(), +- "running": len(sched.running), "waiting": len(sched.waiting), +- "tot": int(getattr(scheduler_output, "total_num_scheduled_tokens", 0) or 0), "reqs": rows}) ++ rows = [ ++ [ ++ rid, ++ r.num_prompt_tokens, ++ r.num_computed_tokens, ++ r.num_output_tokens, ++ _cached(r), ++ r.num_preemptions, ++ len(spec.get(rid, ())), ++ _status(r), ++ _blocks(sched, rid), ++ int(nsched.get(rid, 0)), ++ ] ++ for rid, r in sched.requests.items() ++ ] ++ _write( ++ { ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "step": _state["step"], ++ "total": pool.num_gpu_blocks, ++ "free": pool.get_num_free_blocks(), ++ "running": len(sched.running), ++ "waiting": len(sched.waiting), ++ "tot": int( ++ getattr(scheduler_output, "total_num_scheduled_tokens", 0) or 0 ++ ), ++ "reqs": rows, ++ } ++ ) + except Exception as e: # noqa: BLE001 + _error(e) +diff --git a/vllm/v1/ple_offload/test_fault.py b/vllm/v1/ple_offload/test_fault.py +index 3546a1f..9ae74f7 100644 +--- a/vllm/v1/ple_offload/test_fault.py ++++ b/vllm/v1/ple_offload/test_fault.py +@@ -2,8 +2,12 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Compatibility exports for the top-level test-only fault singleton.""" + ++from vllm_ple_test_fault import ( ++ _FAULT as _FAULT, ++) + from vllm_ple_test_fault import ( # noqa: F401 + FaultSpec as FaultSpec, +- _FAULT as _FAULT, ++) ++from vllm_ple_test_fault import ( + _parse_fault_spec as _parse_fault_spec, + ) +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index e95653e..adca930 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -174,9 +174,7 @@ SITE_NAMES = { + _MODE = os.environ.get("VLLM_E2_GUARD_MODE", "off").strip().lower() + if _MODE not in ("raise", "warn", "off"): + # fail closed: a set-but-broken gate must not silently disarm the instrument +- raise ValueError( +- f"VLLM_E2_GUARD_MODE must be raise|warn|off, got {_MODE!r}" +- ) ++ raise ValueError(f"VLLM_E2_GUARD_MODE must be raise|warn|off, got {_MODE!r}") + + + def _parse_interval(text: str) -> float: +@@ -241,7 +239,7 @@ class NonBlockingReader: + """ + + def __init__(self) -> None: +- self._streams: dict[torch.device, "torch.cuda.Stream"] = {} ++ self._streams: dict[torch.device, torch.cuda.Stream] = {} + # key -> [device_tensor, pinned_host_tensor | None, event | None, inflight] + self._jobs: dict[object, list] = {} + self._lock = threading.Lock() +@@ -266,7 +264,9 @@ class NonBlockingReader: + if dev not in self._streams: + with torch.cuda.device(dev): + self._streams[dev] = torch.cuda.Stream(device=dev) +- host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True) ++ host = torch.empty( ++ tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True ++ ) + with torch.cuda.device(dev): + event = torch.cuda.Event() + self._jobs[key] = [tensor, host, event, False] +@@ -333,7 +333,8 @@ class NonBlockingReader: + + + _nb_reader = NonBlockingReader() +-# device -> the copy in which a count increase was first seen (reported on the next tick) ++# device -> the copy in which a count increase was first seen ++# (reported on the next tick) + _confirm_pending: dict[torch.device, list] = {} + + # Sites that arm the pending raise even under VLLM_E2_GUARD_MODE=warn. warn is +@@ -442,12 +443,11 @@ def start_reader() -> None: + if _reader_started or _MODE == "off": + return + _reader_started = True +- threading.Thread( +- target=_reader, name="bounds-guard-reader", daemon=True +- ).start() ++ threading.Thread(target=_reader, name="bounds-guard-reader", daemon=True).start() + logger.info( + "bounds-guard reader started (mode=%s, interval=%.3gs).", +- _MODE, _INTERVAL, ++ _MODE, ++ _INTERVAL, + ) + + +@@ -531,7 +531,11 @@ def _poll_once() -> None: + # by which time the maxima of that hit are written). The fatal + # decision is count-based either way; only its report waits a tick. + pend = _confirm_pending.get(device) +- rising = [site for site, (n, _r, _v) in enumerate(rows) if n > _seen.get((device, site), 0)] ++ rising = [ ++ site ++ for site, (n, _r, _v) in enumerate(rows) ++ if n > _seen.get((device, site), 0) ++ ] + if pend is None: + if rising: + _confirm_pending[device] = rows +@@ -644,21 +648,18 @@ def drain(*, raise_on_fatal: bool = False) -> dict[str, list[int]]: + _report_rows(device, rows) + for site, (n_masked, _r, _v) in enumerate(rows): + if n_masked and (_MODE == "raise" or site in _ALWAYS_FATAL_SITES): +- fatal.append( +- f"{SITE_NAMES.get(site, site)} on {device}: {n_masked}" +- ) ++ fatal.append(f"{SITE_NAMES.get(site, site)} on {device}: {n_masked}") + snap = snapshot() + if fatal: + logger.error( + "bounds-guard FINAL DRAIN: fatal-site hits %s (snapshot %s)", +- fatal, snap, ++ fatal, ++ snap, + ) + if raise_on_fatal: + raise RuntimeError(f"bounds-guard FINAL DRAIN: fatal-site hits {fatal}") + else: +- logger.info( +- "bounds-guard FINAL DRAIN: no fatal-site hits (snapshot %s)", snap +- ) ++ logger.info("bounds-guard FINAL DRAIN: no fatal-site hits (snapshot %s)", snap) + return snap + + +diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py +index a6972f1..eeeadae 100644 +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -832,7 +832,7 @@ def postprocess_mamba_fused_kernel( + if src_block_idx == dest_block_idx and accept_token_bias == 0: + return + +- if BT_ROW_IS_REQ: ++ if BT_ROW_IS_REQ: # noqa: SIM108 -- a ternary would drop this comment + # Persistent per-req-slot table: the row is the request's own slot, so + # it stays correct however many steps separate the gather from this + # launch (the PP-deferred post-process is enqueued at the top of +diff --git a/vllm_ple_test_fault/__init__.py b/vllm_ple_test_fault/__init__.py +index e8fe835..41830e0 100644 +--- a/vllm_ple_test_fault/__init__.py ++++ b/vllm_ple_test_fault/__init__.py +@@ -145,16 +145,12 @@ class FaultSpec: + if bytes_sent is not None or frame_size is None: + return False + suffix = ( +- f" bytes_written={int(bytes_written)}" +- f" frame_size={int(frame_size)}" ++ f" bytes_written={int(bytes_written)} frame_size={int(frame_size)}" + ) + elif bytes_sent is not None: + if frame_size is None: + return False +- suffix = ( +- f" bytes_sent={int(bytes_sent)}" +- f" frame_size={int(frame_size)}" +- ) ++ suffix = f" bytes_sent={int(bytes_sent)} frame_size={int(frame_size)}" + elif frame_size is not None: + return False + message = ( +@@ -306,7 +302,9 @@ def _parse_fault_spec(value: str | None) -> FaultSpec | None: + if value is None: + return None + if not value or value != value.strip(): +- raise ValueError(f"{_ENV_NAME}: spec must be non-empty with no outer whitespace") ++ raise ValueError( ++ f"{_ENV_NAME}: spec must be non-empty with no outer whitespace" ++ ) + + fields: dict[str, str] = {} + for item in value.split(","): +@@ -355,8 +353,7 @@ def _parse_fault_spec(value: str | None) -> FaultSpec | None: + + if role in {"offload", "bootstrap"} and not has_any_rank: + raise ValueError( +- f"{_ENV_NAME}: {role} cut {cut!r} requires rank=any; " +- "it has no TP/PP rank" ++ f"{_ENV_NAME}: {role} cut {cut!r} requires rank=any; it has no TP/PP rank" + ) + if role == "gpu": + if pp_rank is not None and pp_rank != 0: +-- +2.43.0 + diff --git a/patches/0029-rebase-remove-v0.30.0-merge-artefacts-import-order.patch b/patches/0029-rebase-remove-v0.30.0-merge-artefacts-import-order.patch new file mode 100644 index 0000000..eab9f07 --- /dev/null +++ b/patches/0029-rebase-remove-v0.30.0-merge-artefacts-import-order.patch @@ -0,0 +1,1072 @@ +From 81807002fb36e2b9eab05728ea3bdf54fd12c280 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 04:28:34 -0400 +Subject: [PATCH 029/103] rebase: remove v0.30.0 merge artefacts; import order + +- Drop the PLE-offload half of ple_layer.py that upstream #54371 deleted and a + 3-way merge resurrected (it no longer imported on v0.30.0). +- Drop the structured-output helper duplicate that v0.30.0 already provides, + and an unused import. +- Re-sort the import blocks the merges left unsorted (ruff I001). +--- + tests/v1/core/test_deferred_block_free.py | 3 +- + vllm/models/qwen4_exp/nvidia/model.py | 7 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 2 +- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 882 ---------------------- + vllm/models/qwen4_exp/nvidia/qsa.py | 3 +- + vllm/v1/core/sched/scheduler.py | 4 +- + vllm/v1/structured_output/utils.py | 13 - + 7 files changed, 8 insertions(+), 906 deletions(-) + +diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py +index c45f5a7..e94161b 100644 +--- a/tests/v1/core/test_deferred_block_free.py ++++ b/tests/v1/core/test_deferred_block_free.py +@@ -16,14 +16,13 @@ import time + from unittest.mock import PropertyMock, patch + + import pytest +- + import torch + + from vllm.config import VllmConfig + from vllm.v1.core.block_pool import BlockPool + from vllm.v1.core.kv_cache_utils import BlockHashWithGroupId, KVCacheBlockCopy +-from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.core.sched.output import SchedulerOutput ++from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.outputs import ModelRunnerOutput + from vllm.v1.request import RequestStatus + +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 269bbbd..816a4e2 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -2,16 +2,16 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Inference-only Qwen4Exp model.""" + ++import os + from collections.abc import Iterable + from itertools import islice + +-import os +- + import torch + from torch import nn + + from vllm.config import VllmConfig + from vllm.distributed import get_pp_group ++from vllm.logger import init_logger + from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + ) +@@ -28,8 +28,8 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( + ) + from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.model_executor.layers.vocab_parallel_embedding import ( +- PinnedHostVocabParallelEmbedding, + ParallelLMHead, ++ PinnedHostVocabParallelEmbedding, + VocabParallelEmbedding, + should_use_host_embed_table, + ) +@@ -85,7 +85,6 @@ from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm + from .ple_layer import Qwen4ExpPLELayer + from .qsa import Qwen4ExpQSAAttention, load_qsa_static_kv_scales +-from vllm.logger import init_logger + + logger = init_logger(__name__) + +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 0ad96a2..2df3868 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -28,8 +28,8 @@ from vllm.model_executor.layers.layernorm import GemmaRMSNorm + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.model_executor.layers.logits_processor import LogitsProcessor + from vllm.model_executor.layers.vocab_parallel_embedding import ( +- PinnedHostVocabParallelEmbedding, + ParallelLMHead, ++ PinnedHostVocabParallelEmbedding, + VocabParallelEmbedding, + should_use_host_embed_table, + ) +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index bec35e6..2ec8790 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -4,8 +4,6 @@ + + from collections.abc import Sequence + +-import os +- + import torch + from torch import nn + +@@ -65,886 +63,6 @@ class Qwen4ExpPLEGroupedNorm(nn.Module): + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + +-class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase): +- """FP8 PLE embedding with one global checkpoint scale.""" +- +- def create_weights( +- self, +- layer: nn.Module, +- input_size_per_partition: int, +- output_partition_sizes: list[int], +- input_size: int, +- output_size: int, +- params_dtype: torch.dtype, +- **extra_weight_attrs, +- ) -> None: +- del input_size, output_size, params_dtype +- weight_loader = extra_weight_attrs.get("weight_loader") +- weight = create_fp8_weight_parameter( +- sum(output_partition_sizes), input_size_per_partition, weight_loader +- ) +- layer.register_parameter("weight", weight) +- +- weight_scale = create_fp8_scale_parameter( +- PerTensorScaleParameter, +- output_partition_sizes, +- input_size_per_partition, +- None, +- weight_loader, +- scale_dtype=torch.float32, +- ) +- layer.register_parameter("weight_scale", weight_scale) +- +- def process_weights_after_loading(self, layer: nn.Module) -> None: +- """Reject FP8 PLE checkpoints without a global scale.""" +- sentinel = torch.finfo(torch.float32).min +- if torch.any(layer.weight_scale == sentinel): +- raise ValueError("FP8 PLE checkpoint is missing its global scale") +- +- def apply( +- self, +- layer: nn.Module, +- x: torch.Tensor, +- bias: torch.Tensor | None = None, +- ) -> torch.Tensor: +- raise NotImplementedError("PLE FP8 weights only support embedding lookup") +- +- def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: +- return F.embedding(input_, layer.weight) +- +- +-_NVFP4_BLOCK_SIZE = 16 +- +- +-class Qwen4ExpPLENVFp4EmbeddingMethod(QuantizeMethodBase): +- """NVFP4 PLE embedding kept packed at runtime. +- +- The table stays as uint8 codes with FP8 block scales (16 values per +- block) and one FP32 global scale. A lookup only gathers the packed rows; +- the PLE layer dequantizes them after lookup. +- """ +- +- def create_weights( +- self, +- layer: nn.Module, +- input_size_per_partition: int, +- output_partition_sizes: list[int], +- input_size: int, +- output_size: int, +- params_dtype: torch.dtype, +- **extra_weight_attrs, +- ) -> None: +- del input_size, output_size, extra_weight_attrs +- if input_size_per_partition % _NVFP4_BLOCK_SIZE: +- raise ValueError( +- "NVFP4 PLE embedding requires the embedding dim to be a " +- f"multiple of {_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}" +- ) +- self.params_dtype = params_dtype +- self.head_dim = input_size_per_partition +- self.packed_row_width = ( +- input_size_per_partition // 2 +- + input_size_per_partition // _NVFP4_BLOCK_SIZE +- ) +- rows = sum(output_partition_sizes) +- weight = nn.Parameter( +- torch.empty(rows, input_size_per_partition // 2, dtype=torch.uint8), +- requires_grad=False, +- ) +- set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) +- layer.register_parameter("weight", weight) +- +- weight_scale = nn.Parameter( +- torch.empty( +- rows, +- input_size_per_partition // _NVFP4_BLOCK_SIZE, +- dtype=torch.float8_e4m3fn, +- ), +- requires_grad=False, +- ) +- set_weight_attrs(weight_scale, {"input_dim": 1, "output_dim": 0}) +- layer.register_parameter("weight_scale", weight_scale) +- +- weight_scale_2 = nn.Parameter( +- torch.empty((), dtype=torch.float32), requires_grad=False +- ) +- layer.register_parameter("weight_scale_2", weight_scale_2) +- self._lut: torch.Tensor | None = None +- +- def process_weights_after_loading(self, layer: nn.Module) -> None: +- # Build the lookup table before CUDA graph capture. +- self._lut = torch.tensor( +- _FP4_VALUES, dtype=torch.float32, device=layer.weight.device +- ) +- +- def apply( +- self, +- layer: nn.Module, +- x: torch.Tensor, +- bias: torch.Tensor | None = None, +- ) -> torch.Tensor: +- raise NotImplementedError("PLE NVFP4 weights only support embedding lookup") +- +- def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: +- """Gather packed rows for deferred dequantization.""" +- codes = F.embedding(input_, layer.weight) +- scales = F.embedding(input_, layer.weight_scale) +- return torch.cat((codes, scales.view(torch.uint8)), dim=-1) +- +- def dequantize( +- self, +- packed_rows: torch.Tensor, +- scale_2: torch.Tensor, +- output_dtype: torch.dtype, +- ) -> torch.Tensor: +- """Dequantize rows containing NVFP4 codes and block scales.""" +- lut = self._lut +- if lut is None or lut.device != packed_rows.device: +- lut = torch.tensor( +- _FP4_VALUES, dtype=torch.float32, device=packed_rows.device +- ) +- self._lut = lut +- return _dequant_nvfp4_rows( +- packed_rows, self.head_dim, scale_2, output_dtype, lut +- ) +- +- +-def _ple_dtype_is_fp8(ple_embedding_dtype: object) -> bool: +- """Return whether the model config declares FP8 PLE checkpoint weights.""" +- +- if ple_embedding_dtype is None: +- return False +- if isinstance(ple_embedding_dtype, torch.dtype): +- return ple_embedding_dtype == torch.float8_e4m3fn +- return str(ple_embedding_dtype).rsplit(".", 1)[-1] == "float8_e4m3fn" +- +- +-def _get_ple_embedding_quant_method( +- quant_config: QuantizationConfig | None, +- prefix: str, +- ple_embedding_dtype: object = None, +-) -> QuantizeMethodBase | None: +- """Select a packed PLE embedding method for quantized checkpoint shards.""" +- +- if isinstance(quant_config, ModelOptMixedPrecisionConfig): +- if quant_config._resolve_quant_algo(prefix) == "FP8": +- return Qwen4ExpPLEFp8EmbeddingMethod() +- return None +- +- if isinstance(quant_config, Fp8Config): +- if not quant_config.is_checkpoint_fp8_serialized: +- return None +- ignored_layers = quant_config.ignored_layers +- if is_layer_skipped( +- prefix, +- ignored_layers, +- quant_config.packed_modules_mapping, +- match_mode=quant_config.ignored_layers_match_mode, +- ): +- return None +- # PLE checkpoint shards form one runtime embedding parameter. +- shard_prefix = f"{prefix}.shard_" +- if any(name.startswith(shard_prefix) for name in ignored_layers): +- return None +- return Qwen4ExpPLEFp8EmbeddingMethod() +- +- if isinstance(quant_config, ModelOptNvFp4Config): +- if not quant_config.is_checkpoint_nvfp4_serialized: +- return None +- if not quant_config.is_layer_excluded(prefix): +- logger.info_once("PLE embedding %s uses the runtime NVFP4 method", prefix) +- return Qwen4ExpPLENVFp4EmbeddingMethod() +- if _ple_dtype_is_fp8(ple_embedding_dtype): +- logger.info_once( +- "Excluded ModelOpt PLE embedding %s uses the runtime FP8 method", +- prefix, +- ) +- return Qwen4ExpPLEFp8EmbeddingMethod() +- +- # Checkpoints pairing quantized experts with a global-scale FP8 PLE table +- # (the RadixArk Flash-Next table grafted onto a CompressedTensors checkpoint) +- # list the PLE under exclude_modules, so the FP8-ness is visible only in the +- # tensors' dtype / the config's ple_embedding_dtype, not in the quant config +- # type dispatched above. Without this branch the offload worker treats the +- # table as BF16 and materializes ~96 GiB in host RAM (OOM-killed on the +- # 96 GB host). Honour either the config key (the -e1 +- # checkpoint copy) or the explicit VLLM_PLE_FP8_GLOBAL_SCALE=1 opt-in the +- # previous release uses. Placement after the typed dispatch is the +- # position already proven in service. +- if _ple_dtype_is_fp8(ple_embedding_dtype) or os.environ.get( +- "VLLM_PLE_FP8_GLOBAL_SCALE" +- ) == "1": +- logger.info_once( +- "PLE embedding %s: config declares an FP8 table with quant_config %s; " +- "using the runtime FP8 method (global weight_scale)", +- prefix, +- type(quant_config).__name__, +- ) +- return Qwen4ExpPLEFp8EmbeddingMethod() +- +- return None +- +- +-_FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +- +- +-def _dequant_nvfp4_codes( +- packed: torch.Tensor, +- scale: torch.Tensor, +- scale_2: torch.Tensor, +- lut: torch.Tensor | None = None, +-) -> torch.Tensor: +- """Unpack NVFP4 codes with FP8 block scales and an FP32 global scale.""" +- half = packed.shape[-1] +- codes = torch.stack([packed & 0xF, (packed >> 4) & 0xF], dim=-1).reshape( +- *packed.shape[:-1], half * 2 +- ) +- if lut is None: +- lut = torch.tensor(_FP4_VALUES, dtype=torch.float32, device=packed.device) +- magnitude = lut[(codes & 0x7).long()] +- sign = ((codes >> 3) & 1).to(torch.float32) +- fp4 = (magnitude * (1 - 2 * sign)).reshape( +- *packed.shape[:-1], -1, _NVFP4_BLOCK_SIZE +- ) +- output = fp4 * scale.float().unsqueeze(-1) * scale_2.float() +- return output.reshape(*packed.shape[:-1], half * 2) +- +- +-def _dequant_nvfp4_rows( +- packed_rows: torch.Tensor, +- head_dim: int, +- scale_2: torch.Tensor, +- output_dtype: torch.dtype, +- lut: torch.Tensor, +-) -> torch.Tensor: +- """Dequantize rows containing NVFP4 codes followed by block scales.""" +- half = head_dim // 2 +- codes = packed_rows[..., :half] +- scales = packed_rows[..., half:].contiguous().view(torch.float8_e4m3fn) +- return _dequant_nvfp4_codes(codes, scales, scale_2, lut).to(output_dtype) +- +- +-def _get_shared_nvfp4_outer_scale( +- outer_scales: dict[int, torch.Tensor], +-) -> torch.Tensor: +- """Validate and return the global scale shared by NVFP4 PLE shards.""" +- first_index, reference = next(iter(outer_scales.items())) +- reference = reference.reshape(()) +- for shard_index, outer_scale in outer_scales.items(): +- if not torch.equal(outer_scale.reshape(()), reference): +- raise ValueError( +- "NVFP4 PLE shards must share the same global scale, but " +- f"shards {first_index} and {shard_index} differ" +- ) +- return reference +- +- +-class Qwen4ExpNGramEmbedding(PleOffloadLayer): +- _offload_quant_method: QuantizeMethodBase | None = None +- _MASK64 = (1 << 64) - 1 +- _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 +- _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 +- _SPLITMIX_M2 = 0x94D049BB133111EB +- _PLE_LAYER_PRIME = 10007 +- +- @classmethod +- def _splitmix64(cls, value: int) -> int: +- """Mix an integer into a deterministic unsigned 64-bit value.""" +- value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64 +- value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 +- value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 +- return (value ^ (value >> 31)) & cls._MASK64 +- +- @staticmethod +- def _is_prime_64(value: int) -> bool: +- """Return whether a 64-bit integer is prime.""" +- if value < 2: +- return False +- for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): +- if value % prime == 0: +- return value == prime +- exponent = value - 1 +- shifts = 0 +- while exponent % 2 == 0: +- exponent //= 2 +- shifts += 1 +- for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): +- if base % value == 0: +- continue +- witness = pow(base, exponent, value) +- if witness in (1, value - 1): +- continue +- for _ in range(shifts - 1): +- witness = pow(witness, 2, value) +- if witness == value - 1: +- break +- else: +- return False +- return True +- +- @classmethod +- def _nth_prime_after(cls, start: int, count: int) -> int: +- """Return the ``count``-th prime strictly greater than ``start``.""" +- prime = int(start) +- for _ in range(count): +- candidate = prime + 1 +- if candidate <= 2: +- prime = 2 +- continue +- if candidate % 2 == 0: +- candidate += 1 +- while not cls._is_prime_64(candidate): +- candidate += 2 +- prime = candidate +- return prime +- +- @classmethod +- def _make_layer_multipliers( +- cls, +- *, +- ngram_size: int, +- unigram_vocab_size: int, +- seed: int, +- ple_dense_layer_id: int, +- ) -> list[int]: +- """Build deterministic hash multipliers for one PLE layer.""" +- max_multiplier = ((1 << 63) - 1) // unigram_vocab_size +- half_bound = max(1, max_multiplier // 2) +- base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id +- multipliers = [] +- for index in range(ngram_size): +- value = base_seed + cls._SPLITMIX_GAMMA * (index + 1) +- multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1) +- return multipliers +- +- @classmethod +- def _make_vocab_layout( +- cls, +- *, +- ngram_vocab_size_base: int, +- ngram_heads: int, +- ple_dense_layer_id: int, +- ) -> tuple[list[int], list[int], int]: +- """Build per-head vocabulary sizes, offsets, and total row count.""" +- sizes: list[int] = [] +- offsets: list[int] = [] +- offset = 0 +- for local_head in range(ngram_heads): +- global_head = ple_dense_layer_id * ngram_heads + local_head +- size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) +- sizes.append(size) +- offsets.append(offset) +- offset += size +- return sizes, offsets, offset +- +- def __init__( +- self, +- config: Qwen4ExpTextConfig, +- embedding_dim: int, +- ple_dense_layer_id: int, +- prefix: str, +- layer_name: str, +- quant_config: QuantizationConfig | None = None, +- params_dtype: torch.dtype | None = None, +- ) -> None: +- super().__init__() +- self.layer_name = layer_name +- self.embedding_dim = embedding_dim +- self.ngram_size = int(config.ngram_size) +- self.heads_per_ngram = int(config.heads_per_ngram) +- self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram +- if self.ngram_size < 2: +- raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") +- if self.heads_per_ngram <= 0: +- raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") +- if embedding_dim % self.ngram_heads: +- raise ValueError( +- "ple_embed_dim must be divisible by total ngram heads: " +- f"{embedding_dim} % {self.ngram_heads} != 0" +- ) +- self.head_dim = embedding_dim // self.ngram_heads +- self.eos_token_id = int(config.eos_token_id) +- self.unigram_vocab_size = int(config.vocab_size) +- self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) +- if self.split_ngram_parts <= 0: +- raise ValueError("split_ngram_parts must be positive") +- +- multipliers = self._make_layer_multipliers( +- ngram_size=self.ngram_size, +- unigram_vocab_size=self.unigram_vocab_size, +- seed=int(getattr(config, "seed", 1234)), +- ple_dense_layer_id=ple_dense_layer_id, +- ) +- self.register_buffer( +- "layer_multipliers", +- torch.tensor(multipliers, dtype=torch.long), +- persistent=True, +- ) +- +- sizes, offsets, total_vocab_size = self._make_vocab_layout( +- ngram_vocab_size_base=int(config.ngram_vocab_size_base), +- ngram_heads=self.ngram_heads, +- ple_dense_layer_id=ple_dense_layer_id, +- ) +- self.register_buffer( +- "ngram_heads_vocab_sizes", +- torch.tensor(sizes, dtype=torch.long), +- persistent=True, +- ) +- self.register_buffer( +- "ngram_heads_offsets", +- torch.tensor(offsets, dtype=torch.long), +- persistent=True, +- ) +- divisor = int(config.make_ngram_vocab_size_divisible_by) +- padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor +- self.ngram_embedding = PLEVocabParallelEmbedding( +- padded_vocab_size, +- self.head_dim, +- params_dtype=params_dtype, +- padding_size=divisor, +- prefix=f"{prefix}.ngram_embedding", +- quant_method=_get_ple_embedding_quant_method( +- quant_config, +- f"{prefix}.ngram_embedding", +- getattr(config, "ple_embedding_dtype", None), +- ), +- ) +- +- @staticmethod +- def _shift_precompute( +- tokens: torch.Tensor, eos_token_id: int +- ) -> tuple[torch.Tensor, torch.Tensor]: +- if tokens.dim() != 2: +- raise ValueError("tokens must be a 2D tensor") +- batch_size, seq_len = tokens.shape +- positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) +- eos_positions = torch.where(tokens == eos_token_id, positions, -1) +- previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values +- previous_eos = torch.cat( +- [ +- eos_positions.new_full((batch_size, 1), -1), +- previous_eos_inclusive[:, :-1], +- ], +- dim=1, +- ) +- return positions, positions.unsqueeze(0) - previous_eos - 1 +- +- @staticmethod +- def _shift_apply( +- tokens: torch.Tensor, +- positions: torch.Tensor, +- position_in_segment: torch.Tensor, +- shift: int, +- eos_token_id: int, +- ) -> torch.Tensor: +- if shift == 0: +- return tokens +- source = positions - shift +- gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) +- shifted = tokens.gather(1, gather_indices) +- valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) +- return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) +- +- def compute_ngram_ids( +- self, +- input_ids: torch.Tensor, +- query_start_loc: torch.Tensor, +- ngram_context: torch.Tensor, +- output: torch.Tensor | None = None, +- ) -> torch.Tensor: +- """Compute n-gram embedding indices for the current request layout.""" +- input_ids = input_ids.reshape(-1) +- num_reqs = query_start_loc.numel() - 1 +- num_tokens = input_ids.shape[0] +- +- if input_ids.is_cuda: +- return ple_ngram_ids( +- input_ids=input_ids, +- query_start_loc=query_start_loc, +- ngram_context=ngram_context, +- layer_multipliers=self.layer_multipliers, +- ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes, +- ngram_heads_offsets=self.ngram_heads_offsets, +- eos_token_id=self.eos_token_id, +- heads_per_ngram=self.heads_per_ngram, +- output=output, +- ) +- input_ids = input_ids.long() +- query_start_loc = query_start_loc.long() +- if num_reqs <= 0: +- raise ValueError("PLE requires at least one request") +- +- if is_offload_process(): +- max_seq_len = max( +- 1, +- int((query_start_loc[1:] - query_start_loc[:-1]).max().item()), +- ) +- # CUDA graphs can pad input_ids beyond the valid request layout. +- # Keep those rows from overwriting the final real token. +- num_valid_tokens = min(int(query_start_loc[-1].item()), num_tokens) +- else: +- max_seq_len = num_tokens +- num_valid_tokens = num_tokens +- +- positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64) +- packed = torch.full( +- (num_reqs, max_seq_len), +- self.eos_token_id, +- device=input_ids.device, +- dtype=torch.int64, +- ) +- request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 +- request_indices.clamp_(max=num_reqs - 1) +- columns = (positions - query_start_loc[request_indices]).clamp( +- 0, packed.shape[1] - 1 +- ) +- packed[request_indices[:num_valid_tokens], columns[:num_valid_tokens]] = ( +- input_ids[:num_valid_tokens] +- ) +- ngram_context = ngram_context[:num_reqs].to( +- device=input_ids.device, dtype=torch.long +- ) +- +- context = torch.cat([ngram_context, packed], dim=-1) +- positions_2d, position_in_segment = self._shift_precompute( +- context, self.eos_token_id +- ) +- shifted = [context] +- for shift in range(1, self.ngram_size): +- shifted.append( +- self._shift_apply( +- context, +- positions_2d, +- position_in_segment, +- shift, +- self.eos_token_id, +- ) +- ) +- adjusted_columns = columns + self.ngram_size - 1 +- id_blocks = [] +- for ngram in range(2, self.ngram_size + 1): +- start = (ngram - 2) * self.heads_per_ngram +- end = start + self.heads_per_ngram +- mixed = shifted[0] * self.layer_multipliers[0] +- for index in range(1, ngram): +- mixed = torch.bitwise_xor( +- mixed, shifted[index] * self.layer_multipliers[index] +- ) +- sizes = self.ngram_heads_vocab_sizes[start:end] +- offsets = self.ngram_heads_offsets[start:end] +- ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets +- id_blocks.append(ids[request_indices, adjusted_columns]) +- return torch.cat(id_blocks, dim=-1) +- +- def forward_impl( # type: ignore[override] +- self, +- hidden_states: torch.Tensor, +- input_ids: torch.Tensor, +- query_start_loc: torch.Tensor, +- ngram_context: torch.Tensor, +- output_buffer: torch.Tensor | None = None, +- ) -> torch.Tensor: +- del hidden_states +- if input_ids.is_cuda: +- # Keep request-dependent ID generation outside piecewise graphs. +- ngram_ids = input_ids.new_empty( +- (input_ids.numel(), self.ngram_heads), dtype=torch.long +- ) +- torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids( +- input_ids, +- query_start_loc, +- ngram_context, +- ngram_ids, +- self.layer_name, +- ) +- else: +- ngram_ids = self.compute_ngram_ids( +- input_ids, +- query_start_loc, +- ngram_context, +- ) +- num_tokens = input_ids.reshape(-1).shape[0] +- if output_buffer is not None: +- quant_method = getattr(self.ngram_embedding, "quant_method", None) +- if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): +- output_dim = self.get_offload_output_dim(self.embedding_dim) +- output = output_buffer[:num_tokens, :output_dim] +- output.copy_(self.ngram_embedding(ngram_ids).flatten(-2)) +- return output +- output = output_buffer[:num_tokens, : self.embedding_dim] +- torch.index_select( +- self.ngram_embedding.weight, +- 0, +- ngram_ids.reshape(-1), +- out=output.reshape(-1, self.head_dim), +- ) +- return output +- return self.ngram_embedding(ngram_ids).flatten(-2) +- +- def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: +- """Keep quantized lookup results in their embedding storage dtype.""" +- embedding = getattr(self, "ngram_embedding", None) +- weight = getattr(embedding, "weight", None) +- if weight is not None: +- return weight.dtype +- if isinstance(self._offload_quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): +- return torch.uint8 +- if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod): +- return torch.float8_e4m3fn +- return default_dtype +- +- def get_offload_output_dim(self, default_dim: int) -> int: +- """Keep NVFP4 lookup rows packed while transferring them to the GPU.""" +- quant_method = getattr( +- getattr(self, "ngram_embedding", None), "quant_method", None +- ) +- if quant_method is None: +- quant_method = self._offload_quant_method +- if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): +- if default_dim % _NVFP4_BLOCK_SIZE: +- raise ValueError( +- "NVFP4 PLE output dim must be a multiple of " +- f"{_NVFP4_BLOCK_SIZE}, got {default_dim}" +- ) +- return default_dim // 2 + default_dim // _NVFP4_BLOCK_SIZE +- return default_dim +- +- def initialize_dummy_offload_metadata(self, device: torch.device) -> None: +- """Initialize quantization metadata skipped by the dummy loader.""" +- quant_method = self._offload_quant_method +- if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): +- self.register_buffer( +- "_offload_weight_scale", +- torch.ones((), dtype=torch.float32, device=device), +- persistent=False, +- ) +- elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): +- self.register_buffer( +- "_offload_weight_scale_2", +- torch.ones((), dtype=torch.float32, device=device), +- persistent=False, +- ) +- self.register_buffer( +- "_offload_nvfp4_lut", +- torch.tensor(_FP4_VALUES, dtype=torch.float32, device=device), +- persistent=False, +- ) +- +- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: +- """Load hash buffers and checkpoint-split embedding rows.""" +- +- # Prefix grouping may invoke this method repeatedly for one module. +- # Treat the global FP8 scale as incrementally loaded state. +- # GPU workers retain only dequantization metadata. The CPU process owns +- # the embedding table and transfers quantized lookup rows unchanged. +- if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): +- retained: set[str] = set() +- quant_method = self._offload_quant_method +- if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): +- for name, loaded_weight in weights: +- if name != "ngram_embedding.weight_scale": +- continue +- self.register_buffer( +- "_offload_weight_scale", +- loaded_weight.to( +- device=torch.accelerator.current_accelerator() +- ), +- persistent=False, +- ) +- retained.add(name) +- elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod): +- outer_scales: dict[int, torch.Tensor] = {} +- for name, loaded_weight in weights: +- if not name.startswith( +- "ngram_embedding.shard_" +- ) or not name.endswith(".weight_scale_2"): +- continue +- shard_text = name[ +- len("ngram_embedding.shard_") : -len(".weight_scale_2") +- ] +- shard_index = int(shard_text) +- outer_scales[shard_index] = loaded_weight +- retained.add(name) +- if not retained: +- raise ValueError( +- "NVFP4 PLE offload checkpoint is missing its global scale" +- ) +- scale_2 = _get_shared_nvfp4_outer_scale(outer_scales).to( +- device=torch.accelerator.current_accelerator() +- ) +- self.register_buffer( +- "_offload_weight_scale_2", scale_2, persistent=False +- ) +- self.register_buffer( +- "_offload_nvfp4_lut", +- torch.tensor( +- _FP4_VALUES, dtype=torch.float32, device=scale_2.device +- ), +- persistent=False, +- ) +- else: +- for _ in weights: +- pass +- return retained +- +- persistent_buffers = { +- "layer_multipliers": self.layer_multipliers, +- "ngram_heads_offsets": self.ngram_heads_offsets, +- "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, +- } +- loaded: set[str] = set() +- regular_weights: list[tuple[str, torch.Tensor]] = [] +- shard_prefix = "ngram_embedding.shard_" +- quant_method = getattr(self.ngram_embedding, "quant_method", None) +- nvfp4_runtime = isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod) +- packed_codes: dict[int, torch.Tensor] = {} +- packed_scales: dict[int, torch.Tensor] = {} +- packed_outer_scales: dict[int, torch.Tensor] = {} +- +- for name, loaded_weight in weights: +- leaf_name = name.rsplit(".", 1)[-1] +- if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": +- continue +- if name in persistent_buffers: +- buffer = persistent_buffers[name] +- if buffer.shape != loaded_weight.shape: +- raise ValueError( +- f"Shape mismatch for {name}: expected " +- f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" +- ) +- buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) +- loaded.add(name) +- continue +- +- if nvfp4_runtime and name.startswith(shard_prefix): +- suffix = name[len(shard_prefix) :] +- for ending, sink in ( +- (".weight_scale_2", packed_outer_scales), +- (".weight_scale", packed_scales), +- (".weight", packed_codes), +- ): +- if not suffix.endswith(ending): +- continue +- shard_text = suffix[: -len(ending)] +- shard_index = int(shard_text) +- if shard_index >= self.split_ngram_parts: +- raise ValueError( +- f"PLE embedding shard index {shard_index} exceeds " +- f"split_ngram_parts={self.split_ngram_parts}" +- ) +- sink[shard_index] = loaded_weight +- break +- else: +- regular_weights.append((name, loaded_weight)) +- continue +- +- if ( +- not nvfp4_runtime +- and name.startswith(shard_prefix) +- and name.endswith(".weight") +- ): +- shard_text = name[len(shard_prefix) : -len(".weight")] +- if not shard_text.isdigit(): +- regular_weights.append((name, loaded_weight)) +- continue +- shard_index = int(shard_text) +- if shard_index >= self.split_ngram_parts: +- raise ValueError( +- f"PLE embedding shard index {shard_index} exceeds " +- f"split_ngram_parts={self.split_ngram_parts}" +- ) +- embedding = self.ngram_embedding +- shard_size = ( +- embedding.org_vocab_size + self.split_ngram_parts - 1 +- ) // self.split_ngram_parts +- checkpoint_start = shard_index * shard_size +- expected_rows = max( +- 0, +- min(shard_size, embedding.org_vocab_size - checkpoint_start), +- ) +- expected_shape = (expected_rows, embedding.embedding_dim) +- if tuple(loaded_weight.shape) != expected_shape: +- raise ValueError( +- f"Shape mismatch for PLE embedding shard {shard_index}: " +- f"expected {expected_shape}, got " +- f"{tuple(loaded_weight.shape)}" +- ) +- embedding.weight.weight_loader( +- embedding.weight, +- loaded_weight, +- checkpoint_start=checkpoint_start, +- ) +- loaded.add("ngram_embedding.weight") +- continue +- regular_weights.append((name, loaded_weight)) +- +- if nvfp4_runtime: +- if not packed_codes: +- raise ValueError("NVFP4 PLE checkpoint contains no packed shards") +- for shard_index in packed_codes: +- if ( +- shard_index not in packed_scales +- or shard_index not in packed_outer_scales +- ): +- raise ValueError( +- f"NVFP4 PLE shard {shard_index} is missing its scale tensors" +- ) +- shared_outer_scale = _get_shared_nvfp4_outer_scale( +- { +- shard_index: packed_outer_scales[shard_index] +- for shard_index in packed_codes +- } +- ) +- +- for shard_index, codes in packed_codes.items(): +- embedding = self.ngram_embedding +- shard_size = ( +- embedding.org_vocab_size + self.split_ngram_parts - 1 +- ) // self.split_ngram_parts +- checkpoint_start = shard_index * shard_size +- expected_rows = max( +- 0, min(shard_size, embedding.org_vocab_size - checkpoint_start) +- ) +- expected_codes_shape = (expected_rows, embedding.embedding_dim // 2) +- if tuple(codes.shape) != expected_codes_shape: +- raise ValueError( +- f"Shape mismatch for packed PLE shard {shard_index}: " +- f"expected {expected_codes_shape}, got {tuple(codes.shape)}" +- ) +- tp_bounds = { +- "checkpoint_start": checkpoint_start, +- "tp_start": embedding.shard_indices.org_vocab_start_index, +- "tp_end": embedding.shard_indices.org_vocab_end_index, +- } +- scales = packed_scales[shard_index] +- expected_scales_shape = ( +- expected_rows, +- embedding.embedding_dim // _NVFP4_BLOCK_SIZE, +- ) +- if tuple(scales.shape) != expected_scales_shape: +- raise ValueError( +- f"Scale shape mismatch for packed PLE shard {shard_index}: " +- f"expected {expected_scales_shape}, got {tuple(scales.shape)}" +- ) +- copy_ple_embedding_shard_(embedding.weight.data, codes, **tp_bounds) +- copy_ple_embedding_shard_(embedding.weight_scale.data, scales, **tp_bounds) +- +- if nvfp4_runtime: +- self.ngram_embedding.weight_scale_2.data.copy_(shared_outer_scale) +- loaded.update( +- { +- "ngram_embedding.weight", +- "ngram_embedding.weight_scale", +- "ngram_embedding.weight_scale_2", +- } +- ) +- +- if regular_weights: +- loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) +- return loaded +- +- + class Qwen4ExpPLELayer(nn.Module, MambaBase): + def __init__( + self, +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 91f34de..87c4b3c 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -22,6 +22,7 @@ from vllm.config import VllmConfig + from vllm.config.cache import CacheDType + from vllm.distributed import get_tensor_model_parallel_world_size + from vllm.forward_context import get_forward_context ++from vllm.logger import init_logger + from vllm.model_executor.layers.attention.attention import ( + set_default_quant_scales, + ) +@@ -29,7 +30,6 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.model_executor.layers.layernorm import GemmaRMSNorm + from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear + from vllm.model_executor.layers.quantization import QuantizationConfig +-from vllm.logger import init_logger + from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope + from vllm.model_executor.models.qwen3_next import Qwen3NextAttention + from vllm.platforms import current_platform +@@ -63,7 +63,6 @@ from ..common.qsa_cache import QSAForwardMetadata + from . import model + from .indexer_qsa import QSAIndexer + +- + logger = init_logger(__name__) + + # --- FP8 K/V scale calibration collector ----------------------- +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index ae9eeb6..3b942cb 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -39,6 +39,7 @@ from vllm.v1.core.encoder_cache_manager import ( + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager + from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector + from vllm.v1.core.kv_cache_utils import KVCacheBlock ++from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log + from vllm.v1.core.sched.interface import PauseState, SchedulerInterface + from vllm.v1.core.sched.output import ( + CachedRequestData, +@@ -70,11 +71,10 @@ from vllm.v1.metrics.stats import ( + SchedulerStats, + ) + from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +-from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log + from vllm.v1.request import Request, RequestStatus, StreamingUpdate + from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup + from vllm.v1.spec_decode.metrics import SpecDecodingStats +-from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager ++from vllm.v1.structured_output import StructuredOutputManager + from vllm.v1.structured_output.utils import strip_speculative_padding + from vllm.v1.utils import record_function_or_nullcontext + +diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py +index 7c4af7a..92122f8 100644 +--- a/vllm/v1/structured_output/utils.py ++++ b/vllm/v1/structured_output/utils.py +@@ -573,16 +573,3 @@ def choice_as_grammar(choice: list[str]) -> str: + escaped_choices = (escape_ebnf_string(c) for c in choice) + grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices) + return grammar +- +-def strip_speculative_padding(token_ids: list[int]) -> list[int]: +- """Drop speculative-decoding padding from a token block. +- +- ngram and other speculative backends pad rejected draft positions with a +- -1 sentinel. Structured-output grammars treat every entry as a real token +- id, so the sentinels (and everything after the first one) are removed here, +- before tokens reach any backend, rather than inside a single backend. +- """ +- for i, token_id in enumerate(token_ids): +- if token_id < 0: +- return token_ids[:i] +- return token_ids +-- +2.43.0 + diff --git a/patches/0030-e2_guards-executable-graph-profile-and-breakable-cud.patch b/patches/0030-e2_guards-executable-graph-profile-and-breakable-cud.patch new file mode 100644 index 0000000..b9807ee --- /dev/null +++ b/patches/0030-e2_guards-executable-graph-profile-and-breakable-cud.patch @@ -0,0 +1,1101 @@ +From 00abe3b259622c748b332d94269be0fd8db44bd0 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 08:37:51 -0400 +Subject: [PATCH 030/103] e2_guards: executable graph-profile and + breakable-cudagraph guards + +New vllm/v1/e2_guards.py, called from the target and MTP constructors on every +rank before profiling or capture: the breakable-cudagraph pin is asserted +(three readings of the switch), a graph profile that cannot execute on v0.30.0 +is refused, and VLLM_E2_EXPECTED_CUDAGRAPH_MODE is checked by name. Tests +cover the drafter. An inert knob is no longer hashed into the cache key. +--- + .../test_ple_fp8_selector_checkpoint.py | 276 ++++++++++++++++++ + tests/test_envs.py | 15 +- + tests/v1/test_e2_guards.py | 215 ++++++++++++++ + vllm/envs.py | 63 +++- + vllm/models/qwen4_exp/nvidia/model.py | 8 + + vllm/models/qwen4_exp/nvidia/mtp.py | 5 + + vllm/v1/e2_guards.py | 218 ++++++++++++++ + vllm/v1/ple_offload/test_fault.py | 13 - + vllm/v1/worker/gpu_worker.py | 112 ------- + 9 files changed, 792 insertions(+), 133 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_ple_fp8_selector_checkpoint.py + create mode 100644 tests/v1/test_e2_guards.py + create mode 100644 vllm/v1/e2_guards.py + delete mode 100644 vllm/v1/ple_offload/test_fault.py + +diff --git a/tests/models/qwen4_exp/test_ple_fp8_selector_checkpoint.py b/tests/models/qwen4_exp/test_ple_fp8_selector_checkpoint.py +new file mode 100644 +index 0000000..815a22f +--- /dev/null ++++ b/tests/models/qwen4_exp/test_ple_fp8_selector_checkpoint.py +@@ -0,0 +1,276 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The shipped Flash-Next checkpoint must still select the FP8 PLE method on v0.30.0. ++ ++Why this file exists (vLLM 0.30 rebase). Our ++fork used to force the FP8 PLE embedding method from a 21-line catch-all branch ++in ``_get_ple_embedding_quant_method``. That function lived in the half of ++``vllm/models/qwen4_exp/nvidia/ple_layer.py`` that upstream #54371 deleted, so ++the branch is gone from the 0.30 base and the environment switch that fed it ++(``VLLM_PLE_FP8_GLOBAL_SCALE``) is now inert. ++ ++v0.30.0 selects instead in ``Qwen4ExpPLEEmbeddingMethod.from_quant_config`` ++(``ngram_embedding.py:168-208``), which is stricter: it takes the FP8 path from ++``text_config.ple_embedding_dtype == "float8_e4m3fn"`` (``:174-175``), and for a ++quantization config that is neither ``Fp8Config`` nor a ModelOpt variant it ++**raises** ``NotImplementedError`` (``:186-190``). Our checkpoint is ++compressed-tensors, so without that key the model would not build at all. ++ ++The first version of this file was rejected for reading a plain dict and ++never exercising the resolution chain, for naming the wrong engine prefix, and ++for conflating a dequantization check with a scale-loading check. This version ++walks the real chain -- ++``Qwen4ExpConfig(**config.json)`` (``models/qwen4_exp/config.py:231``) ++-> ``get_hf_text_config()`` (``transformers_utils/config.py:1142``) ++-> ``getattr(text_config, "ple_embedding_dtype")`` as ++``ngram_embedding.py:694-699`` reads it -- and keeps the three questions apart. ++ ++Fixture provenance: the published checkpoint's ``config.json``, the ++published v2.0.1 card. Excerpted verbatim for every field the ++selector or the resolver touches; ``config_groups`` keeps one representative ++group and ``ignore`` its first nine entries. ++""" ++ ++import pytest ++ ++from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 ++ CompressedTensorsConfig, ++) ++from vllm.model_executor.layers.quantization.fp8 import Fp8Config ++from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( ++ Qwen4ExpPLEEmbeddingMethod, ++ Qwen4ExpPLEFp8EmbeddingMethod, ++ Qwen4ExpPLEUnquantizedEmbeddingMethod, ++) ++ ++# ``ple_layer_ids = [2]`` selects the decoder layer whose ``layer_idx + 1`` is 2 ++# (``nvidia/model.py:303``), i.e. zero-based layer **1** -- not layer 2. The ++# multimodal entry class wraps the text model under ``language_model`` ++# (``nvidia/model.py:1111-1114``); the PLE layer prefix is ``.ple`` ++# (``:309-314``), the embedding ``.ple_embedding`` (``ple_layer.py:102``) and ++# the table ``.ngram_embedding`` (``ngram_embedding.py``). ++_PLE_PREFIX_MM = "language_model.model.layers.1.ple.ple_embedding.ngram_embedding" ++_PLE_PREFIX_TEXT = "model.layers.1.ple.ple_embedding.ngram_embedding" ++ ++# --------------------------------------------------------------------------- ++# Verbatim excerpt of the shipped checkpoint's config.json. ++# --------------------------------------------------------------------------- ++SHIPPED_CONFIG = { ++ "architectures": ["Qwen4ExpForConditionalGeneration"], ++ "model_type": "qwen4_exp", ++ "image_token_id": 248056, ++ "video_token_id": 248057, ++ "vision_start_token_id": 248053, ++ "vision_end_token_id": 248054, ++ "tie_word_embeddings": False, ++ "text_config": { ++ "model_type": "qwen4_exp_text", ++ # The key the whole selection now hangs on. ++ "ple_embedding_dtype": "float8_e4m3fn", ++ "ple_layer_ids": [2], ++ }, ++ "quantization_config": { ++ "quant_method": "compressed-tensors", ++ "format": "pack-quantized", ++ "quantization_status": "compressed", ++ "kv_cache_scheme": None, ++ "sparsity_config": {}, ++ "transform_config": {}, ++ "config_groups": { ++ "group_0": { ++ "targets": ["Linear"], ++ "weights": { ++ "num_bits": 4, ++ "type": "int", ++ "symmetric": True, ++ "strategy": "group", ++ "group_size": 128, ++ "dynamic": False, ++ "observer": "minmax", ++ "observer_kwargs": {}, ++ "actorder": None, ++ "block_structure": None, ++ }, ++ }, ++ }, ++ # The PLE is EXCLUDED from the quantization config -- which is exactly ++ # why the type-based dispatch cannot see that the table is FP8. ++ "ignore": [ ++ "lm_head", ++ "re:.*embed_tokens.*", ++ "re:.*\\.ple\\..*", ++ "re:.*visual\\..*", ++ "re:.*\\.gate$", ++ "re:.*hyper_connection.*", ++ "re:.*indexer.*", ++ "re:.*shared_expert.*", ++ "re:.*\\.self_attn\\..*", ++ ], ++ }, ++} ++ ++ ++def _quant_config() -> CompressedTensorsConfig: ++ import copy ++ ++ return CompressedTensorsConfig.from_config( ++ copy.deepcopy(SHIPPED_CONFIG["quantization_config"]) ++ ) ++ ++ ++def _resolved_text_config(): ++ """Walk the real chain the engine walks, not a dict lookup. ++ ++ ``ModelConfig`` stores ``get_hf_text_config(self.hf_config)`` ++ (``config/model.py:641``); that calls ``PretrainedConfig.get_text_config()`` ++ (``transformers_utils/config.py:1142``); ``Qwen4ExpConfig.__init__`` builds ++ the nested ``Qwen4ExpTextConfig`` from the ``text_config`` dict ++ (``models/qwen4_exp/config.py:231``). ++ """ ++ import copy ++ ++ from vllm.models.qwen4_exp.config import Qwen4ExpConfig ++ from vllm.transformers_utils.config import get_hf_text_config ++ ++ raw = copy.deepcopy(SHIPPED_CONFIG) ++ raw.pop("architectures", None) ++ raw.pop("model_type", None) ++ return get_hf_text_config(Qwen4ExpConfig(**raw)) ++ ++ ++def _embedding_dtype_as_the_engine_reads_it(text_config) -> str | None: ++ """``ngram_embedding.py:694-699`` -> ``getattr(config, "ple_embedding_dtype")``.""" ++ return getattr(text_config, "ple_embedding_dtype", None) ++ ++ ++# -------------------------------------------------------------------------- ++# 1. the resolution chain ++# -------------------------------------------------------------------------- ++ ++ ++def test_the_key_survives_the_real_text_config_resolution() -> None: ++ """THE CHAIN: nested dict -> Qwen4ExpConfig -> get_hf_text_config -> getattr.""" ++ text_config = _resolved_text_config() ++ assert _embedding_dtype_as_the_engine_reads_it(text_config) == "float8_e4m3fn" ++ assert list(text_config.ple_layer_ids) == [2] ++ ++ ++def test_a_top_level_only_key_is_lost_on_the_text_path() -> None: ++ """RESIDUAL RISK, made executable. ++ ++ A checkpoint copy that carried the key at the top level of a NESTED config ++ would not be seen by ``hf_text_config``, and the same checkpoint would then ++ hit the raise. This is a known caveat of the 0.30 port. ++ It says nothing about a FLAT config, which ``Qwen4ExpConfig`` converts into ++ a text config itself -- do not generalise it. ++ """ ++ import copy ++ ++ from vllm.models.qwen4_exp.config import Qwen4ExpConfig ++ from vllm.transformers_utils.config import get_hf_text_config ++ ++ raw = copy.deepcopy(SHIPPED_CONFIG) ++ raw.pop("architectures", None) ++ raw.pop("model_type", None) ++ raw["ple_embedding_dtype"] = raw["text_config"].pop("ple_embedding_dtype") ++ text_config = get_hf_text_config(Qwen4ExpConfig(**raw)) ++ assert _embedding_dtype_as_the_engine_reads_it(text_config) is None ++ ++ ++# -------------------------------------------------------------------------- ++# 2. the selector ++# -------------------------------------------------------------------------- ++ ++ ++@pytest.mark.parametrize("prefix", [_PLE_PREFIX_MM, _PLE_PREFIX_TEXT]) ++def test_shipped_checkpoint_selects_the_fp8_ple_method(prefix: str) -> None: ++ """THE PROPERTY: compressed-tensors + the resolved key -> FP8 PLE, no raise. ++ ++ Parametrised over both prefixes because the config-key branch precedes every ++ prefix-dependent check; if that ever stops being true, this fails. ++ """ ++ method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( ++ _quant_config(), ++ prefix, ++ _embedding_dtype_as_the_engine_reads_it(_resolved_text_config()), ++ ) ++ assert isinstance(method, Qwen4ExpPLEFp8EmbeddingMethod) ++ ++ ++def test_without_the_config_key_the_shipped_checkpoint_raises() -> None: ++ """FAILING CONTROL: the key is load-bearing, not decorative. ++ ++ Drop it and the same checkpoint hits ``ngram_embedding.py:186-190``. If this ++ stops raising, upstream has widened the selector and the offline conclusion ++ must be re-read -- it does not mean we are safer. ++ """ ++ with pytest.raises(NotImplementedError): ++ Qwen4ExpPLEEmbeddingMethod.from_quant_config( ++ _quant_config(), _PLE_PREFIX_MM, None ++ ) ++ ++ ++def test_the_key_beats_an_otherwise_unquantized_verdict() -> None: ++ """The key short-circuits before the ignore-list check. ++ ++ ``re:.*\\.ple\\..*`` puts our PLE in the ignore list. On an ``Fp8Config`` ++ that resolves to the *unquantized* method -- shown here as the contrast -- ++ so the key must win, or an FP8 table would be read as bf16. ++ """ ++ fp8_config = Fp8Config( ++ is_checkpoint_fp8_serialized=True, ignored_layers=[_PLE_PREFIX_MM] ++ ) ++ assert isinstance( ++ Qwen4ExpPLEEmbeddingMethod.from_quant_config(fp8_config, _PLE_PREFIX_MM, None), ++ Qwen4ExpPLEUnquantizedEmbeddingMethod, ++ ) ++ assert isinstance( ++ Qwen4ExpPLEEmbeddingMethod.from_quant_config( ++ fp8_config, _PLE_PREFIX_MM, "float8_e4m3fn" ++ ), ++ Qwen4ExpPLEFp8EmbeddingMethod, ++ ) ++ ++ ++# -------------------------------------------------------------------------- ++# 3. the scale -- a SEPARATE question from selection ++# -------------------------------------------------------------------------- ++ ++ ++def test_fp8_ple_method_rejects_an_unloaded_global_scale() -> None: ++ """Sentinel rejection only. This does NOT prove checkpoint scale LOADING. ++ ++ ``create_fp8_scale_parameter`` fills the parameter with the float32 sentinel ++ and the weight loader is expected to overwrite it; proving that the loader ++ does so needs a checkpoint shard, which belongs to a reference-host loader test, ++ not here. ++ """ ++ import torch ++ from torch import nn ++ ++ method = Qwen4ExpPLEFp8EmbeddingMethod() ++ unloaded = nn.Module() ++ unloaded.weight_scale = nn.Parameter( ++ torch.full((1,), torch.finfo(torch.float32).min, dtype=torch.float32), ++ requires_grad=False, ++ ) ++ with pytest.raises(ValueError, match="missing its global scale"): ++ method.process_weights_after_loading(unloaded) ++ ++ ++def test_fp8_ple_dequantization_applies_the_scale() -> None: ++ """Dequantization arithmetic only, with a hand-set scale. Not loading.""" ++ import torch ++ from torch import nn ++ ++ method = Qwen4ExpPLEFp8EmbeddingMethod() ++ loaded = nn.Module() ++ loaded.weight_scale = nn.Parameter( ++ torch.tensor([0.25], dtype=torch.float32), requires_grad=False ++ ) ++ method.process_weights_after_loading(loaded) # must not raise ++ ++ embeddings = torch.tensor([[1.0, 2.0]], dtype=torch.float32) ++ torch.testing.assert_close( ++ method.dequantize(loaded, embeddings, torch.float32), embeddings * 0.25 ++ ) +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 3e65518..920229d 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -610,7 +610,6 @@ class TestVllmMaxNSequences: + + _FORK_HASHED = [ + ("VLLM_HOST_EMBED_TABLE", "1", True), +- ("VLLM_PLE_FP8_GLOBAL_SCALE", "1", True), + ("VLLM_QSA_KV_CLIP_COUNT", "7", "7"), + ("VLLM_QSA_KV_COLLECT", "/var/tmp/qsa-absmax", "/var/tmp/qsa-absmax"), + # Capture-configuration factor read in C++ (topk.cu exact_medium_enabled); +@@ -622,6 +621,10 @@ _FORK_IGNORED = [ + ("VLLM_QSA_KV_SCALES", "/var/tmp/scales.json", "/var/tmp/scales.json"), + # "WARN" also checks the .strip().lower() mirror of the read site. + ("VLLM_E2_GUARD_MODE", " WARN ", "warn"), ++ ("VLLM_E2_ALLOW_BREAKABLE", "1", True), ++ ("VLLM_E2_EXPECTED_CUDAGRAPH_MODE", " FULL ", "FULL"), ++ # INERT: no read site on the 0.30 base, so not a cache factor. ++ ("VLLM_PLE_FP8_GLOBAL_SCALE", "1", True), + ("VLLM_E2_GUARD_INTERVAL", "0.25", "0.25"), + ("VLLM_E2_MEMSTATS", "2", "2"), + ("VLLM_E2_MEMSTATS_DIR", "/var/tmp/memstats", "/var/tmp/memstats"), +@@ -695,7 +698,15 @@ def test_unaudited_host_side_knob_stays_listed_until_it_is_classified(name: str) + exists to prevent, so it fails here. + """ + source = Path(envs.__file__).read_text(encoding="utf-8") +- listed = name in source ++ # Scan ONLY the delimited block. Scanning the whole file can never pass ++ # once a knob is legitimately registered, because the registry entry ++ # contains the knob's own name -- the gate would block the very act it ++ # exists to sequence. ++ start = source.index("# --8<-- [start:fork-unaudited-knobs]") ++ end = source.index("# --8<-- [end:fork-unaudited-knobs]") ++ assert start < end, "the unaudited-knob block markers are inverted" ++ block = source[start:end] ++ listed = name in block + if name in environment_variables: + assert not listed, ( + f"{name} is registered but still named in the envs.py TODO block: " +diff --git a/tests/v1/test_e2_guards.py b/tests/v1/test_e2_guards.py +new file mode 100644 +index 0000000..a3b6839 +--- /dev/null ++++ b/tests/v1/test_e2_guards.py +@@ -0,0 +1,215 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Focused cases for the fork's startup graph-profile guard. ++ ++An earlier version was rejected for accepting values upstream ++treats as off, rejecting valid ``CompilationMode.NONE`` boots, swallowing read ++errors, and missing the drafter. Each of those is a case below. ++ ++These use a plain stand-in config rather than a real ``VllmConfig``: the guard ++reads exactly three attributes, and building a real config needs a checkpoint. ++""" ++ ++from types import SimpleNamespace ++ ++import pytest ++ ++from vllm.config import CompilationMode, CUDAGraphMode ++from vllm.v1 import e2_guards ++from vllm.v1.e2_guards import ( ++ ALLOW_BREAKABLE_ENV, ++ BREAKABLE_ENV, ++ EXPECTED_MODE_ENV, ++ assert_graph_profile_is_executable, ++ breakable_resolved, ++) ++ ++ ++def _config( ++ mode: CompilationMode = CompilationMode.NONE, ++ cudagraph_mode: CUDAGraphMode = CUDAGraphMode.FULL, ++ architectures: tuple[str, ...] = ("Qwen4ExpForCausalLM",), ++): ++ return SimpleNamespace( ++ model_config=SimpleNamespace(architectures=list(architectures)), ++ compilation_config=SimpleNamespace(mode=mode, cudagraph_mode=cudagraph_mode), ++ ) ++ ++ ++@pytest.fixture(autouse=True) ++def _clean_env(monkeypatch): ++ for name in (BREAKABLE_ENV, ALLOW_BREAKABLE_ENV, EXPECTED_MODE_ENV): ++ monkeypatch.delenv(name, raising=False) ++ # Keep the envs accessor from dominating the raw reading in these cases. ++ monkeypatch.setattr(e2_guards, "breakable_resolved", breakable_resolved) ++ ++ ++# --------------------------------------------------------------- the parser ++ ++ ++@pytest.mark.parametrize("raw", ["0", "2", "true", "True", "FALSE", "", "01"]) ++def test_only_the_exact_string_one_is_on(monkeypatch, raw: str) -> None: ++ """Match envs.py:937 exactly. A stricter parser is a FALSE REJECTION.""" ++ monkeypatch.setenv(BREAKABLE_ENV, raw) ++ monkeypatch.setattr( ++ "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False ++ ) ++ assert breakable_resolved() is False ++ ++ ++def test_the_raw_reading_wins_when_envs_is_stale(monkeypatch) -> None: ++ """_maybe_enable_breakable_cudagraph mutates os.environ after envs caches.""" ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ monkeypatch.setattr( ++ "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False ++ ) ++ assert breakable_resolved() is True ++ ++ ++def test_the_envs_reading_wins_when_the_raw_value_is_absent(monkeypatch) -> None: ++ monkeypatch.delenv(BREAKABLE_ENV, raising=False) ++ monkeypatch.setattr( ++ "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", True, raising=False ++ ) ++ assert breakable_resolved() is True ++ ++ ++# ------------------------------------------------------------- the rejections ++ ++ ++def test_an_unpinned_launch_fails(monkeypatch) -> None: ++ """Absence is the state the upstream auto-enable feeds on.""" ++ with pytest.raises(RuntimeError, match="is not set"): ++ assert_graph_profile_is_executable(_config()) ++ ++ ++def test_breakable_on_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ with pytest.raises(RuntimeError, match="resolved ON"): ++ assert_graph_profile_is_executable(_config()) ++ ++ ++def test_piecewise_with_vllm_compile_and_breakable_off_fails(monkeypatch) -> None: ++ """The BLOCKER: upstream accepts this at config time and dies at capture.""" ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ with pytest.raises(RuntimeError, match="cudagraph_utils.py:605"): ++ assert_graph_profile_is_executable( ++ _config( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ ) ++ ) ++ ++ ++def test_expected_mode_mismatch_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL_DECODE_ONLY") ++ with pytest.raises(RuntimeError, match="EXPECTED_CUDAGRAPH_MODE"): ++ assert_graph_profile_is_executable(_config(cudagraph_mode=CUDAGraphMode.FULL)) ++ ++ ++# ------------------------------------------------------------- the permissions ++ ++ ++def test_compilation_none_with_full_graphs_is_permitted(monkeypatch) -> None: ++ """FALSE REJECTION control. ++ ++ On v0.30.0's NVIDIA Qwen4Exp path NONE is a normal resolved state -- #55272 ++ removed the compile wrappers -- so treating NONE as evidence of the ++ auto-enable, which the first version of this guard did, rejects the one ++ profile that is actually executable. ++ """ ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ assert_graph_profile_is_executable( ++ _config(mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.FULL) ++ ) ++ ++ ++def test_none_plus_full_decode_only_is_permitted(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ assert_graph_profile_is_executable( ++ _config( ++ mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY ++ ) ++ ) ++ ++ ++def test_eager_is_permitted(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ assert_graph_profile_is_executable( ++ _config(mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.NONE) ++ ) ++ ++ ++def test_expected_mode_match_passes(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL") ++ assert_graph_profile_is_executable(_config(cudagraph_mode=CUDAGraphMode.FULL)) ++ ++ ++# ------------------------------------------------------------ fail-closed paths ++ ++ ++def test_a_missing_compilation_config_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ cfg = SimpleNamespace( ++ model_config=SimpleNamespace(architectures=["Qwen4ExpForCausalLM"]), ++ compilation_config=None, ++ ) ++ with pytest.raises(RuntimeError, match="no compilation_config"): ++ assert_graph_profile_is_executable(cfg) ++ ++ ++def test_a_missing_model_config_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ cfg = SimpleNamespace(model_config=None, compilation_config=SimpleNamespace()) ++ with pytest.raises(RuntimeError, match="no model_config"): ++ assert_graph_profile_is_executable(cfg) ++ ++ ++def test_an_unreadable_switch_fails_rather_than_reporting_off(monkeypatch) -> None: ++ """A guard that cannot read its own input must not answer False.""" ++ monkeypatch.delenv(BREAKABLE_ENV, raising=False) ++ ++ class _Boom: ++ def __getattr__(self, name): ++ raise RuntimeError("envs unavailable") ++ ++ monkeypatch.setitem(__import__("sys").modules, "vllm.envs", _Boom()) ++ with pytest.raises(RuntimeError): ++ breakable_resolved() ++ ++ ++# ------------------------------------------------------------------ the opt-out ++ ++ ++def test_the_opt_out_is_loud_and_skips_the_checks(monkeypatch, caplog) -> None: ++ monkeypatch.setenv(ALLOW_BREAKABLE_ENV, "1") ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ with caplog.at_level("WARNING"): ++ assert_graph_profile_is_executable( ++ _config( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ ) ++ ) ++ assert ALLOW_BREAKABLE_ENV in caplog.text ++ # The opt-out log line must say it is not a promotion (production) boot. ++ assert "promotion" in caplog.text.lower() ++ ++ ++# -------------------------------------------------------------- the draft role ++ ++ ++def test_the_draft_role_is_checked_independently(monkeypatch) -> None: ++ """Qwen4ExpMTP builds Qwen4ExpMultiTokenPredictor, not CausalLM.""" ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ with pytest.raises(RuntimeError, match=r"e2-guard\[draft\]"): ++ assert_graph_profile_is_executable( ++ _config( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.PIECEWISE, ++ architectures=("Qwen4ExpMTP",), ++ ), ++ role="draft", ++ ) +diff --git a/vllm/envs.py b/vllm/envs.py +index 588ca59..734a95c 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -333,11 +333,29 @@ if TYPE_CHECKING: + # hashed anyway: a spurious cache miss costs a recompile, a missed one costs + # a wrong graph. + # +- # COMPILE/CAPTURE FACTORS (hashed): +- # AOT: yes. CAPTURE: yes. PLE table is FP8-with-global-scale -> +- # Qwen4ExpPLEFp8EmbeddingMethod instead of the checkpoint's method: +- # different weight dtype and a different embedding forward. ++ # *** INERT ON THE 0.30 BASE. *** ++ # NO READ SITE, and therefore NOT HASHED -- moved out of the compile ++ # factors into ignored_factors below. Its only reader was the 21-line ++ # catch-all branch in _get_ple_embedding_quant_method, which lived in the ++ # half of models/qwen4_exp/nvidia/ple_layer.py that #54371 deleted. With no ++ # reader it cannot change an artifact, so hashing it would only invalidate ++ # caches for nothing; and on this uncompiled NVIDIA path (#55272) there may ++ # be no AOT artifact to key at all. It gives NO capture protection either. ++ # v0.30.0 selects the FP8 PLE method from the CONFIG KEY ++ # text_config.ple_embedding_dtype == "float8_e4m3fn" ++ # (ngram_embedding.py:174-175) -- see ++ # tests/models/qwen4_exp/test_ple_fp8_selector_checkpoint.py. ++ # Setting this on the 0.30 base does NOTHING, and NOTHING WARNS: validate_environ ++ # (envs.py, below) warns only about names that are NOT registered, and this ++ # comment lives in a TYPE_CHECKING block that never runs. It is kept ++ # registered so a later PLE transport change can reland its reader and ++ # reclassify it, or delete it. DO NOT assume it works. ++ # When its reader returns: AOT yes, CAPTURE yes -- the FP8-with-global-scale ++ # table selects Qwen4ExpPLEFp8EmbeddingMethod over the checkpoint's method, ++ # a different weight dtype and a different embedding forward. + VLLM_PLE_FP8_GLOBAL_SCALE: bool = False ++ # ++ # COMPILE/CAPTURE FACTORS (hashed): + # CAPTURE: yes, established -- clip counting allocates persistent per-layer + # int64 counters and its increment ops are CAPTURED into the cudagraph (see + # the graph-safety note in models/qwen4_exp/nvidia/qsa.py); unset, those ops +@@ -398,16 +416,34 @@ if TYPE_CHECKING: + # cross-change knob table). Each mirrors its read site's + # parse; if a read site changes, the entry here must be reconciled. + # (Both read sites below are now live.) ++ # *** INERT ON THE 0.30 BASE: NO READ SITE. *** ++ # Its reader is the slot-ownership assertion code (an earlier commit), which edits ++ # vllm/v1/ple_offload/connector.py -- a module that does not exist on ++ # v0.30.0. Setting this on the 0.30 base does NOTHING, and NOTHING WARNS ++ # (validate_environ only reports UNregistered names). Already an ignored ++ # factor, correctly. Kept registered so the contract is recorded; it ++ # becomes live again only together with the pull transport ++ # (the PLE transport choice). + # AOT: no. CAPTURE: no -- strictly host-side and outside capture. + # off|warn|raise (1 = warn), normalised by _e2_diagnostic_mode like the + # rest of the family. The mamba ownership check moved this gate off a plain + # boolean; reconciled here. +- # Read site returns with the deferred PLE transport. + VLLM_E2_PLE_FLAG_CHECK: str = "off" + # AOT: no. CAPTURE: no -- strictly host-side. off|warn|raise (1 = warn). + # the mamba slot-ownership check, v1/worker/mamba_utils.py:74. + VLLM_E2_MAMBA_SLOT_OWNERSHIP: str = "off" ++ # Opt-out for the fork startup guards in vllm/v1/e2_guards.py, read at ++ # Qwen4ExpForCausalLM construction. AOT: no. CAPTURE: no -- it changes ++ # nothing itself; it only permits a boot whose identity IS different ++ # (VLLM_USE_BREAKABLE_CUDAGRAPH, which upstream already registers and ++ # hashes). Ignored factor for exactly that reason. ++ VLLM_E2_ALLOW_BREAKABLE: bool = False ++ # Expected RESOLVED cudagraph mode by name, asserted by the same guard when ++ # set (e.g. "FULL", "FULL_DECODE_ONLY"). AOT: no. CAPTURE: no -- it asserts ++ # an identity, it does not choose one. Ignored factor. ++ VLLM_E2_EXPECTED_CUDAGRAPH_MODE: str = "" + # ++ # --8<-- [start:fork-unaudited-knobs] + # --- TODO: knobs that are NOT in this tree at all ---------- + # The tie-determinism toggles and the BSM pin live in a separate + # patch stack, not in v2.0.1 and not in this tree: +@@ -428,8 +464,11 @@ if TYPE_CHECKING: + # identities from its actual read site, hash it unless exclusion is proven, + # and delete it from this block. tests/test_envs.py holds a test that fails + # if a name is registered while still listed here, so the classification +- # cannot be skipped silently. ++ # cannot be skipped silently. That test reads ONLY the delimited block ++ # below/above -- scanning the whole file could never pass, because a ++ # registry entry necessarily contains its own name. + # ---------------------------------------------------------------------- ++ # --8<-- [end:fork-unaudited-knobs] + VLLM_LOG_MODEL_INSPECTION: bool = False + VLLM_DEBUG_MFU_METRICS: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False +@@ -2209,6 +2248,10 @@ environment_variables: dict[str, Callable[[], Any]] = { + # Capture-configuration factor, read in C++; raw passthrough on purpose. + "MERLIN_TOPK_EXACT_MEDIUM": lambda: os.getenv("MERLIN_TOPK_EXACT_MEDIUM", None), + # Runtime only; all of the following are in compile_factors.ignored_factors. ++ "VLLM_E2_ALLOW_BREAKABLE": lambda: os.getenv("VLLM_E2_ALLOW_BREAKABLE", "") == "1", ++ "VLLM_E2_EXPECTED_CUDAGRAPH_MODE": lambda: os.getenv( ++ "VLLM_E2_EXPECTED_CUDAGRAPH_MODE", "" ++ ).strip(), + "VLLM_QSA_KV_SCALES": lambda: os.getenv("VLLM_QSA_KV_SCALES", "").strip(), + "VLLM_E2_GUARD_MODE": lambda: ( + os.getenv("VLLM_E2_GUARD_MODE", "off").strip().lower() +@@ -2592,6 +2635,14 @@ def compile_factors() -> dict[str, object]: + "VLLM_PLE_TEST_FAULT", + "VLLM_PLE_TEST_FAULT_GATE", + "VLLM_PLE_TEST_FAULT_LOG", ++ # Fork startup guard: the opt-out permits a different identity and the ++ # expected-mode pin asserts one; neither IS one. ++ "VLLM_E2_ALLOW_BREAKABLE", ++ "VLLM_E2_EXPECTED_CUDAGRAPH_MODE", ++ # INERT on the 0.30 base: no read site, so it cannot change any artifact. ++ # Hashing it would invalidate caches for no behavioural change. Restore ++ # it to the compile factors in the same commit that relands its reader. ++ "VLLM_PLE_FP8_GLOBAL_SCALE", + # Diagnostic knobs, both strictly host-side and outside capture, so + # neither AOT nor capture identity depends on them. Hashing + # them would invalidate every compile cache the moment a diagnostic +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 816a4e2..8c29268 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -78,6 +78,7 @@ from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum ++from vllm.v1.e2_guards import assert_graph_profile_is_executable + from vllm.v1.kv_cache_interface import MambaSpec + + from ..config import Qwen4ExpConfig +@@ -786,6 +787,13 @@ class Qwen4ExpForCausalLM( + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() ++ # V0.30.0 auto-enables breakable cudagraphs ++ # for this architecture when VLLM_USE_BREAKABLE_CUDAGRAPH is unset, and ++ # VLLM_COMPILE + a PIECEWISE mode dies at capture because #55272 left no ++ # compiled submodule. Fail here rather than let a reference-host run qualify an ++ # identity nobody chose. Qwen4ExpForConditionalGeneration constructs ++ # this class; Qwen4ExpMTP does not, and guards itself. ++ assert_graph_profile_is_executable(vllm_config, role="target") + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 2df3868..468fc3a 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -49,6 +49,7 @@ from vllm.sequence import IntermediateTensors + from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) ++from vllm.v1.e2_guards import assert_graph_profile_is_executable + + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +@@ -395,6 +396,10 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ++ # The drafter is NOT reached through Qwen4ExpForCausalLM -- ++ # this class builds Qwen4ExpMultiTokenPredictor directly, from a config ++ # rebuilt at mtp.py:159 -- so its graph profile is checked here. ++ assert_graph_profile_is_executable(vllm_config, role="draft") + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + cache_config = vllm_config.cache_config +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +new file mode 100644 +index 0000000..9971ccd +--- /dev/null ++++ b/vllm/v1/e2_guards.py +@@ -0,0 +1,218 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Fork startup guards that fail closed on a graph profile we never qualified. ++ ++These run at model construction, on every rank, for the target model and for ++the MTP drafter, before profiling or capture. They exist because a *documented* ++launch flag is not a mechanism: the first reference-host run of a rebase will otherwise ++qualify a compile/capture identity nobody chose ++(see below). ++ ++What this module does NOT do: choose a graph profile. That is the maintainer's ++call and it is open (the graph-profile choice). What it does is refuse the ++two states that are silently wrong on vLLM v0.30.0 -- ++ ++ * ``VLLM_USE_BREAKABLE_CUDAGRAPH`` left unset, which lets ++ ``VllmConfig._maybe_enable_breakable_cudagraph`` turn it on for this fork's ++ architectures and force ``CompilationMode.NONE``; and ++ * a PIECEWISE-bearing cudagraph mode combined with ``VLLM_COMPILE`` and ++ breakable-off, which upstream accepts at config time and then rejects at ++ capture, because #55272 removed every ``@support_torch_compile`` wrapper ++ from the NVIDIA Qwen4Exp model and MTP. ++ ++-- and to make the resolved profile visible and assertable. ++ ++Timing limit, stated honestly: this runs at model construction, i.e. AFTER ++``VllmConfig.__post_init__`` has already called ++``_maybe_enable_breakable_cudagraph`` (``config/vllm.py:1561``) and ++``_normalize_piecewise_cudagraph_mode`` (``:1673``, ``:1844``). It is early ++enough to stop profiling and capture; it cannot un-mutate ``os.environ`` or ++recover the originally requested cudagraph mode. Requiring the pin to be ++*explicitly present* is the part of "validate before auto-enable" that is ++checkable from here: absence is exactly the state the auto-enable feeds on. ++""" ++ ++import os ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++BREAKABLE_ENV = "VLLM_USE_BREAKABLE_CUDAGRAPH" ++#: Opt-out for every guard in this module. Recorded in the launch identity and ++#: PROHIBITED in production profiles; a production launch must not set it. ++ALLOW_BREAKABLE_ENV = "VLLM_E2_ALLOW_BREAKABLE" ++#: Optional expected resolved cudagraph mode, by name (e.g. "FULL", ++#: "FULL_DECODE_ONLY"). When set, a mismatch is an error. This is how a launch ++#: profile pins the graph identity once the maintainer has chosen one. ++EXPECTED_MODE_ENV = "VLLM_E2_EXPECTED_CUDAGRAPH_MODE" ++ ++#: The architectures vLLM v0.30.0 auto-enables breakable cudagraphs for, out of ++#: ``vllm/config/vllm.py:DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES``, that are ++#: ours. Recorded so the guard's message can name the upstream mechanism. ++OUR_AUTO_BREAKABLE_ARCHITECTURES = frozenset( ++ { ++ "Qwen4ExpForCausalLM", ++ "Qwen4ExpForConditionalGeneration", ++ "Qwen4ExpMTP", ++ } ++) ++ ++ ++def _upstream_truth(raw: str | None) -> bool: ++ """Upstream's exact semantics for the breakable switch. ++ ++ ``envs.py:937`` is ``os.environ.get(..., "0") == "1"``. Nothing else is on: ++ ``"2"``, ``"true"``, ``"FALSE"`` and ``""`` are all off. Using one parser ++ that matches it is the point -- a stricter local parser would reject boots ++ upstream treats as off, which is a false rejection, not extra safety. ++ """ ++ return raw == "1" ++ ++ ++def breakable_resolved() -> bool: ++ """Whether breakable cudagraphs are on, as the engine will see it. ++ ++ Two readings of ONE source, not three independent sources: ++ ``vllm.compilation.breakable_cudagraph.is_breakable_cudagraph_enabled()`` ++ is literally ``bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)`` ++ (``breakable_cudagraph.py:53``), so it adds nothing over the envs read. ++ The envs read is kept beside the raw one because ++ ``_maybe_enable_breakable_cudagraph`` mutates ``os.environ`` at config time ++ while the envs accessor is cached, so the two can legitimately disagree; ++ either saying "on" means on. ++ ++ No exception is swallowed. A guard that cannot read its own input must fail ++ the boot, not report False. ++ """ ++ if _upstream_truth(os.environ.get(BREAKABLE_ENV)): ++ return True ++ import vllm.envs as envs ++ ++ return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH) ++ ++ ++def _profile(vllm_config) -> tuple[list[str], object, object]: ++ """Resolved (architectures, compilation mode, cudagraph mode). Fails closed.""" ++ model_config = getattr(vllm_config, "model_config", None) ++ if model_config is None: ++ raise RuntimeError("e2-guard: vllm_config has no model_config to check") ++ compilation_config = getattr(vllm_config, "compilation_config", None) ++ if compilation_config is None: ++ raise RuntimeError("e2-guard: vllm_config has no compilation_config to check") ++ architectures = list(getattr(model_config, "architectures", None) or []) ++ return ( ++ architectures, ++ compilation_config.mode, ++ compilation_config.cudagraph_mode, ++ ) ++ ++ ++def assert_graph_profile_is_executable(vllm_config, *, role: str = "target") -> None: ++ """Refuse the two silently-wrong startup states on v0.30.0. ++ ++ ``role`` is ``"target"`` or ``"draft"``; both are checked, because ++ ``Qwen4ExpMTP`` builds ``Qwen4ExpMultiTokenPredictor`` rather than ++ ``Qwen4ExpForCausalLM`` (``mtp.py:397``, ``:411``) from a config rebuilt at ++ ``mtp.py:159``, so the drafter's profile can differ from the target's. ++ ++ Intentional eager and intentional ``CompilationMode.NONE`` are PERMITTED. ++ On v0.30.0's NVIDIA Qwen4Exp path ``NONE`` is a normal resolved state, not ++ evidence of the auto-enable: #55272 removed the compile wrappers, so ++ ``_normalize_piecewise_cudagraph_mode`` legitimately lands there. ++ """ ++ from vllm.config import CompilationMode ++ architectures, mode, cudagraph_mode = _profile(vllm_config) ++ ++ if os.environ.get(ALLOW_BREAKABLE_ENV, "") == "1": ++ logger.warning( ++ "e2-guard[%s]: %s=1 -- guards disabled by request. architectures=%s " ++ "breakable=%s compilation_mode=%s cudagraph_mode=%s. This boot is " ++ "NOT a promotion-grade boot; a production launch must not set this flag.", ++ role, ++ ALLOW_BREAKABLE_ENV, ++ architectures, ++ os.environ.get(BREAKABLE_ENV), ++ mode, ++ cudagraph_mode, ++ ) ++ return ++ ++ # 1. The pin must be explicit. Absence is the state _maybe_enable_breakable_ ++ # cudagraph feeds on (config/vllm.py:758: "not in os.environ"). ++ if BREAKABLE_ENV not in os.environ: ++ raise RuntimeError( ++ f"e2-guard[{role}]: {BREAKABLE_ENV} is not set. vLLM v0.30.0 lists " ++ f"{sorted(OUR_AUTO_BREAKABLE_ARCHITECTURES)} in " ++ "DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES and turns the switch on " ++ "by itself when it is unset (vllm/config/vllm.py:756-774), forcing " ++ "CompilationMode.NONE with it. Pin it in the launcher: " ++ f"export {BREAKABLE_ENV}=0 (or =1 deliberately), or set " ++ f"{ALLOW_BREAKABLE_ENV}=1 to take an unpinned boot on purpose." ++ ) ++ ++ breakable = breakable_resolved() ++ ++ # 2. Breakable on is an identity this fork has never qualified. ++ if breakable: ++ raise RuntimeError( ++ f"e2-guard[{role}]: {BREAKABLE_ENV} resolved ON for " ++ f"{architectures or ''} " ++ f"(compilation_mode={mode}, cudagraph_mode={cudagraph_mode}). That " ++ "is not the identity the v2.0.1 results were measured on. Export " ++ f"{BREAKABLE_ENV}=0, or set {ALLOW_BREAKABLE_ENV}=1 to take it " ++ "deliberately." ++ ) ++ ++ # 3. PIECEWISE with breakable off needs a torch-compiled submodule, and ++ # #55272 removed every @support_torch_compile wrapper from ++ # vllm/models/qwen4_exp/nvidia/*. Upstream only rewrites the mode when ++ # compilation mode is NOT VLLM_COMPILE ++ # (config/vllm.py:1010-1030), so VLLM_COMPILE + PIECEWISE + breakable-off ++ # survives config time and then dies at ++ # gpu/cudagraph_utils.py:605 in the middle of capture. Fail here instead, ++ # with the candidates named. ++ if ( ++ cudagraph_mode.requires_piecewise_compilation() ++ and mode is CompilationMode.VLLM_COMPILE ++ ): ++ raise RuntimeError( ++ f"e2-guard[{role}]: cudagraph_mode={cudagraph_mode} needs piecewise " ++ "capture, compilation mode is VLLM_COMPILE and breakable cudagraphs " ++ "are off -- but vLLM PR #55272 removed every @support_torch_compile " ++ "wrapper from vllm/models/qwen4_exp/nvidia/, so " ++ "has_compiled_submodule() is False and capture will raise at " ++ "vllm/v1/worker/gpu/cudagraph_utils.py:605. Executable candidates: " ++ "(a) breakable off + -O0/CompilationMode.NONE, which normalises " ++ "FULL_AND_PIECEWISE to FULL_DECODE_ONLY (changed prefill " ++ "behaviour -- must be qualified); (b) breakable off + an explicit " ++ "cudagraph_mode=FULL or NONE; (c) breakable PIECEWISE, i.e. " ++ f"{BREAKABLE_ENV}=1, which is a different identity. See " ++ "the graph-profile choice." ++ ) ++ ++ # 4. Optional exact pin of the resolved mode, for a launch profile. ++ expected = os.environ.get(EXPECTED_MODE_ENV, "").strip() ++ if expected: ++ actual = getattr(cudagraph_mode, "name", str(cudagraph_mode)) ++ if actual != expected: ++ raise RuntimeError( ++ f"e2-guard[{role}]: resolved cudagraph_mode={actual} but " ++ f"{EXPECTED_MODE_ENV}={expected}. The launch profile and the " ++ "engine disagree about the graph identity." ++ ) ++ ++ logger.info( ++ "e2-guard[%s]: architectures=%s breakable=OFF compilation_mode=%s " ++ "cudagraph_mode=%s expected=%s. Recorded, not certified -- this guard " ++ "checks startup states, not that the profile was qualified.", ++ role, ++ architectures, ++ mode, ++ cudagraph_mode, ++ expected or "", ++ ) ++ ++ ++#: Back-compatible name for the earlier call sites. ++assert_breakable_cudagraph_pinned = assert_graph_profile_is_executable +diff --git a/vllm/v1/ple_offload/test_fault.py b/vllm/v1/ple_offload/test_fault.py +deleted file mode 100644 +index 9ae74f7..0000000 +--- a/vllm/v1/ple_offload/test_fault.py ++++ /dev/null +@@ -1,13 +0,0 @@ +-# SPDX-License-Identifier: Apache-2.0 +-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +-"""Compatibility exports for the top-level test-only fault singleton.""" +- +-from vllm_ple_test_fault import ( +- _FAULT as _FAULT, +-) +-from vllm_ple_test_fault import ( # noqa: F401 +- FaultSpec as FaultSpec, +-) +-from vllm_ple_test_fault import ( +- _parse_fault_spec as _parse_fault_spec, +-) +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 949f299..09f91f9 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -231,118 +231,6 @@ class Worker(WorkerBase): + # Resolved lazily on first sleep/wake; persists worker-process state. + self._sleep_mode_backend: SleepModeBackend | None = None + +- def _has_ple_layers(self) -> bool: +- """Return whether this model configuration constructs PLE layers.""" +- if not envs.VLLM_PLE_CPU_OFFLOAD: +- return False +- text_config = self.model_config.hf_text_config +- ple_layer_ids = getattr(text_config, "ple_layer_ids", None) +- if not ple_layer_ids: +- return False +- # PP (vllm#54709, carried from the earlier overlay): only the pipeline rank +- # that owns a PLE layer takes part in the offload (rank 0 with the +- # default partition). Other ranks build PPMissingLayer there and must +- # skip the connector entirely. +- pp_size = self.parallel_config.pipeline_parallel_size +- if pp_size > 1: +- from vllm.distributed.utils import get_pp_indices +- +- tp_size = self.parallel_config.tensor_parallel_size +- pp_rank = (self.rank // tp_size) % pp_size +- start, end = get_pp_indices( +- int(text_config.num_hidden_layers), pp_rank, pp_size +- ) +- return any(start <= int(i) - 1 < end for i in ple_layer_ids) +- return True +- +- def _validate_ple_offload_config(self) -> None: +- """Reject unsupported PLE offload execution modes.""" +- parallel_config = self.parallel_config +- unsupported = [] +- if not self.use_v2_model_runner: +- unsupported.append("model runner V1") +- if not current_platform.is_cuda(): +- unsupported.append(f"device={current_platform.device_type}") +- if parallel_config.nnodes != 1: +- unsupported.append(f"nnodes={parallel_config.nnodes}") +- if parallel_config.data_parallel_backend != "mp": +- unsupported.append(f"DP backend={parallel_config.data_parallel_backend}") +- if ( +- parallel_config.data_parallel_size_local +- != parallel_config.data_parallel_size +- ): +- unsupported.append( +- "non-local DP " +- f"({parallel_config.data_parallel_size_local}/" +- f"{parallel_config.data_parallel_size} local ranks)" +- ) +- # PP allowed (vllm#54709): PLE layers must live on pipeline rank 0 +- # (enforced placement-based in models/config.py); num_workers stays +- # dp_size * tp_size = the ranks of the owning stage. +- if parallel_config.prefill_context_parallel_size != 1: +- unsupported.append(f"PCP={parallel_config.prefill_context_parallel_size}") +- if parallel_config.decode_context_parallel_size != 1: +- unsupported.append(f"DCP={parallel_config.decode_context_parallel_size}") +- if parallel_config.use_ubatching: +- unsupported.append("ubatching/DBO") +- if self.model_config.architecture not in { +- "Qwen4ExpForCausalLM", +- "Qwen4ExpForConditionalGeneration", +- }: +- unsupported.append("architecture") +- if self.vllm_config.weight_transfer_config is not None: +- unsupported.append("weight transfer") +- +- if unsupported: +- raise ValueError( +- "VLLM_PLE_CPU_OFFLOAD does not support the requested " +- "configuration. Unsupported settings: " +- f"{', '.join(unsupported)}" +- ) +- +- def spawn_ple_offload(self) -> None: +- """Spawn one node-local PLE CPU worker from DP0/TP0.""" +- if ( +- not self._ple_offload_enabled +- or self.rank != 0 +- or self.parallel_config.data_parallel_rank != 0 +- ): +- return +- +- from vllm.v1.ple_offload.worker import PleOffloadWorker +- +- ipc_addr = self.parallel_config._ple_offload_ipc_path +- if not ipc_addr: +- raise RuntimeError("PLE offload IPC address was not initialized") +- dp_size = self.parallel_config.data_parallel_size +- tp_size = self.parallel_config.tensor_parallel_size +- num_workers = dp_size * tp_size +- logger.info( +- "PleOffload: spawning worker " +- "(rank=%d, local_rank=%d, dp_size=%d, tp_size=%d, " +- "num_workers=%d, ipc_addr=%s).", +- self.rank, +- self.local_rank, +- dp_size, +- tp_size, +- num_workers, +- ipc_addr, +- ) +- self._ple_offload_worker_handle = PleOffloadWorker.make_process( +- self.vllm_config, +- num_workers, +- ipc_addr, +- ) +- +- def wait_ple_offload_ready(self) -> None: +- """Wait until weights and every TP registration are ready.""" +- if self._ple_offload_worker_handle is None: +- return +- +- from vllm.v1.ple_offload.worker import PleOffloadWorker +- +- PleOffloadWorker.wait_for_ready(self._ple_offload_worker_handle) +- + @property + def sleep_mode_backend(self) -> "SleepModeBackend": + if self._sleep_mode_backend is None: +-- +2.43.0 + diff --git a/patches/0031-ple_offload-carry-the-host-mapped-pull-transport-ont.patch b/patches/0031-ple_offload-carry-the-host-mapped-pull-transport-ont.patch new file mode 100644 index 0000000..42b6e4d --- /dev/null +++ b/patches/0031-ple_offload-carry-the-host-mapped-pull-transport-ont.patch @@ -0,0 +1,7871 @@ +From 4c793b234fff0b909893dfdbcccd78b98d6fc4e0 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 00:01:58 -0400 +Subject: [PATCH 031/103] ple_offload: carry the host-mapped pull transport + onto v0.30.0 + +Restores the fork's PLE offload package (not upstream) and its pull +transport: the offload process holds no CUDA context; the captured op waits on +a device-mapped host flag and copies host->GPU on the model stream; publication +is fenced; releases are synchronous across dummy/capture work; failures close +the rank. Includes the test-only fault-injection cuts for the transport. +--- + probe/faultcut_gate.py | 46 + + tests/v1/ple_offload/test_correspondence.py | 197 +++ + tests/v1/ple_offload/test_fatal_rpc.py | 145 ++ + .../test_fault_bootstrap_disarmed.py | 80 + + .../ple_offload/test_fault_injector_basics.py | 260 ++++ + .../ple_offload/test_fault_injector_gates.py | 468 ++++++ + .../v1/ple_offload/test_fault_wire_frames.py | 96 ++ + .../v1/ple_offload/test_ipc_socket_unlink.py | 133 ++ + tests/v1/ple_offload/test_pending_rpc.py | 81 + + .../v1/ple_offload/test_progress_deadline.py | 216 +++ + tests/v1/ple_offload/test_shm_unlink.py | 397 +++++ + tests/v1/ple_offload/test_shutdown_order.py | 171 +++ + .../ple_offload/test_spawn_bootstrap_cuda.py | 352 +++++ + tests/v1/ple_offload/test_test_fault.py | 293 ++++ + .../ple_offload/test_transport_bounded_io.py | 343 +++++ + .../test_transport_fatal_helper.py | 387 +++++ + .../layers/ple_offload_layer.py | 433 ++++++ + vllm/v1/executor/multiproc_executor.py | 31 +- + vllm/v1/ple_offload/__init__.py | 2 + + vllm/v1/ple_offload/connector.py | 855 +++++++++++ + vllm/v1/ple_offload/fatal.py | 153 ++ + vllm/v1/ple_offload/progress.py | 314 ++++ + vllm/v1/ple_offload/protocol.py | 50 + + vllm/v1/ple_offload/worker.py | 1298 +++++++++++++++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 140 +- + vllm/v1/worker/gpu/model_runner.py | 140 ++ + vllm/v1/worker/gpu_worker.py | 145 ++ + vllm_ple_offload_bootstrap/__init__.py | 210 +++ + vllm_ple_offload_bootstrap/__main__.py | 7 + + 29 files changed, 7429 insertions(+), 14 deletions(-) + create mode 100755 probe/faultcut_gate.py + create mode 100644 tests/v1/ple_offload/test_correspondence.py + create mode 100644 tests/v1/ple_offload/test_fatal_rpc.py + create mode 100644 tests/v1/ple_offload/test_fault_bootstrap_disarmed.py + create mode 100644 tests/v1/ple_offload/test_fault_injector_basics.py + create mode 100644 tests/v1/ple_offload/test_fault_injector_gates.py + create mode 100644 tests/v1/ple_offload/test_fault_wire_frames.py + create mode 100644 tests/v1/ple_offload/test_ipc_socket_unlink.py + create mode 100644 tests/v1/ple_offload/test_pending_rpc.py + create mode 100644 tests/v1/ple_offload/test_progress_deadline.py + create mode 100644 tests/v1/ple_offload/test_shm_unlink.py + create mode 100644 tests/v1/ple_offload/test_shutdown_order.py + create mode 100644 tests/v1/ple_offload/test_spawn_bootstrap_cuda.py + create mode 100644 tests/v1/ple_offload/test_test_fault.py + create mode 100644 tests/v1/ple_offload/test_transport_bounded_io.py + create mode 100644 tests/v1/ple_offload/test_transport_fatal_helper.py + create mode 100644 vllm/model_executor/layers/ple_offload_layer.py + create mode 100644 vllm/v1/ple_offload/__init__.py + create mode 100644 vllm/v1/ple_offload/connector.py + create mode 100644 vllm/v1/ple_offload/fatal.py + create mode 100644 vllm/v1/ple_offload/progress.py + create mode 100644 vllm/v1/ple_offload/protocol.py + create mode 100644 vllm/v1/ple_offload/worker.py + create mode 100644 vllm_ple_offload_bootstrap/__init__.py + create mode 100644 vllm_ple_offload_bootstrap/__main__.py + +diff --git a/probe/faultcut_gate.py b/probe/faultcut_gate.py +new file mode 100755 +index 0000000..3eb74b5 +--- /dev/null ++++ b/probe/faultcut_gate.py +@@ -0,0 +1,46 @@ ++#!/usr/bin/env python3 ++"""CPU-testable validation for external gates excluded from GPU campaigns.""" ++ ++from collections.abc import Callable ++ ++ ++def validate_partial_send_gate( ++ *, ++ pid: int, ++ expected_start: int, ++ observed_start: int, ++ alive: bool, ++ bytes_sent: int, ++ frame_size: int, ++ completed: bool, ++ child_not_reading: bool, ++) -> tuple[bool, str]: ++ if pid <= 1 or not alive: ++ return False, "target is stale or not alive" ++ if expected_start <= 0 or observed_start != expected_start: ++ return False, "target PID was reused" ++ if frame_size <= 0 or not 0 < bytes_sent < frame_size: ++ return False, "send is zero-byte or complete" ++ if completed: ++ return False, "send already completed" ++ if not child_not_reading: ++ return False, "child non-reading state was not acknowledged" ++ return True, "valid partial send" ++ ++ ++def deliver_checked_signal( ++ *, ++ pid: int, ++ expected_start: int, ++ read_start: Callable[[int], int], ++ alive: Callable[[int], bool], ++ send: Callable[[int, int], bool], ++ signal_number: int, ++) -> tuple[bool, str]: ++ if pid <= 1 or not alive(pid): ++ return False, "target is stale or not alive" ++ if read_start(pid) != expected_start: ++ return False, "target PID was reused" ++ if not send(pid, signal_number): ++ return False, "signal delivery failed" ++ return True, "signal delivered" +diff --git a/tests/v1/ple_offload/test_correspondence.py b/tests/v1/ple_offload/test_correspondence.py +new file mode 100644 +index 0000000..e8d8b9e +--- /dev/null ++++ b/tests/v1/ple_offload/test_correspondence.py +@@ -0,0 +1,197 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++from types import SimpleNamespace ++from typing import Any ++ ++import pytest ++import torch ++ ++from vllm.config import CUDAGraphMode ++from vllm.model_executor.layers.ple_offload_layer import ( ++ _record_ple_copy_geometry, ++) ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload.protocol import PleOffloadRequest ++from vllm.v1.ple_offload.worker import ( ++ PleOffloadInputBuffers, ++ PleOffloadRunner, ++) ++from vllm.v1.worker.gpu.cudagraph_utils import ( ++ BatchExecutionDescriptor, ++ CudaGraphManager, ++) ++from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( ++ SpeculatorCudaGraphManager, ++) ++from vllm.v1.worker.gpu.model_runner import GPUModelRunner ++ ++ ++class _ConnectorForValidation: ++ _input_ids_buf = torch.empty(4, dtype=torch.int32) ++ _query_start_loc_buf = torch.empty(3, dtype=torch.int32) ++ _rowdump_n = None ++ ++ def _signal_dummy_outputs_impl(self, num_tokens: int) -> None: ++ raise AssertionError("must fail before dummy publication") ++ ++ def _launch(self, num_reqs: int, num_tokens: int) -> None: ++ raise AssertionError("must fail before request publication") ++ ++ ++def test_request_consumer_count_mismatch_precedes_publication() -> None: ++ connector = _ConnectorForValidation() ++ with pytest.raises(RuntimeError, match="request/consumer"): ++ PleOffloadConnector._prepare_forward_impl( ++ connector, ++ 1, ++ 3, ++ False, ++ consumer_num_tokens=4, ++ ) ++ ++ ++def _runner(query_start_loc: list[int]) -> PleOffloadRunner: ++ runner = PleOffloadRunner.__new__(PleOffloadRunner) ++ flag = torch.zeros(1, dtype=torch.int32) ++ target = SimpleNamespace( ++ host_output_buffer=torch.zeros(4, 2), ++ host_flag=flag, ++ ) ++ runner._worker_targets = {0: {"ple": [target]}} ++ runner._scratch_bufs = {0: {"ple": torch.empty(4, 2)}} ++ query = torch.zeros(3, dtype=torch.int32) ++ query[: len(query_start_loc)] = torch.tensor( ++ query_start_loc, dtype=torch.int32 ++ ) ++ runner._input_bufs = { ++ 0: PleOffloadInputBuffers( ++ input_ids_buf=torch.zeros(4, dtype=torch.int32), ++ query_start_loc_buf=query, ++ ngram_context_buf=torch.zeros(2, 1, dtype=torch.int32), ++ ) ++ } ++ return runner ++ ++ ++@pytest.mark.parametrize( ++ ("ple_request", "query_start_loc", "match"), ++ [ ++ (PleOffloadRequest(dp_rank=0, num_tokens=3, num_reqs=2), [0, 2, 1], ++ "query_start_loc"), ++ (PleOffloadRequest(dp_rank=0, num_tokens=5, num_reqs=1), [0, 1], ++ "num_tokens"), ++ (PleOffloadRequest(dp_rank=0, num_tokens=1, num_reqs=3), [0, 1, 1], ++ "num_reqs"), ++ (PleOffloadRequest(dp_rank=9, num_tokens=1, num_reqs=1), [0, 1], ++ "dp_rank=9"), ++ ], ++) ++def test_invalid_request_fails_before_publication( ++ ple_request: PleOffloadRequest, ++ query_start_loc: list[int], ++ match: str, ++) -> None: ++ runner = _runner(query_start_loc) ++ flags = [ ++ target.host_flag ++ for layers in runner._worker_targets.values() ++ for targets in layers.values() ++ for target in targets ++ ] ++ with pytest.raises(RuntimeError, match=match): ++ runner._validate_requests_before_publication([ple_request]) ++ assert all(int(flag[0]) == 0 for flag in flags) ++ ++ ++class _SelectedGraphManager: ++ def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int: ++ return 12 ++ ++ ++class _PublicationGuard: ++ called = False ++ ++ def prepare_forward(self, *args: Any, **kwargs: Any) -> None: ++ self.called = True ++ raise AssertionError("must fail before connector publication") ++ ++ ++def test_selected_full_graph_copy_count_precedes_publication() -> None: ++ runner = GPUModelRunner.__new__(GPUModelRunner) ++ connector = _PublicationGuard() ++ runner._ple_offload_connector = connector ++ runner.cudagraph_manager = _SelectedGraphManager() ++ input_batch = SimpleNamespace( ++ num_reqs=1, ++ num_tokens_after_padding=8, ++ input_ids=torch.empty(8, dtype=torch.int32), ++ ) ++ desc = BatchExecutionDescriptor( ++ cg_mode=CUDAGraphMode.FULL, ++ num_tokens=8, ++ num_reqs=1, ++ ) ++ ++ with pytest.raises(RuntimeError, match="selected-forward"): ++ GPUModelRunner._prepare_ple_forward( ++ runner, input_batch, desc, dummy_run=False ++ ) ++ assert connector.called is False ++ ++ ++def test_production_capture_callback_rejects_actual_ple_copy_mismatch() -> None: ++ manager = CudaGraphManager.__new__(CudaGraphManager) ++ manager._ple_full_copy_geometry = {} ++ manager._expected_ple_copy_layers = ("ple",) ++ manager._allow_ple_full_graph = True ++ desc = BatchExecutionDescriptor( ++ cg_mode=CUDAGraphMode.FULL, ++ num_tokens=8, ++ num_reqs=1, ++ ) ++ ++ with pytest.raises(RuntimeError, match="actual PLE copy geometry"): ++ # This is the production capture scope. The injected observation uses ++ # the same hook called by the real wait-copy custom-op implementation. ++ with manager.capture_ple_copy_geometry(desc): ++ _record_ple_copy_geometry("ple", 12, 16) ++ assert desc not in manager._ple_full_copy_geometry ++ ++ ++def test_speculator_full_capture_refuses_observed_ple_op() -> None: ++ manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager) ++ manager._ple_full_copy_geometry = {} ++ manager._expected_ple_copy_layers = () ++ manager._allow_ple_full_graph = False ++ desc = BatchExecutionDescriptor( ++ cg_mode=CUDAGraphMode.FULL, ++ num_tokens=8, ++ num_reqs=1, ++ ) ++ ++ with pytest.raises(RuntimeError, match="speculator/draft"): ++ with manager.capture_ple_copy_geometry(desc): ++ _record_ple_copy_geometry("unexpected-draft-ple", 8, 16) ++ assert desc not in manager._ple_full_copy_geometry ++ ++ ++def test_eager_correspondence_uses_actual_padded_input_view() -> None: ++ runner = GPUModelRunner.__new__(GPUModelRunner) ++ connector = _PublicationGuard() ++ runner._ple_offload_connector = connector ++ runner.cudagraph_manager = None ++ input_batch = SimpleNamespace( ++ num_reqs=1, ++ num_tokens_after_padding=8, ++ input_ids=torch.empty(12, dtype=torch.int32), ++ ) ++ desc = BatchExecutionDescriptor( ++ cg_mode=CUDAGraphMode.NONE, ++ num_tokens=8, ++ num_reqs=1, ++ ) ++ with pytest.raises(RuntimeError, match="selected-forward"): ++ GPUModelRunner._prepare_ple_forward( ++ runner, input_batch, desc, dummy_run=False ++ ) ++ assert connector.called is False +diff --git a/tests/v1/ple_offload/test_fatal_rpc.py b/tests/v1/ple_offload/test_fatal_rpc.py +new file mode 100644 +index 0000000..1c86f9b +--- /dev/null ++++ b/tests/v1/ple_offload/test_fatal_rpc.py +@@ -0,0 +1,145 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import multiprocessing ++from typing import Any ++ ++import pytest ++ ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload.fatal import ( ++ fail_ple_execute_model, ++ ple_fail_closed, ++) ++ ++ ++class _Progress: ++ sequence = 7 ++ ++ def begin(self, **_: Any) -> int: ++ return self.sequence ++ ++ def complete(self) -> None: ++ return ++ ++ ++class _RaisingProgress(_Progress): ++ def begin(self, **_: Any) -> int: ++ raise RuntimeError("begin injection") ++ ++ ++class _InjectedConnector: ++ rank = 3 ++ tp_rank = 1 ++ _layers = {"ple": object()} ++ _progress = _Progress() ++ _input_ids_buf = None ++ _query_start_loc_buf = None ++ ++ def __init__(self, where: str) -> None: ++ self.where = where ++ if where == "begin": ++ self._progress = _RaisingProgress() ++ ++ def _fail_closed(self, where: str, error: BaseException) -> None: ++ ple_fail_closed( ++ where=where, ++ rank=self.rank, ++ layer="ple", ++ sequence=self._progress.sequence, ++ error=error, ++ ) ++ ++ def _prepare_forward_impl(self, *args: Any, **kwargs: Any) -> None: ++ if self.where == "prepare": ++ raise RuntimeError("prepare injection") ++ ++ def _release_outputs_impl(self) -> None: ++ if self.where == "release": ++ raise RuntimeError("release injection") ++ ++ ++class _Runner: ++ def __init__(self, where: str) -> None: ++ self._ple_offload_connector = _InjectedConnector(where) ++ self.where = where ++ ++ @fail_ple_execute_model ++ def execute_model(self) -> None: ++ connector = self._ple_offload_connector ++ PleOffloadConnector.prepare_forward( ++ connector, ++ 1, ++ 1, ++ False, ++ consumer_num_tokens=1, ++ ) ++ if self.where == "forward": ++ raise RuntimeError("forward injection") ++ PleOffloadConnector.release_outputs(connector) ++ ++ ++def _rpc_loop(where: str, successor_count: Any) -> None: ++ runner = _Runner(where) ++ try: ++ runner.execute_model() ++ except Exception: ++ # Mirrors WorkerProc._execute_worker_rpc: this must never regain control. ++ pass ++ with successor_count.get_lock(): ++ successor_count.value += 1 ++ ++ ++@pytest.mark.parametrize("where", ["prepare", "forward", "release"]) ++def test_ple_exception_kills_workerproc_style_rpc( ++ where: str, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ successor_count = context.Value("i", 0) ++ process = context.Process( ++ target=_rpc_loop, ++ args=(where, successor_count), ++ ) ++ process.start() ++ # Spawned children import vLLM cold (measured 4-8 s on the reference host); the bound is ++ # about start-up, not the fatal path, which exits immediately. ++ process.join(timeout=120) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ assert successor_count.value == 0 ++ ++ ++def _rpc_loop_connector_only(where: str, successor_count: Any) -> None: ++ # No execute_model decorator: only the connector's own fatal wrappers can ++ # stop the RPC loop from regaining control. (The decorated test above ++ # passes on the decorator alone, so it cannot vouch for these wrappers.) ++ connector = _InjectedConnector(where) ++ try: ++ if where == "begin": ++ PleOffloadConnector.begin_forward( ++ connector, dummy_run=False, num_tokens=1 ++ ) ++ elif where == "prepare": ++ PleOffloadConnector.prepare_forward( ++ connector, 1, 1, False, consumer_num_tokens=1 ++ ) ++ else: ++ PleOffloadConnector.release_outputs(connector) ++ except Exception: ++ pass ++ with successor_count.get_lock(): ++ successor_count.value += 1 ++ ++ ++@pytest.mark.parametrize("where", ["begin", "prepare", "release"]) ++def test_connector_wrappers_are_fatal_without_decorator(where: str) -> None: ++ context = multiprocessing.get_context("spawn") ++ successor_count = context.Value("i", 0) ++ process = context.Process( ++ target=_rpc_loop_connector_only, ++ args=(where, successor_count), ++ ) ++ process.start() ++ process.join(timeout=120) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ assert successor_count.value == 0 +diff --git a/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py +new file mode 100644 +index 0000000..900851a +--- /dev/null ++++ b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py +@@ -0,0 +1,80 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Disarmed-path check: a disarmed PLE child writes nothing but the real READY.""" ++ ++import os ++import subprocess ++import sys ++from multiprocessing.connection import Connection ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import pytest ++ ++ ++def test_partial_ready_writers_are_inert_unless_armed_for_that_cut(monkeypatch): ++ import vllm_ple_offload_bootstrap as bootstrap ++ ++ sender = Mock(return_value=(8, 16)) ++ monkeypatch.setattr(bootstrap, "_FAULT", None) ++ for cut in ("launch_partial_ready", "launch_partial_ready_header"): ++ assert bootstrap._write_and_ack_partial(99, cut, sender) is False ++ sender.assert_not_called() ++ ++ other = SimpleNamespace(cut="exc_before_prepare", fires=Mock(return_value=True)) ++ monkeypatch.setattr(bootstrap, "_FAULT", other) ++ assert bootstrap._write_and_ack_partial(99, "launch_partial_ready", sender) is False ++ sender.assert_not_called() ++ other.fires.assert_not_called() ++ ++ ++# The launch payload's config field is the protocol-0 encoding of None (b"N."). ++_CHILD = r""" ++import os, sys, types ++write_fd = int(sys.argv[1]) ++os.environ.pop("VLLM_PLE_TEST_FAULT", None) ++import vllm_ple_offload_bootstrap as bootstrap ++assert bootstrap._FAULT is None ++bootstrap._recv_launch_payload = lambda fd: { ++ "vllm_config_payload": b"N.", ++ "num_workers": 1, ++ "ipc_addr": "ipc:///unused", ++} ++ ++def proc_main(*, vllm_config, num_workers, ipc_addr, ready_pipe): ++ assert vllm_config is None ++ ready_pipe.send({"status": "READY", "layer_names": ["ple"]}) ++ os._exit(0) ++ ++fake = types.ModuleType("vllm.v1.ple_offload.worker") ++fake.PleOffloadWorker = types.SimpleNamespace(proc_main=proc_main) ++sys.modules["vllm.v1.ple_offload.worker"] = fake ++death_read, death_write = os.pipe() ++sys.argv = ["bootstrap", str(write_fd), str(death_read)] ++bootstrap.main() ++""" ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="inherited fds are POSIX-only") ++def test_disarmed_bootstrap_control_stream_is_exactly_one_ready_frame(): ++ read_fd, write_fd = os.pipe() ++ env = dict(os.environ) ++ env.pop("VLLM_PLE_TEST_FAULT", None) ++ proc = subprocess.Popen( ++ [sys.executable, "-c", _CHILD, str(write_fd)], ++ pass_fds=(write_fd,), ++ env=env, ++ ) ++ os.close(write_fd) ++ reader = Connection(read_fd) ++ try: ++ assert proc.wait(timeout=120) == 0 ++ assert reader.poll(5), "no READY frame on the control stream" ++ assert reader.recv() == {"status": "READY", "layer_names": ["ple"]} ++ # Nothing may follow (or precede, which would have broken the frame). ++ with pytest.raises(EOFError): ++ reader.recv_bytes() ++ finally: ++ reader.close() ++ if proc.poll() is None: ++ proc.kill() ++ proc.wait(timeout=120) +diff --git a/tests/v1/ple_offload/test_fault_injector_basics.py b/tests/v1/ple_offload/test_fault_injector_basics.py +new file mode 100644 +index 0000000..0a9fd23 +--- /dev/null ++++ b/tests/v1/ple_offload/test_fault_injector_basics.py +@@ -0,0 +1,260 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import os ++import queue ++import subprocess ++import sys ++import threading ++from collections import deque ++from pathlib import Path ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import pytest ++ ++import vllm_ple_test_fault as test_fault ++ ++ ++def _spec(monkeypatch, tmp_path, value): ++ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) ++ spec = test_fault._parse_fault_spec(value) ++ assert spec is not None ++ return spec ++ ++ ++@pytest.mark.parametrize( ++ "value", ++ [ ++ "cut=park_before_publish,ordinal=1,tp=0", ++ "cut=launch_silent_ready,ordinal=1,pp=0", ++ "cut=drop_request,ordinal=1,tp=1,pp=0", ++ "cut=exc_after_prepare,ordinal=1,tp=0,pp=1", ++ ], ++) ++def test_impossible_role_selectors_are_rejected(monkeypatch, tmp_path, value): ++ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) ++ with pytest.raises(ValueError): ++ test_fault._parse_fault_spec(value) ++ ++ ++def test_armed_import_requires_dedicated_log(tmp_path): ++ env = dict(os.environ) ++ env["VLLM_PLE_TEST_FAULT"] = ( ++ "cut=exc_before_prepare,ordinal=1,rank=any" ++ ) ++ env.pop("VLLM_PLE_TEST_FAULT_LOG", None) ++ proc = subprocess.run( ++ [sys.executable, "-c", "import vllm_ple_test_fault"], ++ env=env, ++ text=True, ++ capture_output=True, ++ timeout=120, ++ ) ++ assert proc.returncode != 0 ++ assert "VLLM_PLE_TEST_FAULT_LOG is required" in proc.stderr ++ ++ ++def test_failed_marker_write_suppresses_action(monkeypatch, tmp_path): ++ spec = _spec( ++ monkeypatch, ++ tmp_path, ++ "cut=exc_after_prepare,ordinal=1,tp=0,pp=0", ++ ) ++ monkeypatch.setattr(test_fault.os, "write", Mock(side_effect=BlockingIOError)) ++ assert not spec.fires( ++ "exc_after_prepare", ++ role="gpu", ++ unit="runner-forward", ++ eligible=True, ++ tp_rank=0, ++ pp_rank=0, ++ dp_rank=0, ++ transaction_id="sequence:1", ++ ) ++ assert spec._state.diagnostic_failed ++ assert not spec._state.fired ++ ++ ++@pytest.mark.parametrize("order", ["bootstrap-first", "package-first"]) ++@pytest.mark.parametrize("armed", [False, True]) ++def test_fresh_interpreter_import_identity(tmp_path, order, armed): ++ env = dict(os.environ) ++ env.pop("VLLM_PLE_TEST_FAULT", None) ++ env.pop("VLLM_PLE_TEST_FAULT_LOG", None) ++ if armed: ++ env["VLLM_PLE_TEST_FAULT"] = ( ++ "cut=launch_silent_ready,ordinal=1,rank=any" ++ ) ++ env["VLLM_PLE_TEST_FAULT_LOG"] = str(tmp_path / f"{order}.log") ++ script = f""" ++import sys ++if {order!r} == "bootstrap-first": ++ import vllm_ple_offload_bootstrap as bootstrap ++ absent_before_package = "vllm_ple_test_fault" not in sys.modules ++ import vllm_ple_test_fault as package ++else: ++ import vllm_ple_test_fault as package ++ import vllm_ple_offload_bootstrap as bootstrap ++ absent_before_package = False ++assert bootstrap._FAULT is package._FAULT ++assert getattr(sys.modules["vllm_ple_test_fault"], "_FAULT") is package._FAULT ++if {armed!r}: ++ assert package._FAULT is not None ++else: ++ assert package._FAULT is None ++ if {order!r} == "bootstrap-first": ++ assert absent_before_package ++""" ++ subprocess.run( ++ [sys.executable, "-c", script], ++ env=env, ++ check=True, ++ timeout=120, ++ ) ++ ++ ++class _IncompleteEvent: ++ def query(self): ++ return False ++ ++ ++class _FatalObserved(BaseException): ++ pass ++ ++ ++@pytest.mark.parametrize( ++ ("record_completion", "complete_submission", "expected"), ++ [ ++ (False, False, "host-submission"), ++ (True, False, "gpu-completion"), ++ (True, True, "gpu-completion"), ++ ], ++) ++def test_real_progress_monitor_observed_state( ++ monkeypatch, ++ record_completion, ++ complete_submission, ++ expected, ++): ++ import vllm.v1.ple_offload.progress as progress_module ++ ++ observed = [] ++ reached = threading.Event() ++ ++ def fatal(**kwargs): ++ observed.append(kwargs) ++ reached.set() ++ raise _FatalObserved ++ ++ monkeypatch.setattr(progress_module, "ple_fail_closed", fatal) ++ monitor = progress_module.PleProgressMonitor( ++ rank=0, ++ layer_names=("ple",), ++ serving_deadline_s=0.02, ++ startup_deadline_s=0.02, ++ ) ++ monitor.begin(startup=False, num_tokens=1) ++ if record_completion: ++ monitor.track_completion(_IncompleteEvent()) ++ if complete_submission: ++ monitor.complete() ++ assert reached.wait(2) ++ assert expected in observed[0]["where"] ++ monitor.close() ++ ++ ++def test_drop_retains_event_and_next_pool_one_launch_fails_closed( ++ monkeypatch, ++ tmp_path, ++): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ spec = _spec( ++ monkeypatch, ++ tmp_path, ++ "cut=drop_request,ordinal=1,tp=0,pp=0", ++ ) ++ monkeypatch.setattr(connector_module, "_FAULT", spec) ++ ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector._progress = SimpleNamespace(sequence=7) ++ connector._dropped_events = [] ++ connector._d2h_event_pool = queue.Queue(maxsize=1) ++ event = Mock() ++ connector._d2h_event_pool.put_nowait(event) ++ connector._request_queue = Mock() ++ connector._enqueue_cuda_inputs = Mock() ++ ++ connector._launch( ++ num_reqs=1, ++ num_tokens=2, ++ serving=True, ++ transaction_id="sequence:7@begin:1.000000000", ++ correlation_id="sequence:7", ++ ) ++ assert connector._dropped_events == [event] ++ connector._request_queue.put_nowait.assert_not_called() ++ ++ with pytest.raises( ++ RuntimeError, ++ match="more requests than configured concurrent batches", ++ ): ++ connector._launch( ++ num_reqs=1, ++ num_tokens=2, ++ serving=True, ++ transaction_id="sequence:8@begin:2.000000000", ++ correlation_id="sequence:8", ++ ) ++ ++ ++def test_async_output_cut_uses_real_sample_boundary(monkeypatch, tmp_path): ++ import vllm.v1.worker.gpu.model_runner as model_runner_module ++ ++ spec = _spec( ++ monkeypatch, ++ tmp_path, ++ "cut=exc_async_output,ordinal=1,tp=0,pp=0", ++ ) ++ monkeypatch.setattr(model_runner_module, "_FAULT", spec) ++ ++ class FatalReached(Exception): ++ pass ++ ++ connector = SimpleNamespace( ++ tp_rank=0, ++ pp_rank=0, ++ dp_rank=0, ++ completion_is_pending=Mock(return_value=True), ++ _fail_closed=Mock(side_effect=FatalReached), ++ ) ++ runner = object.__new__(model_runner_module.GPUModelRunner) ++ runner._ple_offload_connector = connector ++ runner.is_last_pp_rank = False ++ runner.pp_handler = Mock() ++ runner.execute_model_state = SimpleNamespace( ++ input_batch=object(), ++ attn_metadata=None, ++ slot_mappings_by_layer=None, ++ hidden_states=None, ++ aux_hidden_states=None, ++ dp_sync=None, ++ finished_req_ids=[], ++ ec_connector_output=None, ++ routed_experts=None, ++ cudagraph_stats=None, ++ ple_serving=True, ++ ple_sequence=11, ++ ple_transaction="sequence:11@begin:1.000000000", ++ ) ++ ++ with pytest.raises(FatalReached): ++ runner.sample_tokens(None) ++ assert connector._fail_closed.call_args.args[0] == "sample_tokens" ++ error = connector._fail_closed.call_args.args[1] ++ assert isinstance(error, RuntimeError) ++ assert "exc_async_output" in str(error) +diff --git a/tests/v1/ple_offload/test_fault_injector_gates.py b/tests/v1/ple_offload/test_fault_injector_gates.py +new file mode 100644 +index 0000000..70e0d3e +--- /dev/null ++++ b/tests/v1/ple_offload/test_fault_injector_gates.py +@@ -0,0 +1,468 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ast ++import contextlib ++import importlib.util ++import os ++import pickle ++import socket ++import struct ++import subprocess ++import sys ++import threading ++import time ++from pathlib import Path ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import pytest ++ ++import vllm_ple_test_fault as test_fault ++ ++ ++class _SequenceMustNotBeRead: ++ @property ++ def sequence(self): ++ raise AssertionError("disarmed path read progress.sequence") ++ ++ def track_completion(self, _event): ++ pass ++ ++ def complete(self): ++ pass ++ ++ ++def _spec(monkeypatch, tmp_path, text): ++ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) ++ value = test_fault._parse_fault_spec(text) ++ assert value is not None ++ return value ++ ++ ++def test_disarmed_sender_deadline_does_not_enter_logger(monkeypatch): ++ import vllm.v1.ple_offload.worker as worker ++ ++ monkeypatch.setattr(worker, "_FAULT", None) ++ monkeypatch.setattr( ++ worker.logger, ++ "info", ++ Mock(side_effect=AssertionError("logger entered")), ++ ) ++ left, right = socket.socketpair() ++ try: ++ with pytest.raises(TimeoutError): ++ worker._send_framed_payload( ++ left.fileno(), ++ b"payload", ++ time.monotonic() - 1, ++ ) ++ worker.logger.info.assert_not_called() ++ finally: ++ left.close() ++ right.close() ++ ++ ++def test_disarmed_launch_does_not_read_progress_sequence(monkeypatch): ++ import queue ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ monkeypatch.setattr(connector_module, "_FAULT", None) ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector._progress = _SequenceMustNotBeRead() ++ connector._d2h_event_pool = queue.Queue() ++ connector._d2h_event_pool.put_nowait(Mock()) ++ connector._enqueue_cuda_inputs = Mock() ++ connector._request_queue = Mock() ++ ++ connector._launch( ++ 1, ++ 2, ++ serving=True, ++ transaction_id=None, ++ correlation_id=None, ++ ) ++ pending = connector._request_queue.put_nowait.call_args.args[0] ++ assert pending.transaction_id is None ++ assert pending.correlation_id is None ++ ++ ++def test_disarmed_release_does_not_read_progress_sequence(monkeypatch): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ monkeypatch.setattr(connector_module, "_FAULT", None) ++ monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) ++ event = Mock() ++ stream = Mock() ++ monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: event) ++ monkeypatch.setattr( ++ connector_module.torch.cuda, ++ "current_stream", ++ lambda _device: stream, ++ ) ++ ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector.device = SimpleNamespace(index=0) ++ connector._layers = {} ++ connector._after_dummy = False ++ connector._progress = _SequenceMustNotBeRead() ++ ++ connector._release_outputs_impl(serving=False) ++ event.record.assert_called_once_with(stream) ++ ++ ++def test_sample_tokens_connector_rank_fatalizes(): ++ from vllm.v1.ple_offload.fatal import fail_ple_execute_model ++ ++ class Reached(Exception): ++ pass ++ ++ class Runner: ++ _ple_offload_connector = SimpleNamespace( ++ _fail_closed=Mock(side_effect=Reached) ++ ) ++ ++ @fail_ple_execute_model ++ def sample_tokens(self): ++ raise ValueError("sample failure") ++ ++ with pytest.raises(Reached): ++ Runner().sample_tokens() ++ ++ ++def test_sample_tokens_non_connector_rank_rethrows(): ++ from vllm.v1.ple_offload.fatal import fail_ple_execute_model ++ ++ class Runner: ++ _ple_offload_connector = None ++ ++ @fail_ple_execute_model ++ def sample_tokens(self): ++ raise ValueError("sample failure") ++ ++ with pytest.raises(ValueError, match="sample failure"): ++ Runner().sample_tokens() ++ ++ ++def test_profile_is_never_eligible_even_with_phase_any( ++ monkeypatch, ++ tmp_path, ++): ++ spec = _spec( ++ monkeypatch, ++ tmp_path, ++ "cut=exc_after_prepare,ordinal=1,tp=0,pp=0,phase=any", ++ ) ++ assert not spec.fires( ++ "exc_after_prepare", ++ role="gpu", ++ unit="runner-forward", ++ eligible=False, ++ tp_rank=0, ++ pp_rank=0, ++ dp_rank=0, ++ transaction_id="sequence:1@begin:1.000000000", ++ correlation_id="sequence:1", ++ ) ++ assert spec._state.count == 0 ++ ++ ++def test_offload_counts_only_serving_and_preserves_correlation( ++ monkeypatch, ++ tmp_path, ++): ++ import vllm.v1.ple_offload.worker as worker ++ ++ spec = _spec( ++ monkeypatch, ++ tmp_path, ++ "cut=park_before_publish,ordinal=1,rank=any", ++ ) ++ monkeypatch.setattr(worker, "_FAULT", spec) ++ runner = object.__new__(worker.PleOffloadRunner) ++ runner._validate_requests_before_publication = Mock() ++ runner._worker_targets = {0: {}} ++ runner._input_bufs = {} ++ runner._scratch_bufs = {} ++ runner._layers = {} ++ runner._clamp_input_ids = False ++ runner._step = 0 ++ runner._request_index = 0 ++ ++ request = worker.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1) ++ for _ in range(3): ++ runner._handle_requests_impl( ++ [worker._ReceivedPleOffloadRequest(request, False, None, None)] ++ ) ++ assert runner._request_index == 0 ++ ++ reached = threading.Event() ++ monkeypatch.setattr( ++ worker.threading, ++ "Event", ++ lambda: SimpleNamespace(wait=lambda: reached.set()), ++ ) ++ runner._handle_requests_impl( ++ [ ++ worker._ReceivedPleOffloadRequest( ++ request, ++ True, ++ "sequence:17@begin:4.000000000", ++ "sequence:17", ++ ) ++ ] ++ ) ++ assert reached.is_set() ++ line = (tmp_path / "fault.log").read_text() ++ assert "transaction=request:1@sequence:17@begin:4.000000000" in line ++ assert "correlation=sequence:17" in line ++ ++ ++def test_cut6_requires_same_incomplete_completion(monkeypatch): ++ import vllm.v1.ple_offload.progress as progress ++ ++ class Event: ++ complete = False ++ ++ def query(self): ++ return self.complete ++ ++ monitor = progress.PleProgressMonitor( ++ rank=0, ++ layer_names=("ple",), ++ serving_deadline_s=60, ++ startup_deadline_s=60, ++ ) ++ event = Event() ++ try: ++ sequence = monitor.begin(startup=False, num_tokens=1) ++ monitor.track_completion(event) ++ monitor.complete() ++ assert monitor.completion_is_pending(sequence) ++ event.complete = True ++ assert not monitor.completion_is_pending(sequence) ++ assert not monitor.completion_is_pending(sequence + 1) ++ finally: ++ monitor.close() ++ ++ ++def test_partial_ack_requires_complete_required_prefix(monkeypatch): ++ import vllm_ple_offload_bootstrap as bootstrap ++ ++ fault = SimpleNamespace( ++ cut="launch_partial_ready_header", ++ fires=Mock(return_value=True), ++ ) ++ monkeypatch.setattr(bootstrap, "_FAULT", fault) ++ assert not bootstrap._write_and_ack_partial( ++ 1, ++ "launch_partial_ready_header", ++ lambda _fd: (1, 4), ++ ) ++ fault.fires.assert_not_called() ++ ++ assert bootstrap._write_and_ack_partial( ++ 1, ++ "launch_partial_ready_header", ++ lambda _fd: (2, 4), ++ ) ++ assert fault.fires.call_args.kwargs["bytes_written"] == 2 ++ assert fault.fires.call_args.kwargs["frame_size"] == 4 ++ ++ ++@pytest.mark.timeout(180) ++@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX pass_fds") ++@pytest.mark.parametrize( ++ ("cut", "expected_transaction"), ++ [ ++ ("launch_block_reconstruction", "reducer:configuration"), ++ ( ++ "launch_block_worker_import", ++ "module:vllm.v1.ple_offload.worker", ++ ), ++ ], ++) ++def test_gate_is_inside_reducer_or_loader_and_parent_death_wins( ++ tmp_path, ++ cut, ++ expected_transaction, ++): ++ gate = tmp_path / "gate" ++ os.mkfifo(gate) ++ gate_keeper = os.open(gate, os.O_RDWR | os.O_NONBLOCK) ++ control_parent, control_child = socket.socketpair() ++ death_parent, death_child = socket.socketpair() ++ fault_log = tmp_path / "fault.log" ++ env = dict(os.environ) ++ env["VLLM_PLE_TEST_FAULT"] = f"cut={cut},ordinal=1,rank=any" ++ env["VLLM_PLE_TEST_FAULT_LOG"] = str(fault_log) ++ env["VLLM_PLE_TEST_FAULT_GATE"] = str(gate) ++ ++ config_payload = pickle.dumps(None, protocol=pickle.HIGHEST_PROTOCOL) ++ if cut == "launch_block_reconstruction": ++ config_payload = test_fault._FaultReconstructionPayload(config_payload) ++ launch = pickle.dumps( ++ { ++ "vllm_config_payload": config_payload, ++ "num_workers": 1, ++ "ipc_addr": "ipc:///unused-fault", ++ }, ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ ++ process = subprocess.Popen( ++ [ ++ sys.executable, ++ "-m", ++ "vllm_ple_offload_bootstrap", ++ str(control_child.fileno()), ++ str(death_child.fileno()), ++ ], ++ env=env, ++ pass_fds=(control_child.fileno(), death_child.fileno()), ++ ) ++ control_child.close() ++ death_child.close() ++ try: ++ control_parent.sendall(struct.pack("!Q", len(launch)) + launch) ++ deadline = time.monotonic() + 30 ++ while time.monotonic() < deadline: ++ if fault_log.exists() and expected_transaction in fault_log.read_text(): ++ break ++ assert process.poll() is None ++ time.sleep(0.05) ++ else: ++ pytest.fail("operation-entry acknowledgement was not emitted") ++ ++ assert process.poll() is None ++ death_parent.close() ++ assert process.wait(timeout=120) != 0 ++ finally: ++ with contextlib.suppress(ProcessLookupError): ++ process.kill() ++ process.wait(timeout=120) ++ control_parent.close() ++ with contextlib.suppress(OSError): ++ death_parent.close() ++ os.close(gate_keeper) ++ ++ ++def _load_gate_helper(): ++ root = Path(__file__).resolve().parents[3] ++ path = root / "probe" / "faultcut_gate.py" ++ spec = importlib.util.spec_from_file_location("faultcut_gate", path) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ spec.loader.exec_module(module) ++ return module ++ ++ ++@pytest.mark.parametrize( ++ "change", ++ [ ++ {"alive": False}, ++ {"observed_start": 88}, ++ {"bytes_sent": 0}, ++ {"bytes_sent": 100}, ++ {"completed": True}, ++ {"child_not_reading": False}, ++ ], ++) ++def test_external_partial_send_gate_rejections(change): ++ gate = _load_gate_helper() ++ values = { ++ "pid": 12, ++ "expected_start": 77, ++ "observed_start": 77, ++ "alive": True, ++ "bytes_sent": 4, ++ "frame_size": 100, ++ "completed": False, ++ "child_not_reading": True, ++ } ++ values.update(change) ++ ok, _reason = gate.validate_partial_send_gate(**values) ++ assert not ok ++ ++ ++def test_external_gate_rejects_failed_signal(): ++ gate = _load_gate_helper() ++ ok, reason = gate.deliver_checked_signal( ++ pid=12, ++ expected_start=77, ++ read_start=lambda _pid: 77, ++ alive=lambda _pid: True, ++ send=lambda _pid, _sig: False, ++ signal_number=9, ++ ) ++ assert not ok ++ assert "delivery failed" in reason ++ ++ ++def test_every_fault_call_has_one_boolop_guard_and_no_prior_state_read(): ++ root = Path(__file__).resolve().parents[3] ++ expected = { ++ "vllm/v1/ple_offload/connector.py": 3, ++ "vllm/v1/ple_offload/worker.py": 1, ++ "vllm/v1/worker/gpu/model_runner.py": 4, ++ "vllm_ple_offload_bootstrap/__init__.py": 1, ++ } ++ ++ for relative, wanted in expected.items(): ++ tree = ast.parse((root / relative).read_text(encoding="utf-8")) ++ parents = {} ++ for parent in ast.walk(tree): ++ for child in ast.iter_child_nodes(parent): ++ parents[child] = parent ++ ++ calls = [ ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "fires" ++ and isinstance(node.func.value, ast.Name) ++ and node.func.value.id == "_FAULT" ++ ] ++ assert len(calls) == wanted, relative ++ for call in calls: ++ owner = call ++ while owner in parents and not isinstance(owner, ast.BoolOp): ++ owner = parents[owner] ++ assert isinstance(owner, ast.BoolOp), (relative, call.lineno) ++ assert isinstance(owner.op, ast.And), (relative, call.lineno) ++ ++ guard_index = None ++ call_index = None ++ for index, value in enumerate(owner.values): ++ if ( ++ isinstance(value, ast.Compare) ++ and isinstance(value.left, ast.Name) ++ and value.left.id == "_FAULT" ++ and len(value.ops) == 1 ++ and isinstance(value.ops[0], ast.IsNot) ++ and len(value.comparators) == 1 ++ and isinstance(value.comparators[0], ast.Constant) ++ and value.comparators[0].value is None ++ ): ++ guard_index = index ++ if call in ast.walk(value): ++ call_index = index ++ assert guard_index is not None ++ assert call_index is not None ++ assert guard_index < call_index ++ ++ forbidden = {"sequence", "has_pending_completions", ++ "completion_is_pending", "fault_transaction"} ++ before_guard = owner.values[:guard_index] ++ assert not any( ++ isinstance(node, ast.Attribute) and node.attr in forbidden ++ for value in before_guard ++ for node in ast.walk(value) ++ ) +diff --git a/tests/v1/ple_offload/test_fault_wire_frames.py b/tests/v1/ple_offload/test_fault_wire_frames.py +new file mode 100644 +index 0000000..b864355 +--- /dev/null ++++ b/tests/v1/ple_offload/test_fault_wire_frames.py +@@ -0,0 +1,96 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Disarmed-path checks: the disarmed PLE request wire format stays single-frame.""" ++ ++import contextlib ++import queue ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import msgspec ++import pytest ++ ++ ++def _request(): ++ from vllm.v1.ple_offload.protocol import PleOffloadRequest ++ ++ return PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1) ++ ++ ++def test_disarmed_offload_receive_uses_single_frame(monkeypatch): ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ monkeypatch.setattr(worker_module, "_FAULT", None) ++ request = _request() ++ socket = Mock() ++ socket.recv.return_value = msgspec.msgpack.encode(request) ++ socket.recv_multipart.side_effect = AssertionError("multipart read while disarmed") ++ ++ assert worker_module._receive_offload_request(socket) == request ++ socket.recv.assert_called_once_with(0) ++ ++ ++def test_armed_offload_receive_requires_metadata_frame(monkeypatch): ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ monkeypatch.setattr(worker_module, "_FAULT", object()) ++ request = _request() ++ metadata = {"serving": True, "transaction": "request:1", "correlation": "sequence:3"} ++ socket = Mock() ++ socket.recv.side_effect = AssertionError("single-frame read while armed") ++ socket.recv_multipart.return_value = [ ++ msgspec.msgpack.encode(request), ++ msgspec.msgpack.encode(metadata), ++ ] ++ ++ received = worker_module._receive_offload_request(socket, 1) ++ assert received.request == request ++ assert received.serving is True ++ assert received.transaction_id == "request:1" ++ assert received.correlation_id == "sequence:3" ++ socket.recv_multipart.assert_called_once_with(1) ++ ++ socket.recv_multipart.return_value = [msgspec.msgpack.encode(request)] ++ with pytest.raises(RuntimeError): ++ worker_module._receive_offload_request(socket) ++ ++ ++@pytest.mark.parametrize("armed", [False, True]) ++def test_connector_publication_frames_follow_arming(monkeypatch, armed): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ fault = SimpleNamespace(fires=lambda *args, **kwargs: False) if armed else None ++ monkeypatch.setattr(connector_module, "_FAULT", fault) ++ monkeypatch.setattr( ++ connector_module.torch.accelerator, ++ "device_index", ++ lambda _index: contextlib.nullcontext(), ++ ) ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector.device = SimpleNamespace(index=0) ++ connector._d2h_event_pool = queue.Queue() ++ event = Mock() ++ request = _request() ++ pending = connector_module._PendingPleOffloadRequest( ++ request, event, True, "sequence:4@begin:1.000000000", "sequence:4" ++ ) ++ socket = Mock() ++ ++ connector._process_request(pending, socket) ++ ++ event.synchronize.assert_called_once() ++ if armed: ++ socket.send.assert_not_called() ++ payload, metadata = socket.send_multipart.call_args.args[0] ++ assert msgspec.msgpack.decode(metadata) == { ++ "serving": True, ++ "transaction": "sequence:4@begin:1.000000000", ++ "correlation": "sequence:4", ++ } ++ else: ++ socket.send_multipart.assert_not_called() ++ (payload,) = socket.send.call_args.args ++ assert payload == msgspec.msgpack.encode(request) ++ assert connector._d2h_event_pool.get_nowait() is event +diff --git a/tests/v1/ple_offload/test_ipc_socket_unlink.py b/tests/v1/ple_offload/test_ipc_socket_unlink.py +new file mode 100644 +index 0000000..92fbc42 +--- /dev/null ++++ b/tests/v1/ple_offload/test_ipc_socket_unlink.py +@@ -0,0 +1,133 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import contextlib ++import os ++from pathlib import Path ++import signal ++import subprocess ++import sys ++import tempfile ++import time ++ ++import pytest ++ ++from vllm.v1.ple_offload.worker import PleOffloadWorkerHandle ++ ++ ++pytestmark = pytest.mark.skipif( ++ sys.platform != "linux", reason="requires Linux filesystem ipc sockets" ++) ++ ++ ++_STUB_CHILD = r""" ++import pathlib ++import signal ++import sys ++import threading ++ ++import zmq ++ ++from vllm.v1.ple_offload.worker import _unlink_ipc_socket ++ ++ipc_addr, mode, marker = sys.argv[1:] ++stopping = threading.Event() ++if mode == "graceful": ++ signal.signal(signal.SIGTERM, lambda _signum, _frame: stopping.set()) ++else: ++ signal.signal(signal.SIGTERM, signal.SIG_IGN) ++ ++context = zmq.Context() ++socket = context.socket(zmq.PULL) ++socket.bind(ipc_addr) ++try: ++ stopping.wait() ++finally: ++ socket.close(linger=0) ++ context.term() ++ if mode == "graceful": ++ _unlink_ipc_socket(ipc_addr) ++ pathlib.Path(marker).touch() ++""" ++ ++ ++def _start_stub(ipc_addr: str, mode: str, marker: Path) -> subprocess.Popen: ++ return subprocess.Popen( ++ [sys.executable, "-c", _STUB_CHILD, ipc_addr, mode, str(marker)], ++ stdout=subprocess.DEVNULL, ++ stderr=subprocess.PIPE, ++ text=True, ++ ) ++ ++ ++def _wait_for_socket(proc: subprocess.Popen, path: Path) -> None: ++ deadline = time.monotonic() + 30 ++ while time.monotonic() < deadline: ++ if path.exists(): ++ return ++ if proc.poll() is not None: ++ stderr = proc.stderr.read() if proc.stderr is not None else "" ++ pytest.fail(f"stub child exited before bind: {stderr}") ++ time.sleep(0.01) ++ pytest.fail("stub child did not bind its ipc socket") ++ ++ ++def _cleanup_process(proc: subprocess.Popen, socket_path: Path) -> None: ++ if proc.poll() is None: ++ proc.kill() ++ proc.wait(timeout=5) ++ with contextlib.suppress(FileNotFoundError): ++ os.unlink(socket_path) ++ if proc.stderr is not None: ++ proc.stderr.close() ++ ++ ++def test_handle_graceful_close_removes_child_bound_socket() -> None: ++ with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir: ++ socket_path = Path(temp_dir) / "owned.sock" ++ marker = Path(temp_dir) / "child-cleaned" ++ ipc_addr = f"ipc://{socket_path}" ++ proc = _start_stub(ipc_addr, "graceful", marker) ++ try: ++ handle = PleOffloadWorkerHandle( ++ proc=proc, ++ death_writer=None, ++ ready_pipe_reader=None, ++ launch_deadline=time.monotonic() + 30, ++ ipc_addr=ipc_addr, ++ ) ++ _wait_for_socket(proc, socket_path) ++ ++ handle.close() ++ ++ assert proc.returncode == 0 ++ assert marker.exists(), "the child ran its post-context unlink" ++ assert not socket_path.exists() ++ finally: ++ _cleanup_process(proc, socket_path) ++ ++ ++def test_handle_close_removes_socket_after_child_sigkill() -> None: ++ with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir: ++ socket_path = Path(temp_dir) / "owned.sock" ++ marker = Path(temp_dir) / "unused" ++ ipc_addr = f"ipc://{socket_path}" ++ proc = _start_stub(ipc_addr, "park", marker) ++ try: ++ handle = PleOffloadWorkerHandle( ++ proc=proc, ++ death_writer=None, ++ ready_pipe_reader=None, ++ launch_deadline=time.monotonic() + 30, ++ ipc_addr=ipc_addr, ++ ) ++ _wait_for_socket(proc, socket_path) ++ proc.kill() ++ proc.wait(timeout=5) ++ assert proc.returncode == -signal.SIGKILL ++ assert socket_path.exists() ++ ++ handle.close() ++ ++ assert not socket_path.exists() ++ finally: ++ _cleanup_process(proc, socket_path) +diff --git a/tests/v1/ple_offload/test_pending_rpc.py b/tests/v1/ple_offload/test_pending_rpc.py +new file mode 100644 +index 0000000..ef3c0fd +--- /dev/null ++++ b/tests/v1/ple_offload/test_pending_rpc.py +@@ -0,0 +1,81 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import multiprocessing ++from types import SimpleNamespace ++from typing import Any ++ ++import pytest ++ ++import vllm.v1.executor.multiproc_executor as executor_module ++from vllm.v1.executor.multiproc_executor import WorkerProc ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++ ++class _BrokenAsyncOutput: ++ def get_output(self) -> Any: ++ raise RuntimeError("injected asynchronous output failure") ++ ++ ++class _Progress: ++ sequence = 9 ++ ++ def __init__(self, pending: bool) -> None: ++ self.has_pending_completions = pending ++ ++ ++class _Connector: ++ rank = 2 ++ _layers = {"ple": object()} ++ ++ def __init__(self, pending: bool) -> None: ++ self._progress = _Progress(pending) ++ ++ @property ++ def has_pending_completions(self) -> bool: ++ return self._progress.has_pending_completions ++ ++ def _fail_closed(self, where: str, error: BaseException) -> None: ++ ple_fail_closed( ++ where=where, ++ rank=self.rank, ++ layer="ple", ++ sequence=self._progress.sequence, ++ error=error, ++ ) ++ ++ ++def _async_output_failure_case(pending: bool) -> None: ++ connector = _Connector(pending) ++ worker = SimpleNamespace( ++ model_runner=SimpleNamespace(_ple_offload_connector=connector) ++ ) ++ proc = WorkerProc.__new__(WorkerProc) ++ proc.worker = SimpleNamespace(worker=worker) ++ proc.worker_response_mq = None ++ ++ # Exercise the real async-output exception recovery branch without needing ++ # to construct the production dataclass's CUDA-backed fields. ++ executor_module.AsyncModelRunnerOutput = _BrokenAsyncOutput ++ proc.enqueue_output(_BrokenAsyncOutput()) ++ ++ ++@pytest.mark.parametrize( ++ ("pending", "must_die"), ++ [(True, True), (False, False)], ++) ++def test_async_output_failure_obeys_pending_ple_ownership( ++ pending: bool, ++ must_die: bool, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ process = context.Process( ++ target=_async_output_failure_case, ++ args=(pending,), ++ ) ++ process.start() ++ process.join(timeout=120) ++ assert not process.is_alive() ++ if must_die: ++ assert process.exitcode not in (None, 0) ++ else: ++ assert process.exitcode == 0 +diff --git a/tests/v1/ple_offload/test_progress_deadline.py b/tests/v1/ple_offload/test_progress_deadline.py +new file mode 100644 +index 0000000..5e7ca19 +--- /dev/null ++++ b/tests/v1/ple_offload/test_progress_deadline.py +@@ -0,0 +1,216 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import multiprocessing ++import os ++import threading ++from pathlib import Path ++from typing import Any ++ ++import pytest ++import torch ++ ++import vllm.v1.ple_offload.connector as connector_module ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload.progress import PleProgressMonitor ++ ++ ++class _PendingEvent: ++ def record(self, stream: Any) -> None: ++ stream.recorded = True ++ ++ def query(self) -> bool: ++ return False ++ ++ ++class _BrokenEvent(_PendingEvent): ++ def __init__(self, query_entered: Any) -> None: ++ self.query_entered = query_entered ++ ++ def query(self) -> bool: ++ self.query_entered.set() ++ raise RuntimeError("injected event query failure") ++ ++ ++class _FakeStream: ++ recorded = False ++ ++ def synchronize(self) -> None: ++ return ++ ++ ++class _FakeLayer: ++ def release_offloaded_output(self, stream: Any) -> None: ++ stream.reset_written = True ++ ++ ++def _silent_producer_case(block_input: bool, started: Any) -> None: ++ monitor = PleProgressMonitor( ++ rank=2, ++ layer_names=("ple",), ++ serving_deadline_s=0.2, ++ startup_deadline_s=1.0, ++ ) ++ monitor.begin(startup=False, num_tokens=32) ++ started.set() ++ ++ if block_input: ++ # Models the request thread stuck in d2h_done_event.synchronize(). ++ threading.Thread( ++ target=threading.Event().wait, ++ daemon=True, ++ ).start() ++ ++ # Neither the fake producer nor the test terminates the process. The ++ # independent progress monitor must do so. ++ threading.Event().wait() ++ ++ ++@pytest.mark.parametrize("block_input", [False, True]) ++def test_progress_deadline_terminates_alive_silent_worker( ++ block_input: bool, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ started = context.Event() ++ process = context.Process( ++ target=_silent_producer_case, ++ args=(block_input, started), ++ ) ++ process.start() ++ # Cold vLLM import in the spawned child takes 4-8 s on the reference host; the deadline ++ # itself (0.2 s) is only charged after `started` is set. ++ assert started.wait(timeout=120) ++ process.join(timeout=10) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ ++ ++def _pending_completion_case(use_release_wrapper: bool, started: Any) -> None: ++ monitor = PleProgressMonitor( ++ rank=2, ++ layer_names=("ple",), ++ serving_deadline_s=0.2, ++ startup_deadline_s=1.0, ++ ) ++ if use_release_wrapper: ++ stream = _FakeStream() ++ connector_module.torch.cuda.current_stream = lambda device=None: stream ++ connector_module.torch.cuda.Event = _PendingEvent ++ connector_module.torch.cuda.is_current_stream_capturing = lambda: False ++ connector_module._rowdump.enabled = lambda: False ++ ++ connector = PleOffloadConnector.__new__(PleOffloadConnector) ++ connector.device = torch.device("cpu") ++ connector.rank = 2 ++ connector.tp_rank = 0 ++ connector._layers = {"ple": _FakeLayer()} ++ connector._progress = monitor ++ connector._rowdump_n = None ++ connector._after_dummy = False ++ monitor.begin(startup=False, num_tokens=8) ++ connector.release_outputs() ++ assert stream.recorded ++ assert stream.reset_written ++ else: ++ monitor.begin(startup=False, num_tokens=8) ++ monitor.track_completion(_PendingEvent()) ++ monitor.complete() ++ started.set() ++ threading.Event().wait() ++ ++ ++@pytest.mark.parametrize("use_release_wrapper", [False, True]) ++def test_pending_gpu_completion_expires_after_release_returns( ++ use_release_wrapper: bool, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ started = context.Event() ++ process = context.Process( ++ target=_pending_completion_case, ++ args=(use_release_wrapper, started), ++ ) ++ process.start() ++ assert started.wait(timeout=120) ++ process.join(timeout=10) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ ++ ++@pytest.mark.parametrize("value", [float("inf"), float("nan"), -1.0, 1e12]) ++def test_constructor_rejects_invalid_deadlines(value: float) -> None: ++ with pytest.raises(ValueError, match="finite"): ++ PleProgressMonitor( ++ rank=0, ++ layer_names=("ple",), ++ serving_deadline_s=value, ++ startup_deadline_s=1.0, ++ ) ++ with pytest.raises(ValueError, match="finite"): ++ PleProgressMonitor( ++ rank=0, ++ layer_names=("ple",), ++ serving_deadline_s=1.0, ++ startup_deadline_s=value, ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "name", ++ ["VLLM_PLE_PROGRESS_DEADLINE_S", "VLLM_PLE_STARTUP_DEADLINE_S"], ++) ++@pytest.mark.parametrize("value", ["inf", "nan", "-1", "1e12"]) ++def test_environment_rejects_invalid_deadlines( ++ monkeypatch: pytest.MonkeyPatch, ++ name: str, ++ value: str, ++) -> None: ++ monkeypatch.setenv("VLLM_PLE_PROGRESS_DEADLINE_S", "1") ++ monkeypatch.setenv("VLLM_PLE_STARTUP_DEADLINE_S", "1") ++ monkeypatch.setenv(name, value) ++ with pytest.raises(ValueError, match="finite"): ++ PleProgressMonitor.from_environment(rank=0, layer_names=("ple",)) ++ ++ ++def _monitor_exception_case( ++ started: Any, ++ query_entered: Any, ++ diagnostic_path: str, ++) -> None: ++ diagnostic_fd = os.open(diagnostic_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) ++ os.dup2(diagnostic_fd, 2) ++ os.close(diagnostic_fd) ++ monitor = PleProgressMonitor( ++ rank=4, ++ layer_names=("ple",), ++ serving_deadline_s=10.0, ++ startup_deadline_s=10.0, ++ ) ++ monitor.begin(startup=False, num_tokens=8) ++ # Set before arming the failing event: the monitor thread's os._exit must not ++ # race a multiprocessing Event.set() that holds the shared condition lock. ++ started.set() ++ monitor.track_completion(_BrokenEvent(query_entered)) ++ monitor.complete() ++ threading.Event().wait() ++ ++ ++def test_monitor_thread_exception_is_fatal(tmp_path: Path) -> None: ++ context = multiprocessing.get_context("spawn") ++ started = context.Event() ++ query_entered = context.Event() ++ diagnostic_path = tmp_path / "monitor-fatal.txt" ++ process = context.Process( ++ target=_monitor_exception_case, ++ args=(started, query_entered, str(diagnostic_path)), ++ ) ++ process.start() ++ assert started.wait(timeout=120) ++ assert query_entered.wait(timeout=10) ++ process.join(timeout=10) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ diagnostic = diagnostic_path.read_text(encoding="utf-8") ++ assert "where=progress-monitor-thread" in diagnostic ++ assert ( ++ "error=RuntimeError:injected event query failure" ++ in diagnostic ++ ) +diff --git a/tests/v1/ple_offload/test_shm_unlink.py b/tests/v1/ple_offload/test_shm_unlink.py +new file mode 100644 +index 0000000..600a7bb +--- /dev/null ++++ b/tests/v1/ple_offload/test_shm_unlink.py +@@ -0,0 +1,397 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import os ++from multiprocessing.reduction import ForkingPickler ++from pathlib import Path ++import subprocess ++import sys ++ ++import pytest ++import torch ++import torch.multiprocessing as torch_mp ++ ++from vllm.v1.ple_offload import connector as connector_module ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload.protocol import PleOffloadRegistration ++ ++ ++def _require_file_system_shm() -> None: ++ if ( ++ os.name != "posix" ++ or not os.path.isdir("/dev/shm") ++ or not os.access("/dev/shm", os.W_OK) ++ ): ++ pytest.skip("POSIX /dev/shm is unavailable") ++ if "file_system" not in torch_mp.get_all_sharing_strategies(): ++ pytest.skip("torch file_system sharing is unavailable") ++ ++ ++def _shm_path(tensor: torch.Tensor) -> Path: ++ _, filename, _ = tensor.untyped_storage()._share_filename_cpu_() ++ return Path("/dev/shm") / os.fsdecode(filename).lstrip("/") ++ ++ ++def _make_connector_storage() -> tuple[PleOffloadConnector, list[torch.Tensor]]: ++ original_strategy = torch_mp.get_sharing_strategy() ++ torch_mp.set_sharing_strategy("file_system") ++ try: ++ host_outputs = { ++ "layer.0": torch.arange(12, dtype=torch.float32) ++ .reshape(3, 4) ++ .share_memory_(), ++ "layer.1": torch.arange(8, dtype=torch.float32) ++ .reshape(2, 4) ++ .share_memory_(), ++ } ++ host_flags = { ++ name: torch.zeros(16, dtype=torch.int32).share_memory_() ++ for name in host_outputs ++ } ++ input_ids = torch.arange(16, dtype=torch.int32).share_memory_() ++ query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_() ++ ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3) ++ ngram_context.share_memory_() ++ ++ registration = PleOffloadRegistration( ++ worker_id=0, ++ tp_rank=0, ++ dp_rank=0, ++ host_output_buffers=host_outputs, ++ host_flag_tensors=host_flags, ++ input_ids_buf=input_ids, ++ query_start_loc_buf=query_start_loc, ++ ngram_context_buf=ngram_context, ++ ) ++ # Match producer serialization and worker reconstruction. loads() ++ # consumes the transfer reference added by dumps(), leaving storage ++ # destruction able to reach the allocator's zero-counter branch. ++ payload = ForkingPickler.dumps(registration) ++ rebuilt_registration = ForkingPickler.loads(payload) ++ assert isinstance(rebuilt_registration, PleOffloadRegistration) ++ del rebuilt_registration ++ finally: ++ torch_mp.set_sharing_strategy(original_strategy) ++ ++ connector = PleOffloadConnector.__new__(PleOffloadConnector) ++ connector._host_output_bufs = host_outputs ++ connector._host_flags = host_flags ++ connector._reserved_host_storage_filenames = set() ++ connector._input_ids_buf = input_ids ++ connector._query_start_loc_buf = query_start_loc ++ connector._ngram_context_buf = ngram_context ++ tensors = [ ++ *host_outputs.values(), ++ *host_flags.values(), ++ input_ids, ++ query_start_loc, ++ ngram_context, ++ ] ++ assert len(tensors) == 7 ++ return connector, tensors ++ ++ ++def test_release_host_storage_unlinks_names_but_keeps_mappings() -> None: ++ _require_file_system_shm() ++ connector, tensors = _make_connector_storage() ++ paths = [_shm_path(tensor) for tensor in tensors] ++ before = [tensor.clone() for tensor in tensors] ++ assert len(set(paths)) == 7 ++ assert all(path.exists() for path in paths) ++ ++ connector.release_host_storage() ++ connector.release_host_storage() ++ ++ assert all(not path.exists() for path in paths) ++ for tensor, expected in zip(tensors, before): ++ assert torch.equal(tensor, expected) ++ tensor.add_(1) ++ assert torch.equal(tensor, expected + 1) ++ ++ ++def test_without_release_helper_names_remain(monkeypatch) -> None: ++ _require_file_system_shm() ++ connector, tensors = _make_connector_storage() ++ paths = [_shm_path(tensor) for tensor in tensors] ++ ++ with monkeypatch.context() as patch: ++ patch.setattr( ++ connector_module, ++ "_unlink_filename_shared_storage", ++ lambda _tensor, _reservations: None, ++ ) ++ connector.release_host_storage() ++ ++ # Do not manually unlink these while their unreserved storages are live. ++ # Normal storage destruction removes them after the test frame is released. ++ assert all(path.exists() for path in paths) ++ ++ ++def test_release_host_storage_does_not_convert_fd_shared_storage() -> None: ++ if "file_descriptor" not in torch_mp.get_all_sharing_strategies(): ++ pytest.skip("torch file_descriptor sharing is unavailable") ++ original_strategy = torch_mp.get_sharing_strategy() ++ torch_mp.set_sharing_strategy("file_descriptor") ++ try: ++ fd_tensor = torch.arange(8, dtype=torch.int32).share_memory_() ++ finally: ++ torch_mp.set_sharing_strategy(original_strategy) ++ ++ connector = PleOffloadConnector.__new__(PleOffloadConnector) ++ connector._host_output_bufs = {"layer.0": fd_tensor} ++ connector._host_flags = {} ++ connector._reserved_host_storage_filenames = set() ++ connector._input_ids_buf = torch.arange(2, dtype=torch.int32) ++ connector._query_start_loc_buf = torch.arange(2, dtype=torch.int32) ++ connector._ngram_context_buf = None ++ expected = fd_tensor.clone() ++ ++ connector.release_host_storage() ++ ++ assert fd_tensor.untyped_storage()._get_shared_fd() >= 0 ++ assert torch.equal(fd_tensor, expected) ++ ++ ++class _FakeFilenameStorage: ++ def __init__(self, filename: bytes) -> None: ++ self.device = torch.device("cpu") ++ self.filename = filename ++ self.increfs = 0 ++ self.decrefs = 0 ++ ++ def is_shared(self) -> bool: ++ return True ++ ++ def _get_shared_fd(self) -> int: ++ raise RuntimeError("filename-backed storage has no shared fd") ++ ++ def _share_filename_cpu_(self) -> tuple[bytes, bytes, int]: ++ return b"manager", self.filename, 4 ++ ++ def _shared_incref(self) -> None: ++ self.increfs += 1 ++ ++ def _shared_decref(self) -> "_FakeFilenameStorage": ++ self.decrefs += 1 ++ return self ++ ++ ++class _FakeTensor: ++ def __init__(self, storage: _FakeFilenameStorage) -> None: ++ self.storage = storage ++ ++ def untyped_storage(self) -> _FakeFilenameStorage: ++ return self.storage ++ ++ ++def _make_fake_connector(tensor: _FakeTensor) -> PleOffloadConnector: ++ connector = PleOffloadConnector.__new__(PleOffloadConnector) ++ # Repeat the same storage in several connector positions to cover aliases. ++ connector._host_output_bufs = {"layer.0": tensor} ++ connector._host_flags = {"layer.0": tensor} ++ connector._reserved_host_storage_filenames = set() ++ connector._input_ids_buf = tensor ++ connector._query_start_loc_buf = tensor ++ connector._ngram_context_buf = None ++ return connector ++ ++ ++def test_release_host_storage_reserves_each_name_once(monkeypatch) -> None: ++ storage = _FakeFilenameStorage(b"/ple-test-reservation") ++ connector = _make_fake_connector(_FakeTensor(storage)) ++ unlinked: list[str] = [] ++ monkeypatch.setattr(connector_module.os, "unlink", unlinked.append) ++ ++ connector.release_host_storage() ++ connector.release_host_storage() ++ ++ assert storage.increfs == 1 ++ assert storage.decrefs == 0 ++ assert unlinked == ["/dev/shm/ple-test-reservation"] ++ ++ ++def test_release_host_storage_rolls_back_on_unlink_failure(monkeypatch) -> None: ++ storage = _FakeFilenameStorage(b"/ple-test-rollback") ++ connector = _make_fake_connector(_FakeTensor(storage)) ++ ++ def fail_unlink(_path: str) -> None: ++ raise PermissionError("test unlink failure") ++ ++ monkeypatch.setattr(connector_module.os, "unlink", fail_unlink) ++ ++ connector.release_host_storage() ++ ++ # The fake connector aliases one storage in four positions. A rolled-back ++ # reservation is retried by the next alias in the same pass (by design: the ++ # name is still present), so every reservation must be balanced by a ++ # rollback and nothing may stay reserved. ++ assert storage.increfs >= 1 ++ assert storage.increfs == storage.decrefs ++ assert connector._reserved_host_storage_filenames == set() ++ ++ ++_BALANCED_LIFETIME_SCRIPT = r""" ++import gc ++import os ++from multiprocessing.reduction import ForkingPickler ++from pathlib import Path ++ ++import torch ++import torch.multiprocessing as torch_mp ++ ++from vllm.v1.ple_offload.connector import ( ++ PleOffloadConnector, ++ _PROCESS_LIFETIME_REGISTERED_STORAGE, ++) ++from vllm.v1.ple_offload.protocol import PleOffloadRegistration ++ ++ ++def shm_path(tensor): ++ _, filename, _ = tensor.untyped_storage()._share_filename_cpu_() ++ return Path("/dev/shm") / os.fsdecode(filename).lstrip("/") ++ ++ ++class Progress: ++ def close(self): ++ pass ++ ++ ++original_strategy = torch_mp.get_sharing_strategy() ++torch_mp.set_sharing_strategy("file_system") ++try: ++ host_outputs = { ++ "layer.0": torch.arange(12, dtype=torch.float32) ++ .reshape(3, 4) ++ .share_memory_(), ++ "layer.1": torch.arange(8, dtype=torch.float32) ++ .reshape(2, 4) ++ .share_memory_(), ++ } ++ host_flags = { ++ name: torch.zeros(16, dtype=torch.int32).share_memory_() ++ for name in host_outputs ++ } ++ input_ids = torch.arange(16, dtype=torch.int32).share_memory_() ++ query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_() ++ ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3) ++ ngram_context.share_memory_() ++ ++ registration = PleOffloadRegistration( ++ worker_id=1, ++ tp_rank=1, ++ dp_rank=0, ++ host_output_buffers=host_outputs, ++ host_flag_tensors=host_flags, ++ input_ids_buf=input_ids, ++ query_start_loc_buf=query_start_loc, ++ ngram_context_buf=ngram_context, ++ ) ++ payload = ForkingPickler.dumps(registration) ++ rebuilt_registration = ForkingPickler.loads(payload) ++finally: ++ torch_mp.set_sharing_strategy(original_strategy) ++ ++connector = PleOffloadConnector.__new__(PleOffloadConnector) ++connector._host_output_bufs = host_outputs ++connector._host_flags = host_flags ++connector._reserved_host_storage_filenames = set() ++connector._input_ids_buf = input_ids ++connector._query_start_loc_buf = query_start_loc ++connector._ngram_context_buf = ngram_context ++ ++# Exercise the production close ordering without CUDA or transport setup. ++connector._request_thread = None ++connector._pinned_input_buffers = [] ++connector._mapped_host_buffers = [] ++connector._d2h_event_pool = object() ++connector._registration_socket = None ++connector._zmq_ctx = None ++connector._progress = Progress() ++ ++tensors = [ ++ *host_outputs.values(), ++ *host_flags.values(), ++ input_ids, ++ query_start_loc, ++ ngram_context, ++] ++assert len(tensors) == 7 ++assert len({tensor.untyped_storage()._weak_ref() for tensor in tensors}) == 7 ++ ++# TP != 0 creates and serializes its input buffers without CUDA registration ++# or the process-lifetime reference used by TP0. ++unpinned_tp_input = input_ids ++assert not connector._pinned_input_buffers ++assert all( ++ unpinned_tp_input is not registered ++ for registered in _PROCESS_LIFETIME_REGISTERED_STORAGE ++) ++ ++paths = [shm_path(tensor) for tensor in tensors] ++before = [tensor.clone() for tensor in tensors] ++assert all(path.exists() for path in paths) ++ ++connector.close() ++connector.close() ++ ++assert all(not path.exists() for path in paths) ++for tensor, expected in zip(tensors, before): ++ assert torch.equal(tensor, expected) ++ tensor.add_(1) ++ assert torch.equal(tensor, expected + 1) ++del tensor, expected ++ ++# Remove every original and reconstructed tensor reference. With an unreserved ++# allocator this reaches counter zero after manual unlink and aborts the child. ++connector._host_output_bufs = {} ++connector._host_flags = {} ++connector._input_ids_buf = None ++connector._query_start_loc_buf = None ++connector._ngram_context_buf = None ++del connector ++del rebuilt_registration ++del registration ++del payload ++del unpinned_tp_input ++del tensors ++del before ++del host_outputs ++del host_flags ++del input_ids ++del query_start_loc ++del ngram_context ++gc.collect() ++ ++assert all(not path.exists() for path in paths) ++""" ++ ++ ++def test_close_balanced_filename_storage_lifetime_exits_cleanly() -> None: ++ _require_file_system_shm() ++ ++ # Resolve from the connector actually imported by this pytest process. This ++ # lets the red command run this patched test against an unpatched worktree. ++ import_root = Path(connector_module.__file__).resolve().parents[3] ++ env = os.environ.copy() ++ inherited_pythonpath = env.get("PYTHONPATH") ++ env["PYTHONPATH"] = ( ++ os.pathsep.join((str(import_root), inherited_pythonpath)) ++ if inherited_pythonpath ++ else str(import_root) ++ ) ++ ++ completed = subprocess.run( ++ [sys.executable, "-c", _BALANCED_LIFETIME_SCRIPT], ++ cwd=import_root, ++ env=env, ++ capture_output=True, ++ text=True, ++ timeout=60, ++ check=False, ++ ) ++ ++ assert completed.returncode == 0, ( ++ "balanced PLE storage destruction failed " ++ f"(exit code {completed.returncode})\n" ++ f"stdout:\n{completed.stdout}\n" ++ f"stderr:\n{completed.stderr}" ++ ) +diff --git a/tests/v1/ple_offload/test_shutdown_order.py b/tests/v1/ple_offload/test_shutdown_order.py +new file mode 100644 +index 0000000..101c3da +--- /dev/null ++++ b/tests/v1/ple_offload/test_shutdown_order.py +@@ -0,0 +1,171 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import contextlib ++from types import SimpleNamespace ++from typing import Any ++ ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload import worker as ple_worker ++from vllm.v1.worker import gpu_worker ++ ++ ++def test_shutdown_marks_ple_stopping_before_teardown(monkeypatch) -> None: ++ worker = gpu_worker.Worker.__new__(gpu_worker.Worker) ++ worker._ple_offload_stopping = False ++ order: list[tuple[str, bool]] = [] ++ ++ def record(name: str): ++ def call() -> None: ++ order.append((name, worker._ple_offload_stopping)) ++ ++ return call ++ ++ monkeypatch.setattr(gpu_worker.gc, "unfreeze", record("gc")) ++ monkeypatch.setattr( ++ gpu_worker, ++ "ensure_kv_transfer_shutdown", ++ record("kv"), ++ ) ++ monkeypatch.setattr( ++ gpu_worker, ++ "ensure_ec_transfer_shutdown", ++ record("ec"), ++ ) ++ monkeypatch.setattr( ++ gpu_worker.current_platform, ++ "is_cuda_alike", ++ lambda: False, ++ ) ++ ++ worker.profiler = None ++ worker.weight_transfer_engine = None ++ worker.elastic_ep_executor = SimpleNamespace(shutdown=record("elastic")) ++ worker._ple_offload_worker_handle = SimpleNamespace(close=record("offload")) ++ worker.model_runner = SimpleNamespace(shutdown=record("runner")) ++ ++ worker.shutdown() ++ ++ assert order ++ assert all(stopping for _, stopping in order) ++ assert [name for name, _ in order][:2] == ["gc", "kv"] ++ ++ ++class _RequestThread: ++ def __init__(self, order: list[str], *, stops: bool) -> None: ++ self.order = order ++ self.stops = stops ++ self.alive = True ++ ++ def is_alive(self) -> bool: ++ return self.alive ++ ++ def join(self, timeout: float) -> None: ++ self.order.append("thread.join") ++ if self.stops: ++ self.alive = False ++ ++ ++def _record(order: list[str], name: str): ++ def call(*_args: Any, **_kwargs: Any) -> None: ++ order.append(name) ++ ++ return call ++ ++ ++def _make_connector_for_close( ++ order: list[str], *, thread_stops: bool ++) -> PleOffloadConnector: ++ connector = PleOffloadConnector.__new__(PleOffloadConnector) ++ connector._request_thread = _RequestThread(order, stops=thread_stops) ++ connector._request_queue = SimpleNamespace(put=_record(order, "queue.stop")) ++ connector._pinned_input_buffers = [] ++ connector._mapped_host_buffers = [] ++ connector._d2h_event_pool = object() ++ connector._registration_socket = SimpleNamespace( ++ close=_record(order, "socket.close") ++ ) ++ connector._zmq_ctx = SimpleNamespace(term=_record(order, "context.term")) ++ connector._progress = SimpleNamespace(close=_record(order, "progress.close")) ++ connector.release_host_storage = _record(order, "storage.release") ++ return connector ++ ++ ++def test_connector_close_releases_storage_last() -> None: ++ order: list[str] = [] ++ connector = _make_connector_for_close(order, thread_stops=True) ++ ++ connector.close() ++ ++ assert order[-1] == "storage.release" ++ assert order.index("thread.join") < order.index("context.term") ++ assert order.index("context.term") < order.index("storage.release") ++ ++ ++def test_connector_close_defers_storage_release_if_thread_does_not_stop() -> None: ++ order: list[str] = [] ++ connector = _make_connector_for_close(order, thread_stops=False) ++ ++ connector.close() ++ ++ assert "storage.release" not in order ++ assert "socket.close" not in order ++ assert "context.term" not in order ++ assert order[-1] == "progress.close" ++ ++ ++def test_offload_child_unlinks_owned_ipc_after_context_term(monkeypatch) -> None: ++ order: list[str] = [] ++ ipc_addr = "ipc:///tmp/ple-owned-test.sock" ++ pull_socket = SimpleNamespace( ++ bind=_record(order, "socket.bind"), ++ close=_record(order, "socket.close"), ++ ) ++ zmq_context = SimpleNamespace( ++ socket=lambda _kind: pull_socket, ++ term=_record(order, "context.term"), ++ ) ++ ++ class FakeRunner: ++ layer_names: list[str] = [] ++ ++ def __init__(self, _config: object) -> None: ++ pass ++ ++ def accept_registrations(self, _socket: object, _count: int) -> None: ++ order.append("registrations.accept") ++ ++ def busy_loop(self, _socket: object, _event: object) -> None: ++ order.append("runner.busy_loop") ++ ++ ready_pipe = SimpleNamespace( ++ send_bytes=_record(order, "ready.send"), ++ close=_record(order, "ready.close"), ++ ) ++ monkeypatch.setattr(ple_worker, "decorate_logs", lambda _name: None) ++ monkeypatch.setattr(ple_worker, "mark_as_offload_worker", lambda: None) ++ monkeypatch.setattr(ple_worker, "_prepare_cuda_free_process", lambda: None) ++ monkeypatch.setattr(ple_worker, "_init_offload_distributed", lambda: None) ++ monkeypatch.setattr( ++ ple_worker, ++ "set_current_vllm_config", ++ lambda _config: contextlib.nullcontext(), ++ ) ++ monkeypatch.setattr(ple_worker, "PleOffloadRunner", FakeRunner) ++ monkeypatch.setattr(ple_worker.zmq, "Context", lambda: zmq_context) ++ monkeypatch.setattr(ple_worker.signal, "signal", lambda *_args: None) ++ monkeypatch.setattr(ple_worker.torch.cuda, "is_initialized", lambda: False) ++ monkeypatch.setattr( ++ ple_worker, ++ "_unlink_ipc_socket", ++ lambda address: order.append(f"ipc.unlink:{address}"), ++ ) ++ ++ ple_worker.PleOffloadWorker.proc_main( ++ SimpleNamespace(), ++ num_workers=1, ++ ipc_addr=ipc_addr, ++ ready_pipe=ready_pipe, ++ ) ++ ++ assert order.index("socket.close") < order.index("context.term") ++ assert order[-1] == f"ipc.unlink:{ipc_addr}" +diff --git a/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py +new file mode 100644 +index 0000000..c709a93 +--- /dev/null ++++ b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py +@@ -0,0 +1,352 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import ctypes ++import ctypes.util ++import inspect ++import json ++import multiprocessing ++import os ++import pickle ++import signal ++import socket ++import subprocess ++import sys ++import threading ++import time ++from multiprocessing.connection import Connection ++from pathlib import Path ++ ++import pytest ++ ++def _native_cuda_probe(output_path: str) -> None: ++ """Run during config unpickling, before the bootstrap imports vLLM.""" ++ library = ctypes.util.find_library("cudart") ++ if library is None: ++ os._exit(77) ++ cudart = ctypes.CDLL(library) ++ result = int(cudart.cudaFree(ctypes.c_void_p())) ++ Path(output_path).write_text( ++ json.dumps( ++ { ++ "cuda_visible_devices": os.environ.get( ++ "CUDA_VISIBLE_DEVICES" ++ ), ++ "cuda_result": result, ++ "torch_imported": "torch" in sys.modules, ++ "vllm_imported": "vllm" in sys.modules, ++ } ++ ), ++ encoding="utf-8", ++ ) ++ os._exit(0) ++ ++ ++class _ProbePayload: ++ def __init__(self, output_path: str) -> None: ++ self.output_path = output_path ++ ++ def __reduce__(self): ++ # The bootstrap interpreter cannot import this test module, so rebuild ++ # the probe from its source text through a builtin instead of by name. ++ source = ( ++ "import ctypes, ctypes.util, json, os, sys\n" ++ "from pathlib import Path\n" ++ + inspect.getsource(_native_cuda_probe) ++ + f"_native_cuda_probe({self.output_path!r})\n" ++ ) ++ return exec, (source, {}) ++ ++ ++def test_spawn_bootstrap_sets_visibility_before_native_cuda( ++ tmp_path: Path, ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ if ctypes.util.find_library("cudart") is None: ++ pytest.skip("CUDA runtime is not installed") ++ ++ original = os.environ.get("CUDA_VISIBLE_DEVICES") ++ output_path = tmp_path / "probe.json" ++ ++ # Use the production Popen launcher. Config reconstruction executes the ++ # probe in the fresh bootstrap interpreter before worker imports. ++ with monkeypatch.context() as child_settings: ++ child_settings.setenv("CUDA_VISIBLE_DEVICES", "0") ++ handle = PleOffloadWorker.make_process( ++ _ProbePayload(str(output_path)), 0, "unused" ++ ) ++ try: ++ returncode = handle.proc.wait(timeout=120) ++ finally: ++ handle.close() ++ ++ if returncode == 77: ++ pytest.skip("spawned child could not load the CUDA runtime") ++ assert returncode == 0 ++ ++ result = json.loads(output_path.read_text(encoding="utf-8")) ++ assert result["cuda_visible_devices"] == "0" ++ assert result["cuda_result"] == 0 ++ assert result["torch_imported"] is False ++ assert result["vllm_imported"] is False ++ assert os.environ.get("CUDA_VISIBLE_DEVICES") == original ++ ++ ++class _FreshInterpreterProbe: ++ def __init__(self, output_path: str) -> None: ++ self.output_path = output_path ++ ++ def __reduce__(self): ++ expression = ( ++ "(__import__('pathlib').Path(" ++ + repr(self.output_path) ++ + ").write_text(__import__('json').dumps({" ++ "'cuda_visible_devices': __import__('os').environ.get(" ++ "'CUDA_VISIBLE_DEVICES')," ++ "'torch_imported': 'torch' in __import__('sys').modules," ++ "'vllm_imported': 'vllm' in __import__('sys').modules" ++ "}), encoding='utf-8'), None)[1]" ++ ) ++ return eval, (expression,) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++def test_fresh_interpreter_has_child_environment_before_vllm_import( ++ tmp_path: Path, ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ # This test module has unconditional top-level imports. The dedicated ++ # bootstrap interpreter must not import it as __mp_main__. ++ original = os.environ.get("CUDA_VISIBLE_DEVICES") ++ output_path = tmp_path / "fresh-interpreter.json" ++ ++ with monkeypatch.context() as child_settings: ++ child_settings.setenv("CUDA_VISIBLE_DEVICES", "child-only") ++ handle = PleOffloadWorker.make_process( ++ _FreshInterpreterProbe(str(output_path)), 0, "unused" ++ ) ++ try: ++ handle.proc.wait(timeout=120) ++ finally: ++ handle.close() ++ ++ assert output_path.exists() ++ result = json.loads(output_path.read_text(encoding="utf-8")) ++ assert result["cuda_visible_devices"] == "child-only" ++ assert result["torch_imported"] is False ++ assert result["vllm_imported"] is False ++ assert os.environ.get("CUDA_VISIBLE_DEVICES") == original ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++def test_launch_payload_send_has_absolute_deadline() -> None: ++ reader, sender = socket.socketpair() ++ sender.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024) ++ code = ( ++ "import sys,time\n" ++ "from vllm.v1.ple_offload.worker import _send_framed_payload\n" ++ "payload=b'x'*(64<<20)\n" ++ "start=time.monotonic()\n" ++ "try:\n" ++ " _send_framed_payload(int(sys.argv[1]),payload,start+0.25)\n" ++ "except TimeoutError:\n" ++ " print('TIMED_OUT',time.monotonic()-start,flush=True)\n" ++ "else:\n" ++ " raise AssertionError('blocked reader unexpectedly accepted payload')\n" ++ ) ++ process = subprocess.Popen( ++ [sys.executable, "-c", code, str(sender.fileno())], ++ pass_fds=(sender.fileno(),), ++ close_fds=True, ++ stdout=subprocess.PIPE, ++ stderr=subprocess.PIPE, ++ text=True, ++ ) ++ sender.close() ++ try: ++ stdout, stderr = process.communicate(timeout=120) ++ finally: ++ reader.close() ++ if process.poll() is None: ++ process.terminate() ++ process.wait(timeout=5) ++ assert process.returncode == 0, stderr ++ # The injector logs the send entry through vLLM's logger on stdout; read only ++ # the probe's own result line. ++ result = [line for line in stdout.splitlines() if line.startswith('TIMED_OUT ')] ++ assert len(result) == 1, stdout ++ marker, elapsed = result[0].split() ++ assert marker == "TIMED_OUT" ++ assert float(elapsed) < 5 ++ ++ ++class _BlockingConfigProbe: ++ def __init__(self, output_path: str) -> None: ++ self.output_path = output_path ++ ++ def __reduce__(self): ++ expression = ( ++ "(__import__('pathlib').Path(" ++ + repr(self.output_path) ++ + ").write_text(str(__import__('os').getpid()), encoding='utf-8')," ++ "__import__('threading').Event().wait())[1]" ++ ) ++ return eval, (expression,) ++ ++ ++def _launch_blocked_config(output_path: str, child_pids) -> None: ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ handle = PleOffloadWorker.make_process( ++ _BlockingConfigProbe(output_path), 0, "unused" ++ ) ++ child_pids.put(handle.proc.pid) ++ # The test kills this launcher parent. Do not close the handle here. ++ threading.Event().wait() ++ ++ ++def _pid_has_exited(pid: int) -> bool: ++ stat_path = Path(f"/proc/{pid}/stat") ++ if not stat_path.exists(): ++ return True ++ try: ++ return stat_path.read_text(encoding="utf-8").split()[2] == "Z" ++ except FileNotFoundError: ++ return True ++ ++ ++@pytest.mark.skipif(sys.platform != "linux", reason="uses Linux process state") ++def test_parent_death_during_config_reconstruction_exits_child( ++ tmp_path: Path, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ child_pids = context.Queue() ++ marker = tmp_path / "blocked-config.pid" ++ parent = context.Process( ++ target=_launch_blocked_config, ++ args=(str(marker), child_pids), ++ ) ++ parent.start() ++ child_pid = child_pids.get(timeout=120) ++ deadline = time.monotonic() + 120 ++ while not marker.exists() and time.monotonic() < deadline: ++ time.sleep(0.05) ++ assert marker.exists() ++ ++ os.kill(parent.pid, signal.SIGKILL) ++ parent.join(timeout=120) ++ assert not parent.is_alive() ++ ++ deadline = time.monotonic() + 10 ++ while not _pid_has_exited(child_pid) and time.monotonic() < deadline: ++ time.sleep(0.05) ++ if not _pid_has_exited(child_pid): ++ # Harness cleanup is performed only after the autonomous-exit assertion ++ # has failed. ++ os.kill(child_pid, signal.SIGKILL) ++ pytest.fail("PLE child survived launcher-parent death during unpickle") ++ ++ ++class _ExecFdProbe: ++ def __init__( ++ self, ++ output_path: str, ++ inherited: list[tuple[int, int]], ++ ) -> None: ++ self.output_path = output_path ++ self.inherited = inherited ++ ++ def __reduce__(self): ++ helper = ( ++ "import json,os,sys\n" ++ "from pathlib import Path\n" ++ "pairs=json.loads(sys.argv[3])\n" ++ "retained=[]\n" ++ "for fd,inode in pairs:\n" ++ " try:\n" ++ " retained.append(os.fstat(fd).st_ino==inode)\n" ++ " except OSError:\n" ++ " retained.append(False)\n" ++ "Path(sys.argv[1]).write_text(json.dumps({" ++ "'noninheritable_before_exec':sys.argv[2]=='True'," ++ "'closed_after_exec':not any(retained)}),encoding='utf-8')\n" ++ ) ++ pairs = json.dumps(self.inherited) ++ expression = ( ++ "__import__('os').execv(__import__('sys').executable," ++ "[__import__('sys').executable,'-c'," ++ + repr(helper) ++ + "," ++ + repr(self.output_path) ++ + ",str(all(not __import__('os').get_inheritable(int(fd)) " ++ "for fd in __import__('sys').argv[1:]))," ++ + repr(pairs) ++ + "])" ++ ) ++ return eval, (expression,) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++def test_actual_bootstrap_restores_cloexec_before_helper_exec( ++ tmp_path: Path, ++) -> None: ++ from vllm.v1.ple_offload.worker import _send_framed_payload ++ ++ control_parent, control_child = socket.socketpair() ++ death_parent, death_child = socket.socketpair() ++ inherited = [ ++ (control_child.fileno(), os.fstat(control_child.fileno()).st_ino), ++ (death_child.fileno(), os.fstat(death_child.fileno()).st_ino), ++ ] ++ output_path = tmp_path / "bootstrap-fds.json" ++ config_payload = pickle.dumps( ++ _ExecFdProbe(str(output_path), inherited), ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ payload = pickle.dumps( ++ { ++ "vllm_config_payload": config_payload, ++ "num_workers": 0, ++ "ipc_addr": "unused", ++ }, ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ control = Connection(control_parent.detach()) ++ death_writer = Connection(death_parent.detach()) ++ process = subprocess.Popen( ++ [ ++ sys.executable, ++ "-m", ++ "vllm_ple_offload_bootstrap", ++ str(control_child.fileno()), ++ str(death_child.fileno()), ++ ], ++ close_fds=True, ++ pass_fds=(control_child.fileno(), death_child.fileno()), ++ ) ++ control_child.close() ++ death_child.close() ++ try: ++ _send_framed_payload( ++ control.fileno(), payload, time.monotonic() + 10 ++ ) ++ assert process.wait(timeout=120) == 0 ++ finally: ++ control.close() ++ death_writer.close() ++ if process.poll() is None: ++ process.terminate() ++ try: ++ process.wait(timeout=5) ++ except subprocess.TimeoutExpired: ++ process.kill() ++ process.wait(timeout=5) ++ ++ result = json.loads(output_path.read_text(encoding="utf-8")) ++ assert result == { ++ "noninheritable_before_exec": True, ++ "closed_after_exec": True, ++ } +diff --git a/tests/v1/ple_offload/test_test_fault.py b/tests/v1/ple_offload/test_test_fault.py +new file mode 100644 +index 0000000..3dc8470 +--- /dev/null ++++ b/tests/v1/ple_offload/test_test_fault.py +@@ -0,0 +1,293 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import contextlib ++import os ++import queue ++import sys ++from dataclasses import FrozenInstanceError ++from pathlib import Path ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import pytest ++ ++ ++@pytest.fixture ++def fault_module(monkeypatch, tmp_path): ++ monkeypatch.delenv("VLLM_PLE_TEST_FAULT", raising=False) ++ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) ++ import vllm_ple_test_fault as test_fault ++ ++ return test_fault ++ ++ ++def test_disarmed_import_is_none(fault_module): ++ assert fault_module._FAULT is None ++ ++ ++@pytest.mark.parametrize( ++ "value", ++ [ ++ "cut=exc_before_prepare,ordinal=60,rank=any", ++ "cut=exc_after_prepare,ordinal=2,tp=1", ++ "cut=exc_in_release,ordinal=3,pp=0", ++ "cut=park_before_publish,ordinal=4,rank=any,phase=any", ++ ], ++) ++def test_valid_specs(fault_module, value): ++ spec = fault_module._parse_fault_spec(value) ++ assert spec is not None ++ assert spec.ordinal > 0 ++ ++ ++@pytest.mark.parametrize( ++ "value", ++ [ ++ "", ++ " cut=exc_in_release,ordinal=1,rank=any", ++ "cut=unknown,ordinal=1,rank=any", ++ "cut=exc_in_release,ordinal=0,rank=any", ++ "cut=exc_in_release,ordinal=01,rank=any", ++ "cut=exc_in_release,ordinal=1", ++ "cut=exc_in_release,ordinal=1,rank=tp1", ++ "cut=exc_in_release,ordinal=1,rank=any,tp=0", ++ "cut=exc_in_release,ordinal=1,tp=-1", ++ "cut=exc_in_release,ordinal=1,tp=0,phase=startup", ++ "cut=exc_in_release,ordinal=1,tp=0,extra=x", ++ "cut=exc_in_release,cut=exc_after_release,ordinal=1,tp=0", ++ ], ++) ++def test_malformed_specs_raise(fault_module, value): ++ with pytest.raises(ValueError): ++ fault_module._parse_fault_spec(value) ++ ++ ++def test_one_shot_rank_phase_and_ordinal(fault_module, monkeypatch): ++ spec = fault_module._parse_fault_spec( ++ "cut=exc_after_prepare,ordinal=2,tp=1,pp=0" ++ ) ++ assert spec is not None ++ ++ monkeypatch.setattr(fault_module.os, "getpid", lambda: 123) ++ monkeypatch.setattr(fault_module, "_pid_start_time", lambda _pid=None: 789) ++ monkeypatch.setattr(fault_module.time, "monotonic", lambda: 12.5) ++ monkeypatch.setattr(fault_module.time, "time", lambda: 456.25) ++ ++ def fires(cut, tp_rank=1, eligible=True): ++ return spec.fires( ++ cut, ++ role="gpu", ++ unit="runner-forward", ++ eligible=eligible, ++ tp_rank=tp_rank, ++ pp_rank=0, ++ dp_rank=0, ++ transaction_id="sequence:9@begin:12.000000000", ++ correlation_id="sequence:9", ++ ) ++ ++ assert not fires("other") ++ assert not fires("exc_after_prepare", tp_rank=0) ++ assert not fires("exc_after_prepare", eligible=False) ++ assert not fires("exc_after_prepare") ++ assert fires("exc_after_prepare") ++ assert not fires("exc_after_prepare") ++ line = Path(os.environ["VLLM_PLE_TEST_FAULT_LOG"]).read_text() ++ assert line == ( ++ "PLE_TEST_FAULT v=1 cut=exc_after_prepare ordinal=2 count=2 " ++ "unit=runner-forward role=gpu tp=1 pp=0 dp=0 pid=123 " ++ "pid_start=789 transaction=sequence:9@begin:12.000000000 " ++ "correlation=sequence:9 monotonic=12.500000000 " ++ "wall=456.250000000\n" ++ ) ++ with pytest.raises(FrozenInstanceError): ++ spec.ordinal = 3 ++ ++ ++def test_disarmed_connector_sites_continue(fault_module, monkeypatch): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ monkeypatch.setattr(connector_module, "_FAULT", None) ++ ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector.device = SimpleNamespace(index=0) ++ connector._d2h_event_pool = queue.Queue() ++ event = Mock() ++ connector._d2h_event_pool.put_nowait(event) ++ connector._request_queue = Mock() ++ connector._progress = SimpleNamespace(sequence=0) ++ connector._enqueue_cuda_inputs = Mock() ++ ++ connector._launch( ++ num_reqs=1, ++ num_tokens=2, ++ serving=True, ++ transaction_id=None, ++ correlation_id=None, ++ ) ++ connector._enqueue_cuda_inputs.assert_called_once() ++ connector._request_queue.put_nowait.assert_called_once() ++ ++ staged = connector_module._PendingPleOffloadRequest( ++ connector_module.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1), ++ event, ++ True, ++ "sequence:0", ++ ) ++ socket = Mock() ++ monkeypatch.setattr( ++ connector_module.torch.accelerator, ++ "device_index", ++ lambda _index: contextlib.nullcontext(), ++ ) ++ connector._process_request(staged, socket) ++ event.synchronize.assert_called_once() ++ socket.send.assert_called_once() ++ ++ ++def test_disarmed_release_site_continues(fault_module, monkeypatch): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ monkeypatch.setattr(connector_module, "_FAULT", None) ++ monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) ++ ++ stream = Mock() ++ completion = Mock() ++ monkeypatch.setattr( ++ connector_module.torch.cuda, ++ "current_stream", ++ lambda _device: stream, ++ ) ++ monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: completion) ++ ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.tp_rank = 0 ++ connector.pp_rank = 0 ++ connector.device = SimpleNamespace(index=0) ++ connector._layers = {} ++ connector._after_dummy = False ++ connector._progress = Mock() ++ ++ connector._release_outputs_impl() ++ completion.record.assert_called_once_with(stream) ++ connector._progress.track_completion.assert_called_once_with(completion) ++ ++ ++def test_disarmed_worker_park_site_continues(fault_module, monkeypatch): ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ monkeypatch.setattr(worker_module, "_FAULT", None) ++ runner = object.__new__(worker_module.PleOffloadRunner) ++ runner._validate_requests_before_publication = Mock() ++ runner._worker_targets = {0: {}} ++ runner._input_bufs = {} ++ runner._scratch_bufs = {} ++ runner._layers = {} ++ runner._clamp_input_ids = False ++ runner._step = 0 ++ runner._request_index = 0 ++ request = SimpleNamespace(dp_rank=0) ++ ++ runner._handle_requests_impl([request]) ++ runner._validate_requests_before_publication.assert_called_once_with([request]) ++ assert runner._step == 1 ++ ++ ++def test_all_injection_sites_use_the_disarmed_short_circuit(): ++ root = Path(__file__).resolve().parents[3] ++ expected = { ++ "vllm/v1/ple_offload/connector.py": 3, ++ "vllm/v1/ple_offload/worker.py": 1, ++ "vllm/v1/worker/gpu/model_runner.py": 4, ++ "vllm_ple_offload_bootstrap/__init__.py": 1, ++ } ++ for relative, count in expected.items(): ++ lines = (root / relative).read_text(encoding="utf-8").splitlines() ++ sites = [i for i, line in enumerate(lines) if "_FAULT.fires(" in line] ++ assert len(sites) == count, relative ++ for i in sites: ++ window = "\n".join(lines[max(0, i - 6): i + 1]) ++ # The None check must precede fires() in the same condition, so a ++ # disarmed process never calls into the fault module. ++ assert "_FAULT is not None" in window, (relative, i + 1) ++ assert window.index("_FAULT is not None") < window.index("_FAULT.fires("), (relative, i + 1) ++ ++ ++class _FatalReached(Exception): ++ pass ++ ++ ++def test_exc_in_release_reaches_real_connector_fatal_wrapper( ++ fault_module, monkeypatch ++): ++ import vllm.v1.ple_offload.connector as connector_module ++ ++ armed = fault_module._parse_fault_spec( ++ "cut=exc_in_release,ordinal=1,tp=1,pp=0" ++ ) ++ assert armed is not None ++ monkeypatch.setattr(connector_module, "_FAULT", armed) ++ ++ fatal = Mock(side_effect=_FatalReached) ++ monkeypatch.setattr(connector_module, "ple_fail_closed", fatal) ++ ++ connector = object.__new__(connector_module.PleOffloadConnector) ++ connector.rank = 1 ++ connector.tp_rank = 1 ++ connector.pp_rank = 0 ++ connector.dp_rank = 0 ++ connector._layers = {} ++ connector._after_dummy = False ++ connector._progress = SimpleNamespace(sequence=9) ++ ++ with pytest.raises(_FatalReached): ++ connector.release_outputs( ++ serving=True, ++ transaction_id="sequence:9@begin:1.000000000", ++ correlation_id="sequence:9", ++ ) ++ ++ assert fatal.call_args.kwargs["where"] == "release_outputs" ++ assert fatal.call_args.kwargs["rank"] == 1 ++ assert isinstance(fatal.call_args.kwargs["error"], RuntimeError) ++ ++ ++@pytest.mark.timeout(180) ++@pytest.mark.skipif( ++ sys.platform == "win32", ++ reason="production PLE launcher requires POSIX pass_fds", ++) ++@pytest.mark.parametrize("cut", ["launch_partial_ready", "launch_partial_ready_header"]) ++def test_partial_ready_times_out_and_terminates_production_child( ++ monkeypatch, ++ tmp_path, ++ cut, ++): ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ monkeypatch.setenv( ++ "VLLM_PLE_TEST_FAULT", ++ f"cut={cut},ordinal=1,rank=any", ++ ) ++ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) ++ monkeypatch.setattr( ++ worker_module.envs, ++ "VLLM_PLE_OFFLOAD_READY_TIMEOUT", ++ 1.0, ++ ) ++ ++ handle = worker_module.PleOffloadWorker.make_process( ++ object(), # The bootstrap fault fires before config reconstruction. ++ num_workers=1, ++ ipc_addr="ipc:///unused-ple-test", ++ ) ++ try: ++ with pytest.raises(TimeoutError): ++ worker_module.PleOffloadWorker.wait_for_ready(handle) ++ assert handle.proc.poll() is not None ++ finally: ++ handle.close(graceful=False) +diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py +new file mode 100644 +index 0000000..fb37122 +--- /dev/null ++++ b/tests/v1/ple_offload/test_transport_bounded_io.py +@@ -0,0 +1,343 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import multiprocessing ++import os ++import signal ++import struct ++import subprocess ++import sys ++import time ++from pathlib import Path ++from types import SimpleNamespace ++from typing import Any ++ ++import pytest ++ ++ ++class _BlockingStrError(Exception): ++ def __str__(self) -> str: ++ __import__("threading").Event().wait() ++ return "unreachable" ++ ++ ++def _blocking_fatal_formatter_case(marker_path: str) -> None: ++ from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++ Path(marker_path).write_text("armed", encoding="utf-8") ++ try: ++ raise _BlockingStrError() ++ except BaseException as error: ++ ple_fail_closed( ++ where="blocking-formatter", ++ rank=1, ++ layer="ple", ++ sequence=3, ++ error=error, ++ ) ++ ++ ++def test_fatal_helper_does_not_call_blocking_exception_str( ++ tmp_path: Path, ++) -> None: ++ context = multiprocessing.get_context("spawn") ++ marker = tmp_path / "fatal-formatter-armed" ++ process = context.Process( ++ target=_blocking_fatal_formatter_case, ++ args=(str(marker),), ++ ) ++ process.start() ++ try: ++ deadline = time.monotonic() + 120 ++ while not marker.exists() and time.monotonic() < deadline: ++ time.sleep(0.05) ++ assert marker.exists() ++ ++ process.join(timeout=5) ++ assert not process.is_alive() ++ assert process.exitcode not in (None, 0) ++ finally: ++ if process.is_alive(): ++ process.terminate() ++ process.join(timeout=120) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++@pytest.mark.parametrize( ++ ("behavior", "ready_timeout", "expected_exception"), ++ [ ++ pytest.param("partial", "1", TimeoutError, id="partial-frame"), ++ pytest.param("silent", "1", TimeoutError, id="no-response"), ++ pytest.param("ready", "120", None, id="complete-ready"), ++ pytest.param("invalid", "120", RuntimeError, id="invalid-layer-names"), ++ ], ++) ++def test_wait_for_ready_receives_complete_bounded_frame( ++ monkeypatch: pytest.MonkeyPatch, ++ behavior: str, ++ ready_timeout: str, ++ expected_exception: type[BaseException] | None, ++) -> None: ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ real_popen = subprocess.Popen ++ spawned: list[subprocess.Popen] = [] ++ child_code = ( ++ "import os,signal,struct,sys,time\n" ++ "from multiprocessing.connection import Connection\n" ++ "signal.alarm(10)\n" ++ "fd=int(sys.argv[1])\n" ++ "behavior=sys.argv[3]\n" ++ "if behavior=='partial':\n" ++ " os.write(fd,struct.pack('!i',128)+b'partial')\n" ++ "elif behavior=='ready':\n" ++ " Connection(fd).send({'status':'READY','layer_names':['ple']})\n" ++ "elif behavior=='invalid':\n" ++ " Connection(fd).send({'status':'READY','layer_names':'ple'})\n" ++ "time.sleep(120)\n" ++ ) ++ ++ def launch_fixture(argv: list[str], **kwargs: Any) -> subprocess.Popen: ++ process = real_popen( ++ [ ++ sys.executable, ++ "-c", ++ child_code, ++ argv[-2], ++ argv[-1], ++ behavior, ++ ], ++ **kwargs, ++ ) ++ spawned.append(process) ++ return process ++ ++ monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", ready_timeout) ++ monkeypatch.setattr(worker_module.subprocess, "Popen", launch_fixture) ++ ++ handle = worker_module.PleOffloadWorker.make_process(None, 0, "unused") ++ try: ++ started = time.monotonic() ++ if expected_exception is None: ++ worker_module.PleOffloadWorker.wait_for_ready(handle) ++ assert handle.ready_pipe_reader is None ++ else: ++ with pytest.raises(expected_exception): ++ worker_module.PleOffloadWorker.wait_for_ready(handle) ++ if expected_exception is TimeoutError: ++ assert time.monotonic() - started < 5 ++ assert spawned[0].poll() is not None ++ finally: ++ handle.close(graceful=False) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++def test_make_process_large_payload_send_has_absolute_deadline( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ import vllm.v1.ple_offload.worker as worker_module ++ ++ real_popen = subprocess.Popen ++ spawned: list[subprocess.Popen] = [] ++ ++ def launch_nonreader(argv: list[str], **kwargs: Any) -> subprocess.Popen: ++ process = real_popen( ++ [ ++ sys.executable, ++ "-c", ++ "import signal,time; signal.alarm(120); time.sleep(120)", ++ argv[-2], ++ argv[-1], ++ ], ++ **kwargs, ++ ) ++ spawned.append(process) ++ return process ++ ++ monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "5") ++ monkeypatch.setattr(worker_module.subprocess, "Popen", launch_nonreader) ++ ++ started = time.monotonic() ++ with pytest.raises(TimeoutError, match="transmission"): ++ worker_module.PleOffloadWorker.make_process( ++ b"x" * (64 << 20), ++ 0, ++ "unused", ++ ) ++ ++ assert time.monotonic() - started < 20 ++ assert spawned ++ assert spawned[0].poll() is not None ++ ++ ++@pytest.mark.parametrize( ++ "identities", ++ [ ++ pytest.param(("shared", "shared"), id="duplicate"), ++ pytest.param(("",), id="unbound"), ++ ], ++) ++def test_expected_ple_layers_are_bound_and_unique_before_publication( ++ monkeypatch: pytest.MonkeyPatch, ++ identities: tuple[str, ...], ++) -> None: ++ from vllm.config import CUDAGraphMode ++ from vllm.v1.worker.gpu.cudagraph_utils import ( ++ CudaGraphManager, ++ ModelCudaGraphManager, ++ ) ++ ++ manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) ++ manager.graphs = {} ++ manager.use_breakable_cg = False ++ manager.cudagraph_mode = CUDAGraphMode.FULL ++ ++ modules = [ ++ ( ++ f"module_{index}", ++ SimpleNamespace( ++ _is_cpu_offloaded=True, ++ _host_output_ptr=1, ++ _ple_offload_name=identity, ++ ), ++ ) ++ for index, identity in enumerate(identities) ++ ] ++ model = SimpleNamespace(named_modules=lambda: iter(modules)) ++ ++ def publish_graph(self: Any, *_args: Any, **_kwargs: Any) -> None: ++ self.graphs["published"] = object() ++ ++ monkeypatch.setattr(CudaGraphManager, "capture", publish_graph) ++ ++ with pytest.raises(RuntimeError, match="distinct, bound"): ++ ModelCudaGraphManager.capture( ++ manager, ++ model=model, ++ model_state=object(), ++ input_buffers=object(), ++ intermediate_tensors=None, ++ block_tables=object(), ++ attn_groups=[], ++ kv_cache_config=object(), ++ ) ++ assert manager.graphs == {} ++ ++ ++def test_real_wait_copy_geometry_rejected_before_graph_publication( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ import torch ++ ++ import vllm.model_executor.layers.ple_offload_layer as ple_layer_module ++ from vllm.config import CUDAGraphMode ++ from vllm.v1.worker.gpu.cudagraph_utils import ( ++ BatchExecutionDescriptor, ++ CudaGraphManager, ++ ) ++ ++ manager = CudaGraphManager.__new__(CudaGraphManager) ++ manager.graphs = {} ++ manager._ple_full_copy_geometry = {} ++ manager._expected_ple_copy_layers = ("ple",) ++ manager._allow_ple_full_graph = True ++ desc = BatchExecutionDescriptor( ++ cg_mode=CUDAGraphMode.FULL, ++ num_tokens=8, ++ num_reqs=1, ++ ) ++ ++ success = SimpleNamespace(value=0) ++ monkeypatch.setattr( ++ ple_layer_module.torch.cuda, ++ "current_stream", ++ lambda: SimpleNamespace(cuda_stream=1), ++ ) ++ monkeypatch.setattr( ++ ple_layer_module.cuda_driver, ++ "CUstream", ++ lambda value: value, ++ ) ++ monkeypatch.setattr( ++ ple_layer_module.cuda_driver, ++ "CUdeviceptr", ++ lambda value: value, ++ ) ++ monkeypatch.setattr( ++ ple_layer_module.cuda_driver, ++ "cuStreamWaitValue32", ++ lambda *_args: (success,), ++ ) ++ monkeypatch.setattr( ++ ple_layer_module.cuda_driver, ++ "cuMemcpyHtoDAsync", ++ lambda *_args: (success,), ++ ) ++ ++ gpu_output_buffer = torch.empty(12, 4, dtype=torch.float32) ++ hidden_states = torch.empty(12, 1, dtype=torch.float32) ++ ++ def capture_then_publish() -> None: ++ with manager.capture_ple_copy_geometry(desc): ++ ple_layer_module._ple_offload_wait_copy_impl( ++ gpu_output_buffer, ++ hidden_states, ++ 1, ++ 2, ++ 16, ++ 12, ++ "ple", ++ ) ++ manager.graphs[desc] = object() ++ ++ with pytest.raises(RuntimeError, match="actual PLE copy geometry"): ++ capture_then_publish() ++ assert desc not in manager._ple_full_copy_geometry ++ assert desc not in manager.graphs ++ ++ ++class _CleanupConfig: ++ def __init__(self, started_path: str, cleaned_path: str) -> None: ++ self.started_path = started_path ++ self.cleaned_path = cleaned_path ++ ++ def __reduce__(self): ++ source = ( ++ "import signal,threading\n" ++ "from pathlib import Path\n" ++ "from vllm.v1.ple_offload.worker import PleOffloadWorker\n" ++ f"_started=Path({self.started_path!r})\n" ++ f"_cleaned=Path({self.cleaned_path!r})\n" ++ "def _cleanup_proc_main(vllm_config,num_workers,ipc_addr,ready_pipe):\n" ++ " _stop=threading.Event()\n" ++ " signal.signal(signal.SIGTERM,lambda *_args:_stop.set())\n" ++ " _started.write_text('started',encoding='utf-8')\n" ++ " try:\n" ++ " _stop.wait()\n" ++ " finally:\n" ++ " _cleaned.write_text('cleaned',encoding='utf-8')\n" ++ "PleOffloadWorker.proc_main=staticmethod(_cleanup_proc_main)\n" ++ ) ++ return exec, (source, {}) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") ++def test_normal_close_runs_child_cleanup(tmp_path: Path) -> None: ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ started = tmp_path / "cleanup-started" ++ cleaned = tmp_path / "cleanup-finished" ++ handle = PleOffloadWorker.make_process( ++ _CleanupConfig(str(started), str(cleaned)), ++ 0, ++ "unused", ++ ) ++ try: ++ deadline = time.monotonic() + 120 ++ while not started.exists() and time.monotonic() < deadline: ++ time.sleep(0.05) ++ assert started.exists() ++ finally: ++ handle.close() ++ ++ assert handle.proc.returncode == 0 ++ assert cleaned.read_text(encoding="utf-8") == "cleaned" +diff --git a/tests/v1/ple_offload/test_transport_fatal_helper.py b/tests/v1/ple_offload/test_transport_fatal_helper.py +new file mode 100644 +index 0000000..4e5d5c6 +--- /dev/null ++++ b/tests/v1/ple_offload/test_transport_fatal_helper.py +@@ -0,0 +1,387 @@ ++# SPDX-License-Identifier: Apache-2.0 ++ ++import os ++import signal ++import subprocess ++import sys ++import time ++from pathlib import Path ++ ++import pytest ++ ++ ++def _run_bounded_child( ++ source: str, ++ tmp_path: Path, ++ *arguments: str, ++ expected_exitcode: int = 0, ++) -> None: ++ marker = tmp_path / f"started-{time.monotonic_ns()}" ++ process = subprocess.Popen( ++ [sys.executable, "-c", source, str(marker), *arguments], ++ stderr=subprocess.DEVNULL, ++ start_new_session=os.name == "posix", ++ ) ++ stayed_alive = False ++ try: ++ deadline = time.monotonic() + 120 ++ while not marker.exists() and time.monotonic() < deadline: ++ time.sleep(0.05) ++ assert marker.exists(), ( ++ "child did not finish cold startup within 120 seconds" ++ ) ++ process.wait(timeout=5) ++ except subprocess.TimeoutExpired: ++ stayed_alive = True ++ finally: ++ if process.poll() is None: ++ if os.name == "posix": ++ os.killpg(process.pid, signal.SIGKILL) ++ else: ++ process.terminate() ++ process.wait(timeout=120) ++ ++ assert not stayed_alive, "child stayed alive after the bounded operation" ++ assert process.returncode == expected_exitcode ++ ++ ++_FATAL_SOURCE = r""" ++import os ++import sys ++import threading ++from pathlib import Path ++ ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++scenario = sys.argv[2] ++ ++if scenario == "blocking-name-property": ++ class Meta(type): ++ @property ++ def __name__(cls): ++ threading.Event().wait() ++ ++ class Error(Exception, metaclass=Meta): ++ pass ++ ++ error = Error() ++elif scenario == "blocking-str-descriptor": ++ class Descriptor: ++ def __get__(self, instance, owner): ++ threading.Event().wait() ++ ++ class Error(Exception): ++ __str__ = Descriptor() ++ ++ error = Error() ++elif scenario == "blocking-name-format": ++ class BlockingName(str): ++ def __format__(self, format_spec): ++ threading.Event().wait() ++ ++ class Meta(type): ++ @property ++ def __name__(cls): ++ return BlockingName("Error") ++ ++ class Error(Exception, metaclass=Meta): ++ pass ++ ++ error = Error() ++elif scenario == "blocking-str-method": ++ class Error(Exception): ++ def __str__(self): ++ threading.Event().wait() ++ ++ error = Error() ++elif scenario == "overridden-getattribute": ++ class Error(Exception): ++ def __getattribute__(self, name): ++ threading.Event().wait() ++ ++ error = Error("detail") ++elif scenario == "str-subclass-argument": ++ class BlockingString(str): ++ def __format__(self, format_spec): ++ threading.Event().wait() ++ ++ def __str__(self): ++ threading.Event().wait() ++ ++ error = RuntimeError(BlockingString("detail")) ++elif scenario == "huge-string": ++ error = RuntimeError("x" * (32 << 20)) ++else: ++ raise AssertionError(scenario) ++ ++Path(sys.argv[1]).write_text("armed", encoding="utf-8") ++ple_fail_closed( ++ where="transport-fatal", ++ rank=1, ++ layer="ple", ++ sequence=3, ++ error=error, ++) ++""" ++ ++ ++@pytest.mark.parametrize( ++ "scenario", ++ [ ++ "blocking-name-property", ++ "blocking-str-descriptor", ++ "blocking-name-format", ++ ], ++) ++def test_fatal_helper_avoids_exception_controlled_descriptors( ++ tmp_path: Path, ++ scenario: str, ++) -> None: ++ _run_bounded_child( ++ _FATAL_SOURCE, ++ tmp_path, ++ scenario, ++ expected_exitcode=1, ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "scenario", ++ [ ++ "blocking-str-method", ++ "overridden-getattribute", ++ "str-subclass-argument", ++ "huge-string", ++ ], ++) ++def test_fatal_helper_hard_exit_controls( ++ tmp_path: Path, ++ scenario: str, ++) -> None: ++ _run_bounded_child( ++ _FATAL_SOURCE, ++ tmp_path, ++ scenario, ++ expected_exitcode=1, ++ ) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only") ++def test_wait_for_ready_rejects_non_string_status_without_comparison( ++ tmp_path: Path, ++) -> None: ++ source = r""" ++import subprocess ++import sys ++import threading ++import time ++from pathlib import Path ++ ++import vllm.v1.ple_offload.worker as worker_module ++ ++class Status: ++ def __eq__(self, other): ++ threading.Event().wait() ++ ++ def __ne__(self, other): ++ threading.Event().wait() ++ ++class Reader: ++ def close(self): ++ return None ++ ++sleeper = subprocess.Popen( ++ [sys.executable, "-c", "import time; time.sleep(120)"] ++) ++handle = worker_module.PleOffloadWorkerHandle( ++ proc=sleeper, ++ death_writer=None, ++ ready_pipe_reader=Reader(), ++ launch_deadline=time.monotonic() + 120, ++) ++worker_module._recv_bounded_connection_message = ( ++ lambda reader, deadline: { ++ "status": Status(), ++ "layer_names": ["ple"], ++ } ++) ++Path(sys.argv[1]).write_text("armed", encoding="utf-8") ++rejected = False ++closed = False ++reader_cleared = False ++try: ++ try: ++ worker_module.PleOffloadWorker.wait_for_ready(handle) ++ except RuntimeError: ++ rejected = True ++ # Postconditions are read before the emergency cleanup below can act. ++ # Stay inside the parent's 5 second bounded wait so a miss reports here. ++ close_deadline = time.monotonic() + 3 ++ while sleeper.poll() is None and time.monotonic() < close_deadline: ++ time.sleep(0.05) ++ closed = sleeper.poll() is not None ++ reader_cleared = handle.ready_pipe_reader is None ++finally: ++ if sleeper.poll() is None: ++ sleeper.kill() ++ sleeper.wait(timeout=120) ++ ++if not rejected: ++ raise AssertionError("invalid response status was not rejected") ++if not closed: ++ raise AssertionError("invalid response did not close the child") ++if not reader_cleared: ++ raise AssertionError("invalid response did not clear the ready reader") ++""" ++ _run_bounded_child(source, tmp_path) ++ ++ ++@pytest.mark.parametrize("identity", ["duplicate", "unbound"]) ++def test_profile_breakable_initializer_follows_identity_validation( ++ tmp_path: Path, ++ identity: str, ++) -> None: ++ source = r""" ++import os ++import sys ++import threading ++from pathlib import Path ++from types import SimpleNamespace ++ ++from vllm.v1.worker.gpu.cudagraph_utils import ModelCudaGraphManager ++ ++identity = sys.argv[2] ++if identity == "duplicate": ++ identities = ("shared", "shared") ++else: ++ identities = ("",) ++ ++modules = [ ++ ( ++ "module_" + str(index), ++ SimpleNamespace( ++ _is_cpu_offloaded=True, ++ _host_output_ptr=1, ++ _ple_offload_name=value, ++ ), ++ ) ++ for index, value in enumerate(identities) ++] ++model = SimpleNamespace(named_modules=lambda: iter(modules)) ++manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) ++spy = Path(sys.argv[1] + ".initializer") ++ ++def blocking_initializer(model): ++ spy.write_text("called", encoding="utf-8") ++ threading.Event().wait() ++ ++manager.init_breakable_cg_runner = blocking_initializer ++Path(sys.argv[1]).write_text("armed", encoding="utf-8") ++ ++prepare = getattr(manager, "_prepare_ple_capture", None) ++if prepare is None: ++ # This is the pre-fix profiling order: initialize before validation. ++ manager.init_breakable_cg_runner(model) ++else: ++ try: ++ prepare(model, initialize_breakable=True) ++ except RuntimeError as error: ++ if "distinct, bound" not in str(error): ++ raise ++ else: ++ raise AssertionError("invalid identities were accepted") ++ ++if spy.exists(): ++ raise AssertionError("initializer ran before identity rejection") ++""" ++ _run_bounded_child(source, tmp_path, identity) ++ ++ ++@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only") ++@pytest.mark.parametrize("status", ["READY", "FAILURE"]) ++def test_oversized_startup_response_becomes_bounded_failure( ++ tmp_path: Path, ++ status: str, ++) -> None: ++ source = r""" ++import socket ++import sys ++from multiprocessing.connection import Connection ++from pathlib import Path ++ ++import vllm.v1.ple_offload.worker as worker_module ++ ++parent_socket, child_socket = socket.socketpair() ++raw_writer = Connection(child_socket.detach()) ++writer_type = getattr( ++ worker_module, ++ "_BoundedStartupResponseWriter", ++ None, ++) ++writer = raw_writer if writer_type is None else writer_type(raw_writer) ++ ++if sys.argv[2] == "READY": ++ message = { ++ "status": "READY", ++ "layer_names": ["x" * (2 << 20)], ++ } ++else: ++ message = { ++ "status": "FAILURE", ++ "error": "x" * (2 << 20), ++ } ++ ++Path(sys.argv[1]).write_text("armed", encoding="utf-8") ++writer.send(message) ++""" ++ _run_bounded_child(source, tmp_path, status) ++ ++ ++def test_startup_response_payload_at_cap_is_accepted(tmp_path: Path) -> None: ++ source = r""" ++import pickle ++import sys ++import threading ++from pathlib import Path ++ ++import vllm.v1.ple_offload.worker as worker_module ++ ++writer_type = getattr( ++ worker_module, ++ "_BoundedStartupResponseWriter", ++ None, ++) ++Path(sys.argv[1]).write_text("armed", encoding="utf-8") ++if writer_type is None: ++ threading.Event().wait() ++ ++cap = worker_module._PLE_READY_MAX_PAYLOAD_BYTES ++name = "x" * cap ++for unused in range(8): ++ message = {"status": "READY", "layer_names": [name]} ++ payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL) ++ if len(payload) == cap: ++ break ++ next_size = len(name) + cap - len(payload) ++ if next_size < 0: ++ raise AssertionError("cannot construct exact-cap payload") ++ name = "x" * next_size ++else: ++ raise AssertionError("exact-cap payload did not converge") ++ ++class Writer: ++ def __init__(self): ++ self.payload = None ++ ++ def send_bytes(self, payload): ++ self.payload = payload ++ ++ def close(self): ++ return None ++ ++raw_writer = Writer() ++writer_type(raw_writer).send(message) ++if raw_writer.payload != payload or len(raw_writer.payload) != cap: ++ raise AssertionError("exact-cap startup response was not preserved") ++""" ++ _run_bounded_child(source, tmp_path) +diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py +new file mode 100644 +index 0000000..1b67578 +--- /dev/null ++++ b/vllm/model_executor/layers/ple_offload_layer.py +@@ -0,0 +1,433 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Base class for PLE layers that support cross-process CPU offload. ++ ++Synchronization between the offload process and the GPU main stream is ++handled by :class:`CpuGpuSemaphore`, which uses CUDA's ++``cuStreamWaitValue32`` / ``cuStreamWriteValue32`` stream memory operations. ++ ++Decode-step timeline, host-mapped PULL design ++-------------------------------------------------------- ++Offload process (CPU, no CUDA) GPU main stream (forward) ++------------------------------ ------------------------------ ++forward_impl() on CPU ... other GPU ops ... ++wait host flag == 0 WaitValue32(host flag == 1) <- in graph ++rows -> shared pinned host slot cuMemcpyHtoDAsync(slot -> GPU) <- in graph ++publish fence; host flag = 1 return _gpu_output_buffer[:n] ++ ... model consumes output ... ++ WriteValue32(host flag = 0) <- release_outputs ++ ++``release_outputs`` performs the sole reset. Safety requires serialized local ++forwards on one ordered model stream: a new transaction cannot begin until the ++prior release has been enqueued. If any prepared forward is abandoned or its ++correspondence becomes uncertain, the worker must terminate; resetting, ++retrying, or continuing could reuse stale rows or overwrite live input. ++""" ++ ++import contextlib ++import functools ++import threading ++from abc import ABC, abstractmethod ++from collections.abc import Callable, Iterator ++from typing import Any, cast ++ ++import torch ++from cuda.bindings import driver as cuda_driver ++from cuda.bindings.driver import CUstreamWaitValue_flags ++from torch import nn ++ ++import vllm.envs as envs ++from vllm.utils.torch_utils import direct_register_custom_op ++ ++# Module-level flag set to True inside the offload subprocess. ++# Because the offload process and GPU worker processes are separate OS ++# processes (spawned via multiprocessing), each has its own memory space. ++# A plain module-level bool is sufficient -- no thread-local storage needed. ++_offload_worker_flag = False ++ ++# Native FULL capture runs on one Python thread. The custom-op implementation ++# appends the exact arguments from every wait-copy invocation while that thread ++# is inside the graph manager's capture scope. ++_ple_copy_capture = threading.local() ++ ++ ++@contextlib.contextmanager ++def observe_ple_copy_geometry() -> Iterator[list[tuple[str, int, int]]]: ++ """Collect ``(layer, rows, bytes)`` from actual wait-copy executions.""" ++ if getattr(_ple_copy_capture, "records", None) is not None: ++ raise RuntimeError("nested PLE copy-geometry observation is unsupported") ++ records: list[tuple[str, int, int]] = [] ++ _ple_copy_capture.records = records ++ try: ++ yield records ++ finally: ++ del _ple_copy_capture.records ++ ++ ++def _record_ple_copy_geometry( ++ layer_name: str, ++ num_tokens: int, ++ row_bytes: int, ++) -> None: ++ records = getattr(_ple_copy_capture, "records", None) ++ if records is not None: ++ rows = int(num_tokens) ++ records.append((str(layer_name), rows, rows * int(row_bytes))) ++ ++ ++def is_offload_process() -> bool: ++ """Return True inside the dedicated PLE CPU-offload subprocess.""" ++ return _offload_worker_flag ++ ++ ++def mark_as_offload_worker() -> None: ++ """Mark the current process as the dedicated PLE offload worker.""" ++ global _offload_worker_flag ++ _offload_worker_flag = True ++ ++ ++def _cuda_check(result: Any, operation: str) -> Any: ++ """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls.""" ++ error = result[0] if isinstance(result, tuple) else result ++ if error.value != 0: ++ raise RuntimeError(f"{operation} failed: {error}") ++ return result ++ ++ ++class CpuGpuSemaphore: ++ """Cross-process semaphore backed by a one-element int32 CUDA tensor. ++ ++ The flag is stored in a regular GPU tensor that is shared through ++ PyTorch's CUDA IPC mechanism. Both processes obtain device pointers that ++ map to the same physical GPU memory, so a stream-memory wait in the GPU ++ worker observes the write issued by the offload process. ++ """ ++ ++ RESET_VALUE = 0 ++ DONE_VALUE = 1 ++ ++ def __init__(self, device: torch.device) -> None: ++ self._flag_tensor = torch.zeros(1, dtype=torch.int32, device=device) ++ self._flag_ptr: int | None = None ++ ++ @classmethod ++ def from_ipc_tensor(cls, flag_tensor: torch.Tensor) -> "CpuGpuSemaphore": ++ """Construct a semaphore from a CUDA tensor received through IPC.""" ++ semaphore = cls.__new__(cls) ++ semaphore._flag_tensor = flag_tensor ++ semaphore._flag_ptr = None ++ return semaphore ++ ++ @classmethod ++ def from_host_mapped( ++ cls, host_flag_tensor: torch.Tensor, device_ptr: int ++ ) -> "CpuGpuSemaphore": ++ """Semaphore on a page-locked, device-mapped shared host int32. ++ ++ ``host_flag_tensor`` keeps the shared storage alive; ``device_ptr`` is ++ its device address (``cuMemHostGetDevicePointer``). Stream memory ops ++ on it are executed by the GPU against host memory, so a CPU process ++ with no CUDA context can signal it with a plain store. ++ """ ++ semaphore = cls.__new__(cls) ++ semaphore._flag_tensor = host_flag_tensor ++ semaphore._flag_ptr = int(device_ptr) ++ return semaphore ++ ++ @property ++ def flag_tensor(self) -> torch.Tensor: ++ """Return the tensor that backs the semaphore flag.""" ++ return self._flag_tensor ++ ++ @property ++ def flag_ptr(self) -> int: ++ """Device address used by the stream memory operations.""" ++ if self._flag_ptr is not None: ++ return self._flag_ptr ++ return self._flag_tensor.data_ptr() ++ ++ def reset(self, stream: torch.cuda.Stream | None = None) -> None: ++ """Enqueue ``WriteValue32(flag=0)`` on ``stream``.""" ++ if stream is None: ++ stream = torch.cuda.current_stream() ++ _cuda_check( ++ cuda_driver.cuStreamWriteValue32( ++ cuda_driver.CUstream(stream.cuda_stream), ++ cuda_driver.CUdeviceptr(self.flag_ptr), ++ self.RESET_VALUE, ++ 0, ++ ), ++ "CpuGpuSemaphore.reset", ++ ) ++ ++ def signal(self, stream: torch.cuda.Stream | None = None) -> None: ++ """Enqueue ``WriteValue32(flag=1)`` on ``stream``.""" ++ if stream is None: ++ stream = torch.cuda.current_stream() ++ _cuda_check( ++ cuda_driver.cuStreamWriteValue32( ++ cuda_driver.CUstream(stream.cuda_stream), ++ cuda_driver.CUdeviceptr(self.flag_ptr), ++ self.DONE_VALUE, ++ 0, ++ ), ++ "CpuGpuSemaphore.signal", ++ ) ++ ++ def wait_reset(self, stream: torch.cuda.Stream | None = None) -> None: ++ """Enqueue ``WaitValue32(flag==0)`` on ``stream``.""" ++ if stream is None: ++ stream = torch.cuda.current_stream() ++ _cuda_check( ++ cuda_driver.cuStreamWaitValue32( ++ cuda_driver.CUstream(stream.cuda_stream), ++ cuda_driver.CUdeviceptr(self.flag_ptr), ++ self.RESET_VALUE, ++ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, ++ ), ++ "CpuGpuSemaphore.wait_reset", ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# Custom op: vllm::ple_offload_wait ++# --------------------------------------------------------------------------- ++# Wraps cuStreamWaitValue32 as a torch custom op so that ++# torch.compile(fullgraph=True) treats the CUDA Driver call as an opaque node. ++# The hidden_states argument creates a dynamo data-dependency edge that prevents ++# the wait from being reordered before preceding GPU work. ++# --------------------------------------------------------------------------- ++ ++ ++def _ple_offload_wait_impl( ++ sem_flag_tensor: torch.Tensor, ++ gpu_output_buffer: torch.Tensor, ++ hidden_states: torch.Tensor, ++) -> None: ++ """Wait for the CPU result without releasing its output buffer.""" ++ stream = torch.cuda.current_stream() ++ cuda_stream = cuda_driver.CUstream(stream.cuda_stream) ++ dev_ptr = cuda_driver.CUdeviceptr(sem_flag_tensor.data_ptr()) ++ _cuda_check( ++ cuda_driver.cuStreamWaitValue32( ++ cuda_stream, ++ dev_ptr, ++ CpuGpuSemaphore.DONE_VALUE, ++ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, ++ ), ++ "cuStreamWaitValue32(done)", ++ ) ++ ++ ++def _ple_offload_wait_fake( ++ sem_flag_tensor: torch.Tensor, ++ gpu_output_buffer: torch.Tensor, ++ hidden_states: torch.Tensor, ++) -> None: ++ """Represent the side-effect-only wait during dynamo tracing.""" ++ pass ++ ++ ++direct_register_custom_op( ++ op_name="ple_offload_wait", ++ op_func=_ple_offload_wait_impl, ++ mutates_args=["gpu_output_buffer"], ++ fake_impl=_ple_offload_wait_fake, ++) ++ ++ ++# --------------------------------------------------------------------------- ++# Custom op: vllm::ple_offload_wait_copy (host-mapped pull design) ++# --------------------------------------------------------------------------- ++# The offload process writes the result into a page-locked shared host buffer ++# and stores flag=1 in a page-locked shared host int32; it holds no CUDA ++# context. On the CURRENT stream (captured into the graph like the wait above): ++# cuStreamWaitValue32(host flag == 1) -> GPU polls host memory ++# cuMemcpyAsync(gpu buffer <- host buffer, num_tokens * row_bytes) ++# The copy engine reads host memory coherently after the flag was observed, ++# and the consumers read device memory as before. The flag is NOT reset here: ++# ``release_outputs`` resets it after the forward (unchanged semantics), so ++# capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``. ++# Pointers are passed as ints so no CPU tensor enters the compiled graph. ++# --------------------------------------------------------------------------- ++ ++ ++def _ple_offload_wait_copy_impl( ++ gpu_output_buffer: torch.Tensor, ++ hidden_states: torch.Tensor, ++ flag_ptr: int, ++ host_ptr: int, ++ row_bytes: int, ++ num_tokens: int, ++ layer_name: str, ++) -> None: ++ """Wait for the host-side result, then DMA it into the GPU buffer.""" ++ if not 0 <= int(num_tokens) <= gpu_output_buffer.shape[0] or int(row_bytes) != ( ++ gpu_output_buffer.shape[1] * gpu_output_buffer.element_size() ++ ): ++ raise RuntimeError( ++ f"PLE wait_copy out of bounds: {num_tokens} rows x {row_bytes} B vs " ++ f"GPU buffer {tuple(gpu_output_buffer.shape)} {gpu_output_buffer.dtype}" ++ ) ++ nbytes = int(num_tokens) * int(row_bytes) ++ _record_ple_copy_geometry(layer_name, num_tokens, row_bytes) ++ stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) ++ _cuda_check( ++ cuda_driver.cuStreamWaitValue32( ++ stream, ++ cuda_driver.CUdeviceptr(flag_ptr), ++ CpuGpuSemaphore.DONE_VALUE, ++ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, ++ ), ++ "cuStreamWaitValue32(host flag done)", ++ ) ++ if nbytes > 0: ++ # host_ptr is this worker's host address of page-locked (registered) ++ # memory, so the copy is a true async DMA and is capturable. ++ _cuda_check( ++ cuda_driver.cuMemcpyHtoDAsync( ++ cuda_driver.CUdeviceptr(gpu_output_buffer.data_ptr()), ++ host_ptr, ++ nbytes, ++ stream, ++ ), ++ "cuMemcpyHtoDAsync(PLE host -> GPU)", ++ ) ++ ++ ++def _ple_offload_wait_copy_fake( ++ gpu_output_buffer: torch.Tensor, ++ hidden_states: torch.Tensor, ++ flag_ptr: int, ++ host_ptr: int, ++ row_bytes: int, ++ num_tokens: int, ++ layer_name: str, ++) -> None: ++ """Represent the side-effect-only wait + copy during dynamo tracing.""" ++ pass ++ ++ ++direct_register_custom_op( ++ op_name="ple_offload_wait_copy", ++ op_func=_ple_offload_wait_copy_impl, ++ mutates_args=["gpu_output_buffer"], ++ fake_impl=_ple_offload_wait_copy_fake, ++) ++ ++ ++class PleOffloadLayer(nn.Module, ABC): ++ """Base class for embedding-like PLE layers that can run on CPU. ++ ++ Subclasses implement :meth:`forward_impl`, which contains the regular ++ on-device computation. In a GPU worker with PLE offload enabled, subclass ++ constructors are skipped so their large weights are never allocated. The ++ CPU process constructs a complete copy and owns those weights instead. ++ """ ++ ++ _is_cpu_offloaded = False ++ _gpu_output_buffer: torch.Tensor ++ _sem: CpuGpuSemaphore ++ # Host-mapped pull design: host address of the shared output buffer and ++ # the bytes per output row. ``None`` keeps the IPC push behaviour. ++ _host_output_ptr: int | None = None ++ _host_row_bytes: int = 0 ++ _ple_offload_name: str = "" ++ ++ def __init_subclass__(cls, **kwargs: object) -> None: ++ """Skip subclass initialization in cross-process GPU-worker mode.""" ++ super().__init_subclass__(**kwargs) ++ original_init_obj = cls.__dict__.get("__init__") ++ if original_init_obj is None or not callable(original_init_obj): ++ return ++ original_init = cast(Callable[..., None], original_init_obj) ++ ++ @functools.wraps(original_init) ++ def guarded_init( ++ self: "PleOffloadLayer", *args: object, **kwargs: object ++ ) -> None: ++ if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): ++ nn.Module.__init__(self) ++ return ++ original_init(self, *args, **kwargs) ++ ++ cls.__init__ = guarded_init # type: ignore[method-assign, assignment] ++ ++ @classmethod ++ def get_target_device(cls) -> torch.device: ++ """Return CPU for the offload process and the active GPU otherwise.""" ++ if envs.VLLM_PLE_CPU_OFFLOAD: ++ return torch.device("cpu") ++ return torch.device("cuda", torch.accelerator.current_device_index()) ++ ++ @abstractmethod ++ def forward_impl( ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ *args: object, ++ **kwargs: object, ++ ) -> torch.Tensor: ++ """Execute the actual embedding computation on the owning device.""" ++ raise NotImplementedError ++ ++ def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: ++ """Return the dtype used by cross-process output buffers.""" ++ return default_dtype ++ ++ def get_offload_output_dim(self, default_dim: int) -> int: ++ """Return the last dimension used by cross-process output buffers.""" ++ return default_dim ++ ++ def initialize_dummy_offload_metadata(self, device: torch.device) -> None: ++ """Initialize metadata needed by dummy offload execution.""" ++ ++ def setup_cross_process_offload( ++ self, ++ gpu_output_buffer: torch.Tensor, ++ semaphore: CpuGpuSemaphore, ++ host_output_ptr: int | None = None, ++ host_row_bytes: int = 0, ++ ) -> None: ++ """Configure the GPU-worker placeholder with its offload resources.""" ++ self._is_cpu_offloaded = True ++ self._gpu_output_buffer = gpu_output_buffer ++ self._sem = semaphore ++ self._host_output_ptr = host_output_ptr ++ self._host_row_bytes = int(host_row_bytes) ++ ++ def forward( ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ *args: object, ++ **kwargs: object, ++ ) -> torch.Tensor: ++ """Wait for an offloaded result or delegate to ``forward_impl``.""" ++ if self._is_cpu_offloaded: ++ if self._host_output_ptr is not None: ++ torch.ops.vllm.ple_offload_wait_copy( ++ self._gpu_output_buffer, ++ hidden_states, ++ self._sem.flag_ptr, ++ self._host_output_ptr, ++ self._host_row_bytes, ++ input_ids.shape[0], ++ self._ple_offload_name, ++ ) ++ else: ++ torch.ops.vllm.ple_offload_wait( ++ self._sem.flag_tensor, ++ self._gpu_output_buffer, ++ hidden_states, ++ ) ++ return self._gpu_output_buffer[: input_ids.shape[0]] ++ return self.forward_impl(hidden_states, input_ids, *args, **kwargs) ++ ++ def release_offloaded_output( ++ self, ++ stream: torch.cuda.Stream | None = None, ++ ) -> None: ++ """Mark the cross-process output buffer reusable on ``stream``.""" ++ if self._is_cpu_offloaded: ++ self._sem.reset(stream) +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index 1ccd23a..e27a18b 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -352,9 +352,10 @@ class MultiprocExecutor(Executor): + proc = next(h.proc for h in workers if h.proc.sentinel == died[0]) + logger.error( + "Worker proc %s died unexpectedly (exit code: %s), " +- "shutting down executor.", ++ "shutting down executor. executor_shutdown_monotonic=%.9f", + proc.name, + proc.exitcode, ++ time.monotonic(), + ) + _self.shutdown() + callback = _self.failure_callback +@@ -1034,6 +1035,28 @@ class WorkerProc: + SUCCESS = auto() + FAILURE = auto() + ++ def _fail_ple_rpc_if_pending( ++ self, where: str, error: BaseException ++ ) -> None: ++ """Fail closed if RPC recovery would abandon queued PLE GPU work.""" ++ owners = (self.worker, getattr(self.worker, "worker", None)) ++ for owner in owners: ++ if owner is None: ++ continue ++ runner = getattr(owner, "model_runner", None) ++ connector = getattr(runner, "_ple_offload_connector", None) ++ if connector is None: ++ continue ++ try: ++ pending = connector.has_pending_completions ++ except BaseException as state_error: ++ connector._fail_closed( ++ f"{where}:pending-state", state_error ++ ) ++ if pending: ++ connector._fail_closed(where, error) ++ return ++ + def enqueue_output(self, output: Any): + """Prepares output from the worker and enqueues it to the + worker_response_mq. If the output is an Exception, it is +@@ -1044,6 +1067,9 @@ class WorkerProc: + output = output.get_output() + except Exception as e: + logger.exception("Error getting async model runner output") ++ self._fail_ple_rpc_if_pending( ++ "async-output-copy", e ++ ) + output = e + + if isinstance(output, Exception): +@@ -1104,6 +1130,9 @@ class WorkerProc: + if output_rank is None or self.rank == output_rank: + self.handle_output(output) + except Exception as e: ++ self._fail_ple_rpc_if_pending( ++ f"worker-rpc:{method}", e ++ ) + # Notes have been introduced in python 3.11 + if hasattr(e, "add_note"): + e.add_note(traceback.format_exc()) +diff --git a/vllm/v1/ple_offload/__init__.py b/vllm/v1/ple_offload/__init__.py +new file mode 100644 +index 0000000..208f01a +--- /dev/null ++++ b/vllm/v1/ple_offload/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py +new file mode 100644 +index 0000000..52bc462 +--- /dev/null ++++ b/vllm/v1/ple_offload/connector.py +@@ -0,0 +1,855 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Exchange PLE data between a GPU worker and the CPU-offload process.""" ++ ++import os ++import queue ++import threading ++from dataclasses import dataclass ++from multiprocessing.reduction import ForkingPickler ++from typing import Any ++ ++import msgspec ++import torch ++from vllm.v1.ple_offload import rowdump as _rowdump ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++from vllm.v1.ple_offload.progress import PleProgressMonitor ++import torch.nn as nn ++import zmq ++from cuda.bindings import driver as cuda_driver ++ ++from vllm.config import VllmConfig ++from vllm.distributed.parallel_state import get_dp_group, get_pp_group, get_tp_group ++from vllm.logger import init_logger ++from vllm.model_executor.layers.ple_offload_layer import ( ++ CpuGpuSemaphore, ++ PleOffloadLayer, ++) ++from vllm.v1.ple_offload.protocol import ( ++ PleOffloadRegistration, ++ PleOffloadRequest, ++) ++from vllm.v1.utils import record_function_or_nullcontext ++from vllm_ple_test_fault import _FAULT ++ ++logger = init_logger(__name__) ++ ++# CUDA registration alone does not own Python storage. Captured graphs retain ++# raw addresses, so every successfully registered host allocation remains ++# strongly referenced until process exit. ++_PROCESS_LIFETIME_REGISTERED_STORAGE: list[torch.Tensor] = [] ++ ++ ++def _unlink_filename_shared_storage( ++ tensor: torch.Tensor, ++ reserved_filenames: set[bytes], ++) -> None: ++ """Reserve and unlink one already filename-shared CPU storage.""" ++ try: ++ storage = tensor.untyped_storage() ++ if storage.device.type != "cpu" or not storage.is_shared(): ++ return ++ try: ++ storage._get_shared_fd() ++ except RuntimeError: ++ # Filename-shared storages do not expose a shared fd. Only now is ++ # it safe to query the filename: doing so on ordinary or fd-shared ++ # storage would convert it to a new filename-backed allocation. ++ pass ++ else: ++ return ++ _, filename, _ = storage._share_filename_cpu_() ++ if not isinstance(filename, bytes): ++ filename = os.fsencode(filename) ++ path = os.path.join("/dev/shm", os.fsdecode(filename).lstrip("/")) ++ except Exception: ++ logger.warning( ++ "Failed to inspect a PLE shared-memory storage during teardown", ++ exc_info=True, ++ ) ++ return ++ ++ # RefcountedMapAllocator unlinks the name when its shared counter reaches ++ # zero. Reserve one count before manual unlink so later storage destruction ++ # cannot try to unlink the already-removed name. The filename identifies ++ # the shared counter and also deduplicates tensor aliases and repeated close. ++ if filename in reserved_filenames: ++ return ++ try: ++ storage._shared_incref() ++ except Exception: ++ logger.warning( ++ "Failed to reserve the PLE shared-memory counter for %s", ++ path, ++ exc_info=True, ++ ) ++ return ++ reserved_filenames.add(filename) ++ ++ try: ++ os.unlink(path) ++ except FileNotFoundError: ++ # Another teardown participant already removed the name. Retain the ++ # reservation because this allocator must still avoid a later unlink. ++ pass ++ except OSError as error: ++ # The name remains available, so restore the allocator's original ++ # counter and permit a later release attempt to retry. ++ try: ++ storage._shared_decref() ++ except Exception: ++ # The reservation state is uncertain after a failed private call. ++ # Keep it marked to avoid adding another reservation on retry. ++ logger.warning( ++ "Failed to roll back the PLE shared-memory counter for %s", ++ path, ++ exc_info=True, ++ ) ++ else: ++ reserved_filenames.discard(filename) ++ logger.warning("Failed to unlink PLE shared memory %s: %s", path, error) ++ ++ ++@dataclass(frozen=True) ++class _PendingPleOffloadRequest: ++ """Bind request metadata to its MRV2 D2H completion event.""" ++ ++ request: PleOffloadRequest ++ d2h_done_event: torch.cuda.Event ++ eligible: bool = False ++ transaction_id: str | None = None ++ correlation_id: str | None = None ++ ++ ++def _cuda_check(result: Any, operation: str) -> Any: ++ """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls.""" ++ error = result[0] if isinstance(result, tuple) else result ++ if error.value != 0: ++ raise RuntimeError(f"{operation} failed: {error}") ++ return result ++ ++ ++class PleOffloadConnector: ++ """Connect a model runner V2 GPU worker to the PLE CPU worker.""" ++ ++ def __init__( ++ self, ++ vllm_config: VllmConfig, ++ model: nn.Module, ++ device: torch.device, ++ ipc_addr: str, ++ *, ++ input_ids_source: torch.Tensor, ++ query_start_loc_source: torch.Tensor, ++ ngram_context_source: torch.Tensor | None, ++ ) -> None: ++ self.device = device ++ self.rank = int(vllm_config.parallel_config.rank) ++ self.dp_rank = get_dp_group().rank_in_group ++ self.tp_rank = get_tp_group().rank_in_group ++ self.pp_rank = get_pp_group().rank_in_group ++ # Host-mapped pull design: per layer a shared host output buffer and a ++ # shared host int32 flag (element 0), filled/stored by the offload ++ # process with plain CPU writes. Page-locked + device-mapped here only ++ # after registration, because serialising for registration may move ++ # the shared storage to a new address. ++ self._host_output_bufs: dict[str, torch.Tensor] = {} ++ self._host_flags: dict[str, torch.Tensor] = {} ++ self._reserved_host_storage_filenames: set[bytes] = set() ++ self._gpu_output_bufs: dict[str, torch.Tensor] = {} ++ self._layers = self._setup_layers(vllm_config, model) ++ self._progress = PleProgressMonitor.from_environment( ++ rank=self.rank, ++ layer_names=tuple(self._layers), ++ ) ++ ++ # TP0 registers shared buffers with CUDA for asynchronous D2H copies. ++ scheduler_config = vllm_config.scheduler_config ++ self._input_ids_buf = torch.empty( ++ scheduler_config.max_num_batched_tokens, ++ dtype=torch.int32, ++ device="cpu", ++ ).share_memory_() ++ self._query_start_loc_buf = torch.empty( ++ scheduler_config.max_num_seqs + 1, ++ dtype=torch.int32, ++ device="cpu", ++ ).share_memory_() ++ self._ngram_context_buf = None ++ config = vllm_config.model_config.hf_text_config ++ ngram_context_len = int(config.ngram_size) - 1 ++ if ngram_context_len > 0: ++ self._ngram_context_buf = torch.empty( ++ scheduler_config.max_num_seqs, ++ ngram_context_len, ++ dtype=torch.int32, ++ device="cpu", ++ ).share_memory_() ++ ++ # Runner input allocations are address-stable, so bind them once and ++ # pass only batch sizes through the per-forward request queue. ++ self._input_ids_source = input_ids_source ++ self._query_start_loc_source = query_start_loc_source ++ self._ngram_context_source = ngram_context_source ++ self._validate_input_sources() ++ ++ self._pinned_input_buffers: list[torch.Tensor] = [] ++ # Device-mapped host outputs/flags referenced by raw address from ++ # captured graphs: never unregistered while the process lives ++ # (unregistering under a pending graph wait unmaps it). ++ self._mapped_host_buffers: list[torch.Tensor] = [] ++ # True while dummy/capture work (seeds, dummy forwards, capture ++ # cleanup resets) has run since the last real forward. Releases in ++ # that state complete synchronously, and the next real forward first ++ # passes a TP barrier, so no rank can publish a real request while ++ # another rank still has a dummy/cleanup reset pending. ++ # In this tree real-path kernel warm-ups run BEFORE ++ # capture and capture adds seeds/resets after them, so this must hold ++ # at every dummy->real transition, not only at the first real one. ++ self._after_dummy = True ++ self._request_queue: queue.Queue[_PendingPleOffloadRequest | None] = ( ++ queue.Queue(maxsize=vllm_config.max_concurrent_batches) ++ ) ++ self._dropped_events: list[torch.cuda.Event] = [] ++ self._request_thread: threading.Thread | None = None ++ self._request_thread_ready = threading.Event() ++ self._zmq_ctx: zmq.Context | None = None ++ self._registration_socket: zmq.Socket | None = None ++ self._d2h_event_pool: queue.Queue[torch.cuda.Event] | None = None ++ ++ try: ++ self._zmq_ctx = zmq.Context() ++ self._registration_socket = self._zmq_ctx.socket(zmq.PUSH) ++ self._registration_socket.connect(ipc_addr) ++ payload = self._serialize_registration(vllm_config) ++ ++ # Serialising moved the shared storages to their final addresses. ++ # Page-lock/map exactly those, THEN publish, so a received ++ # registration always means "storage usable". ++ with torch.accelerator.device_index(self.device.index): ++ self._map_host_outputs(vllm_config) ++ if self.tp_rank == 0: ++ self._pin_input_buffers() ++ self._send_registration(payload, ipc_addr) ++ ++ if self.tp_rank == 0: ++ with torch.accelerator.device_index(self.device.index): ++ self._d2h_event_pool = queue.Queue( ++ maxsize=vllm_config.max_concurrent_batches ++ ) ++ for _ in range(vllm_config.max_concurrent_batches): ++ self._d2h_event_pool.put_nowait(torch.cuda.Event()) ++ self._start_request_thread(ipc_addr) ++ except Exception: ++ self.close() ++ raise ++ ++ def _setup_layers( ++ self, ++ vllm_config: VllmConfig, ++ model: nn.Module, ++ ) -> dict[str, PleOffloadLayer]: ++ """Attach output buffers and semaphores to GPU PLE placeholders.""" ++ layers = { ++ name: module ++ for name, module in model.named_modules() ++ if isinstance(module, PleOffloadLayer) ++ } ++ if not layers: ++ raise RuntimeError( ++ "VLLM_PLE_CPU_OFFLOAD is enabled, but the model has no PleOffloadLayer" ++ ) ++ ++ config = vllm_config.model_config.hf_text_config ++ max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens ++ for name, layer in layers.items(): ++ if vllm_config.load_config.load_format == "dummy": ++ layer.initialize_dummy_offload_metadata(self.device) ++ out_dim = layer.get_offload_output_dim(int(config.ple_embed_dim)) ++ out_dtype = layer.get_offload_output_dtype(vllm_config.model_config.dtype) ++ # The captured forward DMAs the host result into this device ++ # buffer after the host flag is set; consumers read it as before. ++ self._gpu_output_bufs[name] = torch.empty( ++ max_num_tokens, out_dim, dtype=out_dtype, device=self.device ++ ) ++ # Zero-initialised so dummy/profile forwards copy finite values. ++ self._host_output_bufs[name] = torch.zeros( ++ max_num_tokens, out_dim, dtype=out_dtype, device="cpu" ++ ).share_memory_() ++ # 16 x int32 = one 64-byte line; element 0 is the flag. ++ self._host_flags[name] = torch.zeros( ++ 16, dtype=torch.int32, device="cpu" ++ ).share_memory_() ++ return layers ++ ++ def _map_host_outputs(self, vllm_config: VllmConfig) -> None: ++ """Page-lock + device-map host outputs/flags and arm the placeholders.""" ++ for name, layer in self._layers.items(): ++ layer._ple_offload_name = name ++ host_out = self._host_output_bufs[name] ++ host_flag = self._host_flags[name] ++ for buffer in (host_out, host_flag): ++ self._register_host_buffer(buffer, device_map=True) ++ flag_dev_ptr = _cuda_check( ++ cuda_driver.cuMemHostGetDevicePointer(host_flag.data_ptr(), 0), ++ "cuMemHostGetDevicePointer(PLE host flag)", ++ )[1] ++ out_dev_ptr = _cuda_check( ++ cuda_driver.cuMemHostGetDevicePointer(host_out.data_ptr(), 0), ++ "cuMemHostGetDevicePointer(PLE host output)", ++ )[1] ++ row_bytes = host_out.shape[1] * host_out.element_size() ++ gpu_buf = self._gpu_output_bufs[name] ++ if ( ++ gpu_buf.shape[0] != host_out.shape[0] ++ or gpu_buf.shape[1] * gpu_buf.element_size() != row_bytes ++ ): ++ raise RuntimeError(f"PLE host/GPU buffer geometry mismatch for {name}") ++ # Stream memops use the flag's DEVICE pointer; the H2D copy uses ++ # this worker's HOST pointer (no raw addresses cross ++ # processes; the offload writes through its own mapping). ++ layer.setup_cross_process_offload( ++ gpu_buf, ++ CpuGpuSemaphore.from_host_mapped(host_flag, int(flag_dev_ptr)), ++ host_output_ptr=host_out.data_ptr(), ++ host_row_bytes=row_bytes, ++ ) ++ logger.info( ++ "PleOffload: host-mapped pull for %s (rows %d x %d B, flag dev " ++ "0x%x, out dev 0x%x, host==dev %s)", ++ name, ++ host_out.shape[0], ++ row_bytes, ++ int(flag_dev_ptr), ++ int(out_dev_ptr), ++ int(out_dev_ptr) == host_out.data_ptr(), ++ ) ++ ++ def _pin_input_buffers(self) -> None: ++ """Page-lock shared input allocations without replacing their storage.""" ++ buffers = [self._input_ids_buf, self._query_start_loc_buf] ++ if self._ngram_context_buf is not None: ++ buffers.append(self._ngram_context_buf) ++ for buffer in buffers: ++ self._register_host_buffer(buffer, device_map=False) ++ ++ def _register_host_buffer(self, buffer: torch.Tensor, device_map: bool) -> None: ++ """Page-lock (and optionally device-map) a shared CPU allocation.""" ++ if buffer.device.type != "cpu" or not buffer.is_shared(): ++ raise RuntimeError("PLE host buffers must be shared CPU tensors") ++ if not buffer.is_contiguous(): ++ raise RuntimeError("PLE host buffers must be contiguous") ++ flags = cuda_driver.CU_MEMHOSTREGISTER_PORTABLE ++ if device_map: ++ flags |= cuda_driver.CU_MEMHOSTREGISTER_DEVICEMAP ++ _cuda_check( ++ cuda_driver.cuMemHostRegister( ++ buffer.data_ptr(), ++ buffer.numel() * buffer.element_size(), ++ flags, ++ ), ++ "cuMemHostRegister(PLE host buffer)", ++ ) ++ _PROCESS_LIFETIME_REGISTERED_STORAGE.append(buffer) ++ if device_map: ++ self._mapped_host_buffers.append(buffer) ++ else: ++ self._pinned_input_buffers.append(buffer) ++ if not buffer.is_pinned(): ++ raise RuntimeError("CUDA did not page-lock a PLE host buffer") ++ ++ def _unpin_input_buffers(self) -> None: ++ """Release CUDA registrations after the request thread has stopped.""" ++ for buffer in reversed(self._pinned_input_buffers): ++ try: ++ _cuda_check( ++ cuda_driver.cuMemHostUnregister(buffer.data_ptr()), ++ "cuMemHostUnregister(PLE input buffer)", ++ ) ++ except RuntimeError: ++ logger.exception("Failed to unregister a PLE input buffer") ++ self._pinned_input_buffers.clear() ++ ++ def _serialize_registration(self, vllm_config: VllmConfig) -> Any: ++ """Build and serialise the registration (may move shared storages).""" ++ # Each GPU worker owns distinct host output buffers and flags, while ++ # TP0's shared inputs become the request source for its DP rank. ++ registration = PleOffloadRegistration( ++ worker_id=( ++ self.dp_rank * vllm_config.parallel_config.world_size ++ + vllm_config.parallel_config.rank ++ ), ++ tp_rank=self.tp_rank, ++ dp_rank=self.dp_rank, ++ host_output_buffers=dict(self._host_output_bufs), ++ host_flag_tensors=dict(self._host_flags), ++ input_ids_buf=self._input_ids_buf, ++ query_start_loc_buf=self._query_start_loc_buf, ++ ngram_context_buf=self._ngram_context_buf, ++ ) ++ ++ # ForkingPickler transmits tensors through shared-memory and CUDA IPC. ++ import torch.multiprocessing as torch_mp ++ ++ original_strategy = torch_mp.get_sharing_strategy() ++ torch_mp.set_sharing_strategy("file_system") ++ try: ++ payload = ForkingPickler.dumps(registration) ++ finally: ++ torch_mp.set_sharing_strategy(original_strategy) ++ return payload ++ ++ def _send_registration(self, payload: Any, ipc_addr: str) -> None: ++ """Publish the serialised registration (storages already mapped).""" ++ assert self._registration_socket is not None ++ self._registration_socket.send(payload) ++ ++ logger.info( ++ "PleOffload: registered %d PleOffloadLayer(s) " ++ "(dp_rank=%d, tp_rank=%d, ipc_addr=%s): %s", ++ len(self._layers), ++ self.dp_rank, ++ self.tp_rank, ++ ipc_addr, ++ sorted(self._layers), ++ ) ++ ++ def _start_request_thread(self, ipc_addr: str) -> None: ++ """Start the thread that publishes batches after inputs are ready.""" ++ self._request_thread = threading.Thread( ++ target=self._request_loop, ++ args=(ipc_addr,), ++ name=f"ple-offload-dp{self.dp_rank}", ++ daemon=True, ++ ) ++ self._request_thread.start() ++ if not self._request_thread_ready.wait(timeout=10): ++ raise RuntimeError("Timed out starting the PLE request thread") ++ ++ def _request_loop(self, ipc_addr: str) -> None: ++ """Wait for staged inputs, then notify the CPU worker.""" ++ socket: zmq.Socket | None = None ++ try: ++ if self._zmq_ctx is None: ++ raise RuntimeError("PLE ZMQ context closed before thread startup") ++ socket = self._zmq_ctx.socket(zmq.PUSH) ++ socket.connect(ipc_addr) ++ self._request_thread_ready.set() ++ while True: ++ request = self._request_queue.get() ++ if request is None: ++ return ++ self._process_request(request, socket) ++ except Exception: ++ logger.exception("PLE request thread failed") ++ os._exit(1) ++ finally: ++ self._request_thread_ready.set() ++ if socket is not None: ++ socket.close(linger=0) ++ ++ def _process_request( ++ self, pending: _PendingPleOffloadRequest, socket: zmq.Socket ++ ) -> None: ++ """Wait for one staged batch and publish its request.""" ++ request = pending.request ++ event_pool = self._d2h_event_pool ++ event = pending.d2h_done_event ++ assert event_pool is not None, "PLE D2H event pool is not initialized" ++ if _FAULT is not None and _FAULT.fires( ++ "block_input_completion", ++ role="gpu", ++ unit="connector-request", ++ eligible=pending.eligible, ++ tp_rank=self.tp_rank, ++ pp_rank=self.pp_rank, ++ dp_rank=self.dp_rank, ++ transaction_id=pending.transaction_id, ++ correlation_id=pending.correlation_id, ++ ): ++ threading.Event().wait() ++ ++ with torch.accelerator.device_index(self.device.index): ++ event.synchronize() ++ ++ payload = msgspec.msgpack.encode(request) ++ if _FAULT is None: ++ socket.send(payload) ++ else: ++ # Armed runs only: correlation metadata rides as a second frame. ++ metadata = msgspec.msgpack.encode( ++ { ++ "serving": pending.eligible, ++ "transaction": pending.transaction_id, ++ "correlation": pending.correlation_id, ++ } ++ ) ++ socket.send_multipart([payload, metadata]) ++ event_pool.put_nowait(event) ++ ++ def _validate_input_sources(self) -> None: ++ """Validate fixed runner sources against shared input buffers.""" ++ sources = [ ++ ("input_ids", self._input_ids_source, self._input_ids_buf), ++ ( ++ "query_start_loc", ++ self._query_start_loc_source, ++ self._query_start_loc_buf, ++ ), ++ ] ++ if (self._ngram_context_source is None) != (self._ngram_context_buf is None): ++ raise ValueError("PLE ngram_context source and buffer must match") ++ if self._ngram_context_source is not None: ++ assert self._ngram_context_buf is not None ++ sources.append( ++ ( ++ "ngram_context", ++ self._ngram_context_source, ++ self._ngram_context_buf, ++ ) ++ ) ++ ++ for name, source, buffer in sources: ++ if ( ++ source.device != self.device ++ or source.dtype != buffer.dtype ++ or source.ndim != buffer.ndim ++ or source.shape[0] < buffer.shape[0] ++ or source.shape[1:] != buffer.shape[1:] ++ ): ++ raise ValueError(f"PLE {name} source is incompatible") ++ ++ def _enqueue_cuda_inputs( ++ self, ++ request: PleOffloadRequest, ++ d2h_done_event: torch.cuda.Event, ++ ) -> None: ++ """Stage inputs on the model stream and record D2H completion.""" ++ with ( ++ record_function_or_nullcontext("ple_offload: enqueue_cuda_inputs"), ++ torch.accelerator.device_index(self.device.index), ++ ): ++ stream = torch.cuda.current_stream(self.device) ++ self._input_ids_buf[: request.num_tokens].copy_( ++ self._input_ids_source[: request.num_tokens], ++ non_blocking=True, ++ ) ++ self._query_start_loc_buf[: request.num_reqs + 1].copy_( ++ self._query_start_loc_source[: request.num_reqs + 1], ++ non_blocking=True, ++ ) ++ if self._ngram_context_buf is not None: ++ assert self._ngram_context_source is not None ++ self._ngram_context_buf[: request.num_reqs].copy_( ++ self._ngram_context_source[: request.num_reqs], ++ non_blocking=True, ++ ) ++ d2h_done_event.record(stream) ++ ++ def _launch( ++ self, ++ num_reqs: int, ++ num_tokens: int, ++ *, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ """Stage or queue one batch for request publication.""" ++ # Inputs are replicated across TP ranks. One request per DP rank drives ++ # the CPU result fan-out to every registered TP output buffer. ++ if self.tp_rank != 0: ++ return ++ ++ request = PleOffloadRequest( ++ dp_rank=self.dp_rank, ++ num_tokens=num_tokens, ++ num_reqs=num_reqs, ++ ) ++ assert self._d2h_event_pool is not None, "PLE D2H event pool is not initialized" ++ try: ++ d2h_done_event = self._d2h_event_pool.get_nowait() ++ except queue.Empty as exc: ++ raise RuntimeError( ++ "PLE has more requests than configured concurrent batches" ++ ) from exc ++ self._enqueue_cuda_inputs(request, d2h_done_event) ++ if _FAULT is not None and _FAULT.fires( ++ "drop_request", ++ role="gpu", ++ unit="connector-request", ++ eligible=serving, ++ tp_rank=self.tp_rank, ++ pp_rank=self.pp_rank, ++ dp_rank=self.dp_rank, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ): ++ # Deliberately retain the event: the fault is loss after D2H, not ++ # premature event-pool reuse while that D2H may still be in flight. ++ self._dropped_events.append(d2h_done_event) ++ return ++ self._request_queue.put_nowait( ++ _PendingPleOffloadRequest( ++ request, ++ d2h_done_event, ++ eligible=serving, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ) ++ ) ++ ++ def _fail_closed(self, where: str, error: BaseException) -> None: ++ ple_fail_closed( ++ where=where, ++ rank=self.rank, ++ layer=",".join(self._layers), ++ sequence=self._progress.sequence, ++ error=error, ++ ) ++ ++ @property ++ def has_pending_completions(self) -> bool: ++ return self._progress.has_pending_completions ++ ++ def begin_forward(self, *, dummy_run: bool, num_tokens: int) -> int: ++ """Begin the bounded transaction before any runner fault boundary.""" ++ try: ++ return self._progress.begin( ++ startup=dummy_run, ++ num_tokens=num_tokens, ++ ) ++ except Exception as error: ++ self._fail_closed("prepare_forward", error) ++ ++ def fault_transaction(self, sequence: int) -> str: ++ """Read diagnostic state only when the caller has established arming.""" ++ return self._progress.describe_submission(sequence) ++ ++ def completion_is_pending(self, sequence: int) -> bool: ++ """Return exact-sequence incomplete completion state for cut 6.""" ++ return self._progress.completion_is_pending(sequence) ++ ++ def prepare_forward( ++ self, ++ num_reqs: int, ++ num_tokens: int, ++ dummy_run: bool, ++ *, ++ consumer_num_tokens: int | None = None, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ """Prepare one PLE transaction; any ordinary exception is fatal.""" ++ try: ++ self._prepare_forward_impl( ++ num_reqs, ++ num_tokens, ++ dummy_run, ++ consumer_num_tokens=consumer_num_tokens, ++ serving=serving, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ) ++ except Exception as error: ++ self._fail_closed("prepare_forward", error) ++ ++ def _prepare_forward_impl( ++ self, ++ num_reqs: int, ++ num_tokens: int, ++ dummy_run: bool, ++ *, ++ consumer_num_tokens: int | None = None, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ """Submit real inputs or satisfy the PLE wait for a dummy forward.""" ++ if consumer_num_tokens is None: ++ consumer_num_tokens = num_tokens ++ if int(num_tokens) != int(consumer_num_tokens): ++ raise RuntimeError( ++ "PLE request/consumer token-count mismatch: " ++ f"request={num_tokens}, consumer={consumer_num_tokens}" ++ ) ++ if not 0 <= int(num_tokens) <= self._input_ids_buf.shape[0]: ++ raise RuntimeError( ++ f"PLE token count {num_tokens} exceeds connector capacity " ++ f"{self._input_ids_buf.shape[0]}" ++ ) ++ if not 0 <= int(num_reqs) < self._query_start_loc_buf.shape[0]: ++ raise RuntimeError( ++ f"PLE request count {num_reqs} exceeds connector capacity " ++ f"{self._query_start_loc_buf.shape[0] - 1}" ++ ) ++ self._rowdump_n = None if dummy_run else num_tokens ++ if dummy_run: ++ self._signal_dummy_outputs_impl(num_tokens) ++ return ++ if self._after_dummy: ++ # Every dummy/cleanup reset on THIS rank already completed ++ # (synchronous releases); the barrier extends that to every ++ # stage-0 TP rank before TP0 can publish (CPU group, no GPU work; ++ # all TP ranks run the same forwards in the same order). ++ get_tp_group().barrier() ++ self._after_dummy = False ++ logger.info( ++ "PleOffload: dummy->real transition (dp_rank=%d, tp_rank=%d): " ++ "releases were synchronous, TP barrier passed", ++ self.dp_rank, ++ self.tp_rank, ++ ) ++ self._launch( ++ num_reqs, ++ num_tokens, ++ serving=serving, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ) ++ ++ def signal_dummy_outputs(self, num_tokens: int) -> None: ++ """Locally satisfy PLE waits for dummy and capture forwards.""" ++ self._progress.begin(startup=True, num_tokens=num_tokens) ++ self._signal_dummy_outputs_impl(num_tokens) ++ ++ def _signal_dummy_outputs_impl(self, num_tokens: int) -> None: ++ """Signal dummy outputs after the progress transaction is armed.""" ++ # Dummy and capture forwards do not send CPU requests, but every PLE ++ # placeholder still waits for a completed output semaphore. ++ self._after_dummy = True ++ stream = torch.cuda.current_stream(self.device) ++ for layer in self._layers.values(): ++ layer._gpu_output_buffer[:num_tokens].zero_() ++ layer._sem.signal(stream) ++ ++ def release_outputs( ++ self, ++ *, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ """Release one PLE transaction; any ordinary exception is fatal.""" ++ try: ++ self._release_outputs_impl( ++ serving=serving, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ) ++ self._progress.complete() ++ except Exception as error: ++ self._fail_closed("release_outputs", error) ++ ++ def _release_outputs_impl( ++ self, ++ *, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ """Mark GPU output buffers reusable after the model consumes them.""" ++ if _FAULT is not None and _FAULT.fires( ++ "exc_in_release", ++ role="gpu", ++ unit="connector-release", ++ eligible=serving, ++ tp_rank=self.tp_rank, ++ pp_rank=self.pp_rank, ++ dp_rank=self.dp_rank, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ): ++ raise RuntimeError("PLE test fault: exc_in_release") ++ n = getattr(self, "_rowdump_n", None) ++ self._rowdump_n = None ++ if _rowdump.enabled() and n is not None and not torch.cuda.is_current_stream_capturing(): ++ torch.cuda.current_stream(self.device).synchronize() ++ seq = getattr(self, "_rowdump_seq", 0) ++ self._rowdump_seq = seq + 1 ++ for layer in self._layers.values(): ++ _rowdump.write(f"gpu-tp{self.tp_rank}", 0, self.tp_rank, seq, ++ self._input_ids_source[:n], layer._gpu_output_buffer[:n]) ++ # Reset only after the consumer forward so the CPU worker cannot ++ # overwrite an output that a GPU PLE placeholder may still read. ++ stream = torch.cuda.current_stream(self.device) ++ for layer in self._layers.values(): ++ layer.release_offloaded_output(stream) ++ # Retain the serving deadline until the model stream has actually ++ # executed every reset above. The monitor thread is the sole query site. ++ completion_event = torch.cuda.Event() ++ completion_event.record(stream) ++ self._progress.track_completion(completion_event) ++ # After dummy/capture work (profile, warm-up sizes, capture seed, ++ # adaptive-verification dummy replays, capture cleanup): complete the ++ # reset before returning, so it cannot still be pending when a real ++ # request lets the offload publish. Real-forward releases stay async. ++ # Never inside a stream capture. ++ if self._after_dummy and not torch.cuda.is_current_stream_capturing(): ++ stream.synchronize() ++ ++ def release_host_storage(self) -> None: ++ """Unlink host-storage names while retaining mappings and tensors.""" ++ buffers = [ ++ *self._host_output_bufs.values(), ++ *self._host_flags.values(), ++ self._input_ids_buf, ++ self._query_start_loc_buf, ++ ] ++ if self._ngram_context_buf is not None: ++ buffers.append(self._ngram_context_buf) ++ for buffer in buffers: ++ _unlink_filename_shared_storage( ++ buffer, ++ self._reserved_host_storage_filenames, ++ ) ++ ++ def close(self) -> None: ++ """Stop transport and unlink names, retaining host mappings.""" ++ request_thread = self._request_thread ++ if request_thread is not None and request_thread.is_alive(): ++ try: ++ self._request_queue.put(None, timeout=5) ++ except queue.Full: ++ logger.error("Timed out stopping the PLE request thread") ++ request_thread.join(timeout=5) ++ if request_thread is not None and request_thread.is_alive(): ++ # The thread may still access the registered buffers or ZMQ context. ++ # It may also still race another process opening a shared name, so ++ # this deferred-cleanup branch must not unlink host storage. ++ logger.error("PLE request thread did not stop; deferring resource cleanup") ++ self._progress.close() ++ return ++ self._request_thread = None ++ ++ # Inputs, outputs, and flags remain registered and strongly referenced. ++ # A captured graph or delayed request can still contain their addresses; ++ # the CUDA driver releases registrations when the process exits. ++ registered = len(self._pinned_input_buffers) + len(self._mapped_host_buffers) ++ if registered: ++ logger.debug( ++ "PleOffload: %d registered host buffers retained until exit", ++ registered, ++ ) ++ self._d2h_event_pool = None ++ if self._registration_socket is not None: ++ self._registration_socket.close(linger=0) ++ self._registration_socket = None ++ if self._zmq_ctx is not None: ++ self._zmq_ctx.term() ++ self._zmq_ctx = None ++ self._progress.close() ++ ++ # __init__'s failure handler can reach here before the child has rebuilt ++ # a registration. Startup is already failing in that case; if a payload ++ # was sent, a late rebuild may also observe the removed name. Successful ++ # startup accepts every registration before an external close begins. ++ # ++ # NOTE: Graceful shutdown unlinks the PLE shared-memory names and the ++ # PLE ipc socket. The remaining 17 /tmp UUID sockets per boot are ++ # upstream vLLM VLLM_RPC_BASE_PATH behaviour; a tmpfiles/launch-profile ++ # sweep is recommended. ++ # Keep this last: unlinking preserves every existing CPU/CUDA mapping, ++ # but no transport user may still need to open a name at this point. ++ self.release_host_storage() +diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py +new file mode 100644 +index 0000000..a80cad4 +--- /dev/null ++++ b/vllm/v1/ple_offload/fatal.py +@@ -0,0 +1,153 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Fail-closed helpers for the PLE single-slot transport.""" ++ ++import functools ++import os ++import time ++import traceback ++from collections.abc import Callable ++from typing import Any, NoReturn, TypeVar ++ ++_F = TypeVar("_F", bound=Callable[..., Any]) ++ ++_TRUSTED_ERROR_LABELS = ( ++ (RuntimeError, "RuntimeError"), ++ (TimeoutError, "TimeoutError"), ++ (ValueError, "ValueError"), ++ (OSError, "OSError"), ++ (EOFError, "EOFError"), ++ (BrokenPipeError, "BrokenPipeError"), ++ (AssertionError, "AssertionError"), ++ (KeyError, "KeyError"), ++ (IndexError, "IndexError"), ++ (TypeError, "TypeError"), ++ (MemoryError, "MemoryError"), ++ (KeyboardInterrupt, "KeyboardInterrupt"), ++) ++ ++ ++def ple_fail_closed( ++ *, ++ where: str, ++ rank: int, ++ layer: str, ++ sequence: int, ++ error: BaseException, ++) -> NoReturn: ++ """Emit one best-effort nonblocking diagnostic and terminate immediately.""" ++ try: ++ os.set_blocking(2, False) ++ error_type = type(error) ++ error_name = "Exception" ++ trusted_error = False ++ for trusted_type, trusted_name in _TRUSTED_ERROR_LABELS: ++ if error_type is trusted_type: ++ error_name = trusted_name ++ trusted_error = True ++ break ++ ++ safe_where = where[:128] if type(where) is str else "unknown" ++ safe_layer = layer[:256] if type(layer) is str else "unknown" ++ safe_rank = rank if type(rank) is int else -1 ++ safe_sequence = sequence if type(sequence) is int else -1 ++ detail_suffix = "" ++ if trusted_error: ++ args = BaseException.args.__get__(error, error_type) ++ if ( ++ type(args) is tuple ++ and len(args) == 1 ++ and type(args[0]) is str ++ ): ++ detail = args[0][:512].replace("\n", "\\n") ++ detail_suffix = f":{detail}" ++ message = ( ++ f"PLE fatal: where={safe_where} rank={safe_rank} layer={safe_layer} " ++ f"sequence={safe_sequence} error={error_name}{detail_suffix} " ++ f"pid={os.getpid()} " ++ f"monotonic={time.monotonic():.9f} wall={time.time():.9f}\n" ++ ) ++ os.write(2, message.encode("utf-8", "backslashreplace")[:2048]) ++ except BaseException: ++ pass ++ # THE TRACEBACK, unconditionally, in its own try so it cannot cost us the ++ # line above. ++ # ++ # This path os._exit(1)s, which runs no handlers and flushes nothing, so ++ # whatever is not written here is gone. The line above names the TYPE only ++ # when the type is one of _TRUSTED_ERROR_LABELS and flattens everything ++ # else to the literal string "Exception" -- and the flattened set includes ++ # the two that matter most here: torch.AcceleratorError (a RuntimeError ++ # SUBCLASS, so the `is` comparison misses it) and Triton CompilationError. ++ # An earlier boot died exactly that way: both PP0 ranks ++ # reported "error=Exception" and there was no traceback anywhere in the ++ # log, in any rank, so the cause could not be named at all and the next ++ # step was a re-run rather than a fix. ++ # ++ # A fail-closed path is where evidence matters most, not least: it fires ++ # once, it kills the process, and it is the only witness. Printing the type ++ # of a wrapped exception without the frames that raised it is not a ++ # diagnostic. ++ try: ++ # Three-arg form: it does not depend on being inside an active except ++ # block, and this function is called from one frame further out than ++ # the raise. ++ tb = "".join( ++ traceback.format_exception(type(error), error, error.__traceback__) ++ ) ++ raw = tb.encode("utf-8", "backslashreplace") ++ # Truncate from the END, not the start. format_exception puts the ++ # frames first and the exception TYPE AND MESSAGE last, so a head-side ++ # cut throws away the one line that names what went wrong and keeps the ++ # import-time frames nobody needs. A long chained traceback is exactly ++ # the case where that matters. ++ clipped = len(raw) > 8192 ++ raw = raw[-8192:] ++ os.write( ++ 2, ++ b"PLE fatal traceback (tail) follows:\n" if clipped ++ else b"PLE fatal traceback follows:\n", ++ ) ++ os.write(2, raw) ++ if not raw.endswith(b"\n"): ++ os.write(2, b"\n") ++ except BaseException: ++ # Never silent HERE of all places: without this the reader cannot tell ++ # "the wrapper predates the traceback fix" from "the traceback itself ++ # could not be produced", and those call for different next steps. ++ # Still guarded, because os._exit below must run whatever happens: an ++ # escape from this handler would leave the worker alive in the state ++ # the fail-closed path exists to end. ++ try: ++ os.write(2, b"PLE fatal: traceback unavailable\n") ++ except BaseException: ++ pass ++ os._exit(1) ++ ++ ++def fail_ple_execute_model(func: _F) -> _F: ++ """Fail a non-shutdown worker-step escape only on PLE-owning ranks. ++ ++ This includes ``sample_tokens``, the second half of the split execute-model ++ step. A rank without a PLE connector rethrows the original exception. ++ """ ++ ++ @functools.wraps(func) ++ def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any: ++ try: ++ return func(self, *args, **kwargs) ++ except SystemExit: ++ # WorkerProc's signal handler uses SystemExit for clean shutdown. ++ raise ++ except BaseException as error: ++ connector = getattr(self, "_ple_offload_connector", None) ++ if connector is None: ++ model_runner = getattr(self, "model_runner", None) ++ connector = getattr( ++ model_runner, "_ple_offload_connector", None ++ ) ++ if connector is None: ++ raise ++ connector._fail_closed(func.__name__, error) ++ ++ return wrapped # type: ignore[return-value] +diff --git a/vllm/v1/ple_offload/progress.py b/vllm/v1/ple_offload/progress.py +new file mode 100644 +index 0000000..bf88291 +--- /dev/null ++++ b/vllm/v1/ple_offload/progress.py +@@ -0,0 +1,314 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Independent progress deadline for PLE transactions.""" ++ ++import math ++import os ++import threading ++import time ++from collections import deque ++from dataclasses import dataclass ++from typing import Any ++ ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++_DEFAULT_SERVING_DEADLINE_S = 300.0 ++_DEFAULT_STARTUP_DEADLINE_S = 1800.0 ++_MAX_PROGRESS_DEADLINE_S = 86400.0 ++_POLL_INTERVAL_S = 0.1 ++ ++ ++def _validate_deadline(value: float, name: str) -> float: ++ value = float(value) ++ if ( ++ not math.isfinite(value) ++ or value <= 0 ++ or value > _MAX_PROGRESS_DEADLINE_S ++ ): ++ raise ValueError( ++ f"{name} must be finite, greater than zero, and no greater than " ++ f"{_MAX_PROGRESS_DEADLINE_S:g} seconds" ++ ) ++ return value ++ ++ ++def _environment_deadline(name: str, default: float) -> float: ++ value = float(os.environ.get(name, str(default))) ++ return _validate_deadline(value, name) ++ ++ ++@dataclass(frozen=True) ++class _Submission: ++ sequence: int ++ started: float ++ startup: bool ++ num_tokens: int ++ deadline: float ++ depth: int ++ ++ ++@dataclass(frozen=True) ++class _PendingCompletion: ++ sequence: int ++ started: float ++ startup: bool ++ num_tokens: int ++ deadline: float ++ event: Any ++ ++ ++class PleProgressMonitor: ++ """Bound host submission and the resulting ordered GPU completions.""" ++ ++ def __init__( ++ self, ++ *, ++ rank: int, ++ layer_names: tuple[str, ...], ++ serving_deadline_s: float, ++ startup_deadline_s: float, ++ ) -> None: ++ serving_deadline_s = _validate_deadline( ++ serving_deadline_s, "serving_deadline_s" ++ ) ++ startup_deadline_s = _validate_deadline( ++ startup_deadline_s, "startup_deadline_s" ++ ) ++ self._rank = rank ++ self._layer = ",".join(layer_names) ++ self._serving_deadline_s = serving_deadline_s ++ self._startup_deadline_s = startup_deadline_s ++ self._condition = threading.Condition() ++ self._submission: _Submission | None = None ++ self._pending: deque[_PendingCompletion] = deque() ++ self._sequence = 0 ++ self._stopping = False ++ self._thread = threading.Thread( ++ target=self._run, ++ name=f"ple-progress-rank{rank}", ++ daemon=True, ++ ) ++ self._thread.start() ++ ++ @classmethod ++ def from_environment( ++ cls, *, rank: int, layer_names: tuple[str, ...] ++ ) -> "PleProgressMonitor": ++ return cls( ++ rank=rank, ++ layer_names=layer_names, ++ serving_deadline_s=_environment_deadline( ++ "VLLM_PLE_PROGRESS_DEADLINE_S", ++ _DEFAULT_SERVING_DEADLINE_S, ++ ), ++ startup_deadline_s=_environment_deadline( ++ "VLLM_PLE_STARTUP_DEADLINE_S", ++ _DEFAULT_STARTUP_DEADLINE_S, ++ ), ++ ) ++ ++ @property ++ def sequence(self) -> int: ++ with self._condition: ++ return self._sequence ++ ++ @property ++ def has_pending_completions(self) -> bool: ++ with self._condition: ++ return bool(self._pending) ++ ++ def describe_submission(self, sequence: int) -> str: ++ """Return the armed transaction descriptor for the active sequence.""" ++ with self._condition: ++ active = self._submission ++ if active is None or active.sequence != sequence: ++ raise RuntimeError("PLE fault correlation has no active submission") ++ return ( ++ f"sequence:{active.sequence}" ++ f"@begin:{active.started:.9f}" ++ ) ++ ++ def completion_is_pending(self, sequence: int) -> bool: ++ """Whether this serving sequence has a retained, incomplete event. ++ ++ This method is called only behind the armed fault guard. An event that ++ has completed on the GPU but has not yet been dequeued is not counted. ++ """ ++ with self._condition: ++ pending = next( ++ ( ++ item ++ for item in self._pending ++ if item.sequence == sequence and not item.startup ++ ), ++ None, ++ ) ++ if pending is None: ++ return False ++ return not bool(pending.event.query()) ++ ++ def begin(self, *, startup: bool, num_tokens: int) -> int: ++ now = time.monotonic() ++ budget = ( ++ self._startup_deadline_s if startup else self._serving_deadline_s ++ ) ++ with self._condition: ++ active = self._submission ++ if active is not None: ++ if not startup or not active.startup: ++ raise RuntimeError( ++ "Overlapping PLE host submissions violate the serialized " ++ "single-slot protocol" ++ ) ++ sequence = active.sequence ++ self._submission = _Submission( ++ sequence=sequence, ++ started=active.started, ++ startup=True, ++ num_tokens=num_tokens, ++ deadline=now + budget, ++ depth=active.depth + 1, ++ ) ++ else: ++ self._sequence += 1 ++ sequence = self._sequence ++ self._submission = _Submission( ++ sequence=sequence, ++ started=now, ++ startup=startup, ++ num_tokens=num_tokens, ++ deadline=now + budget, ++ depth=1, ++ ) ++ self._condition.notify_all() ++ return sequence ++ ++ def track_completion(self, event: Any) -> None: ++ """Retain a release event until the monitor observes GPU completion.""" ++ if not callable(getattr(event, "query", None)): ++ raise TypeError("PLE completion event must provide query()") ++ with self._condition: ++ active = self._submission ++ if active is None: ++ raise RuntimeError("PLE completion has no matching submission") ++ self._pending.append( ++ _PendingCompletion( ++ sequence=active.sequence, ++ started=active.started, ++ startup=active.startup, ++ num_tokens=active.num_tokens, ++ deadline=active.deadline, ++ event=event, ++ ) ++ ) ++ self._condition.notify_all() ++ ++ def complete(self) -> None: ++ with self._condition: ++ active = self._submission ++ if active is None or active.depth <= 0: ++ raise RuntimeError("PLE release has no matching transaction") ++ if active.depth == 1: ++ self._submission = None ++ else: ++ self._submission = _Submission( ++ sequence=active.sequence, ++ started=active.started, ++ startup=active.startup, ++ num_tokens=active.num_tokens, ++ deadline=active.deadline, ++ depth=active.depth - 1, ++ ) ++ self._condition.notify_all() ++ ++ def close(self) -> None: ++ with self._condition: ++ self._stopping = True ++ self._condition.notify_all() ++ self._thread.join(timeout=1) ++ ++ def _run(self) -> None: ++ try: ++ self._monitor() ++ except BaseException as error: ++ ple_fail_closed( ++ where="progress-monitor-thread", ++ rank=self._rank, ++ layer=self._layer, ++ sequence=self._sequence, ++ error=error, ++ ) ++ ++ def _monitor(self) -> None: ++ while True: ++ with self._condition: ++ if self._stopping: ++ return ++ submission = self._submission ++ pending = self._pending[0] if self._pending else None ++ if submission is None and pending is None: ++ self._condition.wait(timeout=_POLL_INTERVAL_S) ++ continue ++ ++ # CUDA completion is observed only here. Never synchronize and never ++ # use the host flag as evidence that the queued reset completed. ++ if pending is not None and pending.event.query(): ++ with self._condition: ++ if self._pending and self._pending[0] is pending: ++ self._pending.popleft() ++ self._condition.notify_all() ++ continue ++ ++ now = time.monotonic() ++ if pending is not None and now >= pending.deadline: ++ with self._condition: ++ if not self._pending or self._pending[0] is not pending: ++ continue ++ mode = "startup" if pending.startup else "serving" ++ ple_fail_closed( ++ where=( ++ f"progress-deadline:{mode}:gpu-completion:" ++ f"tokens={pending.num_tokens}" ++ ), ++ rank=self._rank, ++ layer=self._layer, ++ sequence=pending.sequence, ++ error=TimeoutError("PLE GPU completion made no bounded progress"), ++ ) ++ ++ if submission is not None and now >= submission.deadline: ++ with self._condition: ++ if self._submission is not submission: ++ continue ++ mode = "startup" if submission.startup else "serving" ++ ple_fail_closed( ++ where=( ++ f"progress-deadline:{mode}:host-submission:" ++ f"tokens={submission.num_tokens}" ++ ), ++ rank=self._rank, ++ layer=self._layer, ++ sequence=submission.sequence, ++ error=TimeoutError("PLE host submission made no bounded progress"), ++ ) ++ ++ deadlines = [] ++ if pending is not None: ++ deadlines.append(pending.deadline) ++ if submission is not None: ++ deadlines.append(submission.deadline) ++ timeout = _POLL_INTERVAL_S ++ if deadlines: ++ timeout = min( ++ timeout, ++ max(0.0, min(deadlines) - time.monotonic()), ++ ) ++ with self._condition: ++ if self._stopping: ++ return ++ if ( ++ self._submission is not submission ++ or (self._pending[0] if self._pending else None) is not pending ++ ): ++ continue ++ self._condition.wait(timeout=timeout) +diff --git a/vllm/v1/ple_offload/protocol.py b/vllm/v1/ple_offload/protocol.py +new file mode 100644 +index 0000000..87eba13 +--- /dev/null ++++ b/vllm/v1/ple_offload/protocol.py +@@ -0,0 +1,50 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""IPC message definitions for PLE CPU offload (host-mapped pull design). ++ ++The offload process owns no CUDA context. Every GPU worker owns, per PLE ++layer, a shared host output buffer and a shared host int32 flag, both ++page-locked and device-mapped in the worker's context. The offload process ++waits for ``flag == 0``, writes the layer result into the host buffer and ++stores ``flag = 1`` with plain CPU stores; the GPU's captured forward waits on ++the host flag, DMAs host -> GPU on the model stream, and ``release_outputs`` ++resets the flag after the forward. ++""" ++ ++from dataclasses import dataclass ++ ++import msgspec ++import torch ++ ++# --------------------------------------------------------------------------- ++# IPC message dataclasses ++# --------------------------------------------------------------------------- ++ ++ ++@dataclass ++class PleOffloadRegistration: ++ """Sent once from each GPU worker during offload setup.""" ++ ++ worker_id: int ++ tp_rank: int ++ dp_rank: int ++ # Shared CPU tensors per PLE layer, filled by the offload process. ++ host_output_buffers: dict[str, torch.Tensor] ++ # Shared CPU int32 tensors per PLE layer; element 0 is the ready flag. ++ host_flag_tensors: dict[str, torch.Tensor] ++ # CPU tensors are allocated in shared memory and registered once. ++ input_ids_buf: torch.Tensor ++ query_start_loc_buf: torch.Tensor ++ ngram_context_buf: torch.Tensor | None ++ ++ ++@dataclass ++class PleOffloadRequest: ++ """Sent by each DP rank's TP rank zero at every inference step.""" ++ ++ dp_rank: int ++ num_tokens: int ++ num_reqs: int ++ ++ ++_PLE_OFFLOAD_REQUEST_DECODER = msgspec.msgpack.Decoder(PleOffloadRequest) +diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py +new file mode 100644 +index 0000000..43caf80 +--- /dev/null ++++ b/vllm/v1/ple_offload/worker.py +@@ -0,0 +1,1298 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Dedicated CPU-offload process for PLE embedding layers. ++ ++This module implements a standalone process that: ++1. Loads only the :class:`PleOffloadLayer` weights into CPU memory. ++2. Accepts per-step computation requests from GPU worker processes. ++3. Runs ``forward_impl()`` on CPU, writes the rows into every TP worker's ++ shared host output buffer for the requesting DP rank and sets that ++ worker's shared host flag (host-mapped pull design: this process holds no ++ CUDA context; the GPU waits on the host flag and DMAs host -> GPU inside ++ its captured forward). ++ ++The TP workers within one DP rank receive identical inputs, so the CPU result ++is computed once per DP rank and fanned out to all of its TP ranks. ++ ++Class structure mirrors the GPU worker pattern in multiproc_executor.py: ++ ++ PleOffloadWorkerHandle -- handle held by the spawning GPU worker ++ PleOffloadWorker -- process lifecycle and READY handshake ++ PleOffloadRunner -- owns weights and serves inference requests ++""" ++ ++import contextlib ++import os ++import pickle ++import select ++import signal ++import socket ++import struct ++import subprocess ++import sys ++import tempfile ++import threading ++import time ++from collections.abc import Iterable ++from dataclasses import dataclass ++from multiprocessing.connection import Connection ++from typing import Any, cast ++ ++import msgspec ++import torch ++import torch.distributed as dist ++import zmq ++ ++import vllm.envs as envs ++from vllm.config import VllmConfig, set_current_vllm_config ++from vllm.distributed.parallel_state import ( ++ ensure_model_parallel_initialized, ++ init_distributed_environment, ++) ++from vllm.logger import init_logger ++from vllm.model_executor.layers.ple_offload_layer import ( ++ PleOffloadLayer, ++ mark_as_offload_worker, ++) ++from vllm.model_executor.model_loader import get_model_loader ++from vllm.model_executor.model_loader.default_loader import DefaultModelLoader ++from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader ++from vllm.model_executor.model_loader.utils import ( ++ initialize_model, ++ process_weights_after_loading, ++) ++from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights ++from vllm.utils.system_utils import decorate_logs ++from vllm.utils.torch_utils import set_default_torch_dtype ++from vllm.v1.ple_offload.protocol import ( ++ _PLE_OFFLOAD_REQUEST_DECODER, ++ PleOffloadRegistration, ++ PleOffloadRequest, ++) ++from vllm.v1.utils import record_function_or_nullcontext ++from vllm_ple_test_fault import ( ++ _FAULT, ++ _FaultReconstructionPayload, ++) ++ ++logger = init_logger(__name__) ++ ++# Publication fence for the host-mapped pull: a locked read-modify-write (the ++# uncontended lock acquire) is a full barrier on x86 and also drains ++# write-combining buffers, so every row store of a result is globally visible ++# before the flag store that tells the GPU to DMA it (do not rely ++# on "memcpy uses no non-temporal stores at this size"). ++_PUBLISH_FENCE = threading.Lock() ++ ++ ++def _unlink_ipc_socket(ipc_addr: str) -> None: ++ """Best-effort unlink of the exact filesystem IPC endpoint PLE owns.""" ++ if not ipc_addr.startswith("ipc://"): ++ return ++ path = ipc_addr.removeprefix("ipc://") ++ # PLE receives an absolute UUID path from get_open_zmq_ipc_path(). Avoid ++ # treating abstract or relative ZMQ addresses as filesystem ownership. ++ if not os.path.isabs(path): ++ return ++ try: ++ os.unlink(path) ++ except FileNotFoundError: ++ pass ++ except OSError as error: ++ logger.warning("Failed to unlink PLE IPC socket %s: %s", path, error) ++ ++ ++def _publish_fence() -> None: ++ with _PUBLISH_FENCE: ++ pass ++ ++ ++def _offload_child_environment() -> dict[str, str]: ++ """Build the offload environment without mutating the parent.""" ++ child_env = dict(os.environ) ++ child_env["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1" ++ child_env.setdefault("CUTLASS_CUDA_DEVICE_ID", "0") ++ child_env.pop("VLLM_PP_LAYER_PARTITION", None) ++ ++ home_device = child_env.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip() ++ if home_device: ++ parent_visible = child_env.get("CUDA_VISIBLE_DEVICES", "").strip() ++ try: ++ physical = ( ++ parent_visible.split(",")[int(home_device)].strip() ++ if parent_visible ++ else home_device ++ ) ++ except (IndexError, ValueError) as error: ++ raise ValueError("Invalid VLLM_PLE_OFFLOAD_HOME_DEVICE") from error ++ child_env["CUDA_VISIBLE_DEVICES"] = physical ++ return child_env ++ ++ ++def _prepare_cuda_free_process() -> None: ++ """Keep the PLE offload process free of any CUDA context. ++ ++ The host-mapped pull design issues no CUDA work here, yet model ++ construction reached CUDA incidentally (found by the offload-startup ++ replays): ++ 1. ``VLLM_DISABLE_SHARED_EXPERTS_STREAM=1``: each MoE layer's shared-experts ++ runner calls ``aux_stream()`` -> ``torch.cuda.Stream()``, whose C++ ++ constructor creates the primary context without calling ++ ``torch.cuda._lazy_init`` (``torch.cuda.is_initialized()`` stays False). ++ This is the context the replay caught (stack dumped inside ++ ``Stream.__new__``). This process never runs an MoE layer, so it never ++ needs the stream. vLLM reads the variable when the layer is built; this ++ process does not enable vLLM's env cache. ++ 2. ``CUTLASS_CUDA_DEVICE_ID``: CUTLASS python (``cutlass_cppgen``) runs ++ ``initialize_cuda_context()`` -> ``cudart.cudaFree(0)`` at import unless ++ it is set. Not observed on this import path; kept as a cheap defence. ++ 3. The torch guard (``_install_cuda_guard``). ++ """ ++ os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1" ++ os.environ.setdefault("CUTLASS_CUDA_DEVICE_ID", "0") ++ _install_cuda_guard() ++ ++ ++def _install_cuda_guard() -> None: ++ """Make torch-level CUDA initialization raise in the PLE offload process. ++ ++ An offload-startup check found the host-mapped pull still got a 322 MiB context ++ on its home card: importing the Qwen4Exp model imports the vendored FLA ++ ops, whose utils call triton's ``get_current_target()`` at import time, ++ i.e. ``torch.cuda.current_device()`` -> lazy init. FLA catches ++ ``RuntimeError`` and falls back to "cpu" (and then skips its Hopper ++ probe), so raising here keeps this process context-free. Any other call ++ site fails loudly at offload startup with its stack logged, instead of ++ silently taking VRAM on a stage-0 card. The GPU workers are untouched. ++ ++ torch also has C++ entry points that create the context without calling ++ ``torch.cuda._lazy_init``: stream construction (``torch.cuda.Stream``) and ++ the pinned host allocator (``Tensor.pin_memory``). Both raise here too; ++ this process needs neither (all of its buffers are plain CPU memory). ++ ``torch.cuda.is_initialized()`` only reflects the Python-level path, so ++ the per-PID GPU sampler, not that flag, is the authority on contexts. ++ """ ++ import traceback ++ ++ if torch.cuda.is_initialized(): ++ # Fail closed: a context that already exists would stay on a ++ # stage-0 card for the whole run. ++ raise RuntimeError( ++ "CUDA is already initialized in the PLE offload process before its " ++ "guard; the host-mapped pull offload must stay CUDA-free" ++ ) ++ ++ def _refuse(what: str) -> None: ++ stack = "".join(traceback.format_stack(limit=12)[:-2]) ++ logger.info( ++ "PLE offload CUDA guard: blocked a CUDA initialization attempt " ++ "(%s; the caller must fall back to CPU or fail); call site:\n%s", ++ what, ++ stack, ++ ) ++ raise RuntimeError( ++ f"CUDA is disabled in the PLE offload process (host-mapped pull): {what}" ++ ) ++ ++ def _blocked_lazy_init() -> None: ++ _refuse("torch.cuda lazy init") ++ ++ def _blocked_stream(cls, *args, **kwargs): ++ _refuse("torch.cuda.Stream construction") ++ ++ def _blocked_pin_memory(self, *args, **kwargs): ++ _refuse("Tensor.pin_memory") ++ ++ torch.cuda._lazy_init = _blocked_lazy_init # type: ignore[assignment] ++ torch.cuda.Stream.__new__ = staticmethod(_blocked_stream) # type: ignore[assignment] ++ torch.Tensor.pin_memory = _blocked_pin_memory # type: ignore[assignment] ++ logger.info( ++ "PLE offload CUDA guard installed (torch CUDA lazy init, stream " ++ "construction and pin_memory raise)" ++ ) ++ ++ ++@contextlib.contextmanager ++def _platform_device_is_meta(): ++ """During the offload's meta-model construction, report the meta device. ++ ++ Some layer constructors pass an explicit ``current_platform.current_device()`` ++ (e.g. the GDN output norm in qwen_gdn_linear_attn.py), which overrides the ++ ``torch.device("meta")`` context and would allocate on CUDA. In the offload ++ process nothing but the CPU PLE layers is ever materialized, so those ++ tensors belong on meta too. The patch is an instance attribute removed on ++ exit, restoring the platform's own method. ++ """ ++ from vllm.platforms import current_platform ++ ++ current_platform.current_device = lambda *args, **kwargs: torch.device("meta") ++ try: ++ yield ++ finally: ++ with contextlib.suppress(AttributeError): ++ del current_platform.current_device ++ ++ ++# Bounded wait for the GPU to release a result (seconds). A timed-out ownership ++# transfer is fatal: the offload process raises and exits, never overwrites. ++_PLE_WAIT_TIMEOUT_S = float(os.environ.get("VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", "120")) ++ ++ ++@dataclass ++class PleOffloadOutputTarget: ++ """Host output destination and ready flag for one TP worker. ++ ++ Both are shared CPU tensors owned by the GPU worker, page-locked and ++ device-mapped in that worker's context. This process only reads/writes ++ them with plain CPU loads/stores (no CUDA): wait for ``flag == 0`` (the GPU ++ released the previous result), write the rows, then store ``flag = 1``. ++ """ ++ ++ tp_rank: int ++ host_output_buffer: torch.Tensor ++ host_flag: torch.Tensor # int32, element 0 is the flag ++ ++ ++@dataclass ++class PleOffloadInputBuffers: ++ """Shared-memory input buffers registered for one DP rank.""" ++ ++ input_ids_buf: torch.Tensor # int32 (max_num_tokens,) ++ query_start_loc_buf: torch.Tensor # int32 (max_num_reqs + 1,) ++ ngram_context_buf: torch.Tensor | None # int32 (max_num_reqs, ngram_context_len) ++ ++ ++@dataclass(frozen=True) ++class _ReceivedPleOffloadRequest: ++ request: PleOffloadRequest ++ serving: bool ++ transaction_id: str | None ++ correlation_id: str | None ++ ++ ++def _decode_offload_request( ++ frames: list[bytes], ++) -> _ReceivedPleOffloadRequest: ++ if len(frames) != 2: ++ raise RuntimeError("PLE offload request must contain two frames") ++ request = _PLE_OFFLOAD_REQUEST_DECODER.decode(frames[0]) ++ metadata = msgspec.msgpack.decode(frames[1]) ++ if type(metadata) is not dict or set(metadata) != { ++ "serving", ++ "transaction", ++ "correlation", ++ }: ++ raise RuntimeError("PLE offload request metadata is invalid") ++ serving = metadata["serving"] ++ transaction = metadata["transaction"] ++ correlation = metadata["correlation"] ++ if type(serving) is not bool: ++ raise RuntimeError("PLE offload serving metadata is invalid") ++ if transaction is not None and type(transaction) is not str: ++ raise RuntimeError("PLE offload transaction metadata is invalid") ++ if correlation is not None and type(correlation) is not str: ++ raise RuntimeError("PLE offload correlation metadata is invalid") ++ return _ReceivedPleOffloadRequest( ++ request=request, ++ serving=serving, ++ transaction_id=transaction, ++ correlation_id=correlation, ++ ) ++ ++ ++def _receive_offload_request( ++ socket: zmq.Socket, flags: int = 0 ++) -> _ReceivedPleOffloadRequest | PleOffloadRequest: ++ """Disarmed: the unchanged single-frame request. Armed: request + metadata.""" ++ if _FAULT is None: ++ return _PLE_OFFLOAD_REQUEST_DECODER.decode(socket.recv(flags)) ++ return _decode_offload_request(socket.recv_multipart(flags)) ++ ++ ++def _send_framed_payload(fd: int, payload: bytes, deadline: float) -> None: ++ """Send one launch payload without blocking beyond an absolute deadline.""" ++ frame = memoryview(struct.pack("!Q", len(payload)) + payload) ++ offset = 0 ++ os.set_blocking(fd, False) ++ try: ++ while offset < len(frame): ++ remaining = deadline - time.monotonic() ++ if remaining <= 0: ++ raise TimeoutError("PLE launch payload transmission timed out") ++ _, writable, _ = select.select([], [fd], [], remaining) ++ if not writable: ++ raise TimeoutError("PLE launch payload transmission timed out") ++ try: ++ sent = os.write(fd, frame[offset:]) ++ except BlockingIOError: ++ continue ++ if sent <= 0: ++ raise BrokenPipeError("PLE launch control socket closed") ++ offset += sent ++ if ( ++ _FAULT is not None ++ and _FAULT.cut == "launch_park_before_payload_read" ++ ): ++ _FAULT.acknowledge_launch_send( ++ bytes_sent=offset, ++ frame_size=len(frame), ++ ) ++ finally: ++ with contextlib.suppress(OSError): ++ os.set_blocking(fd, True) ++ ++ ++_PLE_READY_MAX_PAYLOAD_BYTES = 1 << 20 ++ ++class _BoundedStartupResponseWriter: ++ """Preflight startup responses against the parent receiver's frame cap.""" ++ ++ def __init__(self, writer: Connection) -> None: ++ self._writer = writer ++ ++ def send(self, message: Any) -> None: ++ payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL) ++ if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES: ++ payload = pickle.dumps( ++ { ++ "status": "FAILURE", ++ "error": "PLE startup response exceeded size limit", ++ }, ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES: ++ raise RuntimeError( ++ "Bounded PLE startup failure response exceeds its frame cap" ++ ) ++ self._writer.send_bytes(payload) ++ ++ def close(self) -> None: ++ self._writer.close() ++ ++ ++def _read_exact_before_deadline( ++ fd: int, ++ size: int, ++ deadline: float, ++) -> bytes: ++ """Read exactly ``size`` bytes without crossing an absolute deadline.""" ++ chunks = bytearray() ++ while len(chunks) < size: ++ remaining = deadline - time.monotonic() ++ if remaining <= 0: ++ raise TimeoutError( ++ "PLE offload worker did not become ready before its " ++ "absolute launch deadline" ++ ) ++ readable, _, _ = select.select([fd], [], [], remaining) ++ if not readable: ++ raise TimeoutError( ++ "PLE offload worker did not become ready before its " ++ "absolute launch deadline" ++ ) ++ try: ++ chunk = os.read(fd, size - len(chunks)) ++ except BlockingIOError: ++ continue ++ if not chunk: ++ raise EOFError("PLE startup response was truncated") ++ chunks.extend(chunk) ++ return bytes(chunks) ++ ++ ++def _recv_bounded_connection_message( ++ reader: Connection, ++ deadline: float, ++) -> Any: ++ """Receive one bounded multiprocessing Connection frame by ``deadline``.""" ++ fd = reader.fileno() ++ os.set_blocking(fd, False) ++ try: ++ header = _read_exact_before_deadline(fd, 4, deadline) ++ size = struct.unpack("!i", header)[0] ++ if size == -1: ++ size = struct.unpack( ++ "!Q", _read_exact_before_deadline(fd, 8, deadline) ++ )[0] ++ elif size < 0: ++ raise RuntimeError("PLE startup response has an invalid frame size") ++ if size > _PLE_READY_MAX_PAYLOAD_BYTES: ++ raise RuntimeError( ++ f"PLE startup response exceeds " ++ f"{_PLE_READY_MAX_PAYLOAD_BYTES} bytes" ++ ) ++ frame = _read_exact_before_deadline(fd, size, deadline) ++ finally: ++ with contextlib.suppress(OSError): ++ os.set_blocking(fd, True) ++ try: ++ return pickle.loads(frame) ++ except BaseException as error: ++ raise RuntimeError( ++ "PLE startup response is not a valid serialized message" ++ ) from error ++ ++ ++@dataclass ++class PleOffloadWorkerHandle: ++ """Resources owned by the GPU worker that spawned the offload process.""" ++ ++ proc: subprocess.Popen ++ death_writer: Connection | None ++ ready_pipe_reader: Connection | None ++ launch_deadline: float ++ ipc_addr: str | None = None ++ ++ def close(self, *, graceful: bool = True) -> None: ++ """Release all process resources. Safe to call more than once.""" ++ if self.ready_pipe_reader is not None: ++ self.ready_pipe_reader.close() ++ self.ready_pipe_reader = None ++ if graceful and self.proc.poll() is None: ++ # Keep parent-death supervision connected while SIGTERM lets ++ # proc_main run its shutdown event and finally cleanup. ++ self.proc.terminate() ++ try: ++ self.proc.wait(timeout=5) ++ except subprocess.TimeoutExpired: ++ pass ++ if self.death_writer is not None: ++ self.death_writer.close() ++ self.death_writer = None ++ if not graceful and self.proc.poll() is None: ++ # Expired launch: retain the immediate EOF -> TERM path. ++ self.proc.terminate() ++ try: ++ self.proc.wait(timeout=5) ++ except subprocess.TimeoutExpired: ++ pass ++ if self.proc.poll() is None: ++ self.proc.kill() ++ try: ++ self.proc.wait(timeout=5) ++ except subprocess.TimeoutExpired: ++ pass ++ # The child normally removes the endpoint in proc_main's finally. ++ # Once it has exited, this exact PLE-owned address is safe to remove as ++ # a fallback for SIGKILL or an interrupted child cleanup. ++ if self.proc.poll() is not None and self.ipc_addr is not None: ++ _unlink_ipc_socket(self.ipc_addr) ++ ++ ++def _init_offload_distributed() -> None: ++ """Initialize the single-rank Gloo world required by TP-aware layers.""" ++ if dist.is_initialized(): ++ return ++ ++ # VocabParallelEmbedding reads the TP process group during construction. ++ # The offload process owns the full embedding table, so it uses an isolated ++ # TP1/PP1 Gloo world and never joins the GPU workers' NCCL groups. ++ store_dir = tempfile.mkdtemp(prefix="vllm_ple_offload_") ++ init_distributed_environment( ++ world_size=1, ++ rank=0, ++ distributed_init_method=f"file://{store_dir}/store", ++ local_rank=0, ++ backend="gloo", ++ ) ++ # initialize_model_parallel reads the active VllmConfig in the current ++ # vLLM version. Explicitly configure DP1/TP1/PP1 to match the isolated ++ # world, regardless of any DP environment variables inherited from the GPU ++ # worker. The real DP/TP configuration is used later for model construction, ++ # registration, and request routing. ++ offload_config = VllmConfig() ++ offload_parallel_config = offload_config.parallel_config ++ offload_parallel_config.data_parallel_size = 1 ++ offload_parallel_config.data_parallel_size_local = 1 ++ offload_parallel_config.data_parallel_rank = 0 ++ offload_parallel_config.data_parallel_rank_local = 0 ++ offload_parallel_config.data_parallel_index = 0 ++ offload_parallel_config.tensor_parallel_size = 1 ++ offload_parallel_config.pipeline_parallel_size = 1 ++ offload_parallel_config.prefill_context_parallel_size = 1 ++ offload_parallel_config.decode_context_parallel_size = 1 ++ offload_parallel_config.world_size = 1 ++ offload_parallel_config.nnodes = 1 ++ offload_parallel_config.node_rank = 0 ++ with set_current_vllm_config(offload_config): ++ ensure_model_parallel_initialized( ++ tensor_model_parallel_size=1, ++ pipeline_model_parallel_size=1, ++ backend="gloo", ++ ) ++ logger.info( ++ "Distributed environment initialized (backend=gloo, rank=0, world_size=1)." ++ ) ++ ++ ++class PleOffloadWorker: ++ """Manage process creation, READY handshake, and the child entry point.""" ++ ++ READY_STR = "READY" ++ ++ @staticmethod ++ def make_process( ++ vllm_config: VllmConfig, ++ num_workers: int, ++ ipc_addr: str, ++ ) -> PleOffloadWorkerHandle: ++ """Start one CPU offload process in a fresh interpreter.""" ++ launch_deadline = ( ++ time.monotonic() + float(envs.VLLM_PLE_OFFLOAD_READY_TIMEOUT) ++ ) ++ config_payload = pickle.dumps( ++ vllm_config, protocol=pickle.HIGHEST_PROTOCOL ++ ) ++ child_env = _offload_child_environment() ++ ++ # READY and parent-death supervision must never compete for reads. ++ parent_control, child_control = socket.socketpair() ++ parent_death, child_death = socket.socketpair() ++ ready_reader = Connection(parent_control.detach()) ++ death_writer = Connection(parent_death.detach()) ++ try: ++ proc = subprocess.Popen( ++ [ ++ sys.executable, ++ "-m", ++ "vllm_ple_offload_bootstrap", ++ str(child_control.fileno()), ++ str(child_death.fileno()), ++ ], ++ env=child_env, ++ close_fds=True, ++ pass_fds=(child_control.fileno(), child_death.fileno()), ++ ) ++ except BaseException: ++ ready_reader.close() ++ death_writer.close() ++ raise ++ finally: ++ child_control.close() ++ child_death.close() ++ ++ handle = PleOffloadWorkerHandle( ++ proc=proc, ++ death_writer=death_writer, ++ ready_pipe_reader=ready_reader, ++ launch_deadline=launch_deadline, ++ ipc_addr=ipc_addr, ++ ) ++ try: ++ config_for_launch: bytes | _FaultReconstructionPayload = config_payload ++ if ( ++ _FAULT is not None ++ and _FAULT.cut == "launch_block_reconstruction" ++ ): ++ config_for_launch = _FaultReconstructionPayload(config_payload) ++ launch_state = { ++ "vllm_config_payload": config_for_launch, ++ "num_workers": num_workers, ++ "ipc_addr": ipc_addr, ++ } ++ if ( ++ _FAULT is not None ++ and _FAULT.cut == "launch_park_before_payload_read" ++ ): ++ # Larger than ordinary AF_UNIX socket buffers: an acknowledged ++ # non-reading child leaves the parent inside framed transmission. ++ launch_state["_test_padding"] = b"\0" * (8 << 20) ++ launch_payload = pickle.dumps( ++ launch_state, ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ _send_framed_payload( ++ ready_reader.fileno(), launch_payload, launch_deadline ++ ) ++ except BaseException: ++ handle.close(graceful=False) ++ raise ++ ++ return handle ++ ++ @staticmethod ++ def wait_for_ready(handle: PleOffloadWorkerHandle) -> None: ++ """Wait until weights and all GPU registrations are ready to serve.""" ++ reader = handle.ready_pipe_reader ++ if reader is None: ++ return ++ try: ++ message = _recv_bounded_connection_message( ++ reader, handle.launch_deadline ++ ) ++ if type(message) is not dict: ++ raise RuntimeError("PLE offload worker sent an invalid response") ++ status = dict.get(message, "status") ++ if type(status) is not str: ++ raise RuntimeError( ++ "PLE offload worker sent an invalid response status" ++ ) ++ if status == PleOffloadWorker.READY_STR: ++ pass ++ elif status == "FAILURE": ++ error_detail = dict.get(message, "error") ++ if type(error_detail) is not str: ++ error_detail = "unknown error" ++ raise RuntimeError( ++ "PLE offload worker failed during startup: " ++ f"{error_detail}" ++ ) ++ else: ++ raise RuntimeError("PLE offload worker sent unsupported status") ++ layer_names = dict.get(message, "layer_names") ++ if ( ++ type(layer_names) is not list ++ or any( ++ type(name) is not str for name in layer_names ++ ) ++ ): ++ raise RuntimeError( ++ "PLE offload worker sent invalid READY layer_names" ++ ) ++ logger.info( ++ "Worker ready - %d PleOffloadLayer(s): %s", ++ len(layer_names), ++ layer_names, ++ ) ++ except EOFError as error: ++ handle.close(graceful=False) ++ raise RuntimeError("PLE offload worker exited during startup") from error ++ except BaseException: ++ handle.close(graceful=False) ++ raise ++ finally: ++ if handle.ready_pipe_reader is not None: ++ handle.ready_pipe_reader.close() ++ handle.ready_pipe_reader = None ++ ++ @staticmethod ++ def proc_main( ++ vllm_config: VllmConfig, ++ num_workers: int, ++ ipc_addr: str, ++ ready_pipe: Connection, ++ ) -> None: ++ """Load PLE weights, accept registrations, and run the request loop.""" ++ decorate_logs("PleOffloadWorker") ++ ready_writer = _BoundedStartupResponseWriter(ready_pipe) ++ shutdown_event = threading.Event() ++ ++ def handle_signal(_signum: int, _frame: object) -> None: ++ shutdown_event.set() ++ ++ signal.signal(signal.SIGTERM, handle_signal) ++ signal.signal(signal.SIGINT, handle_signal) ++ ++ zmq_context: zmq.Context | None = None ++ pull_socket: zmq.Socket | None = None ++ owns_ipc_socket = False ++ try: ++ # The flag lets PleOffloadLayer subclasses execute their complete ++ # constructors instead of becoming empty GPU-worker placeholders. ++ mark_as_offload_worker() ++ ++ # The bootstrap established child-only visibility before importing ++ # torch or vLLM. Retain the Python-level guards for known CUDA ++ # initialization paths and fail if initialization already occurred. ++ _prepare_cuda_free_process() ++ home_device = os.environ.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip() ++ if home_device: ++ logger.info( ++ "PLE offload bootstrap visibility is CUDA_VISIBLE_DEVICES=%s " ++ "(host-mapped pull: no CUDA context expected)", ++ os.environ.get("CUDA_VISIBLE_DEVICES", ""), ++ ) ++ ++ # Initialize Gloo before installing the real VllmConfig. This keeps ++ # the CPU process in an isolated rank-zero, world-size-one group. ++ # The offload worker is a PP=1 world; a VLLM_PP_LAYER_PARTITION set ++ # for the GPU workers (e.g. "25,23") would make make_layers raise ++ # "len(partitions) does not match pp_size=1" while the meta model ++ # is built here. Drop it in this process only. ++ os.environ.pop("VLLM_PP_LAYER_PARTITION", None) ++ _init_offload_distributed() ++ ++ # Model components read the active VllmConfig while the meta model ++ # is constructed, so keep the context around runner initialization. ++ with set_current_vllm_config(vllm_config): ++ runner = PleOffloadRunner(vllm_config) ++ ++ zmq_context = zmq.Context() ++ pull_socket = zmq_context.socket(zmq.PULL) ++ pull_socket.bind(ipc_addr) ++ owns_ipc_socket = True ++ logger.info( ++ "Bound IPC address %s; waiting for %d GPU worker registration(s).", ++ ipc_addr, ++ num_workers, ++ ) ++ ++ # READY means that the process can immediately serve requests. Wait ++ # for every DP/TP worker to register before notifying the parent. ++ runner.accept_registrations(pull_socket, num_workers) ++ # Checked: the host-mapped pull design must not create a context. ++ logger.info( ++ "PLE offload host-mapped pull: CUDA initialized in this process = %s", ++ torch.cuda.is_initialized(), ++ ) ++ ready_writer.send( ++ { ++ "status": PleOffloadWorker.READY_STR, ++ "layer_names": sorted(runner.layer_names), ++ } ++ ) ++ ready_writer.close() ++ ready_writer = None # type: ignore[assignment] ++ ++ runner.busy_loop(pull_socket, shutdown_event) ++ except Exception as error: ++ logger.exception("Unexpected failure in PLE offload worker.") ++ if ready_writer is not None: ++ with contextlib.suppress(Exception): ++ ready_writer.send( ++ {"status": "FAILURE", "error": repr(error)[:4096]} ++ ) ++ raise ++ finally: ++ if pull_socket is not None: ++ pull_socket.close(linger=0) ++ if zmq_context is not None: ++ zmq_context.term() ++ if owns_ipc_socket: ++ _unlink_ipc_socket(ipc_addr) ++ if ready_writer is not None: ++ ready_writer.close() ++ ++ ++class PleOffloadRunner: ++ """Own all discovered PLE tables and serve every local DP rank.""" ++ ++ def __init__(self, vllm_config: VllmConfig) -> None: ++ self.vllm_config = vllm_config ++ self._clamp_input_ids = ( ++ getattr(vllm_config, "speculative_config", None) is not None ++ ) ++ # name -> PleOffloadLayer (CPU) ++ self._layers: dict[str, PleOffloadLayer] = {} ++ # dp_rank -> layer_name -> one destination per TP rank ++ self._worker_targets: dict[int, dict[str, list[PleOffloadOutputTarget]]] = {} ++ # Each (dp_rank, layer_name) pair owns a plain CPU scratch buffer that ++ # forward_impl writes; the rows are then copied into every TP worker's ++ # shared host buffer. Nothing in this process is pinned or on a GPU. ++ self._scratch_bufs: dict[int, dict[str, torch.Tensor]] = {} ++ self._shutdown_event: threading.Event | None = None ++ self._step = 0 # number of handled request batches (diagnostics) ++ self._request_index = 0 ++ # Shared-memory inputs are registered once per DP rank by TP rank zero. ++ self._input_bufs: dict[int, PleOffloadInputBuffers] = {} ++ self._load_weights() ++ ++ @property ++ def layer_names(self) -> list[str]: ++ """Return PleOffloadLayer names in model traversal order.""" ++ return list(self._layers) ++ ++ def _load_weights(self) -> None: ++ """Load only :class:`PleOffloadLayer` subtrees into CPU memory. ++ ++ Strategy: ++ 1. Build the entire model on ``meta`` so non-offloaded parameters use no ++ physical memory. PleOffloadLayer constructors explicitly target CPU. ++ 2. Discover all PleOffloadLayer modules from the complete model. ++ 3. Stream the checkpoint through a prefix filter so only matching PLE ++ tensors are materialized and passed to ``model.load_weights``. ++ 4. Run post-load processing only on the CPU-owned PLE subtrees. ++ """ ++ model_config = self.vllm_config.model_config ++ load_config = self.vllm_config.load_config ++ ++ # Step 1: build complete structure, while only PLE subtrees allocate CPU ++ # memory. All transformer, MoE, and vision parameters remain on meta. ++ logger.info("Initializing model structure for PLE weight discovery ...") ++ model_dtype = cast(torch.dtype, model_config.dtype) ++ with ( ++ set_default_torch_dtype(model_dtype), ++ torch.device("meta"), ++ _platform_device_is_meta(), ++ ): ++ model = initialize_model( ++ vllm_config=self.vllm_config, ++ model_config=model_config, ++ ) ++ ++ # Step 2: preserve named_modules DFS order so CPU execution follows the ++ # same layer order as the GPU model forward. ++ offload_layers = { ++ name: module ++ for name, module in model.named_modules() ++ if isinstance(module, PleOffloadLayer) ++ } ++ if not offload_layers: ++ raise RuntimeError( ++ "VLLM_PLE_CPU_OFFLOAD is enabled, but no PleOffloadLayer " ++ "was found in the initialized model" ++ ) ++ logger.info( ++ "Found %d PleOffloadLayer(s): %s", ++ len(offload_layers), ++ sorted(offload_layers), ++ ) ++ offload_prefixes = tuple(f"{name}." for name in offload_layers) ++ ++ # Step 3: filter checkpoint tensors before model.load_weights(). The ++ # conditional-generation checkpoint uses HF names such as ++ # ``model.language_model.*`` while named_modules exposes mapped vLLM ++ # names such as ``language_model.model.*``. Apply the model mapper only ++ # for matching, then yield the original pair so load_weights performs ++ # its normal single mapping pass. ++ mapper = getattr(model, "hf_to_vllm_mapper", None) ++ matched_checkpoint_tensors = 0 ++ ++ def offload_only_iter( ++ weights: Iterable[tuple[str, torch.Tensor]], ++ ) -> Iterable[tuple[str, torch.Tensor]]: ++ nonlocal matched_checkpoint_tensors ++ for weight_name, tensor in weights: ++ mapped_name: str | None = weight_name ++ if mapper is not None: ++ mapped_names = mapper.apply_list([weight_name]) ++ mapped_name = mapped_names[0] if mapped_names else None ++ if mapped_name is not None and mapped_name.startswith(offload_prefixes): ++ matched_checkpoint_tensors += 1 ++ yield weight_name, tensor ++ ++ loader = get_model_loader(load_config) ++ if isinstance(loader, DummyModelLoader): ++ logger.info( ++ "Initializing dummy weights for %d PleOffloadLayer(s) ...", ++ len(offload_layers), ++ ) ++ for layer in offload_layers.values(): ++ initialize_dummy_weights(layer, model_config) ++ elif isinstance(loader, DefaultModelLoader): ++ all_weights = loader.get_all_weights(model_config, model) ++ loaded_params = model.load_weights(offload_only_iter(all_weights)) ++ if matched_checkpoint_tensors == 0: ++ raise RuntimeError( ++ "PLE offload checkpoint filter matched no weights for " ++ f"layers: {sorted(offload_layers)}" ++ ) ++ ++ expected_offload_params = { ++ f"{layer_name}.{param_name}" ++ for layer_name, layer in offload_layers.items() ++ for param_name, _ in layer.named_parameters() ++ } ++ loaded_offload_entries = { ++ name for name in loaded_params if name.startswith(offload_prefixes) ++ } ++ loaded_expected_params = expected_offload_params.intersection(loaded_params) ++ missing_offload_params = sorted( ++ expected_offload_params.difference(loaded_expected_params) ++ ) ++ if missing_offload_params: ++ raise RuntimeError( ++ "PLE offload checkpoint did not load all materialized " ++ f"parameters: {missing_offload_params}" ++ ) ++ logger.info( ++ "PLE offload matched %d checkpoint tensor(s), loaded %d " ++ "offload entries, and verified %d/%d materialized " ++ "parameter(s) for layers: %s", ++ matched_checkpoint_tensors, ++ len(loaded_offload_entries), ++ len(loaded_expected_params), ++ len(expected_offload_params), ++ sorted(offload_layers), ++ ) ++ else: ++ raise NotImplementedError( ++ "PLE offload requires the default or dummy model loader, got " ++ f"{type(loader).__name__}" ++ ) ++ ++ # Step 4: post-load processing is restricted to CPU-owned PLE modules; ++ # the remainder of the model is still on meta and must not be visited. ++ for layer in offload_layers.values(): ++ process_weights_after_loading(layer, model_config, torch.device("cpu")) ++ ++ self._layers.update(offload_layers) ++ del model ++ logger.info("PLE weight loading complete.") ++ ++ def accept_registrations( ++ self, ++ pull_socket: zmq.Socket, ++ num_workers: int, ++ ) -> None: ++ """Receive every local DP/TP worker's IPC and shared-memory buffers.""" ++ logger.info("Waiting for %d GPU worker registration(s) ...", num_workers) ++ registrations: list[PleOffloadRegistration] = [] ++ for index in range(num_workers): ++ item = pickle.loads(pull_socket.recv()) ++ if not isinstance(item, PleOffloadRegistration): ++ raise RuntimeError( ++ "Expected PleOffloadRegistration during setup, got " ++ f"{type(item).__name__} ({index + 1}/{num_workers})" ++ ) ++ registrations.append(item) ++ logger.info( ++ "GPU worker %d registered (dp_rank=%d, tp_rank=%d, layers=%s).", ++ item.worker_id, ++ item.dp_rank, ++ item.tp_rank, ++ sorted(item.host_output_buffers), ++ ) ++ ++ dp_size = self.vllm_config.parallel_config.data_parallel_size ++ tp_size = self.vllm_config.parallel_config.tensor_parallel_size ++ if num_workers != dp_size * tp_size: ++ raise RuntimeError( ++ f"Expected {dp_size * tp_size} registrations for DP={dp_size}, " ++ f"TP={tp_size}, got {num_workers}" ++ ) ++ ++ registrations_by_dp: dict[int, list[PleOffloadRegistration]] = {} ++ for registration in registrations: ++ registrations_by_dp.setdefault(registration.dp_rank, []).append( ++ registration ++ ) ++ if set(registrations_by_dp) != set(range(dp_size)): ++ raise RuntimeError( ++ f"Expected DP ranks {set(range(dp_size))}, " ++ f"got {set(registrations_by_dp)}" ++ ) ++ for dp_rank, dp_registrations in registrations_by_dp.items(): ++ tp_ranks = {registration.tp_rank for registration in dp_registrations} ++ if tp_ranks != set(range(tp_size)): ++ raise RuntimeError( ++ f"DP rank {dp_rank} expected TP ranks {set(range(tp_size))}, " ++ f"got {tp_ranks}" ++ ) ++ ++ for registration in registrations: ++ if set(registration.host_output_buffers) != set(self.layer_names) or set( ++ registration.host_flag_tensors ++ ) != set(self.layer_names): ++ raise RuntimeError( ++ "Registered PLE layers do not match CPU layers: " ++ f"registered={sorted(registration.host_output_buffers)} / " ++ f"{sorted(registration.host_flag_tensors)}, " ++ f"cpu={sorted(self.layer_names)}" ++ ) ++ targets_for_dp = self._worker_targets.setdefault(registration.dp_rank, {}) ++ for layer_name, host_buffer in registration.host_output_buffers.items(): ++ host_flag = registration.host_flag_tensors[layer_name] ++ for tensor in (host_buffer, host_flag): ++ if tensor.device.type != "cpu" or not tensor.is_shared(): ++ raise RuntimeError( ++ f"PLE host buffer for {layer_name} (dp=" ++ f"{registration.dp_rank}, tp={registration.tp_rank}) " ++ "is not a shared CPU tensor" ++ ) ++ if host_flag.dtype != torch.int32 or host_flag.numel() < 1: ++ raise RuntimeError(f"PLE host flag for {layer_name} must be int32") ++ target = PleOffloadOutputTarget( ++ tp_rank=registration.tp_rank, ++ host_output_buffer=host_buffer, ++ host_flag=host_flag, ++ ) ++ targets_for_dp.setdefault(layer_name, []).append(target) ++ # All TP ranks in one DP group receive the same input, so buffers ++ # registered by TP rank zero are sufficient for that DP rank. ++ if registration.tp_rank == 0: ++ self._input_bufs[registration.dp_rank] = PleOffloadInputBuffers( ++ input_ids_buf=registration.input_ids_buf, ++ query_start_loc_buf=registration.query_start_loc_buf, ++ ngram_context_buf=registration.ngram_context_buf, ++ ) ++ ++ if set(self._input_bufs) != set(range(dp_size)): ++ raise RuntimeError( ++ "TP rank zero did not register PLE input buffers for every DP " ++ f"rank: expected={set(range(dp_size))}, got={set(self._input_bufs)}" ++ ) ++ ++ config = self.vllm_config.model_config.hf_text_config ++ max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens ++ for dp_rank, layer_targets in self._worker_targets.items(): ++ self._scratch_bufs[dp_rank] = {} ++ for layer_name, targets in layer_targets.items(): ++ if len(targets) != tp_size: ++ raise RuntimeError( ++ f"PLE layer {layer_name} for DP rank {dp_rank} received " ++ f"{len(targets)} targets, expected {tp_size}" ++ ) ++ targets.sort(key=lambda target: target.tp_rank) ++ out_dim = self._layers[layer_name].get_offload_output_dim( ++ int(config.ple_embed_dim) ++ ) ++ out_dtype = self._layers[layer_name].get_offload_output_dtype( ++ self.vllm_config.model_config.dtype ++ ) ++ for target in targets: ++ shape = tuple(target.host_output_buffer.shape) ++ if shape != (max_tokens, out_dim) or ( ++ target.host_output_buffer.dtype != out_dtype ++ ): ++ raise RuntimeError( ++ f"PLE host buffer for {layer_name} (dp={dp_rank}, " ++ f"tp={target.tp_rank}) is {shape} " ++ f"{target.host_output_buffer.dtype}, expected " ++ f"{(max_tokens, out_dim)} {out_dtype}" ++ ) ++ self._scratch_bufs[dp_rank][layer_name] = torch.empty( ++ max_tokens, out_dim, dtype=out_dtype, device="cpu" ++ ) ++ logger.info( ++ "Registrations complete (dp_size=%d, tp_size=%d, layers=%s).", ++ dp_size, ++ tp_size, ++ sorted(self.layer_names), ++ ) ++ ++ @torch.inference_mode() ++ def busy_loop( ++ self, ++ pull_socket: zmq.Socket, ++ shutdown_event: threading.Event, ++ ) -> None: ++ """Decode and batch available requests by DP rank until shutdown.""" ++ logger.info("Busy-loop started.") ++ self._shutdown_event = shutdown_event ++ poller = zmq.Poller() ++ poller.register(pull_socket, zmq.POLLIN) ++ while not shutdown_event.is_set(): ++ if pull_socket not in dict(poller.poll(timeout=100)): ++ continue ++ ++ requests = [] ++ try: ++ requests.append(_receive_offload_request(pull_socket)) ++ while True: ++ requests.append( ++ _receive_offload_request(pull_socket, zmq.NOBLOCK) ++ ) ++ except zmq.Again: ++ pass ++ except (msgspec.DecodeError, RuntimeError) as error: ++ raise RuntimeError("Unexpected PLE offload request") from error ++ ++ self._handle_requests(requests) ++ ++ def _handle_requests( ++ self, requests: list[_ReceivedPleOffloadRequest] ++ ) -> None: ++ """Run requests layer-first so each DP rank can resume promptly.""" ++ with record_function_or_nullcontext("ple_offload: handle_requests"): ++ self._handle_requests_impl(requests) ++ ++ def _validate_requests_before_publication( ++ self, requests: list[PleOffloadRequest] ++ ) -> None: ++ """Validate a drained request batch before computation or publication.""" ++ for request in requests: ++ dp_rank = request.dp_rank ++ if dp_rank not in self._worker_targets or dp_rank not in self._input_bufs: ++ raise RuntimeError( ++ f"No PLE buffers or output targets for dp_rank={dp_rank}; " ++ "refusing request" ++ ) ++ ++ input_bufs = self._input_bufs[dp_rank] ++ if not 0 < request.num_tokens <= input_bufs.input_ids_buf.numel(): ++ raise RuntimeError( ++ f"PLE num_tokens={request.num_tokens} is outside capacity " ++ f"{input_bufs.input_ids_buf.numel()} for dp_rank={dp_rank}" ++ ) ++ max_reqs = input_bufs.query_start_loc_buf.numel() - 1 ++ if not 0 < request.num_reqs <= max_reqs: ++ raise RuntimeError( ++ f"PLE num_reqs={request.num_reqs} is outside capacity " ++ f"{max_reqs} for dp_rank={dp_rank}" ++ ) ++ if ( ++ input_bufs.query_start_loc_buf.dtype != torch.int32 ++ or input_bufs.query_start_loc_buf.dim() != 1 ++ ): ++ raise RuntimeError( ++ f"PLE query_start_loc for dp_rank={dp_rank} must be 1-D int32" ++ ) ++ if ( ++ input_bufs.ngram_context_buf is not None ++ and request.num_reqs > input_bufs.ngram_context_buf.shape[0] ++ ): ++ raise RuntimeError( ++ f"PLE num_reqs={request.num_reqs} exceeds ngram context " ++ f"capacity for dp_rank={dp_rank}" ++ ) ++ ++ query_start_loc = input_bufs.query_start_loc_buf[ ++ : request.num_reqs + 1 ++ ].tolist() ++ if ( ++ query_start_loc[0] != 0 ++ or query_start_loc[-1] > request.num_tokens ++ or any(a > b for a, b in zip(query_start_loc, query_start_loc[1:])) ++ ): ++ raise RuntimeError( ++ f"Invalid PLE query_start_loc for dp_rank={dp_rank}: " ++ f"{query_start_loc}, num_tokens={request.num_tokens}" ++ ) ++ ++ for layer_name, targets in self._worker_targets[dp_rank].items(): ++ scratch = self._scratch_bufs[dp_rank][layer_name] ++ if request.num_tokens > scratch.shape[0] or any( ++ request.num_tokens > target.host_output_buffer.shape[0] ++ for target in targets ++ ): ++ raise RuntimeError( ++ f"PLE num_tokens={request.num_tokens} exceeds output " ++ f"capacity for {layer_name}, dp_rank={dp_rank}" ++ ) ++ ++ def _handle_requests_impl( ++ self, ++ requests: list[_ReceivedPleOffloadRequest | PleOffloadRequest], ++ ) -> None: ++ """Process one batch of PLE requests.""" ++ received = [ ++ item ++ if isinstance(item, _ReceivedPleOffloadRequest) ++ else _ReceivedPleOffloadRequest(item, False, None, None) ++ for item in requests ++ ] ++ plain_requests = [item.request for item in received] ++ self._validate_requests_before_publication(plain_requests) ++ received_by_dp: dict[int, _ReceivedPleOffloadRequest] = {} ++ for item in received: ++ request = item.request ++ if request.dp_rank not in self._worker_targets: ++ raise RuntimeError( ++ f"No PLE output targets for dp_rank={request.dp_rank}; " ++ "refusing request" ++ ) ++ if request.dp_rank in received_by_dp: ++ # One input slot and one output slot per DP rank: a second ++ # request in the same batch would alias them. Fatal, never ++ # silently dropped. ++ raise RuntimeError( ++ f"Two PLE requests for dp_rank={request.dp_rank} in one batch " ++ f"(step {self._step}); the single-slot protocol is violated" ++ ) ++ received_by_dp[request.dp_rank] = item ++ self._step += 1 ++ if _FAULT is not None: ++ for item in received_by_dp.values(): ++ if not item.serving: ++ continue ++ self._request_index += 1 ++ transaction = f"request:{self._request_index}" ++ if item.transaction_id is not None: ++ transaction = f"{transaction}@{item.transaction_id}" ++ for cut in ("park_before_publish", "park_before_publish_ack"): ++ if _FAULT is not None and _FAULT.fires( ++ cut, ++ role="offload", ++ unit="offload-request", ++ eligible=True, ++ dp_rank=item.request.dp_rank, ++ transaction_id=transaction, ++ correlation_id=item.correlation_id, ++ ): ++ # No multiprocessing or fault-state lock is held here. ++ threading.Event().wait() ++ ++ requests_by_dp = { ++ dp_rank: item.request ++ for dp_rank, item in received_by_dp.items() ++ } ++ ++ # Speculative placeholders are not vocabulary IDs. Normalize each DP ++ # input once before all PLE layers consume the shared buffer. ++ if self._clamp_input_ids: ++ for dp_rank, request in requests_by_dp.items(): ++ self._input_bufs[dp_rank].input_ids_buf[ ++ : request.num_tokens ++ ].clamp_min_(0) ++ ++ for layer_name, layer in self._layers.items(): ++ for dp_rank, request in requests_by_dp.items(): ++ targets = self._worker_targets[dp_rank][layer_name] ++ ++ input_bufs = self._input_bufs[dp_rank] ++ ngram_context = ( ++ input_bufs.ngram_context_buf[: request.num_reqs] ++ if input_bufs.ngram_context_buf is not None ++ else None ++ ) ++ result = layer.forward_impl( ++ input_bufs.input_ids_buf[: request.num_tokens], ++ input_bufs.input_ids_buf[: request.num_tokens], ++ input_bufs.query_start_loc_buf[: request.num_reqs + 1], ++ ngram_context, ++ output_buffer=self._scratch_bufs[dp_rank][layer_name], ++ ) ++ if _rowdump.enabled(): ++ _seq = getattr(self, "_rowdump_seq", 0) ++ self._rowdump_seq = _seq + 1 ++ _rowdump.write(f"offload-dp{dp_rank}", 1, dp_rank, _seq, ++ input_bufs.input_ids_buf[: request.num_tokens], result) ++ ++ # The result is identical on every TP rank in this DP group. ++ # A worker's host buffer may be overwritten only after its GPU ++ # released the previous result (``release_outputs`` stores ++ # flag=0 on the model stream after the forward, which is also ++ # after the in-graph DMA read it). Then write rows, then flag=1. ++ scratch = self._scratch_bufs[dp_rank][layer_name] ++ if result.dim() != 2 or result.shape[0] != request.num_tokens or ( ++ result.shape[1] != scratch.shape[1] ++ ) or result.dtype != scratch.dtype or result.device.type != "cpu": ++ raise RuntimeError( ++ f"PLE result {tuple(result.shape)} {result.dtype} " ++ f"{result.device} for {layer_name} does not " ++ f"match the request ({request.num_tokens} rows x " ++ f"{scratch.shape[1]} {scratch.dtype} CPU); refusing to " ++ f"publish (step {self._step})" ++ ) ++ slices = tuple(slice(0, size) for size in result.shape) ++ for target in targets: ++ self._wait_flag_released(target, layer_name, dp_rank) ++ target.host_output_buffer[slices].copy_(result[slices]) ++ _publish_fence() ++ target.host_flag[0] = 1 ++ ++ def _wait_flag_released( ++ self, target: PleOffloadOutputTarget, layer_name: str, dp_rank: int ++ ) -> None: ++ """Spin until the GPU reset this target's flag (previous result used).""" ++ flag = target.host_flag ++ if int(flag[0]) == 0: ++ return ++ spins = 0 ++ warned = False ++ start = time.monotonic() ++ while int(flag[0]) != 0: ++ spins += 1 ++ if self._shutdown_event is not None and self._shutdown_event.is_set(): ++ raise RuntimeError("PLE offload shutting down while waiting for GPU") ++ if spins > 200: ++ time.sleep(0.0001) ++ waited = time.monotonic() - start ++ if not warned and waited > 5.0: ++ warned = True ++ logger.warning( ++ "PLE offload waited >5 s for the GPU to release %s " ++ "(dp=%d, tp=%d, step %d); flag=%d", ++ layer_name, ++ dp_rank, ++ target.tp_rank, ++ self._step, ++ int(flag[0]), ++ ) ++ if waited > _PLE_WAIT_TIMEOUT_S: ++ raise RuntimeError( ++ f"PLE offload: GPU did not release {layer_name} (dp={dp_rank}, " ++ f"tp={target.tp_rank}, step {self._step}) within " ++ f"{_PLE_WAIT_TIMEOUT_S:.0f} s; refusing to overwrite" ++ ) +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 9657c28..0e132b6 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -1,9 +1,10 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++import contextlib + import gc + import itertools + from collections import defaultdict +-from collections.abc import Callable, Iterable ++from collections.abc import Callable, Iterable, Iterator + from contextlib import ExitStack + from dataclasses import dataclass, replace + from itertools import groupby, product +@@ -30,6 +31,9 @@ from vllm.distributed.parallel_state import ( + ) + from vllm.forward_context import BatchDescriptor, set_forward_context + from vllm.logger import init_logger ++from vllm.model_executor.layers.ple_offload_layer import ( ++ observe_ple_copy_geometry, ++) + from vllm.model_executor.offloader.base import get_offloader + from vllm.platforms import current_platform + from vllm.sequence import IntermediateTensors +@@ -170,6 +174,14 @@ class CudaGraphManager: + self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map() + + self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} ++ # FULL graphs bake every PLE wait/copy into their native graph. Metadata ++ # comes from the actual custom-op implementation during native capture. ++ self._ple_full_copy_geometry: dict[ ++ BatchExecutionDescriptor, tuple[tuple[str, int, int], ...] ++ ] = {} ++ self._expected_ple_copy_layers: tuple[str, ...] = () ++ # Only the target-model manager is allowed to capture PLE operations. ++ self._allow_ple_full_graph = False + self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None + + self._graphs_captured = False +@@ -469,15 +481,16 @@ class CudaGraphManager: + if self._capture_mem_samples is not None: + torch.accelerator.synchronize() + free_before = torch.accelerator.get_memory_info()[0] +- with torch.cuda.graph( +- graph, self.pool, stream=self._capture_stream(desc) +- ): +- forward_fn(CUDAGraphMode.NONE) +- # Join offloader's copy stream after forward to avoid +- # unjoined stream error. The last layer's start_prefetch +- # forks copy_stream, but wait_prefetch only happens in +- # the next forward pass. +- get_offloader().join_after_forward() ++ with self.capture_ple_copy_geometry(desc): ++ with torch.cuda.graph( ++ graph, self.pool, stream=self._capture_stream(desc) ++ ): ++ forward_fn(CUDAGraphMode.NONE) ++ # Join offloader's copy stream after forward to avoid ++ # unjoined stream error. The last layer's start_prefetch ++ # forks copy_stream, but wait_prefetch only happens in ++ # the next forward pass. ++ get_offloader().join_after_forward() + if self._capture_mem_samples is not None: + torch.accelerator.synchronize() + free_after = torch.accelerator.get_memory_info()[0] +@@ -493,6 +506,78 @@ class CudaGraphManager: + {desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0} + ) + ++ @contextlib.contextmanager ++ def capture_ple_copy_geometry( ++ self, desc: BatchExecutionDescriptor ++ ) -> Iterator[None]: ++ """Production capture callback for actual PLE custom-op geometry.""" ++ with observe_ple_copy_geometry() as geometry: ++ yield ++ ++ if geometry and not self._allow_ple_full_graph: ++ raise RuntimeError( ++ "PLE wait-copy reached an unsupported FULL graph path; " ++ "speculator/draft FULL graphs may not contain PLE operations" ++ ) ++ if geometry: ++ self.record_ple_copy_geometry(desc, geometry) ++ elif self._expected_ple_copy_layers: ++ raise RuntimeError( ++ f"PLE copy geometry was not observed while capturing FULL graph {desc}" ++ ) ++ ++ def record_ple_copy_geometry( ++ self, ++ desc: BatchExecutionDescriptor, ++ geometry: Iterable[tuple[str, int, int]], ++ ) -> None: ++ if desc.cg_mode != CUDAGraphMode.FULL: ++ raise ValueError("PLE capture geometry is recorded only for FULL graphs") ++ observed = tuple( ++ (str(layer), int(rows), int(nbytes)) ++ for layer, rows, nbytes in geometry ++ ) ++ if not observed: ++ raise RuntimeError(f"PLE FULL graph {desc} observed no wait-copy operations") ++ observed_layers = tuple(layer for layer, _, _ in observed) ++ expected_layers = self._expected_ple_copy_layers ++ if expected_layers and sorted(observed_layers) != sorted(expected_layers): ++ raise RuntimeError( ++ f"PLE FULL graph {desc} observed layers {observed_layers}, " ++ f"expected {expected_layers}" ++ ) ++ inconsistent = [ ++ (layer, rows, nbytes) ++ for layer, rows, nbytes in observed ++ if rows != desc.num_tokens or nbytes < 0 ++ ] ++ if inconsistent: ++ raise RuntimeError( ++ f"PLE FULL graph {desc} actual PLE copy geometry " ++ f"{inconsistent} does not match descriptor rows={desc.num_tokens}" ++ ) ++ previous = self._ple_full_copy_geometry.get(desc) ++ if previous is not None and previous != observed: ++ raise RuntimeError( ++ f"PLE FULL graph {desc} changed copy geometry from " ++ f"{previous} to {observed}" ++ ) ++ self._ple_full_copy_geometry[desc] = observed ++ ++ def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int: ++ try: ++ geometry = self._ple_full_copy_geometry[desc] ++ except KeyError as error: ++ raise RuntimeError( ++ f"PLE copy geometry was not recorded for FULL graph {desc}" ++ ) from error ++ rows = {num_rows for _, num_rows, _ in geometry} ++ if len(rows) != 1: ++ raise RuntimeError( ++ f"PLE FULL graph {desc} has inconsistent per-layer rows: {geometry}" ++ ) ++ return rows.pop() ++ + def dispatch( + self, + num_reqs: int, +@@ -532,6 +617,11 @@ class CudaGraphManager: + f"Expected FULL mode, got {desc.cg_mode}" + ) + assert desc in self.graphs, f"No cudagraph for {desc}" ++ if ( ++ desc in self._ple_full_copy_geometry ++ and not self._allow_ple_full_graph ++ ): ++ raise RuntimeError("refusing unsupported PLE speculator FULL replay") + # Sync offloader before replay - needed when transitioning from + # eager/piecewise to full cudagraph (e.g., prefill → decode). + # The previous eager iteration's start_prefetch may have queued +@@ -577,11 +667,32 @@ class ModelCudaGraphManager(CudaGraphManager): + varlen_decode=varlen_decode, + ubatch_runner=ubatch_runner, + ) ++ self._allow_ple_full_graph = True + self.hidden_states: torch.Tensor | None = None + self.aux_hidden_states: list[torch.Tensor] = [] + self.use_aux_hidden_state_outputs = False + self.intermediate_tensors: IntermediateTensors | None = None + ++ def _prepare_ple_capture( ++ self, model: nn.Module, *, initialize_breakable: bool ++ ) -> None: ++ expected_ple_copy_layers = tuple( ++ str(getattr(module, "_ple_offload_name", name)) ++ for name, module in model.named_modules() ++ if getattr(module, "_is_cpu_offloaded", False) ++ and getattr(module, "_host_output_ptr", None) is not None ++ ) ++ if ( ++ "" in expected_ple_copy_layers ++ or len(set(expected_ple_copy_layers)) != len(expected_ple_copy_layers) ++ ): ++ raise RuntimeError( ++ "PLE FULL capture requires distinct, bound connector layer identities" ++ ) ++ self._expected_ple_copy_layers = expected_ple_copy_layers ++ if initialize_breakable: ++ self.init_breakable_cg_runner(model) ++ + def capture( + self, + model: nn.Module, +@@ -599,8 +710,9 @@ class ModelCudaGraphManager(CudaGraphManager): + ) -> None: + """Capture CUDA graphs for model forward pass.""" + self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs +- if self.use_breakable_cg: +- self.init_breakable_cg_runner(model) ++ self._prepare_ple_capture( ++ model, initialize_breakable=self.use_breakable_cg ++ ) + + if self.cudagraph_mode.has_piecewise_cudagraphs() and not ( + self.use_breakable_cg or has_compiled_submodule(model) +@@ -901,7 +1013,9 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: + if manager.use_breakable_cg: + # Create the breakable runner before the wrapper pool swap so + # its pool is covered as well. +- manager.init_breakable_cg_runner(runner.model) ++ manager._prepare_ple_capture( ++ runner.model, initialize_breakable=True ++ ) + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 744fe5b..e2ddc52 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -77,6 +77,9 @@ from vllm.v1.outputs import ( + ModelRunnerOutput, + RoutedExpertsTensors, + ) ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++from vllm.v1.ple_offload.fatal import fail_ple_execute_model ++from vllm.v1.ple_offload.test_fault import _FAULT + from vllm.v1.watermarking import create_watermarker + from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler + from vllm.v1.watermarking.spec_decode import ( +@@ -1689,6 +1692,45 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.ec_connector.no_forward(scheduler_output).ec_connector_output, + ) + ++ def _prepare_ple_forward( ++ self, ++ input_batch: InputBatch, ++ batch_desc: BatchExecutionDescriptor, ++ dummy_run: bool, ++ *, ++ serving: bool = False, ++ transaction_id: str | None = None, ++ correlation_id: str | None = None, ++ ) -> None: ++ connector = self._ple_offload_connector ++ assert connector is not None ++ request_num_tokens = int(input_batch.num_tokens_after_padding) ++ if batch_desc.cg_mode == CUDAGraphMode.FULL: ++ manager = self.cudagraph_manager ++ if manager is None: ++ raise RuntimeError("PLE FULL replay has no CUDA graph manager") ++ copy_num_tokens = manager.ple_copy_rows(batch_desc) ++ else: ++ # PIECEWISE/eager execute the Python op with this actual padded view. ++ copy_num_tokens = int(input_batch.input_ids.shape[0]) ++ ++ if request_num_tokens != copy_num_tokens: ++ raise RuntimeError( ++ "PLE request/selected-forward copy-count mismatch: " ++ f"request={request_num_tokens}, copy={copy_num_tokens}, " ++ f"mode={batch_desc.cg_mode}" ++ ) ++ connector.prepare_forward( ++ input_batch.num_reqs, ++ request_num_tokens, ++ dummy_run, ++ consumer_num_tokens=copy_num_tokens, ++ serving=serving, ++ transaction_id=transaction_id, ++ correlation_id=correlation_id, ++ ) ++ ++ @fail_ple_execute_model + @torch.inference_mode() + def execute_model( + self, +@@ -1956,6 +1998,54 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ) + self.step_timing.forward_start() + ++ # prepare_inputs has finalized MRV2's GPU buffers. Stage D2H on this ++ # stream before the PLE placeholder can wait for CPU output. ++ connector = self._ple_offload_connector ++ ple_serving = False ++ ple_sequence = None ++ ple_transaction = None ++ ple_correlation = None ++ if connector is not None: ++ ple_serving = not dummy_run and not is_profile ++ ple_sequence = connector.begin_forward( ++ dummy_run=dummy_run, ++ num_tokens=int(input_batch.num_tokens_after_padding), ++ ) ++ if _FAULT is not None and ple_serving: ++ ple_correlation = f"sequence:{ple_sequence}" ++ ple_transaction = connector.fault_transaction(ple_sequence) ++ if _FAULT is not None and _FAULT.fires( ++ "exc_before_prepare", ++ role="gpu", ++ unit="runner-forward", ++ eligible=ple_serving, ++ tp_rank=connector.tp_rank, ++ pp_rank=connector.pp_rank, ++ dp_rank=connector.dp_rank, ++ transaction_id=ple_transaction, ++ correlation_id=ple_correlation, ++ ): ++ raise RuntimeError("PLE test fault: exc_before_prepare") ++ self._prepare_ple_forward( ++ input_batch, ++ batch_desc, ++ dummy_run, ++ serving=ple_serving, ++ transaction_id=ple_transaction, ++ correlation_id=ple_correlation, ++ ) ++ if _FAULT is not None and _FAULT.fires( ++ "exc_after_prepare", ++ role="gpu", ++ unit="runner-forward", ++ eligible=ple_serving, ++ tp_rank=connector.tp_rank, ++ pp_rank=connector.pp_rank, ++ dp_rank=connector.dp_rank, ++ transaction_id=ple_transaction, ++ correlation_id=ple_correlation, ++ ): ++ raise RuntimeError("PLE test fault: exc_after_prepare") + connector_kwargs = dict( + scheduler_output=scheduler_output, + request_state_indices=input_batch.idx_mapping, +@@ -2013,6 +2103,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # Eager (NONE): call the raw model directly. + model_output = self.model(**model_inputs) + ++ # The model has consumed every PLE output, so the CPU worker may reuse ++ # the registered buffers for the next request. ++ if connector is not None: ++ connector.release_outputs( ++ serving=ple_serving, ++ transaction_id=ple_transaction, ++ correlation_id=ple_correlation, ++ ) ++ if _FAULT is not None and _FAULT.fires( ++ "exc_after_release", ++ role="gpu", ++ unit="runner-forward", ++ eligible=ple_serving, ++ tp_rank=connector.tp_rank, ++ pp_rank=connector.pp_rank, ++ dp_rank=connector.dp_rank, ++ transaction_id=ple_transaction, ++ correlation_id=ple_correlation, ++ ): ++ raise RuntimeError("PLE test fault: exc_after_release") + self.kv_connector.finish_forward() + + if self.is_last_pp_rank: +@@ -2047,6 +2157,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ec_connector_output=ec_connector_output, + routed_experts=routed_experts, + cudagraph_stats=cudagraph_stats, ++ ple_serving=ple_serving, ++ ple_sequence=ple_sequence, ++ ple_transaction=ple_transaction, + ) + + if not self.is_last_pp_rank: +@@ -2058,6 +2171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ) + return None + ++ @fail_ple_execute_model + @torch.inference_mode() + @step_eplb_after() + def sample_tokens( +@@ -2077,6 +2191,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ec_connector_output = self.execute_model_state.ec_connector_output + routed_experts = self.execute_model_state.routed_experts + cudagraph_stats = self.execute_model_state.cudagraph_stats ++ ple_serving = self.execute_model_state.ple_serving ++ ple_sequence = self.execute_model_state.ple_sequence ++ ple_transaction = self.execute_model_state.ple_transaction + self.execute_model_state = None + + if not self.is_last_pp_rank: +@@ -2084,6 +2201,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # IntermediateTensors instead of final hidden states. Receive the + # sampled tokens broadcast from the last rank and update local state. + assert self.pp_handler is not None ++ connector = self._ple_offload_connector ++ if ( ++ connector is not None ++ and _FAULT is not None ++ and ple_serving ++ and ple_sequence is not None ++ and connector.completion_is_pending(ple_sequence) ++ and _FAULT.fires( ++ "exc_async_output", ++ role="gpu", ++ unit="runner-output-relay", ++ eligible=True, ++ tp_rank=connector.tp_rank, ++ pp_rank=connector.pp_rank, ++ dp_rank=connector.dp_rank, ++ transaction_id=ple_transaction, ++ correlation_id=f"sequence:{ple_sequence}", ++ ) ++ ): ++ raise RuntimeError("PLE test fault: exc_async_output") + all_decode_next = self.pp_handler.receive(input_batch) + # Optimistically update num_computed_tokens for entire batch here. + # Will be adjusted for rejections if necessary in update_requests. +@@ -2379,6 +2516,9 @@ class ExecuteModelState(NamedTuple): + ec_connector_output: ECConnectorOutput | None + routed_experts: RoutedExpertsTensors | None + cudagraph_stats: CUDAGraphStat | None ++ ple_serving: bool ++ ple_sequence: int | None ++ ple_transaction: str | None + + + class BatchReqState(NamedTuple): +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 09f91f9..43bc705 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -87,6 +87,7 @@ from vllm.v1.outputs import ( + DraftTokenIds, + ModelRunnerOutput, + ) ++from vllm.v1.ple_offload.fatal import fail_ple_execute_model + from vllm.v1.utils import compute_iteration_details, report_usage_stats + from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel +@@ -231,6 +232,148 @@ class Worker(WorkerBase): + # Resolved lazily on first sleep/wake; persists worker-process state. + self._sleep_mode_backend: SleepModeBackend | None = None + ++ def _has_ple_layers(self) -> bool: ++ """Return whether this model configuration constructs PLE layers.""" ++ if not envs.VLLM_PLE_CPU_OFFLOAD: ++ return False ++ text_config = self.model_config.hf_text_config ++ ple_layer_ids = getattr(text_config, "ple_layer_ids", None) ++ if not ple_layer_ids: ++ return False ++ # PP (vllm#54709, carried from the earlier overlay): only the pipeline rank ++ # that owns a PLE layer takes part in the offload (rank 0 with the ++ # default partition). Other ranks build PPMissingLayer there and must ++ # skip the connector entirely. ++ pp_size = self.parallel_config.pipeline_parallel_size ++ if pp_size > 1: ++ from vllm.distributed.utils import get_pp_indices ++ ++ tp_size = self.parallel_config.tensor_parallel_size ++ pp_rank = (self.rank // tp_size) % pp_size ++ start, end = get_pp_indices( ++ int(text_config.num_hidden_layers), pp_rank, pp_size ++ ) ++ return any(start <= int(i) - 1 < end for i in ple_layer_ids) ++ return True ++ ++ def _validate_ple_offload_config(self) -> None: ++ """Reject unsupported PLE offload execution modes.""" ++ parallel_config = self.parallel_config ++ unsupported = [] ++ if not self.use_v2_model_runner: ++ unsupported.append("model runner V1") ++ if not current_platform.is_cuda(): ++ unsupported.append(f"device={current_platform.device_type}") ++ if parallel_config.nnodes != 1: ++ unsupported.append(f"nnodes={parallel_config.nnodes}") ++ if parallel_config.data_parallel_backend != "mp": ++ unsupported.append(f"DP backend={parallel_config.data_parallel_backend}") ++ if ( ++ parallel_config.data_parallel_size_local ++ != parallel_config.data_parallel_size ++ ): ++ unsupported.append( ++ "non-local DP " ++ f"({parallel_config.data_parallel_size_local}/" ++ f"{parallel_config.data_parallel_size} local ranks)" ++ ) ++ # PP allowed (vllm#54709): PLE layers must live on pipeline rank 0 ++ # (enforced placement-based in models/config.py); num_workers stays ++ # dp_size * tp_size = the ranks of the owning stage. ++ if parallel_config.prefill_context_parallel_size != 1: ++ unsupported.append(f"PCP={parallel_config.prefill_context_parallel_size}") ++ if parallel_config.decode_context_parallel_size != 1: ++ unsupported.append(f"DCP={parallel_config.decode_context_parallel_size}") ++ if parallel_config.use_ubatching: ++ unsupported.append("ubatching/DBO") ++ if self.model_config.architecture not in { ++ "Qwen4ExpForCausalLM", ++ "Qwen4ExpForConditionalGeneration", ++ }: ++ unsupported.append("architecture") ++ if self.vllm_config.weight_transfer_config is not None: ++ unsupported.append("weight transfer") ++ ++ if unsupported: ++ raise ValueError( ++ "VLLM_PLE_CPU_OFFLOAD does not support the requested " ++ "configuration. Unsupported settings: " ++ f"{', '.join(unsupported)}" ++ ) ++ ++ def spawn_ple_offload(self) -> None: ++ """Spawn one node-local PLE CPU worker from DP0/TP0.""" ++ if ( ++ not self._ple_offload_enabled ++ or self.rank != 0 ++ or self.parallel_config.data_parallel_rank != 0 ++ ): ++ return ++ ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ ipc_addr = self.parallel_config._ple_offload_ipc_path ++ if not ipc_addr: ++ raise RuntimeError("PLE offload IPC address was not initialized") ++ dp_size = self.parallel_config.data_parallel_size ++ tp_size = self.parallel_config.tensor_parallel_size ++ num_workers = dp_size * tp_size ++ logger.info( ++ "PleOffload: spawning worker " ++ "(rank=%d, local_rank=%d, dp_size=%d, tp_size=%d, " ++ "num_workers=%d, ipc_addr=%s).", ++ self.rank, ++ self.local_rank, ++ dp_size, ++ tp_size, ++ num_workers, ++ ipc_addr, ++ ) ++ self._ple_offload_worker_handle = PleOffloadWorker.make_process( ++ self.vllm_config, ++ num_workers, ++ ipc_addr, ++ ) ++ self._start_ple_offload_watchdog() ++ ++ def _start_ple_offload_watchdog(self) -> None: ++ """Fail this worker if the PLE offload process exits unexpectedly. ++ ++ Host-mapped pull: a dead producer would leave the captured forward ++ waiting on its host flag forever; turn that into a fast, visible ++ engine failure instead of a hang. ``shutdown`` sets the stop marker ++ before closing the handle, so a normal close is not a failure. ++ """ ++ import threading ++ ++ handle = self._ple_offload_worker_handle ++ self._ple_offload_stopping = False ++ ++ def watch() -> None: ++ returncode = handle.proc.wait() ++ if getattr(self, "_ple_offload_stopping", False): ++ return ++ logger.error( ++ "PLE offload process %s exited unexpectedly (exitcode %s); " ++ "failing this worker so the engine does not hang on a PLE wait", ++ handle.proc.pid, ++ returncode, ++ ) ++ os._exit(1) ++ ++ threading.Thread( ++ target=watch, name="ple-offload-watchdog", daemon=True ++ ).start() ++ ++ def wait_ple_offload_ready(self) -> None: ++ """Wait until weights and every TP registration are ready.""" ++ if self._ple_offload_worker_handle is None: ++ return ++ ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ PleOffloadWorker.wait_for_ready(self._ple_offload_worker_handle) ++ + @property + def sleep_mode_backend(self) -> "SleepModeBackend": + if self._sleep_mode_backend is None: +@@ -1152,6 +1295,7 @@ class Worker(WorkerBase): + ) -> ModelRunnerOutput | AsyncModelRunnerOutput: + return self.model_runner.sample_tokens(grammar_output) + ++ @fail_ple_execute_model + @torch.inference_mode() + @with_gpu_sync_check + def execute_model( +@@ -1487,6 +1631,7 @@ class Worker(WorkerBase): + self.model_runner.reset_lora_state() + + def shutdown(self) -> None: ++ self._ple_offload_stopping = True + gc.unfreeze() + + # Bounds-guard instrumentation: synchronous final read of the guard counters, so +diff --git a/vllm_ple_offload_bootstrap/__init__.py b/vllm_ple_offload_bootstrap/__init__.py +new file mode 100644 +index 0000000..60ddda4 +--- /dev/null ++++ b/vllm_ple_offload_bootstrap/__init__.py +@@ -0,0 +1,210 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Minimal PLE child bootstrap; intentionally imports no torch or vLLM.""" ++ ++import importlib.abc ++import importlib.machinery ++import os ++import pickle ++import struct ++import sys ++import threading ++from multiprocessing.connection import Connection ++from types import ModuleType ++from typing import Any ++ ++_FAULT = None ++if os.environ.get("VLLM_PLE_TEST_FAULT") is not None: ++ from vllm_ple_test_fault import ( ++ _FAULT as _FAULT, ++ _block_on_harness_gate, ++ ) ++ ++ ++def _send_partial_ready_frame(fd: int) -> tuple[int, int]: ++ """Write the complete header and an incomplete body in one syscall.""" ++ body = pickle.dumps( ++ {"status": "READY", "layer_names": []}, ++ protocol=pickle.HIGHEST_PROTOCOL, ++ ) ++ frame = struct.pack("!i", len(body)) + body ++ fragment = frame[: 4 + max(1, len(body) // 2)] ++ try: ++ return os.write(fd, fragment), len(frame) ++ except OSError: ++ return 0, len(frame) ++ ++ ++def _send_partial_ready_header(fd: int) -> tuple[int, int]: ++ """Write fewer than the four required header bytes in one syscall.""" ++ header = struct.pack("!i", 1) ++ try: ++ return os.write(fd, header[:2]), len(header) ++ except OSError: ++ return 0, len(header) ++ ++ ++def _read_exact(fd: int, size: int) -> bytes: ++ chunks = bytearray() ++ while len(chunks) < size: ++ chunk = os.read(fd, size - len(chunks)) ++ if not chunk: ++ raise EOFError("PLE launch payload was truncated") ++ chunks.extend(chunk) ++ return bytes(chunks) ++ ++ ++def _recv_launch_payload(fd: int) -> dict: ++ size = struct.unpack("!Q", _read_exact(fd, 8))[0] ++ payload = pickle.loads(_read_exact(fd, size)) ++ if not isinstance(payload, dict): ++ raise RuntimeError("invalid PLE bootstrap payload") ++ return payload ++ ++ ++def _fires( ++ cut: str, ++ *, ++ unit: str = "bootstrap-launch", ++ transaction: str = "launch:1", ++ bytes_written: int | None = None, ++ frame_size: int | None = None, ++) -> bool: ++ return _FAULT is not None and _FAULT.fires( ++ cut, ++ role="bootstrap", ++ unit=unit, ++ eligible=True, ++ transaction_id=transaction, ++ correlation_id="launch:1", ++ bytes_written=bytes_written, ++ frame_size=frame_size, ++ ) ++ ++ ++def _write_and_ack_partial( ++ fd: int, ++ cut: str, ++ sender: Any, ++) -> bool: ++ """A short or failed prefix write cannot acknowledge or enter the park.""" ++ if _FAULT is None or _FAULT.cut != cut: ++ # Disarmed or armed for another cut: the control stream is untouched. ++ return False ++ written, frame_size = sender(fd) ++ if written <= 0 or written >= frame_size: ++ return False ++ expected = 2 if cut == "launch_partial_ready_header" else None ++ if expected is not None and written != expected: ++ return False ++ if expected is None and written <= 4: ++ return False ++ return _fires( ++ cut, ++ bytes_written=written, ++ frame_size=frame_size, ++ ) ++ ++ ++class _FaultImportLoader(importlib.abc.Loader): ++ def __init__(self, delegate: importlib.abc.Loader) -> None: ++ self._delegate = delegate ++ ++ def create_module(self, spec): ++ create = getattr(self._delegate, "create_module", None) ++ return None if create is None else create(spec) ++ ++ def exec_module(self, module: ModuleType) -> None: ++ if _fires( ++ "launch_block_worker_import", ++ unit="bootstrap-import", ++ transaction="module:vllm.v1.ple_offload.worker", ++ ): ++ _block_on_harness_gate() ++ self._delegate.exec_module(module) ++ ++ ++class _FaultImportFinder(importlib.abc.MetaPathFinder): ++ _TARGET = "vllm.v1.ple_offload.worker" ++ ++ def find_spec(self, fullname: str, path, target=None): ++ if fullname != self._TARGET: ++ return None ++ spec = importlib.machinery.PathFinder.find_spec(fullname, path, target) ++ if spec is None or spec.loader is None: ++ raise ImportError(f"cannot find {fullname}") ++ if not isinstance(spec.loader, importlib.abc.Loader): ++ raise ImportError(f"{fullname} has no executable loader") ++ spec.loader = _FaultImportLoader(spec.loader) ++ return spec ++ ++ ++def _install_import_gate() -> None: ++ target = _FaultImportFinder._TARGET ++ if target in sys.modules: ++ raise RuntimeError("PLE worker import gate was a cache hit") ++ sys.meta_path.insert(0, _FaultImportFinder()) ++ ++ ++def main() -> None: ++ """Receive launch state before importing torch or vLLM.""" ++ if len(sys.argv) != 3: ++ raise SystemExit("expected inherited PLE control and parent-death fds") ++ control_fd = int(sys.argv[1]) ++ death_fd = int(sys.argv[2]) ++ if control_fd == death_fd: ++ raise RuntimeError("PLE control and parent-death fds must be distinct") ++ ++ os.set_inheritable(control_fd, False) ++ os.set_inheritable(death_fd, False) ++ death_pipe = Connection(death_fd) ++ ++ def monitor_parent() -> None: ++ try: ++ death_pipe.recv() ++ except (EOFError, OSError): ++ pass ++ os._exit(1) ++ ++ threading.Thread( ++ target=monitor_parent, ++ daemon=True, ++ name="PleOffloadBootstrapDeathMonitor", ++ ).start() ++ ++ if _fires("launch_park_before_payload_read"): ++ threading.Event().wait() ++ ++ # An armed reconstruction payload invokes its reducer inside this loads(). ++ payload = _recv_launch_payload(control_fd) ++ ++ if _write_and_ack_partial( ++ control_fd, ++ "launch_partial_ready", ++ _send_partial_ready_frame, ++ ): ++ threading.Event().wait() ++ ++ if _write_and_ack_partial( ++ control_fd, ++ "launch_partial_ready_header", ++ _send_partial_ready_header, ++ ): ++ threading.Event().wait() ++ ++ if _fires("launch_silent_ready"): ++ threading.Event().wait() ++ ++ ready_pipe = Connection(control_fd) ++ vllm_config = pickle.loads(payload["vllm_config_payload"]) ++ ++ if _FAULT is not None and _FAULT.cut == "launch_block_worker_import": ++ _install_import_gate() ++ from vllm.v1.ple_offload.worker import PleOffloadWorker ++ ++ PleOffloadWorker.proc_main( ++ vllm_config=vllm_config, ++ num_workers=int(payload["num_workers"]), ++ ipc_addr=str(payload["ipc_addr"]), ++ ready_pipe=ready_pipe, ++ ) +diff --git a/vllm_ple_offload_bootstrap/__main__.py b/vllm_ple_offload_bootstrap/__main__.py +new file mode 100644 +index 0000000..5b3ca57 +--- /dev/null ++++ b/vllm_ple_offload_bootstrap/__main__.py +@@ -0,0 +1,7 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Fresh-interpreter entry point for the PLE CPU-offload worker.""" ++ ++from vllm_ple_offload_bootstrap import main ++ ++main() +-- +2.43.0 + diff --git a/patches/0032-worker-boundary-allocator-telemetry-and-an-opt-in-PL.patch b/patches/0032-worker-boundary-allocator-telemetry-and-an-opt-in-PL.patch new file mode 100644 index 0000000..c937e48 --- /dev/null +++ b/patches/0032-worker-boundary-allocator-telemetry-and-an-opt-in-PL.patch @@ -0,0 +1,365 @@ +From 300e10fef5c7b4e9aeae0838ab086f5b7af315e9 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 00:01:58 -0400 +Subject: [PATCH 032/103] worker: boundary allocator telemetry and an opt-in + PLE row dump + +Env-gated diagnostics, inert unless set. Allocator telemetry is read only at +step boundaries, armed after warm-up and capture. +--- + tests/v1/worker/test_e2_memstats.py | 113 +++++++++++++++++++++++ + vllm/v1/ple_offload/rowdump.py | 41 +++++++++ + vllm/v1/ple_offload/worker.py | 1 + + vllm/v1/worker/e2_memstats.py | 136 ++++++++++++++++++++++++++++ + vllm/v1/worker/gpu_worker.py | 5 + + 5 files changed, 296 insertions(+) + create mode 100644 tests/v1/worker/test_e2_memstats.py + create mode 100644 vllm/v1/ple_offload/rowdump.py + create mode 100644 vllm/v1/worker/e2_memstats.py + +diff --git a/tests/v1/worker/test_e2_memstats.py b/tests/v1/worker/test_e2_memstats.py +new file mode 100644 +index 0000000..c405187 +--- /dev/null ++++ b/tests/v1/worker/test_e2_memstats.py +@@ -0,0 +1,113 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""CPU/GPU-free schema and call-order tests for vllm/v1/worker/e2_memstats.py (debug-only telemetry).""" ++from __future__ import annotations ++ ++import importlib.util ++import os ++import sys ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++import torch ++ ++MODULE_PATH = Path(__file__).resolve().parents[3] / "vllm" / "v1" / "worker" / "e2_memstats.py" ++HOST_CURRENT_KEYS = { ++ "active_bytes.current", ++ "active_requests.current", ++ "allocated_bytes.current", ++ "allocations.current", ++} ++ ++ ++def load_disabled_module(): ++ name = "e2_memstats_under_test" ++ sys.modules.pop(name, None) ++ spec = importlib.util.spec_from_file_location(name, MODULE_PATH) ++ if spec is None or spec.loader is None: ++ raise RuntimeError(f"cannot load {MODULE_PATH}") ++ module = importlib.util.module_from_spec(spec) ++ with mock.patch.dict(os.environ, {"VLLM_E2_MEMSTATS": "", "VLLM_E2_MEMSTATS_DIR": ""}): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++class MemstatsSchemaTest(unittest.TestCase): ++ def test_disabled_import_has_no_cuda_or_thread_side_effect(self): ++ initialized = torch.cuda.is_initialized() ++ module = load_disabled_module() ++ self.assertFalse(module.enabled()) ++ self.assertEqual(module._started, []) ++ self.assertEqual(torch.cuda.is_initialized(), initialized) ++ ++ def test_close_metrics_schema_and_call_order(self): ++ module = load_disabled_module() ++ initialized = torch.cuda.is_initialized() ++ calls: list[str] = [] ++ device_stats = { ++ "allocated_bytes.all.current": 11, ++ "reserved_bytes.all.current": 12, ++ "allocated_bytes.all.peak": 21, ++ "reserved_bytes.all.peak": 22, ++ "num_alloc_retries": 3, ++ "num_ooms": 0, ++ } ++ host_stats = { ++ "active_bytes.current": 31, ++ "active_requests.current": 32, ++ "allocated_bytes.current": 33, ++ "allocations.current": 34, ++ "allocated_bytes.peak": 99, ++ } ++ ++ def memory_stats(dev): ++ self.assertEqual(dev, 7) ++ calls.append("memory_stats") ++ return device_stats ++ ++ def reset_peak_memory_stats(dev): ++ self.assertEqual(dev, 7) ++ calls.append("reset_peak_memory_stats") ++ ++ def mem_get_info(dev): ++ self.assertEqual(dev, 7) ++ calls.append("mem_get_info") ++ return 41, 42 ++ ++ def host_memory_stats(): ++ calls.append("host_memory_stats") ++ return host_stats ++ ++ with ( ++ mock.patch.object(module.torch.cuda, "memory_stats", side_effect=memory_stats), ++ mock.patch.object(module.torch.cuda, "reset_peak_memory_stats", side_effect=reset_peak_memory_stats), ++ mock.patch.object(module.torch.cuda, "mem_get_info", side_effect=mem_get_info), ++ mock.patch.object(module.torch.cuda, "host_memory_stats", side_effect=host_memory_stats), ++ ): ++ record = module._close_metrics(7) ++ ++ # The device read and the peak reset stay adjacent; the host read comes after both. ++ self.assertEqual(calls, ["memory_stats", "reset_peak_memory_stats", "mem_get_info", "host_memory_stats"]) ++ self.assertEqual( ++ set(record), ++ { ++ "allocated_bytes.all.current", ++ "reserved_bytes.all.current", ++ "alloc_peak_w", ++ "res_peak_w", ++ "retries", ++ "ooms", ++ "free", ++ "total", ++ "host", ++ }, ++ ) ++ self.assertEqual(set(record["host"]), HOST_CURRENT_KEYS) ++ self.assertEqual(record["host"]["allocated_bytes.current"], 33) ++ self.assertEqual(record["allocated_bytes.all.current"], 11) ++ self.assertEqual(record["reserved_bytes.all.current"], 12) ++ self.assertEqual(torch.cuda.is_initialized(), initialized) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/vllm/v1/ple_offload/rowdump.py b/vllm/v1/ple_offload/rowdump.py +new file mode 100644 +index 0000000..95fd308 +--- /dev/null ++++ b/vllm/v1/ple_offload/rowdump.py +@@ -0,0 +1,41 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Debug-only PLE row dump for byte-level transport checks (off unless VLLM_PLE_ROW_DUMP=). ++ ++Each process appends framed records to /-pid.bin: ++ b"PLER" | kind u8 (0 = GPU worker: rows its forward consumed; 1 = offload process: rows it published) ++ | rank u8 | reserved u16 | seq u32 | n u32 | row_bytes u32 | ids int64[n] | rows uint8[n * row_bytes] ++""" ++ ++import os ++import struct ++import threading ++ ++import torch ++ ++_DIR = os.environ.get("VLLM_PLE_ROW_DUMP", "") ++_LOCK = threading.Lock() ++_FILES: dict = {} ++MAGIC = b"PLER" ++ ++ ++def enabled() -> bool: ++ return bool(_DIR) ++ ++ ++def write(tag: str, kind: int, rank: int, seq: int, ids: torch.Tensor, rows: torch.Tensor) -> None: ++ ids_cpu = ids.detach().to(device="cpu", dtype=torch.int64).contiguous() ++ rows_cpu = rows.detach().to(device="cpu").contiguous() ++ n = int(ids_cpu.shape[0]) ++ if rows_cpu.shape[0] != n: ++ raise RuntimeError(f"PLE row dump: {rows_cpu.shape[0]} rows for {n} ids") ++ raw = rows_cpu.view(torch.uint8).numpy().tobytes() if n else b"" ++ row_bytes = len(raw) // n if n else 0 ++ record = (MAGIC + struct.pack(" and ++VLLM_E2_MEMSTATS_DIR=. ++ ++WHY THIS IS NOT A SAMPLER ANY MORE. An earlier version read torch.cuda.memory_stats() and ++reset_peak_memory_stats() every VLLM_E2_MEMSTATS seconds. Both take the per-device caching-allocator mutex and the ++Python bindings do not release the GIL, so the old contract's claim that this thread "takes no lock shared with the ++model path" was FALSE. Three one-boot runs at one KV pin showed a monotone dose-response: sampling at 4 Hz wedged a ++stage-0 worker during a 131072-token row; at 1.1 Hz it survived that row and wedged during a 260000-token row; with the ++thread disabled the run completed every row. Causation is not established at n=1 per run ++(single runs only), but polling has no place inside a measured region, and it was ++never necessary: PyTorch maintains the allocator high-water marks itself, so reading them once per window is exact. ++ ++WHAT THIS THREAD DOES NOW. Between window boundaries it makes NO CUDA calls and touches NO allocator state: its tick ++is a plain read of /window.ctl, which holds " ". The gate driver crosses a boundary only while the ++server is quiet (no request running). On a new seq the thread performs these boundary-only allocator reads: ++ 1. memory_stats() -> the closing window's exact high-water marks and cumulative retry/OOM counters ++ 2. reset_peak_memory_stats() -> immediately after the read, with no file I/O in the gap ++ 3. mem_get_info() -> device free/total AT THAT QUIESCENT INSTANT ++ 4. host_memory_stats() -> pinned-host allocator current values, outside the read-to-reset gap ++then writes the "close" record for the old window and the "ack" record for the new one. ++ ++WHAT THIS MEASURES. ++ - A row's allocator peak is EXACT for that row's window (not a sampled maximum; it cannot miss a transient). ++ - Device free/total is observed ONLY at quiescent boundaries. No minimum-free figure exists for the inside of a ++ row, and none may be reported. ++ - Retry and OOM counters are cumulative and never reset; per-row deltas come from consecutive close records. ++ - There is NO interval-level peak. A qualifying interval is narrower than its row and is identified only ++ afterwards from the scheduler log, so the peak attributable to it is the ROW-level high-water mark, which ++ over-includes the prefill ramp. This loss of resolution is inherent to the design. ++ ++Records (one JSON per line): kind start | close | ack | error; rank, local, dev, pid; window, seq; and for close: ++allocated_bytes.all.current, reserved_bytes.all.current, alloc_peak_w, res_peak_w, retries, ooms, free, total, ++run_alloc_peak, run_res_peak, and a host object containing every host_memory_stats() key ending in ".current". ++Records are written with synchronous file I/O, never inside the read-to-reset gap. Nothing ++here raises into the worker. ++""" ++import json ++import os ++import threading ++import time ++ ++import torch ++ ++_SEG = float(os.environ.get("VLLM_E2_MEMSTATS", "0") or 0) ++_DIR = os.environ.get("VLLM_E2_MEMSTATS_DIR", "") ++_started = [] ++VERSION = "e2_memstats/5-boundary-current-host" ++ ++ ++def _close_metrics(dev): ++ """Read one quiescent boundary; kept separate for CPU-only testing.""" ++ s = torch.cuda.memory_stats(dev) ++ ap = int(s.get("allocated_bytes.all.peak", 0)) ++ rp = int(s.get("reserved_bytes.all.peak", 0)) ++ torch.cuda.reset_peak_memory_stats(dev) # no file I/O in the gap ++ free, total = torch.cuda.mem_get_info(dev) ++ host_stats = torch.cuda.host_memory_stats() ++ host = { ++ key: int(value) ++ for key, value in host_stats.items() ++ if key.endswith(".current") ++ } ++ return { ++ "allocated_bytes.all.current": int( ++ s.get("allocated_bytes.all.current", 0) ++ ), ++ "reserved_bytes.all.current": int( ++ s.get("reserved_bytes.all.current", 0) ++ ), ++ "alloc_peak_w": ap, ++ "res_peak_w": rp, ++ "retries": int(s.get("num_alloc_retries", 0)), ++ "ooms": int(s.get("num_ooms", 0)), ++ "free": int(free), ++ "total": int(total), ++ "host": host, ++ } ++ ++ ++def enabled() -> bool: ++ return _SEG > 0 and bool(_DIR) ++ ++ ++def ensure_started(worker, armed_by: str = "compile_or_warm_up_model") -> None: ++ if _started: ++ return ++ dev = torch.cuda.current_device() # captured on the worker's own thread: every call below passes it explicitly ++ t = threading.Thread(target=_loop, name="e2-memstats", daemon=True, ++ args=(dev, getattr(worker, "rank", -1), getattr(worker, "local_rank", -1), armed_by)) ++ _started.append(t) ++ t.start() ++ ++ ++def _ctl(): ++ try: ++ with open(os.path.join(_DIR, "window.ctl")) as f: ++ seq, name = f.read().split(None, 1) ++ return int(seq), name.strip() ++ except (OSError, ValueError): ++ return None ++ ++ ++def _loop(dev, rank, local, armed_by): ++ os.makedirs(_DIR, exist_ok=True) ++ f = open(os.path.join(_DIR, f"memstats-rank{rank}-pid{os.getpid()}.jsonl"), "a", buffering=1) ++ base = {"rank": rank, "local": local, "dev": dev, "pid": os.getpid()} ++ ++ def w(rec): ++ f.write(json.dumps({**base, **rec}, separators=(",", ":")) + "\n") ++ ++ run_ap = run_rp = 0 ++ seq, name = _ctl() or (0, "serve") ++ torch.cuda.reset_peak_memory_stats(dev) ++ w({"kind": "start", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "armed_by": armed_by, ++ "window": name, "seq": seq, "version": VERSION, "poll_s": _SEG}) ++ while True: ++ time.sleep(_SEG) ++ try: ++ c = _ctl() # file read only: no CUDA, no allocator lock, no GIL-held call ++ if c is None or c[0] == seq: ++ continue ++ # Boundary. The driver only crosses one while the server is quiet, so this is the single point at which ++ # this thread touches the allocator at all. ++ metrics = _close_metrics(dev) ++ ap = metrics["alloc_peak_w"] ++ rp = metrics["res_peak_w"] ++ run_ap, run_rp = max(run_ap, ap), max(run_rp, rp) ++ w({"kind": "close", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, ++ "seq": seq, **metrics, ++ "run_alloc_peak": run_ap, "run_res_peak": run_rp}) ++ seq, name = c ++ w({"kind": "ack", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, "seq": seq}) ++ except Exception as e: # noqa: BLE001 - telemetry must never raise into the worker ++ w({"kind": "error", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, ++ "seq": seq, "error": f"{type(e).__name__}: {e}"[:300]}) +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 43bc705..cafeb73 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -89,6 +89,7 @@ from vllm.v1.outputs import ( + ) + from vllm.v1.ple_offload.fatal import fail_ple_execute_model + from vllm.v1.utils import compute_iteration_details, report_usage_stats ++from vllm.v1.worker import e2_memstats as _e2_memstats + from vllm.v1.worker import guard_telemetry + from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel + from vllm.v1.worker.startup_plan import ( +@@ -1098,6 +1099,10 @@ class Worker(WorkerBase): + # intra-op parallelism. + set_torch_threads_for_runtime() + ++ # Allocator telemetry: armed only after warm-up and CUDA-graph capture (env-gated). ++ if _e2_memstats.enabled(): ++ _e2_memstats.ensure_started(self) ++ + return CompilationTimes( + language_model=self.compilation_config.compilation_time, + encoder=self.compilation_config.encoder_compilation_time, +-- +2.43.0 + diff --git a/patches/0033-ple-integrate-the-pull-transport-on-v0.30.0-with-ups.patch b/patches/0033-ple-integrate-the-pull-transport-on-v0.30.0-with-ups.patch new file mode 100644 index 0000000..242d84f --- /dev/null +++ b/patches/0033-ple-integrate-the-pull-transport-on-v0.30.0-with-ups.patch @@ -0,0 +1,1509 @@ +From d8155750abf53c12e2ec699d2d119bdc371419f1 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 00:15:23 -0400 +Subject: [PATCH 033/103] ple: integrate the pull transport on v0.30.0 with + upstream UVA refused + +- VLLM_E2_PLE_PULL_TRANSPORT replaces the fork's old switch name, which + v0.30.0 reuses for its UVA backend with the opposite default; strict parser; + compile/capture factor. +- The pull transport is reachable on v0.30.0 through a pull embedding class + that inherits upstream's without editing it; the upstream UVA backend is + provably off while the pull switch is on. +- resolve_ple_mode()/assert_ple_mode() log and assert the effective PLE mode + for target and drafter, and refuse a mixed configuration. +- Producer log lines for the PLE backend and weight dtype. +- CPU tests for mode resolution and the mixed-config refusal. +--- + tests/test_envs.py | 3 + + .../ple_offload/test_fault_injector_basics.py | 2 - + .../ple_offload/test_transport_bounded_io.py | 3 - + tests/v1/test_e2_ple_mode.py | 338 +++++++++++++++++ + vllm/config/engram.py | 9 +- + vllm/config/parallel.py | 10 +- + vllm/config/vllm.py | 5 + + vllm/envs.py | 31 ++ + .../layers/ple_offload_layer.py | 4 +- + vllm/models/qwen4_exp/nvidia/model.py | 3 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 3 +- + .../qwen4_exp/nvidia/ngram_embedding.py | 28 +- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 41 ++- + vllm/models/qwen4_exp/nvidia/ple_pull.py | 342 ++++++++++++++++++ + vllm/v1/e2_guards.py | 146 ++++++++ + vllm/v1/executor/multiproc_executor.py | 6 + + vllm/v1/executor/uniproc_executor.py | 4 + + vllm/v1/ple_offload/connector.py | 2 +- + vllm/v1/ple_offload/test_fault.py | 13 + + vllm/v1/ple_offload/worker.py | 2 +- + vllm/v1/worker/gpu/model_runner.py | 33 ++ + vllm/v1/worker/gpu_worker.py | 34 +- + 22 files changed, 1033 insertions(+), 29 deletions(-) + create mode 100644 tests/v1/test_e2_ple_mode.py + create mode 100644 vllm/models/qwen4_exp/nvidia/ple_pull.py + create mode 100644 vllm/v1/ple_offload/test_fault.py + +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 920229d..891954b 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -615,6 +615,9 @@ _FORK_HASHED = [ + # Capture-configuration factor read in C++ (topk.cu exact_medium_enabled); + # the Python entry is a raw passthrough, never a second parser. + ("MERLIN_TOPK_EXACT_MEDIUM", "0", "0"), ++ # PLE pull transport: selects a different PLE module class and captured ++ # forward, so it is hashed. Strict "0"/"1" parser (envs._e2_strict_switch). ++ ("VLLM_E2_PLE_PULL_TRANSPORT", "1", True), + ] + + _FORK_IGNORED = [ +diff --git a/tests/v1/ple_offload/test_fault_injector_basics.py b/tests/v1/ple_offload/test_fault_injector_basics.py +index 0a9fd23..a47796d 100644 +--- a/tests/v1/ple_offload/test_fault_injector_basics.py ++++ b/tests/v1/ple_offload/test_fault_injector_basics.py +@@ -6,8 +6,6 @@ import queue + import subprocess + import sys + import threading +-from collections import deque +-from pathlib import Path + from types import SimpleNamespace + from unittest.mock import Mock + +diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py +index fb37122..748fd46 100644 +--- a/tests/v1/ple_offload/test_transport_bounded_io.py ++++ b/tests/v1/ple_offload/test_transport_bounded_io.py +@@ -1,9 +1,6 @@ + # SPDX-License-Identifier: Apache-2.0 + + import multiprocessing +-import os +-import signal +-import struct + import subprocess + import sys + import time +diff --git a/tests/v1/test_e2_ple_mode.py b/tests/v1/test_e2_ple_mode.py +new file mode 100644 +index 0000000..4017771 +--- /dev/null ++++ b/tests/v1/test_e2_ple_mode.py +@@ -0,0 +1,338 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""PLE mode resolution and the guard's refusal of mixed configs. ++ ++CPU-only. Imports nothing beyond ``vllm.envs`` and ``vllm.v1.e2_guards`` (plus ++the standard library), so it runs under ``pytest --noconftest`` on a CPU-only machine ++without pyzmq or a CUDA build. Configs are stand-ins: the resolver reads only ++``engram_config.cpu_offload`` and ``model_config.hf_text_config.ple_layer_ids``. ++ ++Every refusal has a FAILING CONTROL: the same input with the one deciding ++factor flipped must pass, or a deliberately lax implementation must be shown ++to accept what the real one rejects -- otherwise the refusal test could pass ++because everything raises. ++""" ++ ++import re ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++ ++import vllm.envs as envs ++from vllm.v1 import e2_guards ++from vllm.v1.e2_guards import ( ++ PLE_MODE_DEVICE, ++ PLE_MODE_NONE, ++ PLE_MODE_PULL, ++ PLE_MODE_UVA, ++ PULL_TRANSPORT_ENV, ++ UPSTREAM_UVA_ENV, ++ assert_ple_mode, ++ check_upstream_uva_env_against_pull, ++ pull_transport_cpu_offload_default, ++ pull_transport_enabled, ++ resolve_ple_mode, ++ verify_engram_against_pull, ++) ++ ++REPO = Path(__file__).resolve().parents[2] ++ ++ ++def _clear_envs_cache() -> None: ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ ++ ++@pytest.fixture(autouse=True) ++def _clean_env(monkeypatch): ++ monkeypatch.delenv(PULL_TRANSPORT_ENV, raising=False) ++ monkeypatch.delenv(UPSTREAM_UVA_ENV, raising=False) ++ _clear_envs_cache() ++ yield ++ _clear_envs_cache() ++ ++ ++def _set(monkeypatch, name: str, value: str | None) -> None: ++ if value is None: ++ monkeypatch.delenv(name, raising=False) ++ else: ++ monkeypatch.setenv(name, value) ++ _clear_envs_cache() ++ ++ ++def _config(cpu_offload: bool | None, ple_layer_ids=(2,)): ++ return SimpleNamespace( ++ engram_config=( ++ None if cpu_offload is None else SimpleNamespace(cpu_offload=cpu_offload) ++ ), ++ model_config=SimpleNamespace( ++ hf_text_config=SimpleNamespace(ple_layer_ids=list(ple_layer_ids)) ++ ), ++ ) ++ ++ ++# ------------------------------------------------------ the strict switch ++ ++ ++@pytest.mark.parametrize( ++ ("raw", "expected"), [(None, False), ("0", False), ("1", True)] ++) ++def test_switch_accepts_exactly_unset_0_1(monkeypatch, raw, expected) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, raw) ++ assert pull_transport_enabled() is expected ++ ++ ++@pytest.mark.parametrize( ++ "raw", ["true", "True", "yes", "on", "", " 1", "1 ", "01", "2"] ++) ++def test_switch_rejects_every_other_spelling(monkeypatch, raw) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, raw) ++ with pytest.raises(ValueError, match=PULL_TRANSPORT_ENV): ++ pull_transport_enabled() ++ ++ ++def test_switch_rejection_control_a_lax_parser_accepts_true(monkeypatch) -> None: ++ """Failing control: the E2 base's parser (``.lower() in ("true", "1")``) ++ would have turned the transport ON for "true"; the registered one must not. ++ Without this, the rejection test above could pass on a parser that raises ++ for everything -- so "1" is re-checked through the same path.""" ++ ++ def e2_base_parser(raw: str) -> bool: ++ return raw.lower() in ("true", "1") ++ ++ assert e2_base_parser("true") is True ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "true") ++ with pytest.raises(ValueError): ++ envs._e2_strict_switch(PULL_TRANSPORT_ENV) ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "1") ++ assert envs._e2_strict_switch(PULL_TRANSPORT_ENV) is True ++ ++ ++def test_switch_is_registered_and_hashed(monkeypatch) -> None: ++ """Registered, live, and NOT in compile_factors' deny-list. ++ ++ compile_factors() itself is exercised by tests/test_envs.py's ++ _FORK_HASHED row; on a CPU-only machine without pyzmq an unrelated getter makes it ++ raise, so here the deny-list is read from source instead. ++ """ ++ assert PULL_TRANSPORT_ENV in envs.environment_variables ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ assert envs.environment_variables[PULL_TRANSPORT_ENV]() is True ++ src = _src("vllm/envs.py") ++ start = src.index(" ignored_factors: set[str] = {") ++ end = src.index("\n }\n", start) ++ assert PULL_TRANSPORT_ENV not in src[start:end] ++ # Control: a knob that IS ignored is found by the same slice. ++ assert '"VLLM_E2_ALLOW_BREAKABLE"' in src[start:end] ++ ++ ++# ------------------------------------------- the old name while pull is on ++ ++ ++@pytest.mark.parametrize("old", [None, "0"]) ++def test_old_name_unset_or_zero_is_accepted_with_pull_on(monkeypatch, old) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ _set(monkeypatch, UPSTREAM_UVA_ENV, old) ++ check_upstream_uva_env_against_pull() ++ ++ ++def test_old_name_one_with_pull_on_is_the_double_offload(monkeypatch) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ _set(monkeypatch, UPSTREAM_UVA_ENV, "1") ++ with pytest.raises(RuntimeError, match="both transports"): ++ check_upstream_uva_env_against_pull() ++ ++ ++@pytest.mark.parametrize("old", ["2", "true", "00", ""]) ++def test_old_name_other_values_with_pull_on_are_refused(monkeypatch, old) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ _set(monkeypatch, UPSTREAM_UVA_ENV, old) ++ with pytest.raises(ValueError, match=UPSTREAM_UVA_ENV): ++ check_upstream_uva_env_against_pull() ++ ++ ++@pytest.mark.parametrize("pull", [None, "0"]) ++def test_old_name_control_pull_off_leaves_it_to_upstream(monkeypatch, pull) -> None: ++ """Failing control: the SAME VLLM_PLE_CPU_OFFLOAD=1 passes with pull off.""" ++ _set(monkeypatch, PULL_TRANSPORT_ENV, pull) ++ _set(monkeypatch, UPSTREAM_UVA_ENV, "1") ++ check_upstream_uva_env_against_pull() ++ ++ ++# ------------------------------------------------ the Engram default ++ ++ ++@pytest.mark.parametrize("upstream_default", [True, False]) ++def test_default_is_upstream_verbatim_with_pull_off(monkeypatch, upstream_default): ++ assert pull_transport_cpu_offload_default(upstream_default) is upstream_default ++ ++ ++def test_default_is_false_with_pull_on_even_if_upstream_says_true(monkeypatch): ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ assert pull_transport_cpu_offload_default(True) is False ++ ++ ++def test_default_with_pull_on_still_refuses_the_old_name(monkeypatch): ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ _set(monkeypatch, UPSTREAM_UVA_ENV, "1") ++ with pytest.raises(RuntimeError): ++ pull_transport_cpu_offload_default(True) ++ ++ ++# ------------------------------------------------ mode resolution ++ ++ ++@pytest.mark.parametrize( ++ ("pull", "cpu_offload", "mode"), ++ [ ++ (None, None, PLE_MODE_NONE), ++ (None, True, PLE_MODE_UVA), ++ (None, False, PLE_MODE_DEVICE), ++ ("0", True, PLE_MODE_UVA), ++ ("1", False, PLE_MODE_PULL), ++ ("1", None, PLE_MODE_PULL), ++ ], ++) ++def test_mode_resolution_table(monkeypatch, pull, cpu_offload, mode) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, pull) ++ assert resolve_ple_mode(_config(cpu_offload)) == mode ++ ++ ++def test_mixed_config_is_refused(monkeypatch) -> None: ++ """Pull on + an explicit --engram-config cpu_offload=true: no fallback.""" ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ with pytest.raises(RuntimeError, match="both own the PLE table"): ++ resolve_ple_mode(_config(True)) ++ with pytest.raises(RuntimeError, match="cpu_offload=True"): ++ verify_engram_against_pull(SimpleNamespace(cpu_offload=True)) ++ ++ ++def test_mixed_config_control_same_engram_with_pull_off_is_uva(monkeypatch) -> None: ++ """Failing control: the identical EngramConfig resolves (to UVA) and ++ verifies cleanly once the one deciding factor, the switch, is off.""" ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "0") ++ assert resolve_ple_mode(_config(True)) == PLE_MODE_UVA ++ verify_engram_against_pull(SimpleNamespace(cpu_offload=True)) ++ ++ ++def test_mixed_config_control_pull_on_with_cpu_offload_false_is_pull(monkeypatch): ++ """Failing control on the other factor: pull on, cpu_offload flipped.""" ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ assert resolve_ple_mode(_config(False)) == PLE_MODE_PULL ++ verify_engram_against_pull(SimpleNamespace(cpu_offload=False)) ++ ++ ++# ------------------------------------------------ the guard line ++ ++ ++_GUARD_LINE = re.compile( ++ r"e2-guard\[(?P\w+)\]: ple_mode=(?P\w+) pull_transport=(?P

[01]) " ++ r"engram_cpu_offload=(?PTrue|False|None) VLLM_PLE_CPU_OFFLOAD=(?P\S+) " ++ r"ple_layer_ids=(?P\[[^\]]*\])\." ++) ++ ++ ++def _captured_guard_line(monkeypatch, config, role: str) -> str: ++ lines: list[str] = [] ++ monkeypatch.setattr( ++ e2_guards.logger, "info", lambda fmt, *args: lines.append(fmt % args) ++ ) ++ assert_ple_mode(config, role=role) ++ assert len(lines) == 1, lines ++ return lines[0] ++ ++ ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_guard_line_names_the_mode_for_both_roles(monkeypatch, role) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ ids = (2,) if role == "target" else () # the MTP drafter forces PLE off ++ line = _captured_guard_line(monkeypatch, _config(False, ids), role) ++ m = _GUARD_LINE.fullmatch(line) ++ assert m, line ++ assert m["role"] == role and m["mode"] == PLE_MODE_PULL and m["p"] == "1" ++ assert m["co"] == "False" and m["old"] == "" ++ ++ ++def test_guard_refuses_the_drafter_too(monkeypatch) -> None: ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ with pytest.raises(RuntimeError): ++ assert_ple_mode(_config(True, ()), role="draft") ++ ++ ++def test_guard_ignores_the_breakable_opt_out(monkeypatch) -> None: ++ """VLLM_E2_ALLOW_BREAKABLE=1 must not disarm the PLE refusal.""" ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "1") ++ monkeypatch.setenv(e2_guards.ALLOW_BREAKABLE_ENV, "1") ++ with pytest.raises(RuntimeError): ++ assert_ple_mode(_config(True)) ++ ++ ++# --------------------------------------- source-level wiring (no imports) ++ ++ ++def _src(rel: str) -> str: ++ return (REPO / rel).read_text(encoding="utf-8") ++ ++ ++def test_uva_constructor_and_selector_refuse_the_pull_switch() -> None: ++ """ngram_embedding.py cannot be imported CPU-only; check the two ++ refusal sites exist in source, before the UVA table is allocated.""" ++ src = _src("vllm/models/qwen4_exp/nvidia/ngram_embedding.py") ++ init = src.index("class Qwen4ExpPLEPinnedHostEmbedding") ++ refuse = src.index("if envs.VLLM_E2_PLE_PULL_TRANSPORT:", init) ++ alloc = src.index("get_accelerator_view_from_cpu_tensor(self.weight)", init) ++ assert refuse < alloc ++ assert "if use_uva and envs.VLLM_E2_PLE_PULL_TRANSPORT:" in src ++ ++ ++def test_no_fork_code_reads_the_old_name() -> None: ++ offenders = [] ++ for rel in ( ++ "vllm/model_executor/layers/ple_offload_layer.py", ++ "vllm/v1/ple_offload/worker.py", ++ "vllm/v1/ple_offload/connector.py", ++ "vllm/v1/worker/gpu_worker.py", ++ "vllm/v1/executor/multiproc_executor.py", ++ "vllm/v1/executor/uniproc_executor.py", ++ "vllm/config/parallel.py", ++ "vllm/models/qwen4_exp/nvidia/ple_layer.py", ++ "vllm/models/qwen4_exp/nvidia/ple_pull.py", ++ ): ++ if "VLLM_PLE_CPU_OFFLOAD" in _src(rel): ++ offenders.append(rel) ++ assert not offenders ++ ++ ++_SMOKE_PLE = re.compile( ++ r"Initialized PLE embedding (?P\S+?): quantization_method=(?P" ++ r"[A-Za-z0-9_]+), weight_dtype=(?P[A-Za-z0-9_.]+), weight_device=(?P" ++ r"[A-Za-z0-9_:]+), pinned=(?PTrue|False)" ++) # the startup smoke's parser, copied verbatim ++ ++ ++@pytest.mark.parametrize( ++ "rel", ++ [ ++ "vllm/models/qwen4_exp/nvidia/ple_pull.py", ++ "vllm/models/qwen4_exp/nvidia/ngram_embedding.py", ++ ], ++) ++def test_producer_line_format_matches_the_smoke(rel) -> None: ++ src = _src(rel) ++ fmt = ( ++ "Initialized PLE embedding %s: quantization_method=%s, " ++ "weight_dtype=%s, weight_device=%s, pinned=%s, backend=%s" ++ ) ++ joined = re.sub(r'"\s*\n\s*"', "", src) ++ assert fmt in joined ++ line = fmt % ( ++ "model.layers.1.ple.ple_embedding.ngram_embedding", ++ "Qwen4ExpPLEFp8EmbeddingMethod", ++ "torch.float8_e4m3fn", ++ "cpu", ++ False, ++ "E2PullTransport", ++ ) ++ m = _SMOKE_PLE.search(line) ++ assert m and m["dtype"] == "torch.float8_e4m3fn" and m["pin"] == "False" +diff --git a/vllm/config/engram.py b/vllm/config/engram.py +index 5abb2bd..56b5df1 100644 +--- a/vllm/config/engram.py ++++ b/vllm/config/engram.py +@@ -24,7 +24,14 @@ _NGRAM_LAYER_FIELDS = { + + + def _default_cpu_offload() -> bool: +- return envs.VLLM_PLE_CPU_OFFLOAD ++ # Fork: with VLLM_E2_PLE_PULL_TRANSPORT=1 the default is False and a ++ # VLLM_PLE_CPU_OFFLOAD value that would enable UVA fails the boot. With the ++ # pull transport off, the upstream default below is returned unchanged. ++ from vllm.v1.e2_guards import pull_transport_cpu_offload_default ++ ++ if not envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ return envs.VLLM_PLE_CPU_OFFLOAD ++ return pull_transport_cpu_offload_default(False) + + + def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: +diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py +index 0213e06..7e8b9c4 100644 +--- a/vllm/config/parallel.py ++++ b/vllm/config/parallel.py +@@ -17,7 +17,7 @@ from vllm.config.fault_tolerance import FaultToleranceConfig + from vllm.config.utils import config + from vllm.logger import init_logger + from vllm.platforms import current_platform +-from vllm.utils.network_utils import get_open_ports_list ++from vllm.utils.network_utils import get_open_ports_list, get_open_zmq_ipc_path + + if TYPE_CHECKING: + from ray.runtime_env import RuntimeEnv +@@ -350,6 +350,10 @@ class ParallelConfig: + """Port of the coordination TCPStore. Can be set by the API server; workers + connect as clients to exchange self-picked group ports at runtime.""" + ++ _ple_offload_ipc_path: str = "" ++ """Node-local ZMQ IPC address for the fork's PLE pull-transport offload ++ process. Set only when VLLM_E2_PLE_PULL_TRANSPORT=1.""" ++ + decode_context_parallel_size: int = Field(default=1, ge=1) + """Number of ranks that shard the decode KV cache. DCP does not expand + the process world size. Without PCP, DCP reuses TP ranks. With PCP, DCP +@@ -498,6 +502,9 @@ class ParallelConfig: + "The FT system assumes one AsyncMPClient manages all engines." + ) + ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT and not self._ple_offload_ipc_path: ++ self._ple_offload_ipc_path = get_open_zmq_ipc_path() ++ + if self.all2all_backend in ["pplx", "naive"]: + logger.warning( + "The '%s' all2all backend has been removed. " +@@ -855,6 +862,7 @@ class ParallelConfig: + "worker_extension_cls", + "_api_process_count", + "_api_process_rank", ++ "_ple_offload_ipc_path", + # NUMA binding is per-rank host-side memory locality; it does + # not affect collective-communication semantics. When numa_bind + # is enabled with auto-detection, each DP rank stores its own +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index 8e08469..bc9fbbe 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -1249,6 +1249,11 @@ class VllmConfig: + self.engram_config.verify_model_config(model_config) + self.engram_config.verify_parallel_config(self.parallel_config) + self.engram_config.verify_load_config(self.load_config) ++ # Fork: an explicit cpu_offload=True beside the pull transport is ++ # a double offload; refused here, before any worker starts. ++ from vllm.v1.e2_guards import verify_engram_against_pull ++ ++ verify_engram_against_pull(self.engram_config) + logger.info_once("Resolved Engram configuration: %s", str(self.engram_config)) + + def __post_init__(self): +diff --git a/vllm/envs.py b/vllm/envs.py +index 734a95c..35047dd 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -376,6 +376,18 @@ if TYPE_CHECKING: + # to decide behaviour. Registered with the medium-path exact restart + # so the cache-factor contract is complete. + MERLIN_TOPK_EXACT_MEDIUM: str | None = None ++ # The fork's PLE pull transport -- the host-mapped, ++ # fail-closed offload process in vllm/v1/ple_offload/. Renamed from the E2 ++ # base's VLLM_PLE_CPU_OFFLOAD, which vLLM v0.30.0 reuses for a DIFFERENT ++ # subsystem with the OPPOSITE default (upstream: pinned-host UVA lookup, ++ # default ON). The two subsystems no longer share a switch. ++ # Strict: unset or "0" = off, "1" = on, anything else is a boot error. ++ # On: the PLE table lives in the offload process, the GPU placeholder waits ++ # on a host flag and DMAs rows inside the captured forward, and upstream's ++ # UVA backend is refused (vllm/v1/e2_guards.py). Off: the tree is upstream. ++ # AOT: yes (different module class and forward). CAPTURE: yes (the captured ++ # graph contains the wait/copy ops instead of a table lookup). Hashed. ++ VLLM_E2_PLE_PULL_TRANSPORT: bool = False + # + # RUNTIME-ONLY (registered for validate_environ + a single cached read): + # Static K/V scale sidecar: fill_()ed into existing fp32 device buffers +@@ -768,6 +780,21 @@ def _resolve_rust_cli_path() -> str | None: + return raw + + ++def _e2_strict_switch(name: str) -> bool: ++ """Parse a fork on/off switch that accepts exactly ``"0"`` and ``"1"``. ++ ++ Unset means off. Anything else -- ``"true"``, ``"yes"``, ``""``, ``" 1"``, ++ ``"01"`` -- is a boot error: a switch that selects a different transport ++ must not be satisfiable by a spelling nobody meant. ++ """ ++ raw = os.environ.get(name) ++ if raw is None or raw == "0": ++ return False ++ if raw == "1": ++ return True ++ raise ValueError(f'{name} must be exactly "0" or "1" (or unset), got {raw!r}') ++ ++ + def _e2_diagnostic_mode(name: str) -> str: + """Normalize an E2 diagnostic gate to off|warn|raise. + +@@ -2247,6 +2274,10 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_QSA_KV_COLLECT": lambda: os.getenv("VLLM_QSA_KV_COLLECT", "").strip(), + # Capture-configuration factor, read in C++; raw passthrough on purpose. + "MERLIN_TOPK_EXACT_MEDIUM": lambda: os.getenv("MERLIN_TOPK_EXACT_MEDIUM", None), ++ # PLE pull transport switch; strict "0"/"1" (see the TYPE_CHECKING block). ++ "VLLM_E2_PLE_PULL_TRANSPORT": lambda: _e2_strict_switch( ++ "VLLM_E2_PLE_PULL_TRANSPORT" ++ ), + # Runtime only; all of the following are in compile_factors.ignored_factors. + "VLLM_E2_ALLOW_BREAKABLE": lambda: os.getenv("VLLM_E2_ALLOW_BREAKABLE", "") == "1", + "VLLM_E2_EXPECTED_CUDAGRAPH_MODE": lambda: os.getenv( +diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py +index 1b67578..d1f6683 100644 +--- a/vllm/model_executor/layers/ple_offload_layer.py ++++ b/vllm/model_executor/layers/ple_offload_layer.py +@@ -346,7 +346,7 @@ class PleOffloadLayer(nn.Module, ABC): + def guarded_init( + self: "PleOffloadLayer", *args: object, **kwargs: object + ) -> None: +- if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT and not is_offload_process(): + nn.Module.__init__(self) + return + original_init(self, *args, **kwargs) +@@ -356,7 +356,7 @@ class PleOffloadLayer(nn.Module, ABC): + @classmethod + def get_target_device(cls) -> torch.device: + """Return CPU for the offload process and the active GPU otherwise.""" +- if envs.VLLM_PLE_CPU_OFFLOAD: ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: + return torch.device("cpu") + return torch.device("cuda", torch.accelerator.current_device_index()) + +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 8c29268..3e18f1c 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -78,7 +78,7 @@ from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +-from vllm.v1.e2_guards import assert_graph_profile_is_executable ++from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode + from vllm.v1.kv_cache_interface import MambaSpec + + from ..config import Qwen4ExpConfig +@@ -794,6 +794,7 @@ class Qwen4ExpForCausalLM( + # identity nobody chose. Qwen4ExpForConditionalGeneration constructs + # this class; Qwen4ExpMTP does not, and guards itself. + assert_graph_profile_is_executable(vllm_config, role="target") ++ assert_ple_mode(vllm_config, role="target") + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 468fc3a..e7d1255 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -49,7 +49,7 @@ from vllm.sequence import IntermediateTensors + from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) +-from vllm.v1.e2_guards import assert_graph_profile_is_executable ++from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode + + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +@@ -400,6 +400,7 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + # this class builds Qwen4ExpMultiTokenPredictor directly, from a config + # rebuilt at mtp.py:159 -- so its graph profile is checked here. + assert_graph_profile_is_executable(vllm_config, role="draft") ++ assert_ple_mode(vllm_config, role="draft") + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + cache_config = vllm_config.cache_config +diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +index 84c4e7c..573d1e3 100644 +--- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py ++++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +@@ -10,6 +10,7 @@ import torch + import torch.nn.functional as F + from torch import nn + ++import vllm.envs as envs + from vllm.compilation.breakable_cudagraph import eager_break_during_capture + from vllm.config import get_current_vllm_config + from vllm.distributed import get_dp_group, get_etp_group, get_tp_group +@@ -402,6 +403,13 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): + max_total_tokens: int = 0, + data_parallel_rank: int = 0, + ) -> None: ++ # Fork: never build the UVA table beside the E2 pull transport -- ++ # that would be a second offload of the same PLE table. ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ raise RuntimeError( ++ "Qwen4ExpPLEPinnedHostEmbedding (upstream UVA PLE backend) " ++ "refused: VLLM_E2_PLE_PULL_TRANSPORT=1 owns the PLE table" ++ ) + if not is_uva_available(): + raise RuntimeError("Engram CPU offload requires UVA support") + super().__init__( +@@ -700,10 +708,15 @@ class Qwen4ExpNGramEmbedding(nn.Module): + if params_dtype is None: + params_dtype = torch.get_default_dtype() + engram_config = get_current_vllm_config().engram_config ++ use_uva = engram_config is not None and engram_config.cpu_offload ++ if use_uva and envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ # Fork: no silent fallback between transports. ++ raise RuntimeError( ++ "PLE backend selection: cpu_offload=True (upstream UVA) with " ++ "VLLM_E2_PLE_PULL_TRANSPORT=1 -- refusing a double offload" ++ ) + embedding_cls = ( +- Qwen4ExpPLEPinnedHostEmbedding +- if engram_config is not None and engram_config.cpu_offload +- else Qwen4ExpPLEDeviceEmbedding ++ Qwen4ExpPLEPinnedHostEmbedding if use_uva else Qwen4ExpPLEDeviceEmbedding + ) + self.ngram_embedding = embedding_cls( + padded_vocab_size, +@@ -717,14 +730,21 @@ class Qwen4ExpNGramEmbedding(nn.Module): + data_parallel_rank=data_parallel_rank, + ) + weight = self.ngram_embedding.weight ++ # Fork: backend= names the class (or, inside the pull transport's ++ # offload process, the table it serves), so a smoke need not derive ++ # the backend from pinned=. Fields before it are upstream's, unchanged. ++ backend = type(self.ngram_embedding).__name__ ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ backend = f"E2PullTransport.offload_table:{backend}" + logger.info( + "Initialized PLE embedding %s: quantization_method=%s, " +- "weight_dtype=%s, weight_device=%s, pinned=%s", ++ "weight_dtype=%s, weight_device=%s, pinned=%s, backend=%s", + embedding_prefix, + type(embedding_quant_method).__name__, + weight.dtype, + weight.device, + weight.is_pinned(), ++ backend, + ) + + @staticmethod +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index 2ec8790..d50781b 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -7,6 +7,7 @@ from collections.abc import Sequence + import torch + from torch import nn + ++import vllm.envs as envs + from vllm.compilation.breakable_cudagraph import eager_break_during_capture + from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config + from vllm.forward_context import get_forward_context +@@ -93,16 +94,32 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.activation = "silu" +- self.ple_embedding = Qwen4ExpNGramEmbedding( +- config, +- int(config.ple_embed_dim), +- self.ple_dense_layer_id, +- vllm_config.scheduler_config.max_num_batched_tokens, +- data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, +- prefix=f"{prefix}.ple_embedding", +- quant_config=quant_config, +- params_dtype=model_config.dtype, +- ) ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ # Fork: the E2 pull transport owns the table in a CPU-only ++ # offload process; upstream's device/UVA backends are never built. ++ from .ple_pull import Qwen4ExpNGramEmbeddingPull ++ ++ self.ple_embedding = Qwen4ExpNGramEmbeddingPull.build( ++ config, ++ int(config.ple_embed_dim), ++ self.ple_dense_layer_id, ++ vllm_config.scheduler_config.max_num_batched_tokens, ++ data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, ++ prefix=f"{prefix}.ple_embedding", ++ quant_config=quant_config, ++ params_dtype=model_config.dtype, ++ ) ++ else: ++ self.ple_embedding = Qwen4ExpNGramEmbedding( ++ config, ++ int(config.ple_embed_dim), ++ self.ple_dense_layer_id, ++ vllm_config.scheduler_config.max_num_batched_tokens, ++ data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, ++ prefix=f"{prefix}.ple_embedding", ++ quant_config=quant_config, ++ params_dtype=model_config.dtype, ++ ) + # The PLE cache is TP-replicated, so this merged projection is too. + self.kv_proj = MergedColumnParallelLinear( + int(config.ple_embed_dim), +@@ -147,6 +164,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + ) -> torch.Tensor: + """Dequantize PLE lookup output.""" + ++ dequantize_pulled = getattr(self.ple_embedding, "dequantize_pulled", None) ++ if dequantize_pulled is not None: ++ # Fork pull transport: no local table; FP8 scale kept here. ++ return dequantize_pulled(embeddings, output_dtype) + return self.ple_embedding.ngram_embedding.dequantize( + embeddings, + output_dtype, +diff --git a/vllm/models/qwen4_exp/nvidia/ple_pull.py b/vllm/models/qwen4_exp/nvidia/ple_pull.py +new file mode 100644 +index 0000000..f52000c +--- /dev/null ++++ b/vllm/models/qwen4_exp/nvidia/ple_pull.py +@@ -0,0 +1,342 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The fork's PLE pull transport, re-seated on vLLM v0.30.0. ++ ++v0.30.0 split the E2 base's offload-aware ``Qwen4ExpNGramEmbedding`` into an ++upstream-only module (``ngram_embedding.py``, device or pinned-host UVA table). ++This module puts the E2 pull transport back WITHOUT editing that upstream ++class: ``Qwen4ExpNGramEmbeddingPull`` inherits it and ``PleOffloadLayer``. ++ ++* In a GPU worker, ``PleOffloadLayer.__init_subclass__`` skips the constructor: ++ the module is an empty placeholder with no table. ``PleOffloadConnector`` ++ binds it to a host-mapped flag/buffer; its forward is the captured ++ ``ple_offload_wait_copy`` op. Only the FP8 global scale is kept on the GPU. ++* In the CPU-only offload process the constructor runs in full, on CPU. The ++ table is upstream's ``Qwen4ExpPLEDeviceEmbedding`` allocated on CPU; ++ ``forward_impl`` computes n-gram IDs and gathers raw (FP8) rows. ++ ++Selected only when ``VLLM_E2_PLE_PULL_TRANSPORT=1``; with the switch off this ++module is never imported by the model and the tree is upstream. ++ ++Supported table formats: FP8 (global scale) and unquantized -- the two ++``Qwen4ExpPLEEmbeddingMethod`` families v0.30.0 itself supports. The E2 NVFP4 ++path is not carried (v0.30.0 has no NVFP4 PLE method to select). ++""" ++ ++from collections.abc import Iterable ++ ++import torch ++ ++from vllm.config import get_current_vllm_config ++from vllm.logger import init_logger ++from vllm.model_executor.layers.ple_offload_layer import ( ++ PleOffloadLayer, ++ is_offload_process, ++) ++from vllm.model_executor.layers.quantization.base_config import QuantizationConfig ++from vllm.transformers_utils.configs.qwen4_exp import Qwen4ExpTextConfig ++ ++from .ngram_embedding import ( ++ Qwen4ExpNGramEmbedding, ++ Qwen4ExpPLEEmbeddingMethod, ++ Qwen4ExpPLEFp8EmbeddingMethod, ++ Qwen4ExpPLEUnquantizedEmbeddingMethod, ++) ++ ++logger = init_logger(__name__) ++ ++#: The backend name this module reports in the PLE producer log line. ++PULL_BACKEND_NAME = "E2PullTransport" ++ ++ ++class Qwen4ExpNGramEmbeddingPull(PleOffloadLayer, Qwen4ExpNGramEmbedding): ++ """Upstream n-gram embedding whose table lives in the offload process.""" ++ ++ #: GPU placeholder: the method the offload process will use for the table, ++ #: resolved with upstream's own selector (set by ``build``). ++ _offload_quant_method: Qwen4ExpPLEEmbeddingMethod | None = None ++ ++ def __init__( ++ self, ++ config: Qwen4ExpTextConfig, ++ embedding_dim: int, ++ ple_dense_layer_id: int, ++ max_total_tokens: int, ++ *, ++ data_parallel_rank: int, ++ prefix: str, ++ quant_config: QuantizationConfig | None = None, ++ params_dtype: torch.dtype | None = None, ++ ) -> None: ++ # Runs only in the offload process (guarded by PleOffloadLayer). ++ engram_config = get_current_vllm_config().engram_config ++ if engram_config is not None and engram_config.cpu_offload: ++ raise RuntimeError( ++ "PLE pull transport: the offload process resolved " ++ "cpu_offload=True; refusing to build upstream's UVA table here." ++ ) ++ Qwen4ExpNGramEmbedding.__init__( ++ self, ++ config, ++ embedding_dim, ++ ple_dense_layer_id, ++ max_total_tokens, ++ data_parallel_rank=data_parallel_rank, ++ prefix=prefix, ++ quant_config=quant_config, ++ params_dtype=params_dtype, ++ ) ++ ++ @classmethod ++ def build( ++ cls, ++ config: Qwen4ExpTextConfig, ++ embedding_dim: int, ++ ple_dense_layer_id: int, ++ max_total_tokens: int, ++ *, ++ data_parallel_rank: int, ++ prefix: str, ++ quant_config: QuantizationConfig | None, ++ params_dtype: torch.dtype | None, ++ ) -> "Qwen4ExpNGramEmbeddingPull": ++ """Construct the placeholder (GPU) or the full CPU module (offload).""" ++ # The offload process builds the surrounding model on meta; this ++ # subtree must own real CPU storage. In a GPU worker the constructor is ++ # skipped, so the device context allocates nothing. ++ with torch.device(cls.get_target_device()): ++ module = cls( ++ config, ++ embedding_dim, ++ ple_dense_layer_id, ++ max_total_tokens, ++ data_parallel_rank=data_parallel_rank, ++ prefix=prefix, ++ quant_config=quant_config, ++ params_dtype=params_dtype, ++ ) ++ embedding_prefix = f"{prefix}.ngram_embedding" ++ method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( ++ quant_config, ++ embedding_prefix, ++ getattr(config, "ple_embedding_dtype", None), ++ ) ++ if not isinstance( ++ method, ++ (Qwen4ExpPLEFp8EmbeddingMethod, Qwen4ExpPLEUnquantizedEmbeddingMethod), ++ ): ++ raise NotImplementedError( ++ f"PLE pull transport does not support {type(method).__name__}" ++ ) ++ module._offload_quant_method = method ++ module._pull_params_dtype = params_dtype or torch.get_default_dtype() ++ module._pull_embedding_prefix = embedding_prefix ++ if not is_offload_process(): ++ # The GPU-rank producer line, in upstream's format plus backend= ++ # (the offload process logs its own table line from upstream's ++ # constructor). The table lives in the offload process's CPU memory ++ # and is not pinned; rows cross via host-mapped buffers. ++ logger.info( ++ "Initialized PLE embedding %s: quantization_method=%s, " ++ "weight_dtype=%s, weight_device=%s, pinned=%s, backend=%s", ++ embedding_prefix, ++ type(method).__name__, ++ module.get_offload_output_dtype(module._pull_params_dtype), ++ "cpu", ++ False, ++ PULL_BACKEND_NAME, ++ ) ++ return module ++ ++ # ------------------------------------------------------------- formats ++ ++ def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: ++ """Rows cross the transport in their table storage dtype.""" ++ embedding = getattr(self, "ngram_embedding", None) ++ weight = getattr(embedding, "weight", None) ++ if weight is not None: ++ return weight.dtype ++ if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod): ++ return torch.float8_e4m3fn ++ return getattr(self, "_pull_params_dtype", default_dtype) ++ ++ def get_offload_output_dim(self, default_dim: int) -> int: ++ return default_dim ++ ++ def initialize_dummy_offload_metadata(self, device: torch.device) -> None: ++ """Dummy loads skip checkpoint scales; give FP8 a unit scale.""" ++ if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod): ++ self.register_buffer( ++ "_offload_weight_scale", ++ torch.ones((), dtype=torch.float32, device=device), ++ persistent=False, ++ ) ++ ++ def dequantize_pulled( ++ self, embeddings: torch.Tensor, output_dtype: torch.dtype ++ ) -> torch.Tensor: ++ """GPU-side dequantization of rows DMA'd from the offload process. ++ ++ Same arithmetic as upstream ``Qwen4ExpPLEFp8EmbeddingMethod.dequantize`` ++ (and the E2 ``_dequantize_embeddings`` FP8 branch). ++ """ ++ method = self._offload_quant_method ++ if isinstance(method, Qwen4ExpPLEUnquantizedEmbeddingMethod): ++ return embeddings ++ weight_scale = getattr(self, "_offload_weight_scale", None) ++ if weight_scale is None: ++ raise RuntimeError( ++ "FP8 PLE pull transport is missing its global scale on the GPU" ++ ) ++ if weight_scale.device != embeddings.device: ++ raise RuntimeError("FP8 PLE embedding scale must be on the output device") ++ return embeddings.to(output_dtype) * weight_scale.to(output_dtype) ++ ++ # ------------------------------------------------------------- execution ++ ++ def start_prefetch( ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ ) -> None: ++ """No prefetch: rows are published by the offload process.""" ++ return None ++ ++ def _compute_ngram_ids_offload( ++ self, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ ) -> torch.Tensor: ++ """CPU n-gram IDs, carried unchanged from the E2 offload path. ++ ++ Differs from upstream's CPU branch in two E2 fixes: the packed layout ++ is (num_reqs, max_seq_len) rather than (num_reqs, num_tokens), and ++ CUDA-graph padding rows beyond ``query_start_loc[-1]`` are not written, ++ so they cannot overwrite a request's final real token. ++ """ ++ input_ids = input_ids.reshape(-1).long() ++ query_start_loc = query_start_loc.long() ++ num_reqs = query_start_loc.numel() - 1 ++ num_tokens = input_ids.shape[0] ++ if num_reqs <= 0: ++ raise ValueError("PLE requires at least one request") ++ max_seq_len = max( ++ 1, ++ int((query_start_loc[1:] - query_start_loc[:-1]).max().item()), ++ ) ++ num_valid_tokens = min(int(query_start_loc[-1].item()), num_tokens) ++ ++ positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64) ++ packed = torch.full( ++ (num_reqs, max_seq_len), ++ self.eos_token_id, ++ device=input_ids.device, ++ dtype=torch.int64, ++ ) ++ request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 ++ request_indices.clamp_(max=num_reqs - 1) ++ columns = (positions - query_start_loc[request_indices]).clamp( ++ 0, packed.shape[1] - 1 ++ ) ++ packed[request_indices[:num_valid_tokens], columns[:num_valid_tokens]] = ( ++ input_ids[:num_valid_tokens] ++ ) ++ ngram_context = ngram_context[:num_reqs].to( ++ device=input_ids.device, dtype=torch.long ++ ) ++ ++ context = torch.cat([ngram_context, packed], dim=-1) ++ positions_2d, position_in_segment = self._shift_precompute( ++ context, self.eos_token_id ++ ) ++ shifted = [context] ++ for shift in range(1, self.ngram_size): ++ shifted.append( ++ self._shift_apply( ++ context, ++ positions_2d, ++ position_in_segment, ++ shift, ++ self.eos_token_id, ++ ) ++ ) ++ adjusted_columns = columns + self.ngram_size - 1 ++ id_blocks = [] ++ for ngram in range(2, self.ngram_size + 1): ++ start = (ngram - 2) * self.heads_per_ngram ++ end = start + self.heads_per_ngram ++ mixed = shifted[0] * self.layer_multipliers[0] ++ for index in range(1, ngram): ++ mixed = torch.bitwise_xor( ++ mixed, shifted[index] * self.layer_multipliers[index] ++ ) ++ sizes = self.ngram_heads_vocab_sizes[start:end] ++ offsets = self.ngram_heads_offsets[start:end] ++ ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets ++ id_blocks.append(ids[request_indices, adjusted_columns]) ++ return torch.cat(id_blocks, dim=-1) ++ ++ def forward_impl( # type: ignore[override] ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ output_buffer: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ """Offload process only: gather raw table rows for these tokens.""" ++ del hidden_states ++ if not is_offload_process(): ++ raise RuntimeError( ++ "PLE pull transport: forward_impl ran outside the offload " ++ "process -- the GPU placeholder was never bound by the connector" ++ ) ++ ngram_ids = self._compute_ngram_ids_offload( ++ input_ids, query_start_loc, ngram_context ++ ) ++ num_tokens = input_ids.reshape(-1).shape[0] ++ weight = self.ngram_embedding.weight ++ if output_buffer is not None: ++ output = output_buffer[:num_tokens, : self.embedding_dim] ++ torch.index_select( ++ weight, ++ 0, ++ ngram_ids.reshape(-1), ++ out=output.reshape(-1, self.head_dim), ++ ) ++ return output ++ return torch.index_select(weight, 0, ngram_ids.reshape(-1)).reshape( ++ num_tokens, self.embedding_dim ++ ) ++ ++ # ------------------------------------------------------------- loading ++ ++ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ++ """Offload process: upstream loading. GPU worker: keep the FP8 scale. ++ ++ The GPU placeholder owns no table; every other checkpoint tensor for ++ this subtree is consumed and dropped here (the offload process loads ++ the table itself). Prefix grouping may call this more than once. ++ """ ++ if is_offload_process(): ++ return Qwen4ExpNGramEmbedding.load_weights(self, weights) ++ retained: set[str] = set() ++ fp8 = isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod) ++ for name, loaded_weight in weights: ++ if fp8 and name == "ngram_embedding.weight_scale": ++ self.register_buffer( ++ "_offload_weight_scale", ++ loaded_weight.reshape(()).to( ++ device=torch.accelerator.current_accelerator(), ++ dtype=torch.float32, ++ ), ++ persistent=False, ++ ) ++ retained.add(name) ++ return retained ++ ++ ++__all__ = ["PULL_BACKEND_NAME", "Qwen4ExpNGramEmbeddingPull"] +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +index 9971ccd..12a0d48 100644 +--- a/vllm/v1/e2_guards.py ++++ b/vllm/v1/e2_guards.py +@@ -214,5 +214,151 @@ def assert_graph_profile_is_executable(vllm_config, *, role: str = "target") -> + ) + + ++# -------------------------------------------------------------------------- ++# PLE transport: carry the pull transport, UVA disabled. ++# -------------------------------------------------------------------------- ++ ++#: The fork's pull-transport switch (strict "0"/"1"; registered and hashed in ++#: vllm/envs.py). Renamed from the E2 base's VLLM_PLE_CPU_OFFLOAD. ++PULL_TRANSPORT_ENV = "VLLM_E2_PLE_PULL_TRANSPORT" ++#: Upstream v0.30.0's switch for the pinned-host UVA PLE backend. Same name as ++#: the E2 base's pull switch, opposite default (upstream: "1"). ++UPSTREAM_UVA_ENV = "VLLM_PLE_CPU_OFFLOAD" ++ ++ ++def pull_transport_enabled() -> bool: ++ """Whether the fork's pull transport is on. Strict; raises on bad values. ++ ++ One source: the registered accessor ``envs.VLLM_E2_PLE_PULL_TRANSPORT``, ++ whose parser (``envs._e2_strict_switch``) accepts exactly unset/"0"/"1". ++ """ ++ import vllm.envs as envs ++ ++ return bool(envs.VLLM_E2_PLE_PULL_TRANSPORT) ++ ++ ++def check_upstream_uva_env_against_pull() -> None: ++ """Fail loud if the old name would turn upstream's UVA path on. ++ ++ With the pull transport on, ``VLLM_PLE_CPU_OFFLOAD`` may be unset (the ++ fork then supplies ``cpu_offload=False`` as the Engram default) or exactly ++ ``"0"``. ``"1"`` -- the E2-era spelling of "enable our offload" -- would ++ enable upstream's pinned-host UVA backend beside ours: a double offload. ++ Any other value is rejected rather than interpreted. With the pull ++ transport off this function does nothing: the name is upstream's. ++ """ ++ if not pull_transport_enabled(): ++ return ++ raw = os.environ.get(UPSTREAM_UVA_ENV) ++ if raw is None or raw == "0": ++ return ++ if raw == "1": ++ raise RuntimeError( ++ f"e2-guard: {PULL_TRANSPORT_ENV}=1 and {UPSTREAM_UVA_ENV}=1. On " ++ f"vLLM v0.30.0 {UPSTREAM_UVA_ENV} is UPSTREAM's switch for the " ++ "pinned-host UVA PLE backend (default on), not the E2 pull " ++ "transport, so this launch asks for both transports at once. Unset " ++ f"{UPSTREAM_UVA_ENV} or export {UPSTREAM_UVA_ENV}=0 " ++ "(the pull transport replaces that switch)." ++ ) ++ raise ValueError( ++ f"e2-guard: {PULL_TRANSPORT_ENV}=1 and {UPSTREAM_UVA_ENV}={raw!r}. " ++ f'With the pull transport on, {UPSTREAM_UVA_ENV} must be unset or ' ++ 'exactly "0"; any other value is refused rather than guessed.' ++ ) ++ ++ ++def pull_transport_cpu_offload_default(upstream_default: bool) -> bool: ++ """The Engram ``cpu_offload`` default, as this fork resolves it. ++ ++ Pull off: ``upstream_default`` unchanged (the tree behaves as upstream). ++ Pull on: the old-name check above, then ``False`` -- upstream's UVA ++ backend is never the default beside the pull transport. ++ """ ++ if not pull_transport_enabled(): ++ return upstream_default ++ check_upstream_uva_env_against_pull() ++ return False ++ ++ ++def verify_engram_against_pull(engram_config) -> None: ++ """Refuse an explicit ``cpu_offload=True`` while the pull transport is on. ++ ++ ``--engram-config '{"cpu_offload": true}'`` takes precedence over every ++ default, so it is checked on the resolved object, not the environment. ++ """ ++ if not pull_transport_enabled(): ++ return ++ check_upstream_uva_env_against_pull() ++ if engram_config is not None and bool(engram_config.cpu_offload): ++ raise RuntimeError( ++ f"e2-guard: {PULL_TRANSPORT_ENV}=1 but the resolved EngramConfig " ++ "has cpu_offload=True (upstream's pinned-host UVA PLE backend). " ++ "Both transports would own the PLE table. Drop cpu_offload from " ++ "--engram-config (or set it false), or turn the pull transport off." ++ ) ++ ++ ++#: Effective PLE modes, as the guard names them. ++PLE_MODE_PULL = "pull" # E2 pull transport (VLLM_E2_PLE_PULL_TRANSPORT=1) ++PLE_MODE_UVA = "uva" # upstream Qwen4ExpPLEPinnedHostEmbedding ++PLE_MODE_DEVICE = "device" # upstream Qwen4ExpPLEDeviceEmbedding ++PLE_MODE_NONE = "none" # no Engram config resolved (no PLE layers / not CUDA) ++ ++ ++def resolve_ple_mode(vllm_config) -> str: ++ """The PLE transport this process will build, or an error if mixed. ++ ++ Reads the two inputs that decide it -- the pull switch (strict) and the ++ RESOLVED ``engram_config.cpu_offload`` -- and nothing else. Mixed states ++ raise; there is no fallback from one transport to another. ++ """ ++ pull = pull_transport_enabled() ++ check_upstream_uva_env_against_pull() ++ engram_config = getattr(vllm_config, "engram_config", None) ++ cpu_offload = ( ++ None if engram_config is None else bool(engram_config.cpu_offload) ++ ) ++ if pull: ++ if cpu_offload: ++ raise RuntimeError( ++ f"e2-guard: {PULL_TRANSPORT_ENV}=1 with a resolved Engram " ++ "cpu_offload=True: the pull transport and upstream's UVA " ++ "backend would both own the PLE table." ++ ) ++ return PLE_MODE_PULL ++ if cpu_offload is None: ++ return PLE_MODE_NONE ++ return PLE_MODE_UVA if cpu_offload else PLE_MODE_DEVICE ++ ++ ++def assert_ple_mode(vllm_config, *, role: str = "target") -> str: ++ """Startup guard: resolve, assert and log the effective PLE mode. ++ ++ Runs beside ``assert_graph_profile_is_executable`` at model construction ++ for ``role="target"`` and ``role="draft"`` (the MTP drafter carries the ++ target's EngramConfig and the same process environment, so a mixed ++ configuration fails there too). NOT disabled by VLLM_E2_ALLOW_BREAKABLE: ++ that opt-out is about graph identity, not about owning the PLE table twice. ++ """ ++ mode = resolve_ple_mode(vllm_config) ++ model_config = getattr(vllm_config, "model_config", None) ++ text_config = getattr(model_config, "hf_text_config", None) ++ ple_layer_ids = list(getattr(text_config, "ple_layer_ids", None) or []) ++ engram_config = getattr(vllm_config, "engram_config", None) ++ logger.info( ++ "e2-guard[%s]: ple_mode=%s pull_transport=%d engram_cpu_offload=%s " ++ "%s=%s ple_layer_ids=%s.", ++ role, ++ mode, ++ 1 if mode == PLE_MODE_PULL else 0, ++ None if engram_config is None else bool(engram_config.cpu_offload), ++ UPSTREAM_UVA_ENV, ++ os.environ.get(UPSTREAM_UVA_ENV, ""), ++ ple_layer_ids, ++ ) ++ return mode ++ ++ + #: Back-compatible name for the earlier call sites. + assert_breakable_cudagraph_pinned = assert_graph_profile_is_executable +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index e27a18b..fc744cb 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -713,6 +713,10 @@ class WorkerProc: + + # Load model + self.worker.init_device() ++ # PLE pull transport: DP0/TP0 starts the CPU-only PLE process before ++ # load_model, so its weight load overlaps ours; READY is awaited after. ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ self.worker.spawn_ple_offload() + # Update process title now that parallel groups are initialized + self.setup_proc_title_and_log_prefix( + enable_ep=vllm_config.parallel_config.enable_expert_parallel +@@ -721,6 +725,8 @@ class WorkerProc: + self.worker.elastic_ep_execute("load_model") + else: + self.worker.load_model() ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ self.worker.wait_ple_offload_ready() + + scheduler_config = vllm_config.scheduler_config + self.use_async_scheduling = scheduler_config.async_scheduling +diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py +index dc798cd..eeb0786 100644 +--- a/vllm/v1/executor/uniproc_executor.py ++++ b/vllm/v1/executor/uniproc_executor.py +@@ -67,11 +67,15 @@ class UniProcExecutor(Executor): + + self.driver_worker.init_worker(all_kwargs=[kwargs]) + self.driver_worker.init_device() ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ self.driver_worker.spawn_ple_offload() + + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: + self.driver_worker.elastic_ep_execute("load_model") + else: + self.driver_worker.load_model() ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ self.driver_worker.wait_ple_offload_ready() + current_platform.update_block_size_for_backend(self.vllm_config) + + def _distributed_args(self) -> tuple[str, int, int]: +diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py +index 52bc462..2779edf 100644 +--- a/vllm/v1/ple_offload/connector.py ++++ b/vllm/v1/ple_offload/connector.py +@@ -257,7 +257,7 @@ class PleOffloadConnector: + } + if not layers: + raise RuntimeError( +- "VLLM_PLE_CPU_OFFLOAD is enabled, but the model has no PleOffloadLayer" ++ "VLLM_E2_PLE_PULL_TRANSPORT is enabled, but the model has no PleOffloadLayer" + ) + + config = vllm_config.model_config.hf_text_config +diff --git a/vllm/v1/ple_offload/test_fault.py b/vllm/v1/ple_offload/test_fault.py +new file mode 100644 +index 0000000..9ae74f7 +--- /dev/null ++++ b/vllm/v1/ple_offload/test_fault.py +@@ -0,0 +1,13 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Compatibility exports for the top-level test-only fault singleton.""" ++ ++from vllm_ple_test_fault import ( ++ _FAULT as _FAULT, ++) ++from vllm_ple_test_fault import ( # noqa: F401 ++ FaultSpec as FaultSpec, ++) ++from vllm_ple_test_fault import ( ++ _parse_fault_spec as _parse_fault_spec, ++) +diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py +index 24beebd..a5f039e 100644 +--- a/vllm/v1/ple_offload/worker.py ++++ b/vllm/v1/ple_offload/worker.py +@@ -830,7 +830,7 @@ class PleOffloadRunner: + } + if not offload_layers: + raise RuntimeError( +- "VLLM_PLE_CPU_OFFLOAD is enabled, but no PleOffloadLayer " ++ "VLLM_E2_PLE_PULL_TRANSPORT is enabled, but no PleOffloadLayer " + "was found in the initialized model" + ) + logger.info( +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index e2ddc52..2d80a37 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -398,6 +398,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + + # For transferring state from execute_model to subsequent sample_tokens call. + self.execute_model_state: ExecuteModelState | None = None ++ # PLE pull transport: set by _setup_ple_offload() from Worker.load_model ++ # when VLLM_E2_PLE_PULL_TRANSPORT=1 and this rank owns a PLE layer. ++ # None otherwise, and every hook below is then a no-op (upstream path). ++ self._ple_offload_connector: PleOffloadConnector | None = None + + # Expert parallelism load balancer. + self.eplb = EPLBController(self.parallel_config, self.device) +@@ -429,6 +433,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): + tasks.extend(PoolingRunner.get_supported_tasks(self.model)) + return tuple(tasks) + ++ def _setup_ple_offload(self, ipc_addr: str) -> None: ++ """Initialize PLE offload after the model state is available.""" ++ # PLE placeholders are created by model loading, so CUDA output ++ # registration cannot happen in the runner constructor. ++ query_start_loc_source = getattr(self.model_state, "ple_query_start_loc", None) ++ ngram_context_source = getattr(self.model_state, "ngram_context", None) ++ if not isinstance(query_start_loc_source, torch.Tensor): ++ raise RuntimeError("PLE offload requires a query_start_loc source") ++ if not isinstance(ngram_context_source, torch.Tensor): ++ raise RuntimeError("PLE offload requires an ngram_context source") ++ self._ple_offload_connector = PleOffloadConnector( ++ self.vllm_config, ++ self.model, ++ self.device, ++ ipc_addr, ++ input_ids_source=self.input_buffers.input_ids, ++ query_start_loc_source=query_start_loc_source, ++ ngram_context_source=ngram_context_source, ++ ) ++ + def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: + time_before_load = time.perf_counter() + if load_dummy_weights: +@@ -1065,6 +1089,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): + torch.accelerator.empty_cache() + start_free_gpu_memory = torch.accelerator.get_memory_info()[0] + ++ # Decoder capture calls the model outside execute_model. Keep dummy ++ # PLE outputs signaled until all graphs capture their semaphore wait. ++ if self._ple_offload_connector is not None and capture_decoder: ++ self._ple_offload_connector.signal_dummy_outputs(self.max_num_tokens) + with self.maybe_setup_dummy_loras(self.lora_config): + if capture_encoder: + self.model_state.encoder_runner.capture() +@@ -1100,6 +1128,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.adaptive_verification.set_initial_cost_curves(timings) + self.kv_connector.reset_capture_state() + ++ if self._ple_offload_connector is not None and capture_decoder: ++ self._ple_offload_connector.release_outputs() + end_free_gpu_memory = torch.accelerator.get_memory_info()[0] + + if not profile_only: +@@ -2442,6 +2472,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + torch.accelerator.synchronize() ++ if self._ple_offload_connector is not None: ++ self._ple_offload_connector.close() ++ self._ple_offload_connector = None + self.cudagraph_manager = None + self.fast_prefill = None + if hasattr(self, "kv_caches"): +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index cafeb73..847837d 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -226,6 +226,22 @@ class Worker(WorkerBase): + self.profiler_config = vllm_config.profiler_config + + self.use_v2_model_runner = vllm_config.use_v2_model_runner ++ # PLE pull transport (VLLM_E2_PLE_PULL_TRANSPORT). Off: both stay ++ # None/False and every PLE hook in this class is a no-op. ++ self._ple_offload_worker_handle: Any | None = None ++ self._ple_offload_stopping = False ++ self._ple_offload_enabled = self._has_ple_layers() ++ if envs.VLLM_E2_PLE_PULL_TRANSPORT: ++ if self._ple_offload_enabled: ++ self._validate_ple_offload_config() ++ elif self.rank == 0 and self.parallel_config.data_parallel_rank == 0: ++ text_config = self.model_config.hf_text_config ++ logger.warning( ++ "VLLM_E2_PLE_PULL_TRANSPORT is enabled, but this rank owns " ++ "no PLE layer (ple_layer_ids=%s); no PLE offload connector " ++ "is created here.", ++ getattr(text_config, "ple_layer_ids", None), ++ ) + + # Device handles of the previous step's PP intermediate-tensor send. + self._pp_send_work: list[Handle] = [] +@@ -235,7 +251,7 @@ class Worker(WorkerBase): + + def _has_ple_layers(self) -> bool: + """Return whether this model configuration constructs PLE layers.""" +- if not envs.VLLM_PLE_CPU_OFFLOAD: ++ if not envs.VLLM_E2_PLE_PULL_TRANSPORT: + return False + text_config = self.model_config.hf_text_config + ple_layer_ids = getattr(text_config, "ple_layer_ids", None) +@@ -297,7 +313,7 @@ class Worker(WorkerBase): + + if unsupported: + raise ValueError( +- "VLLM_PLE_CPU_OFFLOAD does not support the requested " ++ "VLLM_E2_PLE_PULL_TRANSPORT does not support the requested " + "configuration. Unsupported settings: " + f"{', '.join(unsupported)}" + ) +@@ -651,6 +667,12 @@ class Worker(WorkerBase): + ): + self.model_runner.load_model(load_dummy_weights=load_dummy_weights) + ++ # Model loading constructs the PLE placeholders whose buffers are ++ # registered with the offload process by the model runner V2 connector. ++ if self._ple_offload_enabled: ++ model_runner = cast("GPUModelRunnerV2", self.model_runner) ++ model_runner._setup_ple_offload(self.parallel_config._ple_offload_ipc_path) ++ + if has_ec_transfer(): + get_ec_transfer().start_worker_services() + +@@ -666,6 +688,10 @@ class Worker(WorkerBase): + self.model_runner.update_config(overrides) + + def reload_weights(self, *args, **kwargs) -> None: ++ if self._ple_offload_enabled: ++ raise NotImplementedError( ++ "Weight reload is not supported with the PLE pull transport" ++ ) + with set_current_vllm_config(self.vllm_config): + self.model_runner.reload_weights(*args, **kwargs) + +@@ -1661,6 +1687,10 @@ class Worker(WorkerBase): + + self.elastic_ep_executor.shutdown() + ++ if self._ple_offload_worker_handle is not None: ++ self._ple_offload_worker_handle.close() ++ self._ple_offload_worker_handle = None ++ + # Release GPU resources held by the model runner so that memory + # can be reclaimed when running in-process + if model_runner := getattr(self, "model_runner", None): +-- +2.43.0 + diff --git a/patches/0034-qwen4_exp-restore-torch.compile-on-the-NVIDIA-path-e.patch b/patches/0034-qwen4_exp-restore-torch.compile-on-the-NVIDIA-path-e.patch new file mode 100644 index 0000000..7cf92ca --- /dev/null +++ b/patches/0034-qwen4_exp-restore-torch.compile-on-the-NVIDIA-path-e.patch @@ -0,0 +1,2204 @@ +From 3ffb58395fdba31f5f816aa97d811dce837275c4 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 04:45:12 -0400 +Subject: [PATCH 034/103] qwen4_exp: restore torch.compile on the NVIDIA path; + exact breakable pin; runtime graph-mode assertion + +- torch.compile restored on the NVIDIA Qwen4Exp target and MTP so the + production profile (VLLM_COMPILE + FULL_AND_PIECEWISE, breakable off) runs. +- The production pin requires VLLM_USE_BREAKABLE_CUDAGRAPH to be exactly "0". +- assert_resolved_cudagraph_modes() runs at the top of capture_model. +- CPU tests for the combined profile, the exact pin and the resolved modes. +--- + tests/test_config.py | 7 +- + tests/v1/test_compile_profile_guard.py | 241 ++++++++++ + tests/v1/test_e2_guards.py | 217 +++++---- + tests/v1/test_ple_pull_compile_profile.py | 440 ++++++++++++++++++ + vllm/config/compilation.py | 7 +- + vllm/config/vllm.py | 9 +- + vllm/models/qwen4_exp/nvidia/model.py | 40 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 29 +- + .../qwen4_exp/nvidia/ngram_embedding.py | 213 ++++++++- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 43 +- + vllm/models/qwen4_exp/nvidia/qsa.py | 91 +++- + vllm/v1/e2_guards.py | 393 +++++++++++----- + vllm/v1/worker/gpu/model_runner.py | 7 + + 13 files changed, 1476 insertions(+), 261 deletions(-) + create mode 100644 tests/v1/test_compile_profile_guard.py + create mode 100644 tests/v1/test_ple_pull_compile_profile.py + +diff --git a/tests/test_config.py b/tests/test_config.py +index f75592f..20afe75 100644 +--- a/tests/test_config.py ++++ b/tests/test_config.py +@@ -584,11 +584,12 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( + ("DeepseekV32MTPModel", True, False), + ("GlmMoeDsaForCausalLM", False, True), + ("GlmMoeDsaForCausalLM", True, False), +- ("Qwen4ExpForCausalLM", False, True), ++ # Fork, compile-profile guard: compile restored, so Qwen4Exp is NOT auto-breakable. ++ ("Qwen4ExpForCausalLM", False, False), + ("Qwen4ExpForCausalLM", True, False), +- ("Qwen4ExpForConditionalGeneration", False, True), ++ ("Qwen4ExpForConditionalGeneration", False, False), + ("Qwen4ExpForConditionalGeneration", True, False), +- ("Qwen4ExpMTP", False, True), ++ ("Qwen4ExpMTP", False, False), + ("Qwen4ExpMTP", True, False), + ], + ) +diff --git a/tests/v1/test_compile_profile_guard.py b/tests/v1/test_compile_profile_guard.py +new file mode 100644 +index 0000000..e65dbfc +--- /dev/null ++++ b/tests/v1/test_compile_profile_guard.py +@@ -0,0 +1,241 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Compile-profile guard: the resolved graph profile for Qwen4Exp is VLLM_COMPILE + ++FULL_AND_PIECEWISE with the Qwen4Exp splitting ops present -- and it is NOT ++with upstream v0.30.0's breakable default (the failing control). ++ ++CPU-only. The arch-dependent steps of ``VllmConfig.__post_init__`` run as the ++real unbound methods on a stand-in config (the pattern upstream's own ++``tests/test_config.py::test_breakable_cudagraph_platform_default`` uses), ++because a real ``VllmConfig`` for this architecture needs a checkpoint. The ++splitting ops come from the real ``CompilationConfig.set_splitting_ops_for_v1``. ++""" ++ ++import ast ++import os ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++ ++import vllm.config.vllm as vllm_config_module ++from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, VllmConfig ++from vllm.config.vllm import default_breakable_cudagraph_architectures ++from vllm.v1 import e2_guards ++from vllm.v1.e2_guards import ( ++ ALLOW_BREAKABLE_ENV, ++ BREAKABLE_ENV, ++ EXPECTED_MODE_ENV, ++ REQUIRED_SPLITTING_OPS, ++ assert_graph_profile_is_executable, ++) ++ ++QWEN4EXP_ARCHITECTURES = ( ++ "Qwen4ExpForCausalLM", ++ "Qwen4ExpForConditionalGeneration", ++ "Qwen4ExpMTP", ++) ++#: What upstream v0.30.0 ships (#55272 added the three Qwen4Exp names). ++UPSTREAM_V0300_DEFAULT = frozenset( ++ vllm_config_module.DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES ++) | frozenset(QWEN4EXP_ARCHITECTURES) ++ ++_NVIDIA = Path(vllm_config_module.__file__).resolve().parents[1] / ( ++ "models/qwen4_exp/nvidia" ++) ++ ++ ++@pytest.fixture(autouse=True) ++def _cuda_like_platform(monkeypatch): ++ from vllm.platforms import current_platform ++ ++ for name in (BREAKABLE_ENV, ALLOW_BREAKABLE_ENV, EXPECTED_MODE_ENV): ++ monkeypatch.delenv(name, raising=False) ++ monkeypatch.setattr(current_platform, "is_rocm", lambda: False) ++ # The auto-enable writes os.environ and then asks ++ # is_breakable_cudagraph_enabled(); read os.environ live with envs.py's ++ # exact semantics so a cached envs attribute cannot mask that write. ++ monkeypatch.setattr( ++ "vllm.compilation.breakable_cudagraph.is_breakable_cudagraph_enabled", ++ lambda: os.environ.get(BREAKABLE_ENV, "0") == "1", ++ ) ++ monkeypatch.setattr( ++ e2_guards, "breakable_resolved", lambda: os.environ.get(BREAKABLE_ENV) == "1" ++ ) ++ default_breakable_cudagraph_architectures.cache_clear() ++ yield ++ # _maybe_enable_breakable_cudagraph writes os.environ directly. ++ os.environ.pop(BREAKABLE_ENV, None) ++ default_breakable_cudagraph_architectures.cache_clear() ++ ++ ++def _resolve(architecture: str) -> SimpleNamespace: ++ """The arch-dependent part of VllmConfig.__post_init__, in its order. ++ ++ Starts from the production request (-O2 default mode VLLM_COMPILE, ++ FULL_AND_PIECEWISE) and applies: breakable auto-enable (config/vllm.py ++ _maybe_enable_breakable_cudagraph), piecewise normalisation ++ (_normalize_piecewise_cudagraph_mode), splitting-op resolution ++ (CompilationConfig.set_splitting_ops_for_v1). ++ """ ++ cfg = SimpleNamespace( ++ model_config=SimpleNamespace(architectures=[architecture]), ++ compilation_config=CompilationConfig( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ ), ++ ) ++ cfg._uses_breakable_cudagraph_by_default = lambda: ( ++ VllmConfig._uses_breakable_cudagraph_by_default(cfg) ++ ) ++ enabled = VllmConfig._maybe_enable_breakable_cudagraph(cfg) ++ VllmConfig._normalize_piecewise_cudagraph_mode( ++ cfg, breakable_cudagraph_enabled=enabled ++ ) ++ cfg.compilation_config.set_splitting_ops_for_v1(all2all_backend="naive") ++ cfg.breakable = enabled ++ return cfg ++ ++ ++# ---------------------------------------------------------------- this tree ++ ++ ++@pytest.mark.parametrize("architecture", QWEN4EXP_ARCHITECTURES) ++def test_resolved_profile_is_compile_full_and_piecewise(architecture) -> None: ++ cfg = _resolve(architecture) ++ cc = cfg.compilation_config ++ assert cfg.breakable is False ++ assert os.environ.get(BREAKABLE_ENV) is None # not auto-set ++ assert cc.mode == CompilationMode.VLLM_COMPILE ++ assert cc.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE ++ missing = [op for op in REQUIRED_SPLITTING_OPS if op not in cc.splitting_ops] ++ assert missing == [] ++ ++ ++@pytest.mark.parametrize("role", ["target", "draft"]) ++@pytest.mark.parametrize("architecture", QWEN4EXP_ARCHITECTURES) ++def test_guard_accepts_the_resolved_profile(monkeypatch, architecture, role) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ assert_graph_profile_is_executable(_resolve(architecture), role=role) ++ ++ ++def test_qwen4exp_is_not_in_the_breakable_default_set() -> None: ++ assert not ( ++ set(QWEN4EXP_ARCHITECTURES) & set(default_breakable_cudagraph_architectures()) ++ ) ++ ++ ++# ------------------------------------------------ failing control (upstream) ++ ++ ++@pytest.mark.parametrize("architecture", QWEN4EXP_ARCHITECTURES) ++def test_control_upstream_default_resolves_to_breakable_none( ++ monkeypatch, architecture ++) -> None: ++ """FAILING CONTROL: with v0.30.0's default set the same request resolves ++ to breakable-on + CompilationMode.NONE, i.e. not the qualified profile.""" ++ monkeypatch.setattr( ++ vllm_config_module, ++ "DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES", ++ UPSTREAM_V0300_DEFAULT, ++ ) ++ default_breakable_cudagraph_architectures.cache_clear() ++ cfg = _resolve(architecture) ++ assert cfg.breakable is True ++ assert os.environ.get(BREAKABLE_ENV) == "1" ++ assert cfg.compilation_config.mode == CompilationMode.NONE ++ assert cfg.compilation_config.mode != CompilationMode.VLLM_COMPILE ++ ++ ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_control_guard_rejects_upstream_default(monkeypatch, role) -> None: ++ monkeypatch.setattr( ++ vllm_config_module, ++ "DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES", ++ UPSTREAM_V0300_DEFAULT, ++ ) ++ default_breakable_cudagraph_architectures.cache_clear() ++ cfg = _resolve("Qwen4ExpForCausalLM") ++ with pytest.raises(RuntimeError, match=rf"e2-guard\[{role}\].*resolved ON"): ++ assert_graph_profile_is_executable(cfg, role=role) ++ ++ ++def test_control_compilation_none_is_rejected( ++ monkeypatch, ++) -> None: ++ """-O0 / CompilationMode.NONE (the eager-graph alternative) normalises ++ FULL_AND_PIECEWISE to FULL_DECODE_ONLY; the guard must reject it.""" ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ cfg = _resolve("Qwen4ExpForCausalLM") ++ cfg.compilation_config.mode = CompilationMode.NONE ++ VllmConfig._normalize_piecewise_cudagraph_mode( ++ cfg, breakable_cudagraph_enabled=False ++ ) ++ assert cfg.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY ++ with pytest.raises(RuntimeError, match="qualified profile"): ++ assert_graph_profile_is_executable(cfg) ++ ++ ++def test_control_missing_splitting_op_is_rejected(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ cfg = _resolve("Qwen4ExpForCausalLM") ++ cfg.compilation_config.splitting_ops.remove("vllm::qwen4_exp_compute_ple_ngram_ids") ++ with pytest.raises(RuntimeError, match="qwen4_exp_compute_ple_ngram_ids"): ++ assert_graph_profile_is_executable(cfg) ++ ++ ++# --------------------------------------------- source structure (no import) ++ ++ ++def _decorated_classes(source: str) -> set[str]: ++ out = set() ++ for node in ast.walk(ast.parse(source)): ++ if isinstance(node, ast.ClassDef): ++ for deco in node.decorator_list: ++ target = deco.func if isinstance(deco, ast.Call) else deco ++ if getattr(target, "id", None) == "support_torch_compile": ++ out.add(node.name) ++ return out ++ ++ ++def _registered_ops(source: str) -> set[str]: ++ ops = set() ++ for node in ast.walk(ast.parse(source)): ++ if ( ++ isinstance(node, ast.Call) ++ and getattr(node.func, "id", None) == "direct_register_custom_op" ++ ): ++ for kw in node.keywords: ++ if kw.arg == "op_name" and isinstance(kw.value, ast.Constant): ++ ops.add(f"vllm::{kw.value.value}") ++ return ops ++ ++ ++def test_compile_wrappers_are_present_on_model_and_drafter() -> None: ++ model = (_NVIDIA / "model.py").read_text(encoding="utf-8") ++ mtp = (_NVIDIA / "mtp.py").read_text(encoding="utf-8") ++ assert "Qwen4ExpModel" in _decorated_classes(model) ++ assert {"Qwen4ExpMultiTokenPredictor", "Qwen4ExpMTP"} <= _decorated_classes(mtp) ++ ++ ++def test_control_wrapper_check_discriminates() -> None: ++ """The structural check fails on the #55272 shape (decorator removed).""" ++ model = (_NVIDIA / "model.py").read_text(encoding="utf-8") ++ stripped = model.replace("@support_torch_compile(", "@_removed_by_55272(") ++ assert "Qwen4ExpModel" not in _decorated_classes(stripped) ++ ++ ++def test_every_required_qwen4exp_split_is_registered_by_the_nvidia_model() -> None: ++ registered = set() ++ for name in ("qsa.py", "ple_layer.py", "ngram_embedding.py"): ++ registered |= _registered_ops((_NVIDIA / name).read_text(encoding="utf-8")) ++ ours = {op for op in REQUIRED_SPLITTING_OPS if "qwen4_exp" in op} ++ assert ours <= registered, ours - registered ++ # And each is a split in the default list, not only in the guard. ++ assert ours <= set(CompilationConfig._attention_ops) ++ ++ ++def test_guard_module_reads_the_same_list(monkeypatch) -> None: ++ """The guard's list and the config's list must not drift apart.""" ++ assert set(REQUIRED_SPLITTING_OPS) <= set(CompilationConfig._attention_ops) ++ assert e2_guards.REQUIRED_CUDAGRAPH_MODE_NAME == "FULL_AND_PIECEWISE" +diff --git a/tests/v1/test_e2_guards.py b/tests/v1/test_e2_guards.py +index a3b6839..1a72f4a 100644 +--- a/tests/v1/test_e2_guards.py ++++ b/tests/v1/test_e2_guards.py +@@ -2,14 +2,16 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Focused cases for the fork's startup graph-profile guard. + +-An earlier version was rejected for accepting values upstream +-treats as off, rejecting valid ``CompilationMode.NONE`` boots, swallowing read +-errors, and missing the drafter. Each of those is a case below. ++Compile-profile guard: the qualified profile is VLLM_COMPILE + ++FULL_AND_PIECEWISE with breakable capture off and the Qwen4Exp splitting ops ++present; everything else fails loud unless the logged opt-out is set. The ++parser cases are kept from the first version. + + These use a plain stand-in config rather than a real ``VllmConfig``: the guard +-reads exactly three attributes, and building a real config needs a checkpoint. ++reads four attributes, and building a real config needs a checkpoint. + """ + ++import sys + from types import SimpleNamespace + + import pytest +@@ -20,19 +22,26 @@ from vllm.v1.e2_guards import ( + ALLOW_BREAKABLE_ENV, + BREAKABLE_ENV, + EXPECTED_MODE_ENV, ++ REQUIRED_SPLITTING_OPS, ++ assert_compiled_submodule, + assert_graph_profile_is_executable, + breakable_resolved, + ) + + + def _config( +- mode: CompilationMode = CompilationMode.NONE, +- cudagraph_mode: CUDAGraphMode = CUDAGraphMode.FULL, ++ mode: CompilationMode = CompilationMode.VLLM_COMPILE, ++ cudagraph_mode: CUDAGraphMode = CUDAGraphMode.FULL_AND_PIECEWISE, + architectures: tuple[str, ...] = ("Qwen4ExpForCausalLM",), ++ splitting_ops: list[str] | None = None, + ): ++ if splitting_ops is None: ++ splitting_ops = list(REQUIRED_SPLITTING_OPS) + return SimpleNamespace( + model_config=SimpleNamespace(architectures=list(architectures)), +- compilation_config=SimpleNamespace(mode=mode, cudagraph_mode=cudagraph_mode), ++ compilation_config=SimpleNamespace( ++ mode=mode, cudagraph_mode=cudagraph_mode, splitting_ops=splitting_ops ++ ), + ) + + +@@ -40,8 +49,7 @@ def _config( + def _clean_env(monkeypatch): + for name in (BREAKABLE_ENV, ALLOW_BREAKABLE_ENV, EXPECTED_MODE_ENV): + monkeypatch.delenv(name, raising=False) +- # Keep the envs accessor from dominating the raw reading in these cases. +- monkeypatch.setattr(e2_guards, "breakable_resolved", breakable_resolved) ++ monkeypatch.setattr("vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False) + + + # --------------------------------------------------------------- the parser +@@ -51,103 +59,107 @@ def _clean_env(monkeypatch): + def test_only_the_exact_string_one_is_on(monkeypatch, raw: str) -> None: + """Match envs.py:937 exactly. A stricter parser is a FALSE REJECTION.""" + monkeypatch.setenv(BREAKABLE_ENV, raw) +- monkeypatch.setattr( +- "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False +- ) + assert breakable_resolved() is False + + + def test_the_raw_reading_wins_when_envs_is_stale(monkeypatch) -> None: + """_maybe_enable_breakable_cudagraph mutates os.environ after envs caches.""" + monkeypatch.setenv(BREAKABLE_ENV, "1") +- monkeypatch.setattr( +- "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False +- ) + assert breakable_resolved() is True + + + def test_the_envs_reading_wins_when_the_raw_value_is_absent(monkeypatch) -> None: +- monkeypatch.delenv(BREAKABLE_ENV, raising=False) +- monkeypatch.setattr( +- "vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", True, raising=False +- ) ++ monkeypatch.setattr("vllm.envs.VLLM_USE_BREAKABLE_CUDAGRAPH", True, raising=False) + assert breakable_resolved() is True + + +-# ------------------------------------------------------------- the rejections ++def test_an_unreadable_switch_fails_rather_than_reporting_off(monkeypatch) -> None: ++ """A guard that cannot read its own input must not answer False.""" + ++ class _Boom: ++ def __getattr__(self, name): ++ raise RuntimeError("envs unavailable") + +-def test_an_unpinned_launch_fails(monkeypatch) -> None: +- """Absence is the state the upstream auto-enable feeds on.""" +- with pytest.raises(RuntimeError, match="is not set"): +- assert_graph_profile_is_executable(_config()) ++ boom = _Boom() ++ # `import vllm.envs as envs` resolves through the parent attribute first, ++ # so replace both the sys.modules entry and the attribute. ++ monkeypatch.setitem(sys.modules, "vllm.envs", boom) ++ monkeypatch.setattr(sys.modules["vllm"], "envs", boom, raising=False) ++ with pytest.raises(RuntimeError): ++ breakable_resolved() + + +-def test_breakable_on_fails(monkeypatch) -> None: +- monkeypatch.setenv(BREAKABLE_ENV, "1") +- with pytest.raises(RuntimeError, match="resolved ON"): +- assert_graph_profile_is_executable(_config()) ++# ------------------------------------------------------------- the accept + + +-def test_piecewise_with_vllm_compile_and_breakable_off_fails(monkeypatch) -> None: +- """The BLOCKER: upstream accepts this at config time and dies at capture.""" ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_the_qualified_profile_passes(monkeypatch, role) -> None: + monkeypatch.setenv(BREAKABLE_ENV, "0") +- with pytest.raises(RuntimeError, match="cudagraph_utils.py:605"): +- assert_graph_profile_is_executable( +- _config( +- mode=CompilationMode.VLLM_COMPILE, +- cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, +- ) +- ) ++ assert_graph_profile_is_executable(_config(), role=role) + + +-def test_expected_mode_mismatch_fails(monkeypatch) -> None: ++def test_expected_mode_match_passes(monkeypatch) -> None: + monkeypatch.setenv(BREAKABLE_ENV, "0") +- monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL_DECODE_ONLY") +- with pytest.raises(RuntimeError, match="EXPECTED_CUDAGRAPH_MODE"): +- assert_graph_profile_is_executable(_config(cudagraph_mode=CUDAGraphMode.FULL)) ++ monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL_AND_PIECEWISE") ++ assert_graph_profile_is_executable(_config()) + + +-# ------------------------------------------------------------- the permissions ++# ------------------------------------------------------------- the rejections + + +-def test_compilation_none_with_full_graphs_is_permitted(monkeypatch) -> None: +- """FALSE REJECTION control. ++def test_an_unpinned_launch_fails() -> None: ++ with pytest.raises(RuntimeError, match="is not set"): ++ assert_graph_profile_is_executable(_config()) + +- On v0.30.0's NVIDIA Qwen4Exp path NONE is a normal resolved state -- #55272 +- removed the compile wrappers -- so treating NONE as evidence of the +- auto-enable, which the first version of this guard did, rejects the one +- profile that is actually executable. +- """ +- monkeypatch.setenv(BREAKABLE_ENV, "0") +- assert_graph_profile_is_executable( +- _config(mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.FULL) +- ) ++ ++def test_breakable_on_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ with pytest.raises(RuntimeError, match="resolved ON"): ++ assert_graph_profile_is_executable(_config()) + + +-def test_none_plus_full_decode_only_is_permitted(monkeypatch) -> None: ++@pytest.mark.parametrize("role", ["target", "draft"]) ++@pytest.mark.parametrize( ++ ("mode", "cudagraph_mode"), ++ [ ++ (CompilationMode.NONE, CUDAGraphMode.FULL), ++ (CompilationMode.NONE, CUDAGraphMode.FULL_DECODE_ONLY), ++ (CompilationMode.NONE, CUDAGraphMode.NONE), # eager ++ (CompilationMode.VLLM_COMPILE, CUDAGraphMode.PIECEWISE), ++ (CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL), ++ (CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL_DECODE_ONLY), ++ (CompilationMode.STOCK_TORCH_COMPILE, CUDAGraphMode.FULL_AND_PIECEWISE), ++ ], ++) ++def test_every_other_profile_fails(monkeypatch, mode, cudagraph_mode, role) -> None: + monkeypatch.setenv(BREAKABLE_ENV, "0") +- assert_graph_profile_is_executable( +- _config( +- mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY ++ with pytest.raises(RuntimeError, match=rf"e2-guard\[{role}\].*qualified profile"): ++ assert_graph_profile_is_executable( ++ _config(mode=mode, cudagraph_mode=cudagraph_mode), role=role + ) +- ) + + +-def test_eager_is_permitted(monkeypatch) -> None: ++@pytest.mark.parametrize("op", REQUIRED_SPLITTING_OPS) ++def test_each_missing_splitting_op_fails(monkeypatch, op) -> None: + monkeypatch.setenv(BREAKABLE_ENV, "0") +- assert_graph_profile_is_executable( +- _config(mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.NONE) +- ) ++ ops = [o for o in REQUIRED_SPLITTING_OPS if o != op] ++ with pytest.raises(RuntimeError, match=op): ++ assert_graph_profile_is_executable(_config(splitting_ops=ops)) + + +-def test_expected_mode_match_passes(monkeypatch) -> None: ++def test_unresolved_splitting_ops_fail(monkeypatch) -> None: + monkeypatch.setenv(BREAKABLE_ENV, "0") +- monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL") +- assert_graph_profile_is_executable(_config(cudagraph_mode=CUDAGraphMode.FULL)) ++ cfg = _config() ++ cfg.compilation_config.splitting_ops = None ++ with pytest.raises(RuntimeError, match="unresolved"): ++ assert_graph_profile_is_executable(cfg) + + +-# ------------------------------------------------------------ fail-closed paths ++def test_expected_mode_mismatch_fails(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ monkeypatch.setenv(EXPECTED_MODE_ENV, "FULL_DECODE_ONLY") ++ with pytest.raises(RuntimeError, match="EXPECTED_CUDAGRAPH_MODE"): ++ assert_graph_profile_is_executable(_config()) + + + def test_a_missing_compilation_config_fails(monkeypatch) -> None: +@@ -167,49 +179,50 @@ def test_a_missing_model_config_fails(monkeypatch) -> None: + assert_graph_profile_is_executable(cfg) + + +-def test_an_unreadable_switch_fails_rather_than_reporting_off(monkeypatch) -> None: +- """A guard that cannot read its own input must not answer False.""" +- monkeypatch.delenv(BREAKABLE_ENV, raising=False) ++# ----------------------------------------------------- the compiled submodule + +- class _Boom: +- def __getattr__(self, name): +- raise RuntimeError("envs unavailable") + +- monkeypatch.setitem(__import__("sys").modules, "vllm.envs", _Boom()) +- with pytest.raises(RuntimeError): +- breakable_resolved() ++def test_an_uncompiled_module_fails() -> None: ++ """The #55272 shape: a plain module, no torch.compile wrapper.""" ++ import torch.nn as nn + ++ with pytest.raises(RuntimeError, match=r"e2-guard\[draft\].*cudagraph_utils"): ++ assert_compiled_submodule(nn.Linear(1, 1), _config(), role="draft") + +-# ------------------------------------------------------------------ the opt-out + ++def test_a_disabled_wrapper_fails() -> None: ++ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + +-def test_the_opt_out_is_loud_and_skips_the_checks(monkeypatch, caplog) -> None: +- monkeypatch.setenv(ALLOW_BREAKABLE_ENV, "1") +- monkeypatch.setenv(BREAKABLE_ENV, "1") +- with caplog.at_level("WARNING"): +- assert_graph_profile_is_executable( +- _config( +- mode=CompilationMode.VLLM_COMPILE, +- cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, +- ) +- ) +- assert ALLOW_BREAKABLE_ENV in caplog.text +- # The opt-out log line must say it is not a promotion (production) boot. +- assert "promotion" in caplog.text.lower() ++ fake = object.__new__(type("Wrapped", (TorchCompileWithNoGuardsWrapper,), {})) ++ fake.do_not_compile = True ++ with pytest.raises(RuntimeError, match="do_not_compile=True"): ++ assert_compiled_submodule(fake, _config()) + + +-# -------------------------------------------------------------- the draft role ++def test_a_live_wrapper_passes() -> None: ++ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + ++ fake = object.__new__(type("Wrapped", (TorchCompileWithNoGuardsWrapper,), {})) ++ fake.do_not_compile = False ++ assert_compiled_submodule(fake, _config()) + +-def test_the_draft_role_is_checked_independently(monkeypatch) -> None: +- """Qwen4ExpMTP builds Qwen4ExpMultiTokenPredictor, not CausalLM.""" +- monkeypatch.setenv(BREAKABLE_ENV, "0") +- with pytest.raises(RuntimeError, match=r"e2-guard\[draft\]"): +- assert_graph_profile_is_executable( +- _config( +- mode=CompilationMode.VLLM_COMPILE, +- cudagraph_mode=CUDAGraphMode.PIECEWISE, +- architectures=("Qwen4ExpMTP",), +- ), +- role="draft", +- ) ++ ++# ------------------------------------------------------------------ the opt-out ++ ++ ++def test_the_opt_out_is_loud_and_skips_the_checks(monkeypatch) -> None: ++ messages: list[str] = [] ++ monkeypatch.setattr( ++ e2_guards.logger, ++ "warning", ++ lambda msg, *args: messages.append(msg % args), ++ ) ++ monkeypatch.setenv(ALLOW_BREAKABLE_ENV, "1") ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ assert_graph_profile_is_executable( ++ _config(mode=CompilationMode.NONE, cudagraph_mode=CUDAGraphMode.FULL) ++ ) ++ text = "\n".join(messages) ++ assert ALLOW_BREAKABLE_ENV in text ++ # The opt-out log line must say it is not a promotion (production) boot. ++ assert "promotion" in text.lower() +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +new file mode 100644 +index 0000000..f35926d +--- /dev/null ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -0,0 +1,440 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Combined pull-transport + compile-profile tree: profile guard, exact pin, graph modes, plus the ++combined wiring (both guards on target and draft, pull dispatch). ++ ++CPU-only. Configs are stand-ins except where the real ``CompilationConfig`` ++resolver is the thing under test. Every refusal has a FAILING CONTROL: the ++same input with the one deciding factor flipped passes, or upstream's own ++reading is shown to accept what the fork's guard rejects. ++""" ++ ++import ast ++import os ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++ ++import vllm.envs as envs ++from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode ++from vllm.v1 import e2_guards ++from vllm.v1.e2_guards import ( ++ ALLOW_BREAKABLE_ENV, ++ BREAKABLE_ENV, ++ EXPECTED_MODE_ENV, ++ PLE_MODE_PULL, ++ PULL_TRANSPORT_ENV, ++ REQUIRED_SPLITTING_OPS, ++ UPSTREAM_UVA_ENV, ++ assert_graph_profile_is_executable, ++ assert_ple_mode, ++ assert_resolved_cudagraph_modes, ++ breakable_resolved, ++) ++ ++REPO = Path(__file__).resolve().parents[2] ++_NVIDIA = REPO / "vllm/models/qwen4_exp/nvidia" ++ ++ ++def _clear_envs_cache() -> None: ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ ++ ++@pytest.fixture(autouse=True) ++def _clean_env(monkeypatch): ++ for name in ( ++ BREAKABLE_ENV, ++ ALLOW_BREAKABLE_ENV, ++ EXPECTED_MODE_ENV, ++ PULL_TRANSPORT_ENV, ++ UPSTREAM_UVA_ENV, ++ ): ++ monkeypatch.delenv(name, raising=False) ++ _clear_envs_cache() ++ yield ++ _clear_envs_cache() ++ ++ ++def _production_config( ++ architecture: str = "Qwen4ExpForCausalLM", ++ cudagraph_mode: CUDAGraphMode = CUDAGraphMode.FULL_AND_PIECEWISE, ++): ++ """The production request, resolved by the real CompilationConfig.""" ++ cc = CompilationConfig( ++ mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=cudagraph_mode ++ ) ++ cc.set_splitting_ops_for_v1(all2all_backend="naive") ++ return SimpleNamespace( ++ model_config=SimpleNamespace( ++ architectures=[architecture], ++ hf_text_config=SimpleNamespace(ple_layer_ids=[2]), ++ ), ++ compilation_config=cc, ++ engram_config=SimpleNamespace(cpu_offload=False), ++ ) ++ ++ ++def _production_env(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "1") ++ _clear_envs_cache() ++ ++ ++# --------------------------------------------------- the compile profile ++ ++ ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_combined_guards_accept_the_production_profile(monkeypatch, role) -> None: ++ """compile + FULL_AND_PIECEWISE + breakable off + pull on: both guards pass. ++ ++ The pull-transport port alone rejected this profile (its inherited guard); the combined tree ++ carries the compile-profile graph guard and the pull transport's PLE guard side by side. ++ """ ++ _production_env(monkeypatch) ++ cfg = _production_config() ++ missing = [ ++ op ++ for op in REQUIRED_SPLITTING_OPS ++ if op not in cfg.compilation_config.splitting_ops ++ ] ++ assert missing == [] ++ assert_graph_profile_is_executable(cfg, role=role) ++ assert assert_ple_mode(cfg, role=role) == PLE_MODE_PULL ++ ++ ++def test_control_combined_profile_with_uva_requested_is_refused(monkeypatch) -> None: ++ """Same profile, cpu_offload=True resolved: the PLE guard refuses.""" ++ _production_env(monkeypatch) ++ cfg = _production_config() ++ cfg.engram_config.cpu_offload = True ++ assert_graph_profile_is_executable(cfg) # the graph half still passes ++ with pytest.raises(RuntimeError, match="both own the PLE table"): ++ assert_ple_mode(cfg, role="draft") ++ ++ ++def test_constructors_run_both_guards_then_the_compiled_submodule_check() -> None: ++ """Target and draft constructors: profile, PLE, then compiled submodule.""" ++ for rel, cls_name, role in ( ++ ("model.py", "Qwen4ExpForCausalLM", "target"), ++ ("mtp.py", "Qwen4ExpMTP", "draft"), ++ ): ++ tree = ast.parse((_NVIDIA / rel).read_text(encoding="utf-8")) ++ cls = next( ++ n ++ for n in ast.walk(tree) ++ if isinstance(n, ast.ClassDef) and n.name == cls_name ++ ) ++ init = next( ++ n ++ for n in cls.body ++ if isinstance(n, ast.FunctionDef) and n.name == "__init__" ++ ) ++ calls = [] ++ for node in ast.walk(init): ++ if ( ++ isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Name) ++ and node.func.id.startswith("assert_") ++ ): ++ roles = [kw.value.value for kw in node.keywords if kw.arg == "role"] ++ calls.append((node.lineno, node.func.id, roles)) ++ calls.sort() ++ assert [c[1] for c in calls] == [ ++ "assert_graph_profile_is_executable", ++ "assert_ple_mode", ++ "assert_compiled_submodule", ++ ], (rel, calls) ++ assert all(c[2] == [role] for c in calls), (rel, calls) ++ ++ ++# ---------------------------------------------------- the exact "0" pin ++ ++ ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_the_production_pin_zero_passes(monkeypatch, role) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "0") ++ assert_graph_profile_is_executable(_production_config(), role=role) ++ ++ ++@pytest.mark.parametrize("raw", ["2", "true", "True", "false", "", "00", " 0", "no"]) ++@pytest.mark.parametrize("role", ["target", "draft"]) ++def test_a_pin_other_than_exactly_zero_is_refused(monkeypatch, raw, role) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, raw) ++ with pytest.raises(RuntimeError, match=rf"e2-guard\[{role}\].*exactly '0'"): ++ assert_graph_profile_is_executable(_production_config(), role=role) ++ ++ ++@pytest.mark.parametrize("raw", ["2", "true", "", "00"]) ++def test_control_upstream_reads_those_values_as_off(monkeypatch, raw) -> None: ++ """FAILING CONTROL for the lax reading: upstream (and breakable_resolved) ++ calls each of these OFF, so a guard built on resolution alone would have ++ accepted them as a qualified pin -- the compile-profile guard's earlier behaviour.""" ++ monkeypatch.setenv(BREAKABLE_ENV, raw) ++ monkeypatch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", False, raising=False) ++ assert breakable_resolved() is False ++ ++ ++def test_one_is_still_reported_as_resolved_on(monkeypatch) -> None: ++ monkeypatch.setenv(BREAKABLE_ENV, "1") ++ with pytest.raises(RuntimeError, match="resolved ON"): ++ assert_graph_profile_is_executable(_production_config()) ++ ++ ++def test_the_debug_opt_out_is_separate_and_logged(monkeypatch) -> None: ++ """A non-"0" value under the logged opt-out boots (non-production).""" ++ messages: list[str] = [] ++ monkeypatch.setattr( ++ e2_guards.logger, "warning", lambda msg, *args: messages.append(msg % args) ++ ) ++ monkeypatch.setenv(BREAKABLE_ENV, "2") ++ monkeypatch.setenv(ALLOW_BREAKABLE_ENV, "1") ++ assert_graph_profile_is_executable(_production_config()) ++ assert any(ALLOW_BREAKABLE_ENV in m and "promotion" in m.lower() for m in messages) ++ ++ ++# -------------------------------------- resolved modes before capture ++ ++ ++def _manager(mode): ++ return None if mode is None else SimpleNamespace(cudagraph_mode=mode) ++ ++ ++def _ar_speculator(prefill, decode): ++ """AutoRegressive / MTP speculator shape (prefill + decode managers).""" ++ return SimpleNamespace( ++ prefill_cudagraph_manager=_manager(prefill), ++ decode_cudagraph_manager=_manager(decode), ++ ) ++ ++ ++F_AND_P = CUDAGraphMode.FULL_AND_PIECEWISE ++FDO = CUDAGraphMode.FULL_DECODE_ONLY ++ ++ ++def test_resolved_production_modes_pass() -> None: ++ cfg = _production_config() ++ assert_resolved_cudagraph_modes( ++ cfg, _manager(F_AND_P), _ar_speculator(F_AND_P, FDO) ++ ) ++ assert_resolved_cudagraph_modes(cfg, _manager(F_AND_P), None) ++ # Multi-module MTP: one manager, FULL_DECODE_ONLY by design. ++ assert_resolved_cudagraph_modes( ++ cfg, _manager(F_AND_P), SimpleNamespace(cudagraph_manager=_manager(FDO)) ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "target", [CUDAGraphMode.PIECEWISE, CUDAGraphMode.NONE, FDO, CUDAGraphMode.FULL] ++) ++def test_a_downgraded_target_mode_is_refused(target) -> None: ++ with pytest.raises(RuntimeError, match=r"e2-guard\[runtime\].*'target'"): ++ assert_resolved_cudagraph_modes( ++ _production_config(), _manager(target), _ar_speculator(F_AND_P, FDO) ++ ) ++ ++ ++@pytest.mark.parametrize( ++ ("prefill", "decode"), ++ [ ++ (CUDAGraphMode.PIECEWISE, FDO), ++ (FDO, FDO), ++ (F_AND_P, CUDAGraphMode.NONE), ++ (F_AND_P, F_AND_P), ++ (None, FDO), ++ ], ++) ++def test_a_downgraded_draft_mode_is_refused(prefill, decode) -> None: ++ with pytest.raises(RuntimeError, match=r"e2-guard\[runtime\].*draft\."): ++ assert_resolved_cudagraph_modes( ++ _production_config(), _manager(F_AND_P), _ar_speculator(prefill, decode) ++ ) ++ ++ ++def test_a_missing_target_manager_or_unknown_speculator_is_refused() -> None: ++ with pytest.raises(RuntimeError, match="no manager"): ++ assert_resolved_cudagraph_modes(_production_config(), None, None) ++ with pytest.raises(RuntimeError, match="unrecognised speculator"): ++ assert_resolved_cudagraph_modes( ++ _production_config(), _manager(F_AND_P), SimpleNamespace() ++ ) ++ ++ ++def test_other_architectures_are_untouched() -> None: ++ cfg = _production_config(architecture="Qwen3NextForCausalLM") ++ assert_resolved_cudagraph_modes(cfg, _manager(CUDAGraphMode.NONE), None) ++ ++ ++def test_the_runtime_opt_out_is_logged(monkeypatch) -> None: ++ messages: list[str] = [] ++ monkeypatch.setattr( ++ e2_guards.logger, "warning", lambda msg, *args: messages.append(msg % args) ++ ) ++ monkeypatch.setenv(ALLOW_BREAKABLE_ENV, "1") ++ assert_resolved_cudagraph_modes( ++ _production_config(), _manager(CUDAGraphMode.PIECEWISE), None ++ ) ++ assert any("e2-guard[runtime]" in m for m in messages) ++ ++ ++def _resolve_with(cfg, **kwargs): ++ from vllm.v1.attention.backend import AttentionCGSupport ++ ++ args = dict( ++ min_cg_support=AttentionCGSupport.ALWAYS, ++ min_cg_attn_backend="stand-in", ++ uniform_decode_query_len=4, # MTP K=3 ++ use_v2_model_runner=True, ++ piecewise_capture_available=True, ++ ) ++ args.update(kwargs) ++ if isinstance(args["min_cg_support"], str): ++ args["min_cg_support"] = getattr(AttentionCGSupport, args["min_cg_support"]) ++ return cfg.compilation_config.resolve_cudagraph_mode_and_sizes(**args) ++ ++ ++def test_real_resolver_keeps_the_production_mode_and_the_guard_passes() -> None: ++ cfg = _production_config() ++ resolved = _resolve_with(cfg) ++ assert resolved == F_AND_P ++ assert_resolved_cudagraph_modes( ++ cfg, _manager(resolved), _ar_speculator(resolved, FDO) ++ ) ++ ++ ++@pytest.mark.parametrize( ++ ("kwargs", "downgraded"), ++ [ ++ # spec-decode FULL decode with a backend below UNIFORM_BATCH ++ ({"min_cg_support": "UNIFORM_SINGLE_TOKEN_DECODE"}, CUDAGraphMode.PIECEWISE), ++ # no compiled submodule and no breakable graphs: piecewise half dropped ++ ({"piecewise_capture_available": False}, CUDAGraphMode.FULL_DECODE_ONLY), ++ ], ++) ++def test_control_real_resolver_downgrade_is_caught(kwargs, downgraded) -> None: ++ """FAILING CONTROL driven by the real resolver: the construction guard ++ passed on the configured mode, the resolver downgraded it, and only the ++ runtime assertion sees that.""" ++ cfg = _production_config() ++ os.environ[BREAKABLE_ENV] = "0" ++ try: ++ assert_graph_profile_is_executable(cfg) # configured mode: accepted ++ finally: ++ os.environ.pop(BREAKABLE_ENV, None) ++ resolved = _resolve_with(cfg, **kwargs) ++ assert resolved == downgraded ++ with pytest.raises(RuntimeError, match=r"e2-guard\[runtime\]"): ++ assert_resolved_cudagraph_modes(cfg, _manager(resolved), None) ++ ++ ++def _capture_model_calls(source: str) -> list[tuple[int, str]]: ++ tree = ast.parse(source) ++ fn = next( ++ n ++ for n in ast.walk(tree) ++ if isinstance(n, ast.FunctionDef) and n.name == "capture_model" ++ ) ++ events = [] ++ for node in ast.walk(fn): ++ if isinstance(node, ast.Return): ++ events.append((node.lineno, "return")) ++ elif isinstance(node, ast.Call): ++ f = node.func ++ name = f.id if isinstance(f, ast.Name) else getattr(f, "attr", "") ++ if name in ("assert_resolved_cudagraph_modes", "capture"): ++ events.append((node.lineno, name)) ++ return sorted(events) ++ ++ ++def _guard_precedes_everything(events: list[tuple[int, str]]) -> bool: ++ return bool(events) and events[0][1] == "assert_resolved_cudagraph_modes" ++ ++ ++def test_v2_runner_asserts_before_any_capture_or_early_return() -> None: ++ source = (REPO / "vllm/v1/worker/gpu/model_runner.py").read_text(encoding="utf-8") ++ assert _guard_precedes_everything(_capture_model_calls(source)) ++ ++ ++def test_control_runner_check_discriminates() -> None: ++ source = (REPO / "vllm/v1/worker/gpu/model_runner.py").read_text(encoding="utf-8") ++ stripped = source.replace("assert_resolved_cudagraph_modes(\n", "_removed_guard(\n") ++ assert stripped != source ++ assert not _guard_precedes_everything(_capture_model_calls(stripped)) ++ ++ ++# ------------------------------------ pull dispatch ++ ++ ++def _class(path: Path, name: str) -> ast.ClassDef: ++ tree = ast.parse(path.read_text(encoding="utf-8")) ++ return next( ++ n for n in ast.walk(tree) if isinstance(n, ast.ClassDef) and n.name == name ++ ) ++ ++ ++def test_pull_module_dispatches_through_ple_offload_layer_forward() -> None: ++ """PleOffloadLayer is first in the MRO and the pull class defines no ++ forward, so the GPU placeholder runs PleOffloadLayer.forward (the ++ ple_offload_wait_copy op) -- never upstream's device/UVA forward.""" ++ cls = _class(_NVIDIA / "ple_pull.py", "Qwen4ExpNGramEmbeddingPull") ++ assert [b.id for b in cls.bases] == ["PleOffloadLayer", "Qwen4ExpNGramEmbedding"] ++ methods = {n.name for n in cls.body if isinstance(n, ast.FunctionDef)} ++ assert "forward" not in methods ++ assert {"start_prefetch", "forward_impl", "dequantize_pulled"} <= methods ++ ++ ++def test_pull_module_is_built_without_layer_name() -> None: ++ """Only the upstream branch in ple_layer.py gets the compile-profile layer_name.""" ++ tree = ast.parse((_NVIDIA / "ple_layer.py").read_text(encoding="utf-8")) ++ seen = {} ++ for node in ast.walk(tree): ++ if not isinstance(node, ast.Call): ++ continue ++ f = node.func ++ name = ( ++ f"{f.value.id}.{f.attr}" ++ if isinstance(f, ast.Attribute) and isinstance(f.value, ast.Name) ++ else getattr(f, "id", "") ++ ) ++ if name in ("Qwen4ExpNGramEmbeddingPull.build", "Qwen4ExpNGramEmbedding"): ++ seen[name] = {kw.arg for kw in node.keywords} ++ assert "layer_name" not in seen["Qwen4ExpNGramEmbeddingPull.build"] ++ assert "layer_name" in seen["Qwen4ExpNGramEmbedding"] ++ ++ ++def test_device_ops_refuse_a_pull_module() -> None: ++ """The compile-profile op resolver refuses anything carrying the pull marker.""" ++ tree = ast.parse((_NVIDIA / "ngram_embedding.py").read_text(encoding="utf-8")) ++ fn = next( ++ n ++ for n in ast.walk(tree) ++ if isinstance(n, ast.FunctionDef) and n.name == "_ple_embedding_of" ++ ) ++ assert "dequantize_pulled" in ast.unparse(fn) ++ ++ ++def test_wait_copy_op_keeps_its_fake_impl() -> None: ++ tree = ast.parse( ++ (REPO / "vllm/model_executor/layers/ple_offload_layer.py").read_text( ++ encoding="utf-8" ++ ) ++ ) ++ regs = {} ++ for node in ast.walk(tree): ++ if isinstance(node, ast.Call) and getattr(node.func, "id", "") == ( ++ "direct_register_custom_op" ++ ): ++ kws = {kw.arg: kw.value for kw in node.keywords} ++ regs[kws["op_name"].value] = kws ++ wait_copy = regs["ple_offload_wait_copy"] ++ assert ( ++ getattr(wait_copy.get("fake_impl"), "id", None) == "_ple_offload_wait_copy_fake" ++ ) ++ ++ ++def test_wait_copy_is_not_a_splitting_op() -> None: ++ """Add it to splitting_ops only if a reference-host experiment shows ++ piecewise capture of the pull wait/copy is unsafe. Pinned so a change is ++ deliberate.""" ++ assert "vllm::ple_offload_wait_copy" not in CompilationConfig._attention_ops +diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py +index 2cd5fae..8cc2973 100644 +--- a/vllm/config/compilation.py ++++ b/vllm/config/compilation.py +@@ -775,7 +775,12 @@ class CompilationConfig: + "vllm::mamba_mixer2", + "vllm::mamba_mixer", + "vllm::short_conv", +- # Qwen4Exp's AMD backend still uses these splitting ops. ++ # Qwen4Exp: the AMD backend and (fork, compile-profile guard) the compile-restored ++ # NVIDIA backend. #55272 dropped the NVIDIA-only ngram-ids split; the ++ # prefetch pair is the UVA PLE side-stream fork/join (#54371). ++ "vllm::qwen4_exp_compute_ple_ngram_ids", ++ "vllm::qwen4_exp_ple_start_prefetch", ++ "vllm::qwen4_exp_ple_finalize_prefetch", + "vllm::qwen4_exp_ple_short_conv", + "vllm::qwen4_exp_qsa_with_output", + "vllm::linear_attention", +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index bc9fbbe..67ba137 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -97,9 +97,12 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( + "KimiLinearForCausalLM", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", +- "Qwen4ExpForCausalLM", +- "Qwen4ExpForConditionalGeneration", +- "Qwen4ExpMTP", ++ # Fork (compile-profile guard): Qwen4ExpForCausalLM, ++ # Qwen4ExpForConditionalGeneration and Qwen4ExpMTP are deliberately ++ # NOT here. #55272 added them when it removed torch.compile from the ++ # NVIDIA path; this tree restores compile, so the resolved profile for ++ # them is VLLM_COMPILE + FULL_AND_PIECEWISE with breakable capture off ++ # (vllm/v1/e2_guards.py asserts it at model construction). + } + ) + +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 3e18f1c..2b99063 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -9,6 +9,7 @@ from itertools import islice + import torch + from torch import nn + ++from vllm.compilation.decorators import support_torch_compile + from vllm.config import VllmConfig + from vllm.distributed import get_pp_group + from vllm.logger import init_logger +@@ -78,7 +79,11 @@ from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +-from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode ++from vllm.v1.e2_guards import ( ++ assert_compiled_submodule, ++ assert_graph_profile_is_executable, ++ assert_ple_mode, ++) + from vllm.v1.kv_cache_interface import MambaSpec + + from ..config import Qwen4ExpConfig +@@ -486,6 +491,25 @@ class Qwen4ExpMixtureOfExperts(MixtureOfExperts): + moe.experts.update_expert_map() + + ++# Compile-profile guard: torch.compile restored on the NVIDIA ++# path, reverting #55272's removal. Every request-dependent region is a custom ++# op listed in CompilationConfig._attention_ops (qsa_with_output, ++# ple_short_conv, compute_ple_ngram_ids, ple_{start,finalize}_prefetch), so ++# VLLM_COMPILE + FULL_AND_PIECEWISE is executable again with breakable capture ++# off. query_start_loc / ngram_context are fixed-capacity since #55272 ++# (model_state.py); they stay marked dynamic as on v2.0.1 and only ever reach ++# the opaque PLE ops. ++@support_torch_compile( ++ dynamic_arg_dims={ ++ "input_ids": 0, ++ "positions": -1, ++ "intermediate_tensors": 0, ++ "inputs_embeds": 0, ++ "query_start_loc": 0, ++ "ngram_context": 0, ++ "deepstack_input_embeds": 0, ++ } ++) + class Qwen4ExpModel(nn.Module): + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER + +@@ -787,12 +811,11 @@ class Qwen4ExpForCausalLM( + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() +- # V0.30.0 auto-enables breakable cudagraphs +- # for this architecture when VLLM_USE_BREAKABLE_CUDAGRAPH is unset, and +- # VLLM_COMPILE + a PIECEWISE mode dies at capture because #55272 left no +- # compiled submodule. Fail here rather than let a reference-host run qualify an +- # identity nobody chose. Qwen4ExpForConditionalGeneration constructs +- # this class; Qwen4ExpMTP does not, and guards itself. ++ # Compile-profile guard: the only qualified profile on this tree is VLLM_COMPILE + ++ # FULL_AND_PIECEWISE with breakable capture off (the v2.0.1 identity). ++ # Fail at construction, before profiling or capture, on anything else. ++ # Qwen4ExpForConditionalGeneration constructs this class; Qwen4ExpMTP ++ # does not, and guards itself. + assert_graph_profile_is_executable(vllm_config, role="target") + assert_ple_mode(vllm_config, role="target") + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config +@@ -809,6 +832,9 @@ class Qwen4ExpForCausalLM( + self.model = Qwen4ExpModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) ++ # has_compiled_submodule() at gpu/cudagraph_utils.py:605 is what capture ++ # actually tests; check the same property here. ++ assert_compiled_submodule(self.model, vllm_config, role="target") + # PP ownership: logits are computed on the + # last stage only; earlier stages carried an unused 0.59 GiB head. + if get_pp_group().is_last_rank: +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index e7d1255..5640474 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -19,6 +19,7 @@ import regex as re + import torch + from torch import nn + ++from vllm.compilation.decorators import support_torch_compile + from vllm.config import VllmConfig, replace, set_current_vllm_config + from vllm.distributed import get_pp_group + from vllm.model_executor.layers.fused_moe.utils import ( +@@ -49,7 +50,11 @@ from vllm.sequence import IntermediateTensors + from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) +-from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode ++from vllm.v1.e2_guards import ( ++ assert_compiled_submodule, ++ assert_graph_profile_is_executable, ++ assert_ple_mode, ++) + + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +@@ -167,6 +172,18 @@ def _make_draft_vllm_config( + return draft_vllm_config + + ++# Compile restored on the drafter too (reverts #55272 here), as ++# on v2.0.1: both classes wrapped; the outer one compiles, the inner is ++# inlined while Dynamo traces. ++@support_torch_compile( ++ dynamic_arg_dims={ ++ "input_ids": 0, ++ "positions": -1, ++ "intermediate_tensors": 0, ++ "inputs_embeds": 0, ++ "hidden_states": 0, ++ } ++) + class Qwen4ExpMultiTokenPredictor(nn.Module): + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER + +@@ -382,6 +399,15 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + return loader.load_weights(weights, mapper=mapper) + + ++@support_torch_compile( ++ dynamic_arg_dims={ ++ "input_ids": 0, ++ "positions": -1, ++ "intermediate_tensors": 0, ++ "inputs_embeds": 0, ++ "hidden_states": 0, ++ } ++) + class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], +@@ -418,6 +444,7 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "mtp"), + ) ++ assert_compiled_submodule(self.model, vllm_config, role="draft") + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: +diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +index 573d1e3..e6bd44e 100644 +--- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py ++++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +@@ -42,7 +42,10 @@ from vllm.transformers_utils.configs.qwen4_exp import ( + ) + from vllm.triton_utils import tl, triton + from vllm.utils.platform_utils import is_uva_available +-from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor ++from vllm.utils.torch_utils import ( ++ direct_register_custom_op, ++ get_accelerator_view_from_cpu_tensor, ++) + + from ..common.ple import PLEVocabParallelEmbedding + from .ops.ple import ple_ngram_ids +@@ -531,6 +534,15 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): + ) + output.copy_(embeddings.flatten(-2)) + ++ @property ++ def prefetch_buffer(self) -> torch.Tensor: ++ """Side-stream lookup destination (declared mutated by the start op).""" ++ return self._prefetch_buffer ++ ++ def new_prefetch_output(self, num_tokens: int) -> torch.Tensor: ++ """Graph-owned output storage for the finalize op.""" ++ return self._prefetch_buffer.new_empty((num_tokens, self._output_dim)) ++ + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Finish the pinned lookup into graph-owned output storage.""" + output = self._prefetch_buffer.new_empty( +@@ -646,10 +658,16 @@ class Qwen4ExpNGramEmbedding(nn.Module): + *, + data_parallel_rank: int, + prefix: str, ++ layer_name: str | None = None, + quant_config: QuantizationConfig | None = None, + params_dtype: torch.dtype | None = None, + ) -> None: + super().__init__() ++ # Compile-profile guard: the owning Qwen4ExpPLELayer's static_forward_context key. ++ # When set, request-dependent work runs through the custom ops at the ++ # bottom of this file so VLLM_COMPILE never traces it (#55272 removed ++ # them together with @support_torch_compile). None = direct calls. ++ self.layer_name = layer_name + self.embedding_dim = embedding_dim + self.ngram_size = int(config.ngram_size) + self.heads_per_ngram = int(config.heads_per_ngram) +@@ -864,10 +882,39 @@ class Qwen4ExpNGramEmbedding(nn.Module): + ngram_context: torch.Tensor, + ) -> torch.Tensor: + embedding = self.ngram_embedding ++ if self.layer_name is None: ++ if embedding.supports_prefetch: ++ return embedding(hidden_states) ++ ngram_ids = self.compute_ngram_ids( ++ input_ids, query_start_loc, ngram_context ++ ) ++ return self.ngram_embedding(ngram_ids).flatten(-2) ++ + if embedding.supports_prefetch: +- return embedding(hidden_states) +- ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) +- return self.ngram_embedding(ngram_ids).flatten(-2) ++ output = embedding.new_prefetch_output(hidden_states.shape[0]) ++ torch.ops.vllm.qwen4_exp_ple_finalize_prefetch(output, self.layer_name) ++ return output ++ input_ids = input_ids.reshape(-1) ++ # Request-layout-dependent ID generation stays outside PIECEWISE graphs ++ # (a splitting op), as on v2.0.1; the op writes graph-owned storage. ++ ngram_ids = input_ids.new_empty( ++ (input_ids.numel(), self.ngram_heads), dtype=torch.long ++ ) ++ torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids( ++ input_ids, ++ query_start_loc, ++ ngram_context, ++ ngram_ids, ++ self.layer_name, ++ ) ++ # The table lookup is opaque so Inductor never sees (or autotunes ++ # against) the full PLE table -- #55272's ~50 GB activation peak. ++ output = ngram_ids.new_empty( ++ (ngram_ids.shape[0], self.ngram_heads * self.head_dim), ++ dtype=embedding.weight.dtype, ++ ) ++ torch.ops.vllm.qwen4_exp_ple_device_lookup(ngram_ids, output, self.layer_name) ++ return output + + def start_prefetch( + self, +@@ -880,6 +927,32 @@ class Qwen4ExpNGramEmbedding(nn.Module): + embedding = self.ngram_embedding + if not embedding.supports_prefetch: + return ++ if self.layer_name is not None: ++ # Forks onto the prefetch side stream; the join is in the finalize ++ # op. Both are splitting ops, so neither end of the fork is ever ++ # inside a piecewise capture. ++ torch.ops.vllm.qwen4_exp_ple_start_prefetch( ++ hidden_states, ++ input_ids, ++ query_start_loc, ++ ngram_context, ++ embedding.prefetch_buffer, ++ self.layer_name, ++ ) ++ return ++ self._start_prefetch_eager( ++ hidden_states, input_ids, query_start_loc, ngram_context ++ ) ++ ++ def _start_prefetch_eager( ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ ) -> None: ++ """Compute IDs and launch the pinned lookup (runs outside compile).""" ++ embedding = self.ngram_embedding + ngram_ids = self.compute_ngram_ids( + input_ids, + query_start_loc, +@@ -954,6 +1027,138 @@ class Qwen4ExpNGramEmbedding(nn.Module): + return loaded + + ++# --- compile support ------------------------------------------------ ++# The four ops below are the NVIDIA PLE's boundary with torch.compile. Each ++# resolves its Qwen4ExpPLELayer through the forward context by name, as ++# v2.0.1's qwen4_exp_compute_ple_ngram_ids did before #55272 removed it. ++# Splitting ops (vllm/config/compilation.py ``_attention_ops``): ++# qwen4_exp_compute_ple_ngram_ids, qwen4_exp_ple_start_prefetch, ++# qwen4_exp_ple_finalize_prefetch. ++# qwen4_exp_ple_device_lookup is opaque but NOT a split: a gather plus the ETP ++# all-reduce is capturable; it only has to be hidden from Inductor. ++ ++ ++def _ple_embedding_of(layer_name: str) -> Qwen4ExpNGramEmbedding: ++ layer = get_forward_context().no_compile_layers[layer_name] ++ embedding = getattr(layer, "ple_embedding", None) ++ if not isinstance(embedding, Qwen4ExpNGramEmbedding): ++ raise TypeError(f"{layer_name} does not own a Qwen4Exp PLE embedding") ++ if hasattr(type(embedding), "dequantize_pulled"): ++ # Fork (pull transport + compile profile): the pull transport's GPU placeholder owns no ++ # table and dispatches through PleOffloadLayer.forward; it must never ++ # reach these device/UVA ops (ple_layer.py builds it without a ++ # layer_name). Same marker ple_layer.py uses to select dequantisation. ++ raise TypeError( ++ f"{layer_name}: PLE pull transport module routed through an " ++ "upstream device/UVA PLE op" ++ ) ++ return embedding ++ ++ ++def qwen4_exp_compute_ple_ngram_ids( ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ output: torch.Tensor, ++ layer_name: str, ++) -> None: ++ """Compute request-dependent PLE n-gram IDs outside piecewise graphs.""" ++ _ple_embedding_of(layer_name).compute_ngram_ids( ++ input_ids, query_start_loc, ngram_context, output ++ ) ++ ++ ++def qwen4_exp_ple_device_lookup( ++ ngram_ids: torch.Tensor, ++ output: torch.Tensor, ++ layer_name: str, ++) -> None: ++ """Look up device-resident PLE rows; keeps the table out of Inductor.""" ++ embedding = _ple_embedding_of(layer_name) ++ output.copy_(embedding.ngram_embedding(ngram_ids).flatten(-2)) ++ ++ ++def qwen4_exp_ple_start_prefetch( ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ prefetch_buffer: torch.Tensor, ++ layer_name: str, ++) -> None: ++ """Start the pinned-host (UVA) lookup on the side stream.""" ++ embedding = _ple_embedding_of(layer_name) ++ owned = embedding.ngram_embedding.prefetch_buffer ++ if prefetch_buffer.data_ptr() != owned.data_ptr(): ++ raise RuntimeError(f"{layer_name}: PLE prefetch buffer identity changed") ++ embedding._start_prefetch_eager( ++ hidden_states, input_ids, query_start_loc, ngram_context ++ ) ++ ++ ++def qwen4_exp_ple_finalize_prefetch(output: torch.Tensor, layer_name: str) -> None: ++ """Join the side stream and write the pinned lookup into ``output``.""" ++ pinned = _ple_embedding_of(layer_name).ngram_embedding ++ pinned._finalize_prefetch(pinned.prefetch_buffer, output) ++ ++ ++def _compute_ple_ngram_ids_fake( ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ output: torch.Tensor, ++ layer_name: str, ++) -> None: ++ return None ++ ++ ++def _ple_device_lookup_fake( ++ ngram_ids: torch.Tensor, output: torch.Tensor, layer_name: str ++) -> None: ++ return None ++ ++ ++def _ple_start_prefetch_fake( ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ ngram_context: torch.Tensor, ++ prefetch_buffer: torch.Tensor, ++ layer_name: str, ++) -> None: ++ return None ++ ++ ++def _ple_finalize_prefetch_fake(output: torch.Tensor, layer_name: str) -> None: ++ return None ++ ++ ++direct_register_custom_op( ++ op_name="qwen4_exp_compute_ple_ngram_ids", ++ op_func=qwen4_exp_compute_ple_ngram_ids, ++ mutates_args=["output"], ++ fake_impl=_compute_ple_ngram_ids_fake, ++) ++direct_register_custom_op( ++ op_name="qwen4_exp_ple_device_lookup", ++ op_func=qwen4_exp_ple_device_lookup, ++ mutates_args=["output"], ++ fake_impl=_ple_device_lookup_fake, ++) ++direct_register_custom_op( ++ op_name="qwen4_exp_ple_start_prefetch", ++ op_func=qwen4_exp_ple_start_prefetch, ++ mutates_args=["prefetch_buffer"], ++ fake_impl=_ple_start_prefetch_fake, ++) ++direct_register_custom_op( ++ op_name="qwen4_exp_ple_finalize_prefetch", ++ op_func=qwen4_exp_ple_finalize_prefetch, ++ mutates_args=["output"], ++ fake_impl=_ple_finalize_prefetch_fake, ++) ++ ++ + __all__ = [ + "Qwen4ExpNGramEmbedding", + ] +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index d50781b..94d908e 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -21,6 +21,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( + from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) ++from vllm.utils.torch_utils import direct_register_custom_op + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionBackend, +@@ -97,6 +98,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + if envs.VLLM_E2_PLE_PULL_TRANSPORT: + # Fork: the E2 pull transport owns the table in a CPU-only + # offload process; upstream's device/UVA backends are never built. ++ # Pull transport with the compile-profile guard: no layer_name -- the pull module never routes ++ # through the n-gram custom ops. Its GPU placeholder dispatches ++ # through PleOffloadLayer.forward (the ple_offload_wait_copy op) and ++ # owns no table; the offload process calls forward_impl directly. + from .ple_pull import Qwen4ExpNGramEmbeddingPull + + self.ple_embedding = Qwen4ExpNGramEmbeddingPull.build( +@@ -117,6 +122,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + vllm_config.scheduler_config.max_num_batched_tokens, + data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, + prefix=f"{prefix}.ple_embedding", ++ layer_name=prefix, + quant_config=quant_config, + params_dtype=model_config.dtype, + ) +@@ -446,11 +452,44 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + self.norm_key.eps, + ) + # State routing depends on runtime request metadata and remains outside +- # the graph; short convolution accumulates into gated_output. +- self._short_conv(conv_input, gated_output, hidden_states) ++ # the graph; short convolution accumulates into gated_output. Compile-profile guard: ++ # a splitting op again (as on v2.0.1), so VLLM_COMPILE never traces it. ++ torch.ops.vllm.qwen4_exp_ple_short_conv( ++ conv_input, gated_output, hidden_states, self.prefix ++ ) + return gated_output + + ++def qwen4_exp_ple_short_conv( ++ inputs: torch.Tensor, ++ residual_output: torch.Tensor, ++ outer_residual: torch.Tensor, ++ layer_name: str, ++) -> None: ++ """PLE short convolution + state routing, outside piecewise graphs.""" ++ layer = get_forward_context().no_compile_layers[layer_name] ++ if not isinstance(layer, Qwen4ExpPLELayer): ++ raise TypeError(f"{layer_name} is not a Qwen4Exp PLE layer") ++ layer._short_conv(inputs, residual_output, outer_residual) ++ ++ ++def _qwen4_exp_ple_short_conv_fake( ++ inputs: torch.Tensor, ++ residual_output: torch.Tensor, ++ outer_residual: torch.Tensor, ++ layer_name: str, ++) -> None: ++ return None ++ ++ ++direct_register_custom_op( ++ op_name="qwen4_exp_ple_short_conv", ++ op_func=qwen4_exp_ple_short_conv, ++ mutates_args=["residual_output"], ++ fake_impl=_qwen4_exp_ple_short_conv_fake, ++) ++ ++ + __all__ = [ + "Qwen4ExpPLEGroupedNorm", + "Qwen4ExpPLELayer", +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 87c4b3c..73dd07a 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -38,6 +38,10 @@ from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, + ) + from vllm.utils.torch_utils import ( ++ LayerNameType, ++ _encode_layer_name, ++ _resolve_layer_name, ++ direct_register_custom_op, + kv_cache_dtype_str_to_dtype, + ) + from vllm.v1.attention.backend import ( +@@ -820,17 +824,35 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + if _qsa_collect_dir: + _qsa_collect_absmax(self.layer_name, key, value) + attn_output = torch.empty_like(query) +- # Keep the index projection outside the eager break. ++ # Keep the index projection outside the eager break / splitting op. + projected_qk, _ = self.indexer.index_qk_proj(hidden_states) +- self._run_qsa( +- projected_qk, +- positions, +- query, +- key, +- value, +- attn_output, +- gate, +- ) ++ # Compile-profile guard: the QSA transaction reads per-step attention metadata from ++ # the forward context, so under VLLM_COMPILE it must be opaque to ++ # Dynamo and a piecewise split point (compilation.py _attention_ops). ++ # Restores the op #55272 removed, with #55309's fused output gate. ++ encoded_layer_name = _encode_layer_name(self.layer_name) ++ if current_platform.opaque_attention_op(): ++ torch.ops.vllm.qwen4_exp_qsa_with_output( ++ projected_qk, ++ positions, ++ query, ++ key, ++ value, ++ attn_output, ++ gate, ++ encoded_layer_name, ++ ) ++ else: ++ qwen4_exp_qsa_with_output( ++ projected_qk, ++ positions, ++ query, ++ key, ++ value, ++ attn_output, ++ gate, ++ encoded_layer_name, ++ ) + flat_output = attn_output.view(num_tokens, -1) + output, _ = self.o_proj(flat_output) + return output +@@ -972,10 +994,59 @@ def load_qsa_static_kv_scales( + return sorted(validated) + + ++def qwen4_exp_qsa_with_output( ++ projected_qk: torch.Tensor, ++ positions: torch.Tensor, ++ query: torch.Tensor, ++ key: torch.Tensor, ++ value: torch.Tensor, ++ output: torch.Tensor, ++ output_gate: torch.Tensor, ++ layer_name: LayerNameType, ++) -> None: ++ """Run the complete QSA state/update/attend transaction.""" ++ ++ layer_name = _resolve_layer_name(layer_name) ++ layer = get_forward_context().no_compile_layers[layer_name] ++ if not isinstance(layer, Qwen4ExpQSAAttention): ++ raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner") ++ layer._run_qsa( ++ projected_qk, ++ positions, ++ query, ++ key, ++ value, ++ output, ++ output_gate, ++ ) ++ ++ ++def qwen4_exp_qsa_with_output_fake( ++ projected_qk: torch.Tensor, ++ positions: torch.Tensor, ++ query: torch.Tensor, ++ key: torch.Tensor, ++ value: torch.Tensor, ++ output: torch.Tensor, ++ output_gate: torch.Tensor, ++ layer_name: LayerNameType, ++) -> None: ++ del projected_qk, positions, query, key, value, output, output_gate, layer_name ++ ++ ++direct_register_custom_op( ++ op_name="qwen4_exp_qsa_with_output", ++ op_func=qwen4_exp_qsa_with_output, ++ mutates_args=["output"], ++ fake_impl=qwen4_exp_qsa_with_output_fake, ++) ++ ++ + __all__ = [ + "QSAIndexer", + "Qwen4ExpQSAAttention", + "Qwen4ExpQSAFlashAttentionBackend", + "Qwen4ExpQSAFlashAttentionImpl", + "load_qsa_static_kv_scales", ++ "qwen4_exp_qsa_with_output", + ] +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +index 12a0d48..1917885 100644 +--- a/vllm/v1/e2_guards.py ++++ b/vllm/v1/e2_guards.py +@@ -3,33 +3,30 @@ + """Fork startup guards that fail closed on a graph profile we never qualified. + + These run at model construction, on every rank, for the target model and for +-the MTP drafter, before profiling or capture. They exist because a *documented* +-launch flag is not a mechanism: the first reference-host run of a rebase will otherwise +-qualify a compile/capture identity nobody chose +-(see below). +- +-What this module does NOT do: choose a graph profile. That is the maintainer's +-call and it is open (the graph-profile choice). What it does is refuse the +-two states that are silently wrong on vLLM v0.30.0 -- +- +- * ``VLLM_USE_BREAKABLE_CUDAGRAPH`` left unset, which lets +- ``VllmConfig._maybe_enable_breakable_cudagraph`` turn it on for this fork's +- architectures and force ``CompilationMode.NONE``; and +- * a PIECEWISE-bearing cudagraph mode combined with ``VLLM_COMPILE`` and +- breakable-off, which upstream accepts at config time and then rejects at +- capture, because #55272 removed every ``@support_torch_compile`` wrapper +- from the NVIDIA Qwen4Exp model and MTP. +- +--- and to make the resolved profile visible and assertable. +- +-Timing limit, stated honestly: this runs at model construction, i.e. AFTER +-``VllmConfig.__post_init__`` has already called +-``_maybe_enable_breakable_cudagraph`` (``config/vllm.py:1561``) and +-``_normalize_piecewise_cudagraph_mode`` (``:1673``, ``:1844``). It is early +-enough to stop profiling and capture; it cannot un-mutate ``os.environ`` or +-recover the originally requested cudagraph mode. Requiring the pin to be +-*explicitly present* is the part of "validate before auto-enable" that is +-checkable from here: absence is exactly the state the auto-enable feeds on. ++the MTP drafter, before profiling or capture. ++ ++Compile-profile guard: this tree restores torch.compile on the ++NVIDIA Qwen4Exp path (reverting vLLM #55272) and removes Qwen4Exp from ++``DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES``. The ONE qualified profile is the ++v2.0.1 production identity: ++ ++ * ``VLLM_USE_BREAKABLE_CUDAGRAPH`` explicitly pinned to exactly ``0``, and ++ resolved OFF; ++ * ``CompilationMode.VLLM_COMPILE``; ++ * ``CUDAGraphMode.FULL_AND_PIECEWISE``; ++ * every Qwen4Exp request-dependent custom op in ``splitting_ops`` ++ (:data:`REQUIRED_SPLITTING_OPS`); ++ * a live ``@support_torch_compile`` submodule ++ (:func:`assert_compiled_submodule`), which is the property ++ ``vllm/v1/worker/gpu/cudagraph_utils.py:605`` actually tests at capture. ++ ++Anything else -- eager, NONE + FULL, FULL_DECODE_ONLY, breakable PIECEWISE -- ++fails loud. ``VLLM_E2_ALLOW_BREAKABLE=1`` is the single, logged opt-out for ++deliberate non-production boots (debug, probes); a production launch must not set it. ++ ++Timing limit: this runs AFTER ``VllmConfig.__post_init__`` resolved the ++profile. It stops the boot before profiling and capture; it cannot recover ++the originally requested profile if the config layer rewrote it. + """ + + import os +@@ -42,14 +39,32 @@ BREAKABLE_ENV = "VLLM_USE_BREAKABLE_CUDAGRAPH" + #: Opt-out for every guard in this module. Recorded in the launch identity and + #: PROHIBITED in production profiles; a production launch must not set it. + ALLOW_BREAKABLE_ENV = "VLLM_E2_ALLOW_BREAKABLE" +-#: Optional expected resolved cudagraph mode, by name (e.g. "FULL", +-#: "FULL_DECODE_ONLY"). When set, a mismatch is an error. This is how a launch +-#: profile pins the graph identity once the maintainer has chosen one. ++#: Optional expected resolved cudagraph mode, by name. When set, a mismatch is ++#: an error (in addition to the fixed requirement below). + EXPECTED_MODE_ENV = "VLLM_E2_EXPECTED_CUDAGRAPH_MODE" + +-#: The architectures vLLM v0.30.0 auto-enables breakable cudagraphs for, out of +-#: ``vllm/config/vllm.py:DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES``, that are +-#: ours. Recorded so the guard's message can name the upstream mechanism. ++#: The qualified profile (the compile-profile guard). ++REQUIRED_COMPILATION_MODE_NAME = "VLLM_COMPILE" ++REQUIRED_CUDAGRAPH_MODE_NAME = "FULL_AND_PIECEWISE" ++#: The only raw VLLM_USE_BREAKABLE_CUDAGRAPH value accepted as the production ++#: pin. Resolution semantics stay upstream's (_upstream_truth). ++PRODUCTION_BREAKABLE_PIN = "0" ++ ++#: Custom ops that must stay graph breaks on the NVIDIA Qwen4Exp path. Each ++#: reads per-step request metadata (or forks/joins a side stream) and would be ++#: baked into a piecewise graph otherwise. ++REQUIRED_SPLITTING_OPS = ( ++ "vllm::qwen4_exp_qsa_with_output", ++ "vllm::qwen4_exp_ple_short_conv", ++ "vllm::qwen4_exp_compute_ple_ngram_ids", ++ "vllm::qwen4_exp_ple_start_prefetch", ++ "vllm::qwen4_exp_ple_finalize_prefetch", ++ "vllm::qwen_gdn_attention_core", ++ "vllm::qwen_gdn_attention_core_fused_norm_packed", ++) ++ ++#: Upstream v0.30.0 auto-enabled breakable cudagraphs for these; this tree ++#: does not (vllm/config/vllm.py). Kept to name the mechanism in messages. + OUR_AUTO_BREAKABLE_ARCHITECTURES = frozenset( + { + "Qwen4ExpForCausalLM", +@@ -62,10 +77,7 @@ OUR_AUTO_BREAKABLE_ARCHITECTURES = frozenset( + def _upstream_truth(raw: str | None) -> bool: + """Upstream's exact semantics for the breakable switch. + +- ``envs.py:937`` is ``os.environ.get(..., "0") == "1"``. Nothing else is on: +- ``"2"``, ``"true"``, ``"FALSE"`` and ``""`` are all off. Using one parser +- that matches it is the point -- a stricter local parser would reject boots +- upstream treats as off, which is a false rejection, not extra safety. ++ ``envs.py:937`` is ``os.environ.get(..., "0") == "1"``. Nothing else is on. + """ + return raw == "1" + +@@ -73,17 +85,8 @@ def _upstream_truth(raw: str | None) -> bool: + def breakable_resolved() -> bool: + """Whether breakable cudagraphs are on, as the engine will see it. + +- Two readings of ONE source, not three independent sources: +- ``vllm.compilation.breakable_cudagraph.is_breakable_cudagraph_enabled()`` +- is literally ``bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)`` +- (``breakable_cudagraph.py:53``), so it adds nothing over the envs read. +- The envs read is kept beside the raw one because +- ``_maybe_enable_breakable_cudagraph`` mutates ``os.environ`` at config time +- while the envs accessor is cached, so the two can legitimately disagree; +- either saying "on" means on. +- +- No exception is swallowed. A guard that cannot read its own input must fail +- the boot, not report False. ++ The raw environment and the (cached) envs accessor are two readings of one ++ source; either saying "on" means on. No exception is swallowed. + """ + if _upstream_truth(os.environ.get(BREAKABLE_ENV)): + return True +@@ -92,8 +95,12 @@ def breakable_resolved() -> bool: + return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH) + + +-def _profile(vllm_config) -> tuple[list[str], object, object]: +- """Resolved (architectures, compilation mode, cudagraph mode). Fails closed.""" ++def _name(value) -> str: ++ return getattr(value, "name", str(value)) ++ ++ ++def _profile(vllm_config) -> tuple[list[str], object, object, list[str] | None]: ++ """Resolved (architectures, mode, cudagraph mode, splitting ops).""" + model_config = getattr(vllm_config, "model_config", None) + if model_config is None: + raise RuntimeError("e2-guard: vllm_config has no model_config to check") +@@ -105,112 +112,242 @@ def _profile(vllm_config) -> tuple[list[str], object, object]: + architectures, + compilation_config.mode, + compilation_config.cudagraph_mode, ++ getattr(compilation_config, "splitting_ops", None), ++ ) ++ ++ ++def _opted_out(role: str, detail: str) -> bool: ++ if os.environ.get(ALLOW_BREAKABLE_ENV, "") != "1": ++ return False ++ logger.warning( ++ "e2-guard[%s]: %s=1 -- guards disabled by request (%s). This boot is " ++ "NOT a promotion-grade boot; a production launch must not set this flag.", ++ role, ++ ALLOW_BREAKABLE_ENV, ++ detail, + ) ++ return True + + + def assert_graph_profile_is_executable(vllm_config, *, role: str = "target") -> None: +- """Refuse the two silently-wrong startup states on v0.30.0. ++ """Require the qualified profile: VLLM_COMPILE + FULL_AND_PIECEWISE. + + ``role`` is ``"target"`` or ``"draft"``; both are checked, because + ``Qwen4ExpMTP`` builds ``Qwen4ExpMultiTokenPredictor`` rather than +- ``Qwen4ExpForCausalLM`` (``mtp.py:397``, ``:411``) from a config rebuilt at +- ``mtp.py:159``, so the drafter's profile can differ from the target's. +- +- Intentional eager and intentional ``CompilationMode.NONE`` are PERMITTED. +- On v0.30.0's NVIDIA Qwen4Exp path ``NONE`` is a normal resolved state, not +- evidence of the auto-enable: #55272 removed the compile wrappers, so +- ``_normalize_piecewise_cudagraph_mode`` legitimately lands there. ++ ``Qwen4ExpForCausalLM``, from a config rebuilt in ``mtp.py``. + """ +- from vllm.config import CompilationMode +- architectures, mode, cudagraph_mode = _profile(vllm_config) +- +- if os.environ.get(ALLOW_BREAKABLE_ENV, "") == "1": +- logger.warning( +- "e2-guard[%s]: %s=1 -- guards disabled by request. architectures=%s " +- "breakable=%s compilation_mode=%s cudagraph_mode=%s. This boot is " +- "NOT a promotion-grade boot; a production launch must not set this flag.", +- role, +- ALLOW_BREAKABLE_ENV, +- architectures, +- os.environ.get(BREAKABLE_ENV), +- mode, +- cudagraph_mode, +- ) ++ architectures, mode, cudagraph_mode, splitting_ops = _profile(vllm_config) ++ detail = ( ++ f"architectures={architectures} breakable={os.environ.get(BREAKABLE_ENV)} " ++ f"compilation_mode={_name(mode)} cudagraph_mode={_name(cudagraph_mode)}" ++ ) ++ if _opted_out(role, detail): + return + +- # 1. The pin must be explicit. Absence is the state _maybe_enable_breakable_ +- # cudagraph feeds on (config/vllm.py:758: "not in os.environ"). +- if BREAKABLE_ENV not in os.environ: ++ # 1. The pin must be explicit, so a future upstream merge that re-adds us ++ # to DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES cannot flip it silently. ++ raw_breakable = os.environ.get(BREAKABLE_ENV) ++ if raw_breakable is None: + raise RuntimeError( +- f"e2-guard[{role}]: {BREAKABLE_ENV} is not set. vLLM v0.30.0 lists " +- f"{sorted(OUR_AUTO_BREAKABLE_ARCHITECTURES)} in " +- "DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES and turns the switch on " +- "by itself when it is unset (vllm/config/vllm.py:756-774), forcing " +- "CompilationMode.NONE with it. Pin it in the launcher: " +- f"export {BREAKABLE_ENV}=0 (or =1 deliberately), or set " +- f"{ALLOW_BREAKABLE_ENV}=1 to take an unpinned boot on purpose." ++ f"e2-guard[{role}]: {BREAKABLE_ENV} is not set. Pin it in the " ++ f"launcher: export {BREAKABLE_ENV}=0, or set {ALLOW_BREAKABLE_ENV}=1 " ++ "to take an unpinned boot on purpose." + ) + +- breakable = breakable_resolved() ++ # 2. Breakable capture is not the qualified identity. ++ if raw_breakable == "1" or breakable_resolved(): ++ raise RuntimeError( ++ f"e2-guard[{role}]: {BREAKABLE_ENV} resolved ON ({detail}). The " ++ "qualified profile is breakable-off VLLM_COMPILE + " ++ f"FULL_AND_PIECEWISE. Export {BREAKABLE_ENV}=0, or set " ++ f"{ALLOW_BREAKABLE_ENV}=1 to take it deliberately." ++ ) + +- # 2. Breakable on is an identity this fork has never qualified. +- if breakable: ++ # 2b. The production pin is the exact string "0". ++ # Upstream's accessor reads anything but "1" as off, so "2", "true", ++ # "" or "00" would resolve off -- but a mistyped launcher value must ++ # not be recorded as a qualified breakable-off profile. The logged ++ # debug opt-out (VLLM_E2_ALLOW_BREAKABLE=1, handled above) is separate. ++ if raw_breakable != PRODUCTION_BREAKABLE_PIN: + raise RuntimeError( +- f"e2-guard[{role}]: {BREAKABLE_ENV} resolved ON for " +- f"{architectures or ''} " +- f"(compilation_mode={mode}, cudagraph_mode={cudagraph_mode}). That " +- "is not the identity the v2.0.1 results were measured on. Export " +- f"{BREAKABLE_ENV}=0, or set {ALLOW_BREAKABLE_ENV}=1 to take it " +- "deliberately." ++ f"e2-guard[{role}]: {BREAKABLE_ENV}={raw_breakable!r} is not the " ++ f"production pin; it must be exactly {PRODUCTION_BREAKABLE_PIN!r}. " ++ "Upstream would read it as off, but the launch profile does not " ++ f"say what was meant. Export {BREAKABLE_ENV}=0, or set " ++ f"{ALLOW_BREAKABLE_ENV}=1 for a deliberate non-production boot." + ) + +- # 3. PIECEWISE with breakable off needs a torch-compiled submodule, and +- # #55272 removed every @support_torch_compile wrapper from +- # vllm/models/qwen4_exp/nvidia/*. Upstream only rewrites the mode when +- # compilation mode is NOT VLLM_COMPILE +- # (config/vllm.py:1010-1030), so VLLM_COMPILE + PIECEWISE + breakable-off +- # survives config time and then dies at +- # gpu/cudagraph_utils.py:605 in the middle of capture. Fail here instead, +- # with the candidates named. ++ # 3. The profile itself: mode AND cudagraph mode, by name. + if ( +- cudagraph_mode.requires_piecewise_compilation() +- and mode is CompilationMode.VLLM_COMPILE ++ _name(mode) != REQUIRED_COMPILATION_MODE_NAME ++ or _name(cudagraph_mode) != REQUIRED_CUDAGRAPH_MODE_NAME + ): + raise RuntimeError( +- f"e2-guard[{role}]: cudagraph_mode={cudagraph_mode} needs piecewise " +- "capture, compilation mode is VLLM_COMPILE and breakable cudagraphs " +- "are off -- but vLLM PR #55272 removed every @support_torch_compile " +- "wrapper from vllm/models/qwen4_exp/nvidia/, so " +- "has_compiled_submodule() is False and capture will raise at " +- "vllm/v1/worker/gpu/cudagraph_utils.py:605. Executable candidates: " +- "(a) breakable off + -O0/CompilationMode.NONE, which normalises " +- "FULL_AND_PIECEWISE to FULL_DECODE_ONLY (changed prefill " +- "behaviour -- must be qualified); (b) breakable off + an explicit " +- "cudagraph_mode=FULL or NONE; (c) breakable PIECEWISE, i.e. " +- f"{BREAKABLE_ENV}=1, which is a different identity. See " +- "the graph-profile choice." ++ f"e2-guard[{role}]: resolved compilation_mode={_name(mode)} " ++ f"cudagraph_mode={_name(cudagraph_mode)}; the qualified profile is " ++ f"compilation_mode={REQUIRED_COMPILATION_MODE_NAME} " ++ f"cudagraph_mode={REQUIRED_CUDAGRAPH_MODE_NAME} (compile-profile guard). Check " ++ "-O / --compilation-config / --enforce-eager / TORCH_COMPILE_DISABLE, " ++ f"or set {ALLOW_BREAKABLE_ENV}=1 for a deliberate non-production boot." + ) + +- # 4. Optional exact pin of the resolved mode, for a launch profile. ++ # 4. The graph breaks that make PIECEWISE correct on this model. ++ if splitting_ops is None: ++ raise RuntimeError( ++ f"e2-guard[{role}]: compilation_config.splitting_ops is unresolved; " ++ "cannot confirm the Qwen4Exp graph breaks." ++ ) ++ missing = [op for op in REQUIRED_SPLITTING_OPS if op not in splitting_ops] ++ if missing: ++ raise RuntimeError( ++ f"e2-guard[{role}]: splitting_ops is missing {missing}. Each of " ++ "these reads per-step request metadata (or forks/joins a side " ++ "stream) and would be frozen into a piecewise CUDA graph." ++ ) ++ ++ # 5. Optional exact pin of the resolved mode, for a launch profile. + expected = os.environ.get(EXPECTED_MODE_ENV, "").strip() +- if expected: +- actual = getattr(cudagraph_mode, "name", str(cudagraph_mode)) +- if actual != expected: +- raise RuntimeError( +- f"e2-guard[{role}]: resolved cudagraph_mode={actual} but " +- f"{EXPECTED_MODE_ENV}={expected}. The launch profile and the " +- "engine disagree about the graph identity." +- ) ++ if expected and _name(cudagraph_mode) != expected: ++ raise RuntimeError( ++ f"e2-guard[{role}]: resolved cudagraph_mode={_name(cudagraph_mode)} " ++ f"but {EXPECTED_MODE_ENV}={expected}. The launch profile and the " ++ "engine disagree about the graph identity." ++ ) + + logger.info( + "e2-guard[%s]: architectures=%s breakable=OFF compilation_mode=%s " +- "cudagraph_mode=%s expected=%s. Recorded, not certified -- this guard " +- "checks startup states, not that the profile was qualified.", ++ "cudagraph_mode=%s splitting_ops=+%d required. Recorded, not " ++ "certified -- this checks startup state, not qualification.", + role, + architectures, +- mode, +- cudagraph_mode, +- expected or "", ++ _name(mode), ++ _name(cudagraph_mode), ++ len(REQUIRED_SPLITTING_OPS), ++ ) ++ ++ ++def assert_compiled_submodule(module, vllm_config, *, role: str = "target") -> None: ++ """Require a live ``@support_torch_compile`` wrapper on ``module``. ++ ++ This is the property ``has_compiled_submodule()`` checks before PIECEWISE ++ capture (``gpu/cudagraph_utils.py:132``, ``:605``). Checking it at ++ construction turns a mid-capture failure into a startup one. ++ """ ++ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper ++ ++ compiled = isinstance(module, TorchCompileWithNoGuardsWrapper) and not getattr( ++ module, "do_not_compile", True ++ ) ++ if compiled: ++ return ++ detail = ( ++ f"{type(module).__name__} is_wrapper=" ++ f"{isinstance(module, TorchCompileWithNoGuardsWrapper)} " ++ f"do_not_compile={getattr(module, 'do_not_compile', '')}" ++ ) ++ if _opted_out(role, detail): ++ return ++ raise RuntimeError( ++ f"e2-guard[{role}]: {detail}: no live torch.compile wrapper, so " ++ "PIECEWISE capture would raise at gpu/cudagraph_utils.py:605. The " ++ "@support_torch_compile decorator (compile-profile guard) is missing or disabled." ++ ) ++ ++ ++# -------------------------------------------------------------------------- ++# Resolved runtime graph modes. ++# -------------------------------------------------------------------------- ++ ++#: Target architectures whose V2 capture is checked. The construction guard ++#: above lives in the NVIDIA Qwen4Exp model; this tree targets sm_86 only. ++QWEN4EXP_TARGET_ARCHITECTURES = frozenset( ++ {"Qwen4ExpForCausalLM", "Qwen4ExpForConditionalGeneration"} ++) ++ ++#: What each V2 cudagraph manager may hold immediately before capture. ++#: The construction guard checks ``compilation_config.cudagraph_mode``; the ++#: V2 runner later calls ``resolve_cudagraph_mode_and_sizes`` (and, with ++#: adaptive verification, ``resolve_adaptive_cudagraph_mode``), which may ++#: downgrade to FULL_DECODE_ONLY / PIECEWISE / NONE or drop the piecewise half ++#: (attention CG support, spec-decode uniform batches, no compiled submodule). ++#: No such downgrade is allowed for the target. The drafter's managers are ++#: derived from the target's resolved mode by the speculator itself: ++#: * autoregressive / MTP (``AutoRegressiveSpeculator.init_cudagraph_manager``): ++#: prefill = the target's mode unchanged; decode = FULL_DECODE_ONLY ++#: ("PIECEWISE cudagraphs are not supported for draft decodes"); ++#: * multi-module MTP (``MultiModuleMTPSpeculator``): one manager, ++#: FULL_AND_PIECEWISE -> FULL_DECODE_ONLY by design. ++#: Those decode-specific modes are the only drafter deviations accepted. ++EXPECTED_RESOLVED_MODES = { ++ "target": REQUIRED_CUDAGRAPH_MODE_NAME, ++ "draft.prefill_cudagraph_manager": REQUIRED_CUDAGRAPH_MODE_NAME, ++ "draft.decode_cudagraph_manager": "FULL_DECODE_ONLY", ++ "draft.cudagraph_manager": "FULL_DECODE_ONLY", ++} ++_DRAFT_MANAGER_ATTRS = ( ++ "prefill_cudagraph_manager", ++ "decode_cudagraph_manager", ++ "cudagraph_manager", ++) ++ ++ ++def assert_resolved_cudagraph_modes( ++ vllm_config, target_manager, speculator=None ++) -> None: ++ """Fail loud if the V2 runner is about to capture a downgraded mode. ++ ++ Called from ``GPUModelRunner.capture_model`` (``vllm/v1/worker/gpu``) ++ before any capture, including the ``profile_only`` capture and the ++ "nothing to capture" early return, so a resolver downgrade to NONE cannot ++ skip it. Applies only to a Qwen4Exp target; any other model is untouched. ++ ``VLLM_E2_ALLOW_BREAKABLE=1`` skips it, logged, like the other graph ++ guards. ++ """ ++ model_config = getattr(vllm_config, "model_config", None) ++ architectures = list(getattr(model_config, "architectures", None) or []) ++ if not QWEN4EXP_TARGET_ARCHITECTURES & set(architectures): ++ return ++ ++ observed: dict[str, str] = {} ++ if target_manager is None: ++ observed["target"] = "" ++ else: ++ observed["target"] = _name(getattr(target_manager, "cudagraph_mode", None)) ++ if speculator is not None: ++ found = False ++ for attr in _DRAFT_MANAGER_ATTRS: ++ if not hasattr(speculator, attr): ++ continue ++ found = True ++ manager = getattr(speculator, attr) ++ observed[f"draft.{attr}"] = ( ++ "" ++ if manager is None ++ else _name(getattr(manager, "cudagraph_mode", None)) ++ ) ++ if not found: ++ observed["draft"] = f"" ++ ++ detail = " ".join(f"{key}={value}" for key, value in observed.items()) ++ if _opted_out("runtime", detail): ++ return ++ wrong = { ++ key: value ++ for key, value in observed.items() ++ if EXPECTED_RESOLVED_MODES.get(key) != value ++ } ++ if wrong: ++ expected = {key: EXPECTED_RESOLVED_MODES.get(key, "") for key in wrong} ++ raise RuntimeError( ++ f"e2-guard[runtime]: resolved cudagraph modes {wrong} differ from " ++ f"the qualified profile {expected} ({detail}). The V2 runner's " ++ "resolver downgraded the graph identity after the construction " ++ "guard passed (attention CG support, spec-decode uniform batch, " ++ "or no compiled submodule); see the resolver's warning above. Set " ++ f"{ALLOW_BREAKABLE_ENV}=1 only for a deliberate non-production boot." ++ ) ++ logger.info( ++ "e2-guard[runtime]: resolved cudagraph modes before capture: %s.", detail + ) + + +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 2d80a37..67ba37f 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -66,6 +66,7 @@ from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture + from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib + from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d + from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput ++from vllm.v1.e2_guards import assert_resolved_cudagraph_modes + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + MambaSpec, +@@ -1069,6 +1070,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): + @torch.inference_mode() + def capture_model(self, *, profile_only: bool = False) -> int: + assert self.cudagraph_manager is not None ++ # Fork (0.30 rebase): the construction guard saw the configured ++ # mode; this asserts what the resolver actually handed the target and ++ # drafter managers, before any capture and before the early return. ++ assert_resolved_cudagraph_modes( ++ self.vllm_config, self.cudagraph_manager, self.speculator ++ ) + capture_encoder = ( + self.model_state.supports_mm_inputs + and self.model_state.encoder_runner.has_cudagraph() +-- +2.43.0 + diff --git a/patches/0035-ple-the-pull-wait-copy-op-is-a-graph-splitting-op.patch b/patches/0035-ple-the-pull-wait-copy-op-is-a-graph-splitting-op.patch new file mode 100644 index 0000000..97b1e3c --- /dev/null +++ b/patches/0035-ple-the-pull-wait-copy-op-is-a-graph-splitting-op.patch @@ -0,0 +1,67 @@ +From 236b01309da7e3d934cd9876cf929b2ea68c8e54 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 07:16:45 -0400 +Subject: [PATCH 035/103] ple: the pull wait/copy op is a graph splitting op + +With the wait/copy inside a compiled piece, Inductor's post-grad +auto-functionalization decomposition failed on the stage-0 ranks. +--- + tests/v1/test_ple_pull_compile_profile.py | 14 +++++++++----- + vllm/config/compilation.py | 6 ++++++ + vllm/v1/e2_guards.py | 1 + + 3 files changed, 16 insertions(+), 5 deletions(-) + +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index f35926d..4f7b175 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -433,8 +433,12 @@ def test_wait_copy_op_keeps_its_fake_impl() -> None: + ) + + +-def test_wait_copy_is_not_a_splitting_op() -> None: +- """Add it to splitting_ops only if a reference-host experiment shows +- piecewise capture of the pull wait/copy is unsafe. Pinned so a change is +- deliberate.""" +- assert "vllm::ple_offload_wait_copy" not in CompilationConfig._attention_ops ++def test_wait_copy_is_a_splitting_op() -> None: ++ """Resolved on the reference host: inside an ++ Inductor graph the mutating wait/copy failed post-grad ("auto_functionalized ++ was not removed") on both stage-0 ranks, so it is a split point and the ++ guard requires it.""" ++ from vllm.v1.e2_guards import REQUIRED_SPLITTING_OPS ++ ++ assert "vllm::ple_offload_wait_copy" in CompilationConfig._attention_ops ++ assert "vllm::ple_offload_wait_copy" in REQUIRED_SPLITTING_OPS +diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py +index 8cc2973..f11bea4 100644 +--- a/vllm/config/compilation.py ++++ b/vllm/config/compilation.py +@@ -783,6 +783,12 @@ class CompilationConfig: + "vllm::qwen4_exp_ple_finalize_prefetch", + "vllm::qwen4_exp_ple_short_conv", + "vllm::qwen4_exp_qsa_with_output", ++ # Fork (pull transport with the compile-profile guard, found on the reference host): the PLE pull transport's wait+DMA ++ # mutates a module-owned buffer; inside an Inductor graph it failed ++ # post-grad with "auto_functionalized was not removed" on the stage-0 ++ # ranks. As a split point it runs between pieces, like attention. ++ "vllm::ple_offload_wait_copy", ++ "vllm::ple_offload_wait", + "vllm::linear_attention", + "vllm::qwen_gdn_attention_core", + "vllm::qwen_gdn_attention_core_fused_norm_packed", +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +index 1917885..5f0ed88 100644 +--- a/vllm/v1/e2_guards.py ++++ b/vllm/v1/e2_guards.py +@@ -61,6 +61,7 @@ REQUIRED_SPLITTING_OPS = ( + "vllm::qwen4_exp_ple_finalize_prefetch", + "vllm::qwen_gdn_attention_core", + "vllm::qwen_gdn_attention_core_fused_norm_packed", ++ "vllm::ple_offload_wait_copy", + ) + + #: Upstream v0.30.0 auto-enabled breakable cudagraphs for these; this tree +-- +2.43.0 + diff --git a/patches/0036-kv-restore-PP-fix-54793-no-KV-tensors-for-another-st.patch b/patches/0036-kv-restore-PP-fix-54793-no-KV-tensors-for-another-st.patch new file mode 100644 index 0000000..08012f1 --- /dev/null +++ b/patches/0036-kv-restore-PP-fix-54793-no-KV-tensors-for-another-st.patch @@ -0,0 +1,276 @@ +From 2179328624d5718a369f599920d285040c572a9d Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 07:33:48 -0400 +Subject: [PATCH 036/103] kv: restore PP fix #54793 -- no KV tensors for + another stage's layers + +A KV cache group with no layer on a pipeline stage was still given tensors +there, and allocate_kv_cache raised StopIteration on the second stage. The +fix was lost in the v0.30.0 rebase. Test drives the real get_kv_cache_configs +with the production layout. +--- + tests/v1/test_pp_kv_projection.py | 231 ++++++++++++++++++++++++++++++ + vllm/v1/core/kv_cache_utils.py | 9 +- + 2 files changed, 239 insertions(+), 1 deletion(-) + create mode 100644 tests/v1/test_pp_kv_projection.py + +diff --git a/tests/v1/test_pp_kv_projection.py b/tests/v1/test_pp_kv_projection.py +new file mode 100644 +index 0000000..11cb4ef +--- /dev/null ++++ b/tests/v1/test_pp_kv_projection.py +@@ -0,0 +1,231 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Observed on a PP=2 boot: PP1 workers died in ``allocate_kv_cache`` ++(``vllm/v1/worker/utils.py:452``) with StopIteration. ++ ++Root cause: ``_project_kv_cache_groups_to_worker`` keeps the GLOBAL ++``UniformTypeKVCacheSpecs`` for a group that has no layer on a PP stage, and ++v0.30.0's ``get_kv_cache_config_from_groups`` builds tensors from ++``group_spec.kv_cache_specs`` without checking ``group.layer_names``. On ++Flash-Next the layer-2 PLE short-conv state is a one-layer group of its own ++(SHORT_CONV MambaSpec, PP0 only), so PP1 received a tensor for ++``language_model.model.layers.2.ple`` that is in none of its groups. The E2 ++base carried the fix (#54793); the v0.30.0 rebase dropped it. ++ ++CPU-only. Specs are built from the model's shipped config (48 layers, QSA every ++4th, TP2, PP 25/23, MTP K=3, fp8 KV, block 3200). FAILING CONTROL: on ++an earlier commit (filter absent) ``test_flash_next_pp1_every_tensor_is_grouped`` and ++``test_empty_projected_uniform_group_gets_no_tensor`` fail; the no-PLE control ++shows the PLE group is the trigger. ++""" ++ ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator ++from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum ++from vllm.v1.core import kv_cache_utils as ku ++from vllm.v1.kv_cache_interface import ( ++ CircularBufferSpec, ++ FullAttentionSpec, ++ KVCacheGroupSpec, ++ KVCacheLayout, ++ MambaSpec, ++ MLAAttentionSpec, ++ UniformTypeKVCacheSpecs, ++ get_kv_quant_mode, ++) ++ ++TP = 2 ++BLOCK = 3200 # "Setting attention block size to 3200 tokens" (both stages) ++K = 3 # MTP num_speculative_tokens ++ATTN_PAGE = BLOCK * 1 * 256 * 2 # fp8 K+V, 1 KV head per TP rank, head_dim 256 ++ ++_MAIN = FullAttentionSpec( ++ block_size=BLOCK, ++ num_kv_heads=2 // TP, ++ head_size=256, ++ head_size_v=256, ++ dtype=torch.uint8, ++ kv_quant_mode=get_kv_quant_mode("fp8_e4m3"), ++) ++# QSAKeyStateCache: ring = 4 * cdiv(4 + K, 4) = 8, head 128 + MRoPE tail 12. ++_RING = CircularBufferSpec( ++ block_size=8, num_kv_heads=1, head_size=140, head_size_v=0, dtype=torch.bfloat16 ++) ++# QSACompressedKeyCache: bf16 indexer, compress ratio 4. ++_COMP = MLAAttentionSpec( ++ block_size=BLOCK, ++ num_kv_heads=1, ++ head_size=128, ++ dtype=torch.bfloat16, ++ tokens_per_state=4, ++) ++_GDN = MambaSpec( ++ shapes=tuple( ++ MambaStateShapeCalculator.gated_delta_net_state_shape( ++ TP, 16, 48, 128, 128, 4, K ++ ) ++ ), ++ dtypes=(torch.bfloat16, torch.float32), ++ block_size=BLOCK, ++ page_size_padded=ATTN_PAGE, ++ mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ++ num_speculative_blocks=K, ++) ++# Qwen4ExpPLELayer: hc_hidden 2560*4, conv_state_len (4-1)*3, TP-replicated. ++_PLE = MambaSpec( ++ shapes=tuple( ++ MambaStateShapeCalculator.short_conv_state_shape( ++ tp_world_size=1, intermediate_size=10240, conv_kernel=10, num_spec=K ++ ) ++ ), ++ dtypes=(torch.bfloat16,), ++ block_size=BLOCK, ++ page_size_padded=ATTN_PAGE, ++ mamba_type=MambaAttentionBackendEnum.SHORT_CONV, ++ tp_replicated=True, ++ num_speculative_blocks=K, ++) ++ ++ ++def _qsa(prefix: str, specs: dict) -> None: ++ # Registration order in Qwen4ExpQSAAttention.__init__: indexer caches first. ++ specs[f"{prefix}.self_attn.indexer.raw_key_cache"] = _RING ++ specs[f"{prefix}.self_attn.indexer.compressed_key_cache"] = _COMP ++ specs[f"{prefix}.self_attn.attn"] = _MAIN ++ ++ ++def _stage(lo: int, hi: int, *, mtp: bool, ple: bool = True) -> dict: ++ specs: dict = {} ++ for i in range(lo, hi): ++ prefix = f"language_model.model.layers.{i}" ++ if i == 2 and ple: # ple_layer_ids=[2]; built before the mixer ++ specs[f"{prefix}.ple"] = _PLE ++ if i % 4 == 3: # full_attention_interval=4 -> QSA ++ _qsa(prefix, specs) ++ else: ++ specs[f"{prefix}.linear_attn"] = _GDN ++ if mtp: ++ _qsa("mtp.layers.48", specs) ++ return specs ++ ++ ++def _config(): ++ cache = SimpleNamespace( ++ kv_cache_layout="BLNHC", # reference host: "Using BLNHC KV cache layout." ++ num_gpu_blocks_override=None, ++ prefix_cache_retention_interval=0, ++ mamba_cache_mode="none", ++ get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLNHC, ++ ) ++ return SimpleNamespace( ++ cache_config=cache, ++ speculative_config=SimpleNamespace( ++ num_speculative_tokens=K, ++ method="mtp", ++ use_multi_module_mtp=lambda: False, ++ use_eagle=lambda: True, ++ use_eagle_block_drop=lambda: True, ++ ), ++ scheduler_config=SimpleNamespace( ++ disable_hybrid_kv_cache_manager=False, max_num_encoder_input_tokens=0 ++ ), ++ attention_config=SimpleNamespace(hisparse_config=None), ++ model_config=SimpleNamespace( ++ original_max_model_len=262144, ++ max_model_len=262144, ++ hf_config=SimpleNamespace(model_type="qwen4_exp"), ++ get_total_num_hidden_layers=lambda: 48, ++ ), ++ device_config=SimpleNamespace(device_type="cuda"), ++ parallel_config=SimpleNamespace( ++ pipeline_parallel_size=2, decode_context_parallel_size=1 ++ ), ++ ) ++ ++ ++def _worker_configs(*, ple: bool = True): ++ pp0 = _stage(0, 25, mtp=False, ple=ple) # VLLM_PP_LAYER_PARTITION=25,23 ++ pp1 = _stage(25, 48, mtp=True, ple=ple) ++ specs = [pp0, pp0, pp1, pp1] # TP2 x PP2 ++ configs = ku.get_kv_cache_configs(_config(), specs, [4_100_000_000] * 4) ++ return specs, configs ++ ++ ++def _allocate_lookup(config) -> None: ++ """Verbatim group lookup of allocate_kv_cache (vllm/v1/worker/utils.py:452). ++ ++ The module itself needs triton at import, so the lookup is replayed here. ++ """ ++ for tensor in config.kv_cache_tensors: ++ layer_name = tensor.layers[0] ++ next( ++ (group_id, group) ++ for group_id, group in enumerate(config.kv_cache_groups) ++ if layer_name in group.layer_names ++ ) ++ ++ ++def test_flash_next_pp1_has_an_empty_uniform_group(): ++ """Guard that the shape really exercises the path (else the test below ++ could not fail): PP1 projects the PLE group to zero layers while keeping ++ the global UniformTypeKVCacheSpecs.""" ++ _, configs = _worker_configs() ++ for rank in (2, 3): ++ empty = [g for g in configs[rank].kv_cache_groups if not g.layer_names] ++ assert len(empty) == 1 ++ assert isinstance(empty[0].kv_cache_spec, UniformTypeKVCacheSpecs) ++ assert list(empty[0].kv_cache_spec.kv_cache_specs) == [ ++ "language_model.model.layers.2.ple" ++ ] ++ for rank in (0, 1): ++ assert all(g.layer_names for g in configs[rank].kv_cache_groups) ++ ++ ++def test_flash_next_pp1_every_tensor_is_grouped(): ++ specs, configs = _worker_configs() ++ for rank, config in enumerate(configs): ++ tensor_layers = [n for t in config.kv_cache_tensors for n in t.layers] ++ assert sorted(tensor_layers) == sorted(specs[rank]), f"rank {rank}" ++ _allocate_lookup(config) # StopIteration on PP1 before the fix ++ ++ ++def test_no_ple_control_has_no_empty_group(): ++ """Control: remove the PLE state and no group projects empty, so the ++ old code would also have passed -- the PLE group is the trigger.""" ++ _, configs = _worker_configs(ple=False) ++ for config in configs: ++ assert all(g.layer_names for g in config.kv_cache_groups) ++ _allocate_lookup(config) ++ ++ ++def test_empty_projected_uniform_group_gets_no_tensor(): ++ """Minimal form (the #54793 unit test, config stand-in for offline runs).""" ++ spec = _MAIN ++ groups = [ ++ KVCacheGroupSpec( ++ [], ++ UniformTypeKVCacheSpecs( ++ block_size=spec.block_size, kv_cache_specs={"remote.layer": spec} ++ ), ++ ), ++ KVCacheGroupSpec(["local.layer"], spec), ++ ] ++ config = ku.get_kv_cache_config_from_groups( ++ _config(), groups, available_memory=spec.page_size_bytes * 2 ++ ) ++ assert [t.layers for t in config.kv_cache_tensors] == [["local.layer"]] ++ ++ ++@pytest.mark.parametrize("rank", [0, 1]) ++def test_pp0_layout_unchanged_by_filter(rank): ++ """Groups that are non-empty on a worker keep every layer: the filter only ++ drops layers the worker does not own.""" ++ specs, configs = _worker_configs() ++ grouped = {n for g in configs[rank].kv_cache_groups for n in g.layer_names} ++ assert grouped == set(specs[rank]) ++ tensor_layers = [n for t in configs[rank].kv_cache_tensors for n in t.layers] ++ assert sorted(tensor_layers) == sorted(grouped) +diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py +index 28478b6..d1f827d 100644 +--- a/vllm/v1/core/kv_cache_utils.py ++++ b/vllm/v1/core/kv_cache_utils.py +@@ -1768,7 +1768,14 @@ def get_kv_cache_config_from_groups( + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): +- layers_by_spec[spec].append(layer_name) ++ # Fork (E2 PP fix #54793, dropped in the v0.30.0 rebase, restored ++ # the vLLM 0.30 rebase): _project_kv_cache_groups_to_worker keeps the GLOBAL ++ # UniformTypeKVCacheSpecs for a group with no layer on this PP ++ # stage, so without this filter the worker gets a tensor for ++ # another stage's layer and allocate_kv_cache's group lookup ++ # raises StopIteration (Flash-Next PP1: the layer-2 PLE state). ++ if layer_name in group.layer_names: ++ layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + +-- +2.43.0 + diff --git a/patches/0037-kv-restore-PP-fix-54795-PP-stages-need-only-share-a-.patch b/patches/0037-kv-restore-PP-fix-54795-PP-stages-need-only-share-a-.patch new file mode 100644 index 0000000..fb36060 --- /dev/null +++ b/patches/0037-kv-restore-PP-fix-54795-PP-stages-need-only-share-a-.patch @@ -0,0 +1,104 @@ +From 5d513783b79de125c9ffb96d928919c706777f3e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 07:36:14 -0400 +Subject: [PATCH 037/103] kv: restore PP fix #54795 -- PP stages need only + share a KV layout + +--- + tests/v1/test_pp_kv_projection.py | 35 ++++++++++++++++++++++++++++- + vllm/v1/attention/backends/utils.py | 23 ++++++++++++++----- + 2 files changed, 51 insertions(+), 7 deletions(-) + +diff --git a/tests/v1/test_pp_kv_projection.py b/tests/v1/test_pp_kv_projection.py +index 11cb4ef..2c699ab 100644 +--- a/tests/v1/test_pp_kv_projection.py ++++ b/tests/v1/test_pp_kv_projection.py +@@ -1,6 +1,6 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +-"""Observed on a PP=2 boot: PP1 workers died in ``allocate_kv_cache`` ++"""A reference-host boot of the vLLM 0.30 rebase: PP1 workers died in ``allocate_kv_cache`` + (``vllm/v1/worker/utils.py:452``) with StopIteration. + + Root cause: ``_project_kv_cache_groups_to_worker`` keeps the GLOBAL +@@ -229,3 +229,36 @@ def test_pp0_layout_unchanged_by_filter(rank): + assert grouped == set(specs[rank]) + tensor_layers = [n for t in configs[rank].kv_cache_tensors for n in t.layers] + assert sorted(tensor_layers) == sorted(grouped) ++ ++ ++# --- #54795 (restored with #54793): PP stages may report different layout sets --- ++ ++def _layout_cfg(): ++ from types import SimpleNamespace ++ ++ cache = SimpleNamespace(kv_cache_layout=None) ++ cache.get_resolved_kv_cache_layout = lambda: cache.kv_cache_layout ++ return SimpleNamespace(cache_config=cache) ++ ++ ++def test_resolve_kv_cache_layout_intersects_pp_stage_preferences(monkeypatch): ++ import vllm.v1.attention.backends.utils as butils ++ from vllm.v1.kv_cache_interface import KVCacheLayout ++ ++ monkeypatch.setattr("vllm.envs.VLLM_KV_CACHE_LAYOUT", None) ++ monkeypatch.setattr(butils, "get_kv_connector_cache_layout", lambda cfg: None) ++ cfg = _layout_cfg() ++ layout = butils.resolve_kv_cache_layout( ++ cfg, [["LBNHC", "LBHNC", "BLNHC", "BLHNC"], ["BLHNC", "BLNHC"]] ++ ) ++ assert layout is KVCacheLayout.BLNHC ++ assert cfg.cache_config.kv_cache_layout == "BLNHC" ++ ++ ++def test_resolve_kv_cache_layout_rejects_disjoint_pp_stage_preferences(monkeypatch): ++ import vllm.v1.attention.backends.utils as butils ++ ++ monkeypatch.setattr("vllm.envs.VLLM_KV_CACHE_LAYOUT", None) ++ monkeypatch.setattr(butils, "get_kv_connector_cache_layout", lambda cfg: None) ++ with pytest.raises(ValueError, match="share no supported KV cache layout"): ++ butils.resolve_kv_cache_layout(_layout_cfg(), [["LBNHC"], ["BLHNC"]]) +diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py +index e656343..c2e7adf 100644 +--- a/vllm/v1/attention/backends/utils.py ++++ b/vllm/v1/attention/backends/utils.py +@@ -259,8 +259,9 @@ def resolve_kv_cache_layout( + """Resolve one KV cache layout for the whole model. + + Runs once in the engine core. Every worker reports the layouts its backends +- support, most preferred first (``get_supported_kv_cache_layouts``); all +- ranks run the same backends, so their lists must agree. Specs mixing HNC ++ support, most preferred first (``get_supported_kv_cache_layouts``); ranks ++ may expose different sets under pipeline parallelism, but their lists must ++ share at least one layout. Specs mixing HNC + shapes narrow the candidates to block-compact layouts. An explicit + ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails, + with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the +@@ -277,10 +278,20 @@ def resolve_kv_cache_layout( + assert supported_layouts and all(supported_layouts), ( + "No worker reported supported KV cache layouts." + ) +- assert all(names == supported_layouts[0] for names in supported_layouts[1:]), ( +- f"Workers disagree on supported KV cache layouts: {supported_layouts}." +- ) +- candidates = [_layout_from_name(name) for name in supported_layouts[0]] ++ # Fork: restore PP fix #54795 from the E2 base, dropped ++ # in the v0.30.0 rebase. Under pipeline parallelism the stages run ++ # different backends (PP0 carries the PLE short-conv, PP1 the drafter), so ++ # their lists need only share a layout, not be identical. ++ common_layout_names = [ ++ name ++ for name in supported_layouts[0] ++ if all(name in names for names in supported_layouts[1:]) ++ ] ++ if not common_layout_names: ++ raise ValueError( ++ f"Workers share no supported KV cache layout: {supported_layouts}." ++ ) ++ candidates = [_layout_from_name(name) for name in common_layout_names] + + # A block-compact layout means the block is densely packed in memory, so any mix of + # specs can re-interpret HNC with different sizes as long as the total number of +-- +2.43.0 + diff --git a/patches/0038-e2_guards-the-draft-guard-reports-the-draft-model-s-.patch b/patches/0038-e2_guards-the-draft-guard-reports-the-draft-model-s-.patch new file mode 100644 index 0000000..0e91528 --- /dev/null +++ b/patches/0038-e2_guards-the-draft-guard-reports-the-draft-model-s-.patch @@ -0,0 +1,85 @@ +From 0ba64cc82a5962faf443b63e821399a82e77997e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 07:53:15 -0400 +Subject: [PATCH 038/103] e2_guards: the draft guard reports the draft model's + own architecture + +--- + tests/v1/test_ple_pull_compile_profile.py | 24 +++++++++++++++++++++++ + vllm/v1/e2_guards.py | 20 ++++++++++++++++--- + 2 files changed, 41 insertions(+), 3 deletions(-) + +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index 4f7b175..d5bf87e 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -442,3 +442,27 @@ def test_wait_copy_is_a_splitting_op() -> None: + + assert "vllm::ple_offload_wait_copy" in CompilationConfig._attention_ops + assert "vllm::ple_offload_wait_copy" in REQUIRED_SPLITTING_OPS ++ ++ ++def test_draft_guard_reports_the_draft_architecture() -> None: ++ """On the reference host, the draft guard line printed the TARGET ++ architecture (the drafter is built with the target's VllmConfig), so S ++ could not prove the drafter's identity. The draft role must read ++ speculative_config.draft_model_config; the target role is unchanged.""" ++ from types import SimpleNamespace ++ ++ from vllm.v1 import e2_guards ++ ++ cfg = SimpleNamespace( ++ model_config=SimpleNamespace(architectures=["Qwen4ExpForConditionalGeneration"]), ++ compilation_config=SimpleNamespace( ++ mode="VLLM_COMPILE", cudagraph_mode="FULL_AND_PIECEWISE", splitting_ops=[] ++ ), ++ speculative_config=SimpleNamespace( ++ draft_model_config=SimpleNamespace(architectures=["Qwen4ExpMTP"]) ++ ), ++ ) ++ assert e2_guards._profile(cfg, "target")[0] == ["Qwen4ExpForConditionalGeneration"] ++ assert e2_guards._profile(cfg, "draft")[0] == ["Qwen4ExpMTP"] ++ no_spec = SimpleNamespace(**{**vars(cfg), "speculative_config": None}) ++ assert e2_guards._profile(no_spec, "draft")[0] == ["Qwen4ExpForConditionalGeneration"] +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +index 5f0ed88..7624ae1 100644 +--- a/vllm/v1/e2_guards.py ++++ b/vllm/v1/e2_guards.py +@@ -100,9 +100,23 @@ def _name(value) -> str: + return getattr(value, "name", str(value)) + + +-def _profile(vllm_config) -> tuple[list[str], object, object, list[str] | None]: +- """Resolved (architectures, mode, cudagraph mode, splitting ops).""" ++def _profile( ++ vllm_config, role: str = "target" ++) -> tuple[list[str], object, object, list[str] | None]: ++ """Resolved (architectures, mode, cudagraph mode, splitting ops). ++ ++ For ``role="draft"`` the architectures are the DRAFT model's ++ (``speculative_config.draft_model_config``): the drafter is constructed ++ with the target's VllmConfig, whose model_config names the target. The reference host ++ logged the target architecture on the draft line before this fix. ++ The graph profile is the engine's single CompilationConfig for both roles. ++ """ + model_config = getattr(vllm_config, "model_config", None) ++ if role == "draft": ++ spec = getattr(vllm_config, "speculative_config", None) ++ draft_config = getattr(spec, "draft_model_config", None) ++ if draft_config is not None: ++ model_config = draft_config + if model_config is None: + raise RuntimeError("e2-guard: vllm_config has no model_config to check") + compilation_config = getattr(vllm_config, "compilation_config", None) +@@ -137,7 +151,7 @@ def assert_graph_profile_is_executable(vllm_config, *, role: str = "target") -> + ``Qwen4ExpMTP`` builds ``Qwen4ExpMultiTokenPredictor`` rather than + ``Qwen4ExpForCausalLM``, from a config rebuilt in ``mtp.py``. + """ +- architectures, mode, cudagraph_mode, splitting_ops = _profile(vllm_config) ++ architectures, mode, cudagraph_mode, splitting_ops = _profile(vllm_config, role) + detail = ( + f"architectures={architectures} breakable={os.environ.get(BREAKABLE_ENV)} " + f"compilation_mode={_name(mode)} cudagraph_mode={_name(cudagraph_mode)}" +-- +2.43.0 + diff --git a/patches/0039-ple-the-pull-op-resolves-live-pointers-by-layer-name.patch b/patches/0039-ple-the-pull-op-resolves-live-pointers-by-layer-name.patch new file mode 100644 index 0000000..17e383b --- /dev/null +++ b/patches/0039-ple-the-pull-op-resolves-live-pointers-by-layer-name.patch @@ -0,0 +1,188 @@ +From fc34b9371ed0900ff970fc7631d41a13bf49ee22 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:18:03 -0400 +Subject: [PATCH 039/103] ple: the pull op resolves live pointers by layer name + +The wait/copy op took host pointers as Python-int arguments, which were baked +into the AOT compile cache as the tracing process's addresses; the first +forward after a cache-load boot then failed. The op now carries only +process-stable values and resolves pointers from a per-process registry. +--- + .../ple_offload/test_transport_bounded_io.py | 2 +- + tests/v1/test_ple_pull_compile_profile.py | 68 +++++++++++++++++++ + .../layers/ple_offload_layer.py | 45 ++++++++++-- + 3 files changed, 107 insertions(+), 8 deletions(-) + +diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py +index 748fd46..5a2aa19 100644 +--- a/tests/v1/ple_offload/test_transport_bounded_io.py ++++ b/tests/v1/ple_offload/test_transport_bounded_io.py +@@ -275,7 +275,7 @@ def test_real_wait_copy_geometry_rejected_before_graph_publication( + + def capture_then_publish() -> None: + with manager.capture_ple_copy_geometry(desc): +- ple_layer_module._ple_offload_wait_copy_impl( ++ ple_layer_module._ple_offload_wait_copy_raw( + gpu_output_buffer, + hidden_states, + 1, +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index d5bf87e..d5c47a9 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -466,3 +466,71 @@ def test_draft_guard_reports_the_draft_architecture() -> None: + assert e2_guards._profile(cfg, "draft")[0] == ["Qwen4ExpMTP"] + no_spec = SimpleNamespace(**{**vars(cfg), "speculative_config": None}) + assert e2_guards._profile(no_spec, "draft")[0] == ["Qwen4ExpForConditionalGeneration"] ++ ++ ++def _import_ple_offload_layer(monkeypatch): ++ """Import the layer module offline without cuda-python: stub only the ++ two names it imports at module load (the tests never reach the driver).""" ++ import sys ++ import types ++ ++ try: ++ import cuda.bindings.driver # noqa: F401 ++ except ImportError: ++ cuda = types.ModuleType("cuda") ++ bindings = types.ModuleType("cuda.bindings") ++ driver = types.ModuleType("cuda.bindings.driver") ++ driver.CUstreamWaitValue_flags = types.SimpleNamespace( ++ CU_STREAM_WAIT_VALUE_EQ=types.SimpleNamespace(value=1) ++ ) ++ cuda.bindings = bindings ++ bindings.driver = driver ++ for name, mod in (("cuda", cuda), ("cuda.bindings", bindings), ++ ("cuda.bindings.driver", driver)): ++ monkeypatch.setitem(sys.modules, name, mod) ++ from vllm.model_executor.layers import ple_offload_layer ++ ++ return ple_offload_layer ++ ++ ++def test_wait_copy_op_carries_no_raw_pointers(monkeypatch) -> None: ++ """On the reference host, the first boot that LOADED the AOT compile ++ cache failed cuStreamWaitValue32 with CUDA_ERROR_INVALID_VALUE -- the flag ++ and host pointers were Python-int op arguments, specialised into the saved ++ graph as the previous process's addresses. The op must carry only ++ process-stable values and resolve the live pointers by layer name.""" ++ import inspect ++ ++ m = _import_ple_offload_layer(monkeypatch) ++ ++ params = list(inspect.signature(m._ple_offload_wait_copy_impl).parameters) ++ assert params == ["gpu_output_buffer", "hidden_states", "num_tokens", "layer_name"] ++ fake = list(inspect.signature(m._ple_offload_wait_copy_fake).parameters) ++ assert fake == params ++ for banned in ("flag_ptr", "host_ptr", "row_bytes"): ++ assert banned not in params ++ ++ ++def test_wait_copy_resolves_live_pointers_by_layer_name(monkeypatch) -> None: ++ from types import SimpleNamespace ++ ++ import torch ++ ++ m = _import_ple_offload_layer(monkeypatch) ++ seen = {} ++ monkeypatch.setattr(m, "_ple_offload_wait_copy_raw", ++ lambda buf, hs, flag, host, rb, n, name: seen.update( ++ flag=flag, host=host, rb=rb, n=n, name=name)) ++ layer = SimpleNamespace(_sem=SimpleNamespace(flag_ptr=0x1000), ++ _host_output_ptr=0x2000, _host_row_bytes=16) ++ monkeypatch.setitem(m._PULL_LAYERS, "L2", layer) ++ buf = torch.empty(4, 4) ++ m._ple_offload_wait_copy_impl(buf, buf, 3, "L2") ++ assert seen == {"flag": 0x1000, "host": 0x2000, "rb": 16, "n": 3, "name": "L2"} ++ layer._sem.flag_ptr = 0x3000 # a new process / re-armed layer: no stale value ++ m._ple_offload_wait_copy_impl(buf, buf, 3, "L2") ++ assert seen["flag"] == 0x3000 ++ import pytest ++ ++ with pytest.raises(RuntimeError, match="no armed pull layer"): ++ m._ple_offload_wait_copy_impl(buf, buf, 3, "missing") +diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py +index d1f6683..ff581b2 100644 +--- a/vllm/model_executor/layers/ple_offload_layer.py ++++ b/vllm/model_executor/layers/ple_offload_layer.py +@@ -248,11 +248,46 @@ direct_register_custom_op( + # and the consumers read device memory as before. The flag is NOT reset here: + # ``release_outputs`` resets it after the forward (unchanged semantics), so + # capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``. +-# Pointers are passed as ints so no CPU tensor enters the compiled graph. ++# No CPU tensor and no raw address enters the compiled graph: the op resolves ++# the live pointers by layer name (see _PULL_LAYERS). + # --------------------------------------------------------------------------- + + ++# Fork (found on the reference host): the host flag / host output pointers and ++# the row size are NOT op arguments any more. As Python ints they were ++# specialised into the compiled graph as constants, so a boot that LOADED the ++# AOT compile cache (no Dynamo pass) waited on the previous process's addresses: ++# cuStreamWaitValue32 -> CUDA_ERROR_INVALID_VALUE on the first forward. The op ++# now carries only process-stable values (buffer, token count, layer name) and ++# resolves the live pointers from this registry on every call. ++_PULL_LAYERS: dict[str, "PleOffloadLayer"] = {} ++ ++ + def _ple_offload_wait_copy_impl( ++ gpu_output_buffer: torch.Tensor, ++ hidden_states: torch.Tensor, ++ num_tokens: int, ++ layer_name: str, ++) -> None: ++ """Resolve this process's live pointers for ``layer_name``, then wait+DMA.""" ++ layer = _PULL_LAYERS.get(layer_name) ++ if layer is None or layer._host_output_ptr is None: ++ raise RuntimeError( ++ f"PLE wait_copy: no armed pull layer named {layer_name!r} in this " ++ f"process (armed: {sorted(_PULL_LAYERS)})" ++ ) ++ _ple_offload_wait_copy_raw( ++ gpu_output_buffer, ++ hidden_states, ++ layer._sem.flag_ptr, ++ layer._host_output_ptr, ++ layer._host_row_bytes, ++ num_tokens, ++ layer_name, ++ ) ++ ++ ++def _ple_offload_wait_copy_raw( + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, + flag_ptr: int, +@@ -298,9 +333,6 @@ def _ple_offload_wait_copy_impl( + def _ple_offload_wait_copy_fake( + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, +- flag_ptr: int, +- host_ptr: int, +- row_bytes: int, + num_tokens: int, + layer_name: str, + ) -> None: +@@ -395,6 +427,8 @@ class PleOffloadLayer(nn.Module, ABC): + self._sem = semaphore + self._host_output_ptr = host_output_ptr + self._host_row_bytes = int(host_row_bytes) ++ if host_output_ptr is not None: ++ _PULL_LAYERS[self._ple_offload_name] = self + + def forward( + self, +@@ -409,9 +443,6 @@ class PleOffloadLayer(nn.Module, ABC): + torch.ops.vllm.ple_offload_wait_copy( + self._gpu_output_buffer, + hidden_states, +- self._sem.flag_ptr, +- self._host_output_ptr, +- self._host_row_bytes, + input_ids.shape[0], + self._ple_offload_name, + ) +-- +2.43.0 + diff --git a/patches/0040-topk-Python-reader-for-the-exact-restart-counters.patch b/patches/0040-topk-Python-reader-for-the-exact-restart-counters.patch new file mode 100644 index 0000000..5248809 --- /dev/null +++ b/patches/0040-topk-Python-reader-for-the-exact-restart-counters.patch @@ -0,0 +1,346 @@ +From 8be2e3ab6246beafa10aa7ed2de9591c1b4cfc37 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:31:06 -0400 +Subject: [PATCH 040/103] topk: Python reader for the exact-restart counters + +Reads persistent_topk_exact_counters into {slot: value}, logs one line on a +cadence, and degrades to a one-time warning when the op is absent. +--- + tests/kernels/test_top_k_per_row.py | 60 +++++++ + vllm/envs.py | 12 ++ + .../qwen4_exp/nvidia/ops/topk_counters.py | 167 ++++++++++++++++++ + vllm/v1/worker/gpu/model_runner.py | 13 ++ + 4 files changed, 252 insertions(+) + create mode 100644 vllm/models/qwen4_exp/nvidia/ops/topk_counters.py + +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index 70332b8..b0564e8 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1,11 +1,13 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++import re + + import numpy as np + import pytest + import torch + ++import vllm.envs as envs + from vllm.platforms import current_platform + from vllm.utils.torch_utils import set_random_seed + +@@ -2090,3 +2092,61 @@ def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: + + with pytest.raises(Exception): + op(ok, 1, 0) # reset without sync is not a complete handoff ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_topk_exact_counters_reader_log_line(monkeypatch) -> None: ++ """The Python reader must produce the exact line the smoke greps. ++ ++ Runs the op path directly -- no engine, no model -- so a failure here is ++ the reader's, not a serving regression. Asserting the format with a regex ++ is the point: the smoke keys on these field names, and a reader that ++ prints something else is as good as no reader at all. ++ """ ++ from vllm.models.qwen4_exp.nvidia.ops import topk_counters ++ ++ torch.set_default_device("cuda:0") ++ direct = _read_exact_counters() ++ if direct is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: no " ++ "persistent_topk_exact_counters op. A skip is not a pass." ++ ) ++ ++ monkeypatch.setenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "1") ++ monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", str(_EXACT_COUNTER_SLOTS)) ++ assert envs.MERLIN_TOPK_EXACT_COUNTERS_LOG == 1 ++ ++ counters = topk_counters.read() ++ assert list(counters) == [ ++ "rows", ++ "overflow", ++ "fallbacks", ++ "faults", ++ "switch", ++ ], f"unexpected slot names: {list(counters)}" ++ assert counters == dict(zip(counters, direct, strict=True)), ( ++ "the reader disagrees with a direct op call" ++ ) ++ ++ line = topk_counters.log_line("PP0_TP0") ++ assert re.fullmatch( ++ r"topk-exact-counters rank=PP0_TP0 rows=\d+ overflow=\d+ " ++ r"fallbacks=\d+ faults=\d+ switch=\d+", ++ line, ++ ), f"log line does not match the documented format: {line!r}" ++ ++ # The cadence gate must also be able to emit -- a reader that never fires ++ # from maybe_log() is the same defect one level up. ++ topk_counters._last_log_monotonic = 0.0 ++ expected = line.replace("PP0_TP0", topk_counters.rank_tag()) ++ assert topk_counters.maybe_log() == expected ++ # ... and must then stay silent until the period elapses. ++ assert topk_counters.maybe_log() == "" ++ ++ # Unknown slot counts must degrade to positional names, never mislabel. ++ assert topk_counters.slot_names(7) == tuple(f"slot{i}" for i in range(7)) ++ if getattr(torch.ops._C, "persistent_topk_exact_counters_slots", None) is None: ++ monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", "7") ++ assert topk_counters.slot_count() == 7 +diff --git a/vllm/envs.py b/vllm/envs.py +index 35047dd..5923dc5 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -506,6 +506,7 @@ if TYPE_CHECKING: + VLLM_GPU_NIC_PCIE_MAPPING: str = "" + VLLM_NIC_SELECTION_VARS: str = "" + VLLM_ENABLE_HPC_OPS: bool = False ++ MERLIN_TOPK_EXACT_COUNTERS_LOG: int = 0 + + + def get_default_cache_root(): +@@ -2459,6 +2460,13 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_E2_MAMBA_SLOT_OWNERSHIP": lambda: _e2_diagnostic_mode( + "VLLM_E2_MAMBA_SLOT_OWNERSHIP" + ), ++ # Seconds between `topk-exact-counters` lines from each GPU worker ++ # (0, the default, disables the reader entirely). The line reports the ++ # medium-path exact-restart counters; see ++ # vllm/models/qwen4_exp/nvidia/ops/topk_counters.py. ++ "MERLIN_TOPK_EXACT_COUNTERS_LOG": lambda: int( ++ os.getenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "0") ++ ), + } + + +@@ -2683,6 +2691,10 @@ def compile_factors() -> dict[str, object]: + # configuration factor and is hashed conservatively. + "VLLM_E2_PLE_FLAG_CHECK", + "VLLM_E2_MAMBA_SLOT_OWNERSHIP", ++ # Logging cadence for an already-compiled kernel's counters: it gates a ++ # host-side timer and a snapshot op, and is read by no compiled code, so ++ # hashing it would invalidate every compile cache for a log interval. ++ "MERLIN_TOPK_EXACT_COUNTERS_LOG", + # ------------------------------------------------------------------ + } + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +new file mode 100644 +index 0000000..02c1720 +--- /dev/null ++++ b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +@@ -0,0 +1,167 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Python reader for the medium-path exact-restart top-k counters. ++ ++The counters themselves are device-global uint32s owned by ++`csrc/libtorch_stable/topk.cu`; `torch.ops._C.persistent_topk_exact_counters` ++snapshots them into an int32 CUDA tensor. Without this module nothing in a ++running engine ever calls that op, so the instrument records into a buffer no ++one reads -- which measures nothing. ++ ++SLOT LAYOUT is build-dependent, so nothing here hard-codes five. `slot_count()` ++asks the build first (a `..._slots` op, should a later build expose one) and ++otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5); names come from a ++table keyed by that length: ++ ++ 5 rows overflow fallbacks faults switch ++ 10 rows overflow fallbacks faults rows_large ties_cut ties_dropped ++ faults_large switch switch_large ++ ++An unknown length degrades to `slot` names rather than mislabelling. ++ ++HOW THE SMOKE READS IT. Set `MERLIN_TOPK_EXACT_COUNTERS_LOG=` on the ++engine; each GPU worker then emits, at most that often and once more at ++shutdown, one line into the ordinary engine log: ++ ++ topk-exact-counters rank=PP0_TP1 rows=... overflow=... fallbacks=... ++ faults=... switch=... ++ ++Counts are cumulative per worker process, so the smoke takes the LATEST line ++per rank -- `grep topk-exact-counters | ...` keyed on the `rank=` ++field -- and reads `faults` as the pass/fail number (it must stay 0) with ++`rows` as the denominator. Earlier lines are only useful as a rate. ++ ++COST. `read()` passes `sync=1`, i.e. a full device synchronize, because a ++stream-ordered snapshot can miss increments from work in flight on another ++stream. That is acceptable ONLY because the cadence is coarse (>= 30 s in ++practice): the time check happens before the op call, so a step that is not ++due for a log touches nothing but `time.monotonic()`. The read is also refused ++inside a cudagraph capture -- the op allocates and synchronizes, which would ++abort the capture. ++""" ++ ++import os ++import time ++ ++import torch ++ ++import vllm.envs as envs ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_OP_NAME = "persistent_topk_exact_counters" ++ ++_MEDIUM_NAMES = ("rows", "overflow", "fallbacks", "faults", "switch") ++_LARGEDET_NAMES = ( ++ "rows", ++ "overflow", ++ "fallbacks", ++ "faults", ++ "rows_large", ++ "ties_cut", ++ "ties_dropped", ++ "faults_large", ++ "switch", ++ "switch_large", ++) ++_SLOT_NAMES: dict[int, tuple[str, ...]] = {5: _MEDIUM_NAMES, 10: _LARGEDET_NAMES} ++ ++# The medium-path line contract. Any further name is appended after these, in slot ++# order, so a 10-slot build still starts with exactly the medium-path line. ++_LINE_ORDER = _MEDIUM_NAMES ++ ++_warned_absent = False ++_warned_failed = False ++_last_log_monotonic = 0.0 ++ ++ ++def _get_op(name: str): ++ try: ++ return getattr(torch.ops._C, name, None) ++ except (AttributeError, RuntimeError): ++ return None ++ ++ ++def slot_count() -> int: ++ """How many int32 slots this build's counter array carries.""" ++ query = _get_op(_OP_NAME + "_slots") ++ if query is not None: ++ return int(query()) ++ return int(os.getenv("MERLIN_TOPK_COUNTER_SLOTS", "5")) ++ ++ ++def slot_names(count: int) -> tuple[str, ...]: ++ names = _SLOT_NAMES.get(count) ++ return names if names is not None else tuple(f"slot{i}" for i in range(count)) ++ ++ ++def read(reset: bool = False, sync: bool = True) -> dict[str, int]: ++ """Snapshot the counters as {slot name: value}; {} when unreadable.""" ++ global _warned_absent ++ op = _get_op(_OP_NAME) ++ if op is None: ++ if not _warned_absent: ++ _warned_absent = True ++ logger.warning( ++ "torch.ops._C.%s is not in this build; the top-k exact-restart " ++ "counters cannot be read.", ++ _OP_NAME, ++ ) ++ return {} ++ if reset and not sync: ++ raise ValueError("reset requires sync: zeroing races in-flight atomics") ++ if not (torch.cuda.is_available() and torch.cuda.is_initialized()): ++ return {} ++ if torch.cuda.is_current_stream_capturing(): ++ # The op allocates and synchronizes; calling it here aborts the capture. ++ return {} ++ count = slot_count() ++ device = torch.device("cuda", torch.cuda.current_device()) ++ out = torch.zeros(count, dtype=torch.int32, device=device) ++ op(out, 1 if reset else 0, 1 if sync else 0) ++ return dict(zip(slot_names(count), out.tolist(), strict=True)) ++ ++ ++def rank_tag() -> str: ++ from vllm.distributed.parallel_state import get_pp_group, get_tp_group ++ ++ try: ++ return f"PP{get_pp_group().rank_in_group}_TP{get_tp_group().rank_in_group}" ++ except Exception: ++ # Pre-init or single-process: say so rather than claim rank 0. ++ return "PP?_TP?" ++ ++ ++def log_line(tag: str) -> str: ++ """The one-line form the smoke greps; "" when the counters are unreadable.""" ++ counters = read() ++ if not counters: ++ return "" ++ order = [name for name in _LINE_ORDER if name in counters] ++ order += [name for name in counters if name not in _LINE_ORDER] ++ fields = " ".join(f"{name}={counters[name]}" for name in order) ++ return f"topk-exact-counters rank={tag} {fields}" ++ ++ ++def maybe_log(force: bool = False) -> str: ++ """Emit the line if due (or forced). Time check first: when not due this ++ costs one clock read and never touches the device.""" ++ global _last_log_monotonic, _warned_failed ++ period = envs.MERLIN_TOPK_EXACT_COUNTERS_LOG ++ if period <= 0: ++ return "" ++ now = time.monotonic() ++ if not force and now - _last_log_monotonic < period: ++ return "" ++ _last_log_monotonic = now ++ try: ++ line = log_line(rank_tag()) ++ except Exception as exc: # never let an instrument break a step or shutdown ++ if not _warned_failed: ++ _warned_failed = True ++ logger.warning("top-k exact-restart counter read failed: %s", exc) ++ return "" ++ if line: ++ logger.info("%s", line) ++ return line +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 67ba37f..b5b8a70 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -55,6 +55,7 @@ from vllm.model_executor.offloader import ( + set_offloader, + ) + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry ++from vllm.models.qwen4_exp.nvidia.ops import topk_counters + from vllm.multimodal import MULTIMODAL_REGISTRY + from vllm.multimodal.encoder_budget import ( + MultiModalBudget, +@@ -1784,6 +1785,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # A Python bool; adds no device sync to the step. Late and + # rank-local -- see guard_telemetry.check_pending. + guard_telemetry.check_pending() ++ # Post-step read of the top-k exact-restart counters, i.e. ++ # everything the previous steps accumulated. Off unless ++ # MERLIN_TOPK_EXACT_COUNTERS_LOG > 0, and then a clock read on ++ # every step but the one per cadence that is actually due. It sits ++ # here rather than at the tail of sample_tokens because that method ++ # returns early on non-last PP ranks, so a tail hook would silently ++ # never fire on them. Outside any capture: this is the engine ++ # thread between steps (read() re-checks). ++ topk_counters.maybe_log() + # Update the request states. + self.update_pp_decode_requests() + self.finish_requests(scheduler_output) +@@ -2479,6 +2489,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + torch.accelerator.synchronize() ++ # The last window of a run is the one a cadence log always ++ # misses; emit it before the device state goes away. ++ topk_counters.maybe_log(force=True) + if self._ple_offload_connector is not None: + self._ple_offload_connector.close() + self._ple_offload_connector = None +-- +2.43.0 + diff --git a/patches/0041-topk-barrier-the-histogram-round-before-the-leader-p.patch b/patches/0041-topk-barrier-the-histogram-round-before-the-leader-p.patch new file mode 100644 index 0000000..37bfdde --- /dev/null +++ b/patches/0041-topk-barrier-the-histogram-round-before-the-leader-p.patch @@ -0,0 +1,40 @@ +From 090e8263f6d923c61fa6bc328e5a7b395aff5ae5 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:31:18 -0400 +Subject: [PATCH 041/103] topk: barrier the histogram round before the leader + publishes arrival + +In radix_topk's large-path round loop the leader published arrival with no +block barrier after the histogram merge, so a peer CTA could observe an +incomplete histogram. Inherited by earlier releases; fixed below the placement +work so both arms of any comparison carry it. +--- + csrc/libtorch_stable/persistent_topk.cuh | 12 ++++++++++++ + 1 file changed, 12 insertions(+) + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 8bfb4b9..c5d27d9 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -1179,6 +1179,18 @@ __device__ void radix_topk(const float* __restrict__ row_input, + } + } + ++ // The arrival publication below is made by thread 0 ++ // ALONE, and `red_release`'s `fence.acq_rel.gpu` orders only the issuing ++ // thread's own prior accesses. Without this barrier thread 0 could signal ++ // arrival while other warps of this CTA had not yet merged their bins into ++ // `current_hist`, or (in CTA 0) had not finished zeroing `next_hist`. A ++ // peer CTA that observes the completed arrival count then leaves its own ++ // `wait_ge` -- whose `__syncthreads()` is AFTER arrival, so it orders ++ // nothing here -- and reads a half-written histogram, which changes the ++ // pivot. The barrier must be between the producing stores and the leader's ++ // publication; program order in one thread is not cross-thread ordering. ++ __syncthreads(); ++ + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } +-- +2.43.0 + diff --git a/patches/0042-topk-deterministic-placement-and-tie-membership-for-.patch b/patches/0042-topk-deterministic-placement-and-tie-membership-for-.patch new file mode 100644 index 0000000..8961c3a --- /dev/null +++ b/patches/0042-topk-deterministic-placement-and-tie-membership-for-.patch @@ -0,0 +1,788 @@ +From 5ff40ff3254e92df316c73efa58bba20a12a7d35 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:31:18 -0400 +Subject: [PATCH 042/103] topk: deterministic placement and tie membership for + the large-path radix_topk + +Above RADIX_THRESHOLD a row is split across a group of CTAs. The pivot was +already order-independent; cross-CTA placement, intra-CTA placement and tie +membership were arrival-order races. They are now deterministic behind +MERLIN_TOPK_DET_LARGE (default on). +--- + csrc/libtorch_stable/persistent_topk.cuh | 326 ++++++++++++++++++++++- + csrc/libtorch_stable/topk.cu | 68 ++++- + tests/kernels/test_top_k_per_row.py | 238 ++++++++++++++++- + 3 files changed, 617 insertions(+), 15 deletions(-) + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index c5d27d9..8dd3f7a 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -147,6 +147,11 @@ struct PersistentTopKParams { + // nullptr; the kernel only ever atomicAdds into it, never reads back. + bool exact_medium; + unsigned int* exact_counters; ++ // MERLIN_TOPK_DET_LARGE, ANDed on the host with the capacity gate ++ // 2 * ctas_per_group <= RADIX (the per-CTA counts are staged in one 256-slot ++ // RadixRowState histogram row). Resolved once per process, so it is a ++ // launch constant and a cudagraph captures one value. ++ bool det_large; + }; + + // ============================================================================ +@@ -468,7 +473,16 @@ enum TopKExactCounter : int { + kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop + kCtrMediumExactFallback = 2, // rows that took the exact restart + kCtrMediumExactFault = 3, // exact restart could not complete +- kNumTopKExactCounters = 4, ++ // Large path (multi-CTA radix_topk, seq_len > RADIX_THRESHOLD). ++ // Slots 4-7 are only ever touched by the deterministic emission, so they ++ // stay 0 in the MERLIN_TOPK_DET_LARGE=0 control arm: `[9] == 1 && [4] == 0` ++ // while large rows are known to have run means the host-side capacity gate ++ // (2 * ctas_per_group <= RADIX) turned the switch off for that launch. ++ kCtrLargeDetRows = 4, // large rows emitted deterministically ++ kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut ++ kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts ++ kCtrLargeFault = 7, // large-path invariant violation ++ kNumTopKExactCounters = 8, + }; + + // Exact, bounded-memory re-selection for a single medium-path row. +@@ -1074,6 +1088,300 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, + // coordinated radix select via global-memory histograms and barriers. + // ============================================================================ + ++// ============================================================================ ++// Deterministic placement and tie membership for the large path ++// ++// `radix_topk`'s pivot is already an exact, order-independent function of the ++// row: every quantity feeding it is an integer histogram add and a ++// single-threaded descending bin walk. Three things after the ++// pivot are not: ++// ++// (a) cross-CTA placement — a CTA's base offset into `row_output` is its ++// arrival order at `state->output_counter`; ++// (b) intra-CTA placement — a slot inside that base is arrival order at a ++// shared counter; ++// (c) cross-CTA tie membership — which pivot-equal elements survive the ++// `pos < TopK` cut is a race across CTAs. ++// ++// (a) and (b) change only the permutation of the selected list, but sparse ++// attention consumes that list as an ordered gather, so a permutation forks ++// the hidden state. (c) changes the set. All three are replaced here by a ++// rank derived from (cta_in_group, element index): ++// ++// * each CTA counts its own `> pivot` and `== pivot` elements out of the ++// chunk already resident in shared memory (no new global reads); ++// * the counts are published into `RadixRowState::histogram[radix_iter % 3]` ++// — the slot round 3 has just finished reading — so the workspace ++// contract is unchanged; ++// * every CTA reads all 2*ctas_per_group counts and derives its own base by ++// an exclusive prefix over `cta_in_group`, in one fixed order, so all ++// CTAs agree without further communication; ++// * the tie cut keeps the HIGHEST global indices (most-recent-wins, the ++// `MERLIN_TIE_RECENT` policy, measured at +3.4 % tokens/step ++// over oldest-wins). Chunks are handed out in ascending index order ++// (`my_chunk_start = cta_in_group * chunk_size`), so "the highest ++// `equal_needed` pivot-equal indices" is exactly "walk CTAs from the top ++// down", which the per-CTA equal counts already answer; ++// * emission is a prefix sum over ascending element index — no atomic ++// anywhere on the placement path — so the row is strictly ascending in ++// index and bit-reproducible. ++// ++// Cost per row: one extra inter-CTA barrier, two extra ballots and one ++// 32-lane shuffle scan per 1,024-element tile, 2 global stores and ++// 2*ctas_per_group global loads. No second pass over the row in global ++// memory: the pivot-equal count that the tie cut needs is produced by the same ++// shared-memory pass v2.0.1 already runs to count `> pivot`. ++// ============================================================================ ++template ++__device__ __noinline__ void radix_emit_deterministic( ++ int32_t* __restrict__ row_output, uint32_t ordered_pivot, ++ uint32_t equal_needed, uint32_t my_chunk_start, uint32_t actual_chunk_size, ++ uint32_t* local_histogram, uint32_t* suffix_sum, uint32_t* shared_ordered, ++ RadixRowState* state, uint32_t cta_in_group, uint32_t ctas_per_group, ++ int& barrier_phase, uint32_t radix_iter, uint32_t tx, ++ unsigned int* counters) { ++ constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u; ++ // Staging for the published per-CTA counts. The host gates the switch on ++ // 2 * ctas_per_group <= RADIX, so these two halves never collide. ++ constexpr uint32_t kEqStage = RADIX / 2u; ++ ++ // ---- per-CTA counts, out of the resident chunk ------------------------- ++ if (tx < 2u) suffix_sum[tx] = 0u; ++ __syncthreads(); ++ ++ uint32_t my_gt = 0u; ++ uint32_t my_eq = 0u; ++ for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { ++ const uint32_t ord = shared_ordered[i]; ++ if (ord > ordered_pivot) { ++ my_gt++; ++ } else if (ord == ordered_pivot) { ++ my_eq++; ++ } ++ } ++#pragma unroll ++ for (int off = 16; off > 0; off >>= 1) { ++ my_gt += __shfl_down_sync(0xffffffffu, my_gt, off); ++ my_eq += __shfl_down_sync(0xffffffffu, my_eq, off); ++ } ++ if ((tx & 31u) == 0u) { ++ // Integer counts: the atomic's arrival order cannot change the sum, so ++ // this is a reduction, not a placement. ++ if (my_gt > 0u) atomicAdd(&suffix_sum[0], my_gt); ++ if (my_eq > 0u) atomicAdd(&suffix_sum[1], my_eq); ++ } ++ __syncthreads(); ++ const uint32_t local_gt = suffix_sum[0]; ++ const uint32_t local_eq = suffix_sum[1]; ++ ++ // ---- barrier X1: the histogram slot is free ---------------------------- ++ // Round 3's `current_hist` is state->histogram[(radix_iter*4+3) % 3], i.e. ++ // histogram[radix_iter % 3] (4r+3 == r mod 3). Every CTA copies it into ++ // shared memory AFTER round 3's own barrier and there is no inter-CTA ++ // barrier between that copy and here, so a fast CTA could otherwise ++ // overwrite the slot while a slow CTA is still reading it. This barrier is ++ // the one the slot-reuse argument requires. ++ if (tx == 0) { ++ red_release(&state->arrival_counter, 1); ++ } ++ wait_ge(&state->arrival_counter, ++ (barrier_phase + 1) * static_cast(ctas_per_group), tx); ++ barrier_phase++; ++ __syncthreads(); ++ ++ uint32_t* cta_counts = state->histogram[radix_iter % 3u]; ++ if (tx == 0) { ++ cta_counts[cta_in_group] = local_gt; ++ cta_counts[ctas_per_group + cta_in_group] = local_eq; ++ } ++ ++ // ---- barrier X2: the counts are published ------------------------------ ++ // red_release issues `fence.acq_rel.gpu` before the arrival add and wait_ge ++ // reads with `ld.global.acquire.gpu`, so the two plain stores above are ++ // visible at device scope to every CTA that leaves this barrier. This is ++ // the same release/acquire pairing the shipped kernel already relies on for ++ // CTA-0's plain-store zeroing of `next_hist`. ++ if (tx == 0) { ++ red_release(&state->arrival_counter, 1); ++ } ++ wait_ge(&state->arrival_counter, ++ (barrier_phase + 1) * static_cast(ctas_per_group), tx); ++ barrier_phase++; ++ __syncthreads(); ++ ++ // ---- bases and the tie cut --------------------------------------------- ++ for (uint32_t c = tx; c < ctas_per_group; c += kThreadsPerBlock) { ++ local_histogram[c] = cta_counts[c]; ++ local_histogram[kEqStage + c] = cta_counts[ctas_per_group + c]; ++ } ++ __syncthreads(); ++ ++ if (tx == 0) { ++ uint32_t gt_base = 0u; ++ uint32_t gt_total = 0u; ++ for (uint32_t c = 0; c < ctas_per_group; ++c) { ++ const uint32_t g = local_histogram[c]; ++ if (c < cta_in_group) gt_base += g; ++ gt_total += g; ++ } ++ ++ // Most-recent-wins: walk CTAs from the top down, handing each one as many ++ // of the `equal_needed` slots as the CTAs above it left unclaimed. keep[] ++ // is written over eq[] in place (eq[c] is read before keep[c] is stored). ++ uint32_t eq_above = 0u; ++ for (int c = static_cast(ctas_per_group) - 1; c >= 0; --c) { ++ const uint32_t e = local_histogram[kEqStage + c]; ++ const uint32_t want = ++ (eq_above >= equal_needed) ? 0u : (equal_needed - eq_above); ++ local_histogram[kEqStage + c] = (want < e) ? want : e; ++ eq_above += e; ++ } ++ const uint32_t eq_total = eq_above; ++ ++ uint32_t eq_keep_base = 0u; ++ for (uint32_t c = 0; c < cta_in_group; ++c) { ++ eq_keep_base += local_histogram[kEqStage + c]; ++ } ++ const uint32_t my_keep = local_histogram[kEqStage + cta_in_group]; ++ ++ // Invariants, all three derived from the medium-path restart design and re-derivable ++ // here: the radix walk guarantees #{key > pivot} == TopK - equal_needed ++ // and equal_needed <= #{key == pivot}, and the two together make the ++ // emission a bijection onto [0, TopK). Every CTA evaluates this from the ++ // same published integers, so the verdict is group-uniform and needs no ++ // further communication. ++ uint32_t fault = 0u; ++ if (gt_total + equal_needed != static_cast(TopK)) fault = 1u; ++ if (equal_needed > eq_total) fault = 1u; ++ if (my_keep > local_eq) fault = 1u; ++ ++ suffix_sum[2] = gt_base; ++ suffix_sum[3] = eq_keep_base; ++ // Keep the local pivot-equal elements whose ascending rank within this ++ // chunk is >= cut_rank, i.e. the `my_keep` highest-indexed of them. ++ suffix_sum[4] = local_eq - my_keep; ++ suffix_sum[5] = fault; ++ ++ if (counters != nullptr && cta_in_group == 0u) { ++ atomicAdd(&counters[kCtrLargeDetRows], 1u); ++ if (eq_total > equal_needed) { ++ atomicAdd(&counters[kCtrLargeTiesCut], 1u); ++ atomicAdd(&counters[kCtrLargeTiesDropped], eq_total - equal_needed); ++ } ++ if (fault != 0u) atomicAdd(&counters[kCtrLargeFault], 1u); ++ } ++ } ++ __syncthreads(); ++ ++ const uint32_t gt_base = suffix_sum[2]; ++ const uint32_t eq_keep_base = suffix_sum[3]; ++ const uint32_t cut_rank = suffix_sum[4]; ++ const uint32_t row_fault = suffix_sum[5]; ++ ++ // ---- emission: rank by index, no atomics ------------------------------- ++ if (row_fault == 0u) { ++ // Scratch for the tile scan. The staged per-CTA counts above are dead by ++ // now (thread 0 consumed them before the barrier), so local_histogram is ++ // free again; 33 slots each at 0 and at 64 stay inside RADIX. ++ uint32_t* warp_gt = local_histogram; ++ uint32_t* warp_eq = local_histogram + 64; ++ const uint32_t lane = tx & 31u; ++ const uint32_t warp = tx >> 5; ++ uint32_t gt_run = 0u; ++ uint32_t eq_run = 0u; ++ ++ // Block-uniform trip count: every thread runs every tile, so the ballots ++ // and the barriers inside are reached by the whole block. ++ for (uint32_t base = 0u; base < actual_chunk_size; ++ base += kThreadsPerBlock) { ++ const uint32_t i = base + tx; ++ const bool in_range = (i < actual_chunk_size); ++ const uint32_t ord = in_range ? shared_ordered[i] : 0u; ++ const bool f_gt = in_range && (ord > ordered_pivot); ++ const bool f_eq = in_range && (ord == ordered_pivot); ++ ++ const uint32_t bgt = __ballot_sync(0xffffffffu, f_gt); ++ const uint32_t beq = __ballot_sync(0xffffffffu, f_eq); ++ const uint32_t lane_mask = (1u << lane) - 1u; ++ const uint32_t gt_in_warp = static_cast(__popc(bgt & lane_mask)); ++ const uint32_t eq_in_warp = static_cast(__popc(beq & lane_mask)); ++ if (lane == 0u) { ++ warp_gt[warp] = static_cast(__popc(bgt)); ++ warp_eq[warp] = static_cast(__popc(beq)); ++ } ++ __syncthreads(); ++ ++ if (tx < 32u) { ++ const uint32_t cg = (tx < kWarpsPerBlock) ? warp_gt[tx] : 0u; ++ const uint32_t ce = (tx < kWarpsPerBlock) ? warp_eq[tx] : 0u; ++ uint32_t sg = cg; ++ uint32_t se = ce; ++#pragma unroll ++ for (int off = 1; off < 32; off <<= 1) { ++ const uint32_t ng = __shfl_up_sync(0xffffffffu, sg, off); ++ const uint32_t ne = __shfl_up_sync(0xffffffffu, se, off); ++ if (tx >= static_cast(off)) { ++ sg += ng; ++ se += ne; ++ } ++ } ++ warp_gt[tx] = sg - cg; // inclusive -> exclusive ++ warp_eq[tx] = se - ce; ++ if (tx == 31u) { ++ warp_gt[32] = sg; // tile totals (lanes >= kWarpsPerBlock add 0) ++ warp_eq[32] = se; ++ } ++ } ++ __syncthreads(); ++ ++ const uint32_t gt_rank = gt_run + warp_gt[warp] + gt_in_warp; ++ const uint32_t eq_rank = eq_run + warp_eq[warp] + eq_in_warp; ++ const uint32_t tile_gt = warp_gt[32]; ++ const uint32_t tile_eq = warp_eq[32]; ++ ++ if (f_gt || (f_eq && eq_rank >= cut_rank)) { ++ // Selected elements before this one, in ascending index order: ++ // all `> pivot` ones, plus the pivot-equal ones past the cut. ++ const uint32_t kept_eq_before = ++ (eq_rank > cut_rank) ? (eq_rank - cut_rank) : 0u; ++ const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before; ++ if (slot < static_cast(TopK)) { ++ row_output[slot] = static_cast(my_chunk_start + i); ++ } else if (counters != nullptr) { ++ // Unreachable while the three invariants above hold; bounded store ++ // plus a counter rather than an out-of-range write. ++ atomicAdd(&counters[kCtrLargeFault], 1u); ++ } ++ } ++ ++ __syncthreads(); // before warp_gt / warp_eq are rewritten ++ gt_run += tile_gt; ++ eq_run += tile_eq; ++ } ++ // No thread may start the next row's chunk load — which writes ++ // shared_ordered — until every warp has finished reading it here. ++ __syncthreads(); ++ } else { ++ // ---- fail loud ------------------------------------------------------ ++ // The predicate is group-uniform (same published integers, same fixed ++ // summation order in every CTA), so no CTA of this group emitted anything ++ // and CTA 0 can fill the row alone. -1 is the kernel's own padding value ++ // and the consumer rejects it, so a faulted row attends nothing instead of ++ // quietly attending the wrong blocks. ++ if (cta_in_group == 0u) { ++ for (uint32_t i = tx; i < static_cast(TopK); ++ i += kThreadsPerBlock) { ++ row_output[i] = -1; ++ } ++ } ++ __threadfence(); ++ __syncthreads(); ++#if MERLIN_TOPK_EXACT_TRAP ++ __trap(); ++#endif ++ } ++} ++ + // ============================================================================ + // Multi-CTA cooperative RadixTopK for a single large row. + // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 +@@ -1087,7 +1395,8 @@ __device__ void radix_topk(const float* __restrict__ row_input, + uint32_t* shared_scalars, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, + uint32_t ctas_per_group, int& barrier_phase, +- uint32_t radix_iter, uint32_t tx) { ++ uint32_t radix_iter, uint32_t tx, bool det_large, ++ unsigned int* counters) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; +@@ -1241,6 +1550,16 @@ __device__ void radix_topk(const float* __restrict__ row_input, + // -- Count local > pivot elements -- + const uint32_t ordered_pivot = shared_scalars[0]; + ++ // Deterministic placement + tie membership. Everything below this ++ // branch is v2.0.1 verbatim and is what MERLIN_TOPK_DET_LARGE=0 runs. ++ if (det_large) { ++ radix_emit_deterministic( ++ row_output, ordered_pivot, shared_scalars[1], my_chunk_start, ++ actual_chunk_size, local_histogram, suffix_sum, shared_ordered, state, ++ cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters); ++ return; ++ } ++ + if (tx == 0) suffix_sum[0] = 0; + __syncthreads(); + +@@ -1389,7 +1708,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + radix_topk( + row_input, row_output, seq_len, my_chunk_start, chunk_size, + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, +- cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx); ++ cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, ++ params.det_large, params.exact_counters); + radix_iter++; + } + } +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 3419344..001c782 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -4,6 +4,7 @@ + #include + #include + #include ++#include + #include + #include + +@@ -36,6 +37,24 @@ bool exact_medium_enabled() { + return enabled; + } + ++// --------------------------------------------------------------------------- ++// MERLIN_TOPK_DET_LARGE switch. ++// ++// Same contract as the medium-path switch: resolved once per process, so it is a ++// launch constant and a cudagraph captures one value. Default ON on this ++// branch. MERLIN_TOPK_DET_LARGE=0 is the A/B control arm and runs the ++// v2.0.1 `radix_topk` tail verbatim -- the deterministic emission sits behind ++// one `if (det_large)` in front of it and touches nothing when off, including ++// the large-path counters, which is what makes the control arm the real baseline. ++// --------------------------------------------------------------------------- ++bool det_large_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_DET_LARGE"); ++ return (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0); ++ }(); ++ return enabled; ++} ++ + // Process-lifetime device counters, one block per device. The kernel only + // ever atomicAdds into them, so the launch path needs no host sync and the + // values accumulate across cudagraph replays. +@@ -340,6 +359,30 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + params.ctas_per_group = ctas_per_group; + params.max_seq_len = static_cast(max_seq_len); + params.exact_medium = exact_medium_enabled(); ++ // The per-CTA selected / pivot-equal counts are staged in ONE ++ // 256-slot RadixRowState histogram row (the slot round 3 has just finished ++ // reading), which holds 2 * ctas_per_group uint32s. That is the whole ++ // reason the workspace contract is unchanged, and it is also the only ++ // limit: above 128 CTAs per group the deterministic path cannot be ++ // expressed there, so the launch degrades to the v2.0.1 tail rather than ++ // failing. Unreachable at any shape this kernel serves (ctas_per_group is ++ // 2-3 across our 135K-262K band and <= 16 at the smallest chunk size), but ++ // a silent loss of determinism is exactly what this branch exists to stop, ++ // so say it once. ++ const bool det_large_fits = ++ 2u * ctas_per_group <= static_cast(P::RADIX); ++ params.det_large = det_large_enabled() && det_large_fits; ++ if (det_large_enabled() && !det_large_fits) { ++ static std::atomic det_large_warned{false}; ++ if (!det_large_warned.exchange(true)) { ++ std::fprintf(stderr, ++ "[vllm] MERLIN_TOPK_DET_LARGE is on but ctas_per_group=%u " ++ "needs %u of %d staging slots; the large path falls back " ++ "to nondeterministic placement at this shape (counter [4] " ++ "stays 0).\n", ++ ctas_per_group, 2u * ctas_per_group, P::RADIX); ++ } ++ } + // Allocate only off the capture path; in the ordinary engine lifecycle the + // eager warmup lands here first, long before capture. The explicit, + // quiescent-phase entry point is the reader op. +@@ -444,12 +487,17 @@ void persistent_topk(const torch::stable::Tensor& logits, + // counters. This is also the explicit pre-capture init path: call it once per + // device before any cudagraph capture. + // +-// `out` is an int32 CUDA tensor with at least 5 elements: ++// `out` is an int32 CUDA tensor with at least 10 elements: + // [0] medium-path rows entered (denominator) + // [1] rows where a 4,096-slot guard would drop a candidate + // [2] rows that took the exact restart (0 in the control arm) + // [3] exact restarts that could not complete (must stay 0) +-// [4] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0 ++// [4] large-path rows emitted deterministically (0 in the control arm) ++// [5] large-path rows whose pivot-equal group was cut at the k-th value ++// [6] large-path pivot-equal elements discarded by those cuts ++// [7] large-path invariant violations (must stay 0) ++// [8] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0 ++// [9] 1 if MERLIN_TOPK_DET_LARGE is on, else 0 + // + // COHERENCY -- stated precisely, because the earlier wording overclaimed. + // With `sync == 0` the snapshot is ordered ONLY against work already enqueued +@@ -484,9 +532,9 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + STD_TORCH_CHECK(out.is_contiguous(), + "out must be contiguous (expanded or strided views would be " + "written out of bounds or into the wrong elements)"); +- STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 1, ++ STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 2, + "out must hold at least ", +- vllm::persistent::kNumTopKExactCounters + 1, " elements"); ++ vllm::persistent::kNumTopKExactCounters + 2, " elements"); + STD_TORCH_CHECK( + reset == 0 || sync != 0, + "persistent_topk_exact_counters: reset requires sync != 0, otherwise the " +@@ -531,10 +579,16 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + } + + // Static storage: the source of an async H2D copy must outlive the copy. +- static const int32_t enabled = exact_medium_enabled() ? 1 : 0; ++ // Two flags: [kNumTopKExactCounters] = the medium-path switch, ++ // [kNumTopKExactCounters + 1] = the large-path switch. The large ++ // flag reports the ENV value, not the per-launch effective value: a launch ++ // whose ctas_per_group overflowed the staging row still reports 1 here and ++ // leaves counter [4] at 0, which is how that degradation is diagnosed. ++ static const int32_t enabled[2] = {exact_medium_enabled() ? 1 : 0, ++ det_large_enabled() ? 1 : 0}; + STD_TORCH_CHECK( +- cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, &enabled, +- sizeof(int32_t), cudaMemcpyHostToDevice, ++ cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, enabled, ++ 2 * sizeof(int32_t), cudaMemcpyHostToDevice, + stream) == cudaSuccess, + "exact counter flag write failed"); + +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index b0564e8..daf8ff6 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1849,11 +1849,29 @@ def test_sparse_indexer_topk_backend_resolution() -> None: + # 3. two runs return bit-identical index tensors, in ascending index order. + # --------------------------------------------------------------------------- + +-_EXACT_COUNTER_SLOTS = 5 ++# Slot map of torch.ops._C.persistent_topk_exact_counters (see ++# csrc/libtorch_stable/topk.cu). The large path appended four counters and a ++# second switch flag, so the medium-path "enabled" flag moved from [4] to [8]. ++_SLOT_MEDIUM_ROWS = 0 ++_SLOT_MEDIUM_OVERFLOW = 1 ++_SLOT_MEDIUM_FALLBACK = 2 ++_SLOT_MEDIUM_FAULT = 3 ++_SLOT_LARGE_ROWS = 4 ++_SLOT_LARGE_TIES_CUT = 5 ++_SLOT_LARGE_TIES_DROPPED = 6 ++_SLOT_LARGE_FAULT = 7 ++_SLOT_EXACT_MEDIUM_ON = 8 ++_SLOT_DET_LARGE_ON = 9 ++_EXACT_COUNTER_SLOTS = 10 + + + def _read_exact_counters() -> list[int] | None: +- """Cumulative [rows, overflow, fallbacks, faults, enabled], or None.""" ++ """Cumulative counters + the two switch flags, or None. ++ ++ [0..3] medium: rows, overflow, exact fallbacks, faults. ++ [4..7] large: deterministic rows, tie cuts, elements cut, faults. ++ [8] MERLIN_TOPK_EXACT_MEDIUM, [9] MERLIN_TOPK_DET_LARGE. ++ """ + op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) + if op is None: + return None +@@ -1899,7 +1917,7 @@ def test_persistent_topk_medium_overflow_is_exact(top_k: int, seq_len: int) -> N + "no persistent_topk_exact_counters op, so the exact path cannot be " + "shown to have run. A skip here does not count as a pass." + ) +- if before[4] == 0: ++ if before[_SLOT_EXACT_MEDIUM_ON] == 0: + pytest.skip( + "MERLIN_TOPK_EXACT_MEDIUM=0: control arm. Exactness is not claimed " + "here, but a qualification run must also execute the arm with the " +@@ -1969,7 +1987,7 @@ def test_persistent_topk_medium_overflow_is_bit_reproducible() -> None: + "no persistent_topk_exact_counters op, so the exact path cannot be " + "shown to have run. A skip here does not count as a pass." + ) +- if counters[4] == 0: ++ if counters[_SLOT_EXACT_MEDIUM_ON] == 0: + pytest.skip( + "MERLIN_TOPK_EXACT_MEDIUM=0: control arm. Emission order is not " + "claimed here; a qualification run must also execute the arm with " +@@ -2020,7 +2038,7 @@ def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: + "QUALIFICATION FAILURE if seen on the branch build: no " + "persistent_topk_exact_counters op. A skip is not a pass." + ) +- if before[4] == 0: ++ if before[_SLOT_EXACT_MEDIUM_ON] == 0: + pytest.skip( + "MERLIN_TOPK_EXACT_MEDIUM=0: control arm, the tie rule is not " + "claimed here." +@@ -2150,3 +2168,213 @@ def test_topk_exact_counters_reader_log_line(monkeypatch) -> None: + if getattr(torch.ops._C, "persistent_topk_exact_counters_slots", None) is None: + monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", "7") + assert topk_counters.slot_count() == 7 ++ ++ ++# --------------------------------------------------------------------------- ++# Large-path (multi-CTA radix_topk) deterministic placement and ties ++# ++# Above RADIX_THRESHOLD = 32,768 compressed columns a row is split across a ++# group of CTAs. In v2.0.1 the pivot is exact and order-independent, but ++# everything after it is a race: a CTA's base offset into the row is its ++# arrival order at `state->output_counter`, its slot inside that base is ++# arrival order at a shared counter, and which pivot-equal elements survive the ++# `pos < TopK` cut is decided by whoever gets there first. The first two ++# permute the selected list — which sparse attention consumes as an ordered ++# gather, so they fork the hidden state — and the third changes the set. ++# ++# The row below is built so all three are visible at once: ++# * 200 elements strictly above the pivot, scattered over the whole row, ++# * 700 elements exactly AT the pivot, also scattered, so the tie group spans ++# every CTA chunk and several index high-byte buckets, ++# * top_k = 512, so exactly 312 of the 700 must survive and 388 must be cut. ++# ++# `torch.topk`'s own index choice inside a 700-way tie is arbitrary, so it ++# cannot supply the expectation. The reference is computed from the shipped ++# rule instead — most-recent-wins, i.e. the 312 HIGHEST indices — and ++# torch.topk is used for the weaker check it can make, the value multiset. ++# --------------------------------------------------------------------------- ++ ++LARGEDET_SEQ_LEN = 40960 # > RADIX_THRESHOLD (32768): the multi-CTA radix path ++LARGEDET_TOP_K = 512 ++LARGEDET_ABOVE = 200 # elements strictly greater than the pivot ++LARGEDET_TIE = 700 # elements exactly equal to the pivot ++LARGEDET_KEPT = LARGEDET_TOP_K - LARGEDET_ABOVE # 312 of the 700 survive ++LARGEDET_DROPPED = LARGEDET_TIE - LARGEDET_KEPT # 388 are cut ++ ++ ++def _build_largedet_tie_row(num_rows: int) -> tuple[torch.Tensor, torch.Tensor]: ++ """Identical rows with a deliberate multi-CTA tie at the cutoff. ++ ++ Returns (logits, expected_indices) where expected_indices is the ascending ++ int64 index list the deterministic kernel must return for every row. ++ """ ++ gen = torch.Generator(device="cpu").manual_seed(19) ++ perm = torch.randperm(LARGEDET_SEQ_LEN, generator=gen, device="cpu") ++ above_idx = perm[:LARGEDET_ABOVE] ++ tie_idx = perm[LARGEDET_ABOVE : LARGEDET_ABOVE + LARGEDET_TIE] ++ ++ # Filler stays in [0, 0.5) so nothing outside the two planted groups can ++ # reach the pivot; the planted values are exact binary fractions, so the ++ # 700 tie elements are bit-identical as float32 and really do tie. ++ row = torch.rand( ++ LARGEDET_SEQ_LEN, generator=gen, dtype=torch.float32, device="cpu" ++ ) * 0.5 ++ row[tie_idx] = 2.0 ++ row[above_idx] = 3.0 ++ ++ logits = row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ ++ kept_tie = torch.sort(tie_idx).values[-LARGEDET_KEPT:] ++ expected = torch.sort(torch.cat([above_idx, kept_tie])).values.to(torch.int64) ++ assert expected.numel() == LARGEDET_TOP_K ++ return logits, expected ++ ++ ++def _run_largedet_row(logits: torch.Tensor) -> torch.Tensor: ++ num_rows, seq_len = logits.shape ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ indices = torch.empty((num_rows, LARGEDET_TOP_K), dtype=torch.int32, device="cuda") ++ _run_topk_backend( ++ "persistent_topk", logits, lengths, indices, LARGEDET_TOP_K, seq_len ++ ) ++ torch.accelerator.synchronize() ++ return indices ++ ++ ++def _largedet_counters_or_skip(require_on: bool) -> list[int]: ++ counters = _read_exact_counters() ++ if counters is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: this build has " ++ "no persistent_topk_exact_counters op, so the large path cannot be " ++ "shown to have run deterministically. A skip is not a pass." ++ ) ++ if require_on and counters[_SLOT_DET_LARGE_ON] == 0: ++ pytest.skip( ++ "MERLIN_TOPK_DET_LARGE=0: control arm. Determinism is not claimed " ++ "here; a qualification run must also execute the arm with the " ++ "switch on." ++ ) ++ if not require_on and counters[_SLOT_DET_LARGE_ON] != 0: ++ pytest.skip( ++ "MERLIN_TOPK_DET_LARGE=1: this is the control-arm test and must be " ++ "run a second time with the switch off. A qualification run needs " ++ "BOTH arms." ++ ) ++ return counters ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_tie_keeps_highest_indices() -> None: ++ """Multi-CTA tie at the cutoff: most-recent-wins, ascending, reproducible.""" ++ torch.set_default_device("cuda:0") ++ num_rows = 8 # <= 32 keeps the dispatch on persistent_topk ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ logits, expected = _build_largedet_tie_row(num_rows) ++ first = _run_largedet_row(logits).clone() ++ mid = _read_exact_counters() ++ second = _run_largedet_row(logits) ++ after = _read_exact_counters() ++ ++ # 1. The instrument fired, with EXACT deltas. A lower bound would hide ++ # double counting, and a green assertion below means nothing unless the ++ # deterministic path is what produced it. ++ assert mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_rows, ( ++ "the deterministic large path did not run on every row: " ++ f"{mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS]} of {num_rows}" ++ ) ++ assert mid[_SLOT_LARGE_TIES_CUT] - before[_SLOT_LARGE_TIES_CUT] == num_rows, ( ++ "the tie cut did not fire; this row is built so it must" ++ ) ++ assert ( ++ mid[_SLOT_LARGE_TIES_DROPPED] - before[_SLOT_LARGE_TIES_DROPPED] ++ == num_rows * LARGEDET_DROPPED ++ ), "the number of discarded pivot-equal elements is not the expected 388/row" ++ assert after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0, ( ++ "a large-path invariant fired; the output below is not meaningful" ++ ) ++ ++ ref_values = torch.topk(logits[0], LARGEDET_TOP_K).values.sort().values ++ ++ for row in range(num_rows): ++ got = first[row].to(torch.int64) ++ ++ # 2. Ascending in index. This is what makes the Python-side ++ # MERLIN_QSA_SORT a no-op on these rows rather than a correction. ++ assert bool((got[1:] - got[:-1] > 0).all()), ( ++ f"row {row}: large-path output is not ascending in index" ++ ) ++ ++ # 3. The set is the analytic one: everything above the pivot, plus the ++ # 312 HIGHEST-indexed of the 700 pivot-equal elements. ++ assert torch.equal(got.cpu(), expected), ( ++ f"row {row}: selected set is not most-recent-wins at the cutoff" ++ ) ++ ++ # 4. The weaker check torch.topk can make: same value multiset. ++ got_values = logits[row][got].sort().values ++ assert torch.equal(got_values, ref_values), ( ++ f"row {row}: selected values are not torch.topk's values" ++ ) ++ ++ # 5. Bit identity across two launches on identical input -- the property ++ # that separates this path from the arrival-ordered one at 135K. ++ assert torch.equal(first, second), ( ++ "two identical launches returned different index tensors" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_control_arm_is_untouched() -> None: ++ """MERLIN_TOPK_DET_LARGE=0 must be the real v2.0.1 baseline. ++ ++ Run this with the switch OFF. Two things are asserted and one deliberately ++ is not: ++ ++ * the large-path counters do not move at all -- the deterministic path adds no ++ instruction to the control arm, so an A/B comparison measures the ++ kernel change and not an added atomic (as on the medium path: a ++ control arm that quietly runs a different algorithm is not a control); ++ * the returned VALUES are still exactly torch.topk's, because v2.0.1's ++ pivot is already exact -- only placement and tie membership are racy; ++ * the index SET and the emission ORDER are NOT asserted, because their ++ being unreproducible is precisely the defect the large-path change fixes. Asserting ++ either way here would be a test that cannot fail, or a flake. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows = 8 ++ ++ before = _largedet_counters_or_skip(require_on=False) ++ ++ logits, _ = _build_largedet_tie_row(num_rows) ++ out = _run_largedet_row(logits) ++ after = _read_exact_counters() ++ ++ for slot, name in ( ++ (_SLOT_LARGE_ROWS, "rows"), ++ (_SLOT_LARGE_TIES_CUT, "ties_cut"), ++ (_SLOT_LARGE_TIES_DROPPED, "ties_dropped"), ++ (_SLOT_LARGE_FAULT, "faults"), ++ ): ++ assert after[slot] - before[slot] == 0, ( ++ f"control arm touched the large-path {name} counter; it is not the " ++ "v2.0.1 baseline" ++ ) ++ ++ ref_values = torch.topk(logits[0], LARGEDET_TOP_K).values.sort().values ++ for row in range(num_rows): ++ got = out[row].to(torch.int64) ++ assert int(got.min()) >= 0 and int(got.max()) < LARGEDET_SEQ_LEN, ( ++ f"row {row}: control arm returned an out-of-range index" ++ ) ++ assert got.unique().numel() == LARGEDET_TOP_K, ( ++ f"row {row}: control arm returned a duplicate index" ++ ) ++ got_values = logits[row][got].sort().values ++ assert torch.equal(got_values, ref_values), ( ++ f"row {row}: control-arm values are not torch.topk's values" ++ ) +-- +2.43.0 + diff --git a/patches/0043-topk-large-path-hardening.patch b/patches/0043-topk-large-path-hardening.patch new file mode 100644 index 0000000..2b82e40 --- /dev/null +++ b/patches/0043-topk-large-path-hardening.patch @@ -0,0 +1,1194 @@ +From fcc48b03f454729fbf5b29090ba228112ad1d153 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:31:20 -0400 +Subject: [PATCH 043/103] topk: large-path hardening + +- Inline-PTX acquire/release helpers declare their memory effects. +- The fault verdict is aggregated across the CTA group before emission, so a + faulted row is failed closed by every CTA. +- 64-bit device counters and an int64 reader (counter protocol v2). +- Tests: adversarial distributions, slot-residue coverage, an sm_90 gate for + tests that need it. +--- + csrc/libtorch_stable/persistent_topk.cuh | 267 +++++++++--- + csrc/libtorch_stable/topk.cu | 110 ++++- + tests/kernels/test_top_k_per_row.py | 407 +++++++++++++++++- + .../qwen4_exp/nvidia/ops/topk_counters.py | 18 +- + 4 files changed, 713 insertions(+), 89 deletions(-) + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 8dd3f7a..fcbcc7b 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -146,12 +146,19 @@ struct PersistentTopKParams { + // is a process-lifetime device buffer of kNumTopKExactCounters uint32s, or + // nullptr; the kernel only ever atomicAdds into it, never reads back. + bool exact_medium; +- unsigned int* exact_counters; ++ unsigned long long* exact_counters; + // MERLIN_TOPK_DET_LARGE, ANDed on the host with the capacity gate + // 2 * ctas_per_group <= RADIX (the per-CTA counts are staged in one 256-slot + // RadixRowState histogram row). Resolved once per process, so it is a + // launch constant and a cudagraph captures one value. + bool det_large; ++ // TEST ONLY. -1 (the default) disables it; any ++ // other value names the `cta_in_group` whose CTA-LOCAL fault predicate is ++ // forced, so a fault originating outside CTA 0 can be exercised. Resolved ++ // once per process from MERLIN_TOPK_LARGE_FAULT_INJECT, so it is a launch ++ // constant like the two switches. Setting it also suppresses __trap(), because ++ // the point of injecting is to observe the trap-off row-failure policy. ++ int fault_inject_cta; + }; + + // ============================================================================ +@@ -468,6 +475,14 @@ __device__ __noinline__ void histogram_2048_topk( + // no host sync and no host-visible allocation on the launch path, so this is + // safe to record inside a cudagraph capture and on every replay. Read from + // Python with `torch.ops._C.persistent_topk_exact_counters`. ++// ++// The counters are `unsigned long long`, not `unsigned int`, ++// and the reader tensor is int64. A u32 counter copied into a SIGNED int32 ++// tensor goes negative, then wraps: at 65,536 columns with k=512 an all-equal ++// row drops 65,024 ties, so slot 6 crosses the signed boundary after 33,027 ++// rows and wraps after 66,053, and a delta taken across that boundary is a ++// wrong number that looks like a measurement. `atomicAdd(unsigned long long*)` ++// is native on sm_86, so the width costs nothing on the increment path. + enum TopKExactCounter : int { + kCtrMediumRows = 0, // denominator: medium-path rows entered + kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop +@@ -477,7 +492,7 @@ enum TopKExactCounter : int { + // Slots 4-7 are only ever touched by the deterministic emission, so they + // stay 0 in the MERLIN_TOPK_DET_LARGE=0 control arm: `[9] == 1 && [4] == 0` + // while large rows are known to have run means the host-side capacity gate +- // (2 * ctas_per_group <= RADIX) turned the switch off for that launch. ++ // (4 * ctas_per_group <= RADIX) turned the switch off for that launch. + kCtrLargeDetRows = 4, // large rows emitted deterministically + kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut + kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts +@@ -494,7 +509,7 @@ enum TopKExactCounter : int { + template + __device__ __noinline__ void medium_exact_restart( + const float* __restrict__ scores, int32_t* __restrict__ output, +- uint32_t length, unsigned int* __restrict__ counters) { ++ uint32_t length, unsigned long long* __restrict__ counters) { + static_assert(kThreadsPerBlock >= RADIX, + "the single-thread bin walk assumes RADIX <= blockDim.x"); + +@@ -759,7 +774,7 @@ __device__ __noinline__ void medium_exact_restart( + // ---- fail loud --------------------------------------------------------- + if (row_fault != 0u) { + if (tx == 0 && counters != nullptr) { +- atomicAdd(&counters[kCtrMediumExactFault], 1u); ++ atomicAdd(&counters[kCtrMediumExactFault], 1ull); + } + // -1 is the kernel's own padding value and the consumer rejects it + // (`valid = ... & (token >= 0)`), so a faulted row attends nothing and is +@@ -790,7 +805,7 @@ template + __device__ __noinline__ void histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, + int logits_offset, int seq_len, bool exact_medium, +- unsigned int* __restrict__ counters) { ++ unsigned long long* __restrict__ counters) { + // No default arguments on purpose: a new call site must state whether it + // wants the exact path, rather than silently inheriting the broken one. + // All shared state lives in dynamic shared memory to avoid static +@@ -815,7 +830,7 @@ __device__ __noinline__ void histogram_256_topk( + + if (thread_id == 0) { + shared_overflow_counted = 0; +- if (counters != nullptr) atomicAdd(&counters[kCtrMediumRows], 1u); ++ if (counters != nullptr) atomicAdd(&counters[kCtrMediumRows], 1ull); + } + + if (thread_id < RADIX + 1) { +@@ -891,7 +906,7 @@ __device__ __noinline__ void histogram_256_topk( + // Past the early return the threshold bin WILL be collected, so the M1 + // prediction is now a real overflow event. + if (thread_id == 0 && shared_overflow_counted == 1 && counters != nullptr) { +- atomicAdd(&counters[kCtrMediumOverflow], 1u); ++ atomicAdd(&counters[kCtrMediumOverflow], 1ull); + } + + __syncthreads(); +@@ -951,11 +966,11 @@ __device__ __noinline__ void histogram_256_topk( + if (__builtin_expect(raw_buffered > MAX_BUFFERED_ITEMS, 0)) { + if (thread_id == 0 && shared_overflow_counted == 0) { + shared_overflow_counted = 1; +- if (counters != nullptr) atomicAdd(&counters[kCtrMediumOverflow], 1u); ++ if (counters != nullptr) atomicAdd(&counters[kCtrMediumOverflow], 1ull); + } + if (exact_medium) { + if (thread_id == 0 && counters != nullptr) { +- atomicAdd(&counters[kCtrMediumExactFallback], 1u); ++ atomicAdd(&counters[kCtrMediumExactFallback], 1ull); + } + __syncthreads(); // no thread may clobber medium_smem before this + medium_exact_restart(logits + logits_offset, output_indices, +@@ -1036,24 +1051,55 @@ __device__ __noinline__ void histogram_256_topk( + // Inter-CTA sync primitives + // ============================================================================ + ++// `asm volatile` alone is NOT a compiler memory barrier. It ++// stops the statement being deleted or duplicated and keeps volatile asms in ++// order relative to one another, but it says nothing about the surrounding ++// loads and stores, so nvcc was free to sink a publication store past the ++// release fence or hoist a payload load above the acquire load. NVIDIA ++// documents exactly this in "Inline PTX Assembly: Incorrect Optimization". ++// Every helper below therefore declares a `"memory"` clobber, and the fence is ++// fused into the SAME asm statement as the atomic it orders so that no ++// optimizer pass can separate them. ++// ++// Why inline PTX and not `cuda::atomic_ref` / `__threadfence()`: ++// * `__threadfence()` compiles to `membar.gl` on this toolchain (verified: ++// nvcc 12.8 -arch=sm_86 -ptx), which is the sequentially-consistent device ++// fence, strictly stronger than the `fence.acq_rel.gpu` this protocol ++// needs. Swapping it in would change the emitted ordering primitive on ++// the hottest barrier in the kernel, which is a performance change smuggled ++// into a correctness fix. ++// * `cuda::atomic_ref` expresses the same ++// acquire/release semantics, but it would pull libcu++ into a header that ++// is compiled for every arch this extension builds, and `red` (a reduction ++// with no return value) has no atomic_ref spelling -- `fetch_add` returns ++// the old value and costs the return path. ++// So: keep the PTX, and make it honest about its memory effects. + __device__ __forceinline__ int ld_acquire(int* ptr) { + int state = 0; + #if (__CUDA_ARCH__ >= 700) + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) +- : "l"(ptr)); ++ : "l"(ptr) ++ : "memory"); + #else +- asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); ++ asm volatile("ld.cg.global.b32 %0, [%1];\n" ++ : "=r"(state) ++ : "l"(ptr) ++ : "memory"); + #endif + return state; + } + + __device__ __forceinline__ void red_release(int* ptr, int val) { + #if (__CUDA_ARCH__ >= 700) +- asm volatile("fence.acq_rel.gpu;\n"); +- asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" +- : +- : "l"(ptr), "r"(val)); ++ // One statement: the fence and the arrival add cannot be separated, and the ++ // clobber forbids moving any access across either of them. ++ asm volatile( ++ "fence.acq_rel.gpu;\n" ++ "red.relaxed.gpu.global.add.s32 [%0], %1;\n" ++ : ++ : "l"(ptr), "r"(val) ++ : "memory"); + #else + __threadfence(); + atomicAdd(ptr, val); +@@ -1062,8 +1108,12 @@ __device__ __forceinline__ void red_release(int* ptr, int val) { + + __device__ __forceinline__ void st_release(int* ptr, int val) { + #if (__CUDA_ARCH__ >= 700) +- asm volatile("fence.acq_rel.gpu;\n"); +- asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val)); ++ asm volatile( ++ "fence.acq_rel.gpu;\n" ++ "st.release.gpu.global.b32 [%0], %1;\n" ++ : ++ : "l"(ptr), "r"(val) ++ : "memory"); + #else + __threadfence(); + atomicExch(ptr, val); +@@ -1126,11 +1176,20 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, + // anywhere on the placement path — so the row is strictly ascending in + // index and bit-reproducible. + // +-// Cost per row: one extra inter-CTA barrier, two extra ballots and one +-// 32-lane shuffle scan per 1,024-element tile, 2 global stores and +-// 2*ctas_per_group global loads. No second pass over the row in global ++// * fault handling is fail-closed for the WHOLE GROUP: each CTA publishes a ++// verdict word and every CTA ORs all of them (X3) before any CTA emits, ++// and the emission bound is re-aggregated the same way (X4), so a fault ++// raised in CTA 1 still produces an all-`-1` row and exactly one counter ++// increment. ++// ++// Cost per row: FOUR inter-CTA barriers where the v2.0.1 tail runs one (X1 ++// slot reuse, X2 counts, X3 fault verdict, X4 emission verdict), two extra ++// ballots and one 32-lane shuffle scan per 1,024-element tile, 4 global stores ++// and 4*ctas_per_group global loads. No second pass over the row in global + // memory: the pivot-equal count that the tie cut needs is produced by the same +-// shared-memory pass v2.0.1 already runs to count `> pivot`. ++// shared-memory pass v2.0.1 already runs to count `> pivot`. X3/X4 are the ++// price of a group-uniform fail-closed contract; at ctas_per_group 2-3 each is ++// one release add and a short spin. + // ============================================================================ + template + __device__ __noinline__ void radix_emit_deterministic( +@@ -1139,14 +1198,27 @@ __device__ __noinline__ void radix_emit_deterministic( + uint32_t* local_histogram, uint32_t* suffix_sum, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, uint32_t ctas_per_group, + int& barrier_phase, uint32_t radix_iter, uint32_t tx, +- unsigned int* counters) { ++ unsigned long long* counters, int fault_inject_cta) { + constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u; + // Staging for the published per-CTA counts. The host gates the switch on +- // 2 * ctas_per_group <= RADIX, so these two halves never collide. ++ // 4 * ctas_per_group <= RADIX, so these regions never collide. + constexpr uint32_t kEqStage = RADIX / 2u; +- ++ // Two more regions of the same published histogram row ++ // carry the per-CTA fault verdicts -- the pre-emission one at [2C, 3C) and ++ // the emission-bound one at [3C, 4C). They are SEPARATE regions on purpose: ++ // reusing [2C, 3C) for the second round would let a fast CTA overwrite a word ++ // a slow CTA had not yet read. Each CTA writes only its own word, so nothing ++ // has to pre-zero them. Host gate: 4 * ctas_per_group <= RADIX. ++ const uint32_t kFaultStage = 2u * ctas_per_group; ++ const uint32_t kEmitStage = 3u * ctas_per_group; ++ ++ // suffix_sum scratch map (shared, RADIX words): ++ // [0] gt count [1] eq count [2] gt_base [3] eq_keep_base [4] cut_rank ++ // [5] group fault verdict [6] this CTA's emission-bound violation ++ // [7] group emission-bound verdict [8] ties dropped on this row + // ---- per-CTA counts, out of the resident chunk ------------------------- + if (tx < 2u) suffix_sum[tx] = 0u; ++ if (tx == 6u) suffix_sum[6] = 0u; + __syncthreads(); + + uint32_t my_gt = 0u; +@@ -1200,7 +1272,11 @@ __device__ __noinline__ void radix_emit_deterministic( + // reads with `ld.global.acquire.gpu`, so the two plain stores above are + // visible at device scope to every CTA that leaves this barrier. This is + // the same release/acquire pairing the shipped kernel already relies on for +- // CTA-0's plain-store zeroing of `next_hist`. ++ // CTA-0's plain-store zeroing of `next_hist`. Caveat: that ++ // argument is about the GPU memory model and only holds if the COMPILER also ++ // emits the stores before the fence; the `"memory"` clobber on red_release ++ // is what makes that true, and the PTX for this site is the evidence quoted ++ // in the commit message. + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } +@@ -1244,31 +1320,75 @@ __device__ __noinline__ void radix_emit_deterministic( + } + const uint32_t my_keep = local_histogram[kEqStage + cta_in_group]; + +- // Invariants, all three derived from the medium-path restart design and re-derivable +- // here: the radix walk guarantees #{key > pivot} == TopK - equal_needed +- // and equal_needed <= #{key == pivot}, and the two together make the +- // emission a bijection onto [0, TopK). Every CTA evaluates this from the +- // same published integers, so the verdict is group-uniform and needs no +- // further communication. ++ // Invariants. The radix walk guarantees #{key > pivot} == TopK - ++ // equal_needed and equal_needed <= #{key == pivot}, and the two together ++ // make the emission a bijection onto [0, TopK). ++ // ++ // The first two ARE group-uniform -- every CTA derives ++ // them from the same published integers in the same fixed order -- but the ++ // third, `my_keep > local_eq`, compares a value that came back through ++ // global memory against this CTA's own shared-memory count, and a local ++ // inconsistency is exactly what it exists to detect. It can therefore be ++ // true in CTA 1 and false in CTA 0, and the previous code let CTA 0 emit a ++ // perfectly ordinary-looking row while CTA 1 silently took the fault arm. ++ // The verdict is now published and ORed across the group below (X3) before ++ // ANY CTA emits. The fourth check is new and bounds this CTA's highest ++ // emission slot ahead of time, so the per-element `slot < TopK` test in the ++ // loop is a backstop rather than the only guard. + uint32_t fault = 0u; + if (gt_total + equal_needed != static_cast(TopK)) fault = 1u; + if (equal_needed > eq_total) fault = 1u; +- if (my_keep > local_eq) fault = 1u; ++ uint32_t my_keep_checked = my_keep; ++ // Test-only: force the CTA-LOCAL predicate in one named CTA, so a fault ++ // that originates outside CTA 0 can be observed. Off (-1) in production; ++ // the host resolves it once per process, so it is a launch constant. ++ if (fault_inject_cta >= 0 && ++ static_cast(cta_in_group) == fault_inject_cta) { ++ my_keep_checked = local_eq + 1u; ++ } ++ if (my_keep_checked > local_eq) fault = 1u; ++ if (gt_base + eq_keep_base + local_gt + my_keep > ++ static_cast(TopK)) { ++ fault = 1u; ++ } + + suffix_sum[2] = gt_base; + suffix_sum[3] = eq_keep_base; + // Keep the local pivot-equal elements whose ascending rank within this + // chunk is >= cut_rank, i.e. the `my_keep` highest-indexed of them. +- suffix_sum[4] = local_eq - my_keep; +- suffix_sum[5] = fault; ++ suffix_sum[4] = (my_keep <= local_eq) ? (local_eq - my_keep) : 0u; ++ suffix_sum[8] = (eq_total > equal_needed) ? (eq_total - equal_needed) : 0u; ++ ++ // Publish this CTA's verdict. Own word, so no pre-zeroing is needed. ++ cta_counts[kFaultStage + cta_in_group] = fault; ++ } ++ ++ // ---- barrier X3: the fault verdicts are published ---------------------- ++ // Unconditional, so the per-CTA barrier count stays group-uniform. ++ if (tx == 0) { ++ red_release(&state->arrival_counter, 1); ++ } ++ wait_ge(&state->arrival_counter, ++ (barrier_phase + 1) * static_cast(ctas_per_group), tx); ++ barrier_phase++; ++ __syncthreads(); ++ ++ if (tx == 0) { ++ uint32_t row_fault = 0u; ++ for (uint32_t c = 0; c < ctas_per_group; ++c) { ++ row_fault |= cta_counts[kFaultStage + c]; ++ } ++ suffix_sum[5] = row_fault; + + if (counters != nullptr && cta_in_group == 0u) { +- atomicAdd(&counters[kCtrLargeDetRows], 1u); +- if (eq_total > equal_needed) { +- atomicAdd(&counters[kCtrLargeTiesCut], 1u); +- atomicAdd(&counters[kCtrLargeTiesDropped], eq_total - equal_needed); ++ atomicAdd(&counters[kCtrLargeDetRows], 1ull); ++ // Only meaningful on a healthy row: on a faulted one the tie arithmetic ++ // is what is under suspicion, so do not record it as a measurement. ++ if (row_fault == 0u && suffix_sum[8] > 0u) { ++ atomicAdd(&counters[kCtrLargeTiesCut], 1ull); ++ atomicAdd(&counters[kCtrLargeTiesDropped], ++ static_cast(suffix_sum[8])); + } +- if (fault != 0u) atomicAdd(&counters[kCtrLargeFault], 1u); + } + } + __syncthreads(); +@@ -1347,10 +1467,13 @@ __device__ __noinline__ void radix_emit_deterministic( + const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before; + if (slot < static_cast(TopK)) { + row_output[slot] = static_cast(my_chunk_start + i); +- } else if (counters != nullptr) { +- // Unreachable while the three invariants above hold; bounded store +- // plus a counter rather than an out-of-range write. +- atomicAdd(&counters[kCtrLargeFault], 1u); ++ } else { ++ // An out-of-range slot used to bump a counter and ++ // carry on, so with trapping ON it neither trapped nor produced the ++ // promised all-`-1` row. It is now recorded as a row failure and ++ // goes through the SAME policy as the pre-emission faults below: no ++ // store here, and the row is filled with -1 after the group agrees. ++ suffix_sum[6] = 1u; + } + } + +@@ -1361,23 +1484,58 @@ __device__ __noinline__ void radix_emit_deterministic( + // No thread may start the next row's chunk load — which writes + // shared_ordered — until every warp has finished reading it here. + __syncthreads(); +- } else { +- // ---- fail loud ------------------------------------------------------ +- // The predicate is group-uniform (same published integers, same fixed +- // summation order in every CTA), so no CTA of this group emitted anything +- // and CTA 0 can fill the row alone. -1 is the kernel's own padding value +- // and the consumer rejects it, so a faulted row attends nothing instead of +- // quietly attending the wrong blocks. ++ } ++ ++ // ---- barrier X4: the emission verdicts are published ------------------- ++ // Unconditional (both arms of the branch above reach it), so the per-CTA ++ // barrier count stays group-uniform. On the healthy path this costs one ++ // inter-CTA barrier per large row and nothing else: every word written here ++ // is this CTA's own. ++ if (tx == 0) { ++ cta_counts[kEmitStage + cta_in_group] = suffix_sum[6]; ++ red_release(&state->arrival_counter, 1); ++ } ++ wait_ge(&state->arrival_counter, ++ (barrier_phase + 1) * static_cast(ctas_per_group), tx); ++ barrier_phase++; ++ __syncthreads(); ++ ++ if (tx == 0) { ++ uint32_t emit_fault = 0u; ++ for (uint32_t c = 0; c < ctas_per_group; ++c) { ++ emit_fault |= cta_counts[kEmitStage + c]; ++ } ++ suffix_sum[7] = emit_fault; ++ } ++ __syncthreads(); ++ ++ // ---- fail closed, for the whole group ---------------------------------- ++ // Either verdict is now group-uniform, so every CTA agrees on whether this ++ // row failed, and CTA 0 can overwrite the row alone -- including the case ++ // where some other CTA had already emitted before an out-of-range slot was ++ // discovered. -1 is the kernel's own padding value and the consumer rejects ++ // it, so a faulted row attends nothing instead of quietly attending the ++ // wrong blocks. The counter is incremented exactly ONCE per failed row. ++ if (row_fault != 0u || suffix_sum[7] != 0u) { + if (cta_in_group == 0u) { + for (uint32_t i = tx; i < static_cast(TopK); + i += kThreadsPerBlock) { + row_output[i] = -1; + } ++ if (tx == 0 && counters != nullptr) { ++ atomicAdd(&counters[kCtrLargeFault], 1ull); ++ } + } + __threadfence(); + __syncthreads(); + #if MERLIN_TOPK_EXACT_TRAP +- __trap(); ++ // Fault injection is a test mode and deliberately suppresses the trap: ++ // the point of injecting is to OBSERVE the trap-off policy (all-`-1` row, ++ // counter += 1), which a context-killing __trap() would make unobservable ++ // without a second build. ++ if (fault_inject_cta < 0) { ++ __trap(); ++ } + #endif + } + } +@@ -1396,7 +1554,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, + RadixRowState* state, uint32_t cta_in_group, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t radix_iter, uint32_t tx, bool det_large, +- unsigned int* counters) { ++ unsigned long long* counters, int fault_inject_cta) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; +@@ -1556,7 +1714,8 @@ __device__ void radix_topk(const float* __restrict__ row_input, + radix_emit_deterministic( + row_output, ordered_pivot, shared_scalars[1], my_chunk_start, + actual_chunk_size, local_histogram, suffix_sum, shared_ordered, state, +- cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters); ++ cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters, ++ fault_inject_cta); + return; + } + +@@ -1709,7 +1868,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + row_input, row_output, seq_len, my_chunk_start, chunk_size, + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, +- params.det_large, params.exact_counters); ++ params.det_large, params.exact_counters, params.fault_inject_cta); + radix_iter++; + } + } +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 001c782..07d166a 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -55,6 +55,30 @@ bool det_large_enabled() { + return enabled; + } + ++// --------------------------------------------------------------------------- ++// MERLIN_TOPK_LARGE_FAULT_INJECT. TEST ONLY. ++// ++// -1 (the default, and the value of an unset or negative variable) is off. ++// Any other value names the `cta_in_group` whose CTA-LOCAL fault predicate the ++// large path forces, which is the only way to reach the fail-closed arm: no ++// legal input can make a CTA's own pivot-equal count disagree with the one it ++// published, so those branches are otherwise untestable (as with the ++// medium-path restart's own fault path). Resolved once per process, so ++// it is a launch constant and a cudagraph captures one value. Setting it also ++// suppresses __trap(), because the point of injecting is to OBSERVE the ++// trap-off row-failure policy (all -1 row, counter += 1) without a second ++// build. ++// --------------------------------------------------------------------------- ++int large_fault_inject_cta() { ++ static const int cta = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_LARGE_FAULT_INJECT"); ++ if (raw == nullptr || raw[0] == '\0') return -1; ++ const int v = std::atoi(raw); ++ return (v < 0) ? -1 : v; ++ }(); ++ return cta; ++} ++ + // Process-lifetime device counters, one block per device. The kernel only + // ever atomicAdds into them, so the launch path needs no host sync and the + // values accumulate across cudagraph replays. +@@ -85,12 +109,12 @@ enum CounterState : int { + // a second host thread may be publishing them from the slow path. Release on + // publish, acquire on read, and the pointer is always published BEFORE the + // state that advertises it. +-std::atomic g_exact_counters[kMaxCounterDevices] = {}; ++std::atomic g_exact_counters[kMaxCounterDevices] = {}; + std::atomic g_exact_counters_state[kMaxCounterDevices] = {}; + std::mutex g_exact_counters_mutex; + + // Lookup only. Never allocates, never blocks, safe from any context. +-unsigned int* get_exact_counters(int device_index) { ++unsigned long long* get_exact_counters(int device_index) { + if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr; + if (g_exact_counters_state[device_index].load(std::memory_order_acquire) != + kCountersReady) { +@@ -105,7 +129,7 @@ unsigned int* get_exact_counters(int device_index) { + // `stream_is_capturing` only proves the CALLING stream is idle of capture, so + // the intended entry point is the reader op called explicitly before capture; + // the non-capturing launch path calls this as a fallback. +-unsigned int* ensure_exact_counters(int device_index, cudaStream_t stream) { ++unsigned long long* ensure_exact_counters(int device_index, cudaStream_t stream) { + if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr; + + const int state = +@@ -129,8 +153,8 @@ unsigned int* ensure_exact_counters(int device_index, cudaStream_t stream) { + } + + constexpr size_t bytes = +- vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int); +- unsigned int* ptr = nullptr; ++ vllm::persistent::kNumTopKExactCounters * sizeof(unsigned long long); ++ unsigned long long* ptr = nullptr; + if (cudaMalloc(&ptr, bytes) == cudaSuccess && ptr != nullptr) { + // Host acquire/release orders the HOST state word; it says + // nothing about the device. cudaMemset on device memory is asynchronous +@@ -361,17 +385,20 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + params.exact_medium = exact_medium_enabled(); + // The per-CTA selected / pivot-equal counts are staged in ONE + // 256-slot RadixRowState histogram row (the slot round 3 has just finished +- // reading), which holds 2 * ctas_per_group uint32s. That is the whole +- // reason the workspace contract is unchanged, and it is also the only +- // limit: above 128 CTAs per group the deterministic path cannot be ++ // reading). Two more per-CTA regions were added for the ++ // group-uniform fault verdicts, so the row now holds 4 * ctas_per_group ++ // uint32s. That is still the whole reason the workspace contract is ++ // unchanged, and it is also the only ++ // limit: above 64 CTAs per group the deterministic path cannot be + // expressed there, so the launch degrades to the v2.0.1 tail rather than + // failing. Unreachable at any shape this kernel serves (ctas_per_group is + // 2-3 across our 135K-262K band and <= 16 at the smallest chunk size), but + // a silent loss of determinism is exactly what this branch exists to stop, + // so say it once. + const bool det_large_fits = +- 2u * ctas_per_group <= static_cast(P::RADIX); ++ 4u * ctas_per_group <= static_cast(P::RADIX); + params.det_large = det_large_enabled() && det_large_fits; ++ params.fault_inject_cta = large_fault_inject_cta(); + if (det_large_enabled() && !det_large_fits) { + static std::atomic det_large_warned{false}; + if (!det_large_warned.exchange(true)) { +@@ -380,7 +407,18 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + "needs %u of %d staging slots; the large path falls back " + "to nondeterministic placement at this shape (counter [4] " + "stays 0).\n", +- ctas_per_group, 2u * ctas_per_group, P::RADIX); ++ ctas_per_group, 4u * ctas_per_group, P::RADIX); ++ } ++ } ++ if (params.fault_inject_cta >= 0) { ++ static std::atomic inject_warned{false}; ++ if (!inject_warned.exchange(true)) { ++ std::fprintf(stderr, ++ "[vllm] MERLIN_TOPK_LARGE_FAULT_INJECT=%d is set: the " ++ "large path FORCES a fault in that CTA of every " ++ "deterministic large row and returns an all -1 row. This " ++ "is a test mode and must not be set in production.\n", ++ params.fault_inject_cta); + } + } + // Allocate only off the capture path; in the ordinary engine lifecycle the +@@ -487,7 +525,31 @@ void persistent_topk(const torch::stable::Tensor& logits, + // counters. This is also the explicit pre-capture init path: call it once per + // device before any cudagraph capture. + // +-// `out` is an int32 CUDA tensor with at least 10 elements: ++// PROTOCOL VERSION 2 -- BREAKING. `out` is an **int64** ++// CUDA tensor with at least 10 elements; int32 is REJECTED, not silently ++// accepted. The device counters are 64-bit unsigned and are copied D2D ++// without conversion, so the reader's element width must match the device's. ++// ++// Why the break rather than a reinterpreting reader: the counters were u32 ++// copied into a SIGNED int32 tensor. At 65,536 columns with k=512 an ++// all-equal row drops 65,024 ties, so [6] goes negative in Python after 33,027 ++// rows and wraps modulo 2^32 after 66,053 -- and a subtraction across the ++// signed boundary gives a wrong delta, which is worse than a missing number ++// because it looks like a measurement. Keeping u32 and reinterpreting in the ++// reader would need an enforced per-window increment bound plus quiescent ++// resets, i.e. a contract every future caller has to remember. 64-bit device ++// counters make the failure unreachable instead: `atomicAdd(unsigned long ++// long*)` is native on sm_86 (it is a single ATOM.E.ADD.S64, not a CAS loop), ++// so the instrument costs the same, and a 64-bit slot cannot be driven to the ++// signed boundary by any run this kernel will ever see. ++// ++// Everything that reads these counters must move together: the Python reader ++// `vllm/models/qwen4_exp/nvidia/ops/topk_counters.py` now allocates int64, and ++// the medium-path tests allocate int64. An int32 caller gets a loud error rather ++// than half a number. The OPERATOR SCHEMA is unchanged (dtype is not part of ++// it), so a stale caller fails at the check, not at the binding. ++// ++// Slots: + // [0] medium-path rows entered (denominator) + // [1] rows where a 4,096-slot guard would drop a candidate + // [2] rows that took the exact restart (0 in the control arm) +@@ -521,13 +583,17 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + int64_t sync) { + #ifndef USE_ROCM + STD_TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor"); +- STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Int, +- "out must be int32"); ++ STD_TORCH_CHECK( ++ out.scalar_type() == torch::headeronly::ScalarType::Long, ++ "out must be int64: the exact-restart counters are 64-bit as of the " ++ "large-path counter protocol change, and an int32 buffer would be filled " ++ "from half of each counter. Allocate torch.zeros(10, dtype=torch.int64)." ++ " (int32 was protocol version 1 and is deliberately rejected.)"); + // Element count alone is not enough. A CUDA scalar expanded +- // to five elements (stride 0) satisfies numel but owns one element's +- // storage, and the five consecutive raw int32 writes below would run past +- // it; an ordinary strided view would receive the values in the wrong slots. +- // Require a contiguous 1-D buffer before taking the pointer. ++ // to ten elements (stride 0) satisfies numel but owns one element's storage, ++ // and the raw int64 writes below would run past it; an ordinary strided view ++ // would receive the values in the wrong slots. Require a contiguous 1-D ++ // buffer before taking the pointer. + STD_TORCH_CHECK(out.dim() == 1, "out must be 1-D"); + STD_TORCH_CHECK(out.is_contiguous(), + "out must be contiguous (expanded or strided views would be " +@@ -547,10 +613,10 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + "persistent_topk_exact_counters must not be called during " + "cudagraph capture"); + constexpr size_t bytes = +- vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int); ++ vllm::persistent::kNumTopKExactCounters * sizeof(unsigned long long); + +- int32_t* dst = out.mutable_data_ptr(); +- unsigned int* src = ++ int64_t* dst = out.mutable_data_ptr(); ++ unsigned long long* src = + ensure_exact_counters(static_cast(out.get_device_index()), stream); + + // An unavailable buffer used to be reported as four zeros +@@ -584,11 +650,11 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + // flag reports the ENV value, not the per-launch effective value: a launch + // whose ctas_per_group overflowed the staging row still reports 1 here and + // leaves counter [4] at 0, which is how that degradation is diagnosed. +- static const int32_t enabled[2] = {exact_medium_enabled() ? 1 : 0, ++ static const int64_t enabled[2] = {exact_medium_enabled() ? 1 : 0, + det_large_enabled() ? 1 : 0}; + STD_TORCH_CHECK( + cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, enabled, +- 2 * sizeof(int32_t), cudaMemcpyHostToDevice, ++ 2 * sizeof(int64_t), cudaMemcpyHostToDevice, + stream) == cudaSuccess, + "exact counter flag write failed"); + +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index daf8ff6..9c775b0 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1,6 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++import os + import re + + import numpy as np +@@ -1080,10 +1081,20 @@ def run_large_context_topk_test( + Torch: {torch_vals.sort(descending=True)[0][:10]}""" + + +-@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+") + @pytest.mark.parametrize("top_k", [512, 1024, 2048]) + def test_cooperative_topk_cs2(top_k: int) -> None: +- """The 64-row dispatch uses the two-CTA cooperative kernel.""" ++ """The 64-row dispatch uses the two-CTA cooperative kernel. ++ ++ This gate used to be a bare "CUDA available" check, ++ which then called `torch.ops._C.cooperative_topk` -- an operation the ++ sm_86 build does not contain -- and failed. That is a test-gating defect, ++ not an sm_86 engine dependency: the sm_86 QSA route has no transitive call ++ to `cooperative_topk` (the C++ occupancy fallback calls ++ `top_k_per_row_decode`, a different operation). Use the same SM90+ gate the ++ `cooperative_topk` backend parameter already uses, and do NOT rebuild an ++ unused architecture to make these three pass. ++ """ + run_large_context_topk_test( + batch_size=64, + seq_lens=[65536] * 64, +@@ -1875,7 +1886,9 @@ def _read_exact_counters() -> list[int] | None: + op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) + if op is None: + return None +- out = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") ++ # int64: the large-path determinism change widened the device counters to 64 bits ++ # and the op now rejects int32. See topk.cu's protocol-version-2 note. ++ out = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int64, device="cuda") + # sync=1: the counters are device-global but the snapshot is issued on one + # stream, so without a device synchronize it can miss increments from work + # still in flight elsewhere. The tests compare before/after deltas, so a +@@ -2073,13 +2086,19 @@ def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") + @torch.inference_mode() + def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: +- """The counter op writes five raw int32s; it must refuse unusable layouts. ++ """The counter op writes ten raw int64s; it must refuse unusable layouts. + +- A CUDA scalar expanded to five elements satisfies a numel check while ++ A CUDA scalar expanded to ten elements satisfies a numel check while + owning one element of storage, and a strided view would receive the values + in the wrong slots. Both must be rejected before the raw pointer is taken. + `reset` without `sync` must also be rejected: zeroing while kernels are in + flight drops their atomics silently. ++ ++ The dtype case is a large-path protocol break: the device ++ counters are 64-bit, so an int32 buffer would be filled from half of each ++ counter. It must be REJECTED, not quietly accepted -- that rejection is ++ how a stale caller finds out, and asserting it here is what makes the ++ break a contract rather than a comment. + """ + torch.set_default_device("cuda:0") + op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) +@@ -2089,22 +2108,26 @@ def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: + "persistent_topk_exact_counters op. A skip is not a pass." + ) + +- ok = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") ++ ok = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int64, device="cuda") + op(ok, 0, 1) # the control: this layout must be accepted + +- expanded = torch.zeros(1, dtype=torch.int32, device="cuda").expand( ++ wrong_dtype = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") ++ with pytest.raises(Exception): ++ op(wrong_dtype, 0, 1) ++ ++ expanded = torch.zeros(1, dtype=torch.int64, device="cuda").expand( + _EXACT_COUNTER_SLOTS + ) + with pytest.raises(Exception): + op(expanded, 0, 1) + +- strided = torch.zeros(2 * _EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda")[ ++ strided = torch.zeros(2 * _EXACT_COUNTER_SLOTS, dtype=torch.int64, device="cuda")[ + ::2 + ] + with pytest.raises(Exception): + op(strided, 0, 1) + +- too_small = torch.zeros(_EXACT_COUNTER_SLOTS - 1, dtype=torch.int32, device="cuda") ++ too_small = torch.zeros(_EXACT_COUNTER_SLOTS - 1, dtype=torch.int64, device="cuda") + with pytest.raises(Exception): + op(too_small, 0, 1) + +@@ -2378,3 +2401,369 @@ def test_persistent_topk_large_control_arm_is_untouched() -> None: + assert torch.equal(got_values, ref_values), ( + f"row {row}: control-arm values are not torch.topk's values" + ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_fault_in_nonzero_cta_fails_closed() -> None: ++ """A fault raised OUTSIDE CTA 0 must still fail the whole row closed. ++ ++ `my_keep > local_eq` is a CTA-LOCAL predicate -- it ++ compares a value that came back through global memory against this CTA's ++ own shared count -- so it can be true in CTA 1 and false in CTA 0. Before ++ the fix CTA 1 took the fault arm while CTA 0 emitted an ordinary-looking ++ row and recorded nothing: with trapping off, the promised all -1 row was ++ never produced. ++ ++ No legal input can make a CTA disagree with itself, so the fault arm is ++ only reachable through the test-only injector. Arm it for the whole ++ process: `MERLIN_TOPK_LARGE_FAULT_INJECT` is read into a function-local ++ static on the first launch, exactly like the two switches, so it CANNOT be ++ monkeypatched inside a session that has already run a top-k kernel. Run ++ this test in its own process: ++ ++ MERLIN_TOPK_DET_LARGE=1 MERLIN_TOPK_LARGE_FAULT_INJECT=1 pytest -k large_fault_in_nonzero_cta ++ ++ Setting the injector also suppresses __trap(), so this is the trap-off arm ++ of the policy. The trap-on arm kills the CUDA context by construction and ++ belongs in a disposable-process GPU run, not here. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows = 8 ++ ++ inject = os.getenv("MERLIN_TOPK_LARGE_FAULT_INJECT") ++ if inject is None or inject.strip() in ("", "-1"): ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build with the fault " ++ "cells unrun: MERLIN_TOPK_LARGE_FAULT_INJECT is not armed, so the " ++ "group-uniform fail-closed path is not exercised. A skip is not a " ++ "pass; run this test in its own process with the variable set to a " ++ "NONZERO CTA index." ++ ) ++ assert int(inject) > 0, ( ++ "this test exists to prove a fault originating outside CTA 0 fails the " ++ f"row closed; MERLIN_TOPK_LARGE_FAULT_INJECT={inject} is CTA 0" ++ ) ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ logits, _ = _build_largedet_tie_row(num_rows) ++ out = _run_largedet_row(logits) ++ after = _read_exact_counters() ++ ++ # 1. Every row is fully -1. Not "starts with" and not "contains": the whole ++ # row, because a partially emitted row is the defect. ++ assert int(out.min()) == -1 and int(out.max()) == -1, ( ++ "a row that faulted in a nonzero CTA was not filled with -1: " ++ f"min={int(out.min())} max={int(out.max())}" ++ ) ++ ++ # 2. The fault counter moves exactly once per failed row -- not once per ++ # CTA, and not once per out-of-range element. ++ assert after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == num_rows, ( ++ "the large-path fault counter did not increment exactly once per row: " ++ f"{after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT]} of {num_rows}" ++ ) ++ ++ # 3. The row was still counted as entered, and the tie measurement was NOT ++ # recorded: on a faulted row the tie arithmetic is what is under ++ # suspicion, so recording it would contaminate the measurement. ++ assert after[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_rows, ( ++ "faulted rows must still be counted as large-path rows" ++ ) ++ assert after[_SLOT_LARGE_TIES_CUT] - before[_SLOT_LARGE_TIES_CUT] == 0, ( ++ "a faulted row recorded a tie-cut measurement" ++ ) ++ assert after[_SLOT_LARGE_TIES_DROPPED] - before[_SLOT_LARGE_TIES_DROPPED] == 0, ( ++ "a faulted row recorded dropped-tie counts" ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# Large-path determinism -- expanded coverage (more shapes and ++# adversarial distributions). ++# ++# The original pair of tests used one shape (8 rows x 40,960, k=512) and one ++# distribution. That gives ONE large row per group, ++# so it never exercises the `histogram[radix_iter % 3]` slot reuse that the ++# X1 barrier exists to protect, and that a single planted tie shape is not a ++# distribution class. Both are addressed below: ++# ++# * `_largedet_reference` derives the expected index set from the SHIPPED RULE for ++# any row, so every distribution gets an exact set assertion rather than ++# torch.topk's arbitrary choice inside a tie. The rule: take every element ++# strictly above the k-th largest VALUE, then the highest-indexed of those ++# equal to it. float32 equality is exactly the kernel's ordered-uint32 ++# equality (the transform is monotonic and injective), so the reference is ++# exact, not approximate. ++# * five distributions: all-equal, no cutoff tie, cross-CTA ties, a cutoff ++# landing inside a single 1,024-element tile, and narrow distinct scores ++# that all share one coarse bin. ++# * rows 1 / 4 / 8 / 32, which straddle both shared-memory budget steps in ++# topk.cu (<=4, <=8, >8) and therefore different chunk sizes and ++# ctas_per_group values; k = 512 and 2,048. ++# * a ragged cell mixing trivial / decode / medium / large rows in one launch. ++# * a slot-residue cell whose group geometry forces >= 3 large rows per group. ++# --------------------------------------------------------------------------- ++ ++LARGEDET_RADIX_THRESHOLD = 32768 # persistent_topk.cuh: above this, the large path ++ ++ ++def _largedet_reference(values: torch.Tensor, top_k: int) -> torch.Tensor: ++ """The deterministic large path's own rule, as ascending int64 indices. ++ ++ Everything strictly above the k-th largest value, plus the HIGHEST-indexed ++ of the elements equal to it (most-recent-wins). Computed on CPU float32. ++ float32 equality is exactly the kernel's ordered-uint32 equality, because ++ the transform is monotonic and injective, so this reference is exact. ++ """ ++ assert values.dtype == torch.float32 and values.dim() == 1 ++ pivot = torch.sort(values, descending=True).values[top_k - 1] ++ above = (values > pivot).nonzero(as_tuple=False).flatten() ++ equal = (values == pivot).nonzero(as_tuple=False).flatten() ++ need = top_k - int(above.numel()) ++ assert need >= 0, "more elements above the pivot than k; reference is wrong" ++ kept = equal[equal.numel() - need :] if need > 0 else equal[:0] ++ expected = torch.sort(torch.cat([above, kept])).values.to(torch.int64) ++ assert expected.numel() == top_k ++ return expected ++ ++ ++def _largedet_distribution(kind: str, seq_len: int, top_k: int, seed: int) -> torch.Tensor: ++ """One CPU float32 row. All values are positive, so no -0.0/+0.0 hazard.""" ++ gen = torch.Generator(device="cpu").manual_seed(seed) ++ ++ if kind == "all_equal": ++ # Every element ties. gt_total == 0 and equal_needed == k, so the whole ++ # selection is the tie cut and both halves of the cut walk run. ++ return torch.full((seq_len,), 1.5, dtype=torch.float32) ++ ++ if kind == "no_cutoff_tie": ++ # All distinct: the pivot group has exactly one element, so the tie cut ++ # is a no-op and only the placement arithmetic is under test. ++ vals = torch.randperm(seq_len, generator=gen).to(torch.float32) + 1.0 ++ return vals / seq_len + 0.5 ++ ++ if kind == "narrow_distinct": ++ # Distinct, but inside one coarse bin: the radix walk's own adversarial ++ # class (the medium path's #51782 shape, at large width). ++ steps = torch.randperm(seq_len, generator=gen).to(torch.float32) ++ return 1.0 + steps * (1.0 / (1 << 22)) ++ ++ row = torch.rand(seq_len, generator=gen, dtype=torch.float32) * 0.5 ++ ++ if kind == "cross_cta_ties": ++ # The tie group is scattered over the whole row by a seeded permutation, ++ # so it spans every CTA chunk and several index high-byte buckets, and ++ # the global cut lands mid-chunk. ++ above = top_k // 3 ++ tie = top_k # > k - above, so the cut is strictly inside the tie group ++ perm = torch.randperm(seq_len, generator=gen) ++ row[perm[:above]] = 3.0 ++ row[perm[above : above + tie]] = 2.0 ++ return row ++ ++ if kind == "cutoff_in_tile": ++ # The whole tie group is one contiguous 1,024-element window, i.e. one ++ # tile of one CTA, so the cut falls inside a single ballot/scan tile. ++ above = top_k - 512 ++ tie = 1024 ++ start = (seq_len // 2) & ~1023 ++ perm = torch.randperm(seq_len, generator=gen) ++ picked = perm[perm >= start + tie][:above] ++ if int(picked.numel()) < above: # pathological seed: use the head ++ picked = torch.arange(above) ++ row[picked] = 3.0 ++ row[start : start + tie] = 2.0 ++ return row ++ ++ raise ValueError(f"unknown large-path distribution: {kind}") ++ ++ ++def _largedet_launch(logits: torch.Tensor, lengths: torch.Tensor, top_k: int): ++ indices = torch.empty((logits.shape[0], top_k), dtype=torch.int32, device="cuda") ++ _run_topk_backend( ++ "persistent_topk", logits, lengths, indices, top_k, int(lengths.max()) ++ ) ++ torch.accelerator.synchronize() ++ return indices ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize( ++ "kind", ++ [ ++ "all_equal", ++ "no_cutoff_tie", ++ "cross_cta_ties", ++ "cutoff_in_tile", ++ "narrow_distinct", ++ ], ++) ++@pytest.mark.parametrize("num_rows", [1, 4, 8, 32]) ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_persistent_topk_large_distributions( ++ kind: str, num_rows: int, top_k: int ++) -> None: ++ """Exact set, exact order, exact counters and bit identity, per class.""" ++ torch.set_default_device("cuda:0") ++ seq_len = 40960 # > RADIX_THRESHOLD: the multi-CTA large path ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ cpu_row = _largedet_distribution(kind, seq_len, top_k, seed=19) ++ expected = _largedet_reference(cpu_row, top_k) ++ logits = cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ ++ first = _largedet_launch(logits, lengths, top_k).clone() ++ mid = _read_exact_counters() ++ second = _largedet_launch(logits, lengths, top_k) ++ ++ assert mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_rows, ( ++ "the deterministic large path did not run on every row: " ++ f"{mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS]} of {num_rows}" ++ ) ++ assert mid[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0, ( ++ "a large-path invariant fired; nothing below is meaningful" ++ ) ++ ++ ref_values = torch.topk(logits[0], top_k).values.sort().values ++ for row in range(num_rows): ++ got = first[row].to(torch.int64) ++ # Raw order and selected set are asserted SEPARATELY: a ++ # sorted-set comparison alone cannot see a placement regression. ++ assert bool((got[1:] - got[:-1] > 0).all()), ( ++ f"{kind} row {row}: output is not ascending in index" ++ ) ++ assert int(got.min()) >= 0 and int(got.max()) < seq_len, ( ++ f"{kind} row {row}: index out of range" ++ ) ++ assert got.unique().numel() == top_k, f"{kind} row {row}: duplicate index" ++ assert torch.equal(got.cpu(), expected), ( ++ f"{kind} row {row}: selected set is not the most-recent-wins set" ++ ) ++ assert torch.equal(logits[row][got].sort().values, ref_values), ( ++ f"{kind} row {row}: selected values are not torch.topk's values" ++ ) ++ ++ assert torch.equal(first, second), ( ++ f"{kind}: two identical launches returned different index tensors" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_persistent_topk_large_ragged_mixed_paths(top_k: int) -> None: ++ """Trivial, decode, medium and large rows in ONE launch. ++ ++ A group that handles a short row between two large ones must not advance ++ `radix_iter` -- that is what makes the histogram-slot reuse argument hold ++ across interleaved traffic. Large rows get the exact set assertion; short ++ rows only get the value-multiset check, because their tie policy belongs to ++ the decode/medium paths, not to this branch. ++ """ ++ torch.set_default_device("cuda:0") ++ lengths_list = [ ++ top_k, # trivial: length == k ++ 40960, # large ++ 2048, # decode/short ++ 49152, # large ++ 12288, # medium ++ 40960, # large ++ LARGEDET_RADIX_THRESHOLD, # exactly at the threshold: NOT large ++ 57344, # large ++ ] ++ num_rows = len(lengths_list) ++ seq_len = max(lengths_list) ++ num_large = sum(1 for n in lengths_list if n > LARGEDET_RADIX_THRESHOLD) ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ cpu_row = _largedet_distribution("cross_cta_ties", seq_len, top_k, seed=1901) ++ logits = cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ lengths = torch.tensor(lengths_list, dtype=torch.int32, device="cuda") ++ ++ first = _largedet_launch(logits, lengths, top_k).clone() ++ mid = _read_exact_counters() ++ second = _largedet_launch(logits, lengths, top_k) ++ ++ assert mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_large, ( ++ "the large path ran on the wrong number of rows: " ++ f"{mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS]} of {num_large}" ++ " -- check the RADIX_THRESHOLD boundary row" ++ ) ++ assert mid[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0 ++ ++ for row, length in enumerate(lengths_list): ++ got = first[row].to(torch.int64) ++ prefix = cpu_row[:length] ++ ref_values = torch.topk(prefix, top_k).values.sort().values ++ got_values = prefix[got.cpu()].sort().values ++ assert torch.equal(got_values, ref_values), ( ++ f"row {row} (length {length}): values are not torch.topk's values" ++ ) ++ if length > LARGEDET_RADIX_THRESHOLD: ++ assert torch.equal(got.cpu(), _largedet_reference(prefix, top_k)), ( ++ f"row {row} (length {length}): not the most-recent-wins set" ++ ) ++ ++ assert torch.equal(first, second), ( ++ "two identical ragged launches returned different index tensors" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_slot_residue_reuse() -> None: ++ """Every group must process >= 3 large rows, covering all 3 histogram slots. ++ ++ `RadixRowState::histogram` is triple-buffered over four radix rounds, so ++ row iteration `r` reuses `histogram[r % 3]` for its final round AND for the ++ per-CTA count publication. The X1 barrier is what stops a fast CTA ++ overwriting that slot while a slow CTA still reads it, and the reuse is ++ only exercised when a group handles SEVERAL large rows in a row -- three, ++ to cover residues 0, 1 and 2. The original eight-row test gave one row per ++ group and never reached it. ++ ++ The geometry is chosen, not hoped for. topk.cu sets ++ `num_groups = min(max_resident_ctas / ctas_per_group, num_rows)`, so with ++ 32 rows of 262,144 columns: at `num_rows > 8` the chunk budget gives ++ ~24,824 elements, hence `ctas_per_group = ceil(262144/24824) = 11`; on a ++ 3090 (82 SMs, occupancy 1 for a 1,024-thread CTA at this smem) ++ `max_resident_ctas = 81`, so `num_groups = 7` and the round-robin gives ++ every group 4 or 5 large rows. If a device produces a wider grid the ++ coverage claim weakens but nothing here becomes wrong: the assertions are ++ on the property (exact set, exact order, bit identity), not on the geometry. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows = 32 ++ seq_len = 262144 ++ top_k = 512 ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ cpu_row = _largedet_distribution("cross_cta_ties", seq_len, top_k, seed=1903) ++ expected = _largedet_reference(cpu_row, top_k) ++ logits = cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ ++ first = _largedet_launch(logits, lengths, top_k).clone() ++ mid = _read_exact_counters() ++ second = _largedet_launch(logits, lengths, top_k) ++ ++ assert mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_rows ++ assert mid[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0, ( ++ "a large-path invariant fired during slot reuse" ++ ) ++ for row in range(num_rows): ++ got = first[row].to(torch.int64) ++ assert bool((got[1:] - got[:-1] > 0).all()), f"row {row}: not ascending" ++ assert torch.equal(got.cpu(), expected), ( ++ f"row {row}: slot reuse changed the selected set" ++ ) ++ assert torch.equal(first, second), ( ++ "two identical launches over reused histogram slots differed" ++ ) +diff --git a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +index 02c1720..08ffab8 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +@@ -2,12 +2,20 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Python reader for the medium-path exact-restart top-k counters. + +-The counters themselves are device-global uint32s owned by ++The counters themselves are device-global uint64s owned by + `csrc/libtorch_stable/topk.cu`; `torch.ops._C.persistent_topk_exact_counters` +-snapshots them into an int32 CUDA tensor. Without this module nothing in a ++snapshots them into an **int64** CUDA tensor. Without this module nothing in a + running engine ever calls that op, so the instrument records into a buffer no + one reads -- which measures nothing. + ++DTYPE IS PART OF THE PROTOCOL (large-path determinism change). Version 1 was u32 ++device counters read into a signed int32 tensor, which goes negative after ++~33k all-equal rows at k=512 and wraps after ~66k, so a delta taken across the ++boundary is a wrong number that looks like a measurement. Version 2 widens ++both sides, and the op REJECTS an int32 buffer rather than filling it from half ++of each counter. An older build whose op still demands int32 will therefore ++raise here; that is the intended signal, not a regression to paper over. ++ + SLOT LAYOUT is build-dependent, so nothing here hard-codes five. `slot_count()` + asks the build first (a `..._slots` op, should a later build expose one) and + otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5); names come from a +@@ -84,7 +92,7 @@ def _get_op(name: str): + + + def slot_count() -> int: +- """How many int32 slots this build's counter array carries.""" ++ """How many int64 slots this build's counter array carries.""" + query = _get_op(_OP_NAME + "_slots") + if query is not None: + return int(query()) +@@ -118,7 +126,9 @@ def read(reset: bool = False, sync: bool = True) -> dict[str, int]: + return {} + count = slot_count() + device = torch.device("cuda", torch.cuda.current_device()) +- out = torch.zeros(count, dtype=torch.int32, device=device) ++ # int64: the device counters are 64-bit and the op refuses any other ++ # width. See the module docstring -- this is protocol version 2. ++ out = torch.zeros(count, dtype=torch.int64, device=device) + op(out, 1 if reset else 0, 1 if sync else 0) + return dict(zip(slot_names(count), out.tolist(), strict=True)) + +-- +2.43.0 + diff --git a/patches/0044-topk-determinism-toggles-for-greedy-repeatability.patch b/patches/0044-topk-determinism-toggles-for-greedy-repeatability.patch new file mode 100644 index 0000000..3ff1dc6 --- /dev/null +++ b/patches/0044-topk-determinism-toggles-for-greedy-repeatability.patch @@ -0,0 +1,2043 @@ +From fd7899111c69153b283e55fc7b8915aa2cf68def Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:35:40 -0400 +Subject: [PATCH 044/103] topk: determinism toggles for greedy repeatability + +- MERLIN_FULL_K: deterministic Marlin MoE split-K (after unmerged upstream #46639). +- MERLIN_QSA_SORT: canonical order for the QSA selected blocks. +- MERLIN_TIE_DET: deterministic tie-break at the persistent_topk cutoff. +- MERLIN_TIE_RECENT: cutoff ties keep the highest indices. +- Decode candidate truncation is detected; one tie direction across the + medium path; the kernel switches are registered. +--- + .../moe/marlin_moe_wna16/kernel.h | 3 +- + .../moe/marlin_moe_wna16/marlin_template.h | 12 +- + .../moe/marlin_moe_wna16/ops.cu | 40 +- + csrc/libtorch_stable/ops.h | 4 + + csrc/libtorch_stable/persistent_topk.cuh | 392 +++++++++++++++--- + csrc/libtorch_stable/topk.cu | 183 +++++++- + csrc/libtorch_stable/torch_bindings.cpp | 3 + + tests/kernels/test_top_k_per_row.py | 234 ++++++++++- + tests/test_envs.py | 199 +++++++-- + vllm/envs.py | 151 +++++-- + .../qwen4_exp/nvidia/ops/qsa_indexer.py | 41 ++ + .../qwen4_exp/nvidia/ops/topk_counters.py | 43 +- + 12 files changed, 1141 insertions(+), 164 deletions(-) + +diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +index 25c57f3..15f8018 100644 +--- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h ++++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +@@ -20,7 +20,8 @@ + const int32_t *__restrict__ num_tokens_past_padded_ptr, \ + const float *__restrict__ topk_weights_ptr, int top_k, \ + bool mul_topk_weights, int prob_m, int prob_n, int prob_k, int *locks, \ +- bool has_bias, bool use_atomic_add, bool use_fp32_reduce ++ bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \ ++ bool use_full_k + + namespace MARLIN_NAMESPACE_NAME { + template gridDim.x) { + part2_mn_tiles = global_mn_tiles % gridDim.x; +- if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x; ++ if (!use_full_k && part2_mn_tiles * 3 <= gridDim.x) ++ part2_mn_tiles += gridDim.x; + part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; + } + +- int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); ++ int iters = ++ use_full_k ? k_tiles : div_ceil(k_tiles * part2_mn_tiles, gridDim.x); + + if constexpr (group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { +diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +index 48de07d..2c6ad76 100644 +--- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu ++++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +@@ -33,6 +33,35 @@ + #include + + #include "libtorch_stable/torch_utils.h" ++#include ++ ++namespace { ++// ------------------------------------------------------------------------- ++// MERLIN_FULL_K: the CUDA half of the unmerged upstream PR #46639 ++// (`use_full_k`), which makes the Marlin MoE split-K reduction order ++// independent of how many blocks the launch was split across. ++// ++// Our own toggle rather than upstream's VLLM_BATCH_INVARIANT, because that ++// one aborts engine startup on GDN models (upstream #42960) and so can never ++// reach this kernel on Flash-Next. ++// ++// Resolved ONCE per process, not per launch, so it is a launch constant in ++// exactly the sense topk.cu's exact_medium_enabled() is: a cudagraph bakes in ++// the value that was live at capture time and every replay agrees with it. ++// Default OFF -- enabling it by default is a separate release decision. ++// ++// This std::getenv is the SINGLE source of truth. vllm/envs.py carries a raw ++// passthrough entry for validate_environ and the compile-factor hash only; ++// nothing in Python may parse this name a second time. ++// ------------------------------------------------------------------------- ++bool merlin_full_k() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_FULL_K"); ++ return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ }(); ++ return enabled; ++} ++} // namespace + + #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ +@@ -183,7 +212,8 @@ exec_config_t determine_exec_config( + const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m, + int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks, + bool m_block_size_8, int num_bits, int group_size, bool has_zp, +- bool is_zp_float, bool is_a_8bit, int stages, int max_shared_mem, int sms) { ++ bool is_zp_float, bool is_a_8bit, int stages, int max_shared_mem, int sms, ++ bool use_full_k) { + exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; + thread_config_t* thread_configs = thread_m_blocks > 1 + ? large_batch_thread_configs +@@ -228,7 +258,8 @@ exec_config_t determine_exec_config( + else + allow_count = max(min(allow_count, 2), 1); + +- if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) { ++ if (!use_full_k && ++ prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) { + allow_count = + max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1); + } +@@ -254,6 +285,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int blocks_per_sm, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { ++ const bool use_full_k = merlin_full_k(); + int thread_m_blocks = div_ceil(moe_block_size, 16); + bool m_block_size_8 = moe_block_size == 8; + bool is_a_8bit = a_type.size_bits() == 8; +@@ -331,7 +363,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + exec_cfg = determine_exec_config( + a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts, + top_k, thread_m_blocks, m_block_size_8, num_bits, group_size, has_zp, +- is_zp_float, is_a_8bit, stages, max_shared_mem, sms); ++ is_zp_float, is_a_8bit, stages, max_shared_mem, sms, use_full_k); + thread_tfg = exec_cfg.tb_cfg; + } + +@@ -383,7 +415,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, + topk_weights_ptr, top_k, mul_topk_weights, prob_m, prob_n, prob_k, locks, +- has_bias, use_atomic_add, use_fp32_reduce); ++ has_bias, use_atomic_add, use_fp32_reduce, use_full_k); + // clang-format on + } + +diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h +index ea81216..1156801 100644 +--- a/csrc/libtorch_stable/ops.h ++++ b/csrc/libtorch_stable/ops.h +@@ -478,6 +478,10 @@ void persistent_topk(const torch::stable::Tensor& logits, + void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + int64_t sync); + ++// The slot count of THIS build's counter array, so the Python ++// reader never has to guess it from an environment variable. ++int64_t persistent_topk_exact_counters_slots(); ++ + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index fcbcc7b..f8ce859 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -54,6 +54,19 @@ __device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + } + ++// THE definition of the effective tie direction, called by ++// both the host (to report it) and the device (to obey it), so the reported ++// policy and the executed policy cannot drift. ++// ++// `MERLIN_TIE_DET=0` is the shipped default and leaves the buffered paths on ++// their original racy atomicAdd rule, which has no direction; the medium-path exact ++// restart is deterministic whatever happens, so it keeps the most-recent-wins ++// direction the medium-path restart shipped and was measured with. ++__host__ __device__ __forceinline__ bool effective_tie_recent(bool det_ties, ++ bool tie_recent) { ++ return det_ties ? tie_recent : true; ++} ++ + __device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); +@@ -147,10 +160,27 @@ struct PersistentTopKParams { + // nullptr; the kernel only ever atomicAdds into it, never reads back. + bool exact_medium; + unsigned long long* exact_counters; ++ // MERLIN_TIE_DET: deterministic tie-break at the k-th cutoff. ++ // Resolved once on the host, like exact_medium above, so it is a launch ++ // constant and is baked into a cudagraph at capture. ++ bool det_ties; ++ // MERLIN_TIE_RECENT: ranks ties by DESCENDING element index, so ++ // the surviving blocks are the most recent rather than the oldest. Both ++ // rules are deterministic; they differ in which tied blocks survive, and ++ // measurably in draft acceptance at depth. Only read when det_ties is set. ++ bool tie_recent; ++ // MERLIN_TOPK_EXACT_DECODE: on an overflowing ++ // DECODE row, take the same exact restart the medium path takes instead of ++ // ranking a truncated candidate set. Detection and the [10] counter are ++ // unconditional; only the repair is gated. Default OFF -- the repair costs ++ // four full-row rescans and the overflow incidence on real traffic has ++ // never been measured, which is what the counter is for. ++ bool exact_decode; + // MERLIN_TOPK_DET_LARGE, ANDed on the host with the capacity gate +- // 2 * ctas_per_group <= RADIX (the per-CTA counts are staged in one 256-slot +- // RadixRowState histogram row). Resolved once per process, so it is a +- // launch constant and a cudagraph captures one value. ++ // 4 * ctas_per_group <= RADIX (the per-CTA counts and the two fault-verdict ++ // regions are staged in one 256-slot RadixRowState histogram row). Resolved ++ // once per process, so it is a launch constant and a cudagraph captures one ++ // value. + bool det_large; + // TEST ONLY. -1 (the default) disables it; any + // other value names the `cta_in_group` whose CTA-LOCAL fault predicate is +@@ -177,10 +207,17 @@ __device__ __forceinline__ uint32_t decode_bin(float x) { + return key >> 5; + } + ++// Returns TRUE when this row overflowed its candidate buffer AND the caller ++// asked for the exact repair, in which case nothing has been emitted that the ++// restart will not overwrite and the caller must run it. ++// The restart itself lives with the medium path further down and is called ++// from the dispatch site, so this function needs no forward declaration of it ++// -- some host compilers reject one for a `__device__` template. + template +-__device__ __noinline__ void histogram_2048_topk( ++__device__ __noinline__ bool histogram_2048_topk( + const float* __restrict__ logits, int32_t* __restrict__ output_indices, +- int32_t seq_len) { ++ int32_t seq_len, bool det_ties, bool tie_recent, ++ unsigned long long* __restrict__ counters, bool exact_decode) { + extern __shared__ int decode_smem[]; + const int tx = threadIdx.x; + const int lane = tx & 31; +@@ -193,7 +230,18 @@ __device__ __noinline__ void histogram_2048_topk( + constexpr int MAX_ITEMS_PER_THREAD = + (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; + +- enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; ++ // sOVF claims one of the two scalars that were already ++ // being zeroed by the `tx < 8` loop below and never used, so the decode ++ // shared-memory budget is unchanged. ++ enum : int { ++ sTHR = 0, ++ sOUT = 1, ++ sREF = 2, ++ sFIN = 3, ++ sBUF0 = 4, ++ sBUF1 = 5, ++ sOVF = 6 ++ }; + + // ---- Initialize scalars (prevents stale data from prior rows) ---- + if (tx < 8) { +@@ -322,7 +370,7 @@ __device__ __noinline__ void histogram_2048_topk( + __syncthreads(); + + int remaining_k = TopK - decode_smem[SBASE + sOUT]; +- if (remaining_k <= 0) return; ++ if (remaining_k <= 0) return false; + + // If all buffered elements fit, output them all (common for short seqs) + const int raw_buf0 = decode_smem[SBASE + sBUF0]; +@@ -333,7 +381,7 @@ __device__ __noinline__ void histogram_2048_topk( + output_indices[base + i] = bufs[0][i]; + } + __syncthreads(); +- return; ++ return false; + } + + // ---- Phase 3: Deferred refinement (rare path) ---- +@@ -372,6 +420,42 @@ __device__ __noinline__ void histogram_2048_topk( + const int dst = src ^ 1; + + const int raw_buf = decode_smem[SBASE + sBUF0 + src]; ++ ++ // ---- DECODE TRUNCATION, DETECTED ------------------ ++ // ++ // `DBUF` is 3,708 and the Phase-2 append guard drops everything past it, ++ // but the atomicAdd that hands out slots sits OUTSIDE that guard -- so ++ // `raw_buf` keeps counting and a post-hoc test here is exact, exactly as ++ // the medium path's M2 is. Dropped candidates are invisible to ++ // refinement AND to the deterministic tie ranking further down: ranking ++ // orders the survivors, it cannot recall a winner that was never ++ // buffered. The failing class: an 8,192-element row inside one coarse bin ++ // with one element larger at FP32 but in the same half-precision bin; a ++ // legal append order drops that winner and the row is silently wrong. ++ // ++ // The historical zero-overflow measurements cover a workload, not a ++ // capacity. Until this counter has been read on the real traffic, the ++ // determinism claim for the decode path is conditional on it reading 0. ++ if (__builtin_expect(raw_buf > DBUF, 0)) { ++ if (tx == 0 && decode_smem[SBASE + sOVF] == 0) { ++ decode_smem[SBASE + sOVF] = 1; // once per row, not once per pass ++ if (counters != nullptr) atomicAdd(&counters[kCtrDecodeOverflow], 1ull); ++ } ++ if (exact_decode) { ++ // Hand the row back for the same repair the medium path takes: ++ // throw the buffered state away and re-select from scratch with a ++ // full-row radix rescan. The restart writes all TopK slots (or ++ // fills the row with -1 and faults), so Phase 2's `> threshold` ++ // emissions are overwritten rather than double-counted. `raw_buf` ++ // is block-uniform and `exact_decode` is a launch constant, so this ++ // return is taken by the whole block -- which is what makes it safe ++ // to leave a function that still owns the shared-memory window. ++ __syncthreads(); ++ return true; ++ } ++ __syncthreads(); ++ } ++ + const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; + + compute_suffix_sum(); +@@ -385,6 +469,9 @@ __device__ __noinline__ void histogram_2048_topk( + __syncthreads(); + + const int ref_thr = decode_smem[SBASE + sREF]; ++ // Read before any thread decrements sFIN, so every thread agrees on ++ // how many of the tied candidates survive the cutoff. ++ const int fin_total = decode_smem[SBASE + sFIN]; + remaining_k -= refine[0][ref_thr + 1]; + const int bit_offset = 24 - pass * 8; + +@@ -416,8 +503,35 @@ __device__ __noinline__ void histogram_2048_topk( + output_indices[pos] = idx; + } else if (bin == ref_thr) { + if (pass == 3) { +- const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); +- if (slot > 0) output_indices[TopK - slot] = idx; ++ if (det_ties) { ++ // Every candidate that reaches here has an identical 32-bit value, ++ // and the racy branch below hands out slots by atomicAdd -- first ++ // warp to arrive wins, so WHICH tied block survives the cutoff is ++ // a race. Rank by element index instead and place at a position ++ // derived from that rank, so the surviving set is arrival-order ++ // independent. No atomics on this path. ++ // ++ // O(t^2) over the exactly-tied set t. Exact 32-bit ties are rare ++ // (19 occurrences across 12 measured requests, one element each), so ++ // t is small in practice -- but that is measurement at one shape, ++ // not a bound. ++ int rank = 0; ++ for (int j = 0; j < num_buffered; ++j) { ++ const int jdx = bufs[src][j]; ++ const bool ahead = tie_recent ? (jdx > idx) : (jdx < idx); ++ if (!ahead) continue; ++ const uint32_t jf = convert_to_uint32_v2(logits[jdx]); ++ if (((jf >> bit_offset) & 0xFF) == static_cast(ref_thr)) { ++ ++rank; ++ } ++ } ++ if (rank < fin_total) { ++ output_indices[TopK - fin_total + rank] = idx; ++ } ++ } else { ++ const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); ++ if (slot > 0) output_indices[TopK - slot] = idx; ++ } + } else { + const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); + if (__builtin_expect(bp < DBUF, 1)) { +@@ -430,6 +544,7 @@ __device__ __noinline__ void histogram_2048_topk( + } + __syncthreads(); + } ++ return false; + } + + // ============================================================================ +@@ -457,10 +572,28 @@ __device__ __noinline__ void histogram_2048_topk( + // selected keys in output order, so that would re-introduce exactly the + // nondeterminism MERLIN_QSA_SORT / MERLIN_TIE_DET / MERLIN_TIE_RECENT exist + // to remove. Ours emits in **ascending index order** via a block scan (no +-// atomics anywhere on the placement path) and cuts ties at the k-th value by +-// **most-recent wins** (the highest element indices are kept), matching the +-// rule measured on the reference host (+3.4 % tokens/step at 12K over +-// oldest-wins). ++// atomics anywhere on the placement path). ++// ++// TIE DIRECTION IS ONE POLICY, DECIDED CENTRALLY. The medium-path restart ++// hard-coded most-recent-wins here while the buffered decode/medium paths took ++// their direction from `tie_recent`, so a medium row that grew from 4,096 to ++// 4,097 candidates silently changed tie direction without the configuration ++// changing. The direction is now a parameter and the ONE place that resolves ++// it is `effective_tie_recent()` in topk.cu: ++// ++// det_ties ? tie_recent : true ++// ++// i.e. `MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=0` is oldest-wins on the whole ++// medium path (buffered AND restart), and with `MERLIN_TIE_DET=0` the restart ++// keeps its shipped most-recent-wins default (the direction that was ++// measured at +3.4 % tokens/step at 12K over oldest-wins). ++// ++// The deterministic large path takes the SAME value. ++// `radix_emit_deterministic` receives `effective_tie_recent(det_ties, ++// tie_recent)` from the dispatch site and mirrors both halves of its cut -- ++// the cross-CTA walk and the within-chunk window -- so all five paths ++// (decode-buffered, decode-restart, medium-buffered, medium-restart, large) ++// obey one decision and the host's policy line prints it for all five. + // ============================================================================ + + // Abort the kernel when the exact fallback cannot complete, instead of +@@ -483,6 +616,23 @@ __device__ __noinline__ void histogram_2048_topk( + // rows and wraps after 66,053, and a delta taken across that boundary is a + // wrong number that looks like a measurement. `atomicAdd(unsigned long long*)` + // is native on sm_86, so the width costs nothing on the increment path. ++// ++// SLOT MAP (the tie-determinism + large-path integration, 11 slots). The tie change numbered its decode ++// counter AROUND the large-path block rather than after the medium-path one, so the merged layout is ++// the union of the two branches with nothing renumbered and no external probe ++// silently re-reading a row count as a switch -- the failure a merged layout ++// must avoid. ++// ++// [0] medium rows [1] medium overflow [2] medium exact fallback ++// [3] medium exact fault ++// [4]-[7] large path (rows_large, ties_cut, ties_dropped, faults_large) ++// [8] MERLIN_TOPK_EXACT_MEDIUM [9] MERLIN_TOPK_DET_LARGE ++// [10] decode overflow ++// ++// Slots 8-9 are written by the HOST after the device copy; the device never ++// touches them, and they are slots INSIDE the array, not two elements appended ++// after it. A stale five- or ten-element reader buffer is rejected outright by ++// the length check in the reader op, which is the fail-loud direction. + enum TopKExactCounter : int { + kCtrMediumRows = 0, // denominator: medium-path rows entered + kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop +@@ -497,7 +647,11 @@ enum TopKExactCounter : int { + kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut + kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts + kCtrLargeFault = 7, // large-path invariant violation +- kNumTopKExactCounters = 8, ++ kSlotSwitchMedium = 8, // host-written flag, NOT a device counter ++ kSlotSwitchLarge = 9, // host-written flag, NOT a device counter ++ kCtrDecodeOverflow = 10, // decode rows whose buffered-candidate count ++ // exceeded DBUF before deterministic ranking ++ kNumTopKExactCounters = 11, // total slots: device buffer AND reader output + }; + + // Exact, bounded-memory re-selection for a single medium-path row. +@@ -506,10 +660,17 @@ enum TopKExactCounter : int { + // row, so correctness does not depend on how many values share a coarse bin. + // Entered only after an overflow has been detected, which is 0.1-0.2 % of + // medium rows at 38K-51K prompt tokens. ++// ++// `tie_recent` is the EFFECTIVE tie direction for this launch, resolved once ++// on the host (see the header comment above): true keeps the highest element ++// indices among the pivot-equal group, false keeps the lowest. It is a launch ++// constant and block-uniform, so every walk and the emission predicate below ++// agree on it without a barrier. + template + __device__ __noinline__ void medium_exact_restart( + const float* __restrict__ scores, int32_t* __restrict__ output, +- uint32_t length, unsigned long long* __restrict__ counters) { ++ uint32_t length, unsigned long long* __restrict__ counters, ++ bool tie_recent) { + static_assert(kThreadsPerBlock >= RADIX, + "the single-thread bin walk assumes RADIX <= blockDim.x"); + +@@ -521,7 +682,8 @@ __device__ __noinline__ void medium_exact_restart( + uint32_t prefix; // ordered-uint radix prefix, grows 8 bits per round + uint32_t remaining; // how many of the pivot-equal group must be kept + uint32_t tie_total; // how many elements have key == pivot in total +- uint32_t tie_cut; // lowest element index kept among the pivot-equal ++ uint32_t tie_cut; // pivot-equal cut: lowest kept index under ++ // most-recent-wins, highest kept under oldest-wins + uint32_t fault; // set by the walks; block-wide, read after a barrier + }; + +@@ -541,7 +703,11 @@ __device__ __noinline__ void medium_exact_restart( + smem->prefix = 0; + smem->remaining = static_cast(TopK); + smem->tie_total = 0; +- smem->tie_cut = 0; ++ // The cut is a bound on the kept index range, so its identity element is ++ // direction-dependent. When tie_total == equal_count the walks below are ++ // skipped entirely and this value is what the emission predicate reads: ++ // `idx >= 0` keeps the whole group, and so does `idx <= 0xFFFFFFFF`. ++ smem->tie_cut = tie_recent ? 0u : 0xFFFFFFFFu; + // The index-space tie walk below is two 8-bit rounds, so it resolves + // indices below 65,536. The medium path serves seq_len <= 32,768, but + // fail loudly rather than silently mis-cut if that ever changes. +@@ -626,11 +792,18 @@ __device__ __noinline__ void medium_exact_restart( + // the end, but wasted work on the fault path and harder to debug. The + // condition is block-uniform (shared, read after the barrier above). + if (smem->fault == 0u) { +- // ---- tie cut at the k-th value: most-recent wins -------------------- ++ // ---- tie cut at the k-th value -------------------------------------- + // + // Keep the `equal_count` HIGHEST element indices among the pivot-equal +- // group. Finding the cut index turns the tie rule into a per-element +- // predicate, which is what lets the emission below be atomic-free. ++ // group when `tie_recent`, the `equal_count` LOWEST when not. Finding the ++ // cut index turns the tie rule into a per-element predicate, which is what ++ // lets the emission below be atomic-free. ++ // ++ // The two index walks below are the only code that has ++ // to know the direction; they differ from each other only in which end of ++ // the index-bin array they start from, so `count_above` means "population ++ // on the far side of the walk" in both directions and the crossing test is ++ // unchanged. + // + // Exact 32-bit ties at the cutoff are rare (19 occurrences across + // 12 requests, one element each), so the common case is tie_total == +@@ -651,9 +824,12 @@ __device__ __noinline__ void medium_exact_restart( + } + __syncthreads(); + if (tx == 0) { ++ const int first_bin = tie_recent ? RADIX - 1 : 0; ++ const int bin_step = tie_recent ? -1 : 1; + uint32_t count_above = 0; + bool found = false; +- for (int bin = RADIX - 1; bin >= 0; --bin) { ++ for (int n = 0; n < RADIX; ++n) { ++ const int bin = first_bin + n * bin_step; + const uint32_t count = smem->histogram[bin]; + if (count_above + count >= need) { + smem->tie_cut = static_cast(bin); // high byte, for now +@@ -684,9 +860,12 @@ __device__ __noinline__ void medium_exact_restart( + } + __syncthreads(); + if (tx == 0) { ++ const int first_bin = tie_recent ? RADIX - 1 : 0; ++ const int bin_step = tie_recent ? -1 : 1; + uint32_t count_above = 0; + bool found = false; +- for (int bin = RADIX - 1; bin >= 0; --bin) { ++ for (int n = 0; n < RADIX; ++n) { ++ const int bin = first_bin + n * bin_step; + const uint32_t count = smem->histogram[bin]; + if (count_above + count >= need) { + smem->tie_cut = (high_byte << 8) | static_cast(bin); +@@ -711,7 +890,9 @@ __device__ __noinline__ void medium_exact_restart( + + // ---- emission: ascending index order, no atomics --------------------- + // +- // selected(i) = ordered(i) > pivot || (ordered(i) == pivot && i >= tie_cut) ++ // selected(i) = ordered(i) > pivot || ++ // (ordered(i) == pivot && (tie_recent ? i >= tie_cut ++ // : i <= tie_cut)) + // + // is a pure function of the element, so the output slot is a prefix sum + // over the row rather than an arrival rank. Tiles run in ascending base +@@ -726,8 +907,8 @@ __device__ __noinline__ void medium_exact_restart( + uint32_t flag = 0; + if (idx < length) { + const uint32_t ordered = convert_to_uint32_v2(scores[idx]); +- flag = +- (ordered > pivot || (ordered == pivot && idx >= tie_cut)) ? 1u : 0u; ++ const bool keep_tie = tie_recent ? (idx >= tie_cut) : (idx <= tie_cut); ++ flag = (ordered > pivot || (ordered == pivot && keep_tie)) ? 1u : 0u; + } + uint32_t rank = 0; + uint32_t total = 0; +@@ -805,7 +986,8 @@ template + __device__ __noinline__ void histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, + int logits_offset, int seq_len, bool exact_medium, +- unsigned long long* __restrict__ counters) { ++ unsigned long long* __restrict__ counters, bool det_ties, ++ bool tie_recent) { + // No default arguments on purpose: a new call site must state whether it + // wants the exact path, rather than silently inheriting the broken one. + // All shared state lives in dynamic shared memory to avoid static +@@ -973,8 +1155,14 @@ __device__ __noinline__ void histogram_256_topk( + atomicAdd(&counters[kCtrMediumExactFallback], 1ull); + } + __syncthreads(); // no thread may clobber medium_smem before this ++ // The restart takes the SAME effective direction as ++ // the buffered branch below, so crossing the 4,096-candidate capacity ++ // no longer flips the tie policy. `det_ties == false` means the ++ // buffered branch is the racy atomicAdd rule with no direction at all, ++ // and the restart then keeps the medium path's shipped most-recent-wins. + medium_exact_restart(logits + logits_offset, output_indices, +- static_cast(seq_len), counters); ++ static_cast(seq_len), counters, ++ effective_tie_recent(det_ties, tie_recent)); + return; + } + __syncthreads(); +@@ -1017,6 +1205,9 @@ __device__ __noinline__ void histogram_256_topk( + } + __syncthreads(); + ++ // Read before any thread decrements shared_final_k; see the decode ++ // path's fin_total for the reason. ++ const int fin_total_m = shared_final_k; + for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { + const int idx = buffered_indices[src_buffer][i]; + const float logit_value = logits[idx + logits_offset]; +@@ -1027,9 +1218,35 @@ __device__ __noinline__ void histogram_256_topk( + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + if (pass == 3) { +- const int slot = atomicAdd(&shared_final_k, -1); +- if (slot > 0) { +- output_indices[TopK - slot] = idx; ++ if (det_ties) { ++ // Same rule as the decode path: every candidate here is bit-equal, ++ // so rank by element index rather than by warp arrival order. ++ // ++ // NOTE the medium-path restart interaction: when exact_medium is on AND this row ++ // overflowed the candidate buffer, medium_exact_restart() has ++ // already returned and this branch is never reached. Since the fix, ++ // that path is handed THIS same `tie_recent`, so the ++ // medium path has one tie direction either side of the capacity ++ // boundary. ++ int rank = 0; ++ for (int j = 0; j < num_buffered; ++j) { ++ const int jdx = buffered_indices[src_buffer][j]; ++ const bool ahead = tie_recent ? (jdx > idx) : (jdx < idx); ++ if (!ahead) continue; ++ const uint32_t jf = ++ convert_to_uint32_v2(logits[jdx + logits_offset]); ++ if (static_cast((jf >> bit_offset) & 0xFF) == threshold_bin) { ++ ++rank; ++ } ++ } ++ if (rank < fin_total_m) { ++ output_indices[TopK - fin_total_m + rank] = idx; ++ } ++ } else { ++ const int slot = atomicAdd(&shared_final_k, -1); ++ if (slot > 0) { ++ output_indices[TopK - slot] = idx; ++ } + } + } else { + const int buffer_pos = +@@ -1166,12 +1383,16 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, + // * every CTA reads all 2*ctas_per_group counts and derives its own base by + // an exclusive prefix over `cta_in_group`, in one fixed order, so all + // CTAs agree without further communication; +-// * the tie cut keeps the HIGHEST global indices (most-recent-wins, the +-// `MERLIN_TIE_RECENT` policy, measured at +3.4 % tokens/step +-// over oldest-wins). Chunks are handed out in ascending index order ++// * the tie cut keeps the HIGHEST global indices under most-recent-wins -- ++// the `MERLIN_TIE_RECENT` policy, measured at +3.4 % ++// tokens/step over oldest-wins, and the shipped default -- or the LOWEST ++// under oldest-wins. Chunks are handed out in ascending index order + // (`my_chunk_start = cta_in_group * chunk_size`), so "the highest + // `equal_needed` pivot-equal indices" is exactly "walk CTAs from the top +-// down", which the per-CTA equal counts already answer; ++// down" and the oldest-wins rule is the same walk mirrored; the per-CTA ++// equal counts already answer both. The direction is the launch constant ++// `effective_tie_recent(det_ties, tie_recent)`, shared with the medium ++// path; + // * emission is a prefix sum over ascending element index — no atomic + // anywhere on the placement path — so the row is strictly ascending in + // index and bit-reproducible. +@@ -1198,7 +1419,7 @@ __device__ __noinline__ void radix_emit_deterministic( + uint32_t* local_histogram, uint32_t* suffix_sum, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, uint32_t ctas_per_group, + int& barrier_phase, uint32_t radix_iter, uint32_t tx, +- unsigned long long* counters, int fault_inject_cta) { ++ unsigned long long* counters, int fault_inject_cta, bool tie_recent) { + constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u; + // Staging for the published per-CTA counts. The host gates the switch on + // 4 * ctas_per_group <= RADIX, so these regions never collide. +@@ -1216,6 +1437,8 @@ __device__ __noinline__ void radix_emit_deterministic( + // [0] gt count [1] eq count [2] gt_base [3] eq_keep_base [4] cut_rank + // [5] group fault verdict [6] this CTA's emission-bound violation + // [7] group emission-bound verdict [8] ties dropped on this row ++ // [9] my_keep (the pivot-equal window width; with [4] it spells the cut ++ // in BOTH tie directions -- see the tie_recent note above) + // ---- per-CTA counts, out of the resident chunk ------------------------- + if (tx < 2u) suffix_sum[tx] = 0u; + if (tx == 6u) suffix_sum[6] = 0u; +@@ -1301,18 +1524,38 @@ __device__ __noinline__ void radix_emit_deterministic( + gt_total += g; + } + +- // Most-recent-wins: walk CTAs from the top down, handing each one as many +- // of the `equal_needed` slots as the CTAs above it left unclaimed. keep[] +- // is written over eq[] in place (eq[c] is read before keep[c] is stored). +- uint32_t eq_above = 0u; +- for (int c = static_cast(ctas_per_group) - 1; c >= 0; --c) { +- const uint32_t e = local_histogram[kEqStage + c]; +- const uint32_t want = +- (eq_above >= equal_needed) ? 0u : (equal_needed - eq_above); +- local_histogram[kEqStage + c] = (want < e) ? want : e; +- eq_above += e; ++ // Tie-determinism + large-path integration: the walk direction IS the tie ++ // policy. Most-recent-wins walks CTAs from the top down, handing each one ++ // as many of the `equal_needed` slots as the CTAs above it left unclaimed; ++ // oldest-wins is the same walk mirrored, from the bottom up. Chunks are ++ // handed out in ascending index order (`my_chunk_start = cta_in_group * ++ // chunk_size`), so "the highest `equal_needed` pivot-equal indices" is ++ // exactly "the top CTAs first" and the lowest is exactly the reverse. ++ // `tie_recent` is the EFFECTIVE direction, resolved once on the host by ++ // `effective_tie_recent(det_ties, tie_recent)` -- the same value the ++ // medium restart takes and the same value the host's policy line prints, ++ // so the three paths cannot drift. It is a launch constant and this walk ++ // is single-threaded, so the branch costs nothing that matters. keep[] is ++ // written over eq[] in place (eq[c] is read before keep[c] is stored). ++ uint32_t eq_seen = 0u; ++ if (tie_recent) { ++ for (int c = static_cast(ctas_per_group) - 1; c >= 0; --c) { ++ const uint32_t e = local_histogram[kEqStage + c]; ++ const uint32_t want = ++ (eq_seen >= equal_needed) ? 0u : (equal_needed - eq_seen); ++ local_histogram[kEqStage + c] = (want < e) ? want : e; ++ eq_seen += e; ++ } ++ } else { ++ for (uint32_t c = 0; c < ctas_per_group; ++c) { ++ const uint32_t e = local_histogram[kEqStage + c]; ++ const uint32_t want = ++ (eq_seen >= equal_needed) ? 0u : (equal_needed - eq_seen); ++ local_histogram[kEqStage + c] = (want < e) ? want : e; ++ eq_seen += e; ++ } + } +- const uint32_t eq_total = eq_above; ++ const uint32_t eq_total = eq_seen; + + uint32_t eq_keep_base = 0u; + for (uint32_t c = 0; c < cta_in_group; ++c) { +@@ -1354,9 +1597,15 @@ __device__ __noinline__ void radix_emit_deterministic( + + suffix_sum[2] = gt_base; + suffix_sum[3] = eq_keep_base; +- // Keep the local pivot-equal elements whose ascending rank within this +- // chunk is >= cut_rank, i.e. the `my_keep` highest-indexed of them. +- suffix_sum[4] = (my_keep <= local_eq) ? (local_eq - my_keep) : 0u; ++ // The kept pivot-equal elements of this chunk are the `my_keep` whose ++ // ascending rank within the chunk lies in [cut_rank, cut_rank + my_keep). ++ // Most-recent-wins puts that window at the TOP of the chunk ++ // (cut_rank = local_eq - my_keep, so the upper bound is vacuous and the ++ // emission predicate is byte-for-byte the rule the large path shipped); oldest-wins ++ // puts it at the BOTTOM (cut_rank = 0, and the upper bound is what binds). ++ suffix_sum[4] = ++ tie_recent ? ((my_keep <= local_eq) ? (local_eq - my_keep) : 0u) : 0u; ++ suffix_sum[9] = my_keep; + suffix_sum[8] = (eq_total > equal_needed) ? (eq_total - equal_needed) : 0u; + + // Publish this CTA's verdict. Own word, so no pre-zeroing is needed. +@@ -1396,6 +1645,7 @@ __device__ __noinline__ void radix_emit_deterministic( + const uint32_t gt_base = suffix_sum[2]; + const uint32_t eq_keep_base = suffix_sum[3]; + const uint32_t cut_rank = suffix_sum[4]; ++ const uint32_t keep_width = suffix_sum[9]; + const uint32_t row_fault = suffix_sum[5]; + + // ---- emission: rank by index, no atomics ------------------------------- +@@ -1459,11 +1709,20 @@ __device__ __noinline__ void radix_emit_deterministic( + const uint32_t tile_gt = warp_gt[32]; + const uint32_t tile_eq = warp_eq[32]; + +- if (f_gt || (f_eq && eq_rank >= cut_rank)) { +- // Selected elements before this one, in ascending index order: +- // all `> pivot` ones, plus the pivot-equal ones past the cut. +- const uint32_t kept_eq_before = ++ if (f_gt || (f_eq && eq_rank >= cut_rank && ++ eq_rank - cut_rank < keep_width)) { ++ // Selected elements before this one, in ascending index order: all ++ // `> pivot` ones, plus the pivot-equal ones already inside the cut ++ // window. The clamp to `keep_width` is what makes the expression work ++ // in BOTH directions: under most-recent-wins the window ends at the ++ // last pivot-equal element of the chunk, so `eq_rank - cut_rank` can ++ // never reach `keep_width` and the clamp is dead; under oldest-wins ++ // the window ends early and every pivot-equal element after it must ++ // stop contributing, or a later `> pivot` element is placed past the ++ // end of the row. ++ uint32_t kept_eq_before = + (eq_rank > cut_rank) ? (eq_rank - cut_rank) : 0u; ++ if (kept_eq_before > keep_width) kept_eq_before = keep_width; + const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before; + if (slot < static_cast(TopK)) { + row_output[slot] = static_cast(my_chunk_start + i); +@@ -1554,7 +1813,8 @@ __device__ void radix_topk(const float* __restrict__ row_input, + RadixRowState* state, uint32_t cta_in_group, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t radix_iter, uint32_t tx, bool det_large, +- unsigned long long* counters, int fault_inject_cta) { ++ unsigned long long* counters, int fault_inject_cta, ++ bool tie_recent) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; +@@ -1715,7 +1975,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, + row_output, ordered_pivot, shared_scalars[1], my_chunk_start, + actual_chunk_size, local_histogram, suffix_sum, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters, +- fault_inject_cta); ++ fault_inject_cta, tie_recent); + return; + } + +@@ -1854,10 +2114,23 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + row_output[i] = (i < seq_len) ? static_cast(i) : -1; + } + } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { +- histogram_2048_topk(row_input, row_output, seq_len); ++ // An overflowing decode row comes back here for ++ // the same exact restart the medium path takes. The return is ++ // block-uniform and nothing has been emitted that the restart does ++ // not overwrite. With MERLIN_TOPK_EXACT_DECODE off this is always ++ // false and the decode path is byte-identical to the medium-path restart's. ++ const bool decode_needs_exact = histogram_2048_topk( ++ row_input, row_output, seq_len, params.det_ties, ++ params.tie_recent, params.exact_counters, params.exact_decode); ++ if (decode_needs_exact) { ++ medium_exact_restart( ++ row_input, row_output, seq_len, params.exact_counters, ++ effective_tie_recent(params.det_ties, params.tie_recent)); ++ } + } else { + histogram_256_topk(row_input, row_output, 0, seq_len, +- params.exact_medium, params.exact_counters); ++ params.exact_medium, params.exact_counters, ++ params.det_ties, params.tie_recent); + } + } + continue; +@@ -1868,7 +2141,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + row_input, row_output, seq_len, my_chunk_start, chunk_size, + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, +- params.det_large, params.exact_counters, params.fault_inject_cta); ++ params.det_large, params.exact_counters, params.fault_inject_cta, ++ effective_tie_recent(params.det_ties, params.tie_recent)); + radix_iter++; + } + } +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 07d166a..2d2c81d 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -29,6 +29,41 @@ namespace { + // behaviour exactly (the counters still record, so the control arm reports the + // overflow incidence it is the control for). + // --------------------------------------------------------------------------- ++// ------------------------------------------------------------------------- ++// MERLIN_TIE_DET: deterministic tie-break at the k-th cutoff in ++// persistent_topk. Resolved once per process, exactly like ++// exact_medium_enabled() below, so it is a launch constant that a cudagraph ++// bakes in at capture. ++// ++// Our own toggle rather than upstream's VLLM_BATCH_INVARIANT, which aborts ++// engine startup on GDN models (upstream #42960) and so can never reach this ++// kernel on Flash-Next. Default OFF -- promotion is a release decision. ++// ++// This std::getenv is the SINGLE source of truth; vllm/envs.py carries a raw ++// passthrough for validate_environ and the compile-factor hash only. ++// ------------------------------------------------------------------------- ++// MERLIN_TIE_RECENT -- the DIRECTION of the deterministic tie-break above. ++// Measured at 12k tokens on the reference host: oldest-wins costs -7.9 % decode, ++// most-recent-wins is +1.5 %. Both deterministic. Also the direction ++// medium_exact_restart() already implements, so setting this alongside ++// MERLIN_TIE_DET makes the medium path agree with itself on overflowing ++// rows. Only read when MERLIN_TIE_DET is set. Default OFF. ++bool tie_recent_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TIE_RECENT"); ++ return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ }(); ++ return enabled; ++} ++ ++bool det_ties_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TIE_DET"); ++ return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ }(); ++ return enabled; ++} ++ + bool exact_medium_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TOPK_EXACT_MEDIUM"); +@@ -55,6 +90,23 @@ bool det_large_enabled() { + return enabled; + } + ++// MERLIN_TOPK_EXACT_DECODE: repair an overflowing decode ++// row with the same exact restart the medium path uses. Resolved once per ++// process like the others, so it is a launch constant a cudagraph bakes in. ++// ++// Default OFF, unlike MERLIN_TOPK_EXACT_MEDIUM. The medium path's overflow ++// incidence was measured (0.1-0.2 % of medium rows at 38-51K) before that ++// switch defaulted on; the decode path's has never been measured at all, ++// which is precisely what the [10] counter now exists to establish. Detection ++// is unconditional and free; the repair waits on a number. ++bool exact_decode_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_EXACT_DECODE"); ++ return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ }(); ++ return enabled; ++} ++ + // --------------------------------------------------------------------------- + // MERLIN_TOPK_LARGE_FAULT_INJECT. TEST ONLY. + // +@@ -79,6 +131,50 @@ int large_fault_inject_cta() { + return cta; + } + ++// --------------------------------------------------------------------------- ++// Say the effective tie policy out loud, once. ++// ++// "TIE_DET=1 TIE_RECENT=0" used to mean three different things on the three ++// paths, and nothing printed the difference. One stderr line at the first ++// launch of the process naming the policy the kernel will actually apply on ++// EVERY path. It is a plain fprintf on the host: no CUDA call, so it is safe ++// from a capturing thread, and the `static` initialiser makes it exactly once ++// per process regardless of how many workers share the log. ++// ++// The line covers all FIVE paths. The large ++// path is no longer a fixed rule: `radix_emit_deterministic` takes the same ++// `effective_tie_recent()` value as the medium restart, so `large=` reports ++// the real direction when MERLIN_TOPK_DET_LARGE is on and the inherited ++// v2.0.1 atomic-arrival race when it is off. The flag printed here is the ++// ENV value; a launch whose ctas_per_group overflows the staging row degrades ++// to atomic-arrival for that launch and leaves counter [4] at 0, which is the ++// separate warning below. ++// --------------------------------------------------------------------------- ++void report_tie_policy_once() { ++ static const bool reported = [] { ++ const bool det = det_ties_enabled(); ++ const bool recent = ++ vllm::persistent::effective_tie_recent(det, tie_recent_enabled()); ++ const char* dir = recent ? "most-recent-wins" : "oldest-wins"; ++ std::fprintf( ++ stderr, ++ "[merlin] persistent_topk tie policy: MERLIN_TIE_DET=%d " ++ "MERLIN_TIE_RECENT=%d MERLIN_TOPK_EXACT_MEDIUM=%d " ++ "MERLIN_TOPK_EXACT_DECODE=%d MERLIN_TOPK_DET_LARGE=%d -> " ++ "decode-buffered=%s decode-restart=%s medium-buffered=%s " ++ "medium-restart=%s large=%s\n", ++ det ? 1 : 0, tie_recent_enabled() ? 1 : 0, ++ exact_medium_enabled() ? 1 : 0, exact_decode_enabled() ? 1 : 0, ++ det_large_enabled() ? 1 : 0, det ? dir : "atomic-arrival", ++ exact_decode_enabled() ? dir : "n/a (detect-only; counter [10])", ++ det ? dir : "atomic-arrival", ++ exact_medium_enabled() ? dir : "n/a (exact restart disabled)", ++ det_large_enabled() ? dir : "atomic-arrival (non-deterministic)"); ++ return true; ++ }(); ++ (void)reported; ++} ++ + // Process-lifetime device counters, one block per device. The kernel only + // ever atomicAdds into them, so the launch path needs no host sync and the + // values accumulate across cudagraph replays. +@@ -421,6 +517,10 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + params.fault_inject_cta); + } + } ++ params.det_ties = det_ties_enabled(); ++ params.tie_recent = tie_recent_enabled(); ++ params.exact_decode = exact_decode_enabled(); ++ report_tie_policy_once(); + // Allocate only off the capture path; in the ordinary engine lifecycle the + // eager warmup lands here first, long before capture. The explicit, + // quiescent-phase entry point is the reader op. +@@ -526,9 +626,10 @@ void persistent_topk(const torch::stable::Tensor& logits, + // device before any cudagraph capture. + // + // PROTOCOL VERSION 2 -- BREAKING. `out` is an **int64** +-// CUDA tensor with at least 10 elements; int32 is REJECTED, not silently +-// accepted. The device counters are 64-bit unsigned and are copied D2D +-// without conversion, so the reader's element width must match the device's. ++// CUDA tensor with EXACTLY kNumTopKExactCounters (11) elements; int32, and any ++// other length, are REJECTED rather than silently accepted. The device counters are 64-bit unsigned ++// and are copied D2D without conversion, so the reader's element width must ++// match the device's. + // + // Why the break rather than a reinterpreting reader: the counters were u32 + // copied into a SIGNED int32 tensor. At 65,536 columns with k=512 an +@@ -543,11 +644,20 @@ void persistent_topk(const torch::stable::Tensor& logits, + // so the instrument costs the same, and a 64-bit slot cannot be driven to the + // signed boundary by any run this kernel will ever see. + // ++// The layout is the UNION of the tie-determinism and large-path slot ++// maps -- the large path's ten slots plus the decode restart's [10] decode_overflow, eleven in all, ++// with nothing renumbered. The two switch flags are slots INSIDE the array at ++// [8] and [9], not two elements appended after it, so a probe sized to the ++// old ten- or five-element layout is rejected by the length check below rather ++// than being filled with the wrong meanings. Ask the build for its own answer ++// with `torch.ops._C.persistent_topk_exact_counters_slots()`. ++// + // Everything that reads these counters must move together: the Python reader +-// `vllm/models/qwen4_exp/nvidia/ops/topk_counters.py` now allocates int64, and +-// the medium-path tests allocate int64. An int32 caller gets a loud error rather +-// than half a number. The OPERATOR SCHEMA is unchanged (dtype is not part of +-// it), so a stale caller fails at the check, not at the binding. ++// `vllm/models/qwen4_exp/nvidia/ops/topk_counters.py` allocates int64 and ++// eleven slots, and the medium- and large-path tests do too. An int32 caller gets a loud ++// error rather than half a number. The OPERATOR SCHEMA is unchanged (dtype ++// and length are not part of it), so a stale caller fails at the check, not at ++// the binding. + // + // Slots: + // [0] medium-path rows entered (denominator) +@@ -558,8 +668,11 @@ void persistent_topk(const torch::stable::Tensor& logits, + // [5] large-path rows whose pivot-equal group was cut at the k-th value + // [6] large-path pivot-equal elements discarded by those cuts + // [7] large-path invariant violations (must stay 0) +-// [8] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0 +-// [9] 1 if MERLIN_TOPK_DET_LARGE is on, else 0 ++// [8] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0 (host-written) ++// [9] 1 if MERLIN_TOPK_DET_LARGE is on, else 0 (host-written) ++// [10] decode rows whose buffered candidates exceeded DBUF (3,708) ++// before deterministic ranking -- MUST be read before any decode-path ++// determinism claim + // + // COHERENCY -- stated precisely, because the earlier wording overclaimed. + // With `sync == 0` the snapshot is ordered ONLY against work already enqueued +@@ -587,7 +700,8 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + out.scalar_type() == torch::headeronly::ScalarType::Long, + "out must be int64: the exact-restart counters are 64-bit as of the " + "large-path counter protocol change, and an int32 buffer would be filled " +- "from half of each counter. Allocate torch.zeros(10, dtype=torch.int64)." ++ "from half of each counter. Allocate " ++ "torch.zeros(persistent_topk_exact_counters_slots(), dtype=torch.int64)." + " (int32 was protocol version 1 and is deliberately rejected.)"); + // Element count alone is not enough. A CUDA scalar expanded + // to ten elements (stride 0) satisfies numel but owns one element's storage, +@@ -598,9 +712,19 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + STD_TORCH_CHECK(out.is_contiguous(), + "out must be contiguous (expanded or strided views would be " + "written out of bounds or into the wrong elements)"); +- STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 2, +- "out must hold at least ", +- vllm::persistent::kNumTopKExactCounters + 2, " elements"); ++ // EXACTLY the build's slot count, not "at least". `>=` accepted an ++ // 11-element buffer on the ten-slot large-path build and quietly left the extra ++ // element reading 0 -- a caller compiled against one layout got a plausible ++ // answer from a build with another, which is the failure mode the ++ // self-describing `..._slots` op exists to prevent. Ask the build, size the ++ // buffer to what it says, and a mismatch is an error rather than a wrong ++ // number. ++ STD_TORCH_CHECK(out.numel() == vllm::persistent::kNumTopKExactCounters, ++ "out must hold exactly ", ++ vllm::persistent::kNumTopKExactCounters, ++ " elements (this build's persistent_topk_exact_counters_" ++ "slots(); the switch flags are slots inside the array, not " ++ "appended after it), got ", out.numel()); + STD_TORCH_CHECK( + reset == 0 || sync != 0, + "persistent_topk_exact_counters: reset requires sync != 0, otherwise the " +@@ -645,15 +769,20 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + } + + // Static storage: the source of an async H2D copy must outlive the copy. +- // Two flags: [kNumTopKExactCounters] = the medium-path switch, +- // [kNumTopKExactCounters + 1] = the large-path switch. The large +- // flag reports the ENV value, not the per-launch effective value: a launch +- // whose ctas_per_group overflowed the staging row still reports 1 here and +- // leaves counter [4] at 0, which is how that degradation is diagnosed. ++ // Written AFTER the device copy, which has already zeroed both slots: they ++ // are slots [8] and [9] INSIDE the array and the device never touches them. ++ // The large flag reports the ENV value, not the per-launch effective value: ++ // a launch whose ctas_per_group overflowed the staging row still reports 1 ++ // here and leaves counter [4] at 0, which is how that degradation is ++ // diagnosed. kSlotSwitchLarge == kSlotSwitchMedium + 1, so one copy does ++ // both. ++ static_assert(vllm::persistent::kSlotSwitchLarge == ++ vllm::persistent::kSlotSwitchMedium + 1, ++ "the two host-written flags must stay adjacent for this copy"); + static const int64_t enabled[2] = {exact_medium_enabled() ? 1 : 0, + det_large_enabled() ? 1 : 0}; + STD_TORCH_CHECK( +- cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, enabled, ++ cudaMemcpyAsync(dst + vllm::persistent::kSlotSwitchMedium, enabled, + 2 * sizeof(int64_t), cudaMemcpyHostToDevice, + stream) == cudaSuccess, + "exact counter flag write failed"); +@@ -670,3 +799,19 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); + #endif + } ++ ++// How many int32 slots THIS build's counter array carries. ++// ++// The Python reader used to guess from MERLIN_TOPK_COUNTER_SLOTS (default 5) ++// and name the slots by that length. The decode restart makes the array 11 and the large path makes ++// it 10, so a guess is now wrong on at least one branch, and a reader that ++// allocates the wrong length is rejected by the length check above -- the ++// instrument goes dark exactly when it starts mattering. Ask the build ++// instead. No CUDA call, no allocation, safe from anywhere. ++int64_t persistent_topk_exact_counters_slots() { ++#ifndef USE_ROCM ++ return static_cast(vllm::persistent::kNumTopKExactCounters); ++#else ++ return 0; ++#endif ++} +diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp +index 44c5f95..53f4577 100644 +--- a/csrc/libtorch_stable/torch_bindings.cpp ++++ b/csrc/libtorch_stable/torch_bindings.cpp +@@ -590,6 +590,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + // Read (and optionally reset) the medium-path exact-restart counters. + ops.def( + "persistent_topk_exact_counters(Tensor! out, int reset, int sync) -> ()"); ++ ops.def("persistent_topk_exact_counters_slots() -> int"); + + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( +@@ -866,6 +867,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); + ops.impl("persistent_topk_exact_counters", + TORCH_BOX(&persistent_topk_exact_counters)); ++ ops.impl("persistent_topk_exact_counters_slots", ++ TORCH_BOX(&persistent_topk_exact_counters_slots)); + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); + #endif +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index 9c775b0..fa6ee84 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1081,7 +1081,7 @@ def run_large_context_topk_test( + Torch: {torch_vals.sort(descending=True)[0][:10]}""" + + +-@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+") ++@pytest.mark.skipif(not _has_device_capability(90), reason="cooperative_topk is sm_90+ only") + @pytest.mark.parametrize("top_k", [512, 1024, 2048]) + def test_cooperative_topk_cs2(top_k: int) -> None: + """The 64-row dispatch uses the two-CTA cooperative kernel. +@@ -1861,8 +1861,13 @@ def test_sparse_indexer_topk_backend_resolution() -> None: + # --------------------------------------------------------------------------- + + # Slot map of torch.ops._C.persistent_topk_exact_counters (see +-# csrc/libtorch_stable/topk.cu). The large path appended four counters and a +-# second switch flag, so the medium-path "enabled" flag moved from [4] to [8]. ++# csrc/libtorch_stable/topk.cu). Eleven slots, the ++# union of the tie-determinism and large-path determinism layouts with nothing renumbered -- the large path appended four ++# large-path counters and a second switch flag (so the medium-path "enabled" flag ++# moved from [4] to [8]) and the tie-determinism change numbered its decode counter AROUND that ++# block at [10]. Slots 8-9 are `kSlotSwitchMedium` / `kSlotSwitchLarge` in ++# the kernel: host-written flags, not device counters. ++_EXACT_COUNTER_SLOTS = 11 + _SLOT_MEDIUM_ROWS = 0 + _SLOT_MEDIUM_OVERFLOW = 1 + _SLOT_MEDIUM_FALLBACK = 2 +@@ -1873,7 +1878,7 @@ _SLOT_LARGE_TIES_DROPPED = 6 + _SLOT_LARGE_FAULT = 7 + _SLOT_EXACT_MEDIUM_ON = 8 + _SLOT_DET_LARGE_ON = 9 +-_EXACT_COUNTER_SLOTS = 10 ++_SLOT_DECODE_OVERFLOW = 10 + + + def _read_exact_counters() -> list[int] | None: +@@ -1882,6 +1887,7 @@ def _read_exact_counters() -> list[int] | None: + [0..3] medium: rows, overflow, exact fallbacks, faults. + [4..7] large: deterministic rows, tie cuts, elements cut, faults. + [8] MERLIN_TOPK_EXACT_MEDIUM, [9] MERLIN_TOPK_DET_LARGE. ++ [10] decode overflow. + """ + op = getattr(torch.ops._C, "persistent_topk_exact_counters", None) + if op is None: +@@ -2022,9 +2028,10 @@ def test_persistent_topk_medium_overflow_is_bit_reproducible() -> None: + + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("want_recent", [True, False]) + @torch.inference_mode() +-def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: +- """Most-recent-wins at the cutoff, with an analytically known answer. ++def test_persistent_topk_medium_overflow_tie_direction(want_recent: bool) -> None: ++ """The exact restart's tie cut, in BOTH directions, analytically known. + + Random float32 values essentially never tie at the k-th place, so the + exactness tests above never exercise the tie rule at all -- they pass +@@ -2037,14 +2044,39 @@ def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: + index-radix tie rounds actually run, + * the cut spans several high-byte index buckets and lands mid-bucket. + +- Under most-recent-wins the answer is forced: the top_k HIGHEST indices, +- emitted ascending. `torch.topk` cannot supply this expectation -- with all ++ Both directions are forced answers: the top_k HIGHEST indices under ++ most-recent-wins, the top_k LOWEST under oldest-wins, emitted ascending ++ either way. `torch.topk` cannot supply either expectation -- with all + values equal its index choice is arbitrary -- so the reference is computed + from the rule itself. ++ ++ The direction is a launch constant resolved once per ++ process from `MERLIN_TIE_DET` / `MERLIN_TIE_RECENT`, so the arm that does ++ not match this worker's environment cannot be exercised here and skips ++ LOUDLY naming the environment that would run it. Both arms must be run, ++ in two fresh worker processes -- a green run of one arm is half the gate. + """ + torch.set_default_device("cuda:0") + num_rows, seq_len, width, top_k = 4, 12345, 12352, 2048 + ++ # effective_tie_recent(det_ties, tie_recent) in persistent_topk.cuh, in ++ # Python. Keep the two in step: this is the reported policy, and the ++ # kernel's launcher prints the same resolution to stderr at first use. ++ det_ties = bool(int(os.environ.get("MERLIN_TIE_DET", "0") or "0")) ++ tie_recent = bool(int(os.environ.get("MERLIN_TIE_RECENT", "0") or "0")) ++ effective_recent = tie_recent if det_ties else True ++ if effective_recent != want_recent: ++ pytest.skip( ++ f"this worker resolves the tie direction to " ++ f"{'most-recent' if effective_recent else 'oldest'}-wins " ++ f"(MERLIN_TIE_DET={int(det_ties)} " ++ f"MERLIN_TIE_RECENT={int(tie_recent)}); run the " ++ f"{'most-recent' if want_recent else 'oldest'}-wins arm in a fresh " ++ f"process with " ++ + ("MERLIN_TIE_DET=0 (or =1 with MERLIN_TIE_RECENT=1)" ++ if want_recent else "MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=0") ++ ) ++ + before = _read_exact_counters() + if before is None: + pytest.skip( +@@ -2068,17 +2100,134 @@ def test_persistent_topk_medium_overflow_tie_keeps_highest_indices() -> None: + assert after[2] - before[2] == num_rows, "the exact restart did not run" + assert after[3] == before[3], "the exact fallback reported a fault" + +- # The cut must genuinely exercise both tie rounds, or the test is weaker +- # than it looks: several high-byte buckets, and not on a bucket boundary. +- cut = seq_len - top_k +- assert (cut >> 8) != ((seq_len - 1) >> 8), "tie cut spans only one bucket" +- assert cut % 256 != 0, "tie cut sits on a bucket boundary" ++ # The cut must genuinely exercise both index-radix rounds, or the test is ++ # weaker than it looks. Each direction has its own cut index and its own ++ # strength, stated separately rather than asserted as if they were one. ++ if want_recent: ++ # Highest top_k indices: cut at seq_len - top_k, several high-byte ++ # buckets away from the top and landing mid-bucket. ++ lo, hi = seq_len - top_k, seq_len ++ cut = lo ++ assert (cut >> 8) != ((seq_len - 1) >> 8), "tie cut spans only one bucket" ++ assert cut % 256 != 0, "tie cut sits on a bucket boundary" ++ else: ++ # Lowest top_k indices: cut at top_k - 1. top_k is a whole number of ++ # 256-index buckets, so this one necessarily lands on the LAST slot of ++ # its bucket -- which makes round B walk all 256 low bins before it ++ # crosses, the opposite corner from the arm above. ++ lo, hi = 0, top_k ++ cut = hi - 1 ++ assert (cut >> 8) != 0, "tie cut spans only one bucket" ++ assert cut % 256 == 255, "the oldest-wins cut should end its bucket" ++ expected = torch.arange(lo, hi, dtype=torch.int32, device="cuda") ++ rule = "most-recent-wins" if want_recent else "oldest-wins" ++ for row in range(num_rows): ++ assert torch.equal(indices[row], expected), ( ++ f"row {row}: {rule} should select indices " ++ f"[{lo}, {hi}) in ascending order; got " ++ f"{indices[row][:4].tolist()}..{indices[row][-4:].tolist()}" ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# DECODE candidate truncation, the adversarial class. ++# ++# The decode path buffers the coarse threshold bin's candidates into a ++# 3,708-slot shared buffer (`DBUF`) and refines only what it retained. ++# MERLIN_TIE_DET ranks the survivors deterministically; it cannot recall a ++# winner that was never buffered. The historical zero-overflow measurements ++# establish workload coverage, not a capacity proof. ++# --------------------------------------------------------------------------- ++ + +- expected = torch.arange(cut, seq_len, dtype=torch.int32, device="cuda") ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_decode_overflow_is_detected() -> None: ++ """The failing control for the decode-overflow counter. ++ ++ The adversarial row: 8,192 elements that all land in ONE coarse ++ half-precision bin, so the coarse pass separates nothing and every element ++ becomes a buffered candidate -- 8,192 of them into 3,708 slots. Within ++ that bin three FP32 levels, all rounding to the same half: ++ ++ * indices [0, 511) the highest value -> the 511 true leaders, ++ * index seq_len - 1 the middle value -> the 512th true winner, ++ * everything else the lowest value. ++ ++ Element order does the rest: the 512th winner sits at the END of the row, ++ so it is appended after capacity is gone and is dropped before any ranking ++ runs. The returned set then holds the 511 leaders plus an arbitrary ++ element of the bulk. ++ ++ This test asserts the COUNTER, not the repair: it is the control that ++ fails on a build without the detection (there is no slot 10 to read, and ++ the op rejects the 11-element buffer), and it is what makes the decode ++ path's determinism claim conditional on a measured number rather than on ++ an unmeasured assumption. Exactness is asserted only when the repair is ++ actually enabled, since MERLIN_TOPK_EXACT_DECODE is a launch constant this ++ process cannot change. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows, seq_len, top_k = 4, 8192, 512 ++ leaders = top_k - 1 # 511; the 512th winner is the tail element ++ ++ before = _read_exact_counters() ++ if before is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: no " ++ "persistent_topk_exact_counters op. A skip is not a pass." ++ ) ++ ++ # Three FP32 levels a few ULPs apart. Half-precision spacing near 1.5 is ++ # 2**-10, about 8,000x these gaps, so all three round to the same half and ++ # therefore to the same coarse bin -- asserted below rather than assumed. ++ base = np.float32(1.5) ++ mid = np.nextafter(np.nextafter(base, np.float32(np.inf)), np.float32(np.inf)) ++ high = np.nextafter(np.nextafter(mid, np.float32(np.inf)), np.float32(np.inf)) ++ assert base < mid < high ++ assert ( ++ np.float16(base) == np.float16(mid) == np.float16(high) ++ ), "the three levels must share one half-precision bin or the row is not adversarial" ++ ++ logits = torch.full((num_rows, seq_len), float(base), dtype=torch.float32) ++ logits[:, :leaders] = float(high) ++ logits[:, seq_len - 1] = float(mid) ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32) ++ indices = torch.empty((num_rows, top_k), dtype=torch.int32) ++ _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, seq_len) ++ torch.accelerator.synchronize() ++ ++ after = _read_exact_counters() ++ fired = after[_SLOT_DECODE_OVERFLOW] - before[_SLOT_DECODE_OVERFLOW] ++ assert fired == num_rows, ( ++ f"the decode candidate buffer overflowed on {num_rows} rows by " ++ f"construction (8,192 candidates into 3,708 slots) but the counter " ++ f"moved by {fired}; the truncation is happening unseen" ++ ) ++ ++ # Whatever the arm, the row must still be a well-formed selection. ++ for row in range(num_rows): ++ got = indices[row].tolist() ++ assert len(set(got)) == top_k, f"row {row}: duplicate indices" ++ assert all(0 <= i < seq_len for i in got), f"row {row}: index out of range" ++ ++ if not bool(int(os.environ.get("MERLIN_TOPK_EXACT_DECODE", "0") or "0")): ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_DECODE=0 (the default): detection asserted " ++ "above, the repair is not enabled in this process, so exactness " ++ "is NOT claimed for this row. Run this test again in a fresh " ++ "process with MERLIN_TOPK_EXACT_DECODE=1 for the repaired arm -- " ++ "a green detect-only run is half the gate." ++ ) ++ ++ # Repaired arm: the exact answer is forced and contains the tail winner. ++ expected = torch.tensor( ++ list(range(leaders)) + [seq_len - 1], dtype=torch.int32, device="cuda" ++ ) + for row in range(num_rows): + assert torch.equal(indices[row], expected), ( +- f"row {row}: most-recent-wins should select indices " +- f"[{cut}, {seq_len}) in ascending order; got " ++ f"row {row}: the exact decode restart should select " ++ f"[0, {leaders}) plus {seq_len - 1}; got " + f"{indices[row][:4].tolist()}..{indices[row][-4:].tolist()}" + ) + +@@ -2111,6 +2260,16 @@ def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: + ok = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int64, device="cuda") + op(ok, 0, 1) # the control: this layout must be accepted + ++ slots_op = getattr(torch.ops._C, "persistent_topk_exact_counters_slots", None) ++ assert slots_op is not None, ( ++ "the build must describe its own counter layout; a guessed length is " ++ "how a stale probe reads the wrong slots" ++ ) ++ assert int(slots_op()) == _EXACT_COUNTER_SLOTS, ( ++ f"build reports {int(slots_op())} slots, test expects " ++ f"{_EXACT_COUNTER_SLOTS}" ++ ) ++ + wrong_dtype = torch.zeros(_EXACT_COUNTER_SLOTS, dtype=torch.int32, device="cuda") + with pytest.raises(Exception): + op(wrong_dtype, 0, 1) +@@ -2131,6 +2290,20 @@ def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None: + with pytest.raises(Exception): + op(too_small, 0, 1) + ++ # Too LARGE is rejected too. The check used to be `>=`, so a buffer sized ++ # for one layout was accepted by a build with another and the surplus ++ # slots read 0 -- a plausible answer from the wrong protocol. ++ too_large = torch.zeros(_EXACT_COUNTER_SLOTS + 1, dtype=torch.int64, device="cuda") ++ with pytest.raises(Exception): ++ op(too_large, 0, 1) ++ ++ # This case matters: a stale medium-path five-element probe must ++ # be REJECTED, not filled with the wrong meanings. It is the same length ++ # check as above, asserted at the length a real stale caller would use. ++ medium_sized = torch.zeros(5, dtype=torch.int64, device="cuda") ++ with pytest.raises(Exception): ++ op(medium_sized, 0, 1) ++ + with pytest.raises(Exception): + op(ok, 1, 0) # reset without sync is not a complete handoff + +@@ -2159,24 +2332,49 @@ def test_topk_exact_counters_reader_log_line(monkeypatch) -> None: + monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", str(_EXACT_COUNTER_SLOTS)) + assert envs.MERLIN_TOPK_EXACT_COUNTERS_LOG == 1 + ++ # The build now exposes the slot count, so the env fallback above is ++ # dead on this branch -- assert the build is what answers, because a ++ # reader that guessed a stale length would be rejected by the op. ++ assert topk_counters.slot_count() == _EXACT_COUNTER_SLOTS, ( ++ "the build and the test disagree about the counter layout" ++ ) + counters = topk_counters.read() ++ # The dict is in SLOT order, which is NOT the line order: the line puts the ++ # five medium-path fields first (see `_LINE_ORDER`), the array puts the large path's ++ # large-path block at 4-7 and the two switch flags at 8-9. Asserting the ++ # first five dict keys against the medium-path names is what fails on any build ++ # carrying the large-path block, so assert the whole layout instead. + assert list(counters) == [ + "rows", + "overflow", + "fallbacks", + "faults", ++ "rows_large", ++ "ties_cut", ++ "ties_dropped", ++ "faults_large", + "switch", +- ], f"unexpected slot names: {list(counters)}" ++ "switch_large", ++ "decode_overflow", ++ ], f"unexpected slot layout: {list(counters)}" + assert counters == dict(zip(counters, direct, strict=True)), ( + "the reader disagrees with a direct op call" + ) + + line = topk_counters.log_line("PP0_TP0") ++ # The documented medium-path prefix is asserted verbatim -- the smoke keys on it ++ # -- and any later slot may follow as `name=`. + assert re.fullmatch( + r"topk-exact-counters rank=PP0_TP0 rows=\d+ overflow=\d+ " +- r"fallbacks=\d+ faults=\d+ switch=\d+", ++ r"fallbacks=\d+ faults=\d+ switch=\d+(?: [a-z_0-9]+=\d+)*", + line, + ), f"log line does not match the documented format: {line!r}" ++ assert " decode_overflow=" in line, ( ++ f"the decode-overflow counter is not in the line: {line!r}" ++ ) ++ assert " switch_large=" in line, ( ++ f"the large-path switch flag is not in the line: {line!r}" ++ ) + + # The cadence gate must also be able to emit -- a reader that never fires + # from maybe_log() is the same defect one level up. +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 891954b..4e8513e 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + import os ++import re + from pathlib import Path + from unittest.mock import patch + +@@ -603,6 +604,8 @@ class TestVllmMaxNSequences: + # rests on not changing the value between capture and replay in one process. + # These tests check the AOT side, which is the side this file can check. + # ++# The determinism toggles, and the gate that keeps unaudited knobs out. ++# + # Each row is (name, value to set, value the getter must then return). The + # expected value matters: a getter that ignored the environment would still + # satisfy a bare membership check, so every row proves the entry is live. +@@ -646,17 +649,35 @@ _FORK_IGNORED = [ + ("VLLM_E2_MAMBA_SLOT_OWNERSHIP", " RAISE ", "raise"), + ] + +-# Not in this repository at all: a reference-host-side patch stack. Named in +-# the TODO block in envs.py and deliberately unregistered until someone reads +-# their real read sites and classifies both identities. +-_FORK_UNAUDITED = [ +- "MERLIN_FULL_K", +- "MERLIN_QSA_SORT", +- "MERLIN_TIE_DET", +- "MERLIN_TIE_RECENT", +- "MERLIN_PIN_BSM", ++# Marker fence around the unaudited-knob block in envs.py. The gate below scans ++# BETWEEN these, never the whole file -- see its docstring for why. ++_UNAUDITED_START = "# --8<-- [start:unaudited-host-side-knobs]" ++_UNAUDITED_END = "# --8<-- [end:unaudited-host-side-knobs]" ++ ++# All four are compile/capture factors and must be hashed. Three are raw ++# passthroughs of a value that C++ parses; MERLIN_QSA_SORT is parsed here. ++_TIEDET_HASHED = [ ++ ("MERLIN_FULL_K", "1", "1"), ++ ("MERLIN_TIE_DET", "1", "1"), ++ ("MERLIN_TIE_RECENT", "1", "1"), ++ ("MERLIN_QSA_SORT", "1", True), + ] + ++# The persistent_topk kernel switches. Raw passthroughs ++# of values C++ parses, hashed on the same conservative rule as the four above. ++# They are NOT in _TIEDET_HASHED because "unset reads falsy" says nothing useful ++# about them -- MERLIN_TOPK_EXACT_MEDIUM's C++ default is ON, so asserting it ++# under test_tiedet_toggle_defaults_off would be a test that cannot fail in the ++# direction that matters. ++_TIEDET_KERNEL_SWITCHES = [ ++ ("MERLIN_TOPK_EXACT_MEDIUM", "0", "0"), ++ ("MERLIN_TOPK_EXACT_DECODE", "1", "1"), ++ ("MERLIN_TOPK_DET_LARGE", "1", "1"), ++] ++ ++# Named in the envs.py fenced block, deliberately NOT registered. ++_TIEDET_UNAUDITED = ["MERLIN_PIN_BSM"] ++ + + def _fresh_env_value(name: str): + if hasattr(envs.__getattr__, "cache_clear"): +@@ -678,6 +699,43 @@ def test_fork_knob_is_a_cache_factor( + assert factors[name] == expected + + ++def _unaudited_names(source: str) -> set[str]: ++ """The knob names listed between the two fence markers. ++ ++ Identifiers, not a substring scan: `"MERLIN_FULL_K" in block` is true for a ++ fence whose prose merely mentions the name, which is the same class of ++ defect one level down. Raises if the fence is missing or inverted, so ++ deleting a marker fails the gate instead of silently emptying it. ++ """ ++ start = source.find(_UNAUDITED_START) ++ end = source.find(_UNAUDITED_END) ++ assert start != -1, f"missing marker {_UNAUDITED_START!r} in envs.py" ++ assert end != -1, f"missing marker {_UNAUDITED_END!r} in envs.py" ++ assert start < end, "unaudited-knob markers are inverted in envs.py" ++ block = source[start + len(_UNAUDITED_START) : end] ++ return set(re.findall(r"\b[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+\b", block)) ++ ++ ++def _envs_source() -> str: ++ return Path(envs.__file__).read_text(encoding="utf-8") ++ ++ ++@pytest.mark.parametrize( ++ ("name", "value", "expected"), _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES ++) ++def test_tiedet_toggle_is_a_cache_factor( ++ name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch ++): ++ """A knob that can change a compiled or captured artifact must be hashed, ++ and the hashed value must follow the environment.""" ++ assert name in environment_variables ++ monkeypatch.setenv(name, value) ++ assert _fresh_env_value(name) == expected ++ factors = envs.compile_factors() ++ assert name in factors ++ assert factors[name] == expected ++ ++ + @pytest.mark.parametrize(("name", "value", "expected"), _FORK_IGNORED) + def test_fork_runtime_knob_is_not_a_cache_factor( + name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch +@@ -691,41 +749,128 @@ def test_fork_runtime_knob_is_not_a_cache_factor( + assert name not in envs.compile_factors() + + +-@pytest.mark.parametrize("name", _FORK_UNAUDITED) ++@pytest.mark.parametrize(("name", "value", "expected"), _TIEDET_HASHED) ++def test_tiedet_toggle_defaults_off( ++ name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch ++): ++ """Promotion is a release decision. Unset, every toggle must read as off.""" ++ monkeypatch.delenv(name, raising=False) ++ assert not _fresh_env_value(name) ++ ++ ++@pytest.mark.parametrize( ++ "name", ++ [n for n, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES] + _TIEDET_UNAUDITED, ++) + def test_unaudited_host_side_knob_stays_listed_until_it_is_classified(name: str): +- """The TODO block in envs.py is a gate, not a note. ++ """The fenced block in envs.py is a gate, not a note. + + Registering one of these without classifying it -- which means deleting it + from that block once its AOT and capture identity are written down against + its real read site -- is the exact failure the fork's env rule + exists to prevent, so it fails here. ++ ++ The scan is scoped to the fenced block, not to the whole file. Scanning the ++ whole file is the defect of an earlier version: a legitimately ++ registered knob NECESSARILY appears in envs.py -- as its declaration and as ++ its `environment_variables` key -- so `name in source` is True for every ++ registered knob and the `assert not listed` branch could never pass. The ++ gate could only ever be satisfied by leaving the knob unregistered, which ++ inverts what it was written to enforce. + """ +- source = Path(envs.__file__).read_text(encoding="utf-8") +- # Scan ONLY the delimited block. Scanning the whole file can never pass +- # once a knob is legitimately registered, because the registry entry +- # contains the knob's own name -- the gate would block the very act it +- # exists to sequence. +- start = source.index("# --8<-- [start:fork-unaudited-knobs]") +- end = source.index("# --8<-- [end:fork-unaudited-knobs]") +- assert start < end, "the unaudited-knob block markers are inverted" +- block = source[start:end] +- listed = name in block ++ listed = name in _unaudited_names(_envs_source()) + if name in environment_variables: + assert not listed, ( +- f"{name} is registered but still named in the envs.py TODO block: " +- "classify its AOT and capture identity from its real read site and " +- "remove it from that block" ++ f"{name} is registered but still named in the envs.py unaudited " ++ "block: classify its AOT and capture identity from its real read " ++ "site and remove it from that block" + ) + else: + assert listed, ( +- f"{name} is neither registered nor named in the envs.py TODO block; " +- "every fork knob must be accounted for" ++ f"{name} is neither registered nor named in the envs.py unaudited " ++ "block; every fork knob must be accounted for" + ) + + ++def test_unaudited_gate_discriminates_both_worlds(): ++ """Control for the gate above: it must be able to FAIL. ++ ++ A membership test that reports the same thing in the registered and the ++ unregistered world measures nothing -- which is precisely how the previous ++ whole-file version passed unnoticed. So assert the scan reports differently on ++ a source that lists a knob inside the fence and one that names it only ++ outside, and assert the real envs.py is in the state the gate claims. ++ """ ++ listed = "\n".join( ++ [_UNAUDITED_START, " # MERLIN_EXAMPLE", _UNAUDITED_END, ""] ++ ) ++ registered = "\n".join( ++ [ ++ " # MERLIN_EXAMPLE was classified and removed from the fence.", ++ _UNAUDITED_START, ++ _UNAUDITED_END, ++ " MERLIN_EXAMPLE: str | None = None", ++ ' "MERLIN_EXAMPLE": lambda: os.getenv("MERLIN_EXAMPLE", None),', ++ "", ++ ] ++ ) ++ assert "MERLIN_EXAMPLE" in _unaudited_names(listed) ++ assert "MERLIN_EXAMPLE" not in _unaudited_names(registered) ++ # The whole-file scan the defective version used cannot tell them apart, ++ # and neither can a substring scan of the fence. ++ assert "MERLIN_EXAMPLE" in listed and "MERLIN_EXAMPLE" in registered ++ ++ # And the live file really is in both states at once, one knob each way. ++ names = _unaudited_names(_envs_source()) ++ assert "MERLIN_FULL_K" in _envs_source() ++ assert "MERLIN_FULL_K" not in names ++ assert "MERLIN_PIN_BSM" in names ++ ++ # A missing fence must fail loudly rather than scan nothing. ++ with pytest.raises(AssertionError): ++ _unaudited_names("no markers here") ++ ++ ++def test_exact_trap_is_a_build_define_not_an_env_var(): ++ """MERLIN_TOPK_EXACT_TRAP must stay unregistered, for a stated reason. ++ ++ The kernel switches must be registered. This ++ one is not a switch: persistent_topk.cuh consumes it with `#ifndef` at ++ build time and nothing anywhere calls getenv on it, so an envs.py entry ++ would advertise a runtime knob that does not exist. Pin both halves -- ++ absent from the registry, and present in the header as a preprocessor ++ define -- so a later "tidy-up" that registers it has to argue with this. ++ """ ++ assert "MERLIN_TOPK_EXACT_TRAP" not in environment_variables ++ header = ( ++ Path(envs.__file__).parents[1] ++ / "csrc" ++ / "libtorch_stable" ++ / "persistent_topk.cuh" ++ ) ++ if not header.is_file(): # source tree not beside an installed wheel ++ return ++ text = header.read_text(encoding="utf-8") ++ assert "#ifndef MERLIN_TOPK_EXACT_TRAP" in text ++ assert "getenv(\"MERLIN_TOPK_EXACT_TRAP\"" not in text ++ ++ ++def test_tiedet_toggles_are_not_ignored_factors(): ++ """Negative control for the table above: a toggle silently added to ++ ignored_factors would keep every test green while dropping it from the ++ cache key.""" ++ monkeypatch_free_factors = envs.compile_factors() ++ for name, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES: ++ assert name in environment_variables, name ++ assert name in monkeypatch_free_factors, name ++ for name in _TIEDET_UNAUDITED: ++ assert name not in environment_variables, name ++ assert name not in monkeypatch_free_factors, name ++ ++ + def test_fork_cache_factor_lists_are_disjoint_and_registered(): + """Negative control for the two tables above.""" +- hashed = {n for n, _, _ in _FORK_HASHED} ++ hashed = {n for n, _, _ in _FORK_HASHED} | {n for n, _, _ in _TIEDET_HASHED} + ignored = {n for n, _, _ in _FORK_IGNORED} + assert not hashed & ignored + for name in hashed | ignored: +diff --git a/vllm/envs.py b/vllm/envs.py +index 5923dc5..960d92b 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -455,32 +455,12 @@ if TYPE_CHECKING: + # an identity, it does not choose one. Ignored factor. + VLLM_E2_EXPECTED_CUDAGRAPH_MODE: str = "" + # +- # --8<-- [start:fork-unaudited-knobs] +- # --- TODO: knobs that are NOT in this tree at all ---------- +- # The tie-determinism toggles and the BSM pin live in a separate +- # patch stack, not in v2.0.1 and not in this tree: +- # +- # MERLIN_FULL_K MERLIN_QSA_SORT MERLIN_TIE_DET MERLIN_TIE_RECENT +- # MERLIN_PIN_BSM +- # +- # They are therefore UNAUDITED: nothing in this repository shows what they +- # change, and this file will not classify code it cannot read. The rule they +- # fall under is the standing one -- an os.environ (or std::getenv) knob that +- # changes a compiled or captured artifact is not in the cache key unless it +- # is registered here. From what is known of them (Marlin split-K, a +- # persistent_topk ordering change and a cutoff tie-break) they are very +- # likely capture factors and possibly AOT factors, which is exactly why a +- # guess is not good enough. +- # +- # WHEN THE PATCHES LAND, for each name: register it below, label both +- # identities from its actual read site, hash it unless exclusion is proven, +- # and delete it from this block. tests/test_envs.py holds a test that fails +- # if a name is registered while still listed here, so the classification +- # cannot be skipped silently. That test reads ONLY the delimited block +- # below/above -- scanning the whole file could never pass, because a +- # registry entry necessarily contains its own name. +- # ---------------------------------------------------------------------- +- # --8<-- [end:fork-unaudited-knobs] ++ # The unaudited-knob fence that stood here listed the ++ # tie-determinism toggles as not-yet-in-tree. They landed with the ++ # v2.2 kernel port: the four are registered and ++ # classified in the "determinism toggles" block below, and the ++ # one knob still unclassified sits in that block's own fence, which ++ # tests/test_envs.py reads as a set of identifiers. + VLLM_LOG_MODEL_INSPECTION: bool = False + VLLM_DEBUG_MFU_METRICS: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False +@@ -507,6 +487,112 @@ if TYPE_CHECKING: + VLLM_NIC_SELECTION_VARS: str = "" + VLLM_ENABLE_HPC_OPS: bool = False + MERLIN_TOPK_EXACT_COUNTERS_LOG: int = 0 ++ # ++ # --- determinism toggles ----------------------------------------- ++ # The three nondeterminism mechanisms (FULL_K, QSA_SORT, TIE_DET/TIE_RECENT). ++ # Every one of them is ++ # OFF by default: promotion is the maintainer's decision. ++ # ++ # Why these are registered at all: validate_environ() only inspects names ++ # beginning with VLLM_, so a MERLIN_ name never trips the unknown-variable ++ # warning. Registration here exists for ONE reason -- compile_factors() can ++ # only hash names it knows, and compile_factors() is a deny-list, so ++ # registering a name hashes it unless ignored_factors says otherwise. All ++ # four are deliberately absent from ignored_factors. ++ # ++ # CAPTURE: yes, all four. Each resolves once per process and the resolved ++ # value then reaches the GPU as a kernel launch argument, a launch ++ # configuration, or an extra op in the traced sequence -- so a cudagraph ++ # bakes in whatever was live at capture time. Flipping one between boots ++ # changes the captured graph, which is exactly what a cache key is for. ++ # AOT: not established for the three CUDA toggles (the Marlin and topk ++ # kernels are compiled into the extension at build time, not by the AOT ++ # compiler); established for MERLIN_QSA_SORT, which adds torch ops to the ++ # traced graph. Hashed conservatively throughout, the same rule ++ # already applied to MERLIN_TOPK_EXACT_MEDIUM. ++ # ++ # NOT A PYTHON READ SITE (FULL_K / TIE_DET / TIE_RECENT): each is resolved ++ # once by a std::getenv in C++ -- merlin_full_k() in ++ # csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu and merlin_det_ties() / ++ # merlin_tie_recent() in csrc/libtorch_stable/topk.cu -- and that is the ++ # single source of truth. The entries below are raw passthroughs so nothing ++ # can mistake them for a second parser; do not use them to decide ++ # behaviour. A raw passthrough hashes "0" and "00" differently even though ++ # the C++ treats them alike: over-invalidation, which is the safe direction. ++ # ++ # Marlin MoE split-K reduction order; the CUDA half of unmerged upstream ++ # PR #46639 (`use_full_k`). ++ MERLIN_FULL_K: str | None = None ++ # persistent_topk tie-break at the k-th cutoff, and its direction. ++ MERLIN_TIE_DET: str | None = None ++ MERLIN_TIE_RECENT: str | None = None ++ # ++ # A REAL Python read site, unlike the three above: parsed here and read ++ # once at import by models/qwen4_exp/nvidia/ops/qsa_indexer.py, which adds ++ # a torch.sort + two torch.where ops to the QSA selection path. ++ MERLIN_QSA_SORT: bool = False ++ # ++ # --- persistent_topk kernel switches ------------------------------ ++ # Same classification as the three raw passthroughs above and registered ++ # for the same single reason: compile_factors() is a deny-list, so a name ++ # it does not know cannot be hashed, and a Ray deployment that propagates ++ # registered variables can otherwise drop an explicitly-disabled switch ++ # and leave a worker on the default. Each is resolved ONCE per process by ++ # a std::getenv in csrc/libtorch_stable/topk.cu -- that is the single ++ # source of truth -- and the resolved value reaches the GPU as a kernel ++ # launch argument, so a cudagraph bakes in whatever was live at capture. ++ # CAPTURE: yes, all three. AOT: not established (these kernels are built ++ # into the extension, not by the AOT compiler); hashed conservatively ++ # anyway, same rule as above. None is in ignored_factors. ++ # ++ # Raw passthroughs, so "unset" reads as None here. That is NOT the same as ++ # "off" for the C++: MERLIN_TOPK_EXACT_MEDIUM defaults ON. Read the getter ++ # in topk.cu, not this default, for the shipped behaviour. ++ # ++ # Take the exact restart on an overflowing MEDIUM row: ++ # MERLIN_TOPK_EXACT_MEDIUM. Declared ONCE, in the cross-change knob ++ # block above (same raw passthrough, same hashing); the 0.30 tree already ++ # carried that entry, so it is not repeated here. ++ # The same restart on an overflowing DECODE row. ++ MERLIN_TOPK_EXACT_DECODE: str | None = None ++ # Deterministic placement and tie membership on the LARGE path. ++ # Registered here ahead of its code, exactly as the four toggles above ++ # were, so the full tree is complete; its read site is ++ # det_large_enabled() in topk.cu on the large path. Setting it without that ++ # code does nothing, which is why it is a passthrough and not a parsed ++ # bool that would imply this tree can act on it. ++ MERLIN_TOPK_DET_LARGE: str | None = None ++ # ++ # NOT REGISTERED, deliberately: MERLIN_TOPK_EXACT_TRAP. It is a ++ # preprocessor define consumed by `#ifndef` in persistent_topk.cuh at ++ # BUILD time (-D), not a process environment variable -- nothing in the ++ # tree calls getenv on it. Registering it would advertise a runtime knob ++ # that does not exist. Its default (trap ON) is an open maintainer decision, ++ # not something to settle by ++ # quietly making it settable. ++ # ++ # --- UNAUDITED: reference-host knobs with no read site in this tree ---------- ++ # A knob named inside the fence below is one this repository cannot ++ # classify, because the code that reads it is not in it. Registering one ++ # without first deleting it from the fence -- which means writing down its ++ # AOT and capture identity against its real read site -- is the exact ++ # failure the fork's env rule exists to prevent, and ++ # tests/test_envs.py::test_unaudited_host_side_knob_stays_listed_until_it_is_classified ++ # fails on it. ++ # ++ # The fence holds NOTHING BUT the names, one per line. Prose belongs out ++ # here: the gate reads the fence as a set of identifiers, and a knob name ++ # mentioned in passing inside it would otherwise read as a listing. ++ # ++ # The one entry, MERLIN_PIN_BSM, pins Marlin's block_size_m to 64 (Python, ++ # no rebuild). It is an ALTERNATIVE to the split-K toggle above, not a ++ # complement: on its own it fixed only the 74/102-token cells and left ++ # 1050 upward broken, while use_full_k fixed all of them with block_size_m ++ # still varying. The tie-determinism change promotes use_full_k and keeps this one off, so its ++ # read site was never brought across and it stays unclassified. ++ # --8<-- [start:unaudited-host-side-knobs] ++ # MERLIN_PIN_BSM ++ # --8<-- [end:unaudited-host-side-knobs] + + + def get_default_cache_root(): +@@ -2467,6 +2553,19 @@ environment_variables: dict[str, Callable[[], Any]] = { + "MERLIN_TOPK_EXACT_COUNTERS_LOG": lambda: int( + os.getenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "0") + ), ++ # Determinism toggles. See the declaration block above for the AOT / ++ # capture classification and for why three of these are raw passthroughs. ++ # All four are absent from ignored_factors on purpose, so all four are ++ # hashed into the compile cache key. ++ "MERLIN_FULL_K": lambda: os.getenv("MERLIN_FULL_K", None), ++ "MERLIN_TIE_DET": lambda: os.getenv("MERLIN_TIE_DET", None), ++ "MERLIN_TIE_RECENT": lambda: os.getenv("MERLIN_TIE_RECENT", None), ++ "MERLIN_QSA_SORT": lambda: bool(int(os.getenv("MERLIN_QSA_SORT", "0"))), ++ # persistent_topk kernel switches. Raw passthroughs ++ # of values C++ parses; see the declaration block above. ++ # MERLIN_TOPK_EXACT_MEDIUM is registered once, earlier in this dict. ++ "MERLIN_TOPK_EXACT_DECODE": lambda: os.getenv("MERLIN_TOPK_EXACT_DECODE", None), ++ "MERLIN_TOPK_DET_LARGE": lambda: os.getenv("MERLIN_TOPK_DET_LARGE", None), + } + + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +index a46f915..555e445 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +@@ -23,6 +23,46 @@ _TOPK_WORKSPACE_BYTES = 1024 * 1024 + _DECODE_BLOCK_N = 64 + + ++# --------------------------------------------------------------------------- ++# MERLIN_QSA_SORT: canonicalise the ORDER of the selected block indices. ++# ++# persistent_topk assigns output slots by atomicAdd, so two identical requests ++# get the same selected SET in a different ORDER (measured on the reference host: matching ++# sorted hashes, differing raw hashes). Sparse attention then accumulates over ++# those blocks in whatever order it was handed, and fp32 addition is not ++# associative, so the arithmetic moves. Sorting is semantically a no-op -- the ++# same blocks are attended either way -- and makes the accumulation order ++# canonical. ++# ++# Padding (negative sentinels) is pushed to the END so the valid prefix is ++# preserved: `_expand_qsa_indices_kernel` truncates by POSITION ++# (`complete_blocks = min(visible_blocks, BLOCK_TOPK)`), and position is never ++# compared against a score anywhere in the chain, so there is no relevance-rank ++# contract for the sort to break. ++# ++# Branch-free and sync-free on purpose. `.any()` / `.item()` raise ++# cudaErrorStreamCaptureUnsupported inside graph capture; that cost one boot ++# on the reference host and is why this version has no Python-visible branch on data. ++# ++# Resolved once at import so it is a launch constant: a cudagraph captures the ++# op sequence that was live at capture time. Default OFF -- promotion is a release ++# decision, which is the maintainer's. ++# --------------------------------------------------------------------------- ++_QSA_SORT = envs.MERLIN_QSA_SORT ++ ++ ++def _qsa_canonicalise(block_indices: torch.Tensor) -> None: ++ """Sort selected block indices ascending, padding last. In place, shape preserved.""" ++ if not _QSA_SORT: ++ return ++ big = torch.iinfo(block_indices.dtype).max ++ tmp = torch.where( ++ block_indices < 0, torch.full_like(block_indices, big), block_indices ++ ) ++ tmp = torch.sort(tmp, dim=-1).values ++ block_indices.copy_(torch.where(tmp == big, torch.full_like(tmp, -1), tmp)) ++ ++ + @triton.jit + def _qsa_mqa_paged_uniform_kernel( + q_ptr, +@@ -565,6 +605,7 @@ def _topk( + block_topk, + logits.shape[1], + ) ++ _qsa_canonicalise(block_indices) + + + def qsa_select_paged_decode( +diff --git a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +index 08ffab8..d7db4d7 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +@@ -17,22 +17,40 @@ of each counter. An older build whose op still demands int32 will therefore + raise here; that is the intended signal, not a regression to paper over. + + SLOT LAYOUT is build-dependent, so nothing here hard-codes five. `slot_count()` +-asks the build first (a `..._slots` op, should a later build expose one) and +-otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5); names come from a +-table keyed by that length: ++asks the build first -- the tie-determinism change added the `..._slots` op this was written to look ++for, so on a build that has it the env variable is never consulted -- and ++otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5). The op demands ++EXACTLY that many elements, so a guessed length is an error, not a short read; ++names come from a table keyed by that length: + + 5 rows overflow fallbacks faults switch + 10 rows overflow fallbacks faults rows_large ties_cut ties_dropped + faults_large switch switch_large ++ 11 the ten above, plus decode_overflow + + An unknown length degrades to `slot` names rather than mislabelling. + ++WHY 11 AND NOT 10. The tie-determinism change numbered its one new counter around the large-path block ++instead of after the medium-path one, so the integration renumbers nothing: slots 4-7 are ++the large-path counters, 8-9 the two switch flags, 10 the decode overflow. ++On the integrated build every one of the eleven is live; on a single-branch ++build the other branch's slots are structurally zero and the line still ++reports them by name -- they are the slots they will hold, not a ++mislabelling. `decode_overflow` is the number that has to be read before any ++claim that the DECODE path is deterministic: it counts rows whose buffered ++candidates exceeded the 3,708-slot capacity before the tie ranking ran, which ++ranking cannot repair. ++ + HOW THE SMOKE READS IT. Set `MERLIN_TOPK_EXACT_COUNTERS_LOG=` on the + engine; each GPU worker then emits, at most that often and once more at + shutdown, one line into the ordinary engine log: + + topk-exact-counters rank=PP0_TP1 rows=... overflow=... fallbacks=... +- faults=... switch=... ++ faults=... switch=... [rows_large=... ... decode_overflow=...] ++ ++The first five fields are the medium-path line and keep their order on every build; ++any further slot is appended after them, so an 11-slot build reads as the medium-path ++line plus the four large-path fields, `switch_large` and `decode_overflow`. + + Counts are cumulative per worker process, so the smoke takes the LATEST line + per rank -- `grep topk-exact-counters | ...` keyed on the `rank=` +@@ -73,7 +91,14 @@ _LARGEDET_NAMES = ( + "switch", + "switch_large", + ) +-_SLOT_NAMES: dict[int, tuple[str, ...]] = {5: _MEDIUM_NAMES, 10: _LARGEDET_NAMES} ++# The decode-overflow counter comes after the large-path block; see the module ++# docstring for why it is slot 10 and not slot 4. ++_TIEDET_NAMES = _LARGEDET_NAMES + ("decode_overflow",) ++_SLOT_NAMES: dict[int, tuple[str, ...]] = { ++ 5: _MEDIUM_NAMES, ++ 10: _LARGEDET_NAMES, ++ 11: _TIEDET_NAMES, ++} + + # The medium-path line contract. Any further name is appended after these, in slot + # order, so a 10-slot build still starts with exactly the medium-path line. +@@ -92,7 +117,13 @@ def _get_op(name: str): + + + def slot_count() -> int: +- """How many int64 slots this build's counter array carries.""" ++ """How many int64 slots this build's counter array carries. ++ ++ The build is asked first and the environment is only a fallback for builds ++ older than the slot-count op: a guessed length that does not match the extension is ++ rejected by the op's own length check, so the instrument would go dark ++ exactly when the layout changed -- which is when it matters most. ++ """ + query = _get_op(_OP_NAME + "_slots") + if query is not None: + return int(query()) +-- +2.43.0 + diff --git a/patches/0045-topk-follow-up-fixes-to-the-persistent-top-k-kernels.patch b/patches/0045-topk-follow-up-fixes-to-the-persistent-top-k-kernels.patch new file mode 100644 index 0000000..992f659 --- /dev/null +++ b/patches/0045-topk-follow-up-fixes-to-the-persistent-top-k-kernels.patch @@ -0,0 +1,1082 @@ +From 950a2096aeac1e891d2199c40f0ad121ea925edd Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:35:50 -0400 +Subject: [PATCH 045/103] topk: follow-up fixes to the persistent top-k kernels + +- Counter enum declared above the decode template (two-phase lookup). +- Emission-fault flag set atomically. +- Close the legacy large tie loop with a block barrier (inherited). +- Large-path fault injectors registered as hashed passthroughs. +- The tensor-free counter-slots op registered for CompositeExplicitAutograd. +- Tie measurements recorded only on rows neither fault verdict failed. +- Tests scoped to the properties the build promises. +--- + csrc/libtorch_stable/persistent_topk.cuh | 199 +++++++---- + csrc/libtorch_stable/topk.cu | 41 ++- + csrc/libtorch_stable/torch_bindings.cpp | 8 +- + tests/kernels/test_top_k_per_row.py | 403 +++++++++++++++++++++-- + tests/test_envs.py | 46 ++- + vllm/envs.py | 37 +++ + 6 files changed, 642 insertions(+), 92 deletions(-) + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index f8ce859..85d217f 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -189,6 +189,87 @@ struct PersistentTopKParams { + // constant like the two switches. Setting it also suppresses __trap(), because + // the point of injecting is to observe the trap-off row-failure policy. + int fault_inject_cta; ++ // Which fault the injector forces in that CTA. ++ // false (MERLIN_TOPK_LARGE_FAULT_KIND unset or `pre`): the PRE-EMISSION ++ // predicate `my_keep > local_eq`, caught at X3. Emission is skipped ++ // entirely on that row, so this kind cannot reach the X4 code at all. ++ // true (`emit`): the EMISSION-BOUND slot check inside the placement loop, ++ // caught at X4. This is the only way to exercise `suffix_sum[6]`, and ++ // it exercises it the way the hardware would: every emitting thread of ++ // the named CTA discovers an invalid slot at once, which is exactly the ++ // concurrency the atomic exists for. ++ // Resolved once per process on the host, so it is a launch constant. ++ bool fault_inject_emit; ++}; ++ ++// ============================================================================ ++// Device-global exact-path counters -- DECLARED HERE, ABOVE EVERY USER. ++// ++// This enum used to live beside the medium-path restart, ++// ~200 lines BELOW the decode template that increments ++// `kCtrDecodeOverflow`. `counters` is a plain `unsigned long long*` ++// parameter, so that enumerator is a NON-DEPENDENT name: two-phase lookup ++// resolves it when the template is DEFINED, not when it is instantiated, ++// and the Linux build rejected the tree with ++// persistent_topk.cuh(442): identifier "kCtrDecodeOverflow" is undefined. ++// An offline `nvcc -cubin` on a non-Linux host does NOT reproduce that on its own: nvcc ++// can drive cicc in Microsoft mode (`--ms_c++17 --microsoft_version=`), and that mode ++// defers non-dependent lookup to instantiation, where the enum is already ++// visible. `-Xcudafe --dep_name` restores strict two-phase lookup and the ++// offline harness then reproduces the Linux diagnostic verbatim. ++// ++// The fix is the declaration ORDER, not a literal: every increment still ++// names its enumerator, so a slot renumbering stays a one-line edit here. ++// ============================================================================ ++ ++// Device-global counters. Plain device atomics on a process-lifetime buffer: ++// no host sync and no host-visible allocation on the launch path, so this is ++// safe to record inside a cudagraph capture and on every replay. Read from ++// Python with `torch.ops._C.persistent_topk_exact_counters`. ++// ++// The counters are `unsigned long long`, not `unsigned int`, ++// and the reader tensor is int64. A u32 counter copied into a SIGNED int32 ++// tensor goes negative, then wraps: at 65,536 columns with k=512 an all-equal ++// row drops 65,024 ties, so slot 6 crosses the signed boundary after 33,027 ++// rows and wraps after 66,053, and a delta taken across that boundary is a ++// wrong number that looks like a measurement. `atomicAdd(unsigned long long*)` ++// is native on sm_86, so the width costs nothing on the increment path. ++// ++// SLOT MAP (the tie-determinism + large-path integration, 11 slots). The tie change numbered its decode ++// counter AROUND the large-path block rather than after the medium-path one, so the merged layout is ++// the union of the two branches with nothing renumbered and no external probe ++// silently re-reading a row count as a switch -- the failure a merged layout ++// must avoid. ++// ++// [0] medium rows [1] medium overflow [2] medium exact fallback ++// [3] medium exact fault ++// [4]-[7] large path (rows_large, ties_cut, ties_dropped, faults_large) ++// [8] MERLIN_TOPK_EXACT_MEDIUM [9] MERLIN_TOPK_DET_LARGE ++// [10] decode overflow ++// ++// Slots 8-9 are written by the HOST after the device copy; the device never ++// touches them, and they are slots INSIDE the array, not two elements appended ++// after it. A stale five- or ten-element reader buffer is rejected outright by ++// the length check in the reader op, which is the fail-loud direction. ++enum TopKExactCounter : int { ++ kCtrMediumRows = 0, // denominator: medium-path rows entered ++ kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop ++ kCtrMediumExactFallback = 2, // rows that took the exact restart ++ kCtrMediumExactFault = 3, // exact restart could not complete ++ // Large path (multi-CTA radix_topk, seq_len > RADIX_THRESHOLD). ++ // Slots 4-7 are only ever touched by the deterministic emission, so they ++ // stay 0 in the MERLIN_TOPK_DET_LARGE=0 control arm: `[9] == 1 && [4] == 0` ++ // while large rows are known to have run means the host-side capacity gate ++ // (4 * ctas_per_group <= RADIX) turned the switch off for that launch. ++ kCtrLargeDetRows = 4, // large rows emitted deterministically ++ kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut ++ kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts ++ kCtrLargeFault = 7, // large-path invariant violation ++ kSlotSwitchMedium = 8, // host-written flag, NOT a device counter ++ kSlotSwitchLarge = 9, // host-written flag, NOT a device counter ++ kCtrDecodeOverflow = 10, // decode rows whose buffered-candidate count ++ // exceeded DBUF before deterministic ranking ++ kNumTopKExactCounters = 11, // total slots: device buffer AND reader output + }; + + // ============================================================================ +@@ -604,56 +685,6 @@ __device__ __noinline__ bool histogram_2048_topk( + #define MERLIN_TOPK_EXACT_TRAP 1 + #endif + +-// Device-global counters. Plain device atomics on a process-lifetime buffer: +-// no host sync and no host-visible allocation on the launch path, so this is +-// safe to record inside a cudagraph capture and on every replay. Read from +-// Python with `torch.ops._C.persistent_topk_exact_counters`. +-// +-// The counters are `unsigned long long`, not `unsigned int`, +-// and the reader tensor is int64. A u32 counter copied into a SIGNED int32 +-// tensor goes negative, then wraps: at 65,536 columns with k=512 an all-equal +-// row drops 65,024 ties, so slot 6 crosses the signed boundary after 33,027 +-// rows and wraps after 66,053, and a delta taken across that boundary is a +-// wrong number that looks like a measurement. `atomicAdd(unsigned long long*)` +-// is native on sm_86, so the width costs nothing on the increment path. +-// +-// SLOT MAP (the tie-determinism + large-path integration, 11 slots). The tie change numbered its decode +-// counter AROUND the large-path block rather than after the medium-path one, so the merged layout is +-// the union of the two branches with nothing renumbered and no external probe +-// silently re-reading a row count as a switch -- the failure a merged layout +-// must avoid. +-// +-// [0] medium rows [1] medium overflow [2] medium exact fallback +-// [3] medium exact fault +-// [4]-[7] large path (rows_large, ties_cut, ties_dropped, faults_large) +-// [8] MERLIN_TOPK_EXACT_MEDIUM [9] MERLIN_TOPK_DET_LARGE +-// [10] decode overflow +-// +-// Slots 8-9 are written by the HOST after the device copy; the device never +-// touches them, and they are slots INSIDE the array, not two elements appended +-// after it. A stale five- or ten-element reader buffer is rejected outright by +-// the length check in the reader op, which is the fail-loud direction. +-enum TopKExactCounter : int { +- kCtrMediumRows = 0, // denominator: medium-path rows entered +- kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop +- kCtrMediumExactFallback = 2, // rows that took the exact restart +- kCtrMediumExactFault = 3, // exact restart could not complete +- // Large path (multi-CTA radix_topk, seq_len > RADIX_THRESHOLD). +- // Slots 4-7 are only ever touched by the deterministic emission, so they +- // stay 0 in the MERLIN_TOPK_DET_LARGE=0 control arm: `[9] == 1 && [4] == 0` +- // while large rows are known to have run means the host-side capacity gate +- // (4 * ctas_per_group <= RADIX) turned the switch off for that launch. +- kCtrLargeDetRows = 4, // large rows emitted deterministically +- kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut +- kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts +- kCtrLargeFault = 7, // large-path invariant violation +- kSlotSwitchMedium = 8, // host-written flag, NOT a device counter +- kSlotSwitchLarge = 9, // host-written flag, NOT a device counter +- kCtrDecodeOverflow = 10, // decode rows whose buffered-candidate count +- // exceeded DBUF before deterministic ranking +- kNumTopKExactCounters = 11, // total slots: device buffer AND reader output +-}; +- + // Exact, bounded-memory re-selection for a single medium-path row. + // + // `scores` is already offset to the row. Every radix round rescans the whole +@@ -1419,7 +1450,8 @@ __device__ __noinline__ void radix_emit_deterministic( + uint32_t* local_histogram, uint32_t* suffix_sum, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, uint32_t ctas_per_group, + int& barrier_phase, uint32_t radix_iter, uint32_t tx, +- unsigned long long* counters, int fault_inject_cta, bool tie_recent) { ++ unsigned long long* counters, int fault_inject_cta, ++ bool fault_inject_emit, bool tie_recent) { + constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u; + // Staging for the published per-CTA counts. The host gates the switch on + // 4 * ctas_per_group <= RADIX, so these regions never collide. +@@ -1585,7 +1617,7 @@ __device__ __noinline__ void radix_emit_deterministic( + // Test-only: force the CTA-LOCAL predicate in one named CTA, so a fault + // that originates outside CTA 0 can be observed. Off (-1) in production; + // the host resolves it once per process, so it is a launch constant. +- if (fault_inject_cta >= 0 && ++ if (fault_inject_cta >= 0 && !fault_inject_emit && + static_cast(cta_in_group) == fault_inject_cta) { + my_keep_checked = local_eq + 1u; + } +@@ -1622,22 +1654,22 @@ __device__ __noinline__ void radix_emit_deterministic( + barrier_phase++; + __syncthreads(); + ++ // Tie measurement staged in thread 0's register; recorded only after BOTH ++ // verdicts (X3 here, X4 after emission) have cleared the row. Recording it ++ // here, behind X3 alone, counted tie cuts for rows an emission-bound fault ++ // later failed closed (reference host, FAULT_KIND=emit: ties_cut +8 on 8 ++ // all -1 rows). ++ uint32_t staged_ties_dropped = 0u; + if (tx == 0) { + uint32_t row_fault = 0u; + for (uint32_t c = 0; c < ctas_per_group; ++c) { + row_fault |= cta_counts[kFaultStage + c]; + } + suffix_sum[5] = row_fault; ++ staged_ties_dropped = suffix_sum[8]; + + if (counters != nullptr && cta_in_group == 0u) { + atomicAdd(&counters[kCtrLargeDetRows], 1ull); +- // Only meaningful on a healthy row: on a faulted one the tie arithmetic +- // is what is under suspicion, so do not record it as a measurement. +- if (row_fault == 0u && suffix_sum[8] > 0u) { +- atomicAdd(&counters[kCtrLargeTiesCut], 1ull); +- atomicAdd(&counters[kCtrLargeTiesDropped], +- static_cast(suffix_sum[8])); +- } + } + } + __syncthreads(); +@@ -1648,6 +1680,17 @@ __device__ __noinline__ void radix_emit_deterministic( + const uint32_t keep_width = suffix_sum[9]; + const uint32_t row_fault = suffix_sum[5]; + ++ // TEST ONLY: force the emission-bound violation in one ++ // named CTA. Block-uniform (both operands are launch constants and ++ // `cta_in_group` is block-uniform), so it changes no control flow the whole ++ // block does not take together. It deliberately fires for EVERY emitting ++ // thread of that CTA rather than one: the defect under repair is a ++ // concurrent write to `suffix_sum[6]`, and a single-writer injection would ++ // pass with the broken code. ++ const bool force_emit_fault = ++ fault_inject_emit && fault_inject_cta >= 0 && ++ static_cast(cta_in_group) == fault_inject_cta; ++ + // ---- emission: rank by index, no atomics ------------------------------- + if (row_fault == 0u) { + // Scratch for the tile scan. The staged per-CTA counts above are dead by +@@ -1724,7 +1767,7 @@ __device__ __noinline__ void radix_emit_deterministic( + (eq_rank > cut_rank) ? (eq_rank - cut_rank) : 0u; + if (kept_eq_before > keep_width) kept_eq_before = keep_width; + const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before; +- if (slot < static_cast(TopK)) { ++ if (slot < static_cast(TopK) && !force_emit_fault) { + row_output[slot] = static_cast(my_chunk_start + i); + } else { + // An out-of-range slot used to bump a counter and +@@ -1732,7 +1775,16 @@ __device__ __noinline__ void radix_emit_deterministic( + // promised all-`-1` row. It is now recorded as a row failure and + // goes through the SAME policy as the pre-emission faults below: no + // store here, and the row is filled with -1 after the group agrees. +- suffix_sum[6] = 1u; ++ // ++ // ATOMIC, not a plain store. Every thread whose ++ // slot is out of range reaches this line, and an invariant failure ++ // is precisely the situation in which several of them do at once -- ++ // so the plain `suffix_sum[6] = 1u` was a concurrent non-atomic ++ // write inside the mechanism that exists to detect the failure. It ++ // is a shared-memory word and the flag is monotonic, so `atomicOr` ++ // is the whole fix; the healthy path never executes it. The medium ++ // restart already uses atomicOr for the same job. ++ atomicOr(&suffix_sum[6], 1u); + } + } + +@@ -1765,6 +1817,14 @@ __device__ __noinline__ void radix_emit_deterministic( + emit_fault |= cta_counts[kEmitStage + c]; + } + suffix_sum[7] = emit_fault; ++ // Only meaningful on a healthy row: on a faulted one the tie arithmetic ++ // is what is under suspicion, so do not record it as a measurement. ++ if (counters != nullptr && cta_in_group == 0u && row_fault == 0u && ++ emit_fault == 0u && staged_ties_dropped > 0u) { ++ atomicAdd(&counters[kCtrLargeTiesCut], 1ull); ++ atomicAdd(&counters[kCtrLargeTiesDropped], ++ static_cast(staged_ties_dropped)); ++ } + } + __syncthreads(); + +@@ -1814,7 +1874,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t radix_iter, uint32_t tx, bool det_large, + unsigned long long* counters, int fault_inject_cta, +- bool tie_recent) { ++ bool fault_inject_emit, bool tie_recent) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; +@@ -1975,7 +2035,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, + row_output, ordered_pivot, shared_scalars[1], my_chunk_start, + actual_chunk_size, local_histogram, suffix_sum, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters, +- fault_inject_cta, tie_recent); ++ fault_inject_cta, fault_inject_emit, tie_recent); + return; + } + +@@ -2029,6 +2089,18 @@ __device__ void radix_topk(const float* __restrict__ row_input, + } + } + } ++ // BASELINE repair (inherited from v2.0.1): the tie loop ++ // above reads shared_ordered, and the next row this CTA processes -- a ++ // large row's Stage-1 chunk load, or a decode/medium row reusing the same ++ // dynamic shared memory -- writes it before its first barrier. Without a ++ // closing barrier a fast warp can overwrite entries a slow warp is still ++ // comparing against the pivot, changing tie MEMBERSHIP, not only order. ++ // Every thread of the block reaches this point (the loop bound is ++ // block-uniform and there is no early return in the legacy tail), so the ++ // barrier is uniform. Applies to both comparison arms: the legacy path is ++ // the MERLIN_TOPK_DET_LARGE=0 control, so the control now carries this ++ // repair too, like the histogram-publication barrier (an earlier commit). ++ __syncthreads(); + } + + // ============================================================================ +@@ -2142,6 +2214,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, + params.det_large, params.exact_counters, params.fault_inject_cta, ++ params.fault_inject_emit, + effective_tie_recent(params.det_ties, params.tie_recent)); + radix_iter++; + } +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 2d2c81d..56b0af4 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -6,6 +6,7 @@ + #include + #include + #include ++#include + #include + + #include "ops.h" +@@ -131,6 +132,31 @@ int large_fault_inject_cta() { + return cta; + } + ++// --------------------------------------------------------------------------- ++// MERLIN_TOPK_LARGE_FAULT_KIND. TEST ONLY, companion to ++// the variable above; ignored unless MERLIN_TOPK_LARGE_FAULT_INJECT names a ++// CTA. ++// ++// unset / "pre" (default) force the PRE-EMISSION predicate, caught at X3. ++// "emit" force the EMISSION-BOUND slot check, caught at X4. ++// ++// The two are mutually exclusive by construction: a pre-emission fault makes ++// `row_fault` nonzero, the emission block is skipped wholesale, and the X4 ++// code -- including `suffix_sum[6]`, the word whose concurrent non-atomic ++// write this change is about -- is never reached. The shipped injector could ++// therefore only ever exercise the arm that was already correct, which is why ++// the X4 arm had never been exercised. Resolved once per process like ++// every other switch here, so it is a launch constant a cudagraph bakes in. ++// --------------------------------------------------------------------------- ++bool large_fault_inject_emit() { ++ static const bool emit = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_LARGE_FAULT_KIND"); ++ if (raw == nullptr || raw[0] == '\0') return false; ++ return std::strncmp(raw, "emit", 5) == 0; ++ }(); ++ return emit; ++} ++ + // --------------------------------------------------------------------------- + // Say the effective tie policy out loud, once. + // +@@ -495,6 +521,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + 4u * ctas_per_group <= static_cast(P::RADIX); + params.det_large = det_large_enabled() && det_large_fits; + params.fault_inject_cta = large_fault_inject_cta(); ++ params.fault_inject_emit = large_fault_inject_emit(); + if (det_large_enabled() && !det_large_fits) { + static std::atomic det_large_warned{false}; + if (!det_large_warned.exchange(true)) { +@@ -510,11 +537,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + static std::atomic inject_warned{false}; + if (!inject_warned.exchange(true)) { + std::fprintf(stderr, +- "[vllm] MERLIN_TOPK_LARGE_FAULT_INJECT=%d is set: the " +- "large path FORCES a fault in that CTA of every " +- "deterministic large row and returns an all -1 row. This " +- "is a test mode and must not be set in production.\n", +- params.fault_inject_cta); ++ "[vllm] MERLIN_TOPK_LARGE_FAULT_INJECT=%d kind=%s is " ++ "set: the large path FORCES a %s fault in that CTA of " ++ "every deterministic large row and returns an all -1 row. " ++ "This is a test mode and must not be set in " ++ "production.\n", ++ params.fault_inject_cta, ++ params.fault_inject_emit ? "emit" : "pre", ++ params.fault_inject_emit ? "emission-bound (X4)" ++ : "pre-emission (X3)"); + } + } + params.det_ties = det_ties_enabled(); +diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp +index 53f4577..3b23bc1 100644 +--- a/csrc/libtorch_stable/torch_bindings.cpp ++++ b/csrc/libtorch_stable/torch_bindings.cpp +@@ -867,8 +867,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); + ops.impl("persistent_topk_exact_counters", + TORCH_BOX(&persistent_topk_exact_counters)); +- ops.impl("persistent_topk_exact_counters_slots", +- TORCH_BOX(&persistent_topk_exact_counters_slots)); + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); + #endif +@@ -945,6 +943,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { + ops.impl("cutlass_scaled_mm_supports_fp4", + TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); + #endif ++ // Tensor-free: the counter layout the BUILD carries. Registered here, not ++ // in the CUDA block -- with no tensor argument there is no dispatch key, so ++ // a CUDA-only registration can never be called (reference host: ++ // NotImplementedError "no fallback function is registered"). ++ ops.impl("persistent_topk_exact_counters_slots", ++ TORCH_BOX(&persistent_topk_exact_counters_slots)); + } + + // Cache ops +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index fa6ee84..2b45b3b 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -2423,11 +2423,15 @@ LARGEDET_KEPT = LARGEDET_TOP_K - LARGEDET_ABOVE # 312 of the 700 survive + LARGEDET_DROPPED = LARGEDET_TIE - LARGEDET_KEPT # 388 are cut + + +-def _build_largedet_tie_row(num_rows: int) -> tuple[torch.Tensor, torch.Tensor]: ++def _build_largedet_tie_row( ++ num_rows: int, tie_recent: bool | None = None ++) -> tuple[torch.Tensor, torch.Tensor]: + """Identical rows with a deliberate multi-CTA tie at the cutoff. + + Returns (logits, expected_indices) where expected_indices is the ascending +- int64 index list the deterministic kernel must return for every row. ++ int64 index list the deterministic kernel must return for every row, in the ++ tie direction this launch actually runs (an earlier oracle hard-coded ++ most-recent-wins). + """ + gen = torch.Generator(device="cpu").manual_seed(19) + perm = torch.randperm(LARGEDET_SEQ_LEN, generator=gen, device="cpu") +@@ -2445,7 +2449,10 @@ def _build_largedet_tie_row(num_rows: int) -> tuple[torch.Tensor, torch.Tensor]: + + logits = row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() + +- kept_tie = torch.sort(tie_idx).values[-LARGEDET_KEPT:] ++ if tie_recent is None: ++ tie_recent = _largedet_effective_tie_recent() ++ sorted_tie = torch.sort(tie_idx).values ++ kept_tie = sorted_tie[-LARGEDET_KEPT:] if tie_recent else sorted_tie[:LARGEDET_KEPT] + expected = torch.sort(torch.cat([above_idx, kept_tie])).values.to(torch.int64) + assert expected.numel() == LARGEDET_TOP_K + return logits, expected +@@ -2488,7 +2495,7 @@ def _largedet_counters_or_skip(require_on: bool) -> list[int]: + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") + @torch.inference_mode() + def test_persistent_topk_large_tie_keeps_highest_indices() -> None: +- """Multi-CTA tie at the cutoff: most-recent-wins, ascending, reproducible.""" ++ """Multi-CTA tie at the cutoff: the effective direction, ascending, reproducible.""" + torch.set_default_device("cuda:0") + num_rows = 8 # <= 32 keeps the dispatch on persistent_topk + +@@ -2530,9 +2537,12 @@ def test_persistent_topk_large_tie_keeps_highest_indices() -> None: + ) + + # 3. The set is the analytic one: everything above the pivot, plus the +- # 312 HIGHEST-indexed of the 700 pivot-equal elements. ++ # 312 highest- or lowest-indexed of the 700 pivot-equal elements, ++ # whichever the effective direction says. + assert torch.equal(got.cpu(), expected), ( +- f"row {row}: selected set is not most-recent-wins at the cutoff" ++ f"row {row}: selected set is not the " ++ f"{'most-recent' if _largedet_effective_tie_recent() else 'oldest'}" ++ "-wins set at the cutoff" + ) + + # 4. The weaker check torch.topk can make: same value multiset. +@@ -2638,6 +2648,12 @@ def test_persistent_topk_large_fault_in_nonzero_cta_fails_closed() -> None: + "pass; run this test in its own process with the variable set to a " + "NONZERO CTA index." + ) ++ if os.getenv("MERLIN_TOPK_LARGE_FAULT_KIND", "pre").strip() == "emit": ++ pytest.skip( ++ "this is the PRE-emission (X3) fault test; the process has " ++ "MERLIN_TOPK_LARGE_FAULT_KIND=emit, which the emission-bound test " ++ "covers. Run this one in a process with KIND unset or 'pre'." ++ ) + assert int(inject) > 0, ( + "this test exists to prove a fault originating outside CTA 0 fails the " + f"row closed; MERLIN_TOPK_LARGE_FAULT_INJECT={inject} is CTA 0" +@@ -2677,6 +2693,96 @@ def test_persistent_topk_large_fault_in_nonzero_cta_fails_closed() -> None: + ) + + ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_emission_fault_fails_closed() -> None: ++ """The EMISSION-BOUND (X4) fault arm, from CTA 0 and from a nonzero CTA. ++ ++ `suffix_sum[6]` -- the per-CTA emission-violation flag -- ++ was written with a plain `suffix_sum[6] = 1u`. Every thread whose computed ++ slot is out of range reaches that line, and an invariant failure is exactly ++ the situation in which many of them do at once, so the mechanism that ++ exists to DETECT the failure contained a concurrent non-atomic write of its ++ own. It is now `atomicOr(&suffix_sum[6], 1u)`. ++ ++ The shipped injector could not reach this code at all. It forces the ++ PRE-EMISSION predicate, which makes `row_fault` nonzero; the whole emission ++ block is then skipped and `suffix_sum[6]` is never touched. So the X4 arm ++ had never been exercised, and the existing all -1 result proved only the ++ X3 arm. `MERLIN_TOPK_LARGE_FAULT_KIND=emit` selects the emission-bound ++ predicate instead, and forces it for every emitting thread of the named ++ CTA -- a single-writer injection would pass against the broken code and so ++ would measure nothing. ++ ++ Both values of the CTA index are in scope here, unlike the X3 cell: CTA 0 ++ is the CTA that performs the -1 fill, so "the fill happens even when the ++ discovering CTA is the filling CTA" and "…even when it is not" are two ++ different claims. Run this in its own process, once per CTA index -- the ++ injector is a function-local static baked in at the first launch: ++ ++ MERLIN_TOPK_DET_LARGE=1 MERLIN_TOPK_LARGE_FAULT_INJECT=0 \ ++ MERLIN_TOPK_LARGE_FAULT_KIND=emit pytest -k large_emission_fault ++ MERLIN_TOPK_DET_LARGE=1 MERLIN_TOPK_LARGE_FAULT_INJECT=1 \ ++ MERLIN_TOPK_LARGE_FAULT_KIND=emit pytest -k large_emission_fault ++ ++ Injection suppresses __trap(), so this is the trap-off arm of the policy. ++ The trap-on arm kills the CUDA context by construction and stays in a disposable ++ process (test_isolated_diag_trap_on_kills_only_the_child), not here. ++ """ ++ torch.set_default_device("cuda:0") ++ num_rows = 8 ++ ++ inject = os.getenv("MERLIN_TOPK_LARGE_FAULT_INJECT") ++ kind = (os.getenv("MERLIN_TOPK_LARGE_FAULT_KIND") or "").strip() ++ if inject is None or inject.strip() in ("", "-1") or kind != "emit": ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build with the fault " ++ "cells unrun: the emission-bound arm needs " ++ "MERLIN_TOPK_LARGE_FAULT_INJECT= AND " ++ "MERLIN_TOPK_LARGE_FAULT_KIND=emit. Without the kind the injector " ++ "forces the PRE-emission fault, which skips the emission block " ++ "entirely and cannot exercise suffix_sum[6]. A skip is not a pass." ++ ) ++ assert int(inject) >= 0 ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ logits, _ = _build_largedet_tie_row(num_rows) ++ out = _run_largedet_row(logits) ++ after = _read_exact_counters() ++ ++ # 1. The whole row is -1. The X4 arm's specific hazard is a PARTIAL row: ++ # the CTAs that did not fault have already stored real indices by the ++ # time the violation is discovered, so "contains -1" would pass on the ++ # defect. The fill must overwrite every slot. ++ assert int(out.min()) == -1 and int(out.max()) == -1, ( ++ "an emission-bound fault did not fill the row with -1: " ++ f"min={int(out.min())} max={int(out.max())}; " ++ f"{int((out != -1).sum())} of {out.numel()} slots survived" ++ ) ++ ++ # 2. EXACTLY one fault per row. Not one per CTA and -- the reason the ++ # injector fires in every emitting thread -- not one per out-of-range ++ # element. ++ delta = after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] ++ assert delta == num_rows, ( ++ "the large-path fault counter did not increment exactly once per " ++ f"failed row: {delta} of {num_rows}" ++ ) ++ ++ # 3. The rows were still entered, and no tie measurement was recorded off ++ # a faulted row. ++ assert after[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_rows, ( ++ "faulted rows must still be counted as large-path rows" ++ ) ++ assert after[_SLOT_LARGE_TIES_CUT] - before[_SLOT_LARGE_TIES_CUT] == 0, ( ++ "a faulted row recorded a tie-cut measurement" ++ ) ++ assert after[_SLOT_LARGE_TIES_DROPPED] - before[_SLOT_LARGE_TIES_DROPPED] == 0, ( ++ "a faulted row recorded dropped-tie counts" ++ ) ++ ++ + # --------------------------------------------------------------------------- + # Large-path determinism -- expanded coverage (more shapes and + # adversarial distributions). +@@ -2707,27 +2813,120 @@ def test_persistent_topk_large_fault_in_nonzero_cta_fails_closed() -> None: + LARGEDET_RADIX_THRESHOLD = 32768 # persistent_topk.cuh: above this, the large path + + +-def _largedet_reference(values: torch.Tensor, top_k: int) -> torch.Tensor: ++# --------------------------------------------------------------------------- ++# The EFFECTIVE policy, read from the same place the kernel ++# reads it. ++# ++# Every switch in topk.cu is a function-local static filled from std::getenv on ++# the first launch, so the environment this process started with IS the policy ++# the kernel applies -- provided nothing mutates os.environ after a top-k ++# kernel has run. Do not monkeypatch these inside a session; run each arm in ++# its own process, exactly as the fault cells do. ++# --------------------------------------------------------------------------- ++def _largedet_env_on(name: str, default: bool) -> bool: ++ raw = os.getenv(name) ++ if raw is None or raw.strip() == "": ++ return default ++ try: ++ return int(raw) != 0 ++ except ValueError: ++ return default ++ ++ ++def _largedet_effective_tie_recent() -> bool: ++ """`effective_tie_recent(det_ties, tie_recent)` = det_ties ? tie_recent : true. ++ ++ With MERLIN_TIE_DET=0 (the shipped default) the buffered paths keep their ++ racy atomic rule, which has no direction, and every deterministic emission ++ -- including the large path -- uses most-recent-wins. ++ """ ++ if not _largedet_env_on("MERLIN_TIE_DET", False): ++ return True ++ return _largedet_env_on("MERLIN_TIE_RECENT", False) ++ ++ ++def _buffered_order_is_canonical() -> bool: ++ """Do the DECODE and MEDIUM paths promise a deterministic SELECTED SET? ++ ++ Only with MERLIN_TIE_DET=1, and only the set: TIE_DET makes cutoff-tie ++ MEMBERSHIP deterministic, but the buffered paths still place every ++ selected element by atomicAdd arrival, so the RAW output order stays ++ launch-dependent. Canonical order is MERLIN_QSA_SORT's job (a Python-side ++ sort in qsa_indexer), not the kernel's -- the root-cause analysis split the ++ two mechanisms exactly this way. On the reference host, asserting raw equality ++ here failed on 3 short rows in both tie directions while values, ++ membership and tie policy all passed. Callers compare sorted indices. ++ The LARGE path is separate: it is canonical in raw order whenever ++ MERLIN_TOPK_DET_LARGE is on. ++ """ ++ return _largedet_env_on("MERLIN_TIE_DET", False) ++ ++ ++def _largedet_reference_impl( ++ values: torch.Tensor, top_k: int, tie_recent: bool | None = None ++) -> torch.Tensor: + """The deterministic large path's own rule, as ascending int64 indices. + +- Everything strictly above the k-th largest value, plus the HIGHEST-indexed +- of the elements equal to it (most-recent-wins). Computed on CPU float32. +- float32 equality is exactly the kernel's ordered-uint32 equality, because +- the transform is monotonic and injective, so this reference is exact. ++ Everything strictly above the k-th largest value, plus -- in the direction ++ the launch actually runs -- the HIGHEST-indexed (most-recent-wins) or the ++ LOWEST-indexed (oldest-wins) of the elements equal to it. Computed on CPU ++ float32; float32 equality is exactly the kernel's ordered-uint32 equality, ++ because the transform is monotonic and injective, so this is exact. ++ ++ This oracle used to hard-code most-recent-wins, so a ++ CORRECT oldest-wins run (MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=0) failed every ++ distribution cell, and the oldest-wins behaviour the integration added had ++ no oracle at all. `tie_recent=None` resolves it from the environment the ++ kernel resolves it from; pass an explicit value only to test the oracle. + """ + assert values.dtype == torch.float32 and values.dim() == 1 ++ if tie_recent is None: ++ tie_recent = _largedet_effective_tie_recent() + pivot = torch.sort(values, descending=True).values[top_k - 1] + above = (values > pivot).nonzero(as_tuple=False).flatten() + equal = (values == pivot).nonzero(as_tuple=False).flatten() + need = top_k - int(above.numel()) + assert need >= 0, "more elements above the pivot than k; reference is wrong" +- kept = equal[equal.numel() - need :] if need > 0 else equal[:0] ++ if need == 0: ++ kept = equal[:0] ++ elif tie_recent: ++ kept = equal[equal.numel() - need :] ++ else: ++ kept = equal[:need] + expected = torch.sort(torch.cat([above, kept])).values.to(torch.int64) + assert expected.numel() == top_k + return expected + + +-def _largedet_distribution(kind: str, seq_len: int, top_k: int, seed: int) -> torch.Tensor: ++def _assert_valid_topk_membership( ++ values: torch.Tensor, got: torch.Tensor, top_k: int, where: str ++) -> None: ++ """Policy-CORRECT membership, for rows whose tie choice is not canonical. ++ ++ Where the output order is not promised, the set still is, ++ up to the choice WITHIN the pivot-equal group. Asserted: exactly `top_k` ++ distinct in-range indices; every selected value >= the pivot; and every ++ element strictly above the pivot selected. That leaves free exactly the one ++ degree of freedom the buffered paths have with MERLIN_TIE_DET=0, and ++ nothing else. ++ """ ++ assert got.numel() == top_k, f"{where}: {got.numel()} indices, want {top_k}" ++ assert int(got.min()) >= 0 and int(got.max()) < values.numel(), ( ++ f"{where}: index out of range" ++ ) ++ assert got.unique().numel() == top_k, f"{where}: duplicate index" ++ pivot = torch.sort(values, descending=True).values[top_k - 1] ++ assert bool((values[got] >= pivot).all()), ( ++ f"{where}: a selected element is below the k-th largest value" ++ ) ++ above = (values > pivot).nonzero(as_tuple=False).flatten() ++ missing = int((~torch.isin(above, got)).sum()) ++ assert missing == 0, ( ++ f"{where}: {missing} elements strictly above the pivot were not selected" ++ ) ++ ++ ++def _largedet_distribution_impl(kind: str, seq_len: int, top_k: int, seed: int) -> torch.Tensor: + """One CPU float32 row. All values are positive, so no -0.0/+0.0 hazard.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + +@@ -2778,6 +2977,24 @@ def _largedet_distribution(kind: str, seq_len: int, top_k: int, seed: int) -> to + raise ValueError(f"unknown large-path distribution: {kind}") + + ++def _largedet_reference(*args, **kwargs) -> torch.Tensor: ++ """CPU reference, whatever the global default device is. ++ ++ Earlier tests in this file call torch.set_default_device("cuda:0") and ++ never reset it, so factory calls inside the helpers would otherwise land ++ on the GPU and fail the torch.equal against a .cpu() kernel output ++ (seen on the reference host) -- a harness failure that would read as a kernel one. ++ """ ++ with torch.device("cpu"): ++ return _largedet_reference_impl(*args, **kwargs).cpu() ++ ++ ++def _largedet_distribution(*args, **kwargs) -> torch.Tensor: ++ """CPU fixture row, whatever the global default device is (see above).""" ++ with torch.device("cpu"): ++ return _largedet_distribution_impl(*args, **kwargs).cpu() ++ ++ + def _largedet_launch(logits: torch.Tensor, lengths: torch.Tensor, top_k: int): + indices = torch.empty((logits.shape[0], top_k), dtype=torch.int32, device="cuda") + _run_topk_backend( +@@ -2840,7 +3057,9 @@ def test_persistent_topk_large_distributions( + ) + assert got.unique().numel() == top_k, f"{kind} row {row}: duplicate index" + assert torch.equal(got.cpu(), expected), ( +- f"{kind} row {row}: selected set is not the most-recent-wins set" ++ f"{kind} row {row}: selected set is not the " ++ f"{'most-recent' if _largedet_effective_tie_recent() else 'oldest'}" ++ "-wins set" + ) + assert torch.equal(logits[row][got].sort().values, ref_values), ( + f"{kind} row {row}: selected values are not torch.topk's values" +@@ -2862,6 +3081,18 @@ def test_persistent_topk_large_ragged_mixed_paths(top_k: int) -> None: + across interleaved traffic. Large rows get the exact set assertion; short + rows only get the value-multiset check, because their tie policy belongs to + the decode/medium paths, not to this branch. ++ ++ The raw `torch.equal(first, second)` over the WHOLE ++ tensor is gone. The decode and medium rows in this launch emit by atomicAdd ++ arrival order unless MERLIN_TIE_DET is on, so a legitimate buffered-order ++ difference on a short row used to fail a test presented as large-path ++ qualification -- a misattribution, and on the other side a flake. Raw ++ equality is now asserted per row: unconditionally on the large rows, whose ++ canonical ascending order DET_LARGE promises, and on the short rows only ++ when TIE_DET makes that promise for them too. Where it is not promised, the ++ short rows are still held to the value multiset and to policy-correct ++ membership (a valid top-k SET: no duplicates, in range, every selected ++ value >= the pivot and every strictly-greater element selected). + """ + torch.set_default_device("cuda:0") + lengths_list = [ +@@ -2895,6 +3126,7 @@ def test_persistent_topk_large_ragged_mixed_paths(top_k: int) -> None: + ) + assert mid[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0 + ++ canonical_short = _buffered_order_is_canonical() + for row, length in enumerate(lengths_list): + got = first[row].to(torch.int64) + prefix = cpu_row[:length] +@@ -2903,14 +3135,29 @@ def test_persistent_topk_large_ragged_mixed_paths(top_k: int) -> None: + assert torch.equal(got_values, ref_values), ( + f"row {row} (length {length}): values are not torch.topk's values" + ) ++ _assert_valid_topk_membership(prefix, got.cpu(), top_k, f"row {row}") + if length > LARGEDET_RADIX_THRESHOLD: ++ assert bool((got[1:] - got[:-1] > 0).all()), ( ++ f"row {row} (length {length}): large-path output is not " ++ "ascending in index" ++ ) + assert torch.equal(got.cpu(), _largedet_reference(prefix, top_k)), ( +- f"row {row} (length {length}): not the most-recent-wins set" ++ f"row {row} (length {length}): not the " ++ f"{'most-recent' if _largedet_effective_tie_recent() else 'oldest'}" ++ "-wins set" ++ ) ++ assert torch.equal(first[row], second[row]), ( ++ f"row {row} (length {length}): the LARGE path returned a " ++ "different index tensor on an identical relaunch" ++ ) ++ elif canonical_short: ++ assert torch.equal( ++ first[row].sort().values, second[row].sort().values ++ ), ( ++ f"row {row} (length {length}): MERLIN_TIE_DET=1 promises a " ++ "deterministic selected set, but two identical launches " ++ "selected different indices" + ) +- +- assert torch.equal(first, second), ( +- "two identical ragged launches returned different index tensors" +- ) + + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@@ -2965,3 +3212,121 @@ def test_persistent_topk_large_slot_residue_reuse() -> None: + assert torch.equal(first, second), ( + "two identical launches over reused histogram slots differed" + ) ++ ++ ++# Large-path reuse geometry, computed for the reference 3090 and asserted before use. ++# ++# topk.cu: ctas_per_group = ceil(stride / max_chunk_elements); at >8 rows and ++# 262,144 columns the chunk budget gives ~24,824 elements, so ++# ctas_per_group = 11. On a 3090 (82 SMs; occupancy 1 for a 1,024-thread CTA at ++# this shared-memory size) the cooperative headroom rule leaves ++# max_resident_ctas = 81, hence num_groups = min(81 // 11, num_rows) = 7. ++# ++# persistent_topk.cuh assigns rows `row_idx = group_id + iter * num_groups`, so ++# with num_groups = 7 the ITERATION index of a row is `row_idx // 7` and its ++# GROUP is `row_idx % 7`. Laying the classes out in blocks of 7 therefore gives ++# EVERY group the same five-iteration sequence, and the sequence is the one ++# the coverage needs. ++LARGEDET_REUSE_GROUPS = 7 ++LARGEDET_REUSE_SMS = 82 # RTX 3090; the geometry above is derived from this ++LARGEDET_REUSE_LARGE_LEN = 262144 ++LARGEDET_REUSE_CLASSES = ["large", "decode", "medium", "large", "large"] ++LARGEDET_REUSE_LENGTHS = {"large": LARGEDET_REUSE_LARGE_LEN, "decode": 2048, "medium": 12288} ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@torch.inference_mode() ++def test_persistent_topk_large_within_group_mixed_reuse() -> None: ++ """One group, in order: large, decode, medium, large, large. ++ ++ The eight-row mixed test gives each group only ONE row on ++ this lane, so it shows that a ragged launch returns correct rows -- not ++ that a group which handled a short row BETWEEN two large ones still reuses ++ `RadixRowState::histogram[radix_iter % 3]` correctly. Those are different ++ claims. `radix_iter` only advances on large rows, so the third large row ++ here lands on the slot the first one used, with a decode row and a medium ++ row having run through the same group in between; that is the geometry the ++ X1 publication barrier and the slot-residue argument are actually about. ++ ++ Three large rows per group, not two, so all three histogram residues are ++ covered, and the short rows sit between them rather than after them. ++ ++ The geometry is asserted, not hoped for: if the device is not the 82-SM ++ 3090 the constants above were derived for, the test SKIPS rather than ++ quietly degrading to one large row per group and reporting a pass. ++ """ ++ torch.set_default_device("cuda:0") ++ ++ sms = torch.cuda.get_device_properties(0).multi_processor_count ++ if sms != LARGEDET_REUSE_SMS: ++ pytest.skip( ++ f"this cell's group geometry is derived for a {LARGEDET_REUSE_SMS}-SM " ++ f"RTX 3090; this device has {sms} SMs, so num_groups would not be " ++ f"{LARGEDET_REUSE_GROUPS} and the within-group interleaving would not " ++ "be established. Re-derive the constants before claiming this " ++ "coverage on other hardware -- a skip is not a pass." ++ ) ++ ++ classes = [c for c in LARGEDET_REUSE_CLASSES for _ in range(LARGEDET_REUSE_GROUPS)] ++ lengths_list = [LARGEDET_REUSE_LENGTHS[c] for c in classes] ++ num_rows = len(lengths_list) ++ seq_len = LARGEDET_REUSE_LARGE_LEN ++ top_k = 512 ++ num_large = sum(1 for n in lengths_list if n > LARGEDET_RADIX_THRESHOLD) ++ assert num_large == 3 * LARGEDET_REUSE_GROUPS ++ ++ before = _largedet_counters_or_skip(require_on=True) ++ ++ cpu_row = _largedet_distribution("cross_cta_ties", seq_len, top_k, seed=1907) ++ logits = cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ lengths = torch.tensor(lengths_list, dtype=torch.int32, device="cuda") ++ ++ first = _largedet_launch(logits, lengths, top_k).clone() ++ mid = _read_exact_counters() ++ second = _largedet_launch(logits, lengths, top_k) ++ ++ assert mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS] == num_large, ( ++ "the large path ran on the wrong number of rows: " ++ f"{mid[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS]} of {num_large}" ++ ) ++ assert mid[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT] == 0, ( ++ "a large-path invariant fired across the interleaved reuse" ++ ) ++ ++ canonical_short = _buffered_order_is_canonical() ++ expected_large = _largedet_reference(cpu_row, top_k) ++ ++ for row, length in enumerate(lengths_list): ++ got = first[row].to(torch.int64) ++ prefix = cpu_row[:length] ++ ref_values = torch.topk(prefix, top_k).values.sort().values ++ assert torch.equal(prefix[got.cpu()].sort().values, ref_values), ( ++ f"row {row} (length {length}, iter {row // LARGEDET_REUSE_GROUPS}, " ++ f"group {row % LARGEDET_REUSE_GROUPS}): values are not torch.topk's" ++ ) ++ _assert_valid_topk_membership( ++ prefix, got.cpu(), top_k, ++ f"row {row} (length {length})", ++ ) ++ if length > LARGEDET_RADIX_THRESHOLD: ++ assert bool((got[1:] - got[:-1] > 0).all()), ( ++ f"row {row} (iter {row // LARGEDET_REUSE_GROUPS}): large-path " ++ "output is not ascending in index" ++ ) ++ assert torch.equal(got.cpu(), expected_large), ( ++ f"row {row} (iter {row // LARGEDET_REUSE_GROUPS}, group " ++ f"{row % LARGEDET_REUSE_GROUPS}): the selected set changed across " ++ "histogram-slot reuse with short rows interleaved" ++ ) ++ assert torch.equal(first[row], second[row]), ( ++ f"row {row}: the LARGE path was not reproducible across the " ++ "interleaved reuse" ++ ) ++ elif canonical_short: ++ assert torch.equal( ++ first[row].sort().values, second[row].sort().values ++ ), ( ++ f"row {row} (length {length}): MERLIN_TIE_DET=1 promises a " ++ "deterministic selected set, but two identical launches " ++ "selected different indices" ++ ) +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 4e8513e..11700fa 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -675,6 +675,19 @@ _TIEDET_KERNEL_SWITCHES = [ + ("MERLIN_TOPK_DET_LARGE", "1", "1"), + ] + ++# The large-path fault injectors. TEST-ONLY knobs, but ++# both are capture-baked -- the values become kernel launch arguments and a ++# cudagraph captured with the injector armed returns an all -1 row on every ++# replay, with nothing re-reading the environment to notice it was unset since. ++# So they are registered and hashed on the same conservative rule, and are kept ++# in their own table because they are not switches: unset means "no injection", ++# and a default-off assertion on them would be testing the absence of a test ++# mode rather than a shipped behaviour. ++_LARGEDET_FAULT_INJECTORS = [ ++ ("MERLIN_TOPK_LARGE_FAULT_INJECT", "1", "1"), ++ ("MERLIN_TOPK_LARGE_FAULT_KIND", "emit", "emit"), ++] ++ + # Named in the envs.py fenced block, deliberately NOT registered. + _TIEDET_UNAUDITED = ["MERLIN_PIN_BSM"] + +@@ -721,7 +734,8 @@ def _envs_source() -> str: + + + @pytest.mark.parametrize( +- ("name", "value", "expected"), _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES ++ ("name", "value", "expected"), ++ _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES + _LARGEDET_FAULT_INJECTORS, + ) + def test_tiedet_toggle_is_a_cache_factor( + name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch +@@ -760,7 +774,8 @@ def test_tiedet_toggle_defaults_off( + + @pytest.mark.parametrize( + "name", +- [n for n, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES] + _TIEDET_UNAUDITED, ++ [n for n, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES + _LARGEDET_FAULT_INJECTORS] ++ + _TIEDET_UNAUDITED, + ) + def test_unaudited_host_side_knob_stays_listed_until_it_is_classified(name: str): + """The fenced block in envs.py is a gate, not a note. +@@ -792,6 +807,31 @@ def test_unaudited_host_side_knob_stays_listed_until_it_is_classified(name: str) + ) + + ++@pytest.mark.parametrize(("name", "value", "expected"), _LARGEDET_FAULT_INJECTORS) ++def test_largedet_fault_injector_changes_the_cache_key( ++ name: str, value: str, expected, monkeypatch: pytest.MonkeyPatch ++): ++ """Arming an injector must move the compile factors, not just appear in them. ++ ++ MERLIN_TOPK_LARGE_FAULT_INJECT was once unregistered and ++ capture-baked: a cudagraph captured with it armed keeps returning all -1 ++ rows on replay whatever the environment says afterwards. Registration is ++ only worth something if the hashed value FOLLOWS the environment, so assert ++ both worlds -- unset reads None and absent-from-armed, armed reads the raw ++ string -- rather than membership alone, which a dead getter would satisfy. ++ """ ++ monkeypatch.delenv(name, raising=False) ++ assert _fresh_env_value(name) is None ++ off = envs.compile_factors() ++ monkeypatch.setenv(name, value) ++ assert _fresh_env_value(name) == expected ++ on = envs.compile_factors() ++ assert off[name] != on[name], ( ++ f"{name} is registered but its hashed value does not follow the " ++ "environment, so an injected build and a clean build share a cache key" ++ ) ++ ++ + def test_unaudited_gate_discriminates_both_worlds(): + """Control for the gate above: it must be able to FAIL. + +@@ -860,7 +900,7 @@ def test_tiedet_toggles_are_not_ignored_factors(): + ignored_factors would keep every test green while dropping it from the + cache key.""" + monkeypatch_free_factors = envs.compile_factors() +- for name, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES: ++ for name, _, _ in _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES + _LARGEDET_FAULT_INJECTORS: + assert name in environment_variables, name + assert name in monkeypatch_free_factors, name + for name in _TIEDET_UNAUDITED: +diff --git a/vllm/envs.py b/vllm/envs.py +index 960d92b..b604b39 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -563,6 +563,35 @@ if TYPE_CHECKING: + # bool that would imply this tree can act on it. + MERLIN_TOPK_DET_LARGE: str | None = None + # ++ # The large-path FAULT INJECTORS. TEST ++ # ONLY, and registered for exactly the same reason as the switches above, ++ # not because anything should set them in production: ++ # ++ # * both are resolved once per process by a std::getenv in topk.cu ++ # (large_fault_inject_cta(), large_fault_inject_emit()) and the ++ # resolved values reach the GPU as kernel launch arguments, so a ++ # cudagraph BAKES THEM IN. An injected build and a clean build differ ++ # in the captured graph while sharing a cache key -- which is the ++ # failure mode: a graph captured under injection returns an all -1 row ++ # for every deterministic large row, and on replay nothing re-reads the ++ # environment to notice the variable has since been unset. FAULT_INJECT ++ # was once "unregistered and capture-baked"; hashing ++ # them is what stops a cached artifact crossing between the arms. ++ # * compile_factors() is a deny-list, so a name it does not know cannot ++ # be hashed at all. Neither is in ignored_factors. ++ # ++ # Raw passthroughs, like the switches: the C++ getters are the single ++ # source of truth (INJECT parses with atoi and treats <0 and unset alike; ++ # KIND compares against the literal "emit"). Do not parse them here. ++ # ++ # Which CTA of the group has its fault predicate forced; unset / negative ++ # is off, and setting it also suppresses __trap(). ++ MERLIN_TOPK_LARGE_FAULT_INJECT: str | None = None ++ # Which fault: unset / "pre" forces the pre-emission predicate caught at ++ # X3, "emit" forces the emission-bound slot check caught at X4. Ignored ++ # unless MERLIN_TOPK_LARGE_FAULT_INJECT names a CTA. ++ MERLIN_TOPK_LARGE_FAULT_KIND: str | None = None ++ # + # NOT REGISTERED, deliberately: MERLIN_TOPK_EXACT_TRAP. It is a + # preprocessor define consumed by `#ifndef` in persistent_topk.cuh at + # BUILD time (-D), not a process environment variable -- nothing in the +@@ -2566,6 +2595,14 @@ environment_variables: dict[str, Callable[[], Any]] = { + # MERLIN_TOPK_EXACT_MEDIUM is registered once, earlier in this dict. + "MERLIN_TOPK_EXACT_DECODE": lambda: os.getenv("MERLIN_TOPK_EXACT_DECODE", None), + "MERLIN_TOPK_DET_LARGE": lambda: os.getenv("MERLIN_TOPK_DET_LARGE", None), ++ # Large-path fault injectors. TEST ONLY, hashed for ++ # the reason in the declaration block: both are capture-baked. ++ "MERLIN_TOPK_LARGE_FAULT_INJECT": lambda: os.getenv( ++ "MERLIN_TOPK_LARGE_FAULT_INJECT", None ++ ), ++ "MERLIN_TOPK_LARGE_FAULT_KIND": lambda: os.getenv( ++ "MERLIN_TOPK_LARGE_FAULT_KIND", None ++ ), + } + + +-- +2.43.0 + diff --git a/patches/0046-topk-diagnostic-only-build-defines-and-qualification.patch b/patches/0046-topk-diagnostic-only-build-defines-and-qualification.patch new file mode 100644 index 0000000..9e0524d --- /dev/null +++ b/patches/0046-topk-diagnostic-only-build-defines-and-qualification.patch @@ -0,0 +1,1378 @@ +From ba83d8036e0c5b169d309d286a6bf1cda9b7fd65 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:35:52 -0400 +Subject: [PATCH 046/103] topk: diagnostic-only build defines and qualification + test cells + +Compile-time defines (default 0, device code unchanged by default) for the +delayed-publication stress, the pre-repair control and a forced-fault cell; +test cells for the OFF-arm controls, the QSA_SORT contract, FULL_K attribution +on the Marlin MoE GEMM and legacy-tail repeated-row membership. +--- + csrc/libtorch_stable/persistent_topk.cuh | 85 +++ + csrc/libtorch_stable/topk.cu | 25 + + tests/kernels/moe/merlin_full_k_child.py | 128 +++++ + tests/kernels/moe/test_merlin_full_k.py | 188 +++++++ + tests/kernels/test_merlin_qsa_sort.py | 161 ++++++ + tests/kernels/test_top_k_per_row.py | 689 +++++++++++++++++++++++ + 6 files changed, 1276 insertions(+) + create mode 100644 tests/kernels/moe/merlin_full_k_child.py + create mode 100644 tests/kernels/moe/test_merlin_full_k.py + create mode 100644 tests/kernels/test_merlin_qsa_sort.py + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 85d217f..749d2d9 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -685,6 +685,74 @@ __device__ __noinline__ bool histogram_2048_topk( + #define MERLIN_TOPK_EXACT_TRAP 1 + #endif + ++// ============================================================================ ++// DIAGNOSTIC BUILDS ONLY -- never set in a release or serving build. ++// ++// Three compile-time defines exist so the reference host can run the two large-path ++// diagnostic cells. Every one defaults to 0 and every use ++// sits behind `#if`, so the default build's device code is identical to the ++// tree without them (offline, real TU, sm_86, against an earlier commit: SASS identical, ++// cubin identical after masking nvcc's per-compile internal-linkage TU hash, ++// which differs even between two compiles of the SAME source -- see the ++// diagnostic-build notes). A build carrying any of them prints a ++// `[merlin] persistent_topk DIAGNOSTIC BUILD` line on its first launch ++// (topk.cu), which the smoke's startup-warning scan must reject. ++// ++// MERLIN_TOPK_DIAG_DELAY= ++// In `radix_topk`'s round loop, every NON-LEADER warp (warp != 0) of ++// every NON-LEADER CTA (cta_in_group != 0) spins clock64 ticks ++// before merging its shared bins into the group's global histogram. ++// Warp 0 holds thread 0, the arrival publisher, so it is never delayed: ++// the delay opens exactly the window the histogram-publication barrier closes -- a ++// publisher that announces arrival before its own CTA's other warps have ++// produced. CTA 0 is never delayed, so it reads the histogram as soon ++// as the arrival count completes. ++// MERLIN_TOPK_DIAG_PREREPAIR=1 ++// Compile out the histogram-publication barrier. This is the ++// PRE-REPAIR control for the delayed-publication stress: with a delay it ++// is expected to return wrong or faulted rows; the repaired build with ++// the same delay must not. Only meaningful with MERLIN_TOPK_DIAG_DELAY. ++// Do NOT run it with MERLIN_TOPK_DET_LARGE=0: the legacy tail bounds no ++// emission slot, so a disagreeing pivot writes out of bounds there. ++// MERLIN_TOPK_DIAG_FORCE_FAULT=1 (+ MERLIN_TOPK_DIAG_FORCE_FAULT_CTA, 1) ++// Force the CTA-local X3 predicate in that CTA on every deterministic ++// large row, exactly as MERLIN_TOPK_LARGE_FAULT_INJECT does -- but WITHOUT ++// suppressing `__trap()`. The trap-ON fault arm: the process's CUDA ++// context dies. Run only in a disposable subprocess. ++// ============================================================================ ++#ifndef MERLIN_TOPK_DIAG_DELAY ++ #define MERLIN_TOPK_DIAG_DELAY 0 ++#endif ++#ifndef MERLIN_TOPK_DIAG_PREREPAIR ++ #define MERLIN_TOPK_DIAG_PREREPAIR 0 ++#endif ++#ifndef MERLIN_TOPK_DIAG_FORCE_FAULT ++ #define MERLIN_TOPK_DIAG_FORCE_FAULT 0 ++#endif ++#ifndef MERLIN_TOPK_DIAG_FORCE_FAULT_CTA ++ #define MERLIN_TOPK_DIAG_FORCE_FAULT_CTA 1 ++#endif ++#define MERLIN_TOPK_DIAG_ANY \ ++ ((MERLIN_TOPK_DIAG_DELAY) != 0 || (MERLIN_TOPK_DIAG_PREREPAIR) != 0 || \ ++ (MERLIN_TOPK_DIAG_FORCE_FAULT) != 0) ++#if MERLIN_TOPK_DIAG_FORCE_FAULT && !MERLIN_TOPK_EXACT_TRAP ++ #error "MERLIN_TOPK_DIAG_FORCE_FAULT is the trap-ON cell; it needs MERLIN_TOPK_EXACT_TRAP=1" ++#endif ++#if MERLIN_TOPK_DIAG_FORCE_FAULT && (MERLIN_TOPK_DIAG_DELAY || MERLIN_TOPK_DIAG_PREREPAIR) ++ #error "build the trap-ON and the delayed-publication diagnostics separately" ++#endif ++ ++#if MERLIN_TOPK_DIAG_DELAY ++// Busy-wait on the SM clock. `volatile` sink so the loop is not elided. ++__device__ __noinline__ void merlin_topk_diag_spin(long long cycles) { ++ const long long start = clock64(); ++ volatile long long now = start; ++ while (now - start < cycles) { ++ now = clock64(); ++ } ++} ++#endif ++ + // Exact, bounded-memory re-selection for a single medium-path row. + // + // `scores` is already offset to the row. Every radix round rescans the whole +@@ -1621,6 +1689,13 @@ __device__ __noinline__ void radix_emit_deterministic( + static_cast(cta_in_group) == fault_inject_cta) { + my_keep_checked = local_eq + 1u; + } ++#if MERLIN_TOPK_DIAG_FORCE_FAULT ++ // DIAGNOSTIC BUILD ONLY: the same CTA-local predicate, forced at compile ++ // time, so the trap below is NOT suppressed (fault_inject_cta stays -1). ++ if (cta_in_group == static_cast(MERLIN_TOPK_DIAG_FORCE_FAULT_CTA)) { ++ my_keep_checked = local_eq + 1u; ++ } ++#endif + if (my_keep_checked > local_eq) fault = 1u; + if (gt_base + eq_keep_base + local_gt + my_keep > + static_cast(TopK)) { +@@ -1954,6 +2029,14 @@ __device__ void radix_topk(const float* __restrict__ row_input, + } + __syncthreads(); + ++#if MERLIN_TOPK_DIAG_DELAY ++ // DIAGNOSTIC BUILD ONLY (see the define block): hold back the producers ++ // that are NOT the arrival publisher, in every CTA that is not the leader. ++ if (cta_in_group != 0u && (tx >> 5) != 0u) { ++ merlin_topk_diag_spin(static_cast(MERLIN_TOPK_DIAG_DELAY)); ++ } ++#endif ++ + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + if (local_histogram[i] > 0) { + atomicAdd(¤t_hist[i], local_histogram[i]); +@@ -1976,7 +2059,9 @@ __device__ void radix_topk(const float* __restrict__ row_input, + // nothing here -- and reads a half-written histogram, which changes the + // pivot. The barrier must be between the producing stores and the leader's + // publication; program order in one thread is not cross-thread ordering. ++#if !MERLIN_TOPK_DIAG_PREREPAIR + __syncthreads(); ++#endif + + if (tx == 0) { + red_release(&state->arrival_counter, 1); +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 56b0af4..b52a34b 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -201,6 +201,28 @@ void report_tie_policy_once() { + (void)reported; + } + ++#if MERLIN_TOPK_DIAG_ANY ++// DIAGNOSTIC BUILD ONLY (persistent_topk.cuh, "DIAGNOSTIC BUILDS ONLY"). Say ++// so once per process, on stderr, before the first launch: the diagnostic-build tests ++// key on this line to prove which build they ran against, and the smoke's ++// startup-warning scan must reject any process that prints it. ++void report_diag_build_once() { ++ static const bool reported = [] { ++ std::fprintf(stderr, ++ "[merlin] persistent_topk DIAGNOSTIC BUILD: " ++ "MERLIN_TOPK_DIAG_DELAY=%lld MERLIN_TOPK_DIAG_PREREPAIR=%d " ++ "MERLIN_TOPK_DIAG_FORCE_FAULT=%d " ++ "MERLIN_TOPK_DIAG_FORCE_FAULT_CTA=%d -- NOT FOR SERVING\n", ++ static_cast(MERLIN_TOPK_DIAG_DELAY), ++ static_cast(MERLIN_TOPK_DIAG_PREREPAIR), ++ static_cast(MERLIN_TOPK_DIAG_FORCE_FAULT), ++ static_cast(MERLIN_TOPK_DIAG_FORCE_FAULT_CTA)); ++ return true; ++ }(); ++ (void)reported; ++} ++#endif ++ + // Process-lifetime device counters, one block per device. The kernel only + // ever atomicAdds into them, so the launch path needs no host sync and the + // values accumulate across cudagraph replays. +@@ -552,6 +574,9 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + params.tie_recent = tie_recent_enabled(); + params.exact_decode = exact_decode_enabled(); + report_tie_policy_once(); ++#if MERLIN_TOPK_DIAG_ANY ++ report_diag_build_once(); ++#endif + // Allocate only off the capture path; in the ordinary engine lifecycle the + // eager warmup lands here first, long before capture. The explicit, + // quiescent-phase entry point is the reader op. +diff --git a/tests/kernels/moe/merlin_full_k_child.py b/tests/kernels/moe/merlin_full_k_child.py +new file mode 100644 +index 0000000..3a42008 +--- /dev/null ++++ b/tests/kernels/moe/merlin_full_k_child.py +@@ -0,0 +1,128 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Child process for the MERLIN_FULL_K attribution cell. ++ ++NOT a test module (no `test_` prefix). Run by tests/kernels/moe/ ++test_merlin_full_k.py once per arm -- MERLIN_FULL_K=0, MERLIN_FULL_K=1 and, ++optionally, the untouched v2.0.1 install's interpreter -- because the switch ++is a function-local static read on the first Marlin launch and so is fixed ++for a process's lifetime. It deliberately imports only what v2.0.1 also ++ships, so the same file runs against the production wheel. ++ ++Calls the Marlin MoE GEMM (`ops.moe_wna16_marlin_gemm`) DIRECTLY on fixed, ++seeded inputs -- one expert, top_k 1, W4A16 uint4b8 group 128, bf16, K 2048, ++moe_block_size 16, `use_fp32_reduce=True, use_atomic_add=False` as the ++serving path passes them -- and saves every output to argv[1]: ++ ++ cases per (config, geometry): TWO launches on identical inputs. With the ++ explicit config (thread_k = thread_n = 128, blocks_per_sm = 1) the ++ grid is exactly `sms` CTAs, so the geometries put the tile count ++ (= num_tokens_past_padded / 16 * N / 128) below, equal to and above ++ one grid, the last two above-grid shapes landing on both stream-K ++ branches (remainder * 3 > grid: data-parallel + split tail; ++ remainder * 3 <= grid: remainder + one full grid of two-tile SK). ++ The auto config gets the same shapes; its grid is the heuristic's. ++ series per config: one weight matrix, token prefixes of 16..96 rows, so ++ the SAME first 16 tokens are computed at 1..6 m-blocks -- tile ++ counts T..6T with T = sms // 2 (41 on a 3090: 41, 82 = one grid, ++ 123, 164, 205, 246). This is the batch-invariance property FULL_K ++ exists for: a token's output must not depend on how many other ++ tokens share the launch. ++""" ++ ++import os ++import sys ++ ++import torch ++ ++ ++def main(out_path: str) -> None: ++ from vllm import _custom_ops as ops ++ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( ++ moe_align_block_size, ++ ) ++ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( ++ marlin_make_workspace_new, ++ ) ++ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( ++ marlin_quantize, ++ ) ++ from vllm.scalar_type import scalar_types ++ ++ dev = torch.device("cuda:0") ++ torch.cuda.set_device(dev) ++ sms = torch.cuda.get_device_properties(dev).multi_processor_count ++ dtype = torch.bfloat16 ++ qtype = scalar_types.uint4b8 ++ K, GROUP, BLOCK = 2048, 128, 16 ++ TN = 128 # thread_n of the explicit config ++ configs = {"explicit": (128, 128, 1), "auto": (-1, -1, -1)} ++ ++ gen = torch.Generator(device="cpu").manual_seed(46639) ++ a_full = (torch.randn(96, K, generator=gen) / 10).to(dtype).to(dev) ++ ++ def weights(n: int): ++ w = (torch.randn(K, n, generator=gen) / 10).to(dtype).to(dev) ++ _, q, s, _, _, _ = marlin_quantize(w, qtype, GROUP, False) ++ return q.unsqueeze(0).contiguous(), s.unsqueeze(0).contiguous() ++ ++ def gemm(m: int, n: int, q, s, cfg): ++ tk, tn, bps = cfg ++ a = a_full[:m].contiguous() ++ topk_ids = torch.zeros((m, 1), dtype=torch.int32, device=dev) ++ topk_w = torch.ones((m, 1), dtype=torch.float32, device=dev) ++ sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, BLOCK, 1) ++ ws = marlin_make_workspace_new(dev, 4) ++ out = ops.moe_wna16_marlin_gemm( ++ a, None, q, None, s, None, None, None, None, None, ws, ++ sorted_ids, expert_ids, ntpp, topk_w, ++ moe_block_size=BLOCK, top_k=1, mul_topk_weights=False, ++ b_q_type=qtype, size_m=m, size_n=n, size_k=K, is_k_full=True, ++ use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, ++ thread_k=tk, thread_n=tn, blocks_per_sm=bps, ++ ) # fmt: skip ++ torch.cuda.synchronize() ++ tiles = int(ntpp.item()) // BLOCK * (n // TN) ++ return out.float().cpu(), tiles ++ ++ half = sms // 2 ++ geometries = { ++ # name: n_tiles at one m-block (M = 16) ++ "below": sms - 2, ++ "equal": sms, ++ "above_dp_plus_split": 2 * sms + half, # remainder*3 > grid ++ "above_two_tile_sk": 2 * sms + 3, # remainder*3 <= grid ++ } ++ result = { ++ "sms": sms, ++ "env_MERLIN_FULL_K": os.environ.get("MERLIN_FULL_K"), ++ "torch": str(torch.__version__), ++ "cases": {}, ++ "series": {}, ++ } ++ for gname, n_tiles in geometries.items(): ++ n = TN * n_tiles ++ q, s = weights(n) ++ for cname, cfg in configs.items(): ++ o1, tiles = gemm(16, n, q, s, cfg) ++ o2, _ = gemm(16, n, q, s, cfg) ++ result["cases"][f"{cname}/{gname}"] = { ++ "tiles": tiles, ++ "grid_explicit": sms if cname == "explicit" else None, ++ "first": o1, ++ "second": o2, ++ } ++ n = TN * half ++ q, s = weights(n) ++ for cname, cfg in configs.items(): ++ series = {} ++ for m in (16, 32, 48, 64, 80, 96): ++ o, tiles = gemm(m, n, q, s, cfg) ++ series[m] = {"tiles": tiles, "out": o} ++ result["series"][cname] = series ++ torch.save(result, out_path) ++ print(f"FULL_K_CHILD_DONE sms={sms} env={result['env_MERLIN_FULL_K']}") ++ ++ ++if __name__ == "__main__": ++ main(sys.argv[1]) +diff --git a/tests/kernels/moe/test_merlin_full_k.py b/tests/kernels/moe/test_merlin_full_k.py +new file mode 100644 +index 0000000..88904b0 +--- /dev/null ++++ b/tests/kernels/moe/test_merlin_full_k.py +@@ -0,0 +1,188 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""MERLIN_FULL_K attribution cell. ++ ++The Marlin MoE GEMM called DIRECTLY on fixed inputs (see ++merlin_full_k_child.py for the shapes and why) in separate child processes, ++because `MERLIN_FULL_K` is read once per process: ++ ++ arm OFF MERLIN_FULL_K=0, this build ++ arm ON MERLIN_FULL_K=1, this build ++ arm V201 optional: the untouched v2.0.1 install, when ++ MERLIN_FULL_K_BASELINE_PYTHON (and, if needed, ++ MERLIN_FULL_K_BASELINE_PYTHONPATH) name its interpreter ++ ++Asserted: ++ * ON: two launches on identical inputs are bit-identical, every geometry ++ (below / equal / above one grid, both stream-K branches), explicit and ++ auto config; ++ * ON, explicit config: the same 16 tokens computed inside launches of 1..6 ++ m-blocks (tile counts T..6T, T = sms // 2) are bit-identical -- the ++ batch-invariance claim FULL_K is for; ++ * ON vs OFF: numerically close (bf16 tolerance); ++ * the toggle is OBSERVABLE: ON differs from OFF somewhere, or OFF breaks ++ the series invariance somewhere -- otherwise the cell cannot attribute ++ anything to FULL_K and says so; ++ * V201 (when run): bit-identical to same-build OFF, i.e. OFF really is the ++ v2.0.1 algorithm ("FULL_K=0 is not the production binary"). ++Reported, NOT asserted: OFF two-launch variability and OFF series drift. ++Set MERLIN_FULL_K_REPORT_DIR to also write the report as JSON. ++""" ++ ++import json ++import os ++import subprocess ++import sys ++from pathlib import Path ++ ++import pytest ++import torch ++ ++CHILD = Path(__file__).with_name("merlin_full_k_child.py") ++ ++pytestmark = pytest.mark.skipif( ++ not torch.cuda.is_available(), reason="needs a CUDA GPU (sm_86)" ++) ++ ++ ++def _run_child(out_dir: Path, tag: str, full_k: str | None, python: str | None = None, ++ pythonpath: str | None = None) -> dict: ++ env = dict(os.environ) ++ env.pop("MERLIN_FULL_K", None) ++ if full_k is not None: ++ env["MERLIN_FULL_K"] = full_k ++ if pythonpath is not None: ++ env["PYTHONPATH"] = pythonpath ++ out = out_dir / f"full_k_{tag}.pt" ++ proc = subprocess.run( ++ [python or sys.executable, str(CHILD), str(out)], ++ env=env, ++ capture_output=True, ++ text=True, ++ timeout=3600, ++ ) ++ assert proc.returncode == 0 and out.exists(), ( ++ f"FULL_K child '{tag}' failed rc={proc.returncode}\n" ++ f"{proc.stdout[-1500:]}\n{proc.stderr[-3000:]}" ++ ) ++ return torch.load(out, weights_only=True) ++ ++ ++@pytest.fixture(scope="module") ++def arms(tmp_path_factory) -> dict: ++ d = tmp_path_factory.mktemp("full_k") ++ return {"off": _run_child(d, "off", "0"), "on": _run_child(d, "on", "1"), ++ "dir": d} ++ ++ ++def _ndiff(a: torch.Tensor, b: torch.Tensor) -> int: ++ return int((a != b).sum()) ++ ++ ++def _report(arms: dict) -> dict: ++ off, on = arms["off"], arms["on"] ++ rep = {"sms": on["sms"], "cases": {}, "series": {}} ++ for name, c_on in on["cases"].items(): ++ c_off = off["cases"][name] ++ rep["cases"][name] = { ++ "tiles": c_on["tiles"], ++ "off_two_launch_diff": _ndiff(c_off["first"], c_off["second"]), ++ "on_two_launch_diff": _ndiff(c_on["first"], c_on["second"]), ++ "on_vs_off_diff": _ndiff(c_on["first"], c_off["first"]), ++ "on_vs_off_maxabs": float((c_on["first"] - c_off["first"]).abs().max()), ++ "numel": c_on["first"].numel(), ++ } ++ for cfg, s_on in on["series"].items(): ++ s_off = off["series"][cfg] ++ ref_on, ref_off = s_on[16]["out"], s_off[16]["out"] ++ rep["series"][cfg] = { ++ m: { ++ "tiles": s_on[m]["tiles"], ++ "on_prefix_diff_vs_m16": _ndiff(s_on[m]["out"][:16], ref_on), ++ "off_prefix_diff_vs_m16": _ndiff(s_off[m]["out"][:16], ref_off), ++ "on_vs_off_diff": _ndiff(s_on[m]["out"], s_off[m]["out"]), ++ } ++ for m in s_on ++ } ++ print("\n[full-K isolated case] MERLIN_FULL_K report:\n" + json.dumps(rep, indent=1)) ++ out_dir = os.environ.get("MERLIN_FULL_K_REPORT_DIR") ++ if out_dir: ++ Path(out_dir).mkdir(parents=True, exist_ok=True) ++ (Path(out_dir) / "full_k_report.json").write_text(json.dumps(rep, indent=1)) ++ return rep ++ ++ ++def test_isolated_full_k_on_is_bit_stable_and_batch_invariant(arms) -> None: ++ on = arms["on"] ++ assert on["env_MERLIN_FULL_K"] == "1" and arms["off"]["env_MERLIN_FULL_K"] == "0" ++ rep = _report(arms) ++ sms = on["sms"] ++ # the geometry claims the docstring makes, checked on the explicit config ++ tiles = {n.split("/", 1)[1]: c["tiles"] for n, c in on["cases"].items() ++ if n.startswith("explicit/")} ++ assert tiles["below"] < sms == tiles["equal"] < tiles["above_dp_plus_split"] ++ assert (tiles["above_dp_plus_split"] % sms) * 3 > sms ++ assert 0 < (tiles["above_two_tile_sk"] % sms) * 3 <= sms ++ ++ for name, c in on["cases"].items(): ++ assert torch.equal(c["first"], c["second"]), ( ++ f"MERLIN_FULL_K=1 {name} ({c['tiles']} tiles): two launches on " ++ f"identical inputs differ in {rep['cases'][name]['on_two_launch_diff']}" ++ " elements" ++ ) ++ ser = on["series"]["explicit"] ++ for m, e in ser.items(): ++ assert torch.equal(e["out"][:16], ser[16]["out"]), ( ++ f"MERLIN_FULL_K=1, explicit config: the first 16 tokens at M={m} " ++ f"({e['tiles']} tiles) differ from M=16 ({ser[16]['tiles']} tiles) in " ++ f"{rep['series']['explicit'][m]['on_prefix_diff_vs_m16']} elements --" ++ " the output depends on the batch" ++ ) ++ ++ ++def test_isolated_full_k_on_vs_off_close_and_observable(arms) -> None: ++ rep = _report(arms) ++ off, on = arms["off"], arms["on"] ++ for name, c in on["cases"].items(): ++ torch.testing.assert_close( ++ c["first"], off["cases"][name]["first"], atol=2e-2, rtol=2e-2, ++ msg=lambda m, name=name: f"{name}: ON vs OFF not close. {m}", ++ ) ++ for cfg, s in on["series"].items(): ++ for m, e in s.items(): ++ torch.testing.assert_close(e["out"], off["series"][cfg][m]["out"], ++ atol=2e-2, rtol=2e-2) ++ differs = sum(v["on_vs_off_diff"] for v in rep["cases"].values()) + sum( ++ v["on_vs_off_diff"] + v["off_prefix_diff_vs_m16"] ++ for s in rep["series"].values() for v in s.values() ++ ) ++ assert differs > 0, ( ++ "ON and OFF are bit-identical on every shape and OFF is batch-invariant " ++ "too: this cell cannot attribute anything to MERLIN_FULL_K at these " ++ "shapes (not a kernel failure -- the instrument does not discriminate)" ++ ) ++ ++ ++def test_isolated_full_k_off_is_the_untouched_v201_algorithm(arms) -> None: ++ py = os.environ.get("MERLIN_FULL_K_BASELINE_PYTHON") ++ if not py: ++ pytest.skip( ++ "third arm not run: set MERLIN_FULL_K_BASELINE_PYTHON (and " ++ "MERLIN_FULL_K_BASELINE_PYTHONPATH) to the untouched v2.0.1 " ++ "install. A skip is not a pass." ++ ) ++ base = _run_child(arms["dir"], "v201", None, python=py, ++ pythonpath=os.environ.get("MERLIN_FULL_K_BASELINE_PYTHONPATH")) ++ off = arms["off"] ++ for name, c in base["cases"].items(): ++ # first launches only: OFF's second launch may legitimately vary ++ assert torch.equal(c["first"], off["cases"][name]["first"]), ( ++ f"{name}: the untouched v2.0.1 Marlin and same-build MERLIN_FULL_K=0 " ++ f"differ in {_ndiff(c['first'], off['cases'][name]['first'])} " ++ "elements -- OFF is not the production algorithm" ++ ) ++ for cfg, s in base["series"].items(): ++ for m, e in s.items(): ++ assert torch.equal(e["out"], off["series"][cfg][m]["out"]), ( ++ f"series {cfg} M={m}: v2.0.1 and same-build OFF differ" ++ ) +diff --git a/tests/kernels/test_merlin_qsa_sort.py b/tests/kernels/test_merlin_qsa_sort.py +new file mode 100644 +index 0000000..7496e3c +--- /dev/null ++++ b/tests/kernels/test_merlin_qsa_sort.py +@@ -0,0 +1,161 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""MERLIN_QSA_SORT attribution cell. ++ ++`_qsa_canonicalise` (qsa_indexer.py) sorts the selected block indices of each ++row ascending, pushes negative padding to the END as -1, in place, shape ++preserved, and does it without a host sync so it can sit inside a captured ++cudagraph. These tests pin exactly that contract on CPU and GPU: ++ ++ * permutations of a valid row come back as the one ascending row; ++ * mixed padding (negatives interleaved anywhere) -> valid ascending prefix, ++ then -1s, count of valid entries preserved; ++ * all padding -> all -1; ++ * the switch OFF is a byte-for-byte no-op; ++ * CUDA: no host synchronisation (torch.cuda sync-debug mode "error"), and ++ `torch.cuda.graph` capture SUCCEEDS and the replay equals the eager result ++ on fresh data copied into the static input. ++ ++The module-level `_QSA_SORT` is resolved once at import from the environment, ++so the tests set it with monkeypatch on the module, not through the env. ++""" ++ ++import pytest ++import torch ++ ++qsa_indexer = pytest.importorskip( ++ "vllm.models.qwen4_exp.nvidia.ops.qsa_indexer", ++ reason="qsa_indexer needs triton (the reference-host venv has it)", ++) ++ ++DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) ++ ++ ++def _reference(rows: torch.Tensor) -> torch.Tensor: ++ """Per row: valid (>= 0) entries ascending, then -1 for every negative.""" ++ out = torch.full_like(rows, -1) ++ for i in range(rows.shape[0]): ++ valid = rows[i][rows[i] >= 0].sort().values ++ out[i, : valid.numel()] = valid ++ return out ++ ++ ++@pytest.fixture ++def sort_on(monkeypatch): ++ monkeypatch.setattr(qsa_indexer, "_QSA_SORT", True) ++ ++ ++def _rows(kind: str, n_rows: int, width: int, dtype: torch.dtype, seed: int): ++ gen = torch.Generator(device="cpu").manual_seed(seed) ++ base = torch.stack( ++ [torch.randperm(4 * width, generator=gen)[:width] for _ in range(n_rows)] ++ ).to(dtype) ++ if kind == "permutation": ++ return base ++ if kind == "mixed_padding": ++ # negatives interleaved anywhere, several spellings of "invalid" ++ mask = torch.rand(n_rows, width, generator=gen) < 0.35 ++ pad = torch.tensor([-1, -2, -7], dtype=dtype)[ ++ torch.randint(0, 3, (n_rows, width), generator=gen) ++ ] ++ rows = torch.where(mask, pad, base) ++ rows[0, :] = -1 # one fully padded row among mixed ones ++ rows[1, :] = base[1] # and one with no padding at all ++ return rows ++ if kind == "all_padding": ++ return torch.full((n_rows, width), -1, dtype=dtype) ++ raise ValueError(kind) ++ ++ ++@pytest.mark.parametrize("device", DEVICES) ++@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) ++@pytest.mark.parametrize("kind", ["permutation", "mixed_padding", "all_padding"]) ++def test_isolated_qsa_canonicalise_contract(sort_on, device, dtype, kind) -> None: ++ rows = _rows(kind, n_rows=8, width=512, dtype=dtype, seed=3).to(device) ++ expected = _reference(rows.cpu()) ++ ptr = rows.data_ptr() ++ qsa_indexer._qsa_canonicalise(rows) ++ assert rows.data_ptr() == ptr, "not in place" ++ assert rows.dtype == dtype and rows.shape == expected.shape ++ assert torch.equal(rows.cpu(), expected), f"{kind}: not the canonical row" ++ ++ ++@pytest.mark.parametrize("device", DEVICES) ++def test_isolated_qsa_canonicalise_permutation_invariance(sort_on, device) -> None: ++ """Any permutation of one selected set gives the SAME output bytes.""" ++ gen = torch.Generator(device="cpu").manual_seed(5) ++ selected = torch.randperm(100_000, generator=gen)[:2048].to(torch.int32) ++ selected[::7] = -1 # padding scattered through ++ outs = [] ++ for s in range(6): ++ g = torch.Generator(device="cpu").manual_seed(100 + s) ++ row = selected[torch.randperm(selected.numel(), generator=g)] ++ row = row.view(1, -1).to(device).contiguous() ++ qsa_indexer._qsa_canonicalise(row) ++ outs.append(row.cpu()) ++ for o in outs[1:]: ++ assert torch.equal(o, outs[0]), "permutations of one set disagree" ++ ++ ++@pytest.mark.parametrize("device", DEVICES) ++def test_isolated_qsa_canonicalise_higher_rank(sort_on, device) -> None: ++ """Sorts along the LAST dim only; leading dims untouched.""" ++ rows = _rows("mixed_padding", n_rows=12, width=64, dtype=torch.int32, seed=9) ++ rows3 = rows.view(3, 4, 64).to(device) ++ qsa_indexer._qsa_canonicalise(rows3) ++ assert torch.equal(rows3.cpu().view(12, 64), _reference(rows)) ++ ++ ++@pytest.mark.parametrize("device", DEVICES) ++def test_isolated_qsa_canonicalise_off_is_a_noop(monkeypatch, device) -> None: ++ monkeypatch.setattr(qsa_indexer, "_QSA_SORT", False) ++ rows = _rows("mixed_padding", n_rows=4, width=128, dtype=torch.int32, seed=1) ++ rows = rows.to(device) ++ before = rows.clone() ++ qsa_indexer._qsa_canonicalise(rows) ++ assert torch.equal(rows, before), "MERLIN_QSA_SORT=0 must not touch the row" ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") ++def test_isolated_qsa_canonicalise_no_host_sync(sort_on) -> None: ++ """The branch-free design claim: no device->host synchronisation.""" ++ rows = _rows("mixed_padding", n_rows=8, width=2048, dtype=torch.int32, seed=11) ++ rows = rows.cuda() ++ expected = _reference(rows.cpu()) ++ torch.cuda.synchronize() ++ prev = torch.cuda.get_sync_debug_mode() ++ torch.cuda.set_sync_debug_mode("error") ++ try: ++ qsa_indexer._qsa_canonicalise(rows) ++ finally: ++ torch.cuda.set_sync_debug_mode(prev) ++ torch.cuda.synchronize() ++ assert torch.equal(rows.cpu(), expected) ++ ++ ++@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") ++@pytest.mark.parametrize("width", [512, 2048]) ++def test_isolated_qsa_canonicalise_under_cudagraph_capture(sort_on, width) -> None: ++ """Capture must SUCCEED (no .item()/.any() sync) and replay must be exact.""" ++ static = _rows("mixed_padding", 8, width, torch.int32, seed=21).cuda() ++ # warm up on a side stream, as vLLM does before capture ++ s = torch.cuda.Stream() ++ s.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(s): ++ qsa_indexer._qsa_canonicalise(static.clone()) ++ torch.cuda.current_stream().wait_stream(s) ++ torch.cuda.synchronize() ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): # raises if anything in the op syncs ++ qsa_indexer._qsa_canonicalise(static) ++ ++ for seed in (22, 23, 24): ++ fresh = _rows("mixed_padding", 8, width, torch.int32, seed=seed) ++ fresh[2] = fresh[2].flip(0) ++ static.copy_(fresh.cuda()) ++ graph.replay() ++ torch.cuda.synchronize() ++ assert torch.equal(static.cpu(), _reference(fresh)), ( ++ f"replay on seed {seed} differs from the eager contract" ++ ) +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index 2b45b3b..58bcb7f 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -3330,3 +3330,692 @@ def test_persistent_topk_large_within_group_mixed_reuse() -> None: + "deterministic selected set, but two identical launches " + "selected different indices" + ) ++ ++ ++# =========================================================================== ++# Cells that need a GPU host or a diagnostic build, prepared ++# offline. Every test here is named ++# `test_isolated_*` so one `-k isolated` selects them; each SKIPS outside the process ++# environment (or the diagnostic build) it exists for, and every skip names ++# the environment that runs it. A skip is never a pass. ++# ++# cell 1 OFF-arm oracle controls that CAN fail (medium + decode), and the ++# ON arms on the identical inputs ++# cell 4 MERLIN_TIE_DET non-overflow cutoff ties, both directions, oracle ++# cell 5 delayed-producer publication stress (MERLIN_TOPK_DIAG_DELAY build, ++# repaired vs MERLIN_TOPK_DIAG_PREREPAIR=1) ++# cell 6 trap-ON fault in a disposable subprocess ++# (MERLIN_TOPK_DIAG_FORCE_FAULT build) ++# Cells 2 (FULL_K) and 3 (QSA_SORT) live in tests/kernels/moe/ ++# test_merlin_full_k.py and tests/kernels/test_merlin_qsa_sort.py. ++# =========================================================================== ++ ++ISOLATED_MEDIUM_CAP = 4096 # MAX_BUFFERED_ITEMS, persistent_topk.cuh ++ISOLATED_DECODE_DBUF = 3708 # DBUF in histogram_2048_topk ++ISOLATED_THREADS = 1024 # kThreadsPerBlock ++ISOLATED_NUM_ROWS = 4 ++ ++ ++def _isolated_bins_or_fail() -> None: ++ """The fixtures rest on these bin facts; assert them rather than assume.""" ++ ++ def half_key(v: float) -> int: ++ h = int(np.float16(np.float32(v)).view(np.uint16)) ++ return (h | 0x8000) if not (h & 0x8000) else (~h) & 0xFFFF ++ ++ # medium coarse bin = half key >> 8; decode bin = half key >> 5 ++ assert half_key(1.0) >> 8 == half_key(1.01) >> 8 == 0xBC ++ assert half_key(2.0) >> 8 > 0xBC > half_key(0.5) >> 8 ++ assert half_key(1.0) >> 5 == half_key(1.01) >> 5 ++ assert half_key(2.0) >> 5 > half_key(1.0) >> 5 > half_key(0.5) >> 5 ++ assert np.float32(1.01) > np.float32(1.0) ++ ++ ++def _isolated_medium_drop_fixture(top_k: int) -> tuple[torch.Tensor, torch.Tensor]: ++ """A medium row whose 4,096-slot overflow PROVABLY drops a true winner. ++ ++ `histogram_256_topk`'s collection loop hands buffer slots out with one ++ per-ELEMENT atomicAdd; thread t visits idx = t, t + 1024, t + 2048, ... in ++ program order, so ITS appends are ordered, and nothing else is. The ++ fixture puts threshold-bin candidates (1.0, coarse bin 0xBC) in the first ++ `per_thread` visits of `threads` threads and makes each such thread's LAST ++ candidate a true winner (1.01: same coarse bin, larger at FP32). The ++ candidate count exceeds 4,096, so the append that lands last in the global ++ atomic order gets a slot >= 4,096 and is dropped -- and whichever thread ++ issued it, it is that thread's last append, i.e. a winner. Every legal ++ interleaving drops at least one winner: the OFF arm is wrong on every row, ++ deterministically. ++ """ ++ _isolated_bins_or_fail() ++ seq_len = 32768 # the widest medium row (<= RADIX_THRESHOLD) ++ if top_k == 2048: ++ threads, per_thread, above = 1024, 5, 512 ++ elif top_k == 512: ++ threads, per_thread, above = 512, 9, 0 ++ else: ++ raise ValueError(top_k) ++ with torch.device("cpu"): ++ row = torch.full((seq_len,), 0.5, dtype=torch.float32) ++ t = torch.arange(threads) ++ for j in range(per_thread): ++ row[t + ISOLATED_THREADS * j] = 1.0 ++ winners = t + ISOLATED_THREADS * (per_thread - 1) ++ row[winners] = float(np.float32(1.01)) ++ start = ISOLATED_THREADS * per_thread ++ row[start : start + above] = 2.0 ++ candidates = threads * per_thread ++ assert candidates > ISOLATED_MEDIUM_CAP, "the fixture must overflow" ++ assert threads <= top_k - above, "every winner must belong in the exact top-k" ++ assert int((row == 1.0).sum()) + threads == candidates ++ return row, winners ++ ++ ++def _isolated_decode_drop_fixture(top_k: int) -> tuple[torch.Tensor, torch.Tensor]: ++ """A decode row whose 3,708-slot overflow PROVABLY drops a true winner. ++ ++ The decode collection is warp-aggregated: per (iteration, sub-element) ++ step, lane 0 reserves one contiguous range for the whole ballot and the ++ lane of rank `r` writes slot `base + r` only if it is < DBUF. A warp's ++ reservations are in program order (iteration 0 sub 0..3, then iteration ++ 1); across warps there is no order. The fixture fills every element of ++ iteration 0 (indices [0, 4096): 32 warps x 4 sub-steps x 32 lanes) with ++ threshold-bin candidates and makes lane 31 of each warp's LAST candidate ++ step (iteration 0, sub 3) a true winner. 4,096 > 3,708 candidates, so the ++ reservation that lands last globally starts at 4,064 >= 3,708 and is ++ dropped whole -- and it is some warp's last one, whose top lane is a ++ winner. Every legal interleaving drops at least one winner. ++ """ ++ _isolated_bins_or_fail() ++ seq_len = 8192 # the widest decode row (<= HIST2048_THRESHOLD) ++ above = 64 ++ with torch.device("cpu"): ++ row = torch.full((seq_len,), 0.5, dtype=torch.float32) ++ row[:4096] = 1.0 ++ warps = torch.arange(32) ++ winners = 4 * (warps * 32 + 31) + 3 ++ row[winners] = float(np.float32(1.01)) ++ row[4096 : 4096 + above] = 2.0 ++ assert 4096 > ISOLATED_DECODE_DBUF ++ assert 32 <= top_k - above ++ return row, winners ++ ++ ++def _isolated_launch(row: torch.Tensor, top_k: int, num_rows: int = ISOLATED_NUM_ROWS): ++ seq_len = row.numel() ++ logits = row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") ++ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") ++ _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, seq_len) ++ torch.accelerator.synchronize() ++ return indices.cpu().to(torch.int64) ++ ++ ++def _isolated_counters_or_skip() -> list[int]: ++ counters = _read_exact_counters() ++ if counters is None: ++ pytest.skip( ++ "QUALIFICATION FAILURE if seen on the branch build: no " ++ "persistent_topk_exact_counters op. A skip is not a pass." ++ ) ++ return counters ++ ++ ++def _isolated_assert_off_arm_failed( ++ row: torch.Tensor, got: torch.Tensor, winners: torch.Tensor, top_k: int, arm: str ++) -> None: ++ with torch.device("cpu"): ++ ref_values = torch.topk(row, top_k).values.sort().values ++ for r in range(got.shape[0]): ++ g = got[r] ++ # Well-formed: the failure under test is a WRONG top-k, not garbage. ++ assert int(g.min()) >= 0 and int(g.max()) < row.numel(), f"{arm} row {r}" ++ assert g.unique().numel() == top_k, f"{arm} row {r}: duplicate index" ++ missing = int((~torch.isin(winners, g)).sum()) ++ assert missing >= 1, ( ++ f"{arm} row {r}: the OFF arm kept every true winner. The " ++ "fixture is built so that no legal append order can do that -- " ++ "the control did not exercise the defect, so it proves nothing." ++ ) ++ assert not torch.equal(row[g].sort().values, ref_values), ( ++ f"{arm} row {r}: OFF-arm values equal torch.topk's although " ++ f"{missing} winner(s) are missing" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_isolated_medium_off_arm_drops_a_true_winner(top_k: int) -> None: ++ """Cell 1 (medium control): MERLIN_TOPK_EXACT_MEDIUM=0 must be WRONG here. ++ ++ The same-input, oracle-sensitive control that was missing (the ++ existing OFF test skips before constructing an overflow row). Asserts the ++ overflow counted, no restart taken, and a returned set that misses a true ++ winner on every row. Its ON twin below runs the identical input. ++ """ ++ torch.set_default_device("cuda:0") ++ before = _isolated_counters_or_skip() ++ if before[_SLOT_EXACT_MEDIUM_ON] != 0: ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_MEDIUM=1 in this process: this is the OFF-arm " ++ "control. Run it in a fresh process with MERLIN_TOPK_EXACT_MEDIUM=0." ++ " A skip is not a pass." ++ ) ++ row, winners = _isolated_medium_drop_fixture(top_k) ++ got = _isolated_launch(row, top_k) ++ after = _read_exact_counters() ++ n = ISOLATED_NUM_ROWS ++ assert after[_SLOT_MEDIUM_ROWS] - before[_SLOT_MEDIUM_ROWS] == n ++ assert after[_SLOT_MEDIUM_OVERFLOW] - before[_SLOT_MEDIUM_OVERFLOW] == n, ( ++ "the overflow counter must fire once per row in the control arm too" ++ ) ++ assert after[_SLOT_MEDIUM_FALLBACK] == before[_SLOT_MEDIUM_FALLBACK], ( ++ "the OFF arm took the exact restart" ++ ) ++ assert after[_SLOT_MEDIUM_FAULT] == before[_SLOT_MEDIUM_FAULT] ++ _isolated_assert_off_arm_failed(row, got, winners, top_k, "medium OFF") ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_isolated_medium_on_arm_is_exact_on_the_drop_fixture(top_k: int) -> None: ++ """Cell 1 (medium treatment): the identical input with the restart ON is exact.""" ++ torch.set_default_device("cuda:0") ++ before = _isolated_counters_or_skip() ++ if before[_SLOT_EXACT_MEDIUM_ON] == 0: ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_MEDIUM=0 in this process: this is the ON arm. " ++ "Run it with the switch on (the default). A skip is not a pass." ++ ) ++ row, winners = _isolated_medium_drop_fixture(top_k) ++ expected = _largedet_reference(row, top_k) # the restart's effective direction ++ got = _isolated_launch(row, top_k) ++ after = _read_exact_counters() ++ n = ISOLATED_NUM_ROWS ++ assert after[_SLOT_MEDIUM_OVERFLOW] - before[_SLOT_MEDIUM_OVERFLOW] == n ++ assert after[_SLOT_MEDIUM_FALLBACK] - before[_SLOT_MEDIUM_FALLBACK] == n ++ assert after[_SLOT_MEDIUM_FAULT] == before[_SLOT_MEDIUM_FAULT] ++ for r in range(n): ++ assert bool(torch.isin(winners, got[r]).all()), f"row {r}: a winner is missing" ++ assert torch.equal(got[r], expected), ( ++ f"row {r}: the exact restart did not return the oracle set in " ++ "ascending order" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_isolated_decode_off_arm_drops_a_true_winner(top_k: int) -> None: ++ """Cell 1 (decode control): MERLIN_TOPK_EXACT_DECODE=0 must be WRONG here. ++ ++ Replaces the detect-only test as the control; there the tail ++ winner is dropped only if warp scheduling happens to append it last: here ++ the drop is forced for every interleaving (see the fixture). ++ """ ++ torch.set_default_device("cuda:0") ++ if _largedet_env_on("MERLIN_TOPK_EXACT_DECODE", False): ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_DECODE=1 in this process: this is the OFF-arm " ++ "control. Run it in a process with the variable unset or 0. A skip " ++ "is not a pass." ++ ) ++ before = _isolated_counters_or_skip() ++ row, winners = _isolated_decode_drop_fixture(top_k) ++ got = _isolated_launch(row, top_k) ++ after = _read_exact_counters() ++ fired = after[_SLOT_DECODE_OVERFLOW] - before[_SLOT_DECODE_OVERFLOW] ++ assert fired == ISOLATED_NUM_ROWS, f"decode overflow counter moved by {fired}" ++ assert after[_SLOT_MEDIUM_FAULT] == before[_SLOT_MEDIUM_FAULT] ++ _isolated_assert_off_arm_failed(row, got, winners, top_k, "decode OFF") ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@torch.inference_mode() ++def test_isolated_decode_on_arm_is_exact_on_the_drop_fixture(top_k: int) -> None: ++ """Cell 1 (decode treatment): the identical input with the restart ON.""" ++ torch.set_default_device("cuda:0") ++ if not _largedet_env_on("MERLIN_TOPK_EXACT_DECODE", False): ++ pytest.skip( ++ "MERLIN_TOPK_EXACT_DECODE unset/0: this is the ON arm. Run it in a " ++ "fresh process with MERLIN_TOPK_EXACT_DECODE=1. A skip is not a pass." ++ ) ++ before = _isolated_counters_or_skip() ++ row, winners = _isolated_decode_drop_fixture(top_k) ++ expected = _largedet_reference(row, top_k) ++ got = _isolated_launch(row, top_k) ++ after = _read_exact_counters() ++ fired = after[_SLOT_DECODE_OVERFLOW] - before[_SLOT_DECODE_OVERFLOW] ++ assert fired == ISOLATED_NUM_ROWS, f"decode overflow counter moved by {fired}" ++ assert after[_SLOT_MEDIUM_FAULT] == before[_SLOT_MEDIUM_FAULT], ( ++ "the decode restart shares the medium restart's fault counter" ++ ) ++ for r in range(ISOLATED_NUM_ROWS): ++ assert bool(torch.isin(winners, got[r]).all()), f"row {r}: a winner is missing" ++ assert torch.equal(got[r], expected), ( ++ f"row {r}: the exact decode restart did not return the oracle set " ++ "in ascending order" ++ ) ++ ++ ++def _isolated_tie_fixture(path: str) -> tuple[torch.Tensor, int, int]: ++ """A NON-overflowing row with an exact 32-bit tie AT the k-th value. ++ ++ Returns (row, top_k, need). `above` elements at 3.0 and `tie` elements at ++ 2.0, scattered by a seeded permutation over the whole row so the tie group ++ spans every warp; filler in [0, 0.5). The tie group is one coarse bin and ++ fits the buffer (decode 700 <= 3,708; medium 2,048 <= 4,096), so the ++ BUFFERED path -- not a restart -- ranks it at pass 3. ++ """ ++ if path == "decode": ++ seq_len, top_k, above, tie = 6000, 512, 200, 700 ++ elif path == "medium": ++ seq_len, top_k, above, tie = 20000, 2048, 683, 2048 ++ else: ++ raise ValueError(path) ++ with torch.device("cpu"): ++ gen = torch.Generator(device="cpu").manual_seed(4242) ++ row = torch.rand(seq_len, generator=gen, dtype=torch.float32) * 0.5 ++ perm = torch.randperm(seq_len, generator=gen) ++ row[perm[:above]] = 3.0 ++ row[perm[above : above + tie]] = 2.0 ++ need = top_k - above ++ assert 0 < need < tie, "the cut must fall strictly inside the tie group" ++ return row, top_k, need ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++@pytest.mark.parametrize("path", ["decode", "medium"]) ++@pytest.mark.parametrize("want_recent", [True, False]) ++@torch.inference_mode() ++def test_isolated_tie_det_nonoverflow_cutoff_direction(path: str, want_recent: bool) -> None: ++ """Cell 4: MERLIN_TIE_DET's buffered cutoff tie, both directions, oracle. ++ ++ Why: the mixed-path test's sorted-set equality of two launches is ++ a sampled repeatability check; it does not show that cutoff ties follow ++ the configured direction. Here the kept tied indices are FORCED -- the ++ `need` highest (most-recent-wins) or lowest (oldest-wins) indices of the ++ tie group -- and `_largedet_reference` states them. Compared as SETS (sorted): ++ TIE_DET fixes membership; raw order is atomic arrival on these paths. ++ """ ++ torch.set_default_device("cuda:0") ++ if not _largedet_env_on("MERLIN_TIE_DET", False): ++ pytest.skip( ++ "MERLIN_TIE_DET unset/0: the buffered paths have no tie direction. " ++ "Run with MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=1 (most-recent) and " ++ "MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=0 (oldest), each in a fresh " ++ "process. A skip is not a pass." ++ ) ++ if _largedet_effective_tie_recent() != want_recent: ++ pytest.skip( ++ f"this process resolves {'oldest' if want_recent else 'most-recent'}" ++ f"-wins; the {'most-recent' if want_recent else 'oldest'}-wins arm " ++ "runs in the other MERLIN_TIE_RECENT process." ++ ) ++ row, top_k, need = _isolated_tie_fixture(path) ++ expected = _largedet_reference(row, top_k, tie_recent=want_recent) ++ other = _largedet_reference(row, top_k, tie_recent=not want_recent) ++ assert not torch.equal(expected, other), "the fixture cannot tell the directions apart" ++ ++ before = _isolated_counters_or_skip() ++ first = _isolated_launch(row, top_k) ++ second = _isolated_launch(row, top_k) ++ after = _read_exact_counters() ++ ++ assert after[_SLOT_DECODE_OVERFLOW] == before[_SLOT_DECODE_OVERFLOW], ( ++ "the tie fixture overflowed the decode buffer; it must not" ++ ) ++ assert after[_SLOT_MEDIUM_OVERFLOW] == before[_SLOT_MEDIUM_OVERFLOW], ( ++ "the tie fixture overflowed the medium buffer; it must not" ++ ) ++ assert after[_SLOT_MEDIUM_FALLBACK] == before[_SLOT_MEDIUM_FALLBACK] ++ medium_rows = after[_SLOT_MEDIUM_ROWS] - before[_SLOT_MEDIUM_ROWS] ++ assert medium_rows == (2 * ISOLATED_NUM_ROWS if path == "medium" else 0), ( ++ f"the {path} fixture was dispatched to the wrong path ({medium_rows} " ++ "medium rows)" ++ ) ++ rule = "most-recent" if want_recent else "oldest" ++ for launch, got in (("first", first), ("second", second)): ++ for r in range(ISOLATED_NUM_ROWS): ++ assert torch.equal(got[r].sort().values, expected), ( ++ f"{path} {launch} launch row {r}: the {rule}-wins cutoff set is " ++ f"wrong ({int((~torch.isin(expected, got[r])).sum())} of the " ++ f"{top_k} oracle indices missing; the cut keeps {need})" ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# Cells 5 and 6: DIAGNOSTIC BUILDS. Both run the kernel in a CHILD process ++# (this file executed as a script), so a trap or a corrupted context cannot ++# take the pytest process with it, and both key on the build's own ++# `[merlin] persistent_topk DIAGNOSTIC BUILD:` stderr line to prove which ++# build the child ran. On a normal build they SKIP after a one-row probe. ++# --------------------------------------------------------------------------- ++ ++_ISOLATED_DIAG_RE = re.compile( ++ r"\[merlin\] persistent_topk DIAGNOSTIC BUILD: MERLIN_TOPK_DIAG_DELAY=(\d+) " ++ r"MERLIN_TOPK_DIAG_PREREPAIR=(\d+) MERLIN_TOPK_DIAG_FORCE_FAULT=(\d+) " ++ r"MERLIN_TOPK_DIAG_FORCE_FAULT_CTA=(\d+)" ++) ++_ISOLATED_RESULT_TAG = "ISOLATED_DIAG_RESULT " ++ ++ ++def _isolated_spawn(mode: str, *args: str, drop_env: tuple[str, ...] = ()): ++ import json ++ import subprocess ++ import sys ++ ++ env = dict(os.environ) ++ for name in drop_env: ++ env.pop(name, None) ++ proc = subprocess.run( ++ [sys.executable, os.path.abspath(__file__), "--isolated-diag-child", mode, *args], ++ env=env, ++ capture_output=True, ++ text=True, ++ timeout=1800, ++ ) ++ result = None ++ for line in proc.stdout.splitlines(): ++ if line.startswith(_ISOLATED_RESULT_TAG): ++ result = json.loads(line[len(_ISOLATED_RESULT_TAG) :]) ++ m = _ISOLATED_DIAG_RE.search(proc.stderr) ++ banner = tuple(int(x) for x in m.groups()) if m else None ++ return proc, result, banner ++ ++ ++def _isolated_diag_banner_or_skip(want: str) -> tuple[int, ...]: ++ proc, _, banner = _isolated_spawn("probe") ++ if banner is None: ++ pytest.skip( ++ f"not a diagnostic build (no DIAGNOSTIC BUILD line from the probe " ++ f"child, rc={proc.returncode}); the {want} cell runs only against " ++ "the diagnostic .so built with the MERLIN_TOPK_DIAG_* defines. A skip is " ++ "not a pass." ++ ) ++ return banner ++ ++ ++def _isolated_tail(text: str, n: int = 1500) -> str: ++ return text[-n:] if text else "" ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++def test_isolated_diag_delayed_publication_matches_oracle() -> None: ++ """Cell 5: delayed non-leader histogram publication, repeated large rows. ++ ++ Build with -DMERLIN_TOPK_DIAG_DELAY= (repaired) and additionally ++ -DMERLIN_TOPK_DIAG_PREREPAIR=1 (the earlier arrival barrier compiled out). The ++ child runs repeated large rows (four shapes, ctas_per_group 2..11) and ++ compares every row with the oracle. EXPECTED: PASS on the repaired ++ diagnostic build; FAIL on the pre-repair one (wrong set, all -1 rows with ++ faults_large > 0, or a trapped child) -- that failure is the control that ++ shows the stress can see the defect the barrier repairs. ++ """ ++ if not _largedet_env_on("MERLIN_TOPK_DET_LARGE", True): ++ pytest.skip( ++ "MERLIN_TOPK_DET_LARGE=0: the legacy tail does not bound its " ++ "emission slot, so a pre-repair pivot disagreement would write out " ++ "of bounds. This cell runs on the deterministic path only." ++ ) ++ delay, prerepair, force_fault, _ = _isolated_diag_banner_or_skip("publication") ++ if delay == 0 or force_fault: ++ pytest.skip( ++ f"diagnostic build without the publication delay (delay={delay}, " ++ f"force_fault={force_fault}); build the delay variant." ++ ) ++ proc, res, banner = _isolated_spawn("publication", "12") ++ tag = ( ++ f"[{'PRE-REPAIR' if prerepair else 'REPAIRED'} diagnostic build, " ++ f"MERLIN_TOPK_DIAG_DELAY={delay}]" ++ ) ++ assert banner is not None and banner[0] == delay, f"{tag} child ran another build" ++ assert proc.returncode == 0 and res is not None and "error" not in res, ( ++ f"{tag} the stress child died (rc={proc.returncode}): " ++ f"{(res or {}).get('error')}\n--- child stderr tail ---\n" ++ f"{_isolated_tail(proc.stderr)}" ++ ) ++ assert res["rows_large"] == res["rows_launched"], ( ++ f"{tag} the deterministic large path counted {res['rows_large']} of " ++ f"{res['rows_launched']} rows; the stress did not run where it claims" ++ ) ++ assert res["faults_large"] == 0 and res["all_minus_one_rows"] == 0, ( ++ f"{tag} {res['faults_large']} large-path faults, " ++ f"{res['all_minus_one_rows']} all -1 rows under delayed publication: " ++ f"{res['per_shape']}" ++ ) ++ assert res["mismatched_rows"] == 0, ( ++ f"{tag} {res['mismatched_rows']} of {res['rows_launched']} rows differ " ++ f"from the oracle under delayed publication: {res['per_shape']}" ++ ) ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++def test_isolated_diag_trap_on_kills_only_the_child() -> None: ++ """Cell 6: trap-ON fault arm, in a disposable subprocess. ++ ++ Build with -DMERLIN_TOPK_DIAG_FORCE_FAULT=1 (CTA 1 by default). The child ++ runs deterministic large rows with the runtime injector UNSET, so nothing ++ suppresses `__trap()`: the child must die with a CUDA launch failure and ++ never print its no-trap marker. The parent then proves its own context is ++ intact with a kernel that does not touch the large path (a large row in ++ this process would trap too -- run this test alone, `-k trap_on`). ++ """ ++ _, _, force_fault, cta = _isolated_diag_banner_or_skip("trap-on") ++ if not force_fault: ++ pytest.skip("diagnostic build without MERLIN_TOPK_DIAG_FORCE_FAULT") ++ proc, res, banner = _isolated_spawn( ++ "trap", ++ drop_env=("MERLIN_TOPK_LARGE_FAULT_INJECT", "MERLIN_TOPK_LARGE_FAULT_KIND"), ++ ) ++ text = proc.stdout + proc.stderr ++ assert banner is not None and banner[2] == 1, "the trap child ran another build" ++ assert "ISOLATED_NO_TRAP" not in proc.stdout, ( ++ f"the forced X3 fault in CTA {cta} did NOT trap: the child finished " ++ f"its launch normally. Result: {res}" ++ ) ++ assert proc.returncode != 0, "the trap child exited 0" ++ assert re.search( ++ r"unspecified launch failure|illegal instruction|device-side|" ++ r"CUDA error|AcceleratorError|cudaErrorLaunchFailure", ++ text, ++ ), f"the child died without a CUDA error:\n{_isolated_tail(text)}" ++ # The parent survives: its own context still runs kernels. ++ x = torch.arange(4096, dtype=torch.float32, device="cuda") ++ assert float(x.sum().item()) == 4095 * 4096 / 2 ++ torch.accelerator.synchronize() ++ ++ ++@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") ++def test_isolated_legacy_repeated_rows_membership() -> None: ++ """Repeated-row control on the LEGACY large tail. ++ ++ Run in its own process with MERLIN_TOPK_DET_LARGE=0 on a NORMAL build. ++ The deterministic path's delayed-publication cell cannot cover this arm ++ (a disagreeing pivot under the diagnostic delay would write out of bounds ++ in the legacy tail), so this is the legacy tail at production timing: ++ repeated large rows, every row held to policy-correct membership. There is ++ no failing control -- the pre-repair legacy tail is not safe to run. ++ """ ++ if _largedet_env_on("MERLIN_TOPK_DET_LARGE", True): ++ pytest.skip( ++ "MERLIN_TOPK_DET_LARGE on: this is the legacy-tail control. Run it " ++ "in a fresh process with MERLIN_TOPK_DET_LARGE=0. A skip is not a pass." ++ ) ++ proc, res, banner = _isolated_spawn("legacy", "12") ++ assert banner is None, ( ++ f"diagnostic build {banner}: the legacy control runs on the normal build only" ++ ) ++ assert proc.returncode == 0 and res is not None and "error" not in res, ( ++ f"the legacy child died (rc={proc.returncode}): {(res or {}).get('error')}\n" ++ f"--- child stderr tail ---\n{_isolated_tail(proc.stderr)}" ++ ) ++ assert res["rows_large"] == 0, ( ++ f"the deterministic large path counted {res['rows_large']} rows with " ++ "MERLIN_TOPK_DET_LARGE=0; the legacy tail did not run" ++ ) ++ assert res["faults_large"] == 0 and res["all_minus_one_rows"] == 0, ( ++ f"{res['faults_large']} large faults, {res['all_minus_one_rows']} all -1 " ++ f"rows on the legacy tail: {res['per_shape']}" ++ ) ++ assert res["invalid_rows"] == 0, ( ++ f"{res['invalid_rows']} of {res['rows_launched']} legacy rows fail " ++ f"membership (first: {res['first_invalid']}): {res['per_shape']}" ++ ) ++ ++ ++def _isolated_diag_child(argv: list[str]) -> None: ++ """Child entry point for cells 5 and 6. Prints one ISOLATED_DIAG_RESULT line.""" ++ import json ++ import sys ++ ++ mode = argv[0] ++ result: dict = {"mode": mode} ++ ++ def emit(code: int) -> None: ++ print(_ISOLATED_RESULT_TAG + json.dumps(result), flush=True) ++ sys.stderr.flush() ++ os._exit(code) ++ ++ try: ++ import vllm._custom_ops # noqa: F401 (loads the extensions) ++ ++ if mode == "probe": ++ # One decode row reaches the launcher, so the build's banner (if ++ # any) is printed; it touches no large row, so it cannot trap. ++ with torch.device("cpu"): ++ probe = torch.linspace(0, 1, 4096) ++ _isolated_launch(probe, 512, num_rows=1) ++ emit(0) ++ ++ if mode == "trap": ++ logits, _ = _build_largedet_tie_row(8) ++ _run_largedet_row(logits) # synchronizes: a trap raises here ++ print("ISOLATED_NO_TRAP", flush=True) ++ emit(0) ++ ++ if mode == "publication": ++ iters = int(argv[1]) if len(argv) > 1 else 12 ++ shapes = [ ++ # (distribution, seq_len, num_rows, top_k) ++ ("cross_cta_ties", 262144, 32, 512), ++ ("cross_cta_ties", 131072, 8, 2048), ++ ("narrow_distinct", 65536, 4, 2048), ++ ("all_equal", 40960, 4, 512), ++ ] ++ before = _read_exact_counters() ++ launched = mismatched = minus_one = 0 ++ per_shape = {} ++ for kind, seq_len, num_rows, top_k in shapes: ++ cpu_row = _largedet_distribution(kind, seq_len, top_k, seed=1905) ++ expected = _largedet_reference(cpu_row, top_k) ++ logits = ( ++ cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ ) ++ lengths = torch.full( ++ (num_rows,), seq_len, dtype=torch.int32, device="cuda" ++ ) ++ bad = neg = 0 ++ for _ in range(iters): ++ got = _largedet_launch(logits, lengths, top_k).cpu().to(torch.int64) ++ for r in range(num_rows): ++ if bool((got[r] == -1).all()): ++ neg += 1 ++ if not torch.equal(got[r], expected): ++ bad += 1 ++ launched += num_rows ++ per_shape[f"{kind}/{seq_len}x{num_rows}/k{top_k}"] = { ++ "mismatched": bad, ++ "all_minus_one": neg, ++ "rows": iters * num_rows, ++ } ++ mismatched += bad ++ minus_one += neg ++ after = _read_exact_counters() ++ result.update( ++ rows_launched=launched, ++ mismatched_rows=mismatched, ++ all_minus_one_rows=minus_one, ++ rows_large=after[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS], ++ faults_large=after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT], ++ per_shape=per_shape, ++ ) ++ emit(0) ++ ++ if mode == "legacy": ++ # Repeated large rows on the LEGACY tail ++ # (MERLIN_TOPK_DET_LARGE=0, normal build). Raw order and the tie ++ # choice are not canonical there, so each row is held to ++ # policy-correct membership (_assert_valid_topk_membership); the ++ # reuse race the closing barrier repairs would show as ++ # out-of-range/duplicate indices, below-pivot picks or all -1 rows. ++ iters = int(argv[1]) if len(argv) > 1 else 12 ++ shapes = [ ++ ("cross_cta_ties", 262144, 32, 512), ++ ("cross_cta_ties", 131072, 8, 2048), ++ ("narrow_distinct", 65536, 4, 2048), ++ ("all_equal", 40960, 4, 512), ++ ("no_cutoff_tie", 262144, 8, 2048), ++ ] ++ before = _read_exact_counters() ++ launched = invalid = minus_one = 0 ++ per_shape = {} ++ first_bad = None ++ for kind, seq_len, num_rows, top_k in shapes: ++ cpu_row = _largedet_distribution(kind, seq_len, top_k, seed=1906) ++ logits = ( ++ cpu_row.to("cuda").unsqueeze(0).repeat(num_rows, 1).contiguous() ++ ) ++ lengths = torch.full( ++ (num_rows,), seq_len, dtype=torch.int32, device="cuda" ++ ) ++ bad = neg = 0 ++ for it in range(iters): ++ got = _largedet_launch(logits, lengths, top_k).cpu().to(torch.int64) ++ for r in range(num_rows): ++ if bool((got[r] == -1).all()): ++ neg += 1 ++ try: ++ with torch.device("cpu"): ++ _assert_valid_topk_membership( ++ cpu_row, got[r], top_k, f"{kind} it{it} r{r}" ++ ) ++ except AssertionError as exc: ++ bad += 1 ++ if first_bad is None: ++ first_bad = str(exc)[:300] ++ launched += num_rows ++ per_shape[f"{kind}/{seq_len}x{num_rows}/k{top_k}"] = { ++ "invalid": bad, ++ "all_minus_one": neg, ++ "rows": iters * num_rows, ++ } ++ invalid += bad ++ minus_one += neg ++ after = _read_exact_counters() ++ result.update( ++ rows_launched=launched, ++ invalid_rows=invalid, ++ all_minus_one_rows=minus_one, ++ first_invalid=first_bad, ++ rows_large=after[_SLOT_LARGE_ROWS] - before[_SLOT_LARGE_ROWS], ++ faults_large=after[_SLOT_LARGE_FAULT] - before[_SLOT_LARGE_FAULT], ++ per_shape=per_shape, ++ ) ++ emit(0) ++ ++ result["error"] = f"unknown mode {mode!r}" ++ emit(2) ++ except Exception as exc: # noqa: BLE001 -- report, then die ++ result["error"] = f"{type(exc).__name__}: {exc}"[:800] ++ emit(3) ++ ++ ++if __name__ == "__main__": ++ import sys ++ ++ if len(sys.argv) >= 3 and sys.argv[1] == "--isolated-diag-child": ++ _isolated_diag_child(sys.argv[2:]) ++ raise SystemExit("usage: test_top_k_per_row.py --isolated-diag-child [args]") +-- +2.43.0 + diff --git a/patches/0047-worker-e2-identity-and-rank-tagged-e2-toggles-startu.patch b/patches/0047-worker-e2-identity-and-rank-tagged-e2-toggles-startu.patch new file mode 100644 index 0000000..2b54c32 --- /dev/null +++ b/patches/0047-worker-e2-identity-and-rank-tagged-e2-toggles-startu.patch @@ -0,0 +1,1358 @@ +From d7e28d037aa462c6d7e53a22478767759710995b Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 08:41:02 -0400 +Subject: [PATCH 047/103] worker: e2-identity and rank-tagged e2-toggles + startup lines + +Each worker logs its launch identity and, after warm-up and capture, one +e2-toggles line with the effective toggle state read back from the C++ +extensions through two new tensor-free query ops. Both extensions must be +rebuilt. +--- + .../moe/marlin_moe_wna16/ops.cu | 31 +- + csrc/libtorch_stable/ops.h | 4 + + csrc/libtorch_stable/topk.cu | 89 +++- + csrc/libtorch_stable/torch_bindings.cpp | 5 + + tests/test_merlin_e2_toggles.py | 420 +++++++++++++++++ + vllm/v1/worker/e2_identity.py | 148 ++++++ + vllm/v1/worker/e2_toggles.py | 444 ++++++++++++++++++ + vllm/v1/worker/gpu/model_runner.py | 7 +- + vllm/v1/worker/gpu_worker.py | 13 + + 9 files changed, 1154 insertions(+), 7 deletions(-) + create mode 100644 tests/test_merlin_e2_toggles.py + create mode 100644 vllm/v1/worker/e2_identity.py + create mode 100644 vllm/v1/worker/e2_toggles.py + +diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +index 2c6ad76..0a6f5de 100644 +--- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu ++++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +@@ -33,6 +33,7 @@ + #include + + #include "libtorch_stable/torch_utils.h" ++#include + #include + + namespace { +@@ -53,14 +54,31 @@ namespace { + // This std::getenv is the SINGLE source of truth. vllm/envs.py carries a raw + // passthrough entry for validate_environ and the compile-factor hash only; + // nothing in Python may parse this name a second time. ++// ++// The initialiser also ++// publishes the value it cached into g_cached_full_k (-1 = not initialised ++// yet; nothing else writes it), and marlin_moe_full_k_state() below only ++// LOADS it -- asking never fills the cache. Host code only. + // ------------------------------------------------------------------------- ++std::atomic g_cached_full_k{-1}; ++ + bool merlin_full_k() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_FULL_K"); +- return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ const bool value = ++ (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ g_cached_full_k.store(value ? 1 : 0, std::memory_order_release); ++ return value; + }(); + return enabled; + } ++ ++// MERLIN_FULL_K as this process cached it: -1 = not initialised yet (no Marlin ++// MoE launch has happened in this process), 0 = off, 1 = on. Tensor-free, no ++// CUDA call. ++int64_t marlin_moe_full_k_state() { ++ return static_cast(g_cached_full_k.load(std::memory_order_acquire)); ++} + } // namespace + + #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ +@@ -736,3 +754,14 @@ torch::stable::Tensor moe_wna16_marlin_gemm( + STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { + m.impl("moe_wna16_marlin_gemm", TORCH_BOX(&moe_wna16_marlin_gemm)); + } ++ ++// v2.2 -- defined and registered in this TU so the query exists exactly when ++// the static it reads does. Tensor-free, so CompositeExplicitAutograd (a CUDA ++// registration could never be dispatched to without a tensor argument). ++STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { ++ m.def("marlin_moe_full_k_state() -> int"); ++} ++ ++STABLE_TORCH_LIBRARY_IMPL(_moe_C, CompositeExplicitAutograd, m) { ++ m.impl("marlin_moe_full_k_state", TORCH_BOX(&marlin_moe_full_k_state)); ++} +diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h +index 1156801..32d0673 100644 +--- a/csrc/libtorch_stable/ops.h ++++ b/csrc/libtorch_stable/ops.h +@@ -482,6 +482,10 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset, + // reader never has to guess it from an environment variable. + int64_t persistent_topk_exact_counters_slots(); + ++// The persistent_topk switches as this ++// process cached them, packed; see topk.cu for the layout. ++int64_t persistent_topk_toggle_state(); ++ + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index b52a34b..3c90c25 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -43,6 +43,29 @@ namespace { + // This std::getenv is the SINGLE source of truth; vllm/envs.py carries a raw + // passthrough for validate_environ and the compile-factor hash only. + // ------------------------------------------------------------------------- ++// ------------------------------------------------------------------------- ++// A READ-ONLY MIRROR of the ++// switches as this process CACHED them, for persistent_topk_toggle_state(). ++// ++// Each switch below is a function-local static filled once by its lambda. The ++// lambda now also stores the value it cached into the matching atomic here ++// (publish_cached), and nothing else ever writes these atomics, so a mirror ++// either says "not initialised yet" (-1) or holds exactly the cached value. ++// The query only LOADS them: it never calls a getter, so asking cannot fill a ++// cache from an environment that differs from the one the kernel saw. The ++// parse expressions are unchanged; this is host code only. ++// ------------------------------------------------------------------------- ++std::atomic g_cached_tie_recent{-1}; ++std::atomic g_cached_tie_det{-1}; ++std::atomic g_cached_exact_medium{-1}; ++std::atomic g_cached_det_large{-1}; ++std::atomic g_cached_exact_decode{-1}; ++ ++bool publish_cached(std::atomic& mirror, bool value) { ++ mirror.store(value ? 1 : 0, std::memory_order_release); ++ return value; ++} ++ + // MERLIN_TIE_RECENT -- the DIRECTION of the deterministic tie-break above. + // Measured at 12k tokens on the reference host: oldest-wins costs -7.9 % decode, + // most-recent-wins is +1.5 %. Both deterministic. Also the direction +@@ -52,7 +75,9 @@ namespace { + bool tie_recent_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TIE_RECENT"); +- return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ return publish_cached( ++ g_cached_tie_recent, ++ (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false); + }(); + return enabled; + } +@@ -60,7 +85,9 @@ bool tie_recent_enabled() { + bool det_ties_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TIE_DET"); +- return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ return publish_cached( ++ g_cached_tie_det, ++ (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false); + }(); + return enabled; + } +@@ -68,7 +95,9 @@ bool det_ties_enabled() { + bool exact_medium_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TOPK_EXACT_MEDIUM"); +- return (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0); ++ return publish_cached( ++ g_cached_exact_medium, ++ (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0)); + }(); + return enabled; + } +@@ -86,7 +115,9 @@ bool exact_medium_enabled() { + bool det_large_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TOPK_DET_LARGE"); +- return (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0); ++ return publish_cached( ++ g_cached_det_large, ++ (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0)); + }(); + return enabled; + } +@@ -103,7 +134,9 @@ bool det_large_enabled() { + bool exact_decode_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TOPK_EXACT_DECODE"); +- return (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false; ++ return publish_cached( ++ g_cached_exact_decode, ++ (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false); + }(); + return enabled; + } +@@ -871,3 +904,49 @@ int64_t persistent_topk_exact_counters_slots() { + return 0; + #endif + } ++ ++// The persistent_topk ++// switches AS THIS PROCESS CACHED THEM -- the values every launch of this ++// process uses -- read from the mirrors above. Tensor-free, no CUDA call, no ++// allocation, never initialises a cache (see publish_cached), so it is safe ++// from anywhere and asking changes nothing. ++// ++// Packed into one int64 so the op needs no tensor and no dispatch key. ++// 2 bits per field, field i at bits [2i, 2i+1]: 0 = off, 1 = on, ++// 2 = the static has not been initialised in this process yet. ++// field 0 MERLIN_TIE_DET (det_ties_enabled) ++// field 1 MERLIN_TIE_RECENT as parsed (tie_recent_enabled) ++// field 2 effective direction, 1 = most-recent (effective_tie_recent) ++// field 3 MERLIN_TOPK_EXACT_MEDIUM (exact_medium_enabled) ++// field 4 MERLIN_TOPK_EXACT_DECODE (exact_decode_enabled) ++// field 5 MERLIN_TOPK_DET_LARGE (det_large_enabled) ++// bits [48, 55] = field count (6), bits [56, 62] = layout version (1). A ++// reader must check both and refuse any other layout rather than guess. ++int64_t persistent_topk_toggle_state() { ++#ifndef USE_ROCM ++ const auto field = [](int cached) -> int64_t { ++ return cached < 0 ? 2 : (cached != 0 ? 1 : 0); ++ }; ++ const int det = g_cached_tie_det.load(std::memory_order_acquire); ++ const int recent = g_cached_tie_recent.load(std::memory_order_acquire); ++ // The kernel's own rule, from the kernel's own header, on the cached values. ++ const int effective = ++ (det < 0 || recent < 0) ++ ? -1 ++ : (vllm::persistent::effective_tie_recent(det != 0, recent != 0) ++ ? 1 ++ : 0); ++ int64_t state = 0; ++ state |= field(det) << 0; ++ state |= field(recent) << 2; ++ state |= field(effective) << 4; ++ state |= field(g_cached_exact_medium.load(std::memory_order_acquire)) << 6; ++ state |= field(g_cached_exact_decode.load(std::memory_order_acquire)) << 8; ++ state |= field(g_cached_det_large.load(std::memory_order_acquire)) << 10; ++ state |= int64_t{6} << 48; ++ state |= int64_t{1} << 56; ++ return state; ++#else ++ return 0; // layout version 0: no persistent_topk on ROCm ++#endif ++} +diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp +index 3b23bc1..2ecc945 100644 +--- a/csrc/libtorch_stable/torch_bindings.cpp ++++ b/csrc/libtorch_stable/torch_bindings.cpp +@@ -591,6 +591,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + ops.def( + "persistent_topk_exact_counters(Tensor! out, int reset, int sync) -> ()"); + ops.def("persistent_topk_exact_counters_slots() -> int"); ++ // v2.2 -- the switches as this process cached them (topk.cu for the layout). ++ ops.def("persistent_topk_toggle_state() -> int"); + + #ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( +@@ -949,6 +951,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { + // NotImplementedError "no fallback function is registered"). + ops.impl("persistent_topk_exact_counters_slots", + TORCH_BOX(&persistent_topk_exact_counters_slots)); ++ // Tensor-free for the same reason as the slot count above. ++ ops.impl("persistent_topk_toggle_state", ++ TORCH_BOX(&persistent_topk_toggle_state)); + } + + // Cache ops +diff --git a/tests/test_merlin_e2_toggles.py b/tests/test_merlin_e2_toggles.py +new file mode 100644 +index 0000000..3cd3c1f +--- /dev/null ++++ b/tests/test_merlin_e2_toggles.py +@@ -0,0 +1,420 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""CPU tests for vllm/v1/worker/e2_toggles.py. ++ ++Parsing parity with the C++ (glibc atoi) and Python (int()) readers including ++odd spellings; the effective map for the toggles-ON and knob-absent boots; ++once-only emission; emission with the counter log OFF; deferral while a C++ ++static is uninitialised; the n-a paths; and a line the smoke can tokenise. ++ ++ pytest --noconftest tests/test_merlin_e2_toggles.py ++""" ++ ++import ctypes ++import ctypes.util ++import logging ++import os ++import sys ++ ++import pytest ++ ++from vllm.v1.worker import e2_toggles as t ++ ++ ++# ------------------------------------------------------------ atoi parity ---- ++ATOI_CASES = [ ++ (b"1", 1), ++ (b"0", 0), ++ (b"00", 0), ++ (b"01", 1), ++ (b"true", 0), # FULL_K=true reads as OFF (the release note's warning) ++ (b"yes", 0), ++ (b"on", 0), ++ (b" 1", 1), ++ (b"\t\n\v\f\r1", 1), ++ (b"+1", 1), ++ (b"-1", -1), ++ (b"-0", 0), ++ (b"+ 1", 0), # sign then space: no digits ++ (b"1abc", 1), ++ (b"1 ", 1), ++ (b"0x1", 0), ++ (b"1.9", 1), ++ (b"2147483647", 2147483647), ++ (b"2147483648", -2147483648), # (int) keeps the low 32 bits ++ (b"4294967296", 0), # 2^32 -> 0: set but reads OFF ++ (b"4294967297", 1), ++ (b"99999999999999999999", -1), # strtol clamps to LONG_MAX -> low 32 = -1 ++ (b"-99999999999999999999", 0), # LONG_MIN -> low 32 = 0 ++ (b"\xc2\xa01", 0), # NBSP is not C-locale whitespace ++ (b"\xef\xbc\x91", 0), # fullwidth digit one: not a C digit ++ (b"", 0), ++] ++ ++ ++@pytest.mark.parametrize("raw,want", ATOI_CASES) ++def test_c_atoi_table(raw, want): ++ assert t.c_atoi(raw) == want ++ ++ ++def _glibc(): ++ if not sys.platform.startswith("linux"): ++ return None ++ name = ctypes.util.find_library("c") ++ if not name: ++ return None ++ libc = ctypes.CDLL(name) ++ gnu = getattr(libc, "gnu_get_libc_version", None) ++ if gnu is None: ++ return None ++ libc.atoi.restype = ctypes.c_int ++ libc.atoi.argtypes = [ctypes.c_char_p] ++ return libc ++ ++ ++@pytest.mark.skipif(_glibc() is None, reason="parity against the real atoi needs glibc") ++@pytest.mark.parametrize("raw,want", ATOI_CASES) ++def test_c_atoi_matches_glibc(raw, want): ++ assert t.c_atoi(raw) == _glibc().atoi(raw) ++ ++ ++@pytest.mark.parametrize( ++ "raw,unset_value,want", ++ [ ++ (None, False, False), ++ (b"", False, False), # empty = unset for the default-OFF getters ++ (None, True, True), ++ (b"", True, True), # EXACT_MEDIUM= / DET_LARGE= keep the default ON ++ (b"0", True, False), ++ (b"1", False, True), ++ (b"true", False, False), ++ (b"true", True, False), # a non-numeric value turns a default-ON switch OFF ++ (b"4294967296", False, False), ++ (b"-1", False, True), ++ ], ++) ++def test_cxx_switch(raw, unset_value, want): ++ assert t.cxx_switch(raw, unset_value) is want ++ ++ ++@pytest.mark.parametrize( ++ "raw,want", ++ [ ++ (None, "0"), ++ ("0", "0"), ++ ("1", "1"), ++ ("2", "1"), ++ (" 1 ", "1"), ++ ("+1", "1"), ++ ("-0", "0"), ++ ("1_0", "1"), ++ ("１", "1"), # Python int() accepts non-ASCII digits ++ ("true", "invalid"), # envs.py raises: the engine would not boot ++ ("", "invalid"), ++ ("0x1", "invalid"), ++ ], ++) ++def test_qsa_sort_value(raw, want): ++ assert t.qsa_sort_value(raw) == want ++ ++ ++def test_qsa_sort_value_matches_envs_reader(monkeypatch): ++ """The mirror IS the reader: same int() call as vllm/envs.py:2372.""" ++ import vllm.envs as envs ++ ++ for raw in ("0", "1", " 1 ", "1_0", "１", "-0"): ++ monkeypatch.setenv("MERLIN_QSA_SORT", raw) ++ assert t.qsa_sort_value(raw) == ("1" if envs.MERLIN_QSA_SORT else "0") ++ for raw in ("true", ""): ++ monkeypatch.setenv("MERLIN_QSA_SORT", raw) ++ with pytest.raises(ValueError): ++ _ = envs.MERLIN_QSA_SORT ++ assert t.qsa_sort_value(raw) == "invalid" ++ ++ ++def test_effective_tie_recent(): ++ assert t.effective_tie_recent(False, False) is True ++ assert t.effective_tie_recent(False, True) is True ++ assert t.effective_tie_recent(True, False) is False ++ assert t.effective_tie_recent(True, True) is True ++ ++ ++# ------------------------------------------------------- effective maps ---- ++def _env(values): ++ def cxx(name): ++ v = values.get(name) ++ return None if v is None else v.encode() ++ ++ return cxx, values.get ++ ++ ++ON = { ++ "MERLIN_FULL_K": "1", ++ "MERLIN_QSA_SORT": "1", ++ "MERLIN_TIE_DET": "1", ++ "MERLIN_TIE_RECENT": "1", ++} ++ON_MAP = { ++ "full_k": "1", ++ "qsa_sort": "1", ++ "tie_det": "1", ++ "tie_recent": "1", ++ "tie_recent_effective": "1", ++ "exact_medium": "1", ++ "exact_decode": "0", ++ "det_large": "1", ++} ++ABSENT_MAP = { ++ "full_k": "0", ++ "qsa_sort": "0", ++ "tie_det": "0", ++ "tie_recent": "0", ++ "tie_recent_effective": "1", ++ "exact_medium": "1", ++ "exact_decode": "0", ++ "det_large": "1", ++} ++ ++ ++def test_effective_map_toggles_on(): ++ raw, eff = t.environ_view(*_env(ON)) ++ assert eff == ON_MAP ++ assert raw["MERLIN_TOPK_EXACT_MEDIUM"] == "unset" ++ assert raw["MERLIN_FULL_K"] == "1" ++ ++ ++def test_effective_map_knob_absent(): ++ raw, eff = t.environ_view(*_env({})) ++ assert eff == ABSENT_MAP ++ assert set(raw.values()) == {"unset"} ++ ++ ++def test_effective_map_odd_spellings(): ++ _, eff = t.environ_view( ++ *_env( ++ { ++ "MERLIN_FULL_K": "true", # C++ atoi -> OFF ++ "MERLIN_TIE_DET": " 1", ++ "MERLIN_TIE_RECENT": "0", ++ "MERLIN_TOPK_EXACT_MEDIUM": "", # empty -> default ON ++ "MERLIN_TOPK_DET_LARGE": "off", # atoi("off") == 0 -> OFF ++ "MERLIN_QSA_SORT": "true", ++ } ++ ) ++ ) ++ assert eff["full_k"] == "0" ++ assert eff["tie_det"] == "1" ++ assert eff["tie_recent_effective"] == "0" # TIE_DET=1 TIE_RECENT=0 = oldest-wins ++ assert eff["exact_medium"] == "1" ++ assert eff["det_large"] == "0" ++ assert eff["qsa_sort"] == "invalid" ++ ++ ++def test_encode_raw(): ++ assert t.encode_raw(None) == "unset" ++ assert t.encode_raw(b"1") == "1" ++ assert t.encode_raw(b"") == '""' ++ assert t.encode_raw(b"unset") == '"unset"' ++ assert t.encode_raw(b"1 2") == '"1\\x202"' ++ assert t.encode_raw(b'a"b') == '"a\\x22b"' ++ for v in (b"", b" x ", b"\xff", b"a\tb"): ++ assert " " not in t.encode_raw(v) ++ ++ ++# ----------------------------------------------------------- cpp decode ---- ++def _pack(fields, version=1, count=6): ++ state = 0 ++ for i, f in enumerate(fields): ++ state |= f << (2 * i) ++ return state | (count << 48) | (version << 56) ++ ++ ++def test_decode_topk_state(): ++ assert t.decode_topk_state(_pack([1, 1, 1, 1, 0, 1])) == { ++ "tie_det": "1", ++ "tie_recent": "1", ++ "tie_recent_effective": "1", ++ "exact_medium": "1", ++ "exact_decode": "0", ++ "det_large": "1", ++ } ++ assert t.decode_topk_state(_pack([2, 2, 2, 1, 2, 2]))["tie_det"] == "uninit" ++ assert t.decode_topk_state(_pack([3, 0, 0, 0, 0, 0]))["tie_det"] == "invalid" ++ assert isinstance(t.decode_topk_state(_pack([0] * 6, version=2)), str) ++ assert isinstance(t.decode_topk_state(_pack([0] * 6, count=5)), str) ++ assert isinstance(t.decode_topk_state(0), str) # ROCm build: version 0 ++ ++ ++def test_decode_full_k_state(): ++ assert [t.decode_full_k_state(v) for v in (-1, 0, 1, 7)] == [ ++ "uninit", ++ "0", ++ "1", ++ "invalid", ++ ] ++ ++ ++# ------------------------------------------------------------- emission ---- ++CPP_ON = {k: ON_MAP[k] for k in t.CPP_ORDER} ++ ++ ++class _Capture(logging.Handler): ++ def __init__(self): ++ super().__init__() ++ self.lines = [] ++ ++ def emit(self, record): ++ self.lines.append(record.getMessage()) ++ ++ ++@pytest.fixture ++def capture(): ++ t._reset_for_tests() ++ handler = _Capture() ++ t.logger.addHandler(handler) ++ old = t.logger.level ++ t.logger.setLevel(logging.INFO) ++ yield handler ++ t.logger.removeHandler(handler) ++ t.logger.setLevel(old) ++ t._reset_for_tests() ++ ++ ++def _kw(cpp=CPP_ON, env=None, device=None, qsa="1", tag="PP1_TP0"): ++ return dict( ++ tag_fn=lambda: tag, ++ getenv_fn=lambda: (_env(ON if env is None else env)[0], "libc-getenv"), ++ cpp_fn=lambda: cpp, ++ device_fn=lambda: device or {"switch": "1", "switch_large": "1"}, ++ qsa_fn=lambda: qsa, ++ ) ++ ++ ++def parse(line): ++ """How the startup smoke tokenises the line (sections; key=value tokens).""" ++ tokens = line.split() ++ assert tokens[0].startswith("e2-toggles[") and tokens[0].endswith("]:") ++ out, section = {"tag": tokens[0][len("e2-toggles[") : -2]}, "raw" ++ for tok in tokens[1:]: ++ if tok.endswith(":"): ++ section = tok[:-1] ++ continue ++ key, sep, value = tok.partition("=") ++ out.setdefault(section, {})[key if sep else tok] = value if sep else "" ++ return out ++ ++ ++def test_emits_once_per_process(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ line = t.arm(**_kw()) ++ assert line and capture.lines == [line] ++ assert t.maybe_emit(**_kw()) == "" ++ assert t.maybe_emit(force=True, **_kw()) == "" ++ assert t.arm(**_kw()) == "" ++ assert len([m for m in capture.lines if m.startswith("e2-toggles[")]) == 1 ++ ++ ++def test_not_before_arm(capture): ++ assert t.maybe_emit(**_kw()) == "" ++ assert capture.lines == [] ++ ++ ++def test_fires_with_counter_log_off(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ for value in (None, "0"): ++ t._reset_for_tests() ++ if value is None: ++ monkeypatch.delenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", raising=False) ++ else: ++ monkeypatch.setenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", value) ++ assert t.arm(**_kw()).startswith("e2-toggles[PP1_TP0]: pid=") ++ ++ ++def test_line_contents_toggles_on(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ p = parse(t.arm(**_kw())) ++ assert p["tag"] == "PP1_TP0" ++ assert p["raw"]["pid"] == str(os.getpid()) ++ assert p["raw"]["MERLIN_FULL_K"] == "1" ++ assert p["raw"]["MERLIN_TOPK_DET_LARGE"] == "unset" ++ eff = dict(p["effective"]) ++ assert eff.pop("source") == "libc-getenv" ++ assert eff == ON_MAP ++ cpp = dict(p["cpp-cache"]) ++ assert cpp.pop("source") == "cpp-cache" ++ assert cpp == CPP_ON ++ assert p["python"] == {"qsa_sort": "1", "source": "qsa_indexer._QSA_SORT"} ++ assert p["device"] == { ++ "switch": "1", ++ "switch_large": "1", ++ "source": "counter-slots-8-9", ++ } ++ ++ ++def test_line_knob_absent(capture, monkeypatch): ++ for k in t.RAW_ORDER: ++ monkeypatch.delenv(k, raising=False) ++ cpp = {k: ABSENT_MAP[k] for k in t.CPP_ORDER} ++ p = parse(t.arm(**_kw(cpp=cpp, env={}, qsa="0"))) ++ eff = dict(p["effective"]) ++ eff.pop("source") ++ assert eff == ABSENT_MAP ++ ++ ++def test_cpp_op_absent_is_n_a_and_emits_at_once(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ line = t.arm(**_kw(cpp="op-absent:_C.persistent_topk_toggle_state")) ++ p = parse(line) ++ assert p["cpp-cache"] == { ++ "n-a": "", ++ "reason": "op-absent:_C.persistent_topk_toggle_state", ++ } ++ ++ ++def test_defers_while_a_static_is_uninit(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ pending = dict(CPP_ON, full_k="uninit") ++ assert t.arm(**_kw(cpp=pending)) == "" ++ assert t.maybe_emit(**_kw(cpp=pending)) == "" ++ notes = [m for m in capture.lines if m.startswith("e2-toggles deferred on PP1_TP0")] ++ assert len(notes) == 1 and "full_k" in notes[0] ++ assert not any(m.startswith("e2-toggles[") for m in capture.lines) ++ line = t.maybe_emit(**_kw()) # the next step: every static filled ++ assert parse(line)["cpp-cache"]["full_k"] == "1" ++ ++ ++def test_forced_at_shutdown_while_uninit(capture, monkeypatch): ++ for k, v in ON.items(): ++ monkeypatch.setenv(k, v) ++ pending = dict(CPP_ON, full_k="uninit") ++ t.arm(**_kw(cpp=pending)) ++ line = t.maybe_emit(force=True, **_kw(cpp=pending)) ++ assert parse(line)["cpp-cache"]["full_k"] == "uninit" ++ ++ ++def test_never_raises(capture): ++ def boom(): ++ raise RuntimeError("no groups") ++ ++ kw = _kw() ++ kw["cpp_fn"] = boom ++ assert t.arm(**kw) == "" ++ assert capture.lines ++ assert capture.lines[-1].startswith("e2-toggles[PP1_TP0]: unavailable") ++ assert t.maybe_emit(**_kw()) == "" # an unavailable line is still the one line ++ ++ ++def test_default_sources_do_not_raise_on_cpu(monkeypatch): ++ """The real cpp/device/qsa probes on a CPU build: n-a, never an exception.""" ++ assert isinstance(t.cpp_cache(), (str, dict)) ++ assert set(t.device_switches()) == {"switch", "switch_large"} ++ assert t.qsa_sort_live() in ("0", "1", "n-a") ++ read, source = t._libc_getenv() ++ assert source in ("libc-getenv", "os-environ") ++ monkeypatch.setenv("MERLIN_TIE_DET", "7") ++ assert read("MERLIN_TIE_DET") == b"7" +diff --git a/vllm/v1/worker/e2_identity.py b/vllm/v1/worker/e2_identity.py +new file mode 100644 +index 0000000..e457b09 +--- /dev/null ++++ b/vllm/v1/worker/e2_identity.py +@@ -0,0 +1,148 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Per-rank launch-identity log line (log-only; changes no behaviour). ++ ++After model loading -- and after the PLE offload placeholders are armed -- every ++GPU worker rank logs one line per model role it holds:: ++ ++ e2-identity[PP

_TP]: ple_backend=@ ple_method= ++ ple_weight_dtype= breakable= compile_mode= ++ cudagraph_mode= role= ++ ++(one line in the log). ``role=draft`` is logged only where a draft model is ++loaded (the last PP stage). Every value is read from the live objects: ++ ++* ``ple_backend``: the class of each PLE layer's ``ple_embedding`` module, and ++ the transport that serves its rows on this rank, read from the placeholder's ++ armed state -- ``host-mapped-pull`` (``_is_cpu_offloaded`` with a host output ++ pointer), ``ipc-push`` (offloaded, no host pointer) or ``local:`` (the ++ rank owns the table). ++* ``ple_method``: the embedding's live ``quant_method``, or, on an offloaded ++ placeholder, the ``_offload_quant_method`` the rank dequantizes with. ++* ``ple_weight_dtype``: the table's dtype where the rank owns it; on an ++ offloaded placeholder, the dtype of the armed transfer buffer, into which the ++ offload process copies the table's rows unchanged ++ (``get_offload_output_dtype``). ++* ``breakable``: ``is_breakable_cudagraph_enabled()``, the switch the capture ++ path itself reads. ++* ``compile_mode`` / ``cudagraph_mode``: the role's ``CompilationConfig`` ++ (config level, i.e. before KV-cache-time capture-mode resolution). ++ ++A role with no PLE layer reports ``none`` for the three PLE fields; several ++distinct values on one rank are joined with ``,``. ++""" ++ ++from typing import Any ++ ++from torch import nn ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_NONE = "none" ++ ++ ++def _join(values: set[str]) -> str: ++ return ",".join(sorted(values)) if values else _NONE ++ ++ ++def _enum_name(value: Any) -> str: ++ if value is None: ++ return _NONE ++ return str(getattr(value, "name", value)) ++ ++ ++def _ple_embeddings(model: nn.Module | None) -> list[nn.Module]: ++ """The ``ple_embedding`` module of every PLE layer in ``model``.""" ++ if model is None: ++ return [] ++ found = [] ++ for module in model.modules(): ++ embedding = module.__dict__.get("_modules", {}).get("ple_embedding") ++ if isinstance(embedding, nn.Module): ++ found.append(embedding) ++ return found ++ ++ ++def _transport(embedding: nn.Module) -> str: ++ if getattr(embedding, "_is_cpu_offloaded", False): ++ if getattr(embedding, "_host_output_ptr", None) is not None: ++ return "host-mapped-pull" ++ return "ipc-push" ++ weight = getattr(getattr(embedding, "ngram_embedding", None), "weight", None) ++ if weight is not None: ++ return f"local:{weight.device.type}" ++ return "local:unknown" ++ ++ ++def _method(embedding: nn.Module) -> str: ++ table = getattr(embedding, "ngram_embedding", None) ++ method = getattr(table, "quant_method", None) ++ if method is None: ++ method = getattr(embedding, "_offload_quant_method", None) ++ return type(method).__name__ if method is not None else _NONE ++ ++ ++def _weight_dtype(embedding: nn.Module) -> str: ++ weight = getattr(getattr(embedding, "ngram_embedding", None), "weight", None) ++ if weight is not None: ++ return str(weight.dtype) ++ if getattr(embedding, "_is_cpu_offloaded", False): ++ buffer = getattr(embedding, "_gpu_output_buffer", None) ++ if buffer is not None: ++ return str(buffer.dtype) ++ return "unknown" ++ ++ ++def _rank_tag() -> str: ++ from vllm.distributed.parallel_state import get_pp_group, get_tp_group ++ ++ return f"PP{get_pp_group().rank_in_group}_TP{get_tp_group().rank_in_group}" ++ ++ ++def _draft(model_runner: Any) -> tuple[nn.Module | None, Any]: ++ """The loaded draft model and its config (V2 speculator or V1 drafter).""" ++ for name in ("speculator", "drafter"): ++ holder = getattr(model_runner, name, None) ++ model = getattr(holder, "model", None) ++ if isinstance(model, nn.Module): ++ return model, getattr(holder, "vllm_config", None) ++ return None, None ++ ++ ++def identity_line( ++ tag: str, model: nn.Module | None, vllm_config: Any, role: str ++) -> str: ++ from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled ++ ++ embeddings = _ple_embeddings(model) ++ backends = {f"{type(e).__name__}@{_transport(e)}" for e in embeddings} ++ compilation = getattr(vllm_config, "compilation_config", None) ++ return ( ++ f"e2-identity[{tag}]: " ++ f"ple_backend={_join(backends)} " ++ f"ple_method={_join({_method(e) for e in embeddings})} " ++ f"ple_weight_dtype={_join({_weight_dtype(e) for e in embeddings})} " ++ f"breakable={'ON' if is_breakable_cudagraph_enabled() else 'OFF'} " ++ f"compile_mode={_enum_name(getattr(compilation, 'mode', None))} " ++ f"cudagraph_mode={_enum_name(getattr(compilation, 'cudagraph_mode', None))} " ++ f"role={role}" ++ ) ++ ++ ++def log_launch_identity(model_runner: Any, vllm_config: Any) -> None: ++ """Log the identity line(s) for this rank. Never raises.""" ++ tag = "PP?_TP?" ++ try: ++ tag = _rank_tag() ++ roles = [("target", model_runner.get_model(), vllm_config)] ++ draft_model, draft_config = _draft(model_runner) ++ if draft_model is not None: ++ roles.append(("draft", draft_model, draft_config or vllm_config)) ++ for role, model, config in roles: ++ logger.info("%s", identity_line(tag, model, config, role)) ++ except Exception as error: # log-only: identity must never break a boot ++ logger.warning( ++ "e2-identity[%s]: unavailable (%s: %s)", tag, type(error).__name__, error ++ ) +diff --git a/vllm/v1/worker/e2_toggles.py b/vllm/v1/worker/e2_toggles.py +new file mode 100644 +index 0000000..0ae7ad2 +--- /dev/null ++++ b/vllm/v1/worker/e2_toggles.py +@@ -0,0 +1,444 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Per-rank determinism-toggle line (log-only; changes no behaviour). ++ ++Why this exists: the effective state of the ++four determinism toggles and the three persistent_topk switches must be ++attributable to a RANK and a PID. The kernel's own stderr line ++(``[merlin] persistent_topk tie policy: ...``, topk.cu report_tie_policy_once) ++is printed once per process with no tag, and FULL_K / QSA_SORT have no log ++producer at all. Every GPU worker therefore logs exactly one line:: ++ ++ e2-toggles[PP

_TP]: pid= MERLIN_FULL_K= MERLIN_QSA_SORT= ++ MERLIN_TIE_DET= MERLIN_TIE_RECENT= MERLIN_TOPK_EXACT_MEDIUM= ++ MERLIN_TOPK_EXACT_DECODE= MERLIN_TOPK_DET_LARGE= ++ effective: full_k=<0|1> qsa_sort=<0|1|invalid> tie_det=<0|1> tie_recent=<0|1> ++ tie_recent_effective=<0|1> exact_medium=<0|1> exact_decode=<0|1> ++ det_large=<0|1> source= ++ cpp-cache: > ++ python: qsa_sort=<0|1|n-a> source=qsa_indexer._QSA_SORT ++ device: switch=<0|1|n-a> switch_large=<0|1|n-a> source=counter-slots-8-9 ++ ++(one line in the log). ```` is ``unset``, a plain token, or a double-quoted ++string with every byte outside ``[!-~]`` (and ``"`` / ``\\``) hex-escaped, so ++the line always splits on whitespace. ++ ++WHAT EACH SECTION IS ++* ``effective:`` -- the environment at emission time, parsed with the SAME ++ rules the readers use (C++ ``std::getenv`` + ``std::atoi != 0`` for the six ++ C++ names, ``bool(int(...))`` for QSA_SORT; file:line of each below). The ++ C++ names are read through libc ``getenv`` -- the ``environ`` that ++ ``std::getenv`` reads -- not ``os.environ``, which can diverge from it after ++ an ``os.putenv``. A MIRROR of the parse, not the kernel's state. It exists ++ to report the rule; nothing may branch on it (vllm/envs.py "NOT A PYTHON ++ READ SITE" still holds: this module decides nothing). ++* ``cpp-cache:`` -- the kernel's switches as the C++ statics cached them, from ++ ``torch.ops._C.persistent_topk_toggle_state`` and ++ ``torch.ops._moe_C.marlin_moe_full_k_state`` (query ops that only LOAD a ++ mirror the static initialisers publish; they never fill a cache). ``n-a`` ++ with a reason on a build without those ops -- the extensions this release ++ was first cut against (earlier v2.2 builds, not published) have none. ++* ``python:`` -- ``qsa_indexer._QSA_SORT``, the module constant the QSA sort ++ actually branches on (looked up in ``sys.modules``, never imported here; ++ ``n-a`` if the module was never loaded in this process). ++* ``device:`` -- counter slots [8]/[9], which the counter op fills from the ++ C++ ``exact_medium_enabled()`` / ``det_large_enabled()`` statics ++ (topk.cu persistent_topk_exact_counters). ++ ++WHEN. Once per worker, after ``compile_or_warm_up_model`` has run the MRV2 ++kernel warmup and cudagraph capture (``arm()``), i.e. after this process's ++first persistent_topk and Marlin MoE launches -- NOT at the counter reader's ++first read: an integration boot logged its first counter line BEFORE ++the first tie-policy line in the engine log, because the ++first ``maybe_log`` runs at the head of the warmup's first execute step, before ++that step's forward. When the cpp-cache ops exist and a static is still ++uninitialised at ``arm()``, the line is deferred to the head of the first ++execute step where every static is filled, and forced out at shutdown ++otherwise (with ``uninit`` values, which the smoke FAILs). It does not depend on ++``MERLIN_TOPK_EXACT_COUNTERS_LOG``: it fires with the counter log off. ++ ++WHAT THIS DOES NOT PROVE. Without the cpp-cache ops the C++ state is inferred, ++not read: for FULL_K, TIE_DET, TIE_RECENT and EXACT_DECODE the line is the ++environment at emission time run through the C++ parser's rules, and the only ++C++-derived evidence is ``device:`` (EXACT_MEDIUM and DET_LARGE) plus the ++untagged stderr policy line. The residual gap is an environment mutated ++between the kernel's first ``getenv`` and this read. Nothing in this tree does ++that (by grep of csrc/): no ``setenv``/``putenv``/``unsetenv`` under csrc/, and ++the only runtime ``os.environ.update`` calls in vllm/ are the RunAI streamer's ++fixed ``RUNAI_STREAMER_*`` keys (model_executor/model_loader/ ++runai_streamer_loader.py:73) and the Ray executor's ``override_env_vars`` ++(v1/executor/ray_utils.py:180), which the multiproc executor never calls. With the ++cpp-cache ops that gap closes for every C++ name. Even then the line proves the ++CONFIGURATION each rank ran, not that a switch changed any output: that is the ++job of an output-equality run (identical sends -> one distinct output). ++""" ++ ++from __future__ import annotations ++ ++import os ++import sys ++from collections.abc import Callable ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++# (environment name, effective key, value when unset or empty) -- the C++ ++# readers, each `static const bool enabled = [] { ... }();` once per process: ++# MERLIN_FULL_K csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu:57-63 ++# MERLIN_TIE_DET csrc/libtorch_stable/topk.cu:60-66 ++# MERLIN_TIE_RECENT csrc/libtorch_stable/topk.cu:52-58 ++# MERLIN_TOPK_EXACT_MEDIUM csrc/libtorch_stable/topk.cu:68-74 (unset -> ON) ++# MERLIN_TOPK_EXACT_DECODE csrc/libtorch_stable/topk.cu:103-109 ++# MERLIN_TOPK_DET_LARGE csrc/libtorch_stable/topk.cu:86-92 (unset -> ON) ++CXX_SWITCHES: tuple[tuple[str, str, bool], ...] = ( ++ ("MERLIN_FULL_K", "full_k", False), ++ ("MERLIN_TIE_DET", "tie_det", False), ++ ("MERLIN_TIE_RECENT", "tie_recent", False), ++ ("MERLIN_TOPK_EXACT_MEDIUM", "exact_medium", True), ++ ("MERLIN_TOPK_EXACT_DECODE", "exact_decode", False), ++ ("MERLIN_TOPK_DET_LARGE", "det_large", True), ++) ++QSA_SORT_NAME = "MERLIN_QSA_SORT" ++RAW_ORDER = ( ++ "MERLIN_FULL_K", ++ "MERLIN_QSA_SORT", ++ "MERLIN_TIE_DET", ++ "MERLIN_TIE_RECENT", ++ "MERLIN_TOPK_EXACT_MEDIUM", ++ "MERLIN_TOPK_EXACT_DECODE", ++ "MERLIN_TOPK_DET_LARGE", ++) ++EFFECTIVE_ORDER = ( ++ "full_k", ++ "qsa_sort", ++ "tie_det", ++ "tie_recent", ++ "tie_recent_effective", ++ "exact_medium", ++ "exact_decode", ++ "det_large", ++) ++CPP_ORDER = tuple(k for k in EFFECTIVE_ORDER if k != "qsa_sort") ++ ++_QSA_MODULE = "vllm.models.qwen4_exp.nvidia.ops.qsa_indexer" ++_C_SPACE = b" \t\n\v\f\r" # isspace() in the "C" locale ++_LONG_MIN, _LONG_MAX = -(2**63), 2**63 - 1 # LP64 long, the reference host's glibc ++ ++ ++# --------------------------------------------------------------- parsers ---- ++def c_atoi(raw: bytes) -> int: ++ """glibc ``atoi``: ``(int) strtol(raw, NULL, 10)`` (glibc stdlib/atoi.c). ++ ++ Leading C-locale whitespace, one optional sign, then decimal digits up to ++ the first non-digit; no digits -> 0. ``strtol`` clamps to the ``long`` ++ range and the ``(int)`` conversion keeps the low 32 bits (GCC's documented ++ modulo behaviour). So ``"true"`` -> 0, ``"1abc"`` -> 1, ``"0x1"`` -> 0, ++ ``" +1"`` -> 1, ``"+ 1"`` -> 0 and ``"4294967296"`` -> 0. ++ """ ++ i, n = 0, len(raw) ++ while i < n and raw[i] in _C_SPACE: ++ i += 1 ++ negative = False ++ if i < n and raw[i] in b"+-": ++ negative = raw[i] == ord("-") ++ i += 1 ++ value, digits = 0, 0 ++ while i < n and 0x30 <= raw[i] <= 0x39: ++ value = value * 10 + (raw[i] - 0x30) ++ digits += 1 ++ i += 1 ++ if digits == 0: ++ return 0 ++ if negative: ++ value = -value ++ value = max(_LONG_MIN, min(_LONG_MAX, value)) ++ value &= 0xFFFFFFFF ++ return value - (1 << 32) if value >= (1 << 31) else value ++ ++ ++def cxx_switch(raw: bytes | None, unset_value: bool) -> bool: ++ """The C++ getter body, both shapes: ++ default OFF ``(raw != nullptr && raw[0] != '\\0') ? (atoi(raw) != 0) : false`` ++ default ON ``(raw == nullptr || raw[0] == '\\0') ? true : (atoi(raw) != 0)`` ++ An environment value cannot contain NUL, so "empty" is exactly ++ ``raw[0] == '\\0'``.""" ++ if raw is None or raw == b"": ++ return unset_value ++ return c_atoi(raw) != 0 ++ ++ ++def qsa_sort_value(raw: str | None) -> str: ++ """vllm/envs.py:2372 ``bool(int(os.getenv("MERLIN_QSA_SORT", "0")))``, ++ read once at import by qsa_indexer.py:51 (``_QSA_SORT``). ``int()`` raises ++ on ``"true"`` and on ``""`` -- that engine would not have booted -- which is ++ reported as ``invalid`` rather than guessed. Python's own ``int`` is used, ++ so whitespace, ``+``/``-``, ``1_0`` and non-ASCII digits parse exactly as ++ the reader parses them.""" ++ try: ++ return "1" if bool(int("0" if raw is None else raw)) else "0" ++ except ValueError: ++ return "invalid" ++ ++ ++def effective_tie_recent(det_ties: bool, tie_recent: bool) -> bool: ++ """csrc/libtorch_stable/persistent_topk.cuh:65-68: ``det_ties ? tie_recent ++ : true`` -- with TIE_DET off the restart paths keep most-recent-wins.""" ++ return tie_recent if det_ties else True ++ ++ ++# --------------------------------------------------------------- sources ---- ++def _libc_getenv() -> tuple[Callable[[str], bytes | None], str]: ++ """libc ``getenv``: the ``environ`` the C++ ``std::getenv`` calls read. ++ ++ Falls back to ``os.environb`` / ``os.environ`` -- and says so in the line's ++ ``source=`` -- where there is no C library to ask (an offline run off Linux). ++ """ ++ try: ++ import ctypes ++ ++ getenv = ctypes.CDLL(None).getenv ++ getenv.restype = ctypes.c_char_p ++ getenv.argtypes = [ctypes.c_char_p] ++ ++ def read(name: str) -> bytes | None: ++ return getenv(name.encode("ascii")) ++ ++ return read, "libc-getenv" ++ except (OSError, AttributeError, TypeError): ++ pass ++ environb = getattr(os, "environb", None) ++ if environb is not None: ++ return (lambda name: environb.get(name.encode("ascii"))), "os-environ" ++ ++ def read_environ(name: str) -> bytes | None: ++ value = os.environ.get(name) ++ return None if value is None else value.encode("utf-8", "surrogateescape") ++ ++ return read_environ, "os-environ" ++ ++ ++def _printable(b: int) -> bool: ++ return 0x21 <= b <= 0x7E and b not in b'"\\' ++ ++ ++def encode_raw(raw: bytes | None) -> str: ++ """``unset``, a plain token, or a quoted, hex-escaped string -- never a ++ space, so the line splits on whitespace. The literal value ``unset`` is ++ quoted so it cannot pass for an absent variable.""" ++ if raw is None: ++ return "unset" ++ if raw and raw != b"unset" and all(_printable(b) for b in raw): ++ return raw.decode("ascii") ++ return '"' + "".join(chr(b) if _printable(b) else f"\\x{b:02x}" for b in raw) + '"' ++ ++ ++def environ_view( ++ cxx_getenv: Callable[[str], bytes | None], ++ py_getenv: Callable[[str], str | None], ++) -> tuple[dict[str, str], dict[str, str]]: ++ """(raw value by name, effective map) from one read of the environment.""" ++ raw: dict[str, str] = {} ++ flags: dict[str, bool] = {} ++ for name, key, unset_value in CXX_SWITCHES: ++ value = cxx_getenv(name) ++ raw[name] = encode_raw(value) ++ flags[key] = cxx_switch(value, unset_value) ++ qsa_raw = py_getenv(QSA_SORT_NAME) ++ raw[QSA_SORT_NAME] = encode_raw( ++ None if qsa_raw is None else qsa_raw.encode("utf-8", "surrogateescape") ++ ) ++ eff = {key: "1" if flags[key] else "0" for key in flags} ++ eff["qsa_sort"] = qsa_sort_value(qsa_raw) ++ eff["tie_recent_effective"] = ( ++ "1" if effective_tie_recent(flags["tie_det"], flags["tie_recent"]) else "0" ++ ) ++ return raw, {key: eff[key] for key in EFFECTIVE_ORDER} ++ ++ ++_TOPK_STATE_FIELDS = ( ++ "tie_det", ++ "tie_recent", ++ "tie_recent_effective", ++ "exact_medium", ++ "exact_decode", ++ "det_large", ++) ++ ++ ++def decode_topk_state(state: int) -> dict[str, str] | str: ++ """Unpack ``persistent_topk_toggle_state()``: 2 bits per field (0 off, ++ 1 on, 2 not initialised), field count at bits 48-55 (6), layout version at ++ bits 56-62 (1). Any other layout is refused with a reason string.""" ++ version, count = (state >> 56) & 0x7F, (state >> 48) & 0xFF ++ if version != 1 or count != len(_TOPK_STATE_FIELDS): ++ return f"unknown-layout-v{version}-{count}-fields" ++ names = {0: "0", 1: "1", 2: "uninit"} ++ return { ++ key: names.get((state >> (2 * i)) & 3, "invalid") ++ for i, key in enumerate(_TOPK_STATE_FIELDS) ++ } ++ ++ ++def decode_full_k_state(state: int) -> str: ++ """``marlin_moe_full_k_state()``: -1 not initialised, 0 off, 1 on.""" ++ return {-1: "uninit", 0: "0", 1: "1"}.get(int(state), "invalid") ++ ++ ++def _op(namespace: str, name: str): ++ try: ++ import torch ++ ++ return getattr(getattr(torch.ops, namespace), name, None) ++ except (AttributeError, RuntimeError, ImportError): ++ return None ++ ++ ++def cpp_cache() -> dict[str, str] | str: ++ """The C++ cached switches, or a reason string when the build cannot say.""" ++ topk_op = _op("_C", "persistent_topk_toggle_state") ++ full_k_op = _op("_moe_C", "marlin_moe_full_k_state") ++ missing = [ ++ name ++ for name, op in ( ++ ("_C.persistent_topk_toggle_state", topk_op), ++ ("_moe_C.marlin_moe_full_k_state", full_k_op), ++ ) ++ if op is None ++ ] ++ if missing: ++ return "op-absent:" + ",".join(missing) ++ decoded = decode_topk_state(int(topk_op())) ++ if isinstance(decoded, str): ++ return decoded ++ out = dict(decoded, full_k=decode_full_k_state(int(full_k_op()))) ++ return {key: out[key] for key in CPP_ORDER} ++ ++ ++def qsa_sort_live() -> str: ++ module = sys.modules.get(_QSA_MODULE) ++ if module is None or not hasattr(module, "_QSA_SORT"): ++ return "n-a" ++ return "1" if bool(module._QSA_SORT) else "0" ++ ++ ++def device_switches() -> dict[str, str]: ++ """Counter slots [8]/[9] (switch / switch_large). One synchronizing read ++ (topk_counters.read refuses inside a capture), so it is only called from ++ arm() and the head of an execute step.""" ++ try: ++ from vllm.models.qwen4_exp.nvidia.ops import topk_counters ++ ++ counters = topk_counters.read() ++ except Exception: ++ counters = {} ++ return { ++ key: str(counters[key]) if key in counters else "n-a" ++ for key in ("switch", "switch_large") ++ } ++ ++ ++# ------------------------------------------------------------------ line ---- ++def toggles_line( ++ tag: str, ++ pid: int, ++ raw: dict[str, str], ++ effective: dict[str, str], ++ env_source: str, ++ cpp: dict[str, str] | str, ++ qsa_live: str, ++ device: dict[str, str], ++) -> str: ++ parts = [f"e2-toggles[{tag}]: pid={pid}"] ++ parts += [f"{name}={raw[name]}" for name in RAW_ORDER] ++ parts.append("effective:") ++ parts += [f"{key}={effective[key]}" for key in EFFECTIVE_ORDER] ++ parts.append(f"source={env_source}") ++ parts.append("cpp-cache:") ++ if isinstance(cpp, str): ++ parts += ["n-a", f"reason={cpp}"] ++ else: ++ parts += [f"{key}={cpp[key]}" for key in CPP_ORDER] ++ parts.append("source=cpp-cache") ++ parts += ["python:", f"qsa_sort={qsa_live}", "source=qsa_indexer._QSA_SORT"] ++ parts.append("device:") ++ parts += [f"{key}={device[key]}" for key in ("switch", "switch_large")] ++ parts.append("source=counter-slots-8-9") ++ return " ".join(parts) ++ ++ ++def _rank_tag() -> str: ++ from vllm.distributed.parallel_state import get_pp_group, get_tp_group ++ ++ return f"PP{get_pp_group().rank_in_group}_TP{get_tp_group().rank_in_group}" ++ ++ ++# ----------------------------------------------------------------- state ---- ++_armed = False ++_emitted = False ++_deferral_noted = False ++ ++ ++def _reset_for_tests() -> None: ++ global _armed, _emitted, _deferral_noted ++ _armed = _emitted = _deferral_noted = False ++ ++ ++def arm(**kwargs) -> str: ++ """Once per worker, at the end of compile_or_warm_up_model (warmup and ++ capture done): emit now, unless a C++ static is still uninitialised.""" ++ global _armed ++ _armed = True ++ return maybe_emit(**kwargs) ++ ++ ++def maybe_emit( ++ force: bool = False, ++ *, ++ tag_fn: Callable[[], str] = _rank_tag, ++ getenv_fn: Callable[[], tuple[Callable[[str], bytes | None], str]] = _libc_getenv, ++ cpp_fn: Callable[[], dict[str, str] | str] = cpp_cache, ++ device_fn: Callable[[], dict[str, str]] = device_switches, ++ qsa_fn: Callable[[], str] = qsa_sort_live, ++) -> str: ++ """Emit the line at most once per process. One bool check when already ++ emitted or not yet armed, so it can sit at the head of every execute step. ++ Never raises. Returns the line when this call emitted it, else "".""" ++ global _emitted, _deferral_noted ++ if _emitted or not _armed: ++ return "" ++ tag = "PP?_TP?" ++ try: ++ tag = tag_fn() ++ cpp = cpp_fn() ++ pending = ( ++ sorted(k for k, v in cpp.items() if v == "uninit") ++ if isinstance(cpp, dict) ++ else [] ++ ) ++ if pending and not force: ++ if not _deferral_noted: ++ _deferral_noted = True ++ # Deliberately not 'e2-toggles[': the smoke counts those. ++ logger.info( ++ "e2-toggles deferred on %s: C++ cache not initialised for %s", ++ tag, ++ ",".join(pending), ++ ) ++ return "" ++ cxx_getenv, env_source = getenv_fn() ++ raw, effective = environ_view(cxx_getenv, os.environ.get) ++ line = toggles_line( ++ tag, os.getpid(), raw, effective, env_source, cpp, qsa_fn(), device_fn() ++ ) ++ _emitted = True ++ logger.info("%s", line) ++ return line ++ except Exception as error: # log-only: must never break a boot or a step ++ _emitted = True ++ logger.warning( ++ "e2-toggles[%s]: unavailable (%s: %s)", tag, type(error).__name__, error ++ ) ++ return "" +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index b5b8a70..f37f741 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -88,7 +88,7 @@ from vllm.v1.watermarking.spec_decode import ( + create_speculative_target_watermarker, + speculative_target_watermark_key, + ) +-from vllm.v1.worker import guard_telemetry ++from vllm.v1.worker import e2_toggles, guard_telemetry + from vllm.v1.worker.block_table import get_block_table_width + from vllm.v1.worker.cp_utils import check_attention_cp_compatibility + from vllm.v1.worker.gpu import pcp_manager as pcp +@@ -1794,6 +1794,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # never fire on them. Outside any capture: this is the engine + # thread between steps (read() re-checks). + topk_counters.maybe_log() ++ # v2.2: the e2-toggles line if arm() had to defer it (a C++ static ++ # still uninitialised). One bool check once it has been emitted. ++ e2_toggles.maybe_emit() + # Update the request states. + self.update_pp_decode_requests() + self.finish_requests(scheduler_output) +@@ -2492,6 +2495,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): + # The last window of a run is the one a cadence log always + # misses; emit it before the device state goes away. + topk_counters.maybe_log(force=True) ++ # v2.2: a still-deferred e2-toggles line goes out now, 'uninit' and all. ++ e2_toggles.maybe_emit(force=True) + if self._ple_offload_connector is not None: + self._ple_offload_connector.close() + self._ple_offload_connector = None +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 847837d..4edfb76 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -673,6 +673,12 @@ class Worker(WorkerBase): + model_runner = cast("GPUModelRunnerV2", self.model_runner) + model_runner._setup_ple_offload(self.parallel_config._ple_offload_ipc_path) + ++ # Log-only: this rank's PLE / capture identity for target and drafter, ++ # read after the offload placeholders are armed (v1/worker/e2_identity.py). ++ from vllm.v1.worker.e2_identity import log_launch_identity ++ ++ log_launch_identity(self.model_runner, self.vllm_config) ++ + if has_ec_transfer(): + get_ec_transfer().start_worker_services() + +@@ -986,6 +992,13 @@ class Worker(WorkerBase): + with self._get_cudagraph_capture_context(): + cuda_graph_memory_bytes = self.model_runner.capture_model() + ++ # Log-only: this rank's determinism toggles, once, now that the warmup ++ # and capture above have launched persistent_topk and Marlin MoE, so ++ # their C++ statics are filled (v1/worker/e2_toggles.py). Never raises. ++ from vllm.v1.worker import e2_toggles ++ ++ e2_toggles.arm() ++ + # Compare actual vs estimated CUDA graph memory (if we did profiling) + if ( + hasattr(self, "cudagraph_memory_estimate") +-- +2.43.0 + diff --git a/patches/0048-ple-the-pull-registry-refuses-unbound-or-clashing-na.patch b/patches/0048-ple-the-pull-registry-refuses-unbound-or-clashing-na.patch new file mode 100644 index 0000000..44a97d0 --- /dev/null +++ b/patches/0048-ple-the-pull-registry-refuses-unbound-or-clashing-na.patch @@ -0,0 +1,203 @@ +From e42e2eee792338708274acd3d0dc5446bfb3951e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 09:45:20 -0400 +Subject: [PATCH 048/103] ple: the pull registry refuses unbound or clashing + names; test fixes + +A schema walk fails on any pointer-like int parameter of a custom op on the +PLE path. The FULL_K test child calls the Marlin MoE GEMM by keyword against +its v0.30.0 signature. +--- + tests/kernels/moe/merlin_full_k_child.py | 28 +++++++-- + .../ple_offload/test_transport_bounded_io.py | 17 ++++-- + tests/v1/test_ple_pull_compile_profile.py | 61 +++++++++++++++++++ + .../layers/ple_offload_layer.py | 16 ++++- + 4 files changed, 111 insertions(+), 11 deletions(-) + +diff --git a/tests/kernels/moe/merlin_full_k_child.py b/tests/kernels/moe/merlin_full_k_child.py +index 3a42008..d8cc365 100644 +--- a/tests/kernels/moe/merlin_full_k_child.py ++++ b/tests/kernels/moe/merlin_full_k_child.py +@@ -56,6 +56,11 @@ def main(out_path: str) -> None: + qtype = scalar_types.uint4b8 + K, GROUP, BLOCK = 2048, 128, 16 + TN = 128 # thread_n of the explicit config ++ import inspect ++ ++ _sig = inspect.signature(ops.moe_wna16_marlin_gemm) ++ GEMM_PARAMS = list(_sig.parameters) ++ GEMM_DEFAULTED = {k for k, v in _sig.parameters.items() if v.default is not inspect.Parameter.empty} + configs = {"explicit": (128, 128, 1), "auto": (-1, -1, -1)} + + gen = torch.Generator(device="cpu").manual_seed(46639) +@@ -63,7 +68,8 @@ def main(out_path: str) -> None: + + def weights(n: int): + w = (torch.randn(K, n, generator=gen) / 10).to(dtype).to(dev) +- _, q, s, _, _, _ = marlin_quantize(w, qtype, GROUP, False) ++ # 0.28 returns 6 items (4th positional = act_order); 0.30 returns 3 (#54809) ++ _, q, s = marlin_quantize(w, qtype, GROUP)[:3] + return q.unsqueeze(0).contiguous(), s.unsqueeze(0).contiguous() + + def gemm(m: int, n: int, q, s, cfg): +@@ -73,14 +79,24 @@ def main(out_path: str) -> None: + topk_w = torch.ones((m, 1), dtype=torch.float32, device=dev) + sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, BLOCK, 1) + ws = marlin_make_workspace_new(dev, 4) +- out = ops.moe_wna16_marlin_gemm( +- a, None, q, None, s, None, None, None, None, None, ws, +- sorted_ids, expert_ids, ntpp, topk_w, ++ # Keyword call against the op's REAL signature: upstream #54809 (in vLLM ++ # 0.30) removed act-order (g_idx, perm, is_k_full) from this op, so the ++ # same child runs on the 0.28.1rc1 fork and on the 0.30 fork. ++ kw = dict( ++ input=a, output=None, b_qweight=q, b_bias=None, b_scales=s, ++ a_scales=None, global_scale=None, b_qzeros=None, workspace=ws, ++ sorted_token_ids=sorted_ids, expert_ids=expert_ids, ++ num_tokens_past_padded=ntpp, topk_weights=topk_w, + moe_block_size=BLOCK, top_k=1, mul_topk_weights=False, +- b_q_type=qtype, size_m=m, size_n=n, size_k=K, is_k_full=True, ++ b_q_type=qtype, size_m=m, size_n=n, size_k=K, + use_atomic_add=False, use_fp32_reduce=True, is_zp_float=False, + thread_k=tk, thread_n=tn, blocks_per_sm=bps, +- ) # fmt: skip ++ ) ++ if "is_k_full" in GEMM_PARAMS: ++ kw.update(g_idx=None, perm=None, is_k_full=True) ++ missing = sorted(set(GEMM_PARAMS) - set(kw) - GEMM_DEFAULTED) ++ assert not missing, f"moe_wna16_marlin_gemm needs {missing}: child out of date" ++ out = ops.moe_wna16_marlin_gemm(**kw) + torch.cuda.synchronize() + tiles = int(ntpp.item()) // BLOCK * (n // TN) + return out.float().cpu(), tiles +diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py +index 5a2aa19..78a9db0 100644 +--- a/tests/v1/ple_offload/test_transport_bounded_io.py ++++ b/tests/v1/ple_offload/test_transport_bounded_io.py +@@ -273,14 +273,23 @@ def test_real_wait_copy_geometry_rejected_before_graph_publication( + gpu_output_buffer = torch.empty(12, 4, dtype=torch.float32) + hidden_states = torch.empty(12, 1, dtype=torch.float32) + ++ # Go through the op's real entry point, which resolves ++ # the live pointers by layer name from the per-process registry. ++ from types import SimpleNamespace ++ ++ monkeypatch.setitem( ++ ple_layer_module._PULL_LAYERS, ++ "ple", ++ SimpleNamespace( ++ _sem=SimpleNamespace(flag_ptr=1), _host_output_ptr=2, _host_row_bytes=16 ++ ), ++ ) ++ + def capture_then_publish() -> None: + with manager.capture_ple_copy_geometry(desc): +- ple_layer_module._ple_offload_wait_copy_raw( ++ ple_layer_module._ple_offload_wait_copy_impl( + gpu_output_buffer, + hidden_states, +- 1, +- 2, +- 16, + 12, + "ple", + ) +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index d5c47a9..0f3ab58 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -534,3 +534,64 @@ def test_wait_copy_resolves_live_pointers_by_layer_name(monkeypatch) -> None: + + with pytest.raises(RuntimeError, match="no armed pull layer"): + m._ple_offload_wait_copy_impl(buf, buf, 3, "missing") ++ ++ ++def test_pull_registry_refuses_unbound_and_clashing_names(monkeypatch) -> None: ++ """setup_cross_process_offload must not register under ++ the default '' name nor let a second layer overwrite a first one's ++ entry (the op would resolve the wrong pointers). Re-arming the SAME layer ++ is allowed.""" ++ import pytest ++ import torch ++ ++ m = _import_ple_offload_layer(monkeypatch) ++ monkeypatch.setattr(m, "_PULL_LAYERS", {}) ++ ++ class _L(m.PleOffloadLayer): ++ def forward_impl(self, hidden_states, input_ids, *a, **k): # pragma: no cover ++ return hidden_states ++ ++ def armed(name): ++ layer = _L.__new__(_L) ++ torch.nn.Module.__init__(layer) ++ layer._ple_offload_name = name ++ return layer ++ ++ sem = object() ++ buf = torch.empty(2, 2) ++ with pytest.raises(RuntimeError, match="no _ple_offload_name"): ++ armed("").setup_cross_process_offload(buf, sem, host_output_ptr=1, host_row_bytes=8) ++ a = armed("L1") ++ a.setup_cross_process_offload(buf, sem, host_output_ptr=1, host_row_bytes=8) ++ a.setup_cross_process_offload(buf, sem, host_output_ptr=2, host_row_bytes=8) # re-arm ok ++ assert m._PULL_LAYERS["L1"] is a ++ with pytest.raises(RuntimeError, match="two layers armed"): ++ armed("L1").setup_cross_process_offload(buf, sem, host_output_ptr=3, host_row_bytes=8) ++ ++ ++def test_nvidia_qwen4exp_ops_carry_no_pointer_like_int_params() -> None: ++ """A schema walk -- no custom op registered by the ++ NVIDIA qwen4_exp / PLE offload path may take an int parameter named like a ++ pointer/address (the AOT cache bakes Python ints as constants).""" ++ import ast ++ import pathlib ++ ++ root = pathlib.Path(__file__).resolve().parents[2] / "vllm" ++ files = [root / "model_executor/layers/ple_offload_layer.py", ++ *sorted((root / "models/qwen4_exp/nvidia").rglob("*.py")), ++ *sorted((root / "v1/ple_offload").rglob("*.py"))] ++ bad = [] ++ for f in files: ++ tree = ast.parse(f.read_text(encoding="utf-8")) ++ funcs = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} ++ for node in ast.walk(tree): ++ if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "direct_register_custom_op": ++ kw = {k.arg: k.value for k in node.keywords} ++ fn = getattr(kw.get("op_func"), "id", None) ++ if fn not in funcs: ++ continue ++ for a in funcs[fn].args.args: ++ ann = ast.unparse(a.annotation) if a.annotation is not None else "" ++ if ann == "int" and any(t in a.arg for t in ("ptr", "addr", "pointer")): ++ bad.append(f"{f.name}:{fn}({a.arg})") ++ assert not bad, f"pointer-like int op params (AOT-baking hazard): {bad}" +diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py +index ff581b2..cc8bd73 100644 +--- a/vllm/model_executor/layers/ple_offload_layer.py ++++ b/vllm/model_executor/layers/ple_offload_layer.py +@@ -428,7 +428,21 @@ class PleOffloadLayer(nn.Module, ABC): + self._host_output_ptr = host_output_ptr + self._host_row_bytes = int(host_row_bytes) + if host_output_ptr is not None: +- _PULL_LAYERS[self._ple_offload_name] = self ++ # Hardening: the registry key is the name the ++ # traced graph carries, so it must be a real, unique layer name. ++ name = self._ple_offload_name ++ if name == "": ++ raise RuntimeError( ++ "PLE pull: setup_cross_process_offload on a layer with no " ++ "_ple_offload_name; the op would resolve nothing by name" ++ ) ++ prior = _PULL_LAYERS.get(name) ++ if prior is not None and prior is not self: ++ raise RuntimeError( ++ f"PLE pull: two layers armed under the name {name!r}; the op " ++ "would resolve the wrong layer's pointers" ++ ) ++ _PULL_LAYERS[name] = self + + def forward( + self, +-- +2.43.0 + diff --git a/patches/0049-qsa-tag-every-raw-key-ring-row-and-validate-it-befor.patch b/patches/0049-qsa-tag-every-raw-key-ring-row-and-validate-it-befor.patch new file mode 100644 index 0000000..9cadd8d --- /dev/null +++ b/patches/0049-qsa-tag-every-raw-key-ring-row-and-validate-it-befor.patch @@ -0,0 +1,3035 @@ +From d56bf8be4cd2f25ae2539a3de137fa4422bf7346 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 09:58:46 -0400 +Subject: [PATCH 049/103] qsa: tag every raw-key ring row and validate it + before pooling + +The fused pre-indexer could read a stale or uncommitted raw-key ring row and +pass its int64 tail to RoPE as a position (illegal address). The writer now +tags each row with logical_position + 1; the reader checks the tag of every +source row and, on a mismatch, masks the pooled load, the RoPE-tail read and +the compressed-slot store. New guard site 11 (always fatal under warn). The +KV pool size is unchanged. +--- + .../models/qwen4_exp/test_qsa_pre_indexer.py | 163 +++++- + tests/models/qwen4_exp/test_qsa_reference.py | 6 +- + .../qwen4_exp/test_qsa_ring_tag_contract.py | 371 ++++++++++++++ + .../v1/core/test_qsa_ring_resume_alignment.py | 266 ++++++++++ + tests/v1/core/test_ring_tag_clear.py | 367 ++++++++++++++ + tests/v1/test_pp_kv_projection.py | 5 +- + tests/v1/worker/gpu/ring_tag_replay.py | 467 ++++++++++++++++++ + .../worker/gpu/test_guard_telemetry_limits.py | 8 +- + .../gpu/test_qsa_ring_skip_telemetry.py | 11 +- + tests/v1/worker/gpu/test_qsa_ring_tags_gpu.py | 311 ++++++++++++ + .../gpu/test_qsa_rope_position_guard.py | 13 +- + vllm/models/qwen4_exp/amd/indexer_qsa.py | 9 + + vllm/models/qwen4_exp/common/qsa_cache.py | 76 ++- + vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 14 + + .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 118 ++++- + vllm/v1/core/kv_cache_manager.py | 15 + + vllm/v1/core/sched/output.py | 7 + + vllm/v1/core/sched/scheduler.py | 24 + + vllm/v1/core/single_type_kv_cache_manager.py | 46 ++ + vllm/v1/worker/gpu/model_runner.py | 10 + + vllm/v1/worker/gpu_model_runner.py | 9 + + vllm/v1/worker/guard_telemetry.py | 73 ++- + vllm/v1/worker/ring_tags.py | 63 +++ + 23 files changed, 2420 insertions(+), 32 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_qsa_ring_tag_contract.py + create mode 100644 tests/v1/core/test_qsa_ring_resume_alignment.py + create mode 100644 tests/v1/core/test_ring_tag_clear.py + create mode 100644 tests/v1/worker/gpu/ring_tag_replay.py + create mode 100644 tests/v1/worker/gpu/test_qsa_ring_tags_gpu.py + create mode 100644 vllm/v1/worker/ring_tags.py + +diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer.py b/tests/models/qwen4_exp/test_qsa_pre_indexer.py +index 8d4df3d..13b9a8d 100644 +--- a/tests/models/qwen4_exp/test_qsa_pre_indexer.py ++++ b/tests/models/qwen4_exp/test_qsa_pre_indexer.py +@@ -6,9 +6,12 @@ import pytest + import torch + + from vllm.models.qwen4_exp.common.qsa_cache import ( ++ QSA_RING_TAG_WORD, ++ QSA_RING_TAIL_ELEMENTS, + canonical_qsa_rope_positions, + circular_qsa_slot_mapping, + compressed_qsa_slot_mapping, ++ qsa_ring_row_tags, + ) + from vllm.models.qwen4_exp.nvidia.indexer_qsa import apply_qsa_rope + from vllm.models.qwen4_exp.nvidia.ops.qsa import ( +@@ -20,6 +23,7 @@ from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import ( + ) + from vllm.platforms import current_platform + from vllm.triton_utils import HAS_TRITON ++from vllm.v1.worker import guard_telemetry + + requires_qsa_kernels = pytest.mark.skipif( + not current_platform.is_cuda() or not HAS_TRITON, +@@ -195,7 +199,9 @@ def test_qsa_fused_pre_indexer_matches_unfused( + (work_requests, local_work), dim=1 + ) + +- raw_width = D + 12 if cache_rope_positions else D ++ # Every ring row is key + four int64 words (pos_t, pos_h, pos_w, tag), ++ # with or without cached RoPE positions. ++ raw_width = D + QSA_RING_TAIL_ELEMENTS + # Match vLLM's padded-page cache layout: rows are contiguous, while physical + # blocks have a larger stride than their logical contents. + raw_page_elements = state_size * raw_width +@@ -217,14 +223,17 @@ def test_qsa_fused_pre_indexer_matches_unfused( + row = fused_raw[block, position % state_size, 0] + row[:D] = torch.randn(D, dtype=torch.bfloat16, device=device) + if cache_rope_positions: +- row[ROPE_POS_OFFSET:].view(torch.int64).copy_( ++ row[ROPE_POS_OFFSET:].view(torch.int64)[:3].copy_( + torch.tensor( + [position, position // 7 + 3, position // 13 + 11], + dtype=torch.int64, + device=device, + ) + ) ++ # the history was committed by this request: its tag names it ++ row[ROPE_POS_OFFSET:].view(torch.int64)[QSA_RING_TAG_WORD] = position + 1 + unfused_raw = fused_raw.clone() ++ unfused_raw_history = fused_raw.clone() # pre-launch ring, for the reference + compressed_page_elements = COMP_PAGE * D + fused_compressed_storage = torch.zeros( + num_compressed_blocks, +@@ -246,6 +255,8 @@ def test_qsa_fused_pre_indexer_matches_unfused( + k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 + + fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) ++ guard_telemetry.prealloc(torch.device(device, torch.cuda.current_device())) ++ guard_telemetry.reset_for_test() + qsa_pre_indexer( + projected_qk[:, : HQ * D], + projected_qk[:, HQ * D :], +@@ -275,11 +286,9 @@ def test_qsa_fused_pre_indexer_matches_unfused( + unfused_query = apply_qsa_rope(rope, positions, unfused_query) + + raw_keys = unfused_raw[..., :D] +- rope_positions = ( +- unfused_raw[..., ROPE_POS_OFFSET:].view(torch.int64) +- if cache_rope_positions +- else None +- ) ++ raw_tail = unfused_raw[..., ROPE_POS_OFFSET:].view(torch.int64) ++ rope_positions = raw_tail[..., :3] if cache_rope_positions else None ++ raw_tags = raw_tail[..., QSA_RING_TAG_WORD : QSA_RING_TAG_WORD + 1] + pooled, first_positions = qsa_compress_groups_with_ratio( + projected_qk[:, HQ * D :].reshape(-1, 1, D), + position_rows, +@@ -303,6 +312,9 @@ def test_qsa_fused_pre_indexer_matches_unfused( + qsa_store_cache_rows(raw_keys, raw_slots, projected_qk[:, HQ * D :]) + if rope_positions is not None: + qsa_store_cache_rows(rope_positions, raw_slots, position_rows) ++ qsa_store_cache_rows( ++ raw_tags, raw_slots, qsa_ring_row_tags(logical_positions).view(-1, 1, 1) ++ ) + + if indexer_dtype == torch.float8_e4m3fn: + # Both paths round the same ~bf16 intermediates to e4m3. Bitwise +@@ -320,3 +332,140 @@ def test_qsa_fused_pre_indexer_matches_unfused( + torch.testing.assert_close( + fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL + ) ++ # Committed history only, so the uncommitted-row site stays at 0... ++ torch.cuda.synchronize() ++ ring_site = guard_telemetry.counter( ++ guard_telemetry.SITE_QSA_RING_UNCOMMITTED, torch.device(device) ++ ) ++ assert ring_site.tolist() == [0, 0, 0] ++ # ...and every committed row now carries its tag (the fused writer's). ++ fused_tags = fused_raw[..., ROPE_POS_OFFSET:].view(torch.int64)[ ++ ..., QSA_RING_TAG_WORD ++ ] ++ committed = raw_slots >= 0 ++ slot_blocks = raw_slots[committed] // state_size ++ slot_rows = raw_slots[committed] % state_size ++ assert torch.equal( ++ fused_tags[slot_blocks, slot_rows, 0], ++ qsa_ring_row_tags(logical_positions[committed]), ++ ) ++ ++ # INDEPENDENT reference: V1 and V2 could agree on the same wrong read, ++ # so every completed group is recomputed here in plain torch from the ++ # inputs alone -- the history rows as written above, this chunk's keys -- ++ # and compared with the fused compressed cache. ++ reference = _reference_compressed_rows( ++ rope.cos_sin_cache, ++ k_weight, ++ projected_qk[:, HQ * D :], ++ positions, ++ logical_positions, ++ query_start_loc, ++ compressed_slots, ++ unfused_raw_history, ++ raw_block_table, ++ state_size, ++ mrope=mrope, ++ cache_rope_positions=cache_rope_positions, ++ is_2d_positions=is_2d_positions, ++ ) ++ for slot, expected in reference.items(): ++ block, row = divmod(slot, COMP_PAGE) ++ got = fused_compressed[block, row, 0] ++ if indexer_dtype == torch.float8_e4m3fn: ++ assert_fp8_within_one_ulp(got, expected.to(indexer_dtype)) ++ else: ++ torch.testing.assert_close( ++ got, expected.to(indexer_dtype), rtol=RTOL, atol=ATOL ++ ) ++ assert reference, "no completed group: the reference compared nothing" ++ ++ ++def _reference_rope_rows( ++ cos_sin_cache: torch.Tensor, ++ x: torch.Tensor, ++ pos_t: int, ++ pos_h: int, ++ pos_w: int, ++ mrope: bool, ++) -> torch.Tensor: ++ """NeoX rotary on the first half of D, torch only (independent of both the ++ fused _norm_rope and apply_qsa_rope). MRoPE picks the axis per rotary pair ++ the way Qwen interleaves them: h on pairs 1,4,7.. up to 3*MROPE_H, w on ++ pairs 2,5,8.. up to 3*MROPE_W, t elsewhere.""" ++ half, quarter = D // 2, D // 4 ++ table = cos_sin_cache.float() ++ pairs = torch.arange(quarter, device=x.device) ++ if mrope: ++ h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_SECTION[1]) ++ w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_SECTION[2]) ++ axis_pos = torch.where( ++ h_mask, ++ torch.full_like(pairs, pos_h), ++ torch.where(w_mask, torch.full_like(pairs, pos_w), pos_t), ++ ) ++ else: ++ axis_pos = torch.full_like(pairs, pos_t) ++ cos = table[axis_pos, pairs] ++ sin = table[axis_pos, quarter + pairs] ++ r0, r1 = x[:quarter], x[quarter:half] ++ return torch.cat((r0 * cos - r1 * sin, r1 * cos + r0 * sin, x[half:])) ++ ++ ++def _reference_compressed_rows( ++ cos_sin_cache, ++ k_weight, ++ keys, ++ positions, ++ logical_positions, ++ query_start_loc, ++ compressed_slots, ++ history, ++ raw_block_table, ++ state_size, ++ *, ++ mrope, ++ cache_rope_positions, ++ is_2d_positions, ++) -> dict[int, torch.Tensor]: ++ """{compressed slot: expected row} for every completed group, recomputed ++ from the inputs: mean of the group's four raw keys (historical ones from ++ the pre-launch ring, rounded to bf16 like the model), Gemma RMSNorm, then ++ RoPE at the group's first token.""" ++ out: dict[int, torch.Tensor] = {} ++ starts = query_start_loc.tolist() ++ for request in range(len(starts) - 1): ++ q0, q1 = starts[request], starts[request + 1] ++ chunk_start = int(logical_positions[q0]) ++ block = int(raw_block_table[request, 0]) ++ for token in range(q0, q1): ++ slot = int(compressed_slots[token]) ++ if slot < 0: ++ continue ++ end = int(logical_positions[token]) ++ rows, first = [], None ++ for position in range(end - CR + 1, end + 1): ++ if position >= chunk_start: ++ t = q0 + position - chunk_start ++ rows.append(keys[t].float()) ++ pos3 = ( ++ [int(positions[a, t]) for a in range(3)] ++ if is_2d_positions ++ else [int(positions[t])] * 3 ++ ) ++ else: ++ ring_row = history[block, position % state_size, 0] ++ rows.append(ring_row[:D].float()) ++ tail = ring_row[ROPE_POS_OFFSET:].view(torch.int64) ++ pos3 = [int(v) for v in tail[:3]] ++ if not cache_rope_positions: ++ # without cached coordinates the kernel rotates the group ++ # by its first LOGICAL position on every axis ++ pos3 = [position] * 3 ++ if first is None: ++ first = pos3 ++ pooled = (torch.stack(rows).sum(0) / CR).to(torch.bfloat16).float() ++ normed = pooled * torch.rsqrt(pooled.pow(2).mean() + EPS) ++ normed = (normed * (k_weight.float() + 1.0)).to(torch.bfloat16).float() ++ out[slot] = _reference_rope_rows(cos_sin_cache, normed, *first, mrope=mrope) ++ return out +diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py +index 16294a3..29f7bce 100644 +--- a/tests/models/qwen4_exp/test_qsa_reference.py ++++ b/tests/models/qwen4_exp/test_qsa_reference.py +@@ -56,6 +56,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection( + ), + compressed_key_cache=SimpleNamespace(kv_cache=torch.empty(0)), + use_fused_pre_indexer=True, ++ layer_id=0, # the fused launch names the layer in its record + rotary_emb=SimpleNamespace(cos_sin_cache=torch.empty(0)), + q_layernorm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6), + k_layernorm=SimpleNamespace(weight=torch.ones(1)), +@@ -381,13 +382,14 @@ def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None: + ), + compress_ratio=4, + ) +- raw_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) ++ # Every ring row carries the four-word int64 tail (+16 bf16). ++ raw_view = torch.empty(2, 1, 8, 64 + 16, dtype=torch.bfloat16) + compressed_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) + + raw_cache.bind_kv_cache(raw_view) + compressed_cache.bind_kv_cache(compressed_view) + +- assert raw_cache.kv_cache.shape == (2, 8, 1, 64) ++ assert raw_cache.kv_cache.shape == (2, 8, 1, 64 + 16) + assert compressed_cache.kv_cache.shape == (2, 8, 1, 64) + assert raw_cache.kv_cache.data_ptr() == raw_view.data_ptr() + assert compressed_cache.kv_cache.data_ptr() == compressed_view.data_ptr() +diff --git a/tests/models/qwen4_exp/test_qsa_ring_tag_contract.py b/tests/models/qwen4_exp/test_qsa_ring_tag_contract.py +new file mode 100644 +index 0000000..6714070 +--- /dev/null ++++ b/tests/models/qwen4_exp/test_qsa_ring_tag_contract.py +@@ -0,0 +1,371 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""v2.2 ring-row tag fix, CPU half: the raw-key ring row contract. ++ ++The pre-indexer used to read a historical ring row's int64 position tail with ++nothing tying the row to the request or the position (the ring-residue ++illegal-address class). The ring-row tag fix grows the tail to four int64 words ++(pos_t, pos_h, pos_w, tag = logical_position + 1), validates the tag of EVERY ++historical row before pooling, and clears the tag words of a ring block when it ++is claimed. This file pins, without Triton or a GPU: ++ ++ * the layout: spec width D+16 (was D+12 with MRoPE, D without), the V1 ++ position view (3 int64, same bytes as before) and the tag view (word 3); ++ * the claim clear touches the tag words of the named blocks and nothing else; ++ * the torch reference of the reader's rule, including the every-row case ++ (only the 2nd or 3rd row of a group bad) that guard 9 cannot see; ++ * the new telemetry site: declared, named, fatal, coherent record packing; ++ * the production KV-pool impact of the wider row (0 blocks at 4.1e9). ++ ++FAILING CONTROL: on the unmodified tree (without this fix) every ++test here fails -- the constants/helpers do not exist (ImportError at ++collection) and the spec width is 140, not 144. ++""" ++ ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.models.qwen4_exp.common import qsa_cache ++from vllm.models.qwen4_exp.common.qsa_cache import ( ++ QSA_RING_TAG_WORD, ++ QSA_RING_TAIL_ELEMENTS, ++ QSA_RING_TAIL_WORDS, ++ qsa_ring_group_rows_valid, ++ qsa_ring_row_tags, ++) ++from vllm.v1.kv_cache_interface import CircularBufferSpec ++from vllm.v1.worker import guard_telemetry as gt ++ ++D = 128 ++CR = 4 ++STATE = 8 ++ ++ ++def _key_cache(*, rope: bool, head: int = D, block_size: int = 3200): ++ return qsa_cache.QSAKeyStateCache( ++ head_size=head, ++ dtype=torch.bfloat16, ++ cache_rope_positions=rope, ++ cache_config=SimpleNamespace(block_size=block_size), ++ prefix=f"ringtag.raw.{rope}.{head}.{block_size}", ++ vllm_config=SimpleNamespace( ++ compilation_config=SimpleNamespace(static_forward_context={}) ++ ), ++ compress_ratio=CR, ++ ) ++ ++ ++# -------------------------------------------------------------------------- ++# layout: spec width, strides, views ++# -------------------------------------------------------------------------- ++ ++ ++@pytest.mark.parametrize("rope", [True, False], ids=["mrope", "text"]) ++def test_spec_width_is_key_plus_four_int64_words(rope): ++ cache = _key_cache(rope=rope) ++ assert QSA_RING_TAIL_WORDS == 4 and QSA_RING_TAIL_ELEMENTS == 16 ++ assert cache.head_size == D + 16 ++ # production: MTP K=3 -> ring capacity 8, row 144 bf16 = 288 bytes ++ spec = cache.get_kv_cache_spec(SimpleNamespace(num_speculative_tokens=3)) ++ assert isinstance(spec, CircularBufferSpec) ++ assert spec.head_size == 144 and spec.block_size == STATE ++ assert spec.page_size_bytes == STATE * 144 * 2 == 2304 ++ ++ ++def test_tail_starts_on_an_int64_boundary_for_odd_key_widths(): ++ # a key width that is not a multiple of 4 bf16 is padded before the tail ++ cache = _key_cache(rope=True, head=126, block_size=32) ++ assert cache.rope_position_offset == 128 ++ assert cache.head_size == 128 + 16 ++ ++ ++def _bound(rope: bool, num_blocks: int = 3): ++ cache = _key_cache(rope=rope, block_size=32) ++ # the unified [blocks, 1, states, width] view, with a padded block stride ++ # like the real packed pool (block stride > states * width) ++ width = cache.head_size ++ storage = torch.zeros(num_blocks, STATE * width + 32, dtype=torch.bfloat16) ++ unified = torch.as_strided( ++ storage, (num_blocks, 1, STATE, width), (STATE * width + 32, 0, width, 1) ++ ) ++ cache.bind_kv_cache(unified) ++ return cache ++ ++ ++def test_views_alias_the_right_bytes(): ++ cache = _bound(rope=True) ++ assert cache.key_cache.shape == (3, STATE, 1, D) ++ assert cache.rope_position_cache.shape == (3, STATE, 1, 3) ++ assert cache.rope_position_cache.dtype == torch.int64 ++ assert cache.ring_tag_cache.shape == (3, STATE, 1, 1) ++ # same storage, word 3 of the tail ++ cache.rope_position_cache[1, 5, 0] = torch.tensor([11, 22, 33]) ++ cache.ring_tag_cache[1, 5, 0, 0] = 44 ++ tail_bytes = cache.kv_cache[1, 5, 0, D:].contiguous().view(torch.int64) ++ assert tail_bytes.tolist() == [11, 22, 33, 44] ++ # the position view is the pre-tag view (first 3 words): V1 stores into ++ # it unchanged and never reaches the tag ++ assert cache.rope_position_cache.stride(-1) == 1 ++ ++ ++def test_text_ring_has_a_tag_but_no_position_view(): ++ cache = _bound(rope=False) ++ assert cache.rope_position_cache is None ++ assert cache.ring_tag_cache is not None ++ cache.ring_tag_cache[2, 7, 0, 0] = 9 ++ assert cache.kv_cache[2, 7, 0, D:].contiguous().view(torch.int64)[3] == 9 ++ ++ ++def test_bind_rejects_the_pre_tag_width(): ++ cache = _key_cache(rope=True, block_size=32) ++ old = torch.zeros(2, 1, STATE, D + 12, dtype=torch.bfloat16) ++ with pytest.raises(ValueError, match="width"): ++ cache.bind_kv_cache(old) ++ ++ ++# -------------------------------------------------------------------------- ++# the claim clear ++# -------------------------------------------------------------------------- ++ ++ ++def test_claim_clear_zeroes_only_the_tag_words_of_the_named_blocks(): ++ cache = _bound(rope=True, num_blocks=4) ++ # every byte non-zero, then remember it ++ cache.kv_cache.view(torch.int16).fill_(0x1234) ++ storage_before = cache.kv_cache.clone() ++ cache.clear_ring_row_tags(torch.tensor([1, 3])) ++ after = cache.kv_cache ++ tags = cache.ring_tag_cache[..., 0] ++ assert torch.all(tags[[1, 3]] == 0) ++ assert torch.all(tags[[0, 2]] != 0) ++ # nothing but the tag words changed ++ changed = after.view(torch.int16) != storage_before.view(torch.int16) ++ changed_cols = changed.nonzero()[:, -1].unique().tolist() ++ tag_cols = list(range(D + 4 * QSA_RING_TAG_WORD, D + 4 * QSA_RING_TAG_WORD + 4)) ++ assert changed_cols == tag_cols ++ assert changed.nonzero()[:, 0].unique().tolist() == [1, 3] ++ ++ ++def test_claim_clear_before_bind_is_loud(): ++ cache = _key_cache(rope=True, block_size=32) ++ with pytest.raises(RuntimeError, match="before the KV cache was bound"): ++ cache.clear_ring_row_tags(torch.tensor([0])) ++ ++ ++# -------------------------------------------------------------------------- ++# the reader's rule (torch reference of the kernel check) ++# -------------------------------------------------------------------------- ++ ++ ++def _group(chunk_start: int): ++ """The first group of a chunk starting at chunk_start (the only one that ++ can cross): its four source positions and which of them are historical.""" ++ first_boundary = ((chunk_start + CR) // CR) * CR - 1 ++ positions = torch.arange(first_boundary - CR + 1, first_boundary + 1) ++ return positions, positions < chunk_start ++ ++ ++def _committed_ring(positions): ++ """found tags for rows this request committed at exactly these positions""" ++ return qsa_ring_row_tags(positions) ++ ++ ++def test_tags_are_logical_position_plus_one_and_never_zero(): ++ tags = qsa_ring_row_tags(torch.tensor([0, 1, 52933])) ++ assert tags.dtype == torch.int64 ++ assert tags.tolist() == [1, 2, 52934] ++ ++ ++@pytest.mark.parametrize("chunk_start", [52933, 52934, 52935, 52936]) ++def test_committed_history_validates(chunk_start): ++ positions, historical = _group(chunk_start) ++ ok = qsa_ring_group_rows_valid( ++ _committed_ring(positions)[None], positions[None], historical[None] ++ ) ++ assert ok.tolist() == [True] ++ ++ ++@pytest.mark.parametrize("bad_row", [1, 2], ids=["row2", "row3"]) ++def test_every_row_is_checked_not_only_the_first(bad_row): ++ """First row valid, only the 2nd or 3rd row of the group stale -- an ++ in-range position and key from the previous lap of the ring (the row ++ holds position p-8). Guard 9 is blind to this (the position is in ++ range); a first-row-only check is blind to it too.""" ++ chunk_start = 52935 # 3 historical rows: 52932, 52933, 52934 ++ positions, historical = _group(chunk_start) ++ assert historical.tolist() == [True, True, True, False] ++ found = _committed_ring(positions).clone() ++ found[bad_row] = positions[bad_row] - STATE + 1 # previous lap's tag ++ ok = qsa_ring_group_rows_valid(found[None], positions[None], historical[None]) ++ assert ok.tolist() == [False] ++ # ...and the first row alone would have passed ++ first_only = qsa_ring_group_rows_valid( ++ found[None, :1], positions[None, :1], historical[None, :1] ++ ) ++ assert first_only.tolist() == [True] ++ ++ ++@pytest.mark.parametrize( ++ "found_tag", ++ [0, 0x3E803E803E803E80, -1, 7], ++ ids=["cleared", "bf16-filler-residue", "negative-residue", "other-position"], ++) ++def test_uncommitted_rows_fail(found_tag): ++ positions, historical = _group(52934) ++ found = torch.full_like(positions, found_tag) ++ ok = qsa_ring_group_rows_valid(found[None], positions[None], historical[None]) ++ assert ok.tolist() == [False] ++ ++ ++def test_in_chunk_rows_are_never_checked(): ++ # aligned chunk start: no historical row, garbage in the ring is irrelevant ++ positions, historical = _group(52936) ++ assert not historical.any() ++ found = torch.full_like(positions, 0x3E803E803E803E80) ++ ok = qsa_ring_group_rows_valid(found[None], positions[None], historical[None]) ++ assert ok.tolist() == [True] ++ ++ ++def test_previous_occupant_same_position_needs_the_claim_clear(): ++ """The collision the claim clear exists for: the previous occupant ++ committed the SAME logical positions into the same rows, so its tags are ++ exactly the ones the new owner expects. The reader alone cannot tell; ++ after the clear the rows read 0 and fail.""" ++ positions, historical = _group(52934) ++ found = _committed_ring(positions) # left by the previous occupant ++ assert qsa_ring_group_rows_valid( ++ found[None], positions[None], historical[None] ++ ).tolist() == [True] ++ cache = _bound(rope=True, num_blocks=2) ++ for p, tag in zip(positions.tolist(), found.tolist()): ++ cache.ring_tag_cache[1, p % STATE, 0, 0] = tag ++ cache.clear_ring_row_tags(torch.tensor([1])) ++ after = cache.ring_tag_cache[1, positions % STATE, 0, 0] ++ assert qsa_ring_group_rows_valid( ++ after[None], positions[None], historical[None] ++ ).tolist() == [False] ++ ++ ++# -------------------------------------------------------------------------- ++# telemetry site ++# -------------------------------------------------------------------------- ++ ++ ++def test_site_is_declared_named_and_fatal(): ++ site = gt.SITE_QSA_RING_UNCOMMITTED ++ assert site == 11 and gt.NUM_SITES == 12 ++ assert "UNCOMMITTED" in gt.SITE_NAMES[site] ++ assert site in gt._ALWAYS_FATAL_SITES ++ ctr = gt.counter(site, torch.device("cpu")) ++ assert ctr.dtype == torch.int64 and ctr.numel() == gt.SLOTS_PER_SITE ++ assert ctr.data_ptr() != gt.counter(gt.SITE_QSA_RING_SKIP, "cpu").data_ptr() ++ ++ ++@pytest.mark.parametrize( ++ "layer,row,pos", [(3, 0, 52932), (48, 255, 262143), (-1, 7, 0)] ++) ++def test_record_packing_round_trips(layer, row, pos): ++ biased = gt.encode_ring_record(layer, row, pos) + gt.GUARD_BIAS ++ assert 0 < biased <= (1 << 63) - 1 ++ assert gt.decode_ring_record(biased) == { ++ "layer": layer, ++ "request_row": row, ++ "source_position": pos, ++ } ++ assert gt.decode_ring_record(0) is None ++ ++ ++def test_reader_reports_one_coherent_record_and_arms(monkeypatch): ++ gt.reset_for_test() ++ msgs = [] ++ monkeypatch.setattr(gt.logger, "warning", lambda fmt, m: msgs.append(m)) ++ rows = [[0, 0, 0] for _ in range(gt.NUM_SITES)] ++ rows[gt.SITE_QSA_RING_UNCOMMITTED] = [ ++ 2, ++ gt.encode_ring_record(7, 1, 52933) + gt.GUARD_BIAS, ++ 0 + gt.GUARD_BIAS, # the row held tag 0: cleared, never committed ++ ] ++ gt._report_rows(torch.device("cpu"), rows) ++ assert len(msgs) == 1 ++ text = msgs[0] ++ assert "2 uncommitted" in text ++ assert "'layer': 7" in text and "'request_row': 1" in text ++ assert "'source_position': 52933" in text ++ assert "held tag 0" in text ++ with pytest.raises(RuntimeError, match="uncommitted"): ++ gt.check_pending() ++ gt.reset_for_test() ++ ++ ++# -------------------------------------------------------------------------- ++# KV capacity: the wider row at the production shape ++# -------------------------------------------------------------------------- ++ ++ ++def test_production_pool_blocks_unchanged_by_the_tag_word(): ++ """Qwen3.8-Flash-Next at TP2 x PP2 (25/23), MTP K=3, fp8 KV, block 3200, ++ --kv-cache-memory 4100000000: 317 blocks with the D+12 row and with the ++ D+16 row. The ring is a one-block-per-request state group whose pages ++ (6-7 layers x 2,304 B) sit in a 11.5-12.9 MB block, so four more bf16 per ++ row cannot change the block count or the per-request block need; the ++ reported pool (806,792 tokens on the reference host) therefore stays. (This offline ++ stand-in reports 839,390 for both widths: it does not annotate the draft ++ group, so its blocks-per-request differ from the reference host's; the DELTA is the ++ measurement.)""" ++ import tests.v1.test_pp_kv_projection as prod ++ from vllm.v1.core import kv_cache_utils as ku ++ ++ real = _key_cache(rope=True).get_kv_cache_spec( ++ SimpleNamespace(num_speculative_tokens=3) ++ ) ++ assert real.head_size == 144 # the row this tree actually allocates ++ results = {} ++ for name, ring in ( ++ ( ++ # instrument check: an absurdly wide ring DOES move the capacity, ++ # so equality below is a measurement, not a probe that ignores ++ # the ring (the ring group splits and costs blocks per request) ++ "sensitivity", ++ CircularBufferSpec( ++ block_size=8, ++ num_kv_heads=1, ++ head_size=262144, ++ head_size_v=0, ++ dtype=torch.bfloat16, ++ ), ++ ), ++ ( ++ "pre_tag", ++ CircularBufferSpec( ++ block_size=8, ++ num_kv_heads=1, ++ head_size=140, ++ head_size_v=0, ++ dtype=torch.bfloat16, ++ ), ++ ), ++ ("tagged", real), ++ ): ++ saved = prod._RING ++ prod._RING = ring ++ try: ++ specs = [ ++ prod._stage(0, 25, mtp=False), ++ prod._stage(0, 25, mtp=False), ++ prod._stage(25, 48, mtp=True), ++ prod._stage(25, 48, mtp=True), ++ ] ++ cfg = prod._config() ++ configs = ku.get_kv_cache_configs(cfg, specs, [4_100_000_000] * 4) ++ sched = ku.generate_scheduler_kv_cache_config(configs) ++ results[name] = ( ++ [c.num_blocks for c in configs], ++ ku.get_kv_cache_capacity(cfg, sched)[0], ++ ) ++ finally: ++ prod._RING = saved ++ assert results["sensitivity"][1] < results["pre_tag"][1] ++ assert results["pre_tag"][0] == [317] * 4 ++ assert results["tagged"] == results["pre_tag"] +diff --git a/tests/v1/core/test_qsa_ring_resume_alignment.py b/tests/v1/core/test_qsa_ring_resume_alignment.py +new file mode 100644 +index 0000000..9bc5cf4 +--- /dev/null ++++ b/tests/v1/core/test_qsa_ring_resume_alignment.py +@@ -0,0 +1,266 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Ring-row tag fix: can a prefix-cache hit make a request's ++FIRST QSA group read raw-key ring rows this request never wrote? ++ ++The ring (CircularBufferSpec) is not prefix-cacheable: a request that resumes ++from a cache hit at position k starts with an empty ring. The pre-indexer's ++first group reads ring rows iff k is not a multiple of the compression ratio R ++(rows [k - k % R, k)). The ring-row tag fix makes such a read rank-fatal, so a legitimate ++unaligned resume would be a false positive. This file pins, on the local ++(no KV connector) path the model runs in production, that k % R == 0 always: ++ ++ 1. the ring never takes part in hit lookup (prefix_cacheable False, skipped ++ by the coordinator, its finder returns 0) -- kv_cache_coordinator.py ++ attention_groups loop, single_type_kv_cache_manager.py ++ CircularBufferManager.find_longest_cache_hit; ++ 2. every hit length is a multiple of the coordinator's cache-hit alignment ++ (scheduler block, or the hash block under Mamba "align" partial hits), ++ including the full-prompt "recompute the last token" cap and the EAGLE ++ block drop -- kv_cache_manager.py get_computed_blocks; ++ 3. boot refuses any alignment that is not a multiple of every ++ prefix-cacheable tokens_per_state, i.e. of R -- kv_cache_utils.py ++ resolve_kv_cache_block_sizes; ++ 4. preemption restarts from 0 -- scheduler.py _preempt_request. ++ ++The one path that CAN resume unaligned is an external KV connector ++(num_external_computed_tokens is arbitrary; a full-prompt async load resets ++to num_tokens - 1, scheduler.py _update_waiting_for_remote_kv). There the ring ++history must come from the connector: a connector that transfers the ring ++brings the producer's committed rows and tags (and the tag fix no longer clears the ++tags of a block an async load is filling); one that does not has no history ++to give, and failing closed is the correct outcome -- the recurrent GDN state ++at k cannot be rewound to an aligned k' < k. Production runs no connector. ++""" ++ ++import ast ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.sampling_params import SamplingParams ++from vllm.utils.hashing import sha256 ++from vllm.v1.core import kv_cache_utils as ku ++from vllm.v1.core.kv_cache_manager import KVCacheManager ++from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash ++from vllm.v1.kv_cache_interface import ( ++ CircularBufferSpec, ++ FullAttentionSpec, ++ KVCacheConfig, ++ KVCacheGroupSpec, ++ MambaSpec, ++ MLAAttentionSpec, ++) ++from vllm.v1.request import Request ++ ++pytestmark = pytest.mark.cpu_test ++ ++R = 4 # QSA compress ratio == MLAAttentionSpec.tokens_per_state ++BLOCK = 16 # stands in for the production 3200 (both multiples of R) ++REPO = Path(__file__).resolve().parents[3] ++ ++ ++def first_group_ring_rows(chunk_start: int, ratio: int = R) -> list[int]: ++ """Ring rows the pre-indexer's first group reads for a chunk starting at ++ chunk_start -- the kernel's arithmetic (qsa_pre_indexer.py: first_boundary, ++ first_position, source_in_chunk).""" ++ first_boundary = ((chunk_start + ratio) // ratio) * ratio - 1 ++ first_position = first_boundary - (ratio - 1) ++ return [p for p in range(first_position, first_position + ratio) if p < chunk_start] ++ ++ ++def test_the_checker_can_fail(): ++ """Negative control for everything below: an unaligned resume (the ++ connector full-prompt reset to num_tokens - 1) DOES read ring rows.""" ++ assert first_group_ring_rows(52934) == [52932, 52933] ++ assert first_group_ring_rows(53 - 1) == [] # 52 is aligned ++ assert first_group_ring_rows(54 - 1) == [52] ++ for k in range(0, 64, R): ++ assert first_group_ring_rows(k) == [] ++ ++ ++def _config() -> KVCacheConfig: ++ return KVCacheConfig( ++ num_blocks=256, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ KVCacheGroupSpec( ++ ["attn"], ++ FullAttentionSpec( ++ block_size=BLOCK, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ), ++ ), ++ KVCacheGroupSpec( ++ ["comp"], ++ MLAAttentionSpec( ++ block_size=BLOCK, ++ num_kv_heads=1, ++ head_size=1, ++ dtype=torch.float32, ++ tokens_per_state=R, ++ ), ++ ), ++ KVCacheGroupSpec( ++ ["ring"], ++ CircularBufferSpec( ++ block_size=8, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ), ++ ), ++ ], ++ ) ++ ++ ++def _request(rid: str, tokens: list[int], hash_block: int) -> Request: ++ params = SamplingParams(max_tokens=4) ++ params.update_from_generation_config({}, eos_token_id=100) ++ return Request( ++ request_id=rid, ++ prompt_token_ids=tokens, ++ mm_features=None, ++ sampling_params=params, ++ pooling_params=None, ++ block_hasher=get_request_block_hasher(hash_block, sha256), ++ ) ++ ++ ++def _mamba_align_group() -> KVCacheGroupSpec: ++ return KVCacheGroupSpec( ++ ["gdn"], ++ MambaSpec( ++ shapes=((4,),), ++ dtypes=(torch.float32,), ++ block_size=BLOCK, ++ mamba_cache_mode="align", ++ ), ++ ) ++ ++ ++def _sweep_hits(*, use_eagle: bool, hash_block: int, mamba_align: bool) -> set[int]: ++ """Every hit length the real KVCacheManager returns for prompts of every ++ length in a range, for both a full-prompt repeat (capped to ++ num_tokens - 1) and a prompt continuation.""" ++ init_none_hash(sha256) ++ config = _config() ++ if mamba_align: ++ config.kv_cache_groups.append(_mamba_align_group()) ++ assert not config.kv_cache_groups[2].kv_cache_spec.prefix_cacheable ++ hits = set() ++ for length in range(BLOCK + 1, 5 * BLOCK + 3): ++ mgr = KVCacheManager( ++ config, ++ max_model_len=4096, ++ scheduler_block_size=BLOCK, ++ hash_block_size=hash_block, ++ enable_caching=True, ++ use_eagle=use_eagle, ++ ) ++ prompt = list(range(1000, 1000 + length)) ++ first = _request("a", prompt, hash_block) ++ blocks, hit, _ = mgr.get_computed_blocks(first) ++ assert hit == 0 ++ assert mgr.allocate_slots(first, length, hit, blocks) is not None ++ mgr.free(first) ++ for rid, tokens in (("b", prompt), ("c", prompt + [7, 8, 9])): ++ req = _request(rid, tokens, hash_block) ++ blocks, hit, _ = mgr.get_computed_blocks(req) ++ hits.add(hit) ++ # the ring contributes no cached block: its history starts empty ++ assert len(blocks.blocks[2]) == 0 ++ return hits ++ ++ ++@pytest.mark.parametrize( ++ "use_eagle,hash_block,mamba_align", ++ [ ++ (False, BLOCK, False), ++ (True, BLOCK, False), ++ # Mamba "align" fine-grained (partial) hits: lengths 20, 24, 28, ... ++ # (EAGLE + Mamba needs the eagle-group annotation this stand-in does ++ # not model -- it returns no hit at all, so it is not swept here.) ++ (False, R, True), ++ ], ++ ids=["block", "block-eagle", "partial-hits"], ++) ++def test_every_local_prefix_hit_resumes_on_a_group_boundary( ++ use_eagle, hash_block, mamba_align ++): ++ hits = _sweep_hits( ++ use_eagle=use_eagle, hash_block=hash_block, mamba_align=mamba_align ++ ) ++ for hit in hits: ++ assert hit % R == 0 and first_group_ring_rows(hit) == [], hit ++ # the sweep really produced non-trivial hits (else it proves nothing) ++ assert max(hits) >= 3 * BLOCK ++ if mamba_align: # ...and really partial ones ++ assert any(hit % BLOCK for hit in hits) ++ ++ ++def test_failing_control_a_finer_hash_block_would_open_the_path(): ++ """What the boot check below closes: with partial hits on a hash block of ++ 2 (not a multiple of R) the SAME sweep returns hits that make the first ++ group read ring rows this request never wrote.""" ++ hits = _sweep_hits(use_eagle=False, hash_block=2, mamba_align=True) ++ assert any(first_group_ring_rows(hit) for hit in hits), sorted(hits) ++ ++ ++def _boot_config(prefix_match_unit): ++ return SimpleNamespace( ++ cache_config=SimpleNamespace( ++ block_size=BLOCK, ++ enable_prefix_caching=True, ++ prefix_match_unit=prefix_match_unit, ++ ), ++ parallel_config=SimpleNamespace(decode_context_parallel_size=1), ++ kv_transfer_config=None, ++ ) ++ ++ ++def _boot_groups(): ++ return KVCacheConfig( ++ num_blocks=64, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ *_config().kv_cache_groups, ++ KVCacheGroupSpec( ++ ["gdn"], ++ MambaSpec( ++ shapes=((4,),), ++ dtypes=(torch.float32,), ++ block_size=BLOCK, ++ mamba_cache_mode="align", ++ ), ++ ), ++ ], ++ ) ++ ++ ++def test_boot_refuses_a_hit_alignment_finer_than_the_qsa_group(): ++ """With Mamba "align" partial hits the alignment drops to the hash block; ++ a hash block that is not a multiple of R would allow an unaligned resume, ++ and boot refuses it.""" ++ with pytest.raises(ValueError, match="per-state compression"): ++ ku.resolve_kv_cache_block_sizes(_boot_groups(), _boot_config(2)) ++ # control: an R-multiple hash block is accepted and stays R-aligned ++ sched, hash_bs = ku.resolve_kv_cache_block_sizes(_boot_groups(), _boot_config(4)) ++ assert sched == BLOCK and hash_bs % R == 0 ++ ++ ++def _function(path: Path, name: str) -> ast.FunctionDef: ++ tree = ast.parse(path.read_text(encoding="utf-8")) ++ for node in ast.walk(tree): ++ if isinstance(node, ast.FunctionDef) and node.name == name: ++ return node ++ raise AssertionError(f"{name} not found in {path}") ++ ++ ++def test_preemption_restarts_from_zero(): ++ fn = _function(REPO / "vllm/v1/core/sched/scheduler.py", "_preempt_request") ++ assigns = [ ++ ast.unparse(n) ++ for n in ast.walk(fn) ++ if isinstance(n, ast.Assign) ++ and any("num_computed_tokens" in ast.unparse(t) for t in n.targets) ++ ] ++ assert "request.num_computed_tokens = 0" in assigns +diff --git a/tests/v1/core/test_ring_tag_clear.py b/tests/v1/core/test_ring_tag_clear.py +new file mode 100644 +index 0000000..0c5b3d5 +--- /dev/null ++++ b/tests/v1/core/test_ring_tag_clear.py +@@ -0,0 +1,367 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""v2.2 ring-row tag fix, manager half: the unconditional, ordered ring tag clear. ++ ++`_record_new_block_ids` keeps `CircularBufferSpec` out of whole-block zeroing ++(single_type_kv_cache_manager.py:117-121), so a re-claimed ring block keeps ++its previous occupant's bytes. The tag fix does not change that; it records EVERY ++ring claim on a separate list, unconditionally (never under ++`needs_kv_cache_zeroing`), ships it as `SchedulerOutput.new_ring_block_ids_to_clear` ++and the runner clears the rows' TAG words right after the zeroer and before ++copies and the forward -- i.e. before the new owner's first ring access. ++ ++FAILING CONTROL: on the unmodified tree (without this fix) every ++test here fails -- `take_new_ring_block_ids`, the SchedulerOutput field and ++`vllm.v1.worker.ring_tags` do not exist, and nothing calls a ring clear. ++""" ++ ++import ast ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.v1.core.block_pool import BlockPool ++from vllm.v1.core.kv_cache_manager import KVCacheManager ++from vllm.v1.core.sched.output import SchedulerOutput ++from vllm.v1.core.single_type_kv_cache_manager import ( ++ CircularBufferManager, ++ FullAttentionManager, ++) ++from vllm.v1.kv_cache_interface import CircularBufferSpec, FullAttentionSpec ++from vllm.v1.worker.ring_tags import clear_ring_row_tags, ring_tag_layers ++ ++pytestmark = pytest.mark.cpu_test ++ ++REPO = Path(__file__).resolve().parents[3] ++ ++ ++def _ring_manager(pool: BlockPool, *, zeroing: bool) -> CircularBufferManager: ++ spec = CircularBufferSpec( ++ block_size=8, num_kv_heads=1, head_size=144, dtype=torch.bfloat16 ++ ) ++ return CircularBufferManager( ++ spec, ++ block_pool=pool, ++ enable_caching=False, ++ kv_cache_group_id=0, ++ scheduler_block_size=8, ++ needs_kv_cache_zeroing=zeroing, ++ max_admission_blocks_per_request=1, ++ ) ++ ++ ++@pytest.mark.parametrize("zeroing", [False, True], ids=["no-zeroing", "zeroing"]) ++def test_every_ring_claim_is_recorded_whatever_the_zeroing_gate(zeroing): ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=8) ++ mgr = _ring_manager(pool, zeroing=zeroing) ++ (a,) = mgr.allocate_new_blocks("a", num_tokens=100, num_tokens_main_model=100) ++ (b,) = mgr.allocate_new_blocks("b", num_tokens=3, num_tokens_main_model=3) ++ # later steps of the same request claim nothing new ++ assert mgr.allocate_new_blocks("a", 200, 200) == [] ++ assert mgr.take_new_ring_block_ids() == [a.block_id, b.block_id] ++ # drained: a second take is empty ++ assert mgr.take_new_ring_block_ids() == [] ++ # the ring stays OUT of whole-block zeroing (unchanged behaviour) ++ assert not mgr.records_new_block_ids ++ assert mgr.take_new_block_ids() == [] ++ ++ ++def test_external_claim_is_recorded_too(): ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=8) ++ mgr = _ring_manager(pool, zeroing=False) ++ mgr.allocate_external_computed_blocks("x", 16, 16) ++ (block,) = mgr.req_to_blocks["x"] ++ assert mgr.take_new_ring_block_ids() == [block.block_id] ++ ++ ++def test_a_block_freed_before_the_drain_is_not_cleared(): ++ """Claimed then freed in the same scheduling pass (the request was ++ preempted after its allocation): the block may belong to another group ++ by the time the worker runs, so the clear must not touch it.""" ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=8) ++ mgr = _ring_manager(pool, zeroing=False) ++ (a,) = mgr.allocate_new_blocks("a", 8, 8) ++ (b,) = mgr.allocate_new_blocks("b", 8, 8) ++ mgr.free("a") ++ assert mgr.take_new_ring_block_ids() == [b.block_id] ++ ++ ++def test_a_block_freed_and_reclaimed_in_one_pass_is_cleared_once(): ++ pool = BlockPool(num_gpu_blocks=2, enable_caching=False, hash_block_size=8) ++ mgr = _ring_manager(pool, zeroing=False) # one usable block (0 is null) ++ (a,) = mgr.allocate_new_blocks("a", 8, 8) ++ mgr.free("a") ++ (b,) = mgr.allocate_new_blocks("b", 8, 8) ++ assert b.block_id == a.block_id # same physical block, new owner ++ assert mgr.take_new_ring_block_ids() == [b.block_id] ++ ++ ++def test_non_ring_managers_report_nothing(): ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=4) ++ spec = FullAttentionSpec( ++ block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ) ++ mgr = FullAttentionManager( ++ spec, ++ block_pool=pool, ++ enable_caching=False, ++ kv_cache_group_id=0, ++ scheduler_block_size=4, ++ needs_kv_cache_zeroing=True, ++ ) ++ mgr.allocate_new_blocks("a", 16, 16) ++ assert mgr.take_new_ring_block_ids() == [] ++ assert mgr.take_new_block_ids() # attention still goes to the zeroer ++ ++ ++def test_kv_cache_manager_aggregates_every_group(): ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=8) ++ ring_a, ring_b = ( ++ _ring_manager(pool, zeroing=False), ++ _ring_manager(pool, zeroing=True), ++ ) ++ (x,) = ring_a.allocate_new_blocks("r", 8, 8) ++ (y,) = ring_b.allocate_new_blocks("r", 8, 8) ++ fake = SimpleNamespace( ++ coordinator=SimpleNamespace(single_type_managers=(ring_a, ring_b)) ++ ) ++ assert KVCacheManager.take_new_ring_block_ids(fake) == [x.block_id, y.block_id] ++ assert KVCacheManager.take_new_ring_block_ids(fake) == [] ++ ++ ++def test_scheduler_output_carries_the_field_default_none(): ++ field = SchedulerOutput.__dataclass_fields__["new_ring_block_ids_to_clear"] ++ assert field.default is None ++ ++ ++# -------------------------------------------------------------------------- ++# wiring, checked on the source (the scheduler and runners need a model / ++# Triton to construct; the behavioural twin of the scheduler check is ++# test_scheduler_ships_ring_claims_without_kv_zeroing below) ++# -------------------------------------------------------------------------- ++ ++ ++def _function(path: Path, name: str) -> ast.FunctionDef: ++ tree = ast.parse(path.read_text(encoding="utf-8")) ++ for node in ast.walk(tree): ++ if isinstance(node, ast.FunctionDef) and node.name == name: ++ return node ++ raise AssertionError(f"{name} not found in {path}") ++ ++ ++def test_scheduler_drains_ring_claims_unconditionally(): ++ path = REPO / "vllm/v1/core/sched/scheduler.py" ++ fn = _function(path, "schedule") ++ found = False ++ for node in ast.walk(fn): ++ if isinstance(node, ast.keyword) and node.arg == "new_ring_block_ids_to_clear": ++ src = ast.unparse(node.value) ++ assert "_get_new_ring_block_ids_to_clear" in src ++ assert "needs_kv_cache_zeroing" not in src ++ found = True ++ assert found, "SchedulerOutput is built without new_ring_block_ids_to_clear" ++ body = _function(path, "_get_new_ring_block_ids_to_clear").body ++ if isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): ++ body = body[1:] # the docstring may name the gate it is not under ++ drain = "\n".join(ast.unparse(stmt) for stmt in body) ++ assert "take_new_ring_block_ids" in drain ++ assert "needs_kv_cache_zeroing" not in drain ++ ++ ++def _drain_scheduler(ring_ids, skip): ++ from vllm.v1.core.sched.scheduler import Scheduler ++ ++ sched = Scheduler.__new__(Scheduler) ++ sched.kv_cache_manager = SimpleNamespace( ++ take_new_ring_block_ids=lambda: list(ring_ids) ++ ) ++ sched._skip_ring_clear_block_ids = set(skip) ++ return sched ++ ++ ++def test_drain_ships_every_claim_and_skips_only_async_load_targets(): ++ sched = _drain_scheduler([3, 5, 9], skip={5}) ++ assert sched._get_new_ring_block_ids_to_clear() == [3, 9] ++ # the skip set is one-shot, like _skip_zero_block_ids ++ assert sched._skip_ring_clear_block_ids == set() ++ assert _drain_scheduler([], skip=set())._get_new_ring_block_ids_to_clear() is None ++ ++ ++def test_async_kv_load_marks_the_request_ring_blocks_for_skipping(): ++ """The async-load branch of schedule() records the request's ring blocks ++ in _skip_ring_clear_block_ids OUTSIDE the needs_kv_cache_zeroing guard ++ (an out-of-band write of the producer's ring must not be followed by a ++ tag clear that wipes its committed tags).""" ++ fn = _function(REPO / "vllm/v1/core/sched/scheduler.py", "schedule") ++ hits = [] ++ for node in ast.walk(fn): ++ if isinstance(node, ast.Call) and "_skip_ring_clear_block_ids.update" in ( ++ ast.unparse(node.func) ++ ): ++ hits.append(ast.unparse(node)) ++ assert hits and all("get_ring_block_ids" in h for h in hits) ++ for node in ast.walk(fn): ++ if isinstance(node, ast.If) and "needs_kv_cache_zeroing" in ast.unparse( ++ node.test ++ ): ++ assert "_skip_ring_clear_block_ids" not in ast.unparse(node) ++ ++ ++def test_ring_block_ids_of_a_request(): ++ pool = BlockPool(num_gpu_blocks=10, enable_caching=False, hash_block_size=8) ++ ring = _ring_manager(pool, zeroing=False) ++ (x,) = ring.allocate_new_blocks("r", 8, 8) ++ assert ring.get_ring_block_ids("r") == [x.block_id] ++ assert ring.get_ring_block_ids("absent") == [] ++ fake = SimpleNamespace(coordinator=SimpleNamespace(single_type_managers=(ring,))) ++ assert KVCacheManager.get_ring_block_ids(fake, "r") == [x.block_id] ++ ++ ++@pytest.mark.parametrize( ++ "path,func,zero_call,copy_call", ++ [ ++ ( ++ "vllm/v1/worker/gpu/model_runner.py", ++ "update_requests", ++ "zero_block_ids", ++ "copy_kv_cache_blocks_inplace", ++ ), ++ ( ++ "vllm/v1/worker/gpu_model_runner.py", ++ "_update_states", ++ "_zero_block_ids", ++ "copy_kv_cache_blocks_inplace", ++ ), ++ ], ++ ids=["v2-runner", "v1-runner"], ++) ++def test_runner_clears_after_the_zeroer_and_before_copies( ++ path, func, zero_call, copy_call ++): ++ fn = _function(REPO / path, func) ++ lines: dict[str, int] = {} ++ for node in ast.walk(fn): ++ if isinstance(node, ast.Call): ++ name = ast.unparse(node.func).rsplit(".", 1)[-1] ++ lines[name] = min(lines.get(name, node.lineno), node.lineno) ++ assert "clear_ring_row_tags" in lines, f"{path}:{func} never clears ring tags" ++ assert lines[zero_call] < lines["clear_ring_row_tags"] < lines[copy_call] ++ # and the clear is not nested under the zeroing branch ++ for node in ast.walk(fn): ++ if isinstance(node, ast.If) and "new_block_ids_to_zero" in ast.unparse( ++ node.test ++ ): ++ inner = {ast.unparse(n) for n in ast.walk(node)} ++ assert not any("clear_ring_row_tags" in s for s in inner) ++ ++ ++# -------------------------------------------------------------------------- ++# the worker helper ++# -------------------------------------------------------------------------- ++ ++ ++class _TagLayer: ++ def __init__(self, num_blocks: int, prefix: str = "layer"): ++ self.prefix = prefix ++ self.ring_tag_cache = torch.arange( ++ 1, num_blocks * 8 + 1, dtype=torch.int64 ++ ).view(num_blocks, 8, 1, 1) ++ self.cleared: list[list[int]] = [] ++ ++ def clear_ring_row_tags(self, block_ids: torch.Tensor) -> None: ++ self.cleared.append(block_ids.tolist()) ++ self.ring_tag_cache.index_fill_(0, block_ids, 0) ++ ++ ++def test_helper_clears_every_tag_layer_and_skips_others(): ++ a, b = _TagLayer(4, "a"), _TagLayer(4, "b") ++ ctx = {"a": a, "attn": object(), "b": b} ++ layers = ring_tag_layers(ctx) ++ assert layers == [a, b] ++ assert clear_ring_row_tags(layers, [2, 0], torch.device("cpu")) == 2 ++ for layer in (a, b): ++ assert layer.cleared == [[2, 0]] ++ assert torch.all(layer.ring_tag_cache[[0, 2]] == 0) ++ assert torch.all(layer.ring_tag_cache[[1, 3]] != 0) ++ ++ ++def test_helper_is_a_noop_without_ids_or_layers(): ++ a = _TagLayer(2) ++ assert clear_ring_row_tags([a], [], torch.device("cpu")) == 0 ++ assert clear_ring_row_tags([], [1], torch.device("cpu")) == 0 ++ assert a.cleared == [] ++ ++ ++def test_helper_fails_loud_on_a_block_outside_the_pool(): ++ a = _TagLayer(2) ++ with pytest.raises(ValueError, match="outside the pool"): ++ clear_ring_row_tags([a], [1, 2], torch.device("cpu")) ++ assert a.cleared == [] ++ ++ ++def test_helper_fails_loud_on_an_unbound_layer(): ++ a = _TagLayer(2) ++ a.ring_tag_cache = None ++ with pytest.raises(RuntimeError, match="no bound ring_tag_cache"): ++ clear_ring_row_tags([a], [1], torch.device("cpu")) ++ ++ ++def test_clear_then_commit_orders_correctly_on_a_real_ring(): ++ """Claim -> clear -> the new owner's commit: the commit survives (the ++ clear never lands after it), and an uncommitted row reads tag 0.""" ++ from vllm.models.qwen4_exp.common import qsa_cache ++ ++ cache = qsa_cache.QSAKeyStateCache( ++ head_size=128, ++ dtype=torch.bfloat16, ++ cache_rope_positions=True, ++ cache_config=SimpleNamespace(block_size=32), ++ prefix="ringtag.order", ++ vllm_config=SimpleNamespace( ++ compilation_config=SimpleNamespace(static_forward_context={}) ++ ), ++ compress_ratio=4, ++ ) ++ unified = torch.full((3, 1, 8, 144), 0.25, dtype=torch.bfloat16) ++ cache.bind_kv_cache(unified) ++ residue = int(cache.ring_tag_cache[1, 0, 0, 0]) ++ assert residue != 0 ++ # step N: scheduler claimed block 1; the runner clears before the forward ++ clear_ring_row_tags(ring_tag_layers({"l": cache}), [1], torch.device("cpu")) ++ # step N forward: the new owner commits positions 100, 101 ++ for position in (100, 101): ++ cache.ring_tag_cache[1, position % 8, 0, 0] = position + 1 ++ tags = cache.ring_tag_cache[1, :, 0, 0].tolist() ++ assert tags[100 % 8] == 101 and tags[101 % 8] == 102 ++ assert [t for i, t in enumerate(tags) if i not in (100 % 8, 101 % 8)] == [0] * 6 ++ # other blocks untouched ++ assert int(cache.ring_tag_cache[0, 0, 0, 0]) == residue ++ ++ ++# -------------------------------------------------------------------------- ++# behavioural scheduler check (needs the HF config of the test model: runs on ++# the reference host / CI, not offline without the model) ++# -------------------------------------------------------------------------- ++ ++ ++def test_scheduler_ships_ring_claims_without_kv_zeroing(): ++ from .utils import create_requests, create_scheduler ++ ++ spec = CircularBufferSpec( ++ block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ) ++ scheduler = create_scheduler(kv_cache_spec=spec) ++ assert not scheduler.needs_kv_cache_zeroing ++ requests = create_requests(num_requests=2, num_tokens=10) ++ scheduler.add_request(requests[0]) ++ out0 = scheduler.schedule() ++ blocks0 = scheduler.kv_cache_manager.get_block_ids(requests[0].request_id)[0] ++ assert out0.new_ring_block_ids_to_clear == list(blocks0) ++ assert out0.new_block_ids_to_zero is None ++ scheduler.add_request(requests[1]) ++ out1 = scheduler.schedule() ++ blocks1 = scheduler.kv_cache_manager.get_block_ids(requests[1].request_id)[0] ++ # only the NEW claim, never the first request's block again ++ assert out1.new_ring_block_ids_to_clear == list(blocks1) +diff --git a/tests/v1/test_pp_kv_projection.py b/tests/v1/test_pp_kv_projection.py +index 2c699ab..2f8cc59 100644 +--- a/tests/v1/test_pp_kv_projection.py ++++ b/tests/v1/test_pp_kv_projection.py +@@ -51,9 +51,10 @@ _MAIN = FullAttentionSpec( + dtype=torch.uint8, + kv_quant_mode=get_kv_quant_mode("fp8_e4m3"), + ) +-# QSAKeyStateCache: ring = 4 * cdiv(4 + K, 4) = 8, head 128 + MRoPE tail 12. ++# QSAKeyStateCache: ring = 4 * cdiv(4 + K, 4) = 8, head 128 + int64 tail 16 ++# (pos_t, pos_h, pos_w, tag; the ring-row tag fix -- was 12 before the tag word). + _RING = CircularBufferSpec( +- block_size=8, num_kv_heads=1, head_size=140, head_size_v=0, dtype=torch.bfloat16 ++ block_size=8, num_kv_heads=1, head_size=144, head_size_v=0, dtype=torch.bfloat16 + ) + # QSACompressedKeyCache: bf16 indexer, compress ratio 4. + _COMP = MLAAttentionSpec( +diff --git a/tests/v1/worker/gpu/ring_tag_replay.py b/tests/v1/worker/gpu/ring_tag_replay.py +new file mode 100644 +index 0000000..bd09d85 +--- /dev/null ++++ b/tests/v1/worker/gpu/ring_tag_replay.py +@@ -0,0 +1,467 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Ring-row tag fix reference-host harness: the ring-residue variants, before and after the fix. ++ ++Derived from an earlier crash-replay harness. ++The decode row under test (52,934 computed, ++4 scheduled -> chunk start 2 mod 4, so its first group reads two ring rows), ++real metadata from build_qsa_metadata_triton, one process per variant (an ++illegal address is sticky and would void every later variant). ++ ++It runs unchanged on BOTH trees and says which one it is on: ++ * fixed (tag fix): rows are D+16 wide with a tag word; SITE_QSA_RING_UNCOMMITTED ++ exists; layer_id is passed. ++ * unfixed (pre-fix kernels, an earlier commit): rows D+12, no tag. ++ ++--disable-guard9 swaps _norm_rope for a clone WITHOUT the guard-9 clamp ++before the first launch (Triton resolves the callee from the kernel module's ++globals at compile time). The `guard9_probe` variant proves the swap took: ++it feeds a valid-tag row whose position is one past the cos/sin table, which ++guard 9 counts (> 0) and the clone does not (== 0). ++ ++Expected outcomes (ACCEPT column of the summary): ++ variant unfixed, guard 9 off unfixed, guard 9 on fixed ++ ring_realistic clean clean clean, site 0 ++ ring_uncommitted_row FAULT (memcheck/IMA) guard9 > 0 (masked) request_error ++ ring_untouched_residue FAULT or silent* guard9 >= 0* request_error ++ ring_every_row_2/_3 silent (blind) silent (blind) request_error ++ guard9_probe guard9 == 0 guard9 > 0 guard9 > 0, site 0 ++* untouched residue is whatever the allocator left; the crash replay measured it ++ reaching the rope load (11 memcheck reads), but a run where it decodes in ++ range is the silent-corruption outcome, not a clean one. ++"request_error" = no fault, SITE_QSA_RING_UNCOMMITTED count > 0, its record ++names request row 0 and a historical source position, guard 9 count 0, and ++the group's compressed slot left unwritten. ++ ++ python ring_tag_replay.py --self-test # CPU, geometry only ++ python ring_tag_replay.py [--disable-guard9] # all variants, isolated ++ compute-sanitizer --tool memcheck python ring_tag_replay.py --only X ++""" ++ ++from __future__ import annotations ++ ++import argparse ++import json ++import os ++import subprocess ++import sys ++ ++D = 128 ++HQ = 4 ++CR = 4 ++STATE = 8 ++COMP_PAGE = 800 ++NUM_BLOCKS = 317 ++TABLE_WIDTH = 82 ++MAX_MODEL_LEN = 262144 ++MROPE = (11, 11, 10) ++LAYER = 3 ++VICTIM_RING, FRESH_RING = 5, 60 ++ ++# (computed, scheduled) of the request under test ++SHAPES = { ++ "ring_realistic": (52934, 4), ++ "ring_uncommitted_row": (52934, 4), ++ "ring_untouched_residue": (52934, 4), ++ "ring_every_row_2": (52935, 4), ++ "ring_every_row_3": (52935, 4), ++ "guard9_probe": (52934, 4), ++} ++VARIANTS = list(SHAPES) ++ ++ ++def historical_positions(computed: int) -> list[int]: ++ """Ring rows the first group of a chunk starting at `computed` reads.""" ++ first = computed - computed % CR ++ return list(range(first, computed)) ++ ++ ++def self_test() -> int: ++ ok = True ++ ++ def check(name, cond): ++ nonlocal ok ++ ok = ok and bool(cond) ++ print(f" {'PASS' if cond else 'FAIL'} {name}") ++ ++ check("the decode row reads two ring rows", historical_positions(52934) == [52932, 52933]) ++ check( ++ "every-row shape reads three", ++ historical_positions(52935) == [52932, 52933, 52934], ++ ) ++ check("rows 4,5 for the decode row", [p % STATE for p in (52932, 52933)] == [4, 5]) ++ # the crash replay's victim committed positions 100/101, i.e. rows 4/5 -- the ++ # very rows the reader picks. Here the victim sits at 96/97 (rows 0/1) so ++ # the reader's rows stay residue, which is the case being asked about. ++ check("victim rows 96,97 -> 0,1, disjoint from 4,5", [96 % 8, 97 % 8] == [0, 1]) ++ check("the group ends inside the chunk", (52934 - 52934 % 4 + 3) < 52934 + 4) ++ print(f"self-test: {'PASS' if ok else 'FAIL'}") ++ return 0 if ok else 1 ++ ++ ++def _unguarded_norm_rope(): ++ """_norm_rope with the guard-9 block and clamp removed, same signature.""" ++ from vllm.triton_utils import tl, triton ++ ++ @triton.jit ++ def _norm_rope( ++ x, ++ pos_t, ++ pos_h, ++ pos_w, ++ cos_sin_ptr, ++ cos_sin_stride, ++ cos_sin_rows, ++ guard_ctr, ++ norm_weight_ptr, ++ eps, ++ IS_MROPE: tl.constexpr, ++ MROPE_H: tl.constexpr, ++ MROPE_W: tl.constexpr, ++ ): ++ TILE_T: tl.constexpr = x.shape[0] ++ TILE_H: tl.constexpr = x.shape[1] ++ DD: tl.constexpr = x.shape[2] ++ ROWS: tl.constexpr = TILE_T * TILE_H ++ HALF: tl.constexpr = DD // 2 ++ QUARTER: tl.constexpr = DD // 4 ++ pairs = tl.arange(0, QUARTER) ++ if IS_MROPE: ++ h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_H) ++ w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_W) ++ t_mask = ~(h_mask | w_mask) ++ base = cos_sin_ptr + pairs[None, :] ++ pos_rows = (pos_t, pos_h, pos_w) ++ axis_masks = (t_mask, h_mask, w_mask) ++ cos = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) ++ sin = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) ++ for axis in tl.static_range(3): ++ cos += tl.load( ++ base + pos_rows[axis][:, None] * cos_sin_stride, ++ mask=axis_masks[axis][None, :], ++ other=0, ++ ) ++ sin += tl.load( ++ base + pos_rows[axis][:, None] * cos_sin_stride + QUARTER, ++ mask=axis_masks[axis][None, :], ++ other=0, ++ ) ++ else: ++ cos = tl.load( ++ cos_sin_ptr + pos_t[:, None] * cos_sin_stride + pairs[None, :] ++ ) ++ sin = tl.load( ++ cos_sin_ptr + pos_t[:, None] * cos_sin_stride + QUARTER + pairs[None, :] ++ ) ++ cos = tl.reshape( ++ tl.broadcast_to(cos[:, None, :], (TILE_T, TILE_H, QUARTER)), (ROWS, QUARTER) ++ ) ++ sin = tl.reshape( ++ tl.broadcast_to(sin[:, None, :], (TILE_T, TILE_H, QUARTER)), (ROWS, QUARTER) ++ ) ++ x = tl.reshape(x, (ROWS, DD)).to(tl.float32) ++ weight = tl.load(norm_weight_ptr + tl.arange(0, DD)).to(tl.float32) + 1.0 ++ rrms = tl.rsqrt(tl.sum(x * x, axis=1) / DD + eps) ++ y = (x * rrms[:, None] * weight[None, :]).to(cos.dtype) ++ rotated, passthrough = tl.split( ++ tl.permute(tl.reshape(y, (ROWS, 2, HALF)), (0, 2, 1)) ++ ) ++ r0, r1 = tl.split( ++ tl.permute(tl.reshape(rotated, (ROWS, 2, QUARTER)), (0, 2, 1)) ++ ) ++ out0 = r0 * cos - r1 * sin ++ out1 = r1 * cos + r0 * sin ++ rotated = tl.reshape(tl.permute(tl.join(out0, out1), (0, 2, 1)), (ROWS, HALF)) ++ result = tl.reshape( ++ tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), (ROWS, DD) ++ ) ++ return tl.reshape(result, (TILE_T, TILE_H, DD)) ++ ++ return _norm_rope ++ ++ ++def run_variant(variant: str, disable_guard9: bool) -> dict: ++ import torch ++ ++ from vllm.models.qwen4_exp.common import qsa_cache as qc ++ from vllm.models.qwen4_exp.nvidia.ops import qsa_pre_indexer as qpi ++ from vllm.v1.attention.backend import CommonAttentionMetadata ++ from vllm.v1.worker import guard_telemetry as gt ++ ++ fixed = hasattr(gt, "SITE_QSA_RING_UNCOMMITTED") ++ width = D + (16 if fixed else 12) ++ rec: dict = { ++ "variant": variant, ++ "tree": "fixed" if fixed else "unfixed", ++ "guard9": "disabled" if disable_guard9 else "present", ++ } ++ if disable_guard9: ++ qpi._norm_rope = _unguarded_norm_rope() ++ device = torch.device("cuda:0") ++ torch.manual_seed(20260923) ++ gt.prealloc(device) ++ gt.reset_for_test() ++ ++ computed, sched = SHAPES[variant] ++ seq = computed + sched ++ qsl = torch.tensor([0, sched], dtype=torch.int32) ++ ring_tab = torch.full((1, 1), VICTIM_RING, dtype=torch.int32, device=device) ++ ++ def md_for(table): ++ return CommonAttentionMetadata( ++ query_start_loc=qsl.to(device), ++ query_start_loc_cpu=qsl, ++ seq_lens=torch.tensor([seq], dtype=torch.int32, device=device), ++ num_reqs=1, ++ num_actual_tokens=sched, ++ max_query_len=sched, ++ max_seq_len=seq, ++ block_table_tensor=table, ++ slot_mapping=torch.zeros(sched, dtype=torch.int64, device=device), ++ ) ++ ++ def bufs(): ++ return [ ++ torch.empty(sched, dtype=d, device=device) ++ for d in (torch.int32, torch.int64, torch.int32, torch.int64) ++ ] ++ ++ _t, logical, _v, state_slots = qc.build_qsa_metadata_triton( ++ md_for(ring_tab), ++ *bufs(), ++ storage_block_size=STATE, ++ compress_ratio=1, ++ circular_buffer_size=STATE, ++ k_work_metadata_buffer=None, ++ request_capacity=None, ++ ) ++ comp_tab = ( ++ ((torch.arange(TABLE_WIDTH, device=device) % (NUM_BLOCKS - 1)) + 1) ++ .to(torch.int32) ++ .reshape(1, -1) ++ ) ++ kwb = torch.empty( ++ max((sched + CR - 1) // CR, 1), 2, dtype=torch.int32, device=device ++ ) ++ _t2, _l2, _v2, comp_slots = qc.build_qsa_metadata_triton( ++ md_for(comp_tab), ++ *bufs(), ++ storage_block_size=COMP_PAGE, ++ compress_ratio=CR, ++ circular_buffer_size=0, ++ k_work_metadata_buffer=kwb, ++ request_capacity=8, ++ ) ++ ++ shape = (NUM_BLOCKS, STATE, 1, width) ++ if variant == "ring_untouched_residue": ++ dirty = torch.empty(shape, dtype=torch.bfloat16, device=device) ++ dirty.normal_(mean=0.0, std=6.0e4) ++ del dirty ++ ring = torch.empty(shape, dtype=torch.bfloat16, device=device) ++ else: ++ ring = torch.zeros(shape, dtype=torch.bfloat16, device=device) ++ # 64 rows of slack past the table the kernel is told about, so the ++ # guard9_probe read one row past the end stays inside this allocation ++ cos_sin = torch.randn( ++ (MAX_MODEL_LEN + 64, D // 2), dtype=torch.bfloat16, device=device ++ )[:MAX_MODEL_LEN] ++ comp = torch.full( ++ (NUM_BLOCKS, COMP_PAGE, 1, D), 7.0, dtype=torch.bfloat16, device=device ++ ) ++ ones = torch.ones(D, dtype=torch.bfloat16, device=device) ++ ++ def write_row(position: int, pos_value: int | None, tag: int | None): ++ row = ring[VICTIM_RING, position % STATE, 0] ++ row[:D].normal_() ++ tail = row[D:].view(torch.int64) ++ tail[:3] = pos_value if pos_value is not None else position ++ if fixed: ++ tail[3] = position + 1 if tag is None else tag ++ ++ hist = historical_positions(computed) ++ rec["historical_positions"] = hist ++ if variant == "ring_realistic": ++ for p in hist: # this request committed them ++ write_row(p, 1234, None) ++ elif variant == "ring_uncommitted_row": ++ ring[VICTIM_RING].fill_(0.25) # residue everywhere... ++ for p in (96, 97): # ...but the previous occupant's ++ write_row(p, 1234, None) # rows 0-1 (its own positions) ++ elif variant.startswith("ring_every_row"): ++ bad = hist[1] if variant.endswith("_2") else hist[2] ++ for p in hist: ++ write_row(p, 1234, None) ++ write_row(bad - STATE, 1234, None) # same row, previous lap ++ rec["stale_position"] = bad ++ elif variant == "guard9_probe": ++ for p in hist: ++ write_row(p, MAX_MODEL_LEN, None) # one past the table, valid tag ++ ++ kw = dict(compress_ratio=CR, mrope_section=MROPE, rope_pos_offset=D) ++ if fixed: ++ kw["layer_id"] = LAYER ++ if variant == "ring_untouched_residue": ++ # the victim step commits rows 0-1 through the kernel itself ++ v_logical = torch.tensor([96, 97], dtype=torch.int64, device=device) ++ qpi.qsa_pre_indexer( ++ torch.randn((2, HQ * D), dtype=torch.bfloat16, device=device), ++ torch.randn((2, D), dtype=torch.bfloat16, device=device), ++ v_logical.unsqueeze(0).expand(3, -1).contiguous(), ++ cos_sin, ++ ones, ++ ones, ++ 1e-6, ++ torch.empty((2, HQ, D), dtype=torch.bfloat16, device=device), ++ ring, ++ torch.tensor( ++ [VICTIM_RING * STATE, VICTIM_RING * STATE + 1], ++ dtype=torch.int64, ++ device=device, ++ ), ++ ring_tab, ++ torch.tensor([0, 2], dtype=torch.int32, device=device), ++ v_logical, ++ comp, ++ torch.full((2,), -1, dtype=torch.int64, device=device), ++ torch.zeros(1, 2, dtype=torch.int32, device=device), ++ **kw, ++ ) ++ torch.cuda.synchronize() ++ gt.reset_for_test() ++ ++ q = torch.randn((sched, HQ * D), dtype=torch.bfloat16, device=device) ++ k = torch.randn((sched, D), dtype=torch.bfloat16, device=device) ++ pos = logical.to(torch.int64).unsqueeze(0).expand(3, -1).contiguous() ++ q_out = torch.empty((sched, HQ, D), dtype=torch.bfloat16, device=device) ++ qpi.qsa_pre_indexer( ++ q, ++ k, ++ pos, ++ cos_sin, ++ ones, ++ ones, ++ 1e-6, ++ q_out, ++ ring, ++ state_slots, ++ ring_tab, ++ md_for(ring_tab).query_start_loc, ++ logical, ++ comp, ++ comp_slots, ++ kwb, ++ **kw, ++ ) ++ torch.cuda.synchronize() ++ ++ g9 = gt.counter(gt.SITE_QSA_ROPE_POS, device).tolist() ++ rec["guard9_count"] = g9[0] ++ if fixed: ++ site = gt.counter(gt.SITE_QSA_RING_UNCOMMITTED, device).tolist() ++ rec["ring_site_count"] = site[0] ++ rec["ring_record"] = gt.decode_ring_record(site[1]) ++ rec["ring_found_tag"] = gt.decode_slot(site[2]) ++ first_boundary = computed - computed % CR + CR - 1 ++ tok = first_boundary - computed ++ slot = int(comp_slots[tok]) ++ rec["group_slot_written"] = bool( ++ (comp[slot // COMP_PAGE, slot % COMP_PAGE, 0] != 7.0).any() ++ ) ++ rec["status"] = "ran" ++ return rec ++ ++ ++def verdict(rec: dict) -> str: ++ v, tree, g9off = rec["variant"], rec.get("tree"), rec.get("guard9") == "disabled" ++ if rec["status"] != "ran": ++ return "FAULT" ++ if tree == "fixed": ++ if v in ("ring_realistic",): ++ ok = ( ++ rec["ring_site_count"] == 0 ++ and rec["guard9_count"] == 0 ++ and rec["group_slot_written"] ++ ) ++ return "clean" if ok else "UNEXPECTED" ++ if v == "guard9_probe": ++ want = 0 if g9off else 1 ++ ok = rec["ring_site_count"] == 0 and ( ++ (rec["guard9_count"] > 0) == bool(want) ++ ) ++ return ( ++ f"guard9 {'off' if g9off else 'on'} confirmed" if ok else "UNEXPECTED" ++ ) ++ r = rec.get("ring_record") or {} ++ ok = ( ++ rec["ring_site_count"] > 0 ++ and rec["guard9_count"] == 0 ++ and not rec["group_slot_written"] ++ and r.get("request_row") == 0 ++ and r.get("layer") == LAYER ++ and r.get("source_position") in rec["historical_positions"] ++ ) ++ return "request_error" if ok else "UNEXPECTED" ++ # unfixed tree ++ if v == "guard9_probe": ++ ok = (rec["guard9_count"] == 0) if g9off else (rec["guard9_count"] > 0) ++ return f"guard9 {'off' if g9off else 'on'} confirmed" if ok else "UNEXPECTED" ++ if v.startswith("ring_every_row") or v == "ring_realistic": ++ return "silent" if rec["guard9_count"] == 0 else "UNEXPECTED" ++ return "masked by guard 9" if rec["guard9_count"] > 0 else "silent" ++ ++ ++def main(argv) -> int: ++ p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) ++ p.add_argument("--self-test", action="store_true") ++ p.add_argument("--only", default="") ++ p.add_argument("--disable-guard9", action="store_true") ++ p.add_argument( ++ "--in-process", ++ action="store_true", ++ help="run the named variant(s) here (child mode)", ++ ) ++ args = p.parse_args(argv) ++ if args.self_test: ++ os.environ["CUDA_VISIBLE_DEVICES"] = "" ++ return self_test() ++ names = [n for n in args.only.split(",") if n] or VARIANTS ++ for n in names: ++ if n not in SHAPES: ++ print(f"unknown variant {n!r}; known: {', '.join(VARIANTS)}") ++ return 2 ++ if args.in_process: ++ for n in names: ++ rec = run_variant(n, args.disable_guard9) ++ rec["accept"] = verdict(rec) ++ print(json.dumps(rec, sort_keys=True), flush=True) ++ return 0 ++ rc = 0 ++ for n in names: ++ child = [ ++ sys.executable, ++ os.path.abspath(__file__), ++ "--in-process", ++ "--only", ++ n, ++ ] + (["--disable-guard9"] if args.disable_guard9 else []) ++ out = subprocess.run(child, capture_output=True, text=True) ++ lines = [ln for ln in out.stdout.splitlines() if ln.startswith("{")] ++ if out.returncode != 0 or not lines: ++ rec = { ++ "variant": n, ++ "status": "fault", ++ "accept": "FAULT", ++ "exit": out.returncode, ++ "stderr_tail": out.stderr.strip().splitlines()[-4:], ++ } ++ print(json.dumps(rec, sort_keys=True), flush=True) ++ else: ++ for ln in lines: ++ print(ln, flush=True) ++ if json.loads(ln).get("accept") == "UNEXPECTED": ++ rc = 1 ++ return rc ++ ++ ++if __name__ == "__main__": ++ sys.exit(main(sys.argv[1:])) +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_limits.py b/tests/v1/worker/gpu/test_guard_telemetry_limits.py +index 9fa7536..09d2865 100644 +--- a/tests/v1/worker/gpu/test_guard_telemetry_limits.py ++++ b/tests/v1/worker/gpu/test_guard_telemetry_limits.py +@@ -256,10 +256,16 @@ def test_zero_hits_is_not_bitwise_equivalence() -> None: + guard set is still a bounds-check set of the size this argument was + written against; if a site is added, this argument is re-read. + """ +- assert gt.NUM_SITES == 11 ++ assert gt.NUM_SITES == 12 + assert set(gt.SITE_NAMES) == set(range(gt.NUM_SITES)) + # The one non-bounds site, which is telemetry and not a masked fault. + assert "SKIPPED" in gt.SITE_NAMES[gt.SITE_QSA_RING_SKIP] ++ # Re-read for the ring-row tag fix (site 11): the one CONTENT check. It sees an in-range ++ # but wrong ring row only when the row's tag is wrong (never committed for ++ # that position); a row committed for the right position with the wrong ++ # speculative candidate's key still passes. So point 1 above stands for ++ # every other site and narrows, not vanishes, for this one. ++ assert "UNCOMMITTED" in gt.SITE_NAMES[gt.SITE_QSA_RING_UNCOMMITTED] + + + # --------------------------------------------------------------------------- +diff --git a/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +index 84f109d..6e07e6d 100644 +--- a/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py ++++ b/tests/v1/worker/gpu/test_qsa_ring_skip_telemetry.py +@@ -24,7 +24,7 @@ COMPRESS_RATIO = 4 + HEAD_DIM = 128 + N_HEADS = 4 + STATE_SIZE = 8 # ring capacity +-ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++ROW_WIDTH = 144 # head dim + the int64 tail (pos_t, pos_h, pos_w, tag) + MAX_MODEL_LEN = 4096 + NUM_BLOCKS = 8 + RING_CAPACITY = NUM_BLOCKS * STATE_SIZE +@@ -122,7 +122,7 @@ def test_site_is_declared_named_and_not_fatal(): + that is fatal would abort a boot on the older behaviour it merely + observes. Both are the wrong instrument.""" + site = guard_telemetry.SITE_QSA_RING_SKIP +- assert site == 10 and guard_telemetry.NUM_SITES == 11 ++ assert site == 10 and guard_telemetry.NUM_SITES == 12 # the ring-row tag fix added site 11 + assert site in guard_telemetry.SITE_NAMES + assert "SKIPPED" in guard_telemetry.SITE_NAMES[site] + assert site not in guard_telemetry._ALWAYS_FATAL_SITES +@@ -174,12 +174,13 @@ def test_negative_control_valid_slots_count_zero(): + torch.cuda.synchronize() + ctr = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, device) + assert ctr.tolist() == [0, 0, 0] +- # and the commit happened: every token's raw key and position tail ++ # and the commit happened: every token's raw key and position tail, and ++ # the tag word = logical position + 1 + for t in range(kw["k"].shape[0]): + row = kw["state_cache"][0, t, 0, :] + assert torch.equal(row[:HEAD_DIM], kw["k"][t]) +- tail = row[HEAD_DIM : HEAD_DIM + 12].contiguous().view(torch.int64) +- assert tail.tolist() == [t, t, t] ++ tail = row[HEAD_DIM : HEAD_DIM + 16].contiguous().view(torch.int64) ++ assert tail.tolist() == [t, t, t, t + 1] + + + @pytest.mark.gpu +diff --git a/tests/v1/worker/gpu/test_qsa_ring_tags_gpu.py b/tests/v1/worker/gpu/test_qsa_ring_tags_gpu.py +new file mode 100644 +index 0000000..8987e52 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_qsa_ring_tags_gpu.py +@@ -0,0 +1,311 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""v2.2 ring-row tag fix, kernel half (GPU): the pre-indexer's ring-row tag contract. ++ ++Writer: every committed ring row gets tag = logical_position + 1 in word 3 of ++its int64 tail. Reader: every HISTORICAL source row of a pooled group must ++carry exactly source_position + 1; otherwise the group's pooled load, its RoPE ++tail read and its compressed-slot store are masked and SITE_QSA_RING_UNCOMMITTED ++records one coherent first-offender record (layer, request row, source position +++ the tag found there). ++ ++Each positive case is paired with its clean twin in the same test (same batch, ++history committed correctly) that must count 0 and write the slot -- a counter ++that reads the same in both worlds measures nothing. ++ ++FAILING CONTROL: on the unmodified tree (an earlier commit) the module fails at ++collection (QSA_RING_TAIL_ELEMENTS / QSA_RING_TAG_WORD do not exist). Were the ++constants shimmed, the old launcher also has no ``layer_id`` (TypeError), no ++tag writer (the writer test reads tag 0) and no reader check (the every-row ++test counts nothing and writes the slot). ++""" ++ ++import pytest ++import torch ++ ++from vllm.models.qwen4_exp.common.qsa_cache import ( ++ QSA_RING_TAG_WORD, ++ QSA_RING_TAIL_ELEMENTS, ++) ++from vllm.v1.worker import guard_telemetry as gt ++ ++pytestmark = pytest.mark.gpu ++ ++CR = 4 ++D = 128 ++HQ = 4 ++STATE = 8 ++WIDTH = D + QSA_RING_TAIL_ELEMENTS ++NUM_BLOCKS = 8 ++COMP_PAGE = 800 ++MAX_POS = 4096 ++LAYER = 7 ++SENTINEL = 7.0 ++MROPE = (11, 11, 10) ++ ++ ++def _device(): ++ return torch.device("cuda", torch.cuda.current_device()) ++ ++ ++def _pos3(position: int) -> list[int]: ++ # distinct, in-range MRoPE coordinates per position ++ return [position, position // 7 + 3, position // 13 + 11] ++ ++ ++class _Batch: ++ """Requests (ring block, chunk_start, query_len) as one decode/prefill batch.""" ++ ++ def __init__(self, requests, device): ++ self.device = device ++ self.requests = requests ++ logical, slots, comp, qsl, work = [], [], [], [0], [] ++ comp_next = 0 ++ self.group_slot: dict[tuple[int, int], int] = {} ++ for r, (block, start, qlen) in enumerate(requests): ++ for i in range(qlen): ++ p = start + i ++ logical.append(p) ++ slots.append(block * STATE + p % STATE if i >= qlen - STATE else -1) ++ if (p + 1) % CR == 0: ++ comp.append(comp_next) ++ self.group_slot[(r, p)] = comp_next ++ comp_next += 1 ++ else: ++ comp.append(-1) ++ qsl.append(qsl[-1] + qlen) ++ groups = (start + qlen) // CR - start // CR ++ work += [[r, w] for w in range(max(groups, 1))] ++ T = len(logical) ++ self.T = T ++ self.logical = torch.tensor(logical, dtype=torch.int64, device=device) ++ self.positions = ( ++ torch.tensor([_pos3(p) for p in logical], dtype=torch.int64) ++ .t() ++ .contiguous() ++ .to(device) ++ ) ++ self.state_slots = torch.tensor(slots, dtype=torch.int64, device=device) ++ self.comp_slots = torch.tensor(comp, dtype=torch.int64, device=device) ++ self.qsl = torch.tensor(qsl, dtype=torch.int32, device=device) ++ self.state_table = torch.tensor( ++ [[b] for b, _s, _q in requests], dtype=torch.int32, device=device ++ ) ++ self.work = torch.tensor(work, dtype=torch.int32, device=device) ++ g = torch.Generator(device="cpu").manual_seed(1234) ++ self.q = torch.randn(T, HQ * D, generator=g).to(device, torch.bfloat16) ++ self.k = torch.randn(T, D, generator=g).to(device, torch.bfloat16) ++ self.cos_sin = torch.randn(MAX_POS, D // 2, generator=g).to( ++ device, torch.bfloat16 ++ ) ++ self.ring = torch.zeros( ++ NUM_BLOCKS, STATE, 1, WIDTH, dtype=torch.bfloat16, device=device ++ ) ++ self.comp = torch.full( ++ (NUM_BLOCKS, COMP_PAGE, 1, D), SENTINEL, dtype=torch.bfloat16, device=device ++ ) ++ self.q_out = torch.empty(T, HQ, D, dtype=torch.bfloat16, device=device) ++ ++ def commit(self, block: int, position: int, *, tag: int | None = None): ++ """What a correct earlier step of this request leaves in the ring.""" ++ row = self.ring[block, position % STATE, 0] ++ g = torch.Generator(device="cpu").manual_seed(position) ++ row[:D] = torch.randn(D, generator=g).to(self.device, torch.bfloat16) ++ tail = row[D:].view(torch.int64) ++ tail[:3] = torch.tensor(_pos3(position), device=self.device) ++ tail[QSA_RING_TAG_WORD] = position + 1 if tag is None else tag ++ ++ def tag(self, block: int, row: int) -> int: ++ return int(self.ring[block, row, 0, D:].view(torch.int64)[QSA_RING_TAG_WORD]) ++ ++ def run(self, *, rope=True): ++ from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import qsa_pre_indexer ++ ++ qsa_pre_indexer( ++ self.q, ++ self.k, ++ self.positions if rope else self.logical, ++ self.cos_sin, ++ torch.ones(D, dtype=torch.bfloat16, device=self.device), ++ torch.ones(D, dtype=torch.bfloat16, device=self.device), ++ 1e-6, ++ self.q_out, ++ self.ring, ++ self.state_slots, ++ self.state_table, ++ self.qsl, ++ self.logical, ++ self.comp, ++ self.comp_slots, ++ self.work, ++ compress_ratio=CR, ++ mrope_section=MROPE if rope else None, ++ rope_pos_offset=D if rope else None, ++ layer_id=LAYER, ++ ) ++ torch.cuda.synchronize() ++ ++ def slot_written(self, request: int, end_position: int) -> bool: ++ slot = self.group_slot[(request, end_position)] ++ row = self.comp[slot // COMP_PAGE, slot % COMP_PAGE, 0] ++ return not bool(torch.all(row == SENTINEL)) ++ ++ ++def _site(): ++ return gt.counter(gt.SITE_QSA_RING_UNCOMMITTED, _device()).tolist() ++ ++ ++def _rope_site(): ++ return gt.counter(gt.SITE_QSA_ROPE_POS, _device()).tolist() ++ ++ ++@pytest.fixture(autouse=True) ++def _fresh_counters(): ++ gt.prealloc(_device()) ++ gt.reset_for_test() ++ yield ++ gt.reset_for_test() ++ ++ ++# -------------------------------------------------------------------------- ++# writer ++# -------------------------------------------------------------------------- ++ ++ ++@pytest.mark.parametrize("rope", [True, False], ids=["mrope", "text"]) ++@pytest.mark.parametrize("qlen", [3, 6, 13], ids=["no-group", "short", "long"]) ++def test_writer_tags_every_committed_row(rope, qlen): ++ b = _Batch([(2, 40, qlen)], _device()) ++ b.run(rope=rope) ++ committed = list(range(40 + max(qlen - STATE, 0), 40 + qlen)) ++ for p in committed: ++ assert b.tag(2, p % STATE) == p + 1, f"row for position {p}" ++ if rope: ++ tail = b.ring[2, p % STATE, 0, D:].view(torch.int64) ++ assert tail[:3].tolist() == _pos3(p) ++ untouched = set(range(STATE)) - {p % STATE for p in committed} ++ for row in untouched: ++ assert b.tag(2, row) == 0 ++ assert _site() == [0, 0, 0] ++ ++ ++# -------------------------------------------------------------------------- ++# reader: the every-row control and its clean twin ++# -------------------------------------------------------------------------- ++ ++ ++def _three_historical_rows(device): ++ # chunk start 1003 = 3 mod 4: the first group is 1000..1003 and 1000, ++ # 1001, 1002 all come from the ring ++ b = _Batch([(3, 1003, 4)], device) ++ for p in (1000, 1001, 1002): ++ b.commit(3, p) ++ return b ++ ++ ++def test_clean_twin_counts_zero_and_writes_the_slot(): ++ b = _three_historical_rows(_device()) ++ b.run() ++ assert _site() == [0, 0, 0] ++ assert b.slot_written(0, 1003) ++ assert _rope_site()[0] == 0 ++ ++ ++@pytest.mark.parametrize("bad", [1001, 1002], ids=["row2", "row3"]) ++def test_only_the_second_or_third_row_stale_is_flagged(bad): ++ """First row valid; ONE later row holds an in-range stale position and key ++ from the previous lap of the ring (it was committed for bad - 8). Guard 9 ++ cannot see it (the position is in range) and neither could a first-row ++ check. The group must not be pooled and its slot must stay unwritten.""" ++ b = _three_historical_rows(_device()) ++ b.commit(3, bad - STATE) # overwrites the row with the previous lap ++ assert b.tag(3, bad % STATE) == bad - STATE + 1 ++ b.run() ++ n, rec, found = _site() ++ assert n == 1 ++ assert gt.decode_ring_record(rec) == { ++ "layer": LAYER, ++ "request_row": 0, ++ "source_position": bad, ++ } ++ assert gt.decode_slot(found) == bad - STATE + 1 ++ assert not b.slot_written(0, 1003) ++ assert _rope_site()[0] == 0 # guard 9 blind, as argued ++ # the Q path is independent of the ring and still ran ++ assert torch.isfinite(b.q_out.float()).all() ++ ++ ++@pytest.mark.parametrize( ++ "residue", ++ ["bf16-filler", "cleared", "negative"], ++) ++def test_uncommitted_history_is_flagged_not_faulted(residue): ++ """The production ring-residue shape: the rows the reader needs were never committed ++ by this request (allocator residue / a claimed block after the tag clear / ++ a negative int64). Before the ring-row tag fix the filler case decoded to position ++ ~0x3E803E803E803E80 and faulted in _norm_rope (guard 9 disabled) or was ++ clamped and counted (guard 9).""" ++ b = _Batch([(5, 52934 % 1024, 4)], _device()) # chunk start 2 mod 4 ++ start = 52934 % 1024 ++ first = start - start % CR ++ if residue == "bf16-filler": ++ b.ring[5].fill_(0.25) ++ elif residue == "negative": ++ b.ring[5].view(torch.int16).fill_(-1) ++ # "cleared": zeros, which is what the claim clear leaves ++ b.run() ++ n, rec, found = _site() ++ assert n == start - first # both historical rows ++ assert gt.decode_ring_record(rec)["source_position"] == first ++ expect_found = { ++ "bf16-filler": 0x3E803E803E803E80, ++ "cleared": 0, ++ "negative": -1, ++ }[residue] ++ assert gt.decode_slot(found) == expect_found ++ assert not b.slot_written(0, first + CR - 1) ++ assert _rope_site()[0] == 0 # the garbage tail was never read ++ ++ ++def test_missing_ring_block_is_flagged(): ++ """A historical row with no ring block at all (state table entry -1): ++ before the ring-row tag fix the load was masked and the group silently pooled zeros.""" ++ b = _Batch([(1, 1002, 4)], _device()) ++ b.state_table.fill_(-1) ++ b.run() ++ n, rec, found = _site() ++ assert n == 2 ++ assert gt.decode_ring_record(rec)["source_position"] == 1000 ++ assert gt.decode_slot(found) == 0 ++ assert not b.slot_written(0, 1003) ++ ++ ++def test_one_poisoned_request_leaves_its_neighbour_intact(): ++ """Containment inside the launch: only the poisoned request's group is ++ masked; the clean request's group is pooled and stored as normal.""" ++ b = _Batch([(2, 1002, 4), (4, 2002, 4)], _device()) ++ b.commit(2, 1000) ++ b.commit(2, 1001) ++ b.ring[4].fill_(0.25) # request 1's history was never committed ++ b.run() ++ n, rec, _found = _site() ++ assert n == 2 ++ assert gt.decode_ring_record(rec)["request_row"] == 1 ++ assert b.slot_written(0, 1003) ++ assert not b.slot_written(1, 2003) ++ ++ ++def test_first_offender_wins_and_the_count_takes_all(): ++ """Two poisoned requests: the count covers every flagged row, the record ++ is ONE of them whole (never request A's row with request B's tag).""" ++ b = _Batch([(2, 1002, 4), (4, 2003, 4)], _device()) ++ b.ring[2].fill_(0.25) ++ b.commit(4, 2000) ++ b.commit(4, 2001, tag=7) # stale: committed for another position ++ b.commit(4, 2002) ++ b.run() ++ n, rec, found = _site() ++ assert n == 3 # 2 rows of request 0 + 1 row of request 1 ++ record = gt.decode_ring_record(rec) ++ pair = (record["request_row"], record["source_position"], gt.decode_slot(found)) ++ assert pair in {(0, 1000, 0x3E803E803E803E80), (1, 2001, 7)} +diff --git a/tests/v1/worker/gpu/test_qsa_rope_position_guard.py b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +index 0628dd2..037a20f 100644 +--- a/tests/v1/worker/gpu/test_qsa_rope_position_guard.py ++++ b/tests/v1/worker/gpu/test_qsa_rope_position_guard.py +@@ -34,7 +34,7 @@ COMPRESS_RATIO = 4 + HEAD_DIM = 128 + N_HEADS = 4 + STATE_SIZE = 8 # ring capacity +-ROW_WIDTH = 140 # head dim + the int64 position triple in the tail ++ROW_WIDTH = 144 # head dim + the int64 tail (pos_t, pos_h, pos_w, tag) + MAX_MODEL_LEN = 4096 + NUM_BLOCKS = 8 + +@@ -242,6 +242,14 @@ def test_ring_filler_bytes_are_masked_and_counted(): + ) + # ...and the row the reader picks must be one carrying the filler. + assert 0 <= (first_position % STATE_SIZE) < STATE_SIZE ++ # Guard 9 is now the SECOND net. A row whose tag does not name its ++ # logical position is rejected before its tail is read, so the filler ++ # would never reach _norm_rope and this control would count 0. Give the ++ # historical rows VALID tags (bad coordinates with valid tags) so the ++ # garbage position tail is what reaches the guard. ++ for position in range(first_position, chunk_start): ++ tail = kw["state_cache"][0, position % STATE_SIZE, 0, HEAD_DIM:] ++ tail.view(torch.int64)[3] = position + 1 + + kw["k_work_metadata"][:, 0] = 0 + kw["k_work_metadata"][:, 1] = 0 +@@ -280,6 +288,9 @@ def test_ring_filler_bytes_are_masked_and_counted(): + assert worst >= MAX_MODEL_LEN or worst < 0, ( + f"guard fired but recorded an in-range position {worst}" + ) ++ # the tags were valid, so the ring-tag site must not have claimed this one ++ ring = guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_UNCOMMITTED, device) ++ assert int(ring[0].item()) == 0 + + + @pytest.mark.gpu +diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py +index 6b7d8f9..1391bdb 100644 +--- a/vllm/models/qwen4_exp/amd/indexer_qsa.py ++++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py +@@ -123,6 +123,15 @@ class QSAIndexer(nn.Module): + eps=float(getattr(config, "rms_norm_eps", 1e-6)), + ) + ++ # QSAKeyStateCache rows now carry ++ # a tag word the NVIDIA pre-indexer writes and validates; this path ++ # neither writes nor checks it, so it would run the widened layout ++ # with no ring validation at all. Refuse it rather than share it. ++ raise NotImplementedError( ++ "Qwen4Exp QSA on the AMD path is not supported on this fork: the " ++ "ring-row tag contract is implemented for the NVIDIA " ++ "pre-indexer only (vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py)." ++ ) + cache_config = vllm_config.cache_config + cache_prefix = f"{prefix}." if prefix else "" + self.raw_key_cache = QSAKeyStateCache( +diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py +index e806dc0..fbc517f 100644 +--- a/vllm/models/qwen4_exp/common/qsa_cache.py ++++ b/vllm/models/qwen4_exp/common/qsa_cache.py +@@ -44,6 +44,42 @@ from vllm.v1.kv_cache_interface import ( + MLAAttentionSpec, + ) + ++# Ring-row tag contract. Every raw-key ring row is the key (bf16), padded to a ++# whole int64, followed by a tail of four int64 words: ++# (pos_t, pos_h, pos_w, tag), tag = logical_position + 1. ++# The tail exists on every ring, with or without cached RoPE positions: the tag ++# is what lets the reader prove the row it is about to pool was committed by ++# this request for exactly that logical position. 0 is never a valid tag (so a ++# cleared row can never validate) and the manager clears the tag word of every ++# newly claimed ring block before its new owner's first step. ++QSA_RING_TAIL_WORDS = 4 ++QSA_RING_TAG_WORD = 3 ++QSA_RING_BF16_PER_INT64 = 4 ++QSA_RING_TAIL_ELEMENTS = QSA_RING_TAIL_WORDS * QSA_RING_BF16_PER_INT64 ++ ++ ++def qsa_ring_row_tags(logical_positions: torch.Tensor) -> torch.Tensor: ++ """The tag the writer stores for each committed row (int64).""" ++ return logical_positions.to(torch.int64) + 1 ++ ++ ++def qsa_ring_group_rows_valid( ++ found_tags: torch.Tensor, ++ source_positions: torch.Tensor, ++ historical: torch.Tensor, ++) -> torch.Tensor: ++ """Reference for the pre-indexer's per-group check (torch, any device). ++ ++ ``found_tags``/``source_positions``/``historical`` are ``[groups, ratio]``: ++ the tag word read from each source row's ring slot, the row's logical ++ position, and whether the row must come from the ring (it precedes the ++ chunk start). Returns a ``[groups]`` bool: True when EVERY historical row ++ carries ``source_position + 1``. In-chunk rows are never checked. ++ """ ++ expected = qsa_ring_row_tags(source_positions) ++ bad = historical & (found_tags.to(torch.int64) != expected) ++ return ~bad.any(dim=-1) ++ + + def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: + """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" +@@ -809,34 +845,51 @@ class _QSAStateCache(nn.Module, AttentionLayerBase): + + + class QSAKeyStateCache(_QSAStateCache): +- """Raw BF16 key, optionally followed by exact int64 MRoPE positions.""" ++ """Raw BF16 key followed by the int64 ring tail (pos_t, pos_h, pos_w, tag). + +- _BF16_PER_INT64 = 4 ++ The position words are consumed only when ``cache_rope_positions``; the tag ++ word is always written and always validated.""" ++ ++ _BF16_PER_INT64 = QSA_RING_BF16_PER_INT64 + _NUM_ROPE_AXES = 3 + + def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: + key_head_size = int(kwargs.pop("head_size")) + self.key_head_size = key_head_size + self.cache_rope_positions = bool(cache_rope_positions) ++ # Element offset of the int64 tail (the name predates the tag word). + self.rope_position_offset = ( + (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 + ) * self._BF16_PER_INT64 +- storage_head_size = key_head_size +- if self.cache_rope_positions: +- storage_head_size = self.rope_position_offset + ( +- self._NUM_ROPE_AXES * self._BF16_PER_INT64 +- ) ++ storage_head_size = self.rope_position_offset + QSA_RING_TAIL_ELEMENTS + super().__init__(head_size=storage_head_size, **kwargs) ++ self.rope_position_cache: torch.Tensor | None = None ++ self.ring_tag_cache: torch.Tensor | None = None + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + super().bind_kv_cache(kv_cache) + qsa_cache = self.kv_cache + self.key_cache = qsa_cache[..., : self.key_head_size] ++ tail_end = self.rope_position_offset + QSA_RING_TAIL_ELEMENTS ++ tail = qsa_cache[..., self.rope_position_offset : tail_end].view(torch.int64) ++ assert tail.shape[-1] == QSA_RING_TAIL_WORDS + if self.cache_rope_positions: +- position_tail = qsa_cache[..., self.rope_position_offset :] +- self.rope_position_cache = position_tail.view(torch.int64) ++ self.rope_position_cache = tail[..., : self._NUM_ROPE_AXES] + else: + self.rope_position_cache = None ++ self.ring_tag_cache = tail[..., QSA_RING_TAG_WORD : QSA_RING_TAG_WORD + 1] ++ ++ def clear_ring_row_tags(self, block_ids: torch.Tensor) -> None: ++ """Zero the tag word of every row of ``block_ids`` (ring-row tag claim clear). ++ ++ Only the tag word: a claimed ring block's keys and positions stay as ++ the previous occupant left them, and the zero tag is what makes every ++ such row fail validation until this request commits it.""" ++ if self.ring_tag_cache is None: ++ raise RuntimeError( ++ f"{self.prefix}: ring tag clear before the KV cache was bound" ++ ) ++ self.ring_tag_cache.index_fill_(0, block_ids, 0) + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # Hold the open group's committed keys plus every row a speculative +@@ -874,6 +927,9 @@ class QSACompressedKeyCache(_QSAStateCache): + + + __all__ = [ ++ "QSA_RING_TAG_WORD", ++ "QSA_RING_TAIL_ELEMENTS", ++ "QSA_RING_TAIL_WORDS", + "QSACompressedKeyCache", + "QSAForwardMetadata", + "QSAKeyStateCache", +@@ -882,4 +938,6 @@ __all__ = [ + "canonical_qsa_rope_positions", + "circular_qsa_slot_mapping", + "compressed_qsa_slot_mapping", ++ "qsa_ring_group_rows_valid", ++ "qsa_ring_row_tags", + ] +diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +index bbdc5e4..ea1865f 100644 +--- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +@@ -22,6 +22,7 @@ from ..common.qsa_cache import ( + QSAForwardMetadata, + QSAKeyStateCache, + canonical_qsa_rope_positions, ++ qsa_ring_row_tags, + ) + from .ops.qsa_pre_indexer import qsa_pre_indexer + +@@ -317,6 +318,7 @@ class QSAIndexer(nn.Module): + if raw_key_state_cache.rope_position_cache is not None + else None + ), ++ layer_id=self.layer_id, + ) + else: + # Unfused reference path +@@ -383,6 +385,18 @@ class QSAIndexer(nn.Module): + raw_metadata.slot_mapping, + position_rows, + ) ++ # The unfused writer commits the same tag word as the fused ++ # one, so both paths leave byte-identical rings. (This reference ++ # reader does not validate tags; the fused kernel does.) ++ ring_tag_cache = getattr(raw_key_state_cache, "ring_tag_cache", None) ++ if ring_tag_cache is not None: ++ qsa_store_cache_rows( ++ ring_tag_cache, ++ raw_metadata.slot_mapping, ++ qsa_ring_row_tags(raw_metadata.logical_positions[:num_tokens]).view( ++ -1, 1, 1 ++ ), ++ ) + + if self.skip_topk: + if out is None: +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +index 912bef3..cc5ab36 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py +@@ -4,6 +4,10 @@ + + import torch + ++from vllm.models.qwen4_exp.common.qsa_cache import ( ++ QSA_RING_TAG_WORD, ++ QSA_RING_TAIL_ELEMENTS, ++) + from vllm.triton_utils import tl, triton + from vllm.v1.worker import guard_telemetry + +@@ -12,6 +16,14 @@ from vllm.v1.worker import guard_telemetry + # from a @triton.jit body -- an earlier build shipped one and every guarded kernel died + # with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. + _GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) ++# Int64 word of the ring tail that carries `logical_position + 1`. ++_RING_TAG_WORD = tl.constexpr(QSA_RING_TAG_WORD) ++# Ring-row tag record packing and the encodable range of the found-tag slot. ++_REC_POS_BITS = tl.constexpr(guard_telemetry.RING_RECORD_POS_BITS) ++_REC_REQ_BITS = tl.constexpr(guard_telemetry.RING_RECORD_REQ_BITS) ++_REC_LAYER_BITS = tl.constexpr(guard_telemetry.RING_RECORD_LAYER_BITS) ++_GUARD_MAX_ENC = tl.constexpr(guard_telemetry.GUARD_MAX_ENCODABLE) ++_GUARD_MIN_ENC = tl.constexpr(guard_telemetry.GUARD_MIN_ENCODABLE) + + + +@@ -163,6 +175,7 @@ def _norm_rope( + "num_compressed_blocks", + "num_k_work", + "num_requests", ++ "layer_id", + ] + ) + def _qsa_pre_indexer_kernel( +@@ -202,6 +215,8 @@ def _qsa_pre_indexer_kernel( + num_requests, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] + ring_skip_ctr, # ring-skip telemetry, persistent int64[3], SITE_QSA_RING_SKIP ++ ring_uncommitted_ctr, # ring-row tag fix, persistent int64[3], SITE_QSA_RING_UNCOMMITTED ++ layer_id, # ring-row tag fix: named in the uncommitted-row record only + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, +@@ -345,6 +360,64 @@ def _qsa_pre_indexer_kernel( + + safe_state_block * state_cache_stride_block + + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token + ) ++ # Ring-row tag fix: EVERY historical source row must carry the tag the ++ # writer stores for exactly this logical position, before anything ++ # reads its key or its position tail. A missing ring block (state ++ # table entry out of range) reads tag 0 and fails the same way -- ++ # before the ring-row tag fix it pooled zeros silently. In-chunk rows come from ++ # this launch's own k and are never checked. ++ historical = valid & ~source_in_chunk ++ source_tags = tl.load( ++ ( ++ state_cache_ptr ++ + safe_state_block * state_cache_stride_block ++ + (source_positions % STATE_SIZE) * state_cache_stride_token ++ + D ++ ).to(tl.pointer_type(tl.int64)) ++ + _RING_TAG_WORD, ++ mask=historical & state_block_valid, ++ other=0, ++ ) ++ row_bad = historical & (source_tags != source_positions.to(tl.int64) + 1) ++ num_bad_rows = tl.sum(row_bad.to(tl.int32), axis=0) ++ group_ok = num_bad_rows == 0 ++ if num_bad_rows > 0: ++ # One coherent record per site, first offender wins: the CAS ++ # claims slot 1 and only its winner writes slot 2, so both ++ # fields always describe the same row. ++ bad_position = tl.min( ++ tl.where(row_bad, source_positions.to(tl.int64), 1 << 62), ++ axis=0, ++ ) ++ bad_tag = tl.sum( ++ tl.where( ++ row_bad & (source_positions.to(tl.int64) == bad_position), ++ source_tags, ++ 0, ++ ), ++ axis=0, ++ ) ++ record = ( ++ ( ++ ((layer_id + 1).to(tl.int64) & ((1 << _REC_LAYER_BITS) - 1)) ++ << (_REC_REQ_BITS + _REC_POS_BITS) ++ ) ++ | ( ++ (request.to(tl.int64) & ((1 << _REC_REQ_BITS) - 1)) ++ << _REC_POS_BITS ++ ) ++ | (bad_position & ((1 << _REC_POS_BITS) - 1)) ++ ) ++ tl.atomic_add(ring_uncommitted_ctr, num_bad_rows.to(tl.int64)) ++ zero64 = bad_position * 0 ++ prev = tl.atomic_cas( ++ ring_uncommitted_ctr + 1, zero64, record + _GUARD_BIAS ++ ) ++ if prev == 0: ++ clamped_tag = tl.minimum( ++ tl.maximum(bad_tag, _GUARD_MIN_ENC), _GUARD_MAX_ENC ++ ) ++ tl.atomic_xchg(ring_uncommitted_ctr + 2, clamped_tag + _GUARD_BIAS) + # Only the first completed group can cross the chunk boundary. Select + # historical rows from the ring without issuing two masked loads. + source_base = tl.where(source_in_chunk[:, None], current_base, cached_base) +@@ -353,9 +426,11 @@ def _qsa_pre_indexer_kernel( + source_valid = tl.where( + source_in_chunk, source_tokens_valid, state_block_valid + ) ++ # A failed group loads nothing (the mask, not a tl.where on the ++ # result, is what suppresses the memory operation). + source = tl.load( + source_base + dims[None, :], +- mask=valid & source_valid[:, None], ++ mask=valid & group_ok & source_valid[:, None], + other=0.0, + ).to(tl.float32) + # Match the unfused path's BF16 pooled tensor before RMSNorm. +@@ -403,7 +478,9 @@ def _qsa_pre_indexer_kernel( + + (first_position % STATE_SIZE) * state_cache_stride_token + + D + ).to(tl.pointer_type(tl.int64)) +- load_cached_position = ~first_in_chunk & state_block_valid ++ # The first row's tag was validated above with the rest; ++ # a failed group never reads its tail (pos stays 0, in range). ++ load_cached_position = ~first_in_chunk & state_block_valid & group_ok + cached_pos_t = tl.load(tail, mask=load_cached_position, other=0) + cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0) + cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0) +@@ -431,13 +508,16 @@ def _qsa_pre_indexer_kernel( + ) + compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64) + compressed_row = compressed_slot % COMP_PAGE_SIZE ++ # A group built from an uncommitted row must never become ++ # usable compressed state. The slot stays unwritten; the site is ++ # fatal because the selection kernels will still read it. + tl.store( + compressed_cache_ptr + + compressed_block * compressed_cache_stride_block + + compressed_row * compressed_cache_stride_token + + dims, + tl.reshape(y, (D,)), +- mask=valid, ++ mask=valid & group_ok, + ) + + if work_in_request == 0: +@@ -508,6 +588,20 @@ def _qsa_pre_indexer_kernel( + tl.store(tail, pos_t.to(tl.int64), mask=valid_slot) + tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot) + tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) ++ # Ring-row tag fix: the tag names the logical position this row now ++ # holds, taken from the authoritative per-token metadata (not ++ # re-derived from chunk arithmetic), stored after key + tail. ++ # Every committed row gets one, with or without RoPE positions. ++ token_logical = tl.load( ++ logical_positions_ptr + tl.maximum(token, 0), ++ mask=valid_slot, ++ other=-1, ++ ) ++ tl.store( ++ (state_row + D).to(tl.pointer_type(tl.int64)) + _RING_TAG_WORD, ++ token_logical.to(tl.int64) + 1, ++ mask=valid_slot, ++ ) + + + def qsa_pre_indexer( +@@ -531,8 +625,13 @@ def qsa_pre_indexer( + compress_ratio: int, + mrope_section: tuple[int, int, int] | None, + rope_pos_offset: int | None, ++ layer_id: int = -1, + ) -> None: +- """Normalize Q, compress K, then update the circular raw state.""" ++ """Normalize Q, compress K, then update the circular raw state. ++ ++ ``state_cache`` rows are ``[key D bf16 | 4 int64: pos_t, pos_h, pos_w, tag]`` ++ (``QSAKeyStateCache``, ring-row tag fix); the tail is present with or without ++ ``rope_pos_offset``. ``layer_id`` only labels the uncommitted-row record.""" + num_tokens = q.shape[0] + if num_tokens == 0: + return +@@ -551,6 +650,15 @@ def qsa_pre_indexer( + is_k_mrope = bool(mrope_section) + cache_has_rope_pos = rope_pos_offset is not None + assert rope_pos_offset is None or rope_pos_offset == head_dim ++ # The kernel addresses the tail at +head_dim and the tag at word 3 of ++ # it. A ring without the four-word tail (the older D or D+12 layouts) ++ # would put the tag store into the NEXT row, so refuse it outright. ++ assert state_cache.shape[-1] == head_dim + QSA_RING_TAIL_ELEMENTS, ( ++ f"QSA ring row width {state_cache.shape[-1]} != head_dim + " ++ f"{QSA_RING_TAIL_ELEMENTS} (key + pos_t/h/w + tag)" ++ ) ++ assert state_cache.dtype == torch.bfloat16 ++ assert state_cache.stride(1) % 4 == 0 and state_cache.stride(0) % 4 == 0 + if is_2d_positions: + assert positions.shape == (3, num_tokens) + assert is_k_mrope +@@ -604,6 +712,8 @@ def qsa_pre_indexer( + min(query_start_loc.numel() - 1, state_block_table.shape[0]), + guard_telemetry.counter(guard_telemetry.SITE_QSA_PRE_INDEXER, q.device), + guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, q.device), ++ guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_UNCOMMITTED, q.device), ++ int(layer_id), + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, +diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py +index 675ca1e..c0b94f0 100644 +--- a/vllm/v1/core/kv_cache_manager.py ++++ b/vllm/v1/core/kv_cache_manager.py +@@ -844,6 +844,21 @@ class KVCacheManager: + ids.extend(mgr.take_new_block_ids()) + return ids + ++ def take_new_ring_block_ids(self) -> list[int]: ++ """Drain ring blocks claimed since the last call (ring-row tag clear). ++ Independent of KV zeroing: every claim is returned.""" ++ ids: list[int] = [] ++ for mgr in self.coordinator.single_type_managers: ++ ids.extend(mgr.take_new_ring_block_ids()) ++ return ids ++ ++ def get_ring_block_ids(self, request_id: str) -> list[int]: ++ """The request's ring (CircularBufferSpec) block ids, all groups.""" ++ ids: list[int] = [] ++ for mgr in self.coordinator.single_type_managers: ++ ids.extend(mgr.get_ring_block_ids(request_id)) ++ return ids ++ + def get_zeroing_block_ids_in_range( + self, request_id: str, start_token: int, end_token: int + ) -> list[int]: +diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py +index 03bb2cd..afcaf42 100644 +--- a/vllm/v1/core/sched/output.py ++++ b/vllm/v1/core/sched/output.py +@@ -291,6 +291,13 @@ class SchedulerOutput: + # preventing stale NaN/data from corrupting attention or SSM computation. + new_block_ids_to_zero: list[int] | None = None + ++ # v2.2 ring-row tag fix: ring (CircularBufferSpec) blocks claimed during this ++ # scheduling step. Independent of new_block_ids_to_zero and of ++ # needs_kv_cache_zeroing: the worker clears only the rows' tag words, after ++ # the zeroer and before any copy or forward of this step, i.e. before the ++ # new owner's first read or write of the ring. ++ new_ring_block_ids_to_clear: list[int] | None = None ++ + # CoW copies to apply after zeroing new blocks and before forward. + kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 3b942cb..30704d0 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -366,6 +366,11 @@ class Scheduler(SchedulerInterface): + # Blocks that async KV loads will overwrite this step, skipped from + # zeroing since the zeroing could race the out-of-band write. + self._skip_zero_block_ids: set[int] = set() ++ # Ring blocks an async KV load is about to fill. Their tag clear ++ # is skipped for the same reason as zeroing (it could land after the ++ # out-of-band write and wipe the producer's committed tags). Tracked ++ # whatever needs_kv_cache_zeroing says, like the clear itself. ++ self._skip_ring_clear_block_ids: set[int] = set() + self.need_mamba_block_aligned_split = ( + self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" + ) +@@ -1276,6 +1281,12 @@ class Scheduler(SchedulerInterface): + num_computed_tokens, + ) + ) ++ # Same race for the ring's tag clear; the load ++ # brings the producer's rows and tags (or none, and the ++ # reader then fails closed on the missing history). ++ self._skip_ring_clear_block_ids.update( ++ self.kv_cache_manager.get_ring_block_ids(request.request_id) ++ ) + continue + + self.running.append(request) +@@ -1461,6 +1472,8 @@ class Scheduler(SchedulerInterface): + finished_req_ids=self.finished_req_ids, + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), + new_block_ids_to_zero=self._get_new_block_ids_to_zero(), ++ # Drained every step, never gated on KV zeroing. ++ new_ring_block_ids_to_clear=self._get_new_ring_block_ids_to_clear(), + has_sync_kv_loads=has_sync_kv_loads, + kv_cache_block_copies=pending_kv_cache_block_copies, + kv_connector_block_state=kv_connector_block_state, +@@ -1501,6 +1514,17 @@ class Scheduler(SchedulerInterface): + ) -> KVConnectorMetadata: + return connector.build_connector_meta(scheduler_output) + ++ def _get_new_ring_block_ids_to_clear(self) -> list[int] | None: ++ """Ring claims of this step whose tag words the worker clears. ++ Unconditional (not under needs_kv_cache_zeroing); only blocks an async ++ KV load is about to write are left out.""" ++ ids = self.kv_cache_manager.take_new_ring_block_ids() ++ if self._skip_ring_clear_block_ids: ++ skip = self._skip_ring_clear_block_ids ++ ids = [b for b in ids if b not in skip] ++ skip.clear() ++ return ids or None ++ + def _get_new_block_ids_to_zero(self) -> list[int] | None: + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 611e1ce..8242e25 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -120,6 +120,12 @@ class SingleTypeKVCacheManager(ABC): + and not isinstance(kv_cache_spec, CircularBufferSpec) + ) + self.new_block_ids: list[int] = [] ++ # v2.2 ring-row tag fix: ring blocks claimed since the last drain. The ring stays ++ # OUT of whole-block zeroing above; instead every claim is recorded ++ # UNCONDITIONALLY (never under needs_kv_cache_zeroing) so the worker ++ # clears the rows' tag words before the new owner's first step. Only ++ # CircularBufferManager appends here. ++ self.new_ring_block_ids: list[int] = [] + + # Mapping from request ID to blocks to track the blocks allocated + # for each request, so that we can free the blocks when the request +@@ -413,6 +419,17 @@ class SingleTypeKVCacheManager(ABC): + self.new_block_ids = [] + return ids + ++ def take_new_ring_block_ids(self) -> list[int]: ++ """Drain ring blocks claimed since the last call; empty for ++ every manager but CircularBufferManager.""" ++ ids = self.new_ring_block_ids ++ self.new_ring_block_ids = [] ++ return ids ++ ++ def get_ring_block_ids(self, request_id: str) -> list[int]: ++ """The request's ring block ids; empty for non-ring managers.""" ++ return [] ++ + def take_pending_cow_copies( + self, + ) -> list[tuple[KVCacheBlock, KVCacheBlock]]: +@@ -1203,8 +1220,37 @@ class CircularBufferManager(FullAttentionManager): + req_blocks.extend(new_blocks) + if self._record_new_block_ids: + self.new_block_ids.extend(block.block_id for block in new_blocks) ++ # Unconditional -- not gated on needs_kv_cache_zeroing. ++ self.new_ring_block_ids.extend(block.block_id for block in new_blocks) + return new_blocks + ++ def take_new_ring_block_ids(self) -> list[int]: ++ """Drain the ring blocks claimed since the last call, keeping only ++ blocks a request of this manager still holds, in claim order, once. ++ ++ A block claimed and then freed within the same scheduling pass (the ++ request was preempted after its allocation) may already belong to ++ another group by the time the worker acts; clearing its tag words ++ there would write into that group's page, so it is dropped here. ++ """ ++ claimed = self.new_ring_block_ids ++ self.new_ring_block_ids = [] ++ if not claimed: ++ return [] ++ held = { ++ block.block_id for blocks in self.req_to_blocks.values() for block in blocks ++ } ++ out: list[int] = [] ++ seen: set[int] = set() ++ for block_id in claimed: ++ if block_id in held and block_id not in seen: ++ seen.add(block_id) ++ out.append(block_id) ++ return out ++ ++ def get_ring_block_ids(self, request_id: str) -> list[int]: ++ return [block.block_id for block in self.req_to_blocks.get(request_id, ())] ++ + def get_num_blocks_to_allocate( + self, + request_id: str, +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index f37f741..92ae922 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -179,6 +179,7 @@ from vllm.v1.worker.gpu.ubatch_utils import ( + maybe_build_ubatch_runner, + ) + from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin ++from vllm.v1.worker.ring_tags import clear_ring_row_tags, ring_tag_layers + from vllm.v1.worker.utils import ( + KVBlockZeroer, + clear_layer_kv_caches, +@@ -1291,6 +1292,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): + assert self.kv_block_zeroer is not None + self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + ++ # v2.2 ring-row tag fix: clear the tag words of newly claimed ring blocks. Not ++ # under the zeroing gate; after the zeroer, before copies and forward. ++ if scheduler_output.new_ring_block_ids_to_clear: ++ clear_ring_row_tags( ++ ring_tag_layers(self.compilation_config.static_forward_context), ++ scheduler_output.new_ring_block_ids_to_clear, ++ self.device, ++ ) ++ + # Apply copy-on-write block copies for partial prefix-cache hits, after + # zeroing new blocks and before the forward pass reads them. + if scheduler_output.kv_cache_block_copies: +diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py +index 743c593..910c72c 100644 +--- a/vllm/v1/worker/gpu_model_runner.py ++++ b/vllm/v1/worker/gpu_model_runner.py +@@ -222,6 +222,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch + from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper + from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin + from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin ++from vllm.v1.worker.ring_tags import clear_ring_row_tags, ring_tag_layers + from vllm.v1.worker.ubatch_utils import ( + UBatchSlices, + check_ubatch_thresholds, +@@ -1243,6 +1244,14 @@ class GPUModelRunner( + # stale NaN/data from corrupting attention or SSM computation. + if scheduler_output.new_block_ids_to_zero: + self._zero_block_ids(scheduler_output.new_block_ids_to_zero) ++ # v2.2 ring-row tag fix: clear the tag words of newly claimed ring blocks. Not ++ # under the zeroing gate; after the zeroer, before copies and forward. ++ if scheduler_output.new_ring_block_ids_to_clear: ++ clear_ring_row_tags( ++ ring_tag_layers(self.compilation_config.static_forward_context), ++ scheduler_output.new_ring_block_ids_to_clear, ++ self.device, ++ ) + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index adca930..63c450f 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -146,7 +146,53 @@ SITE_QSA_ROPE_POS = 9 + # The skip itself is the earlier behaviour; the count is the instrument + # for the v2.2 ring-row reconciliation. Telemetry only: warn. + SITE_QSA_RING_SKIP = 10 +-NUM_SITES = 11 ++# The pre-indexer found a historical raw-key ring row whose tag is ++# not `logical_position + 1` -- a row this request never committed for that ++# position (claim residue, a skipped suffix commit, a stale mapping). The ++# group's pooled load, its RoPE-tail read and its compressed-slot store are all ++# MASKED, so the compressed slot is left unwritten: durable damage, fatal. ++# Unlike the other sites the three slots form ONE coherent record, claimed ++# first-offender-wins with atomic_cas: slot 0 = rows flagged (all offenders), ++# slot 1 = packed (layer + 1, batch request row, source logical position) of ++# the first offender, slot 2 = the tag that offender's row actually held. ++# Decode slot 1 with decode_ring_record(). ++SITE_QSA_RING_UNCOMMITTED = 11 ++NUM_SITES = 12 ++ ++# Bit layout of the SITE_QSA_RING_UNCOMMITTED record (before GUARD_BIAS). ++RING_RECORD_POS_BITS = 32 ++RING_RECORD_REQ_BITS = 16 ++RING_RECORD_LAYER_BITS = 14 ++ ++ ++def encode_ring_record(layer_id: int, request_row: int, position: int) -> int: ++ """Python twin of the kernel's packing (tests; the kernel inlines it).""" ++ layer = (int(layer_id) + 1) & ((1 << RING_RECORD_LAYER_BITS) - 1) ++ row = int(request_row) & ((1 << RING_RECORD_REQ_BITS) - 1) ++ pos = int(position) & ((1 << RING_RECORD_POS_BITS) - 1) ++ return ( ++ (layer << (RING_RECORD_REQ_BITS + RING_RECORD_POS_BITS)) ++ | (row << RING_RECORD_POS_BITS) ++ | pos ++ ) ++ ++ ++def decode_ring_record(biased: int) -> dict[str, int] | None: ++ """Unpack slot 1 of SITE_QSA_RING_UNCOMMITTED. ``layer`` is -1 when the ++ launcher was not given one. None = nothing recorded (or not yet visible).""" ++ value = decode_slot(biased) ++ if value is None: ++ return None ++ pos_mask = (1 << RING_RECORD_POS_BITS) - 1 ++ req_mask = (1 << RING_RECORD_REQ_BITS) - 1 ++ layer_mask = (1 << RING_RECORD_LAYER_BITS) - 1 ++ return { ++ "layer": ((value >> (RING_RECORD_REQ_BITS + RING_RECORD_POS_BITS)) & layer_mask) ++ - 1, ++ "request_row": (value >> RING_RECORD_POS_BITS) & req_mask, ++ "source_position": value & pos_mask, ++ } ++ + + SITE_NAMES = { + SITE_MAMBA_COPY: "_copy_mamba_state_block (block-table row/column)", +@@ -165,6 +211,10 @@ SITE_NAMES = { + SITE_QSA_RING_SKIP: "_qsa_pre_indexer_kernel suffix commit SKIPPED " + "(state slot < 0 or >= ring capacity; slot 1 = token index, " + "slot 2 = state slot value)", ++ SITE_QSA_RING_UNCOMMITTED: "_qsa_pre_indexer_kernel historical ring row " ++ "UNCOMMITTED (tag != logical_position + 1; group load, RoPE tail and " ++ "compressed-slot store masked; slot 1 = packed first-offender record, " ++ "slot 2 = the tag its row held)", + } + + +@@ -373,6 +423,14 @@ _ALWAYS_FATAL_SITES: set[int] = set() + # whatever VLLM_E2_GUARD_MODE says. (bounds-guard instrumentation, rope site) + _ALWAYS_FATAL_SITES.add(SITE_QSA_ROPE_POS) + ++# An uncommitted ring row leaves the group's compressed slot ++# UNWRITTEN (the store is masked, there is no correct value to store), and the ++# selection kernels will read that slot on every later step. That is durable ++# damage by the same argument as the rope site, so it stops the rank under ++# warn too. The failure boundary is the rank, not the request: per-request ++# containment (discard + quarantine + TP/PP coordination) is not implemented. ++_ALWAYS_FATAL_SITES.add(SITE_QSA_RING_UNCOMMITTED) ++ + + def _normalize(device: torch.device) -> torch.device: + """torch.device("cuda") and torch.device("cuda:0") are different dict keys. +@@ -479,6 +537,19 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None: + _seen[(device, site)] = n_masked + row = decode_slot(biased_row) + val = decode_slot(biased_val) ++ if site == SITE_QSA_RING_UNCOMMITTED: ++ # One coherent first-offender record, not independent maxima. ++ record = decode_ring_record(biased_row) ++ msg = ( ++ f"ring-row tag guard masked {n_masked} uncommitted historical " ++ f"ring row(s) in {SITE_NAMES[site]} on {device}; first " ++ f"offender {record}, its row held tag {val} (expected " ++ f"source_position + 1). The group was not pooled and its " ++ f"compressed slot was NOT written." ++ ) ++ logger.warning("%s", msg) ++ _arm(msg) ++ continue + msg = ( + f"bounds guard masked {n_masked} out-of-range access(es) " + f"in {SITE_NAMES.get(site, site)} on {device} " +diff --git a/vllm/v1/worker/ring_tags.py b/vllm/v1/worker/ring_tags.py +new file mode 100644 +index 0000000..ef0a8b7 +--- /dev/null ++++ b/vllm/v1/worker/ring_tags.py +@@ -0,0 +1,63 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""v2.2 ring-row tag fix: worker side of the ring-row tag clear. ++ ++The scheduler reports every newly claimed ring (``CircularBufferSpec``) block ++in ``SchedulerOutput.new_ring_block_ids_to_clear``, whether or not KV zeroing ++is on. Before the step's forward, the runner hands them to every layer that ++keeps a tagged ring (``clear_ring_row_tags``; today ``QSAKeyStateCache``), ++which zeroes only the tag word of each row. A zero tag never validates, so the ++new owner cannot pool a row it did not commit itself, whatever the previous ++occupant (or another group that held the block) left there. ++ ++ORDER. The runners call this right after the KV zeroer and before the CoW ++copies and the forward, on the same stream, inside the same execute_model as ++the claim's first step. So: previous owner's work done (deferred free) -> ++claim -> tag clear -> new owner's first read/write. ++""" ++ ++from collections.abc import Iterable, Mapping ++ ++import torch ++ ++ ++def ring_tag_layers(static_forward_context: Mapping[str, object]) -> list: ++ """Layers of this worker that keep a tagged ring.""" ++ return [ ++ layer ++ for layer in static_forward_context.values() ++ if callable(getattr(layer, "clear_ring_row_tags", None)) ++ ] ++ ++ ++def clear_ring_row_tags( ++ layers: Iterable, ++ block_ids: list[int], ++ device: torch.device, ++) -> int: ++ """Zero the tag word of every row of ``block_ids`` in every layer. ++ ++ Returns the number of layers cleared. Raises (fail loud) on a block id ++ outside a layer's pool: an index_fill_ with a bad index is a device-side ++ assert, far from its cause.""" ++ if not block_ids: ++ return 0 ++ layers = list(layers) ++ if not layers: ++ return 0 ++ lo, hi = min(block_ids), max(block_ids) ++ index = torch.tensor(block_ids, dtype=torch.long, device=device) ++ for layer in layers: ++ tags = getattr(layer, "ring_tag_cache", None) ++ if tags is None: ++ raise RuntimeError( ++ f"ring tag clear: layer {getattr(layer, 'prefix', layer)!r} has " ++ "no bound ring_tag_cache" ++ ) ++ if lo < 0 or hi >= tags.shape[0]: ++ raise ValueError( ++ f"ring tag clear: block ids [{lo}, {hi}] outside the pool of " ++ f"{tags.shape[0]} blocks ({getattr(layer, 'prefix', layer)!r})" ++ ) ++ layer.clear_ring_row_tags(index) ++ return len(layers) +-- +2.43.0 + diff --git a/patches/0050-sched-drop-the-one-victim-preemption-guard-that-upst.patch b/patches/0050-sched-drop-the-one-victim-preemption-guard-that-upst.patch new file mode 100644 index 0000000..a814934 --- /dev/null +++ b/patches/0050-sched-drop-the-one-victim-preemption-guard-that-upst.patch @@ -0,0 +1,130 @@ +From 4fcb598510da1f2f4fad1724fdb88c34d1301e07 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 10:00:36 -0400 +Subject: [PATCH 050/103] sched: drop the one-victim preemption guard that + upstream #49675 made dead + +The cascade test now expects zero needless victims. +--- + tests/v1/core/test_deferred_block_free.py | 43 +++++++++++++++-------- + vllm/v1/core/sched/scheduler.py | 18 ++++------ + 2 files changed, 36 insertions(+), 25 deletions(-) + +diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py +index e94161b..e5e91a3 100644 +--- a/tests/v1/core/test_deferred_block_free.py ++++ b/tests/v1/core/test_deferred_block_free.py +@@ -646,13 +646,18 @@ def test_ring_preempt_withholds_the_ring_block_until_the_fence(): + + + def test_deferred_victim_stops_the_preemption_cascade(): +- """Under deferred free a preempted victim returns no blocks +- until its in-flight step is processed, so the allocation retry must not +- keep preempting. Two 48-token requests fill a 6-block pool; the first +- request's speculative decode (token 49) needs a 4th block. Expected: exactly ONE +- victim (the last running request), the requester stays RUNNING and +- unscheduled, and after the in-flight step is processed it gets its block. +- Without the fix this preempts both requests in one schedule() call. ++ """Restated for upstream #49675 (v0.30.0): under deferred free ++ a preempted victim returns no blocks until its in-flight step is ++ processed, so preempting it cannot satisfy the allocation. Two 48-token ++ requests fill a 6-block pool; the first request's speculative decode ++ (token 49) needs a 4th block. Contract: ZERO victims while the fence is ++ pending -- both requests stay RUNNING, nothing is deferred, the requester ++ is unscheduled -- and once the in-flight step is processed ordinary ++ preemption resumes with an IMMEDIATE free (the victim's blocks never enter ++ the deferred queue) and the requester gets its block. ++ Without the fix this preempted both requests in one schedule() call; our v2 ++ post-hoc guard stopped at exactly one deferred victim; #49675's ++ pre-check (`_request_blocks_can_be_freed(victim)`) makes it zero. + """ + with patch.object( + VllmConfig, +@@ -680,24 +685,34 @@ def test_deferred_victim_stops_the_preemption_cascade(): + assert pool.get_num_free_blocks() == 0 + + # step 2 (step 1 in flight): the first request's decode crosses into a +- # 4th block -> allocation fails -> the loop preempts the LAST running one. ++ # 4th block -> allocation fails -> the only candidate victim (the LAST ++ # running one) is fenced behind step 1, so nobody is preempted. + first, second = scheduler.running[0], scheduler.running[1] + out1 = scheduler.schedule() +- assert second.status == RequestStatus.PREEMPTED ++ assert not out1.preempted_req_ids, "a fenced victim was preempted" ++ assert second.status == RequestStatus.RUNNING, "a fenced victim was preempted" + assert first.status == RequestStatus.RUNNING, "the cascade reached the requester" +- assert first in scheduler.running +- assert len(scheduler.deferred_frees) == 1 ++ assert scheduler.running == [first, second] ++ assert first.num_preemptions == 0 and second.num_preemptions == 0 ++ assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == 0 + assert first.request_id not in out1.num_scheduled_tokens ++ assert out1.total_num_scheduled_tokens == 0 + +- # step 1's output processed: the victim's blocks drain, the requester +- # is scheduled on the next call. ++ # step 1's output processed: the fence has passed, so the retry preempts ++ # the last running request with an immediate free and the requester is ++ # scheduled in the same call. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert not scheduler.deferred_frees +- assert pool.get_num_free_blocks() == 3 ++ assert pool.get_num_free_blocks() == 0 + out2 = scheduler.schedule() ++ assert out2.preempted_req_ids == {second.request_id} ++ assert second.status == RequestStatus.PREEMPTED ++ assert second.num_preemptions == 1 ++ assert not scheduler.deferred_frees, "a free past the fence was deferred" + assert first.request_id in out2.num_scheduled_tokens + assert first.status == RequestStatus.RUNNING ++ assert first.num_preemptions == 0 + + + def _make_pool_with_cached_blocks(num_blocks: int, block_size: int = 16): +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 30704d0..1402e8e 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -787,6 +787,13 @@ class Scheduler(SchedulerInterface): + preempted_req = self.running[-1] + + # A deferred free will not help with immediate allocation. ++ # (Upstream #49675. It subsumes our earlier post-hoc guard, ++ # which stopped after ONE deferred victim: a victim that ++ # passes this check is freed immediately by ++ # _preempt_request, so deferred_frees cannot grow below. ++ # The requester stays running, unscheduled, and retries ++ # after the fence drains; with every running request in ++ # flight (PP / async) that is zero victims, not the batch.) + if not self._request_blocks_can_be_freed(preempted_req): + break + +@@ -818,7 +825,6 @@ class Scheduler(SchedulerInterface): + else: + preempted_req = self.running.pop() + +- num_deferred_before = len(self.deferred_frees) + self._preempt_request( + preempted_req, + scheduled_timestamp, +@@ -828,16 +834,6 @@ class Scheduler(SchedulerInterface): + if preempted_req == request: + # No more request to preempt. Cannot schedule this request. + break +- if len(self.deferred_frees) > num_deferred_before: +- # The ring-skip fix: the victim's blocks are fenced +- # behind its in-flight step and return to the pool only +- # when that step's output is processed, so retrying the +- # allocation now cannot succeed on their account and would +- # preempt the next request for nothing -- with every running +- # request in flight (PP / async) the whole batch. Stop at one +- # deferred victim; this request stays running, unscheduled, +- # and retries next step after the drain. +- break + + if new_blocks is None: + # Cannot schedule this request. +-- +2.43.0 + diff --git a/patches/0051-ple-tests-debug-gated-pull-flag-check-diagnostics-re.patch b/patches/0051-ple-tests-debug-gated-pull-flag-check-diagnostics-re.patch new file mode 100644 index 0000000..37988bc --- /dev/null +++ b/patches/0051-ple-tests-debug-gated-pull-flag-check-diagnostics-re.patch @@ -0,0 +1,1015 @@ +From dcd053cd9472cb967328501e99c2b09078b365d0 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 10:01:40 -0400 +Subject: [PATCH 051/103] ple, tests: debug-gated pull-flag check; diagnostics + report that they ran + +The pull-transport flag check is wired in (rank-gated, off by default, warns when enabled). Both +diagnostic checks log that they ran. Test repairs for helpers v0.30.0 renamed. +--- + .../kernels/mamba/test_precopy_mamba_align.py | 158 ++++++++++++ + .../v1/ple_offload/test_flag_check_wiring.py | 240 ++++++++++++++++++ + .../ple_offload/test_stale_flag_assertion.py | 135 ++++++++++ + .../worker/gpu/test_mamba_slot_ownership.py | 40 +++ + .../gpu/test_pp_deferred_mamba_block_table.py | 11 +- + .../layers/ple_offload_layer.py | 85 +++++++ + vllm/v1/ple_offload/connector.py | 151 +++++++++++ + vllm/v1/worker/mamba_utils.py | 20 ++ + 8 files changed, 836 insertions(+), 4 deletions(-) + create mode 100644 tests/v1/ple_offload/test_flag_check_wiring.py + create mode 100644 tests/v1/ple_offload/test_stale_flag_assertion.py + +diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py +index 1229511..460c4b9 100644 +--- a/tests/kernels/mamba/test_precopy_mamba_align.py ++++ b/tests/kernels/mamba/test_precopy_mamba_align.py +@@ -21,7 +21,11 @@ The kernel must also no-op when ``src_col < 0`` (fresh request) or + + from __future__ import annotations + ++import ast ++import inspect + import os ++import sys ++import textwrap + from collections.abc import Callable + from types import SimpleNamespace + from typing import Any +@@ -168,6 +172,7 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first + @_parametrize("token_bias", [0, 1, 2]) + @_parametrize("has_idx_mapping", [True, False]) + @_parametrize("temporal_tiles", [1, _TEMPORAL_TILES]) ++@_gpu_host_tier + @_cuda_required + def test_precopy_matches_v1_copy_specs( + num_reqs, +@@ -213,6 +218,11 @@ def test_precopy_matches_v1_copy_specs( + ) + bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) ++ # The bounds-guard instrumentation added block_table_num_rows + guard_ctr after the row ++ # stride. This launch predated them and passed 15 positionals to a ++ # 17-positional kernel, so it could never bind. ++ # test_every_launch_in_this_file_binds_the_kernel_signature pins it GPU-free. ++ guard_ctr = torch.zeros(3, dtype=torch.int64, device=device) + grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) + precopy_mamba_align_fused_kernel[grid]( + dst_col, +@@ -220,6 +230,8 @@ def test_precopy_matches_v1_copy_specs( + bias, + bt_ptrs, + bt.stride(0), ++ bt.size(0), ++ guard_ctr, + base, + blk_stride, + elem, +@@ -240,6 +252,152 @@ def test_precopy_matches_v1_copy_specs( + for layer in range(NUM_LAYERS): + torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) + torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) ++ # Every row is < bt.size(0) here, so a non-zero count is a guard false ++ # positive (or an out-of-range row this setup does not create). ++ assert int(guard_ctr[0]) == 0, "in-range rows must not trip the bounds guard" ++ ++ ++# The Triton kernels launched in this file, by the name the launch uses. ++_LAUNCHED_KERNELS = { ++ "precopy_mamba_align_fused_kernel": precopy_mamba_align_fused_kernel, ++ "postprocess_mamba_fused_kernel": postprocess_mamba_fused_kernel, ++} ++ ++ ++def _launch_sites_in_this_file() -> list[tuple[int, str, int, list[str], list]]: ++ """``(lineno, kernel, n_positional, keyword_names, positional_exprs)`` for ++ every ``[grid](...)`` call in this module's source.""" ++ tree = ast.parse(inspect.getsource(sys.modules[__name__])) ++ sites = [] ++ for node in ast.walk(tree): ++ if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Subscript)): ++ continue ++ target = node.func.value ++ if not (isinstance(target, ast.Name) and target.id in _LAUNCHED_KERNELS): ++ continue ++ assert not any(isinstance(a, ast.Starred) for a in node.args), ( ++ f"line {node.lineno}: a *args launch cannot be checked statically" ++ ) ++ sites.append( ++ ( ++ node.lineno, ++ target.id, ++ len(node.args), ++ [kw.arg for kw in node.keywords], ++ node.args, ++ ) ++ ) ++ return sites ++ ++ ++def launch_arity_findings(sites) -> list[str]: ++ """Every way a launch site fails to bind its kernel's signature. ++ ++ The rule this file follows: runtime parameters are passed POSITIONALLY and ++ all of them are passed; ``tl.constexpr`` parameters (read from the kernel's ++ own source annotations, since ``block_size`` is a lower-case constexpr) are ++ passed by keyword. A launch one runtime parameter ++ short does not fail until Triton binds it on a GPU, which is how ++ ``test_precopy_matches_v1_copy_specs`` stayed broken from that change onwards ++ without anyone seeing it (it was never run). Checked here with no device. ++ """ ++ findings = [] ++ for lineno, name, n_pos, kw_names, pos_exprs in sites: ++ kernel = _LAUNCHED_KERNELS[name] ++ fn = getattr(kernel, "fn", kernel) ++ sig = inspect.signature(fn) ++ [fdef] = [ ++ n ++ for n in ast.parse(textwrap.dedent(inspect.getsource(fn))).body ++ if isinstance(n, ast.FunctionDef) ++ ] ++ constexpr = { ++ a.arg ++ for a in fdef.args.args ++ if a.annotation is not None ++ and ast.unparse(a.annotation).endswith("constexpr") ++ } ++ assert constexpr, f"{name}: no tl.constexpr parameters found in source" ++ runtime = [p for p in sig.parameters if p not in constexpr] ++ if n_pos != len(runtime): ++ findings.append( ++ f"line {lineno}: {name} launched with {n_pos} positional " ++ f"argument(s), kernel declares {len(runtime)} runtime " ++ f"parameter(s) {runtime}" ++ ) ++ continue ++ try: ++ sig.bind(*([None] * n_pos), **dict.fromkeys(kw_names)) ++ except TypeError as error: ++ findings.append(f"line {lineno}: {name}: {error}") ++ continue ++ # The bounds-guard counter slot must be handed a guard counter, not ++ # whatever tensor happened to land in that position. ++ if "guard_ctr" in runtime: ++ expr = pos_exprs[runtime.index("guard_ctr")] ++ text = ast.unparse(expr) ++ if "guard" not in text: ++ findings.append( ++ f"line {lineno}: {name}: guard_ctr position receives " ++ f"{text!r}, not a guard counter" ++ ) ++ return findings ++ ++ ++def test_every_launch_in_this_file_binds_the_kernel_signature(): ++ """GPU-free: each kernel launch in this module binds its real signature. ++ ++ An earlier ``test_precopy_matches_v1_copy_specs`` passed 15 positionals ++ to a kernel that has taken 17 since the bounds-guard instrumentation added ++ ``block_table_num_rows`` and ``guard_ctr``, and nothing noticed because the ++ GPU tier never ran it. FAILING CONTROL: on the unrepaired file this reports ++ the precopy launch as 15 vs 17 (see ``test_arity_check_discriminates``). ++ """ ++ sites = _launch_sites_in_this_file() ++ kernels_seen = {name for _l, name, *_rest in sites} ++ assert kernels_seen == set(_LAUNCHED_KERNELS), ( ++ f"expected launches of {sorted(_LAUNCHED_KERNELS)}, found " ++ f"{sorted(kernels_seen)}: the scan is not seeing this file's launches" ++ ) ++ assert launch_arity_findings(sites) == [] ++ ++ ++def test_arity_check_discriminates(): ++ """The check fires on the pre-repair launch shape (the failing control). ++ ++ Rebuilds the exact pre-repair precopy launch -- the row stride followed ++ directly by the state metadata, 15 positionals -- and one with the right ++ arity but a non-guard tensor in the counter slot. Both must be reported, ++ or the green result above is a test that cannot fail. ++ """ ++ pre_repair = ast.parse( ++ "precopy_mamba_align_fused_kernel[grid](dst_col, src_col, bias, " ++ "bt_ptrs, bt.stride(0), base, blk_stride, elem, inner, width, group, " ++ "drc, drs, idx_mapping, num_reqs, COPY_BLOCK_SIZE=1024, " ++ "CONV_STATE_DIM_FIRST=False, HAS_IDX_MAPPING=True, TEMPORAL_TILES=1)", ++ mode="eval", ++ ).body ++ wrong_slot = ast.parse( ++ "precopy_mamba_align_fused_kernel[grid](dst_col, src_col, bias, " ++ "bt_ptrs, bt.stride(0), bt.size(0), base, base, blk_stride, elem, " ++ "inner, width, group, drc, drs, idx_mapping, num_reqs, " ++ "COPY_BLOCK_SIZE=1024, CONV_STATE_DIM_FIRST=False)", ++ mode="eval", ++ ).body ++ ++ def _site(call): ++ return ( ++ call.lineno, ++ call.func.value.id, ++ len(call.args), ++ [kw.arg for kw in call.keywords], ++ call.args, ++ ) ++ ++ [arity] = launch_arity_findings([_site(pre_repair)]) ++ assert "15 positional" in arity and "17 runtime" in arity, arity ++ [slot] = launch_arity_findings([_site(wrong_slot)]) ++ assert "guard_ctr position receives 'base'" in slot, slot + + + def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch): +diff --git a/tests/v1/ple_offload/test_flag_check_wiring.py b/tests/v1/ple_offload/test_flag_check_wiring.py +new file mode 100644 +index 0000000..2ee5a0b +--- /dev/null ++++ b/tests/v1/ple_offload/test_flag_check_wiring.py +@@ -0,0 +1,240 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The PLE flag pre-launch check is WIRED on this tree, not just registered. ++ ++Why this file exists (PLE diagnostics port to 0.30). On the pre-fix kernel branch, ++an earlier commit had ``VLLM_E2_PLE_FLAG_CHECK`` registered in ``envs.py`` -- parsed, ++normalized, kept out of the compile factors -- but the connector code that ++reads it had been dropped by the rebase (the PLE-diagnostics changes were replayed onto ++the 0.30 base before the PLE offload package was restored there). A diagnostics run ++with the knob on would have reported "zero findings" from an instrument that ++never ran. ``test_stale_flag_assertion.py`` pins the pure helpers; it cannot ++catch this, because the helpers can exist while nothing calls them. ++ ++These tests drive the REAL ``PleOffloadConnector._prepare_forward_impl`` with ++the transport's I/O (``_launch``, the stream synchronize) replaced, and the ++REAL ``stale_offload_flags`` reading real ``CpuGpuSemaphore`` flags. GPU-free. ++ ++FAILING CONTROL: on that earlier commit every test in this file fails -- ++``_PLE_FLAG_CHECK_MODE`` / ``_check_flags_clear_before_launch`` do not exist and ++``VLLM_E2_PLE_FLAG_CHECK`` has no reader outside ``envs.py``. ++""" ++ ++import importlib ++import re ++from pathlib import Path ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++import vllm ++import vllm.envs as envs ++from vllm.model_executor.layers.ple_offload_layer import ( ++ CpuGpuSemaphore, ++ PleOffloadLayer, ++) ++from vllm.v1.ple_offload import connector as connector_mod ++from vllm.v1.ple_offload.connector import PleOffloadConnector ++ ++HIDDEN = 4 ++ ++ ++class _ProbeLayer(PleOffloadLayer): ++ def forward_impl(self, hidden_states, input_ids, *args, **kwargs): ++ return hidden_states ++ ++ ++def _layer(flag_value: int) -> _ProbeLayer: ++ layer = _ProbeLayer() ++ layer.setup_cross_process_offload( ++ gpu_output_buffer=torch.zeros(8, HIDDEN, dtype=torch.bfloat16), ++ semaphore=CpuGpuSemaphore.from_ipc_tensor( ++ torch.tensor([flag_value], dtype=torch.int32) ++ ), ++ ) ++ return layer ++ ++ ++class _Harness: ++ """A real connector object with only its I/O edges replaced.""" ++ ++ def __init__(self, monkeypatch, *, mode: str, tp_rank: int, flag: int): ++ monkeypatch.setattr(connector_mod, "_PLE_FLAG_CHECK_MODE", mode) ++ self.syncs = 0 ++ self.launches: list[int] = [] ++ self.warnings: list[str] = [] ++ self.infos: list[str] = [] ++ ++ def _sync(): ++ self.syncs += 1 ++ ++ monkeypatch.setattr( ++ torch.cuda, ++ "current_stream", ++ lambda device=None: SimpleNamespace(synchronize=_sync), ++ ) ++ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) ++ monkeypatch.setattr( ++ connector_mod.logger, ++ "warning", ++ lambda fmt, *args: self.warnings.append(fmt % args), ++ ) ++ monkeypatch.setattr( ++ connector_mod.logger, ++ "info", ++ lambda fmt, *args: self.infos.append(fmt % args), ++ ) ++ conn = object.__new__(PleOffloadConnector) ++ conn._input_ids_buf = torch.zeros(64, dtype=torch.int32) ++ conn._query_start_loc_buf = torch.zeros(9, dtype=torch.int32) ++ conn._after_dummy = False ++ conn._rowdump_n = None ++ conn.device = torch.device("cpu") ++ conn.dp_rank, conn.tp_rank, conn.pp_rank = 0, tp_rank, 0 ++ conn._layers = {"ple.0": _layer(CpuGpuSemaphore.RESET_VALUE)} ++ conn._layers["ple.1"] = _layer(flag) ++ ++ def _launch(num_reqs, num_tokens, **kwargs): ++ self.launches.append(num_tokens) ++ ++ conn._launch = _launch ++ self.conn = conn ++ ++ def real_forward(self, num_tokens: int = 7) -> None: ++ self.conn._prepare_forward_impl(2, num_tokens, False) ++ ++ ++STALE = CpuGpuSemaphore.DONE_VALUE ++CLEAR = CpuGpuSemaphore.RESET_VALUE ++ ++ ++def test_raise_mode_refuses_to_publish_into_a_stale_flag(monkeypatch): ++ h = _Harness(monkeypatch, mode="raise", tp_rank=0, flag=STALE) ++ with pytest.raises(RuntimeError, match=r"#53899.*\('ple\.1', 1\)"): ++ h.real_forward() ++ assert h.launches == [], "published into a stale flag" ++ assert h.syncs == 1, "the read was not ordered behind the model stream" ++ ++ ++def test_warn_mode_reports_and_still_launches(monkeypatch): ++ h = _Harness(monkeypatch, mode="warn", tp_rank=0, flag=STALE) ++ h.real_forward(num_tokens=11) ++ assert h.launches == [11] ++ assert len(h.warnings) == 1 and "('ple.1', 1)" in h.warnings[0] ++ ++ ++def test_off_mode_does_no_work_on_the_launch_path(monkeypatch): ++ """Production default: no synchronize, no read, no report.""" ++ h = _Harness(monkeypatch, mode="off", tp_rank=0, flag=STALE) ++ h.real_forward() ++ assert h.launches == [7] ++ assert h.syncs == 0 and h.warnings == [] ++ ++ ++def test_clear_flags_on_the_publishing_rank_are_silent(monkeypatch): ++ h = _Harness(monkeypatch, mode="raise", tp_rank=0, flag=CLEAR) ++ h.real_forward() ++ assert h.launches == [7] and h.warnings == [] and h.syncs == 1 ++ ++ ++def test_a_non_publishing_rank_is_not_judged(monkeypatch): ++ """TP1's flag may legitimately carry THIS step's fan-out.""" ++ h = _Harness(monkeypatch, mode="raise", tp_rank=1, flag=STALE) ++ h.real_forward() ++ assert h.launches == [7] and h.warnings == [] and h.syncs == 0 ++ ++ ++def test_dummy_forwards_are_not_checked(monkeypatch): ++ h = _Harness(monkeypatch, mode="raise", tp_rank=0, flag=STALE) ++ dummies: list[int] = [] ++ h.conn._signal_dummy_outputs_impl = dummies.append ++ h.conn._prepare_forward_impl(1, 5, True) ++ assert dummies == [5] and h.syncs == 0 and h.launches == [] ++ ++ ++def _liveness(h: _Harness) -> list[str]: ++ return [line for line in h.infos if line.startswith("PLE flag check")] ++ ++ ++def test_the_publishing_rank_reports_that_the_check_ran(monkeypatch): ++ """Liveness: first check and every heartbeat carry a running count. ++ ++ Zero findings on a diagnostics run is evidence only if the check ran; this line is ++ what a diagnostics run greps. Heartbeat 2: lines at checks 1 and 2, the ++ third (stale, warn) counted and reported by its own warning. ++ """ ++ monkeypatch.setattr(connector_mod, "_E2_CHECK_HEARTBEAT", 2) ++ h = _Harness(monkeypatch, mode="warn", tp_rank=0, flag=CLEAR) ++ h.real_forward() ++ h.real_forward() ++ h.conn._layers["ple.1"] = _layer(STALE) ++ h.real_forward() ++ tail = "(dp_rank=0, tp_rank=0, pp_rank=0)" ++ assert _liveness(h) == [ ++ f"PLE flag check (mode=warn): 1 checks run, 0 with findings {tail}", ++ f"PLE flag check (mode=warn): 2 checks run, 0 with findings {tail}", ++ ] ++ assert h.conn._flag_checks_run == 3 and h.conn._flag_check_findings == 1 ++ assert len(h.warnings) == 1 ++ ++ ++def test_a_non_publishing_rank_says_once_that_it_does_not_judge(monkeypatch): ++ h = _Harness(monkeypatch, mode="warn", tp_rank=1, flag=STALE) ++ h.real_forward() ++ h.real_forward() ++ assert _liveness(h) == [ ++ "PLE flag check (mode=warn): non-publishing rank, not judged here by " ++ "design (dp_rank=0, tp_rank=1, pp_rank=0)" ++ ] ++ ++ ++def test_off_mode_logs_no_liveness(monkeypatch): ++ h = _Harness(monkeypatch, mode="off", tp_rank=0, flag=CLEAR) ++ h.real_forward() ++ assert _liveness(h) == [] ++ ++ ++@pytest.mark.parametrize( ++ ("raw", "mode"), [(None, "off"), ("1", "warn"), ("raise", "raise")] ++) ++def test_the_connector_mode_is_the_registry_value(monkeypatch, raw, mode): ++ """The module-level mode is read through ``envs``, so the value the ++ registry reports and the value the check applies cannot drift.""" ++ if raw is None: ++ monkeypatch.delenv("VLLM_E2_PLE_FLAG_CHECK", raising=False) ++ else: ++ monkeypatch.setenv("VLLM_E2_PLE_FLAG_CHECK", raw) ++ try: ++ reloaded = importlib.reload(connector_mod) ++ assert mode == reloaded._PLE_FLAG_CHECK_MODE ++ assert mode == envs.VLLM_E2_PLE_FLAG_CHECK ++ finally: ++ monkeypatch.delenv("VLLM_E2_PLE_FLAG_CHECK", raising=False) ++ importlib.reload(connector_mod) ++ ++ ++def _diagnostic_knobs() -> list[str]: ++ source = Path(envs.__file__).read_text(encoding="utf-8") ++ return sorted(set(re.findall(r'_e2_diagnostic_mode\(\s*"(\w+)"', source))) ++ ++ ++def test_every_e2_diagnostic_knob_has_a_reader_outside_envs(): ++ """A registered diagnostic with no reader reports "zero findings" forever. ++ ++ Class guard for the defect this port fixes: registration and the read site ++ travelled in different commits and a rebase kept one without the other. ++ """ ++ knobs = _diagnostic_knobs() ++ assert {"VLLM_E2_PLE_FLAG_CHECK", "VLLM_E2_MAMBA_SLOT_OWNERSHIP"} <= set(knobs) ++ root = Path(vllm.__file__).parent ++ readers: dict[str, list[str]] = {k: [] for k in knobs} ++ for path in root.rglob("*.py"): ++ if path.name == "envs.py" and path.parent == root: ++ continue ++ text = path.read_text(encoding="utf-8", errors="replace") ++ for knob in knobs: ++ if f"envs.{knob}" in text: ++ readers[knob].append(str(path.relative_to(root))) ++ unread = [k for k, files in readers.items() if not files] ++ assert unread == [], f"registered but never read via envs: {unread}" +diff --git a/tests/v1/ple_offload/test_stale_flag_assertion.py b/tests/v1/ple_offload/test_stale_flag_assertion.py +new file mode 100644 +index 0000000..c663642 +--- /dev/null ++++ b/tests/v1/ple_offload/test_stale_flag_assertion.py +@@ -0,0 +1,135 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The pull transport's staleness precondition, read back instead of argued. ++ ++vllm#53899 (GB10/sm_121, TP=1, PIECEWISE and FULL_DECODE_ONLY): the capture ++path signalled dummy PLE outputs, the first real step's ++``cuStreamWaitValue32(flag == DONE_VALUE)`` was satisfied by that DUMMY signal, ++its release reset the flag, and the worker's late copy raised it again for the ++next step -- so every forward consumed the PREVIOUS step's PLE rows for the ++life of the server. ``cudagraph_mode=NONE`` was correct. ++ ++This tree holds the precondition structurally (``_signal_dummy_outputs_impl`` ++-> dummy forward -> ``_release_outputs_impl`` resets every layer and ++synchronizes while ``_after_dummy`` is set -> TP barrier at the dummy->real ++transition), but nothing ever read the flag back. ``stale_offload_flags`` is ++that read, and these tests pin its contract. ++ ++No GPU is required: ``CpuGpuSemaphore.from_ipc_tensor`` accepts any int32 ++tensor, and only ``reset``/``signal``/``wait_reset`` (never called here) touch ++the CUDA driver. The layer is a real ``PleOffloadLayer`` subclass configured ++through the real ``setup_cross_process_offload``, not a stand-in. ++""" ++ ++import torch ++ ++from vllm.model_executor.layers.ple_offload_layer import ( ++ CpuGpuSemaphore, ++ PleOffloadLayer, ++ flag_check_is_authoritative, ++ stale_offload_flags, ++) ++ ++HIDDEN = 4 ++ ++ ++class _ProbeLayer(PleOffloadLayer): ++ """Minimal real subclass: the base class is abstract on forward_impl.""" ++ ++ def forward_impl( ++ self, ++ hidden_states: torch.Tensor, ++ input_ids: torch.Tensor, ++ *args: object, ++ **kwargs: object, ++ ) -> torch.Tensor: ++ return hidden_states ++ ++ ++def _layer(flag_value: int) -> _ProbeLayer: ++ layer = _ProbeLayer() ++ flag = torch.tensor([flag_value], dtype=torch.int32) ++ layer.setup_cross_process_offload( ++ gpu_output_buffer=torch.zeros(8, HIDDEN, dtype=torch.bfloat16), ++ semaphore=CpuGpuSemaphore.from_ipc_tensor(flag), ++ ) ++ return layer ++ ++ ++def test_cleared_flags_are_not_reported() -> None: ++ """The correct world: every semaphore clear before a real launch.""" ++ layers = { ++ "ple.0": _layer(CpuGpuSemaphore.RESET_VALUE), ++ "ple.1": _layer(CpuGpuSemaphore.RESET_VALUE), ++ } ++ assert stale_offload_flags(layers) == [] ++ ++ ++def test_a_leftover_signal_is_reported_with_its_layer() -> None: ++ """The #53899 world: a signal no forward consumed survives into a launch. ++ ++ This is the failing control for the pre-launch assertion. Without it the ++ check would be one that cannot fail: a ++ "same prompt twice" comparison passes in the broken world, because the ++ lag is constant and produces two bit-identical output classes. ++ """ ++ layers = { ++ "ple.0": _layer(CpuGpuSemaphore.RESET_VALUE), ++ "ple.1": _layer(CpuGpuSemaphore.DONE_VALUE), ++ } ++ assert stale_offload_flags(layers) == [("ple.1", CpuGpuSemaphore.DONE_VALUE)] ++ ++ ++def test_every_stale_layer_is_named_not_just_the_first() -> None: ++ """Attribution matters: #53899's fix resets EVERY layer's semaphore.""" ++ layers = {f"ple.{i}": _layer(CpuGpuSemaphore.DONE_VALUE) for i in range(3)} ++ assert stale_offload_flags(layers) == [ ++ ("ple.0", 1), ++ ("ple.1", 1), ++ ("ple.2", 1), ++ ] ++ ++ ++def test_a_value_that_is_neither_reset_nor_done_is_reported_verbatim() -> None: ++ """A torn or foreign write is a violation too, and its value is evidence. ++ ++ The flag is binary by construction, so any third value means something ++ other than this transport wrote it. Reporting the raw value distinguishes ++ that from an ordinary leftover signal. ++ """ ++ layers = {"ple.0": _layer(7)} ++ assert stale_offload_flags(layers) == [("ple.0", 7)] ++ ++ ++def test_a_layer_without_cross_process_offload_is_skipped() -> None: ++ """Non-offloaded layers have no ``_sem``; reading one would be an error.""" ++ layers = {"ple.0": _ProbeLayer(), "ple.1": _layer(CpuGpuSemaphore.DONE_VALUE)} ++ assert stale_offload_flags(layers) == [("ple.1", CpuGpuSemaphore.DONE_VALUE)] ++ ++ ++# --------------------------------------------------------------------------- ++# Which rank may judge the flag at all ++# --------------------------------------------------------------------------- ++ ++ ++def test_only_the_publishing_rank_may_judge_the_flag() -> None: ++ """``flag == 0`` before a launch is an invariant on TP0 ONLY. ++ ++ ``_launch`` publishes only on TP0. The offload worker then fans the result ++ out to EVERY registered target, each gated only on that target's own flag ++ having been reset. So this interleaving is legal and healthy: ++ ++ 1. TP0 checks its flag (clear) and publishes step N. ++ 2. The worker fills TP1's buffer and sets TP1's flag for step N. ++ 3. TP1 reaches its own pre-launch check and reads DONE. ++ ++ A rank-blind check calls step 3 a stale signal and terminates a healthy ++ TP2 run. The dummy->real TP barrier is before the checks and does not ++ order this. A non-publishing rank cannot separate "this step's fan-out ++ arrived early" from "last step's signal was never consumed", because the ++ flag is binary and carries no transaction identity -- so it is skipped ++ rather than checked weakly. ++ """ ++ assert flag_check_is_authoritative(0) is True ++ for tp_rank in (1, 2, 3, 7): ++ assert flag_check_is_authoritative(tp_rank) is False, tp_rank +diff --git a/tests/v1/worker/gpu/test_mamba_slot_ownership.py b/tests/v1/worker/gpu/test_mamba_slot_ownership.py +index 89131ef..1532bfa 100644 +--- a/tests/v1/worker/gpu/test_mamba_slot_ownership.py ++++ b/tests/v1/worker/gpu/test_mamba_slot_ownership.py +@@ -159,6 +159,46 @@ def test_an_identity_mapping_hides_the_defect_from_this_check( + assert scenario.ctx.last_slot_ownership_violations == [] + + ++def test_the_check_reports_that_it_ran(monkeypatch: pytest.MonkeyPatch) -> None: ++ """Liveness (PLE diagnostics on 0.30): zero findings must be told apart from no run. ++ ++ A zero-findings result is evidence only if ++ the check executed, so it counts the launches it judged and logs the count ++ on the first check and every ``_E2_CHECK_HEARTBEAT``. Here: two green ++ launches then one red, heartbeat 2 -> lines at checks 1 and 2 only, the ++ red one counted but not logged by the heartbeat. ++ """ ++ _arm(monkeypatch, "warn") ++ monkeypatch.setattr(mamba_utils_mod, "_E2_CHECK_HEARTBEAT", 2) ++ lines: list[str] = [] ++ monkeypatch.setattr( ++ mamba_utils_mod.logger, ++ "info", ++ lambda fmt, *args: lines.append(fmt % args), ++ ) ++ scenario = _build_scenario(slot_mapping=_PERM) ++ _run_deferred(scenario, monkeypatch) ++ _run_deferred(scenario, monkeypatch) ++ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False) ++ _run_deferred(scenario, monkeypatch) ++ ++ ctx = scenario.ctx ++ assert ctx.slot_ownership_checks_run == 3 ++ assert ctx.slot_ownership_checks_with_findings == 1 ++ liveness = [ln for ln in lines if ln.startswith("Mamba slot-ownership check")] ++ assert liveness == [ ++ "Mamba slot-ownership check (mode=warn): 1 checks run, 0 with findings", ++ "Mamba slot-ownership check (mode=warn): 2 checks run, 0 with findings", ++ ] ++ ++ ++def test_off_mode_counts_nothing(monkeypatch: pytest.MonkeyPatch) -> None: ++ _arm(monkeypatch, "off") ++ scenario = _build_scenario(slot_mapping=_PERM) ++ _run_deferred(scenario, monkeypatch) ++ assert scenario.ctx.slot_ownership_checks_run == 0 ++ ++ + # --------------------------------------------------------------------------- + # The three properties, directly on the pure checker + # --------------------------------------------------------------------------- +diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +index df904ff..3f08d4a 100644 +--- a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py ++++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py +@@ -622,16 +622,19 @@ def mask_from_generation_change( + + captured: dict[str, object] = {} + +- def fake_async_copy_to_gpu(array, device=None): ++ # v0.30.0 (#56888) renamed the H2D helper pp_utils uses from ++ # buffer_utils.async_copy_to_gpu to torch_utils.async_tensor_h2d; the ++ # handler still stages the masked mapping through it, so it is the seam. ++ def fake_async_tensor_h2d(array, device=None, dtype=None, out=None): + captured["np"] = np.asarray(array).copy() + return torch.as_tensor(np.asarray(array), dtype=torch.int32) + +- original = pp_utils_mod.async_copy_to_gpu +- pp_utils_mod.async_copy_to_gpu = fake_async_copy_to_gpu ++ original = pp_utils_mod.async_tensor_h2d ++ pp_utils_mod.async_tensor_h2d = fake_async_tensor_h2d + try: + outputs = handler.get_prev_sampled_outputs() + finally: +- pp_utils_mod.async_copy_to_gpu = original ++ pp_utils_mod.async_tensor_h2d = original + + assert outputs is not None, "the whole batch was excluded" + return [int(v) for v in outputs["idx_mapping"].tolist()] +diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py +index cc8bd73..6eab455 100644 +--- a/vllm/model_executor/layers/ple_offload_layer.py ++++ b/vllm/model_executor/layers/ple_offload_layer.py +@@ -476,3 +476,88 @@ class PleOffloadLayer(nn.Module, ABC): + """Mark the cross-process output buffer reusable on ``stream``.""" + if self._is_cpu_offloaded: + self._sem.reset(stream) ++ ++ ++# --------------------------------------------------------------------------- ++# Pre-launch staleness check for the pull transport (#53899 class) ++# --------------------------------------------------------------------------- ++# WHAT IT TESTS. The transport's correctness rests on one precondition: when a ++# real request is published, every layer's semaphore must be CLEAR, so the ++# ``cuStreamWaitValue32(flag == DONE_VALUE)`` inside the forward (line 279 for ++# the pull path, line 215 for the IPC push path) can only be satisfied by the ++# copy THIS step publishes. If a leftover ``DONE_VALUE`` survives into a real ++# step, that wait is satisfied immediately by the PREVIOUS step's signal and ++# the forward consumes the previous step's PLE rows. That is vllm#53899 on ++# GB10/sm_121 exactly: "every forward consumes the previous step's PLE ++# outputs", for the life of the server, with ``cudagraph_mode=NONE`` correct. ++# ++# In this tree the precondition is held structurally -- the dummy signal is ++# consumed by the dummy forward, ``_release_outputs_impl`` resets every layer ++# and synchronizes the model stream while ``_after_dummy`` is set, and the ++# dummy->real transition takes a TP barrier -- but nothing ever READ the flag ++# back, so the structural argument had no instrument. This is that instrument. ++# ++# WHY NOT AN EPOCH WORD. The stronger fix is a monotone epoch: the launch ++# writes ``epoch = step``, the worker signals ``epoch``, and the placeholder ++# waits ``EQ(epoch)``, so a stale signal carries the previous epoch and cannot ++# satisfy the wait. It is not implementable on this transport as it stands. ++# The compare value is an IMMEDIATE OPERAND of ``cuStreamWaitValue32`` ++# (``_ple_offload_wait_copy_impl``, this file, line 279), and that call is ++# issued inside the model forward, which is captured: PIECEWISE and ++# FULL_AND_PIECEWISE both bake the wait into the graph with the value frozen at ++# capture time. Making the compare value change per replay needs the captured ++# node's parameters rewritten before each launch ++# (``cuGraphExecBatchMemOpNodeSetParams``); nothing in this tree's capture path ++# exposes the graph exec handle to the connector, and the wait is reached ++# through a torch custom op inside a compiled region. An epoch word would ++# therefore work only with ``cudagraph_mode=NONE``, i.e. only in the ++# configuration where #53899 does not reproduce. So: binary flag, read back. ++# ++# The check is HOST-SIDE and needs a stream synchronize to be ordered against ++# a pending reset, so it is never run inside a capture and is off by default. ++# --------------------------------------------------------------------------- ++ ++ ++def flag_check_is_authoritative(tp_rank: int) -> bool: ++ """Whether ``flag == RESET_VALUE`` before a launch is an INVARIANT here. ++ ++ Only on the publishing rank (TP0). Why: the offload worker fans a ++ result out to every registered TP target, each gated only on that target's ++ own flag having been reset, so a non-publishing rank's flag can legitimately ++ carry THIS step's signal before that rank reaches its own pre-launch check. ++ Reading a set flag there and calling it stale terminates a healthy run. ++ ++ On TP0 the invariant holds by construction: the worker learns of a request ++ only from ``_launch``, which runs after the check, and the progress monitor ++ refuses overlapping serving submissions, so the previous transaction has ++ been released and (after a stream synchronize) its reset has retired. ++ ++ A non-publishing rank cannot tell "this step's fan-out arrived early" from ++ "last step's signal was never consumed", because the flag is BINARY and ++ carries no transaction identity -- the same representational limit that ++ rules out an epoch word. So those ranks are skipped, not checked weakly. ++ """ ++ return tp_rank == 0 ++ ++ ++def stale_offload_flags( ++ layers: "dict[str, PleOffloadLayer]", ++) -> list[tuple[str, int]]: ++ """Return ``(layer_name, flag_value)`` for every layer whose flag is set. ++ ++ An empty list is the precondition holding. A non-empty list means at least ++ one semaphore still carries a signal that no forward consumed, so the next ++ ``wait`` would be satisfied by the previous step's copy. ++ ++ Pure and side-effect free: the caller owns the ordering (synchronize the ++ model stream first, and never call this inside a stream capture). Layers ++ that are not cross-process offloaded have no semaphore and are skipped. ++ """ ++ stale: list[tuple[str, int]] = [] ++ for name, layer in layers.items(): ++ if not getattr(layer, "_is_cpu_offloaded", False): ++ continue ++ value = int(layer._sem.flag_tensor.reshape(-1)[0].item()) ++ if value != CpuGpuSemaphore.RESET_VALUE: ++ stale.append((name, value)) ++ return stale +diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py +index 2779edf..0d52651 100644 +--- a/vllm/v1/ple_offload/connector.py ++++ b/vllm/v1/ple_offload/connector.py +@@ -18,12 +18,15 @@ import torch.nn as nn + import zmq + from cuda.bindings import driver as cuda_driver + ++import vllm.envs as envs + from vllm.config import VllmConfig + from vllm.distributed.parallel_state import get_dp_group, get_pp_group, get_tp_group + from vllm.logger import init_logger + from vllm.model_executor.layers.ple_offload_layer import ( + CpuGpuSemaphore, + PleOffloadLayer, ++ flag_check_is_authoritative, ++ stale_offload_flags, + ) + from vllm.v1.ple_offload.protocol import ( + PleOffloadRegistration, +@@ -34,6 +37,30 @@ from vllm_ple_test_fault import _FAULT + + logger = init_logger(__name__) + ++# PLE diagnostics / vllm#53899 class: read the pull-transport semaphore back before a real ++# launch and refuse to publish into a flag that is already set. Costs a ++# model-stream synchronize, so it is a diagnostic, not a hot-path guard, and is ++# OFF by default -- the same split VLLM_E2_ZERO_OWNERSHIP_CHECK established in ++# the scheduler (VLLM_E2_ZERO_OWNERSHIP_CHECK). ++# ++# off | warn | raise (1 is accepted as warn). Default off. The default action on ++# a violation is WARN, not raise: this diagnostic could ++# terminate a healthy run, and a diagnostic that kills the qualification boot ++# intended to establish safety is worse than no diagnostic. raise is for a ++# targeted repro. ++# ++# Read through the envs registry, not os.getenv: the registry owns the ++# normalization, so the value reported by ++# `vllm.envs` and the value this check applies cannot drift apart. It is ++# registered as a non-compile factor there. A set-but-unrecognised value ++# raises on first read -- fail closed, because a broken gate that silently ++# disarms its own instrument is the failure mode these checks exist to avoid. ++_PLE_FLAG_CHECK_MODE: str = envs.VLLM_E2_PLE_FLAG_CHECK ++# Liveness cadence for the diagnostic checks: the first check and every Nth ++# are logged with a running count, so "zero findings" can be told apart from ++# "never ran" in a reference-host log. ++_E2_CHECK_HEARTBEAT = 4096 ++ + # CUDA registration alone does not own Python storage. Captured graphs retain + # raw addresses, so every successfully registered host allocation remains + # strongly referenced until process exit. +@@ -701,6 +728,8 @@ class PleOffloadConnector: + self.dp_rank, + self.tp_rank, + ) ++ if _PLE_FLAG_CHECK_MODE != "off": ++ self._check_flags_clear_before_launch(num_tokens) + self._launch( + num_reqs, + num_tokens, +@@ -709,6 +738,128 @@ class PleOffloadConnector: + correlation_id=correlation_id, + ) + ++ def _check_flags_clear_before_launch(self, num_tokens: int) -> None: ++ """Test the transport's staleness precondition, off by default. ++ ++ A real launch publishes rows that the forward will pick up with ++ ``cuStreamWaitValue32(flag == DONE_VALUE)`` ++ (``ple_offload_layer.py:279``). That wait is only meaningful if the ++ flag is CLEAR when we publish: a leftover ``DONE_VALUE`` satisfies it ++ from the PREVIOUS step's copy and the forward silently consumes the ++ previous step's PLE outputs -- vllm#53899, where the semaphore stayed ++ one step ahead for the life of the server and ``cudagraph_mode=NONE`` ++ was the only correct configuration. ++ ++ Nothing here can be inferred from ``_after_dummy``: that is a bool ++ tracking a transition, not the flag's value. The structural argument ++ (dummy signal consumed by the dummy forward, ``_release_outputs_impl`` ++ resetting every layer and synchronizing while ``_after_dummy`` is set, ++ a TP barrier at the transition) holds only while every ++ ``prepare_forward`` is matched by a ``release_outputs``. An exception ++ path, an aborted step, or a future capture that signals without a ++ consuming forward leaves the flag at 1, and the #53899 thread's own ++ trap applies: identical sequential requests produce two bit-identical ++ classes, so a "same prompt twice" check passes in the broken world ++ (an instrument must discriminate both worlds). ++ ++ PUBLISHING RANK ONLY. ``flag == 0`` before a launch is ++ NOT an invariant on a non-publishing rank, and the first version of ++ this check could terminate a HEALTHY TP2 run because of it. ``_launch`` ++ publishes only on TP0 (``if self.tp_rank != 0: return``), and the ++ offload worker then fans the result out to EVERY registered target, ++ each gated only on that target's own flag having been reset ++ (``worker.py``, ``for target in targets: _wait_flag_released(target) ++ ... target.host_flag[0] = 1``). So this legal interleaving exists: ++ ++ 1. TP0 checks its flag (clear), publishes step N. ++ 2. The worker fills TP1's buffer and sets TP1's flag for step N. ++ 3. TP1 reaches its own check for step N and reads ``DONE``. ++ 4. TP1 reports a stale signal -- for a signal that belongs to the ++ CURRENT step. ++ ++ The dummy->real TP barrier is before the checks and does not order this. ++ On TP0 the property does hold, and by construction rather than by ++ timing: the worker learns of a request only from ``_launch``, which is ++ after this check, and ``PleProgressMonitor.begin`` refuses overlapping ++ serving submissions ("Overlapping PLE host submissions violate the ++ serialized single-slot protocol"), so at most one transaction is ++ outstanding and the stream synchronize below has already retired its ++ release. A non-publishing rank cannot distinguish "this step's fan-out ++ arrived early" from "last step's signal was never consumed" with a ++ BINARY flag -- that is the same representational limit that rules out ++ an epoch word (see ``ple_offload_layer.py``), so those ranks are ++ skipped rather than checked weakly. ++ ++ VLLM_E2_PLE_FLAG_CHECK=off|warn|raise, default off, default ACTION warn. ++ ++ SCOPE. The skip below tests whether a capture ++ is ACTIVE, which is not the same as "this step's forward is eager". A ++ captured replay is not an active capture, and ``prepare_forward`` is ++ host Python that runs ahead of the forward either way, so this check ++ DOES run for steps whose forward is a graph replay. What it does not do ++ is run *inside* a capture, where a host synchronize is illegal. The ++ earlier "eager only" wording confused those two lifecycle states. ++ ++ LIVENESS (PLE diagnostics on 0.30). A zero-findings result is only evidence if the ++ check demonstrably RAN, so every rank says once what it does with the ++ check, and the publishing rank logs a running count of checks and ++ findings on its first check and every ``_E2_CHECK_HEARTBEAT`` checks ++ after. A reference-host test arm greps these lines; their absence voids the arm. ++ """ ++ if not flag_check_is_authoritative(self.tp_rank): ++ if not getattr(self, "_flag_check_skip_logged", False): ++ self._flag_check_skip_logged = True ++ logger.info( ++ "PLE flag check (mode=%s): non-publishing rank, not judged " ++ "here by design (dp_rank=%d, tp_rank=%d, pp_rank=%d)", ++ _PLE_FLAG_CHECK_MODE, ++ self.dp_rank, ++ self.tp_rank, ++ self.pp_rank, ++ ) ++ return ++ if torch.cuda.is_current_stream_capturing(): ++ return ++ stream = torch.cuda.current_stream(self.device) ++ # Order the read against any reset still queued on the model stream. ++ # Without this the check reports the pre-reset value and fires on a ++ # correct run -- a guard that fires on correct metadata is worse than ++ # no guard. ++ stream.synchronize() ++ stale = stale_offload_flags(self._layers) ++ runs = self._flag_checks_run = getattr(self, "_flag_checks_run", 0) + 1 ++ findings = getattr(self, "_flag_check_findings", 0) + (1 if stale else 0) ++ self._flag_check_findings = findings ++ if runs == 1 or runs % _E2_CHECK_HEARTBEAT == 0: ++ logger.info( ++ "PLE flag check (mode=%s): %d checks run, %d with findings " ++ "(dp_rank=%d, tp_rank=%d, pp_rank=%d)", ++ _PLE_FLAG_CHECK_MODE, ++ runs, ++ findings, ++ self.dp_rank, ++ self.tp_rank, ++ self.pp_rank, ++ ) ++ if not stale: ++ return ++ message = ( ++ "PLE pull transport: the output semaphore was already signalled " ++ "before this publishing rank published a real request, so the " ++ "forward's cuStreamWaitValue32(flag == 1) would be satisfied by " ++ "the PREVIOUS step's copy and this forward would consume the " ++ "previous step's PLE rows (vllm#53899 one-step-lag class). Stale " ++ f"layers (name, flag): {stale}; expected " ++ f"{CpuGpuSemaphore.RESET_VALUE} for all {len(self._layers)} " ++ f"layer(s). num_tokens={num_tokens}, dp_rank={self.dp_rank}, " ++ f"tp_rank={self.tp_rank}, pp_rank={self.pp_rank}." ++ ) ++ if _PLE_FLAG_CHECK_MODE == "raise": ++ # Reaches _fail_closed through prepare_forward, which is the ++ # transport's stated contract for an uncertain correspondence. ++ raise RuntimeError(message) ++ logger.warning("%s", message) ++ + def signal_dummy_outputs(self, num_tokens: int) -> None: + """Locally satisfy PLE waits for dummy and capture forwards.""" + self._progress.begin(startup=True, num_tokens=num_tokens) +diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py +index eeeadae..a89383b 100644 +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -75,6 +75,9 @@ ALIGN_BT_ROW_IS_REQ = True + # registered as a non-compile factor there, and an unrecognised value raises + # on first read rather than silently disarming the instrument. + _SLOT_OWNERSHIP_MODE: str = envs.VLLM_E2_MAMBA_SLOT_OWNERSHIP ++# Liveness cadence: first check and every Nth are logged with a running count ++# (see MambaSpecDecodeGPUContext.slot_ownership_checks_run). ++_E2_CHECK_HEARTBEAT = 4096 + + _NULL_BLOCK_ID = 0 + +@@ -1230,6 +1233,11 @@ class MambaSpecDecodeGPUContext: + last_slot_ownership_violations: list[str] = dataclasses.field( + default_factory=list + ) ++ # Liveness for the same check (diagnostics on 0.30): how many launches it has ++ # judged and how many had findings. Logged on the first check and every ++ # _E2_CHECK_HEARTBEAT after, so a zero-findings run can prove it ran. ++ slot_ownership_checks_run: int = 0 ++ slot_ownership_checks_with_findings: int = 0 + + @classmethod + def create( +@@ -1932,6 +1940,18 @@ class MambaSpecDecodeGPUContext: + block_size=int(self.block_size), + ) + self.last_slot_ownership_violations = violations ++ self.slot_ownership_checks_run += 1 ++ if violations: ++ self.slot_ownership_checks_with_findings += 1 ++ runs = self.slot_ownership_checks_run ++ if runs == 1 or runs % _E2_CHECK_HEARTBEAT == 0: ++ logger.info( ++ "Mamba slot-ownership check (mode=%s): %d checks run, %d with " ++ "findings", ++ _SLOT_OWNERSHIP_MODE, ++ runs, ++ self.slot_ownership_checks_with_findings, ++ ) + if not violations: + return violations + message = ( +-- +2.43.0 + diff --git a/patches/0052-code-quality-on-v0.30.0-row-dump-removed-lint-knob-r.patch b/patches/0052-code-quality-on-v0.30.0-row-dump-removed-lint-knob-r.patch new file mode 100644 index 0000000..f842b14 --- /dev/null +++ b/patches/0052-code-quality-on-v0.30.0-row-dump-removed-lint-knob-r.patch @@ -0,0 +1,1669 @@ +From e67478a0a8d4f487b63ff2cb44674c4114257fad Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 10:06:57 -0400 +Subject: [PATCH 052/103] code quality on v0.30.0: row dump removed, lint, knob + read-site census + +The debug PLE row dump is deleted; ruff findings on fork lines cleared; a +read-site census test pins which fork knobs are read and where. +--- + tests/kernels/test_top_k_per_row.py | 30 ++- + tests/test_envs.py | 186 +++++++++++++++++- + tests/test_merlin_e2_toggles.py | 1 - + tests/v1/ple_offload/test_correspondence.py | 51 +++-- + tests/v1/ple_offload/test_fatal_rpc.py | 15 +- + .../ple_offload/test_fault_injector_basics.py | 8 +- + .../ple_offload/test_fault_injector_gates.py | 14 +- + .../v1/ple_offload/test_fault_wire_frames.py | 6 +- + .../v1/ple_offload/test_ipc_socket_unlink.py | 3 +- + .../v1/ple_offload/test_progress_deadline.py | 16 +- + tests/v1/ple_offload/test_rowdump_deleted.py | 49 +++++ + tests/v1/ple_offload/test_shm_unlink.py | 4 +- + tests/v1/ple_offload/test_shutdown_order.py | 2 +- + .../ple_offload/test_spawn_bootstrap_cuda.py | 15 +- + tests/v1/ple_offload/test_test_fault.py | 20 +- + .../ple_offload/test_transport_bounded_io.py | 5 +- + .../test_transport_fatal_helper.py | 4 +- + tests/v1/test_ple_pull_compile_profile.py | 33 +++- + tests/v1/worker/test_e2_memstats.py | 43 +++- + vllm/envs.py | 47 +++-- + .../qwen4_exp/nvidia/ops/qsa_indexer.py | 2 +- + vllm/v1/ple_offload/connector.py | 18 +- + vllm/v1/ple_offload/fatal.py | 19 +- + vllm/v1/ple_offload/progress.py | 15 +- + vllm/v1/ple_offload/rowdump.py | 41 ---- + vllm/v1/ple_offload/worker.py | 19 +- + vllm/v1/worker/e2_memstats.py | 176 +++++++++++------ + vllm/v1/worker/gpu/cudagraph_utils.py | 5 +- + vllm_ple_offload_bootstrap/__init__.py | 6 +- + 29 files changed, 566 insertions(+), 287 deletions(-) + create mode 100644 tests/v1/ple_offload/test_rowdump_deleted.py + delete mode 100644 vllm/v1/ple_offload/rowdump.py + +diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py +index 58bcb7f..e0c3ac6 100644 +--- a/tests/kernels/test_top_k_per_row.py ++++ b/tests/kernels/test_top_k_per_row.py +@@ -1081,7 +1081,9 @@ def run_large_context_topk_test( + Torch: {torch_vals.sort(descending=True)[0][:10]}""" + + +-@pytest.mark.skipif(not _has_device_capability(90), reason="cooperative_topk is sm_90+ only") ++@pytest.mark.skipif( ++ not _has_device_capability(90), reason="cooperative_topk is sm_90+ only" ++) + @pytest.mark.parametrize("top_k", [512, 1024, 2048]) + def test_cooperative_topk_cs2(top_k: int) -> None: + """The 64-row dispatch uses the two-CTA cooperative kernel. +@@ -2185,9 +2187,10 @@ def test_persistent_topk_decode_overflow_is_detected() -> None: + mid = np.nextafter(np.nextafter(base, np.float32(np.inf)), np.float32(np.inf)) + high = np.nextafter(np.nextafter(mid, np.float32(np.inf)), np.float32(np.inf)) + assert base < mid < high +- assert ( +- np.float16(base) == np.float16(mid) == np.float16(high) +- ), "the three levels must share one half-precision bin or the row is not adversarial" ++ assert np.float16(base) == np.float16(mid) == np.float16(high), ( ++ "the three levels must share one half-precision bin or the row is not " ++ "adversarial" ++ ) + + logits = torch.full((num_rows, seq_len), float(base), dtype=torch.float32) + logits[:, :leaders] = float(high) +@@ -2630,7 +2633,8 @@ def test_persistent_topk_large_fault_in_nonzero_cta_fails_closed() -> None: + monkeypatched inside a session that has already run a top-k kernel. Run + this test in its own process: + +- MERLIN_TOPK_DET_LARGE=1 MERLIN_TOPK_LARGE_FAULT_INJECT=1 pytest -k large_fault_in_nonzero_cta ++ export MERLIN_TOPK_DET_LARGE=1 MERLIN_TOPK_LARGE_FAULT_INJECT=1 ++ pytest -k large_fault_in_nonzero_cta + + Setting the injector also suppresses __trap(), so this is the trap-off arm + of the policy. The trap-on arm kills the CUDA context by construction and +@@ -2926,7 +2930,9 @@ def _assert_valid_topk_membership( + ) + + +-def _largedet_distribution_impl(kind: str, seq_len: int, top_k: int, seed: int) -> torch.Tensor: ++def _largedet_distribution_impl( ++ kind: str, seq_len: int, top_k: int, seed: int ++) -> torch.Tensor: + """One CPU float32 row. All values are positive, so no -0.0/+0.0 hazard.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + +@@ -3435,8 +3441,8 @@ def _isolated_decode_drop_fixture(top_k: int) -> tuple[torch.Tensor, torch.Tenso + winners = 4 * (warps * 32 + 31) + 3 + row[winners] = float(np.float32(1.01)) + row[4096 : 4096 + above] = 2.0 +- assert 4096 > ISOLATED_DECODE_DBUF +- assert 32 <= top_k - above ++ assert ISOLATED_DECODE_DBUF < 4096 ++ assert top_k - above >= 32 + return row, winners + + +@@ -3630,7 +3636,9 @@ def _isolated_tie_fixture(path: str) -> tuple[torch.Tensor, int, int]: + @pytest.mark.parametrize("path", ["decode", "medium"]) + @pytest.mark.parametrize("want_recent", [True, False]) + @torch.inference_mode() +-def test_isolated_tie_det_nonoverflow_cutoff_direction(path: str, want_recent: bool) -> None: ++def test_isolated_tie_det_nonoverflow_cutoff_direction( ++ path: str, want_recent: bool ++) -> None: + """Cell 4: MERLIN_TIE_DET's buffered cutoff tie, both directions, oracle. + + Why: the mixed-path test's sorted-set equality of two launches is +@@ -3657,7 +3665,9 @@ def test_isolated_tie_det_nonoverflow_cutoff_direction(path: str, want_recent: b + row, top_k, need = _isolated_tie_fixture(path) + expected = _largedet_reference(row, top_k, tie_recent=want_recent) + other = _largedet_reference(row, top_k, tie_recent=not want_recent) +- assert not torch.equal(expected, other), "the fixture cannot tell the directions apart" ++ assert not torch.equal(expected, other), ( ++ "the fixture cannot tell the directions apart" ++ ) + + before = _isolated_counters_or_skip() + first = _isolated_launch(row, top_k) +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 11700fa..0db9be3 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -644,9 +644,15 @@ _FORK_IGNORED = [ + ("VLLM_PLE_TEST_FAULT", "cut=exc_in_release,ordinal=3,rank=any", None), + ("VLLM_PLE_TEST_FAULT_GATE", "/var/tmp/gate", "/var/tmp/gate"), + ("VLLM_PLE_TEST_FAULT_LOG", "/var/tmp/fault.log", "/var/tmp/fault.log"), +- # Read sites arrive with the slot-ownership assertions. ++ # PLE diagnostics gates, both read through vllm.envs (connector.py / mamba_utils.py); ++ # test_read_site_status_matches_the_tables pins that the readers exist. + ("VLLM_E2_PLE_FLAG_CHECK", "1", "warn"), + ("VLLM_E2_MAMBA_SLOT_OWNERSHIP", " RAISE ", "raise"), ++ # Medium-path counter reader: a host-side log cadence, and the older slot-count ++ # fallback that this build never consults. Registered and ignored; before ++ # the read-site census neither classification was asserted anywhere. ++ ("MERLIN_TOPK_EXACT_COUNTERS_LOG", "30", 30), ++ ("MERLIN_TOPK_COUNTER_SLOTS", "11", "11"), + ] + + # Marker fence around the unaudited-knob block in envs.py. The gate below scans +@@ -915,3 +921,181 @@ def test_fork_cache_factor_lists_are_disjoint_and_registered(): + assert not hashed & ignored + for name in hashed | ignored: + assert name in environment_variables, name ++ ++ ++# --------------------------------------------------------------------------- ++# Read-site census on 0.30. Every table above is a list somebody wrote; none ++# of them is checked against the READ SITES. That is how the 0.30 release tree ++# came to read VLLM_PLE_ROW_DUMP (restored by the pull-transport carry after it was deleted ++# it) and MERLIN_TOPK_COUNTER_SLOTS (the medium-path counter reader) with neither classified, ++# and how MERLIN_TOPK_EXACT_COUNTERS_LOG sat in ignored_factors with no test ++# asserting it. The census closes the loop in both directions: ++# ++# source -> tables: every fork-family name the tree mentions (Python, C++ and ++# the two PLE bootstrap packages) is hashed, ignored, fenced as unaudited, ++# a build-time define, or upstream's own registration. ++# registry -> tables: every fork-family name registered in envs.py has its ++# cache-factor classification asserted by one of the tables. ++# ++# A knob that changes compiled code but is missing from the cache key is a ++# correctness bug on an AOT cache-load boot (the loaded artifact skips guard ++# evaluation), so "unclassified" must fail here rather than on the reference host. ++# --------------------------------------------------------------------------- ++ ++_FORK_FAMILY = re.compile( ++ r"\b(?:MERLIN_[A-Z0-9_]+|VLLM_(?:E2|PLE|QSA)_[A-Z0-9_]+|VLLM_HOST_EMBED_TABLE)\b" ++) ++ ++# Preprocessor defines consumed at BUILD time (-D / #ifndef / #define), never a ++# getenv; test_exact_trap_is_a_build_define_not_an_env_var pins the reasoning ++# for the first, and test_build_defines_are_not_env_reads re-checks them all. ++_BUILD_DEFINES = [ ++ "MERLIN_TOPK_EXACT_TRAP", ++ "MERLIN_TOPK_DIAG_DELAY", ++ "MERLIN_TOPK_DIAG_PREREPAIR", ++ "MERLIN_TOPK_DIAG_FORCE_FAULT", ++ "MERLIN_TOPK_DIAG_FORCE_FAULT_CTA", ++ "MERLIN_TOPK_DIAG_ANY", ++] ++ ++# In the fork's name family but registered and hashed by UPSTREAM v0.30.0 ++# (the pinned-host UVA PLE switch the PLE pull transport refuses). ++_UPSTREAM_OWNED = ["VLLM_PLE_CPU_OFFLOAD"] ++ ++_CENSUS_ROOTS = ("vllm", "csrc", "vllm_ple_offload_bootstrap", "vllm_ple_test_fault") ++_CENSUS_SUFFIXES = {".py", ".cu", ".cuh", ".cpp", ".h", ".c"} ++ ++ ++def _all_tabled() -> set[str]: ++ tables = _FORK_HASHED + _FORK_IGNORED + _TIEDET_HASHED + _TIEDET_KERNEL_SWITCHES ++ return {n for n, _, _ in tables + _LARGEDET_FAULT_INJECTORS} ++ ++ ++def _classified() -> set[str]: ++ return ( ++ _all_tabled() | set(_TIEDET_UNAUDITED) | set(_BUILD_DEFINES) | set(_UPSTREAM_OWNED) ++ ) ++ ++ ++def _census(sources) -> dict[str, set[str]]: ++ """{fork-family name: files mentioning it} over (relpath, text) pairs.""" ++ found: dict[str, set[str]] = {} ++ for rel, text in sources: ++ for name in _FORK_FAMILY.findall(text): ++ found.setdefault(name, set()).add(rel) ++ return found ++ ++ ++def _tree_sources(): ++ root = Path(envs.__file__).resolve().parents[1] ++ envs_py = Path(envs.__file__).resolve() ++ for sub in _CENSUS_ROOTS: ++ base = root / sub ++ if not base.is_dir(): ++ continue ++ for path in base.rglob("*"): ++ if path.suffix not in _CENSUS_SUFFIXES or path.resolve() == envs_py: ++ continue ++ text = path.read_text(encoding="utf-8", errors="replace") ++ yield str(path.relative_to(root)), text ++ ++ ++def test_every_fork_knob_in_the_source_is_classified(): ++ found = _census(_tree_sources()) ++ assert found, "census scanned nothing: the source tree is not beside envs.py" ++ missing = {n: sorted(f) for n, f in found.items() if n not in _classified()} ++ assert not missing, ( ++ "fork knobs read or named in the tree but in no classification table " ++ "(hash it, ignore it with a reason, fence it, or delete the reader): " ++ f"{missing}" ++ ) ++ ++ ++def test_every_registered_fork_knob_has_an_asserted_classification(): ++ registered = {n for n in environment_variables if _FORK_FAMILY.fullmatch(n)} ++ untested = sorted(registered - _all_tabled() - set(_UPSTREAM_OWNED)) ++ assert not untested, ( ++ f"registered fork knobs whose hashed/ignored status no test asserts: {untested}" ++ ) ++ ++ ++def test_build_defines_are_not_env_reads(): ++ """A name classified as a build define must never be read at run time.""" ++ reads = re.compile(r"(?:getenv|environ(?:\.get)?)\s*[\(\[]\s*[\"'](\w+)[\"']") ++ runtime: set[str] = set() ++ for _, text in _tree_sources(): ++ runtime.update(reads.findall(text)) ++ assert not runtime & set(_BUILD_DEFINES), runtime & set(_BUILD_DEFINES) ++ for name in _BUILD_DEFINES: ++ assert name not in environment_variables, name ++ ++ ++def test_census_discriminates_both_worlds(): ++ """Control: the census must flag an unclassified reader and pass a ++ classified one -- a scan that reports the same thing in both measures ++ nothing.""" ++ unknown = [("x.py", 'os.environ.get("VLLM_PLE_EXAMPLE_KNOB", "")')] ++ known = [("x.py", 'os.environ.get("VLLM_QSA_KV_CLIP_COUNT", "")')] ++ cxx = [("x.cu", 'const char* raw = std::getenv("MERLIN_EXAMPLE_SWITCH");')] ++ assert set(_census(unknown)) - _classified() == {"VLLM_PLE_EXAMPLE_KNOB"} ++ assert set(_census(known)) - _classified() == set() ++ assert set(_census(cxx)) - _classified() == {"MERLIN_EXAMPLE_SWITCH"} ++ # Upstream names outside the family never enter the census. ++ assert _census([("x.py", 'os.getenv("VLLM_USE_AOT_COMPILE")')]) == {} ++ ++ ++# Read sites: the census above proves every knob is CLASSIFIED; it says ++# nothing about whether a knob is READ. The envs.py comments do, and on ++# the earlier base they were right that VLLM_E2_PLE_FLAG_CHECK had no reader. The diagnostics change ++# relanded its reader (connector.py), so the tables now say which knobs must ++# have a read site and which are INERT, and this test checks both against the ++# tree. A regression either way (a reader dropped by a carry, or an inert knob ++# silently wired) fails here instead of in a reference-host arm that reports "zero ++# findings" from an instrument that never ran. ++_MUST_BE_READ = { ++ "VLLM_E2_PLE_FLAG_CHECK": "vllm/v1/ple_offload/connector.py", ++ "VLLM_E2_MAMBA_SLOT_OWNERSHIP": "vllm/v1/worker/mamba_utils.py", ++} ++_INERT_NO_READ_SITE = ["VLLM_PLE_FP8_GLOBAL_SCALE"] ++ ++_READ_SITE = re.compile( ++ r"\benvs\.(\w+)\b" ++ r"|(?:getenv|environ(?:\.get)?)\s*[\(\[]\s*[\"'](\w+)[\"']" ++) ++ ++ ++def _read_sites(sources) -> dict[str, set[str]]: ++ """{name: files that READ it} (envs. or an environ/getenv read).""" ++ sites: dict[str, set[str]] = {} ++ for rel, text in sources: ++ for attr, raw in _READ_SITE.findall(text): ++ sites.setdefault(attr or raw, set()).add(rel.replace("\\", "/")) ++ return sites ++ ++ ++def test_read_site_status_matches_the_tables(): ++ sites = _read_sites(_tree_sources()) ++ for name, where in _MUST_BE_READ.items(): ++ assert where in sites.get(name, set()), ( ++ f"{name} must be read in {where}; read sites found: " ++ f"{sorted(sites.get(name, set()))}" ++ ) ++ for name in _INERT_NO_READ_SITE: ++ assert name not in sites, ( ++ f"{name} is documented INERT but is read in {sorted(sites[name])}: " ++ "move it to the compile factors or out of the inert list" ++ ) ++ # Both tables are classified knobs, and disjoint. ++ assert set(_MUST_BE_READ) <= _classified() ++ assert set(_INERT_NO_READ_SITE) <= _classified() ++ assert not set(_MUST_BE_READ) & set(_INERT_NO_READ_SITE) ++ ++ ++def test_read_site_scan_discriminates_both_worlds(): ++ """Control: a mention in a docstring is not a read; envs.X and getenv are.""" ++ mention = [("a.py", '"""VLLM_E2_PLE_FLAG_CHECK=off|warn|raise"""')] ++ via_envs = [("a.py", "_M: str = envs.VLLM_E2_PLE_FLAG_CHECK")] ++ via_getenv = [("a.py", 'os.getenv("VLLM_PLE_FP8_GLOBAL_SCALE", "0")')] ++ assert "VLLM_E2_PLE_FLAG_CHECK" not in _read_sites(mention) ++ assert _read_sites(via_envs)["VLLM_E2_PLE_FLAG_CHECK"] == {"a.py"} ++ assert _read_sites(via_getenv)["VLLM_PLE_FP8_GLOBAL_SCALE"] == {"a.py"} +diff --git a/tests/test_merlin_e2_toggles.py b/tests/test_merlin_e2_toggles.py +index 3cd3c1f..2ded113 100644 +--- a/tests/test_merlin_e2_toggles.py ++++ b/tests/test_merlin_e2_toggles.py +@@ -20,7 +20,6 @@ import pytest + + from vllm.v1.worker import e2_toggles as t + +- + # ------------------------------------------------------------ atoi parity ---- + ATOI_CASES = [ + (b"1", 1), +diff --git a/tests/v1/ple_offload/test_correspondence.py b/tests/v1/ple_offload/test_correspondence.py +index e8d8b9e..b7eaf9c 100644 +--- a/tests/v1/ple_offload/test_correspondence.py ++++ b/tests/v1/ple_offload/test_correspondence.py +@@ -20,16 +20,15 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, + ) ++from vllm.v1.worker.gpu.model_runner import GPUModelRunner + from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( + SpeculatorCudaGraphManager, + ) +-from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + + class _ConnectorForValidation: + _input_ids_buf = torch.empty(4, dtype=torch.int32) + _query_start_loc_buf = torch.empty(3, dtype=torch.int32) +- _rowdump_n = None + + def _signal_dummy_outputs_impl(self, num_tokens: int) -> None: + raise AssertionError("must fail before dummy publication") +@@ -60,9 +59,7 @@ def _runner(query_start_loc: list[int]) -> PleOffloadRunner: + runner._worker_targets = {0: {"ple": [target]}} + runner._scratch_bufs = {0: {"ple": torch.empty(4, 2)}} + query = torch.zeros(3, dtype=torch.int32) +- query[: len(query_start_loc)] = torch.tensor( +- query_start_loc, dtype=torch.int32 +- ) ++ query[: len(query_start_loc)] = torch.tensor(query_start_loc, dtype=torch.int32) + runner._input_bufs = { + 0: PleOffloadInputBuffers( + input_ids_buf=torch.zeros(4, dtype=torch.int32), +@@ -76,14 +73,14 @@ def _runner(query_start_loc: list[int]) -> PleOffloadRunner: + @pytest.mark.parametrize( + ("ple_request", "query_start_loc", "match"), + [ +- (PleOffloadRequest(dp_rank=0, num_tokens=3, num_reqs=2), [0, 2, 1], +- "query_start_loc"), +- (PleOffloadRequest(dp_rank=0, num_tokens=5, num_reqs=1), [0, 1], +- "num_tokens"), +- (PleOffloadRequest(dp_rank=0, num_tokens=1, num_reqs=3), [0, 1, 1], +- "num_reqs"), +- (PleOffloadRequest(dp_rank=9, num_tokens=1, num_reqs=1), [0, 1], +- "dp_rank=9"), ++ ( ++ PleOffloadRequest(dp_rank=0, num_tokens=3, num_reqs=2), ++ [0, 2, 1], ++ "query_start_loc", ++ ), ++ (PleOffloadRequest(dp_rank=0, num_tokens=5, num_reqs=1), [0, 1], "num_tokens"), ++ (PleOffloadRequest(dp_rank=0, num_tokens=1, num_reqs=3), [0, 1, 1], "num_reqs"), ++ (PleOffloadRequest(dp_rank=9, num_tokens=1, num_reqs=1), [0, 1], "dp_rank=9"), + ], + ) + def test_invalid_request_fails_before_publication( +@@ -133,9 +130,7 @@ def test_selected_full_graph_copy_count_precedes_publication() -> None: + ) + + with pytest.raises(RuntimeError, match="selected-forward"): +- GPUModelRunner._prepare_ple_forward( +- runner, input_batch, desc, dummy_run=False +- ) ++ GPUModelRunner._prepare_ple_forward(runner, input_batch, desc, dummy_run=False) + assert connector.called is False + + +@@ -150,11 +145,13 @@ def test_production_capture_callback_rejects_actual_ple_copy_mismatch() -> None: + num_reqs=1, + ) + +- with pytest.raises(RuntimeError, match="actual PLE copy geometry"): +- # This is the production capture scope. The injected observation uses +- # the same hook called by the real wait-copy custom-op implementation. +- with manager.capture_ple_copy_geometry(desc): +- _record_ple_copy_geometry("ple", 12, 16) ++ # capture_ple_copy_geometry is the production capture scope. The injected ++ # observation uses the same hook the real wait-copy custom-op calls. ++ with ( ++ pytest.raises(RuntimeError, match="actual PLE copy geometry"), ++ manager.capture_ple_copy_geometry(desc), ++ ): ++ _record_ple_copy_geometry("ple", 12, 16) + assert desc not in manager._ple_full_copy_geometry + + +@@ -169,9 +166,11 @@ def test_speculator_full_capture_refuses_observed_ple_op() -> None: + num_reqs=1, + ) + +- with pytest.raises(RuntimeError, match="speculator/draft"): +- with manager.capture_ple_copy_geometry(desc): +- _record_ple_copy_geometry("unexpected-draft-ple", 8, 16) ++ with ( ++ pytest.raises(RuntimeError, match="speculator/draft"), ++ manager.capture_ple_copy_geometry(desc), ++ ): ++ _record_ple_copy_geometry("unexpected-draft-ple", 8, 16) + assert desc not in manager._ple_full_copy_geometry + + +@@ -191,7 +190,5 @@ def test_eager_correspondence_uses_actual_padded_input_view() -> None: + num_reqs=1, + ) + with pytest.raises(RuntimeError, match="selected-forward"): +- GPUModelRunner._prepare_ple_forward( +- runner, input_batch, desc, dummy_run=False +- ) ++ GPUModelRunner._prepare_ple_forward(runner, input_batch, desc, dummy_run=False) + assert connector.called is False +diff --git a/tests/v1/ple_offload/test_fatal_rpc.py b/tests/v1/ple_offload/test_fatal_rpc.py +index 1c86f9b..2b10492 100644 +--- a/tests/v1/ple_offload/test_fatal_rpc.py ++++ b/tests/v1/ple_offload/test_fatal_rpc.py +@@ -80,10 +80,12 @@ class _Runner: + + def _rpc_loop(where: str, successor_count: Any) -> None: + runner = _Runner(where) +- try: ++ try: # noqa: SIM105 + runner.execute_model() + except Exception: +- # Mirrors WorkerProc._execute_worker_rpc: this must never regain control. ++ # Mirrors WorkerProc._execute_worker_rpc: this must never regain ++ # control. SIM105 declined -- the point of the fake is to have the ++ # same shape as the code it stands in for. + pass + with successor_count.get_lock(): + successor_count.value += 1 +@@ -100,8 +102,9 @@ def test_ple_exception_kills_workerproc_style_rpc( + args=(where, successor_count), + ) + process.start() +- # Spawned children import vLLM cold (measured 4-8 s on the reference host); the bound is +- # about start-up, not the fatal path, which exits immediately. ++ # Spawned children import vLLM cold (measured 4-8 s on the reference ++ # host); the bound is about start-up, not the fatal path, which exits ++ # immediately. + process.join(timeout=120) + assert not process.is_alive() + assert process.exitcode not in (None, 0) +@@ -115,9 +118,7 @@ def _rpc_loop_connector_only(where: str, successor_count: Any) -> None: + connector = _InjectedConnector(where) + try: + if where == "begin": +- PleOffloadConnector.begin_forward( +- connector, dummy_run=False, num_tokens=1 +- ) ++ PleOffloadConnector.begin_forward(connector, dummy_run=False, num_tokens=1) + elif where == "prepare": + PleOffloadConnector.prepare_forward( + connector, 1, 1, False, consumer_num_tokens=1 +diff --git a/tests/v1/ple_offload/test_fault_injector_basics.py b/tests/v1/ple_offload/test_fault_injector_basics.py +index a47796d..c804729 100644 +--- a/tests/v1/ple_offload/test_fault_injector_basics.py ++++ b/tests/v1/ple_offload/test_fault_injector_basics.py +@@ -38,9 +38,7 @@ def test_impossible_role_selectors_are_rejected(monkeypatch, tmp_path, value): + + def test_armed_import_requires_dedicated_log(tmp_path): + env = dict(os.environ) +- env["VLLM_PLE_TEST_FAULT"] = ( +- "cut=exc_before_prepare,ordinal=1,rank=any" +- ) ++ env["VLLM_PLE_TEST_FAULT"] = "cut=exc_before_prepare,ordinal=1,rank=any" + env.pop("VLLM_PLE_TEST_FAULT_LOG", None) + proc = subprocess.run( + [sys.executable, "-c", "import vllm_ple_test_fault"], +@@ -81,9 +79,7 @@ def test_fresh_interpreter_import_identity(tmp_path, order, armed): + env.pop("VLLM_PLE_TEST_FAULT", None) + env.pop("VLLM_PLE_TEST_FAULT_LOG", None) + if armed: +- env["VLLM_PLE_TEST_FAULT"] = ( +- "cut=launch_silent_ready,ordinal=1,rank=any" +- ) ++ env["VLLM_PLE_TEST_FAULT"] = "cut=launch_silent_ready,ordinal=1,rank=any" + env["VLLM_PLE_TEST_FAULT_LOG"] = str(tmp_path / f"{order}.log") + script = f""" + import sys +diff --git a/tests/v1/ple_offload/test_fault_injector_gates.py b/tests/v1/ple_offload/test_fault_injector_gates.py +index 70e0d3e..116236c 100644 +--- a/tests/v1/ple_offload/test_fault_injector_gates.py ++++ b/tests/v1/ple_offload/test_fault_injector_gates.py +@@ -65,6 +65,7 @@ def test_disarmed_sender_deadline_does_not_enter_logger(monkeypatch): + + def test_disarmed_launch_does_not_read_progress_sequence(monkeypatch): + import queue ++ + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) +@@ -94,7 +95,6 @@ def test_disarmed_release_does_not_read_progress_sequence(monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) +- monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) + event = Mock() + stream = Mock() + monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: event) +@@ -124,9 +124,7 @@ def test_sample_tokens_connector_rank_fatalizes(): + pass + + class Runner: +- _ple_offload_connector = SimpleNamespace( +- _fail_closed=Mock(side_effect=Reached) +- ) ++ _ple_offload_connector = SimpleNamespace(_fail_closed=Mock(side_effect=Reached)) + + @fail_ple_execute_model + def sample_tokens(self): +@@ -458,8 +456,12 @@ def test_every_fault_call_has_one_boolop_guard_and_no_prior_state_read(): + assert call_index is not None + assert guard_index < call_index + +- forbidden = {"sequence", "has_pending_completions", +- "completion_is_pending", "fault_transaction"} ++ forbidden = { ++ "sequence", ++ "has_pending_completions", ++ "completion_is_pending", ++ "fault_transaction", ++ } + before_guard = owner.values[:guard_index] + assert not any( + isinstance(node, ast.Attribute) and node.attr in forbidden +diff --git a/tests/v1/ple_offload/test_fault_wire_frames.py b/tests/v1/ple_offload/test_fault_wire_frames.py +index b864355..a237eaa 100644 +--- a/tests/v1/ple_offload/test_fault_wire_frames.py ++++ b/tests/v1/ple_offload/test_fault_wire_frames.py +@@ -34,7 +34,11 @@ def test_armed_offload_receive_requires_metadata_frame(monkeypatch): + + monkeypatch.setattr(worker_module, "_FAULT", object()) + request = _request() +- metadata = {"serving": True, "transaction": "request:1", "correlation": "sequence:3"} ++ metadata = { ++ "serving": True, ++ "transaction": "request:1", ++ "correlation": "sequence:3", ++ } + socket = Mock() + socket.recv.side_effect = AssertionError("single-frame read while armed") + socket.recv_multipart.return_value = [ +diff --git a/tests/v1/ple_offload/test_ipc_socket_unlink.py b/tests/v1/ple_offload/test_ipc_socket_unlink.py +index 92fbc42..0216d01 100644 +--- a/tests/v1/ple_offload/test_ipc_socket_unlink.py ++++ b/tests/v1/ple_offload/test_ipc_socket_unlink.py +@@ -2,18 +2,17 @@ + + import contextlib + import os +-from pathlib import Path + import signal + import subprocess + import sys + import tempfile + import time ++from pathlib import Path + + import pytest + + from vllm.v1.ple_offload.worker import PleOffloadWorkerHandle + +- + pytestmark = pytest.mark.skipif( + sys.platform != "linux", reason="requires Linux filesystem ipc sockets" + ) +diff --git a/tests/v1/ple_offload/test_progress_deadline.py b/tests/v1/ple_offload/test_progress_deadline.py +index 5e7ca19..ed825b4 100644 +--- a/tests/v1/ple_offload/test_progress_deadline.py ++++ b/tests/v1/ple_offload/test_progress_deadline.py +@@ -76,8 +76,9 @@ def test_progress_deadline_terminates_alive_silent_worker( + args=(block_input, started), + ) + process.start() +- # Cold vLLM import in the spawned child takes 4-8 s on the reference host; the deadline +- # itself (0.2 s) is only charged after `started` is set. ++ # Cold vLLM import in the spawned child takes 4-8 s on the reference ++ # host; the deadline itself (0.2 s) is only charged after `started` is ++ # set. + assert started.wait(timeout=120) + process.join(timeout=10) + assert not process.is_alive() +@@ -96,7 +97,6 @@ def _pending_completion_case(use_release_wrapper: bool, started: Any) -> None: + connector_module.torch.cuda.current_stream = lambda device=None: stream + connector_module.torch.cuda.Event = _PendingEvent + connector_module.torch.cuda.is_current_stream_capturing = lambda: False +- connector_module._rowdump.enabled = lambda: False + + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector.device = torch.device("cpu") +@@ -104,7 +104,6 @@ def _pending_completion_case(use_release_wrapper: bool, started: Any) -> None: + connector.tp_rank = 0 + connector._layers = {"ple": _FakeLayer()} + connector._progress = monitor +- connector._rowdump_n = None + connector._after_dummy = False + monitor.begin(startup=False, num_tokens=8) + connector.release_outputs() +@@ -175,7 +174,9 @@ def _monitor_exception_case( + query_entered: Any, + diagnostic_path: str, + ) -> None: +- diagnostic_fd = os.open(diagnostic_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) ++ diagnostic_fd = os.open( ++ diagnostic_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600 ++ ) + os.dup2(diagnostic_fd, 2) + os.close(diagnostic_fd) + monitor = PleProgressMonitor( +@@ -210,7 +211,4 @@ def test_monitor_thread_exception_is_fatal(tmp_path: Path) -> None: + assert process.exitcode not in (None, 0) + diagnostic = diagnostic_path.read_text(encoding="utf-8") + assert "where=progress-monitor-thread" in diagnostic +- assert ( +- "error=RuntimeError:injected event query failure" +- in diagnostic +- ) ++ assert "error=RuntimeError:injected event query failure" in diagnostic +diff --git a/tests/v1/ple_offload/test_rowdump_deleted.py b/tests/v1/ple_offload/test_rowdump_deleted.py +new file mode 100644 +index 0000000..97216e2 +--- /dev/null ++++ b/tests/v1/ple_offload/test_rowdump_deleted.py +@@ -0,0 +1,49 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The PLE row dump stays deleted on the 0.30 tree. ++ ++An earlier cleanup deleted ``vllm/v1/ple_offload/rowdump.py`` and its call ++sites: it was a byte-level push-vs-pull diff for a push transport that no ++longer exists, and nothing reads the ``b"PLER"`` records it writes. The pull-transport ++carry (an earlier commit) restored it verbatim from the E2 base, so the 0.30 ++release tree read ``VLLM_PLE_ROW_DUMP`` again -- an unregistered ``VLLM_`` name ++that validate_environ reports as unknown, on the per-step publish path. ++ ++Source-level on purpose: importing the connector needs cuda-python, which the ++GPU-free tier does not have, and the property is "no reader and no call site", ++which the source states directly. ++""" ++ ++from pathlib import Path ++ ++import pytest ++ ++_ROOT = Path(__file__).resolve().parents[3] ++_PKG = _ROOT / "vllm" / "v1" / "ple_offload" ++ ++ ++def _sources(): ++ for sub in ("vllm", "vllm_ple_offload_bootstrap", "vllm_ple_test_fault"): ++ base = _ROOT / sub ++ if base.is_dir(): ++ yield from base.rglob("*.py") ++ ++ ++def test_rowdump_module_is_gone(): ++ assert _PKG.is_dir(), f"PLE offload package missing at {_PKG}" ++ assert not (_PKG / "rowdump.py").exists() ++ ++ ++@pytest.mark.parametrize("name", ["connector.py", "worker.py"]) ++def test_transport_has_no_rowdump_call_site(name: str): ++ text = (_PKG / name).read_text(encoding="utf-8") ++ assert "rowdump" not in text, f"{name} still imports or calls the row dump" ++ ++ ++def test_nothing_reads_the_row_dump_env(): ++ readers = [ ++ str(p.relative_to(_ROOT)) ++ for p in _sources() ++ if "VLLM_PLE_ROW_DUMP" in p.read_text(encoding="utf-8", errors="replace") ++ ] ++ assert not readers, f"VLLM_PLE_ROW_DUMP still read by {readers}" +diff --git a/tests/v1/ple_offload/test_shm_unlink.py b/tests/v1/ple_offload/test_shm_unlink.py +index 600a7bb..bdb3a98 100644 +--- a/tests/v1/ple_offload/test_shm_unlink.py ++++ b/tests/v1/ple_offload/test_shm_unlink.py +@@ -1,10 +1,10 @@ + # SPDX-License-Identifier: Apache-2.0 + + import os +-from multiprocessing.reduction import ForkingPickler +-from pathlib import Path + import subprocess + import sys ++from multiprocessing.reduction import ForkingPickler ++from pathlib import Path + + import pytest + import torch +diff --git a/tests/v1/ple_offload/test_shutdown_order.py b/tests/v1/ple_offload/test_shutdown_order.py +index 101c3da..0f73a68 100644 +--- a/tests/v1/ple_offload/test_shutdown_order.py ++++ b/tests/v1/ple_offload/test_shutdown_order.py +@@ -4,8 +4,8 @@ import contextlib + from types import SimpleNamespace + from typing import Any + +-from vllm.v1.ple_offload.connector import PleOffloadConnector + from vllm.v1.ple_offload import worker as ple_worker ++from vllm.v1.ple_offload.connector import PleOffloadConnector + from vllm.v1.worker import gpu_worker + + +diff --git a/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py +index c709a93..1a0d21e 100644 +--- a/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py ++++ b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py +@@ -18,6 +18,7 @@ from pathlib import Path + + import pytest + ++ + def _native_cuda_probe(output_path: str) -> None: + """Run during config unpickling, before the bootstrap imports vLLM.""" + library = ctypes.util.find_library("cudart") +@@ -28,9 +29,7 @@ def _native_cuda_probe(output_path: str) -> None: + Path(output_path).write_text( + json.dumps( + { +- "cuda_visible_devices": os.environ.get( +- "CUDA_VISIBLE_DEVICES" +- ), ++ "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "cuda_result": result, + "torch_imported": "torch" in sys.modules, + "vllm_imported": "vllm" in sys.modules, +@@ -176,7 +175,7 @@ def test_launch_payload_send_has_absolute_deadline() -> None: + assert process.returncode == 0, stderr + # The injector logs the send entry through vLLM's logger on stdout; read only + # the probe's own result line. +- result = [line for line in stdout.splitlines() if line.startswith('TIMED_OUT ')] ++ result = [line for line in stdout.splitlines() if line.startswith("TIMED_OUT ")] + assert len(result) == 1, stdout + marker, elapsed = result[0].split() + assert marker == "TIMED_OUT" +@@ -282,9 +281,7 @@ class _ExecFdProbe: + + "," + + repr(self.output_path) + + ",str(all(not __import__('os').get_inheritable(int(fd)) " +- "for fd in __import__('sys').argv[1:]))," +- + repr(pairs) +- + "])" ++ "for fd in __import__('sys').argv[1:]))," + repr(pairs) + "])" + ) + return eval, (expression,) + +@@ -330,9 +327,7 @@ def test_actual_bootstrap_restores_cloexec_before_helper_exec( + control_child.close() + death_child.close() + try: +- _send_framed_payload( +- control.fileno(), payload, time.monotonic() + 10 +- ) ++ _send_framed_payload(control.fileno(), payload, time.monotonic() + 10) + assert process.wait(timeout=120) == 0 + finally: + control.close() +diff --git a/tests/v1/ple_offload/test_test_fault.py b/tests/v1/ple_offload/test_test_fault.py +index 3dc8470..a199207 100644 +--- a/tests/v1/ple_offload/test_test_fault.py ++++ b/tests/v1/ple_offload/test_test_fault.py +@@ -64,9 +64,7 @@ def test_malformed_specs_raise(fault_module, value): + + + def test_one_shot_rank_phase_and_ordinal(fault_module, monkeypatch): +- spec = fault_module._parse_fault_spec( +- "cut=exc_after_prepare,ordinal=2,tp=1,pp=0" +- ) ++ spec = fault_module._parse_fault_spec("cut=exc_after_prepare,ordinal=2,tp=1,pp=0") + assert spec is not None + + monkeypatch.setattr(fault_module.os, "getpid", lambda: 123) +@@ -153,7 +151,6 @@ def test_disarmed_release_site_continues(fault_module, monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) +- monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) + + stream = Mock() + completion = Mock() +@@ -210,25 +207,24 @@ def test_all_injection_sites_use_the_disarmed_short_circuit(): + sites = [i for i, line in enumerate(lines) if "_FAULT.fires(" in line] + assert len(sites) == count, relative + for i in sites: +- window = "\n".join(lines[max(0, i - 6): i + 1]) ++ window = "\n".join(lines[max(0, i - 6) : i + 1]) + # The None check must precede fires() in the same condition, so a + # disarmed process never calls into the fault module. + assert "_FAULT is not None" in window, (relative, i + 1) +- assert window.index("_FAULT is not None") < window.index("_FAULT.fires("), (relative, i + 1) ++ assert window.index("_FAULT is not None") < window.index("_FAULT.fires("), ( ++ relative, ++ i + 1, ++ ) + + + class _FatalReached(Exception): + pass + + +-def test_exc_in_release_reaches_real_connector_fatal_wrapper( +- fault_module, monkeypatch +-): ++def test_exc_in_release_reaches_real_connector_fatal_wrapper(fault_module, monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + +- armed = fault_module._parse_fault_spec( +- "cut=exc_in_release,ordinal=1,tp=1,pp=0" +- ) ++ armed = fault_module._parse_fault_spec("cut=exc_in_release,ordinal=1,tp=1,pp=0") + assert armed is not None + monkeypatch.setattr(connector_module, "_FAULT", armed) + +diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py +index 78a9db0..e42591a 100644 +--- a/tests/v1/ple_offload/test_transport_bounded_io.py ++++ b/tests/v1/ple_offload/test_transport_bounded_io.py +@@ -275,8 +275,9 @@ def test_real_wait_copy_geometry_rejected_before_graph_publication( + + # Go through the op's real entry point, which resolves + # the live pointers by layer name from the per-process registry. +- from types import SimpleNamespace +- ++ # (SimpleNamespace is the module-level import: a function-local import ++ # here made the name local to the whole function, so the earlier ++ # `success = SimpleNamespace(...)` raised UnboundLocalError -- ruff F823.) + monkeypatch.setitem( + ple_layer_module._PULL_LAYERS, + "ple", +diff --git a/tests/v1/ple_offload/test_transport_fatal_helper.py b/tests/v1/ple_offload/test_transport_fatal_helper.py +index 4e5d5c6..f1abd27 100644 +--- a/tests/v1/ple_offload/test_transport_fatal_helper.py ++++ b/tests/v1/ple_offload/test_transport_fatal_helper.py +@@ -27,9 +27,7 @@ def _run_bounded_child( + deadline = time.monotonic() + 120 + while not marker.exists() and time.monotonic() < deadline: + time.sleep(0.05) +- assert marker.exists(), ( +- "child did not finish cold startup within 120 seconds" +- ) ++ assert marker.exists(), "child did not finish cold startup within 120 seconds" + process.wait(timeout=5) + except subprocess.TimeoutExpired: + stayed_alive = True +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index 0f3ab58..d19ec53 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -465,7 +465,9 @@ def test_draft_guard_reports_the_draft_architecture() -> None: + assert e2_guards._profile(cfg, "target")[0] == ["Qwen4ExpForConditionalGeneration"] + assert e2_guards._profile(cfg, "draft")[0] == ["Qwen4ExpMTP"] + no_spec = SimpleNamespace(**{**vars(cfg), "speculative_config": None}) +- assert e2_guards._profile(no_spec, "draft")[0] == ["Qwen4ExpForConditionalGeneration"] ++ assert e2_guards._profile(no_spec, "draft")[0] == [ ++ "Qwen4ExpForConditionalGeneration" ++ ] + + + def _import_ple_offload_layer(monkeypatch): +@@ -560,13 +562,19 @@ def test_pull_registry_refuses_unbound_and_clashing_names(monkeypatch) -> None: + sem = object() + buf = torch.empty(2, 2) + with pytest.raises(RuntimeError, match="no _ple_offload_name"): +- armed("").setup_cross_process_offload(buf, sem, host_output_ptr=1, host_row_bytes=8) ++ armed("").setup_cross_process_offload( ++ buf, sem, host_output_ptr=1, host_row_bytes=8 ++ ) + a = armed("L1") + a.setup_cross_process_offload(buf, sem, host_output_ptr=1, host_row_bytes=8) +- a.setup_cross_process_offload(buf, sem, host_output_ptr=2, host_row_bytes=8) # re-arm ok ++ a.setup_cross_process_offload( ++ buf, sem, host_output_ptr=2, host_row_bytes=8 ++ ) # re-arm ok + assert m._PULL_LAYERS["L1"] is a + with pytest.raises(RuntimeError, match="two layers armed"): +- armed("L1").setup_cross_process_offload(buf, sem, host_output_ptr=3, host_row_bytes=8) ++ armed("L1").setup_cross_process_offload( ++ buf, sem, host_output_ptr=3, host_row_bytes=8 ++ ) + + + def test_nvidia_qwen4exp_ops_carry_no_pointer_like_int_params() -> None: +@@ -577,21 +585,28 @@ def test_nvidia_qwen4exp_ops_carry_no_pointer_like_int_params() -> None: + import pathlib + + root = pathlib.Path(__file__).resolve().parents[2] / "vllm" +- files = [root / "model_executor/layers/ple_offload_layer.py", +- *sorted((root / "models/qwen4_exp/nvidia").rglob("*.py")), +- *sorted((root / "v1/ple_offload").rglob("*.py"))] ++ files = [ ++ root / "model_executor/layers/ple_offload_layer.py", ++ *sorted((root / "models/qwen4_exp/nvidia").rglob("*.py")), ++ *sorted((root / "v1/ple_offload").rglob("*.py")), ++ ] + bad = [] + for f in files: + tree = ast.parse(f.read_text(encoding="utf-8")) + funcs = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + for node in ast.walk(tree): +- if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "direct_register_custom_op": ++ if ( ++ isinstance(node, ast.Call) ++ and getattr(node.func, "id", "") == "direct_register_custom_op" ++ ): + kw = {k.arg: k.value for k in node.keywords} + fn = getattr(kw.get("op_func"), "id", None) + if fn not in funcs: + continue + for a in funcs[fn].args.args: + ann = ast.unparse(a.annotation) if a.annotation is not None else "" +- if ann == "int" and any(t in a.arg for t in ("ptr", "addr", "pointer")): ++ if ann == "int" and any( ++ t in a.arg for t in ("ptr", "addr", "pointer") ++ ): + bad.append(f"{f.name}:{fn}({a.arg})") + assert not bad, f"pointer-like int op params (AOT-baking hazard): {bad}" +diff --git a/tests/v1/worker/test_e2_memstats.py b/tests/v1/worker/test_e2_memstats.py +index c405187..bf83e5b 100644 +--- a/tests/v1/worker/test_e2_memstats.py ++++ b/tests/v1/worker/test_e2_memstats.py +@@ -1,5 +1,7 @@ + # SPDX-License-Identifier: Apache-2.0 +-"""CPU/GPU-free schema and call-order tests for vllm/v1/worker/e2_memstats.py (debug-only telemetry).""" ++"""CPU/GPU-free schema and call-order tests for vllm/v1/worker/e2_memstats.py ++(debug-only telemetry).""" ++ + from __future__ import annotations + + import importlib.util +@@ -11,7 +13,9 @@ from unittest import mock + + import torch + +-MODULE_PATH = Path(__file__).resolve().parents[3] / "vllm" / "v1" / "worker" / "e2_memstats.py" ++MODULE_PATH = ( ++ Path(__file__).resolve().parents[3] / "vllm" / "v1" / "worker" / "e2_memstats.py" ++) + HOST_CURRENT_KEYS = { + "active_bytes.current", + "active_requests.current", +@@ -27,7 +31,9 @@ def load_disabled_module(): + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {MODULE_PATH}") + module = importlib.util.module_from_spec(spec) +- with mock.patch.dict(os.environ, {"VLLM_E2_MEMSTATS": "", "VLLM_E2_MEMSTATS_DIR": ""}): ++ with mock.patch.dict( ++ os.environ, {"VLLM_E2_MEMSTATS": "", "VLLM_E2_MEMSTATS_DIR": ""} ++ ): + spec.loader.exec_module(module) + return module + +@@ -79,15 +85,34 @@ class MemstatsSchemaTest(unittest.TestCase): + return host_stats + + with ( +- mock.patch.object(module.torch.cuda, "memory_stats", side_effect=memory_stats), +- mock.patch.object(module.torch.cuda, "reset_peak_memory_stats", side_effect=reset_peak_memory_stats), +- mock.patch.object(module.torch.cuda, "mem_get_info", side_effect=mem_get_info), +- mock.patch.object(module.torch.cuda, "host_memory_stats", side_effect=host_memory_stats), ++ mock.patch.object( ++ module.torch.cuda, "memory_stats", side_effect=memory_stats ++ ), ++ mock.patch.object( ++ module.torch.cuda, ++ "reset_peak_memory_stats", ++ side_effect=reset_peak_memory_stats, ++ ), ++ mock.patch.object( ++ module.torch.cuda, "mem_get_info", side_effect=mem_get_info ++ ), ++ mock.patch.object( ++ module.torch.cuda, "host_memory_stats", side_effect=host_memory_stats ++ ), + ): + record = module._close_metrics(7) + +- # The device read and the peak reset stay adjacent; the host read comes after both. +- self.assertEqual(calls, ["memory_stats", "reset_peak_memory_stats", "mem_get_info", "host_memory_stats"]) ++ # The device read and the peak reset stay adjacent; the host read ++ # comes after both. ++ self.assertEqual( ++ calls, ++ [ ++ "memory_stats", ++ "reset_peak_memory_stats", ++ "mem_get_info", ++ "host_memory_stats", ++ ], ++ ) + self.assertEqual( + set(record), + { +diff --git a/vllm/envs.py b/vllm/envs.py +index b604b39..0dca963 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -339,8 +339,10 @@ if TYPE_CHECKING: + # catch-all branch in _get_ple_embedding_quant_method, which lived in the + # half of models/qwen4_exp/nvidia/ple_layer.py that #54371 deleted. With no + # reader it cannot change an artifact, so hashing it would only invalidate +- # caches for nothing; and on this uncompiled NVIDIA path (#55272) there may +- # be no AOT artifact to key at all. It gives NO capture protection either. ++ # caches for nothing. (The AOT artifact is real on this tree: #55272 ++ # removed torch.compile from the NVIDIA Qwen4Exp path, the compile-profile guard restored it, ++ # and the vLLM 0.30 rebase ran AOT cache-load boots -- so the reason is "no ++ # reader", not "no artifact".) It gives NO capture protection either. + # v0.30.0 selects the FP8 PLE method from the CONFIG KEY + # text_config.ple_embedding_dtype == "float8_e4m3fn" + # (ngram_embedding.py:174-175) -- see +@@ -428,14 +430,12 @@ if TYPE_CHECKING: + # cross-change knob table). Each mirrors its read site's + # parse; if a read site changes, the entry here must be reconciled. + # (Both read sites below are now live.) +- # *** INERT ON THE 0.30 BASE: NO READ SITE. *** +- # Its reader is the slot-ownership assertion code (an earlier commit), which edits +- # vllm/v1/ple_offload/connector.py -- a module that does not exist on +- # v0.30.0. Setting this on the 0.30 base does NOTHING, and NOTHING WARNS +- # (validate_environ only reports UNregistered names). Already an ignored +- # factor, correctly. Kept registered so the contract is recorded; it +- # becomes live again only together with the pull transport +- # (the PLE transport choice). ++ # LIVE: read once at ++ # import as envs.VLLM_E2_PLE_FLAG_CHECK by vllm/v1/ple_offload/connector.py ++ # (the PLE pull-transport flag check, rank-gated, off by default). On ++ # the 0.30 base it had NO read site: the PLE-diagnostics half ++ # had not been relanded after the PLE pull-transport port restored connector.py. The ++ # census in tests/test_envs.py now asserts the read site exists. + # AOT: no. CAPTURE: no -- strictly host-side and outside capture. + # off|warn|raise (1 = warn), normalised by _e2_diagnostic_mode like the + # rest of the family. The mamba ownership check moved this gate off a plain +@@ -487,6 +487,9 @@ if TYPE_CHECKING: + VLLM_NIC_SELECTION_VARS: str = "" + VLLM_ENABLE_HPC_OPS: bool = False + MERLIN_TOPK_EXACT_COUNTERS_LOG: int = 0 ++ # Pre-tie-determinism fallback for the counter reader's slot count; runtime-only ++ # (ignored factor), never consulted on this build. See the dict entry. ++ MERLIN_TOPK_COUNTER_SLOTS: str | None = None + # + # --- determinism toggles ----------------------------------------- + # The three nondeterminism mechanisms (FULL_K, QSA_SORT, TIE_DET/TIE_RECENT). +@@ -556,11 +559,13 @@ if TYPE_CHECKING: + # The same restart on an overflowing DECODE row. + MERLIN_TOPK_EXACT_DECODE: str | None = None + # Deterministic placement and tie membership on the LARGE path. +- # Registered here ahead of its code, exactly as the four toggles above +- # were, so the full tree is complete; its read site is +- # det_large_enabled() in topk.cu on the large path. Setting it without that +- # code does nothing, which is why it is a passthrough and not a parsed +- # bool that would imply this tree can act on it. ++ # LIVE on this tree since the v2.2 kernel port: the ++ # read site is det_large_enabled() in csrc/libtorch_stable/topk.cu, a ++ # std::getenv resolved once per process, and the resolved value reaches ++ # the large-path radix_topk as a launch argument, so a cudagraph bakes it ++ # in. CAPTURE: yes. AOT: not established (built into the extension, not by ++ # the AOT compiler); hashed conservatively like the switches above. A raw ++ # passthrough because the C++ getter is the single source of truth. + MERLIN_TOPK_DET_LARGE: str | None = None + # + # The large-path FAULT INJECTORS. TEST +@@ -606,8 +611,8 @@ if TYPE_CHECKING: + # without first deleting it from the fence -- which means writing down its + # AOT and capture identity against its real read site -- is the exact + # failure the fork's env rule exists to prevent, and +- # tests/test_envs.py::test_unaudited_host_side_knob_stays_listed_until_it_is_classified +- # fails on it. ++ # test_unaudited_host_side_knob_stays_listed_until_it_is_classified (in ++ # tests/test_envs.py) fails on it. + # + # The fence holds NOTHING BUT the names, one per line. Prose belongs out + # here: the gate reads the fence as a set of identifiers, and a knob name +@@ -2582,6 +2587,11 @@ environment_variables: dict[str, Callable[[], Any]] = { + "MERLIN_TOPK_EXACT_COUNTERS_LOG": lambda: int( + os.getenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "0") + ), ++ # Fallback slot count for the counter reader on a build WITHOUT the ++ # persistent_topk_exact_counters_slots op (pre-tie-determinism). This tree registers ++ # that op, so topk_counters.slot_count() never consults the variable here. ++ # Raw passthrough: the reader's int(os.getenv(..., "5")) is the parser. ++ "MERLIN_TOPK_COUNTER_SLOTS": lambda: os.getenv("MERLIN_TOPK_COUNTER_SLOTS", None), + # Determinism toggles. See the declaration block above for the AOT / + # capture classification and for why three of these are raw passthroughs. + # All four are absent from ignored_factors on purpose, so all four are +@@ -2831,6 +2841,9 @@ def compile_factors() -> dict[str, object]: + # host-side timer and a snapshot op, and is read by no compiled code, so + # hashing it would invalidate every compile cache for a log interval. + "MERLIN_TOPK_EXACT_COUNTERS_LOG", ++ # Host-side reader fallback (buffer length for a snapshot op outside ++ # capture); dead on a build that answers the _slots query. ++ "MERLIN_TOPK_COUNTER_SLOTS", + # ------------------------------------------------------------------ + } + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +index 555e445..8e5da35 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +@@ -52,7 +52,7 @@ _QSA_SORT = envs.MERLIN_QSA_SORT + + + def _qsa_canonicalise(block_indices: torch.Tensor) -> None: +- """Sort selected block indices ascending, padding last. In place, shape preserved.""" ++ """Sort selected block indices ascending, padding last. In place, same shape.""" + if not _QSA_SORT: + return + big = torch.iinfo(block_indices.dtype).max +diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py +index 0d52651..2cf05bb 100644 +--- a/vllm/v1/ple_offload/connector.py ++++ b/vllm/v1/ple_offload/connector.py +@@ -11,9 +11,6 @@ from typing import Any + + import msgspec + import torch +-from vllm.v1.ple_offload import rowdump as _rowdump +-from vllm.v1.ple_offload.fatal import ple_fail_closed +-from vllm.v1.ple_offload.progress import PleProgressMonitor + import torch.nn as nn + import zmq + from cuda.bindings import driver as cuda_driver +@@ -28,6 +25,8 @@ from vllm.model_executor.layers.ple_offload_layer import ( + flag_check_is_authoritative, + stale_offload_flags, + ) ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++from vllm.v1.ple_offload.progress import PleProgressMonitor + from vllm.v1.ple_offload.protocol import ( + PleOffloadRegistration, + PleOffloadRequest, +@@ -284,7 +283,8 @@ class PleOffloadConnector: + } + if not layers: + raise RuntimeError( +- "VLLM_E2_PLE_PULL_TRANSPORT is enabled, but the model has no PleOffloadLayer" ++ "VLLM_E2_PLE_PULL_TRANSPORT is enabled, but the model has no " ++ "PleOffloadLayer" + ) + + config = vllm_config.model_config.hf_text_config +@@ -711,7 +711,6 @@ class PleOffloadConnector: + f"PLE request count {num_reqs} exceeds connector capacity " + f"{self._query_start_loc_buf.shape[0] - 1}" + ) +- self._rowdump_n = None if dummy_run else num_tokens + if dummy_run: + self._signal_dummy_outputs_impl(num_tokens) + return +@@ -913,15 +912,6 @@ class PleOffloadConnector: + correlation_id=correlation_id, + ): + raise RuntimeError("PLE test fault: exc_in_release") +- n = getattr(self, "_rowdump_n", None) +- self._rowdump_n = None +- if _rowdump.enabled() and n is not None and not torch.cuda.is_current_stream_capturing(): +- torch.cuda.current_stream(self.device).synchronize() +- seq = getattr(self, "_rowdump_seq", 0) +- self._rowdump_seq = seq + 1 +- for layer in self._layers.values(): +- _rowdump.write(f"gpu-tp{self.tp_rank}", 0, self.tp_rank, seq, +- self._input_ids_source[:n], layer._gpu_output_buffer[:n]) + # Reset only after the consumer forward so the CPU worker cannot + # overwrite an output that a GPU PLE placeholder may still read. + stream = torch.cuda.current_stream(self.device) +diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py +index a80cad4..90ab4a0 100644 +--- a/vllm/v1/ple_offload/fatal.py ++++ b/vllm/v1/ple_offload/fatal.py +@@ -54,11 +54,7 @@ def ple_fail_closed( + detail_suffix = "" + if trusted_error: + args = BaseException.args.__get__(error, error_type) +- if ( +- type(args) is tuple +- and len(args) == 1 +- and type(args[0]) is str +- ): ++ if type(args) is tuple and len(args) == 1 and type(args[0]) is str: + detail = args[0][:512].replace("\n", "\\n") + detail_suffix = f":{detail}" + message = ( +@@ -105,7 +101,8 @@ def ple_fail_closed( + raw = raw[-8192:] + os.write( + 2, +- b"PLE fatal traceback (tail) follows:\n" if clipped ++ b"PLE fatal traceback (tail) follows:\n" ++ if clipped + else b"PLE fatal traceback follows:\n", + ) + os.write(2, raw) +@@ -118,9 +115,13 @@ def ple_fail_closed( + # Still guarded, because os._exit below must run whatever happens: an + # escape from this handler would leave the worker alive in the state + # the fail-closed path exists to end. +- try: ++ try: # noqa: SIM105 + os.write(2, b"PLE fatal: traceback unavailable\n") + except BaseException: ++ # SIM105 declined: an earlier revision made this helper call no ++ # exception-defined code on the dying path. contextlib.suppress ++ # runs a Python __exit__ between here and os._exit below; a bare ++ # except keeps that window empty. + pass + os._exit(1) + +@@ -143,9 +144,7 @@ def fail_ple_execute_model(func: _F) -> _F: + connector = getattr(self, "_ple_offload_connector", None) + if connector is None: + model_runner = getattr(self, "model_runner", None) +- connector = getattr( +- model_runner, "_ple_offload_connector", None +- ) ++ connector = getattr(model_runner, "_ple_offload_connector", None) + if connector is None: + raise + connector._fail_closed(func.__name__, error) +diff --git a/vllm/v1/ple_offload/progress.py b/vllm/v1/ple_offload/progress.py +index bf88291..9ea0406 100644 +--- a/vllm/v1/ple_offload/progress.py ++++ b/vllm/v1/ple_offload/progress.py +@@ -20,11 +20,7 @@ _POLL_INTERVAL_S = 0.1 + + def _validate_deadline(value: float, name: str) -> float: + value = float(value) +- if ( +- not math.isfinite(value) +- or value <= 0 +- or value > _MAX_PROGRESS_DEADLINE_S +- ): ++ if not math.isfinite(value) or value <= 0 or value > _MAX_PROGRESS_DEADLINE_S: + raise ValueError( + f"{name} must be finite, greater than zero, and no greater than " + f"{_MAX_PROGRESS_DEADLINE_S:g} seconds" +@@ -123,10 +119,7 @@ class PleProgressMonitor: + active = self._submission + if active is None or active.sequence != sequence: + raise RuntimeError("PLE fault correlation has no active submission") +- return ( +- f"sequence:{active.sequence}" +- f"@begin:{active.started:.9f}" +- ) ++ return f"sequence:{active.sequence}@begin:{active.started:.9f}" + + def completion_is_pending(self, sequence: int) -> bool: + """Whether this serving sequence has a retained, incomplete event. +@@ -149,9 +142,7 @@ class PleProgressMonitor: + + def begin(self, *, startup: bool, num_tokens: int) -> int: + now = time.monotonic() +- budget = ( +- self._startup_deadline_s if startup else self._serving_deadline_s +- ) ++ budget = self._startup_deadline_s if startup else self._serving_deadline_s + with self._condition: + active = self._submission + if active is not None: +diff --git a/vllm/v1/ple_offload/rowdump.py b/vllm/v1/ple_offload/rowdump.py +deleted file mode 100644 +index 95fd308..0000000 +--- a/vllm/v1/ple_offload/rowdump.py ++++ /dev/null +@@ -1,41 +0,0 @@ +-# SPDX-License-Identifier: Apache-2.0 +-"""Debug-only PLE row dump for byte-level transport checks (off unless VLLM_PLE_ROW_DUMP=

). +- +-Each process appends framed records to /-pid.bin: +- b"PLER" | kind u8 (0 = GPU worker: rows its forward consumed; 1 = offload process: rows it published) +- | rank u8 | reserved u16 | seq u32 | n u32 | row_bytes u32 | ids int64[n] | rows uint8[n * row_bytes] +-""" +- +-import os +-import struct +-import threading +- +-import torch +- +-_DIR = os.environ.get("VLLM_PLE_ROW_DUMP", "") +-_LOCK = threading.Lock() +-_FILES: dict = {} +-MAGIC = b"PLER" +- +- +-def enabled() -> bool: +- return bool(_DIR) +- +- +-def write(tag: str, kind: int, rank: int, seq: int, ids: torch.Tensor, rows: torch.Tensor) -> None: +- ids_cpu = ids.detach().to(device="cpu", dtype=torch.int64).contiguous() +- rows_cpu = rows.detach().to(device="cpu").contiguous() +- n = int(ids_cpu.shape[0]) +- if rows_cpu.shape[0] != n: +- raise RuntimeError(f"PLE row dump: {rows_cpu.shape[0]} rows for {n} ids") +- raw = rows_cpu.view(torch.uint8).numpy().tobytes() if n else b"" +- row_bytes = len(raw) // n if n else 0 +- record = (MAGIC + struct.pack(" TERM path. + self.proc.terminate() +- try: ++ with contextlib.suppress(subprocess.TimeoutExpired): + self.proc.wait(timeout=5) +- except subprocess.TimeoutExpired: +- pass + if self.proc.poll() is None: + self.proc.kill() +- try: ++ with contextlib.suppress(subprocess.TimeoutExpired): + self.proc.wait(timeout=5) +- except subprocess.TimeoutExpired: +- pass + # The child normally removes the endpoint in proc_main's finally. + # Once it has exited, this exact PLE-owned address is safe to remove as + # a fallback for SIGKILL or an interrupted child cleanup. +@@ -1234,12 +1227,6 @@ class PleOffloadRunner: + ngram_context, + output_buffer=self._scratch_bufs[dp_rank][layer_name], + ) +- if _rowdump.enabled(): +- _seq = getattr(self, "_rowdump_seq", 0) +- self._rowdump_seq = _seq + 1 +- _rowdump.write(f"offload-dp{dp_rank}", 1, dp_rank, _seq, +- input_bufs.input_ids_buf[: request.num_tokens], result) +- + # The result is identical on every TP rank in this DP group. + # A worker's host buffer may be overwritten only after its GPU + # released the previous result (``release_outputs`` stores +diff --git a/vllm/v1/worker/e2_memstats.py b/vllm/v1/worker/e2_memstats.py +index 586f8fa..9fcbfcf 100644 +--- a/vllm/v1/worker/e2_memstats.py ++++ b/vllm/v1/worker/e2_memstats.py +@@ -1,40 +1,57 @@ + # SPDX-License-Identifier: Apache-2.0 +-"""Debug-only allocator telemetry for E2, BOUNDARY-ONLY. Off unless VLLM_E2_MEMSTATS= and +-VLLM_E2_MEMSTATS_DIR=. +- +-WHY THIS IS NOT A SAMPLER ANY MORE. An earlier version read torch.cuda.memory_stats() and +-reset_peak_memory_stats() every VLLM_E2_MEMSTATS seconds. Both take the per-device caching-allocator mutex and the +-Python bindings do not release the GIL, so the old contract's claim that this thread "takes no lock shared with the +-model path" was FALSE. Three one-boot runs at one KV pin showed a monotone dose-response: sampling at 4 Hz wedged a +-stage-0 worker during a 131072-token row; at 1.1 Hz it survived that row and wedged during a 260000-token row; with the +-thread disabled the run completed every row. Causation is not established at n=1 per run +-(single runs only), but polling has no place inside a measured region, and it was +-never necessary: PyTorch maintains the allocator high-water marks itself, so reading them once per window is exact. +- +-WHAT THIS THREAD DOES NOW. Between window boundaries it makes NO CUDA calls and touches NO allocator state: its tick +-is a plain read of /window.ctl, which holds " ". The gate driver crosses a boundary only while the +-server is quiet (no request running). On a new seq the thread performs these boundary-only allocator reads: +- 1. memory_stats() -> the closing window's exact high-water marks and cumulative retry/OOM counters +- 2. reset_peak_memory_stats() -> immediately after the read, with no file I/O in the gap ++"""Debug-only allocator telemetry for E2, BOUNDARY-ONLY. Off unless ++VLLM_E2_MEMSTATS= and VLLM_E2_MEMSTATS_DIR=. ++ ++WHY THIS IS NOT A SAMPLER ANY MORE. An earlier version read ++torch.cuda.memory_stats() and reset_peak_memory_stats() every VLLM_E2_MEMSTATS ++seconds. Both take the per-device caching-allocator mutex and the Python ++bindings do not release the GIL, so the old contract's claim that this thread ++"takes no lock shared with the model path" was FALSE. Three one-boot arms at a ++4.1e9 KV pin showed a monotone dose-response: sampling at 4 Hz wedged a stage-0 ++worker during a 131K thinking-on row; at 1.1 Hz it survived that row and wedged during ++a 260K row; with the thread disabled the run completed every row. Causation ++is not established at n=1 per arm (a known limit of ++that capacity probe), but polling has no place ++inside a measured region, and it was never necessary: PyTorch maintains the ++allocator high-water marks itself, so reading them once per window is exact. ++ ++WHAT THIS THREAD DOES NOW. Between window boundaries it makes NO CUDA calls and ++touches NO allocator state: its tick is a plain read of /window.ctl, which ++holds " ". The gate driver crosses a boundary only while the server ++is quiet (no request running). On a new seq the thread performs these ++boundary-only allocator reads: ++ 1. memory_stats() -> the closing window's exact high-water marks ++ and cumulative retry/OOM counters ++ 2. reset_peak_memory_stats() -> immediately after the read, with no file I/O ++ in the gap + 3. mem_get_info() -> device free/total AT THAT QUIESCENT INSTANT +- 4. host_memory_stats() -> pinned-host allocator current values, outside the read-to-reset gap +-then writes the "close" record for the old window and the "ack" record for the new one. ++ 4. host_memory_stats() -> pinned-host allocator current values, outside ++ the read-to-reset gap ++then writes the "close" record for the old window and the "ack" record for the ++new one. + + WHAT THIS MEASURES. +- - A row's allocator peak is EXACT for that row's window (not a sampled maximum; it cannot miss a transient). +- - Device free/total is observed ONLY at quiescent boundaries. No minimum-free figure exists for the inside of a +- row, and none may be reported. +- - Retry and OOM counters are cumulative and never reset; per-row deltas come from consecutive close records. +- - There is NO interval-level peak. A qualifying interval is narrower than its row and is identified only +- afterwards from the scheduler log, so the peak attributable to it is the ROW-level high-water mark, which +- over-includes the prefill ramp. This loss of resolution is inherent to the design. +- +-Records (one JSON per line): kind start | close | ack | error; rank, local, dev, pid; window, seq; and for close: +-allocated_bytes.all.current, reserved_bytes.all.current, alloc_peak_w, res_peak_w, retries, ooms, free, total, +-run_alloc_peak, run_res_peak, and a host object containing every host_memory_stats() key ending in ".current". +-Records are written with synchronous file I/O, never inside the read-to-reset gap. Nothing +-here raises into the worker. ++ - A row's allocator peak is EXACT for that row's window (not a sampled ++ maximum; it cannot miss a transient). ++ - Device free/total is observed ONLY at quiescent boundaries. No minimum-free ++ figure exists for the inside of a row, and none may be reported. ++ - Retry and OOM counters are cumulative and never reset; per-row deltas come ++ from consecutive close records. ++ - There is NO interval-level peak. A qualifying interval is narrower ++ than its row and is identified only afterwards from the scheduler log, so ++ the peak attributable to it is the ROW-level high-water mark, which ++ over-includes the prefill ramp. Report this loss of resolution ++ explicitly. ++ ++Records (one JSON per line): kind start | close | ack | error; rank, local, ++dev, pid; window, seq; and for close: allocated_bytes.all.current, ++reserved_bytes.all.current, alloc_peak_w, res_peak_w, retries, ooms, free, ++total, run_alloc_peak, run_res_peak, and a host object containing every ++host_memory_stats() key ending in ".current". Records are written with ++synchronous file I/O, never inside the read-to-reset gap. Nothing here raises ++into the worker. + """ ++ + import json + import os + import threading +@@ -57,17 +74,11 @@ def _close_metrics(dev): + free, total = torch.cuda.mem_get_info(dev) + host_stats = torch.cuda.host_memory_stats() + host = { +- key: int(value) +- for key, value in host_stats.items() +- if key.endswith(".current") ++ key: int(value) for key, value in host_stats.items() if key.endswith(".current") + } + return { +- "allocated_bytes.all.current": int( +- s.get("allocated_bytes.all.current", 0) +- ), +- "reserved_bytes.all.current": int( +- s.get("reserved_bytes.all.current", 0) +- ), ++ "allocated_bytes.all.current": int(s.get("allocated_bytes.all.current", 0)), ++ "reserved_bytes.all.current": int(s.get("reserved_bytes.all.current", 0)), + "alloc_peak_w": ap, + "res_peak_w": rp, + "retries": int(s.get("num_alloc_retries", 0)), +@@ -85,9 +96,20 @@ def enabled() -> bool: + def ensure_started(worker, armed_by: str = "compile_or_warm_up_model") -> None: + if _started: + return +- dev = torch.cuda.current_device() # captured on the worker's own thread: every call below passes it explicitly +- t = threading.Thread(target=_loop, name="e2-memstats", daemon=True, +- args=(dev, getattr(worker, "rank", -1), getattr(worker, "local_rank", -1), armed_by)) ++ dev = ( ++ torch.cuda.current_device() ++ ) # captured on the worker's own thread: every call below passes it explicitly ++ t = threading.Thread( ++ target=_loop, ++ name="e2-memstats", ++ daemon=True, ++ args=( ++ dev, ++ getattr(worker, "rank", -1), ++ getattr(worker, "local_rank", -1), ++ armed_by, ++ ), ++ ) + _started.append(t) + t.start() + +@@ -103,7 +125,13 @@ def _ctl(): + + def _loop(dev, rank, local, armed_by): + os.makedirs(_DIR, exist_ok=True) +- f = open(os.path.join(_DIR, f"memstats-rank{rank}-pid{os.getpid()}.jsonl"), "a", buffering=1) ++ # noqa SIM115: the handle lives as long as the telemetry daemon thread ++ # that writes to it; there is no scope to close it in. ++ f = open( # noqa: SIM115 ++ os.path.join(_DIR, f"memstats-rank{rank}-pid{os.getpid()}.jsonl"), ++ "a", ++ buffering=1, ++ ) + base = {"rank": rank, "local": local, "dev": dev, "pid": os.getpid()} + + def w(rec): +@@ -112,25 +140,61 @@ def _loop(dev, rank, local, armed_by): + run_ap = run_rp = 0 + seq, name = _ctl() or (0, "serve") + torch.cuda.reset_peak_memory_stats(dev) +- w({"kind": "start", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "armed_by": armed_by, +- "window": name, "seq": seq, "version": VERSION, "poll_s": _SEG}) ++ w( ++ { ++ "kind": "start", ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "armed_by": armed_by, ++ "window": name, ++ "seq": seq, ++ "version": VERSION, ++ "poll_s": _SEG, ++ } ++ ) + while True: + time.sleep(_SEG) + try: +- c = _ctl() # file read only: no CUDA, no allocator lock, no GIL-held call ++ c = _ctl() # file read only: no CUDA, no allocator lock, no GIL-held call + if c is None or c[0] == seq: + continue +- # Boundary. The driver only crosses one while the server is quiet, so this is the single point at which +- # this thread touches the allocator at all. ++ # Boundary. The driver only crosses one while the server is ++ # quiet, so this is the single point at which this thread ++ # touches the allocator at all. + metrics = _close_metrics(dev) + ap = metrics["alloc_peak_w"] + rp = metrics["res_peak_w"] + run_ap, run_rp = max(run_ap, ap), max(run_rp, rp) +- w({"kind": "close", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, +- "seq": seq, **metrics, +- "run_alloc_peak": run_ap, "run_res_peak": run_rp}) ++ w( ++ { ++ "kind": "close", ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "window": name, ++ "seq": seq, ++ **metrics, ++ "run_alloc_peak": run_ap, ++ "run_res_peak": run_rp, ++ } ++ ) + seq, name = c +- w({"kind": "ack", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, "seq": seq}) ++ w( ++ { ++ "kind": "ack", ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "window": name, ++ "seq": seq, ++ } ++ ) + except Exception as e: # noqa: BLE001 - telemetry must never raise into the worker +- w({"kind": "error", "t": round(time.time(), 4), "m": round(time.monotonic(), 4), "window": name, +- "seq": seq, "error": f"{type(e).__name__}: {e}"[:300]}) ++ w( ++ { ++ "kind": "error", ++ "t": round(time.time(), 4), ++ "m": round(time.monotonic(), 4), ++ "window": name, ++ "seq": seq, ++ "error": f"{type(e).__name__}: {e}"[:300], ++ } ++ ) +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 0e132b6..91d4d92 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -481,7 +481,10 @@ class CudaGraphManager: + if self._capture_mem_samples is not None: + torch.accelerator.synchronize() + free_before = torch.accelerator.get_memory_info()[0] +- with self.capture_ple_copy_geometry(desc): ++ # noqa SIM117 declined: merging this with the inner ++ # torch.cuda.graph block would re-indent upstream ++ # lines this fork has no business reformatting. ++ with self.capture_ple_copy_geometry(desc): # noqa: SIM117 + with torch.cuda.graph( + graph, self.pool, stream=self._capture_stream(desc) + ): +diff --git a/vllm_ple_offload_bootstrap/__init__.py b/vllm_ple_offload_bootstrap/__init__.py +index 60ddda4..73f017c 100644 +--- a/vllm_ple_offload_bootstrap/__init__.py ++++ b/vllm_ple_offload_bootstrap/__init__.py +@@ -17,6 +17,8 @@ _FAULT = None + if os.environ.get("VLLM_PLE_TEST_FAULT") is not None: + from vllm_ple_test_fault import ( + _FAULT as _FAULT, ++ ) ++ from vllm_ple_test_fault import ( + _block_on_harness_gate, + ) + +@@ -160,9 +162,11 @@ def main() -> None: + death_pipe = Connection(death_fd) + + def monitor_parent() -> None: +- try: ++ try: # noqa: SIM105 + death_pipe.recv() + except (EOFError, OSError): ++ # SIM105 declined: this module's contract is the docstring above -- ++ # its import list is the audited surface of the pre-torch child. + pass + os._exit(1) + +-- +2.43.0 + diff --git a/patches/0053-e2_guards-draft-identity-and-draft-PLE-guard-report-.patch b/patches/0053-e2_guards-draft-identity-and-draft-PLE-guard-report-.patch new file mode 100644 index 0000000..1f9291c --- /dev/null +++ b/patches/0053-e2_guards-draft-identity-and-draft-PLE-guard-report-.patch @@ -0,0 +1,554 @@ +From e5f6387064eadbeaba0ae923f0141ef9f638b9af Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 11:11:01 -0400 +Subject: [PATCH 053/103] e2_guards: draft identity and draft PLE guard report + the drafter's own state + +--- + tests/v1/test_ple_pull_compile_profile.py | 234 ++++++++++++++++++++++ + vllm/v1/e2_guards.py | 37 +++- + vllm/v1/worker/e2_identity.py | 150 ++++++++++++-- + vllm/v1/worker/gpu_worker.py | 11 +- + 4 files changed, 411 insertions(+), 21 deletions(-) + +diff --git a/tests/v1/test_ple_pull_compile_profile.py b/tests/v1/test_ple_pull_compile_profile.py +index d19ec53..97f4e60 100644 +--- a/tests/v1/test_ple_pull_compile_profile.py ++++ b/tests/v1/test_ple_pull_compile_profile.py +@@ -11,6 +11,7 @@ reading is shown to accept what the fork's guard rejects. + + import ast + import os ++import re + from pathlib import Path + from types import SimpleNamespace + +@@ -610,3 +611,236 @@ def test_nvidia_qwen4exp_ops_carry_no_pointer_like_int_params() -> None: + ): + bad.append(f"{f.name}:{fn}({a.arg})") + assert not bad, f"pointer-like int op params (AOT-baking hazard): {bad}" ++ ++ ++# ------------------------------------------ the e2-identity draft line ++ ++#: The startup smoke's anchored parser for the identity line. ++_IDENTITY_LINE = re.compile( ++ r"e2-identity\[(?PPP\d+_TP\d+)\]: ple_backend=(?P\S+) " ++ r"ple_method=(?P\S+) ple_weight_dtype=(?P\S+) " ++ r"breakable=(?PON|OFF) " ++ r"compile_mode=(?P\S+) cudagraph_mode=(?P\S+) " ++ r"role=(?Ptarget|draft)\s*$" ++) ++ ++ ++def _compiled_draft(live: bool): ++ """A draft model whose only compile wrapper is live (or not).""" ++ import torch.nn as nn ++ ++ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper ++ ++ class _Wrapped(nn.Module, TorchCompileWithNoGuardsWrapper): ++ def __init__(self) -> None: ++ nn.Module.__init__(self) ++ ++ inner = _Wrapped() ++ inner.do_not_compile = not live ++ inner.compilation_config = SimpleNamespace(mode=CompilationMode.VLLM_COMPILE) ++ outer = nn.Module() ++ outer.model = inner ++ return outer ++ ++ ++def _shared_config(): ++ """The ONE engine CompilationConfig, as the drafter's holder carries it.""" ++ return SimpleNamespace( ++ compilation_config=SimpleNamespace( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ ) ++ ) ++ ++ ++def _resolved_speculator(model, prefill, decode): ++ """V2 autoregressive / MTP speculator after init_cudagraph_manager.""" ++ return SimpleNamespace( ++ model=model, ++ vllm_config=_shared_config(), ++ prefill_cudagraph_manager=( ++ None if prefill is None else SimpleNamespace(cudagraph_mode=prefill) ++ ), ++ decode_cudagraph_manager=( ++ None if decode is None else SimpleNamespace(cudagraph_mode=decode) ++ ), ++ ) ++ ++ ++def _logged_identity(monkeypatch, fn, *args) -> list[str]: ++ from vllm.v1.worker import e2_identity ++ ++ lines: list[str] = [] ++ monkeypatch.setattr(e2_identity, "_rank_tag", lambda: "PP1_TP0") ++ monkeypatch.setattr( ++ e2_identity.logger, "info", lambda fmt, *a: lines.append(fmt % a) ++ ) ++ monkeypatch.setattr( ++ e2_identity.logger, "warning", lambda fmt, *a: lines.append("WARN " + fmt % a) ++ ) ++ fn(*args) ++ return lines ++ ++ ++def test_draft_identity_reports_the_drafters_resolved_graph_state(monkeypatch) -> None: ++ """The draft line read compile/cudagraph from the target's ++ shared VllmConfig (FULL_AND_PIECEWISE), while the V2 MTP speculator runs ++ prefill FULL_AND_PIECEWISE + decode FULL_DECODE_ONLY (the modes ++ e2-guard[runtime] logs). It must report the managers and the live wrapper, ++ in the smoke-anchored format.""" ++ from vllm.v1.worker import e2_identity ++ ++ runner = SimpleNamespace( ++ speculator=_resolved_speculator( ++ _compiled_draft(live=True), ++ CUDAGraphMode.FULL_AND_PIECEWISE, ++ CUDAGraphMode.FULL_DECODE_ONLY, ++ ) ++ ) ++ lines = _logged_identity(monkeypatch, e2_identity.log_draft_identity, runner) ++ assert len(lines) == 1, lines ++ m = _IDENTITY_LINE.search(lines[0]) ++ assert m, lines[0] ++ assert m["role"] == "draft" and m["tag"] == "PP1_TP0" ++ assert m["cmode"] == "VLLM_COMPILE" ++ assert m["gmode"] == "prefill:FULL_AND_PIECEWISE,decode:FULL_DECODE_ONLY" ++ assert m["backend"] == m["method"] == m["dtype"] == "none" ++ ++ ++def test_draft_identity_compile_mode_is_the_live_wrapper(monkeypatch) -> None: ++ """Same shared config (VLLM_COMPILE); only the draft wrapper's liveness ++ differs -- the line must follow the wrapper, not the config.""" ++ from vllm.v1.worker import e2_identity ++ ++ runner = SimpleNamespace( ++ speculator=_resolved_speculator( ++ _compiled_draft(live=False), ++ CUDAGraphMode.FULL_AND_PIECEWISE, ++ CUDAGraphMode.FULL_DECODE_ONLY, ++ ) ++ ) ++ (line,) = _logged_identity(monkeypatch, e2_identity.log_draft_identity, runner) ++ m = _IDENTITY_LINE.search(line) ++ assert m and m["cmode"] == "NONE", line ++ ++ ++@pytest.mark.parametrize("shape", ["single", "unbuilt", "v1"]) ++def test_draft_identity_other_speculator_shapes(monkeypatch, shape) -> None: ++ """Single-manager speculator -> bare mode; manager not built -> unresolved; ++ V1 drafter -> its dispatcher's mode. All still parse.""" ++ from vllm.v1.worker import e2_identity ++ ++ if shape == "single": ++ holder = SimpleNamespace( ++ model=_compiled_draft(live=True), ++ vllm_config=_shared_config(), ++ cudagraph_manager=SimpleNamespace( ++ cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY ++ ), ++ ) ++ runner, want = SimpleNamespace(speculator=holder), "FULL_DECODE_ONLY" ++ elif shape == "unbuilt": ++ holder = _resolved_speculator(_compiled_draft(live=True), None, None) ++ runner = SimpleNamespace(speculator=holder) ++ want = "prefill:unresolved,decode:unresolved" ++ else: ++ holder = SimpleNamespace( ++ model=_compiled_draft(live=True), ++ vllm_config=_shared_config(), ++ cudagraph_dispatcher=SimpleNamespace( ++ keys_initialized=True, cudagraph_mode=CUDAGraphMode.PIECEWISE ++ ), ++ ) ++ runner, want = SimpleNamespace(drafter=holder), "PIECEWISE" ++ (line,) = _logged_identity(monkeypatch, e2_identity.log_draft_identity, runner) ++ m = _IDENTITY_LINE.search(line) ++ assert m and m["role"] == "draft" and m["gmode"] == want, line ++ ++ ++def test_target_line_at_load_draft_line_after_kv_init(monkeypatch) -> None: ++ """Load time logs the target line only (config level, unchanged); the ++ draft line is logged by log_draft_identity, which the worker calls after ++ initialize_kv_cache, where the drafter's managers are resolved.""" ++ from vllm.v1.worker import e2_identity ++ ++ runner = SimpleNamespace( ++ get_model=lambda: None, ++ speculator=_resolved_speculator( ++ _compiled_draft(live=True), ++ CUDAGraphMode.FULL_AND_PIECEWISE, ++ CUDAGraphMode.FULL_DECODE_ONLY, ++ ), ++ ) ++ lines = _logged_identity( ++ monkeypatch, e2_identity.log_launch_identity, runner, _shared_config() ++ ) ++ assert len(lines) == 1, lines ++ m = _IDENTITY_LINE.search(lines[0]) ++ assert m and m["role"] == "target", lines ++ assert (m["cmode"], m["gmode"]) == ("VLLM_COMPILE", "FULL_AND_PIECEWISE") ++ # No drafter on this rank (not the last PP stage): no draft line, no warning. ++ assert ( ++ _logged_identity(monkeypatch, e2_identity.log_draft_identity, SimpleNamespace()) ++ == [] ++ ) ++ # The worker wiring: log_draft_identity follows initialize_kv_cache. ++ src = (REPO / "vllm/v1/worker/gpu_worker.py").read_text(encoding="utf-8") ++ body = src[src.index("def initialize_from_config") :] ++ body = body[: body.index("\n def ", 1)] ++ assert "log_draft_identity(self.model_runner)" in body ++ assert body.index("initialize_kv_cache(") < body.index( ++ "log_draft_identity(self.model_runner)" ++ ) ++ ++ ++# ------------------------------------------ draft-role PLE layer ids ++ ++ ++def _ple_config(ids, *, num_hidden_layers=4, mtp_layers=1): ++ def text(): ++ return SimpleNamespace( ++ ple_layer_ids=list(ids), ++ num_hidden_layers=num_hidden_layers, ++ mtp_num_hidden_layers=mtp_layers, ++ ) ++ ++ return SimpleNamespace( ++ engram_config=SimpleNamespace(cpu_offload=False), ++ model_config=SimpleNamespace(hf_text_config=text()), ++ speculative_config=SimpleNamespace( ++ draft_model_config=SimpleNamespace(hf_text_config=text()) ++ ), ++ ) ++ ++ ++def _ple_guard_ids(monkeypatch, config, role) -> str: ++ lines: list[str] = [] ++ monkeypatch.setattr( ++ e2_guards.logger, "info", lambda fmt, *args: lines.append(fmt % args) ++ ) ++ assert_ple_mode(config, role=role) ++ (line,) = lines ++ return line.rsplit("ple_layer_ids=", 1)[1] ++ ++ ++def test_draft_ple_guard_reports_the_drafters_own_ple_layers(monkeypatch) -> None: ++ """The reference host logged ple_layer_ids=[2] on the DRAFT guard ++ line -- the target's ids, read from the target's model_config. The MTP ++ drafter's layers (num_hidden_layers + i) carry no PLE block, so the draft ++ line must say []; the target line keeps [2].""" ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "1") ++ _clear_envs_cache() ++ cfg = _ple_config([2]) ++ assert _ple_guard_ids(monkeypatch, cfg, "target") == "[2]." ++ assert _ple_guard_ids(monkeypatch, cfg, "draft") == "[]." ++ ++ ++def test_draft_ple_ids_follow_the_mtp_layer_range(monkeypatch) -> None: ++ """Discrimination: an id that DOES name an MTP layer (num_hidden_layers+1) ++ is reported on the draft line -- the rule is the drafter's own layers, not ++ 'always empty'. (Qwen4ExpTextConfig rejects such ids; this is a stand-in.)""" ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "1") ++ _clear_envs_cache() ++ cfg = _ple_config([2, 5], num_hidden_layers=4, mtp_layers=1) ++ assert _ple_guard_ids(monkeypatch, cfg, "draft") == "[5]." ++ assert _ple_guard_ids(monkeypatch, cfg, "target") == "[2, 5]." +diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py +index 7624ae1..947dada 100644 +--- a/vllm/v1/e2_guards.py ++++ b/vllm/v1/e2_guards.py +@@ -484,6 +484,38 @@ def resolve_ple_mode(vllm_config) -> str: + return PLE_MODE_UVA if cpu_offload else PLE_MODE_DEVICE + + ++def _role_ple_layer_ids(vllm_config, role: str) -> list[int]: ++ """The 1-based PLE layer ids the role's own decoder layers carry. ++ ++ ``role="target"``: ``model_config.hf_text_config.ple_layer_ids``. ++ ``role="draft"``: the drafter is constructed with the target's VllmConfig ++ (the trap an earlier commit fixed for the architectures), so the target's ++ model_config names the TARGET's PLE layers. The drafter's decoder layers ++ are ``num_hidden_layers + i`` for ``i < mtp_num_hidden_layers`` ++ (``Qwen4ExpMultiTokenPredictor``), and a layer builds a PLE block only if ++ ``layer_idx + 1`` is in ``ple_layer_ids`` (``Qwen4ExpDecoderLayer``). The ++ draft line reports exactly that intersection, read from ++ ``speculative_config.draft_model_config`` -- empty for this checkpoint, ++ whose ids must name backbone layers (``Qwen4ExpTextConfig``). ++ """ ++ model_config = getattr(vllm_config, "model_config", None) ++ if role == "draft": ++ spec = getattr(vllm_config, "speculative_config", None) ++ draft_config = getattr(spec, "draft_model_config", None) ++ if draft_config is not None: ++ model_config = draft_config ++ text_config = getattr(model_config, "hf_text_config", None) ++ ids = list(getattr(text_config, "ple_layer_ids", None) or []) ++ if role != "draft": ++ return ids ++ num_hidden_layers = getattr(text_config, "num_hidden_layers", None) ++ if num_hidden_layers is None: ++ return ids ++ num_mtp = int(getattr(text_config, "mtp_num_hidden_layers", 1)) ++ own = range(int(num_hidden_layers) + 1, int(num_hidden_layers) + num_mtp + 1) ++ return [layer_id for layer_id in ids if layer_id in own] ++ ++ + def assert_ple_mode(vllm_config, *, role: str = "target") -> str: + """Startup guard: resolve, assert and log the effective PLE mode. + +@@ -492,11 +524,10 @@ def assert_ple_mode(vllm_config, *, role: str = "target") -> str: + target's EngramConfig and the same process environment, so a mixed + configuration fails there too). NOT disabled by VLLM_E2_ALLOW_BREAKABLE: + that opt-out is about graph identity, not about owning the PLE table twice. ++ The logged ``ple_layer_ids`` are the role's own (:func:`_role_ple_layer_ids`). + """ + mode = resolve_ple_mode(vllm_config) +- model_config = getattr(vllm_config, "model_config", None) +- text_config = getattr(model_config, "hf_text_config", None) +- ple_layer_ids = list(getattr(text_config, "ple_layer_ids", None) or []) ++ ple_layer_ids = _role_ple_layer_ids(vllm_config, role) + engram_config = getattr(vllm_config, "engram_config", None) + logger.info( + "e2-guard[%s]: ple_mode=%s pull_transport=%d engram_cpu_offload=%s " +diff --git a/vllm/v1/worker/e2_identity.py b/vllm/v1/worker/e2_identity.py +index e457b09..e2a9527 100644 +--- a/vllm/v1/worker/e2_identity.py ++++ b/vllm/v1/worker/e2_identity.py +@@ -9,8 +9,10 @@ GPU worker rank logs one line per model role it holds:: + ple_weight_dtype= breakable= compile_mode= + cudagraph_mode= role= + +-(one line in the log). ``role=draft`` is logged only where a draft model is +-loaded (the last PP stage). Every value is read from the live objects: ++(one line in the log). ``role=target`` is logged after model loading; ++``role=draft`` is logged only where a draft model is loaded (the last PP ++stage), once the drafter's graph modes are resolved (after the KV cache is ++initialised -- see below). Every value is read from the live objects: + + * ``ple_backend``: the class of each PLE layer's ``ple_embedding`` module, and + the transport that serves its rows on this rank, read from the placeholder's +@@ -25,8 +27,28 @@ loaded (the last PP stage). Every value is read from the live objects: + (``get_offload_output_dtype``). + * ``breakable``: ``is_breakable_cudagraph_enabled()``, the switch the capture + path itself reads. +-* ``compile_mode`` / ``cudagraph_mode``: the role's ``CompilationConfig`` +- (config level, i.e. before KV-cache-time capture-mode resolution). ++* ``compile_mode`` / ``cudagraph_mode`` for ``role=target``: the engine's ++ ``CompilationConfig`` (config level, i.e. before KV-cache-time capture-mode ++ resolution). ++* ``compile_mode`` / ``cudagraph_mode`` for ``role=draft``: the DRAFTER's own ++ resolved state, not the config. The drafter is built with the target's ++ VllmConfig and there is one ``CompilationConfig`` in the engine, so the ++ config fields name the target's profile (the trap an earlier commit fixed for ++ the draft guard's architectures). What the drafter actually runs is: ++ ++ - ``compile_mode``: the mode of the draft model's live ++ ``@support_torch_compile`` wrapper(s) (``do_not_compile`` false -- the ++ property the capture path tests, ``gpu/cudagraph_utils.py`` ++ ``has_compiled_submodule``); ``NONE`` when no wrapper is live. ++ - ``cudagraph_mode``: the mode each speculator cudagraph manager holds after ++ ``init_cudagraph_manager`` (V2 runner): ``prefill:,decode:`` for the ++ autoregressive / MTP speculator, the bare mode for a single-manager ++ speculator; for a V1 drafter, its dispatcher's mode. These are the values ++ ``e2-guard[runtime]`` checks before capture. ``unresolved`` if a manager ++ does not exist yet. ++ ++ Hence the draft line is logged from ``Worker.initialize_from_config`` after ++ ``initialize_kv_cache``, not at load time. + + A role with no PLE layer reports ``none`` for the three PLE fields; several + distinct values on one rank are joined with ``,``. +@@ -102,46 +124,142 @@ def _rank_tag() -> str: + + + def _draft(model_runner: Any) -> tuple[nn.Module | None, Any]: +- """The loaded draft model and its config (V2 speculator or V1 drafter).""" ++ """The loaded draft model and its holder (V2 speculator or V1 drafter).""" + for name in ("speculator", "drafter"): + holder = getattr(model_runner, name, None) + model = getattr(holder, "model", None) + if isinstance(model, nn.Module): +- return model, getattr(holder, "vllm_config", None) ++ return model, holder + return None, None + + ++#: Written where a draft graph mode has no resolved source yet. ++UNRESOLVED = "unresolved" ++#: Speculator cudagraph-manager attributes (the set ``e2-guard[runtime]`` ++#: reads), with the phase label each carries on the draft line. ++_DRAFT_MANAGER_PHASES = ( ++ ("prefill_cudagraph_manager", "prefill"), ++ ("decode_cudagraph_manager", "decode"), ++ ("cudagraph_manager", None), ++) ++ ++ ++def _draft_compile_mode(model: nn.Module) -> str: ++ """Compile mode of the draft model's live ``@support_torch_compile`` wrappers. ++ ++ ``NONE`` when no wrapper is live: the drafter then runs eager, whatever ++ the shared ``CompilationConfig.mode`` says. ++ """ ++ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper ++ ++ modes = { ++ _enum_name(getattr(getattr(m, "compilation_config", None), "mode", None)) ++ for m in model.modules() ++ if isinstance(m, TorchCompileWithNoGuardsWrapper) ++ and not getattr(m, "do_not_compile", True) ++ } ++ return _join(modes) if modes else "NONE" ++ ++ ++def _draft_cudagraph_mode(holder: Any) -> str: ++ """The drafter's resolved cudagraph mode(s), read from its managers.""" ++ parts = [] ++ for attr, phase in _DRAFT_MANAGER_PHASES: ++ if not hasattr(holder, attr): ++ continue ++ manager = getattr(holder, attr) ++ mode = ( ++ UNRESOLVED ++ if manager is None ++ else _enum_name(getattr(manager, "cudagraph_mode", None)) ++ ) ++ parts.append(mode if phase is None else f"{phase}:{mode}") ++ if parts: ++ return ",".join(parts) ++ # V1 drafter: one dispatcher, keyed in _check_and_update_cudagraph_mode. ++ dispatcher = getattr(holder, "cudagraph_dispatcher", None) ++ if dispatcher is not None and getattr(dispatcher, "keys_initialized", False): ++ return _enum_name(getattr(dispatcher, "cudagraph_mode", None)) ++ return UNRESOLVED ++ ++ + def identity_line( +- tag: str, model: nn.Module | None, vllm_config: Any, role: str ++ tag: str, ++ model: nn.Module | None, ++ vllm_config: Any, ++ role: str, ++ *, ++ compile_mode: str | None = None, ++ cudagraph_mode: str | None = None, + ) -> str: ++ """One identity line. ``compile_mode`` / ``cudagraph_mode`` override the ++ config-level values (the draft role passes its resolved state).""" + from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled + + embeddings = _ple_embeddings(model) + backends = {f"{type(e).__name__}@{_transport(e)}" for e in embeddings} + compilation = getattr(vllm_config, "compilation_config", None) ++ if compile_mode is None: ++ compile_mode = _enum_name(getattr(compilation, "mode", None)) ++ if cudagraph_mode is None: ++ cudagraph_mode = _enum_name(getattr(compilation, "cudagraph_mode", None)) + return ( + f"e2-identity[{tag}]: " + f"ple_backend={_join(backends)} " + f"ple_method={_join({_method(e) for e in embeddings})} " + f"ple_weight_dtype={_join({_weight_dtype(e) for e in embeddings})} " + f"breakable={'ON' if is_breakable_cudagraph_enabled() else 'OFF'} " +- f"compile_mode={_enum_name(getattr(compilation, 'mode', None))} " +- f"cudagraph_mode={_enum_name(getattr(compilation, 'cudagraph_mode', None))} " ++ f"compile_mode={compile_mode} " ++ f"cudagraph_mode={cudagraph_mode} " + f"role={role}" + ) + + ++def draft_identity_line(tag: str, draft_model: nn.Module, holder: Any) -> str: ++ """The ``role=draft`` line: PLE fields from the draft modules, graph fields ++ from the drafter's resolved state (never the shared CompilationConfig).""" ++ return identity_line( ++ tag, ++ draft_model, ++ getattr(holder, "vllm_config", None), ++ "draft", ++ compile_mode=_draft_compile_mode(draft_model), ++ cudagraph_mode=_draft_cudagraph_mode(holder), ++ ) ++ ++ + def log_launch_identity(model_runner: Any, vllm_config: Any) -> None: +- """Log the identity line(s) for this rank. Never raises.""" ++ """Log this rank's ``role=target`` line after model load. Never raises. ++ ++ The ``role=draft`` line comes later, from :func:`log_draft_identity`, once ++ the drafter's graph modes exist. ++ """ ++ tag = "PP?_TP?" ++ try: ++ tag = _rank_tag() ++ logger.info( ++ "%s", identity_line(tag, model_runner.get_model(), vllm_config, "target") ++ ) ++ except Exception as error: # log-only: identity must never break a boot ++ logger.warning( ++ "e2-identity[%s]: unavailable (%s: %s)", tag, type(error).__name__, error ++ ) ++ ++ ++def log_draft_identity(model_runner: Any) -> None: ++ """Log this rank's ``role=draft`` line, if it holds a drafter. Never raises. ++ ++ Call after ``initialize_kv_cache``: that is where the V2 runner resolves ++ the target's mode and ``init_cudagraph_manager`` derives the drafter's ++ (and where the V1 runner keys the drafter's dispatcher). ++ """ + tag = "PP?_TP?" + try: ++ draft_model, holder = _draft(model_runner) ++ if draft_model is None: ++ return + tag = _rank_tag() +- roles = [("target", model_runner.get_model(), vllm_config)] +- draft_model, draft_config = _draft(model_runner) +- if draft_model is not None: +- roles.append(("draft", draft_model, draft_config or vllm_config)) +- for role, model, config in roles: +- logger.info("%s", identity_line(tag, model, config, role)) ++ logger.info("%s", draft_identity_line(tag, draft_model, holder)) + except Exception as error: # log-only: identity must never break a boot + logger.warning( + "e2-identity[%s]: unavailable (%s: %s)", tag, type(error).__name__, error +diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py +index 4edfb76..8cb9d18 100644 +--- a/vllm/v1/worker/gpu_worker.py ++++ b/vllm/v1/worker/gpu_worker.py +@@ -673,8 +673,9 @@ class Worker(WorkerBase): + model_runner = cast("GPUModelRunnerV2", self.model_runner) + model_runner._setup_ple_offload(self.parallel_config._ple_offload_ipc_path) + +- # Log-only: this rank's PLE / capture identity for target and drafter, +- # read after the offload placeholders are armed (v1/worker/e2_identity.py). ++ # Log-only: this rank's PLE / capture identity for the target, read ++ # after the offload placeholders are armed (v1/worker/e2_identity.py). ++ # The drafter's line follows initialize_kv_cache (initialize_from_config). + from vllm.v1.worker.e2_identity import log_launch_identity + + log_launch_identity(self.model_runner, self.vllm_config) +@@ -935,6 +936,12 @@ class Worker(WorkerBase): + ), + ) + ++ # Log-only: the drafter's identity line, now that initialize_kv_cache ++ # has resolved its cudagraph managers (v1/worker/e2_identity.py). ++ from vllm.v1.worker.e2_identity import log_draft_identity ++ ++ log_draft_identity(self.model_runner) ++ + if self.model_config.enable_return_routed_experts: + self.model_runner.init_routed_experts_capturer() + +-- +2.43.0 + diff --git a/patches/0054-ple_offload-the-dying-path-traceback-runs-no-excepti.patch b/patches/0054-ple_offload-the-dying-path-traceback-runs-no-excepti.patch new file mode 100644 index 0000000..32c6d72 --- /dev/null +++ b/patches/0054-ple_offload-the-dying-path-traceback-runs-no-excepti.patch @@ -0,0 +1,286 @@ +From 0ed590631391b088b17c35ffb5c54fca34dfc0b8 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 12:30:20 -0400 +Subject: [PATCH 054/103] ple_offload: the dying-path traceback runs no + exception code + +ple_fail_closed formatted its traceback with traceback.format_exception, which +runs exception-defined __str__/__getattribute__; if one blocked, the rank never +reached os._exit(1) and kept its GPUs. A safe formatter reads frames and plain +string arguments only. +--- + tests/v1/ple_offload/test_test_fault.py | 11 +- + .../test_transport_fatal_helper.py | 69 ++++++++++ + vllm/v1/ple_offload/fatal.py | 125 +++++++++++++++++- + 3 files changed, 196 insertions(+), 9 deletions(-) + +diff --git a/tests/v1/ple_offload/test_test_fault.py b/tests/v1/ple_offload/test_test_fault.py +index a199207..ace9fab 100644 +--- a/tests/v1/ple_offload/test_test_fault.py ++++ b/tests/v1/ple_offload/test_test_fault.py +@@ -270,11 +270,12 @@ def test_partial_ready_times_out_and_terminates_production_child( + f"cut={cut},ordinal=1,rank=any", + ) + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) +- monkeypatch.setattr( +- worker_module.envs, +- "VLLM_PLE_OFFLOAD_READY_TIMEOUT", +- 1.0, +- ) ++ # setenv, not setattr(envs, ...): envs serves this name from a module ++ # __getattr__, so monkeypatch's undo would setattr the pre-test default ++ # (600.0) back as a REAL module attribute that shadows __getattr__ for the ++ # rest of the session -- tests/test_envs.py then reads 600.0 whatever the ++ # environment says (a test-ordering failure seen on the reference host). ++ monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "1") + + handle = worker_module.PleOffloadWorker.make_process( + object(), # The bootstrap fault fires before config reconstruction. +diff --git a/tests/v1/ple_offload/test_transport_fatal_helper.py b/tests/v1/ple_offload/test_transport_fatal_helper.py +index f1abd27..16aaaf9 100644 +--- a/tests/v1/ple_offload/test_transport_fatal_helper.py ++++ b/tests/v1/ple_offload/test_transport_fatal_helper.py +@@ -107,6 +107,18 @@ elif scenario == "str-subclass-argument": + threading.Event().wait() + + error = RuntimeError(BlockingString("detail")) ++elif scenario == "blocking-str-cause": ++ class Cause(Exception): ++ def __str__(self): ++ threading.Event().wait() ++ ++ try: ++ try: ++ raise Cause() ++ except Cause as cause: ++ raise RuntimeError("outer") from cause ++ except RuntimeError as outer: ++ error = outer + elif scenario == "huge-string": + error = RuntimeError("x" * (32 << 20)) + else: +@@ -149,6 +161,7 @@ def test_fatal_helper_avoids_exception_controlled_descriptors( + "blocking-str-method", + "overridden-getattribute", + "str-subclass-argument", ++ "blocking-str-cause", + "huge-string", + ], + ) +@@ -164,6 +177,62 @@ def test_fatal_helper_hard_exit_controls( + ) + + ++_FATAL_TRACEBACK_SOURCE = r""" ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++ ++class AcceleratorLikeError(RuntimeError): ++ pass ++ ++ ++def launch_kernel(): ++ raise ValueError("root cause detail") ++ ++ ++def step(): ++ try: ++ launch_kernel() ++ except ValueError as cause: ++ raise AcceleratorLikeError("CUDA error: an illegal memory access") from cause ++ ++ ++try: ++ step() ++except BaseException as error: ++ ple_fail_closed( ++ where="transport-traceback", ++ rank=0, ++ layer="ple", ++ sequence=7, ++ error=error, ++ ) ++""" ++ ++ ++def test_fatal_helper_traceback_names_subclass_frames_and_cause() -> None: ++ """The dying-path traceback still carries what the early boot lacked: the ++ real (subclass) type, its message, the raising frames and the chain.""" ++ completed = subprocess.run( ++ [sys.executable, "-c", _FATAL_TRACEBACK_SOURCE], ++ stdout=subprocess.DEVNULL, ++ stderr=subprocess.PIPE, ++ timeout=120, ++ ) ++ stderr = completed.stderr.decode("utf-8", "replace") ++ assert completed.returncode == 1, stderr ++ assert "PLE fatal: where=transport-traceback" in stderr ++ assert "PLE fatal traceback follows:" in stderr ++ assert "AcceleratorLikeError: CUDA error: an illegal memory access" in stderr ++ assert "ValueError: root cause detail" in stderr ++ assert "direct cause of the following exception" in stderr ++ assert "in launch_kernel" in stderr ++ assert "in step" in stderr ++ # Interpreter order: the cause's block precedes the outer exception's. ++ assert stderr.index("ValueError: root cause detail") < stderr.index( ++ "AcceleratorLikeError: CUDA error" ++ ) ++ ++ + @pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only") + def test_wait_for_ready_rejects_non_string_status_without_comparison( + tmp_path: Path, +diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py +index 90ab4a0..a96515f 100644 +--- a/vllm/v1/ple_offload/fatal.py ++++ b/vllm/v1/ple_offload/fatal.py +@@ -5,8 +5,8 @@ + import functools + import os + import time +-import traceback + from collections.abc import Callable ++from types import TracebackType + from typing import Any, NoReturn, TypeVar + + _F = TypeVar("_F", bound=Callable[..., Any]) +@@ -26,6 +26,122 @@ _TRUSTED_ERROR_LABELS = ( + (KeyboardInterrupt, "KeyboardInterrupt"), + ) + ++# Slot readers for the dying path. Each is the C getter of the base type, ++# called directly, so no __getattribute__, property, metaclass attribute or ++# descriptor defined by the exception (or its class) can run. The traceback ++# module is NOT safe here: traceback.format_exception calls str(error) (and ++# str() on every chained exception) and reads error.__traceback__ through the ++# instance's own __getattribute__, so an exception whose __str__ or ++# __getattribute__ blocks kept the process alive forever -- the fail-closed ++# path hung instead of dying (found by the blocking-exception transport tests). ++_EXC_ARGS = BaseException.args.__get__ ++_EXC_TRACEBACK = BaseException.__traceback__.__get__ ++_EXC_CAUSE = BaseException.__cause__.__get__ ++_EXC_CONTEXT = BaseException.__context__.__get__ ++_EXC_SUPPRESS_CONTEXT = BaseException.__suppress_context__.__get__ ++_TYPE_QUALNAME = type.__dict__["__qualname__"].__get__ ++_TYPE_MODULE = type.__dict__["__module__"].__get__ ++_MAX_CHAIN = 16 ++_MAX_ARG_CHARS = 1024 ++ ++ ++def _safe_type_label(cls: type) -> str: ++ """``module.qualname`` of ``cls`` without running class-defined code.""" ++ try: ++ qualname = _TYPE_QUALNAME(cls) ++ module = _TYPE_MODULE(cls) ++ except BaseException: ++ return "" ++ if type(qualname) is not str: ++ qualname = "" ++ if type(module) is str and module != "builtins": ++ return f"{module}.{qualname}" ++ return qualname ++ ++ ++def _safe_exception_line(error: BaseException) -> str: ++ """Type label plus the exact-``str``/``int`` args; nothing else is shown.""" ++ label = _safe_type_label(type(error)) ++ args = _EXC_ARGS(error) ++ if type(args) is not tuple or len(args) == 0: ++ return label ++ parts = [] ++ for arg in args[:4]: ++ if type(arg) is str: ++ parts.append(arg[:_MAX_ARG_CHARS]) ++ elif type(arg) is int and -(10**18) < arg < 10**18: ++ parts.append(str(arg)) ++ else: ++ parts.append(f"<{_safe_type_label(type(arg))}>") ++ if len(args) > 4: ++ parts.append("...") ++ if len(args) == 1: ++ return f"{label}: {parts[0]}" ++ return f"{label}: ({', '.join(parts)})" ++ ++ ++def _safe_frames(error: BaseException) -> list[str]: ++ lines = [] ++ tb = _EXC_TRACEBACK(error) ++ while type(tb) is TracebackType: ++ code = tb.tb_frame.f_code ++ filename = code.co_filename ++ name = code.co_qualname ++ lineno = tb.tb_lineno ++ lines.append( ++ f' File "{filename if type(filename) is str else "?"}", ' ++ f"line {lineno if type(lineno) is int else '?'}, " ++ f"in {name if type(name) is str else '?'}\n" ++ ) ++ tb = tb.tb_next ++ return lines ++ ++ ++def _safe_format_exception(error: BaseException) -> str: ++ """A traceback.format_exception look-alike that runs no exception code. ++ ++ Frames come from the traceback objects (file, line, function; no source ++ lines, which would mean file I/O on the dying path); the message is the ++ exception's exact-``str`` args, never ``str(error)``. Chains follow ++ ``__cause__`` / ``__context__`` like the interpreter, bounded and ++ cycle-safe. ++ """ ++ chain: list[tuple[BaseException, str]] = [] ++ seen: set[int] = set() ++ current: BaseException | None = error ++ link = "" ++ while current is not None and len(chain) < _MAX_CHAIN: ++ if id(current) in seen: ++ break ++ seen.add(id(current)) ++ chain.append((current, link)) ++ cause = _EXC_CAUSE(current) ++ if cause is not None: ++ current = cause ++ link = ( ++ "\nThe above exception was the direct cause of the " ++ "following exception:\n\n" ++ ) ++ elif _EXC_SUPPRESS_CONTEXT(current) is not True: ++ current = _EXC_CONTEXT(current) ++ link = ( ++ "\nDuring handling of the above exception, another " ++ "exception occurred:\n\n" ++ ) ++ else: ++ current = None ++ out: list[str] = [] ++ # Innermost first, as the interpreter prints; each outer exception's link ++ # text sits between it and the exception it chains to. ++ for index in range(len(chain) - 1, -1, -1): ++ exc, _ = chain[index] ++ out.append("Traceback (most recent call last):\n") ++ out.extend(_safe_frames(exc)) ++ out.append(_safe_exception_line(exc) + "\n") ++ if index > 0: ++ out.append(chain[index][1]) ++ return "".join(out) ++ + + def ple_fail_closed( + *, +@@ -88,9 +204,10 @@ def ple_fail_closed( + # Three-arg form: it does not depend on being inside an active except + # block, and this function is called from one frame further out than + # the raise. +- tb = "".join( +- traceback.format_exception(type(error), error, error.__traceback__) +- ) ++ # _safe_format_exception, not traceback.format_exception: see the note ++ # above _EXC_ARGS -- the stdlib formatter calls str(error), which is ++ # exception-defined code and can block the dying process forever. ++ tb = _safe_format_exception(error) + raw = tb.encode("utf-8", "backslashreplace") + # Truncate from the END, not the start. format_exception puts the + # frames first and the exception TYPE AND MESSAGE last, so a head-side +-- +2.43.0 + diff --git a/patches/0055-probe-TP2-captured-replay-sentinel-keeps-the-graph-s.patch b/patches/0055-probe-TP2-captured-replay-sentinel-keeps-the-graph-s.patch new file mode 100644 index 0000000..0e66905 --- /dev/null +++ b/patches/0055-probe-TP2-captured-replay-sentinel-keeps-the-graph-s.patch @@ -0,0 +1,409 @@ +From 579390199295b0b0447555198eaef7ddf22d2555 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 15:51:26 -0400 +Subject: [PATCH 055/103] probe: TP2 captured-replay sentinel keeps the graph's + static ids alive + +Test-harness lifetime fix; no served code. +--- + probe/tp2_replay_sentinel.py | 388 +++++++++++++++++++++++++++++++++++ + 1 file changed, 388 insertions(+) + create mode 100644 probe/tp2_replay_sentinel.py + +diff --git a/probe/tp2_replay_sentinel.py b/probe/tp2_replay_sentinel.py +new file mode 100644 +index 0000000..bf5edc0 +--- /dev/null ++++ b/probe/tp2_replay_sentinel.py +@@ -0,0 +1,388 @@ ++"""Compiled TP2 captured-replay host sentinel (reference host), exact bytes. ++ ++What it establishes, on the 0.30 tree at TP=2 over the production collective ++path, for PinnedHostVocabParallelEmbedding (VLLM_HOST_EMBED_TABLE=1 layer): ++ ++ E eager TP2 forward -> EXACT base bytes ++ C torch.compile(fullgraph) TP2 forward -> EXACT base bytes ++ G compiled + CUDA-graph captured replay -> EXACT base bytes ++ M host rows written AFTER capture, G replayed -> EXACT mutated bytes ++ (sentinel S0 on rank 0's shard, S1 on rank 1's shard, so each rank's own ++ UVA read AND the all-reduce that carries it to the other rank are covered) ++ R host rows restored, G replayed -> EXACT base bytes (two-way) ++ X REPLAY_SENTINEL_CYCLES more mutate/restore cycles, each with an eager (compiled, ++ not captured) forward -- i.e. an eager all-reduce -- between the host write ++ and the replay: the production interleave of eager and graph collectives ++ P per-rank PRE-all-reduce partial, its own captured graph (no collective), ++ replayed after every host write: splits a UVA staleness from an ++ all-reduce fault if R/X ever fail ++ I the captured graph's static input ids are still IDS after every replay ++ Z POSITIVE CONTROL for the observed R failure signature: static ids overwritten ++ with an id outside both shards -> the replay must give EXACT all-zero ++ bytes (sha ad7facb2586fc6e9), then ids restored -> EXACT base again ++ K FAILING CONTROL: the same compiled+captured path with the lookup table ++ swapped for a DEVICE copy, host rows mutated after capture -> the replay ++ must still give the BASE bytes, i.e. the M check would FAIL in the ++ "capture baked a device copy" world. If K sees the mutation, the ++ instrument does not discriminate and M proves nothing. ++ ++FIX (sentinel fix): the first reference-host run failed R on ++both ranks with sha ad7facb2586fc6e9 = EXACT all-zero output. Cause: capture() ++cloned the static input ids into a local that was never returned, so the ++caching allocator freed that block while the graph still read it (use after ++free of a graph input). G and M passed because no CUDA allocation happened ++between capture and those replays; the eager compiled(ids) call of ++M_compiled_eager reallocated the block and wrote bf16 data over it, which read ++as int32 ids lands outside both shards -> every row masked on both ranks -> ++all-zero all-reduce. capture() now returns and the caller holds static_ids for ++the graph's lifetime; arm I asserts they are intact, arm Z proves the ++signature. REPLAY_SENTINEL_AR=nccl disables the custom all-reduce (discriminating arm). ++ ++What it does NOT establish (stated, as the pytest sentinel's docstring does): ++ * It is stock Inductor + a manual torch.cuda.CUDAGraph under vLLM's ++ graph_capture() context -- not VllmBackend's piecewise split inside the ++ served Flash-Next forward. The served-completion half is not covered here. ++ * Observing the mutation proves the replay DEPENDS ON MUTABLE HOST DATA, not ++ exclusive UVA residency (a per-launch host->device copy would also pass). ++ ++EXPECTED BYTES -- derived, not observed. The table is a closed form: ++ table[t, j] = bf16((t * HIDDEN + j) % 251) ++Every value 0..250 is an integer below 2**8, so it is EXACT in bf16 (8 ++significant bits) and its bf16 bit pattern is the top 16 bits of the float32. ++Sentinels: S0 := bf16 max (0x7F7F), S1 := bf16 lowest (0xFF7F). The TP2 ++forward masks the other rank's rows to +0.0 and all-reduces: x + (+0.0) == x ++exactly for every value here (no -0.0 appears, so the signed-zero caveat of the ++real TP path cannot arise). So the output for IDS is the concatenation of the ++table rows, byte for byte. `python tp2_replay_sentinel.py --derive` computes ++the digests with numpy only (no torch, no GPU); the reference-host run recomputes them the ++same way and must match the constants below before any arm is judged. ++ ++Run (reference host, from the wrapper script): torchrun --nproc-per-node 2 ++Env: REPLAY_SENTINEL_AR=custom|nccl (default custom), REPLAY_SENTINEL_CYCLES (default 3), ++ REPLAY_SENTINEL_OUT= for the per-rank JSON. ++""" ++ ++from __future__ import annotations ++ ++import hashlib ++import json ++import os ++import sys ++import time ++ ++import numpy as np ++ ++VOCAB = 4096 ++HIDDEN = 256 ++SHARD = VOCAB // 2 ++S0 = 1234 # rank 0 shard [0, 2048) ++S1 = 3000 # rank 1 shard [2048, 4096) ++IDS = [0, 1, S0, S1, VOCAB - 1, S0, S1, 2048] ++BF16_MAX = 0x7F7F ++BF16_LOWEST = 0xFF7F ++OUT_OF_SHARD_ID = 1 << 30 # outside both shards and the (empty) added vocab ++ ++# Derived offline with `--derive` (numpy only) and cross-checked ++# by an independent torch-CPU emulation of the TP2 forward (two shards, mask to ++# +0.0, bf16 sum): both give these digests. The reference-host run re-derives and refuses ++# to judge if they disagree. ZERO is the R_restored sha (all +0.0 bytes). ++EXPECT_BASE_SHA256 = "2de516cb758dbf3475b4d48a474647b27c728421dba7cecb78f0b57786d75bad" ++EXPECT_MUT_SHA256 = "e03d704a409642d0178963db202af4381958ecbbcf4f485e8b1325d154b9de04" ++EXPECT_ZERO_SHA256 = "ad7facb2586fc6e966c004d7d1d16b024f5805ff7cb47c7a85dabd8b48892ca7" ++ ++ ++def _bf16_bits_of_int(values: np.ndarray) -> np.ndarray: ++ f32 = values.astype(np.float32) ++ assert np.all(f32 == values), "closed form left the exact-integer range" ++ return (f32.view(np.uint32) >> 16).astype(np.uint16) ++ ++ ++def closed_form_rows(mutated: bool, owner: int | None = None) -> np.ndarray: ++ """uint16 bf16 bit patterns, shape [len(IDS), HIDDEN]. ++ ++ owner=None: the all-reduced output. owner=r: rank r's pre-all-reduce ++ partial (rows outside rank r's shard are +0.0). ++ """ ++ rows = [] ++ for t in IDS: ++ if owner is not None and not (owner * SHARD <= t < (owner + 1) * SHARD): ++ rows.append(np.zeros(HIDDEN, dtype=np.uint16)) ++ elif mutated and t == S0: ++ rows.append(np.full(HIDDEN, BF16_MAX, dtype=np.uint16)) ++ elif mutated and t == S1: ++ rows.append(np.full(HIDDEN, BF16_LOWEST, dtype=np.uint16)) ++ else: ++ k = (t * HIDDEN + np.arange(HIDDEN, dtype=np.int64)) % 251 ++ rows.append(_bf16_bits_of_int(k)) ++ return np.stack(rows) ++ ++ ++def digest(rows_u16: np.ndarray) -> str: ++ return hashlib.sha256(rows_u16.astype(" dict[str, str]: ++ base, mut = closed_form_rows(False), closed_form_rows(True) ++ assert digest(base) != digest(mut), "sentinel does not change the bytes" ++ d = {"base": digest(base), "mut": digest(mut), ++ "zero": digest(np.zeros_like(base))} ++ for r in (0, 1): ++ d[f"part{r}_base"] = digest(closed_form_rows(False, owner=r)) ++ d[f"part{r}_mut"] = digest(closed_form_rows(True, owner=r)) ++ assert d[f"part{r}_base"] != d[f"part{r}_mut"], f"rank {r} partial blind" ++ # the bf16 constants are what torch.finfo reports (checked on the reference host too) ++ return d ++ ++ ++def _log(tag: str, msg: str) -> None: ++ rank = os.environ.get("RANK", "-") ++ print(f"{time.strftime('%H:%M:%S', time.gmtime())} REPLAY_SENTINEL r{rank} {tag} {msg}", ++ flush=True) ++ ++ ++def main() -> int: # noqa: C901 -- one linear protocol, kept in one place ++ import torch ++ import torch.nn.functional as F ++ ++ from vllm.config import VllmConfig, set_current_vllm_config ++ from vllm.distributed.parallel_state import ( ++ ensure_model_parallel_initialized, ++ get_tp_group, ++ graph_capture, ++ init_distributed_environment, ++ set_custom_all_reduce, ++ ) ++ ++ derived = derive() ++ if (derived["base"], derived["mut"], derived["zero"]) != ( ++ EXPECT_BASE_SHA256, EXPECT_MUT_SHA256, EXPECT_ZERO_SHA256 ++ ): ++ _log("FAIL", f"derivation drift: derived={derived}") ++ return 2 ++ info = torch.finfo(torch.bfloat16) ++ max_bits = int(torch.tensor([info.max], dtype=torch.bfloat16).view(torch.int16)[0]) & 0xFFFF ++ low_bits = int(torch.tensor([info.min], dtype=torch.bfloat16).view(torch.int16)[0]) & 0xFFFF ++ if (max_bits, low_bits) != (BF16_MAX, BF16_LOWEST): ++ _log("FAIL", f"bf16 max/lowest bits {max_bits:#x}/{low_bits:#x}") ++ return 2 ++ ++ ar_want = os.environ.get("REPLAY_SENTINEL_AR", "custom") ++ if ar_want not in ("custom", "nccl"): ++ _log("FAIL", f"REPLAY_SENTINEL_AR={ar_want!r} (custom|nccl)") ++ return 2 ++ cycles = int(os.environ.get("REPLAY_SENTINEL_CYCLES", "3")) ++ set_custom_all_reduce(ar_want == "custom") ++ ++ rank = int(os.environ["RANK"]) ++ local_rank = int(os.environ["LOCAL_RANK"]) ++ world = int(os.environ["WORLD_SIZE"]) ++ assert world == 2, "TP2 sentinel" ++ torch.cuda.set_device(local_rank) ++ device = torch.device("cuda", local_rank) ++ results: dict[str, object] = {"rank": rank, "derived": derived} ++ ok = True ++ part_base, part_mut = derived[f"part{rank}_base"], derived[f"part{rank}_mut"] ++ ++ def check(arm: str, got: bytes, want: str) -> bool: ++ sha = hashlib.sha256(got).hexdigest() ++ good = sha == want ++ results[arm] = {"sha256": sha, "want": want, "pass": good} ++ _log("PASS" if good else "FAIL", f"arm={arm} sha={sha[:16]} want={want[:16]}") ++ return good ++ ++ with set_current_vllm_config(VllmConfig()): ++ init_distributed_environment( ++ world_size=world, ++ rank=rank, ++ distributed_init_method="env://", ++ local_rank=local_rank, ++ backend="nccl", ++ ) ++ ensure_model_parallel_initialized(2, 1) ++ import vllm.model_executor.layers.vocab_parallel_embedding as vpe ++ ++ if not vpe.PIN_MEMORY: ++ _log("FAIL", "pinned memory unavailable") ++ return 2 ++ layer = vpe.PinnedHostVocabParallelEmbedding( ++ VOCAB, HIDDEN, params_dtype=torch.bfloat16 ++ ) ++ si = layer.shard_indices ++ start, end = si.org_vocab_start_index, si.org_vocab_end_index ++ _log("INFO", f"shard [{start},{end}) weight={tuple(layer.weight.shape)} " ++ f"pinned={layer.weight.is_pinned()} uva={layer.accelerator_weight_view is not None}") ++ assert (start, end) == ((0, SHARD) if rank == 0 else (SHARD, VOCAB)) ++ assert layer.weight.is_pinned() and layer.accelerator_weight_view is not None ++ # Fill the shard from the closed form (values exact in bf16). ++ t = torch.arange(start, end, dtype=torch.int64).unsqueeze(1) ++ j = torch.arange(HIDDEN, dtype=torch.int64).unsqueeze(0) ++ layer.weight.data[: end - start].copy_(((t * HIDDEN + j) % 251).to(torch.bfloat16)) ++ ++ tp = get_tp_group() ++ ca = getattr(tp.device_communicator, "ca_comm", None) ++ ar_path = "custom" if (ca is not None and not ca.disabled) else "nccl" ++ results["allreduce_path"] = ar_path ++ results["allreduce_requested"] = ar_want ++ _log("INFO", f"all-reduce path at capture: {ar_path} (requested {ar_want})") ++ if ar_path != ar_want: ++ _log("FAIL", f"all-reduce path {ar_path} != requested {ar_want}: arm void") ++ return 2 ++ ++ ids = torch.tensor(IDS, dtype=torch.int32, device=device) ++ ids_bytes = ids.cpu().numpy().tobytes() ++ ++ def to_bytes(x: torch.Tensor) -> bytes: ++ return x.contiguous().view(torch.uint8).cpu().numpy().tobytes() ++ ++ def set_sentinels(on: bool) -> None: ++ for s, bits in ((S0, BF16_MAX), (S1, BF16_LOWEST)): ++ if start <= s < end: ++ row = layer.weight.data[s - start] ++ if on: ++ row.view(torch.int16).fill_( ++ bits - 0x10000 if bits >= 0x8000 else bits ++ ) ++ else: ++ k = (s * HIDDEN + torch.arange(HIDDEN)) % 251 ++ row.copy_(k.to(torch.bfloat16)) ++ torch.distributed.barrier(group=tp.cpu_group) ++ ++ def partial(x: torch.Tensor) -> torch.Tensor: ++ # PinnedHostVocabParallelEmbedding.forward minus the all-reduce ++ masked, mask = vpe.get_masked_input_and_mask( ++ x, start, end, si.num_org_vocab_padding, ++ si.added_vocab_start_index, si.added_vocab_end_index, ++ ) ++ out = F.embedding(masked.long(), layer._lookup_weight(x)) ++ out.masked_fill_(mask.unsqueeze(-1), 0) ++ return out ++ ++ # E: eager ++ ok &= check("E_eager", to_bytes(layer(ids)), EXPECT_BASE_SHA256) ++ ++ # C: compiled, not captured ++ compiled = torch.compile(layer, fullgraph=True, dynamic=False) ++ for _ in range(2): ++ out_c = compiled(ids) ++ torch.cuda.synchronize() ++ ok &= check("C_compiled", to_bytes(out_c), EXPECT_BASE_SHA256) ++ ++ def capture(fn): ++ # static_ids is RETURNED: the graph reads it on every replay, so ++ # the caller must keep it alive for the graph's lifetime (the ++ # observed R failure was this block being freed and reallocated). ++ static_ids = ids.clone() ++ graph = torch.cuda.CUDAGraph() ++ with graph_capture(device) as ctx: ++ with torch.cuda.stream(ctx.stream): ++ for _ in range(2): # warm-up off the capture ++ fn(static_ids) ++ ctx.stream.synchronize() ++ with torch.cuda.graph(graph, stream=ctx.stream): ++ static_out = fn(static_ids) ++ return graph, static_out, static_ids ++ ++ # P: the pre-all-reduce partial, its own graph (no collective inside) ++ # A P-setup failure voids only the P arms (logged VOID), never the run. ++ try: ++ compiled_p = torch.compile(partial, fullgraph=True, dynamic=False) ++ graph_p, out_p, ids_p = capture(compiled_p) ++ except Exception as exc: # noqa: BLE001 -- diagnostic arm only ++ _log("VOID", f"arm=P setup failed: {type(exc).__name__}: {str(exc)[:160]}") ++ graph_p = out_p = ids_p = None ++ results["P_available"] = graph_p is not None ++ ++ def check_partial(arm: str, want: str) -> bool: ++ if graph_p is None: ++ return True ++ graph_p.replay() ++ torch.cuda.synchronize() ++ return check(arm, to_bytes(out_p), want) ++ ++ # G / M / R: compiled + captured, sentinel written after capture ++ graph, static_out, static_ids = capture(compiled) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check("G_captured", to_bytes(static_out), EXPECT_BASE_SHA256) ++ ok &= check_partial("P_G", part_base) ++ set_sentinels(True) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check("M_mutated_after_capture", to_bytes(static_out), EXPECT_MUT_SHA256) ++ ok &= check_partial("P_M", part_mut) ++ # the compiled-but-not-captured path must see it too (eager all-reduce) ++ ok &= check("M_compiled_eager", to_bytes(compiled(ids)), EXPECT_MUT_SHA256) ++ set_sentinels(False) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check("R_restored", to_bytes(static_out), EXPECT_BASE_SHA256) ++ ok &= check_partial("P_R", part_base) ++ ok &= check("I_ids_intact_R", to_bytes(static_ids), hashlib.sha256(ids_bytes).hexdigest()) ++ ++ # X: more cycles, an eager all-reduce between every host write and replay ++ for c in range(cycles): ++ set_sentinels(True) ++ ok &= check(f"X{c}_eager_mut", to_bytes(compiled(ids)), EXPECT_MUT_SHA256) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check(f"X{c}_replay_mut", to_bytes(static_out), EXPECT_MUT_SHA256) ++ ok &= check_partial(f"X{c}_P_mut", part_mut) ++ set_sentinels(False) ++ ok &= check(f"X{c}_eager_base", to_bytes(compiled(ids)), EXPECT_BASE_SHA256) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check(f"X{c}_replay_base", to_bytes(static_out), EXPECT_BASE_SHA256) ++ ok &= check_partial(f"X{c}_P_base", part_base) ++ ok &= check("I_ids_intact_X", to_bytes(static_ids), hashlib.sha256(ids_bytes).hexdigest()) ++ ++ # Z: positive control for the observed signature (ids outside both shards) ++ static_ids.fill_(OUT_OF_SHARD_ID) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check("Z_ids_out_of_shard", to_bytes(static_out), EXPECT_ZERO_SHA256) ++ static_ids.copy_(ids) ++ graph.replay() ++ torch.cuda.synchronize() ++ ok &= check("Z_ids_restored", to_bytes(static_out), EXPECT_BASE_SHA256) ++ del graph, static_ids, graph_p, ids_p ++ ++ # K: failing control -- a device copy baked in before capture. ++ torch._dynamo.reset() ++ uva = layer._uva_weight ++ layer._uva_weight = uva.clone() # device-resident, NOT host-mapped ++ try: ++ compiled_k = torch.compile(layer, fullgraph=True, dynamic=False) ++ graph_k, out_k, ids_k = capture(compiled_k) ++ set_sentinels(True) ++ graph_k.replay() ++ torch.cuda.synchronize() ++ sha_k = hashlib.sha256(to_bytes(out_k)).hexdigest() ++ discriminates = sha_k == EXPECT_BASE_SHA256 and sha_k != EXPECT_MUT_SHA256 ++ results["K_device_copy_control"] = { ++ "sha256": sha_k, ++ "sees_mutation": sha_k == EXPECT_MUT_SHA256, ++ "pass": discriminates, ++ } ++ _log("PASS" if discriminates else "FAIL", ++ f"arm=K_device_copy_control sha={sha_k[:16]} (must equal base " ++ f"{EXPECT_BASE_SHA256[:16]}: a device copy cannot see the host write)") ++ ok &= discriminates ++ del graph_k, ids_k ++ finally: ++ set_sentinels(False) ++ layer._uva_weight = uva ++ ++ results["pass"] = bool(ok) ++ out = os.environ.get("REPLAY_SENTINEL_OUT") ++ if out: ++ with open(f"{out}/rank{rank}.json", "w") as fh: ++ json.dump(results, fh, indent=1) ++ _log("DONE", f"pass={bool(ok)} ar={ar_path} cycles={cycles}") ++ return 0 if ok else 1 ++ ++ ++if __name__ == "__main__": ++ if "--derive" in sys.argv: ++ print(json.dumps(derive())) ++ sys.exit(0) ++ sys.exit(main()) +-- +2.43.0 + diff --git a/patches/0056-topk-in-kernel-canonical-emission-MERLIN_TOPK_SORTED.patch b/patches/0056-topk-in-kernel-canonical-emission-MERLIN_TOPK_SORTED.patch new file mode 100644 index 0000000..4d14fc2 --- /dev/null +++ b/patches/0056-topk-in-kernel-canonical-emission-MERLIN_TOPK_SORTED.patch @@ -0,0 +1,2026 @@ +From 8523d8499726e91371d584fb09bfffcd0bf6fb01 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 23 Sep 2026 18:59:25 -0400 +Subject: [PATCH 056/103] topk: in-kernel canonical emission + (MERLIN_TOPK_SORTED_EMIT) + +Emits the selected indices in canonical ascending order inside the kernel so +the Python QSA sort can be skipped. Output-equivalent to the Python sort; +reference model, property tests and GPU arms included. Default off in this +commit. +--- + csrc/libtorch_stable/persistent_topk.cuh | 227 ++++- + csrc/libtorch_stable/topk.cu | 72 +- + tests/kernels/merlin_sorted_emit_arm.py | 152 +++ + tests/kernels/test_merlin_topk_sorted_emit.py | 904 ++++++++++++++++++ + tests/test_envs.py | 10 +- + tests/test_merlin_e2_toggles.py | 47 +- + vllm/envs.py | 26 + + .../qwen4_exp/nvidia/ops/qsa_indexer.py | 82 +- + vllm/v1/worker/e2_toggles.py | 68 +- + 9 files changed, 1552 insertions(+), 36 deletions(-) + create mode 100644 tests/kernels/merlin_sorted_emit_arm.py + create mode 100644 tests/kernels/test_merlin_topk_sorted_emit.py + +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 749d2d9..0b1b8b5 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -200,6 +200,17 @@ struct PersistentTopKParams { + // concurrency the atomic exists for. + // Resolved once per process on the host, so it is a launch constant. + bool fault_inject_emit; ++ // MERLIN_TOPK_SORTED_EMIT: every row leaves ++ // the kernel in the canonical order MERLIN_QSA_SORT used to impose in ++ // Python -- ascending index, padding last as -1 -- so the Python sort can be ++ // dropped. Paths that already emit that order by construction (trivial, ++ // medium-path restart, deterministic large path) are untouched; the paths that ++ // place by atomic arrival (decode-buffered, medium-buffered, the legacy ++ // large tail) re-emit their OWN selected set through sorted_emit_row(). ++ // Selection is never re-derived, so the selected set cannot change. ++ // Resolved once per process on the host, so it is a launch constant a ++ // cudagraph bakes in. Default OFF. ++ bool sorted_emit; + }; + + // ============================================================================ +@@ -272,6 +283,149 @@ enum TopKExactCounter : int { + kNumTopKExactCounters = 11, // total slots: device buffer AND reader output + }; + ++// ============================================================================ ++// MERLIN_TOPK_SORTED_EMIT -- in-kernel ascending emission of a row's OWN ++// selected set. Declared here, above every caller (decode/medium dispatch, ++// the legacy large tail, and the standalone fallback kernel), for the ++// two-phase-lookup reason given above the counter enum. ++// ++// CONTRACT. On return, row[0, TopK) is byte-identical to what ++// `_qsa_canonicalise` (vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py) makes ++// of the row it was handed: ++// ++// tmp = where(row < 0, INT32_MAX, row); tmp = sort(tmp) ++// row = where(tmp == INT32_MAX, -1, tmp) ++// ++// i.e. the non-negative entries ascending, then -1 for every negative entry ++// (and, to the byte, for any entry equal to INT32_MAX -- the Python spelling ++// maps those to -1 too). It is a function of the MULTISET of values in the ++// row and nothing else, so it cannot change which blocks were selected: the ++// selection is whatever the path already wrote, including a truncated or ++// racy one. ++// ++// WHAT WAS TAKEN FROM UPSTREAM #55122 (head b2312b2de7, det_select_row, the ++// emission half). The placement idea: a per-index membership flag, a ++// BlockScan exclusive prefix over it with BLOCKED ownership (each thread owns ++// consecutive indices, so scan order is index order), and a store straight to ++// the ascending position -- no atomics on the placement path. WHAT IS NOT ++// TAKEN: #55122 re-derives membership with a full-row radix rescan and its ++// own lowest-index tie rule, and canonicalises -0.0 to +0.0. Adopting either ++// would change the selected set against production (most-recent-wins, ++// +0.0 > -0.0), so membership here comes from the path's own emission ++// instead: a 32,768-bit shared bitmap, one word per thread. Thread tx owns ++// indices [32 tx, 32 tx + 32); its exclusive prefix of popcounts is the ++// ascending position of its first set bit. ++// ++// FAST PATH. Every non-negative entry is < 32,768 (= RADIX_THRESHOLD, so every ++// single-CTA path qualifies) and no entry repeats. Cost: one smem word ++// cleared and at most TopK/1024 global loads per thread, TopK shared ++// atomicOr's with no two on the same bit, one 1,024-wide BlockScan, TopK ++// global stores, four block barriers. No global atomics. ++// ++// FALLBACK. An out-of-range value or a duplicate (bitmap atomicOr returns the ++// bit already set) cannot be expressed as a bitmap, so the row is ranked ++// instead: rank(i) = #{j : key_j < key_i} + #{j < i : key_j == key_i}, a stable ++// counting sort that equals torch.sort on ANY int32 multiset. O(TopK^2 / 1024) ++// shared reads per thread; reached only by the legacy large tail (indices past ++// 32,768, control arm only) and by rows whose unfilled slots hold stale ++// values (buffered overflow with its exact repair switched off). ++// ++// SHARED MEMORY. Aliases the dynamic window from byte 0. The caller's shared ++// state is dead on entry -- the leading barrier makes sure no thread is still ++// reading it -- and the trailing barrier keeps the next row's loader off this ++// row's bitmap. Size is static_asserted against kSmemMedium, the floor every ++// persistent_topk launch allocates. ++// ============================================================================ ++constexpr uint32_t kSortedEmitBitmapWords = static_cast(kThreadsPerBlock); ++constexpr uint32_t kSortedEmitBitmapBits = 32u * kSortedEmitBitmapWords; ++static_assert(kSortedEmitBitmapWords == 1024u, ++ "one bitmap word per thread: the scan below is one item per thread"); ++static_assert(RADIX_THRESHOLD <= kSortedEmitBitmapBits, ++ "every single-CTA path's indices must fit the fast-path bitmap"); ++ ++template ++struct SortedEmitSmem { ++ typename cub::BlockScan::TempStorage scan; ++ uint32_t bitmap[kSortedEmitBitmapWords]; ++ int32_t vals[TopK]; ++ uint32_t fallback; ++}; ++ ++template ++__host__ __device__ constexpr size_t sorted_emit_smem_bytes() { ++ return (sizeof(SortedEmitSmem) + 15) & ~size_t(15); ++} ++ ++template ++__device__ __noinline__ void sorted_emit_row(int32_t* __restrict__ row) { ++ static_assert(sorted_emit_smem_bytes() <= kSmemMedium, ++ "sorted emission must fit the smallest persistent_topk window"); ++ static_assert(alignof(SortedEmitSmem) <= 16, ++ "SortedEmitSmem over-aligned for the dynamic window"); ++ using BlockScan = cub::BlockScan; ++ extern __shared__ __align__(16) char sorted_emit_smem[]; ++ auto* s = reinterpret_cast*>(sorted_emit_smem); ++ const uint32_t tx = threadIdx.x; ++ ++ // E0: every emission store of the path is issued and no thread still reads ++ // the path's shared state (block-scope ordering of the global stores too). ++ __syncthreads(); ++ s->bitmap[tx] = 0u; ++ if (tx == 0) s->fallback = 0u; ++ __syncthreads(); ++ ++ // E1: membership. Reads of `row` all happen before E2; writes all after. ++ for (uint32_t i = tx; i < static_cast(TopK); i += kThreadsPerBlock) { ++ const int32_t v = row[i]; ++ s->vals[i] = v; ++ if (v >= 0) { ++ const uint32_t u = static_cast(v); ++ if (u < kSortedEmitBitmapBits) { ++ const uint32_t bit = 1u << (u & 31u); ++ if ((atomicOr(&s->bitmap[u >> 5], bit) & bit) != 0u) { ++ atomicOr(&s->fallback, 1u); // duplicate value ++ } ++ } else { ++ atomicOr(&s->fallback, 1u); // beyond the bitmap ++ } ++ } ++ } ++ __syncthreads(); // E2 ++ ++ if (s->fallback == 0u) { // block-uniform: written only before E2 ++ uint32_t word = s->bitmap[tx]; ++ uint32_t pos = 0u; ++ uint32_t total = 0u; ++ BlockScan(s->scan).ExclusiveSum(static_cast(__popc(word)), pos, ++ total); ++ while (word != 0u) { ++ const uint32_t b = static_cast(__ffs(word) - 1); ++ row[pos++] = static_cast((tx << 5) | b); ++ word &= word - 1u; ++ } ++ for (uint32_t i = total + tx; i < static_cast(TopK); ++ i += kThreadsPerBlock) { ++ row[i] = -1; ++ } ++ } else { ++ constexpr int32_t kBig = 0x7FFFFFFF; ++ for (uint32_t i = tx; i < static_cast(TopK); ++ i += kThreadsPerBlock) { ++ const int32_t vi = s->vals[i]; ++ const int32_t ki = vi < 0 ? kBig : vi; ++ uint32_t rank = 0u; ++ for (uint32_t j = 0; j < static_cast(TopK); ++j) { ++ const int32_t vj = s->vals[j]; ++ const int32_t kj = vj < 0 ? kBig : vj; ++ rank += (kj < ki || (kj == ki && j < i)) ? 1u : 0u; ++ } ++ row[rank] = (ki == kBig) ? -1 : ki; ++ } ++ } ++ // E3: the next row (or the next path in this row) reuses the window. ++ __syncthreads(); ++} ++ + // ============================================================================ + // Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) + // Uses 11-bit half-precision bins for fine granularity. +@@ -1081,8 +1235,13 @@ __device__ __noinline__ void medium_exact_restart( + // by: DarkSharpness + // which at the same time is an optimized topk kernel copied from tilelang + // kernel ++// ++// Returns TRUE when the row was handed to medium_exact_restart(), whose ++// emission is already ascending (MERLIN_TOPK_SORTED_EMIT skips those rows); ++// FALSE when the buffered path placed the row by atomic arrival. The return ++// value is block-uniform (every return below is taken by the whole block). + template +-__device__ __noinline__ void histogram_256_topk( ++__device__ __noinline__ bool histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, + int logits_offset, int seq_len, bool exact_medium, + unsigned long long* __restrict__ counters, bool det_ties, +@@ -1181,7 +1340,7 @@ __device__ __noinline__ void histogram_256_topk( + } + } + __syncthreads(); +- return; ++ return false; + } + + // Past the early return the threshold bin WILL be collected, so the M1 +@@ -1262,7 +1421,7 @@ __device__ __noinline__ void histogram_256_topk( + medium_exact_restart(logits + logits_offset, output_indices, + static_cast(seq_len), counters, + effective_tie_recent(det_ties, tie_recent)); +- return; ++ return true; + } + __syncthreads(); + } +@@ -1361,6 +1520,7 @@ __device__ __noinline__ void histogram_256_topk( + } + __syncthreads(); + } ++ return false; + } + + // ============================================================================ +@@ -1949,7 +2109,8 @@ __device__ void radix_topk(const float* __restrict__ row_input, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t radix_iter, uint32_t tx, bool det_large, + unsigned long long* counters, int fault_inject_cta, +- bool fault_inject_emit, bool tie_recent) { ++ bool fault_inject_emit, bool tie_recent, ++ bool sorted_emit) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; +@@ -2186,6 +2347,30 @@ __device__ void radix_topk(const float* __restrict__ row_input, + // the MERLIN_TOPK_DET_LARGE=0 control, so the control now carries this + // repair too, like the histogram-publication barrier (an earlier commit). + __syncthreads(); ++ ++ // MERLIN_TOPK_SORTED_EMIT on the legacy tail (MERLIN_TOPK_DET_LARGE=0, or ++ // the host's staging-capacity gate): the row was written by EVERY CTA of the ++ // group, by arrival order, so CTA 0 may only re-emit it after all of them ++ // have finished. One more inter-CTA barrier -- the same red_release / ++ // wait_ge pair as every other barrier here, taken by every CTA of the group ++ // because `sorted_emit` is a launch constant, so barrier_phase stays ++ // group-uniform. The closing __syncthreads() just above orders each CTA's ++ // stores before its thread 0's release fence. Indices here run past ++ // 32,768, so sorted_emit_row() takes its counting-sort fallback: this is the ++ // control arm, never the production large path (which is the deterministic one and is ++ // already ascending). ++ if (sorted_emit) { ++ if (tx == 0) { ++ red_release(&state->arrival_counter, 1); ++ } ++ wait_ge(&state->arrival_counter, ++ (barrier_phase + 1) * static_cast(ctas_per_group), tx); ++ barrier_phase++; ++ __syncthreads(); ++ if (cta_in_group == 0) { ++ sorted_emit_row(row_output); ++ } ++ } + } + + // ============================================================================ +@@ -2280,15 +2465,25 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + row_input, row_output, seq_len, params.det_ties, + params.tie_recent, params.exact_counters, params.exact_decode); + if (decode_needs_exact) { ++ // Already ascending (restart emission): no sorted re-emission. + medium_exact_restart( + row_input, row_output, seq_len, params.exact_counters, + effective_tie_recent(params.det_ties, params.tie_recent)); ++ } else if (params.sorted_emit) { ++ // Decode-buffered placement is by atomic arrival. ++ sorted_emit_row(row_output); + } + } else { +- histogram_256_topk(row_input, row_output, 0, seq_len, +- params.exact_medium, params.exact_counters, +- params.det_ties, params.tie_recent); ++ const bool restarted = histogram_256_topk( ++ row_input, row_output, 0, seq_len, params.exact_medium, ++ params.exact_counters, params.det_ties, params.tie_recent); ++ if (!restarted && params.sorted_emit) { ++ // Medium-buffered placement is by atomic arrival. ++ sorted_emit_row(row_output); ++ } + } ++ // The trivial branch above writes [0, seq_len) then -1 in index ++ // order, which is already the canonical order. + } + continue; + } +@@ -2300,11 +2495,27 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, + params.det_large, params.exact_counters, params.fault_inject_cta, + params.fault_inject_emit, +- effective_tie_recent(params.det_ties, params.tie_recent)); ++ effective_tie_recent(params.det_ties, params.tie_recent), ++ params.sorted_emit); + radix_iter++; + } + } + ++// MERLIN_TOPK_SORTED_EMIT for the two host-side fallbacks that bypass ++// persistent_topk_kernel (FilteredTopK, needs >= 128 KiB opt-in smem, and ++// top_k_per_row_decode, needs a cooperative launch that does not fit). Both ++// are unreachable on a 3090 at the shapes Flash-Next serves (the fallbacks ++// never dispatch there), but with the knob on the Python sort is gone, so the op's ++// contract -- canonical order on EVERY dispatch -- is kept here rather than ++// assumed. One block per row, launched on the same stream right after the ++// fallback, so it is stream-ordered after the fallback's stores. ++template ++__global__ void __launch_bounds__(kThreadsPerBlock) ++ sorted_emit_rows_kernel(int32_t* __restrict__ output, uint32_t num_rows) { ++ if (blockIdx.x >= num_rows) return; ++ sorted_emit_row(output + static_cast(blockIdx.x) * TopK); ++} ++ + } // namespace persistent + + // ============================================================================ +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 3c90c25..060f7a6 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -60,6 +60,7 @@ std::atomic g_cached_tie_det{-1}; + std::atomic g_cached_exact_medium{-1}; + std::atomic g_cached_det_large{-1}; + std::atomic g_cached_exact_decode{-1}; ++std::atomic g_cached_sorted_emit{-1}; + + bool publish_cached(std::atomic& mirror, bool value) { + mirror.store(value ? 1 : 0, std::memory_order_release); +@@ -141,6 +142,30 @@ bool exact_decode_enabled() { + return enabled; + } + ++// --------------------------------------------------------------------------- ++// MERLIN_TOPK_SORTED_EMIT: every persistent_topk ++// row leaves the op in the canonical order the Python MERLIN_QSA_SORT imposes ++// (ascending index, padding last as -1), so qsa_indexer.py can drop that sort. ++// See sorted_emit_row() in persistent_topk.cuh for the contract; the selected ++// SET is never touched. Resolved once per process like every switch here, so ++// it is a launch constant a cudagraph bakes in. Default OFF: the maintainer ++// promotes after reference-host qualification. ++// ++// This std::getenv is the single source of truth. Python does not parse the ++// variable to decide anything: qsa_indexer.py asks THIS cache, through ++// persistent_topk_toggle_state() field 6, after its first launch -- so the ++// Python sort is skipped only when the kernel that ran says it sorted. ++// --------------------------------------------------------------------------- ++bool sorted_emit_enabled() { ++ static const bool enabled = [] { ++ const char* raw = std::getenv("MERLIN_TOPK_SORTED_EMIT"); ++ return publish_cached( ++ g_cached_sorted_emit, ++ (raw != nullptr && raw[0] != '\0') ? (std::atoi(raw) != 0) : false); ++ }(); ++ return enabled; ++} ++ + // --------------------------------------------------------------------------- + // MERLIN_TOPK_LARGE_FAULT_INJECT. TEST ONLY. + // +@@ -219,16 +244,21 @@ void report_tie_policy_once() { + stderr, + "[merlin] persistent_topk tie policy: MERLIN_TIE_DET=%d " + "MERLIN_TIE_RECENT=%d MERLIN_TOPK_EXACT_MEDIUM=%d " +- "MERLIN_TOPK_EXACT_DECODE=%d MERLIN_TOPK_DET_LARGE=%d -> " ++ "MERLIN_TOPK_EXACT_DECODE=%d MERLIN_TOPK_DET_LARGE=%d " ++ "MERLIN_TOPK_SORTED_EMIT=%d -> " + "decode-buffered=%s decode-restart=%s medium-buffered=%s " +- "medium-restart=%s large=%s\n", ++ "medium-restart=%s large=%s emission=%s\n", + det ? 1 : 0, tie_recent_enabled() ? 1 : 0, + exact_medium_enabled() ? 1 : 0, exact_decode_enabled() ? 1 : 0, +- det_large_enabled() ? 1 : 0, det ? dir : "atomic-arrival", ++ det_large_enabled() ? 1 : 0, sorted_emit_enabled() ? 1 : 0, ++ det ? dir : "atomic-arrival", + exact_decode_enabled() ? dir : "n/a (detect-only; counter [10])", + det ? dir : "atomic-arrival", + exact_medium_enabled() ? dir : "n/a (exact restart disabled)", +- det_large_enabled() ? dir : "atomic-arrival (non-deterministic)"); ++ det_large_enabled() ? dir : "atomic-arrival (non-deterministic)", ++ sorted_emit_enabled() ++ ? "ascending-in-kernel (Python QSA sort skipped)" ++ : "arrival-order on buffered paths (Python QSA sort decides)"); + return true; + }(); + (void)reported; +@@ -393,6 +423,19 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + max_smem_per_block = device_prop->sharedMemPerBlockOptin; + } + ++ // MERLIN_TOPK_SORTED_EMIT on the two fallbacks that bypass ++ // persistent_topk_kernel: re-emit every row ascending on the same stream. ++ // A plain launch with a fixed configuration, so it is capture-safe. ++ const auto sorted_emit_fallback_rows = [&]() { ++ if (!sorted_emit_enabled()) return; ++ report_tie_policy_once(); ++ P::sorted_emit_rows_kernel ++ <<(num_rows), P::kThreadsPerBlock, ++ P::sorted_emit_smem_bytes(), stream>>>( ++ output.mutable_data_ptr(), ++ static_cast(num_rows)); ++ }; ++ + if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + cudaError_t status = + vllm::FilteredTopKRaggedTransform( +@@ -401,6 +444,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + static_cast(TopK), static_cast(stride), stream); + STD_TORCH_CHECK(status == cudaSuccess, + "FilteredTopK failed: ", cudaGetErrorString(status)); ++ sorted_emit_fallback_rows(); + } else { + STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + STD_TORCH_CHECK( +@@ -497,6 +541,10 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1; + top_k_per_row_decode(logits, next_n, lengths, output, num_rows, + logits.stride(0), logits.stride(1), TopK); ++ sorted_emit_fallback_rows(); ++ const cudaError_t fb_err = cudaGetLastError(); ++ STD_TORCH_CHECK(fb_err == cudaSuccess, "sorted emission failed: ", ++ cudaGetErrorString(fb_err)); + return; + } + cudaError_t status = +@@ -508,6 +556,10 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + static_cast(stride), stream); + STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ", + cudaGetErrorString(status)); ++ sorted_emit_fallback_rows(); ++ const cudaError_t fb_err = cudaGetLastError(); ++ STD_TORCH_CHECK(fb_err == cudaSuccess, "sorted emission failed: ", ++ cudaGetErrorString(fb_err)); + return; + } + +@@ -606,6 +658,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, + params.det_ties = det_ties_enabled(); + params.tie_recent = tie_recent_enabled(); + params.exact_decode = exact_decode_enabled(); ++ params.sorted_emit = sorted_emit_enabled(); + report_tie_policy_once(); + #if MERLIN_TOPK_DIAG_ANY + report_diag_build_once(); +@@ -920,8 +973,12 @@ int64_t persistent_topk_exact_counters_slots() { + // field 3 MERLIN_TOPK_EXACT_MEDIUM (exact_medium_enabled) + // field 4 MERLIN_TOPK_EXACT_DECODE (exact_decode_enabled) + // field 5 MERLIN_TOPK_DET_LARGE (det_large_enabled) +-// bits [48, 55] = field count (6), bits [56, 62] = layout version (1). A ++// field 6 MERLIN_TOPK_SORTED_EMIT (sorted_emit_enabled) ++// bits [48, 55] = field count (7), bits [56, 62] = layout version (2). A + // reader must check both and refuse any other layout rather than guess. ++// Layout 2 = layout 1 plus field 6; fields 0-5 are unchanged, so ++// a reader that knows both can decode either. Field 6 is what ++// qsa_indexer.py reads to decide whether its Python sort may be skipped. + int64_t persistent_topk_toggle_state() { + #ifndef USE_ROCM + const auto field = [](int cached) -> int64_t { +@@ -943,8 +1000,9 @@ int64_t persistent_topk_toggle_state() { + state |= field(g_cached_exact_medium.load(std::memory_order_acquire)) << 6; + state |= field(g_cached_exact_decode.load(std::memory_order_acquire)) << 8; + state |= field(g_cached_det_large.load(std::memory_order_acquire)) << 10; +- state |= int64_t{6} << 48; +- state |= int64_t{1} << 56; ++ state |= field(g_cached_sorted_emit.load(std::memory_order_acquire)) << 12; ++ state |= int64_t{7} << 48; ++ state |= int64_t{2} << 56; + return state; + #else + return 0; // layout version 0: no persistent_topk on ROCm +diff --git a/tests/kernels/merlin_sorted_emit_arm.py b/tests/kernels/merlin_sorted_emit_arm.py +new file mode 100644 +index 0000000..700dab1 +--- /dev/null ++++ b/tests/kernels/merlin_sorted_emit_arm.py +@@ -0,0 +1,152 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""One ARM of the MERLIN_TOPK_SORTED_EMIT GPU equivalence test (GPU only). ++ ++Not a test module: test_merlin_topk_sorted_emit.py runs this file in a fresh ++subprocess per arm, because every persistent_topk switch is a C++ static ++resolved once per process. The caller sets the MERLIN_* environment. ++ ++ python merlin_sorted_emit_arm.py OUT.pt [--graph] ++ ++Writes {"cells": {name: int32 [rows, k] CPU tensor}, "meta": {...}} with the ++RAW kernel output (no Python sort), the exact counters after the run and the ++toggle_state fields. Inputs are seeded and generated on the CPU, so every arm ++sees the same bytes. The output buffer is PRE-FILLED with a fixed pattern ++before each launch, so a slot a path leaves unwritten (overflow without its ++repair) holds the same stale bytes in every arm. ++ ++Cell names starting with `racy_` have a launch-order-dependent SET in the ++production arm (decode overflow with MERLIN_TOPK_EXACT_DECODE=0), so the ++cross-arm equality skips them; the canonical-order check still covers them. ++""" ++ ++import sys ++import zlib ++ ++import torch ++ ++import vllm._custom_ops # noqa: F401 (loads vllm._C) ++from vllm.v1.worker.e2_toggles import decode_topk_state ++ ++WORKSPACE = 1024 * 1024 ++ ++ ++def _row(kind: str, n: int, gen: torch.Generator) -> torch.Tensor: ++ if kind == "ties": ++ levels = torch.tensor([0.0, 0.0, 0.0, 0.5, 1.0, 2.0]) ++ return levels[torch.randint(0, 6, (n,), generator=gen)] ++ if kind == "allequal": ++ return torch.full((n,), 1.25) ++ if kind == "signedzero": ++ x = torch.where( ++ torch.rand(n, generator=gen) < 0.5, torch.tensor(0.0), torch.tensor(-0.0) ++ ) ++ hot = torch.rand(n, generator=gen) < 0.02 ++ x[hot] = torch.rand(int(hot.sum()), generator=gen) + 1.0 ++ return x ++ if kind == "neginf": ++ x = torch.randn(n, generator=gen) ++ x[torch.rand(n, generator=gen) < 0.7] = float("-inf") ++ return x ++ if kind == "distinct": ++ return torch.randperm(n, generator=gen).float() * 1e-3 ++ if kind == "onebin": ++ return 1.0 + torch.rand(n, generator=gen) * 0.01 ++ raise ValueError(kind) ++ ++ ++# (name, k, row lengths, kind). Lengths are in BLOCKS; 32,768 / 32,769 are the ++# 131,072 / 131,076-token medium/large boundary at compress 4. ++CELLS = [ ++ ("trivial_300", 512, [300] * 4, "ties"), ++ ("trivial_512", 512, [512] * 4, "ties"), ++ ("decode_600_ties", 512, [600] * 4, "ties"), ++ ("decode_1024_ties", 512, [1024] * 4, "ties"), # 4K tokens ++ ("decode_4096_signedzero", 512, [4096] * 4, "signedzero"), ++ ("decode_3000_neginf", 512, [3000] * 4, "neginf"), ++ ("racy_decode_8192_allequal", 512, [8192] * 4, "allequal"), ++ ("racy_decodeovf_7000_onebin", 512, [7000] * 4, "onebin"), ++ ("medium_8193_ties", 512, [8193] * 4, "ties"), ++ ("medium_16384_ties", 512, [16384] * 4, "ties"), # 65K tokens ++ ("medium_12288_onebin", 512, [12288] * 4, "onebin"), # medium-path restart ++ ("medium_32768_allequal", 512, [32768] * 4, "allequal"), # 131,072 tokens ++ ("medium_32768_signedzero", 512, [32768] * 4, "signedzero"), ++ ("large_32769_ties", 512, [32769] * 2, "ties"), # 131,076 tokens ++ ("large_33750_ties", 512, [33750] * 2, "ties"), # 135K tokens ++ ("large_50000_allequal", 512, [50000] * 2, "allequal"), ++ ("large_65536_signedzero", 512, [65536] * 2, "signedzero"), ++ ("large_40000_distinct", 512, [40000] * 2, "distinct"), ++ ("mixed_batch_ties", 512, [100, 600, 8192, 8193, 32768, 32769, 40000, 512], "ties"), ++ ("k2048_decode_ties", 2048, [6000] * 2, "ties"), ++ ("k2048_medium_allequal", 2048, [20000] * 2, "allequal"), ++ ("k2048_large_ties", 2048, [45000] * 2, "ties"), ++] ++ ++ ++def _inputs(name, k, lengths, kind): ++ # crc32, not hash(): str hashing is salted per process, and every arm ++ # must see the same bytes. ++ gen = torch.Generator().manual_seed(zlib.crc32(name.encode())) ++ stride = max(lengths) ++ logits = torch.full((len(lengths), stride), float("-inf")) ++ for r, n in enumerate(lengths): ++ logits[r, :n] = _row(kind, n, gen) ++ prefill = torch.stack( ++ [torch.randperm(4 * k, generator=gen)[:k] - k // 2 for _ in lengths] ++ ).to(torch.int32) ++ return logits, torch.tensor(lengths, dtype=torch.int32), prefill ++ ++ ++def _launch(logits, lengths, out, ws, k): ++ torch.ops._C.persistent_topk(logits, lengths, out, ws, k, int(lengths.max())) ++ ++ ++def main() -> int: ++ out_path = sys.argv[1] ++ graph = "--graph" in sys.argv[2:] ++ dev = torch.device("cuda") ++ ws = torch.empty(WORKSPACE, dtype=torch.uint8, device=dev) ++ cells, replays_ok = {}, True ++ for name, k, lengths, kind in CELLS: ++ logits, lens, prefill = _inputs(name, k, lengths, kind) ++ logits, lens = logits.to(dev), lens.to(dev) ++ out = prefill.to(dev).clone() ++ _launch(logits, lens, out, ws, k) ++ torch.cuda.synchronize() ++ cells[name] = out.cpu().clone() ++ if graph: ++ static_out = prefill.to(dev).clone() ++ g = torch.cuda.CUDAGraph() ++ s = torch.cuda.Stream() ++ s.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(s): ++ _launch(logits, lens, static_out, ws, k) # warm on a side stream ++ torch.cuda.current_stream().wait_stream(s) ++ with torch.cuda.graph(g): ++ _launch(logits, lens, static_out, ws, k) ++ static_out.copy_(prefill.to(dev)) ++ g.replay() ++ torch.cuda.synchronize() ++ if not name.startswith("racy_"): ++ replays_ok &= torch.equal(static_out.cpu(), cells[name]) ++ counters = torch.zeros( ++ torch.ops._C.persistent_topk_exact_counters_slots(), ++ dtype=torch.int64, ++ device=dev, ++ ) ++ torch.ops._C.persistent_topk_exact_counters(counters, 0, 1) ++ toggles = decode_topk_state(int(torch.ops._C.persistent_topk_toggle_state())) ++ meta = { ++ "counters": counters.cpu().tolist(), ++ "toggle_sorted_emit": toggles["sorted_emit"] ++ if isinstance(toggles, dict) ++ else toggles, ++ "toggles": str(toggles), ++ "graph_replays_match_eager": bool(replays_ok), ++ } ++ torch.save({"cells": cells, "meta": meta}, out_path) ++ return 0 ++ ++ ++if __name__ == "__main__": ++ sys.exit(main()) +diff --git a/tests/kernels/test_merlin_topk_sorted_emit.py b/tests/kernels/test_merlin_topk_sorted_emit.py +new file mode 100644 +index 0000000..1876fae +--- /dev/null ++++ b/tests/kernels/test_merlin_topk_sorted_emit.py +@@ -0,0 +1,904 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""MERLIN_TOPK_SORTED_EMIT: in-kernel canonical emission. ++ ++The claim under test: with the knob on, every row persistent_topk returns is ++BYTE-IDENTICAL to what the knob-off kernel returns after the Python ++MERLIN_QSA_SORT (`qsa_indexer._canonical`: ascending, padding last as -1), on ++every path the tree dispatches -- trivial, decode (buffered and exact restart), ++medium (buffered and exact restart), large (deterministic and the legacy ++tail) and the host fallbacks -- with the SAME selected set, so production ++outputs cannot change. ++ ++Two halves: ++ ++CPU (runnable offline, no GPU): ++ * a reference model of each path's selection + tie rule + emission LAYOUT ++ (the arrival-order layouts are modelled as seeded permutations; the ++ ascending layouts -- trivial, restart, large deterministic -- are ports of the kernel's ++ own slot arithmetic, including the host's chunking); ++ * a port of sorted_emit_row() (bitmap words, popcount exclusive scan, bit ++ walk, -1 fill, counting-sort fallback); ++ * property tests: sorted_emit_row(emission) == _canonical(emission) ++ byte-for-byte on random rows with heavy ties, +-0.0, -inf, and lengths ++ across 512/513, 8192/8193, 32,768/32,769 (131,072/131,076 tokens), for ++ the paths the kernel re-emits; native emission == _canonical(native) for ++ the paths it skips; and every path's set == the exact top-k under the ++ effective tie rule; ++ * pathological rows (duplicates, stale, out-of-range, INT32_MAX, odd ++ negatives) through the fallback; ++ * FAILING CONTROLS: a canonicaliser without duplicate detection, without the ++ -1 fill, the identity (sort skipped), and a flipped tie rule each FAIL the ++ same checks -- so the checks can fail; ++ * the Python gate: kernel_sorted_emit() reads toggle_state field 6, latches, ++ refuses layout 1, never latches `uninit`; _topk skips the sort only then. ++ ++GPU (reference host; written here, run on a GPU host): the knob is a ++process-lifetime C++ static, so each arm is a SUBPROCESS ++(`merlin_sorted_emit_arm.py`). Knob-on raw output must equal knob-off output ++after `_canonical`, cell by cell, as int32 bytes; the failing controls (knob ++off without the sort, and the tie rule flipped) must DIFFER. ++ ++ pytest --noconftest tests/kernels/test_merlin_topk_sorted_emit.py ++""" ++ ++from __future__ import annotations ++ ++import json ++import os ++import subprocess ++import sys ++from pathlib import Path ++ ++import numpy as np ++import pytest ++import torch ++ ++qsa_indexer = pytest.importorskip( ++ "vllm.models.qwen4_exp.nvidia.ops.qsa_indexer", ++ reason="qsa_indexer needs triton (the reference-host venv has it; offline runs use the shim)", ++) ++ ++INT32_MAX = 0x7FFFFFFF ++ ++# Kernel constants (csrc/libtorch_stable/persistent_topk.cuh). ++RADIX_THRESHOLD = 32768 ++HIST2048_THRESHOLD = 8192 ++MAX_BUFFERED_ITEMS = 4096 ++DBUF = (8192 - 8 - 2 * (256 + 128)) // 2 # 3708 ++BITMAP_WORDS = 1024 ++BITMAP_BITS = 32 * BITMAP_WORDS ++K_SMEM_MEDIUM = 35968 ++K_FIXED_SMEM_LARGE = ((256 + 256 + 5) * 4 + 15) & ~15 ++SMEM_OPTIN_3090 = 101376 ++ ++ ++def canonical(row: np.ndarray) -> np.ndarray: ++ """The REAL Python reference: qsa_indexer._canonical on an int32 row.""" ++ t = torch.from_numpy(np.ascontiguousarray(row, dtype=np.int32)).view(1, -1) ++ return qsa_indexer._canonical(t).view(-1).numpy() ++ ++ ++# ============================================================ keys ========== ++def ordered_key(x: np.ndarray) -> np.ndarray: ++ """convert_to_uint32_v2: +0.0 orders ABOVE -0.0 (we do not adopt #55122's ++ signed-zero canonicalisation).""" ++ bits = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32) ++ neg = (bits & np.uint32(0x80000000)) != 0 ++ return np.where(neg, ~bits, bits | np.uint32(0x80000000)).astype(np.uint32) ++ ++ ++def effective_tie_recent(det_ties: bool, tie_recent: bool) -> bool: ++ return tie_recent if det_ties else True ++ ++ ++# ======================================================= dispatch =========== ++def path_of(seq_len: int, top_k: int) -> str: ++ if seq_len > RADIX_THRESHOLD: ++ return "large" ++ if seq_len <= top_k: ++ return "trivial" ++ if seq_len <= HIST2048_THRESHOLD: ++ return "decode" ++ return "medium" ++ ++ ++# ====================================================== reference =========== ++def reference_select(keys: np.ndarray, top_k: int, keep_recent: bool): ++ """Exact top-k by key; among elements EQUAL to the k-th key keep the ++ highest indices (most-recent-wins) or the lowest (oldest-wins). ++ Returns (selected index set as sorted int64, pivot key, #eq kept, #eq).""" ++ n = keys.size ++ assert n >= top_k ++ pivot = np.sort(keys)[::-1][top_k - 1] ++ gt = np.nonzero(keys > pivot)[0] ++ eq = np.nonzero(keys == pivot)[0] ++ need = top_k - gt.size ++ kept_eq = eq[-need:] if keep_recent else eq[:need] ++ return np.sort(np.concatenate([gt, kept_eq])), pivot, need, eq.size ++ ++ ++# ================================================= path emissions =========== ++def emit_trivial(seq_len: int, top_k: int) -> np.ndarray: ++ """persistent_topk_kernel trivial branch, verbatim.""" ++ i = np.arange(top_k) ++ return np.where(i < seq_len, i, -1).astype(np.int32) ++ ++ ++def emit_buffered(keys, top_k, det_ties, tie_recent, rng, candidates=None): ++ """Decode-/medium-buffered emission LAYOUT: winners by warp arrival ++ (a seeded permutation) in slots [0, TopK - fin); a partially kept pivot ++ group at [TopK - fin + rank] (rank by index in the tie direction when ++ TIE_DET, arrival order otherwise). `candidates` is the decode path's ++ padded range: n_vec * 4 elements, -inf past seq_len.""" ++ if candidates is not None: ++ keys = candidates ++ n = keys.size ++ pivot = np.sort(keys)[::-1][top_k - 1] ++ gt = np.nonzero(keys > pivot)[0] ++ eq = np.nonzero(keys == pivot)[0] ++ need = top_k - gt.size ++ out = np.full(top_k, -7, dtype=np.int64) # never survives: all slots written ++ if need == eq.size: # whole pivot group kept: no pass-3 tie cut ++ winners = np.concatenate([gt, eq]) ++ out[:] = rng.permutation(winners) ++ return out.astype(np.int32) ++ out[: top_k - need] = rng.permutation(gt) ++ if det_ties: ++ ranked = eq[::-1] if tie_recent else eq # rank 0 = first kept ++ out[top_k - need :] = ranked[:need] ++ else: # atomicAdd(&sFIN, -1): first `need` arrivals keep a slot ++ out[top_k - need :] = rng.permutation(eq)[:need] ++ assert n >= top_k ++ return out.astype(np.int32) ++ ++ ++def emit_restart(keys: np.ndarray, top_k: int, tie_recent: bool) -> np.ndarray: ++ """medium_exact_restart(): four 8-bit radix rounds, the two-round index ++ tie walk, then flag + exclusive prefix over 1,024-wide tiles in index ++ order. A port of the kernel's arithmetic, not a restatement of the rule.""" ++ length = keys.size ++ prefix, remaining, tie_total = 0, top_k, 0 ++ masks = [0, 0xFF000000, 0xFFFF0000, 0xFFFFFF00] ++ for rnd in range(4): ++ shift = 24 - 8 * rnd ++ sel = (keys & np.uint32(masks[rnd])) == np.uint32(prefix) ++ hist = np.bincount( ++ (keys[sel] >> np.uint32(shift)) & np.uint32(0xFF), minlength=256 ++ ) ++ count_above = 0 ++ for b in range(255, -1, -1): ++ if count_above + hist[b] >= remaining: ++ remaining -= count_above ++ prefix |= b << shift ++ if rnd == 3: ++ tie_total = int(hist[b]) ++ break ++ count_above += int(hist[b]) ++ else: ++ raise AssertionError("no crossing bin") ++ pivot, equal_count = np.uint32(prefix), remaining ++ tie_cut = 0 if tie_recent else 0xFFFFFFFF ++ if tie_total > equal_count: ++ eq_idx = np.nonzero(keys == pivot)[0] ++ need = equal_count ++ hist = np.bincount((eq_idx >> 8) & 0xFF, minlength=256) ++ order = range(255, -1, -1) if tie_recent else range(256) ++ count_above = 0 ++ for b in order: ++ if count_above + hist[b] >= need: ++ high, need = b, need - count_above ++ break ++ count_above += int(hist[b]) ++ in_high = eq_idx[((eq_idx >> 8) & 0xFF) == high] ++ hist = np.bincount(in_high & 0xFF, minlength=256) ++ count_above = 0 ++ for b in order: ++ if count_above + hist[b] >= need: ++ tie_cut = (high << 8) | b ++ break ++ count_above += int(hist[b]) ++ out = np.full(top_k, -9, dtype=np.int64) ++ base_rank = 0 ++ idx_all = np.arange(length) ++ keep_tie = idx_all >= tie_cut if tie_recent else idx_all <= tie_cut ++ flag = (keys > pivot) | ((keys == pivot) & keep_tie) ++ for base in range(0, -(-length // 1024) * 1024, 1024): ++ f = flag[base : base + 1024] ++ rank = np.cumsum(f) - f # BlockScan ExclusiveSum ++ for t in np.nonzero(f)[0]: ++ pos = base_rank + rank[t] ++ assert pos < top_k, "restart fault path" ++ out[pos] = base + t ++ base_rank += int(f.sum()) ++ assert base_rank == top_k ++ return out.astype(np.int32) ++ ++ ++def large_geometry(stride: int, num_rows: int) -> tuple[int, int]: ++ """launch_persistent_topk's ctas_per_group / chunk_size, 3090 opt-in.""" ++ if num_rows <= 4: ++ eff = min(SMEM_OPTIN_3090, K_SMEM_MEDIUM) ++ elif num_rows <= 8: ++ eff = min(SMEM_OPTIN_3090, 48 * 1024) ++ else: ++ eff = SMEM_OPTIN_3090 ++ max_chunk = (eff - K_FIXED_SMEM_LARGE) // 4 ++ vec = 4 if stride % 4 == 0 else (2 if stride % 2 == 0 else 1) ++ max_chunk = (max_chunk // vec) * vec ++ max_chunk = max(max_chunk, vec * 1024) ++ ctas = -(-stride // max_chunk) ++ chunk = -(-stride // ctas) ++ chunk = -(-chunk // vec) * vec ++ chunk = min(chunk, max_chunk) ++ return ctas, chunk ++ ++ ++def emit_large_det(keys, top_k, tie_recent, stride, num_rows) -> np.ndarray: ++ """radix_emit_deterministic(): per-CTA gt/eq counts, the directional CTA ++ walk handing out `equal_needed`, cut_rank/keep_width, and the slot ++ gt_base + eq_keep_base + gt_rank + kept_eq_before.""" ++ seq_len = keys.size ++ ctas, chunk = large_geometry(stride, num_rows) ++ assert 4 * ctas <= 256, "host gate would route to the legacy tail" ++ pivot = np.sort(keys)[::-1][top_k - 1] ++ equal_needed = top_k - int((keys > pivot).sum()) ++ chunks = [keys[c * chunk : min((c + 1) * chunk, seq_len)] for c in range(ctas)] ++ gt = [int((ch > pivot).sum()) for ch in chunks] ++ eq = [int((ch == pivot).sum()) for ch in chunks] ++ keep = [0] * ctas ++ seen = 0 ++ order = range(ctas - 1, -1, -1) if tie_recent else range(ctas) ++ for c in order: ++ want = max(equal_needed - seen, 0) ++ keep[c] = min(want, eq[c]) ++ seen += eq[c] ++ out = np.full(top_k, -11, dtype=np.int64) ++ for c, ch in enumerate(chunks): ++ gt_base, eq_keep_base = sum(gt[:c]), sum(keep[:c]) ++ cut_rank = (eq[c] - keep[c]) if tie_recent else 0 ++ f_gt = (ch > pivot).astype(np.int64) ++ f_eq = (ch == pivot).astype(np.int64) ++ gt_rank = np.cumsum(f_gt) - f_gt # exclusive, ascending element index ++ eq_rank = np.cumsum(f_eq) - f_eq ++ sel = (f_gt == 1) | ( ++ (f_eq == 1) & (eq_rank >= cut_rank) & (eq_rank - cut_rank < keep[c]) ++ ) ++ kept_before = np.clip(eq_rank - cut_rank, 0, keep[c]) ++ slot = gt_base + eq_keep_base + gt_rank + kept_before ++ assert (slot[sel] < top_k).all(), "X4 fault path" ++ out[slot[sel]] = c * chunk + np.nonzero(sel)[0] ++ return out.astype(np.int32) ++ ++ ++def emit_large_legacy(keys, top_k, rng) -> np.ndarray: ++ """v2.0.1 tail: > pivot by per-CTA arrival, then pivot-equal by global ++ arrival until TopK -- set racy on ties, order racy everywhere.""" ++ pivot = np.sort(keys)[::-1][top_k - 1] ++ gt = np.nonzero(keys > pivot)[0] ++ eq = np.nonzero(keys == pivot)[0] ++ need = top_k - gt.size ++ return np.concatenate([rng.permutation(gt), rng.permutation(eq)[:need]]).astype( ++ np.int32 ++ ) ++ ++ ++# ================================================ sorted_emit_row() ========= ++def sorted_emit_row( ++ row: np.ndarray, *, dup_check=True, fill=True, force_fast=False ++) -> np.ndarray: ++ """Port of persistent_topk.cuh sorted_emit_row(). The keyword switches ++ exist ONLY to build the failing controls.""" ++ top_k = row.size ++ vals = row.astype(np.int64) ++ bitmap = [0] * BITMAP_WORDS ++ fallback = False ++ for v in vals: # E1 (the atomicOr order cannot matter: OR is commutative) ++ if v >= 0: ++ if v < BITMAP_BITS: ++ bit = 1 << (int(v) & 31) ++ if bitmap[int(v) >> 5] & bit and dup_check: ++ fallback = True ++ bitmap[int(v) >> 5] |= bit ++ elif not force_fast: ++ fallback = True ++ out = row.astype(np.int64).copy() # slots the kernel does not write keep bytes ++ if not fallback: ++ pops = [bin(w).count("1") for w in bitmap] ++ excl = np.concatenate([[0], np.cumsum(pops)[:-1]]) ++ total = int(sum(pops)) ++ for tx, w in enumerate(bitmap): # one bitmap word per thread ++ pos = int(excl[tx]) ++ while w: ++ b = (w & -w).bit_length() - 1 # __ffs - 1 ++ if pos < top_k: ++ out[pos] = (tx << 5) | b ++ pos += 1 ++ w &= w - 1 ++ if fill: ++ out[total:] = -1 ++ else: ++ key = np.where(vals < 0, INT32_MAX, vals) ++ idx = np.arange(top_k) ++ less = (key[None, :] < key[:, None]).sum(axis=1) ++ tie_before = ( ++ (key[None, :] == key[:, None]) & (idx[None, :] < idx[:, None]) ++ ).sum(axis=1) ++ rank = less + tie_before ++ out[rank] = np.where(key == INT32_MAX, -1, key) ++ return out.astype(np.int32) ++ ++ ++# ================================================== the harness ============= ++def make_row(seq_len: int, kind: str, rng: np.random.Generator) -> np.ndarray: ++ if kind == "heavy_ties": # relu(q.k) sums: many exact zeros, few levels ++ levels = np.array([0.0, 0.0, 0.0, 0.5, 1.0, 2.0], dtype=np.float32) ++ return levels[rng.integers(0, levels.size, seq_len)] ++ if kind == "signed_zero": # +0.0 / -0.0 at the cutoff ++ x = np.where(rng.random(seq_len) < 0.5, np.float32(0.0), np.float32(-0.0)) ++ x = x.astype(np.float32) ++ hot = rng.random(seq_len) < 0.02 ++ x[hot] = rng.random(int(hot.sum())).astype(np.float32) + 1.0 ++ return x ++ if kind == "neg_inf": # masked blocks ++ x = rng.standard_normal(seq_len).astype(np.float32) ++ x[rng.random(seq_len) < 0.7] = -np.inf ++ return x ++ if kind == "all_equal": ++ return np.full(seq_len, 1.25, dtype=np.float32) ++ if kind == "distinct": ++ return rng.permutation(seq_len).astype(np.float32) * np.float32(1e-3) ++ if kind == "one_coarse_bin": # medium / decode overflow class ++ return (1.0 + rng.random(seq_len) * 0.01).astype(np.float32) ++ raise ValueError(kind) ++ ++ ++def emissions( ++ seq_len, ++ top_k, ++ kind, ++ det_ties, ++ tie_recent, ++ rng, ++ *, ++ num_rows=1, ++ exact_medium=True, ++ exact_decode=False, ++ det_large=True, ++): ++ """(path, native emission, re-emitted by sorted_emit_row?) for one row.""" ++ x = make_row(seq_len, kind, rng) ++ keys = ordered_key(x) ++ path = path_of(seq_len, top_k) ++ eff = effective_tie_recent(det_ties, tie_recent) ++ if path == "trivial": ++ return path, x, emit_trivial(seq_len, top_k), False ++ if path == "decode": ++ n_pad = -(-seq_len // 4) * 4 ++ cand = ordered_key( ++ np.concatenate([x, np.full(n_pad - seq_len, -np.inf, dtype=np.float32)]) ++ ) ++ top_bin_pop = _half_bin_population_at_threshold(cand, top_k) ++ if top_bin_pop > DBUF and exact_decode: ++ return "decode-restart", x, emit_restart(keys, top_k, eff), False ++ if top_bin_pop > DBUF: ++ return "decode-overflow", x, None, True # modelled separately ++ return ( ++ path, ++ x, ++ emit_buffered(keys, top_k, det_ties, tie_recent, rng, candidates=cand), ++ True, ++ ) ++ if path == "medium": ++ coarse_pop = _coarse_bin_population_at_threshold(keys, x, top_k) ++ if coarse_pop > MAX_BUFFERED_ITEMS and exact_medium: ++ return "medium-restart", x, emit_restart(keys, top_k, eff), False ++ if coarse_pop > MAX_BUFFERED_ITEMS: ++ return "medium-overflow", x, None, True ++ return path, x, emit_buffered(keys, top_k, det_ties, tie_recent, rng), True ++ ctas, _ = large_geometry(seq_len, num_rows) ++ if det_large and 4 * ctas <= 256: ++ return path, x, emit_large_det(keys, top_k, eff, seq_len, num_rows), False ++ return "large-legacy", x, emit_large_legacy(keys, top_k, rng), True ++ ++ ++def _half_key16(x: np.ndarray) -> np.ndarray: ++ h = x.astype(np.float16).view(np.uint16) ++ return np.where(h & 0x8000, ~h, h | 0x8000).astype(np.uint16) ++ ++ ++def _half_bin_population_at_threshold(cand_keys: np.ndarray, top_k: int) -> int: ++ x = cand_keys.copy() ++ # back to float for the half bin: invert ordered_key ++ neg = (x & np.uint32(0x80000000)) == 0 ++ bits = np.where(neg, ~x, x & np.uint32(0x7FFFFFFF)).astype(np.uint32) ++ f = bits.view(np.float32) ++ b = _half_key16(f) >> 5 ++ hist = np.bincount(b, minlength=2048) ++ suf = np.cumsum(hist[::-1])[::-1] ++ thr = int(np.nonzero(suf >= top_k)[0].max()) ++ return int(hist[thr]) if suf[thr] - hist[thr] < top_k else 0 ++ ++ ++def _coarse_bin_population_at_threshold(keys, x, top_k) -> int: ++ """histogram_256_topk's M1: threshold = the bin with suffix > TopK and ++ next suffix <= TopK (strict, unlike decode); a next suffix of exactly TopK ++ resolves the row without collecting, so nothing overflows.""" ++ b = _half_key16(x) >> 8 ++ hist = np.bincount(b, minlength=256) ++ suf = np.cumsum(hist[::-1])[::-1] ++ thr = int(np.nonzero(suf > top_k)[0].max()) ++ above = suf[thr] - hist[thr] ++ return int(hist[thr]) if above != top_k else 0 ++ ++ ++def assert_bytes(a: np.ndarray, b: np.ndarray, what: str): ++ a32 = np.ascontiguousarray(a, dtype=np.int32) ++ b32 = np.ascontiguousarray(b, dtype=np.int32) ++ assert a32.tobytes() == b32.tobytes(), what ++ ++ ++# Lengths in BLOCKS (QSA k=512 blocks, compress 4): 131,072 tokens = 32,768 ++# blocks (medium's last), 131,076 = 32,769 (large's first). ++BOUNDARY_LENGTHS = [ ++ 1, ++ 300, ++ 511, ++ 512, ++ 513, ++ 700, ++ 4096, ++ 8191, ++ 8192, ++ 8193, ++ 12000, ++ 32767, ++ 32768, ++ 32769, ++ 32770, ++ 33750, ++ 50000, ++] ++KINDS = [ ++ "heavy_ties", ++ "signed_zero", ++ "neg_inf", ++ "all_equal", ++ "distinct", ++ "one_coarse_bin", ++] ++TIE_CONFIGS = [(True, True), (True, False), (False, False)] # (TIE_DET, RECENT) ++ ++ ++def _cases(top_k): ++ for seq_len in BOUNDARY_LENGTHS: ++ for kind in KINDS: ++ if seq_len > 8192 and kind in ("distinct",) and seq_len % 3: ++ continue # keep the grid affordable ++ yield seq_len, kind ++ ++ ++@pytest.mark.parametrize("top_k", [512, 2048]) ++@pytest.mark.parametrize("det_ties,tie_recent", TIE_CONFIGS) ++@pytest.mark.parametrize( ++ "exact_decode,det_large", ++ [(False, True), (True, False)], # production / the two other arms ++) ++def test_sorted_emission_equals_python_sort_every_path( ++ top_k, det_ties, tie_recent, exact_decode, det_large ++): ++ """THE property: knob-on bytes == knob-off bytes after `_canonical`, and ++ the set is the exact top-k under the effective rule (deterministic ++ configurations) -- on every path at every boundary length.""" ++ rng = np.random.default_rng(1000 * top_k + 10 * det_ties + tie_recent) ++ seen_paths = set() ++ for seq_len, kind in _cases(top_k): ++ path, x, native, reemitted = emissions( ++ seq_len, ++ top_k, ++ kind, ++ det_ties, ++ tie_recent, ++ rng, ++ exact_decode=exact_decode, ++ det_large=det_large, ++ ) ++ seen_paths.add(path) ++ if native is None: ++ continue # overflow-without-repair rows: see the stale-slot test ++ knob_off_after_sort = canonical(native) ++ knob_on = sorted_emit_row(native) if reemitted else native ++ assert_bytes( ++ knob_on, knob_off_after_sort, f"{path} len={seq_len} kind={kind} k={top_k}" ++ ) ++ # set check against the independent reference ++ if path == "trivial": ++ continue ++ keys = ordered_key(x) ++ if path == "decode": ++ n_pad = -(-seq_len // 4) * 4 ++ keys = ordered_key( ++ np.concatenate([x, np.full(n_pad - seq_len, -np.inf, dtype=np.float32)]) ++ ) ++ racy = (path in ("decode", "medium") and not det_ties) or path == "large-legacy" ++ ref, pivot, need, n_eq = reference_select( ++ keys, top_k, effective_tie_recent(det_ties, tie_recent) ++ ) ++ got = np.sort(knob_on.astype(np.int64)) ++ if racy and need < n_eq: ++ # racy tie membership: only the value multiset is determined ++ assert np.array_equal(np.sort(keys[got]), np.sort(keys[ref])) ++ else: ++ assert np.array_equal(got, ref), f"{path} len={seq_len} kind={kind}" ++ want = {"trivial", "decode", "medium", "medium-restart"} ++ want |= {"decode-restart", "large-legacy"} if exact_decode else {"large"} ++ assert want <= seen_paths, seen_paths ++ ++ ++@pytest.mark.parametrize("det_ties,tie_recent", TIE_CONFIGS) ++def test_restart_and_large_det_are_already_canonical(det_ties, tie_recent): ++ """The paths the kernel does NOT re-emit must already be canonical, or the ++ knob would leave them unsorted once the Python sort is gone.""" ++ rng = np.random.default_rng(7) ++ eff = effective_tie_recent(det_ties, tie_recent) ++ for seq_len in (513, 8193, 20000, 32768): ++ for kind in ("heavy_ties", "all_equal", "one_coarse_bin", "signed_zero"): ++ keys = ordered_key(make_row(seq_len, kind, rng)) ++ out = emit_restart(keys, 512, eff) ++ assert_bytes(out, canonical(out), f"restart {seq_len} {kind}") ++ for seq_len, num_rows in ( ++ (32769, 1), ++ (33750, 4), ++ (50000, 8), ++ (65536, 32), ++ (131072, 1), ++ ): ++ for kind in ("heavy_ties", "all_equal", "signed_zero", "distinct"): ++ keys = ordered_key(make_row(seq_len, kind, rng)) ++ out = emit_large_det(keys, 512, eff, seq_len, num_rows) ++ assert_bytes(out, canonical(out), f"large {seq_len} {kind}") ++ ref, *_ = reference_select(keys, 512, eff) ++ assert np.array_equal(out.astype(np.int64), ref) ++ ++ ++def test_mixed_batch_rows_are_independent(): ++ """A launch mixes paths row by row; each row's re-emission reads only its ++ own TopK slots.""" ++ rng = np.random.default_rng(3) ++ lens = [100, 600, 8192, 8193, 32768, 32769, 40000, 512] ++ rows = [] ++ for L in lens: ++ path, x, native, reemit = emissions( ++ L, 512, "heavy_ties", True, True, rng, num_rows=len(lens) ++ ) ++ rows.append(sorted_emit_row(native) if reemit else native) ++ batch = np.stack(rows) ++ t = torch.from_numpy(batch.astype(np.int32)) ++ assert torch.equal(qsa_indexer._canonical(t), t) ++ ++ ++# ------------------------------------------------------ pathological rows ---- ++@pytest.mark.parametrize("top_k", [512, 1024, 2048]) ++def test_pathological_rows_take_the_fallback_and_still_match(top_k): ++ rng = np.random.default_rng(top_k) ++ cases = { ++ "duplicates": rng.integers(0, 64, top_k), ++ "stale_mixed": np.concatenate( ++ [ ++ rng.permutation(8192)[: top_k // 2], ++ rng.integers(-5, 40000, top_k - top_k // 2), ++ ] ++ ), ++ "beyond_bitmap": rng.integers(30000, 300000, top_k), ++ "int32_max": np.where( ++ rng.random(top_k) < 0.1, INT32_MAX, rng.permutation(top_k * 4)[:top_k] ++ ), ++ "odd_negatives": np.where( ++ rng.random(top_k) < 0.3, ++ rng.integers(-100, 0, top_k), ++ rng.permutation(top_k * 8)[:top_k], ++ ), ++ "all_padding": np.full(top_k, -1), ++ "int32_min": np.where( ++ rng.random(top_k) < 0.2, -(2**31), rng.permutation(top_k * 2)[:top_k] ++ ), ++ "exact_32768": np.concatenate([[32767, 32768], np.arange(top_k - 2)]), ++ } ++ for name, row in cases.items(): ++ row = row.astype(np.int32) ++ rng.shuffle(row) ++ assert_bytes(sorted_emit_row(row), canonical(row), name) ++ ++ ++@pytest.mark.parametrize("path", ["decode-overflow", "medium-overflow"]) ++def test_overflow_without_repair_keeps_stale_bytes_semantics(path): ++ """EXACT_DECODE=0 / EXACT_MEDIUM=0 overflow can leave slots the buffered ++ path never wrote this launch; they hold whatever the output tensor held. ++ The Python sort sorts those stale values too, and so must we.""" ++ rng = np.random.default_rng(11) ++ for trial in range(40): ++ written = rng.permutation(8192)[: rng.integers(100, 512)] ++ stale = rng.integers(-3, 70000, 512 - written.size) ++ row = np.concatenate([written, stale]).astype(np.int32) ++ rng.shuffle(row) ++ assert_bytes(sorted_emit_row(row), canonical(row), f"{path} #{trial}") ++ ++ ++# ------------------------------------------------------- failing controls ---- ++def _any_mismatch(fn, rows) -> bool: ++ for row in rows: ++ try: ++ if fn(row).tobytes() != canonical(row).astype(np.int32).tobytes(): ++ return True ++ except (IndexError, AssertionError): ++ return True ++ return False ++ ++ ++def test_failing_control_no_duplicate_detection(): ++ rng = np.random.default_rng(5) ++ rows = [rng.integers(0, 64, 512).astype(np.int32) for _ in range(4)] ++ assert _any_mismatch(lambda r: sorted_emit_row(r, dup_check=False), rows) ++ assert not _any_mismatch(sorted_emit_row, rows) ++ ++ ++def test_failing_control_no_padding_fill(): ++ rng = np.random.default_rng(6) ++ rows = [ ++ np.where(rng.random(512) < 0.3, -1, rng.permutation(4096)[:512]).astype( ++ np.int32 ++ ) ++ for _ in range(4) ++ ] ++ assert _any_mismatch(lambda r: sorted_emit_row(r, fill=False), rows) ++ assert not _any_mismatch(sorted_emit_row, rows) ++ ++ ++def test_failing_control_sort_skipped(): ++ """Knob off WITHOUT the Python sort must be caught on buffered paths.""" ++ rng = np.random.default_rng(8) ++ rows = [] ++ for L in (600, 4096, 8193, 20000): ++ _, _, native, reemit = emissions(L, 512, "heavy_ties", True, True, rng) ++ assert reemit ++ rows.append(native) ++ assert _any_mismatch(lambda r: r, rows) ++ ++ ++def test_failing_control_out_of_range_forced_fast(): ++ rng = np.random.default_rng(9) ++ rows = [rng.integers(30000, 300000, 512).astype(np.int32) for _ in range(3)] ++ assert _any_mismatch(lambda r: sorted_emit_row(r, force_fast=True), rows) ++ ++ ++def test_failing_control_tie_rule_flipped(): ++ """The set check fails if the kernel kept the OTHER end of the tie group ++ (upstream #55122's lowest-index rule) where production keeps the highest.""" ++ rng = np.random.default_rng(10) ++ keys = ordered_key(make_row(40000, "all_equal", rng)) ++ ours, *_ = reference_select(keys, 512, True) ++ flipped = emit_large_det(keys, 512, False, 40000, 1) ++ assert not np.array_equal(np.sort(flipped.astype(np.int64)), ours) ++ keys = ordered_key(make_row(8193, "all_equal", rng)) ++ flipped = emit_restart(keys, 512, False) ++ ours, *_ = reference_select(keys, 512, True) ++ assert not np.array_equal(np.sort(flipped.astype(np.int64)), ours) ++ ++ ++def test_signed_zero_canonicalisation_would_change_the_set(): ++ """Why #55122's +-0.0 change is NOT adopted: with the pivot at zero, +0.0 ++ currently outranks -0.0 regardless of index; treating them as a tie hands ++ the choice to the index rule and changes the SET.""" ++ # +0.0 at the OLD end, -0.0 at the recent end, the cutoff inside zero. ++ x = np.array([0.0] * 100 + [-0.0] * 600 + [1.0] * 400, dtype=np.float32) ++ ours, *_ = reference_select(ordered_key(x), 512, True) ++ canon = ordered_key(np.where(x == 0, np.float32(0.0), x)) # -0 -> +0 ++ theirs, *_ = reference_select(canon, 512, True) ++ assert not np.array_equal(ours, theirs) ++ assert set(range(0, 100)) <= set(ours.tolist()) # all +0.0 kept today ++ assert not set(range(0, 100)) & set(theirs.tolist()) # all lost after ++ ++ ++# ------------------------------------------------------------ Python gate ---- ++def _pack(fields, version, count): ++ s = 0 ++ for i, f in enumerate(fields): ++ s |= f << (2 * i) ++ return s | (count << 48) | (version << 56) ++ ++ ++@pytest.fixture ++def fresh_gate(monkeypatch): ++ monkeypatch.setattr(qsa_indexer, "_kernel_sorted_emit_state", None) ++ yield ++ qsa_indexer._kernel_sorted_emit_state = None ++ ++ ++class _FakeOps: ++ def __init__(self, state): ++ self.calls = 0 ++ self._state = state ++ ++ def persistent_topk_toggle_state(self): ++ self.calls += 1 ++ return self._state() if callable(self._state) else self._state ++ ++ ++def _install(monkeypatch, fake): ++ monkeypatch.setattr(torch.ops, "_C", fake, raising=False) ++ ++ ++def test_gate_reads_field_six_and_latches(fresh_gate, monkeypatch): ++ fake = _FakeOps(_pack([1, 1, 1, 1, 0, 1, 1], 2, 7)) ++ _install(monkeypatch, fake) ++ assert qsa_indexer.kernel_sorted_emit() is True ++ assert qsa_indexer.kernel_sorted_emit() is True ++ assert fake.calls == 1 # latched: one host query per process ++ ++ ++def test_gate_off_and_layout_one(fresh_gate, monkeypatch): ++ _install(monkeypatch, _FakeOps(_pack([1, 1, 1, 1, 0, 1, 0], 2, 7))) ++ assert qsa_indexer.kernel_sorted_emit() is False ++ qsa_indexer._kernel_sorted_emit_state = None ++ _install(monkeypatch, _FakeOps(_pack([1, 1, 1, 1, 0, 1], 1, 6))) ++ assert qsa_indexer.kernel_sorted_emit() is False # old build: cannot sort ++ ++ ++def test_gate_does_not_latch_uninit(fresh_gate, monkeypatch): ++ states = iter([_pack([2] * 7, 2, 7), _pack([1, 1, 1, 1, 0, 1, 1], 2, 7)]) ++ fake = _FakeOps(lambda: next(states)) ++ _install(monkeypatch, fake) ++ assert qsa_indexer.kernel_sorted_emit() is False ++ assert qsa_indexer._kernel_sorted_emit_state is None ++ assert qsa_indexer.kernel_sorted_emit() is True ++ ++ ++def test_gate_op_absent(fresh_gate, monkeypatch): ++ monkeypatch.setattr(torch.ops, "_C", object(), raising=False) ++ assert qsa_indexer.kernel_sorted_emit() is False ++ ++ ++class _TopkOps(_FakeOps): ++ """Fake _C: persistent_topk writes a fixed UNSORTED row.""" ++ ++ def __init__(self, state, row): ++ super().__init__(state) ++ self.row = row ++ ++ def persistent_topk(self, logits, lengths, out, ws, k, maxlen): ++ out.copy_(self.row.expand_as(out)) ++ ++ cooperative_topk = persistent_topk ++ ++ ++def _run_topk(monkeypatch, state, sort_on, row): ++ fake = _TopkOps(state, row) ++ _install(monkeypatch, fake) ++ monkeypatch.setattr(qsa_indexer, "_QSA_SORT", sort_on) ++ out = torch.empty(2, row.numel(), dtype=torch.int32) ++ qsa_indexer._topk( ++ torch.zeros(2, 64), ++ torch.zeros(2, dtype=torch.int32), ++ 4 * row.numel(), ++ 4, ++ out, ++ torch.empty(8, dtype=torch.uint8), ++ ) ++ return out ++ ++ ++def test_topk_skips_python_sort_only_when_kernel_says_sorted(fresh_gate, monkeypatch): ++ row = torch.tensor([5, -1, 3, 9, 0, -1, 7, 1], dtype=torch.int32) ++ on = _pack([1, 1, 1, 1, 0, 1, 1], 2, 7) ++ off = _pack([1, 1, 1, 1, 0, 1, 0], 2, 7) ++ # knob on: the kernel's bytes pass through untouched (the sort is skipped) ++ assert torch.equal(_run_topk(monkeypatch, on, True, row)[0], row) ++ qsa_indexer._kernel_sorted_emit_state = None ++ # knob off + QSA_SORT: Python sorts ++ got = _run_topk(monkeypatch, off, True, row)[0] ++ assert torch.equal(got, qsa_indexer._canonical(row.view(1, -1))[0]) ++ ++ ++def test_debug_check_asserts_on_unsorted(monkeypatch): ++ monkeypatch.setattr(qsa_indexer, "_SORTED_EMIT_CHECK", True) ++ ok = torch.tensor([[0, 3, 7, -1, -1]], dtype=torch.int32) ++ qsa_indexer._check_sorted_emission(ok) # canonical: passes ++ for bad in ([[3, 0, 7, -1, -1]], [[0, 3, -1, 7, -1]], [[0, 3, 7, -2, -1]]): ++ with pytest.raises(RuntimeError): ++ qsa_indexer._check_sorted_emission(torch.tensor(bad, dtype=torch.int32)) ++ ++ ++# ===================================================== GPU only ============== ++_ARM = Path(__file__).with_name("merlin_sorted_emit_arm.py") ++_PROD = { ++ "MERLIN_TIE_DET": "1", ++ "MERLIN_TIE_RECENT": "1", ++ "MERLIN_TOPK_EXACT_MEDIUM": "1", ++ "MERLIN_TOPK_DET_LARGE": "1", ++ "MERLIN_TOPK_EXACT_DECODE": "0", ++} ++ ++ ++def _run_arm(tmp_path, name, env_over, extra_args=()): ++ out = tmp_path / f"{name}.pt" ++ env = {k: v for k, v in os.environ.items() if not k.startswith("MERLIN_")} ++ env.update(env_over) ++ proc = subprocess.run( ++ [sys.executable, str(_ARM), str(out), *extra_args], ++ env=env, ++ capture_output=True, ++ text=True, ++ timeout=1800, ++ ) ++ assert proc.returncode == 0, proc.stdout[-4000:] + proc.stderr[-4000:] ++ return torch.load(out, weights_only=True), proc.stderr ++ ++ ++needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU only") ++ ++ ++@needs_cuda ++def test_gpu_knob_on_equals_knob_off_plus_python_sort(tmp_path): ++ off, _ = _run_arm(tmp_path, "off", dict(_PROD, MERLIN_TOPK_SORTED_EMIT="0")) ++ on, err = _run_arm(tmp_path, "on", dict(_PROD, MERLIN_TOPK_SORTED_EMIT="1")) ++ assert "MERLIN_TOPK_SORTED_EMIT=1" in err and "ascending-in-kernel" in err ++ assert on["meta"]["toggle_sorted_emit"] == "1" ++ assert off["meta"]["toggle_sorted_emit"] == "0" ++ assert on["meta"]["counters"][3] == 0 and on["meta"]["counters"][7] == 0 ++ differs_unsorted = 0 ++ for cell, raw_off in off["cells"].items(): ++ raw_on = on["cells"][cell] ++ assert torch.equal(raw_on, qsa_indexer._canonical(raw_on)), cell ++ if cell.startswith("racy_"): ++ continue # launch-order-dependent SET in this arm (see the arm) ++ want = qsa_indexer._canonical(raw_off) ++ assert raw_on.dtype == torch.int32 and want.dtype == torch.int32 ++ assert raw_on.numpy().tobytes() == want.numpy().tobytes(), cell ++ differs_unsorted += int(not torch.equal(raw_off, want)) ++ # FAILING CONTROL (sort skipped): the knob-off raw bytes are NOT canonical ++ # on the buffered cells, so the equality above is not vacuous. ++ assert differs_unsorted > 0 ++ ++ ++@needs_cuda ++def test_gpu_failing_control_tie_rule_flip(tmp_path): ++ ref, _ = _run_arm(tmp_path, "ref", dict(_PROD, MERLIN_TOPK_SORTED_EMIT="1")) ++ flip, _ = _run_arm( ++ tmp_path, ++ "flip", ++ dict(_PROD, MERLIN_TOPK_SORTED_EMIT="1", MERLIN_TIE_RECENT="0"), ++ ) ++ tie_cells = [c for c in ref["cells"] if "tie" in c or "all_equal" in c] ++ assert tie_cells ++ assert any(not torch.equal(ref["cells"][c], flip["cells"][c]) for c in tie_cells) ++ ++ ++@needs_cuda ++@pytest.mark.parametrize( ++ "arm", ++ [ ++ {"MERLIN_TOPK_DET_LARGE": "0"}, # legacy large tail: re-emitted ++ {"MERLIN_TIE_DET": "0"}, # racy buffered ties ++ {"MERLIN_TOPK_EXACT_MEDIUM": "0"}, # medium overflow unrepaired ++ {"MERLIN_TOPK_EXACT_DECODE": "1"}, # decode overflow restarted ++ ], ++) ++def test_gpu_non_production_arms_are_canonical(tmp_path, arm): ++ """Where the SET is racy the two arms may legitimately disagree, so the ++ claim is only: the knob-on row is already canonical.""" ++ on, _ = _run_arm(tmp_path, "arm", dict(_PROD, MERLIN_TOPK_SORTED_EMIT="1", **arm)) ++ for cell, raw in on["cells"].items(): ++ assert torch.equal(raw, qsa_indexer._canonical(raw)), (arm, cell) ++ ++ ++@needs_cuda ++def test_gpu_cudagraph_replay_knob_on(tmp_path): ++ on, _ = _run_arm( ++ tmp_path, ++ "graph", ++ dict(_PROD, MERLIN_TOPK_SORTED_EMIT="1"), ++ extra_args=("--graph",), ++ ) ++ for cell, raw in on["cells"].items(): ++ assert torch.equal(raw, qsa_indexer._canonical(raw)), cell ++ assert on["meta"]["graph_replays_match_eager"], json.dumps(on["meta"])[:2000] +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 0db9be3..9ddb51e 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -660,13 +660,19 @@ _FORK_IGNORED = [ + _UNAUDITED_START = "# --8<-- [start:unaudited-host-side-knobs]" + _UNAUDITED_END = "# --8<-- [end:unaudited-host-side-knobs]" + +-# All four are compile/capture factors and must be hashed. Three are raw +-# passthroughs of a value that C++ parses; MERLIN_QSA_SORT is parsed here. ++# All of these are compile/capture factors and must be hashed. Four are raw ++# passthroughs of a value that C++ parses; MERLIN_QSA_SORT and ++# MERLIN_TOPK_SORTED_EMIT_CHECK are parsed here. + _TIEDET_HASHED = [ + ("MERLIN_FULL_K", "1", "1"), + ("MERLIN_TIE_DET", "1", "1"), + ("MERLIN_TIE_RECENT", "1", "1"), + ("MERLIN_QSA_SORT", "1", True), ++ # In-kernel canonical emission: raw passthrough of the C++ ++ # getter (topk.cu sorted_emit_enabled, default OFF), and its Python-parsed ++ # debug check. Both change the captured sequence, so both are hashed. ++ ("MERLIN_TOPK_SORTED_EMIT", "1", "1"), ++ ("MERLIN_TOPK_SORTED_EMIT_CHECK", "1", True), + ] + + # The persistent_topk kernel switches. Raw passthroughs +diff --git a/tests/test_merlin_e2_toggles.py b/tests/test_merlin_e2_toggles.py +index 2ded113..10d465f 100644 +--- a/tests/test_merlin_e2_toggles.py ++++ b/tests/test_merlin_e2_toggles.py +@@ -162,6 +162,7 @@ ON_MAP = { + "exact_medium": "1", + "exact_decode": "0", + "det_large": "1", ++ "sorted_emit": "0", + } + ABSENT_MAP = { + "full_k": "0", +@@ -172,6 +173,7 @@ ABSENT_MAP = { + "exact_medium": "1", + "exact_decode": "0", + "det_large": "1", ++ "sorted_emit": "0", + } + + +@@ -209,6 +211,17 @@ def test_effective_map_odd_spellings(): + assert eff["qsa_sort"] == "invalid" + + ++def test_effective_map_sorted_emit(): ++ """MERLIN_TOPK_SORTED_EMIT is a default-OFF C++ switch (topk.cu ++ sorted_emit_enabled): atoi rules, empty = unset = OFF.""" ++ for raw, want in (("1", "1"), ("0", "0"), ("", "0"), ("true", "0"), (" 1", "1")): ++ raw_map, eff = t.environ_view(*_env(dict(ON, MERLIN_TOPK_SORTED_EMIT=raw))) ++ assert eff["sorted_emit"] == want, raw ++ raw_map, eff = t.environ_view(*_env(ON)) ++ assert raw_map["MERLIN_TOPK_SORTED_EMIT"] == "unset" ++ assert eff["sorted_emit"] == "0" ++ ++ + def test_encode_raw(): + assert t.encode_raw(None) == "unset" + assert t.encode_raw(b"1") == "1" +@@ -229,6 +242,7 @@ def _pack(fields, version=1, count=6): + + + def test_decode_topk_state(): ++ # Layout 1 (builds before sorted emission): six fields, no sorted emission. + assert t.decode_topk_state(_pack([1, 1, 1, 1, 0, 1])) == { + "tie_det": "1", + "tie_recent": "1", +@@ -236,7 +250,16 @@ def test_decode_topk_state(): + "exact_medium": "1", + "exact_decode": "0", + "det_large": "1", ++ "sorted_emit": "n-a", + } ++ # Layout 2: field 6 = MERLIN_TOPK_SORTED_EMIT. ++ v2 = t.decode_topk_state(_pack([1, 1, 1, 1, 0, 1, 1], version=2, count=7)) ++ assert v2["sorted_emit"] == "1" and v2["det_large"] == "1" ++ v2 = t.decode_topk_state(_pack([1, 1, 1, 1, 0, 1, 2], version=2, count=7)) ++ assert v2["sorted_emit"] == "uninit" ++ assert isinstance(t.decode_topk_state(_pack([0] * 7, version=1, count=7)), str) ++ assert isinstance(t.decode_topk_state(_pack([0] * 8, version=2, count=8)), str) ++ assert isinstance(t.decode_topk_state(_pack([0] * 7, version=3, count=7)), str) + assert t.decode_topk_state(_pack([2, 2, 2, 1, 2, 2]))["tie_det"] == "uninit" + assert t.decode_topk_state(_pack([3, 0, 0, 0, 0, 0]))["tie_det"] == "invalid" + assert isinstance(t.decode_topk_state(_pack([0] * 6, version=2)), str) +@@ -279,13 +302,14 @@ def capture(): + t._reset_for_tests() + + +-def _kw(cpp=CPP_ON, env=None, device=None, qsa="1", tag="PP1_TP0"): ++def _kw(cpp=CPP_ON, env=None, device=None, qsa="1", tag="PP1_TP0", skip="0"): + return dict( + tag_fn=lambda: tag, + getenv_fn=lambda: (_env(ON if env is None else env)[0], "libc-getenv"), + cpp_fn=lambda: cpp, + device_fn=lambda: device or {"switch": "1", "switch_large": "1"}, + qsa_fn=lambda: qsa, ++ sorted_skip_fn=lambda: skip, + ) + + +@@ -345,7 +369,11 @@ def test_line_contents_toggles_on(capture, monkeypatch): + cpp = dict(p["cpp-cache"]) + assert cpp.pop("source") == "cpp-cache" + assert cpp == CPP_ON +- assert p["python"] == {"qsa_sort": "1", "source": "qsa_indexer._QSA_SORT"} ++ assert p["python"] == { ++ "qsa_sort": "1", ++ "sorted_emit_skip": "0", ++ "source": "qsa_indexer._QSA_SORT", ++ } + assert p["device"] == { + "switch": "1", + "switch_large": "1", +@@ -363,6 +391,20 @@ def test_line_knob_absent(capture, monkeypatch): + assert eff == ABSENT_MAP + + ++def test_line_sorted_emit_on(capture, monkeypatch): ++ """The knob-on arm the reference host reads: raw, effective, cpp-cache and the ++ Python skip all say sorted emission.""" ++ env = dict(ON, MERLIN_TOPK_SORTED_EMIT="1") ++ for k, v in env.items(): ++ monkeypatch.setenv(k, v) ++ cpp = dict(CPP_ON, sorted_emit="1") ++ p = parse(t.arm(**_kw(cpp=cpp, env=env, skip="1"))) ++ assert p["raw"]["MERLIN_TOPK_SORTED_EMIT"] == "1" ++ assert p["effective"]["sorted_emit"] == "1" ++ assert p["cpp-cache"]["sorted_emit"] == "1" ++ assert p["python"]["sorted_emit_skip"] == "1" ++ ++ + def test_cpp_op_absent_is_n_a_and_emits_at_once(capture, monkeypatch): + for k, v in ON.items(): + monkeypatch.setenv(k, v) +@@ -413,6 +455,7 @@ def test_default_sources_do_not_raise_on_cpu(monkeypatch): + assert isinstance(t.cpp_cache(), (str, dict)) + assert set(t.device_switches()) == {"switch", "switch_large"} + assert t.qsa_sort_live() in ("0", "1", "n-a") ++ assert t.sorted_emit_skip_live() in ("0", "1", "n-a") + read, source = t._libc_getenv() + assert source in ("libc-getenv", "os-environ") + monkeypatch.setenv("MERLIN_TIE_DET", "7") +diff --git a/vllm/envs.py b/vllm/envs.py +index 0dca963..f13f442 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -597,6 +597,26 @@ if TYPE_CHECKING: + # unless MERLIN_TOPK_LARGE_FAULT_INJECT names a CTA. + MERLIN_TOPK_LARGE_FAULT_KIND: str | None = None + # ++ # MERLIN_TOPK_SORTED_EMIT -- persistent_topk ++ # emits every row in the canonical order MERLIN_QSA_SORT imposes ++ # (ascending index, padding last as -1), and qsa_indexer.py then skips its ++ # Python sort. Same classification as the kernel switches above: resolved ++ # ONCE per process by a std::getenv in csrc/libtorch_stable/topk.cu ++ # (sorted_emit_enabled) -- the single source of truth -- and the resolved ++ # value reaches the GPU as a kernel launch argument AND removes three torch ++ # ops from the captured sequence, so a cudagraph bakes it in. CAPTURE: yes. ++ # AOT: not established for the kernel half; hashed conservatively (a fork ++ # env knob must be a compile factor: an os.environ knob that is not ++ # registered is not in the cache key). Raw passthrough: Python never parses ++ # it to decide anything -- qsa_indexer.py reads the C++ cache through ++ # persistent_topk_toggle_state() field 6. Not in ignored_factors. ++ MERLIN_TOPK_SORTED_EMIT: str | None = None ++ # Debug companion, a REAL Python read site (qsa_indexer._SORTED_EMIT_CHECK, ++ # read once at import): when the kernel sorted, assert on device that the ++ # row equals its canonical form instead of trusting it. Adds ops to the ++ # captured sequence, so it is hashed too. Default off. ++ MERLIN_TOPK_SORTED_EMIT_CHECK: bool = False ++ # + # NOT REGISTERED, deliberately: MERLIN_TOPK_EXACT_TRAP. It is a + # preprocessor define consumed by `#ifndef` in persistent_topk.cuh at + # BUILD time (-D), not a process environment variable -- nothing in the +@@ -2613,6 +2633,12 @@ environment_variables: dict[str, Callable[[], Any]] = { + "MERLIN_TOPK_LARGE_FAULT_KIND": lambda: os.getenv( + "MERLIN_TOPK_LARGE_FAULT_KIND", None + ), ++ # In-kernel canonical emission and its debug check; see the ++ # declaration block. Both hashed: both change the captured sequence. ++ "MERLIN_TOPK_SORTED_EMIT": lambda: os.getenv("MERLIN_TOPK_SORTED_EMIT", None), ++ "MERLIN_TOPK_SORTED_EMIT_CHECK": lambda: bool( ++ int(os.getenv("MERLIN_TOPK_SORTED_EMIT_CHECK", "0")) ++ ), + } + + +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +index 8e5da35..b9e88b0 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +@@ -50,17 +50,84 @@ _DECODE_BLOCK_N = 64 + # --------------------------------------------------------------------------- + _QSA_SORT = envs.MERLIN_QSA_SORT + ++# --------------------------------------------------------------------------- ++# MERLIN_TOPK_SORTED_EMIT: persistent_topk emits ++# every row in the canonical order below by itself (sorted_emit_row() in ++# csrc/libtorch_stable/persistent_topk.cuh), so the Python sort is skipped. ++# ++# The decision is NOT a Python parse of the variable. The C++ getter ++# (topk.cu sorted_emit_enabled) is the single source of truth and publishes ++# what it cached into persistent_topk_toggle_state() field 6; that is read ++# once, host-side, right after this process's first persistent_topk launch ++# (the static is filled by then) and latched. So the sort is dropped only when ++# the kernel that actually ran says it sorted; an older build (layout 1, no ++# field 6) or an absent op keeps the Python sort. cooperative_topk (SM90+) has ++# no sorted emission and always keeps the Python sort. ++# ++# MERLIN_TOPK_SORTED_EMIT_CHECK=1 (debug, default off): instead of trusting ++# the kernel, assert on device that the row already equals its canonical form. ++# torch._assert_async is capture-safe (no host sync); a violation is a device ++# assert, i.e. a dead context -- which is the point of a debug check. ++# --------------------------------------------------------------------------- ++_SORTED_EMIT_CHECK = envs.MERLIN_TOPK_SORTED_EMIT_CHECK ++_kernel_sorted_emit_state: bool | None = None # None: not read yet + +-def _qsa_canonicalise(block_indices: torch.Tensor) -> None: +- """Sort selected block indices ascending, padding last. In place, same shape.""" +- if not _QSA_SORT: +- return ++ ++def _canonical(block_indices: torch.Tensor) -> torch.Tensor: ++ """Ascending, padding last as -1 -- THE definition of the canonical order, ++ shared by the Python sort and the debug check (and the reference the ++ kernel's sorted emission is tested against).""" + big = torch.iinfo(block_indices.dtype).max + tmp = torch.where( + block_indices < 0, torch.full_like(block_indices, big), block_indices + ) + tmp = torch.sort(tmp, dim=-1).values +- block_indices.copy_(torch.where(tmp == big, torch.full_like(tmp, -1), tmp)) ++ return torch.where(tmp == big, torch.full_like(tmp, -1), tmp) ++ ++ ++def _qsa_canonicalise(block_indices: torch.Tensor) -> None: ++ """Sort selected block indices ascending, padding last. In place, same shape.""" ++ if not _QSA_SORT: ++ return ++ block_indices.copy_(_canonical(block_indices)) ++ ++ ++def kernel_sorted_emit() -> bool: ++ """True iff this process's persistent_topk emits canonical order. ++ ++ Host-only (a tensor-free query op; no CUDA call), so it is safe inside a ++ capture. Latches once the C++ static is known; an uninitialised static ++ (no launch yet) answers False without latching. ++ """ ++ global _kernel_sorted_emit_state ++ if _kernel_sorted_emit_state is not None: ++ return _kernel_sorted_emit_state ++ op = getattr(torch.ops._C, "persistent_topk_toggle_state", None) ++ if op is None: ++ _kernel_sorted_emit_state = False ++ return False ++ from vllm.v1.worker.e2_toggles import decode_topk_state ++ ++ decoded = decode_topk_state(int(op())) ++ if isinstance(decoded, str) or "sorted_emit" not in decoded: ++ # Unknown layout, or layout 1 (no field 6): this build cannot sort. ++ _kernel_sorted_emit_state = False ++ return False ++ value = decoded["sorted_emit"] ++ if value == "uninit": ++ return False # not latched: ask again after a launch ++ _kernel_sorted_emit_state = value == "1" ++ return _kernel_sorted_emit_state ++ ++ ++def _check_sorted_emission(block_indices: torch.Tensor) -> None: ++ """Debug only: device-side assert that the kernel's row is canonical.""" ++ if not _SORTED_EMIT_CHECK: ++ return ++ torch._assert_async( ++ torch.eq(block_indices, _canonical(block_indices)).all(), ++ "MERLIN_TOPK_SORTED_EMIT: persistent_topk row is not in canonical order", ++ ) + + + @triton.jit +@@ -605,6 +672,11 @@ def _topk( + block_topk, + logits.shape[1], + ) ++ # MERLIN_TOPK_SORTED_EMIT: the kernel already emitted the canonical order, ++ # so the Python sort would be a no-op -- skip it (optionally check it). ++ if not use_cooperative_topk and kernel_sorted_emit(): ++ _check_sorted_emission(block_indices) ++ return + _qsa_canonicalise(block_indices) + + +diff --git a/vllm/v1/worker/e2_toggles.py b/vllm/v1/worker/e2_toggles.py +index 0ae7ad2..056a905 100644 +--- a/vllm/v1/worker/e2_toggles.py ++++ b/vllm/v1/worker/e2_toggles.py +@@ -12,13 +12,15 @@ producer at all. Every GPU worker therefore logs exactly one line:: + e2-toggles[PP

_TP]: pid= MERLIN_FULL_K= MERLIN_QSA_SORT= + MERLIN_TIE_DET= MERLIN_TIE_RECENT= MERLIN_TOPK_EXACT_MEDIUM= + MERLIN_TOPK_EXACT_DECODE= MERLIN_TOPK_DET_LARGE= ++ MERLIN_TOPK_SORTED_EMIT= + effective: full_k=<0|1> qsa_sort=<0|1|invalid> tie_det=<0|1> tie_recent=<0|1> + tie_recent_effective=<0|1> exact_medium=<0|1> exact_decode=<0|1> +- det_large=<0|1> source= ++ det_large=<0|1> sorted_emit=<0|1> source= + cpp-cache: > +- python: qsa_sort=<0|1|n-a> source=qsa_indexer._QSA_SORT ++ exact_medium=... exact_decode=... det_large=... sorted_emit=... ++ source=cpp-cache | n-a reason=> ++ python: qsa_sort=<0|1|n-a> sorted_emit_skip=<0|1|n-a> ++ source=qsa_indexer._QSA_SORT + device: switch=<0|1|n-a> switch_large=<0|1|n-a> source=counter-slots-8-9 + + (one line in the log). ```` is ``unset``, a plain token, or a double-quoted +@@ -42,7 +44,11 @@ WHAT EACH SECTION IS + was first cut against (earlier v2.2 builds, not published) have none. + * ``python:`` -- ``qsa_indexer._QSA_SORT``, the module constant the QSA sort + actually branches on (looked up in ``sys.modules``, never imported here; +- ``n-a`` if the module was never loaded in this process). ++ ``n-a`` if the module was never loaded in this process), and ++ ``sorted_emit_skip`` -- ``qsa_indexer._kernel_sorted_emit_state``, i.e. ++ whether the Python sort is being SKIPPED because the kernel reported ++ MERLIN_TOPK_SORTED_EMIT through ``persistent_topk_toggle_state`` field 6 ++ (``n-a`` until the first QSA top-k has run and latched it). + * ``device:`` -- counter slots [8]/[9], which the counter op fills from the + C++ ``exact_medium_enabled()`` / ``det_large_enabled()`` statics + (topk.cu persistent_topk_exact_counters). +@@ -93,6 +99,7 @@ logger = init_logger(__name__) + # MERLIN_TOPK_EXACT_MEDIUM csrc/libtorch_stable/topk.cu:68-74 (unset -> ON) + # MERLIN_TOPK_EXACT_DECODE csrc/libtorch_stable/topk.cu:103-109 + # MERLIN_TOPK_DET_LARGE csrc/libtorch_stable/topk.cu:86-92 (unset -> ON) ++# MERLIN_TOPK_SORTED_EMIT csrc/libtorch_stable/topk.cu sorted_emit_enabled + CXX_SWITCHES: tuple[tuple[str, str, bool], ...] = ( + ("MERLIN_FULL_K", "full_k", False), + ("MERLIN_TIE_DET", "tie_det", False), +@@ -100,6 +107,7 @@ CXX_SWITCHES: tuple[tuple[str, str, bool], ...] = ( + ("MERLIN_TOPK_EXACT_MEDIUM", "exact_medium", True), + ("MERLIN_TOPK_EXACT_DECODE", "exact_decode", False), + ("MERLIN_TOPK_DET_LARGE", "det_large", True), ++ ("MERLIN_TOPK_SORTED_EMIT", "sorted_emit", False), + ) + QSA_SORT_NAME = "MERLIN_QSA_SORT" + RAW_ORDER = ( +@@ -110,6 +118,7 @@ RAW_ORDER = ( + "MERLIN_TOPK_EXACT_MEDIUM", + "MERLIN_TOPK_EXACT_DECODE", + "MERLIN_TOPK_DET_LARGE", ++ "MERLIN_TOPK_SORTED_EMIT", + ) + EFFECTIVE_ORDER = ( + "full_k", +@@ -120,6 +129,7 @@ EFFECTIVE_ORDER = ( + "exact_medium", + "exact_decode", + "det_large", ++ "sorted_emit", + ) + CPP_ORDER = tuple(k for k in EFFECTIVE_ORDER if k != "qsa_sort") + +@@ -265,21 +275,29 @@ _TOPK_STATE_FIELDS = ( + "exact_medium", + "exact_decode", + "det_large", ++ "sorted_emit", + ) ++# Known layouts: version -> field count. Layout 2 is layout 1 ++# plus field 6 (MERLIN_TOPK_SORTED_EMIT); fields 0-5 are unchanged. ++_TOPK_STATE_LAYOUTS = {1: 6, 2: 7} + + + def decode_topk_state(state: int) -> dict[str, str] | str: + """Unpack ``persistent_topk_toggle_state()``: 2 bits per field (0 off, +- 1 on, 2 not initialised), field count at bits 48-55 (6), layout version at +- bits 56-62 (1). Any other layout is refused with a reason string.""" ++ 1 on, 2 not initialised), field count at bits 48-55, layout version at ++ bits 56-62 -- (1, 6) or (2, 7). Any other layout is refused with a reason ++ string. A layout-1 build has no sorted emission, reported as ``n-a``.""" + version, count = (state >> 56) & 0x7F, (state >> 48) & 0xFF +- if version != 1 or count != len(_TOPK_STATE_FIELDS): ++ if _TOPK_STATE_LAYOUTS.get(version) != count: + return f"unknown-layout-v{version}-{count}-fields" + names = {0: "0", 1: "1", 2: "uninit"} +- return { ++ out = { + key: names.get((state >> (2 * i)) & 3, "invalid") +- for i, key in enumerate(_TOPK_STATE_FIELDS) ++ for i, key in enumerate(_TOPK_STATE_FIELDS[:count]) + } ++ for key in _TOPK_STATE_FIELDS[count:]: ++ out[key] = "n-a" ++ return out + + + def decode_full_k_state(state: int) -> str: +@@ -324,6 +342,17 @@ def qsa_sort_live() -> str: + return "1" if bool(module._QSA_SORT) else "0" + + ++def sorted_emit_skip_live() -> str: ++ """``qsa_indexer._kernel_sorted_emit_state``: 1 = the Python sort is being ++ skipped because the kernel reported sorted emission, 0 = not skipped, ++ ``n-a`` = not latched yet (or the module was never loaded).""" ++ module = sys.modules.get(_QSA_MODULE) ++ state = getattr(module, "_kernel_sorted_emit_state", None) if module else None ++ if state is None: ++ return "n-a" ++ return "1" if state else "0" ++ ++ + def device_switches() -> dict[str, str]: + """Counter slots [8]/[9] (switch / switch_large). One synchronizing read + (topk_counters.read refuses inside a capture), so it is only called from +@@ -350,6 +379,7 @@ def toggles_line( + cpp: dict[str, str] | str, + qsa_live: str, + device: dict[str, str], ++ sorted_skip: str = "n-a", + ) -> str: + parts = [f"e2-toggles[{tag}]: pid={pid}"] + parts += [f"{name}={raw[name]}" for name in RAW_ORDER] +@@ -362,7 +392,12 @@ def toggles_line( + else: + parts += [f"{key}={cpp[key]}" for key in CPP_ORDER] + parts.append("source=cpp-cache") +- parts += ["python:", f"qsa_sort={qsa_live}", "source=qsa_indexer._QSA_SORT"] ++ parts += [ ++ "python:", ++ f"qsa_sort={qsa_live}", ++ f"sorted_emit_skip={sorted_skip}", ++ "source=qsa_indexer._QSA_SORT", ++ ] + parts.append("device:") + parts += [f"{key}={device[key]}" for key in ("switch", "switch_large")] + parts.append("source=counter-slots-8-9") +@@ -402,6 +437,7 @@ def maybe_emit( + cpp_fn: Callable[[], dict[str, str] | str] = cpp_cache, + device_fn: Callable[[], dict[str, str]] = device_switches, + qsa_fn: Callable[[], str] = qsa_sort_live, ++ sorted_skip_fn: Callable[[], str] = sorted_emit_skip_live, + ) -> str: + """Emit the line at most once per process. One bool check when already + emitted or not yet armed, so it can sit at the head of every execute step. +@@ -431,7 +467,15 @@ def maybe_emit( + cxx_getenv, env_source = getenv_fn() + raw, effective = environ_view(cxx_getenv, os.environ.get) + line = toggles_line( +- tag, os.getpid(), raw, effective, env_source, cpp, qsa_fn(), device_fn() ++ tag, ++ os.getpid(), ++ raw, ++ effective, ++ env_source, ++ cpp, ++ qsa_fn(), ++ device_fn(), ++ sorted_skip_fn(), + ) + _emitted = True + logger.info("%s", line) +-- +2.43.0 + diff --git a/patches/0057-engine-executor-drain-order-shutdown.patch b/patches/0057-engine-executor-drain-order-shutdown.patch new file mode 100644 index 0000000..ca426bf --- /dev/null +++ b/patches/0057-engine-executor-drain-order-shutdown.patch @@ -0,0 +1,1965 @@ +From fc7b101d3fa07d5bd097b1a880de15c208664d48 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 24 Sep 2026 01:36:08 -0400 +Subject: [PATCH 057/103] engine, executor: drain-order shutdown + +- EngineCore survives a late SIGTERM during teardown (bounded deferral) so + every worker drains its guard telemetry. +- A worker exits right after its orderly shutdown. +- A worker failure racing shutdown stays a failure (exit 1). +- Rank 0 waits for the PLE child's whole thread group, bounded, instead of its + reap; drain-aware failure ladder with a bounded grace for survivors whose + shutdown already started. +- tests: the sorted-emit graph arm takes max_len on the host before capture. +--- + tests/kernels/merlin_sorted_emit_arm.py | 16 +- + .../test_engine_core_teardown_signal.py | 194 ++++++++ + .../test_multiproc_executor_parent_death.py | 2 +- + .../test_multiproc_failed_worker_ladder.py | 212 ++++++++ + .../test_multiproc_ple_exit_without_reap.py | 344 +++++++++++++ + .../test_multiproc_worker_exit_after_drain.py | 456 ++++++++++++++++++ + vllm/v1/engine/core.py | 98 +++- + vllm/v1/executor/multiproc_executor.py | 226 ++++++++- + vllm/v1/ple_offload/worker.py | 118 +++++ + 9 files changed, 1640 insertions(+), 26 deletions(-) + create mode 100644 tests/v1/engine/test_engine_core_teardown_signal.py + create mode 100644 tests/v1/executor/test_multiproc_failed_worker_ladder.py + create mode 100644 tests/v1/executor/test_multiproc_ple_exit_without_reap.py + create mode 100644 tests/v1/executor/test_multiproc_worker_exit_after_drain.py + +diff --git a/tests/kernels/merlin_sorted_emit_arm.py b/tests/kernels/merlin_sorted_emit_arm.py +index 700dab1..5141f28 100644 +--- a/tests/kernels/merlin_sorted_emit_arm.py ++++ b/tests/kernels/merlin_sorted_emit_arm.py +@@ -97,8 +97,11 @@ def _inputs(name, k, lengths, kind): + return logits, torch.tensor(lengths, dtype=torch.int32), prefill + + +-def _launch(logits, lengths, out, ws, k): +- torch.ops._C.persistent_topk(logits, lengths, out, ws, k, int(lengths.max())) ++def _launch(logits, lengths, out, ws, k, max_len): ++ # max_len is a host int computed BEFORE any graph capture: int(t.max()) on ++ # a CUDA tensor is a D2H sync, which is illegal inside torch.cuda.graph ++ # (cudaErrorStreamCaptureUnsupported). ++ torch.ops._C.persistent_topk(logits, lengths, out, ws, k, max_len) + + + def main() -> int: +@@ -109,9 +112,10 @@ def main() -> int: + cells, replays_ok = {}, True + for name, k, lengths, kind in CELLS: + logits, lens, prefill = _inputs(name, k, lengths, kind) ++ max_len = int(lens.max()) # CPU tensor: no device sync + logits, lens = logits.to(dev), lens.to(dev) + out = prefill.to(dev).clone() +- _launch(logits, lens, out, ws, k) ++ _launch(logits, lens, out, ws, k, max_len) + torch.cuda.synchronize() + cells[name] = out.cpu().clone() + if graph: +@@ -119,11 +123,11 @@ def main() -> int: + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) +- with torch.cuda.stream(s): +- _launch(logits, lens, static_out, ws, k) # warm on a side stream ++ with torch.cuda.stream(s): # warm on a side stream ++ _launch(logits, lens, static_out, ws, k, max_len) + torch.cuda.current_stream().wait_stream(s) + with torch.cuda.graph(g): +- _launch(logits, lens, static_out, ws, k) ++ _launch(logits, lens, static_out, ws, k, max_len) + static_out.copy_(prefill.to(dev)) + g.replay() + torch.cuda.synchronize() +diff --git a/tests/v1/engine/test_engine_core_teardown_signal.py b/tests/v1/engine/test_engine_core_teardown_signal.py +new file mode 100644 +index 0000000..94a5255 +--- /dev/null ++++ b/tests/v1/engine/test_engine_core_teardown_signal.py +@@ -0,0 +1,194 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""EngineCore must survive a late SIGTERM while it tears down its executor. ++ ++Workers carry PR_SET_PDEATHSIG=SIGKILL, so an EngineCore that dies before ++model_executor.shutdown() has reaped them SIGKILLs every worker before its ++bounds-guard FINAL DRAIN runs. Observed on an earlier-build shutdown: the ++process-group SIGTERM started the shutdown, the API server's process manager ++sent a second SIGTERM in the same second, the finally block had already reset ++SIGTERM to SIG_DFL, EngineCore died before "[shutdown] Executor: waiting for ++worker exit", and FINAL DRAIN printed on 2 of 4 ranks. ++ ++The test drives the real EngineCoreProc.run_engine_core with a fake engine whose ++shutdown() receives that second SIGTERM mid-teardown. Unfixed code: the child ++dies with -SIGTERM before teardown completes (the failing control). ++""" ++ ++import contextlib ++import multiprocessing ++import os ++import signal ++import sys ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++pytestmark = pytest.mark.skipif( ++ not sys.platform.startswith("linux"), reason="fork + POSIX signals" ++) ++ ++ ++def _child_run(conn, mode: str = "once") -> None: ++ import logging ++ ++ from vllm.v1.engine import core ++ ++ if mode == "once-hang": ++ core._TEARDOWN_SIGNAL_BOUND_S = 1.0 ++ ++ run_engine_core = core.EngineCoreProc.run_engine_core ++ ++ records: list[str] = [] ++ ++ class _Capture(logging.Handler): ++ def emit(self, record): ++ records.append(record.getMessage()) ++ ++ core.logger.addHandler(_Capture(level=logging.WARNING)) ++ ++ class _FakeEngineCore: ++ def __init__(self, *args, **kwargs): ++ import queue ++ ++ self.input_queue = queue.Queue() ++ self.shutdown_state = None ++ self.vllm_config = kwargs["vllm_config"] ++ ++ def run_busy_loop(self): ++ # As after the first (process-group) SIGTERM: the busy loop exits. ++ raise SystemExit ++ ++ def has_work(self): ++ return False ++ ++ def shutdown(self): ++ # The process manager's SIGTERM, delivered mid-teardown. ++ conn.send(("teardown-start", None)) ++ os.kill(os.getpid(), signal.SIGTERM) ++ if mode == "twice-hang": ++ time.sleep(0.2) ++ conn.send(("second-signal", None)) ++ os.kill(os.getpid(), signal.SIGTERM) ++ if mode in ("twice-hang", "once-hang"): ++ time.sleep(3600) # teardown hangs (e.g. distributed cleanup) ++ time.sleep(0.3) # the executor waiting for its workers ++ conn.send(("teardown-done", None)) ++ ++ core.EngineCoreProc = _FakeEngineCore ++ core.maybe_register_config_serialize_by_value = lambda: None ++ core.set_process_title = lambda *a, **k: None ++ core.maybe_init_worker_tracer = lambda *a, **k: None ++ core.decorate_logs = lambda *a, **k: None ++ ++ vllm_config = SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ data_parallel_size=1, ++ numa_bind=False, ++ data_parallel_index=0, ++ reconfigure_for_independent_dp_rank=lambda: None, ++ ), ++ kv_transfer_config=None, ++ model_config=SimpleNamespace(is_moe=False), ++ shutdown_timeout=60, ++ ) ++ with contextlib.suppress(SystemExit): ++ run_engine_core(vllm_config=vllm_config) ++ conn.send( ++ ( ++ "after", ++ { ++ "sigterm_is_default": signal.getsignal(signal.SIGTERM) ++ is signal.SIG_DFL, ++ "sigint_is_default": signal.getsignal(signal.SIGINT) is signal.SIG_DFL, ++ "warnings": records, ++ }, ++ ) ++ ) ++ conn.close() ++ ++ ++def test_late_sigterm_does_not_kill_engine_core_teardown(): ++ ctx = multiprocessing.get_context("fork") ++ parent, child = ctx.Pipe(duplex=False) ++ proc = ctx.Process(target=_child_run, args=(child,)) ++ proc.start() ++ child.close() ++ events = {} ++ deadline = time.monotonic() + 60 ++ while time.monotonic() < deadline: ++ if not parent.poll(0.1): ++ # after exit, drain whatever was sent before it ++ if not proc.is_alive() and not parent.poll(0): ++ break ++ continue ++ try: ++ name, payload = parent.recv() ++ except EOFError: ++ break ++ events[name] = payload ++ if name == "after": ++ break ++ proc.join(30) ++ assert "teardown-start" in events, f"fake teardown never ran: {events}" ++ assert proc.exitcode != -signal.SIGTERM, ( ++ "EngineCore died of the SIGTERM delivered during teardown " ++ "(workers would be SIGKILLed before their FINAL DRAIN)" ++ ) ++ assert "teardown-done" in events, f"teardown did not complete: {events}" ++ after = events.get("after") ++ assert after is not None, f"run_engine_core did not return: {events}" ++ assert proc.exitcode == 0, proc.exitcode ++ # default dispositions are restored once teardown is over ++ assert after["sigterm_is_default"] and after["sigint_is_default"], after ++ # fail loud: the deferred signal is logged ++ assert any( ++ "SIGTERM during resource teardown deferred" in m for m in after["warnings"] ++ ), after["warnings"] ++ ++ ++def _run_hanging_teardown(mode: str, bound_s: float): ++ ctx = multiprocessing.get_context("fork") ++ parent, child = ctx.Pipe(duplex=False) ++ proc = ctx.Process(target=_child_run, args=(child, mode)) ++ start = time.monotonic() ++ proc.start() ++ child.close() ++ proc.join(bound_s) ++ elapsed = time.monotonic() - start ++ alive = proc.is_alive() ++ if alive: ++ proc.kill() ++ proc.join(10) ++ events = [] ++ while True: ++ try: ++ if not parent.poll(0.2): ++ break ++ events.append(parent.recv()[0]) ++ except (EOFError, OSError): ++ break ++ return alive, proc.exitcode, elapsed, events ++ ++ ++def test_second_late_signal_ends_a_hung_teardown(): ++ # The deferral must be bounded. The first late SIGTERM is ++ # deferred (the drain-race fix); a second one while teardown hangs ends ++ # the process by that signal. Unfixed code (every late signal deferred) ++ # stays alive until killed (the failing control). ++ alive, exitcode, elapsed, events = _run_hanging_teardown("twice-hang", 15) ++ assert "teardown-start" in events and "second-signal" in events, events ++ assert not alive, f"hung teardown survived a second SIGTERM ({elapsed:.1f}s)" ++ assert exitcode == -signal.SIGTERM, exitcode ++ assert "teardown-done" not in events, events ++ ++ ++def test_teardown_watchdog_bounds_a_single_deferred_signal(): ++ # One deferred SIGTERM, then teardown hangs: the watchdog ends the process ++ # (exit code 1) once _TEARDOWN_SIGNAL_BOUND_S (1 s here) has passed. ++ alive, exitcode, elapsed, events = _run_hanging_teardown("once-hang", 15) ++ assert "teardown-start" in events, events ++ assert not alive, f"hung teardown outlived the watchdog ({elapsed:.1f}s)" ++ assert exitcode == 1, exitcode ++ assert "teardown-done" not in events, events +diff --git a/tests/v1/executor/test_multiproc_executor_parent_death.py b/tests/v1/executor/test_multiproc_executor_parent_death.py +index 111fca6..2ea2755 100644 +--- a/tests/v1/executor/test_multiproc_executor_parent_death.py ++++ b/tests/v1/executor/test_multiproc_executor_parent_death.py +@@ -87,7 +87,7 @@ def _run_fake_worker_main( + time.sleep(3600) + raise AssertionError(f"unknown block mode: {block_mode}") + +- def shutdown(self): ++ def shutdown(self, *, destroy_distributed: bool = True): + pass + + executor_module.WorkerProc = FakeWorkerProc +diff --git a/tests/v1/executor/test_multiproc_failed_worker_ladder.py b/tests/v1/executor/test_multiproc_failed_worker_ladder.py +new file mode 100644 +index 0000000..9c59113 +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_failed_worker_ladder.py +@@ -0,0 +1,212 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""After a worker died, parked survivors are killed fast; draining ones are not. ++ ++A drain fault-injection run (earlier build): the PLE fault killed rank 0. One ++survivor drained and exited at once. Two ranks sat in a collective with the ++dead rank, inside C code where neither the grace nor the SIGTERM reaches them, ++and the executor ran its full 5 s grace + 4 s SIGTERM wait before SIGKILL. ++Workers were gone 17.6 s after the fault, vs 8.7 s on an earlier build, where the ++engine died on the late SIGTERM and PDEATHSIG killed the survivors at once. ++the drain-order change keeps the engine alive for the drain, so the failure ladder has ++to be short on its own. ++ ++A known hazard: a blanket 1 s cap can cut a survivor that is still ++draining. Each worker now publishes a shared byte: 0 = shutdown not started, ++1 = started (FINAL DRAIN in progress), 2 = done. Only workers with no drain ++signal get the 1 s ladder. ++ ++A later drain fault-injection run: a survivor whose shutdown had STARTED blocked ++inside it on the dead rank (no FINAL DRAIN line) and got the full orderly ++ladder; workers gone 17.6 s after the fault. Started shutdowns now get a ++bounded 2 s grace on the failure path (every measured FINAL DRAIN < 0.1 s), ++then SIGTERM, then SIGKILL after 1 s. ++""" ++ ++import contextlib ++import signal ++import subprocess ++import sys ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++pytestmark = [ ++ pytest.mark.skip_global_cleanup, ++ pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals"), ++] ++ ++# Parked in a collective: ignores SIGTERM, never exits on its own. ++_PARKED = ( ++ "import signal, time\n" ++ "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" ++ "print('ready', flush=True)\n" ++ "time.sleep(600)\n" ++) ++# Draining: default SIGTERM (would die -15), exits 0 by itself after 1.2 s -- ++# 12x the longest FINAL DRAIN measured on the reference host, inside the 2 s failure grace. ++_DRAINING_S = 1.2 ++_DRAINING = ( ++ "import os, time\n" ++ "print('ready', flush=True)\n" ++ f"time.sleep({_DRAINING_S})\n" ++ "os._exit(0)\n" ++) ++ ++ ++class _Proc: ++ """The BaseProcess surface _ensure_worker_termination uses.""" ++ ++ def __init__(self, script): ++ self.popen = subprocess.Popen( ++ [sys.executable, "-c", script], stdout=subprocess.PIPE, text=True ++ ) ++ assert self.popen.stdout.readline().strip() == "ready" ++ ++ def is_alive(self): ++ return self.popen.poll() is None ++ ++ def terminate(self): ++ self.popen.send_signal(signal.SIGTERM) ++ ++ def kill(self): ++ self.popen.kill() ++ ++ ++class _Handle: ++ def __init__(self, script, state): ++ self.proc = _Proc(script) ++ self.death_writer = None ++ self.worker_response_mq = None ++ # None = no state published (older handle): treated as parked. ++ self.shutdown_state = None if state is None else SimpleNamespace(value=state) ++ ++ ++def _executor(workers, *, failed): ++ from vllm.v1.executor.multiproc_executor import MultiprocExecutor ++ ++ executor = object.__new__(MultiprocExecutor) ++ executor.workers = workers ++ executor.is_failed = failed ++ executor.rpc_broadcast_mq = None ++ executor.response_mqs = [] ++ return executor ++ ++ ++def _shutdown(specs, *, failed): ++ workers = [_Handle(script, state) for script, state in specs] ++ try: ++ start = time.monotonic() ++ _executor(workers, failed=failed).shutdown() ++ elapsed = time.monotonic() - start ++ # The ladder's last step sends SIGKILL without waiting for the reap. ++ for w in workers: ++ with contextlib.suppress(subprocess.TimeoutExpired): ++ w.proc.popen.wait(2) ++ codes = [w.proc.popen.returncode for w in workers] ++ finally: ++ for w in workers: ++ if w.proc.is_alive(): ++ w.proc.kill() ++ w.proc.popen.wait(10) ++ return elapsed, codes ++ ++ ++def test_failure_shutdown_kills_parked_survivors_fast(): ++ elapsed, codes = _shutdown([(_PARKED, 0), (_PARKED, None)], failed=True) ++ assert all(c is not None for c in codes), codes ++ assert elapsed < 3.0, ( ++ f"a failure-triggered shutdown took {elapsed:.1f} s to kill ranks parked " ++ "with the dead rank (the orderly 5 s grace + 4 s SIGTERM ladder)" ++ ) ++ ++ ++def test_failure_shutdown_lets_a_draining_survivor_finish(): ++ elapsed, codes = _shutdown([(_DRAINING, 1), (_PARKED, 0)], failed=True) ++ assert codes[0] == 0, ( ++ f"the draining survivor was signalled (exit {codes[0]}) instead of " ++ "finishing its FINAL DRAIN" ++ ) ++ assert codes[1] is not None ++ # The parked survivor is SIGKILLed at ~2 s; the drainer was never signalled. ++ assert _DRAINING_S - 0.3 <= elapsed < 3.0, elapsed ++ ++ ++def test_failure_shutdown_ends_a_survivor_stuck_in_its_shutdown(): ++ # Started its shutdown (state 1), then blocked on the dead rank: ignores ++ # SIGTERM like a rank inside a device synchronize, never exits by itself. ++ elapsed, codes = _shutdown([(_PARKED, 1), (_PARKED, 0)], failed=True) ++ assert all(c is not None for c in codes), codes ++ assert codes[0] == -signal.SIGKILL, codes ++ assert elapsed < 2.0 + 1.0 + 0.6, ( ++ f"a failure-triggered shutdown took {elapsed:.1f} s to end a survivor " ++ "stuck inside its started shutdown (bound: 2 s grace + 1 s SIGTERM wait)" ++ ) ++ ++ ++def test_orderly_shutdown_keeps_its_grace(monkeypatch): ++ monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", "2") ++ elapsed, codes = _shutdown([(_PARKED, 0), (_PARKED, 0)], failed=False) ++ assert all(c is not None for c in codes), codes ++ # Unchanged: the configured grace, then the 4 s SIGTERM wait, then SIGKILL. ++ assert elapsed >= 2.0 + 4.0 - 0.5, elapsed ++ ++ ++def test_worker_main_publishes_shutdown_started_then_done(monkeypatch): ++ """worker_main marks 1 before WorkerProc.shutdown() (drain) and 2 after.""" ++ import multiprocessing ++ import os ++ ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ seen = [] ++ state = multiprocessing.RawValue("b", 0) ++ ++ class _Inner: ++ def shutdown(self): ++ seen.append(("drain", state.value)) ++ ++ def fake_init(self, *args, **kwargs): ++ self.rank = 0 ++ self.worker = _Inner() ++ self.rpc_broadcast_mq = None ++ self.worker_response_mq = SimpleNamespace( ++ export_handle=lambda: None, ++ wait_until_ready=lambda: None, ++ shutdown=lambda: None, ++ ) ++ self.peer_response_handles = [] ++ ++ def busy_loop(self): ++ seen.append(("busy", state.value)) ++ raise ValueError("boom") # an unexpected failure: full teardown path ++ ++ monkeypatch.setattr(executor_module.WorkerProc, "__init__", fake_init) ++ monkeypatch.setattr(executor_module.WorkerProc, "worker_busy_loop", busy_loop) ++ monkeypatch.setattr(executor_module, "destroy_model_parallel", lambda: None) ++ monkeypatch.setattr( ++ executor_module, "destroy_distributed_environment", lambda: None ++ ) ++ monkeypatch.setattr(executor_module, "maybe_init_worker_tracer", lambda **k: None) ++ monkeypatch.setattr(executor_module, "set_worker_net_device", lambda *a, **k: None) ++ monkeypatch.setattr( ++ executor_module, "_install_worker_parent_death_signal", lambda pid: None ++ ) ++ monkeypatch.setattr(executor_module.signal, "signal", lambda *a, **k: None) ++ ready_reader, ready_writer = multiprocessing.Pipe(duplex=False) ++ executor_module.WorkerProc.worker_main( ++ vllm_config=SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, numa_bind=False ++ ) ++ ), ++ ready_pipe=ready_writer, ++ death_pipe=None, ++ shutdown_state=state, ++ inherited_fds=[], ++ expected_parent_pid=os.getpid(), ++ ) ++ ready_reader.close() ++ assert seen == [("busy", 0), ("drain", 1)], seen ++ assert state.value == 2 +diff --git a/tests/v1/executor/test_multiproc_ple_exit_without_reap.py b/tests/v1/executor/test_multiproc_ple_exit_without_reap.py +new file mode 100644 +index 0000000..4bab227 +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_ple_exit_without_reap.py +@@ -0,0 +1,344 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Rank 0 waits for its PLE offload child to leave userspace, not for its reap. ++ ++A drain run on an earlier build: on every F cycle three ranks exited ++right after their FINAL DRAIN, but Worker_PP0_TP0, the rank that owns the PLE ++offload process, sat in Worker.shutdown() -> PleOffloadWorkerHandle.close() ++until the executor's grace ran out (SIGTERM at +5 s, SIGKILL at +9 s). A ++read-only /proc sample showed the cause. The offload child (RSS ~51 GB) ++finished its userspace, and its leader went Z at ~+2.7 s. One thread stayed R ++with PF_EXITING set, in-kernel, freeing the address space, so the child was ++not reapable until ~+8.6 s. ++ ++A known hazard: a Z leader alone does not prove the child left ++userspace, because another thread can still run. close() therefore waits for ++every task to be Z/X or PF_EXITING. ++ ++Covered here: ++- the predicate itself, on a fake procfs; ++- close() does not wait for the reap once the thread group is in exit (fake ++ procfs over a live, unreaped child); the earlier build waits for the reap and fails this; ++- close() does wait while a non-leader thread still runs userspace (real ++ child whose leader exits alone); the earlier commit fails this; ++- the bound: a child that ignores SIGTERM is SIGKILLed within the bound, and ++ the bound fits inside the executor's grace; ++- the default close() still reaps. ++""" ++ ++import contextlib ++import platform ++import signal ++import socket ++import subprocess ++import sys ++import tempfile ++import threading ++import time ++from multiprocessing.connection import Connection ++from pathlib import Path ++ ++import pytest ++ ++pytestmark = pytest.mark.skip_global_cleanup ++ ++linux_only = pytest.mark.skipif( ++ sys.platform != "linux" or platform.machine() not in ("x86_64", "aarch64"), ++ reason="needs Linux /proc and a raw SYS_exit number", ++) ++ ++_SYS_EXIT = {"x86_64": 60, "aarch64": 93}.get(platform.machine(), -1) ++ ++# On SIGTERM: run the cleanup (unlink + marker, like proc_main's finally), then ++# the leader thread exits alone (raw SYS_exit) while a second thread keeps ++# running userspace for LINGER seconds. mode=ignore: SIGTERM is ignored. ++_STUB_CHILD = r""" ++import ctypes ++import os ++import pathlib ++import signal ++import sys ++import threading ++import time ++ ++endpoint, marker, linger_s, sys_exit, mode = sys.argv[1:] ++stop = threading.Event() ++if mode == "ignore": ++ signal.signal(signal.SIGTERM, signal.SIG_IGN) ++else: ++ signal.signal(signal.SIGTERM, lambda _s, _f: stop.set()) ++pathlib.Path(endpoint).touch() ++print("ready", flush=True) ++if mode == "ignore": ++ time.sleep(600) ++stop.wait() ++os.unlink(endpoint) ++pathlib.Path(marker).touch() ++ ++ ++def linger(): ++ time.sleep(float(linger_s)) # userspace: NOT in the exit path yet ++ os._exit(0) ++ ++ ++threading.Thread(target=linger).start() ++ctypes.CDLL(None).syscall(int(sys_exit), 0) ++""" ++ ++ ++def _stat_line(pid, comm, state, flags): ++ # Field 9 (flags) is the 7th field after the parenthesised comm. ++ return f"{pid} ({comm}) {state} 1 1 1 0 -1 {flags} 0 0 0 0\n" ++ ++ ++def _fake_task_tree(root: Path, pid: int, tasks): ++ for tid, comm, state, flags in tasks: ++ d = root / str(pid) / "task" / str(tid) ++ d.mkdir(parents=True) ++ (d / "stat").write_text(_stat_line(tid, comm, state, flags)) ++ ++ ++def _start_handle(tmp: Path, linger_s: float, mode: str = "graceful"): ++ from vllm.v1.ple_offload.worker import PleOffloadWorkerHandle ++ ++ endpoint = tmp / "ple.sock" ++ marker = tmp / "cleanup.done" ++ proc = subprocess.Popen( ++ [ ++ sys.executable, ++ "-c", ++ _STUB_CHILD, ++ str(endpoint), ++ str(marker), ++ str(linger_s), ++ str(_SYS_EXIT), ++ mode, ++ ], ++ stdout=subprocess.PIPE, ++ text=True, ++ ) ++ assert proc.stdout.readline().strip() == "ready" ++ # gpu_worker._start_ple_offload_watchdog: a thread blocked in proc.wait() ++ # holds Popen._waitpid_lock for as long as the child is unreapable. ++ watchdog = threading.Thread(target=proc.wait, daemon=True) ++ watchdog.start() ++ parent_death, child_death = socket.socketpair() ++ handle = PleOffloadWorkerHandle( ++ proc=proc, ++ death_writer=Connection(parent_death.detach()), ++ ready_pipe_reader=None, ++ launch_deadline=time.monotonic() + 60, ++ ipc_addr=f"ipc://{endpoint}", ++ ) ++ return handle, proc, watchdog, endpoint, marker, child_death ++ ++ ++def _finish(proc, watchdog, child_end, join_s=15): ++ with contextlib.suppress(ProcessLookupError): ++ if proc.returncode is None: ++ proc.send_signal(signal.SIGKILL) ++ watchdog.join(join_s) ++ child_end.close() ++ ++ ++@pytest.mark.parametrize( ++ "tasks, expected", ++ [ ++ # Observed shape (a shutdown sample): leader Z, one task R with PF_EXITING. ++ ([(1, "python3.13", "Z", 0x40810C), (2, "python3.13", "R", 0x40044C)], True), ++ # Leader Z, a thread still in userspace (no PF_EXITING). ++ ([(1, "python3.13", "Z", 0x40810C), (2, "python3.13", "S", 0x400040)], False), ++ ([(1, "python3.13", "S", 0x400100)], False), ++ ([(1, "odd) comm (x", "X", 0)], True), ++ ([], True), ++ ], ++) ++def test_thread_group_predicate(tmp_path, tasks, expected): ++ from vllm.v1.ple_offload.worker import _child_thread_group_exiting ++ ++ pid = 4242 ++ if tasks: ++ _fake_task_tree(tmp_path, pid, tasks) ++ else: ++ (tmp_path / str(pid) / "task").mkdir(parents=True) ++ proc = type("P", (), {"pid": pid, "returncode": None})() ++ assert _child_thread_group_exiting(proc, str(tmp_path)) is expected ++ ++ ++@linux_only ++def test_exit_imminent_close_does_not_wait_for_the_reap(monkeypatch, tmp_path): ++ import vllm.v1.ple_offload.worker as ple_worker ++ ++ monkeypatch.setattr(ple_worker, "_PROCESS_EXIT_IMMINENT", True, raising=False) ++ with tempfile.TemporaryDirectory() as d: ++ # A live child that will not exit on SIGTERM stands in for the 51 GB ++ # child: unreapable for as long as the test wants. ++ handle, proc, watchdog, endpoint, marker, child_end = _start_handle( ++ Path(d), 0, mode="ignore" ++ ) ++ try: ++ # procfs says: leader Z, the other task R with PF_EXITING. ++ fake = tmp_path / "proc" ++ _fake_task_tree( ++ fake, ++ proc.pid, ++ [(proc.pid, "python3.13", "Z", 0x40810C), (1, "x", "R", 0x40044C)], ++ ) ++ monkeypatch.setattr(ple_worker, "_PROC_ROOT", str(fake), raising=False) ++ start = time.monotonic() ++ handle.close() ++ elapsed = time.monotonic() - start ++ assert elapsed < 1.0, ( ++ f"close() waited {elapsed:.1f} s for the offload child's reap " ++ "although its thread group was in the kernel exit path; the " ++ "executor's 5 s grace expires first on the reference host" ++ ) ++ assert proc.poll() is None, "close() reaped or killed the child" ++ assert handle.death_writer is None ++ finally: ++ _finish(proc, watchdog, child_end) ++ ++ ++@linux_only ++def test_exit_imminent_close_waits_for_a_userspace_thread(monkeypatch): ++ import vllm.v1.ple_offload.worker as ple_worker ++ ++ monkeypatch.setattr(ple_worker, "_PROCESS_EXIT_IMMINENT", True, raising=False) ++ linger_s = 1.5 ++ with tempfile.TemporaryDirectory() as d: ++ handle, proc, watchdog, endpoint, marker, child_end = _start_handle( ++ Path(d), linger_s ++ ) ++ try: ++ start = time.monotonic() ++ handle.close() ++ elapsed = time.monotonic() - start ++ assert elapsed >= linger_s - 0.2, ( ++ f"close() returned after {elapsed:.2f} s, while a child thread " ++ f"was still running userspace ({linger_s} s): a Z leader is " ++ "not the whole thread group" ++ ) ++ assert elapsed < linger_s + 1.5, elapsed ++ assert marker.exists(), "child cleanup (proc_main finally) never ran" ++ assert not endpoint.exists() ++ finally: ++ _finish(proc, watchdog, child_end) ++ assert proc.returncode == 0 ++ ++ ++@linux_only ++def test_exit_imminent_wait_is_bounded_inside_the_grace(monkeypatch): ++ import vllm.envs as envs ++ import vllm.v1.ple_offload.worker as ple_worker ++ ++ monkeypatch.setattr(ple_worker, "_PROCESS_EXIT_IMMINENT", True, raising=False) ++ bound = ple_worker._EXIT_IMMINENT_CHILD_WAIT_S ++ # Rank 0 must be gone before the executor's grace; leave >= 1 s for the ++ # rest of its shutdown and its own exit. ++ assert bound <= envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS - 1.0 ++ with tempfile.TemporaryDirectory() as d: ++ handle, proc, watchdog, endpoint, marker, child_end = _start_handle( ++ Path(d), 0, mode="ignore" ++ ) ++ try: ++ start = time.monotonic() ++ handle.close() ++ elapsed = time.monotonic() - start ++ assert bound - 0.1 <= elapsed < bound + 1.0, elapsed ++ watchdog.join(5) ++ assert proc.returncode == -signal.SIGKILL ++ finally: ++ _finish(proc, watchdog, child_end) ++ ++ ++@linux_only ++def test_default_close_still_reaps(monkeypatch): ++ import vllm.v1.ple_offload.worker as ple_worker ++ ++ monkeypatch.setattr(ple_worker, "_PROCESS_EXIT_IMMINENT", False, raising=False) ++ with tempfile.TemporaryDirectory() as d: ++ handle, proc, watchdog, endpoint, marker, child_end = _start_handle( ++ Path(d), 1.0 ++ ) ++ try: ++ handle.close() ++ # Unchanged full-teardown path: the child is reaped before return. ++ assert proc.returncode is not None ++ assert marker.exists() and not endpoint.exists() ++ finally: ++ _finish(proc, watchdog, child_end) ++ ++ ++class _FakeInner: ++ def __init__(self): ++ self.calls = 0 ++ ++ def shutdown(self): ++ self.calls += 1 ++ ++ ++@pytest.mark.parametrize("destroy_distributed", [False, True]) ++def test_worker_shutdown_marks_exit_imminent_only_on_the_exit_path( ++ monkeypatch, destroy_distributed ++): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ import vllm.v1.ple_offload.worker as ple_worker ++ ++ monkeypatch.setattr(ple_worker, "_PROCESS_EXIT_IMMINENT", False, raising=False) ++ monkeypatch.setattr(executor_module, "destroy_model_parallel", lambda: None) ++ monkeypatch.setattr( ++ executor_module, "destroy_distributed_environment", lambda: None ++ ) ++ proc = object.__new__(executor_module.WorkerProc) ++ proc.rpc_broadcast_mq = None ++ proc.worker_response_mq = None ++ proc.worker = _FakeInner() ++ proc.shutdown(destroy_distributed=destroy_distributed) ++ assert proc.worker.calls == 1 ++ assert ple_worker._PROCESS_EXIT_IMMINENT is (not destroy_distributed) ++ ++ ++_FATAL_CHILD = r""" ++import os ++import sys ++import time ++ ++import vllm.v1.ple_offload.worker as ple_worker ++from vllm.v1.ple_offload.fatal import ple_fail_closed ++ ++ ++def forbidden(*args, **kwargs): ++ os.write(1, b"CLOSE-CALLED\n") ++ time.sleep(10) # what a reap wait would cost ++ ++ ++ple_worker.PleOffloadWorkerHandle.close = forbidden ++ple_worker.PleOffloadWorkerHandle._close_without_reap = forbidden ++t0 = time.monotonic() ++os.write(1, b"FIRE %.6f\n" % t0) ++try: ++ raise RuntimeError("PLE test fault: exc_after_prepare") ++except RuntimeError as error: ++ ple_fail_closed(where="execute_model", rank=0, layer="l", sequence=9, error=error) ++os.write(1, b"RETURNED\n") ++""" ++ ++ ++def test_ple_fatal_path_exits_without_closing_or_reaping_the_offload_child(): ++ """Drain fault-injection test: rank 0 died ~8.3 s after its PLE fatal line. The fatal path ++ must be a straight os._exit(1): no Worker.shutdown(), no ++ PleOffloadWorkerHandle.close() (terminate + reap wait). The offload child ++ is ended by its own death-pipe monitor when this process exits.""" ++ start = time.monotonic() ++ proc = subprocess.run( ++ [sys.executable, "-c", _FATAL_CHILD], ++ capture_output=True, ++ text=True, ++ timeout=120, ++ ) ++ assert proc.returncode == 1, (proc.returncode, proc.stdout, proc.stderr[-500:]) ++ assert "CLOSE-CALLED" not in proc.stdout ++ assert "RETURNED" not in proc.stdout ++ assert "FIRE " in proc.stdout ++ assert "PLE fatal: where=execute_model" in proc.stderr ++ assert time.monotonic() - start < 60 +diff --git a/tests/v1/executor/test_multiproc_worker_exit_after_drain.py b/tests/v1/executor/test_multiproc_worker_exit_after_drain.py +new file mode 100644 +index 0000000..7694a62 +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_worker_exit_after_drain.py +@@ -0,0 +1,456 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""A worker whose orderly shutdown has finished must exit, not linger. ++ ++A fault-injection run with the drain-order deferred-signal fix: every rank printed ++its bounds-guard FINAL DRAIN within ~1 s of the executor closing the death pipes, ++but then stayed alive past VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS (5 s); the ++executor logged two WARNINGs and SIGKILLed the workers at ~9 s. After ++Worker.shutdown() (drain, instrumentation flushes, PLE cleanup) WorkerProc ++destroys the model-parallel groups, and PyNcclCommunicator.destroy() joins its ++ncclCommAbort thread for up to 5 s per communicator while the CUDA graphs that ++captured NCCL ops are still alive. The executor's SIGTERM is then a no-op (the ++worker's handler raises SystemExit only once and the death-pipe monitor already ++set shutdown_requested), so the ladder runs to SIGKILL. ++ ++The fake below reproduces that: destroy_model_parallel blocks _SLOW_DESTROY_S. ++Unfixed code pays it before the process can exit (the failing controls); fixed ++code runs WorkerProc.shutdown(destroy_distributed=False) and exits right after. ++""" ++ ++import logging ++import multiprocessing ++import os ++import sys ++import threading ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++# CPU-only fakes: no process groups or device memory to clean up afterwards. ++pytestmark = pytest.mark.skip_global_cleanup ++ ++_SLOW_DESTROY_S = 6.0 # > the executor's 5 s grace, like the NCCL destroy join ++_EXIT_BOUND_S = 3.0 # reference-host bound: all workers gone <= 3 s after the drain ++ ++ ++class _FakeMQ: ++ def __init__(self, entered: threading.Event | None = None): ++ self.entered = entered ++ self.stop = threading.Event() ++ self.shutting_down = False ++ ++ def export_handle(self): ++ return None ++ ++ def wait_until_ready(self): ++ pass ++ ++ def shutdown(self): ++ self.shutting_down = True ++ self.stop.set() ++ ++ def dequeue(self, indefinite=True): ++ if self.entered is not None: ++ self.entered.set() ++ # shm_broadcast.acquire_read raises RuntimeError("cancelled") once the ++ # death-pipe monitor has called shutdown(). ++ self.stop.wait() ++ raise RuntimeError("cancelled") ++ ++ ++class _FakeInnerWorker: ++ """Stands in for gpu_worker.Worker; shutdown() is where FINAL DRAIN runs.""" ++ ++ def __init__(self, record): ++ self.record = record ++ ++ def shutdown(self): ++ self.record("drain") ++ ++ ++def _patch_executor(executor_module, record, entered, slow_destroy_s, set_=setattr): ++ real_worker_proc = executor_module.WorkerProc ++ ++ def fake_init(self, *args, **kwargs): ++ self.rank = 0 ++ self.worker = _FakeInnerWorker(record) ++ self.rpc_broadcast_mq = _FakeMQ(entered) ++ self.worker_response_mq = _FakeMQ() ++ self.peer_response_handles = [] ++ self.use_async_scheduling = False ++ ++ def slow_destroy_model_parallel(): ++ record("destroy_model_parallel") ++ time.sleep(slow_destroy_s) ++ ++ set_(real_worker_proc, "__init__", fake_init) ++ set_(executor_module, "destroy_model_parallel", slow_destroy_model_parallel) ++ set_( ++ executor_module, ++ "destroy_distributed_environment", ++ lambda: record("destroy_distributed_environment"), ++ ) ++ set_(executor_module, "maybe_init_worker_tracer", lambda **kwargs: None) ++ set_(executor_module, "set_worker_net_device", lambda *args, **kwargs: None) ++ ++ ++def _fake_config(): ++ return SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, ++ numa_bind=False, ++ ) ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# In-process: the real worker_main / WorkerProc.shutdown, os._exit intercepted. ++# --------------------------------------------------------------------------- ++ ++ ++class _Exited(BaseException): ++ def __init__(self, code): ++ super().__init__(code) ++ self.code = code ++ ++ ++def _run_worker_main_in_process( ++ monkeypatch, ++ *, ++ busy_loop_error=None, ++ fail_after_shutdown_request=False, ++ slow_destroy_s=_SLOW_DESTROY_S, ++): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ events: list[tuple[str, float]] = [] ++ ++ def record(name): ++ events.append((name, time.monotonic())) ++ ++ entered = threading.Event() ++ _patch_executor( ++ executor_module, record, entered, slow_destroy_s, monkeypatch.setattr ++ ) ++ if busy_loop_error is not None: ++ ++ def failing_busy_loop(self): ++ if fail_after_shutdown_request: ++ # Idle until the death-pipe monitor has set shutdown_requested ++ # and shut the queues down, then fail for a reason of its own. ++ entered.set() ++ assert self.rpc_broadcast_mq.stop.wait(30) ++ raise busy_loop_error ++ ++ monkeypatch.setattr( ++ executor_module.WorkerProc, "worker_busy_loop", failing_busy_loop ++ ) ++ # worker_main arms PDEATHSIG and installs SIGTERM/SIGINT handlers; keep ++ # both off the pytest process. ++ monkeypatch.setattr( ++ executor_module, "_install_worker_parent_death_signal", lambda pid: None ++ ) ++ monkeypatch.setattr(executor_module.signal, "signal", lambda *a, **k: None) ++ monkeypatch.setattr(logging, "shutdown", lambda *a, **k: None) ++ # The orderly path marks this (pytest) process exit-imminent for the PLE ++ # handle; restore the flag so later PLE close() tests see a live process. ++ ple_worker = sys.modules.get("vllm.v1.ple_offload.worker") ++ if ple_worker is not None: ++ monkeypatch.setattr( ++ ple_worker, ++ "_PROCESS_EXIT_IMMINENT", ++ getattr(ple_worker, "_PROCESS_EXIT_IMMINENT", False), ++ raising=False, ++ ) ++ ++ def fake_exit(code): ++ record("os._exit") ++ raise _Exited(code) ++ ++ monkeypatch.setattr(os, "_exit", fake_exit) ++ ++ ctx = multiprocessing.get_context() ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ death_reader, death_writer = ctx.Pipe(duplex=False) ++ ++ def close_death_pipe_when_idle(): ++ # The executor's shutdown(): close the death pipe once the worker idles. ++ if entered.wait(30): ++ record("death_pipe_closed") ++ death_writer.close() ++ ++ closer = threading.Thread(target=close_death_pipe_when_idle, daemon=True) ++ closer.start() ++ ++ exit_code = None ++ start = time.monotonic() ++ try: ++ executor_module.WorkerProc.worker_main( ++ vllm_config=_fake_config(), ++ ready_pipe=ready_writer, ++ death_pipe=death_reader, ++ inherited_fds=[], ++ expected_parent_pid=os.getpid(), ++ ) ++ except _Exited as exited: ++ exit_code = exited.code ++ except SystemExit as exited: ++ record("SystemExit") ++ exit_code = exited.code ++ finally: ++ entered.set() # a failing busy loop never dequeues ++ closer.join(5) ++ assert ready_reader.recv()["status"] == "READY" ++ ready_reader.close() ++ return events, exit_code, time.monotonic() - start ++ ++ ++def test_orderly_shutdown_exits_right_after_drain(monkeypatch): ++ events, exit_code, _ = _run_worker_main_in_process(monkeypatch) ++ names = [name for name, _ in events] ++ assert "drain" in names, f"FINAL DRAIN never ran: {names}" ++ assert exit_code is not None, ( ++ "worker_main returned instead of exiting the process after its " ++ f"orderly shutdown (the process would linger in teardown): {names}" ++ ) ++ assert exit_code == 0 ++ assert "destroy_model_parallel" not in names, ( ++ f"the exit-only NCCL destroy ran before the process exit: {names}" ++ ) ++ t = dict(events) ++ assert names.index("drain") < names.index("os._exit"), names ++ assert t["os._exit"] - t["drain"] <= _EXIT_BOUND_S, events ++ ++ ++def test_unexpected_busy_loop_failure_keeps_full_teardown(monkeypatch): ++ # Not an orderly shutdown (no shutdown request): unchanged behaviour -- ++ # full WorkerProc.shutdown incl. distributed destroy, then a normal return. ++ events, exit_code, _ = _run_worker_main_in_process( ++ monkeypatch, busy_loop_error=ValueError("boom"), slow_destroy_s=0.0 ++ ) ++ names = [name for name, _ in events] ++ assert exit_code is None, names ++ assert names.count("drain") == 1, names ++ assert "destroy_model_parallel" in names, names ++ assert "destroy_distributed_environment" in names, names ++ ++ ++def test_failure_racing_shutdown_is_not_an_orderly_exit(monkeypatch): ++ # A real busy-loop failure that happens while a shutdown ++ # request is pending must not take the orderly exit-0 path. Unfixed code ++ # treats ANY exception as orderly once shutdown_requested is set (the ++ # failing control: exit code 0, no distributed teardown). ++ events, exit_code, _ = _run_worker_main_in_process( ++ monkeypatch, ++ busy_loop_error=ValueError("boom"), ++ fail_after_shutdown_request=True, ++ slow_destroy_s=0.0, ++ ) ++ names = [name for name, _ in events] ++ assert "death_pipe_closed" in names, names ++ assert exit_code not in (None, 0), ( ++ f"a worker failure that raced the shutdown request exited {exit_code!r} " ++ f"(hidden as a clean exit): {names}" ++ ) ++ assert "os._exit" not in names, f"took the orderly exit path: {names}" ++ assert names.count("drain") == 1, names ++ assert "destroy_model_parallel" in names, names ++ assert "destroy_distributed_environment" in names, names ++ ++ ++def test_busy_loop_raises_shutdown_type_only_for_its_own_cancellation(): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ proc = executor_module.WorkerProc.__new__(executor_module.WorkerProc) ++ proc.rpc_broadcast_mq = _FakeMQ() ++ proc.rpc_broadcast_mq.shutdown() ++ with pytest.raises(executor_module.WorkerShutdownRequested): ++ proc.worker_busy_loop() ++ ++ class _NotShutDownMQ(_FakeMQ): ++ def dequeue(self, indefinite=True): ++ raise RuntimeError("cancelled") # but nobody shut this queue down ++ ++ proc.rpc_broadcast_mq = _NotShutDownMQ() ++ with pytest.raises(RuntimeError) as info: ++ proc.worker_busy_loop() ++ assert not isinstance(info.value, executor_module.WorkerShutdownRequested) ++ ++ ++# --------------------------------------------------------------------------- ++# Real processes: the executor's own grace/SIGTERM/SIGKILL ladder. ++# --------------------------------------------------------------------------- ++ ++ ++def _worker_child(ready_pipe, death_pipe, events_pipe, parent_pid): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ def record(name): ++ events_pipe.send((name, time.monotonic())) ++ ++ _patch_executor(executor_module, record, None, _SLOW_DESTROY_S) ++ executor_module.WorkerProc.worker_main( ++ vllm_config=_fake_config(), ++ ready_pipe=ready_pipe, ++ death_pipe=death_pipe, ++ inherited_fds=[], ++ expected_parent_pid=parent_pid, ++ ) ++ ++ ++def test_executor_reaps_drained_workers_within_grace(): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ ctx = multiprocessing.get_context("spawn") ++ procs, deaths, readers = [], [], [] ++ for _ in range(2): # two ranks, as the executor sees them ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ death_reader, death_writer = ctx.Pipe(duplex=False) ++ events_reader, events_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_worker_child, ++ args=(ready_writer, death_reader, events_writer, os.getpid()), ++ daemon=True, ++ ) ++ proc.start() ++ ready_writer.close() ++ death_reader.close() ++ events_writer.close() ++ assert ready_reader.poll(180), "worker child never became READY" ++ assert ready_reader.recv()["status"] == "READY" ++ ready_reader.close() ++ procs.append(proc) ++ deaths.append(death_writer) ++ readers.append(events_reader) ++ ++ warnings: list[str] = [] ++ ++ class _Capture(logging.Handler): ++ def emit(self, record): ++ warnings.append(record.getMessage()) ++ ++ handler = _Capture(level=logging.WARNING) ++ executor_module.logger.addHandler(handler) ++ try: ++ start = time.monotonic() ++ for death_writer in deaths: # MultiprocExecutor.shutdown() ++ death_writer.close() ++ executor_module.MultiprocExecutor._ensure_worker_termination(procs) ++ elapsed = time.monotonic() - start ++ finally: ++ executor_module.logger.removeHandler(handler) ++ for proc in procs: ++ if proc.is_alive(): ++ proc.kill() ++ proc.join(10) ++ ++ for rank, reader in enumerate(readers): ++ names = [] ++ # The child has exited: read what it sent until EOF (Windows reports a ++ # closed pipe as BrokenPipeError). ++ while True: ++ try: ++ if not reader.poll(0.5): ++ break ++ names.append(reader.recv()[0]) ++ except (EOFError, OSError): ++ break ++ assert "drain" in names, f"rank {rank} never drained: {names}" ++ assert not warnings, f"executor escalated past the grace period: {warnings}" ++ assert elapsed <= _EXIT_BOUND_S, ( ++ f"workers took {elapsed:.1f}s to exit after the death pipes closed" ++ ) ++ assert [proc.exitcode for proc in procs] == [0, 0] ++ ++ ++def _stuck_worker_child(ready_pipe, death_pipe, events_pipe, parent_pid): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ def record(name): ++ events_pipe.send((name, time.monotonic())) ++ ++ _patch_executor(executor_module, record, None, 0.0) ++ ++ def stuck_busy_loop(self): ++ # Never leaves the loop and ignores the queue shutdown (a hung RPC). ++ record("busy") ++ while True: ++ time.sleep(0.05) ++ ++ executor_module.WorkerProc.worker_busy_loop = stuck_busy_loop ++ executor_module.WorkerProc.worker_main( ++ vllm_config=_fake_config(), ++ ready_pipe=ready_pipe, ++ death_pipe=death_pipe, ++ inherited_fds=[], ++ expected_parent_pid=parent_pid, ++ ) ++ ++ ++def test_ladder_bounds_a_worker_stuck_in_its_busy_loop(monkeypatch): ++ """Documents (does not change) the bounded ladder. ++ ++ A READY worker that never leaves its busy loop after the death pipe ++ closes: the death-pipe monitor has already set shutdown_requested, so the ++ executor's SIGTERM is a no-op in the worker's handler and the ladder ends ++ in SIGKILL -- that rank's FINAL DRAIN does not run. The bound is the ++ property: grace + 4 s, never longer. ++ """ ++ import vllm.envs as envs ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ grace_s = 1 ++ monkeypatch.setattr( ++ envs, "VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", grace_s, raising=False ++ ) ++ ctx = multiprocessing.get_context("spawn") ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ death_reader, death_writer = ctx.Pipe(duplex=False) ++ events_reader, events_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_stuck_worker_child, ++ args=(ready_writer, death_reader, events_writer, os.getpid()), ++ daemon=True, ++ ) ++ proc.start() ++ ready_writer.close() ++ death_reader.close() ++ events_writer.close() ++ assert ready_reader.poll(180), "worker child never became READY" ++ assert ready_reader.recv()["status"] == "READY" ++ ready_reader.close() ++ ++ warnings: list[str] = [] ++ ++ class _Capture(logging.Handler): ++ def emit(self, record): ++ warnings.append(record.getMessage()) ++ ++ handler = _Capture(level=logging.WARNING) ++ executor_module.logger.addHandler(handler) ++ try: ++ start = time.monotonic() ++ death_writer.close() ++ executor_module.MultiprocExecutor._ensure_worker_termination([proc]) ++ elapsed = time.monotonic() - start ++ finally: ++ executor_module.logger.removeHandler(handler) ++ if proc.is_alive(): ++ proc.kill() ++ proc.join(10) ++ names = [] ++ while True: ++ try: ++ if not events_reader.poll(0.5): ++ break ++ names.append(events_reader.recv()[0]) ++ except (EOFError, OSError): ++ break ++ assert "busy" in names, names ++ assert "drain" not in names, names # the documented gap: SIGKILL, no drain ++ assert any("after grace period" in w for w in warnings), warnings ++ assert any("sending SIGKILL" in w for w in warnings), warnings ++ assert proc.exitcode == -9, proc.exitcode ++ assert elapsed <= grace_s + 4 + 2, f"ladder took {elapsed:.1f}s" +diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py +index 110f56d..72da32e 100644 +--- a/vllm/v1/engine/core.py ++++ b/vllm/v1/engine/core.py +@@ -4,12 +4,13 @@ import gc + import os + import queue + import signal ++import sys + import threading + import time + from collections import defaultdict, deque + from collections.abc import Callable, Generator, Sequence + from concurrent.futures import Future +-from contextlib import ExitStack, contextmanager ++from contextlib import ExitStack, contextmanager, suppress + from enum import IntEnum + from functools import partial + from inspect import isclass, signature +@@ -102,6 +103,76 @@ HANDSHAKE_TIMEOUT_MINS = 5 + _R = TypeVar("_R") # Return type for collective_rpc + + ++# How long EngineCore teardown may run after it deferred a late SIGTERM/SIGINT ++# before the process exits anyway. Covers the whole of engine_core.shutdown(), ++# not only the executor's worker ladder (grace + 4 s): a hang in later teardown ++# (scheduler, distributed cleanup) must not make the process unkillable. ++_TEARDOWN_SIGNAL_BOUND_S = 30.0 ++ ++ ++class _TeardownSignalGuard: ++ """EngineCore teardown: defer ONE late SIGTERM/SIGINT, bounded. ++ ++ Dying on the first late signal would SIGKILL every worker ++ (PR_SET_PDEATHSIG) before its shutdown drain; see the finally block of ++ EngineCoreProc.run_engine_core. The deferral is bounded two ways: ++ - a SECOND late signal restores the default disposition and re-delivers ++ itself, so the process dies of it (the maintainer's escalation); ++ - after the first deferred signal a watchdog ends the process with ++ exit code 1 if teardown is still running _TEARDOWN_SIGNAL_BOUND_S later. ++ """ ++ ++ def __init__(self, bound_s: float | None = None) -> None: ++ self.bound_s = _TEARDOWN_SIGNAL_BOUND_S if bound_s is None else bound_s ++ self.deferred = 0 ++ self._done = threading.Event() ++ self._watchdog: threading.Thread | None = None ++ ++ def __call__(self, signum, frame) -> None: ++ name = signal.Signals(signum).name ++ if self._done.is_set(): ++ return ++ self.deferred += 1 ++ if self.deferred == 1: ++ logger.warning( ++ "[shutdown] EngineCore: %s during resource teardown deferred; " ++ "workers finish their shutdown (FINAL DRAIN) first (a second " ++ "signal, or %.0f s of teardown, ends the process)", ++ name, ++ self.bound_s, ++ ) ++ self._watchdog = threading.Thread( ++ target=self._expire, ++ daemon=True, ++ name="EngineCoreTeardownWatchdog", ++ ) ++ self._watchdog.start() ++ return ++ logger.warning( ++ "[shutdown] EngineCore: second late signal (%s) during resource " ++ "teardown; not deferring again", ++ name, ++ ) ++ signal.signal(signum, signal.SIG_DFL) ++ os.kill(os.getpid(), signum) ++ ++ def _expire(self) -> None: ++ if self._done.wait(self.bound_s): ++ return ++ logger.error( ++ "[shutdown] EngineCore: resource teardown still running %.0f s " ++ "after a deferred signal; exiting (code 1)", ++ self.bound_s, ++ ) ++ for stream in (sys.stdout, sys.stderr): ++ with suppress(Exception): ++ stream.flush() ++ os._exit(1) ++ ++ def disarm(self) -> None: ++ self._done.set() ++ ++ + class EngineCore: + """Inner loop of vLLM's Engine.""" + +@@ -1371,12 +1442,29 @@ class EngineCoreProc(EngineCore): + engine_core._send_engine_dead() + raise e + finally: +- signal.signal(signal.SIGTERM, signal.SIG_DFL) +- signal.signal(signal.SIGINT, signal.SIG_DFL) ++ # Workers carry PR_SET_PDEATHSIG=SIGKILL (fail-closed PLE teardown), ++ # so if this process dies before model_executor.shutdown() has reaped ++ # them, every worker is SIGKILLed mid-shutdown and its bounds-guard ++ # FINAL DRAIN never runs. The API server's process manager sends a ++ # second SIGTERM right after the signal that started this shutdown; ++ # with SIG_DFL here it killed teardown before the executor waited for ++ # the workers. Defer ONE late signal until teardown ends; a second ++ # signal or _TEARDOWN_SIGNAL_BOUND_S of teardown after the deferral ++ # still ends the process (_TeardownSignalGuard), as do the ++ # executor's terminate/kill ladder and the process manager's timeout ++ # (kill_process_tree -> SIGKILL). ++ teardown_guard = _TeardownSignalGuard() ++ signal.signal(signal.SIGTERM, teardown_guard) ++ signal.signal(signal.SIGINT, teardown_guard) + if signal_callback is not None: + signal_callback.stop() +- if engine_core is not None: +- engine_core.shutdown() ++ try: ++ if engine_core is not None: ++ engine_core.shutdown() ++ finally: ++ signal.signal(signal.SIGTERM, signal.SIG_DFL) ++ signal.signal(signal.SIGINT, signal.SIG_DFL) ++ teardown_guard.disarm() + if clean_shutdown: + from vllm.platforms import current_platform + +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index fc744cb..0c5af34 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -114,6 +114,65 @@ def _install_worker_parent_death_signal(expected_parent_pid: int) -> None: + os._exit(1) + + ++# Grace and post-SIGTERM wait, after a worker's death, for survivors that never ++# started their shutdown (see MultiprocExecutor._ensure_worker_termination). ++_FAILED_WORKER_WAIT_S = 1.0 ++# After a worker's death, the bounded grace for survivors whose shutdown has ++# STARTED (FINAL DRAIN in progress). Every FINAL DRAIN measured on the reference host took ++# < 0.1 s (drain-path tests: the drain line lands within ~0.05 s of the ++# shutdown request); 2 s is 20x that, so a draining rank is not cut ++# (a shorter cap did), while one blocked inside its shutdown on the dead rank is ended. ++_FAILED_DRAIN_GRACE_S = 2.0 ++# Per-worker shared byte (WorkerProcHandle.shutdown_state): written by the ++# worker in worker_main's finally, read by the executor's termination ladder. ++_SHUTDOWN_NOT_STARTED = 0 ++_SHUTDOWN_STARTED = 1 ++_SHUTDOWN_DONE = 2 ++ ++ ++class WorkerShutdownRequested(RuntimeError): ++ """The busy loop's dequeue was cancelled by a shutdown request. ++ ++ Raised only by WorkerProc.worker_busy_loop when the RPC broadcast queue ++ reports its own cancellation (shm_broadcast: RuntimeError("cancelled") ++ after MessageQueue.shutdown(), which the death-pipe monitor calls once the ++ parent closes the pipe). worker_main takes the orderly exit path on this ++ type only; any other exception is a failure, even one that races a ++ shutdown request. Subclasses RuntimeError so existing catchers still match. ++ """ ++ ++ ++def _exit_worker_process(exit_code: int) -> None: ++ """End a worker process whose orderly shutdown has already run. ++ ++ Called after WorkerProc.shutdown(destroy_distributed=False): the FINAL ++ DRAIN, the instrumentation flushes and the PLE offload / shared-memory ++ cleanup are done. What would remain is process-exit-only teardown -- ++ destroying the NCCL communicators (PyNcclCommunicator.destroy joins its ++ ncclCommAbort thread for up to 5 s per communicator while CUDA graphs that ++ captured NCCL ops are still alive) and interpreter finalization (CUDA ++ context, allocator pools). The driver and kernel reclaim all of it at ++ process exit, and without this the executor's 5 s grace expired and the ++ worker was SIGKILLed at ~9 s anyway. Flush the logs, ++ then leave. ++ """ ++ logger.info( ++ "[shutdown] WorkerProc: shutdown complete; exiting process (code %d) " ++ "without distributed teardown", ++ exit_code, ++ ) ++ try: ++ import logging ++ ++ logging.shutdown() ++ except Exception: # noqa: BLE001 -- the exit must happen regardless ++ pass ++ for stream in (sys.stdout, sys.stderr): ++ with suppress(Exception): ++ stream.flush() ++ os._exit(exit_code) ++ ++ + class FutureWrapper(Future): + def __init__( + self, +@@ -491,10 +550,36 @@ class MultiprocExecutor(Executor): + return future if non_block else future.result() + + @staticmethod +- def _ensure_worker_termination(worker_procs: list[BaseProcess]): ++ def _ensure_worker_termination( ++ worker_procs: list[BaseProcess], ++ *, ++ after_failure: bool = False, ++ shutdown_states: list[Any] | None = None, ++ ): + """Ensure that all worker processes are terminated. Assumes workers have + received termination requests. Waits for processing, then sends +- termination and kill signals if needed.""" ++ termination and kill signals if needed. ++ ++ after_failure: a worker already died. Survivors split in two. Those ++ that never started WorkerProc.shutdown() are parked in a collective ++ with the dead rank (NCCL / gloo, in C, where neither the grace nor a ++ SIGTERM reaches them -- in a fault-injection test two ranks ran the full ++ 5 s + 4 s ladder to SIGKILL): they get _FAILED_WORKER_WAIT_S to start ++ their shutdown, then SIGTERM, then SIGKILL after as long again. Those ++ whose shutdown_states entry shows the shutdown started (FINAL DRAIN, ++ flushes) get _FAILED_DRAIN_GRACE_S from the start of the ladder, then ++ SIGTERM, then SIGKILL after _FAILED_WORKER_WAIT_S: a started shutdown ++ can itself block on the dead rank (in a drain fault-injection test one ++ survivor had started its shutdown and never logged its FINAL DRAIN -- ++ the drain's device synchronize waits on work paired with the dead ++ rank -- and ran the full orderly ladder, workers gone 17.6 s after the ++ fault). Before the drain-order fix the engine died on the late SIGTERM and ++ PDEATHSIG killed all survivors at once. A worker with no state (None, ++ or no shutdown_states) counts as parked. The orderly ladder (no ++ failure) is unchanged. ++ """ ++ grace_s = float(envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS) ++ term_wait_s = 4.0 + + def wait_for_termination(procs, timeout): + if not time: +@@ -516,9 +601,42 @@ class MultiprocExecutor(Executor): + "[shutdown] Executor: waiting for worker exit count=%d", + initial_count, + ) +- if wait_for_termination( +- active_procs(), timeout=envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS +- ): ++ if after_failure and time: ++ start = time.time() ++ states = list(shutdown_states or [None] * len(worker_procs)) ++ ++ def parked() -> list[BaseProcess]: ++ return [ ++ proc ++ for proc, state in zip(worker_procs, states) ++ if proc.is_alive() ++ and (state is None or state.value == _SHUTDOWN_NOT_STARTED) ++ ] ++ ++ if not wait_for_termination(active_procs(), _FAILED_WORKER_WAIT_S): ++ stuck = parked() ++ if stuck: ++ logger.warning( ++ "[shutdown] Executor: workers parked after a worker " ++ "failure (shutdown not started); sending SIGTERM count=%d", ++ len(stuck), ++ ) ++ for p in stuck: ++ p.terminate() ++ if not wait_for_termination(stuck, _FAILED_WORKER_WAIT_S): ++ stuck = [p for p in stuck if p.is_alive()] ++ logger.warning( ++ "[shutdown] Executor: parked workers still running " ++ "after SIGTERM; sending SIGKILL count=%d", ++ len(stuck), ++ ) ++ for p in stuck: ++ p.kill() ++ # Workers that started their shutdown get a BOUNDED grace, counted ++ # from the start of this ladder, then SIGTERM, then SIGKILL. ++ grace_s = max(0.0, _FAILED_DRAIN_GRACE_S - (time.time() - start)) ++ term_wait_s = _FAILED_WORKER_WAIT_S ++ if wait_for_termination(active_procs(), timeout=grace_s): + logger.info_once("[shutdown] Executor: all workers exited gracefully") + return + +@@ -531,7 +649,7 @@ class MultiprocExecutor(Executor): + ) + for p in remaining: + p.terminate() +- if not wait_for_termination(active_procs(), 4): ++ if not wait_for_termination(active_procs(), term_wait_s): + # Send SIGKILL if still running + remaining = active_procs() + logger.warning( +@@ -559,7 +677,13 @@ class MultiprocExecutor(Executor): + if w.death_writer is not None: + w.death_writer.close() + w.death_writer = None +- self._ensure_worker_termination([w.proc for w in workers]) ++ self._ensure_worker_termination( ++ [w.proc for w in workers], ++ after_failure=getattr(self, "is_failed", False), ++ shutdown_states=[ ++ getattr(w, "shutdown_state", None) for w in workers ++ ], ++ ) + + for w in workers: + # Shutdown response queues +@@ -610,6 +734,7 @@ class UnreadyWorkerProcHandle: + rank: int + ready_pipe: Connection + death_writer: Connection | None = None ++ shutdown_state: Any = None + + + @dataclass +@@ -623,6 +748,8 @@ class WorkerProcHandle: + # `peer_worker_response_mqs[i]` + peer_worker_response_mqs: list[MessageQueue | None] + death_writer: Connection | None = None ++ # Shared byte: _SHUTDOWN_NOT_STARTED / _STARTED / _DONE (worker writes). ++ shutdown_state: Any = None + + @classmethod + def from_unready_handle( +@@ -637,6 +764,7 @@ class WorkerProcHandle: + worker_response_mq=worker_response_mq, + peer_worker_response_mqs=peer_worker_response_mqs, + death_writer=unready_handle.death_writer, ++ shutdown_state=unready_handle.shutdown_state, + ) + + +@@ -775,6 +903,8 @@ class WorkerProc: + ready_reader, ready_writer = context.Pipe(duplex=False) + # Death pipe to let child detect parent process exit + death_reader, death_writer = context.Pipe(duplex=False) ++ # One shared byte: has this worker started / finished its shutdown? ++ shutdown_state = context.RawValue("b", _SHUTDOWN_NOT_STARTED) + if inherited_fds is not None: + inherited_fds = inherited_fds.copy() + inherited_fds.extend((ready_reader.fileno(), death_writer.fileno())) +@@ -787,6 +917,7 @@ class WorkerProc: + "input_shm_handle": input_shm_handle, + "ready_pipe": ready_writer, + "death_pipe": death_reader, ++ "shutdown_state": shutdown_state, + "shared_worker_lock": shared_worker_lock, + "is_driver_worker": is_driver_worker, + # Have the worker close parent end of this worker's pipes too +@@ -811,7 +942,9 @@ class WorkerProc: + death_reader.close() + # Keep death_writer open in parent - when parent exits, + # death_reader in child will get EOFError +- return UnreadyWorkerProcHandle(proc, rank, ready_reader, death_writer) ++ return UnreadyWorkerProcHandle( ++ proc, rank, ready_reader, death_writer, shutdown_state ++ ) + + @staticmethod + def wait_for_response_handle_ready( +@@ -872,7 +1005,14 @@ class WorkerProc: + + return cast(list[WorkerProcHandle], ready_proc_handles) + +- def shutdown(self): ++ def shutdown(self, *, destroy_distributed: bool = True): ++ if not destroy_distributed: ++ # worker_main os._exits right after this: the PLE offload handle ++ # (rank 0 owns the child) must not wait for the child's reap, ++ # which trails its exit by seconds for a large host table. ++ ple_worker = sys.modules.get("vllm.v1.ple_offload.worker") ++ if ple_worker is not None: ++ ple_worker.mark_process_exit_imminent() + if self.rpc_broadcast_mq is not None: + self.rpc_broadcast_mq.shutdown() + if self.worker_response_mq is not None: +@@ -880,8 +1020,12 @@ class WorkerProc: + self.worker.shutdown() + self.rpc_broadcast_mq = None + self.worker_response_mq = None +- destroy_model_parallel() +- destroy_distributed_environment() ++ # worker_main passes False when the process exits right after this ++ # (see _exit_worker_process): the NCCL destroy is exit-only work that ++ # can block for 5 s per communicator. ++ if destroy_distributed: ++ destroy_model_parallel() ++ destroy_distributed_environment() + + def monitor_death_pipe(self, death_pipe, shutdown_requested: threading.Event): + if death_pipe is None: +@@ -948,8 +1092,20 @@ class WorkerProc: + set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"]) + + worker = None ++ # Set once a READY worker leaves its busy loop because shutdown was ++ # requested (death pipe EOF or SIGTERM/SIGINT): the process then exits ++ # straight after WorkerProc.shutdown (see _exit_worker_process). ++ orderly_shutdown = False ++ # Set when a READY worker's busy loop FAILED while a shutdown request ++ # was pending: full teardown, then exit code 1 (never the exit-0 path). ++ failed_during_shutdown = False + ready_writer = kwargs.pop("ready_pipe") + death_pipe = kwargs.pop("death_pipe", None) ++ shutdown_state = kwargs.pop("shutdown_state", None) ++ ++ def mark_shutdown(value: int) -> None: ++ if shutdown_state is not None: ++ shutdown_state.value = value + + # Close inherited pipes from parent (incl. other worker pipes) + # Explicitly passing in existing pipes and closing them makes the pipe +@@ -996,7 +1152,7 @@ class WorkerProc: + + worker.worker_busy_loop() + +- except Exception: ++ except Exception as e: + # NOTE: if an Exception arises in busy_loop, we send + # a FAILURE message over the MQ RPC to notify the Executor, + # which triggers system shutdown. +@@ -1004,10 +1160,20 @@ class WorkerProc: + + if ready_writer is not None: + logger.exception("WorkerProc failed to start.") +- elif shutdown_requested.is_set(): ++ elif isinstance(e, WorkerShutdownRequested) and ( ++ shutdown_requested.is_set() ++ ): ++ # The busy loop left BECAUSE of the shutdown request (its ++ # dequeue was cancelled), not because something failed. + logger.debug_once( + "[shutdown] WorkerProc: exiting after shutdown request" + ) ++ orderly_shutdown = True ++ elif shutdown_requested.is_set(): ++ # A real failure that raced a shutdown request: keep it a ++ # failure -- full teardown, then a non-zero exit. ++ logger.exception("WorkerProc failed during shutdown.") ++ failed_during_shutdown = True + else: + logger.exception("WorkerProc failed.") + +@@ -1023,6 +1189,7 @@ class WorkerProc: + logger.debug_once( + "[shutdown] WorkerProc: terminated by shutdown signal" + ) ++ orderly_shutdown = ready_writer is None + else: + logger.warning("WorkerProc was terminated") + # SystemExit must never be ignored +@@ -1035,7 +1202,26 @@ class WorkerProc: + death_pipe.close() + # Clean up once worker exits busy loop + if worker is not None: ++ # Tells the executor's ladder this rank is draining, not parked ++ # in a collective (see _ensure_worker_termination). ++ mark_shutdown(_SHUTDOWN_STARTED) ++ if orderly_shutdown: ++ # Everything that must happen before exit (FINAL DRAIN, ++ # instrumentation flushes, PLE cleanup) is in ++ # worker.shutdown(); then exit instead of paying the NCCL ++ # destroy + finalization past the executor's grace. ++ exit_code = 0 ++ try: ++ worker.shutdown(destroy_distributed=False) ++ except BaseException: ++ logger.exception("WorkerProc shutdown failed") ++ exit_code = 1 ++ mark_shutdown(_SHUTDOWN_DONE) ++ _exit_worker_process(exit_code) + worker.shutdown() ++ mark_shutdown(_SHUTDOWN_DONE) ++ if failed_during_shutdown: ++ sys.exit(1) + + class ResponseStatus(Enum): + SUCCESS = auto() +@@ -1116,8 +1302,20 @@ class WorkerProc: + def worker_busy_loop(self): + """Main busy loop for Multiprocessing Workers""" + assert self.rpc_broadcast_mq is not None ++ mq = self.rpc_broadcast_mq + while True: +- self._execute_worker_rpc(self.rpc_broadcast_mq.dequeue(indefinite=True)) ++ # No local holds the request across the next dequeue: its payload ++ # must be released first (test_multiproc_executor). ++ self._execute_worker_rpc(self._dequeue_worker_rpc(mq)) ++ ++ @staticmethod ++ def _dequeue_worker_rpc(mq: MessageQueue): ++ try: ++ return mq.dequeue(indefinite=True) ++ except RuntimeError as e: ++ if getattr(mq, "shutting_down", False) and str(e) == "cancelled": ++ raise WorkerShutdownRequested("RPC broadcast queue shut down") from e ++ raise + + def _execute_worker_rpc( + self, +diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py +index 3d29407..608f364 100644 +--- a/vllm/v1/ple_offload/worker.py ++++ b/vllm/v1/ple_offload/worker.py +@@ -102,6 +102,86 @@ def _unlink_ipc_socket(ipc_addr: str) -> None: + logger.warning("Failed to unlink PLE IPC socket %s: %s", path, error) + + ++# Set once by a GPU worker process whose orderly shutdown is about to end in ++# os._exit (multiproc_executor._exit_worker_process). PleOffloadWorkerHandle ++# then waits for the offload child's whole thread group to be in the kernel's ++# exit path, not for its REAP. In a drain fault-injection test the ~51 GB ++# child's leader went Z at ~+2.7 s, but one thread stayed R -- PF_EXITING set, ++# in-kernel, freeing the address space -- until the reap at ~+8.6 s. Rank 0 ++# waiting for that reap ran the executor's grace -> SIGTERM -> SIGKILL ladder on ++# every cycle. A read-only sample of one such shutdown showed ++# the same thing: from Z to reap (6.0 s), the only live task was R with ++# PF_EXITING, the leader's fd table was already closed (0 /dev/nvidia fds, ++# down from 13), and the child was never an nvidia-smi compute app, before or ++# after. It holds no GPU memory (host-mapped pull: no CUDA context). What is ++# left is kernel-side host-memory teardown, and nothing in this process ++# needs to wait for it. ++_PROCESS_EXIT_IMMINENT = False ++# Bound for the child's userspace cleanup (its proc_main finally) on that path; ++# it took ~2.7 s on the reference host. It has to fit inside the executor's 5 s ++# VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS with room for rank 0's own exit. After it, ++# the child is SIGKILLed (userspace ends at once) and not waited for. ++_EXIT_IMMINENT_CHILD_WAIT_S = 3.5 ++# procfs mount the predicate reads (tests point it at a fake tree). ++_PROC_ROOT = "/proc" ++# task_struct flag: set by do_exit() -> exit_signals() before exit_mm(); a task ++# carrying it never returns to userspace. ++_PF_EXITING = 0x4 ++ ++ ++def mark_process_exit_imminent() -> None: ++ """This process os._exits right after its shutdown; see close().""" ++ global _PROCESS_EXIT_IMMINENT ++ _PROCESS_EXIT_IMMINENT = True ++ ++ ++def _task_in_exit(stat: bytes) -> bool: ++ """One /proc//task//stat: Z/X, or PF_EXITING (kernel exit path).""" ++ # Fields after the parenthesised comm (which may contain spaces/parens): ++ # [0]=state ... [6]=flags (stat field 9). ++ fields = stat[stat.rfind(b")") + 2 :].split() ++ if not fields: ++ return True ++ if fields[0] in (b"Z", b"X"): ++ return True ++ try: ++ return bool(int(fields[6]) & _PF_EXITING) ++ except (IndexError, ValueError): ++ return False ++ ++ ++def _child_thread_group_exiting( ++ proc: subprocess.Popen, proc_root: str | None = None ++) -> bool: ++ """True once EVERY task of ``proc`` is exited or in the kernel exit path. ++ ++ A zombie leader alone is not enough (a thread can still run userspace ++ after the leader exits). Never reaps: the child is unreaped, so its PID ++ cannot be reused and //task is its own. Off Linux, only a ++ reaped child counts. ++ """ ++ if proc.returncode is not None: ++ return True ++ task_dir = os.path.join(proc_root or _PROC_ROOT, str(proc.pid), "task") ++ try: ++ tids = os.listdir(task_dir) ++ except FileNotFoundError: ++ return True ++ except OSError: ++ return proc.poll() is not None ++ for tid in tids: ++ try: ++ with open(os.path.join(task_dir, tid, "stat"), "rb") as fh: ++ stat = fh.read() ++ except FileNotFoundError: ++ continue # this task finished exiting between listdir and open ++ except OSError: ++ return False ++ if not _task_in_exit(stat): ++ return False ++ return True ++ ++ + def _publish_fence() -> None: + with _PUBLISH_FENCE: + pass +@@ -448,6 +528,9 @@ class PleOffloadWorkerHandle: + if self.ready_pipe_reader is not None: + self.ready_pipe_reader.close() + self.ready_pipe_reader = None ++ if graceful and _PROCESS_EXIT_IMMINENT: ++ self._close_without_reap() ++ return + if graceful and self.proc.poll() is None: + # Keep parent-death supervision connected while SIGTERM lets + # proc_main run its shutdown event and finally cleanup. +@@ -472,6 +555,41 @@ class PleOffloadWorkerHandle: + if self.proc.poll() is not None and self.ipc_addr is not None: + _unlink_ipc_socket(self.ipc_addr) + ++ def _close_without_reap(self) -> None: ++ """close() for a parent that os._exits next: wait (bounded) until every ++ child task has left userspace, never for the reap. ++ ++ SIGTERM lets proc_main's finally unlink the IPC socket, as in the ++ graceful path; parent-death supervision stays connected meanwhile. ++ Once the child has run exit, the address-space teardown that keeps it ++ unreapable needs nothing from this process: the orphan is reaped by ++ the subreaper/init after this process exits. ++ """ ++ if not _child_thread_group_exiting(self.proc): ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(self.proc.pid, signal.SIGTERM) ++ deadline = time.monotonic() + _EXIT_IMMINENT_CHILD_WAIT_S ++ while not _child_thread_group_exiting(self.proc): ++ if time.monotonic() >= deadline: ++ logger.warning( ++ "PLE offload process %d still running %.1f s after " ++ "SIGTERM; killing it without waiting (parent exits)", ++ self.proc.pid, ++ _EXIT_IMMINENT_CHILD_WAIT_S, ++ ) ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(self.proc.pid, signal.SIGKILL) ++ break ++ time.sleep(0.01) ++ exited = _child_thread_group_exiting(self.proc) ++ if self.death_writer is not None: ++ self.death_writer.close() ++ self.death_writer = None ++ # Same fallback as close(): the child has run exit, so this exact ++ # PLE-owned address is no longer in use. ++ if exited and self.ipc_addr is not None: ++ _unlink_ipc_socket(self.ipc_addr) ++ + + def _init_offload_distributed() -> None: + """Initialize the single-rank Gloo world required by TP-aware layers.""" +-- +2.43.0 + diff --git a/patches/0058-executor-detect-a-fail-closed-fault-or-death-on-any-.patch b/patches/0058-executor-detect-a-fail-closed-fault-or-death-on-any-.patch new file mode 100644 index 0000000..04cda71 --- /dev/null +++ b/patches/0058-executor-detect-a-fail-closed-fault-or-death-on-any-.patch @@ -0,0 +1,2100 @@ +From 36980a4bc6f822c2458301327003cedb7c6d183c Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 24 Sep 2026 07:58:38 -0400 +Subject: [PATCH 058/103] executor: detect a fail-closed fault or death on any + rank + +A fault on a non-zero TP rank could hang the engine: the worker's sentinel pipe +was inherited by a helper process and stayed open, so the monitor never saw +the exit. Earlier releases carry the same monitor and fail-closed exit path. + +- Worker exit is detected by pidfd, with the sentinel as fallback, before and + after READY and for every worker through the whole initialisation window. +- A per-worker fault pipe: ple_fail_closed writes one byte before os._exit and + the monitor starts the failure ladder at once; the fd stays registered + through worker teardown. +- worker_main marks inherited pipes close-on-exec. +- Any unexpected busy-loop failure exits 1; the termination ladder uses + time.monotonic(); forced-exit paths no longer wait on logging. +- Tests on real fork/spawn workers for rank 0 and rank 1, with and without an + fd-inheriting helper. +--- + .../test_engine_core_teardown_signal.py | 35 ++ + tests/v1/executor/test_executor.py | 2 + + .../test_multiproc_failed_worker_ladder.py | 56 ++- + .../executor/test_multiproc_fatal_any_rank.py | 343 +++++++++++++++ + .../test_multiproc_monitor_any_rank.py | 226 ++++++++++ + .../test_multiproc_monitor_combined.py | 147 +++++++ + .../test_multiproc_ready_and_inherit.py | 355 +++++++++++++++ + .../test_multiproc_worker_exit_after_drain.py | 28 +- + vllm/v1/engine/core.py | 53 ++- + vllm/v1/executor/multiproc_executor.py | 410 ++++++++++++++---- + vllm/v1/ple_offload/connector.py | 3 +- + vllm/v1/ple_offload/fatal.py | 35 ++ + 12 files changed, 1587 insertions(+), 106 deletions(-) + create mode 100644 tests/v1/executor/test_multiproc_fatal_any_rank.py + create mode 100644 tests/v1/executor/test_multiproc_monitor_any_rank.py + create mode 100644 tests/v1/executor/test_multiproc_monitor_combined.py + create mode 100644 tests/v1/executor/test_multiproc_ready_and_inherit.py + +diff --git a/tests/v1/engine/test_engine_core_teardown_signal.py b/tests/v1/engine/test_engine_core_teardown_signal.py +index 94a5255..822c735 100644 +--- a/tests/v1/engine/test_engine_core_teardown_signal.py ++++ b/tests/v1/engine/test_engine_core_teardown_signal.py +@@ -35,6 +35,9 @@ def _child_run(conn, mode: str = "once") -> None: + + from vllm.v1.engine import core + ++ blocked_log = mode.endswith("-blocked-log") ++ if blocked_log: ++ mode = mode[: -len("-blocked-log")] + if mode == "once-hang": + core._TEARDOWN_SIGNAL_BOUND_S = 1.0 + +@@ -48,6 +51,17 @@ def _child_run(conn, mode: str = "once") -> None: + + core.logger.addHandler(_Capture(level=logging.WARNING)) + ++ class _Blocking(logging.Handler): ++ """A stderr pipe nobody drains: the forced-exit lines never return.""" ++ ++ def emit(self, record): ++ message = record.getMessage() ++ if "second late signal" in message or "still running" in message: ++ time.sleep(3600) ++ ++ if blocked_log: ++ core.logger.addHandler(_Blocking(level=logging.WARNING)) ++ + class _FakeEngineCore: + def __init__(self, *args, **kwargs): + import queue +@@ -184,6 +198,27 @@ def test_second_late_signal_ends_a_hung_teardown(): + assert "teardown-done" not in events, events + + ++def test_second_late_signal_is_not_held_by_a_blocked_log(): ++ # An earlier second-signal path logged BEFORE restoring the default ++ # disposition and re-delivering; a blocked handler / full pipe held it. ++ alive, exitcode, elapsed, events = _run_hanging_teardown( ++ "twice-hang-blocked-log", 15 ++ ) ++ assert "second-signal" in events, events ++ assert not alive, f"a blocked log held the second-signal exit ({elapsed:.1f}s)" ++ assert exitcode == -signal.SIGTERM, exitcode ++ ++ ++def test_teardown_watchdog_is_not_held_by_a_blocked_log(): ++ # An earlier watchdog logged and flushed before os._exit(1). ++ alive, exitcode, elapsed, events = _run_hanging_teardown( ++ "once-hang-blocked-log", 15 ++ ) ++ assert "teardown-start" in events, events ++ assert not alive, f"a blocked log held the watchdog exit ({elapsed:.1f}s)" ++ assert exitcode == 1, exitcode ++ ++ + def test_teardown_watchdog_bounds_a_single_deferred_signal(): + # One deferred SIGTERM, then teardown hangs: the watchdog ends the process + # (exit code 1) once _TEARDOWN_SIGNAL_BOUND_S (1 s here) has passed. +diff --git a/tests/v1/executor/test_executor.py b/tests/v1/executor/test_executor.py +index a21a8da..8e5d9fe 100644 +--- a/tests/v1/executor/test_executor.py ++++ b/tests/v1/executor/test_executor.py +@@ -82,6 +82,8 @@ def test_multiproc_executor_worker_termination_timeout( + monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", str(timeout)) + clock = _FakeClock() + monkeypatch.setattr(multiproc_executor_module.time, "time", clock.time) ++ # The ladder times its waits with time.monotonic(). ++ monkeypatch.setattr(multiproc_executor_module.time, "monotonic", clock.time) + monkeypatch.setattr(multiproc_executor_module.time, "sleep", clock.sleep) + executor = MultiprocExecutor.__new__(MultiprocExecutor) + proc = _FakeProcess(clock, exits_at=exits_at) +diff --git a/tests/v1/executor/test_multiproc_failed_worker_ladder.py b/tests/v1/executor/test_multiproc_failed_worker_ladder.py +index 9c59113..c67d593 100644 +--- a/tests/v1/executor/test_multiproc_failed_worker_ladder.py ++++ b/tests/v1/executor/test_multiproc_failed_worker_ladder.py +@@ -133,6 +133,35 @@ def test_failure_shutdown_lets_a_draining_survivor_finish(): + assert _DRAINING_S - 0.3 <= elapsed < 3.0, elapsed + + ++def test_failure_ladder_ignores_a_wall_clock_step(monkeypatch): ++ """An earlier ladder timed its graces with time.time(). ++ ++ A wall-clock step forward (NTP, maintainer) mid-ladder made every wait ++ expire at once, so a survivor still inside its FINAL DRAIN was SIGTERMed. ++ """ ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ real_time = time.time ++ calls = {"n": 0} ++ ++ def stepping_time(): ++ # Steps forward an hour between any two reads. ++ calls["n"] += 1 ++ return real_time() + 3600.0 * calls["n"] ++ ++ monkeypatch.setattr( ++ executor_module, ++ "time", ++ SimpleNamespace(time=stepping_time, monotonic=time.monotonic, sleep=time.sleep), ++ ) ++ elapsed, codes = _shutdown([(_DRAINING, 1), (_PARKED, 0)], failed=True) ++ assert codes[0] == 0, ( ++ f"a wall-clock step cut the draining survivor (exit {codes[0]})" ++ ) ++ assert codes[1] is not None ++ assert _DRAINING_S - 0.3 <= elapsed < 3.0, elapsed ++ ++ + def test_failure_shutdown_ends_a_survivor_stuck_in_its_shutdown(): + # Started its shutdown (state 1), then blocked on the dead rank: ignores + # SIGTERM like a rank inside a device synchronize, never exits by itself. +@@ -195,18 +224,21 @@ def test_worker_main_publishes_shutdown_started_then_done(monkeypatch): + ) + monkeypatch.setattr(executor_module.signal, "signal", lambda *a, **k: None) + ready_reader, ready_writer = multiprocessing.Pipe(duplex=False) +- executor_module.WorkerProc.worker_main( +- vllm_config=SimpleNamespace( +- parallel_config=SimpleNamespace( +- assigned_physical_gpu_ids=None, numa_bind=False +- ) +- ), +- ready_pipe=ready_writer, +- death_pipe=None, +- shutdown_state=state, +- inherited_fds=[], +- expected_parent_pid=os.getpid(), +- ) ++ # An unexpected busy-loop failure ends in exit code 1 after the teardown. ++ with pytest.raises(SystemExit) as exited: ++ executor_module.WorkerProc.worker_main( ++ vllm_config=SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, numa_bind=False ++ ) ++ ), ++ ready_pipe=ready_writer, ++ death_pipe=None, ++ shutdown_state=state, ++ inherited_fds=[], ++ expected_parent_pid=os.getpid(), ++ ) ++ assert exited.value.code == 1 + ready_reader.close() + assert seen == [("busy", 0), ("drain", 1)], seen + assert state.value == 2 +diff --git a/tests/v1/executor/test_multiproc_fatal_any_rank.py b/tests/v1/executor/test_multiproc_fatal_any_rank.py +new file mode 100644 +index 0000000..bf1a0ff +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_fatal_any_rank.py +@@ -0,0 +1,343 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""A fail-closed fault on ANY rank reaches the executor before the process is gone. ++ ++Rank-1 fault test (earlier build, TP2xPP2): the PLE ++test fault cut=exc_in_release fired on PP0_TP1 (rank 1). ple_fail_closed wrote ++its line and traceback and os._exit(1)ed. The engine then sat in ++mq.dequeue for 300 s ("No available shared memory broadcast block found in 60 ++seconds" x4) and was killed by the harness; "Worker proc ... died ++unexpectedly" never appeared. The rank showed Z in ps but was never reaped. ++ ++The executor's monitor (start_worker_monitor, stock upstream code) waits on ++the workers' process sentinels only. A sentinel closes when the kernel has ++torn down the WHOLE thread group -- for a worker holding a CUDA context that ++is seconds after os._exit (F1/F2/F3/F5 on rank 0 in the same series: 8.0 s ++between the fatal line and the monitor's log line), and when a thread of the ++dying rank waits in-kernel on a live peer it does not close while the peer ++lives. The fault test's rank 1 became reapable only after rank 0's own 300 s progress ++deadline killed rank 0 (rank 3's gloo recv from rank 1 broke at +2.0 s, ++two seconds after rank 0's fatal at +0.0 s). The same leader-Z / ++thread-in-exit state was sampled on the reference host for the PLE offload child ++(test_multiproc_ple_exit_without_reap.py). ++ ++The fix: a per-worker fault pipe. ple_fail_closed writes one byte on it ++right before os._exit; the monitor waits on it next to the sentinels and ++starts the failure ladder at once, which ends the peers -- and that is what ++frees the stuck teardown. ++ ++The stuck-teardown state is modelled here without a GPU: the worker's leader ++thread leaves alone (raw SYS_exit) while another thread keeps the thread ++group alive, so its sentinel stays open and it is not reapable -- exactly ++what the executor saw in the rank-1 fault test. ++ ++Failing controls on the earlier build: test_fatal_on_any_rank_wakes_the_ ++monitor_before_its_teardown (the monitor stays blocked) and ++test_worker_main_registers_the_fault_pipe (nothing registers it). ++""" ++ ++import ctypes ++import multiprocessing ++import os ++import platform ++import signal ++import sys ++import threading ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++pytestmark = [ ++ pytest.mark.skip_global_cleanup, ++ pytest.mark.skipif( ++ sys.platform != "linux" or platform.machine() not in ("x86_64", "aarch64"), ++ reason="needs Linux process sentinels and a raw SYS_exit number", ++ ), ++] ++ ++_SYS_EXIT = {"x86_64": 60, "aarch64": 93}.get(platform.machine(), -1) ++_LINGER_S = 30.0 ++_MONITOR_WAIT_S = 5.0 ++ ++ ++def _leader_exit(code: int) -> None: ++ """Only the calling thread leaves; the thread group (and its fds) stay.""" ++ ctypes.CDLL(None).syscall(_SYS_EXIT, code) ++ ++ ++def _worker(mode: str, fault_pipe, linger_s: float) -> None: ++ from vllm.v1.ple_offload import fatal ++ ++ # Absent on the earlier build: that is the failing control, not an error. ++ register = getattr(fatal, "set_fatal_notify_fd", None) ++ if register is not None and fault_pipe is not None: ++ register(fault_pipe.fileno()) ++ if mode == "idle": ++ time.sleep(600) ++ # Non-daemon: keeps the thread group alive after the leader leaves. ++ threading.Thread(target=time.sleep, args=(linger_s,)).start() ++ if mode == "fatal-stuck": ++ # ple_fail_closed ends with os._exit(1); here that exit is the ++ # leader-only one, so the process stays unreapable like the fault test's rank 1. ++ os._exit = _leader_exit # type: ignore[assignment] ++ fatal.ple_fail_closed( ++ where="release_outputs", ++ rank=1, ++ layer="language_model.model.layers.1.ple.ple_embedding", ++ sequence=52, ++ error=RuntimeError("PLE test fault: exc_in_release"), ++ ) ++ if mode == "orderly-close": ++ fault_pipe.close() ++ time.sleep(linger_s) ++ os._exit(1 if mode == "crash" else 0) ++ ++ ++class _Executor: ++ """The surface start_worker_monitor's closure uses.""" ++ ++ def __init__(self, workers): ++ self.workers = workers ++ self.is_failed = False ++ self.shutting_down = False ++ self.shutdown_calls = 0 ++ self.callback_calls = 0 ++ self.failure_callback = self._callback ++ ++ def _callback(self): ++ self.callback_calls += 1 ++ ++ def shutdown(self): ++ self.shutdown_calls += 1 ++ ++ ++def _spawn(mode: str, rank: int, linger_s: float = _LINGER_S): ++ ctx = multiprocessing.get_context("spawn") ++ fault_reader, fault_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_worker, ++ args=(mode, fault_writer, linger_s), ++ name=f"VllmWorker-{rank}", ++ daemon=True, ++ ) ++ proc.start() ++ fault_writer.close() ++ return SimpleNamespace( ++ proc=proc, ++ rank=rank, ++ fault_reader=fault_reader, ++ death_writer=None, ++ worker_response_mq=None, ++ shutdown_state=None, ++ ) ++ ++ ++def _start_monitor(executor) -> threading.Thread: ++ from vllm.v1.executor.multiproc_executor import MultiprocExecutor ++ ++ thread = threading.Thread( ++ target=MultiprocExecutor.start_worker_monitor, ++ args=(executor,), ++ kwargs={"inline": True}, ++ daemon=True, ++ name="MultiprocWorkerMonitor-test", ++ ) ++ thread.start() ++ return thread ++ ++ ++@pytest.fixture ++def workers(): ++ spawned = [] ++ yield spawned ++ for w in spawned: ++ # SIGKILL ends the lingering thread too; then the group is reapable. ++ if w.proc.is_alive(): ++ os.kill(w.proc.pid, signal.SIGKILL) ++ w.proc.join(10) ++ w.fault_reader.close() ++ ++ ++def test_fatal_on_any_rank_wakes_the_monitor_before_its_teardown(workers): ++ """Rank-1 fault: rank 1 fails closed, stays unreapable; the monitor must not wait.""" ++ workers.append(_spawn("idle", 0)) ++ workers.append(_spawn("fatal-stuck", 1)) ++ executor = _Executor(workers) ++ thread = _start_monitor(executor) ++ thread.join(_MONITOR_WAIT_S) ++ assert not thread.is_alive(), ( ++ f"monitor still blocked {_MONITOR_WAIT_S} s after a fail-closed fault on " ++ "rank 1 whose process teardown has not finished (leader gone, thread " ++ "group alive, sentinel open): this is the rank-1 fault hang" ++ ) ++ # It returned while the rank was still not reapable -- the point. ++ assert workers[1].proc.is_alive() ++ assert executor.is_failed ++ assert executor.shutdown_calls == 1 ++ assert executor.callback_calls == 1 ++ assert executor.failure_callback is None ++ ++ ++def test_plain_death_is_still_detected_by_the_sentinel(workers): ++ workers.append(_spawn("idle", 0)) ++ workers.append(_spawn("crash", 1, linger_s=0.0)) ++ executor = _Executor(workers) ++ thread = _start_monitor(executor) ++ thread.join(_MONITOR_WAIT_S) ++ assert not thread.is_alive() ++ assert executor.is_failed ++ assert executor.shutdown_calls == 1 ++ assert workers[1].proc.exitcode == 1 ++ ++ ++def test_orderly_close_of_the_fault_pipe_is_not_a_fault(workers): ++ """EOF on the fault pipe (worker_main's finally) must not start the ladder.""" ++ orderly_s = 1.5 ++ workers.append(_spawn("orderly-close", 0, linger_s=orderly_s)) ++ workers.append(_spawn("idle", 1)) ++ executor = _Executor(workers) ++ thread = _start_monitor(executor) ++ thread.join(orderly_s / 2) ++ assert thread.is_alive(), "monitor treated the closed fault pipe as a fault" ++ assert executor.shutdown_calls == 0 ++ thread.join(_MONITOR_WAIT_S) ++ assert not thread.is_alive(), "sentinel death not seen after the fault-pipe EOF" ++ assert executor.is_failed ++ assert workers[0].proc.exitcode == 0 ++ ++ ++def test_worker_main_registers_the_fault_pipe(monkeypatch): ++ """worker_main hands the fault pipe to the fatal module, and unregisters it.""" ++ import vllm.v1.executor.multiproc_executor as executor_module ++ from vllm.v1.ple_offload import fatal ++ ++ seen = {} ++ ++ class _Inner: ++ def shutdown(self): ++ pass ++ ++ def fake_init(self, *args, **kwargs): ++ self.rank = 1 ++ self.worker = _Inner() ++ self.rpc_broadcast_mq = None ++ self.worker_response_mq = SimpleNamespace( ++ export_handle=lambda: None, ++ wait_until_ready=lambda: None, ++ shutdown=lambda: None, ++ ) ++ self.peer_response_handles = [] ++ ++ def busy_loop(self): ++ seen["registered_fd"] = getattr(fatal, "_FATAL_NOTIFY_FD", None) ++ raise ValueError("boom") ++ ++ monkeypatch.setattr(executor_module.WorkerProc, "__init__", fake_init) ++ monkeypatch.setattr(executor_module.WorkerProc, "worker_busy_loop", busy_loop) ++ monkeypatch.setattr(executor_module, "destroy_model_parallel", lambda: None) ++ monkeypatch.setattr( ++ executor_module, "destroy_distributed_environment", lambda: None ++ ) ++ monkeypatch.setattr(executor_module, "maybe_init_worker_tracer", lambda **k: None) ++ monkeypatch.setattr(executor_module, "set_worker_net_device", lambda *a, **k: None) ++ monkeypatch.setattr( ++ executor_module, "_install_worker_parent_death_signal", lambda pid: None ++ ) ++ monkeypatch.setattr(executor_module.signal, "signal", lambda *a, **k: None) ++ ready_reader, ready_writer = multiprocessing.Pipe(duplex=False) ++ fault_reader, fault_writer = multiprocessing.Pipe(duplex=False) ++ fault_fd = fault_writer.fileno() ++ # An unexpected busy-loop failure ends in exit code 1 after the teardown ++ # (any-rank fix). ++ with pytest.raises(SystemExit) as exited: ++ executor_module.WorkerProc.worker_main( ++ vllm_config=SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, numa_bind=False ++ ) ++ ), ++ ready_pipe=ready_writer, ++ death_pipe=None, ++ fault_pipe=fault_writer, ++ shutdown_state=None, ++ inherited_fds=[], ++ expected_parent_pid=os.getpid(), ++ ) ++ assert exited.value.code == 1 ++ ready_reader.close() ++ assert seen.get("registered_fd") == fault_fd, seen ++ assert getattr(fatal, "_FATAL_NOTIFY_FD", "unset") is None ++ # Closed on the way out: EOF, which the monitor does not count as a fault. ++ assert fault_reader.poll(1.0), "worker_main left its fault pipe open" ++ with pytest.raises(EOFError): ++ fault_reader.recv_bytes() ++ fault_reader.close() ++ ++ ++def test_fail_closed_during_teardown_still_sends_the_fault_byte(monkeypatch): ++ """The fault fd stays registered through worker teardown. ++ ++ worker_main's finally unregistered and closed the fault pipe BEFORE ++ worker.shutdown(); a PLE fail-closed inside that teardown (the drain waits ++ on the PLE transport) could not send its byte and the hang window reopened. ++ """ ++ import vllm.v1.executor.multiproc_executor as executor_module ++ from vllm.v1.ple_offload import fatal ++ ++ class _Inner: ++ def shutdown(self): ++ # ple_fail_closed's notify, from inside the teardown. ++ fatal.notify_parent_fatal() ++ ++ def fake_init(self, *args, **kwargs): ++ self.rank = 1 ++ self.worker = _Inner() ++ self.rpc_broadcast_mq = None ++ self.worker_response_mq = SimpleNamespace( ++ export_handle=lambda: None, ++ wait_until_ready=lambda: None, ++ shutdown=lambda: None, ++ ) ++ self.peer_response_handles = [] ++ ++ def busy_loop(self): ++ raise ValueError("boom") ++ ++ monkeypatch.setattr(executor_module.WorkerProc, "__init__", fake_init) ++ monkeypatch.setattr(executor_module.WorkerProc, "worker_busy_loop", busy_loop) ++ monkeypatch.setattr(executor_module, "destroy_model_parallel", lambda: None) ++ monkeypatch.setattr( ++ executor_module, "destroy_distributed_environment", lambda: None ++ ) ++ monkeypatch.setattr(executor_module, "maybe_init_worker_tracer", lambda **k: None) ++ monkeypatch.setattr(executor_module, "set_worker_net_device", lambda *a, **k: None) ++ monkeypatch.setattr( ++ executor_module, "_install_worker_parent_death_signal", lambda pid: None ++ ) ++ monkeypatch.setattr(executor_module.signal, "signal", lambda *a, **k: None) ++ ready_reader, ready_writer = multiprocessing.Pipe(duplex=False) ++ fault_reader, fault_writer = multiprocessing.Pipe(duplex=False) ++ with pytest.raises(SystemExit): ++ executor_module.WorkerProc.worker_main( ++ vllm_config=SimpleNamespace( ++ parallel_config=SimpleNamespace( ++ assigned_physical_gpu_ids=None, numa_bind=False ++ ) ++ ), ++ ready_pipe=ready_writer, ++ death_pipe=None, ++ fault_pipe=fault_writer, ++ shutdown_state=None, ++ inherited_fds=[], ++ expected_parent_pid=os.getpid(), ++ ) ++ ready_reader.close() ++ assert fault_reader.poll(1.0), "no byte and no EOF on the fault pipe" ++ data = os.read(fault_reader.fileno(), 1) ++ fault_reader.close() ++ assert data == b"F", ( ++ f"a fail-closed inside worker teardown sent {data!r}: the fault fd was " ++ "already unregistered/closed" ++ ) ++ assert getattr(fatal, "_FATAL_NOTIFY_FD", "unset") is None +diff --git a/tests/v1/executor/test_multiproc_monitor_any_rank.py b/tests/v1/executor/test_multiproc_monitor_any_rank.py +new file mode 100644 +index 0000000..d03366a +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_monitor_any_rank.py +@@ -0,0 +1,226 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The worker monitor must see ANY rank's exit, even with an fd inheritor. ++ ++The rank-1 fault test (an earlier build): the PLE fault killed PP0_TP1 (rank 1) in ++release_outputs. The worker sat as a zombie, GPU 1 emptied, but the executor ++never logged "died unexpectedly" and EngineCore hung in the RPC broadcast ++("No available shared memory broadcast block"). The monitor waited on ++proc.sentinel, a pipe whose write end the worker holds -- and the worker's ++torch_shm_manager child (started by the PLE host-buffer registration under the ++"file_system" sharing strategy) inherited that write end and outlived it, ++kept alive by rank 0's PLE child mapping rank 1's buffers. Rank-0 faults only ++looked detected because rank 0's death took the PLE child, and so the ++managers, down with it. ++ ++The inheritor here is a forked grandchild that sleeps; it holds every fd the ++worker had, the sentinel's write end included. ++""" ++ ++import contextlib ++import inspect ++import multiprocessing ++import os ++import signal ++import sys ++import threading ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++pytestmark = [ ++ pytest.mark.skip_global_cleanup, ++ pytest.mark.skipif(sys.platform != "linux", reason="fork + pidfd"), ++] ++ ++# A healthy worker detects a death within this bound on the earlier build's rank-0 path ++# (sentinel fires at exit); a pidfd fires at the exit itself. ++_DETECT_BOUND_S = 5.0 ++_INHERITOR_LIFETIME_S = 60.0 ++ ++ ++def _healthy_worker(): ++ time.sleep(600) ++ ++ ++def _dying_worker(inheritor_pid, spawn_inheritor): ++ if spawn_inheritor: ++ pid = os.fork() ++ if pid == 0: ++ # Stand-in for torch_shm_manager: holds the inherited fds. ++ time.sleep(_INHERITOR_LIFETIME_S) ++ os._exit(0) ++ inheritor_pid.value = pid ++ time.sleep(0.5) ++ # Stand-in for ple_fail_closed: no cleanup, straight out. ++ os._exit(1) ++ ++ ++def _executor(procs): ++ from vllm.v1.executor.multiproc_executor import MultiprocExecutor ++ ++ executor = object.__new__(MultiprocExecutor) ++ executor.workers = [ ++ SimpleNamespace( ++ proc=p, ++ rank=rank, ++ death_writer=None, ++ worker_response_mq=None, ++ shutdown_state=None, ++ ) ++ for rank, p in enumerate(procs) ++ ] ++ executor.is_failed = False ++ executor.rpc_broadcast_mq = None ++ executor.response_mqs = [] ++ failed = threading.Event() ++ executor.failure_callback = failed.set ++ return executor, failed ++ ++ ++@pytest.mark.parametrize("spawn_inheritor", [False, True]) ++@pytest.mark.parametrize("dying_rank", [0, 1]) ++def test_worker_exit_detected_on_any_rank(dying_rank, spawn_inheritor): ++ ctx = multiprocessing.get_context("fork") ++ inheritor_pid = ctx.Value("i", 0) ++ procs = [] ++ for rank in range(2): ++ if rank == dying_rank: ++ proc = ctx.Process( ++ target=_dying_worker, ++ args=(inheritor_pid, spawn_inheritor), ++ name=f"VllmWorker-{rank}", ++ daemon=True, ++ ) ++ else: ++ proc = ctx.Process( ++ target=_healthy_worker, name=f"VllmWorker-{rank}", daemon=True ++ ) ++ proc.start() ++ procs.append(proc) ++ ++ executor, failed = _executor(procs) ++ try: ++ start = time.monotonic() ++ executor.start_worker_monitor() ++ detected = failed.wait(_DETECT_BOUND_S + 5.0) ++ elapsed = time.monotonic() - start ++ assert detected, ( ++ f"rank {dying_rank} exited (inheritor={spawn_inheritor}) but the " ++ f"monitor never reported it within {_DETECT_BOUND_S + 5.0:.0f} s; " ++ f"dying proc alive per waitpid: {procs[dying_rank].is_alive()}" ++ ) ++ assert elapsed <= _DETECT_BOUND_S, elapsed ++ assert executor.is_failed ++ assert procs[dying_rank].exitcode == 1 ++ # The failure ladder ended the healthy survivor too. ++ assert not any(p.is_alive() for p in procs) ++ finally: ++ for p in procs: ++ if p.is_alive(): ++ p.kill() ++ p.join(timeout=5) ++ if inheritor_pid.value: ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(inheritor_pid.value, signal.SIGKILL) ++ ++ ++def test_exit_waitables_close_opened_pidfds_when_a_later_proc_fails(monkeypatch): ++ """_worker_exit_waitables is exception-safe.""" ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ opened: list[int] = [] ++ real_open = os.pidfd_open ++ ++ def recording_open(pid): ++ fd = real_open(pid) ++ opened.append(fd) ++ return fd ++ ++ class _Broken: ++ pid = None # no pidfd: falls back to the sentinel, which raises ++ ++ @property ++ def sentinel(self): ++ raise ValueError("process not started") ++ ++ class _Self: ++ pid = os.getpid() ++ sentinel = -1 ++ ++ monkeypatch.setattr(executor_module.os, "pidfd_open", recording_open) ++ with pytest.raises(ValueError): ++ executor_module._worker_exit_waitables([_Self(), _Self(), _Broken()]) ++ assert len(opened) == 2, opened ++ leaked = [] ++ for fd in opened: ++ try: ++ os.fstat(fd) ++ leaked.append(fd) ++ except OSError: ++ pass ++ for fd in leaked: ++ os.close(fd) ++ assert not leaked, f"pidfds left open after the failure: {leaked}" ++ ++ ++def test_exit_waitables_close_the_pidfd_of_the_failing_iteration(monkeypatch): ++ """A failure between opening a pidfd and recording it. ++ ++ An exception is injected (via a trace function) on the line that records ++ the (waitable, proc, pidfd) tuple, after the pidfd for that proc was ++ opened; that pidfd must be closed too, not only the earlier ones. ++ """ ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ opened: list[int] = [] ++ real_open = os.pidfd_open ++ ++ def recording_open(pid): ++ fd = real_open(pid) ++ opened.append(fd) ++ return fd ++ ++ class _Self: ++ pid = os.getpid() ++ sentinel = -1 ++ ++ monkeypatch.setattr(executor_module.os, "pidfd_open", recording_open) ++ code = executor_module._worker_exit_waitables.__code__ ++ append_lines = { ++ code.co_firstlineno + i ++ for i, line in enumerate( ++ inspect.getsource(executor_module._worker_exit_waitables).splitlines() ++ ) ++ if "waitables.append(" in line ++ } ++ assert len(append_lines) == 1, append_lines ++ ++ class _Injected(Exception): ++ pass ++ ++ def tracer(frame, event, arg): ++ if frame.f_code is not code: ++ return None ++ ++ def local(frame, event, arg): ++ if event == "line" and frame.f_lineno in append_lines and opened: ++ raise _Injected ++ return local ++ ++ return local ++ ++ sys.settrace(tracer) ++ try: ++ with pytest.raises(_Injected): ++ executor_module._worker_exit_waitables([_Self()]) ++ finally: ++ sys.settrace(None) ++ assert len(opened) == 1, opened ++ try: ++ os.fstat(opened[0]) ++ except OSError: ++ return ++ os.close(opened[0]) ++ pytest.fail("the pidfd opened in the failing iteration was left open") +diff --git a/tests/v1/executor/test_multiproc_monitor_combined.py b/tests/v1/executor/test_multiproc_monitor_combined.py +new file mode 100644 +index 0000000..d06af4e +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_monitor_combined.py +@@ -0,0 +1,147 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The worker monitor's two wake sources together: fault byte + process exit. ++ ++The rank-1 fault hang had two root causes. First fix: the sentinel's ++write end was inherited by the worker's torch_shm_manager child, so the exit ++never showed on it (fix: wait on a pidfd). Second fix: the dying ++rank's kernel teardown waits on a live peer, so no exit-based signal fires ++(fix: the fail-closed path writes one byte on a fault pipe before os._exit). ++The merged monitor must take whichever arrives first and must not count a ++fault pipe's EOF as a fault. Each case below fails on one of the two parents. ++""" ++ ++import contextlib ++import multiprocessing ++import os ++import signal ++import sys ++import threading ++import time ++from types import SimpleNamespace ++ ++import pytest ++ ++pytestmark = [ ++ pytest.mark.skip_global_cleanup, ++ pytest.mark.skipif(sys.platform != "linux", reason="fork + pidfd"), ++] ++ ++_BOUND_S = 2.0 ++ ++ ++def _healthy(fault_writer): ++ time.sleep(600) ++ ++ ++def _fault_byte_then_stuck(fault_writer): ++ # ple_fail_closed's notify, then a teardown that does not finish (the ++ # process stays alive, so neither its pidfd nor its sentinel fires). ++ os.write(fault_writer.fileno(), b"F") ++ time.sleep(600) ++ ++ ++def _silent_death_with_inheritor(fault_writer, inheritor_pid): ++ # A non-fail-closed death (no byte) whose fd-holding grandchild keeps the ++ # sentinel AND the fault pipe's write end open (no EOF either). ++ pid = os.fork() ++ if pid == 0: ++ time.sleep(60) ++ os._exit(0) ++ inheritor_pid.value = pid ++ time.sleep(0.3) ++ os._exit(1) ++ ++ ++def _orderly_eof_then_exit(fault_writer): ++ # worker_main's finally closes the fault pipe (EOF, no byte) and the ++ # process exits: a death, not a reported fault. ++ fault_writer.close() ++ time.sleep(0.3) ++ os._exit(0) ++ ++ ++def _run(dying_target, extra_args=()): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ from vllm.v1.executor.multiproc_executor import MultiprocExecutor ++ ++ ctx = multiprocessing.get_context("fork") ++ procs, handles = [], [] ++ for rank in range(2): ++ fault_reader, fault_writer = ctx.Pipe(duplex=False) ++ target = dying_target if rank == 1 else _healthy ++ args = (fault_writer, *extra_args) if rank == 1 else (fault_writer,) ++ proc = ctx.Process( ++ target=target, args=args, name=f"VllmWorker-{rank}", daemon=True ++ ) ++ proc.start() ++ fault_writer.close() ++ procs.append(proc) ++ handles.append( ++ SimpleNamespace( ++ proc=proc, ++ rank=rank, ++ death_writer=None, ++ worker_response_mq=None, ++ shutdown_state=None, ++ fault_reader=fault_reader, ++ ) ++ ) ++ executor = object.__new__(MultiprocExecutor) ++ executor.workers = handles ++ executor.is_failed = False ++ executor.rpc_broadcast_mq = None ++ executor.response_mqs = [] ++ failed = threading.Event() ++ executor.failure_callback = failed.set ++ ++ messages: list[str] = [] ++ real_error = executor_module.logger.error ++ ++ def capture(msg, *args, **kwargs): ++ messages.append(msg % args if args else msg) ++ return real_error(msg, *args, **kwargs) ++ ++ executor_module.logger.error = capture ++ try: ++ start = time.monotonic() ++ executor.start_worker_monitor() ++ detected = failed.wait(_BOUND_S + 8.0) ++ elapsed = time.monotonic() - start ++ finally: ++ executor_module.logger.error = real_error ++ for p in procs: ++ if p.is_alive(): ++ p.kill() ++ p.join(5) ++ return detected, elapsed, messages ++ ++ ++def test_fault_byte_wins_while_the_process_is_still_alive(): ++ detected, elapsed, messages = _run(_fault_byte_then_stuck) ++ assert detected, "fault byte written, process alive: monitor never woke" ++ assert elapsed <= _BOUND_S, elapsed ++ assert any("VllmWorker-1 reported a fatal fault" in m for m in messages), messages ++ ++ ++def test_silent_death_with_an_fd_inheritor_is_seen_by_its_exit(): ++ inheritor_pid = multiprocessing.get_context("fork").Value("i", 0) ++ try: ++ detected, elapsed, messages = _run( ++ _silent_death_with_inheritor, (inheritor_pid,) ++ ) ++ finally: ++ if inheritor_pid.value: ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(inheritor_pid.value, signal.SIGKILL) ++ assert detected, "worker exited (no fault byte, inherited fds): not seen" ++ assert elapsed <= _BOUND_S + 1.0, elapsed ++ assert any("VllmWorker-1 died unexpectedly" in m for m in messages), messages ++ ++ ++def test_fault_pipe_eof_is_not_a_fault(): ++ detected, elapsed, messages = _run(_orderly_eof_then_exit) ++ assert detected, "worker exited after closing its fault pipe: not seen" ++ assert elapsed <= _BOUND_S + 1.0, elapsed ++ assert any("VllmWorker-1 died unexpectedly" in m for m in messages), messages ++ assert not any("reported a fatal fault" in m for m in messages), messages +diff --git a/tests/v1/executor/test_multiproc_ready_and_inherit.py b/tests/v1/executor/test_multiproc_ready_and_inherit.py +new file mode 100644 +index 0000000..2ef20cd +--- /dev/null ++++ b/tests/v1/executor/test_multiproc_ready_and_inherit.py +@@ -0,0 +1,355 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Pre-READY worker death, and the inherited-pipe leak behind the rank-1 fault hang. ++ ++The rank-1 fault test: a PLE-owning worker's torch_shm_manager child (fork and ++program load, from the PLE host-buffer registration) inherited the worker's ++end of the sentinel pipe and outlived it, so the executor never saw rank 1 ++die. The same inheritance covers the READY pipe: a worker dying between the ++registration and READY left WorkerProc.wait_for_ready waiting for an EOF that ++never came. ++""" ++ ++import contextlib ++import multiprocessing ++import multiprocessing.connection ++import os ++import signal ++import subprocess ++import sys ++import threading ++import time ++ ++import pytest ++ ++pytestmark = [ ++ pytest.mark.skip_global_cleanup, ++ pytest.mark.skipif(sys.platform != "linux", reason="fork + pidfd"), ++] ++ ++_BOUND_S = 5.0 ++ ++ ++def _die_before_ready(inheritor_pid): ++ # Stand-in for torch_shm_manager: a forked child holding every fd, ++ # the READY pipe's write end included. ++ pid = os.fork() ++ if pid == 0: ++ time.sleep(60) ++ os._exit(0) ++ inheritor_pid.value = pid ++ time.sleep(0.3) ++ os._exit(1) # died during init, READY never sent ++ ++ ++def test_wait_for_ready_sees_a_worker_die_before_ready(): ++ from vllm.v1.executor.multiproc_executor import ( ++ UnreadyWorkerProcHandle, ++ WorkerProc, ++ ) ++ ++ ctx = multiprocessing.get_context("fork") ++ inheritor_pid = ctx.Value("i", 0) ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_die_before_ready, args=(inheritor_pid,), name="VllmWorker-1" ++ ) ++ proc.start() ++ ready_writer.close() ++ handle = UnreadyWorkerProcHandle(proc, 1, ready_reader) ++ ++ outcome: dict = {} ++ ++ def run(): ++ try: ++ WorkerProc.wait_for_ready([handle]) ++ outcome["result"] = "returned" ++ except BaseException as exc: # noqa: BLE001 ++ outcome["result"] = exc ++ ++ waiter = threading.Thread(target=run, daemon=True) ++ try: ++ start = time.monotonic() ++ waiter.start() ++ waiter.join(_BOUND_S + 5.0) ++ elapsed = time.monotonic() - start ++ assert not waiter.is_alive(), ( ++ "wait_for_ready still waiting for READY from a worker that exited " ++ f"{elapsed:.1f} s ago (an inheritor holds the pipe's write end)" ++ ) ++ assert isinstance(outcome.get("result"), Exception), outcome ++ assert "initialization failed" in str(outcome["result"]) ++ assert elapsed <= _BOUND_S, elapsed ++ finally: ++ if inheritor_pid.value: ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(inheritor_pid.value, signal.SIGKILL) ++ waiter.join(10) ++ proc.join(5) ++ ++ ++def _spawned_worker_starts_a_helper(helper_pid): ++ from vllm.v1.executor import multiproc_executor ++ ++ # The worker_main step under test; absent on unfixed code. ++ guard = getattr( ++ multiproc_executor, "_make_inherited_pipes_close_on_exec", lambda: None ++ ) ++ guard() ++ # Stand-in for torch_shm_manager: a helper program started with no fd ++ # closing, as libshm does. It outlives every bound below (killed in the ++ # test's cleanup), so a fired sentinel cannot be the helper's own exit. ++ helper = subprocess.Popen(["sleep", "3600"], close_fds=False) ++ helper_pid.value = helper.pid ++ os._exit(1) ++ ++ ++def _proc_state(pid): ++ """/proc//stat state letter, or None if the pid is gone.""" ++ try: ++ with open(f"/proc/{pid}/stat") as f: ++ stat = f.read() ++ except OSError: ++ return None ++ return stat[stat.rindex(")") + 2 :].split()[0] ++ ++ ++def _running(pid): ++ state = _proc_state(pid) ++ return state is not None and state not in ("Z", "X") ++ ++ ++def test_running_check_rejects_a_zombie(): ++ """os.kill(pid, 0) cannot tell a zombie from a live process; _running can.""" ++ pid = os.fork() ++ if pid == 0: ++ os._exit(0) ++ try: ++ deadline = time.monotonic() + 5 ++ while _proc_state(pid) != "Z" and time.monotonic() < deadline: ++ time.sleep(0.01) ++ assert _proc_state(pid) == "Z" ++ os.kill(pid, 0) # succeeds on the zombie: the old check passed here ++ assert not _running(pid) ++ finally: ++ os.waitpid(pid, 0) ++ assert _running(os.getpid()) ++ ++ ++# The worker exits right after publishing the helper's pid; its sentinel must ++# fire within this bound while the helper is provably still alive. ++_SENTINEL_BOUND_S = 2.0 ++ ++ ++def test_started_helper_does_not_inherit_the_sentinel(): ++ ctx = multiprocessing.get_context("spawn") ++ helper_pid = ctx.Value("i", 0) ++ proc = ctx.Process( ++ target=_spawned_worker_starts_a_helper, args=(helper_pid,), daemon=True ++ ) ++ proc.start() ++ try: ++ # Spawn + import can take a while; the bound starts at the helper's ++ # start (the worker exits right after it). ++ deadline = time.monotonic() + 120 ++ while not helper_pid.value and time.monotonic() < deadline: ++ time.sleep(0.05) ++ assert helper_pid.value, "the helper never started" ++ fired = multiprocessing.connection.wait( ++ [proc.sentinel], timeout=_SENTINEL_BOUND_S ++ ) ++ # The helper must still be RUNNING (not gone, and not a zombie -- ++ # os.kill(pid, 0) succeeds on a zombie too). ++ assert _running(helper_pid.value), ( ++ f"helper {helper_pid.value} is not running: " ++ f"state={_proc_state(helper_pid.value)}" ++ ) ++ assert fired, ( ++ f"the worker exited but its sentinel did not fire within " ++ f"{_SENTINEL_BOUND_S} s while the helper it started is alive: the " ++ "helper inherited the sentinel pipe's write end" ++ ) ++ proc.join(5) ++ assert proc.exitcode == 1 ++ finally: ++ if proc.is_alive(): ++ proc.kill() ++ if helper_pid.value: ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(helper_pid.value, signal.SIGKILL) ++ ++ ++def test_guard_marks_inherited_pipes_only(): ++ from vllm.v1.executor import multiproc_executor ++ ++ guard = multiproc_executor._make_inherited_pipes_close_on_exec ++ stdio = {} ++ for fd in (0, 1, 2): ++ with contextlib.suppress(OSError): ++ stdio[fd] = os.get_inheritable(fd) ++ reader, writer = os.pipe() ++ try: ++ os.set_inheritable(reader, True) ++ os.set_inheritable(writer, True) ++ guard() ++ assert not os.get_inheritable(reader) ++ assert not os.get_inheritable(writer) ++ for fd, before in stdio.items(): ++ assert os.get_inheritable(fd) == before, fd ++ finally: ++ os.close(reader) ++ os.close(writer) ++ ++ ++# ---- Pre-READY waiting keeps the sentinel fallback and the fault pipe ++ ++_PRE_READY_BOUND_S = 3.0 ++ ++ ++def _exits_with_ready_writer_held_by_a_helper(ready_writer, helper_pid): ++ # A helper program inherits ONLY the READY pipe's write end (fork's ++ # sentinel end is close-on-exec), so the READY pipe never EOFs while the ++ # sentinel does fire at the exit. ++ helper = subprocess.Popen(["sleep", "60"], pass_fds=(ready_writer.fileno(),)) ++ helper_pid.value = helper.pid ++ os._exit(1) ++ ++ ++def _fault_byte_before_ready(fault_writer): ++ # ple_fail_closed's notify during init; the process then stays alive ++ # (its teardown is pinned), so no exit signal fires either. ++ os.write(fault_writer.fileno(), b"F") ++ time.sleep(60) ++ ++ ++def _wait_for_ready_in_thread(handles): ++ from vllm.v1.executor.multiproc_executor import WorkerProc ++ ++ outcome: dict = {} ++ ++ def run(): ++ try: ++ WorkerProc.wait_for_ready(handles) ++ outcome["result"] = "returned" ++ except BaseException as exc: # noqa: BLE001 ++ outcome["result"] = exc ++ ++ waiter = threading.Thread(target=run, daemon=True) ++ start = time.monotonic() ++ waiter.start() ++ waiter.join(_PRE_READY_BOUND_S) ++ return waiter, outcome, time.monotonic() - start ++ ++ ++def test_pre_ready_death_is_seen_by_the_sentinel_when_no_pidfd(monkeypatch): ++ import vllm.v1.executor.multiproc_executor as executor_module ++ from vllm.v1.executor.multiproc_executor import UnreadyWorkerProcHandle ++ ++ def no_pidfd(pid): ++ raise OSError(38, "pidfd_open unavailable") ++ ++ monkeypatch.setattr(executor_module.os, "pidfd_open", no_pidfd, raising=False) ++ ctx = multiprocessing.get_context("fork") ++ helper_pid = ctx.Value("i", 0) ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_exits_with_ready_writer_held_by_a_helper, ++ args=(ready_writer, helper_pid), ++ name="VllmWorker-1", ++ ) ++ proc.start() ++ ready_writer.close() ++ handle = UnreadyWorkerProcHandle(proc, 1, ready_reader) ++ try: ++ waiter, outcome, elapsed = _wait_for_ready_in_thread([handle]) ++ assert not waiter.is_alive(), ( ++ f"no pidfd: wait_for_ready dropped the sentinel fallback and is " ++ f"still waiting {elapsed:.1f} s after the worker exited" ++ ) ++ assert isinstance(outcome.get("result"), Exception), outcome ++ finally: ++ if helper_pid.value: ++ with contextlib.suppress(ProcessLookupError): ++ os.kill(helper_pid.value, signal.SIGKILL) ++ proc.join(5) ++ ++ ++def test_pre_ready_fault_byte_fails_initialization(): ++ from vllm.v1.executor.multiproc_executor import UnreadyWorkerProcHandle ++ ++ ctx = multiprocessing.get_context("fork") ++ ready_reader, ready_writer = ctx.Pipe(duplex=False) ++ fault_reader, fault_writer = ctx.Pipe(duplex=False) ++ proc = ctx.Process( ++ target=_fault_byte_before_ready, args=(fault_writer,), name="VllmWorker-1" ++ ) ++ proc.start() ++ handle = UnreadyWorkerProcHandle(proc, 1, ready_reader, None, None, fault_reader) ++ try: ++ waiter, outcome, elapsed = _wait_for_ready_in_thread([handle]) ++ assert not waiter.is_alive(), ( ++ "a fault byte arrived before READY but wait_for_ready is still " ++ f"waiting after {elapsed:.1f} s" ++ ) ++ assert isinstance(outcome.get("result"), Exception), outcome ++ finally: ++ proc.kill() ++ proc.join(5) ++ # unblock a still-waiting pre-fix thread ++ ready_writer.close() ++ fault_writer.close() ++ ++ ++# ---- A READY worker that faults or dies while a sibling initialises ++ ++ ++def _ready_then_fault(ready_writer, fault_writer): ++ ready_writer.send({"status": "READY"}) ++ time.sleep(0.2) ++ os.write(fault_writer.fileno(), b"F") # fail-closed after READY ++ time.sleep(60) # teardown pinned: no exit signal ++ ++ ++def _never_ready(ready_writer): ++ time.sleep(60) # still initialising ++ ++ ++def test_ready_worker_fault_while_sibling_initialises_fails_init(monkeypatch): ++ from vllm.v1.executor.multiproc_executor import ( ++ UnreadyWorkerProcHandle, ++ WorkerProc, ++ ) ++ ++ # READY handling beyond the status check needs real MQ handles; the ++ # response contents are not what this test is about. ++ monkeypatch.setattr( ++ WorkerProc, ++ "wait_for_response_handle_ready", ++ staticmethod(lambda response, handle: handle), ++ ) ++ ctx = multiprocessing.get_context("fork") ++ ra, wa = ctx.Pipe(duplex=False) ++ fa, fwa = ctx.Pipe(duplex=False) ++ rb, wb = ctx.Pipe(duplex=False) ++ fb, fwb = ctx.Pipe(duplex=False) ++ pa = ctx.Process(target=_ready_then_fault, args=(wa, fwa), name="VllmWorker-0") ++ pb = ctx.Process(target=_never_ready, args=(wb,), name="VllmWorker-1") ++ pa.start() ++ pb.start() ++ handles = [ ++ UnreadyWorkerProcHandle(pa, 0, ra, None, None, fa), ++ UnreadyWorkerProcHandle(pb, 1, rb, None, None, fb), ++ ] ++ try: ++ waiter, outcome, elapsed = _wait_for_ready_in_thread(handles) ++ assert not waiter.is_alive(), ( ++ "worker 0 was READY and then reported a fatal fault while worker 1 " ++ f"was still initialising; wait_for_ready still waiting at {elapsed:.1f} s" ++ ) ++ assert isinstance(outcome.get("result"), Exception), outcome ++ finally: ++ for p in (pa, pb): ++ p.kill() ++ p.join(5) ++ for w in (wa, fwa, wb, fwb): ++ w.close() +diff --git a/tests/v1/executor/test_multiproc_worker_exit_after_drain.py b/tests/v1/executor/test_multiproc_worker_exit_after_drain.py +index 7694a62..2881193 100644 +--- a/tests/v1/executor/test_multiproc_worker_exit_after_drain.py ++++ b/tests/v1/executor/test_multiproc_worker_exit_after_drain.py +@@ -225,13 +225,18 @@ def test_orderly_shutdown_exits_right_after_drain(monkeypatch): + + + def test_unexpected_busy_loop_failure_keeps_full_teardown(monkeypatch): +- # Not an orderly shutdown (no shutdown request): unchanged behaviour -- +- # full WorkerProc.shutdown incl. distributed destroy, then a normal return. ++ # Not an orderly shutdown: full WorkerProc.shutdown incl. distributed ++ # destroy, then exit 1. An earlier build returned normally (process exit ++ # 0) whenever shutdown_requested was not yet set at the classification -- ++ # which a death-pipe EOF landing just after the check makes a race. + events, exit_code, _ = _run_worker_main_in_process( + monkeypatch, busy_loop_error=ValueError("boom"), slow_destroy_s=0.0 + ) + names = [name for name, _ in events] +- assert exit_code is None, names ++ assert exit_code == 1, ( ++ f"an unexpected busy-loop failure exited {exit_code!r}, not 1: {names}" ++ ) ++ assert "os._exit" not in names, f"took the orderly exit path: {names}" + assert names.count("drain") == 1, names + assert "destroy_model_parallel" in names, names + assert "destroy_distributed_environment" in names, names +@@ -454,3 +459,20 @@ def test_ladder_bounds_a_worker_stuck_in_its_busy_loop(monkeypatch): + assert any("sending SIGKILL" in w for w in warnings), warnings + assert proc.exitcode == -9, proc.exitcode + assert elapsed <= grace_s + 4 + 2, f"ladder took {elapsed:.1f}s" ++ ++ ++def test_worker_main_runs_the_inherited_pipe_guard(monkeypatch): ++ # Rank-1 fault: worker_main marks its inherited pipes close-on-exec before ++ # anything can start torch_shm_manager (see ++ # test_multiproc_ready_and_inherit.py for the helper itself). ++ import vllm.v1.executor.multiproc_executor as executor_module ++ ++ called = [] ++ monkeypatch.setattr( ++ executor_module, ++ "_make_inherited_pipes_close_on_exec", ++ lambda: called.append(True), ++ raising=False, ++ ) ++ _run_worker_main_in_process(monkeypatch) ++ assert called, "worker_main never ran the inherited-pipe guard" +diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py +index 72da32e..dec44e1 100644 +--- a/vllm/v1/engine/core.py ++++ b/vllm/v1/engine/core.py +@@ -110,6 +110,24 @@ _R = TypeVar("_R") # Return type for collective_rpc + _TEARDOWN_SIGNAL_BOUND_S = 30.0 + + ++# How long the guard's forced-exit paths wait for their diagnostic line. The ++# line is written from a helper thread; a blocked handler, a held logging lock ++# or a full stderr pipe costs at most this, never the exit. ++_TEARDOWN_EXIT_LOG_BOUND_S = 0.5 ++ ++ ++def _log_bounded(emit: Callable[[], None], bound_s: float) -> None: ++ """Run ``emit`` (logging, flushes) on a daemon thread; wait <= bound_s.""" ++ try: ++ thread = threading.Thread( ++ target=emit, daemon=True, name="EngineCoreTeardownLog" ++ ) ++ thread.start() ++ thread.join(bound_s) ++ except BaseException: # noqa: BLE001 -- the caller's exit must still run ++ pass ++ ++ + class _TeardownSignalGuard: + """EngineCore teardown: defer ONE late SIGTERM/SIGINT, bounded. + +@@ -148,25 +166,34 @@ class _TeardownSignalGuard: + ) + self._watchdog.start() + return +- logger.warning( +- "[shutdown] EngineCore: second late signal (%s) during resource " +- "teardown; not deferring again", +- name, +- ) ++ # Default disposition first, so a third signal also ends the process ++ # while the (bounded) diagnostic line is written. + signal.signal(signum, signal.SIG_DFL) ++ _log_bounded( ++ lambda: logger.warning( ++ "[shutdown] EngineCore: second late signal (%s) during " ++ "resource teardown; not deferring again", ++ name, ++ ), ++ _TEARDOWN_EXIT_LOG_BOUND_S, ++ ) + os.kill(os.getpid(), signum) + + def _expire(self) -> None: + if self._done.wait(self.bound_s): + return +- logger.error( +- "[shutdown] EngineCore: resource teardown still running %.0f s " +- "after a deferred signal; exiting (code 1)", +- self.bound_s, +- ) +- for stream in (sys.stdout, sys.stderr): +- with suppress(Exception): +- stream.flush() ++ ++ def emit() -> None: ++ logger.error( ++ "[shutdown] EngineCore: resource teardown still running %.0f s " ++ "after a deferred signal; exiting (code 1)", ++ self.bound_s, ++ ) ++ for stream in (sys.stdout, sys.stderr): ++ with suppress(Exception): ++ stream.flush() ++ ++ _log_bounded(emit, _TEARDOWN_EXIT_LOG_BOUND_S) + os._exit(1) + + def disarm(self) -> None: +diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py +index 0c5af34..d7beb9a 100644 +--- a/vllm/v1/executor/multiproc_executor.py ++++ b/vllm/v1/executor/multiproc_executor.py +@@ -5,6 +5,7 @@ import os + import pickle + import queue + import signal ++import stat + import sys + import threading + import time +@@ -71,6 +72,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput + from vllm.v1.executor.abstract import Executor, FailureCallback + from vllm.v1.executor.vllm_net_devices import set_worker_net_device + from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput ++from vllm.v1.ple_offload.fatal import set_fatal_notify_fd + from vllm.v1.worker.worker_base import WorkerWrapperBase + + logger = init_logger(__name__) +@@ -97,13 +99,16 @@ def _install_worker_parent_death_signal(expected_parent_pid: int) -> None: + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int +- if prctl( +- _PR_SET_PDEATHSIG, +- int(_WORKER_PARENT_DEATH_SIGNAL), +- 0, +- 0, +- 0, +- ) != 0: ++ if ( ++ prctl( ++ _PR_SET_PDEATHSIG, ++ int(_WORKER_PARENT_DEATH_SIGNAL), ++ 0, ++ 0, ++ 0, ++ ) ++ != 0 ++ ): + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + +@@ -173,6 +178,86 @@ def _exit_worker_process(exit_code: int) -> None: + os._exit(exit_code) + + ++def _worker_exit_waitables( ++ procs: Sequence[BaseProcess], ++) -> list[tuple[Any, BaseProcess, int | None]]: ++ """What the worker monitor waits on to learn that a worker process exited. ++ ++ A pidfd (Linux >= 5.3) where possible: it becomes readable when the process ++ itself exits. proc.sentinel is only the read end of a pipe whose write end ++ the worker holds, so it reports the exit of the LAST process holding that ++ write end -- and anything the worker fork+execs inherits it. Every PLE- ++ owning worker does: registering its host buffers under the "file_system" ++ sharing strategy starts a torch_shm_manager child, which keeps the write ++ end, and which lives while any client -- rank 0's PLE child, for a TP1 ++ worker's buffers -- still maps them. In a rank-1 fault test PP0_TP1 died ++ in release_outputs, sat as a zombie, its sentinel never fired and the ++ engine hung; rank 0's death only looked detected because its PLE child ++ (and so both managers) died with it, ~8 s later. ++ ++ Returns (waitable, proc, pidfd-to-close-or-None) per process, falling back ++ to the sentinel where no pidfd can be opened. ++ """ ++ pidfd_open = getattr(os, "pidfd_open", None) ++ waitables: list[tuple[Any, BaseProcess, int | None]] = [] ++ try: ++ for proc in procs: ++ pidfd: int | None = None ++ pid = getattr(proc, "pid", None) ++ if pidfd_open is not None and pid is not None: ++ try: ++ pidfd = pidfd_open(pid) ++ except OSError: ++ # ESRCH (already reaped) or no pidfd support: the sentinel. ++ pidfd = None ++ try: ++ waitable = pidfd if pidfd is not None else proc.sentinel ++ waitables.append((waitable, proc, pidfd)) ++ except BaseException: ++ # Not yet in `waitables`: the outer handler would miss it. ++ if pidfd is not None: ++ os.close(pidfd) ++ raise ++ except BaseException: ++ # Nothing is handed out on failure: close what was opened. ++ for _, _, opened in waitables: ++ if opened is not None: ++ with suppress(OSError): ++ os.close(opened) ++ raise ++ return waitables ++ ++ ++def _make_inherited_pipes_close_on_exec() -> None: ++ """Keep this worker's inherited pipe fds out of the programs it execs. ++ ++ Under the spawn start method multiprocessing hands the worker its end of ++ the sentinel pipe (and the ready / death pipes) as INHERITABLE fds, so any ++ program the worker fork+execs keeps them open: the torch_shm_manager that ++ the PLE host-buffer registration starts did, and outlived the worker, so ++ neither the sentinel nor the ready pipe reported its death (rank-1 fault injection). ++ Nothing this process execs needs them -- multiprocessing children get ++ their fds passed explicitly, subprocess closes them by default -- so mark ++ every inherited pipe above stderr close-on-exec. The executor's pidfd ++ wait is the primary detection; this only stops the leak at its source. ++ """ ++ if not sys.platform.startswith("linux"): ++ return ++ try: ++ fds = [int(name) for name in os.listdir("/proc/self/fd")] ++ except (OSError, ValueError): ++ return ++ for fd in fds: ++ if fd <= 2: ++ continue ++ try: ++ if stat.S_ISFIFO(os.fstat(fd).st_mode) and os.get_inheritable(fd): ++ os.set_inheritable(fd, False) ++ except OSError: ++ # The listing's own directory fd, or one closed since. ++ continue ++ ++ + class FutureWrapper(Future): + def __init__( + self, +@@ -304,6 +389,7 @@ class MultiprocExecutor(Executor): + if inherited_fds is not None: + inherited_fds.append(unready_worker_handle.death_writer.fileno()) + inherited_fds.append(unready_worker_handle.ready_pipe.fileno()) ++ inherited_fds.append(unready_worker_handle.fault_reader.fileno()) + + # Workers must be created before wait_for_ready to avoid + # deadlock, since worker.init_device() does a device sync. +@@ -400,22 +486,78 @@ class MultiprocExecutor(Executor): + # Monitors worker process liveness. If any die unexpectedly, + # logs an error, shuts down the executor and invokes the failure + # callback to inform the engine. ++ # ++ # Two wake sources per worker: ++ # - its exit: a pidfd, readable when the process itself has exited ++ # (fallback: proc.sentinel, which a fork+exec'd child holding the ++ # write end -- torch_shm_manager from the PLE registration -- keeps ++ # silent; see _worker_exit_waitables); ++ # - its fault pipe: one byte from the fail-closed path right before ++ # os._exit, so the failure ladder starts -- and ends the peers -- ++ # without waiting for the process teardown, which can itself wait ++ # on a live peer. A worker closing its end on the way out (EOF, no ++ # byte) is not a fault; its exit follows. ++ # In a rank-1 fault test a fault on PP0_TP1 was never detected and the ++ # engine hung on the response queue for 300 s. ++ # Opened here, before the thread starts, so a worker that dies early ++ # is still a zombie (pidfds work on zombies) rather than reaped. ++ exit_waitables = _worker_exit_waitables([h.proc for h in workers]) ++ + def monitor_workers(): +- sentinels = [h.proc.sentinel for h in workers] +- died = multiprocessing.connection.wait(sentinels) ++ exit_to_handle = { ++ w: h for (w, _, _), h in zip(exit_waitables, workers) ++ } ++ fault_to_handle = { ++ h.fault_reader: h ++ for h in workers ++ if getattr(h, "fault_reader", None) is not None ++ } ++ waiting: list[Any] = list(exit_to_handle) + list(fault_to_handle) ++ handle = None ++ reported_fatal = False ++ try: ++ while handle is None: ++ ready = multiprocessing.connection.wait(waiting) ++ for obj in ready: ++ if obj in exit_to_handle: ++ handle = exit_to_handle[obj] ++ break ++ try: ++ data = os.read(obj.fileno(), 1) ++ except OSError: ++ data = b"" ++ if data: ++ handle = fault_to_handle[obj] ++ reported_fatal = True ++ break ++ waiting.remove(obj) ++ finally: ++ for _, _, pidfd in exit_waitables: ++ if pidfd is not None: ++ with suppress(OSError): ++ os.close(pidfd) + _self = self_ref() + if not _self or getattr(_self, "shutting_down", False): + logger.debug("MultiprocWorkerMonitor: shutdown already initiated") + return + _self.is_failed = True +- proc = next(h.proc for h in workers if h.proc.sentinel == died[0]) +- logger.error( +- "Worker proc %s died unexpectedly (exit code: %s), " +- "shutting down executor. executor_shutdown_monotonic=%.9f", +- proc.name, +- proc.exitcode, +- time.monotonic(), +- ) ++ proc = handle.proc ++ if reported_fatal: ++ logger.error( ++ "Worker proc %s reported a fatal fault (fail-closed) and is " ++ "exiting; shutting down executor without waiting for its " ++ "process teardown. executor_shutdown_monotonic=%.9f", ++ proc.name, ++ time.monotonic(), ++ ) ++ else: ++ logger.error( ++ "Worker proc %s died unexpectedly (exit code: %s), " ++ "shutting down executor. executor_shutdown_monotonic=%.9f", ++ proc.name, ++ proc.exitcode, ++ time.monotonic(), ++ ) + _self.shutdown() + callback = _self.failure_callback + if callback is not None: +@@ -423,9 +565,16 @@ class MultiprocExecutor(Executor): + callback() + + if not inline: +- Thread( +- target=monitor_workers, daemon=True, name="MultiprocWorkerMonitor" +- ).start() ++ try: ++ Thread( ++ target=monitor_workers, daemon=True, name="MultiprocWorkerMonitor" ++ ).start() ++ except BaseException: ++ for _, _, pidfd in exit_waitables: ++ if pidfd is not None: ++ with suppress(OSError): ++ os.close(pidfd) ++ raise + return + + monitor_workers() +@@ -586,8 +735,9 @@ class MultiprocExecutor(Executor): + # If we are in late stage shutdown, the interpreter may replace + # `time` with `None`. + return all(not proc.is_alive() for proc in procs) +- start_time = time.time() +- while time.time() - start_time < timeout: ++ # Monotonic: a wall-clock step must not cut or stretch the grace. ++ start_time = time.monotonic() ++ while time.monotonic() - start_time < timeout: + if all(not proc.is_alive() for proc in procs): + return True + time.sleep(0.1) +@@ -602,7 +752,7 @@ class MultiprocExecutor(Executor): + initial_count, + ) + if after_failure and time: +- start = time.time() ++ start = time.monotonic() + states = list(shutdown_states or [None] * len(worker_procs)) + + def parked() -> list[BaseProcess]: +@@ -634,7 +784,7 @@ class MultiprocExecutor(Executor): + p.kill() + # Workers that started their shutdown get a BOUNDED grace, counted + # from the start of this ladder, then SIGTERM, then SIGKILL. +- grace_s = max(0.0, _FAILED_DRAIN_GRACE_S - (time.time() - start)) ++ grace_s = max(0.0, _FAILED_DRAIN_GRACE_S - (time.monotonic() - start)) + term_wait_s = _FAILED_WORKER_WAIT_S + if wait_for_termination(active_procs(), timeout=grace_s): + logger.info_once("[shutdown] Executor: all workers exited gracefully") +@@ -735,6 +885,7 @@ class UnreadyWorkerProcHandle: + ready_pipe: Connection + death_writer: Connection | None = None + shutdown_state: Any = None ++ fault_reader: Connection | None = None + + + @dataclass +@@ -750,6 +901,8 @@ class WorkerProcHandle: + death_writer: Connection | None = None + # Shared byte: _SHUTDOWN_NOT_STARTED / _STARTED / _DONE (worker writes). + shutdown_state: Any = None ++ # Read end of the worker's fault pipe (see start_worker_monitor). ++ fault_reader: Connection | None = None + + @classmethod + def from_unready_handle( +@@ -765,6 +918,7 @@ class WorkerProcHandle: + peer_worker_response_mqs=peer_worker_response_mqs, + death_writer=unready_handle.death_writer, + shutdown_state=unready_handle.shutdown_state, ++ fault_reader=unready_handle.fault_reader, + ) + + +@@ -903,11 +1057,16 @@ class WorkerProc: + ready_reader, ready_writer = context.Pipe(duplex=False) + # Death pipe to let child detect parent process exit + death_reader, death_writer = context.Pipe(duplex=False) ++ # Fault pipe: the worker's fail-closed path writes one byte before ++ # os._exit; the executor's monitor waits on it next to the sentinel. ++ fault_reader, fault_writer = context.Pipe(duplex=False) + # One shared byte: has this worker started / finished its shutdown? + shutdown_state = context.RawValue("b", _SHUTDOWN_NOT_STARTED) + if inherited_fds is not None: + inherited_fds = inherited_fds.copy() +- inherited_fds.extend((ready_reader.fileno(), death_writer.fileno())) ++ inherited_fds.extend( ++ (ready_reader.fileno(), death_writer.fileno(), fault_reader.fileno()) ++ ) + process_kwargs = { + "expected_parent_pid": os.getpid(), + "vllm_config": vllm_config, +@@ -917,6 +1076,7 @@ class WorkerProc: + "input_shm_handle": input_shm_handle, + "ready_pipe": ready_writer, + "death_pipe": death_reader, ++ "fault_pipe": fault_writer, + "shutdown_state": shutdown_state, + "shared_worker_lock": shared_worker_lock, + "is_driver_worker": is_driver_worker, +@@ -940,10 +1100,11 @@ class WorkerProc: + # Close child ends of pipes here in the parent + ready_writer.close() + death_reader.close() ++ fault_writer.close() + # Keep death_writer open in parent - when parent exits, + # death_reader in child will get EOFError + return UnreadyWorkerProcHandle( +- proc, rank, ready_reader, death_writer, shutdown_state ++ proc, rank, ready_reader, death_writer, shutdown_state, fault_reader + ) + + @staticmethod +@@ -980,10 +1141,59 @@ class WorkerProc: + ready_proc_handles: list[WorkerProcHandle | None] = [None] * len( + unready_proc_handles + ) +- while pipes: +- ready = multiprocessing.connection.wait(pipes.keys()) +- for pipe in ready: +- assert isinstance(pipe, Connection) ++ # A worker that dies before READY need not close its ready pipe: a ++ # child it fork+exec'd (torch_shm_manager, started by the PLE ++ # registration) can hold the write end, so no EOF ever comes. Wait on ++ # each pending worker's exit (pidfd, else its sentinel -- see ++ # _worker_exit_waitables) and on its fault pipe (one byte from the ++ # fail-closed path; EOF is not a fault), as the post-READY monitor does. ++ exit_waitables = _worker_exit_waitables( ++ [h.proc for h in unready_proc_handles] ++ ) ++ exits = { ++ w: handle ++ for (w, _, _), handle in zip(exit_waitables, unready_proc_handles) ++ } ++ faults = { ++ handle.fault_reader: handle ++ for handle in unready_proc_handles ++ if getattr(handle, "fault_reader", None) is not None ++ } ++ try: ++ while pipes: ++ waiting: list[Any] = [*pipes.keys()] ++ # EVERY worker's exit and fault pipe, READY or not: the live ++ # monitor starts only after all are READY, so a READY worker ++ # that faults or dies while a sibling initialises must fail ++ # initialization here. ++ waiting += list(exits) ++ waiting += list(faults) ++ ready = multiprocessing.connection.wait(waiting) ++ WorkerProc._receive_ready( ++ ready, pipes, exits, faults, ready_proc_handles, e ++ ) ++ finally: ++ for _, _, pidfd in exit_waitables: ++ if pidfd is not None: ++ with suppress(OSError): ++ os.close(pidfd) ++ ++ return cast(list[WorkerProcHandle], ready_proc_handles) ++ ++ @staticmethod ++ def _receive_ready( ++ ready: list[Any], ++ pipes: dict[Connection, UnreadyWorkerProcHandle], ++ exits: dict[Any, UnreadyWorkerProcHandle], ++ faults: dict[Any, UnreadyWorkerProcHandle], ++ ready_proc_handles: list[WorkerProcHandle | None], ++ e: Exception, ++ ) -> None: ++ # Ready pipes first: a READY sent before an exit or fault still counts. ++ ordered = sorted(ready, key=lambda obj: 0 if obj in pipes else 1) ++ for obj in ordered: ++ if obj in pipes: ++ pipe = obj + try: + # Wait until the WorkerProc is ready. + unready_proc_handle = pipes.pop(pipe) +@@ -992,18 +1202,50 @@ class WorkerProc: + raise e + + idx = unready_proc_handle.rank % len(ready_proc_handles) +- ready_proc_handles[idx] = WorkerProc.wait_for_response_handle_ready( +- response, unready_proc_handle ++ ready_proc_handles[idx] = ( ++ WorkerProc.wait_for_response_handle_ready( ++ response, unready_proc_handle ++ ) + ) + except EOFError: + e.__suppress_context__ = True + raise e from None +- + finally: + # Close connection. + pipe.close() +- +- return cast(list[WorkerProcHandle], ready_proc_handles) ++ continue ++ handle = exits.get(obj) if obj in exits else faults.get(obj) ++ if handle is None: ++ continue ++ pending = handle.ready_pipe in pipes ++ if pending and handle.ready_pipe.poll(0): ++ continue # its READY is queued: read it next round ++ phase = ( ++ "before READY" ++ if pending ++ else "after READY, while other workers were initializing" ++ ) ++ if obj in exits: ++ logger.error( ++ "Worker proc %s exited (code %s) %s", ++ handle.proc.name, ++ handle.proc.exitcode, ++ phase, ++ ) ++ raise e ++ try: ++ data = os.read(obj.fileno(), 1) ++ except OSError: ++ data = b"" ++ if data: ++ logger.error( ++ "Worker proc %s reported a fatal fault (fail-closed) %s", ++ handle.proc.name, ++ phase, ++ ) ++ raise e ++ # EOF on the fault pipe is not a fault; its exit follows. ++ del faults[obj] + + def shutdown(self, *, destroy_distributed: bool = True): + if not destroy_distributed: +@@ -1096,11 +1338,21 @@ class WorkerProc: + # requested (death pipe EOF or SIGTERM/SIGINT): the process then exits + # straight after WorkerProc.shutdown (see _exit_worker_process). + orderly_shutdown = False +- # Set when a READY worker's busy loop FAILED while a shutdown request +- # was pending: full teardown, then exit code 1 (never the exit-0 path). +- failed_during_shutdown = False ++ # Set when a READY worker's busy loop FAILED (anything but its own ++ # shutdown-request cancellation): full teardown, then exit code 1. ++ # Never keyed to shutdown_requested -- the death-pipe monitor can set ++ # it just after the check, which made a failure exit 0. ++ busy_loop_failed = False + ready_writer = kwargs.pop("ready_pipe") + death_pipe = kwargs.pop("death_pipe", None) ++ # Fault pipe: registered for the PLE fail-closed path ++ # (ple_offload.fatal.notify_parent_fatal); kept open for the life of ++ # the process so the byte written right before os._exit reaches the ++ # executor's monitor. Closed only on the orderly way out (EOF, which ++ # the monitor does not count as a fault). ++ fault_pipe = kwargs.pop("fault_pipe", None) ++ if fault_pipe is not None: ++ set_fatal_notify_fd(fault_pipe.fileno()) + shutdown_state = kwargs.pop("shutdown_state", None) + + def mark_shutdown(value: int) -> None: +@@ -1116,6 +1368,8 @@ class WorkerProc: + os.close(fd) + except Exception as e: + logger.warning("Error closing inherited connection: %s: %s", type(e), e) ++ # Before anything below can fork+exec a helper (torch_shm_manager). ++ _make_inherited_pipes_close_on_exec() + + try: + # Initialize tracer +@@ -1169,13 +1423,14 @@ class WorkerProc: + "[shutdown] WorkerProc: exiting after shutdown request" + ) + orderly_shutdown = True +- elif shutdown_requested.is_set(): +- # A real failure that raced a shutdown request: keep it a +- # failure -- full teardown, then a non-zero exit. +- logger.exception("WorkerProc failed during shutdown.") +- failed_during_shutdown = True + else: +- logger.exception("WorkerProc failed.") ++ # A real failure, whether or not a shutdown request is pending ++ # (yet): full teardown, then a non-zero exit. ++ if shutdown_requested.is_set(): ++ logger.exception("WorkerProc failed during shutdown.") ++ else: ++ logger.exception("WorkerProc failed.") ++ busy_loop_failed = True + + # The parent sends a SIGTERM to all worker processes if + # any worker dies. Set this value so we don't re-throw +@@ -1200,36 +1455,43 @@ class WorkerProc: + ready_writer.close() + if death_pipe is not None: + death_pipe.close() +- # Clean up once worker exits busy loop +- if worker is not None: +- # Tells the executor's ladder this rank is draining, not parked +- # in a collective (see _ensure_worker_termination). +- mark_shutdown(_SHUTDOWN_STARTED) +- if orderly_shutdown: +- # Everything that must happen before exit (FINAL DRAIN, +- # instrumentation flushes, PLE cleanup) is in +- # worker.shutdown(); then exit instead of paying the NCCL +- # destroy + finalization past the executor's grace. +- exit_code = 0 +- try: +- worker.shutdown(destroy_distributed=False) +- except BaseException: +- logger.exception("WorkerProc shutdown failed") +- exit_code = 1 ++ # The fault fd stays registered through the whole teardown: a PLE ++ # fail-closed inside worker.shutdown() must still reach the ++ # executor. Unregistered after it; on the orderly ++ # os._exit path the process exit closes it. ++ try: ++ # Clean up once worker exits busy loop ++ if worker is not None: ++ # Tells the executor's ladder this rank is draining, not parked ++ # in a collective (see _ensure_worker_termination). ++ mark_shutdown(_SHUTDOWN_STARTED) ++ if orderly_shutdown: ++ # Everything that must happen before exit (FINAL DRAIN, ++ # instrumentation flushes, PLE cleanup) is in ++ # worker.shutdown(); then exit instead of paying the NCCL ++ # destroy + finalization past the executor's grace. ++ exit_code = 0 ++ try: ++ worker.shutdown(destroy_distributed=False) ++ except BaseException: ++ logger.exception("WorkerProc shutdown failed") ++ exit_code = 1 ++ mark_shutdown(_SHUTDOWN_DONE) ++ _exit_worker_process(exit_code) ++ worker.shutdown() + mark_shutdown(_SHUTDOWN_DONE) +- _exit_worker_process(exit_code) +- worker.shutdown() +- mark_shutdown(_SHUTDOWN_DONE) +- if failed_during_shutdown: +- sys.exit(1) ++ if busy_loop_failed: ++ sys.exit(1) ++ finally: ++ if fault_pipe is not None: ++ set_fatal_notify_fd(None) ++ fault_pipe.close() + + class ResponseStatus(Enum): + SUCCESS = auto() + FAILURE = auto() + +- def _fail_ple_rpc_if_pending( +- self, where: str, error: BaseException +- ) -> None: ++ def _fail_ple_rpc_if_pending(self, where: str, error: BaseException) -> None: + """Fail closed if RPC recovery would abandon queued PLE GPU work.""" + owners = (self.worker, getattr(self.worker, "worker", None)) + for owner in owners: +@@ -1242,9 +1504,7 @@ class WorkerProc: + try: + pending = connector.has_pending_completions + except BaseException as state_error: +- connector._fail_closed( +- f"{where}:pending-state", state_error +- ) ++ connector._fail_closed(f"{where}:pending-state", state_error) + if pending: + connector._fail_closed(where, error) + return +@@ -1259,9 +1519,7 @@ class WorkerProc: + output = output.get_output() + except Exception as e: + logger.exception("Error getting async model runner output") +- self._fail_ple_rpc_if_pending( +- "async-output-copy", e +- ) ++ self._fail_ple_rpc_if_pending("async-output-copy", e) + output = e + + if isinstance(output, Exception): +@@ -1334,9 +1592,7 @@ class WorkerProc: + if output_rank is None or self.rank == output_rank: + self.handle_output(output) + except Exception as e: +- self._fail_ple_rpc_if_pending( +- f"worker-rpc:{method}", e +- ) ++ self._fail_ple_rpc_if_pending(f"worker-rpc:{method}", e) + # Notes have been introduced in python 3.11 + if hasattr(e, "add_note"): + e.add_note(traceback.format_exc()) +diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py +index 2cf05bb..6aa4e04 100644 +--- a/vllm/v1/ple_offload/connector.py ++++ b/vllm/v1/ple_offload/connector.py +@@ -25,7 +25,7 @@ from vllm.model_executor.layers.ple_offload_layer import ( + flag_check_is_authoritative, + stale_offload_flags, + ) +-from vllm.v1.ple_offload.fatal import ple_fail_closed ++from vllm.v1.ple_offload.fatal import notify_parent_fatal, ple_fail_closed + from vllm.v1.ple_offload.progress import PleProgressMonitor + from vllm.v1.ple_offload.protocol import ( + PleOffloadRegistration, +@@ -469,6 +469,7 @@ class PleOffloadConnector: + self._process_request(request, socket) + except Exception: + logger.exception("PLE request thread failed") ++ notify_parent_fatal() + os._exit(1) + finally: + self._request_thread_ready.set() +diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py +index a96515f..32831aa 100644 +--- a/vllm/v1/ple_offload/fatal.py ++++ b/vllm/v1/ple_offload/fatal.py +@@ -44,6 +44,40 @@ _TYPE_MODULE = type.__dict__["__module__"].__get__ + _MAX_CHAIN = 16 + _MAX_ARG_CHARS = 1024 + ++# Write end of the per-worker fault pipe (WorkerProc.make_worker_process); the ++# executor's worker monitor waits on the read end next to the process ++# sentinels. None outside a worker process (tests, the offload child). ++_FATAL_NOTIFY_FD: int | None = None ++ ++ ++def set_fatal_notify_fd(fd: int | None) -> None: ++ """Register the fault pipe worker_main received from the executor.""" ++ global _FATAL_NOTIFY_FD ++ _FATAL_NOTIFY_FD = fd ++ ++ ++def notify_parent_fatal() -> None: ++ """Tell the executor this worker is dying, before the process is gone. ++ ++ The executor's only other signal is the process sentinel, which the kernel ++ closes when the LAST task of the worker leaves -- after the address-space ++ teardown, not at os._exit. On a rank whose GPU work is paired with a live ++ peer that teardown can wait for the peer's exit (in a rank-1 fault test, a ++ fail-closed fault on PP0_TP1 left its leader Z and unreapable for 300 s ++ while rank 0 waited on it in turn; the engine hung on the response ++ queue). One byte here wakes the monitor at once so the failure ladder ends ++ the peers, which is what frees the teardown. Nonblocking, never raises: ++ the caller os._exits right after, whatever happens here. ++ """ ++ fd = _FATAL_NOTIFY_FD ++ if fd is None: ++ return ++ try: ++ os.set_blocking(fd, False) ++ os.write(fd, b"F") ++ except BaseException: ++ pass ++ + + def _safe_type_label(cls: type) -> str: + """``module.qualname`` of ``cls`` without running class-defined code.""" +@@ -240,6 +274,7 @@ def ple_fail_closed( + # runs a Python __exit__ between here and os._exit below; a bare + # except keeps that window empty. + pass ++ notify_parent_fatal() + os._exit(1) + + +-- +2.43.0 + diff --git a/patches/0059-kv-share-token-to-request-mappings-across-KV-cache-g.patch b/patches/0059-kv-share-token-to-request-mappings-across-KV-cache-g.patch new file mode 100644 index 0000000..dc4363b --- /dev/null +++ b/patches/0059-kv-share-token-to-request-mappings-across-KV-cache-g.patch @@ -0,0 +1,43 @@ +From 22e10be45aebad0611c2a55ebf6318e08ad2db0b Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 17 Sep 2026 20:14:36 -0700 +Subject: [PATCH 059/103] kv: share token-to-request mappings across KV cache + groups (port of upstream #57102) + +Port of the upstream vLLM change #57102 (Apache-2.0). KV cache groups that use the same token-to-request mapping now share one +instead of rebuilding it per group. +--- + vllm/v1/worker/gpu/attn_utils.py | 3 +++ + 1 file changed, 3 insertions(+) + +diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py +index eccfdba..89f54d4 100644 +--- a/vllm/v1/worker/gpu/attn_utils.py ++++ b/vllm/v1/worker/gpu/attn_utils.py +@@ -426,6 +426,7 @@ def build_attn_metadata( + seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] + + attn_metadata: dict[str, Any] = {} ++ token_to_req_indices: torch.Tensor | None = None + num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) + for i in range(num_kv_cache_groups): + if not attn_groups[i]: +@@ -472,6 +473,7 @@ def build_attn_metadata( + mm_req_doc_ranges=mm_req_doc_ranges, + rswa_prefix_lens=rswa_prefix_lens, + req_idx=req_idx, ++ _token_to_req_indices_cache=token_to_req_indices, + **common_attn_metadata_extra_kwargs, + ) + +@@ -497,6 +499,7 @@ def build_attn_metadata( + ) + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = metadata ++ token_to_req_indices = common_attn_metadata._token_to_req_indices_cache + return attn_metadata + + +-- +2.43.0 + diff --git a/patches/0060-spec-decode-annotate-MTP-draft-KV-cache-groups-posit.patch b/patches/0060-spec-decode-annotate-MTP-draft-KV-cache-groups-posit.patch new file mode 100644 index 0000000..bb6980f --- /dev/null +++ b/patches/0060-spec-decode-annotate-MTP-draft-KV-cache-groups-posit.patch @@ -0,0 +1,254 @@ +From 1608d7132324e5f2f133f2b61e6a7405d0472020 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 21 Sep 2026 20:31:26 -0400 +Subject: [PATCH 060/103] spec-decode: annotate MTP draft KV cache groups + positionally on the hybrid path (port of upstream #55390) + +Port of the upstream vLLM change #55390 (Apache-2.0). On the hybrid grouping path the MTP draft's KV cache group is identified by +position. The only conflict was a docstring; upstream's text was taken. +--- + tests/v1/core/test_kv_cache_utils.py | 96 ++++++++++++++++++++++++++-- + vllm/v1/core/kv_cache_utils.py | 59 ++++++++++++----- + 2 files changed, 134 insertions(+), 21 deletions(-) + +diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py +index c6354cd..a2231d5 100644 +--- a/tests/v1/core/test_kv_cache_utils.py ++++ b/tests/v1/core/test_kv_cache_utils.py +@@ -4226,12 +4226,16 @@ def _deepseek_v4_specs(model_version="deepseek_v4"): + } + + +-def test_deepseek_v4_draft_group_annotated_on_packed_path(): ++@pytest.mark.parametrize( ++ ("method", "model_type"), ++ [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], ++) ++def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec + # carries no draft marker and only the positional rule can find it. This + # pins the pre-existing behaviour that the unified annotator must preserve. + groups = get_kv_cache_groups( +- _spec_decode_grouping_config(method="mtp", model_type="deepseek_v4"), ++ _spec_decode_grouping_config(method=method, model_type=model_type), + _deepseek_v4_specs(model_version=None), + ) + +@@ -4240,13 +4244,93 @@ def test_deepseek_v4_draft_group_annotated_on_packed_path(): + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +-def test_deepseek_v4_annotation_requires_model_type(): +- # The positional rule is only sound for DeepseekV4, where the draft layer +- # is known to be registered last. Without that model gate nothing may be +- # flagged, however the grouping happens to fall out. ++def test_trailing_layer_fallback_applies_to_any_mtp_model(): ++ # The positional rule is sound for every MTP drafter, not just DeepseekV4: ++ # MTP blocks reuse the target's decoder layer (no spec marker) and always ++ # register after every target layer. The model_type must not gate it. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="other"), + _deepseek_v4_specs(), + ) + ++ flagged = [g for g in groups if g.is_eagle_group] ++ assert len(flagged) == 1 ++ assert "model.layers.3.self_attn.attn" in flagged[0].layer_names ++ ++ ++def _qwen3_5_hybrid_specs(with_mtp_layer: bool): ++ """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with ++ the MTP drafter's full-attn layer (spec-identical to the target's) ++ registered last.""" ++ specs = {} ++ idx = 0 ++ for _ in range(2): ++ for _ in range(3): ++ specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( ++ mamba_cache_mode="align" ++ ) ++ idx += 1 ++ specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() ++ idx += 1 ++ if with_mtp_layer: ++ specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() ++ return specs ++ ++ ++def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): ++ # A hybrid mamba + full-attention model with an MTP drafter that is ++ # spec-indistinguishable from the target reaches the general multi-group ++ # path. The trailing-layer rule must locate the draft group there so the ++ # Mamba groups are not swept up by the flag-all consumer fallback. ++ groups = get_kv_cache_groups( ++ _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), ++ _qwen3_5_hybrid_specs(with_mtp_layer=True), ++ ) ++ ++ flagged = [g for g in groups if g.is_eagle_group] ++ assert len(flagged) == 1 ++ assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names ++ for group in groups: ++ if any( ++ isinstance(spec, MambaSpec) ++ for spec in iter_layer_specs(group.kv_cache_spec) ++ ): ++ assert not group.is_eagle_group ++ assert "could be identified as the draft model's" not in caplog_vllm.text ++ ++ ++@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) ++def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): ++ # Other drafters are not covered by the trailing-layer rule, so an ++ # unidentifiable hybrid draft must still warn. ++ groups = get_kv_cache_groups( ++ _spec_decode_grouping_config(method=method, model_type="qwen3_5"), ++ _qwen3_5_hybrid_specs(with_mtp_layer=True), ++ ) ++ + assert not any(g.is_eagle_group for g in groups) ++ assert "could be identified as the draft model's" in caplog_vllm.text ++ ++ ++def test_trailing_layer_fallback_requires_exact_partition(): ++ # If the groups do not partition the layers exactly (e.g. a caller that ++ # dropped or duplicated layers), the positional rule is meaningless and ++ # must not fire. ++ from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups ++ ++ specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) ++ config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") ++ groups = get_kv_cache_groups(config, specs) ++ for g in groups: ++ g.is_eagle_group = False ++ # Remove one layer from its group: no longer an exact partition. ++ trimmed = [ ++ KVCacheGroupSpec( ++ [n for n in g.layer_names if n != "model.layers.0.linear_attn"], ++ g.kv_cache_spec, ++ ) ++ for g in groups ++ ] ++ _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) ++ ++ assert not any(g.is_eagle_group for g in trimmed) +diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py +index d1f827d..ad179e6 100644 +--- a/vllm/v1/core/kv_cache_utils.py ++++ b/vllm/v1/core/kv_cache_utils.py +@@ -2113,16 +2113,24 @@ def _get_packed_kv_cache_groups( + vllm_config, + kv_cache_spec, + groups, +- use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config), ++ use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +-def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: ++def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: ++ """Whether the drafter's KV layers can be located positionally. ++ ++ MTP drafters register their KV layers after the target layers, without a ++ distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. ++ The annotator separately checks that the groups partition the layers exactly. ++ """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return False ++ if spec_config.method == "mtp": ++ return True + model_config = vllm_config.model_config + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", +@@ -2130,11 +2138,24 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: + ) + + ++def _groups_partition_layers_exactly( ++ kv_cache_spec: dict[str, KVCacheSpec], ++ kv_cache_groups: list[KVCacheGroupSpec], ++) -> bool: ++ """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. ++ ++ The trailing-layer draft fallback is only meaningful when the last ++ registered layer is guaranteed to appear in exactly one group. ++ """ ++ names = [name for group in kv_cache_groups for name in group.layer_names] ++ return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) ++ ++ + def _annotate_eagle_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +- use_deepseek_v4_fallback: bool = False, ++ use_trailing_layer_fallback: bool = False, + ) -> None: + """Flag the KV cache groups that hold drafter attention layers. + +@@ -2147,14 +2168,14 @@ def _annotate_eagle_groups( + spec merging, wherever grouping happens to land. It is sufficient but + not necessary: a drafter whose spec is indistinguishable from the + target's cannot be found this way. +- 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block +- reuses the target's own decoder layer and so carries no spec marker. Its +- draft attention layer is always the last registered layer, so flag whichever +- group holds it. This rule is only valid where the groups partition +- exactly the layers of ``kv_cache_spec``, which is true on the packed +- grouping path and not in general; other callers must leave +- ``use_deepseek_v4_fallback`` False. The caller gates this fallback on +- the configured model type. ++ 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose ++ MTP block reuses the target's own decoder layer and so carries no spec ++ marker. Their draft attention layers always register after every ++ target layer, so flag whichever group holds the last registered layer. ++ This rule is only valid where the groups partition exactly the layers ++ of ``kv_cache_spec``, which is re-checked here before applying it. ++ When a drafter's trailing caches span several groups, this rule flags ++ only the group holding the very last layer and must be generalized. + FIXME(yifan): avoid/generalize this hacky check. + + Args: +@@ -2162,8 +2183,9 @@ def _annotate_eagle_groups( + kv_cache_spec: The kv cache spec of each attention layer, in layer + registration order. Only read by rule 2. + kv_cache_groups: Groups to annotate in place. +- use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed +- group. ++ use_trailing_layer_fallback: Enable rule 2. Callers gate this on ++ ``_uses_trailing_mtp_layers``. ++ + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle_block_drop(): +@@ -2176,7 +2198,9 @@ def _annotate_eagle_groups( + ): + group.is_eagle_group = True + +- if not use_deepseek_v4_fallback: ++ if not use_trailing_layer_fallback: ++ return ++ if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): + return + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: +@@ -2346,7 +2370,12 @@ def get_kv_cache_groups( + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) + groups.append(KVCacheGroupSpec([name], aligned)) + +- _annotate_eagle_groups(vllm_config, kv_cache_spec, groups) ++ _annotate_eagle_groups( ++ vllm_config, ++ kv_cache_spec, ++ groups, ++ use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), ++ ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + +-- +2.43.0 + diff --git a/patches/0061-gdn-fix-stateless-first-chunk-classification-port-of.patch b/patches/0061-gdn-fix-stateless-first-chunk-classification-port-of.patch new file mode 100644 index 0000000..f63f6ff --- /dev/null +++ b/patches/0061-gdn-fix-stateless-first-chunk-classification-port-of.patch @@ -0,0 +1,147 @@ +From 108e2c9e5b62d1990796ee3f710d39f5063b3759 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 22 Sep 2026 09:11:49 +0800 +Subject: [PATCH 061/103] gdn: fix stateless first-chunk classification (port + of upstream #51565) + +Port of the upstream vLLM change #51565 (Apache-2.0). A Gated DeltaNet prefill whose first chunk carries no prior state is +classified correctly. +--- + .../v1/attention/test_gdn_metadata_builder.py | 83 +++++++++++++++++++ + vllm/v1/attention/backends/gdn_attn.py | 27 +++++- + 2 files changed, 109 insertions(+), 1 deletion(-) + +diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py +index 221f933..b4d3acf 100644 +--- a/tests/v1/attention/test_gdn_metadata_builder.py ++++ b/tests/v1/attention/test_gdn_metadata_builder.py +@@ -221,3 +221,86 @@ def test_full_cudagraph_spec_metadata_uses_request_count(): + assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,) + assert meta.num_accepted_tokens is not None + assert meta.num_accepted_tokens.shape == (batch.batch_size,) ++ ++ ++def _build_non_spec( ++ batch: BatchSpec, ++ is_prefilling: list[bool] | None, ++ full_cuda_graph: bool = False, ++): ++ common_attn_metadata = create_common_attn_metadata( ++ batch, BLOCK_SIZE, DEVICE ++ ).replace( ++ is_prefilling=None ++ if is_prefilling is None ++ else torch.tensor(is_prefilling, dtype=torch.bool) ++ ) ++ builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph) ++ return builder, common_attn_metadata, builder.build(0, common_attn_metadata) ++ ++ ++@pytest.mark.parametrize( ++ ("seq_len", "query_len", "is_prefilling", "num_prefills"), ++ [ ++ pytest.param(1, 1, True, 1, id="first-chunk"), ++ pytest.param(65, 1, True, 0, id="resumed-chunk"), ++ pytest.param(0, 0, True, 0, id="padding"), ++ pytest.param(1, 1, None, 0, id="missing-prefill-flag"), ++ ], ++) ++def test_one_token_chunk_classification( ++ seq_len: int, ++ query_len: int, ++ is_prefilling: bool | None, ++ num_prefills: int, ++): ++ """Only a real first chunk with a prefill flag needs state initialization.""" ++ _, _, meta = _build_non_spec( ++ BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]), ++ is_prefilling=None if is_prefilling is None else [False, is_prefilling], ++ ) ++ ++ assert meta.num_prefills == num_prefills ++ assert meta.num_decodes == 2 - num_prefills ++ assert meta.num_prefill_tokens == num_prefills ++ assert meta.num_decode_tokens == 1 + query_len - num_prefills ++ if num_prefills: ++ assert meta.has_initial_state is not None ++ assert meta.has_initial_state.tolist() == [True, False] ++ else: ++ assert meta.has_initial_state is None ++ ++ ++def test_one_token_first_chunk_excludes_padding(): ++ """Neither padding requests nor padding tokens count as prefill work.""" ++ common = create_common_attn_metadata( ++ BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]), ++ BLOCK_SIZE, ++ DEVICE, ++ ).replace( ++ is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool), ++ num_actual_tokens=4, ++ ) ++ meta = _create_gdn_builder().build(0, common) ++ ++ assert meta.num_decodes == 1 ++ assert meta.num_prefills == 1 ++ assert meta.num_decode_tokens == 1 ++ assert meta.num_prefill_tokens == 1 ++ ++ ++def test_cudagraph_capture_batch_stays_decode_only(): ++ """Capture rows have no history, but must still select decode kernels.""" ++ batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4) ++ builder, common_attn_metadata, _ = _build_non_spec( ++ batch, [False] * 4, full_cuda_graph=True ++ ) ++ meta = builder.build_for_cudagraph_capture(common_attn_metadata) ++ ++ assert meta.num_prefills == 0 ++ assert meta.num_decodes == 4 ++ assert meta.has_initial_state is None ++ staged = meta.non_spec_state_indices_tensor ++ assert staged is not None ++ assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr() ++ torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0]) +diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py +index 23daff5..dfc04a8 100644 +--- a/vllm/v1/attention/backends/gdn_attn.py ++++ b/vllm/v1/attention/backends/gdn_attn.py +@@ -247,9 +247,34 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + ) + + if spec_sequence_masks is None: ++ # V2 already excludes prefills from full decode graphs via ++ # has_prefill. Classify first chunks as prefills to mask recycled ++ # state; resumed one-token chunks can still use the decode kernels. ++ assert m.seq_lens_cpu_upper_bound is not None ++ query_lens_cpu = query_start_loc_cpu.diff() ++ no_prior_state = (query_lens_cpu > 0) & ( ++ m.seq_lens_cpu_upper_bound <= query_lens_cpu ++ ) ++ # Capture batches also have seq_len == query_len, but are not ++ # prefills. ++ if m.is_prefilling is not None: ++ no_prior_state &= m.is_prefilling ++ else: ++ no_prior_state = torch.zeros_like(no_prior_state) + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( +- split_decodes_and_prefills(m, decode_threshold=1) ++ split_decodes_and_prefills( ++ m.replace(is_prefilling=no_prior_state), ++ decode_threshold=1, ++ treat_short_extends_as_decodes=False, ++ ) + ) ++ # Exclude trailing padding from both prefill counts. ++ if num_prefills: ++ num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum()) ++ num_prefill_tokens = ( ++ int(query_start_loc_cpu[num_decodes + num_prefills]) ++ - num_decode_tokens ++ ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None +-- +2.43.0 + diff --git a/patches/0062-cudagraph-capture-prefill-kernels-for-mixed-FULL-gra.patch b/patches/0062-cudagraph-capture-prefill-kernels-for-mixed-FULL-gra.patch new file mode 100644 index 0000000..6d8cdd0 --- /dev/null +++ b/patches/0062-cudagraph-capture-prefill-kernels-for-mixed-FULL-gra.patch @@ -0,0 +1,248 @@ +From f26f0f4a9cf4e746112161a64838ed2d80ae18b3 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 24 Sep 2026 04:05:07 -0500 +Subject: [PATCH 062/103] cudagraph: capture prefill kernels for mixed FULL + graphs (port of upstream #58275) + +Port of the upstream vLLM change #58275 (Apache-2.0). Mixed prefill/decode batches under FULL CUDA graphs capture the prefill kernels +they run. +--- + .../compile/fullgraph/test_full_cudagraph.py | 40 +++++---- + tests/v1/cudagraph/test_cudagraph_manager.py | 86 +++++++++++++++++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 6 +- + .../autoregressive/cudagraph_utils.py | 1 + + 4 files changed, 117 insertions(+), 16 deletions(-) + +diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py +index c2af7be..5d010aa 100644 +--- a/tests/compile/fullgraph/test_full_cudagraph.py ++++ b/tests/compile/fullgraph/test_full_cudagraph.py +@@ -37,7 +37,7 @@ def temporary_environ(env_vars): + model_backends_full_cudagraph = [] + + # deepseek-ai/DeepSeek-V2-Lite with MLA +-MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"] ++MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA", "FlashInferMLA"] + for mla_backend in MLA_backends: + model_backends_full_cudagraph.append( + ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend]) +@@ -54,27 +54,28 @@ for backend_config in other_backend_configs: + @pytest.fixture(scope="class") + def llm_pair(request): + model, backend_config, use_inductor_graph_partition = request.param +- backend_config.comp_config["use_inductor_graph_partition"] = ( +- use_inductor_graph_partition +- ) +- + if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): + pytest.skip("Inductor graph partition only supported in torch>=2.9") + ++ backend = AttentionBackendEnum[backend_config.attention_config["backend"]] ++ if current_platform.is_rocm() and backend not in ( ++ AttentionBackendEnum.TRITON_ATTN, ++ AttentionBackendEnum.ROCM_ATTN, ++ ): ++ pytest.skip(f"{backend_config.name} is a CUDA-only attention backend") ++ if backend == AttentionBackendEnum.ROCM_ATTN and not current_platform.is_rocm(): ++ pytest.skip("ROCM_ATTN requires ROCm") ++ + # Dynamically skip test if GPU capability is not met +- if ( +- backend_config.specific_gpu_arch +- and backend_config.specific_gpu_arch != current_platform.get_device_capability() ++ if backend_config.specific_gpu_arch and ( ++ not current_platform.is_cuda() ++ or backend_config.specific_gpu_arch != current_platform.get_device_capability() + ): + if backend_config.specific_gpu_arch == (9, 0): + pytest.skip("Only Hopper GPUs support FA3 and FlashMLA") + elif backend_config.specific_gpu_arch == (10, 0): + pytest.skip("Only Blackwell GPUs support Cutlass MLA") + +- # FlashInfer is not supported on ROCm +- if backend_config == AttentionBackendEnum.FLASHINFER and current_platform.is_rocm(): +- pytest.skip("FlashInfer is not supported on ROCm") +- + env_vars = { + # Force native sampler to avoid potential nondeterminism in FlashInfer + # when per-request generators are not used in V1. +@@ -87,7 +88,11 @@ def llm_pair(request): + trust_remote_code=True, + max_model_len=1024, + max_num_seqs=128, +- compilation_config=CompilationConfig(**backend_config.comp_config), ++ attention_config=backend_config.attention_config, ++ compilation_config=CompilationConfig( ++ **backend_config.comp_config, ++ use_inductor_graph_partition=use_inductor_graph_partition, ++ ), + generation_config="vllm", + seed=42, + ) +@@ -97,8 +102,10 @@ def llm_pair(request): + trust_remote_code=True, + max_model_len=1024, + max_num_seqs=128, ++ attention_config=backend_config.attention_config, + compilation_config=CompilationConfig( +- cudagraph_mode=CUDAGraphMode.PIECEWISE ++ cudagraph_mode=CUDAGraphMode.PIECEWISE, ++ use_inductor_graph_partition=use_inductor_graph_partition, + ), + generation_config="vllm", + seed=42, +@@ -118,7 +125,10 @@ def llm_pair(request): + @pytest.mark.parametrize( + "llm_pair", + [ +- pytest.param((model, backend_config, use_inductor_graph_partition)) ++ pytest.param( ++ (model, backend_config, use_inductor_graph_partition), ++ id=f"{backend_config.name}-partition={use_inductor_graph_partition}", ++ ) + for model, backend_config in model_backends_full_cudagraph + for use_inductor_graph_partition in [True, False] + ], +diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py +index c907be3..76c4732 100644 +--- a/tests/v1/cudagraph/test_cudagraph_manager.py ++++ b/tests/v1/cudagraph/test_cudagraph_manager.py +@@ -22,6 +22,9 @@ from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + from vllm.v1.worker.gpu.input_batch import InputBuffers + from vllm.v1.worker.gpu.pcp_manager import PCPManager ++from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( ++ SpeculatorCudaGraphManager, ++) + + pytestmark = pytest.mark.cpu_test + +@@ -165,6 +168,89 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings(): + block_tables.get_dummy_slot_mappings.assert_not_called() + + ++@pytest.mark.parametrize( ++ "full_cudagraph,max_query_len,expected_max_query_len", ++ [(False, None, None), (True, None, 8), (True, 1, 1), (True, 3, 3)], ++) ++def test_capture_query_length_covers_replay_batches( ++ full_cudagraph, max_query_len, expected_max_query_len ++): ++ """Mixed FULL graphs must capture prefill kernels as well as decode kernels.""" ++ num_tokens = num_reqs = 8 ++ buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu")) ++ block_tables = MagicMock() ++ block_tables.cp_size = 1 ++ block_tables.get_dummy_block_tables.return_value = () ++ block_tables.get_dummy_slot_mappings.return_value = torch.empty( ++ 0, num_tokens, dtype=torch.int64 ++ ) ++ model_state = MagicMock() ++ kv_cache_config = KVCacheConfig( ++ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[] ++ ) ++ ++ gpu_cudagraph_utils.prepare_inputs_to_capture( ++ num_reqs, ++ num_tokens, ++ model_state, ++ buffers, ++ block_tables, ++ [], ++ kv_cache_config, ++ full_cudagraph=full_cudagraph, ++ max_query_len=max_query_len, ++ ) ++ ++ batch = model_state.prepare_attn.call_args.args[0] ++ assert batch.num_scheduled_tokens.max() == 1 ++ assert batch.max_query_len == expected_max_query_len ++ ++ ++@pytest.mark.parametrize( ++ "uniform_token_count,max_query_len,expected_max_query_len", ++ [(1, None, 1), (None, 3, 3), (None, None, 8)], ++) ++def test_speculator_capture_preserves_decode_query_bounds( ++ monkeypatch, uniform_token_count, max_query_len, expected_max_query_len ++): ++ """Draft decode graphs keep narrow queries while mixed graphs cover prefill.""" ++ num_tokens = num_reqs = 8 ++ desc = BatchExecutionDescriptor( ++ CUDAGraphMode.FULL, ++ num_tokens, ++ num_reqs, ++ uniform_token_count=uniform_token_count, ++ max_query_len=max_query_len, ++ ) ++ manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager) ++ manager.max_num_reqs = num_reqs ++ manager.dp_size = 1 ++ buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu")) ++ block_tables = MagicMock() ++ block_tables.cp_size = 1 ++ block_tables.get_dummy_block_tables.return_value = () ++ block_tables.get_dummy_slot_mappings.return_value = torch.empty( ++ 0, num_tokens, dtype=torch.int64 ++ ) ++ model_state = MagicMock() ++ kv_cache_config = KVCacheConfig( ++ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[] ++ ) ++ ++ def capture(self, create_forward_fn, progress_bar_desc): ++ for warmup in (True, False): ++ create_forward_fn(desc, warmup)(desc.cg_mode) ++ ++ monkeypatch.setattr(gpu_cudagraph_utils.CudaGraphManager, "capture", capture) ++ manager.capture( ++ MagicMock(), model_state, buffers, block_tables, [], kv_cache_config ++ ) ++ ++ assert model_state.prepare_attn.call_count == 2 ++ for call in model_state.prepare_attn.call_args_list: ++ assert call.args[0].max_query_len == expected_max_query_len ++ ++ + _DECODE_QUERY_LEN = 3 + + +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 91d4d92..921c562 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -818,7 +818,7 @@ class ModelCudaGraphManager(CudaGraphManager): + attn_groups, + kv_cache_config, + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, +- max_query_len=desc.max_query_len, ++ max_query_len=desc.max_query_len or desc.uniform_token_count, + pcp_manager=pcp_manager, + ) + +@@ -889,6 +889,10 @@ def prepare_inputs_to_capture( + max_query_len: int | None = None, + pcp_manager: "PCPManager | None" = None, + ) -> AttentionState: ++ if full_cudagraph and max_query_len is None: ++ # Mixed graphs can replay a single prefill spanning the entire batch, ++ # even when the dummy batch distributes one token to each request. ++ max_query_len = num_tokens + input_batch = InputBatch.make_dummy( + num_reqs, num_tokens, input_buffers, max_query_len=max_query_len + ) +diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +index ef3b6e2..fe3bdef 100644 +--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +@@ -57,6 +57,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager): + attn_groups, + kv_cache_config, + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, ++ max_query_len=desc.max_query_len or desc.uniform_token_count, + ) + + return lambda cg_mode: forward_fn( +-- +2.43.0 + diff --git a/patches/0063-mamba-restore-prompt-tail-prefix-cache-hits-with-MTP.patch b/patches/0063-mamba-restore-prompt-tail-prefix-cache-hits-with-MTP.patch new file mode 100644 index 0000000..7f2294a --- /dev/null +++ b/patches/0063-mamba-restore-prompt-tail-prefix-cache-hits-with-MTP.patch @@ -0,0 +1,61 @@ +From f4efb5ef63426dad77a45b5a9787fd05fbac5638 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Fri, 25 Sep 2026 02:46:42 +0300 +Subject: [PATCH 063/103] mamba: restore prompt-tail prefix-cache hits with MTP + (port of upstream #58368) + +Port of the upstream vLLM change #58368 (Apache-2.0). With MTP enabled, a prompt's tail block can again be served from the prefix cache. +--- + .../test_mamba_eagle_resume_checkpoint.py | 22 +++++++++++++++++++ + vllm/v1/core/single_type_kv_cache_manager.py | 2 +- + 2 files changed, 23 insertions(+), 1 deletion(-) + +diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +index 059b400..87a056a 100644 +--- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py ++++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +@@ -196,6 +196,28 @@ def test_sibling_resumes_below_the_block_grid_when_the_prefix_ends_early(): + assert hit == resume, f"expected the resume point at {resume}, got {hit}" + + ++@pytest.mark.parametrize("mtp_draft_group_identified", [False, True]) ++def test_mamba_prefix_cache_drops_hash_block(mtp_draft_group_identified): ++ block_size, hash_block_size = 1024, 64 ++ manager = _manager( ++ block_size, ++ hash_block_size, ++ eagle_group=0 if mtp_draft_group_identified else None, ++ ) ++ prompt_a = PREFIX[:1600] ++ hit_after_mtp_block_drop = len(prompt_a) - hash_block_size ++ ++ _prefill( ++ manager, ++ _stub(manager, block_size, hash_block_size), ++ make_request("A", prompt_a, hash_block_size, sha256), ++ ) ++ ++ tokens_after_a = [-1] * 500 ++ hit = _sibling_hit(manager, len(prompt_a), tokens_after_a, hash_block_size) ++ assert hit == hit_after_mtp_block_drop ++ ++ + # -------------------------------------------------------------------------- + # Invariants the two clauses have to keep + +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 8242e25..76ba691 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -2126,7 +2126,7 @@ class MambaManager(SingleTypeKVCacheManager): + latest_prompt_hash_boundary = ( + request.num_prompt_tokens // hash_block_size + ) * hash_block_size +- if self.use_eagle: ++ if self.drop_eagle_checkpoint_block: + # Eagle groups match one hash unit past the candidate and drop it, + # so register the tail one unit lower. + latest_prompt_hash_boundary = max( +-- +2.43.0 + diff --git a/patches/0064-mamba-reuse-Mamba-GDN-metadata-across-KV-cache-group.patch b/patches/0064-mamba-reuse-Mamba-GDN-metadata-across-KV-cache-group.patch new file mode 100644 index 0000000..d58aa43 --- /dev/null +++ b/patches/0064-mamba-reuse-Mamba-GDN-metadata-across-KV-cache-group.patch @@ -0,0 +1,362 @@ +From 26f4ff1480ab7307aaad9d27309fb5f5a85061e3 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 28 Sep 2026 06:36:58 +0200 +Subject: [PATCH 064/103] mamba: reuse Mamba/GDN metadata across KV cache + groups (port of upstream #58762) + +Port of the upstream vLLM change #58762 (Apache-2.0). Mamba/GDN attention metadata is built once and reused by groups that share it. +Test-only conflicts with upstream changes not in v0.30.0 were resolved by taking only this change's tests. +--- + tests/models/kimi_k3/test_kda_metadata.py | 1 + + .../v1/attention/test_gdn_metadata_builder.py | 63 +++++++++ + vllm/v1/attention/backends/gdn_attn.py | 124 +++++++++++++++--- + vllm/v1/worker/gpu/attn_utils.py | 18 ++- + 4 files changed, 186 insertions(+), 20 deletions(-) + +diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py +index c3f4238..8ff7aff 100644 +--- a/tests/models/kimi_k3/test_kda_metadata.py ++++ b/tests/models/kimi_k3/test_kda_metadata.py +@@ -49,6 +49,7 @@ PRUNED_METADATA_FIELDS = { + "prefill_state_indices", + "prefill_has_initial_state", + "spec_sequence_masks", ++ "spec_sequence_masks_cpu", + "flashinfer_prefill_query_start_loc", + "flashinfer_prefill_seq_order", + } +diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py +index b4d3acf..44dee49 100644 +--- a/tests/v1/attention/test_gdn_metadata_builder.py ++++ b/tests/v1/attention/test_gdn_metadata_builder.py +@@ -21,6 +21,7 @@ from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, + ) ++from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor + from vllm.v1.kv_cache_interface import MambaSpec + + BLOCK_SIZE = 16 +@@ -75,6 +76,17 @@ GDN_BUILD_TEST_CASES = { + expected_num_prefill_tokens=0, + expected_num_spec_decodes=0, + ), ++ # No speculative config, decode alongside prefill ++ "regular_decode_with_prefill": GDNBuildTestCase( ++ seq_lens=[40, 100], ++ query_lens=[1, 50], ++ num_decode_draft_tokens=None, ++ num_speculative_tokens=0, ++ expected_num_decodes=1, ++ expected_num_prefills=1, ++ expected_num_prefill_tokens=50, ++ expected_num_spec_decodes=0, ++ ), + # Multi-token prefill alongside spec decode — no decode to reclassify + "spec_decode_with_real_prefill": GDNBuildTestCase( + seq_lens=[100, 20], +@@ -155,9 +167,12 @@ def _build( + builder: GDNAttentionMetadataBuilder, + batch_spec: BatchSpec, + num_decode_draft_tokens: list[int] | None = None, ++ block_table: torch.Tensor | None = None, + ) -> GDNAttentionMetadata: + """Build GDN attention metadata, optionally with spec-decode kwargs.""" + common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE) ++ if block_table is not None: ++ common = common.replace(block_table_tensor=block_table) + kwargs: dict = {} + if num_decode_draft_tokens is not None: + kwargs["num_decode_draft_tokens_cpu"] = torch.tensor( +@@ -184,6 +199,54 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase): + assert meta.num_spec_decodes == test_case.expected_num_spec_decodes + + ++@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) ++@pytest.mark.parametrize("full_cuda_graph", [False, True]) ++@pytest.mark.parametrize( ++ "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() ++) ++def test_update_block_table_matches_build( ++ test_case: GDNBuildTestCase, full_cuda_graph: bool, mamba_cache_mode: str ++): ++ """update_block_table() on another group's metadata matches build().""" ++ batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens) ++ src, dst, ref = ( ++ _create_gdn_builder(test_case.num_speculative_tokens, full_cuda_graph) ++ for _ in range(3) ++ ) ++ for builder in (src, dst, ref): ++ builder.vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode ++ common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE) ++ if mamba_cache_mode == "align": ++ # MRV2 precomputes these for every Mamba group each step. ++ dst.mamba_aligned_state_indices = ref.mamba_aligned_state_indices = ( ++ mamba_get_block_table_tensor( ++ common.block_table_tensor, common.seq_lens, ref.kv_cache_spec, "align" ++ ) ++ ) ++ draft_tokens = test_case.num_decode_draft_tokens ++ expected = _build(ref, batch, draft_tokens, common.block_table_tensor) ++ source = _build(src, batch, draft_tokens) ++ fields = ( ++ "spec_state_indices_tensor", ++ "non_spec_state_indices_tensor", ++ "prefill_state_indices", ++ ) ++ source_indices = [getattr(source, f) for f in fields] ++ source_indices = [t if t is None else t.clone() for t in source_indices] ++ meta = dst.update_block_table( ++ source, common.block_table_tensor, common.slot_mapping ++ ) ++ ++ for field, source_index in zip(fields, source_indices): ++ actual = getattr(meta, field) ++ torch.testing.assert_close(actual, getattr(expected, field)) ++ # The source group's indices are untouched. ++ torch.testing.assert_close(getattr(source, field), source_index) ++ # FULL graph state indices land in this group's own buffers. ++ if full_cuda_graph and meta.num_prefills == 0 and actual is not None: ++ assert actual.data_ptr() == getattr(dst, field).data_ptr() ++ ++ + def test_has_initial_state_after_reclassification(): + """After reclassification, num_prefills > 0 so the prefill kernel path + should compute has_initial_state. For the reclassified request with +diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py +index dfc04a8..f11e2a9 100644 +--- a/vllm/v1/attention/backends/gdn_attn.py ++++ b/vllm/v1/attention/backends/gdn_attn.py +@@ -2,7 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Backend for GatedDeltaNet attention.""" + +-from dataclasses import dataclass ++from dataclasses import dataclass, replace + from typing import Literal + + import torch +@@ -60,6 +60,7 @@ class GDNAttentionMetadata: + None # shape: [batch - num_spec_decodes,] + ) + spec_sequence_masks: torch.Tensor | None = None # shape: [batch,] ++ spec_sequence_masks_cpu: torch.Tensor | None = None # shape: [batch,] + spec_token_indx: torch.Tensor | None = None + non_spec_token_indx: torch.Tensor | None = None + +@@ -113,6 +114,18 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + self.use_full_cuda_graph: bool = ( + self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ) ++ # update_block_table() keeps the source group's batch-level FULL graph ++ # buffers, so only MRV2, which also reuses metadata at capture, may use ++ # it. ++ self.supports_update_block_table = ( ++ vllm_config.use_v2_model_runner ++ and device.type == "cuda" ++ # Not isinstance: KDA's RecoverSSM/checkpoint metadata is per group. ++ and type(self) is GDNAttentionMetadataBuilder ++ ) ++ if self.supports_update_block_table: ++ # Opts into MRV2's CUDA-only aligned-index precompute. ++ self.mamba_aligned_state_indices: torch.Tensor | None = None + + self.decode_cudagraph_max_bs: int = ( + self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1) +@@ -219,12 +232,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None +- block_table_tensor = mamba_get_block_table_tensor( +- m.block_table_tensor, +- m.seq_lens, +- self.kv_cache_spec, +- self.vllm_config.cache_config.mamba_cache_mode, +- ) ++ aligned_state_indices = getattr(self, "mamba_aligned_state_indices", None) ++ if aligned_state_indices is not None: ++ block_table_tensor = aligned_state_indices[: m.num_reqs] ++ else: ++ block_table_tensor = mamba_get_block_table_tensor( ++ m.block_table_tensor, ++ m.seq_lens, ++ self.kv_cache_spec, ++ self.vllm_config.cache_config.mamba_cache_mode, ++ ) + + spec_sequence_masks_cpu: torch.Tensor | None = None + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: +@@ -451,12 +468,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + # metadata below is indexed by request. + batch_size = m.num_reqs + +- if ( +- self.use_full_cuda_graph +- and num_prefills == 0 +- and num_decodes == 0 +- and num_spec_decodes <= self.decode_cudagraph_max_bs +- and num_spec_decode_tokens <= self.decode_cudagraph_max_bs ++ if self._stage_spec_decode( ++ num_prefills, num_decodes, num_spec_decodes, num_spec_decode_tokens + ): + assert spec_sequence_masks is not None + self.spec_state_indices_tensor[:num_spec_decodes].copy_( +@@ -497,12 +510,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + num_accepted_tokens[num_spec_decodes:].fill_(1) + +- if ( +- self.use_full_cuda_graph +- and num_prefills == 0 +- and num_spec_decodes == 0 +- and num_decodes <= self.decode_cudagraph_max_bs +- ): ++ if self._stage_decode(num_prefills, num_decodes, num_spec_decodes): + self.non_spec_state_indices_tensor[:num_decodes].copy_( + non_spec_state_indices_tensor, non_blocking=True + ) +@@ -537,6 +545,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, ++ spec_sequence_masks_cpu=spec_sequence_masks_cpu, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, +@@ -546,6 +555,83 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] + ) + return attn_metadata + ++ def _stage_spec_decode( ++ self, ++ num_prefills: int, ++ num_decodes: int, ++ num_spec_decodes: int, ++ num_spec_decode_tokens: int, ++ ) -> bool: ++ """Whether spec-decode metadata goes into the FULL cudagraph buffers.""" ++ return ( ++ self.use_full_cuda_graph ++ and num_prefills == 0 ++ and num_decodes == 0 ++ and num_spec_decodes <= self.decode_cudagraph_max_bs ++ and num_spec_decode_tokens <= self.decode_cudagraph_max_bs ++ ) ++ ++ def _stage_decode( ++ self, num_prefills: int, num_decodes: int, num_spec_decodes: int ++ ) -> bool: ++ """Whether decode metadata goes into the FULL cudagraph buffers.""" ++ return ( ++ self.use_full_cuda_graph ++ and num_prefills == 0 ++ and num_spec_decodes == 0 ++ and num_decodes <= self.decode_cudagraph_max_bs ++ ) ++ ++ def update_block_table( ++ self, ++ metadata: GDNAttentionMetadata, ++ blk_table: torch.Tensor, ++ slot_mapping: torch.Tensor, ++ ) -> GDNAttentionMetadata: ++ """Re-gather this group's state indices. The other fields are ++ batch-level and stay shared with ``metadata``.""" ++ m = metadata ++ if self.vllm_config.cache_config.mamba_cache_mode == "align": ++ assert self.mamba_aligned_state_indices is not None ++ blk_table = self.mamba_aligned_state_indices ++ masks = m.spec_sequence_masks_cpu ++ spec_indices = non_spec_indices = prefill_indices = None ++ if masks is None: ++ non_spec_indices = blk_table[:, 0] ++ if m.num_prefills > 0: ++ prefill_indices = non_spec_indices[m.num_decodes :] ++ else: ++ spec_indices = blk_table[masks, : self.num_spec + 1] ++ if m.num_prefills > 0: ++ non_spec_indices = prefill_indices = blk_table[~masks, 0] ++ ++ if self._stage_spec_decode( ++ m.num_prefills, m.num_decodes, m.num_spec_decodes, m.num_spec_decode_tokens ++ ): ++ assert m.spec_state_indices_tensor is not None ++ batch_size = m.spec_state_indices_tensor.shape[0] ++ self.spec_state_indices_tensor[: m.num_spec_decodes].copy_( ++ spec_indices, non_blocking=True ++ ) ++ spec_indices = self.spec_state_indices_tensor[:batch_size] ++ spec_indices[m.num_spec_decodes :].fill_(NULL_BLOCK_ID) ++ ++ if self._stage_decode(m.num_prefills, m.num_decodes, m.num_spec_decodes): ++ assert m.non_spec_state_indices_tensor is not None ++ batch_size = m.non_spec_state_indices_tensor.shape[0] ++ self.non_spec_state_indices_tensor[: m.num_decodes].copy_( ++ non_spec_indices, non_blocking=True ++ ) ++ non_spec_indices = self.non_spec_state_indices_tensor[:batch_size] ++ non_spec_indices[m.num_decodes :].fill_(NULL_BLOCK_ID) ++ ++ return replace( ++ m, ++ spec_state_indices_tensor=spec_indices, ++ non_spec_state_indices_tensor=non_spec_indices, ++ prefill_state_indices=prefill_indices, ++ ) ++ + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ): +diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py +index 89f54d4..cdf0eb8 100644 +--- a/vllm/v1/worker/gpu/attn_utils.py ++++ b/vllm/v1/worker/gpu/attn_utils.py +@@ -29,6 +29,7 @@ from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheConfig, + KVCacheSpec, ++ MambaSpec, + UniformTypeKVCacheSpecs, + ) + from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata +@@ -427,6 +428,10 @@ def build_attn_metadata( + + attn_metadata: dict[str, Any] = {} + token_to_req_indices: torch.Tensor | None = None ++ # Mamba groups with the same spec and builder differ only in their state ++ # indices, so later groups re-gather those from the first group's metadata. ++ # Also at capture, so FULL graphs share the batch-level buffers. ++ cached_metadata: dict[tuple[KVCacheSpec, type], Any] = {} + num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) + for i in range(num_kv_cache_groups): + if not attn_groups[i]: +@@ -479,7 +484,16 @@ def build_attn_metadata( + + for attn_group in attn_groups[i]: + attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx) +- if for_cudagraph_capture: ++ reuse_key = None ++ if attn_metadata_builder.supports_update_block_table and isinstance( ++ attn_group.kv_cache_spec, MambaSpec ++ ): ++ reuse_key = (attn_group.kv_cache_spec, type(attn_metadata_builder)) ++ if reuse_key in cached_metadata: ++ metadata = attn_metadata_builder.update_block_table( ++ cached_metadata[reuse_key], block_table, slot_mapping ++ ) ++ elif for_cudagraph_capture: + metadata = attn_metadata_builder.build_for_cudagraph_capture( + common_attn_metadata + ) +@@ -497,6 +511,8 @@ def build_attn_metadata( + common_attn_metadata=common_attn_metadata, + **attn_metadata_extra_kwargs, + ) ++ if reuse_key is not None and reuse_key not in cached_metadata: ++ cached_metadata[reuse_key] = metadata + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = metadata + token_to_req_indices = common_attn_metadata._token_to_req_indices_cache +-- +2.43.0 + diff --git a/patches/0065-ple-reduce-PLE-metadata-construction-overhead-port-o.patch b/patches/0065-ple-reduce-PLE-metadata-construction-overhead-port-o.patch new file mode 100644 index 0000000..5c84f4b --- /dev/null +++ b/patches/0065-ple-reduce-PLE-metadata-construction-overhead-port-o.patch @@ -0,0 +1,421 @@ +From b24ac2b5bbf91bd73342d275f39aa9b6bb3f338a Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 28 Sep 2026 20:46:22 +0800 +Subject: [PATCH 065/103] ple: reduce PLE metadata construction overhead (port + of upstream #58114) + +Port of the upstream vLLM change #58114 (Apache-2.0). Applied cleanly, including the NVIDIA and AMD PLE layers; the pull transport +does not read the changed metadata fields. +--- + tests/models/qwen4_exp/test_ple.py | 125 ++++++++++++++++++ + vllm/models/qwen4_exp/amd/ple_layer.py | 9 +- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 3 +- + vllm/v1/attention/backends/mamba_attn.py | 20 +-- + vllm/v1/attention/backends/short_conv_attn.py | 82 +++--------- + 5 files changed, 161 insertions(+), 78 deletions(-) + +diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py +index e7d0e72..281275e 100644 +--- a/tests/models/qwen4_exp/test_ple.py ++++ b/tests/models/qwen4_exp/test_ple.py +@@ -14,6 +14,9 @@ from torch.nn import functional as F + import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module + import vllm.model_executor.parameter as parameter_module + import vllm.models.qwen4_exp.nvidia.ngram_embedding as ngram_embedding_module ++from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata ++from vllm.config import SpeculativeConfig, VllmConfig ++from vllm.config.compilation import CUDAGraphMode + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMixedPrecisionConfig, +@@ -35,8 +38,125 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( + from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer + from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadata, ++ PleShortConvAttentionMetadataBuilder, + ) + from vllm.v1.attention.backends.utils import NULL_BLOCK_ID ++from vllm.v1.kv_cache_interface import MambaSpec ++ ++ ++@pytest.mark.parametrize("device", ["cpu", "cuda"]) ++@pytest.mark.parametrize("full_graph", [False, True]) ++@pytest.mark.parametrize( ++ "query_lens,drafts,spec_rows,non_spec_rows,token_order", ++ [ ++ pytest.param([4, 2, 0], [3, 1, -1], [0, 1], [], None, id="spec-padding"), ++ pytest.param( ++ [3, 1, 2, 5, 0], ++ [2, -1, 1, -1, -1], ++ [0, 2], ++ [1, 3], ++ [0, 1, 2, 4, 5, 3, 6, 7, 8, 9, 10], ++ id="mixed-padding", ++ ), ++ pytest.param([1, 5], [-1, -1], [], [0, 1], None, id="non-spec"), ++ ], ++) ++def test_ple_metadata_preserves_request_and_token_order( ++ device, full_graph, query_lens, drafts, spec_rows, non_spec_rows, token_order ++): ++ """Preserve state routing, acceptance and prefill offsets without GPU sync.""" ++ if device == "cuda" and not torch.cuda.is_available(): ++ pytest.skip("CUDA required") ++ device = torch.device(device) ++ config = VllmConfig() ++ config.speculative_config = SpeculativeConfig( ++ method="ngram", num_speculative_tokens=3 ++ ) ++ config.compilation_config.cudagraph_mode = ( ++ CUDAGraphMode.FULL if full_graph else CUDAGraphMode.NONE ++ ) ++ config.compilation_config.max_cudagraph_capture_size = 32 ++ builder = PleShortConvAttentionMetadataBuilder( ++ MambaSpec(block_size=16, shapes=((16, 64),), dtypes=(torch.float32,)), ++ ["ple"], ++ config, ++ device, ++ ) ++ common = create_common_attn_metadata( ++ BatchSpec(seq_lens=[q + 16 for q in query_lens], query_lens=query_lens), ++ 16, ++ device, ++ arange_block_indices=True, ++ ) ++ common.is_prefilling = torch.tensor( ++ [q > 1 and d < 0 for q, d in zip(query_lens, drafts)] ++ ) ++ # Token padding is independent of the sum of per-request query lengths. ++ common.num_actual_tokens += 3 ++ accepted = torch.arange(1, len(query_lens) + 1, dtype=torch.int32, device=device) ++ drafts_cpu = torch.tensor(drafts, dtype=torch.int32) ++ sync_mode = torch.cuda.get_sync_debug_mode() if device.type == "cuda" else None ++ if sync_mode is not None: ++ torch.cuda.set_sync_debug_mode("error") ++ try: ++ metadata = builder.build( ++ 0, ++ common, ++ num_accepted_tokens=accepted, ++ num_decode_draft_tokens_cpu=drafts_cpu, ++ ) ++ finally: ++ if sync_mode is not None: ++ torch.cuda.set_sync_debug_mode(sync_mode) ++ ++ def check(actual, expected): ++ torch.testing.assert_close( ++ actual, torch.tensor(expected, dtype=actual.dtype, device=device) ++ ) ++ ++ assert metadata.num_spec_decodes == len(spec_rows) ++ assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows) ++ assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows) ++ if spec_rows: ++ torch.testing.assert_close( ++ metadata.spec_state_indices_tensor[: len(spec_rows)], ++ common.block_table_tensor[spec_rows, 0], ++ ) ++ torch.testing.assert_close( ++ metadata.num_accepted_tokens[: len(spec_rows)], accepted[spec_rows] ++ ) ++ check( ++ metadata.spec_query_start_loc[: len(spec_rows) + 1], ++ [0, *accumulate(query_lens[i] for i in spec_rows)], ++ ) ++ if full_graph and not non_spec_rows: ++ check(metadata.spec_state_indices_tensor[len(spec_rows) :], [NULL_BLOCK_ID]) ++ check(metadata.num_accepted_tokens[len(spec_rows) :], [1]) ++ check( ++ metadata.spec_query_start_loc[len(spec_rows) + 1 :], [sum(query_lens)] ++ ) ++ assert ( ++ metadata.spec_state_indices_tensor.data_ptr() ++ == builder.spec_state_indices_tensor.data_ptr() ++ ) ++ if non_spec_rows: ++ torch.testing.assert_close( ++ metadata.state_indices_tensor, common.block_table_tensor[non_spec_rows, 0] ++ ) ++ prefill_lens = [query_lens[i] for i in non_spec_rows if query_lens[i] > 1] ++ check(metadata.query_start_loc_p, [0, *accumulate(prefill_lens)]) ++ check(metadata.has_initial_states_p, [True] * len(prefill_lens)) ++ if token_order is not None: ++ check( ++ torch.cat((metadata.spec_token_indx, metadata.non_spec_token_indx)), ++ token_order, ++ ) ++ else: ++ assert metadata.spec_token_indx is None ++ assert metadata.non_spec_token_indx is None ++ assert metadata.nums_dict is None ++ assert metadata.batch_ptr is None ++ assert metadata.token_chunk_offset_ptr is None + + + def _mock_etp_group( +@@ -1501,6 +1621,11 @@ def _make_conv_metadata( + else None + ), + non_spec_query_start_loc=(non_spec_query_start_loc if has_non_spec else None), ++ query_start_loc_p=( ++ non_spec_query_start_loc[case.num_decodes :] - case.num_decodes ++ if num_prefills ++ else None ++ ), + has_initial_states_p=( + torch.arange(num_prefills, device=device) % 2 == 0 if num_prefills else None + ), +diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py +index fc494be..c686e34 100644 +--- a/vllm/models/qwen4_exp/amd/ple_layer.py ++++ b/vllm/models/qwen4_exp/amd/ple_layer.py +@@ -616,14 +616,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + num_decode_tokens: int, + num_prefill_tokens: int, + ) -> torch.Tensor: +- # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill) +- # requests and equals ``query_start_loc`` when spec-decode is inactive. +- non_spec_query_start_loc = metadata.non_spec_query_start_loc +- if non_spec_query_start_loc is None: ++ query_start_loc_p = metadata.query_start_loc_p ++ if query_start_loc_p is None: + raise ValueError("query_start_loc is required for prefill short-conv") +- query_start_loc_p = ( +- non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens +- ) + # The metadata builder guarantees that the prefill query offsets start + # at 0 and end at num_prefill_tokens. Avoid reading those values here, + # since doing so would force a device-to-host synchronization. +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index 94d908e..4925e03 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -320,10 +320,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + token_indices=token_indices_d, + ) + +- query_start_loc = metadata.non_spec_query_start_loc ++ query_start_loc = metadata.query_start_loc_p + if query_start_loc is None: + raise ValueError("query_start_loc is required for prefill short-conv") +- query_start_loc = query_start_loc[-num_prefills - 1 :] - num_decode_tokens + has_initial_states = metadata.has_initial_states_p + if has_initial_states is None: + raise ValueError( +diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py +index e0b5fab..817773a 100644 +--- a/vllm/v1/attention/backends/mamba_attn.py ++++ b/vllm/v1/attention/backends/mamba_attn.py +@@ -94,6 +94,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): + + # Will be disabled if speculative decoding is used + supports_update_block_table: bool = True ++ needs_causal_conv1d_metadata: bool = True + + def __init__( + self, +@@ -602,10 +603,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): + if num_computed_tokens is None: + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() + +- query_start_loc_p_cpu = ( +- common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :] +- - num_decode_tokens +- ) + query_start_loc_p = ( + common_attn_metadata.query_start_loc[-num_prefills - 1 :] + - num_decode_tokens +@@ -614,12 +611,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): + num_computed_tokens[num_reqs - num_prefills : num_reqs] > 0 + ) + +- nums_dict, batch_ptr, token_chunk_offset_ptr = ( +- compute_causal_conv1d_metadata( +- query_start_loc_p_cpu, +- device=common_attn_metadata.query_start_loc.device, ++ if self.needs_causal_conv1d_metadata: ++ query_start_loc_p_cpu = ( ++ common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :] ++ - num_decode_tokens ++ ) ++ nums_dict, batch_ptr, token_chunk_offset_ptr = ( ++ compute_causal_conv1d_metadata( ++ query_start_loc_p_cpu, ++ device=common_attn_metadata.query_start_loc.device, ++ ) + ) +- ) + + if self.vllm_config.cache_config.mamba_cache_mode == "all": + assert num_computed_tokens is not None +diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py +index 5a0c6f7..f2cc3da 100644 +--- a/vllm/v1/attention/backends/short_conv_attn.py ++++ b/vllm/v1/attention/backends/short_conv_attn.py +@@ -18,7 +18,6 @@ from vllm.v1.attention.backends.mamba_attn import ( + ) + from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, +- compute_causal_conv1d_metadata, + mamba_get_block_table_tensor, + ) + from vllm.v1.kv_cache_interface import MambaSpec +@@ -108,6 +107,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + reorder_batch_threshold: int = 1 + supports_update_block_table = False ++ needs_causal_conv1d_metadata = False + + def __init__( + self, +@@ -281,18 +281,10 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + spec_sequence_masks_cpu, device=query_start_loc.device + ) + +- # For causal_conv1d (non-spec prefill Triton kernel metadata). +- nums_dict = None +- batch_ptr = None +- token_chunk_offset_ptr = None + has_initial_states_p = None + has_initial_states_d = None + num_computed_tokens_p = None +- # Original request indices of the non-spec requests, ordered +- # [decodes, prefills]. Used to gather per-request data consistently. +- non_spec_req_idx_cpu: torch.Tensor | None = None + +- query_lens = torch.diff(query_start_loc) + # Per-request classification by mask, NOT by position. With + # spec-decode, the front decode group can contain both spec-decode + # requests and plain non-spec single-token decodes. Spec requests are +@@ -317,37 +309,32 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + else 0 + ) + +- # Original request indices grouped as +- # [spec | non-spec decode | non-spec prefill]; each group keeps the +- # original (already reordered) relative order via a stable nonzero. +- spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] +- decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] +- prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] +- non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) +- spec_req_idx = async_tensor_h2d(spec_req_idx_cpu, device=query_start_loc.device) ++ assert num_accepted_tokens is not None + non_spec_req_idx: torch.Tensor | None = None + + if num_decodes == 0 and num_prefills == 0: +- # Pure speculative-decode batch: all real tokens are spec tokens. +- spec_token_indx = torch.arange( +- num_spec_decode_tokens, +- dtype=torch.int32, +- device=query_start_loc.device, +- ) +- non_spec_token_indx = torch.empty( +- 0, dtype=torch.int32, device=query_start_loc.device +- ) +- spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] ++ # All real requests are speculative; zero-length padding is trailing. ++ spec_token_indx = None ++ non_spec_token_indx = None ++ spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, 0] ++ num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] + non_spec_state_indices_tensor = None + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None +- non_spec_query_start_loc_cpu = None + else: + # Mixed batch: build a per-token group key consistent with the + # request grouping above (spec=0 | decode=1 | prefill=2) and a + # stable sort, so tokens of each request stay contiguous and in + # request order. This yields spec tokens first, then the non-spec + # [decode, prefill] tokens. ++ query_lens = torch.diff(query_start_loc) ++ spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] ++ decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] ++ prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] ++ non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) ++ spec_req_idx = async_tensor_h2d( ++ spec_req_idx_cpu, device=query_start_loc.device ++ ) + non_spec_req_idx = async_tensor_h2d( + non_spec_req_idx_cpu, device=query_start_loc.device + ) +@@ -360,14 +347,17 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + dtype=torch.int64, + device=query_start_loc.device, + ) +- req_group[spec_req_idx] = 0 +- req_group[decode_req_idx] = 1 +- token_group = torch.repeat_interleave(req_group, query_lens) ++ req_group.index_fill_(0, spec_req_idx, 0) ++ req_group.index_fill_(0, decode_req_idx, 1) ++ token_group = torch.repeat_interleave( ++ req_group, query_lens, output_size=int(query_start_loc_cpu[-1]) ++ ) + token_perm = torch.argsort(token_group, stable=True) + spec_token_indx = token_perm[:num_spec_decode_tokens] + non_spec_token_indx = token_perm[num_spec_decode_tokens:] + + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] ++ num_accepted_tokens = num_accepted_tokens[spec_req_idx] + non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, +@@ -385,22 +375,8 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + dim=0, + out=non_spec_query_start_loc[1:], + ) +- non_spec_query_start_loc_cpu = torch.zeros( +- num_decodes + num_prefills + 1, dtype=torch.int32 +- ) +- torch.cumsum( +- query_lens_cpu[non_spec_req_idx_cpu], +- dim=0, +- out=non_spec_query_start_loc_cpu[1:], +- ) + +- assert num_accepted_tokens is not None +- # Accepted-token counts must follow the same request order as the +- # speculative state indices. +- num_accepted_tokens = num_accepted_tokens[spec_req_idx] +- +- # Compute the conv-state slots for the non-spec decode/prefill split, +- # plus the initial-state masks and Triton causal_conv1d metadata. ++ # Compute the conv-state slots for the non-spec decode/prefill split. + if non_spec_state_indices_tensor is None: + state_indices_tensor = block_table_tensor[:0, 0] + else: +@@ -431,20 +407,9 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + ] + has_initial_states_p = num_computed_tokens_p > 0 + assert non_spec_query_start_loc is not None +- assert non_spec_query_start_loc_cpu is not None + query_start_loc_p = ( + non_spec_query_start_loc[num_decodes:] - num_decode_tokens + ) +- query_start_loc_p_cpu = ( +- non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens +- ) +- if query_start_loc.device.type != "cpu": +- nums_dict, batch_ptr, token_chunk_offset_ptr = ( +- compute_causal_conv1d_metadata( +- query_start_loc_p_cpu, +- device=query_start_loc.device, +- ) +- ) + + # Prepare persistent tensors for CUDA graph capture and replay. + # ``m.num_actual_tokens`` is already padded by the model runner. +@@ -509,9 +474,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, +- nums_dict=nums_dict, +- batch_ptr=batch_ptr, +- token_chunk_offset_ptr=token_chunk_offset_ptr, + query_start_loc_p=query_start_loc_p, + query_start_loc_d=query_start_loc_d, + state_indices_tensor_p=state_indices_tensor_p, +-- +2.43.0 + diff --git a/patches/0066-cudagraph-allow-FULL-decode-graphs-for-one-token-pro.patch b/patches/0066-cudagraph-allow-FULL-decode-graphs-for-one-token-pro.patch new file mode 100644 index 0000000..e06bea1 --- /dev/null +++ b/patches/0066-cudagraph-allow-FULL-decode-graphs-for-one-token-pro.patch @@ -0,0 +1,454 @@ +From f0949d57a607a8882d0cba5a9937fa03bb82eefc Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 28 Sep 2026 12:48:52 -0700 +Subject: [PATCH 066/103] cudagraph: allow FULL decode graphs for one-token + prompt tails (port of upstream #58400) + +Port of the upstream vLLM change #58400 (Apache-2.0). A prompt tail of one token can run under the FULL decode graph. The +speculative-row rule of v0.30.0 (num_scheduled_tokens == drafts + 1) is kept, because the upstream rule depends +on a later change not in this base. +--- + tests/v1/spec_decode/test_dynamic_sd_cug.py | 12 ++-- + tests/v1/worker/test_gpu_batch_ordering.py | 69 +++++++++++++++++-- + tests/v1/worker/test_gpu_batch_shard.py | 1 + + .../worker/test_mamba_hybrid_model_state.py | 1 + + vllm/v1/worker/gpu/input_batch.py | 8 +++ + vllm/v1/worker/gpu/model_runner.py | 49 ++++++++++--- + .../worker/gpu/model_states/mamba_hybrid.py | 9 ++- + vllm/v1/worker/gpu/pcp_manager.py | 5 +- + .../spec_decode/autoregressive/speculator.py | 2 +- + .../multi_module_mtp/speculator.py | 2 +- + vllm/v1/worker/gpu/ubatch_utils.py | 5 ++ + vllm/v1/worker/utils.py | 6 +- + 12 files changed, 140 insertions(+), 29 deletions(-) + +diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py +index a5cbd11..1ca8486 100644 +--- a/tests/v1/spec_decode/test_dynamic_sd_cug.py ++++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py +@@ -128,7 +128,7 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch) + num_reqs, + num_tokens, + max_query_len, +- has_prefill=False, ++ decode_graph_eligible=True, + ) + + # The scheduler should mark every one of these shapes as a uniform +@@ -266,14 +266,14 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch): + # The batch shape is indistinguishable from a full batch of spec decodes. + assert ( + get_uniform_decode_token_count( +- num_reqs, num_tokens, decode_query_len, has_prefill=False ++ num_reqs, num_tokens, decode_query_len, decode_graph_eligible=True + ) + == decode_query_len + ) + + # Two of the requests are decode_query_len tokens into a longer prompt. + uniform_tok_count = get_uniform_decode_token_count( +- num_reqs, num_tokens, decode_query_len, has_prefill=True ++ num_reqs, num_tokens, decode_query_len, decode_graph_eligible=False + ) + assert uniform_tok_count is None + desc = manager.dispatch( +@@ -287,7 +287,7 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch): + + # The same shape with every request decoding still gets its FULL graph. + uniform_tok_count = get_uniform_decode_token_count( +- num_reqs, num_tokens, decode_query_len, has_prefill=False ++ num_reqs, num_tokens, decode_query_len, decode_graph_eligible=True + ) + assert uniform_tok_count == decode_query_len + desc = manager.dispatch( +@@ -341,7 +341,7 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch): + num_reqs, + num_tokens, + max_query_len, +- has_prefill=False, ++ decode_graph_eligible=True, + ) + assert uniform_tok_count == max_query_len + +@@ -407,7 +407,7 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch): + num_reqs, + num_tokens, + max_query_len, +- has_prefill=False, ++ decode_graph_eligible=True, + ) + assert uniform_tok_count == max_query_len + +diff --git a/tests/v1/worker/test_gpu_batch_ordering.py b/tests/v1/worker/test_gpu_batch_ordering.py +index 5037dcc..db9e080 100644 +--- a/tests/v1/worker/test_gpu_batch_ordering.py ++++ b/tests/v1/worker/test_gpu_batch_ordering.py +@@ -56,6 +56,7 @@ def _make_runner( + runner: Any = GPUModelRunner.__new__(GPUModelRunner) + runner.decode_query_len = decode_query_len + runner.adaptive_verification = None ++ runner.pcp_manager = None + runner.req_states = SimpleNamespace( + req_id_to_index={req_id: i for i, req_id in enumerate(req_states)}, + # The runner keeps this as min(num_computed_tokens, prefill_len). +@@ -128,6 +129,65 @@ def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode(): + assert _uniform_token_count({"p0": (0, 8)}, 8) is None + + ++def test_one_new_prompt_token_over_context_runs_as_decode(): ++ # A prompt tail padded with placeholder drafts and a one-token mid-prompt ++ # chunk both compute like decodes. A fresh one-token prompt and a real ++ # multi-token chunk don't. ++ batch = {"d0": (16, 16), "tail": (128, 129), "mid": (64, 129)} ++ batch.update({"fresh": (0, 1), "chunk": (120, 129)}) ++ runner = _make_runner(batch, decode_query_len=4) ++ scheduler_output = SimpleNamespace( ++ num_scheduled_tokens={"d0": 4, "tail": 4, "mid": 1, "fresh": 1, "chunk": 4}, ++ total_num_scheduled_tokens=14, ++ scheduled_spec_decode_tokens={"d0": [1, 2, 3], "tail": [-1] * 3}, ++ ) ++ state, _ = runner.gather_batch_req_state(scheduler_output, False) ++ assert state is not None ++ runs_as_decode = dict(zip(state.req_ids, state.prefill_runs_as_decode_np)) ++ assert runs_as_decode == { ++ "d0": False, ++ "tail": True, ++ "mid": True, ++ "fresh": False, ++ "chunk": False, ++ } ++ ++ ++def test_one_token_prompt_tail_with_prior_context_is_uniform_decode(): ++ # Replaying the last prompt token over existing context (full prefix-cache ++ # hit or P/D KV import) computes exactly like a decode. ++ batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)} ++ assert _uniform_token_count(batch, 1) == 1 ++ ++ # A fresh 1-token prompt alongside it has no prior state. ++ batch["p1"] = (0, 1) ++ assert _uniform_token_count(batch, 1) is None ++ ++ # Only one-token tails qualify, not K+1 chunks with prior context. ++ assert _uniform_token_count({"d0": (16, 16), "p0": (120, 129)}, 8) is None ++ ++ ++def test_padded_prompt_tail_with_prior_context_is_uniform_decode(): ++ # The scheduler pads a one-token prompt tail with K placeholder drafts to ++ # keep the K+1 spec-decode shape; it verifies like any other decode. ++ batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)} ++ runner = _make_runner(batch, decode_query_len=4) ++ scheduler_output = SimpleNamespace( ++ num_scheduled_tokens={req_id: 4 for req_id in batch}, ++ total_num_scheduled_tokens=12, ++ scheduled_spec_decode_tokens={req_id: [-1] * 3 for req_id in batch}, ++ ) ++ state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) ++ assert state.has_prefill and state.decode_graph_eligible ++ assert uniform_tok_count == 4 ++ ++ # Without the placeholder drafts, it is a real 4-token prompt chunk. ++ del scheduler_output.scheduled_spec_decode_tokens["p0"] ++ state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) ++ assert not state.decode_graph_eligible ++ assert uniform_tok_count is None ++ ++ + def test_dummy_batches_stay_uniform_decode(): + # Dummy batches have no request state to consult and must stay classified + # by shape alone; rejecting them would make capture and replay disagree. +@@ -185,6 +245,7 @@ def test_uniform_decode_uses_state_index_not_batch_position(): + runner: Any = GPUModelRunner.__new__(GPUModelRunner) + runner.decode_query_len = 8 + runner.adaptive_verification = None ++ runner.pcp_manager = None + # State arrays in state-index order: a prefilling request, then two decodes. + runner.req_states = SimpleNamespace( + req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2}, +@@ -214,11 +275,11 @@ def test_uniform_decode_uses_state_index_not_batch_position(): + + + def test_uniform_decode_predicate(): +- # Shape and prefill state must both pass. +- assert get_uniform_decode_token_count(2, 16, 8, False) == 8 +- assert get_uniform_decode_token_count(2, 16, 8, True) is None ++ # Shape and decode-graph eligibility must both pass. ++ assert get_uniform_decode_token_count(2, 16, 8, True) == 8 ++ assert get_uniform_decode_token_count(2, 16, 8, False) is None + # 12 tokens over 2 requests is no shared query length. +- assert get_uniform_decode_token_count(2, 12, 8, False) is None ++ assert get_uniform_decode_token_count(2, 12, 8, True) is None + + + def test_no_speculator_dispatches_on_query_length_alone(): +diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py +index 1309686..64c7187 100644 +--- a/tests/v1/worker/test_gpu_batch_shard.py ++++ b/tests/v1/worker/test_gpu_batch_shard.py +@@ -96,6 +96,7 @@ def _make_batch( + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), + has_prefill=False, ++ decode_graph_eligible=True, + input_ids=torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE), + positions=torch.arange(num_tokens, dtype=torch.int64, device=DEVICE), + is_padding=torch.zeros(num_tokens, dtype=torch.bool, device=DEVICE), +diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py +index 0cdd91e..caec567 100644 +--- a/tests/v1/worker/test_mamba_hybrid_model_state.py ++++ b/tests/v1/worker/test_mamba_hybrid_model_state.py +@@ -37,6 +37,7 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non + seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32), + seq_lens=torch.tensor([1537], dtype=torch.int32), + is_prefilling_np=torch.tensor([False]).numpy(), ++ prefill_runs_as_decode_np=None, + dcp_local_seq_lens=None, + positions=positions, + prompt_lens=torch.tensor([1024], dtype=torch.int32), +diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py +index d3c65bd..cdeb6e1 100644 +--- a/vllm/v1/worker/gpu/input_batch.py ++++ b/vllm/v1/worker/gpu/input_batch.py +@@ -86,6 +86,8 @@ class InputBatch: + is_prefilling_np: np.ndarray + # == np.any(is_prefilling_np) + has_prefill: bool ++ # No prefills, or only prefill rows that run as decodes. ++ decode_graph_eligible: bool + + # [num_tokens_after_padding] + input_ids: torch.Tensor +@@ -118,6 +120,11 @@ class InputBatch: + # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata). + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + ++ # [num_reqs] prefilling rows that schedule one new prompt token (excluding ++ # drafts) over existing context and so compute exactly like decodes. ++ # None if there are no prefills. ++ prefill_runs_as_decode_np: np.ndarray | None = None ++ + @classmethod + def make_dummy( + cls, +@@ -200,6 +207,7 @@ class InputBatch: + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), + has_prefill=False, ++ decode_graph_eligible=True, + input_ids=input_ids, + positions=positions, + is_padding=is_padding, +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 92ae922..246fe4c 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -1324,7 +1324,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + if dummy_run: + # Dummy batches are uniform by construction. + return None, get_uniform_decode_token_count( +- num_reqs, num_toks, max_query_len, has_prefill=False ++ num_reqs, num_toks, max_query_len, decode_graph_eligible=True + ) + + draft_tokens = scheduler_output.scheduled_spec_decode_tokens +@@ -1343,24 +1343,48 @@ class GPUModelRunner(LoRAModelRunnerMixin): + idx_mapping_np + ] + is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np ++ has_prefill = bool(is_prefilling_np.any()) + +- if self.adaptive_verification is not None and draft_tokens: +- num_toks = self.adaptive_verification.get_num_tokens( +- num_tokens_per_req, draft_tokens ++ num_draft_tokens_np = None ++ if draft_tokens: ++ num_drafts_iter = (len(draft_tokens.get(req_id, ())) for req_id in req_ids) ++ num_draft_tokens_np = np.fromiter( ++ num_drafts_iter, dtype=np.int32, count=num_reqs ++ ) ++ if self.adaptive_verification is not None: ++ num_toks = self.adaptive_verification.get_num_tokens( ++ num_tokens_per_req, draft_tokens ++ ) ++ ++ prefill_runs_as_decode_np = None ++ decode_graph_eligible = not has_prefill ++ if has_prefill and self.pcp_manager is None: ++ # One new prompt token over existing context (possibly padded with ++ # placeholder drafts) runs as a decode. PCP partitions prefill rows. ++ num_new_tokens = num_scheduled_tokens ++ if num_draft_tokens_np is not None: ++ num_new_tokens = num_new_tokens - num_draft_tokens_np ++ one_token_ext = (num_new_tokens == 1) & (num_computed_prefill_tokens_np > 0) ++ prefill_runs_as_decode_np = is_prefilling_np & one_token_ext ++ decode_graph_eligible = bool( ++ (prefill_runs_as_decode_np | ~is_prefilling_np).all() + ) + + batch_state = BatchReqState( + req_ids=req_ids, + num_scheduled_tokens=num_scheduled_tokens, + num_tokens=num_toks, ++ num_draft_tokens_np=num_draft_tokens_np, + idx_mapping_np=idx_mapping_np, + prefill_len_np=prefill_len_np, + num_computed_prefill_tokens_np=num_computed_prefill_tokens_np, + is_prefilling_np=is_prefilling_np, +- has_prefill=bool(is_prefilling_np.any()), ++ has_prefill=has_prefill, ++ prefill_runs_as_decode_np=prefill_runs_as_decode_np, ++ decode_graph_eligible=decode_graph_eligible, + ) + return batch_state, get_uniform_decode_token_count( +- num_reqs, num_toks, max_query_len, batch_state.has_prefill ++ num_reqs, num_toks, max_query_len, decode_graph_eligible + ) + + def prepare_inputs( +@@ -1401,11 +1425,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): + num_reqs, dtype=torch.int32, device=self.device + ) + else: +- num_draft_tokens_per_req = np.fromiter( +- (len(draft_tokens.get(req_id, ())) for req_id in req_ids), +- dtype=np.int32, +- count=num_reqs, +- ) ++ num_draft_tokens_per_req = batch_req_state.num_draft_tokens_np ++ if num_draft_tokens_per_req is None: ++ num_draft_tokens_per_req = np.zeros(num_reqs, dtype=np.int32) + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step + total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens +@@ -1543,6 +1565,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): + num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np, + is_prefilling_np=batch_req_state.is_prefilling_np, + has_prefill=batch_req_state.has_prefill, ++ prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np, ++ decode_graph_eligible=batch_req_state.decode_graph_eligible, + input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], + positions=self.input_buffers.positions[:num_tokens_after_padding], + is_padding=self.input_buffers.is_padding[:num_tokens_after_padding], +@@ -2597,11 +2621,14 @@ class BatchReqState(NamedTuple): + # May be less than scheduler_output.total_num_scheduled_tokens: + # adaptive verification trims the draft budget before running. + num_tokens: int ++ num_draft_tokens_np: np.ndarray | None # [num_reqs], None if no drafts + idx_mapping_np: np.ndarray # [num_reqs] + prefill_len_np: np.ndarray # [num_reqs] + num_computed_prefill_tokens_np: np.ndarray # [num_reqs] + is_prefilling_np: np.ndarray # [num_reqs] + has_prefill: bool ++ prefill_runs_as_decode_np: np.ndarray | None # [num_reqs] ++ decode_graph_eligible: bool + + + def sort_batch_req_ids( +diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +index 0a8b5d7..6b6e624 100644 +--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py ++++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +@@ -275,10 +275,13 @@ class MambaHybridModelState(DefaultModelState): + else: + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() + ++ is_prefilling_np = input_batch.is_prefilling_np ++ if input_batch.prefill_runs_as_decode_np is not None: ++ # A prompt tail the scheduler padded with placeholder drafts must run ++ # as a spec-decode row: the prefill kernels can't roll them back. ++ is_prefilling_np = is_prefilling_np & ~input_batch.prefill_runs_as_decode_np + is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") +- is_prefilling[: input_batch.num_reqs] = torch.from_numpy( +- input_batch.is_prefilling_np +- ) ++ is_prefilling[: input_batch.num_reqs] = torch.from_numpy(is_prefilling_np) + # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens + # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only + # compute them during actual (non-capture) forward execution. +diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py +index 376e0f2..ea9648d 100644 +--- a/vllm/v1/worker/gpu/pcp_manager.py ++++ b/vllm/v1/worker/gpu/pcp_manager.py +@@ -593,6 +593,7 @@ class PCPManager: + ) + local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_) + local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np ++ local_has_prefill = bool(local_is_prefilling_np.any()) + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32) + seq_lens_cpu_upper_bound_np[:num_local_reqs] = ( + local_start_pos_np + local_num_scheduled_tokens +@@ -643,7 +644,9 @@ class PCPManager: + prefill_len_np=local_prefill_len_np, + num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np, + is_prefilling_np=local_is_prefilling_np, +- has_prefill=bool(local_is_prefilling_np.any()), ++ has_prefill=local_has_prefill, ++ decode_graph_eligible=not local_has_prefill, ++ prefill_runs_as_decode_np=None, + input_ids=input_buffers.input_ids[:num_local_tokens_padded], + positions=input_buffers.positions[:num_local_tokens_padded], + is_padding=is_padding, +diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +index 3d243b6..52d389b 100644 +--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +@@ -289,7 +289,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): + # the target model during FULL cudagraph. + num_tokens, + max_query_len, +- input_batch.has_prefill, ++ input_batch.decode_graph_eligible, + ) + prefill_batch_desc, prefill_batch_sync = dispatch_cg_and_sync_dp( + self.prefill_cudagraph_manager, +diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +index 1c6ae31..e755a65 100644 +--- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +@@ -189,7 +189,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): + # When all requests are decoding (no true prefills), each has + # num_speculative_steps + 1 tokens, enabling FULL graph replay. + uniform_token_count = get_uniform_decode_token_count( +- num_reqs, num_tokens, max_query_len, input_batch.has_prefill ++ num_reqs, num_tokens, max_query_len, input_batch.decode_graph_eligible + ) + batch_desc, batch_sync = dispatch_cg_and_sync_dp( + self.cudagraph_manager, +diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py +index e3cade2..a28d1c2 100644 +--- a/vllm/v1/worker/gpu/ubatch_utils.py ++++ b/vllm/v1/worker/gpu/ubatch_utils.py +@@ -181,6 +181,11 @@ def _slice_input_batch( + req_start:req_stop + ], + is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop], ++ prefill_runs_as_decode_np=( ++ None ++ if input_batch.prefill_runs_as_decode_np is None ++ else input_batch.prefill_runs_as_decode_np[req_start:req_stop] ++ ), + input_ids=input_batch.input_ids[tok_start:tok_stop], + positions=input_batch.positions[tok_start:tok_stop], + is_padding=input_batch.is_padding[tok_start:tok_stop], +diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py +index 71c1a2d..61698bf 100644 +--- a/vllm/v1/worker/utils.py ++++ b/vllm/v1/worker/utils.py +@@ -764,10 +764,12 @@ def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> + + + def get_uniform_decode_token_count( +- num_reqs: int, num_tokens: int, max_query_len: int, has_prefill: bool ++ num_reqs: int, num_tokens: int, max_query_len: int, decode_graph_eligible: bool + ) -> int | None: + """Per-request token count of a uniform decode batch, or None.""" +- if not has_prefill and is_uniform_query_len(num_reqs, num_tokens, max_query_len): ++ if decode_graph_eligible and is_uniform_query_len( ++ num_reqs, num_tokens, max_query_len ++ ): + return max_query_len + return None + +-- +2.43.0 + diff --git a/patches/0067-spec-decode-reject-draft-slots-that-were-never-propo.patch b/patches/0067-spec-decode-reject-draft-slots-that-were-never-propo.patch new file mode 100644 index 0000000..8f13e22 --- /dev/null +++ b/patches/0067-spec-decode-reject-draft-slots-that-were-never-propo.patch @@ -0,0 +1,218 @@ +From ec0e5762517c1d0cd6a4d2e7a5303ce028d5752a Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 28 Sep 2026 13:07:15 -0700 +Subject: [PATCH 067/103] spec-decode: reject draft slots that were never + proposed (port of upstream #58784) + +Port of the upstream vLLM change #58784 (Apache-2.0). Draft slots the drafter never filled are rejected instead of being verified. +--- + .../test_rejection_sampler_utils.py | 80 +++++++++++++++++++ + .../gpu/spec_decode/rejection_sampler.py | 80 ++++++++++++++++++- + 2 files changed, 158 insertions(+), 2 deletions(-) + +diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py +index a64b0eb..02aa542 100644 +--- a/tests/v1/spec_decode/test_rejection_sampler_utils.py ++++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + import math ++from types import SimpleNamespace + + import pytest + import torch +@@ -10,6 +11,10 @@ from vllm.v1.watermarking import GumbelWatermarker + from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample + from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample ++from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( ++ RejectionSampler, ++ gather_draft_sampled, ++) + from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + rejection_sample, + ) +@@ -675,6 +680,81 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool): + ) + + ++def test_verify_rejects_unproposed_drafts(): ++ """Zero drafts padded onto a step starting in the prefill must be rejected.""" ++ torch.manual_seed(0) ++ device = "cuda" ++ num_trials = 256 ++ K = 8 ++ ++ # Row 0: token 0 is very unlikely. Rows >= 1 follow the zero inputs and ++ # predict token 0. The never-written draft logits row is uniform. ++ target_logits_1d = torch.randn(VOCAB_SIZE, device=device) ++ target_logits_1d[0] = -20.0 ++ inputs = _build_rejection_sample_inputs( ++ target_logits_1d, ++ torch.zeros(VOCAB_SIZE, device=device), ++ K, ++ temperature=1.0, ++ num_trials=num_trials, ++ ) ++ target_rest = torch.randn(VOCAB_SIZE, device=device) ++ target_rest[0] = 20.0 ++ inputs["target_logits"].view(num_trials, K + 1, VOCAB_SIZE)[:, 1:] = target_rest ++ inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = 0 ++ ++ # Each request's first logit is the last prefill token, as in the first ++ # decode step after a P/D remote KV load. ++ first_pos = inputs["pos"].view(num_trials, K + 1)[:, 0] ++ prefill_len = first_pos + 1 ++ rejection_sampler = object.__new__(RejectionSampler) ++ rejection_sampler.sampler = SimpleNamespace( ++ apply_sampling_params=lambda logits, *args: logits, ++ sampling_states=SimpleNamespace( ++ temperature=SimpleNamespace(gpu=inputs["temperature"]), ++ seeds=SimpleNamespace(gpu=inputs["seed"]), ++ ), ++ use_fp64_gumbel=False, ++ ) ++ rejection_sampler.num_speculative_steps = K ++ rejection_sampler.synthetic_conditional_rates = None ++ rejection_sampler.use_block_verification = True ++ rejection_sampler.watermark_key = None ++ ++ def verify() -> tuple[torch.Tensor, torch.Tensor]: ++ draft_sampled, pos = gather_draft_sampled( ++ inputs["draft_sampled"], ++ inputs["pos"], ++ torch.arange(num_trials * (K + 1), device=device), ++ inputs["expanded_idx_mapping"], ++ inputs["expanded_local_pos"], ++ prefill_len, ++ ) ++ _, sampled, num_sampled = rejection_sampler._verify( ++ inputs["target_logits"], ++ inputs["draft_logits"], ++ draft_sampled, ++ pos, ++ inputs["cu_num_logits"], ++ inputs["idx_mapping"], ++ None, ++ inputs["expanded_idx_mapping"], ++ inputs["expanded_local_pos"], ++ None, ++ ) ++ return sampled, num_sampled ++ ++ sampled, num_sampled = verify() ++ assert (num_sampled == 1).all() ++ # Resampled from the target at row 0, not from a residual. ++ assert (sampled[:, 0] != 0).all() ++ ++ # Once past the prefill, block verification accepts the same zero drafts. ++ prefill_len.copy_(first_pos) ++ _, num_sampled = verify() ++ assert (num_sampled == K + 1).all() ++ ++ + @pytest.mark.parametrize( + "num_speculative_steps,temperature", + [ +diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +index 6c07e6a..af7be08 100644 +--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py ++++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +@@ -74,6 +74,76 @@ def _flatten_sampled_kernel( + tl.store(flat_sampled_ptr + start_idx + i, token_id) + + ++@triton.jit ++def _gather_draft_sampled_kernel( ++ # [num_logits] ++ draft_sampled_ptr, ++ # [num_logits] ++ pos_ptr, ++ # [num_tokens] ++ input_ids_ptr, ++ # [num_tokens] ++ positions_ptr, ++ # [num_logits] ++ logits_indices_ptr, ++ # [num_logits] ++ expanded_idx_mapping_ptr, ++ # [num_logits] ++ expanded_local_pos_ptr, ++ # [max_num_reqs] ++ prefill_len_ptr, ++ num_logits, ++ BLOCK_SIZE: tl.constexpr, ++): ++ block = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) ++ mask = block < num_logits ++ token_idx = tl.load(logits_indices_ptr + block, mask=mask) ++ token_id = tl.load(input_ids_ptr + token_idx, mask=mask) ++ pos = tl.load(positions_ptr + token_idx, mask=mask) ++ req_state_idx = tl.load(expanded_idx_mapping_ptr + block, mask=mask) ++ local_pos = tl.load(expanded_local_pos_ptr + block, mask=mask) ++ prefill_len = tl.load(prefill_len_ptr + req_state_idx, mask=mask) ++ # Drafts are only proposed after sampling past the prefill, so any draft ++ # slots in a step starting within the prefill are placeholders (e.g. the ++ # padded first step after a P/D remote KV load). ++ is_placeholder = (local_pos > 0) & (pos - local_pos < prefill_len) ++ token_id = tl.where(is_placeholder, -1, token_id) ++ tl.store(draft_sampled_ptr + block, token_id, mask=mask) ++ tl.store(pos_ptr + block, pos, mask=mask) ++ ++ ++def gather_draft_sampled( ++ input_ids: torch.Tensor, ++ positions: torch.Tensor, ++ logits_indices: torch.Tensor, ++ expanded_idx_mapping: torch.Tensor, ++ expanded_local_pos: torch.Tensor, ++ prefill_len: torch.Tensor, ++) -> tuple[torch.Tensor, torch.Tensor]: ++ """Gather the input token and position of each logits row. ++ ++ Draft rows of requests that have not yet sampled past their prefill are ++ set to -1 so that the rejection kernels reject them. ++ """ ++ num_logits = logits_indices.shape[0] ++ draft_sampled = input_ids.new_empty(num_logits) ++ pos = positions.new_empty(num_logits) ++ BLOCK_SIZE = 1024 ++ _gather_draft_sampled_kernel[(triton.cdiv(num_logits, BLOCK_SIZE),)]( ++ draft_sampled, ++ pos, ++ input_ids, ++ positions, ++ logits_indices, ++ expanded_idx_mapping, ++ expanded_local_pos, ++ prefill_len, ++ num_logits, ++ BLOCK_SIZE=BLOCK_SIZE, ++ ) ++ return draft_sampled, pos ++ ++ + class RejectionSampler: + def __init__( + self, +@@ -300,8 +370,14 @@ class RejectionSampler: + # that num_nans is computed before applying penalties and temperature. + num_nans = get_num_nans(logits) if self.sampler.compute_nans else None + +- draft_sampled = input_batch.input_ids[input_batch.logits_indices] +- pos = input_batch.positions[input_batch.logits_indices] ++ draft_sampled, pos = gather_draft_sampled( ++ input_batch.input_ids, ++ input_batch.positions, ++ input_batch.logits_indices, ++ input_batch.expanded_idx_mapping, ++ input_batch.expanded_local_pos, ++ self.sampler.req_states.prefill_len.gpu, ++ ) + + # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the + # accepted length before it reaches a permissive bitmask row. Only +-- +2.43.0 + diff --git a/patches/0068-guard-the-bounds-guard-telemetry-code-default-is-war.patch b/patches/0068-guard-the-bounds-guard-telemetry-code-default-is-war.patch new file mode 100644 index 0000000..3c886d4 --- /dev/null +++ b/patches/0068-guard-the-bounds-guard-telemetry-code-default-is-war.patch @@ -0,0 +1,226 @@ +From 0f95d8a93025684b4a36b6e57ccd908b589e812e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 05:05:23 -0400 +Subject: [PATCH 068/103] guard: the bounds-guard telemetry code default is + warn + +VLLM_E2_GUARD_MODE now defaults to warn in code, matching what the shipped launchers set; v2 through v2.2.0 +defaulted to off. The comment above it, the MODES docstring, the vllm/envs.py default and its annotation follow. +tests/v1/worker/gpu/test_guard_default_mode.py checks that code, comment and envs.py agree, from the sources, so +it runs without CUDA, and fails on the earlier source. +--- + .../v1/worker/gpu/test_guard_default_mode.py | 86 +++++++++++++++++++ + .../test_guard_telemetry_defaults_drain.py | 10 +-- + vllm/envs.py | 4 +- + vllm/v1/worker/guard_telemetry.py | 22 ++--- + 4 files changed, 105 insertions(+), 17 deletions(-) + create mode 100644 tests/v1/worker/gpu/test_guard_default_mode.py + +diff --git a/tests/v1/worker/gpu/test_guard_default_mode.py b/tests/v1/worker/gpu/test_guard_default_mode.py +new file mode 100644 +index 0000000..4dc4516 +--- /dev/null ++++ b/tests/v1/worker/gpu/test_guard_default_mode.py +@@ -0,0 +1,86 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The guard's CODE default is warn (since v2.5), and the code, its comment and ++vllm/envs.py agree. ++ ++Source-level pins: they parse the files instead of importing them, so they run on a machine with no ++CUDA, no zmq and no working vllm import. The behavioural check at the end imports the ++module only when that works. ++""" ++ ++import ast ++import importlib ++import os ++import re ++import sys ++ ++import pytest ++ ++_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), *([".."] * 4))) ++_GT = os.path.join(_ROOT, "vllm", "v1", "worker", "guard_telemetry.py") ++_ENVS = os.path.join(_ROOT, "vllm", "envs.py") ++ ++ ++def _read(path): ++ with open(path, encoding="utf-8") as fh: ++ return fh.read() ++ ++ ++def _getenv_defaults(src, name): ++ """every default literal passed to os.environ.get / os.getenv for `name`""" ++ out = [] ++ for node in ast.walk(ast.parse(src)): ++ if not isinstance(node, ast.Call) or len(node.args) < 2: ++ continue ++ fn = ast.unparse(node.func) ++ if fn not in ("os.environ.get", "os.getenv"): ++ continue ++ a0, a1 = node.args[0], node.args[1] ++ if isinstance(a0, ast.Constant) and a0.value == name and isinstance(a1, ast.Constant): ++ out.append(a1.value) ++ return out ++ ++ ++def test_code_default_is_warn(): ++ d = _getenv_defaults(_read(_GT), "VLLM_E2_GUARD_MODE") ++ assert d == ["warn"], d ++ ++ ++def test_envs_default_is_warn_everywhere(): ++ src = _read(_ENVS) ++ assert _getenv_defaults(src, "VLLM_E2_GUARD_MODE") == ["warn"] ++ assert re.search(r'VLLM_E2_GUARD_MODE: str = "warn"', src) ++ assert 'VLLM_E2_GUARD_MODE: str = "off"' not in src ++ ++ ++def test_comment_and_docstring_agree_with_the_default(): ++ src = _read(_GT) ++ line = next(ln for ln in src.splitlines() if ln.startswith("_MODE = os.environ.get(")) ++ i = src.splitlines().index(line) ++ above = "\n".join(src.splitlines()[max(0, i - 8) : i]) ++ assert "Release defaults" in above and "warn, 0.05 s" in above ++ assert "v2.5" in above ++ # the MODES block must name warn as the default and must not call off the default ++ assert re.search(r"warn\s+\(DEFAULT since v2\.5", src) ++ assert not re.search(r"off\s+\(DEFAULT", src) ++ # the interval default stays 0.05 (the comment says so) ++ assert _getenv_defaults(src, "VLLM_E2_GUARD_INTERVAL") == ["0.05"] ++ ++ ++def test_failing_control_the_old_default_would_be_caught(): ++ """the pin must fail on the earlier source (default off)""" ++ old = _read(_GT).replace('"VLLM_E2_GUARD_MODE", "warn"', '"VLLM_E2_GUARD_MODE", "off"') ++ assert old != _read(_GT) ++ assert _getenv_defaults(old, "VLLM_E2_GUARD_MODE") == ["off"] ++ assert _getenv_defaults(old, "VLLM_E2_GUARD_MODE") != ["warn"] ++ ++ ++def test_module_default_with_env_absent(monkeypatch): ++ monkeypatch.delenv("VLLM_E2_GUARD_MODE", raising=False) ++ monkeypatch.delenv("VLLM_E2_GUARD_INTERVAL", raising=False) ++ sys.modules.pop("vllm.v1.worker.guard_telemetry", None) ++ try: ++ gt = importlib.import_module("vllm.v1.worker.guard_telemetry") ++ except Exception as e: # noqa: BLE001 (no CUDA / zmq on this host: the source pins above stand in) ++ pytest.skip(f"vllm.v1.worker.guard_telemetry not importable here: {e!r}") ++ assert gt._MODE == "warn" +diff --git a/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +index 606292e..32a99c1 100644 +--- a/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py ++++ b/tests/v1/worker/gpu/test_guard_telemetry_defaults_drain.py +@@ -2,8 +2,8 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Bounds-guard instrumentation: release defaults, reader loss and the final drain. + +-1. Release defaults are off / 0.05 with both variables absent (earlier +- builds defaulted to warn), and a ++1. Release defaults are warn / 0.05 with both variables absent (v2 through ++ v2.2.0 defaulted to off), and a + non-finite or non-positive interval is rejected at import. + 2. A reader that loses a device arms the stop (the fatal-site contract is + otherwise silently void). +@@ -32,7 +32,7 @@ def _fresh(monkeypatch, **env): + + def test_release_defaults_with_env_absent(monkeypatch): + gt = _fresh(monkeypatch) +- assert gt._MODE == "off" # v2: default off ++ assert gt._MODE == "warn" # code default warn since v2.5 + assert pytest.approx(0.05) == gt._INTERVAL + + +@@ -45,7 +45,7 @@ def test_bad_interval_rejected_at_import(monkeypatch, bad): + def test_reader_loss_arms_the_stop(monkeypatch): + gt = _fresh( + monkeypatch, VLLM_E2_GUARD_MODE="warn" +- ) # the reader exists only when opted in (default off) ++ ) # explicit, so the test does not depend on the default + device = torch.device("cuda:0") + gt.prealloc(device) + gt.reset_for_test() +@@ -65,7 +65,7 @@ def test_reader_loss_arms_the_stop(monkeypatch): + def test_drain_reports_last_step_hit_and_raises_for_fatal_site(monkeypatch): + gt = _fresh( + monkeypatch, VLLM_E2_GUARD_MODE="warn" +- ) # warn mode (opt-in): a plain site must NOT raise ++ ) # warn mode: a plain site must NOT raise + device = torch.device("cuda:0") + gt.prealloc(device) + gt.reset_for_test() +diff --git a/vllm/envs.py b/vllm/envs.py +index f13f442..4356854 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -399,7 +399,7 @@ if TYPE_CHECKING: + # Guard telemetry: mode raise|warn|off and the reader poll interval in + # seconds. counter() returns a real tensor in every mode, so the Triton + # arguments and the captured graph do not depend on either. +- VLLM_E2_GUARD_MODE: str = "off" ++ VLLM_E2_GUARD_MODE: str = "warn" + VLLM_E2_GUARD_INTERVAL: str = "0.05" + # Debug allocator telemetry (poll seconds + output directory) and the + # scheduler allocation log path. +@@ -2426,7 +2426,7 @@ environment_variables: dict[str, Callable[[], Any]] = { + ).strip(), + "VLLM_QSA_KV_SCALES": lambda: os.getenv("VLLM_QSA_KV_SCALES", "").strip(), + "VLLM_E2_GUARD_MODE": lambda: ( +- os.getenv("VLLM_E2_GUARD_MODE", "off").strip().lower() ++ os.getenv("VLLM_E2_GUARD_MODE", "warn").strip().lower() + ), + "VLLM_E2_GUARD_INTERVAL": lambda: os.getenv("VLLM_E2_GUARD_INTERVAL", "0.05"), + "VLLM_E2_MEMSTATS": lambda: os.getenv("VLLM_E2_MEMSTATS", "0"), +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index 63c450f..b02a7c3 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -40,13 +40,13 @@ MODES -- VLLM_E2_GUARD_MODE: + serving product than a loud stop. Read check_pending below before + using it -- the raise is late and it is local to one rank. + Targeted repro boots. +- warn log only -- EXCEPT for the sites in _ALWAYS_FATAL_SITES, which arm +- the raise whatever the mode says, because their masked value has +- already been written somewhere durable. USE THIS FOR LONG STRESS +- BOOTS. +- off (DEFAULT since v2: the reader thread +- and the QSA clip reader cost ~0.3 ms per decode step on TP2xPP2; opt in +- with VLLM_E2_GUARD_MODE=warn for diagnosis) ++ warn (DEFAULT since v2.5) log only -- EXCEPT for the sites in ++ _ALWAYS_FATAL_SITES, which arm the raise whatever the mode says, ++ because their masked value has already been written somewhere durable. ++ USE THIS FOR LONG STRESS BOOTS. ++ off (the default in v2 through v2.2.0: the reader thread and the QSA clip ++ reader cost ~0.3 ms per decode step on TP2xPP2; set ++ VLLM_E2_GUARD_MODE=off to skip the reader) + no reader thread. The kernels still mask and still count; nothing is + read -- INCLUDING the sites in _ALWAYS_FATAL_SITES. Use only when you + have decided you do not want the evidence. +@@ -218,10 +218,12 @@ SITE_NAMES = { + } + + +-# Release defaults: off, 0.05 s. ++# Release defaults: warn, 0.05 s. + # Earlier revisions shipped raise / 0.5 while every qualification boot set both + # variables explicitly, so the defaults were never the thing being tested. +-_MODE = os.environ.get("VLLM_E2_GUARD_MODE", "off").strip().lower() ++# v2 through v2.2.0 shipped the code default off; v2.5 puts ++# it back to warn so the code default, this comment and the shipped launchers agree. ++_MODE = os.environ.get("VLLM_E2_GUARD_MODE", "warn").strip().lower() + if _MODE not in ("raise", "warn", "off"): + # fail closed: a set-but-broken gate must not silently disarm the instrument + raise ValueError(f"VLLM_E2_GUARD_MODE must be raise|warn|off, got {_MODE!r}") +@@ -491,7 +493,7 @@ def start_reader() -> None: + host, and a device-to-host copy from a second thread while a capture is + active is exactly the call class cudaStreamCaptureModeGlobal -- the default + for torch.cuda.graph -- rejects. The QSA clip-counter reader survives it in +- practice, but that one is maintainer-gated and this one can be enabled by any user, so ++ practice, but that one is opt-in and this one is on by default, so + the precedent is evidence, not proof. Starting after capture deletes the + question rather than arguing it. + +-- +2.43.0 + diff --git a/patches/0069-qsa-start-the-clip-counter-reader-after-capture-not-.patch b/patches/0069-qsa-start-the-clip-counter-reader-after-capture-not-.patch new file mode 100644 index 0000000..0c277b0 --- /dev/null +++ b/patches/0069-qsa-start-the-clip-counter-reader-after-capture-not-.patch @@ -0,0 +1,324 @@ +From 22b00ba2e2239b52923cb555ae6e24198bb59835 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 05:08:33 -0400 +Subject: [PATCH 069/103] qsa: start the clip-counter reader after capture, not + on the first forward + +The FP8 KV clip counter's reader thread used to start on the first eager forward, before CUDA graph capture, so +its device-to-host copies could overlap a capture. _qsa_clip_count now only allocates the per-layer counter; the +reader is started by a hook that guard_telemetry runs after warmup and capture (or at once if capture is already +done). Counting is unchanged. Tests cover the order, idempotence and a failing control that restores the old start. +--- + .../test_qsa_clip_reader_lifecycle.py | 170 ++++++++++++++++++ + vllm/models/qwen4_exp/nvidia/qsa.py | 48 +++-- + vllm/v1/worker/guard_telemetry.py | 30 ++++ + 3 files changed, 234 insertions(+), 14 deletions(-) + create mode 100644 tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py + +diff --git a/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py +new file mode 100644 +index 0000000..a81feb7 +--- /dev/null ++++ b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py +@@ -0,0 +1,170 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The QSA clip-counter READER starts after warmup + cudagraph capture, not on the first forward (since v2.5). ++ ++Before, `_qsa_clip_count` started the `qsa-clip-reader` thread on the first eager forward -- before capture -- so its ++device-to-host copies could run on a second thread while a capture was active (the call class ++cudaStreamCaptureModeGlobal rejects). That is why `COUNTERS` shipped off. Now the starter is registered with ++guard_telemetry.register_post_capture_hook and run by guard_telemetry.start_reader() (gpu_worker, after capture). ++ ++qsa.py cannot be imported on a machine without zmq/CUDA, so the clip-counter block is cut out of the source ++text, written to a temporary module (with a small prelude) and imported next to the real guard_telemetry; the reader ++loop itself is replaced by a stub that records "I ran". The failing control re-inserts the OLD behaviour into the same ++block and must flip. ++""" ++ ++import ast ++import importlib.util ++import itertools ++import os ++import threading ++import time ++ ++import pytest ++import torch ++ ++_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), *([".."] * 3))) ++_QSA = os.path.join(_ROOT, "vllm", "models", "qwen4_exp", "nvidia", "qsa.py") ++_WORKER = os.path.join(_ROOT, "vllm", "v1", "worker", "gpu_worker.py") ++_GT_PATH = os.path.join(_ROOT, "vllm", "v1", "worker", "guard_telemetry.py") ++_PRELUDE = ( ++ "import logging, math, os, threading, time\n" ++ "import torch\n" ++ "logger = logging.getLogger('qsa-clip-test')\n" ++ "Qwen4ExpQSAAttention = object\n" ++) ++_seq = itertools.count() ++ ++ ++def _src(path): ++ with open(path, encoding="utf-8") as fh: ++ return fh.read() ++ ++ ++def _block(src): ++ a = src.index('_QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT"') ++ b = src.index("_QSA_FP8_CACHE_DTYPES = (") ++ return src[a:b] ++ ++ ++@pytest.fixture ++def gt(monkeypatch): ++ pytest.importorskip("vllm.v1.worker.guard_telemetry", reason="guard_telemetry needs torch + a vllm import") ++ from vllm.v1.worker import guard_telemetry as g ++ ++ monkeypatch.setattr(g, "_post_capture_hooks", []) ++ monkeypatch.setattr(g, "_post_capture_done", False) ++ monkeypatch.setattr(g, "_MODE", "off") # no guard reader thread in this test ++ monkeypatch.setattr(g, "_reader_started", False) ++ return g ++ ++ ++class _Layer: ++ layer_name = "layers.0.attn" ++ _k_scale = torch.tensor(1.0 / 448.0) # ceiling = 1.0 ++ _v_scale = torch.tensor(1.0 / 448.0) ++ ++ ++def _load_block(monkeypatch, tmp_path, env, transform=None): ++ if env is None: ++ monkeypatch.delenv("VLLM_QSA_KV_CLIP_COUNT", raising=False) ++ else: ++ monkeypatch.setenv("VLLM_QSA_KV_CLIP_COUNT", env) ++ text = _block(_src(_QSA)) ++ if transform: ++ text = transform(text) ++ path = tmp_path / ("qsa_clip_block_%d.py" % next(_seq)) ++ path.write_text(_PRELUDE + text, encoding="utf-8") ++ spec = importlib.util.spec_from_file_location(path.stem, str(path)) ++ mod = importlib.util.module_from_spec(spec) ++ spec.loader.exec_module(mod) ++ ran = threading.Event() ++ mod._qsa_clip_reader = ran.set # the loop itself is not under test ++ return mod, ran ++ ++ ++def _forward(mod): ++ key = torch.tensor([[[2.0, 0.5]], [[0.1, 3.0]], [[5.0, 5.0]]]) # 3 tokens ++ value = torch.zeros_like(key) ++ slot = torch.tensor([0, 1, -1]) # the last row is padding: never counted ++ mod._qsa_clip_count(_Layer(), key, value, slot) ++ ++ ++def test_first_forward_does_not_start_the_reader(monkeypatch, tmp_path, gt): ++ mod, ran = _load_block(monkeypatch, tmp_path, "5") ++ _forward(mod) ++ time.sleep(0.05) ++ assert not ran.is_set() ++ assert mod._qsa_clip_reader_started is False ++ assert "_qsa_clip_start_reader" in [f.__name__ for f in gt._post_capture_hooks] ++ assert set(mod._qsa_clip_counters) == {"layers.0.attn"} ++ ++ ++def test_reader_starts_once_after_capture(monkeypatch, tmp_path, gt): ++ mod, ran = _load_block(monkeypatch, tmp_path, "5") ++ _forward(mod) ++ gt.start_reader() # what gpu_worker calls after capture ++ assert ran.wait(2.0) ++ assert mod._qsa_clip_reader_started is True ++ gt.start_reader() # idempotent ++ assert len([x for x in threading.enumerate() if x.name == "qsa-clip-reader"]) <= 1 ++ ++ ++def test_counting_is_unchanged(monkeypatch, tmp_path, gt): ++ mod, _ = _load_block(monkeypatch, tmp_path, "5") ++ _forward(mod) ++ # ceiling 1.0: row0 has one element over (2.0), row1 has one (3.0), row2 is padding (slot -1), not counted ++ assert mod._qsa_clip_counters["layers.0.attn"].tolist() == [2, 0] ++ ++ ++def test_no_hook_when_the_env_is_unset(monkeypatch, tmp_path, gt): ++ mod, _ = _load_block(monkeypatch, tmp_path, None) ++ assert mod._qsa_clip_interval == 0.0 ++ assert gt._post_capture_hooks == [] ++ ++ ++def test_late_registration_starts_at_once(monkeypatch, tmp_path, gt): ++ gt.run_post_capture_hooks() # capture already finished ++ mod, _ = _load_block(monkeypatch, tmp_path, "5") # qsa imported late: the starter runs during the import ++ assert mod._qsa_clip_reader_started is True # (the thread was the real reader, asleep for its 5 s interval) ++ ++ ++def test_a_failing_starter_does_not_stop_the_others(gt): ++ seen = [] ++ ++ def boom(): ++ raise RuntimeError("boom") ++ ++ gt.register_post_capture_hook(boom) ++ gt.register_post_capture_hook(lambda: seen.append(1)) ++ gt.start_reader() ++ assert seen == [1] ++ ++ ++def test_failing_control_the_old_path_started_the_reader_before_capture(monkeypatch, tmp_path, gt): ++ """re-insert the earlier behaviour (reader started inside the first forward): the observation must flip""" ++ ++ def old_path(text): ++ marker = " _qsa_clip_counters[layer.layer_name] = ctr" ++ assert marker in text ++ return text.replace(marker, marker + "\n _qsa_clip_start_reader()", 1) ++ ++ mod, ran = _load_block(monkeypatch, tmp_path, "5", transform=old_path) ++ _forward(mod) # first (pre-capture) forward, no start_reader() call ++ assert ran.wait(2.0), "control did not flip: the old path should start the reader on the first forward" ++ assert mod._qsa_clip_reader_started is True ++ ++ ++def test_source_pins(): ++ qsa = _src(_QSA) ++ tree = ast.parse(qsa) ++ fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "_qsa_clip_count") ++ assert "threading.Thread" not in ast.unparse(fn), "the clip reader must not be started from the forward" ++ assert "register_post_capture_hook(_qsa_clip_start_reader)" in qsa ++ gtree = ast.parse(_src(_GT_PATH)) ++ sr = next(n for n in ast.walk(gtree) if isinstance(n, ast.FunctionDef) and n.name == "start_reader") ++ stmts = [s for s in sr.body if not (isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant))] ++ assert isinstance(stmts[0], ast.Global) ++ assert ast.unparse(stmts[1]) == "run_post_capture_hooks()", "hooks must run before the guard-mode early return" ++ w = _src(_WORKER) # the worker calls start_reader after capture_model ++ assert w.index("self.model_runner.capture_model(") < w.index("guard_telemetry.start_reader()") +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 73dd07a..467f7ad 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -150,6 +150,13 @@ def _qsa_collect_absmax( + # graph replay keeps counting. A daemon reader thread .item()s the counters on + # a cadence and logs: WARNING on any new clips, INFO heartbeat every 10th pass + # so a quiet log still proves the instrument is alive. ++# ++# Since v2.5 the reader thread is started AFTER warmup and cudagraph capture ++# (guard_telemetry.register_post_capture_hook, run by start_reader() from ++# gpu_worker), never from the first eager forward. Before, it started inside the ++# first forward, i.e. before capture, and its device-to-host copies could run on ++# a second thread while a capture was active -- the call class ++# cudaStreamCaptureModeGlobal rejects. That is why COUNTERS shipped off. + _QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT" + _qsa_clip_raw = os.environ.get(_QSA_CLIP_ENV, "").strip() + try: +@@ -222,31 +229,44 @@ def _qsa_clip_reader() -> None: + logger.exception("QSA clip reader pass failed; thread continues") + + ++def _qsa_clip_start_reader() -> None: ++ """Start the reader thread. Runs after warmup + capture (post-capture hook); idempotent.""" ++ global _qsa_clip_reader_started ++ with _qsa_clip_lock: ++ if _qsa_clip_reader_started: ++ return ++ _qsa_clip_reader_started = True ++ threading.Thread( ++ target=_qsa_clip_reader, ++ name="qsa-clip-reader", ++ daemon=True, ++ ).start() ++ logger.info( ++ "QSA clip counter reader started after capture (read interval %.0fs, %d layer(s) registered)", ++ _qsa_clip_interval, ++ len(_qsa_clip_counters), ++ ) ++ ++ ++if _qsa_clip_interval > 0.0: ++ from vllm.v1.worker.guard_telemetry import register_post_capture_hook # noqa: E402 ++ ++ register_post_capture_hook(_qsa_clip_start_reader) ++ ++ + def _qsa_clip_count( + layer: Qwen4ExpQSAAttention, + key: torch.Tensor, + value: torch.Tensor, + slot_mapping: torch.Tensor | None = None, + ) -> None: +- global _qsa_clip_reader_started + ctr = _qsa_clip_counters.get(layer.layer_name) + if ctr is None: + # First forward is eager (profiling/warmup precede capture), so the +- # allocation never happens inside cudagraph capture. ++ # allocation never happens inside cudagraph capture. The READER is not ++ # started here: _qsa_clip_start_reader runs after capture. + ctr = torch.zeros(2, dtype=torch.int64, device=key.device) + _qsa_clip_counters[layer.layer_name] = ctr # registered with the reader by the reader thread (never during capture) +- with _qsa_clip_lock: +- if not _qsa_clip_reader_started: +- _qsa_clip_reader_started = True +- threading.Thread( +- target=_qsa_clip_reader, +- name="qsa-clip-reader", +- daemon=True, +- ).start() +- logger.info( +- "QSA clip counter armed (read interval %.0fs)", +- _qsa_clip_interval, +- ) + # Thresholds are 1-D so the comparison promotes to float32; a 0-dim float32 + # tensor promotes like a Python scalar and the threshold would be rounded to + # the BF16 grid of K/V, undercounting values just above the calibrated ceiling. +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index b02a7c3..19cf966 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -486,9 +486,38 @@ def prealloc(device: torch.device) -> None: + _nb_reader.register(device, _buffers[device]) + + ++# Readers that other modules own but that must obey the same rule: nothing reads a device tensor from a second ++# thread while a cudagraph capture is active. A module registers its starter here at import; start_reader() runs ++# them after capture, and one registered after that point (a late import) starts at once. Before v2.5 the QSA ++# clip-counter reader started on the first eager forward, i.e. BEFORE capture. ++_post_capture_hooks: list = [] ++_post_capture_done = False ++ ++ ++def register_post_capture_hook(fn) -> None: ++ """Run fn() once after warmup and cudagraph capture (immediately if that already happened).""" ++ if _post_capture_done: ++ fn() ++ else: ++ _post_capture_hooks.append(fn) ++ ++ ++def run_post_capture_hooks() -> None: ++ global _post_capture_done ++ _post_capture_done = True ++ hooks, _post_capture_hooks[:] = list(_post_capture_hooks), [] ++ for fn in hooks: ++ try: ++ fn() ++ except Exception: # noqa: BLE001 -- one broken starter must not stop the others or the engine ++ logger.exception("post-capture reader start failed: %r", fn) ++ ++ + def start_reader() -> None: + """Start the cadence reader. Called AFTER warmup and cudagraph capture. + ++ Also runs the post-capture hooks of the other readers (register_post_capture_hook), whatever the guard mode. ++ + The reader must not exist during capture. It copies device counters to the + host, and a device-to-host copy from a second thread while a capture is + active is exactly the call class cudaStreamCaptureModeGlobal -- the default +@@ -500,6 +529,7 @@ def start_reader() -> None: + Between prealloc() and this call the kernels still count; nothing reads. + """ + global _reader_started ++ run_post_capture_hooks() + if _reader_started or _MODE == "off": + return + _reader_started = True +-- +2.43.0 + diff --git a/patches/0070-qsa-keep-the-clip-counter-start-line-s-log-token.patch b/patches/0070-qsa-keep-the-clip-counter-start-line-s-log-token.patch new file mode 100644 index 0000000..0f44828 --- /dev/null +++ b/patches/0070-qsa-keep-the-clip-counter-start-line-s-log-token.patch @@ -0,0 +1,46 @@ +From cde56de93bd067484ced3193b07fa0f93e94578d Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 05:39:19 -0400 +Subject: [PATCH 070/103] qsa: keep the clip-counter start line's log token + +The post-capture start line keeps the token "QSA clip counter armed (read interval Ns)" that log checks look for; +the previous change had reworded it. A source test pins the token. +--- + tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py | 3 +++ + vllm/models/qwen4_exp/nvidia/qsa.py | 5 ++++- + 2 files changed, 7 insertions(+), 1 deletion(-) + +diff --git a/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py +index a81feb7..7de02bc 100644 +--- a/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py ++++ b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py +@@ -161,6 +161,9 @@ def test_source_pins(): + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "_qsa_clip_count") + assert "threading.Thread" not in ast.unparse(fn), "the clip reader must not be started from the forward" + assert "register_post_capture_hook(_qsa_clip_start_reader)" in qsa ++ # the serving smoke greps r"QSA clip counter armed \(read interval [0-9.]+s\)" (an earlier wording of this line ++ # failed it); the post-capture start line must keep that token. ++ assert "QSA clip counter armed (read interval %.0fs)" in qsa + gtree = ast.parse(_src(_GT_PATH)) + sr = next(n for n in ast.walk(gtree) if isinstance(n, ast.FunctionDef) and n.name == "start_reader") + stmts = [s for s in sr.body if not (isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant))] +diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py +index 467f7ad..a10f0a7 100644 +--- a/vllm/models/qwen4_exp/nvidia/qsa.py ++++ b/vllm/models/qwen4_exp/nvidia/qsa.py +@@ -241,8 +241,11 @@ def _qsa_clip_start_reader() -> None: + name="qsa-clip-reader", + daemon=True, + ).start() ++ # The leading "QSA clip counter armed (read interval Ns)" is the token the serving smoke greps ++ # (v2.2.0 logged it from the first forward); keep it verbatim when editing this line. + logger.info( +- "QSA clip counter reader started after capture (read interval %.0fs, %d layer(s) registered)", ++ "QSA clip counter armed (read interval %.0fs); reader started after capture " ++ "(%d layer(s) registered)", + _qsa_clip_interval, + len(_qsa_clip_counters), + ) +-- +2.43.0 + diff --git a/patches/0071-recoverssm-share-RecoverSSM-commit-kernels-and-helpe.patch b/patches/0071-recoverssm-share-RecoverSSM-commit-kernels-and-helpe.patch new file mode 100644 index 0000000..5216ed0 --- /dev/null +++ b/patches/0071-recoverssm-share-RecoverSSM-commit-kernels-and-helpe.patch @@ -0,0 +1,730 @@ +From 7d1df38791dfe3678ae8389e29c3fd3db3ea586b Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 27 Sep 2026 00:07:04 +0200 +Subject: [PATCH 071/103] recoverssm: share RecoverSSM commit kernels and + helpers across models (port of upstream #58863, a80e41d690) + +Port of the upstream vLLM change #58863 (Apache-2.0). The commit-plan and conv-window compaction kernels move from the Kimi-K3 model +into model_executor/layers/mamba/ops/recoverssm.py with shared helpers; kernel source unchanged. +--- + tests/models/kimi_k3/test_kda.py | 2 +- + tests/models/kimi_k3/test_kda_metadata.py | 6 +- + tests/test_config.py | 4 +- + tests/v1/cudagraph/test_cudagraph_manager.py | 8 +- + vllm/config/cache.py | 2 +- + vllm/model_executor/layers/mamba/abstract.py | 2 +- + .../layers/mamba/ops/recoverssm.py | 196 ++++++++++++++++++ + .../layers/mamba/recoverssm_utils.py | 25 +++ + vllm/models/kimi_k3/nvidia/kda.py | 6 +- + vllm/models/kimi_k3/nvidia/kda_metadata.py | 2 +- + vllm/models/kimi_k3/nvidia/model.py | 4 +- + vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 178 +--------------- + .../attention/backends/recoverssm_metadata.py | 38 ++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- + .../worker/gpu/model_states/mamba_hybrid.py | 2 +- + 15 files changed, 285 insertions(+), 192 deletions(-) + create mode 100644 vllm/model_executor/layers/mamba/ops/recoverssm.py + create mode 100644 vllm/model_executor/layers/mamba/recoverssm_utils.py + +diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py +index 84a5e7e..3fbcc60 100644 +--- a/tests/models/kimi_k3/test_kda.py ++++ b/tests/models/kimi_k3/test_kda.py +@@ -93,7 +93,7 @@ def test_kda_recoverssm_config_state_layout(): + cache_config=SimpleNamespace( + mamba_cache_dtype="auto", + mamba_ssm_cache_dtype="auto", +- use_kda_recoverssm=True, ++ use_recoverssm=True, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + speculative_config=SimpleNamespace(num_speculative_tokens=2), +diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py +index 8ff7aff..8f6f561 100644 +--- a/tests/models/kimi_k3/test_kda_metadata.py ++++ b/tests/models/kimi_k3/test_kda_metadata.py +@@ -119,7 +119,7 @@ def _make_builder( + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.use_replayssm = use_recoverssm +- vllm_config.cache_config.use_kda_recoverssm = use_recoverssm ++ vllm_config.cache_config.use_recoverssm = use_recoverssm + vllm_config.cache_config.prefix_match_unit = prefix_match_unit + builder = builder_cls( + kv_cache_spec=MambaSpec( +@@ -161,7 +161,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): + cache_config=SimpleNamespace( + mamba_cache_dtype="auto", + mamba_ssm_cache_dtype="auto", +- use_kda_recoverssm=True, ++ use_recoverssm=True, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + speculative_config=SimpleNamespace(num_speculative_tokens=2), +@@ -194,7 +194,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): + ), + cache_config=SimpleNamespace( + mamba_cache_mode="align", +- use_kda_recoverssm=True, ++ use_recoverssm=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), +diff --git a/tests/test_config.py b/tests/test_config.py +index 20afe75..bdd527c 100644 +--- a/tests/test_config.py ++++ b/tests/test_config.py +@@ -102,7 +102,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): + config = SimpleNamespace( + cache_config=SimpleNamespace( + use_replayssm=True, +- use_kda_recoverssm=False, ++ use_recoverssm=False, + mamba_cache_mode="none", + ), + num_speculative_tokens=3, +@@ -121,7 +121,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): + + VllmConfig.validate_mamba_cached_kernel(config) + assert config.cache_config.use_replayssm +- assert config.cache_config.use_kda_recoverssm ++ assert config.cache_config.use_recoverssm + + config.cache_config.mamba_cache_mode = "align" + VllmConfig.validate_mamba_cached_kernel(config) +diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py +index 76c4732..27bc729 100644 +--- a/tests/v1/cudagraph/test_cudagraph_manager.py ++++ b/tests/v1/cudagraph/test_cudagraph_manager.py +@@ -259,7 +259,7 @@ def _create_decode_vllm_config( + num_speculative_tokens: int = 0, + dynamic_spec_schedule: list[tuple[int, int, int]] | None = None, + max_num_seqs: int = 8, +- use_kda_recoverssm: bool = False, ++ use_recoverssm: bool = False, + ) -> MagicMock: + compilation_config = CompilationConfig( + cudagraph_mode="FULL_AND_PIECEWISE", +@@ -274,7 +274,7 @@ def _create_decode_vllm_config( + max_num_seqs=max_num_seqs + ) + vllm_config.parallel_config = ParallelConfig() +- vllm_config.cache_config = SimpleNamespace(use_kda_recoverssm=use_kda_recoverssm) ++ vllm_config.cache_config = SimpleNamespace(use_recoverssm=use_recoverssm) + vllm_config.num_speculative_tokens = num_speculative_tokens + if dynamic_spec_schedule is None: + vllm_config.speculative_config = None +@@ -293,7 +293,7 @@ def _make_spec_decode_manager( + num_speculative_tokens: int = 0, + dynamic_spec_schedule: list[tuple[int, int, int]] | None = None, + max_num_seqs: int = 8, +- use_kda_recoverssm: bool = False, ++ use_recoverssm: bool = False, + ) -> gpu_cudagraph_utils.CudaGraphManager: + monkeypatch.setattr( + gpu_cudagraph_utils, +@@ -311,7 +311,7 @@ def _make_spec_decode_manager( + num_speculative_tokens=num_speculative_tokens, + dynamic_spec_schedule=dynamic_spec_schedule, + max_num_seqs=max_num_seqs, +- use_kda_recoverssm=use_kda_recoverssm, ++ use_recoverssm=use_recoverssm, + ), + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, +diff --git a/vllm/config/cache.py b/vllm/config/cache.py +index 9bd8987..6dca830 100644 +--- a/vllm/config/cache.py ++++ b/vllm/config/cache.py +@@ -201,7 +201,7 @@ class CacheConfig: + or FlashInfer mamba backend; standard (non-speculative) decode only. In align + mode flushes are most efficient when mamba_block_size is a multiple of + replayssm_buffer_len, but this is not required.""" +- use_kda_recoverssm: bool = field(default=False, init=False) ++ use_recoverssm: bool = field(default=False, init=False) + """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" + + # Will be set after profiling. +diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py +index 10ed7d0..977081b 100644 +--- a/vllm/model_executor/layers/mamba/abstract.py ++++ b/vllm/model_executor/layers/mamba/abstract.py +@@ -80,7 +80,7 @@ class MambaBase(AttentionLayerBase): + # never writes the baseline's per-draft-token state slots. + num_speculative_blocks=( + 0 +- if vllm_config.cache_config.use_kda_recoverssm ++ if vllm_config.cache_config.use_recoverssm + else vllm_config.num_speculative_tokens + ), + ) +diff --git a/vllm/model_executor/layers/mamba/ops/recoverssm.py b/vllm/model_executor/layers/mamba/ops/recoverssm.py +new file mode 100644 +index 0000000..073bdd5 +--- /dev/null ++++ b/vllm/model_executor/layers/mamba/ops/recoverssm.py +@@ -0,0 +1,196 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Model-independent RecoverSSM Triton kernels shared by every recurrent-state ++implementation. ++ ++RecoverSSM verifies a speculative window from a read-only checkpoint and, after ++sampling, commits the accepted prefix once. The per-model parts (the KDA and GDN verify ++kernels and their state reconstruction) live with their models; the kernels they share ++are here: ++ ++- ``prepare_commit_plan_kernel``: per request, the commit length and the destination ++ state blocks (final and, in align mode, the block-boundary block) from the accepted ++ counts and the block table; ++- ``compact_conv_state_kernel``: moves the accepted suffix of an extended short-conv ++ window to the front of the destination block (and writes the boundary window in align ++ mode).""" ++ ++from vllm.triton_utils import tl, triton ++ ++ ++@triton.heuristics( ++ { ++ "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, ++ "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, ++ } ++) ++@triton.jit ++def prepare_commit_plan_kernel( ++ num_accepted_ptr, ++ request_indices_ptr, ++ state_indices_ptr, ++ query_start_loc_ptr, ++ block_table_ptr, ++ num_computed_ptr, ++ commit_lens_ptr, ++ final_state_indices_ptr, ++ boundary_state_indices_ptr, ++ boundary_recovery_lens_ptr, ++ null_block_id, ++ mamba_block_size, ++ block_table_width, ++ stride_num_accepted, ++ stride_request_indices, ++ stride_state_indices, ++ stride_query_start_loc, ++ stride_block_table_row, ++ stride_block_table_col, ++ stride_num_computed, ++ SPEC_QUERY_LEN: tl.constexpr, ++ HAS_REQUEST_INDICES: tl.constexpr, ++ ALIGN_MODE: tl.constexpr, ++): ++ spec_idx = tl.program_id(0) ++ source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( ++ tl.int64 ++ ) ++ request_idx = spec_idx ++ if HAS_REQUEST_INDICES: ++ request_idx = tl.load( ++ request_indices_ptr + spec_idx * stride_request_indices ++ ).to(tl.int64) ++ num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( ++ tl.int32 ++ ) ++ bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) ++ eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( ++ tl.int64 ++ ) ++ query_len = (eos - bos).to(tl.int32) ++ commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) ++ commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) ++ ++ final_state_idx = source_state_idx ++ boundary_state_idx = null_block_id ++ boundary_recovery_len = 0 ++ if ALIGN_MODE: ++ num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( ++ tl.int32 ++ ) ++ final_num_computed = num_computed + commit_len ++ final_state_col = tl.minimum( ++ final_num_computed // mamba_block_size, block_table_width - 1 ++ ) ++ final_state_idx = tl.load( ++ block_table_ptr ++ + request_idx * stride_block_table_row ++ + final_state_col * stride_block_table_col ++ ).to(tl.int64) ++ next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size ++ crosses_boundary = final_num_computed >= next_boundary ++ boundary_recovery_len = next_boundary - num_computed ++ boundary_state_idx = tl.load( ++ block_table_ptr ++ + request_idx * stride_block_table_row ++ + (next_boundary // mamba_block_size - 1) * stride_block_table_col, ++ mask=crosses_boundary, ++ other=null_block_id, ++ ).to(tl.int64) ++ valid = (source_state_idx > null_block_id) & (commit_len > 0) ++ tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) ++ tl.store( ++ final_state_indices_ptr + spec_idx, ++ tl.where(valid, final_state_idx, null_block_id), ++ ) ++ tl.store( ++ boundary_state_indices_ptr + spec_idx, ++ tl.where(valid, boundary_state_idx, null_block_id), ++ ) ++ tl.store( ++ boundary_recovery_lens_ptr + spec_idx, ++ tl.where(valid, boundary_recovery_len, 0), ++ ) ++ ++ ++@triton.jit ++def compact_conv_state_kernel( ++ conv_state_ref_ptr, ++ conv_state_base_addrs_ptr, ++ conv_state_block_strides_ptr, ++ conv_state_dim_strides_ptr, ++ conv_state_token_strides_ptr, ++ state_indices_ptr, ++ commit_lens_ptr, ++ final_state_indices_ptr, ++ boundary_state_indices_ptr, ++ boundary_recovery_lens_ptr, ++ null_block_id, ++ conv_dim, ++ conv_history_len, ++ stride_state_indices, ++ BLOCK_D: tl.constexpr, ++ BLOCK_HISTORY: tl.constexpr, ++ ALIGN_MODE: tl.constexpr, ++): ++ pid_d = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ pid_l = tl.program_id(2) ++ source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( ++ tl.int64 ++ ) ++ if source_state_idx <= null_block_id: ++ return ++ ++ commit_len = tl.load(commit_lens_ptr + pid_b) ++ if commit_len == 0: ++ return ++ final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) ++ boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) ++ boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) ++ ++ if final_state_idx <= null_block_id: ++ return ++ ++ base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) ++ block_stride = tl.load(conv_state_block_strides_ptr + pid_l) ++ dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) ++ token_stride = tl.load(conv_state_token_strides_ptr + pid_l) ++ conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) ++ source_ptr = conv_state_ptr + source_state_idx * block_stride ++ final_ptr = conv_state_ptr + final_state_idx * block_stride ++ ++ offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) ++ offs_h = tl.arange(0, BLOCK_HISTORY) ++ mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) ++ final_values = tl.load( ++ source_ptr ++ + offs_d[:, None] * dim_stride ++ + (commit_len - 1 + offs_h[None, :]) * token_stride, ++ mask=mask, ++ ) ++ if ALIGN_MODE: ++ boundary_values = tl.load( ++ source_ptr ++ + offs_d[:, None] * dim_stride ++ + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, ++ mask=mask & (boundary_state_idx > null_block_id), ++ ) ++ boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride ++ tl.store( ++ boundary_ptr ++ + offs_d[:, None] * dim_stride ++ + offs_h[None, :] * token_stride, ++ boundary_values, ++ mask=mask & (boundary_state_idx > null_block_id), ++ ) ++ tl.store( ++ final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, ++ final_values, ++ mask=mask, ++ ) ++ ++ ++__all__ = [ ++ "compact_conv_state_kernel", ++ "prepare_commit_plan_kernel", ++] +diff --git a/vllm/model_executor/layers/mamba/recoverssm_utils.py b/vllm/model_executor/layers/mamba/recoverssm_utils.py +new file mode 100644 +index 0000000..1f3998e +--- /dev/null ++++ b/vllm/model_executor/layers/mamba/recoverssm_utils.py +@@ -0,0 +1,25 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Model-independent RecoverSSM Python helpers: the activation predicate shared by ++layers, models and backends, and argument validation that raises ``ValueError`` ++before any raw-pointer kernel launch.""" ++ ++ ++def recoverssm_require(cond: bool, msg: str, component: str = "RecoverSSM") -> None: ++ if not cond: ++ raise ValueError(f"{component}: {msg}") ++ ++ ++def uses_recoverssm(cache_config, num_speculative_tokens: int) -> bool: ++ """True when this model instance runs RecoverSSM for its recurrent states. ++ ++ VllmConfig sets ``cache_config.use_recoverssm`` for every model whose recurrent ++ states support it (Kimi-K3 KDA, Qwen GDN/PLE). Without speculative tokens there is ++ nothing to recover, so the stock path runs.""" ++ return ( ++ bool(getattr(cache_config, "use_recoverssm", False)) ++ and num_speculative_tokens > 0 ++ ) ++ ++ ++__all__ = ["recoverssm_require", "uses_recoverssm"] +diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py +index 9232694..27de6ad 100644 +--- a/vllm/models/kimi_k3/nvidia/kda.py ++++ b/vllm/models/kimi_k3/nvidia/kda.py +@@ -547,7 +547,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, + ) +- if self.cache_config.use_kda_recoverssm: ++ if self.cache_config.use_recoverssm: + return MambaStateDtypeCalculator.append_kda_recoverssm_record( + base_dtypes, self.model_config.dtype + ) +@@ -563,7 +563,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) +- if self.cache_config.use_kda_recoverssm: ++ if self.cache_config.use_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + base_shapes, + self.num_heads, +@@ -582,7 +582,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + run_gemm_rs_ar: bool = False, + ) -> None: + super().__init__(config, vllm_config, prefix) +- self.use_recoverssm = self.cache_config.use_kda_recoverssm ++ self.use_recoverssm = self.cache_config.use_recoverssm + if self.cache_config.use_replayssm and not self.use_recoverssm: + raise ValueError( + "Kimi-K3 supports --use-replayssm only with speculative decoding" +diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py +index 0b395b0..4999868 100644 +--- a/vllm/models/kimi_k3/nvidia/kda_metadata.py ++++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py +@@ -327,7 +327,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + isinstance(additional_config, dict) + and additional_config.get("kda_prefill_backend") == "flashinfer" + ) +- self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm ++ self.use_recoverssm = vllm_config.cache_config.use_recoverssm + self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 + self.recoverssm_num_accepted_tokens: torch.Tensor | None = None + self.recoverssm_context: KDARecoverSSMCommitContext | None = None +diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py +index 6a26ce7..c41a0fc 100644 +--- a/vllm/models/kimi_k3/nvidia/model.py ++++ b/vllm/models/kimi_k3/nvidia/model.py +@@ -1676,7 +1676,7 @@ class KimiLinearForCausalLM( + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) +- if vllm_config.cache_config.use_kda_recoverssm: ++ if vllm_config.cache_config.use_recoverssm: + dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( + dtypes, vllm_config.model_config.dtype + ) +@@ -1701,7 +1701,7 @@ class KimiLinearForCausalLM( + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) +- if vllm_config.cache_config.use_kda_recoverssm: ++ if vllm_config.cache_config.use_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + shapes, + hf_config.linear_attn_config["num_heads"], +diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +index c603ec7..bca9e52 100644 +--- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py ++++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +@@ -9,6 +9,12 @@ from typing import Any + import torch + + from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first ++from vllm.model_executor.layers.mamba.ops.recoverssm import ( ++ compact_conv_state_kernel as _compact_conv_state_kernel, ++) ++from vllm.model_executor.layers.mamba.ops.recoverssm import ( ++ prepare_commit_plan_kernel as _prepare_commit_plan_kernel, ++) + from vllm.triton_utils import tl, triton + from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +@@ -223,178 +229,6 @@ def _kda_recoverssm_verify_kernel( + ) + + +-@triton.heuristics( +- { +- "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, +- "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, +- } +-) +-@triton.jit +-def _prepare_commit_plan_kernel( +- num_accepted_ptr, +- request_indices_ptr, +- state_indices_ptr, +- query_start_loc_ptr, +- block_table_ptr, +- num_computed_ptr, +- commit_lens_ptr, +- final_state_indices_ptr, +- boundary_state_indices_ptr, +- boundary_recovery_lens_ptr, +- null_block_id, +- mamba_block_size, +- block_table_width, +- stride_num_accepted, +- stride_request_indices, +- stride_state_indices, +- stride_query_start_loc, +- stride_block_table_row, +- stride_block_table_col, +- stride_num_computed, +- SPEC_QUERY_LEN: tl.constexpr, +- HAS_REQUEST_INDICES: tl.constexpr, +- ALIGN_MODE: tl.constexpr, +-): +- spec_idx = tl.program_id(0) +- source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( +- tl.int64 +- ) +- request_idx = spec_idx +- if HAS_REQUEST_INDICES: +- request_idx = tl.load( +- request_indices_ptr + spec_idx * stride_request_indices +- ).to(tl.int64) +- num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( +- tl.int32 +- ) +- bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) +- eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( +- tl.int64 +- ) +- query_len = (eos - bos).to(tl.int32) +- commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) +- commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) +- +- final_state_idx = source_state_idx +- boundary_state_idx = null_block_id +- boundary_recovery_len = 0 +- if ALIGN_MODE: +- num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( +- tl.int32 +- ) +- final_num_computed = num_computed + commit_len +- final_state_col = tl.minimum( +- final_num_computed // mamba_block_size, block_table_width - 1 +- ) +- final_state_idx = tl.load( +- block_table_ptr +- + request_idx * stride_block_table_row +- + final_state_col * stride_block_table_col +- ).to(tl.int64) +- next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size +- crosses_boundary = final_num_computed >= next_boundary +- boundary_recovery_len = next_boundary - num_computed +- boundary_state_idx = tl.load( +- block_table_ptr +- + request_idx * stride_block_table_row +- + (next_boundary // mamba_block_size - 1) * stride_block_table_col, +- mask=crosses_boundary, +- other=null_block_id, +- ).to(tl.int64) +- valid = (source_state_idx > null_block_id) & (commit_len > 0) +- tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) +- tl.store( +- final_state_indices_ptr + spec_idx, +- tl.where(valid, final_state_idx, null_block_id), +- ) +- tl.store( +- boundary_state_indices_ptr + spec_idx, +- tl.where(valid, boundary_state_idx, null_block_id), +- ) +- tl.store( +- boundary_recovery_lens_ptr + spec_idx, +- tl.where(valid, boundary_recovery_len, 0), +- ) +- +- +-@triton.jit +-def _compact_conv_state_kernel( +- conv_state_ref_ptr, +- conv_state_base_addrs_ptr, +- conv_state_block_strides_ptr, +- conv_state_dim_strides_ptr, +- conv_state_token_strides_ptr, +- state_indices_ptr, +- commit_lens_ptr, +- final_state_indices_ptr, +- boundary_state_indices_ptr, +- boundary_recovery_lens_ptr, +- null_block_id, +- conv_dim, +- conv_history_len, +- stride_state_indices, +- BLOCK_D: tl.constexpr, +- BLOCK_HISTORY: tl.constexpr, +- ALIGN_MODE: tl.constexpr, +-): +- pid_d = tl.program_id(0) +- pid_b = tl.program_id(1) +- pid_l = tl.program_id(2) +- source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( +- tl.int64 +- ) +- if source_state_idx <= null_block_id: +- return +- +- commit_len = tl.load(commit_lens_ptr + pid_b) +- if commit_len == 0: +- return +- final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) +- boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) +- boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) +- +- if final_state_idx <= null_block_id: +- return +- +- base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) +- block_stride = tl.load(conv_state_block_strides_ptr + pid_l) +- dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) +- token_stride = tl.load(conv_state_token_strides_ptr + pid_l) +- conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) +- source_ptr = conv_state_ptr + source_state_idx * block_stride +- final_ptr = conv_state_ptr + final_state_idx * block_stride +- +- offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) +- offs_h = tl.arange(0, BLOCK_HISTORY) +- mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) +- final_values = tl.load( +- source_ptr +- + offs_d[:, None] * dim_stride +- + (commit_len - 1 + offs_h[None, :]) * token_stride, +- mask=mask, +- ) +- if ALIGN_MODE: +- boundary_values = tl.load( +- source_ptr +- + offs_d[:, None] * dim_stride +- + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, +- mask=mask & (boundary_state_idx > null_block_id), +- ) +- boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride +- tl.store( +- boundary_ptr +- + offs_d[:, None] * dim_stride +- + offs_h[None, :] * token_stride, +- boundary_values, +- mask=mask & (boundary_state_idx > null_block_id), +- ) +- tl.store( +- final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, +- final_values, +- mask=mask, +- ) +- +- + @triton.jit + def _commit_kda_state_kernel( + state_ref_ptr, +diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py +index 7e4cca6..b1339e9 100644 +--- a/vllm/v1/attention/backends/recoverssm_metadata.py ++++ b/vllm/v1/attention/backends/recoverssm_metadata.py +@@ -24,3 +24,41 @@ class RecoverSSMMetadata(abc.ABC): + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + raise NotImplementedError ++ ++ ++def recoverssm_spec_rows( ++ is_prefilling_cpu: torch.Tensor, ++ query_start_loc_cpu: torch.Tensor, ++ num_decode_draft_tokens_cpu: torch.Tensor, ++ num_reqs: int, ++) -> tuple[torch.Tensor, torch.Tensor]: ++ """Which batch rows take the RecoverSSM spec path, shared by the RecoverSSM metadata ++ builders (GDN, PLE). ++ ++ Every active decode row does (a draft-less step is a window of one token), plus ++ every row that has drafts. Only the first ``num_reqs`` rows are real; padding rows ++ never take the spec path. Returns the draft-count vector to hand to the base builder ++ (1 for spec rows, -1 otherwise, padding included) and the spec rows in batch order, ++ which is also the order of the base builder's spec state indices.""" ++ query_lens = query_start_loc_cpu[: num_reqs + 1].diff() ++ active_decode = (~is_prefilling_cpu[:num_reqs]) & (query_lens > 0) ++ spec_mask = (num_decode_draft_tokens_cpu[:num_reqs] >= 0) | active_decode ++ drafts = torch.full_like(num_decode_draft_tokens_cpu, -1) ++ drafts[:num_reqs] = torch.where(spec_mask, 1, -1).to(drafts.dtype) ++ return drafts, spec_mask.nonzero().flatten() ++ ++ ++def recoverssm_request_indices( ++ rows: torch.Tensor, num_spec_decodes: int, device: torch.device ++) -> torch.Tensor | None: ++ """None when the spec rows are exactly 0..n-1, else their batch rows (the commit ++ plan's request mapping). Raises if the base builder counted a different number of ++ spec rows than the classification above.""" ++ if rows.numel() != num_spec_decodes: ++ raise ValueError( ++ f"RecoverSSM: {rows.numel()} spec rows classified, but the builder " ++ f"produced {num_spec_decodes}" ++ ) ++ if num_spec_decodes == 0 or int(rows[-1]) == num_spec_decodes - 1: ++ return None ++ return rows.to(torch.int32).to(device, non_blocking=True) +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 921c562..a143a3d 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -353,7 +353,7 @@ class CudaGraphManager: + + # recoverSSM cannot capture a dummy query wider than its workspace. + if mixed_mode and ( +- not self.vllm_config.cache_config.use_kda_recoverssm ++ not self.vllm_config.cache_config.use_recoverssm + or num_tokens <= max_decode_tokens + ): + # for PIECEWISE graphs there is no limit on requests when replaying +diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +index 6b6e624..de0e1ce 100644 +--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py ++++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +@@ -98,7 +98,7 @@ class MambaHybridModelState(DefaultModelState): + # the persistent per-req-slot tables, not the per-step gather. + self._runner_block_tables: Any = None + self.recoverssm = ( +- RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ++ RecoverSSMState() if self.cache_config.use_recoverssm else None + ) + if self._align_mode: + self._mamba_state_idx_gpu = torch.zeros( +-- +2.43.0 + diff --git a/patches/0072-recoverssm-RecoverSSM-for-Qwen-GDN-and-the-PLE-short.patch b/patches/0072-recoverssm-RecoverSSM-for-Qwen-GDN-and-the-PLE-short.patch new file mode 100644 index 0000000..0087e5e --- /dev/null +++ b/patches/0072-recoverssm-RecoverSSM-for-Qwen-GDN-and-the-PLE-short.patch @@ -0,0 +1,2715 @@ +From f1902d20ba112511ed8ad1ad042f68e484e74d58 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 27 Sep 2026 00:07:04 +0200 +Subject: [PATCH 072/103] recoverssm: RecoverSSM for Qwen GDN and the PLE short + conv (port of upstream #58863, 03cae7837f) + +Port of the upstream vLLM change #58863 (Apache-2.0). With --use-replayssm and speculative decoding, Gated DeltaNet layers verify the +draft window from a read-only checkpoint and keep a per-token replay record; after sampling one commit replays +the accepted tokens and writes the recurrent state once. The PLE short conv gets the same treatment. +--- + tests/config/test_recoverssm_config.py | 173 +++++ + tests/kernels/mamba/test_recoverssm_gdn.py | 601 +++++++++++++++ + tests/kernels/mamba/test_recoverssm_ple.py | 153 ++++ + vllm/config/cache.py | 2 +- + vllm/config/compilation.py | 13 +- + vllm/config/vllm.py | 29 +- + .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 132 +++- + .../layers/mamba/gdn/recoverssm_gdn.py | 695 ++++++++++++++++++ + vllm/model_executor/models/interfaces.py | 1 + + vllm/models/qwen4_exp/nvidia/model.py | 59 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 6 +- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 8 + + vllm/v1/attention/backends/gdn_recoverssm.py | 181 +++++ + vllm/v1/attention/backends/mamba_attn.py | 7 +- + vllm/v1/attention/backends/ple_recoverssm.py | 326 ++++++++ + 15 files changed, 2352 insertions(+), 34 deletions(-) + create mode 100644 tests/config/test_recoverssm_config.py + create mode 100644 tests/kernels/mamba/test_recoverssm_gdn.py + create mode 100644 tests/kernels/mamba/test_recoverssm_ple.py + create mode 100644 vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py + create mode 100644 vllm/v1/attention/backends/gdn_recoverssm.py + create mode 100644 vllm/v1/attention/backends/ple_recoverssm.py + +diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py +new file mode 100644 +index 0000000..21791aa +--- /dev/null ++++ b/tests/config/test_recoverssm_config.py +@@ -0,0 +1,173 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Configuration-level regression tests for Qwen4Exp GDN/PLE RecoverSSM selection. ++ ++``--use-replayssm`` with speculative decoding selects RecoverSSM for Qwen4Exp, as it ++does for Kimi-K3 KDA. The MTP drafter's derived config shares ``cache_config`` and must ++keep the choice. These tests run the real validator function against stand-in configs. ++""" ++ ++from types import SimpleNamespace ++from typing import Any, cast ++ ++import pytest ++ ++from vllm.config.mamba import MambaBackendEnum ++from vllm.config.vllm import VllmConfig ++from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm ++from vllm.v1.attention.backend import AttentionCGSupport ++ ++TARGET = "Qwen4ExpForConditionalGeneration" ++DRAFTER = "Qwen4ExpMTP" ++ ++ ++def _validator(): ++ validators = VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined] ++ return validators["validate_mamba_cached_kernel"].func ++ ++ ++def _config( ++ arch: str, ++ cache_config=None, ++ num_spec: int = 3, ++ use_replayssm: bool = True, ++ supports_replayssm: bool = True, ++ **overrides, ++): ++ cfg = SimpleNamespace( ++ num_speculative_tokens=num_spec, ++ model_config=SimpleNamespace( ++ architecture=arch, supports_replayssm=supports_replayssm ++ ), ++ cache_config=cache_config ++ or SimpleNamespace( ++ use_recoverssm=False, ++ use_replayssm=use_replayssm, ++ mamba_cache_mode="align", ++ ), ++ mamba_config=SimpleNamespace( ++ enable_stochastic_rounding=False, backend=MambaBackendEnum.TRITON ++ ), ++ use_v2_model_runner=True, ++ parallel_config=SimpleNamespace(pipeline_parallel_size=1), ++ kv_transfer_config=None, ++ ) ++ for key, value in overrides.items(): ++ target, _, attr = key.partition("__") ++ if attr: ++ setattr(getattr(cfg, target), attr, value) ++ else: ++ setattr(cfg, target, value) ++ return cfg ++ ++ ++def test_validator_is_registered(): ++ assert "validate_mamba_cached_kernel" in ( ++ VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined] ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM"] ++) ++def test_replayssm_with_spec_selects_recoverssm(arch): ++ cfg = _config(arch) ++ _validator()(cfg) ++ assert cfg.cache_config.use_recoverssm is True ++ assert uses_recoverssm(cfg.cache_config, 3) ++ ++ ++def test_derived_drafter_keeps_the_choice_and_runs_the_checks(): ++ target = _config(TARGET) ++ _validator()(target) ++ drafter = _config(DRAFTER, cache_config=target.cache_config) ++ _validator()(drafter) ++ assert target.cache_config.use_recoverssm is True ++ bad = _config( ++ DRAFTER, ++ cache_config=target.cache_config, ++ parallel_config__pipeline_parallel_size=2, ++ ) ++ with pytest.raises(ValueError, match="pipeline_parallel_size"): ++ _validator()(bad) ++ ++ ++@pytest.mark.parametrize( ++ ("override", "value", "match"), ++ [ ++ ("parallel_config__pipeline_parallel_size", 2, "pipeline_parallel_size"), ++ ("mamba_config__enable_stochastic_rounding", True, "stochastic"), ++ ("mamba_config__backend", "flashinfer", "mamba-backend triton"), ++ ("use_v2_model_runner", False, "V2_MODEL_RUNNER"), ++ ("cache_config__mamba_cache_mode", "all", "none and align"), ++ ], ++) ++def test_runtime_checks_reject(override, value, match): ++ cfg = _config(TARGET, **{override: value}) ++ with pytest.raises(ValueError, match=match): ++ _validator()(cfg) ++ ++ ++def test_flag_unset_keeps_the_stock_path(): ++ cfg = _config(TARGET, use_replayssm=False) ++ _validator()(cfg) ++ assert cfg.cache_config.use_recoverssm is False ++ assert not uses_recoverssm(cfg.cache_config, 3) ++ ++ ++def test_qwen4exp_replayssm_requires_speculative_decoding(): ++ with pytest.raises(ValueError, match="requires speculative decoding"): ++ _validator()(_config(TARGET, num_spec=0)) ++ ++ ++def test_architecture_without_replayssm_support_is_rejected(): ++ cfg = _config("Qwen3NextForCausalLM", supports_replayssm=False) ++ with pytest.raises(ValueError, match="not supported for architecture"): ++ _validator()(cfg) ++ ++ ++def test_uses_recoverssm_predicate(): ++ on = SimpleNamespace(use_recoverssm=True) ++ assert uses_recoverssm(on, 3) ++ assert not uses_recoverssm(on, 0) ++ assert not uses_recoverssm(SimpleNamespace(use_recoverssm=False), 3) ++ assert not uses_recoverssm(SimpleNamespace(), 3) ++ ++ ++def test_builders_fall_back_to_piecewise(): ++ from vllm.v1.attention.backends.gdn_recoverssm import GDNRecoverSSMMetadataBuilder ++ from vllm.v1.attention.backends.ple_recoverssm import PleRecoverSSMMetadataBuilder ++ ++ for builder in (GDNRecoverSSMMetadataBuilder, PleRecoverSSMMetadataBuilder): ++ assert ( ++ builder.get_cudagraph_support(cast(Any, None), cast(Any, None)) ++ == AttentionCGSupport.NEVER ++ ) ++ ++ ++@pytest.mark.parametrize("compile_mode", ["VLLM_COMPILE", "NONE"]) ++@pytest.mark.parametrize(("breakable", "expected"), [("1", "PIECEWISE"), ("0", "NONE")]) ++def test_never_support_falls_back_to_piecewise_with_breakable_cudagraphs( ++ monkeypatch, compile_mode, breakable, expected ++): ++ """A backend without full-graph support must keep piecewise graphs when breakable ++ CUDA graphs split at attention (splitting_ops is empty then, and torch.compile may ++ be off).""" ++ from vllm.config.compilation import ( ++ CompilationConfig, ++ CompilationMode, ++ CUDAGraphMode, ++ ) ++ ++ monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", breakable) ++ compilation_config = CompilationConfig( ++ mode=CompilationMode[compile_mode], ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ splitting_ops=[], ++ ) ++ resolved = compilation_config.resolve_cudagraph_mode_and_sizes( ++ AttentionCGSupport.NEVER, ++ "GDNRecoverSSMAttentionBackend", ++ uniform_decode_query_len=4, ++ ) ++ assert resolved.name == expected +diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py +new file mode 100644 +index 0000000..1e0cbcf +--- /dev/null ++++ b/tests/kernels/mamba/test_recoverssm_gdn.py +@@ -0,0 +1,601 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""GDN RecoverSSM: verify against the native fused kernel, commit against the native ++per-token states, and the boundary checks (malformed metadata must raise ValueError ++before any kernel reads or writes out of bounds).""" ++ ++from typing import Any, cast ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( ++ GDNRecoverSSMCommitContext, ++ gdn_recoverssm_verify, ++) ++from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import ( ++ fused_sigmoid_gating_delta_rule_update, ++) ++ ++pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") ++H, HV, K, V, T = 16, 48, 128, 128, 4 ++QLENS = (4, 4, 3) ++ ++ ++def _inputs(nb: int = 10) -> dict: ++ torch.manual_seed(0) ++ dev = "cuda" ++ tot = sum(QLENS) ++ qsl = torch.tensor([0, 4, 8, 11], dtype=torch.int32, device=dev) ++ return dict( ++ A_log=(torch.rand(HV, device=dev) * 2 - 1).float(), ++ a=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16), ++ b=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16), ++ dt_bias=(torch.randn(HV, device=dev) * 0.5).bfloat16(), ++ q=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16), ++ k=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16), ++ v=torch.randn(1, tot, HV, V, device=dev, dtype=torch.bfloat16), ++ checkpoint_state=(torch.randn(nb, HV, V, K, device=dev) * 0.05).float(), ++ replay_cache=torch.zeros(nb, HV, T, V + K + 1, device=dev), ++ query_start_loc=qsl, ++ state_indices=torch.tensor([2, 5, 7], dtype=torch.int32, device=dev), ++ spec_query_len=T, ++ ) ++ ++ ++def _verify(inp: dict, **override): ++ d = {**inp, **override} ++ return gdn_recoverssm_verify( ++ d["A_log"], ++ d["a"], ++ d["b"], ++ d["dt_bias"], ++ d["q"], ++ d["k"], ++ d["v"], ++ checkpoint_state=d["checkpoint_state"], ++ replay_cache=d["replay_cache"], ++ query_start_loc=d["query_start_loc"], ++ state_indices=d["state_indices"], ++ spec_query_len=d["spec_query_len"], ++ out=d.get("out"), ++ ) ++ ++ ++def _native(inp: dict): ++ return fused_sigmoid_gating_delta_rule_update( ++ A_log=inp["A_log"], ++ a=inp["a"], ++ b=inp["b"], ++ dt_bias=inp["dt_bias"], ++ q=inp["q"], ++ k=inp["k"], ++ v=inp["v"], ++ initial_state=inp["checkpoint_state"].clone(), ++ inplace_final_state=False, ++ cu_seqlens=inp["query_start_loc"].long(), ++ ssm_state_indices=inp["state_indices"].long(), ++ use_qk_l2norm_in_kernel=True, ++ ) ++ ++ ++def _context(inp: dict, ckpt: torch.Tensor) -> GDNRecoverSSMCommitContext: ++ nb = ckpt.shape[0] ++ conv = [torch.zeros(nb, 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)] ++ return GDNRecoverSSMCommitContext.from_tensors( ++ conv, [ckpt], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 ++ ) ++ ++ ++def test_verify_matches_native_and_keeps_the_checkpoint(): ++ inp = _inputs() ++ before = inp["checkpoint_state"].clone() ++ out = _verify(inp) ++ ref, _ = _native(inp) ++ torch.testing.assert_close( ++ out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0 ++ ) ++ assert torch.equal(inp["checkpoint_state"], before) ++ ++ ++@pytest.mark.parametrize("accepted", [1, 2, 3, 4]) ++def test_commit_matches_native_per_token_state(accepted): ++ inp = _inputs() ++ _verify(inp) ++ _, states = _native(inp) ++ ckpt = inp["checkpoint_state"].clone() ++ ctx = _context(inp, ckpt) ++ acc = torch.tensor( ++ [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" ++ ) ++ ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) ++ for i, blk in enumerate(inp["state_indices"].tolist()): ++ tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 ++ torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) ++ others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()] ++ assert torch.equal(ckpt[others], inp["checkpoint_state"][others]) ++ ++ ++@pytest.mark.parametrize( ++ "case", ["short_qsl", "capacity", "hv_lt_h", "a_log", "out", "replay", "device"] ++) ++def test_verify_rejects_malformed_inputs(case): ++ inp = _inputs() ++ over, match = { ++ "short_qsl": ( ++ dict(query_start_loc=inp["query_start_loc"][:-1]), ++ "query metadata", ++ ), ++ "capacity": ( ++ dict( ++ state_indices=inp["state_indices"][:2], ++ query_start_loc=inp["query_start_loc"][:3], ++ ), ++ "activation capacity", ++ ), ++ "hv_lt_h": ( ++ dict( ++ v=inp["v"][:, :, :8].contiguous(), ++ a=inp["a"][:, :8].contiguous(), ++ b=inp["b"][:, :8].contiguous(), ++ ), ++ "positive multiple", ++ ), ++ "a_log": (dict(A_log=torch.zeros(HV + 1, device="cuda")), "A_log"), ++ "out": ( ++ dict( ++ out=torch.empty(1, 11, HV, V + 1, device="cuda", dtype=torch.bfloat16) ++ ), ++ "output shape", ++ ), ++ "replay": ( ++ dict(replay_cache=torch.zeros(10, HV, T, V + K, device="cuda")), ++ "replay buffer", ++ ), ++ "device": (dict(A_log=inp["A_log"].cpu()), "same device"), ++ }[case] ++ with pytest.raises(ValueError, match=match): ++ _verify(inp, **over) ++ ++ ++def test_commit_rejects_partial_or_short_metadata(): ++ inp = _inputs() ++ ctx = _context(inp, inp["checkpoint_state"].clone()) ++ acc = torch.ones(3, dtype=torch.int32, device="cuda") ++ bt = torch.zeros(3, 4, dtype=torch.int32, device="cuda") ++ with pytest.raises(ValueError, match="align metadata is incomplete"): ++ ctx.commit( ++ acc, ++ inp["state_indices"], ++ inp["query_start_loc"], ++ block_table=bt, ++ mamba_block_size=16, ++ ) ++ with pytest.raises(ValueError, match="commit metadata"): ++ ctx.commit(acc, inp["state_indices"], inp["query_start_loc"][:-1]) ++ with pytest.raises(ValueError, match="request mapping"): ++ ctx.commit( ++ acc, ++ inp["state_indices"], ++ inp["query_start_loc"], ++ request_indices=torch.zeros(2, dtype=torch.int32, device="cuda"), ++ ) ++ ++ ++def test_context_rejects_inconsistent_layers(): ++ inp = _inputs() ++ ck = inp["checkpoint_state"] ++ conv = torch.zeros(ck.shape[0], 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16) ++ with pytest.raises(ValueError, match="at least one layer"): ++ GDNRecoverSSMCommitContext.from_tensors( ++ [], [], [], spec_query_len=T, max_num_reqs=8 ++ ) ++ with pytest.raises(ValueError, match="differ"): ++ GDNRecoverSSMCommitContext.from_tensors( ++ [conv], [ck, ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 ++ ) ++ with pytest.raises(ValueError, match="block count"): ++ GDNRecoverSSMCommitContext.from_tensors( ++ [conv[:5]], [ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 ++ ) ++ ++ ++# ------------------------------------------------------------------ more cases: ++# mapping, align, extreme gates ++from vllm.v1.attention.backends.recoverssm_metadata import ( # noqa: E402 ++ recoverssm_request_indices, ++ recoverssm_spec_rows, ++) ++ ++# (A_log range, a scale, b offset): drive g = -exp(A_log) * softplus(a + dt_bias) ++# and beta = sigmoid(b) ++GATES = { ++ "typical": (0.0, 1.0, 0.0), ++ "g_near_0": (-12.0, 1.0, 0.0), ++ "g_very_negative": (2.5, 4.0, 0.0), ++ "beta_near_0": (0.0, 1.0, -12.0), ++ "beta_near_1": (0.0, 1.0, 12.0), ++} ++ ++ ++def _make(qlens, nb=16, seed=0, gate="typical"): ++ g = torch.Generator(device="cuda").manual_seed(seed) ++ rnd = lambda *s: torch.randn(*s, device="cuda", generator=g) # noqa: E731 ++ a_log_shift, ascale, boff = GATES[gate] ++ tot = sum(qlens) ++ qsl = torch.tensor( ++ [0, *torch.tensor(qlens).cumsum(0).tolist()], dtype=torch.int32, device="cuda" ++ ) ++ return dict( ++ A_log=( ++ torch.rand(HV, device="cuda", generator=g) * 2 - 1 + a_log_shift ++ ).float(), ++ a=(rnd(tot, HV) * ascale).bfloat16(), ++ b=(rnd(tot, HV) + boff).bfloat16(), ++ dt_bias=(rnd(HV) * 0.5).bfloat16(), ++ q=rnd(1, tot, H, K).bfloat16(), ++ k=rnd(1, tot, H, K).bfloat16(), ++ v=rnd(1, tot, HV, V).bfloat16(), ++ checkpoint_state=(rnd(nb, HV, V, K) * 0.05).float(), ++ replay_cache=torch.zeros(nb, HV, T, V + K + 1, device="cuda"), ++ query_start_loc=qsl, ++ spec_query_len=T, ++ ) ++ ++ ++def test_spec_rows_mixed_batch(): ++ # rows: 0 spec K=3 | 1 prefill | 2 draft-less decode | 3 spec K=1 | 4 padded ++ # zero-query | 5,6 beyond num_reqs ++ is_prefilling = torch.tensor([False, True, False, False, False, False, False]) ++ qsl = torch.tensor([0, 4, 104, 105, 107, 107], dtype=torch.int32) ++ drafts_in = torch.tensor( ++ [3, -1, -1, 1, -1, 2, 0], dtype=torch.int32 ++ ) # 5, 6: garbage padding ++ drafts, rows = recoverssm_spec_rows(is_prefilling, qsl, drafts_in, num_reqs=5) ++ assert rows.tolist() == [0, 2, 3] ++ assert drafts.tolist() == [1, -1, 1, 1, -1, -1, -1] ++ idx = recoverssm_request_indices(rows, 3, torch.device("cuda")) ++ assert idx is not None and idx.tolist() == [0, 2, 3] ++ assert ( ++ recoverssm_request_indices(torch.tensor([0, 1, 2]), 3, torch.device("cuda")) ++ is None ++ ) ++ with pytest.raises(ValueError, match="spec rows classified"): ++ recoverssm_request_indices(rows, 2, torch.device("cuda")) ++ ++ ++def test_commit_maps_noncontiguous_requests_in_align_mode(): ++ """Spec rows 0, 2, 3 of a 5-request batch; accepted counts, block table and ++ num_computed are indexed by request row. Every spec request must commit its own ++ state into its own block.""" ++ qlens = [4, 1, 2] # the spec rows' windows (K=3, draft-less, K=1) ++ inp = _make(qlens, nb=40) ++ bs = 8 ++ bt = torch.arange(1, 1 + 5 * 6, dtype=torch.int32, device="cuda").reshape( ++ 5, 6 ++ ) # request row -> blocks ++ nc = torch.tensor( ++ [5, 0, 16, 9, 0], dtype=torch.int32, device="cuda" ++ ) # per request row ++ req_rows = torch.tensor([0, 2, 3], dtype=torch.int32, device="cuda") ++ src = torch.stack([bt[r, int(nc[r]) // bs] for r in req_rows.tolist()]).contiguous() ++ inp["state_indices"] = src ++ _verify(inp) ++ _, states = _native(inp) ++ ckpt = inp["checkpoint_state"].clone() ++ ctx = _context(inp, ckpt) ++ acc = torch.tensor( ++ [3, 9, 1, 2, 9], dtype=torch.int32, device="cuda" ++ ) # rows 1, 4 must be ignored ++ ctx.commit( ++ acc, ++ src, ++ inp["query_start_loc"], ++ request_indices=req_rows, ++ block_table=bt, ++ num_computed_tokens=nc, ++ mamba_block_size=bs, ++ ) ++ for i, r in enumerate(req_rows.tolist()): ++ n = min(int(acc[r]), qlens[i]) ++ c0 = int(nc[r]) ++ final_blk = int(bt[r, (c0 + n) // bs]) ++ tok = int(inp["query_start_loc"][i]) + n - 1 ++ torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6) ++ ++ ++# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on ++# it / crossing and continuing / accepted = 1 without crossing / the full window ++# crossing after 3 ++ALIGN_CASES = [ ++ (3, 2, 4), ++ (2, 3, 4), ++ (0, 4, 4), ++ (6, 4, 4), ++ (1, 1, 4), ++ (5, 4, 8), ++ (7, 4, 8), ++ (4, 4, 8), ++] ++ ++ ++@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES) ++def test_align_boundary_and_final_state_match_native(nc, accepted, bs): ++ inp = _make([T], nb=12) ++ bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6) ++ src = bt[0, nc // bs].reshape(1).contiguous() ++ inp["state_indices"] = src ++ _verify(inp) ++ _, states = _native(inp) ++ ckpt = inp["checkpoint_state"].clone() ++ ctx = _context(inp, ckpt) ++ ctx.commit( ++ torch.tensor([accepted], dtype=torch.int32, device="cuda"), ++ src, ++ inp["query_start_loc"], ++ block_table=bt, ++ num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"), ++ mamba_block_size=bs, ++ ) ++ next_boundary = (nc // bs + 1) * bs ++ final_blk = int(bt[0, (nc + accepted) // bs]) ++ if ( ++ nc + accepted >= next_boundary ++ ): # crosses: the boundary state goes into the source block ++ torch.testing.assert_close( ++ ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6 ++ ) ++ torch.testing.assert_close( ++ ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6 ++ ) ++ ++ ++@pytest.mark.parametrize("gate", list(GATES)) ++@pytest.mark.parametrize("seed", [0, 1, 2]) ++def test_commit_matches_native_under_extreme_gates(gate, seed): ++ inp = _make(list(QLENS), nb=10, seed=seed, gate=gate) ++ inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") ++ _verify(inp) ++ _, states = _native(inp) ++ for accepted in (1, 2, 3, 4): ++ ckpt = inp["checkpoint_state"].clone() ++ ctx = _context(inp, ckpt) ++ acc = torch.tensor( ++ [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" ++ ) ++ ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) ++ for i, blk in enumerate(inp["state_indices"].tolist()): ++ tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 ++ torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) ++ ++ ++@pytest.mark.parametrize("batch", [1, 2, 4]) ++def test_verify_on_strided_views_writes_in_place(batch): ++ """The layer passes views of the packed conv output (token stride = the packed ++ width) and the layer's output buffer; result and replay record must equal the ++ contiguous-copy path.""" ++ torch.manual_seed(batch) ++ H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6 ++ T = batch * SQ ++ qd, vd = H * K, HV * V ++ packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16) ++ mixed = packed[:, : 2 * qd + vd] ++ a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) ++ b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) ++ A_log = torch.randn(HV, device="cuda") ++ dt_bias = torch.randn(HV, device="cuda") ++ checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05 ++ qsl = torch.arange(0, T + 1, SQ, device="cuda", dtype=torch.int32) ++ si = torch.arange(1, batch + 1, device="cuda", dtype=torch.int32) ++ ++ def run(q, k, v, out=None): ++ replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda") ++ o = gdn_recoverssm_verify( ++ A_log, ++ a, ++ b, ++ dt_bias, ++ q, ++ k, ++ v, ++ checkpoint_state=checkpoint, ++ replay_cache=replay, ++ query_start_loc=qsl, ++ state_indices=si, ++ spec_query_len=SQ, ++ out=out, ++ ) ++ return o, replay ++ ++ q, k, v = ( ++ x.contiguous().view(1, T, -1, d) ++ for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V)) ++ ) ++ want_out, want_replay = run(q, k, v) ++ ++ core = torch.zeros(T + 3, HV, V, device="cuda", dtype=torch.bfloat16) ++ views = ( ++ mixed[:, :qd].view(1, T, -1, K), ++ mixed[:, qd : 2 * qd].view(1, T, -1, K), ++ mixed[:, 2 * qd :].view(1, T, -1, V), ++ ) ++ got_out, got_replay = run(*views, out=core[:T].unsqueeze(0)) ++ assert got_out.data_ptr() == core.data_ptr() ++ assert torch.equal(core[:T].unsqueeze(0), want_out) ++ assert torch.equal(got_replay, want_replay) ++ ++ ++def test_verify_on_strided_views_with_variable_query_lengths(): ++ """Token-strided views with ragged windows (query lengths 4, 1, 3), as MTP ++ scheduling produces them.""" ++ torch.manual_seed(7) ++ H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6 ++ qsl_list = [0, 4, 5, 8] ++ T = qsl_list[-1] ++ qd, vd = H * K, HV * V ++ packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16) ++ mixed = packed[:, : 2 * qd + vd] ++ a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) ++ b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) ++ A_log = torch.randn(HV, device="cuda") ++ dt_bias = torch.randn(HV, device="cuda") ++ checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05 ++ qsl = torch.tensor(qsl_list, device="cuda", dtype=torch.int32) ++ si = torch.tensor([1, 3, 5], device="cuda", dtype=torch.int32) ++ ++ def run(q, k, v, out=None): ++ replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda") ++ o = gdn_recoverssm_verify( ++ A_log, ++ a, ++ b, ++ dt_bias, ++ q, ++ k, ++ v, ++ checkpoint_state=checkpoint, ++ replay_cache=replay, ++ query_start_loc=qsl, ++ state_indices=si, ++ spec_query_len=SQ, ++ out=out, ++ ) ++ return o, replay ++ ++ q, k, v = ( ++ x.contiguous().view(1, T, -1, d) ++ for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V)) ++ ) ++ want_out, want_replay = run(q, k, v) ++ core = torch.zeros(T + 2, HV, V, device="cuda", dtype=torch.bfloat16) ++ views = ( ++ mixed[:, :qd].view(1, T, -1, K), ++ mixed[:, qd : 2 * qd].view(1, T, -1, K), ++ mixed[:, 2 * qd :].view(1, T, -1, V), ++ ) ++ _, got_replay = run(*views, out=core[:T].unsqueeze(0)) ++ assert torch.equal(core[:T].unsqueeze(0), want_out) ++ assert torch.equal(got_replay, want_replay) ++ ++ ++def test_layer_verify_qkv_views_and_fallback(): ++ """The layer's guard: views for a [tokens, q|k|v] conv output with a contiguous ++ last dimension, the contiguous-copy fallback for any other layout; same values ++ either way.""" ++ from types import SimpleNamespace ++ ++ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( ++ QwenGatedDeltaNetAttention as Layer, ++ ) ++ ++ H, HV, K, V, T = 4, 8, 16, 16, 5 ++ layer = SimpleNamespace( ++ key_dim=H * K, value_dim=HV * V, tp_size=1, head_k_dim=K, head_v_dim=V ++ ) ++ layer.rearrange_mixed_qkv = lambda x: Layer.rearrange_mixed_qkv(cast(Any, layer), x) ++ width = 2 * H * K + HV * V ++ packed = torch.randn(T, width + HV * V, dtype=torch.bfloat16) # q|k|v|z, like qkvz ++ strided = packed[:, :width] ++ q, k, v, is_view = Layer._recoverssm_verify_qkv(cast(Any, layer), strided) ++ assert is_view ++ assert q.data_ptr() == packed.data_ptr() ++ rq, rk, rv = Layer.rearrange_mixed_qkv(cast(Any, layer), strided) ++ assert torch.equal(q, rq) and torch.equal(k, rk) and torch.equal(v, rv) ++ ++ # Column-major storage: last dimension not contiguous -> copies, same values. ++ col_major = strided.t().contiguous().t() ++ assert col_major.stride(-1) != 1 ++ q2, k2, v2, is_view2 = Layer._recoverssm_verify_qkv(cast(Any, layer), col_major) ++ assert not is_view2 ++ assert torch.equal(q2, rq) and torch.equal(k2, rk) and torch.equal(v2, rv) ++ ++ ++# ------------------------------------------------------------- BF16 recurrent state ++# RecoverSSM upcasts the BF16 checkpoint, carries the accepted window in FP32 and rounds ++# once when it writes the committed state. The reference is the FP32 trajectory from the ++# same (already BF16-rounded) checkpoint; the committed state must equal its BF16 ++# rounding up to one BF16 ulp, and be no further from it than the native path is. ++BF16_ULP = 2**-7 ++ ++ ++def _bf16_case(qlens, nb=12, seed=0): ++ inp = _make(qlens, nb=nb, seed=seed) ++ inp["checkpoint_state"] = inp["checkpoint_state"].bfloat16() ++ ref = dict(inp) ++ ref["checkpoint_state"] = inp["checkpoint_state"].float() ++ return inp, ref ++ ++ ++def test_verify_bf16_state_matches_native(): ++ inp, _ = _bf16_case(list(QLENS)) ++ inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") ++ before = inp["checkpoint_state"].clone() ++ out = _verify(inp) ++ ref, _ = _native(inp) ++ torch.testing.assert_close( ++ out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0 ++ ) ++ assert torch.equal(inp["checkpoint_state"], before) ++ ++ ++@pytest.mark.parametrize("accepted", [1, 2, 3, 4]) ++@pytest.mark.parametrize("seed", [0, 1]) ++def test_commit_bf16_state_rounds_once(accepted, seed): ++ inp, ref_inp = _bf16_case(list(QLENS), seed=seed) ++ idx = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") ++ inp["state_indices"] = ref_inp["state_indices"] = idx ++ _verify(inp) ++ _, native_states = _native(inp) ++ _, ref_states = _native(ref_inp) ++ ckpt = inp["checkpoint_state"].clone() ++ _context(inp, ckpt).commit( ++ torch.tensor( ++ [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" ++ ), ++ idx, ++ inp["query_start_loc"], ++ ) ++ for i, blk in enumerate(idx.tolist()): ++ tok = int(inp["query_start_loc"][i]) + min(accepted, QLENS[i]) - 1 ++ ref = ref_states[tok].float() ++ got = ckpt[blk].float() ++ torch.testing.assert_close( ++ got, ref.bfloat16().float(), rtol=BF16_ULP, atol=1e-5 ++ ) ++ err_rssm = (got - ref).abs().mean() ++ err_native = (native_states[tok].bfloat16().float() - ref).abs().mean() ++ assert err_rssm <= err_native * 1.05 + 1e-7, (err_rssm, err_native) ++ ++ ++@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES) ++def test_align_boundary_bf16_state(nc, accepted, bs): ++ inp, ref_inp = _bf16_case([T]) ++ bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6) ++ src = bt[0, nc // bs].reshape(1).contiguous() ++ inp["state_indices"] = ref_inp["state_indices"] = src ++ _verify(inp) ++ _, ref_states = _native(ref_inp) ++ ckpt = inp["checkpoint_state"].clone() ++ _context(inp, ckpt).commit( ++ torch.tensor([accepted], dtype=torch.int32, device="cuda"), ++ src, ++ inp["query_start_loc"], ++ block_table=bt, ++ num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"), ++ mamba_block_size=bs, ++ ) ++ next_boundary = (nc // bs + 1) * bs ++ final_blk = int(bt[0, (nc + accepted) // bs]) ++ if nc + accepted >= next_boundary: ++ torch.testing.assert_close( ++ ckpt[int(src)].float(), ++ ref_states[next_boundary - nc - 1].float().bfloat16().float(), ++ rtol=BF16_ULP, ++ atol=1e-5, ++ ) ++ torch.testing.assert_close( ++ ckpt[final_blk].float(), ++ ref_states[accepted - 1].float().bfloat16().float(), ++ rtol=BF16_ULP, ++ atol=1e-5, ++ ) +diff --git a/tests/kernels/mamba/test_recoverssm_ple.py b/tests/kernels/mamba/test_recoverssm_ple.py +new file mode 100644 +index 0000000..ec11ad3 +--- /dev/null ++++ b/tests/kernels/mamba/test_recoverssm_ple.py +@@ -0,0 +1,153 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""PLE short-conv RecoverSSM against the sequential reference. ++ ++Reference: the stock decode path fed the accepted tokens one at a time. ++RecoverSSM: one spec-mode conv over the whole window from the checkpoint ++(offset 0), then ``_PleConvCommit`` compacts the accepted suffix to the front. ++The conv history is pure data movement, so the states must be bit-identical, ++and the outputs of the accepted tokens must match as well. ++""" ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first ++from vllm.models.qwen4_exp.nvidia.ops.ple import ple_conv ++from vllm.v1.attention.backends.ple_recoverssm import _PleConvCommit ++ ++C, KSIZE, DIL, SQ, NB = 512, 4, 3, 4, 8 ++STATE_LEN = (KSIZE - 1) * DIL ++WIDTH = STATE_LEN + SQ - 1 ++DEV = "cuda" ++ ++ ++def _state(): ++ if is_conv_state_dim_first(): ++ raw = torch.zeros(NB, C, WIDTH, dtype=torch.bfloat16, device=DEV) ++ return raw, raw ++ raw = torch.zeros(NB, WIDTH, C, dtype=torch.bfloat16, device=DEV) ++ return raw, raw.transpose(-1, -2) ++ ++ ++def _i32(xs): ++ return torch.tensor(xs, dtype=torch.int32, device=DEV) ++ ++ ++def _setup(seed, blocks): ++ g = torch.Generator(device=DEV).manual_seed(seed) ++ w = torch.randn(C, KSIZE, generator=g, device=DEV).to(torch.bfloat16) ++ x = torch.randn(len(blocks) * SQ, C, generator=g, device=DEV).to(torch.bfloat16) ++ outer = torch.randn_like(x) ++ raw, view = _state() ++ for b in blocks: ++ view[b, :, :STATE_LEN] = torch.randn(C, STATE_LEN, generator=g, device=DEV).to( ++ torch.bfloat16 ++ ) ++ return w, x, outer, raw, view ++ ++ ++def _sequential(w, x, outer, view, block, first, count): ++ """Stock decode, one token at a time, on a copy of the checkpoint block.""" ++ ref = view.clone() ++ res = torch.zeros(count, C, dtype=torch.bfloat16, device=DEV) ++ for i in range(count): ++ t = first + i ++ ple_conv( ++ inputs=x[t : t + 1], ++ residual=res[i : i + 1], ++ conv_state=ref, ++ conv_weights=w, ++ state_indices=_i32([block]), ++ outer_residual=outer[t : t + 1], ++ mode="decode", ++ dilation=DIL, ++ ) ++ return ref[block, :, :STATE_LEN].clone(), res ++ ++ ++def _spec(w, x, outer, view, blocks): ++ res = torch.zeros(len(blocks) * SQ, C, dtype=torch.bfloat16, device=DEV) ++ ple_conv( ++ inputs=x, ++ residual=res, ++ conv_state=view, ++ conv_weights=w, ++ state_indices=_i32(blocks), ++ outer_residual=outer, ++ mode="spec", ++ dilation=DIL, ++ query_start_loc=_i32(list(range(0, len(blocks) * SQ + 1, SQ))), ++ num_accepted_tokens=_i32([1] * len(blocks)), ++ spec_query_len=SQ, ++ ) ++ return res ++ ++ ++@pytest.mark.parametrize("accepted", [1, 2, SQ]) ++@pytest.mark.parametrize("seed", [0, 1]) ++def test_ple_commit_matches_sequential(accepted, seed): ++ block = 3 ++ w, x, outer, raw, view = _setup(seed, [block]) ++ want_state, want_out = _sequential(w, x, outer, view, block, 0, accepted) ++ history = torch.cat([view[block, :, :STATE_LEN], x[:accepted].T], dim=1) ++ ++ got_out = _spec(w, x, outer, view, [block]) ++ ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) ++ ctx.commit(_i32([accepted]), _i32([block]), _i32([0, SQ])) ++ ++ got_state = view[block, :, :STATE_LEN] ++ assert torch.equal(got_state, want_state) ++ assert torch.equal(got_state, history[:, -STATE_LEN:]) ++ assert torch.equal(got_out[:accepted], want_out) ++ ++ ++@pytest.mark.parametrize( ++ ("num_computed", "accepted", "crosses"), ++ [(6, 3, True), (6, 2, True), (1, 2, False), (7, SQ, True)], ++) ++def test_ple_commit_align_boundary(num_computed, accepted, crosses): ++ block_size, src, nxt = 8, 2, 5 ++ w, x, outer, raw, view = _setup(7, [src]) ++ boundary_len = block_size - num_computed ++ want_final, _ = _sequential(w, x, outer, view, src, 0, accepted) ++ want_boundary, _ = _sequential(w, x, outer, view, src, 0, min(boundary_len, SQ)) ++ ++ _spec(w, x, outer, view, [src]) ++ ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) ++ ctx.commit( ++ _i32([accepted]), ++ _i32([src]), ++ _i32([0, SQ]), ++ block_table=_i32([[src, nxt]]), ++ num_computed_tokens=_i32([num_computed]), ++ mamba_block_size=block_size, ++ ) ++ ++ final_block = nxt if num_computed + accepted >= block_size else src ++ assert torch.equal(view[final_block, :, :STATE_LEN], want_final) ++ if crosses: ++ assert torch.equal(view[src, :, :STATE_LEN], want_boundary) ++ ++ ++def test_ple_commit_non_contiguous_request_rows(): ++ blocks, accepted = [1, 6], [3, 1] ++ w, x, outer, raw, view = _setup(11, blocks) ++ want = [ ++ _sequential(w, x, outer, view, b, i * SQ, a)[0] ++ for i, (b, a) in enumerate(zip(blocks, accepted)) ++ ] ++ ++ _spec(w, x, outer, view, blocks) ++ # Spec rows 0 and 1 belong to requests 1 and 3 of a four-request batch. ++ counts = _i32([9, accepted[0], 9, accepted[1]]) ++ ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) ++ ctx.commit( ++ counts, ++ _i32(blocks), ++ _i32([0, SQ, 2 * SQ]), ++ request_indices=_i32([1, 3]), ++ ) ++ ++ for b, ref in zip(blocks, want): ++ assert torch.equal(view[b, :, :STATE_LEN], ref) +diff --git a/vllm/config/cache.py b/vllm/config/cache.py +index 6dca830..e276dda 100644 +--- a/vllm/config/cache.py ++++ b/vllm/config/cache.py +@@ -202,7 +202,7 @@ class CacheConfig: + mode flushes are most efficient when mamba_block_size is a multiple of + replayssm_buffer_len, but this is not required.""" + use_recoverssm: bool = field(default=False, init=False) +- """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" ++ """Whether RecoverSSM speculative decode is used (Kimi-K3 KDA; Qwen4Exp GDN and PLE).""" + + # Will be set after profiling. + num_gpu_blocks: int | None = field(default=None, init=False) +diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py +index f11bea4..791975b 100644 +--- a/vllm/config/compilation.py ++++ b/vllm/config/compilation.py +@@ -1446,10 +1446,15 @@ class CompilationConfig: + f"with {min_cg_attn_backend} backend (support: " + f"{min_cg_support})" + ) +- if self.mode == CompilationMode.VLLM_COMPILE and ( +- self.splitting_ops_contain_attention() +- or self.use_inductor_graph_partition +- ): ++ if ( ++ self.mode == CompilationMode.VLLM_COMPILE ++ and ( ++ self.splitting_ops_contain_attention() ++ or self.use_inductor_graph_partition ++ ) ++ ) or envs.VLLM_USE_BREAKABLE_CUDAGRAPH: ++ # Breakable CUDA graphs break at the attention kernels at runtime, ++ # with or without torch.compile. + msg += ( + "; setting cudagraph_mode=PIECEWISE because " + "attention is compiled piecewise" +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index 67ba137..f215ab3 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -68,6 +68,13 @@ else: + + logger = init_logger(__name__) + ++# Architectures whose GDN (and PLE) layers run RecoverSSM under --use-replayssm. ++_QWEN4EXP_RECOVERSSM_ARCHS = ( ++ "Qwen4ExpForCausalLM", ++ "Qwen4ExpForConditionalGeneration", ++ "Qwen4ExpMTP", ++) ++ + # TODO(rocm): These models are either unsupported by MRV2 or slower with + # MRV2 on AMD GPUs. + ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset( +@@ -3218,21 +3225,35 @@ class VllmConfig: + @model_validator(mode="after") + def validate_mamba_cached_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm: +- self.cache_config.use_kda_recoverssm = False ++ self.cache_config.use_recoverssm = False + return self +- self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 ++ self.cache_config.use_recoverssm = self.num_speculative_tokens > 0 + + if self.model_config is not None and not self.model_config.supports_replayssm: + raise ValueError( + "--use-replayssm is not supported for architecture " + f"{self.model_config.architecture!r}" + ) +- if self.cache_config.use_kda_recoverssm: ++ if ( ++ self.model_config is not None ++ and self.model_config.architecture in _QWEN4EXP_RECOVERSSM_ARCHS ++ and not self.cache_config.use_recoverssm ++ ): ++ # Qwen4Exp's GDN and PLE layers have no plain ReplaySSM path; with them the ++ # flag selects RecoverSSM, which needs speculative decoding. ++ raise ValueError( ++ "--use-replayssm on Qwen4Exp selects GDN RecoverSSM and requires " ++ "speculative decoding (num_speculative_tokens > 0)" ++ ) ++ if self.cache_config.use_recoverssm: + if self.model_config is not None and self.model_config.architecture not in ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", ++ *_QWEN4EXP_RECOVERSSM_ARCHS, + ): +- raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA") ++ raise ValueError( ++ "RecoverSSM is only supported for Kimi-K3 KDA and Qwen4Exp GDN" ++ ) + if self.mamba_config.enable_stochastic_rounding: + raise ValueError( + "RecoverSSM supports bfloat16/float32 " +diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +index bbbb1c8..2d6b675 100644 +--- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py ++++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +@@ -368,10 +368,15 @@ class ChunkGatedDeltaRule(CustomOp): + + @PluggableLayer.register("qwen_gated_delta_net_attention") + class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): +- def get_state_shape( +- self, +- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: +- return MambaStateShapeCalculator.gated_delta_net_state_shape( ++ def _uses_gdn_recoverssm(self) -> bool: ++ """One predicate for the replay state's shape, dtype and the attention ++ backend.""" ++ from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm ++ ++ return uses_recoverssm(self.cache_config, self.num_spec) ++ ++ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: ++ shapes = MambaStateShapeCalculator.gated_delta_net_state_shape( + self.tp_size, + self.num_k_heads, + self.num_v_heads, +@@ -380,6 +385,32 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + self.conv_kernel_size, + self.num_spec, + ) ++ if self._uses_gdn_recoverssm(): ++ # per-token replay record [HV, spec_query_len, V + K + 1] fp32 ++ shapes = ( ++ *shapes, ++ ( ++ self.num_v_heads // self.tp_size, ++ self.num_spec + 1, ++ self.head_v_dim + self.head_k_dim + 1, ++ ), ++ ) ++ return shapes ++ ++ def get_state_dtype(self) -> tuple[torch.dtype, ...]: ++ dtypes = super().get_state_dtype() ++ if self._uses_gdn_recoverssm(): ++ dtypes = (*dtypes, torch.float32) ++ return dtypes ++ ++ def get_attn_backend(self): ++ if self._uses_gdn_recoverssm(): ++ from vllm.v1.attention.backends.gdn_recoverssm import ( ++ GDNRecoverSSMAttentionBackend, ++ ) ++ ++ return GDNRecoverSSMAttentionBackend ++ return super().get_attn_backend() + + def __init__( + self, +@@ -525,6 +556,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + self.gdn_decode_kernel = "CPU" + + self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" ++ # RecoverSSM verify replaces the per-draft snapshot path (Triton decode ++ # only) ++ self.use_gdn_recoverssm = self._uses_gdn_recoverssm() ++ if self.use_gdn_recoverssm: ++ self.gdn_decode_kernel = "triton" ++ self.enable_fused_gdn_decode = False ++ self.enable_packed_recurrent_decode = False ++ logger.info_once( ++ "GDN RecoverSSM speculative verify active (spec_query_len %d)", ++ self.num_spec + 1, ++ ) + logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) + + compilation_config = get_current_vllm_config().compilation_config +@@ -535,7 +577,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + def _fused_gdn_decode_unsupported_reason( + self, vllm_config: VllmConfig + ) -> str | None: +- conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() ++ conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2] + if ( + self.gqa_interleaved_layout + or self.head_k_dim != 128 +@@ -798,6 +840,25 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + + return mixed_qkv_out, z_out, b_out, a_out + ++ def _recoverssm_verify_qkv(self, mixed_qkv): ++ """Q/K/V for the RecoverSSM verify. ++ ++ The verify kernel reads token-strided heads, so for a 2-D [tokens, q|k|v] conv ++ output with a contiguous last dimension it gets views instead of copies. Any ++ other layout falls back to the contiguous copies of rearrange_mixed_qkv(). ++ Returns (query, key, value, is_view).""" ++ if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: ++ return (*self.rearrange_mixed_qkv(mixed_qkv), False) ++ q_dim = self.key_dim // self.tp_size ++ v_dim = self.value_dim // self.tp_size ++ num_tokens = mixed_qkv.shape[0] ++ query = mixed_qkv[:, :q_dim].view(1, num_tokens, -1, self.head_k_dim) ++ key = mixed_qkv[:, q_dim : 2 * q_dim].view(1, num_tokens, -1, self.head_k_dim) ++ value = mixed_qkv[:, 2 * q_dim : 2 * q_dim + v_dim].view( ++ 1, num_tokens, -1, self.head_v_dim ++ ) ++ return query, key, value, True ++ + def rearrange_mixed_qkv(self, mixed_qkv): + """Split packed qkv into contiguous (1, seq, heads, dim) tensors. + +@@ -1105,7 +1166,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + dtype = qkv_or_qkvz.dtype + num_k_heads = self.num_k_heads // self.tp_size + num_v_heads = self.num_v_heads // self.tp_size +- _, state_dtype = self.get_state_dtype() ++ state_dtype = self.get_state_dtype()[1] # 3 dtypes with the replay record + + # All kernels use BT = chunk_size, so a single pass with T = chunk_size + # is sufficient to populate every autotuner cache. Mirror the real +@@ -1365,7 +1426,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + ], + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, +- max_query_len=spec_state_indices_tensor.size(-1), ++ max_query_len=( ++ self.num_spec + 1 ++ if getattr(self, "use_gdn_recoverssm", False) ++ else spec_state_indices_tensor.size(-1) ++ ), + validate_data=False, + ) + +@@ -1402,7 +1467,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + else: + mixed_qkv_non_spec = None + +- query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) ++ if ( ++ spec_sequence_masks is not None ++ and mixed_qkv_spec is not None ++ and self.use_gdn_recoverssm ++ ): ++ query_spec, key_spec, value_spec, verify_views = ( ++ self._recoverssm_verify_qkv(mixed_qkv_spec) ++ ) ++ else: ++ verify_views = False ++ query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + + # Split mixed non-spec-decode+prefill to process independently + split_non_spec = ( +@@ -1465,7 +1540,43 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + # 2. Recurrent attention + + # 2.1: Process the multi-query part +- if spec_sequence_masks is not None: ++ if spec_sequence_masks is not None and getattr( ++ self, "use_gdn_recoverssm", False ++ ): ++ # verify the window off the checkpoint; the commit after sampling ++ # writes the state once ++ from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( ++ gdn_recoverssm_verify, ++ ) ++ ++ _n = attn_metadata.num_spec_decodes ++ assert spec_query_start_loc is not None ++ assert spec_state_indices_tensor is not None ++ core_attn_out_spec = gdn_recoverssm_verify( ++ self.A_log, ++ a_spec, ++ b_spec, ++ self.dt_bias, ++ query_spec, ++ key_spec, ++ value_spec, ++ checkpoint_state=ssm_state, ++ replay_cache=self_kv_cache[2], ++ query_start_loc=spec_query_start_loc[: _n + 1], ++ state_indices=spec_state_indices_tensor[:_n, 0], ++ spec_query_len=self.num_spec + 1, ++ use_qk_l2norm_in_kernel=True, ++ # spec-only batch: write straight into the layer output ++ out=( ++ core_attn_out[:num_actual_tokens].unsqueeze(0) ++ if verify_views ++ and mixed_qkv_non_spec is None ++ and query_spec.shape[1] == num_actual_tokens ++ else None ++ ), ++ ) ++ last_recurrent_state = None ++ elif spec_sequence_masks is not None: + core_attn_out_spec, last_recurrent_state = ( + fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, +@@ -1584,7 +1695,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) + elif spec_sequence_masks is not None: +- core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) ++ if core_attn_out_spec.data_ptr() != core_attn_out.data_ptr(): ++ core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + else: + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + +diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +new file mode 100644 +index 0000000..d347bef +--- /dev/null ++++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +@@ -0,0 +1,695 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""RecoverSSM for Gated DeltaNet (Qwen-style GDN): speculative verify from one ++checkpoint without per-draft state snapshots, and accepted-state recovery after ++sampling. ++ ++GDN port of the Kimi-K3 KDA RecoverSSM path ++(vllm/models/kimi_k3/nvidia/ops/recoverssm.py). Differences from KDA: the gate is a ++scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped ++over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested. ++The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t ++k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one ++fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 + ++sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the ++block-boundary state in align mode).""" ++ ++from collections.abc import Sequence ++from dataclasses import dataclass ++from typing import Any ++ ++import torch ++ ++from vllm.model_executor.layers.mamba.ops.recoverssm import ( ++ compact_conv_state_kernel, ++ prepare_commit_plan_kernel, ++) ++from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require ++from vllm.triton_utils import tl, triton ++from vllm.v1.attention.backends.utils import NULL_BLOCK_ID ++ ++ ++def _require(cond: bool, msg: str) -> None: ++ recoverssm_require(cond, msg, "GDN RecoverSSM") ++ ++ ++@triton.jit ++def _gdn_recoverssm_verify_kernel( ++ q_ptr, ++ k_ptr, ++ v_ptr, ++ a_ptr, ++ b_ptr, ++ A_log_ptr, ++ dt_bias_ptr, ++ state_ptr, ++ replay_ptr, ++ out_ptr, ++ query_start_loc_ptr, ++ state_indices_ptr, ++ scale, ++ softplus_beta, ++ softplus_threshold, ++ null_block_id, ++ stride_q_token, ++ stride_k_token, ++ stride_v_token, ++ stride_a_token, ++ stride_b_token, ++ stride_state_block, ++ stride_state_head, ++ stride_state_v, ++ stride_state_k, ++ stride_replay_block, ++ stride_replay_head, ++ stride_replay_pos, ++ stride_replay_dim, ++ stride_out_token, ++ stride_qsl, ++ stride_si, ++ H: tl.constexpr, ++ HV: tl.constexpr, ++ K: tl.constexpr, ++ V: tl.constexpr, ++ BK: tl.constexpr, ++ BV: tl.constexpr, ++ SPEC_QUERY_LEN: tl.constexpr, ++ USE_QK_L2NORM: tl.constexpr, ++): ++ pid_v = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ pid_hv = tl.program_id(2) ++ pid_h = pid_hv // (HV // H) ++ bos = tl.load(query_start_loc_ptr + pid_b * stride_qsl).to(tl.int64) ++ eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_qsl).to(tl.int64) ++ query_len = eos - bos ++ state_idx = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64) ++ offs_k = tl.arange(0, BK) ++ offs_v = pid_v * BV + tl.arange(0, BV) ++ mask_k = offs_k < K ++ mask_v = offs_v < V ++ mask_state = mask_v[:, None] & mask_k[None, :] ++ if state_idx <= null_block_id: ++ for t in tl.static_range(SPEC_QUERY_LEN): ++ tl.store( ++ out_ptr + (bos + t) * stride_out_token + pid_hv * V + offs_v, ++ tl.zeros([BV], dtype=tl.float32), ++ mask=(t < query_len) & mask_v, ++ ) ++ return ++ state = tl.load( ++ state_ptr ++ + state_idx * stride_state_block ++ + pid_hv * stride_state_head ++ + offs_v[:, None] * stride_state_v ++ + offs_k[None, :] * stride_state_k, ++ mask=mask_state, ++ other=0.0, ++ ).to(tl.float32) ++ neg_A = -tl.exp(tl.load(A_log_ptr + pid_hv).to(tl.float32)) ++ dt_bias = tl.load(dt_bias_ptr + pid_hv).to(tl.float32) ++ rec_base = ( ++ replay_ptr + state_idx * stride_replay_block + pid_hv * stride_replay_head ++ ) ++ for t in tl.static_range(SPEC_QUERY_LEN): ++ valid = t < query_len ++ tok = bos + t ++ q = tl.load( ++ q_ptr + tok * stride_q_token + pid_h * K + offs_k, ++ mask=valid & mask_k, ++ other=0.0, ++ ).to(tl.float32) ++ k = tl.load( ++ k_ptr + tok * stride_k_token + pid_h * K + offs_k, ++ mask=valid & mask_k, ++ other=0.0, ++ ).to(tl.float32) ++ v = tl.load( ++ v_ptr + tok * stride_v_token + pid_hv * V + offs_v, ++ mask=valid & mask_v, ++ other=0.0, ++ ).to(tl.float32) ++ x = ( ++ tl.load(a_ptr + tok * stride_a_token + pid_hv, mask=valid, other=0.0).to( ++ tl.float32 ++ ) ++ + dt_bias ++ ) ++ sp = tl.where( ++ softplus_beta * x <= softplus_threshold, ++ (1 / softplus_beta) * tl.log(1 + tl.exp(softplus_beta * x)), ++ x, ++ ) ++ g = neg_A * sp ++ beta = tl.sigmoid( ++ tl.load(b_ptr + tok * stride_b_token + pid_hv, mask=valid, other=0.0).to( ++ tl.float32 ++ ) ++ ) ++ if USE_QK_L2NORM: ++ q = q * tl.rsqrt(tl.sum(q * q) + 1e-6) ++ k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) ++ q = q * scale ++ new_state = state * tl.exp(g) ++ corr = (v - tl.sum(new_state * k[None, :], 1)) * beta ++ new_state = new_state + corr[:, None] * k[None, :] ++ state = tl.where(valid, new_state, state) ++ o = tl.sum(state * q[None, :], 1) ++ tl.store( ++ out_ptr + tok * stride_out_token + pid_hv * V + offs_v, ++ o.to(out_ptr.dtype.element_ty), ++ mask=valid & mask_v, ++ ) ++ rec = rec_base + t * stride_replay_pos ++ tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v) ++ if pid_v == 0: ++ tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k) ++ tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid) ++ ++ ++def gdn_recoverssm_verify( ++ A_log: torch.Tensor, ++ a: torch.Tensor, ++ b: torch.Tensor, ++ dt_bias: torch.Tensor, ++ q: torch.Tensor, ++ k: torch.Tensor, ++ v: torch.Tensor, ++ *, ++ checkpoint_state: torch.Tensor, ++ replay_cache: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ state_indices: torch.Tensor, ++ spec_query_len: int, ++ scale: float | None = None, ++ softplus_beta: float = 1.0, ++ softplus_threshold: float = 20.0, ++ use_qk_l2norm_in_kernel: bool = True, ++ out: torch.Tensor | None = None, ++) -> torch.Tensor: ++ """Verify a GDN speculative window against its checkpoint without modifying the ++ checkpoint. q, k: [1, T, H, K]; v: [1, T, HV, V]; a, b: [T, HV] (or [1, T, HV]); ++ checkpoint_state: [blocks, HV, V, K]; replay_cache: [blocks, HV, spec_query_len, V + ++ K + 1] fp32. Returns out [1, T, HV, V].""" ++ _require( ++ q.ndim == 4 and q.shape[0] == 1, "q must have shape [1, tokens, heads, dim]" ++ ) ++ _, total, H, K = q.shape ++ _require(k.shape == q.shape, "q and k shapes differ") ++ _require( ++ v.ndim == 4 and tuple(v.shape[:2]) == (1, total), ++ "v must have shape [1, tokens, value heads, dim]", ++ ) ++ HV, V = v.shape[-2], v.shape[-1] ++ _require( ++ H > 0 and HV >= H and HV % H == 0, ++ "value heads must be a positive multiple of the q/k heads", ++ ) ++ _require( ++ a.numel() == total * HV and b.numel() == total * HV, ++ "gate or beta shape is incompatible", ++ ) ++ a2 = a.reshape(total, HV) ++ b2 = b.reshape(total, HV) ++ _require( ++ a2.stride(1) == 1 and b2.stride(1) == 1, ++ "gate and beta heads must be contiguous", ++ ) ++ _require( ++ tuple(A_log.shape) == (HV,) and tuple(dt_bias.shape) == (HV,), ++ "A_log or dt_bias shape is incompatible", ++ ) ++ _require( ++ all(t_.stride(-1) == 1 and t_.stride(-2) == K for t_ in (q, k)), ++ "q and k heads must be contiguous", ++ ) ++ _require(v.stride(-1) == 1 and v.stride(-2) == V, "v heads must be contiguous") ++ _require(checkpoint_state.ndim == 4, "checkpoint must be four-dimensional") ++ nb = checkpoint_state.shape[0] ++ _require( ++ tuple(checkpoint_state.shape[1:]) == (HV, V, K), ++ "checkpoint shape is incompatible", ++ ) ++ _rec = (nb, HV, spec_query_len, V + K + 1) ++ _require(tuple(replay_cache.shape) == _rec, f"replay buffer needs shape {_rec}") ++ _require(replay_cache.dtype == torch.float32, "replay buffer must use float32") ++ _require(state_indices.ndim == 1, "state indices must be one-dimensional") ++ _require( ++ query_start_loc.ndim == 1 ++ and query_start_loc.shape[0] == state_indices.shape[0] + 1, ++ "query metadata is incompatible", ++ ) ++ _require( ++ total <= state_indices.shape[0] * spec_query_len, ++ "speculative decode input exceeds its activation capacity", ++ ) ++ _dev = q.device ++ _require( ++ all( ++ t_.device == _dev ++ for t_ in ( ++ k, ++ v, ++ a, ++ b, ++ A_log, ++ dt_bias, ++ checkpoint_state, ++ replay_cache, ++ query_start_loc, ++ state_indices, ++ ) ++ if t_ is not None ++ ) ++ and (out is None or out.device == _dev), ++ "inputs must be on the same device", ++ ) ++ if scale is None: ++ scale = K**-0.5 ++ if out is None: ++ out = torch.empty(1, total, HV, V, dtype=v.dtype, device=v.device) ++ _require( ++ tuple(out.shape) == (1, total, HV, V) and out.stride()[2:] == (V, 1), ++ "output shape or layout is incompatible", ++ ) ++ batch = state_indices.shape[0] ++ if total == 0 or batch == 0: ++ return out ++ BK = triton.next_power_of_2(K) ++ BV = min(triton.next_power_of_2(V), 32) ++ grid = (triton.cdiv(V, BV), batch, HV) ++ _gdn_recoverssm_verify_kernel[grid]( ++ q, ++ k, ++ v, ++ a2, ++ b2, ++ A_log, ++ dt_bias, ++ checkpoint_state, ++ replay_cache, ++ out, ++ query_start_loc, ++ state_indices, ++ scale, ++ softplus_beta, ++ softplus_threshold, ++ NULL_BLOCK_ID, ++ q.stride(1), ++ k.stride(1), ++ v.stride(1), ++ a2.stride(0), ++ b2.stride(0), ++ checkpoint_state.stride(0), ++ checkpoint_state.stride(1), ++ checkpoint_state.stride(2), ++ checkpoint_state.stride(3), ++ replay_cache.stride(0), ++ replay_cache.stride(1), ++ replay_cache.stride(2), ++ replay_cache.stride(3), ++ out.stride(1), ++ query_start_loc.stride(0), ++ state_indices.stride(0), ++ H=H, ++ HV=HV, ++ K=K, ++ V=V, ++ BK=BK, ++ BV=BV, ++ SPEC_QUERY_LEN=spec_query_len, ++ USE_QK_L2NORM=use_qk_l2norm_in_kernel, ++ num_warps=4, ++ num_stages=2, ++ ) ++ return out ++ ++ ++@triton.jit ++def _commit_gdn_state_kernel( ++ state_ref_ptr, ++ state_base_addrs_ptr, ++ state_block_strides_ptr, ++ replay_ref_ptr, ++ replay_base_addrs_ptr, ++ replay_block_strides_ptr, ++ state_indices_ptr, ++ commit_lens_ptr, ++ final_state_indices_ptr, ++ boundary_state_indices_ptr, ++ boundary_recovery_lens_ptr, ++ null_block_id, ++ stride_state_head, ++ stride_state_v, ++ stride_state_k, ++ stride_replay_head, ++ stride_replay_pos, ++ stride_replay_dim, ++ stride_si, ++ K: tl.constexpr, ++ V: tl.constexpr, ++ BK: tl.constexpr, ++ BV: tl.constexpr, ++ NUM_HEADS: tl.constexpr, ++ ALIGN_MODE: tl.constexpr, ++): ++ pid_v = tl.program_id(0) ++ pid_b = tl.program_id(1) ++ pid_lh = tl.program_id(2) ++ pid_l = pid_lh // NUM_HEADS ++ pid_h = pid_lh % NUM_HEADS ++ src = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64) ++ if src <= null_block_id: ++ return ++ n = tl.load(commit_lens_ptr + pid_b) ++ if n == 0: ++ return ++ final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) ++ boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) ++ boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b) ++ if final_idx <= null_block_id: ++ return ++ state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( ++ tl.pointer_type(state_ref_ptr.dtype.element_ty) ++ ) ++ sbs = tl.load(state_block_strides_ptr + pid_l) ++ replay_ptr = tl.load(replay_base_addrs_ptr + pid_l).to( ++ tl.pointer_type(replay_ref_ptr.dtype.element_ty) ++ ) ++ rbs = tl.load(replay_block_strides_ptr + pid_l) ++ rec_base = replay_ptr + src * rbs + pid_h * stride_replay_head ++ offs_k = tl.arange(0, BK) ++ offs_v = pid_v * BV + tl.arange(0, BV) ++ mask_k = offs_k < K ++ mask_v = offs_v < V ++ mask_state = mask_v[:, None] & mask_k[None, :] ++ sp = offs_v[:, None] * stride_state_v + offs_k[None, :] * stride_state_k ++ s0 = tl.load( ++ state_ptr + src * sbs + pid_h * stride_state_head + sp, ++ mask=mask_state, ++ other=0.0, ++ ).to(tl.float32) ++ fin_decay = 1.0 ++ fin_corr = tl.zeros([BV, BK], tl.float32) ++ bnd_decay = 1.0 ++ bnd_corr = tl.zeros([BV, BK], tl.float32) ++ for r in range(n): ++ t = n - r - 1 ++ rec = rec_base + t * stride_replay_pos ++ c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0) ++ kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0) ++ g = tl.load(rec + (V + K) * stride_replay_dim) ++ upd = c[:, None] * kk[None, :] ++ dec = tl.exp(g) ++ fin_corr += upd * fin_decay ++ fin_decay *= dec ++ if ALIGN_MODE: ++ before = t < boundary_len ++ bnd_corr += tl.where(before, upd * bnd_decay, 0.0) ++ bnd_decay *= tl.where(before, dec, 1.0) ++ if ALIGN_MODE: ++ bs = s0 * bnd_decay + bnd_corr ++ tl.store( ++ state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp, ++ bs.to(state_ref_ptr.dtype.element_ty), ++ mask=mask_state & (boundary_idx > null_block_id), ++ ) ++ fs = s0 * fin_decay + fin_corr ++ tl.store( ++ state_ptr + final_idx * sbs + pid_h * stride_state_head + sp, ++ fs.to(state_ref_ptr.dtype.element_ty), ++ mask=mask_state, ++ ) ++ ++ ++@dataclass ++class GDNRecoverSSMCommitContext: ++ conv_states: tuple[torch.Tensor, ...] ++ conv_state_base_addrs: torch.Tensor ++ conv_state_block_strides: torch.Tensor ++ conv_state_dim_strides: torch.Tensor ++ conv_state_token_strides: torch.Tensor ++ conv_history_len: int ++ checkpoints: tuple[torch.Tensor, ...] ++ state_base_addrs: torch.Tensor ++ state_block_strides: torch.Tensor ++ replays: tuple[torch.Tensor, ...] ++ replay_base_addrs: torch.Tensor ++ replay_block_strides: torch.Tensor ++ commit_lens: torch.Tensor ++ final_state_indices: torch.Tensor ++ boundary_state_indices: torch.Tensor ++ boundary_recovery_lens: torch.Tensor ++ spec_query_len: int ++ ++ @classmethod ++ def from_tensors( ++ cls, ++ conv_states: Sequence[torch.Tensor], ++ checkpoints: Sequence[torch.Tensor], ++ replays: Sequence[torch.Tensor], ++ *, ++ spec_query_len: int, ++ max_num_reqs: int, ++ conv_dim_first: bool = True, ++ ) -> "GDNRecoverSSMCommitContext": ++ _require(len(checkpoints) > 0, "commit requires at least one layer") ++ _require( ++ len(conv_states) == len(checkpoints) == len(replays), ++ "conv, state and replay lists differ", ++ ) ++ if not conv_dim_first: ++ conv_states = [s.transpose(-1, -2) for s in conv_states] ++ ref = checkpoints[0] ++ _require(ref.ndim == 4, "checkpoint must be four-dimensional") ++ nb, HV, V, K = ref.shape ++ _dev = ref.device ++ for s in checkpoints: ++ _require( ++ s.shape == ref.shape ++ and s.dtype == ref.dtype ++ and s.stride()[1:] == ref.stride()[1:] ++ and s.device == _dev, ++ "layers need matching checkpoints", ++ ) ++ _rec = (nb, HV, spec_query_len, V + K + 1) ++ for r in replays: ++ _require( ++ tuple(r.shape) == _rec ++ and r.dtype == torch.float32 ++ and r.stride()[1:] == replays[0].stride()[1:] ++ and r.device == _dev, ++ f"layers need matching float32 replay buffers of shape {_rec}", ++ ) ++ conv_ref = conv_states[0] ++ _require( ++ conv_ref.ndim == 3 and conv_ref.shape[0] == nb, ++ "conv state must be [blocks, dim, window] with the checkpoint's " ++ "block count", ++ ) ++ for c in conv_states: ++ _require( ++ c.shape == conv_ref.shape ++ and c.dtype == conv_ref.dtype ++ and c.stride()[1:] == conv_ref.stride()[1:] ++ and c.device == _dev, ++ "layers need matching conv states", ++ ) ++ conv_dim, conv_len = conv_ref.shape[1:] ++ hist = conv_len - spec_query_len + 1 ++ _require(hist > 0, "conv state is shorter than its window") ++ dev = ref.device ++ addr = lambda ts: torch.tensor( ++ [t.data_ptr() for t in ts], dtype=torch.int64, device=dev ++ ) ++ bstr = lambda ts: torch.tensor( ++ [t.stride(0) for t in ts], dtype=torch.int64, device=dev ++ ) ++ z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev) ++ return cls( ++ tuple(conv_states), ++ addr(conv_states), ++ bstr(conv_states), ++ torch.tensor( ++ [s.stride(1) for s in conv_states], dtype=torch.int64, device=dev ++ ), ++ torch.tensor( ++ [s.stride(2) for s in conv_states], dtype=torch.int64, device=dev ++ ), ++ hist, ++ tuple(checkpoints), ++ addr(checkpoints), ++ bstr(checkpoints), ++ tuple(replays), ++ addr(replays), ++ bstr(replays), ++ z(), ++ z(), ++ z(), ++ z(), ++ spec_query_len, ++ ) ++ ++ @classmethod ++ def create( ++ cls, layers: Sequence[Any], *, spec_query_len: int, max_num_reqs: int ++ ) -> "GDNRecoverSSMCommitContext": ++ from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first ++ ++ if any(len(layer.kv_cache) != 3 for layer in layers): ++ raise ValueError("GDN RecoverSSM pages must contain conv, state and replay") ++ return cls.from_tensors( ++ [layer.kv_cache[0] for layer in layers], ++ [layer.kv_cache[1] for layer in layers], ++ [layer.kv_cache[2] for layer in layers], ++ spec_query_len=spec_query_len, ++ max_num_reqs=max_num_reqs, ++ conv_dim_first=is_conv_state_dim_first(), ++ ) ++ ++ def commit( ++ self, ++ num_accepted_tokens: torch.Tensor, ++ state_indices: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ request_indices: torch.Tensor | None = None, ++ block_table: torch.Tensor | None = None, ++ num_computed_tokens: torch.Tensor | None = None, ++ mamba_block_size: int | None = None, ++ commit_conv: bool = True, ++ ) -> None: ++ """Fold accepted GDN and convolution inputs into every layer's checkpoint.""" ++ batch = state_indices.shape[0] ++ if batch == 0: ++ return ++ _require(state_indices.ndim == 1, "state indices must be one-dimensional") ++ _require( ++ batch <= self.commit_lens.shape[0], "commit batch exceeds its plan capacity" ++ ) ++ _require( ++ query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1, ++ "commit metadata is incompatible", ++ ) ++ _require( ++ request_indices is None or request_indices.shape[0] >= batch, ++ "request mapping is too short", ++ ) ++ _require( ++ num_accepted_tokens.ndim == 1 ++ and (request_indices is not None or num_accepted_tokens.shape[0] >= batch), ++ "accepted-token counts are too short", ++ ) ++ _align = (block_table, num_computed_tokens, mamba_block_size) ++ _require( ++ all(x is None for x in _align) or all(x is not None for x in _align), ++ "align metadata is incomplete", ++ ) ++ _require( ++ mamba_block_size is None or mamba_block_size >= self.spec_query_len, ++ "align block size must cover one speculative window", ++ ) ++ _require( ++ block_table is None or block_table.ndim == 2, ++ "block table must be two-dimensional", ++ ) ++ _dev = self.checkpoints[0].device ++ _require( ++ all( ++ t_ is None or t_.device == _dev ++ for t_ in ( ++ num_accepted_tokens, ++ state_indices, ++ query_start_loc, ++ request_indices, ++ block_table, ++ num_computed_tokens, ++ ) ++ ), ++ "commit inputs must be on the same device", ++ ) ++ bt_stride = (0, 0) if block_table is None else block_table.stride() ++ prepare_commit_plan_kernel[(batch,)]( ++ num_accepted_tokens, ++ request_indices, ++ state_indices, ++ query_start_loc, ++ block_table, ++ num_computed_tokens, ++ self.commit_lens, ++ self.final_state_indices, ++ self.boundary_state_indices, ++ self.boundary_recovery_lens, ++ NULL_BLOCK_ID, ++ mamba_block_size or 1, ++ block_table.shape[1] if block_table is not None else 1, ++ num_accepted_tokens.stride(0), ++ request_indices.stride(0) if request_indices is not None else 0, ++ state_indices.stride(0), ++ query_start_loc.stride(0), ++ bt_stride[0], ++ bt_stride[1], ++ 0 if num_computed_tokens is None else num_computed_tokens.stride(0), ++ SPEC_QUERY_LEN=self.spec_query_len, ++ num_warps=1, ++ ) ++ num_layers = len(self.checkpoints) ++ if commit_conv: ++ conv_ref = self.conv_states[0] ++ conv_dim = conv_ref.shape[1] ++ compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)]( ++ conv_ref, ++ self.conv_state_base_addrs, ++ self.conv_state_block_strides, ++ self.conv_state_dim_strides, ++ self.conv_state_token_strides, ++ state_indices, ++ self.commit_lens, ++ self.final_state_indices, ++ self.boundary_state_indices, ++ self.boundary_recovery_lens, ++ NULL_BLOCK_ID, ++ conv_dim, ++ self.conv_history_len, ++ state_indices.stride(0), ++ BLOCK_D=256, ++ BLOCK_HISTORY=triton.next_power_of_2(self.conv_history_len), ++ ALIGN_MODE=block_table is not None, ++ num_warps=4, ++ ) ++ ref = self.checkpoints[0] ++ _, HV, V, K = ref.shape ++ BK = triton.next_power_of_2(K) ++ BV = min(triton.next_power_of_2(V), 32) ++ rr = self.replays[0] ++ _commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)]( ++ ref, ++ self.state_base_addrs, ++ self.state_block_strides, ++ rr, ++ self.replay_base_addrs, ++ self.replay_block_strides, ++ state_indices, ++ self.commit_lens, ++ self.final_state_indices, ++ self.boundary_state_indices, ++ self.boundary_recovery_lens, ++ NULL_BLOCK_ID, ++ ref.stride(1), ++ ref.stride(2), ++ ref.stride(3), ++ rr.stride(1), ++ rr.stride(2), ++ rr.stride(3), ++ state_indices.stride(0), ++ K=K, ++ V=V, ++ BK=BK, ++ BV=BV, ++ NUM_HEADS=HV, ++ ALIGN_MODE=block_table is not None, ++ num_warps=4, ++ num_stages=2, ++ ) ++ ++ ++__all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"] +diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py +index 2e79001..90f7714 100644 +--- a/vllm/model_executor/models/interfaces.py ++++ b/vllm/model_executor/models/interfaces.py +@@ -85,6 +85,7 @@ MambaStateShapes: TypeAlias = ( + | tuple[tuple[int, int, int]] + | tuple[tuple[int, int], tuple[int, int]] + | tuple[tuple[int, int], tuple[int, int, int]] ++ | tuple[tuple[int, int], tuple[int, int, int], tuple[int, int, int]] + | tuple[ + tuple[int, int], + tuple[int, int, int], +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index 2b99063..fe416d4 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -37,11 +37,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, ++ MambaStateShapes, + MixtureOfExperts, + MultiModalEmbeddings, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, ++ SupportsReplaySSM, + _require_is_multimodal, + ) + from vllm.model_executor.models.qwen3_5 import ( +@@ -791,6 +793,7 @@ class Qwen4ExpForCausalLM( + SupportsPP, + Qwen4ExpMixtureOfExperts, + IsHybrid, ++ SupportsReplaySSM, + ): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], +@@ -911,20 +914,38 @@ class Qwen4ExpForCausalLM( + num_spec=num_spec, + ) + ++ @classmethod ++ def _uses_gdn_recoverssm(cls, vllm_config: VllmConfig) -> bool: ++ """One predicate for the replay record's shape AND dtype (they must ++ agree for the MambaSpec).""" ++ num_spec = ( ++ vllm_config.speculative_config.num_speculative_tokens ++ if vllm_config.speculative_config ++ else 0 ++ ) ++ from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm ++ ++ return uses_recoverssm(vllm_config.cache_config, num_spec or 0) ++ + @classmethod + def get_gdn_mamba_state_dtype_from_config( + cls, vllm_config: VllmConfig +- ) -> tuple[torch.dtype, torch.dtype]: +- return MambaStateDtypeCalculator.gated_delta_net_state_dtype( +- vllm_config.model_config.dtype, +- vllm_config.cache_config.mamba_cache_dtype, +- vllm_config.cache_config.mamba_ssm_cache_dtype, ++ ) -> tuple[torch.dtype, ...]: ++ dtypes: tuple[torch.dtype, ...] = ( ++ MambaStateDtypeCalculator.gated_delta_net_state_dtype( ++ vllm_config.model_config.dtype, ++ vllm_config.cache_config.mamba_cache_dtype, ++ vllm_config.cache_config.mamba_ssm_cache_dtype, ++ ) + ) ++ if cls._uses_gdn_recoverssm(vllm_config): # replay record, fp32 ++ dtypes = (*dtypes, torch.float32) ++ return dtypes + + @classmethod + def get_gdn_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig +- ) -> tuple[tuple[int, int], tuple[int, int]]: ++ ) -> MambaStateShapes: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size +@@ -933,7 +954,7 @@ class Qwen4ExpForCausalLM( + if vllm_config.speculative_config + else 0 + ) +- return MambaStateShapeCalculator.gated_delta_net_state_shape( ++ shapes = MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, +@@ -942,18 +963,30 @@ class Qwen4ExpForCausalLM( + hf_config.linear_conv_kernel_dim, + num_spec, + ) ++ if cls._uses_gdn_recoverssm( ++ vllm_config ++ ): # replay record [HV, spec_query_len, V + K + 1] ++ shapes = ( ++ *shapes, ++ ( ++ hf_config.linear_num_value_heads // tp_size, ++ num_spec + 1, ++ hf_config.linear_value_head_dim + hf_config.linear_key_head_dim + 1, ++ ), ++ ) ++ return shapes + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, +- ) -> tuple[torch.dtype, torch.dtype]: ++ ) -> tuple[torch.dtype, ...]: + return cls.get_gdn_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig +- ) -> tuple[tuple[int, int], tuple[int, int]]: ++ ) -> MambaStateShapes: + return cls.get_gdn_mamba_state_shape_from_config(vllm_config) + + @classmethod +@@ -1043,6 +1076,7 @@ class Qwen4ExpForConditionalGeneration( + Qwen3_5ForConditionalGeneration, + HasInnerState, + Qwen4ExpMixtureOfExperts, ++ SupportsReplaySSM, + ): + """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" + +@@ -1252,14 +1286,15 @@ class Qwen4ExpForConditionalGeneration( + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, +- ) -> tuple[torch.dtype, torch.dtype]: ++ ) -> tuple[torch.dtype, ...]: + return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) + + @classmethod +- def get_mamba_state_shape_from_config( ++ def get_mamba_state_shape_from_config( # type: ignore[override] + cls, + vllm_config: VllmConfig, +- ) -> tuple[tuple[int, int], tuple[int, int]]: ++ ) -> MambaStateShapes: ++ # GDN RecoverSSM adds a third state (the replay record) to Qwen3.5's two. + return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) + + @classmethod +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 5640474..76c0d7c 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -35,7 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + should_use_host_embed_table, + ) + from vllm.model_executor.model_loader.utils import configure_quant_config +-from vllm.model_executor.models.interfaces import SupportsPP ++from vllm.model_executor.models.interfaces import SupportsPP, SupportsReplaySSM + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + from vllm.model_executor.models.utils import ( + AutoWeightsLoader, +@@ -408,7 +408,9 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + "hidden_states": 0, + } + ) +-class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): ++class Qwen4ExpMTP( ++ nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts, SupportsReplaySSM ++): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index 4925e03..af8f528 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -203,6 +203,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + return True + + def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: ++ from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm ++ ++ if uses_recoverssm(self.cache_config, self.num_spec_tokens): ++ from vllm.v1.attention.backends.ple_recoverssm import ( ++ PleRecoverSSMAttentionBackend, ++ ) ++ ++ return PleRecoverSSMAttentionBackend + return PleShortConvAttentionBackend + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: +diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py +new file mode 100644 +index 0000000..a45142a +--- /dev/null ++++ b/vllm/v1/attention/backends/gdn_recoverssm.py +@@ -0,0 +1,181 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""GDN attention metadata for RecoverSSM speculative decode (Qwen-style Gated DeltaNet). ++ ++Wraps the generic GDN builder: every active decode row runs through the spec path (a ++draft-less step is a T=1 window), the spec state indices keep only the checkpoint ++column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv window ++already hold the accepted state), and the metadata carries what the post-sampling commit ++needs. Mirrors the Kimi-K3 KDA builder's RecoverSSM handling.""" ++ ++from dataclasses import dataclass, field, fields ++from typing import Any ++ ++import torch ++ ++from vllm.logger import init_logger ++from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata ++from vllm.v1.attention.backends.gdn_attn import ( ++ GDNAttentionBackend, ++ GDNAttentionMetadata, ++ GDNAttentionMetadataBuilder, ++) ++from vllm.v1.attention.backends.recoverssm_metadata import ( ++ RecoverSSMMetadata, ++ RecoverSSMPostprocessMetadata, ++ recoverssm_request_indices, ++ recoverssm_spec_rows, ++) ++ ++logger = init_logger(__name__) ++ ++ ++@dataclass ++class GDNRecoverSSMCommitMetadata: ++ state_indices: torch.Tensor # [num_spec_decodes, 1] ++ query_start_loc: torch.Tensor # [num_spec_decodes + 1] ++ request_indices: ( ++ torch.Tensor | None ++ ) # batch rows of the spec requests (None: rows 0..n-1) ++ block_table: torch.Tensor | None # align mode ++ num_computed_tokens: torch.Tensor | None ++ block_size: int | None ++ ++ ++@dataclass ++class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata): ++ recoverssm_commit: GDNRecoverSSMCommitMetadata | None = None ++ recoverssm_context: Any = field(default=None, repr=False, compare=False) ++ ++ def commit_recoverssm_state( ++ self, num_accepted_tokens: torch.Tensor ++ ) -> RecoverSSMPostprocessMetadata | None: ++ c = self.recoverssm_commit ++ if c is None: ++ return None ++ n = self.num_spec_decodes ++ self.recoverssm_context.commit( ++ num_accepted_tokens, ++ c.state_indices[:n, 0], ++ c.query_start_loc[: n + 1], ++ request_indices=c.request_indices, ++ block_table=c.block_table, ++ num_computed_tokens=c.num_computed_tokens, ++ mamba_block_size=c.block_size, ++ ) ++ if c.block_table is None: ++ return None ++ assert c.block_size is not None and c.num_computed_tokens is not None ++ return RecoverSSMPostprocessMetadata( ++ num_spec_decodes=n, ++ request_indices=c.request_indices, ++ block_table=c.block_table, ++ num_computed_tokens=c.num_computed_tokens, ++ block_size=c.block_size, ++ ) ++ ++ ++class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): ++ # The verify/commit metadata is not staged for FULL graph replay yet; declaring ++ # NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise). ++ _cudagraph_support = AttentionCGSupport.NEVER ++ ++ def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: ++ super().__init__(kv_cache_spec, layer_names, vllm_config, device) ++ # The builder may be created before the runner resolves the CUDA graph mode. ++ self.use_full_cuda_graph = False ++ max_reqs = vllm_config.scheduler_config.max_num_seqs ++ self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device) ++ self._context: Any = None ++ self._logged = False ++ ++ def _get_context(self): ++ if self._context is None: ++ from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( ++ GDNRecoverSSMCommitContext, ++ ) ++ ++ fc = self.vllm_config.compilation_config.static_forward_context ++ layers = [fc[name] for name in self.layer_names] ++ self._context = GDNRecoverSSMCommitContext.create( ++ layers, ++ spec_query_len=1 + self.num_spec, ++ max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, ++ ) ++ logger.info( ++ "GDN RecoverSSM commit context: %d layers, spec_query_len %d", ++ len(layers), ++ 1 + self.num_spec, ++ ) ++ return self._context ++ ++ def build( # type: ignore[override] ++ self, ++ common_prefix_len: int, ++ common_attn_metadata: CommonAttentionMetadata, ++ num_accepted_tokens: torch.Tensor | None = None, ++ num_decode_draft_tokens_cpu: torch.Tensor | None = None, ++ fast_build: bool = False, ++ ) -> "GDNRecoverSSMMetadata": ++ m = common_attn_metadata ++ rows = None ++ if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: ++ if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": ++ raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask") ++ # Positive counts only defeat the base builder's "no drafts -> regular ++ # decode" shortcut: RecoverSSM keeps every post-prefill row on the spec path ++ # (extended conv window). ++ num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows( ++ m.is_prefilling, ++ m.query_start_loc_cpu, ++ num_decode_draft_tokens_cpu, ++ m.num_reqs, ++ ) ++ meta = super().build( ++ common_prefix_len, ++ m, ++ num_accepted_tokens=num_accepted_tokens, ++ num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, ++ fast_build=fast_build, ++ ) ++ base = {f.name: getattr(meta, f.name) for f in fields(meta)} ++ commit = None ++ if meta.num_spec_decodes > 0: ++ n = meta.num_spec_decodes ++ assert meta.spec_state_indices_tensor is not None ++ base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1] ++ base["num_accepted_tokens"] = self._ones[:n] ++ if rows is None: ++ raise ValueError("GDN RecoverSSM: spec rows without draft counts") ++ request_indices = recoverssm_request_indices( ++ rows, n, m.query_start_loc.device ++ ) ++ align = self.vllm_config.cache_config.mamba_cache_mode == "align" ++ commit = GDNRecoverSSMCommitMetadata( ++ state_indices=base["spec_state_indices_tensor"], ++ query_start_loc=meta.spec_query_start_loc, ++ request_indices=request_indices, ++ block_table=m.block_table_tensor if align else None, ++ num_computed_tokens=m.compute_num_computed_tokens() if align else None, ++ block_size=self.kv_cache_spec.block_size if align else None, ++ ) ++ if not self._logged: ++ self._logged = True ++ logger.info( ++ "GDN RecoverSSM path taken: %d spec rows, align=%s", n, align ++ ) ++ return GDNRecoverSSMMetadata( ++ **base, ++ recoverssm_commit=commit, ++ recoverssm_context=self._get_context() if commit is not None else None, ++ ) ++ ++ ++class GDNRecoverSSMAttentionBackend(GDNAttentionBackend): ++ @staticmethod ++ def get_name() -> str: ++ return "GDN_RECOVERSSM" ++ ++ @staticmethod ++ def get_builder_cls() -> type[GDNRecoverSSMMetadataBuilder]: ++ return GDNRecoverSSMMetadataBuilder +diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py +index 817773a..29352b3 100644 +--- a/vllm/v1/attention/backends/mamba_attn.py ++++ b/vllm/v1/attention/backends/mamba_attn.py +@@ -110,7 +110,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): + self.compilation_config = vllm_config.compilation_config + self.num_spec_tokens: int = vllm_config.num_speculative_tokens + self.use_spec_decode = self.num_spec_tokens > 0 +- self.use_replayssm = vllm_config.cache_config.use_replayssm ++ # RecoverSSM also rides on --use-replayssm but keeps its own metadata; the ++ # Mamba2 ReplaySSM ring-buffer workspaces below are not used then. ++ self.use_replayssm = ( ++ vllm_config.cache_config.use_replayssm ++ and not vllm_config.cache_config.use_recoverssm ++ ) + self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len + self.use_flashinfer_replayssm = ( + self.use_replayssm +diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py +new file mode 100644 +index 0000000..7402001 +--- /dev/null ++++ b/vllm/v1/attention/backends/ple_recoverssm.py +@@ -0,0 +1,326 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""RecoverSSM protocol for the Qwen4Exp PLE short-conv layer. ++ ++When the GDN layers run RecoverSSM, the runner resets the shared num_accepted_tokens to ++1 after each commit (align mode). The PLE dilated short conv keeps an extended window in ++one state block and reads it at num_accepted - 1, so it must follow the same protocol: ++every active decode row takes the spec path, the conv reads at offset 0, and after ++sampling the accepted window is compacted to the front (the shared RecoverSSM compaction ++kernel; in align mode it also writes the block-boundary window).""" ++ ++from dataclasses import dataclass, field, fields ++from typing import Any ++ ++import torch ++ ++from vllm.logger import init_logger ++from vllm.model_executor.layers.mamba.ops.recoverssm import ( ++ compact_conv_state_kernel, ++ prepare_commit_plan_kernel, ++) ++from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require ++from vllm.triton_utils import triton ++from vllm.v1.attention.backend import AttentionCGSupport ++from vllm.v1.attention.backends.recoverssm_metadata import ( ++ RecoverSSMMetadata, ++ RecoverSSMPostprocessMetadata, ++ recoverssm_request_indices, ++ recoverssm_spec_rows, ++) ++from vllm.v1.attention.backends.short_conv_attn import ( ++ PleShortConvAttentionBackend, ++ PleShortConvAttentionMetadata, ++ PleShortConvAttentionMetadataBuilder, ++) ++from vllm.v1.attention.backends.utils import NULL_BLOCK_ID ++ ++logger = init_logger(__name__) ++ ++ ++def _require(cond: bool, msg: str) -> None: ++ recoverssm_require(cond, msg, "PLE RecoverSSM") ++ ++ ++class _PleConvCommit: ++ def __init__(self, conv_states, spec_query_len: int, max_num_reqs: int): ++ from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first ++ ++ _require(len(conv_states) > 0, "PLE commit requires at least one layer") ++ if not is_conv_state_dim_first(): ++ conv_states = [ ++ s.transpose(-1, -2) for s in conv_states ++ ] # -> [blocks, C, W] ++ self.conv_states = list(conv_states) ++ ref = self.conv_states[0] ++ _require(ref.ndim == 3, "PLE conv state must be [blocks, dim, window]") ++ dev = ref.device ++ for s in self.conv_states: ++ _require( ++ s.shape == ref.shape ++ and s.dtype == ref.dtype ++ and s.stride()[1:] == ref.stride()[1:] ++ and s.device == dev, ++ "PLE layers need matching conv states", ++ ) ++ self.conv_dim, conv_len = ref.shape[1], ref.shape[2] ++ self.hist = conv_len - spec_query_len + 1 ++ _require(self.hist > 0, "PLE conv state is shorter than its window") ++ self.spec_query_len = spec_query_len ++ t = lambda xs: torch.tensor(xs, dtype=torch.int64, device=dev) ++ self.base = t([s.data_ptr() for s in self.conv_states]) ++ self.bstride = t([s.stride(0) for s in self.conv_states]) ++ self.dstride = t([s.stride(1) for s in self.conv_states]) ++ self.tstride = t([s.stride(2) for s in self.conv_states]) ++ z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev) ++ self.commit_lens, self.final_idx, self.bnd_idx, self.bnd_len = ( ++ z(), ++ z(), ++ z(), ++ z(), ++ ) ++ ++ def commit( ++ self, ++ num_accepted_tokens, ++ state_indices, ++ query_start_loc, ++ request_indices=None, ++ block_table=None, ++ num_computed_tokens=None, ++ mamba_block_size=None, ++ ): ++ batch = state_indices.shape[0] ++ if batch == 0: ++ return ++ _require(state_indices.ndim == 1, "state indices must be one-dimensional") ++ _require( ++ batch <= self.commit_lens.shape[0], ++ "PLE commit batch exceeds its plan capacity", ++ ) ++ _require( ++ query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1, ++ "PLE commit metadata is incompatible", ++ ) ++ _require( ++ request_indices is None or request_indices.shape[0] >= batch, ++ "PLE request mapping is too short", ++ ) ++ _require( ++ num_accepted_tokens.ndim == 1 ++ and (request_indices is not None or num_accepted_tokens.shape[0] >= batch), ++ "PLE accepted-token counts are too short", ++ ) ++ _align = (block_table, num_computed_tokens, mamba_block_size) ++ _require( ++ all(x is None for x in _align) or all(x is not None for x in _align), ++ "PLE align metadata is incomplete", ++ ) ++ _require( ++ mamba_block_size is None or mamba_block_size >= self.spec_query_len, ++ "PLE align block size must cover one speculative window", ++ ) ++ _require( ++ block_table is None or block_table.ndim == 2, ++ "PLE block table must be two-dimensional", ++ ) ++ _dev = self.conv_states[0].device ++ _require( ++ all( ++ t_ is None or t_.device == _dev ++ for t_ in ( ++ num_accepted_tokens, ++ state_indices, ++ query_start_loc, ++ request_indices, ++ block_table, ++ num_computed_tokens, ++ ) ++ ), ++ "PLE commit inputs must be on the same device", ++ ) ++ bt = (0, 0) if block_table is None else block_table.stride() ++ prepare_commit_plan_kernel[(batch,)]( ++ num_accepted_tokens, ++ request_indices, ++ state_indices, ++ query_start_loc, ++ block_table, ++ num_computed_tokens, ++ self.commit_lens, ++ self.final_idx, ++ self.bnd_idx, ++ self.bnd_len, ++ NULL_BLOCK_ID, ++ mamba_block_size or 1, ++ block_table.shape[1] if block_table is not None else 1, ++ num_accepted_tokens.stride(0), ++ request_indices.stride(0) if request_indices is not None else 0, ++ state_indices.stride(0), ++ query_start_loc.stride(0), ++ bt[0], ++ bt[1], ++ 0 if num_computed_tokens is None else num_computed_tokens.stride(0), ++ SPEC_QUERY_LEN=self.spec_query_len, ++ num_warps=1, ++ ) ++ compact_conv_state_kernel[ ++ (triton.cdiv(self.conv_dim, 256), batch, len(self.conv_states)) ++ ]( ++ self.conv_states[0], ++ self.base, ++ self.bstride, ++ self.dstride, ++ self.tstride, ++ state_indices, ++ self.commit_lens, ++ self.final_idx, ++ self.bnd_idx, ++ self.bnd_len, ++ NULL_BLOCK_ID, ++ self.conv_dim, ++ self.hist, ++ state_indices.stride(0), ++ BLOCK_D=256, ++ BLOCK_HISTORY=triton.next_power_of_2(self.hist), ++ ALIGN_MODE=block_table is not None, ++ num_warps=4, ++ ) ++ ++ ++@dataclass ++class PleRecoverSSMCommitMetadata: ++ state_indices: torch.Tensor # [n] ++ query_start_loc: torch.Tensor # [n + 1] ++ request_indices: torch.Tensor | None ++ block_table: torch.Tensor | None ++ num_computed_tokens: torch.Tensor | None ++ block_size: int | None ++ ++ ++@dataclass ++class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata): ++ recoverssm_commit: PleRecoverSSMCommitMetadata | None = None ++ recoverssm_context: Any = field(default=None, repr=False, compare=False) ++ ++ def commit_recoverssm_state(self, num_accepted_tokens): ++ c = self.recoverssm_commit ++ if c is None: ++ return None ++ n = self.num_spec_decodes ++ self.recoverssm_context.commit( ++ num_accepted_tokens, ++ c.state_indices[:n], ++ c.query_start_loc[: n + 1], ++ request_indices=c.request_indices, ++ block_table=c.block_table, ++ num_computed_tokens=c.num_computed_tokens, ++ mamba_block_size=c.block_size, ++ ) ++ if c.block_table is None: ++ return None ++ assert c.block_size is not None and c.num_computed_tokens is not None ++ return RecoverSSMPostprocessMetadata( ++ num_spec_decodes=n, ++ request_indices=c.request_indices, ++ block_table=c.block_table, ++ num_computed_tokens=c.num_computed_tokens, ++ block_size=c.block_size, ++ ) ++ ++ ++class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): ++ # Same as the GDN RecoverSSM builder: PIECEWISE only for now. ++ _cudagraph_support = AttentionCGSupport.NEVER ++ ++ def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: ++ super().__init__(kv_cache_spec, layer_names, vllm_config, device) ++ self.use_full_cuda_graph = False ++ self._ones = torch.ones( ++ vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device ++ ) ++ self._ctx: _PleConvCommit | None = None ++ self._logged = False ++ ++ def _get_ctx(self): ++ if self._ctx is None: ++ fc = self.vllm_config.compilation_config.static_forward_context ++ convs = [fc[name].kv_cache[0] for name in self.layer_names] ++ ctx = _PleConvCommit( ++ convs, 1 + self.num_spec, self.vllm_config.scheduler_config.max_num_seqs ++ ) ++ self._ctx = ctx ++ logger.info( ++ "PLE RecoverSSM commit context: %d layers, history %d", ++ len(convs), ++ ctx.hist, ++ ) ++ return self._ctx ++ ++ def build( ++ self, ++ common_prefix_len, ++ common_attn_metadata, ++ fast_build=False, ++ *, ++ num_accepted_tokens=None, ++ num_decode_draft_tokens_cpu=None, ++ **kwargs, ++ ): # type: ignore[override] ++ m = common_attn_metadata ++ rows = None ++ if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: ++ if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": ++ raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask") ++ # the same classification as the GDN builder, so both commits map the same ++ # requests ++ num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows( ++ m.is_prefilling, ++ m.query_start_loc_cpu, ++ num_decode_draft_tokens_cpu, ++ m.num_reqs, ++ ) ++ meta = super().build( ++ common_prefix_len, ++ m, ++ fast_build, ++ num_accepted_tokens=num_accepted_tokens, ++ num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, ++ **kwargs, ++ ) ++ if ( ++ not isinstance(meta, PleShortConvAttentionMetadata) ++ or meta.num_spec_decodes == 0 ++ ): ++ return meta ++ base = {f.name: getattr(meta, f.name) for f in fields(meta)} ++ n = meta.num_spec_decodes ++ base["num_accepted_tokens"] = self._ones[:n] ++ if rows is None: ++ raise ValueError("PLE RecoverSSM: spec rows without draft counts") ++ req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device) ++ align = self.vllm_config.cache_config.mamba_cache_mode == "align" ++ commit = PleRecoverSSMCommitMetadata( ++ state_indices=meta.spec_state_indices_tensor, ++ query_start_loc=meta.spec_query_start_loc, ++ request_indices=req_idx, ++ block_table=m.block_table_tensor if align else None, ++ num_computed_tokens=m.compute_num_computed_tokens() if align else None, ++ block_size=self.kv_cache_spec.block_size if align else None, ++ ) ++ if not self._logged: ++ self._logged = True ++ logger.info("PLE RecoverSSM path taken: %d spec rows, align=%s", n, align) ++ return PleRecoverSSMMetadata( ++ **base, recoverssm_commit=commit, recoverssm_context=self._get_ctx() ++ ) ++ ++ ++class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend): ++ @staticmethod ++ def get_name() -> str: ++ return "PLE_SHORT_CONV_RECOVERSSM" ++ ++ @staticmethod ++ def get_builder_cls() -> type[PleRecoverSSMMetadataBuilder]: ++ return PleRecoverSSMMetadataBuilder +-- +2.43.0 + diff --git a/patches/0073-recoverssm-align-keep-a-boundary-ending-state-in-its.patch b/patches/0073-recoverssm-align-keep-a-boundary-ending-state-in-its.patch new file mode 100644 index 0000000..5d42e29 --- /dev/null +++ b/patches/0073-recoverssm-align-keep-a-boundary-ending-state-in-its.patch @@ -0,0 +1,244 @@ +From 57a0d446925c54448a4bfbc7879d6556a77101cc Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 27 Sep 2026 11:06:54 +0200 +Subject: [PATCH 073/103] recoverssm: align: keep a boundary-ending state in + its block (port of upstream #58863, 5567cc1b25) + +Port of the upstream vLLM change #58863 (Apache-2.0). A window whose accepted tokens end exactly on a block boundary now commits into +the block it ends in, not the next, possibly unallocated, column. +--- + tests/kernels/mamba/test_recoverssm_gdn.py | 9 +- + tests/kernels/mamba/test_recoverssm_ple.py | 2 +- + .../kimi_k3/test_recoverssm_boundary.py | 129 ++++++++++++++++++ + .../layers/mamba/ops/recoverssm.py | 6 +- + vllm/v1/worker/gpu/model_states/recoverssm.py | 4 +- + 5 files changed, 143 insertions(+), 7 deletions(-) + create mode 100644 tests/models/kimi_k3/test_recoverssm_boundary.py + +diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py +index 1e0cbcf..ce16f79 100644 +--- a/tests/kernels/mamba/test_recoverssm_gdn.py ++++ b/tests/kernels/mamba/test_recoverssm_gdn.py +@@ -299,14 +299,15 @@ def test_commit_maps_noncontiguous_requests_in_align_mode(): + for i, r in enumerate(req_rows.tolist()): + n = min(int(acc[r]), qlens[i]) + c0 = int(nc[r]) +- final_blk = int(bt[r, (c0 + n) // bs]) ++ final_blk = int(bt[r, max(c0 + n - 1, 0) // bs]) + tok = int(inp["query_start_loc"][i]) + n - 1 + torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6) + + + # (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on + # it / crossing and continuing / accepted = 1 without crossing / the full window +-# crossing after 3 ++# crossing after 3. A window ending exactly on a boundary keeps its state in the block ++# holding its last token: the next block may still be unallocated. + ALIGN_CASES = [ + (3, 2, 4), + (2, 3, 4), +@@ -338,7 +339,7 @@ def test_align_boundary_and_final_state_match_native(nc, accepted, bs): + mamba_block_size=bs, + ) + next_boundary = (nc // bs + 1) * bs +- final_blk = int(bt[0, (nc + accepted) // bs]) ++ final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs]) + if ( + nc + accepted >= next_boundary + ): # crosses: the boundary state goes into the source block +@@ -585,7 +586,7 @@ def test_align_boundary_bf16_state(nc, accepted, bs): + mamba_block_size=bs, + ) + next_boundary = (nc // bs + 1) * bs +- final_blk = int(bt[0, (nc + accepted) // bs]) ++ final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs]) + if nc + accepted >= next_boundary: + torch.testing.assert_close( + ckpt[int(src)].float(), +diff --git a/tests/kernels/mamba/test_recoverssm_ple.py b/tests/kernels/mamba/test_recoverssm_ple.py +index ec11ad3..6e2e1f0 100644 +--- a/tests/kernels/mamba/test_recoverssm_ple.py ++++ b/tests/kernels/mamba/test_recoverssm_ple.py +@@ -124,7 +124,7 @@ def test_ple_commit_align_boundary(num_computed, accepted, crosses): + mamba_block_size=block_size, + ) + +- final_block = nxt if num_computed + accepted >= block_size else src ++ final_block = nxt if num_computed + accepted > block_size else src + assert torch.equal(view[final_block, :, :STATE_LEN], want_final) + if crosses: + assert torch.equal(view[src, :, :STATE_LEN], want_boundary) +diff --git a/tests/models/kimi_k3/test_recoverssm_boundary.py b/tests/models/kimi_k3/test_recoverssm_boundary.py +new file mode 100644 +index 0000000..8023e44 +--- /dev/null ++++ b/tests/models/kimi_k3/test_recoverssm_boundary.py +@@ -0,0 +1,129 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Align-mode RecoverSSM state columns at recurrent block boundaries. ++ ++A speculative window whose accepted tokens end exactly on a recurrent block ++boundary must keep its state in the block holding the last committed token. ++The following block can still be unallocated, so selecting it would drop the ++commit while the bookkeeping recorded the state as advanced. ++""" ++ ++import pytest ++import torch ++ ++from vllm.models.kimi_k3.nvidia.ops.recoverssm import _prepare_commit_plan_kernel ++from vllm.platforms import current_platform ++from vllm.v1.attention.backends.utils import NULL_BLOCK_ID ++from vllm.v1.worker.gpu.model_states.recoverssm import ( ++ _postprocess_recoverssm_align_kernel, ++) ++ ++pytestmark = pytest.mark.skipif( ++ not current_platform.is_cuda_alike(), ++ reason="The RecoverSSM kernels require a CUDA-alike device.", ++) ++ ++BLOCK_SIZE = 2048 ++SPEC_QUERY_LEN = 4 ++BLOCKS = [100, 101, 102, 103, 104, 105] ++ ++# (computed before the step, committed tokens, allocated block columns, ++# final column, boundary column or None) ++CASES = [ ++ # 8,188 + 4 ends exactly on the 8,192 boundary; column 4 is unallocated. ++ (8188, 4, 4, 3, 3), ++ # The window crosses the boundary: state moves on, checkpoint stays behind. ++ (8190, 4, 5, 4, 3), ++ # The window stays inside column 3. ++ (8180, 4, 4, 3, None), ++ # The first token after a boundary belongs to the new column. ++ (8192, 1, 5, 4, None), ++ # Single-token commit ending on the first boundary. ++ (2047, 1, 1, 0, 0), ++] ++ ++ ++@pytest.mark.parametrize( ++ "num_computed,commit_len,allocated,final_col,boundary_col", CASES ++) ++def test_commit_plan_keeps_boundary_state_in_last_computed_block( ++ num_computed, commit_len, allocated, final_col, boundary_col ++): ++ device = current_platform.device_type ++ columns = [ ++ block if col < allocated else NULL_BLOCK_ID for col, block in enumerate(BLOCKS) ++ ] ++ block_table = torch.tensor([columns], dtype=torch.int32, device=device) ++ num_accepted = torch.tensor([commit_len], dtype=torch.int32, device=device) ++ state_indices = torch.tensor( ++ [BLOCKS[(num_computed - 1) // BLOCK_SIZE]], dtype=torch.int32, device=device ++ ) ++ query_start_loc = torch.tensor( ++ [0, SPEC_QUERY_LEN], dtype=torch.int32, device=device ++ ) ++ computed = torch.tensor([num_computed], dtype=torch.int32, device=device) ++ commit_lens, final_indices, boundary_indices, recovery_lens = ( ++ torch.full((1,), -7, dtype=torch.int32, device=device) for _ in range(4) ++ ) ++ ++ _prepare_commit_plan_kernel[(1,)]( ++ num_accepted, ++ None, ++ state_indices, ++ query_start_loc, ++ block_table, ++ computed, ++ commit_lens, ++ final_indices, ++ boundary_indices, ++ recovery_lens, ++ NULL_BLOCK_ID, ++ BLOCK_SIZE, ++ block_table.shape[1], ++ num_accepted.stride(0), ++ 0, ++ state_indices.stride(0), ++ query_start_loc.stride(0), ++ block_table.stride(0), ++ block_table.stride(1), ++ computed.stride(0), ++ SPEC_QUERY_LEN=SPEC_QUERY_LEN, ++ num_warps=1, ++ ) ++ ++ assert commit_lens.item() == commit_len ++ assert final_indices.item() == BLOCKS[final_col] ++ expected_boundary = NULL_BLOCK_ID if boundary_col is None else BLOCKS[boundary_col] ++ assert boundary_indices.item() == expected_boundary ++ if boundary_col is not None: ++ next_boundary = (num_computed // BLOCK_SIZE + 1) * BLOCK_SIZE ++ assert recovery_lens.item() == next_boundary - num_computed ++ ++ ++@pytest.mark.parametrize( ++ "num_computed,commit_len,allocated,final_col,boundary_col", CASES ++) ++def test_postprocess_records_the_committed_state_column( ++ num_computed, commit_len, allocated, final_col, boundary_col ++): ++ device = current_platform.device_type ++ idx_mapping = torch.tensor([2], dtype=torch.int32, device=device) ++ num_sampled = torch.tensor([commit_len], dtype=torch.int32, device=device) ++ computed = torch.tensor([num_computed], dtype=torch.int32, device=device) ++ state_idx = torch.full((4,), -7, dtype=torch.int32, device=device) ++ num_accepted = torch.full((4,), -7, dtype=torch.int32, device=device) ++ ++ _postprocess_recoverssm_align_kernel[(1,)]( ++ idx_mapping, ++ num_sampled, ++ None, ++ computed, ++ state_idx, ++ num_accepted, ++ MAMBA_BLOCK_SIZE=BLOCK_SIZE, ++ BLOCK_TABLE_WIDTH=len(BLOCKS), ++ ) ++ ++ assert state_idx[2].item() == final_col ++ assert num_accepted[2].item() == 1 ++ assert state_idx[[0, 1, 3]].tolist() == [-7, -7, -7] +diff --git a/vllm/model_executor/layers/mamba/ops/recoverssm.py b/vllm/model_executor/layers/mamba/ops/recoverssm.py +index 073bdd5..9379688 100644 +--- a/vllm/model_executor/layers/mamba/ops/recoverssm.py ++++ b/vllm/model_executor/layers/mamba/ops/recoverssm.py +@@ -78,8 +78,12 @@ def prepare_commit_plan_kernel( + tl.int32 + ) + final_num_computed = num_computed + commit_len ++ # The committed state belongs to the block holding the last committed ++ # token. A window ending exactly on a boundary stays in that block; the ++ # next block may still be unallocated. + final_state_col = tl.minimum( +- final_num_computed // mamba_block_size, block_table_width - 1 ++ tl.maximum(final_num_computed - 1, 0) // mamba_block_size, ++ block_table_width - 1, + ) + final_state_idx = tl.load( + block_table_ptr +diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py +index cfbe1a4..b02a6c1 100644 +--- a/vllm/v1/worker/gpu/model_states/recoverssm.py ++++ b/vllm/v1/worker/gpu/model_states/recoverssm.py +@@ -91,10 +91,12 @@ def _postprocess_recoverssm_align_kernel( + return + num_sampled = tl.load(num_sampled_ptr + batch_idx) + num_computed = tl.load(num_computed_ptr + batch_idx) ++ # Match the commit plan: the running state lives in the block holding the ++ # last computed token, so the next step migrates it across a boundary. + tl.store( + state_idx_ptr + req_state_idx, + tl.minimum( +- (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, ++ tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE, + BLOCK_TABLE_WIDTH - 1, + ), + ) +-- +2.43.0 + diff --git a/patches/0074-recoverssm-release-the-RecoverSSM-step-after-CUDA-gr.patch b/patches/0074-recoverssm-release-the-RecoverSSM-step-after-CUDA-gr.patch new file mode 100644 index 0000000..82ce37b --- /dev/null +++ b/patches/0074-recoverssm-release-the-RecoverSSM-step-after-CUDA-gr.patch @@ -0,0 +1,90 @@ +From aaffb7483252d9ff8a6bb395ceba4f1d87f5a981 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 29 Sep 2026 18:16:59 -0700 +Subject: [PATCH 074/103] recoverssm: release the RecoverSSM step after CUDA + graph memory profiling (port of upstream #58863, fdcc43cf3e) + +Port of the upstream vLLM change #58863 (Apache-2.0). The step recorded during graph-capture dummy runs held the profiling KV cache +alive after profiling; it is now released at teardown. +--- + ...gpu_model_runner_v2_cudagraph_profiling.py | 33 +++++++++++++++++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 5 +++ + vllm/v1/worker/gpu/model_states/recoverssm.py | 5 +++ + 3 files changed, 43 insertions(+) + +diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +index 7419120..c912053 100644 +--- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py ++++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +@@ -439,3 +439,36 @@ def test_teardown_profiling_state_clears_mamba_align_metadata(monkeypatch): + assert runner.model_state._mamba_ctx is None + assert runner.model_state._mamba_group_ids == [] + assert runner.model_state._mamba_spec is None ++ ++ ++def test_teardown_profiling_state_releases_recoverssm_step(monkeypatch): ++ """The recorded RecoverSSM step must not outlive the profiling KV cache. ++ ++ PIECEWISE capture runs are prepared with ``for_capture=False``, so ++ ``RecoverSSMState`` records their step. Its metadata holds commit contexts ++ that reference the profiling KV cache; keeping it alive past teardown keeps ++ that cache allocated and the real KV cache allocation runs out of memory. ++ """ ++ from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState ++ ++ recoverssm = RecoverSSMState() ++ recoverssm._step = (object(),) # type: ignore[assignment] ++ runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) ++ runner.compilation_config = SimpleNamespace(static_forward_context={}) ++ runner.model_state = SimpleNamespace( ++ supports_mm_inputs=False, recoverssm=recoverssm ++ ) ++ runner.cache_config = SimpleNamespace(num_gpu_blocks=1) ++ runner.kv_caches = [] ++ runner.attn_groups = [] ++ runner.kv_cache_config = SimpleNamespace() ++ runner.cudagraph_manager = object() ++ runner.lora_config = None ++ runner.maybe_remove_all_loras = lambda _: None ++ ++ monkeypatch.setattr(cgu.torch.accelerator, "synchronize", lambda: None) ++ monkeypatch.setattr(cgu.torch.accelerator, "empty_cache", lambda: None) ++ ++ cgu._teardown_profiling_state(runner) ++ ++ assert recoverssm._step is None +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index a143a3d..b28286a 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -1114,6 +1114,11 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + torch.accelerator.synchronize() + if hasattr(runner.model_state, "_mamba_ctx"): + runner.model_state._mamba_ctx = None ++ # A PIECEWISE profiling capture records its RecoverSSM step (only FULL ++ # captures skip it), and that step's commit contexts reference the ++ # profiling KV cache. ++ if (recoverssm := getattr(runner.model_state, "recoverssm", None)) is not None: ++ recoverssm.reset() + # Invalidate the align-mode Mamba group metadata cached from the + # profiling KVCacheConfig: the real (e.g. PP-projected) config may + # place Mamba layers into a different group layout, so it must be +diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py +index b02a6c1..3cfbece 100644 +--- a/vllm/v1/worker/gpu/model_states/recoverssm.py ++++ b/vllm/v1/worker/gpu/model_states/recoverssm.py +@@ -16,6 +16,11 @@ class RecoverSSMState: + def __init__(self) -> None: + self._step: tuple[RecoverSSMMetadata, ...] | None = None + ++ def reset(self) -> None: ++ """Drop the recorded step. Its metadata holds the commit contexts, which ++ reference the KV cache tensors they were built for.""" ++ self._step = None ++ + def record_step( + self, + attn_metadata: dict[str, Any], +-- +2.43.0 + diff --git a/patches/0075-recoverssm-read-the-GDN-RecoverSSM-flag-defensively-.patch b/patches/0075-recoverssm-read-the-GDN-RecoverSSM-flag-defensively-.patch new file mode 100644 index 0000000..1502105 --- /dev/null +++ b/patches/0075-recoverssm-read-the-GDN-RecoverSSM-flag-defensively-.patch @@ -0,0 +1,28 @@ +From 1954885e52446a5979dd514c004fe5d76bb5ed12 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 29 Sep 2026 18:40:22 -0700 +Subject: [PATCH 075/103] recoverssm: read the GDN RecoverSSM flag defensively + in _forward_core (port of upstream #58863, c1d2da8773) + +Port of the upstream vLLM change #58863 (Apache-2.0). One check read use_gdn_recoverssm directly; it now uses getattr with a False +default like its neighbours, so layers built without __init__ in tests work. +--- + vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py | 2 +- + 1 file changed, 1 insertion(+), 1 deletion(-) + +diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +index 2d6b675..80b1344 100644 +--- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py ++++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +@@ -1470,7 +1470,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + if ( + spec_sequence_masks is not None + and mixed_qkv_spec is not None +- and self.use_gdn_recoverssm ++ and getattr(self, "use_gdn_recoverssm", False) + ): + query_spec, key_spec, value_spec, verify_views = ( + self._recoverssm_verify_qkv(mixed_qkv_spec) +-- +2.43.0 + diff --git a/patches/0076-recoverssm-one-fold-per-program-in-the-GDN-RecoverSS.patch b/patches/0076-recoverssm-one-fold-per-program-in-the-GDN-RecoverSS.patch new file mode 100644 index 0000000..4b701f5 --- /dev/null +++ b/patches/0076-recoverssm-one-fold-per-program-in-the-GDN-RecoverSS.patch @@ -0,0 +1,108 @@ +From f5b942010a1cce6c0807bcd9cbcc36c1f6bec792 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 29 Sep 2026 20:19:14 -0700 +Subject: [PATCH 076/103] recoverssm: one fold per program in the GDN + RecoverSSM state commit (port of upstream #58863, af0d2b45d7) + +Port of the upstream vLLM change #58863 (Apache-2.0). Each program folds either the final or the block-boundary state, not both, +which halves the live accumulator tiles. +--- + .../layers/mamba/gdn/recoverssm_gdn.py | 51 +++++++++---------- + 1 file changed, 25 insertions(+), 26 deletions(-) + +diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +index d347bef..a95d579 100644 +--- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py ++++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +@@ -353,6 +353,10 @@ def _commit_gdn_state_kernel( + NUM_HEADS: tl.constexpr, + ALIGN_MODE: tl.constexpr, + ): ++ # In align mode the second half of axis 0 folds the block-boundary state ++ # (the window's first boundary_len tokens) instead of the final state, so every ++ # program runs one fold and keeps one tile live. ++ NUM_V_TILES: tl.constexpr = (V + BV - 1) // BV + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_lh = tl.program_id(2) +@@ -365,10 +369,17 @@ def _commit_gdn_state_kernel( + if n == 0: + return + final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) +- boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) +- boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b) + if final_idx <= null_block_id: + return ++ dst_idx = final_idx ++ # ALIGN_MODE is a constexpr: kept apart so non-align builds drop the branch. ++ if ALIGN_MODE: # noqa: SIM102 ++ if pid_v >= NUM_V_TILES: ++ pid_v -= NUM_V_TILES ++ dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) ++ if dst_idx <= null_block_id: ++ return ++ n = tl.load(boundary_recovery_lens_ptr + pid_b) + state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( + tl.pointer_type(state_ref_ptr.dtype.element_ty) + ) +@@ -389,35 +400,19 @@ def _commit_gdn_state_kernel( + mask=mask_state, + other=0.0, + ).to(tl.float32) +- fin_decay = 1.0 +- fin_corr = tl.zeros([BV, BK], tl.float32) +- bnd_decay = 1.0 +- bnd_corr = tl.zeros([BV, BK], tl.float32) ++ decay = 1.0 ++ corr = tl.zeros([BV, BK], tl.float32) + for r in range(n): + t = n - r - 1 + rec = rec_base + t * stride_replay_pos + c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0) + kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0) + g = tl.load(rec + (V + K) * stride_replay_dim) +- upd = c[:, None] * kk[None, :] +- dec = tl.exp(g) +- fin_corr += upd * fin_decay +- fin_decay *= dec +- if ALIGN_MODE: +- before = t < boundary_len +- bnd_corr += tl.where(before, upd * bnd_decay, 0.0) +- bnd_decay *= tl.where(before, dec, 1.0) +- if ALIGN_MODE: +- bs = s0 * bnd_decay + bnd_corr +- tl.store( +- state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp, +- bs.to(state_ref_ptr.dtype.element_ty), +- mask=mask_state & (boundary_idx > null_block_id), +- ) +- fs = s0 * fin_decay + fin_corr ++ corr += (c[:, None] * kk[None, :]) * decay ++ decay *= tl.exp(g) + tl.store( +- state_ptr + final_idx * sbs + pid_h * stride_state_head + sp, +- fs.to(state_ref_ptr.dtype.element_ty), ++ state_ptr + dst_idx * sbs + pid_h * stride_state_head + sp, ++ (s0 * decay + corr).to(state_ref_ptr.dtype.element_ty), + mask=mask_state, + ) + +@@ -659,9 +654,13 @@ class GDNRecoverSSMCommitContext: + ref = self.checkpoints[0] + _, HV, V, K = ref.shape + BK = triton.next_power_of_2(K) +- BV = min(triton.next_power_of_2(V), 32) ++ # 16-row tiles: the commit streams each state once, and smaller tiles ++ # keep more of them in flight (about 13% more bandwidth than 32 on H200). ++ BV = min(triton.next_power_of_2(V), 16) + rr = self.replays[0] +- _commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)]( ++ align = block_table is not None ++ num_v_tiles = triton.cdiv(V, BV) * (2 if align else 1) ++ _commit_gdn_state_kernel[(num_v_tiles, batch, num_layers * HV)]( + ref, + self.state_base_addrs, + self.state_block_strides, +-- +2.43.0 + diff --git a/patches/0077-recoverssm-FULL-cudagraphs-for-uniform-spec-decode-b.patch b/patches/0077-recoverssm-FULL-cudagraphs-for-uniform-spec-decode-b.patch new file mode 100644 index 0000000..93e0b8e --- /dev/null +++ b/patches/0077-recoverssm-FULL-cudagraphs-for-uniform-spec-decode-b.patch @@ -0,0 +1,206 @@ +From 9caab704196d6854658181e1e995a2869bc2d38e Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 30 Sep 2026 07:11:03 +0200 +Subject: [PATCH 077/103] recoverssm: FULL cudagraphs for uniform spec-decode + batches (port of upstream #58863, 3aa3d5b586) + +Port of the upstream vLLM change #58863 (Apache-2.0). The GDN and PLE RecoverSSM metadata builders declare UNIFORM_BATCH, so +--use-replayssm resolves FULL_AND_PIECEWISE; graph padding rows get a null state slot and write nothing. +--- + tests/config/test_recoverssm_config.py | 4 +-- + tests/kernels/mamba/test_recoverssm_gdn.py | 28 +++++++++++++++++++ + vllm/v1/attention/backends/gdn_recoverssm.py | 29 ++++++++++++++++---- + vllm/v1/attention/backends/ple_recoverssm.py | 23 +++++++++++++--- + 4 files changed, 72 insertions(+), 12 deletions(-) + +diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py +index 21791aa..27e26d9 100644 +--- a/tests/config/test_recoverssm_config.py ++++ b/tests/config/test_recoverssm_config.py +@@ -134,14 +134,14 @@ def test_uses_recoverssm_predicate(): + assert not uses_recoverssm(SimpleNamespace(), 3) + + +-def test_builders_fall_back_to_piecewise(): ++def test_builders_support_uniform_batch_full_graphs(): + from vllm.v1.attention.backends.gdn_recoverssm import GDNRecoverSSMMetadataBuilder + from vllm.v1.attention.backends.ple_recoverssm import PleRecoverSSMMetadataBuilder + + for builder in (GDNRecoverSSMMetadataBuilder, PleRecoverSSMMetadataBuilder): + assert ( + builder.get_cudagraph_support(cast(Any, None), cast(Any, None)) +- == AttentionCGSupport.NEVER ++ == AttentionCGSupport.UNIFORM_BATCH + ) + + +diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py +index ce16f79..724a41d 100644 +--- a/tests/kernels/mamba/test_recoverssm_gdn.py ++++ b/tests/kernels/mamba/test_recoverssm_gdn.py +@@ -600,3 +600,31 @@ def test_align_boundary_bf16_state(nc, accepted, bs): + rtol=BF16_ULP, + atol=1e-5, + ) ++ ++ ++def test_verify_skips_cuda_graph_padding_rows(): ++ """FULL graph replay pads the batch with null-state, zero-length rows (the GDN ++ builder's staging); the verify must give the real rows the same output and replay ++ records as the unpadded call, and touch nothing for the padding.""" ++ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID ++ ++ inp = _make([4, 2, 4], nb=12, seed=3) ++ real = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") ++ ref_cache = inp["replay_cache"].clone() ++ ref = _verify(inp, state_indices=real, replay_cache=ref_cache) ++ ++ before = inp["checkpoint_state"].clone() ++ pad_cache = inp["replay_cache"].clone() ++ padded_idx = torch.tensor( ++ [2, 5, 7, NULL_BLOCK_ID, NULL_BLOCK_ID], dtype=torch.int32, device="cuda" ++ ) ++ padded_qsl = torch.tensor([0, 4, 6, 10, 10, 10], dtype=torch.int32, device="cuda") ++ out = _verify( ++ inp, ++ state_indices=padded_idx, ++ query_start_loc=padded_qsl, ++ replay_cache=pad_cache, ++ ) ++ assert torch.equal(out, ref) ++ assert torch.equal(pad_cache, ref_cache) ++ assert torch.equal(inp["checkpoint_state"], before) +diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py +index a45142a..8b6349a 100644 +--- a/vllm/v1/attention/backends/gdn_recoverssm.py ++++ b/vllm/v1/attention/backends/gdn_recoverssm.py +@@ -76,18 +76,19 @@ class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + + + class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): +- # The verify/commit metadata is not staged for FULL graph replay yet; declaring +- # NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise). +- _cudagraph_support = AttentionCGSupport.NEVER ++ # Uniform spec-decode batches replay as FULL graphs: the base builder stages the ++ # state indices, query offsets and accepted counts into persistent buffers and pads ++ # the batch with null-state, zero-length rows, which the verify kernel skips. The ++ # commit runs after sampling, outside the graph, on the staged tensors. ++ _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) +- # The builder may be created before the runner resolves the CUDA graph mode. +- self.use_full_cuda_graph = False + max_reqs = vllm_config.scheduler_config.max_num_seqs + self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device) + self._context: Any = None + self._logged = False ++ self._capturing = False + + def _get_context(self): + if self._context is None: +@@ -120,6 +121,9 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + m = common_attn_metadata + rows = None + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: ++ if self._capturing and m.is_prefilling is None: ++ # CUDA graph capture: dummy decode batches, every row a spec window. ++ m = m.replace(is_prefilling=torch.zeros(m.num_reqs, dtype=torch.bool)) + if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": + raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask") + # Positive counts only defeat the base builder's "no drafts -> regular +@@ -144,7 +148,11 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + n = meta.num_spec_decodes + assert meta.spec_state_indices_tensor is not None + base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1] +- base["num_accepted_tokens"] = self._ones[:n] ++ # One accepted token per row: the checkpoint already holds the accepted ++ # state. Sized like the (possibly graph-padded) state indices. ++ base["num_accepted_tokens"] = self._ones[ ++ : base["spec_state_indices_tensor"].shape[0] ++ ] + if rows is None: + raise ValueError("GDN RecoverSSM: spec rows without draft counts") + request_indices = recoverssm_request_indices( +@@ -170,6 +178,15 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + recoverssm_context=self._get_context() if commit is not None else None, + ) + ++ def build_for_cudagraph_capture( ++ self, common_attn_metadata: CommonAttentionMetadata ++ ): ++ self._capturing = True ++ try: ++ return super().build_for_cudagraph_capture(common_attn_metadata) ++ finally: ++ self._capturing = False ++ + + class GDNRecoverSSMAttentionBackend(GDNAttentionBackend): + @staticmethod +diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py +index 7402001..9b43a6f 100644 +--- a/vllm/v1/attention/backends/ple_recoverssm.py ++++ b/vllm/v1/attention/backends/ple_recoverssm.py +@@ -230,17 +230,19 @@ class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata): + + + class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): +- # Same as the GDN RecoverSSM builder: PIECEWISE only for now. +- _cudagraph_support = AttentionCGSupport.NEVER ++ # Same as the GDN RecoverSSM builder: uniform spec-decode batches replay as FULL ++ # graphs on the base builder's staged, null-padded buffers; the commit runs after ++ # sampling, outside the graph. ++ _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) +- self.use_full_cuda_graph = False + self._ones = torch.ones( + vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device + ) + self._ctx: _PleConvCommit | None = None + self._logged = False ++ self._capturing = False + + def _get_ctx(self): + if self._ctx is None: +@@ -270,6 +272,9 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + m = common_attn_metadata + rows = None + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: ++ if self._capturing and m.is_prefilling is None: ++ # CUDA graph capture: dummy decode batches, every row a spec window. ++ m = m.replace(is_prefilling=torch.zeros(m.num_reqs, dtype=torch.bool)) + if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": + raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask") + # the same classification as the GDN builder, so both commits map the same +@@ -295,7 +300,10 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + return meta + base = {f.name: getattr(meta, f.name) for f in fields(meta)} + n = meta.num_spec_decodes +- base["num_accepted_tokens"] = self._ones[:n] ++ # sized like the (possibly graph-padded) state indices ++ assert meta.spec_state_indices_tensor is not None ++ num_rows = meta.spec_state_indices_tensor.shape[0] ++ base["num_accepted_tokens"] = self._ones[:num_rows] + if rows is None: + raise ValueError("PLE RecoverSSM: spec rows without draft counts") + req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device) +@@ -315,6 +323,13 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + **base, recoverssm_commit=commit, recoverssm_context=self._get_ctx() + ) + ++ def build_for_cudagraph_capture(self, common_attn_metadata): ++ self._capturing = True ++ try: ++ return super().build_for_cudagraph_capture(common_attn_metadata) ++ finally: ++ self._capturing = False ++ + + class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend): + @staticmethod +-- +2.43.0 + diff --git a/patches/0078-recoverssm-align-commit-fold-final-states-before-bou.patch b/patches/0078-recoverssm-align-commit-fold-final-states-before-bou.patch new file mode 100644 index 0000000..b924296 --- /dev/null +++ b/patches/0078-recoverssm-align-commit-fold-final-states-before-bou.patch @@ -0,0 +1,100 @@ +From c2a86c85f32040d36d0837657994478c731101a7 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Wed, 30 Sep 2026 07:35:13 +0200 +Subject: [PATCH 078/103] recoverssm: align commit: fold final states before + boundary states (port of upstream #58863, 3388ba1a25) + +Port of the upstream vLLM change #58863 (Apache-2.0). The final-state folds run before the boundary-state folds, because in align +mode a boundary state can be written into the slot a final fold still reads. +--- + .../layers/mamba/gdn/recoverssm_gdn.py | 51 ++++++++++--------- + 1 file changed, 28 insertions(+), 23 deletions(-) + +diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +index a95d579..1404edc 100644 +--- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py ++++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +@@ -351,12 +351,12 @@ def _commit_gdn_state_kernel( + BK: tl.constexpr, + BV: tl.constexpr, + NUM_HEADS: tl.constexpr, +- ALIGN_MODE: tl.constexpr, ++ BOUNDARY: tl.constexpr, + ): +- # In align mode the second half of axis 0 folds the block-boundary state +- # (the window's first boundary_len tokens) instead of the final state, so every +- # program runs one fold and keeps one tile live. +- NUM_V_TILES: tl.constexpr = (V + BV - 1) // BV ++ # Every program runs one fold and keeps one tile live. In align mode the ++ # block-boundary state (the window's first boundary_len tokens) is folded by a ++ # second launch (BOUNDARY): a crossing window's boundary state goes into its ++ # source block, which the first launch's final-state fold must still read. + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_lh = tl.program_id(2) +@@ -372,14 +372,13 @@ def _commit_gdn_state_kernel( + if final_idx <= null_block_id: + return + dst_idx = final_idx +- # ALIGN_MODE is a constexpr: kept apart so non-align builds drop the branch. +- if ALIGN_MODE: # noqa: SIM102 +- if pid_v >= NUM_V_TILES: +- pid_v -= NUM_V_TILES +- dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) +- if dst_idx <= null_block_id: +- return +- n = tl.load(boundary_recovery_lens_ptr + pid_b) ++ if BOUNDARY: ++ dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) ++ # A window ending exactly on the boundary keeps its final state in that ++ # block: the first launch already wrote it (and may have overwritten src). ++ if dst_idx <= null_block_id or dst_idx == final_idx: ++ return ++ n = tl.load(boundary_recovery_lens_ptr + pid_b) + state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( + tl.pointer_type(state_ref_ptr.dtype.element_ty) + ) +@@ -659,8 +658,7 @@ class GDNRecoverSSMCommitContext: + BV = min(triton.next_power_of_2(V), 16) + rr = self.replays[0] + align = block_table is not None +- num_v_tiles = triton.cdiv(V, BV) * (2 if align else 1) +- _commit_gdn_state_kernel[(num_v_tiles, batch, num_layers * HV)]( ++ args = ( + ref, + self.state_base_addrs, + self.state_block_strides, +@@ -680,15 +678,22 @@ class GDNRecoverSSMCommitContext: + rr.stride(2), + rr.stride(3), + state_indices.stride(0), +- K=K, +- V=V, +- BK=BK, +- BV=BV, +- NUM_HEADS=HV, +- ALIGN_MODE=block_table is not None, +- num_warps=4, +- num_stages=2, + ) ++ grid = (triton.cdiv(V, BV), batch, num_layers * HV) ++ # Final states first: a boundary state can land in the window's source ++ # slot, which the final-state fold reads (stream order separates them). ++ for boundary in (False, True) if align else (False,): ++ _commit_gdn_state_kernel[grid]( ++ *args, ++ K=K, ++ V=V, ++ BK=BK, ++ BV=BV, ++ NUM_HEADS=HV, ++ BOUNDARY=boundary, ++ num_warps=4, ++ num_stages=2, ++ ) + + + __all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"] +-- +2.43.0 + diff --git a/patches/0079-recoverssm-GDN-commit-forward-with-the-verify-step-s.patch b/patches/0079-recoverssm-GDN-commit-forward-with-the-verify-step-s.patch new file mode 100644 index 0000000..5beb205 --- /dev/null +++ b/patches/0079-recoverssm-GDN-commit-forward-with-the-verify-step-s.patch @@ -0,0 +1,124 @@ +From e6f5849032deddd8eee331d99242d3ab646198f1 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 20:14:32 +0200 +Subject: [PATCH 079/103] recoverssm: GDN: commit forward with the verify + step's own decay, bit-identical to per-token (port of upstream #58863, + 514102a1db) + +Port of the upstream vLLM change #58863 (Apache-2.0). The verify kernel records the decay it applied and the commit replays the +accepted tokens forward with the same update, so the committed state equals the per-token state bit for bit. +--- + tests/kernels/mamba/test_recoverssm_gdn.py | 14 ++++------- + .../layers/mamba/gdn/recoverssm_gdn.py | 24 ++++++++++--------- + 2 files changed, 18 insertions(+), 20 deletions(-) + +diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py +index 724a41d..59ddb7f 100644 +--- a/tests/kernels/mamba/test_recoverssm_gdn.py ++++ b/tests/kernels/mamba/test_recoverssm_gdn.py +@@ -111,7 +111,7 @@ def test_commit_matches_native_per_token_state(accepted): + ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) + for i, blk in enumerate(inp["state_indices"].tolist()): + tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 +- torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) ++ assert torch.equal(ckpt[blk], states[tok]) + others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()] + assert torch.equal(ckpt[others], inp["checkpoint_state"][others]) + +@@ -301,7 +301,7 @@ def test_commit_maps_noncontiguous_requests_in_align_mode(): + c0 = int(nc[r]) + final_blk = int(bt[r, max(c0 + n - 1, 0) // bs]) + tok = int(inp["query_start_loc"][i]) + n - 1 +- torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6) ++ assert torch.equal(ckpt[final_blk], states[tok]) + + + # (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on +@@ -343,12 +343,8 @@ def test_align_boundary_and_final_state_match_native(nc, accepted, bs): + if ( + nc + accepted >= next_boundary + ): # crosses: the boundary state goes into the source block +- torch.testing.assert_close( +- ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6 +- ) +- torch.testing.assert_close( +- ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6 +- ) ++ assert torch.equal(ckpt[int(src)], states[next_boundary - nc - 1]) ++ assert torch.equal(ckpt[final_blk], states[accepted - 1]) + + + @pytest.mark.parametrize("gate", list(GATES)) +@@ -367,7 +363,7 @@ def test_commit_matches_native_under_extreme_gates(gate, seed): + ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) + for i, blk in enumerate(inp["state_indices"].tolist()): + tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 +- torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) ++ assert torch.equal(ckpt[blk], states[tok]) + + + @pytest.mark.parametrize("batch", [1, 2, 4]) +diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +index 1404edc..8adc094 100644 +--- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py ++++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +@@ -9,7 +9,7 @@ GDN port of the Kimi-K3 KDA RecoverSSM path + scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped + over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested. + The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t +-k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one ++k_t) (V floats), the normalised key (K floats) and the decay exp(g_t) (1 float) in one + fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 + + sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the + block-boundary state in align mode).""" +@@ -150,7 +150,8 @@ def _gdn_recoverssm_verify_kernel( + q = q * tl.rsqrt(tl.sum(q * q) + 1e-6) + k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + q = q * scale +- new_state = state * tl.exp(g) ++ decay = tl.exp(g) ++ new_state = state * decay + corr = (v - tl.sum(new_state * k[None, :], 1)) * beta + new_state = new_state + corr[:, None] * k[None, :] + state = tl.where(valid, new_state, state) +@@ -164,7 +165,8 @@ def _gdn_recoverssm_verify_kernel( + tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v) + if pid_v == 0: + tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k) +- tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid) ++ # the decay as multiplied here, so the commit replays this exact float ++ tl.store(rec + (V + K) * stride_replay_dim, decay, mask=valid) + + + def gdn_recoverssm_verify( +@@ -399,19 +401,19 @@ def _commit_gdn_state_kernel( + mask=mask_state, + other=0.0, + ).to(tl.float32) +- decay = 1.0 +- corr = tl.zeros([BV, BK], tl.float32) +- for r in range(n): +- t = n - r - 1 ++ # Forward over the accepted tokens with the verify step's own update, so the committed ++ # state is the per-token state bit for bit (the decay and corr stored are the floats ++ # the verify kernel used). ++ h = s0 ++ for t in range(n): + rec = rec_base + t * stride_replay_pos + c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0) + kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0) +- g = tl.load(rec + (V + K) * stride_replay_dim) +- corr += (c[:, None] * kk[None, :]) * decay +- decay *= tl.exp(g) ++ decay = tl.load(rec + (V + K) * stride_replay_dim) ++ h = tl.fma(c[:, None], kk[None, :], h * decay) + tl.store( + state_ptr + dst_idx * sbs + pid_h * stride_state_head + sp, +- (s0 * decay + corr).to(state_ref_ptr.dtype.element_ty), ++ h.to(state_ref_ptr.dtype.element_ty), + mask=mask_state, + ) + +-- +2.43.0 + diff --git a/patches/0080-recoverssm-allow-pipeline-parallelism-for-Qwen4Exp-o.patch b/patches/0080-recoverssm-allow-pipeline-parallelism-for-Qwen4Exp-o.patch new file mode 100644 index 0000000..cba2f25 --- /dev/null +++ b/patches/0080-recoverssm-allow-pipeline-parallelism-for-Qwen4Exp-o.patch @@ -0,0 +1,1017 @@ +From 9917a9883e710ddae023d31392c0d5804d2b3b6c Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 21:15:44 -0400 +Subject: [PATCH 080/103] recoverssm: allow pipeline parallelism for Qwen4Exp + on the V2 runner + +Upstream RecoverSSM commits a step in the same step's postprocess, hence PP=1. Under pipeline parallelism a +non-last rank learns a step's accepted counts only when the PP broadcast arrives pp_size steps later, while later +builds overwrite the persistent commit inputs. The commit is now deferred with its own copy of those inputs, and +the PP=1 check is lifted for Qwen4Exp on Model Runner V2. Config and CPU tests cover the deferral. +--- + tests/config/test_recoverssm_config.py | 13 +- + tests/config/test_recoverssm_pp_config.py | 55 +++ + .../worker/gpu/test_recoverssm_pp_deferral.py | 408 ++++++++++++++++++ + vllm/config/vllm.py | 22 +- + vllm/v1/attention/backends/gdn_recoverssm.py | 34 +- + vllm/v1/attention/backends/ple_recoverssm.py | 23 +- + .../attention/backends/recoverssm_metadata.py | 41 +- + vllm/v1/worker/gpu/cudagraph_utils.py | 4 + + vllm/v1/worker/gpu/model_runner.py | 11 + + vllm/v1/worker/gpu/model_states/interface.py | 13 + + .../worker/gpu/model_states/mamba_hybrid.py | 14 + + vllm/v1/worker/gpu/model_states/recoverssm.py | 68 +++ + vllm/v1/worker/gpu/pp_utils.py | 31 +- + 13 files changed, 727 insertions(+), 10 deletions(-) + create mode 100644 tests/config/test_recoverssm_pp_config.py + create mode 100644 tests/v1/worker/gpu/test_recoverssm_pp_deferral.py + +diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py +index 27e26d9..2fba7eb 100644 +--- a/tests/config/test_recoverssm_config.py ++++ b/tests/config/test_recoverssm_config.py +@@ -83,19 +83,26 @@ def test_derived_drafter_keeps_the_choice_and_runs_the_checks(): + drafter = _config(DRAFTER, cache_config=target.cache_config) + _validator()(drafter) + assert target.cache_config.use_recoverssm is True +- bad = _config( ++ # Fork: PP > 1 is lifted for Qwen4Exp on Model Runner V2; the ++ # drafter's derived config runs the same check and passes it. ++ pp2 = _config( + DRAFTER, + cache_config=target.cache_config, + parallel_config__pipeline_parallel_size=2, + ) +- with pytest.raises(ValueError, match="pipeline_parallel_size"): ++ _validator()(pp2) ++ bad = _config( ++ DRAFTER, ++ cache_config=target.cache_config, ++ mamba_config__backend="flashinfer", ++ ) ++ with pytest.raises(ValueError, match="mamba-backend triton"): + _validator()(bad) + + + @pytest.mark.parametrize( + ("override", "value", "match"), + [ +- ("parallel_config__pipeline_parallel_size", 2, "pipeline_parallel_size"), + ("mamba_config__enable_stochastic_rounding", True, "stochastic"), + ("mamba_config__backend", "flashinfer", "mamba-backend triton"), + ("use_v2_model_runner", False, "V2_MODEL_RUNNER"), +diff --git a/tests/config/test_recoverssm_pp_config.py b/tests/config/test_recoverssm_pp_config.py +new file mode 100644 +index 0000000..95117f9 +--- /dev/null ++++ b/tests/config/test_recoverssm_pp_config.py +@@ -0,0 +1,55 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The RecoverSSM pipeline-parallel lift at the config level. ++ ++Upstream (#58863) rejects RecoverSSM for pipeline_parallel_size > 1. The fork ++lifts that for Qwen4Exp GDN/PLE on Model Runner V2, where non-last PP ranks ++defer the commit with a snapshot of its inputs. Kimi-K3 KDA keeps the PP=1 ++check (its metadata has no snapshot), and Model Runner V1 has no deferral. ++ ++Failing control: on the unlifted port (upstream check), test_pp2_qwen4exp_v2_* ++raises "RecoverSSM currently requires pipeline_parallel_size=1". ++""" ++ ++import pytest ++ ++from tests.config.test_recoverssm_config import _config, _validator ++ ++TARGET = "Qwen4ExpForConditionalGeneration" ++ ++ ++@pytest.mark.parametrize("pp", [2, 4]) ++@pytest.mark.parametrize( ++ "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM", "Qwen4ExpMTP"] ++) ++def test_pp2_qwen4exp_v2_is_accepted(arch, pp): ++ cfg = _config(arch, parallel_config__pipeline_parallel_size=pp) ++ _validator()(cfg) ++ assert cfg.cache_config.use_recoverssm is True ++ ++ ++def test_pp2_kda_keeps_the_pp1_requirement(): ++ cfg = _config( ++ "KimiK3ForConditionalGeneration", parallel_config__pipeline_parallel_size=2 ++ ) ++ with pytest.raises(ValueError, match="only for Qwen4Exp"): ++ _validator()(cfg) ++ ++ ++def test_pp2_v1_runner_is_rejected(): ++ # none mode, so the align-mode V2 check does not fire first ++ cfg = _config(TARGET, parallel_config__pipeline_parallel_size=2) ++ cfg.cache_config.mamba_cache_mode = "none" ++ cfg.use_v2_model_runner = False ++ with pytest.raises(ValueError, match="pipeline_parallel_size > 1 requires"): ++ _validator()(cfg) ++ ++ ++def test_pp2_still_runs_the_other_checks(): ++ cfg = _config( ++ TARGET, ++ parallel_config__pipeline_parallel_size=2, ++ mamba_config__backend="flashinfer", ++ ) ++ with pytest.raises(ValueError, match="mamba-backend triton"): ++ _validator()(cfg) +diff --git a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +new file mode 100644 +index 0000000..9420aaa +--- /dev/null ++++ b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +@@ -0,0 +1,408 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The RecoverSSM commit under pipeline parallelism (CPU, no GPU). ++ ++A non-last PP rank verifies a speculative window at step T but receives the ++accepted counts only when the PP broadcast of step T is consumed, pp_size steps ++later (``PPHandler.get_prev_sampled_outputs`` in ``update_pp_decode_requests``). ++Upstream #58863 keeps one recorded step (``RecoverSSMState._step``) and commits ++it in the same step's postprocess, so it requires pipeline_parallel_size=1: ++under PP the deferred postprocess of step T would commit the metadata of the ++most recent step instead, and that metadata's spec state indices, query offsets ++and batch-order block table are persistent buffers the next build overwrites. ++ ++The fork detaches the step in ``sample_tokens`` (a snapshot whose commit inputs ++are clones), keeps it on step T's pending PP entry, restores it right before ++step T's deferred postprocess, and masks rows whose request was freed in ++between (their blocks may already belong to another request). ++ ++Every scenario here has a failing control that flips: the upstream single-slot ++flow and a snapshot that aliases the persistent buffers both commit the wrong ++step; a deferred commit without the freed-row mask commits into a freed slot; ++a teardown without the PP purge keeps the profiling KV cache alive. ++""" ++ ++import ast ++import gc ++import inspect ++import textwrap ++import weakref ++from collections import deque ++from dataclasses import dataclass, fields ++from types import SimpleNamespace ++ ++import numpy as np ++import pytest ++import torch ++ ++from vllm.v1.attention.backends import gdn_recoverssm as gdn_mod ++from vllm.v1.attention.backends.gdn_recoverssm import ( ++ GDNRecoverSSMCommitMetadata, ++ GDNRecoverSSMMetadata, ++) ++from vllm.v1.attention.backends.ple_recoverssm import ( ++ PleRecoverSSMCommitMetadata, ++ PleRecoverSSMMetadata, ++) ++from vllm.v1.worker.gpu import pp_utils as pp_utils_mod ++from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState ++ ++MAX_REQS = 8 ++SPEC_LEN = 4 # 1 + K, K = 3 (prod) ++ ++ ++class _RecordingContext: ++ """Stands in for GDNRecoverSSMCommitContext: records what each commit ++ reads, at commit time (clones, like the kernels reading device memory).""" ++ ++ def __init__(self) -> None: ++ self.commits: list[dict] = [] ++ ++ def commit(self, num_accepted_tokens, state_indices, query_start_loc, **kw): ++ req = kw.get("request_indices") ++ rows = ( ++ list(range(state_indices.shape[0])) ++ if req is None ++ else [int(r) for r in req.tolist()] ++ ) ++ self.commits.append( ++ dict( ++ state_indices=state_indices.tolist(), ++ query_start_loc=query_start_loc.tolist(), ++ num_accepted=[ ++ int(num_accepted_tokens[r]) ++ if r < num_accepted_tokens.shape[0] ++ else None ++ for r in rows ++ ], ++ block_table=None ++ if kw.get("block_table") is None ++ else kw["block_table"].clone(), ++ ) ++ ) ++ ++ ++class _Buffers: ++ """The builder's persistent, graph-staged buffers (one set per builder).""" ++ ++ def __init__(self) -> None: ++ self.spec_state_indices = torch.full((MAX_REQS, SPEC_LEN), -1, dtype=torch.int32) ++ self.spec_query_start_loc = torch.zeros(MAX_REQS + 1, dtype=torch.int32) ++ self.block_table = torch.zeros(MAX_REQS, 6, dtype=torch.int32) ++ ++ ++def _gdn_meta(buf: _Buffers, slots, ctx, *, align=False) -> GDNRecoverSSMMetadata: ++ """Build step metadata the way GDNRecoverSSMMetadataBuilder.build does: the ++ commit points at views of the persistent buffers, written in place.""" ++ n = len(slots) ++ buf.spec_state_indices[:n, 0] = torch.tensor(slots, dtype=torch.int32) ++ buf.spec_query_start_loc[: n + 1] = torch.arange( ++ 0, (n + 1) * SPEC_LEN, SPEC_LEN, dtype=torch.int32 ++ ) ++ for i, s in enumerate(slots): ++ buf.block_table[i] = s * 10 + torch.arange(6, dtype=torch.int32) ++ commit = GDNRecoverSSMCommitMetadata( ++ state_indices=buf.spec_state_indices[:, :1], ++ query_start_loc=buf.spec_query_start_loc, ++ request_indices=None, ++ block_table=buf.block_table[:n] if align else None, ++ num_computed_tokens=torch.full((n,), 100, dtype=torch.int32) ++ if align ++ else None, ++ block_size=64 if align else None, ++ ) ++ required = dict( ++ num_prefills=0, ++ num_prefill_tokens=0, ++ num_decodes=0, ++ num_decode_tokens=0, ++ num_spec_decodes=n, ++ num_spec_decode_tokens=n * SPEC_LEN, ++ num_actual_tokens=n * SPEC_LEN, ++ ) ++ kwargs = { ++ f.name: required.get(f.name) ++ for f in fields(GDNRecoverSSMMetadata) ++ if f.init and f.name not in ("recoverssm_commit", "recoverssm_context") ++ } ++ return GDNRecoverSSMMetadata( ++ **kwargs, recoverssm_commit=commit, recoverssm_context=ctx ++ ) ++ ++ ++def _record(state: RecoverSSMState, meta) -> None: ++ group = SimpleNamespace(layer_names=["model.layers.0.linear_attn"]) ++ state.record_step( ++ {"model.layers.0.linear_attn": meta}, [[group]], for_capture=False ++ ) ++ ++ ++def _handler() -> pp_utils_mod.PPHandler: ++ h = object.__new__(pp_utils_mod.PPHandler) ++ h.is_last_rank = False ++ h.device = torch.device("cpu") ++ h.main_stream = SimpleNamespace(wait_event=lambda event: None) ++ h.req_idx_gen_np = np.zeros(MAX_REQS, dtype=np.int32) ++ h.queue = deque([None, None]) # pp_size = 2, pre-seeded like __init__ ++ return h ++ ++ ++def _receive(h: pp_utils_mod.PPHandler, idx_mapping, num_sampled) -> None: ++ """What PPHandler.receive leaves in the queue (minus the NCCL broadcast).""" ++ idx_np = np.asarray(idx_mapping, dtype=np.int64) ++ n = len(idx_mapping) ++ h.queue[-1] = pp_utils_mod.PendingRecv( ++ event=object(), ++ sampled_tokens=torch.zeros(n, SPEC_LEN, dtype=torch.int64), ++ num_sampled=torch.tensor(num_sampled, dtype=torch.int32), ++ num_rejected=torch.zeros(n, dtype=torch.int32), ++ idx_mapping=torch.tensor(idx_mapping, dtype=torch.int32), ++ idx_mapping_np=idx_np, ++ need_sampled_mask=np.ones(n, dtype=bool), ++ gen_at_receive_np=h.req_idx_gen_np[idx_np].copy(), ++ ) ++ ++ ++@pytest.fixture ++def h2d(monkeypatch): ++ monkeypatch.setattr( ++ pp_utils_mod, ++ "async_tensor_h2d", ++ lambda a, device=None, dtype=None, out=None: torch.as_tensor( ++ np.asarray(a), dtype=torch.int32 ++ ), ++ ) ++ ++ ++def _run_two_microbatches(mode: str, *, free_slot: int | None = None): ++ """Non-last rank, pp_size 2: micro-batch A at step T (req slots 0, 1 with ++ state blocks 3, 5), micro-batch B at T+1 (req slots 2, 4, 6 with blocks 7, ++ 9, 11), then step T+2 consumes A's broadcast. Returns the commit A got. ++ ++ mode: "fork" (detach/attach/restore), "upstream" (#58863 single slot) or ++ "alias" (fork flow, but the snapshot aliases the persistent buffers).""" ++ ctx = _RecordingContext() ++ buf = _Buffers() # one builder: both micro-batches share its buffers ++ state = RecoverSSMState() ++ h = _handler() ++ ++ def sample_tokens(idx_mapping, num_sampled): ++ _receive(h, idx_mapping, num_sampled) ++ if mode != "upstream": ++ h.attach_deferred_state(state.detach_step()) ++ ++ # step T: micro-batch A (execute_model first consumes a pre-seeded None) ++ assert h.get_prev_sampled_outputs() is None ++ _record(state, _gdn_meta(buf, [3, 5], ctx)) ++ sample_tokens([0, 1], [2, 4]) ++ # step T+1: micro-batch B; its build overwrites the builder's buffers ++ assert h.get_prev_sampled_outputs() is None ++ _record(state, _gdn_meta(buf, [7, 9, 11], ctx)) ++ sample_tokens([2, 4, 6], [1, 1, 1]) ++ if free_slot is not None: ++ h.on_req_idx_freed(free_slot) # A's request freed before T+2 ++ # step T+2: update_pp_decode_requests consumes A's entry ++ out = h.get_prev_sampled_outputs() ++ assert out is not None ++ deferred = out.pop("deferred_state") ++ if mode != "upstream": ++ state.restore_step(deferred) ++ state.commit_step( ++ out["num_sampled"], ++ out["idx_mapping"], ++ state_indices=None, ++ num_accepted_tokens=torch.ones(MAX_REQS, dtype=torch.int32), ++ ) ++ assert ctx.commits, "no commit ran" ++ return ctx.commits[-1] ++ ++ ++A_COMMIT = dict(state_indices=[3, 5], num_accepted=[2, 4]) ++ ++ ++def test_deferred_commit_commits_its_own_step(h2d): ++ got = _run_two_microbatches("fork") ++ assert got["state_indices"] == A_COMMIT["state_indices"] ++ assert got["num_accepted"] == A_COMMIT["num_accepted"] ++ assert got["query_start_loc"] == [0, 4, 8] ++ ++ ++def test_control_upstream_single_slot_commits_the_wrong_step(h2d): ++ # #58863 as-is under PP: step T's counts land on micro-batch B's blocks. ++ got = _run_two_microbatches("upstream") ++ assert got["state_indices"] != A_COMMIT["state_indices"] ++ assert got["state_indices"] == [7, 9, 11] ++ ++ ++def test_control_aliasing_snapshot_commits_the_wrong_blocks(h2d, monkeypatch): ++ # Detaching without cloning keeps views of the builder's persistent ++ # buffers: the next micro-batch's build rewrites them before the commit. ++ monkeypatch.setattr( ++ GDNRecoverSSMMetadata, "snapshot_for_deferred_commit", lambda self: self ++ ) ++ got = _run_two_microbatches("alias") ++ assert got["state_indices"] != A_COMMIT["state_indices"] ++ ++ ++def test_freed_row_commits_nothing(h2d): ++ # Request slot 1 (block 5) freed between verify and commit: count 0 makes ++ # the commit plan skip it (commit_len 0), slot 0 still commits. ++ got = _run_two_microbatches("fork", free_slot=1) ++ assert got["state_indices"] == A_COMMIT["state_indices"] ++ assert got["num_accepted"] == [2, 0] ++ ++ ++def test_control_unmasked_deferred_commit_writes_a_freed_slot(h2d, monkeypatch): ++ real = RecoverSSMState.restore_step ++ ++ def restore_without_mask(self, deferred): ++ real(self, deferred) ++ self._step_is_deferred = False # what a commit without the mask sees ++ ++ monkeypatch.setattr(RecoverSSMState, "restore_step", restore_without_mask) ++ got = _run_two_microbatches("fork", free_slot=1) ++ assert got["num_accepted"] == [2, 4] # block 5 would be written ++ ++ ++def test_snapshot_clones_every_commit_input_gdn_align(): ++ ctx = _RecordingContext() ++ buf = _Buffers() ++ meta = _gdn_meta(buf, [3, 5], ctx, align=True) ++ snap = meta.snapshot_for_deferred_commit() ++ before = ( ++ snap.recoverssm_commit.state_indices.clone(), ++ snap.recoverssm_commit.query_start_loc.clone(), ++ snap.recoverssm_commit.block_table.clone(), ++ ) ++ _gdn_meta(buf, [7, 9], ctx, align=True) # next build rewrites the buffers ++ c = snap.recoverssm_commit ++ assert torch.equal(c.state_indices, before[0]) ++ assert torch.equal(c.query_start_loc, before[1]) ++ assert torch.equal(c.block_table, before[2]) ++ # control: the live metadata's views did change ++ assert not torch.equal(meta.recoverssm_commit.state_indices[:2], before[0]) ++ assert not torch.equal(meta.recoverssm_commit.block_table, before[2]) ++ for t in (c.state_indices, c.query_start_loc, c.block_table): ++ assert all(t.data_ptr() != b.data_ptr() for b in vars(buf).values()) ++ ++ ++def test_snapshot_ple(): ++ sidx = torch.tensor([3, 5, -1, -1], dtype=torch.int32) ++ qsl = torch.tensor([0, 4, 8, 8, 8], dtype=torch.int32) ++ commit = PleRecoverSSMCommitMetadata( ++ state_indices=sidx, ++ query_start_loc=qsl, ++ request_indices=None, ++ block_table=None, ++ num_computed_tokens=None, ++ block_size=None, ++ ) ++ names = [ ++ f.name ++ for f in fields(PleRecoverSSMMetadata) ++ if f.init and f.name not in ("recoverssm_commit", "recoverssm_context") ++ ] ++ meta = PleRecoverSSMMetadata( ++ **{n: None for n in names}, ++ recoverssm_commit=commit, ++ recoverssm_context=_RecordingContext(), ++ ) ++ meta.num_spec_decodes = 2 ++ snap = meta.snapshot_for_deferred_commit() ++ sidx.fill_(9) ++ qsl.fill_(0) ++ assert snap.recoverssm_commit.state_indices.tolist() == [3, 5] ++ assert snap.recoverssm_commit.query_start_loc.tolist() == [0, 4, 8] ++ ++ ++def test_attach_without_a_pending_entry_raises(): ++ h = _handler() ++ h.queue[-1] = None # receive() skipped the broadcast for this step ++ with pytest.raises(RuntimeError, match="without sampled outputs"): ++ h.attach_deferred_state(object()) ++ h.attach_deferred_state(None) # nothing to keep: fine ++ ++ ++def test_restored_step_without_counts_raises(): ++ state = RecoverSSMState() ++ ctx = _RecordingContext() ++ _record(state, _gdn_meta(_Buffers(), [3], ctx)) ++ state.restore_step(state.detach_step()) ++ with pytest.raises(RuntimeError, match="without sampled counts"): ++ state.commit_step( ++ 0, torch.zeros(1, dtype=torch.int32), state_indices=None, ++ num_accepted_tokens=torch.ones(1, dtype=torch.int32), ++ ) ++ ++ ++def test_detach_of_a_step_without_spec_rows_is_none(): ++ state = RecoverSSMState() ++ meta = _gdn_meta(_Buffers(), [3], _RecordingContext()) ++ meta.recoverssm_commit = None ++ _record(state, meta) ++ assert state.detach_step() is None ++ ++ ++@dataclass ++class _Holder: ++ kv: torch.Tensor ++ ++ ++def test_teardown_drops_the_deferred_states_and_their_kv_cache(): ++ """Profiling-state leak guard (the ~26 GiB on H200 class, #58863 commit 4): ++ a deferred step on a pending PP entry references the KV cache through its ++ commit context; the teardown must release it.""" ++ h = _handler() ++ kv = torch.zeros(1024) ++ ref = weakref.ref(kv) ++ _receive(h, [0], [1]) ++ h.attach_deferred_state(_Holder(kv)) ++ del kv ++ gc.collect() ++ assert ref() is not None # control: held by the pending entry ++ assert h.drop_deferred_states() == 1 ++ gc.collect() ++ assert ref() is None ++ assert h.queue[-1] is not None # the sampled-token entry itself stays ++ ++ ++def test_teardown_calls_the_pp_purge(): ++ from vllm.v1.worker.gpu import cudagraph_utils ++ ++ src = inspect.getsource(cudagraph_utils._teardown_profiling_state) ++ assert "drop_deferred_states()" in src ++ assert "recoverssm.reset()" in src ++ ++ ++def test_context_binding_guard(): ++ good = torch.zeros(4) ++ ctx = SimpleNamespace(checkpoints=(good,)) ++ gdn_mod._check_context_binding(ctx, [good]) ++ with pytest.raises(RuntimeError, match="stale context"): ++ gdn_mod._check_context_binding(ctx, [torch.zeros(4)]) ++ ++ ++def _calls_in(fn) -> list[str]: ++ tree = ast.parse(textwrap.dedent(inspect.getsource(fn))) ++ names = [] ++ for node in ast.walk(tree): ++ if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): ++ names.append((node.lineno, node.col_offset, node.func.attr)) ++ return [n for _, _, n in sorted(names)] ++ ++ ++def test_runner_wiring_order(): ++ """The runner detaches after receive() and before the immediate prefill- ++ chunk postprocess (which would otherwise drop the step), and restores ++ before the deferred postprocess.""" ++ from vllm.v1.worker.gpu.model_runner import GPUModelRunner ++ ++ calls = _calls_in(GPUModelRunner.sample_tokens) ++ i_recv = calls.index("receive") ++ i_attach = calls.index("attach_deferred_state") ++ i_detach = calls.index("detach_deferred_state") ++ i_post = calls.index("postprocess_state") ++ assert i_recv < i_attach and i_recv < i_detach < i_post and i_attach < i_post ++ ++ calls = _calls_in(GPUModelRunner.update_pp_decode_requests) ++ assert calls.index("restore_deferred_state") < calls.index("postprocess_sampled") +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index f215ab3..d07579e 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -3272,9 +3272,25 @@ class VllmConfig: + "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + ) + if self.parallel_config.pipeline_parallel_size > 1: +- raise ValueError( +- "RecoverSSM currently requires pipeline_parallel_size=1" +- ) ++ # Non-last PP ranks get the accepted counts pp_size steps after ++ # their verify. Model Runner V2 defers the commit with a ++ # snapshot of its inputs (RecoverSSMState.detach_step) and ++ # masks freed rows; only the Qwen4Exp GDN/PLE metadata ++ # implement that snapshot (Kimi-K3 KDA does not). ++ if ( ++ self.model_config is None ++ or self.model_config.architecture ++ not in _QWEN4EXP_RECOVERSSM_ARCHS ++ ): ++ raise ValueError( ++ "RecoverSSM with pipeline_parallel_size > 1 is supported " ++ "only for Qwen4Exp GDN/PLE" ++ ) ++ if not self.use_v2_model_runner: ++ raise ValueError( ++ "RecoverSSM with pipeline_parallel_size > 1 requires " ++ "VLLM_USE_V2_MODEL_RUNNER=1" ++ ) + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError("RecoverSSM requires --mamba-backend triton") + elif self.cache_config.mamba_cache_mode == "all": +diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py +index 8b6349a..e1c12ac 100644 +--- a/vllm/v1/attention/backends/gdn_recoverssm.py ++++ b/vllm/v1/attention/backends/gdn_recoverssm.py +@@ -8,7 +8,7 @@ column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv w + already hold the accepted state), and the metadata carries what the post-sampling commit + needs. Mirrors the Kimi-K3 KDA builder's RecoverSSM handling.""" + +-from dataclasses import dataclass, field, fields ++from dataclasses import dataclass, field, fields, replace + from typing import Any + + import torch +@@ -23,6 +23,7 @@ from vllm.v1.attention.backends.gdn_attn import ( + from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, ++ clone_commit_tensors, + recoverssm_request_indices, + recoverssm_spec_rows, + ) +@@ -74,6 +75,32 @@ class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + block_size=c.block_size, + ) + ++ def snapshot_for_deferred_commit(self) -> "GDNRecoverSSMMetadata": ++ c = self.recoverssm_commit ++ if c is None: ++ return self ++ return replace( ++ self, ++ recoverssm_commit=clone_commit_tensors(c, self.num_spec_decodes), ++ ) ++ ++ ++def _check_context_binding(context: Any, checkpoints: list[torch.Tensor]) -> None: ++ """A cached commit context must point at the layers' live KV cache. ++ ++ The context stores raw base addresses of every layer's conv, state and ++ replay tensors. One that outlived its cache (the profiling KV cache is ++ torn down and re-allocated before serving) would keep that cache alive ++ and, worse, commit into freed memory. Builders are re-created with the ++ real cache, so a mismatch is a lifecycle bug: fail loudly.""" ++ bound = [t.data_ptr() for t in context.checkpoints] ++ live = [t.data_ptr() for t in checkpoints] ++ if bound != live: ++ raise RuntimeError( ++ "GDN RecoverSSM commit context is bound to a different KV cache than " ++ "its layers (stale context after a KV cache re-allocation)" ++ ) ++ + + class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + # Uniform spec-decode batches replay as FULL graphs: the base builder stages the +@@ -91,6 +118,11 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + self._capturing = False + + def _get_context(self): ++ if self._context is not None: ++ fc = self.vllm_config.compilation_config.static_forward_context ++ _check_context_binding( ++ self._context, [fc[name].kv_cache[1] for name in self.layer_names] ++ ) + if self._context is None: + from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( + GDNRecoverSSMCommitContext, +diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py +index 9b43a6f..9e44ad8 100644 +--- a/vllm/v1/attention/backends/ple_recoverssm.py ++++ b/vllm/v1/attention/backends/ple_recoverssm.py +@@ -9,7 +9,7 @@ every active decode row takes the spec path, the conv reads at offset 0, and aft + sampling the accepted window is compacted to the front (the shared RecoverSSM compaction + kernel; in align mode it also writes the block-boundary window).""" + +-from dataclasses import dataclass, field, fields ++from dataclasses import dataclass, field, fields, replace + from typing import Any + + import torch +@@ -25,6 +25,7 @@ from vllm.v1.attention.backend import AttentionCGSupport + from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, ++ clone_commit_tensors, + recoverssm_request_indices, + recoverssm_spec_rows, + ) +@@ -228,6 +229,15 @@ class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata): + block_size=c.block_size, + ) + ++ def snapshot_for_deferred_commit(self) -> "PleRecoverSSMMetadata": ++ c = self.recoverssm_commit ++ if c is None: ++ return self ++ return replace( ++ self, ++ recoverssm_commit=clone_commit_tensors(c, self.num_spec_decodes), ++ ) ++ + + class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + # Same as the GDN RecoverSSM builder: uniform spec-decode batches replay as FULL +@@ -245,6 +255,17 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + self._capturing = False + + def _get_ctx(self): ++ if self._ctx is not None: ++ fc = self.vllm_config.compilation_config.static_forward_context ++ live = [fc[name].kv_cache[0].data_ptr() for name in self.layer_names] ++ if [s.data_ptr() for s in self._ctx.conv_states] != live: ++ # Same lifecycle guard as the GDN builder: never commit into a ++ # KV cache that was torn down (profiling) and re-allocated. ++ raise RuntimeError( ++ "PLE RecoverSSM commit context is bound to a different KV " ++ "cache than its layers (stale context after a KV cache " ++ "re-allocation)" ++ ) + if self._ctx is None: + fc = self.vllm_config.compilation_config.static_forward_context + convs = [fc[name].kv_cache[0] for name in self.layer_names] +diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py +index b1339e9..135b3eb 100644 +--- a/vllm/v1/attention/backends/recoverssm_metadata.py ++++ b/vllm/v1/attention/backends/recoverssm_metadata.py +@@ -2,7 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + import abc +-from dataclasses import dataclass ++from dataclasses import dataclass, replace + + import torch + +@@ -25,6 +25,25 @@ class RecoverSSMMetadata(abc.ABC): + ) -> RecoverSSMPostprocessMetadata | None: + raise NotImplementedError + ++ def has_recoverssm_commit(self) -> bool: ++ """Whether this step's metadata carries a commit (any spec row).""" ++ return getattr(self, "recoverssm_commit", None) is not None ++ ++ def snapshot_for_deferred_commit(self) -> "RecoverSSMMetadata": ++ """A copy whose commit inputs no longer alias per-step buffers. ++ ++ Under pipeline parallelism a non-last rank learns the accepted counts ++ pp_size steps after its verify (the PP sampled-token broadcast), and ++ the steps in between rebuild the builders' persistent, graph-staged ++ buffers (spec state indices, query offsets, the batch-order block ++ table). The deferred commit must read this step's values, so they are ++ cloned here, on the current stream, before the next build overwrites ++ them. Implementations that do not support the PP deferral raise.""" ++ raise NotImplementedError( ++ f"{type(self).__name__} does not support the pipeline-parallel " ++ "RecoverSSM commit deferral" ++ ) ++ + + def recoverssm_spec_rows( + is_prefilling_cpu: torch.Tensor, +@@ -62,3 +81,23 @@ def recoverssm_request_indices( + if num_spec_decodes == 0 or int(rows[-1]) == num_spec_decodes - 1: + return None + return rows.to(torch.int32).to(device, non_blocking=True) ++ ++ ++def clone_commit_tensors(commit, num_spec_decodes: int): ++ """Clone the commit inputs a deferred (PP) commit reads, trimmed to the real ++ spec rows. Works for every RecoverSSM commit dataclass with the fields ++ state_indices, query_start_loc, request_indices, block_table and ++ num_computed_tokens.""" ++ n = num_spec_decodes ++ ++ def c(t: torch.Tensor | None) -> torch.Tensor | None: ++ return None if t is None else t.clone() ++ ++ return replace( ++ commit, ++ state_indices=commit.state_indices[:n].clone(), ++ query_start_loc=commit.query_start_loc[: n + 1].clone(), ++ request_indices=c(commit.request_indices), ++ block_table=c(commit.block_table), ++ num_computed_tokens=c(commit.num_computed_tokens), ++ ) +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index b28286a..924acee 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -1119,6 +1119,10 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + # profiling KV cache. + if (recoverssm := getattr(runner.model_state, "recoverssm", None)) is not None: + recoverssm.reset() ++ # Under PP the non-last ranks keep detached RecoverSSM steps on the pending ++ # PP entries; they reference the profiling KV cache just the same. ++ if (pp_handler := getattr(runner, "pp_handler", None)) is not None: ++ pp_handler.drop_deferred_states() + # Invalidate the align-mode Mamba group metadata cached from the + # profiling KVCacheConfig: the real (e.g. PP-projected) config may + # place Mamba layers into a different group layout, so it must be +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 246fe4c..1876f4b 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -1202,6 +1202,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): + self.req_states.draft_tokens + ) + if outputs is not None: ++ # Restore the step's deferred model state (RecoverSSM commits) ++ # so this postprocess commits the step that sampled them. ++ self.model_state.restore_deferred_state( ++ outputs.pop("deferred_state", None) ++ ) + self.postprocess_sampled(**outputs) + + def add_requests(self, scheduler_output: SchedulerOutput) -> None: +@@ -2306,6 +2311,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): + ): + raise RuntimeError("PLE test fault: exc_async_output") + all_decode_next = self.pp_handler.receive(input_batch) ++ # The accepted counts of this step arrive pp_size steps later; keep ++ # the state their postprocess needs with this step's pending entry ++ # (before the immediate prefill-chunk postprocess below). ++ self.pp_handler.attach_deferred_state( ++ self.model_state.detach_deferred_state() ++ ) + # Optimistically update num_computed_tokens for entire batch here. + # Will be adjusted for rejections if necessary in update_requests. + self.postprocess_num_computed_tokens(input_batch) +diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py +index 8ce5121..6f41111 100644 +--- a/vllm/v1/worker/gpu/model_states/interface.py ++++ b/vllm/v1/worker/gpu/model_states/interface.py +@@ -163,6 +163,19 @@ class ModelState(ABC): + ) -> None: + return None + ++ def detach_deferred_state(self) -> Any: ++ """Non-last PP rank, after this step's forward: hand over any state the ++ step's postprocess needs once its sampled tokens arrive (pp_size steps ++ later). The PP handler keeps it with the step's pending entry. None by ++ default (nothing deferred).""" ++ return None ++ ++ def restore_deferred_state(self, state: Any) -> None: ++ """Counterpart of `detach_deferred_state`, called right before the ++ deferred postprocess of the step that detached `state`.""" ++ if state is not None: ++ raise RuntimeError(f"{type(self).__name__} cannot restore {state!r}") ++ + @abstractmethod + def prepare_inputs_embeds( + self, +diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +index de0e1ce..b31a0ce 100644 +--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py ++++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +@@ -359,6 +359,20 @@ class MambaHybridModelState(DefaultModelState): + ) + return attn_metadata + ++ def detach_deferred_state(self) -> Any: ++ # Non-last PP rank: the RecoverSSM commit of this step has to wait for ++ # the accepted counts of the PP broadcast; snapshot its inputs now. ++ if self.recoverssm is None: ++ return None ++ return self.recoverssm.detach_step() ++ ++ def restore_deferred_state(self, state: Any) -> None: ++ if state is None: ++ return ++ if self.recoverssm is None: ++ raise RuntimeError("deferred RecoverSSM state without RecoverSSM") ++ self.recoverssm.restore_step(state) ++ + def postprocess_state( + self, + idx_mapping: torch.Tensor, +diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py +index 3cfbece..2aafa25 100644 +--- a/vllm/v1/worker/gpu/model_states/recoverssm.py ++++ b/vllm/v1/worker/gpu/model_states/recoverssm.py +@@ -1,25 +1,77 @@ + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + ++from dataclasses import dataclass + from typing import Any + + import torch + ++from vllm.logger import init_logger + from vllm.triton_utils import tl, triton + from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata + from vllm.v1.worker.utils import AttentionGroup + ++logger = init_logger(__name__) ++ ++ ++@dataclass(frozen=True) ++class DeferredRecoverSSMStep: ++ """One verify step's RecoverSSM commits, detached on a non-last PP rank. ++ ++ The rank verified the window at step T but receives the accepted counts ++ only when the PP broadcast of step T is consumed (pp_size steps later). ++ The metadata is a snapshot (`snapshot_for_deferred_commit`), so the steps ++ in between cannot overwrite its commit inputs. It rides on the PP ++ handler's pending entry for step T and is restored right before that ++ entry's postprocess.""" ++ ++ metadata: tuple[RecoverSSMMetadata, ...] ++ + + class RecoverSSMState: + """Coordinates RecoverSSM metadata between attention and postprocessing.""" + + def __init__(self) -> None: + self._step: tuple[RecoverSSMMetadata, ...] | None = None ++ # True while `_step` is a restored PP-deferred step: its commit must ++ # skip rows whose request was freed since the verify (their state ++ # blocks may already belong to another request). ++ self._step_is_deferred = False + + def reset(self) -> None: + """Drop the recorded step. Its metadata holds the commit contexts, which + reference the KV cache tensors they were built for.""" + self._step = None ++ self._step_is_deferred = False ++ ++ def detach_step(self) -> DeferredRecoverSSMStep | None: ++ """Take this step's commits for a deferred (non-last PP rank) commit. ++ ++ Returns None when the step has nothing to commit (no spec rows).""" ++ step = self._step ++ self._step = None ++ self._step_is_deferred = False ++ if step is None: ++ return None ++ live = tuple(m for m in step if m.has_recoverssm_commit()) ++ if not live: ++ return None ++ return DeferredRecoverSSMStep( ++ tuple(m.snapshot_for_deferred_commit() for m in live) ++ ) ++ ++ def restore_step(self, deferred: DeferredRecoverSSMStep | None) -> None: ++ """Make a detached step the one the next `commit_step` commits.""" ++ if self._step is not None: ++ # Every non-last-rank step detaches its record in sample_tokens, so ++ # only a capture-time record (for_capture=False under PIECEWISE) ++ # can be here; it has no sampled tokens and is never committed. ++ logger.warning_once( ++ "RecoverSSM: dropping a recorded step that was never committed " ++ "(expected only after a PIECEWISE capture)" ++ ) ++ self._step = None if deferred is None else deferred.metadata ++ self._step_is_deferred = deferred is not None + + def record_step( + self, +@@ -49,10 +101,26 @@ class RecoverSSMState: + num_accepted_tokens: torch.Tensor, + ) -> None: + step = self._step ++ deferred = self._step_is_deferred + self._step = None ++ self._step_is_deferred = False + if isinstance(num_sampled, int) or step is None: ++ if deferred and step is not None: ++ raise RuntimeError( ++ "RecoverSSM: a deferred PP commit was restored without " ++ "sampled counts" ++ ) + return + ++ if deferred: ++ # Under PP the consume filter marks freed/non-sampling rows with -1 ++ # in idx_mapping. A zero count makes the commit plan skip the row ++ # (commit_len 0): no state, boundary or conv write into blocks the ++ # request no longer owns. ++ num_sampled = torch.where( ++ idx_mapping >= 0, num_sampled, torch.zeros_like(num_sampled) ++ ) ++ + for metadata in step: + postprocess_meta = metadata.commit_recoverssm_state(num_sampled) + if postprocess_meta is None: +diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py +index ab1b991..0a54342 100644 +--- a/vllm/v1/worker/gpu/pp_utils.py ++++ b/vllm/v1/worker/gpu/pp_utils.py +@@ -3,7 +3,7 @@ + """Pipeline Parallelism utils for V2 Model Runner.""" + + from collections import deque +-from dataclasses import dataclass ++from dataclasses import dataclass, replace + + import numpy as np + import torch +@@ -32,6 +32,9 @@ class PendingRecv: + # detect requests aborted since then. + gen_at_receive_np: np.ndarray # [num_reqs] + draft_tokens: torch.Tensor | None = None # [num_reqs, num_speculative_steps] ++ # Model-state data the step's postprocess needs once the sampled tokens ++ # arrive (RecoverSSM: the snapshot of the step's deferred commits). ++ deferred_state: object | None = None + + + def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: +@@ -158,8 +161,34 @@ class PPHandler: + num_sampled=slot.num_sampled, + num_rejected=slot.num_rejected, + idx_mapping=idx_mapping, ++ deferred_state=slot.deferred_state, + ) + ++ def attach_deferred_state(self, state: object | None) -> None: ++ """Attach model-state data to the entry `receive` just reserved for this ++ step. Called right after `receive` in the same step.""" ++ assert not self.is_last_rank ++ if state is None: ++ return ++ slot = self.queue[-1] ++ if slot is None: ++ # receive() skipped the broadcast: no row of this step samples, so ++ # no row can have a speculative window to commit. ++ raise RuntimeError( ++ "deferred model state for a step without sampled outputs" ++ ) ++ self.queue[-1] = replace(slot, deferred_state=state) ++ ++ def drop_deferred_states(self) -> int: ++ """Release every pending entry's deferred state (profiling teardown: the ++ state references the KV cache it was built for). Returns the count.""" ++ dropped = 0 ++ for i, slot in enumerate(self.queue): ++ if slot is not None and slot.deferred_state is not None: ++ self.queue[i] = replace(slot, deferred_state=None) ++ dropped += 1 ++ return dropped ++ + def broadcast_drafts( + self, draft_tokens: torch.Tensor, input_batch: InputBatch + ) -> None: +-- +2.43.0 + diff --git a/patches/0081-recoverssm-tests-kernel-level-test-of-the-PP-deferre.patch b/patches/0081-recoverssm-tests-kernel-level-test-of-the-PP-deferre.patch new file mode 100644 index 0000000..c000bc0 --- /dev/null +++ b/patches/0081-recoverssm-tests-kernel-level-test-of-the-PP-deferre.patch @@ -0,0 +1,251 @@ +From a70876f4aacbad42d7c4e608b439e5af425578f5 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 21:19:19 -0400 +Subject: [PATCH 081/103] recoverssm, tests: kernel-level test of the + PP-deferred commit + +Runs the real Triton verify and commit kernels (CUDA, or CPU under TRITON_INTERPRET=1): verify A, verify B, then +commit A deferred. Committed states and conv windows equal the PP=1 order bit for bit, in none and align mode. +--- + .../test_recoverssm_pp_deferral_kernels.py | 229 ++++++++++++++++++ + 1 file changed, 229 insertions(+) + create mode 100644 tests/kernels/mamba/test_recoverssm_pp_deferral_kernels.py + +diff --git a/tests/kernels/mamba/test_recoverssm_pp_deferral_kernels.py b/tests/kernels/mamba/test_recoverssm_pp_deferral_kernels.py +new file mode 100644 +index 0000000..96dd8d7 +--- /dev/null ++++ b/tests/kernels/mamba/test_recoverssm_pp_deferral_kernels.py +@@ -0,0 +1,229 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The PP-deferred RecoverSSM commit with the real Triton kernels. ++ ++A non-last PP rank verifies micro-batch A, then verifies micro-batch B (whose ++build rewrites the builder's persistent buffers), and only then commits A with ++A's accepted counts. The committed states must equal the PP=1 order (verify A, ++commit A) bit for bit; a row freed in between must not be written. ++ ++Runs on CUDA, or on CPU under the Triton interpreter ++(``CUDA_VISIBLE_DEVICES= TRITON_INTERPRET=1``), which is how a GPU-less host runs it. ++ ++Failing controls (they flip): a snapshot that aliases the persistent buffers ++commits A's counts into B's blocks; a deferred commit without the freed-row ++mask writes the freed block. ++""" ++ ++import os ++from dataclasses import fields ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( ++ GDNRecoverSSMCommitContext, ++ gdn_recoverssm_verify, ++) ++from vllm.v1.attention.backends.gdn_recoverssm import ( ++ GDNRecoverSSMCommitMetadata, ++ GDNRecoverSSMMetadata, ++) ++from vllm.v1.worker.gpu.model_states.recoverssm import ( ++ DeferredRecoverSSMStep, ++ RecoverSSMState, ++) ++ ++if torch.cuda.is_available(): ++ DEV = "cuda" ++elif os.environ.get("TRITON_INTERPRET") == "1": ++ DEV = "cpu" ++else: ++ pytest.skip("needs CUDA or TRITON_INTERPRET=1", allow_module_level=True) ++ ++H, HV, K, V, T = 2, 4, 32, 32, 4 # T = 1 + K_spec (K_spec = 3, prod) ++NB, CONV_DIM, MAX_REQS, BLOCK = 16, 8, 8, 8 ++A_BLOCKS, B_BLOCKS = [2, 5], [7, 9, 11] ++A_ACC = [2, 3] ++ ++ ++class _Pool: ++ def __init__(self, seed: int = 0) -> None: ++ g = torch.Generator(device="cpu").manual_seed(seed) ++ r = lambda *s: torch.randn(*s, generator=g) # noqa: E731 ++ self.ckpt = (r(NB, HV, V, K) * 0.05).float().to(DEV) ++ self.replay = torch.zeros(NB, HV, T, V + K + 1, device=DEV) ++ self.conv = (r(NB, CONV_DIM, 3 + T - 1)).to(torch.bfloat16).to(DEV) ++ self.A_log = (torch.rand(HV, generator=g) * 2 - 1).float().to(DEV) ++ self.dt_bias = (r(HV) * 0.5).bfloat16().to(DEV) ++ self._g = g ++ ++ def window(self, n: int) -> dict: ++ r = lambda *s: torch.randn(*s, generator=self._g) # noqa: E731 ++ tot = n * T ++ return dict( ++ a=r(tot, HV).bfloat16().to(DEV), ++ b=r(tot, HV).bfloat16().to(DEV), ++ q=r(1, tot, H, K).bfloat16().to(DEV), ++ k=r(1, tot, H, K).bfloat16().to(DEV), ++ v=r(1, tot, HV, V).bfloat16().to(DEV), ++ ) ++ ++ def context(self) -> GDNRecoverSSMCommitContext: ++ return GDNRecoverSSMCommitContext.from_tensors( ++ [self.conv], [self.ckpt], [self.replay], spec_query_len=T, ++ max_num_reqs=MAX_REQS, ++ ) ++ ++ ++class _Buffers: ++ """The builder's persistent, graph-staged buffers.""" ++ ++ def __init__(self) -> None: ++ self.sidx = torch.full((MAX_REQS, T), -1, dtype=torch.int32, device=DEV) ++ self.qsl = torch.zeros(MAX_REQS + 1, dtype=torch.int32, device=DEV) ++ self.bt = torch.zeros(MAX_REQS, 4, dtype=torch.int32, device=DEV) ++ ++ ++def _build(buf, blocks, ctx, *, align) -> GDNRecoverSSMMetadata: ++ n = len(blocks) ++ buf.sidx[:n, 0] = torch.tensor(blocks, dtype=torch.int32, device=DEV) ++ buf.qsl[: n + 1] = torch.arange(0, (n + 1) * T, T, dtype=torch.int32, device=DEV) ++ for i, b in enumerate(blocks): ++ # align, 6 computed tokens: column 0 = block b holds tokens 0..7 and the ++ # running state (the checkpoint), column 1 = block b + 1 (tokens 8..15). ++ # Accepting 2 ends exactly on the boundary (final stays in b); accepting ++ # 3 crosses it: final -> b + 1, boundary state -> b, the source slot ++ # (the fold-ordering case of #58863 commit 6). ++ buf.bt[i] = torch.tensor([b, b + 1, 0, 0], dtype=torch.int32, device=DEV) ++ commit = GDNRecoverSSMCommitMetadata( ++ state_indices=buf.sidx[:, :1], ++ query_start_loc=buf.qsl, ++ request_indices=None, ++ block_table=buf.bt[:n] if align else None, ++ # 6 computed tokens: a window accepting >= 2 crosses the boundary at 8 ++ num_computed_tokens=torch.full((n,), 6, dtype=torch.int32, device=DEV) ++ if align ++ else None, ++ block_size=BLOCK if align else None, ++ ) ++ req = dict( ++ num_prefills=0, num_prefill_tokens=0, num_decodes=0, num_decode_tokens=0, ++ num_spec_decodes=n, num_spec_decode_tokens=n * T, num_actual_tokens=n * T, ++ ) ++ kw = { ++ f.name: req.get(f.name) ++ for f in fields(GDNRecoverSSMMetadata) ++ if f.init and f.name not in ("recoverssm_commit", "recoverssm_context") ++ } ++ return GDNRecoverSSMMetadata(**kw, recoverssm_commit=commit, recoverssm_context=ctx) ++ ++ ++def _verify(pool: _Pool, buf: _Buffers, n: int, w: dict) -> None: ++ gdn_recoverssm_verify( ++ pool.A_log, w["a"], w["b"], pool.dt_bias, w["q"], w["k"], w["v"], ++ checkpoint_state=pool.ckpt, replay_cache=pool.replay, ++ query_start_loc=buf.qsl[: n + 1], state_indices=buf.sidx[:n, 0], ++ spec_query_len=T, ++ ) ++ ++ ++def _record(state, meta): ++ from types import SimpleNamespace ++ ++ g = SimpleNamespace(layer_names=["l0"]) ++ state.record_step({"l0": meta}, [[g]], for_capture=False) ++ ++ ++def _commit(state, idx_mapping): ++ acc = torch.tensor(A_ACC, dtype=torch.int32, device=DEV) ++ state.commit_step( ++ acc, ++ torch.tensor(idx_mapping, dtype=torch.int32, device=DEV), ++ state_indices=torch.zeros(MAX_REQS, dtype=torch.int32, device=DEV), ++ num_accepted_tokens=torch.ones(MAX_REQS, dtype=torch.int32, device=DEV), ++ ) ++ ++ ++def _pp1_reference(align: bool) -> _Pool: ++ pool, buf, state = _Pool(), _Buffers(), RecoverSSMState() ++ wa = pool.window(2) ++ _record(state, _build(buf, A_BLOCKS, pool.context(), align=align)) ++ _verify(pool, buf, 2, wa) ++ _commit(state, [0, 1]) ++ return pool ++ ++ ++def _pp2(align: bool, *, alias=False, idx_mapping=(0, 1), mask=True): ++ pool, buf, state = _Pool(), _Buffers(), RecoverSSMState() ++ wa = pool.window(2) ++ wb = pool.window(3) ++ ctx = pool.context() ++ # step T: verify A, detach its commit (sample_tokens on a non-last rank) ++ meta_a = _build(buf, A_BLOCKS, ctx, align=align) ++ _record(state, meta_a) ++ _verify(pool, buf, 2, wa) ++ deferred = state.detach_step() ++ assert deferred is not None ++ if alias: # control: a detach that keeps the live views (no snapshot) ++ deferred = DeferredRecoverSSMStep((meta_a,)) ++ # step T+1: micro-batch B's build rewrites the buffers, B verifies ++ _record(state, _build(buf, B_BLOCKS, ctx, align=align)) ++ _verify(pool, buf, 3, wb) ++ state.detach_step() # B goes to its own pending entry ++ before = (pool.ckpt.clone(), pool.conv.clone()) ++ # step T+2: A's broadcast consumed -> restore, then the deferred commit ++ state.restore_step(deferred) ++ if not mask: ++ state._step_is_deferred = False ++ _commit(state, list(idx_mapping)) ++ return pool, before ++ ++ ++def _blocks_for(align: bool, blocks): ++ # every block a row's commit may write (see _build) ++ if not align: ++ return list(blocks) ++ return list(blocks) + [b + 1 for b in blocks] ++ ++ ++@pytest.mark.parametrize("align", [False, True]) ++def test_deferred_commit_is_bit_identical_to_pp1(align): ++ ref = _pp1_reference(align) ++ pool, before = _pp2(align) ++ touched = _blocks_for(align, A_BLOCKS) ++ for b in touched: ++ assert torch.equal(pool.ckpt[b], ref.ckpt[b]), b ++ assert torch.equal(pool.conv[b], ref.conv[b]), b ++ others = [b for b in range(NB) if b not in touched] ++ assert torch.equal(pool.ckpt[others], before[0][others]) ++ # and A's blocks did change (the commit really ran) ++ assert any(not torch.equal(pool.ckpt[b], before[0][b]) for b in A_BLOCKS) ++ ++ ++@pytest.mark.parametrize("align", [False, True]) ++def test_control_aliasing_snapshot_breaks_it(align): ++ ref = _pp1_reference(align) ++ pool, _ = _pp2(align, alias=True) ++ assert any( ++ not torch.equal(pool.ckpt[b], ref.ckpt[b]) for b in _blocks_for(align, A_BLOCKS) ++ ) ++ ++ ++@pytest.mark.parametrize("align", [False, True]) ++def test_freed_row_is_not_written(align): ++ ref = _pp1_reference(align) ++ pool, before = _pp2(align, idx_mapping=(0, -1)) ++ live = _blocks_for(align, A_BLOCKS[:1]) ++ freed = _blocks_for(align, A_BLOCKS[1:]) ++ for b in live: ++ assert torch.equal(pool.ckpt[b], ref.ckpt[b]) ++ for b in freed: ++ assert torch.equal(pool.ckpt[b], before[0][b]), b ++ assert torch.equal(pool.conv[b], before[1][b]), b ++ ++ ++@pytest.mark.parametrize("align", [False, True]) ++def test_control_unmasked_commit_writes_the_freed_block(align): ++ pool, before = _pp2(align, idx_mapping=(0, -1), mask=False) ++ assert not torch.equal(pool.ckpt[A_BLOCKS[1]], before[0][A_BLOCKS[1]]) +-- +2.43.0 + diff --git a/patches/0082-recoverssm-tests-CPU-harness-for-the-kernel-suites-o.patch b/patches/0082-recoverssm-tests-CPU-harness-for-the-kernel-suites-o.patch new file mode 100644 index 0000000..d7db67c --- /dev/null +++ b/patches/0082-recoverssm-tests-CPU-harness-for-the-kernel-suites-o.patch @@ -0,0 +1,79 @@ +From 029f835d8c5becced0382796057a3152bc2b41a3 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 21:27:02 -0400 +Subject: [PATCH 082/103] recoverssm, tests: CPU harness for the kernel suites + on a GPU-less host + +tests/merlin_cpu/mk_cpu_tests.sh makes CPU copies of the CUDA-only RecoverSSM kernel tests (device strings and +skip marks only; it refuses a copy with residual cuda strings). tests/merlin_cpu/triton_interp_fix.py is a pytest +plugin that works around a Triton 3.7 interpreter bug with non-constexpr arguments. +--- + tests/merlin_cpu/mk_cpu_tests.sh | 25 +++++++++++++++++++++++++ + tests/merlin_cpu/triton_interp_fix.py | 23 +++++++++++++++++++++++ + 2 files changed, 48 insertions(+) + create mode 100644 tests/merlin_cpu/mk_cpu_tests.sh + create mode 100644 tests/merlin_cpu/triton_interp_fix.py + +diff --git a/tests/merlin_cpu/mk_cpu_tests.sh b/tests/merlin_cpu/mk_cpu_tests.sh +new file mode 100644 +index 0000000..c4df611 +--- /dev/null ++++ b/tests/merlin_cpu/mk_cpu_tests.sh +@@ -0,0 +1,25 @@ ++#!/usr/bin/env bash ++# CPU copies of CUDA-only RecoverSSM kernel tests, for a GPU-less build host. ++# The kernels then run under the Triton interpreter: ++# ++# tests/merlin_cpu/mk_cpu_tests.sh ++# cd && CUDA_VISIBLE_DEVICES= TRITON_INTERPRET=1 PYTHONPATH=.:

\ ++# python -m pytest --noconftest -p triton_interp_fix /test_*_cpu.py ++# ++# CUDA_VISIBLE_DEVICES= (empty) keeps vllm.triton_utils from disabling Triton when no GPU ++# driver is active; TRITON_INTERPRET=1 runs every kernel on CPU tensors. Only device ++# strings and CUDA skip marks are rewritten; the test bodies are unchanged. ++set -euo pipefail ++tree=$1; out=$2; mkdir -p "$out" ++for f in tests/kernels/mamba/test_recoverssm_gdn.py tests/kernels/mamba/test_recoverssm_ple.py \ ++ tests/models/kimi_k3/test_recoverssm_boundary.py; do ++ [ -f "$tree/$f" ] || { echo "skip (absent): $f"; continue; } ++ b=$(basename "$f" .py) ++ sed -e 's/^pytestmark = pytest.mark.skipif(/pytestmark = [] and pytest.mark.skipif(/' \ ++ -e 's/device="cuda"/device="cpu"/g' -e "s/device='cuda'/device='cpu'/g" \ ++ -e 's/torch.device("cuda")/torch.device("cpu")/g' -e 's/^DEV = "cuda"/DEV = "cpu"/' \ ++ -e 's/dev = "cuda"/dev = "cpu"/g' -e 's/\.cuda()/.cpu()/g' \ ++ -e 's/current_platform.device_type/"cpu"/g' \ ++ "$tree/$f" > "$out/${b}_cpu.py" ++ if grep -nE "\"cuda\"|'cuda'" "$out/${b}_cpu.py"; then echo "residual cuda refs in ${b}_cpu.py" >&2; exit 1; fi ++done +diff --git a/tests/merlin_cpu/triton_interp_fix.py b/tests/merlin_cpu/triton_interp_fix.py +new file mode 100644 +index 0000000..542d946 +--- /dev/null ++++ b/tests/merlin_cpu/triton_interp_fix.py +@@ -0,0 +1,23 @@ ++"""pytest plugin (-p triton_interp_fix): Triton 3.7 interpreter shim for CPU runs. ++ ++The interpreter converts a non-constexpr Python bool argument with an int32 ++numpy array and then rejects it against the int1 type ("numpy data itemsize ++(32 bits) exceeds dtype primitive_bitwidth (1 bits)"). Kernels on the stock ++reference paths (e.g. ple_conv) pass such bools. Convert them as np.bool_. ++Harness-only: it changes nothing in the code under test.""" ++ ++import numpy as np ++import triton.language as tl ++from triton.runtime import interpreter as _interp ++ ++_orig = _interp._implicit_cvt ++ ++ ++def _implicit_cvt(arg): ++ if isinstance(arg, bool): ++ handle = _interp.TensorHandle(np.array([arg], dtype=np.bool_), tl.int1) ++ return tl.tensor(handle, tl.int1) ++ return _orig(arg) ++ ++ ++_interp._implicit_cvt = _implicit_cvt +-- +2.43.0 + diff --git a/patches/0083-worker-report-a-profiling-KV-cache-that-survives-the.patch b/patches/0083-worker-report-a-profiling-KV-cache-that-survives-the.patch new file mode 100644 index 0000000..622884f --- /dev/null +++ b/patches/0083-worker-report-a-profiling-KV-cache-that-survives-the.patch @@ -0,0 +1,160 @@ +From 6a9af861a09b0082f8e0e1d15eb046881e9abe87 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Thu, 1 Oct 2026 21:29:59 -0400 +Subject: [PATCH 083/103] worker: report a profiling KV cache that survives the + teardown + +_teardown_profiling_state measures the storage behind the profiling KV caches on entry and the allocated memory +before and after each release, and logs any profiling KV cache that is still alive afterwards. +--- + .../worker/gpu/test_recoverssm_pp_deferral.py | 77 +++++++++++++++++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 40 ++++++++++ + 2 files changed, 117 insertions(+) + +diff --git a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +index 9420aaa..7fa3162 100644 +--- a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py ++++ b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +@@ -406,3 +406,80 @@ def test_runner_wiring_order(): + + calls = _calls_in(GPUModelRunner.update_pp_decode_requests) + assert calls.index("restore_deferred_state") < calls.index("postprocess_sampled") ++ ++ ++def test_still_held_lower_bound(): ++ from vllm.v1.worker.gpu.cudagraph_utils import _profiling_kv_still_held ++ ++ mib = 1 << 20 ++ assert _profiling_kv_still_held(300 * mib, 1000 * mib, 700 * mib) == 0 ++ assert _profiling_kv_still_held(300 * mib, 1000 * mib, 1000 * mib) == 300 * mib ++ assert _profiling_kv_still_held(300 * mib, 1000 * mib, 900 * mib) == 200 * mib ++ # other frees count against the cache first: never a false alarm ++ assert _profiling_kv_still_held(300 * mib, 1000 * mib, 100 * mib) == 0 ++ ++ ++def _fake_runner(kv, pp_handler): ++ model_state = SimpleNamespace( ++ _mamba_ctx=object(), ++ recoverssm=RecoverSSMState(), ++ _mamba_group_ids=[1], ++ _mamba_spec=object(), ++ supports_mm_inputs=False, ++ ) ++ return SimpleNamespace( ++ model_state=model_state, ++ pp_handler=pp_handler, ++ kv_caches=[kv], ++ attn_groups=[[object()]], ++ kv_cache_config=object(), ++ cudagraph_manager=object(), ++ compilation_config=SimpleNamespace(static_forward_context={}), ++ model=None, ++ cache_config=SimpleNamespace(num_gpu_blocks=8), ++ lora_config=None, ++ maybe_remove_all_loras=lambda cfg: None, ++ ) ++ ++ ++@pytest.mark.parametrize("case", ["clean", "stray_holder", "no_pp_purge"]) ++def test_teardown_reports_a_held_profiling_kv_cache(monkeypatch, case): ++ """End to end on the real _teardown_profiling_state with a fake runner and ++ a fake allocator whose allocated bytes include the KV cache while anything ++ still references it. The PP-queue deferred step is released by the purge; ++ a stray holder, or the same deferred step without the purge (controls), ++ is reported.""" ++ from vllm.v1.worker.gpu import cudagraph_utils as cu ++ ++ kv_bytes = 350 << 20 ++ kv = torch.zeros(4) ++ ref = weakref.ref(kv) ++ monkeypatch.setattr(cu, "_profiling_kv_cache_bytes", lambda caches: kv_bytes) ++ monkeypatch.setattr( ++ torch.cuda, ++ "memory_allocated", ++ lambda *a, **k: (1 << 30) + (kv_bytes if ref() is not None else 0), ++ ) ++ monkeypatch.setattr(torch.accelerator, "synchronize", lambda *a, **k: None) ++ monkeypatch.setattr(torch.accelerator, "empty_cache", lambda *a, **k: None) ++ monkeypatch.setattr(cu, "release_hisparse_profiling_cache", lambda fc: None) ++ warnings: list[str] = [] ++ monkeypatch.setattr( ++ cu.logger, "warning", lambda msg, *a: warnings.append(msg % a) ++ ) ++ ++ h = _handler() ++ _receive(h, [0], [1]) ++ h.attach_deferred_state(_Holder(kv)) # what a deferred RecoverSSM step holds ++ stray = [_Holder(kv)] if case == "stray_holder" else [] ++ if case == "no_pp_purge": ++ monkeypatch.setattr(h, "drop_deferred_states", lambda: 0) ++ runner = _fake_runner(kv, h) ++ del kv ++ cu._teardown_profiling_state(runner) ++ if case != "clean": ++ assert len(warnings) == 1 and "350.0 MiB" in warnings[0], warnings ++ else: ++ assert warnings == [] ++ assert ref() is None ++ del stray +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 924acee..1cdceb3 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -1108,10 +1108,35 @@ def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: + runner.cache_config.num_gpu_blocks = minimal_config.num_blocks + + ++def _profiling_kv_cache_bytes(kv_caches: Iterable[Any]) -> int: ++ """Bytes of the distinct CUDA storages behind the runner's KV cache views.""" ++ storages: dict[int, int] = {} ++ for cache in kv_caches: ++ for t in cache if isinstance(cache, (list, tuple)) else (cache,): ++ if isinstance(t, torch.Tensor) and t.is_cuda: ++ st = t.untyped_storage() ++ storages[st.data_ptr()] = st.nbytes() ++ return sum(storages.values()) ++ ++ ++def _profiling_kv_still_held( ++ kv_bytes: int, allocated_before: int, allocated_after: int ++) -> int: ++ """Lower bound on the profiling KV cache bytes still allocated after the ++ teardown: whatever the teardown freed counts against the cache first, so a ++ positive result means something still references the cache (the RecoverSSM ++ commit-context class of leak, ~26 GiB on H200; against a sub-GiB per-card ++ per-card cushion even a small one matters). One-sided: never a false alarm, ++ may miss a leak masked by other frees.""" ++ return max(kv_bytes - max(allocated_before - allocated_after, 0), 0) ++ ++ + def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + """Release the profiling KV cache and captured graphs while keeping model + weights, so the real ``initialize_kv_cache`` starts from a clean slate.""" + torch.accelerator.synchronize() ++ kv_bytes = _profiling_kv_cache_bytes(getattr(runner, "kv_caches", ())) ++ allocated_before = torch.cuda.memory_allocated() if kv_bytes else 0 + if hasattr(runner.model_state, "_mamba_ctx"): + runner.model_state._mamba_ctx = None + # A PIECEWISE profiling capture records its RecoverSSM step (only FULL +@@ -1154,3 +1179,18 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + runner.maybe_remove_all_loras(runner.lora_config) + gc.collect() + torch.accelerator.empty_cache() ++ if kv_bytes: ++ # last, after every release above (one-sided check) ++ gc.collect() ++ torch.accelerator.synchronize() ++ held = _profiling_kv_still_held( ++ kv_bytes, allocated_before, torch.cuda.memory_allocated() ++ ) ++ if held > (1 << 20): ++ logger.warning( ++ "Profiling teardown: at least %.1f MiB of the %.1f MiB profiling KV " ++ "cache is still allocated; something still references it (it will " ++ "be missing from the real KV cache / free-memory cushion)", ++ held / (1 << 20), ++ kv_bytes / (1 << 20), ++ ) +-- +2.43.0 + diff --git a/patches/0084-worker-gate-the-graph-fallback-and-fail-on-profiling.patch b/patches/0084-worker-gate-the-graph-fallback-and-fail-on-profiling.patch new file mode 100644 index 0000000..231378d --- /dev/null +++ b/patches/0084-worker-gate-the-graph-fallback-and-fail-on-profiling.patch @@ -0,0 +1,87 @@ +From 62bc4977ff293ee2c726b8007744cc78685f499a Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Fri, 2 Oct 2026 04:07:45 +0100 +Subject: [PATCH 084/103] worker: gate the graph fallback and fail on profiling + leaks + +The breakable-CUDA-graph fallback in the compilation config now applies only with --use-replayssm, and a +profiling KV cache that survives the teardown fails the boot instead of only being logged +(VLLM_FAIL_ON_PROFILING_KV_LEAK=0 turns it back into a warning). +--- + .../worker/gpu/test_recoverssm_pp_deferral.py | 11 +++++++++++ + vllm/config/compilation.py | 5 ++++- + vllm/v1/worker/gpu/cudagraph_utils.py | 17 ++++++++++++----- + 3 files changed, 27 insertions(+), 6 deletions(-) + +diff --git a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +index 7fa3162..aab545a 100644 +--- a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py ++++ b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py +@@ -483,3 +483,14 @@ def test_teardown_reports_a_held_profiling_kv_cache(monkeypatch, case): + assert warnings == [] + assert ref() is None + del stray ++ ++ ++@pytest.mark.parametrize("case", ["request_id_reuse", "prefix_cache_shared_block"]) ++def test_lifetime_scenario_is_explicitly_not_synthetic(case): ++ """Guard the CPU protocol suite from being mistaken for MRV2 scheduling coverage. ++ ++ The scheduler and prefix-cache lifetime regression is exercised by the ++ GPU integration runs, which are not part of this tree. This CPU suite ++ only validates the PP queue protocol. ++ """ ++ pytest.skip(f"{case} requires the TP2xPP2 MRV2 integration harness") +diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py +index 791975b..bd63f99 100644 +--- a/vllm/config/compilation.py ++++ b/vllm/config/compilation.py +@@ -1452,7 +1452,10 @@ class CompilationConfig: + self.splitting_ops_contain_attention() + or self.use_inductor_graph_partition + ) +- ) or envs.VLLM_USE_BREAKABLE_CUDAGRAPH: ++ ) or ( ++ envs.VLLM_USE_BREAKABLE_CUDAGRAPH ++ and self.cache_config.use_replayssm ++ ): + # Breakable CUDA graphs break at the attention kernels at runtime, + # with or without torch.compile. + msg += ( +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 1cdceb3..00eee6b 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -3,6 +3,7 @@ + import contextlib + import gc + import itertools ++import os + from collections import defaultdict + from collections.abc import Callable, Iterable, Iterator + from contextlib import ExitStack +@@ -1187,10 +1188,16 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + kv_bytes, allocated_before, torch.cuda.memory_allocated() + ) + if held > (1 << 20): +- logger.warning( ++ message = ( + "Profiling teardown: at least %.1f MiB of the %.1f MiB profiling KV " + "cache is still allocated; something still references it (it will " +- "be missing from the real KV cache / free-memory cushion)", +- held / (1 << 20), +- kv_bytes / (1 << 20), +- ) ++ "be missing from the real KV cache / free-memory cushion)" ++ ) % (held / (1 << 20), kv_bytes / (1 << 20)) ++ # Fail boot by default: this memory is absent from the real KV ++ # pool and can invalidate capacity planning. Operators may opt out ++ # temporarily while diagnosing with VLLM_FAIL_ON_PROFILING_KV_LEAK=0. ++ if os.getenv("VLLM_FAIL_ON_PROFILING_KV_LEAK", "1").lower() not in ( ++ "0", "false", "no" ++ ): ++ raise RuntimeError(message) ++ logger.warning(message) +-- +2.43.0 + diff --git a/patches/0085-recoverssm-use-replayssm-without-MTP-falls-back-to-t.patch b/patches/0085-recoverssm-use-replayssm-without-MTP-falls-back-to-t.patch new file mode 100644 index 0000000..55436ff --- /dev/null +++ b/patches/0085-recoverssm-use-replayssm-without-MTP-falls-back-to-t.patch @@ -0,0 +1,81 @@ +From a42c2e52b5765a7bcd896e97cecdb18a0f3a2764 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 4 Oct 2026 07:33:26 -0400 +Subject: [PATCH 085/103] recoverssm: --use-replayssm without MTP falls back to + the stock GDN path + +On Qwen4Exp the flag selects RecoverSSM, which needs speculative decoding. With num_speculative_tokens=0 the +validator refused to start; it now logs a warning, clears use_replayssm (the no-flag path) and continues, so one +launch line serves with and without MTP. Other architectures keep the error. Config tests cover both. +--- + tests/config/test_recoverssm_config.py | 27 +++++++++++++++++++++++--- + vllm/config/vllm.py | 14 +++++++++---- + 2 files changed, 34 insertions(+), 7 deletions(-) + +diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py +index 2fba7eb..352b8bf 100644 +--- a/tests/config/test_recoverssm_config.py ++++ b/tests/config/test_recoverssm_config.py +@@ -122,9 +122,30 @@ def test_flag_unset_keeps_the_stock_path(): + assert not uses_recoverssm(cfg.cache_config, 3) + + +-def test_qwen4exp_replayssm_requires_speculative_decoding(): +- with pytest.raises(ValueError, match="requires speculative decoding"): +- _validator()(_config(TARGET, num_spec=0)) ++@pytest.mark.parametrize( ++ "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM"] ++) ++def test_qwen4exp_replayssm_without_spec_falls_back_to_stock(arch): ++ # Fork (v2.5): no speculative decoding -> the flag is dropped, not refused. ++ cfg = _config(arch, num_spec=0) ++ _validator()(cfg) ++ assert cfg.cache_config.use_replayssm is False ++ assert cfg.cache_config.use_recoverssm is False ++ assert not uses_recoverssm(cfg.cache_config, 0) ++ # The fallback is the stock path: the RecoverSSM runtime checks no longer apply. ++ pp2 = _config( ++ arch, num_spec=0, mamba_config__backend="flashinfer", use_v2_model_runner=False ++ ) ++ _validator()(pp2) ++ assert pp2.cache_config.use_replayssm is False ++ ++ ++def test_fallback_is_qwen4exp_only(): ++ # Other ReplaySSM architectures keep their plain ReplaySSM path without drafts. ++ cfg = _config("NemotronHForCausalLM", num_spec=0, use_v2_model_runner=False) ++ _validator()(cfg) ++ assert cfg.cache_config.use_replayssm is True ++ assert cfg.cache_config.use_recoverssm is False + + + def test_architecture_without_replayssm_support_is_rejected(): +diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py +index d07579e..2538e80 100644 +--- a/vllm/config/vllm.py ++++ b/vllm/config/vllm.py +@@ -3240,11 +3240,17 @@ class VllmConfig: + and not self.cache_config.use_recoverssm + ): + # Qwen4Exp's GDN and PLE layers have no plain ReplaySSM path; with them the +- # flag selects RecoverSSM, which needs speculative decoding. +- raise ValueError( +- "--use-replayssm on Qwen4Exp selects GDN RecoverSSM and requires " +- "speculative decoding (num_speculative_tokens > 0)" ++ # flag selects RecoverSSM, which needs speculative decoding. Fork (v2.5): ++ # without speculative decoding there are no drafts to recover from, so the ++ # flag falls back to the stock path (as if it were not passed) instead of ++ # refusing to start; one launch line then serves with and without MTP. ++ logger.warning( ++ "--use-replayssm on Qwen4Exp selects GDN RecoverSSM, which needs " ++ "speculative decoding; num_speculative_tokens=0, so the stock GDN " ++ "path is used (as without --use-replayssm)" + ) ++ self.cache_config.use_replayssm = False ++ return self + if self.cache_config.use_recoverssm: + if self.model_config is not None and self.model_config.architecture not in ( + "KimiLinearForCausalLM", +-- +2.43.0 + diff --git a/patches/0086-sched-keep-the-running-request-s-Mamba-prefix-checkp.patch b/patches/0086-sched-keep-the-running-request-s-Mamba-prefix-checkp.patch new file mode 100644 index 0000000..66c2085 --- /dev/null +++ b/patches/0086-sched-keep-the-running-request-s-Mamba-prefix-checkp.patch @@ -0,0 +1,957 @@ +From c58638c919ba9acc7624a4e439129e6ccd394152 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Fri, 2 Oct 2026 04:09:20 -0400 +Subject: [PATCH 086/103] sched: keep the running request's Mamba prefix + checkpoint referenced until it finishes + +In Mamba align mode a state block leaves the block table once the request has computed one block past it. When +that block is hashed it is the request's prefix checkpoint, and releasing it let a long concurrent session evict +it, so the next turn re-prefilled the whole conversation. The manager now keeps the newest hashed checkpoint +referenced until the request finishes, and admission reserves that block. VLLM_E2_MAMBA_RETAIN_CHECKPOINT +(default 1) switches it off. Scheduler tests cover eviction, admission, preemption and freeing. +--- + tests/test_envs.py | 23 + + .../core/test_mamba_checkpoint_retention.py | 563 ++++++++++++++++++ + tests/v1/core/test_prefix_caching.py | 22 +- + vllm/envs.py | 25 +- + vllm/v1/core/kv_cache_coordinator.py | 22 + + vllm/v1/core/single_type_kv_cache_manager.py | 109 +++- + 6 files changed, 753 insertions(+), 11 deletions(-) + create mode 100644 tests/v1/core/test_mamba_checkpoint_retention.py + +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 9ddb51e..bcd7b33 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -648,6 +648,9 @@ _FORK_IGNORED = [ + # test_read_site_status_matches_the_tables pins that the readers exist. + ("VLLM_E2_PLE_FLAG_CHECK", "1", "warn"), + ("VLLM_E2_MAMBA_SLOT_OWNERSHIP", " RAISE ", "raise"), ++ # Mamba prefix-checkpoint retention (v2.5): scheduler-side, default ON; "0" ++ # turns it off. Read in single_type_kv_cache_manager.py (MambaManager). ++ ("VLLM_E2_MAMBA_RETAIN_CHECKPOINT", "0", False), + # Medium-path counter reader: a host-side log cadence, and the older slot-count + # fallback that this build never consults. Registered and ignored; before + # the read-site census neither classification was asserted anywhere. +@@ -1061,6 +1064,7 @@ def test_census_discriminates_both_worlds(): + _MUST_BE_READ = { + "VLLM_E2_PLE_FLAG_CHECK": "vllm/v1/ple_offload/connector.py", + "VLLM_E2_MAMBA_SLOT_OWNERSHIP": "vllm/v1/worker/mamba_utils.py", ++ "VLLM_E2_MAMBA_RETAIN_CHECKPOINT": "vllm/v1/core/single_type_kv_cache_manager.py", + } + _INERT_NO_READ_SITE = ["VLLM_PLE_FP8_GLOBAL_SCALE"] + +@@ -1105,3 +1109,22 @@ def test_read_site_scan_discriminates_both_worlds(): + assert "VLLM_E2_PLE_FLAG_CHECK" not in _read_sites(mention) + assert _read_sites(via_envs)["VLLM_E2_PLE_FLAG_CHECK"] == {"a.py"} + assert _read_sites(via_getenv)["VLLM_PLE_FP8_GLOBAL_SCALE"] == {"a.py"} ++ ++ ++def test_mamba_retain_checkpoint_switch_defaults_on_and_is_strict( ++ monkeypatch: pytest.MonkeyPatch, ++): ++ """Unset is ON (the fix ships enabled); "0" is the only way to the old ++ release-on-retire behaviour, and a misspelt value is a boot error rather ++ than a silent arm flip.""" ++ name = "VLLM_E2_MAMBA_RETAIN_CHECKPOINT" ++ monkeypatch.delenv(name, raising=False) ++ assert _fresh_env_value(name) is True ++ monkeypatch.setenv(name, "1") ++ assert _fresh_env_value(name) is True ++ monkeypatch.setenv(name, "0") ++ assert _fresh_env_value(name) is False ++ for bad in ("off", "false", "", " 0"): ++ monkeypatch.setenv(name, bad) ++ with pytest.raises(ValueError, match=name): ++ _fresh_env_value(name) +diff --git a/tests/v1/core/test_mamba_checkpoint_retention.py b/tests/v1/core/test_mamba_checkpoint_retention.py +new file mode 100644 +index 0000000..1a4ebe5 +--- /dev/null ++++ b/tests/v1/core/test_mamba_checkpoint_retention.py +@@ -0,0 +1,563 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Align-mode Mamba prefix-checkpoint retention (v2.5, HF discussion #2). ++ ++In align mode a recurrent-state block leaves the request's block table as soon ++as the request has computed one block past it (``remove_skipped_blocks``). When ++that block is hashed it is the request's GDN prefix checkpoint (under ++``prefix_cache_retention_interval=0``, the prompt's replay boundary), and the ++old code returned it to the free queue there and then -- mid-request, long ++before the request's own attention blocks, which stay referenced until it ++finishes. Other requests' prefills pop the free queue from the head, so the ++checkpoint was evicted first; the hybrid hit is the shortest group's, so that ++one block zeroed the next turn's reuse (measured: FP32 state, 8 x ~80K agent ++conversations, cached fraction 0.000 with 0 preemptions). ++ ++``VLLM_E2_MAMBA_RETAIN_CHECKPOINT=1`` (default) keeps the newest hashed retired ++state referenced until the request is freed, then releases it after the ++request's attention blocks. ``=0`` is the old behaviour and is the FAILING ++control in every test below. ++ ++These tests drive ``KVCacheManager`` through the calls the scheduler makes: ++``get_computed_blocks`` -> ``allocate_slots`` per step (chunk ends from the real ++``Scheduler._mamba_block_aligned_split``, MTP lookahead slots, draft tokens) -> ++``free`` / ``pop_blocks_for_free``. CPU only. ++""" ++ ++from collections.abc import Callable ++from dataclasses import replace ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm import envs ++from vllm.sampling_params import SamplingParams ++from vllm.utils.hashing import sha256 ++from vllm.v1.core.kv_cache_manager import KVCacheManager ++from vllm.v1.core.kv_cache_utils import ( ++ KVCacheBlock, ++ get_request_block_hasher, ++ init_none_hash, ++) ++from vllm.v1.core.sched.scheduler import Scheduler ++from vllm.v1.core.single_type_kv_cache_manager import MambaManager ++from vllm.v1.kv_cache_interface import ( ++ FullAttentionSpec, ++ KVCacheConfig, ++ KVCacheGroupSpec, ++ MambaSpec, ++) ++from vllm.v1.request import Request ++ ++pytestmark = pytest.mark.cpu_test ++ ++BLOCK = 32 ++NUM_SPEC = 3 # MTP K=3, as deployed ++SWITCH = "VLLM_E2_MAMBA_RETAIN_CHECKPOINT" ++# Group 0 = full attention; groups 1, 2 = two GDN-like align-mode Mamba groups ++# (the deployment has several; a hit needs the checkpoint in EVERY one). ++ATTN = 0 ++MAMBA_GROUPS = (1, 2) ++ ++ ++@pytest.fixture(autouse=True) ++def _init_hash(): ++ init_none_hash(sha256) ++ ++ ++@pytest.fixture(params=["1", "0"], ids=["retain-on", "retain-off-control"]) ++def retain(request, monkeypatch) -> bool: ++ monkeypatch.setenv(SWITCH, request.param) ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ return request.param == "1" ++ ++ ++def _config(num_blocks: int) -> KVCacheConfig: ++ mamba = MambaSpec( ++ block_size=BLOCK, ++ shapes=((1, 1),), ++ dtypes=(torch.float32,), ++ mamba_cache_mode="align", ++ num_speculative_blocks=NUM_SPEC, ++ ) ++ return KVCacheConfig( ++ num_blocks=num_blocks, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ KVCacheGroupSpec( ++ ["attn"], ++ FullAttentionSpec( ++ block_size=BLOCK, num_kv_heads=1, head_size=1, dtype=torch.float16 ++ ), ++ ), ++ KVCacheGroupSpec(["gdn_a"], mamba), ++ KVCacheGroupSpec(["gdn_b"], mamba), ++ ], ++ ) ++ ++ ++def _manager(num_blocks: int, num_prefill_lookahead: int = 0) -> KVCacheManager: ++ # retention interval 0 = the deployed default (cache.py): keep only the ++ # replay-boundary (and shared-prefix) states. ++ config = replace(_config(num_blocks), prefix_cache_retention_interval=0) ++ return KVCacheManager( ++ config, ++ max_model_len=8192, ++ scheduler_block_size=BLOCK, ++ hash_block_size=BLOCK, ++ enable_caching=True, ++ use_eagle=True, ++ num_prefill_lookahead=num_prefill_lookahead, ++ ) ++ ++ ++def _mamba(manager: KVCacheManager, gid: int) -> MambaManager: ++ mgr = manager.coordinator.single_type_managers[gid] ++ assert isinstance(mgr, MambaManager) ++ return mgr ++ ++ ++def _request(rid: str, tokens: list[int]) -> Request: ++ params = SamplingParams(max_tokens=4096) ++ params.update_from_generation_config({}, eos_token_id=100000) ++ return Request( ++ request_id=rid, ++ prompt_token_ids=tokens, ++ sampling_params=params, ++ pooling_params=None, ++ block_hasher=get_request_block_hasher(BLOCK, sha256), ++ ) ++ ++ ++def _split(request: Request, num_new: int, num_local_hit: int = 0) -> int: ++ """The real `Scheduler._mamba_block_aligned_split` on a stub self.""" ++ stub = SimpleNamespace( ++ block_size=BLOCK, ++ cache_config=SimpleNamespace(block_size=BLOCK), ++ use_eagle_block_drop=True, ++ max_num_scheduled_tokens=8192, ++ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), ++ mamba_partial_cache_hit=False, ++ mamba_fine_grained_prefix_cache=False, ++ hash_block_size=BLOCK, ++ mamba_has_prefill_checkpoint_blocks=False, ++ mamba_prefill_checkpoint_alignment=None, ++ ) ++ return Scheduler._mamba_block_aligned_split(stub, request, num_new, num_local_hit) ++ ++ ++def _prefill(manager: KVCacheManager, request: Request, budget: int = 8192) -> int: ++ """Prefill to the end of the prompt in scheduler-sized chunks; returns the ++ local prefix-hit length the request was admitted with.""" ++ computed_blocks, num_hit, _ = manager.get_computed_blocks(request) ++ first = True ++ while request.num_computed_tokens < request.num_tokens: ++ start = request.num_computed_tokens + (num_hit if first else 0) ++ remaining = request.num_tokens - start ++ num_new = _split( ++ request, min(remaining, budget), num_hit if first else 0 ++ ) ++ assert num_new > 0 ++ blocks = manager.allocate_slots( ++ request, ++ num_new, ++ num_new_computed_tokens=num_hit if first else 0, ++ new_computed_blocks=computed_blocks if first else None, ++ num_lookahead_tokens=NUM_SPEC, ++ ) ++ assert blocks is not None, "prefill chunk did not fit" ++ request.num_computed_tokens = start + num_new ++ manager.new_step_starts() ++ first = False ++ return num_hit ++ ++ ++def _decode( ++ manager: KVCacheManager, ++ request: Request, ++ steps: int, ++ accepted: Callable[[int], int] = lambda step: 0, ++ on_step: Callable[[], None] | None = None, ++) -> None: ++ """MTP K=3 decode: each step schedules the last token + K drafts and ++ reserves K lookahead slots; ``accepted(step)`` drafts are accepted.""" ++ next_tok = 50000 + len(request.output_token_ids) ++ if request.num_computed_tokens == request.num_tokens: ++ request.append_output_token_ids(next_tok) # sampled by the prefill ++ next_tok += 1 ++ for step in range(steps): ++ assert request.num_computed_tokens == request.num_tokens - 1 ++ request.spec_token_ids = [next_tok + i for i in range(NUM_SPEC)] ++ blocks = manager.allocate_slots( ++ request, 1 + NUM_SPEC, num_lookahead_tokens=NUM_SPEC ++ ) ++ assert blocks is not None, "decode step did not fit" ++ n_acc = accepted(step) ++ request.spec_token_ids = [] ++ request.num_computed_tokens += 1 + n_acc ++ request.append_output_token_ids(list(range(next_tok, next_tok + n_acc + 1))) ++ next_tok += n_acc + 1 ++ manager.new_step_starts() ++ if on_step is not None: ++ on_step() ++ ++ ++def _tokens(seed: int, n: int) -> list[int]: ++ return [seed * 10000 + i for i in range(n)] ++ ++ ++def _hashed_cached(manager: KVCacheManager, block: KVCacheBlock) -> bool: ++ return block.block_hash is not None and block.ref_cnt == 0 ++ ++ ++def _free_queue_ids(manager: KVCacheManager) -> list[int]: ++ return [b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()] ++ ++ ++def _replay_checkpoint(manager: KVCacheManager, request: Request, gid: int): ++ """The group's state block hashed at the request's replay boundary.""" ++ [boundary] = set(manager.coordinator.get_replay_boundaries(request)) ++ block_hash = request.block_hashes[boundary // BLOCK - 1] ++ cached = manager.block_pool.get_cached_block(block_hash, [gid]) ++ assert cached is not None, "the replay checkpoint was never registered" ++ return boundary, cached[0] ++ ++ ++# -------------------------------------------------------------------------- ++# 1. The multi-turn eviction scenario (Hugging Face discussion #2). ++# -------------------------------------------------------------------------- ++def test_concurrent_prefill_cannot_evict_running_requests_checkpoint(retain): ++ """Conversation A decodes; conversation D finishes meanwhile; a cold ++ prefill B then needs a few more blocks than the uncached free ones. LRU ++ order decides the victim: with the old release-on-retire, A's checkpoint ++ (released early in A's decode) is the OLDEST cached free block, so B takes ++ it and A's next turn gets no reuse at all although A's whole attention ++ prefix survives. With retention, B takes D's least-recent blocks instead ++ and A's next turn resumes at its replay boundary. ++ """ ++ manager = _manager(num_blocks=64) ++ pool = manager.block_pool ++ ++ # A: 150-token prompt -> replay boundary 96 (eagle drop): state block 2. ++ a = _request("A", _tokens(1, 150)) ++ assert _prefill(manager, a) == 0 ++ _decode(manager, a, steps=4) # first decode step retires state@96 ++ ckpts = {gid: _replay_checkpoint(manager, a, gid)[1] for gid in MAMBA_GROUPS} ++ for gid, block in ckpts.items(): ++ assert block.block_hash is not None ++ if retain: ++ assert block.ref_cnt == 1, "checkpoint must stay referenced" ++ assert block.block_id not in _free_queue_ids(manager) ++ else: ++ # The exposure the fix removes: cached-free while A still runs. ++ assert _hashed_cached(manager, block) ++ ++ # D: a short conversation that runs and finishes while A decodes. ++ d = _request("D", _tokens(2, 200)) ++ _prefill(manager, d) ++ _decode(manager, d, steps=2) ++ manager.free(d) ++ ++ uncached_free = sum( ++ blk.block_hash is None ++ for blk in pool.free_block_queue.get_all_free_blocks() ++ ) ++ cached_free = pool.get_num_free_blocks() - uncached_free ++ # Precondition: the two worlds differ only by A's checkpoints being in the ++ # free queue (off) or held (on); the uncached supply is identical. ++ assert cached_free >= 4 ++ ++ # B: a cold prefill that pops every uncached free block plus two cached. ++ taken = pool.get_new_blocks(uncached_free + 2) ++ ++ a_ckpt_lost = any(blk.block_hash is None for blk in ckpts.values()) ++ assert a_ckpt_lost is (not retain) ++ pool.free_blocks(taken) ++ ++ # A finishes; its next turn (prompt + output + new user text) arrives. ++ manager.free(a) ++ a2 = _request("A2", list(a.all_token_ids) + _tokens(3, 40)) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ if retain: ++ assert num_hit == 96 ++ else: ++ # FAILING CONTROL: every attention block of A is still cached, yet ++ # the hybrid hit is zero because one Mamba group lost its state. ++ assert all( ++ pool.get_cached_block(a.block_hashes[i], [ATTN]) for i in range(3) ++ ) ++ assert num_hit == 0 ++ ++ ++def test_real_cold_prefill_evicts_checkpoint_only_without_retention(retain): ++ """Same scenario with B as a real request prefilled through allocate_slots ++ (attention + both Mamba groups + MTP scratch), not a raw pool pop.""" ++ manager = _manager(num_blocks=48) ++ a = _request("A", _tokens(1, 150)) ++ _prefill(manager, a) ++ _decode(manager, a, steps=4) ++ ++ d = _request("D", _tokens(2, 200)) ++ _prefill(manager, d) ++ _decode(manager, d, steps=2) ++ manager.free(d) ++ ++ pool = manager.block_pool ++ cached_free = [ ++ b for b in pool.free_block_queue.get_all_free_blocks() if b.block_hash ++ ] ++ # B (448 tokens, sized on this 48-block pool) needs more than the ++ # uncached supply but fits in what is free in both worlds. ++ b = _request("B", _tokens(4, 448)) ++ _prefill(manager, b) ++ _decode(manager, b, steps=1) ++ # Precondition, both worlds: B's prefill did reach the cached blocks. ++ evicted = [blk for blk in cached_free if blk.block_hash is None] ++ assert evicted ++ ++ manager.free(a) ++ a2 = _request("A2", list(a.all_token_ids) + _tokens(3, 40)) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ assert num_hit == (96 if retain else 0) ++ ++ ++# -------------------------------------------------------------------------- ++# 2. Accounting. ++# -------------------------------------------------------------------------- ++def _held_blocks(manager: KVCacheManager) -> set[int]: ++ held: set[int] = set() ++ for mgr in manager.coordinator.single_type_managers: ++ for blocks in mgr.req_to_blocks.values(): ++ held.update(b.block_id for b in blocks if not b.is_null) ++ if isinstance(mgr, MambaManager): ++ held.update(b.block_id for _, b in mgr._retained_checkpoints.values()) ++ return held ++ ++ ++def test_retained_checkpoint_is_counted_by_the_pool_and_admission(retain): ++ manager = _manager(num_blocks=64) ++ pool = manager.block_pool ++ a = _request("A", _tokens(1, 150)) ++ _prefill(manager, a) ++ ++ def check_pool(): ++ # Every block is free, in a table, or retained -- never two of these. ++ held = _held_blocks(manager) ++ free = set(_free_queue_ids(manager)) ++ assert not held & free ++ assert len(held) + len(free) == pool.num_gpu_blocks - 1 # null block ++ for gid in MAMBA_GROUPS: ++ # Referenced hashed states outside A's table: at most the one ++ # retained checkpoint, and none at all without retention. ++ mgr = _mamba(manager, gid) ++ in_table = {b.block_id for b in mgr.req_to_blocks["A"]} ++ outside = [ ++ bid ++ for bid in held - in_table ++ if pool.blocks[bid].block_hash is not None ++ and bid in {b.block_id for _, b in mgr._retained_checkpoints.values()} ++ ] ++ assert len(outside) <= (1 if retain else 0) ++ ++ _decode(manager, a, steps=6, on_step=check_pool) ++ for gid in MAMBA_GROUPS: ++ assert ("A" in _mamba(manager, gid)._retained_checkpoints) is retain ++ ++ # Admission (full-sequence gate): a fresh request reserves one retained ++ # checkpoint per Mamba group; a request already holding one does not. ++ fresh = _request("F", _tokens(5, 150)) ++ empty = manager.empty_kv_cache_blocks.blocks ++ ++ def admission(req: Request, computed: int) -> int: ++ return manager.coordinator.get_num_blocks_to_allocate( ++ request_id=req.request_id, ++ num_tokens=req.num_tokens, ++ new_computed_blocks=empty, ++ num_encoder_tokens=0, ++ total_computed_tokens=computed, ++ num_local_computed_tokens=computed, ++ num_tokens_main_model=req.num_tokens, ++ apply_admission_cap=True, ++ ) ++ ++ def per_step(req: Request) -> int: ++ return manager.coordinator.get_num_blocks_to_allocate( ++ request_id=req.request_id, ++ num_tokens=req.num_tokens, ++ new_computed_blocks=empty, ++ num_encoder_tokens=0, ++ total_computed_tokens=0, ++ num_local_computed_tokens=0, ++ num_tokens_main_model=req.num_tokens, ++ ) ++ ++ extra = len(MAMBA_GROUPS) if retain else 0 ++ assert admission(fresh, 0) == per_step(fresh) + extra ++ # A already retains: only a swap remains, so no extra reservation. ++ assert admission(a, a.num_computed_tokens) == per_step_running(manager, a) ++ ++ ++def per_step_running(manager: KVCacheManager, req: Request) -> int: ++ return manager.coordinator.get_num_blocks_to_allocate( ++ request_id=req.request_id, ++ num_tokens=req.num_tokens, ++ new_computed_blocks=manager.empty_kv_cache_blocks.blocks, ++ num_encoder_tokens=0, ++ total_computed_tokens=req.num_computed_tokens, ++ num_local_computed_tokens=req.num_computed_tokens, ++ num_tokens_main_model=req.num_tokens, ++ apply_admission_cap=False, ++ ) ++ ++ ++def test_admission_refuses_a_request_whose_checkpoints_would_overcommit(retain): ++ """Size the pool so the full-sequence gate passes by exactly the retained ++ blocks' margin: with retention it must refuse (the request would hold ++ those blocks for its lifetime), without it admit.""" ++ req = _request("R", _tokens(6, 150)) ++ probe = _manager(num_blocks=256) ++ need = probe.coordinator.get_num_blocks_to_allocate( ++ request_id="R", ++ num_tokens=req.num_tokens, ++ new_computed_blocks=probe.empty_kv_cache_blocks.blocks, ++ num_encoder_tokens=0, ++ total_computed_tokens=0, ++ num_local_computed_tokens=0, ++ num_tokens_main_model=req.num_tokens, ++ apply_admission_cap=True, ++ ) ++ # The managers' own estimate excludes the retained term, so the pool size ++ # is the same in both worlds. ++ base_need = sum( ++ mgr.get_num_blocks_to_allocate( ++ "R", req.num_tokens, [], 0, 0, req.num_tokens, apply_admission_cap=True ++ ) ++ for mgr in probe.coordinator.single_type_managers ++ ) ++ assert need == base_need + (len(MAMBA_GROUPS) if retain else 0) ++ manager = _manager(num_blocks=base_need + 1) # + null block ++ admitted = manager.allocate_slots(req, BLOCK, full_sequence_must_fit=True) ++ assert (admitted is None) is retain ++ ++ ++# -------------------------------------------------------------------------- ++# 3. Finish path. ++# -------------------------------------------------------------------------- ++@pytest.mark.parametrize("deferred", [False, True], ids=["free", "pop_blocks_for_free"]) ++def test_finish_releases_checkpoint_after_attention_blocks(retain, deferred): ++ manager = _manager(num_blocks=64) ++ pool = manager.block_pool ++ a = _request("A", _tokens(1, 150)) ++ _prefill(manager, a) ++ _decode(manager, a, steps=4) ++ attn_ids = [b.block_id for b in manager.coordinator.single_type_managers[ ++ ATTN].req_to_blocks["A"]] ++ ckpt_ids = { ++ gid: _replay_checkpoint(manager, a, gid)[1].block_id for gid in MAMBA_GROUPS ++ } ++ queue_before = _free_queue_ids(manager) ++ ++ if deferred: ++ # Async-connector path: the scheduler holds the list, frees it later. ++ blocks = manager.pop_blocks_for_free(a) ++ if retain: ++ assert [b.block_id for b in blocks[-2:]] == list(ckpt_ids.values()) ++ pool.free_blocks(blocks) ++ else: ++ manager.free(a) ++ ++ for gid, bid in ckpt_ids.items(): ++ mgr = _mamba(manager, gid) ++ assert "A" not in mgr._retained_checkpoints # nothing leaks ++ block = pool.blocks[bid] ++ assert block.ref_cnt == 0 and block.block_hash is not None # evictable ++ queue = _free_queue_ids(manager) ++ pos = {bid: i for i, bid in enumerate(queue)} ++ newest_attn = max(pos[b] for b in attn_ids if pool.blocks[b].block_hash) ++ if retain: ++ # Freed last: the checkpoints are the request's most recent entries, ++ # so its attention prefix is evicted before them, never after. ++ assert all(pos[bid] > newest_attn for bid in ckpt_ids.values()) ++ assert all(bid not in queue_before for bid in ckpt_ids.values()) ++ else: ++ # Control: they were already queued mid-request, ahead of attention. ++ assert all(bid in queue_before for bid in ckpt_ids.values()) ++ assert all(pos[bid] < newest_attn for bid in ckpt_ids.values()) ++ ++ # Both worlds: with nothing evicted, the next turn hits. ++ a2 = _request("A2", list(a.all_token_ids) + _tokens(3, 40)) ++ assert manager.get_computed_blocks(a2)[1] == 96 ++ ++ ++@pytest.mark.parametrize("drop_release", [False, True], ids=["real", "leak-control"]) ++def test_preempted_request_releases_its_checkpoint(retain, drop_release, monkeypatch): ++ """free() is also the preemption path. Control: if the retained ++ checkpoint were not handed back on free, the pool would lose one block per ++ Mamba group per request -- and the test must see that.""" ++ if drop_release: ++ monkeypatch.setattr(MambaManager, "pop_retained_checkpoints", lambda *a: []) ++ manager = _manager(num_blocks=64) ++ a = _request("A", _tokens(1, 150)) ++ _prefill(manager, a) ++ _decode(manager, a, steps=2) ++ manager.free(a) # the scheduler's preemption path ++ leaked = 63 - manager.block_pool.get_num_free_blocks() # 63 = all but null ++ assert leaked == (len(MAMBA_GROUPS) if retain and drop_release else 0) ++ ++ ++# -------------------------------------------------------------------------- ++# 4. MTP K=3 + chunked prefill on a block-aligned boundary. ++# -------------------------------------------------------------------------- ++@pytest.mark.parametrize("num_prefill_lookahead", [0, NUM_SPEC]) ++def test_mtp_chunked_prefill_aligned_prompt_retains_newest_boundary_only( ++ retain, num_prefill_lookahead ++): ++ """128-token prompt = 4 blocks exactly: under the EAGLE/MTP drop BOTH ++ 64 (identical resend) and 96 (longer sibling) are replay boundaries, so ++ two hashed states retire. Chunks of 40 tokens are cut to block ends by the ++ real split; decode accepts 0..3 drafts per step and crosses several block ++ boundaries, relocating the K=3 scratch blocks. At most one checkpoint is ++ held per group at every step, it is the newest (state@96), and state@64 ++ goes back to the free queue still hashed. Then every free block is popped ++ while the request still runs (a burst of concurrent prefills) and the ++ longer sibling asks for the prefix after the request finishes. ++ """ ++ manager = _manager(num_blocks=96, num_prefill_lookahead=num_prefill_lookahead) ++ pool = manager.block_pool ++ a = _request("A", _tokens(7, 4 * BLOCK)) ++ _prefill(manager, a, budget=40) ++ ++ def at_most_one(): ++ for gid in MAMBA_GROUPS: ++ held = _mamba(manager, gid)._retained_checkpoints ++ assert len(held) <= 1 ++ if held: ++ idx, block = held["A"] ++ assert block.block_hash is not None and block.ref_cnt >= 1 ++ ++ _decode( ++ manager, a, steps=24, accepted=lambda s: s % (NUM_SPEC + 1), on_step=at_most_one ++ ) ++ assert a.num_computed_tokens > 5 * BLOCK # decode crossed two block ends ++ ++ lo_hash, hi_hash = a.block_hashes[1], a.block_hashes[2] # state@64, @96 ++ for gid in MAMBA_GROUPS: ++ [lo] = pool.get_cached_block(lo_hash, [gid]) ++ [hi] = pool.get_cached_block(hi_hash, [gid]) ++ assert _hashed_cached(manager, lo) # superseded / never held: ages ++ if retain: ++ idx, held = _mamba(manager, gid)._retained_checkpoints["A"] ++ assert idx == 2 and held is hi ++ assert hi.ref_cnt == 1 ++ else: ++ assert _hashed_cached(manager, hi) ++ ++ # Burst: concurrent prefills drain the whole free queue while A runs. ++ burst = pool.get_new_blocks(pool.get_num_free_blocks()) ++ manager.free(a) ++ sibling = _request("S", list(a.all_token_ids)[: 4 * BLOCK] + _tokens(8, BLOCK)) ++ _, num_hit, _ = manager.get_computed_blocks(sibling) ++ assert num_hit == (3 * BLOCK if retain else 0) ++ pool.free_blocks(burst) +diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py +index fa70126..5cd93df 100644 +--- a/tests/v1/core/test_prefix_caching.py ++++ b/tests/v1/core/test_prefix_caching.py +@@ -1033,8 +1033,17 @@ def test_hisparse_events_report_host_and_device_placement(): + assert removed.group_idx == 0 + + +-def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks(): +- """Draining unused connector offers must not retain old recurrent states.""" ++@pytest.mark.parametrize("retain_checkpoint", ["1", "0"]) ++def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks( ++ retain_checkpoint: str, monkeypatch: pytest.MonkeyPatch ++): ++ """Draining unused connector offers must not retain old recurrent states. ++ ++ v2.5: with checkpoint retention on (the default) the request keeps ++ exactly ONE retired hashed state referenced -- the newest -- and nothing ++ older; with it off every retired state is released. ++ """ ++ monkeypatch.setenv("VLLM_E2_MAMBA_RETAIN_CHECKPOINT", retain_checkpoint) + manager = make_kv_cache_manager( + KVCacheConfig( + num_blocks=16, +@@ -1067,7 +1076,14 @@ def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks(): + manager.remove_skipped_blocks(request.request_id, end, 64) + manager.new_step_starts() + +- assert all(block.ref_cnt == 0 for block in old_blocks[:-1]) ++ # old_blocks[-1] is the running state; [-2] the newest retired one. ++ if retain_checkpoint == "1": ++ assert all(block.ref_cnt == 0 for block in old_blocks[:-2]) ++ assert old_blocks[-2].ref_cnt == 1 ++ mamba = manager.coordinator.single_type_managers[0] ++ assert mamba._retained_checkpoints[request.request_id][1] is old_blocks[-2] ++ else: ++ assert all(block.ref_cnt == 0 for block in old_blocks[:-1]) + + + def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): +diff --git a/vllm/envs.py b/vllm/envs.py +index 4356854..4159f34 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -444,6 +444,14 @@ if TYPE_CHECKING: + # AOT: no. CAPTURE: no -- strictly host-side. off|warn|raise (1 = warn). + # the mamba slot-ownership check, v1/worker/mamba_utils.py:74. + VLLM_E2_MAMBA_SLOT_OWNERSHIP: str = "off" ++ # v2.5: keep each running request's newest hashed align-mode Mamba ++ # state (its prefix-cache checkpoint) referenced until the request ends, ++ # instead of returning it to the free queue the moment decode passes it ++ # (single_type_kv_cache_manager.MambaManager._retire_state_blocks). Strict ++ # "0"/"1", default ON; "0" restores the old release-on-retire behaviour ++ # for A/B. AOT: no. CAPTURE: no -- scheduler-side block bookkeeping only; ++ # the worker sees the same block tables either way. ++ VLLM_E2_MAMBA_RETAIN_CHECKPOINT: bool = True + # Opt-out for the fork startup guards in vllm/v1/e2_guards.py, read at + # Qwen4ExpForCausalLM construction. AOT: no. CAPTURE: no -- it changes + # nothing itself; it only permits a boot whose identity IS different +@@ -921,15 +929,17 @@ def _resolve_rust_cli_path() -> str | None: + return raw + + +-def _e2_strict_switch(name: str) -> bool: ++def _e2_strict_switch(name: str, default: bool = False) -> bool: + """Parse a fork on/off switch that accepts exactly ``"0"`` and ``"1"``. + +- Unset means off. Anything else -- ``"true"``, ``"yes"``, ``""``, ``" 1"``, ++ Unset means ``default`` (off unless the knob says otherwise). Anything else -- ``"true"``, ``"yes"``, ``""``, ``" 1"``, + ``"01"`` -- is a boot error: a switch that selects a different transport + must not be satisfiable by a spelling nobody meant. + """ + raw = os.environ.get(name) +- if raw is None or raw == "0": ++ if raw is None: ++ return default ++ if raw == "0": + return False + if raw == "1": + return True +@@ -2600,6 +2610,11 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_E2_MAMBA_SLOT_OWNERSHIP": lambda: _e2_diagnostic_mode( + "VLLM_E2_MAMBA_SLOT_OWNERSHIP" + ), ++ # Mamba prefix-checkpoint retention (v2.5; default ON, "0" turns it off). ++ # Strict "0"/"1"; read when each MambaManager is constructed. ++ "VLLM_E2_MAMBA_RETAIN_CHECKPOINT": lambda: _e2_strict_switch( ++ "VLLM_E2_MAMBA_RETAIN_CHECKPOINT", default=True ++ ), + # Seconds between `topk-exact-counters` lines from each GPU worker + # (0, the default, disables the reader entirely). The line reports the + # medium-path exact-restart counters; see +@@ -2863,6 +2878,10 @@ def compile_factors() -> dict[str, object]: + # configuration factor and is hashed conservatively. + "VLLM_E2_PLE_FLAG_CHECK", + "VLLM_E2_MAMBA_SLOT_OWNERSHIP", ++ # Scheduler-side block retention (v2.5). The scheduler is never ++ # compiled and the worker's tables are identical in both worlds, so an ++ # A/B flip must reuse the same AOT cache. ++ "VLLM_E2_MAMBA_RETAIN_CHECKPOINT", + # Logging cadence for an already-compiled kernel's counters: it gates a + # host-side timer and a snapshot op, and is read by no compiled code, so + # hashing it would invalidate every compile cache for a log interval. +diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py +index e3f31a6..ca8c86e 100644 +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -228,6 +228,17 @@ class KVCacheCoordinator(ABC): + num_tokens_main_model, + apply_admission_cap=apply_admission_cap, + ) ++ if apply_admission_cap: ++ # Full-sequence admission: a retained Mamba checkpoint is ++ # held outside the block table for the request's ++ # lifetime, so reserve it here (v2.5). Per-step ++ # allocation needs no term: a held block is simply not in ++ # the free queue that check reads. ++ num_blocks_to_allocate += ( ++ manager.get_num_retained_checkpoint_blocks_to_reserve( ++ request_id, num_tokens_main_model ++ ) ++ ) + return num_blocks_to_allocate + + def allocate_new_computed_blocks( +@@ -386,8 +397,15 @@ class KVCacheCoordinator(ABC): + Args: + request_id: The request ID. + """ ++ # Retained Mamba checkpoints go back last, after every group's table, ++ # so a finished request's checkpoint is never evicted ahead of the ++ # attention prefix it resumes (v2.5). ++ retained: list[KVCacheBlock] = [] + for manager in self.single_type_managers: ++ retained.extend(manager.pop_retained_checkpoints(request_id)) + manager.free(request_id) ++ if retained: ++ self.block_pool.free_blocks(retained) + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + """ +@@ -405,12 +423,16 @@ class KVCacheCoordinator(ABC): + reversed, then concatenated group by group). + """ + blocks: list[KVCacheBlock] = [] ++ retained: list[KVCacheBlock] = [] + for manager in self.single_type_managers: ++ retained.extend(manager.pop_retained_checkpoints(request_id)) + # Reverse within each group so the concatenation is already in + # eviction order, matching `free()`, which frees each group in its + # own reversed `free_blocks()` call. Reversing the flat list + # instead would emit whole groups in the wrong order. + blocks.extend(reversed(manager.pop_blocks_for_free(request_id))) ++ # Retained checkpoints last, as in `free()`. ++ blocks.extend(retained) + return blocks + + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 76ba691..b88330d 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -6,6 +6,7 @@ from collections import defaultdict + from collections.abc import Sequence + from typing import TYPE_CHECKING, ClassVar + ++from vllm import envs + from vllm.distributed.kv_events import MEDIUM_CPU + from vllm.logger import init_logger + from vllm.utils.math_utils import cdiv +@@ -597,6 +598,22 @@ class SingleTypeKVCacheManager(ABC): + # Free blocks in reverse order so that the tail blocks are freed first. + self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) + ++ def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]: ++ """Pop the prefix-cache checkpoints held for ``request_id`` outside its ++ block table, without freeing them. Only align-mode Mamba holds any ++ (see ``MambaManager._retire_state_blocks``). The coordinator frees them ++ after every group's table, so they are the request's last blocks to ++ become eviction candidates.""" ++ return [] ++ ++ def get_num_retained_checkpoint_blocks_to_reserve( ++ self, request_id: str, num_tokens: int ++ ) -> int: ++ """Blocks the admission gate must add for checkpoints this request ++ will hold outside its block table (0 for every manager but align-mode ++ Mamba with retention on).""" ++ return 0 ++ + @abstractmethod + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ +@@ -1528,6 +1545,80 @@ class MambaManager(SingleTypeKVCacheManager): + # connector; a request that finishes first hands off this table + # source directly. + self._producer_partial_tail_reqs: dict[str, tuple[KVCacheBlock, int]] = {} ++ # v2.5. In align mode a state block leaves the block table (it is ++ # nulled and released) as soon as the request computes one block past ++ # it -- `_remove_blocks_in_range` / `remove_skipped_blocks`. When that ++ # block carries a prefix-cache hash it is the request's GDN checkpoint ++ # (the replay boundary under retention interval 0), and releasing it ++ # mid-request puts it in the free queue AHEAD of the request's own ++ # attention blocks, which stay referenced until finish. Concurrent ++ # prefills then pop it first; since the hybrid hit is the shortest ++ # group's, that single block zeroes the next turn's reuse. With ++ # retention on, the newest hashed retired state stays referenced here ++ # (at most one per request per group) and is released only when the ++ # request is freed, after its attention blocks (coordinator.free). ++ self._retain_checkpoint: bool = ( ++ self.mamba_cache_mode == "align" ++ and self.enable_caching ++ and envs.VLLM_E2_MAMBA_RETAIN_CHECKPOINT ++ ) ++ # {request_id: (logical block index, block)} ++ self._retained_checkpoints: dict[str, tuple[int, KVCacheBlock]] = {} ++ if self.mamba_cache_mode == "align" and self.enable_caching: ++ # The A/B evidence line: boot logs must show which arm ran. ++ logger.info_once( ++ "Mamba align prefix-checkpoint retention: %s " ++ "(VLLM_E2_MAMBA_RETAIN_CHECKPOINT)", ++ "on" if self._retain_checkpoint else "off", ++ ) ++ ++ def _retire_state_blocks( ++ self, request_id: str, retired: list[tuple[int, KVCacheBlock]] ++ ) -> None: ++ """Release align-mode state blocks the request no longer reads. ++ ++ ``retired`` holds (logical index, block) pairs already nulled out of ++ the block table, in the eviction order the caller wants. With ++ retention on, the newest hashed one is kept referenced as the request's ++ checkpoint instead, and the checkpoint it supersedes (if any) is ++ released now -- so the request never holds more than one. ++ """ ++ if self._retain_checkpoint: ++ newest: tuple[int, KVCacheBlock] | None = None ++ for idx, block in retired: ++ if block.block_hash is not None and ( ++ newest is None or idx > newest[0] ++ ): ++ newest = (idx, block) ++ held = self._retained_checkpoints.get(request_id) ++ if newest is not None and (held is None or newest[0] > held[0]): ++ self._retained_checkpoints[request_id] = newest ++ retired = [entry for entry in retired if entry[1] is not newest[1]] ++ if held is not None: ++ # Superseded: it ages from here like any cached free block. ++ retired.append(held) ++ if retired: ++ self.block_pool.free_blocks([block for _, block in retired]) ++ ++ def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]: ++ if self.mamba_cache_mode != "align": ++ return [] ++ held = self._retained_checkpoints.pop(request_id, None) ++ return [] if held is None else [held[1]] ++ ++ def get_num_retained_checkpoint_blocks_to_reserve( ++ self, request_id: str, num_tokens: int ++ ) -> int: ++ # A request that already holds a checkpoint only ever swaps it. One ++ # that does not will hold one once a hashed state retires, which needs ++ # at least one state block behind the running one. ++ if ( ++ not self._retain_checkpoint ++ or request_id in self._retained_checkpoints ++ or num_tokens <= self.block_size ++ ): ++ return 0 ++ return 1 + + @classmethod + def find_longest_cache_hit( +@@ -1678,15 +1769,14 @@ class MambaManager(SingleTypeKVCacheManager): + last_block = min(last_block, len(blocks)) + if first_block >= last_block: + return +- freed: list[KVCacheBlock] = [] ++ freed: list[tuple[int, KVCacheBlock]] = [] + # Mamba prefill leaves null gaps between states awaiting retirement. + for i in range(last_block - 1, first_block - 1, -1): + if blocks[i].is_null: + continue +- freed.append(blocks[i]) ++ freed.append((i, blocks[i])) + blocks[i] = self._null_block +- if freed: +- self.block_pool.free_blocks(freed) ++ self._retire_state_blocks(request_id, freed) + self._num_retired_blocks[request_id] = last_block + + def remove_skipped_blocks( +@@ -1716,8 +1806,11 @@ class MambaManager(SingleTypeKVCacheManager): + ): + blocks = self.req_to_blocks[request_id] + if blocks[last_state_block_idx] != self._null_block: +- self.block_pool.free_blocks([blocks[last_state_block_idx]]) ++ block = blocks[last_state_block_idx] + blocks[last_state_block_idx] = self._null_block ++ self._retire_state_blocks( ++ request_id, [(last_state_block_idx, block)] ++ ) + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ +@@ -2022,6 +2115,12 @@ class MambaManager(SingleTypeKVCacheManager): + for entry in self._pending_boundary_state_offloads + if entry[0] != request_id + ] ++ # The coordinator pops the retained checkpoint first and frees it ++ # after every group. A direct caller still gets it back, first in ++ # allocation order, so the caller's reversed free releases it last. ++ retained = self.pop_retained_checkpoints(request_id) ++ if retained: ++ return retained + super().pop_blocks_for_free(request_id) + return super().pop_blocks_for_free(request_id) + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: +-- +2.43.0 + diff --git a/patches/0087-sched-don-t-double-count-the-checkpoint-of-a-prefix-.patch b/patches/0087-sched-don-t-double-count-the-checkpoint-of-a-prefix-.patch new file mode 100644 index 0000000..f0c85d0 --- /dev/null +++ b/patches/0087-sched-don-t-double-count-the-checkpoint-of-a-prefix-.patch @@ -0,0 +1,164 @@ +From f7ed2e656d90d9cf203288995be8a8138e1aefd6 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Fri, 2 Oct 2026 06:23:33 -0400 +Subject: [PATCH 087/103] sched: don't double-count the checkpoint of a + prefix-hit admission + +The admission gate reserved one retained-checkpoint block per Mamba group for every request not holding one. A +follow-up turn with a prefix hit already brings that block, so it was counted twice; the gate now counts it once. +--- + .../core/test_mamba_checkpoint_retention.py | 78 +++++++++++++++++++ + vllm/v1/core/kv_cache_coordinator.py | 2 +- + vllm/v1/core/single_type_kv_cache_manager.py | 23 +++++- + 3 files changed, 100 insertions(+), 3 deletions(-) + +diff --git a/tests/v1/core/test_mamba_checkpoint_retention.py b/tests/v1/core/test_mamba_checkpoint_retention.py +index 1a4ebe5..9d0d440 100644 +--- a/tests/v1/core/test_mamba_checkpoint_retention.py ++++ b/tests/v1/core/test_mamba_checkpoint_retention.py +@@ -561,3 +561,81 @@ def test_mtp_chunked_prefill_aligned_prompt_retains_newest_boundary_only( + _, num_hit, _ = manager.get_computed_blocks(sibling) + assert num_hit == (3 * BLOCK if retain else 0) + pool.free_blocks(burst) ++ ++ ++# -------------------------------------------------------------------------- ++# 5. A follow-up turn's admission must not double-count its checkpoint. ++# -------------------------------------------------------------------------- ++def _followup_after_finish(manager: KVCacheManager): ++ a = _request("A", _tokens(1, 150)) ++ _prefill(manager, a) ++ _decode(manager, a, steps=4) ++ manager.free(a) ++ a2 = _request("A2", list(a.all_token_ids) + _tokens(3, 40)) ++ computed, num_hit, _ = manager.get_computed_blocks(a2) ++ assert num_hit == 96 ++ return a2, computed, num_hit ++ ++ ++def _admission_need(manager, req, computed, num_hit) -> int: ++ return manager.coordinator.get_num_blocks_to_allocate( ++ request_id=req.request_id, ++ num_tokens=req.num_tokens, ++ new_computed_blocks=computed.blocks, ++ num_encoder_tokens=0, ++ total_computed_tokens=num_hit, ++ num_local_computed_tokens=num_hit, ++ num_tokens_main_model=req.num_tokens, ++ apply_admission_cap=True, ++ ) ++ ++ ++def test_hit_admission_reserves_no_extra_checkpoint_block(retain): ++ """The hit state IS the block the follow-up will hold as its checkpoint, ++ and get_num_blocks_to_allocate already counts it, so the full-sequence ++ gate must add nothing for it -- with retention on or off.""" ++ manager = _manager(num_blocks=64) ++ a2, computed, num_hit = _followup_after_finish(manager) ++ assert _admission_need(manager, a2, computed, num_hit) == _managers_need( ++ manager, a2, computed, num_hit ++ ) ++ ++ ++def _managers_need(manager, req, computed, num_hit) -> int: ++ """The per-group estimates alone, without the retained-checkpoint term.""" ++ return sum( ++ mgr.get_num_blocks_to_allocate( ++ req.request_id, req.num_tokens, computed.blocks[gid], num_hit, ++ num_hit, req.num_tokens, apply_admission_cap=True, ++ ) ++ for gid, mgr in enumerate(manager.coordinator.single_type_managers) ++ ) ++ ++ ++def test_hit_followup_is_admitted_when_it_exactly_fits(retain): ++ """Leave exactly the follow-up's real need free (uncached blocks only are ++ taken, so its cache stays intact). It must be admitted with its hit. The ++ previous version reserved one more block per Mamba group here ++ and refused it -- the run that failing control is measured on.""" ++ manager = _manager(num_blocks=64) ++ pool = manager.block_pool ++ a2, computed, num_hit = _followup_after_finish(manager) ++ need = _managers_need(manager, a2, computed, num_hit) ++ uncached = sum( ++ b.block_hash is None for b in pool.free_block_queue.get_all_free_blocks() ++ ) ++ surplus = pool.get_num_free_blocks() - need ++ assert 0 < surplus <= uncached ++ held = pool.get_new_blocks(surplus) ++ assert pool.get_num_free_blocks() == need ++ num_new = _split(a2, a2.num_tokens - num_hit, num_hit) ++ blocks = manager.allocate_slots( ++ a2, ++ num_new, ++ num_new_computed_tokens=num_hit, ++ new_computed_blocks=computed, ++ num_lookahead_tokens=NUM_SPEC, ++ full_sequence_must_fit=True, ++ ) ++ assert blocks is not None ++ pool.free_blocks(held) +diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py +index ca8c86e..4b2839d 100644 +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -236,7 +236,7 @@ class KVCacheCoordinator(ABC): + # the free queue that check reads. + num_blocks_to_allocate += ( + manager.get_num_retained_checkpoint_blocks_to_reserve( +- request_id, num_tokens_main_model ++ request_id, num_tokens_main_model, new_computed_blocks[i] + ) + ) + return num_blocks_to_allocate +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index b88330d..0ec6bca 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -607,7 +607,10 @@ class SingleTypeKVCacheManager(ABC): + return [] + + def get_num_retained_checkpoint_blocks_to_reserve( +- self, request_id: str, num_tokens: int ++ self, ++ request_id: str, ++ num_tokens: int, ++ new_computed_blocks: Sequence[KVCacheBlock] = (), + ) -> int: + """Blocks the admission gate must add for checkpoints this request + will hold outside its block table (0 for every manager but align-mode +@@ -1607,7 +1610,10 @@ class MambaManager(SingleTypeKVCacheManager): + return [] if held is None else [held[1]] + + def get_num_retained_checkpoint_blocks_to_reserve( +- self, request_id: str, num_tokens: int ++ self, ++ request_id: str, ++ num_tokens: int, ++ new_computed_blocks: Sequence[KVCacheBlock] = (), + ) -> int: + # A request that already holds a checkpoint only ever swaps it. One + # that does not will hold one once a hashed state retires, which needs +@@ -1618,6 +1624,19 @@ class MambaManager(SingleTypeKVCacheManager): + or num_tokens <= self.block_size + ): + return 0 ++ # A prefix hit already brings the block that becomes the held ++ # checkpoint -- the hit state, counted by get_num_blocks_to_allocate as ++ # an evictable computed block (or already referenced elsewhere) -- and ++ # the same holds for a hashed state still in the request's table. ++ # Reserving another one double-counts every hit admission by one block ++ # per group, which only delays the follow-up turns the fix is for. ++ if any(not blk.is_null for blk in new_computed_blocks): ++ return 0 ++ if any( ++ not blk.is_null and blk.block_hash is not None ++ for blk in self.req_to_blocks.get(request_id, ()) ++ ): ++ return 0 + return 1 + + @classmethod +-- +2.43.0 + diff --git a/patches/0088-kv-side-cache-layout-fold-the-PLE-state-into-a-GDN-g.patch b/patches/0088-kv-side-cache-layout-fold-the-PLE-state-into-a-GDN-g.patch new file mode 100644 index 0000000..e35d2ae --- /dev/null +++ b/patches/0088-kv-side-cache-layout-fold-the-PLE-state-into-a-GDN-g.patch @@ -0,0 +1,1728 @@ +From e2d6962242963fb82f4b55ae460bd549bd2b4152 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Fri, 2 Oct 2026 09:50:32 -0400 +Subject: [PATCH 088/103] kv: side-cache layout -- fold the PLE state into a + GDN group, pack QSA rings + +The packed pool strides every block by the widest group, and every group draws whole blocks. The PLE short-conv +state now shares a GDN cache group, and the QSA raw-key rings are packed as equal slots inside shared blocks, so +neither holds a whole block per request. VLLM_E2_SIDE_CACHE_LAYOUT (default 1) selects the old layout with 0. +On its own the layout raises the Flash-Next pool from 806,792 to 853,994 tokens at the same budget. +--- + tests/test_envs.py | 23 + + tests/v1/core/test_side_cache_layout.py | 882 +++++++++++++++++++ + tests/v1/test_pp_kv_projection.py | 20 +- + vllm/envs.py | 15 + + vllm/v1/core/kv_cache_coordinator.py | 12 + + vllm/v1/core/kv_cache_utils.py | 335 ++++++- + vllm/v1/core/sched/scheduler.py | 4 + + vllm/v1/core/single_type_kv_cache_manager.py | 112 ++- + vllm/v1/kv_cache_interface.py | 23 +- + vllm/v1/worker/gpu/model_runner.py | 1 + + vllm/v1/worker/gpu_model_runner.py | 1 + + vllm/v1/worker/utils.py | 11 + + 12 files changed, 1426 insertions(+), 13 deletions(-) + create mode 100644 tests/v1/core/test_side_cache_layout.py + +diff --git a/tests/test_envs.py b/tests/test_envs.py +index bcd7b33..af5e39d 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -621,6 +621,10 @@ _FORK_HASHED = [ + # PLE pull transport: selects a different PLE module class and captured + # forward, so it is hashed. Strict "0"/"1" parser (envs._e2_strict_switch). + ("VLLM_E2_PLE_PULL_TRANSPORT", "1", True), ++ # v2.5 side-cache layout: changes the ring views' block stride and which ++ # block table the PLE state is read through, so ON/OFF must not share an ++ # AOT cache. Strict "0"/"1", default ON; "0" selects the old layout. ++ ("VLLM_E2_SIDE_CACHE_LAYOUT", "0", False), + ] + + _FORK_IGNORED = [ +@@ -1065,6 +1069,7 @@ _MUST_BE_READ = { + "VLLM_E2_PLE_FLAG_CHECK": "vllm/v1/ple_offload/connector.py", + "VLLM_E2_MAMBA_SLOT_OWNERSHIP": "vllm/v1/worker/mamba_utils.py", + "VLLM_E2_MAMBA_RETAIN_CHECKPOINT": "vllm/v1/core/single_type_kv_cache_manager.py", ++ "VLLM_E2_SIDE_CACHE_LAYOUT": "vllm/v1/core/kv_cache_utils.py", + } + _INERT_NO_READ_SITE = ["VLLM_PLE_FP8_GLOBAL_SCALE"] + +@@ -1128,3 +1133,21 @@ def test_mamba_retain_checkpoint_switch_defaults_on_and_is_strict( + monkeypatch.setenv(name, bad) + with pytest.raises(ValueError, match=name): + _fresh_env_value(name) ++ ++ ++def test_side_cache_layout_switch_defaults_on_and_is_strict( ++ monkeypatch: pytest.MonkeyPatch, ++): ++ """Unset is ON (the layout fix ships enabled); "0" is the only way back to ++ the old one-group-each layout, and a misspelt value is a boot error.""" ++ name = "VLLM_E2_SIDE_CACHE_LAYOUT" ++ monkeypatch.delenv(name, raising=False) ++ assert _fresh_env_value(name) is True ++ monkeypatch.setenv(name, "1") ++ assert _fresh_env_value(name) is True ++ monkeypatch.setenv(name, "0") ++ assert _fresh_env_value(name) is False ++ for bad in ("off", "false", "", " 1"): ++ monkeypatch.setenv(name, bad) ++ with pytest.raises(ValueError, match=name): ++ _fresh_env_value(name) +diff --git a/tests/v1/core/test_side_cache_layout.py b/tests/v1/core/test_side_cache_layout.py +new file mode 100644 +index 0000000..9ebba77 +--- /dev/null ++++ b/tests/v1/core/test_side_cache_layout.py +@@ -0,0 +1,882 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""v2.5 side-cache layout (``VLLM_E2_SIDE_CACHE_LAYOUT``, default ON). ++ ++Flash-Next's packed block-outermost KV pool strides every block by the widest ++group (12,902,400 B on PP1 at FP32 SSM state) and every group draws whole ++blocks from it. The TP-replicated PLE short-conv state (one layer, 2 + K blocks ++per request in align mode) and the QSA raw-key ring (one block per request) ++hold ~1.2 MB of real state but took 6 of the 103 blocks a 262K request needs. ++The layout fix folds the PLE group into a GDN group that has a spare page on ++every PP stage, and packs several requests' rings into one pool block. ++ ++What these tests pin (CPU only): ++ ++1. Pool size. On the reference host's per-layer specs (TP2 x PP2, partition 25/23, MTP ++ K=3, fp8 KV, BLNHC, pin 4.1e9) the block count and both stages' strides are ++ unchanged and a 262K request needs 97 blocks instead of 103 at FP32 state ++ (173 instead of 179 at FP16). FAILING CONTROL: switch "0" reproduces the ++ logged pools (806,792 and 894,804 tokens) block for block. ++2. Allocation / eviction through ``KVCacheManager`` as the scheduler drives it ++ (aligned prefill chunks, MTP K=3 decode with draft acceptance, finish, ++ preemption, the deferred ``pop_blocks_for_free`` path, prefix hits with the ++ retained checkpoints). The PLE state's block table evolves exactly as ++ it did in its own group, hits are identical, and N rings share one pool ++ block. FAILING CONTROLS: with the switch off the same pool cannot admit the ++ extra request, and every ring takes its own block. ++3. Worker views. The ring slots and the folded PLE page lie inside their pool ++ block and overlap nothing they must not, and the KV zeroer zeroes exactly ++ the same bytes of every block in both layouts. FAILING CONTROL: a zeroer ++ built without the slot-group skip zeroes different bytes. ++""" ++ ++from dataclasses import replace ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm import envs ++from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator ++from vllm.sampling_params import SamplingParams ++from vllm.utils.hashing import sha256 ++from vllm.utils.math_utils import cdiv ++from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum ++from vllm.v1.core import kv_cache_utils as ku ++from vllm.v1.core.block_pool import BlockPool ++from vllm.v1.core.kv_cache_manager import KVCacheManager ++from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash ++from vllm.v1.core.sched.scheduler import Scheduler ++from vllm.v1.core.single_type_kv_cache_manager import ( ++ CircularBufferManager, ++ RingSlotPool, ++) ++from vllm.v1.kv_cache_interface import ( ++ CircularBufferSpec, ++ FullAttentionSpec, ++ KVCacheLayout, ++ MambaSpec, ++ MLAAttentionSpec, ++ UniformTypeKVCacheSpecs, ++ get_kv_quant_mode, ++) ++from vllm.v1.request import Request ++ ++pytestmark = pytest.mark.cpu_test ++ ++SWITCH = "VLLM_E2_SIDE_CACHE_LAYOUT" ++TP = 2 ++K = 3 # MTP num_speculative_tokens, as deployed ++PIN = 4_100_000_000 # --kv-cache-memory on every rank ++MAX_LEN = 262_144 ++ ++ ++def _set_switch(monkeypatch, value: str) -> None: ++ monkeypatch.setenv(SWITCH, value) ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ ++ ++@pytest.fixture(params=["1", "0"], ids=["layout-on", "layout-off-control"]) ++def layout(request, monkeypatch) -> bool: ++ _set_switch(monkeypatch, request.param) ++ return request.param == "1" ++ ++ ++@pytest.fixture(autouse=True) ++def _init_hash(): ++ init_none_hash(sha256) ++ ++ ++# --------------------------------------------------------------------------- ++# Flash-Next per-layer specs (the reference fixture, in align mode, any block size ++# and SSM dtype). Boot log: "Setting attention block size to 3200 tokens" at FP32 ++# state, 1664 at FP16; every Mamba page padded to the attention page. ++# --------------------------------------------------------------------------- ++def _flash_next_specs( ++ block: int = 3200, ++ ssm_dtype: torch.dtype = torch.float32, ++ *, ++ tiny_state: bool = False, ++) -> tuple[dict, dict]: ++ attn_page = block * 1 * 256 * 2 # fp8 K+V, 1 KV head per TP rank ++ main = FullAttentionSpec( ++ block_size=block, ++ num_kv_heads=2 // TP, ++ head_size=256, ++ head_size_v=256, ++ dtype=torch.uint8, ++ kv_quant_mode=get_kv_quant_mode("fp8_e4m3"), ++ ) ++ ring = CircularBufferSpec( ++ block_size=8, num_kv_heads=1, head_size=144, head_size_v=0, dtype=torch.bfloat16 ++ ) ++ comp = MLAAttentionSpec( ++ block_size=block, ++ num_kv_heads=1, ++ head_size=128, ++ dtype=torch.bfloat16, ++ tokens_per_state=4, ++ ) ++ if tiny_state: ++ # Small blocks for scheduler tests: the real states exceed the page. ++ gdn_shapes: tuple = ((4,),) ++ gdn_dtypes: tuple = (torch.float32,) ++ ple_shapes: tuple = ((2,),) ++ else: ++ gdn_shapes = tuple( ++ MambaStateShapeCalculator.gated_delta_net_state_shape( ++ TP, 16, 48, 128, 128, 4, K ++ ) ++ ) ++ gdn_dtypes = (torch.bfloat16, ssm_dtype) ++ ple_shapes = tuple( ++ MambaStateShapeCalculator.short_conv_state_shape( ++ tp_world_size=1, intermediate_size=10240, conv_kernel=10, num_spec=K ++ ) ++ ) ++ gdn = MambaSpec( ++ shapes=gdn_shapes, ++ dtypes=gdn_dtypes, ++ block_size=block, ++ page_size_padded=attn_page, ++ mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ++ mamba_cache_mode="align", ++ num_speculative_blocks=K, ++ ) ++ ple = MambaSpec( ++ shapes=ple_shapes, ++ dtypes=(torch.bfloat16,), ++ block_size=block, ++ page_size_padded=attn_page, ++ mamba_type=MambaAttentionBackendEnum.SHORT_CONV, ++ mamba_cache_mode="align", ++ tp_replicated=True, ++ num_speculative_blocks=K, ++ ) ++ ++ def qsa(prefix: str, specs: dict) -> None: ++ specs[f"{prefix}.self_attn.indexer.raw_key_cache"] = ring ++ specs[f"{prefix}.self_attn.indexer.compressed_key_cache"] = comp ++ specs[f"{prefix}.self_attn.attn"] = main ++ ++ def stage(lo: int, hi: int, mtp: bool) -> dict: ++ specs: dict = {} ++ for i in range(lo, hi): ++ prefix = f"language_model.model.layers.{i}" ++ if i == 2: ++ specs[f"{prefix}.ple"] = ple ++ if i % 4 == 3: ++ qsa(prefix, specs) ++ else: ++ specs[f"{prefix}.linear_attn"] = gdn ++ if mtp: ++ qsa("mtp.layers.48", specs) ++ return specs ++ ++ return stage(0, 25, False), stage(25, 48, True) ++ ++ ++def _vllm_config( ++ *, ++ override: int | None = None, ++ max_num_seqs: int = 8, ++ kv_transfer_config=None, ++ max_model_len: int = MAX_LEN, ++): ++ cache = SimpleNamespace( ++ kv_cache_layout="BLNHC", ++ num_gpu_blocks_override=override, ++ prefix_cache_retention_interval=0, ++ mamba_cache_mode="align", ++ get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLNHC, ++ ) ++ return SimpleNamespace( ++ cache_config=cache, ++ speculative_config=SimpleNamespace( ++ num_speculative_tokens=K, ++ method="mtp", ++ use_multi_module_mtp=lambda: False, ++ use_eagle=lambda: True, ++ use_eagle_block_drop=lambda: True, ++ ), ++ scheduler_config=SimpleNamespace( ++ disable_hybrid_kv_cache_manager=False, ++ max_num_encoder_input_tokens=0, ++ max_num_seqs=max_num_seqs, ++ ), ++ attention_config=SimpleNamespace(hisparse_config=None), ++ model_config=SimpleNamespace( ++ original_max_model_len=max_model_len, ++ max_model_len=max_model_len, ++ hf_config=SimpleNamespace(model_type="qwen4_exp"), ++ get_total_num_hidden_layers=lambda: 48, ++ ), ++ device_config=SimpleNamespace(device_type="cuda"), ++ parallel_config=SimpleNamespace( ++ pipeline_parallel_size=2, decode_context_parallel_size=1 ++ ), ++ kv_transfer_config=kv_transfer_config, ++ ) ++ ++ ++def _configs(block=3200, ssm_dtype=torch.float32, *, pin=PIN, **kw): ++ pp0, pp1 = _flash_next_specs(block, ssm_dtype, tiny_state=kw.pop("tiny", False)) ++ cfg = _vllm_config(**kw) ++ specs = [pp0, pp0, pp1, pp1] # TP2 x PP2 ++ return cfg, specs, ku.get_kv_cache_configs(cfg, specs, [pin] * 4) ++ ++ ++def _first_spec(group): ++ spec = group.kv_cache_spec ++ return spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ++ ++ ++def _request_blocks(cfg, config) -> int: ++ """Blocks one max_model_len request holds, as the boot's capacity line ++ counts them (shared ring pages excluded).""" ++ total = 0 ++ for group in config.kv_cache_groups: ++ if group.slots_per_block > 1: ++ continue ++ total += cdiv( ++ group.kv_cache_spec.max_memory_usage_bytes(cfg), ++ group.kv_cache_spec.page_size_bytes, ++ ) ++ return total ++ ++ ++# --------------------------------------------------------------------------- ++# 1. Pool size (blocks freed per request), FP32 and FP16 state. ++# --------------------------------------------------------------------------- ++# (block, ssm dtype, stage strides, blocks, per-request OFF, ON, pool OFF, ON) ++_POOL_CASES = { ++ "fp32": ( ++ 3200, ++ torch.float32, ++ (11_468_800, 12_902_400), ++ 317, ++ 103, ++ 97, ++ 806_792, ++ 853_994, ++ ), ++ "fp16": ( ++ 1664, ++ torch.float16, ++ (5_963_776, 6_709_248), ++ 611, ++ 179, ++ 173, ++ 894_804, ++ 924_322, ++ ), ++} ++ ++ ++@pytest.mark.parametrize("case", sorted(_POOL_CASES)) ++def test_pool_blocks_per_request(case, layout): ++ block, dtype, strides, blocks, req_off, req_on, pool_off, pool_on = _POOL_CASES[ ++ case ++ ] ++ cfg, specs, configs = _configs(block, dtype) ++ for rank, config in enumerate(configs): ++ # The pool itself does not move: same block count, same stride on ++ # both stages (PP1 binds), so not one byte of pin is spent. ++ assert config.num_blocks == blocks, f"rank {rank}" ++ assert ( ++ ku._get_kv_cache_bytes_per_block(config.kv_cache_groups) ++ == strides[rank // 2] ++ ) ++ per_request = _request_blocks(cfg, configs[0]) ++ capacity, _ = ku.get_kv_cache_capacity(cfg, configs[0]) ++ if layout: ++ assert per_request == req_on ++ assert req_off - per_request == 6 # 5 PLE + 1 ring blocks freed ++ assert capacity == pool_on ++ else: ++ # FAILING CONTROL: the old layout, exactly as the reference host logged it ++ # ("GPU KV cache size: 806,792 tokens" / "894,804 tokens"). ++ assert per_request == req_off ++ assert capacity == pool_off ++ # Every rank agrees, and so does the scheduler's config, which is what ++ # update_kv_cache_capacity logs at boot ("GPU KV cache size: N tokens"). ++ assert {ku.get_kv_cache_capacity(cfg, c)[0] for c in configs} == {capacity} ++ sched = ku.generate_scheduler_kv_cache_config(configs) ++ assert ku.get_kv_cache_capacity(cfg, sched)[0] == capacity ++ ++ ++def test_ple_rides_in_a_gdn_group_with_a_spare_page(layout): ++ _, specs, configs = _configs() ++ ple = "language_model.model.layers.2.ple" ++ for rank, config in enumerate(configs): ++ tensor_layers = sorted(n for t in config.kv_cache_tensors for n in t.layers) ++ assert tensor_layers == sorted(specs[rank]), f"rank {rank}" ++ owner = [g for g in config.kv_cache_groups if ple in g.layer_names] ++ mamba_groups = [ ++ g for g in config.kv_cache_groups if isinstance(_first_spec(g), MambaSpec) ++ ] ++ if rank >= 2: ++ assert not owner # PLE lives on PP0 only ++ if layout: ++ assert len(mamba_groups) == 3 ++ assert all(g.layer_names for g in config.kv_cache_groups) ++ if rank < 2: ++ [group] = owner ++ gdn = [n for n in group.layer_names if n.endswith("linear_attn")] ++ assert len(gdn) == 6 and group.layer_names[-1] == ple ++ else: ++ # FAILING CONTROL: its own group, empty on PP1 (the reference shape). ++ assert len(mamba_groups) == 4 ++ if rank < 2: ++ assert owner[0].layer_names == [ple] ++ ++ ++def test_ring_slots_divide_both_stage_strides(): ++ _, _, configs = _configs() ++ rings = [ ++ (rank, g) ++ for rank, c in enumerate(configs) ++ for g in c.kv_cache_groups ++ if isinstance(_first_spec(g), CircularBufferSpec) ++ ] ++ assert {g.slots_per_block for _, g in rings} == {800} ++ for rank, config in enumerate(configs): ++ stride = ku._get_kv_cache_bytes_per_block(config.kv_cache_groups) ++ for tensor in config.kv_cache_tensors: ++ if tensor.slots_per_block > 1: ++ assert tensor.block_stride * tensor.slots_per_block == stride ++ assert tensor.block_stride % 256 == 0 ++ assert tensor.block_stride >= len(tensor.layers) * 2304 ++ assert config.num_blocks_of(tensor) == config.num_blocks * 800 ++ ++ ++def test_kv_connector_keeps_the_old_layout(monkeypatch): ++ """NIXL / offload / P/D address groups themselves: layout stays off.""" ++ _set_switch(monkeypatch, "1") ++ cfg, _, configs = _configs(kv_transfer_config=SimpleNamespace()) ++ assert _request_blocks(cfg, configs[0]) == 103 ++ assert all(g.slots_per_block == 1 for g in configs[0].kv_cache_groups) ++ ++ ++def test_fold_never_widens_a_stage(monkeypatch): ++ """The PP-aware target choice. Folding into GDN group 0 (7 pages on PP0) ++ would widen PP0's block to 8 pages; the pass must pick a 6-page group. ++ CONTROL: the forced group-0 fold really does widen PP0.""" ++ _set_switch(monkeypatch, "0") ++ pp0, pp1 = _flash_next_specs() ++ specs = [pp0, pp0, pp1, pp1] ++ merged: dict = {} ++ for s in specs: ++ merged.update(s) ++ groups = ku.get_kv_cache_groups(_vllm_config(), merged) ++ base = ku._worker_strides(groups, specs) ++ folded, notes = ku._fold_replicated_mamba_groups(groups, specs) ++ assert ku._worker_strides(folded, specs) == base ++ assert len(folded) == len(groups) - 1 and "into Mamba group 1" in notes[0] ++ ++ ple_index = next( ++ i ++ for i, g in enumerate(groups) ++ if any(n.endswith(".ple") for n in g.layer_names) ++ ) ++ g0 = groups[0] ++ forced = UniformTypeKVCacheSpecs( ++ block_size=g0.kv_cache_spec.block_size, ++ kv_cache_specs={ ++ **g0.kv_cache_spec.kv_cache_specs, ++ **groups[ple_index].kv_cache_spec.kv_cache_specs, ++ }, ++ ) ++ wrong = [ ++ replace( ++ g0, ++ layer_names=g0.layer_names + groups[ple_index].layer_names, ++ kv_cache_spec=forced, ++ ) ++ ] ++ wrong += [g for i, g in enumerate(groups) if i not in (0, ple_index)] ++ widened = ku._worker_strides(wrong, specs) ++ assert widened[0] == base[0] + 1_638_400 # PP0: 8 pages instead of 7 ++ ++ ++# --------------------------------------------------------------------------- ++# 2. Allocation / eviction through KVCacheManager (tiny states, block 32). ++# --------------------------------------------------------------------------- ++BLOCK = 32 ++PLE = "language_model.model.layers.2.ple" ++ ++ ++SCHED_MAX_LEN = 512 # small, so small test pools pass the boot memory check ++ ++ ++def _manager(num_blocks: int) -> KVCacheManager: ++ _, _, configs = _configs( ++ BLOCK, tiny=True, override=num_blocks, max_model_len=SCHED_MAX_LEN ++ ) ++ config = ku.generate_scheduler_kv_cache_config(configs) ++ return KVCacheManager( ++ replace(config, prefix_cache_retention_interval=0), ++ max_model_len=SCHED_MAX_LEN, ++ scheduler_block_size=BLOCK, ++ hash_block_size=BLOCK, ++ enable_caching=True, ++ use_eagle=True, ++ ) ++ ++ ++def _group_of(manager: KVCacheManager, layer: str) -> int: ++ return next( ++ i ++ for i, g in enumerate(manager.kv_cache_config.kv_cache_groups) ++ if layer in g.layer_names ++ ) ++ ++ ++def _ring_gid(manager: KVCacheManager) -> int: ++ return next( ++ i ++ for i, m in enumerate(manager.coordinator.single_type_managers) ++ if type(m) is CircularBufferManager ++ ) ++ ++ ++def _ring(manager: KVCacheManager) -> CircularBufferManager: ++ mgr = manager.coordinator.single_type_managers[_ring_gid(manager)] ++ assert isinstance(mgr, CircularBufferManager) ++ return mgr ++ ++ ++def _request(rid: str, tokens: list[int]) -> Request: ++ params = SamplingParams(max_tokens=4096) ++ params.update_from_generation_config({}, eos_token_id=100000) ++ return Request( ++ request_id=rid, ++ prompt_token_ids=tokens, ++ sampling_params=params, ++ pooling_params=None, ++ block_hasher=get_request_block_hasher(BLOCK, sha256), ++ ) ++ ++ ++def _tokens(seed: int, n: int) -> list[int]: ++ return [seed * 10000 + i for i in range(n)] ++ ++ ++def _split(request: Request, num_new: int, num_local_hit: int = 0) -> int: ++ stub = SimpleNamespace( ++ block_size=BLOCK, ++ cache_config=SimpleNamespace(block_size=BLOCK), ++ use_eagle_block_drop=True, ++ max_num_scheduled_tokens=8192, ++ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), ++ mamba_partial_cache_hit=False, ++ mamba_fine_grained_prefix_cache=False, ++ hash_block_size=BLOCK, ++ mamba_has_prefill_checkpoint_blocks=False, ++ mamba_prefill_checkpoint_alignment=None, ++ ) ++ return Scheduler._mamba_block_aligned_split(stub, request, num_new, num_local_hit) ++ ++ ++def _prefill(manager, request, budget=40, trace=None, may_stall=False) -> int | None: ++ """Prefill in scheduler-sized chunks; None if admission is refused (or, ++ with ``may_stall``, if a later chunk does not fit -- the request is then ++ freed, as a scheduler would preempt it).""" ++ computed, num_hit, _ = manager.get_computed_blocks(request) ++ first = True ++ while request.num_computed_tokens < request.num_tokens: ++ start = request.num_computed_tokens + (num_hit if first else 0) ++ num_new = _split( ++ request, min(request.num_tokens - start, budget), num_hit if first else 0 ++ ) ++ blocks = manager.allocate_slots( ++ request, ++ num_new, ++ num_new_computed_tokens=num_hit if first else 0, ++ new_computed_blocks=computed if first else None, ++ num_lookahead_tokens=K, ++ full_sequence_must_fit=first, ++ ) ++ if blocks is None: ++ if not first: ++ assert may_stall, "a later prefill chunk did not fit" ++ manager.free(request) ++ return None ++ request.num_computed_tokens = start + num_new ++ manager.new_step_starts() ++ first = False ++ if trace is not None: ++ trace() ++ return num_hit ++ ++ ++def _decode(manager, request, steps, accepted=lambda step: 0, trace=None) -> None: ++ next_tok = 50000 + len(request.output_token_ids) ++ if request.num_computed_tokens == request.num_tokens: ++ request.append_output_token_ids(next_tok) ++ next_tok += 1 ++ for step in range(steps): ++ request.spec_token_ids = [next_tok + i for i in range(K)] ++ assert manager.allocate_slots(request, 1 + K, num_lookahead_tokens=K) ++ n_acc = accepted(step) ++ request.spec_token_ids = [] ++ request.num_computed_tokens += 1 + n_acc ++ request.append_output_token_ids(list(range(next_tok, next_tok + n_acc + 1))) ++ next_tok += n_acc + 1 ++ manager.new_step_starts() ++ if trace is not None: ++ trace() ++ ++ ++def _pattern(blocks) -> tuple: ++ """A block table with ids renamed by first appearance (null = -1), so two ++ layouts can be compared position by position.""" ++ names: dict[int, int] = {} ++ return tuple( ++ -1 if b.is_null else names.setdefault(b.block_id, len(names)) for b in blocks ++ ) ++ ++ ++def _trajectory(manager: KVCacheManager, accept_seed: int = 7) -> list: ++ """Run A (prefill + MTP decode across block boundaries, varying accepts), ++ finish it, resume it as A2 (prefix hit), run cold B beside it, free B via ++ the deferred path, resume it as B2, and record, after every step, the ++ block-table pattern of the attention group, of the PLE layer's group and ++ of every GDN group, the ring table width, and every hit length.""" ++ groups = manager.kv_cache_config.kv_cache_groups ++ ple_gid = _group_of(manager, PLE) ++ gdn_gids = [ ++ i ++ for i, g in enumerate(groups) ++ if any(n.endswith("linear_attn") for n in g.layer_names) ++ ] ++ attn_gid = _group_of(manager, "language_model.model.layers.3.self_attn.attn") ++ ring_gid = _ring_gid(manager) ++ log: list = [] ++ live: list[Request] = [] ++ ++ def trace(): ++ for req in live: ++ blocks = manager.coordinator.get_blocks(req.request_id) ++ if not blocks[attn_gid]: ++ continue ++ log.append( ++ ( ++ req.request_id, ++ _pattern(blocks[attn_gid]), ++ _pattern(blocks[ple_gid]), ++ tuple(sorted(_pattern(blocks[g]) for g in gdn_gids)), ++ len(blocks[ring_gid]), ++ ) ++ ) ++ ++ a = _request("A", _tokens(1, 150)) ++ live.append(a) ++ log.append(("hitA", _prefill(manager, a, trace=trace))) ++ _decode(manager, a, 30, accepted=lambda s: (s * accept_seed) % 4, trace=trace) ++ manager.free(a) ++ live.remove(a) ++ a2 = _request("A2", list(a.all_token_ids) + _tokens(3, 40)) ++ live.append(a2) ++ log.append(("hitA2", _prefill(manager, a2, trace=trace))) ++ b = _request("B", _tokens(4, 120)) ++ live.append(b) ++ log.append(("hitB", _prefill(manager, b, trace=trace))) ++ _decode(manager, a2, 6, accepted=lambda s: s % 4, trace=trace) ++ _decode(manager, b, 3, trace=trace) ++ popped = manager.pop_blocks_for_free(b) # deferred free (the block-free fence) ++ live.remove(b) ++ manager.block_pool.free_blocks(popped) # the scheduler's fenced drain ++ b2 = _request("B2", list(b.all_token_ids) + _tokens(5, 10)) ++ live.append(b2) ++ log.append(("hitB2", _prefill(manager, b2, trace=trace))) ++ for req in (a2, b2): ++ manager.free(req) ++ return log ++ ++ ++def test_ple_block_table_and_hits_match_the_old_layout(monkeypatch): ++ """The PLE layer's block table (null/live positions, which entries share a ++ block across steps) and every hybrid hit length are the same as when it ++ had its own group: its kernels index the same state positions, so its ++ outputs cannot change. A large pool keeps eviction out of the comparison. ++ The retained checkpoints are on (the default) in both runs.""" ++ runs = {} ++ for value in ("0", "1"): ++ _set_switch(monkeypatch, value) ++ manager = _manager(num_blocks=400) ++ assert envs.VLLM_E2_MAMBA_RETAIN_CHECKPOINT ++ runs[value] = _trajectory(manager) ++ # Nothing leaked: only the null block is not free. ++ assert manager.block_pool.get_num_free_blocks() == 399 ++ assert len(runs["1"]) > 50 ++ assert runs["1"] == runs["0"] ++ hits = { ++ e[0]: e[1] for e in runs["1"] if isinstance(e[0], str) and e[0][:3] == "hit" ++ } ++ assert hits["hitA"] == 0 and hits["hitB"] == 0 ++ assert hits["hitA2"] > 0 and hits["hitB2"] > 0 # real hits were compared ++ # CONTROL: the comparison is not vacuous -- a different accept pattern ++ # gives different tables, and the comparator sees it. ++ manager = _manager(num_blocks=400) ++ assert _trajectory(manager, accept_seed=1) != runs["1"] ++ ++ ++def _held(manager: KVCacheManager) -> set[int]: ++ ids: set[int] = set() ++ for mgr in manager.coordinator.single_type_managers: ++ pool = getattr(mgr, "slot_pool", None) ++ if pool is not None: ++ ids |= set(pool._pages) ++ continue ++ for blocks in mgr.req_to_blocks.values(): ++ ids |= {b.block_id for b in blocks if not b.is_null} ++ held = getattr(mgr, "_retained_checkpoints", {}) ++ ids |= {blk.block_id for _, blk in held.values()} ++ return ids ++ ++ ++def _owner_count(manager: KVCacheManager) -> int: ++ """Owner references to pool blocks, counting shared prefix blocks once ++ per holder. Equal to len(_held) iff no block has two exclusive owners.""" ++ total = 0 ++ for mgr in manager.coordinator.single_type_managers: ++ pool = getattr(mgr, "slot_pool", None) ++ if pool is not None: ++ total += pool.num_pages ++ continue ++ total += len( ++ { ++ b.block_id ++ for blocks in mgr.req_to_blocks.values() ++ for b in blocks ++ if not b.is_null ++ } ++ ) ++ total += len(getattr(mgr, "_retained_checkpoints", {})) ++ return total ++ ++ ++def test_rings_share_one_pool_block_and_pool_partitions(layout): ++ manager = _manager(num_blocks=400) ++ pool = manager.block_pool ++ ring = _ring(manager) ++ reqs = [_request(f"R{i}", _tokens(10 + i, 70)) for i in range(5)] ++ for req in reqs: ++ assert _prefill(manager, req) == 0 ++ _decode(manager, req, 2) ++ # Partition: every non-null block is free XOR held, by one owner. ++ assert pool.get_num_free_blocks() + len(_held(manager)) == 399 ++ assert _owner_count(manager) == len(_held(manager)) ++ ring_ids = [ring.get_ring_block_ids(r.request_id)[0] for r in reqs] ++ assert len(set(ring_ids)) == 5 ++ claimed = manager.take_new_ring_block_ids() ++ assert claimed == ring_ids # tag clears name exactly the claimed ids ++ if layout: ++ slots = ring.slot_pool.slots_per_block ++ assert slots > 5 ++ assert ring.slot_pool.num_pages == 1 # five rings, one pool block ++ assert len({rid // slots for rid in ring_ids}) == 1 ++ assert sorted(rid % slots for rid in ring_ids) == [0, 1, 2, 3, 4] ++ else: ++ # FAILING CONTROL: one pool block per ring. ++ assert ring.slot_pool is None ++ assert all(pool.blocks[rid].ref_cnt == 1 for rid in ring_ids) ++ # Free in a mixed order, one through the deferred path. ++ manager.free(reqs[1]) ++ popped = manager.pop_blocks_for_free(reqs[3]) ++ if layout: ++ # The popped slot is not reusable before the fence: still live. ++ assert ring.slot_pool.num_live_slots == 4 ++ manager.block_pool.free_blocks(popped) ++ if layout: ++ assert ring.slot_pool.num_live_slots == 3 ++ # A new request reuses the lowest free slot of the same page. ++ r5 = _request("R5", _tokens(30, 70)) ++ _prefill(manager, r5) ++ assert ring.get_ring_block_ids("R5") == [ring_ids[1]] ++ manager.free(r5) ++ for req in (reqs[0], reqs[2], reqs[4]): ++ manager.free(req) ++ assert pool.get_num_free_blocks() == 399 ++ if layout: ++ assert ring.slot_pool.num_pages == 0 ++ ++ ++def _running_footprint(value: str, prompt: int) -> int: ++ with pytest.MonkeyPatch.context() as mp: ++ _set_switch(mp, value) ++ manager = _manager(num_blocks=400) ++ req = _request("X", _tokens(99, prompt)) ++ _prefill(manager, req) ++ _decode(manager, req, 3) ++ return 399 - manager.block_pool.get_num_free_blocks() ++ ++ ++def test_freed_blocks_admit_requests_the_old_layout_refuses(layout): ++ """Admission under pool pressure, through the full-sequence gate. The ++ per-request footprint is measured, not assumed; the pool is sized to six ++ requests of the NEW footprint plus less than one request's saving. ON ++ admits all six; the old layout cannot.""" ++ on, off = _running_footprint("1", 160), _running_footprint("0", 160) ++ # 2 + K PLE state blocks (retained checkpoint included) and the ring block, ++ # minus the one shared ring page the new layout holds for everyone. ++ assert off - on == (2 + K) + 1 - 1 ++ n = 6 ++ margin = off - on - 1 ++ manager = _manager(num_blocks=1 + n * on + margin) # + the null block ++ admitted = [] ++ for i in range(n): ++ req = _request(f"C{i}", _tokens(200 + i, 160)) ++ if _prefill(manager, req, may_stall=True) is None: ++ break ++ _decode(manager, req, 3) ++ admitted.append(req) ++ if layout: ++ assert len(admitted) == n ++ else: ++ # FAILING CONTROL: the same pool fits fewer conversations. ++ assert len(admitted) < n ++ for req in admitted: ++ manager.free(req) ++ assert manager.block_pool.get_num_free_blocks() == n * on + margin ++ ++ ++def test_ring_slot_pool_unit(): ++ pool = BlockPool(num_gpu_blocks=6, enable_caching=True, hash_block_size=16) ++ slots = RingSlotPool(pool, slots_per_block=3) ++ got = slots.get_new_blocks(4) ++ assert [b.block_id // 3 for b in got] == [got[0].block_id // 3] * 3 + [ ++ got[3].block_id // 3 ++ ] ++ assert slots.num_pages == 2 and pool.get_num_free_blocks() == 3 ++ # Routed through the main pool's free path (as the scheduler drain does). ++ pool.free_blocks([got[3]]) ++ assert slots.num_pages == 1 and pool.get_num_free_blocks() == 4 ++ pool.free_blocks(got[:2]) ++ assert slots.has_free_slot() and slots.num_pages == 1 ++ with pytest.raises(AssertionError): ++ slots.free_blocks([got[0]]) # double free is loud ++ pool.free_blocks([got[2]]) ++ assert slots.num_pages == 0 and pool.get_num_free_blocks() == 5 ++ ++ ++# --------------------------------------------------------------------------- ++# 3. Worker views and KV zeroing (CPU allocation of a 3-block pool). ++# --------------------------------------------------------------------------- ++def _views(monkeypatch, value: str, rank: int): ++ from vllm.v1.worker.utils import allocate_kv_cache ++ ++ _set_switch(monkeypatch, value) ++ # 24 blocks (~300 MB on CPU): one 3200-token request (22 blocks in the old ++ # layout) must still pass the boot's memory check. ++ _, _, configs = _configs(override=24, max_model_len=3200) ++ config = configs[rank] ++ caches = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLNHC) ++ return config, caches ++ ++ ++def _byte_range(view: torch.Tensor, index: int, base: int) -> tuple[int, int]: ++ page = view[index] ++ start = page.data_ptr() - base ++ span = page.element_size() * ( ++ 1 + sum((s - 1) * st for s, st in zip(page.shape, page.stride())) ++ ) ++ return start, start + span ++ ++ ++@pytest.mark.parametrize("rank", [0, 2], ids=["pp0", "pp1"]) ++def test_slot_and_ple_views_stay_inside_their_block(monkeypatch, rank): ++ config, caches = _views(monkeypatch, "1", rank) ++ stride = ku._get_kv_cache_bytes_per_block(config.kv_cache_groups) ++ base = min(v.data_ptr() for v in caches.values()) ++ rings = sorted(n for n in caches if n.endswith("raw_key_cache")) ++ slots = 800 ++ assert all(caches[n].shape[0] == 24 * slots for n in rings) ++ for pool_block in (1, 2): ++ spans = [] ++ for slot in (0, 1, slots - 1): ++ for name in rings: ++ lo, hi = _byte_range(caches[name], pool_block * slots + slot, base) ++ assert pool_block * stride <= lo < hi <= (pool_block + 1) * stride ++ spans.append((lo, hi)) ++ spans.sort() ++ assert all(a[1] <= b[0] for a, b in zip(spans, spans[1:])), "slots overlap" ++ if rank == 0: ++ group = next(g for g in config.kv_cache_groups if PLE in g.layer_names) ++ lo, hi = _byte_range(caches[PLE], 1, base) ++ assert stride <= lo < hi <= 2 * stride ++ for name in group.layer_names[:-1]: # the GDN layers of that group ++ for state in ( ++ caches[name] if isinstance(caches[name], list) else [caches[name]] ++ ): ++ glo, ghi = _byte_range(state, 1, base) ++ assert ghi <= lo or glo >= hi, f"PLE page overlaps {name}" ++ ++ ++def _zeroed_offsets(config, caches, skip: bool) -> tuple[set, int]: ++ """Block-relative byte ranges the KV zeroer would clear, from its own ++ segment table (built on CPU tensors; the kernel is not launched).""" ++ from vllm.v1.worker.utils import KVBlockZeroer ++ ++ groups = [] ++ for gid, group in enumerate(config.kv_cache_groups): ++ for name in group.layer_names: ++ spec = group.kv_cache_spec ++ if isinstance(spec, UniformTypeKVCacheSpecs): ++ spec = spec.kv_cache_specs[name] ++ groups.append( ++ SimpleNamespace( ++ kv_cache_spec=spec, kv_cache_group_id=gid, layer_names=[name] ++ ) ++ ) ++ kernel_block_sizes = [_first_spec(g).block_size for g in config.kv_cache_groups] ++ ctx = {name: SimpleNamespace(kv_cache=view) for name, view in caches.items()} ++ zeroer = KVBlockZeroer( ++ torch.device("cpu"), ++ attn_groups_iter=iter(groups), ++ kernel_block_sizes=kernel_block_sizes, ++ static_forward_context=ctx, ++ num_blocks=config.num_blocks, ++ skip_kv_cache_group_ids=config.slot_packed_group_ids if skip else (), ++ ) ++ addrs, strides, pages, *_ = zeroer._meta ++ num_segments = len(addrs) ++ base = min(v.data_ptr() for v in caches.values()) ++ stride = ku._get_kv_cache_bytes_per_block(config.kv_cache_groups) ++ covered: list[tuple[int, int]] = [] ++ for addr, page in zip(addrs.tolist(), pages.tolist()): ++ lo = (addr - base) % stride ++ covered.append((lo, lo + 4 * page)) ++ # Merge into a canonical union of intervals. ++ covered.sort() ++ merged: list[list[int]] = [] ++ for lo, hi in covered: ++ if merged and lo <= merged[-1][1]: ++ merged[-1][1] = max(merged[-1][1], hi) ++ else: ++ merged.append([lo, hi]) ++ return {tuple(m) for m in merged}, num_segments ++ ++ ++@pytest.mark.parametrize("rank", [0, 2], ids=["pp0", "pp1"]) ++def test_zeroed_bytes_per_block_are_unchanged(monkeypatch, rank): ++ old_config, old_caches = _views(monkeypatch, "0", rank) ++ new_config, new_caches = _views(monkeypatch, "1", rank) ++ old, old_segments = _zeroed_offsets(old_config, old_caches, skip=True) ++ new, new_segments = _zeroed_offsets(new_config, new_caches, skip=True) ++ assert new == old ++ # The old ring segments were redundant (inside the first attention page); ++ # the new table drops exactly those. ++ assert new_segments <= old_segments ++ # FAILING CONTROL: without the slot-group skip every ring slot becomes a ++ # segment. On PP0 that zeroes bytes the old layout never zeroed (past the ++ # attention pages, inside the 7th GDN page); on PP1 the attention pages ++ # already cover the block, so it is "only" 5,600 redundant segments. ++ unskipped, unskipped_segments = _zeroed_offsets(new_config, new_caches, skip=False) ++ assert unskipped_segments > new_segments + 4000 # ~800 per ring layer ++ if rank == 0: ++ assert unskipped != old ++ else: ++ assert unskipped == old +diff --git a/tests/v1/test_pp_kv_projection.py b/tests/v1/test_pp_kv_projection.py +index 2f8cc59..cbf5956 100644 +--- a/tests/v1/test_pp_kv_projection.py ++++ b/tests/v1/test_pp_kv_projection.py +@@ -170,10 +170,22 @@ def _allocate_lookup(config) -> None: + ) + + +-def test_flash_next_pp1_has_an_empty_uniform_group(): ++def _side_cache_layout(monkeypatch, value: str) -> None: ++ monkeypatch.setenv("VLLM_E2_SIDE_CACHE_LAYOUT", value) ++ from vllm import envs ++ ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ ++ ++def test_flash_next_pp1_has_an_empty_uniform_group(monkeypatch): + """Guard that the shape really exercises the path (else the test below + could not fail): PP1 projects the PLE group to zero layers while keeping +- the global UniformTypeKVCacheSpecs.""" ++ the global UniformTypeKVCacheSpecs. v2.5: that shape exists only with the ++ side-cache layout off (VLLM_E2_SIDE_CACHE_LAYOUT=0, or an engine it does ++ not apply to); with it on the PLE state rides in a GDN group, so no group ++ projects empty (test_side_cache_layout.py).""" ++ _side_cache_layout(monkeypatch, "0") + _, configs = _worker_configs() + for rank in (2, 3): + empty = [g for g in configs[rank].kv_cache_groups if not g.layer_names] +@@ -186,7 +198,9 @@ def test_flash_next_pp1_has_an_empty_uniform_group(): + assert all(g.layer_names for g in configs[rank].kv_cache_groups) + + +-def test_flash_next_pp1_every_tensor_is_grouped(): ++@pytest.mark.parametrize("layout", ["0", "1"]) ++def test_flash_next_pp1_every_tensor_is_grouped(monkeypatch, layout): ++ _side_cache_layout(monkeypatch, layout) + specs, configs = _worker_configs() + for rank, config in enumerate(configs): + tensor_layers = [n for t in config.kv_cache_tensors for n in t.layers] +diff --git a/vllm/envs.py b/vllm/envs.py +index 4159f34..7401ed7 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -452,6 +452,16 @@ if TYPE_CHECKING: + # for A/B. AOT: no. CAPTURE: no -- scheduler-side block bookkeeping only; + # the worker sees the same block tables either way. + VLLM_E2_MAMBA_RETAIN_CHECKPOINT: bool = True ++ # v2.5 side-cache layout (Flash-Next PLE + QSA ring): stop charging each ++ # tiny side cache a whole packed pool block per request. The TP-replicated ++ # PLE short-conv state group is folded into a GDN group that has a spare ++ # page on every PP stage, and the one-block-per-request QSA ring group packs ++ # several requests' rings into one pool block. Strict "0"/"1", default ON; ++ # "0" restores the old one-group-each layout for A/B. Read in ++ # vllm/v1/core/kv_cache_utils.py (get_kv_cache_configs). AOT: yes, hashed -- ++ # it changes the ring views' block stride and which block table the PLE ++ # state is read through, so ON and OFF must not share a compile cache. ++ VLLM_E2_SIDE_CACHE_LAYOUT: bool = True + # Opt-out for the fork startup guards in vllm/v1/e2_guards.py, read at + # Qwen4ExpForCausalLM construction. AOT: no. CAPTURE: no -- it changes + # nothing itself; it only permits a boot whose identity IS different +@@ -2615,6 +2625,11 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_E2_MAMBA_RETAIN_CHECKPOINT": lambda: _e2_strict_switch( + "VLLM_E2_MAMBA_RETAIN_CHECKPOINT", default=True + ), ++ # Side-cache layout (v2.5; PLE state folded into a GDN group + QSA ring slot ++ # packing), default ON. Strict "0"/"1"; read once per get_kv_cache_configs call. ++ "VLLM_E2_SIDE_CACHE_LAYOUT": lambda: _e2_strict_switch( ++ "VLLM_E2_SIDE_CACHE_LAYOUT", default=True ++ ), + # Seconds between `topk-exact-counters` lines from each GPU worker + # (0, the default, disables the reader entirely). The line reports the + # medium-path exact-restart counters; see +diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py +index 4b2839d..7537ccd 100644 +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -14,6 +14,7 @@ from vllm.v1.core.kv_cache_utils import ( + dcp_world_size_for_kv_cache_spec, + ) + from vllm.v1.core.single_type_kv_cache_manager import ( ++ CircularBufferManager, + CrossAttentionManager, + MambaManager, + SingleTypeKVCacheManager, +@@ -151,6 +152,17 @@ class KVCacheCoordinator(ABC): + ) + for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) + ) ++ # v2.5 side-cache layout: a slot-packed ring group shares each pool ++ # block between several requests (kv_cache_utils._choose_ring_slots). ++ for kv_cache_group, manager in zip( ++ self.kv_cache_config.kv_cache_groups, self.single_type_managers ++ ): ++ if kv_cache_group.slots_per_block > 1: ++ assert type(manager) is CircularBufferManager, ( ++ f"slot packing is only defined for the QSA ring, got " ++ f"{type(manager).__name__}" ++ ) ++ manager.enable_slot_packing(kv_cache_group.slots_per_block) + # Match Mamba checkpoints to Eagle's attention replay boundary. + if use_eagle: + for manager in self.single_type_managers: +diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py +index ad179e6..2d81658 100644 +--- a/vllm/v1/core/kv_cache_utils.py ++++ b/vllm/v1/core/kv_cache_utils.py +@@ -1082,7 +1082,19 @@ def get_max_concurrency_for_kv_cache_config( + """ + num_blocks_per_request = 0 + host_blocks_per_request = 0 ++ # v2.5 side-cache layout: a slot-packed ring group shares each pool block ++ # between `slots_per_block` requests, so it costs cdiv(max_num_seqs, slots) ++ # blocks for the whole pool instead of one block per request. ++ shared_blocks = 0 ++ max_num_seqs = getattr( ++ getattr(vllm_config, "scheduler_config", None), "max_num_seqs", 1 ++ ) ++ if not isinstance(max_num_seqs, int) or max_num_seqs < 1: ++ max_num_seqs = 1 + for group in kv_cache_config.kv_cache_groups: ++ if group.slots_per_block > 1 and not group.host_resident: ++ shared_blocks += cdiv(max_num_seqs, group.slots_per_block) ++ continue + required = cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, +@@ -1091,7 +1103,9 @@ def get_max_concurrency_for_kv_cache_config( + host_blocks_per_request += required + else: + num_blocks_per_request += required +- limits = [kv_cache_config.num_blocks / num_blocks_per_request] ++ limits = [ ++ max(kv_cache_config.num_blocks - shared_blocks, 0) / num_blocks_per_request ++ ] + if host_blocks_per_request: + assert kv_cache_config.hisparse_host_num_blocks is not None + limits.append( +@@ -1779,14 +1793,36 @@ def get_kv_cache_config_from_groups( + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + ++ # v2.5 side-cache layout: a slot-packed group views every pool block as ++ # `slots` equal slots, so its block stride is one slot and its view has ++ # num_blocks * slots blocks. Pages of its layers must fit in one slot. ++ slots = group.slots_per_block ++ group_block_stride = interleaved_block_stride ++ group_num_blocks = num_blocks ++ if slots > 1: ++ assert layout.is_block_outermost, ( ++ "Slot-packed KV cache groups need a block-outermost layout" ++ ) ++ assert bytes_per_block % slots == 0, (bytes_per_block, slots) ++ group_block_stride = bytes_per_block // slots ++ group_num_blocks = num_blocks * slots ++ group_bytes = sum( ++ spec.page_size_bytes * len(names) ++ for spec, names in layers_by_spec.items() ++ ) ++ assert group_bytes <= group_block_stride, ( ++ f"slot-packed group needs {group_bytes} B per slot, slot is " ++ f"{group_block_stride} B" ++ ) ++ + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, +- num_blocks, ++ group_num_blocks, + len(layer_names), + layout, +- fixed_strides=(None, interleaved_block_stride, None, None, None), ++ fixed_strides=(None, group_block_stride, None, None, None), + ) + offset = ( + byte_offset +@@ -1800,6 +1836,7 @@ def get_kv_cache_config_from_groups( + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, ++ slots_per_block=slots, + ) + ) + byte_offset += len(layer_names) * spec.page_size_bytes +@@ -2650,6 +2687,289 @@ def _project_kv_cache_groups_to_worker( + return projected_groups + + ++# --------------------------------------------------------------------------- ++# v2.5 side-cache layout (VLLM_E2_SIDE_CACHE_LAYOUT, default ON). ++# ++# The packed block-outermost layout strides every pool block by the widest ++# group, and every group draws whole blocks from the one pool. Flash-Next has ++# two side caches whose real state is tiny next to that stride (12.9 MB on the ++# binding PP stage): the TP-replicated PLE short-conv state (one layer, 2 + K ++# blocks per request in align mode, a 1.6 MB padded page) and the QSA raw-key ++# ring (one block per request, 16 KB for all of a stage's layers). Together ++# they hold 6 of the 103 blocks a 262K request needs. Two layout-only changes ++# remove that cost without changing any value a kernel reads: ++# ++# 1. PLE fold. The replicated Mamba group joins a sharded Mamba group with the ++# same scheduling parameters whose block has a spare page on every worker ++# (Flash-Next: a GDN group holds 6 of PP0's 7 pages), so the PLE state rides ++# in that group's blocks under the block table it would have had on its ++# own. No worker's stride changes, so the pool block count does not change. ++# 2. Ring slot packing. The ring group splits every pool block into S equal ++# slots and the ring manager hands out slots; a pool block goes back to the ++# pool when its last slot is released. The ring then costs cdiv(requests, S) ++# blocks in all instead of one per request. ++# ++# Off (=0) restores one group each, block for block. ++# --------------------------------------------------------------------------- ++ ++# Slot stride alignment for packed ring views (bytes) and the largest S tried. ++_RING_SLOT_ALIGN = 256 ++_RING_MAX_SLOTS = 1024 ++ ++ ++def _side_cache_layout_enabled(vllm_config: VllmConfig) -> bool: ++ """Whether the v2.5 side-cache layout may apply to this engine. ++ ++ KV connectors (NIXL, offload, P/D) address groups and blocks themselves ++ (the NIXL worker requires exactly one single-layer replicated group), and ++ HiSparse has its own pool layout, so those engines keep the old layout. ++ """ ++ if not envs.VLLM_E2_SIDE_CACHE_LAYOUT: ++ return False ++ if getattr(vllm_config, "kv_transfer_config", None) is not None: ++ return False ++ attention_config = getattr(vllm_config, "attention_config", None) ++ if getattr(attention_config, "hisparse_config", None) is not None: ++ return False ++ parallel_config = getattr(vllm_config, "parallel_config", None) ++ if getattr(parallel_config, "decode_context_parallel_size", 1) not in (1, None): ++ return False ++ try: ++ layout = vllm_config.cache_config.get_resolved_kv_cache_layout() ++ except ValueError: ++ # Unresolved layout: only reachable from callers that bypass the ++ # engine core; nothing to pack without a block-outermost layout. ++ return False ++ return bool(layout.is_block_outermost) ++ ++ ++def _group_layer_specs(group: KVCacheGroupSpec) -> list[KVCacheSpec]: ++ spec = group.kv_cache_spec ++ if isinstance(spec, UniformTypeKVCacheSpecs): ++ return [spec.kv_cache_specs[name] for name in group.layer_names] ++ return [spec] * len(group.layer_names) ++ ++ ++def _mamba_schedule_key(spec: MambaSpec) -> tuple: ++ """Everything the scheduler's MambaManager and the worker's align-mode ++ bookkeeping read from a group's spec. Two Mamba groups with equal keys get ++ identical block tables for every request.""" ++ return ( ++ spec.block_size, ++ spec.page_size_bytes, ++ spec.num_speculative_blocks, ++ spec.num_prefill_checkpoint_blocks, ++ spec.prefill_checkpoint_alignment, ++ spec.mamba_cache_mode, ++ ) ++ ++ ++def _worker_strides( ++ groups: list[KVCacheGroupSpec], worker_specs: list[dict[str, KVCacheSpec]] ++) -> list[int]: ++ """Pool block bytes of every worker for this global grouping.""" ++ strides = [] ++ for spec in worker_specs: ++ projected = _project_kv_cache_groups_to_worker(groups, spec) ++ if not any(g.layer_names for g in projected): ++ strides.append(0) ++ continue ++ strides.append(_get_kv_cache_bytes_per_block(projected)) ++ return strides ++ ++ ++def _fold_replicated_mamba_groups( ++ groups: list[KVCacheGroupSpec], ++ worker_specs: list[dict[str, KVCacheSpec]], ++) -> tuple[list[KVCacheGroupSpec], list[str]]: ++ """Fold each TP-replicated Mamba group into a compatible sharded one. ++ ++ A target holds only sharded Mamba layers with the same scheduling key and ++ EAGLE flag, and the merged group must not widen the block of any worker ++ (checked on every worker's projection, so PP stages are covered). Among ++ valid targets the lowest index wins. Returns the new group list and one ++ log note per replicated group. ++ """ ++ groups = list(groups) ++ notes: list[str] = [] ++ base_strides = _worker_strides(groups, worker_specs) ++ index = 0 ++ while index < len(groups): ++ source = groups[index] ++ layer_specs = _group_layer_specs(source) ++ if not ( ++ layer_specs ++ and all( ++ isinstance(spec, MambaSpec) and spec.tp_replicated ++ for spec in layer_specs ++ ) ++ and len({_mamba_schedule_key(spec) for spec in layer_specs}) == 1 ++ ): ++ index += 1 ++ continue ++ key = _mamba_schedule_key(cast(MambaSpec, layer_specs[0])) ++ chosen = None ++ for target_index, target in enumerate(groups): ++ if target_index == index or target.is_eagle_group != source.is_eagle_group: ++ continue ++ target_specs = _group_layer_specs(target) ++ if not target_specs or not all( ++ isinstance(spec, MambaSpec) ++ and not spec.tp_replicated ++ and _mamba_schedule_key(spec) == key ++ for spec in target_specs ++ ): ++ continue ++ merged_specs: dict[str, KVCacheSpec] = dict( ++ zip(target.layer_names, target_specs) ++ ) ++ merged_specs.update(zip(source.layer_names, layer_specs)) ++ # Built directly: is_uniform_type rejects mixed tp_replicated on ++ # purpose (NIXL), and connectors never reach this path. ++ merged = KVCacheGroupSpec( ++ target.layer_names + source.layer_names, ++ UniformTypeKVCacheSpecs( ++ block_size=target_specs[0].block_size, ++ kv_cache_specs=merged_specs, ++ ), ++ is_eagle_group=target.is_eagle_group, ++ ) ++ trial = [ ++ merged if i == target_index else g ++ for i, g in enumerate(groups) ++ if i != index ++ ] ++ strides = _worker_strides(trial, worker_specs) ++ if all(new <= old for new, old in zip(strides, base_strides)): ++ chosen = (target_index, merged, trial) ++ break ++ if chosen is None: ++ notes.append( ++ f"replicated Mamba group {index} ({len(source.layer_names)} " ++ "layer(s)) kept separate: no sharded group has a spare page on " ++ "every worker" ++ ) ++ index += 1 ++ continue ++ target_index, merged, trial = chosen ++ notes.append( ++ f"folded replicated Mamba group {index} " ++ f"({', '.join(source.layer_names)}) into Mamba group " ++ f"{target_index} (now {len(merged.layer_names)} layers); per-worker " ++ f"block bytes unchanged {sorted(set(base_strides))}" ++ ) ++ groups = trial ++ # The group after the folded one now sits at `index`. ++ return groups, notes ++ ++ ++def _is_ring_group(group: KVCacheGroupSpec) -> bool: ++ spec = group.kv_cache_spec ++ layer_specs = ( ++ list(spec.kv_cache_specs.values()) ++ if isinstance(spec, UniformTypeKVCacheSpecs) ++ else [spec] ++ ) ++ return bool(layer_specs) and all( ++ type(layer_spec) is CircularBufferSpec for layer_spec in layer_specs ++ ) ++ ++ ++def _choose_ring_slots( ++ projected_groups_per_worker: list[list[KVCacheGroupSpec]], ++) -> dict[int, int]: ++ """Slots per pool block for every ring group (group index -> S >= 2). ++ ++ S must cut every worker's block stride into equal slots that are ++ `_RING_SLOT_ALIGN`-aligned and hold that worker's ring pages; the largest ++ such S up to `_RING_MAX_SLOTS` is taken. Block ids are global, so S is the ++ same on every worker. A ring group with no valid S keeps one pool block ++ per request. ++ """ ++ chosen: dict[int, int] = {} ++ if not projected_groups_per_worker or not projected_groups_per_worker[0]: ++ return chosen ++ num_groups = len(projected_groups_per_worker[0]) ++ strides = [ ++ _get_kv_cache_bytes_per_block(groups) ++ if any(g.layer_names for g in groups) ++ else 0 ++ for groups in projected_groups_per_worker ++ ] ++ for group_index in range(num_groups): ++ if not _is_ring_group(projected_groups_per_worker[0][group_index]): ++ continue ++ needs = [] ++ for groups, stride in zip(projected_groups_per_worker, strides): ++ group = groups[group_index] ++ if not group.layer_names: ++ continue ++ ring_bytes = sum(spec.page_size_bytes for spec in _group_layer_specs(group)) ++ needs.append((stride, ring_bytes)) ++ if not needs: ++ continue ++ for slots in range(_RING_MAX_SLOTS, 1, -1): ++ if all( ++ stride % slots == 0 ++ and (stride // slots) % _RING_SLOT_ALIGN == 0 ++ and stride // slots >= ring_bytes ++ for stride, ring_bytes in needs ++ ): ++ chosen[group_index] = slots ++ break ++ return chosen ++ ++ ++def _apply_side_cache_layout( ++ vllm_config: VllmConfig, ++ global_kv_cache_groups: list[KVCacheGroupSpec], ++ kv_cache_specs: list[dict[str, KVCacheSpec]], ++) -> tuple[list[KVCacheGroupSpec], list[list[KVCacheGroupSpec]]]: ++ """Fold side Mamba groups, project per worker, then pick ring slots. ++ ++ Returns the (possibly folded) global groups and every worker's projection, ++ with `slots_per_block` set on packed ring groups. ++ """ ++ relevant = any( ++ (isinstance(spec, MambaSpec) and spec.tp_replicated) ++ or type(spec) is CircularBufferSpec ++ for worker_spec in kv_cache_specs ++ for spec in worker_spec.values() ++ ) ++ enabled = relevant and _side_cache_layout_enabled(vllm_config) ++ notes: list[str] = [] ++ if enabled: ++ global_kv_cache_groups, notes = _fold_replicated_mamba_groups( ++ global_kv_cache_groups, kv_cache_specs ++ ) ++ projected = [ ++ _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) ++ for worker_spec in kv_cache_specs ++ ] ++ if enabled and all( ++ _glm5_next_tensor_layout(groups) is None for groups in projected if groups ++ ): ++ for group_index, slots in _choose_ring_slots(projected).items(): ++ for groups in projected: ++ groups[group_index].slots_per_block = slots ++ notes.append( ++ f"QSA ring group {group_index} packs {slots} requests per pool block" ++ ) ++ if relevant and envs.VLLM_E2_SIDE_CACHE_LAYOUT and not enabled: ++ notes.append( ++ "not applicable to this engine (KV connector, HiSparse, DCP or a " ++ "layer-outermost layout)" ++ ) ++ if relevant: ++ # The A/B evidence line: boot logs must show which arm ran. ++ logger.info( ++ "Side-cache layout: %s (VLLM_E2_SIDE_CACHE_LAYOUT)%s", ++ "on" if enabled else "off", ++ "".join(f"; {note}" for note in notes), ++ ) ++ return global_kv_cache_groups, projected ++ ++ + def get_kv_cache_configs( + vllm_config: VllmConfig, + kv_cache_specs: list[dict[str, KVCacheSpec]], +@@ -2725,10 +3045,11 @@ def get_kv_cache_configs( + # If original_max_model_len was -1, automatically + # determine the maximum model length that fits in available GPU memory. + # We use per-worker projected groups to account for PP sharding. +- projected_groups_per_worker = [ +- _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) +- for worker_spec in kv_cache_specs +- ] ++ # v2.5: the side-cache layout (PLE fold, ring slots) is applied here, where ++ # every worker's layers are known, so no PP stage's block can widen. ++ global_kv_cache_groups, projected_groups_per_worker = _apply_side_cache_layout( ++ vllm_config, global_kv_cache_groups, kv_cache_specs ++ ) + + # If `num_gpu_blocks_override` is set, the cache size that will actually + # be allocated is decoupled from the profiled `available_memory`: +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index 1402e8e..dc8f033 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -1554,6 +1554,10 @@ class Scheduler(SchedulerInterface): + pending = set(block_ids) + owners: dict[int, set[str]] = {} + for mgr in self.kv_cache_manager.coordinator.single_type_managers: ++ if getattr(mgr, "slot_pool", None) is not None: ++ # v2.5 slot-packed ring: its block ids are slot ids, not pool ++ # block ids, and its pages are never zeroed. ++ continue + for req_id, blocks in mgr.req_to_blocks.items(): + for blk in blocks: + if blk.block_id in pending: +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 0ec6bca..ca390ef 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -1227,16 +1227,119 @@ class SlidingWindowManager(SingleTypeKVCacheManager): + return 0 + + ++class _RingSlotPage: ++ """One pool block shared by up to ``slots`` requests' rings.""" ++ ++ __slots__ = ("block", "free_slots", "live") ++ ++ def __init__(self, block: KVCacheBlock, slots: int) -> None: ++ self.block = block ++ # Highest first, so pop() hands out the lowest free slot. ++ self.free_slots = list(range(slots - 1, -1, -1)) ++ self.live = 0 ++ ++ ++class RingSlotPool: ++ """v2.5 side-cache layout: ring blocks as slots of shared pool blocks. ++ ++ The QSA ring needs one block per request for the request's lifetime, but ++ its pages for every layer of a stage are ~16 KB against a 12.9 MB pool ++ block. This pool takes whole blocks from the main ``BlockPool`` and cuts ++ each into ``slots_per_block`` slots; a slot's block id is ++ ``pool_block_id * slots_per_block + slot``, which is exactly the index of ++ the slot in the worker's ring view (``KVCacheTensor.slots_per_block``). ++ ++ The slot blocks it hands out carry ``pool=self``, so every existing free ++ path -- ``SingleTypeKVCacheManager.free``, the coordinator's deferred ++ ``pop_blocks_for_free`` and the scheduler's fenced drain, which all call ++ ``BlockPool.free_blocks`` -- routes them back here, at the time the old ++ per-request block would have gone back (the deferred free is kept). A pool ++ block returns to the main pool when its last slot is freed. Rings are never ++ prefix-cached, so a slot block carries no hash. ++ """ ++ ++ def __init__(self, block_pool: BlockPool, slots_per_block: int) -> None: ++ assert slots_per_block > 1 ++ self.block_pool = block_pool ++ self.slots_per_block = slots_per_block ++ # Pool block id -> page, in claim order (the oldest page fills first). ++ self._pages: dict[int, _RingSlotPage] = {} ++ ++ @property ++ def num_pages(self) -> int: ++ """Pool blocks currently held for rings.""" ++ return len(self._pages) ++ ++ @property ++ def num_live_slots(self) -> int: ++ return sum(page.live for page in self._pages.values()) ++ ++ def has_free_slot(self) -> bool: ++ return any(page.free_slots for page in self._pages.values()) ++ ++ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ++ out: list[KVCacheBlock] = [] ++ for _ in range(num_blocks): ++ page = next((p for p in self._pages.values() if p.free_slots), None) ++ if page is None: ++ (pool_block,) = self.block_pool.get_new_blocks(1) ++ page = _RingSlotPage(pool_block, self.slots_per_block) ++ self._pages[pool_block.block_id] = page ++ slot = page.free_slots.pop() ++ page.live += 1 ++ out.append( ++ KVCacheBlock( ++ page.block.block_id * self.slots_per_block + slot, ++ ref_cnt=1, ++ pool=self, # type: ignore[arg-type] ++ ) ++ ) ++ return out ++ ++ def free_blocks(self, ordered_blocks) -> None: ++ released: list[KVCacheBlock] = [] ++ for block in ordered_blocks: ++ assert block.pool is self, "not a ring slot of this pool" ++ block.ref_cnt -= 1 ++ assert block.ref_cnt == 0, "ring slots have a single owner" ++ page_id, slot = divmod(block.block_id, self.slots_per_block) ++ page = self._pages[page_id] ++ assert slot not in page.free_slots, f"slot {block.block_id} freed twice" ++ page.free_slots.append(slot) ++ page.free_slots.sort(reverse=True) ++ page.live -= 1 ++ if page.live == 0: ++ del self._pages[page_id] ++ released.append(page.block) ++ if released: ++ self.block_pool.free_blocks(released) ++ ++ + class CircularBufferManager(FullAttentionManager): + """Claims the ring's single block per request; prefix caching disabled.""" + + supports_fine_grained_hash_lookup: ClassVar[bool] = False + ++ # v2.5 side-cache layout: set by the coordinator for a slot-packed group. ++ _slot_pool: RingSlotPool | None = None ++ ++ def enable_slot_packing(self, slots_per_block: int) -> None: ++ """Claim ring slots of shared pool blocks instead of whole blocks.""" ++ assert not self.req_to_blocks, "enable slot packing before any claim" ++ self._slot_pool = RingSlotPool(self.block_pool, slots_per_block) ++ ++ @property ++ def slot_pool(self) -> RingSlotPool | None: ++ return self._slot_pool ++ + def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + if req_blocks: + return [] +- new_blocks = self.block_pool.get_new_blocks(1) ++ if self._slot_pool is not None: ++ new_blocks = self._slot_pool.get_new_blocks(1) ++ else: ++ new_blocks = self.block_pool.get_new_blocks(1) + req_blocks.extend(new_blocks) + if self._record_new_block_ids: + self.new_block_ids.extend(block.block_id for block in new_blocks) +@@ -1281,7 +1384,12 @@ class CircularBufferManager(FullAttentionManager): + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: +- return 0 if self.req_to_blocks.get(request_id) else 1 ++ if self.req_to_blocks.get(request_id): ++ return 0 ++ # A free slot in a page the ring already holds costs no pool block. ++ if self._slot_pool is not None and self._slot_pool.has_free_slot(): ++ return 0 ++ return 1 + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int +diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py +index 8cd867e..8931b95 100644 +--- a/vllm/v1/kv_cache_interface.py ++++ b/vllm/v1/kv_cache_interface.py +@@ -1380,6 +1380,11 @@ class KVCacheTensor: + block_stride: int + offset: int = 0 # byte offset of layers[0]'s block 0 + host_resident: bool = False ++ # v2.5 side-cache layout: a slot-packed ring tensor views every pool block ++ # as ``slots_per_block`` equal slots (``block_stride`` is then one slot), so ++ # its block ids are ``pool_block * slots_per_block + slot``. 1 = one view ++ # block per pool block, the ordinary case. ++ slots_per_block: int = 1 + + + class KVCacheGroupRole(str, Enum): +@@ -1408,6 +1413,10 @@ class KVCacheGroupSpec: + # source group after a connector load. + enable_kv_transfer: bool = True + role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT ++ # v2.5 side-cache layout: > 1 when this one-block-per-request ring group ++ # shares each pool block between that many requests (see ++ # ``CircularBufferManager.enable_slot_packing``); 1 = one pool block each. ++ slots_per_block: int = 1 + + + @dataclass +@@ -1494,6 +1503,18 @@ class KVCacheConfig: + ) + return tuple(block_ids[group_id] for group_id in self.transfer_group_ids) + ++ @property ++ def slot_packed_group_ids(self) -> tuple[int, ...]: ++ """Groups whose blocks are slots of shared pool blocks (v2.5 side-cache ++ layout). Their views are not one block per pool block, so whole-block ++ zeroing must skip them; their bytes lie inside pool blocks the ring ++ manager owns and that no other group writes.""" ++ return tuple( ++ group_id ++ for group_id, group in enumerate(self.kv_cache_groups) ++ if group.slots_per_block > 1 ++ ) ++ + @property + def host_group_ids(self) -> tuple[int, ...]: + return tuple( +@@ -1505,7 +1526,7 @@ class KVCacheConfig: + def num_blocks_of(self, tensor: KVCacheTensor) -> int: + """Number of blocks addressable by the pool backing ``tensor``.""" + if not tensor.host_resident: +- return self.num_blocks ++ return self.num_blocks * tensor.slots_per_block + assert self.hisparse_host_num_blocks is not None + return self.hisparse_host_num_blocks + +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 1876f4b..d2bb002 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -843,6 +843,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + kernel_block_sizes=self.kernel_block_sizes, + static_forward_context=self.compilation_config.static_forward_context, + num_blocks=self.kv_cache_config.num_blocks, ++ skip_kv_cache_group_ids=self.kv_cache_config.slot_packed_group_ids, + ) + + @torch.inference_mode() +diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py +index 910c72c..010f110 100644 +--- a/vllm/v1/worker/gpu_model_runner.py ++++ b/vllm/v1/worker/gpu_model_runner.py +@@ -1173,6 +1173,7 @@ class GPUModelRunner( + runner_only_attn_layers=self.runner_only_attn_layers, + static_forward_context=self.compilation_config.static_forward_context, + num_blocks=self.kv_cache_config.num_blocks, ++ skip_kv_cache_group_ids=self.kv_cache_config.slot_packed_group_ids, + ) + + def _zero_block_ids(self, block_ids: list[int]) -> None: +diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py +index 61698bf..c09a0f7 100644 +--- a/vllm/v1/worker/utils.py ++++ b/vllm/v1/worker/utils.py +@@ -138,6 +138,7 @@ class KVBlockZeroer: + static_forward_context: dict[str, Any], + num_blocks: int, + runner_only_attn_layers: set[str] | None = None, ++ skip_kv_cache_group_ids: Iterable[int] = (), + ) -> None: + """Precompute the absolute-address table for the Triton zeroing kernel. + +@@ -156,6 +157,13 @@ class KVBlockZeroer: + physical block stride and zeroed page span remain independent. + + Only AttentionSpec layers are processed; Mamba layers are skipped. ++ ++ ``skip_kv_cache_group_ids`` (v2.5 side-cache layout): slot-packed ring ++ groups, whose views hold several slots per pool block. Their claims ++ never enter whole-block zeroing (the ring-tag path clears their tag words instead), ++ and with one view block per pool block their pages sat inside the ++ first attention page, so skipping them leaves the zeroed bytes of ++ every block unchanged. + """ + self.device = device + # Bounds-guard instrumentation: pool bound for the guard in _zero_kv_blocks_kernel. +@@ -174,12 +182,15 @@ class KVBlockZeroer: + seg_block_strides: list[int] = [] + seg_page_sizes: list[int] = [] + ++ skip_group_ids = set(skip_kv_cache_group_ids) + for group in attn_groups_iter: + spec = group.kv_cache_spec + if not isinstance(spec, AttentionSpec): + continue + if group.kv_cache_group_id >= len(kernel_block_sizes): + continue ++ if group.kv_cache_group_id in skip_group_ids: ++ continue + kernel_bs = kernel_block_sizes[group.kv_cache_group_id] + assert spec.block_size % kernel_bs == 0 + for layer_name in group.layer_names: +-- +2.43.0 + diff --git a/patches/0089-ple-VLLM_E2_PLE_PULL_TRANSPORT-defaults-to-on-when-u.patch b/patches/0089-ple-VLLM_E2_PLE_PULL_TRANSPORT-defaults-to-on-when-u.patch new file mode 100644 index 0000000..747af82 --- /dev/null +++ b/patches/0089-ple-VLLM_E2_PLE_PULL_TRANSPORT-defaults-to-on-when-u.patch @@ -0,0 +1,126 @@ +From 46b2370779cd9b41d4107b06ca916d1bc4f00345 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 4 Oct 2026 09:07:17 -0400 +Subject: [PATCH 089/103] ple: VLLM_E2_PLE_PULL_TRANSPORT defaults to on when + unset + +With the switch unset, the PLE took upstream's UVA backend, whose Triton lookup kernel needs fp8e4nv, which sm_86 +does not have, so a launcher without the switch failed at profile_run. Unset now means on, the shipped value; +0 still selects upstream's path. +--- + tests/v1/test_e2_ple_mode.py | 27 +++++++++++++++++++----- + vllm/envs.py | 7 ++++-- + vllm/models/qwen4_exp/nvidia/ple_pull.py | 5 +++-- + 3 files changed, 30 insertions(+), 9 deletions(-) + +diff --git a/tests/v1/test_e2_ple_mode.py b/tests/v1/test_e2_ple_mode.py +index 4017771..a614a0c 100644 +--- a/tests/v1/test_e2_ple_mode.py ++++ b/tests/v1/test_e2_ple_mode.py +@@ -76,7 +76,8 @@ def _config(cpu_offload: bool | None, ple_layer_ids=(2,)): + + + @pytest.mark.parametrize( +- ("raw", "expected"), [(None, False), ("0", False), ("1", True)] ++ # Fork (v2.5): unset = ON (the shipped transport); "0" selects upstream. ++ ("raw", "expected"), [(None, True), ("0", False), ("1", True)] + ) + def test_switch_accepts_exactly_unset_0_1(monkeypatch, raw, expected) -> None: + _set(monkeypatch, PULL_TRANSPORT_ENV, raw) +@@ -109,6 +110,19 @@ def test_switch_rejection_control_a_lax_parser_accepts_true(monkeypatch) -> None + assert envs._e2_strict_switch(PULL_TRANSPORT_ENV) is True + + ++def test_switch_unset_defaults_on_and_zero_still_selects_upstream(monkeypatch) -> None: ++ """Fork (v2.5, knob-absent boot): unset = the shipped pull transport, because the ++ upstream UVA path compiles an fp8e4nv kernel sm_86 cannot build. "0" keeps ++ selecting upstream explicitly; spellings stay strict.""" ++ monkeypatch.delenv(PULL_TRANSPORT_ENV, raising=False) ++ assert envs.VLLM_E2_PLE_PULL_TRANSPORT is True ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "0") ++ assert envs.VLLM_E2_PLE_PULL_TRANSPORT is False ++ monkeypatch.setenv(PULL_TRANSPORT_ENV, "true") ++ with pytest.raises(ValueError): ++ envs.VLLM_E2_PLE_PULL_TRANSPORT ++ ++ + def test_switch_is_registered_and_hashed(monkeypatch) -> None: + """Registered, live, and NOT in compile_factors' deny-list. + +@@ -152,7 +166,7 @@ def test_old_name_other_values_with_pull_on_are_refused(monkeypatch, old) -> Non + check_upstream_uva_env_against_pull() + + +-@pytest.mark.parametrize("pull", [None, "0"]) ++@pytest.mark.parametrize("pull", ["0"]) # v2.5: unset is pull ON + def test_old_name_control_pull_off_leaves_it_to_upstream(monkeypatch, pull) -> None: + """Failing control: the SAME VLLM_PLE_CPU_OFFLOAD=1 passes with pull off.""" + _set(monkeypatch, PULL_TRANSPORT_ENV, pull) +@@ -165,6 +179,7 @@ def test_old_name_control_pull_off_leaves_it_to_upstream(monkeypatch, pull) -> N + + @pytest.mark.parametrize("upstream_default", [True, False]) + def test_default_is_upstream_verbatim_with_pull_off(monkeypatch, upstream_default): ++ _set(monkeypatch, PULL_TRANSPORT_ENV, "0") # v2.5: unset is pull ON + assert pull_transport_cpu_offload_default(upstream_default) is upstream_default + + +@@ -186,10 +201,12 @@ def test_default_with_pull_on_still_refuses_the_old_name(monkeypatch): + @pytest.mark.parametrize( + ("pull", "cpu_offload", "mode"), + [ +- (None, None, PLE_MODE_NONE), +- (None, True, PLE_MODE_UVA), +- (None, False, PLE_MODE_DEVICE), ++ # v2.5: unset = pull ON; the upstream modes need an explicit "0". ++ (None, None, PLE_MODE_PULL), ++ (None, False, PLE_MODE_PULL), ++ ("0", None, PLE_MODE_NONE), + ("0", True, PLE_MODE_UVA), ++ ("0", False, PLE_MODE_DEVICE), + ("1", False, PLE_MODE_PULL), + ("1", None, PLE_MODE_PULL), + ], +diff --git a/vllm/envs.py b/vllm/envs.py +index 7401ed7..9fe769a 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -389,7 +389,10 @@ if TYPE_CHECKING: + # UVA backend is refused (vllm/v1/e2_guards.py). Off: the tree is upstream. + # AOT: yes (different module class and forward). CAPTURE: yes (the captured + # graph contains the wait/copy ops instead of a table lookup). Hashed. +- VLLM_E2_PLE_PULL_TRANSPORT: bool = False ++ # UNSET = ON (since v2.5), the shipped value: upstream's UVA backend (the off ++ # path) compiles an fp8e4nv Triton kernel that sm_86 cannot build, so a launcher ++ # without this switch failed at profile_run. "0" still selects upstream's path. ++ VLLM_E2_PLE_PULL_TRANSPORT: bool = True + # + # RUNTIME-ONLY (registered for validate_environ + a single cached read): + # Static K/V scale sidecar: fill_()ed into existing fp32 device buffers +@@ -2437,7 +2440,7 @@ environment_variables: dict[str, Callable[[], Any]] = { + "MERLIN_TOPK_EXACT_MEDIUM": lambda: os.getenv("MERLIN_TOPK_EXACT_MEDIUM", None), + # PLE pull transport switch; strict "0"/"1" (see the TYPE_CHECKING block). + "VLLM_E2_PLE_PULL_TRANSPORT": lambda: _e2_strict_switch( +- "VLLM_E2_PLE_PULL_TRANSPORT" ++ "VLLM_E2_PLE_PULL_TRANSPORT", default=True + ), + # Runtime only; all of the following are in compile_factors.ignored_factors. + "VLLM_E2_ALLOW_BREAKABLE": lambda: os.getenv("VLLM_E2_ALLOW_BREAKABLE", "") == "1", +diff --git a/vllm/models/qwen4_exp/nvidia/ple_pull.py b/vllm/models/qwen4_exp/nvidia/ple_pull.py +index f52000c..25d9640 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_pull.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_pull.py +@@ -15,8 +15,9 @@ class: ``Qwen4ExpNGramEmbeddingPull`` inherits it and ``PleOffloadLayer``. + table is upstream's ``Qwen4ExpPLEDeviceEmbedding`` allocated on CPU; + ``forward_impl`` computes n-gram IDs and gathers raw (FP8) rows. + +-Selected only when ``VLLM_E2_PLE_PULL_TRANSPORT=1``; with the switch off this +-module is never imported by the model and the tree is upstream. ++Selected when ``VLLM_E2_PLE_PULL_TRANSPORT`` is ``1`` or unset (the fork default ++since v2.5); with ``=0`` this module is never imported by the model and the ++tree is upstream. + + Supported table formats: FP8 (global scale) and unquantized -- the two + ``Qwen4ExpPLEEmbeddingMethod`` families v0.30.0 itself supports. The E2 NVFP4 +-- +2.43.0 + diff --git a/patches/0090-envs-the-MERLIN_-kernel-toggles-default-to-the-shipp.patch b/patches/0090-envs-the-MERLIN_-kernel-toggles-default-to-the-shipp.patch new file mode 100644 index 0000000..16d1a37 --- /dev/null +++ b/patches/0090-envs-the-MERLIN_-kernel-toggles-default-to-the-shipp.patch @@ -0,0 +1,201 @@ +From 5e39153dbfb58d752f93ce708c230db0f3987772 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Sun, 4 Oct 2026 09:23:24 -0400 +Subject: [PATCH 090/103] envs: the MERLIN_* kernel toggles default to the + shipped value when unset + +MERLIN_FULL_K, MERLIN_QSA_SORT, MERLIN_TIE_DET, MERLIN_TIE_RECENT and MERLIN_TOPK_SORTED_EMIT defaulted to off in +code while every shipped launcher set them to 1, so a launch without them ran a different kernel configuration. +vllm/env_override.py now sets each to 1 when unset; 0 still turns one off. tests/test_merlin_toggle_defaults.py +covers it. +--- + .../moe/marlin_moe_wna16/ops.cu | 2 +- + csrc/libtorch_stable/persistent_topk.cuh | 4 +- + csrc/libtorch_stable/topk.cu | 8 ++-- + tests/test_merlin_toggle_defaults.py | 48 +++++++++++++++++++ + vllm/env_override.py | 16 +++++++ + vllm/envs.py | 2 +- + .../qwen4_exp/nvidia/ops/qsa_indexer.py | 4 +- + 7 files changed, 74 insertions(+), 10 deletions(-) + create mode 100644 tests/test_merlin_toggle_defaults.py + +diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +index 0a6f5de..12a6a2c 100644 +--- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu ++++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +@@ -49,7 +49,7 @@ namespace { + // Resolved ONCE per process, not per launch, so it is a launch constant in + // exactly the sense topk.cu's exact_medium_enabled() is: a cudagraph bakes in + // the value that was live at capture time and every replay agrees with it. +-// Default OFF -- enabling it by default is a separate release decision. ++// Code default off; vllm/env_override.py sets 1 (the shipped value) when unset. + // + // This std::getenv is the SINGLE source of truth. vllm/envs.py carries a raw + // passthrough entry for validate_environ and the compile-factor hash only; +diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh +index 0b1b8b5..c86e63d 100644 +--- a/csrc/libtorch_stable/persistent_topk.cuh ++++ b/csrc/libtorch_stable/persistent_topk.cuh +@@ -58,7 +58,7 @@ __device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + // both the host (to report it) and the device (to obey it), so the reported + // policy and the executed policy cannot drift. + // +-// `MERLIN_TIE_DET=0` is the shipped default and leaves the buffered paths on ++// `MERLIN_TIE_DET=0` (the code default; launchers ship 1) leaves the buffered paths on + // their original racy atomicAdd rule, which has no direction; the medium-path exact + // restart is deterministic whatever happens, so it keeps the most-recent-wins + // direction the medium-path restart shipped and was measured with. +@@ -209,7 +209,7 @@ struct PersistentTopKParams { + // large tail) re-emit their OWN selected set through sorted_emit_row(). + // Selection is never re-derived, so the selected set cannot change. + // Resolved once per process on the host, so it is a launch constant a +- // cudagraph bakes in. Default OFF. ++ // cudagraph bakes in. Code default off; vllm/env_override.py sets 1 (the shipped value) when unset. + bool sorted_emit; + }; + +diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu +index 060f7a6..3200966 100644 +--- a/csrc/libtorch_stable/topk.cu ++++ b/csrc/libtorch_stable/topk.cu +@@ -38,7 +38,7 @@ namespace { + // + // Our own toggle rather than upstream's VLLM_BATCH_INVARIANT, which aborts + // engine startup on GDN models (upstream #42960) and so can never reach this +-// kernel on Flash-Next. Default OFF -- promotion is a release decision. ++// kernel on Flash-Next. Code default off; vllm/env_override.py sets 1 (the shipped value) when unset. + // + // This std::getenv is the SINGLE source of truth; vllm/envs.py carries a raw + // passthrough for validate_environ and the compile-factor hash only. +@@ -72,7 +72,7 @@ bool publish_cached(std::atomic& mirror, bool value) { + // most-recent-wins is +1.5 %. Both deterministic. Also the direction + // medium_exact_restart() already implements, so setting this alongside + // MERLIN_TIE_DET makes the medium path agree with itself on overflowing +-// rows. Only read when MERLIN_TIE_DET is set. Default OFF. ++// rows. Only read when MERLIN_TIE_DET is set; code default OFF; vllm/env_override.py sets 1 (the shipped value) when unset. + bool tie_recent_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MERLIN_TIE_RECENT"); +@@ -148,8 +148,8 @@ bool exact_decode_enabled() { + // (ascending index, padding last as -1), so qsa_indexer.py can drop that sort. + // See sorted_emit_row() in persistent_topk.cuh for the contract; the selected + // SET is never touched. Resolved once per process like every switch here, so +-// it is a launch constant a cudagraph bakes in. Default OFF: the maintainer +-// promotes after reference-host qualification. ++// it is a launch constant a cudagraph bakes in. Code default OFF; vllm/env_override.py ++// sets 1 (the shipped value) when unset. + // + // This std::getenv is the single source of truth. Python does not parse the + // variable to decide anything: qsa_indexer.py asks THIS cache, through +diff --git a/tests/test_merlin_toggle_defaults.py b/tests/test_merlin_toggle_defaults.py +new file mode 100644 +index 0000000..bb17582 +--- /dev/null ++++ b/tests/test_merlin_toggle_defaults.py +@@ -0,0 +1,48 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Fork (v2.5): the MERLIN_* kernel toggles default to the shipped value (1) when ++unset, and an explicit "0" is left alone. Each case runs `import vllm` in a ++fresh interpreter, since env_override applies the defaults at import time.""" ++ ++import json ++import os ++import subprocess ++import sys ++ ++import pytest ++ ++NAMES = ( ++ "MERLIN_FULL_K", ++ "MERLIN_QSA_SORT", ++ "MERLIN_TIE_DET", ++ "MERLIN_TIE_RECENT", ++ "MERLIN_TOPK_SORTED_EMIT", ++) ++ ++ ++def _env_after_import(overrides: dict) -> dict: ++ env = {k: v for k, v in os.environ.items() if k not in NAMES} ++ env.update(overrides) ++ code = ( ++ "import json, os, vllm, vllm.envs as e;" ++ f"print(json.dumps({{n: os.environ.get(n) for n in {NAMES!r}}}" ++ " | {'QSA_SORT_ENVS': e.MERLIN_QSA_SORT}))" ++ ) ++ out = subprocess.run( ++ [sys.executable, "-c", code], env=env, capture_output=True, text=True, check=True ++ ) ++ return json.loads(out.stdout.strip().splitlines()[-1]) ++ ++ ++def test_unset_toggles_default_to_the_shipped_value(): ++ got = _env_after_import({}) ++ assert {n: got[n] for n in NAMES} == dict.fromkeys(NAMES, "1") ++ assert got["QSA_SORT_ENVS"] is True ++ ++ ++@pytest.mark.parametrize("name", NAMES) ++def test_explicit_zero_is_kept(name): ++ got = _env_after_import({name: "0"}) ++ assert got[name] == "0" ++ others = [n for n in NAMES if n != name] ++ assert all(got[n] == "1" for n in others) +diff --git a/vllm/env_override.py b/vllm/env_override.py +index a6b3c37..fc15a40 100644 +--- a/vllm/env_override.py ++++ b/vllm/env_override.py +@@ -182,6 +182,22 @@ os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1") + # debug layout (see https://github.com/vllm-project/vllm/issues/41410). + os.environ.setdefault("TILELANG_CLEANUP_TEMP_FILES", "1") + ++# Fork (v2.5, knob-absent boot): the MERLIN_* kernel toggles are read by ++# getenv in C++ (FULL_K, TIE_DET, TIE_RECENT, TOPK_SORTED_EMIT) and by ++# vllm.envs (QSA_SORT) with OFF as the code default, while every shipped ++# launcher sets them to 1. Unset now means the shipped value; an explicit "0" ++# still turns a toggle off. Set here because this module runs on `import vllm`, ++# before any reader, and worker processes inherit the environment. ++_SHIPPED_KERNEL_TOGGLE_DEFAULTS = ( ++ "MERLIN_FULL_K", ++ "MERLIN_QSA_SORT", ++ "MERLIN_TIE_DET", ++ "MERLIN_TIE_RECENT", ++ "MERLIN_TOPK_SORTED_EMIT", ++) ++for _name in _SHIPPED_KERNEL_TOGGLE_DEFAULTS: ++ os.environ.setdefault(_name, "1") ++ + # =================================================== + # torch 2.9 Inductor PythonWrapperCodegen monkeypatch + # =================================================== +diff --git a/vllm/envs.py b/vllm/envs.py +index 9fe769a..7b55618 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -515,7 +515,7 @@ if TYPE_CHECKING: + # --- determinism toggles ----------------------------------------- + # The three nondeterminism mechanisms (FULL_K, QSA_SORT, TIE_DET/TIE_RECENT). + # Every one of them is +- # OFF by default: promotion is the maintainer's decision. ++ # OFF by default in this file; vllm/env_override.py sets each to 1 (the shipped value) when unset. + # + # Why these are registered at all: validate_environ() only inspects names + # beginning with VLLM_, so a MERLIN_ name never trips the unknown-variable +diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +index b9e88b0..45ace70 100644 +--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py ++++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +@@ -45,8 +45,8 @@ _DECODE_BLOCK_N = 64 + # on the reference host and is why this version has no Python-visible branch on data. + # + # Resolved once at import so it is a launch constant: a cudagraph captures the +-# op sequence that was live at capture time. Default OFF -- promotion is a release +-# decision, which is the maintainer's. ++# op sequence that was live at capture time. Code default OFF; vllm/env_override.py ++# sets 1 (the shipped value) when unset. + # --------------------------------------------------------------------------- + _QSA_SORT = envs.MERLIN_QSA_SORT + +-- +2.43.0 + diff --git a/patches/0091-config-envs-pass-use_replayssm-to-the-breakable-grap.patch b/patches/0091-config-envs-pass-use_replayssm-to-the-breakable-grap.patch new file mode 100644 index 0000000..f83bfb5 --- /dev/null +++ b/patches/0091-config-envs-pass-use_replayssm-to-the-breakable-grap.patch @@ -0,0 +1,219 @@ +From 2b39cee8bb05379b3186073a4562ce22fbdcfaf3 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 5 Oct 2026 00:13:13 -0400 +Subject: [PATCH 091/103] config, envs: pass use_replayssm to the + breakable-graph fallback; register VLLM_FAIL_ON_PROFILING_KV_LEAK + +CompilationConfig.resolve_cudagraph_mode_and_sizes read self.cache_config.use_replayssm, but CompilationConfig has +no cache_config, so with VLLM_USE_BREAKABLE_CUDAGRAPH=1 the fallback branch raised AttributeError. Both model +runners now pass cache_config.use_replayssm as an argument. The shipped launchers set VLLM_USE_BREAKABLE_CUDAGRAPH=0 +and never reach this branch. + +VLLM_FAIL_ON_PROFILING_KV_LEAK was read with a raw os.getenv, so setting it printed the unknown-variable warning. It +is now registered in vllm.envs with the same spellings (unset fails the boot; 0, false or no warn instead), listed +as an ignored compile factor and read through vllm.envs. Tests cover both. +--- + tests/config/test_recoverssm_config.py | 27 ++++++++++++++++++++++++++ + tests/test_envs.py | 19 ++++++++++++++++++ + vllm/config/compilation.py | 4 +++- + vllm/envs.py | 13 +++++++++++++ + vllm/v1/worker/gpu/cudagraph_utils.py | 6 ++---- + vllm/v1/worker/gpu/model_runner.py | 1 + + vllm/v1/worker/gpu_model_runner.py | 1 + + 7 files changed, 66 insertions(+), 5 deletions(-) + +diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py +index 352b8bf..50d5aa0 100644 +--- a/tests/config/test_recoverssm_config.py ++++ b/tests/config/test_recoverssm_config.py +@@ -197,5 +197,32 @@ def test_never_support_falls_back_to_piecewise_with_breakable_cudagraphs( + AttentionCGSupport.NEVER, + "GDNRecoverSSMAttentionBackend", + uniform_decode_query_len=4, ++ use_replayssm=True, + ) + assert resolved.name == expected ++ ++ ++def test_breakable_fallback_needs_replayssm_and_never_reads_cache_config(monkeypatch): ++ """The breakable-graph branch reads the RecoverSSM flag from its argument: ++ CompilationConfig has no cache_config, and reading one raised AttributeError ++ whenever VLLM_USE_BREAKABLE_CUDAGRAPH=1. Without RecoverSSM the branch is not ++ taken, so the unsupported FULL decode mode falls back to NONE.""" ++ from vllm.config.compilation import ( ++ CompilationConfig, ++ CompilationMode, ++ CUDAGraphMode, ++ ) ++ ++ monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") ++ compilation_config = CompilationConfig( ++ mode=CompilationMode.NONE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ splitting_ops=[], ++ ) ++ assert not hasattr(compilation_config, "cache_config") ++ resolved = compilation_config.resolve_cudagraph_mode_and_sizes( ++ AttentionCGSupport.NEVER, ++ "GDNRecoverSSMAttentionBackend", ++ uniform_decode_query_len=4, ++ ) ++ assert resolved.name == "NONE" +diff --git a/tests/test_envs.py b/tests/test_envs.py +index af5e39d..4871fda 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -655,6 +655,8 @@ _FORK_IGNORED = [ + # Mamba prefix-checkpoint retention (v2.5): scheduler-side, default ON; "0" + # turns it off. Read in single_type_kv_cache_manager.py (MambaManager). + ("VLLM_E2_MAMBA_RETAIN_CHECKPOINT", "0", False), ++ # Profiling-teardown leak policy: fails the boot unless "0"/"false"/"no". ++ ("VLLM_FAIL_ON_PROFILING_KV_LEAK", "no", False), + # Medium-path counter reader: a host-side log cadence, and the older slot-count + # fallback that this build never consults. Registered and ignored; before + # the read-site census neither classification was asserted anywhere. +@@ -1070,6 +1072,7 @@ _MUST_BE_READ = { + "VLLM_E2_MAMBA_SLOT_OWNERSHIP": "vllm/v1/worker/mamba_utils.py", + "VLLM_E2_MAMBA_RETAIN_CHECKPOINT": "vllm/v1/core/single_type_kv_cache_manager.py", + "VLLM_E2_SIDE_CACHE_LAYOUT": "vllm/v1/core/kv_cache_utils.py", ++ "VLLM_FAIL_ON_PROFILING_KV_LEAK": "vllm/v1/worker/gpu/cudagraph_utils.py", + } + _INERT_NO_READ_SITE = ["VLLM_PLE_FP8_GLOBAL_SCALE"] + +@@ -1151,3 +1154,19 @@ def test_side_cache_layout_switch_defaults_on_and_is_strict( + monkeypatch.setenv(name, bad) + with pytest.raises(ValueError, match=name): + _fresh_env_value(name) ++ ++ ++def test_fail_on_profiling_kv_leak_is_registered_and_defaults_to_fail( ++ monkeypatch: pytest.MonkeyPatch, ++): ++ """Registered (no unknown-variable warning), unset fails the boot, and the three ++ documented spellings downgrade it to a warning.""" ++ name = "VLLM_FAIL_ON_PROFILING_KV_LEAK" ++ assert name in envs.environment_variables ++ monkeypatch.delenv(name, raising=False) ++ assert _fresh_env_value(name) is True ++ for off in ("0", "false", "No"): ++ monkeypatch.setenv(name, off) ++ assert _fresh_env_value(name) is False ++ monkeypatch.setenv(name, "1") ++ assert _fresh_env_value(name) is True +diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py +index bd63f99..8538c17 100644 +--- a/vllm/config/compilation.py ++++ b/vllm/config/compilation.py +@@ -1401,6 +1401,7 @@ class CompilationConfig: + max_num_reqs: int | None = None, + is_profiling: bool = False, + piecewise_capture_available: bool = True, ++ use_replayssm: bool = False, + ) -> CUDAGraphMode: + from vllm.v1.attention.backend import AttentionCGSupport + +@@ -1454,7 +1455,8 @@ class CompilationConfig: + ) + ) or ( + envs.VLLM_USE_BREAKABLE_CUDAGRAPH +- and self.cache_config.use_replayssm ++ # CompilationConfig has no cache_config; the runner passes the flag. ++ and use_replayssm + ): + # Breakable CUDA graphs break at the attention kernels at runtime, + # with or without torch.compile. +diff --git a/vllm/envs.py b/vllm/envs.py +index 7b55618..4a40974 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -465,6 +465,11 @@ if TYPE_CHECKING: + # it changes the ring views' block stride and which block table the PLE + # state is read through, so ON and OFF must not share a compile cache. + VLLM_E2_SIDE_CACHE_LAYOUT: bool = True ++ # Profiling-teardown leak policy (vllm/v1/worker/gpu/cudagraph_utils.py): a profiling ++ # KV cache still allocated after teardown fails the boot by default; "0", "false" or ++ # "no" downgrade it to a warning while diagnosing. AOT: no. CAPTURE: no -- it only ++ # decides whether a host-side check raises. ++ VLLM_FAIL_ON_PROFILING_KV_LEAK: bool = True + # Opt-out for the fork startup guards in vllm/v1/e2_guards.py, read at + # Qwen4ExpForCausalLM construction. AOT: no. CAPTURE: no -- it changes + # nothing itself; it only permits a boot whose identity IS different +@@ -2633,6 +2638,11 @@ environment_variables: dict[str, Callable[[], Any]] = { + "VLLM_E2_SIDE_CACHE_LAYOUT": lambda: _e2_strict_switch( + "VLLM_E2_SIDE_CACHE_LAYOUT", default=True + ), ++ # Profiling-teardown leak policy: fail the boot unless "0"/"false"/"no". ++ "VLLM_FAIL_ON_PROFILING_KV_LEAK": lambda: os.getenv( ++ "VLLM_FAIL_ON_PROFILING_KV_LEAK", "1" ++ ).strip().lower() ++ not in ("0", "false", "no"), + # Seconds between `topk-exact-counters` lines from each GPU worker + # (0, the default, disables the reader entirely). The line reports the + # medium-path exact-restart counters; see +@@ -2900,6 +2910,9 @@ def compile_factors() -> dict[str, object]: + # compiled and the worker's tables are identical in both worlds, so an + # A/B flip must reuse the same AOT cache. + "VLLM_E2_MAMBA_RETAIN_CHECKPOINT", ++ # Boot-failure policy of a host-side profiling-teardown check; read by no ++ # compiled code. ++ "VLLM_FAIL_ON_PROFILING_KV_LEAK", + # Logging cadence for an already-compiled kernel's counters: it gates a + # host-side timer and a snapshot op, and is read by no compiled code, so + # hashing it would invalidate every compile cache for a log interval. +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 00eee6b..63b9811 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -3,7 +3,6 @@ + import contextlib + import gc + import itertools +-import os + from collections import defaultdict + from collections.abc import Callable, Iterable, Iterator + from contextlib import ExitStack +@@ -15,6 +14,7 @@ import torch + import torch.nn as nn + from tqdm import tqdm + ++import vllm.envs as envs + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +@@ -1196,8 +1196,6 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: + # Fail boot by default: this memory is absent from the real KV + # pool and can invalidate capacity planning. Operators may opt out + # temporarily while diagnosing with VLLM_FAIL_ON_PROFILING_KV_LEAK=0. +- if os.getenv("VLLM_FAIL_ON_PROFILING_KV_LEAK", "1").lower() not in ( +- "0", "false", "no" +- ): ++ if envs.VLLM_FAIL_ON_PROFILING_KV_LEAK: + raise RuntimeError(message) + logger.warning(message) +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index d2bb002..bc0e37c 100644 +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -787,6 +787,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + max_num_reqs=self.max_num_reqs, + is_profiling=is_profiling, + piecewise_capture_available=piecewise_capture_available, ++ use_replayssm=self.cache_config.use_replayssm, + ) + self.cudagraph_manager = ModelCudaGraphManager( + self.vllm_config, +diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py +index 010f110..a9dc802 100644 +--- a/vllm/v1/worker/gpu_model_runner.py ++++ b/vllm/v1/worker/gpu_model_runner.py +@@ -7221,6 +7221,7 @@ class GPUModelRunner( + kv_cache_config=self.kv_cache_config, + max_num_reqs=self.max_num_reqs, + is_profiling=is_profiling, ++ use_replayssm=self.cache_config.use_replayssm, + ) + # Trigger cudagraph dispatching keys initialization after + # resolved cudagraph mode. +-- +2.43.0 + diff --git a/patches/0092-gdn-keep-the-fused-packed-forward-under-RecoverSSM.patch b/patches/0092-gdn-keep-the-fused-packed-forward-under-RecoverSSM.patch new file mode 100644 index 0000000..8262833 --- /dev/null +++ b/patches/0092-gdn-keep-the-fused-packed-forward-under-RecoverSSM.patch @@ -0,0 +1,153 @@ +From 8acf03d55b9ad23e243a42cf337740d37226c084 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Mon, 5 Oct 2026 06:36:02 -0400 +Subject: [PATCH 092/103] gdn: keep the fused packed forward under RecoverSSM + +RecoverSSM's verify needs the Triton decode kernel for its speculative rows, and the previous code met that by turning +off the fused CUDA decode for the whole layer. That also turned off the fused packed forward, which carries every +prefill chunk, so prefill ran the unfused path. The decode kernel stays Triton and the packed recurrent decode stays +off; the fused packed forward stays on. A small helper makes the choice, with CPU tests for each combination. +--- + tests/kernels/mamba/test_gdn_decode_paths.py | 73 +++++++++++++++++++ + .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 36 +++++++-- + 2 files changed, 103 insertions(+), 6 deletions(-) + create mode 100644 tests/kernels/mamba/test_gdn_decode_paths.py + +diff --git a/tests/kernels/mamba/test_gdn_decode_paths.py b/tests/kernels/mamba/test_gdn_decode_paths.py +new file mode 100644 +index 0000000..a442ef6 +--- /dev/null ++++ b/tests/kernels/mamba/test_gdn_decode_paths.py +@@ -0,0 +1,73 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""GDN decode-path selection with and without RecoverSSM (CPU only). ++ ++RecoverSSM needs the Triton decode kernel for its spec rows, but must not turn ++off the fused packed forward, which carries prefill and verify alike. ++""" ++ ++import types ++from typing import cast ++from unittest.mock import MagicMock, patch ++ ++import pytest ++import torch ++ ++from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as mod ++from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( ++ QwenGatedDeltaNetAttention, ++) ++from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata ++ ++PREFIX = "model.layers.0.linear_attn" ++ ++ ++@pytest.mark.parametrize( ++ "kernel,recoverssm,packed,expected", ++ [ ++ pytest.param("cuda", True, True, ("triton", True, False), id="recoverssm-cuda"), ++ pytest.param("triton", True, True, ("triton", False, False), id="recoverssm-triton"), ++ pytest.param("cuda", False, True, ("cuda", True, True), id="cuda"), ++ pytest.param("triton", False, False, ("triton", False, False), id="triton"), ++ pytest.param("CPU", False, False, ("CPU", False, False), id="cpu"), ++ ], ++) ++def test_decode_paths(kernel, recoverssm, packed, expected) -> None: ++ assert QwenGatedDeltaNetAttention._decode_paths(kernel, recoverssm, packed) == expected ++ ++ ++def test_recoverssm_keeps_fused_forward_off_the_fused_mtp_kernel() -> None: ++ """With the Triton decode kernel (RecoverSSM), a spec-decode batch on the fused ++ packed forward goes through _forward_core and the gated norm, never through ++ the fused CUDA MTP decode.""" ++ layer = types.SimpleNamespace( ++ prefix=PREFIX, ++ gdn_decode_kernel="triton", ++ num_k_heads=2, ++ num_v_heads=4, ++ kv_cache=(None, torch.empty(1, dtype=torch.float32)), ++ ) ++ for name in ("_forward_core_fused_norm", "_can_use_fused_gdn_mtp_decode"): ++ setattr(layer, name, types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer)) ++ layer._forward_core = MagicMock() ++ layer._rms_norm_gated_cuda = MagicMock() ++ layer._forward_core_decode_spec_fused_norm = MagicMock() ++ layer._warmup_prefill_kernels = MagicMock() ++ meta = MagicMock(spec=GDNAttentionMetadata) ++ meta.spec_sequence_masks = object() ++ meta.num_decodes = 0 ++ meta.num_spec_decodes = 2 ++ meta.num_prefills = 0 ++ meta.num_actual_tokens = 8 ++ meta.spec_state_indices_tensor = torch.zeros(2, 4, dtype=torch.int32) ++ ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta}) ++ t = torch.zeros(8, 4, 2) ++ with patch.object(mod, "get_forward_context", return_value=ctx): ++ QwenGatedDeltaNetAttention._forward_core_fused_norm( ++ cast(QwenGatedDeltaNetAttention, layer), ++ mixed_qkv=torch.zeros(8, 16), b=torch.zeros(8, 4), a=torch.zeros(8, 4), ++ output_gate=t, core_attn_out=t, ++ ) ++ layer._forward_core_decode_spec_fused_norm.assert_not_called() ++ layer._forward_core.assert_called_once() ++ layer._rms_norm_gated_cuda.assert_called_once() +diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +index 80b1344..cc7ba1c 100644 +--- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py ++++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +@@ -555,14 +555,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + elif current_platform.is_cpu(): + self.gdn_decode_kernel = "CPU" + +- self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" +- # RecoverSSM verify replaces the per-draft snapshot path (Triton decode +- # only) + self.use_gdn_recoverssm = self._uses_gdn_recoverssm() ++ ( ++ self.gdn_decode_kernel, ++ self.enable_fused_gdn_decode, ++ self.enable_packed_recurrent_decode, ++ ) = self._decode_paths( ++ self.gdn_decode_kernel, ++ self.use_gdn_recoverssm, ++ self.enable_packed_recurrent_decode, ++ ) + if self.use_gdn_recoverssm: +- self.gdn_decode_kernel = "triton" +- self.enable_fused_gdn_decode = False +- self.enable_packed_recurrent_decode = False + logger.info_once( + "GDN RecoverSSM speculative verify active (spec_query_len %d)", + self.num_spec + 1, +@@ -574,6 +577,27 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + ++ @staticmethod ++ def _decode_paths( ++ gdn_decode_kernel: str, ++ use_recoverssm: bool, ++ packed_recurrent_decode: bool, ++ ) -> tuple[str, bool, bool]: ++ """(decode kernel, fused packed forward, packed recurrent decode). ++ ++ The fused packed forward is on whenever the CUDA decode kernel is ++ eligible. RecoverSSM's verify replaces the per-draft snapshot path and ++ needs the Triton decode kernel for its spec rows, so with RecoverSSM the ++ decode kernel is Triton and the packed recurrent decode is off. The ++ fused packed forward stays on: it calls _forward_core for prefill and ++ verify rows alike and only fuses the input split and the gated norm. ++ Turning it off sends every prefill chunk down the unfused path. ++ """ ++ enable_fused_forward = gdn_decode_kernel == "cuda" ++ if use_recoverssm: ++ return "triton", enable_fused_forward, False ++ return gdn_decode_kernel, enable_fused_forward, packed_recurrent_decode ++ + def _fused_gdn_decode_unsupported_reason( + self, vllm_config: VllmConfig + ) -> str | None: +-- +2.43.0 + diff --git a/patches/0093-mamba-hold-the-partial-tail-CoW-checkpoint-until-the.patch b/patches/0093-mamba-hold-the-partial-tail-CoW-checkpoint-until-the.patch new file mode 100644 index 0000000..2bb3691 --- /dev/null +++ b/patches/0093-mamba-hold-the-partial-tail-CoW-checkpoint-until-the.patch @@ -0,0 +1,783 @@ +From 2c86f0cf169ab1274238260bca55a7dabfffd7d0 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 6 Oct 2026 02:49:51 -0400 +Subject: [PATCH 093/103] mamba: hold the partial-tail CoW checkpoint until the + request is freed + +With --prefix-match-unit below the block size, a follow-up turn that +resumes at a partial tail (mid-block) never hashes a block-boundary GDN +state. Under retention interval 0 the only boundary kept is the replay +boundary, which lies below the resume point. The request's own tail is +published on its running block and moved into a CoW block by the next +chunk, and the scheduler releases both copy endpoints once the copy has +run. That block then sat at ref 0 in the free queue for the rest of the +request, ahead of its attention blocks, and concurrent prefills popped +it. The hybrid hit is the shortest group's, so the next turn resumed at +0. Under a long-context multi-turn agent load, 7 of 61 follow-ups at +unit 64 (6 at 256/512/1024) came back with cached_tokens 0 on 80-100K +prompts whose whole attention prefix was still cached. Every one +followed a partial-tail resume, none followed a block-aligned one, and +unit 4096 never missed. + +The existing retention (VLLM_E2_MAMBA_RETAIN_CHECKPOINT) only held +hashed blocks retired from the table, so it never covered that copy. +Now, under the same switch: + +- The running-request CoW branch holds the moved tail block in + _retained_checkpoints with a ref of its own; the copy retention pair + is unchanged. The new-request branch holds the hit source after + _apply_cow for the same reason: a request that stops before its own + tail still leaves the state it resumed from. +- _retained_checkpoints is keyed by block_hash_num_tokens instead of + block index, so a sub-block tail key orders correctly against full + block states; _retire_state_blocks uses the same key. The newest + position supersedes and releases the older one, so a request still + holds at most one block per group. +- When a producer moves its tail hashes into its CoW block, another + request holding that source (a sibling that resumed at the tail while + it was the producer's live block) follows the hashes to the copy. +- Admission reserves one extra block per Mamba group for a partial-tail + hit: the held source and the moved tail coexist at the tail step. + +The hold is released on every free path (finish, abort, preemption, +deferred pop_blocks_for_free) by the coordinator, last, after the +request's attention blocks, as before. When the hold is active (switch +on, hash unit below the Mamba block) MambaManager logs one startup line +saying so, with the hash unit and block size. + +Tests (CPU): a replay of the miss that fails on v2.5.0 with the +retention switch on (next-turn hit 0 instead of the tail) and passes +here; the structural variant; admission accounting; no-leak checks for +a request freed with its copy still queued, in both release orders; and +the moved-hash follow. The retention suite's one block-index assertion +now reads the token position. +--- + .../test_partial_prefix_cache_hits.py | 499 ++++++++++++++++++ + .../core/test_mamba_checkpoint_retention.py | 5 +- + vllm/v1/core/kv_cache_coordinator.py | 5 +- + vllm/v1/core/single_type_kv_cache_manager.py | 111 +++- + 4 files changed, 606 insertions(+), 14 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index 4798e0c..ee876b0 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -2532,3 +2532,502 @@ def test_boundary_states_offered_past_prompt_for_resumed_prefill(): + offered = [b for _, _, b in drain_boundary_state_offloads(manager).get("0", [])] + assert 2 * block_size in offered + assert req0.num_prompt_tokens < 2 * block_size ++ ++ ++# --------------------------------------------------------------------------- ++# A partial-tail resume must leave its own tail checkpoint referenced ++# (the tail-CoW total miss, measured at --prefix-match-unit 64/256/512/1024). ++# --------------------------------------------------------------------------- ++TAIL_BLOCK = 32 # mamba block ++TAIL_HASH = 8 # prefix-match unit (full attention hashes at the unit) ++TAIL_SPEC = 3 # MTP K=3, as deployed ++TAIL_SWITCH = "VLLM_E2_MAMBA_RETAIN_CHECKPOINT" ++TAIL_ATTN = 0 # group ids ++TAIL_MAMBA = 1 ++ ++ ++@pytest.fixture(params=["1", "0"], ids=["retain-on", "retain-off-control"]) ++def tail_retain(request, monkeypatch) -> bool: ++ from vllm import envs ++ ++ monkeypatch.setenv(TAIL_SWITCH, request.param) ++ if hasattr(envs.__getattr__, "cache_clear"): ++ envs.__getattr__.cache_clear() ++ return request.param == "1" ++ ++ ++def _tail_manager(num_blocks: int, use_eagle: bool = True): ++ """Full attention at the hash unit + one align-mode Mamba group at the ++ block, eagle/MTP drop on, retention interval 0 (the deployed default).""" ++ kv_cache_config = KVCacheConfig( ++ num_blocks=num_blocks, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ KVCacheGroupSpec( ++ ["full"], ++ FullAttentionSpec( ++ block_size=TAIL_HASH, ++ num_kv_heads=1, ++ head_size=1, ++ dtype=torch.float32, ++ ), ++ ), ++ KVCacheGroupSpec( ++ ["mamba"], ++ MambaSpec( ++ block_size=TAIL_BLOCK, ++ shapes=(1, 1), ++ dtypes=(torch.float32,), ++ mamba_cache_mode="align", ++ num_speculative_blocks=TAIL_SPEC, ++ ), ++ ), ++ ], ++ ) ++ manager = make_kv_cache_manager( ++ kv_cache_config=kv_cache_config, ++ max_model_len=8192, ++ enable_caching=True, ++ hash_block_size=TAIL_HASH, ++ use_eagle=use_eagle, ++ retention_interval=0, ++ ) ++ assert manager.coordinator.enable_partial_hash_hits ++ return manager ++ ++ ++def _tail_split( ++ request, num_new: int, num_local_hit: int = 0, use_eagle: bool = True ++) -> int: ++ """The real ``Scheduler._mamba_block_aligned_split`` on a stub self.""" ++ stub = SimpleNamespace( ++ block_size=TAIL_BLOCK, ++ cache_config=SimpleNamespace(block_size=TAIL_BLOCK), ++ use_eagle_block_drop=use_eagle, ++ max_num_scheduled_tokens=8192, ++ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), ++ mamba_partial_cache_hit=True, ++ mamba_fine_grained_prefix_cache=False, ++ hash_block_size=TAIL_HASH, ++ mamba_has_prefill_checkpoint_blocks=False, ++ mamba_prefill_checkpoint_alignment=None, ++ ) ++ return Scheduler._mamba_block_aligned_split(stub, request, num_new, num_local_hit) ++ ++ ++def _tail_drain_copies(manager) -> None: ++ """What the scheduler does once a step's CoW copies have run: release ++ both endpoints' copy retentions.""" ++ _, retained = manager.take_kv_cache_block_copies() ++ manager.block_pool.free_blocks(retained) ++ ++ ++def _tail_prefill(manager, request) -> int: ++ """Prefill to the prompt end in the chunks the scheduler would cut ++ (block ends, the partial-tail stop); returns the admitted hit length.""" ++ computed_blocks, num_hit, _ = manager.get_computed_blocks(request) ++ first = True ++ while request.num_computed_tokens < request.num_tokens: ++ start = request.num_computed_tokens + (num_hit if first else 0) ++ num_new = _tail_split( ++ request, request.num_tokens - start, num_hit if first else 0 ++ ) ++ assert num_new > 0 ++ blocks = manager.allocate_slots( ++ request, ++ num_new, ++ num_new_computed_tokens=num_hit if first else 0, ++ new_computed_blocks=computed_blocks if first else None, ++ num_lookahead_tokens=TAIL_SPEC, ++ ) ++ assert blocks is not None, "prefill chunk did not fit" ++ request.num_computed_tokens = start + num_new ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ first = False ++ return num_hit ++ ++ ++def _tail_decode(manager, request, steps: int) -> None: ++ """MTP K=3 decode steps, no drafts accepted.""" ++ next_tok = 50000 + len(request.output_token_ids) ++ if request.num_computed_tokens == request.num_tokens: ++ request.append_output_token_ids(next_tok) ++ next_tok += 1 ++ for _ in range(steps): ++ assert request.num_computed_tokens == request.num_tokens - 1 ++ request.spec_token_ids = [next_tok + i for i in range(TAIL_SPEC)] ++ blocks = manager.allocate_slots( ++ request, 1 + TAIL_SPEC, num_lookahead_tokens=TAIL_SPEC ++ ) ++ assert blocks is not None, "decode step did not fit" ++ request.spec_token_ids = [] ++ request.num_computed_tokens += 1 ++ request.append_output_token_ids([next_tok]) ++ next_tok += 1 ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ ++ ++def _tail_tokens(seed: int, n: int) -> list[int]: ++ return [seed * 10000 + i for i in range(n)] ++ ++ ++def _tail_position(num_prompt_tokens: int) -> int: ++ """Where the prompt's partial-tail state is keyed (eagle shifts it one ++ unit below the last hash boundary).""" ++ return num_prompt_tokens // TAIL_HASH * TAIL_HASH - TAIL_HASH ++ ++ ++def _tail_followup(prev, user_tokens: int, seed: int): ++ """The next turn: the previous turn's prompt + its output + new text, ++ sized so that its tail lands strictly inside a block above the last ++ cacheable block end (the deployed case; a block-aligned tail is a ++ full-block key, not a partial one).""" ++ while True: ++ prompt = len(prev.all_token_ids) + user_tokens ++ tail = _tail_position(prompt) ++ last_block_end = prompt - prompt % TAIL_BLOCK - TAIL_BLOCK ++ if tail % TAIL_BLOCK != 0 and last_block_end < tail: ++ break ++ user_tokens += 1 ++ return make_request( ++ prev.request_id + "+", ++ list(prev.all_token_ids) + _tail_tokens(seed, user_tokens), ++ TAIL_HASH, ++ sha256, ++ ) ++ ++ ++def _tail_mamba(manager): ++ return manager.coordinator.single_type_managers[TAIL_MAMBA] ++ ++ ++def _tail_free_ids(manager) -> set[int]: ++ return {b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()} ++ ++ ++def _tail_assert_pool_clean(manager, num_blocks: int) -> None: ++ pool = manager.block_pool ++ assert pool.get_num_free_blocks() == num_blocks - 1 # all but the null block ++ assert all(blk.ref_cnt == 0 for blk in pool.blocks[1:]) ++ mamba = _tail_mamba(manager) ++ assert not mamba._retained_checkpoints ++ assert not mamba._partial_hit_reqs ++ assert not mamba.req_to_blocks ++ ++ ++def test_partial_tail_resume_next_turn_hits_after_concurrent_allocation(tail_retain): ++ """Behaviour only (the observed miss): A1 resumed at A0's partial ++ tail; while A1 decodes, concurrent prefills pop every free block; A1 ++ finishes; A2 must resume at A1's tail. Before the fix A2's hit is 0 -- ++ a total miss on a prompt whose whole attention prefix is still cached.""" ++ manager = _tail_manager(96) ++ pool = manager.block_pool ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_prefill(manager, a0) ++ _tail_decode(manager, a0, steps=4) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ ++ a1 = _tail_followup(a0, 26, seed=2) ++ num_hit = _tail_prefill(manager, a1) ++ assert num_hit == _tail_position(a0.num_prompt_tokens) # a partial tail ++ assert num_hit % TAIL_BLOCK != 0 ++ _tail_decode(manager, a1, steps=4) ++ pool.free_blocks(pool.get_new_blocks(pool.get_num_free_blocks())) # burst ++ manager.free(a1) ++ _tail_drain_copies(manager) ++ ++ a2 = _tail_followup(a1, 20, seed=3) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ assert all( ++ pool.get_cached_block(a1.block_hashes[i], [TAIL_ATTN]) ++ for i in range(a1.num_prompt_tokens // TAIL_HASH) ++ ) ++ assert num_hit == (_tail_position(a1.num_prompt_tokens) if tail_retain else 0) ++ ++ ++def test_partial_tail_resume_keeps_its_own_tail_under_concurrent_allocation( ++ tail_retain, ++): ++ """Turn 0 (A0) is a cold prefill; turn 1 (A1) resumes at A0's partial tail ++ (mid-block). A1 hashes no boundary state -- its replay boundary is below ++ its resume point -- so its own tail, moved into a CoW block by the chunk ++ after the tail stop, is the only GDN state its next turn can hit. Before ++ the fix the scheduler's post-copy release left that block at ref 0 in the ++ free queue while A1 still decoded; a concurrent prefill popped it and the ++ next turn (A2) got a TOTAL miss although every attention block of A1 was ++ still cached. With retention on, A1 holds the moved tail until it is ++ freed and A2 resumes at it. ``retain-off-control`` is the old behaviour ++ (and v2.5.0 fails the retain-on case the same way). ++ """ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks) ++ pool = manager.block_pool ++ mamba = _tail_mamba(manager) ++ ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ assert _tail_prefill(manager, a0) == 0 ++ _tail_decode(manager, a0, steps=4) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ ++ # A1 resumes at A0's tail: a sub-block hit, so its first chunk copies the ++ # hit state into a private block (CoW). ++ a1 = _tail_followup(a0, 26, seed=2) ++ a1_tail = _tail_position(a1.num_prompt_tokens) ++ num_hit = _tail_prefill(manager, a1) ++ assert num_hit == _tail_position(a0.num_prompt_tokens) ++ assert num_hit % TAIL_BLOCK != 0 ++ [tail_block] = pool.get_cached_block( ++ a1.block_hashes[a1_tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ # The tail state was moved out of A1's table into the CoW block, and A1 ++ # keyed no full-block boundary state of its own. ++ assert tail_block not in mamba.req_to_blocks.get(a1.request_id, []) ++ assert tail_block.block_hash_num_tokens == a1_tail ++ for position in range(TAIL_BLOCK, a1_tail, TAIL_BLOCK): ++ if position > num_hit: ++ assert ( ++ pool.get_cached_block( ++ a1.block_hashes[position // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ is None ++ ) ++ _tail_decode(manager, a1, steps=4) ++ ++ if tail_retain: ++ assert mamba._retained_checkpoints[a1.request_id] == (a1_tail, tail_block) ++ assert tail_block.ref_cnt == 1 ++ assert tail_block.block_id not in _tail_free_ids(manager) ++ else: ++ # The exposure: cached-free while A1 still runs. ++ assert tail_block.ref_cnt == 0 and tail_block.block_hash is not None ++ assert tail_block.block_id in _tail_free_ids(manager) ++ ++ # A burst of concurrent prefills pops every free block while A1 decodes. ++ taken = pool.get_new_blocks(pool.get_num_free_blocks()) ++ assert (tail_block.block_hash is not None) is tail_retain ++ pool.free_blocks(taken) ++ ++ # A1 finishes; its next turn arrives. ++ manager.free(a1) ++ _tail_drain_copies(manager) ++ a2 = _tail_followup(a1, 20, seed=3) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ # Every attention block of A1's prompt is still cached in both worlds. ++ assert all( ++ pool.get_cached_block(a1.block_hashes[i], [TAIL_ATTN]) ++ for i in range(a1.num_prompt_tokens // TAIL_HASH) ++ ) ++ if tail_retain: ++ assert num_hit == a1_tail ++ else: ++ # FAILING CONTROL: total miss with the whole attention prefix cached. ++ assert num_hit == 0 ++ ++ # The next turn runs to completion on what it got; nothing leaks. ++ _tail_prefill(manager, a2) ++ _tail_decode(manager, a2, steps=2) ++ manager.free(a2) ++ _tail_drain_copies(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_partial_tail_hit_admission_reserves_the_tail_swap(tail_retain): ++ """Full-sequence admission of a partial-tail follow-up reserves one block ++ per Mamba group above the per-step estimate (the held hit source and the ++ moved tail coexist for the swap); a block-aligned follow-up reserves none. ++ """ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks) ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_prefill(manager, a0) ++ _tail_decode(manager, a0, steps=2) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ ++ a1 = _tail_followup(a0, 26, seed=2) ++ computed, num_hit, _ = manager.get_computed_blocks(a1) ++ assert num_hit % TAIL_BLOCK != 0 ++ ++ def need(apply_admission_cap: bool) -> int: ++ return manager.coordinator.get_num_blocks_to_allocate( ++ request_id=a1.request_id, ++ num_tokens=a1.num_tokens, ++ new_computed_blocks=computed.blocks, ++ num_encoder_tokens=0, ++ total_computed_tokens=num_hit, ++ num_local_computed_tokens=num_hit, ++ num_tokens_main_model=a1.num_tokens, ++ apply_admission_cap=apply_admission_cap, ++ ) ++ ++ assert need(True) == need(False) + (1 if tail_retain else 0) ++ ++ ++@pytest.mark.parametrize("deferred", [False, True], ids=["free", "pop_blocks_for_free"]) ++@pytest.mark.parametrize("stop_at", ["after-hit-copy", "after-tail-copy"]) ++def test_partial_tail_resume_freed_with_queued_copy_leaks_nothing( ++ tail_retain, deferred, stop_at ++): ++ """A request that resumed at a partial tail is freed (finish, abort or ++ preemption) while its CoW copy is still queued -- before the scheduler ++ has released the copy retentions. Whatever order the retentions and the ++ request's hold are released in, every block ends at ref 0 and the pool ++ returns to baseline; the superseded hit source ages instead of leaking. ++ """ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks) ++ pool = manager.block_pool ++ mamba = _tail_mamba(manager) ++ ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_prefill(manager, a0) ++ _tail_decode(manager, a0, steps=2) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ ++ a1 = _tail_followup(a0, 26, seed=2) ++ a1_tail = _tail_position(a1.num_prompt_tokens) ++ computed_blocks, num_hit, _ = manager.get_computed_blocks(a1) ++ assert 0 < num_hit < a1_tail and num_hit % TAIL_BLOCK != 0 ++ [source] = pool.get_cached_block( ++ a0.block_hashes[num_hit // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ assert source.ref_cnt == 0 ++ ++ # First chunk: the hit source is copied into A1's private block. ++ num_new = _tail_split(a1, a1.num_tokens - num_hit, num_hit) ++ assert manager.allocate_slots( ++ a1, num_new, num_hit, computed_blocks, num_lookahead_tokens=TAIL_SPEC ++ ) ++ a1.num_computed_tokens = num_hit + num_new ++ manager.new_step_starts() ++ if tail_retain: ++ assert mamba._retained_checkpoints[a1.request_id] == (num_hit, source) ++ if stop_at == "after-tail-copy": ++ _tail_drain_copies(manager) ++ # Run to the chunk that moves A1's own tail into a CoW block. ++ while a1.num_computed_tokens <= a1_tail: ++ num_new = _tail_split(a1, a1.num_tokens - a1.num_computed_tokens) ++ assert manager.allocate_slots(a1, num_new, num_lookahead_tokens=TAIL_SPEC) ++ a1.num_computed_tokens += num_new ++ manager.new_step_starts() ++ if a1.num_computed_tokens <= a1_tail: ++ _tail_drain_copies(manager) ++ [tail_block] = pool.get_cached_block( ++ a1.block_hashes[a1_tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ if tail_retain: ++ assert mamba._retained_checkpoints[a1.request_id] == (a1_tail, tail_block) ++ assert source.ref_cnt == 0 # superseded: ages, not leaked ++ assert source.block_hash is not None ++ ++ # Freed with the copy still queued. ++ if deferred: ++ pool.free_blocks(manager.pop_blocks_for_free(a1)) ++ else: ++ manager.free(a1) ++ assert a1.request_id not in mamba._retained_checkpoints ++ assert all(blk.ref_cnt >= 0 for blk in pool.blocks[1:]) ++ # The copy's retentions are released afterwards. ++ _tail_drain_copies(manager) ++ assert source.ref_cnt == 0 ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_consumer_hold_follows_the_producers_moved_tail(tail_retain): ++ """Without the eagle drop a sibling can resume at a producer's tail while ++ that tail is still the producer's LIVE running block (published at the ++ tail stop, hit before the producer's next chunk). The consumer's hold ++ then pins the producer's block; when the producer's next chunk moves the ++ hashes into its CoW block, the hold must follow them -- it pins the ++ hashed state, never another request's running block -- and both ++ requests must still free cleanly. ++ """ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks, use_eagle=False) ++ pool = manager.block_pool ++ mamba = _tail_mamba(manager) ++ ++ producer = make_request("P", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ tail = producer.num_prompt_tokens // TAIL_HASH * TAIL_HASH # 144, no shift ++ computed, num_hit, _ = manager.get_computed_blocks(producer) ++ while producer.num_computed_tokens < tail: ++ # The real split gets the whole remainder and stops at the tail. ++ num_new = _tail_split( ++ producer, producer.num_tokens - producer.num_computed_tokens, use_eagle=False ++ ) ++ assert producer.num_computed_tokens + num_new <= tail ++ assert manager.allocate_slots( ++ producer, ++ num_new, ++ num_new_computed_tokens=num_hit, ++ new_computed_blocks=computed, ++ num_lookahead_tokens=TAIL_SPEC, ++ ) ++ producer.num_computed_tokens += num_new ++ computed, num_hit = None, 0 ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ [live] = pool.get_cached_block( ++ producer.block_hashes[tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ assert live in mamba.req_to_blocks["P"] # the tail is keyed on the live block ++ ++ consumer = make_request( ++ "C", list(producer.prompt_token_ids) + _tail_tokens(2, 30), TAIL_HASH, sha256 ++ ) ++ c_computed, c_hit, _ = manager.get_computed_blocks(consumer) ++ assert c_hit == tail ++ num_new = _tail_split( ++ consumer, consumer.num_tokens - c_hit, c_hit, use_eagle=False ++ ) ++ assert manager.allocate_slots( ++ consumer, num_new, c_hit, c_computed, num_lookahead_tokens=TAIL_SPEC ++ ) ++ consumer.num_computed_tokens = c_hit + num_new ++ if tail_retain: ++ assert mamba._retained_checkpoints["C"] == (tail, live) ++ ++ # The producer's next chunk moves the tail hashes into its CoW block. ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ num_new = _tail_split( ++ producer, producer.num_tokens - producer.num_computed_tokens, use_eagle=False ++ ) ++ assert manager.allocate_slots(producer, num_new, num_lookahead_tokens=TAIL_SPEC) ++ producer.num_computed_tokens += num_new ++ [moved] = pool.get_cached_block( ++ producer.block_hashes[tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ assert moved is not live and live.block_hash is None ++ if tail_retain: ++ assert mamba._retained_checkpoints["C"] == (tail, moved) ++ assert mamba._retained_checkpoints["P"] == (tail, moved) ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ # The live block is referenced by the producer's table alone again. ++ assert live.ref_cnt == 1 ++ ++ _tail_decode(manager, producer, steps=2) ++ manager.free(producer) ++ _tail_drain_copies(manager) ++ while consumer.num_computed_tokens < consumer.num_tokens: ++ num_new = _tail_split( ++ consumer, ++ consumer.num_tokens - consumer.num_computed_tokens, ++ use_eagle=False, ++ ) ++ assert manager.allocate_slots( ++ consumer, num_new, num_lookahead_tokens=TAIL_SPEC ++ ) ++ consumer.num_computed_tokens += num_new ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ _tail_decode(manager, consumer, steps=2) ++ manager.free(consumer) ++ _tail_drain_copies(manager) ++ _tail_assert_pool_clean(manager, num_blocks) +diff --git a/tests/v1/core/test_mamba_checkpoint_retention.py b/tests/v1/core/test_mamba_checkpoint_retention.py +index 9d0d440..502e68b 100644 +--- a/tests/v1/core/test_mamba_checkpoint_retention.py ++++ b/tests/v1/core/test_mamba_checkpoint_retention.py +@@ -548,8 +548,9 @@ def test_mtp_chunked_prefill_aligned_prompt_retains_newest_boundary_only( + [hi] = pool.get_cached_block(hi_hash, [gid]) + assert _hashed_cached(manager, lo) # superseded / never held: ages + if retain: +- idx, held = _mamba(manager, gid)._retained_checkpoints["A"] +- assert idx == 2 and held is hi ++ # Keyed by the hashed token position, not block index. ++ position, held = _mamba(manager, gid)._retained_checkpoints["A"] ++ assert position == 3 * BLOCK and held is hi + assert hi.ref_cnt == 1 + else: + assert _hashed_cached(manager, hi) +diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py +index 7537ccd..b58aa96 100644 +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -248,7 +248,10 @@ class KVCacheCoordinator(ABC): + # the free queue that check reads. + num_blocks_to_allocate += ( + manager.get_num_retained_checkpoint_blocks_to_reserve( +- request_id, num_tokens_main_model, new_computed_blocks[i] ++ request_id, ++ num_tokens_main_model, ++ new_computed_blocks[i], ++ num_local_computed_tokens, + ) + ) + return num_blocks_to_allocate +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index ca390ef..649f0a4 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -611,6 +611,7 @@ class SingleTypeKVCacheManager(ABC): + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock] = (), ++ num_local_computed_tokens: int = 0, + ) -> int: + """Blocks the admission gate must add for checkpoints this request + will hold outside its block table (0 for every manager but align-mode +@@ -1673,7 +1674,24 @@ class MambaManager(SingleTypeKVCacheManager): + and self.enable_caching + and envs.VLLM_E2_MAMBA_RETAIN_CHECKPOINT + ) +- # {request_id: (logical block index, block)} ++ # {request_id: (hashed token position, block)}: the newest hashed ++ # state the request resumed from or produced, held with one ref of ++ # its own (ordered by ``block_hash_num_tokens`` rather than ++ # block index so a sub-block partial-tail key sorts correctly against ++ # full-block boundary states). ++ # ++ # With ``--prefix-match-unit`` below the block size a request ++ # that resumed at a partial tail (mid-block) never hashes a boundary ++ # state: the replay boundary lies below its resume point. Its only ++ # reusable GDN state is its own partial-tail checkpoint, which the ++ # next chunk moves into a CoW block in no block table; the scheduler ++ # drops the copy's retention refs once the copy has run, so that ++ # block sat at ref 0 in the free queue for the rest of the request, ++ # ahead of its attention blocks, and concurrent prefills evicted it ++ # (measured: 7/61 follow-ups with cached_tokens 0 at unit 64, every ++ # one after a partial-tail resume; none at unit 4096). The CoW paths ++ # therefore hold their hashed endpoint here too: the moved tail for ++ # a running request, and the hit source a new request copies from. + self._retained_checkpoints: dict[str, tuple[int, KVCacheBlock]] = {} + if self.mamba_cache_mode == "align" and self.enable_caching: + # The A/B evidence line: boot logs must show which arm ran. +@@ -1682,6 +1700,16 @@ class MambaManager(SingleTypeKVCacheManager): + "(VLLM_E2_MAMBA_RETAIN_CHECKPOINT)", + "on" if self._retain_checkpoint else "off", + ) ++ if ( ++ self._retain_checkpoint ++ and self.block_size > self.block_pool.hash_block_size ++ ): ++ logger.info_once( ++ "Mamba align: each request holds its partial-tail CoW checkpoint " ++ "until it is freed (hash unit %d, block %d)", ++ self.block_pool.hash_block_size, ++ self.block_size, ++ ) + + def _retire_state_blocks( + self, request_id: str, retired: list[tuple[int, KVCacheBlock]] +@@ -1696,13 +1724,16 @@ class MambaManager(SingleTypeKVCacheManager): + """ + if self._retain_checkpoint: + newest: tuple[int, KVCacheBlock] | None = None +- for idx, block in retired: +- if block.block_hash is not None and ( +- newest is None or idx > newest[0] +- ): +- newest = (idx, block) ++ for _, block in retired: ++ if block.block_hash is None: ++ continue ++ position = block.block_hash_num_tokens ++ assert position is not None ++ if newest is None or position > newest[0]: ++ newest = (position, block) + held = self._retained_checkpoints.get(request_id) + if newest is not None and (held is None or newest[0] > held[0]): ++ # The retired block's table ref becomes the hold's ref. + self._retained_checkpoints[request_id] = newest + retired = [entry for entry in retired if entry[1] is not newest[1]] + if held is not None: +@@ -1711,6 +1742,43 @@ class MambaManager(SingleTypeKVCacheManager): + if retired: + self.block_pool.free_blocks([block for _, block in retired]) + ++ def _hold_checkpoint(self, request_id: str, block: KVCacheBlock) -> bool: ++ """Hold ``block`` (a hashed state outside the request's table) as the ++ request's prefix checkpoint, taking one extra ref for the hold. ++ ++ Only the newest position is held: an older one is left alone (it ages ++ from wherever it is), a superseded one is released here. The hold's ++ ref is returned by ``pop_retained_checkpoints`` -> ``free_blocks``. ++ Returns whether the block is now held. ++ """ ++ if not self._retain_checkpoint or block.is_null or block.block_hash is None: ++ return False ++ position = block.block_hash_num_tokens ++ assert position is not None ++ held = self._retained_checkpoints.get(request_id) ++ if held is not None and (held[1] is block or held[0] >= position): ++ return False ++ self.block_pool.touch([block]) ++ self._retained_checkpoints[request_id] = (position, block) ++ if held is not None: ++ self.block_pool.free_blocks([held[1]]) ++ return True ++ ++ def _follow_moved_hash( ++ self, owner_id: str, source_block: KVCacheBlock, cow_block: KVCacheBlock ++ ) -> None: ++ """``owner_id`` just moved ``source_block``'s hashes to ``cow_block`` ++ and keeps writing ``source_block``. A request that resumed at that ++ tail while it was still the owner's live block holds the source for ++ those hashes; move its hold to the copy, so it pins the hashed state ++ and not another request's running block. ++ """ ++ for request_id, (position, block) in list(self._retained_checkpoints.items()): ++ if request_id != owner_id and block is source_block: ++ self.block_pool.touch([cow_block]) ++ self._retained_checkpoints[request_id] = (position, cow_block) ++ self.block_pool.free_blocks([source_block]) ++ + def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]: + if self.mamba_cache_mode != "align": + return [] +@@ -1722,16 +1790,25 @@ class MambaManager(SingleTypeKVCacheManager): + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock] = (), ++ num_local_computed_tokens: int = 0, + ) -> int: + # A request that already holds a checkpoint only ever swaps it. One + # that does not will hold one once a hashed state retires, which needs + # at least one state block behind the running one. +- if ( +- not self._retain_checkpoint +- or request_id in self._retained_checkpoints +- or num_tokens <= self.block_size +- ): ++ if not self._retain_checkpoint or request_id in self._retained_checkpoints: + return 0 ++ # A sub-block (partial-tail) hit holds the hit source only ++ # until the request's own tail is moved into a CoW block that is held ++ # in its place; at that step both are referenced, so the request's ++ # lifetime peak is one block above what the hit brought (whatever the ++ # prompt length: the tail CoW happens inside the first block too). ++ if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): ++ return 1 ++ if num_tokens <= self.block_size: ++ # No state block retires behind the running one, but a prompt ++ # whose tail is keyed inside the block still moves that tail into ++ # a held CoW block when the hash unit is finer than the block. ++ return int(self.block_pool.hash_block_size < self.block_size) + # A prefix hit already brings the block that becomes the held + # checkpoint -- the hit state, counted by get_num_blocks_to_allocate as + # an evictable computed block (or already referenced elsewhere) -- and +@@ -2166,6 +2243,7 @@ class MambaManager(SingleTypeKVCacheManager): + # overwrites source_block. + assert req_blocks[block_idx] is source_block + self.block_pool.move_block_hashes(source_block, cow_block) ++ self._follow_moved_hash(request_id, source_block, cow_block) + self._pending_cow_copies.append((source_block, cow_block)) + source_block.ref_cnt += 1 + producer_tail = self._producer_partial_tail_reqs.pop( +@@ -2189,9 +2267,20 @@ class MambaManager(SingleTypeKVCacheManager): + # The moved entry is only filled by this step's + # copy, so defer same-step hits on it. + self.cached_blocks_this_step.add(cow_block.block_hash) ++ # The moved tail is the request's newest ++ # reusable state and lives in no block table; hold it ++ # past the copy's retention so the request's next ++ # turn can resume at it (see __init__). ++ self._hold_checkpoint(request_id, cow_block) + else: + self._apply_cow(request_id, block_idx, source_block, cow_block) + returned_blocks = [cow_block] + returned_blocks ++ # The hit source keeps the hash while the ++ # request advances its private copy; hold it until ++ # the request's own tail supersedes it, so a request ++ # that stops early (preempted, aborted, block-aligned ++ # prompt) still leaves the state it resumed from. ++ self._hold_checkpoint(request_id, source_block) + req_blocks.extend(new_blocks) + self._allocated_block_reqs.add(request_id) + self._partial_hit_reqs.pop(request_id, None) +-- +2.43.0 + diff --git a/patches/0094-mamba-key-the-prompt-tail-when-it-lands-on-a-block-e.patch b/patches/0094-mamba-key-the-prompt-tail-when-it-lands-on-a-block-e.patch new file mode 100644 index 0000000..68db145 --- /dev/null +++ b/patches/0094-mamba-key-the-prompt-tail-when-it-lands-on-a-block-e.patch @@ -0,0 +1,247 @@ +From 9eead3f065e4db6894fcc324425b30ca739117ec Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 6 Oct 2026 02:49:51 -0400 +Subject: [PATCH 094/103] mamba: key the prompt tail when it lands on a block + end + +With a hash unit finer than the block, the prompt's tail key lands +exactly on a block end once every block/unit turns. The tail key is the +position the next turn resumes at: the last unit boundary, one unit +lower under the eagle drop. It is then a full-block key, so the +partial-tail registration does nothing, and under retention interval 0 +the mask kept only the replay boundary one block below it. That state +was never hashed. The next turn therefore resumed one block lower after +a cold turn, or one turn lower (at the held source) after a partial-tail +resume, re-prefilling the previous turn's extension. + +MambaManager.cache_blocks now adds that block-aligned tail to the +reachable boundaries, so the block ending there is hashed like any +boundary state, retired into the held checkpoint and freed last. At the +block unit the tail is the replay boundary itself and nothing changes. +When the hash unit is below the Mamba block, MambaManager logs one +startup line saying the block-aligned tail key is hashed. + +Tests: a cold and a partial-tail-resume case whose tail lands on a block +end, under a concurrent burst, with the no-leak check. Both fail on the +previous commit (the tail state is not keyed and the next turn resumes +lower). A block-unit case asserts the hashed set is exactly the replay +boundary. +--- + .../test_partial_prefix_cache_hits.py | 130 +++++++++++++++++- + vllm/v1/core/single_type_kv_cache_manager.py | 35 +++++ + 2 files changed, 161 insertions(+), 4 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index ee876b0..5247e08 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -2679,16 +2679,16 @@ def _tail_position(num_prompt_tokens: int) -> int: + return num_prompt_tokens // TAIL_HASH * TAIL_HASH - TAIL_HASH + + +-def _tail_followup(prev, user_tokens: int, seed: int): ++def _tail_followup(prev, user_tokens: int, seed: int, aligned_tail: bool = False): + """The next turn: the previous turn's prompt + its output + new text, + sized so that its tail lands strictly inside a block above the last +- cacheable block end (the deployed case; a block-aligned tail is a +- full-block key, not a partial one).""" ++ cacheable block end (the common case), or -- ``aligned_tail`` -- exactly ++ on a block end (a full-block key, not a partial one).""" + while True: + prompt = len(prev.all_token_ids) + user_tokens + tail = _tail_position(prompt) + last_block_end = prompt - prompt % TAIL_BLOCK - TAIL_BLOCK +- if tail % TAIL_BLOCK != 0 and last_block_end < tail: ++ if (tail % TAIL_BLOCK == 0) is aligned_tail and last_block_end < tail: + break + user_tokens += 1 + return make_request( +@@ -3031,3 +3031,125 @@ def test_consumer_hold_follows_the_producers_moved_tail(tail_retain): + manager.free(consumer) + _tail_drain_copies(manager) + _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++@pytest.mark.parametrize("resume", ["cold", "partial-tail"]) ++def test_block_aligned_tail_key_is_hashed_and_held(tail_retain, resume): ++ """The prompt's tail key can land exactly on a block end (1 in ++ block/unit turns). It is then a full-block key: `_cache_partial_tail_block` ++ registers nothing, and under retention interval 0 the mask kept only the ++ replay boundary one block below it, so that state was never hashed. The ++ next turn then resumed one block (cold) or one turn (partial-tail resume: ++ at the held source) lower. Now the tail block is hashed, retired into the ++ hold like any boundary state, survives a concurrent burst, and the next ++ turn resumes at the tail itself. At the block unit nothing changes: the ++ tail is the replay boundary. ++ """ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks) ++ pool = manager.block_pool ++ mamba = _tail_mamba(manager) ++ ++ if resume == "cold": ++ # 136 tokens: tail 128 = 4 blocks exactly; replay boundary 96. ++ a1 = make_request("A1", _tail_tokens(1, 136), TAIL_HASH, sha256) ++ below = 3 * TAIL_BLOCK ++ assert _tail_prefill(manager, a1) == 0 ++ else: ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_prefill(manager, a0) ++ _tail_decode(manager, a0, steps=4) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ a1 = _tail_followup(a0, 13, seed=2, aligned_tail=True) ++ below = _tail_position(a0.num_prompt_tokens) # the held source ++ assert _tail_prefill(manager, a1) == below ++ assert below % TAIL_BLOCK != 0 ++ tail = _tail_position(a1.num_prompt_tokens) ++ assert tail % TAIL_BLOCK == 0 and tail > below ++ _tail_decode(manager, a1, steps=4) ++ ++ # The tail state is keyed (without the block-end tail key the next turn ++ # resumed at `below` instead). ++ cached = pool.get_cached_block(a1.block_hashes[tail // TAIL_HASH - 1], [TAIL_MAMBA]) ++ assert cached is not None, f"tail state@{tail} is not keyed (block-end tail)" ++ [tail_block] = cached ++ assert tail_block.block_hash_num_tokens == tail ++ if tail_retain: ++ assert mamba._retained_checkpoints[a1.request_id] == (tail, tail_block) ++ assert tail_block.ref_cnt == 1 ++ else: ++ assert tail_block.ref_cnt == 0 ++ ++ # Concurrent prefills pop every free block while A1 decodes. ++ taken = pool.get_new_blocks(pool.get_num_free_blocks()) ++ assert (tail_block.block_hash is not None) is tail_retain ++ pool.free_blocks(taken) ++ ++ manager.free(a1) ++ _tail_drain_copies(manager) ++ a2 = _tail_followup(a1, 20, seed=3) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ assert num_hit == (tail if tail_retain else 0) ++ ++ _tail_prefill(manager, a2) ++ _tail_decode(manager, a2, steps=2) ++ manager.free(a2) ++ _tail_drain_copies(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_block_unit_replay_boundary_is_unchanged(): ++ """At the block unit (U = block) the tail helper is inert and the hashed ++ boundary set is exactly the replay boundary, as before.""" ++ hash_block_size = block_size = 32 ++ kv_cache_config = KVCacheConfig( ++ num_blocks=64, ++ kv_cache_tensors=[], ++ kv_cache_groups=[ ++ KVCacheGroupSpec( ++ ["full"], ++ FullAttentionSpec( ++ block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 ++ ), ++ ), ++ KVCacheGroupSpec( ++ ["mamba"], ++ MambaSpec( ++ block_size=block_size, ++ shapes=(1, 1), ++ dtypes=(torch.float32,), ++ mamba_cache_mode="align", ++ num_speculative_blocks=TAIL_SPEC, ++ ), ++ ), ++ ], ++ ) ++ manager = make_kv_cache_manager( ++ kv_cache_config=kv_cache_config, ++ max_model_len=8192, ++ enable_caching=True, ++ hash_block_size=hash_block_size, ++ use_eagle=True, ++ retention_interval=0, ++ ) ++ mamba = manager.coordinator.single_type_managers[TAIL_MAMBA] ++ # 136 tokens: tail key 96 (eagle) = the replay boundary; the scheduler ++ # cuts the prefill there (last cacheable block end), then the rest. ++ request = make_request("R", _tail_tokens(1, 136), hash_block_size, sha256) ++ [boundary] = set(manager.coordinator.get_replay_boundaries(request)) ++ assert boundary == 96 ++ if hasattr(mamba, "_block_aligned_prompt_tail"): ++ assert mamba._block_aligned_prompt_tail(request) == 0 ++ for chunk in (96, 40): ++ assert manager.allocate_slots(request, chunk, num_lookahead_tokens=TAIL_SPEC) ++ request.num_computed_tokens += chunk ++ manager.new_step_starts() ++ hashed = [ ++ i * block_size ++ for i in range(1, 136 // block_size + 1) ++ if manager.block_pool.get_cached_block( ++ request.block_hashes[i - 1], [TAIL_MAMBA] ++ ) ++ ] ++ assert hashed == [boundary] +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index 649f0a4..cf2b716 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -1710,6 +1710,17 @@ class MambaManager(SingleTypeKVCacheManager): + self.block_pool.hash_block_size, + self.block_size, + ) ++ if ( ++ self.mamba_cache_mode == "align" ++ and self.enable_caching ++ and self.block_size > self.block_pool.hash_block_size ++ ): ++ logger.info_once( ++ "Mamba align: a prompt tail key that lands on a block end is " ++ "hashed too (hash unit %d, block %d)", ++ self.block_pool.hash_block_size, ++ self.block_size, ++ ) + + def _retire_state_blocks( + self, request_id: str, retired: list[tuple[int, KVCacheBlock]] +@@ -2356,6 +2367,14 @@ class MambaManager(SingleTypeKVCacheManager): + replay_boundaries: Sequence[int], + ) -> None: + num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0) ++ tail = self._block_aligned_prompt_tail(request) ++ if tail: ++ # With a hash unit finer than the block, the prompt's tail ++ # key (where the next turn resumes; `_cache_partial_tail_block`) ++ # can land exactly on a block end. It is then a full-block key, ++ # not a partial one, and the replay-boundary mask alone (one block ++ # below it under the eagle drop) would never hash that state. ++ replay_boundaries = (*replay_boundaries, tail) + super().cache_blocks( + request, + num_tokens, +@@ -2397,6 +2416,22 @@ class MambaManager(SingleTypeKVCacheManager): + def new_step_starts(self) -> None: + self.cached_blocks_this_step.clear() + ++ def _block_aligned_prompt_tail(self, request: Request) -> int: ++ """The prompt's tail key when it falls on a block end, else 0. ++ ++ Only with a hash unit finer than the block (at the block unit the ++ tail IS the replay boundary, which the mask already keeps). The key ++ is placed exactly as `_cache_partial_tail_block` places a sub-block ++ tail, eagle shift included. ++ """ ++ hash_block_size = self.block_pool.hash_block_size ++ if self.mamba_cache_mode != "align" or self.block_size <= hash_block_size: ++ return 0 ++ tail = request.num_prompt_tokens // hash_block_size * hash_block_size ++ if self.drop_eagle_checkpoint_block: ++ tail = max(tail - hash_block_size, 0) ++ return tail if tail % self.block_size == 0 else 0 ++ + def _cache_partial_tail_block( + self, + request: Request, +-- +2.43.0 + diff --git a/patches/0095-scheduler-pin-a-waiting-request-s-prefix-cache-hit-u.patch b/patches/0095-scheduler-pin-a-waiting-request-s-prefix-cache-hit-u.patch new file mode 100644 index 0000000..c5859c0 --- /dev/null +++ b/patches/0095-scheduler-pin-a-waiting-request-s-prefix-cache-hit-u.patch @@ -0,0 +1,591 @@ +From 9ebb3348c4c3ef2120b3e7aa1b0f10099788f8ba Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 6 Oct 2026 02:49:52 -0400 +Subject: [PATCH 095/103] scheduler: pin a waiting request's prefix-cache hit + until it is admitted + +A follow-up turn that waits for admission holds no reference on its +cached prefix. With a prefix unit finer than the block, its whole hit +depends on every one of those blocks: the resume key sits in the top +attention block, and the only GDN state is the held tail. So the first +of them that another request's allocation popped from the free queue +turned the next turn into a cold prefill; that was usually its own top +block, which is freed first. Under a long-context multi-turn agent load +(8 conversations, 80-100K-token prompts) this gave 4-5 total misses per +300 s at unit 64/512 and none at unit 4096. All of them happened while +the pool was at 88-99 % with one or two conversations waiting, and the +conversation that missed was usually the next victim too. + +Now the manager can pin a waiter's hit: + +- pin_waiting_hit looks the request up and takes a ref on every + non-null hit block in every group. It also pins the block holding the + EAGLE/MTP proof unit: one unit past the key, dropped from the returned + hit, and a different block when the key ends a block. The pinned + blocks stay out of the free queue. +- release_pinned_hit drops a pin; yield_pinned_hits drops the most + recently pinned waiters' first. +- allocate_slots releases the request's pin when it succeeds (its table + holds the hit from then on), free releases it (abort, finish), and + reset_prefix_cache yields every pin before resetting. + +Pinned blocks are referenced, so admission and the watermark see them +as used. + +The scheduler pins the front of the waiting queue after each scheduling +pass: arrival order, at most MAX_PINNED_WAITERS = 8, only requests that +have not computed anything yet, never with a KV connector. It does this +also while the running set is full, when a waiter is not even looked +up. A pin never costs anyone blocks. A running request that fails to +allocate yields one waiter's pin and retries before any preemption +victim is chosen. A head-of-line waiter whose admission fails yields the +pins of the waiters behind it and retries once. When pins are active +the scheduler logs one startup line saying so, with the bound, the hash +unit and the block size. + +On a closed-loop CPU model of that load (6-8 conversations, 1-2 +waiting, pool 88-100 %), total misses went from 3/1/1/5 to 2/0/0/0 with +every waiter pinned. The remaining two sit at 99.2 % use, where the +running set needed the pinned blocks themselves (two yields, 65 +blocks). Unit 4096 is unaffected except where a waiter was losing its +boundary block. + +Tests: +- A pinned waiter survives another follow-up's admission and a burst, + and resumes at its tail (a total miss without the pin). +- Pool full and a running request needs a block: the pin yields, nobody + is preempted, and the waiter is still admitted later. +- Pins are released on abort, admission, yield and reset without + leaks; the bound is respected; yield order is newest first. +- The scheduler's pass pins in arrival order up to the bound and skips + partly computed requests and connector setups. +--- + .../test_partial_prefix_cache_hits.py | 223 +++++++++++++++++- + vllm/v1/core/kv_cache_manager.py | 112 +++++++++ + vllm/v1/core/sched/scheduler.py | 70 ++++++ + 3 files changed, 401 insertions(+), 4 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index 5247e08..f5de3e7 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -2556,7 +2556,9 @@ def tail_retain(request, monkeypatch) -> bool: + return request.param == "1" + + +-def _tail_manager(num_blocks: int, use_eagle: bool = True): ++def _tail_manager( ++ num_blocks: int, use_eagle: bool = True, attn_block: int = TAIL_HASH ++): + """Full attention at the hash unit + one align-mode Mamba group at the + block, eagle/MTP drop on, retention interval 0 (the deployed default).""" + kv_cache_config = KVCacheConfig( +@@ -2566,7 +2568,7 @@ def _tail_manager(num_blocks: int, use_eagle: bool = True): + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( +- block_size=TAIL_HASH, ++ block_size=attn_block, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, +@@ -2679,7 +2681,9 @@ def _tail_position(num_prompt_tokens: int) -> int: + return num_prompt_tokens // TAIL_HASH * TAIL_HASH - TAIL_HASH + + +-def _tail_followup(prev, user_tokens: int, seed: int, aligned_tail: bool = False): ++def _tail_followup( ++ prev, user_tokens: int, seed: int, aligned_tail: bool = False, suffix: str = "+" ++): + """The next turn: the previous turn's prompt + its output + new text, + sized so that its tail lands strictly inside a block above the last + cacheable block end (the common case), or -- ``aligned_tail`` -- exactly +@@ -2692,7 +2696,7 @@ def _tail_followup(prev, user_tokens: int, seed: int, aligned_tail: bool = False + break + user_tokens += 1 + return make_request( +- prev.request_id + "+", ++ prev.request_id + suffix, + list(prev.all_token_ids) + _tail_tokens(seed, user_tokens), + TAIL_HASH, + sha256, +@@ -3153,3 +3157,214 @@ def test_block_unit_replay_boundary_is_unchanged(): + ) + ] + assert hashed == [boundary] ++ ++ ++# --------------------------------------------------------------------------- ++# A waiting follow-up's prefix-cache hit is pinned until admission. ++# --------------------------------------------------------------------------- ++def _pinned(manager) -> dict: ++ return getattr(manager, "_pinned_hits", {}) ++ ++ ++def _tail_run(manager, request, decode_steps: int = 2) -> int: ++ hit = _tail_prefill(manager, request) ++ _tail_decode(manager, request, steps=decode_steps) ++ return hit ++ ++ ++def test_pinned_waiter_survives_concurrent_allocation(tail_retain): ++ """A1 finished and its next turn A2 waits (a slot is busy). Another ++ follow-up B is admitted meanwhile and its CoW/growth allocations pop the ++ free queue. Unpinned (the previous commit): A1's top attention block goes first ++ and A2 is a total miss. Pinned: A2's hit (attention prefix + the GDN ++ hold) is referenced, B's allocations take other blocks, A2 resumes at ++ A1's tail; the pin is released by A2's own admission and nothing leaks. ++ """ ++ num_blocks = 64 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ b0 = make_request("B0", _tail_tokens(2, 150), TAIL_HASH, sha256) ++ _tail_run(manager, b0) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ a1 = _tail_followup(a0, 26, seed=3) ++ _tail_run(manager, a1) ++ manager.free(a1) ++ _tail_drain_copies(manager) ++ manager.free(b0) ++ _tail_drain_copies(manager) ++ ++ # A2 arrives and must wait; B1 (older in the queue) is admitted. ++ a2 = _tail_followup(a1, 20, seed=4) ++ a1_tail = _tail_position(a1.num_prompt_tokens) ++ pinned = manager.pin_waiting_hit(a2) if hasattr(manager, "pin_waiting_hit") else 0 ++ assert pinned == ( ++ len([b for g in manager.get_computed_blocks(a2)[0].blocks for b in g if not b.is_null]) ++ ) ++ assert manager.pin_waiting_hit(a2) == 0 # idempotent ++ hold = _tail_mamba(manager) ++ hit_blocks = _pinned(manager)[a2.request_id] ++ assert all(b.ref_cnt == 1 for b in hit_blocks) ++ assert set(b.block_id for b in hit_blocks).isdisjoint(_tail_free_ids(manager)) ++ ++ # Pool pressure from B1's admission and a burst: everything free goes. ++ b1 = _tail_followup(b0, 26, seed=5) ++ _tail_run(manager, b1) ++ taken = pool.get_new_blocks(pool.get_num_free_blocks()) ++ assert all(b.block_hash is not None for b in hit_blocks) ++ pool.free_blocks(taken) ++ ++ # A2 is admitted: its hit is intact; the admission releases the pin. ++ computed, num_hit, _ = manager.get_computed_blocks(a2) ++ assert num_hit == a1_tail ++ _tail_prefill(manager, a2) ++ assert a2.request_id not in _pinned(manager) ++ # Still cached: A2's table or CoW sources, never evicted. ++ assert all(b.block_hash is not None for b in hit_blocks) ++ _tail_decode(manager, a2, steps=2) ++ manager.free(a2) ++ manager.free(b1) ++ _tail_drain_copies(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ del hold ++ ++ ++def test_pins_yield_before_a_running_request_fails_and_waiter_is_admitted_later( ++ tail_retain, ++): ++ """Pool exactly full: a running request needs a block, the free queue is ++ empty, the only reclaimable blocks are a waiter's pin. The pin yields ++ (``yield_pinned_hits``), the running request allocates without anyone ++ being preempted, and the waiter is still admitted when a slot frees ++ (possibly with a smaller hit). Nothing leaks.""" ++ num_blocks = 48 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ a1 = _tail_followup(a0, 26, seed=3) ++ a1_tail = _tail_position(a1.num_prompt_tokens) ++ assert manager.pin_waiting_hit(a1) > 0 ++ pinned_ids = {b.block_id for b in _pinned(manager)[a1.request_id]} ++ ++ # R runs and consumes the whole free queue (a cold prefill), then needs ++ # more: the pin is the only thing left to reclaim. ++ r = make_request("R", _tail_tokens(7, 2 * TAIL_BLOCK), TAIL_HASH, sha256) ++ _tail_run(manager, r, decode_steps=0) ++ taken = pool.get_new_blocks(pool.get_num_free_blocks()) ++ assert pool.get_num_free_blocks() == 0 ++ r.append_output_token_ids(777) ++ r.spec_token_ids = [778, 779, 780] ++ blocks = manager.allocate_slots(r, 1 + TAIL_SPEC, num_lookahead_tokens=TAIL_SPEC) ++ if blocks is None: ++ # The scheduler's step: yield pins, retry, no preemption. ++ released = manager.yield_pinned_hits() ++ assert released == len(pinned_ids) ++ assert a1.request_id not in _pinned(manager) ++ blocks = manager.allocate_slots(r, 1 + TAIL_SPEC, num_lookahead_tokens=TAIL_SPEC) ++ assert blocks is not None ++ r.spec_token_ids = [] ++ r.num_computed_tokens += 1 ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ pool.free_blocks(taken) ++ ++ # Slot frees: the waiter is admitted (hit may have shrunk, never stuck). ++ manager.free(r) ++ _tail_drain_copies(manager) ++ computed, num_hit, _ = manager.get_computed_blocks(a1) ++ assert 0 <= num_hit <= a1_tail ++ _tail_run(manager, a1) ++ manager.free(a1) ++ _tail_drain_copies(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ assert not _pinned(manager) ++ ++ ++def test_pin_released_on_abort_admission_and_yield_without_leak(tail_retain): ++ num_blocks = 64 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ free_baseline = pool.get_num_free_blocks() ++ ++ # Abort while pinned. ++ w1 = _tail_followup(a0, 20, seed=3, suffix="w1") ++ assert manager.pin_waiting_hit(w1) > 0 ++ assert pool.get_num_free_blocks() < free_baseline ++ manager.free(w1) # abort path ++ assert not _pinned(manager) and pool.get_num_free_blocks() == free_baseline ++ ++ # Bound: at most MAX_PINNED_WAITERS pins; one more is refused. ++ w2 = _tail_followup(a0, 21, seed=4, suffix="w2") ++ w3 = _tail_followup(a0, 22, seed=5, suffix="w3") ++ assert manager.pin_waiting_hit(w2) > 0 ++ assert manager.pin_waiting_hit(w3) > 0 ++ extra = [ ++ _tail_followup(a0, 24 + i, seed=10 + i, suffix=f"x{i}") ++ for i in range(manager.MAX_PINNED_WAITERS - 1) ++ ] ++ for req in extra[:-1]: ++ assert manager.pin_waiting_hit(req) > 0 ++ assert manager.pin_waiting_hit(extra[-1]) == 0 ++ assert manager.num_pinned_waiters == manager.MAX_PINNED_WAITERS ++ # Yield one at a time, newest first, keeping the head. ++ assert manager.yield_pinned_hits(keep=w2.request_id, limit=1) > 0 ++ assert list(_pinned(manager))[-1] == extra[-3].request_id ++ assert manager.yield_pinned_hits(keep=w2.request_id) > 0 ++ assert list(_pinned(manager)) == [w2.request_id] ++ # Admission releases the head's pin; the hit it brought is referenced. ++ _tail_run(manager, w2) ++ assert not _pinned(manager) ++ manager.free(w2) ++ _tail_drain_copies(manager) ++ assert pool.get_num_free_blocks() == free_baseline ++ # reset_prefix_cache yields pins first and succeeds. ++ assert manager.pin_waiting_hit(w3) > 0 ++ assert manager.reset_prefix_cache() ++ assert not _pinned(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_scheduler_pins_the_front_of_the_waiting_queue(tail_retain): ++ """The scheduler's pass pins waiting follow-ups in arrival order up to ++ the manager's bound, skips requests that are not fresh lookups (partly ++ computed, async KV) and does nothing with a connector or caching off.""" ++ manager = _tail_manager(96, attn_block=TAIL_BLOCK) ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ waiters = [ ++ _tail_followup(a0, 20 + i, seed=30 + i, suffix=f"q{i}") ++ for i in range(manager.MAX_PINNED_WAITERS + 2) ++ ] ++ waiters[1].num_computed_tokens = 8 # resumed async load: not a lookup ++ cold = make_request("COLD", _tail_tokens(9, 40), TAIL_HASH, sha256) ++ stub = SimpleNamespace( ++ kv_cache_manager=manager, connector=None, waiting=[cold, *waiters] ++ ) ++ Scheduler._pin_waiting_hits(stub) ++ pinned = list(_pinned(manager)) ++ assert waiters[1].request_id not in pinned ++ assert cold.request_id not in pinned # nothing cached for it ++ assert pinned[0] == waiters[0].request_id ++ assert len(pinned) == manager.MAX_PINNED_WAITERS ++ assert waiters[-1].request_id not in pinned # over the bound ++ # A second pass is idempotent; with a connector it is a no-op. ++ Scheduler._pin_waiting_hits(stub) ++ assert list(_pinned(manager)) == pinned ++ manager.yield_pinned_hits() ++ stub.connector = object() ++ Scheduler._pin_waiting_hits(stub) ++ assert not _pinned(manager) ++ _tail_assert_pool_clean(manager, 96) +diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py +index c0b94f0..dac8b8f 100644 +--- a/vllm/v1/core/kv_cache_manager.py ++++ b/vllm/v1/core/kv_cache_manager.py +@@ -226,6 +226,21 @@ class KVCacheManager: + tuple(() for _ in range(self.num_kv_cache_groups)) + ) + ++ # Prefix-cache hits pinned for requests waiting for admission ++ # ({request_id: blocks}). A waiting follow-up turn holds no ref on its ++ # cached prefix, and with a prefix unit finer than the block its ++ # whole hit depends on every one of those blocks (the resume key sits ++ # in the top block); under pool pressure other requests' allocations ++ # popped them from the free queue and the turn re-prefilled from 0 ++ # (under a long-context multi-turn agent load: 4-5 total misses per ++ # 300 s at unit 64/512, none at 4096). Pinned blocks stay out of the ++ # free queue, so they count against admission and the watermark like ++ # any referenced block; they ++ # are released on admission (the table takes over), on free, and ++ # whenever a running request would otherwise fail to allocate ++ # (``yield_pinned_hits``), so a pin never causes a preemption. ++ self._pinned_hits: dict[str, list[KVCacheBlock]] = {} ++ + @property + def usage(self) -> float: + """Get the KV cache usage. +@@ -585,6 +600,7 @@ class KVCacheManager: + # P/D: delay caching blocks if we have to recv from + # remote. Update state for locally cached blocks. + if delay_cache_blocks: ++ self.release_pinned_hit(request.request_id) + return self.create_kv_cache_blocks(new_blocks) + + # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens +@@ -598,6 +614,8 @@ class KVCacheManager: + ) + self.coordinator.cache_blocks(request, num_tokens_to_cache) + ++ # Admitted: the table now references the hit, the pin is redundant. ++ self.release_pinned_hit(request.request_id) + return self.create_kv_cache_blocks(new_blocks) + + def free(self, request: Request) -> None: +@@ -608,8 +626,100 @@ class KVCacheManager: + Args: + request: The request to free the blocks. + """ ++ self.release_pinned_hit(request.request_id) + self.coordinator.free(request.request_id) + ++ # Waiting requests whose prefix-cache hit is pinned at once (arrival ++ # order). Pins never block anyone (they yield before a running request ++ # fails to allocate), so the bound only caps the lookup work per step. ++ MAX_PINNED_WAITERS = 8 ++ ++ @property ++ def num_pinned_waiters(self) -> int: ++ return len(self._pinned_hits) ++ ++ @property ++ def num_pinned_blocks(self) -> int: ++ return sum(len(blocks) for blocks in self._pinned_hits.values()) ++ ++ def pin_waiting_hit(self, request: Request) -> int: ++ """Pin the request's current prefix-cache hit while it waits for ++ admission: every non-null hit block in every group (the attention ++ prefix up to the key and the GDN state it resumes from) takes a ref ++ and leaves the free queue. Idempotent. Returns the number of blocks ++ pinned now (0 if nothing hit, already pinned, caching is off, or the ++ waiter bound is reached).""" ++ request_id = request.request_id ++ if ( ++ request_id in self._pinned_hits ++ or not self.enable_caching ++ or not self.prefix_cache_lookup_enabled(request) ++ or len(self._pinned_hits) >= self.MAX_PINNED_WAITERS ++ ): ++ return 0 ++ computed, num_hit, _ = self.get_computed_blocks(request) ++ blocks = [ ++ block ++ for group in computed.blocks ++ for block in group ++ if not block.is_null ++ ] ++ if num_hit == 0 or not blocks: ++ return 0 ++ # The EAGLE/MTP proof: fine-grained attention hits match one hash ++ # unit past the key and drop it, so the hit also depends on the block ++ # holding that unit -- a different block when the key ends a block ++ # (and never part of the returned hit). Pin it too. ++ unit = self.block_pool.hash_block_size ++ proof_idx = num_hit // unit ++ if proof_idx < len(request.block_hashes): ++ seen = {block.block_id for block in blocks} ++ for group_id in range(self.num_kv_cache_groups): ++ cached = self.block_pool.get_cached_block( ++ request.block_hashes[proof_idx], [group_id] ++ ) ++ if cached and cached[0].block_id not in seen: ++ seen.add(cached[0].block_id) ++ blocks.append(cached[0]) ++ self.block_pool.touch(blocks) ++ self._pinned_hits[request_id] = blocks ++ return len(blocks) ++ ++ def pinned_hit_blocks(self, request_id: str) -> list[KVCacheBlock]: ++ """The blocks pinned for ``request_id`` (empty if none).""" ++ return list(self._pinned_hits.get(request_id, ())) ++ ++ def release_pinned_hit(self, request_id: str) -> int: ++ """Drop the request's pin (admission, abort, finish). Returns the ++ number of blocks released.""" ++ blocks = self._pinned_hits.pop(request_id, None) ++ if not blocks: ++ return 0 ++ # Reversed: the prefix stays an eviction candidate longer than the ++ # tail, as after a free. ++ self.block_pool.free_blocks(reversed(blocks)) ++ return len(blocks) ++ ++ def yield_pinned_hits( ++ self, keep: str | None = None, limit: int | None = None ++ ) -> int: ++ """Drop pins (except ``keep``'s) so a running request can allocate, ++ or a waiter ahead in line can be admitted: the most recently pinned ++ waiters first, at most ``limit`` of them per call (all by default; ++ the scheduler passes 1 and retries, so a decode step that needs a ++ block costs one waiter its pin, not every waiter). Returns the number ++ of blocks released; 0 means there was nothing to yield.""" ++ released = 0 ++ yielded = 0 ++ for request_id in reversed(list(self._pinned_hits)): ++ if request_id == keep: ++ continue ++ if limit is not None and yielded >= limit: ++ break ++ released += self.release_pinned_hit(request_id) ++ yielded += 1 ++ return released ++ + def remove_skipped_blocks( + self, + request_id: str, +@@ -661,6 +771,8 @@ class KVCacheManager: + bool: True if the prefix cache is successfully reset, + False otherwise. + """ ++ # Pinned waiters' hits are about to be invalidated anyway. ++ self.yield_pinned_hits() + if not self.coordinator.reset_prefix_cache(): + return False + if self.log_stats: +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index dc8f033..fea77b1 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -348,6 +348,8 @@ class Scheduler(SchedulerInterface): + if self.connector is not None: + self.connector.bind_kv_cache_manager(self.kv_cache_manager) + ++ self._log_prefix_cache_features() ++ + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 + self.use_v2_model_runner = vllm_config.use_v2_model_runner + # Scheduler iteration counter. Drives the V2+PP+async decode-throttle +@@ -559,6 +561,45 @@ class Scheduler(SchedulerInterface): + end = min((s for s in stops if start < s < end), default=end) + return max(end - start, 0) + ++ def _log_prefix_cache_features(self) -> None: ++ """One startup line per prefix-cache retention feature that is active ++ (nothing when it is off).""" ++ manager = self.kv_cache_manager ++ if manager.enable_caching and self.connector is None: ++ logger.info( ++ "Prefix cache: waiting requests' cached prefixes are pinned " ++ "until admission (up to %d waiters; hash unit %d, block %d)", ++ manager.MAX_PINNED_WAITERS, ++ self.hash_block_size, ++ self.block_size, ++ ) ++ ++ def _pin_waiting_hits(self) -> None: ++ """Pin the prefix-cache hits of the requests at the front of the ++ waiting queue (arrival order, bounded by the manager) so the ++ admissions and decode growth of other requests cannot evict them ++ while they wait -- also while the running set is full, when they are ++ not even looked up. Released on admission, free, and whenever a ++ running request would otherwise fail to allocate. ++ """ ++ manager = self.kv_cache_manager ++ if ( ++ not manager.enable_caching ++ or self.connector is not None ++ or manager.num_pinned_waiters >= manager.MAX_PINNED_WAITERS ++ ): ++ return ++ for request in self.waiting: ++ if manager.num_pinned_waiters >= manager.MAX_PINNED_WAITERS: ++ break ++ if ( ++ request.num_computed_tokens != 0 ++ or request.status ++ not in (RequestStatus.WAITING, RequestStatus.PREEMPTED) ++ ): ++ continue ++ manager.pin_waiting_hit(request) ++ + def _get_local_prefix_cache_hit( + self, request: Request + ) -> tuple[KVCacheBlocks, int, int, bool]: +@@ -776,6 +817,12 @@ class Scheduler(SchedulerInterface): + ): + break + ++ # Waiting requests' pinned prefix hits yield ++ # before anyone is preempted (a pin never costs a running ++ # request its blocks); retry with them in the free queue. ++ if self.kv_cache_manager.yield_pinned_hits(limit=1): ++ continue ++ + # The request cannot be scheduled. + # Preempt the lowest-priority request. + if self.policy == SchedulingPolicy.PRIORITY: +@@ -1211,6 +1258,26 @@ class Scheduler(SchedulerInterface): + has_scheduled_reqs=bool(self.running), + ) + ++ if new_blocks is None and self.kv_cache_manager.yield_pinned_hits( ++ keep=request_id ++ ): ++ # Pins of waiters behind this one yield to it ++ # (arrival order); its own pin stays, those blocks are ++ # its hit. One retry with them back in the free queue. ++ new_blocks = self.kv_cache_manager.allocate_slots( ++ request, ++ num_new_tokens, ++ num_new_computed_tokens=num_new_local_computed_tokens, ++ new_computed_blocks=new_computed_blocks, ++ num_lookahead_tokens=effective_lookahead_tokens, ++ num_external_computed_tokens=num_external_computed_tokens, ++ delay_cache_blocks=load_kv_async, ++ num_encoder_tokens=num_encoder_tokens, ++ full_sequence_must_fit=self.scheduler_reserve_full_isl, ++ reserved_blocks=reserved_blocks, ++ has_scheduled_reqs=bool(self.running), ++ ) ++ + if new_blocks is None: + # The request cannot be scheduled. + +@@ -1338,6 +1405,9 @@ class Scheduler(SchedulerInterface): + if step_skipped_waiting: + self.skipped_waiting.prepend_requests(step_skipped_waiting) + ++ # Whoever is still waiting keeps its prefix-cache hit. ++ self._pin_waiting_hits() ++ + # DP prefill balancing: on a step that admitted prefills (release), + # record whether it was capacity-bound. + if not defer_prefills: +-- +2.43.0 + diff --git a/patches/0096-scheduler-release-a-waiter-s-pin-on-the-deferred-fre.patch b/patches/0096-scheduler-release-a-waiter-s-pin-on-the-deferred-fre.patch new file mode 100644 index 0000000..bf26e27 --- /dev/null +++ b/patches/0096-scheduler-release-a-waiter-s-pin-on-the-deferred-fre.patch @@ -0,0 +1,116 @@ +From 7b88239d5efa36dc3aa394e880c7f011cb732a2f Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 6 Oct 2026 02:49:52 -0400 +Subject: [PATCH 096/103] scheduler: release a waiter's pin on the deferred + free path too + +The scheduler takes pop_blocks_for_free when a step may still be in +flight, for example when a preempted request is aborted before its +fence drained. A waiting request's pin is never written by any step, so +it is now released there at once. Before, it stayed until the next +yield or reset, keeping its blocks pinned and a waiter slot taken. + +pin_waiting_hit touches each block in the pool it belongs to; private +pool groups route like free_blocks. Its docstring states the known gap: +a pin is a snapshot of the hit at pin time. It is not refreshed if the +hit grows while the request waits, or if a producer moves a pinned live +block's hash into its copy. The admission lookup still sees whatever is +cached, so the cost is a smaller hit, never a wrong one. +MAX_PINNED_WAITERS goes from 8 to 16. + +Test: a pinned waiter popped through the deferred path leaves no pin, +no waiter slot and the pool at its baseline. +--- + .../test_partial_prefix_cache_hits.py | 23 +++++++++++++++++++ + vllm/v1/core/kv_cache_manager.py | 23 +++++++++++++++---- + 2 files changed, 42 insertions(+), 4 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index f5de3e7..75d6477 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -3368,3 +3368,26 @@ def test_scheduler_pins_the_front_of_the_waiting_queue(tail_retain): + Scheduler._pin_waiting_hits(stub) + assert not _pinned(manager) + _tail_assert_pool_clean(manager, 96) ++ ++ ++def test_deferred_abort_of_a_pinned_waiter_releases_the_pin(tail_retain): ++ """pop_blocks_for_free (the deferred-free path a scheduler takes when a ++ step may still be in flight) releases a waiter's pin at once: the pin ++ is never written by any step, and leaving it would keep the blocks ++ pinned and a waiter slot taken until the next yield or reset.""" ++ manager = _tail_manager(64, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ manager.free(a0) ++ _tail_drain_copies(manager) ++ baseline = pool.get_num_free_blocks() ++ w = _tail_followup(a0, 20, seed=3) ++ assert manager.pin_waiting_hit(w) > 0 ++ assert pool.get_num_free_blocks() < baseline ++ deferred = manager.pop_blocks_for_free(w) ++ assert deferred == [] # a waiter has no table blocks ++ assert not _pinned(manager) and manager.num_pinned_waiters == 0 ++ assert pool.get_num_free_blocks() == baseline ++ pool.free_blocks(deferred) ++ _tail_assert_pool_clean(manager, 64) +diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py +index dac8b8f..3dbda90 100644 +--- a/vllm/v1/core/kv_cache_manager.py ++++ b/vllm/v1/core/kv_cache_manager.py +@@ -631,8 +631,9 @@ class KVCacheManager: + + # Waiting requests whose prefix-cache hit is pinned at once (arrival + # order). Pins never block anyone (they yield before a running request +- # fails to allocate), so the bound only caps the lookup work per step. +- MAX_PINNED_WAITERS = 8 ++ # fails to allocate), so the bound only caps the lookup work per step; ++ # twice the deployed max_num_seqs covers every waiter seen under load. ++ MAX_PINNED_WAITERS = 16 + + @property + def num_pinned_waiters(self) -> int: +@@ -648,7 +649,15 @@ class KVCacheManager: + prefix up to the key and the GDN state it resumes from) takes a ref + and leaves the free queue. Idempotent. Returns the number of blocks + pinned now (0 if nothing hit, already pinned, caching is off, or the +- waiter bound is reached).""" ++ waiter bound is reached). ++ ++ Known gap: a pin is a snapshot of the hit at pin time. It is not ++ refreshed if the hit grows while the request waits (a sibling ++ computing more of the same prefix) or if a producer moves the hash of ++ a pinned live block into its CoW copy (a resume at a running ++ producer's tail, unreachable under the EAGLE drop); the admission ++ lookup still finds whatever is cached then, so the only cost is a ++ smaller hit, never a wrong one.""" + request_id = request.request_id + if ( + request_id in self._pinned_hits +@@ -681,7 +690,10 @@ class KVCacheManager: + if cached and cached[0].block_id not in seen: + seen.add(cached[0].block_id) + blocks.append(cached[0]) +- self.block_pool.touch(blocks) ++ # Blocks of a private-pool group (hisparse host tier) are touched in ++ # their own pool, as free_blocks routes them. ++ for block in blocks: ++ (block.pool or self.block_pool).touch([block]) + self._pinned_hits[request_id] = blocks + return len(blocks) + +@@ -752,6 +764,9 @@ class KVCacheManager: + The request's blocks in eviction order (each group's blocks + reversed, then concatenated group by group). + """ ++ # A pin is never written by an in-flight step (the waiter has no ++ # table), so it is released now, not with the deferred blocks. ++ self.release_pinned_hit(request.request_id) + return self.coordinator.pop_blocks_for_free(request.request_id) + + def evict_blocks(self, block_ids: set[int]) -> None: +-- +2.43.0 + diff --git a/patches/0097-kv-cache-hold-a-finished-turn-s-prefix-until-its-nex.patch b/patches/0097-kv-cache-hold-a-finished-turn-s-prefix-until-its-nex.patch new file mode 100644 index 0000000..ddb6a48 --- /dev/null +++ b/patches/0097-kv-cache-hold-a-finished-turn-s-prefix-until-its-nex.patch @@ -0,0 +1,1454 @@ +From 81a766855cebf5cd133d95be829b50bd71166ef7 Mon Sep 17 00:00:00 2001 +From: halt95 +Date: Tue, 6 Oct 2026 02:49:52 -0400 +Subject: [PATCH 097/103] kv cache: hold a finished turn's prefix until its + next turn is admitted + +With a prefix unit finer than the block, a hybrid follow-up resumes from +a single GDN state at the previous prompt's tail, and that key sits in +the top attention block. Losing any one block of that prefix is a total +miss. The waiting-request pins only protect a hit that is still whole +when the follow-up is first looked up. Tracing the remaining misses +under a long-context multi-turn agent load shows they lost their top +attention block before that, in the gap between the previous turn +finishing and the next turn arriving. In that gap nothing referenced +the prefix. LRU then consumed it ahead of dead blocks freed later +(other conversations' handed-over CoW sources and superseded states). + +A finished turn (stop, length, repetition, or an abort after decoding) +now keeps one extra ref on each of its hashed blocks, including its GDN +holds, as a grace pin. The ref is taken on both free paths, the +immediate free and the deferred pop_blocks_for_free behind an in-flight +step, so the prefix never passes through the free queue. A grace is +released when: + +- the next turn is admitted on it. That is the grace holding the + deepest hit block of every group, with a hit that reaches its prompt + to within two prefix units, so a new turn sharing only a system + prompt releases nothing. The table takes the blocks over, with no + double ref. +- anyone needs blocks. Graces yield before waiter pins, oldest first and + one per failed allocation, so they never cause a preemption. +- the bound (16) is exceeded. The oldest grace goes. +- reset_prefix_cache runs. + +A head waiter whose admission fails pins its own hit first. It then +yields graces and other pins one at a time, and only when everything +yieldable covers its shortfall. Otherwise yielding would drop other +conversations' prefixes and still not admit it. The shortfall is +measured after the pin: the pin takes the proof block past the hit (and +any hit block the lookup did not return) out of the free queue, so the +shortfall of the failed attempt can be short by those blocks and the +guard would yield another conversation's grace for nothing. Its own +pinned blocks outside the hit it is admitted on, the proof block past +it, go back first: once the hit is taken over they protect nothing, and +they cost nobody else anything. At an exact fit they are the blocks it +needs, and nothing else ever yields a head waiter's own pin, so without +this a request that needs the whole pool waited forever. + +While any pin or grace is held the scheduler reports no common prefix +to cascade attention. Those refs can make a block's count equal the +number of running requests when only some of them share it, and cascade +attention would then read it for all of them. + +Graces and waiter pins now follow a producer's GDN hash into its CoW +copy, as holds already do. Otherwise they would keep a second ref on a +running state block that its owner must own exclusively. Graces apply +only to fine-grained hybrid hits (coordinator.enable_partial_hash_hits) +and are off with a KV connector; a block-sized unit (the 4096 lane) is +unchanged. When graces are active the scheduler logs one startup line +saying so, with the bound, the hash unit and the block size. + +Tests: the arrival-gap reproduction fails on the previous commit (hit 0 +instead of 168, both free paths). It passes here. Further tests cover +the handover refcounts, yield order and bound, finish statuses, a +running request at an exactly full pool, yieldable accounting, +moved-hash following, the pin's extras contract, and six runs on the +real scheduler: among them a head waiter whose pin raises its shortfall +(it gives its proof block back and the one grace that covers the rest +yields), a request that needs the whole pool (admitted on its own proof +block; it waited forever before), and the common-prefix count under a +grace (none reported; the count returns once the pins are gone). +--- + .../test_partial_prefix_cache_hits.py | 818 ++++++++++++++++++ + vllm/v1/core/kv_cache_manager.py | 302 ++++++- + vllm/v1/core/sched/scheduler.py | 99 ++- + vllm/v1/core/single_type_kv_cache_manager.py | 10 +- + 4 files changed, 1210 insertions(+), 19 deletions(-) + +diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +index 75d6477..17e1cbc 100644 +--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py ++++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +@@ -2713,6 +2713,7 @@ def _tail_free_ids(manager) -> set[int]: + + def _tail_assert_pool_clean(manager, num_blocks: int) -> None: + pool = manager.block_pool ++ assert not getattr(manager, "_grace_pins", {}) # released explicitly + assert pool.get_num_free_blocks() == num_blocks - 1 # all but the null block + assert all(blk.ref_cnt == 0 for blk in pool.blocks[1:]) + mamba = _tail_mamba(manager) +@@ -3391,3 +3392,820 @@ def test_deferred_abort_of_a_pinned_waiter_releases_the_pin(tail_retain): + assert pool.get_num_free_blocks() == baseline + pool.free_blocks(deferred) + _tail_assert_pool_clean(manager, 64) ++ ++ ++# --------------------------------------------------------------------------- ++# A finished turn's hashed blocks stay referenced ("grace") until the ++# conversation's next turn is admitted on them. ++# --------------------------------------------------------------------------- ++def _graces(manager) -> dict: ++ return getattr(manager, "_grace_pins", {}) ++ ++ ++def _finish(request, status=None) -> None: ++ from vllm.v1.request import RequestStatus ++ ++ request.status = status or RequestStatus.FINISHED_STOPPED ++ ++ ++def _finish_and_free(manager, request, path: str = "immediate", status=None): ++ """The scheduler's _free_request_blocks for a finished turn: free now, or ++ (a step still in flight) pop now and free the list at the fence.""" ++ _finish(request, status) ++ if path == "immediate": ++ manager.free(request) ++ else: ++ deferred = manager.pop_blocks_for_free(request) ++ grace = _graces(manager).get(request.request_id, []) ++ assert all(b.ref_cnt >= 2 for b in grace) # the table's + the grace's ++ manager.block_pool.free_blocks(deferred) # the fence drains ++ _tail_drain_copies(manager) ++ ++ ++def _table_and_hold(manager, request) -> list: ++ blocks = [b for g in manager.get_blocks(request.request_id).blocks for b in g] ++ held = _tail_mamba(manager)._retained_checkpoints.get(request.request_id) ++ if held is not None: ++ blocks.append(held[1]) ++ return [b for b in blocks if not b.is_null] ++ ++ ++@pytest.mark.parametrize("free_path", ["immediate", "deferred"]) ++def test_finished_turn_keeps_its_prefix_through_the_arrival_gap( ++ tail_retain, free_path ++): ++ """The residual miss at a fine prefix unit: a conversation's turn A1 ++ has finished and its next turn A2 has not been looked up yet. Meanwhile ++ another conversation's follow-up B1 is admitted, which frees blocks that ++ are dead from then on (B0's CoW source, superseded states) -- AFTER A1's ++ prefix -- and the running set then consumes about that many blocks. LRU ++ (without the grace) consumes A1's older prefix first, its top attention block ++ first, while the dead blocks survive: A2 is a total miss. With the grace ++ A1's hashed blocks stay referenced, the dead blocks go, A2 resumes at ++ A1's tail, and A2's admission hands the grace over (no double ref).""" ++ num_blocks = 64 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ _finish_and_free(manager, a0, free_path) ++ a1 = _tail_followup(a0, 26, seed=3) ++ _tail_run(manager, a1) ++ a1_ids = {id(b) for b in _table_and_hold(manager, a1) if b.block_hash is not None} ++ _finish_and_free(manager, a1, free_path) ++ a1_tail = _tail_position(a1.num_prompt_tokens) ++ ++ b0 = make_request("B0", _tail_tokens(2, 150), TAIL_HASH, sha256) ++ _tail_run(manager, b0) ++ _finish_and_free(manager, b0, free_path) ++ b1 = _tail_followup(b0, 26, seed=5) ++ _tail_run(manager, b1) # B1 stays running ++ ++ # The running set consumes as many blocks as are free besides A1's. ++ free_now = pool.free_block_queue.get_all_free_blocks() ++ dead = [b for b in free_now if id(b) not in a1_ids] ++ taken = pool.get_new_blocks(len(dead)) ++ ++ a2 = _tail_followup(a1, 20, seed=4) ++ _, num_hit, _ = manager.get_computed_blocks(a2) ++ if tail_retain: ++ assert num_hit == a1_tail ++ # (Retain off = the no-hold control: A1's tail state is a CoW copy ++ # nobody holds, so it is among the "dead" blocks and the hit drops.) ++ pool.free_blocks(taken) ++ ++ # Admission hands the grace over: the table holds the hit, nothing else. ++ _tail_prefill(manager, a2) ++ if tail_retain: ++ assert a1.request_id not in _graces(manager) ++ attn = manager.coordinator.single_type_managers[TAIL_ATTN] ++ for block in attn.req_to_blocks[a2.request_id]: ++ if id(block) in a1_ids: ++ assert block.ref_cnt == 1 ++ _tail_decode(manager, a2, steps=2) ++ _finish_and_free(manager, a2, free_path) ++ _finish_and_free(manager, b1, free_path) ++ assert {a2.request_id, b1.request_id} <= set(_graces(manager)) ++ assert manager.yield_pinned_hits() > 0 ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_grace_pins_yield_before_waiter_pins_oldest_first_and_are_bounded( ++ tail_retain, ++): ++ num_blocks = 192 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ finished = [] ++ for i in range(3): ++ req = make_request(f"F{i}", _tail_tokens(10 + i, 150), TAIL_HASH, sha256) ++ _tail_run(manager, req) ++ _finish_and_free(manager, req) ++ finished.append(req) ++ assert list(_graces(manager)) == ["F0", "F1", "F2"] ++ for req in finished: ++ grace = manager.grace_pinned_blocks(req.request_id) ++ assert grace and all(b.block_hash is not None for b in grace) ++ assert all(b.ref_cnt == 1 for b in grace) ++ ++ # A waiter pinned on F2's prefix: graces yield first, oldest first. ++ w = _tail_followup(finished[2], 20, seed=5) ++ assert manager.pin_waiting_hit(w) > 0 ++ f0 = manager.grace_pinned_blocks("F0") ++ assert manager.yield_pinned_hits(limit=1) == len(f0) ++ assert all(b.ref_cnt == 0 for b in f0) # back in the free queue ++ assert list(_graces(manager)) == ["F1", "F2"] ++ assert w.request_id in _pinned(manager) ++ assert manager.yield_pinned_hits(keep=w.request_id, limit=2) > 0 ++ assert not _graces(manager) and w.request_id in _pinned(manager) ++ # The waiter's hit outlived F2's grace: its own pin holds it. ++ assert all(b.ref_cnt == 1 for b in _pinned(manager)[w.request_id]) ++ ++ # Bound: beyond MAX_GRACE_PINS the oldest grace goes. ++ for i in range(manager.MAX_GRACE_PINS + 2): ++ req = make_request( ++ f"G{i}", _tail_tokens(100 + i, 2 * TAIL_BLOCK + 9), TAIL_HASH, sha256 ++ ) ++ _tail_run(manager, req, decode_steps=1) ++ _finish_and_free(manager, req) ++ assert manager.num_grace_pins == manager.MAX_GRACE_PINS ++ assert list(_graces(manager))[0] == "G2" ++ ++ # An aborted waiter (no output) takes no grace; reset releases all. ++ from vllm.v1.request import RequestStatus ++ ++ _finish(w, RequestStatus.FINISHED_ABORTED) ++ manager.free(w) ++ assert w.request_id not in _graces(manager) and not _pinned(manager) ++ assert manager.reset_prefix_cache() ++ assert not _graces(manager) and not _pinned(manager) ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++@pytest.mark.parametrize("free_path", ["immediate", "deferred"]) ++def test_grace_only_for_turns_that_finished_with_output(tail_retain, free_path): ++ """Abort after decoding / stop / length: grace. Error, ignored, a ++ preemption (frees while RUNNING), graces off (connector) and a ++ block-sized unit: plain free. Nothing leaks.""" ++ from vllm.v1.request import RequestStatus ++ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ ++ def turn(name, seed, status): ++ req = make_request( ++ name, _tail_tokens(seed, 2 * TAIL_BLOCK + 9), TAIL_HASH, sha256 ++ ) ++ _tail_run(manager, req, decode_steps=1) ++ _finish_and_free(manager, req, free_path, status) ++ return req ++ ++ graced = [ ++ turn("stop", 1, RequestStatus.FINISHED_STOPPED), ++ turn("len", 2, RequestStatus.FINISHED_LENGTH_CAPPED), ++ turn("abort", 3, RequestStatus.FINISHED_ABORTED), ++ ] ++ plain = [ ++ turn("err", 4, RequestStatus.FINISHED_ERROR), ++ turn("ign", 5, RequestStatus.FINISHED_IGNORED), ++ turn("pre", 6, RequestStatus.RUNNING), # a preemption's free ++ ] ++ assert list(_graces(manager)) == [r.request_id for r in graced] ++ for req in plain: ++ assert req.request_id not in _graces(manager) ++ held = manager.num_grace_pinned_blocks ++ assert held > 0 ++ assert pool.get_num_free_blocks() == num_blocks - 1 - held ++ manager.enable_grace_pins = False # what the scheduler sets with a connector ++ turn("conn", 7, RequestStatus.FINISHED_STOPPED) ++ assert "conn" not in _graces(manager) ++ manager.enable_grace_pins = True ++ assert manager.yield_pinned_hits() == held ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ # A block-sized prefix unit (the U=4096 lane): graces are off. ++ coarse = make_kv_cache_manager( ++ kv_cache_config=manager.kv_cache_config, ++ max_model_len=8192, ++ enable_caching=True, ++ hash_block_size=TAIL_BLOCK, ++ use_eagle=True, ++ retention_interval=0, ++ ) ++ assert not coarse.enable_grace_pins ++ ++ ++def test_grace_pin_never_stops_a_running_request_at_a_full_pool(tail_retain): ++ """Pool exactly full (graces + running): a running request that needs a ++ block gets one after yielding one grace (the oldest), no preemption, ++ the newer graces stay; nothing leaks.""" ++ num_blocks = 40 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ pool = manager.block_pool ++ r = make_request("R", _tail_tokens(7, 2 * TAIL_BLOCK), TAIL_HASH, sha256) ++ _tail_run(manager, r, decode_steps=0) ++ for i in range(3): ++ req = make_request( ++ f"F{i}", _tail_tokens(20 + i, 2 * TAIL_BLOCK + 9), TAIL_HASH, sha256 ++ ) ++ _tail_run(manager, req, decode_steps=1) ++ _finish_and_free(manager, req) ++ taken = pool.get_new_blocks(pool.get_num_free_blocks()) ++ assert pool.get_num_free_blocks() == 0 and manager.num_grace_pins == 3 ++ r.append_output_token_ids(777) ++ r.spec_token_ids = [778, 779, 780] ++ blocks = manager.allocate_slots(r, 1 + TAIL_SPEC, num_lookahead_tokens=TAIL_SPEC) ++ assert blocks is None and manager.last_alloc_shortfall > 0 ++ # The scheduler's running loop: yield one, retry; never preempt. ++ yields = 0 ++ while blocks is None: ++ assert manager.yield_pinned_hits(limit=1) > 0, "nothing left to yield" ++ yields += 1 ++ blocks = manager.allocate_slots( ++ r, 1 + TAIL_SPEC, num_lookahead_tokens=TAIL_SPEC ++ ) ++ assert yields == 1 and manager.last_alloc_shortfall == 0 ++ assert list(_graces(manager)) == ["F1", "F2"] ++ r.spec_token_ids = [] ++ r.num_computed_tokens += 1 ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ pool.free_blocks(taken) ++ _finish_and_free(manager, r) ++ manager.yield_pinned_hits() ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_yieldable_blocks_exclude_the_head_waiters_own_hit(tail_retain): ++ """num_yieldable_blocks counts only blocks a release would return to the ++ free queue: never the head waiter's own pinned hit.""" ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks, attn_block=TAIL_BLOCK) ++ a0 = make_request("A0", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ _tail_run(manager, a0) ++ _finish_and_free(manager, a0) ++ grace = manager.grace_pinned_blocks("A0") ++ assert manager.num_yieldable_blocks() == len(grace) ++ a1 = _tail_followup(a0, 26, seed=3) ++ assert manager.pin_waiting_hit(a1) > 0 ++ own = {id(b) for b in _pinned(manager)[a1.request_id]} ++ others = [b for b in grace if id(b) not in own] ++ assert manager.num_yieldable_blocks(keep=a1.request_id) == len(others) ++ assert manager.num_yieldable_blocks() == len({id(b) for b in grace} | own) ++ _tail_run(manager, a1) # admission releases both the pin and A0's grace ++ assert not _graces(manager) and not _pinned(manager) ++ _finish_and_free(manager, a1) ++ manager.yield_pinned_hits() ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++def test_pins_and_graces_follow_a_producers_moved_tail(tail_retain): ++ """Without the eagle drop a sibling can resume at a producer's LIVE tail ++ block. A waiter pin (and, retain on, a finished consumer's grace via its ++ hold) on that block must follow the hash into the producer's CoW copy: ++ otherwise it leaves ref 2 on a running state block its owner keeps ++ writing, and protects nothing.""" ++ from vllm.v1.request import RequestStatus ++ ++ num_blocks = 96 ++ manager = _tail_manager(num_blocks, use_eagle=False) ++ pool = manager.block_pool ++ ++ producer = make_request("P", _tail_tokens(1, 150), TAIL_HASH, sha256) ++ tail = producer.num_prompt_tokens // TAIL_HASH * TAIL_HASH ++ computed, num_hit, _ = manager.get_computed_blocks(producer) ++ while producer.num_computed_tokens < tail: ++ num_new = _tail_split( ++ producer, ++ producer.num_tokens - producer.num_computed_tokens, ++ use_eagle=False, ++ ) ++ assert manager.allocate_slots( ++ producer, num_new, num_hit, computed, num_lookahead_tokens=TAIL_SPEC ++ ) ++ producer.num_computed_tokens += num_new ++ computed, num_hit = None, 0 ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ [live] = pool.get_cached_block( ++ producer.block_hashes[tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ ++ # A consumer admitted on the live tail finishes (aborted after one output ++ # token) before the producer moves on; retain on, its hold of the live ++ # block becomes part of its grace. ++ consumer = make_request( ++ "C", list(producer.prompt_token_ids) + _tail_tokens(2, 30), TAIL_HASH, sha256 ++ ) ++ c_computed, c_hit, _ = manager.get_computed_blocks(consumer) ++ assert c_hit == tail ++ num_new = _tail_split( ++ consumer, consumer.num_tokens - c_hit, c_hit, use_eagle=False ++ ) ++ assert manager.allocate_slots( ++ consumer, num_new, c_hit, c_computed, num_lookahead_tokens=TAIL_SPEC ++ ) ++ consumer.num_computed_tokens = c_hit + num_new ++ consumer.append_output_token_ids(4242) ++ _finish(consumer, RequestStatus.FINISHED_ABORTED) ++ manager.free(consumer) ++ assert (live in manager.grace_pinned_blocks("C")) == tail_retain ++ # And a waiter pinned on the same live tail. ++ w = make_request( ++ "W", list(producer.prompt_token_ids) + _tail_tokens(3, 30), TAIL_HASH, sha256 ++ ) ++ assert manager.pin_waiting_hit(w) > 0 ++ assert live in _pinned(manager)["W"] ++ ++ # The producer's next chunk moves the tail hash into its CoW block. ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ num_new = _tail_split( ++ producer, producer.num_tokens - producer.num_computed_tokens, use_eagle=False ++ ) ++ assert manager.allocate_slots(producer, num_new, num_lookahead_tokens=TAIL_SPEC) ++ producer.num_computed_tokens += num_new ++ [moved] = pool.get_cached_block( ++ producer.block_hashes[tail // TAIL_HASH - 1], [TAIL_MAMBA] ++ ) ++ assert moved is not live and live.block_hash is None ++ assert live not in manager.grace_pinned_blocks("C") ++ assert live not in _pinned(manager)["W"] and moved in _pinned(manager)["W"] ++ if tail_retain: ++ assert moved in manager.grace_pinned_blocks("C") ++ manager.new_step_starts() ++ _tail_drain_copies(manager) ++ assert live.ref_cnt == 1 # the producer's table alone ++ ++ _tail_decode(manager, producer, steps=2) ++ _finish_and_free(manager, producer) ++ _, w_hit, _ = manager.get_computed_blocks(w) ++ assert w_hit == tail ++ manager.free(w) # abort the waiter (no output: no grace) ++ manager.yield_pinned_hits() ++ _tail_assert_pool_clean(manager, num_blocks) ++ ++ ++# --- the real scheduler (full attention, block 16; graces forced on, as the ++# --- fine-grained hybrid lane gets them by default) ------------------------- ++def _sched_request(request_id: str, tokens: list[int], max_tokens: int): ++ from vllm.sampling_params import SamplingParams ++ from vllm.v1.core.kv_cache_utils import get_request_block_hasher ++ from vllm.v1.request import Request ++ ++ params = SamplingParams(ignore_eos=True, max_tokens=max_tokens) ++ params.update_from_generation_config({}, 50256) ++ return Request( ++ request_id=request_id, ++ prompt_token_ids=tokens, ++ sampling_params=params, ++ pooling_params=None, ++ block_hasher=get_request_block_hasher(16, sha256), ++ ) ++ ++ ++def _sched_step(scheduler): ++ from vllm.v1.outputs import ModelRunnerOutput ++ ++ out = scheduler.schedule() ++ req_ids = list(out.num_scheduled_tokens) ++ scheduler.update_from_output( ++ out, ++ ModelRunnerOutput( ++ req_ids=req_ids, ++ req_id_to_index={r: i for i, r in enumerate(req_ids)}, ++ sampled_token_ids=[[7] for _ in req_ids], ++ logprobs=None, ++ prompt_logprobs_dict={}, ++ pooler_output=[], ++ ), ++ ) ++ return out ++ ++ ++def _grace_scheduler(num_blocks: int): ++ import os ++ ++ from tests.v1.core.utils import create_scheduler ++ ++ scheduler = create_scheduler( ++ model=os.environ.get("VLLM_TEST_DEFER_FREE_MODEL", "facebook/opt-125m"), ++ enable_prefix_caching=True, ++ num_blocks=num_blocks, ++ block_size=16, ++ max_num_seqs=4, ++ ) ++ scheduler.kv_cache_manager.enable_grace_pins = True ++ return scheduler ++ ++ ++def _sched_finish(scheduler, request) -> None: ++ scheduler.add_request(request) ++ while not request.is_finished(): ++ _sched_step(scheduler) ++ ++ ++def test_scheduler_running_request_yields_graces_never_preempts(): ++ scheduler = _grace_scheduler(num_blocks=12) # 11 usable ++ manager = scheduler.kv_cache_manager ++ pool = manager.block_pool ++ a = _sched_request("A", list(range(1000, 1048)), max_tokens=2) ++ _sched_finish(scheduler, a) ++ assert manager.num_grace_pinned_blocks == 3 # the 3 full prompt blocks ++ assert pool.get_num_free_blocks() == 8 ++ ++ r = _sched_request("R", list(range(2000, 2128)), max_tokens=4) # 8 blocks ++ scheduler.add_request(r) ++ out = _sched_step(scheduler) # admitted on the 8 free blocks ++ assert "R" in out.num_scheduled_tokens and manager.num_grace_pins == 1 ++ assert pool.get_num_free_blocks() == 0 ++ out = _sched_step(scheduler) # token 129 needs a 9th block ++ assert "R" in out.num_scheduled_tokens ++ assert r.num_preemptions == 0 and not out.preempted_req_ids ++ assert manager.num_grace_pins == 0 # A's grace yielded instead ++ while not r.is_finished(): ++ _sched_step(scheduler) ++ assert list(manager._grace_pins) == ["R"] ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 11 ++ ++ ++def test_scheduler_head_admission_yields_graces_only_when_they_cover_it(): ++ scheduler = _grace_scheduler(num_blocks=12) # 11 usable ++ manager = scheduler.kv_cache_manager ++ pool = manager.block_pool ++ a = _sched_request("A", list(range(1000, 1048)), max_tokens=2) ++ _sched_finish(scheduler, a) # grace 3, free 8 ++ r = _sched_request("R", list(range(2000, 2112)), max_tokens=6) # 7 blocks ++ scheduler.add_request(r) ++ _sched_step(scheduler) ++ assert pool.get_num_free_blocks() == 1 ++ # W needs 4 blocks; R's next decode takes the free one first, so 4 > 0 ++ # free + 3 yieldable: no yield, W waits, A keeps its grace, nobody is ++ # preempted. ++ w = _sched_request("W", list(range(3000, 3064)), max_tokens=2) ++ scheduler.add_request(w) ++ out = _sched_step(scheduler) ++ assert "W" not in out.num_scheduled_tokens ++ assert manager.grace_pinned_blocks("A") and r.num_preemptions == 0 ++ while not r.is_finished(): ++ _sched_step(scheduler) ++ assert "W" not in scheduler.running or r.is_finished() ++ # R finished: grace 7, free 1 (R's unhashed tail). W (head) needs 4: ++ # the graces cover it, the OLDEST (A's) yields first and is enough. ++ out = _sched_step(scheduler) ++ assert "W" in out.num_scheduled_tokens ++ assert "A" not in manager._grace_pins and "R" in manager._grace_pins ++ assert w.num_preemptions == 0 and r.num_preemptions == 0 ++ while not w.is_finished(): ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 11 ++ ++ ++def _proof_pin_scene(): ++ """99 usable blocks, EAGLE drop on. A (80 tokens) and B (16) finished; R ++ took every free block and its first decode step made A's grace (the ++ oldest) yield: R consumed A's 5th block, A's first four stay cached in ++ the free queue. Returns (scheduler, manager, pool, a, a4, r): free 4, ++ B's 1-block grace the only other pin.""" ++ scheduler = _grace_scheduler(num_blocks=100) ++ manager = scheduler.kv_cache_manager ++ manager.coordinator.eagle_group_ids = {0} # lookups drop the last hit block ++ pool = manager.block_pool ++ a = _sched_request("A", list(range(1000, 1080)), max_tokens=2) ++ _sched_finish(scheduler, a) # grace: its 5 prompt blocks ++ b = _sched_request("B", list(range(2000, 2016)), max_tokens=2) ++ _sched_finish(scheduler, b) # grace: 1 block, newer than A's ++ assert pool.get_num_free_blocks() == 93 ++ r = _sched_request("R", list(range(3000, 4488)), max_tokens=4) ++ scheduler.add_request(r) ++ out = _sched_step(scheduler) ++ assert out.num_scheduled_tokens["R"] == 1488 and pool.get_num_free_blocks() == 0 ++ _sched_step(scheduler) ++ assert list(manager._grace_pins) == ["B"] and r.num_preemptions == 0 ++ assert pool.get_num_free_blocks() == 4 ++ [a4] = pool.get_cached_block(a.block_hashes[3], [0]) ++ assert a4.ref_cnt == 0 ++ assert pool.get_cached_block(a.block_hashes[4], [0]) is None ++ return scheduler, manager, pool, a, a4, r ++ ++ ++def test_scheduler_head_waiter_measures_its_shortfall_after_pinning(): ++ """The head waiter's own pin can raise its shortfall. With the EAGLE drop ++ its hit ends one block below what is cached, and that proof block is ++ pinned along with the hit; when it was free, the failed admission counted ++ it as free and the pin takes it out of the queue. The yield decision uses ++ the shortfall measured after the pin, and the head gives its proof block ++ back first: it is admitted in the same step on exactly the blocks it ++ pinned plus the one grace that covers the rest.""" ++ scheduler, manager, pool, a, a4, r = _proof_pin_scene() ++ # H resends A's prompt: 4 cached blocks, the drop returns 3 and the 4th ++ # is the proof block. Admission needs 2 new blocks + the 3 free hit ++ # blocks against 4 free: short by 1 before the pin, by 2 after it (the ++ # pin took the proof block too). The proof block comes back first, B's ++ # grace covers the other one: H is admitted on its 48-token hit. Before ++ # the recompute B's grace yielded against the stale shortfall of 1 and ++ # H still waited. ++ h = _sched_request("H", list(a.prompt_token_ids), max_tokens=2) ++ scheduler.add_request(h) ++ out = _sched_step(scheduler) ++ assert out.num_scheduled_tokens["H"] == 32 ++ assert not manager.pinned_hit_blocks("H") and not manager._grace_pins ++ assert pool.get_num_free_blocks() == 0 and r.num_preemptions == 0 ++ while not (r.is_finished() and h.is_finished()): ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 99 ++ ++ ++def test_scheduler_head_waiter_keeps_every_pin_when_nothing_covers_it(): ++ """Same scene, H one block longer: 3 new blocks needed, short by 3 after ++ the pin. Its proof block and B's grace together give 2: nothing is ++ released, nothing yields. B keeps its grace, H keeps its whole pin, and ++ once R has finished H is admitted on the hit the pin protected.""" ++ scheduler, manager, pool, a, a4, r = _proof_pin_scene() ++ h = _sched_request("H", list(a.prompt_token_ids) + list(range(5000, 5016)), 2) ++ scheduler.add_request(h) ++ out = _sched_step(scheduler) ++ assert "H" not in out.num_scheduled_tokens ++ assert manager.grace_pinned_blocks("B"), "B's grace yielded for nothing" ++ assert len(manager.pinned_hit_blocks("H")) == 4 and a4.ref_cnt == 1 ++ assert pool.get_num_free_blocks() == 0 and r.num_preemptions == 0 ++ while not r.is_finished(): ++ _sched_step(scheduler) ++ out = _sched_step(scheduler) # R's grace covers it now; the oldest yields ++ assert out.num_scheduled_tokens["H"] == 96 - 48 ++ assert not manager.pinned_hit_blocks("H") ++ while not h.is_finished(): ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 99 ++ ++ ++def test_scheduler_head_waiter_gives_its_proof_block_back_at_an_exact_fit(): ++ """No running request. A finished (grace on its 3 blocks). H resends A's ++ prompt + 1535 tokens: 99 blocks, the whole usable pool. Its hit is 3 ++ blocks, the EAGLE drop returns 2 and the 3rd is pinned as the proof. H ++ needs 97 new blocks; the pool has 96 free plus the proof block, held by ++ A's grace and H's own pin. Nothing else ever yields H's own pin, so H ++ gives the proof block back itself, A's grace is then yieldable and H is ++ admitted. Before that H waited forever: shortfall 1, nothing yieldable.""" ++ scheduler = _grace_scheduler(num_blocks=100) # 99 usable ++ manager = scheduler.kv_cache_manager ++ manager.coordinator.eagle_group_ids = {0} ++ pool = manager.block_pool ++ a = _sched_request("A", list(range(1000, 1048)), max_tokens=2) ++ _sched_finish(scheduler, a) ++ assert pool.get_num_free_blocks() == 96 and manager.num_grace_pinned_blocks == 3 ++ ++ # One output token: the prompt's last block has room for it, so the ++ # whole request fits the pool exactly. ++ h = _sched_request("H", list(a.prompt_token_ids) + list(range(5000, 6535)), 1) ++ assert len(h.block_hashes) == 98 and h.num_prompt_tokens == 1583 ++ scheduler.add_request(h) ++ out = _sched_step(scheduler) ++ assert out.num_scheduled_tokens.get("H") == 1583 - 32, ( ++ f"H not admitted: free={pool.get_num_free_blocks()} " ++ f"own pin={len(manager.pinned_hit_blocks('H'))} " ++ f"graces={list(manager._grace_pins)} " ++ f"shortfall={manager.last_alloc_shortfall}" ++ ) ++ assert h.is_finished() and not manager.pinned_hit_blocks("H") ++ assert list(manager._grace_pins) == ["H"] # A's yielded, H's taken ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 99 ++ ++ ++def test_scheduler_pins_never_report_a_common_prefix(): ++ """A finished on a 2-block system prefix and its grace holds it. R1 ++ (system + own) and R2 (unrelated) run. The system blocks have ref 2 = ++ R1's table + A's grace, equal to the number of running requests, which ++ the per-block count would read as a prefix common to all of them. With ++ a pin or grace held the scheduler reports no common prefix, so cascade ++ attention cannot read R2 against blocks it does not share.""" ++ scheduler = _grace_scheduler(num_blocks=64) ++ manager = scheduler.kv_cache_manager ++ system = list(range(500, 532)) ++ a = _sched_request("A", system + list(range(1000, 1064)), max_tokens=2) ++ _sched_finish(scheduler, a) # prompt 96: R1's 32-token hit is no follow-up ++ assert len(manager.grace_pinned_blocks("A")) == 6 ++ r1 = _sched_request("R1", system + list(range(2000, 2032)), max_tokens=8) ++ r2 = _sched_request("R2", list(range(3000, 3048)), max_tokens=8) ++ scheduler.add_request(r1) ++ scheduler.add_request(r2) ++ out = _sched_step(scheduler) ++ assert "R1" in out.num_scheduled_tokens and "R2" in out.num_scheduled_tokens ++ assert "A" in manager._grace_pins ++ assert out.num_common_prefix_blocks == [0], out.num_common_prefix_blocks ++ # Without any pin the count is what it was: R2 gone, R3 joins the prefix, ++ # every running request shares the system blocks. ++ from vllm.v1.request import RequestStatus ++ ++ scheduler.finish_requests("R2", RequestStatus.FINISHED_ABORTED) ++ manager.yield_pinned_hits() # A's grace and R2's ++ r3 = _sched_request("R3", system + list(range(4000, 4032)), max_tokens=8) ++ scheduler.add_request(r3) ++ out = _sched_step(scheduler) ++ assert "R3" in out.num_scheduled_tokens and not manager._grace_pins ++ assert out.num_common_prefix_blocks == [2] ++ while not (r1.is_finished() and r3.is_finished()): ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ assert manager.block_pool.get_num_free_blocks() == 63 ++ ++ ++def test_release_pinned_extras_drops_only_the_blocks_outside_the_hit(): ++ """A waiter's pin = its hit + the proof block past it. Only the proof ++ block counts as yieldable for the waiter itself, and only when its hit is ++ given; releasing the extras drops exactly that block, once, and leaves a ++ block another holder still references out of the free queue.""" ++ scheduler = _grace_scheduler(num_blocks=100) ++ manager = scheduler.kv_cache_manager ++ manager.coordinator.eagle_group_ids = {0} ++ pool = manager.block_pool ++ a = _sched_request("A", list(range(1000, 1048)), max_tokens=2) ++ _sched_finish(scheduler, a) # grace on A1..A3 ++ h = _sched_request("H", list(a.prompt_token_ids) + list(range(5000, 5064)), 2) ++ hit, num_hit, _ = manager.get_computed_blocks(h) ++ assert num_hit == 32 and len(hit.blocks[0]) == 2 # A3 dropped ++ assert manager.pin_waiting_hit(h) == 3 # A1, A2 + the proof A3 ++ [a3] = pool.get_cached_block(a.block_hashes[2], [0]) ++ assert a3.ref_cnt == 2 # A's grace + H's pin ++ assert manager.num_yieldable_blocks(keep="H") == 0 ++ assert manager.num_yieldable_blocks(keep="H", keep_hit=hit.blocks) == 1 ++ assert manager.release_pinned_extras("H", hit.blocks) == 1 ++ assert manager.release_pinned_extras("H", hit.blocks) == 0 ++ assert [b.block_id for b in manager.pinned_hit_blocks("H")] == [ ++ b.block_id for b in hit.blocks[0] ++ ] ++ assert a3.ref_cnt == 1 and pool.get_num_free_blocks() == 96 # A's grace holds it ++ assert manager.num_yieldable_blocks(keep="H") == 1 # now A's grace can go ++ manager.yield_pinned_hits(keep="H") ++ assert a3.ref_cnt == 0 and pool.get_num_free_blocks() == 97 ++ manager.release_pinned_hit("H") ++ assert pool.get_num_free_blocks() == 99 ++ ++ ++def _check_pool_invariants(scheduler) -> None: ++ """Every block's refs are accounted for by a table, a deferred free, a ++ waiter pin or a grace; a referenced block is never in the free queue; ++ every ref-0 block is in it; nothing is pinned twice by one holder.""" ++ manager = scheduler.kv_cache_manager ++ pool = manager.block_pool ++ expected: dict[int, int] = {} ++ holders = [ ++ blocks ++ for group in manager.coordinator.single_type_managers ++ for blocks in group.req_to_blocks.values() ++ ] ++ holders += [blocks for _, blocks in scheduler.deferred_frees] ++ holders += list(manager._pinned_hits.values()) ++ holders += list(manager._grace_pins.values()) ++ for blocks in holders: ++ seen: set[int] = set() ++ for block in blocks: ++ if block.is_null: ++ continue ++ assert block.block_id not in seen, f"block {block.block_id} held twice" ++ seen.add(block.block_id) ++ expected[block.block_id] = expected.get(block.block_id, 0) + 1 ++ free_ids = _free_block_ids(manager) ++ num_ref0 = 0 ++ for block in pool.blocks: ++ if block.is_null: ++ continue ++ assert block.ref_cnt == expected.get(block.block_id, 0), ( ++ f"block {block.block_id}: ref {block.ref_cnt}, " ++ f"holders {expected.get(block.block_id, 0)}" ++ ) ++ if block.ref_cnt == 0: ++ num_ref0 += 1 ++ assert block.block_id in free_ids, f"block {block.block_id} leaked" ++ else: ++ assert block.block_id not in free_ids, f"block {block.block_id} in queue" ++ assert num_ref0 == pool.get_num_free_blocks() == len(free_ids) ++ ++ ++def test_scheduler_random_mix_keeps_pool_invariants_and_progress(): ++ """Seeded random request mixes on small pools (fresh, shared-prefix, ++ follow-up and exact-fit prompts, aborts, finishes, the EAGLE drop on or ++ off). After every step the pool accounting closes, and a head waiter ++ that fits an idle pool is admitted within three passes.""" ++ import random ++ ++ from vllm.v1.request import RequestStatus ++ ++ for seed in range(200): ++ rng = random.Random(seed) ++ num_blocks = rng.randint(16, 40) ++ usable = num_blocks - 1 ++ scheduler = _grace_scheduler(num_blocks=num_blocks) ++ manager = scheduler.kv_cache_manager ++ pool = manager.block_pool ++ if rng.random() < 0.7: ++ manager.coordinator.eagle_group_ids = {0} ++ prompts: list[list[int]] = [] # every prompt seen, for prefixes ++ finished_prompts: list[list[int]] = [] # with their outputs ++ live: dict[str, object] = {} ++ next_id = 0 ++ stuck = 0 ++ for _ in range(80): ++ if len(live) < 4 and rng.random() < 0.6: ++ kind = rng.choice(["fresh", "prefix", "follow", "exact", "exact"]) ++ max_tokens = rng.randint(1, 6) ++ if kind == "fresh" or not prompts: ++ n = rng.randint(1, max(1, usable * 16 - max_tokens)) ++ tokens = [rng.randint(1, 50000) for _ in range(n)] ++ elif kind == "prefix": ++ base = rng.choice(prompts) ++ cut = rng.randint(1, len(base)) ++ tokens = base[:cut] + [ ++ rng.randint(1, 50000) for _ in range(rng.randint(1, 48)) ++ ] ++ elif kind == "follow" and finished_prompts: ++ tokens = rng.choice(finished_prompts) + [ ++ rng.randint(1, 50000) for _ in range(rng.randint(1, 64)) ++ ] ++ else: # exact: a cached prompt extended to fill the pool ++ base = rng.choice(finished_prompts or prompts) ++ max_tokens = 1 ++ room = usable * 16 - 1 - len(base) ++ tokens = base + [rng.randint(1, 50000) for _ in range(max(room, 1))] ++ limit = usable * 16 - max_tokens ++ if len(tokens) > limit: ++ tokens = tokens[:limit] ++ if not tokens: ++ continue ++ req = _sched_request(f"q{seed}-{next_id}", tokens, max_tokens) ++ next_id += 1 ++ prompts.append(list(tokens)) ++ scheduler.add_request(req) ++ live[req.request_id] = req ++ if live and rng.random() < 0.08: ++ victim = rng.choice(list(live)) ++ scheduler.finish_requests(victim, RequestStatus.FINISHED_ABORTED) ++ idle_head = not scheduler.running and len(scheduler.waiting) > 0 ++ head = scheduler.waiting.peek_request() if idle_head else None ++ out = _sched_step(scheduler) ++ _check_pool_invariants(scheduler) ++ if head is not None: ++ need = -(-(head.num_prompt_tokens + head.max_tokens) // 16) ++ if need <= usable and head.request_id not in out.num_scheduled_tokens: ++ stuck += 1 ++ assert stuck < 3, ( ++ f"seed {seed}: head {head.request_id} ({need} blocks of " ++ f"{usable}) not admitted on an idle pool: free " ++ f"{pool.get_num_free_blocks()} own pin " ++ f"{len(manager.pinned_hit_blocks(head.request_id))} " ++ f"graces {manager.num_grace_pinned_blocks} shortfall " ++ f"{manager.last_alloc_shortfall}" ++ ) ++ else: ++ stuck = 0 ++ for request_id, req in list(live.items()): ++ if req.is_finished(): ++ if req.num_output_tokens > 0: ++ finished_prompts.append(list(req.all_token_ids)) ++ del live[request_id] ++ for request_id in list(live): ++ scheduler.finish_requests(request_id, RequestStatus.FINISHED_ABORTED) ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ _check_pool_invariants(scheduler) ++ assert pool.get_num_free_blocks() == usable, f"seed {seed}: leak" ++ ++ ++def test_scheduler_follow_up_is_admitted_on_the_grace_and_takes_it_over(): ++ """A's next turn is admitted on A's grace and takes it over (no double ++ ref). A new conversation that only shares the system prompt with the ++ finished turns A and B releases neither grace.""" ++ scheduler = _grace_scheduler(num_blocks=32) ++ manager = scheduler.kv_cache_manager ++ pool = manager.block_pool ++ system = list(range(500, 532)) # 2 shared blocks ++ a = _sched_request("A", system + list(range(1000, 1048)), max_tokens=2) ++ _sched_finish(scheduler, a) ++ b = _sched_request("B", system + list(range(2000, 2048)), max_tokens=2) ++ _sched_finish(scheduler, b) # B's admission on the system blocks keeps A's ++ grace = manager.grace_pinned_blocks("A") ++ assert len(grace) == 5 and len(manager.grace_pinned_blocks("B")) == 5 ++ ++ c = _sched_request("C", system + list(range(3000, 3040)), max_tokens=2) ++ scheduler.add_request(c) ++ out = _sched_step(scheduler) ++ assert out.num_scheduled_tokens["C"] == c.num_prompt_tokens - 32 # system hit ++ assert set(manager._grace_pins) == {"A", "B"} ++ ++ a2 = _sched_request("A2", list(a.all_token_ids) + list(range(4000, 4030)), 2) ++ scheduler.add_request(a2) ++ out = _sched_step(scheduler) ++ assert out.num_scheduled_tokens["A2"] == a2.num_prompt_tokens - 80 # hit 80 ++ assert "A" not in manager._grace_pins and "B" in manager._grace_pins ++ shared = {id(blk) for blk in manager.grace_pinned_blocks("B")} ++ for blk in grace: # A2's table (+ C's table or grace, B's grace: system) ++ assert blk.ref_cnt == 1 + 2 * (id(blk) in shared) ++ while not (a2.is_finished() and c.is_finished()): ++ _sched_step(scheduler) ++ manager.yield_pinned_hits() ++ assert pool.get_num_free_blocks() == 31 +diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py +index 3dbda90..c0bc6c2 100644 +--- a/vllm/v1/core/kv_cache_manager.py ++++ b/vllm/v1/core/kv_cache_manager.py +@@ -240,6 +240,37 @@ class KVCacheManager: + # whenever a running request would otherwise fail to allocate + # (``yield_pinned_hits``), so a pin never causes a preemption. + self._pinned_hits: dict[str, list[KVCacheBlock]] = {} ++ # A finished request's hashed blocks (its attention prefix and ++ # the GDN states it holds) stay referenced as a "grace" pin until the ++ # conversation's next turn is admitted on them. With waiter pins alone, ++ # every remaining total miss had lost its top attention block BEFORE ++ # the next turn was first looked up -- in the gap between the ++ # previous turn's free and the next turn's arrival, or while its ++ # admission kept failing -- when nothing referenced the prefix and ++ # the running set's allocations popped it; once the hybrid hit is 0 ++ # a waiter pin has nothing to pin. Grace pins yield first (oldest ++ # first) whenever anyone needs blocks, so they never cost a running ++ # request or an admission anything a free queue would have given. ++ # A grace is one extra ref per block, taken on whichever free path ++ # the request takes (immediate or deferred behind an in-flight step). ++ self._grace_pins: dict[str, list[KVCacheBlock]] = {} ++ self._grace_block_ids: dict[str, frozenset[int]] = {} ++ self._grace_prompt_tokens: dict[str, int] = {} ++ # Only for fine-grained hybrid hits (align Mamba hashed finer than its ++ # block): there the next turn's key sits in the top attention block ++ # and one lost block is a total miss; at a block-sized unit the top ++ # block is a spare and nothing changes. The scheduler also turns ++ # graces off when a KV connector is configured (as for waiter pins). ++ self.enable_grace_pins = ( ++ self.enable_caching and self.coordinator.enable_partial_hash_hits ++ ) ++ # Shortfall (blocks) of the last allocate_slots that returned None. ++ self.last_alloc_shortfall = 0 ++ # A producer that moves its live GDN tail's hash to a CoW copy keeps ++ # writing the source; pins and graces follow the hash to the copy. ++ for manager in self.coordinator.single_type_managers: ++ if isinstance(manager, MambaManager): ++ manager.moved_hash_listeners.append(self._pins_follow_moved_hash) + + @property + def usage(self) -> float: +@@ -536,6 +567,9 @@ class KVCacheManager: + ) + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > self.block_pool.get_num_free_blocks(): ++ self.last_alloc_shortfall = ( ++ required_blocks - self.block_pool.get_num_free_blocks() ++ ) + return None + + num_tokens_main_model = total_computed_tokens + num_new_tokens +@@ -575,7 +609,9 @@ class KVCacheManager: + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > available_blocks: + # Cannot allocate new blocks ++ self.last_alloc_shortfall = required_blocks - available_blocks + return None ++ self.last_alloc_shortfall = 0 + + if ( + new_computed_block_list is not self.empty_kv_cache_blocks.blocks +@@ -601,6 +637,10 @@ class KVCacheManager: + # remote. Update state for locally cached blocks. + if delay_cache_blocks: + self.release_pinned_hit(request.request_id) ++ if self._grace_pins and num_new_computed_tokens > 0: ++ self._release_grace_for_hit( ++ new_computed_block_list, num_local_computed_tokens ++ ) + return self.create_kv_cache_blocks(new_blocks) + + # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens +@@ -614,8 +654,13 @@ class KVCacheManager: + ) + self.coordinator.cache_blocks(request, num_tokens_to_cache) + +- # Admitted: the table now references the hit, the pin is redundant. ++ # Admitted: the table now references the hit, the pins are redundant. + self.release_pinned_hit(request.request_id) ++ if self._grace_pins and num_new_computed_tokens > 0: ++ # ... and so does the grace pin of the turn it resumes. ++ self._release_grace_for_hit( ++ new_computed_block_list, num_local_computed_tokens ++ ) + return self.create_kv_cache_blocks(new_blocks) + + def free(self, request: Request) -> None: +@@ -627,7 +672,220 @@ class KVCacheManager: + request: The request to free the blocks. + """ + self.release_pinned_hit(request.request_id) +- self.coordinator.free(request.request_id) ++ if not self._grace_eligible(request): ++ self.coordinator.free(request.request_id) ++ return ++ # Same blocks and order as coordinator.free (the deferred path frees ++ # this list as-is too); the grace takes its ref before the table's ++ # goes, so the hashed blocks never pass through the free queue. ++ blocks = self.coordinator.pop_blocks_for_free(request.request_id) ++ self._take_grace_pin(request, blocks) ++ self.block_pool.free_blocks(blocks) ++ ++ # Finished requests whose hashed blocks stay pinned for the next turn ++ # (oldest released first beyond this). One per conversation is the ++ # need; twice the deployed max_num_seqs. ++ MAX_GRACE_PINS = 16 ++ ++ @property ++ def num_grace_pins(self) -> int: ++ return len(self._grace_pins) ++ ++ @property ++ def num_grace_pinned_blocks(self) -> int: ++ return sum(len(blocks) for blocks in self._grace_pins.values()) ++ ++ def grace_pinned_blocks(self, request_id: str) -> list[KVCacheBlock]: ++ return list(self._grace_pins.get(request_id, ())) ++ ++ def _grace_eligible(self, request: Request) -> bool: ++ """A turn that finished with output (stop, length, repetition, or ++ an abort after decoding) may have a next turn that resumes it. A ++ preemption (still RUNNING), an aborted waiter (no output), an error ++ or an ignored request leaves nothing worth holding.""" ++ return ( ++ self.enable_grace_pins ++ and self.enable_caching ++ and self.prefix_cache_lookup_enabled(request) ++ and request.is_finished() ++ and request.status ++ not in (RequestStatus.FINISHED_ERROR, RequestStatus.FINISHED_IGNORED) ++ and request.num_output_tokens > 0 ++ ) ++ ++ def _take_grace_pin(self, request: Request, blocks: list[KVCacheBlock]) -> int: ++ """Take one extra ref on every hashed block of a finished request's ++ just-popped blocks (``blocks`` is in eviction order: each group's ++ table tail first, the held GDN states last). The caller then drops ++ the table refs, now or after the in-flight fence; the hashed blocks ++ stay out of the free queue until the grace is released. Unhashed ++ blocks (running state, MTP scratch, an unkeyed last block) go back ++ with the table refs as before. Returns the number of blocks held.""" ++ keep: list[KVCacheBlock] = [] ++ seen: set[int] = set() ++ for block in blocks: ++ if block.is_null or block.block_hash is None or id(block) in seen: ++ continue ++ # Still referenced by the request's table or hold: the grace ref ++ # never moves a block in or out of a free queue (any pool). ++ assert block.ref_cnt > 0 ++ seen.add(id(block)) ++ keep.append(block) ++ if not keep: ++ return 0 ++ request_id = request.request_id ++ self.release_grace_pin(request_id) # a repeated free: replace ++ for block in keep: ++ block.ref_cnt += 1 ++ self._grace_pins[request_id] = keep ++ self._grace_block_ids[request_id] = frozenset(seen) ++ self._grace_prompt_tokens[request_id] = request.num_prompt_tokens ++ while len(self._grace_pins) > self.MAX_GRACE_PINS: ++ self.release_grace_pin(next(iter(self._grace_pins))) ++ return len(keep) ++ ++ def release_grace_pin(self, request_id: str) -> int: ++ """Return a finished request's grace-pinned blocks to the pool (in ++ the eviction order a free would have used). Returns the count.""" ++ blocks = self._grace_pins.pop(request_id, None) ++ self._grace_block_ids.pop(request_id, None) ++ self._grace_prompt_tokens.pop(request_id, None) ++ if not blocks: ++ return 0 ++ self.block_pool.free_blocks(blocks) ++ return len(blocks) ++ ++ def _release_grace_for_hit( ++ self, ++ computed_blocks: tuple[Sequence[KVCacheBlock], ...], ++ num_hit_tokens: int, ++ ) -> None: ++ """An admitted hit now references its blocks through the table, so ++ the grace of the turn it resumes has done its job. That turn is a ++ grace that holds the deepest hit block of every group whose deepest ++ hit block any grace holds (the GDN state it resumes from, the ++ attention block holding the key), and whose prompt the hit reaches ++ to within two prefix units (the EAGLE-dropped tail key). A new turn ++ that only shares a system prompt with finished conversations ++ releases none of theirs.""" ++ candidates: set[str] | None = None ++ for group in computed_blocks: ++ deepest = next((b for b in reversed(group) if not b.is_null), None) ++ if deepest is None: ++ continue ++ owners = { ++ request_id ++ for request_id, block_ids in self._grace_block_ids.items() ++ if id(deepest) in block_ids ++ } ++ if owners: ++ candidates = owners if candidates is None else candidates & owners ++ if not candidates: ++ return ++ reach = num_hit_tokens + 2 * self.block_pool.hash_block_size ++ for request_id in candidates: ++ if reach >= self._grace_prompt_tokens.get(request_id, 0): ++ self.release_grace_pin(request_id) ++ ++ def _pins_follow_moved_hash( ++ self, source_block: KVCacheBlock, cow_block: KVCacheBlock ++ ) -> None: ++ """A Mamba producer moved ``source_block``'s hashes to ``cow_block`` ++ and keeps writing ``source_block``. A grace or waiter pin on the ++ source follows to the copy (as a hold does), so it protects the ++ hashed state and never leaves an extra ref on a block its owner must ++ hold exclusively (speculative relocation asserts on that).""" ++ for request_id, blocks in self._grace_pins.items(): ++ if id(source_block) in self._grace_block_ids[request_id]: ++ self._swap_pinned_block(blocks, source_block, cow_block) ++ self._grace_block_ids[request_id] = frozenset( ++ id(block) for block in blocks ++ ) ++ for blocks in self._pinned_hits.values(): ++ self._swap_pinned_block(blocks, source_block, cow_block) ++ ++ def _swap_pinned_block( ++ self, ++ blocks: list[KVCacheBlock], ++ source_block: KVCacheBlock, ++ cow_block: KVCacheBlock, ++ ) -> None: ++ for i, block in enumerate(blocks): ++ if block is source_block: ++ if any(other is cow_block for other in blocks): ++ del blocks[i] ++ else: ++ assert cow_block.ref_cnt > 0 ++ cow_block.ref_cnt += 1 ++ blocks[i] = cow_block ++ self.block_pool.free_blocks([source_block]) ++ return ++ ++ def num_yieldable_blocks( ++ self, ++ keep: str | None = None, ++ keep_hit: tuple[Sequence[KVCacheBlock], ...] | None = None, ++ ) -> int: ++ """Blocks that would return to the free queue if every grace and ++ every waiter pin except ``keep``'s were released now. With ++ ``keep_hit`` (the hit ``keep`` is being admitted on), ``keep``'s ++ pinned blocks outside that hit (the proof block past it) count too: ++ ``release_pinned_extras`` can give them back. Blocks of that hit ++ never count: freed, they are taken straight back by the admission ++ (a hit that grew past the pin snapshot would otherwise be counted ++ through the grace that still holds it).""" ++ pins: dict[int, list] = {} ++ holders = [*self._grace_pins.values()] + [ ++ blocks ++ for request_id, blocks in self._pinned_hits.items() ++ if request_id != keep ++ ] ++ hit_ids: set[int] = set() ++ if keep is not None and keep_hit is not None: ++ holders.append(self._pinned_extras(keep, keep_hit)) ++ hit_ids = {id(block) for group in keep_hit for block in group} ++ for blocks in holders: ++ for block in blocks: ++ entry = pins.setdefault(id(block), [block, 0]) ++ entry[1] += 1 ++ return sum( ++ 1 ++ for block_id, (block, n) in pins.items() ++ if block.ref_cnt == n and block_id not in hit_ids ++ ) ++ ++ def _pinned_extras( ++ self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...] ++ ) -> list[KVCacheBlock]: ++ hit_ids = {id(block) for group in hit for block in group} ++ return [ ++ block ++ for block in self._pinned_hits.get(request_id, ()) ++ if id(block) not in hit_ids ++ ] ++ ++ def release_pinned_extras( ++ self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...] ++ ) -> int: ++ """Drop from the request's pin the blocks outside ``hit``: the proof ++ block past it, and any hit block a later lookup added. The request ++ is being admitted on ``hit`` now, so they protect nothing it still ++ needs; they cost it blocks. Returns the number dropped.""" ++ extras = self._pinned_extras(request_id, hit) ++ if not extras: ++ return 0 ++ extra_ids = {id(block) for block in extras} ++ rest = [ ++ block ++ for block in self._pinned_hits[request_id] ++ if id(block) not in extra_ids ++ ] ++ if rest: ++ self._pinned_hits[request_id] = rest ++ else: ++ del self._pinned_hits[request_id] ++ self.block_pool.free_blocks(extras) ++ return len(extras) + + # Waiting requests whose prefix-cache hit is pinned at once (arrival + # order). Pins never block anyone (they yield before a running request +@@ -653,11 +911,10 @@ class KVCacheManager: + + Known gap: a pin is a snapshot of the hit at pin time. It is not + refreshed if the hit grows while the request waits (a sibling +- computing more of the same prefix) or if a producer moves the hash of +- a pinned live block into its CoW copy (a resume at a running +- producer's tail, unreachable under the EAGLE drop); the admission +- lookup still finds whatever is cached then, so the only cost is a +- smaller hit, never a wrong one.""" ++ computing more of the same prefix); the admission lookup still finds ++ whatever is cached then, so the only cost is a smaller hit, never a ++ wrong one. A producer moving the hash of a pinned live GDN block into ++ its CoW copy is followed (``_pins_follow_moved_hash``).""" + request_id = request.request_id + if ( + request_id in self._pinned_hits +@@ -716,13 +973,21 @@ class KVCacheManager: + self, keep: str | None = None, limit: int | None = None + ) -> int: + """Drop pins (except ``keep``'s) so a running request can allocate, +- or a waiter ahead in line can be admitted: the most recently pinned +- waiters first, at most ``limit`` of them per call (all by default; +- the scheduler passes 1 and retries, so a decode step that needs a +- block costs one waiter its pin, not every waiter). Returns the number +- of blocks released; 0 means there was nothing to yield.""" ++ or a waiter ahead in line can be admitted: grace pins of finished ++ turns first (oldest first), then the most recently pinned waiters, at ++ most ``limit`` pins per call (all by default; the scheduler passes 1 ++ and retries, so a decode step that needs a block costs one pin, not ++ every pin). Returns the number of blocks released; 0 means there was ++ nothing to yield.""" + released = 0 + yielded = 0 ++ # Grace pins first (oldest finished conversation first), then the ++ # waiters' pins (most recently pinned first). ++ for request_id in list(self._grace_pins): ++ if limit is not None and yielded >= limit: ++ return released ++ released += self.release_grace_pin(request_id) ++ yielded += 1 + for request_id in reversed(list(self._pinned_hits)): + if request_id == keep: + continue +@@ -767,7 +1032,12 @@ class KVCacheManager: + # A pin is never written by an in-flight step (the waiter has no + # table), so it is released now, not with the deferred blocks. + self.release_pinned_hit(request.request_id) +- return self.coordinator.pop_blocks_for_free(request.request_id) ++ blocks = self.coordinator.pop_blocks_for_free(request.request_id) ++ # A finished turn whose last step is still in flight: the grace takes ++ # its ref now; the table refs go when the fence drains. ++ if self._grace_eligible(request): ++ self._take_grace_pin(request, blocks) ++ return blocks + + def evict_blocks(self, block_ids: set[int]) -> None: + """evict blocks from the prefix cache by their block IDs. +@@ -827,6 +1097,12 @@ class KVCacheManager: + list[int]: The number of common prefix blocks for each kv cache + group. + """ ++ if self._pinned_hits or self._grace_pins: ++ # A pinned or graced block carries a ref that is not a running ++ # request's, so the ref count can equal the running count for a ++ # block only some of them share. Report no common prefix (cascade ++ # attention is an optimisation; a wrong prefix would be read). ++ return [0] * len(self.kv_cache_config.kv_cache_groups) + return self.coordinator.get_num_common_prefix_blocks(running_request_id) + + def take_events(self) -> list[KVCacheEvent]: +diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py +index fea77b1..2ad6f1c 100644 +--- a/vllm/v1/core/sched/scheduler.py ++++ b/vllm/v1/core/sched/scheduler.py +@@ -347,6 +347,9 @@ class Scheduler(SchedulerInterface): + # Bind after construction so connectors can access the cache manager. + if self.connector is not None: + self.connector.bind_kv_cache_manager(self.kv_cache_manager) ++ # Finished turns' grace pins are off with a connector, as ++ # waiter pins are (_pin_waiting_hits). ++ self.kv_cache_manager.enable_grace_pins = False + + self._log_prefix_cache_features() + +@@ -573,6 +576,15 @@ class Scheduler(SchedulerInterface): + self.hash_block_size, + self.block_size, + ) ++ if manager.enable_grace_pins: ++ logger.info( ++ "Prefix cache: a finished turn's cached prefix stays referenced " ++ "until its next turn is admitted (up to %d turns; hash unit %d, " ++ "block %d)", ++ manager.MAX_GRACE_PINS, ++ self.hash_block_size, ++ self.block_size, ++ ) + + def _pin_waiting_hits(self) -> None: + """Pin the prefix-cache hits of the requests at the front of the +@@ -1258,12 +1270,89 @@ class Scheduler(SchedulerInterface): + has_scheduled_reqs=bool(self.running), + ) + +- if new_blocks is None and self.kv_cache_manager.yield_pinned_hits( +- keep=request_id ++ # Finished turns' grace pins (oldest first), then the ++ # pins of waiters behind this one, yield to it one at a time ++ # (arrival order). Its own hit is pinned first (a turn that ++ # arrived this step is not pinned yet and its hit may sit in ++ # a grace that yields), so only blocks it does not resume ++ # from go back. And only when releasing them all would cover ++ # the shortfall: otherwise yielding costs other conversations ++ # their prefix and still does not admit this one. ++ kv_cache_manager = self.kv_cache_manager ++ yield_for_head = new_blocks is None and ( ++ kv_cache_manager.last_alloc_shortfall > 0 ++ and ( ++ kv_cache_manager.num_grace_pins > 0 ++ or kv_cache_manager.num_pinned_waiters > 0 ++ ) ++ ) ++ if yield_for_head: ++ if ( ++ self.connector is None ++ and request.num_computed_tokens == 0 ++ and kv_cache_manager.pin_waiting_hit(request) > 0 ++ ): ++ # The pin just took blocks out of the free queue that ++ # the failed attempt counted as free: the proof block ++ # past the hit, and any hit block the lookup above did ++ # not return. Recompute the shortfall on the pinned ++ # state before measuring it against the others' pins. ++ new_blocks = self.kv_cache_manager.allocate_slots( ++ request, ++ num_new_tokens, ++ num_new_computed_tokens=num_new_local_computed_tokens, ++ new_computed_blocks=new_computed_blocks, ++ num_lookahead_tokens=effective_lookahead_tokens, ++ num_external_computed_tokens=num_external_computed_tokens, ++ delay_cache_blocks=load_kv_async, ++ num_encoder_tokens=num_encoder_tokens, ++ full_sequence_must_fit=self.scheduler_reserve_full_isl, ++ reserved_blocks=reserved_blocks, ++ has_scheduled_reqs=bool(self.running), ++ ) ++ yield_for_head = new_blocks is None and ( ++ kv_cache_manager.last_alloc_shortfall ++ <= kv_cache_manager.num_yieldable_blocks( ++ keep=request_id, keep_hit=new_computed_blocks.blocks ++ ) ++ ) ++ if ( ++ yield_for_head ++ and kv_cache_manager.release_pinned_extras( ++ request_id, new_computed_blocks.blocks ++ ) ++ > 0 ++ ): ++ # Its own pinned blocks outside the hit it is admitted ++ # on (the proof block past it) go first: they cost ++ # nobody else anything, and at an exact fit they are ++ # the blocks it needs. Retry on them before the others ++ # yield. ++ new_blocks = self.kv_cache_manager.allocate_slots( ++ request, ++ num_new_tokens, ++ num_new_computed_tokens=num_new_local_computed_tokens, ++ new_computed_blocks=new_computed_blocks, ++ num_lookahead_tokens=effective_lookahead_tokens, ++ num_external_computed_tokens=num_external_computed_tokens, ++ delay_cache_blocks=load_kv_async, ++ num_encoder_tokens=num_encoder_tokens, ++ full_sequence_must_fit=self.scheduler_reserve_full_isl, ++ reserved_blocks=reserved_blocks, ++ has_scheduled_reqs=bool(self.running), ++ ) ++ # Coverage is re-read before every yield: a retry can fail ++ # by more than the attempt before it (the full-sequence gate ++ # passes, the lookahead gate does not). ++ while ( ++ yield_for_head ++ and new_blocks is None ++ and kv_cache_manager.last_alloc_shortfall ++ <= kv_cache_manager.num_yieldable_blocks( ++ keep=request_id, keep_hit=new_computed_blocks.blocks ++ ) ++ and kv_cache_manager.yield_pinned_hits(keep=request_id, limit=1) + ): +- # Pins of waiters behind this one yield to it +- # (arrival order); its own pin stays, those blocks are +- # its hit. One retry with them back in the free queue. + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, +diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py +index cf2b716..24a72bb 100644 +--- a/vllm/v1/core/single_type_kv_cache_manager.py ++++ b/vllm/v1/core/single_type_kv_cache_manager.py +@@ -3,7 +3,7 @@ + import itertools + from abc import ABC, abstractmethod + from collections import defaultdict +-from collections.abc import Sequence ++from collections.abc import Callable, Sequence + from typing import TYPE_CHECKING, ClassVar + + from vllm import envs +@@ -1693,6 +1693,12 @@ class MambaManager(SingleTypeKVCacheManager): + # therefore hold their hashed endpoint here too: the moved tail for + # a running request, and the hit source a new request copies from. + self._retained_checkpoints: dict[str, tuple[int, KVCacheBlock]] = {} ++ # Callbacks (source_block, cow_block) run after a running ++ # request moved its live tail's hash to the CoW copy, for refs held ++ # outside this manager (the KVCacheManager's grace and waiter pins). ++ self.moved_hash_listeners: list[ ++ Callable[[KVCacheBlock, KVCacheBlock], None] ++ ] = [] + if self.mamba_cache_mode == "align" and self.enable_caching: + # The A/B evidence line: boot logs must show which arm ran. + logger.info_once( +@@ -1789,6 +1795,8 @@ class MambaManager(SingleTypeKVCacheManager): + self.block_pool.touch([cow_block]) + self._retained_checkpoints[request_id] = (position, cow_block) + self.block_pool.free_blocks([source_block]) ++ for listener in self.moved_hash_listeners: ++ listener(source_block, cow_block) + + def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]: + if self.mamba_cache_mode != "align": +-- +2.43.0 + diff --git a/patches/0098-guard_telemetry-skip-reader-polls-while-the-engine-i.patch b/patches/0098-guard_telemetry-skip-reader-polls-while-the-engine-i.patch new file mode 100644 index 0000000..0fd452d --- /dev/null +++ b/patches/0098-guard_telemetry-skip-reader-polls-while-the-engine-i.patch @@ -0,0 +1,147 @@ +From 58fc55a0d99fcea188de6045e9a41c1720fdcaa6 Mon Sep 17 00:00:00 2001 +From: Hinata +Date: Thu, 8 Oct 2026 22:19:21 +0100 +Subject: [PATCH 098/103] guard_telemetry: skip reader polls while the engine + is idle + +The bounds-guard reader issues a side-stream device->host copy of the +counters every VLLM_E2_GUARD_INTERVAL (0.05 s) forever, including when no +request is in flight. That periodic device work keeps every GPU out of its +idle P-state: on 4x RTX 3090 with GUARD=warn (the v2.5 default) all four +cards sit in P2 at 1695/9501 MHz, 450 W total, vs 86 W in P8 with GUARD=off. + +The counters can only move inside execute_model, and check_pending() runs at +the top of every execute_model, so it now bumps a step counter that the +reader uses as an activity signal. The reader keeps its cadence while steps +arrive and for VLLM_E2_GUARD_IDLE_TAIL seconds (default 2 s) after the last +one, never skips while a count increase awaits its confirming copy, and +issues no device work at all once idle. A runner that never calls +check_pending (V1) never enables the skip. VLLM_E2_GUARD_IDLE_TAIL=inf +restores always-poll; VLLM_E2_GUARD_STATS= logs polls vs idle skips. + +Measured on the same boot: idle 85-88 W (P8); under load 598-599 polls per +30 s (full cadence, 0 skips) over ~61k steps; prefill/decode and an 80-request +chat replay unchanged; 0 Xid, 0 guard warnings. +--- + vllm/v1/worker/guard_telemetry.py | 82 +++++++++++++++++++++++++++++-- + 1 file changed, 78 insertions(+), 4 deletions(-) + +diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py +index 19cf966..199c3eb 100644 +--- a/vllm/v1/worker/guard_telemetry.py ++++ b/vllm/v1/worker/guard_telemetry.py +@@ -250,6 +250,46 @@ def _parse_interval(text: str) -> float: + + _INTERVAL = _parse_interval(os.environ.get("VLLM_E2_GUARD_INTERVAL", "0.05")) + ++ ++def _parse_seconds(name: str, default: str, allow_inf: bool) -> float: ++ """A non-negative number of seconds (inf allowed when allow_inf), or ValueError at import.""" ++ text = os.environ.get(name, default) ++ try: ++ value = float(text) ++ except ValueError as exc: ++ raise ValueError(f"{name} must be seconds, got {text!r}") from exc ++ if math.isnan(value) or value < 0 or (math.isinf(value) and not allow_inf): ++ raise ValueError(f"{name} must be a number >= 0, got {text!r}") ++ return value ++ ++ ++# IDLE SKIP. The counters can only move while a guarded kernel runs, i.e. inside ++# execute_model, and check_pending() runs at the top of every execute_model -- so it ++# doubles as an activity signal (_step_seq). Once no step has started for ++# _IDLE_TAIL seconds, and no count increase is waiting for its confirming copy, the ++# reader stops issuing side-stream copies until the next step starts. ++# ++# Why: without it the reader issues a device->host copy every _INTERVAL forever, ++# idle or not, and that keeps every GPU out of its idle P-state. Measured on 4x RTX ++# 3090 with no requests in flight: guard=warn 450 W total, all four in P2 at ++# 1695/9501 MHz; guard=off (no reader) 86 W, P8 at 210/405 MHz. ++# ++# Detection is unchanged while serving: every step resets the tail. When activity ++# resumes, the first copy consumed is the one issued on the last polled tick ++# (counters are monotone, so nothing is lost) and the next tick reads fresh values: ++# at most one extra _INTERVAL of latency on the first step after an idle period. ++# The tail (default 2 s = 40 ticks at 0.05 s) covers asynchronously queued GPU work ++# of the last steps and the one-tick confirm delay. ++# ++# A runner that never calls check_pending (the V1 runner -- see WHICH RUNNER ++# above) never sets _activity_tracked, so the skip never engages there. ++# VLLM_E2_GUARD_IDLE_TAIL=inf restores the always-poll behaviour. ++# VLLM_E2_GUARD_STATS= logs polls vs idle skips at that cadence (0 = off). ++_IDLE_TAIL = _parse_seconds("VLLM_E2_GUARD_IDLE_TAIL", "2.0", allow_inf=True) ++_STATS_EVERY = _parse_seconds("VLLM_E2_GUARD_STATS", "0", allow_inf=False) ++_step_seq = 0 ++_activity_tracked = False ++ + _buffers: dict[torch.device, torch.Tensor] = {} + # Devices the reader has given up on. They stay in _buffers: see _reader. + _read_disabled: set[torch.device] = set() +@@ -535,18 +575,49 @@ def start_reader() -> None: + _reader_started = True + threading.Thread(target=_reader, name="bounds-guard-reader", daemon=True).start() + logger.info( +- "bounds-guard reader started (mode=%s, interval=%.3gs).", ++ "bounds-guard reader started (mode=%s, interval=%.3gs, idle tail=%gs).", + _MODE, + _INTERVAL, ++ _IDLE_TAIL, + ) + + + def _reader() -> None: +- """Off-engine-thread cadence reader; drain() is the only other counter read.""" ++ """Off-engine-thread cadence reader; drain() is the only other counter read. ++ ++ Idle skip: see _IDLE_TAIL. A skipped tick issues no device work at all. ++ """ ++ last_seq = -1 ++ last_active = time.monotonic() ++ polled = skipped = 0 ++ stats_t0 = last_active + while True: + try: + time.sleep(_INTERVAL) +- _poll_once() ++ now = time.monotonic() ++ seq = _step_seq ++ if seq != last_seq: ++ last_seq, last_active = seq, now ++ if ( ++ _activity_tracked ++ and not _confirm_pending ++ and now - last_active > _IDLE_TAIL ++ ): ++ skipped += 1 ++ else: ++ _poll_once() ++ polled += 1 ++ if _STATS_EVERY and now - stats_t0 >= _STATS_EVERY: ++ logger.info( ++ "bounds-guard reader: %d polls, %d idle skips in the last %.0fs " ++ "(steps seen: %d)", ++ polled, ++ skipped, ++ now - stats_t0, ++ seq, ++ ) ++ polled = skipped = 0 ++ stats_t0 = now + except Exception: # noqa: BLE001 -- the daemon itself must survive + # A reader that dies is an instrument + # that is not there, and _ALWAYS_FATAL_SITES depends on it. Arm +@@ -709,7 +780,10 @@ def check_pending() -> None: + targeted repro where losing the group is acceptable and stopping is the + point. + """ +- global _pending ++ global _pending, _step_seq, _activity_tracked ++ # Activity signal for the reader's idle skip (see _IDLE_TAIL): one int per step. ++ _step_seq += 1 ++ _activity_tracked = True + if _pending is None: + return + msg = _pending +-- +2.43.0 + diff --git a/patches/0099-qwen4_exp-mix-LocalArgmaxMixin-into-the-MTP-draft-he.patch b/patches/0099-qwen4_exp-mix-LocalArgmaxMixin-into-the-MTP-draft-he.patch new file mode 100644 index 0000000..1432b02 --- /dev/null +++ b/patches/0099-qwen4_exp-mix-LocalArgmaxMixin-into-the-MTP-draft-he.patch @@ -0,0 +1,79 @@ +From cf89be37d15d2b60d08f6be1b032f8a1dd71bce9 Mon Sep 17 00:00:00 2001 +From: Hinata +Date: Sat, 10 Oct 2026 13:35:32 +0700 +Subject: [PATCH 099/103] qwen4_exp: mix LocalArgmaxMixin into the MTP draft + head + +Qwen4ExpMTP picks each MTP draft token by materialising full-vocab logits, +even though the tree already provides the O(2 * tp_size) path for exactly +this case: LocalArgmaxMixin builds the argmax from per-rank (value, index) +pairs instead of all-gathering a vocab-sized tensor. + +The mixin needs only self.logits_processor and self.lm_head, both already +present on this class, and it is the established idiom for draft heads +(Qwen3_5MTP, the llama/qwen3 Eagle drafters, deepseek_eagle3). Wiring it in +is what makes use_local_argmax_reduction usable with this model: without +get_top_tokens the speculative-config validator refuses to start. + +Applied to both platform implementations (nvidia, amd) to keep them in step. +No functional change unless the flag is enabled. +--- + vllm/models/qwen4_exp/amd/mtp.py | 4 ++-- + vllm/models/qwen4_exp/nvidia/mtp.py | 12 ++++++++++-- + 2 files changed, 12 insertions(+), 4 deletions(-) + +diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py +index 40ee6a6..f9637ae 100644 +--- a/vllm/models/qwen4_exp/amd/mtp.py ++++ b/vllm/models/qwen4_exp/amd/mtp.py +@@ -33,7 +33,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, + ) + from vllm.model_executor.model_loader.utils import configure_quant_config +-from vllm.model_executor.models.interfaces import SupportsPP ++from vllm.model_executor.models.interfaces import LocalArgmaxMixin, SupportsPP + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + from vllm.model_executor.models.utils import ( + AutoWeightsLoader, +@@ -383,7 +383,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + "hidden_states": 0, + } + ) +-class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): ++class Qwen4ExpMTP(LocalArgmaxMixin, nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 76c0d7c..01fb69c 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -35,7 +35,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( + should_use_host_embed_table, + ) + from vllm.model_executor.model_loader.utils import configure_quant_config +-from vllm.model_executor.models.interfaces import SupportsPP, SupportsReplaySSM ++from vllm.model_executor.models.interfaces import ( ++ LocalArgmaxMixin, ++ SupportsPP, ++ SupportsReplaySSM, ++) + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + from vllm.model_executor.models.utils import ( + AutoWeightsLoader, +@@ -409,7 +413,11 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + } + ) + class Qwen4ExpMTP( +- nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts, SupportsReplaySSM ++ LocalArgmaxMixin, ++ nn.Module, ++ SupportsPP, ++ Qwen4ExpMixtureOfExperts, ++ SupportsReplaySSM, + ): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], +-- +2.43.0 + diff --git a/patches/0100-Register-the-guard-idle-tail-and-stats-knobs-as-runt.patch b/patches/0100-Register-the-guard-idle-tail-and-stats-knobs-as-runt.patch new file mode 100644 index 0000000..a410232 --- /dev/null +++ b/patches/0100-Register-the-guard-idle-tail-and-stats-knobs-as-runt.patch @@ -0,0 +1,68 @@ +From 4854c9c4cb90d8381b8454dc03a4224b27badae5 Mon Sep 17 00:00:00 2001 +From: Hinata +Date: Sat, 10 Oct 2026 13:30:25 +0700 +Subject: [PATCH 100/103] Register the guard idle-tail and stats knobs as + runtime-only + +Patch 0098 reads VLLM_E2_GUARD_IDLE_TAIL and VLLM_E2_GUARD_STATS straight +from os.environ in guard_telemetry.py. Every fork knob must be classified, +and both are reader-thread cadence: no kernel argument and no captured +graph depends on them, so they go on the runtime-only deny-list in +environment_variables and in the ignored table in tests/test_envs.py. +--- + tests/test_envs.py | 4 ++++ + vllm/envs.py | 9 +++++++++ + 2 files changed, 13 insertions(+) + +diff --git a/tests/test_envs.py b/tests/test_envs.py +index 4871fda..e87d0d4 100644 +--- a/tests/test_envs.py ++++ b/tests/test_envs.py +@@ -636,6 +636,10 @@ _FORK_IGNORED = [ + # INERT: no read site on the 0.30 base, so not a cache factor. + ("VLLM_PLE_FP8_GLOBAL_SCALE", "1", True), + ("VLLM_E2_GUARD_INTERVAL", "0.25", "0.25"), ++ # Same reader thread: the idle tail ("inf" = always poll) and the stats ++ # cadence only change when the side reader polls and how often it logs. ++ ("VLLM_E2_GUARD_IDLE_TAIL", "inf", "inf"), ++ ("VLLM_E2_GUARD_STATS", "10", "10"), + ("VLLM_E2_MEMSTATS", "2", "2"), + ("VLLM_E2_MEMSTATS_DIR", "/var/tmp/memstats", "/var/tmp/memstats"), + ("VLLM_E2_SCHED_LOG", "/var/tmp/sched.jsonl", "/var/tmp/sched.jsonl"), +diff --git a/vllm/envs.py b/vllm/envs.py +index 4a40974..4b8dcf6 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -404,6 +404,9 @@ if TYPE_CHECKING: + # arguments and the captured graph do not depend on either. + VLLM_E2_GUARD_MODE: str = "warn" + VLLM_E2_GUARD_INTERVAL: str = "0.05" ++ # Reader-thread idle tail ("inf" = always poll) and stats log cadence. ++ VLLM_E2_GUARD_IDLE_TAIL: str = "2.0" ++ VLLM_E2_GUARD_STATS: str = "0" + # Debug allocator telemetry (poll seconds + output directory) and the + # scheduler allocation log path. + VLLM_E2_MEMSTATS: str = "0" +@@ -2457,6 +2460,8 @@ environment_variables: dict[str, Callable[[], Any]] = { + os.getenv("VLLM_E2_GUARD_MODE", "warn").strip().lower() + ), + "VLLM_E2_GUARD_INTERVAL": lambda: os.getenv("VLLM_E2_GUARD_INTERVAL", "0.05"), ++ "VLLM_E2_GUARD_IDLE_TAIL": lambda: os.getenv("VLLM_E2_GUARD_IDLE_TAIL", "2.0"), ++ "VLLM_E2_GUARD_STATS": lambda: os.getenv("VLLM_E2_GUARD_STATS", "0"), + "VLLM_E2_MEMSTATS": lambda: os.getenv("VLLM_E2_MEMSTATS", "0"), + "VLLM_E2_MEMSTATS_DIR": lambda: os.getenv("VLLM_E2_MEMSTATS_DIR", ""), + "VLLM_E2_SCHED_LOG": lambda: os.getenv("VLLM_E2_SCHED_LOG", ""), +@@ -2867,6 +2872,10 @@ def compile_factors() -> dict[str, object]: + # raise, and the interval is that thread's poll cadence. + "VLLM_E2_GUARD_MODE", + "VLLM_E2_GUARD_INTERVAL", ++ # Same thread: the idle tail and the stats cadence only change when the ++ # side reader polls and how often it logs. No kernel, no graph. ++ "VLLM_E2_GUARD_IDLE_TAIL", ++ "VLLM_E2_GUARD_STATS", + # Debug telemetry cadence and output locations. + "VLLM_E2_MEMSTATS", + "VLLM_E2_MEMSTATS_DIR", +-- +2.43.0 + diff --git a/patches/0101-auto_gptq-skip-the-Marlin-workspace-when-CUDA-is-not.patch b/patches/0101-auto_gptq-skip-the-Marlin-workspace-when-CUDA-is-not.patch new file mode 100644 index 0000000..fb61ec1 --- /dev/null +++ b/patches/0101-auto_gptq-skip-the-Marlin-workspace-when-CUDA-is-not.patch @@ -0,0 +1,37 @@ +From 9d428fa5e8b11873e5af5452cbad76dee8715de1 Mon Sep 17 00:00:00 2001 +From: Hinata +Date: Sat, 10 Oct 2026 17:31:17 +0000 +Subject: [PATCH 101/103] auto_gptq: skip the Marlin workspace when CUDA is not + initialized + +The PLE offload worker builds weights with CUDA forbidden; its guard +(vllm/v1/ple_offload/worker.py) refuses torch.cuda lazy init. AutoGPTQ +create_weights calls marlin_make_workspace_new for modular MoE experts, +which initializes CUDA and kills the offload process at load time. +Marlin creates the workspace lazily at apply time when the layer has +none, so the eager workspace is only needed where CUDA is live. Skip it +everywhere else; native AutoRound checkpoints then load in the CUDA-free +worker. +--- + vllm/model_executor/layers/quantization/auto_gptq.py | 5 ++++- + 1 file changed, 4 insertions(+), 1 deletion(-) + +diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py +index 893f85b..770b25e 100644 +--- a/vllm/model_executor/layers/quantization/auto_gptq.py ++++ b/vllm/model_executor/layers/quantization/auto_gptq.py +@@ -604,7 +604,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): + self.experts_cls, FusedMoEExpertsModular + ): + device = layer.w13_qweight.device +- layer.workspace = marlin_make_workspace_new(device, 4) ++ # PLE offload workers run CUDA-free; the Marlin kernel creates the ++ # workspace lazily at apply time when it is missing. ++ if device.type == "cuda" and torch.cuda.is_initialized(): ++ layer.workspace = marlin_make_workspace_new(device, 4) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + def replace_or_register(name: str, val: torch.Tensor | None): +-- +2.43.0 + diff --git a/patches/0102-qwen4_exp-add-QWEN38_ABLITERATION-runtime-refusal-di.patch b/patches/0102-qwen4_exp-add-QWEN38_ABLITERATION-runtime-refusal-di.patch new file mode 100644 index 0000000..2202eff --- /dev/null +++ b/patches/0102-qwen4_exp-add-QWEN38_ABLITERATION-runtime-refusal-di.patch @@ -0,0 +1,220 @@ +From c0d73eb28886bf2119856a9e66637a8eb532906e Mon Sep 17 00:00:00 2001 +From: Hinata +Date: Sun, 11 Oct 2026 03:47:14 +0000 +Subject: [PATCH 102/103] qwen4_exp: add QWEN38_ABLITERATION runtime + refusal-direction removal + +--- + vllm/models/qwen4_exp/nvidia/abliteration.py | 88 ++++++++++++++++++++ + vllm/models/qwen4_exp/nvidia/model.py | 7 +- + vllm/models/qwen4_exp/nvidia/mtp.py | 4 +- + vllm/models/qwen4_exp/nvidia/ple_layer.py | 3 + + 4 files changed, 100 insertions(+), 2 deletions(-) + create mode 100644 vllm/models/qwen4_exp/nvidia/abliteration.py + +diff --git a/vllm/models/qwen4_exp/nvidia/abliteration.py b/vllm/models/qwen4_exp/nvidia/abliteration.py +new file mode 100644 +index 0000000..840168b +--- /dev/null ++++ b/vllm/models/qwen4_exp/nvidia/abliteration.py +@@ -0,0 +1,88 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Optional refusal-direction removal ("abliteration") applied at runtime. ++ ++``QWEN38_ABLITERATION`` names a JSON file with a ``refusal_direction`` list, or a safetensors file with a ++``refusal_direction`` tensor, of length ``hidden_size``. A bare name such as ``Qwen3.8-Flash-Next-Uncensored`` resolves to ++``/opt/qwen38/abliteration/.json``. ++With it set, every residual-stream write loses its component along the unit direction r: ++ ++ x <- x - r (r . x) ++ ++applied to the attention and linear-attention block outputs, the MoE block output (routed plus shared experts), ++the PLE value projection and the text token embeddings. For one direction this equals the weight edit ++``W' = W - r (r^T W)`` on ``self_attn.o_proj``, ``linear_attn.out_proj``, every expert's and the shared expert's ++``down_proj``, ``ple.value_proj`` and the rows of ``embed_tokens`` (Arditi et al., 2024): the routed and shared ++expert outputs are summed with scalar weights before the residual write, so projecting their sum is the same as ++editing each down projection. The INT4 experts and every other weight stay untouched. Unset, nothing is loaded and ++no work is added. ++""" ++ ++import json ++import math ++import os ++ ++import torch ++ ++ENV = "QWEN38_ABLITERATION" ++BUILTIN_DIR = "/opt/qwen38/abliteration" ++ ++ ++def direction_path() -> str | None: ++ value = os.environ.get(ENV, "").strip() ++ if value in ("", "0"): ++ return None ++ if os.sep not in value and not value.endswith((".json", ".safetensors")): ++ value = os.path.join(BUILTIN_DIR, f"{value}.json") ++ return value ++ ++ ++def load_direction(hidden_size: int) -> torch.Tensor | None: ++ """The unit refusal direction as a float32 CPU tensor, or None when the switch is off.""" ++ path = direction_path() ++ if path is None: ++ return None ++ if path.endswith(".json"): ++ with open(path) as f: ++ r = torch.tensor(json.load(f)["refusal_direction"], dtype=torch.float32) ++ else: ++ from safetensors import safe_open ++ ++ with safe_open(path, framework="pt", device="cpu") as f: ++ r = f.get_tensor("refusal_direction") ++ r = r.to(torch.float32).flatten() ++ if r.numel() != hidden_size: ++ raise ValueError(f"{ENV}: {path} holds {r.numel()} values; the model's hidden size is {hidden_size}") ++ norm = float(r.norm()) ++ if not math.isfinite(norm) or norm == 0.0: ++ raise ValueError(f"{ENV}: {path} holds a zero or non-finite direction") ++ return r / norm ++ ++ ++def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None: ++ """Attach the direction as a non-persistent buffer ``_abliteration_r`` (None when the switch is off). ++ ++ Call inside the model constructor so the buffer lands on the model's device before CUDA-graph capture. The ++ CPU-only PLE worker process never runs these forwards and gets no buffer. ++ """ ++ from vllm.model_executor.layers.ple_offload_layer import is_offload_process ++ ++ r = None if is_offload_process() else load_direction(hidden_size) ++ if r is None: ++ module._abliteration_r = None ++ return ++ buf = torch.empty(hidden_size, dtype=dtype) ++ buf.copy_(r) ++ module.register_buffer("_abliteration_r", buf, persistent=False) ++ ++ ++def project_(x: torch.Tensor, r: torch.Tensor | None) -> torch.Tensor: ++ """In place: remove the component along r from x's last dimension. No-op when r is None.""" ++ if r is None: ++ return x ++ width = x.shape[-1] ++ if x.is_contiguous(): ++ flat = x.view(-1, width) ++ flat.addr_(flat @ r, r, alpha=-1.0) # rank-1 update, no full-size temporary ++ else: ++ x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape)) ++ return x +diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py +index fe416d4..b4fe373 100644 +--- a/vllm/models/qwen4_exp/nvidia/model.py ++++ b/vllm/models/qwen4_exp/nvidia/model.py +@@ -89,6 +89,7 @@ from vllm.v1.e2_guards import ( + from vllm.v1.kv_cache_interface import MambaSpec + + from ..config import Qwen4ExpConfig ++from . import abliteration + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm + from .ple_layer import Qwen4ExpPLELayer +@@ -384,6 +385,7 @@ class Qwen4ExpDecoderLayer(nn.Module): + hc_config, + prefix=maybe_prefix(prefix, "mlp_hyper_connection"), + ) ++ abliteration.register(self, config.hidden_size, model_config.dtype) + + def forward( + self, +@@ -434,12 +436,14 @@ class Qwen4ExpDecoderLayer(nn.Module): + ) + else: + raise ValueError("Invalid layer_type") ++ abliteration.project_(attn_out, self._abliteration_r) + + mlp_hc = self.mlp_hyper_connection + hidden_states, block_input, injection = mlp_hc.combine_and_mix( + hidden_states, attn_out, injection + ) + mlp_out = self.mlp(block_input) ++ abliteration.project_(mlp_out, self._abliteration_r) + return hidden_states, mlp_out, injection + + +@@ -554,6 +558,7 @@ class Qwen4ExpModel(nn.Module): + self.embed_tokens = embedding_cls(self.vocab_size, config.hidden_size) + else: + self.embed_tokens = PPMissingLayer() ++ abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype) + + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: + layer_idx = extract_layer_index(prefix) +@@ -615,7 +620,7 @@ class Qwen4ExpModel(nn.Module): + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: +- return self.embed_tokens(input_ids) ++ return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r) + + @staticmethod + def _start_layer_ple_prefetch( +diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py +index 01fb69c..39a0ff4 100644 +--- a/vllm/models/qwen4_exp/nvidia/mtp.py ++++ b/vllm/models/qwen4_exp/nvidia/mtp.py +@@ -60,6 +60,7 @@ from vllm.v1.e2_guards import ( + assert_ple_mode, + ) + ++from . import abliteration + from .hyperconnection import GatedResidual, HyperConnectionConfig + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm + from .model import ( +@@ -220,6 +221,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + else VocabParallelEmbedding + ) + self.embed_tokens = embedding_cls(self.vocab_size, self.hidden_size) ++ abliteration.register(self, self.hidden_size, model_config.dtype) + draft_vllm_config = _make_draft_vllm_config( + vllm_config, + self.mtp_start_layer_idx, +@@ -312,7 +314,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): + buffer[:num_rows].copy_(selected) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: +- return self.embed_tokens(input_ids) ++ return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r) + + def forward( + self, +diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py +index af8f528..443e9dd 100644 +--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py ++++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py +@@ -28,6 +28,7 @@ from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadata, + ) + ++from . import abliteration + from .ngram_embedding import Qwen4ExpNGramEmbedding + from .ops.ple import ple_conv, ple_gate + +@@ -136,6 +137,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + prefix=f"{prefix}.kv_proj", + disable_tp=True, + ) ++ abliteration.register(self, self.hidden_size, model_config.dtype) + norm_args = ( + self.hc_hidden_size, + config.rms_norm_eps, +@@ -449,6 +451,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): + embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) + kv, _ = self.kv_proj(embeddings) + key, value = kv.split(self.kv_proj.output_sizes, dim=-1) ++ abliteration.project_(value, self._abliteration_r) + gated_output, conv_input = ple_gate( + key, + value, +-- +2.43.0 + diff --git a/patches/0103-qwen4_exp-abliteration-always-registers-the-buffer-z.patch b/patches/0103-qwen4_exp-abliteration-always-registers-the-buffer-z.patch new file mode 100644 index 0000000..b838da3 --- /dev/null +++ b/patches/0103-qwen4_exp-abliteration-always-registers-the-buffer-z.patch @@ -0,0 +1,1168 @@ +From 975db60922a1f8d0a38752724c13f09b6d229c1c Mon Sep 17 00:00:00 2001 +From: MOVZX +Date: Sun, 11 Oct 2026 05:04:35 +0000 +Subject: [PATCH 103/103] qwen4_exp: abliteration always registers the buffer + (zeros when off) + +register() used to set a plain attribute None when QWEN38_ABLITERATION was +off. The AOT compile cache keeps graphs compiled with the switch on, and +those graphs resolve _abliteration_r as a registered buffer, so an off boot +on a warm cache died with KeyError: '_abliteration_r'. Register the buffer +unconditionally, filled with zeros when the switch is off: projecting along +a zero direction is an exact no-op, ON and OFF compile to the same graph, +and toggling the switch never invalidates the compile cache. + +Co-Authored-By: Hinata +--- + .build-products.sha256 | 1103 ++++++++++++++++++ + vllm/models/qwen4_exp/nvidia/abliteration.py | 14 +- + 2 files changed, 1109 insertions(+), 8 deletions(-) + create mode 100644 .build-products.sha256 + +diff --git a/.build-products.sha256 b/.build-products.sha256 +new file mode 100644 +index 0000000..f16f6c0 +--- /dev/null ++++ b/.build-products.sha256 +@@ -0,0 +1,1103 @@ ++e7f29df31aac49fafaea1bb0adbf83a5367b462d6749e1daf4597fee72d04289 vllm/_C_stable_libtorch.abi3.so ++9fee3d26b1f3196ec86b9ad2ccd229d6548974abb542a0c457f67a29ff050f42 vllm/_moe_C_stable_libtorch.abi3.so ++b08c79bd5c8dcab59c608617c74dcd10bd43bc26f8695421524bbab2ca371d9e vllm/cumem_allocator.abi3.so ++5e3f94012dd486f0fffd17c1f46ef03a3ea93cfd2f792cf8a0e62b980e546bfb vllm/fs_io_C.abi3.so ++b81276f64735dec9a8b5388f73e66b07abe2a139a256ed0590925450c203a6b8 vllm/spinloop.abi3.so ++75b470a05ac3baa9f80df664fa5bccdb52fba6016caed7a85f03b2c7725b0510 vllm/third_party/flashmla/flash_mla_interface.py ++2221f7d078405274da96b3aae6ab5b1126f50e93ce6f44403335e492a1e968fc vllm/third_party/fmha_sm100/__init__.py ++c4c58ed675b5906511e0c46a118af9c731e6a2197bbd4a6ba11fb1b49efb992d vllm/third_party/fmha_sm100/api.py ++bb6663b4ae1542075f799885edc18fe2f67306b30a2e9e0948cdc12761e8c023 vllm/third_party/fmha_sm100/bench_utils.py ++7dba06c345f89f4cac265434ed9b2af7716bcc599e34b8f975679beae06491a9 vllm/third_party/fmha_sm100/csrc/fmha_sm100_inst.jinja ++e1ed55f897cecfb42eada62face49c09b258d32c2d68579d5bce617b7b38a7d7 vllm/third_party/fmha_sm100/csrc/fmha_sm100_params.h ++3a93e431c54f514a9df5092c37182bf5eda3fe732dc703694bfed6ece99f3123 vllm/third_party/fmha_sm100/csrc/fmha_sm100_plan.cu ++5d7e382745068e7fd2e70a9e58719e4bde568e2d0506727f8a5de206271356a3 vllm/third_party/fmha_sm100/csrc/fmha_sm100_plan.h ++21577998a732694a4a0acca9982c366e554f84434faa0ca5dc60989ad926e7f8 vllm/third_party/fmha_sm100/csrc/fmha_sm100_reduction.cu ++8ac7601bc7d195531c800a238e914c61b8bb1997e937bce63b4e0be9c5df925f vllm/third_party/fmha_sm100/csrc/fmha_sm100_reduction.h ++a3cdbb33e90c31013089830f4568fa172e10a06f5b9098aa1029cbfe10b3c4de vllm/third_party/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja ++c6a174cd5d392bca154c17a119da4bd2ff1116371a20609b38f4eea27fa43597 vllm/third_party/fmha_sm100/csrc/gmem_bounds_check.h ++5924328cbe568095ea3a537e420060e8dfcbc02b753ffa2ceffe81f2c2b3d78c vllm/third_party/fmha_sm100/csrc/include/allocator.h ++2c2f6d2a96710c7418797643ab108c62cce944273b63420447030aacd0f1f4b7 vllm/third_party/fmha_sm100/csrc/include/cutlass_utils.cuh ++83126766f5b3ef24586eca68bcd3ebe45b1f8b1c5952fb55940421618175f271 vllm/third_party/fmha_sm100/csrc/include/exception.h ++707fe88459a3714a4ea0ed8f37d275533c23d903d6d042fa5b8c44ea878fe61c vllm/third_party/fmha_sm100/csrc/include/fmha.hpp ++f6b57982d00359fc5bc847c73929e84df99c529c6610be664e6bbc5f9599d029 vllm/third_party/fmha_sm100/csrc/include/fmha_common.hpp ++c99ec665411e0853fd4c7797eb3d3299788898703febac63edda57161d21eda4 vllm/third_party/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh ++001378d8d663b88219e2ab0b911cfd8ab467c38ed0a6b8de01e2c575a8993878 vllm/third_party/fmha_sm100/csrc/include/fmha_fusion.hpp ++bcfc6809b04ed7e9fafc01182af562611d801e43070092691a3ceeb40094e762 vllm/third_party/fmha_sm100/csrc/include/fmha_options.hpp ++3d07689226c4363880a14f4d694b60ca90e679c017f087bb34fbb3550d21a1e4 vllm/third_party/fmha_sm100/csrc/include/fmha_tile_scheduler.hpp ++f72afc2a4340ce80aa13870a3648117e0999b78565cb5264809fb5570a9e63da vllm/third_party/fmha_sm100/csrc/include/gather_tensor.hpp ++197f8bd34e279d8da44ef6ec5f6f64be73541dcbed8c72cccd26354d4af30222 vllm/third_party/fmha_sm100/csrc/include/gpu_trace.h ++aecde488616cbf433347534ba3d3862d7576eb35d50e3640f425e31dc95ef7a0 vllm/third_party/fmha_sm100/csrc/include/mask.cuh ++5f1b780b8f18d0c7d56aa69e0f6b1a6f6ed412114334dcd3d711bf4afe70c4e3 vllm/third_party/fmha_sm100/csrc/include/plan.cuh ++a90bdb34e36ba546d8bf91d25514f203f7eba497b3f482476f9178afb6285bf4 vllm/third_party/fmha_sm100/csrc/include/pow_2.hpp ++b8bb02cf51086afe5b1c37f9fc83d601a4598d67aa7c1c554bcffe699f25683b vllm/third_party/fmha_sm100/csrc/include/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp ++12449d0dfba043ba99269aff4b13c8e277e1f00dd9e61d535e289c43a822b236 vllm/third_party/fmha_sm100/csrc/include/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp ++47f9a5e8a37f5a0dec85a4c133e0a93d3764e22c14fd7d105b23456c9ad89682 vllm/third_party/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp ++b48eb681e970121e966c88f897fe722f6efcd59f42d2b9c6d297ec7d1f5a91b8 vllm/third_party/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp ++2dbf2fdd3b5067ec4e4a5c4459a8f88a85ef7aaa41bb573622ddf5bd8daef162 vllm/third_party/fmha_sm100/csrc/include/sm100_fmha_reduction.hpp ++31ee62916871c887ae7439443e262a64933a7b93e9545272e71bf5bc436d5033 vllm/third_party/fmha_sm100/csrc/include/sparse_topk_select.cuh ++010132b30fbfe4a4693b0da103d6dbc30e43d45521bafaa155b58e5f3e2e1ffd vllm/third_party/fmha_sm100/csrc/include/utils.cuh ++dc5ee4a41db56908846eae8ca96b76a37dedff38c400681e7faaf51dc78d83b6 vllm/third_party/fmha_sm100/csrc/include/vec_dtypes.cuh ++676f9fe2dbd6b46cc79d78a15c603c93b30141e3c4771b1afc8fdc631e61e251 vllm/third_party/fmha_sm100/csrc/sparse_topk_select.cu ++e77134f7fc39637bcdd490e6371ee140014dbadf9a670627b673c9737a7e9b86 vllm/third_party/fmha_sm100/csrc/tvm_ffi_utils.h ++1bc347f48cc306d14e899a190ea039558281d7fb6ef23e90acf9b1ad6c8b8457 vllm/third_party/fmha_sm100/cute/Makefile ++44bf50c9d2af96e8a288b44b5c4d08cc5d2c554c6fbc854c3ffe80c705f9295b vllm/third_party/fmha_sm100/cute/README.md ++5f7a95fcc574f659fc67ce0d542bbbab2f0a0d074a0514a58792e3299bc235b2 vllm/third_party/fmha_sm100/cute/example.py ++10df94ac974772414706232a56f5ff5d314554e25ba1d473c351ad08bc3a59b6 vllm/third_party/fmha_sm100/cute/fp4_indexer_interface.py ++31703d7cbd695071081b61dee168f7c6bf0fc62cbb819470d72a6647ac7e0f47 vllm/third_party/fmha_sm100/cute/interface.py ++b160d82f7811555051f8966d8e02750a985227c1d022208c784e606946f2eada vllm/third_party/fmha_sm100/cute/pytest.ini ++d6378f2db2678238bc00d7e70cae8f086f36f4031277e5ce32a615b89e695f26 vllm/third_party/fmha_sm100/cute/quantize.py ++336404acc7a51e1bb7a962ebcc57ece1fcbd688416d96e396c54da18b30e4a29 vllm/third_party/fmha_sm100/cute/requirements.txt ++bd9a73e4bb34dbb333c672e64b99beb7a4ce400691b300f1dd486bfa320650e7 vllm/third_party/fmha_sm100/cute/sparse_index_utils.py ++65d0a9ce76e5abd520192360234849a035e993e7af72f1a3bf7186b61c43fe73 vllm/third_party/fmha_sm100/cute/src/__init__.py ++65d0a9ce76e5abd520192360234849a035e993e7af72f1a3bf7186b61c43fe73 vllm/third_party/fmha_sm100/cute/src/common/__init__.py ++c9ad4e1c4e8bab1fe96fd5211fef81bbc6bdf07ba199d974f380ba72059d1f5b vllm/third_party/fmha_sm100/cute/src/common/aot_cache.py ++1a876f870c1a7fd8b20c0fc0efc56840c311467e994a7ab660fa2cafb2b64951 vllm/third_party/fmha_sm100/cute/src/common/barrier.py ++05825809e350f5c6158566657e3bf422035a36a967a03db59d7de0000e69441e vllm/third_party/fmha_sm100/cute/src/common/blackwell_helpers.py ++f26e6afb3b304987be276717e057a0887cbfe5f2a09b37d3f040164a9126fd4c vllm/third_party/fmha_sm100/cute/src/common/block_info.py ++6c4b726fc3bb67b8c828d8cde444a59e1e1656f75ff2faf965f43103abd2eb30 vllm/third_party/fmha_sm100/cute/src/common/copy_utils.py ++9ddf2f208fabe3d2a4d7ce636b7f9533a770265bcca82923301461d161076a4c vllm/third_party/fmha_sm100/cute/src/common/cute_dsl_utils.py ++9eab7a32dcc0b7bba9942e3e91cce005f8a7e201cdb3e412a2e7d1d53b8c67cf vllm/third_party/fmha_sm100/cute/src/common/fast_math.py ++2fbef9d57a60398a65c11e5789225dc15a85d007633f0887c5ee8e3f11e85cb0 vllm/third_party/fmha_sm100/cute/src/common/mask.py ++0b53ea05da7a08d3c0f7169d436c419dfe30bd0944f77492ffb094f8b6414240 vllm/third_party/fmha_sm100/cute/src/common/mma_sm100_desc.py ++e64b4988efa13fcd1f8d347bf7b0ac95419f9b63f2869716e96259acdc88e5b4 vllm/third_party/fmha_sm100/cute/src/common/named_barrier.py ++1ff64fa5bfeac939bf194fc5a0f82262c754377a759a534785464887a31328dd vllm/third_party/fmha_sm100/cute/src/common/pack_gqa.py ++8f4ffab2d4f5039b845402c45ff19a4216162227bee8ba46b1e016f1a4221db9 vllm/third_party/fmha_sm100/cute/src/common/paged_kv.py ++30815fa030c3f05b3ffd241247e27323fd0c275a8f1a6976f7b46d6c2daa3d15 vllm/third_party/fmha_sm100/cute/src/common/pipeline.py ++9e3fe993b720047216de6b456f7e376acbe15d170f3883896b5f15db4846de19 vllm/third_party/fmha_sm100/cute/src/common/seqlen_info.py ++d6756e56a74c638eaeae6ca840d34db2863f7d0836f84afbbce2b4a8015caaeb vllm/third_party/fmha_sm100/cute/src/common/softmax.py ++c776629ab254d69676155313f9c62f36aa6411abd1a3622f97350c6914368bb2 vllm/third_party/fmha_sm100/cute/src/common/tile_scheduler.py ++f4a3c2bdb120e7e962d0ec3314f6099933eac1ad4421b87c4d06d6f37f1cac24 vllm/third_party/fmha_sm100/cute/src/common/tma_utils.py ++474e68bdbafb2a61099fe128ac35262eba9cd61cc781d8997511b5d199907724 vllm/third_party/fmha_sm100/cute/src/common/utils.py ++250a50b4be7c7eca3c0765f0be7d1755ebd5a888e7935e705501345141862c2b vllm/third_party/fmha_sm100/cute/src/sm100/__init__.py ++9fcd178faa1332085a4e1afbf2f6efbcd19c8e3a7b3254f4a176a5d7adb8c966 vllm/third_party/fmha_sm100/cute/src/sm100/build_k2q_csr/__init__.py ++67c5cc46bd24f4fca6b70034e331e20d43de80f5d2f9d399955bdac74fc39e6a vllm/third_party/fmha_sm100/cute/src/sm100/build_k2q_csr/build_k2q_csr.cu ++427b6d5e57f75116c735e733e855e8f63261d3ec09e9e30cf92aa0727ed2c606 vllm/third_party/fmha_sm100/cute/src/sm100/decode_schedule.py ++28fa7bcd65d48bedd425c1554a3b72fc8d2b38c4100ae4913e8babec2eb2e93c vllm/third_party/fmha_sm100/cute/src/sm100/fp4_indexer.py ++034baadade27e58df89848316fd373c64f6c2ac62bd85678b05f91a048a53a6a vllm/third_party/fmha_sm100/cute/src/sm100/fwd/__init__.py ++69b615bcbeaacd1fd87446870c3dbd5e65300549590f5477701b1cd51dc65510 vllm/third_party/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py ++e9ae56a026259de916b73e1e5d25c4bfb5a03d4e324c522fa82f336a7da064a8 vllm/third_party/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py ++9d1c3eecbc822f512ea148ed400230018b94813b6864c200daf30f122a00d752 vllm/third_party/fmha_sm100/cute/src/sm100/fwd/combine.py ++5d025dc2f2d09b6f51c15a9566f0acb449d3c7e76152bc261fa45c5baef9a51f vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/__init__.py ++ab4e35f985460a6a90fc5493cb014894d52fd000b44a6109996bc8b8c8e66c21 vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py ++007d878dc6a57994cc41866a73cb94e5134045fbb793d724794329c1b1c7cb35 vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/build_decode_schedule/__init__.py ++d3b60e5ab6cc97887fb4714913ff6317e0030049b48e445b0b18a6fc0d54bd50 vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/build_decode_schedule/build_decode_schedule.cu ++6012bdd6bdf0f06484dd4f56df4d10c2114845b94ebd8c276a37f120f77b8513 vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/combine.py ++ea67a3147fd0eb3ea6a66857af72a5c18822ccb0a1e549398eb354df9492d147 vllm/third_party/fmha_sm100/cute/src/sm100/fwd_decode/tile_scheduler.py ++ad04914b94b5eeac425fa7ff8ddabde057a8795588f27727c4bbcebb4740cdc4 vllm/third_party/fmha_sm100/cute/src/sm100/prepare_k2q_csr.py ++55e49e022a9439b21be5b33f9975413332375a6ce37836c6ad13e58d204470ac vllm/third_party/fmha_sm100/cute/src/sm100/prepare_scheduler.py ++6b494f6b9ce9c8aa70a3bee9c18cb4c196aef9b8682893bc6077a83ae0114263 vllm/third_party/fmha_sm100/cute/test_fp4_indexer.py ++957f74fc52b4b585470a1c80736213b0df4c2b13df124c1122f1f5cd0305c0a4 vllm/third_party/fmha_sm100/cute/test_sparse_atten.py ++2b11f2d6dc6d00c3374b68b0cfbd69c6e3529611a6e56b86e38d025f42c69a0b vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/axpby.hpp ++28f97f81b281fea7ccde08a6e1f082f63384e741943b64d304df2e00bc2e5965 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/clear.hpp ++e03e105882db8a5df784bfb08a817b8d90789b297fd0ad6e113314486c34965a vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/cooperative_copy.hpp ++d0de07433f3be215a196acc0cfc969db16ca6b162193634ee0eb502378aee029 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/cooperative_gemm.hpp ++3b47a251bbfef136850e52000fe955d8341fe54b770e4814e32094ba72713417 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/copy.hpp ++e06904e6b029b9a96f4806abd9d7a2006052bff8249fe54d010dc3f682580bce vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/fill.hpp ++db90b96cc1e9b76f1dbbbda3f455367fe2e90cb60a95d91940dccffdb162680b vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/functional.hpp ++25272aefdc06ac7b90ac8b638b651730d9ffd0b9e89b6bc56440deef21e2e0ea vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/gemm.hpp ++4faa5471b2712b90de71825644025212246a5cdd499fae27f4da6c1ec71e38d7 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/prefer.hpp ++d708f5dad04c2da8138d2e88e120c75ccab34b3da7702821b9db29b7f977ed55 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/prefetch.hpp ++36d7105ddd3e1adf8fa2248b111efd93bda03855024c08bd3e1ec06869d97d28 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/tensor_algorithms.hpp ++9e5cf5f24893e95f8eb6f40029f663b230db415da5228fa8fa3190c8b061bfd9 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/tensor_reduce.hpp ++4e1e2fceff0dee3cb073aae9bb3b0073b357d30b2a5ec77f2b67fb65b86ce395 vllm/third_party/fmha_sm100/cutlass/include/cute/algorithm/tuple_algorithms.hpp ++f413d655b21b6b17edcbdeb0c9fc3fa0ec404c2f9750b65092955ecf78d907f6 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/cluster_sm100.hpp ++d70030bec8a3ebb36949408c3ee0951124aa1d1d0b610ac5682c1bc3bda2d5d6 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/cluster_sm90.hpp ++0a53de16c74f09b6e0f5f155bbe9fd324b4e216eb926d008041c91b0f59c9e10 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/config.hpp ++ec31401d27eb23ace72211c1345178934c6a242a36adc8e2f7196bef08920eea vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy.hpp ++43f813e1e4541bc682112e7791a9e38d44cabbbd90875858d249b9c79df4c50e vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm100.hpp ++e205145d40d887c1aea5682de7f10c325fc752dee3cf1ca10a4aea1a58076f24 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm100_tma.hpp ++523dc2ea1257bb0efff64f626c3aadcfc623e02cbdef249ab8ab57d6475796a5 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm50.hpp ++6a46942a2e85cbb2c542ccbcaffad76f85c4afde033437b290b9be32867b1155 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm75.hpp ++f98ebf662eb95419bf2dfcbcbe9b9e54ba0fd965b8d65f5a850f56c71b158b94 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm80.hpp ++7668c6f9304866f65655555c39f5e10fa6ffb02df964a8e7918ac877d41cc0db vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm90.hpp ++c45306bfa5b32d32d658d8f73b7dce1519a916b0253975965769c0d95c436474 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm90_desc.hpp ++d31756a7105dc875751fdef875be1b9f2bc008111d0c45734fed39465c12c636 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/copy_sm90_tma.hpp ++8f65943f681e16ed11df4bc7bcf6c97ebd808b6b2d950f6cc5ef29805a154029 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma.hpp ++68306540dfc897b8247e456a271479c66e15c74de9d204997d67daf48470f5a3 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm100.hpp ++762d488e441c4949fb2273d510ad629efaeaefbf0064025b4cb7a51cff27194a vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm100_desc.hpp ++7ba0cef9f7fe47148851c8c259df4943f6b1ffa1ba5ef268878596a375a54b1e vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm100_umma.hpp ++4de2816faaf3176326eb9fea837f5ea06dbb21d2b49dacd8f9a8f33d89c98133 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm120.hpp ++9926f677ef823178f653c9e6c6d810ccec16a6f34ab781901a57797c2388a724 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm120_sparse.hpp ++33e3511cf8884660145c946cb75c05763bfafdcb2932b4763ceedf9b59bbcffa vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm61.hpp ++c7ecaf1ab57557ed6c4baab62032b27d372b6f71138c5b2adf6da9f53402b302 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm70.hpp ++68f7a26da6c89c3b219afc5f65a2e44da83c8f7bff8f9491799c27257dcafdd5 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm75.hpp ++4d1315c12cccc6652a1dced6d6cb2d13a0390da5829978b5e01895571b677a60 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm80.hpp ++fd7e5afd5b45bfbb33558d6b2e3718b8e4ffe370228460e91bf4a7ca7167b2ee vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm89.hpp ++eede4c08af6c2d79a6a956072c91a902026dec8112fa8f43392fd11a2dab9de5 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90.hpp ++4ae842109e111094fe5fd5276e46bb99f5446f97e806b50ecd55ac47dc9fa658 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90_desc.hpp ++bbe378e69891013e93ae348850a80d39ea299249d0211758c5d849cd46b4c17a vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90_gmma.hpp ++af6c9ea73dce26b0e19e66d9d8596fea671f2a2b730d6b9a1579a513ccc647e1 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90_gmma_ext.hpp ++b933601f83fc8fa796f127f099fc0e23ff11051b6a2d8f9b2a5d8d87bd5329ed vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90_gmma_sparse.hpp ++d9208b9dfcfefe4a4bd26e343f5d6b164c4916b6673bf559767660764af951e0 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/mma_sm90_gmma_sparse_ext.hpp ++b566b7e35f5237c9caacf7e8d8a99231c31cadb2f167eff1a1ab6fdee72d31cf vllm/third_party/fmha_sm100/cutlass/include/cute/arch/simd_sm100.hpp ++2452d93263bfbdf2fe0523d57267fc2bb111148072b6752307894049737a52e4 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/tmem_allocator_sm100.hpp ++6c4f78507adbfc43094eb1770d60c96ce3e02cc53c5d1b49edb0a8444c6ad7f2 vllm/third_party/fmha_sm100/cutlass/include/cute/arch/util.hpp ++8a6f6bf2d92b64111202cdd36f5c60fa4a10954d6f3d469d7109c3b82caefb31 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_atom.hpp ++f1434137f94707650fba379eb53b8c307273917d33329e2688f556e5e5d162ca vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits.hpp ++cd8afe8362df160458da38c2f4075773feb465bc69039bfff4deb6fbbb7bfc24 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm100.hpp ++d0e41c9af3e02261bf4497bac6487394bd24465d6098480c564fc907f0cc70bc vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm100_im2col.hpp ++dc0febbca9f9c62be2e30ba1aacf978af6753e6ca7ec4f54842b29705e14dfa2 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm100_tma.hpp ++312058966e3779166e8eb164cd9b5845b3a716f91fd5590c347b04b8c6a103b7 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm50.hpp ++58fee5b1eaaf0d19e4aa7b44600e8749ac54d34ffb2cfd82a46d547744f06eed vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm75.hpp ++9d768405d7ddcb1612572cb7439fb522fb96f3121337bfc0bc188e4d5ebd1e78 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm80.hpp ++843805a1b64a5c9c2d21871cb96e5b49f2c2a202bea17c3dd18812264fb22c9b vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm90.hpp ++0e6d4b6ab2c8ee4d8a30ca6f3226489c26da2d606f192b5378306f39f2ccbfc6 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm90_im2col.hpp ++b8d4d0be64ccf683b57c5ce8ce39d40b7dd9ca50a34581bb6a8b35f34de6c848 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm90_tma.hpp ++f3c78286745e7e3151a5194852a9efd2742e4fa7cb0047b8b221451491b23474 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/copy_traits_sm90_tma_swizzle.hpp ++115be67d45771698fc1d7f833a1afa44897f45132537496620d5f2622d22fd7e vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_atom.hpp ++624700113bb75fbd887cfd9c6d69c06073e8f4dd6d265564a00accd24ffc35e3 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits.hpp ++06f6ac58f8e3672b9b41bcb47e2e47c78bb8ddb2a26c68d5fb584dfc64f12631 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm100.hpp ++084d939b07fd21651841edfeca3edfa679c2fd59279e769d65256d9a865c2bcb vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm120.hpp ++be6f0182446037c966eb98404580457b5434021e07e8044b4236b3034129e456 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm120_sparse.hpp ++9a9d01cbfea0a69525d073bc7b973740d5ba2e0929f76067d54858357d8ef057 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm61.hpp ++23c72c128f8d58cac840089bcf07b120a2f4f6cb1b1b95c3b055dae2c34df8eb vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm70.hpp ++6cc3179ea53ef83d2f12373c2c421fd94dad84afab477876f43f77d24c281b0f vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm75.hpp ++58b972a2ec119bb15db0605a2613cd39ca52e09f630095e35677f5a17770d084 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm80.hpp ++170bd190b81fda2cd6d3a055e84933e7718c36f958f2e7d25f278e0414f815bc vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm89.hpp ++10965f179974148f8febbe9cd771a5f5c742c4d2444d72cad22bc23b5693ec25 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm90.hpp ++956d254e25fbc5a0f400424bdfecb5a5857d258a0da3839bfa971f0fc56f19d3 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm90_gmma.hpp ++d0e1e9736ad8eadda9ef36268a9e92d1488d37579566c4936f3a9ccdffdb0e42 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm90_gmma_ext.hpp ++3e4fde935f2c1134dfb99de542358f0d5aa3b23f53473b97b8b102424fdb7205 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm90_gmma_sparse.hpp ++6b5b3fd9c3a115bba3af79fb4c3525b4ad9340dcd4e5211a6bb8daf551f1b986 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/mma_traits_sm90_gmma_sparse_ext.hpp ++e96b8b24c50ec888e805e95205f402a2216ec984a24bc2ddcb21bd8ed083fc02 vllm/third_party/fmha_sm100/cutlass/include/cute/atom/partitioner.hpp ++0abde1dda75c7f8e75f1661ba30815449590da23c7bf86f615d35de5ef86fc3d vllm/third_party/fmha_sm100/cutlass/include/cute/config.hpp ++7b7944102b31e03d868548d44c2555baa7b167a9a42200361269860af0f06260 vllm/third_party/fmha_sm100/cutlass/include/cute/container/alignment.hpp ++9bc447b9ce4ffdb8b9d2a5e37ed0816aa24ea33bc7937fa4c036cebabee1571c vllm/third_party/fmha_sm100/cutlass/include/cute/container/array.hpp ++def0627cc4c07ea19b41f2582c30832ba44daa1cab68ba26418c52c176a0f765 vllm/third_party/fmha_sm100/cutlass/include/cute/container/array_aligned.hpp ++0e28b13e7ddde2fc02ce8009817ef43dfe2d1882f2998d73589c1a1a1b2d621d vllm/third_party/fmha_sm100/cutlass/include/cute/container/array_subbyte.hpp ++44c497437463c5b0e44289b4c0c7d9c69043c3095ab93132ce855ada7158ff8a vllm/third_party/fmha_sm100/cutlass/include/cute/container/bit_field.hpp ++2726a840e42e21c4840127fcb9c492ba99b5de2ee368034e837e0d4da2088443 vllm/third_party/fmha_sm100/cutlass/include/cute/container/cuda_types.hpp ++2b2f700d7eb56c58967aa7d7c279848673b8fc062b59a566f94c4419dfbbd439 vllm/third_party/fmha_sm100/cutlass/include/cute/container/tuple.hpp ++7585622a41a2d0dc5beabeadb8ff9a5b483ba0488ad8f8fab60b6b022090c947 vllm/third_party/fmha_sm100/cutlass/include/cute/container/type_list.hpp ++c24e9f65815221a0c81ddd59dece334ad5169ebd3e82527181cc6859931df9b3 vllm/third_party/fmha_sm100/cutlass/include/cute/int_tuple.hpp ++b2e9a96efed8cf78193435e1b72697ab8e2f908b396f5c2327a65fdd82bf2b98 vllm/third_party/fmha_sm100/cutlass/include/cute/layout.hpp ++54f1fbd500c43c68755c018f8eb1706b43e463e5a35527714e7e9e7eef2d34f1 vllm/third_party/fmha_sm100/cutlass/include/cute/layout_composed.hpp ++fdf85737c21f5ecc167b66a9ce513023102cc62edfbf2d7cf798bc7d83eadc1e vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/arithmetic_tuple.hpp ++2852d6ca5131305f3e6c9c40082a616afd1085b8c1420d31a7a22c5d91803223 vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/complex.hpp ++26853bc58995932ef47745c58bc8500101a16a1e49140a56b9cff99573e8566f vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/int.hpp ++34dcf102cedc0bf4e30891fec07cffee5d2e7ad7f0998ca64dd981cc45fe69d3 vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/integer_sequence.hpp ++4d0f0bb32ed2cbe70a8654f5287a17d9a96deece8e4d64883e28dcba87f8c097 vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/integral_constant.hpp ++0181b139d610bbb5503d6c7a21b6b55392c3912505cca39c9e86ef2422458581 vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/integral_ratio.hpp ++bed15cfe8363236e8c7ac06cf94efeb086ed30309c61fa30716afc770b8c10ff vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/math.hpp ++6e69e574f97e0bd3480cc6547dfc642a5e7c5f5dabe6ae7f7a1e06c30b3161ff vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/numeric_types.hpp ++04d22a2a9a9f7d7a02b744de255a55df30ba12b2105ce0585b2f7f8a10724fa1 vllm/third_party/fmha_sm100/cutlass/include/cute/numeric/real.hpp ++5200e0be69bd7dd2fb4a214bd54100e66ae986aacaf69acb9f22ff788e1bc6d2 vllm/third_party/fmha_sm100/cutlass/include/cute/pointer.hpp ++b5abdc6ed02e90ea2098a7094b13ce06f98abf562b53b5f1acc9992131ff651f vllm/third_party/fmha_sm100/cutlass/include/cute/pointer_base.hpp ++2320c0cecbbae3381e59c99cbb9f926c59bf7bbc9289dd725c9b88feeeb22efa vllm/third_party/fmha_sm100/cutlass/include/cute/pointer_flagged.hpp ++21a60a84ded6bcb2243292ca5ec7a5e495928921071848935e99ebef8659cbca vllm/third_party/fmha_sm100/cutlass/include/cute/pointer_sparse.hpp ++5452d2cd15476b2d0e62fcb739b918f5b2dc1cc7fff1f9c3aaee068f30d51fa2 vllm/third_party/fmha_sm100/cutlass/include/cute/pointer_swizzle.hpp ++e77d5febafa814d8bc6150261e099f3b54dcbb74d653415fd6f4c8c4c68fbeb8 vllm/third_party/fmha_sm100/cutlass/include/cute/stride.hpp ++308f641a775688413cc2c0e2ba5a5b52670589b731e468acb8929dd6ff6c0a34 vllm/third_party/fmha_sm100/cutlass/include/cute/swizzle.hpp ++ca1b8460d0f57d451459e9ae827aa24090b0b7bab0399c018543b0e19a692c1c vllm/third_party/fmha_sm100/cutlass/include/cute/swizzle_layout.hpp ++165ce9484ddf256881662136a0d25004299d87e36986c1de33b27030a0955b5c vllm/third_party/fmha_sm100/cutlass/include/cute/tensor.hpp ++1a364eb5aa03b3a46d6cfcb0932414f7ed53d459cce4ffc6f4a0770342208489 vllm/third_party/fmha_sm100/cutlass/include/cute/tensor_impl.hpp ++aec4c5c59b3c603af1d4b91f2251f4ea2b253d9c4894a2112ed5fdd47c31d41f vllm/third_party/fmha_sm100/cutlass/include/cute/tensor_zip.hpp ++65edce496beb4b954cc64602ccb9b81a11635187b0626c9a95e327a68604f5f5 vllm/third_party/fmha_sm100/cutlass/include/cute/underscore.hpp ++11960fcac4506a9ef06f25cc0cddc3a01c5a9f3a09b06e86fe837d6195093043 vllm/third_party/fmha_sm100/cutlass/include/cute/util/debug.hpp ++209d5a8480c67f82973e27d99b3bb4e0564088892cddfacdbf643e233065ac66 vllm/third_party/fmha_sm100/cutlass/include/cute/util/print.hpp ++7e5ed4dc823eae7df6d9b5009dadf550696491fca4580e6570d1fce8a6c1dbd9 vllm/third_party/fmha_sm100/cutlass/include/cute/util/print_latex.hpp ++ee5066dca499a5fbd933dbd13becb16fb824eb6307e047c516b4fd2e93f85d11 vllm/third_party/fmha_sm100/cutlass/include/cute/util/print_svg.hpp ++8333631555137f6e0661f2c10d34bc13095dba6773fb26a4753c0e05ca20c800 vllm/third_party/fmha_sm100/cutlass/include/cute/util/print_tensor.hpp ++f45583785a90f150c0cf753f7f44c746e52124baaa2746774575a83ab8135212 vllm/third_party/fmha_sm100/cutlass/include/cute/util/type_traits.hpp ++d9f024bd0e12e3c19a5c8619809b84594f9e5c018f5465ae875283b4251f0b2c vllm/third_party/fmha_sm100/cutlass/include/cutlass/aligned_buffer.h ++1d972720299e91a7f4f0b840fb89fd06b474cb10a8eae8df1e2b5b7aff16530b vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/arch.h ++a2f03ef654f1711c9072cf3f219d512804d9ed88616bb998e2e8673599a83757 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/barrier.h ++66b10bb160c54e5ccb4ef8b6c16b9558a5d1501e47a23a7485be6753d08cc8be vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/cache_operation.h ++fa43ab70628f699518e446dd528a0b4a4606432def639342b744ab22eafe07ab vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/config.h ++1c8ec8ec02ea42737ee8aa89abce83bae05b65bd1dae5ad4d2c5ed066b6d0d51 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/grid_dependency_control.h ++e122de5fe050aef21e3004e5fe853da7fc541d5ee52b48136ec34a1772f7ae75 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/memory.h ++205be5c3e8b2f9599eb9104ca7904d309d7c57c97abc60fffa8236ecbc53a647 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/memory_sm75.h ++ac539b45732761f7d41edb7e6b7e35ccf0d552a323056c244f7d3ad845d41411 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/memory_sm80.h ++995520aa528c8c3e8b3f2aceb91fd0f536255dd900eca24184bda0f9af020429 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma.h ++d88fd3db3498ac3237ef14129319f825ef36b6b13ad813faf46f3e5a08be9fe9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm100.h ++3cadc148bca05b97ba8e5a35307b24ceaea1c32ecb580e05076fec8d7e0f8c74 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm50.h ++421c584d982d200c260093a0e34c5dd5bac59d0905f590e1028328bc1511b9f4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm60.h ++60ffa31a2e9dcb3ff97aaf3776c0aacf5714667d68f478842b132fba8261b8a4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm61.h ++bb51d4bc0f1e869ec17299712d36b693955d7ddfe0a1c59d65cd0e0cabc7269c vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm70.h ++1b3b3cf8f6f5cec10e0b1a2025469212e6efd232e5ee6e8a645a88e5ef8142f9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm75.h ++eaf7789b82bb0dacf9b63ad03615479be04eb3e8516f5739aa11812a3d2e00d0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm80.h ++7937995bce66992cf10ab1733eee213ef3ee60f1924d326ac09c6a43cc582e51 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm89.h ++ecc03c51de3c03a5880ba664c4afe08da409f53f2b8d81e2aeb3ab11f9452e48 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sm90.h ++fbf5a7de37560263a03eaa420a4a3d9e46ab472db395147784fe5162de41b536 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sparse_sm80.h ++69928d200ea330b54fdcde6be064da07faac081542fa74c26924df95c1d8d621 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/mma_sparse_sm89.h ++02d929c9a49456c2648c98a22b01e81f01e39e4ff436c4f55eb53f94f67a97f9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/reg_reconfig.h ++0d15befef4a420a6e91219319250c7dcc4be2d7a008021fb0d930e67c7355cea vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/simd.h ++bc773ad062258b29fb7ec11b6b63c49fd8bbd765d7e6643dc7fdb12e1b7605b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/simd_sm60.h ++0dd49c918c06bd6d71dbe558db798317f2b23386e154fd10f4b9e4ef0d665f21 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/simd_sm61.h ++1232c652366854f2bd7f8d8835ac5254ddb7a48680ee13855de9f1007c459a61 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/synclog.hpp ++aed335890d9fe5305d16f0c48bb5441cfadafdf1baeeedcc2e6c6547bddb853c vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/wmma.h ++7d1fe05a79c2429eaf523690f63b45c4e682f1658e99462097b13ff60e85f6b3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/wmma_sm70.h ++f56b628c85b9567b66ad43e2e317aef04157d267b5a876787499278947b180cb vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/wmma_sm72.h ++1e2eba2839d9b4ce6ffbfd9215b926696c0ec2b315d45bcf8e39aedf39a9f3db vllm/third_party/fmha_sm100/cutlass/include/cutlass/arch/wmma_sm75.h ++038bdd0e14c8aca1ef163957175a02f7cdca5db37b330f9f494001d19b52527d vllm/third_party/fmha_sm100/cutlass/include/cutlass/array.h ++d8f77a667e4849f323e3b933c80b3cfe29fc123c8175d699c21cf8d422b699a7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/array_planar_complex.h ++19496f601472bde0d7fb9f4a36b8c49ac357781be2b9180630ab03fb697c6287 vllm/third_party/fmha_sm100/cutlass/include/cutlass/array_subbyte.h ++27971bdaa2e217c0d535481c951e2b37e19456f83224a2e91a030d7b1685f468 vllm/third_party/fmha_sm100/cutlass/include/cutlass/barrier.h ++ef3f7262c17ebbbaeb32ecb267c4529279b82aba085b2c13d844b8d12fb53236 vllm/third_party/fmha_sm100/cutlass/include/cutlass/bfloat16.h ++2e72bb87dd453d49228a64298c25fcf4fac66be1393dc6624b3241b0e6b4267c vllm/third_party/fmha_sm100/cutlass/include/cutlass/blas3.h ++850981f39d7d88e0959fdb8af5669b6764daf1aa357f291632ca6aa5b56008af vllm/third_party/fmha_sm100/cutlass/include/cutlass/blas3_types.h ++468f847a2d356bc5d477b73c0d75a14cd4abaf93a375fb1a6c79e4b4049b6364 vllm/third_party/fmha_sm100/cutlass/include/cutlass/block_striped.h ++5eebc67a5ca150e9fe158b9ef4f9229fe6c666be2527758764b52e2efd9693ce vllm/third_party/fmha_sm100/cutlass/include/cutlass/cluster_launch.hpp ++ee6b42d0da5e22ef5c3d4ad8cf396943608da0c9d935936246c181360457135f vllm/third_party/fmha_sm100/cutlass/include/cutlass/complex.h ++d676e2aff49f1561e3b92aadc0ec720ea38a64354e694e6a2664b1c5232c2d82 vllm/third_party/fmha_sm100/cutlass/include/cutlass/constants.h ++87de30ddd872c4e952e55c450b5a6527f0cff6971eae59320506716129e290d3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/builders/sm100_common.inl ++9e855d7dba6f8928126d12110c4d84ca59a2bb087386bcd51039db51766c1745 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/builders/sm100_umma_builder.inl ++de011804064994a677a860bf0bfb7b26d41f5a8218d505409dae8c2a5d6cde6c vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/builders/sm90_common.inl ++4a244f9fb04fbdcc4dbf765f45181c7a5c94d7aad932b3f511aa65cda61febb3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/builders/sm90_gmma_builder.inl ++e6887291b73157f43e2885a766725c246159cc1d5242ed2c3a0c892b3385f661 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/collective_builder.hpp ++e65e4b78d53e5d42a59305b63977662e5e9250e29a208b25366bd0dc21cae6ab vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/collective_conv.hpp ++ec6226f31637bd70125d7e497e74dff0237a72196742a280131977215104a5c6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/detail.hpp ++7c9bec76e6845f88269f93616b1efe664da42ae28eb550fd476d254ec3784af2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/sm100_implicit_gemm_umma_warpspecialized.hpp ++dc991e58178c4d10c94ebfb0e6e9fffba4b642f025a8c08bb694b313824f8f4d vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp ++a341c409c9ff1eeeb82b70fd20c0b23f61390086a16081f71b435cfd00e2fa15 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/conv2d_problem_size.h ++17e17e6f572a8abe2da85c856a71a72191c6f52ed3c7a0eac902804b12af6738 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/conv3d_problem_size.h ++ea6ca40e3c51f53871b73d6bd8e93a12448a0bbf345c727740e9c5c9e799df92 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/convnd_problem_shape.hpp ++bcd79d6d36415c0082fb1d16329d05ee25b27cab0d30c0f33df8fdd035c9d247 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/convolution.h ++ec6121026dd8a4ea65c1627805f6721fa90d2a07ffd1be93c11b8e140be4cfef vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/detail.hpp ++4c31e7f68e8ad7f924431222f03b65d3371cb75205d5d274b05877dd656cda3a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/device/conv_universal_adapter.hpp ++91f2f85f9fe1f9739874200ee640ed566b7da5950b0c9e1aaffe18aa33c75c6a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/device/direct_convolution.h ++42b645d3cd363a81f874621206ae634cb95be2236da3c287c57e39269405bd92 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/device/implicit_gemm_convolution.h ++a4730aff6ff355d4fb52c3da9677fcd95831289311eebfe310862fd11a98741a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/device/implicit_gemm_convolution_fusion.h ++34dde1a996315ae2c9a8f8c6c62736b1528403746d5a87df7d3ec3bf0f431f25 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/dispatch_policy.hpp ++1a81c00db754e9c347bf46303ee40b2ee1a362b5e5dc756c66385731109d39d9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/conv_universal.hpp ++3dd0a4f2632c48b145b9cdca3db43eeebbeade6c17537cb4dda30197a6648d1a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d.h ++3265a486c5dd41ac753e6a96e70b2847677a90a4bc8a20ab0890d7df228f9bee vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_dgrad.h ++a9e853438d035d6b6347d3bc12fe1156597e9e8c2c36b9c939b7fd01d5848101 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_fprop.h ++481682866356638f2210b8a73e9d0d62f394a2ec1b0e0cae3bfa56855b005cf2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_fprop_fusion.h ++2d273cd7451173873b014a58631aa48aefe4eb7a208fe562bf5157ad615da30b vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_fprop_with_absmax.h ++d6f10c5dc643cff54141fcec7538e5c8344560fbd2bd4cd75f216684e4a0b598 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h ++7759f447022996e387e1482b0ca4d2d09ebfdbbb778534d295bc2e2c5903f0df vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_fprop_with_reduction.h ++4f976c0cdb0397bf932dd87361040a6cd5be9516850f772ed5e3336e12da312b vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_group_fprop.h ++e1aabb2acebf428b734f0f4926950e6a4dca68deb77417b2a62e53164195035b vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_wgrad.h ++7b61e5d8f1032da96a1c00abae2bee16ef40960f45a1a509dbc26c96e0a6eefa vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv2d_wgrad_fusion.h ++9aa4fe4d29217ac0fb114ed8bdcbf66a6fbe40ce1c7690d87ff793036977208c vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv3d_dgrad.h ++ca6550b7bb515e79affc8b34254a80058b3408800dafa5541ca002ee87739038 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv3d_fprop.h ++88770277fccbcc8e627c676506beeb52b2c93b3edef9bdad9eff86ca9dc658a2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv3d_fprop_fusion.h ++4a0c0ca125b80aaf6131760f3378f5e5de1f94b2ddd19b1a657b72cb40e9701e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv3d_fprop_with_broadcast.h ++2dd2dc5303e8e4396634e74c5468085c107d4e65b74e05f7117f00f87691696e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_conv3d_wgrad.h ++aaa163b3df9894a596895e2eb00ae6579460f18e5d71eb8bcacc03828ea4e6db vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_deconv2d.h ++6fbb6b766eeb8ff8a1d4831f9afd7ccbe1ea1e6e9fbe504287374f68bcf31fab vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_deconv2d_with_broadcast.h ++a483ba5c6130abcfff03192a2a325b2f8d63c1fb6a81d532aa2620a316dd3adc vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_deconv3d.h ++c628fbc73f72ad9f13b7ee6c51a35f4fd552992edbbba52c923a17c551176bbd vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_deconv3d_with_broadcast.h ++668b7e8ec99221fe70a629a696458a3308246bd781463c6203238d77d509d88a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/default_depthwise_fprop.h ++fd0f6b816ea6eb7734952823e6e5e7409451eafcac477d2f1ee8d7aac68c42d0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/direct_convolution.h ++156ecd95f818f830463469b692c9d1c5977823fd1feb0e095d836e8c60f4f9fc vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution.h ++e42af3a2b026d2c89f2564ca097ca99a917d28dae95b03ac33e92541e27465e8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h ++14bfd33a7a93bd78abc122c3a207af8d5b6bd7747bbba411d42e1b2f4bf257e7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h ++e1dbe1be3c9c693453e19547a01f1e2896b481023b25b8e323afebfb44267085 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution_with_absmax.h ++75ed67af22aee6acf61836884e3abbcdd2c6c10ed6cf3f755a220bcd93ac0e04 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h ++dabd2a83a2872551dc910425c97096657efc92c4f0790cea40cf5585ea9cff99 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp ++acbefc9fcfb60129334a9299f0e044db0ea49f6c57a4c14acd78507394fc371e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp ++e59d9c8ed47ca2858ea760438afb4d9351b5f7998e7500ec092cdbfe59b2ceaa vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/thread/depthwise_mma.h ++5c625124e339b76e5a99eed6750f2cb74131e6f86a0cdffa036b21fd20e710e9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h ++c958acad86642c74c195e331fcbae662db37f6a8e0065a9655c85b6e84d2270e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h ++943923924d7495a10c89fdc16ffdda79875ed8056766e095e311801a9290405e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h ++fababf526f027545bf572256983f3180e4e39b25edd1f5cd4ccfaa579d4fcb84 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h ++d6931bf1ef4ebccdc20b8836675191a79819e98648a40204afa5e0c40b1c31f9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h ++06582273387c884318b828c3fbb98e2c40a97c94da0bb033dc7a7f46086ec071 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h ++4bba2a024aae7c1dc57769b0227670caeb6b1e0d9049cf66db71c8b03244fa31 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h ++444606e1a9dc0eccf479debdce05610fb36d872a926038c589260e1935fef6ea vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h ++6a9f81a17babf49a2bb036d64caaf4ad463814838cb016832f631f5443abb5d1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h ++8659b4556cba40d2a3b57edc8606ea73b165609696f5313aa1c953d5f0f35845 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h ++7cbfaf7e860d7ab43204ac12a0309071e1c74d58b14428db673fe2fda08cd91c vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h ++b45f7d285954a1d5c60807ec35f6691c43a2f8dbcf6570f41f16951a137c6d5c vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h ++474e924fdd23fa66264b37b74a7d5053eea1eb4c71e5f8f30f47467b686d2859 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_params.h ++2b76cff6a4e799e3335ae5ab09dca08ddd5314a0f71c0fec5cddf58968f5c805 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_tile_iterator.h ++d94bf42017e6660d9bd811ec6698951b0ae724dcc305c9987a1b8ce61c850f55 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h ++8143ec7805cc5e5a749eb6161536e610c4aac0194a751d9b588259a800d2e985 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h ++d0ea11386b667f55a94261d8614d66685e9411400414ff382eb230bc412afbc2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h ++dd8ee03f6c62a5785688550326712e7011c3cc30af9691bff897cc2f102f2291 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h ++d16740e33e4981a02401d40e35675320c99016dc08daff7784b3075116ad7da7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h ++17dbc795b9722b30eb124a5dc3fd36f4d91c644f54227bdf71f29905f6825cac vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h ++7812ca22b0fa81758ab40381cee866565adfa9c05b16793aecb93ece8679069f vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h ++7bc802602d5b6dcaa0c1ef57b5a6fdad1006cb6e9052b8b0097b21feebe903e6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h ++54e1d75ec6570b966bb4d814325e2aa5067f108c5fcf2c302ce7dfe46c54411d vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h ++06448bb7a6a6b83083861d5315177a46a4cf64dc6197a019c83a2e8926236244 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h ++b38a416f9bf9a6c6d48d149dd909033bbe9c47c052077b22709845457ce68fde vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h ++0af9d7f157f613a4cfd6496187277f22b5bde9f48b6c345fd3dff418d51e6942 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h ++ab7dfdab52c6a9ff22e4525028e22213f0dae0962c994a2db88f425c9d226900 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_params.h ++01eb0081ede4d8ade0454105d5babb99a928184589ba4d1f95c25bd2706cca9a vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h ++946ec3af0da2696726f6cdacc0519ad27b42011d12646c1af0a6ee46723539b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h ++0cb3b899b8871aaeb92b226fd7bfa1a1777ef31ff2fba4c3cfb6b46fa31deeb8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h ++8f8e9171bcf34cd57e1ebe1a3b166b8103c09933d8ffd45ae2b6d87d968a16cd vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h ++e96623b61895aed7a5e1b9b99576ccf385f5fb49e0cf2d85c44fb8b25cc6a947 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_direct_conv_params.h ++c4fb254d5dcba5e74f49177b2b1568fcc7e58e30f98a3ef4aced555d7b52c19e vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h ++b5d0dc439b523c058bc2df695e6a22e9399f34aa6dd3ac36e8bbb19f374f3fa8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h ++a5d86324724e09439263e2f23bf3f71cca2496c81869826bc47c57e0921d02ee vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h ++c5da1cddab89349f302dffe2b6012143e4fd4c1b1a07b5fbd90a257239a60898 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h ++001a4e1fac05ee3f8e0aa0474738f290972568fdcef9a9d38a8ecfb5c95985cc vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_fprop_pipelined.h ++17a3b22ae74f7f30766e61aeb410ab2468b05c7a920f0f8f2fe3610dbc523fd5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_mma_base.h ++d9739fd996e8e7d72b05ff5336c58357dc21b527a3e8e0faada18a3c33e8cd89 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h ++3000038ca9feefb5890e0f26320c9a3b8763a096b37ef9593dff784447ea67eb vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h ++b628331404305ff8d486d7de635cf71c4ec11508d4a63a34350509347c76cc05 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/implicit_gemm_multistage.h ++4819466c7e4d9ea46c43bad3ad0f9cabaae03afef35488ecfb379e06635f1118 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/implicit_gemm_pipelined.h ++152f6dd209df00e922e08eec82c5425799fde6305142158780be61fd9bc65803 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h ++7457fbb4558e0d952029cc9a0cad7c42b13ebc5c9c6ac39e9be60d0c8b328108 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h ++e972d619710bd9bd864f45f22d6016aa9a49bdf148a9a6734b5dbdab9b8439ac vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h ++9656c5f1842b9fa01b0737a098d206b02121bbf7b016398da3fb173c3fa99048 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/threadblock/threadblock_swizzle.h ++cd03da0160ab79331cc4aab97c5222668c50cfa285801719d17485fec90c37e7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/warp/mma_depthwise_simt.h ++350c0967cc5cb7db4eabc868af61c5a0246788f1585344c9e74b1974a995bd35 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/warp/mma_depthwise_simt_tile_iterator.h ++69c0a237781499628a3c6022ef4e0c79876ed79a7b836afb09f78bf918c35756 vllm/third_party/fmha_sm100/cutlass/include/cutlass/conv/warp/scale_bias_relu_transform.h ++d215d166dd6b032a3220da5202068e3941606cad87398543b9822918b020372b vllm/third_party/fmha_sm100/cutlass/include/cutlass/coord.h ++47649140f1df9d488a306dc2fed79aaae9e1f7982db9cc7bab25b5f0da81484b vllm/third_party/fmha_sm100/cutlass/include/cutlass/core_io.h ++3dd7abd0bd4da92fc435ef6b0a8c04520aef903c1538da0a86b1bf56a138a685 vllm/third_party/fmha_sm100/cutlass/include/cutlass/cuda_host_adapter.hpp ++b0983633d014d2d7955d6e4963fe3338e8f444ad3ae4ff7096041f4d87fcf82e vllm/third_party/fmha_sm100/cutlass/include/cutlass/cutlass.h ++0e7639f652cbee43a108f42c73161ff5427094f5193b487bd4d824c0b1aa4822 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/blockwise_scale_layout.hpp ++f8eb86224bfd04a9dca3e9d1da7c6078732287e449233e9d3dbb23105f7bde56 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/cluster.hpp ++30a97c79bfa752036478816c031ff2c42bebfd7ef19877beb7ac72eeaaae5cbd vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/collective.hpp ++aac7cf52dc7e8f779d1a676ed016cc87b6fd0a8c19946d888d186fff75f2737e vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/collective/mixed_input_utils.hpp ++721826f1f33161cefbd2e47524f8b194844ea04405b50dc832bc2dd486ed6651 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/collective/moe_stride_utils.hpp ++860600584680f40d19aa24b8ed19342f9abcf452d54920eebdb2696c01346f3d vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/collective/sm103_kernel_type.hpp ++4bf38c035f080d1a0cd49722cd81990f0ae34368ada7efa304f76f7367231513 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/dependent_false.hpp ++f59a5e9fc0f822ce6411d1f6d0979ded68efdb1bf6755d322b93163f31e3b01a vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/helper_macros.hpp ++1e9609d71f2f75014f9b3c625a4b967a5b14ef7882cd969753bd04a780bee23a vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/layout.hpp ++387c801091d8ed49d84c3c306bd4f67b8b41f80a9a4aa6e33e75c1dbf7131d14 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/mainloop_fusion_helper_scale_factor.hpp ++fc52fc1fa6772593e3bc85f70a4cf607e51344ce751a0522518110c92259fb67 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/mma.hpp ++c7afc5429445cffa7457b8d6cb0ef5b696635265340da2500aa926073ed77fc0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/sm100_blockscaled_layout.hpp ++6b36b47ee1f2030d0251c665fe3a53d739684e01470e9743c077e74520dd26b3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/sm100_mixed_dtype_blockwise_layout.hpp ++ce320a1161768648828545ea7e9eff412374514ebb9b59a232d0fa441dc1d1ce vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/sm100_tmem_helper.hpp ++99ac5fa927fcad93e0835e78bbae74448b10928e62d1f609ad0e04e68a7164cd vllm/third_party/fmha_sm100/cutlass/include/cutlass/detail/sm103_blockscaled_layout.hpp ++c5b234f7812d4bee7b2b781851cdc36cafa0fde8799e6b9da0f5ee5630e80363 vllm/third_party/fmha_sm100/cutlass/include/cutlass/device_kernel.h ++236b85ef1962c0132919be87c771a568eb96664514a65ad59609bcb40f488ae3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm100_builder.inl ++c8b4ebae999d9616fd6ba23662550fba75f0aafd1b96fe5a761f222941c86f2b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm103_builder.inl ++1154146c1f2a0f40054d9d9aa24f05e7b7c4a0b0d1987180b768193b0a02cf1b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm120_builder.inl ++d56fee66d1319c7d19aa21fcc67d420b192e1d41155ecd468df0c089634e0c7a vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm120_common.inl ++bf695a19582a2d40f9ebeac464e800134244b362d2b37442c8957207ac01497d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm90_builder.inl ++a3956a3ea24e69dbcd34d7461b10ab9a08548f106025de07130ad24b490a907c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/builders/sm90_common.inl ++abac4795e7c13b21a2f01ad54fa45eb2c1390658d9fe2a29acd5559edcdd3ed2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/collective_builder.hpp ++36e795254040db105259ec4a2832b9ea74a16145d06d6141dd4d4cef849d4034 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/collective_epilogue.hpp ++47554250e74624d6702a8d2c4d9cc70dccbf030d890ae24f81f91e6a2d002c1b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/default_epilogue.hpp ++ae22ed2b9baa82c727a1c6c98478f68d5863f552bbf5a1ab4773be4f04ac9ea9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/default_epilogue_array.hpp ++59611eb5a50334d4a06890cc91d3de5ddb1c0a2c2d387366bb694446e97b9c80 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/detail.hpp ++f9675cb1acc5a089e49ce0c8e7f48d4344c54b12330231c1dc7a2699d6ad12b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/epilogue_tensor_broadcast.hpp ++2b1dfea335d11d904445916d96b451613f13fbf90a2dd15845370d42551c80c1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp ++6ca409145b484cc948d2b516ab56f6c827bf867b33611336fad3f49e78062870 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp ++3263920f0128c941074ab631793dbf10ee9378395b7b7ca2f66195ff70b488d1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp ++ed1083da635b3a89fdd5c0eb898922c0e5421bc455d6f2fc4b027d95642add59 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm100_epilogue_tma_warpspecialized.hpp ++d39524f09c1322ce30889d671a3abd313774eba3f5591c0365f1e34aba976f50 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp ++01d0f976d599dbd949f1b201c6c74a25d9bc6bdbd72d940fbf8714b8608139b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm70_epilogue_vectorized_array.hpp ++3b03d642eae5e9d6d1a950ded36ebf769c5eeb0e09e6a808633e5ca30e8763f4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp ++b852c59063eec6b3383398c7018e0fd94e1a4ffb7a6fb8bca1f6b46a1160beb3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp ++58289eafcb0128e4f61a1c96c6dea35542252ac7ceaeeed0ca0d408cf4433305 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp ++49a1500dec5dc534d9e83734a6d88e32bd5f5b075aa89ff656f7613e45cbe382 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/dispatch_policy.hpp ++468358680c6aa975b397ea7fbe9d1956780297749668e362391b50101eac8b90 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/callbacks.hpp ++e49c43cb9f67ef6aa6845fc35751da61a2d36fb16eead9737af1ac563a339599 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/operations.hpp ++5cf6d006bca63ea20911c1002fd3eca2b21a3768f09a0916f8994eac00a2a439 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp ++1dd70bd030d5098f260d6bbf84982509b87ab392cb887a7bcb9ff65b4ea94508 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm100_visitor_compute_tma_warpspecialized.hpp ++b6ce1d179e0bc852c9677c143876823ed86437a181027c503ede8b0370d7b2bb vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm100_visitor_store_tma_warpspecialized.hpp ++87cfb252459b4ca6085917fbe492cfc762b532744064912b2b6b4c75fc313d1a vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm120_callbacks_tma_warpspecialized.hpp ++5ef61d98f8491770a9da1e6856d5d1a9a892f83f3bee1d0aea175a142ef5d0f5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm120_visitor_store_tma_warpspecialized.hpp ++15938fc8f3d89ed7aa3c2189c57e604a3ce88fe0ac0c449c8aed626799dfa388 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp ++4deecdd5e1f9efe7ee8ed25bf470f448b6e2b620fe58c58336c5ff3f0f252b0f vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp ++6904d958e757dc685a131c7a3fac114419005b1618de834b5067be195e410954 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp ++67e3412b9f838eada32be6904d2462675fc3b383c406543deb1358351c775fc1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp ++f94e2d5853901df1d9fc9a8e43cc71fa7b86f819c2e1d8ed46a71196d67d77ac vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp ++9346d0ad9097db9c61ae1570644541a480aea392211d4293949fed8cd1eebca7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp ++01abb59ffece48adb266eaf1a5b58d422631eb955a59231e2db9fd56f6d8091d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/activation.h ++5ba6550b7c74c546a8b506cc7d742670ed0c47102bfab0a1badf785a2995176e vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/conversion_op.h ++4f4b13ace7dddd4aca67954d03ef00189b6b08726550edfcbce2f41af2013d7c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/detail.hpp ++e70c257c3722979a0ffa0dd73482ca8e7238c273ef3ca6c4269577eab7cfab08 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination.h ++d59049f0cd0cea7db908821aee7ec4c83ddea432df8f8666ff3ab850e03c77f0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h ++13970e2dc281b8de848dfeb499d8fd39802ec45f9ab82cd2b6c08e6be7fb2877 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_bias_relu.h ++648df71fa5cf31297a9cc5ff66784528e353028224b4edf2d2ea58d2f6906446 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_clamp.h ++dfe5bd3afc0a8600a97dc648bcd1fdf15ce8db1c3324e74ef13e7f8baa8360f9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_dgelu.h ++ef85350eca86eec1fe91359fe9b8405d449ef490153f8fd943738aa2cb4aab20 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_drelu.h ++d164ab6bb0623a3a3aef7d3987a789d39f372d136c46855a31ec28baa0a51c38 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h ++08947670037c27d02fe5ace0aac7ec3fa9c256b0391366ac4bdeec579728dbef vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_generic.h ++efac2b274664c00acd2d38f1de0521fc562428b3046d62ef8c24bfe9e293a08b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_generic_with_scaling.h ++82a5f1cbdbd52a19a883a04f1197da5d0ef12f38c6b2e6d0793c4312a3af8f34 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_hardswish.h ++05b91e823a5befb1e2d32267f320ff23aa70a96dc8d3a1aee35cd419a2c0bb32 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_leaky_relu.h ++b382568642170f2dce308085a7a389e7450e9563fa91d5003946c727a085be27 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_params.h ++c5a1c3dcc40e973f95e754c6686fda33e78a49823222857d7438db67969fd186 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_planar_complex.h ++900b0fb6ca3536aa3e9f090187cf67d26605f52ea6fd3177d941bcd67c320344 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_relu.h ++b57723042af39ec52345ef5b71c979b2e7005c6337fdf9581d0b6351195d2929 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_relu0.h ++85f9b683846b0228682afac86ef361c0cd2a4ec9308bb3e4acdc5749e7da439c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_residual_block.h ++5880c8fd533346abaa8f42697a5c91a883a470b3f0f400207592c06fdd053972 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_sigmoid.h ++2aedea01a3bd3ee15875cc3c6cc27c5a7558beb9463089dd8872ef77648830fa vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_silu.h ++f03217a93a96ef9b6bd861e8627fa1cc093dc9b73d6d1d4c424be37143fdbe1d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_tensor_broadcast.hpp ++cf90329b48d2ee0023d02bef9b9d0b6f69279a1ef7b4ed35f8b702942d484949 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/linear_combination_with_elementwise.h ++493532c6161be66adf1a57337472c41115807111e134ba5fb1d4151af53f12e0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/reduction_op.h ++f45767eb8a06a73d84667bafb31f6e9146e51ff4225ffe4ff3cbd8c5f721c74e vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/thread/scale_type.h ++9f9485300f037226a819abc83d6662158a107e2cf0d0704f2fc3c174baef579b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h ++0afebdd56a2ff2825224c000ebc878c4e12e17c9a0c18f9a4712f12adcd2adb0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h ++2913178fde0fdc91efedd9b28c8b4bcc21a3fc8395a91e6841681c8b78da9a7a vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_direct_store.h ++914b87d1a66e53c30b02233c2fa03ab5fc1d8f8ba1d22b3cc6014fde13270344 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_planar_complex.h ++65f45cc87757cf83576eb9657d789333d415145506f304e567f3e624fc019656 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_simt.h ++b6f1428e7cae08344631e52fe9b252888a0ac67b7345661db0ddf7b3d407e4ae vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h ++b2359cedfe08e699a660cf9d4bc3110273369c1eee5be8d51d1549536033f661 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h ++f31c9726b0763c3db6dc170c7222077fae69c50c5e525dd9de8c8b53198d9424 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h ++b84088a0c2578a326306b57e2f339aa7993f572fdc730d91586e7b60d0a9fcc0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_with_absmax.h ++0cad0fb1675e26f5824e9bf0cc91e87be11cf66dd7f3c82517de9ff5c8429a26 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h ++a12c0d0358f82f07570d65e6d82006e7599b73ac5c4fa38870e2a6985122cb63 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_with_reduction.h ++992c3f5b9c8a128ec9c06a7815c5b4b99b5bd698954c13b659b93085bec1571a vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h ++03e0a932d7dd97391cad1beef4fd073666dadabba10a5f1b57784e75d79f8c9a vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_thread_map_simt.h ++4055c599c13cde7d90aefa45f736b0414f4ee197844a1a4a2c4f7138818f4c0c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_thread_map_tensor_op.h ++2d9b37c69b4ffb9c1d452ba97fb5ce5ebde46b2f67c61f0a188f5a9b30495c33 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h ++a71e99c4b567d52ac727e74cca31317140a6ddd96df9237b53e6c7862e737124 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h ++ac009c785511a2bf0ae126063d86f5aa36e5f0088bedb1fef9ca63e6fbd80c85 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h ++68854dd467eded2c517aaa2533a3af2bc899a9172f1d91c44482b22a303a48e9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue.h ++4601f0bdf2c14758979b0ac890231744a6f5428bf0f7707c2e3a0f33c67a838c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_base.h ++7bf621b7cfbe98800b86c1a9d5aa9b6875386d95ff606ed9ceb084c3aca0972d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_base_streamk.h ++bef93f0415680366723c60fac0c68ceb9503e17adfa8c8731a0ad433e289c6af vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_depthwise.h ++6b71e3101384b38eb5651620de4fe84c08da10db25ed7e134c0b9760a1fe1084 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_direct_store.h ++f61443fd73481bf1612269854cb6175cfae644e167ff763acf0a46a8167e799c vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_gemm_k_reduction.h ++164bf5d65d08886b3f25682f2f425c78afc29912c4f761fefa08d1927ca0dfcc vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_planar_complex.h ++a7208d7ba3261b423a123c20a8e28b37b84f7f67fd085decd4ac5725776eff23 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h ++9c360b378fa7bd8c17b69cd0e1c1789221c12f8fce42f33049944c77ae342e7e vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h ++24f0f0e948e78655c0bba6c57c0afd265a826df86185675687a878cad10915d3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_visitor_with_softmax.h ++27a8369fb67dfcd25688c7a7200fb971a6f8f5be2b071558fa6fb129e3099d8b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_absmax.h ++23f5f35b0c73369c39520d794b073d8237d982ec7f75da0a64f22a90ea334460 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h ++4d09e4be3f26d2db5f15b843a4212244b31a3e2ef16e940a711724048ad507cd vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_reduction.h ++b680896b0ed17943650005204e0220ce278598e590a6ff333f2eb30a81b24c88 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_scaling_factor.h ++9afe6fb88c875215ace6ac7e5e1e5026502679b482686868fcd7489bff586c1d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_visitor.h ++31bb21295e90a7b8fb66aae4ed89226886c747e115788015487864f23a1c3c87 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h ++8208740e630bcca16e1c55bb42329864c590bb8ba1e8fbbf2beaa7d4ae73b3d2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/epilogue_workspace.h ++dee2f2104b80ddaa66851a12fa39bec9eeaba7d13fd889926c0a4f8bc6fa03fd vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp ++7d88b13213749257ad309c33f136866e054d5cb49868f3029e918d4746809938 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/fusion/visitor_compute.hpp ++07d9d5705f995783d96df75ba182289573a615a2c98949316e521dc9d5d41da4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp ++b1839d21edde831b405eaf64128cec6c78befe72f572b7256cbe774035bc95f7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp ++ac5b85998b22ce0600fdb03f77fd0947daf0881a46ff2368180a76e5e837b28d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/fusion/visitors.hpp ++6abbc6d2459fdf699c3ea347b3201d42a71e6b5b35e366c84c6205e39069e48f vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/interleaved_epilogue.h ++75bbef3d282aef1374d7adbdc8aa870b46620c6af78ecfbfbcdc9659e960b090 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/output_iterator_parameter.h ++8080c54418faa996e5a2aa2cb7def7ca4e380b5008756913197d55df9ea4b8b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/output_tile_thread_map.h ++ea396417bd649522e1492b5ff532dec40bacc1c2ef051f99a6366398f7a945fa vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator.h ++9a1e957f06f172dd9f7ce19baf2d378b804573c52d2fadc04d37d5b8ad7b8309 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_affine.h ++3f4dc1ee953a67a039d2a8e0baea405305a1c96a36e9348746399593a0e31ee0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h ++e269e000b2f21803e098f54c38b0e5a29dad0141f3941952a204fa16a6aa861f vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_blas3.h ++9d3b6e9c602056ce4f5c84be1dc5abefe8a0db31bd1bc8bb1c56c3094c0d7834 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_conv.h ++17199f7c09ba3d22a4d10534ad47d9ec22ab9904f448a28527752b33a1752416 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h ++e14565ba6023fccbb84b19989418197a2192aaa897c24d2635f6bea502de9887 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_params.h ++af12e2448f5e50bd0f37bbe42f0d582799719dd517b214286ef278b3177feca1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_predicates.h ++a8acd6922bfebd2294cd04dfe113c9bf935115e39ed6f7f0b45f8af99d1d1a1b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h ++0304d3ed2f9881ae8f64c2d0abffbfd277efb37e44165c1fede5f91ef5057038 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator.h ++665b75676311136dcb4ae5f486b27ab1582c200deb8dbe597962b0b5591f5bc5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator_mixed.h ++3f41714b23157def7b2d5211005a9c697fd034ef93619e4ac4f70859f4d8f69b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator_pitch_linear.h ++079e79d86c07ce5740da8b8d22970bdbc84744316e11740dc04c2e112bfeb86d vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h ++f2813506fd1ec8b3c4e644f3cb89c63a7bcbfc280ca0f362df5f406259fa9a25 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h ++b7892fdaec55d678b0b1bc21542de97e2e54b545d85fd798eab09e9fb798fa16 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_simt.h ++26aac0f3bda2172d58852bbd3d7ba549b4199765924baa31f6548f3aa0c8bc8e vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h ++db267b954bd80346ca3091f253518120e90cc659da4a680279a72dd95f7c6b00 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_volta_tensor_op.h ++ed51c4c4b1ebe4b012bada087a05be616d638b9189d925cac0a0eeda73d82090 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h ++ea00a65d7148cd0619cae32d1bea19a6613bd3a1db3bbcd38c7d0cde4bac4fab vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/simt_policy.h ++3c2a3d403f1c00e2074975f30094d25e615336660ad7edf61807e43aa95853db vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tensor_op_policy.h ++91b7064d936671278b976cefcbb5b805d5eedc45c8562a2a87d53a0c2615bcc6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tile_iterator_simt.h ++a60f4fe115eaaa5d907bd6aa1c920f263b5d1034b55375f69aa77f36c4f92de3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h ++626e067923955162090dc8257317644cdfc8e94992c8741378e9748940a7b8df vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h ++5d520da82dae8faa51a5551195ec4538fbc768f700b5060ed3430c48c90aaa59 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tile_iterator_volta_tensor_op.h ++28f650ff35496c39fe76de822e79ec8ac151d0cc80a93f56894bf2767ca7080b vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/tile_iterator_wmma_tensor_op.h ++1319d74d0a253a7b45c5c53600503504f4bc21e4059f5cdcb02005bf13c1e0e3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/volta_tensor_op_policy.h ++6f4aed66ed34c28b0caab2b1d1cd5fa2e2f9af0c478803409fc7ba7cb145e8e8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/epilogue/warp/wmma_tensor_op_policy.h ++0cd6b090cd1ffbe3e9096b15ec5d03765eae28e3beb6ec409f44ae323e6ab48f vllm/third_party/fmha_sm100/cutlass/include/cutlass/exmy_base.h ++db536ef2a1b1f3a2086e0aa838542de200d4df1cb2baf32faee4f8de9412eb25 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/device/detail.hpp ++e9b4c46ab9be69eb9b47e304d63c7d760816f0792101267bb6d1f9150c62f5a3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/device/dist_gemm_universal_wrapper.hpp ++f432c2bcddab633dc22f92d1645bb3be7ef5386df2fa9e67c582ebd3b95e93de vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/device/full_barrier.hpp ++2c1769484a04d666eec88c0202f5a20b7a6fdb867d6862dbafbb65952c577fd3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/kernel/detail.hpp ++5c404feb519d7228e2ac4d5490c867c3c64758986d7102710d3a4157f60e3058 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/kernel/dist_gemm_kernel_wrapper.hpp ++194b2fe25945d72f5ff52ed4feb27b512478d5023124f18cb6caaaa9a44748f1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/kernel/full_barrier.hpp ++9c4b713676d53efa2fea5d1bdef21682418befce4ffa5628123118d897190fc9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/schedules/dist_gemm_1d_schedules.hpp ++231f9cdf069cc25c6c95303173f0033ce436772d703d04e307b222c3145e3d48 vllm/third_party/fmha_sm100/cutlass/include/cutlass/experimental/distributed/schedules/dist_gemm_base_schedule.hpp ++f51dcb1972da008666c2a3bb24ddc20a192c3bb3bc3bc553b29d04fc621a1258 vllm/third_party/fmha_sm100/cutlass/include/cutlass/fast_math.h ++a603f6f38374dc98a4fc18e451d23ebe448b7467f23a4258706ae4a2580c9fa5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/float8.h ++fd47999da1bc4c7add3e9a378539265c5be6f15e478eb035fa25d87e5d0f59dc vllm/third_party/fmha_sm100/cutlass/include/cutlass/float_subbyte.h ++9e460f494422cc738a3d9c224df0ff89079002eec477b99cfcb475033d784c2c vllm/third_party/fmha_sm100/cutlass/include/cutlass/floating_point_nvrtc.h ++440f47886e4677973e3552ca0af605f53476f8f3d722bbf8f47e9203f721bf7f vllm/third_party/fmha_sm100/cutlass/include/cutlass/functional.h ++d281fb8dcf53461c61727be913e7ecedd86f36c0285d40a97a350c91d78c47f8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl ++5673f7edbb9b2a021daeecf12a582077a8877ba06edd216490271f423b186238 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl ++2cd0f7772a8b1eed2e9114af0a334b083e2899f6fd08ff0f90b5078061116779 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl ++c1f67ed5fef152bf8d5010e8ae1a4082911c3febdbdfd2f660b3d730f66e9489 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl ++f933a1b3b09446883dba5af61b733fbe2eebb1806d7f5dc3f933a3cbc51d9203 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl ++6301d7b4eb6bb45e82954418d5b253bdd02c0ef7aa520fa5daa6ab8d64d0bf58 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_common.inl ++66c8944f32f540d3a13ec4658388392aee597fbbe734ea631ed472a50833bd94 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl ++602d1f8d5d794d12fbe7c2c81826300c27e30b4a74923256beed64ad804400d3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl ++4285f6b815f1ba759616644b97e0c02b60038077a7fa9f5cbbada6223c76459a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl ++149d50b99d63b56174ea368aee301a0c0f8fecc768a29a781f0904e1faa73cd6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_pipeline_carveout.inl ++04cd628dcd40ff502a57f066385dcd5a835afaafa97d6413786fefea482f7918 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_simt_builder.inl ++22f59f9a4af142aa8f4a1204621e315bc9eac6aee96064f1df321b017ce4a2a5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl ++4024fdc1a60c1c72c0993eb30f5cb96e14ca6f262b734b3f55f3421fdf501f6c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm100_umma_builder.inl ++52b9fc5cfa2457ac138848da5d418f3ddadd834efb78c878b8a2cfde7eefa24c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl ++8f75ade68b1f809277cd7b6938616bac4291d4bf806d7fbcf3272de68a4642ba vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl ++4bdb5ba9d6b8760c1d743abe9e89483fc2ff26209cc952b7ca8bd137aadd26c8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_blockscaled_sparse_mma_builder.inl ++97278d844a31eec0c828221c4b5cb5cafafab8f37fab69699edae833a07ddcbe vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl ++b47a6c916b6ee17cecad6c39b47a64480075ce229703b635c5b3965e304ee1aa vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_common.inl ++617fd924767e4793d6450878c6baa6f5efd99776b5463afe3e1181280162908f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_mma_builder.inl ++959c11c708e09dfccfbd6e2852e3bc95761916df4491011e73e71c40a2674d5a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl ++c462f487a30b83ee9d18e60871bb78a61dd96c5c4dd53349abb25c0b1e6fb137 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm1xx_common.inl ++046243d4cb7b73482e620eba485ec9a50bfcecfd890ed6bc8bc1fdde2890550f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm1xx_sparse_config.inl ++e8eab1d8c31296092e132d5f0a0111dece4c44525de4a6c49ee0e7213d8c9f65 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm90_common.inl ++b28be78cbc280a7458f25bd9416382d2f60809ebde836f8dd2e59ba01f9c65d9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl ++a70c1888246903369b0971072b30782cf003b8463f8d6f60d990132374ded66a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm90_sparse_config.inl ++2b00d220f9d1a759fd6e41d1e3ba6ce5167c7f69ae6c6d9999bdb20da600f2a6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/builders/sm90_sparse_gmma_builder.inl ++0088a3e7ef77f12200ba90246a00c0980f29a37f964e0c11d0979f87a100acea vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/collective_builder.hpp ++c94ae031ec110d4299ab46f70ef760debfa2ddc4eed8999f12678717a858c399 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/collective_builder_decl.hpp ++d2dbedff13703f457607e0c3a1a65ec9fbb8515ce7816da5becc642572671841 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/collective_mma.hpp ++d41eac3af233524f508116d97e3b6fd542950100707360d42be0881bbfd63d6f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/collective_mma_decl.hpp ++cdf0cd6125e1c3d0d6f83d5b381e0c08a9f51a31d7b1490b812cea98642163e4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/fp8_accumulation.hpp ++341fc2b5c1d49600b9ef6a8b62909aed088acb5ba674771ceedb0819a017d809 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp ++fdbb7bab18438c0ca026535f7699fc223b464e6e1d7e6abe942d1e5243c65bf6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp ++900b60de1710420bef27a325ee4b952374557a83b33167d015f12ea092cdd0c4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp ++6ff8e3a43b8cd560db7b332e45c3236adcbaa98bcdafe45544dcc59e62dbc3ed vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp ++29ebad22fdb0e5c07c78557cae574ac22cd5dde1daab98a26735da50160bafcf vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp ++0bb9b4c519922b84ab5a60de95d17e82643f1670222a78bbb0fbef37ac9e9445 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp ++208040d3d8302674f6687fd8b89f34ee727539ec85716fb8fccb4b425c0eb1dc vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp ++c5317e922f5af59a98fe3505cab800ed2330efeac1ed969665851cdb57e9df34 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp ++1d1041b7d62dd0bb5b9848ecabf6eca63188255d5f9d8ff6b7db406ba3f5e4d6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_rcggemm.hpp ++33aa0e16bb9ca5cc77a4c4f8395fb0f7cc33ab44c52defacddea6386bc1323c4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp ++e34f47efa90ba427516c30e7bd4df71c20e02c6455794d9c7b003fdcc79b698a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp ++280808ced7afe8d692168a321cf88e68e8431736eb00bed30135517cd8c4d865 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_warpspecialized.hpp ++ddb9540b67ab79df2f12e0490610332d2659841d8cf858986d4199e3a5fb67cd vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp ++9f63751f67a7bc738c51c96495d1423b475cf4f418e632567aed4efebed6b1b0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp ++2725a01eba9ff6b51bec30665e8b4237dff6e38febaed04f38d2e108af44a127 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp ++b88891e76af11f7ecf39366dcff411f6de41ef2b465e53bb33794e16ba31fdb8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp ++06f5ab0ad59fa596bf874517a11395c9fd480c0a94c14f6a9a9b5b00d9b522d0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp ++68eb8d909ccb901368b0a894d55f38196978fc4503cefca7ac163742ae4044b1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp ++3da9fc4cc4cc0b1f68fe4cc7ce4d8c50aaaa6d94fed39505b07d9858ec7ce532 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp ++891d42caf7c3458299b3259e6bc82542ecb1f831bbcb8bf760ad64975c608c48 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp ++9ec5e394c26b53b9719cda84bba42438137fa48dbf68e153f8f308878ac3b019 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp ++2461a7162570fcaced101d379ab3f60fc45e23fa603ec7913cd1e94ed7272967 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp ++132711ad3a120e7ad1fd118e7efceeabe97e8b8eaf7ee715179fca408579d7b5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_mma_tma.hpp ++a2e7164550041886dce3a54a49e3f778f46d6f2e8847c4fa629922f10004696c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp ++6ad6440369c390d600219194400b3d87599a455cdf9f0ae430ced3638794da47 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp ++15cd9c2245ce610e02b670f1845a3e6be042094e610a34ea6291838d63e84c1c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm70_mma_twostage.hpp ++30e1cf55e1d4730f46b59ec0f38a10793a57aed267e3b6562e11dda5fa1f4bcb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm80_mma_array_multistage.hpp ++0ee20617304f361518fceda53b0adc31e5fbbbcb8e467a1408b2e8e24e45535d vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm80_mma_multistage.hpp ++ae3a9a82c245c068e30d71dae6f6877bb02080e1914d168a10589107b0f059b3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp ++1739d592a6bcb83f4e2c9c498bc05ee0711cef172957df40641a2cfe25c15c98 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp ++ae27e7bdcea4b8b5257b5daa44c909e26ba6d47cca973868ad53410645444147 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp ++aefb7a4c2eb601ff29c32d5bf0d432dcdad020edcd06c368d80e154bc05793ff vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp ++67bb37a1c886e15e23da5b7b29cea72bfec240e60b8d41c304ea96e7914ec5d7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp ++e2e855ca8d8bf1dba87af3bedf3496975b7bd00bf7bdb7a9d980017d91364f5e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp ++b5f73f094b11a5cab41459b3738879916d1db7b92d8e29e75acf4872b15f4fca vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp ++2abc2647c144d0bff0a970ccfc3f7fe6b3bb7afd365d3e5006f3ff08c27a3baf vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp ++b2460d4ebb9fc6d5cc2ecbec302c0dab78b9e4424402292a6fa02c565e19f3ce vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp ++258a745d1d74afd84374eb66aeabb8f35b4c973b95eeb2c4bafb0a05aaf5da43 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp ++8ffefc316c6dad732b88111b1c2d912adda181c7c91dcb4ffbb48cbad9aa29c1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp ++02603a3be25439042a8f63e95cccb955dbf71b10a6f75027b073de1bdca7361f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp ++87ff633b11d98d88c256e66011174772823723ed3f694bc01482788aed8fd764 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp ++b70e168d6fadb1f834298f60f0f5145307f816cca3f6590f31ac5b37b0fede32 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp ++60b726003a226e2bcb363ff5157ac4a848973f29ccf632c9aa4e04c189f906d6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/base_grouped.h ++3850d5933fff925707b12d5260affc594523859aaadb97ab5d55a91f484076b7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/default_gemm_configuration.h ++6fd4a0660610ebe3cd3f1227cf3fea222e928d63e013fdd9483aee109cbcfb1f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/ell_gemm.h ++09b6452d552d840853a461272655031f5e548d5543d6cc96010fdeafa82a2e32 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm.h ++f1598d660028ffb4f48fb0524e4c7859afafb90652bde5315ebdddcfa3a4a3d0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_array.h ++b308ffaf3d89b51ee60a5f8482343b652d6a68bf1c223b4ecb0d1f74850fa6ad vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_batched.h ++89e9f3fca903eff4d213bbe5ca9a55e369ec79b12133c000be3566705c8c2583 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_complex.h ++60775ad07c45851c55f1a84c96e9e25fbdae29234db18270e2c3b5efbe9b1657 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_grouped.h ++591bfc9f609a10324cbfa7c829b30143e8e43e34b28cd508acef8f2c0e6da241 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_layernorm_mainloop_fusion.h ++86f06c690cfd6f2f5f779fec1dcbf13c39589791f319b021986e184a19389fb3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_sparse.h ++0154c9d165f5b77d2e494624199ee2210f41f28694f334c3b9ebc7d82199cf63 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_sparse_universal.h ++c679422a528715b22e627d20d89473048f2da1ea9f2f5f06cb997c9dec9e0ba8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_sparse_universal_with_absmax.h ++e326672b61fd5d108dfc3055e0dfc1695a324fc6604a92a88ab7329a68a8fcd4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_sparse_with_absmax.h ++f602857a5c9716b017e704ed5c1035be0c917f97937cb7b2d4cc8b93a519cbc9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_sparse_with_visitor.h ++e6364cae26f816491d7389332cc90cc1603017799ea5da0934fdb64b1b266c10 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_splitk_parallel.h ++99147aaf43d012f4adb99c3765b8f7d5bf8e76d71f5071b8aa8a1cb52ade8286 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal.h ++5c7bfb62c1b063b17471dcd98012ebdc5f98fb88f689fc4882bad6787ef94716 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal_adapter.h ++4fb6960565db62922c8642681b1e1aabbda359fbd7dc9ea6386c2305380733bd vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal_base.h ++273909e7ee195e2fd25f7a8af44042ad402dcb1da55a3924699e5dd7a974434b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal_streamk_with_broadcast.h ++373b17062ad2ee1f492ce406ef4927e1117a689c14c54fb3b1a41e2ab3be1743 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal_with_absmax.h ++554b651a76629100aaef995bde7a57cda72cf77a37606910c941288ab766c9d8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_universal_with_broadcast.h ++02689961541a9842cc4c8d36d4c3d3de860a73258561324ae560c16eef1f9a46 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemm_with_k_reduction.h ++c2d975f0baece0dfc287c35a67a35671bddb1e432d5c85c1a10f9af48fb4d212 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemv.h ++cd81f392b0309612e35e408c7d11966ed51580bb1847d879e0ee961b6c8bb417 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/gemv_blockscaled.h ++5fe83bb7d386696015d25064cb232c6eebd673f8d6570126d5e54c7b4bd53242 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/rank_2k.h ++e0ad11b36f255cc03910c0df21db443912ed73b825793dd58574fa11cbc79add vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/rank_2k_grouped.h ++d7156ece53365468fa113e5e31df2e3b28a653668d20e43a1aecbbd7a34d333e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/rank_k.h ++856aa9bf2259d29bd78aa3f8e067aac84da35b015e160a17aac8a7c871d4a91c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/symm.h ++793d8a9bc97a255d406af9ccd25103617167e2ea87afcbe8bd23c40a173105cb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/device/trmm.h ++95fbfc361b69a17f60e34ae9a1900741cb667fc5761038c4782cfdbf2f2b7cd4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/dispatch_policy.hpp ++bc3413d897c767a0aba7dc1f5b1954cb3e1f3367c4fbc4018596b4639827fc6e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/gemm.h ++93833fb86b1a2bf87f0ec4b64d51d2d5f65399ea40ba84cdef2e8b3489a8c144 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/gemm_enumerated_types.h ++d019feb35da34e621ca74192fba4bee973e91770352b4ba7d28ae054b3b87618 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/group_array_problem_shape.hpp ++9755352fe35ed312a6198cd5d526640cfe60938752ade83e6933d445fd37e9c6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_ell_gemm.h ++6bec4b81e9502ed0c7f1af5b8305ee83ffae68c95222c565784e8cdd510aba50 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm.h ++c6f269fb477c6ca20f38fd570079dc80abb88bd948605c8c2c5820d9889a28ff vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_complex.h ++397f99226768c8999ae89265d81e7a410109a0c0e5a8a6e6f9afb4199a50758f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_grouped.h ++586261a89487519bcefd1c86c7d332c4af6e9176dfd4ed12cba7734654ea49df vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_grouped_per_group_scale.h ++648b4c2cc8ceea91bfe3eed7d9f163be57e4b5e187c2d5b8942f100ab3ec4d73 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h ++887e9510a5ca85ac6e0ea696698656fdabb4d71e2bc8ea4b44bbd4dc9f40d91f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h ++2a3b4023f5802423d92487a2e4c2d323a0403527635ea23d19416dfc44d207e2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_planar_complex_universal.h ++3f73d988e6df161e79259dcde374fbd3383104f8fa36dbd928e01f1f6cfec9ea vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_sparse.h ++776c280e6e368be2b6bf578855e641dbeffc94db2269ee998b6bd39d106be0bc vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_sparse_universal.h ++ab3f0b7e4df7fc04761e86720af3799e8796742bf4fd4ec52c42447ddb6a55ee vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_sparse_universal_with_absmax.h ++d873ac5410a97c53ca529e764d95f948c1f4147e7a9ae72217259b23c94bc3d8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_sparse_with_absmax.h ++10b487527db8b60be0c2f53050659146ac4e9ac95c530b9c4eb7b89972edeb2c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_sparse_with_visitor.h ++0403ccfe2032d65e9a9dd444c5ecadb7a8d2d17d7e6c47bef79e93f7c0c6280d vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_splitk_parallel.h ++e5eab86b364fb0fce3b457035a8b9de1c2673c67ae48cd80cf18133f18522e59 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_streamk_with_broadcast.h ++38bc8813c5801a049f6b6d778d6becc3a12470cd32a86c70e748676734c35693 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_universal.h ++221cb62cda94f08cb698fd9f9c0d231cc10334e9871d0aec5737c3cae33c77a7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_universal_with_visitor.h ++6cd04920abfff5dfb989c68281ed9ffa2a6a3e1adb00cec0699beb0e71611005 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_with_absmax.h ++efd0f4f8e377b2d333cdeeeb1401d83c1d65b1d58ef97273e13e344b0a6c31be vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_with_broadcast.h ++c31bcdb36522ce8dcc7a823fa90584c6e86ff527702c39124dbb47b3b33ba03e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_with_k_reduction.h ++655eb1081a835db66d550621e94e8728e8ca0d968131b6ebb9f54c75c47057fa vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemm_with_reduction.h ++3dccacd5836230cd219ecc8afdd8f24f3c7c4124ef1e0f3f2b4979f0f6a97b26 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_gemv.h ++b81b877f9dd4277d27b6e5d2b6c6901cde7c1506eabdf0e4117807dff6f08630 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_2k.h ++aa46bb1be238d8b28eb2f7a64176808a3f14f62892b02fdc0f5c6aaa78fd9187 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_2k_complex.h ++f30e0b0d55dcb537223c31e7209464be88ba164acb4cc06356f6a83a2bbacee9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_2k_grouped.h ++0135956653b8df0f7060a49f41c42ae79267690ac469c445c02c36b1924a8dcc vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_2k_universal.h ++63712f163d7ec04e81fde6641a96f078a840c01fe5363d4bf7df1d066d80dce1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_k.h ++ef0807cce11fdc43087911be2a5352bddd834a3bf9365918b0f3a23bef786c94 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_k_complex.h ++d9e711548109f6d8d4f31667c5dad6c605ddb576aeeea424877bd4f87a142ef5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_rank_k_universal.h ++6fa42ef0a1d2c7fba24094a163fd014176ba86d4ac828a513b7d9386d5c3810e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_symm.h ++0727fa574477edd63836d87a59813023f3d4b019d37f941fd31bcb8e044d15e5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_symm_complex.h ++bb2b75b3aa934e297d55651d02f5d22165cd5ff374e6fbccfbbec344bb047695 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_symm_universal.h ++327b137d5462b3741a07aa57aa955260a4c093bbde14fba2d2be663202e5f7c2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_trmm.h ++70ddf3a2ce56d41341b34e769c1437f7d415913877d9ca4e4238ee08ea269037 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_trmm_complex.h ++99595ffbeac6d25f38a0914648f77fd9563512d597d795f0634f3eff52276e50 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/default_trmm_universal.h ++1c6be29679b7c1bd5301cf0229a2016d45679ff973e0e8747efa2524ba6bb19b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/ell_gemm.h ++2ef84bbf3f223518cba2f97bd6bcfe8b30d9c3fc57519f56f6f56b7663bf21ec vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm.h ++aa8032cd4b63bdffbf358ccc98f06b8f1cffe4730c8bc9a188f1bc566658f8e3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_array.h ++aea47ba7bb13d956f57fa971e6fc2122113fb5e94f743dfd8b83606f445e496e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_batched.h ++c12637231e08519cbaa18fcc739e5db60145cc504bf9995b7c1b8a66d66ab2e4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_grouped.h ++79ec6bb93b4eeb2f806141a99515381ab57e5b6cdc954d500468e98a7eb418d7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_grouped_per_group_scale.h ++0d3544883cebe2ef1ee24fe77d21c824a1d38f5ef10be25445c5fc7f28cc4334 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_grouped_problem_visitor.h ++b8ae2f3533e6d37c12104abcf204e5580207fc41438d0446b6e6c96c379e43cb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h ++5bc748e2671b6fd1206506530435c5d8fba63c2f41405dbc7b6afa670d4f997d vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h ++2687b6e590f5745b98c3bb47ceb11369a5d38f06b249b5794be1c8d09c7b29d4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_params.h ++fcbb4e594a6871e14eb02bc47c815f9190435b0ecffac0f7be8c28c7c4c8b0a3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_pipelined.h ++4f44613cf307a47bda34bbc6a7adcaa6637fb2b6a4a2a47f6b109689015eda1e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_planar_complex.h ++19e4949b01ded3625c28b01714d0de31b88cade692aefae09e71b431160180db vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_planar_complex_array.h ++cfe94c6955795cb616289d1fc08ad7bbe3f00062d7defef6fb0dccc4a593e78a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_sparse_universal.h ++ade6f2666ddd09abc8836d2abbcd11f7f5ace2b049fe77e32ff27f100856cef4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_sparse_universal_with_absmax.h ++37f4a64bb95a845de8b10ed15411d7f148c51b23ef0c21dec53e394812dbeceb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_splitk_parallel.h ++ff9aa6f87f4335714ee14e0dd8851f96dedff5651224b5ad2329b1672ecaa6b3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h ++ba7889f80af1cacb81426ca78997e4cb99c205860eb5ac533e25e734e003c081 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_transpose_operands.h ++1d1fe23d25674e3559e73b568f5d2b0b44c33baa9923e4d5e3d038db8ce79f97 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal.h ++6fd71ed541507628ea7cbc5cd9993d7cb500168d86ec3158fe6f7932bec5e925 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal.hpp ++d40c2a481ad3a1ce385fa14db85f332487efc8fb3784225ae503fb3f9aa2546a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal_decl.h ++da1bff85c6d2fb73c9ac41ba325387aa08ed60f3f1a67d31a5ab8b84245061f3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal_streamk.h ++b3460e90391c6ab475057a6de130b8e79ad8320aa3cc3a938a20869b41fb6bf8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal_with_visitor.h ++97e901fb0f465861c4cdeef53b3992da801b56679c3f8b92204c0463857d89ba vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h ++945b76e0d94e9a63425d4de4928d22f8f98959e4286708e9cd8a1572e7c19b7b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_with_absmax.h ++74d89f0852bf4890587f61e4ce0c43612bceb855d63597dcb5ae9cac08b75327 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h ++30209eebd39bfcba7e74b670fc41cceb059f5c4b0aebc1f4bbed8a0ab0604b35 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemm_with_k_reduction.h ++522ac8d06b3dd86985a69aab2f8b9a39bcdb2e76924689779f4b5ab7ac2b6f58 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemv.h ++07d087978ad64bc2c4c5fed27551665acd6241f097def29ffef24c347617fa0f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemv_batched_strided.h ++b8db28680773ba9119bfb66ed9b30d53bc281a993718e1d84317a4f7584c57b9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/gemv_blockscaled.h ++3629375a2f2837cca3be176dda60afbf206ffcfc29fed264b6664a78ee2d503a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/grouped_problem_visitor.h ++f986082734e04e52fafbe745f6acdc3fa1862455e01e3fe507425d84977fc923 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/params_sparse_base.h ++172ef99d535082469ef8c3b4cca62867d9df730ec5a8a046b143221770282c12 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/params_universal_base.h ++5bf56b48bba7a3ebabc52d42bcd5321acdae92de3b9c351970da0da4eded693f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/rank_2k_grouped.h ++1ebd5ba99443daafb67ed83d64fcd31f3b5cefec7e27cb457ea498740eefa516 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/rank_2k_grouped_problem_visitor.h ++1bc69b8eb25b423e5e590a7d32d080b8eebbcdee7e7aa65a51811424a46b7d76 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/rank_2k_transpose_operands.h ++5b6b3924bf543364f809efe33b79006936828bc7be633e1a8ee1e1638d2dd5d8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/rank_2k_universal.h ++f03b88e863a03fa1a459a3cfc6e76b3fe6632808d0c4f3b4ddf963f2abdb6b69 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/rank_k_universal.h ++620de2b694dea7a6d33f0ad89bbcd001b4200dfbd3ff4c045d338586ed613880 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp ++a259214c43f8ec83da356283e09dd603db762608d96d12b4d9f1d1f14e7a81c0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp ++ef9d86ee75e65133350c09f06a4a54519807f465ffd20e9eef57e9eac844c1ec vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp ++50a8bedc3aaab5437ed70f52f5864bc4abde2946b8f5f2b82ada8c8b683a8159 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp ++88705d0f8c3f2a3841d4e097213772d2e01b8c97fa5c8def7f6c81e73a59228d vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp ++6fb1811d77f61c7dfc20044fafe7a82a3f83e1597f5c335e5a68d774d92206ff vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized.hpp ++d26c324a8cffdfbf2332544897306b0e7baae764eae53586401dd5e04d1ff4b8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp ++e795b23f8c42ff208a3b77225cba28210271a9352b69e2a928c17aa8e4bcd9c8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mixed_input_transform.hpp ++0f048b1eef419c6f546fd7cfd19e100b202684b5ce1070907047fcad3990d1bd vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mma_transform.hpp ++56723c6e12db6b0754cad569fd359f3d1e1156416a24c485afe9992d53de9f71 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp ++d49263fffc4bbd04cf94cad6ba6d2f462a5b4a55e3716d86d3cf5be804fa0c85 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_static_tile_scheduler.hpp ++c60d4784bbcc26c2a12f4a2fdbd08d4d83252f8fcc976f67eb8b2f683d54b2f1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp ++f4f25b67a1e8b88a4f55e72dc050d66f6c6149e1c40115a1cb4751c53705afcf vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp ++6794466d39adf93d3b239ff807c00cce2a9ba3a3d0df771e304964a93cc678ac vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp ++2068f080ac6667a65690f54033bafa92a68b75569cf922d28d6b1e305a1ee4d8 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp ++8c3268ce0bbc53d441e6f7d0582f308ccf1eab6306180412b4fe1e8f53d5b3b1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp ++72471aff9e853f89b0091a0cac7978532071faa8cfd657e654b474d99a19514e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp ++226d7f5c94c212b533ba9f2b1978ba1ad714d2c0b9f4361f9a977039dec3a758 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm70_gemm.hpp ++0a86c9e8a3449de52889110df2d9b35ae35f5ac3969ea149e0e612b9130dcc03 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm70_gemm_array.hpp ++937368e4a0d0dfc80963aac02c2377a006a3b8ba0ee45bd655407ca75af6bca4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp ++b29b69a775f39a2c2904ea467b2aa1dd8d8fb45a48ccc475abb88e580e9b561e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp ++d858b03e4286ce357c60f27688597b0ed892b79bc093858d7e4d3bc85c0cc1d3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_tma.hpp ++9fd7ebcb45fd38ad336645cbe82d4757df4775cc01f1dda5a2a2b57264792c8f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp ++f49505874e5ad19aba9862347495216e1a761e64369292066aeef33813328bcd vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp ++a9754682fefe667b5785fe1698143d74423aef3d13d6a2b1df290e7202b13f5a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp ++0b59907193a7a26efa25e3d5304b1860e6d8807d88f9f0032b5aef4a1c5c9534 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp ++5814e340cd810f876ef610012912e2b519365af551acf1f09144a9c2f0d68a58 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp ++ea9d98466e04bdc630c8b46974e3e838e908f748b576a444cb82caf8d4f0e359 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp ++4a714b9833336f419e9b9e2267bb821a483863db0ee245d5a08c63e37ed3ad08 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp ++ecb81ceaba9a88e07e996a9f1aebcbb150c7534c3fae410c551df1bf48798e5e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp ++7d6544a3a476bc75e93a88b147a326fcc2fc01133254ff99e39c4e47fc164ab0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp ++04dc8fa9b12749d7ddfa51adbd73248f5cfbeddae7e1460116cc6e4c24280b8a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sparse_gemm.h ++ac98673b0972425b561ade60116857c717cd8679a5c9470e5e90a3c8bf01f3fe vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sparse_gemm_with_absmax.h ++95b0e4cb84de98e066dbeba1045012db57f17598fa63691a26669ada2752b45f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/sparse_gemm_with_visitor.h ++cb3110f9b2ee925e185b69ff8aab1682a6b07a43a8643d2d532bf832feb15899 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/static_tile_scheduler.hpp ++04058b7b6dfa4dcf39a627ef12f9fbae10c1d25aa490d5884f6b865cd3c68f07 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/symm_universal.h ++ae25f8ab4b75b895a8419f055dbcf1712b1bef9e4d9b5f64eba176901f2672a1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/tile_scheduler.hpp ++79568a964ef867f7e9a716f7a09a28c0e873b75f028df4bb9c8196c413796738 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/tile_scheduler_detail.hpp ++6d06370136bd2e261583e035bf23cb77295973df553795dd4bd764829a18e50a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/tile_scheduler_params.h ++da347406c11b21849804d294a20fd4c62255472b93eaaf8b23f9947f7baaaff4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/kernel/trmm_universal.h ++130c112fb274b71fefb87d63edd6a964e7fd9e7b56bdffd027a5d9e02bcf36ee vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/thread/mma.h ++d15be778ba0f05283b382ea143efaaff6072b73cb066b963a32ab1dd1c0fc206 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/thread/mma_sm50.h ++b0551820e736d19b5a15efc5b5e095ecb8c72c015d53c5188c20f771448f4174 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/thread/mma_sm60.h ++ad7a2a5dd4e99639fccbdb026343d934dd29d3781979a849114196b732aa2272 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/thread/mma_sm61.h ++1b791b41097c586338557ecca950943d0c40dd96f7a7eabb8be8941127bb33b4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_ell_mma.h ++130bbbbb60a3dbd3a0cb8ba131904a7ca5f2d30be014d1bfaa76f453147902a0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_gemv_core.h ++0ac1ced6ab44558921aec22b8592949f060fbcd50974ec10a74175dce95c5477 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma.h ++5ce37063015a5d7a86e0b5f255a6b7f815c5e7f6ea16ee7e7ca0a5d8f6f88a2f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core.h ++7de1d009e3f6eddfde47f72ee3aefe338ad387ca3f508c4cb269a5160ee7b746 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_simt.h ++34d28ac60196b90785ee88444f67ff954d2cdf2d723dd50dea21558febc7f23f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_sm70.h ++80659cf69513641a3f72b1d9ff9c152c303b9746a8ee32bf5337d8566e1d76e4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_sm75.h ++ab3b955a3e60e6b9a2fb1221bdc5759c2cc68adb5a7f32da29280474421db268 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_sm80.h ++20b2e482bde96a3bc83709cbdd27501815e07b8bdabfe071297921fca4ff5794 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_sparse_sm80.h ++640d6333f2e710decd1e4a7c508ab602cf45e2d02f457321771b02b30aef2f35 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_with_access_size.h ++c1e0dbc8acb3985711e014c55ad76e5f6a8b97c04763a89cd28ba075bc4cbca0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_with_reduction.h ++29a6cb0db7c373bd22d3f5283cefbee9c8c6d5bd9287e09b919752859ce9034f vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_core_wmma.h ++8f684c3a7a806dc6a893b2f6c0d3aa35911220ffe6db638fd0da52067be26170 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h ++8982fce3dfcc42a1e1c3f772ffc29e70f877b9bf16b180769a57b0e221f83b00 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_planar_complex_multistage.h ++a3c4afb3358140b421f1d4eaa86ef2feda729e9e4da2c05724ec4520465047ad vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_planar_complex_pipelined.h ++e15555be54f49462fe68998b9174651242ab7abf0b4b01f8c0b10ff425296f0b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h ++1aacf38f0b05e1c160d6f75ceabd9629d2771e4103a98d961d3a317fdfd129bd vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_mma_with_reduction.h ++f1c4641f5d5aa94113d0863afe2ff7233cd5f2f4d6e50528ce7d0e6cac928bd9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_multistage_mma_complex.h ++b6fad473d3d1c661f22410424731aa65256463e83d5879a98025d9ee0cee418c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core.h ++a03739d0fd1d0c4eac8ffa7d5d50d699017eba359541ba7d8e9620ca5ed7f9ed vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h ++681f2d18e63bc55cc2f2ffa658dbc4bec11acc9aab4fd8d0951ccf1b9be53510 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_multistage_trmm_complex.h ++15836042c234e3d640369e12bb66672657afe1db45bc428c862c689b03d88838 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_sparse_mma.h ++663325e77c3020ff6fbd4b166b97305c5b9901eb19e70731b76824222e25c2f9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/default_trmm.h ++b557b0699393d680ef622a466a353442ecfea165725ec91cc17110fe2494e073 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/ell_mma_multistage.h ++6e3acbcaf297587e876af7fe4c3a093168a270555624ea84feafa4cb4eba550c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/ell_mma_pipelined.h ++f349d3de8286c5800567da6e77c23067e0a1a4554099c5c0c5ebb58db8c03100 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/gemv.h ++3f54c9134212c45a2bb328febcf60c6f3d62aa477b290f214ab7154c442c2906 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/index_remat.h ++b912b7b24e246d0b77428e83a86575143b1191b6334b5d30877498819a40fdb6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_base.h ++235f7f2c260e5f4e488226cb80df790a2881f91e15920bdb135a8cf644616d49 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_blas3_multistage.h ++a6b843a3b81b057f4c5f661c6f470147731153f9b63cdaae510134731a5e0ed6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h ++3f89e410ca96b06e4a56ba19f90399e8fdd2744a28c45d458a60940708b4171b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_multistage.h ++fac90ae04135ccacc0ea48440d20c1059c555ad6355a8608df0e6a74a8a7b4e5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_pipelined.h ++e96095627fce77ea514c727117d34f450b4541bbd90c33386de35872fed1bacb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_planar_complex_base.h ++a4fa9a48763d2526b144a29ad081fc9d7ed473126fc188e4b276092d6863f5c4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_planar_complex_multistage.h ++c0f57cd888e54347615b9c24a7417f7c7808d01d70a4bb72d158a0bbaaa9a0f3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_planar_complex_pipelined.h ++07582ba2aa5ed85c58f1c6a34c1f0082723525520c68feac631b737181f85564 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_singlestage.h ++518e893f1a9486f38386983ea164c28db33c79ca43d0c75d4e1131040ec32904 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h ++3a390fee174e2e6be5be5d502175cc91edaa6b28488d0791a81ea85b062cd6a0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_sparse_base.h ++34f8b8fda11accfc6d4e0b113b8c86a08a0542b04bb0cd67d0b119f97cace00e vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_sparse_multistage.h ++4c97e30b9d59b7d5f95e45d3c5133a488d766f41e474ce451ff20da4436fda22 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/mma_with_reduction_multistage.h ++6d68c3d8abf7aaca1337b36cce3724934bec551d02127e7973486ce916095579 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/threadblock_swizzle.h ++2ed84c51c9752fb6c1af334f86aba49be50fdb4de1c1c0cf740110871e70e6eb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/threadblock/threadblock_swizzle_streamk.h ++6d865922a45cb9bb5e6aa3da983fd169d901236d189e081458b7c56916f004da vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_complex_tensor_op.h ++4a364094e1c91a9889e94ba4209ad460a9807ac3596770113c638c7632ae3c7b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_sparse_tensor_op.h ++46d2d8299db997c3aed1b5e4bed4f2903c06cd3b882ec56c9d1187b5c34bdee2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_tensor_op.h ++b8ff50c262b391d399ed92fc0ceb513cd4ad62a3c8b6a228a53f71305c6c475b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_tensor_op_sm80.h ++776cb088a6847fe33a28722072859deab57d58c00b0c1405eb5960fbfd459e58 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_with_reduction_tensor_op.h ++916de96d75a1ea5aaf4dd812e3e01d3863a842393c6e191eb0c9074137efa303 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/default_mma_wmma_tensor_op.h ++53a77c5e4a5f15daef007c33d928cb40523d73e06047776b930ab3321d8e2bd2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/layernorm_scale_bias_transform.h ++f972411fdcdcf8274baad170b9afb61403053fba68b4ebdb1e3bbfacd9d83563 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma.h ++063ef7bf36e08487e060efe76f03a9a107920b6b99bb60e50ba18bcfa0b675d3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_complex_tensor_op.h ++5dd3d3e0696b4ab4c21577a2233a56cd22cbd8dc6cfdadf18a0ce212d05a2a21 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h ++91ebd0266d2a649a6bf7f4fc8d2d875c172348d9e28affb273b76670d0c0c2b6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h ++4414c02af6fa60ae6b75209ecc20cf95faa763c9790a862dab2c41308163b45b vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h ++f41176552ebdc6e183b3c4e26f7e7e6325950c6706d053609de0c319dbcdabd0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h ++0a27c6943e45127dec62c49cf8185dbe44ea54460ba1a9556eee3939c28989bb vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_mixed_input_tensor_op.h ++ed11b8a67dc7471aa236019fba88c82ece41e85e9b35a0a74cb0fc544f7bfc6c vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_planar_complex.h ++dc2e3b55842dd235f790e6b1c00528ddf0066d0f43a5a609c94da0dc07fc6002 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_simt.h ++0469c741a85400a6702a9910584ae1d2546d5f238458cb2a66b3af97ea1078ab vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_simt_policy.h ++95031c798a5df4e10192dcb37a27b030874a70d0c45459c9544a7b49165cbb67 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_simt_tile_iterator.h ++9aed1a437c480d0bc4fec436d63a3c526c3173ef8bbb8a26cd320354aa342d1a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_sparse_tensor_op.h ++6642025020444930188fee16b9870c7edee249736e2213432aa26dac0a9eee51 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op.h ++1707213d882effb724f8f1642b77c84334566ea188c146b99d39eafb27e8b882 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h ++9ddc0c943eda0a50f7fa32b5b4a01008b4ebc6b22d38ae80b732ef44ea2cb174 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h ++65e570dda86e8c87b0584668afea1409b2996f9f2c7ecdb55a9a334cbde8d121 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_policy.h ++49d9b316e52eacd1365a5061cba568ff2220de967ae6c221335bc851d4f5a865 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_sm70.h ++09afc49bf994244ce9bf34fc2cff1e01f407f680e5b86173a7819deba3e79cf9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h ++87b7b522608a9ae1cecaa7b2913ddbe56b7ffb7d184742526e84f915f256dc1d vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_iterator.h ++9109a6e8716bfaaeae1025c27b52b16fe7d37858e991db44c1dd940b873cb9c2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h ++74c8d2e51aa13e8cfe95a7a38b1bdcd1cf3752aef2d6cc9e5e66d6372a747dc6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h ++fe8183095df6ef40d37dcb434168eb918b500f8499f825e94618166d7817b843 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h ++9f70dfed1d845445488be4b0cdad5e91f5ca4a93c781a158577437e75c0ed2e0 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h ++0e8498ad650b483806a60aad419cea6193f4302851b931eeea7d24e6586ab508 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_tensor_op_wmma.h ++881bbaca182c1fc1ed0f732828d34d415adb0b7cdabc50769249cd2aca234b17 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/mma_with_reduction_tensor_op.h ++6c533f9279901d64bff3295504b7ae9b5a04694abe2bc6cf4deffbe122774ba5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/scale_bias_tile_iterator.h ++e99a2a8319cea2454e03b852dc7bffdb7794704fcbc1d4126853a37f8e2fc53a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/softmax_scale_bias_transform.h ++aea61a730ac29d973ec05311fa2aea55154bdcb15fe8db56a7da00a1a2c38fa4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm/warp/tile_iterator_planar_complex.h ++1b9d463105ef5a9f3b66b7ee4c589e3e9535241ec44d146fe4a12c5209f29ac9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm_coord.h ++a05ee50778a8b7914ca83783f859ec5c09bffc20bf98400c99592f0516eaa45a vllm/third_party/fmha_sm100/cutlass/include/cutlass/gemm_coord.hpp ++1714b4a0f21027966ed1b00bcf774f5bd938a5c703a4413ddb47f424ef1f0585 vllm/third_party/fmha_sm100/cutlass/include/cutlass/half.h ++d18944bf77e5adc4c6039b08944613d5b8ec0aea5a10753a839cd0436cfb40da vllm/third_party/fmha_sm100/cutlass/include/cutlass/integer_subbyte.h ++ef86f98aa102be4671e1e5dc05087d9c3d8321ee41d1389a927553f09ae5d0c6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/kernel_hardware_info.h ++a9ba5f2737017e30d33a25628601b009ea334f27785df14b7164fac8999391d2 vllm/third_party/fmha_sm100/cutlass/include/cutlass/kernel_hardware_info.hpp ++d9cb7d3bea0446c1f6f7a441758ce3357403796ba5fa98f1f1b39c71b3a681e3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/kernel_launch.h ++259da643dfd0498946ea81718676afbd68d2348cac171da6b41bf896f94b0d08 vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/layout.h ++f7ff90051e0b2c7ea6f1295e62a2550e89aec5f8385a26c38ad38890a670341d vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/matrix.h ++bbba15b1896db9e072ee0aa8f86312e22637176b3f6fa5152400440481e4c30c vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/permute.h ++f3d01d0cafed30d817c41ec5c4acc8c89d405bce9854eb7f0cd68f8363ed146a vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/pitch_linear.h ++80a337895d3efc47d400ab573ca7a4763bd0f8c83ec949d47d88654f4b21361f vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/tensor.h ++9299f947d97c8a66c0a39fb2904bdd1513e54743e204e02fd144f770cc023f2e vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/tensor_op_multiplicand_sm70.h ++d4fef95a420e5724b775f854252352fa0447742d19e35fd26e7e70c72051db2f vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/tensor_op_multiplicand_sm75.h ++8ca98951049a3ecd2cf0027d7aaa07f6fddeb01e0c1d9356bdf66707c7231b97 vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/tensor_op_multiplicand_sm80.h ++4c08c6e7935d76ba92318c281abb803b3cf669a76ac87f15881fdc8ab526c90a vllm/third_party/fmha_sm100/cutlass/include/cutlass/layout/vector.h ++1f49e26e9f35de517ffcbc26bd67199f62037ecde1e7404e3a2f329ae2a6a848 vllm/third_party/fmha_sm100/cutlass/include/cutlass/matrix.h ++22b96274790f0f1dfd44aa476c1575c3ba44ac905353efc3325404530e39fe7e vllm/third_party/fmha_sm100/cutlass/include/cutlass/matrix_coord.h ++a5aa67691c6391a462696556df9a50d84746005d52f5980ee57870165edc0722 vllm/third_party/fmha_sm100/cutlass/include/cutlass/matrix_shape.h ++a0ca8f57c5ffb29de746ed5680559acc5f0c2257a22791c44a07c7d70b809be5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/numeric_conversion.h ++8b6948eb7e6e2675e95b57f3c32e90fbfcd53a1c05043305b34ccbec1642bf26 vllm/third_party/fmha_sm100/cutlass/include/cutlass/numeric_size.h ++d0caaa13b790351b7052aef6f0811fd11a06814b56389680c4d32955dd77aaeb vllm/third_party/fmha_sm100/cutlass/include/cutlass/numeric_types.h ++d0095570d60cf1ca77d8d8532b24afdea29d41d08c9be80c37b102996d83bf74 vllm/third_party/fmha_sm100/cutlass/include/cutlass/pipeline/pipeline.hpp ++c62adf8335bfaf94ac405781927d959c8a9972ebdd45eb110a2c15e23c676e60 vllm/third_party/fmha_sm100/cutlass/include/cutlass/pipeline/sm100_pipeline.hpp ++da8830e1780bad31ab2bbb460a4223cde7b6d6afa7de7b2aa19bffb3e4a23922 vllm/third_party/fmha_sm100/cutlass/include/cutlass/pipeline/sm90_pipeline.hpp ++5ae2681cc9098e616dafebe3386bdd5d56c7dce33754beb6f6485716d127c952 vllm/third_party/fmha_sm100/cutlass/include/cutlass/pitch_linear_coord.h ++e930770518ae9c70666162d6bc2af7b8b7d0cd0120980b1534e917eb5a20335e vllm/third_party/fmha_sm100/cutlass/include/cutlass/platform/platform.h ++8ba1b9b35e68878f1192996b13fe33fe636c0ee89160c92a647c50d531e1814e vllm/third_party/fmha_sm100/cutlass/include/cutlass/predicate_vector.h ++521209fca5d47f6a11fecbcbe03af5ad0e84375e63f44d2031614b3b4e1d4bd6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/quaternion.h ++6fd96474fd128d6ee44bfdc9274a9cabf805bd726ff7a1563b70fd3b41fbbf44 vllm/third_party/fmha_sm100/cutlass/include/cutlass/real.h ++ce389e886522f0df573df924b718246a1b251b701b32b4cfea53a945c3b2820b vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/device/reduce_split_k.h ++be4a0a1e39a65f18954ae232b2d890be12c8dd252bc5e0e59e37a34db5ae4565 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/device/tensor_reduce.h ++553f8810f1219f51a142e49fe5a5d0bfc8a9a553a041fe872b269b74907bbc6b vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/device/tensor_reduce_affine_contiguous.h ++b8bf2b3799d731100d986d5f96d37e44b47750ea1b4126a51b2e2315067d69c3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/device/tensor_reduce_affine_strided.h ++71ee31eebddd64d055a1ac341bea5f8cf9551fc58deebb90c43660ae99f37d66 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/kernel/reduce_softmax_final.h ++8eeb4aec22b215a8b1ba85ab0fb7e4dbb3620f222571d268a2ffc9d1335a8895 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/kernel/reduce_split_k.h ++ccb697a588dc1302179e9419d4df1a4e6e9c61c034935da2d476982e7f3f749a vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/kernel/tensor_reduce_affine_contiguous.h ++945ec58286b7d85a26877f082efe63574ce7b201194063280aef6f3ae966ae94 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/kernel/tensor_reduce_affine_strided.h ++7aab1efe95a5d3c504ce88933e1f3cf7d8515d8c4a451be0b58ba30108eb121b vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/thread/reduce.h ++d0eb73d4a9ff0c78f4c3044ce2711b090077175020da639b6d7018290918afc4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/thread/reduction_operators.h ++cf3dd7adf51c771fce3fbee0bc4d50f0ac77c6a119513af0f8e83481a7518706 vllm/third_party/fmha_sm100/cutlass/include/cutlass/reduction/threadblock_swizzle.h ++d19debf8ec0d40966fefca8615e1931729a071484ac90e6e5aca816376aca7b7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/relatively_equal.h ++854eeafb6d774891c40ec2a6a8dffae9277450e300486d2b182b949aade7fe06 vllm/third_party/fmha_sm100/cutlass/include/cutlass/semaphore.h ++e49ab0b69c5d1cff243946d4678dc3bf56e0835d2278235f49de6894feca276b vllm/third_party/fmha_sm100/cutlass/include/cutlass/subbyte_reference.h ++e97dbe66edce5c63d4207355163188fc50c10a86359f51e111478f51f7089214 vllm/third_party/fmha_sm100/cutlass/include/cutlass/tensor_coord.h ++734908cf9b867f28439ae59023fde6b1e36a71f7aacd8ee065146e8d3b23863e vllm/third_party/fmha_sm100/cutlass/include/cutlass/tensor_ref.h ++4185256d3b5a3ab9be5d234bef9f8a616273ba39cbf6ae49243371a0d270b26f vllm/third_party/fmha_sm100/cutlass/include/cutlass/tensor_ref_planar_complex.h ++026cc5dd6c3352da4785365f1184ad17af34c77523df04d447d552bc3ed6bc68 vllm/third_party/fmha_sm100/cutlass/include/cutlass/tensor_view.h ++0e2292afec05131952ca030920ae95ccec39716933a1ad047b349800822047d1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/tensor_view_planar_complex.h ++bcebe5e79eae0e73b60ede2b5c8fe9e50d6a3a4ada57719509214133abe40c52 vllm/third_party/fmha_sm100/cutlass/include/cutlass/tfloat32.h ++281fc2699998893a083737526532b38cc2ca33a5f782a6a30f9631aad50eaf9f vllm/third_party/fmha_sm100/cutlass/include/cutlass/thread/matrix.h ++b124b4c1cb13a7877188f142a39fb88f652ac3ff0bf7a1d8af9faed29b805f96 vllm/third_party/fmha_sm100/cutlass/include/cutlass/trace.h ++4c28dac2c25d8062dc3f4e8b71088da633a76f4a58480492629d1807e16f3b53 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/collective/sm90_wgmma_transpose.hpp ++3f63527830813f9c2593c9a3bd78a53030505a82e2bbe356725479dc6b081163 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/device/transform_universal_adapter.hpp ++d5b95357b7e46a7efddd7ea27033c1c709d836b825862aa1e18e3f3669b3c447 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/kernel/filter_format_transformer.hpp ++598ed04baf204cda5776f2c84db1529ce587f195047c7b3d7eb13ad11756ecd1 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp ++65c539018ef2ee58a33004ddbd0eabf8f445250516c087fc50a0a7e0b0f990b7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/kernel/sparse_gemm_compressor.hpp ++d28cee3534a897b4d2379868a2f4d733d79638802e4d6d63cb35ac27b733aa89 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/pitch_linear_thread_map.h ++81cca89f50a27ad09821c35cd804687f71d3b990158b541db42f399acd32aabf vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/thread/transpose.h ++7f9a91bdc04947599cecede830e1f33ed90d9ee8beb4c446cb0d62d66a1929f5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/thread/unary_op.h ++452ca74528dd650b7633c983ba327f73003cb97973eb5fd96aca712e9876a272 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/ell_iterator.h ++7af6f6a7603a730b63c5fc34ff1f9161f445b4947884bfdf4e33f19e7a99687b vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/ell_predicated_tile_access_iterator.h ++7c96296bf63349b0714ca06c6291bbe899e6a017cb1d5876dded34e8f3d37128 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/ell_predicated_tile_iterator.h ++9a4c9211605c22b1d40d7f5949c6ee9f154f3c9c15a15f6e143d2d2cd7c80e76 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h ++c77cd5a7571b1e821bf5b92b8112fe4b3b48d09dc7fa533fb6b5ea9cb226d4f7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_scale_bias_vector_iterator.h ++bea46e994f412ddabe7ff091d5f632a9fb7dd5edbd74902c4d9dd022f128f318 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_access_iterator.h ++41a6278aa465a774b0b81a445594a8dfd167f2447053d6d7040738dac7f39ab9 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h ++3efa8edd754f72fe97aed17188d5fe77c0df0431e7df162609f40801bf318359 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_access_iterator_params.h ++de5b2eb430cb98a311384bb43cc05afc8e3152294bd5110667aa764bb4fceab4 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h ++f3a7c5f6251bb0ef8b590117c2c1acb3c11d78eda8dfaad00b6f2804a2378006 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_iterator.h ++4a6d88d2beca239f42b2ada64249fa2f9ae51a3f03e4fe149732ba44f404a694 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h ++72c301ae50bc72bb45c0b5508cc2db21527be446eec5606e567839f060db6210 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h ++312504157bf6caa33db9b104499ae657994ce99b17f91eadf59394b4fede0d0e vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/predicated_vector_access_iterator.h ++23ee41cf5a4b448465f24373440eb1d6cc272945ccc4a92b8e33fd735f587398 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h ++46734ebbba2c2f82eec138c61dc90dd7cf84bd8a28564cb5e1b176cdd7b4a221 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_access_iterator.h ++1061192cb8ddfd6515462d2eba76e658ef83d5cd73e8c9515a8567819176c0be vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h ++dc0ab04cd73313c1405b4ec13c6e4050f5651ab9a03a8299eaf56b1c69470347 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h ++d3eb84b73dcddda1720684fe7051dfe98a48624c970261dea51a93810ca7bb24 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h ++ecad65270e7f57998726d42168272eb941b0a6d627205a620a2340bc330409f7 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h ++d37179eebba723577b2b1e6c56ebd2d4addf15c57baabb8baf5446690d072008 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_iterator.h ++1098ae3846e7812ec194a27cce48f699433dc7faadc62bc72c687519bc3d2e9e vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h ++00d925a5f5c4721ba40e946e7351849fe9fe10f0662d17ed391d94bf294f9d13 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h ++8da22500088ab7fe4416a9699caa0396307d7e73a31fb19b5cf8aa2f493a49a3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_iterator_tensor_op.h ++bec085104ab1934791dd94a401b04c5ab77f391d37a67a7fc152c0d21407b9f3 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h ++9306bdbee848ecc5dce74357d933ae34e041197b8e5d220009d684eb9ce7caad vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/threadblock/vector_iterator.h ++bfd6a8365d9d35199500686bacaebc0ce92afb7582aba943b010cc40ef2ba8e5 vllm/third_party/fmha_sm100/cutlass/include/cutlass/transform/warp/vector_fragment_iterator.h ++eb63921e2624d4ce61bbe532c044c78cd630245fdd4dc82ce773e54cf84d1864 vllm/third_party/fmha_sm100/cutlass/include/cutlass/uint128.h ++c3ab76ecc5fff8cc57bef0d7c4406521e1aae480e65a22f70c7c8138cf075578 vllm/third_party/fmha_sm100/cutlass/include/cutlass/uint256.h ++e6cc0b578a93e19a4a61694dcfd7e38f9f0f176ff235072bb44b6946dfaa2fe6 vllm/third_party/fmha_sm100/cutlass/include/cutlass/version.h ++747b26a6980c2e20b7f5f80d907ea40ea6dabeec8ec23be26fa9f5da61c6a63d vllm/third_party/fmha_sm100/cutlass/include/cutlass/wmma_array.h ++24b6c737521332445c0d6b08698234442c2a3a4fd7b0ed87c14488b2be13cb12 vllm/third_party/fmha_sm100/cutlass/include/cutlass/workspace.h ++9d26c7b0b54891432a3002b1794e3252074204f31e7ee92915871b708f812d4f vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/GPU_Clock.hpp ++1e7c3d5e737297bfbe43712d3126cb6474da28c6374805fa059512d37337f21e vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/command_line.h ++d97b42134bba99c0ff970ef9abf846d43172acb6e90e30e5003e16b708a20c42 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/cublas_wrappers.hpp ++738c1268d48a86d16fc39da2491a788c99c20ffd2df7fa51b1365dd322e09b03 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/debug.h ++393c5a96788cff201ff04b399795690b75d80bacbbf7f677216e236acec60cde vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_dump.h ++4711daf704f026268b6cf4dbdce6bee3f5a4f37a8f5c1ca54980e0dd8d1c4217 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_groupnorm.h ++b5e20fa61a97f1608f998904314e750dba4d808aed6d0c549367b453adf268bf vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_layernorm.h ++cd1da0696b2a75cb0c9875681515bdf7fc9cf525db69405a208f383e585be4e2 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_memory.h ++1c660c46ba622c2ec204e81fc8b25d8615ac10e11775fcdad04c60e709d584d0 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_nchw_to_nhwc.h ++5ed4539c6b222dfd56aa849906984d3397638e30dd870048ad3c8715d6a33212 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_nhwc_padding.h ++42697d7cd32571593fc86e40d1837bc19320471fdae93459fd8ec205f60c70f0 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_nhwc_pooling.h ++46383d970ad475541d4639cf4e5a91fde0a599dd5178b4e58d0a7d6a54e94860 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_nhwc_to_nchw.h ++e36813641bc54d1b9a876f1576ad9fe0de3aff26a051582573cf25c250bde4a7 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_rmsnorm.h ++2fb75f3f48d34827ddfc0dff3c96683126a2d1534f304138897150d6e8cbe835 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/device_utils.h ++a2aba46c40f79e39a860e985c7716f657bb02ace754b8d13b1592677b7e1b129 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/distribution.h ++78b2df81acd8713b909614a8e6e5424dbcf66fd4a893759779b332c5b8550a25 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/exceptions.h ++a5dba923744c383c74df8b50200b4cc464c77e748fa52c04d79abd907ec6257c vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/gett_commandline.hpp ++c0b030e7de9cb82422bce4e19b18ccd3f03de512782598473967b152dd2e3c01 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/helper_cuda.hpp ++2429ad181a0253b97debef0cd7122a5f316aa02d364f1e6c26dcd45a608c9358 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/host_reorder.h ++a71fb53dd1554bf97f8e8aff44c8565647c5537214c19b8c8bfac7e039d9c0d5 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/host_tensor.h ++e70eb5fca107897517caf0a63d99761912c41509fcf4ac3ea069fda9299917d7 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/host_tensor_planar_complex.h ++52f1d3128403fcccf964a0806469212d0acb39e1b9192bc88f3c74930e37ca6f vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/host_uncompress.h ++8efff8799aca2804465717f44c8589576cf6ebc302f27b559231523a33724a28 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/index_sequence.h ++2f28e4228c1383185fbe3a12d56df084c378cdd4c2b99a0a257870070fcd76a4 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/mixed_dtype_utils.hpp ++ca7f6b722a87848a53730cf8049991dcd781a5d79b37d7cda683ab6367710650 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/packed_stride.hpp ++7c0d41b17bc2685a3ac854f2f27d0ab1c4d1e534394a36aae11a4c1b683968c6 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/print_error.hpp ++472422dce82251da275c1b96e0bbfaf528e0fd76b6325e878bf762857d2ce7df vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/detail/inner_product.h ++0fae7148d079bc593d53d864ea93b6f91a3916a71f07d1f8479cbb67ffad32b4 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/detail/linear_to_coordinate.h ++16ec9ed7870891d611e9ab4cfca7d8cbec2ff1de081bffeb9ae8c3718368ef00 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/convolution.h ++c618a036fdd24edf56e69e51b1b9527a69244db4787e0db3c19bfa0808da7c8e vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/gemm.h ++b669308d2b98f67d4f666c03684cefdca0695f266fa0df61ca0b13c138fb3478 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/gemm_complex.h ++490631916c77c86cff1c8ba6617e68f9ab59d0315929202c6fb0253d5aae503b vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/gemm_planar_complex.h ++bc0376c361e0ac4d0579db92c6722956a2858a94713dbf8677ba07bc00a92c2f vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/gett.hpp ++b04bde1f1177dfddc2857c2f015736be17e5b1959c297b8b92d818a48a208e6a vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/kernel/gemm.h ++f0d7d809fee1ab9a356cf81f4491631e9b8e33063ca95e45019a25a459c46738 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/kernel/tensor_elementwise.h ++f5aadf07d19dd1085ec93775346f352bc1c572244fc291e7125f231fbf0938c8 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/kernel/tensor_foreach.h ++f993109af4ab23ca1fe3d720e256b33d00c75a832c6affc9d89a9fe5ed2a7efb vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/rank_2k_complex.h ++5e42caddedab231f67e8b207733579a6d772fa48d503cfd8cecf7176089bd736 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/tensor_compare.h ++2986439711a536b7f31bcda6e74510bf0b45b8090e72a41f7de06e4c5c5136c6 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/tensor_fill.h ++5a679c049dee011aacbbb135b08e9d844a089f7cd65a07ebc0c178cf05654485 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/tensor_foreach.h ++c1af84078b0888b69c328eca9de4e5eaeb638c7e51abb2263e611fa3e7cff734 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/tensor_reduce.h ++c86439b29341b5c44970f6d9ad9e8d0c02f002a6ea343246f4184c7000be0537 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/tensor_relu.h ++596745bcf2040a435c6389b645ab0cf66fc107ca5d7cbeadfb83fed2675202c0 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/device/thread/gemm.h ++cd0567922364fcd48e4b0b042496418c5a96ef71b6c9ae916cd4a38fe637f802 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/conv.hpp ++9b24fbf41d5dd41aa8da90c36df651f023a99fd8079996c50fc30792581d2eed vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/convolution.h ++644a52027ead70046d99311884bbcc804e1c372ac0b6fa0a544f769eb3b2dfa2 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/error_metrics.h ++7dcc71c2e7261314164a4e9606fa7b5eb61af4782cac558ee1c515b181fb8b69 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/gemm.h ++28d0a0c23ca3691f6196694018800e1c4f9831a530d41826fd5a2ecd64be7220 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/gemm_complex.h ++47ad4fd268d359a664f401c5e76f45602572e0484c9bb088b305bc4fb8777895 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/gemm_planar_complex.h ++bb64addd4e9186dad4ebc75061cb85a104653c5096ee0d5ea960aec2a3d1fc22 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/gett.hpp ++dc4ecdb59c1f69bf0103da8bae03742e804d39684113f439c03c7e15939b530b vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/rank_2k.h ++c85916223706908d5cccbbbd406da294e5636bcf281c980126498bba1d73cdea vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/rank_2k_complex.h ++e16b9d402089890ba110400ebd624d196861438b887797a1567d9504f2023db8 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/rank_k_complex.h ++e589249239f49751ff0df8d72ecf1f5c466953142d78688cceba9ab9b0de972e vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/symm.h ++73f897bdc0a7c53970b5e026ff1ce8cf9b6d483e34c4d032f94336c913377c0b vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/symm_complex.h ++8e721ddf56b35eaa6232d2d97a6a3d6408231dd73a72479e489b4bf2b0992f1b vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_compare.h ++221e3b9e1b7477a7544cd4e89f67ef569a482585817a1710040e2474b8791fbe vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_compare.hpp ++18c2e45376b789dfccd037954290fbb49a0f94ffab67171aabab16eddfb49bfa vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_copy.h ++16937b10e496309132b53ab035645bf51f4e4bbfb5f08a1641513ed29da7d574 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_elementwise.h ++392fd2d8fa6265de622b47dcbc871e3bb0651bd76cd411d6be3682ab8067730c vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_fill.h ++35ec6d63ac71d4e1388dd60877ecbc3ff933a3e2294ac2e9d580af6875d9d260 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_fill.hpp ++a50b2c8e469cbbad74373df0a9cad1fcabef618be9a2bd0b4aa693643848cb46 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_foreach.h ++2fde203d103d593c804ce0c9d04cb0fd3fdc43a884cd245c99bb6c44c6848de6 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_norm.h ++f859a41c9afec6c067145ba8c91d3d6f6a22e7c66d0244dabaf1ad55a4e8d9c9 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_reduce.h ++09c8fa7916f168fa0a8816c195ebf0cd51549cf46d590581c0d2ae8a730242c5 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/tensor_reduce.hpp ++b3eb3cc35c7dcb5500103cbd63afdb5eaa4a17a8241ab66b00185d41ee6d363c vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/trmm.h ++63c92fad58bf7e1ab569bd46c8677654ed93fbd502e90e5b367bbb82c4c710de vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/reference/host/trmm_complex.h ++9bfb6bcd8cf2b0a6177c589c7ec963f3f7f4b40018ca1f4a30f1a4654ca8f190 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/tensor_view_io.h ++456c182930d498142d2a6185d376d2f1d7798750341a5a055f4abe4a9b624415 vllm/third_party/fmha_sm100/cutlass/tools/util/include/cutlass/util/type_traits.h ++00bdea4e47a235a44c5ccc3dbfab7639161044bd1f7d978be20a6713c21db3d5 vllm/third_party/fmha_sm100/jit.py ++2a82bfd40b2fe842d5bacf0a8df5e3a77d8ca004107c71a1b68460a0b746dd4b vllm/third_party/fmha_sm100/sparse.py ++51a66a1ae50717aead927af2dec30621016d28e87646bb9109cb757fbfdf88db vllm/third_party/fmha_sm100/sparse_fmha_adapter.py ++e3b54fd4d89ec2a646087874356926475f5f5befdc64f222dd77e905aee78052 vllm/third_party/tml_fa4/__init__.py ++cadac170a7da379e27c36d044360480d710cfa88d082dad2b9decd58e4847111 vllm/third_party/tml_fa4/ampere_helpers.py ++767af807aa7521b8d15ee6c9534f4745a7ad2224b0da2b816cc49ccb9ea6c119 vllm/third_party/tml_fa4/barrier.py ++68ed6ae96a18aac82d334d2eadc401fe1c74c0dca5f176eb8398201f174ad2a2 vllm/third_party/tml_fa4/blackwell_helpers.py ++c801f1c2aea5b23531572cbb39fd0c3707f7dc827cb0fb3fa5598701cf753817 vllm/third_party/tml_fa4/block_info.py ++23202c5ab32976bdf2815af8a74b832fcee198b65fa30e64301f6b91a94d70ea vllm/third_party/tml_fa4/block_sparse_utils.py ++3905d3202fb9cf90f27d6f9645eb17605766310a3fb1d0bc1d91fe8a4fd61aba vllm/third_party/tml_fa4/block_sparsity.py ++e2c2eec7b3a529045fe3b830f65013bae242b677a0ef798c43092a45f1438569 vllm/third_party/tml_fa4/blockscaled_utils.py ++ed29d0efe4b0fe049d393407551244712400921ee946d61f0de3902110729b82 vllm/third_party/tml_fa4/cache_utils.py ++da2e1a50dbe03a89c78064948e7d1e25de2930d0e05052ca288e75a09e4b2178 vllm/third_party/tml_fa4/compute_block_sparsity.py ++8da38d559d684dc70bf0b1d8333765952aece5ace9eefa1036e9738e35e63d65 vllm/third_party/tml_fa4/copy_utils.py ++fac5571de88193cd9b9f05684a513de6f1a52e67b92509416bcd7f9619dd0f35 vllm/third_party/tml_fa4/cu_blocks_kernels.py ++21f0674c9f5a99d7c339092e49d5822d387b0a4673888becff1ac073cb5ac619 vllm/third_party/tml_fa4/cute_dsl_ptxas.py ++a9f6d79bda82744aebd992593d6ec6aa45dfea43f638c007bce6f3df2987f0ab vllm/third_party/tml_fa4/cute_dsl_utils.py ++bc71edfa4b1cb7af64b78def09e790ff4c6fe64ab4fcf8e3d9699f44e8d3c0a0 vllm/third_party/tml_fa4/fa_logging.py ++7e5f822bf4d0ba36967558f3ae9a1408b410cb726c31222cd61703569d1be738 vllm/third_party/tml_fa4/fast_math.py ++2327ba3cd0155ef8f286a4edb3772d18120a1b37d33de7be69651ab96dda268a vllm/third_party/tml_fa4/flash_fwd.py ++1114d94ada45d3df7ff950ff2380e7e5782aefb552708e1ba7a4e67c3280ae9d vllm/third_party/tml_fa4/flash_fwd_combine.py ++25d00d622ffd6a51f0ed8c781754f4df12eacef3432ccc881446b0b5d26bc5cf vllm/third_party/tml_fa4/flash_fwd_mla_sm100.py ++6f45c3b9ce6fa7f9255f8ba6d656d0ce2f411d5a4b516169916ad28aefca3e51 vllm/third_party/tml_fa4/flash_fwd_sm100.py ++3797f27c07d6458a78a7a22198c442fe9abe1c4562047e81f81d61df831f885a vllm/third_party/tml_fa4/flash_fwd_sm120.py ++1b640455fb76e8c7719e4137b66c809de641e6986a8cdc181f6216db48f9d0ba vllm/third_party/tml_fa4/flash_fwd_sm90.py ++5a3ec5fdcda160c6a453586194fd4c2249527d560c26dde7a7a099875957e1bb vllm/third_party/tml_fa4/interface.py ++d547d83cfb5dba511fc9b34068cf5379ba1e1e5567a2fb472772b628c02858c1 vllm/third_party/tml_fa4/kernel_info_utils.py ++7521049651eca207d8dcd2aefd6724e859a234f1c667bcd8116f7c3cc3ad5ac7 vllm/third_party/tml_fa4/mask.py ++4ed2b297f7fec34f1519481ef97a317b8f20fa742540d7d5a3e4a271e42445a2 vllm/third_party/tml_fa4/mixed_dtype_gemm.py ++ac6ec28b2f1a4b797c391a7807b88f83b430452d418262c3dfa7b0f3516e39f9 vllm/third_party/tml_fa4/mma_sm100_desc.py ++bf5607c2f28769f6756d62f061538ac4a79b6c23eebabc9ceb7fcf87720376d3 vllm/third_party/tml_fa4/named_barrier.py ++066cfd82e31a49a0616af51c8354c307945020ecea8e9514139e8c4eff1cc484 vllm/third_party/tml_fa4/pack_gqa.py ++af0d174d58c6030847e9829035ae1812957e54faca5f3d1a9f800f4c2599a62a vllm/third_party/tml_fa4/paged_kv.py ++849bce91c72be0a5c7458979e10c8290f45e05253a83600e30fb9540c61ed56a vllm/third_party/tml_fa4/pipeline.py ++1b7fd8f55c47a79e581672aa2cbc41763481d3f2ad23ec458dfb2ee3ea19b53f vllm/third_party/tml_fa4/prepare_scheduler.py ++c200347671d38aef34b26c9d47b2658471c395f7735780034d557ee4df908bf3 vllm/third_party/tml_fa4/seqlen_info.py ++df97fcdebae584a75c14a907637dc65819bdeafcde97ec8a207246f24482704d vllm/third_party/tml_fa4/shearing_bias.py ++b3f28b8239f3503b89bd9e3806a4ccb0b8332b283412f77b69aec55a4d3e6f5d vllm/third_party/tml_fa4/sm100_hd256_2cta_fmha_forward.py ++3db63695e9d36f8fd7ec302844937c6378539941945bb92a8a32a43d0c304099 vllm/third_party/tml_fa4/softmax.py ++1db68083be2e5a0eaf1e9545048cbb947c22ce2899a0f45767b10376a176e7b9 vllm/third_party/tml_fa4/testing.py ++470e5cc8bb6fb36721266eaa5bed61eb75721833f8442d29d761b821243d943d vllm/third_party/tml_fa4/tile_scheduler.py ++db7897a374279261287980a9612b52cd01a9bf079dafa7d6f88ad09bd993d5c7 vllm/third_party/tml_fa4/topk_gather_kv.py ++fbf21187a903aac33c9646922a85de80db4d3b4dd2dd9437acb8fc5b861c3848 vllm/third_party/tml_fa4/utils.py ++e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 vllm/third_party/triton_kernels/__init__.py ++34919681c109b47d3bad6dd2d31a169371cd353e88137feebc1c8c6a49f2c032 vllm/third_party/triton_kernels/compaction.py ++80575b48b4d9969169dcef2adc2888d9fe18a15162a0e9a53190802d6b4d7178 vllm/third_party/triton_kernels/compaction_details/_masked_compaction.py ++56061a61a2d018245b4a6ee3e41ff79d9f6cf089ffc072d1ca95fb5ded1da4a2 vllm/third_party/triton_kernels/matmul_ogs.py ++08787760becc4590c1dc3fa3f98308b78a66e730a19cee9b0d77c1e8a99b45f7 vllm/third_party/triton_kernels/matmul_ogs_details/_common.py ++c192e1d30f60f2c292744f37149f9363d76a825728610783d1eb533be284a921 vllm/third_party/triton_kernels/matmul_ogs_details/_matmul_ogs.py ++96159a4e7ac1cb220e99f5cbe60f991c68b9012b88b555678d607f1b32d2c839 vllm/third_party/triton_kernels/matmul_ogs_details/_p_matmul_ogs.py ++f67a8632e89012e745b1c7fed7f6ee031f0478752644669ccc29425be20bac49 vllm/third_party/triton_kernels/matmul_ogs_details/_reduce_grouped.py ++19e20f1cb9cf1f61bba0b2787f8ecda47183e17773fbadded02bd4ee48f9c719 vllm/third_party/triton_kernels/matmul_ogs_details/opt_flags.py ++9426288afcfd8ba29a3723d0a889f286e83312d212e210ee4bb49b94c521e02f vllm/third_party/triton_kernels/matmul_ogs_details/opt_flags_details/opt_flags_amd.py ++cea9971a3b39fdb90257563bf16beda190368f3b4126f71c48270bccda5635f3 vllm/third_party/triton_kernels/matmul_ogs_details/opt_flags_details/opt_flags_nvidia.py ++74e775d0084578bdc668872de76c6a72450406d2f19e86b2c7db7371ee20e860 vllm/third_party/triton_kernels/numerics.py ++e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 vllm/third_party/triton_kernels/numerics_details/__init__.py ++5596fd01d03991d94cec364342393b41523c629abe7f166d5e9d1229a4c56061 vllm/third_party/triton_kernels/numerics_details/flexpoint.py ++c115715b9add3247084c4576247e88c7f525bf13834d0d7ab41b3d0e9c6c57ce vllm/third_party/triton_kernels/numerics_details/mxfp.py ++6908f774b3c1d949196d7a1b6b4c6ad4b1e2c36f78877a7e35437a675fa29b54 vllm/third_party/triton_kernels/numerics_details/mxfp_details/_downcast_to_mxfp.py ++50298d8b908b1355cd0f82364abe748a540f20ae5f1d9c7df1618bc2a918c56c vllm/third_party/triton_kernels/numerics_details/mxfp_details/_upcast_from_mxfp.py ++0443785ec981dcf9267c727f6b3c07e8a3444888be191443edf6afc858b5f1c3 vllm/third_party/triton_kernels/proton_opts.py ++83c92ebb11a35556c7185e3aeb487802f4e280de158f291e5dddf95407d725fd vllm/third_party/triton_kernels/reduction_details/reduce_bitmatrix.py ++0aa55772883118d0eec936d8058f682f35074089e85fb9cfc87215ac97f11899 vllm/third_party/triton_kernels/routing.py ++13defd26d0da5bf6b2f3cbd945c748096a9293117de8c1d9f13b877836c1a80e vllm/third_party/triton_kernels/routing_details/_expt_data.py ++987c5b9e68ea589a40d07854cb49699943f2193a8ef867e9e7708c3a5066719b vllm/third_party/triton_kernels/routing_details/_routing_compute.py ++e8d94b8912af3ca58c6e4557c8c07bdf1e3a5b3206d55af99ea7fb6f79311e80 vllm/third_party/triton_kernels/specialize.py ++6f22081f81630cdb1eff3515e7f335d2ebd7110ec70c646320e58227758e4a8d vllm/third_party/triton_kernels/swiglu.py ++230fb476dccc603ac8b633003c9a356d297482c071e48ed22e857c1fd1b8f67a vllm/third_party/triton_kernels/swiglu_details/_swiglu.py ++4003467834f8a5c4c80f8f0bc77bc396025e77fbd9ae89083cb438e10a1a1a77 vllm/third_party/triton_kernels/target_info.py ++e08640c83d78a45a01bca8f874b9a69e759976c785ea38047a2b52e0a049a24e vllm/third_party/triton_kernels/tensor.py ++dbfe89fb74e1c377676a299cfcdd08aaead03a2aac160833b478cccf1222cafe vllm/third_party/triton_kernels/tensor_details/layout.py ++2223c16c379d5c47adb29eaf8b5e6a4fcd760a7b864f5591fa68847caffb04a2 vllm/third_party/triton_kernels/tensor_details/layout_details/base.py ++e60a5969b587dbd8e5fa675cf751cf3b5f93f7d43013a70318c0ac0f49a56611 vllm/third_party/triton_kernels/tensor_details/layout_details/blackwell_scale.py ++abf4653fc4b6641f6a79e2af9fbae2d50093a01dbea20c10c05f6b9de931cda8 vllm/third_party/triton_kernels/tensor_details/layout_details/blackwell_value.py ++0dc2bbf97fc70f9b919436411a1ae7557c5676bbbbcf63501b63962ec6088c8a vllm/third_party/triton_kernels/tensor_details/layout_details/cdna4_scale.py ++dd7ed682636b6bfa93e394b9087b208e9fe0a013b202adbeac43e3c034115f48 vllm/third_party/triton_kernels/tensor_details/layout_details/hopper_scale.py ++21e97c9800ff325525d5274344f7e82afcd72af537524b027806df6b31b74faf vllm/third_party/triton_kernels/tensor_details/layout_details/hopper_value.py ++86504afe2d2033068e862e21f41c0e8e53b05dca85544412565f6de6b41d4835 vllm/third_party/triton_kernels/tensor_details/layout_details/strided.py ++a2ba392036876caf15cb2071353ad3cc66e7a2a1048a39fa267d340d1c8e7592 vllm/third_party/triton_kernels/testing.py ++db3baedc20e07d41c38d0bbc89710e2bbf43373d29348388f06b91d75dd58371 vllm/third_party/triton_kernels/topk.py ++e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 vllm/third_party/triton_kernels/topk_details/__init__.py ++4820d74485e8eecfbc62cdbcca6c8b30f95da8e692e623ac3dd4983e6b5467b1 vllm/third_party/triton_kernels/topk_details/_topk_backward.py ++b96cd075eec221f009d9a32139605b9966f28edf6385ee166108f3768a6599f6 vllm/third_party/triton_kernels/topk_details/_topk_forward.py ++51a8760203b7e06d71f6094d57a40ec6a4a8092da57d19f852234013924fc8e2 vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so ++f155f95b83b5f4611a5118d9cceb58e6beb25434c49eebe613d940da7ace07da vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so ++50713012c3a5e8045f368672342aea6eeeb3b3ce0968cb5639b1f158b9499487 vllm/vllm_flash_attn/cute/__init__.py ++cadac170a7da379e27c36d044360480d710cfa88d082dad2b9decd58e4847111 vllm/vllm_flash_attn/cute/ampere_helpers.py ++767af807aa7521b8d15ee6c9534f4745a7ad2224b0da2b816cc49ccb9ea6c119 vllm/vllm_flash_attn/cute/barrier.py ++5bd0302c4cd4cd2c3b7e8b1d01eb49d546475d06897678f5543ea2e0172aae84 vllm/vllm_flash_attn/cute/bench_utils.py ++43bb07eb55d46a402abab8c50984eb39d4cf8f19a11ad79cd0291919614a9101 vllm/vllm_flash_attn/cute/benchmark.py ++ba5e05f40bd3cc788a6a6c32d5d75bf60f1ea04767783115377e434e386c241d vllm/vllm_flash_attn/cute/benchmark_flash_attention_fp8.py ++4228d771518f22d26c9785417dea5a92acab35da807f3151c2b79b81d8ebbb44 vllm/vllm_flash_attn/cute/blackwell_helpers.py ++f0ed9cc670dd7382e9ca65dcbe118b9abed4296ebadcc5d68a43ed33900c4968 vllm/vllm_flash_attn/cute/block_info.py ++9aeba86f906805f93b0c381985333458ae0f7287f8e965e60b50f0903dc2a70f vllm/vllm_flash_attn/cute/block_sparse_utils.py ++5d804b8c0730a84fa016d9a313cbb16fb66a023cb6bdcee7390e8c2e85d84059 vllm/vllm_flash_attn/cute/block_sparsity.py ++fc62f0889f79b770717e9030cb003749a14a7e01b3ada1443062ca11ac0edb9d vllm/vllm_flash_attn/cute/cache_utils.py ++aa508c474fcd03494d7dbbdef5fb2598935ec4447e0d3b5a8f2700c799390188 vllm/vllm_flash_attn/cute/compute_block_sparsity.py ++0f563eeeb6654602d1541f7202b7b88386c4dbdbd24ad6b0825834d65e508c09 vllm/vllm_flash_attn/cute/copy_utils.py ++7a8f4ae3cecf7ff517f05abea31cd95ef3048ba2be47b3175b15cd78af7555c0 vllm/vllm_flash_attn/cute/cu_blocks_kernel.py ++21f0674c9f5a99d7c339092e49d5822d387b0a4673888becff1ac073cb5ac619 vllm/vllm_flash_attn/cute/cute_dsl_ptxas.py ++ea5955c88e692e9704b0f47868eede7b405529e87dcd86e2d30474855a1abbea vllm/vllm_flash_attn/cute/cute_dsl_utils.py ++bc71edfa4b1cb7af64b78def09e790ff4c6fe64ab4fcf8e3d9699f44e8d3c0a0 vllm/vllm_flash_attn/cute/fa_logging.py ++7e5f822bf4d0ba36967558f3ae9a1408b410cb726c31222cd61703569d1be738 vllm/vllm_flash_attn/cute/fast_math.py ++91fdb0e4beb0a3f9611984380090eef1968e4c7c0b3781dd2af20dfd47760f3d vllm/vllm_flash_attn/cute/flash_bwd.py ++f42a8c414eb2f708135f687359059c724b137af9c3ff600ad23899663a113b5b vllm/vllm_flash_attn/cute/flash_bwd_mla_dk_sm100.py ++8f1e74128562799159715e616b562ddd117652cf480f3669e6014f29615f7a80 vllm/vllm_flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py ++16ccd55cb1d00eaf5e943ad3496ad7a5a6c39076dcc55b7d8c59dd5529a2d26e vllm/vllm_flash_attn/cute/flash_bwd_mla_sm100.py ++e22dea9c6c94c508a023a55d1d5b0a66c9474eae93ba620e357fbddc1ebff38c vllm/vllm_flash_attn/cute/flash_bwd_postprocess.py ++73dc44b33bf72e713476e973f786f40d57ecd71a0af77581429395c57ea21750 vllm/vllm_flash_attn/cute/flash_bwd_preprocess.py ++7560cf738c342206cb45eb10f2a76ab2976030d1b55a734cc9e2d462a1bb50e0 vllm/vllm_flash_attn/cute/flash_bwd_sm100.py ++9f5b0b67da797a30978452e86c191545c4f097a8e07123aef049b5c0b050f910 vllm/vllm_flash_attn/cute/flash_bwd_sm120.py ++b78e0e8f587c478431cfda7f84cc25f0e4027fdee20709a71aebabdc28588eb2 vllm/vllm_flash_attn/cute/flash_bwd_sm90.py ++2f953f9be60af867bc5ac3efb0823e440abe153f3309753325bbce9bf49a7ef2 vllm/vllm_flash_attn/cute/flash_fwd.py ++6cbd911fa7b0165cd1f643e1ee5645e74bcd3ce8bd849d579aa42b2de48e8b3d vllm/vllm_flash_attn/cute/flash_fwd_combine.py ++4ec9001144e46c348fd3d5ea475c0907283eb9e3933a33c15676cfb8084640f2 vllm/vllm_flash_attn/cute/flash_fwd_mla_sm100.py ++a080f5c1565a53bd90789ca88133ab055f0fa34e93fe866ed768f39c916667cb vllm/vllm_flash_attn/cute/flash_fwd_sm100.py ++61f3c4c018b8d9e6a3e12694a0a565ce660cd4ddbe68aed060ba71291abd8493 vllm/vllm_flash_attn/cute/flash_fwd_sm120.py ++2d81c74558f84f541206abdf4d30be586d0444360d2bc87f7fcacaa24e85d901 vllm/vllm_flash_attn/cute/flash_fwd_sm90.py ++5d93fe7c60b3865719cf83aec3f79b06f76299fc2634b4f3811d56e7b0d358ed vllm/vllm_flash_attn/cute/hopper_helpers.py ++827fc6a2d4ef8d1cc3ac991d33cab5cbe5a66ec7a82c9636ebf64b61ad780773 vllm/vllm_flash_attn/cute/interface.py ++ea35a2c9f29d82e8c939a7b161c39bd8bf40c43c702958fbfe074205fe5054e7 vllm/vllm_flash_attn/cute/mask.py ++ac6ec28b2f1a4b797c391a7807b88f83b430452d418262c3dfa7b0f3516e39f9 vllm/vllm_flash_attn/cute/mma_sm100_desc.py ++4be91f45f6d800f5059f9026b7f2d030cf755a99eba6718ccd069fe6d64f61ef vllm/vllm_flash_attn/cute/named_barrier.py ++24dfcd04219f2c40128c6ae718c519e514bb23ab17c116c614c477fedd38a181 vllm/vllm_flash_attn/cute/pack_gqa.py ++668bfae057343bba37e414459ad20ad5d4d97d5a78a64f59e7b4ba739f1066e1 vllm/vllm_flash_attn/cute/paged_kv.py ++849bce91c72be0a5c7458979e10c8290f45e05253a83600e30fb9540c61ed56a vllm/vllm_flash_attn/cute/pipeline.py ++95d8425a6ec9c4f1cd9d2390e06af58716a4a518d6baf1949c7810d369f68e42 vllm/vllm_flash_attn/cute/prepare_scheduler.py ++640d2433702635d68d4f0b94e9fa2d6b81da48b7f62056270b1b1fd480c72f30 vllm/vllm_flash_attn/cute/seqlen_info.py ++9ad38e0535bccb21b647bf0de532969445c69f9434b13b3c467b606c0c973600 vllm/vllm_flash_attn/cute/sm100_hd256_2cta_fmha_backward.py ++b2ab2ca01c7313f7d38f19b2ca12d2bf3e7a9e2d19704fd6dddfc8ed136bdfc9 vllm/vllm_flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py ++9cabc2ed2c15ad4f15a3464d0b7982e08bd25838e9906b0ab94ba76c61cce03b vllm/vllm_flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py ++83f83a0c2dfe89804f2843ac5b1d56d42e353cd802eea61382f9a59c8e3ea937 vllm/vllm_flash_attn/cute/sm100_hd256_2cta_fmha_forward.py ++bd0ed6094e9fc48de587098f9b05bccbdd5f450b70d3500759a48b5d436cff9c vllm/vllm_flash_attn/cute/sm90_config_search.py ++fbf2ae0275d10fd66855baacbf9804db8030b233aa02140fd1a728e7a3a79969 vllm/vllm_flash_attn/cute/softmax.py ++3342b4475126666e5bff050316cd2a586852f7c82eb23fd21172470cc40a15ff vllm/vllm_flash_attn/cute/testing.py ++27482cb3ef36aabcc07b115075557f6c18740790b4bd30986f2d52959be23abd vllm/vllm_flash_attn/cute/tile_scheduler.py ++85d31fc88cdf332054439fcc2e2438d2a387b52d421eb458789f767d3cf64cc6 vllm/vllm_flash_attn/cute/topk_gather_kv.py ++2115599e1eefed2f8504f7d1d48088bb7bf3286cf2ea2fa98eb093876267099f vllm/vllm_flash_attn/cute/utils.py ++b826cc3339488815b26f974135a8f53d44149d019b4c82e1ad4240a6ec0e28c4 vllm/vllm_flash_attn/layers/rotary.py ++f1f4bf9a0b28cc544fa882b60628d74b8923e8b47e1f8438a7af43fa7b5061e4 vllm/vllm_flash_attn/ops/triton/rotary.py +diff --git a/vllm/models/qwen4_exp/nvidia/abliteration.py b/vllm/models/qwen4_exp/nvidia/abliteration.py +index 840168b..6ce0528 100644 +--- a/vllm/models/qwen4_exp/nvidia/abliteration.py ++++ b/vllm/models/qwen4_exp/nvidia/abliteration.py +@@ -59,24 +59,22 @@ def load_direction(hidden_size: int) -> torch.Tensor | None: + + + def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None: +- """Attach the direction as a non-persistent buffer ``_abliteration_r`` (None when the switch is off). ++ """Attach the direction as a non-persistent buffer ``_abliteration_r`` (all zeros when the switch is off). + + Call inside the model constructor so the buffer lands on the model's device before CUDA-graph capture. The +- CPU-only PLE worker process never runs these forwards and gets no buffer. ++ buffer always exists, so a compiled graph compiled with the switch on resolves when it runs with the switch ++ off: a zero direction makes the projection an exact no-op. The CPU-only PLE worker process never runs these ++ forwards and gets the zero buffer. + """ + from vllm.model_executor.layers.ple_offload_layer import is_offload_process + + r = None if is_offload_process() else load_direction(hidden_size) +- if r is None: +- module._abliteration_r = None +- return +- buf = torch.empty(hidden_size, dtype=dtype) +- buf.copy_(r) ++ buf = torch.zeros(hidden_size, dtype=dtype) if r is None else r.to(dtype).flatten().contiguous() + module.register_buffer("_abliteration_r", buf, persistent=False) + + + def project_(x: torch.Tensor, r: torch.Tensor | None) -> torch.Tensor: +- """In place: remove the component along r from x's last dimension. No-op when r is None.""" ++ """In place: remove the component along r from x's last dimension. No-op when r is None or all zeros.""" + if r is None: + return x + width = x.shape[-1] +-- +2.43.0 + diff --git a/scripts/build-artifacts.list b/scripts/build-artifacts.list new file mode 100644 index 0000000..d6ed83a --- /dev/null +++ b/scripts/build-artifacts.list @@ -0,0 +1,35 @@ +# Build products the serving tree needs. ./build.sh verifies every path in this list +# after the compile, and hashes them into vllm/.build-products.sha256. +# Lines starting with '#' are skipped. +vllm/_C_stable_libtorch.abi3.so +vllm/_moe_C_stable_libtorch.abi3.so +vllm/_version.py +vllm/cumem_allocator.abi3.so +vllm/fs_io_C.abi3.so +vllm/spinloop.abi3.so +vllm/third_party/flashmla/flash_mla_interface.py +vllm/third_party/fmha_sm100/ +vllm/third_party/tml_fa4/ +vllm/third_party/triton_kernels/ +vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so +vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so +vllm/vllm_flash_attn/cute/ +vllm/vllm_flash_attn/layers/ +vllm/vllm_flash_attn/ops/ + +# Built only when the matching toolchain is present. build.sh warns when cargo is missing, +# and prints the exact command to add it. vllm-rs is the optional Rust frontend +# (envs.py: VLLM_RUST_FRONTEND_PATH); the tool parsers in the shipped serve profile are Python. +#vllm/vllm-rs +#vllm/_rust_tool_parser.abi3.so + +# Built only for the architectures that have kernels. cmake/external_projects/*.cmake gate these +# per arch: flashmla 9.0a/10.0/10.7, deepselect 10.0, flashkda 9.0a/10.0/12.0, qutlass 10.0/12.0, +# deep_gemm (optional vllm._deep_gemm_C, the DeepSeek-V4 path). arch 8.6 (RTX 3090) builds none of +# them, and the Qwen3.8-Flash-Next serve path does not import them. +#vllm/_deepselect_C.abi3.so +#vllm/_flashkda_C.abi3.so +#vllm/_flashmla_C.abi3.so +#vllm/_flashmla_extension_C.abi3.so +#vllm/_qutlass_C.abi3.so +#vllm/third_party/deep_gemm/ diff --git a/scripts/gen-cuda-stub.sh b/scripts/gen-cuda-stub.sh new file mode 100755 index 0000000..af329d3 --- /dev/null +++ b/scripts/gen-cuda-stub.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash + +# Generates the CUDA driver link stub: $1 = a CUDA C header, $2 = the .so to write. +# +# vllm links its C++ kernels against CUDA::cuda_driver, the CUDA driver library +# (libcuda.so.1). The driver package provides it; the CUDA toolkit wheel does +# not, and neither does a build container without a GPU. NVIDIA ships the +# link-time stub for it in the runfile installer, not in the pip wheel, so this +# script builds an equivalent one from the header: one exported symbol per +# driver entry point, no bodies, 87 KB. +# +# Link time only. The soname is libcuda.so.1, so the dynamic linker opens the +# real driver from the host at run time and can never pick the stub. +# +# Writes nothing when the file already has every entry point, so a re-run does +# not invalidate the compiled-ops hash recorded by build.sh. + +set -euo pipefail + +H="${1:?usage: $0 }" +OUT="${2:?usage: $0 }" + +[ -r "$H" ] || { echo "no readable CUDA header at $H"; exit 1; } + +decls(){ + # Calls written as "CUresult CUDAAPI cuFoo(...)", split over lines. + awk '/CUresult[ \t]+CUDAAPI/{cap=1;buf=""} cap{buf=buf" "$0; if(buf~/;/){ + if(match(buf,/CUresult[ \t]+CUDAAPI[ \t]+[A-Za-z_][A-Za-z0-9_]*/)){ + s=substr(buf,RSTART,RLENGTH); sub(/.*CUDAAPI[ \t]+/,"",s); print s} cap=0}}' "$H" + # Versioned names hidden behind a macro: code writes cuFoo, the preprocessor + # rewrites it to cuFoo_v2, so the object file asks the linker for cuFoo_v2. + sed -nE 's/^#define[ \t]+(cu[A-Za-z0-9_]+)[ \t]+(cu[A-Za-z0-9_]+_v[0-9]+)([ \t]|$).*/\2/p' "$H" +} + +want="$(decls | sort -u)" +n="$(printf '%s\n' "$want" | sed '/^$/d' | wc -l)" + +[ "$n" -gt 100 ] || { echo "only $n driver entry points found in $H"; exit 1; } + +if [ -e "$OUT" ]; then + have="$(nm -D --defined-only "$OUT" 2>/dev/null | awk '$2=="T"{print $3}' | sort -u)" + miss="$(LC_ALL=C comm -23 <(printf '%s\n' "$want") <(printf '%s\n' "$have") | sed '/^$/d')" + + if [ -z "$miss" ]; then + echo "driver stub is current: $(basename "$OUT"), $n entry points" + exit 0 + fi + + echo "driver stub is missing $(printf '%s\n' "$miss" | wc -l) entry points, regenerating" +fi + +mkdir -p "$(dirname "$OUT")" +src="$(mktemp --suffix=.c)" + +cleanup(){ rm -f "$src"; } + +trap cleanup EXIT +: > "$src" +printf '%s\n' "$want" | while read -r f; do + [ -n "$f" ] && printf 'unsigned int %s(void){ return 0; }\n' "$f" >> "$src" +done + +cc -shared -fPIC -O0 -o "$OUT" "$src" -Wl,-soname,libcuda.so.1 + +have="$(nm -D --defined-only "$OUT" | awk '$2=="T"{print $3}' | sort -u)" +miss="$(LC_ALL=C comm -23 <(printf '%s\n' "$want") <(printf '%s\n' "$have") | sed '/^$/d')" + +[ -z "$miss" ] || { echo "the stub is still missing:"; printf '%s\n' "$miss"; exit 1; } + +echo "driver stub generated: $OUT ($(printf '%s\n' "$want" | wc -l) entry points, soname libcuda.so.1)" diff --git a/scripts/install-as-service.sh b/scripts/install-as-service.sh new file mode 100755 index 0000000..fced26e --- /dev/null +++ b/scripts/install-as-service.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash + +# One entry point for the bare-metal systemd service: install, start, stop, status, logs, +# update, uninstall. Run from anywhere; the script uses the repo it lives in. +# +# scripts/install-as-service.sh +# +# install write /etc/systemd/system/vllm.service for THIS repo directory, +# write /etc/default/vllm (checkpoint path, model name) when missing, +# install the logrotate policy, daemon-reload, enable +# start|stop|restart|status|logs systemctl for the service +# update ./build.sh in this repo, then restart the service +# uninstall disable + remove the unit, the EnvironmentFile and the logrotate file +# +# Needs root (the log goes to /var/log, the unit goes to /etc/systemd/system). +# The service runs the repo where it lives: no copy, no /opt, no second build tree. +# Env: CKPT MODEL_NAME PORT override what the EnvironmentFile sets. + +set -euo pipefail + +HERE="$(cd "$(dirname "$0")/.." && pwd)" +UNIT=vllm +UNIT_PATH=/etc/systemd/system/$UNIT.service +DEFAULT=/etc/default/$UNIT +LOGROTATE=/etc/logrotate.d/$UNIT +LOG="${LOG:-/var/log/vllm.log}" + +[ "$(id -u)" = 0 ] || { echo "run as root: sudo $0 $*"; exit 1; } +command -v systemctl >/dev/null || { echo "no systemd on this host"; exit 1; } + +cmd="${1:-}"; shift || true +case "$cmd" in + install) + ckpt="${CKPT:-$(sed -n 's/^MODEL_DIR=//p' "$HERE/.env" 2>/dev/null | head -1)}" + ckpt="${ckpt:-/data/models/halt95/Swift1.5-Qwen3.8-Flash-Next-W4A16-Merlin}" + [ -f "$HERE/start.sh" ] || { echo "not the repo root: $HERE"; exit 1; } + # 1. the unit: the shipped template, with /opt/vllm replaced by this directory + sed "s|/opt/vllm|$HERE|g" "$HERE/systemd/$UNIT.service" > "$UNIT_PATH" + echo "==> unit: $UNIT_PATH (WorkingDirectory=$HERE)" + # 2. the EnvironmentFile: create once, never overwrite an edited one + if [ ! -f "$DEFAULT" ]; then + { echo "# Read by $UNIT.service. Uncomment and edit, then: systemctl restart $UNIT" + echo "CKPT=$ckpt" + echo "MODEL_NAME=Scarlet" + echo "# LOG=/var/log/vllm.log" + echo "# CACHE_ROOT=$HERE/.vllm-cache" + echo "# CUDA_VISIBLE_DEVICES=0,1,2,3" + echo "# VLLM_LOG_STATS_INTERVAL=1" + } > "$DEFAULT" + echo "==> wrote $DEFAULT: set CKPT to the checkpoint directory" + else + echo "==> kept $DEFAULT (it exists; CKPT=$(. "$DEFAULT" 2>/dev/null && echo "${CKPT:-unset}" || echo unset))" + fi + # 3. log rotation for $LOG (same policy as the shipped systemd/vllm.logrotate) + sed "s|/var/log/vllm.log|$LOG|g" "$HERE/systemd/$UNIT.logrotate" > "$LOGROTATE" + chmod 644 "$LOGROTATE" + echo "==> logrotate: $LOGROTATE" + systemctl daemon-reload + systemctl enable "$UNIT" + ck="$( . "$DEFAULT" 2>/dev/null && echo "${CKPT:-}" )" + [ -n "$ck" ] && [ -d "$ck" ] || echo "==> next: edit CKPT in $DEFAULT (the checkpoint directory), then sudo $0 start" + echo "==> installed and enabled. Start with: sudo $0 start" + ;; + start) systemctl start "$UNIT"; echo "==> started; a cold start takes ~6 min: sudo $0 logs" ;; + stop) systemctl stop "$UNIT" ;; + restart) systemctl restart "$UNIT" ;; + status) + systemctl status "$UNIT" --no-pager -l | head -12 || : + echo "==> engine log: tail -f $LOG (healthy line: GPU KV cache size)" + ;; + logs) exec tail -n 100 -f "$LOG" ;; + update) + echo "==> build, then restart" + ( cd "$HERE" && ./build.sh ) + systemctl restart "$UNIT" + echo "==> update done" + ;; + uninstall) + printf 'remove the unit, the EnvironmentFile and the logrotate file? [y/N] ' + read -r a + case "$a" in + [yY]*) systemctl disable --now "$UNIT" 2>/dev/null || : + rm -f "$UNIT_PATH" "$DEFAULT" "$LOGROTATE" + systemctl daemon-reload + echo "==> removed (the repo, .venv, .vllm-cache and $LOG are kept)" ;; + *) echo "cancelled" ;; + esac + ;; + *) + sed -n '3,18p' "$0" + exit 1 + ;; +esac diff --git a/scripts/make-e1-config.py b/scripts/make-e1-config.py new file mode 100755 index 0000000..eb788e2 --- /dev/null +++ b/scripts/make-e1-config.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 + +"""Add the FP8 PLE config key the serving tree needs to a Qwen3.8-Flash-Next-W4A16-Merlin checkpoint. + +The serving tree keys the FP8 PLE (n-gram embedding) table on a checkpoint config field instead of the +VLLM_PLE_FP8_GLOBAL_SCALE=1 environment variable the earlier lane used. The field is read from the model's +*text* config (`vllm_config.model_config.hf_text_config`), so it has to live inside the `"text_config"` +object of config.json, not at the top level. Checkpoints published without that key serve at the env-var +scale, which is the wrong scale for this tree. This script inserts +the key as a targeted text edit (the rest of the file is byte-preserved), re-parses the result to prove the +key landed where the engine reads it, writes a new file and renames it into place -- so a hard-linked +config.json is un-linked rather than mutated in place -- and keeps a backup. + +Weights are untouched: every shard can stay a hard link of the original checkpoint. + + python3 scripts/make-e1-config.py /path/to/Qwen3.8-Flash-Next-W4A16-Merlin +""" + +import json, os, re, shutil, sys + +KEY, VAL = "ple_embedding_dtype", "float8_e4m3fn" +d = sys.argv[1] if len(sys.argv) > 1 else sys.exit(__doc__) +p = os.path.join(d, "config.json") + +if not os.path.isfile(p) or not os.path.isfile(os.path.join(d, "model.safetensors.index.json")): + sys.exit(f"{d}: not a checkpoint directory (config.json + model.safetensors.index.json expected)") + +raw = open(p, "rb").read() +cfg = json.loads(raw) +text = cfg.get("text_config") + +if not isinstance(text, dict): + sys.exit("config.json has no text_config object; this is not the Flash-Next checkpoint this script is for") + +if KEY in text and text[KEY] != VAL: + sys.exit(f"text_config.{KEY}={text[KEY]!r} already present; refusing to overwrite") + +stale_top = KEY in cfg # left by the pre-fix version of this script; the engine never read it + +if text.get(KEY) == VAL and not stale_top: + print("already set"); sys.exit(0) + +if os.stat(p).st_nlink > 1: + print(f"note: config.json has {os.stat(p).st_nlink} hard links; writing a new file so the other links stay unchanged") + +s = raw.decode("utf-8") + +if stale_top: + # remove the top-level copy (the pre-fix script wrote it right after the opening brace): the first + # occurrence of the key line that sits before "text_config" is the top-level one + tc = s.index('"text_config"') + mm = None + + for cand in re.finditer(r'\r?\n[ \t]*"' + KEY + r'"\s*:\s*"[^"]*",?', s): + if cand.start() < tc: mm = cand; break + + if mm is None: sys.exit(f"top-level {KEY} present but not in the expected position; remove it by hand") + + s = s[:mm.start()] + s[mm.end():] + + if KEY in json.loads(s): sys.exit("internal error: top-level key still present after removal; nothing written") + + print(f"note: stale top-level {KEY} (from the pre-fix script) removed") + +m = re.search(r'"text_config"\s*:\s*\{', s) + +if not m or s.count('"text_config"') != 1: + sys.exit("could not locate a unique \"text_config\": { in config.json") + +i = m.end() +nl = "\r\n" if "\r\n" in s[:400] else "\n" + +# indentation of the first key inside text_config (the text after the opening brace up to the first quote) +after = s[i:i + 64] +indent = after[: len(after) - len(after.lstrip())].replace("\r", "").replace("\n", "") or " " +new = s if text.get(KEY) == VAL else s[:i] + nl + indent + f'"{KEY}": "{VAL}",' + s[i:] +check = json.loads(new) # must still parse ... + +if check["text_config"].get(KEY) != VAL: # ... and the key must be where the engine reads it + sys.exit("internal error: key did not land in text_config; nothing written") + +other = {k: v for k, v in check.items() if k != "text_config"} + +if other != {k: v for k, v in cfg.items() if k not in ("text_config", KEY)} or \ + {k: v for k, v in check["text_config"].items() if k != KEY} != {k: v for k, v in text.items() if k != KEY}: + sys.exit("internal error: something other than the one key changed; nothing written") + +shutil.copy2(p, p + ".bak-pre-v2") + +tmp = p + ".tmp-v2" + +with open(tmp, "wb") as f: f.write(new.encode("utf-8")) + +os.replace(tmp, p) + +print(f"config.json: text_config.{KEY} = {VAL} in place, no other key (backup config.json.bak-pre-v2)") diff --git a/scripts/probe-untyped-storage.py b/scripts/probe-untyped-storage.py new file mode 100755 index 0000000..9c2ae7c --- /dev/null +++ b/scripts/probe-untyped-storage.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 + +"""Reproduce the UntypedStorage failure from boot #1. + +Mirrors the engine: CUDA initialized, then safetensors_weights_iterator over the +heretic2 shards. Logs the exact tensor name that raises. +""" + +import glob +import sys +import torch + +device = sys.argv[1] if len(sys.argv) > 1 else "cuda:0" + +torch.cuda.set_device(int(device.split(":")[1])) +torch.zeros(1, device=device) + +from safetensors import safe_open +from vllm.model_executor.model_loader.weight_utils import ( + safetensors_weights_iterator, +) + + +files = sorted(glob.glob("/models/model/*.safetensors")) +orig = safe_open.get_tensor + + +def patched(self, name): + try: + return orig(self, name) + except Exception as exc: + print(f"RAISE {self.__class__.__name__} name={name} err={exc}", flush=True) + + raise + + +safe_open.get_tensor = patched +n = 0 + +for name, t in safetensors_weights_iterator(files, False): + n += 1 + + if n % 40000 == 0: + print(n, flush=True) + +print("OK tensors:", n) diff --git a/scripts/requirements-pinned.txt b/scripts/requirements-pinned.txt new file mode 100644 index 0000000..2c02fb6 --- /dev/null +++ b/scripts/requirements-pinned.txt @@ -0,0 +1,196 @@ +agent-detector==2.0.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aiosignal==1.4.0 +annotated-doc==0.0.5 +annotated-types==0.8.0 +anthropic==1.7.0 +anyio==4.15.1 +apache-tvm-ffi==0.1.11 +astor==0.8.1 +attrs==26.1.0 +blake3==1.0.9 +cachetools==7.2.0 +cbor2==6.1.4 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.1 +click==8.5.0 +cloudpickle==3.1.2 +compressed-tensors==0.17.0 +cryptography==50.0.1 +cuda-bindings==13.4.2 +cuda-core==1.2.0 +cuda-pathfinder==1.8.2 +cuda-python==13.4.1 +cuda-tile==1.6.0 +cuda-toolkit==13.0.3.0 +depyf==0.20.0 +detect-installer==0.2.1 +dill==0.4.1 +dnspython==2.8.0 +docstring_parser==0.18.0 +einops==0.8.2 +email-validator==2.3.0 +fastapi==0.136.3 +fastapi-cli==0.0.32 +fastapi-cloud-cli==0.26.0 +fastar==0.12.0 +fastsafetensors==0.4.0 +filelock==4.0.1 +flashinfer-python==0.6.18.post1 +frozenlist==1.8.0 +fsspec==2026.9.0 +googleapis-common-protos==1.75.3 +grpcio==1.84.0 +h11==0.16.0 +hf-xet==1.6.0 +httpcore==1.0.9 +httpcore2==2.13.0 +httptools==0.8.0 +httpx==0.28.1 +httpx2==2.13.0 +huggingface_hub==1.32.0 +humming-kernels==0.1.12 +idna==3.20 +ijson==3.5.1 +instanttensor==0.2.0 +interegular==0.3.3 +Jinja2==3.1.6 +jiter==0.17.0 +jmespath==1.1.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +lark==1.2.2 +llguidance==1.7.6 +llvmlite==0.47.0 +lm-format-enforcer==0.11.3 +loguru==0.7.3 +markdown-it-py==4.2.0 +MarkupSafe==3.0.3 +mcp==2.2.0 +mcp-types==2.2.0 +mdurl==0.1.2 +mistral_common==1.11.7 +ml_dtypes==0.6.0 +model-hosting-container-standards==0.1.16 +mpmath==1.3.0 +msgspec==0.21.1 +multidict==6.9.1 +nccl4py==0.5.0 +networkx==3.7 +ninja==1.13.2 +numba==0.65.0 +numpy==2.4.6 +nvidia-cublas==13.1.1.3 +nvidia-cuda-cccl==13.0.85 +nvidia-cuda-crt==13.0.88 +nvidia-cuda-cupti==13.0.85 +nvidia-cuda-nvcc==13.0.88 +nvidia-cuda-nvdisasm==13.4.92 +nvidia-cuda-nvrtc==13.0.88 +nvidia-cuda-runtime==13.0.96 +nvidia-cudnn-cu13==9.20.0.48 +nvidia-cudnn-frontend==1.29.0 +nvidia-cufft==12.0.0.61 +nvidia-cufile==1.15.1.6 +nvidia-curand==10.4.0.35 +nvidia-cusolver==12.0.4.66 +nvidia-cusparse==12.6.3.3 +nvidia-cusparselt-cu13==0.8.1 +nvidia-cutlass-dsl==4.7.1 +nvidia-cutlass-dsl-libs-base==4.7.1 +nvidia-cutlass-dsl-libs-core==4.7.1 +nvidia-cutlass-dsl-libs-cu12==4.7.1 +nvidia-cutlass-dsl-libs-cu13==4.7.1 +nvidia-ml-py==13.610.43 +nvidia-nccl-cu13==2.29.7 +nvidia-nvjitlink==13.4.92 +nvidia-nvshmem-cu13==3.4.5 +nvidia-nvtx==13.0.85 +nvidia-nvvm==13.0.88 +nvtx==0.2.15 +openai==3.17.0 +openai-harmony==0.0.8 +opencv-python-headless==5.0.0.93 +opentelemetry-api==1.44.0 +opentelemetry-exporter-otlp==1.44.0 +opentelemetry-exporter-otlp-proto-common==1.44.0 +opentelemetry-exporter-otlp-proto-grpc==1.44.0 +opentelemetry-exporter-otlp-proto-http==1.44.0 +opentelemetry-proto==1.44.0 +opentelemetry-sdk==1.44.0 +opentelemetry-semantic-conventions==0.65b0 +opentelemetry-semantic-conventions-ai==0.5.1 +outlines_core==0.2.14 +packaging==26.3 +partial-json-parser==0.2.1.1.post7 +pillow==12.3.0 +prometheus-fastapi-instrumentator==8.1.0 +prometheus_client==0.26.0 +propcache==0.5.4 +protobuf==7.36.2 +psutil==7.2.2 +py-cpuinfo==9.0.0 +pybase64==1.5.0 +pycountry==26.2.16 +pycparser==3.0 +pydantic==2.13.5 +pydantic-extra-types==2.11.1 +pydantic-settings==2.15.0 +pydantic_core==2.46.5 +Pygments==2.21.0 +PyJWT==2.14.0 +pynvvideocodec==2.0.4 +python-dotenv==1.2.3 +python-json-logger==4.2.0 +python-multipart==0.0.32 +PyYAML==6.0.3 +pyzmq==27.2.0 +quack-kernels==0.6.5 +referencing==0.37.0 +regex==2026.9.10 +requests==2.34.2 +rich==15.0.0 +rich-toolkit==0.20.5 +rignore==0.8.1 +rpds-py==2026.6.3 +safetensors==0.8.0 +sentencepiece==0.2.2 +sentry-sdk==2.69.2 +setproctitle==1.3.7 +setuptools==80.10.2 +shellingham==1.5.4 +six==1.17.0 +sniffio==1.3.1 +sse-starlette==3.4.11 +starlette==1.6.0 +supervisor==4.3.0 +sympy==1.14.0 +tabulate==0.10.0 +tiktoken==0.14.0 +tilelang==0.1.12 +tokenizers==0.23.2 +tokenspeed-mla==0.1.8 +tokenspeed-triton==3.8.10.post20260920 +torch==2.13.0 +torch_c_dlpack_ext==0.1.5 +torchaudio==2.11.0 +torchcodec==0.16.0 +torchvision==0.28.0 +tqdm==4.70.1 +transformers==5.18.0 +triton==3.7.1 +truststore==0.10.4 +typer==0.27.2 +typing-inspection==0.4.4 +typing_extensions==4.16.0 +urllib3==2.8.0 +uvicorn==0.53.0 +uvloop==0.22.1 +watchfiles==1.3.0 +websockets==17.1 +wheel==0.48.0 +xgrammar==0.2.7 +yarl==1.25.1 +z3-solver==4.15.4.0 diff --git a/scripts/stdbg-sitecustomize.py b/scripts/stdbg-sitecustomize.py new file mode 100755 index 0000000..8fc2953 --- /dev/null +++ b/scripts/stdbg-sitecustomize.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 + +"""Diagnose heretic2 boot Failure B. + +When PLE_ST_DEBUG_GETTENSOR=1, logs every safetensors tensor read to +/tmp/stdbg-.log so the last line before an UntypedStorage ValueError +names the tensor that fails. Installed via PYTHONPATH as sitecustomize. +""" + +import os +import sys + +if os.environ.get("PLE_ST_DEBUG_GETTENSOR") == "1": + try: + from safetensors import safe_open + + _orig = safe_open.get_tensor + _path = f"/stdbg/stdbg-{os.getpid()}.log" + _fh = open(_path, "a", buffering=1 << 16) + + def _patched(self, name): + try: + t = _orig(self, name) + except Exception as exc: + _fh.write(f"RAISE {name} {type(exc).__name__} {exc}\n") + _fh.flush() + + raise + + _fh.write(f"OK {name} {tuple(t.shape)} {t.dtype}\n") + + return t + + safe_open.get_tensor = _patched + except Exception as exc: # pragma: no cover + sys.stderr.write(f"STDBG patch failed {exc}\n") diff --git a/scripts/topo_select.py b/scripts/topo_select.py new file mode 100755 index 0000000..5c3f764 --- /dev/null +++ b/scripts/topo_select.py @@ -0,0 +1,1310 @@ +#!/usr/bin/env python3 + +# topo_select.py -- chooses the GPU order, the P2P/NCCL settings and (opt-in) a NUMA binding for the +# 4-GPU lane (tensor parallel 2 x pipeline parallel 2) before the serve script launches vLLM. +# +# python3 scripts/topo_select.py --dry-run # show what it would decide on this machine +# +# The serve script calls it only when AUTO_TOPO=1 or AUTO_TOPO=strict (opt-in; unset or 0 keeps the fixed +# settings and never runs this file). Run by hand, it treats an unset AUTO_TOPO as 1. +# CUDA_VISIBLE_DEVICES=GPU-,... picks the lane on a machine with more than 4 GPUs. +# It never creates a CUDA context: it reads NVML, sysfs and /proc only. +"""Topology selector for the 4-GPU serving lane. + +Purpose +------- +The serve script used to hard-code settings for one kind of host (4 GPUs on one NUMA node, PCIe +peer-to-peer working, no NVLink). This module inspects the lane's GPUs and writes a small shell env +file that the serve script sources. It decides: + +* the rank order of the lane's four GPUs (TP pairs (a,b),(c,d); PP edges (a,c),(b,d)), reordering only + when another order is strictly better on a simple link-locality score; +* whether vLLM should verify P2P itself (``VLLM_SKIP_P2P_CHECK``) and the NCCL variables; +* optionally (``AUTO_TOPO_BIND=1``) a ``numactl`` launch prefix binding the server to the lane's NUMA node. + +On the reference layout (4 GPUs, one NUMA node, P2P everywhere, no NVLink, bus order) the result is +byte-identical to the serve script's fixed settings. Any doubt falls back to those fixed settings. + +Usage +----- +``python3 topo_select.py [--env-out PATH] [--json-dir DIR] [--dry-run] [--quiet]`` + +Caller intent is read only from the snapshot the serve script takes before applying its own defaults: +``TOPO_CALLER_SET_=1`` + ``TOPO_CALLER_=`` for V in CUDA_VISIBLE_DEVICES, CUDA_DEVICE_ORDER, +VLLM_SKIP_P2P_CHECK, NCCL_P2P_LEVEL, NCCL_PROTO, PLE_HOME. When run by hand, export those to mimic a caller. + +A caller ``CUDA_VISIBLE_DEVICES`` that is set but empty (or only whitespace) counts as unset, as it did in the +serve script's own ``${CUDA_VISIBLE_DEVICES:-0,1,2,3}`` default. + +Opt-in: the serve script runs this module only when ``AUTO_TOPO`` is ``1`` or ``strict``; with +``AUTO_TOPO`` unset or ``0`` serve never invokes it and keeps its fixed settings. The module's own default is +different on purpose: run by hand (or by anything else), an unset or empty ``AUTO_TOPO`` means ``1``, so +``--dry-run`` shows a decision without extra flags. ``AUTO_TOPO=0`` given to the module itself writes the fixed +settings. + +Controls: ``AUTO_TOPO`` (1 | strict | 0; unset = 1 here, see above), ``AUTO_TOPO_NCCL`` (legacy | default), +``AUTO_TOPO_BIND`` (0 | 1), ``LANE_GPUS`` (auto), ``LANE_CPUS`` (cpulist; set but empty means no CPUs, so binding is skipped), +``TOPO_TIMEOUT`` (seconds, default 20, at most ``threading.TIMEOUT_MAX``). +Diagnostics go to ``--json-dir``, else ``$VLLM_CACHE_ROOT/topology``, else ``$CACHE_ROOT/topology``, else +``~/.cache/flash-next/topology``. The serve script passes ``--json-dir`` explicitly: the ``topology/`` directory under +the same cache root it exports as ``VLLM_CACHE_ROOT``. + +Test-only inputs: ``TOPO_FIXTURE`` (a JSON file replacing every hardware read) and ``TOPO_SYSFS_ROOT`` (a directory +prefixed to the ``/sys`` and ``/proc`` paths the real-hardware adapter reads, so the tests can run that adapter +against a fake sysfs tree and a fake ``pynvml`` module). Neither is meant for production use. + +Exit codes: 0 decision written, 2 usage error, 3 strict-mode refusal, 4 timeout, 5 internal error. +Requires Python >= 3.10; ``pynvml`` (package ``nvidia-ml-py``) is imported only for real-host discovery. +""" + +from __future__ import annotations +import argparse +import dataclasses +import glob +import hashlib +import itertools +import json +import math +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +import threading +import time +from typing import Any + +VERSION = 1 +LEGACY_CVD = "0,1,2,3" +SNAPSHOT_VARS = ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", "VLLM_SKIP_P2P_CHECK", + "NCCL_P2P_LEVEL", "NCCL_PROTO", "PLE_HOME") +ENV_VARS = ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", "VLLM_SKIP_P2P_CHECK", "NCCL_P2P_LEVEL", "NCCL_PROTO") +LEGACY_ENV = {"CUDA_VISIBLE_DEVICES": LEGACY_CVD, "CUDA_DEVICE_ORDER": "PCI_BUS_ID", + "VLLM_SKIP_P2P_CHECK": "1", "NCCL_P2P_LEVEL": "SYS", "NCCL_PROTO": "LL"} +ANCESTORS = ("INTERNAL", "SINGLE", "MULTIPLE", "HOSTBRIDGE", "NODE", "SYSTEM") +IDLE_MEM_MIB = 1024 +CVD_UNSET_EQUIVALENT = re.compile(r"\s*") # set-but-empty or whitespace-only CUDA_VISIBLE_DEVICES == unset +BIND_MIN_CPUS = 8 +BIND_MIN_FREE_KIB = 1.2 * 69 * 1024 * 1024 # 1.2 x the lane's ~69 GiB resident host memory +EXIT_OK, EXIT_USAGE, EXIT_STRICT, EXIT_TIMEOUT, EXIT_INTERNAL = 0, 2, 3, 4, 5 + + +class UsageError(Exception): + pass + + +class StrictRefusal(Exception): + pass + + +class NvmlUnavailable(Exception): + pass + + +# ---------------------------------------------------------------------------------------------------------------- +# small helpers + +def normalise_bus_id(raw: Any) -> str: + """'00000000:01:00.0' / '0000:01:00.0' / bytes -> '00000000:01:00.0' (lowercase, 8-digit domain).""" + if isinstance(raw, bytes): + raw = raw.decode("ascii", "replace") + + s = str(raw).strip().rstrip("\x00").lower() + m = re.fullmatch(r"(?:([0-9a-f]+):)?([0-9a-f]{1,2}):([0-9a-f]{1,2})\.([0-7])", s) + + if not m: + raise ValueError(f"unrecognised PCI bus id {raw!r}") + + dom = int(m.group(1) or "0", 16) + + return f"{dom:08x}:{int(m.group(2), 16):02x}:{int(m.group(3), 16):02x}.{m.group(4)}" + + +def sysfs_bus_id(bus_id: str) -> str: + """Normalised id -> the sysfs spelling '0000:01:00.0'.""" + dom, rest = bus_id.split(":", 1) + + return f"{int(dom, 16):04x}:{rest}" + + +def parse_cpulist(text: str) -> set[int]: + """'0-3,8,10-11' -> {0,1,2,3,8,10,11}. Raises ValueError on malformed input.""" + out: set[int] = set() + text = text.strip() + + if not text: + return out + + for part in text.split(","): + part = part.strip() + m = re.fullmatch(r"(\d+)(?:-(\d+))?(?::(\d+))?", part) + + if not m: + raise ValueError(f"bad cpulist element {part!r}") + + lo = int(m.group(1)) + hi = int(m.group(2)) if m.group(2) else lo + step = int(m.group(3)) if m.group(3) else 1 + + if hi < lo or step < 1: + raise ValueError(f"bad cpulist range {part!r}") + + out.update(range(lo, hi + 1, step)) + + return out + + +def format_cpulist(cpus: set[int]) -> str: + """{0,1,2,3,8} -> '0-3,8'.""" + items = sorted(cpus) + parts = [] + i = 0 + + while i < len(items): + j = i + + while j + 1 < len(items) and items[j + 1] == items[j] + 1: + j += 1 + + parts.append(str(items[i]) if i == j else f"{items[i]}-{items[j]}") + + i = j + 1 + + return ",".join(parts) + + +def host_path(path: str) -> str: + """'/sys/...' or '/proc/...' under TOPO_SYSFS_ROOT when that is set (tests only), else unchanged. On Windows, + where a file name cannot contain ':', the fake tree spells the ':' of PCI addresses as '_'.""" + root = os.environ.get("TOPO_SYSFS_ROOT") + + if not root: + return path + + rel = path.lstrip("/") + + return os.path.join(root, rel.replace(":", "_") if os.name == "nt" else rel) + + +def short_uuid(uuid: str) -> str: + return uuid[:8] + ".." if len(uuid) > 10 else uuid + + +def fmt_class(c: float | None) -> str: + if c is None: + return "?" + return str(int(c)) if float(c).is_integer() else str(c) + + +def atomic_write(path: str, text: str) -> None: + d = os.path.dirname(os.path.abspath(path)) + fd, tmp = tempfile.mkstemp(prefix=".topo_select.", dir=d) + + try: + with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as f: + f.write(text) + + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + +# ---------------------------------------------------------------------------------------------------------------- +# facts + +@dataclasses.dataclass +class Gpu: + uuid: str + bus_id: str # normalised, the sort key + key: Any # fixture index or NVML handle index + ordinal: int = -1 # CUDA ordinal under PCI_BUS_ID order + numa: int | None = None + pcie_gen_max: int | None = None + pcie_gen_cur: int | None = None + pcie_width_cur: int | None = None + pcie_width_max: int | None = None + details_read: bool = False + + +@dataclasses.dataclass +class PairFacts: + ancestor: str | None # one of ANCESTORS, or None = unknown + p2p_read: str | None # "OK" or another status; None = unknown + p2p_write: str | None + nvlink: bool + + @property + def p2p_ok(self) -> bool | None: + vals = (self.p2p_read, self.p2p_write) + + if any(v is None for v in vals): + return None # one direction unknown: the pair is unknown, never "not supported" + + return all(v == "OK" for v in vals) + + +@dataclasses.dataclass +class CpuFacts: + affinity: set[int] | None + mems_allowed: set[int] | None + nodes: dict[int, tuple[set[int] | None, int | None]] # node -> (cpus, mem_free_kib) + + +class Topology: + """Common interface over fixture and NVML facts. Everything beyond identity is read lazily, per GPU/pair.""" + + source = "?" + + def __init__(self) -> None: + self.gpus: list[Gpu] = [] # sorted by bus id + self.warnings: list[str] = [] + self._pairs: dict[tuple[str, str], PairFacts] = {} + + def _finish_identity(self) -> None: + self.gpus.sort(key=lambda g: g.bus_id) + + for k, g in enumerate(self.gpus): + g.ordinal = k + + def details(self, g: Gpu) -> Gpu: + if not g.details_read: + self._read_details(g) + + g.details_read = True + + return g + + def pair(self, a: Gpu, b: Gpu) -> PairFacts: + k = tuple(sorted((a.bus_id, b.bus_id))) + + if k not in self._pairs: + self._pairs[k] = self._read_pair(a, b) + + return self._pairs[k] + + # implemented by subclasses + def _read_details(self, g: Gpu) -> None: ... + def _read_pair(self, a: Gpu, b: Gpu) -> PairFacts: ... + def occupancy(self, g: Gpu) -> tuple[int | None, int | None]: ... + def system_nodes(self) -> list[int] | None: ... + def cpu(self) -> CpuFacts: ... + def numactl_present(self) -> bool: ... + def numactl_dry_run(self, cpus: str, node: int) -> bool: ... + + def resolve_numa(self, raw: int | None) -> int | None: + """sysfs numa_node -> node; -1 on a single-node system means that node, otherwise unknown.""" + if raw is None: + return None + + if raw >= 0: + return raw + + nodes = self.system_nodes() + + if raw == -1 and nodes is not None and len(nodes) == 1: + return nodes[0] + + return None + + +def _opt_int(v: Any) -> int | None: + if v is None or isinstance(v, bool): + return None + try: + return int(v) + except (TypeError, ValueError): + return None + + +class FixtureTopology(Topology): + source = "fixture" + + def __init__(self, data: dict) -> None: + super().__init__() + + self.data = data + self.fixture_name = data.get("name") + + for i, g in enumerate(data.get("gpus") or []): + self.gpus.append(Gpu(uuid=str(g["uuid"]), bus_id=normalise_bus_id(g["bus_id"]), key=i)) + + self._finish_identity() + self._cpu_raw = data.get("cpu") or {} + self._numactl = data.get("numactl") or {} + + def _raw(self, g: Gpu) -> dict: + return self.data["gpus"][g.key] + + def _read_details(self, g: Gpu) -> None: + r = self._raw(g) + g.numa = self.resolve_numa(_opt_int(r.get("numa"))) + g.pcie_gen_max = _opt_int(r.get("pcie_gen_max")) + g.pcie_gen_cur = _opt_int(r.get("pcie_gen_cur")) + g.pcie_width_cur = _opt_int(r.get("pcie_width_cur")) + g.pcie_width_max = _opt_int(r.get("pcie_width_max")) + + def _read_pair(self, a: Gpu, b: Gpu) -> PairFacts: + pairs = self.data.get("pairs") or {} + p = pairs.get(f"{a.key}-{b.key}") or pairs.get(f"{b.key}-{a.key}") + + if p is None: + return PairFacts(None, None, None, False) + + anc = p.get("ancestor") + anc = str(anc).upper() if anc is not None else None + rd, wr = p.get("p2p_read"), p.get("p2p_write") + + return PairFacts(anc if anc in ANCESTORS else None, + None if rd is None else str(rd), None if wr is None else str(wr), + bool(p.get("nvlink", False))) + + def occupancy(self, g: Gpu) -> tuple[int | None, int | None]: + r = self._raw(g) + + return _opt_int(r.get("procs")), _opt_int(r.get("mem_used_mib")) + + def system_nodes(self) -> list[int] | None: + nodes = self._cpu_raw.get("nodes") + + if not isinstance(nodes, dict): + return None + + return sorted(int(k) for k in nodes) + + def cpu(self) -> CpuFacts: + c = self._cpu_raw + + def cl(v: Any) -> set[int] | None: + return None if v is None else parse_cpulist(str(v)) + + nodes = {int(k): (cl(v.get("cpulist")), _opt_int(v.get("mem_free_kib"))) + for k, v in (c.get("nodes") or {}).items()} + + return CpuFacts(cl(c.get("affinity")), cl(c.get("mems_allowed")), nodes) + + def numactl_present(self) -> bool: + return bool(self._numactl.get("present", False)) + + def numactl_dry_run(self, cpus: str, node: int) -> bool: + return bool(self._numactl.get("dry_run_ok", False)) + + +class NvmlTopology(Topology): + """Real-host facts from NVML (no CUDA context), sysfs and /proc.""" + + source = "nvml" + + def __init__(self) -> None: + super().__init__() + + try: + import pynvml # from nvidia-ml-py; imported lazily on purpose + + pynvml.nvmlInit() + except Exception as e: # ImportError, NVMLError, OSError (library missing) + raise NvmlUnavailable(f"{type(e).__name__}: {e}") from e + + self.nv = pynvml + self._handles: dict[int, Any] = {} + self._nvlink_peers: dict[str, set[str] | None] = {} + self._nodes_cache: list[int] | None | bool = False + + try: + for i in range(pynvml.nvmlDeviceGetCount()): + h = pynvml.nvmlDeviceGetHandleByIndex(i) + uuid = pynvml.nvmlDeviceGetUUID(h) + uuid = uuid.decode() if isinstance(uuid, bytes) else str(uuid) + bus = normalise_bus_id(pynvml.nvmlDeviceGetPciInfo(h).busId) + self._handles[i] = h + + self.gpus.append(Gpu(uuid=uuid, bus_id=bus, key=i)) + except Exception as e: + self.close() + + raise NvmlUnavailable(f"device enumeration failed: {type(e).__name__}: {e}") from e + self._finish_identity() + + def close(self) -> None: + try: + self.nv.nvmlShutdown() + except Exception: + pass + + def const(self, name: str, default: int) -> int: + """Some nvidia-ml-py versions define NVML_* constants as 1-tuples; unwrap before passing to NVML.""" + c = getattr(self.nv, name, default) + + return c[0] if isinstance(c, tuple) else c + + def _try(self, fn, *args): + try: + return fn(*args) + except Exception: + return None + + def _read_details(self, g: Gpu) -> None: + h = self._handles[g.key] + nv = self.nv + g.pcie_gen_max = self._try(nv.nvmlDeviceGetMaxPcieLinkGeneration, h) + g.pcie_gen_cur = self._try(nv.nvmlDeviceGetCurrPcieLinkGeneration, h) + g.pcie_width_max = self._try(nv.nvmlDeviceGetMaxPcieLinkWidth, h) + g.pcie_width_cur = self._try(nv.nvmlDeviceGetCurrPcieLinkWidth, h) + raw = None + + try: + with open(host_path(f"/sys/bus/pci/devices/{sysfs_bus_id(g.bus_id)}/numa_node"), + encoding="ascii") as f: + raw = int(f.read().strip()) + except (OSError, ValueError): + raw = None + + g.numa = self.resolve_numa(raw) + + def _ancestor(self, ha, hb) -> str | None: + try: + level = self.nv.nvmlDeviceGetTopologyCommonAncestor(ha, hb) + except Exception: + return None + + level = level[0] if isinstance(level, tuple) else level + names = {self.const(f"NVML_TOPOLOGY_{n}", v): n + for n, v in (("INTERNAL", 0), ("SINGLE", 10), ("MULTIPLE", 20), + ("HOSTBRIDGE", 30), ("NODE", 40), ("SYSTEM", 50))} + + return names.get(level) + + def _p2p(self, ha, hb, index_name: str, index_default: int) -> str | None: + """'OK', another status name, or None when the query itself failed (unknown, not unsupported).""" + try: + st = self.nv.nvmlDeviceGetP2PStatus(ha, hb, self.const(index_name, index_default)) + except Exception: + return None + + st = st[0] if isinstance(st, tuple) else st + + if st == self.const("NVML_P2P_STATUS_OK", 0): + return "OK" + + names = {1: "CNS", 2: "GNS", 3: "TNS", 4: "DISABLED_BY_REGKEY", 5: "NS", 6: "UNKNOWN"} + + return names.get(st, f"STATUS_{st}") + + def _peers(self, g: Gpu) -> set[str] | None: + """Bus ids at the far end of this GPU's active NVLinks (empty on cards without NVLink).""" + if g.bus_id in self._nvlink_peers: + return self._nvlink_peers[g.bus_id] + + nv, h = self.nv, self._handles[g.key] + peers: set[str] = set() + enabled = self.const("NVML_FEATURE_ENABLED", 1) + remote_fn = (getattr(nv, "nvmlDeviceGetNvLinkRemotePciInfo_v2", None) + or getattr(nv, "nvmlDeviceGetNvLinkRemotePciInfo", None)) + + for link in range(self.const("NVML_NVLINK_MAX_LINKS", 18)): + try: + state = nv.nvmlDeviceGetNvLinkState(h, link) + except Exception as e: + code = getattr(e, "value", None) + + if code not in (self.const("NVML_ERROR_NOT_SUPPORTED", 3), + self.const("NVML_ERROR_INVALID_ARGUMENT", 2)): + self.warnings.append(f"NVLink state query failed on {g.uuid} link {link}: {e}") + + break + + state = state[0] if isinstance(state, tuple) else state + + if state != enabled or remote_fn is None: + continue + try: + peers.add(normalise_bus_id(remote_fn(h, link).busId)) + except Exception: + continue + + self._nvlink_peers[g.bus_id] = peers + + return peers + + def _read_pair(self, a: Gpu, b: Gpu) -> PairFacts: + ha, hb = self._handles[a.key], self._handles[b.key] + nvlink = b.bus_id in (self._peers(a) or ()) or a.bus_id in (self._peers(b) or ()) + + return PairFacts(self._ancestor(ha, hb), + self._p2p(ha, hb, "NVML_P2P_CAPS_INDEX_READ", 0), + self._p2p(ha, hb, "NVML_P2P_CAPS_INDEX_WRITE", 1), + nvlink) + + def occupancy(self, g: Gpu) -> tuple[int | None, int | None]: + h = self._handles[g.key] + procs = self._try(self.nv.nvmlDeviceGetComputeRunningProcesses, h) + mem = self._try(self.nv.nvmlDeviceGetMemoryInfo, h) + + return (None if procs is None else len(procs), + None if mem is None else int(mem.used) // (1024 * 1024)) + + def system_nodes(self) -> list[int] | None: + if self._nodes_cache is False: + paths = glob.glob(os.path.join(host_path("/sys/devices/system/node"), "node[0-9]*")) + self._nodes_cache = sorted(int(os.path.basename(p)[4:]) for p in paths) or None + + return self._nodes_cache # type: ignore[return-value] + + def cpu(self) -> CpuFacts: + try: + affinity: set[int] | None = set(os.sched_getaffinity(0)) + except (AttributeError, OSError): + affinity = None + + mems = None + + try: + with open(host_path("/proc/self/status"), encoding="ascii", errors="replace") as f: + for line in f: + if line.startswith("Mems_allowed_list:"): + mems = parse_cpulist(line.split(":", 1)[1]) + except (OSError, ValueError): + mems = None + + nodes: dict[int, tuple[set[int] | None, int | None]] = {} + + for n in self.system_nodes() or []: + base = host_path(f"/sys/devices/system/node/node{n}") + + try: + with open(f"{base}/cpulist", encoding="ascii") as f: + cpus: set[int] | None = parse_cpulist(f.read()) + except (OSError, ValueError): + cpus = None + + free = None + + try: + with open(f"{base}/meminfo", encoding="ascii") as f: + for line in f: + m = re.match(r"Node\s+\d+\s+MemFree:\s+(\d+)\s*kB", line) + + if m: + free = int(m.group(1)) + except OSError: + free = None + + nodes[n] = (cpus, free) + + return CpuFacts(affinity, mems, nodes) + + def numactl_present(self) -> bool: + return shutil.which("numactl") is not None + + def numactl_dry_run(self, cpus: str, node: int) -> bool: + exe = shutil.which("numactl") + + if exe is None: + return False + try: + r = subprocess.run([exe, f"--physcpubind={cpus}", f"--preferred={node}", "true"], + stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + timeout=10, check=False) + + return r.returncode == 0 + except (OSError, subprocess.SubprocessError): + return False + + +def read_facts_fixture(path: str) -> Topology: + """Load a fixture. faults.nvml 'error' behaves like a failed NVML init; 'timeout' hangs (the caller times out).""" + with open(path, encoding="utf-8") as f: + data = json.load(f) + + fault = (data.get("faults") or {}).get("nvml") + + if fault == "error": + raise NvmlUnavailable("fixture fault: nvml error") + if fault == "timeout": + while True: + time.sleep(3600) + + return FixtureTopology(data) + + +def read_facts_nvml() -> Topology: + return NvmlTopology() + + +# ---------------------------------------------------------------------------------------------------------------- +# decisions + +@dataclasses.dataclass +class Controls: + auto_topo: str # "1" | "strict" | "0" + nccl: str # "legacy" | "default" + bind: bool + lane_auto: bool + lane_cpus: str | None + caller: dict[str, str | None] # snapshot: value, or None when the caller did not set it + + +@dataclasses.dataclass +class Decision: + status: str = "legacy" + reason: str = "" + facts_source: str = "?" + visible: int = 0 + lane: list[Gpu] = dataclasses.field(default_factory=list) # final order + input_order: list[Gpu] = dataclasses.field(default_factory=list) + reordered: bool = False + cvd_numeric_ok: bool = False # caller gave numeric tokens or nothing (numeric form may be kept) + auto_lane: bool = False + classes: dict[tuple[int, int], float | None] = dataclasses.field(default_factory=dict) # over final order + p2p: dict[tuple[int, int], PairFacts] = dataclasses.field(default_factory=dict) + candidates: list[dict] = dataclasses.field(default_factory=list) + chosen_score: tuple | None = None + input_score: tuple | None = None + env: dict[str, tuple[str | None, str]] = dataclasses.field(default_factory=dict) # value, provenance + prefix: list[str] = dataclasses.field(default_factory=list) + binding: dict = dataclasses.field(default_factory=lambda: {"enabled": False, "command": None, + "skipped_reason": "not_requested"}) + warnings: list[str] = dataclasses.field(default_factory=list) + + +def resolve_lane(topo: Topology, ctl: Controls, dec: Decision) -> list[Gpu] | None: + """Map the caller's CUDA_VISIBLE_DEVICES (or the visible set) to 4 distinct GPUs, or set a legacy reason.""" + gpus = topo.gpus + cvd = caller_cvd(ctl) + + if cvd is not None: + lane: list[Gpu] = [] + numeric = True + + for tok in (t.strip() for t in cvd.split(",")): + if re.fullmatch(r"\d+", tok): + k = int(tok) + + if k >= len(gpus): + dec.reason = "cvd_unresolved" + + return None + lane.append(gpus[k]) + elif tok.startswith("GPU-"): + numeric = False + t = tok.lower() + exact = [g for g in gpus if g.uuid.lower() == t] + match = exact or [g for g in gpus if g.uuid.lower().startswith(t)] + + if len(match) != 1: + dec.reason = "cvd_unresolved" + + return None + + lane.append(match[0]) + else: + dec.reason = "cvd_unresolved" + + return None + if len({g.bus_id for g in lane}) != len(lane): + dec.reason = "cvd_duplicate" + + return None + if len(lane) != 4: + dec.reason = "lane_not_4" + + return None + dec.cvd_numeric_ok = numeric + + return lane + + dec.cvd_numeric_ok = True + + if len(gpus) == 4: + return list(gpus) + if len(gpus) < 4: + dec.reason = "fewer_than_4_visible" + + return None + + # more than 4 visible, no lane given. Strict mode refuses before anything else, LANE_GPUS=auto included: an idle + # observation is a guess about which GPUs are free, and strict means "do not guess". + if ctl.auto_topo == "strict": + raise StrictRefusal(f"topo: {len(gpus)} GPUs visible and CUDA_VISIBLE_DEVICES is not set; AUTO_TOPO=strict " + "refuses (LANE_GPUS=auto is not used in strict mode). " + "Set CUDA_VISIBLE_DEVICES=GPU-,... (4 GPUs).") + + if ctl.lane_auto: + return auto_lane(topo, ctl, dec) + + dec.reason = "more_than_4_visible" + + dec.warnings.append(f"{len(gpus)} GPUs visible and CUDA_VISIBLE_DEVICES is not set: using GPUs 0,1,2,3 in " + "PCI bus order. Pick the lane explicitly with CUDA_VISIBLE_DEVICES=GPU-,... " + "(4 GPUs; list them with nvidia-smi -L) or set LANE_GPUS=auto.") + + return None + + +def auto_lane(topo: Topology, ctl: Controls, dec: Decision) -> list[Gpu] | None: + """LANE_GPUS=auto: best 4-subset of idle GPUs (all on one NUMA node first, then score, then smallest bus ids).""" + dec.auto_lane = True + eligible = [] + + for g in topo.gpus: + procs, mem = topo.occupancy(g) + + if procs == 0 and mem is not None and mem <= IDLE_MEM_MIB: + eligible.append(g) + + if len(eligible) < 4: + dec.reason = "auto_insufficient_idle" + + return None + + best_key, best = None, None + + for subset in itertools.combinations(eligible, 4): + cls = lane_classes(topo, list(subset)) + + if cls is None: + continue + + top = max(order_score(p, cls) for p in itertools.permutations(range(4))) + numas = {topo.details(g).numa for g in subset} + one_node = len(numas) == 1 and None not in numas + key = (one_node, top, tuple(-int(c, 16) for g in subset for c in re.split(r"[:.]", g.bus_id))) + + if best_key is None or key > best_key: + best_key, best = key, list(subset) + + if best is None: + dec.reason = "auto_facts_unknown" + + return None + + dec.warnings.append("LANE_GPUS=auto picked " + ",".join(g.uuid for g in best) + + " (idle when checked; an idle observation is not a reservation)") + + return best + + +def link_class(p: PairFacts) -> float | None: + """NVLink 5; same switch 4 (P2P) / 2.5; host bridge or same node 3 (P2P) / 2; cross-socket 1; unknown None.""" + if p.nvlink: + return 5 + + if p.ancestor is None or p.p2p_ok is None: + return None + + if p.ancestor == "SYSTEM": + return 1 + + if p.ancestor in ("INTERNAL", "SINGLE", "MULTIPLE"): + return 4 if p.p2p_ok else 2.5 + + return 3 if p.p2p_ok else 2 # HOSTBRIDGE, NODE + + +def lane_classes(topo: Topology, lane: list[Gpu]) -> dict[tuple[int, int], float] | None: + """Link class for all six pairs (indices into lane); None when any class is unknown.""" + out = {} + + for i, j in itertools.combinations(range(4), 2): + c = link_class(topo.pair(lane[i], lane[j])) + + if c is None: + return None + + out[(i, j)] = c + + return out + + +def order_score(order: tuple[int, ...], cls: dict[tuple[int, int], float]) -> tuple: + def c(x: int, y: int) -> float: + return cls[(min(x, y), max(x, y))] + + a, b, cc, d = order + tp = (c(a, b), c(cc, d)) + + return (min(tp), sum(tp), min(c(a, cc), c(b, d))) + + +def choose_order(topo: Topology, lane: list[Gpu], dec: Decision) -> list[Gpu] | None: + """Keep the input order unless a permutation scores strictly higher; None when facts are unknown.""" + cls = lane_classes(topo, lane) + + if cls is None: + return None + + identity = (0, 1, 2, 3) + dec.input_score = order_score(identity, cls) + scored = [(order_score(p, cls), p) for p in itertools.permutations(range(4))] + scored.sort(key=lambda sp: (tuple(-x for x in sp[0]), tuple(lane[i].bus_id for i in sp[1]))) + dec.candidates = [{"order": [lane[i].uuid for i in p], "ordinals": [lane[i].ordinal for i in p], + "score": list(s)} for s, p in scored[:5]] + best_score, best = scored[0] + + if best_score > dec.input_score: + dec.reordered = True + dec.chosen_score = best_score + + return [lane[i] for i in best] + + dec.chosen_score = dec.input_score + + return list(lane) + + +def is_reference(topo: Topology, dec: Decision) -> bool: + lane = dec.lane + + if dec.visible != 4 or not dec.cvd_numeric_ok or dec.reordered or dec.auto_lane: + return False + + for a, b in itertools.combinations(lane, 2): + p = topo.pair(a, b) + + if p.nvlink or p.p2p_ok is not True: + return False + + numas = {g.numa for g in lane} + + return len(numas) == 1 and None not in numas + + +def caller_cvd(ctl: Controls) -> str | None: + """The caller's CUDA_VISIBLE_DEVICES, or None when unset, empty or whitespace-only (the serve script's old + ${CUDA_VISIBLE_DEVICES:-0,1,2,3} default treated an empty value as unset; whitespace is no more usable).""" + cvd = ctl.caller["CUDA_VISIBLE_DEVICES"] + + if cvd is None or CVD_UNSET_EQUIVALENT.fullmatch(cvd): + return None + + return cvd + + +def decide_env(topo: Topology | None, ctl: Controls, dec: Decision) -> None: + """Fill dec.env (value, provenance). Caller-set values always win; CUDA_VISIBLE_DEVICES may be reordered.""" + caller = ctl.caller + cvd = caller_cvd(ctl) + legacy_cvd = cvd if cvd is not None else LEGACY_CVD + cvd_prov = "caller" if cvd is not None else "legacy" + + if dec.status == "legacy" or dec.reason == "reference": + wanted = dict(LEGACY_ENV) + prov = "legacy" + wanted["CUDA_VISIBLE_DEVICES"] = legacy_cvd + else: + prov = "selector" + tp_pairs = ((dec.lane[0], dec.lane[1]), (dec.lane[2], dec.lane[3])) + tp_p2p = any(topo.pair(a, b).p2p_ok is True for a, b in tp_pairs) # type: ignore[union-attr] + wanted = {"CUDA_DEVICE_ORDER": "PCI_BUS_ID", "NCCL_P2P_LEVEL": "SYS", "NCCL_PROTO": "LL", + "VLLM_SKIP_P2P_CHECK": "0" if tp_p2p else "1"} + + if not tp_p2p and ctl.nccl == "default": + wanted["NCCL_PROTO"] = None + if not dec.reordered and dec.cvd_numeric_ok and not dec.auto_lane: + wanted["CUDA_VISIBLE_DEVICES"] = legacy_cvd + else: + wanted["CUDA_VISIBLE_DEVICES"] = ",".join(g.uuid for g in dec.lane) + cvd_prov = "selector" + + for name in ENV_VARS: + if name == "CUDA_VISIBLE_DEVICES": + dec.env[name] = (wanted[name], cvd_prov if wanted[name] == legacy_cvd else "selector") + elif caller[name] is not None: + dec.env[name] = (caller[name], "caller") + else: + dec.env[name] = (wanted[name], prov) + + +def decide_binding(topo: Topology, ctl: Controls, dec: Decision) -> None: + """Tier-1 binding: all lane GPUs on one known node N; every precondition must pass or no prefix.""" + if not ctl.bind: + return + + b = dec.binding + b["enabled"] = True + + def skip(reason: str) -> None: + b["enabled"], b["skipped_reason"] = False, reason + + # The prefix binds the whole server process tree. Never combine it with vLLM's own --numa-bind: that option + # auto-detects from the (already constrained) affinity and refuses to start, or binds a second time. + if dec.status != "ok": + return skip("legacy_decision") + + numas = {g.numa for g in dec.lane} + + if None in numas: + return skip("numa_unknown") + + if len(numas) != 1: + return skip("lane_spans_nodes") + + node = numas.pop() + cpu = topo.cpu() + + if not topo.numactl_present(): + return skip("numactl_missing") + + if cpu.affinity is None: + return skip("affinity_unknown") + + node_cpus, node_free = cpu.nodes.get(node, (None, None)) + + if node_cpus is None: + return skip("node_cpulist_unknown") + + chosen = cpu.affinity & node_cpus + + if ctl.lane_cpus is not None: + try: + chosen &= parse_cpulist(ctl.lane_cpus) + except ValueError: + return skip("lane_cpus_invalid") + + if len(chosen) < BIND_MIN_CPUS: + return skip(f"too_few_cpus({len(chosen)})") + + if cpu.mems_allowed is None or node not in cpu.mems_allowed: + return skip("node_not_in_mems_allowed") + + if node_free is None or node_free < BIND_MIN_FREE_KIB: + return skip("node_mem_free_below_1.2x69GiB") + + cpus = format_cpulist(chosen) + + if not topo.numactl_dry_run(cpus, node): + return skip("numactl_dry_run_failed") + + dec.prefix = ["numactl", f"--physcpubind={cpus}", f"--preferred={node}", "--"] + b["command"] = " ".join(dec.prefix) + b["skipped_reason"] = None + + +def decide(topo_loader, ctl: Controls) -> tuple[Decision, Topology | None]: + """Run discovery and every decision. Runs inside the timeout guard.""" + dec = Decision() + + if ctl.auto_topo == "0": + dec.reason = "auto_topo_off" + + decide_env(None, ctl, dec) + + return dec, None + + order = ctl.caller["CUDA_DEVICE_ORDER"] + + if order is not None and order != "PCI_BUS_ID": + # numeric tokens would mean a different enumeration than the one resolved here: do not touch the lane + dec.reason = "device_order_not_pci_bus_id" + + decide_env(None, ctl, dec) + + return dec, None + + try: + topo = topo_loader() + except NvmlUnavailable as e: + dec.reason = "nvml_unavailable" + + dec.warnings.append(f"NVML unavailable ({e}); using the legacy settings") + + decide_env(None, ctl, dec) + + return dec, None + + dec.facts_source = topo.source + dec.visible = len(topo.gpus) + + if len({g.bus_id for g in topo.gpus}) != len(topo.gpus) or len({g.uuid for g in topo.gpus}) != len(topo.gpus): + dec.reason = "facts_duplicate" + else: + lane = resolve_lane(topo, ctl, dec) + + if lane is not None: + for g in lane: + topo.details(g) + + dec.input_order = list(lane) + final = choose_order(topo, lane, dec) + + if final is None: + dec.lane = list(lane) + dec.reason = "facts_unknown" + else: + dec.lane = final + dec.status = "ok" + + if is_reference(topo, dec): + dec.reason = "reference" + else: + dec.reason = "reordered" if dec.reordered else "kept" + + for (i, a), (j, b) in itertools.combinations(enumerate(dec.lane), 2): + p = topo.pair(a, b) + dec.p2p[(i, j)] = p + dec.classes[(i, j)] = link_class(p) + + decide_env(topo, ctl, dec) + decide_binding(topo, ctl, dec) + dec.warnings.extend(topo.warnings) + + return dec, topo + + +# ---------------------------------------------------------------------------------------------------------------- +# output + +def summary_line(dec: Decision) -> str: + env = {k: v for k, (v, _) in dec.env.items()} + proto = env["NCCL_PROTO"] if env["NCCL_PROTO"] is not None else "unset" + bind = dec.binding + bind_s = "on" if dec.prefix else ("off" if bind["skipped_reason"] == "not_requested" + else f"off ({bind['skipped_reason']})") + + if dec.status != "ok": + return (f"topo: legacy ({dec.reason}) | CUDA_VISIBLE_DEVICES={env['CUDA_VISIBLE_DEVICES']} | " + f"skip_p2p={env['VLLM_SKIP_P2P_CHECK']} | NCCL_PROTO={proto} | bind {bind_s}") + + lane = dec.lane + numas = sorted({str(g.numa) if g.numa is not None else "?" for g in lane}) + ok = sum(1 for p in dec.p2p.values() if p.p2p_ok) + p2p_s = "none" if ok == 0 else ("all" if ok == len(dec.p2p) else f"{ok}/{len(dec.p2p)} pairs") + tp = f"({lane[0].ordinal},{lane[1].ordinal})({lane[2].ordinal},{lane[3].ordinal})" + cls = dec.classes + + return (f"topo: lane 4 GPUs [{' '.join(short_uuid(g.uuid) for g in lane)}] numa{{{','.join(numas)}}} | " + f"TP {tp} classes {fmt_class(cls[(0, 1)])}/{fmt_class(cls[(2, 3)])} | " + f"PP {fmt_class(min(cls[(0, 2)], cls[(1, 3)]))} | " + f"order {'reordered' if dec.reordered else 'kept'} | P2P {p2p_s} | " + f"skip_p2p={env['VLLM_SKIP_P2P_CHECK']} | NCCL_PROTO={proto} | bind {bind_s} | reason={dec.reason}") + + +def render_env(dec: Decision, json_path: str | None) -> str: + lines = [f"# topo_select v{VERSION} status={dec.status} reason={dec.reason}"] + + for name in ENV_VARS: + value = dec.env[name][0] + + lines.append(f"unset {name}" if value is None else f"export {name}={shlex.quote(value)}") + + prefix = " ".join(shlex.quote(w) for w in dec.prefix) + + lines.append(f"export TOPO_LAUNCH_PREFIX={shlex.quote(prefix)}") + lines.append(f"export TOPO_DECISION_JSON={shlex.quote(json_path or '')}") + + return "\n".join(lines) + "\n" + + +def diagnostics(dec: Decision, ctl: Controls) -> dict: + ple = ctl.caller["PLE_HOME"] + + try: + ple_idx = int(ple) if ple is not None else 0 + except ValueError: + ple_idx = -1 + return { + "version": VERSION, + "status": dec.status, + "reason": dec.reason, + "lane": [{"uuid": g.uuid, "bus_id": g.bus_id, "numa": g.numa, "pcie_gen_max": g.pcie_gen_max, + "pcie_width_cur": g.pcie_width_cur, "pcie_width_max": g.pcie_width_max} for g in dec.lane], + "input_order": [g.uuid for g in dec.input_order], + "link_matrix": {f"{i}-{j}": c for (i, j), c in sorted(dec.classes.items())}, + "p2p_matrix": {f"{i}-{j}": {"read": p.p2p_read, "write": p.p2p_write, "nvlink": p.nvlink, + "ancestor": p.ancestor} for (i, j), p in sorted(dec.p2p.items())}, + "candidates": dec.candidates, + "chosen_score": list(dec.chosen_score) if dec.chosen_score else None, + "input_score": list(dec.input_score) if dec.input_score else None, + "env": {k: {"value": v, "provenance": p} for k, (v, p) in dec.env.items()}, + "ple_home_uuid": dec.lane[ple_idx].uuid if 0 <= ple_idx < len(dec.lane) else None, + "p2p_probe_cache_hint": dec.env["CUDA_VISIBLE_DEVICES"][0], + "binding": dict(dec.binding), + "warnings": list(dec.warnings), + "facts_source": dec.facts_source, + } + + +def write_outputs(dec: Decision, ctl: Controls, args: argparse.Namespace) -> None: + """Diagnostics JSON first (non-fatal), then the env file (atomic); or print both for --dry-run.""" + for w in dec.warnings: + print(f"topo: warning: {w}", file=sys.stderr) + + line = summary_line(dec) + + if args.dry_run: + if not args.quiet: + print(line) + + sys.stdout.write(render_env(dec, None)) + + return + + json_path = None + + try: + json_dir = args.json_dir or os.path.join( + os.environ.get("VLLM_CACHE_ROOT") or os.environ.get("CACHE_ROOT") + or os.path.join(os.path.expanduser("~"), ".cache", "flash-next"), + "topology") + + os.makedirs(json_dir, exist_ok=True) + + lane_id = hashlib.sha256(",".join(sorted(g.uuid for g in dec.lane)).encode()).hexdigest()[:12] + stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) + path = os.path.join(json_dir, f"{lane_id}-{stamp}.json") + + atomic_write(path, json.dumps(diagnostics(dec, ctl), indent=2) + "\n") + + json_path = os.path.abspath(path) + except Exception as e: + print(f"topo: could not write diagnostics ({type(e).__name__}: {e}); continuing", file=sys.stderr) + + atomic_write(args.env_out, render_env(dec, json_path)) + + if not args.quiet: + print(line, file=sys.stderr) + + +# ---------------------------------------------------------------------------------------------------------------- +# entry point + +def read_controls(environ: dict[str, str]) -> Controls: + def get(name: str) -> str: + return environ.get(name, "").strip() + + # unset means on HERE: serve (opt-in) only calls this module with AUTO_TOPO=1 or strict, so this + # default only matters when the selector is run by hand. + auto = get("AUTO_TOPO") or "1" + + if auto not in ("1", "0", "strict"): + raise UsageError(f"AUTO_TOPO must be 1, 0 or strict (got {auto!r})") + + nccl = get("AUTO_TOPO_NCCL") or "legacy" + + if nccl not in ("legacy", "default"): + raise UsageError(f"AUTO_TOPO_NCCL must be legacy or default (got {nccl!r})") + + bind = get("AUTO_TOPO_BIND") or "0" + + if bind not in ("0", "1"): + raise UsageError(f"AUTO_TOPO_BIND must be 0 or 1 (got {bind!r})") + + lane = get("LANE_GPUS") + + if lane not in ("", "auto"): + raise UsageError(f"LANE_GPUS supports only 'auto' (got {lane!r}); pick a lane with CUDA_VISIBLE_DEVICES") + + caller: dict[str, str | None] = {} + + for v in SNAPSHOT_VARS: + caller[v] = environ.get(f"TOPO_CALLER_{v}", "") if environ.get(f"TOPO_CALLER_SET_{v}") == "1" else None + + # LANE_CPUS: unset -> no restriction; set (even to '') -> exactly that list (empty -> binding is skipped) + return Controls(auto, nccl, bind == "1", lane == "auto", environ.get("LANE_CPUS"), caller) + + +def main(argv: list[str] | None = None) -> int: + for stream in (sys.stdout, sys.stderr): + try: + stream.reconfigure(errors="replace", newline="\n") # type: ignore[attr-defined] + except Exception: + pass + + ap = argparse.ArgumentParser(description="Choose GPU order, P2P/NCCL settings and optional NUMA binding " + "for the 4-GPU serving lane.") + + ap.add_argument("--env-out", help="write the decision env file here (required unless --dry-run)") + ap.add_argument("--json-dir", help="diagnostics directory (default $VLLM_CACHE_ROOT/topology, else " + "$CACHE_ROOT/topology, else ~/.cache/flash-next/topology)") + ap.add_argument("--dry-run", action="store_true", help="print the summary and env file; write nothing") + ap.add_argument("--quiet", action="store_true", help="no summary line") + + args = ap.parse_args(argv) # exits 2 on bad arguments + + try: + if not args.dry_run and not args.env_out: + raise UsageError("--env-out PATH is required unless --dry-run") + + ctl = read_controls(dict(os.environ)) + raw_timeout = os.environ.get("TOPO_TIMEOUT", "").strip() or "20" + + try: + timeout = float(raw_timeout) + + if not (math.isfinite(timeout) and 0 < timeout <= threading.TIMEOUT_MAX): + raise ValueError + except ValueError: + raise UsageError(f"TOPO_TIMEOUT must be a finite number of seconds above 0 and at most " + f"{threading.TIMEOUT_MAX:g} (got {raw_timeout!r})") from None + except UsageError as e: + print(f"topo_select: {e}", file=sys.stderr) + + return EXIT_USAGE + + fixture = os.environ.get("TOPO_FIXTURE") + loader = (lambda: read_facts_fixture(fixture)) if fixture else read_facts_nvml + + result: dict[str, Any] = {} + + def work() -> None: + topo = None + + try: + result["value"] = decide(loader, ctl) + topo = result["value"][1] + + except BaseException as e: # handed to the main thread + result["error"] = e + + finally: + if isinstance(topo, NvmlTopology): + topo.close() + + try: + t = threading.Thread(target=work, name="topo-discovery", daemon=True) + + t.start() + t.join(timeout) + + alive = t.is_alive() + except Exception as e: + print(f"topo: internal error: discovery thread failed: {type(e).__name__}: {e}", file=sys.stderr) + + return EXIT_INTERNAL + + if alive: + print(f"topo: discovery timed out after {raw_timeout}s", file=sys.stderr) + + sys.stderr.flush() + os._exit(EXIT_TIMEOUT) # a hung NVML call cannot be interrupted; leave without waiting for it + + if "error" in result: + e = result["error"] + + if isinstance(e, StrictRefusal): + print(str(e), file=sys.stderr) + + return EXIT_STRICT + print(f"topo: internal error: {type(e).__name__}: {e}", file=sys.stderr) + + return EXIT_INTERNAL + + dec, _ = result["value"] + + try: + write_outputs(dec, ctl, args) + except Exception as e: + print(f"topo: internal error while writing outputs: {type(e).__name__}: {e}", file=sys.stderr) + + return EXIT_INTERNAL + + return EXIT_OK + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/start.sh b/start.sh new file mode 100755 index 0000000..d1eaf80 --- /dev/null +++ b/start.sh @@ -0,0 +1,275 @@ +#!/usr/bin/env bash + +# Serve Qwen3.8-Flash-Next-W4A16 on 4x RTX 3090. Builds first when the tree or the venv is +# missing, checks the box, then runs the engine. One script for bare metal and for the image. +# +# sudo ./start.sh [checkpoint-dir] [extra vllm args] # build on first run, then serve +# ./build.sh --force # compile again (see ./build.sh) +# +# Shape: TP2 x PP2 + EP, MTP K=3, 262K per request, ~925K-token FP8 KV pool, full 262K context, +# host-resident PLE/embeddings, prefix caching with a 64-token match unit, RecoverSSM, side-cache +# layout. This tree's profile guard boots ONLY the command at the end of this script; +# --enforce-eager, other cudagraph modes and breakable graphs are refused. +# +# Env: +# CKPT / MODEL_MOUNT_PATH checkpoint dir (bare metal: the Swift1.5 Qwen3.8-Flash-Next W4A16 +# build; container: /models/Qwen3.8-Flash-Next). A first argument that starts with +# `-` is a vllm argument, so the default checkpoint is kept. +# TREE VENV serving tree and venv (default: ./vllm and ./.venv, built by build.sh) +# PORT (8000) HOST (127.0.0.1) MODEL_NAME ("name1 name2 ...", first is the health probe's) +# SCALES KV scale sidecar (default: qsa_kv_scales_swift.json inside the checkpoint) +# CACHE_ROOT compile/cudagraph cache (default: .vllm-cache next to this script, on disk) +# PLE_HOME GPU that hosts the PLE offload home (0) +# GUARD bounds-guard telemetry off|warn|raise (warn) +# COUNTERS=1 logging-only counters (top-k exact restarts /30s, FP8 KV clips /300s) +# VLLM_LOG_STATS_INTERVAL seconds between engine stats lines (unset = vLLM default 1 s) +# CUDA_HOME toolkit for first-serve JIT; default: the venv's CUDA 13.0 wheels +# AUTO_TOPO 0 (default, fixed shipped settings) | 1 | strict -> scripts/topo_select.py +# AUTO_TOPO_NCCL=default AUTO_TOPO_BIND=1 LANE_GPUS=auto LANE_CPUS=... TOPO_TIMEOUT=20 +# LOG log file (/var/log/vllm.log). LOG= (empty) keeps the output on the terminal, +# which is what the image does so `docker compose logs` works. +# +# First serve compiles kernels (~6 min, writes $CACHE_ROOT). RAM: the PLE n-gram table is host +# memory, resident floor ~69 GiB. /dev/shm >= 1 GB. One CUDA process per GPU. + +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +TREE="${TREE:-$HERE/vllm}"; VENV="${VENV:-$HERE/.venv}" +CKPT="${CKPT:-${MODEL_MOUNT_PATH:-/data/models/halt95/Swift1.5-Qwen3.8-Flash-Next-W4A16-Merlin}}" +CACHE_ROOT="${CACHE_ROOT:-$HERE/.vllm-cache}" +# Colon-free default: LOG= (set, empty) is the container's "do not redirect". +LOG="${LOG-/var/log/vllm.log}" + +if [ $# -gt 0 ]; then + case "$1" in -*) ;; "") shift ;; *) CKPT="$1"; shift ;; esac +fi +case "$TREE" in /*) ;; *) TREE="$PWD/$TREE" ;; esac +case "$VENV" in /*) ;; *) VENV="$PWD/$VENV" ;; esac + +# --- 1. build when this box has no serving tree or no venv --------------------------------------- +mkdir -p "$CACHE_ROOT" 2>/dev/null || { echo "cannot create $CACHE_ROOT; set CACHE_ROOT="; exit 1; } +if mountpoint -q "$CACHE_ROOT" 2>/dev/null; then cache_kind="RAM"; else cache_kind="disk"; fi +echo "cache: $CACHE_ROOT ($cache_kind)" +if [ ! -f "$TREE/vllm/__init__.py" ] || [ ! -x "$VENV/bin/vllm" ]; then + echo "first run: building the tree and the venv (./build.sh, ~20 min on 64 cores)" + "$HERE/build.sh" +fi + +# --- 2. the checkpoint --------------------------------------------------------------------------- +if [ ! -f "$CKPT/config.json" ]; then + echo "checkpoint not found at $CKPT" + echo "bare metal: sudo ./start.sh /path/to/Qwen3.8-Flash-Next-W4A16-Merlin" + echo "container: -v /path/to/Qwen3.8-Flash-Next-W4A16-Merlin:/models/Qwen3.8-Flash-Next:ro" + exit 1 +fi +if ! "$VENV/bin/python" -c 'import json,sys; c=json.load(open(sys.argv[1])); sys.exit(0 if c.get("text_config",{}).get("ple_embedding_dtype")=="float8_e4m3fn" else 1)' "$CKPT/config.json"; then + if [ -w "$CKPT" ] && [ -w "$CKPT/config.json" ]; then + echo "config.json lacks text_config.ple_embedding_dtype; adding it (weights untouched, backup kept)" + python3 "$HERE/scripts/make-e1-config.py" "$CKPT" + else + echo "config.json lacks text_config.ple_embedding_dtype and the directory is read-only." + echo "Run once on the writable copy: python3 scripts/make-e1-config.py $CKPT" + exit 1 + fi +fi +if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L >/dev/null 2>&1; then + echo "WARNING: no NVIDIA GPU visible (in Docker: --gpus all, and nvidia-ctk runtime configure --runtime=docker)" +fi + +# --- 3. the build products ----------------------------------------------------------------------- +[ -f "$TREE/vllm/__init__.py" ] || { echo "TREE=$TREE is not the serving tree"; exit 1; } +[ -f "$TREE/.build-products" ] && [ -s "$TREE/.build-products.sha256" ] \ + || { echo "TREE=$TREE has no verified build products (./build.sh)"; exit 1; } +(cd "$TREE" && sha256sum -c --quiet --strict .build-products.sha256) \ + || { echo "TREE=$TREE: build products changed since ./build.sh verified them"; exit 1; } +[ -x "$VENV/bin/python" ] || { echo "VENV=$VENV has no bin/python"; exit 1; } + +# A vllm/ in the launch directory shadows $TREE: `python -m` puts CWD before PYTHONPATH. +if [ -f "$PWD/vllm/__init__.py" ] && [ "$(cd "$PWD/vllm" 2>/dev/null && pwd -P)" != "$(cd "$TREE/vllm" && pwd -P)" ]; then + echo "refusing to start: $PWD/vllm would shadow $TREE/vllm; start from another directory"; exit 1 +fi + +# --- 4. first-serve toolchain (FlashInfer JITs against CUDA_HOME, Triton against Python.h) ------- +if [ -z "${CUDA_HOME:-}" ]; then + CUDA_HOME="$("$VENV/bin/python" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')/nvidia/cu13" +fi +export CUDA_HOME PATH="$CUDA_HOME/bin:$PATH" +for t in cc c++ ninja; do command -v "$t" >/dev/null || { echo "first serve needs $t (apt install gcc g++ ninja-build)"; exit 1; }; done +"$VENV/bin/python" -c 'import os,sys,sysconfig; sys.exit(0 if os.path.exists(os.path.join(sysconfig.get_paths()["include"],"Python.h")) else 1)' \ + || { echo "first serve needs the Python 3.13 headers (apt install python3.13-dev)"; exit 1; } +[ -x "$CUDA_HOME/bin/nvcc" ] || { echo "no nvcc at CUDA_HOME=$CUDA_HOME"; exit 1; } +[ -e "$CUDA_HOME/lib64/libcudart.so" ] || { echo "$CUDA_HOME/lib64/libcudart.so missing (./build.sh creates it)"; exit 1; } +nv_rel="$("$CUDA_HOME/bin/nvcc" --version | sed -n 's/.*release \([0-9][0-9]*\.[0-9][0-9]*\),.*/\1/p' || true)" +[ -n "$nv_rel" ] || { echo "no CUDA release from \`$CUDA_HOME/bin/nvcc --version\`"; exit 1; } +drv_rel="$(nvidia-smi 2>/dev/null | sed -n 's/.*CUDA \(UMD \)\{0,1\}Version: *\([0-9][0-9]*\.[0-9][0-9]*\).*/\2/p' | head -1 || true)" +[ -n "$drv_rel" ] || echo "note: driver CUDA version unreadable; the nvcc-vs-driver check is skipped" +if [ -n "$drv_rel" ] && [ "$(printf '%s\n%s\n' "$nv_rel" "$drv_rel" | sort -V | tail -1)" != "$drv_rel" ]; then + echo "nvcc CUDA $nv_rel is newer than the driver's CUDA $drv_rel: first-serve kernels would be rejected."; exit 1 +fi + +# --- 5. topology (opt-in; default = the fixed shipped settings) ---------------------------------- +case "${AUTO_TOPO:-0}" in + 0) topo_on=0 ;; + 1 | strict) topo_on=1 ;; + *) echo "topo: note: AUTO_TOPO='$AUTO_TOPO' is not 0, 1 or strict; the selector stays off"; topo_on=0 ;; +esac +topo_launch_prefix="" +legacy_topo_env() { + if [ "${1:-}" = fallback ]; then + if [ "${TOPO_CALLER_SET_CUDA_VISIBLE_DEVICES:-0}" = 1 ]; then CUDA_VISIBLE_DEVICES="$TOPO_CALLER_CUDA_VISIBLE_DEVICES" + else unset CUDA_VISIBLE_DEVICES; fi + case "${CUDA_VISIBLE_DEVICES:-}" in *[![:space:]]*) ;; *) unset CUDA_VISIBLE_DEVICES ;; esac + export TOPO_LAUNCH_PREFIX="" TOPO_DECISION_JSON=""; topo_launch_prefix="" + fi + export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" + export CUDA_DEVICE_ORDER=PCI_BUS_ID + export VLLM_SKIP_P2P_CHECK=1 NCCL_P2P_LEVEL=SYS NCCL_PROTO=LL +} +topo_env_file_ok() { + local f="$1" i=0 line + local val="([A-Za-z0-9_@%+=:,./-]+|'[^']*'(\"'\"'[^']*')*)" + local -a want=( + "^# topo_select v1 status=(ok|legacy) reason=[A-Za-z0-9_.-]+$" + "^export CUDA_VISIBLE_DEVICES=$val$" + "^export CUDA_DEVICE_ORDER=$val$" + "^export VLLM_SKIP_P2P_CHECK=$val$" + "^(export NCCL_P2P_LEVEL=$val|unset NCCL_P2P_LEVEL)$" + "^(export NCCL_PROTO=$val|unset NCCL_PROTO)$" + "^export TOPO_LAUNCH_PREFIX=('numactl --physcpubind=[0-9,-]+ --preferred=[0-9]+ --'|'')$" + "^export TOPO_DECISION_JSON=$val$") + [ -f "$f" ] || return 1 + while IFS= read -r line || [ -n "$line" ]; do + [ "$i" -lt "${#want[@]}" ] || return 1 + [[ $line =~ ${want[$i]} ]] || return 1 + i=$((i + 1)) + done < "$f" || return 1 + [ "$i" -eq "${#want[@]}" ] +} +topo_unavailable() { + if [ "$AUTO_TOPO" = strict ]; then echo "topo: $1; AUTO_TOPO=strict refuses to guess"; exit 1; fi + echo "topo: $1, using legacy settings"; legacy_topo_env fallback +} +if [ "$topo_on" = 0 ]; then + legacy_topo_env +else + for v in CUDA_VISIBLE_DEVICES CUDA_DEVICE_ORDER VLLM_SKIP_P2P_CHECK NCCL_P2P_LEVEL NCCL_PROTO PLE_HOME; do + if [ "${!v+set}" = set ]; then export "TOPO_CALLER_SET_$v=1" "TOPO_CALLER_$v=${!v}" + else export "TOPO_CALLER_SET_$v=0"; fi + done + export TOPO_LAUNCH_PREFIX="" TOPO_DECISION_JSON="" + topo_timeout=20 + case "${TOPO_TIMEOUT:-}" in + [1-9]|[1-9][0-9]|[1-9][0-9][0-9]) [ "$TOPO_TIMEOUT" -le 600 ] && topo_timeout="$TOPO_TIMEOUT" ;; + esac + topo_json_dir="${CACHE_ROOT:-$HERE/.vllm-cache}/topology" + tmp_topo_env="" + if tmp_topo_env="$(mktemp 2>/dev/null)" && [ -n "$tmp_topo_env" ]; then + trap 'rm -f -- "$tmp_topo_env" 2>/dev/null || :' EXIT + rc=0 + TOPO_TIMEOUT="$topo_timeout" timeout --kill-after=5 "$((topo_timeout + 5))" \ + "$VENV/bin/python" "$HERE/scripts/topo_select.py" --env-out "$tmp_topo_env" --json-dir "$topo_json_dir" || rc=$? + if [ "$rc" = 3 ]; then exit 1 + elif [ "$rc" != 0 ]; then topo_unavailable "selector unavailable (rc=$rc)" + elif ! topo_env_file_ok "$tmp_topo_env"; then topo_unavailable "selector handed back an unusable env file" + else + # shellcheck disable=SC1090 + if ! source "$tmp_topo_env"; then topo_unavailable "could not apply the selector's env file" + else topo_launch_prefix="$TOPO_LAUNCH_PREFIX"; fi + fi + rm -f -- "$tmp_topo_env" 2>/dev/null || :; trap - EXIT + else + topo_unavailable "could not create a temporary file" + fi +fi +if [ "${SERVE_PRINT_ENV_AND_EXIT:-0}" = 1 ]; then + echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-} CUDA_DEVICE_ORDER=${CUDA_DEVICE_ORDER:-} VLLM_SKIP_P2P_CHECK=${VLLM_SKIP_P2P_CHECK:-} NCCL_P2P_LEVEL=${NCCL_P2P_LEVEL-} NCCL_PROTO=${NCCL_PROTO-} PREFIX=$topo_launch_prefix" + exit 0 +fi +# --- end topology -------------------------------------------------------------------------------- + +# --- 6. the serve profile ------------------------------------------------------------------------- +export VLLM_PP_LAYER_PARTITION=25,23 +# the shipped profile sets every tree knob explicitly; drop inherited ones first +for v in $(env | sed -n 's/^\(MERLIN_[A-Z0-9_]*\|VLLM_E2_[A-Z0-9_]*\)=.*/\1/p'); do unset "$v"; done +unset VLLM_PLE_CPU_OFFLOAD VLLM_PLE_FP8_GLOBAL_SCALE VLLM_PLE_TEST_FAULT VLLM_PLE_TEST_FAULT_GATE VLLM_PLE_TEST_FAULT_LOG \ + VLLM_QSA_KV_COLLECT VLLM_QSA_KV_CLIP_COUNT VLLM_E2_ALLOW_BREAKABLE VLLM_E2_EXPECTED_CUDAGRAPH_MODE \ + VLLM_E2_MAMBA_SLOT_OWNERSHIP VLLM_E2_MEMSTATS VLLM_E2_MEMSTATS_DIR VLLM_E2_PLE_FLAG_CHECK \ + VLLM_E2_SCHED_LOG VLLM_E2_ZERO_OWNERSHIP_CHECK VLLM_E2_GUARD_INTERVAL \ + VLLM_PLE_OFFLOAD_WAIT_TIMEOUT VLLM_PLE_PROGRESS_DEADLINE_S VLLM_PLE_STARTUP_DEADLINE_S \ + PYTORCH_CUDA_ALLOC_CONF || true +export VLLM_E2_PLE_PULL_TRANSPORT=1 VLLM_PLE_OFFLOAD_READY_TIMEOUT=1800 +export VLLM_PLE_OFFLOAD_HOME_DEVICE="${PLE_HOME:-0}" +export VLLM_HOST_EMBED_TABLE=1 +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +export VLLM_QSA_KV_SCALES="${SCALES:-$CKPT/qsa_kv_scales_swift.json}" +# engine stats line (throughput, KV usage, queue). Set empty to keep vLLM's default. +if [ -n "${VLLM_LOG_STATS_INTERVAL:-}" ]; then export VLLM_LOG_STATS_INTERVAL; fi +export HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 +# the kernel toggles (the shipped values). These names are the tree's own: envs.py in the +# serving tree registers MERLIN_* for the kernel toggles and VLLM_E2_* for the runtime knobs. +export MERLIN_FULL_K=1 MERLIN_QSA_SORT=1 MERLIN_TIE_DET=1 MERLIN_TIE_RECENT=1 MERLIN_TOPK_SORTED_EMIT=1 +export VLLM_E2_GUARD_MODE="${GUARD:-warn}" +export VLLM_E2_SIDE_CACHE_LAYOUT=1 VLLM_USE_V2_MODEL_RUNNER=1 VLLM_E2_MAMBA_RETAIN_CHECKPOINT=1 +if [ "${COUNTERS:-0}" = 1 ]; then export MERLIN_TOPK_EXACT_COUNTERS_LOG=30 VLLM_QSA_KV_CLIP_COUNT=300; fi +[ -f "$VLLM_QSA_KV_SCALES" ] || { echo "scales sidecar missing: $VLLM_QSA_KV_SCALES (FP8 KV would run at scale 1.0)"; exit 1; } + +# Compile cache: on disk by default, so it survives a reboot. A cold cache costs ~6 min of kernel +# compile, and the cache is only ~160 MB, so disk is the right home for it. Keep ONE root per +# deployment -- a new root costs a fresh compile. CACHE_ROOT=/vllm-cache (a tmpfs) is the opt-in +# for a box where you want it in RAM. +[ -d "$CACHE_ROOT" ] || { mkdir -p "$CACHE_ROOT" 2>/dev/null || { echo "cannot create $CACHE_ROOT"; exit 1; }; } +if df -T "$CACHE_ROOT" 2>/dev/null | grep -q tmpfs; then + echo "note: $CACHE_ROOT is a tmpfs; the compile cache is lost on reboot" +fi +export VLLM_CACHE_ROOT="$CACHE_ROOT" +export FLASHINFER_WORKSPACE_BASE="$CACHE_ROOT/flashinfer" TRITON_CACHE_DIR="$CACHE_ROOT/triton" +export PYTHONPATH="$TREE" VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 + +# host RAM (the n-gram table lives here; a short host OOMs mid-boot) and /dev/shm +mem_kb="$(awk '/^MemTotal:/{print $2}' /proc/meminfo 2>/dev/null || true)" +for f in /sys/fs/cgroup/memory.max /sys/fs/cgroup/memory/memory.limit_in_bytes; do + lim="$(cat "$f" 2>/dev/null || true)" + case "$lim" in ''|*[!0-9]*) continue ;; esac + lim_kb=$((lim / 1024)); if [ -z "${mem_kb:-}" ] || [ "$lim_kb" -lt "$mem_kb" ]; then mem_kb="$lim_kb"; fi +done +if [ -n "${mem_kb:-}" ] && [ "$mem_kb" -lt 100663296 ]; then + echo "WARNING: ${mem_kb} KB visible RAM (< the required 96 GiB); resident floor is ~69 GiB" +fi +shm_kb="$(df -Pk /dev/shm 2>/dev/null | awk 'NR==2{print $2}' || true)" +case "$shm_kb" in *[!0-9]*) shm_kb= ;; esac +if [ -n "${shm_kb:-}" ] && [ "$shm_kb" -lt 1048576 ]; then + echo "WARNING: /dev/shm is $((shm_kb / 1024)) MB; the PLE transport needs >= 1 GB" +fi + +read -r -a NAMES <<< "${MODEL_NAME:-qwen38-flash-next}" +echo "serving $CKPT as ${NAMES[*]} on ${HOST:-127.0.0.1}:${PORT:-8000} (cache $CACHE_ROOT)" +if [ -n "$LOG" ]; then + # stdout + stderr, so the PLE offload child and the engine workers land in the same file. + # systemd/vllm.logrotate rotates it: daily, 100 MB cap, 14 files, gzip. Read it with: tail -f $LOG + touch "$LOG" 2>/dev/null || { echo "cannot write $LOG; run as root or set LOG="; exit 1; } + echo "log: $LOG (rotate: systemd/vllm.logrotate)" + exec >> "$LOG" 2>&1 +fi + +# unquoted on purpose: word-splits into a numactl prefix from the validated selector file, empty when off +exec ${topo_launch_prefix:+$topo_launch_prefix} "$VENV/bin/vllm" serve "$CKPT" \ + --served-model-name "${NAMES[@]}" \ + --host "${HOST:-127.0.0.1}" --port "${PORT:-8000}" \ + --tensor-parallel-size 2 --pipeline-parallel-size 2 --enable-expert-parallel \ + --speculative-config '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \ + --max-model-len 262144 --max-num-seqs 2 --max-num-batched-tokens 4096 \ + --gpu-memory-utilization 0.96 --kv-cache-memory 4100000000 --kv-cache-dtype fp8_e4m3 \ + --limit-mm-per-prompt '{"image":42,"video":0}' --mm-processor-cache-gb 0 \ + --mm-processor-kwargs '{"size":{"longest_edge":4194304,"shortest_edge":65536}}' \ + --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,3,4,6,8,12,16,24,32]}' \ + --enable-prefix-caching --enable-prompt-tokens-details \ + --no-enable-flashinfer-autotune \ + --enable-auto-tool-choice --tool-call-parser qwen3_coder --reasoning-parser qwen3 \ + --override-generation-config '{"temperature": 1.0, "top_p": 0.95, "top_k": 20, "min_p": 0.0, "presence_penalty": 0.0, "repetition_penalty": 1.0}' \ + --default-chat-template-kwargs '{"enable_thinking": true, "reasoning_effort": "low"}' \ + --shutdown-timeout 10 \ + --block-size 4096 --use-replayssm \ + --prefix-match-unit 64 \ + --safetensors-load-strategy prefetch \ + "$@" diff --git a/systemd/vllm.logrotate b/systemd/vllm.logrotate new file mode 100644 index 0000000..7891777 --- /dev/null +++ b/systemd/vllm.logrotate @@ -0,0 +1,16 @@ +# Bare-metal log policy for /var/log/vllm.log (start.sh writes it, LOG= to move it). +# Install: sudo install -m 0644 systemd/vllm.logrotate /etc/logrotate.d/vllm +# The container needs none of this: docker-compose.yml caps the json-file log at 100m x 14. + +/var/log/vllm.log { + daily + maxsize 100M + rotate 14 + compress + delaycompress + missingok + notifempty + copytruncate + maxage 60 + create 0640 root root +} diff --git a/systemd/vllm.service b/systemd/vllm.service new file mode 100644 index 0000000..0b1e593 --- /dev/null +++ b/systemd/vllm.service @@ -0,0 +1,27 @@ +[Unit] +Description=vLLM (Qwen3.8-Flash-Next-W4A16 on 4x RTX 3090, TP2xPP2+EP, MTP, 262K) +After=network-online.target +Wants=network-online.target + +[Service] +Type=simple +User=root +EnvironmentFile=-/etc/default/vllm +# start.sh is the single source of truth; the checkpoint goes first (CKPT in /etc/default/vllm) +WorkingDirectory=/opt/vllm +# The compile cache is $CACHE_ROOT, default /opt/vllm/.vllm-cache (on disk: it survives a +# reboot). A cold cache costs ~6 min of kernel compile at the next start, so do not use tmpfs here. +# Set CACHE_ROOT in /etc/default/vllm to move it. +ExecStart=/opt/vllm/start.sh +# MTP + cudagraphs do not always exit on SIGTERM: give --shutdown-timeout 10 room, then take the cgroup +KillMode=control-group +TimeoutStopSec=20 +Restart=on-failure +RestartSec=15 +LimitMEMLOCK=infinity +LimitNOFILE=1048576 +# start.sh writes the engine and the PLE offload child to $LOG (/var/log/vllm.log by default); +# systemd/vllm.logrotate rotates it. The pre-serve lines stay in the journal: journalctl -u vllm + +[Install] +WantedBy=multi-user.target