From 23e76141451a5e1ba937a4ceccfbe47c9b9289b4 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:28:53 -0300 Subject: [PATCH] Studio: extend the ROCm MIOpen cutoff and Qwen-Image-2.1 query chunking beyond gfx1151 --- .../core/inference/diffusion_qwenimage21_rocm.py | 12 +++++++----- .../tests/test_diffusion_qwenimage21_rocm.py | 4 ++++ studio/backend/tests/test_rocm_miopen_cutoff.py | 14 ++++++++++++-- studio/backend/utils/desktop_shell_env.py | 2 +- studio/backend/utils/hardware/hardware.py | 15 +++++++++++---- 5 files changed, 35 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/diffusion_qwenimage21_rocm.py b/studio/backend/core/inference/diffusion_qwenimage21_rocm.py index 14ac8a2659..9a5f73b936 100644 --- a/studio/backend/core/inference/diffusion_qwenimage21_rocm.py +++ b/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: diff --git a/studio/backend/tests/test_diffusion_qwenimage21_rocm.py b/studio/backend/tests/test_diffusion_qwenimage21_rocm.py index d07bd0b91c..fcde014582 100644 --- a/studio/backend/tests/test_diffusion_qwenimage21_rocm.py +++ b/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), ], diff --git a/studio/backend/tests/test_rocm_miopen_cutoff.py b/studio/backend/tests/test_rocm_miopen_cutoff.py index 539b3e30ff..6aa6bec315 100644 --- a/studio/backend/tests/test_rocm_miopen_cutoff.py +++ b/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), ], diff --git a/studio/backend/utils/desktop_shell_env.py b/studio/backend/utils/desktop_shell_env.py index f956061cb7..9c47d99bf6 100644 --- a/studio/backend/utils/desktop_shell_env.py +++ b/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", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index fd5c7f5496..6e775351bb 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1407,21 +1407,28 @@ def _xpu_device_name_or_placeholder(torch) -> str: return "" +# 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: