From 4d99249733ef02a4ca8c7350d45f6e85dafed880 Mon Sep 17 00:00:00 2001
From: LOGO127 <luozijian0924@gmail.com>
Date: Sat, 5 Sep 2026 06:48:17 +0800
Subject: [PATCH] Add ROCm architecture environment override

---
 bitsandbytes/cuda_specs.py         | 14 ++++++++
 docs/source/errors.mdx             | 10 ++++++
 tests/test_cuda_setup_evaluator.py | 54 +++++++++++++++++++++++++++++-
 3 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/bitsandbytes/cuda_specs.py b/bitsandbytes/cuda_specs.py
index 25ce3cd1e..e990ce3e3 100644
--- a/bitsandbytes/cuda_specs.py
+++ b/bitsandbytes/cuda_specs.py
@@ -1,6 +1,7 @@
 import dataclasses
 from functools import lru_cache
 import logging
+import os
 import platform
 import re
 import subprocess
@@ -84,6 +85,19 @@ def get_rocm_gpu_arch() -> str:
     logger = logging.getLogger(__name__)
     try:
         if torch.version.hip:
+            override = os.environ.get("BNB_ROCM_ARCH")
+            if override:
+                # gcnArchName may include feature flags, e.g.
+                # "gfx90a:sramecc+:xnack-". Only the architecture is used here.
+                arch = override.strip().lower().split(":", 1)[0]
+                if re.fullmatch(r"gfx[a-z\d]+", arch):
+                    return arch
+                logger.warning(
+                    "Ignoring invalid BNB_ROCM_ARCH=%r; expected a value such as 'gfx90a'.",
+                    override,
+                )
+                return "unknown"
+
             # On Windows, use hipinfo.exe; on Linux, use rocminfo
             if platform.system() == "Windows":
                 cmd = ["hipinfo.exe"]
diff --git a/docs/source/errors.mdx b/docs/source/errors.mdx
index 987488770..3ac24fcc4 100644
--- a/docs/source/errors.mdx
+++ b/docs/source/errors.mdx
@@ -45,3 +45,13 @@ If you see an error like `No compatible CUDA library found`, it means no compati
    set BNB_CUDA_VERSION=128
    ```
    The value must be digits only, e.g. `128` for CUDA 12.8 or `72` for ROCm 7.2.
+
+## ROCm architecture detection without `rocminfo`
+
+On Linux, bitsandbytes normally uses `rocminfo` to report the active GPU architecture. If that executable is unavailable in a restricted container, set the architecture before importing bitsandbytes:
+
+```bash
+export BNB_ROCM_ARCH=gfx90a
+```
+
+Use one concrete `gfx` architecture. This runtime override is separate from the `-DBNB_ROCM_ARCH` CMake option used when compiling bitsandbytes.
diff --git a/tests/test_cuda_setup_evaluator.py b/tests/test_cuda_setup_evaluator.py
index 56a52736e..b414f430e 100644
--- a/tests/test_cuda_setup_evaluator.py
+++ b/tests/test_cuda_setup_evaluator.py
@@ -5,7 +5,7 @@
 
 from bitsandbytes.cextension import get_cuda_bnb_library_path
 from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX
-from bitsandbytes.cuda_specs import CUDASpecs
+from bitsandbytes.cuda_specs import CUDASpecs, get_rocm_gpu_arch
 
 
 @pytest.fixture
@@ -133,3 +133,55 @@ def test_override_invalid_format(monkeypatch, cuda120_spec):
     monkeypatch.setenv("BNB_CUDA_VERSION", "12.4")
     with pytest.raises(RuntimeError, match="digits only"):
         get_cuda_bnb_library_path(cuda120_spec)
+
+
+def test_rocm_arch_override_skips_tool_probe(monkeypatch):
+    monkeypatch.setenv("BNB_ROCM_ARCH", "gfx90a:sramecc+:xnack-")
+
+    with (
+        patch("torch.version.hip", "7.0.0"),
+        patch("bitsandbytes.cuda_specs.subprocess.run") as run,
+    ):
+        assert get_rocm_gpu_arch() == "gfx90a"
+
+    run.assert_not_called()
+
+
+def test_rocm_arch_without_override_uses_tool_probe(monkeypatch):
+    monkeypatch.delenv("BNB_ROCM_ARCH", raising=False)
+
+    with (
+        patch("torch.version.hip", "7.0.0"),
+        patch("bitsandbytes.cuda_specs.platform.system", return_value="Linux"),
+        patch("bitsandbytes.cuda_specs.subprocess.run") as run,
+    ):
+        run.return_value.stdout = "  Name: gfx942\n"
+        assert get_rocm_gpu_arch() == "gfx942"
+
+    run.assert_called_once_with(["rocminfo"], capture_output=True, text=True)
+
+
+def test_invalid_rocm_arch_override_is_reported(monkeypatch, caplog):
+    monkeypatch.setenv("BNB_ROCM_ARCH", "90a")
+
+    with (
+        patch("torch.version.hip", "7.0.0"),
+        patch("bitsandbytes.cuda_specs.subprocess.run") as run,
+        caplog.at_level("WARNING"),
+    ):
+        assert get_rocm_gpu_arch() == "unknown"
+
+    run.assert_not_called()
+    assert "BNB_ROCM_ARCH" in caplog.text
+
+
+def test_non_rocm_ignores_arch_override(monkeypatch):
+    monkeypatch.setenv("BNB_ROCM_ARCH", "gfx90a")
+
+    with (
+        patch("torch.version.hip", None),
+        patch("bitsandbytes.cuda_specs.subprocess.run") as run,
+    ):
+        assert get_rocm_gpu_arch() == "unknown"
+
+    run.assert_not_called()