diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..7961f200a2 100644 --- a/studio/backend/core/training/worker.py +++ b/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 diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 14fc0933d0..56d2d1aab8 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/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) diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..cb21999571 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/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) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..dc456450d6 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/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, *, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e033a56a0a..263b71689a 100644 --- a/studio/install_python_stack.py +++ b/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( diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 26ff03505a..752d12f95c 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/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),