fix(studio/rocm): worker BNB/grouped_mm broad gate, install.sh Strix visibility, runtime-only ROCm detection

Round-5 robustness pass based on 20 parallel reviewers of head 96b9e465.

1. studio/backend/core/training/worker.py - BNB version pin / dynamo disable
   / _grouped_mm fallback block was still gated on torch.version.hip alone
   despite the torchao stub block above already using the broad check. AMD
   SDK / Radeon Windows wheels (torch.__version__ contains "rocm" but
   torch.version.hip is None) silently skipped the Windows ROCm runtime
   patches. Aligned to the same broad check (8/20 reviewers).

2. studio/backend/core/training/worker.py - _hip_ver_at_least() now also
   parses the ROCm version out of torch.__version__ (e.g. "2.11.0+rocm7.13.0")
   when torch.version.hip is missing, so the kernel-fix gate is correct for
   SDK / Radeon wheels too.

3. studio/backend/core/training/worker.py - _grouped_mm_safe_impl with
   offs=None now picks torch.bmm/matmul for 3-D inputs instead of always
   calling torch.mm. The real _grouped_mm accepts 3-D batched matmul; the
   prior fallback raised "self must be a matrix" on MoE workloads (2/20).

4. studio/backend/main.py - dropped the HIP_PATH / ROCM_PATH env-var gate
   from the BNB block; probe torch directly. Runtime-only Radeon / AMD SDK
   Windows installs do not set those SDK env vars but still ship ROCm torch
   (5/20 reviewers).

5. install.sh - Strix override now collects every gfx token from
   rocminfo / amd-smi (in enumeration order), then indexes by
   HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-
   Strix dGPU host where the user selected the dGPU does NOT get rerouted
   to the Strix per-gfx index. Mirrors the Python update path (5/20 reviewers).

6. install.sh - Strix detection chain now also probes `amd-smi static --asic`,
   matching the PowerShell installer (1/20). Closes the gap on runtime-only
   Strix hosts where `amd-smi list` does not surface a gfx token.

7. studio/install_python_stack.py - _has_rocm_gpu() now has the sysfs KFD
   topology fallback (/sys/class/kfd/kfd/topology/nodes/*/gpu_id), matching
   install.sh. On minimal package-managed installs without rocminfo /
   amd-smi GUI tools, `studio update` can now detect the GPU and repair the
   venv instead of returning early (2/20).

8. studio/install_python_stack.py - _detect_amd_gfx_codes() now falls back
   to `amd-smi list` and `amd-smi static --asic` when rocminfo is missing
   (2/20). Strix routing on runtime-only Radeon hosts now matches what
   install.sh has done for a while.

9. studio/install_python_stack.py - Strix override now applies even when
   has_hip_torch is True. The whole point of the override is to repair an
   existing broken torch.version.hip == "7.1" install; skipping the
   reinstall left users on the known _grouped_mm segfaulting stack (3/20).

Tests: 231 passed, 1 skipped. sim_5301 30 cases pass. sim_cross 12 pass.

install.sh

@@ -1832,13 +1832,45 @@ esac
 # Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
 case "$TORCH_INDEX_URL" in
     */rocm7.1|*/rocm7.1.*)
-        _strix_gfx=""
+        # Collect every gfx token in rocminfo / amd-smi enumeration order
+        # (skip duplicates), then index by HIP_VISIBLE_DEVICES /
+        # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
+        # where the user selected the dGPU does NOT get rerouted to the
+        # Strix per-gfx index.
+        _gfx_all=""
         if command -v rocminfo >/dev/null 2>&1; then
-            _strix_gfx=$(rocminfo 2>/dev/null | grep -oE 'gfx1151|gfx1150' | head -1)
+            _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++')
         fi
-        if [ -z "$_strix_gfx" ] && command -v amd-smi >/dev/null 2>&1; then
-            _strix_gfx=$(amd-smi list 2>/dev/null | grep -oE 'gfx1151|gfx1150' | head -1)
+        if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
+            _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++')
+            # PowerShell paths also probe `amd-smi static --asic`; mirror it
+            # so a host with hipinfo-less amd-smi reports the gfx target.
+            if [ -z "$_gfx_all" ]; then
+                _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++')
+            fi
         fi
