From e790b934e2abce30ae0699396c50110c5e71d77f Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 30 Sep 2026 12:23:31 -0300 Subject: [PATCH 1/3] Studio: keep Qwen-Image-2.1 image conditioning finite on ROCm --- studio/backend/core/inference/diffusion.py | 4 + .../inference/diffusion_qwenimage21_vision.py | 89 +++++++ .../test_diffusion_qwenimage21_vision.py | 251 ++++++++++++++++++ 3 files changed, 344 insertions(+) create mode 100644 studio/backend/core/inference/diffusion_qwenimage21_vision.py create mode 100644 studio/backend/tests/test_diffusion_qwenimage21_vision.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ac5e74dbc0..73ef0a710a 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -6812,6 +6812,10 @@ def _refuse_pinned_precision() -> None: ) plan = refined_plan + from .diffusion_qwenimage21_vision import configure_vision_attention + + configure_vision_attention(pipe, family = fam.name, device = device, logger = logger) + # Persistent conditioning cache (UNSLOTH_DIFFUSION_COND_CACHE_DIR): repeated prompts skip the # text-encoder forward. After the TE quant so the key reflects the encoders that run; ``base`` # keys the companion repo. diff --git a/studio/backend/core/inference/diffusion_qwenimage21_vision.py b/studio/backend/core/inference/diffusion_qwenimage21_vision.py new file mode 100644 index 0000000000..5e9113388c --- /dev/null +++ b/studio/backend/core/inference/diffusion_qwenimage21_vision.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Keep Qwen-Image-2.1 reference-image encoding finite on measured ROCm GPUs. + +Defaults to gfx1151. UNSLOTH_QWEN_IMAGE_ROCM_VISION_EAGER=0 opts out; +=1 allows testing on other ROCm GPUs. Only the vision SDPA backend is changed. +Query chunks bound temporary attention memory without dropping any keys or values. +""" + +from __future__ import annotations + +import os + + +_ENV = "UNSLOTH_QWEN_IMAGE_ROCM_VISION_EAGER" +_BACKEND = "unsloth_qwen3_vl_eager" +_QUERY_ROWS = 512 + + +def _eager_vision_attention(module, query, key, value, attention_mask, **kwargs): + from transformers.models.qwen3_vl.modeling_qwen3_vl import eager_attention_forward + + # Keep the stock dropout RNG sequence for training; Studio uses this in eval. + if query.shape[-2] <= _QUERY_ROWS or kwargs.get("dropout", 0.0): + return eager_attention_forward(module, query, key, value, attention_mask, **kwargs) + import torch + + outputs = [] + for start in range(0, query.shape[-2], _QUERY_ROWS): + end = start + _QUERY_ROWS + mask = attention_mask + if mask is not None and mask.shape[-2] != 1: + mask = mask[..., start:end, :] + output, weights = eager_attention_forward( + module, query[..., start:end, :], key, value, mask, **kwargs + ) + del weights + outputs.append(output) + # Transformers' attention interface returns [batch, query, heads, head_dim]. + return torch.cat(outputs, dim = 1), None + + +def configure_vision_attention(pipe, *, family, device, logger) -> bool: + """Change only this encoder's vision backend, leaving text and denoising alone.""" + if family != "qwen-image-2.1": + return False + import torch + + if not (getattr(torch.version, "hip", None) or "rocm" in torch.__version__.lower()): + return False + target = torch.device(device) + if target.type != "cuda": + return False + override = (os.environ.get(_ENV) or "auto").strip().lower() + if override in ("0", "false", "off", "no"): + return False + if override not in ("auto", "1", "true", "on", "yes"): + raise ValueError(f"Invalid {_ENV}={override!r}; use auto, 1/true/on/yes or 0/false/off/no") + if override == "auto": + try: + from utils.hardware.hardware import _props_gfx_arch + if _props_gfx_arch(torch.cuda.get_device_properties(target)) != "gfx1151": + return False + except Exception: + return False + encoder = getattr(pipe, "text_encoder", None) + visual = getattr(getattr(encoder, "model", None), "visual", None) + if visual is None or type(visual).__name__ != "Qwen3VLVisionModel": + return False + # Respect explicitly selected non-SDPA backends. This fixes the native SDPA + # path, including Speed Off; it is a correctness fix, not a speed option. + if visual.config._attn_implementation != "sdpa": + return False + # Unlike a speed optimization, silently skipping this can produce blank images. + try: + from transformers import AttentionInterface + AttentionInterface.register(_BACKEND, _eager_vision_attention) + visual.set_attn_implementation(_BACKEND) + except Exception as exc: + raise RuntimeError( + "Could not enable Qwen-Image-2.1 ROCm vision attention. " + "Check the installed Transformers version. " + f"Set {_ENV}=0 to opt out; reference images may produce NaNs." + ) from exc + logger.info( + "diffusion.qwenimage21: using query-chunked eager vision attention on ROCm to avoid NaN image conditioning" + ) + return True diff --git a/studio/backend/tests/test_diffusion_qwenimage21_vision.py b/studio/backend/tests/test_diffusion_qwenimage21_vision.py new file mode 100644 index 0000000000..98119e4e89 --- /dev/null +++ b/studio/backend/tests/test_diffusion_qwenimage21_vision.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: AGPL-3.0-only +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers.models.qwen3_vl.modeling_qwen3_vl") + +from core.inference.diffusion_qwenimage21_vision import ( + configure_vision_attention, + _ENV, + _BACKEND, + _eager_vision_attention, +) + + +@pytest.fixture +def runtime(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "7.13") + monkeypatch.delenv(_ENV, raising = False) + monkeypatch.setattr( + torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = "gfx1151") + ) + visual = type("Qwen3VLVisionModel", (), {})() + visual.config = SimpleNamespace(_attn_implementation = "sdpa") + visual.set_attn_implementation = Mock( + side_effect = lambda value: setattr(visual.config, "_attn_implementation", value) + ) + text = SimpleNamespace(config = SimpleNamespace(_attn_implementation = "sdpa"), visual = visual) + pipe = SimpleNamespace(text_encoder = SimpleNamespace(model = text), transformer = object()) + return pipe, visual + + +def test_only_vision_changes_and_repeated_install_is_noop(runtime): + pipe, visual = runtime + assert configure_vision_attention(pipe, family = "qwen-image-2.1", device = "cuda:0", logger = Mock()) + assert visual.config._attn_implementation == _BACKEND + assert pipe.text_encoder.model.config._attn_implementation == "sdpa" + assert not configure_vision_attention( + pipe, family = "qwen-image-2.1", device = "cuda:0", logger = Mock() + ) + visual.set_attn_implementation.assert_called_once_with(_BACKEND) + + +@pytest.mark.parametrize( + "family, device, arch, override, expected", + [ + ("qwen-image", "cuda", "gfx1151", None, False), + ("qwen-image-2.1", "cpu", "gfx1151", "1", False), + ("qwen-image-2.1", "cuda", "gfx1100", None, False), + ("qwen-image-2.1", "cuda", "gfx1100", "1", True), + ("qwen-image-2.1", "cuda", "gfx1151:xnack-", None, True), + ("qwen-image-2.1", "cuda", "gfx1151", "0", False), + ], +) +def test_scope(runtime, monkeypatch, family, device, arch, override, expected): + monkeypatch.setattr( + torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = arch) + ) + if override is not None: + monkeypatch.setenv(_ENV, override) + assert ( + configure_vision_attention(runtime[0], family = family, device = device, logger = Mock()) + is expected + ) + + +@pytest.mark.parametrize("backend", ["eager", "flash_attention_2", "custom"]) +def test_preserves_other_attention_backends(runtime, backend): + runtime[1].config._attn_implementation = backend + assert not configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + + +def test_cuda_not_rocm(runtime, monkeypatch): + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch, "__version__", "2.11+cu130") + monkeypatch.setenv(_ENV, "1") + assert not configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + + +@pytest.mark.parametrize("lengths", [[16], [7, 9]]) +def test_real_vision_attention_parity_and_scope(monkeypatch, lengths): + from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLVisionConfig, + Qwen3VLTextConfig, + ) + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLModel + + config = Qwen3VLConfig( + vision_config = Qwen3VLVisionConfig( + depth = 1, + hidden_size = 32, + intermediate_size = 64, + num_heads = 4, + out_hidden_size = 32, + deepstack_visual_indexes = [], + ).to_dict(), + text_config = Qwen3VLTextConfig( + hidden_size = 32, + intermediate_size = 64, + num_hidden_layers = 1, + num_attention_heads = 4, + num_key_value_heads = 2, + head_dim = 8, + vocab_size = 64, + ).to_dict(), + ) + config._attn_implementation = "sdpa" + model = Qwen3VLModel(config).eval() + visual = model.visual + text_backend = model.language_model.config._attn_implementation + class_forward = type(visual.blocks[0].attn).forward + flags = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + pipe = SimpleNamespace(text_encoder = SimpleNamespace(model = model)) + x = torch.randn(sum(lengths), 32) + cu = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype = torch.int32) + positions = (torch.ones(sum(lengths), 8), torch.zeros(sum(lengths), 8)) + with torch.no_grad(): + reference = visual.blocks[0].attn(x, cu, position_embeddings = positions) + monkeypatch.setattr(torch.version, "hip", "7.13") + monkeypatch.setenv(_ENV, "1") + assert configure_vision_attention(pipe, family = "qwen-image-2.1", device = "cuda", logger = Mock()) + assert all(block.attn.config._attn_implementation == _BACKEND for block in visual.blocks) + assert model.language_model.config._attn_implementation == text_backend + assert type(visual.blocks[0].attn).forward is class_forward + assert flags == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + with torch.no_grad(): + actual = visual.blocks[0].attn(x, cu, position_embeddings = positions) + torch.testing.assert_close(actual, reference, rtol = 1e-5, atol = 1e-6) + + +@pytest.mark.parametrize("mask_kind", ["none", "broadcast", "per_query"]) +def test_chunked_attention_retains_all_keys_and_values(monkeypatch, mask_kind): + from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen + + module = SimpleNamespace(num_key_value_groups = 1, training = False) + q = torch.randn(1, 4, 1025, 8) + k = torch.randn(1, 4, 1103, 8) + v = torch.randn_like(k) + mask = None + if mask_kind != "none": + mask = torch.zeros(1, 1, 1 if mask_kind == "broadcast" else 1025, 1103) + mask[..., -3:] = -torch.inf + original = qwen.eager_attention_forward + expected, _ = original(module, q, k, v, mask, scaling = 0.25) + chunks = [] + + def tracked(module, query, key, value, attention_mask, **kwargs): + chunks.append(query.shape[-2]) + assert key is k and value is v + return original(module, query, key, value, attention_mask, **kwargs) + + monkeypatch.setattr(qwen, "eager_attention_forward", tracked) + actual, weights = _eager_vision_attention(module, q, k, v, mask, scaling = 0.25) + assert chunks == [512, 512, 1] + assert weights is None + torch.testing.assert_close(actual, expected, rtol = 1e-5, atol = 1e-6) + + +def test_dropout_keeps_stock_dispatch(monkeypatch): + from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen + + native = Mock(return_value = ("output", "weights")) + monkeypatch.setattr(qwen, "eager_attention_forward", native) + q = torch.empty(1, 4, 1025, 8) + assert _eager_vision_attention(None, q, q, q, None, dropout = 0.1) == ("output", "weights") + native.assert_called_once() + + +def test_architecture_gate_uses_selected_device(runtime, monkeypatch): + def properties(device): + return SimpleNamespace(gcnArchName = "gfx1151" if device.index == 1 else "gfx1100") + + probe = Mock(side_effect = properties) + monkeypatch.setattr(torch.cuda, "get_device_properties", probe) + assert not configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda:0", logger = Mock() + ) + assert configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda:1", logger = Mock() + ) + assert [call.args[0] for call in probe.call_args_list] == [ + torch.device("cuda:0"), + torch.device("cuda:1"), + ] + + +@pytest.mark.parametrize("setting", ["1", "true", "on", "yes", " TRUE "]) +def test_boolean_opt_in_on_other_architecture(runtime, monkeypatch, setting): + monkeypatch.setenv(_ENV, setting) + monkeypatch.setattr( + torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = "gfx1100") + ) + assert configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + + +@pytest.mark.parametrize("setting", ["0", "false", "off", "no", " OFF "]) +def test_boolean_opt_out(runtime, monkeypatch, setting): + monkeypatch.setenv(_ENV, setting) + assert not configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + + +@pytest.mark.parametrize("setting", ["", "auto", " AUTO "]) +def test_auto_override_keeps_architecture_gate(runtime, monkeypatch, setting): + monkeypatch.setenv(_ENV, setting) + assert configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + + +def test_invalid_override_is_actionable(runtime, monkeypatch): + monkeypatch.setenv(_ENV, "treu") + with pytest.raises(ValueError, match = _ENV): + configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + runtime[1].set_attn_implementation.assert_not_called() + + +@pytest.mark.parametrize("failure", ["register", "select"]) +def test_install_failure_is_actionable_and_preserves_cause(runtime, monkeypatch, failure): + from transformers import AttentionInterface + + error = RuntimeError("unsupported attention API") + if failure == "register": + monkeypatch.setattr(AttentionInterface, "register", Mock(side_effect = error)) + else: + runtime[1].set_attn_implementation.side_effect = error + with pytest.raises(RuntimeError, match = f"{_ENV}=0") as caught: + configure_vision_attention( + runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + ) + assert caught.value.__cause__ is error From 4db5bff326ad3277289096020835ffaa19662f3f Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:46:28 -0300 Subject: [PATCH 2/3] Gate the vision attention fix on the resolved device target --- studio/backend/core/inference/diffusion.py | 2 +- .../inference/diffusion_qwenimage21_vision.py | 13 ++-- .../test_diffusion_qwenimage21_vision.py | 64 ++++++++++++------- 3 files changed, 46 insertions(+), 33 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 73ef0a710a..3a3e4a3529 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -6814,7 +6814,7 @@ def _refuse_pinned_precision() -> None: from .diffusion_qwenimage21_vision import configure_vision_attention - configure_vision_attention(pipe, family = fam.name, device = device, logger = logger) + configure_vision_attention(pipe, family = fam.name, target = target, logger = logger) # Persistent conditioning cache (UNSLOTH_DIFFUSION_COND_CACHE_DIR): repeated prompts skip the # text-encoder forward. After the TE quant so the key reflects the encoders that run; ``base`` diff --git a/studio/backend/core/inference/diffusion_qwenimage21_vision.py b/studio/backend/core/inference/diffusion_qwenimage21_vision.py index 5e9113388c..38c365b307 100644 --- a/studio/backend/core/inference/diffusion_qwenimage21_vision.py +++ b/studio/backend/core/inference/diffusion_qwenimage21_vision.py @@ -41,16 +41,11 @@ def _eager_vision_attention(module, query, key, value, attention_mask, **kwargs) return torch.cat(outputs, dim = 1), None -def configure_vision_attention(pipe, *, family, device, logger) -> bool: +def configure_vision_attention(pipe, *, family, target, logger) -> bool: """Change only this encoder's vision backend, leaving text and denoising alone.""" if family != "qwen-image-2.1": return False - import torch - - if not (getattr(torch.version, "hip", None) or "rocm" in torch.__version__.lower()): - return False - target = torch.device(device) - if target.type != "cuda": + if getattr(target, "backend", None) != "rocm" or getattr(target, "device", None) != "cuda": return False override = (os.environ.get(_ENV) or "auto").strip().lower() if override in ("0", "false", "off", "no"): @@ -58,9 +53,11 @@ def configure_vision_attention(pipe, *, family, device, logger) -> bool: if override not in ("auto", "1", "true", "on", "yes"): raise ValueError(f"Invalid {_ENV}={override!r}; use auto, 1/true/on/yes or 0/false/off/no") if override == "auto": + import torch try: from utils.hardware.hardware import _props_gfx_arch - if _props_gfx_arch(torch.cuda.get_device_properties(target)) != "gfx1151": + props = torch.cuda.get_device_properties(torch.device(target.torch_device)) + if _props_gfx_arch(props) != "gfx1151": return False except Exception: return False diff --git a/studio/backend/tests/test_diffusion_qwenimage21_vision.py b/studio/backend/tests/test_diffusion_qwenimage21_vision.py index 98119e4e89..056e113640 100644 --- a/studio/backend/tests/test_diffusion_qwenimage21_vision.py +++ b/studio/backend/tests/test_diffusion_qwenimage21_vision.py @@ -7,6 +7,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("transformers.models.qwen3_vl.modeling_qwen3_vl") +from core.inference.diffusion_device import DiffusionDeviceTarget from core.inference.diffusion_qwenimage21_vision import ( configure_vision_attention, _ENV, @@ -15,9 +16,21 @@ ) +def _target(backend = "rocm", ordinal = None): + return DiffusionDeviceTarget( + device = "cpu" if backend == "cpu" else "cuda", + dtype = torch.bfloat16, + backend = backend, + vendor = None, + supports_model_cpu_offload = True, + supports_default_torch_compile = False, + supports_pinned_transfer = True, + ordinal = ordinal, + ) + + @pytest.fixture def runtime(monkeypatch): - monkeypatch.setattr(torch.version, "hip", "7.13") monkeypatch.delenv(_ENV, raising = False) monkeypatch.setattr( torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = "gfx1151") @@ -34,34 +47,38 @@ def runtime(monkeypatch): def test_only_vision_changes_and_repeated_install_is_noop(runtime): pipe, visual = runtime - assert configure_vision_attention(pipe, family = "qwen-image-2.1", device = "cuda:0", logger = Mock()) + assert configure_vision_attention( + pipe, family = "qwen-image-2.1", target = _target(ordinal = 0), logger = Mock() + ) assert visual.config._attn_implementation == _BACKEND assert pipe.text_encoder.model.config._attn_implementation == "sdpa" assert not configure_vision_attention( - pipe, family = "qwen-image-2.1", device = "cuda:0", logger = Mock() + pipe, family = "qwen-image-2.1", target = _target(ordinal = 0), logger = Mock() ) visual.set_attn_implementation.assert_called_once_with(_BACKEND) @pytest.mark.parametrize( - "family, device, arch, override, expected", + "family, backend, arch, override, expected", [ - ("qwen-image", "cuda", "gfx1151", None, False), + ("qwen-image", "rocm", "gfx1151", None, False), ("qwen-image-2.1", "cpu", "gfx1151", "1", False), - ("qwen-image-2.1", "cuda", "gfx1100", None, False), - ("qwen-image-2.1", "cuda", "gfx1100", "1", True), - ("qwen-image-2.1", "cuda", "gfx1151:xnack-", None, True), - ("qwen-image-2.1", "cuda", "gfx1151", "0", False), + ("qwen-image-2.1", "rocm", "gfx1100", None, False), + ("qwen-image-2.1", "rocm", "gfx1100", "1", True), + ("qwen-image-2.1", "rocm", "gfx1151:xnack-", None, True), + ("qwen-image-2.1", "rocm", "gfx1151", "0", False), ], ) -def test_scope(runtime, monkeypatch, family, device, arch, override, expected): +def test_scope(runtime, monkeypatch, family, backend, arch, override, expected): monkeypatch.setattr( torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = arch) ) if override is not None: monkeypatch.setenv(_ENV, override) assert ( - configure_vision_attention(runtime[0], family = family, device = device, logger = Mock()) + configure_vision_attention( + runtime[0], family = family, target = _target(backend), logger = Mock() + ) is expected ) @@ -70,16 +87,14 @@ def test_scope(runtime, monkeypatch, family, device, arch, override, expected): def test_preserves_other_attention_backends(runtime, backend): runtime[1].config._attn_implementation = backend assert not configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) def test_cuda_not_rocm(runtime, monkeypatch): - monkeypatch.setattr(torch.version, "hip", None) - monkeypatch.setattr(torch, "__version__", "2.11+cu130") monkeypatch.setenv(_ENV, "1") assert not configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target("cuda"), logger = Mock() ) @@ -127,9 +142,10 @@ def test_real_vision_attention_parity_and_scope(monkeypatch, lengths): positions = (torch.ones(sum(lengths), 8), torch.zeros(sum(lengths), 8)) with torch.no_grad(): reference = visual.blocks[0].attn(x, cu, position_embeddings = positions) - monkeypatch.setattr(torch.version, "hip", "7.13") monkeypatch.setenv(_ENV, "1") - assert configure_vision_attention(pipe, family = "qwen-image-2.1", device = "cuda", logger = Mock()) + assert configure_vision_attention( + pipe, family = "qwen-image-2.1", target = _target(), logger = Mock() + ) assert all(block.attn.config._attn_implementation == _BACKEND for block in visual.blocks) assert model.language_model.config._attn_implementation == text_backend assert type(visual.blocks[0].attn).forward is class_forward @@ -188,10 +204,10 @@ def properties(device): probe = Mock(side_effect = properties) monkeypatch.setattr(torch.cuda, "get_device_properties", probe) assert not configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda:0", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(ordinal = 0), logger = Mock() ) assert configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda:1", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(ordinal = 1), logger = Mock() ) assert [call.args[0] for call in probe.call_args_list] == [ torch.device("cuda:0"), @@ -206,7 +222,7 @@ def test_boolean_opt_in_on_other_architecture(runtime, monkeypatch, setting): torch.cuda, "get_device_properties", lambda _: SimpleNamespace(gcnArchName = "gfx1100") ) assert configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) @@ -214,7 +230,7 @@ def test_boolean_opt_in_on_other_architecture(runtime, monkeypatch, setting): def test_boolean_opt_out(runtime, monkeypatch, setting): monkeypatch.setenv(_ENV, setting) assert not configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) @@ -222,7 +238,7 @@ def test_boolean_opt_out(runtime, monkeypatch, setting): def test_auto_override_keeps_architecture_gate(runtime, monkeypatch, setting): monkeypatch.setenv(_ENV, setting) assert configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) @@ -230,7 +246,7 @@ def test_invalid_override_is_actionable(runtime, monkeypatch): monkeypatch.setenv(_ENV, "treu") with pytest.raises(ValueError, match = _ENV): configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) runtime[1].set_attn_implementation.assert_not_called() @@ -246,6 +262,6 @@ def test_install_failure_is_actionable_and_preserves_cause(runtime, monkeypatch, runtime[1].set_attn_implementation.side_effect = error with pytest.raises(RuntimeError, match = f"{_ENV}=0") as caught: configure_vision_attention( - runtime[0], family = "qwen-image-2.1", device = "cuda", logger = Mock() + runtime[0], family = "qwen-image-2.1", target = _target(), logger = Mock() ) assert caught.value.__cause__ is error From 8e4ce5bcb675d6438bb3f04a3e50d11b9fa2dac0 Mon Sep 17 00:00:00 2001 From: Daniel Han <23090290+danielhanchen@users.noreply.github.com> Date: Thu, 1 Oct 2026 02:12:25 -0700 Subject: [PATCH 3/3] Tighten Qwen-Image-2.1 vision attention comments --- .../core/inference/diffusion_qwenimage21_vision.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/inference/diffusion_qwenimage21_vision.py b/studio/backend/core/inference/diffusion_qwenimage21_vision.py index 38c365b307..1860f2adc1 100644 --- a/studio/backend/core/inference/diffusion_qwenimage21_vision.py +++ b/studio/backend/core/inference/diffusion_qwenimage21_vision.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 -"""Keep Qwen-Image-2.1 reference-image encoding finite on measured ROCm GPUs. +"""Qwen-Image-2.1 vision SDPA gives NaN conditioning on gfx1151: use query-chunked eager there. -Defaults to gfx1151. UNSLOTH_QWEN_IMAGE_ROCM_VISION_EAGER=0 opts out; -=1 allows testing on other ROCm GPUs. Only the vision SDPA backend is changed. -Query chunks bound temporary attention memory without dropping any keys or values. +UNSLOTH_QWEN_IMAGE_ROCM_VISION_EAGER: auto (gfx1151 only), 1 (any ROCm GPU), 0 (off). """ from __future__ import annotations @@ -21,7 +19,7 @@ def _eager_vision_attention(module, query, key, value, attention_mask, **kwargs): from transformers.models.qwen3_vl.modeling_qwen3_vl import eager_attention_forward - # Keep the stock dropout RNG sequence for training; Studio uses this in eval. + # Dropout: keep the stock RNG sequence. if query.shape[-2] <= _QUERY_ROWS or kwargs.get("dropout", 0.0): return eager_attention_forward(module, query, key, value, attention_mask, **kwargs) import torch @@ -42,7 +40,6 @@ def _eager_vision_attention(module, query, key, value, attention_mask, **kwargs) def configure_vision_attention(pipe, *, family, target, logger) -> bool: - """Change only this encoder's vision backend, leaving text and denoising alone.""" if family != "qwen-image-2.1": return False if getattr(target, "backend", None) != "rocm" or getattr(target, "device", None) != "cuda": @@ -65,11 +62,10 @@ def configure_vision_attention(pipe, *, family, target, logger) -> bool: visual = getattr(getattr(encoder, "model", None), "visual", None) if visual is None or type(visual).__name__ != "Qwen3VLVisionModel": return False - # Respect explicitly selected non-SDPA backends. This fixes the native SDPA - # path, including Speed Off; it is a correctness fix, not a speed option. + # Correctness fix for native SDPA (Speed Off too); explicit non-SDPA backends are kept. if visual.config._attn_implementation != "sdpa": return False - # Unlike a speed optimization, silently skipping this can produce blank images. + # Never skip silently: that renders blank images. try: from transformers import AttentionInterface AttentionInterface.register(_BACKEND, _eager_vision_attention)