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

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

tests/diffusion/attention/test_cudnn_attn.py

@@ -12,9 +12,10 @@ import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
+@pytest.mark.cpu
 def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
     """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
     selected_backends = []
@@ -40,6 +41,7 @@ def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
     assert selected_backends == [(SDPBackend.MATH,)]
 
 
+@pytest.mark.cpu
 def test_explicit_cudnn_rejects_kv_seq_len_one():
     impl = CuDNNAttentionImpl(
         num_heads=2,
@@ -54,6 +56,7 @@ def test_explicit_cudnn_rejects_kv_seq_len_one():
         impl.forward_cuda(query, singleton_kv, singleton_kv)
 
 
+@pytest.mark.cpu
 def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
     selected_backends = []
 
@@ -77,6 +80,7 @@ def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
     assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]
 
 
+@pytest.mark.cpu
 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
     observed: dict[str, Any] = {}
 
@@ -112,6 +116,7 @@ def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
     assert observed["kwargs"]["attn_mask"] is None
 
 
+@pytest.mark.cpu
 def test_cudnn_rejects_invalid_valid_kv_length():
     impl = CuDNNAttentionImpl(
         num_heads=2,
@@ -130,12 +135,97 @@ def test_cudnn_rejects_invalid_valid_kv_length():
 
 
 @pytest.mark.parametrize("head_size", [8, 64, 128, 256])
+@pytest.mark.cpu
 def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
     assert CuDNNAttentionBackend.supports_head_size(head_size)
     assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
 
 
 @pytest.mark.parametrize("head_size", [0, 7, 12, 320])
+@pytest.mark.cpu
 def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
     assert not CuDNNAttentionBackend.supports_head_size(head_size)
     assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("kv_heads", [1, 2])
+@pytest.mark.parametrize("query_offset", [0, 4])
+@pytest.mark.cpu
+def test_piecewise_matches_dense_block_causal_mask(monkeypatch, batch_size, kv_heads, query_offset):
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    torch.manual_seed(42)
+    query = torch.randn(batch_size, 12, 2, 8)
+    key = torch.randn(batch_size, 12, kv_heads, 8)
+    value = torch.randn_like(key)
+    spans = [(2, 5), (7, 10)]
+    positions = torch.arange(12)
+    mask = positions[:, None] >= positions[None, :]
+    for start, end in spans:
+        mask[start:end, start:end] = True
+    expected = torch.nn.functional.scaled_dot_product_attention(
+        query[:, query_offset:].transpose(1, 2),
+        key.transpose(1, 2),
+        value.transpose(1, 2),
+        attn_mask=mask[None, None, query_offset:],
+        scale=0.5,
+        enable_gqa=kv_heads == 1,
+    ).transpose(1, 2)
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    actual = impl.forward_cuda(
+        query[:, query_offset:], key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size)
+    )
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.cpu
+def test_piecewise_preserves_discontinuous_query_ranges(monkeypatch):
+    from vllm_omni.diffusion.attention.backends.abstract import QueryRange
+
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    query, key, value = [torch.randn(2, 12, 2, 8) for _ in range(3)]
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (7, 12)]] * 2)
+    expected = impl.forward_cuda(query, key, value, metadata)
+    metadata.query_ranges = (QueryRange(0, 3, 0), QueryRange(3, 6, 8))
+    actual = impl.forward_cuda(torch.cat([query[:, :3], query[:, 8:11]], dim=1), key, value, metadata)
+    torch.testing.assert_close(actual, torch.cat([expected[:, :3], expected[:, 8:11]], dim=1))
+
+
+@pytest.mark.cpu
+def test_piecewise_rejects_padding_mask(monkeypatch):
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    query = torch.randn(1, 8, 2, 8)
+    metadata = AttentionMetadata(full_attn_spans=[[(2, 8)]], attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    with pytest.raises(ValueError, match="padding mask"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
+@pytest.mark.cpu
+def test_dense_mask_takes_precedence_over_spans(monkeypatch):
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    query, key, value = [torch.randn(2, 8, 2, 8) for _ in range(3)]
+    mask = torch.ones(2, 1, 8, 8, dtype=torch.bool).tril()
+    mask[1, :, :, 3] = False
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask))
+    actual = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask, full_attn_spans=[[(2, 8)]] * 2))
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_cudnn_piecewise_matches_dense_on_device(batch_size):
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    query, key, value = [
+        torch.randn(batch_size, 96, 2, 128, generator=generator, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+    ]
+    spans = [(8, 40), (48, 96)]
+    mask = torch.ones(96, 96, device="cuda", dtype=torch.bool).tril()
+    for start, end in spans:
+        mask[start:end, start:end] = True
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=128, softmax_scale=128**-0.5)
+    actual = impl.forward_cuda(query, key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size))
+    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask[None, None]))
+    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.004)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass
 from types import SimpleNamespace
 
@@ -62,6 +65,7 @@ class E2EOperator:
 @dataclass
 class FakeWorldGroup:
     device_group: object
+    cpu_group: object = None
 
 
 class DummyMixin(DistributedVaeMixin):

tests/diffusion/distributed/test_vae_error_recovery.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as vae_executor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TileOperator:
+    def split(self, tensor):
+        return [vae_executor.TileTask(i, (i,), row[None]) for i, row in enumerate(tensor)], vae_executor.GridSpec(
+            split_dims=(0,), grid_shape=(tensor.shape[0],)
+        )
+
+    def exec(self, task):
+        return task.tensor * 2
+
+    def merge(self, tiles, grid):
+        return torch.cat([tiles[(i,)] for i in range(grid.grid_shape[0])])
+
+
+def run_recovery_worker(rank, init_file):
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo", init_method=f"file://{init_file}", rank=rank, world_size=2, timeout=timedelta(seconds=30)
+    )
+    try:
+        group = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(vae_executor, "get_world_group", lambda: group)
+            executor = vae_executor.DistributedVaeExecutor()
+        executor.set_parallel_size(2)
+        tensor = torch.arange(12, dtype=torch.float32).reshape(4, 3)
+        operator = TileOperator()
+        for stage, failing_rank in [("decode", 1), ("packing", 1), ("merge", 0), ("gather", 1)]:
+            with pytest.MonkeyPatch.context() as patch:
+                if rank == failing_rank:
+
+                    def fail(*args, **kwargs):
+                        raise torch.OutOfMemoryError("injected OOM")
+
+                    target, attribute = {
+                        "decode": (operator, "exec"),
+                        "packing": (executor, "_pack_local_tiles"),
+                        "merge": (operator, "merge"),
+                        "gather": (torch, "empty_like"),
+                    }[stage]
+                    patch.setattr(target, attribute, fail)
+                with pytest.raises(torch.OutOfMemoryError, match="Distributed VAE"):
+                    executor.execute(tensor, operator)
+            torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
+
+        with pytest.raises(RuntimeError, match="Distributed VAE mixed failures failed"):
+            with executor._sync_errors("mixed failures"):
+                if rank == 0:
+                    raise torch.OutOfMemoryError("retryable")
+                raise ValueError("not retryable")
+        torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
+    finally:
+        dist.destroy_process_group()
+
+
+def test_all_ranks_observe_failures_and_can_retry(tmp_path):
+    torch.multiprocessing.spawn(run_recovery_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)

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)