Studio: extend the ROCm MIOpen cutoff and Qwen-Image-2.1 query chunking beyond gfx1151

studio/backend/core/inference/diffusion_qwenimage21_rocm.py

@@ -1,11 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Bound large Qwen-Image-2.1 query batches on measured ROCm GPUs.
+"""Speed up large Qwen-Image-2.1 ROCm attention with query chunks.
 
-On gfx1151, native flash attention slows sharply above 8192 image tokens. Query
-rows are independent for unmasked, noncausal attention, so each chunk can retain
-all keys and values. Prefill, masks and explicit alternative backends stay stock.
+Unmasked, noncausal queries are independent; each chunk keeps all keys and values.
 """
 
 from __future__ import annotations
@@ -17,6 +15,8 @@
 QUERY_CHUNK_ENV = "UNSLOTH_QWEN_IMAGE_ROCM_QUERY_CHUNKING"
 QUERY_CHUNK_SIZE = 512
 MIN_QUERY_TOKENS = 8192
+# gfx1151 and gfx1201 are measured; gfx1200 runs gfx1201's AOTriton kernels and tuning.
+CHUNKED_ARCHES = frozenset({"gfx1151", "gfx1200", "gfx1201"})
 
 
 def _device_supported(target: Any) -> bool:
@@ -35,7 +35,9 @@ def _device_supported(target: Any) -> bool:
     props = torch.cuda.get_device_properties(
         torch.cuda.current_device() if ordinal is None else ordinal
     )
-    return str(getattr(props, "gcnArchName", "")).split(":", 1)[0] == "gfx1151"
+    from utils.hardware.hardware import _props_gfx_arch
+
+    return _props_gfx_arch(props) in CHUNKED_ARCHES
 
 
 def _native_backend(processor: Any) -> bool:

studio/backend/tests/test_diffusion_qwenimage21_rocm.py

@@ -13,8 +13,12 @@
     [
         ("gfx1151", "auto", True),
         ("gfx1151:sramecc-", "auto", True),
+        ("gfx1201", "auto", True),
+        ("gfx1200", "auto", True),
         ("gfx1100", "auto", False),
+        ("gfx1150", "auto", False),
         ("gfx942", "auto", False),
+        ("gfx1201", "0", False),
         ("gfx1100", "1", True),
         ("gfx1151", "0", False),
     ],

studio/backend/tests/test_rocm_miopen_cutoff.py

@@ -19,8 +19,18 @@ def isolated_environment(monkeypatch):
     [
         (["gfx1151"], "1"),
         (["gfx1151:xnack-", "gfx1151"], "1"),
-        (["gfx1100"], None),
-        (["gfx1151", "gfx1100"], None),
+        (["gfx1100"], "1"),
+        (["gfx1102"], "1"),
+        (["gfx1150"], "1"),
+        (["gfx1200"], "1"),
+        (["gfx1201"], "1"),
+        (["gfx1151", "gfx1100", "gfx1201"], "1"),
+        (["gfx1030"], None),
+        (["gfx90a"], None),
+        (["gfx942:sramecc+:xnack-"], None),
+        (["gfx950"], None),
+        (["gfx1250"], None),
+        (["gfx1201", "gfx942"], None),
         (["gfx1151", ""], None),
         ([], None),
     ],

studio/backend/utils/desktop_shell_env.py

@@ -48,7 +48,7 @@
     "MIOPEN_USER_DB_PATH",
     "MIOPEN_CUSTOM_CACHE_DIR",
     "MIOPEN_FIND_MODE",
-    # A shell opt-out (0) or opt-in must beat hardware.py's gfx1151 default.
+    # Preserve shell overrides, including an explicit opt-out (0).
     "MIOPEN_SEARCH_CUTOFF",
     "HIP_VISIBLE_DEVICES",
     "ROCR_VISIBLE_DEVICES",

studio/backend/utils/hardware/hardware.py

@@ -1407,21 +1407,28 @@ def _xpu_device_name_or_placeholder(torch) -> str:
         return "<unavailable>"
 
 
+# RDNA 3/3.5/4 only; gfx1250 is Instinct, so avoid a broad gfx12 prefix.
+_MIOPEN_SEARCH_CUTOFF_ARCH_PREFIXES = ("gfx110", "gfx115", "gfx120")
+
+
 def _configure_rocm_miopen(torch) -> None:
     if "MIOPEN_SEARCH_CUTOFF" in os.environ:
         return
     try:
         count = torch.cuda.device_count()
-        # Process-wide in MIOpen: any untested arch on a mixed host keeps the default.
-        if not count or any(
-            _props_gfx_arch(torch.cuda.get_device_properties(i)) != "gfx1151" for i in range(count)
+        # MIOpen's cutoff is process-wide, so every visible GPU must qualify.
+        if not count or not all(
+            _props_gfx_arch(torch.cuda.get_device_properties(i)).startswith(
+                _MIOPEN_SEARCH_CUTOFF_ARCH_PREFIXES
+            )
+            for i in range(count)
         ):
             return
     except Exception as exc:
         logger.debug("MIOpen search cutoff device probe failed: %s", exc)
         return
     os.environ.setdefault("MIOPEN_SEARCH_CUTOFF", "1")
-    logger.info("ROCm gfx1151: enabled MIOpen search cutoff (MIOPEN_SEARCH_CUTOFF=1)")
+    logger.info("ROCm RDNA: enabled MIOpen search cutoff (MIOPEN_SEARCH_CUTOFF=1)")
 
 
 def _detect_hardware_locked() -> DeviceType: