From 240c0c350027e74dfef00a44eead7154768705fb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 04:46:39 -0700 Subject: [PATCH] Studio: fix WSL Strix Halo GPU on reinstall (ROCDXG drop-in + system HIP before bundle) (#6227) * install.sh: persist ROCm-on-WSL drop-in even when rocminfo already works _maybe_bootstrap_rocm_wsl calls _ensure_rocm_probe_env (which exports a transient HSA_ENABLE_DXG_DETECTION + adds /opt/rocm/bin to PATH on the installer process) right before the "rocminfo enumerates gfx1151 -> already set up, return early" gate. On any reinstall over an existing /opt/rocm -- the common case, since the uninstaller keeps shared ROCm userspace but removes /etc/profile.d/unsloth-rocm-wsl.sh -- that probe env makes rocminfo succeed, so the gate returns 0 WITHOUT ever persisting the drop-in. The transient env dies with the installer, so the next login shell (Studio, llama-server) sees no GPU: torch cuda_avail=False, rocminfo finds nothing, the llama.cpp ROCm prebuilt segfaults on a GPU it can't reach. Factor the drop-in writer into _persist_rocm_wsl_dropin() and call it before the early return so the persistent env is restored whenever librocdxg is present. Idempotent (only writes when the drop-in is missing), gated on librocdxg so it never fires on non-WSL/non-ROCDXG hosts, root-writes or sudo-tees like before. The fast-path branch now reuses the same helper. Reproduced on gfx1151 (Radeon 8060S) under dash (the curl|sh shell): before the fix a reinstall left the drop-in absent and torch cuda_avail False; after, the drop-in is persisted and a fresh login shell reports cuda_avail True. Verified under both dash and bash, and idempotent on re-run. * Studio WSL: load system HIP before a prebuilt's bundled runtime (gfx1151) The lemonade / published llama.cpp ROCm prebuilts bundle their own HIP runtime (libamdhip64) built for bare-metal Linux. In WSL the GPU is reached through the system ROCm's librocdxg bridge over /dev/dxg, which the bundled runtime cannot drive -- it segfaults on the first GPU call. So: - install_llama_prebuilt.py: the prebuilt's llama-quantize/llama-server validation runs with the bundle dir first on LD_LIBRARY_PATH, segfaults (empty stderr), and the install silently falls back to a CPU source build (which on this host can't even build for GPU -- hipcc absent). The Strix Halo WSL user ends up on CPU despite a working GPU. - llama_cpp.py: even if a GPU prebuilt were kept, the serve-time launcher put the bundle dir first too, so it would crash at load. Fix: on a ROCDXG WSL host (gated on /dev/dxg + "microsoft" /proc/version + a librocdxg-providing /opt/rocm), prepend the system ROCm lib dir to LD_LIBRARY_PATH so the WSL-capable libamdhip64 + librocdxg load first, while the bundle still supplies libggml-hip / librocblas with the gfx1151 kernels. Set HSA_ENABLE_DXG_DETECTION=1 alongside. Added _wsl_system_rocm_lib_dirs() to both modules (kept identical so a prebuilt that passed install validation runs the same way at serve time). Strict no-op on bare-metal Linux, NVIDIA, macOS, and Windows. Verified on gfx1151 (Radeon 8060S) in WSL (ROCm 7.2.1 + librocdxg, Adrenalin ROCDXG): before, the lemonade gfx1151 prebuilt segfaulted and the install fell back to a broken CPU build; after, install_llama_prebuilt validates and keeps the GPU prebuilt (source=published, prebuilt_fallback_used=False), and Studio serves Qwen3-1.7B-GGUF at 53 tok/s with the model resident in GPU memory (llama-server device_info: ROCm0 = AMD Radeon 8060S). * tests: cover the WSL ROCDXG drop-in + system-HIP-ordering fixes - _wsl_system_rocm_lib_dirs: no-op without /dev/dxg, on bare-metal Linux, and on WSL without librocdxg; returns the system lib dir on a ROCDXG WSL host. - binary_env: prepends the system ROCm lib dir ahead of the bundle and sets HSA_ENABLE_DXG_DETECTION on WSL; unchanged on bare-metal Linux. - install.sh: _persist_rocm_wsl_dropin exists, is gated on librocdxg, and the rocminfo-already-works early return calls it before returning. - llama_cpp.py: the serve-time launcher prepends the WSL rocm dirs before the bundle dir (mirrors binary_env). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten WSL ROCDXG fix comments (no logic change) Condense the drop-in / system-HIP-ordering comments and docstrings added in this PR. Verified comment-only via AST parse + py_compile + sh/bash -n, the 308-test rocm_support suite, and a dash functional re-run of the bootstrap (drop-in still persisted, env still set). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.sh | 63 +++++---- studio/backend/core/inference/llama_cpp.py | 37 ++++- studio/install_llama_prebuilt.py | 32 +++++ tests/studio/install/test_rocm_support.py | 154 ++++++++++++++++++++- 4 files changed, 259 insertions(+), 27 deletions(-) diff --git a/install.sh b/install.sh index eba92d074..af062ae05 100755 --- a/install.sh +++ b/install.sh @@ -2047,6 +2047,37 @@ _pick_radeon_wheel() { # CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts # the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + # librocdxg), then sources the env it persisted so detection finds the GPU. +# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d +# so non-login Studio/llama launches inherit it. Idempotent (writes only when +# the drop-in is missing); no-op without librocdxg, so never fires off WSL. +# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes +# after this shell on a non-root reinstall. Best-effort either way. +_persist_rocm_wsl_dropin() { + [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0 + _rw_rocm=/opt/rocm + export HSA_ENABLE_DXG_DETECTION=1 + export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 + case ":${PATH}:" in + *":${_rw_rocm}/bin:"*) ;; + *) export PATH="${_rw_rocm}/bin:${PATH}" ;; + esac + export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" + [ -r /etc/profile.d/unsloth-rocm-wsl.sh ] && return 0 + _rw_dropin="$( + printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' + printf 'export HSA_ENABLE_DXG_DETECTION=1\n' + printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' + printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" + printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" + printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' + )" + if [ "$(id -u)" = "0" ]; then + printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true + elif command -v sudo >/dev/null 2>&1; then + printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true + fi +} + _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 @@ -2060,6 +2091,11 @@ _maybe_bootstrap_rocm_wsl() { _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + # rocminfo may work only via the transient env _ensure_rocm_probe_env + # just set, which dies with the installer. Persist the drop-in so login + # shells (Studio, llama.cpp) inherit it -- else a reinstall over an + # existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU. + _persist_rocm_wsl_dropin return 0 fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). @@ -2077,31 +2113,8 @@ _maybe_bootstrap_rocm_wsl() { . /etc/profile.d/unsloth-rocm-wsl.sh || true else # librocdxg present but the env drop-in is gone (e.g. a Studio - # uninstall removed it while keeping shared ROCm). Restore the FULL - # env inline (so rocminfo is on PATH) and recreate the drop-in. - _rw_rocm=/opt/rocm - export HSA_ENABLE_DXG_DETECTION=1 - export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 - export PATH="${_rw_rocm}/bin:${PATH}" - export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" - # Persist the drop-in so later non-login Studio launches get the env - # too. /etc/profile.d is root-owned: a plain redirect fails for a - # non-root reinstall (ROCm would silently disappear after this shell), - # so tee through sudo when not root. Best-effort -- the current shell - # already has the env, so the install proceeds either way. - _rw_dropin="$( - printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' - printf 'export HSA_ENABLE_DXG_DETECTION=1\n' - printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' - printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" - printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" - printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' - )" - if [ "$(id -u)" = "0" ]; then - printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true - elif command -v sudo >/dev/null 2>&1; then - printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true - fi + # uninstall removed it while keeping shared ROCm). Restore the env. + _persist_rocm_wsl_dropin fi return 0 fi diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3cc7514ff..e2de5fc32 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -75,6 +75,33 @@ from state.tool_approvals import ( logger = get_logger(__name__) +def _wsl_system_rocm_lib_dirs() -> "list[str]": + """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL. + + The bundled bare-metal HIP can't drive WSL's /dev/dxg and segfaults on the + first GPU call; the system ROCm libs (libamdhip64 + librocdxg) can, while + the bundle still supplies libggml-hip / librocblas (gfx1151 kernels). + Mirrors install_llama_prebuilt._wsl_system_rocm_lib_dirs so a prebuilt that + passed install validation runs the same at serve time. No-op off a ROCDXG + WSL host (needs /dev/dxg, "microsoft" /proc/version, librocdxg in /opt/rocm). + """ + try: + if not os.path.exists("/dev/dxg"): + return [] + with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: + if "microsoft" not in fh.read().lower(): + return [] + except OSError: + return [] + out: "list[str]" = [] + for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): + if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists( + os.path.join(d, "librocdxg.so.1") + ): + out.append(d) + return out + + # ── Pre-compiled patterns for plan-without-action re-prompt ── # Forward-looking intent signals: the model is describing what it *will* # do rather than giving a final answer. @@ -3789,7 +3816,15 @@ class LlamaCppBackend: # plus CUDA runtime libs (libcudart, libcublas, etc.) import platform - lib_dirs = [binary_dir] + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on + # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which + # validates the prebuilt with this same ordering. + for _wsl_rocm in _wsl_system_rocm_lib_dirs(): + lib_dirs.append(_wsl_rocm) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) _arch = platform.machine() # x86_64, aarch64, etc. # Pip-installed nvidia CUDA runtime libs. The prebuilt diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index e3b2d493f..680b621b1 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -5662,6 +5662,33 @@ def windows_runtime_dirs_for_runtime_line(runtime_line: str | None) -> list[str] return windows_runtime_dirs_for_patterns(patterns) +def _wsl_system_rocm_lib_dirs() -> list[str]: + """System ROCm lib dir(s) for binary_env to load before a prebuilt's HIP. + + A prebuilt bundles a bare-metal HIP runtime that can't drive WSL's /dev/dxg + and segfaults on the first GPU call -> validation fails, install falls back + to a CPU build. Putting the system ROCm libs first loads the WSL-capable + HIP (libamdhip64 + librocdxg) while the bundle still supplies libggml-hip / + librocblas with the gfx1151 kernels. Strict no-op off WSL (needs /dev/dxg, a + "microsoft" /proc/version, and a librocdxg-providing ROCm). + """ + try: + if not os.path.exists("/dev/dxg"): + return [] + with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: + if "microsoft" not in fh.read().lower(): + return [] + except OSError: + return [] + out: list[str] = [] + for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): + if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists( + os.path.join(d, "librocdxg.so.1") + ): + out.append(d) + return out + + def binary_env( binary_path: Path, install_dir: Path, @@ -5683,6 +5710,11 @@ def binary_env( str(install_dir), *linux_runtime_dirs(binary_path), ] + # WSL: system HIP before the bundle's (which segfaults on /dev/dxg). + _wsl_rocm = _wsl_system_rocm_lib_dirs() + if _wsl_rocm: + ld_dirs = [*_wsl_rocm, *ld_dirs] + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") existing = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part] env["LD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*ld_dirs, *existing])) elif host.is_macos: diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 9aa06b6ff..04463fb72 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -11,7 +11,7 @@ import os import subprocess import sys from pathlib import Path -from unittest.mock import MagicMock, patch, PropertyMock +from unittest.mock import MagicMock, mock_open, patch, PropertyMock import pytest @@ -3207,5 +3207,157 @@ class TestRocmGfxForwarding: assert "$script:ROCmGfxArch" in source +# TEST: WSL ROCDXG fixes -- drop-in persistence + system-HIP-before-bundle + + +_INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh" +_LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + + +class TestWslSystemRocmLibDirs: + """install_llama_prebuilt._wsl_system_rocm_lib_dirs: a strict no-op off a + ROCDXG WSL host; returns the system ROCm lib dir there so binary_env can + load the WSL-capable HIP runtime before a prebuilt's bundled one.""" + + def test_empty_without_dev_dxg(self): + with patch("os.path.exists", return_value = False): + assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] + + def test_empty_on_bare_metal_linux(self): + # /dev/dxg present but /proc/version is not a WSL kernel. + with patch("os.path.exists", lambda p: p == "/dev/dxg"): + with patch( + "builtins.open", + mock_open(read_data = "Linux version 6.8.0-generic"), + ): + assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] + + def test_returns_system_lib_on_wsl_with_librocdxg(self): + # Normalize separators: os.path.join uses "\" on the Windows test host + # though the target path is Linux. + def _exists(p): + p = str(p).replace("\\", "/") + return p in ("/dev/dxg", "/opt/rocm/lib/librocdxg.so") + + with patch("os.path.exists", _exists): + with patch( + "builtins.open", + mock_open(read_data = "Linux version 5.15.0-microsoft-standard-WSL2"), + ): + assert prebuilt_mod._wsl_system_rocm_lib_dirs() == ["/opt/rocm/lib"] + + def test_empty_on_wsl_without_librocdxg(self): + # WSL kernel + /dev/dxg but no librocdxg -> not a ROCDXG ROCm install. + with patch("os.path.exists", lambda p: p == "/dev/dxg"): + with patch( + "builtins.open", + mock_open(read_data = "microsoft-standard-WSL2"), + ): + assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] + + +class TestBinaryEnvWslOrdering: + """binary_env must put the system ROCm lib dir AHEAD of the prebuilt's + bundle dir on a ROCDXG WSL host, and set HSA_ENABLE_DXG_DETECTION; no-op + on bare-metal Linux.""" + + @staticmethod + def _linux_host(): + return HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + ) + + def test_wsl_prepends_system_rocm_and_sets_hsa(self, tmp_path): + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + # binary_env's dedupe_existing_dirs drops non-existent dirs, so use a + # real dir to stand in for the system ROCm lib path. + sys_rocm = tmp_path / "sysrocm" + sys_rocm.mkdir() + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]): + with patch.dict(os.environ, {}, clear = True): + env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + ld = env["LD_LIBRARY_PATH"].split(os.pathsep) + # Compare resolved paths (dedupe_existing_dirs calls Path.resolve()). + ld_resolved = [str(Path(p).resolve()) for p in ld] + assert ld_resolved[0] == str(sys_rocm.resolve()) + assert str(binary.parent.resolve()) in ld_resolved + assert ld_resolved.index(str(sys_rocm.resolve())) < ld_resolved.index( + str(binary.parent.resolve()) + ) + assert env.get("HSA_ENABLE_DXG_DETECTION") == "1" + + def test_bare_metal_linux_unchanged(self, tmp_path): + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): + with patch.dict(os.environ, {}, clear = True): + env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + ld = env["LD_LIBRARY_PATH"].split(os.pathsep) + assert ld[0] == str(binary.parent) # bundle dir first, as before + assert "HSA_ENABLE_DXG_DETECTION" not in env + + +class TestInstallShDropinPersistence: + """install.sh must persist the ROCm-on-WSL drop-in even when rocminfo + already enumerates the GPU (via the transient probe env), so a reinstall + over an existing /opt/rocm doesn't leave login shells without the env.""" + + def test_has_persist_helper(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "_persist_rocm_wsl_dropin()" in source + + def test_gate5_early_return_persists_dropin(self): + """The rocminfo-already-works early return must call the persist helper + BEFORE returning.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + # Find the rocminfo gfx1151 gate and assert the persist call precedes + # its `return 0`. + gate = source.find("Name:[[:space:]]*gfx1151") + assert gate != -1 + window = source[gate : gate + 900] + assert "_persist_rocm_wsl_dropin" in window + assert window.find("_persist_rocm_wsl_dropin") < window.find("return 0") + + def test_persist_helper_gated_on_librocdxg(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + body_start = source.find("_persist_rocm_wsl_dropin()") + body = source[body_start : body_start + 1200] + assert "librocdxg.so" in body + assert "profile.d/unsloth-rocm-wsl.sh" in body + + +class TestLlamaCppRuntimeWslOrdering: + """The serve-time launcher must mirror binary_env: system HIP before the + bundle dir on WSL, so a prebuilt that passed install validation runs.""" + + def test_has_wsl_helper(self): + source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") + assert "_wsl_system_rocm_lib_dirs" in source + + def test_prepends_before_binary_dir(self): + source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") + # The Linux lib_dirs build must add WSL rocm dirs before binary_dir. + idx_helper = source.find("for _wsl_rocm in _wsl_system_rocm_lib_dirs()") + idx_binary = source.find("lib_dirs.append(binary_dir)") + assert idx_helper != -1 and idx_binary != -1 + assert idx_helper < idx_binary + + if __name__ == "__main__": pytest.main([__file__, "-v"])