From eff55fb8db25475089d3e9df6c4a9156c6b2e1b6 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 24 Apr 2026 15:50:25 +0000
Subject: [PATCH] Studio: fix ROCm visibility precedence + narrow ROCm child
 env

Two reviewer-flagged correctness bugs in the AMD GPU probe path.

1) ROCm visibility precedence was reversed. torch.cuda enumerates GPUs
   relative to HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm builds,
   but the probe's env-var lookup checked CUDA_VISIBLE_DEVICES first. With
   CUDA_VISIBLE_DEVICES=0,1 and HIP_VISIBLE_DEVICES=6,7 the probe returned
   [(0, ...), (1, ...)] when torch's view was actually [(6, ...), (7, ...)].
   The wrong physical IDs flowed downstream into CUDA_VISIBLE_DEVICES for
   the llama-server subprocess, pinning it to GPUs 0,1 instead of 6,7.

   Fix: branch on torch.version.hip. On ROCm, prefer HIP > ROCR > CUDA
   (matches torch's own ordering). On NVIDIA, use CUDA only -- ignoring
   any HIP/ROCR vars the parent happens to have set.

2) Child env narrowing only set CUDA_VISIBLE_DEVICES. On ROCm, llama-server
   honors HIP/ROCR; if the parent shell exported HIP_VISIBLE_DEVICES=4,5
   and the selector picked just GPU 4, the child still saw both because
   we never narrowed HIP/ROCR. Now we set all three on ROCm so the AMD
   subprocess actually sees the planned subset.

Both branches verified via temp/pr_simulation/sim_5172_rocm_precedence.py
(7/7 cases pass), including the reviewer's verbatim R5 case
(CVD=0,1 + HIP/ROCR=6,7).
---
 studio/backend/core/inference/llama_cpp.py | 41 +++++++++++++++-------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 5583086a83..82ca3f807b 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -609,19 +609,23 @@ class LlamaCppBackend:
                 return []
             if not hasattr(torch.cuda, "mem_get_info"):
                 return []
-            # torch.cuda enumerates GPUs RELATIVE to CUDA_VISIBLE_DEVICES
-            # (or HIP_VISIBLE_DEVICES on ROCm). Downstream we feed these
-            # IDs back into CUDA_VISIBLE_DEVICES for the llama-server
-            # subprocess, so we must translate visible ordinals back to
-            # physical indices first; otherwise launching with
-            # ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
+            # torch.cuda enumerates GPUs RELATIVE to the visibility mask.
+            # On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD
+            # ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES
+            # if HIP is unset). Downstream we feed these IDs back into the
+            # llama-server subprocess as CVD, so we must translate visible
+            # ordinals back to physical indices first; otherwise launching
+            # with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
             # ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs.
             physical_ids: Optional[list[int]] = None
-            cvd = (
-                os.environ.get("CUDA_VISIBLE_DEVICES")
-                or os.environ.get("HIP_VISIBLE_DEVICES")
-                or os.environ.get("ROCR_VISIBLE_DEVICES")
-            )
+            if getattr(torch.version, "hip", None) is not None:
+                cvd = (
+                    os.environ.get("HIP_VISIBLE_DEVICES")
+                    or os.environ.get("ROCR_VISIBLE_DEVICES")
+                    or os.environ.get("CUDA_VISIBLE_DEVICES")
+                )
+            else:
+                cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
             if cvd and cvd.strip():
                 try:
                     physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
@@ -1792,9 +1796,20 @@ class LlamaCppBackend:
                     f"{new_ld}:{existing_ld}" if existing_ld else new_ld
                 )
 
-            # Pin to selected GPU(s) via CUDA_VISIBLE_DEVICES
+            # Pin to selected GPU(s). On ROCm, llama-server (and any torch
+            # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES;
+            # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing
+            # the full HIP/ROCR set the parent inherited.
             if gpu_indices is not None:
-                env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_indices)
+                pinned = ",".join(str(i) for i in gpu_indices)
+                env["CUDA_VISIBLE_DEVICES"] = pinned
+                try:
+                    import torch as _torch
+                    if getattr(_torch.version, "hip", None) is not None:
+                        env["HIP_VISIBLE_DEVICES"] = pinned
+                        env["ROCR_VISIBLE_DEVICES"] = pinned
+                except Exception:
+                    pass
 
             self._stdout_lines = []
             self._process = subprocess.Popen(