From cf68daba2eaa552563dec671086a723856106fe2 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 30 Sep 2026 02:25:10 -0300 Subject: [PATCH] Studio: chunk large Qwen-Image-2.1 attention queries on ROCm --- .../inference/diffusion_qwenimage21_rocm.py | 160 +++++++++++++++ .../backend/core/inference/diffusion_speed.py | 8 + .../tests/test_diffusion_qwenimage21_rocm.py | 192 ++++++++++++++++++ studio/backend/tests/test_diffusion_speed.py | 1 + 4 files changed, 361 insertions(+) create mode 100644 studio/backend/core/inference/diffusion_qwenimage21_rocm.py create mode 100644 studio/backend/tests/test_diffusion_qwenimage21_rocm.py diff --git a/studio/backend/core/inference/diffusion_qwenimage21_rocm.py b/studio/backend/core/inference/diffusion_qwenimage21_rocm.py new file mode 100644 index 000000000000..14ac8a2659bf --- /dev/null +++ b/studio/backend/core/inference/diffusion_qwenimage21_rocm.py @@ -0,0 +1,160 @@ +# 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. + +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. +""" + +from __future__ import annotations + +import functools +import os +from typing import Any + +QUERY_CHUNK_ENV = "UNSLOTH_QWEN_IMAGE_ROCM_QUERY_CHUNKING" +QUERY_CHUNK_SIZE = 512 +MIN_QUERY_TOKENS = 8192 + + +def _device_supported(target: Any) -> bool: + if getattr(target, "backend", None) != "rocm" or getattr(target, "device", None) != "cuda": + return False + setting = (os.environ.get(QUERY_CHUNK_ENV) or "auto").strip().lower() + if setting in ("0", "false", "off", "no"): + return False + import torch + + if not getattr(torch.version, "hip", None): + return False + if setting in ("1", "true", "on", "yes"): + return True + ordinal = getattr(target, "ordinal", None) + 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" + + +def _native_backend(processor: Any) -> bool: + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + backend = processor._attention_backend + if backend is None: + backend, _ = _AttentionBackendRegistry.get_active_backend() + return backend == "native" + + +def _can_chunk(processor, attn, hidden_states, attention_mask, segments, key_valid) -> bool: + import torch + return not ( + segments is not None + or attention_mask is not None + or key_valid is not None + or processor._parallel_config is not None + or hidden_states.device.type != "cuda" + or hidden_states.dtype != torch.bfloat16 + or hidden_states.shape[1] < MIN_QUERY_TOKENS + or attn.heads != 32 + or attn.inner_dim != 4096 + or torch.is_grad_enabled() + or torch.compiler.is_compiling() + or not _native_backend(processor) + ) + + +def _chunk_attention(query: Any, key: Any, value: Any, dispatch: Any, **kwargs: Any) -> Any: + import torch + return torch.cat( + [ + dispatch(query[:, start : start + QUERY_CHUNK_SIZE], key, value, **kwargs) + for start in range(0, query.shape[1], QUERY_CHUNK_SIZE) + ], + dim = 1, + ) + + +@functools.lru_cache(maxsize = 1) +def _processor_class() -> type: + import torch + from diffusers.models.transformers.transformer_qwenimage21 import ( + QwenImage21AttnProcessor, + _qwenimage21_prepare_qkv, + dispatch_attention_fn, + ) + + class ChunkedQwenImage21AttnProcessor(QwenImage21AttnProcessor): + def __call__( + self, + attn, + hidden_states, + attention_mask = None, + rotary_emb = None, + layer_cache = None, + kv_cache_mode = None, + cache_write_slice = None, + segments = None, + key_valid = None, + ): + if not _can_chunk(self, attn, hidden_states, attention_mask, segments, key_valid): + return super().__call__( + attn, + hidden_states, + attention_mask, + rotary_emb, + layer_cache, + kv_cache_mode, + cache_write_slice, + segments, + key_valid, + ) + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + ) + output = _chunk_attention( + query, + key, + value, + dispatch_attention_fn, + attn_mask = None, + dropout_p = 0.0, + backend = self._attention_backend, + parallel_config = self._parallel_config, + ) + output = output[:, :seq_len_q].flatten(2, 3).type_as(query) + return attn.to_out[1](attn.to_out[0](output)) + + return ChunkedQwenImage21AttnProcessor + + +def install( + pipe: Any, + target: Any, + logger: Any = None, +) -> bool: + """Replace only this pipeline's stock Qwen-Image-2.1 processors, before placement.""" + if not _device_supported(target): + return False + transformer = getattr(pipe, "transformer", None) + if type(transformer).__name__ != "QwenImage21Transformer2DModel": + return False + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21AttnProcessor + + cls = _processor_class() + installed = False + for module in transformer.modules(): + processor = getattr(module, "processor", None) + if type(processor) is cls: + installed = True + elif type(processor) is QwenImage21AttnProcessor: + replacement = cls() + replacement.__dict__.update(processor.__dict__) + module.set_processor(replacement) + installed = True + if installed and logger is not None: + logger.info( + "diffusion.qwenimage21: large ROCm decode queries are eligible for 512-row chunks" + ) + return installed diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 0f8bc2469500..0836c85ea1ff 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -445,6 +445,7 @@ def apply_speed_optims( "fused_qkv": False, "compiled": False, "compiled_dequant": False, + "rocm_query_chunks": False, "compiled_vae_decode": False, "cuda_graph": False, } @@ -454,6 +455,13 @@ def apply_speed_optims( if mode == SPEED_OFF: return applied + if getattr(target, "backend", None) == "rocm": + try: + from .diffusion_qwenimage21_rocm import install as install_rocm_query_chunks + applied["rocm_query_chunks"] = install_rocm_query_chunks(pipe, target, logger) + except Exception as exc: # noqa: BLE001 - keep stock attention on unsupported installs + _warn(logger, "qwen-image-2.1 ROCm query chunks", exc) + on_cuda = getattr(target, "device", None) == "cuda" family_allows_compile = bool(getattr(family, "supports_torch_compile", True)) diff --git a/studio/backend/tests/test_diffusion_qwenimage21_rocm.py b/studio/backend/tests/test_diffusion_qwenimage21_rocm.py new file mode 100644 index 000000000000..d07bd0b91cdb --- /dev/null +++ b/studio/backend/tests/test_diffusion_qwenimage21_rocm.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from types import SimpleNamespace + +import pytest + +from core.inference import diffusion_qwenimage21_rocm as chunks + + +@pytest.mark.parametrize( + "architecture,setting,expected", + [ + ("gfx1151", "auto", True), + ("gfx1151:sramecc-", "auto", True), + ("gfx1100", "auto", False), + ("gfx942", "auto", False), + ("gfx1100", "1", True), + ("gfx1151", "0", False), + ], +) +def test_device_scope(monkeypatch, architecture, setting, expected): + torch = pytest.importorskip("torch") + monkeypatch.setenv(chunks.QUERY_CHUNK_ENV, setting) + monkeypatch.setattr(torch.version, "hip", "7.13") + monkeypatch.setattr( + torch.cuda, "get_device_properties", lambda i: SimpleNamespace(gcnArchName = architecture) + ) + target = SimpleNamespace(backend = "rocm", device = "cuda", ordinal = 1) + assert chunks._device_supported(target) is expected + assert not chunks._device_supported(SimpleNamespace(backend = "cuda", device = "cuda")) + assert not chunks._device_supported(SimpleNamespace(backend = "cpu", device = "cpu")) + + +@pytest.mark.parametrize( + "reason", + [ + "eligible", + "cpu", + "fp32", + "small", + "heads", + "mask", + "segments", + "key_valid", + "parallel", + "grad", + "compile", + "backend", + ], +) +def test_runtime_scope(monkeypatch, reason): + torch = pytest.importorskip("torch") + hidden = SimpleNamespace( + device = SimpleNamespace(type = "cuda"), dtype = torch.bfloat16, shape = (1, 8192, 4096) + ) + processor = SimpleNamespace(_parallel_config = None) + attn = SimpleNamespace(heads = 32, inner_dim = 4096) + mask = segments = key_valid = None + if reason == "cpu": + hidden.device.type = "cpu" + if reason == "fp32": + hidden.dtype = torch.float32 + if reason == "small": + hidden.shape = (1, 4096, 4096) + if reason == "heads": + attn.heads = 16 + if reason == "mask": + mask = object() + if reason == "segments": + segments = [] + if reason == "key_valid": + key_valid = object() + if reason == "parallel": + processor._parallel_config = object() + monkeypatch.setattr(torch, "is_grad_enabled", lambda: reason == "grad") + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: reason == "compile") + monkeypatch.setattr(chunks, "_native_backend", lambda p: reason != "backend") + assert chunks._can_chunk(processor, attn, hidden, mask, segments, key_valid) is ( + reason == "eligible" + ) + + +def test_chunks_keep_all_keys_and_values(): + torch = pytest.importorskip("torch") + torch.manual_seed(12) + q = torch.randn(2, 1100, 2, 8) + k = torch.randn(2, 1200, 2, 8) + v = torch.randn_like(k) + calls = [] + + def dispatch(qpart, key, value, **kwargs): + calls.append((qpart.shape[1], key, value, kwargs)) + return torch.nn.functional.scaled_dot_product_attention( + qpart.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), **kwargs + ).transpose(1, 2) + + expected = dispatch(q, k, v, scale = 0.2) + calls.clear() + actual = chunks._chunk_attention(q, k, v, dispatch, scale = 0.2) + torch.testing.assert_close(actual, expected) + assert [c[0] for c in calls] == [512, 512, 76] + assert all(c[1] is k and c[2] is v and c[3] == {"scale": 0.2} for c in calls) + + +def test_processor_output_and_stock_fallback(monkeypatch): + torch = pytest.importorskip("torch") + module = pytest.importorskip("diffusers.models.transformers.transformer_qwenimage21") + torch.manual_seed(8) + attn = module.QwenImage21Attention(dim = 16, heads = 2, dim_head = 8) + stock = module.QwenImage21AttnProcessor() + patched = chunks._processor_class()() + hidden = torch.randn(1, 1100, 16) + with torch.inference_mode(): + expected = stock(attn, hidden) + monkeypatch.setattr(chunks, "_can_chunk", lambda *a: True) + torch.testing.assert_close(patched(attn, hidden), expected) + monkeypatch.setattr(chunks, "_can_chunk", lambda *a: False) + mask = torch.ones(1, 1, 1100, 1100, dtype = torch.bool).tril() + torch.testing.assert_close( + patched(attn, hidden, attention_mask = mask), stock(attn, hidden, attention_mask = mask) + ) + + +def test_install_is_local_and_idempotent(monkeypatch): + torch = pytest.importorskip("torch") + module = pytest.importorskip("diffusers.models.transformers.transformer_qwenimage21") + monkeypatch.setattr(chunks, "_device_supported", lambda target: True) + attn = module.QwenImage21Attention(dim = 16, heads = 2, dim_head = 8) + untouched = module.QwenImage21Attention(dim = 16, heads = 2, dim_head = 8) + transformer = type("QwenImage21Transformer2DModel", (), {"modules": lambda self: [attn]})() + pipe = SimpleNamespace(transformer = transformer) + stock = attn.processor + assert chunks.install(pipe, None) + installed = attn.processor + assert chunks.install(pipe, None) + assert attn.processor is installed + assert type(untouched.processor) is type(stock) + assert type(installed) is chunks._processor_class() + + +def test_native_backend_honors_context_and_explicit_choice(): + dispatch = pytest.importorskip("diffusers.models.attention_dispatch") + processor = SimpleNamespace(_attention_backend = None) + with dispatch.attention_backend("native"): + assert chunks._native_backend(processor) + with dispatch.attention_backend("_native_math"): + assert not chunks._native_backend(processor) + processor._attention_backend = "native" + assert chunks._native_backend(processor) + processor._attention_backend = "_native_math" + assert not chunks._native_backend(processor) + + +def test_cached_decode_preserves_prefix_and_projections(monkeypatch): + torch = pytest.importorskip("torch") + module = pytest.importorskip("diffusers.models.transformers.transformer_qwenimage21") + torch.manual_seed(9) + attn = module.QwenImage21Attention(dim = 16, heads = 2, dim_head = 8) + stock = module.QwenImage21AttnProcessor() + patched = chunks._processor_class()() + prefix = torch.randn(1, 7, 16) + target = torch.randn(1, 1050, 16) + cache = module.QwenImage21KVLayerCache() + with torch.inference_mode(): + # Text and condition-image prefix segments retain the stock causal rules. + joined = torch.cat([prefix, target], dim = 1) + before = stock( + attn, + joined, + layer_cache = cache, + kv_cache_mode = "extract", + cache_write_slice = slice(0, 7), + segments = [(0, 3, True), (3, 7, False)], + ) + other_cache = module.QwenImage21KVLayerCache() + monkeypatch.setattr(chunks, "_can_chunk", lambda *a: False) + after = patched( + attn, + joined, + layer_cache = other_cache, + kv_cache_mode = "extract", + cache_write_slice = slice(0, 7), + segments = [(0, 3, True), (3, 7, False)], + ) + torch.testing.assert_close(after, before, rtol = 0, atol = 0) + torch.testing.assert_close(other_cache.k, cache.k, rtol = 0, atol = 0) + expected = stock(attn, target, layer_cache = cache, kv_cache_mode = "cached") + monkeypatch.setattr(chunks, "_can_chunk", lambda *a: True) + actual = patched(attn, target, layer_cache = other_cache, kv_cache_mode = "cached") + torch.testing.assert_close(actual, expected) + assert other_cache.k.shape[1] == 7 diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9c94c416edac..99346c5e9417 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -434,6 +434,7 @@ def test_speed_off_applies_nothing(monkeypatch): "fused_qkv": False, "compiled": False, "compiled_dequant": False, + "rocm_query_chunks": False, "compiled_vae_decode": False, "fp16_accum": False, "cuda_graph": False,