Studio: forward the resolved AMD gfx arch to the prebuilt installer (#5923)

setup.sh and setup.ps1 resolve the AMD gfx target (rocminfo/hipinfo/amd-smi,
name inference, or UNSLOTH_ROCM_GFX_ARCH) but never passed it to
install_llama_prebuilt.py. The installer re-probed on its own and, on hosts where
hipinfo/amd-smi cannot report the arch (amd-smi-only Linux, HIP-runtime-only
Strix Halo, name-inferred GPUs), left rocm_gfx_target unset. The lemonade HIP
prebuilt selection needs that arch, so those hosts got no GPU prebuilt and fell
back to a source build.

Add a --rocm-gfx argument (defaulting to UNSLOTH_ROCM_GFX_ARCH) and fold it into
the host profile via _apply_host_overrides: a forwarded gfx is authoritative
(setup already applied visible-device selection) and implies has_rocm. setup.sh
and setup.ps1 now forward their resolved arch.

Add unit tests for _normalize_forwarded_gfx and _apply_host_overrides, plus
source checks that both setup scripts forward --rocm-gfx.

studio/install_llama_prebuilt.py

@@ -3085,6 +3085,46 @@ def detect_host() -> HostInfo:
     )
 
 
+def _normalize_forwarded_gfx(value: str | None) -> str | None:
+    """Extract a single gfx token from a forwarded --rocm-gfx / env value.
+    setup.sh/setup.ps1 already picked the active GPU, so take the token as-is
+    without re-applying visible-device selection. Ignore anything malformed."""
+    if not value:
+        return None
+    m = re.search(r"gfx[1-9][0-9a-z]{2,3}", value.lower())
+    return m.group(0) if m else None
+
+
+def _apply_host_overrides(
+    host: HostInfo,
+    *,
+    override_has_rocm: bool = False,
+    override_rocm_gfx: str | None = None,
+    force_cpu: bool = False,
+) -> HostInfo:
+    """Fold setup.sh/setup.ps1's forwarded detection into the host profile.
+    A forwarded gfx (--rocm-gfx or UNSLOTH_ROCM_GFX_ARCH) is authoritative and
+    implies ROCm: the installer's own hipinfo/amd-smi probe can miss the arch on
+    amd-smi-only hosts or when setup inferred it from the GPU name, leaving
+    rocm_gfx_target None and no lemonade prebuilt selected. force_cpu is the
+    opposite explicit signal (arm64 Linux GPU host whose source build failed):
+    drop GPU attributes so the CPU prebuilt for this OS/arch is selected."""
+    if force_cpu:
+        return dataclasses_replace(
+            host,
+            has_usable_nvidia = False,
+            has_physical_nvidia = False,
+            has_rocm = False,
+            rocm_gfx_target = None,
+        )
+    gfx = _normalize_forwarded_gfx(override_rocm_gfx)
+    if gfx:
+        return dataclasses_replace(host, has_rocm = True, rocm_gfx_target = gfx)
+    if override_has_rocm and not host.has_rocm:
+        return dataclasses_replace(host, has_rocm = True)
+    return host
+
+
 def pick_windows_cuda_runtime(host: HostInfo) -> str | None:
     if not host.driver_cuda_version:
         return None
@@ -6460,22 +6500,16 @@ def install_prebuilt(
     *,
     simple_policy: bool = False,
     override_has_rocm: bool = False,
+    override_rocm_gfx: str | None = None,
     force_cpu: bool = False,
 ) -> None:
     host = detect_host()
-    if override_has_rocm and not host.has_rocm:
-        host = dataclasses_replace(host, has_rocm = True)
-    if force_cpu:
-        # Explicit CPU fallback: drop GPU attributes so the CPU prebuilt for this
-        # OS/arch is selected. setup.sh uses this for arm64 Linux GPU hosts whose
-        # source build failed, where no arm64 CUDA prebuilt exists anywhere.
-        host = dataclasses_replace(
-            host,
-            has_usable_nvidia = False,
-            has_physical_nvidia = False,
-            has_rocm = False,
-            rocm_gfx_target = None,
-        )
+    host = _apply_host_overrides(
+        host,
+        override_has_rocm = override_has_rocm,
+        override_rocm_gfx = override_rocm_gfx,
+        force_cpu = force_cpu,
+    )
     choice: AssetChoice | None = None
     try:
         with install_lock(install_lock_path(install_dir)):
@@ -6620,6 +6654,16 @@ def parse_args() -> argparse.Namespace:
             "so the HIP llama.cpp prebuilt is selected even when hipinfo is not on PATH."
         ),
     )
+    parser.add_argument(
+        "--rocm-gfx",
+        default = os.environ.get("UNSLOTH_ROCM_GFX_ARCH"),
+        help = (
+            "Forward the AMD gfx target (e.g. gfx1151) that setup.ps1/setup.sh "
+            "resolved, so the lemonade HIP prebuilt is selected even when the "
+            "installer's own hipinfo/amd-smi probe cannot report it. Implies "
+            "--has-rocm. Defaults to the UNSLOTH_ROCM_GFX_ARCH environment variable."
+        ),
+    )
     parser.add_argument(
         "--cpu-fallback",
         action = "store_true",
@@ -6751,6 +6795,7 @@ def main() -> int:
         published_release_tag = args.published_release_tag or "",
         simple_policy = args.simple_policy,
         override_has_rocm = args.has_rocm,
+        override_rocm_gfx = args.rocm_gfx,
         force_cpu = args.cpu_fallback,
     )
     return EXIT_SUCCESS

