Do not auto-fall back to Vulkan when a HIP device mask filtered the probe

hipinfo is itself a HIP application, and AMD documents HIP_VISIBLE_DEVICES as
"only devices whose index is present in the sequence are visible to HIP", with
that spelling recommended on Windows. Under a mask the Windows probe therefore
enumerates the visible devices, so rocm_gfx_targets is what survived the mask
rather than the physical inventory the auto-Vulkan floor check assumes. A
masked-out gfx1100 next to a visible gfx803 made the check conclude that no AMD
GPU on the box reaches the Windows HIP prebuilt and route the install to Vulkan,
which honours none of these masks and would enumerate the reserved card.

Decline to guess when a mask is set: the physical inventory is unknowable from a
masked probe, so keep the HIP / fork / source path. This only ever turns the
automatic fallback off, never on. The driver-only single-GPU host the fallback
exists for sets no mask, an all-hiding "" / -1 mask is still handled as no active
target rather than a partial view, and an explicit --llama-backend vulkan or
UNSLOTH_LLAMA_BACKEND=vulkan is unaffected.

Reading the physical inventory through an unmasked re-probe would also correct
_pick_rocm_gfx_target, which indexes the token list by the mask value and so
already assumes an unmasked probe. That is pre-existing behaviour on main and is
left alone here.

Author: danielhanchen

studio/backend/tests/test_install_resolve_prebuilt.py

@@ -32,6 +32,19 @@
 UPSTREAM = ilp.UPSTREAM_REPO  # ggml-org/llama.cpp
 
 
+@pytest.fixture(autouse = True)
+def _no_ambient_hip_device_mask(monkeypatch):
+    """These tests describe hosts through HostInfo, not through the environment.
+
+    A HIP visible-device mask inherited from the shell (an ML box commonly exports
+    CUDA_VISIBLE_DEVICES) means the arch probe saw only part of the GPUs, which the
+    Windows auto-Vulkan guard treats as an unknown physical inventory. Clear all
+    three so a host is described by its fields alone; the tests that are about the
+    mask set it explicitly."""
+    for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+        monkeypatch.delenv(_env, raising = False)
+
+
 def _host(**kw):
     base = dict(
         system = "Linux",
@@ -875,6 +888,61 @@ def test_route_to_vulkan_prebuilt_auto_fallback_when_no_amd_gpu_reaches_floor():
     assert routed.has_rocm is False
 
 
+@pytest.mark.parametrize(
+    "mask_env", ["HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"]
+)
+def test_auto_vulkan_declines_when_a_hip_device_mask_filtered_the_probe(mask_env, monkeypatch):
+    # hipinfo is a HIP application, so under a visible-device mask it enumerates only
+    # the devices that survived it: rocm_gfx_targets is then the VISIBLE set and a
+    # HIP-capable card can be masked out of view entirely. "No AMD GPU on this box
+    # reaches the floor" is unprovable there, and Vulkan honours none of these masks,
+    # so the automatic fallback must decline rather than hand it the reserved card.
+    monkeypatch.setenv(mask_env, "1")
+    host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
+    assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
+    routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
+    assert routed is host
+    assert repo == FORK
+    assert persist is None
+
+
+@pytest.mark.parametrize("mask_value", ["", "  ", "-1"])
+def test_auto_vulkan_still_fires_when_the_mask_hides_every_amd_gpu(mask_value, monkeypatch):
+    # Negative control: an empty / "-1" mask is not a partial view, it hides every AMD
+    # GPU, and _pick_rocm_gfx_target already reports that as no active target. Such a
+    # value must not be read as "the probe list may be incomplete" and suppress the
+    # fallback on a host whose arch came from --rocm-gfx.
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", mask_value)
+    host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
+    assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is True
+    _routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
+    assert repo == UPSTREAM
+    assert persist == "vulkan"
+
+
+def test_hip_device_mask_check_takes_the_first_variable_that_is_set(monkeypatch):
+    # Same first-var-wins order as _pick_rocm_gfx_target, which reads the identical
+    # three: HIP_VISIBLE_DEVICES="" means no AMD GPU visible even with a later
+    # CUDA_VISIBLE_DEVICES set, so the two must not disagree about the host.
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "")
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
+    assert ilp._hip_visible_device_mask_set() is False
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
+    assert ilp._hip_visible_device_mask_set() is True
+
+
+def test_masked_probe_suppression_does_not_touch_non_amd_auto_paths(monkeypatch):
+    # The mask says nothing about an Intel iGPU, whose Vulkan auto path is unrelated
+    # and predates this feature.
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
+    host = _host(
+        system = "Windows", is_windows = True, has_intel_gpu = True,
+        has_rocm = False, has_physical_nvidia = False, has_usable_nvidia = False,
+    )
+    _routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
+    assert repo == UPSTREAM
+
+
 def test_route_to_vulkan_prebuilt_hip_masked_host_still_honours_explicit_optin(monkeypatch):
     # The mask guard only suppresses the AUTOMATIC fallback; an explicit opt-in is
     # the user taking responsibility for the Vulkan device mask themselves.

studio/install_llama_prebuilt.py

@@ -6165,6 +6165,24 @@ def _active_rocm_gfx_target(host: HostInfo) -> str | None:
     return None
 
 
+def _hip_visible_device_mask_set() -> bool:
+    """Whether a HIP visible-device mask is filtering this process's view of the GPUs.
+
+    The Windows arch probe is hipinfo, itself a HIP application, and AMD documents these
+    as "only devices whose index is present in the sequence are visible to HIP" (with
+    HIP_VISIBLE_DEVICES the Windows spelling), so under a mask it enumerates the VISIBLE
+    devices, not the physical ones. Same first-var-wins order as _pick_rocm_gfx_target,
+    which reads the identical three. Empty / "-1" hides every AMD GPU, which that
+    function already reports as no active target, so it is not a partial mask here."""
+    for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+        _val = os.environ.get(_env)
+        if _val is None:
+            continue
+        _val = _val.strip()
+        return bool(_val) and _val != "-1"
+    return False
+
+
 def _windows_hip_gfx_targets(published_repo: str | None) -> frozenset[str]:
     """gfx targets the Windows HIP bundle of ``published_repo`` is actually built for.
 
@@ -6217,6 +6235,15 @@ def _should_auto_vulkan_for_amd_windows(host: HostInfo, published_repo: str | No
     # would happily enumerate the HIP-capable card the user deliberately hid, possibly
     # one reserved for another workload. Auto-fall back only when no AMD device on the
     # box can be exposed to HIP; an explicit vulkan opt-in is unaffected.
+    #
+    # Under a mask the probe cannot supply that inventory at all: hipinfo is a HIP
+    # application, so it enumerates only the visible devices and rocm_gfx_targets lists
+    # what survived the mask rather than what is installed. "No AMD GPU here reaches the
+    # floor" is then unprovable, and guessing wrong is the same reserved-card handover,
+    # so decline to guess. The mask is only ever set deliberately, and the driver-only
+    # single-GPU host this fallback exists for does not set one.
+    if _hip_visible_device_mask_set():
+        return False
     targets = list(dict.fromkeys([*_host_rocm_gfx_targets(host), active]))
     return not any(_gfx_is_windows_hip_supported(target, published_repo) for target in targets)