From 1d20869e4ba65b9d59c3598c2c672c3a03fcd14e Mon Sep 17 00:00:00 2001 From: Peuqui Date: Sun, 13 Sep 2026 14:39:01 +0200 Subject: [PATCH 1/2] [Feature][SM75] Build and load a Turing FlashAttention-2 library next to the Volta one Turing GPUs got no FlashAttention backend: the CMake FA2 block fetches vllm-project/flash-attention for 8.0+ and the Volta fork for 7.0, so an sm75 or mixed 7.0;7.5 build shipped no sm75 kernels and the platform fell back to FlashInfer (paged prefill fails on SM75) or TRITON_ATTN. - cmake/external_projects/vllm_flash_attn_sm75.cmake: ExternalProject that builds Peuqui/flash-attention @ 43b9d29c (tag sm75-1cat-2026-09-13) with CUDA_ARCHS=7.5 and VLLM_FA2_OUTPUT_NAME=_vllm_fa2_C_sm75, installed next to the regular library; included only when 7.5 is in CUDA_ARCHS. - setup.py: the extension and the precompiled-wheel extraction entry. - flash_attn_interface: load the FA2 library on first use, chosen per the worker's own device (one library per process); ensure_fa2_library_loaded() for the SM70 backend's operator lookups, which relied on import-time loading (flash_attn_v100, sm70_e4m3_long, sm70_e4m3_scalar). - platforms/cuda.py: FLASH_ATTN, TRITON_ATTN, FLEX_ATTENTION on 7.5, capability read from the worker's own device. - flash_attn backend: capability floor 7.5, fp16 only below 8.0. - docs: FA2 row 7.5+. - tests: library per device, Turing gates, backend priority, sm75 forward against torch SDPA on a 7.5 device. Co-authored-by: Claude Fable 5.1 Signed-off-by: Peuqui --- CMakeLists.txt | 8 ++ .../vllm_flash_attn_sm75.cmake | 53 +++++++++++ docs/design/attention_backends.md | 2 +- setup.py | 5 + .../attention/test_fa2_sm75_forward.py | 74 +++++++++++++++ .../test_cuda_backend_priority_turing.py | 28 ++++++ .../attention/test_flash_attn_turing_gates.py | 43 +++++++++ .../test_fa2_library_per_device.py | 89 +++++++++++++++++ vllm/platforms/cuda.py | 17 +++- vllm/v1/attention/backends/flash_attn.py | 6 +- vllm/v1/attention/backends/flash_attn_v100.py | 10 +- vllm/v1/attention/ops/sm70_e4m3_long.py | 8 +- vllm/v1/attention/ops/sm70_e4m3_scalar.py | 8 +- vllm/vllm_flash_attn/flash_attn_interface.py | 95 +++++++++++++++++-- 14 files changed, 430 insertions(+), 16 deletions(-) create mode 100644 cmake/external_projects/vllm_flash_attn_sm75.cmake create mode 100644 tests/kernels/attention/test_fa2_sm75_forward.py create mode 100644 tests/v1/attention/test_cuda_backend_priority_turing.py create mode 100644 tests/v1/attention/test_flash_attn_turing_gates.py create mode 100644 tests/vllm_flash_attn/test_fa2_library_per_device.py diff --git a/CMakeLists.txt b/CMakeLists.txt index bfc58439a8..6c525c9da1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1572,6 +1572,14 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") "Skipping QuTLASS external project fetch: unsupported CUDA architecture ${CUDA_ARCHS}") endif() + # Turing gets its own FA2 build (see the file for why it is a separate + # project). Independent of which FA2 tree the block below fetches. + cuda_archs_loose_intersection(VLLM_FLASH_ATTN_SM75_ARCHS + "7.5" "${CUDA_ARCHS}") + if(VLLM_FLASH_ATTN_SM75_ARCHS) + include(cmake/external_projects/vllm_flash_attn_sm75.cmake) + endif() + cuda_archs_loose_intersection(VLLM_FLASH_ATTN_ARCHS "8.0+PTX" "${CUDA_ARCHS}") if(VLLM_FLASH_ATTN_ARCHS) diff --git a/cmake/external_projects/vllm_flash_attn_sm75.cmake b/cmake/external_projects/vllm_flash_attn_sm75.cmake new file mode 100644 index 0000000000..ca17411ddb --- /dev/null +++ b/cmake/external_projects/vllm_flash_attn_sm75.cmake @@ -0,0 +1,53 @@ +# FlashAttention-2 for Turing (sm75). +# +# The FA2 sources fetched by vllm_flash_attn.cmake build no sm75 kernels: the +# vllm-project tree targets 8.0+, the Volta fork targets 7.0 only. Turing +# therefore fell through to FlashInfer, whose paged prefill fails on sm75, or +# to TRITON_ATTN. This builds a second FA2 library from a pinned fork that +# enables the sm75 forward path (fp16-only, forward-only) and installs it as +# _vllm_fa2_C_sm75 next to the regular _vllm_fa2_C; flash_attn_interface.py +# loads the one that matches the worker's device. +# +# Why ExternalProject and not a second FetchContent: both FA2 trees define +# the target _vllm_fa2_C and override global CMake functions (see the note in +# vllm_flash_attn.cmake), so they cannot share one configure. The fork's +# CMake names the interpreter Python_EXECUTABLE, hence the translation below. + +include(ExternalProject) + +set(VLLM_FLASH_ATTN_SM75_COMMIT 43b9d29c9aa8e18d9351e7c643dc78ef7e7979fc) # tag sm75-1cat-2026-09-13 +set(VLLM_FLASH_ATTN_SM75_LIB _vllm_fa2_C_sm75.abi3.so) + +if(DEFINED ENV{MAX_JOBS}) + set(VLLM_FLASH_ATTN_SM75_JOBS -j $ENV{MAX_JOBS}) +else() + set(VLLM_FLASH_ATTN_SM75_JOBS "") +endif() + +ExternalProject_Add(vllm-flash-attn-sm75 + GIT_REPOSITORY https://github.com/Peuqui/flash-attention.git + GIT_TAG ${VLLM_FLASH_ATTN_SM75_COMMIT} + GIT_PROGRESS TRUE + GIT_SUBMODULES csrc/cutlass + GIT_SUBMODULES_RECURSE TRUE + CMAKE_ARGS + -DPython_EXECUTABLE=${VLLM_PYTHON_EXECUTABLE} + -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_CUDA_COMPILER=${CMAKE_CUDA_COMPILER} + -DCUDA_ARCHS=7.5 + -DFA2_ENABLED=ON + -DFA3_ENABLED=OFF + -DVLLM_FA2_OUTPUT_NAME=_vllm_fa2_C_sm75 + BUILD_COMMAND ${CMAKE_COMMAND} --build --target _vllm_fa2_C ${VLLM_FLASH_ATTN_SM75_JOBS} + BUILD_BYPRODUCTS /${VLLM_FLASH_ATTN_SM75_LIB} + INSTALL_COMMAND "" +) +ExternalProject_Get_Property(vllm-flash-attn-sm75 BINARY_DIR) + +# setup.py builds and installs extensions by name: target and install +# component are both called _vllm_fa2_C_sm75. +add_custom_target(_vllm_fa2_C_sm75 ALL DEPENDS vllm-flash-attn-sm75) +install(FILES ${BINARY_DIR}/${VLLM_FLASH_ATTN_SM75_LIB} + DESTINATION vllm/vllm_flash_attn + COMPONENT _vllm_fa2_C_sm75) +message(STATUS "vllm-flash-attn-sm75 will be built from Peuqui/flash-attention@${VLLM_FLASH_ATTN_SM75_COMMIT}") diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d5d18f4ef1..505f1e45da 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -163,7 +163,7 @@ Priority is **1 = highest** (tried first). | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x | | `FLASHINFER_SM70` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | +| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥7.5 | | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | diff --git a/setup.py b/setup.py index 3dedb83e66..2dabbc3fb3 100644 --- a/setup.py +++ b/setup.py @@ -859,6 +859,7 @@ def extract_precompiled_and_patch_package( "vllm/_flashmla_extension_C.abi3.so", "vllm/_sparse_flashmla_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so", + "vllm/vllm_flash_attn/_vllm_fa2_C_sm75.abi3.so", "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", "vllm/cumem_allocator.abi3.so", "vllm/spinloop.abi3.so", @@ -1255,6 +1256,10 @@ def _read_requirements(filename: str) -> list[str]: build_sm70_fa2 = _cuda_arch_contains(7, 0) and not _cuda_arch_at_least(8, 0) if _cuda_arch_at_least(8, 0) or build_sm70_fa2: ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) + if _cuda_arch_contains(7, 5): + # Turing FA2 from a pinned fork, next to the regular library + # (cmake/external_projects/vllm_flash_attn_sm75.cmake). + ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C_sm75")) if _cuda_arch_at_least(8, 0): if _cuda_arch_at_least(9, 0) and ( USE_PRECOMPILED_EXTENSIONS diff --git a/tests/kernels/attention/test_fa2_sm75_forward.py b/tests/kernels/attention/test_fa2_sm75_forward.py new file mode 100644 index 0000000000..eca4418f23 --- /dev/null +++ b/tests/kernels/attention/test_fa2_sm75_forward.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The Turing FA2 build (_vllm_fa2_C_sm75) computes the same attention as torch. + +Runs only on a compute capability 7.5 device with the sm75 library installed; +it exercises the loader (library chosen for the device) and the fp16 forward +path for the head sizes the Qwen3.8 and Llama families use. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + + +def _turing_device_index() -> int | None: + if not current_platform.is_cuda(): + return None + for i in range(torch.accelerator.device_count()): + if torch.cuda.get_device_capability(i) == (7, 5): + return i + return None + + +@pytest.mark.parametrize("head_size", [64, 128, 256]) +@pytest.mark.parametrize("seqlen", [1, 8, 333]) +def test_sm75_varlen_forward_matches_torch(head_size: int, seqlen: int): + index = _turing_device_index() + if index is None: + pytest.skip("needs a compute capability 7.5 device") + device = torch.device(f"cuda:{index}") + from vllm.vllm_flash_attn import flash_attn_interface as fai + + if fai._fa2_library_path((7, 5)) is None: + pytest.skip("the sm75 FA2 library is not installed") + torch.accelerator.set_device_index(index) + torch.manual_seed(0) + batch, heads, kv_heads = 2, 8, 2 + q = torch.randn( + batch * seqlen, heads, head_size, device=device, dtype=torch.float16 + ) + k = torch.randn( + batch * seqlen, kv_heads, head_size, device=device, dtype=torch.float16 + ) + v = torch.randn_like(k) + cu = torch.arange(0, (batch + 1) * seqlen, seqlen, device=device, dtype=torch.int32) + + out = fai.flash_attn_varlen_func( + q, + k, + v, + max_seqlen_q=seqlen, + cu_seqlens_q=cu, + max_seqlen_k=seqlen, + cu_seqlens_k=cu, + causal=True, + fa_version=2, + ) + assert fai._fa2_loaded_capability == (7, 5) + + # torch reference: grouped heads expanded, causal SDPA per sequence. + ref = torch.empty_like(q) + rep = heads // kv_heads + for b in range(batch): + s = slice(b * seqlen, (b + 1) * seqlen) + qb = q[s].transpose(0, 1).float() + kb = k[s].repeat_interleave(rep, dim=1).transpose(0, 1).float() + vb = v[s].repeat_interleave(rep, dim=1).transpose(0, 1).float() + ref[s] = ( + torch.nn.functional.scaled_dot_product_attention(qb, kb, vb, is_causal=True) + .transpose(0, 1) + .to(torch.float16) + ) + torch.testing.assert_close(out, ref, atol=2e-3, rtol=2e-3) diff --git a/tests/v1/attention/test_cuda_backend_priority_turing.py b/tests/v1/attention/test_cuda_backend_priority_turing.py new file mode 100644 index 0000000000..66faf8e01c --- /dev/null +++ b/tests/v1/attention/test_cuda_backend_priority_turing.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""On Turing the FA2 build is tried first; FlashInfer is not offered there. + +FlashInfer's paged prefill fails with "invalid argument" on sm75, so the +priority list for (7, 5) is FLASH_ATTN, TRITON_ATTN, FLEX_ATTENTION. +""" + +from vllm.platforms.cuda import _get_backend_priorities +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +def test_turing_tries_the_fa2_build_first(): + priorities = _get_backend_priorities( + use_mla=False, device_capability=DeviceCapability(7, 5) + ) + assert priorities[0] is AttentionBackendEnum.FLASH_ATTN + assert AttentionBackendEnum.FLASHINFER not in priorities + assert AttentionBackendEnum.TRITON_ATTN in priorities + + +def test_ampere_priorities_are_untouched(): + priorities = _get_backend_priorities( + use_mla=False, device_capability=DeviceCapability(8, 0) + ) + assert AttentionBackendEnum.FLASH_ATTN in priorities + assert AttentionBackendEnum.FLASHINFER in priorities diff --git a/tests/v1/attention/test_flash_attn_turing_gates.py b/tests/v1/attention/test_flash_attn_turing_gates.py new file mode 100644 index 0000000000..ea82eb4f3b --- /dev/null +++ b/tests/v1/attention/test_flash_attn_turing_gates.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Turing runs the fp16-only FA2 build: capability floor 7.5, dtype gate.""" + +import pytest +import torch + +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + + +@pytest.mark.parametrize( + ("capability", "expected"), + [((7, 0), False), ((7, 5), True), ((8, 0), True), ((9, 0), True)], +) +def test_capability_floor_is_turing(capability, expected): + assert ( + FlashAttentionBackend.supports_compute_capability(DeviceCapability(*capability)) + is expected + ) + + +def _combination(dtype: torch.dtype, capability: tuple[int, int]) -> str | None: + return FlashAttentionBackend.supports_combination( + head_size=128, + dtype=dtype, + kv_cache_dtype="auto", + block_size=16, + use_mla=False, + has_sink=False, + use_sparse=False, + device_capability=DeviceCapability(*capability), + ) + + +def test_turing_takes_fp16_only(): + assert _combination(torch.float16, (7, 5)) is None + reason = _combination(torch.bfloat16, (7, 5)) + assert reason is not None and "fp16" in reason + + +def test_ampere_keeps_bf16(): + assert _combination(torch.bfloat16, (8, 0)) is None diff --git a/tests/vllm_flash_attn/test_fa2_library_per_device.py b/tests/vllm_flash_attn/test_fa2_library_per_device.py new file mode 100644 index 0000000000..21dcf1dd92 --- /dev/null +++ b/tests/vllm_flash_attn/test_fa2_library_per_device.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FA2 ships one library per architecture; the loader follows the device. + +The regular _vllm_fa2_C build covers the CUDA arch list, a Turing rig gets a +second file _vllm_fa2_C_sm75.abi3.so next to it. Both export the module +_vllm_fa2_C, so a process loads exactly one, on first use, for the +capability of the device the ops run on. +""" + +import sys +from pathlib import Path + +import pytest +import torch + +from vllm.vllm_flash_attn import flash_attn_interface as fai + + +@pytest.fixture(autouse=True) +def _fresh_loader_state(monkeypatch): + monkeypatch.setattr(fai, "_fa2_loaded_capability", None) + monkeypatch.delitem(sys.modules, fai._FA2_MODULE, raising=False) + + +def test_turing_path_only_when_the_file_exists(monkeypatch, tmp_path: Path): + turing = tmp_path / "_vllm_fa2_C_sm75.abi3.so" + monkeypatch.setattr(fai, "_FA2_TURING_PATH", str(turing)) + assert fai._fa2_library_path((7, 5)) is None + turing.write_bytes(b"") + assert fai._fa2_library_path((7, 5)) == str(turing) + + +def test_other_capabilities_take_the_regular_build(monkeypatch, tmp_path: Path): + regular = tmp_path / "_vllm_fa2_C.abi3.so" + monkeypatch.setattr( + fai, "_FA2_DEFAULT_SPEC", type("Spec", (), {"origin": str(regular)})() + ) + assert fai._fa2_library_path((7, 0)) == str(regular) + assert fai._fa2_library_path((8, 0)) == str(regular) + monkeypatch.setattr(fai, "_FA2_DEFAULT_SPEC", None) + assert fai._fa2_library_path((8, 0)) is None + + +def test_loader_picks_the_devices_library_once(monkeypatch, tmp_path: Path): + # A Python file stands in for the extension: the loader only needs a + # module spec it can execute. + stub = tmp_path / "_vllm_fa2_C_sm75.py" + stub.write_text("LOADED_FOR = 'sm75'\n") + asked: list[torch.device] = [] + + def capability(device: torch.device) -> tuple[int, int]: + asked.append(device) + return (7, 5) + + monkeypatch.setattr(fai.torch.cuda, "get_device_capability", capability) + monkeypatch.setattr( + fai, "_fa2_library_path", lambda cap: str(stub) if cap == (7, 5) else None + ) + + fai.load_fa2_library(torch.device("cuda:1")) + assert sys.modules[fai._FA2_MODULE].LOADED_FOR == "sm75" + assert fai._fa2_loaded_capability == (7, 5) + # The second call does not ask the device again: one library per process. + fai.load_fa2_library(torch.device("cuda:0")) + assert asked == [torch.device("cuda:1")] + + +def test_ensure_loads_the_current_devices_library_once(monkeypatch): + loaded: list[torch.device] = [] + monkeypatch.setattr(fai, "_fa2_loaded_capability", None) + monkeypatch.setattr(fai.torch.accelerator, "current_device_index", lambda: 2) + + def load(device: torch.device) -> None: + loaded.append(device) + fai._fa2_loaded_capability = (7, 0) + + monkeypatch.setattr(fai, "load_fa2_library", load) + fai.ensure_fa2_library_loaded() + fai.ensure_fa2_library_loaded() + assert loaded == [torch.device("cuda", 2)] + + +def test_loader_refuses_a_capability_without_library(monkeypatch): + monkeypatch.setattr(fai.torch.cuda, "get_device_capability", lambda device: (7, 5)) + monkeypatch.setattr(fai, "_fa2_library_path", lambda cap: None) + with pytest.raises(ImportError, match="7.5"): + fai.load_fa2_library(torch.device("cuda:0")) + assert fai._fa2_loaded_capability is None diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 88b94bcc59..27187a2dd3 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -157,6 +157,16 @@ def _get_backend_priorities( AttentionBackendEnum.FLEX_ATTENTION, AttentionBackendEnum.TURBOQUANT, ] + if device_capability.major == 7 and device_capability.minor == 5: + # Turing: FLASH_ATTN runs via the sm75 FA2 build (fp16-only, + # cmake/external_projects/vllm_flash_attn_sm75.cmake). + # FlashInfer's paged prefill fails with "invalid argument" on + # SM75; TRITON_ATTN stays as the fallback for non-fp16 models. + return [ + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TRITON_ATTN, + AttentionBackendEnum.FLEX_ATTENTION, + ] return [ AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.FLASHINFER, @@ -327,7 +337,12 @@ def get_attn_backend_cls( attn_selector_config: AttentionSelectorConfig, num_heads: int | None = None, ) -> str: - device_capability = cls.get_device_capability() + # Heterogeneous PP: the backend must match THIS worker's GPU, not + # device 0 of the visibility list (an RTX stage and a V100 stage + # need different backends). Workers have set their device before + # any attention layer is built. + device_id = torch.cuda.current_device() if torch.cuda.is_initialized() else 0 + device_capability = cls.get_device_capability(device_id) assert device_capability is not None # First try checking just the selected backend, if there is one. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index c56c4ee6e1..f76ddedfdc 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -200,7 +200,9 @@ def supports_sink(cls) -> bool: @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - return capability >= DeviceCapability(8, 0) + # SM75 enablement: Turing runs the fp16-only FA2 build + # (see supports_combination for the dtype gate). + return capability >= DeviceCapability(7, 5) @classmethod def supports_combination( @@ -216,6 +218,8 @@ def supports_combination( ) -> str | None: if has_sink and device_capability < DeviceCapability(9, 0): return "sink not supported on compute capability < 9.0" + if device_capability < DeviceCapability(8, 0) and dtype != torch.float16: + return "the sm75 FA2 build is fp16-only" return None diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index a9b98d7c9a..72df95468b 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -1319,8 +1319,14 @@ def _get_sm70_splitd_d256_ops(): "sm70_d256_splitd_n32_paged_fwd", ) with suppress(ImportError): - # Importing the interface loads the bundled FA2 torch library. - from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + # The FA2 library loads on first use, one per process and chosen + # for the worker's device; make sure it is there before the + # operators are resolved. + from vllm.vllm_flash_attn.flash_attn_interface import ( + ensure_fa2_library_loaded, + ) + + ensure_fa2_library_loaded() namespace = getattr(torch.ops, "_vllm_fa2_C", None) if namespace is None or not all( diff --git a/vllm/v1/attention/ops/sm70_e4m3_long.py b/vllm/v1/attention/ops/sm70_e4m3_long.py index 5bc6876c12..a09948f39b 100644 --- a/vllm/v1/attention/ops/sm70_e4m3_long.py +++ b/vllm/v1/attention/ops/sm70_e4m3_long.py @@ -53,8 +53,14 @@ @lru_cache(maxsize=1) def builtin_long_attention(): try: + # The FA2 library loads on first use; resolve it before the lookup. + from vllm.vllm_flash_attn.flash_attn_interface import ( + ensure_fa2_library_loaded, + ) + + ensure_fa2_library_loaded() return getattr(torch.ops._vllm_fa2_C, BUILTIN_OP) - except AttributeError: + except (AttributeError, ImportError): return None diff --git a/vllm/v1/attention/ops/sm70_e4m3_scalar.py b/vllm/v1/attention/ops/sm70_e4m3_scalar.py index 28d0f2fa75..0a0d8041bc 100644 --- a/vllm/v1/attention/ops/sm70_e4m3_scalar.py +++ b/vllm/v1/attention/ops/sm70_e4m3_scalar.py @@ -40,8 +40,14 @@ @lru_cache(maxsize=1) def builtin_scalar_tail_attention(): try: + # The FA2 library loads on first use; resolve it before the lookup. + from vllm.vllm_flash_attn.flash_attn_interface import ( + ensure_fa2_library_loaded, + ) + + ensure_fa2_library_loaded() return getattr(torch.ops._vllm_fa2_C, BUILTIN_SCALAR_OP) - except AttributeError: + except (AttributeError, ImportError): return None diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 33955bb239..0748b4d834 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -4,21 +4,84 @@ # ruff: noqa: E501 +import importlib.util +import os +import sys + import torch -# isort: off -# We need to import the CUDA kernels after importing torch -# Use relative import to support build-from-source installation in vLLM +from vllm.logger import init_logger -try: - from . import _vllm_fa2_C # type: ignore[attr-defined] # noqa: F401 +logger = init_logger(__name__) +# isort: off +# FA2 ships one library per architecture: the CMake build for the target list +# and, on mixed Volta/Turing rigs, a separate Turing build next to it. Both +# export the same module and op namespace, so a process can hold only one of +# them, and the choice has to follow the GPU the ops run on. Importing here +# would decide before a worker has selected its device, so the library loads +# on first use instead (load_fa2_library). +_FA2_MODULE = f"{__package__}._vllm_fa2_C" +_FA2_TURING_CAPABILITY = (7, 5) +_FA2_TURING_PATH = os.path.join(os.path.dirname(__file__), "_vllm_fa2_C_sm75.abi3.so") +_FA2_DEFAULT_SPEC = importlib.util.find_spec(_FA2_MODULE) +_fa2_loaded_capability: tuple[int, int] | None = None + +if _FA2_DEFAULT_SPEC is not None or os.path.exists(_FA2_TURING_PATH): FA2_UNAVAILABLE_REASON = None FA2_AVAILABLE = True -except ImportError as e: - FA2_UNAVAILABLE_REASON = str(e) +else: + FA2_UNAVAILABLE_REASON = f"no {_FA2_MODULE} library is installed" FA2_AVAILABLE = False + +def _fa2_library_path(capability: tuple[int, int]) -> str | None: + """Return the FA2 library installed for ``capability``, or None.""" + if capability == _FA2_TURING_CAPABILITY: + return _FA2_TURING_PATH if os.path.exists(_FA2_TURING_PATH) else None + return _FA2_DEFAULT_SPEC.origin if _FA2_DEFAULT_SPEC is not None else None + + +def load_fa2_library(device: torch.device) -> None: + """Load the FA2 library built for ``device``'s architecture. + + A process holds one FA2 library; the first call decides which. + """ + global _fa2_loaded_capability + if _fa2_loaded_capability is not None: + return + capability = torch.cuda.get_device_capability(device) + path = _fa2_library_path(capability) + if path is None: + raise ImportError( + f"No {_FA2_MODULE} library is installed for compute capability " + f"{capability[0]}.{capability[1]}" + ) + spec = importlib.util.spec_from_file_location(_FA2_MODULE, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[_FA2_MODULE] = module + spec.loader.exec_module(module) + _fa2_loaded_capability = capability + logger.info( + "Loaded FA2 library %s for compute capability %d.%d.", + os.path.basename(path), + *capability, + ) + + +def ensure_fa2_library_loaded() -> None: + """Load the FA2 library for this process's current device, once. + + For code that resolves operators from ``torch.ops._vllm_fa2_C`` before the + first attention call goes through this module (the SM70 backend looks its + prefill and tail operators up at initialisation). Importing this module no + longer loads a library, so those lookups ask here first. + """ + if _fa2_loaded_capability is None: + load_fa2_library(torch.device("cuda", torch.accelerator.current_device_index())) + + try: from . import _vllm_fa3_C # type: ignore[attr-defined] # noqa: F401 @@ -54,8 +117,19 @@ def _is_fa2_supported() -> tuple[bool, str | None]: return False, f"FA2 is unavailable due to: {FA2_UNAVAILABLE_REASON}" from vllm.platforms import current_platform - if not current_platform.has_device_capability(80): - return False, "FA2 is only supported on devices with compute capability >= 8" + # SM75 enablement: Turing runs the fp16-only FA2 build; bf16 inputs are + # rejected by the C++ entry points. + # Mixed rigs: ask this worker's own GPU, not device 0 -- otherwise the + # weakest card in the grid decides for every stage. + device = torch.accelerator.current_device_index() + if not current_platform.has_device_capability(75, device): + return False, "FA2 is only supported on devices with compute capability >= 7.5" + capability = current_platform.get_device_capability(device) + if ( + capability is None + or _fa2_library_path((capability.major, capability.minor)) is None + ): + return False, "no FA2 library is installed for this GPU's architecture" return True, None @@ -297,6 +371,7 @@ def flash_attn_varlen_func( raise NotImplementedError("FA2 does not support s_aux") if num_splits > 1: raise NotImplementedError("FA2 does not support num_splits > 1") + load_fa2_library(q.device) out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd( q, k, @@ -449,6 +524,7 @@ def sparse_attn_func( softmax_scale = q.shape[-1] ** (-0.5) q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + load_fa2_library(q.device) out, softmax_lse = torch.ops._vllm_fa2_C.fwd_sparse( q, k, @@ -535,6 +611,7 @@ def sparse_attn_varlen_func( softmax_scale = q.shape[-1] ** (-0.5) q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + load_fa2_library(q.device) out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd_sparse( q, k, From d22daa743c02aa26330f0c2a9da22d1aad82e2f3 Mon Sep 17 00:00:00 2001 From: Peuqui Date: Sun, 13 Sep 2026 16:58:43 +0200 Subject: [PATCH 2/2] [Test][SM75] Follow the Turing backend priority in the Flash-V100 policy test test_flash_v100_priority_is_sm70_only asserted FLASHINFER second on 7.5; the sm75 list is FLASH_ATTN, TRITON_ATTN, FLEX_ATTENTION and FlashInfer is not offered there (its paged prefill fails with "invalid argument" on SM75). Co-authored-by: Claude Fable 5.1 Signed-off-by: Peuqui --- tests/v1/attention/test_sm70_flash_v100_policy.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 2525878271..364983ed5b 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -217,11 +217,14 @@ def test_flash_v100_priority_is_sm70_only(monkeypatch): device_capability=DeviceCapability(major=7, minor=5), ) + # Turing: FLASH_ATTN via the sm75 FA2 build, then TRITON_ATTN. FlashInfer is + # not offered on 7.5 (its paged prefill fails with "invalid argument"). assert backends[:3] == [ AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.TRITON_ATTN, + AttentionBackendEnum.FLEX_ATTENTION, ] + assert AttentionBackendEnum.FLASHINFER not in backends assert AttentionBackendEnum.FLASH_ATTN_V100 not in backends