studio/setup.ps1

@@ -2483,6 +2483,12 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
         )
         if ($HasROCm) {
             $prebuiltArgs += "--has-rocm"
+            # Forward the resolved gfx arch so the lemonade HIP prebuilt is picked
+            # even when the installer's own probe cannot report it (amd-smi-only
+            # hosts, name-inferred arch).
+            if ($script:ROCmGfxArch) {
+                $prebuiltArgs += @("--rocm-gfx", $script:ROCmGfxArch)
+            }
         }
         if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
             $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)

studio/setup.sh

@@ -857,6 +857,12 @@ else
     if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
         _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
     fi
+    # Forward the gfx arch resolved above so the lemonade HIP prebuilt is picked
+    # even when the installer's own probe cannot report it (amd-smi-only hosts,
+    # name-inferred arch). Implies --has-rocm on the installer side.
+    if [ -n "${_setup_gfx:-}" ]; then
+        _PREBUILT_CMD+=(--rocm-gfx "$_setup_gfx")
+    fi
     _PREBUILT_LOG="$(mktemp)"
     set +e
     if _is_verbose; then

tests/studio/install/test_rocm_support.py

@@ -35,6 +35,8 @@ AssetChoice = prebuilt_mod.AssetChoice
 PrebuiltFallback = prebuilt_mod.PrebuiltFallback
 resolve_upstream_asset_choice = prebuilt_mod.resolve_upstream_asset_choice
 runtime_patterns_for_choice = prebuilt_mod.runtime_patterns_for_choice
+_apply_host_overrides = prebuilt_mod._apply_host_overrides
+_normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx
 
 # install_python_stack.py
 _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
@@ -2598,5 +2600,93 @@ class TestHipSdkInstalledButDeviceInaccessible:
         assert "GPU not ROCm-accessible" in source
 
 
+# =============================================================================
+# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch
+# to install_llama_prebuilt.py so the lemonade HIP prebuilt is selected even when
+# the installer's own hipinfo/amd-smi probe cannot report it.
+# =============================================================================
+
+_SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"
+
+
+class TestNormalizeForwardedGfx:
+    """A forwarded gfx string is reduced to a single clean gfx token."""
+
+    def test_plain_token(self):
+        assert _normalize_forwarded_gfx("gfx1151") == "gfx1151"
+
+    def test_uppercase_normalized(self):
+        assert _normalize_forwarded_gfx("GFX1151") == "gfx1151"
+
+    def test_extracts_from_noise(self):
+        assert _normalize_forwarded_gfx("gcnArchName: gfx942") == "gfx942"
+
+    def test_malformed_is_ignored(self):
+        assert _normalize_forwarded_gfx("not-a-gpu") is None
+
+    def test_empty_and_none(self):
+        assert _normalize_forwarded_gfx("") is None
+        assert _normalize_forwarded_gfx(None) is None
+
+
+class TestApplyHostOverrides:
+    """Forwarded ROCm detection is folded into the host profile correctly."""
+
+    def test_forwarded_gfx_fills_empty_probe(self):
+        # amd-smi-only / name-inferred host: installer probe found no gfx.
+        host = rocm_host(rocm_gfx_target = None)
+        out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
+        assert out.has_rocm is True
+        assert out.rocm_gfx_target == "gfx1151"
+
+    def test_forwarded_gfx_implies_rocm(self):
+        # A CPU-looking host with a forwarded gfx is an AMD host.
+        out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx1200")
+        assert out.has_rocm is True
+        assert out.rocm_gfx_target == "gfx1200"
+
+    def test_forwarded_gfx_is_authoritative(self):
+        # setup already applied visible-device selection; its value wins.
+        host = rocm_host(rocm_gfx_target = "gfx1100")
+        out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
+        assert out.rocm_gfx_target == "gfx1151"
+
+    def test_has_rocm_only_keeps_probe_gfx(self):
+        out = _apply_host_overrides(cpu_host(), override_has_rocm = True)
+        assert out.has_rocm is True
+        assert out.rocm_gfx_target is None
+
+    def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self):
+        out = _apply_host_overrides(
+            cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk"
+        )
+        assert out.has_rocm is True
+        assert out.rocm_gfx_target is None
+
+    def test_no_overrides_leaves_host_unchanged(self):
+        host = nvidia_host()
+        assert _apply_host_overrides(host) is host
+
+
+class TestRocmGfxForwarding:
+    """setup.sh / setup.ps1 forward their resolved gfx; the installer accepts it."""
+
+    def test_installer_exposes_rocm_gfx_arg(self):
+        source = _PREBUILT_PATH.read_text(encoding = "utf-8")
+        assert '"--rocm-gfx"' in source
+        # Defaults to the env override so a standalone run still works.
+        assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source
+
+    def test_setup_sh_forwards_rocm_gfx(self):
+        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
+        assert "--rocm-gfx" in source
+        assert '"$_setup_gfx"' in source
+
+    def test_setup_ps1_forwards_rocm_gfx(self):
+        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
+        assert "--rocm-gfx" in source
+        assert "$script:ROCmGfxArch" in source
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])