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

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

.buildkite/cuda/test-merge.yml

@@ -221,6 +221,18 @@ steps:
           - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_1
 
+      - label: "Diffusion · Qwen Image 2.1 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_21.py
+          - tests/helpers/
+          - vllm_omni/diffusion/models/qwen_image_21/
+          - vllm_omni/diffusion/distributed/autoencoders/
+          - vllm_omni/diffusion/sched/
+          - vllm_omni/entrypoints/openai/
+        commands:
+          - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen_image_21.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_1
+
       - label: "Diffusion · Qwen Image Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen_image.py

.buildkite/cuda/test-ready.yml

@@ -331,6 +331,18 @@ steps:
             "
         mirror_hardwares: h100_2
 
+      - label: "Diffusion · Qwen Image 2.1 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_21.py
+          - tests/helpers/
+          - vllm_omni/diffusion/models/qwen_image_21/
+          - vllm_omni/diffusion/distributed/autoencoders/
+          - vllm_omni/diffusion/sched/
+          - vllm_omni/entrypoints/openai/
+        commands:
+          - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen_image_21.py -m 'core_model and cuda' --run-level 'core_model'
+        mirror_hardwares: h100_1
+
       - label: "Diffusion · Qwen Image Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen_image.py

examples/offline_inference/image_to_image/image_edit.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
 
 """
 Example script for image editing with OmniGen2.
@@ -121,14 +121,14 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
+    parser = argparse.ArgumentParser(description="Edit images with vLLM-Omni.")
     parser.add_argument(
         "--model",
         default="Qwen/Qwen-Image-Edit",
         help=(
             "Diffusion model name or local path. "
-            "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511"
-            "which supports QwenImageEditPlusPipeline."
+            "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509, "
+            "Qwen/Qwen-Image-Edit-2511, or Qwen/Qwen-Image-2.1."
         ),
     )
     parser.add_argument(
@@ -280,7 +280,7 @@ def parse_args() -> argparse.Namespace:
         "--color-format",
         type=str,
         default="RGB",
-        help="For Qwen-Image-Layered, set to RGBA.",
+        help="For Qwen-Image-Layered or Qwen-Image-2.1, set to RGBA to preserve transparency.",
     )
 
     # Cache-DiT specific parameters
@@ -541,6 +541,7 @@ def main():
 
     diffusion_params = OmniDiffusionSamplingParams(
         generator=generator,
+        seed=args.seed,
         true_cfg_scale=args.cfg_scale,
         guidance_scale=args.guidance_scale,
         guidance_scale_2=args.guidance_scale_2,

examples/offline_inference/image_to_image/qwen_image_21_edit.py

@@ -1,262 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-"""
-Example script for image-conditioned generation with Qwen-Image 2.1.
-
-Qwen-Image 2.1 uses a single pipeline (QwenImage21Pipeline) for both
-text-to-image and image-conditioned generation. Pass one or more condition
-images (at most 4) via --image to run the image-conditioned path.
-
-Usage (single image):
-    python qwen_image_21_edit.py \
-        --image input.png \
-        --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
-        --negative-prompt "blurry, low quality, text, watermark" \
-        --output qwen_image_21_edit.png \
-        --num-inference-steps 50 \
-        --cfg-scale 4.0
-
-Usage (multiple images, up to 4):
-    python qwen_image_21_edit.py \
-        --image input1.png input2.png \
-        --prompt "Combine these images into a single scene" \
-        --negative-prompt "blurry, low quality" \
-        --output qwen_image_21_edit.png \
-        --num-inference-steps 50 \
-        --cfg-scale 4.0
-
-Notes:
-    - Qwen-Image 2.1 only supports true classifier-free guidance
-      (--cfg-scale); there is no guidance-distilled guidance_scale.
-      --cfg-scale > 1 only takes effect when --negative-prompt is provided.
-    - --height/--width must be multiples of 32. When omitted, the output size
-      is derived from the last condition image's aspect ratio at ~1024x1024.
-    - Cache acceleration backends (cache_dit / tea_cache) are not supported.
-
-For more options, run:
-    python qwen_image_21_edit.py --help
-"""
-
-import argparse
-import os
-import time
-from pathlib import Path
-
-import torch
-from PIL import Image
-
-from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_extras import (
-    build_image_to_image_prompt,
-    get_model_class_name,
-)
-from vllm_omni.platforms import current_omni_platform
-
-MAX_INPUT_IMAGES = 4
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Image-conditioned generation with Qwen-Image 2.1.")
-    parser.add_argument(
-        "--model",
-        default="Qwen/Qwen-Image-2.1",
-        help="Diffusion model name or local path.",
-    )
-    parser.add_argument(
-        "--image",
-        type=str,
-        nargs="+",
-        required=True,
-        help="Path(s) to input condition image file(s) (PNG, JPG, etc.). At most 4 images.",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="Text prompt describing the desired generation.",
-    )
-    parser.add_argument(
-        "--negative-prompt",
-        type=str,
-        default=None,
-        help="Negative prompt. Required for true classifier-free guidance (--cfg-scale > 1) to take effect.",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=None,
-        metavar="W",
-        help="Output image width in pixels (multiple of 32). Default: None (derived from the condition image).",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=None,
-        metavar="H",
-        help="Output image height in pixels (multiple of 32). Default: None (derived from the condition image).",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for deterministic results.",
-    )
-    parser.add_argument(
-        "--cfg-scale",
-        type=float,
-        default=4.0,
-        help=(
-            "True classifier-free guidance scale (default: 4.0). Enabled by setting cfg_scale > 1 and providing "
-            "a negative prompt. Qwen-Image 2.1 supports true CFG only; there is no guidance_scale."
-        ),
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        default="qwen_image_21_edit.png",
-        help="Path to save the generated image (PNG).",
-    )
-    parser.add_argument(
-        "--num-outputs-per-prompt",
-        type=int,
-        default=1,
-        help="Number of images to generate for the given prompt.",
-    )
-    parser.add_argument(
-        "--num-inference-steps",
-        type=int,
-        default=50,
-        help="Number of denoising steps for the diffusion sampler.",
-    )
-    parser.add_argument(
-        "--ulysses-degree",
-        type=int,
-        default=1,
-        help="Number of GPUs used for ulysses sequence parallelism (ring attention is not supported).",
-    )
-    parser.add_argument(
-        "--cfg-parallel-size",
-        type=int,
-        default=1,
-        choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
-    )
-    parser.add_argument(
-        "--vae-use-tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization (also tiles condition-image encoding).",
-    )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        default=None,
-        help="Disable torch.compile and force eager execution.",
-    )
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    if len(args.image) > MAX_INPUT_IMAGES:
-        raise ValueError(f"Qwen-Image 2.1 supports at most {MAX_INPUT_IMAGES} input images, got {len(args.image)}")
-
-    input_images = []
-    for image_path in args.image:
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Input image not found: {image_path}")
-        input_images.append(Image.open(image_path).convert("RGBA"))
-
-    # Use single image or list based on number of inputs
-    input_image = input_images[0] if len(input_images) == 1 else input_images
-
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-
-    omni_kwargs = {
-        "model": args.model,
-        "ulysses_degree": args.ulysses_degree,
-        "cfg_parallel_size": args.cfg_parallel_size,
-        "vae_use_tiling": args.vae_use_tiling,
-    }
-    if args.enforce_eager is not None:
-        omni_kwargs["enforce_eager"] = args.enforce_eager
-    omni = Omni(**omni_kwargs)
-    model_class_name = get_model_class_name(omni)
-    print("Pipeline loaded")
-
-    print(f"\n{'=' * 60}")
-    print("Generation Configuration:")
-    print(f"  Model: {args.model}")
-    print(f"  Inference steps: {args.num_inference_steps}")
-    print(f"  CFG scale (true CFG): {args.cfg_scale}")
-    if args.height is not None or args.width is not None:
-        print(f"  Output size: {args.width or 'auto'}x{args.height or 'auto'}")
-    if isinstance(input_image, list):
-        print(f"  Number of input images: {len(input_image)}")
-    else:
-        print(f"  Input image size: {input_image.size}")
-    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, cfg_parallel_size={args.cfg_parallel_size}")
-    print(f"{'=' * 60}\n")
-
-    generation_start = time.perf_counter()
-
-    prompt_dict = build_image_to_image_prompt(
-        model_class_name=model_class_name,
-        prompt=args.prompt,
-        negative_prompt=args.negative_prompt,
-        input_image=input_image,
-        height=args.height,
-        width=args.width,
-    )
-
-    diffusion_params = OmniDiffusionSamplingParams(
-        generator=generator,
-        seed=args.seed,
-        true_cfg_scale=args.cfg_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_outputs_per_prompt,
-        height=args.height,
-        width=args.width,
-    )
-
-    outputs = omni.generate(prompt_dict, sampling_params_list=[diffusion_params])
-    generation_end = time.perf_counter()
-    generation_time = generation_end - generation_start
-
-    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
-
-    if not outputs:
-        raise ValueError("No output generated from omni.generate()")
-
-    images = None
-    for output in outputs:
-        images = getattr(output, "images", None)
-        if images:
-            break
-
-    if not images:
-        images = extract_images_from_outputs(outputs)
-
-    if not images:
-        raise ValueError("No images found in request_output")
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    suffix = output_path.suffix or ".png"
-    stem = output_path.stem or "qwen_image_21_edit"
-
-    if len(images) <= 1:
-        images[0].save(output_path)
-        print(f"Saved generated image to {os.path.abspath(output_path)}")
-    else:
-        for idx, img in enumerate(images):
-            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
-            img.save(save_path)
-            print(f"Saved generated image to {os.path.abspath(save_path)}")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh

@@ -1,8 +1,17 @@
-python qwen_image_21_edit.py \
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+set -euo pipefail
+
+python "$(dirname "${BASH_SOURCE[0]}")/image_edit.py" \
     --model Qwen/Qwen-Image-2.1 \
+    --color-format RGBA \
+    --seed 42 \
     --image qwen_bear.png \
     --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
     --negative-prompt "blurry, low quality, text, watermark" \
     --output qwen_image_21_edit.png \
     --num-inference-steps 50 \
-    --cfg-scale 4.0
+    --cfg-scale 4.0 \
+    "$@"

recipes/Qwen/Qwen-Image-2.1.md

@@ -27,7 +27,7 @@ text-to-image.
 
 - Related offline examples:
   [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py),
-  [`examples/offline_inference/image_to_image/qwen_image_21_edit.py`](../../examples/offline_inference/image_to_image/qwen_image_21_edit.py)
+  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
 - Supported models table:
   [`docs/models/supported_models.md`](../../docs/models/supported_models.md)
 - Feature compatibility matrix:
@@ -52,7 +52,10 @@ python examples/offline_inference/text_to_image/text_to_image.py \
 ### Image-conditioned generation (editing)
 
 ```bash
-python examples/offline_inference/image_to_image/qwen_image_21_edit.py \
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --color-format RGBA \
+  --seed 42 \
   --image qwen_bear.png \
   --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
   --negative-prompt "blurry, low quality, text, watermark" \
@@ -70,7 +73,10 @@ bash examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh
 Up to 4 condition images can be passed via `--image`:
 
 ```bash
-python examples/offline_inference/image_to_image/qwen_image_21_edit.py \
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --color-format RGBA \
+  --seed 42 \
   --image input1.png input2.png \
   --prompt "Combine these images into a single scene" \
   --negative-prompt "blurry, low quality" \
@@ -79,6 +85,13 @@ python examples/offline_inference/image_to_image/qwen_image_21_edit.py \
   --cfg-scale 4.0
 ```
 
+Use `--color-format RGBA` to preserve transparency in condition images. The
+shared example otherwise loads inputs as RGB. Add `--width` and `--height`
+(multiples of 32) to set the output size; when omitted, Qwen-Image 2.1 derives
+it from the last condition image's aspect ratio at approximately 1024×1024.
+Use `--num-outputs-per-prompt` to generate multiple images and `--vae-use-tiling`
+to reduce VAE memory usage. The launcher accepts these options as well.
+
 ## Online Serving
 
 Start the server:
@@ -102,6 +115,32 @@ requests wait for that batch to finish before starting their own prefill.
 `--max-num-seqs` controls batch capacity; mixed prefill/decode admission is not
 supported yet.
 
+### CUDA Graph decode
+
+Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising
+decode steps when `enforce_eager=False` (the default). Prefill remains eager.
+This works with both request execution and `--step-execution`.
+
+To disable graph capture:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --enforce-eager --vae-use-tiling
+```
+
+The same control is available as `Omni(..., enforce_eager=True)` and as
+`enforce_eager: true` on a stage in the deployment YAML. It disables both graph
+capture and automatic `torch.compile`. Otherwise, this transformer uses its
+model-specific decode graphs instead of automatic `torch.compile`.
+
+The autoregressive engine's `compilation_config.cudagraph_mode` does not control
+this diffusion path; use `enforce_eager` to force eager execution.
+
+TP/SP/ring parallelism, HSDP, offload/cache hooks, compiled blocks, quantized
+KV caches, dynamic LoRA, and padded text masks fall back to eager decode.
+Graphs keep separate entries for different image layouts and copy request-owned
+prefix K/V into static buffers before replay. Those buffers require additional
+memory; graph entries are bounded by the transformer's cache limit.
+
 ### Verification
 
 For a direct API smoke test:

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)

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

@@ -849,7 +849,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"])))
@@ -1719,7 +1719,7 @@ def test_image_edit_parameter_pass(async_omni_test_client):
         img = Image.open(io.BytesIO(img_bytes))
         assert img.format.lower() == "jpeg"
         assert data["output_format"] == "jpeg"
-        assert data["size"] == "16x24"
+        assert data["size"] == f"{img.width}x{img.height}"
 
 
 def test_image_edit_layers_and_resolution(async_omni_test_client):

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -10,11 +10,13 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import piecewise_attn
 
 
 class CuDNNAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_prefix_kv_slicing: bool = True
+    supports_piecewise_spans: bool = True
 
     # cuDNN 9.5+ FMHA on Blackwell: head_dim divisible by 8 and at most 256
     # for BF16/FP16. Used by automatic platform selection; explicit CUDNN_ATTN
@@ -74,6 +76,23 @@ class CuDNNAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
+        if (
+            attn_metadata is not None
+            and attn_metadata.full_attn_spans is not None
+            and (attn_metadata.attn_mask is None or attn_metadata.attn_mask.ndim != 4)
+        ):
+            if attn_metadata.attn_mask is not None or "valid_kv_length" in attn_metadata.extra:
+                raise ValueError("Piecewise cuDNN attention cannot combine spans with a padding mask.")
+            return piecewise_attn(
+                query,
+                key,
+                value,
+                attn_metadata.full_attn_spans,
+                self.softmax_scale,
+                self._piecewise_attention,
+                query_ranges=attn_metadata.query_ranges,
+            )
+
         attention_mask = None
         if attn_metadata:
             valid_kv_length = attn_metadata.extra.get("valid_kv_length")
@@ -97,6 +116,19 @@ class CuDNNAttentionImpl(AttentionImpl):
                     mask_mode="broadcast_k",
                 )
 
+        return self._attention(query, key, value, attention_mask, self.causal, self.softmax_scale)
+
+    def _piecewise_attention(self, query, key, value, *, causal, softmax_scale):
+        mask = None
+        if causal:
+            # SDPA's causal flag aligns at the top left. Text segments need
+            # bottom-right alignment so every query can also see its prefix.
+            q_len, kv_len = query.shape[1], key.shape[1]
+            mask = torch.ones(q_len, kv_len, dtype=torch.bool, device=query.device).tril(kv_len - q_len)
+            mask = mask[None, None]
+        return self._attention(query, key, value, mask, False, softmax_scale)
+
+    def _attention(self, query, key, value, attention_mask, causal, softmax_scale):
         enable_gqa = query.shape[2] != key.shape[2]
         kv_seq_len = key.shape[1]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
@@ -130,8 +162,8 @@ class CuDNNAttentionImpl(AttentionImpl):
                 value,
                 attn_mask=attention_mask,
                 dropout_p=0.0,
-                is_causal=self.causal,
-                scale=self.softmax_scale,
+                is_causal=causal,
+                scale=softmax_scale,
                 enable_gqa=enable_gqa,
             )
         return output.permute(0, 2, 1, 3)

vllm_omni/diffusion/data.py

@@ -902,11 +902,6 @@ class OmniDiffusionConfig:
 
     # Compilation
     enforce_eager: bool = False
-    # Opt-in: capture fixed-shape KV-cache decode (denoising) steps into CUDA
-    # graphs. Currently implemented by Qwen-Image-2.1's transformer; unsupported
-    # configurations (SP/TP, ring, HSDP, dynamic LoRA, padded masks) log and
-    # fall back to eager decode. Other models ignore this flag.
-    enable_cuda_graph_decode: bool = False
     # Controls the generic compilation path used when a pipeline does not
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 
 import torch
@@ -46,7 +47,11 @@ class DistributedVaeExecutor:
 
     def __init__(self):
         # Use a dedicated process group spanning the complete worker WORLD.
-        self.group = get_world_group().device_group
+        world_group = get_world_group()
+        self.group = world_group.device_group
+        self._error_group = world_group.cpu_group
+        # Reporting an OOM must not require another device allocation.
+        self._error_status = torch.zeros((), dtype=torch.int32, device="cpu")
         self.world_size = dist.get_world_size(self.group)
         self.rank = dist.get_rank(self.group)
         self.parallel_size = 1
@@ -56,8 +61,34 @@ class DistributedVaeExecutor:
         self.parallel_size = parallel_size
         self.parallel_mode = mode
 
+    @contextmanager
+    def _sync_errors(self, stage: str):
+        """Agree on local-stage failures before entering the next collective."""
+        status = 0
+        message = ""
+        try:
+            yield
+        except torch.OutOfMemoryError as error:
+            status = 1
+            message = str(error)
+        except Exception as error:
+            # A non-retryable error takes precedence over OOM on another rank.
+            status = 2
+            message = f"{type(error).__name__}: {error}"
+
+        self._error_status.fill_(status)
+        dist.all_reduce(self._error_status, op=dist.ReduceOp.MAX, group=self._error_group)
+        global_status = self._error_status.item()
+        if global_status == 2:
+            detail = message if status == 2 else "another rank failed"
+            raise RuntimeError(f"Distributed VAE {stage} failed: {detail}")
+        if global_status == 1:
+            detail = message if status == 1 else "another rank ran out of memory"
+            raise torch.OutOfMemoryError(f"Distributed VAE {stage} failed: {detail}")
+
     def gather_tensors(self, tensor: torch.Tensor):
-        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
+        with self._sync_errors("gather allocation"):
+            gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
         dist.all_gather(gather_list, tensor, group=self.group)
         return gather_list if self.rank == 0 else None
 
@@ -66,11 +97,12 @@ class DistributedVaeExecutor:
         return tensor
 
     def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
-        local_tile_max_dims = [0] * output_ndim
-        for _, tile_tensor in local_results:
-            for dim_idx, dim_size in enumerate(tile_tensor.shape):
-                local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
-        local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
+        with self._sync_errors("shape preparation"):
+            local_tile_max_dims = [0] * output_ndim
+            for _, tile_tensor in local_results:
+                for dim_idx, dim_size in enumerate(tile_tensor.shape):
+                    local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
+            local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
         dist.all_reduce(
             local_shape_stat,
             op=dist.ReduceOp.MAX,
@@ -122,49 +154,51 @@ class DistributedVaeExecutor:
     def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
         pp_size = min(self.parallel_size, self.world_size)
 
-        # 1. Split into tiles
-        tiletask_list, grid_spec = operator.split(z)
-        tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
-
-        # 2. local decode
-        assigned = self._balance_tasks(tiletask_list, pp_size)
-        local_tasks = assigned[self.rank] if self.rank < pp_size else []
-        local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
+        with self._sync_errors("tile decode"):
+            tiletask_list, grid_spec = operator.split(z)
+            tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
+            assigned = self._balance_tasks(tiletask_list, pp_size)
+            local_tasks = assigned[self.rank] if self.rank < pp_size else []
+            local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
 
         # 3. compute shape per rank
         global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
 
         # 5. prepare tile tensors
         output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
-        local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
-            local_results, global_padding_shape, grid_spec, z.device, output_dtype
-        )
+        with self._sync_errors("tile packing"):
+            local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
+                local_results, global_padding_shape, grid_spec, z.device, output_dtype
+            )
 
         # 6. gather tiles & meta
         meta_gather = self.gather_tensors(local_meta_tensor)
         tile_gather = self.gather_tensors(local_tile_tensor)
 
-        if self.rank != 0:
-            result = torch.empty(0, device=z.device)  # Dummy return for non-zero ranks
-        else:
-            # 7. reconstruct full tensor (rank 0)
-            coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
-            result = operator.merge(coord_tensor_map, grid_spec)
+        with self._sync_errors("tile merge"):
+            if self.rank != 0:
+                result = torch.empty(0, device=z.device)  # Dummy return for non-zero ranks
+            else:
+                # 7. reconstruct full tensor (rank 0)
+                coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+                result = operator.merge(coord_tensor_map, grid_spec)
 
         if broadcast_result:
             result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
         return result
 
     def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
-        shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
-        if self.rank == 0:
-            shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
+        with self._sync_errors("output shape preparation"):
+            shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
+            if self.rank == 0:
+                shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
         dist.broadcast(shape_tensor, src=0, group=self.group)
 
-        if self.rank != 0:
-            sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
-        else:
-            sync_result = rank0_result
+        with self._sync_errors("output allocation"):
+            if self.rank != 0:
+                sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
+            else:
+                sync_result = rank0_result
         dist.broadcast(sync_result, src=0, group=self.group)
         return sync_result
 

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -9,19 +9,14 @@ default cache stores prefix K/V in freshly allocated tensors, so a captured
 graph would bind addresses that the next request (or a step-mode cache merge)
 no longer owns.
 
-This module keeps the prefix K/V in persistent, fixed-address buffers. Prefill
-copies its K/V into them and re-points the caller's cache dict at the buffers;
-decode then replays a captured graph that reads those buffers plus static input
-buffers (target latents, timestep, RoPE frequencies). One graph is captured per
-``(branch, batch, prefix_len, target_tokens, dtype, device, backend)`` key and
-reused across requests: a new prefill with the same key rewrites the buffers in
-place, which a captured graph reads correctly because it binds addresses, not
-contents.
+Requests retain ownership of their prefix K/V. Decode refreshes separate
+fixed-address graph buffers when the active prefix tensors change. Graphs
+are keyed by branch, batch, prefix length, complete image layout,
+dtype, device and backend, so equal token counts with different RoPE layouts stay separate.
 
-Memory note: only the *prefix* K/V are persistent. The decode attention still
-concatenates prefix K/V with the current step's target K/V inside the captured
-region, and that transient tensor is served by the graph's memory pool, so the
-persistent footprint is the same as the existing dict cache.
+Memory note: graph entries own an additional copy of the prefix K/V. Decode
+attention concatenates that prefix with target K/V inside the captured region;
+transient tensors use the graph memory pool.
 
 Fallbacks (all logged, all eager): CUDA unavailable, sequence/ring/tensor
 parallelism, HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization,
@@ -32,6 +27,7 @@ capture failure.
 
 from __future__ import annotations
 
+import weakref
 from collections import OrderedDict
 from typing import TYPE_CHECKING
 
@@ -39,6 +35,8 @@ import torch
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 
+from vllm_omni.platforms import current_omni_platform
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
         QwenImage21Transformer2DModel,
@@ -67,8 +65,8 @@ class QwenImage21DecodeGraphEntry:
 
     Buffers are created outside inference mode so replay-side ``copy_`` is
     legal regardless of the caller's grad context. ``graph`` is None until the
-    first decode step captures it; ``failed`` disables capture permanently for
-    this key after an exception.
+    first decode step captures it. Failed entries are replaced by a ``None``
+    marker in the manager so their buffers can be released.
     """
 
     def __init__(
@@ -99,15 +97,30 @@ class QwenImage21DecodeGraphEntry:
             # inference_mode, and the captured body takes views of this tensor
             # outside it, which is forbidden for inference tensors.
             self.freqs = freqs.clone()
-        # Decode needs no mask once padding is excluded (see the manager's
-        # prefill gate), so the captured body attends without metadata.
-        self.attn_metadata = None
+            self.attn_metadata = None
         # Per-block cache dicts with the same shape as the legacy protocol.
         self.block_caches = [{branch: {"key": self.k[i], "value": self.v[i]}} for i in range(len(kv_shapes))]
         self.graph: torch.cuda.CUDAGraph | None = None
         self.output: torch.Tensor | None = None
-        self.failed = False
         self.captures = 0
+        self.prefix_sources: list[tuple[weakref.ReferenceType[torch.Tensor], int | None]] = []
+
+    def refresh_prefix(self, kv_cache: list[dict[str, dict[str, torch.Tensor]]]) -> None:
+        sources = [block[self.branch][name] for block in kv_cache for name in ("key", "value")]
+        versions = [None if tensor.is_inference() else tensor._version for tensor in sources]
+        if len(sources) == len(self.prefix_sources) and all(
+            previous() is tensor and version == current_version
+            for (previous, version), tensor, current_version in zip(self.prefix_sources, sources, versions)
+        ):
+            return
+        # Prefill replaces inference-mode tensors; decode treats them as immutable.
+        # Weak references avoid retaining completed requests or mistaking a reused
+        # allocation for the same prefix. Normal tensors also track in-place writes.
+        self.prefix_sources = []
+        for i, block in enumerate(kv_cache):
+            self.k[i].copy_(block[self.branch]["key"])
+            self.v[i].copy_(block[self.branch]["value"])
+        self.prefix_sources = [(weakref.ref(tensor), version) for tensor, version in zip(sources, versions)]
 
     def capture(self, body) -> None:
         """Warm up on a side stream, then capture one decode step.
@@ -119,10 +132,12 @@ class QwenImage21DecodeGraphEntry:
         """
         side = torch.cuda.Stream()
         side.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(side), torch.inference_mode(False), torch.no_grad():
-            for _ in range(2):
-                body()
-        torch.cuda.current_stream().wait_stream(side)
+        try:
+            with torch.cuda.stream(side), torch.inference_mode(False), torch.no_grad():
+                for _ in range(2):
+                    body()
+        finally:
+            torch.cuda.current_stream().wait_stream(side)
 
         graph = torch.cuda.CUDAGraph()
         # The platform-wide pool is shared with every other captured path in
@@ -142,7 +157,7 @@ class QwenImage21DecodeGraphManager:
     def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8):
         self.model = model
         self.max_entries = max_entries
