From 0668459e33d002200b0f841ca6aaaff6eeac0660 Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Mon, 18 May 2026 19:14:25 +0300 Subject: [PATCH 1/6] Force compatible attention backend in vLLM tests on non-Hopper GPUs (issue #36) NOTE: This commit may be temporary. The underlying fix is to rebuild vllm-flash-attn with TORCH_CUDA_ARCH_LIST covering the runtime GPU. Until that happens, this works around a CI/test environment where the installed vllm-flash-attn ships only Hopper (SM 9.0) kernels but the test GPU is Ampere (SM 8.0), which crashed the worker with "no kernel image is available for execution on the device" before any test could assert. Changes: - New helper tests/shared/vllm_attn_backend.py monkey-patches vLLM's attention selector to return FLASHINFER when device capability < (9,0). Override gate via GS_TEST_FORCE_ATTN_BACKEND, backend choice via GS_TEST_ATTN_BACKEND. No-op on Hopper+ so FA3 path stays under test. - _single_switch_worker.py: applies the override around SingleSwitch construction, plus a fail-fast _probe_attention() that runs a 1-token forward at startup and emits a structured fatal JSON if the kernel is unusable, so we get one clear failure instead of cascading BrokenPipeErrors. - test_single_switch.py: _ensure_worker recognizes the fatal JSON, drains stderr, and pytest.fails each test with the same actionable message via a sticky _fatal_startup_error flag. - _model_forward_tests.py / _position_zero_nan_tests.py: install the override at module import (one-shot, subprocess-scoped). --- tests/shared/vllm_attn_backend.py | 106 +++++++++++++++++++++++++ tests/vllm/_model_forward_tests.py | 6 ++ tests/vllm/_position_zero_nan_tests.py | 6 ++ tests/vllm/_single_switch_worker.py | 59 +++++++++++++- tests/vllm/test_single_switch.py | 27 ++++++- 5 files changed, 202 insertions(+), 2 deletions(-) create mode 100644 tests/shared/vllm_attn_backend.py diff --git a/tests/shared/vllm_attn_backend.py b/tests/shared/vllm_attn_backend.py new file mode 100644 index 0000000..2db736a --- /dev/null +++ b/tests/shared/vllm_attn_backend.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Force a compatible attention backend in vLLM tests on non-Hopper GPUs. + +The CI/test environment has a `vllm-flash-attn` build that ships only Hopper +(SM 9.0) kernels. On Ampere/Ada GPUs vLLM still auto-selects ``FLASH_ATTN`` +because validation passes (the build is present and importable), but the very +first kernel launch crashes with:: + + CUDA error (.../hopper/flash_api.cpp:697): + no kernel image is available for execution on the device + +That C-level abort kills the worker process before any Python handler runs, +which is what produces the cascade of ``BrokenPipeError``s in the long-lived +SingleSwitch worker and the opaque "Worker died unexpectedly" message in the +short-lived per-class subprocess tests. + +Until the environment is fixed (rebuild ``vllm-flash-attn`` with +``TORCH_CUDA_ARCH_LIST`` covering the runtime GPU), force a backend whose +kernels are universally compiled. ``FLASHINFER`` validates cleanly for +head_dim=64 / bf16 on SM 8.0 and is the default fallback. ``TRITON_ATTN`` and +``FLEX_ATTENTION`` are also valid choices and can be selected via env. + +Activation rules +---------------- +- Only fires when ``torch.cuda.get_device_capability() < (9, 0)``. On Hopper+, + the FA3 build matches the GPU and we leave selection alone so the FA3 path + remains under test. +- Override the gate (e.g. on a CI runner with a known-broken FA build on + Hopper) by setting ``GS_TEST_FORCE_ATTN_BACKEND=1``. +- Pick a different backend with ``GS_TEST_ATTN_BACKEND=TRITON_ATTN`` (or + ``FLEX_ATTENTION``). Default is ``FLASHINFER``. + +The helper is a no-op outside vLLM/CUDA so importing it from CPU-only test +contexts doesn't break collection. +""" + +from __future__ import annotations + +import os +import sys +from typing import Callable + + +def _should_force(force_flag: str | None, capability: tuple[int, int] | None) -> bool: + if force_flag and force_flag.lower() not in ("0", "false", ""): + return True + if capability is None: + return False + return capability < (9, 0) + + +def force_compatible_attn_backend() -> Callable[[], None] | None: + """Monkey-patch vLLM's attention selector to a fixed, validated backend. + + Returns a restorer callable that undoes the patch, or ``None`` if no patch + was applied (Hopper+, or vLLM not importable, or no CUDA). + """ + try: + import torch + except ImportError: + return None + if not torch.cuda.is_available(): + return None + + capability = torch.cuda.get_device_capability() + force_flag = os.environ.get("GS_TEST_FORCE_ATTN_BACKEND") + if not _should_force(force_flag, capability): + return None + + backend_name = os.environ.get("GS_TEST_ATTN_BACKEND", "FLASHINFER") + + try: + from vllm.v1.attention import selector as _sel + from vllm.v1.attention.backends.registry import AttentionBackendEnum + except ImportError: + return None + + try: + backend_enum = AttentionBackendEnum[backend_name] + except KeyError as exc: + valid = ", ".join(b.name for b in AttentionBackendEnum) + raise ValueError( + f"GS_TEST_ATTN_BACKEND={backend_name!r} is not a known vLLM " + f"attention backend. Valid: {valid}" + ) from exc + + forced_cls = backend_enum.get_class() + original = _sel.get_attn_backend + + def _patched_get_attn_backend(*args, **kwargs): + return forced_cls + + _sel.get_attn_backend = _patched_get_attn_backend + + sys.stderr.write( + f"[vllm_attn_backend] GPU compute capability {capability} < (9, 0): " + f"forcing {backend_name} attention backend (set " + f"GS_TEST_FORCE_ATTN_BACKEND=0 to disable, GS_TEST_ATTN_BACKEND= " + f"to pick a different one).\n" + ) + sys.stderr.flush() + + def _restore() -> None: + _sel.get_attn_backend = original + + return _restore diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index 285a6c6..8537d62 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -42,6 +42,12 @@ def _try_import_vllm(): from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM from granite_switch.vllm.switch.single import SingleSwitch + # Force a working attention backend on non-Hopper GPUs (FA3 build is + # Hopper-only in this environment). One-shot patch for the whole module + # — the subprocess exits when this file's tests finish so no restore. + from tests.shared.vllm_attn_backend import force_compatible_attn_backend + force_compatible_attn_backend() + # ── Constants ──────────────────────────────────────────────────────── BLOCK_SIZE = 16 diff --git a/tests/vllm/_position_zero_nan_tests.py b/tests/vllm/_position_zero_nan_tests.py index 8b5fb90..78f8131 100644 --- a/tests/vllm/_position_zero_nan_tests.py +++ b/tests/vllm/_position_zero_nan_tests.py @@ -47,6 +47,12 @@ def _try_import_vllm(): from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM from granite_switch.vllm.core.decoder import GraniteLoRAEmbeddedAttention + # Force a working attention backend on non-Hopper GPUs (FA3 build is + # Hopper-only in this environment). One-shot patch for the whole module + # — the subprocess exits when this file's tests finish so no restore. + from tests.shared.vllm_attn_backend import force_compatible_attn_backend + force_compatible_attn_backend() + from tests.shared.vllm_distributed import ensure_distributed as _ensure_distributed BLOCK_SIZE = 16 diff --git a/tests/vllm/_single_switch_worker.py b/tests/vllm/_single_switch_worker.py index 0d9d0f9..b5779d1 100644 --- a/tests/vllm/_single_switch_worker.py +++ b/tests/vllm/_single_switch_worker.py @@ -3,6 +3,11 @@ Protocol (JSON-line over stdin/stdout): Startup: prints {"ready": true, "backend_name": "..."} to stdout + OR prints {"fatal": "...", "hint": "...", "backend_name": "..."} and exits + when the auto-selected attention backend's kernel image is incompatible + with this GPU (e.g. FA3 Hopper-only kernels on a non-Hopper card). The + parent reads this and converts it into one clear pytest failure instead + of letting hundreds of subsequent tests cascade into BrokenPipeErrors. Request: {"seq": [...], "num_adapters": N, "control_token_gain": G} Response: {"result": [...]} Error: {"error": "..."} @@ -22,6 +27,7 @@ import torch +from tests.shared.vllm_attn_backend import force_compatible_attn_backend from tests.shared.vllm_distributed import ensure_distributed @@ -57,6 +63,7 @@ def _setup(): old_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.bfloat16) + restore_attn_backend = force_compatible_attn_backend() try: vllm_config = VllmConfig() ensure_distributed(vllm_config) @@ -69,6 +76,8 @@ def _setup(): ) finally: torch.set_default_dtype(old_dtype) + if restore_attn_backend is not None: + restore_attn_backend() attn = switch.attn attn.kv_cache_torch_dtype = torch.bfloat16 @@ -232,8 +241,56 @@ def _query_geometry(harness): } +def _probe_attention(harness): + """Run a 1-token forward to verify the auto-selected attention kernel is usable. + + vLLM picks FLASH_ATTN by default. If the installed vllm-flash-attn was built + only for Hopper (SM 9.0) but this GPU is a different architecture, the very + first kernel launch raises "no kernel image is available for execution on + the device" — and in some builds this kills the worker process at the C + layer before any Python handler runs. Catching it here, on a tiny synthetic + input, lets us emit a structured fatal message before signaling ready. + + Returns None on success, or a dict {"fatal": ..., "hint": ...} on failure. + """ + try: + _run(harness, seq=[0], num_adapters=1, control_token_gain=15.0) + except Exception as exc: + msg = f"{type(exc).__name__}: {exc}" + hint = ( + "Auto-selected attention backend " + f"{harness['backend_name']!r} cannot launch on this GPU. " + "If 'no kernel image is available' appears above, the installed " + "vllm-flash-attn was built for a different SM than the runtime GPU. " + "Either rebuild vllm-flash-attn with TORCH_CUDA_ARCH_LIST covering " + "this GPU, or set VLLM_ATTENTION_BACKEND to FLASHINFER / TRITON_ATTN " + "/ FLEX_ATTENTION before running the suite." + ) + return {"fatal": msg, "hint": hint, "backend_name": harness["backend_name"]} + return None + + def main(): - harness = _setup() + try: + harness = _setup() + except Exception as exc: + # Setup itself blew up — surface it on stdout so the parent can show a + # clean failure instead of "Worker failed to start: ". + msg = { + "fatal": f"{type(exc).__name__}: {exc}", + "hint": "Worker setup failed before attention probe.", + "backend_name": "unknown", + } + sys.stdout.write(json.dumps(msg) + "\n") + sys.stdout.flush() + traceback.print_exc(file=sys.stderr) + return + + probe_failure = _probe_attention(harness) + if probe_failure is not None: + sys.stdout.write(json.dumps(probe_failure) + "\n") + sys.stdout.flush() + return # Signal ready ready_msg = {"ready": True, "backend_name": harness["backend_name"]} diff --git a/tests/vllm/test_single_switch.py b/tests/vllm/test_single_switch.py index da93238..e73e2c8 100644 --- a/tests/vllm/test_single_switch.py +++ b/tests/vllm/test_single_switch.py @@ -47,11 +47,19 @@ _worker_proc = None _worker_lock = threading.Lock() _backend_name = None +# Sticky error: once the worker reports a fatal startup condition (e.g. a +# kernel-image mismatch in the auto-selected attention backend), we keep the +# message here and re-raise it on every subsequent test instead of re-spawning +# a worker that will hit the same crash. This converts what used to be ~hundreds +# of cascading BrokenPipeErrors into one clear, actionable failure per test. +_fatal_startup_error = None def _ensure_worker(): """Lazily start the long-lived worker subprocess.""" - global _worker_proc, _backend_name + global _worker_proc, _backend_name, _fatal_startup_error + if _fatal_startup_error is not None: + pytest.fail(_fatal_startup_error, pytrace=False) if _worker_proc is not None and _worker_proc.poll() is None: return proc = subprocess.Popen( @@ -68,6 +76,23 @@ def _ensure_worker(): stderr = proc.stderr.read() raise RuntimeError(f"Worker failed to start:\n{stderr}") ready = json.loads(ready_line) + if "fatal" in ready: + # Worker detected a non-recoverable environment problem (e.g. the + # auto-selected attention backend has no kernel image for this GPU). + # Drain stderr for context, mark this poison sticky, and fail. + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + stderr_tail = (proc.stderr.read() or "")[-2000:] + backend = ready.get("backend_name", "unknown") + _fatal_startup_error = ( + f"vLLM worker cannot start: {ready['fatal']}\n" + f"Backend: {backend}\n" + f"Hint: {ready.get('hint', '')}\n" + f"--- worker stderr (tail) ---\n{stderr_tail}" + ) + pytest.fail(_fatal_startup_error, pytrace=False) assert ready.get("ready"), f"Unexpected ready message: {ready}" _backend_name = ready.get("backend_name", "unknown") _worker_proc = proc From f8aa8c9380e68c6e7b5688b299f0c76df8b3051d Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Mon, 18 May 2026 19:28:54 +0300 Subject: [PATCH 2/6] Patch get_attn_backend across all vLLM consumer modules, not just selector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit (Still tmp — same caveat as the parent commit on this branch.) Previous attempt only patched vllm.v1.attention.selector.get_attn_backend, but every consumer (Attention, mla_attention, etc.) does "from vllm.v1.attention.selector import get_attn_backend" at module import time, so the consumer's already-bound local reference was untouched and vLLM kept auto-picking FLASH_ATTN. Now iterate the known consumer modules and replace get_attn_backend in each module's namespace. The restorer reverts all of them. If the list of consumers ever drifts, a stderr warning is emitted instead of a silent no-op. --- tests/shared/vllm_attn_backend.py | 48 +++++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/tests/shared/vllm_attn_backend.py b/tests/shared/vllm_attn_backend.py index 2db736a..55e4df6 100644 --- a/tests/shared/vllm_attn_backend.py +++ b/tests/shared/vllm_attn_backend.py @@ -85,22 +85,58 @@ def force_compatible_attn_backend() -> Callable[[], None] | None: ) from exc forced_cls = backend_enum.get_class() - original = _sel.get_attn_backend def _patched_get_attn_backend(*args, **kwargs): return forced_cls - _sel.get_attn_backend = _patched_get_attn_backend + # vLLM's Attention layer and friends do + # from vllm.v1.attention.selector import get_attn_backend + # at module import time, so patching only `selector.get_attn_backend` + # leaves their already-bound local references untouched. Patch each + # known consumer's module namespace too. Missing modules (older/newer + # vLLM) are ignored — at least one of these will be the live one. + consumer_modules = [ + "vllm.v1.attention.selector", + "vllm.model_executor.layers.attention.attention", + "vllm.model_executor.layers.attention.chunked_local_attention", + "vllm.model_executor.layers.attention.cross_attention", + "vllm.model_executor.layers.attention.encoder_only_attention", + "vllm.model_executor.layers.attention.mla_attention", + "vllm.model_executor.layers.attention.static_sink_attention", + ] + + import importlib + + originals: list[tuple[object, object]] = [] + for mod_name in consumer_modules: + try: + mod = importlib.import_module(mod_name) + except ImportError: + continue + if hasattr(mod, "get_attn_backend"): + originals.append((mod, mod.get_attn_backend)) + mod.get_attn_backend = _patched_get_attn_backend + + if not originals: + # Nothing to patch — vLLM internals must have moved. + sys.stderr.write( + "[vllm_attn_backend] WARNING: no get_attn_backend symbols found " + "in known vLLM modules; backend override is a no-op. The test " + "will likely still hit the FA3 kernel-image crash.\n" + ) + sys.stderr.flush() + return None sys.stderr.write( f"[vllm_attn_backend] GPU compute capability {capability} < (9, 0): " - f"forcing {backend_name} attention backend (set " - f"GS_TEST_FORCE_ATTN_BACKEND=0 to disable, GS_TEST_ATTN_BACKEND= " - f"to pick a different one).\n" + f"forcing {backend_name} attention backend in " + f"{len(originals)} vLLM module(s) (set GS_TEST_FORCE_ATTN_BACKEND=0 " + f"to disable, GS_TEST_ATTN_BACKEND= to pick a different one).\n" ) sys.stderr.flush() def _restore() -> None: - _sel.get_attn_backend = original + for mod, original in originals: + mod.get_attn_backend = original return _restore From 8c4e749a7c20954c64f3fca6972dd2544b4fb619 Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Mon, 18 May 2026 20:13:02 +0300 Subject: [PATCH 3/6] Drop attention-backend override; keep only fail-fast diagnostics The override was solving the wrong problem. Root cause turned out to be a stale standalone 'vllm-flash-attn' PyPI package whose Hopper-only .so was being imported in preference to vLLM's bundled FA kernels. Once that package is uninstalled, vLLM's get_flash_attn_version() correctly returns 2 on Ampere/A100 and the bundled FA2 kernel runs. Removing: - tests/shared/vllm_attn_backend.py (and all its imports/calls in worker and inner pytest scripts) Keeping (the genuinely useful part): - _probe_attention() in _single_switch_worker.py: 1-token forward at startup, structured fatal JSON on failure - sticky _fatal_startup_error in test_single_switch.py: one clear pytest.fail per test instead of cascading BrokenPipeErrors when the worker can't start Also updated the worker's hint message to mention the actual fix (uninstall stale vllm-flash-attn) instead of the dead-end workaround. --- tests/shared/vllm_attn_backend.py | 142 ------------------------- tests/vllm/_model_forward_tests.py | 6 -- tests/vllm/_position_zero_nan_tests.py | 6 -- tests/vllm/_single_switch_worker.py | 16 ++- 4 files changed, 6 insertions(+), 164 deletions(-) delete mode 100644 tests/shared/vllm_attn_backend.py diff --git a/tests/shared/vllm_attn_backend.py b/tests/shared/vllm_attn_backend.py deleted file mode 100644 index 55e4df6..0000000 --- a/tests/shared/vllm_attn_backend.py +++ /dev/null @@ -1,142 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Force a compatible attention backend in vLLM tests on non-Hopper GPUs. - -The CI/test environment has a `vllm-flash-attn` build that ships only Hopper -(SM 9.0) kernels. On Ampere/Ada GPUs vLLM still auto-selects ``FLASH_ATTN`` -because validation passes (the build is present and importable), but the very -first kernel launch crashes with:: - - CUDA error (.../hopper/flash_api.cpp:697): - no kernel image is available for execution on the device - -That C-level abort kills the worker process before any Python handler runs, -which is what produces the cascade of ``BrokenPipeError``s in the long-lived -SingleSwitch worker and the opaque "Worker died unexpectedly" message in the -short-lived per-class subprocess tests. - -Until the environment is fixed (rebuild ``vllm-flash-attn`` with -``TORCH_CUDA_ARCH_LIST`` covering the runtime GPU), force a backend whose -kernels are universally compiled. ``FLASHINFER`` validates cleanly for -head_dim=64 / bf16 on SM 8.0 and is the default fallback. ``TRITON_ATTN`` and -``FLEX_ATTENTION`` are also valid choices and can be selected via env. - -Activation rules ----------------- -- Only fires when ``torch.cuda.get_device_capability() < (9, 0)``. On Hopper+, - the FA3 build matches the GPU and we leave selection alone so the FA3 path - remains under test. -- Override the gate (e.g. on a CI runner with a known-broken FA build on - Hopper) by setting ``GS_TEST_FORCE_ATTN_BACKEND=1``. -- Pick a different backend with ``GS_TEST_ATTN_BACKEND=TRITON_ATTN`` (or - ``FLEX_ATTENTION``). Default is ``FLASHINFER``. - -The helper is a no-op outside vLLM/CUDA so importing it from CPU-only test -contexts doesn't break collection. -""" - -from __future__ import annotations - -import os -import sys -from typing import Callable - - -def _should_force(force_flag: str | None, capability: tuple[int, int] | None) -> bool: - if force_flag and force_flag.lower() not in ("0", "false", ""): - return True - if capability is None: - return False - return capability < (9, 0) - - -def force_compatible_attn_backend() -> Callable[[], None] | None: - """Monkey-patch vLLM's attention selector to a fixed, validated backend. - - Returns a restorer callable that undoes the patch, or ``None`` if no patch - was applied (Hopper+, or vLLM not importable, or no CUDA). - """ - try: - import torch - except ImportError: - return None - if not torch.cuda.is_available(): - return None - - capability = torch.cuda.get_device_capability() - force_flag = os.environ.get("GS_TEST_FORCE_ATTN_BACKEND") - if not _should_force(force_flag, capability): - return None - - backend_name = os.environ.get("GS_TEST_ATTN_BACKEND", "FLASHINFER") - - try: - from vllm.v1.attention import selector as _sel - from vllm.v1.attention.backends.registry import AttentionBackendEnum - except ImportError: - return None - - try: - backend_enum = AttentionBackendEnum[backend_name] - except KeyError as exc: - valid = ", ".join(b.name for b in AttentionBackendEnum) - raise ValueError( - f"GS_TEST_ATTN_BACKEND={backend_name!r} is not a known vLLM " - f"attention backend. Valid: {valid}" - ) from exc - - forced_cls = backend_enum.get_class() - - def _patched_get_attn_backend(*args, **kwargs): - return forced_cls - - # vLLM's Attention layer and friends do - # from vllm.v1.attention.selector import get_attn_backend - # at module import time, so patching only `selector.get_attn_backend` - # leaves their already-bound local references untouched. Patch each - # known consumer's module namespace too. Missing modules (older/newer - # vLLM) are ignored — at least one of these will be the live one. - consumer_modules = [ - "vllm.v1.attention.selector", - "vllm.model_executor.layers.attention.attention", - "vllm.model_executor.layers.attention.chunked_local_attention", - "vllm.model_executor.layers.attention.cross_attention", - "vllm.model_executor.layers.attention.encoder_only_attention", - "vllm.model_executor.layers.attention.mla_attention", - "vllm.model_executor.layers.attention.static_sink_attention", - ] - - import importlib - - originals: list[tuple[object, object]] = [] - for mod_name in consumer_modules: - try: - mod = importlib.import_module(mod_name) - except ImportError: - continue - if hasattr(mod, "get_attn_backend"): - originals.append((mod, mod.get_attn_backend)) - mod.get_attn_backend = _patched_get_attn_backend - - if not originals: - # Nothing to patch — vLLM internals must have moved. - sys.stderr.write( - "[vllm_attn_backend] WARNING: no get_attn_backend symbols found " - "in known vLLM modules; backend override is a no-op. The test " - "will likely still hit the FA3 kernel-image crash.\n" - ) - sys.stderr.flush() - return None - - sys.stderr.write( - f"[vllm_attn_backend] GPU compute capability {capability} < (9, 0): " - f"forcing {backend_name} attention backend in " - f"{len(originals)} vLLM module(s) (set GS_TEST_FORCE_ATTN_BACKEND=0 " - f"to disable, GS_TEST_ATTN_BACKEND= to pick a different one).\n" - ) - sys.stderr.flush() - - def _restore() -> None: - for mod, original in originals: - mod.get_attn_backend = original - - return _restore diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index 8537d62..285a6c6 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -42,12 +42,6 @@ def _try_import_vllm(): from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM from granite_switch.vllm.switch.single import SingleSwitch - # Force a working attention backend on non-Hopper GPUs (FA3 build is - # Hopper-only in this environment). One-shot patch for the whole module - # — the subprocess exits when this file's tests finish so no restore. - from tests.shared.vllm_attn_backend import force_compatible_attn_backend - force_compatible_attn_backend() - # ── Constants ──────────────────────────────────────────────────────── BLOCK_SIZE = 16 diff --git a/tests/vllm/_position_zero_nan_tests.py b/tests/vllm/_position_zero_nan_tests.py index 78f8131..8b5fb90 100644 --- a/tests/vllm/_position_zero_nan_tests.py +++ b/tests/vllm/_position_zero_nan_tests.py @@ -47,12 +47,6 @@ def _try_import_vllm(): from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM from granite_switch.vllm.core.decoder import GraniteLoRAEmbeddedAttention - # Force a working attention backend on non-Hopper GPUs (FA3 build is - # Hopper-only in this environment). One-shot patch for the whole module - # — the subprocess exits when this file's tests finish so no restore. - from tests.shared.vllm_attn_backend import force_compatible_attn_backend - force_compatible_attn_backend() - from tests.shared.vllm_distributed import ensure_distributed as _ensure_distributed BLOCK_SIZE = 16 diff --git a/tests/vllm/_single_switch_worker.py b/tests/vllm/_single_switch_worker.py index b5779d1..a53a1fc 100644 --- a/tests/vllm/_single_switch_worker.py +++ b/tests/vllm/_single_switch_worker.py @@ -27,7 +27,6 @@ import torch -from tests.shared.vllm_attn_backend import force_compatible_attn_backend from tests.shared.vllm_distributed import ensure_distributed @@ -63,7 +62,6 @@ def _setup(): old_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.bfloat16) - restore_attn_backend = force_compatible_attn_backend() try: vllm_config = VllmConfig() ensure_distributed(vllm_config) @@ -76,8 +74,6 @@ def _setup(): ) finally: torch.set_default_dtype(old_dtype) - if restore_attn_backend is not None: - restore_attn_backend() attn = switch.attn attn.kv_cache_torch_dtype = torch.bfloat16 @@ -259,12 +255,12 @@ def _probe_attention(harness): msg = f"{type(exc).__name__}: {exc}" hint = ( "Auto-selected attention backend " - f"{harness['backend_name']!r} cannot launch on this GPU. " - "If 'no kernel image is available' appears above, the installed " - "vllm-flash-attn was built for a different SM than the runtime GPU. " - "Either rebuild vllm-flash-attn with TORCH_CUDA_ARCH_LIST covering " - "this GPU, or set VLLM_ATTENTION_BACKEND to FLASHINFER / TRITON_ATTN " - "/ FLEX_ATTENTION before running the suite." + f"{harness['backend_name']!r} crashed during the startup smoke " + "test. If 'no kernel image is available' appears above, the FA " + "kernels in this venv were compiled for a different SM than the " + "runtime GPU. Common cause: a stale standalone 'vllm-flash-attn' " + "PyPI package shadowing vLLM's bundled FA kernels — try " + "`pip uninstall vllm-flash-attn` and re-run." ) return {"fatal": msg, "hint": hint, "backend_name": harness["backend_name"]} return None From 2ca8574e21135a3ba6e827322724d81495b3ac36 Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Mon, 18 May 2026 20:20:13 +0300 Subject: [PATCH 4/6] Gate scheduler_metadata on FA version 3 (fixes FA3 crash on A100) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scheduler_metadata is an FA3-only kwarg. The three test workers were computing and passing it whenever the auto-selected backend was FLASH_ATTN, regardless of which FA version vLLM had picked for the GPU. On A100 (SM 8.0) vLLM picks FA2, but supplying scheduler_metadata forces dispatch into the FA3 Hopper-only kernel inside _vllm_fa3_C.abi3.so, which crashes with: CUDA error (.../hopper/flash_api.cpp:697): no kernel image is available for execution on the device Direct repro: a varlen FA call with scheduler_metadata on A100 reproduces the crash; without it, the same call succeeds. Now gate the metadata computation on get_flash_attn_version() == 3 in _single_switch_worker.py, _model_forward_tests.py, _position_zero_nan_tests.py. This was the actual root cause of issue #36 — not a kernel-image build mismatch. The fail-fast probe stays in place as a regression net. --- tests/vllm/_model_forward_tests.py | 37 +++++++++++++++----------- tests/vllm/_position_zero_nan_tests.py | 37 +++++++++++++++----------- tests/vllm/_single_switch_worker.py | 37 +++++++++++++++----------- 3 files changed, 63 insertions(+), 48 deletions(-) diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index 285a6c6..faf4cf4 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -265,27 +265,32 @@ def _build_metadata(self, seq_len): if backend_name == "FLASH_ATTN": from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata - # FA3 requires scheduler_metadata; compute it when available. + # scheduler_metadata is FA3-only — passing it on FA2 (Ampere/A100) + # forces FA3 kernel dispatch and crashes with "no kernel image + # is available". Only compute it when get_flash_attn_version() == 3 + # (Hopper SM90+). scheduler_metadata = None try: from vllm.v1.attention.backends.fa_utils import ( + get_flash_attn_version, get_scheduler_metadata, ) - first_attn = list(self._attention_map.values())[0] - scheduler_metadata = get_scheduler_metadata( - batch_size=1, - max_seqlen_q=seq_len, - max_seqlen_k=seq_len, - num_heads_q=first_attn.num_heads, - num_heads_kv=first_attn.num_kv_heads, - headdim=first_attn.head_size, - cache_seqlens=seq_lens, - qkv_dtype=torch.bfloat16, - cu_seqlens_q=query_start_loc, - page_size=BLOCK_SIZE, - causal=True, - num_splits=0, - ) + if get_flash_attn_version() == 3: + first_attn = list(self._attention_map.values())[0] + scheduler_metadata = get_scheduler_metadata( + batch_size=1, + max_seqlen_q=seq_len, + max_seqlen_k=seq_len, + num_heads_q=first_attn.num_heads, + num_heads_kv=first_attn.num_kv_heads, + headdim=first_attn.head_size, + cache_seqlens=seq_lens, + qkv_dtype=torch.bfloat16, + cu_seqlens_q=query_start_loc, + page_size=BLOCK_SIZE, + causal=True, + num_splits=0, + ) except ImportError: pass diff --git a/tests/vllm/_position_zero_nan_tests.py b/tests/vllm/_position_zero_nan_tests.py index 8b5fb90..69b3360 100644 --- a/tests/vllm/_position_zero_nan_tests.py +++ b/tests/vllm/_position_zero_nan_tests.py @@ -158,27 +158,32 @@ def _build_metadata(attention_map, seq_len, device): if backend_name == "FLASH_ATTN": from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata - # FA3 requires scheduler_metadata; compute it when available. + # scheduler_metadata is FA3-only — passing it on FA2 (Ampere/A100) + # forces FA3 kernel dispatch and crashes with "no kernel image + # is available". Only compute it when get_flash_attn_version() == 3 + # (Hopper SM90+). scheduler_metadata = None try: from vllm.v1.attention.backends.fa_utils import ( + get_flash_attn_version, get_scheduler_metadata, ) - first_attn = list(attention_map.values())[0] - scheduler_metadata = get_scheduler_metadata( - batch_size=1, - max_seqlen_q=seq_len, - max_seqlen_k=seq_len, - num_heads_q=first_attn.num_heads, - num_heads_kv=first_attn.num_kv_heads, - headdim=first_attn.head_size, - cache_seqlens=seq_lens, - qkv_dtype=torch.bfloat16, - cu_seqlens_q=query_start_loc, - page_size=BLOCK_SIZE, - causal=True, - num_splits=0, - ) + if get_flash_attn_version() == 3: + first_attn = list(attention_map.values())[0] + scheduler_metadata = get_scheduler_metadata( + batch_size=1, + max_seqlen_q=seq_len, + max_seqlen_k=seq_len, + num_heads_q=first_attn.num_heads, + num_heads_kv=first_attn.num_kv_heads, + headdim=first_attn.head_size, + cache_seqlens=seq_lens, + qkv_dtype=torch.bfloat16, + cu_seqlens_q=query_start_loc, + page_size=BLOCK_SIZE, + causal=True, + num_splits=0, + ) except ImportError: pass diff --git a/tests/vllm/_single_switch_worker.py b/tests/vllm/_single_switch_worker.py index a53a1fc..d876dc4 100644 --- a/tests/vllm/_single_switch_worker.py +++ b/tests/vllm/_single_switch_worker.py @@ -123,27 +123,32 @@ def _build_metadata(harness, seq_len): if backend_name == "FLASH_ATTN": from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata - # FA3 requires scheduler_metadata; compute it when available. + # scheduler_metadata is FA3-only — passing it on FA2 (Ampere/A100) + # forces FA3 kernel dispatch and crashes with "no kernel image + # is available". Only compute it when get_flash_attn_version() == 3 + # (Hopper SM90+). scheduler_metadata = None try: from vllm.v1.attention.backends.fa_utils import ( + get_flash_attn_version, get_scheduler_metadata, ) - switch = harness["switch"] - scheduler_metadata = get_scheduler_metadata( - batch_size=1, - max_seqlen_q=seq_len, - max_seqlen_k=seq_len, - num_heads_q=switch.num_heads, - num_heads_kv=switch.num_kv_heads, - headdim=switch.head_dim, - cache_seqlens=seq_lens, - qkv_dtype=torch.bfloat16, - cu_seqlens_q=query_start_loc, - page_size=block_size, - causal=True, - num_splits=0, - ) + if get_flash_attn_version() == 3: + switch = harness["switch"] + scheduler_metadata = get_scheduler_metadata( + batch_size=1, + max_seqlen_q=seq_len, + max_seqlen_k=seq_len, + num_heads_q=switch.num_heads, + num_heads_kv=switch.num_kv_heads, + headdim=switch.head_dim, + cache_seqlens=seq_lens, + qkv_dtype=torch.bfloat16, + cu_seqlens_q=query_start_loc, + page_size=block_size, + causal=True, + num_splits=0, + ) except ImportError: pass From 56855f6bc3313f33b28704030578120d5e9f1735 Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Tue, 19 May 2026 11:56:47 +0300 Subject: [PATCH 5/6] Pin torch to >=2.10.0,<2.11.0 in pyproject.toml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Matches the torch version vllm 0.19.1 hard-pins, so installs with the [vllm] extra are now consistent with what vLLM actually links against. Old floor (>=2.0.0) was loose enough that pip resolution could silently swap torch 2.10 for an ancient version when an unrelated wheel pinned something older — that's exactly what corrupted a working venv during debugging of issue #36. Trade-off: this conflicts with the [vllm20] extra if vLLM 0.20 moves to a different torch. If/when that happens, switch to per-extra torch pins (or drop the upper bound). Note: this does NOT fix issue #36 on its own. The fix for that is the scheduler_metadata gating already on this branch. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 907cb47..648c795 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10,<3.14" dependencies = [ - "torch>=2.0.0", + "torch>=2.10.0,<2.11.0", "transformers>=5.5.1", ] From 13e501fe7581e8701978d97452d4ffbc36cb1ac7 Mon Sep 17 00:00:00 2001 From: AlonMalach Date: Tue, 19 May 2026 12:23:13 +0300 Subject: [PATCH 6/6] rem torch version upperbound limit --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 648c795..fcf0f3f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10,<3.14" dependencies = [ - "torch>=2.10.0,<2.11.0", + "torch>=2.10.0", "transformers>=5.5.1", ]