[New model] Qwen-Image-2.1 support (rebase + review follow-ups) (#8099)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Signed-off-by: SamitHuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: SamitHuang <285365963@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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

@@ -64,6 +64,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

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from types import SimpleNamespace
+
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+    get_qwen_image_21_pre_process_func,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler, StepScheduler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def preprocess(tmp_path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
+    return get_qwen_image_21_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+
+
+def make_request(request_id, sizes, color="red"):
+    return OmniDiffusionRequest(
+        request_id=request_id,
+        prompt={
+            "prompt": "A ceramic teapot",
+            "multi_modal_data": {"image": [Image.new("RGB", size, color) for size in sizes]},
+        },
+        sampling_params=OmniDiffusionSamplingParams(
+            height=1024, width=1024, num_inference_steps=2, true_cfg_scale=1.0, seed=42
+        ),
+    )
+
+
+@pytest.mark.parametrize("scheduler_class", [RequestScheduler, StepScheduler])
+@pytest.mark.parametrize(
+    "first_sizes,second_sizes,compatible",
+    [
+        pytest.param([], [(64, 64)], False, id="text-and-edit"),
+        pytest.param([(64, 64)], [], False, id="edit-and-text"),
+        pytest.param([(64, 64)], [(64, 64), (64, 64)], False, id="image-count"),
+        pytest.param([(64, 64)], [(128, 64)], False, id="image-aspect-ratio"),
+        pytest.param([(64, 64), (128, 64)], [(128, 64), (64, 64)], False, id="slot-order"),
+        pytest.param([], [], True, id="two-text-requests"),
+        pytest.param([(64, 64)], [(64, 64)], True, id="different-image-content"),
+        pytest.param([(64, 64)], [(128, 128)], True, id="same-resized-layout"),
+        pytest.param([(64, 64), (128, 64)], [(128, 128), (256, 128)], True, id="same-multi-image-layout"),
+    ],
+)
+def test_scheduler_groups_by_condition_layout(preprocess, scheduler_class, first_sizes, second_sizes, compatible):
+    first = preprocess(make_request("first", first_sizes))
+    second = preprocess(make_request("second", second_sizes, color="blue"))
+    scheduler = scheduler_class()
+    scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+    scheduler.add_request(first)
+    scheduler.add_request(second)
+
+    scheduled = scheduler.schedule()
+    expected_ids = ["first", "second"] if compatible else ["first"]
+    assert scheduled.scheduled_request_ids == expected_ids
+
+    scheduler.finish_requests(expected_ids, DiffusionRequestStatus.FINISHED_COMPLETED)
+    if not compatible:
+        assert scheduler.schedule().scheduled_request_ids == ["second"]
+        scheduler.finish_requests("second", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert not scheduler.has_requests()
+
+
+@pytest.mark.parametrize(
+    "prompt", ["A teapot", {"prompt": "A teapot"}, {"prompt": "A teapot", "multi_modal_data": {"image": None}}]
+)
+def test_text_input_forms_share_compatibility_key(preprocess, prompt):
+    request = make_request("text", [])
+    request.prompt = prompt
+    processed = preprocess(request)
+    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"]
+
+
+def test_preprocess_normalizes_tuple_images_like_a_list(preprocess):
+    image = Image.new("RGB", (64, 64), "red")
+    tuple_request = make_request("tuple", [])
+    tuple_request.prompt["multi_modal_data"]["image"] = (image,)
+
+    processed = preprocess(tuple_request)
+    reference = preprocess(make_request("list", [(64, 64)]))
+
+    assert processed.batch_compatibility_key == reference.batch_compatibility_key
+    assert processed.batch_compatibility_key[0] == "qwen_image_21"
+    assert len(processed.batch_compatibility_key[1]) == 1
+    info = processed.prompt["additional_information"]
+    assert len(info["prompt_image"]) == 1
+    assert len(info["vae_images"]) == 1
+    assert len(info["input_image_sizes"]) == 1
+
+
+def test_preprocess_rejects_more_images_than_the_model_supports(preprocess):
+    from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+        MAX_QWEN_IMAGE_21_INPUT_IMAGES,
+    )
+
+    request = make_request("too-many", [(64, 64)] * (MAX_QWEN_IMAGE_21_INPUT_IMAGES + 1))
+    with pytest.raises(ValueError, match="At most"):
+        preprocess(request)
+
+
+def test_preprocess_defaults_true_cfg_scale(preprocess):
+    request = make_request("cfg", [])
+    request.sampling_params.true_cfg_scale = None
+    processed = preprocess(request)
+    assert processed.sampling_params.true_cfg_scale == 1.0
+
+    request = make_request("cfg-set", [])
+    request.sampling_params.true_cfg_scale = 4.0
+    assert preprocess(request).sampling_params.true_cfg_scale == 4.0
+
+
+@pytest.mark.parametrize("size,expected_wider", [((128, 64), True), ((64, 128), False)])
+def test_preprocess_derives_output_size_from_condition_aspect(preprocess, size, expected_wider):
+    request = make_request("aspect", [size])
+    request.sampling_params.height = None
+    request.sampling_params.width = None
+
+    processed = preprocess(request)
+
+    height = processed.sampling_params.height
+    width = processed.sampling_params.width
+    assert height % 32 == 0 and width % 32 == 0
+    assert (width > height) is expected_wider
+    info = processed.prompt["additional_information"]
+    assert info["calculated_height"] == height
+    assert info["calculated_width"] == width
+
+
+def test_preprocess_stages_rgba_vae_images_and_resized_prompt_images(preprocess):
+    processed = preprocess(make_request("staged", [(96, 64)]))
+    info = processed.prompt["additional_information"]
+
+    (vae_image,) = info["vae_images"]
+    # RGB input gains an opaque alpha channel for the RGBA VAE; the frame dim is
+    # inserted for the 3D (video-capable) VAE.
+    assert vae_image.shape[1] == 4
+    assert vae_image.shape[2] == 1
+    (prompt_image,) = info["prompt_image"]
+    assert isinstance(prompt_image, Image.Image)
+    (width, height) = info["input_image_sizes"][0]
+    assert prompt_image.size == (width, height)
+    assert processed.batch_compatibility_key == ("qwen_image_21", ((width, height),))

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -0,0 +1,318 @@
+# 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.hooks import HookRegistry, ModelHook
+from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
+from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
+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_model_level_offload_without_staging_disables_decode_graphs():
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True)
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_top_level_offload_hook_without_staging_disables_decode_graphs():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model)
+    registry = HookRegistry.get_or_create(model)
+    # A present-but-empty registry (e.g. after hooks were removed) must not disqualify.
+    assert manager._offload_reason() is None
+    registry.register_hook("sequential_offload", ModelHook())
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_persistent_staging_makes_model_level_offload_graph_eligible():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True)
+    hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu"))
+    HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook)
+    # No staging storage yet: a swap would allocate fresh storage per generation.
+    assert manager._offload_reason() is not None
+    # Staging established: fixed device storage keeps captured pointers valid.
+    hook._stager = object()
+    assert manager._offload_reason() is None
+    # A second, non-offload hook on the top-level module still disqualifies.
+    model._hook_registry.register_hook("teacache", ModelHook())
+    assert manager._offload_reason() is not 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)
+    first_key = weakref.ref(first[0]["cond"]["key"])
+    with torch.inference_mode():
+        register(manager, first)
+        # A second live request aliasing the same key is refused: it keeps its
+        # own cache and decodes eagerly while the owner is in flight.
+        register(manager, second)
+        assert decode(manager, second) is None
+        output = decode(manager, first)
+        expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
+        torch.testing.assert_close(output, expected)
+        # Once the owner's cache is released, the key is free to adopt.
+        del first
+        assert first_key() is None
+        register(manager, second)
+        output = decode(manager, second)
+        expected = (
+            torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2
+        )
+        torch.testing.assert_close(output, expected)
+        original = decode(manager, second).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, second), 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)
+        # A rebuilt cache (e.g. reordered rows) is a different live cache
+        # aliasing the same key: decode refuses the shared entry until the
+        # previous owner's tensors are released.
+        flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}]
+        kwargs["kv_cache"] = flipped
+        assert manager.try_decode(**kwargs) is None
+        del merged
+        register(manager, flipped)
+        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_pipeline.py