-        self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry] = OrderedDict()
+        self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict()
         self._static_eligible: bool | None = None
 
     # ── eligibility ──
@@ -202,17 +217,12 @@ class QwenImage21DecodeGraphManager:
         kv_cache: list[dict[str, dict[str, torch.Tensor]]],
         cache_branch: str,
         prefix_len: int,
-        target_tokens: int,
+        img_shapes: list[tuple[int, int, int]],
         target_freqs: torch.Tensor,
         joint_key_valid: torch.Tensor | None,
         dtype: torch.dtype,
     ) -> None:
-        """Move a freshly prefilled prefix K/V into static buffers.
-
-        The caller's ``kv_cache`` dicts are re-pointed at the static buffers, so
-        the existing decode path reads fixed addresses and a captured graph can
-        be replayed for later steps. Prefill itself always runs eagerly.
-        """
+        """Prepare graph buffers without changing the request's prefix K/V."""
         if not self.eligible() or torch.compiler.is_compiling():
             return
         if _dynamic_lora_wrappers_present(self.model):
@@ -239,44 +249,55 @@ class QwenImage21DecodeGraphManager:
             cache_branch,
             batch_size,
             prefix_len,
-            target_tokens,
+            tuple(tuple(shape) for shape in img_shapes),
             dtype,
             first["key"].device.index,
             self._backend_name(),
         )
 
-        entry = self.entries.get(key)
-        if entry is None:
+        if key not in self.entries:
             while len(self.entries) >= self.max_entries:
-                evicted_key, _ = self.entries.popitem(last=False)
+                evicted_key, evicted_entry = self.entries.popitem(last=False)
+                del evicted_entry
                 logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries)
-            entry = QwenImage21DecodeGraphEntry(
-                branch=cache_branch,
-                batch_size=batch_size,
-                prefix_len=prefix_len,
-                target_tokens=target_tokens,
-                in_channels=self.model.in_channels,
-                dtype=dtype,
-                device=first["key"].device,
-                kv_shapes=[block_cache[cache_branch]["key"].shape for block_cache in kv_cache],
-                kv_dtype=first["key"].dtype,
-                freqs=target_freqs,
-            )
-            self.entries[key] = entry
+            allocation_error = None
+            try:
+                self.entries[key] = QwenImage21DecodeGraphEntry(
+                    branch=cache_branch,
+                    batch_size=batch_size,
+                    prefix_len=prefix_len,
+                    target_tokens=target_freqs.shape[0],
+                    in_channels=self.model.in_channels,
+                    dtype=dtype,
+                    device=first["key"].device,
+                    kv_shapes=[block_cache[cache_branch]["key"].shape for block_cache in kv_cache],
+                    kv_dtype=first["key"].dtype,
+                    freqs=target_freqs,
+                )
+            except torch.OutOfMemoryError as exc:
+                allocation_error = str(exc)
+            if allocation_error is not None:
+                # Keep a bounded failure marker without retaining tensors or the exception traceback.
+                self.entries[key] = None
+                current_omni_platform.empty_cache()
+                logger.warning(
+                    "Qwen-Image-2.1 CUDA graph allocation failed for key=%s: %s. Falling back to eager decode.",
+                    key,
+                    allocation_error,
+                )
+                return
         else:
             self.entries.move_to_end(key)
 
