Bugfix/139 fa3 schedule full cudagraph prefix caching fa3 fixing

.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

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

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

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
 

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

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]

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

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``.