diff --git a/pyproject.toml b/pyproject.toml index 907cb47..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.0.0", + "torch>=2.10.0", "transformers>=5.5.1", ] 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 0d9d0f9..d876dc4 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": "..."} @@ -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: ". + 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