Install Flash Attention on Blackwell when a CUDA >= 12.8 wheel exists

The Studio installer skipped Flash Attention on any Blackwell GPU
(compute capability >= 10.0: sm_100 B200/B100, sm_120 RTX 50 series),
on the assumption that no prebuilt wheels existed for those archs. That
assumption is stale: Dao-AILab has shipped wheels with sm_100 / sm_120
kernels since flash-attn v2.7.3, built for CUDA >= 12.8.

Replace the blanket has_blackwell_gpu() skip in _ensure_flash_attn and
_ensure_flash_attn_for_long_context with a CUDA-version-aware gate. The
normal prebuilt-wheel path already builds the exact wheel URL, HEAD-checks
it, and falls back gracefully, so it now runs on Blackwell unless the
installed torch's CUDA build is older than 12.8 (in which case the
matching wheel would lack Blackwell kernels and crash at kernel launch).

- wheel_utils: probe now reports full cuda_version; add
  cuda_supports_blackwell_flash_attn and should_skip_flash_attn_for_blackwell.
- install_python_stack / training worker: use the new gate.
- Update the corresponding tests.

studio/backend/core/training/worker.py

@@ -48,9 +48,9 @@ from utils.training_runs import build_default_output_dir_name
 from utils.wheel_utils import (
     direct_wheel_url,
     flash_attn_wheel_url,
-    has_blackwell_gpu,
     install_wheel,
     probe_torch_wheel_env,
+    should_skip_flash_attn_for_blackwell,
     url_exists,
 )
 
@@ -1030,10 +1030,10 @@ def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool:
 def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None:
     if not _should_try_runtime_flash_attn_install(max_seq_length):
         return
-    if has_blackwell_gpu():
+    if should_skip_flash_attn_for_blackwell():
         _send_status(
             event_queue,
-            "Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel",
+            "Skipping flash-attn install: Blackwell GPU needs a CUDA >= 12.8 torch build for prebuilt kernels",
         )
         return
 

studio/backend/tests/test_mlx_training_worker_config.py

@@ -37,9 +37,9 @@ def _load_worker_module():
         for name in (
             "direct_wheel_url",
             "flash_attn_wheel_url",
-            "has_blackwell_gpu",
             "install_wheel",
             "probe_torch_wheel_env",
+            "should_skip_flash_attn_for_blackwell",
             "url_exists",
         ):
             setattr(wheel_utils, name, lambda *_args, **_kwargs: None)

studio/backend/tests/test_training_worker_flash_attn.py

@@ -10,6 +10,7 @@ from typing import Any
 from unittest import mock
 
 from core.training import worker
+from utils import wheel_utils
 
 
 def _missing_flash_attn_import():
@@ -59,7 +60,7 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch):
     statuses: list[str] = []
 
     monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
-    monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
+    monkeypatch.setattr(wheel_utils, "has_blackwell_gpu", lambda: False)
     monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
     monkeypatch.setattr(
         worker,
@@ -88,7 +89,7 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
     statuses: list[str] = []
 
     monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
-    monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
+    monkeypatch.setattr(wheel_utils, "has_blackwell_gpu", lambda: False)
     monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
     monkeypatch.setattr(
         worker,
@@ -141,13 +142,18 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch):
     worker._sp.run.assert_not_called()
 
 
-def test_runtime_flash_attn_skips_on_blackwell(monkeypatch):
+def test_runtime_flash_attn_skips_on_blackwell_with_old_cuda(monkeypatch):
     statuses: list[str] = []
     install_mock = mock.Mock()
 
     monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
     monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True)
