Skip to content

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

Merged
antonpibm merged 6 commits into
mainfrom
bugfix/issue-36-fa3-kernel-mismatch
May 19, 2026
Merged

Fix FA3 scheduler_metadata crash on non-Hopper GPUs (issue #36)#39
antonpibm merged 6 commits into
mainfrom
bugfix/issue-36-fa3-kernel-mismatch

Conversation

@AlonMalach

@AlonMalach AlonMalach commented May 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Summary

Closes #36.

The vLLM test suite was failing 374+ tests on A100 with cascading BrokenPipeErrors rooted in:

CUDA error (.../hopper/flash_api.cpp:697):
  no kernel image is available for execution on the device

The diagnosis trail (logged in commits below) ran through a few wrong paths before landing on the actual root cause: the test workers were unconditionally passing scheduler_metadata (an FA3-only kwarg) to FlashAttentionMetadata. On A100, vLLM correctly auto-selects FA2 via get_flash_attn_version(), but supplying that kwarg forced dispatch into vLLM's bundled FA3 Hopper-only kernel (_vllm_fa3_C.abi3.so), which aborted the worker process at the C layer.

Direct repro confirmed it: a varlen FA call with scheduler_metadata crashes on A100; without it, the same call succeeds.

What changed

Real fix (commit 2ca8574):

  • tests/vllm/_single_switch_worker.py, tests/vllm/_model_forward_tests.py, tests/vllm/_position_zero_nan_tests.py — gate the get_scheduler_metadata(...) call on get_flash_attn_version() == 3. On Hopper FA3 still gets its scheduler metadata; on Ampere FA2 runs without it.

Diagnostic regression net (commits 0668459, partially preserved through 8c4e749):

  • tests/vllm/_single_switch_worker.py — _probe_attention() runs a 1-token forward at startup and emits a structured {"fatal", "hint", "backend_name"} JSON to stdout if anything crashes. main() wraps _setup() similarly.
  • tests/vllm/test_single_switch.py — _ensure_worker() recognizes the fatal JSON, drains stderr, stores it in a sticky _fatal_startup_error flag, and pytest.fails every subsequent test with the same message. No more 374-deep BrokenPipeError cascades when the worker can't start.

Environment hardening (56855f6, 13e501f):

  • pyproject.toml — bump torch floor from >=2.0.0 to >=2.10.0. Keeps installs from silently slipping to ancient torch when an unrelated wheel pins something older (which contaminated a working venv during this debug session).

Wrong turns preserved in history

Commits 0668459 and f8aa8c9 tried to monkey-patch vLLM's attention selector to force FLASHINFER on non-Hopper. That path was abandoned (8c4e749) once we confirmed vLLM 0.19's auto-selection of FA2 on A100 was correct — the only bug was the test code over-feeding the FA call. Squash-merge will collapse these.

Files changed

  • tests/vllm/_single_switch_worker.py — FA-version gate + probe + fatal protocol
  • tests/vllm/_model_forward_tests.py — FA-version gate
  • tests/vllm/_position_zero_nan_tests.py — FA-version gate
  • tests/vllm/test_single_switch.py — sticky-error parent
  • pyproject.toml — torch floor bump

AlonMalach added 6 commits May 18, 2026 19:14
@AlonMalach
Force compatible attention backend in vLLM tests on non-Hopper GPUs (…
0668459 
…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).
@AlonMalach
Patch get_attn_backend across all vLLM consumer modules, not just sel…
f8aa8c9 
…ector

(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.
@AlonMalach
Drop attention-backend override; keep only fail-fast diagnostics
8c4e749 
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.
@AlonMalach
Gate scheduler_metadata on FA version 3 (fixes FA3 crash on A100)
2ca8574 
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.
@AlonMalach
Pin torch to >=2.10.0,<2.11.0 in pyproject.toml
56855f6 
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.
@AlonMalach
rem torch version upperbound limit
13e501f
@antonpibm

antonpibm commented May 19, 2026

Copy link
Copy Markdown
Collaborator

Tests passed:
Compose
=========== 243 passed, 1 skipped, 17 warnings in 1329.59s (0:22:09) ===========
vLLM
================= 425 passed, 2 skipped in 2436.38s (0:40:36) ==================
Unit
============================= 121 passed in 10.13s =============================
HF
========== 681 passed, 16 skipped, 56 warnings in 2674.80s (0:44:34) ===========
Integration
============ 3 failed, 24 passed, 66 warnings in 838.18s (0:13:58) =============

Chronically failing integration tests will be covered in #42

@antonpibm
antonpibm merged commit f37cd54 into main May 19, 2026
@antonpibm
antonpibm deleted the bugfix/issue-36-fa3-kernel-mismatch branch May 19, 2026 11:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@antonpibm antonpibm antonpibm approved these changes
@freunda freunda Awaiting requested review from freunda freunda is a code owner
@yairallouche yairallouche Awaiting requested review from yairallouche yairallouche is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

make test-all: vLLM suites fail with FlashAttention CUDA kernel-image mismatch (main + feature branch)

2 participants

@AlonMalach @antonpibm