Fix FA3 scheduler_metadata crash on non-Hopper GPUs (issue #36)

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",
     "transformers>=5.5.1",
 ]
 

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
 

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
 

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": "..."}
@@ -118,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
 
@@ -232,8 +242,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} 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
+
+
 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: <empty>".
+        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"]}

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