-    monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True)
+    monkeypatch.setattr(wheel_utils, "has_blackwell_gpu", lambda: True)
+    monkeypatch.setattr(
+        wheel_utils,
+        "probe_torch_wheel_env",
+        lambda timeout = None: {"cuda_version": "12.6"},
+    )
     monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
     monkeypatch.setattr(
         worker,
@@ -162,6 +168,40 @@ def test_runtime_flash_attn_skips_on_blackwell(monkeypatch):
     assert "Blackwell" in statuses[0]
 
 
+def test_runtime_flash_attn_installs_on_blackwell_with_new_cuda(monkeypatch):
+    statuses: list[str] = []
+
+    monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
+    monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True)
+    monkeypatch.setattr(wheel_utils, "has_blackwell_gpu", lambda: True)
+    monkeypatch.setattr(
+        wheel_utils,
+        "probe_torch_wheel_env",
+        lambda timeout = None: {"cuda_version": "13.0"},
+    )
+    monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
+    monkeypatch.setattr(
+        worker,
+        "flash_attn_wheel_url",
+        lambda env: "https://example.com/fa.whl",
+    )
+    monkeypatch.setattr(worker, "url_exists", lambda url: True)
+    monkeypatch.setattr(
+        worker,
+        "_send_status",
+        lambda queue, message: statuses.append(message),
+    )
+    monkeypatch.setattr(
+        worker,
+        "install_wheel",
+        lambda *args, **kwargs: [("pip", subprocess.CompletedProcess(["pip"], 0, ""))],
+    )
+
+    worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536)
+
+    assert statuses == ["Installing flash-attn for faster training..."]
+
+
 def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch):
     install_mock = mock.Mock(return_value = True)
     monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)

studio/backend/utils/wheel_utils.py

@@ -26,9 +26,11 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele
 def has_blackwell_gpu() -> bool:
     """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell).
 
-    Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels
-    fail to load, so callers use this to skip the flash-attn install path. Cached
-    for the process lifetime; tests mocking nvidia-smi must call
+    flash-attention ships prebuilt wheels with sm_100 / sm_120 kernels since v2.7.3,
+    but only when the wheel was built with CUDA >= 12.8. Callers combine this with
+    ``cuda_supports_blackwell_flash_attn`` (see ``should_skip_flash_attn_for_blackwell``)
+    to decide whether a Blackwell box can use the prebuilt wheel. Cached for the
+    process lifetime; tests mocking nvidia-smi must call
     ``has_blackwell_gpu.cache_clear()`` first.
     """
     exe = shutil.which("nvidia-smi")
@@ -91,6 +93,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
                     "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
                     "'torch_mm': torch_mm, "
                     "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
+                    "'cuda_version': str(torch.version.cuda) if torch.version.cuda else '', "
                     "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', "
                     "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
                     "}))"
@@ -164,6 +167,44 @@ def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None:
     )
 
 
