741bebb36a
Studio already falls back to the torch2.10 flash-attn / causal-conv1d / mamba-ssm wheels when it finds torch 2.11, because upstream publishes no 2.11-tagged builds. torch 2.12 is in exactly the same position, and the same wheels work there, so a 2.12 install currently drops to a source build for no reason. Measured on a B200, python 3.12, fresh uv venv on torch 2.12.1+cu130, wheels installed with --no-deps and torch verified unmoved afterwards, importing the compiled .so directly rather than only the Python package: causal-conv1d 1.6.1 9412 passed / 3888 skipped / 0 failed mamba-ssm 2.3.1 tests/ops, 20 passed flash-attn 2.8.1 splitkv + qkvpacked subset, 848 passed Against a torch 2.10 control the pass/fail/skip counts match and the failing test-ID sets are byte identical. The reuse window is bounded rather than open ended, so the comment now records that. flash-attn v2.8.3.post1's torch2.9 wheel fails to import on torch 2.10 and on torch 2.12 alike, with an undefined symbol out of flash_attn_2_cuda: torch broke extension ABI between 2.9 and 2.10 and has held it from 2.10 through 2.12. A wheel cannot skip a torch minor backwards, so torch 2.13 is deliberately left out of the table until it is measured. The torch2.10 flash-attn pin stays at 2.8.1. v2.8.3 looks like a free upgrade but publishes only 2 of the 8 torch2.10 assets that v2.8.1 does, keeping just cu13/cp312 for x86_64 and aarch64 and dropping every cu12 and every cp313 torch2.10 wheel, while v2.8.3.post1 dropped the torch2.10 assets entirely. Bumping the pin would silently 404 most users back to a source build, so the constant now carries that warning. Tests cover the 2.12 mapping through both direct_wheel_url and the flash-attn URL builder, that reuse only ever targets torch2.10, and that the selected flash-attn version is never a .post release.
395 lines
16 KiB
Python
395 lines
16 KiB
Python
"""Tests for the optional FlashAttention installer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
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 utils import wheel_utils
|
|
|
|
|
|
class TestPrebuiltWheelTorchMapping:
|
|
def test_torch_211_maps_to_torch210(self):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10"
|
|
|
|
def test_torch_212_maps_to_torch210(self):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm("2.12") == "2.10"
|
|
|
|
def test_other_versions_pass_through(self):
|
|
# 2.13 stays unmapped on purpose: a torch minor only joins the reuse
|
|
# table once its wheels have actually been measured.
|
|
for torch_mm in ("2.9", "2.10", "2.13"):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm
|
|
|
|
def test_reuse_never_targets_a_pre_210_wheel(self):
|
|
# torch broke extension ABI between 2.9 and 2.10, so the torch2.9 .so
|
|
# raises "undefined symbol" on 2.10+. Reuse may only point at torch2.10.
|
|
assert set(wheel_utils._PREBUILT_WHEEL_TORCH_MM.values()) == {"2.10"}
|
|
|
|
def test_direct_wheel_url_reuses_torch210_on_211(self):
|
|
# causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the
|
|
# torch2.10 wheel filename just like flash-attn does.
|
|
url = wheel_utils.direct_wheel_url(
|
|
filename_prefix = "causal_conv1d",
|
|
package_version = "1.6.1",
|
|
release_tag = "v1.6.1.post4",
|
|
release_base_url = "https://example.test/download",
|
|
env = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.11",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_direct_wheel_url_reuses_torch210_on_212(self):
|
|
url = wheel_utils.direct_wheel_url(
|
|
filename_prefix = "mamba_ssm",
|
|
package_version = "2.3.1",
|
|
release_tag = "v2.3.1",
|
|
release_base_url = "https://example.test/download",
|
|
env = {
|
|
"python_tag": "cp312",
|
|
"torch_mm": "2.12",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "mamba_ssm-2.3.1+cu13torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" in url
|
|
|
|
|
|
class TestFlashAttnWheelSelection:
|
|
def test_torch_210_maps_to_v281(self):
|
|
# v2.8.1 is the newest release still publishing the full torch2.10 asset
|
|
# matrix (cu12 + cu13, cp312 + cp313, x86_64 + aarch64).
|
|
assert ips._select_flash_attn_version("2.10") == "2.8.1"
|
|
|
|
def test_selected_version_is_never_a_post_release(self):
|
|
# The v2.8.3.post1 respin dropped every torch2.10 asset and stops at
|
|
# torch2.9, whose .so will not load on torch 2.10+. A future "just take
|
|
# the newest release" bump must fail here instead of shipping that.
|
|
for torch_mm in ("2.4", "2.7", "2.9", "2.10"):
|
|
version = ips._select_flash_attn_version(torch_mm)
|
|
assert version is not None
|
|
assert ".post" not in version
|
|
|
|
def test_torch_29_maps_to_v283(self):
|
|
assert ips._select_flash_attn_version("2.9") == "2.8.3"
|
|
|
|
def test_torch_211_has_no_native_version_entry(self):
|
|
# The raw version table has no torch2.11-tagged wheel; the URL builder
|
|
# reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel).
|
|
assert ips._select_flash_attn_version("2.11") is None
|
|
|
|
def test_torch_211_reuses_torch210_wheel(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.11",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_torch_212_reuses_torch210_wheel(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.12",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_exact_wheel_url_uses_full_env_tuple(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "v2.8.1" in url
|
|
assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_missing_cuda_major_disables_wheel_lookup(self):
|
|
assert (
|
|
ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
class TestEnsureFlashAttn:
|
|
def _import_check(self, code: int = 1):
|
|
return subprocess.CompletedProcess(["python", "-c", "import flash_attn"], code)
|
|
|
|
def test_prefers_exact_match_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(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"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
|
|
args, kwargs = install_calls[0]
|
|
assert args == (
|
|
"https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
|
|
)
|
|
assert kwargs["python_executable"] == sys.executable
|
|
assert kwargs["use_uv"] is True
|
|
assert kwargs["uv_needs_system"] is False
|
|
|
|
def test_uv_install_respects_system_flag(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", True),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"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
|
|
_, kwargs = install_calls[0]
|
|
assert kwargs["uv_needs_system"] is True
|
|
|
|
def test_wheel_failure_warns_and_continues(self):
|
|
step_messages: list[tuple[str, str]] = []
|
|
printed_failures: list[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, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [
|
|
("uv", subprocess.CompletedProcess(["uv"], 1, "uv wheel failed")),
|
|
(
|
|
"pip",
|
|
subprocess.CompletedProcess(["pip"], 1, "pip wheel failed"),
|
|
),
|
|
],
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"_print_optional_install_failure",
|
|
side_effect = lambda label, result: printed_failures.append(label),
|
|
),
|
|
mock.patch.object(ips, "_step", side_effect = fake_step),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
assert printed_failures == [
|
|
"Installing flash-attn prebuilt wheel with uv",
|
|
"Installing flash-attn prebuilt wheel with pip",
|
|
]
|
|
assert ("warning", "Continuing without flash-attn") in step_messages
|
|
|
|
def test_wheel_missing_skips_install_at_setup_time(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,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = False),
|
|
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_install_wheel.assert_not_called()
|
|
assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
|
|
|
|
def test_skip_env_disables_setup_install(self):
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.dict(os.environ, {"UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL": "1"}),
|
|
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_probe.assert_not_called()
|
|
mock_install_wheel.assert_not_called()
|
|
|
|
def test_windows_skips_install_without_probing(self):
|
|
# flash-attn is Linux-only: on Windows the installer returns before
|
|
# probing the torch env or resolving a wheel (no Windows wheels are
|
|
# published upstream).
|
|
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, "probe_torch_wheel_env") as mock_probe,
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_probe.assert_not_called()
|
|
mock_install_wheel.assert_not_called()
|
|
|
|
|
|
class TestInstallPythonStackFlashAttnIntegration:
|
|
def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int:
|
|
flash_attn_calls = 0
|
|
|
|
def fake_run(cmd, **kw):
|
|
return subprocess.CompletedProcess(cmd, 0, b"", b"")
|
|
|
|
def count_flash_attn():
|
|
nonlocal flash_attn_calls
|
|
flash_attn_calls += 1
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", no_torch),
|
|
mock.patch.object(ips, "IS_MACOS", is_macos),
|
|
mock.patch.object(ips, "IS_WINDOWS", is_windows),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(ips, "VERBOSE", False),
|
|
mock.patch.object(ips, "_bootstrap_uv", return_value = True),
|
|
mock.patch.object(ips, "_ensure_flash_attn", side_effect = count_flash_attn),
|
|
mock.patch("subprocess.run", side_effect = fake_run),
|
|
mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False),
|
|
mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
|
|
mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
|
|
mock.patch("pathlib.Path.is_dir", return_value = True),
|
|
mock.patch("pathlib.Path.is_file", return_value = True),
|
|
mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False),
|
|
):
|
|
ips.install_python_stack()
|
|
|
|
return flash_attn_calls
|
|
|
|
def test_linux_torch_install_calls_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = False, is_windows = False) == 1
|
|
|
|
def test_no_torch_install_skips_flash_attn_step(self):
|
|
assert self._run_install(no_torch = True, is_macos = False, is_windows = False) == 0
|
|
|
|
def test_macos_install_skips_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = True, is_windows = False) == 0
|
|
|
|
def test_windows_install_skips_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = False, is_windows = True) == 0
|