mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-21 14:00:32 +00:00
4b8447f63a
Linux users on torch 2.10, 2.11 or 2.12 still get upstream's causal-conv1d and mamba-ssm wheels: 205 MB and 351 MB downloads (torch2.10 builds, reused on 2.11 and 2.12), with `-lineinfo` and every arch. This lets the prebuilt wheel workflow build our trimmed versions for those cells, on both linux_x86_64 and linux_aarch64. ## Changes - `linux_x86_64` and `linux_aarch64` gain `extra_torches` 2.10.0 / 2.11.0 / 2.12.0, restricted to causal-conv1d and mamba-ssm by a new `extra_torch_packages` key. flash-attn is not built there: upstream's 2.8.1 for torch 2.10 is already smaller (125 MB, no PTX) than ours. - `_platform_packages(platform, torch)` is the one place both `build_matrix` and the `wheel-name` CLI read, so the CLI now also refuses a package a platform does not build (it used to print a name for flash-attn on win_arm64). - Release notes describe packages with different coverage separately, so Linux does not read as having flash-attn for torch 2.10 to 2.12. - A default dispatch builds exactly what it did before (12 cells); the older torches stay opt-in. ## Rollout The Studio resolver does not ask for these cells yet (`_UNSLOTH_PREBUILT_EXTRA_CELLS` is unchanged, and the existing test keeps it a subset of what the workflow builds). After this merges: dispatch `platforms=linux_x86_64,linux_aarch64`, `packages=causal-conv1d,mamba-ssm`, `torch_versions=2.10.0,2.11.0,2.12.0`, `python_versions=3.11,3.12,3.13` (36 cells), then a Studio change adds the Linux cells to the resolver for those two packages only. Expected sizes from the torch 2.13 / 2.14 Linux builds just published from the same recipe: causal-conv1d about 29 MB (x86_64) / 42 MB (aarch64) download, mamba-ssm about 54 / 79 MB. ## Tests `tests/python/test_prebuilt_cuda_wheels*.py`: 256 passed. The updated matrix and CLI tests fail on main (Linux refused every torch below 2.13); a new test covers the release notes split.
837 lines
37 KiB
Python
837 lines
37 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""The win_amd64 half of .github/workflows/prebuilt-cuda-wheels.yml.
|
|
|
|
Nobody publishes flash-attn, causal-conv1d or mamba-ssm for Windows: not one Windows asset exists
|
|
across every Dao-AILab and state-spaces release. These pin the parts of the Windows legs that can
|
|
be checked without a Windows runner:
|
|
|
|
* the Linux cells are byte-for-byte what they were, only gaining keys;
|
|
* Windows names follow each upstream setup.py's own naming (`cxx11abiTRUE`, as torch reports, and `win_amd64`);
|
|
* the shard selector understands Windows ninja targets, and can never run a targetless ninja;
|
|
* the deadline wrapper ends the whole process tree;
|
|
* SHA256SUMS is rebuilt over every wheel on the release, so a one-platform publish keeps the other;
|
|
* the workflow's Windows steps exist in warm and build alike, are gated on the OS, and keep the
|
|
supply-chain checks (Authenticode on the CUDA installer, a pinned sha256 on ccache).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
REPO = Path(__file__).resolve().parents[2]
|
|
WORKFLOW = REPO / ".github" / "workflows" / "prebuilt-cuda-wheels.yml"
|
|
SCRIPTS = REPO / ".github" / "scripts"
|
|
|
|
|
|
def _load(name: str, path: Path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
prebuilt_wheels = _load("prebuilt_wheels_win", SCRIPTS / "prebuilt_wheels.py")
|
|
shard = _load("prebuilt_wheels_shard_win", SCRIPTS / "prebuilt_wheels_shard.py")
|
|
deadline = _load("prebuilt_wheels_timeout_win", SCRIPTS / "prebuilt_wheels_timeout.py")
|
|
|
|
|
|
def triggers(doc: dict) -> dict:
|
|
return doc.get(True) if True in doc else doc.get("on")
|
|
|
|
|
|
@pytest.fixture(scope = "module")
|
|
def workflow():
|
|
return yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8"))
|
|
|
|
|
|
def _steps(workflow, job):
|
|
return workflow["jobs"][job]["steps"]
|
|
|
|
|
|
def _step(workflow, job, name):
|
|
return next(s for s in _steps(workflow, job) if s.get("name") == name)
|
|
|
|
|
|
# Linux cells as published before Windows existed; changing them changes released wheels.
|
|
LINUX_BEFORE = {
|
|
"flash-attn": {
|
|
"build_env": "FLASH_ATTENTION_FORCE_BUILD=TRUE FLASH_ATTENTION_FORCE_CXX11_ABI=TRUE "
|
|
"FLASH_ATTN_CUDA_ARCHS=80;90;100;120",
|
|
"abi": "TRUE",
|
|
"shards": 8,
|
|
"label": "flash-attn / torch 2.13 / cp313",
|
|
"wheel_name": "flash_attn-2.8.4+cu13torch2.13cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
|
|
},
|
|
"causal-conv1d": {
|
|
"build_env": "CAUSAL_CONV1D_FORCE_BUILD=TRUE",
|
|
"abi": "TRUE",
|
|
"shards": 0,
|
|
"label": "causal-conv1d / torch 2.13 / cp313",
|
|
"wheel_name": "causal_conv1d-1.7.0+cu13torch2.13cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
|
|
},
|
|
"mamba-ssm": {
|
|
"build_env": "MAMBA_FORCE_BUILD=TRUE MAMBA_KEEP_CUDA_BUILD=TRUE",
|
|
"abi": "TRUE",
|
|
"shards": 0,
|
|
"label": "mamba-ssm / torch 2.13 / cp313",
|
|
"wheel_name": "mamba_ssm-2.3.2.post1+cu13torch2.13cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
|
|
},
|
|
}
|
|
|
|
|
|
class TestMatrix:
|
|
def test_linux_cells_are_unchanged(self):
|
|
cells = prebuilt_wheels.build_matrix(torches = "2.13.0", platforms = "linux_x86_64")
|
|
assert [c["package"] for c in cells] == list(prebuilt_wheels.SPECS)
|
|
for cell in cells:
|
|
for key, value in LINUX_BEFORE[cell["package"]].items():
|
|
assert cell[key] == value, (cell["package"], key)
|
|
assert cell["runner"] == "ubuntu-22.04"
|
|
assert cell["pytorch_nvcc"] == "ccache /usr/local/cuda-13.0/bin/nvcc"
|
|
assert cell["smoke_exclude"] == cell["smoke_extra"] == ""
|
|
|
|
def test_linux_comes_first_so_its_cells_keep_their_order(self):
|
|
cells = prebuilt_wheels.build_matrix()
|
|
platforms = [c["platform"] for c in cells]
|
|
assert platforms == sorted(platforms, key = list(prebuilt_wheels.PLATFORMS).index)
|
|
|
|
def test_windows_cells(self):
|
|
cells = prebuilt_wheels.build_matrix(platforms = "win_amd64")
|
|
assert len(cells) == 6
|
|
for cell in cells:
|
|
assert cell["os"] == "windows"
|
|
assert cell["runner"] == "windows-2022"
|
|
assert cell["abi"] == "TRUE"
|
|
assert cell["wheel_name"].endswith("cxx11abiTRUE-cp313-cp313-win_amd64.whl")
|
|
assert cell["label"].endswith(" / win_amd64")
|
|
assert cell["pytorch_nvcc"] == "C:/ccache/ccache.exe C:/cuda/bin/nvcc.exe"
|
|
|
|
def test_windows_flash_attn_names_the_abi_torch_reports(self):
|
|
"""The Windows cu130 torch reports _GLIBCXX_USE_CXX11_ABI=True (measured on windows-2022),
|
|
and FORCE_CXX11_ABI follows the platform's `abi`, so the name and the assert agree."""
|
|
(cell,) = prebuilt_wheels.build_matrix(
|
|
packages = "flash-attn", torches = "2.14.0", platforms = "win_amd64"
|
|
)
|
|
assert "FLASH_ATTENTION_FORCE_CXX11_ABI=TRUE" in cell["build_env"].split()
|
|
assert "FLASH_ATTN_CUDA_ARCHS=80;90;100;120" in cell["build_env"].split()
|
|
assert cell["shards"] == 8
|
|
|
|
def test_only_windows_mamba_swaps_its_smoke_dependencies(self):
|
|
for cell in prebuilt_wheels.build_matrix():
|
|
if cell["os"] == "windows" and cell["package"] == "mamba-ssm":
|
|
pin = prebuilt_wheels.TRITON_WINDOWS[cell["torch_mm"]]
|
|
assert cell["smoke_extra"] == f"triton-windows=={pin}"
|
|
assert set(cell["smoke_exclude"].split()) == {
|
|
"triton",
|
|
"tilelang",
|
|
"apache-tvm-ffi",
|
|
"quack-kernels",
|
|
}
|
|
else:
|
|
assert cell["smoke_extra"] == cell["smoke_exclude"] == ""
|
|
|
|
def test_every_torch_minor_has_a_triton_windows_pin(self):
|
|
"""torch 2.10 / 2.11 pin triton 3.6.0, 2.12 3.7.0, 2.13 3.7.1, 2.14 3.8.x; triton-windows follows."""
|
|
minors = {
|
|
prebuilt_wheels.torch_minor(v)
|
|
for platform, target in prebuilt_wheels.PLATFORMS.items()
|
|
if target["os"] == "windows"
|
|
for v in prebuilt_wheels._platform_torches(platform)
|
|
}
|
|
assert set(prebuilt_wheels.TRITON_WINDOWS) == minors
|
|
assert prebuilt_wheels.TRITON_WINDOWS["2.10"].startswith("3.6.0.post")
|
|
assert prebuilt_wheels.TRITON_WINDOWS["2.11"].startswith("3.6.0.post")
|
|
assert prebuilt_wheels.TRITON_WINDOWS["2.12"].startswith("3.7.0.post")
|
|
assert prebuilt_wheels.TRITON_WINDOWS["2.13"].startswith("3.7.1.post")
|
|
assert prebuilt_wheels.TRITON_WINDOWS["2.14"].startswith("3.8.0.post")
|
|
|
|
@pytest.mark.parametrize("torch_version", ["2.10.0", "2.11.0", "2.12.0"])
|
|
def test_older_torches_are_opt_in(self, torch_version):
|
|
torch_mm = prebuilt_wheels.torch_minor(torch_version)
|
|
cells = prebuilt_wheels.build_matrix(
|
|
torches = torch_version, platforms = "win_amd64", pythons = "3.11,3.12,3.13"
|
|
)
|
|
assert len(cells) == 9
|
|
assert {cell["python_tag"] for cell in cells} == {"cp311", "cp312", "cp313"}
|
|
assert {cell["torch_spec"] for cell in cells} == {f"torch=={torch_version}"}
|
|
for cell in cells:
|
|
assert f"+cu13torch{torch_mm}cxx11abiTRUE-" in cell["wheel_name"]
|
|
if cell["package"] == "mamba-ssm":
|
|
triton = prebuilt_wheels.TRITON_WINDOWS[torch_mm]
|
|
assert cell["smoke_extra"] == f"triton-windows=={triton}"
|
|
with pytest.raises(SystemExit):
|
|
prebuilt_wheels.build_matrix(torches = torch_version, platforms = "win_arm64")
|
|
# Linux builds causal-conv1d and mamba-ssm only: upstream flash-attn 2.8.1 is smaller there.
|
|
for platform in ("linux_x86_64", "linux_aarch64"):
|
|
with pytest.raises(SystemExit):
|
|
prebuilt_wheels.build_matrix(
|
|
torches = torch_version, platforms = platform, packages = "flash-attn"
|
|
)
|
|
linux = prebuilt_wheels.build_matrix(
|
|
torches = torch_version, platforms = platform, pythons = "3.11,3.12,3.13"
|
|
)
|
|
assert len(linux) == 6
|
|
assert {c["package"] for c in linux} == {"causal-conv1d", "mamba-ssm"}
|
|
assert {c["python_tag"] for c in linux} == {"cp311", "cp312", "cp313"}
|
|
for cell in linux:
|
|
assert cell["wheel_name"].endswith(f"-{platform}.whl")
|
|
assert f"+cu13torch{torch_mm}cxx11abiTRUE-" in cell["wheel_name"]
|
|
# A default dispatch builds what it did before.
|
|
assert torch_version not in {cell["torch"] for cell in prebuilt_wheels.build_matrix()}
|
|
both = prebuilt_wheels.build_matrix(torches = f"{torch_version},2.13.0")
|
|
assert {(c["platform"], c["torch"], c["package"]) for c in both} == {
|
|
*(("win_amd64", torch_version, p) for p in prebuilt_wheels.SPECS),
|
|
*(("win_amd64", "2.13.0", p) for p in prebuilt_wheels.SPECS),
|
|
*(("linux_x86_64", "2.13.0", p) for p in prebuilt_wheels.SPECS),
|
|
("linux_x86_64", torch_version, "causal-conv1d"),
|
|
("linux_x86_64", torch_version, "mamba-ssm"),
|
|
}
|
|
|
|
def test_win_amd64_inherits_a_new_global_torch(self, monkeypatch):
|
|
monkeypatch.setattr(prebuilt_wheels, "TORCH_VERSIONS", ("2.13.0", "2.14.0", "2.15.0"))
|
|
assert prebuilt_wheels._platform_torches("win_amd64") == (
|
|
"2.10.0",
|
|
"2.11.0",
|
|
"2.12.0",
|
|
"2.13.0",
|
|
"2.14.0",
|
|
"2.15.0",
|
|
)
|
|
assert prebuilt_wheels._platform_torches("linux_x86_64")[-1] == "2.15.0"
|
|
|
|
def test_pytorch_nvcc_paths_have_no_spaces(self):
|
|
"""torch writes PYTORCH_NVCC into build.ninja verbatim; a space would split a path."""
|
|
for target in prebuilt_wheels.PLATFORMS.values():
|
|
words = target["pytorch_nvcc"].split(" ")
|
|
assert len(words) == 2 and all(words), target
|
|
|
|
@pytest.mark.parametrize(
|
|
"platforms", ["win_arm64ec", "macosx_arm64", "win_amd64; whoami", "$(id)"]
|
|
)
|
|
def test_unknown_platforms_are_refused(self, platforms):
|
|
with pytest.raises(SystemExit):
|
|
prebuilt_wheels.build_matrix(platforms = platforms)
|
|
|
|
def test_the_gpu_runner_only_gets_linux_cells(self):
|
|
cells = prebuilt_wheels.build_matrix()
|
|
gpu = prebuilt_wheels.gpu_matrix(cells)
|
|
assert gpu and all(c["os"] == "linux" for c in gpu)
|
|
assert len(gpu) == len([c for c in cells if c["platform"] == "linux_x86_64"])
|
|
assert prebuilt_wheels.gpu_matrix(prebuilt_wheels.build_matrix(platforms = "win_amd64")) == []
|
|
|
|
def test_a_windows_only_plan_writes_an_empty_gpu_matrix(self, tmp_path):
|
|
output = tmp_path / "output"
|
|
subprocess.run(
|
|
[sys.executable, str(SCRIPTS / "prebuilt_wheels.py"), "matrix", "--github"],
|
|
env = {"GITHUB_OUTPUT": str(output), "UW_PLATFORMS": "win_amd64"},
|
|
check = True,
|
|
capture_output = True,
|
|
)
|
|
values = dict(line.split("=", 1) for line in output.read_text().splitlines())
|
|
assert values["count"] == "6"
|
|
assert values["gpu_count"] == "0"
|
|
assert values["warm_count"] == "16"
|
|
warm = json.loads(values["warm_matrix"])["include"]
|
|
assert {w["platform"] for w in warm} == {"win_amd64"}
|
|
|
|
|
|
class TestNames:
|
|
def test_windows_names_are_what_upstream_setup_py_would_write(self):
|
|
"""flash-attn, causal-conv1d and mamba's get_wheel_url() all format
|
|
`+cu{cuda}torch{mm}cxx11abi{torch._C._GLIBCXX_USE_CXX11_ABI}` and `win_amd64`."""
|
|
assert (
|
|
prebuilt_wheels.wheel_name("flash-attn", "2.13.0", "3.13", "win_amd64")
|
|
== "flash_attn-2.8.4+cu13torch2.13cxx11abiTRUE-cp313-cp313-win_amd64.whl"
|
|
)
|
|
assert (
|
|
prebuilt_wheels.wheel_name("causal-conv1d", "2.14.0", "3.13", "win_amd64")
|
|
== "causal_conv1d-1.7.0+cu13torch2.14cxx11abiTRUE-cp313-cp313-win_amd64.whl"
|
|
)
|
|
assert (
|
|
prebuilt_wheels.wheel_name("mamba-ssm", "2.14.0", "3.12", "win_amd64")
|
|
== "mamba_ssm-2.3.2.post1+cu13torch2.14cxx11abiTRUE-cp312-cp312-win_amd64.whl"
|
|
)
|
|
|
|
def test_the_default_platform_is_still_linux(self):
|
|
assert prebuilt_wheels.wheel_name("flash-attn", "2.13.0", "3.13") == (
|
|
prebuilt_wheels.wheel_name("flash-attn", "2.13.0", "3.13", "linux_x86_64")
|
|
)
|
|
|
|
def test_wheel_name_cli_takes_a_platform(self):
|
|
out = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(SCRIPTS / "prebuilt_wheels.py"),
|
|
"wheel-name",
|
|
"--package",
|
|
"causal-conv1d",
|
|
"--torch",
|
|
"2.13.0",
|
|
"--python",
|
|
"3.13",
|
|
"--platform",
|
|
"win_amd64",
|
|
],
|
|
check = True,
|
|
capture_output = True,
|
|
text = True,
|
|
).stdout.strip()
|
|
assert out == "causal_conv1d-1.7.0+cu13torch2.13cxx11abiTRUE-cp313-cp313-win_amd64.whl"
|
|
|
|
def test_wheel_name_cli_refuses_a_cell_nothing_builds(self):
|
|
def name(package, *platform):
|
|
return subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(SCRIPTS / "prebuilt_wheels.py"),
|
|
"wheel-name",
|
|
"--package",
|
|
package,
|
|
"--torch",
|
|
"2.11.0",
|
|
"--python",
|
|
"3.12",
|
|
*platform,
|
|
],
|
|
capture_output = True,
|
|
text = True,
|
|
)
|
|
|
|
# win_arm64 builds torch 2.14 only; Linux builds no flash-attn for torch 2.10 to 2.12.
|
|
for package, platform in (
|
|
("mamba-ssm", ["--platform", "win_arm64"]),
|
|
("flash-attn", []),
|
|
("flash-attn", ["--platform", "linux_x86_64"]),
|
|
("flash-attn", ["--platform", "linux_aarch64"]),
|
|
):
|
|
result = name(package, *platform)
|
|
assert result.returncode != 0, (package, platform)
|
|
assert "not built for" in result.stderr
|
|
for package, platform in (
|
|
("flash-attn", "win_amd64"),
|
|
("mamba-ssm", "win_amd64"),
|
|
("mamba-ssm", "linux_x86_64"),
|
|
("causal-conv1d", "linux_aarch64"),
|
|
):
|
|
result = name(package, "--platform", platform)
|
|
assert result.returncode == 0, (package, platform)
|
|
assert result.stdout.strip().endswith(
|
|
f"torch2.11cxx11abiTRUE-cp312-cp312-{platform}.whl"
|
|
)
|
|
|
|
def test_round_trip_on_both_platforms(self):
|
|
for platform in prebuilt_wheels.PLATFORMS:
|
|
for package in prebuilt_wheels.SPECS:
|
|
name = prebuilt_wheels.wheel_name(package, "2.14.0", "3.13", platform)
|
|
parsed = prebuilt_wheels.parse_wheel_name(name)
|
|
assert parsed["package"] == package
|
|
assert parsed["platform"] == platform
|
|
assert parsed["abi"] == prebuilt_wheels.PLATFORMS[platform]["abi"]
|
|
assert parsed["version"] == prebuilt_wheels.SPECS[package]["version"]
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"flash_attn-2.8.4+cu13torch2.13cxx11abiFALSE-cp313-cp313-win_amd64.whl",
|
|
"flash_attn-2.8.4+cu13torch2.13cxx11abiFALSE-cp313-cp313-linux_x86_64.whl",
|
|
"flash_attn-2.8.4+cu13torch2.14-cp313-cp313-win_arm64.whl",
|
|
],
|
|
)
|
|
def test_an_abi_the_platform_does_not_have_is_not_ours(self, name):
|
|
assert prebuilt_wheels.parse_wheel_name(name) is None
|
|
|
|
def test_the_version_comes_from_the_file(self):
|
|
parsed = prebuilt_wheels.parse_wheel_name(
|
|
"flash_attn-9.9.9+cu13torch2.13cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
|
|
)
|
|
assert parsed["version"] == "9.9.9"
|
|
|
|
|
|
def _entries(platforms, torches = ("2.13.0", "2.14.0")):
|
|
return [
|
|
(f"{index:064x}", prebuilt_wheels.wheel_name(package, torch, "3.13", platform))
|
|
for index, (platform, torch, package) in enumerate(
|
|
(platform, torch, package)
|
|
for platform in platforms
|
|
for torch in torches
|
|
for package in prebuilt_wheels.SPECS
|
|
)
|
|
]
|
|
|
|
|
|
class TestNotes:
|
|
def test_both_platforms(self):
|
|
notes = prebuilt_wheels.render_notes(
|
|
_entries(["linux_x86_64", "win_amd64"]), tag = "t", repo = "o/r"
|
|
)
|
|
assert notes == (
|
|
"Prebuilt CUDA 13 wheels. Linux x86_64: flash-attn 2.8.4, causal-conv1d 1.7.0 and "
|
|
"mamba-ssm 2.3.2.post1, built for PyTorch 2.13 and 2.14 on Python 3.13. "
|
|
"Windows x86_64: flash-attn 2.8.4, causal-conv1d 1.7.0 and mamba-ssm 2.3.2.post1, "
|
|
"built for PyTorch 2.13 and 2.14 on Python 3.13.\n"
|
|
)
|
|
|
|
def test_packages_with_different_coverage_are_described_apart(self):
|
|
"""Linux torch 2.10 to 2.12 carries causal-conv1d and mamba-ssm but no flash-attn."""
|
|
entries = [
|
|
(f"{index:064x}", prebuilt_wheels.wheel_name(package, torch, "3.13", "linux_x86_64"))
|
|
for index, (package, torch) in enumerate(
|
|
[(pkg, "2.10.0") for pkg in ("causal-conv1d", "mamba-ssm")]
|
|
+ [(pkg, "2.13.0") for pkg in prebuilt_wheels.SPECS]
|
|
)
|
|
]
|
|
notes = prebuilt_wheels.render_notes(entries, tag = "t", repo = "r")
|
|
assert notes == (
|
|
"Prebuilt Linux x86_64 CUDA 13 wheels for flash-attn 2.8.4, built for PyTorch 2.13 on "
|
|
"Python 3.13; causal-conv1d 1.7.0 and mamba-ssm 2.3.2.post1, built for PyTorch 2.10 "
|
|
"and 2.13 on Python 3.13.\n"
|
|
)
|
|
|
|
def test_torch_minors_with_different_interpreters_are_not_merged(self):
|
|
entries = [
|
|
(f"{index:064x}", prebuilt_wheels.wheel_name(package, torch, python, "win_amd64"))
|
|
for index, (package, torch, python) in enumerate(
|
|
[
|
|
(pkg, t, py)
|
|
for t in ("2.10.0", "2.11.0", "2.12.0")
|
|
for py in ("3.11", "3.12", "3.13")
|
|
for pkg in prebuilt_wheels.SPECS
|
|
]
|
|
+ [(pkg, t, "3.13") for t in ("2.13.0", "2.14.0") for pkg in prebuilt_wheels.SPECS]
|
|
)
|
|
]
|
|
notes = prebuilt_wheels.render_notes(entries, tag = "t", repo = "r")
|
|
assert (
|
|
"built for PyTorch 2.10, 2.11 and 2.12 on Python 3.11, 3.12 and 3.13; 2.13 and 2.14 "
|
|
"on Python 3.13." in notes
|
|
)
|
|
|
|
def test_platforms_with_different_coverage_are_not_merged(self):
|
|
entries = _entries(["linux_x86_64"]) + _entries(["win_amd64"], torches = ("2.14.0",))
|
|
notes = prebuilt_wheels.render_notes(entries, tag = "t", repo = "o/r")
|
|
assert "built for PyTorch 2.13 and 2.14 on Python 3.13. Windows" in notes
|
|
assert notes.endswith("built for PyTorch 2.14 on Python 3.13.\n")
|
|
|
|
def test_packages_are_listed_per_platform(self):
|
|
entries = _entries(["linux_x86_64"]) + [
|
|
row for row in _entries(["win_amd64"]) if row[1].startswith("mamba_ssm-")
|
|
]
|
|
notes = prebuilt_wheels.render_notes(entries, tag = "t", repo = "o/r")
|
|
windows = notes.split("Windows x86_64: ", 1)[1]
|
|
assert windows == "mamba-ssm 2.3.2.post1, built for PyTorch 2.13 and 2.14 on Python 3.13.\n"
|
|
|
|
def test_windows_only(self):
|
|
notes = prebuilt_wheels.render_notes(_entries(["win_amd64"]), tag = "t", repo = "o/r")
|
|
assert notes.startswith("Prebuilt Windows x86_64 CUDA 13 wheels for flash-attn 2.8.4")
|
|
|
|
def test_linux_only_reads_as_before(self):
|
|
notes = prebuilt_wheels.render_notes(_entries(["linux_x86_64"]), tag = "t", repo = "o/r")
|
|
assert notes == (
|
|
"Prebuilt Linux x86_64 CUDA 13 wheels for flash-attn 2.8.4, causal-conv1d 1.7.0 and "
|
|
"mamba-ssm 2.3.2.post1, built for PyTorch 2.13 and 2.14 on Python 3.13.\n"
|
|
)
|
|
|
|
|
|
class TestSums:
|
|
def test_every_wheel_sorted_by_name(self):
|
|
entries = _entries(["win_amd64", "linux_x86_64"])
|
|
entries.append(("sha256:" + "f" * 64, "SHA256SUMS"))
|
|
entries.append(("a" * 64, "x.whl.sigstore.json"))
|
|
text = prebuilt_wheels.render_sums(entries)
|
|
lines = text.splitlines()
|
|
assert len(lines) == 12
|
|
names = [line.split(" ", 1)[1] for line in lines]
|
|
assert names == sorted(names)
|
|
assert all(name.endswith(".whl") for name in names)
|
|
assert {n.rsplit("-", 1)[1] for n in names} == {"linux_x86_64.whl", "win_amd64.whl"}
|
|
|
|
def test_the_sha256_prefix_is_dropped(self):
|
|
text = prebuilt_wheels.render_sums([("sha256:" + "AB" * 32, "a.whl")])
|
|
assert text == "ab" * 32 + " a.whl\n"
|
|
|
|
@pytest.mark.parametrize("digest", ["unknown", "sha256:unknown", "-", "", "z" * 64, "a" * 63])
|
|
def test_a_missing_digest_is_refused(self, digest):
|
|
with pytest.raises(SystemExit):
|
|
prebuilt_wheels.render_sums([(digest, "a.whl")])
|
|
|
|
def test_nothing_to_list_is_refused(self):
|
|
with pytest.raises(SystemExit):
|
|
prebuilt_wheels.render_sums([("a" * 64, "SHA256SUMS")])
|
|
|
|
def test_cli(self):
|
|
stdin = f"{'b' * 64} z.whl\n{'a' * 64} a.whl\n"
|
|
out = subprocess.run(
|
|
[sys.executable, str(SCRIPTS / "prebuilt_wheels.py"), "sums"],
|
|
input = stdin,
|
|
check = True,
|
|
capture_output = True,
|
|
text = True,
|
|
).stdout
|
|
assert out == f"{'a' * 64} a.whl\n{'b' * 64} z.whl\n"
|
|
|
|
|
|
WINDOWS_LISTING = "\n".join(
|
|
[
|
|
rf"D:\a\unsloth\unsloth\src\build\temp.win-amd64-cpython-313\Release\csrc\flash_attn\src\k_{i:03d}.obj: cuda_compile"
|
|
for i in range(20)
|
|
]
|
|
+ [
|
|
r"D:\a\unsloth\unsloth\src\build\temp.win-amd64-cpython-313\Release\csrc\flash_attn\flash_api.obj: compile",
|
|
r"D:\a\unsloth\unsloth\src\build\lib.win-amd64-cpython-313\flash_attn_2_cuda.cp313-win_amd64.pyd: link",
|
|
"all: phony",
|
|
]
|
|
)
|
|
|
|
|
|
class TestShards:
|
|
def test_windows_targets_are_whole_paths(self):
|
|
objects = shard.list_objects(WINDOWS_LISTING)
|
|
assert len(objects) == 21
|
|
assert all(o.startswith("D:\\a\\") and o.endswith(".obj") for o in objects)
|
|
|
|
def test_windows_slices_cover_everything_once(self):
|
|
slices = [shard.slice_objects(WINDOWS_LISTING, k, 8) for k in range(8)]
|
|
assert sorted(o for s in slices for o in s) == shard.list_objects(WINDOWS_LISTING)
|
|
|
|
def test_linux_targets_are_unchanged(self):
|
|
listing = "/src/build/temp/a.o: cuda_compile\n/src/build/temp/b.o: compile\nall: phony"
|
|
assert shard.list_objects(listing) == ["/src/build/temp/a.o", "/src/build/temp/b.o"]
|
|
|
|
def test_no_objects_is_an_error_not_a_full_build(self):
|
|
with pytest.raises(SystemExit) as raised:
|
|
shard.slice_objects("all: phony\nfoo.pyd: link", 0, 8)
|
|
assert raised.value.code != 0
|
|
|
|
def test_an_empty_slice_never_runs_ninja(self, tmp_path, monkeypatch):
|
|
"""Two objects across eight shards leaves six empty slices. `ninja -v -j 1` with no
|
|
target would build the default target, the whole extension, in every one of them."""
|
|
calls = []
|
|
|
|
def fake_run(command, **kwargs):
|
|
calls.append(command)
|
|
return types.SimpleNamespace(stdout = "a.obj: cuda_compile\nb.obj: cuda_compile\n")
|
|
|
|
cpp_extension = types.ModuleType("torch.utils.cpp_extension")
|
|
cpp_extension._run_ninja_build = None
|
|
torch = types.ModuleType("torch")
|
|
utils = types.ModuleType("torch.utils")
|
|
torch.utils = utils
|
|
utils.cpp_extension = cpp_extension
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
monkeypatch.setitem(sys.modules, "torch.utils", utils)
|
|
monkeypatch.setitem(sys.modules, "torch.utils.cpp_extension", cpp_extension)
|
|
monkeypatch.setattr(shard.subprocess, "run", fake_run)
|
|
(tmp_path / "setup.py").write_text(
|
|
"import torch.utils.cpp_extension as c\nc._run_ninja_build('.', False, '')\n"
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr(sys, "argv", ["prebuilt_wheels_shard.py", "5", "8"])
|
|
with pytest.raises(SystemExit) as raised:
|
|
shard.main()
|
|
assert raised.value.code == 0
|
|
assert calls == [["ninja", "-t", "targets", "all"]]
|
|
|
|
def test_a_non_empty_slice_names_its_targets(self, tmp_path, monkeypatch):
|
|
calls = []
|
|
|
|
def fake_run(command, **kwargs):
|
|
calls.append(command)
|
|
return types.SimpleNamespace(stdout = "a.obj: cuda_compile\nb.obj: cuda_compile\n")
|
|
|
|
cpp_extension = types.ModuleType("torch.utils.cpp_extension")
|
|
torch = types.ModuleType("torch")
|
|
utils = types.ModuleType("torch.utils")
|
|
torch.utils = utils
|
|
utils.cpp_extension = cpp_extension
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
monkeypatch.setitem(sys.modules, "torch.utils", utils)
|
|
monkeypatch.setitem(sys.modules, "torch.utils.cpp_extension", cpp_extension)
|
|
monkeypatch.setattr(shard.subprocess, "run", fake_run)
|
|
monkeypatch.setenv("MAX_JOBS", "1")
|
|
(tmp_path / "setup.py").write_text(
|
|
"import torch.utils.cpp_extension as c\nc._run_ninja_build('.', False, '')\n"
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr(sys, "argv", ["prebuilt_wheels_shard.py", "1", "8"])
|
|
with pytest.raises(SystemExit):
|
|
shard.main()
|
|
assert calls[-1] == ["ninja", "-v", "-j", "1", "b.obj"]
|
|
|
|
|
|
class TestDeadline:
|
|
@pytest.mark.parametrize(
|
|
"text, seconds", [("300m", 18000), ("90s", 90), ("2h", 7200), ("45", 45)]
|
|
)
|
|
def test_durations(self, text, seconds):
|
|
assert deadline.parse_duration(text) == seconds
|
|
|
|
@pytest.mark.parametrize("text", ["", "m", "-5m", "0", "five"])
|
|
def test_bad_durations(self, text):
|
|
with pytest.raises(SystemExit):
|
|
deadline.parse_duration(text)
|
|
|
|
def test_the_exit_status_passes_through(self):
|
|
assert deadline.run(30, [sys.executable, "-c", "raise SystemExit(3)"]) == 0 + 3
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason = "POSIX process groups")
|
|
def test_the_whole_tree_ends_at_the_deadline(self, tmp_path):
|
|
marker = tmp_path / "grandchild.pid"
|
|
script = (
|
|
"import subprocess, sys, time\n"
|
|
"child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(120)'])\n"
|
|
f"open({str(marker)!r}, 'w').write(str(child.pid))\n"
|
|
"time.sleep(120)\n"
|
|
)
|
|
started = time.monotonic()
|
|
assert deadline.run(2, [sys.executable, "-c", script]) == 124
|
|
assert time.monotonic() - started < 60
|
|
grandchild = int(marker.read_text())
|
|
for _ in range(50):
|
|
try:
|
|
import os
|
|
os.kill(grandchild, 0)
|
|
except ProcessLookupError:
|
|
break
|
|
time.sleep(0.1)
|
|
else:
|
|
pytest.fail("the grandchild outlived the deadline")
|
|
|
|
def test_usage(self):
|
|
assert deadline.main(["x", "5m"]) == 2
|
|
assert deadline.main(["x", "5m", "echo"]) == 2
|
|
|
|
|
|
MSVC_STEP = "Activate the MSVC ${{ matrix.platform == 'win_arm64' && 'arm64' || 'x64' }} environment (Windows)"
|
|
# Step name -> its `if`: the two toolkits are per Windows platform, the rest per OS.
|
|
WINDOWS_SETUP = {
|
|
"Enable long paths (Windows)": "runner.os == 'Windows'",
|
|
"Install the CUDA ${{ matrix.cuda_tag }} toolkit (Windows)": "matrix.platform == 'win_amd64'",
|
|
"Install the CUDA 13.4 toolkit (Windows ARM64)": "matrix.platform == 'win_arm64'",
|
|
MSVC_STEP: "runner.os == 'Windows'",
|
|
"Install ccache (Windows)": "runner.os == 'Windows'",
|
|
"Report memory, page file and disk (Windows)": "runner.os == 'Windows'",
|
|
}
|
|
LINUX_ONLY = (
|
|
"Harden runner (audit)",
|
|
"Free up disk space",
|
|
"Install the CUDA ${{ matrix.cuda_tag }} toolkit",
|
|
"Install ccache",
|
|
"Set up swap space",
|
|
)
|
|
|
|
|
|
class TestWorkflow:
|
|
def test_platforms_input_defaults_to_both(self, workflow):
|
|
inputs = triggers(workflow)["workflow_dispatch"]["inputs"]
|
|
assert inputs["platforms"]["default"] == "linux_x86_64,win_amd64"
|
|
plan = _step(workflow, "plan", "Expand the dispatch inputs into a matrix")
|
|
assert plan["env"]["UW_PLATFORMS"] == "${{ inputs.platforms }}"
|
|
|
|
@pytest.mark.parametrize("job", ["warm", "build"])
|
|
def test_windows_setup_exists_in_both_jobs_and_is_gated(self, workflow, job):
|
|
names = [s.get("name") for s in _steps(workflow, job)]
|
|
for name, gate in WINDOWS_SETUP.items():
|
|
step = _step(workflow, job, name)
|
|
assert step["if"] == gate, (job, name)
|
|
assert step["shell"] == "pwsh", (job, name)
|
|
for name in LINUX_ONLY:
|
|
assert _step(workflow, job, name)["if"] == "runner.os == 'Linux'", (job, name)
|
|
first_checkout = next(i for i, n in enumerate(names) if n and n.startswith("Checkout"))
|
|
assert names.index("Enable long paths (Windows)") < first_checkout
|
|
assert workflow["jobs"][job]["defaults"]["run"]["shell"] == "bash"
|
|
|
|
def test_warm_and_build_run_the_same_windows_setup(self, workflow):
|
|
for name in WINDOWS_SETUP:
|
|
assert _step(workflow, "warm", name) == _step(workflow, "build", name), name
|
|
|
|
def test_the_cuda_installer_is_nvidias_and_signature_checked(self, workflow):
|
|
run = _step(workflow, "build", "Install the CUDA ${{ matrix.cuda_tag }} toolkit (Windows)")[
|
|
"run"
|
|
]
|
|
assert "https://developer.download.nvidia.com/compute/cuda/" in run
|
|
assert "Get-AuthenticodeSignature" in run and "O=NVIDIA Corporation" in run
|
|
assert "'-s'" in run and "'-n'" in run
|
|
assert "New-Item -ItemType Junction -Path 'C:\\cuda'" in run
|
|
for package in ("nvcc", "cudart", "cublas_dev", "cusparse_dev", "cusolver_dev"):
|
|
assert f"'{package}'" in run, package
|
|
assert "visual_studio_integration" not in run and "Display.Driver" not in run
|
|
|
|
def test_ccache_is_pinned_by_sha256(self, workflow):
|
|
run = _step(workflow, "build", "Install ccache (Windows)")["run"]
|
|
assert "ccache-4.14.1-windows-x86_64" in run
|
|
assert "6219f3865ca59aec41ee4b678df171d5d35855ecb2b6dbbbd20690b3a68af7b4" in run
|
|
assert "CCACHE_COMPILERCHECK=content" in run
|
|
|
|
def test_msvc_environment_is_exported_with_distutils_use_sdk(self, workflow):
|
|
step = _step(workflow, "build", MSVC_STEP)
|
|
run = step["run"]
|
|
assert "vswhere.exe" in run and "$env:VCVARS" in run
|
|
assert (
|
|
"'vcvars64.bat'" in step["env"]["VCVARS"]
|
|
and "'vcvarsarm64.bat'" in step["env"]["VCVARS"]
|
|
)
|
|
assert "DISTUTILS_USE_SDK=1" in run
|
|
assert "GITHUB_ENV" in run and "GITHUB_PATH" in run
|
|
|
|
def test_artifacts_and_caches_are_per_platform(self, workflow):
|
|
body = json.dumps(workflow["jobs"])
|
|
for prefix in ("ccache-", "wheel-", "prebuilt-ccache-"):
|
|
for value in (
|
|
v
|
|
for v in body.split('"')
|
|
if v.startswith(prefix) and "${{ matrix" in v and "*" not in v
|
|
):
|
|
assert value.startswith(prefix + "${{ matrix.platform }}"), value
|
|
|
|
def test_the_windows_build_ends_the_whole_tree_at_the_deadline(self, workflow):
|
|
for job, name in (
|
|
("warm", "Compile slice ${{ matrix.shard }} of ${{ matrix.shards }}"),
|
|
("build", "Build the wheel"),
|
|
):
|
|
run = _step(workflow, job, name)["run"]
|
|
assert "prebuilt_wheels_timeout.py" in run
|
|
assert 'timeout --signal=INT --kill-after=120 "$BUILD_TIMEOUT"' in run
|
|
|
|
def test_the_rename_expects_the_cells_platform(self, workflow):
|
|
step = _step(workflow, "build", "Name the wheel the way upstream names it")
|
|
assert step["env"]["WHEEL_PLATFORM"] == "${{ matrix.platform }}"
|
|
assert (
|
|
"PLATFORM" not in step["env"]
|
|
) # vcvars exports Platform=x64; Windows env is case-insensitive
|
|
assert "-${WHEEL_PLATFORM}.whl" in step["run"]
|
|
|
|
def test_the_torch_abi_check_follows_the_cell(self, workflow):
|
|
for job in ("warm", "build"):
|
|
step = next(
|
|
s
|
|
for s in _steps(workflow, job)
|
|
if s.get("name", "") == "Install ${{ matrix.torch_spec }}"
|
|
)
|
|
assert step["env"]["MATRIX_ABI"] == "${{ matrix.abi }}"
|
|
assert 'os.environ["MATRIX_ABI"]' in step["run"]
|
|
|
|
def test_the_windows_smoke_hides_the_toolkit(self, workflow):
|
|
"""A user has no CUDA toolkit; the import must resolve against torch's own DLLs."""
|
|
smoke = _step(workflow, "build", "Smoke test the wheel in a fresh venv")
|
|
run = smoke["run"]
|
|
assert "Scripts/python.exe" in run
|
|
assert "unset CUDA_HOME CUDA_PATH" in run
|
|
assert smoke["env"]["SMOKE_EXCLUDE"] == "${{ matrix.smoke_exclude }}"
|
|
assert smoke["env"]["SMOKE_EXTRA"] == "${{ matrix.smoke_extra }}"
|
|
assert 'os.environ.get("SMOKE_EXCLUDE"' in run and 'os.environ.get("SMOKE_EXTRA"' in run
|
|
|
|
def test_the_gpu_job_only_gets_linux_cells(self, workflow):
|
|
job = workflow["jobs"]["gpu-smoke"]
|
|
assert job["strategy"]["matrix"] == "${{ fromJSON(needs.plan.outputs.gpu_matrix) }}"
|
|
assert "needs.plan.outputs.gpu_count != '0'" in job["if"]
|
|
assert workflow["jobs"]["plan"]["outputs"]["gpu_count"]
|
|
|
|
def test_publish_rebuilds_sha256sums_over_the_whole_release(self, workflow):
|
|
steps = workflow["jobs"]["publish"]["steps"]
|
|
create = next(s for s in steps if s.get("name") == "Create or update the release")
|
|
# This run's SHA256SUMS lists this run only; it must not become the release's.
|
|
assert "wheels/SHA256SUMS" not in create["run"]
|
|
regen = next(s for s in steps if "SHA256SUMS" in s.get("name", ""))
|
|
run = regen["run"]
|
|
assert "prebuilt_wheels.py sums" in run
|
|
assert "--clobber SHA256SUMS" in run
|
|
assert "releases/assets/$id" in run
|
|
assert "sha256:unknown" not in run
|
|
assert run.index("prebuilt_wheels.py sums") < run.index("prebuilt_wheels.py notes")
|
|
|
|
|
|
# The text the patch keys on, from the pinned upstream files.
|
|
_MAMBA_KERNEL = """#include <c10/cuda/CUDAException.h>
|
|
|
|
#ifndef USE_ROCM
|
|
#include <cub/block/block_load.cuh>
|
|
#else
|
|
#include <hipcub/hipcub.hpp>
|
|
#endif
|
|
|
|
constexpr float kLog2e = M_LOG2E;
|
|
|
|
void launch() {
|
|
BOOL_SWITCH(x, kX, [&] {
|
|
BOOL_SWITCH(y, kY, [&] {
|
|
if (kSmemSize >= 48 * 1024) {
|
|
#ifndef USE_ROCM
|
|
C10_CUDA_CHECK(cudaFuncSetAttribute(
|
|
kernel, cuda_branch));
|
|
#else
|
|
C10_CUDA_CHECK(cudaFuncSetAttribute(
|
|
(void *) kernel, rocm_branch));
|
|
#endif
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
void dispatch() {
|
|
#ifndef USE_ROCM
|
|
cuda_dispatch();
|
|
#else
|
|
rocm_dispatch();
|
|
#endif
|
|
}
|
|
"""
|
|
|
|
|
|
class TestMambaMsvcPatch:
|
|
def _tree(
|
|
self,
|
|
tmp_path,
|
|
text = _MAMBA_KERNEL,
|
|
):
|
|
kernels = tmp_path / "csrc" / "selective_scan"
|
|
kernels.mkdir(parents = True)
|
|
for name in ("selective_scan_fwd_kernel.cuh", "selective_scan_bwd_kernel.cuh"):
|
|
(kernels / name).write_text(text, encoding = "utf-8")
|
|
return kernels
|
|
|
|
def _run(self, root):
|
|
return subprocess.run(
|
|
[sys.executable, str(SCRIPTS / "patch_mamba_msvc.py"), str(root)],
|
|
capture_output = True,
|
|
text = True,
|
|
)
|
|
|
|
def test_keeps_the_cuda_branch_and_is_idempotent(self, tmp_path):
|
|
kernels = self._tree(tmp_path)
|
|
assert self._run(tmp_path).returncode == 0
|
|
first = (kernels / "selective_scan_fwd_kernel.cuh").read_text(encoding = "utf-8")
|
|
assert "cuda_branch" in first and "rocm_branch" not in first
|
|
# File-scope USE_ROCM blocks are legal under MSVC and stay.
|
|
assert "hipcub" in first
|
|
# Function-scope USE_ROCM blocks are legal under MSVC and stay.
|
|
assert "rocm_dispatch()" in first
|
|
again = self._run(tmp_path)
|
|
assert again.returncode == 0 and "already patched" in again.stdout
|
|
assert (kernels / "selective_scan_fwd_kernel.cuh").read_text(encoding = "utf-8") == first
|
|
|
|
def test_a_second_in_lambda_block_fails_loudly(self, tmp_path):
|
|
body = _MAMBA_KERNEL.split("void launch()")[1].split("void dispatch()")[0]
|
|
self._tree(tmp_path, _MAMBA_KERNEL + "void launch2()" + body)
|
|
result = self._run(tmp_path)
|
|
assert result.returncode == 1
|
|
assert "found 2" in result.stderr
|
|
|
|
def test_runs_on_windows_mamba_cells_only(self, workflow):
|
|
for job in ("warm", "build"):
|
|
step = next(
|
|
s
|
|
for s in workflow["jobs"][job]["steps"]
|
|
if "patch_mamba_msvc.py" in s.get("run", "")
|
|
)
|
|
assert step["if"] == "runner.os == 'Windows' && matrix.package == 'mamba-ssm'"
|