Skip to content

Commit f6d1ecb

Browse files
0z5aptrendx
0z5a
and
ptrendx
authored
Isolate FlashAttention 2/3 imports so a broken install cannot break TE (#3553)
* Isolate FlashAttention 2/3 imports so a broken install cannot break TE The version gate only inspects distribution metadata, so a flash-attn whose extension fails to load -- missing CUDA dependency, ABI mismatch -- raised straight out of backends.py and made `import transformer_engine.pytorch` fail outright, even for callers that never ask for FlashAttention. FlashAttention 4 already guards its import this way, so give FA2 and FA3 the same treatment: on ImportError leave the module aliases at None, mark the backend unavailable, and warn with the version and the original error. A genuine TE failure still raises; only the optional dependency is contained. The regression test injects a metadata-only flash-attn distribution in a subprocess, because the fault has to be in place before transformer_engine is imported. It fails for both FA2 and FA3 without this change and passes with it. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de> * Fix optional FlashAttention import regression coverage Signed-off-by: Przemek Tredak <ptredak@nvidia.com> --------- Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
1 parent b57646c commit f6d1ecb

3 files changed

Lines changed: 159 additions & 34 deletions

File tree

‎qa/L0_pytorch_unittest/test.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ export NVTE_FLASH_ATTN_V4=0
2929

3030
pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
3131

32+
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_optional_flash_attn_import.xml $TE_PATH/tests/pytorch/test_optional_flash_attn_import.py || test_fail "test_optional_flash_attn_import.py"
3233
NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py"
3334
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py"
3435
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py"

‎tests/pytorch/test_optional_flash_attn_import.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# See LICENSE for license information.
4+
"""A broken FlashAttention install must not take down ``transformer_engine.pytorch``.
5+
6+
The distribution metadata can resolve to a supported version while the extension
7+
itself fails to load (missing CUDA dependency, ABI mismatch). That used to
8+
propagate out of ``backends.py`` and make the whole PyTorch interface
9+
unimportable, even for code that never asks for FlashAttention.
10+
11+
Fault injection has to happen before ``transformer_engine`` is imported, so each
12+
case runs in a fresh interpreter with a throwaway directory first on
13+
``PYTHONPATH``.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
import os
19+
import subprocess
20+
import sys
21+
import tempfile
22+
import textwrap
23+
from pathlib import Path
24+
25+
import pytest
26+
27+
PROBE = textwrap.dedent(
28+
"""
29+
import sys
30+
31+
import transformer_engine.pytorch # noqa: F401 (the import under test)
32+
from transformer_engine.pytorch.attention.dot_product_attention import backends
33+
from transformer_engine.pytorch.attention.dot_product_attention.utils import (
34+
FlashAttentionUtils,
35+
)
36+
37+
if sys.argv[1] == "flash-attn":
38+
assert FlashAttentionUtils.is_installed is False
39+
assert backends.flash_attn_func is None
40+
elif sys.argv[1] == "flash-attn-3":
41+
assert FlashAttentionUtils.v3_is_installed is False
42+
assert backends.flash_attn_func_v3 is None
43+
else:
44+
raise ValueError(f"Unexpected FlashAttention distribution: {sys.argv[1]}")
45+
print("PROBE_OK")
46+
"""
47+
)
48+
49+
50+
def _fake_broken_dist(root: Path, name: str, version: str, module: str, message: str) -> None:
51+
"""A distribution whose metadata resolves but whose module raises ImportError."""
52+
dist = root / f"{name.replace('-', '_')}-{version}.dist-info"
53+
dist.mkdir(parents=True, exist_ok=True)
54+
(dist / "METADATA").write_text(
55+
f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n", encoding="utf-8"
56+
)
57+
(dist / "RECORD").write_text("", encoding="utf-8")
58+
path = root / (module.replace(".", "/") + ".py")
59+
path.parent.mkdir(parents=True, exist_ok=True)
60+
path.write_text(f"raise ImportError({message!r})\n", encoding="utf-8")
61+
62+
63+
@pytest.mark.parametrize(
64+
"name,version,module",
65+
[
66+
("flash-attn", "2.7.4.post1", "flash_attn_2_cuda"),
67+
("flash-attn-3", "3.0.0b1", "flash_attn_interface"),
68+
],
69+
ids=["fa2", "fa3"],
70+
)
71+
def test_broken_flash_attn_does_not_break_te_import(name, version, module):
72+
with tempfile.TemporaryDirectory(prefix="te-broken-fa-") as tmp:
73+
_fake_broken_dist(
74+
Path(tmp), name, version, module, "libcudnn.so.9: cannot open shared object file"
75+
)
76+
env = dict(os.environ)
77+
inherited = env.get("PYTHONPATH", "")
78+
env["PYTHONPATH"] = tmp + (os.pathsep + inherited if inherited else "")
79+
proc = subprocess.run(
80+
[sys.executable, "-c", PROBE, name],
81+
capture_output=True,
82+
text=True,
83+
env=env,
84+
timeout=900,
85+
)
86+
assert proc.returncode == 0, proc.stderr[-3000:]
87+
assert "PROBE_OK" in proc.stdout

‎transformer_engine/pytorch/attention/dot_product_attention/backends.py‎

Lines changed: 71 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -115,19 +115,41 @@
115115
fa_utils.is_installed = True
116116

117117
if fa_utils.is_installed:
118-
from flash_attn_2_cuda import varlen_bwd as flash_attn_cuda_bwd
119-
from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_func
120-
from flash_attn.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd
121-
from flash_attn.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd
122-
from flash_attn.flash_attn_interface import (
123-
_flash_attn_varlen_forward as _flash_attn_varlen_fwd,
124-
)
125-
from flash_attn.flash_attn_interface import (
126-
_flash_attn_varlen_backward as _flash_attn_varlen_bwd,
127-
)
128-
129-
# Setup Flash attention utils
130-
fa_utils.set_flash_attention_version()
118+
# The distribution metadata can resolve to a supported version while the
119+
# extension itself still fails to load (missing CUDA dependency, ABI
120+
# mismatch, ...). Until every required symbol is bound, this backend must
121+
# not be advertised as available: flip the flag back and leave the module
122+
# aliases at their ``None`` defaults so an unusable FlashAttention 2
123+
# cannot break the rest of the PyTorch interface.
124+
try:
125+
from flash_attn_2_cuda import varlen_bwd as flash_attn_cuda_bwd
126+
from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_func
127+
from flash_attn.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd
128+
from flash_attn.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd
129+
from flash_attn.flash_attn_interface import (
130+
_flash_attn_varlen_forward as _flash_attn_varlen_fwd,
131+
)
132+
from flash_attn.flash_attn_interface import (
133+
_flash_attn_varlen_backward as _flash_attn_varlen_bwd,
134+
)
135+
except ImportError as exc:
136+
flash_attn_cuda_bwd = None
137+
flash_attn_func = None
138+
flash_attn_varlen_func = None
139+
_flash_attn_fwd = None
140+
_flash_attn_bwd = None
141+
_flash_attn_varlen_fwd = None
142+
_flash_attn_varlen_bwd = None
143+
fa_utils.is_installed = False
144+
warnings.warn(
145+
f"FlashAttention 2 is installed (version {fa_utils.version}) but cannot be "
146+
f"loaded: {exc}",
147+
RuntimeWarning,
148+
stacklevel=2,
149+
)
150+
else:
151+
# Setup Flash attention utils
152+
fa_utils.set_flash_attention_version()
131153
elif (
132154
torch.cuda.is_available()
133155
and get_device_compute_capability() >= (8, 0)
@@ -157,29 +179,44 @@
157179
_flash_attn_bwd_v3 = None
158180
# pass # only print warning if use_flash_attention_3 = True in get_attention_backend
159181
else:
160-
from flash_attn_interface import flash_attn_func as flash_attn_func_v3
161-
from flash_attn_interface import (
162-
flash_attn_varlen_func as flash_attn_varlen_func_v3,
163-
)
164-
from flash_attn_interface import (
165-
flash_attn_with_kvcache as flash_attn_with_kvcache_v3,
166-
)
167-
from flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3
168-
from flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3
169-
170-
fa_utils.set_flash_attention_3_params()
171-
172-
# Older FA3 releases expose no `softcap` kwarg, so probe the API rather than the version.
173-
# This cannot see a FLASHATTENTION_DISABLE_SOFTCAP build: that still exposes the kwarg and
174-
# rejects a nonzero cap at dispatch.
182+
# Same isolation as FlashAttention 2 above: an importable-but-unusable FA3
183+
# must degrade to "unavailable", not take down the module import.
175184
try:
176-
fa_utils.fa3_supports_softcap = (
177-
"softcap" in inspect.signature(flash_attn_func_v3).parameters
178-
and "softcap" in inspect.signature(flash_attn_varlen_func_v3).parameters
179-
and "softcap" in inspect.signature(flash_attn_with_kvcache_v3).parameters
185+
from flash_attn_interface import flash_attn_func as flash_attn_func_v3
186+
from flash_attn_interface import (
187+
flash_attn_varlen_func as flash_attn_varlen_func_v3,
188+
)
189+
from flash_attn_interface import (
190+
flash_attn_with_kvcache as flash_attn_with_kvcache_v3,
180191
)
181-
except (ValueError, TypeError):
182-
fa_utils.fa3_supports_softcap = False
192+
from flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3
193+
from flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3
194+
except ImportError as exc:
195+
flash_attn_func_v3 = None
196+
flash_attn_varlen_func_v3 = None
197+
flash_attn_with_kvcache_v3 = None
198+
_flash_attn_fwd_v3 = None
199+
_flash_attn_bwd_v3 = None
200+
warnings.warn(
201+
f"FlashAttention 3 is installed (version {fa_utils.fa3_version}) but cannot be "
202+
f"loaded: {exc}",
203+
RuntimeWarning,
204+
stacklevel=2,
205+
)
206+
else:
207+
fa_utils.set_flash_attention_3_params()
208+
209+
# Older FA3 releases expose no `softcap` kwarg, so probe the API rather than the version.
210+
# This cannot see a FLASHATTENTION_DISABLE_SOFTCAP build: that still exposes the kwarg and
211+
# rejects a nonzero cap at dispatch.
212+
try:
213+
fa_utils.fa3_supports_softcap = (
214+
"softcap" in inspect.signature(flash_attn_func_v3).parameters
215+
and "softcap" in inspect.signature(flash_attn_varlen_func_v3).parameters
216+
and "softcap" in inspect.signature(flash_attn_with_kvcache_v3).parameters
217+
)
218+
except (ValueError, TypeError):
219+
fa_utils.fa3_supports_softcap = False
183220

184221
# Try to import Flash Attention v4
185222
try:

0 commit comments

Comments
 (0)