@@ -0,0 +1,402 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU unit tests for QwenImage21Pipeline helpers: prepare_latents duplication
+order, KV-cache assembly across step phases, and the post-process func."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+    QwenImage21Pipeline,
+    get_qwen_image_21_post_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _toy_pipeline() -> QwenImage21Pipeline:
+    """A pipeline shell with the attributes prepare_latents reads, no weights."""
+    pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.vae_scale_factor = 16
+    pipeline.latent_channels = 4
+    return pipeline
+
+
+# ---------------------------------------------------------------------------
+# prepare_latents
+# ---------------------------------------------------------------------------
+
+
+def test_prepare_latents_duplicates_condition_latents_request_major():
+    """batch=2 prompts x n=2 outputs must order condition latents A,A,B,B.
+
+    ``encode_prompt`` expands prompt embeddings request-major (A,A,B,B); the
+    condition latents must follow the same order or generations 2 and 3 swap
+    condition images. ``repeat`` would produce the wrong (A,B,A,B) tile-major
+    order; ``repeat_interleave`` is correct.
+    """
+    pipeline = _toy_pipeline()
+    cond_a = torch.full((1, 4, 1, 8, 8), 1.0)
+    cond_b = torch.full((1, 4, 1, 8, 8), 2.0)
+    stacked = torch.cat([cond_a, cond_b], dim=0)  # one row per request
+
+    latents, image_latents = pipeline.prepare_latents(
+        [stacked],
+        batch_size=4,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    assert latents.shape == (4, 16, 4)  # packed: 4x4 latent tokens, 4 channels
+    assert image_latents.shape == (4, 64, 4)  # packed: 8x8 condition tokens
+    row_means = image_latents.mean(dim=(1, 2))
+    torch.testing.assert_close(row_means, torch.tensor([1.0, 1.0, 2.0, 2.0]))
+
+
+def test_prepare_latents_without_duplication_when_rows_match_batch():
+    pipeline = _toy_pipeline()
+    stacked = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 4, 1, 2, 2).contiguous()
+
+    _, image_latents = pipeline.prepare_latents(
+        [stacked],
+        batch_size=4,
+        num_channels_latents=4,
+        height=32,
+        width=32,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    torch.testing.assert_close(image_latents.mean(dim=(1, 2)), torch.arange(4, dtype=torch.float32))
+
+
+def test_prepare_latents_rejects_non_divisible_duplication():
+    pipeline = _toy_pipeline()
+    stacked = torch.zeros(2, 4, 1, 8, 8)
+
+    with pytest.raises(ValueError, match="Cannot duplicate"):
+        pipeline.prepare_latents(
+            [stacked],
+            batch_size=3,
+            num_channels_latents=4,
+            height=64,
+            width=64,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            generator=torch.Generator().manual_seed(0),
+        )
+
+
+def test_prepare_latents_rejects_generator_list_length_mismatch():
+    pipeline = _toy_pipeline()
+
+    with pytest.raises(ValueError, match="list of generators"):
+        pipeline.prepare_latents(
+            None,
+            batch_size=2,
+            num_channels_latents=4,
+            height=32,
+            width=32,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            generator=[torch.Generator().manual_seed(0)],
+        )
+
+
+def test_prepare_latents_concatenates_multiple_condition_images():
+    pipeline = _toy_pipeline()
+    first = torch.full((2, 4, 1, 4, 4), 1.0)
+    second = torch.full((2, 4, 1, 2, 2), 3.0)
+
+    _, image_latents = pipeline.prepare_latents(
+        [first, second],
+        batch_size=2,
+        num_channels_latents=4,
+        height=32,
+        width=32,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    # Token dim concatenates the packed condition images: 16 + 4 per row.
+    assert image_latents.shape == (2, 20, 4)
+    torch.testing.assert_close(image_latents[:, :16].mean(), torch.tensor(1.0))
+    torch.testing.assert_close(image_latents[:, 16:].mean(), torch.tensor(3.0))
+
+
+# ---------------------------------------------------------------------------
+# KV-cache assembly across step phases
+# ---------------------------------------------------------------------------
+
+
+def _state(cache) -> SimpleNamespace:
+    return SimpleNamespace(extra={"kv_cache": cache})
+
+
+def _decode_cache(rows: int, value: float) -> list[dict]:
+    return [
+        {
+            "cond": {
+                "key": torch.full((rows, 2), value),
+                "value": torch.full((rows, 2), value + 0.5),
+            }
+        }
+    ]
+
+
+def test_kv_cache_phase_classification():
+    assert QwenImage21Pipeline._kv_cache_phase(None) == "none"
+    assert QwenImage21Pipeline._kv_cache_phase([{}]) == "prefill"
+    assert QwenImage21Pipeline._kv_cache_phase(_decode_cache(1, 0.0)) == "decode"
+
+
+def test_assemble_kv_cache_all_none():
+    assert QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(None)]) == (None, False)
+
+
+def test_assemble_kv_cache_prefill_returns_fresh_owned_cache():
+    states = [_state([{} for _ in range(3)]), _state([{} for _ in range(3)])]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
+
+    assert take_ownership is True
+    assert kv_cache == [{}, {}, {}]
+    # Fresh blocks, not aliases of either request's cache: the batched prefill
+    # writes into these and _scatter_kv_cache splits them back per request.
+    for block in kv_cache:
+        assert all(block is not request_block for state in states for request_block in state.extra["kv_cache"])
+
+
+def test_assemble_kv_cache_decode_concatenates_in_request_order():
+    first, second = _decode_cache(2, 1.0), _decode_cache(1, 3.0)
+    states = [_state(first), _state(second)]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
+
+    assert take_ownership is False
+    assert len(kv_cache) == 1
+    cond = kv_cache[0]["cond"]
+    torch.testing.assert_close(cond["key"], torch.tensor([[1.0, 1.0], [1.0, 1.0], [3.0, 3.0]]))
+    torch.testing.assert_close(cond["value"], torch.tensor([[1.5, 1.5], [1.5, 1.5], [3.5, 3.5]]))
+    # Decode only reads the cache: per-request tensors must stay untouched.
+    assert first[0]["cond"]["key"].shape == (2, 2)
+    assert second[0]["cond"]["key"].shape == (1, 2)
+
+
+def test_assemble_kv_cache_decode_merges_every_part():
+    first = [{"cond": {"key": torch.ones(1, 1), "value": torch.ones(1, 1), "key_scale": torch.ones(1)}}]
+    second = [{"cond": {"key": torch.zeros(1, 1), "value": torch.zeros(1, 1), "key_scale": torch.zeros(1)}}]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache([_state(first), _state(second)])
+
+    assert take_ownership is False
+    assert set(kv_cache[0]["cond"]) == {"key", "value", "key_scale"}
+    torch.testing.assert_close(kv_cache[0]["cond"]["key_scale"], torch.tensor([1.0, 0.0]))
+
+
+def test_assemble_kv_cache_mixed_phases_raises():
+    """Defensive backstop: a late prefill must never join a decoding batch.
+
+    The preprocessor sets ``allow_mixed_step_phases = False`` so the
+    StepScheduler defers late requests until the batch finishes decoding; this
+    raise is the last line of defense if a mixed batch is ever assembled.
+    """
+    states = [_state([{}]), _state(_decode_cache(1, 0.0))]
+
+    with pytest.raises(ValueError, match="mixed KV-cache phases"):
+        QwenImage21Pipeline._assemble_kv_cache(states)
+
+    with pytest.raises(ValueError, match="mixed KV-cache phases"):
+        QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(_decode_cache(1, 0.0))])
+
+
+def test_assemble_kv_cache_decode_rejects_mismatched_prefill_prompt_seq_len():
+    """Homogeneous merge invariant: denoise_step must sub-batch before assembling."""
+    short = _state(_decode_cache(1, 1.0))
+    short.extra["prefill_prompt_seq_len"] = 4
+    long = _state(_decode_cache(1, 2.0))
+    long.extra["prefill_prompt_seq_len"] = 8
+
+    with pytest.raises(ValueError, match="sub-batch by prefill_prompt_seq_len"):
+        QwenImage21Pipeline._assemble_kv_cache([short, long])
+
+    # Same baked length remains batchable.
+    same_a = _state(_decode_cache(1, 1.0))
+    same_a.extra["prefill_prompt_seq_len"] = 4
+    same_b = _state(_decode_cache(1, 2.0))
+    same_b.extra["prefill_prompt_seq_len"] = 4
+    _, take_ownership = QwenImage21Pipeline._assemble_kv_cache([same_a, same_b])
+    assert take_ownership is False
+
+
+def test_split_decode_groups_by_prefill_prompt_seq_len():
+    """Routine serving: separate prefills → joint decode must split, not fail."""
+    short = _state(_decode_cache(1, 1.0))
+    short.request_id = "short"
+    short.extra["prefill_prompt_seq_len"] = 4
+    long = _state(_decode_cache(1, 2.0))
+    long.request_id = "long"
+    long.extra["prefill_prompt_seq_len"] = 8
+    short2 = _state(_decode_cache(1, 3.0))
+    short2.request_id = "short2"
+    short2.extra["prefill_prompt_seq_len"] = 4
+
+    states = [short, long, short2]
+    assert QwenImage21Pipeline._needs_decode_seq_len_split(states) is True
+    groups = QwenImage21Pipeline._split_decode_groups(states)
+    assert [[s.request_id for s in g] for g in groups] == [["short", "short2"], ["long"]]
+
+    same = [_state(_decode_cache(1, 1.0)), _state(_decode_cache(1, 2.0))]
+    for s in same:
+        s.extra["prefill_prompt_seq_len"] = 4
+    assert QwenImage21Pipeline._needs_decode_seq_len_split(same) is False
+    assert QwenImage21Pipeline._split_decode_groups(same) == [same]
+
+
+def test_scatter_kv_cache_splits_rows_back_per_request():
+    states = [_state([{}]), _state([{}])]
+    merged = [
+        {
+            "cond": {
+                "key": torch.tensor([[1.0], [2.0], [3.0]]),
+                "value": torch.tensor([[4.0], [5.0], [6.0]]),
+            }
+        }
+    ]
+
+    QwenImage21Pipeline._scatter_kv_cache(states, [1, 2], merged)
+
+    torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[1.0]]))
+    torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["value"], torch.tensor([[4.0]]))
+    torch.testing.assert_close(states[1].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[2.0], [3.0]]))
+    scattered = states[1].extra["kv_cache"][0]["cond"]["key"]
+    assert scattered.is_contiguous()
+    # Must own storage (clone), not a view into the merged prefill tensor —
+    # otherwise step-mode CUDA-graph decode never matches owner identity.
+    assert scattered.data_ptr() != merged[0]["cond"]["key"][1:3].data_ptr()
+    merged[0]["cond"]["key"].fill_(0.0)
+    torch.testing.assert_close(scattered, torch.tensor([[2.0], [3.0]]))
+
+
+def test_encode_prompt_expands_masks_request_major():
+    """Two prompts × n=2 must keep embeds and masks in A,A,B,B order.
+
+    A 2D ``mask.repeat(1, n, 1).view(B*n, S)`` would wrongly produce A,B,A,B.
+    """
+    pipeline = _toy_pipeline()
+    embeds = torch.tensor([[[1.0], [1.0]], [[2.0], [0.0]]])  # [B=2, S=2, D=1]
+    mask = torch.tensor([[True, True], [True, False]])
+    image_pad = torch.tensor([[False, True], [False, False]])
+
+    def _fake_get(prompt, images_per_prompt, max_sequence_length=None, prompt_name="prompt"):
+        del prompt, images_per_prompt, max_sequence_length, prompt_name
+        return embeds.clone(), mask.clone(), image_pad.clone()
+
+    pipeline._get_qwen_prompt_embeds = _fake_get  # type: ignore[method-assign]
+
+    out_embeds, out_mask, out_pad = pipeline.encode_prompt(
+        ["short", "loooooonger"],
+        num_images_per_prompt=2,
+    )
+
+    assert out_embeds.shape[0] == 4
+    # embeds request-major: A,A,B,B
+    torch.testing.assert_close(out_embeds[:, 0, 0], torch.tensor([1.0, 1.0, 2.0, 2.0]))
+    # masks must follow the same order (not A,B,A,B)
+    assert out_mask.tolist() == [
+        [True, True],
+        [True, True],
+        [True, False],
+        [True, False],
+    ]
+    assert out_pad.tolist() == [
+        [False, True],
+        [False, True],
+        [False, False],
+        [False, False],
+    ]
+
+
+def test_resolve_prompt_seq_len_keeps_prefill_padding_after_batch_shrink():
+    """Longest-prompt cancel must not shrink the surviving request's KV layout."""
+    short = _state([{}])
+    short.extra["image_pad_mask"] = torch.tensor([[True, False]])
+    short.extra["prefill_prompt_seq_len"] = 4  # baked in when batched with a longer prompt
+
+    # Current InputBatch only sees the short prompt's native length.
+    assert QwenImage21Pipeline._resolve_prompt_seq_len([short], current_seq_len=2) == 4
+    padded = QwenImage21Pipeline._pad_rows_to_seq_len([short.extra["image_pad_mask"]], 4)
+    assert padded.tolist() == [[True, False, False, False]]
+
+
+# ---------------------------------------------------------------------------
+# post-process func
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def postprocess(tmp_path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
+    return get_qwen_image_21_post_process_func(SimpleNamespace(model=str(tmp_path)))
+
+
+def test_postprocess_tensor_returns_pil_images(postprocess):
+    from PIL import Image
+
+    images = torch.rand(2, 4, 16, 16) * 2 - 1
+
+    result = postprocess(images)
+
+    assert len(result) == 2
+    assert all(isinstance(image, Image.Image) for image in result)
+
+
+def test_postprocess_envelope_unwraps_payload_and_keeps_metadata(postprocess):
+    from PIL import Image
+
+    envelope = {
+        "payload": {"image": torch.rand(1, 4, 16, 16), "seed": 42},
+        "metadata": {"request_id": "abc"},
+    }
+
+    result = postprocess(envelope)
+
+    assert result["metadata"] == {"request_id": "abc"}
+    assert result["payload"]["seed"] == 42
+    assert len(result["payload"]["image"]) == 1
+    assert isinstance(result["payload"]["image"][0], Image.Image)
+
+
+def test_postprocess_envelope_requires_image_payload(postprocess):
+    with pytest.raises(ValueError, match=r"payload\['image'\]"):
+        postprocess({"payload": {"seed": 42}, "metadata": {}})
+
+
+def test_release_captured_graphs_delegates_to_transformer() -> None:
+    pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline)
+    released: list[bool] = []
+    pipeline.transformer = SimpleNamespace(release_captured_graphs=lambda: released.append(True))
+
+    pipeline.release_captured_graphs()
+
+    assert released == [True]
+
+
+def test_postprocess_envelope_normalizes_non_dict_metadata(postprocess):
+    result = postprocess({"payload": {"image": torch.rand(1, 4, 16, 16)}, "metadata": "bogus"})
+
+    assert result["metadata"] == {}

tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU lock for the Qwen-Image-2.1 CFG tail slice.
+
+2.1 prepends condition latents, so the target noise is the tail. The shared
+mixin slices the head. These tests stub ``predict_noise`` and check the
+returned rows are the tail, including rank-0 / rank-1 gather order.
+"""
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image_21 import cfg_parallel as cfg_mod
+from vllm_omni.diffusion.models.qwen_image_21.cfg_parallel import QwenImage21CFGParallelMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+PREFIX = 3
+TARGET = 2
+SCALE = 4.0
+
+
+class _Stub(QwenImage21CFGParallelMixin):
+    def __init__(self) -> None:
+        self.calls: list[dict[str, Any]] = []
+
+    def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+        self.calls.append(kwargs)
+        return kwargs["joint"]
+
+
+class _Gather:
+    def __init__(self, rank_tails: list[torch.Tensor]) -> None:
+        self.rank_tails = rank_tails
+        self.seen: list[torch.Tensor] = []
+
+    def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False) -> list[torch.Tensor]:
+        del dim
+        assert separate_tensors
+        self.seen.append(input_.clone())
+        return [tensor.clone() for tensor in self.rank_tails]
+
+
+def _joint(fill: float) -> torch.Tensor:
+    prefix = torch.full((1, PREFIX, 1), fill)
+    target = torch.full((1, TARGET, 1), fill + 100.0)
+    return torch.cat([prefix, target], dim=1)
+
+
+def _cfg(positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
+    return negative + SCALE * (positive - negative)
+
+
+def test_single_rank_cfg_uses_target_tail_not_condition_head() -> None:
+    pipe = _Stub()
+    positive = _joint(1.0)
+    negative = _joint(2.0)
+    actual = pipe.predict_noise_maybe_with_cfg(
+        do_true_cfg=True,
+        true_cfg_scale=SCALE,
+        positive_kwargs={"joint": positive, "branch": "pos"},
+        negative_kwargs={"joint": negative, "branch": "neg"},
+        output_slice=TARGET,
+    )
+    tail = _cfg(positive[:, -TARGET:], negative[:, -TARGET:])
+    head = _cfg(positive[:, :TARGET], negative[:, :TARGET])
+    torch.testing.assert_close(actual, tail)
+    assert not torch.allclose(actual, head)
+    assert [call["branch"] for call in pipe.calls] == ["pos", "neg"]
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_cfg_parallel_gather_order_uses_target_tails(monkeypatch: pytest.MonkeyPatch, rank: int) -> None:
+    pipe = _Stub()
+    positive = _joint(1.0)
+    negative = _joint(2.0)
+    pos_tail = positive[:, -TARGET:]
+    neg_tail = negative[:, -TARGET:]
+    group = _Gather([pos_tail, neg_tail])
+    monkeypatch.setattr(cfg_mod, "_get_cfg_world_size_or_one", lambda: 2)
+    monkeypatch.setattr(cfg_mod, "get_classifier_free_guidance_rank", lambda: rank)
+    monkeypatch.setattr(cfg_mod, "get_cfg_group", lambda: group)
+
+    actual = pipe.predict_noise_maybe_with_cfg(
+        do_true_cfg=True,
+        true_cfg_scale=SCALE,
+        positive_kwargs={"joint": positive, "branch": "pos"},
+        negative_kwargs={"joint": negative, "branch": "neg"},
+        output_slice=TARGET,
+    )
+
+    expected_branch = "pos" if rank == 0 else "neg"
+    expected_local = pos_tail if rank == 0 else neg_tail
+    assert [call["branch"] for call in pipe.calls] == [expected_branch]
+    torch.testing.assert_close(group.seen[0], expected_local)
+    torch.testing.assert_close(actual, _cfg(pos_tail, neg_tail))
+    head = _cfg(positive[:, :TARGET], negative[:, :TARGET])
+    assert not torch.allclose(actual, head)

tests/diffusion/models/qwen_image_21/test_registry.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""L1 tests for checkpoint-to-pipeline routing of Qwen-Image 2.1."""
+
+import json
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_registry_entry_loads_the_pipeline_class():
+    from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    assert DiffusionModelRegistry._try_load_model_cls("QwenImage21Pipeline") is QwenImage21Pipeline
+
+
+def test_pipeline_declares_request_batching_and_step_execution():
+    from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
+
+    assert QwenImage21Pipeline.supports_request_batch
+    assert QwenImage21Pipeline.supports_step_execution
+
+
+def test_process_funcs_resolve_to_the_pipeline_module():
+    from vllm_omni.diffusion.models.qwen_image_21 import pipeline_qwen_image_21
+    from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
+
+    pre_name = _DIFFUSION_PRE_PROCESS_FUNCS["QwenImage21Pipeline"]
+    post_name = _DIFFUSION_POST_PROCESS_FUNCS["QwenImage21Pipeline"]
+    assert pre_name == "get_qwen_image_21_pre_process_func"
+    assert post_name == "get_qwen_image_21_post_process_func"
+    assert callable(getattr(pipeline_qwen_image_21, pre_name))
+    assert callable(getattr(pipeline_qwen_image_21, post_name))
+
+
+def test_multimodal_metadata_caps_condition_images_at_four():
+    from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES, get_diffusion_model_metadata
+
+    metadata = get_diffusion_model_metadata("QwenImage21Pipeline")
+
+    assert metadata.supports_multimodal_inputs
+    assert metadata.max_multimodal_image_inputs == 4
+    assert metadata.max_multimodal_image_inputs == QWEN_IMAGE_21_MAX_INPUT_IMAGES
+
+
+def test_model_index_resolves_and_enriches(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"_class_name": "QwenImage21Pipeline"}),
+        encoding="utf-8",
+    )
+
+    assert resolve_model_class_name(str(tmp_path)) == "QwenImage21Pipeline"
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+
+    assert config.model_class_name == "QwenImage21Pipeline"
+    assert config.supports_multimodal_inputs
+    assert config.max_multimodal_image_inputs == 4