-        for i, block_cache in enumerate(kv_cache):
-            parts = block_cache[cache_branch]
-            entry.k[i].copy_(parts["key"])
-            entry.v[i].copy_(parts["value"])
-            parts["key"] = entry.k[i]
-            parts["value"] = entry.v[i]
+        entry = self.entries[key]
+        if entry is not None:
+            entry.prefix_sources = []
+
         logger.debug(
             "Registered decode graph entry key=%s (entries=%d, captured=%d)",
             key,
             len(self.entries),
-            sum(1 for e in self.entries.values() if e.graph is not None),
+            sum(1 for e in self.entries.values() if e is not None and e.graph is not None),
         )
 
     # ── decode replay ──
@@ -288,11 +309,20 @@ class QwenImage21DecodeGraphManager:
         timestep: torch.Tensor,
         kv_cache: list[dict[str, dict[str, torch.Tensor]]],
         cache_branch: str,
-        target_tokens: int,
+        img_shapes: list[tuple[int, int, int]],
+        img_mask: torch.Tensor,
+        encoder_hidden_states_mask: torch.Tensor | None,
     ) -> torch.Tensor | None:
         """Replay the captured decode step, or return None to fall back to eager."""
         if not self.eligible() or torch.compiler.is_compiling():
             return None
+        if _dynamic_lora_wrappers_present(self.model):
+            self.clear()
+            return None
+        if encoder_hidden_states_mask is not None:
+            text_positions = ~img_mask[0, : encoder_hidden_states_mask.shape[1]]
+            if not bool(encoder_hidden_states_mask[:, text_positions].all()):
+                return None
         first = kv_cache[0][cache_branch]
         cached_key = first["key"]
         if cached_key.device.type != "cuda":
