Optimize Qwen-Image 2.1 attention and decode graph reuse (#7790)

Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>

tests/diffusion/models/qwen_image_21/test_batch_compatibility.py

@@ -82,3 +82,40 @@ def test_text_input_forms_share_compatibility_key(preprocess, prompt):
     empty_list_request = preprocess(make_request("empty-list", []))
     assert processed.batch_compatibility_key is not None
     assert processed.batch_compatibility_key == empty_list_request.batch_compatibility_key
+
+
+@pytest.mark.parametrize("image_type", ["numpy", "tensor"])
+def test_preprocess_accepts_array_images_without_truth_testing(preprocess, image_type):
+    import numpy as np
+    import torch
+
+    pixels = np.full((32, 32, 4), 255, dtype=np.uint8)
+    image = torch.from_numpy(pixels).permute(2, 0, 1) if image_type == "tensor" else pixels
+    request = make_request("array", [])
+    request.prompt["multi_modal_data"]["image"] = image
+    processed = preprocess(request)
+    assert processed.batch_compatibility_key is not None
+    assert processed.allow_mixed_step_phases is False
+
+
+def test_late_prefill_waits_until_all_decode_requests_finish(preprocess):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    scheduler = StepScheduler()
+    scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+    for request_id in ("a", "b"):
+        scheduler.add_request(preprocess(make_request(request_id, [])))
+    wave = scheduler.schedule()
+    assert wave.scheduled_request_ids == ["a", "b"]
+    scheduler.update_from_output(
+        wave,
+        BatchRunnerOutput.from_list(
+            [RunnerOutput(request_id=request_id, step_index=1, finished=False) for request_id in ("a", "b")]
+        ),
+    )
+    scheduler.add_request(preprocess(make_request("late", [])))
+    assert scheduler.schedule().scheduled_request_ids == ["a", "b"]
+    scheduler.finish_requests("a", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert scheduler.schedule().scheduled_request_ids == ["b"]
+    scheduler.finish_requests("b", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert scheduler.schedule().scheduled_request_ids == ["late"]

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import weakref
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class DecodeModel(torch.nn.Module):
+    in_channels = 4
+
+    def _decode_graph_forward(self, entry):
+        return (
+            entry.hidden.sin()
+            + entry.k[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.v[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.freqs[None, :, :]
+        )
+
+
+@pytest.fixture
+def graph_case(request, monkeypatch):
+    device = "cuda" if request.node.get_closest_marker("cuda") else "cpu"
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2)
+    monkeypatch.setattr(manager, "eligible", lambda: True)
+    monkeypatch.setattr(manager, "_backend_name", lambda: "test")
+    cleanup = []
+    monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True))
+    return manager, device, cleanup
+
+
+def make_cache(device, value):
+    return [
+        {
+            "cond": {
+                "key": torch.full((1, 2, 1, 4), value, device=device),
+                "value": torch.full((1, 2, 1, 4), value + 1, device=device),
+            }
+        }
+    ]
+
+
+def register(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    manager.register_prefill(
+        kv_cache=cache,
+        cache_branch="cond",
+        prefix_len=2,
+        img_shapes=[layout],
+        target_freqs=torch.full((4, 4), float(layout[1]), device=device),
+        joint_key_valid=None,
+        dtype=torch.float32,
+    )
+
+
+def decode(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    return manager.try_decode(
+        hidden_states=torch.ones(1, 4, 4, device=device),
+        timestep=torch.ones(1, device=device),
+        kv_cache=cache,
+        cache_branch="cond",
+        img_shapes=[layout],
+        img_mask=torch.tensor([[False, False, True]], device=device),
+        encoder_hidden_states_mask=None,
+    )
+
+
+@pytest.mark.cpu
+def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+    allocated = []
+
+    def fail_allocation(**kwargs):
+        entry = constructor(**kwargs)
+        allocated.append(weakref.ref(entry.k[0]))
+        raise torch.OutOfMemoryError("injected allocation OOM")
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation)
+    register(manager, cache)
+    register(manager, cache)
+    assert len(allocated) == 1
+    assert allocated[0]() is None
+    assert list(manager.entries.values()) == [None]
+    assert len(cleanup) == 1
+    assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4))
+
+
+@pytest.mark.cpu
+def test_evict_before_allocating_replacement(graph_case, monkeypatch):
+    manager, device, _ = graph_case
+    manager.max_entries = 1
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0])
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+
+    def allocate_after_release(**kwargs):
+        assert old_buffer() is None
+        return constructor(**kwargs)
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release)
+    register(manager, cache, (1, 1, 4))
+    assert len(manager.entries) == 1
+
+
+@pytest.mark.cpu
+def test_registration_preserves_request_cache_and_separates_layouts(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    original = first[0]["cond"]["key"]
+    register(manager, first)
+    register(manager, second)
+    assert first[0]["cond"]["key"] is original
+    assert torch.equal(original, torch.ones_like(original))
+    register(manager, second, (1, 1, 4))
+    assert len(manager.entries) == 2
+    entries = list(manager.entries.values())
+    assert not torch.equal(entries[0].freqs, entries[1].freqs)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+@pytest.mark.parametrize("failure_phase", ["warmup", "capture"])
+def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0])
+    del entry
+
+    original = manager.model._decode_graph_forward
+
+    def fail_capture(entry):
+        output = original(entry)
+        if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing():
+            raise torch.OutOfMemoryError("injected capture OOM")
+        return output
+
+    monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture)
+    assert decode(manager, cache) is None
+    assert entry_ref() is None and buffer_ref() is None
+    assert list(manager.entries.values()) == [None]
+    register(manager, cache)
+    assert decode(manager, cache) is None
+    assert len(cleanup) == 1
+    assert bool((cache[0]["cond"]["key"] == 1).all())
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_replay_uses_active_request_and_layout(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    with torch.inference_mode():
+        register(manager, first)
+        register(manager, second)
+        for cache in (first, second, first):
+            output = decode(manager, cache)
+            expected = (
+                torch.ones_like(output).sin() + cache[0]["cond"]["key"].mean() + cache[0]["cond"]["value"].mean() + 2
+            )
+            torch.testing.assert_close(output, expected)
+        original = decode(manager, first).clone()
+        register(manager, second, (1, 1, 4))
+        other = decode(manager, second, (1, 1, 4))
+        torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1)
+        torch.testing.assert_close(decode(manager, first), original)
+        assert all(entry.captures == 1 for entry in manager.entries.values())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("inference", [False, True])
+def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference):
+    manager, device, _ = graph_case
+    with torch.inference_mode(inference):
+        cache = make_cache(device, 1.0) + make_cache(device, 2.0)
+        register(manager, cache)
+        entry = next(iter(manager.entries.values()))
+        entry.refresh_prefix(cache)
+        versions = [tensor._version for tensor in entry.k + entry.v]
+        entry.refresh_prefix(cache)
+        assert versions == [tensor._version for tensor in entry.k + entry.v]
+        cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0)
+        entry.refresh_prefix(cache)
+        assert bool((entry.v[1] == 7).all())
+        source = weakref.ref(cache[0]["cond"]["key"])
+        del cache
+        assert source() is None
+
+
+@pytest.mark.cpu
+def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry.refresh_prefix(cache)
+    cache[0]["cond"]["value"].add_(3)
+    entry.refresh_prefix(cache)
+    assert bool((entry.v[0] == 5).all())
+    register(manager, cache)
+    assert not entry.prefix_sources
+
+
+@pytest.mark.cpu
+def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    for valid in (True, False):
+        manager.register_prefill(
+            kv_cache=cache,
+            cache_branch="cond",
+            prefix_len=2,
+            img_shapes=[(1, 2, 2)],
+            target_freqs=torch.ones(4, 4),
+            joint_key_valid=torch.full((1, 6), valid),
+            dtype=torch.float32,
+        )
+        assert list(manager.entries.values()) == [entry]
+        assert entry.attn_metadata is None
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case):
+    manager, device, _ = graph_case
+    first = make_cache(device, 1.0)
+    second = make_cache(device, 4.0)
+    merged = [
+        {"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}}
+    ]
+    with torch.inference_mode():
+        register(manager, first)
+        single = decode(manager, first).clone()
+        register(manager, merged)
+        kwargs = dict(
+            hidden_states=torch.ones(2, 4, 4, device=device),
+            timestep=torch.ones(2, device=device),
+            kv_cache=merged,
+            cache_branch="cond",
+            img_shapes=[(1, 2, 2)],
+            img_mask=torch.tensor([[False, False, True]] * 2, device=device),
+            encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
+        )
+        result = manager.try_decode(**kwargs)
+        assert result is not None and result.shape == (2, 4, 4)
+        expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None]
+        torch.testing.assert_close(result, expected)
+        kwargs["kv_cache"] = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}]
+        torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0))
+        kwargs["encoder_hidden_states_mask"][1, 0] = False
+        assert manager.try_decode(**kwargs) is None
+        torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_transformer_numerics.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image_21 import qwen_image_21_transformer as transformer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "device",
+    [pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
+)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_qk_norm_matches_diffusers_rounding(monkeypatch, device, dtype):
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state.get_tp_group", lambda: SimpleNamespace(rank_in_group=0, world_size=1)
+    )
+    monkeypatch.setattr(transformer, "Attention", lambda **kwargs: torch.nn.Identity())
+    attention = transformer.QwenImage21Attention(dim=256, heads=2, dim_head=128).to(device=device, dtype=dtype)
+    values = torch.randn(2, 17, 2, 128, generator=torch.Generator().manual_seed(42)).to(device=device, dtype=dtype)
+    for norm in (attention.norm_q, attention.norm_k):
+        with torch.no_grad():
+            norm.weight.copy_(torch.linspace(0.5, 1.5, 128, device=device, dtype=dtype))
+        variance = values.float().square().mean(-1, keepdim=True)
+        expected = (values * torch.rsqrt(variance + norm.eps)).to(dtype) * norm.weight
+        torch.testing.assert_close(norm(values), expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [pytest.param("meta", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
+)
+def test_rope_frequencies_do_not_depend_on_loader_device(device):
+    reference = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
+    with torch.device(device):
+        loaded = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
+    for expected, actual in zip(reference.freqs, loaded.freqs):
+        assert actual.device.type == "cpu"
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_scheduler_matches_diffusers_on_latent_device():
+    import numpy as np
+    from diffusers import FlowMatchEulerDiscreteScheduler
+
+    from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import QwenImage21Pipeline
+
+    pipeline = object.__new__(QwenImage21Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cuda", 0)
+    pipeline.scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
+    timesteps, _ = pipeline.prepare_timesteps(50, None, 4096)
+    reference = FlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
+    reference.set_timesteps(sigmas=np.linspace(1.0, 1 / 50, 50), device=pipeline.device, mu=1.15)
+    sample = torch.randn(1, 16, 64, generator=torch.Generator().manual_seed(42)).to(
+        device=pipeline.device, dtype=torch.bfloat16
+    )
+    noise = sample.cos()
+    actual = pipeline.scheduler.step(noise, timesteps[0], sample, return_dict=False)[0]
+    expected = reference.step(noise, reference.timesteps[0], sample, return_dict=False)[0]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert pipeline.scheduler.sigmas.device == sample.device
+    assert timesteps.device == sample.device
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_timestep_frequencies_match_cpu_initialization():
+    reference = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
+    with torch.device("cuda"):
+        loaded = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
+    assert loaded.freqs.device.type == "cuda"
+    torch.testing.assert_close(loaded.freqs.cpu(), reference.freqs, rtol=0, atol=0)
+    timesteps = torch.tensor([0.0, 0.984375, 1.0], device="cuda")
+    torch.testing.assert_close(loaded(timesteps), reference.to("cuda")(timesteps), rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_vae.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import math
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21 import (
+    DistributedAutoencoderKLQwenImage21,
+)
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def single_thread():
+    previous = torch.get_num_threads()
+    torch.set_num_threads(1)
+    yield
+    torch.set_num_threads(previous)
+
+
+def _make_vae(cls=AutoencoderKLQwenImage21, **overrides):
+    """Toy VAE with the 2.1 residual architecture; 4x spatial compression."""
+    kwargs = dict(
+        base_dim=16,
+        decoder_base_dim=16,
+        z_dim=4,
+        dim_mult=[1, 2, 2],
+        num_res_blocks=1,
+        temperal_downsample=[False, False],
+        is_residual=True,
+        in_channels=4,
+        out_channels=4,
+        latents_mean=[0.0] * 4,
+        latents_std=[1.0] * 4,
+        scale_factor_spatial=None,
+        scale_factor_temporal=None,
+    )
+    kwargs.update(overrides)
+    return cls(**kwargs)
+
+
+def test_spatial_compression_ratio_derived_from_architecture():
+    # One 2x downsampling stage per `temperal_downsample` entry.
+    vae = _make_vae()
+    assert vae.spatial_compression_ratio == 4
+
+    # The real Qwen-Image-2.1 checkpoint ships scale_factor_spatial=8, but its
+    # four downsampling stages compress 16x; the architecture wins.
+    with torch.device("meta"):
+        vae = AutoencoderKLQwenImage21(
+            base_dim=96,
+            decoder_base_dim=144,
+            z_dim=64,
+            dim_mult=[1, 2, 4, 8, 8],
+            num_res_blocks=2,
+            temperal_downsample=[False, True, True, True],
+            is_residual=True,
+            in_channels=4,
+            out_channels=4,
+            scale_factor_spatial=8,
+        )
+    assert vae.spatial_compression_ratio == 16
+
+
+def test_tiled_decode_output_shape_and_closeness():
+    torch.manual_seed(0)
+    vae = _make_vae()
+    z = torch.randn(1, 4, 1, 16, 16)
+    ratio = vae.spatial_compression_ratio
+
+    with torch.no_grad():
+        ref = vae._decode(z, return_dict=False)[0]
+        vae.enable_tiling(
+            tile_sample_min_height=32,
+            tile_sample_min_width=32,
+            tile_sample_stride_height=24,
+            tile_sample_stride_width=24,
+        )
+        tiled = vae.decode(z, return_dict=False)[0]
+
+    # Regression: with the wrong (config-derived) ratio the tiled output was
+    # cropped to a fraction of the sample size.
+    assert tiled.shape == (1, 4, 1, 16 * ratio, 16 * ratio)
+    assert tiled.shape == ref.shape
+
+    # Tiled decode blends overlapping tiles, so it is not bitwise-identical to
+    # untiled decode; on this toy model the seam error stays small.
+    mse = ((ref - tiled) ** 2).mean().item()
+    psnr = 10 * math.log10(4.0 / mse)
+    assert psnr > 20.0
+
+    # Tiled output is clamped to [-1, 1] like the untiled path.
+    assert tiled.min() >= -1.0 and tiled.max() <= 1.0
+
+
+def test_distributed_tile_hooks_match_tiled_decode():
+    """tile_split -> tile_exec -> tile_merge must reproduce tiled_decode exactly."""
+    torch.manual_seed(0)
+    vae = _make_vae()
+    vae.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_stride_height=24,
+        tile_sample_stride_width=24,
+    )
+    distributed = _make_vae(DistributedAutoencoderKLQwenImage21)
+    distributed.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_stride_height=24,
+        tile_sample_stride_width=24,
+    )
+    distributed.load_state_dict(vae.state_dict())
+
+    z = torch.randn(1, 4, 1, 16, 16)
+    with torch.no_grad():
+        tiled = vae.tiled_decode(z, return_dict=False)[0]
+
+        tasks, grid_spec = DistributedAutoencoderKLQwenImage21.tile_split(distributed, z)
+        # 16x16 latent, 8x8 latent tiles, stride 6 -> range(0, 16, 6) = 3 stops/axis.
+        assert len(tasks) == 9
+        assert grid_spec.grid_shape == (3, 3)
+        assert grid_spec.split_dims == (3, 4)
+        # Every tile carries all frames (T=1 for images).
+        assert all(len(task.tensor) == z.shape[2] for task in tasks)
+
+        coord_tensor_map = {task.grid_coord: distributed.tile_exec(task) for task in tasks}
+        merged = distributed.tile_merge(coord_tensor_map, grid_spec)
+
+    assert merged.shape == tiled.shape
+    torch.testing.assert_close(merged, tiled)
+
+
+def test_tiled_encode_output_shape_and_closeness():
+    """tiled_encode is engaged for image-conditioned generation when tiling is on."""
+    torch.manual_seed(0)
+    vae = _make_vae()
+    x = torch.randn(1, 4, 1, 64, 64)
+
+    with torch.no_grad():
+        ref = vae._encode(x)
+        vae.enable_tiling(
+            tile_sample_min_height=32,
+            tile_sample_min_width=32,
+            tile_sample_stride_height=24,
+            tile_sample_stride_width=24,
+        )
+        tiled = vae._encode(x)
+
+    ratio = vae.spatial_compression_ratio
+    assert tiled.shape == (1, 8, 1, 64 // ratio, 64 // ratio)
+    assert tiled.shape == ref.shape
+    # Blended tile borders deviate from untiled encode but stay small.
+    assert (ref - tiled).abs().max().item() < 1.0
+
+
+def test_encode_decode_round_trip_shape():
+    torch.manual_seed(0)
+    vae = _make_vae()
+    x = torch.randn(1, 4, 1, 32, 32)
+    with torch.no_grad():
+        latents = vae.encode(x, return_dict=False)[0].mode()
+        decoded = vae.decode(latents, return_dict=False)[0]
+    ratio = vae.spatial_compression_ratio
+    assert latents.shape == (1, 4, 1, 32 // ratio, 32 // ratio)
+    assert decoded.shape == x.shape
+
+
+@pytest.mark.parametrize("persistent", [False, True])
+def test_adaptive_oom_restores_tile_configuration(monkeypatch, persistent):
+    vae = _make_vae()
+    vae.enable_tiling(
+        tile_sample_min_height=512,
+        tile_sample_min_width=512,
+        tile_sample_stride_height=384,
+        tile_sample_stride_width=384,
+    )
+    attempts = []
+    z = torch.zeros(1, 4, 1, 4, 4)
+
+    def decode_tile(z, return_dict):
+        attempts.append(vae.tile_sample_min_height)
+        if persistent or vae.tile_sample_min_height > 128:
+            raise torch.OutOfMemoryError("injected tile OOM")
+        return (z,)
+
+    monkeypatch.setattr(vae, "tiled_decode", decode_tile)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+    if persistent:
+        with pytest.raises(torch.OutOfMemoryError):
+            vae._tiled_decode_adaptive(z, return_dict=False)
+    else:
+        assert vae._tiled_decode_adaptive(z, return_dict=False)[0] is z
+    assert attempts == [512, 256, 128]
+    assert (
+        vae.tile_sample_min_height,
+        vae.tile_sample_min_width,
+        vae.tile_sample_stride_height,
+        vae.tile_sample_stride_width,
+    ) == (512, 512, 384, 384)