tests/diffusion/models/qwen_image_21/test_token_metadata.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU lock: adjacent condition images stay separate token blocks.
+
+Two condition images with no text between them are one run of True in
+``image_pad_mask``. Block ids come from ``img_shapes`` token counts, so they
+must not collapse into one bidirectional attention span.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
+    QwenImage21Transformer2DModel,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_adjacent_condition_images_stay_separate_blocks() -> None:
+    # text(2) | cond0(2) | cond1(3) | text(1) | target(2)
+    # cond0 and cond1 are one contiguous True run.
+    image_pad_mask = torch.tensor([False, False, True, True, True, True, True, False, True, True])
+    img_shapes = [(1, 1, 2), (1, 1, 3), (1, 1, 2)]
+
+    image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes)
+
+    assert image_ids.tolist() == [-1, -1, 0, 0, 1, 1, 1, -1, 2, 2]
+    assert target_token_mask.tolist() == [
+        False,
+        False,
+        False,
+        False,
+        False,
+        False,
+        False,
+        False,
+        True,
+        True,
+    ]
+    assert QwenImage21Transformer2DModel._build_full_attn_spans(image_ids) == [
+        (2, 4),
+        (4, 7),
+        (8, 10),
+    ]

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,266 @@
+# 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)
+
+
+def test_adaptive_oom_clears_feat_map_when_retries_exhausted(monkeypatch):
+    """Final min-tile OOM must not leave decoder activations on the device."""
+    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,
+    )
+    vae.clear_cache()
+    z = torch.zeros(1, 4, 1, 4, 4)
+    leaked = torch.ones(2, 2, device="cpu")  # stand-in for a cached activation
+
+    def decode_tile(z, return_dict):
+        # Mimic a partial tiled decode that already wrote into _feat_map before OOM.
+        vae._feat_map[0] = leaked
+        raise torch.OutOfMemoryError("injected tile OOM")
+
+    monkeypatch.setattr(vae, "tiled_decode", decode_tile)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+
+    with pytest.raises(torch.OutOfMemoryError):
+        vae._tiled_decode_adaptive(z, return_dict=False)
+
+    assert all(slot is None for slot in vae._feat_map)
+    assert all(slot is None for slot in vae._enc_feat_map)
+    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)
+
+
+def test_tiled_encode_oom_clears_enc_feat_map(monkeypatch):
+    """Tiled _encode must clear encoder activations if tiled_encode OOMs."""
+    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,
+    )
+    vae.clear_cache()
+    leaked = torch.ones(2, 2, device="cpu")
+    x = torch.zeros(1, 4, 1, 64, 64)
+
+    def encode_tile(x_in):
+        vae._enc_feat_map[0] = leaked
+        raise torch.OutOfMemoryError("injected tiled encode OOM")
+
+    monkeypatch.setattr(vae, "tiled_encode", encode_tile)
+
+    with pytest.raises(torch.OutOfMemoryError):
+        vae._encode(x)
+
+    assert all(slot is None for slot in vae._enc_feat_map)
+    assert all(slot is None for slot in vae._feat_map)