@@ -304,40 +334,36 @@ class QwenImage21DecodeGraphManager:
             cache_branch,
             cached_key.shape[0],
             cached_key.shape[1],
-            target_tokens,
+            tuple(tuple(shape) for shape in img_shapes),
             hidden_states.dtype,
             cached_key.device.index,
             self._backend_name(),
         )
         entry = self.entries.get(key)
-        if entry is None or entry.failed:
+        if entry is None:
             return None
         self.entries.move_to_end(key)
 
-        entry.hidden.copy_(hidden_states[:, -target_tokens:])
+        entry.hidden.copy_(hidden_states[:, -entry.target_tokens :])
         entry.timestep.copy_(timestep)
-        # The caller's cache dicts normally already reference the entry's
-        # buffers (prefill re-pointed them). A step-mode batched decode instead
-        # hands over freshly merged tensors; copy those into the static buffers.
-        for i, block_cache in enumerate(kv_cache):
-            parts = block_cache[cache_branch]
-            if parts["key"] is not entry.k[i]:
-                entry.k[i].copy_(parts["key"])
-                entry.v[i].copy_(parts["value"])
-                parts["key"] = entry.k[i]
-                parts["value"] = entry.v[i]
+        # Graph buffers are scratch space; request caches must never alias them.
+        entry.refresh_prefix(kv_cache)
 
         if entry.graph is None:
+            capture_error = None
             try:
                 entry.capture(lambda: self.model._decode_graph_forward(entry))
             except Exception as exc:
-                entry.failed = True
-                entry.graph = None
+                capture_error = str(exc)
+            if capture_error is not None:
+                self.entries[key] = None
+                del entry
+                current_omni_platform.empty_cache()
                 logger.warning(
                     "Qwen-Image-2.1 CUDA graph capture failed for key=%s: %s. "
                     "Falling back to eager decode for this shape.",
                     key,
-                    exc,
+                    capture_error,
                 )
                 return None
             logger.info("Captured Qwen-Image-2.1 decode graph for key=%s", key)

vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py

@@ -144,7 +144,7 @@ def get_qwen_image_21_pre_process_func(
         if "additional_information" not in prompt:
             prompt["additional_information"] = {}
 
-        if not raw_image:  # None or empty list: pure text-to-image request
+        if raw_image is None or (isinstance(raw_image, list) and len(raw_image) == 0):
             request.batch_compatibility_key = ("qwen_image_21", ())
             request.prompt = prompt
             return request
@@ -744,6 +744,7 @@ class QwenImage21Pipeline(
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,
             num_inference_steps,
+            device=self.device,
             sigmas=sigmas,
             mu=mu,
         )

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -13,6 +13,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import RMSNorm
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -107,7 +108,10 @@ class QwenImage21TemporalTimesteps(nn.Module):
         self.time_factor = time_factor
 
         half = timestep_dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
+        # Compute the fixed table on CPU like Diffusers, then honor the loader's device.
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device="cpu") / half
+        ).to(torch.get_default_device())
         self.register_buffer("freqs", freqs, persistent=False)
 
     def forward(self, timestep: torch.Tensor) -> torch.Tensor:
