install-kernels: add flash-attn and install every kernel by default (#12744)

* install-kernels: add flash-attn and install every kernel by default

`unsloth install-kernels` with no names now installs xformers, flash-attn,
causal_conv1d and mamba_ssm in that order (mamba_ssm's fast path imports
causal_conv1d), each skipped when no wheel matches this torch / CUDA /
Python, so notebooks can call it once.

flash-attn resolves through wheel_utils.flash_attn_wheel_url, which covers
the upstream wheels up to torch 2.12 and the prebuilt-wheels-cu13 builds for
torch 2.13 / 2.14. It is skipped below sm80, where FlashAttention 2 does not
run and Unsloth never enables it, so a T4 does not download it.

* install-kernels: probe the GPU only when a flash-attn wheel exists, bounded, across all GPUs

Resolve the wheel first, so a torch probe that already failed or timed out
skips flash-attn without a second, unbounded torch import. The capability
probe now has a timeout and takes the best visible GPU, so a pre-Ampere
device 0 does not hide an Ampere card next to it.

* install-kernels: decode the GPU probe as UTF-8

studio/backend/tests/test_kernel_install.py

@@ -105,12 +105,19 @@ def test_causal_conv1d_has_no_wheel_off_linux_cuda(env):
 
 
 class _Runner:
-    def __init__(self, loads):
+    def __init__(
+        self,
+        loads,
+        capability = "9 0",
+    ):
         self.loads = list(loads)
+        self.capability = capability
         self.calls = []
 
     def __call__(self, cmd, **kwargs):
         self.calls.append(cmd)
+        if cmd[1] == "-c" and "get_device_capability" in cmd[2]:
+            return SimpleNamespace(returncode = 0, stdout = self.capability)
         if cmd[1] == "-c":
             return SimpleNamespace(returncode = 0 if self.loads.pop(0) else 1, stdout = "")
         return SimpleNamespace(returncode = 0, stdout = "")
@@ -306,3 +313,86 @@ def test_console_entry_drops_an_unwritable_ssl_keylog_file(tmp_path):
     )
     assert result.returncode == 0, result.stderr
     assert "ignoring SSLKEYLOGFILE" in result.stdout + result.stderr
+
+
+@pytest.mark.parametrize(
+    "env, expected",
+    [
+        (
+            _env("2.11.0+cu130", "13.0"),
+            "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",
+        ),
+        (
+            _env("2.13.0+cu130", "13.0"),
+            "https://github.com/unslothai/unsloth/releases/download/prebuilt-wheels-cu13/"
+            "flash_attn-2.8.4+cu13torch2.13cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
+        ),
+        (_env("2.11.0+cu130", "13.0", platform_tag = "win_amd64"), None),
+        (_env("2.11.0+cpu", ""), None),
+    ],
+)
+def test_flash_attn_resolution(env, expected):
+    assert kernel_install.resolve_wheel_url("flash_attn", env) == expected
+
+
+@pytest.mark.parametrize("capability", ["7 5", "None"])
+def test_flash_attn_is_skipped_below_sm80(capability, capsys):
+    run = _Runner([], capability = capability)
+    assert (
+        kernel_install.install_kernel("flash_attn", _COLAB, run = run, exists = lambda url: True) == 0
+    )
+    assert run.installer_calls == []
+    assert "skipping flash_attn" in capsys.readouterr().out
+
+
+def test_flash_attn_installs_on_sm80_and_newer(uv):
+    uv(False)
+    run = _Runner([False, True], capability = "8 0")
+    assert (
+        kernel_install.install_kernel("flash_attn", _COLAB, run = run, exists = lambda url: True) == 0
+    )
+    assert run.installer_calls[0][-1].endswith(
+        "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
+    )
+
+
+@pytest.mark.parametrize(
+    "argv, expected",
+    [
+        ([], ["xformers", "flash_attn", "causal_conv1d", "mamba_ssm"]),
+        (["mamba_ssm", "causal_conv1d", "mamba_ssm"], ["causal_conv1d", "mamba_ssm"]),
+    ],
+)
+def test_main_installs_all_by_default_in_dependency_order(monkeypatch, argv, expected):
+    order = []
+    monkeypatch.setattr(kernel_install, "probe_torch_wheel_env", lambda **kwargs: _COLAB)
+    monkeypatch.setattr(
+        kernel_install, "install_kernel", lambda name, env, dry_run = False: order.append(name) or 0
+    )
+    assert kernel_install.main(argv) == 0
+    assert order == expected
+
+
+def test_flash_attn_without_a_wheel_never_probes_the_gpu():
+    run = _Runner([])
+    assert kernel_install.install_kernel("flash_attn", None, run = run, exists = lambda url: True) == 0
+    assert run.calls == []
+
+
+def test_capability_probe_timeout_counts_as_no_gpu():
+    def run(cmd, **kwargs):
+        raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout"))
+
+    assert kernel_install._gpu_capability(run) is None
+
+
+def test_capability_probe_takes_the_best_visible_gpu():
+    seen = {}
+
+    def run(cmd, **kwargs):
+        seen["check"], seen["timeout"] = cmd[2], kwargs.get("timeout")
+        return SimpleNamespace(returncode = 0, stdout = "9 0\n")
+
+    assert kernel_install._gpu_capability(run) == (9, 0)
+    assert "max(torch.cuda.get_device_capability(i)" in seen["check"] and seen["timeout"]