+        _runtime_gfx=""
+        if [ -n "$_gfx_all" ]; then
+            _vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
+            _idx=0
+            if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then
+                _first=${_vis%%,*}
+                case "$_first" in
+                    ''|*[!0-9]*) _idx=0 ;;
+                    *) _idx=$_first ;;
+                esac
+            fi
+            _runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" '
+                NF { vals[n++] = $0 }
+                END {
+                    if (idx < 0 || idx >= n) idx = 0
+                    if (n > 0) print vals[idx]
+                }')
+        fi
+        _strix_gfx=""
+        case "$_runtime_gfx" in
+            gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
+        esac
         if [ -n "$_strix_gfx" ]; then
             echo "" >&2
             echo "  [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2

studio/backend/core/training/worker.py

@@ -2060,9 +2060,24 @@ def run_training_process(
     global _WINDOWS_ROCM_GROUPED_MM_LIB
     if sys.platform == "win32":
         _torch_for_rocm = sys.modules.get("torch")
-        if _torch_for_rocm is not None and getattr(
-            getattr(_torch_for_rocm, "version", None), "hip", None
-        ):
+        # Broad check: torch.version.hip OR "rocm" in torch.__version__.
+        # AMD SDK / Radeon Windows wheels do not always populate
+        # torch.version.hip; without the broad check the BNB version pin,
+        # dynamo-disable, and _grouped_mm fallback below silently skip
+        # (matches the torchao stub gate above and main.py).
+        _build_version_for_rocm = (
+            getattr(_torch_for_rocm, "__version__", "").lower()
+            if _torch_for_rocm is not None
+            else ""
+        )
+        _is_win_rocm_torch = bool(
+            _torch_for_rocm is not None
+            and (
+                getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
+                or "rocm" in _build_version_for_rocm
+            )
+        )
+        if _is_win_rocm_torch:
             # Disable dynamo (belt-and-suspenders; JitDecomp patch below is the
             # real fix, but keeping dynamo off avoids any other compile paths).
             if "TORCHDYNAMO_DISABLE" not in os.environ:
@@ -2112,12 +2127,24 @@ def run_training_process(
 
             # Parse HIP version for the kernel-fix gate below.
             # torch.version.hip can be "7.13.99004", "7.2.0", etc.
-            # We only need major.minor for the comparison.
+            # AMD SDK / Radeon wheels may leave torch.version.hip unset and
+            # encode the ROCm version in torch.__version__ instead
+            # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
+            # to that string when version.hip is missing.
             def _hip_ver_at_least(major: int, minor: int) -> bool:
+                import re as _re_ver
                 _hip_str = getattr(
                     getattr(_torch_for_rocm, "version", None), "hip", None
                 )
                 if not _hip_str:
+                    _ver_match = _re_ver.search(
+                        r"rocm(\d+)\.(\d+)", _build_version_for_rocm
+                    )
+                    if _ver_match:
+                        return (
+                            int(_ver_match.group(1)),
+                            int(_ver_match.group(2)),
+                        ) >= (major, minor)
                     return False
                 try:
                     _parts = [int(x) for x in str(_hip_str).split(".")[:2]]
@@ -2138,11 +2165,24 @@ def run_training_process(
                     def _grouped_mm_safe_impl(
                         self, mat2, offs = None, bias = None, out_dtype = None
                     ):
-                        """Python mm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
+                        """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
                         _t = _torch_for_rocm
                         if offs is None:
-                            # Simple case: plain matrix multiply.
-                            result = _t.mm(self.contiguous(), mat2.contiguous())
+                            # No offsets: behave like the real op, which
+                            # accepts either (M, K) x (K, N) -> mm, or 3-D
+                            # batched inputs -> bmm. Picking torch.mm
+                            # unconditionally previously raised "self must be
+                            # a matrix" on 3-D MoE workloads.
+                            if self.dim() == 3 and mat2.dim() == 3:
+                                result = _t.bmm(self.contiguous(), mat2.contiguous())
+                            elif self.dim() == 3 and mat2.dim() == 2:
+                                # Broadcast 2-D mat2 across the batch dim.
+                                result = _t.matmul(self.contiguous(), mat2.contiguous())
+                            elif self.dim() == 2 and mat2.dim() == 3:
+                                # Broadcast 2-D self across batch via matmul semantics.
+                                result = _t.matmul(self.contiguous(), mat2.contiguous())
+                            else:
+                                result = _t.mm(self.contiguous(), mat2.contiguous())
                         else:
                             # Grouped case: offs[i] is the exclusive end-row of
                             # group i in `self`; mat2 may be 3-D or 2-D.

studio/backend/main.py

@@ -70,25 +70,22 @@ if sys.platform == "win32":
     # this the server process crashes with "Configured ROCm binary not found".
     # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
     # before any import that pulls in bitsandbytes (mirrors worker.py logic).
-    # Gate on the active torch runtime, not env-var presence -- HIP_PATH /
-    # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
-    # CUDA torch wheel, and setting BNB_ROCM_VERSION there makes bitsandbytes
-    # look for a ROCm DLL that doesn't exist and crash the CUDA backend.
+    # Gate on the active torch runtime only. AMD SDK / Radeon Windows wheels
+    # may not set HIP_PATH / ROCM_PATH, but they do populate torch.version.hip
+    # or encode "rocm" in torch.__version__. A previous version of this gate
+    # required HIP_PATH / ROCM_PATH and silently skipped BNB_ROCM_VERSION for
+    # those wheels.
     _is_rocm_host = False
-    if os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"):
-        try:
-            import torch as _torch_probe
+    try:
+        import torch as _torch_probe
 
-            # Broad check: torch.version.hip OR "rocm" in torch.__version__ --
-            # AMD SDK / Radeon wheels may not populate torch.version.hip but
-            # still encode "rocm" in __version__. Matches worker.py + hardware.py.
-            _is_rocm_host = bool(
-                getattr(getattr(_torch_probe, "version", None), "hip", None)
-                or "rocm" in getattr(_torch_probe, "__version__", "").lower()
-            )
-            del _torch_probe
-        except Exception:
-            pass
+        _is_rocm_host = bool(
+            getattr(getattr(_torch_probe, "version", None), "hip", None)
+            or "rocm" in getattr(_torch_probe, "__version__", "").lower()
+        )
+        del _torch_probe
+    except Exception:
+        pass
     if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ:
         import glob as _glob
 

studio/install_python_stack.py

@@ -405,6 +405,26 @@ def _has_rocm_gpu() -> bool:
         if result.returncode == 0 and result.stdout.strip():
             if check_fn(result.stdout):
                 return True
+    # sysfs KFD topology fallback (Linux only) -- matches install.sh's
+    # runtime-only detection. On minimal package-managed installs (no
+    # rocminfo / no amd-smi GUI tools), the kernel exposes AMD GPUs via
+    # /sys/class/kfd so `studio update` can still detect the GPU and
+    # repair the venv.
+    if sys.platform != "win32":
+        try:
+            kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
+            if os.path.isdir(kfd_nodes):
+                for entry in os.listdir(kfd_nodes):
+                    gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id")
+                    try:
+                        with open(gpu_id_path) as fh:
+                            gpu_id = fh.read().strip()
+                    except OSError:
+                        continue
+                    if gpu_id and gpu_id != "0":  # gpu_id 0 = CPU node
+                        return True
+        except OSError:
+            pass
     return False
 
 
@@ -429,30 +449,40 @@ def _has_usable_nvidia_gpu() -> bool:
 def _detect_amd_gfx_codes() -> list[str]:
     """Return the list of AMD gfx ISA strings visible to ROCm (e.g. ['gfx1151']).
 
-    Parses ``rocminfo`` output for ``ISA Info`` / ``gfx`` entries.  Returns an
-    empty list when rocminfo is not found or no GPU agents are present.
+    Probes rocminfo first, then falls back to ``amd-smi list`` and
+    ``amd-smi static --asic`` for runtime-only Radeon hosts that ship
+    amd-smi but no rocminfo. Returns an empty list when no probe yields
+    a gfx target.
     """
     import re
 
-    exe = shutil.which("rocminfo")
-    if not exe:
-        return []
-    try:
-        result = subprocess.run(
-            [exe],
-            stdout = subprocess.PIPE,
-            stderr = subprocess.DEVNULL,
-            text = True,
-            timeout = 15,
-        )
-    except Exception:
-        return []
-    if result.returncode != 0:
-        return []
-    # Match lines like "  Name:                    gfx1151" or ISA strings
-    # "amdgcn-amd-amdhsa--gfx1151".  Exclude the CPU agent (gfx000).
-    codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", result.stdout.lower())
-    return list(dict.fromkeys(f"gfx{c}" for c in codes))  # deduplicate, preserve order
+    def _extract(text: str) -> list[str]:
+        codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower())
+        return list(dict.fromkeys(f"gfx{c}" for c in codes))
+
+    probes: list[list[str]] = []
+    if shutil.which("rocminfo"):
+        probes.append(["rocminfo"])
+    if shutil.which("amd-smi"):
+        probes.append(["amd-smi", "list"])
+        probes.append(["amd-smi", "static", "--asic"])
+    for cmd in probes:
+        try:
+            result = subprocess.run(
+                cmd,
+                stdout = subprocess.PIPE,
+                stderr = subprocess.DEVNULL,
+                text = True,
+                timeout = 15,
+            )
+        except Exception:
+            continue
+        if result.returncode != 0 or not result.stdout.strip():
+            continue
+        codes = _extract(result.stdout)
+        if codes:
+            return codes
+    return []
 
 
 # Set by _ensure_rocm_torch() on success; suppresses the post-install AMD warning.
@@ -709,13 +739,49 @@ def _ensure_rocm_torch() -> None:
                     f"   skipping AMD per-gfx index override.\n"
                 )
 
-    if not has_hip_torch:
-        if _strix_override_url is not None and _strix_override_pkgs is not None:
-            index_url = _strix_override_url
-            _torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
-            print(f"   Strix ROCm 7.1 override -- installing torch from {index_url}")
+    # Strix override on ROCm 7.1 must fire even when has_hip_torch is True --
+    # an existing torch with `torch.version.hip == "7.1"` is exactly the broken
+    # combo the override is meant to repair, so skipping it leaves users on
+    # the known _grouped_mm segfault.
+    if _strix_override_url is not None and _strix_override_pkgs is not None:
+        index_url = _strix_override_url
+        _torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
+        print(f"   Strix ROCm 7.1 override -- installing torch from {index_url}")
+        pip_install(
+            "ROCm torch (Strix arch-specific)",
+            "--force-reinstall",
+            "--no-cache-dir",
+            _torch_pkg,
+            _vision_pkg,
+            _audio_pkg,
+            "--index-url",
+            index_url,
+            constrain = False,
+        )
+        rocm_torch_ready = True
+    elif not has_hip_torch:
+        # Select best matching wheel tag (newest ROCm version <= installed)
+        tag = next(
+            (
+                t
+                for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
+                if ver >= (maj, mn)
+            ),
+            None,
+        )
+        if tag is None:
+            print(
+                f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- "
+                f"skipping torch reinstall"
+            )
+        else:
+            index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
+            print(f"   ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
+            _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
+                tag, _ROCM_TORCH_PKG_SPECS["_default"]
+            )
             pip_install(
-                "ROCm torch (Strix arch-specific)",
+                f"ROCm torch ({tag})",
                 "--force-reinstall",
                 "--no-cache-dir",
                 _torch_pkg,
@@ -726,39 +792,6 @@ def _ensure_rocm_torch() -> None:
                 constrain = False,
             )
             rocm_torch_ready = True
-        else:
-            # Select best matching wheel tag (newest ROCm version <= installed)
-            tag = next(
-                (
-                    t
-                    for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
-                    if ver >= (maj, mn)
-                ),
-                None,
-            )
-            if tag is None:
-                print(
-                    f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- "
-                    f"skipping torch reinstall"
-                )
-            else:
-                index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
-                print(f"   ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
-                _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
-                    tag, _ROCM_TORCH_PKG_SPECS["_default"]
-                )
-                pip_install(
-                    f"ROCm torch ({tag})",
-                    "--force-reinstall",
-                    "--no-cache-dir",
-                    _torch_pkg,
-                    _vision_pkg,
-                    _audio_pkg,
-                    "--index-url",
-                    index_url,
-                    constrain = False,
-                )
-                rocm_torch_ready = True
 
     # Install bitsandbytes only when torch links against ROCm. Prefers the
     # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls

tests/studio/install/test_rocm_support.py

@@ -2338,7 +2338,7 @@ class TestStrixRocm71Override:
         strix_idx = source.find("_strix_gfx")
         assert strix_idx != -1
         # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx
-        context_before = source[max(0, strix_idx - 600) : strix_idx]
+        context_before = source[max(0, strix_idx - 2400) : strix_idx]
         assert "rocm7.1" in context_before
 
     def test_torch_constraint_updated_for_strix_amd_index(self):