studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420)

* studio: skip flash-attn install on Blackwell GPUs (sm_100+)

Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120,
or sm_121, and the older-arch wheels fail to load on Blackwell. Add a
shared has_blackwell_gpu() helper and gate both the install-time
(install_python_stack._ensure_flash_attn) and runtime
(worker._ensure_flash_attn_for_long_context) paths on it. Detection uses
nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows.

* test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests

prefers_prebuilt_wheel and falls_back_to_pypi exercise the install
paths that the Blackwell guard now short-circuits. Make them explicit
about non-Blackwell so they pass on real Blackwell hosts.

* studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH

- Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a
  single process avoid redundant nvidia-smi spawns. Tests clear the
  cache via setup_method/teardown_method.
- In _ensure_flash_attn, run the NO_TORCH short-circuit before the
  Blackwell check so GGUF-only users (who never install torch anyway)
  do not see a Blackwell warning. Blackwell check still runs above the
  IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see
  the explicit reason rather than a silent OS skip.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: add has_blackwell_gpu to mlx worker test wheel_utils stub

test_mlx_training_worker_config loads worker.py against a hand-rolled
utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol
list so worker's import line resolves.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/training/worker.py

@@ -30,6 +30,7 @@ from utils.hardware import apply_gpu_ids
 from utils.wheel_utils import (
     direct_wheel_url,
     flash_attn_wheel_url,
+    has_blackwell_gpu,
     install_wheel,
     probe_torch_wheel_env,
     url_exists,
@@ -313,6 +314,12 @@ 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():
+        _send_status(
+            event_queue,
+            "Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel",
+        )
+        return
 
     installed = _install_package_wheel_first(
         event_queue = event_queue,

studio/backend/tests/test_mlx_training_worker_config.py

@@ -37,6 +37,7 @@ def _load_worker_module():
         for name in (
             "direct_wheel_url",
             "flash_attn_wheel_url",
+            "has_blackwell_gpu",
             "install_wheel",
             "probe_torch_wheel_env",
             "url_exists",

studio/backend/tests/test_training_worker_flash_attn.py

@@ -37,6 +37,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(builtins, "__import__", _missing_flash_attn_import())
     monkeypatch.setattr(
         worker,
@@ -65,6 +66,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(builtins, "__import__", _missing_flash_attn_import())
     monkeypatch.setattr(
         worker,
@@ -112,6 +114,29 @@ 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):
+    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(worker, "_install_package_wheel_first", install_mock)
+    monkeypatch.setattr(
+        worker,
+        "_send_status",
+        lambda queue, message: statuses.append(message),
+    )
+
+    worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536)
+
+    install_mock.assert_not_called()
+    assert len(statuses) == 1
+    assert "Blackwell" in statuses[0]
+
+
 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

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import functools
 import json
 import logging
 import platform
@@ -22,6 +23,49 @@ FLASH_ATTN_RELEASE_BASE_URL = (
 )
 
 
+@functools.lru_cache(maxsize = 1)
+def has_blackwell_gpu() -> bool:
+    """Return True if any visible NVIDIA GPU has compute capability >= 10.0
+    (Blackwell: sm_100, sm_120, sm_121, ...).
+
+    Dao-AILab does not publish prebuilt flash-attention wheels for these
+    architectures, and the older-arch wheels fail to load on Blackwell, so
+    callers use this gate to skip the flash-attn install/upgrade path.
+
+    Result is cached for the process lifetime since GPU hardware does not
+    change. Tests that mock subprocess/nvidia-smi must call
+    ``has_blackwell_gpu.cache_clear()`` before each invocation.
+    """
+    exe = shutil.which("nvidia-smi")
+    if not exe:
+        return False
+    try:
+        result = subprocess.run(
+            [exe, "--query-gpu=compute_cap", "--format=csv,noheader"],
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            text = True,
+            timeout = 10,
+            env = child_env_without_native_path_secret(),
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return False
+    if result.returncode != 0:
+        return False
+    for line in result.stdout.splitlines():
+        cap = line.strip()
+        if not cap:
+            continue
+        major_part = cap.split(".", 1)[0]
+        try:
+            major = int(major_part)
+        except ValueError:
+            continue
+        if major >= 10:
+            return True
+    return False
+
+
 def linux_wheel_platform_tag() -> str | None:
     machine = platform.machine().lower()
     if sys.platform.startswith("linux"):

studio/install_python_stack.py

@@ -28,6 +28,7 @@ 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,
     url_exists,
@@ -628,10 +629,19 @@ def _flash_attn_install_disabled() -> bool:
 
 
 def _ensure_flash_attn() -> None:
-    if NO_TORCH or IS_WINDOWS or IS_MACOS:
-        return
     if _flash_attn_install_disabled():
         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 (
         subprocess.run(
             [sys.executable, "-c", "import flash_attn"],

tests/python/test_flash_attn_install_python_stack.py

@@ -10,8 +10,133 @@ from unittest import mock
 
 STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
 sys.path.insert(0, str(STUDIO_DIR))
+sys.path.insert(0, str(STUDIO_DIR / "backend"))
 
 import install_python_stack as ips
+from backend.utils import wheel_utils
+
+
+def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess:
+    return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "")
+
+
+class TestHasBlackwellGpu:
+    def setup_method(self):
+        wheel_utils.has_blackwell_gpu.cache_clear()
+
+    def teardown_method(self):
+        wheel_utils.has_blackwell_gpu.cache_clear()
+
+    def test_returns_false_when_nvidia_smi_missing(self):
+        with mock.patch.object(wheel_utils.shutil, "which", return_value = None):
+            assert wheel_utils.has_blackwell_gpu() is False
+
+    def test_returns_true_for_sm_100(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is True
+
+    def test_returns_true_for_sm_120(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is True
+
+    def test_returns_true_for_sm_121(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is True
+
+    def test_returns_false_for_sm_90(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is False
+
+    def test_returns_false_for_sm_89(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is False
+
+    def test_mixed_gpus_with_one_blackwell_returns_true(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess,
+                "run",
+                return_value = _smi_result("8.0\n10.0\n"),
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is True
+
+    def test_returns_false_when_nvidia_smi_fails(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess,
+                "run",
+                return_value = _smi_result("", returncode = 1),
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is False
+
+    def test_returns_false_on_subprocess_timeout(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess,
+                "run",
+                side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10),
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is False
+
+    def test_returns_false_on_malformed_output(self):
+        with (
+            mock.patch.object(
+                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
+            ),
+            mock.patch.object(
+                wheel_utils.subprocess,
+                "run",
+                return_value = _smi_result("not-a-number\n\n"),
+            ),
+        ):
+            assert wheel_utils.has_blackwell_gpu() is False
 
 
 class TestFlashAttnWheelSelection:
@@ -234,6 +359,76 @@ class TestEnsureFlashAttn:
         mock_probe.assert_not_called()
         mock_install_wheel.assert_not_called()
 
+    def test_blackwell_gpu_skips_install_with_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", 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(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_blackwell_gpu_on_windows_emits_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 = 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(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 not any("Blackwell" in msg for _, msg in step_messages)
+
 
 class TestInstallPythonStackFlashAttnIntegration:
     def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int: