From c697ffc382cdc009cc22740b6a64cf31810649d8 Mon Sep 17 00:00:00 2001 From: HaileyStorm Date: Fri, 28 Aug 2026 19:17:43 -0600 Subject: [PATCH] fix(build): honor explicit and versioned CUDA toolkits --- CONTRIBUTING.md | 6 +- docs/install.md | 7 +- freetoken-kernel-cache/build_backend.py | 6 +- python/freetoken/kernel/_toolchain.py | 158 ++++++++++++++--- python/freetoken/server/launch.py | 10 +- setup.py | 13 +- tests/kernels/test_toolchain.py | 219 ++++++++++++++++++++++++ 7 files changed, 385 insertions(+), 34 deletions(-) create mode 100644 tests/kernels/test_toolchain.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e286808d7..447d3976e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,7 +46,11 @@ git clone https://github.com/FlashML-org/FreeToken.git && cd FreeToken uv pip install -e ".[accel]" ``` -See [docs/install.md](docs/install.md) for requirements. CUDA kernels are JIT-compiled on first use and need a CUDA 13 toolkit with `nvcc` on `PATH`. +See [docs/install.md](docs/install.md) for requirements. CUDA kernels are JIT-compiled +on first use. An explicit absolute `CUDA_HOME` must contain an `nvcc` matching +PyTorch's exact CUDA release; otherwise FreeToken checks the matching versioned +`/usr/local/cuda-X.Y` toolkit and then `PATH`. This selection does not change the +supported PyTorch, driver, or GPU matrix. Run the tests with `pytest`. Most tests need an NVIDIA GPU. `-m "not slow"` skips the long kernel sweeps; tests marked `needs_weights` are off unless you point them at a local checkpoint (see [tests/README.md](tests/README.md)). diff --git a/docs/install.md b/docs/install.md index f5205ab3b..4038a03c1 100644 --- a/docs/install.md +++ b/docs/install.md @@ -13,7 +13,12 @@ uv venv && source .venv/bin/activate uv pip install "freetoken[accel]" ``` -CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH. +CUDA kernels are JIT-compiled on first use. An explicit absolute `CUDA_HOME` +must contain an `nvcc` matching PyTorch's exact CUDA release; otherwise +FreeToken checks the matching `/usr/local/cuda-X.Y` toolkit and then `PATH`. +JIT and kernel-cache builds fail if no exact compiler is available. Toolkit +selection does not change the supported PyTorch, driver, or GPU matrix. A +complete prebuilt kernel cache does not require a compiler at server startup. ## Method 2: Install from source diff --git a/freetoken-kernel-cache/build_backend.py b/freetoken-kernel-cache/build_backend.py index 03f17a314..10f9bf216 100644 --- a/freetoken-kernel-cache/build_backend.py +++ b/freetoken-kernel-cache/build_backend.py @@ -42,9 +42,8 @@ def _cuda_version_suffix() -> str: cuda_version = getattr(torch.version, "cuda", None) if not cuda_version: return "" - # The tag advertises torch's CUDA; the cache .so link nvcc's libcudart. - # Only a matching major makes both statements true at once. - _check_toolchain() + # The tag advertises torch's CUDA. The wheel-build hook validates nvcc before + # compiling; metadata and sdist hooks deliberately remain compiler-free. return f"+cu{cuda_version.replace('.', '')}" @@ -74,6 +73,7 @@ def _selected_specs() -> list[str] | None: def _build_jit_cache() -> None: + _check_toolchain() _ensure_freetoken_importable() from freetoken.kernel.aot import compile_and_package_kernels diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb4..ca5eee996 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -6,21 +6,129 @@ from __future__ import annotations -import functools import os import re import shutil import subprocess +import sys +from pathlib import Path ALLOW_MISMATCH_ENV = "FREETOKEN_ALLOW_CUDA_MISMATCH" _TRUE_VALUES = {"1", "true", "yes", "on"} -def _nvcc_path() -> str | None: - from torch.utils.cpp_extension import CUDA_HOME +def _mismatch_allowed() -> bool: + return os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES + + +def _torch_cuda_release() -> tuple[int, int] | None: + try: + import torch + except ModuleNotFoundError: + return None + + cuda = getattr(torch.version, "cuda", None) + if not cuda: + return None + match = re.match(r"^(\d+)\.(\d+)", str(cuda)) + if match is None: + raise RuntimeError(f"torch reports an invalid CUDA release: {cuda!r}") + return int(match.group(1)), int(match.group(2)) + + +def _prepend_path(directory: str) -> None: + entries = [item for item in os.getenv("PATH", "").split(os.pathsep) if item] + entries = [item for item in entries if os.path.abspath(item) != directory] + os.environ["PATH"] = os.pathsep.join([directory, *entries]) + + +def _publish_cuda_home(cuda_home: Path) -> str: + home = str(cuda_home) + os.environ["CUDA_HOME"] = home + _prepend_path(str(cuda_home / "bin")) + + # Normally configure_cuda_toolchain() runs before cpp_extension is imported. + # Keep direct library/JIT callers correct too if another dependency imported it first. + cpp_extension = sys.modules.get("torch.utils.cpp_extension") + if cpp_extension is not None: + cpp_extension.CUDA_HOME = home + return home + + +def _release_error( + nvcc: Path, + expected: tuple[int, int], + actual: tuple[int, int] | None, +) -> str: + wanted = f"{expected[0]}.{expected[1]}" + if actual is None: + return f"{nvcc} does not report a valid CUDA release; torch requires CUDA {wanted}" + return ( + f"nvcc {actual[0]}.{actual[1]} at {nvcc} does not match " + f"torch CUDA {wanted}" + ) + - if CUDA_HOME: - return os.path.join(CUDA_HOME, "bin", "nvcc") +def configure_cuda_toolchain(*, reject_path_mismatch: bool = True) -> str | None: + """Select torch's exact CUDA toolkit in this process before any JIT/build. + + An explicit absolute ``CUDA_HOME`` is authoritative. Otherwise discovery is + deliberately bounded to the exact versioned ``/usr/local`` toolkit and the + current ``PATH``. ``reject_path_mismatch=False`` lets a prebuilt-only server + ignore an unrelated compiler on ``PATH``; actual JIT/build checks stay strict. + We never rewrite system alternatives or shell state. + """ + expected = _torch_cuda_release() + if expected is None: + return None + + explicit = os.getenv("CUDA_HOME", "").strip() + if explicit: + cuda_home = Path(explicit) + if not cuda_home.is_absolute(): + raise RuntimeError(f"CUDA_HOME must be absolute, got {explicit!r}") + nvcc = cuda_home / "bin" / "nvcc" + actual = nvcc_release(str(nvcc)) + if actual is None: + raise RuntimeError(_release_error(nvcc, expected, actual)) + if actual != expected and not _mismatch_allowed(): + raise RuntimeError( + f"{_release_error(nvcc, expected, actual)}; set " + f"{ALLOW_MISMATCH_ENV}=1 to override the explicit CUDA_HOME check" + ) + return _publish_cuda_home(cuda_home) + + exact_home = Path(f"/usr/local/cuda-{expected[0]}.{expected[1]}") + exact_nvcc = exact_home / "bin" / "nvcc" + if nvcc_release(str(exact_nvcc)) == expected: + return _publish_cuda_home(exact_home) + + path_nvcc_raw = shutil.which("nvcc") + if path_nvcc_raw is None: + return None + # Derive the toolkit root from the PATH entry, not its resolved implementation. + # Distro packages commonly expose /usr/bin/nvcc as a symlink into /usr/lib while + # keeping their public CUDA headers and libraries rooted at /usr. + path_nvcc = Path(os.path.abspath(path_nvcc_raw)) + actual = nvcc_release(str(path_nvcc)) + if actual is None: + if not reject_path_mismatch: + return None + raise RuntimeError(_release_error(path_nvcc, expected, actual)) + if actual != expected and not _mismatch_allowed(): + if not reject_path_mismatch: + return None + raise RuntimeError( + f"{_release_error(path_nvcc, expected, actual)}; install " + f"/usr/local/cuda-{expected[0]}.{expected[1]} or set an absolute CUDA_HOME" + ) + return _publish_cuda_home(path_nvcc.parent.parent) + + +def _nvcc_path() -> str | None: + cuda_home = os.getenv("CUDA_HOME", "").strip() + if cuda_home: + return os.path.join(cuda_home, "bin", "nvcc") return shutil.which("nvcc") @@ -36,36 +144,38 @@ def nvcc_release(nvcc: str) -> tuple[int, int] | None: def torch_cuda_major() -> int | None: - import torch - - cuda = getattr(torch.version, "cuda", None) - return int(cuda.split(".")[0]) if cuda else None + release = _torch_cuda_release() + return release[0] if release is not None else None -@functools.cache def check_nvcc_matches_torch() -> None: - """Refuse to nvcc-compile kernels across CUDA majors. + """Refuse to nvcc-compile kernels across CUDA release families. - nvcc-built binaries link libcudart.so.; at runtime only the - torch wheel's own CUDA runtime is guaranteed to be loadable. + configure_cuda_toolchain() also publishes the selected toolkit before the + downstream builder imports or spawns its compiler machinery. """ - if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: - return - torch_major = torch_cuda_major() - if torch_major is None: + configure_cuda_toolchain() + torch_release = _torch_cuda_release() + if torch_release is None: return nvcc = _nvcc_path() if nvcc is None: - return + raise RuntimeError( + f"no nvcc matching torch CUDA {torch_release[0]}.{torch_release[1]} was found; " + "install the exact versioned toolkit or set an absolute CUDA_HOME" + ) release = nvcc_release(nvcc) if release is None: - return - if release[0] != torch_major: + raise RuntimeError( + f"{nvcc} does not report a valid CUDA release for torch CUDA " + f"{torch_release[0]}.{torch_release[1]}" + ) + if release != torch_release and not _mismatch_allowed(): import torch raise RuntimeError( - f"nvcc {release[0]}.{release[1]} would build kernels linking " - f"libcudart.so.{release[0]}, but torch {torch.__version__} ships CUDA " - f"{torch.version.cuda} (libcudart.so.{torch_major}). Install a CUDA " - f"{torch_major}.x toolkit, or set {ALLOW_MISMATCH_ENV}=1 to override." + f"nvcc {release[0]}.{release[1]} would build kernels for a different " + f"CUDA release family, but torch {torch.__version__} ships CUDA " + f"{torch.version.cuda}. Set an exact absolute CUDA_HOME, or set " + f"{ALLOW_MISMATCH_ENV}=1 to override." ) diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index acef55175..310c38375 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -127,7 +127,6 @@ def launch_server( argv: list[str] | None = None, prog: str | None = None, ) -> None: - from .api_server import run_api_server from .args import parse_args server_args, run_shell = parse_args( @@ -135,6 +134,15 @@ def launch_server( run_shell, prog=prog, ) + + # Publish an exact versioned toolkit before multiprocessing spawn. A mismatched + # PATH compiler is ignored here so a complete prebuilt cache can still serve; + # actual JIT/build entry points perform the strict check. + from freetoken.kernel._toolchain import configure_cuda_toolchain + + configure_cuda_toolchain(reject_path_mismatch=False) + + from .api_server import run_api_server logger = init_logger(__name__, "initializer") if server_args.gpu: diff --git a/setup.py b/setup.py index cfe41b7d8..b305a0385 100644 --- a/setup.py +++ b/setup.py @@ -4,18 +4,24 @@ from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension ROOT = Path(__file__).parent -def _check_toolchain() -> None: +def _load_toolchain(): path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - module.check_nvcc_matches_torch() + return module + + +_TOOLCHAIN = _load_toolchain() +_TOOLCHAIN.check_nvcc_matches_torch() + +# Import only after CUDA_HOME/PATH are exact, because cpp_extension caches CUDA_HOME. +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension # noqa: E402 def _cuda_runtime_paths() -> tuple[list[str], list[str]]: @@ -32,7 +38,6 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() -_check_toolchain() setup( diff --git a/tests/kernels/test_toolchain.py b/tests/kernels/test_toolchain.py new file mode 100644 index 000000000..fd7ea6d5f --- /dev/null +++ b/tests/kernels/test_toolchain.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +import importlib.util +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from freetoken.kernel import _toolchain + + +def _load_kernel_cache_backend(): + path = Path(__file__).resolve().parents[2] / "freetoken-kernel-cache" / "build_backend.py" + spec = importlib.util.spec_from_file_location("_freetoken_kernel_cache_backend", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _isolate_environment( + monkeypatch: pytest.MonkeyPatch, + cuda: str | None = "13.0", +) -> None: + monkeypatch.delenv("CUDA_HOME", raising=False) + monkeypatch.delenv(_toolchain.ALLOW_MISMATCH_ENV, raising=False) + monkeypatch.setenv("PATH", "/usr/bin:/bin") + monkeypatch.setattr( + _toolchain, + "_torch_cuda_release", + lambda: None if cuda is None else tuple(map(int, cuda.split("."))), + ) + monkeypatch.setattr(_toolchain.shutil, "which", lambda _name: None) + + +def test_cpu_or_missing_torch_is_a_noop(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch, cuda=None) + + assert _toolchain.configure_cuda_toolchain() is None + assert "CUDA_HOME" not in os.environ + assert os.environ["PATH"] == "/usr/bin:/bin" + + +def test_jit_check_requires_a_valid_exact_compiler(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setattr(_toolchain, "nvcc_release", lambda _path: None) + + with pytest.raises(RuntimeError, match=r"no nvcc matching torch CUDA 13\.0"): + _toolchain.check_nvcc_matches_torch() + + +def test_discovers_exact_versioned_usr_local_toolkit(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setattr( + _toolchain, + "nvcc_release", + lambda path: (13, 0) if path == "/usr/local/cuda-13.0/bin/nvcc" else None, + ) + + assert _toolchain.configure_cuda_toolchain() == "/usr/local/cuda-13.0" + assert os.environ["CUDA_HOME"] == "/usr/local/cuda-13.0" + assert os.environ["PATH"].split(os.pathsep)[0] == "/usr/local/cuda-13.0/bin" + + +def test_explicit_exact_cuda_home_remains_authoritative(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch, cuda="12.8") + monkeypatch.setenv("CUDA_HOME", "/usr/local/cuda-12.8") + monkeypatch.setattr(_toolchain, "nvcc_release", lambda _path: (12, 8)) + + assert _toolchain.configure_cuda_toolchain() == "/usr/local/cuda-12.8" + assert os.environ["PATH"].split(os.pathsep)[0] == "/usr/local/cuda-12.8/bin" + + +def test_explicit_mismatch_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setenv("CUDA_HOME", "/usr/local/cuda-12.8") + monkeypatch.setattr(_toolchain, "nvcc_release", lambda _path: (12, 8)) + + with pytest.raises(RuntimeError, match=r"nvcc 12\.8.*torch CUDA 13\.0"): + _toolchain.configure_cuda_toolchain() + assert os.environ["CUDA_HOME"] == "/usr/local/cuda-12.8" + + +def test_explicit_mismatch_override_preserves_selected_home(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setenv("CUDA_HOME", "/usr/local/cuda-12.8") + monkeypatch.setenv(_toolchain.ALLOW_MISMATCH_ENV, "1") + monkeypatch.setattr(_toolchain, "nvcc_release", lambda _path: (12, 8)) + + assert _toolchain.configure_cuda_toolchain() == "/usr/local/cuda-12.8" + + +def test_auto_discovery_rejects_wrong_path_nvcc(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setattr(_toolchain.shutil, "which", lambda _name: "/usr/bin/nvcc") + monkeypatch.setattr( + _toolchain, + "nvcc_release", + lambda path: None if path == "/usr/local/cuda-13.0/bin/nvcc" else (12, 0), + ) + + with pytest.raises(RuntimeError, match=r"nvcc 12\.0.*torch CUDA 13\.0"): + _toolchain.configure_cuda_toolchain() + assert "CUDA_HOME" not in os.environ + + +def test_prebuilt_server_ignores_mismatched_path_compiler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from freetoken.server.launch import launch_server + + _isolate_environment(monkeypatch) + monkeypatch.setattr(_toolchain.shutil, "which", lambda _name: "/usr/bin/nvcc") + monkeypatch.setattr( + _toolchain, + "nvcc_release", + lambda path: None if path == "/usr/local/cuda-13.0/bin/nvcc" else (12, 8), + ) + + server_args = SimpleNamespace(gpu=()) + observed: list[tuple[object, bool]] = [] + fake_args_module = SimpleNamespace( + parse_args=lambda _argv, run_shell, prog=None: (server_args, run_shell) + ) + fake_api_module = SimpleNamespace( + run_api_server=lambda args, _start, run_shell: observed.append((args, run_shell)) + ) + monkeypatch.setitem(sys.modules, "freetoken.server.args", fake_args_module) + monkeypatch.setitem(sys.modules, "freetoken.server.api_server", fake_api_module) + + launch_server(argv=[], prog="ft serve") + + assert observed == [(server_args, False)] + assert "CUDA_HOME" not in os.environ + with pytest.raises(RuntimeError, match=r"nvcc 12\.8.*torch CUDA 13\.0"): + _toolchain.check_nvcc_matches_torch() + + +def test_updates_already_imported_cpp_extension_cache(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setattr(_toolchain, "nvcc_release", lambda _path: (13, 0)) + + class FakeCppExtension: + CUDA_HOME = None + + fake = FakeCppExtension() + monkeypatch.setitem(sys.modules, "torch.utils.cpp_extension", fake) + + _toolchain.configure_cuda_toolchain() + assert fake.CUDA_HOME == "/usr/local/cuda-13.0" + + +def test_path_exact_toolkit_is_published_from_its_own_nvcc(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + path_nvcc = Path("/opt/cuda-13.0/bin/nvcc") + monkeypatch.setattr(_toolchain.shutil, "which", lambda _name: str(path_nvcc)) + monkeypatch.setattr( + _toolchain, + "nvcc_release", + lambda path: None if path == "/usr/local/cuda-13.0/bin/nvcc" else (13, 0), + ) + + assert _toolchain.configure_cuda_toolchain() == "/opt/cuda-13.0" + + +def test_path_symlink_uses_public_entry_toolkit_root(monkeypatch: pytest.MonkeyPatch) -> None: + _isolate_environment(monkeypatch) + monkeypatch.setattr(_toolchain.shutil, "which", lambda _name: "/usr/bin/nvcc") + monkeypatch.setattr( + _toolchain, + "nvcc_release", + lambda path: None if path == "/usr/local/cuda-13.0/bin/nvcc" else (13, 0), + ) + + assert _toolchain.configure_cuda_toolchain() == "/usr" + assert os.environ["CUDA_HOME"] == "/usr" + + +def test_kernel_cache_metadata_does_not_require_nvcc(monkeypatch: pytest.MonkeyPatch) -> None: + backend = _load_kernel_cache_backend() + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(version=SimpleNamespace(cuda="12.8"))) + + def fail_if_called() -> None: + raise AssertionError("metadata generation must not require a compiler") + + monkeypatch.setattr(backend, "_check_toolchain", fail_if_called) + assert backend._cuda_version_suffix() == "+cu128" + + +def test_kernel_cache_build_checks_toolchain_before_compilation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _load_kernel_cache_backend() + + def reject_toolchain() -> None: + raise RuntimeError("toolchain rejected") + + monkeypatch.setattr(backend, "_check_toolchain", reject_toolchain) + with pytest.raises(RuntimeError, match="toolchain rejected"): + backend._build_jit_cache() + + +def test_server_help_does_not_require_toolchain_resolution( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + from freetoken.server.launch import launch_server + + def fail_if_called() -> None: + raise AssertionError("toolchain resolution must follow argument parsing") + + monkeypatch.setattr(_toolchain, "configure_cuda_toolchain", fail_if_called) + with pytest.raises(SystemExit) as exc_info: + launch_server(argv=["--help"], prog="ft serve") + + assert exc_info.value.code == 0 + assert "usage: ft serve" in capsys.readouterr().out