[Doc] Add Qwen-Image-2.1 recipe to recipes index (#8070)

Signed-off-by: SamitHuang <285365963@qq.com>

.buildkite/cuda/test-merge.yml

@@ -133,6 +133,18 @@ steps:
           - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py tests/e2e/online_serving/test_minicpmo_4_5_window.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: diffusion_qwen_image_function
         commands:

.buildkite/cuda/test-ready.yml

@@ -229,6 +229,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: diffusion_qwen_image_function
         commands:

docs/models/supported_models.md

@@ -32,6 +32,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImage21Pipeline` | Qwen-Image-2.1 (T2I + image-conditioned generation, up to 4 condition images) | `Qwen/Qwen-Image-2.1` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2.1.md) |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |

docs/user_guide/diffusion_features.md

@@ -139,6 +139,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Qwen-Image-2.1**       |     ❌     |     ❌      |   ✅ (Ulysses only)    |       ✅        |         ❓         |          ❌          |   ❓    |             ❓             |      ✅ (decode)      |       ✅        |        ✅         |
 | **SenseNova-U1 / U1.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
@@ -155,6 +156,7 @@ The following tables show which models support each feature:
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 > 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks.
+> 8. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends.
 
 ### VideoGen
 

docs/user_guide/quantization/fp8.md

@@ -1,4 +1,4 @@
-# FP8 Quantization
+# FP8 Quantization
 
 ## Overview
 
@@ -85,6 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
 |-------|-----------|:-------:|:------:|----------------|------------------|------------------|
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
+| Qwen-Image-2.1 | `Qwen/Qwen-Image-2.1` | Yes | Not validated | Skip image-stream MLPs for best fidelity; all-layer FP8 is loadable but composition drifts on some prompts. On Blackwell GB200 FP8 saves ~17% peak memory but is not faster than BF16 at 1024x1024. Text encoder FP8 covers the language model only — the vision tower and LM head stay BF16 | `img_mlp` | ✅︎ (language model only) |
 | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
 | LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
 | LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |

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.
@@ -209,14 +209,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(
@@ -368,7 +368,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
@@ -647,6 +647,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/run_qwen_image_21.sh

@@ -0,0 +1,16 @@
+#!/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'" \
+    --output qwen_image_21_edit.png \
+    --num-inference-steps 50 \
+    --cfg-scale 1.0 \
+    "$@"

examples/offline_inference/text_to_image/text_to_image.py

@@ -89,7 +89,8 @@ def parse_args() -> argparse.Namespace:
         "--model",
         default="Qwen/Qwen-Image",
         help="Diffusion model name or local path. Supported models: "
-        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
+        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-2.1, "
+        "stepfun-ai/NextStep-1.1, "
         "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
         "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
         "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "

recipes/Qwen/Qwen-Image-2.1.md

@@ -0,0 +1,398 @@
+# Qwen-Image-2.1
+
+> Text-to-image and image-conditioned generation with Qwen-Image 2.1
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen-Image-2.1`
+- Task: Text-to-image generation and image-conditioned generation (editing)
+- Mode: Offline inference and online serving with optional step-wise execution
+- Maintainer: Community
+
+> **Model id:** `Qwen/Qwen-Image-2.1` follows the reference diffusers
+> documentation and may change before the official release — the final
+> checkpoint name is subject to the official release.
+
+## When to use this recipe
+
+Use this recipe as a starting point for running `Qwen/Qwen-Image-2.1`. A single
+pipeline class, `QwenImage21Pipeline`, handles both pure text-to-image requests
+(no condition image) and image-conditioned requests with up to 4 condition
+images: prompt and condition images are encoded together by a Qwen3-VL text
+encoder, so editing-style requests go through the same serving path as
+text-to-image.
+
+## References
+
+- 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/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:
+  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
+- Sibling recipes: [`Qwen-Image.md`](./Qwen-Image.md),
+  [`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md)
+
+## Offline Inference
+
+### Text-to-image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --output qwen_image_21_t2i.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+### Image-conditioned generation (editing)
+
+```bash
+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" \
+  --output qwen_image_21_edit.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+or via the bundled launcher:
+
+```bash
+bash examples/offline_inference/image_to_image/run_qwen_image_21.sh
+```
+
+Up to 4 condition images can be passed via `--image`:
+
+```bash
+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" \
+  --output qwen_image_21_multi.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.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:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --port 8091
+```
+
+To enable the step-wise runtime:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni \
+  --port 8091 \
+  --step-execution \
+  --max-num-seqs 8
+```
+
+Step execution batches compatible requests at the same KV-cache phase. The
+pipeline tags every request with a `batch_compatibility_key` (conditioning
+mode plus condition-image layout), so a newly arriving request that would
+otherwise join a batch already past its first denoising step is isolated by
+the scheduler and deferred until the in-flight batch drains — it no longer
+fails the whole batch. `--max-num-seqs` controls batch capacity.
+
+### CUDA Graph decode
+
+Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising
+decode steps when `enable_cuda_graph_decode=True` (the default) and
+`enforce_eager=False`. Prefill remains eager. This works with both request
+execution and `--step-execution`.
+
+To disable graph capture while keeping the rest of the compilation setup:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode
+```
+
+The same control is available as `Omni(..., enable_cuda_graph_decode=False)`
+and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML.
+`--enforce-eager` / `enforce_eager: true` remains the broader switch: it
+disables both graph capture and automatic `torch.compile`. The two
+optimizations stack: the DiT blocks are regionally `torch.compile`'d and
+graph capture records the compiled (fused) kernels, while inductor's own
+cudagraphs stay disabled so the two graph layers never nest. Measured on a
+single GB200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`,
+warmup 1 + median of 3): compile+graph 3.04 s, compile-only 3.2 s,
+graph-only 4.0 s, eager 4.3 s end-to-end per image; combo output vs eager is
+44.5–45.6 dB PSNR (the same magnitude as pure `torch.compile` fusion
+divergence), and graph-only vs eager is bit-identical.
+
+The autoregressive engine's `compilation_config.cudagraph_mode` does not
+control this diffusion path; use `enable_cuda_graph_decode` (or
+`enforce_eager` for everything) to force eager decode execution.
+
+TP/SP/ring parallelism, HSDP, offload/cache hooks, quantized KV caches,
+dynamic LoRA, padded text masks, and a second in-flight request whose cache
+aliases an already-owned graph key 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; each graph key
+has a single live owner at a time, so concurrent same-shape requests never
+overwrite each other's prefix. Those buffers require additional memory; graph
+entries are bounded by the transformer's cache limit.
+
+### Verification
+
+For a direct API smoke test:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen-Image-2.1",
+    "prompt": "A ceramic teapot on a wooden table",
+    "size": "1024x1024",
+    "num_inference_steps": 50,
+    "true_cfg_scale": 1.0,
+    "seed": 42
+  }'
+```
+
+For image-conditioned generation, pass condition images through the
+OpenAI-compatible multimodal endpoint (`/v1/chat/completions` with
+`image_url` content parts, or `/v1/images/edits`), same as
+[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md).
+
+## Parameters
+
+| Parameter | Default | Notes |
+| --- | --- | --- |
+| `true_cfg_scale` | 1.0 | Disables the negative-prompt branch by default. Set above 1 with a `negative_prompt` to enable true CFG. Qwen-Image 2.1 supports true CFG only — there is no `guidance_scale` parameter. |
+| `num_inference_steps` | 50 | Number of denoising steps. |
+| `height` / `width` | 1024×1024 | Must be multiples of 32. For image-conditioned requests, leaving them unset derives the output size from the last condition image's aspect ratio at ~1024×1024. |
+| `negative_prompt` | None | Ignored at the default CFG scale of 1.0; used only when `true_cfg_scale > 1`. |
+| condition images | — | Up to 4 input images per request; more than 4 raises an error. |
+| `seed` | — | Fix for reproducible outputs. |
+
+The default CFG scale of 1.0 runs only the positive-prompt branch at each
+denoising step. To enable negative-prompt guidance, explicitly set
+`--cfg-scale` (offline) or `true_cfg_scale` (API) above 1 and provide a
+negative prompt. This changes the generated output as well as compute cost.
+
+### Prefix KV cache
+
+The transformer caches the text/condition-image prefix across denoising steps
+(a prefill/decode split keyed by CFG branch). The cache is enabled
+automatically when the checkpoint declares `causal_condition: true` and is
+re-created per generation, so no state leaks between requests. It changes the
+compute path — cached prefix vs. full recompute at every step — so exact
+numerical reproduction against a no-cache reference (e.g. the diffusers
+pipeline) requires matching the cache setting on both sides. The prefix KV
+cache is orthogonal to weight quantization and works unchanged under FP8.
+
+### FP8 quantization
+
+Online FP8 quantization of the transformer is supported; the fused
+`to_qkv` projections load correctly, and the prefix KV cache path is
+unaffected. Only block-internal linears are eligible — the boundary
+projections (`img_in`/`txt_in`/`modulation`/`proj_out`/`time_text_embed`) stay
+in BF16 by construction:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --output qwen_image_21_fp8.png \
+  --num-inference-steps 50 --cfg-scale 1.0 \
+  --quantization fp8 --ignored-layers "img_mlp"
+```
+
+`ignored_layers` entries are matched as name patterns for this model (e.g.
+`img_mlp` matches every block's `img_mlp.proj`/`img_mlp.gate_layer`/
+`img_mlp.out`); exact full prefixes like
+`transformer.transformer_blocks.3.attn.to_qkv` also work.
+
+A per-module sensitivity sweep (fixed prompt embeds / latent / timestep, 5
+timesteps across the trajectory, relative L2 of the noise prediction vs BF16)
+ranks the groups: `img_mlp` is the most sensitive (3.8% max error on its own,
+driven by `img_mlp.proj` 2.8% and `img_mlp.out` 2.5%), then `attn.to_out`
+(2.7%), then `attn.to_qkv` (1.6%); errors are largest at the late (low-sigma)
+steps. End-to-end on 4 fixed prompts at 1024x1024, 50 steps, seed 42
+(GB200, single GPU):
+
+| Config | Quantized block linears | Avg PSNR vs BF16 | Min PSNR | Steady-state time/image | Peak memory |
+| --- | --- | --- | --- | --- | --- |
+| BF16 | 0 / 160 | — | — | 7.2 s | 40.0 GB |
+| FP8 all layers | 160 / 160 | 26.1 dB | 19.1 dB | 8.4 s | 33.4 GB |
+| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | 29.7 dB | 21.6 dB | 8.8 s | 38.3 GB |
+
+All-layer FP8 is usable but its composition can drift on some prompts
+(fine detail and layout shift, worst case ~19 dB); keeping `img_mlp` in BF16
+preserves composition noticeably better and is the recommended setting. On
+Blackwell (GB200) FP8 here is a memory optimization, not a speedup: dynamic
+per-token activation quantization costs more than the FP8 GEMM saves at these
+shapes, so BF16 remains the fastest option.
+
+## Quantization
+
+Online FP8 is supported per component via `quantization_config`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image-2.1",
+    quantization_config={
+        # DiT: skip the sensitive image-stream MLPs if quality regresses.
+        "transformer": {"method": "fp8", "ignored_layers": ["img_mlp"]},
+        # Qwen3-VL text encoder: only the language-model linear layers are
+        # quantized; the vision tower and the (unused) LM head stay BF16.
+        "text_encoder": {"method": "fp8"},
+    },
+)
+```
+
+Either component can be quantized on its own. Measured on GB200 at 1024×1024
+(seed 42, 50 steps): text-encoder FP8 lowers peak GPU memory from ~41.0 GiB to
+~34.4 GiB (−6.6 GiB); adding DiT FP8 reaches ~28.1 GiB. Output quality stays
+close to BF16 (T2I PSNR vs BF16 ≈ 28–32 dB with text-encoder FP8; the
+image-conditioned edit path ≈ 36 dB). FP8 saves memory but does not speed up
+generation on this hardware. The edit path routes condition images through the
+BF16 vision tower, so it is unaffected by text-encoder FP8. See
+[`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md)
+for the scope rules.
+
+### FP8 prefix KV storage
+
+The cached prefix K/V can be stored in FP8 E4M3 (with per-token-per-head fp32
+scales), halving the prefix-cache memory — relevant for long prompts (up to
+8192 text tokens) and up to 4 condition images (~4096 latent tokens each).
+Prefill quantizes once; decode dequantizes per step, so attention compute is
+unchanged. Enable it through the model-specific `extras` config (e.g. in the
+deploy YAML's stage `extras` or as an `Omni(...)` engine kwarg):
+
+```yaml
+extras:
+  prefix_kv_cache_dtype: "fp8"
+```
+
+Accepted values: `"fp8"` / `"fp8_e4m3"` (K and V quantized), `"fp8_v"` (V
+only — K stays in the native dtype), and `None` / `"auto"` (default, native
+dtype — behavior unchanged). This is independent of
+`diffusion_kv_cache_dtype`, which quantizes attention Q/K/V *compute* per
+forward pass on supported backends.
+
+Measured on T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0, PSNR vs. the
+bf16 baseline): **`"fp8"` 34.9 dB**, **`"fp8_v"` 40.9 dB** — same composition
+and semantics, with texture-level drift in fine detail for `"fp8"`, and
+visually indistinguishable output for `"fp8_v"`. The error is dominated by
+K quantization: post-RoPE keys are the precision-sensitive half of the
+cache, so `"fp8_v"` buys back ~6 dB at 75% (instead of 50%) of the original
+cache size. The residual error is inherent e4m3 precision accumulated
+coherently over the denoising trajectory; finer scale granularity
+(per-tensor/per-head/per-token were compared) or Hadamard-rotated V did not
+improve it. The saving scales with prefix length — at the limit
+(8192 text tokens + 4 condition images, ~24.6k prefix tokens) the prefix
+cache is ~12.9 GB per CFG branch in bf16, ~6.6 GB in `"fp8"` and ~9.7 GB in
+`"fp8_v"`, which is where this option matters. Treat it as an opt-in for
+memory-bound long-prompt / multi-image workloads, not a free lunch;
+`"fp8_v"` is the better default trade-off when quality matters.
+
+Note: a quantized prefix cache is not CUDA-graph capturable — requests using
+this option stay on the eager decode path (a warning is logged once).
+
+### Why not the scheduler-managed paged KV (`DiffusionKVCacheMode.PAGED_SCHEDULER`)?
+
+Evaluated and deliberately not adopted (see the `qwen21-p0-kv` commit message
+for the full analysis): the paged scheduler targets *cross-request* KV
+management — a statically sized page pool, block sharing/dedup, and eviction —
+whereas this prefix cache lives exactly as long as one generation and is
+sized once per request. Migrating would require routing decode attention
+through the worker paged adapter (writing fresh target K/V into pages every
+step), per-request page tables, and re-validating the block-causal
+piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over
+the per-request dense tensors. The FP8 storage above captures the actual
+memory win at ~1% of the integration cost.
+
+## Quality Comparison (head vs diffusers)
+
+The reference is the diffusers Qwen-Image 2.1 pipeline
+([huggingface/diffusers#14804](https://github.com/huggingface/diffusers/pull/14804)).
+Both sides must pin the same checkpoint, inputs, seed, precision (BF16), and
+hardware.
+
+vLLM-Omni (this recipe):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --seed 42 \
+  --output qwen_image_21_vllm_omni.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+diffusers reference:
+
+```python
+import torch
+from diffusers import QwenImage21Pipeline  # diffusers PR #14804
+
+pipe = QwenImage21Pipeline.from_pretrained(
+    "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
+).to("cuda")
+image = pipe(
+    prompt="A ceramic teapot on a wooden table",
+    num_inference_steps=50,
+    true_cfg_scale=1.0,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+    height=1024,
+    width=1024,
+).images[0]
+image.save("qwen_image_21_diffusers.png")
+```
+
+| Metric | vLLM-Omni | diffusers (#14804) | Status |
+| --- | --- | --- | --- |
+| Output image (1024×1024, 50 steps, seed 42, BF16) | pending measurement | pending measurement | not yet run |
+| PSNR / LPIPS vs reference | pending measurement | — | not yet run |
+| Latency (time/image, same GPU) | pending measurement | pending measurement | not yet run |
+
+The prefix KV cache changes the compute path (cached prefix vs. full recompute
+at every step), so exact numerical parity is not expected; match the cache
+setting on both sides or compare with PSNR/LPIPS rather than exact pixels.
+
+The following capabilities are implemented but the PR-body checkboxes predate
+any attached measurements; treat them as unverified until the A/B above is
+filled in:
+
+- [ ] TP=2/4 output parity vs single-card — pending measurement.
+- [ ] VAE tiling vs non-tiled decode — pending measurement.
+- [x] CUDA graph decode vs eager — measured on GB200 (1024×1024, 50 steps,
+  seed 42, BF16): compile+graph 3.04 s/image vs eager 4.3 s/image (~1.4×),
+  44.5–45.6 dB PSNR vs eager; graph-only is bit-identical to eager. See
+  "CUDA Graph decode" above.
+
+## Known Limitations
+
+- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;
+  `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`.
+- Sequence parallelism supports Ulysses only; ring attention is not supported.
+- Only true CFG is exposed; there is no `guidance_scale` knob.
+- VAE tiling (`--vae-use-tiling`, also implied by `--vae-patch-parallel-size > 1`)
+  decodes in 512px tiles with 384px stride (32 latent pixels per tile at 16x
+  spatial compression). On CUDA OOM the tile size is automatically halved
+  (keeping 25% overlap) and decoding retried, down to 128px tiles. Tile sizes
+  are not yet exposed as engine arguments.

recipes/README.md

@@ -83,6 +83,7 @@ recipes/
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
 | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
+| [`Qwen/Qwen-Image-2.1.md`](./Qwen/Qwen-Image-2.1.md) | Text-to-image and image-conditioned generation (editing) with prefix KV cache and CUDA graph decode | 1x GPU with CPU offload / 1x GB200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |

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

@@ -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,292 @@
+# 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_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]]))
+    assert states[1].extra["kv_cache"][0]["cond"]["key"].is_contiguous()
+
+
+# ---------------------------------------------------------------------------
+# 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_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_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_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/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,97 @@ 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)
+
+    # Offload via a *different* module's hook (the cross-module swap path),
+    # then reload: the same fixed device storage must be reused so that
+    # CUDA-graph-captured weight pointers stay valid.
+    encoder_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+    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"
+
+    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

@@ -45,6 +45,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_multiproc_engine_concurrency.py

@@ -17,7 +17,7 @@ import zmq
 from vllm.v1.engine.exceptions import EngineDeadError
 
 import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
@@ -37,6 +37,38 @@ from vllm_omni.outputs import OmniRequestOutput
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _signal_worker_ready(_rank, _config, writer, _broadcast, wake_event, *_args):
+    assert mp.get_start_method() == "spawn"
+    wake_event.set()
+    writer.send({"status": "ready", "result_handle": "ready"})
+    writer.close()
+
+
+@pytest.mark.skipif("fork" not in mp.get_all_start_methods(), reason="Requires a fork default context")
+def test_executor_spawns_workers_with_compatible_events(mocker):
+    previous_method = mp.get_start_method(allow_none=True)
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = OmniDiffusionConfig(model="test-model", step_execution=True)
+    broadcast = mocker.Mock()
+    broadcast.export_handle.return_value = None
+    mocker.patch.object(executor, "_init_broadcast_queue", return_value=broadcast)
+    mocker.patch.object(executor, "_init_result_queue")
+    mocker.patch.object(executor, "_start_worker_monitor")
+    mocker.patch.object(WorkerProc, "worker_main", new=_signal_worker_ready)
+    try:
+        mp.set_start_method("fork", force=True)
+        executor._init_executor()
+        for process in executor._processes:
+            process.join(timeout=30)
+            assert process.exitcode == 0
+        assert all(event.is_set() for event in executor.wake_events)
+        assert mp.get_start_method() == "fork"
+    finally:
+        if getattr(executor, "_finalizer", None) is not None:
+            executor.shutdown()
+        mp.set_start_method(previous_method, force=True)
+
+
 # ───────────────────────────────────────────── helpers ─────────────────────
 
 

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

tests/model_extras/test_model_extras.py

@@ -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:

vllm_omni/config/omni_config.py

@@ -851,6 +851,10 @@ class _DiffusionConfigProjection:
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
+    # Mirrors OmniDiffusionConfig.enable_cuda_graph_decode; enforced by the
+    # model (Qwen-Image-2.1) together with enforce_eager, which lives in the
+    # shared engine config.
+    enable_cuda_graph_decode: bool = True
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False

vllm_omni/config/stage_config.py

@@ -462,6 +462,10 @@ class StageDeployConfig:
     # Diffusion execution, cache, and VAE behavior.
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
+    # today). Independent of compilation_config.cudagraph_mode;
+    # enforce_eager=True also disables it.
+    enable_cuda_graph_decode: bool | None = None
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None

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

@@ -16,7 +16,7 @@ import diffusers
 import huggingface_hub
 import torch
 from PIL import Image
-from pydantic import Field, model_validator
+from pydantic import model_validator
 from typing_extensions import Self
 from vllm.config.utils import config
 from vllm.logger import init_logger
@@ -955,6 +955,19 @@ class OmniDiffusionConfig:
 
     # Compilation
     enforce_eager: bool = False
+    # Capture fixed-shape KV-cache decode (denoising) steps into CUDA graphs.
+    # Currently only Qwen-Image-2.1's transformer implements this; other models
+    # ignore the flag. This is a per-model diffusion path, independent of the
+    # AR engine's ``compilation_config.cudagraph_mode`` (which does not apply
+    # to diffusion stages). It stacks with ``diffusion_compile_granularity``:
+    # the DiT blocks are still torch.compile'd and graph capture records the
+    # compiled (fused) kernels, while inductor's own cudagraphs stay off.
+    # ``enforce_eager=True`` forces eager decode and
+    # disables capture regardless of this flag; unsupported configurations
+    # (SP/TP, ring, HSDP, dynamic LoRA, padded masks, quantized prefix KV, or
+    # a second in-flight request aliasing the same graph key) log and fall
+    # back to eager decode.
+    enable_cuda_graph_decode: bool = True
     # Controls the generic compilation path used when a pipeline does not
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"
@@ -1099,7 +1112,7 @@ class OmniDiffusionConfig:
     request_batch_max_wait_ms: float = 0.0
 
     # Supplementary model specific parameters
-    extras: dict[str, Any] = Field(default_factory=dict)
+    extras: dict[str, Any] = field(default_factory=dict)
 
     @property
     def is_moe(self) -> bool:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+    _cudnn_deterministic,
+    _unpatchify,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        # mostly copy from AutoencoderKLQwenImage21.tiled_decode
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            sample_height = sample_height // self.config.patch_size
+            sample_width = sample_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+            blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
+        else:
+            blend_height = self.tile_sample_min_height - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width - tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        tiletask_list = []
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                time_list = []
+                for k in range(num_frames):
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": blend_height,
+            "blend_width": blend_width,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        self.clear_cache()
+        time = []
+        with _cudnn_deterministic():
+            for k in range(len(task.tensor)):
+                self._conv_idx = [0]
+                tile = self.post_quant_conv(task.tensor[k])
+                decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+                time.append(decoded)
+        result = torch.cat(time, dim=2)
+        return result
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+        grid_h, grid_w = grid_spec.grid_shape
+        self.clear_cache()
+
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+
+        if self.config.patch_size is not None:
+            dec = _unpatchify(dec, patch_size=self.config.patch_size)
+
+        dec = torch.clamp(dec, min=-1.0, max=1.0)
+        return dec
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, return_dict=return_dict)
+
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

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/executor/multiproc_executor.py

@@ -14,6 +14,7 @@ import time
 import weakref
 from collections.abc import Callable
 from dataclasses import dataclass, field
+from multiprocessing.process import BaseProcess
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, cast
 
@@ -149,7 +150,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     _rpc_wave_id: int = 0
 
     def _init_executor(self) -> None:
-        self._processes: list[mp.Process] = []
+        self._processes: list[BaseProcess] = []
         self._closed = False
         self._is_failed = False
         self._failure_callbacks: list[Callable[[], None]] = []
@@ -158,7 +159,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._rpc_wave_id: int = 0
 
         num_workers = cast(int, self.od_config.num_gpus)
-        self.wake_events = [mp.Event() for _ in range(num_workers)]
+        self._mp_context = mp.get_context("spawn")
+        self.wake_events = [self._mp_context.Event() for _ in range(num_workers)]
 
         self._broadcast_mq = self._init_broadcast_queue(num_workers)
         broadcast_handle = self._broadcast_mq.export_handle()
@@ -345,7 +347,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self,
         broadcast_handle: Handle,
         wake_events: list[Event],
-    ) -> tuple[list[mp.Process], list[Handle]]:
+    ) -> tuple[list[BaseProcess], list[Handle]]:
         od_config = self.od_config
         logger.info("Starting server...")
 
@@ -354,8 +356,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # N-GPU run oversubscribes the host by N x core_count. Honours a
         # user-provided OMP_NUM_THREADS.
         set_multiprocessing_worker_envs()
-        mp.set_start_method("spawn", force=True)
-        processes = []
+        processes: list[BaseProcess] = []
 
         # Extract worker_extension_cls and custom_pipeline_args from od_config
         worker_extension_cls = od_config.worker_extension_cls
@@ -366,9 +367,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         scheduler_pipe_writers = []
 
         for i in range(num_gpus):
-            reader, writer = mp.Pipe(duplex=False)
+            reader, writer = self._mp_context.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
-            process = mp.Process(
+            process = self._mp_context.Process(
                 target=WorkerProc.worker_main,
                 args=(
                     i,  # rank

vllm_omni/diffusion/lora/loader.py

@@ -23,6 +23,11 @@ lora_convert_mapping: dict[str, Callable] = {
     "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
+    # 2.1's single-stream blocks keep the same leaf names the converter's
+    # protected n-grams cover (`attn.to_q/to_k/to_v`, `img_mlp`, ...), so the
+    # 2.0 converter applies. Non-diffusers LoRAs that target the fused `to_qkv`
+    # projection are not supported (the converter would split it into `to.qkv`).
+    "QwenImage21Pipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
     "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
 }

vllm_omni/diffusion/model_metadata.py

@@ -26,6 +26,8 @@ class DiffusionModelMetadata:
 # FLUX.2 Klein supports up to four reference images.
 FLUX2_KLEIN_MAX_INPUT_IMAGES = 4
 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
+# Qwen-Image 2.1 image-conditioned generation caps condition images at 4.
+QWEN_IMAGE_21_MAX_INPUT_IMAGES = 4
 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
 HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
 # Boogu-Image editing (TI2I) supports a single reference image for now.
@@ -41,6 +43,12 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
     ),
+    # Qwen-Image 2.1 handles both text-to-image (no image) and
+    # image-conditioned requests through the same pipeline class.
+    "QwenImage21Pipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=QWEN_IMAGE_21_MAX_INPUT_IMAGES,
+    ),
     "HunyuanImage3Pipeline": DiffusionModelMetadata(
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,

vllm_omni/diffusion/models/qwen_image_21/__init__.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Qwen Image 2.1 diffusion model components.
+
+Lazy by necessity: the pipeline imports the distributed VAE
+(`vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21`),
+which in turn imports the plain VAE module from this package. Eager imports
+here would make that a circular import whenever the distributed module is
+imported first.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    "AutoencoderKLQwenImage21",
+    "QwenImage21Pipeline",
+    "QwenImage21Transformer2DModel",
+    "get_qwen_image_21_post_process_func",
+    "get_qwen_image_21_pre_process_func",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "AutoencoderKLQwenImage21":
+        from .autoencoder_kl_qwenimage21 import AutoencoderKLQwenImage21
+
+        return AutoencoderKLQwenImage21
+    if name == "QwenImage21Transformer2DModel":
+        from .qwen_image_21_transformer import QwenImage21Transformer2DModel
+
+        return QwenImage21Transformer2DModel
+    if name in (
+        "QwenImage21Pipeline",
+        "get_qwen_image_21_post_process_func",
+        "get_qwen_image_21_pre_process_func",
+    ):
+        from .pipeline_qwen_image_21 import (
+            QwenImage21Pipeline,
+            get_qwen_image_21_post_process_func,
+            get_qwen_image_21_pre_process_func,
+        )
+
+        return {
+            "QwenImage21Pipeline": QwenImage21Pipeline,
+            "get_qwen_image_21_post_process_func": get_qwen_image_21_post_process_func,
+            "get_qwen_image_21_pre_process_func": get_qwen_image_21_pre_process_func,
+        }[name]
+    raise AttributeError(name)

vllm_omni/diffusion/models/qwen_image_21/cfg_parallel.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CFG Parallel Mixin for Qwen-Image 2.1.
+
+Differences from the 2.0 mixin (``qwen_image/cfg_parallel.py``):
+- The transformer owns a KV cache for the text/condition-image prefix, so the
+  denoise loop passes a fresh per-generation ``kv_cache`` and names the CFG
+  branch (``cache_branch="cond"/"uncond"``) on every forward.
+- The joint sequence is ``[condition images, ..., target image]`` — the target
+  noise prediction is the *tail* of the output, not the head.
+"""
+
+import logging
+from typing import Any
+
+import torch
+
+from vllm_omni.diffusion.distributed.cfg_parallel import (
+    CFGParallelMixin,
+    _get_cfg_world_size_or_one,
+    _unwrap,
+    _wrap,
+)
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+
+logger = logging.getLogger(__name__)
+
+
+class QwenImage21CFGParallelMixin(CFGParallelMixin, ProgressBarMixin):
+    """Base Mixin class for the Qwen-Image 2.1 pipeline providing shared CFG methods."""
+
+    def predict_noise_maybe_with_cfg(
+        self,
+        do_true_cfg: bool,
+        true_cfg_scale: float,
+        positive_kwargs: dict[str, Any],
+        negative_kwargs: dict[str, Any] | None,
+        cfg_normalize: bool = False,
+        output_slice: int | None = None,
+        kwargs: dict[str, Any] | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        # The base implementation slices `[:, :output_slice]` per branch, which
+        # matches the 2.0 condition-last layout. 2.1 concatenates condition
+        # latents *before* the target, so the target prediction is the tail of
+        # the joint output. Each branch must be tail-sliced *before* the CFG
+        # combine: positive and negative prompts tokenize to different lengths,
+        # so the full joint outputs do not align positionally.
+        def tail_slice(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+            if output_slice is None:
+                return _wrap(pred)
+            return tuple(p[:, -output_slice:] for p in _wrap(pred))
+
+        if do_true_cfg:
+            if _get_cfg_world_size_or_one() > 1:
+                cfg_group = get_cfg_group()
+                cfg_rank = get_classifier_free_guidance_rank()
+                branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
+                local_pred = tail_slice(self.predict_noise(**branch_kwargs))
+                gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred]
+                positive_noise_pred = tuple(g[0] for g in gathered)
+                negative_noise_pred = tuple(g[1] for g in gathered)
+            else:
+                positive_noise_pred = tail_slice(self.predict_noise(**positive_kwargs))
+                negative_noise_pred = tail_slice(self.predict_noise(**negative_kwargs))
+            return self.combine_cfg_noise(
+                positive_noise_pred,
+                negative_noise_pred,
+                true_cfg_scale,
+                cfg_normalize,
+                **({} if kwargs is None else {"kwargs": kwargs}),
+            )
+
+        pred = self.predict_noise(**positive_kwargs)
+        if output_slice is not None:
+            pred = _unwrap(tail_slice(pred))
+        return pred
+
+    def diffuse(
+        self,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        img_mask: torch.Tensor,
+        negative_img_mask: torch.Tensor | None,
+        latents: torch.Tensor,
+        image_latents: torch.Tensor | None,
+        img_shapes: list,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+        true_cfg_scale: float,
+        attention_kwargs: dict[str, Any] | None = None,
+        additional_transformer_kwargs: dict[str, Any] | None = None,
+    ) -> torch.Tensor:
+        """
+        Diffusion loop with optional classifier-free guidance.
+
+        Args:
+            prompt_embeds: Positive prompt embeddings
+            prompt_embeds_mask: Mask for positive prompt
+            negative_prompt_embeds: Negative prompt embeddings
+            negative_prompt_embeds_mask: Mask for negative prompt
+            img_mask: Joint-sequence image-slot mask for the positive branch
+                (VLM image slots plus one slot per 2x2 group of target tokens)
+            negative_img_mask: Same for the negative branch
+            latents: Noise latents to denoise (target image only, packed)
+            image_latents: Packed condition-image latents, prepended to the
+                noise latents (default: None for pure text-to-image)
+            img_shapes: Per-sample list of (frame, height, width) latent-token
+                shapes, condition images first and the target last
+            timesteps: Diffusion timesteps
+            do_true_cfg: Whether to apply CFG
+            true_cfg_scale: CFG scale factor
+            attention_kwargs: Passed through to the transformer (attn_path hint)
+            additional_transformer_kwargs: Extra kwargs for the transformer
+
+        Returns:
+            Denoised latents
+        """
+        # Fresh cache for every generation: the pipeline is a singleton shared
+        # across requests, so cache state must never leak between them.
+        cache_enabled = getattr(self.transformer, "causal_condition", False)
+        kv_cache = [{} for _ in range(len(self.transformer.transformer_blocks))] if cache_enabled else None
+        self.scheduler.set_begin_index(0)
+        self.transformer.do_true_cfg = do_true_cfg
+        additional_transformer_kwargs = additional_transformer_kwargs or {}
+
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                self._current_timestep = t
+
+                # Broadcast timestep to match batch size
+                timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+                # Condition image latents go in front of the noise latents.
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([image_latents, latents], dim=1)
+
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep / 1000,
+                    "encoder_hidden_states_mask": prompt_embeds_mask,
+                    "encoder_hidden_states": prompt_embeds,
+                    "img_shapes": img_shapes,
+                    "img_mask": img_mask,
+                    "attention_kwargs": attention_kwargs,
+                    "kv_cache": kv_cache,
+                    "cache_branch": "cond",
+                    **additional_transformer_kwargs,
+                }
+                if do_true_cfg:
+                    negative_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep / 1000,
+                        "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "img_shapes": img_shapes,
+                        "img_mask": negative_img_mask,
+                        "attention_kwargs": attention_kwargs,
+                        "kv_cache": kv_cache,
+                        "cache_branch": "uncond",
+                        **additional_transformer_kwargs,
+                    }
+                else:
+                    negative_kwargs = None
+
+                # Prefill (first step) returns the full joint sequence and
+                # caches the prefix; decode steps return target tokens only.
+                # Slicing the target tail is correct for both.
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg,
+                    true_cfg_scale,
+                    positive_kwargs,
+                    negative_kwargs,
+                    cfg_normalize=False,
+                    output_slice=latents.size(1),
+                )
+
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+                pbar.update()
+
+        return latents

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -0,0 +1,457 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA Graph capture for Qwen-Image-2.1 KV-cache decode steps.
+
+Decode (denoising steps 1..N) recomputes only the target image's tokens: the
+shape is identical every step and the timestep-independent prefix K/V come from
+the cross-step KV cache. That is the textbook CUDA-graph case, except the
+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.
+
+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.
+
+Each key admits a single live owner, tracked by a weak reference to the
+registering request's prefix K tensor. A second in-flight request whose cache
+aliases the same key is refused graph decode and stays eager until the
+owner's cache is released, so shared static buffers are never overwritten
+under a concurrent owner.
+
+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.
+
+The model runner regionally torch.compile's the transformer blocks even when
+decode graphs are enabled, so capture records the compiled (fused) kernels;
+inductor's own cudagraphs stay off so the two graph layers never stack.
+
+Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload
+without persistent weight staging, sequence/ring/tensor parallelism,
+HSDP/offload hooks, KV-cache quantization, dynamic LoRA wrappers, padded text
+masks (the masked attention path branches on mask contents, which cannot be
+captured), unknown graph key at decode, a second live request aliasing an
+owned graph key, and capture failure.
+"""
+
+from __future__ import annotations
+
+import weakref
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+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,
+    )
+
+logger = init_logger(__name__)
+
+
+def _dynamic_lora_wrappers_present(module: torch.nn.Module) -> bool:
+    """True once ``DiffusionLoRAManager`` has wrapped layers under ``module``.
+
+    A captured graph records the layers and branches that ran at capture time;
+    binding or rescaling an adapter afterwards changes the math without
+    changing the shapes the graph was keyed on, so graphs and dynamic LoRA
+    must not mix. Same reasoning as ``sensenova_u1.paged_decode``.
+    """
+    try:
+        from vllm.lora.layers import BaseLayerWithLoRA
+    except ImportError:  # pragma: no cover - depends on the wheel
+        return False
+    return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules())
+
+
+class QwenImage21DecodeGraphEntry:
+    """Static state for one captured decode graph.
+
+    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 entries are replaced by a ``None``
+    marker in the manager so their buffers can be released.
+    """
+
+    def __init__(
+        self,
+        *,
+        branch: str,
+        batch_size: int,
+        prefix_len: int,
+        target_tokens: int,
+        in_channels: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        kv_shapes: list[torch.Size],
+        kv_dtype: torch.dtype,
+        freqs: torch.Tensor,
+    ):
+        self.branch = branch
+        self.target_tokens = target_tokens
+        self.prefix_len = prefix_len
+        with torch.inference_mode(False):
+            self.hidden = torch.zeros(batch_size, target_tokens, in_channels, dtype=dtype, device=device)
+            self.timestep = torch.zeros(batch_size, dtype=dtype, device=device)
+            self.k = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes]
+            self.v = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes]
+            self.target_token_mask = torch.ones(target_tokens, dtype=torch.bool, device=device)
+            # RoPE frequencies are layout-dependent but step-independent. Clone
+            # into a normal (non-inference) tensor: the prefill runs under
+            # inference_mode, and the captured body takes views of this tensor
+            # outside it, which is forbidden for inference tensors.
+            self.freqs = freqs.clone()
+            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.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.
+
+        Warm-up runs the same computation it will record, allocating cuBLAS
+        workspaces and any lazy buffers outside the capture. Decode writes
+        nothing to the cache (prefix K/V are read-only), so warm-up is
+        idempotent.
+        """
+        side = torch.cuda.Stream()
+        side.wait_stream(torch.cuda.current_stream())
+        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
+        # the tree; entries only keep their *output* alive, whose tiny size
+        # makes reuse safe because the caller clones it before any other graph
+        # on the pool replays.
+        with torch.inference_mode(False), torch.no_grad():
+            with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                self.output = body()
+        self.graph = graph
+        self.captures += 1
+
+
+class QwenImage21DecodeGraphManager:
+    """Owns the static KV buffers and captured graphs of one transformer."""
+
+    def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8, model_level_offload: bool = False):
+        self.model = model
+        self.max_entries = max_entries
+        self.model_level_offload = model_level_offload
+        self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict()
+        # Key -> weak reference to the prefix K tensor of the request that owns
+        # the entry's static buffers. Dead references mean the owning request
+        # finished and its cache was released, so the key is free to adopt.
+        self._owners: dict[tuple, weakref.ReferenceType[torch.Tensor]] = {}
+        self._static_eligible: bool | None = None
+
+    # ── eligibility ──
+
+    def _offload_reason(self) -> str | None:
+        """Reason decode graphs are incompatible with model-level offload, if any.
+
+        Model-level (sequential) offload registers its hook on the top-level
+        transformer module, not on individual blocks. With persistent staging
+        the hook keeps the DiT on fixed device storage and only rebinds
+        ``p.data`` around each generation, and the swap completes in
+        ``pre_forward`` before capture or replay begin, so captured weight
+        pointers stay valid. Without staging every swap allocates fresh
+        storage and replay would read stale pointers, so capture stays
+        disabled and decode falls back to eager.
+        """
+        registry = getattr(self.model, "_hook_registry", None)
+        hooks = registry._hooks if registry is not None else {}
+        if hooks:
+            from vllm_omni.diffusion.offloader.sequential_backend import (
+                SequentialOffloadHook,
+                sequential_offload_staging_active,
+            )
+
+            if set(hooks) <= {SequentialOffloadHook._HOOK_NAME} and sequential_offload_staging_active(self.model):
+                return None
+            return "model module carries offload/cache hooks"
+        if self.model_level_offload:
+            return "model-level CPU offload without persistent staging swaps weight storage after capture"
+        return None
+
+    def _check_static_eligibility(self) -> bool:
+        """One-time checks that cannot change after weights are loaded."""
+        model = self.model
+        reason = self._offload_reason()
+        param = next(model.parameters(), None)
+        if reason is None and (param is None or param.device.type != "cuda"):
+            reason = f"model is not on CUDA (device={None if param is None else param.device})"
+        elif reason is None:
+            parallel_config = getattr(model, "parallel_config", None)
+            if parallel_config is not None:
+                sp = getattr(parallel_config, "sequence_parallel_size", None) or 1
+                if sp > 1 or getattr(parallel_config, "ring_degree", 1) > 1:
+                    reason = "sequence/ring parallelism changes decode shapes per rank"
+                elif getattr(parallel_config, "use_hsdp", False):
+                    reason = "HSDP wraps blocks with gather hooks outside the captured region"
+            if reason is None:
+                from vllm.distributed import get_tensor_model_parallel_world_size
+
+                if get_tensor_model_parallel_world_size() > 1:
+                    reason = "tensor parallelism runs collectives inside the captured region"
+            if reason is None:
+                for block in model.transformer_blocks:
+                    if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"):
+                        reason = "transformer blocks carry offload/cache hooks"
+                        break
+            if reason is None and model.transformer_blocks:
+                attn_layer = model.transformer_blocks[0].attn.attn
+                if getattr(attn_layer, "_kv_cache_dtype", None) is not None:
+                    reason = "KV-cache quantization is enabled"
+
+        if reason is not None:
+            logger.warning_once("Qwen-Image-2.1 CUDA graph decode disabled: %s. Falling back to eager decode.", reason)
+            return False
+        return True
+
+    def eligible(self) -> bool:
+        if self._static_eligible is None:
+            self._static_eligible = self._check_static_eligibility()
+        return self._static_eligible
+
+    def _backend_name(self) -> str:
+        backend = self.model.transformer_blocks[0].attn.attn.attn_backend
+        return backend.get_name() if backend is not None else "custom"
+
+    # ── prefill registration ──
+
+    def register_prefill(
+        self,
+        *,
+        kv_cache: list[dict[str, dict[str, torch.Tensor]]],
+        cache_branch: str,
+        prefix_len: int,
+        img_shapes: list[tuple[int, int, int]],
+        target_freqs: torch.Tensor,
+        joint_key_valid: torch.Tensor | None,
+        dtype: torch.dtype,
+    ) -> None:
+        """Prepare graph buffers without changing the request's prefix K/V.
+
+        The first in-flight request to register a key owns its static buffers;
+        registrations for the same key from another live request are refused
+        so their decode stays eager until the owner completes.
+        """
+        if not self.eligible() or torch.compiler.is_compiling():
+            return
+        if _dynamic_lora_wrappers_present(self.model):
+            if self.entries:
+                logger.warning_once(
+                    "Qwen-Image-2.1 CUDA graph decode: LoRA wrappers appeared; "
+                    "dropping captured graphs and falling back to eager decode."
+                )
+                self.entries.clear()
+            return
+        if joint_key_valid is not None and not bool(joint_key_valid.all()):
+            # The masked decode path (varlen-unpack or 4D mask) branches on mask
+            # contents and cannot be captured; this request stays eager.
+            logger.warning_once(
+                "Qwen-Image-2.1 CUDA graph decode: padded text mask present; this request falls back to eager decode."
+            )
+            return
+
+        first = kv_cache[0].get(cache_branch)
+        if first is None or "key" not in first:
+            return
+        if any(name.endswith("_scale") for name in first):
+            # Quantized (FP8) prefix storage: the static K/V buffers assume a single
+            # dtype and the captured body would not see the dequant scales (a bf16 K +
+            # fp8 V mix would even copy_ fp8 into bf16 buffers, silently dropping the
+            # scales). Keep quantized-cache requests on the eager decode path.
+            logger.warning_once(
+                "Qwen-Image-2.1 CUDA graph decode: quantized prefix KV cache "
+                "(prefix_kv_cache_dtype) is not graph-capturable; falling back to eager decode."
+            )
+            return
+        batch_size = first["key"].shape[0]
+        key = (
+            cache_branch,
+            batch_size,
+            prefix_len,
+            tuple(tuple(shape) for shape in img_shapes),
+            dtype,
+            first["key"].device.index,
+            self._backend_name(),
+        )
+
+        if key not in self.entries:
+            while len(self.entries) >= self.max_entries:
+                evicted_key, evicted_entry = self.entries.popitem(last=False)
+                self._owners.pop(evicted_key, None)
+                del evicted_entry
+                logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries)
+            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)
+
+        entry = self.entries[key]
+        if entry is not None:
+            owner = self._owners.get(key)
+            owner_prefix = owner() if owner is not None else None
+            if owner_prefix is not None and owner_prefix is not first["key"]:
+                # Another in-flight request still owns this entry's static
+                # buffers; refreshing them would overwrite that request's
+                # prefix mid-generation. This request decodes eagerly until
+                # the owner's cache is released and the weak reference dies.
+                logger.warning_once(
+                    "Qwen-Image-2.1 CUDA graph decode: entry key=%s is owned by another in-flight "
+                    "request; this request falls back to eager decode.",
+                    key,
+                )
+                return
+            self._owners[key] = weakref.ref(first["key"])
+            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 is not None and e.graph is not None),
+        )
+
+    # ── decode replay ──
+
+    def try_decode(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        kv_cache: list[dict[str, dict[str, torch.Tensor]]],
+        cache_branch: str,
+        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":
+            # Registration is device-agnostic (unit tests exercise the buffer
+            # machinery on CPU), but capture/replay is CUDA-only: an empty
+            # graph over CPU ops would replay as a stale no-op.
+            return None
+        key = (
+            cache_branch,
+            cached_key.shape[0],
+            cached_key.shape[1],
+            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:
+            return None
+        self.entries.move_to_end(key)
+
+        owner = self._owners.get(key)
+        owner_prefix = owner() if owner is not None else None
+        if owner_prefix is not None and owner_prefix is not cached_key:
+            # The entry is owned by a different in-flight request (see
+            # register_prefill); this request stays on the eager decode path
+            # until the owner's cache is released.
+            return None
+
+        entry.hidden.copy_(hidden_states[:, -entry.target_tokens :])
+        entry.timestep.copy_(timestep)
+        # 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:
+                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,
+                    capture_error,
+                )
+                return None
+            logger.info("Captured Qwen-Image-2.1 decode graph for key=%s", key)
+        entry.graph.replay()
+        assert entry.output is not None
+        # Clone before returning: the output lives in the shared graph pool and
+        # the next replay (of this or another graph on the pool) overwrites it.
+        return entry.output.clone()
+
+    def clear(self) -> None:
+        """Drop all entries, freeing buffers and captured graphs."""
+        self.entries.clear()
+        self._owners.clear()

vllm_omni/diffusion/offloader/module_residency.py

@@ -98,6 +98,12 @@ class PinnedModuleStager:
     A module iterable is treated as one staging group. It uses one copy stream
     and one reusable completion event. Tensors sharing storage keep their
     shapes, strides, offsets, dtypes, and aliases across every transition.
+
+    With ``retain_device_storage`` the device storages survive ``offload``:
+    ``load`` then only rebinds to the same fixed storages (weights are
+    immutable, so no H2D copy repeats). The stable storage addresses keep
+    CUDA-graph-captured weight pointers valid across offload/load swaps, at
+    the cost of keeping the device storage resident while offloaded.
     """
 
     def __init__(
@@ -108,6 +114,7 @@ class PinnedModuleStager:
         pin_memory: bool = True,
         copy_stream: Any | None = None,
         cache_retention: BoundedAllocatorCache | None = None,
+        retain_device_storage: bool = False,
     ) -> None:
         modules = (module,) if isinstance(module, nn.Module) else tuple(module)
         if not modules or not all(isinstance(item, nn.Module) for item in modules):
@@ -117,6 +124,7 @@ class PinnedModuleStager:
         self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream()
         self._ready_event = current_omni_platform.Event()
         self.cache_retention = cache_retention
+        self.retain_device_storage = retain_device_storage
         self.loaded = False
         self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
         self._device_storages: list[torch.Tensor] = []
@@ -234,6 +242,12 @@ class PinnedModuleStager:
     def load(self) -> None:
         if self.loaded:
             return
+        if self._device_storages:
+            # Storage retained from a previous load still holds the immutable
+            # weights; rebinding suffices, no H2D copy repeats.
+            self._bind(self._device_storages)
+            self.loaded = True
+            return
         try:
             self._load_once()
         except torch.OutOfMemoryError:
@@ -263,7 +277,8 @@ class PinnedModuleStager:
             self.loaded = False
             self._release_cache(force=True)
             raise
-        self._device_storages.clear()
+        if not self.retain_device_storage:
+            self._device_storages.clear()
         self.loaded = False
         self._release_cache()
 

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -13,6 +13,7 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
+from .module_residency import PinnedModuleStager
 from .plan_resolver import resolve_offload_plan
 
 logger = init_logger(__name__)
@@ -60,6 +61,7 @@ class SequentialOffloadHook(ModelHook):
         pin_memory: bool = True,
         use_hsdp: bool = False,
         offload_after_context: bool = True,
+        persistent_staging: bool = False,
     ):
         # Modules to offload to CPU before this module runs
         self.offload_targets = offload_targets
@@ -67,6 +69,21 @@ class SequentialOffloadHook(ModelHook):
         self.pin_memory = pin_memory
         self.use_hsdp = use_hsdp
         self.offload_after_context = offload_after_context
+        self.persistent_staging = persistent_staging
+        self._stager: PinnedModuleStager | None = None
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        if self.persistent_staging and not self.use_hsdp:
+            # Fixed staging storage keeps CUDA-graph-captured weight pointers
+            # valid across offload/load swaps. HSDP keeps the plain move
+            # semantics (decode graphs refuse HSDP anyway).
+            self._stager = PinnedModuleStager(
+                module,
+                self.device,
+                pin_memory=self.pin_memory,
+                retain_device_storage=True,
+            )
+        return module
 
     @staticmethod
     def _move_params(
@@ -101,7 +118,23 @@ class SequentialOffloadHook(ModelHook):
                 moved = True
         return moved
 
+    @staticmethod
+    def _module_stager(module: nn.Module) -> PinnedModuleStager | None:
+        """The persistent stager owned by ``module``'s own sequential hook, if any.
+
+        Cross-module swaps call ``_to_cpu(target)`` through the *running*
+        module's hook, so staging state must be resolved on the target's hook
+        to keep its ``loaded`` bookkeeping consistent.
+        """
+        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+        hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None
+        return hook._stager if isinstance(hook, SequentialOffloadHook) else None
+
     def _to_cpu(self, module: nn.Module) -> None:
+        stager = self._module_stager(module)
+        if stager is not None:
+            stager.offload()
+            return
         # XPU's allocator doesn't respect stream dependencies in empty_cache,
         # so non-blocking copies can race with cache eviction. Use blocking
         # copies on XPU to avoid NULL pointer errors during DMA.
@@ -116,6 +149,10 @@ class SequentialOffloadHook(ModelHook):
             current_omni_platform.empty_cache()
 
     def _to_gpu(self, module: nn.Module) -> None:
+        stager = self._module_stager(module)
+        if stager is not None:
+            stager.load()
+            return
         self._move_params(module, self.device, non_blocking=False)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
@@ -147,6 +184,7 @@ def apply_sequential_offload(
     offload_initial_dits: bool = False,
     offload_dit_modules: Collection[nn.Module] | None = None,
     offload_encoder_modules: Collection[nn.Module] | None = None,
+    persistent_dit_staging: bool = False,
 ) -> None:
     """Apply sequential offloading hooks to DiT and encoder modules.
 
@@ -165,6 +203,10 @@ def apply_sequential_offload(
             every DiT for backward compatibility.
         offload_encoder_modules: Encoder/stage modules allowed to move to CPU.
             None selects every supplied module for backward compatibility.
+        persistent_dit_staging: Keep DiT weights on fixed device staging
+            storage across swaps, so CUDA-graph-captured weight pointers stay
+            valid. Encoders keep plain move semantics so their memory is
+            actually freed while the DiT runs.
 
     Example:
         >>> apply_sequential_offload(
@@ -193,6 +235,7 @@ def apply_sequential_offload(
                 pin_memory=pin_memory,
                 use_hsdp=use_hsdp,
                 offload_after_context=id(dit_mod) in selected_dit_ids,
+                persistent_staging=persistent_dit_staging,
             )
             registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
             logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__)
@@ -261,6 +304,15 @@ def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
     return hook
 
 
+def sequential_offload_staging_active(module: nn.Module) -> bool:
+    """True when the module's sequential offload hook serves fixed device storage.
+
+    Fixed staging storage keeps CUDA-graph-captured weight pointers valid
+    across offload/load swaps.
+    """
+    return SequentialOffloadHook._module_stager(module) is not None
+
+
 @contextmanager
 def sequential_offload_component(module: nn.Module) -> Iterator[None]:
     """Activate and release a hooked component called outside ``forward``."""
@@ -362,6 +414,14 @@ class ModelLevelOffloadBackend(OffloadBackend):
                 use_hsdp=self.config.use_hsdp,
                 offload_dit_modules=selected_dits,
                 offload_encoder_modules=selected_encoders,
+                # Fixed DiT staging keeps decode-graph weight pointers valid;
+                # only models that capture weight pointers (Qwen-Image-2.1's
+                # CUDA-graph decode) pay for it — every other model keeps
+                # plain move semantics so offloading actually frees DiT VRAM
+                # while the encoders run.
+                persistent_dit_staging=(
+                    self.device.type == "cuda" and any(getattr(dit, "enable_cuda_graph_decode", False) for dit in dits)
+                ),
             )
         except BaseException:
             try:

vllm_omni/diffusion/registry.py

@@ -54,6 +54,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image_layered",
         "QwenImageLayeredPipeline",
     ),
+    "QwenImage21Pipeline": (
+        "qwen_image_21",
+        "pipeline_qwen_image_21",
+        "QwenImage21Pipeline",
+    ),
     "GlmImagePipeline": (
         "glm_image",
         "pipeline_glm_image",
@@ -413,6 +418,10 @@ _NO_CACHE_ACCELERATION = {
     "Pi0Pipeline",
     "Pi05Pipeline",
     "LingBotWorldCausalDMDPipeline",
+    # Qwen-Image 2.1's transformer carries its own prefix KV cache (keyed by CFG
+    # branch) across denoising steps, which conflicts with cache_dit / tea_cache
+    # step-skipping hooks.
+    "QwenImage21Pipeline",
 }
 
 
@@ -586,6 +595,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "AnimaPipeline": "get_anima_post_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
+    "QwenImage21Pipeline": "get_qwen_image_21_post_process_func",
     "GlmImagePipeline": "get_glm_image_post_process_func",
     "ZImagePipeline": "get_post_process_func",
     "OvisImagePipeline": "get_ovis_image_post_process_func",
@@ -668,6 +678,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "BooguImageTurboPipeline": "get_boogu_image_pre_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
+    "QwenImage21Pipeline": "get_qwen_image_21_pre_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",

vllm_omni/diffusion/request.py

@@ -75,6 +75,8 @@ class OmniDiffusionRequest:
     # A model preprocessor may keep selected requests on the legacy full-forward
     # path even when the engine is globally configured for step execution.
     use_step_execution: bool = True
+    # Whether first and later denoising steps may share a step batch.
+    allow_mixed_step_phases: bool = True
     # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward
     # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric.
     kv_recv_ms: float = 0.0

vllm_omni/diffusion/sched/step_scheduler.py

@@ -13,6 +13,7 @@ from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
 from vllm_omni.diffusion.sched.interface import (
     DiffusionRequestStatus,
     DiffusionSchedulerOutput,
+    SchedulerRequestState,
 )
 
 if TYPE_CHECKING:
@@ -62,6 +63,19 @@ class StepScheduler(BaseScheduler):
         )
         return request_id
 
+    def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
+        if not super()._can_schedule_waiting(state):
+            return False
+
+        is_first_step = self._request_progress[state.request_id].current_step == 0
+        for request_id in self._running:
+            running = self._request_states[request_id]
+            if state.req.allow_mixed_step_phases and running.req.allow_mixed_step_phases:
+                continue
+            if (self._request_progress[request_id].current_step == 0) != is_first_step:
+                return False
+        return True
+
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids
         if not scheduled_request_ids and not sched_output.finished_req_ids:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -239,6 +239,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if model is None:
             return
 
+        if getattr(model, "enable_cuda_graph_decode", False):
+            # Decode-graph models still compile their blocks: graph capture
+            # records the compiled (fused) kernels. Inductor's own cudagraphs
+            # must stay off so the two graph layers never stack.
+            try:
+                import torch._inductor.config as inductor_config
+
+                inductor_config.triton.cudagraphs = False
+            except Exception as e:
+                logger.warning(
+                    "Model runner: could not disable inductor cudagraphs for the "
+                    "compile+CUDA-graph combo (%s); capture may fail and fall back "
+                    "to compiled eager decode.",
+                    e,
+                )
+            logger.info("Model runner: %s combines CUDA graph decode with torch.compile.", attr_name)
+
         compile_granularity = self.od_config.diffusion_compile_granularity
         compile_dynamic = self.od_config.diffusion_compile_dynamic
         try:

vllm_omni/engine/arg_utils.py

@@ -580,6 +580,10 @@ class OrchestratorArgs:
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
+    # today). None defers to the OmniDiffusionConfig default (enabled);
+    # --enforce-eager also disables it.
+    enable_cuda_graph_decode: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
     video_output_transport: dict[str, object] | None = None

vllm_omni/entrypoints/cli/serve.py

@@ -605,6 +605,16 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
+        omni_config_group.add_argument(
+            "--enable-cuda-graph-decode",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Capture fixed-shape KV-cache decode steps into CUDA graphs on models that implement them "
+                "(currently Qwen-Image-2.1). Independent of compilation_config.cudagraph_mode; "
+                "disable with --no-enable-cuda-graph-decode. --enforce-eager also disables it."
+            ),
+        )
         omni_config_group.add_argument(
             "--fa-deterministic",
             dest="fa_deterministic",

vllm_omni/entrypoints/openai/api_server.py

@@ -2239,6 +2239,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},
@@ -2246,6 +2247,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),

vllm_omni/entrypoints/openai/serving_chat.py

@@ -3519,7 +3519,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,

vllm_omni/model_extras/qwen_image_21.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from typing import Any
+
+from PIL import Image
+
+
+def build_image_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    input_image: Image.Image | list[Image.Image],
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    img_prompt: dict[str, Any] = {
+        "prompt": prompt,
+        "multi_modal_data": {"image": input_image},
+    }
+    if height is not None:
+        img_prompt["height"] = height
+    if width is not None:
+        img_prompt["width"] = width
+    if negative_prompt is not None:
+        img_prompt["negative_prompt"] = negative_prompt
+    return img_prompt

vllm_omni/model_extras/registry.py

@@ -75,6 +75,9 @@ from vllm_omni.model_extras.ming_flash_omni import (
 from vllm_omni.model_extras.ming_flash_omni import (
     build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt,
 )
+from vllm_omni.model_extras.qwen_image_21 import (
+    build_image_to_image_prompt as build_qwen_image_21_image_to_image_prompt,
+)
 from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS
 from vllm_omni.model_extras.sensenova_u1 import (
     SENSENOVA_U1_EXTRA_BODY_PARAMS,
@@ -299,6 +302,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt,
         "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt,
     },
+    "QwenImage21Pipeline": {
+        "image_to_image_prompt_builder": build_qwen_image_21_image_to_image_prompt,
+    },
 }
 
 for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"):