Studio: chunk large Qwen-Image-2.1 attention queries on ROCm

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

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))
 

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

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,