studio/backend/utils/kernel_install.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""`unsloth install-kernels xformers causal_conv1d`: prebuilt kernel wheels matched to the resident torch.
+"""`unsloth install-kernels [names]`: prebuilt kernel wheels matched to the resident torch, all by default.
 
 Wheel-only by design: no wheel for this torch / CUDA / Python means nothing is installed and the
 model keeps its torch fallback, never a source build. Installs use --no-deps so torch is never touched.
@@ -23,6 +23,7 @@ from utils.wheel_utils import (
     MAMBA_SSM_RELEASE_BASE_URL,
     MAMBA_SSM_RELEASE_TAG,
     direct_wheel_url,
+    flash_attn_wheel_url,
     install_wheel,
     probe_torch_wheel_env,
     redact_url_credentials,
@@ -32,11 +33,13 @@ from utils.wheel_utils import (
 
 # name -> (pip distribution, check that loads the compiled extension). The package imports are no
 # proof: `import causal_conv1d` never loads it, and `import mamba_ssm` fails on einops under --no-deps.
+# Install order: mamba_ssm's fast path imports causal_conv1d.
 KERNELS = {
     "xformers": (
         "xformers",
         "from xformers._cpp_lib import _register_extensions; _register_extensions()",
     ),
+    "flash_attn": ("flash-attn", "import torch, flash_attn_2_cuda"),
     "causal_conv1d": ("causal-conv1d", "import torch, causal_conv1d_cuda"),
     "mamba_ssm": ("mamba-ssm", "import torch, selective_scan_cuda"),
 }
@@ -56,9 +59,11 @@ def resolve_wheel_url(name: str, env: dict[str, str] | None) -> str | None:
         return None
     if name == "xformers":
         return xformers_wheel_url(env)
-    # causal-conv1d and mamba-ssm publish Linux wheels only.
+    # flash-attn, causal-conv1d and mamba-ssm publish Linux wheels only.
     if not str(env.get("platform_tag") or "").startswith("linux"):
         return None
+    if name == "flash_attn":
+        return flash_attn_wheel_url(env)
     package_version, release_tag, release_base_url = _RELEASES[name]
     return direct_wheel_url(
         filename_prefix = name,
@@ -69,6 +74,31 @@ def resolve_wheel_url(name: str, env: dict[str, str] | None) -> str | None:
     )
 
 
+def _gpu_capability(run: Callable[..., subprocess.CompletedProcess]) -> tuple[int, int] | None:
+    """The best compute capability across the visible GPUs, or None."""
+    check = (
+        "import torch; n = torch.cuda.device_count() if torch.cuda.is_available() else 0; "
+        "print(*max(torch.cuda.get_device_capability(i) for i in range(n))) if n else None"
+    )
+    try:
+        result = run(
+            [sys.executable, "-c", check],
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            text = True,
+            encoding = "utf-8",
+            errors = "replace",
+            timeout = 120,
+        )
+    except subprocess.TimeoutExpired:
+        return None
+    try:
+        major, minor = (int(part) for part in result.stdout.split())
+    except ValueError:
+        return None
+    return major, minor
+
+
 def _loads(check: str, run: Callable[..., subprocess.CompletedProcess]) -> bool:
     result = run(
         [sys.executable, "-c", check], stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL
@@ -100,6 +130,17 @@ def install_kernel(
 ) -> int:
     distribution, check = KERNELS[name]
     url = resolve_wheel_url(name, env)
+    if name == "flash_attn" and url is not None:
+        # FlashAttention 2 needs Ampere or newer, and Unsloth only enables it there.
+        capability = _gpu_capability(run)
+        if capability is None or capability < (8, 0):
+            gpu = (
+                "no CUDA GPU is visible"
+                if capability is None
+                else "the best GPU is sm%d%d" % capability
+            )
+            print(f"Unsloth: skipping flash_attn, which needs sm80 or newer ({gpu}).")
+            return 0
     torch_desc = f"torch {env.get('torch_version')}" if env else "this environment"
     # Only a 404 proves nothing is published; an unreachable check falls through to the install.
     if url is None or exists(url) is False:
@@ -146,14 +187,19 @@ def main(argv: list[str] | None = None) -> int:
         prog = "unsloth install-kernels",
         description = "Install prebuilt kernel wheels matching the installed torch. Never builds from source or changes torch.",
     )
-    parser.add_argument("kernels", nargs = "+", choices = sorted(KERNELS))
+    parser.add_argument(
+        "kernels",
+        nargs = "*",
+        choices = list(KERNELS),
+        help = "kernels to install (default: all, each skipped when no wheel matches)",
+    )
     parser.add_argument(
         "--dry-run", action = "store_true", help = "print the resolved wheel URL and install nothing"
     )
     args = parser.parse_args(argv)
     env = probe_torch_wheel_env(timeout = 120, include_windows = True)
     status = 0
-    for name in dict.fromkeys(args.kernels):
+    for name in [name for name in KERNELS if not args.kernels or name in args.kernels]:
         status = max(status, install_kernel(name, env, dry_run = args.dry_run))
     return status
 

unsloth_cli/__init__.py

@@ -267,7 +267,7 @@ if not _windows_studio_mutation_entry:
         "ignore_unknown_options": True,
         "help_option_names": [],
     },
-    help = "Install prebuilt xformers / causal_conv1d / mamba_ssm wheels matching the installed torch.",
+    help = "Install prebuilt xformers / flash-attn / causal_conv1d / mamba_ssm wheels matching the installed torch (all by default).",
 )
 def install_kernels(ctx: typer.Context):
     # Listed for `unsloth --help`; the console script dispatches before typer is imported.