@@ -294,12 +298,15 @@ class QwenImage21Rope(nn.Module):
         self.theta = theta
         self.axes_dim = axes_dim
 
-        pos_index = torch.arange(8192)
-        neg_index = torch.arange(1024).flip(0) * -1 - 1
-        self.freqs = [
-            torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0)
-            for dim in axes_dim
-        ]
+        # Match Diffusers' CPU initialization even when the loader sets a CUDA
+        # default device; CPU and CUDA trigonometric kernels round differently.
+        with torch.device("cpu"):
+            pos_index = torch.arange(8192)
+            neg_index = torch.arange(1024).flip(0) * -1 - 1
+            self.freqs = [
+                torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0)
+                for dim in axes_dim
+            ]
 
     def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor:
         freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)))
@@ -374,9 +381,10 @@ class QwenImage21Attention(nn.Module):
         self.num_heads = self.to_qkv.num_heads
         self.num_kv_heads = self.to_qkv.num_kv_heads
 
-        # Per-head RMSNorm over the head dim; TP shards whole heads, so the norm is TP-safe.
-        self.norm_q = nn.RMSNorm(dim_head, eps=eps)
-        self.norm_k = nn.RMSNorm(dim_head, eps=eps)
+        # Diffusers normalizes in FP32, then casts before the learned scale.
+        # Keep that rounding order for BF16 Q/K; TP shards whole heads.
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
 
         self.to_out = RowParallelLinear(
             heads * dim_head,
@@ -718,7 +726,6 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         causal_block: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "transformer",
-        enable_cuda_graph_decode: bool | None = None,
         cuda_graph_max_decode_graphs: int = 8,
     ):
         super().__init__()
