diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 65d8008a..6cf73fab 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -77,8 +77,16 @@ jobs: -v -s --tb=short -x \ --cov=granite_switch --cov-report=xml + # COVERAGE REPORTING MUST NOT GATE A PR. `fail_ci_if_error: false` already + # declared that intent, but v4 can fail the step *before* it consults that + # flag -- it exits nonzero on a token/rate-limit problem of its own, which + # is what turned a fully green CPU suite red (the "Run CPU tests" step + # passed on both 3.11 and 3.12; only this step failed). `continue-on-error` + # enforces the intent at the workflow level, where the action cannot + # override it. v5 is the supported line; v4 is deprecated. - name: Upload coverage - uses: codecov/codecov-action@v4 + uses: codecov/codecov-action@v5 + continue-on-error: true with: token: ${{ secrets.CODECOV_TOKEN }} files: coverage.xml diff --git a/CLAUDE.md b/CLAUDE.md index 66fe77d0..0e3481ca 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -471,6 +471,7 @@ This repo uses [pre-commit](https://pre-commit.com/) (ruff, hygiene hooks, SPDX- ## Documentation +- `CONTRIBUTING.md` - Contribution guidelines, uv/dev setup, and DCO sign-off - `docs/GIT_WORKFLOW.md` - Git branching strategy and commit guidelines - `docs/SUPPORTED_MODELS.md` - Model compatibility - `docs/SR_ARCHITECTURE.md` - Shadow Residual dual-stream architecture diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 969eb17c..7d3aa066 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,6 +12,18 @@ curl -LsSf https://astral.sh/uv/install.sh | sh Or via pip: `pip install uv` +> **Keep `uv.lock` at `revision = 3`.** The `revision` field at the top of `uv.lock` is the lockfile +> *format* version, bumped automatically by newer uv releases (uv 0.12.x writes `revision = 5`). It is +> **not** a dependency change, and we pin it at `3` so the committed lockfile stays readable by the uv +> versions the team and CI use. If `git diff uv.lock` shows only `revision = 3` → `revision = 5` (or +> similar) after a `uv sync`/`uv lock`, that bump is spurious — revert it before committing: +> ```bash +> git checkout uv.lock # if revision is the only change +> ``` +> If you also changed dependencies, re-lock and then reset just the first line back to `revision = 3` +> by hand (edit `uv.lock` line 2) so only the real dependency changes land. Do not commit a bare +> `revision` bump. + ## Getting Started 1. Fork the repository diff --git a/README.md b/README.md index 8319dc87..bb696d66 100644 --- a/README.md +++ b/README.md @@ -47,16 +47,18 @@ Other install options depending on your use case: ```bash pip install "granite-switch[compose]" # Compose modular models pip install "granite-switch[hf]" # HuggingFace inference -pip install "granite-switch[vllm27]" # newer vLLM line (0.27.x) pip install "granite-switch[dev]" # Everything ``` Requires Python 3.11+ and PyTorch 2.11+. -> **vLLM version note:** This project requires `transformers>=5.16`, so it pins vLLM to the -> `0.26.x` line (the earliest vLLM whose model code handles the transformers-5.13+ -> `full_attention` layer-type rename) via the default `[vllm]` extra. `[vllm27]` selects the -> newer `0.27.x` line. Both require PyTorch 2.11+ (CUDA 13+). +> **vLLM version note:** `[vllm]` accepts any line from `0.26.x` to `0.30.x`. The floor is +> `transformers>=5.16`: `0.26` is the earliest vLLM whose model code handles the +> transformers-5.13+ `full_attention` layer-type rename. One tree serves all five lines — +> where their APIs differ, the backend checks for the capability rather than the version. All +> of them require PyTorch 2.11+ (CUDA 13+). To pin a single line for development, use the +> `vllm26` / `vllm30` dependency *groups* — the two ends of the range, which is also what CI +> tests on every GPU run. ### Compose a Model diff --git a/docs/AUDIO.md b/docs/AUDIO.md index d6205325..4cca72bb 100644 --- a/docs/AUDIO.md +++ b/docs/AUDIO.md @@ -27,7 +27,7 @@ checkpoint at any `--asr-model` your installed transformers supports (for exampl ```bash # Serving an audio-enabled checkpoint -uv sync --extra vllm --extra audio # or --extra vllm27 --extra audio +uv sync --extra vllm --extra audio # any vLLM 0.26.x - 0.30.x # Development / running the test suite (the dev groups include audio already) uv sync --group dev # vLLM 0.26.x diff --git a/pyproject.toml b/pyproject.toml index 40cb81c2..f967d83e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,8 +20,10 @@ Documentation = "https://github.com/generative-computing/granite-switch/tree/mai [project.optional-dependencies] hf = ["accelerate>=0.20.0"] -vllm = ["vllm>=0.26.0,<0.27.0"] -vllm27 = ["vllm>=0.27.0,<0.28.0"] +# Every line from 0.26 through 0.30 is supported by one tree; the backend uses +# capability checks (hasattr) rather than version comparisons where the two +# differ. Pin a specific line with the per-line dependency groups below. +vllm = ["vllm>=0.26.0,<0.31.0"] compose = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"] build = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"] # Backward compatibility alias for compose # Audio (ASR) decode + resample. Reuse vLLM's own audio deps (unversioned, so it @@ -56,40 +58,72 @@ markers = [ ] [dependency-groups] +# vllm26 AND vllm30 ARE THE TWO CI LEGS -- the ends of the supported range, which +# is why these two and not two others. The `vllm` extra above claims 0.26-0.30, +# and the two bounds are where a break actually lands: a deleted API shows up at +# the ceiling and an API that does not exist yet shows up at the floor. +# 0.27/0.28/0.29 sit strictly between two tested lines, so they add install time +# without adding a failure mode the ends would miss. All five lines were verified +# together out of band (five-leg GPU matrix, every suite green); these two are +# what CI re-checks on every run. vllm26 = ["vllm>=0.26.0,<0.27.0"] +vllm30 = ["vllm>=0.30.0,<0.31.0"] +# vllm27 IS NOT A CI LEG -- it exists so the GPU workflow stays dispatchable +# during this transition. `workflow_dispatch` always executes the workflow file +# from the DEFAULT BRANCH, so until the 0.26/0.30 matrix lands on main, a +# `/gpu-test` run uses main's matrix (legs vllm26 -> dev and vllm27 -> +# dev-vllm27) while checking out the PR's commit. Deleting dev-vllm27 here made +# that leg unresolvable against this tree -- `uv sync --group dev-vllm27` has no +# such group -- and the leg died before pytest started. Keeping the group costs +# nothing and is not a fiction: the `vllm` extra claims 0.26-0.30, so 0.27 is a +# supported line that simply isn't worth its own per-run leg. Removable once the +# matrix rename is on main. vllm27 = ["vllm>=0.27.0,<0.28.0"] # `audio` is included so the audio tests can actually run: the ASR path needs # vLLM's audio deps (av/soundfile/resampy) at runtime, and no group pulled them # in before (integration tests failed with ModuleNotFoundError on a synced pod). dev = ["pytest", "pytest-cov", { include-group = "vllm26" }, "granite-switch[hf,compose,audio]"] dev-vllm27 = ["pytest", "pytest-cov", { include-group = "vllm27" }, "granite-switch[hf,compose,audio]"] -# The newest vLLM line, as an alias. The GPU workflow tests `dev` and `dev-latest` -# (one leg if both pin the same minor), so a vLLM bump is repointing this line. -dev-latest = [{ include-group = "dev-vllm27" }] +dev-vllm30 = ["pytest", "pytest-cov", { include-group = "vllm30" }, "granite-switch[hf,compose,audio]"] +# The newest vLLM line, as an alias. The GPU workflow (derive_legs.py) and +# CLAUDE.md ask for `dev-latest` rather than the concrete `dev-vllmNN`, so moving +# the ceiling is repointing this one line instead of chasing the version through +# every caller. Currently the 0.30 line. +dev-latest = [{ include-group = "dev-vllm30" }] test = ["pytest", "pytest-cov", "bitsandbytes", "optimum-quanto", { include-group = "dev" }] [tool.uv] default-groups = ["vllm26"] +# Only the per-line GROUPS are mutually exclusive now. The `vllm` extra spans +# 0.26-0.30, so it overlaps every group and no longer conflicts with any of them -- +# which is what made `uv sync` unsatisfiable while the extra said 0.30-only and the +# default group said 0.26. +# Every pair that would pull two different vLLM lines into one venv. Three lines +# are declared (0.26 / 0.27 / 0.30) with a bare and a dev- flavour each, so the +# 26-side {vllm26, dev}, the 27-side {vllm27, dev-vllm27} and the 30-side +# {vllm30, dev-vllm30} give 3 x (2 x 2) = 12 cross-line pairs. `test` is absent +# deliberately: it reaches vllm26 through `include-group = "dev"` and uv derives +# that transitively -- main's two-line version listed it nowhere either. conflicts = [ - # group-vs-group + # 0.26 vs 0.27 [{ group = "vllm26" }, { group = "vllm27" }], + [{ group = "vllm26" }, { group = "dev-vllm27" }], [{ group = "dev" }, { group = "vllm27" }], [{ group = "dev" }, { group = "dev-vllm27" }], - [{ group = "dev-vllm27" }, { group = "vllm26" }], + # 0.26 vs 0.30 (dev-latest aliases dev-vllm30, so it conflicts with the 0.26 side too) + [{ group = "vllm26" }, { group = "vllm30" }], + [{ group = "vllm26" }, { group = "dev-vllm30" }], + [{ group = "vllm26" }, { group = "dev-latest" }], + [{ group = "dev" }, { group = "vllm30" }], + [{ group = "dev" }, { group = "dev-vllm30" }], [{ group = "dev" }, { group = "dev-latest" }], - [{ group = "dev-latest" }, { group = "vllm26" }], - # group-vs-extra - [{ group = "vllm26" }, { extra = "vllm27" }], - [{ group = "vllm27" }, { extra = "vllm" }], - [{ group = "vllm27" }, { extra = "tutorials" }], - [{ group = "dev" }, { extra = "vllm27" }], - [{ group = "dev-vllm27" }, { extra = "vllm" }], - [{ group = "dev-vllm27" }, { extra = "tutorials" }], - [{ group = "dev-latest" }, { extra = "vllm" }], - [{ group = "dev-latest" }, { extra = "tutorials" }], - # extra-vs-extra - [{ extra = "vllm" }, { extra = "vllm27" }], - [{ extra = "tutorials" }, { extra = "vllm27" }], + # 0.27 vs 0.30 + [{ group = "vllm27" }, { group = "vllm30" }], + [{ group = "vllm27" }, { group = "dev-vllm30" }], + [{ group = "vllm27" }, { group = "dev-latest" }], + [{ group = "dev-vllm27" }, { group = "vllm30" }], + [{ group = "dev-vllm27" }, { group = "dev-vllm30" }], + [{ group = "dev-vllm27" }, { group = "dev-latest" }], ] [tool.coverage.run] diff --git a/src/granite_switch/vllm/__init__.py b/src/granite_switch/vllm/__init__.py index 56615750..096c8fb6 100644 --- a/src/granite_switch/vllm/__init__.py +++ b/src/granite_switch/vllm/__init__.py @@ -43,6 +43,25 @@ def register(): """ from vllm import ModelRegistry + # FlashAttention 3 computes wrong attention when captured inside a FULL CUDA + # graph on this vLLM line, so for a granite_switch model we report the FA2 + # cudagraph-support level (UNIFORM_BATCH), which downgrades FULL to + # FULL_AND_PIECEWISE (attention runs eagerly, output matches the reference). + # See fa3_schedule and issue #139. + from .fa3_schedule import patch_flash_attn_schedule + + patch_flash_attn_schedule() + + # Separately, size FA3's ahead-of-time schedule to the switch's own attention + # layers (SR's doubled Q-heads, the counting/memory heads). The downgrade + # above only acts under FULL cudagraphs; under enforce_eager FA3 still crashes + # with "scheduler_metadata must have shape (metadata_size)" unless the AOT + # schedule is sized correctly. The two patches touch different seams (one + # wraps get_cudagraph_support, the other __init__) and compose. See issue #139. + from .fa3_aot_schedule import patch_flash_attn_aot_schedule + + patch_flash_attn_aot_schedule() + # Register config with transformers AutoConfig try: from transformers import AutoConfig diff --git a/src/granite_switch/vllm/audio/asr.py b/src/granite_switch/vllm/audio/asr.py index 1314d00b..88eab82b 100644 --- a/src/granite_switch/vllm/audio/asr.py +++ b/src/granite_switch/vllm/audio/asr.py @@ -12,6 +12,8 @@ from __future__ import annotations +import os +import subprocess import threading from collections.abc import Mapping from typing import Any, Union @@ -58,6 +60,66 @@ } +def gpu_is_exclusive_process() -> bool: + """True if any visible GPU is in the ``Exclusive_Process`` compute mode. + + Queried through ``nvidia-smi`` on purpose: it must NOT create a CUDA context + in the calling process. That is the whole point of the caller + (:func:`resolve_driver_device`) -- under ``Exclusive_Process`` only one + process may hold a device's primary context, and in vLLM serving the audio + multimodal preprocessing runs in the *driver* process while the ``EngineCore`` + subprocess already owns the device. A ``torch.cuda`` probe here would create + exactly the parent context we are trying to avoid. + + Honours ``CUDA_VISIBLE_DEVICES`` so a partitioned host is read correctly, and + is conservative: any error (no ``nvidia-smi``, no GPU, parse failure) returns + False, leaving the requested device unchanged. + """ + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=index,compute_mode", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ).stdout + except Exception: + return False + + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + allowed = ( + {v.strip() for v in visible.split(",") if v.strip() != ""} + if visible is not None + else None + ) + for line in out.splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) != 2: + continue + index, mode = parts + if allowed is not None and index not in allowed: + continue + if mode == "Exclusive_Process": + return True + return False + + +def resolve_driver_device(device: str) -> str: + """Device for an ASR transcriber that runs in the vLLM *driver* process. + + ``"cuda"`` is downgraded to ``"cpu"`` when a visible GPU is in + ``Exclusive_Process`` mode: there the driver cannot acquire a second context + on the device the ``EngineCore`` subprocess already holds, so loading the ASR + model onto CUDA in the driver dies with ``cudaErrorDevicesUnavailable`` (and + leaks a partial parent context). The ASR encoder is small and runs on short + clips, so CPU in the driver is the correct, working fallback. Any explicit + non-``"cuda"`` device, and the normal ``Default`` compute mode, pass through + unchanged. + """ + if device == "cuda" and gpu_is_exclusive_process(): + return "cpu" + return device + + def _resolve_torch_dtype(dtype: str | None, device: str) -> Any: """Resolve an ``asr_dtype`` name to a ``torch.dtype``. diff --git a/src/granite_switch/vllm/audio/processor.py b/src/granite_switch/vllm/audio/processor.py index a47c81e0..8cb16738 100644 --- a/src/granite_switch/vllm/audio/processor.py +++ b/src/granite_switch/vllm/audio/processor.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """vLLM multimodal processor for the audio cascade. -``_call_hf_processor`` runs ASR and tokenizes the transcript; -``_get_prompt_updates`` then replaces the ``<|audio|>`` marker with the real -transcript token ids via ``PromptReplacement``, so the scheduler sizes KV for the -runtime-determined length rather than a fixed audio window. +ASR runs in the HF-processor hook (``_call_hf_processor`` on vLLM 0.26 - 0.28, +``_apply_hf_processor_main`` on 0.29+, which deleted the former) and tokenizes the +transcript; ``_get_prompt_updates`` then replaces the ``<|audio|>`` marker with the +real transcript token ids via ``PromptReplacement``, so the scheduler sizes KV for +the runtime-determined length rather than a fixed audio window. Modeled on vLLM 0.19.1's ``ultravox.py``. Audio is answered by the base model — no adapter control tokens are placed, so the switch is not involved. @@ -37,6 +38,7 @@ DEFAULT_ALLOWED_REQUEST_GENERATE_KEYS, DEFAULT_ASR_MODEL_ID, get_transcriber, + resolve_driver_device, resolve_generate_kwargs, ) @@ -89,7 +91,13 @@ def _asr_model_id(self) -> str: def _asr_device(self) -> str: cfg = self.get_hf_config() - return getattr(cfg, "asr_device", "cuda") or "cuda" + requested = getattr(cfg, "asr_device", "cuda") or "cuda" + # The transcriber runs in the vLLM *driver* process (multimodal input + # preprocessing), not the EngineCore subprocess. Under Exclusive_Process + # the driver cannot share the device the engine already holds, so a + # "cuda" request is downgraded to CPU there to avoid cudaErrorDevices- + # Unavailable. Default compute mode and explicit devices pass through. + return resolve_driver_device(requested) def _asr_dtype(self) -> str | None: cfg = self.get_hf_config() @@ -338,19 +346,17 @@ def _reject_reserved_ids(self, ids: Sequence[int]) -> None: f"tokens must not originate from audio content." ) - def _call_hf_processor( + def _audio_features( self, - prompt: str, - mm_data: Mapping[str, object], + audios: Sequence[object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - tokenizer = self.info.get_tokenizer() - audios = mm_data.get("audios", []) or [] + ) -> dict[str, torch.Tensor]: + """Transcribe every clip into the out-of-band audio tensors. + Empty when there is no audio, so the caller can splat it unconditionally. + """ if not audios: - input_ids = tokenizer.encode(prompt, add_special_tokens=False) - return BatchFeature(dict(input_ids=[input_ids]), tensor_type="pt") + return {} # Resolved once, then applied to every audio item in this request. generate_kwargs = resolve_generate_kwargs( @@ -359,19 +365,72 @@ def _call_hf_processor( DEFAULT_ALLOWED_REQUEST_GENERATE_KEYS, ) - input_ids = tokenizer.encode(prompt, add_special_tokens=False) - # Concatenated flat, with per-item sizes to split them back. per_item_ids = [self._transcribe(a, generate_kwargs) for a in audios] sizes = [len(ids) for ids in per_item_ids] flat_ids = [tid for ids in per_item_ids for tid in ids] + return dict( + audio_token_ids=torch.tensor(flat_ids, dtype=torch.long), + audio_num_tokens=torch.tensor(sizes, dtype=torch.long), + ) + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + """vLLM 0.26 - 0.28 only; the base class dropped this hook at 0.29. + + Here the processor owns both halves of the job, so it returns the + tokenized prompt alongside the audio tensors. See + ``_apply_hf_processor_main`` for the 0.29+ split. + """ + tokenizer = self.info.get_tokenizer() + input_ids = tokenizer.encode(prompt, add_special_tokens=False) + features = self._audio_features(mm_data.get("audios", []) or [], mm_kwargs) return BatchFeature( - dict( - input_ids=[input_ids], - audio_token_ids=torch.tensor(flat_ids, dtype=torch.long), - audio_num_tokens=torch.tensor(sizes, dtype=torch.long), - ), + dict(input_ids=[input_ids], **features), + tensor_type="pt", + ) + + def _apply_hf_processor_main(self, **kwargs): + """Produce the audio tensors on whichever hook this vLLM actually calls. + + vLLM 0.29 deleted ``BaseMultiModalProcessor._call_hf_processor`` (and + ``_apply_hf_processor_text_mm``, which was its only caller). Text and + multi-modal data are now processed independently: ``apply()`` tokenizes the + prompt itself via ``_postprocess_prompt`` and ``_apply_hf_processor_main`` + was narrowed to ``(mm_items, hf_processor_mm_kwargs) -> BatchFeature``, + returning the MM tensors ONLY. Left alone, our ``_call_hf_processor`` above + is simply never called on 0.29+ and no transcript is ever produced. + + The two signatures are mutually exclusive, so this dispatches on whether the + base class still has ``_call_hf_processor``: that is the hook whose removal + is the actual problem, and it was removed in the same commit that narrowed + this signature. Both vLLM call sites pass keywords only (0.28 adds + ``prompt`` / ``tokenization_kwargs`` / ``enable_hf_prompt_update``), hence + ``**kwargs`` rather than a signature that could only match one line. + """ + if hasattr(super(), "_call_hf_processor"): # 0.26 - 0.28 + return super()._apply_hf_processor_main(**kwargs) + + mm_items = kwargs["mm_items"] + hf_processor_mm_kwargs = kwargs["hf_processor_mm_kwargs"] + + valid_mm_items = mm_items.select( + {k for k, c in mm_items.get_all_counts().items() if c > 0} + ) + processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + features = self._audio_features( + processor_data.get("audios", []) or [], + hf_processor_mm_kwargs, + ) + # No input_ids: on 0.29+ the prompt is vLLM's business, not ours. + return BatchFeature( + dict(**features, **passthrough_data), tensor_type="pt", ) @@ -397,6 +456,11 @@ def _hf_processor_applies_updates( Only the uncached path consults this hook; the cached path already hardcodes False, which is why audio works with the default ``mm_processor_cache_gb=4`` and breaks under ``--mm-processor-cache-gb 0``. + + Required on vLLM 0.26 - 0.28 and inert on 0.29+, where the base class + deleted the hook along with the ``is_update_applied`` branch it fed: + ``_apply_prompt_updates`` now always runs, which is exactly the False + behaviour asked for here. Kept because one tree has to serve both. """ return False @@ -428,10 +492,31 @@ def replacement(item_idx: int): e = int(starts[item_idx + 1]) return [int(t) for t in all_ids[s:e]] + # TARGET IS A TOKEN-ID LIST, NOT THE MARKER STRING. vLLM narrowed the + # accepted target type at 0.29: + # + # 0.26 - 0.28 UpdateTarget = PromptSeq | PromptIndex # str allowed + # 0.29 - 0.30 UpdateTarget = list[int] | PromptIndex # str dropped + # + # A ``list[int]`` is legal on every one of those lines, so this needs no + # version branch -- only ``str`` stopped being legal. + # + # Passing the string on 0.29+ does not fail loudly where it is given. + # Token matching simply finds nothing, which sends vLLM into + # ``_apply_prompt_updates_via_text`` -- a fallback that only exists from + # 0.29 -- and that calls ``tokenizer.decode(target)`` on the raw string. + # The resulting error surfaces far away and looks like a tokenizer bug: + # ``TypeError: Can't extract `str` to `Vec``` from the Rust tokenizer, + # killing EngineCore during init rather than rejecting the request. + # + # ``_marker_id()`` cannot answer -1 here: ``_validate_marker_count`` + # already requires one marker per audio item, and ``_count_markers`` + # reports 0 for an unregistered marker, so a prompt that reached this + # point with audio items has a registered marker. return [ PromptReplacement( modality="audio", - target=AUDIO_MARKER, + target=[self._marker_id()], replacement=replacement, ) ] diff --git a/src/granite_switch/vllm/decoder/shadow_residual/decoder.py b/src/granite_switch/vllm/decoder/shadow_residual/decoder.py index 68fd7b4f..7703b5e5 100644 --- a/src/granite_switch/vllm/decoder/shadow_residual/decoder.py +++ b/src/granite_switch/vllm/decoder/shadow_residual/decoder.py @@ -20,10 +20,13 @@ left zero (SR does not adapt K/V), which the loader marks as loaded. TP=1 for v1. """ +import contextlib + import torch from torch import nn from vllm.config import VllmConfig from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.attention.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -43,6 +46,45 @@ def _max_lora_rank(config) -> int: return max(config.adapter_ranks) if getattr(config, "adapter_ranks", None) else 0 +@contextlib.contextmanager +def _doubled_padding_mask(m: int): + """Match the forward context's ``is_padding`` to SR's doubled router batch. + + vLLM >= 0.27 passes an ``is_padding`` mask to the topk_softmax/topk_sigmoid + kernels so they can skip the padded rows of the batch. The router reads it off + the forward context and slices it to the number of gating rows + (``fused_moe/router/fused_topk_router.py::_get_padding_mask``), and the kernel + then *requires* the lengths to agree: + + RuntimeError: dispatch_topk_launch, topk_softmax_kernels.cu:785, + is_padding size mismatch, expected: 16384 + + The runner sizes the mask to the padded token count ``m``, but SR routes the + whole ``[2M, H]`` stack in one expert call, so the kernel sees ``2 * m`` gating + rows against an ``m``-entry mask. Doubling the mask is the right answer rather + than a workaround: adapter row ``i`` is the same token as base row ``i``, so it + is padding exactly when the base row is. + + Inert where there is nothing to match: no forward context (unit tests), no + ``is_padding`` field at all (``ForwardContext`` gained it at 0.25.1 — absent at + 0.24.0, which is why this is ``getattr`` and not attribute access), no mask + (``VLLM_MOE_SKIP_PADDING=0``, or 0.25/0.26 where the field exists but no kernel + consumes it), or a mask whose length is not ``m`` — in that last case the + premise of the doubling does not hold, and failing loudly in the kernel beats + silently masking the wrong rows. + """ + ctx = get_forward_context() if is_forward_context_available() else None + mask = getattr(ctx, "is_padding", None) if ctx is not None else None + if mask is None or mask.shape[0] != m: + yield + return + ctx.is_padding = torch.cat([mask, mask], dim=0) + try: + yield + finally: + ctx.is_padding = mask + + class ShadowResidualAttention(nn.Module): """Doubled-Q / base-only-KV attention for one SR decoder layer (SWITCH kernel).""" @@ -303,7 +345,9 @@ def forward(self, positions: torch.Tensor, hs: torch.Tensor) -> torch.Tensor: logits, _ = self.block_sparse_moe.gate(normed[:m]) # [M, E] logits = torch.cat([logits, logits], dim=0) # [2M, E] # FusedMoE modifies its input in place. - mlp_out = self.block_sparse_moe.experts(normed.clone(), logits) + # The doubled rows also need a doubled padding mask (>= 0.27). + with _doubled_padding_mask(m): + mlp_out = self.block_sparse_moe.experts(normed.clone(), logits) if self.has_shared_mlp: # Only ROUTING is shared. The dense shared MLP still runs per-stream diff --git a/src/granite_switch/vllm/fa3_aot_schedule.py b/src/granite_switch/vllm/fa3_aot_schedule.py new file mode 100644 index 00000000..d45e17f7 --- /dev/null +++ b/src/granite_switch/vllm/fa3_aot_schedule.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Size FlashAttention 3's ahead-of-time schedule for the layers it serves. + +vLLM builds one FlashAttention metadata builder per KV-cache group and sizes +FA3's ahead-of-time (AOT) scheduler metadata from the *model config*: its query +heads, KV heads and head size. A Granite Switch checkpoint has attention layers +whose shape differs from the model config: + +- the Shadow-Residual decoder doubles its query heads + (``Attention(2 * num_heads, ...)``) against base-only K/V, but the model + config reports the un-doubled ``num_attention_heads``; and +- the switch's counting and memory heads use their own ``head_size``. + +For such a group FA3's schedule is computed for the wrong shape. Eager steps +raise ``scheduler_metadata must have shape (metadata_size)``; CUDA-graph steps +compute wrong attention silently. It shows only in small, prefix-cached steps +(one conversation or one game served alone), so parity checks on batches of +fresh prompts pass while the served model misreads its context. See +issue #139 for the full trace. + +This patch wraps ``FlashAttentionMetadataBuilder.__init__`` and, for a +``granite_switch`` model, resizes the schedule from the group's *own* layers: + +- Single-shape group: overwrite ``num_heads_q``/``num_heads_kv``/``headdim`` with + the group's real shape. AOT stays on (the FULL-cudagraph buffer is sized from + batch, not shape, so it remains valid). This is the correctness fix for the SR + decoder group and keeps the FULL latency win. +- Genuinely mixed-shape group (rare): disable AOT *consistently* — + ``aot_schedule = False`` **and** tear down the FULL commitment + (``scheduler_metadata = None``, ``max_num_splits = 0``) so the builder state is + self-consistent whatever path ``build()`` takes, regardless of vLLM's internal + ordering. + +Other models are untouched. The patch targets vLLM 0.26-0.30 internals. + +NOTE on reproducing the bug: it only manifests under FlashAttention **3**, which +requires a Hopper GPU (compute capability 9.x). On older GPUs vLLM uses FA2, +whose cudagraph support forces ``cudagraph_mode=FULL`` down to +``FULL_AND_PIECEWISE`` and the bug cannot appear. This patch is still correct to +install everywhere; it is simply inert where FA3 is not used. + +History (issue #139): this ``__init__``-wrapping, schedule-resizing patch was +first written as commit 2c6035e, then shelved in favour of the +``get_cudagraph_support -> UNIFORM_BATCH`` downgrade in ``fa3_schedule.py`` (which +only acts under FULL cudagraphs). The eager TP-integration arms crash with +``scheduler_metadata must have shape (metadata_size)`` under ``enforce_eager``, +where the downgrade is a no-op, so this schedule-resizing patch is reinstated +*alongside* the downgrade — the two touch different seams and compose. +""" + +from __future__ import annotations + +import inspect + +_PATCHED = False + +# Builder instance attributes the patch reads/writes; asserted present so a +# future vLLM refactor fails loudly instead of silently installing a no-op. +_REQUIRED_BUILDER_ATTRS = ("num_heads_q", "num_heads_kv", "headdim", "aot_schedule") + + +def patch_flash_attn_aot_schedule() -> None: + """Install the FA3 AOT-schedule patch (idempotent). + + Soft-returns only when vLLM is genuinely absent (CPU-only / import error). + When vLLM is present but its internals have moved, it logs loudly and does + **not** mark the patch installed — a silent no-op here would reintroduce + wrong tokens under FA3 + FULL cudagraphs and crash eager FA3 steps. + """ + global _PATCHED + if _PATCHED: + return + + try: + from vllm.config import get_layers_from_vllm_config + from vllm.logger import init_logger + from vllm.model_executor.layers.attention.attention import Attention + from vllm.v1.attention.backends import flash_attn as fa + except ImportError: + # vLLM not installed (e.g. CPU-only or a non-vLLM backend). Not an error; + # do NOT mark _PATCHED so a later call in a vLLM-capable process retries. + return + + logger = init_logger(__name__) + + builder = getattr(fa, "FlashAttentionMetadataBuilder", None) + if builder is None: + logger.warning( + "granite_switch: could not find FlashAttentionMetadataBuilder in " + "vllm.v1.attention.backends.flash_attn; FA3 AOT-schedule patch NOT " + "installed. Under FlashAttention 3 this will crash eager steps with " + "'scheduler_metadata must have shape (metadata_size)' and produce " + "wrong tokens under cudagraph_mode=FULL (see issue #139). vLLM " + "internals may have changed; the patch needs updating." + ) + # Not marked installed: a vLLM version bump that restores the symbol + # (or a fixed patch) should get another chance. + return + + if getattr(builder, "_granite_switch_aot_schedule", False): + _PATCHED = True + return + + original = builder.__init__ + original_sig = inspect.signature(original) + + def __init__(self, *args, **kwargs): + # Forward verbatim: never re-positionalise. vLLM 0.26-0.30 construct the + # builder with a mix of positional and keyword arguments that differs by + # line, so a fixed four-positional signature here would raise TypeError at + # engine init on any line that passes one of these by keyword (the same + # reason ``_apply_hf_processor_main`` in audio/processor.py takes **kwargs). + original(self, *args, **kwargs) + + # Resolve the two arguments the correction needs by *name*, binding against + # the real signature so it does not matter whether the caller passed them + # positionally or as keywords. A rename (not a reorder) is the only thing + # that can defeat this, and it surfaces loudly below rather than silently. + try: + bound = original_sig.bind(self, *args, **kwargs) + bound.apply_defaults() + vllm_config = bound.arguments["vllm_config"] + layer_names = bound.arguments["layer_names"] + except (TypeError, KeyError): + logger.warning( + "granite_switch: could not resolve vllm_config/layer_names from " + "FlashAttentionMetadataBuilder.__init__%s; FA3 AOT-schedule " + "correction skipped. vLLM internals may have changed; the patch " + "needs updating (see issue #139).", + original_sig, + ) + return + + hf = getattr(vllm_config.model_config, "hf_config", None) + if getattr(hf, "model_type", None) != "granite_switch": + return + + # The attributes we are about to touch must exist; if vLLM renamed them + # the correction would silently miss, so fail loudly instead. + missing = [a for a in _REQUIRED_BUILDER_ATTRS if not hasattr(self, a)] + if missing: + raise AttributeError( + "granite_switch: FlashAttentionMetadataBuilder is missing " + f"attribute(s) {missing}; the FA3 AOT-schedule patch cannot size " + "the schedule for the switch's own layers. vLLM internals have " + "changed — update fa3_aot_schedule.py." + ) + + if not self.aot_schedule: + # No AOT schedule to size (e.g. FA2 fallback, or already disabled). + return + + layers = get_layers_from_vllm_config(vllm_config, Attention, layer_names) + if not layers: + # This builder serves a non-attention group; nothing to resize. + return + + shapes = {(a.num_heads, a.num_kv_heads, a.head_size) for a in layers.values()} + if len(shapes) == 1: + # One real shape for the whole group — size the schedule to it. This + # corrects the SR doubled-Q decoder group (and any switch group whose + # shape disagrees with the model config) while keeping AOT on. + self.num_heads_q, self.num_heads_kv, self.headdim = shapes.pop() + else: + # Genuinely mixed shapes in one group (rare). No single AOT schedule + # is correct, so disable AOT and tear down the FULL-cudagraph + # commitment consistently, leaving the builder self-consistent for + # whatever path build() takes. + self.aot_schedule = False + if hasattr(self, "scheduler_metadata"): + self.scheduler_metadata = None + if hasattr(self, "max_num_splits"): + self.max_num_splits = 0 + + builder.__init__ = __init__ + builder._granite_switch_aot_schedule = True + _PATCHED = True diff --git a/src/granite_switch/vllm/fa3_schedule.py b/src/granite_switch/vllm/fa3_schedule.py new file mode 100644 index 00000000..2f606b3b --- /dev/null +++ b/src/granite_switch/vllm/fa3_schedule.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Keep FlashAttention 3 attention out of the FULL CUDA graph for Granite Switch. + +On vLLM 0.26 under FlashAttention **3**, serving a Granite Switch checkpoint with +``cudagraph_mode=FULL`` produces attention that diverges — grossly, from the very +first prefill position — from the eager / ``FULL_AND_PIECEWISE`` reference. The +divergence is in the attention op *as captured inside the full CUDA graph*: +turning off FULL (eager), or dropping attention out of the graph +(``FULL_AND_PIECEWISE`` / ``PIECEWISE``), matches the reference exactly. See +issue #139. + +Why FA3 + FULL specifically +--------------------------- +``FlashAttentionMetadataBuilder.get_cudagraph_support`` reports +``AttentionCGSupport.ALWAYS`` when FA3 is active and ``UNIFORM_BATCH`` otherwise +(``vllm/v1/attention/backends/flash_attn.py``). ``ALWAYS`` is what lets +``cudagraph_mode=FULL`` capture attention for mixed prefill-decode batches; with +``UNIFORM_BATCH`` vLLM downgrades ``FULL`` to ``FULL_AND_PIECEWISE`` and runs +attention eagerly. On an FA2 (non-Hopper) box the backend already reports +``UNIFORM_BATCH``, FULL silently downgrades, and the bug never appears — which is +why it had never been seen until an FA3-capable (Hopper, compute 9.x) run. + +The break is **not** Granite-Switch-specific in mechanism: a plain Granite base +served through this same vLLM 0.26 + FA3 build diverges under FULL in exactly the +same way. It is an upstream FA3 full-cudagraph correctness problem on this vLLM +line. But Granite Switch is where it bites in production (FULL was chosen for the +latency of tiny prefix-cached decode steps), so the backend's broad FA3-ALWAYS +promise is wrong *for this model* and we narrow it here. + +The fix +------- +Wrap ``get_cudagraph_support`` so that, **for a granite_switch model only**, the +FlashAttention builder reports ``UNIFORM_BATCH`` instead of ``ALWAYS``. vLLM then +downgrades ``cudagraph_mode=FULL`` to ``FULL_AND_PIECEWISE`` on its own +(logging ``CUDAGraphMode.FULL is not supported ... setting +cudagraph_mode=FULL_AND_PIECEWISE``): the compute-heavy MLP/projection regions +stay in the graph, attention runs eagerly, and output matches the eager +reference token-for-token. Every other model keeps FA3's native ``ALWAYS`` +support untouched. + +An earlier revision of this module tried to *resize* FA3's ahead-of-time (AOT) +schedule — overwriting ``num_heads_q`` for the SR decoder's doubled query heads — +on the theory that a mis-sized schedule was the cause. A GPU reproduction +disproved it: the AOT head count was already being resized correctly, yet FULL +still diverged, and disabling AOT entirely did not help either (attention stayed +captured). The schedule shape was a red herring; the capture of attention itself +is the fault. See issue #139 for the full trace. + +This costs the FULL-vs-FULL_AND_PIECEWISE latency delta on attention launches, +which is the price of correctness until the upstream FA3 full-cudagraph path is +fixed on this vLLM line. ``PIECEWISE``/``FULL_AND_PIECEWISE`` remain fully +captured everywhere except the attention op. + +The patch targets vLLM 0.26-0.30 internals and is inert on any build without an +FA3 ``ALWAYS`` promise (e.g. FA2), where it would be a no-op anyway. +""" + +from __future__ import annotations + +_PATCHED = False + + +def patch_flash_attn_schedule() -> None: + """Install the FA3 full-cudagraph downgrade for granite_switch (idempotent). + + Soft-returns only when vLLM is genuinely absent (CPU-only / import error). + When vLLM is present but its internals have moved, it logs loudly and does + **not** mark the patch installed — a silent no-op here would reintroduce + wrong tokens under FlashAttention 3 + cudagraph_mode=FULL. + """ + global _PATCHED + if _PATCHED: + return + + try: + from vllm.logger import init_logger + from vllm.v1.attention.backends import flash_attn as fa + except ImportError: + # vLLM not installed (e.g. CPU-only or a non-vLLM backend). Not an error; + # do NOT mark _PATCHED so a later call in a vLLM-capable process retries. + return + + logger = init_logger(__name__) + + builder = getattr(fa, "FlashAttentionMetadataBuilder", None) + get_support = getattr(builder, "get_cudagraph_support", None) + if builder is None or get_support is None: + logger.warning( + "granite_switch: could not find " + "FlashAttentionMetadataBuilder.get_cudagraph_support in " + "vllm.v1.attention.backends.flash_attn; the FA3 full-cudagraph " + "downgrade is NOT installed. Under FlashAttention 3 + " + "cudagraph_mode=FULL this will produce wrong tokens (see " + "issue #139). vLLM internals may have " + "changed; the patch needs updating." + ) + # Not marked installed: a vLLM version bump that restores the symbol + # (or a fixed patch) should get another chance. + return + + if getattr(builder, "_granite_switch_cg_support", False): + _PATCHED = True + return + + # UNIFORM_BATCH is the FA2 support level: it tells vLLM the backend cannot + # take attention into a FULL graph for mixed prefill-decode, so FULL is + # downgraded to FULL_AND_PIECEWISE. Resolve the enum defensively so a rename + # fails loudly rather than silently installing a no-op. + cg_support_enum = getattr(fa, "AttentionCGSupport", None) + uniform_batch = getattr(cg_support_enum, "UNIFORM_BATCH", None) + if uniform_batch is None: + logger.warning( + "granite_switch: could not resolve " + "AttentionCGSupport.UNIFORM_BATCH; the FA3 full-cudagraph downgrade " + "is NOT installed and FlashAttention 3 + cudagraph_mode=FULL will " + "produce wrong tokens (see issue #139). " + "vLLM internals may have changed; the patch needs updating." + ) + return + + # get_cudagraph_support is a classmethod; unwrap to the plain function so the + # wrapper can re-bind it as a classmethod cleanly (and call through for every + # non-granite_switch model). + original = get_support.__func__ if hasattr(get_support, "__func__") else get_support + + def get_cudagraph_support(cls, vllm_config, kv_cache_spec): + hf = getattr(vllm_config.model_config, "hf_config", None) + if getattr(hf, "model_type", None) == "granite_switch": + # Force FULL -> FULL_AND_PIECEWISE: attention runs eagerly (correct) + # while the rest of the model stays in the CUDA graph. + return uniform_batch + return original(cls, vllm_config, kv_cache_spec) + + builder.get_cudagraph_support = classmethod(get_cudagraph_support) + builder._granite_switch_cg_support = True + _PATCHED = True diff --git a/tests/composer/_skinning_equivalence_worker_vllm.py b/tests/composer/_skinning_equivalence_worker_vllm.py index 11b5814c..b6f3df72 100644 --- a/tests/composer/_skinning_equivalence_worker_vllm.py +++ b/tests/composer/_skinning_equivalence_worker_vllm.py @@ -176,10 +176,18 @@ def cmd_run(args): # Save logprobs with open(output_path, "w") as f: json.dump({"logprobs": all_logprobs}, f) - print(f" saved logprobs to {output_path}") - + f.flush() + os.fsync(f.fileno()) + print(f" saved logprobs to {output_path}", flush=True) + + # vLLM/NCCL/CUDA background threads can abort during interpreter teardown + # ("terminate called without an active exception" -> SIGABRT), which would + # mask a run whose output is already safely on disk. The output is written + # and fsynced above, and the downstream compare step reads it from disk and + # needs no graceful vLLM shutdown, so skip Python's atexit/static + # destruction entirely rather than race the C++ teardown. del llm - return 0 + os._exit(0) # ── compare mode ────────────────────────────────────────────────── diff --git a/tests/conftest.py b/tests/conftest.py index 0be5df15..e15f0b2c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,35 @@ def pytest_configure(config): + # ── flashinfer sampler: off unless the runner asks for it ───── + # + # MEASURED, not precautionary. On an image shipping CUDA 12.4 + # (/usr/local/cuda-12.4, while the driver advertises 13.0 -- that is the driver + # ceiling, not the toolkit) the vLLM 0.30 engine failed to start 171 times on + # + # nvcc fatal : Unknown option '--compress-mode=size' + # + # flashinfer JIT-compiles its sampling kernels with `--compress-mode`, an nvcc + # option added in CUDA 12.8. Every GPU suite went red downstream of that -- + # including on UPSTREAM's own GraniteMoeHybridForCausalLM, so it is not a + # granite-switch defect, and it masks everything else because the engine never + # reaches a test. + # + # It is version-dependent, which is why this lives here now: vLLM 0.26 is green + # on that image WITHOUT this (test_logs/gputest-v26, 219 passed) and 0.30 is not + # (test_logs/gputest-v30), so raising the CI ceiling to 0.30 is exactly what + # exposes it. The variable is read and honoured on all of 0.26-0.30 + # (vllm/v1/sample/ops/topk_topp_sampler.py). + # + # SCOPED TO THE SAMPLER, and inert for what these tests assert: the attention + # backend resolves to FlashAttention independently, and the equivalence suites + # compare logprobs, which the model produces before any sampling op runs. + # + # setdefault, so an image with CUDA >= 12.8 can export + # VLLM_USE_FLASHINFER_SAMPLER=1 and genuinely exercise that path. The real fix is + # a newer toolkit in the image; this keeps the suites meaningful until then. + os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") + worker_id = os.environ.get("PYTEST_XDIST_WORKER") if worker_id is None: return # not running under xdist diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 00000000..1278ba77 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Integration-suite collection policy: keep the parent pytest process CUDA-free. + +Same invariant, same hazard, same remedy as ``tests/vllm/conftest.py`` -- see +its module docstring for the full story. Under **Exclusive_Process** GPU compute +mode only one process may hold the device's primary context at a time. The +vLLM-serving integration tests (``test_audio_serving_smoke``, +``test_audio_uncached_processor``, ``test_multi_switch_vllm_generate`` ...) do +their CUDA work in a vLLM ``EngineCore`` **subprocess**; that subprocess dies +with ``cudaErrorDevicesUnavailable`` ("CUDA-capable device(s) is/are busy or +unavailable") the instant the parent pytest process is already holding a +context. ``torch.cuda.empty_cache()`` frees cached memory but not the primary +context, so a teardown fixture cannot hand it back -- only the holding process +exiting releases it. + +Two integration modules create a CUDA context *directly in the test process* +(``.cuda()`` / ``device="cuda"`` in a module-scoped fixture or test body), +because they drive a HuggingFace / transformers model in-process with nothing to +delegate to a worker: + +* ``test_adapter_routing_audio_enabled.py`` -- ``.cuda()`` on the HF model +* ``test_asr_ctc_default_gpu.py`` -- ``ASRTranscriber(device="cuda")`` + +Collected in-parent, either one poisons every vLLM-subprocess test ordered after +it (alphabetical collection puts both ahead of the serving tests, which is +exactly how the suite went red). ``ISOLATED_MODULES`` lists them; each runs in +its own fronted ``pytest`` subprocess so its context lives and dies inside that +child and the parent stays clean regardless of collection order. + +To add a file here: it must create a CUDA context directly in the test process, +rather than inside a subprocess worker. +""" + +import os + +# Reuse the proven isolation machinery from the vLLM suite rather than duplicate +# it. ``_IsolatedModuleItem`` shells out to ``pytest `` in a fresh +# process; ``_CHILD_ENV_FLAG`` makes the hook inert inside that child so the real +# tests run (and own the child's context harmlessly). +from tests.vllm.conftest import _CHILD_ENV_FLAG, _IsolatedModuleItem + +# Integration modules that create a CUDA context in-process and must therefore +# run in their own subprocess. Keep this minimal -- prefer the subprocess-worker +# pattern for new GPU tests. +ISOLATED_MODULES = frozenset( + { + "test_adapter_routing_audio_enabled.py", + "test_asr_ctc_default_gpu.py", + } +) + + +def pytest_collection_modifyitems(config, items): + """Front each isolated module as a single subprocess runner. + + In the parent we drop the in-module items (so the parent never imports or + executes their CUDA bodies) and prepend one item per isolated module that + shells out to ``pytest `` in a fresh process. In the spawned child + (flagged via the env var) collection is left untouched. + """ + if os.environ.get(_CHILD_ENV_FLAG): + return # inside the isolated child: run the real tests as collected + + if not items: + return + + kept = [] + isolated_paths = {} # module fspath -> basename, first-seen order preserved + for item in items: + basename = os.path.basename(str(item.path)) + if basename in ISOLATED_MODULES: + isolated_paths.setdefault(str(item.path), basename) + else: + kept.append(item) + + if not isolated_paths: + return + + session = items[0].session + runner_items = [ + _IsolatedModuleItem.from_parent( + session, + name=f"{basename}::", + module_path=path, + ) + for path, basename in isolated_paths.items() + ] + # Run the isolated subprocess modules FIRST, before any in-parent vLLM + # EngineCore subprocess starts. Under Exclusive_Process only one process can + # hold the device: each isolated child must fully exit (releasing the GPU) + # before the serving tests acquire it, and vice versa. Fronting them + # guarantees that ordering regardless of how the rest of the suite is + # collected. + items[:] = runner_items + kept + + +def shutdown_llm(llm): + """Deterministically tear down a vLLM ``LLM`` and release its GPU. + + ``del llm; gc.collect(); torch.cuda.empty_cache()`` is **not** enough. vLLM + runs its ``EngineCore`` in a child process (``VLLM_WORKER_MULTIPROC_METHOD= + spawn``); ``empty_cache()`` frees cached memory but neither tears that + subprocess down nor releases the device's primary CUDA context -- the engine + only exits via its own ``atexit`` hook, i.e. at *pytest* exit, not at a + fixture boundary. + + That gap is fatal for a ``scope="module"`` fixture parametrized over several + checkpoints: when the fixture is rebuilt for the next checkpoint, the + previous engine is still holding the GPU, and under **Exclusive_Process** the + new ``EngineCore`` dies with ``cudaErrorDevicesUnavailable``. Calling the + engine-core shutdown explicitly makes the subprocess exit and hand the device + back before the next engine boots. + + Best-effort and version-tolerant: the private path has moved across vLLM + releases, so every step is guarded and a missing attribute is a no-op (the + ``atexit`` hook remains the backstop). + """ + import gc + + try: + engine = getattr(llm, "llm_engine", None) + engine_core = getattr(engine, "engine_core", None) + shutdown = getattr(engine_core, "shutdown", None) + if callable(shutdown): + shutdown() + except Exception: + pass # best-effort: fall back to del + atexit below + finally: + del llm + gc.collect() + try: + import torch + + torch.cuda.empty_cache() + except Exception: + pass diff --git a/tests/integration/test_audio_serving_smoke.py b/tests/integration/test_audio_serving_smoke.py index 4e493da0..11edcc3a 100644 --- a/tests/integration/test_audio_serving_smoke.py +++ b/tests/integration/test_audio_serving_smoke.py @@ -108,9 +108,7 @@ def served(audio_checkpoint): Tokenizer init stays ON (unlike the argmax-equivalence test): the ASR processor needs it to encode the prompt and the transcript. """ - import gc - - import torch + from tests.integration.conftest import shutdown_llm os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") from vllm import LLM @@ -124,9 +122,10 @@ def served(audio_checkpoint): try: yield {"llm": llm, "config": llm.llm_engine.model_config.hf_config} finally: - del llm - gc.collect() - torch.cuda.empty_cache() + # Shut the EngineCore subprocess down (not just empty_cache) so it hands + # the GPU back before the next parametrized checkpoint boots its engine; + # under Exclusive_Process a leaked engine kills the next one. + shutdown_llm(llm) def _tone(seconds: float = 1.0, freq: float = 440.0): diff --git a/tests/integration/test_audio_uncached_processor.py b/tests/integration/test_audio_uncached_processor.py index af83440b..2eba3f41 100644 --- a/tests/integration/test_audio_uncached_processor.py +++ b/tests/integration/test_audio_uncached_processor.py @@ -134,9 +134,7 @@ def served_uncached(audio_checkpoint): a separate test it would only skip itself, leaving the rest of the module green while silently running the cached path. """ - import gc - - import torch + from tests.integration.conftest import shutdown_llm cache_kwargs = _cache_disabling_kwargs() if not cache_kwargs: @@ -173,9 +171,10 @@ def served_uncached(audio_checkpoint): "tokenizer": llm.get_tokenizer(), } finally: - del llm - gc.collect() - torch.cuda.empty_cache() + # Shut the EngineCore subprocess down (not just empty_cache) so it hands + # the GPU back before the next parametrized checkpoint boots its engine; + # under Exclusive_Process a leaked engine kills the next one. + shutdown_llm(llm) @pytest.fixture(scope="module") diff --git a/tests/integration/test_hf_to_vllm_weights.py b/tests/integration/test_hf_to_vllm_weights.py index 364f8d34..2e22ce9d 100644 --- a/tests/integration/test_hf_to_vllm_weights.py +++ b/tests/integration/test_hf_to_vllm_weights.py @@ -68,6 +68,7 @@ def _try_import_vllm(): # ── Distributed init ───────────────────────────────────────────────── from tests.shared.vllm_distributed import ensure_distributed as _ensure_distributed +from tests.shared.vllm_kv_cache import setup_kv_cache # ── Helpers ────────────────────────────────────────────────────────── @@ -247,15 +248,9 @@ def _setup_vllm_kv_caches(self): def _setup_single_attn(self, attn, layer_name, num_blocks): """Configure a single Attention layer with KV cache.""" - attn.kv_cache_torch_dtype = torch.bfloat16 - cache_shape = attn.attn_backend.get_kv_cache_shape( - num_blocks, - BLOCK_SIZE, - attn.num_kv_heads, - attn.head_size, + kv_cache = setup_kv_cache( + attn, self.vllm_config, num_blocks, BLOCK_SIZE, self.device ) - kv_cache = torch.zeros(cache_shape, device=self.device, dtype=torch.bfloat16) - attn.kv_cache = kv_cache self._kv_caches.append(kv_cache) self._attention_map[layer_name] = attn diff --git a/tests/shared/_vllm_logprobs_worker.py b/tests/shared/_vllm_logprobs_worker.py new file mode 100644 index 00000000..5866e2c8 --- /dev/null +++ b/tests/shared/_vllm_logprobs_worker.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Subprocess worker: run ONE vLLM engine and emit teacher-forced prompt logprobs. + +Invoked as a fresh process per engine so the **pytest process never creates a +CUDA context**:: + + python _vllm_logprobs_worker.py run \ + --model --inputs --out + +This is load-bearing under the GPU host's NVIDIA **Exclusive_Process** compute +mode: only one OS process may hold a device's primary CUDA context at a time, so +if the pytest process touched ``torch.cuda`` (even a ``mem_get_info`` probe) it +would own the device and every ``EngineCore`` child it spawned afterward would die +with ``cudaErrorDevicesUnavailable``. + +The engine lives entirely inside this short-lived process. Because the process +exits after a single engine there is no back-to-back-engine handoff within it, so +the ``gpu_mem_util`` / ``shutdown_llm`` helpers (which DO touch ``torch.cuda``) are +called HERE -- in the process that legitimately holds the context -- rather than in +the pytest process. + +Contract: +- ``inputs.json``: ``{"input_ids": [int, ...], "vocab_size": int, + "llm_kwargs": {...}}``. ``llm_kwargs`` is forwarded verbatim to ``LLM(...)`` and + carries exactly what the caller passed (e.g. ``max_model_len``, ``enforce_eager``, + ``enable_prefix_caching``, ``gpu_memory_utilization``). +- ``--out``: a ``torch.save`` of the dense ``[seq_len-1, vocab_size]`` CPU logprob + tensor, loadable back in the pytest process without CUDA. +""" + +import argparse +import json +import os +import sys + +# Make tests.shared importable when run as a bare subprocess (cwd-independent), +# mirroring tests/vllm/_generation_equivalence_worker.py. +_REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +if _REPO not in sys.path: + sys.path.insert(0, _REPO) + + +def _pin_single_visible_device(): + """Narrow ``CUDA_VISIBLE_DEVICES`` to a single device for this worker. + + Each worker runs exactly one ``tensor_parallel_size=1`` engine, so it never + needs more than one GPU. Leaving two devices visible makes vLLM's distributed + init (``init_worker_distributed_environment`` -> NCCL) set up across both even + at world_size=1, which widens the device-acquisition surface -- and under the + host's NVIDIA Exclusive_Process compute mode that is exactly where engine init + intermittently dies with ``cudaErrorDevicesUnavailable``. Pinning to one device removes the + cross-device NCCL step and halves the surface. + + WHICH device: if the parent (``run_vllm_logprobs``) set ``GS_LOGPROBS_PIN_DEVICE``, + use exactly that entry of ``CUDA_VISIBLE_DEVICES``. The parent round-robins that + index across consecutive engines so that back-to-back non-eager spawns land on + DIFFERENT physical GPUs -- which is the cluster-2 Cause-B fix: the device that + just tore down a non-eager EngineCore gets a full engine-lifetime (~60-90 s) to + relinquish its primary context before the next non-eager spawn hits it, instead + of being re-acquired immediately into the window its own teardown opened. The + recovery gap the host needs is paid in parallelism (the other GPU), not in test + wall-clock. Falls back to the FIRST entry when the parent set nothing (single + engine, or an opted-out run), so a harness that already restricts/orders devices + is respected; or ``"0"`` if ``CUDA_VISIBLE_DEVICES`` is unset. + + Must run BEFORE any ``torch``/``vllm`` import so no context has formed yet. + ``GS_LOGPROBS_NO_PIN=1`` opts out (e.g. to A/B this behaviour). + """ + if os.environ.get("GS_LOGPROBS_NO_PIN"): + return + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible is None: + chosen = "0" + else: + entries = [v.strip() for v in visible.split(",") if v.strip() != ""] + if len(entries) == 0: + return # no devices -- nothing to narrow + # Parent-chosen round-robin slot (modulo the real entry count, so a stale + # or oversized index can never point off the end of the allowlist). + pin_idx_raw = os.environ.get("GS_LOGPROBS_PIN_DEVICE") + if pin_idx_raw is not None and pin_idx_raw.strip() != "": + try: + slot = int(pin_idx_raw) % len(entries) + except ValueError: + slot = 0 + else: + slot = 0 + if len(entries) <= 1 and slot == 0: + return # already a single device -- nothing to narrow + chosen = entries[slot] + os.environ["CUDA_VISIBLE_DEVICES"] = chosen + print(f" pinned CUDA_VISIBLE_DEVICES={chosen} (was {visible!r})") + + +def cmd_run(args): + """Build one vLLM engine, capture dense prompt logprobs, save to --out.""" + # Pin to one GPU BEFORE importing torch/vllm (see _pin_single_visible_device). + _pin_single_visible_device() + # vLLM v1 forks a subprocess for EngineCore; spawn avoids inheriting any CUDA + # state. (This process has no CUDA context yet -- the engine child creates it.) + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + import torch + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + from granite_switch.vllm import register as register_granite_switch + from tests.shared.vllm_equivalence import extract_logprobs_tensor + from tests.shared.vllm_gpu_mem import gpu_mem_util, shutdown_llm + + register_granite_switch() + + with open(args.inputs) as f: + spec = json.load(f) + input_ids_list = spec["input_ids"] + vocab_size = spec["vocab_size"] + llm_kwargs = dict(spec.get("llm_kwargs") or {}) + + # Allow requesting logprobs for all vocab tokens. + llm_kwargs.setdefault("max_logprobs", vocab_size) + + # gpu_memory_utilization: vLLM's default is a fraction of TOTAL memory, which + # assumes an idle card. These suites start engines back to back (one worker + # process each) and the previous one may still be releasing, so size against + # what is actually free. setdefault, so an explicit caller value still wins. + # See tests/shared/vllm_gpu_mem.py. Safe to probe torch.cuda here: this is the + # process that will hold the engine's context anyway. + mem_util = gpu_mem_util() + if mem_util is not None: + llm_kwargs.setdefault("gpu_memory_utilization", mem_util) + + print(f"Creating vLLM LLM for {args.model} with kwargs {sorted(llm_kwargs)}...") + llm = LLM( + model=args.model, + skip_tokenizer_init=True, + dtype="bfloat16", + **llm_kwargs, + ) + + sampling_params = SamplingParams( + max_tokens=1, + temperature=0.0, + prompt_logprobs=vocab_size, + ) + prompt = TokensPrompt(prompt_token_ids=input_ids_list) + outputs = llm.generate(prompt, sampling_params=sampling_params) + + logprobs_tensor = extract_logprobs_tensor(outputs, vocab_size) + torch.save(logprobs_tensor, args.out) + print(f" saved logprobs {tuple(logprobs_tensor.shape)} to {args.out}") + + # Synchronous shutdown + device/memory reclaim, in the process that owns the + # context. The process then exits, which under Exclusive_Process is what + # actually releases the device for the next worker. + shutdown_llm(llm) + del llm + return 0 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="mode", required=True) + + p_run = sub.add_parser("run", help="Run one engine, emit dense prompt logprobs") + p_run.add_argument("--model", required=True, help="Model dir or HF name to load") + p_run.add_argument( + "--inputs", required=True, help="JSON: {input_ids, vocab_size, llm_kwargs}" + ) + p_run.add_argument("--out", required=True, help="Path to torch.save the tensor to") + + args = parser.parse_args() + if args.mode == "run": + return cmd_run(args) + raise SystemExit(f"unknown mode: {args.mode}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/shared/granite4_equivalence.py b/tests/shared/granite4_equivalence.py index d9c3efb9..6ec0f9fb 100644 --- a/tests/shared/granite4_equivalence.py +++ b/tests/shared/granite4_equivalence.py @@ -281,6 +281,20 @@ def get_tolerances(layer_types, long_sequence=False, has_kv_hidden=False): # ~20x the measured 4.77e-7 worst case (comfortable margin) yet ~1000x # tighter than the adapter-path tolerance, so a real logit regression # still trips it. + # + # CAVEAT: 4.77e-7 was measured at num_adapters=0 AND COMPILED, and it is + # only that tight there. Callers that pass has_kv_hidden=False with the + # ADAPTER TIER LIVE -- TestZeroAdapterNoHiding, via + # save_zero_adapter_model's num_adapters=2 with the LoRA weights zeroed -- + # get a much larger reduction-order difference, because the projection + # runs the fused SWITCH kernel plus a zero-valued shrink/expand instead of + # the plain base linear. Measured eager on vLLM 0.26.0 + # (test_logs/floor-eager): 9.9182e-04 / 1.9522e-03 / 5.0116e-04 for + # 4.0-{1b,350m,micro}, with each engine bit-exact against a second + # instance of itself over all 3840 elements -- deterministic, not noise. + # torch.compile lowers both op sequences to equivalent kernels and brings + # it back under 1e-5, which is why those tests must not be run eager; see + # tests/shared/vllm_equivalence.py::run_vllm_logprobs. return (1e-5, 1e-5) else: # Substitute-embedding propagates through attention to visible diff --git a/tests/shared/vllm_equivalence.py b/tests/shared/vllm_equivalence.py index a5629ae7..f28bb85a 100644 --- a/tests/shared/vllm_equivalence.py +++ b/tests/shared/vllm_equivalence.py @@ -16,9 +16,13 @@ """ import gc +import itertools import json import os +import subprocess +import sys import tempfile +import time import torch @@ -189,83 +193,207 @@ def extract_logprobs_tensor(outputs, vocab_size): return tensor -def dump_model_params(model_dir, dump_path): - """Load model via vLLM and dump parameter checksums to a file.""" - os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - - from vllm import LLM - - llm = LLM( - model=model_dir, - skip_tokenizer_init=True, - dtype="bfloat16", +#: The GPU worker that actually stands up each engine (see module docstring). +#: Runs in its OWN process so this -- the pytest -- +#: process never creates a CUDA context. +_LOGPROBS_WORKER = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "_vllm_logprobs_worker.py" +) +#: Per engine: HF checkpoint load + vLLM init + a single teacher-forced generate. +#: Generous; the real gate is correctness, and a hung child should fail loudly. +_WORKER_TIMEOUT = float(os.environ.get("GS_LOGPROBS_WORKER_TIMEOUT", "1200")) +#: Fresh-worker attempts on a transient device-unavailable init failure (see +#: _spawn_logprobs_worker). Each attempt is a brand-new process, which is what +#: lets it ride through the host's sticky-but-transient Exclusive_Process window. +_WORKER_INIT_RETRIES = int(os.environ.get("GS_LOGPROBS_WORKER_RETRIES", "4")) +#: Seconds to let the device settle between those attempts. +_WORKER_RETRY_DELAY = float(os.environ.get("GS_LOGPROBS_WORKER_RETRY_DELAY", "15")) +#: Monotonic counter so each engine run gets its own handoff files inside tmpdir. +#: Also drives the per-engine round-robin device pin (see _visible_device_count / +#: GS_LOGPROBS_PIN_DEVICE in _vllm_logprobs_worker.py): consecutive engines land on +#: alternating GPUs, so a worker whose EngineCore child is still releasing a device +#: does not collide with the next worker trying to acquire the SAME device under +#: Exclusive_Process. (The primary cluster-2 fix is keeping the worker PARENT +#: CUDA-free -- see vllm_gpu_mem.gpu_mem_util; round-robin just spreads the two +#: workers of a pair across the two cards so their teardown/acquire windows never +#: overlap on one device.) +_logprobs_run_counter = itertools.count() +#: Set GS_LOGPROBS_NO_ROUND_ROBIN=1 to disable device round-robin (e.g. to A/B it or +#: on a single-GPU host where it is a no-op anyway). +_ROUND_ROBIN_DEVICES = not os.environ.get("GS_LOGPROBS_NO_ROUND_ROBIN") + + +def _visible_device_count(): + """How many GPUs are visible to the workers, WITHOUT creating a CUDA context. + + The pytest process must stay CUDA-free, so + this never calls ``torch.cuda``. It reads ``CUDA_VISIBLE_DEVICES`` when set + (the authoritative allowlist the workers also parse), else asks ``nvidia-smi``. + Returns ``1`` on any ambiguity so the round-robin degrades to "always device 0" + -- i.e. today's single-device pin -- rather than pointing off the end. + """ + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible is not None: + entries = [v.strip() for v in visible.split(",") if v.strip() != ""] + return max(1, len(entries)) + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ).stdout + n = len([line for line in out.splitlines() if line.strip()]) + return max(1, n) + except Exception: + return 1 + + +def _is_transient_device_error(text: str) -> bool: + """True if worker output shows the Exclusive_Process device-availability race. + + The signature is ``cudaErrorDevicesUnavailable`` ("CUDA-capable device(s) + is/are busy or unavailable"), raised at engine ``init_device`` -- a transient + the host throws under NVIDIA Exclusive_Process compute mode, independent of the + model or weights. Matching the worker's + FULL stdout+stderr (not a wrapped parent exception string) is why this can key + on the real error text, which the old in-process retry could not see. + """ + low = text.lower() + return ( + "cudaerrordevicesunavailable" in low + or "device(s) is/are busy or unavailable" in low ) - # Access model params via internal API - model = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model - checksums = {} - for name, param in model.named_parameters(): - checksums[name] = ( - param.shape, - param.sum().item(), - param.abs().max().item(), - param.float().norm().item(), - ) - - torch.save(checksums, dump_path) - - del llm - gc.collect() - torch.cuda.empty_cache() +def run_vllm_logprobs(model_dir, input_ids_list, vocab_size, *, tmpdir, **llm_kwargs): + """Run one vLLM engine in a subprocess, return its dense prompt logprobs. -def run_vllm_logprobs(model_dir, input_ids_list, vocab_size, **llm_kwargs): - """Load model via vllm.LLM, extract prompt logprobs as dense tensor. + The engine is NOT built in this process. Under the GPU host's + Exclusive_Process compute mode, if this (the pytest) process ever created a + CUDA context -- which ``gpu_mem_util``/``shutdown_llm`` would, via + ``torch.cuda.mem_get_info``/``empty_cache`` -- it would own the device and + every ``EngineCore`` child spawned afterward would die with + ``cudaErrorDevicesUnavailable``. So the engine, and every ``torch.cuda`` call, + live entirely in ``_vllm_logprobs_worker.py``, one fresh process per engine. + This process only writes a JSON spec, shells out, and loads back a CPU tensor. Args: model_dir: path to saved HF model (config.json + model.safetensors) input_ids_list: list of int token IDs vocab_size: vocab size for dense tensor allocation + tmpdir: directory for the subprocess handoff files (the caller's tmp_path) **llm_kwargs: extra kwargs for LLM() constructor (e.g. max_model_len, - gpu_memory_utilization) + enforce_eager), forwarded verbatim to the worker Returns: [seq_len-1, vocab_size] float tensor of log-probabilities """ - # vLLM v1 forks a subprocess for EngineCore; if CUDA was initialized - # in the parent (e.g. by torch.cuda.is_available() probe), fork fails. - # Spawn avoids inheriting the parent's CUDA state. - os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - - from vllm import LLM, SamplingParams - from vllm.inputs import TokensPrompt - - # Allow requesting logprobs for all vocab tokens - llm_kwargs.setdefault("max_logprobs", vocab_size) - - llm = LLM( - model=model_dir, - skip_tokenizer_init=True, - dtype="bfloat16", - **llm_kwargs, - ) - - sampling_params = SamplingParams( - max_tokens=1, - temperature=0.0, - prompt_logprobs=vocab_size, - ) - - prompt = TokensPrompt(prompt_token_ids=input_ids_list) - outputs = llm.generate(prompt, sampling_params=sampling_params) - - logprobs_tensor = extract_logprobs_tensor(outputs, vocab_size) + # enable_prefix_caching: a SYMMETRIC-config default, applied here (no CUDA) so + # both engines of a pair get it. upstream's GraniteMoeHybridForCausalLM declares + # IsHybrid and GraniteSwitchForCausalLM deliberately does not (declaring it + # would size the KV cache for zero attention layers), so the two sides resolve + # DIFFERENT prefix-caching defaults -- an asymmetry in a test whose whole job is + # to have one variable. It also kills a vLLM 0.28-only crash in upstream's + # model: with prefix caching on, 0.28 logs "Mamba cache mode is set to 'align'" + # and then asserts "no mamba layers in the model" (v1/worker/mamba_utils.py). + llm_kwargs.setdefault("enable_prefix_caching", False) + # + # Deliberately NOT defaulting enforce_eager, and NOT a free choice -- it decides + # whether the adapter-present runners can meet get_tolerances' 1e-5 inert-path + # gate. Measured on vLLM 0.26.0, eager, zero-LoRA pair: switch-vs-upstream is a + # real, fully deterministic ~1e-3 difference (each engine bit-exact against a + # second copy of itself), appearing ONLY when the adapter tier is live, because + # the fused SWITCH kernel reduces in a different order than the native linear. + # torch.compile normalizes that order and closes the gap while preserving + # semantics (a real kernel bug still trips 1e-5). The one place eager IS + # required is _granite4_fullsize_tests.py, which asks for it explicitly (an + # unpinned autotuner reddened ~1 run in 7 at full size). max_logprobs and + # gpu_memory_utilization are set in the worker, where torch.cuda is safe to + # touch. Full rationale table: git history of this file. + + run_id = next(_logprobs_run_counter) + spec_path = os.path.join(str(tmpdir), f"logprobs_inputs_{run_id}.json") + out_path = os.path.join(str(tmpdir), f"logprobs_out_{run_id}.pt") + with open(spec_path, "w") as f: + json.dump( + { + "input_ids": list(input_ids_list), + "vocab_size": vocab_size, + "llm_kwargs": llm_kwargs, + }, + f, + ) - del llm - gc.collect() - torch.cuda.empty_cache() + cmd = [ + sys.executable, + _LOGPROBS_WORKER, + "run", + "--model", + str(model_dir), + "--inputs", + spec_path, + "--out", + out_path, + ] + # Bounded retry on the transient device-unavailable race ONLY. Each attempt is + # a FRESH worker process -- the key difference from the old in-process retry: a + # new process gets a clean CUDA context, which is what rides through the host's + # sticky-but-transient Exclusive_Process window (an in-process retry could not, + # and anyway it keyed on a wrapped exception string that never matched). vLLM's + # non-eager init (torch.compile + CUDA-graph capture) is the path that trips + # this; eager init does not, but the adapter-live tests require non-eager for + # the tolerance gate, so retry rather than force eager. Any OTHER failure (a + # real config/weight/logic error) raises immediately -- it is not retried. + # Round-robin the pinned GPU across engines (and across retries of one engine) + # so consecutive spawns do not contend for the same device while a previous + # worker's EngineCore child is still releasing it under Exclusive_Process. A + # retry advances the slot too, moving a failed spawn to the OTHER GPU rather + # than back onto the one that just refused it. No-op on a single-GPU host + # (slot stays 0). This is a secondary robustness measure; the primary fix is + # keeping the worker parent CUDA-free (vllm_gpu_mem.gpu_mem_util). + device_count = _visible_device_count() if _ROUND_ROBIN_DEVICES else 1 + print(f"\n logprobs worker: {' '.join(cmd)}") + for attempt in range(1, _WORKER_INIT_RETRIES + 1): + worker_env = dict(os.environ) + if device_count > 1: + slot = (run_id + attempt - 1) % device_count + worker_env["GS_LOGPROBS_PIN_DEVICE"] = str(slot) + print(f" round-robin device slot {slot} (of {device_count})") + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=_WORKER_TIMEOUT, + env=worker_env, + ) + if result.stdout: + print(result.stdout[-4000:]) + if result.stderr: + print(" STDERR:", result.stderr[-2000:]) + if result.returncode == 0: + break + + combined = f"{result.stdout}\n{result.stderr}" + transient = _is_transient_device_error(combined) + if not transient or attempt == _WORKER_INIT_RETRIES: + raise RuntimeError( + f"vLLM logprobs worker failed for {model_dir} " + f"(exit {result.returncode}" + f"{', transient device-unavailable, retries exhausted' if transient else ''}" + f").\n" + f"STDOUT (tail):\n{result.stdout[-2000:]}\n" + f"STDERR (tail):\n{result.stderr[-2000:]}" + ) + print( + f" worker hit the transient device-unavailable race " + f"(attempt {attempt}/{_WORKER_INIT_RETRIES}); letting the device settle " + f"{_WORKER_RETRY_DELAY:g}s and re-spawning a fresh worker" + ) + time.sleep(_WORKER_RETRY_DELAY) - return logprobs_tensor + # Loaded on CPU: the pytest process must stay CUDA-free. + return torch.load(out_path, map_location="cpu") # ── Full integration pipeline ───────────────────────────────────── @@ -312,12 +440,14 @@ def run_equivalence_integration(cfg_dict, *, seq_len=16, seed=0, tmpdir, **llm_k upstream_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) switch_logprobs = run_vllm_logprobs( switch_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) @@ -371,12 +501,14 @@ def run_zero_adapter_no_hiding_equivalence( upstream_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) switch_logprobs = run_vllm_logprobs( switch_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) @@ -421,12 +553,14 @@ def run_zero_adapter_equivalence(cfg_dict, *, seq_len=16, seed=0, tmpdir, **llm_ upstream_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) switch_logprobs = run_vllm_logprobs( switch_dir, input_ids, vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) @@ -474,12 +608,14 @@ def run_gap_equivalence(cfg_dict, *, seq_len, ctrl_pos, seed=0, tmpdir, **llm_kw upstream_dir, upstream_ids[0].tolist(), vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) switch_logprobs = run_vllm_logprobs( switch_dir, switch_ids[0].tolist(), vocab_size, + tmpdir=tmpdir, **llm_kwargs, ) diff --git a/tests/shared/vllm_gpu_mem.py b/tests/shared/vllm_gpu_mem.py new file mode 100644 index 00000000..36e2a8b3 --- /dev/null +++ b/tests/shared/vllm_gpu_mem.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pick a ``gpu_memory_utilization`` that fits the memory actually free right now. + +vLLM measures the fraction against TOTAL device memory and refuses to start when +that exceeds what is free:: + + Free memory on device cuda:0 (4.6/79.25 GiB) ... is less than desired GPU + memory utilization (0.92, 72.91 GiB) -- vllm/v1/worker/utils.py + +So a fixed value — including vLLM's own default — only works on an idle card, and +these suites do not have one: they start several engines in a row and the previous +engine's subprocess is often still releasing, which is enough to trip the guard +in-pod with the second GPU completely idle. The guard is byte-identical back to +0.26, so this is not a version problem. +""" + +import gc +import os +import subprocess +import time + +#: Ceiling, i.e. what to ask for on an empty card. +GPU_MEM_UTIL = float(os.environ.get("GS_GPU_MEM_UTIL", "0.85")) +#: Below this the card is too full to be worth starting on -- a 3B model plus its +#: KV cache needs roughly a tenth of an 80 GiB card, so this is generous. +GPU_MEM_UTIL_FLOOR = float(os.environ.get("GS_GPU_MEM_UTIL_FLOOR", "0.15")) +#: Free fraction that counts as "the previous engine is gone". Not 1.0: the parent +#: holds a small CUDA context of its own once it has called ``mem_get_info``. +GPU_MEM_RECLAIM_TARGET = float(os.environ.get("GS_GPU_MEM_RECLAIM_TARGET", "0.80")) +#: How long to wait for that. Generous, because the alternative to waiting is a +#: spurious red, and a shut-down engine normally releases within a few seconds. +GPU_MEM_RECLAIM_TIMEOUT = float(os.environ.get("GS_GPU_MEM_RECLAIM_TIMEOUT", "120")) +#: How long to wait for the DEVICE itself (not its memory) to become acquirable by +#: a new CUDA context after an engine exits. See device_is_free / wait_device_free. +GPU_DEVICE_READY_TIMEOUT = float(os.environ.get("GS_GPU_DEVICE_READY_TIMEOUT", "60")) + + +def _nvsmi_free_total_bytes(): + """(free, total) bytes for the worker's visible GPU, via ``nvidia-smi``. + + CUDA-FREE on purpose -- NOT ``torch.cuda.mem_get_info``. ``gpu_mem_util`` runs + in the worker *before* ``LLM()``, and ``LLM()`` spawns the ``EngineCore`` child + that creates the real CUDA context. Under the host's NVIDIA + ``Exclusive_Process`` compute mode only ONE process may hold a device's primary + context at a time, so if this (the parent worker) process created a context + first -- which ``torch.cuda.mem_get_info`` does: it flips + ``torch.cuda.is_initialized()`` to True -- the ``EngineCore`` child's first CUDA + op dies with ``cudaErrorDevicesUnavailable`` at ``init_device`` -> + ``MemorySnapshot`` -> ``get_memory_info``. That is the real cluster-2 Cause B + (the doc's "self-sustaining non-eager window" was a mis-diagnosis; the actual + precondition is deterministic -- parent holds a context before the child spawns; + minimal repro in git history / scratch). ``nvidia-smi`` reads NVML without a + context, so the parent stays CUDA-free. + + Honours ``CUDA_VISIBLE_DEVICES`` by taking the FIRST visible index (the worker + pins to one device before this runs). Returns ``None`` on any failure so the + caller falls back to letting vLLM size itself. + """ + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + index = "0" + if visible is not None: + entries = [v.strip() for v in visible.split(",") if v.strip() != ""] + if not entries: + return None # device explicitly hidden -- no GPU to measure + index = entries[0] + try: + out = subprocess.run( + [ + "nvidia-smi", + f"--id={index}", + "--query-gpu=memory.free,memory.total", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=10, + ).stdout + free_mib, total_mib = (int(x.strip()) for x in out.splitlines()[0].split(",")) + except Exception: + return None + return free_mib * 2**20, total_mib * 2**20 + + +def gpu_mem_util(ceiling: float = GPU_MEM_UTIL, floor: float = GPU_MEM_UTIL_FLOOR): + """90% of the free fraction, capped at ``ceiling``. + + Reads free/total memory via ``nvidia-smi`` (:func:`_nvsmi_free_total_bytes`), + NOT ``torch.cuda`` -- this runs in the worker before ``LLM()`` spawns the + ``EngineCore`` child, and creating a CUDA context here would block that child + under ``Exclusive_Process`` (see :func:`_nvsmi_free_total_bytes`). + + Returns ``None`` when there is no visible GPU to measure (or the query fails), + so a caller can leave the kwarg unset and let vLLM pick. + + Raises ``SystemExit`` rather than returning a doomed value when the result is + below ``floor``: at that point something else is holding the card and a test + failure would be the wrong diagnosis. + """ + ft = _nvsmi_free_total_bytes() + if ft is None: + return None + free, total = ft + util = min(ceiling, 0.9 * free / total) + print( + f" GPU memory: {free / 2**30:.1f}/{total / 2**30:.1f} GiB free -> " + f"gpu_memory_utilization={util:.3f} (ceiling {ceiling:g})" + ) + if util < floor: + raise SystemExit( + f"FATAL: only {free / 2**30:.1f} GiB free on this device, which leaves " + f"gpu_memory_utilization={util:.3f} below the {floor:g} floor. Another " + "process is holding the card; this is not a test failure." + ) + return util + + +def device_is_free() -> bool: + """True if no compute process is holding a visible GPU, per ``nvidia-smi``. + + Queried through ``nvidia-smi`` on purpose, NOT a ``torch.cuda`` probe: this + runs in the **parent** pytest process, and under the host's + ``Exclusive_Process`` compute mode only + one process may hold a device's primary context at a time. A + ``torch.zeros(1, device=...)`` here would create exactly the parent context + that then blocks the next ``EngineCore`` child -- the very failure we are + guarding against. So we observe the device from the outside instead. + + Honours ``CUDA_VISIBLE_DEVICES``. Conservative: any error (no ``nvidia-smi``, + parse failure) returns ``True`` so a probe failure never wedges a run -- the + engine start itself remains the real gate. + """ + try: + out = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=gpu_uuid", + "--format=csv,noheader", + ], + capture_output=True, + text=True, + timeout=5, + ).stdout + except Exception: + return True + + # Map visible indices -> UUIDs so CUDA_VISIBLE_DEVICES is respected. If that + # query fails, fall back to "any running compute app means busy". + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + allowed_uuids = None + if visible is not None: + allowed = {v.strip() for v in visible.split(",") if v.strip() != ""} + try: + idx_out = subprocess.run( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ).stdout + allowed_uuids = set() + for line in idx_out.splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) == 2 and parts[0] in allowed: + allowed_uuids.add(parts[1]) + except Exception: + allowed_uuids = None + + for line in out.splitlines(): + uuid = line.strip() + if not uuid: + continue + if allowed_uuids is not None and uuid not in allowed_uuids: + continue + return False + return True + + +def wait_device_free(*, timeout: float = GPU_DEVICE_READY_TIMEOUT) -> None: + """Block until no compute process holds a visible GPU, or ``timeout`` elapses. + + Complements the memory-reclaim wait in :func:`shutdown_llm`. The observed + evidence is a device that reports its memory fully + free (~78.7/79.2 GiB) yet still fails a *new* engine's first host->device copy + with ``cudaErrorDevicesUnavailable``: the exiting ``EngineCore`` child has + released memory but not yet relinquished the device's primary context. Memory + free != device acquirable. We wait for the child to actually drop off + ``nvidia-smi``'s compute-app list before letting the next engine start. + + A timeout prints and returns rather than raising: like the memory wait, a + spurious red is worse than a bounded wait, and the engine start is still the + authoritative gate. + """ + deadline = time.monotonic() + timeout + while True: + if device_is_free(): + return + if time.monotonic() >= deadline: + print( + f" WARNING: a compute process still holds the device after " + f"{timeout:g}s; continuing and letting the engine start rule on it" + ) + return + time.sleep(1.0) + + +def shutdown_llm( + llm, + *, + target: float = GPU_MEM_RECLAIM_TARGET, + timeout: float = GPU_MEM_RECLAIM_TIMEOUT, + device_ready_timeout: float = GPU_DEVICE_READY_TIMEOUT, +) -> None: + """Stop an ``LLM``'s EngineCore and wait for the device memory to come back. + + ``del llm; gc.collect(); torch.cuda.empty_cache()`` does not do this, which is + why the suites that start engines back to back were tripping ``gpu_mem_util``'s + floor on a card with nothing else running: + + * The weights and KV cache belong to the **EngineCore child process** (v1 forks + one, and these runners set ``VLLM_WORKER_MULTIPROC_METHOD=spawn``). + ``empty_cache()`` drains the *parent's* allocator, which never held them. + * Teardown is registered as a ``weakref.finalize``, so ``del`` schedules it at + collection time rather than performing it, and nothing waits for the child. + + ``engine_core.shutdown()`` is the synchronous path: it chains to + ``CoreEngineProcManager.shutdown``, which terminates *and joins* the children. + Its ``shutdown(timeout=None)`` signature is byte-identical on 0.26 - 0.30, and + it is reached defensively here so a future rename degrades to today's + behaviour rather than an ``AttributeError``. + + The caller's ``llm`` reference does not have to be dropped before calling: + once the child has exited the driver owns the reclaim, so refcounts in this + process no longer gate it. + + A timeout prints and returns instead of raising. ``gpu_mem_util`` is the one + place that decides a card is too full to use, and two components disagreeing + about that is worse than waiting a bounded time and letting it rule. + + After the memory wait, this also waits for the **device** itself to become + acquirable again (:func:`wait_device_free`): under ``Exclusive_Process`` the + exiting ``EngineCore`` child can return its memory before it relinquishes the + device's primary context, so the next engine's first CUDA op fails with + ``cudaErrorDevicesUnavailable`` even on a card that reports fully free. + """ + engine_core = getattr(getattr(llm, "llm_engine", None), "engine_core", None) + shutdown = getattr(engine_core, "shutdown", None) + if shutdown is not None: + shutdown() + + gc.collect() + + import torch + + if not torch.cuda.is_available(): + return + torch.cuda.empty_cache() + + deadline = time.monotonic() + timeout + while True: + free, total = torch.cuda.mem_get_info() + frac = free / total + if frac >= target: + print( + f" GPU memory reclaimed: {free / 2**30:.1f}/{total / 2**30:.1f} GiB " + f"free ({frac:.2f} >= {target:g})" + ) + break + if time.monotonic() >= deadline: + print( + f" WARNING: {free / 2**30:.1f}/{total / 2**30:.1f} GiB free " + f"({frac:.2f}) still below {target:g} after {timeout:g}s; " + "continuing and letting gpu_mem_util rule on it" + ) + break + time.sleep(1.0) + + # Memory can be back while the device is momentarily unacquirable by a NEW + # CUDA context -- the device-availability race in cluster 2. Wait that out too. + wait_device_free(timeout=device_ready_timeout) diff --git a/tests/shared/vllm_kv_cache.py b/tests/shared/vllm_kv_cache.py new file mode 100644 index 00000000..4aa1c735 --- /dev/null +++ b/tests/shared/vllm_kv_cache.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Allocate a real KV cache for an :class:`Attention` layer, on any vLLM 0.26+. + +These harnesses hand the tensor to a real FlashAttention kernel (they build real +``FlashAttentionMetadata``: slot_mapping, block_table, query_start_loc, seq_lens), +so the shape has to be genuinely right, not merely plausible — a wrong one reads +back garbage instead of raising. + +Two APIs, because vLLM 0.29 deleted the first: + +* 0.26 - 0.28: ``attn_backend.get_kv_cache_shape(num_blocks, block_size, + num_kv_heads, head_size)``, a per-backend staticmethod. 29 definitions at 0.28, + ZERO at 0.29 (0.30 keeps one, on flashinfer_mla_sparse_sm90). +* 0.29+: ``compute_layer_kv_cache_shape_bytes(spec, num_blocks)``, a module + function returning ``(B, H, N, C)`` with C in BYTES rather than elements. + +The two are provably equal for dense bf16 attention, which is why this is a port +and not a guess. FlashAttention returns +``(num_blocks, num_kv_heads, block_size, 2 * head_size)``; for a +``FullAttentionSpec`` with no packing overrides, ``num_heads == num_kv_heads``, +``get_num_kernel_states(bs) == bs`` (``tokens_per_state == 1``), and +``state_content_size_bytes == (head_size + head_size_v) * itemsize``, so dividing +the last dim by the itemsize reproduces the old tuple exactly. + +No shipped version offers both APIs for FlashAttention — 0.26-0.28 have only the +backend staticmethod, 0.29+ only the spec function — so that equality is an +argument from the source above, not something a run can confirm. The cross-check +below therefore never fires today; it is there so that a future version which +reintroduces the staticmethod tells us immediately if the two ever disagree. +""" + +import torch + + +def kv_cache_shape(attn, vllm_config, num_blocks: int, block_size: int): + """The KV-cache shape, in ELEMENTS, for one ``Attention`` layer.""" + from vllm.utils.torch_utils import get_dtype_size + + try: + from vllm.v1.kv_cache_interface import compute_layer_kv_cache_shape_bytes + except ImportError: # 0.26 - 0.28 + compute_layer_kv_cache_shape_bytes = None + + legacy = None + if hasattr(attn.attn_backend, "get_kv_cache_shape"): # 0.26 - 0.28 + legacy = tuple( + attn.attn_backend.get_kv_cache_shape( + num_blocks, + block_size, + attn.num_kv_heads, + attn.head_size, + ) + ) + if compute_layer_kv_cache_shape_bytes is None: + return legacy + assert compute_layer_kv_cache_shape_bytes is not None, ( + "this vLLM has neither attn_backend.get_kv_cache_shape nor " + "compute_layer_kv_cache_shape_bytes" + ) + + spec = attn.get_kv_cache_spec(vllm_config) + assert spec is not None, ( + f"{attn.layer_name} contributes no KV cache; it cannot be set up here" + ) + # compute_layer_kv_cache_shape_bytes derives the leading dim from + # spec.block_size (which tracks vllm_config.cache_config.block_size), so a + # harness block_size that disagrees would silently shift every block index. + assert spec.block_size == block_size, ( + f"harness block_size={block_size} but the spec says {spec.block_size}; " + "set cache_config.block_size to match" + ) + *lead, content_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks) + itemsize = get_dtype_size(spec.dtype) + assert content_bytes % itemsize == 0, ( + f"{content_bytes} bytes per cell is not a whole number of " + f"{spec.dtype} elements; this layer's cache is packed, not dense" + ) + shape = (*lead, content_bytes // itemsize) + + if legacy is not None: + assert shape == legacy, ( + f"the 0.29+ spec path gives {shape} but the backend's own " + f"get_kv_cache_shape gives {legacy}" + ) + return shape + + +def setup_kv_cache(attn, vllm_config, num_blocks: int, block_size: int, device): + """Attach a zeroed bf16 KV cache to ``attn`` and return it.""" + attn.kv_cache_torch_dtype = torch.bfloat16 + shape = kv_cache_shape(attn, vllm_config, num_blocks, block_size) + kv_cache = torch.zeros(shape, device=device, dtype=torch.bfloat16) + attn.kv_cache = kv_cache + return kv_cache diff --git a/tests/unit/test_asr.py b/tests/unit/test_asr.py index 38c45a57..7d26d750 100644 --- a/tests/unit/test_asr.py +++ b/tests/unit/test_asr.py @@ -135,7 +135,7 @@ def _fake_pipe_transcriber(self): def test_self_chunks_true_is_single_call(self): t = asr.ASRTranscriber(model_id="x", device="cpu") calls = [] - t._pipeline = lambda inp, **k: (calls.append(len(inp["raw"])) or {"text": "x"}) + t._pipeline = lambda inp, **k: calls.append(len(inp["raw"])) or {"text": "x"} # 70s clip; with self_chunks the whole thing goes in one call. t.transcribe( np.zeros(70 * 16000, dtype=np.float32), @@ -558,3 +558,68 @@ def test_pipeline_kwargs_override_defaults(self): assert kwargs["task"] == "automatic-speech-recognition" assert kwargs["chunk_length_s"] == 15 # overrode the default 30 assert kwargs["batch_size"] == 4 # extra kwarg passed through + + +def _smi(csv_text): + """mock.patch target: a subprocess.run result whose stdout is ``csv_text``.""" + return mock.Mock(return_value=mock.Mock(stdout=csv_text)) + + +class TestExclusiveProcessDetection: + """gpu_is_exclusive_process reads nvidia-smi WITHOUT creating a CUDA context. + + The driver-side ASR transcriber must decide CPU-vs-CUDA before touching the + device: under Exclusive_Process the vLLM EngineCore subprocess already owns + it, so a torch.cuda probe in the driver would be the very fault we avoid. + """ + + def test_exclusive_mode_detected(self): + with mock.patch.object(asr.subprocess, "run", _smi("0, Exclusive_Process\n")): + assert asr.gpu_is_exclusive_process() is True + + def test_default_mode_is_not_exclusive(self): + with mock.patch.object(asr.subprocess, "run", _smi("0, Default\n")): + assert asr.gpu_is_exclusive_process() is False + + def test_nvidia_smi_missing_is_conservative_false(self): + with mock.patch.object( + asr.subprocess, "run", mock.Mock(side_effect=FileNotFoundError) + ): + assert asr.gpu_is_exclusive_process() is False + + def test_only_visible_devices_are_considered(self, monkeypatch): + # GPU 1 is exclusive but not visible; the visible GPU 0 is Default. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + with mock.patch.object( + asr.subprocess, + "run", + _smi("0, Default\n1, Exclusive_Process\n"), + ): + assert asr.gpu_is_exclusive_process() is False + + def test_visible_exclusive_device_is_seen(self, monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + with mock.patch.object( + asr.subprocess, + "run", + _smi("0, Default\n1, Exclusive_Process\n"), + ): + assert asr.gpu_is_exclusive_process() is True + + +class TestResolveDriverDevice: + """cuda -> cpu only under Exclusive_Process; everything else passes through.""" + + def test_cuda_downgraded_under_exclusive(self): + with mock.patch.object(asr, "gpu_is_exclusive_process", lambda: True): + assert asr.resolve_driver_device("cuda") == "cpu" + + def test_cuda_kept_in_default_mode(self): + with mock.patch.object(asr, "gpu_is_exclusive_process", lambda: False): + assert asr.resolve_driver_device("cuda") == "cuda" + + def test_explicit_cpu_never_probes_and_passes_through(self): + # An explicit non-cuda device must not even query compute mode. + probe = mock.Mock(side_effect=AssertionError("should not be called")) + with mock.patch.object(asr, "gpu_is_exclusive_process", probe): + assert asr.resolve_driver_device("cpu") == "cpu" diff --git a/tests/vllm/_fa3_schedule_worker.py b/tests/vllm/_fa3_schedule_worker.py new file mode 100644 index 00000000..5ff10759 --- /dev/null +++ b/tests/vllm/_fa3_schedule_worker.py @@ -0,0 +1,549 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Subprocess worker: FA3 AOT-schedule correctness under FULL CUDA graphs + prefix caching. + +Reproduces and gates the bug tracked in +issue #139: on vLLM 0.26, a +Shadow-Residual (SR) MultiSwitch checkpoint served with +``cudagraph_mode=FULL`` AND prefix caching ON emits tokens outside the +adapter's allowed set and distributions that diverge from the reference — but +only on small, prefix-cached decode steps. Turning off EITHER FULL or prefix +caching matches the reference exactly. + +Root cause (summary): FA3's ahead-of-time schedule is sized from the model +config's attention shape, but SR's decoder doubles its query heads +(``Attention(2 * num_heads, ...)``), so ``num_heads_q`` is half the real value +for that group. Under FULL the mis-sized schedule is frozen into the captured +graph; a 1-token prefix-cached decode is where it goes wrong. See the doc for +the full trace. + +Each mode is a separate subprocess so only one vLLM model is ever resident on +GPU at a time:: + + python worker.py probe + python worker.py build --work-dir + python worker.py run --work-dir --tag --cudagraph --prefix + python worker.py compare --work-dir --ref --cand --label