@@ -0,0 +1,10 @@
|
||||
/vllm
|
||||
/.venv
|
||||
/.git
|
||||
/.pi
|
||||
__pycache__
|
||||
*.pyc
|
||||
/.vllm-cache
|
||||
/artifacts
|
||||
/.env
|
||||
/build-*.log
|
||||
@@ -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
|
||||
+20
@@ -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
|
||||
+103
@@ -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"]
|
||||
@@ -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.
|
||||
@@ -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-<timestamp>.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. |
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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-<tag>.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=<tarball> 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=<path>, 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=<path> 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=<tarball> 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=<tarball> 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/<name>.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)"
|
||||
@@ -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}
|
||||
@@ -0,0 +1,970 @@
|
||||
From 1b2ba827becf485be6c06d2447b1faf2772c42c9 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,294 @@
|
||||
From f36b743fb015941866baa173d0b0efc66f694d98 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
From 8b49ca6aeeccc87962e3976d564932e099f4c573 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,536 @@
|
||||
From 5a84c1bc7dc5e9a54d5a14ce1766adeb6f1dca10 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,927 @@
|
||||
From 609c86adc68243fabc81d603cd424687f829669a Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,101 @@
|
||||
From 8cb229146ba10ea093d647deeaf9a074afaa5bee Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,399 @@
|
||||
From b60dadeaaf41349275337d5d1fcd0b1a86d630c7 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
From 25dd4897d4378bca33e56a2f332723d599ec3e57 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
From b9742140206908ce1fecb5026ccb0a981c8775c0 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
From 612e3d64642c5e3a0edba7d37d53f3b2ae9cd8da Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,263 @@
|
||||
From 430b5ec450421d89201331ebe04450d135e75517 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
From 85f78b1410981293e030f6256c0c16c84e1a12d7 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
From 298ddf906275a3a3743d94f50021c755b658a39f Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
From cda271d5ea26ff3c03c5d202041f9a09552c5ef3 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,444 @@
|
||||
From d90c541cf2da170e17a94cebaf19dbda03fda390 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
From 4823da32e2090433d0f56e40316de3cfcf637d2f Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
From 5db178597cf9d2d250ce7d324bd995b45e91af29 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
From 3191b58705d1a99632ee8e7fcc820eb3323c230b Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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=<path>).
|
||||
+
|
||||
+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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,330 @@
|
||||
From 0271354dc4a882fdcc5ef7eca2f5c55fba812925 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
From de31f552293348bf1b4cf02f394fef97d071031b Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,353 @@
|
||||
From c390570424760dafe3861d388a16105badd7706e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 <thisfile>``
|
||||
|
||||
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
|
||||
|
||||
@@ -0,0 +1,416 @@
|
||||
From 77661daf6fc3e9023c7181fb558c0dcb370e6eff Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,365 @@
|
||||
From 300e10fef5c7b4e9aeae0838ab086f5b7af315e9 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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=<dir>).
|
||||
+
|
||||
+Each process appends framed records to <dir>/<tag>-pid<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("<BBHIII", kind, rank, 0, seq, n, row_bytes)
|
||||
+ + ids_cpu.numpy().tobytes() + raw)
|
||||
+ with _LOCK:
|
||||
+ f = _FILES.get(tag)
|
||||
+ if f is None:
|
||||
+ os.makedirs(_DIR, exist_ok=True)
|
||||
+ f = open(os.path.join(_DIR, f"{tag}-pid{os.getpid()}.bin"), "ab", buffering=0)
|
||||
+ _FILES[tag] = f
|
||||
+ f.write(record)
|
||||
diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py
|
||||
index 43caf80..24beebd 100644
|
||||
--- a/vllm/v1/ple_offload/worker.py
|
||||
+++ b/vllm/v1/ple_offload/worker.py
|
||||
@@ -40,6 +40,7 @@ from typing import Any, cast
|
||||
|
||||
import msgspec
|
||||
import torch
|
||||
+from vllm.v1.ple_offload import rowdump as _rowdump
|
||||
import torch.distributed as dist
|
||||
import zmq
|
||||
|
||||
diff --git a/vllm/v1/worker/e2_memstats.py b/vllm/v1/worker/e2_memstats.py
|
||||
new file mode 100644
|
||||
index 0000000..586f8fa
|
||||
--- /dev/null
|
||||
+++ b/vllm/v1/worker/e2_memstats.py
|
||||
@@ -0,0 +1,136 @@
|
||||
+# SPDX-License-Identifier: Apache-2.0
|
||||
+"""Debug-only allocator telemetry for E2, BOUNDARY-ONLY. Off unless VLLM_E2_MEMSTATS=<ctl poll seconds> and
|
||||
+VLLM_E2_MEMSTATS_DIR=<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 <dir>/window.ctl, which holds "<seq> <name>". 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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,67 @@
|
||||
From 236b01309da7e3d934cd9876cf929b2ea68c8e54 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
From 2179328624d5718a369f599920d285040c572a9d Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
From 5d513783b79de125c9ffb96d928919c706777f3e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
From 0ba64cc82a5962faf443b63e821399a82e77997e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
From fc34b9371ed0900ff970fc7631d41a13bf49ee22 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,346 @@
|
||||
From 8be2e3ab6246beafa10aa7ed2de9591c1b4cfc37 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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<i>` names rather than mislabelling.
|
||||
+
|
||||
+HOW THE SMOKE READS IT. Set `MERLIN_TOPK_EXACT_COUNTERS_LOG=<seconds>` 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 <engine log> | ...` 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
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
From 090e8263f6d923c61fa6bc328e5a7b395aff5ae5 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,788 @@
|
||||
From 5ff40ff3254e92df316c73efa58bba20a12a7d35 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 <int TopK>
|
||||
+__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<int>(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<int>(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<int>(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<uint32_t>(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<uint32_t>(__popc(bgt & lane_mask));
|
||||
+ const uint32_t eq_in_warp = static_cast<uint32_t>(__popc(beq & lane_mask));
|
||||
+ if (lane == 0u) {
|
||||
+ warp_gt[warp] = static_cast<uint32_t>(__popc(bgt));
|
||||
+ warp_eq[warp] = static_cast<uint32_t>(__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<uint32_t>(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<uint32_t>(TopK)) {
|
||||
+ row_output[slot] = static_cast<int32_t>(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<uint32_t>(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<TopK>(
|
||||
+ 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<TopK, VEC_SIZE>(
|
||||
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 <cuda_runtime.h>
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
+#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
@@ -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<uint32_t>(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<uint32_t>(P::RADIX);
|
||||
+ params.det_large = det_large_enabled() && det_large_fits;
|
||||
+ if (det_large_enabled() && !det_large_fits) {
|
||||
+ static std::atomic<bool> 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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,203 @@
|
||||
From e42e2eee792338708274acd3d0dc5446bfb3951e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 '<unbound>' 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("<unbound>").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 == "<unbound>":
|
||||
+ 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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,130 @@
|
||||
From 4fcb598510da1f2f4fad1724fdb88c34d1301e07 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,554 @@
|
||||
From e5f6387064eadbeaba0ae923f0141ef9f638b9af Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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\[(?P<tag>PP\d+_TP\d+)\]: ple_backend=(?P<backend>\S+) "
|
||||
+ r"ple_method=(?P<method>\S+) ple_weight_dtype=(?P<dtype>\S+) "
|
||||
+ r"breakable=(?P<brk>ON|OFF) "
|
||||
+ r"compile_mode=(?P<cmode>\S+) cudagraph_mode=(?P<gmode>\S+) "
|
||||
+ r"role=(?P<role>target|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=<dtype> breakable=<ON|OFF> compile_mode=<mode>
|
||||
cudagraph_mode=<mode> role=<target|draft>
|
||||
|
||||
-(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:<m>,decode:<m>`` 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
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
From 0ed590631391b088b17c35ffb5c54fca34dfc0b8 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 "<unknown type>"
|
||||
+ if type(qualname) is not str:
|
||||
+ qualname = "<unknown type>"
|
||||
+ 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
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
From 579390199295b0b0447555198eaef7ddf22d2555 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 <this file>
|
||||
+Env: REPLAY_SENTINEL_AR=custom|nccl (default custom), REPLAY_SENTINEL_CYCLES (default 3),
|
||||
+ REPLAY_SENTINEL_OUT=<dir> 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("<u2").tobytes()).hexdigest()
|
||||
+
|
||||
+
|
||||
+def derive() -> 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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,43 @@
|
||||
From 22e10be45aebad0611c2a55ebf6318e08ad2db0b Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
From 1608d7132324e5f2f133f2b61e6a7405d0472020 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
From 108e2c9e5b62d1990796ee3f710d39f5063b3759 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
From f26f0f4a9cf4e746112161a64838ed2d80ae18b3 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
From f4efb5ef63426dad77a45b5a9787fd05fbac5638 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,362 @@
|
||||
From 26f4ff1480ab7307aaad9d27309fb5f5a85061e3 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
From b24ac2b5bbf91bd73342d275f39aa9b6bb3f338a Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,454 @@
|
||||
From f0949d57a607a8882d0cba5a9937fa03bb82eefc Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
From ec0e5762517c1d0cd6a4d2e7a5303ce028d5752a Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
From 0f95d8a93025684b4a36b6e57ccd908b589e812e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
From 22b00ba2e2239b52923cb555ae6e24198bb59835 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
From cde56de93bd067484ced3193b07fa0f93e94578d Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,730 @@
|
||||
From 7d1df38791dfe3678ae8389e29c3fd3db3ea586b Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,244 @@
|
||||
From 57a0d446925c54448a4bfbc7879d6556a77101cc Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
From aaffb7483252d9ff8a6bb395ceba4f1d87f5a981 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
From 1954885e52446a5979dd514c004fe5d76bb5ed12 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
From f5b942010a1cce6c0807bcd9cbcc36c1f6bec792 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
From 9caab704196d6854658181e1e995a2869bc2d38e Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
From c2a86c85f32040d36d0837657994478c731101a7 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
From e6f5849032deddd8eee331d99242d3ab646198f1 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,251 @@
|
||||
From a70876f4aacbad42d7c4e608b439e5af425578f5 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
From 029f835d8c5becced0382796057a3152bc2b41a3 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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 <tree> <outdir>
|
||||
+# cd <tree> && CUDA_VISIBLE_DEVICES= TRITON_INTERPRET=1 PYTHONPATH=.:<dir with triton_interp_fix.py> \
|
||||
+# python -m pytest --noconftest -p triton_interp_fix <outdir>/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
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
From 6a9af861a09b0082f8e0e1d15eb046881e9abe87 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
From 62bc4977ff293ee2c726b8007744cc78685f499a Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
From a42c2e52b5765a7bcd896e97cecdb18a0f3a2764 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,957 @@
|
||||
From c58638c919ba9acc7624a4e439129e6ccd394152 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
From f7ed2e656d90d9cf203288995be8a8138e1aefd6 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,126 @@
|
||||
From 46b2370779cd9b41d4107b06ca916d1bc4f00345 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
From 5e39153dbfb58d752f93ce708c230db0f3987772 Mon Sep 17 00:00:00 2001
|
||||
From: halt95 <halt95@users.noreply.github.com>
|
||||
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<int>& 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
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user