@@ -790,14 +797,11 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         # Module boundary where _sp_plan shards the target tokens + their RoPE freqs.
         self.sequence_prepare = QwenImage21SequencePrepare(self.img_in, self.pos_embed)
 
-        # Opt-in CUDA Graph capture of the fixed-shape KV-cache decode steps.
-        # Default off: ``None`` defers to ``od_config.enable_cuda_graph_decode``.
-        if enable_cuda_graph_decode is None:
-            enable_cuda_graph_decode = bool(getattr(od_config, "enable_cuda_graph_decode", False))
-        self.enable_cuda_graph_decode = enable_cuda_graph_decode
+        # Use decode graphs unless the caller requests eager execution.
+        self.enable_cuda_graph_decode = not od_config.enforce_eager
         self._decode_graph_manager = (
             QwenImage21DecodeGraphManager(self, max_entries=cuda_graph_max_decode_graphs)
-            if enable_cuda_graph_decode
+            if self.enable_cuda_graph_decode
             else None
         )
 
@@ -972,7 +976,9 @@ class QwenImage21Transformer2DModel(CachedTransformer):
                 timestep=timestep,
                 kv_cache=kv_cache,
                 cache_branch=cache_branch,
-                target_tokens=math.prod(layout[-1]),
+                img_shapes=layout,
+                img_mask=img_mask,
+                encoder_hidden_states_mask=encoder_hidden_states_mask,
             )
             if graph_output is not None:
                 if not return_dict:
@@ -1039,6 +1045,9 @@ class QwenImage21Transformer2DModel(CachedTransformer):
             joint_key_valid = F.pad(joint_key_valid, (0, sp_padding), value=False)
 
         has_padding = joint_key_valid is not None and not bool(joint_key_valid.all())
+        if not has_padding:
+            # Match Diffusers and use the same mask-free dispatch in eager and graph decode.
+            joint_key_valid = None
         attn_metadata: AttentionMetadata | None = None
         cache_write_len: int | None = None
         if is_decode:
@@ -1082,13 +1091,12 @@ class QwenImage21Transformer2DModel(CachedTransformer):
             and kv_cache is not None
             and self._decode_graph_manager is not None
         ):
-            # Move the freshly written prefix K/V into fixed-address buffers so
-            # decode steps can replay a captured CUDA graph against them.
+            # Allocate graph scratch buffers; the request keeps its own prefix K/V.
             self._decode_graph_manager.register_prefill(
                 kv_cache=kv_cache,
                 cache_branch=cache_branch,
                 prefix_len=cache_write_len,
-                target_tokens=target_freqs.shape[0],
+                img_shapes=layout,
                 target_freqs=target_freqs,
                 joint_key_valid=joint_key_valid,
                 dtype=hidden_states.dtype,
@@ -1113,9 +1121,9 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         """The exact decode-step computation, captured into a CUDA graph.
 
         Mirrors the eager decode path of ``forward``: target tokens only,
-        prefix K/V read from the entry's static buffers, no attention metadata
-        (the registration gate excludes padded masks), and an all-ones
-        modulation mask. All inputs live in the entry's fixed-address buffers.
+        prefix K/V read from the entry's static buffers, the same mask-free
+        attention dispatch as eager decode, and an all-ones modulation mask.
+        All inputs live in the entry's fixed-address buffers.
         """
         hidden_states = self.img_in(entry.hidden)
         timestep = torch.cat([entry.timestep, entry.timestep.new_zeros(1)], dim=0)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -231,6 +231,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if model is None:
             return
 
+        if getattr(model, "enable_cuda_graph_decode", False):
+            logger.info("Model runner: %s uses CUDA graph decode instead of torch.compile.", attr_name)
+            return
+
         compile_granularity = self.od_config.diffusion_compile_granularity
         compile_dynamic = self.od_config.diffusion_compile_dynamic
         try:

vllm_omni/entrypoints/openai/api_server.py

@@ -1605,8 +1605,8 @@ def _build_image_generation_response(
             peak_memory_mb=peak_memory_mb,
         ),
     }
-    if request.size is not None:
-        response_kwargs["size"] = request.size
+    if images:
+        response_kwargs["size"] = f"{images[0].width}x{images[0].height}"
     response = ImageGenerationResponse(**response_kwargs)
     if request.response_format == ResponseFormat.FILE:
         return response.stream_response()
@@ -1839,6 +1839,7 @@ async def generate_images(
 
 @router.post(
     "/v1/images/edits",
+    response_model=None,
     responses={
         HTTPStatus.OK.value: {"model": ImageGenerationResponse},
         HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
@@ -1846,6 +1847,7 @@ async def generate_images(
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
     },
 )
+@with_cancellation
 async def edit_images(
     raw_request: Request,
     image: list[UploadFile] | None = File(None),
@@ -2181,7 +2183,7 @@ async def edit_images(
             created=int(time.time()),
             data=image_data,
             output_format=output_format,
-            size=size_str,
+            size=f"{images[0].width}x{images[0].height}" if images else size_str,
             cot_output=cot_output,
             metrics=_build_image_response_metrics(
                 response_metrics=response_metrics,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -3466,7 +3466,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     chunk = ImageEditImageChunk(
                         data=image_data,
                         output_format=output_format,
-                        size=size,
+                        size=f"{images[0].width}x{images[0].height}",
                         created=created,
                         model=model,
                         metrics=metrics,