diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ac5e74dbc0..3a3e4a3529 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, 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`` # 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..1860f2adc1 --- /dev/null +++ b/studio/backend/core/inference/diffusion_qwenimage21_vision.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Qwen-Image-2.1 vision SDPA gives NaN conditioning on gfx1151: use query-chunked eager there. + +UNSLOTH_QWEN_IMAGE_ROCM_VISION_EAGER: auto (gfx1151 only), 1 (any ROCm GPU), 0 (off). +""" + +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 + + # 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 + + 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, target, logger) -> bool: + if family != "qwen-image-2.1": + return False + 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"): + 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": + import torch + try: + from utils.hardware.hardware import _props_gfx_arch + props = torch.cuda.get_device_properties(torch.device(target.torch_device)) + if _props_gfx_arch(props) != "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 + # Correctness fix for native SDPA (Speed Off too); explicit non-SDPA backends are kept. + if visual.config._attn_implementation != "sdpa": + return False + # Never skip silently: that renders 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..056e113640 --- /dev/null +++ b/studio/backend/tests/test_diffusion_qwenimage21_vision.py @@ -0,0 +1,267 @@ +# 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_device import DiffusionDeviceTarget +from core.inference.diffusion_qwenimage21_vision import ( + configure_vision_attention, + _ENV, + _BACKEND, + _eager_vision_attention, +) + + +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.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", 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", target = _target(ordinal = 0), logger = Mock() + ) + visual.set_attn_implementation.assert_called_once_with(_BACKEND) + + +@pytest.mark.parametrize( + "family, backend, arch, override, expected", + [ + ("qwen-image", "rocm", "gfx1151", None, False), + ("qwen-image-2.1", "cpu", "gfx1151", "1", 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, 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, target = _target(backend), 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", target = _target(), logger = Mock() + ) + + +def test_cuda_not_rocm(runtime, monkeypatch): + monkeypatch.setenv(_ENV, "1") + assert not configure_vision_attention( + runtime[0], family = "qwen-image-2.1", target = _target("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.setenv(_ENV, "1") + 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 + 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", target = _target(ordinal = 0), logger = Mock() + ) + assert configure_vision_attention( + 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"), + 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", target = _target(), 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", target = _target(), 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", target = _target(), 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", target = _target(), 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", target = _target(), logger = Mock() + ) + assert caught.value.__cause__ is error