tests/diffusion/offloader/test_sequential_backend.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for SequentialOffloadBackend."""
 
@@ -284,6 +284,102 @@ def test_component_selective_model_offload_requires_swap_counterpart(component,
         backend.enable(pipeline)
 
 
+def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None:
+    dit = _create_simple_module().to(accelerator_device)
+    encoder = _create_simple_module().to(accelerator_device)
+    original_weight = dit.linear.weight.detach().clone()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=accelerator_device,
+        pin_memory=False,
+        persistent_dit_staging=True,
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook._stager is not None
+    assert encoder_hook._stager is None  # encoders keep plain move semantics
+
+    # Registration snapshots the CPU master and rebinds to it.
+    assert dit.linear.weight.device.type == "cpu"
+
+    dit_hook._to_gpu(dit)
+    first_ptr = dit.linear.weight.data_ptr()
+    assert dit.linear.weight.device == accelerator_device
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+    retained_storage = dit_hook._stager._device_storages[0]
+    retained_ptr = retained_storage.data_ptr()
+
+    # Offload via a *different* module's hook (the cross-module swap path).
+    # The module parameters move back to the CPU master, but the fixed device
+    # storage stays allocated so CUDA-graph-captured weight pointers stay valid.
+    encoder_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+    assert dit_hook._stager._device_storages
+    assert dit_hook._stager._device_storages[0].data_ptr() == retained_ptr
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+
+    # A second offload/load cycle through the DiT hook itself stays put too.
+    dit_hook._to_cpu(dit)
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+
+    remove_sequential_offload([dit, encoder])
+
+
+def test_default_sequential_offload_releases_device_storage(accelerator_device) -> None:
+    """Without persistent staging, offloading keeps plain move semantics."""
+    dit = _create_simple_module().to(accelerator_device)
+    encoder = _create_simple_module().to(accelerator_device)
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=accelerator_device,
+        pin_memory=False,
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook._stager is None  # no retained staging: GPU storage is freed on offload
+
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.device == accelerator_device
+    dit_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+    assert all(parameter.device.type == "cpu" for parameter in dit.parameters())
+
+    remove_sequential_offload([dit, encoder])
+
+
+@pytest.mark.parametrize(
+    ("decode_graphs", "expect_stager"),
+    [(True, True), (False, False)],
+)
+def test_model_level_persistent_staging_gated_on_decode_graph(
+    accelerator_device, decode_graphs: bool, expect_stager: bool
+) -> None:
+    """ModelLevelOffloadBackend only retains DiT staging for decode-graph models."""
+    pipeline = nn.Module()
+    pipeline.transformer = _create_simple_module()
+    if decode_graphs:
+        pipeline.transformer.enable_cuda_graph_decode = True
+    pipeline.text_encoder = _create_simple_module()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        accelerator_device,
+    )
+    backend.enable(pipeline)
+
+    hook = pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert (hook._stager is not None) is expect_stager
+
+    backend.disable()
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()

tests/diffusion/test_diffusion_config_propagation.py

@@ -46,6 +46,21 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
 
+def test_direct_config_extras_are_independent_dicts():
+    first = OmniDiffusionConfig(model="test-model")
+    second = OmniDiffusionConfig(model="test-model")
+
+    assert first.extras.get("prefix_kv_cache_dtype") is None
+    first.extras["prefix_kv_cache_dtype"] = "fp8"
+    assert second.extras == {}
+
+
+def test_direct_config_preserves_explicit_extras():
+    config = OmniDiffusionConfig(model="test-model", extras={"prefix_kv_cache_dtype": "fp8"})
+
+    assert config.extras["prefix_kv_cache_dtype"] == "fp8"
+
+
 class TestParallelConfigPropagation:
     """Core regression tests: parallel_config must survive serialization."""
 

tests/diffusion/test_diffusion_model_runner.py

@@ -521,6 +521,35 @@ def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp):
     ]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_transformer_scopes_inductor_cudagraphs_to_decode_graph(monkeypatch):
+    model = SimpleNamespace(enable_cuda_graph_decode=True)
+    runner = _make_compile_runner(model)
+    compile_calls = []
+
+    def _regionally_compile(target, *args, **kwargs):
+        compile_calls.append(kwargs)
+        return target
+
+    monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+    import torch._inductor.config as inductor_config
+
+    before = inductor_config.triton.cudagraphs
+    DiffusionModelRunner._compile_transformer(runner, "transformer")
+
+    assert inductor_config.triton.cudagraphs == before
+    assert compile_calls == [
+        {
+            "dynamic": True,
+            "options": {
+                "triton.cudagraphs": False,
+                "triton.cudagraph_trees": False,
+            },
+        }
+    ]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_compile_transformer_uses_regional_dynamic_false_config(monkeypatch):

tests/e2e/online_serving/test_qwen_image_21.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
+
+MODEL = os.environ.get("QWEN_IMAGE_21_TEST_MODEL", "Qwen/Qwen-Image-2.1")
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(model=MODEL, server_args=["--vae-use-tiling", "--max-num-seqs", "2"] + extra),
+        id=name,
+        marks=hardware_marks(res={"cuda": "H100"}),
+    )
+    for name, extra in [("request", []), ("step", ["--step-execution"])]
+]
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.advanced_model,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(
+        not os.environ.get("QWEN_IMAGE_21_TEST_MODEL"),
+        reason="Set QWEN_IMAGE_21_TEST_MODEL to a provisioned checkpoint until public weights are available.",
+    ),
+]
+
+
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient):
+    online_client.send_diffusion_request(
+        {
+            "model": omni_server.model,
+            "messages": dummy_messages_from_mix_data(content_text="A red ceramic teapot on a wooden table."),
+            "extra_body": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 4,
+                "seed": 42,
+                "negative_prompt": "blurry, low quality",
+                "true_cfg_scale": 4.0,
+            },
+        }
+    )
+
+
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient):
+    image = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    online_client.send_diffusion_request(
+        {
+            "model": omni_server.model,
+            "messages": dummy_messages_from_mix_data(content_text="Make the image blue.", image_data_url=image),
+            "extra_body": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 4,
+                "seed": 42,
+                "negative_prompt": "blurry, low quality",
+                "true_cfg_scale": 4.0,
+            },
+        }
+    )

tests/entrypoints/openai_api/test_image_server.py

@@ -933,7 +933,7 @@ def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_
     assert payloads[0]["index"] == 0
     assert payloads[1]["delta"] == " done"
     assert payloads[2]["output_format"] == "png"
-    assert payloads[2]["size"] == "16x16"
+    assert payloads[2]["size"] == "32x24"
 
     image_payload = payloads[2]["data"][0]
     img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
@@ -2290,8 +2290,8 @@ def test_extract_images_from_result():
     class EmptyResult:
         pass
 
-    result = EmptyResult()
-    images = _extract_images_from_result(result)
+    empty_result = EmptyResult()
+    images = _extract_images_from_result(empty_result)
     assert images == []
 
     # Test nested batch: [np.array(shape=(3, 64, 64, 3))]
@@ -2301,8 +2301,8 @@ def test_extract_images_from_result():
         def __init__(self):
             self.images = [batch]
 
-    result = BatchResult()
-    images = _extract_images_from_result(result)
+    batch_result = BatchResult()
+    images = _extract_images_from_result(batch_result)
     assert len(images) == 3
     assert all(isinstance(img, Image.Image) for img in images)
     assert all(img.size == (64, 64) for img in images)
@@ -2312,8 +2312,8 @@ def test_extract_images_from_result():
         def __init__(self):
             self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]
 
-    result = DictRequestOutput()
-    images = _extract_images_from_result(result)
+    dict_result = DictRequestOutput()
+    images = _extract_images_from_result(dict_result)
     assert len(images) == 1
     assert isinstance(images[0], Image.Image)
 
@@ -2322,8 +2322,8 @@ def test_extract_images_from_result():
         def __init__(self):
             self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]
 
-    result = AttrRequestOutput()
-    images = _extract_images_from_result(result)
+    attr_result = AttrRequestOutput()
+    images = _extract_images_from_result(attr_result)
     assert len(images) == 1
     assert isinstance(images[0], Image.Image)
     assert images[0].size == (32, 32)

tests/model_extras/test_model_extras.py

@@ -271,7 +271,7 @@ def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected:
 
 
 def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured = {}
+    captured: dict[str, Any] = {}
 
     def fake_policy(*, model, revision=None):
         captured.update(model=model, revision=revision)
@@ -495,6 +495,27 @@ def test_bagel_image_to_image_prompt_builder() -> None:
     assert result["negative_prompt"] == "ugly"
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_qwen_image_21_image_to_image_prompt_builder() -> None:
+    dummy_image = Image.new("RGBA", (64, 64))
+    result = build_image_to_image_prompt(
+        "QwenImage21Pipeline",
+        prompt="make it dance",
+        negative_prompt="blurry",
+        input_image=dummy_image,
+        height=1024,
+        width=768,
+    )
+    assert result == {
+        "prompt": "make it dance",
+        "multi_modal_data": {"image": dummy_image},
+        "height": 1024,
+        "width": 768,
+        "negative_prompt": "blurry",
+    }
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None:
@@ -785,7 +806,7 @@ def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: obj
             return tokenizer
 
     fake_transformers = type(sys)("transformers")
-    fake_transformers.AutoTokenizer = _FakeAutoTokenizer
+    setattr(fake_transformers, "AutoTokenizer", _FakeAutoTokenizer)
     monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
     return calls
 

tests/model_tests/diffusion/test_alignment.py

@@ -25,6 +25,9 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
 EXCLUDED_MODELS = [
     "QwenImageLayeredPipeline",
+    # Dedicated CPU/unit coverage in tests/diffusion/models/qwen_image_21/.
+    # Tiny-model offline suite is not wired yet (2.1 transformer/VAE configs).
+    "QwenImage21Pipeline",
     "GlmImagePipeline",
     "ZImagePipeline",
     "OvisImagePipeline",