+# flash-attention's prebuilt wheels only carry Blackwell (sm_100 / sm_120) SASS when
+# built with CUDA >= 12.8 -- the compute_100 / compute_120 gencode gate added in
+# flash-attention v2.7.3. Wheel filenames encode only the CUDA major (cu12 / cu13),
+# so decide off the installed torch's full CUDA version instead.
+BLACKWELL_MIN_CUDA = (12, 8)
+
+
+def cuda_supports_blackwell_flash_attn(cuda_version: str | None) -> bool:
+    """True when ``cuda_version`` (e.g. '13.0', '12.8') is >= 12.8, the CUDA release
+    from which flash-attention wheels contain sm_100 / sm_120 kernels. False for a
+    missing or unparseable version.
+    """
+    if not cuda_version:
+        return False
+    parts = str(cuda_version).split(".")
+    try:
+        major = int(parts[0])
+        minor = int(parts[1]) if len(parts) > 1 else 0
+    except (IndexError, ValueError):
+        return False
+    return (major, minor) >= BLACKWELL_MIN_CUDA
+
+
+def should_skip_flash_attn_for_blackwell(env: dict[str, str] | None = None) -> bool:
+    """True when a Blackwell GPU is present but the installed torch's CUDA build is
+    older than 12.8, so the matching prebuilt flash-attn wheel would lack sm_100 /
+    sm_120 kernels (installs fine, then crashes at kernel launch). Returns False for
+    non-Blackwell GPUs and for Blackwell on CUDA >= 12.8, letting the normal
+    prebuilt-wheel install path run. Pass a pre-probed ``env`` to avoid re-probing.
+    """
+    if not has_blackwell_gpu():
+        return False
+    if env is None:
+        env = probe_torch_wheel_env()
+    cuda_version = env.get("cuda_version") if env else None
+    return not cuda_supports_blackwell_flash_attn(cuda_version)
+
+
 def install_wheel(
     wheel_url: str,
     *,

studio/install_python_stack.py

@@ -31,9 +31,9 @@ if str(_BACKEND_DIR) not in sys.path:
 from backend.utils.wheel_utils import (
     flash_attn_package_version,
     flash_attn_wheel_url,
-    has_blackwell_gpu,
     install_wheel,
     probe_torch_wheel_env,
+    should_skip_flash_attn_for_blackwell,
     url_exists,
 )
 from backend.utils.uv_path_safety import uv_safe_path as _uv_safe_path
@@ -1759,13 +1759,6 @@ def _ensure_flash_attn() -> None:
         return
     if NO_TORCH:
         return
-    if has_blackwell_gpu():
-        _step(
-            "warning",
-            "Skipping flash-attn: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel",
-            _cyan,
-        )
-        return
     if IS_WINDOWS or IS_MACOS:
         return
     if (
@@ -1779,6 +1772,13 @@ def _ensure_flash_attn() -> None:
         return
 
     env = probe_torch_wheel_env()
+    if should_skip_flash_attn_for_blackwell(env):
+        _step(
+            "warning",
+            "Skipping flash-attn: Blackwell GPU needs a CUDA >= 12.8 torch build for prebuilt kernels",
+            _cyan,
+        )
+        return
     wheel_url = _build_flash_attn_wheel_url(env) if env else None
     if wheel_url and url_exists(wheel_url):
         for installer, wheel_result in install_wheel(

tests/python/test_flash_attn_install_python_stack.py

@@ -111,6 +111,63 @@ class TestHasBlackwellGpu:
             assert wheel_utils.has_blackwell_gpu() is False
 
 
+class TestCudaSupportsBlackwellFlashAttn:
+    def test_cuda_13_supported(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("13.0") is True
+
+    def test_cuda_128_is_the_threshold(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("12.8") is True
+
+    def test_cuda_129_supported(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("12.9") is True
+
+    def test_major_only_supported(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("13") is True
+
+    def test_cuda_126_below_threshold(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("12.6") is False
+
+    def test_cuda_118_below_threshold(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("11.8") is False
+
+    def test_empty_or_none_is_false(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("") is False
+        assert wheel_utils.cuda_supports_blackwell_flash_attn(None) is False
+
+    def test_unparseable_is_false(self):
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("not-a-version") is False
+        assert wheel_utils.cuda_supports_blackwell_flash_attn("12.x") is False
+
+
+class TestShouldSkipFlashAttnForBlackwell:
+    def test_non_blackwell_never_skips(self):
+        with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False):
+            assert wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": ""}) is False
+
+    def test_blackwell_with_cuda_13_does_not_skip(self):
+        with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True):
+            assert (
+                wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": "13.0"}) is False
+            )
+
+    def test_blackwell_with_old_cuda_skips(self):
+        with mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True):
+            assert (
+                wheel_utils.should_skip_flash_attn_for_blackwell({"cuda_version": "12.6"}) is True
+            )
+
+    def test_blackwell_probes_when_env_missing(self):
+        with (
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
+            mock.patch.object(
+                wheel_utils,
+                "probe_torch_wheel_env",
+                return_value = {"cuda_version": "12.6"},
+            ),
+        ):
+            assert wheel_utils.should_skip_flash_attn_for_blackwell() is True
+
+
 class TestFlashAttnWheelSelection:
     def test_torch_210_maps_to_v281(self):
         assert ips._select_flash_attn_version("2.10") == "2.8.1"
@@ -167,6 +224,7 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "IS_MACOS", False),
             mock.patch.object(ips, "USE_UV", True),
             mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False),
             mock.patch.object(
                 ips,
                 "probe_torch_wheel_env",
@@ -206,6 +264,7 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "IS_MACOS", False),
             mock.patch.object(ips, "USE_UV", True),
             mock.patch.object(ips, "UV_NEEDS_SYSTEM", True),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False),
             mock.patch.object(
                 ips,
                 "probe_torch_wheel_env",
@@ -244,6 +303,7 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "IS_MACOS", False),
             mock.patch.object(ips, "USE_UV", True),
             mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False),
             mock.patch.object(
                 ips,
                 "probe_torch_wheel_env",
@@ -297,6 +357,7 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "NO_TORCH", False),
             mock.patch.object(ips, "IS_WINDOWS", False),
             mock.patch.object(ips, "IS_MACOS", False),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = False),
             mock.patch.object(
                 ips,
                 "probe_torch_wheel_env",
@@ -333,7 +394,46 @@ class TestEnsureFlashAttn:
         mock_probe.assert_not_called()
         mock_install_wheel.assert_not_called()
 
-    def test_blackwell_gpu_skips_install_with_warning(self):
+    def test_blackwell_with_supported_cuda_installs_wheel(self):
+        install_calls = []
+
+        def fake_install_wheel(*args, **kwargs):
+            install_calls.append((args, kwargs))
+            return [("uv", subprocess.CompletedProcess(["uv"], 0, ""))]
+
+        with (
+            mock.patch.object(ips, "NO_TORCH", False),
+            mock.patch.object(ips, "IS_WINDOWS", False),
+            mock.patch.object(ips, "IS_MACOS", False),
+            mock.patch.object(ips, "USE_UV", True),
+            mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
+            mock.patch.object(
+                ips,
+                "probe_torch_wheel_env",
+                return_value = {
+                    "python_tag": "cp313",
+                    "torch_mm": "2.10",
+                    "cuda_major": "13",
+                    "cuda_version": "13.0",
+                    "cxx11abi": "TRUE",
+                    "platform_tag": "linux_x86_64",
+                },
+            ),
+            mock.patch.object(ips, "url_exists", return_value = True),
+            mock.patch.object(ips, "install_wheel", side_effect = fake_install_wheel),
+            mock.patch("subprocess.run", return_value = self._import_check()),
+        ):
+            ips._ensure_flash_attn()
+
+        assert len(install_calls) == 1
+        (url,), _ = install_calls[0]
+        assert (
+            url
+            == "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
+        )
+
+    def test_blackwell_with_old_cuda_skips_with_warning(self):
         step_messages: list[tuple[str, str]] = []
 
         def fake_step(
@@ -347,19 +447,32 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "NO_TORCH", False),
             mock.patch.object(ips, "IS_WINDOWS", False),
             mock.patch.object(ips, "IS_MACOS", False),
-            mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
-            mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
+            mock.patch.object(
+                ips,
+                "probe_torch_wheel_env",
+                return_value = {
+                    "python_tag": "cp313",
+                    "torch_mm": "2.10",
+                    "cuda_major": "12",
+                    "cuda_version": "12.6",
+                    "cxx11abi": "TRUE",
+                    "platform_tag": "linux_x86_64",
+                },
+            ),
             mock.patch.object(ips, "install_wheel") as mock_install_wheel,
             mock.patch.object(ips, "_step", side_effect = fake_step),
             mock.patch("subprocess.run", return_value = self._import_check()),
         ):
             ips._ensure_flash_attn()
 
-        mock_probe.assert_not_called()
         mock_install_wheel.assert_not_called()
-        assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
+        assert any(
+            label == "warning" and "Blackwell" in msg and "CUDA" in msg
+            for label, msg in step_messages
+        )
 
-    def test_blackwell_gpu_on_windows_emits_blackwell_warning(self):
+    def test_windows_returns_before_blackwell_gate(self):
         step_messages: list[tuple[str, str]] = []
 
         def fake_step(
@@ -373,33 +486,7 @@ class TestEnsureFlashAttn:
             mock.patch.object(ips, "NO_TORCH", False),
             mock.patch.object(ips, "IS_WINDOWS", True),
             mock.patch.object(ips, "IS_MACOS", False),
-            mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
-            mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
-            mock.patch.object(ips, "install_wheel") as mock_install_wheel,
-            mock.patch.object(ips, "_step", side_effect = fake_step),
-            mock.patch("subprocess.run", return_value = self._import_check()),
-        ):
-            ips._ensure_flash_attn()
-
-        mock_probe.assert_not_called()
-        mock_install_wheel.assert_not_called()
-        assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
-
-    def test_non_blackwell_windows_does_not_emit_blackwell_warning(self):
-        step_messages: list[tuple[str, str]] = []
-
-        def fake_step(
-            label: str,
-            value: str,
-            color_fn = None,
-        ):
-            step_messages.append((label, value))
-
-        with (
-            mock.patch.object(ips, "NO_TORCH", False),
-            mock.patch.object(ips, "IS_WINDOWS", True),
-            mock.patch.object(ips, "IS_MACOS", False),
-            mock.patch.object(ips, "has_blackwell_gpu", return_value = False),
+            mock.patch.object(wheel_utils, "has_blackwell_gpu", return_value = True),
             mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
             mock.patch.object(ips, "install_wheel") as mock_install_wheel,
             mock.patch.object(ips, "_step", side_effect = fake_step),