diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index d69eccc54d..e9941d9e62 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -209,3 +209,154 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) assert seen["repo"] == FORK assert out["repo"] == FORK + + +# Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer +# sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. + + +def _gpu_linux_host(caps): + return _host( + is_linux = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + driver_cuda_version = (13, 1), + compute_caps = caps, + ) + + +def test_host_is_blackwell_includes_datacenter_parts(): + assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100 + assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103 + assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120 + assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121 + assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper + assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere + assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins + + +def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile): + return ilp.PublishedLlamaArtifact( + asset_name = f"app-b9739-linux-x64-{profile}.tar.gz", + install_kind = "linux-cuda", + runtime_line = runtime_line, + coverage_class = "newer", + supported_sms = supported_sms, + min_sm = min_sm, + max_sm = max_sm, + bundle_profile = profile, + rank = 50, + ) + + +def test_linux_blackwell_override_prefers_cuda13_for_datacenter(monkeypatch): + # Both bundles cover sm_100 and torch reports cuda12, so coverage alone can't + # decide -- only the sm_100 Blackwell floor lifts cuda13 to the front. + cuda12 = _linux_cuda_artifact( + "cuda12", ["86", "89", "90", "100", "120"], 86, 120, "cuda12-newer" + ) + cuda13 = _linux_cuda_artifact( + "cuda13", ["86", "89", "90", "100", "103", "120"], 86, 120, "cuda13-newer" + ) + release = ilp.PublishedReleaseBundle( + repo = FORK, + release_tag = "b9739-mix", + upstream_tag = "b9739", + assets = {cuda12.asset_name: "https://x/cuda12", cuda13.asset_name: "https://x/cuda13"}, + artifacts = [cuda12, cuda13], + ) + monkeypatch.setattr( + ilp, + "detected_linux_runtime_lines", + lambda: (["cuda13", "cuda12"], {"cuda13": ["/usr/lib"], "cuda12": ["/usr/lib"]}), + ) + + selection = ilp.linux_cuda_choice_from_release( + _gpu_linux_host(["10.0"]), release, preferred_runtime_line = "cuda12" + ) + assert selection is not None + assert selection.primary.runtime_line == "cuda13" + assert selection.primary.bundle_profile == "cuda13-newer" + + +def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter(): + # B200 (sm_100) on Windows must drop the cuda-12.4 build and keep cuda13. + host = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.0"], + ) + cuda124 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.4-x64.zip", + url = "https://x/124", + source_label = "published", + install_kind = "windows-cuda", + ) + cuda13 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "app-b9739-windows-x64-cuda13-newer.zip", + url = "https://x/13", + source_label = "published", + install_kind = "windows-cuda", + max_sm = 120, + ) + kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda124, cuda13]) + assert [a.name for a in kept] == [cuda13.name] + + +def test_blackwell_min_toolkit_is_sm_aware(): + # Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it. + f = ilp._blackwell_min_toolkit_for_host + assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200 + assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50 + assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300 + assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark + assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins + + +def test_sm103_host_drops_cuda128_windows_build(): + # B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped. + host = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.3"], + ) + cuda128 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.8-x64.zip", + url = "https://x/128", + source_label = "published", + install_kind = "windows-cuda", + ) + cuda129 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.9-x64.zip", + url = "https://x/129", + source_label = "published", + install_kind = "windows-cuda", + ) + kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda128, cuda129]) + assert [a.name for a in kept] == [cuda129.name] + # sm_100 stays on the 12.8 family floor and keeps the same 12.8 build. + b200 = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.0"], + ) + kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) + assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 1823a0f67e..33cb709ba7 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -230,12 +230,15 @@ FORCE_COMPILE_DEFAULT_REF = os.environ.get("UNSLOTH_LLAMA_FORCE_COMPILE_REF", "m _MIN_CUDA_MAJOR = 12 _MAX_PROBE_CUDA_MAJOR = 19 -# Blackwell sm_120 capability thresholds. A host is Blackwell when its highest -# compute capability is at least sm_120; ggml compiles sm_120 only at toolkit -# >= 12.8, so an in-release windows-cuda build at or above that already covers -# Blackwell, while cuda-12.4 does not and is dropped on a Blackwell host. -_BLACKWELL_MIN_SM = 120 +# Blackwell floor is sm_100: data-center parts (B100/B200 sm_100, B300/GB300 +# sm_103) sit below consumer Blackwell (RTX 50 sm_120); the family needs toolkit +# >= 12.8, except sm_103/sm_121 which need 12.9. (120 here wrongly excluded the +# sm_100/103 data-center hosts.) +_BLACKWELL_MIN_SM = 100 _BLACKWELL_MIN_TOOLKIT = (12, 8) +# SMs that need a newer toolkit than the family floor (CUDA 12.9 added native +# sm_103/sm_121 targets; 12.8 covers sm_100/101/120). +_BLACKWELL_SM_MIN_TOOLKIT = {103: (12, 9), 121: (12, 9)} def _cuda_runtime_lines_for_major(major: int) -> list[str]: @@ -3418,18 +3421,18 @@ def windows_cuda_attempts( return attempts -def _windows_cuda_attempt_covers_blackwell(attempt: AssetChoice) -> bool: - """True if an in-release windows-cuda attempt yields a Blackwell sm_120 - capable build. The fork's app-named bundles declare their SM coverage - directly; legacy upstream-named bundles instead encode their CUDA toolkit - minor in the filename (covers Blackwell at toolkit >= 12.8).""" +def _windows_cuda_attempt_covers_blackwell( + attempt: AssetChoice, min_toolkit: tuple[int, int] = _BLACKWELL_MIN_TOOLKIT +) -> bool: + """True if a windows-cuda attempt is Blackwell-capable (app bundles via + declared SMs; legacy upstream bundles via toolkit minor >= min_toolkit: + 12.8 for the family, 12.9 for sm_103/sm_121).""" if attempt.install_kind != "windows-cuda": return False - # Legacy upstream-named bundles encode their toolkit minor; it is the binding - # constraint (a 12.4 toolkit cannot offload sm_120 whatever its metadata says). + # Legacy bundle: the toolkit minor binds (12.4 cannot offload Blackwell). m = re.search(r"-bin-win-cuda-(\d+)\.(\d+)-x64\.zip$", attempt.name) if m is not None: - return (int(m.group(1)), int(m.group(2))) >= _BLACKWELL_MIN_TOOLKIT + return (int(m.group(1)), int(m.group(2))) >= min_toolkit # App-named bundles carry no minor and declare their SM coverage directly. return attempt.max_sm is not None and attempt.max_sm >= _BLACKWELL_MIN_SM @@ -3439,22 +3442,30 @@ def _host_is_blackwell(host: HostInfo) -> bool: return bool(caps) and int(caps[-1]) >= _BLACKWELL_MIN_SM +def _blackwell_min_toolkit_for_host(host: HostInfo) -> tuple[int, int]: + """Minimum CUDA toolkit this Blackwell host needs: 12.8 for the family, + 12.9 if any of its SMs is sm_103/sm_121 (no native target before 12.9).""" + req = _BLACKWELL_MIN_TOOLKIT + for sm in normalize_compute_caps(host.compute_caps): + req = max(req, _BLACKWELL_SM_MIN_TOOLKIT.get(int(sm), _BLACKWELL_MIN_TOOLKIT)) + return req + + def _drop_blackwell_incapable_windows_cuda( host: HostInfo, attempts: list[AssetChoice] ) -> list[AssetChoice]: - """On a Blackwell host, drop windows-cuda attempts that cannot offload - sm_120 (e.g. upstream cuda-12.4, toolkit 12.4). Such a build loads and - passes the functional validator but runs the model on a slow non-native - path (an RTX 5090 measured 7.1 tok/s vs 551.2 on cuda-13.3), so it must - not sit in the fallback chain behind the pin or an in-release cuda13. - Non-cuda attempts (windows-cpu, windows-hip, ...) pass through so the - host still degrades to an honest CPU install when no CUDA 13 exists.""" + """On a Blackwell host, drop windows-cuda attempts that can't offload + Blackwell (e.g. cuda-12.4): they load and validate but run a slow non-native + path (RTX 5090: 7.1 vs 551.2 tok/s on cuda-13.3). Non-cuda attempts pass + through so the host can still fall back to an honest CPU install.""" if not _host_is_blackwell(host): return attempts + min_toolkit = _blackwell_min_toolkit_for_host(host) return [ attempt for attempt in attempts - if attempt.install_kind != "windows-cuda" or _windows_cuda_attempt_covers_blackwell(attempt) + if attempt.install_kind != "windows-cuda" + or _windows_cuda_attempt_covers_blackwell(attempt, min_toolkit) ]