From b656d05e0503c9a430cea209793b0fe8c74c4cc3 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Thu, 8 Oct 2026 12:02:56 -0700 Subject: [PATCH 1/2] Add a prebuilt DeepGEMM option to vLLM fake-quant serving vllm_serve_fakequant.py takes --deepgemm-url / DEEPGEMM_URL: a PEP 503 index of prebuilt DeepGEMM wheels or a single wheel. From an index, pip selects the newest version with a wheel for this Python and CPU, or the version pinned by --deepgemm-version / DEEPGEMM_VERSION, ignoring any installed deep_gemm; nothing is built. The launcher installs the wheel with pip install --target into a cache folder per version and source (a local wheel by its content), under a lock and with a timeout, checks that import deep_gemm loads it with the installed torch before publishing it (vLLM would silently fall back to its own copy), and puts the folder first on PYTHONPATH, so vLLM and its workers import it instead of vLLM's own DeepGEMM, e.g. a DeepGEMM fork with indexer scorer numerics options. Signed-off-by: Shiyang Chen --- CHANGELOG.rst | 1 + examples/vllm_serve/README.md | 46 ++ examples/vllm_serve/vllm_deepgemm_utils.py | 347 +++++++++++ examples/vllm_serve/vllm_serve_fakequant.py | 4 + .../vllm_serve/test_fakequant_example.py | 26 + .../vllm_serve/test_vllm_deepgemm_utils.py | 567 ++++++++++++++++++ 6 files changed, 991 insertions(+) create mode 100644 examples/vllm_serve/vllm_deepgemm_utils.py create mode 100644 tests/examples/vllm_serve/test_vllm_deepgemm_utils.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5c063cd1c07..31eb89c3123 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -33,6 +33,7 @@ Changelog - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. - Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. +- Add a ``--deepgemm-url`` option to ``examples/vllm_serve/vllm_serve_fakequant.py`` that serves with a prebuilt DeepGEMM wheel, such as a DeepGEMM fork's, from a package index, without changing the Python environment; see `examples/vllm_serve/README.md `_. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - vLLM fake-quant serving now runs on pre-quantized checkpoints such as FP8 when the recipe leaves those layers unquantized (for example a KV-cache-only recipe), and on MLA models with an FP8 KV cache; both previously failed during quantization. - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 8eaea64700d..88845f5118f 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -49,6 +49,8 @@ corresponding environment variable when omitted: | `--modelopt-quant-dataset` | `QUANT_DATASET` | Calibration dataset | | `--modelopt-quant-calib-size` | `QUANT_CALIB_SIZE` | Calibration sample count | | `--modelopt-calib-batch-size` | `CALIB_BATCH_SIZE` | Calibration batch size | +| `--deepgemm-url` | `DEEPGEMM_URL` | Prebuilt DeepGEMM to serve with: a package index with its wheels, or a wheel (see [Serve with a prebuilt DeepGEMM](#serve-with-a-prebuilt-deepgemm)) | +| `--deepgemm-version` | `DEEPGEMM_VERSION` | Version to install from the `DEEPGEMM_URL` package index instead of the newest one for the installed torch and CUDA | `QUANT_CFG` and `KV_QUANT_CFG` (and their CLI flags) will be deprecated in a future release. They still work today. For new runs, use a PTQ recipe through `RECIPE_PATH` or @@ -223,6 +225,50 @@ Tracking needs the client: `pip install nvidia-modelopt[mlflow]` (already in thi `--distributed-executor-backend ray` those are forwarded to the workers along with the tracking settings, since a Ray worker starts with a clean environment. +## Serve with a prebuilt DeepGEMM + +vLLM serves with its own copy of DeepGEMM unless another `deep_gemm` comes first on the Python +path. To serve with a DeepGEMM build instead, such as a DeepGEMM fork that publishes prebuilt wheels +to a package index, pass the index as `--deepgemm-url` or `DEEPGEMM_URL`: + +```bash +DEEPGEMM_URL="https:///simple" \ + python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 +``` + +Before vLLM starts, pip selects the newest `deep_gemm` version on the index that has a wheel for +this Python, CPU architecture and glibc (its tag, such as `manylinux_2_38_x86_64`, names the oldest +glibc it runs on), or the version that `--deepgemm-version` or `DEEPGEMM_VERSION` pins, e.g. +`2.8.0.post12` or exactly `2.8.0.post12+g1a2b3c4.torch2.13.cu130`. The wheel must be built for the +installed torch and CUDA: if the index has builds for others, pin one. The launcher installs the +wheel with `pip install --no-deps --target` into `/wheels/---`, +once per version and source and one launch at a time, checks in a subprocess that `import deep_gemm` +loads it with the installed torch, and puts that folder first on `PYTHONPATH`, so that vLLM's +workers import it instead of vLLM's own copy. Nothing is built, and the Python environment does not +change. `` is `MODELOPT_DEEPGEMM_CACHE_DIR`, else `$VLLM_CACHE_ROOT/modelopt/deepgemm`, else +`~/.cache/modelopt/deepgemm`. + +`DEEPGEMM_URL` also takes a single wheel (a path, or an `https://` or `file://` URL), installed the +same way; a local wheel once per content, so a rebuilt wheel of the same version is installed +again. To serve a DeepGEMM that has no published wheel, such as a local branch, build its wheel +where DeepGEMM can be built and pass the file: + +```bash +pip wheel --no-deps --no-build-isolation -w dist/ +DEEPGEMM_URL=dist/deep_gemm--.whl python vllm_serve_fakequant.py ... +``` + +Give the exact file name: a `*` in `DEEPGEMM_URL=` is not expanded. + +Notes: + +- Every installed version adds a folder of tens of MB to `/wheels`, which can be deleted + when no server uses it. +- pip runs with your pip settings, e.g. `PIP_TRUSTED_HOST` for an `http://` index on another host. +- Ray workers that vLLM does not start from the launcher's environment, such as those of a running + Ray cluster, need the installed folder (printed at launch) on their `PYTHONPATH`, on a filesystem + that all nodes share. + ## Load QAT/PTQ model and serve in vLLM (WIP) Step 1: export the model with bf16 weights and quantizer state. To export the model: diff --git a/examples/vllm_serve/vllm_deepgemm_utils.py b/examples/vllm_serve/vllm_deepgemm_utils.py new file mode 100644 index 00000000000..6c8a794a62d --- /dev/null +++ b/examples/vllm_serve/vllm_deepgemm_utils.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Serve with a custom DeepGEMM: a prebuilt wheel, e.g. of a fork with the indexer scorer numerics. + +vLLM imports a ``deep_gemm`` on the Python path before its vendored copy. ``--deepgemm-url`` (or +``$DEEPGEMM_URL``) names a package index with prebuilt DeepGEMM wheels, or one wheel. From an index, +pip selects the newest version that has a wheel for this Python and machine, or the version that +``--deepgemm-version`` (``$DEEPGEMM_VERSION``) pins. Before vLLM starts its workers, the launcher +installs the wheel with ``pip install --target`` into a folder of its cache, once per version and +source and under a lock, checks in a subprocess that ``import deep_gemm`` loads it with the installed +torch, and puts the folder first on the Python path of the server's processes. Nothing is built, +and the Python environment itself does not change. + +Nothing here imports vLLM, so this can be tested without it. +""" + +import argparse +import contextlib +import fcntl +import hashlib +import importlib +import json +import os +import platform +import shutil +import subprocess +import sys +import tempfile +import time +from collections.abc import Iterator +from pathlib import Path +from urllib.parse import unquote, urlsplit + +import torch +from packaging import tags +from packaging.utils import InvalidWheelFilename, canonicalize_name, parse_wheel_filename +from packaging.version import InvalidVersion, Version + +# TODO: make this generic when another prebuilt kernel package is needed. Only the names below, the +# flags and the import check are DeepGEMM's. It only helps packages that vLLM imports by name from +# the Python path; vLLM builds some, such as FlashMLA, into its own wheel. +DEEPGEMM_URL_ENV = "DEEPGEMM_URL" +DEEPGEMM_VERSION_ENV = "DEEPGEMM_VERSION" +# The root of the cache of installed DeepGEMM wheels (see deepgemm_cache_dir). +DEEPGEMM_CACHE_ENV = "MODELOPT_DEEPGEMM_CACHE_DIR" + +_PROJECT = "deep-gemm" # the normalized name of deep_gemm on an index +# Wheels only: nothing is built, e.g. from the unrelated deep_gemm source distribution on PyPI. An +# installed deep_gemm does not count, or pip would select nothing that it satisfies. +_PIP_OPTIONS = ( + "--no-deps", + "--no-input", + "--disable-pip-version-check", + "--ignore-installed", + "--only-binary", + ":all:", +) +# A pip run reads an index and downloads a wheel of tens of MB, while it can hold the cache lock. +_PIP_TIMEOUT = 1800 # seconds +_IMPORT_TIMEOUT = 600 # seconds +# Installs stage in the cache under its lock; staging folders older than this are left by launches +# that were killed while installing. +_STALE_STAGING = 24 * 3600 # seconds +# Run by a fresh interpreter with the installed folder first on the Python path, like the workers. +_IMPORT_CHECK = "import deep_gemm; print(deep_gemm.__file__)" + + +def add_deepgemm_args(parser: argparse.ArgumentParser) -> None: + """Add ``--deepgemm-url`` and ``--deepgemm-version`` to the launcher's parser.""" + parser.add_argument( + "--deepgemm_url", + "--deepgemm-url", + default=None, + help=( + "Serve with the DeepGEMM wheel that pip selects from this package index " + "(https://.../simple), or with this wheel (a path or URL), installed into a cache " + "folder that is put first on the Python path before vLLM starts. Default: " + f"${DEEPGEMM_URL_ENV}." + ), + ) + parser.add_argument( + "--deepgemm_version", + "--deepgemm-version", + default=None, + help=( + "The version to install from the package index of --deepgemm-url instead of the " + "newest one, e.g. 2.8.0.post210 or, exactly, 2.8.0.post210+g6fec872c.torch2.13.cu130. " + f"Default: ${DEEPGEMM_VERSION_ENV}." + ), + ) + + +def resolve_deepgemm_args(args: argparse.Namespace) -> None: + """Install the DeepGEMM that ``--deepgemm-url`` or ``$DEEPGEMM_URL`` names, if any, and use it.""" + url = args.deepgemm_url or os.environ.get(DEEPGEMM_URL_ENV) + version = args.deepgemm_version or os.environ.get(DEEPGEMM_VERSION_ENV) + if not url: + if version: + raise ValueError( + f"{DEEPGEMM_VERSION_ENV} {version} pins a version of the package index that " + f"{DEEPGEMM_URL_ENV} names, which is not set." + ) + return + folder, installed = install_deepgemm(url, version) + prepend_python_path(folder) + print(f"[deepgemm] serving with deep_gemm {installed} from {url}: {folder}") + + +def install_deepgemm(url: str, version: str | None = None) -> tuple[Path, str]: + """Install the DeepGEMM that ``url`` names into a folder of the cache, unless it is there. + + Args: + url: A package index, or a wheel: a path, or an http(s) or file URL. + version: The version to install from the package index ``url``. + + Returns: + The folder, from which ``import deep_gemm`` loads the installed DeepGEMM with the installed + torch, and the installed version. + """ + name = unquote(urlsplit(url).path) if "://" in url else url + if name.endswith(".whl"): + if version is not None: + raise ValueError( + f"{DEEPGEMM_VERSION_ENV} pins a version of a package index, but {DEEPGEMM_URL_ENV} " + f"{url} is a wheel." + ) + installed, requirement = _wheel_version(Path(name).name), [url] + else: + pin = _parse_version(version) if version is not None else None + # A version with its local label is exact: once installed, the index is not needed. + installed = str(pin) if pin is not None and pin.local else _resolve(url, pin) + requirement = ["--index-url", url, f"deep_gemm=={installed}"] + wheels = deepgemm_cache_dir() / "wheels" + folder = wheels / _folder_name(installed, url) + if not folder.is_dir(): # complete, as published by a rename; a read-only cache works too + with _locked(wheels): + if not folder.is_dir(): + print(f"[deepgemm] installing deep_gemm {installed} from {url}", flush=True) + _pip_install(requirement, folder, url, installed) # checks the import first + return folder, installed + _check_import(folder, installed, url) # again: the environment can change + return folder, installed + + +def prepend_python_path(folder: Path) -> None: + """Put ``folder`` first on the Python path of this process and of the processes it starts.""" + path = str(folder) + if path not in sys.path: + sys.path.insert(0, path) + importlib.invalidate_caches() + python_path = os.environ.get("PYTHONPATH") + if not python_path: + os.environ["PYTHONPATH"] = path + elif python_path.split(os.pathsep)[0] != path: + os.environ["PYTHONPATH"] = f"{path}{os.pathsep}{python_path}" + + +def deepgemm_cache_dir() -> Path: + """The root of the cache of installed DeepGEMM wheels. + + ``$MODELOPT_DEEPGEMM_CACHE_DIR``, else ``modelopt/deepgemm`` under ``$VLLM_CACHE_ROOT``, else + ``~/.cache/modelopt/deepgemm``. + """ + if root := os.environ.get(DEEPGEMM_CACHE_ENV): + return Path(root).expanduser().absolute() + if vllm_root := os.environ.get("VLLM_CACHE_ROOT"): + return Path(vllm_root).expanduser().absolute() / "modelopt" / "deepgemm" + return Path.home() / ".cache" / "modelopt" / "deepgemm" + + +def _pip() -> list[str]: + # subprocess: pip has no supported Python API. Safe: an argument list without a shell, and pip + # installs only wheels, from the index or wheel that the user passes and must trust like any + # package source. + return [sys.executable, "-m", "pip", "install", *_PIP_OPTIONS] + + +def _run_pip(command: list[str], what: str, **kwargs) -> subprocess.CompletedProcess: + """Run the pip ``command`` for ``what`` within ``_PIP_TIMEOUT``.""" + try: + return subprocess.run(command, timeout=_PIP_TIMEOUT, check=False, **kwargs) + except subprocess.TimeoutExpired: + raise RuntimeError(f"pip did not finish {what} in {_PIP_TIMEOUT} s.") from None + + +def _parse_version(version: str) -> Version: + try: + return Version(version) + except InvalidVersion: + raise ValueError(f"{DEEPGEMM_VERSION_ENV} {version!r} is no version.") from None + + +def _wheel_version(name: str) -> str: + """The version of the deep_gemm wheel file ``name``.""" + try: + project, version, _, _ = parse_wheel_filename(name) + except InvalidWheelFilename as error: + raise ValueError(f"{name}: {error}") from None + if canonicalize_name(project) != _PROJECT: + raise ValueError(f"{name} is no deep_gemm wheel.") + return str(version) + + +def _resolve(index: str, pin: Version | None) -> str: + """The version of deep_gemm that pip selects on ``index``: the pinned or the newest one.""" + requirement = f"deep_gemm=={pin}" if pin is not None else "deep_gemm" + # pip prints its own errors, such as an http index on another host without --trusted-host. + result = _run_pip( + [*_pip(), "--dry-run", "--quiet", "--report", "-", "--index-url", index, requirement], + f"selecting deep_gemm on {index}", + stdout=subprocess.PIPE, + text=True, + ) + if result.returncode != 0: + raise RuntimeError( + f"pip finds no deep_gemm wheel for {_python_tag()} on {platform.machine()} on {index}, " + f"see its output above. {DEEPGEMM_URL_ENV} takes a package index (https://.../simple) " + "or a wheel; to serve a DeepGEMM source tree, build its wheel with " + "`pip wheel --no-deps --no-build-isolation -w ` and pass that file." + ) + selected = json.loads(result.stdout)["install"] + if len(selected) != 1: # --no-deps and --ignore-installed: deep_gemm only + raise RuntimeError(f"pip selected {len(selected)} packages for deep_gemm on {index}.") + return selected[0]["metadata"]["version"] + + +def _python_tag() -> str: + return f"{tags.interpreter_name()}{tags.interpreter_version()}" + + +def _folder_name(version: str, url: str) -> str: + """The cache folder of ``version`` from ``url``, which no other source of it shares.""" + return f"{version}-{_python_tag()}-{platform.machine()}-{_source_id(url)}" + + +def _source_id(url: str) -> str: + """A local wheel by its content, as a rebuild can keep its version; else the URL.""" + parts = urlsplit(url) + path = Path(unquote(parts.path) if parts.scheme == "file" else url) + if parts.scheme in ("", "file") and path.is_file(): + source = path.read_bytes() + else: + source = url.rstrip("/").encode() + return hashlib.sha256(source).hexdigest()[:12] + + +def _pip_install(requirement: list[str], folder: Path, url: str, version: str) -> None: + """``pip install --target`` into ``folder``: published complete, and only if it imports. + + Everyone can read it. + """ + staging = Path(tempfile.mkdtemp(prefix=".install-", dir=folder.parent)) + try: + command = [*_pip(), "--target", str(staging), *requirement] + if _run_pip(command, f"installing deep_gemm from {url}").returncode != 0: + raise RuntimeError(f"Installing deep_gemm from {url} failed, see the pip output above.") + for directory, _, files in os.walk(staging): + os.chmod(directory, 0o755) # mkdtemp's 0700 would hide it from other users + for path in (os.path.join(directory, name) for name in files): + if not os.path.islink(path): + os.chmod(path, os.stat(path).st_mode | 0o444) + _check_import(staging, version, url) + try: + staging.rename(folder) + except OSError: + if not folder.is_dir(): # else another process published it first + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def _last_line(text: str) -> str: + return (text.strip().splitlines() or [""])[-1] + + +def _check_import(folder: Path, version: str, url: str) -> None: + """Check that ``import deep_gemm`` loads the one installed in ``folder``, with this torch. + + vLLM would fall back to its own copy without an error. + """ + python_path = os.environ.get("PYTHONPATH") + env = { + **os.environ, + "PYTHONPATH": f"{folder}{os.pathsep}{python_path}" if python_path else str(folder), + } + try: + # subprocess: a fresh interpreter, as a worker starts, runs the fixed _IMPORT_CHECK without + # a shell. + result = subprocess.run( + [sys.executable, "-c", _IMPORT_CHECK], + capture_output=True, + text=True, + env=env, + cwd=folder, + timeout=_IMPORT_TIMEOUT, + check=False, + ) + except subprocess.TimeoutExpired: + raise RuntimeError(f"import deep_gemm {version} from {url} timed out.") from None + if result.returncode != 0: + error = _last_line(result.stderr) or "no error message" + raise RuntimeError( + f"deep_gemm {version} from {url} does not import ({error}): it needs a build for this " + f"torch {torch.__version__} with CUDA {torch.version.cuda} and glibc " + f"{platform.libc_ver()[1] or '?'} or older. Pin one with --deepgemm-version or " + f"{DEEPGEMM_VERSION_ENV}." + ) + loaded = Path(_last_line(result.stdout)).resolve() + if not loaded.is_relative_to(folder.resolve()): + raise RuntimeError( + f"import deep_gemm loads {loaded} instead of deep_gemm {version} from {url}: " + "another deep_gemm comes first on the Python path, such as an editable install. " + "Uninstall it." + ) + + +@contextlib.contextmanager +def _locked(directory: Path) -> Iterator[None]: + """Hold an exclusive lock on the cache ``directory``, which it creates. + + Also removes the old staging folders that launches killed while installing left in it. + """ + directory.mkdir(parents=True, exist_ok=True) + with (directory / ".lock").open("a") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + print(f"[deepgemm] waiting for another launch installing into {directory}", flush=True) + fcntl.flock(lock, fcntl.LOCK_EX) + for staging in directory.glob(".install-*"): + with contextlib.suppress(OSError): + if time.time() - staging.stat().st_mtime > _STALE_STAGING: + shutil.rmtree(staging, ignore_errors=True) + yield diff --git a/examples/vllm_serve/vllm_serve_fakequant.py b/examples/vllm_serve/vllm_serve_fakequant.py index 24e692d04ff..e2c334d9ab8 100644 --- a/examples/vllm_serve/vllm_serve_fakequant.py +++ b/examples/vllm_serve/vllm_serve_fakequant.py @@ -64,6 +64,7 @@ import vllm from packaging import version from vllm.entrypoints.cli.main import main as vllm_main +from vllm_deepgemm_utils import add_deepgemm_args, resolve_deepgemm_args from vllm_mlflow_utils import MLFLOW_ENV_VARS, add_mlflow_args, resolve_mlflow_args @@ -284,6 +285,7 @@ def main(): modelopt_parser = FlexibleArgumentParser(add_help=False) _add_fakequant_args(modelopt_parser) add_mlflow_args(modelopt_parser) + add_deepgemm_args(modelopt_parser) modelopt_args, rest_argv = modelopt_parser.parse_known_args(sys.argv[1:]) rest_argv = _default_to_serve(rest_argv) if rest_argv and rest_argv[0] != "serve": @@ -347,6 +349,8 @@ def main(): if not _has_flag(rest_argv, "--moe-backend"): rest_argv = [*rest_argv, "--moe-backend", "triton"] + # Before vLLM starts the workers, which import deep_gemm from the Python path it sets. + resolve_deepgemm_args(modelopt_args) _run_vllm_cli(rest_argv) diff --git a/tests/examples/vllm_serve/test_fakequant_example.py b/tests/examples/vllm_serve/test_fakequant_example.py index d5976cd9e21..2898798f09f 100644 --- a/tests/examples/vllm_serve/test_fakequant_example.py +++ b/tests/examples/vllm_serve/test_fakequant_example.py @@ -67,6 +67,8 @@ def clean_launcher_env(): "MODELOPT_MLFLOW_REQUIRED", "MODELOPT_MLFLOW_COMMAND", "MODELOPT_MLFLOW_RUN_NAME", + "DEEPGEMM_URL", + "DEEPGEMM_VERSION", ): os.environ.pop(key, None) yield @@ -156,6 +158,9 @@ def _launcher_import_modules(): add_mlflow_args=Mock(), resolve_mlflow_args=Mock(), ), + "vllm_deepgemm_utils": SimpleNamespace( + add_deepgemm_args=Mock(), resolve_deepgemm_args=Mock() + ), "vllm.entrypoints.openai.cli_args": SimpleNamespace( make_arg_parser=entrypoints.legacy_arg_parser ), @@ -204,6 +209,27 @@ def test_vllm_serve_main_disables_compile_cache(monkeypatch, clean_launcher_env) vllm_main.assert_called_once_with() +def test_vllm_serve_main_installs_deepgemm_before_vllm_starts(monkeypatch, clean_launcher_env): + """vLLM's workers import deep_gemm from the Python path that the install sets.""" + launcher = _load_fakequant_launcher(monkeypatch) + monkeypatch.setattr(launcher, "resolve_mlflow_args", Mock()) + vllm_main, _, _ = _stub_launcher_runtime(monkeypatch, launcher) + calls = [] + vllm_main.side_effect = lambda: calls.append("vllm") + monkeypatch.setattr( + launcher, "resolve_deepgemm_args", lambda args: calls.append(args.deepgemm_url) + ) + url = "https://example.com/simple" + monkeypatch.setattr( + sys, "argv", ["vllm_serve_fakequant.py", "/models/qwen", "--deepgemm-url", url] + ) + + launcher.main() + + assert calls == [url, "vllm"] + assert sys.argv == ["vllm", "serve", "/models/qwen"] + + def test_vllm_serve_entrypoint_dependency_error_propagates(monkeypatch): """Do not replace a missing parser dependency with a fallback import error.""" modules, _ = _launcher_import_modules() diff --git a/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py new file mode 100644 index 00000000000..40a796640e0 --- /dev/null +++ b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py @@ -0,0 +1,567 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Installing the DeepGEMM that the vLLM fake-quant server imports, into a folder of a cache. + +``vllm_deepgemm_utils`` imports no vLLM. pip and the import check are stubbed, except in the last +tests, which install tiny pure-Python ``deep_gemm`` wheels with the real pip, from a local package +index or a file, and check the import in a real subprocess. +""" + +import argparse +import base64 +import fcntl +import hashlib +import importlib +import json +import os +import platform +import re +import stat +import subprocess +import sys +import time +import zipfile +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "vllm_serve" + +LABEL = "torch2.13.cu130" +NEWEST = f"2.8.0.post5+gbbbbbbb.{LABEL}" +ANY = "py3-none-any" # the tags of a wheel that every Python installs +INDEX = "https://pypi.example.com/simple" +GLIBC = platform.libc_ver()[1] or "?" +PIP = [sys.executable, "-m", "pip", "install", "--no-deps", "--no-input"] +PIP += ["--disable-pip-version-check", "--ignore-installed", "--only-binary", ":all:"] + + +def _isolated_module(monkeypatch, tmp_path): + """``vllm_deepgemm_utils`` with HOME under ``tmp_path``. + + The variables it reads are unset; ``PYTHONPATH`` and ``sys.path``, which it sets, are restored. + """ + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) # also restores sys.path + module = importlib.import_module("vllm_deepgemm_utils") + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for name in ( + module.DEEPGEMM_URL_ENV, + module.DEEPGEMM_VERSION_ENV, + module.DEEPGEMM_CACHE_ENV, + "VLLM_CACHE_ROOT", + ): + monkeypatch.delenv(name, raising=False) + python_path = os.environ.get("PYTHONPATH") + monkeypatch.setenv("PYTHONPATH", python_path or "") + if python_path is None: + monkeypatch.delenv("PYTHONPATH") + return module + + +@pytest.fixture +def deepgemm(monkeypatch, tmp_path): + return _isolated_module(monkeypatch, tmp_path) + + +def _wheels(): + return Path.home() / ".cache/modelopt/deepgemm/wheels" + + +class _Subprocesses: + """Stubbed pip and import check. + + A pip dry run selects ``selected``, pip installs a fake deep_gemm, and the check loads it. + Records each command with its keyword arguments. + """ + + def __init__(self, deepgemm, monkeypatch, returncode=0): + self.calls, self.returncode, self.import_check = [], returncode, None + self.selected = NEWEST + monkeypatch.setattr(deepgemm.subprocess, "run", self) + + def __call__(self, command, **kwargs): + self.calls.append(SimpleNamespace(command=command, kwargs=kwargs)) + if command[1] == "-c": + folder = Path(kwargs["cwd"]) + assert kwargs["env"]["PYTHONPATH"].split(os.pathsep)[0] == str(folder) + if self.import_check is not None: + return self.import_check(command) + stdout = f"{folder / 'deep_gemm/__init__.py'}\n" + return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + if self.returncode != 0: + return subprocess.CompletedProcess(command, self.returncode, stdout="") + if "--dry-run" in command: + report = {"install": [{"metadata": {"name": "deep_gemm", "version": self.selected}}]} + return subprocess.CompletedProcess(command, 0, stdout=json.dumps(report)) + target = _target(command) + (target / "deep_gemm").mkdir() + (target / "deep_gemm/__init__.py").write_text("__version__ = '2.8.0'\n") + (target / "deep_gemm/_C.so").write_bytes(b"\x7fELF") + (target / "deep_gemm/_C.so").chmod(0o700) + return subprocess.CompletedProcess(command, 0) + + def pip(self, dry_run=False): + """The pip installs, or the pip dry runs.""" + return [ + call + for call in self.calls + if call.command[1:4] == ["-m", "pip", "install"] + and ("--dry-run" in call.command) == dry_run + ] + + +def _target(command): + return Path(command[command.index("--target") + 1]) + + +def test_index_installs_the_version_pip_selects_once(deepgemm, monkeypatch): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + folder, version = deepgemm.install_deepgemm(INDEX) + + assert (folder, version) == (_wheels() / deepgemm._folder_name(NEWEST, INDEX), NEWEST) + source = hashlib.sha256(INDEX.encode()).hexdigest()[:12] + assert folder.name == f"{NEWEST}-{deepgemm._python_tag()}-{os.uname().machine}-{source}" + (dry_run,) = subprocesses.pip(dry_run=True) + assert dry_run.command == [ + *PIP, + *["--dry-run", "--quiet", "--report", "-", "--index-url", INDEX, "deep_gemm"], + ] + # Only the report is captured: pip's own errors reach the terminal. + assert dry_run.kwargs["stdout"] == subprocess.PIPE and "stderr" not in dry_run.kwargs + assert dry_run.kwargs["timeout"] == deepgemm._PIP_TIMEOUT + (install,) = subprocesses.pip() + assert install.kwargs["timeout"] == deepgemm._PIP_TIMEOUT + staging = _target(install.command) + assert staging.parent == _wheels() and not staging.exists() # published by a rename + assert install.command == [ + *PIP, + *["--target", str(staging), "--index-url", INDEX, f"deep_gemm=={NEWEST}"], + ] + assert (folder / "deep_gemm/__init__.py").is_file() + assert subprocesses.calls[-1].command == [sys.executable, "-c", deepgemm._IMPORT_CHECK] + + # Installed: pip selects the version again, and only the import is checked. + subprocesses.calls.clear() + assert deepgemm.install_deepgemm(INDEX) == (folder, version) + assert len(subprocesses.pip(dry_run=True)) == 1 and subprocesses.pip() == [] + + +def test_index_version_can_be_pinned(deepgemm, monkeypatch): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + older = f"2.8.0.post3+gaaaaaaa.{LABEL}" + subprocesses.selected = older + # Without its local label: pip selects the build of that version. + assert deepgemm.install_deepgemm(INDEX, "2.8.0.post3")[1] == older + (dry_run,) = subprocesses.pip(dry_run=True) + assert dry_run.command[-1] == "deep_gemm==2.8.0.post3" + assert subprocesses.pip()[-1].command[-1] == f"deep_gemm=={older}" + + # With it: exact, so pip does not select, and once it is installed pip does not run. + subprocesses.calls.clear() + exact = "2.8.0.post6+gccccccc.torch2.14.cu130" + assert deepgemm.install_deepgemm(INDEX, exact.upper()) == ( # normalized + _wheels() / deepgemm._folder_name(exact, INDEX), + exact, + ) + assert subprocesses.pip(dry_run=True) == [] + assert subprocesses.pip()[-1].command[-1] == f"deep_gemm=={exact}" + subprocesses.calls.clear() + deepgemm.install_deepgemm(INDEX, exact) + assert [call.command[1] for call in subprocesses.calls] == ["-c"] + + +@pytest.mark.parametrize("form", ["path", "file", "https"]) +def test_wheel_is_installed_once_per_version(deepgemm, monkeypatch, tmp_path, form): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl" + url = { + "path": str(wheel), + "file": wheel.as_uri(), + "https": f"https://files.example.com/{wheel.name.replace('+', '%2B')}", + }[form] + folder, version = deepgemm.install_deepgemm(url) + assert (folder, version) == (_wheels() / deepgemm._folder_name(NEWEST, url), NEWEST) + assert subprocesses.pip(dry_run=True) == [] # the file name has the version + (install,) = subprocesses.pip() + assert install.command == [*PIP, "--target", str(_target(install.command)), url] + assert deepgemm.install_deepgemm(url) == (folder, version) + assert len(subprocesses.pip()) == 1 + + +def test_sources_of_a_version_install_separately(deepgemm, monkeypatch, tmp_path): + """A rebuilt wheel or another index with the same version is not served from the cache.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl" + wheel.write_bytes(b"first build") + first, _ = deepgemm.install_deepgemm(str(wheel)) + assert deepgemm.install_deepgemm(wheel.as_uri())[0] == first # the same file + wheel.write_bytes(b"rebuilt with the same version") + rebuilt, _ = deepgemm.install_deepgemm(str(wheel)) + index, _ = deepgemm.install_deepgemm(INDEX) + other_index, _ = deepgemm.install_deepgemm("https://other.example.com/simple") + assert len({first, rebuilt, index, other_index}) == 4 + assert all(folder.is_dir() for folder in (first, rebuilt, index, other_index)) + assert len(subprocesses.pip()) == 4 + # Each source reuses its own folder. + assert deepgemm.install_deepgemm(str(wheel))[0] == rebuilt + assert deepgemm.install_deepgemm(f"{INDEX}/")[0] == index + assert len(subprocesses.pip()) == 4 + + +def test_a_build_that_does_not_import_is_not_cached(deepgemm, monkeypatch): + """The import is checked before the install is published, so a fixed build can be retried.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, 1, stdout="", stderr="ImportError: wrong torch\n" + ) + with pytest.raises(RuntimeError, match=r"does not import \(ImportError: wrong torch\)"): + deepgemm.install_deepgemm(INDEX, NEWEST) + assert [path.name for path in _wheels().iterdir()] == [".lock"] # no folder, no staging + subprocesses.import_check = None # fixed, e.g. the environment + folder, _ = deepgemm.install_deepgemm(INDEX, NEWEST) + assert folder.is_dir() and len(subprocesses.pip()) == 2 + + +def test_pip_timeouts_release_the_cache(deepgemm, monkeypatch): + """A pip run that hangs fails the launch, leaves no staging folder and releases the lock.""" + + def hang(command, **kwargs): + assert kwargs["timeout"] == deepgemm._PIP_TIMEOUT + raise subprocess.TimeoutExpired(command, kwargs["timeout"]) + + monkeypatch.setattr(deepgemm.subprocess, "run", hang) + with pytest.raises(RuntimeError, match=r"pip did not finish selecting deep_gemm on .* in"): + deepgemm.install_deepgemm(INDEX) + with pytest.raises(RuntimeError, match=r"pip did not finish installing deep_gemm from .* in"): + deepgemm.install_deepgemm(INDEX, NEWEST) # under the lock + assert [path.name for path in _wheels().iterdir()] == [".lock"] + with (_wheels() / ".lock").open() as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + + +@pytest.mark.parametrize( + ("url", "version", "match"), + [ + (INDEX, "latest", "DEEPGEMM_VERSION 'latest' is no version"), + (f"https://example.com/deep_gemm-2.8.0-{ANY}.whl", "2.8.0", r"pins .* but .* is a wheel"), + ( + f"https://example.com/other-1.0-{ANY}.whl", + None, + r"other-1.0-py3-none-any.whl is no deep_", + ), + ("dist/deep_gemm.whl", None, r"deep_gemm.whl: Invalid wheel filename"), + ], + ids=["version", "wheel-pin", "other-wheel", "wheel-name"], +) +def test_inputs_are_checked_before_pip_runs(deepgemm, monkeypatch, url, version, match): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + with pytest.raises(ValueError, match=match): + deepgemm.install_deepgemm(url, version) + assert subprocesses.calls == [] + + +def test_pip_failures_are_reported(deepgemm, monkeypatch): + """Nothing is built: a git URL or source tree finds no wheel, with a hint to build one.""" + _Subprocesses(deepgemm, monkeypatch, returncode=1) + git = "git+https://git.example.com/group/deepgemm.git@main" + with pytest.raises(RuntimeError) as error: + deepgemm.install_deepgemm(git) + assert str(error.value).startswith( + f"pip finds no deep_gemm wheel for {deepgemm._python_tag()} on {os.uname().machine} on " + f"{git}, see its output above." + ) + assert "build its wheel with `pip wheel --no-deps" in str(error.value) + with pytest.raises(RuntimeError, match=re.escape(f"from {INDEX} failed, see the pip output")): + deepgemm.install_deepgemm(INDEX, NEWEST) + assert [path.name for path in _wheels().iterdir()] == [".lock"] # no folder, no staging + + +def test_version_needs_a_url(deepgemm): + with pytest.raises(ValueError, match=r"DEEPGEMM_VERSION 2.8.0 pins .* which is not set"): + deepgemm.resolve_deepgemm_args( + argparse.Namespace(deepgemm_url=None, deepgemm_version="2.8.0") + ) + + +def test_import_check_names_what_went_wrong(deepgemm, monkeypatch, tmp_path): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr="Traceback ...\nImportError: _C.so: undefined symbol: _ZN3c10\n", + ) + with pytest.raises(RuntimeError) as error: + deepgemm.install_deepgemm(INDEX) + assert str(error.value).startswith( + f"deep_gemm {NEWEST} from {INDEX} does not import (ImportError: _C.so: undefined symbol: " + "_ZN3c10): it needs a build for this torch " + ) + assert str(error.value).endswith( + f"and glibc {GLIBC} or older. Pin one with --deepgemm-version or DEEPGEMM_VERSION." + ) + + elsewhere = tmp_path / "src/deep_gemm/__init__.py" + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, 0, stdout=f"Failed to load legacy kernels\n{elsewhere}\n", stderr="" + ) + with pytest.raises(RuntimeError, match=re.escape(f"loads {elsewhere} instead of deep_gemm")): + deepgemm.install_deepgemm(INDEX) + + +def test_install_holds_the_cache_lock(deepgemm, monkeypatch): + """Concurrent launches sharing the cache wait while one of them installs.""" + subprocesses, locked = _Subprocesses(deepgemm, monkeypatch), [] + + def run(command, **kwargs): + if command[1:3] == ["-m", "pip"]: + with (_wheels() / ".lock").open("a") as other: + try: + fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(other, fcntl.LOCK_UN) + locked.append(False) + except BlockingIOError: + locked.append(True) + return subprocesses(command, **kwargs) + + _wheels().mkdir(parents=True) + monkeypatch.setattr(deepgemm.subprocess, "run", run) + deepgemm.install_deepgemm(INDEX) + assert locked == [False, True] # the dry run, then the install + + +def test_install_rechecks_after_waiting_for_the_lock(deepgemm, monkeypatch, capsys): + """Another launch installed the same version while this one waited: no second install.""" + subprocesses, flock = _Subprocesses(deepgemm, monkeypatch), fcntl.flock + folder = _wheels() / deepgemm._folder_name(NEWEST, INDEX) + + def wait_for_other_launch(file, operation): + if operation & fcntl.LOCK_NB: + (folder / "deep_gemm").mkdir(parents=True) + raise BlockingIOError + return flock(file, operation) + + monkeypatch.setattr(deepgemm.fcntl, "flock", wait_for_other_launch) + assert deepgemm.install_deepgemm(INDEX) == (folder, NEWEST) + assert subprocesses.pip() == [] + assert "waiting for another launch" in capsys.readouterr().out + + +def test_installed_versions_need_no_lock(deepgemm, monkeypatch, tmp_path): + """A cache that this user cannot write, e.g. shared or read-only, serves what it holds.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = str(tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl") + installed = [deepgemm.install_deepgemm(INDEX), deepgemm.install_deepgemm(wheel)] + + def unwritable(directory): + raise PermissionError(13, "Permission denied", str(directory / ".lock")) + + monkeypatch.setattr(deepgemm, "_locked", unwritable) + subprocesses.calls.clear() + assert [deepgemm.install_deepgemm(INDEX), deepgemm.install_deepgemm(wheel)] == installed + assert subprocesses.pip() == [] + + +def test_install_removes_what_killed_installs_left(deepgemm, monkeypatch): + """Installs stage under the lock; staging folders that old are left by killed launches.""" + _Subprocesses(deepgemm, monkeypatch) + stale, recent = _wheels() / ".install-stale", _wheels() / ".install-recent" + for staging in (stale, recent): + (staging / "deep_gemm").mkdir(parents=True) + day_ago = time.time() - deepgemm._STALE_STAGING - 60 + os.utime(stale, (day_ago, day_ago)) + deepgemm.install_deepgemm(INDEX) + assert not stale.exists() and recent.is_dir() + + +def test_installed_folder_is_readable_by_everyone(deepgemm, monkeypatch): + _Subprocesses(deepgemm, monkeypatch) + umask = os.umask(0o077) + try: + folder, _ = deepgemm.install_deepgemm(INDEX) + finally: + os.umask(umask) + modes = { + path.relative_to(folder).as_posix(): stat.S_IMODE(path.stat().st_mode) + for path in [folder, *folder.rglob("*")] + } + assert modes == { + ".": 0o755, + "deep_gemm": 0o755, + "deep_gemm/__init__.py": 0o644, + "deep_gemm/_C.so": 0o744, + } + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({}, "home/.cache/modelopt/deepgemm"), + ({"VLLM_CACHE_ROOT": "vllm"}, "vllm/modelopt/deepgemm"), + ({"VLLM_CACHE_ROOT": "vllm", "MODELOPT_DEEPGEMM_CACHE_DIR": "mine"}, "mine"), + ], + ids=["home", "vllm-cache-root", "override"], +) +def test_cache_dir_follows_the_environment(deepgemm, monkeypatch, tmp_path, env, expected): + _Subprocesses(deepgemm, monkeypatch) + for name, value in env.items(): + monkeypatch.setenv(name, str(tmp_path / value)) + assert deepgemm.deepgemm_cache_dir() == tmp_path / expected + folder, _ = deepgemm.install_deepgemm(INDEX) + assert folder.parent == tmp_path / expected / "wheels" + + +@pytest.mark.parametrize( + ("argv", "env", "expected"), + [ + (["--deepgemm-url", INDEX, "--deepgemm-version", "1"], {"DEEPGEMM_VERSION": "2"}, "1"), + (["--deepgemm_url", INDEX], {"DEEPGEMM_URL": "/other", "DEEPGEMM_VERSION": "2"}, "2"), + ([], {"DEEPGEMM_URL": INDEX}, None), + ([], {}, None), + ], + ids=["flags", "underscores", "env", "none"], +) +def test_flags_override_env_and_the_install_is_used(deepgemm, monkeypatch, argv, env, expected): + """The installed folder comes first on the Python path of the launcher and its processes.""" + for name, value in env.items(): + monkeypatch.setenv(name, value) + monkeypatch.setenv("PYTHONPATH", "/user/path") + installs, folder = [], Path("/cache/wheels/2.8.0.post5-cp312-x86_64") + + def install(url, version): + installs.append((url, version)) + return folder, "2.8.0.post5" + + monkeypatch.setattr(deepgemm, "install_deepgemm", install) + parser = argparse.ArgumentParser() + deepgemm.add_deepgemm_args(parser) + deepgemm.resolve_deepgemm_args(parser.parse_args(argv)) + if not argv and not env: + assert installs == [] and os.environ["PYTHONPATH"] == "/user/path" + assert str(folder) not in sys.path + return + assert installs == [(INDEX, expected)] + assert sys.path[0] == str(folder) + assert os.environ["PYTHONPATH"] == f"{folder}{os.pathsep}/user/path" + deepgemm.prepend_python_path(folder) # again: no change + assert sys.path.count(str(folder)) == 1 + assert os.environ["PYTHONPATH"] == f"{folder}{os.pathsep}/user/path" + + +def test_python_path_changes_do_not_outlive_the_test(tmp_path): + """The fixture's monkeypatch undoes the launcher's Python path changes.""" + environ, path = dict(os.environ), list(sys.path) + with pytest.MonkeyPatch.context() as monkeypatch: + deepgemm = _isolated_module(monkeypatch, tmp_path) + monkeypatch.setattr(deepgemm, "install_deepgemm", lambda url, version: (tmp_path, "2.8.0")) + deepgemm.resolve_deepgemm_args(argparse.Namespace(deepgemm_url="/x", deepgemm_version=None)) + assert sys.path[0] == str(tmp_path) + assert dict(os.environ) == environ and sys.path == path + + +# The real pip installs tiny pure-Python deep_gemm wheels from a file:// index or a file. + + +def _make_wheel(directory, version, tag=ANY): + """A pure-Python ``deep_gemm`` wheel of ``version`` whose package records its version.""" + dist_info = f"deep_gemm-{version}.dist-info" + files = { + "deep_gemm/__init__.py": f"__version__ = {version!r}\n", + f"{dist_info}/METADATA": f"Metadata-Version: 2.1\nName: deep_gemm\nVersion: {version}\n", + f"{dist_info}/WHEEL": ( + f"Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: {tag}\n" + ), + } + record = [] + for name, text in files.items(): + digest = base64.urlsafe_b64encode(hashlib.sha256(text.encode()).digest()).rstrip(b"=") + record.append(f"{name},sha256={digest.decode()},{len(text.encode())}") + files[f"{dist_info}/RECORD"] = "\n".join([*record, f"{dist_info}/RECORD,,", ""]) + path = directory / f"deep_gemm-{version}-{tag}.whl" + with zipfile.ZipFile(path, "w") as wheel: + for name, text in files.items(): + wheel.writestr(name, text) + return path + + +@pytest.fixture +def real_pip(deepgemm, monkeypatch, tmp_path): + """pip without the user's configuration; a local PEP 503 index of tiny wheels, its URL.""" + for name in [name for name in os.environ if name.startswith("PIP_")]: + monkeypatch.delenv(name) + monkeypatch.setenv("PIP_CONFIG_FILE", os.devnull) + files = tmp_path / "files" + files.mkdir() + wheels = [ + _make_wheel(files, f"2.8.0.post3+gaaaaaaa.{LABEL}"), + _make_wheel(files, NEWEST), + _make_wheel(files, f"2.8.0.post7+gddddddd.{LABEL}", tag="py3-none-win_amd64"), + ] + page = tmp_path / "simple" / "deep-gemm" + page.mkdir(parents=True) + links = "".join(f'{wheel.name}
\n' for wheel in wheels) + (page / "index.html").write_text(f"\n\n{links}\n") + return (tmp_path / "simple").as_uri() + + +def test_real_pip_installs_the_newest_wheel_for_this_machine(deepgemm, real_pip): + folder, version = deepgemm.install_deepgemm(real_pip) + assert version == NEWEST # not post7, which is for Windows + assert (folder / "deep_gemm/__init__.py").read_text() == f"__version__ = {NEWEST!r}\n" + assert (folder / f"deep_gemm-{NEWEST}.dist-info/METADATA").is_file() + # The import check ran in a fresh interpreter: an install that does not import fails it. + (folder / "deep_gemm/__init__.py").write_text("raise ImportError('wrong torch')\n") + with pytest.raises(RuntimeError, match=r"does not import \(ImportError: wrong torch\)"): + deepgemm.install_deepgemm(real_pip) + + +def test_real_pip_installs_a_pinned_version_and_the_launcher_uses_it(deepgemm, real_pip, capsys): + older = f"2.8.0.post3+gaaaaaaa.{LABEL}" + deepgemm.resolve_deepgemm_args( + argparse.Namespace(deepgemm_url=real_pip, deepgemm_version="2.8.0.post3") + ) + folder = _wheels() / deepgemm._folder_name(older, real_pip) + assert sys.path[0] == str(folder) and os.environ["PYTHONPATH"].startswith(str(folder)) + result = subprocess.run( + [sys.executable, "-c", "import deep_gemm; print(deep_gemm.__version__)"], + capture_output=True, + text=True, + check=True, + cwd=Path.home().parent, + ) + assert result.stdout == f"{older}\n" + assert f"serving with deep_gemm {older} from {real_pip}: {folder}" in capsys.readouterr().out + + +def test_real_pip_ignores_an_installed_deep_gemm(deepgemm, real_pip, monkeypatch, tmp_path): + """An installed deep_gemm, e.g. in site-packages, neither satisfies nor pins the selection.""" + installed = tmp_path / "site" + with zipfile.ZipFile(_make_wheel(tmp_path, f"2.8.0.post3+gaaaaaaa.{LABEL}")) as wheel: + wheel.extractall(installed) + monkeypatch.setenv("PYTHONPATH", str(installed)) + assert deepgemm.install_deepgemm(real_pip)[1] == NEWEST + folder, version = deepgemm.install_deepgemm(real_pip, "2.8.0.post3") # the installed version + assert version == f"2.8.0.post3+gaaaaaaa.{LABEL}" + assert (folder / "deep_gemm/__init__.py").is_file() + + +def test_real_pip_installs_a_wheel_file(deepgemm, real_pip, tmp_path): + wheel = _make_wheel(tmp_path, NEWEST) + for url in (str(wheel), wheel.as_uri()): + folder, version = deepgemm.install_deepgemm(url) + assert version == NEWEST + assert (folder / "deep_gemm/__init__.py").read_text() == f"__version__ = {NEWEST!r}\n" From 6d317f35da418528eda251c26452f6f4c081996d Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Thu, 8 Oct 2026 15:22:39 -0700 Subject: [PATCH 2/2] Install a DeepGEMM source tree with pip instead of DEEPGEMM_URL Serving a DeepGEMM source tree, such as a local branch, no longer builds a wheel to pass through DEEPGEMM_URL: the README and the launcher's error hint now say to pip install the source into the serving environment, which vLLM then uses directly. DEEPGEMM_URL still takes prebuilt wheels from a package index or a file. Signed-off-by: Shiyang Chen --- examples/vllm_serve/README.md | 13 +++++++------ examples/vllm_serve/vllm_deepgemm_utils.py | 4 ++-- .../examples/vllm_serve/test_vllm_deepgemm_utils.py | 4 ++-- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 88845f5118f..299ad9950d0 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -250,16 +250,17 @@ change. `` is `MODELOPT_DEEPGEMM_CACHE_DIR`, else `$VLLM_CACHE_ROOT/model `DEEPGEMM_URL` also takes a single wheel (a path, or an `https://` or `file://` URL), installed the same way; a local wheel once per content, so a rebuilt wheel of the same version is installed -again. To serve a DeepGEMM that has no published wheel, such as a local branch, build its wheel -where DeepGEMM can be built and pass the file: +again. + +To serve a DeepGEMM source tree, such as a local branch, install it into the serving environment +instead; vLLM then uses it, without `DEEPGEMM_URL`, for every serve from that environment until it +is uninstalled: ```bash -pip wheel --no-deps --no-build-isolation -w dist/ -DEEPGEMM_URL=dist/deep_gemm--.whl python vllm_serve_fakequant.py ... +pip install --no-deps --no-build-isolation +python vllm_serve_fakequant.py ... ``` -Give the exact file name: a `*` in `DEEPGEMM_URL=` is not expanded. - Notes: - Every installed version adds a folder of tens of MB to `/wheels`, which can be deleted diff --git a/examples/vllm_serve/vllm_deepgemm_utils.py b/examples/vllm_serve/vllm_deepgemm_utils.py index 6c8a794a62d..5937eb50b6d 100644 --- a/examples/vllm_serve/vllm_deepgemm_utils.py +++ b/examples/vllm_serve/vllm_deepgemm_utils.py @@ -228,8 +228,8 @@ def _resolve(index: str, pin: Version | None) -> str: raise RuntimeError( f"pip finds no deep_gemm wheel for {_python_tag()} on {platform.machine()} on {index}, " f"see its output above. {DEEPGEMM_URL_ENV} takes a package index (https://.../simple) " - "or a wheel; to serve a DeepGEMM source tree, build its wheel with " - "`pip wheel --no-deps --no-build-isolation -w ` and pass that file." + "or a wheel; to serve a DeepGEMM source tree, install it into the environment with " + "`pip install --no-deps --no-build-isolation ` instead." ) selected = json.loads(result.stdout)["install"] if len(selected) != 1: # --no-deps and --ignore-installed: deep_gemm only diff --git a/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py index 40a796640e0..c6d2e826097 100644 --- a/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py +++ b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py @@ -275,7 +275,7 @@ def test_inputs_are_checked_before_pip_runs(deepgemm, monkeypatch, url, version, def test_pip_failures_are_reported(deepgemm, monkeypatch): - """Nothing is built: a git URL or source tree finds no wheel, with a hint to build one.""" + """Nothing is built: a git URL or source tree finds no wheel, with a hint to install it.""" _Subprocesses(deepgemm, monkeypatch, returncode=1) git = "git+https://git.example.com/group/deepgemm.git@main" with pytest.raises(RuntimeError) as error: @@ -284,7 +284,7 @@ def test_pip_failures_are_reported(deepgemm, monkeypatch): f"pip finds no deep_gemm wheel for {deepgemm._python_tag()} on {os.uname().machine} on " f"{git}, see its output above." ) - assert "build its wheel with `pip wheel --no-deps" in str(error.value) + assert "install it into the environment with `pip install --no-deps" in str(error.value) with pytest.raises(RuntimeError, match=re.escape(f"from {INDEX} failed, see the pip output")): deepgemm.install_deepgemm(INDEX, NEWEST) assert [path.name for path in _wheels().iterdir()] == [".lock"] # no folder, no staging