[Model] Add JoyAI-Image-Edit pipeline support (#4112)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com>

.buildkite/amd/scripts/prewarm_cosmos3_aiter.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Prebuild the BF16 AITER MHA modules used by Cosmos3 on ROCm."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.platforms import current_omni_platform
+
+ATTENTION_VARIANTS = (
+    (True, "mha_fwd_bf16_nbias_mask_nlse_ndropout_nqscale"),
+    (False, "mha_fwd_bf16_nbias_nmask_nlse_ndropout_nqscale"),
+)
+
+
+def main() -> None:
+    if not current_omni_platform.is_rocm():
+        raise RuntimeError("The Cosmos3 AITER prewarm must run on ROCm")
+    if fa.flash_attn_func is None:
+        raise RuntimeError("ROCm FlashAttention/AITER is unavailable")
+
+    cache_dir_value = os.environ.get("AITER_JIT_DIR")
+    if not cache_dir_value:
+        raise RuntimeError("AITER_JIT_DIR must identify the job-local cache")
+    cache_dir = Path(cache_dir_value)
+
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+    batch_size = 1
+    sequence_length = 64
+    num_heads = 8
+    head_dim = 64
+
+    torch.manual_seed(42)
+    query = torch.randn(
+        batch_size,
+        sequence_length,
+        num_heads,
+        head_dim,
+        device=device,
+        dtype=dtype,
+    )
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+
+    for causal, expected_module in ATTENTION_VARIANTS:
+        output = fa.flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=0.0,
+            softmax_scale=head_dim**-0.5,
+            causal=causal,
+        )
+        if isinstance(output, tuple):
+            output = output[0]
+        current_omni_platform.synchronize()
+
+        if output.shape != query.shape:
+            raise RuntimeError(f"Unexpected AITER output shape: {output.shape}")
+        if not torch.isfinite(output).all():
+            raise RuntimeError("AITER prewarm produced non-finite output")
+
+        module_path = cache_dir / f"{expected_module}.so"
+        if not module_path.is_file():
+            raise RuntimeError(f"Expected AITER module was not built: {module_path}")
+        print(f"Prebuilt Cosmos3 AITER module: {module_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

.buildkite/amd/test-amd-merge.yml

@@ -37,24 +37,22 @@ steps:
       mirror_hardwares: [amdproduction]
       grade: Blocking
       timeout_in_minutes: 45
+      # Buildkite's status-based rule cannot prove every exit 134 is a GPU hang.
+      # It retries 134 once; the Kubernetes job retry starts a fresh pod.
+      # Assertions (1), timeouts (124), and OOMs (137) remain non-retryable.
+      retry:
+        automatic:
+          - exit_status: 134
+            limit: 1
       commands:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
         # Keep multi-GPU tests out of this one-GPU pool even when their module
         # also carries the cpu marker for unit tests.
+        # Keep π0.5 CPU units out of these shards: CUDA READY/MERGE already run
+        # the same suite mandatorily, and AMD covers it non-blocking in nightly.
         - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-    # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4,
-    # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane.
-    - label: "Simple · Pi0.5 CPU Test"
-      agent_pool: mi300_1
-      depends_on: amd-build
-      mirror_hardwares: [amdproduction]
-      grade: Blocking
-      timeout_in_minutes: 45
-      commands:
-        - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'"
 
     - label: "Diffusion · AuK VAE GPU Test"
       agent_pool: mi300_1
@@ -128,6 +126,7 @@ steps:
       commands:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA
         - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
         - >-
           pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py

.buildkite/amd/test-amd-nightly.yml

@@ -82,3 +82,89 @@ steps:
               limit: 1
         commands:
           - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
+
+  # π0.5 CPU units are mandatory on CUDA READY/MERGE. On AMD they exercise the
+  # ROCm container/import path with CPU tensors only, so keep them off the
+  # blocking READY/MERGE critical path and retain non-blocking nightly coverage.
+  - group: ":robot_face: Pi0.5 CPU Nightly"
+    steps:
+      - label: "Simple · Pi0.5 CPU Test"
+        agent_pool: mi300_1
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 60
+        commands:
+          - export VLLM_ROCM_USE_AITER=0
+          - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'"
+
+  - group: ":full_moon: LingBot-Video MI300 Nightly"
+    steps:
+      - label: "LingBot-Video Dense CFG-off Smoke"
+        agent_pool: mi300_1
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 90
+        artifact_paths:
+          - artifacts/rocm-lingbot-nightly/**/*
+        commands:
+          - pytest --collect-only -q tests/e2e/online_serving/test_lingbot_video_expansion.py::test_cfg_off -m "full_model and diffusion and rocm and MI325 and cards_1" --run-level full_model
+          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_expansion.py::test_cfg_off -m "full_model and diffusion and rocm and MI325 and cards_1" --run-level full_model --junitxml=/tmp/lingbot-pytest.xml
+          - |
+            artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH:?}/artifacts/rocm-lingbot-nightly"
+            mkdir -p "$$artifact_dir"
+            cp -v /tmp/lingbot-pytest.xml "$$artifact_dir/"
+
+  - group: ":full_moon: Cosmos3 MI300 Nightly"
+    steps:
+      - label: "Cosmos3 T2I Function"
+        agent_pool: mi300_1
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 90
+        artifact_paths:
+          - artifacts/rocm-cosmos3-nightly/**/*
+        commands:
+          # Cosmos3 uses both masked and unmasked BF16 AITER MHA modules. Build
+          # them before server startup, then reuse the same job-local cache
+          # during the smoke test.
+          - export VLLM_ROCM_USE_AITER=1
+          - export AITER_JIT_DIR="/tmp/vllm-omni-cosmos3-aiter-$$BUILDKITE_JOB_ID"
+          - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-cosmos3-torch-extensions-$$BUILDKITE_JOB_ID"
+          - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+          - pytest --collect-only -q tests/e2e/online_serving/test_cosmos3.py::test_text_to_image_001 -m "core_model and rocm and MI325 and cards_1" --run-level full_model
+          - >-
+            timeout --signal=TERM --kill-after=2m 55m
+            python .buildkite/amd/scripts/prewarm_cosmos3_aiter.py
+          - >-
+            timeout --signal=TERM --kill-after=2m 25m
+            pytest -s -v
+            tests/e2e/online_serving/test_cosmos3.py::test_text_to_image_001
+            -m "core_model and rocm and MI325 and cards_1"
+            --run-level full_model
+            --junitxml=/tmp/cosmos3-pytest.xml
+          - |
+            artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH:?}/artifacts/rocm-cosmos3-nightly"
+            mkdir -p "$$artifact_dir"
+            cp -v /tmp/cosmos3-pytest.xml "$$artifact_dir/"
+
+  - group: ":full_moon: HunyuanImage3 MI300 Nightly"
+    steps:
+      - label: "HunyuanImage3 Offline Pixel Accuracy"
+        agent_pool: mi300_4
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 180
+        env:
+          HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
+          HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        artifact_paths:
+          - artifacts/rocm-hunyuanimage3/**/*
+        commands:
+          - pytest --collect-only -q tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py::test_hunyuan_image3_pixel_accuracy_offline -m "full_model and rocm and MI325 and cards_4" --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py::test_hunyuan_image3_pixel_accuracy_offline -m "full_model and rocm and MI325 and cards_4" --run-level full_model --junitxml=/tmp/hunyuanimage3-pytest.xml
+          - |
+            artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH:?}/artifacts/rocm-hunyuanimage3"
+            mkdir -p "$$artifact_dir"
+            cp -v tests/e2e/accuracy/artifacts/HunyuanImage-3_0-Instruct/*.png "$$artifact_dir/"
+            cp -v /tmp/hunyuanimage3-pytest.xml "$$artifact_dir/"

.buildkite/amd/test-amd-ready.yml

@@ -37,24 +37,22 @@ steps:
       mirror_hardwares: [amdproduction]
       grade: Blocking
       timeout_in_minutes: 45
+      # Buildkite's status-based rule cannot prove every exit 134 is a GPU hang.
+      # It retries 134 once; the Kubernetes job retry starts a fresh pod.
+      # Assertions (1), timeouts (124), and OOMs (137) remain non-retryable.
+      retry:
+        automatic:
+          - exit_status: 134
+            limit: 1
       commands:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
         # Keep multi-GPU tests out of this one-GPU pool even when their module
         # also carries the cpu marker for unit tests.
+        # Keep π0.5 CPU units out of these shards: CUDA READY/MERGE already run
+        # the same suite mandatorily, and AMD covers it non-blocking in nightly.
         - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-    # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4,
-    # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane.
-    - label: "Simple · Pi0.5 CPU Test"
-      agent_pool: mi300_1
-      depends_on: amd-build
-      mirror_hardwares: [amdproduction]
-      grade: Blocking
-      timeout_in_minutes: 45
-      commands:
-        - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'"
 
     - label: "Diffusion · AuK VAE GPU Test"
       agent_pool: mi300_1
@@ -128,6 +126,15 @@ steps:
           tests/diffusion/attention/
           -m 'core_model and rocm and MI325 and cards_1'
           --run-level "core_model"
+        # MammothModa2 exercises the same two AITER MHA modules built by the
+        # preceding suite. Keep its default-backend run in this job so a cold
+        # worker compiles each module once instead of timing out the generic
+        # model job.
+        - >-
+          timeout --signal=TERM --kill-after=2m 10m
+          pytest -s -v -n 1
+          tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py
+          --run-level "core_model"
 
     # AR paged attention is real GPU work, including cold AITER compilation.
     # Keep its complete numerical matrix out of the CPU-marked Simple shards.
@@ -216,9 +223,17 @@ steps:
         - >-
           timeout 15m pytest -sv tests/diffusion/models/
           --ignore=tests/diffusion/models/auk/test_auk_vae_fastpath.py
+          --ignore=tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py
           -m 'core_model and cuda and not
           (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
           --run-level "core_model"
+        # Keep the model-specific numerical coverage blocking without paying
+        # the default ROCm backend's cold AITER compilation cost in this lane.
+        - >-
+          timeout 10m env DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA
+          pytest -sv
+          tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py
+          --run-level "core_model"
 
     - label: "Diffusion · Wan22 Test"
       agent_pool: mi300_1
@@ -245,7 +260,7 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
-- label: "ROCm · Engine & Model Executor GPU Coverage (R2-01)"
+- label: "ROCm · Model Executor GPU Coverage (R2-01)"
   agent_pool: mi300_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -258,18 +273,17 @@ steps:
     # artifacts relative to its checkout directory.
     - export R2_01_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-01"
     - mkdir -p "$$R2_01_ARTIFACT_DIR"
+    # The dedicated Engine Test below owns the only selected tests/engine case;
+    # do not duplicate its model startup in this time-bounded coverage lane.
     # The CUDA ready job also covers tests/worker/test_batched_omni_output.py.
-    # Keep the initial ROCm burn-in scoped to the requested directories.
+    # Keep the initial ROCm burn-in scoped to tests/model_executor/.
     # Verbose pytest output and JUnit retain node IDs without a second
     # collection-only pass over the same GPU-backed test tree.
-    # The CUDA/L4-only CosyVoice3 ragged-flow parity test cold-JITs multiple
-    # AITER kernels on ROCm and is covered by the CUDA ready lane.
     - >-
       pytest -s -v -ra
-      tests/engine/ tests/model_executor/
-      -m 'core_model and cuda and not
+      tests/model_executor/
+      -m 'core_model and cuda and not nvidia_only and not
       (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
-      -k 'not test_code2wav_streaming_batch_matches_ragged_flow_numerics'
       --run-level "core_model"
       --durations=0
       --junitxml="$$R2_01_ARTIFACT_DIR/pytest.xml"

.buildkite/common/ci_source_file_dependencies.yml

@@ -115,6 +115,10 @@ x-qwen-image: &qwen_image
   - vllm_omni/diffusion/models/dmd2/
   - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
 
+x-joy-image: &joy_image
+  - vllm_omni/diffusion/models/joy_image/
+  - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
+
 x-hunyuan-image3: &hunyuan_image3
   - vllm_omni/diffusion/models/hunyuan_image3/
   - vllm_omni/model_executor/models/hunyuan_image3/
@@ -149,6 +153,10 @@ x-lingbot-video: &lingbot_video
 x-lingbot-world: &lingbot_world
   - vllm_omni/diffusion/models/lingbot_world/
 
+x-kandinsky6: &kandinsky6
+  - vllm_omni/diffusion/models/kandinsky6/
+  - vllm_omni/model_extras/kandinsky6.py
+
 
 source_file_dependencies:
   # Hardware presets and shared test helpers. Not a job key — do not attach this
@@ -383,6 +391,12 @@ source_file_dependencies:
     - tests/e2e/accuracy/test_qwen_image.py
     - tests/e2e/accuracy/test_diffusers_backend_similarity.py
 
+  diffusion_joy_image_accuracy:
+    - *joy_image
+    - tests/diffusion/models/joy_image/
+    - tests/e2e/accuracy/test_joy_image_edit.py
+    - recipes/JD/JoyAI-Image-Edit.md
+
   diffusion_qwen_image_perf:
     - *qwen_image
     - tests/dfx/perf/scripts/run_diffusion_benchmark.py
@@ -491,6 +505,11 @@ source_file_dependencies:
     - tests/dfx/perf/scripts/run_benchmark.py
     - tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json
 
+  diffusion_kandinsky6_function:
+    - *kandinsky6
+    - tests/e2e/offline_inference/test_kandinsky6.py
+    - tests/e2e/online_serving/test_kandinsky6.py
+
 
   # Invalid-param jobs: model business code only, plus the reliability scripts.
   reliability_invalid_param_h100:

.buildkite/cuda/test-merge.yml

@@ -189,6 +189,15 @@ steps:
             '
         mirror_hardwares: h100_1
 
+      - label: "Diffusion · Kandinsky6 Test"
+        source_file_dependencies: diffusion_kandinsky6_function
+        commands:
+          - |
+            timeout 40m bash -c "
+              pytest -s -v tests/e2e/offline_inference/test_kandinsky6.py tests/e2e/online_serving/test_kandinsky6.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            "
+        mirror_hardwares: h100_1
+
       - label: "Diffusion · Wan22 Test"
         source_file_dependencies: diffusion_wan22_function
         commands:

.buildkite/cuda/test-nightly.yml

@@ -268,6 +268,12 @@ steps:
           - pytest -s -v tests/e2e/accuracy/test_qwen_image.py -m "H100 and B200 and cards_1" --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m "full_model and H100 and B200 and cards_1" --run-level full_model
 
+      - label: ":full_moon: JoyAI-Image-Edit · Accuracy Test"
+        source_file_dependencies: diffusion_joy_image_accuracy
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_joy_image_edit.py -m "full_model and diffusion and H100 and B200 and cards_1" --run-level full_model
+
       - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
         source_file_dependencies: diffusion_hunyuan_image3_accuracy
         timeout_in_minutes: 180

.buildkite/cuda/test-ready.yml

@@ -105,7 +105,7 @@ steps:
 
       - label: "Diffusion · Other Test"
         commands:
-          - timeout 15m pytest -sv tests/diffusion/*.py tests/diffusion/ar_diffusion -m 'core_model and cuda' --run-level "core_model"
+          - timeout 15m pytest -sv tests/diffusion/*.py tests/diffusion/ar_diffusion tests/diffusion/layers/test_adalayernorm_*.py -m 'core_model and cuda' --run-level "core_model"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
@@ -122,7 +122,7 @@ steps:
     commands:
       - |
         timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model'
+                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker/test_gpu_generation_async_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model'
         "
     mirror_hardwares: l4_1
 

.buildkite/intel/pipeline-intel.yml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.30.0"
+          VLLM_VERSION: "v0.31.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

apps/ComfyUI-vLLM-Omni/README.md

@@ -271,10 +271,10 @@ The [WF-05 template](example_workflows/vLLM-Omni%20MiniMax-H3%20Latent%20Mask%20
 Connect a **Latent Mask Editing** node to **Generate Video → latent_edit** to edit a source clip instead of generating from scratch. It uploads the source media and serializes the video/audio noise masks the MiniMax H3 API accepts:
 
 - `source_video` / `source_audio` — the media to edit.
-- `video_mask` — a ComfyUI mask image; `0` preserves a region, `1` regenerates it, fractional values blend. A 2D mask `[H, W]` is applied to every frame; a 3D mask `[T, H, W]` is treated as a temporal mask (one slice per frame) for continuation or extension. It is resized to the video latent grid.
+- `video_mask` — a ComfyUI mask image; `0` preserves a region, `1` regenerates it, fractional values blend. A 2D mask `[H, W]` is applied to every frame; a 3D mask `[T, H, W]` is treated as a temporal mask (one slice per frame) for continuation or extension. The client area-downsamples it by 16 to keep the upload small, and the server resizes it to the video latent grid.
 - `audio_mask` — a scalar in `[0, 1]`; `0` keeps the source audio, `1` regenerates it, fractional values blend.
 
-A non-trivial mask requires its matching source, and a source without a mask is rejected. This node only forwards inputs to the server; the served model must declare latent-mask editing support (MiniMax-H3).
+The node requires at least one mask; the server then enforces the cross-field rules (a source requires a matching mask, and a non-trivial mask requires its source). This node only forwards inputs to the server; the served model must declare latent-mask editing support (MiniMax-H3).
 
 ### TTS (e.g., Qwen TTS series)
 

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -36,7 +36,6 @@ from .comfyui_vllm_omni.nodes import (
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
-    "VLLMOmniMiniMaxH3TemporalMask": VLLMOmniMiniMaxH3TemporalMask,
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
     "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
@@ -46,6 +45,7 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
     "VLLMOmniVideoReferences": VLLMOmniVideoReferences,
     "VLLMOmniLatentMaskEditing": VLLMOmniLatentMaskEditing,
+    "VLLMOmniMiniMaxH3TemporalMask": VLLMOmniMiniMaxH3TemporalMask,
     # === Params ===
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
@@ -59,7 +59,6 @@ NODE_CLASS_MAPPINGS = {
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "VLLMOmniMiniMaxH3TemporalMask": "MiniMax-H3 Temporal Mask",
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
     "VLLMOmniGenerateVideo": "Generate Video",
@@ -69,6 +68,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
     "VLLMOmniVideoReferences": "Video References",
     "VLLMOmniLatentMaskEditing": "Latent Mask Editing",
+    "VLLMOmniMiniMaxH3TemporalMask": "MiniMax-H3 Temporal Mask",
     # === Params ===
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -8,7 +8,7 @@ import torch
 from comfy_api.input import AudioInput, VideoInput
 
 from .utils.api_client import VLLMOmniClient
-from .utils.latent_mask import _align_frame_count, _video_latent_t
+from .utils.latent_mask import _align_frame_count
 from .utils.logger import get_logger
 from .utils.models import lookup_model_spec
 from .utils.types import (
@@ -1133,13 +1133,13 @@ class VLLMOmniMiniMaxH3TemporalMask:
         else:
             raise ValueError(f"Unknown temporal mask mode: {mode}")
         available = min(frames, math.floor(boundary * 24 + 1e-8))
+        # Snap the preserved prefix to whole VAE clips (5 + 17n frames) so every
+        # latent it maps to is fully preserved.
         prefix = 0 if available < 5 else 5 + 17 * ((available - 5) // 17)
-        preserved = _video_latent_t(prefix) if prefix else 0
-        total = _video_latent_t(frames)
-        mask = torch.ones(total, 1, 1)
-        mask[:preserved] = 0
+        # One slice per output frame; the server pools frames to the latent grid.
+        mask = torch.ones(frames, 1, 1)
+        mask[:prefix] = 0
         indices = torch.arange(frames, device=images.device)
         source_indices = (indices * (source_fps / 24)).floor().long().clamp(max=images.shape[0] - 1)
         preview_images = images.index_select(0, source_indices)
-        preview_mask = mask.index_select(0, torch.arange(frames) * total // frames)
-        return mask, 24.0, preview_images, preview_mask
+        return mask, 24.0, preview_images, mask

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -30,7 +30,7 @@ from .format import (
     video_to_base64,
     video_to_bytes,
 )
-from .latent_mask import scalar_mask_to_json, video_mask_to_grid_json
+from .latent_mask import scalar_mask_to_json, video_mask_to_json
 from .logger import get_logger, pretty_printer
 from .models import lookup_model_spec
 from .types import (
@@ -388,6 +388,15 @@ class VLLMOmniClient:
                 content_type="image/png",
             )
 
+        if len(keyframe_images) == 2:
+            for image_filename, image in keyframe_images:
+                form.add_field(
+                    "input_references",
+                    image_tensor_to_png_bytes(image, image_filename),
+                    filename=image_filename,
+                    content_type="image/png",
+                )
+
         # === latent-mask editing (MiniMax H3) ===
         if latent_edit is not None:
             source_video = latent_edit.get("source_video")
@@ -398,21 +407,16 @@ class VLLMOmniClient:
             if video_mask is None and audio_mask is None:
                 raise ValueError("Latent-mask editing requires at least one mask.")
 
-            video_mask_trivial = video_mask is None or bool((video_mask == 1.0).all().item())
-            audio_mask_trivial = audio_mask is None or audio_mask == 1.0
-            if not video_mask_trivial and source_video is None:
-                raise ValueError("A non-trivial video mask requires a source video.")
-            if not audio_mask_trivial and source_audio is None and source_video is None:
-                raise ValueError("A non-trivial audio mask requires a source audio or a source video with audio.")
-
             if source_video is not None:
                 form.add_field(
                     "source_video",
-                    video_to_bytes(source_video, "source.mp4"),
+                    video_to_bytes(source_video),
                     filename="source.mp4",
                     content_type="video/mp4",
                 )
             if source_audio is not None:
+                # The filename extension selects the audio codec (mp3), so it is
+                # meaningful even though aiohttp uses the explicit filename below.
                 form.add_field(
                     "source_audio",
                     audio_to_bytes(source_audio, "source_audio.mp3"),
@@ -420,7 +424,7 @@ class VLLMOmniClient:
                     content_type="audio/mpeg",
                 )
             if video_mask is not None:
-                mask_json = video_mask_to_grid_json(video_mask, width=width, height=height, num_frames=num_frames)
+                mask_json = video_mask_to_json(video_mask)
                 form.add_field(
                     "video_noise_mask",
                     mask_json.encode("utf-8"),
@@ -435,15 +439,6 @@ class VLLMOmniClient:
                     content_type="application/json",
                 )
 
-        if len(keyframe_images) == 2:
-            for image_filename, image in keyframe_images:
-                form.add_field(
-                    "input_references",
-                    image_tensor_to_png_bytes(image, image_filename),
-                    filename=image_filename,
-                    content_type="image/png",
-                )
-
         # === model specific params. Either use a specialized builder, or add flattened fields as-is ===
         if model_params is not None:
             model_params = dict(model_params)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/latent_mask.py

@@ -1,11 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+"""Client-side serialization of MiniMax-H3 latent-edit masks.
+
+A video mask is sent in frame space — a 2D spatial ``[H, W]`` (applied to every
+frame) or a 3D ``[T, H, W]`` (one slice per frame) — and the server resolves it
+to the latent grid. An audio mask is a scalar applied to all time steps.
+"""
+
 import json
+import math
 
 import torch
 import torch.nn.functional as F
 
+# The server rejects mask JSON above 8 MiB, so the spatial axes are
+# area-downsampled by the VAE spatial stride before upload. This depends only on
+# the stride, not on the H3 shape lattice, which the server still owns.
+_VAE_SPATIAL_STRIDE = 16
+_MASK_DECIMALS = 4
+
 
 def scalar_mask_to_json(value: float) -> str:
     if not 0.0 <= value <= 1.0:
@@ -13,6 +27,27 @@ def scalar_mask_to_json(value: float) -> str:
     return str(value)
 
 
+def video_mask_to_json(mask: torch.Tensor) -> str:
+    """Serialize a frame-space video mask to JSON for the ``video_noise_mask`` field.
+
+    ``mask`` is a 2D spatial mask or a 3D frame-space mask. Its spatial axes are
+    area-downsampled by the VAE stride; the server resizes the result to the
+    latent grid, so the client neither floors the canvas nor maps frames to
+    latents.
+    """
+    if mask.ndim not in (2, 3):
+        raise ValueError(f"expected a 2D or 3D mask tensor, got {mask.ndim}D")
+    frames = mask.unsqueeze(0) if mask.ndim == 2 else mask
+    size = tuple(max(1, math.ceil(dim / _VAE_SPATIAL_STRIDE)) for dim in frames.shape[1:])
+    frames = F.interpolate(frames.unsqueeze(1).float(), size=size, mode="area").squeeze(1)
+    if mask.ndim == 2:
+        frames = frames.squeeze(0)
+    # Round in float64 so tolist() emits short decimals rather than float32 noise.
+    return json.dumps(frames.double().round(decimals=_MASK_DECIMALS).tolist(), separators=(",", ":"))
+
+
+# The temporal-mask node snaps its preserve boundary to the H3 frame lattice;
+# this mirrors the server shape planner.
 def _align_frame_count(frame_count: int) -> int:
     if frame_count <= 0:
         return 1
@@ -20,32 +55,3 @@ def _align_frame_count(frame_count: int) -> int:
     while current % 17 != 5:
         current += 1
     return current
-
-
-def _video_latent_t(frame_count: int) -> int:
-    if frame_count <= 5:
-        return 2
-    return ((int(frame_count) - 5) // 17) * 5 + 2
-
-
-def video_mask_to_grid(mask: torch.Tensor, *, width: int, height: int, num_frames: int) -> torch.Tensor:
-    tv = _video_latent_t(_align_frame_count(num_frames))
-    height = int(height) // 32 * 32
-    width = int(width) // 32 * 32
-    gh, gw = height // 16, width // 16
-
-    if mask.ndim == 2:
-        mask = mask.unsqueeze(0)
-    elif mask.ndim != 3:
-        raise ValueError(f"expected a 2D or 3D mask tensor, got {mask.ndim}D")
-
-    grid = F.interpolate(mask.unsqueeze(1).float(), size=(gh, gw), mode="area").squeeze(1)
-    if grid.shape[0] == 1:
-        grid = grid.expand(tv, gh, gw)
-    elif grid.shape[0] != tv:
-        grid = F.interpolate(grid.unsqueeze(0).unsqueeze(0), size=(tv, gh, gw), mode="nearest").squeeze(0).squeeze(0)
-    return grid
-
-
-def video_mask_to_grid_json(mask: torch.Tensor, *, width: int, height: int, num_frames: int) -> str:
-    return json.dumps(video_mask_to_grid(mask, width=width, height=height, num_frames=num_frames).tolist())

apps/ComfyUI-vLLM-Omni/docs/wf05-h3-latent-editing.md

@@ -12,6 +12,6 @@ Edit a source video with MiniMax-H3 using spatial or temporal masks. The workflo
 Open **vLLM-Omni MiniMax-H3 Latent Mask Editing.json** from the example workflows and set the Generate Video URL/model for the case you want to run.
 
 1. Upload your source video in Load Video.
-2. Adjust the case's mask. For removal and inpainting, set Mask Canvas to the video dimensions, Mask Rectangle to the region size, and Combine Masks x/y to its position. For continuation and extension, set MiniMax-H3 Temporal Mask's mode and duration; continuation also uses `preserve_fraction`. Match the duration in Generate Video, and use a duration longer than the source for extension. Mask values are `0` to preserve and `1` to regenerate.
+2. Adjust the case's mask. For removal and inpainting, set Mask Canvas to the video dimensions, Mask Rectangle to the region size, and Combine Masks x/y to its position, or replace that chain with Load Image (as Mask) for an arbitrary-shape mask. For continuation and extension, set MiniMax-H3 Temporal Mask's mode and duration; continuation also uses `preserve_fraction`. Match the duration in Generate Video, and use a duration longer than the source for extension. Mask values are `0` to preserve and `1` to regenerate.
 3. Adjust the prompt to describe the desired content. Set the output dimensions and duration in Generate Video. If editing audio, update the sound description and audio mask (`0` preserves, `1` regenerates the whole clip).
 4. Select the case's Mask Preview Save Video node and choose **Execute to selected output nodes** to preview the mask as a red overlay without inference, or select its Result Save Video node to generate the edited video. Running the entire workflow generates all four cases.

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Latent Mask Editing - Image Mask.json

@@ -1,499 +0,0 @@
-{
-  "id": "latent-mask-editing-7380",
-  "revision": 0,
-  "last_node_id": 7,
-  "last_link_id": 5,
-  "nodes": [
-    {
-      "id": 1,
-      "type": "MarkdownNote",
-      "pos": [
-        40,
-        20
-      ],
-      "size": [
-        420,
-        210
-      ],
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "title": "Latent-mask editing server prerequisite",
-      "properties": {},
-      "widgets_values": [
-        "Requires a vLLM-Omni server that declares latent-mask editing support (MiniMax-H3, EDIT-01). The source clip and video/audio noise masks are uploaded to /v1/videos; 0 preserves, 1 regenerates, fractional blends."
-      ],
-      "color": "#432",
-      "bgcolor": "#000"
-    },
-    {
-      "id": 2,
-      "type": "LoadVideo",
-      "pos": [
-        40,
-        280
-      ],
-      "size": [
-        420,
-        120
-      ],
-      "flags": {},
-      "order": 1,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "file",
-          "name": "file",
-          "type": "COMBO",
-          "widget": {
-            "name": "file"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            1
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LoadVideo"
-      },
-      "widgets_values": [
-        ""
-      ]
-    },
-    {
-      "id": 3,
-      "type": "LoadImageMask",
-      "pos": [
-        40,
-        460
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "image",
-          "type": "STRING",
-          "widget": {
-            "name": "image"
-          },
-          "link": null
-        },
-        {
-          "name": "channel",
-          "type": "STRING",
-          "widget": {
-            "name": "channel"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "mask",
-          "name": "mask",
-          "type": "MASK",
-          "links": [
-            2
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LoadImageMask"
-      },
-      "widgets_values": [
-        "",
-        "red"
-      ]
-    },
-    {
-      "id": 4,
-      "type": "VLLMOmniLatentMaskEditing",
-      "pos": [
-        560,
-        260
-      ],
-      "size": [
-        420,
-        300
-      ],
-      "flags": {},
-      "order": 3,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "source_video",
-          "name": "source_video",
-          "type": "VIDEO",
-          "link": 1,
-          "shape": 7
-        },
-        {
-          "localized_name": "source_audio",
-          "name": "source_audio",
-          "type": "AUDIO",
-          "link": null
-        },
-        {
-          "localized_name": "video_mask",
-          "name": "video_mask",
-          "type": "MASK",
-          "link": 2,
-          "shape": 7
-        },
-        {
-          "localized_name": "audio_mask",
-          "name": "audio_mask",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_mask"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "latent_edit",
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "links": [
-            3
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniLatentMaskEditing"
-      },
-      "widgets_values": [
-        0.5
-      ]
-    },
-    {
-      "id": 5,
-      "type": "VLLMOmniMiniMaxH3Params",
-      "pos": [
-        40,
-        680
-      ],
-      "size": [
-        420,
-        140
-      ],
-      "flags": {},
-      "order": 4,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "audio_flow_shift",
-          "name": "audio_flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_flow_shift"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "flow_shift",
-          "name": "flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "flow_shift"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video params",
-          "name": "video params",
-          "type": "VIDEO_PARAMS",
-          "links": [
-            4
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniMiniMaxH3Params"
-      },
-      "widgets_values": [
-        3.0,
-        12.0
-      ]
-    },
-    {
-      "id": 6,
-      "type": "VLLMOmniGenerateVideo",
-      "pos": [
-        1080,
-        200
-      ],
-      "size": [
-        540,
-        680
-      ],
-      "flags": {},
-      "order": 5,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "frame",
-          "name": "frame",
-          "type": "IMAGE",
-          "link": null,
-          "shape": 7
-        },
-        {
-          "localized_name": "references",
-          "name": "references",
-          "type": "VIDEO_REFERENCES",
-          "link": null,
-          "shape": 7
-        },
-        {
-          "localized_name": "sampling_params",
-          "name": "sampling_params",
-          "type": "SAMPLING_PARAMS",
-          "link": null
-        },
-        {
-          "localized_name": "lora",
-          "name": "lora",
-          "type": "REMOTE_LORA",
-          "link": null
-        },
-        {
-          "localized_name": "model_params",
-          "name": "model_params",
-          "type": "VIDEO_PARAMS",
-          "link": 4,
-          "shape": 7
-        },
-        {
-          "localized_name": "fast_h3",
-          "name": "fast_h3",
-          "type": "FASTH3_DEPLOYMENT",
-          "link": null
-        },
-        {
-          "localized_name": "latent_edit",
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "link": 3,
-          "shape": 7
-        },
-        {
-          "localized_name": "url",
-          "name": "url",
-          "type": "STRING",
-          "widget": {
-            "name": "url"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "model",
-          "name": "model",
-          "type": "STRING",
-          "widget": {
-            "name": "model"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "prompt",
-          "name": "prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "prompt"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "negative_prompt",
-          "name": "negative_prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "negative_prompt"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "width",
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "height",
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "fps",
-          "name": "fps",
-          "type": "INT",
-          "widget": {
-            "name": "fps"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "duration",
-          "name": "duration",
-          "type": "FLOAT",
-          "widget": {
-            "name": "duration"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            5
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniGenerateVideo"
-      },
-      "widgets_values": [
-        "http://127.0.0.1:8000/v1",
-        "MiniMaxAI/MiniMax-H3",
-        "Partially restyle the clip and soundtrack",
-        "",
-        160,
-        120,
-        24,
-        4.458
-      ]
-    },
-    {
-      "id": 7,
-      "type": "SaveVideo",
-      "pos": [
-        1700,
-        300
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "link": 5,
-          "shape": 7
-        },
-        {
-          "localized_name": "filename_prefix",
-          "name": "filename_prefix",
-          "type": "STRING",
-          "widget": {
-            "name": "filename_prefix"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "format",
-          "name": "format",
-          "type": "COMFY_DYNAMICCOMBO_V3",
-          "widget": {
-            "name": "format"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": []
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SaveVideo"
-      },
-      "widgets_values": [
-        "vllm-omni/latent-edit",
-        "mp4"
-      ]
-    }
-  ],
-  "links": [
-    [
-      1,
-      2,
-      0,
-      4,
-      0,
-      "VIDEO"
-    ],
-    [
-      2,
-      3,
-      0,
-      4,
-      2,
-      "MASK"
-    ],
-    [
-      3,
-      4,
-      0,
-      6,
-      6,
-      "LATENT_MASK_EDITING"
-    ],
-    [
-      4,
-      5,
-      0,
-      6,
-      4,
-      "VIDEO_PARAMS"
-    ],
-    [
-      5,
-      6,
-      0,
-      7,
-      0,
-      "VIDEO"
-    ]
-  ],
-  "groups": [],
-  "config": {},
-  "extra": {},
-  "version": 0.4
-}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Latent Mask Editing - Temporal Mask.json

@@ -1,560 +0,0 @@
-{
-  "last_node_id": 9,
-  "last_link_id": 7,
-  "nodes": [
-    {
-      "id": 2,
-      "type": "LoadVideo",
-      "pos": [
-        40,
-        100
-      ],
-      "size": [
-        420,
-        120
-      ],
-      "flags": {},
-      "order": 1,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "file",
-          "type": "STRING",
-          "widget": {
-            "name": "file"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            1
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LoadVideo"
-      },
-      "widgets_values": [
-        ""
-      ]
-    },
-    {
-      "id": 3,
-      "type": "SolidMask",
-      "pos": [
-        40,
-        300
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "value",
-          "type": "FLOAT",
-          "widget": {
-            "name": "value"
-          },
-          "link": null
-        },
-        {
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "mask",
-          "type": "MASK",
-          "links": [
-            2
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SolidMask"
-      },
-      "widgets_values": [
-        0.0,
-        160,
-        120
-      ]
-    },
-    {
-      "id": 4,
-      "type": "SolidMask",
-      "pos": [
-        40,
-        500
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 3,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "value",
-          "type": "FLOAT",
-          "widget": {
-            "name": "value"
-          },
-          "link": null
-        },
-        {
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "mask",
-          "type": "MASK",
-          "links": [
-            3
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SolidMask"
-      },
-      "widgets_values": [
-        1.0,
-        160,
-        120
-      ]
-    },
-    {
-      "id": 5,
-      "type": "BatchMasksNode",
-      "pos": [
-        250,
-        380
-      ],
-      "size": [
-        200,
-        100
-      ],
-      "flags": {},
-      "order": 4,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "masks.mask0",
-          "type": "MASK",
-          "link": 2,
-          "slot_index": 0
-        },
-        {
-          "name": "masks.mask1",
-          "type": "MASK",
-          "link": 3,
-          "slot_index": 1
-        }
-      ],
-      "outputs": [
-        {
-          "name": "MASK",
-          "type": "MASK",
-          "links": [
-            4
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "BatchMasksNode"
-      },
-      "widgets_values": []
-    },
-    {
-      "id": 6,
-      "type": "VLLMOmniLatentMaskEditing",
-      "pos": [
-        450,
-        180
-      ],
-      "size": [
-        400,
-        120
-      ],
-      "flags": {},
-      "order": 5,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "source_video",
-          "type": "VIDEO",
-          "link": 1
-        },
-        {
-          "name": "source_audio",
-          "type": "AUDIO",
-          "link": null
-        },
-        {
-          "name": "video_mask",
-          "type": "MASK",
-          "link": 4
-        },
-        {
-          "name": "audio_mask",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_mask"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "links": [
-            5
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniLatentMaskEditing"
-      },
-      "widgets_values": [
-        -1.0
-      ]
-    },
-    {
-      "id": 7,
-      "type": "VLLMOmniMiniMaxH3Params",
-      "pos": [
-        40,
-        700
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "audio_flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_flow_shift"
-          },
-          "link": null
-        },
-        {
-          "name": "flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "flow_shift"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "video params",
-          "type": "VIDEO_PARAMS",
-          "links": [
-            6
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniMiniMaxH3Params"
-      },
-      "widgets_values": [
-        3.0,
-        12.0
-      ]
-    },
-    {
-      "id": 8,
-      "type": "VLLMOmniGenerateVideo",
-      "pos": [
-        700,
-        180
-      ],
-      "size": [
-        450,
-        380
-      ],
-      "flags": {},
-      "order": 7,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "frame",
-          "type": "IMAGE",
-          "link": null
-        },
-        {
-          "name": "references",
-          "type": "VIDEO_REFERENCES",
-          "link": null
-        },
-        {
-          "name": "sampling_params",
-          "type": "SAMPLING_PARAMS",
-          "link": null
-        },
-        {
-          "name": "lora",
-          "type": "LORA",
-          "link": null
-        },
-        {
-          "name": "model_params",
-          "type": "VIDEO_PARAMS",
-          "link": 6
-        },
-        {
-          "name": "fast_h3",
-          "type": "FAST_H3",
-          "link": null
-        },
-        {
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "link": 5
-        },
-        {
-          "name": "url",
-          "type": "STRING",
-          "widget": {
-            "name": "url"
-          },
-          "link": null
-        },
-        {
-          "name": "model",
-          "type": "STRING",
-          "widget": {
-            "name": "model"
-          },
-          "link": null
-        },
-        {
-          "name": "prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "prompt"
-          },
-          "link": null
-        },
-        {
-          "name": "negative_prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "negative_prompt"
-          },
-          "link": null
-        },
-        {
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        },
-        {
-          "name": "fps",
-          "type": "INT",
-          "widget": {
-            "name": "fps"
-          },
-          "link": null
-        },
-        {
-          "name": "duration",
-          "type": "FLOAT",
-          "widget": {
-            "name": "duration"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            7
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniGenerateVideo"
-      },
-      "widgets_values": [
-        "http://127.0.0.1:8000/v1",
-        "MiniMaxAI/MiniMax-H3",
-        "temporal mask: regenerate the ending",
-        "",
-        160,
-        120,
-        24,
-        4.458
-      ]
-    },
-    {
-      "id": 9,
-      "type": "SaveVideo",
-      "pos": [
-        900,
-        180
-      ],
-      "size": [
-        420,
-        200
-      ],
-      "flags": {},
-      "order": 8,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "video",
-          "type": "VIDEO",
-          "link": 7
-        },
-        {
-          "name": "filename_prefix",
-          "type": "STRING",
-          "widget": {
-            "name": "filename_prefix"
-          },
-          "link": null
-        },
-        {
-          "name": "format",
-          "type": "STRING",
-          "widget": {
-            "name": "format"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "video",
-          "type": "VIDEO",
-          "links": []
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SaveVideo"
-      },
-      "widgets_values": [
-        "temporal_mask",
-        "mp4"
-      ]
-    }
-  ],
-  "links": [
-    [
-      1,
-      2,
-      0,
-      6,
-      0,
-      "VIDEO"
-    ],
-    [
-      2,
-      3,
-      0,
-      5,
-      0,
-      "MASK"
-    ],
-    [
-      3,
-      4,
-      0,
-      5,
-      1,
-      "MASK"
-    ],
-    [
-      4,
-      5,
-      0,
-      6,
-      2,
-      "MASK"
-    ],
-    [
-      5,
-      6,
-      0,
-      8,
-      6,
-      "LATENT_MASK_EDITING"
-    ],
-    [
-      6,
-      7,
-      0,
-      8,
-      4,
-      "VIDEO_PARAMS"
-    ],
-    [
-      7,
-      8,
-      0,
-      9,
-      0,
-      "VIDEO"
-    ]
-  ],
-  "groups": [],
-  "config": {},
-  "extra": {},
-  "version": 0.4
-}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Latent Mask Editing.json

@@ -1,511 +0,0 @@
-{
-  "id": "latent-mask-editing-7380",
-  "revision": 0,
-  "last_node_id": 7,
-  "last_link_id": 5,
-  "nodes": [
-    {
-      "id": 1,
-      "type": "MarkdownNote",
-      "pos": [
-        40,
-        20
-      ],
-      "size": [
-        420,
-        210
-      ],
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "title": "Latent-mask editing server prerequisite",
-      "properties": {},
-      "widgets_values": [
-        "Requires a vLLM-Omni server that declares latent-mask editing support (MiniMax-H3, EDIT-01). The source clip and video/audio noise masks are uploaded to /v1/videos; 0 preserves, 1 regenerates, fractional blends."
-      ],
-      "color": "#432",
-      "bgcolor": "#000"
-    },
-    {
-      "id": 2,
-      "type": "LoadVideo",
-      "pos": [
-        40,
-        280
-      ],
-      "size": [
-        420,
-        120
-      ],
-      "flags": {},
-      "order": 1,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "file",
-          "name": "file",
-          "type": "COMBO",
-          "widget": {
-            "name": "file"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            1
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LoadVideo"
-      },
-      "widgets_values": [
-        ""
-      ]
-    },
-    {
-      "id": 3,
-      "type": "SolidMask",
-      "pos": [
-        40,
-        460
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "value",
-          "name": "value",
-          "type": "FLOAT",
-          "widget": {
-            "name": "value"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "width",
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "height",
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "mask",
-          "name": "mask",
-          "type": "MASK",
-          "links": [
-            2
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SolidMask"
-      },
-      "widgets_values": [
-        1.0,
-        160,
-        120
-      ]
-    },
-    {
-      "id": 4,
-      "type": "VLLMOmniLatentMaskEditing",
-      "pos": [
-        560,
-        260
-      ],
-      "size": [
-        420,
-        300
-      ],
-      "flags": {},
-      "order": 3,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "source_video",
-          "name": "source_video",
-          "type": "VIDEO",
-          "link": 1,
-          "shape": 7
-        },
-        {
-          "localized_name": "source_audio",
-          "name": "source_audio",
-          "type": "AUDIO",
-          "link": null
-        },
-        {
-          "localized_name": "video_mask",
-          "name": "video_mask",
-          "type": "MASK",
-          "link": 2,
-          "shape": 7
-        },
-        {
-          "localized_name": "audio_mask",
-          "name": "audio_mask",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_mask"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "latent_edit",
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "links": [
-            3
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniLatentMaskEditing"
-      },
-      "widgets_values": [
-        0.5
-      ]
-    },
-    {
-      "id": 5,
-      "type": "VLLMOmniMiniMaxH3Params",
-      "pos": [
-        40,
-        680
-      ],
-      "size": [
-        420,
-        140
-      ],
-      "flags": {},
-      "order": 4,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "audio_flow_shift",
-          "name": "audio_flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "audio_flow_shift"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "flow_shift",
-          "name": "flow_shift",
-          "type": "FLOAT",
-          "widget": {
-            "name": "flow_shift"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video params",
-          "name": "video params",
-          "type": "VIDEO_PARAMS",
-          "links": [
-            4
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniMiniMaxH3Params"
-      },
-      "widgets_values": [
-        3.0,
-        12.0
-      ]
-    },
-    {
-      "id": 6,
-      "type": "VLLMOmniGenerateVideo",
-      "pos": [
-        1080,
-        200
-      ],
-      "size": [
-        540,
-        680
-      ],
-      "flags": {},
-      "order": 5,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "frame",
-          "name": "frame",
-          "type": "IMAGE",
-          "link": null,
-          "shape": 7
-        },
-        {
-          "localized_name": "references",
-          "name": "references",
-          "type": "VIDEO_REFERENCES",
-          "link": null,
-          "shape": 7
-        },
-        {
-          "localized_name": "sampling_params",
-          "name": "sampling_params",
-          "type": "SAMPLING_PARAMS",
-          "link": null
-        },
-        {
-          "localized_name": "lora",
-          "name": "lora",
-          "type": "REMOTE_LORA",
-          "link": null
-        },
-        {
-          "localized_name": "model_params",
-          "name": "model_params",
-          "type": "VIDEO_PARAMS",
-          "link": 4,
-          "shape": 7
-        },
-        {
-          "localized_name": "fast_h3",
-          "name": "fast_h3",
-          "type": "FASTH3_DEPLOYMENT",
-          "link": null
-        },
-        {
-          "localized_name": "latent_edit",
-          "name": "latent_edit",
-          "type": "LATENT_MASK_EDITING",
-          "link": 3,
-          "shape": 7
-        },
-        {
-          "localized_name": "url",
-          "name": "url",
-          "type": "STRING",
-          "widget": {
-            "name": "url"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "model",
-          "name": "model",
-          "type": "STRING",
-          "widget": {
-            "name": "model"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "prompt",
-          "name": "prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "prompt"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "negative_prompt",
-          "name": "negative_prompt",
-          "type": "STRING",
-          "widget": {
-            "name": "negative_prompt"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "width",
-          "name": "width",
-          "type": "INT",
-          "widget": {
-            "name": "width"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "height",
-          "name": "height",
-          "type": "INT",
-          "widget": {
-            "name": "height"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "fps",
-          "name": "fps",
-          "type": "INT",
-          "widget": {
-            "name": "fps"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "duration",
-          "name": "duration",
-          "type": "FLOAT",
-          "widget": {
-            "name": "duration"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": [
-            5
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VLLMOmniGenerateVideo"
-      },
-      "widgets_values": [
-        "http://127.0.0.1:8000/v1",
-        "MiniMaxAI/MiniMax-H3",
-        "Partially restyle the clip and soundtrack",
-        "",
-        160,
-        120,
-        24,
-        4.458
-      ]
-    },
-    {
-      "id": 7,
-      "type": "SaveVideo",
-      "pos": [
-        1700,
-        300
-      ],
-      "size": [
-        420,
-        160
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "link": 5,
-          "shape": 7
-        },
-        {
-          "localized_name": "filename_prefix",
-          "name": "filename_prefix",
-          "type": "STRING",
-          "widget": {
-            "name": "filename_prefix"
-          },
-          "link": null
-        },
-        {
-          "localized_name": "format",
-          "name": "format",
-          "type": "COMFY_DYNAMICCOMBO_V3",
-          "widget": {
-            "name": "format"
-          },
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "localized_name": "video",
-          "name": "video",
-          "type": "VIDEO",
-          "links": []
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "SaveVideo"
-      },
-      "widgets_values": [
-        "vllm-omni/latent-edit",
-        "mp4"
-      ]
-    }
-  ],
-  "links": [
-    [
-      1,
-      2,
-      0,
-      4,
-      0,
-      "VIDEO"
-    ],
-    [
-      2,
-      3,
-      0,
-      4,
-      2,
-      "MASK"
-    ],
-    [
-      3,
-      4,
-      0,
-      6,
-      6,
-      "LATENT_MASK_EDITING"
-    ],
-    [
-      4,
-      5,
-      0,
-      6,
-      4,
-      "VIDEO_PARAMS"
-    ],
-    [
-      5,
-      6,
-      0,
-      7,
-      0,
-      "VIDEO"
-    ]
-  ],
-  "groups": [],
-  "config": {},
-  "extra": {},
-  "version": 0.4
-}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax-H3 Latent Mask Editing.json

@@ -2293,10 +2293,10 @@
       "title": "Usage Guide",
       "properties": {},
       "widgets_values": [
-        "Use a compatible MiniMax-H3 latent-edit server. Select a source video and set the server URL/model in each Generate Video node.\n\nSelect one Save Preview or Save Video output, then use Execute to selected output nodes. Running the whole workflow submits all four cases.\n\nMask: 0 preserves source latents; 1 regenerates. Audio mask applies to the whole clip, not individual sounds. Preview overlays are silent and do not run inference."
+        "Use a compatible MiniMax-H3 latent-edit server. Select a source video and set the server URL/model in each Generate Video node.\n\nSelect one Save Preview or Save Video output, then use Execute to selected output nodes. Running the whole workflow submits all four cases.\n\nMask: 0 preserves source latents; 1 regenerates. For an arbitrary-shape spatial mask, replace a case's Mask Rectangle chain with Load Image (as Mask). Audio mask applies to the whole clip, not individual sounds. Preview overlays are silent and do not run inference."
       ],
       "widgets_values_named": {
-        "text": "Use a compatible MiniMax-H3 latent-edit server. Select a source video and set the server URL/model in each Generate Video node.\n\nSelect one Save Preview or Save Video output, then use Execute to selected output nodes. Running the whole workflow submits all four cases.\n\nMask: 0 preserves source latents; 1 regenerates. Audio mask applies to the whole clip, not individual sounds. Preview overlays are silent and do not run inference."
+        "text": "Use a compatible MiniMax-H3 latent-edit server. Select a source video and set the server URL/model in each Generate Video node.\n\nSelect one Save Preview or Save Video output, then use Execute to selected output nodes. Running the whole workflow submits all four cases.\n\nMask: 0 preserves source latents; 1 regenerates. For an arbitrary-shape spatial mask, replace a case's Mask Rectangle chain with Load Image (as Mask). Audio mask applies to the whole clip, not individual sounds. Preview overlays are silent and do not run inference."
       },
       "color": "#432",
       "bgcolor": "#000"
@@ -2823,7 +2823,7 @@
         "Node name for S&R": "VLLMOmniGenerateVideo"
       },
       "widgets_values": [
-        "http://127.0.0.1:8001/v1",
+        "http://127.0.0.1:8000/v1",
         "MiniMaxAI/MiniMax-H3",
         "integrated_multimodal_description: The masked area is a flat, seamless continuation of the surrounding wall, matching its color, texture, and lighting. The wall surface remains consistent throughout the video.\n\noverall_soundscape: \n\nnon_diegetic_music: ",
         "",
@@ -2833,7 +2833,7 @@
         5
       ],
       "widgets_values_named": {
-        "url": "http://127.0.0.1:8001/v1",
+        "url": "http://127.0.0.1:8000/v1",
         "model": "MiniMaxAI/MiniMax-H3",
         "prompt": "integrated_multimodal_description: The masked area is a flat, seamless continuation of the surrounding wall, matching its color, texture, and lighting. The wall surface remains consistent throughout the video.\n\noverall_soundscape: \n\nnon_diegetic_music: ",
         "negative_prompt": "",
@@ -3677,7 +3677,7 @@
         "Node name for S&R": "VLLMOmniGenerateVideo"
       },
       "widgets_values": [
-        "http://127.0.0.1:8001/v1",
+        "http://127.0.0.1:8000/v1",
         "MiniMaxAI/MiniMax-H3",
         "integrated_multimodal_description: A small orange cat rests curled up in the masked area, gently swaying its tail from side to side. Its appearance matches the surrounding anime style, lighting, and shadows.\n\noverall_soundscape:\n\nnon_diegetic_music: ",
         "",
@@ -3687,7 +3687,7 @@
         5
       ],
       "widgets_values_named": {
-        "url": "http://127.0.0.1:8001/v1",
+        "url": "http://127.0.0.1:8000/v1",
         "model": "MiniMaxAI/MiniMax-H3",
         "prompt": "integrated_multimodal_description: A small orange cat rests curled up in the masked area, gently swaying its tail from side to side. Its appearance matches the surrounding anime style, lighting, and shadows.\n\noverall_soundscape:\n\nnon_diegetic_music: ",
         "negative_prompt": "",
@@ -3916,7 +3916,7 @@
         "Node name for S&R": "VLLMOmniGenerateVideo"
       },
       "widgets_values": [
-        "http://127.0.0.1:8001/v1",
+        "http://127.0.0.1:8000/v1",
         "MiniMaxAI/MiniMax-H3",
         "integrated_multimodal_description: The sky turns completely overcast, with the sun hidden and no direct sunlight, as torrential rain pours over the school entrance in dense visible streaks and splashes on the wet pavement. Distant lightning flashes in the sky behind the school buildings.\n\noverall_soundscape: Loud, steady rainfall, splashing water, and gusts of wind.\n\nnon_diegetic_music: None.",
         "",
@@ -3926,7 +3926,7 @@
         5
       ],
       "widgets_values_named": {
-        "url": "http://127.0.0.1:8001/v1",
+        "url": "http://127.0.0.1:8000/v1",
         "model": "MiniMaxAI/MiniMax-H3",
         "prompt": "integrated_multimodal_description: The sky turns completely overcast, with the sun hidden and no direct sunlight, as torrential rain pours over the school entrance in dense visible streaks and splashes on the wet pavement. Distant lightning flashes in the sky behind the school buildings.\n\noverall_soundscape: Loud, steady rainfall, splashing water, and gusts of wind.\n\nnon_diegetic_music: None.",
         "negative_prompt": "",
@@ -4363,7 +4363,7 @@
         "Node name for S&R": "VLLMOmniGenerateVideo"
       },
       "widgets_values": [
-        "http://127.0.0.1:8001/v1",
+        "http://127.0.0.1:8000/v1",
         "MiniMaxAI/MiniMax-H3",
         "integrated_multimodal_description: Cut to the same school on a summer night beneath a deep blue sky dotted with stars. A firework rises from behind the school buildings, leaving a glowing trail; the camera slowly tilts upward to follow it into the sky, ending as it bursts into a large bloom of colorful sparks.\n\noverall_soundscape: Soft summer insect sounds, followed by the firework rising with a whistle and bursting with a distant boom and crackling sparks.\n\nnon_diegetic_music: N/A",
         "",
@@ -4373,7 +4373,7 @@
         10
       ],
       "widgets_values_named": {
-        "url": "http://127.0.0.1:8001/v1",
+        "url": "http://127.0.0.1:8000/v1",
         "model": "MiniMaxAI/MiniMax-H3",
         "prompt": "integrated_multimodal_description: Cut to the same school on a summer night beneath a deep blue sky dotted with stars. A firework rises from behind the school buildings, leaving a glowing trail; the camera slowly tilts upward to follow it into the sky, ending as it bursts into a large bloom of colorful sparks.\n\noverall_soundscape: Soft summer insect sounds, followed by the firework rising with a whistle and bursting with a distant boom and crackling sparks.\n\nnon_diegetic_music: N/A",
         "negative_prompt": "",

benchmarks/diffusion/bench_wan_vae_decode.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
 """Benchmark and exactness check for the Wan VAE decoder fast path.
 
 Loads a diffusers Wan VAE (by default the Cosmos3 one), decodes seeded latents
@@ -30,10 +30,13 @@ return the full video on rank 0 only, so equality and PSNR are computed there.
 from __future__ import annotations
 
 import argparse
+import json
 import math
 import os
+import statistics
 import time
 from collections.abc import Iterable
+from unittest.mock import patch
 
 import torch
 import torch.distributed as dist
@@ -82,6 +85,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--profile", action="store_true", help="Print a torch.profiler CUDA kernel table")
     parser.add_argument("--profile-rows", type=int, default=25)
     parser.add_argument("--save-output", default=None, help="Save the `off` decode output to this .pt path")
+    parser.add_argument(
+        "--first-frame-ablation", action="store_true", help="Compare channels_last with/without first-frame Conv2d"
+    )
+    parser.add_argument("--latents", help="For ablation: saved VAE-input tensor (.pt), after latent mean/std scaling")
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
@@ -254,8 +261,99 @@ def run_level(
     return output.detach(), stats
 
 
+@torch.inference_mode()
+def run_first_frame_ablation(args: argparse.Namespace, dtype: torch.dtype, device: torch.device) -> None:
+    """Same-model ABBA comparison; disable only the first-frame helper in the reference arm."""
+    from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards
+
+    if args.iters < 2 or args.warmup < 0:
+        raise ValueError("Ablation requires --iters >= 2 and --warmup >= 0")
+    vae = load_vae(args, dtype, device, parallel=False)
+    report = install_wan_vae_fastpath(vae, level="channels_last")
+    if not report.installed:
+        raise RuntimeError(f"channels_last installation failed: {report.reason}")
+    latents = make_latents(vae, args, dtype, device)
+    if args.latents:
+        saved = torch.load(args.latents, map_location="cpu", weights_only=True)
+        saved = saved[:, :, : latents.shape[2]]
+        if saved.shape != latents.shape:
+            raise ValueError(f"Saved latent shape {saved.shape} does not match {latents.shape}")
+        latents = saved.to(device=device, dtype=dtype).contiguous()
+    helpers = {"baseline": lambda conv, x, bias: None, "candidate": forwards._first_frame_conv2d}
+    optimized_calls = 0
+
+    def counted(conv, x, bias):
+        nonlocal optimized_calls
+        result = helpers["candidate"](conv, x, bias)
+        optimized_calls += result is not None
+        return result
+
+    with patch.object(forwards, "_first_frame_conv2d", counted):
+        decode(vae, latents)
+        sync()
+    if optimized_calls == 0:
+        raise RuntimeError("The selected workload did not exercise first-frame Conv2d")
+    measurements = []
+    outputs: dict[str, list[torch.Tensor]] = {arm: [] for arm in helpers}
+    for _ in range(args.warmup):
+        for helper in helpers.values():
+            with patch.object(forwards, "_first_frame_conv2d", helper):
+                decode(vae, latents)
+                sync()
+    for iteration in range(args.iters):
+        # Alternating AB and BA pairs give ABBA order without dropping odd runs.
+        arms = ("baseline", "candidate") if iteration % 2 == 0 else ("candidate", "baseline")
+        for arm in arms:
+            with patch.object(forwards, "_first_frame_conv2d", helpers[arm]):
+                sync()
+                torch.accelerator.reset_peak_memory_stats()
+                start = time.perf_counter()
+                output = decode(vae, latents)
+                sync()
+                elapsed = (time.perf_counter() - start) * 1000
+                peak = torch.accelerator.max_memory_allocated()
+            measurements.append({"arm": arm, "wall_ms": elapsed, "peak_allocated_bytes": peak})
+            if len(outputs[arm]) < 2:
+                outputs[arm].append(output.cpu())
+            del output
+    baseline, candidate = outputs["baseline"][0], outputs["candidate"][0]
+    difference = candidate.float() - baseline.float()
+    stats = {}
+    for arm in helpers:
+        times = [row["wall_ms"] for row in measurements if row["arm"] == arm]
+        stats[arm] = {
+            "median_ms": statistics.median(times),
+            "mean_ms": statistics.mean(times),
+            "std_ms": statistics.pstdev(times),
+            "min_ms": min(times),
+            "max_ms": max(times),
+        }
+    result = {
+        "scope": "VAE decode; same-head channels_last first-frame Conv2d ablation",
+        "args": vars(args),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "baseline": stats["baseline"],
+        "candidate": stats["candidate"],
+        "latency_reduction": 1 - stats["candidate"]["median_ms"] / stats["baseline"]["median_ms"],
+        "optimized_calls_per_decode": optimized_calls,
+        "outputs_finite": bool(torch.isfinite(baseline).all() and torch.isfinite(candidate).all()),
+        "self_repeat_equal": {arm: torch.equal(*values) for arm, values in outputs.items()},
+        "candidate_vs_baseline": {
+            "torch_equal": torch.equal(candidate, baseline),
+            "max_abs": difference.abs().max().item(),
+            "relative_l2": (difference.norm() / baseline.float().norm().clamp_min(1e-30)).item(),
+            "psnr_db": psnr(candidate, baseline),
+        },
+        "measurements": measurements,
+    }
+    print("FIRST_FRAME_ABLATION " + json.dumps(result))
+
+
 def main() -> None:
     args = parse_args()
+    if args.latents and not args.first_frame_ablation:
+        raise SystemExit("--latents requires --first-frame-ablation")
     if not torch.cuda.is_available():
         raise SystemExit("CUDA device required")
     rank, world_size = init_distributed()
@@ -267,6 +365,12 @@ def main() -> None:
         )
     device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda")
     dtype = DTYPES[args.dtype]
+    if args.first_frame_ablation:
+        if world_size != 1 or parallel:
+            raise SystemExit("--first-frame-ablation requires one GPU")
+        with torch.backends.cudnn.flags(benchmark=False, allow_tf32=False):
+            run_first_frame_ablation(args, dtype, device)
+        return
     levels: Iterable[str] = [level.strip() for level in args.fast_path.split(",") if level.strip()]
     for level in levels:
         if level not in VAE_FAST_PATH_LEVELS:

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_VERSION=v0.30.0
+ARG VLLM_ASCEND_VERSION=v0.31.0
 ARG VLLM_ASCEND_DEVICE_SUFFIX=
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
 

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_VERSION=v0.30.0
+ARG VLLM_ASCEND_VERSION=v0.31.0
 ARG VLLM_ASCEND_DEVICE_SUFFIX=
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
 

docker/Dockerfile.xpu

@@ -1,4 +1,4 @@
-# Target vLLM version: v0.30.0. Keep VLLM_VERSION aligned with the target vLLM
+# Target vLLM version: v0.31.0. Keep VLLM_VERSION aligned with the target vLLM
 # release on every rebase.
 #
 # This file layers vLLM-Omni on top of the published upstream XPU image, which
@@ -18,7 +18,7 @@
 #
 # .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu
-ARG VLLM_VERSION=v0.30.0
+ARG VLLM_VERSION=v0.31.0
 ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION}
 
 FROM ${VLLM_BASE} AS vllm-omni

docs/configuration/stage_configs.md

@@ -598,8 +598,9 @@ the intended GPU, and warm the complete pipeline before measuring performance.
 
 ## MOSS-TTS Local 1.5 with Model Runner V2
 
-`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
-the shared runtime introduced for Qwen3-TTS:
+`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` defaults to CUDA MRV2 using
+the shared runtime introduced for Qwen3-TTS. To select the C64 profile without
+MPS explicitly:
 
 ```bash
 vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
@@ -607,7 +608,7 @@ vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
 ```
 
 This profile inherits the batching, codec graph buckets and 1-frame/15-frame
-chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
+chunk geometry from `moss_tts_local_v1.yaml`, and selects V2 for both the Local
 Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
 capabilities while retaining its V1 `talker_mtp` implementation, sampling
 defaults and explicit request-seed handling. Codec chunks use the native data
@@ -631,8 +632,12 @@ guarantee that 128 maximum-length prompts fit simultaneously. See the
 [MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
 for activation, backend comparisons, memory requirements and benchmark commands.
 
-Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
-NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
+Omitting `--deploy-config` selects the C128 system profile on CUDA GPUs with
+at least 140 GiB memory and MPS available. Smaller GPUs or a failed memory
+query use C64 with utilization-based memory budgets; missing MPS selects
+C64 without MPS. Selecting `moss_tts_local.yaml` explicitly uses C64 with MPS.
+Use `moss_tts_local_v1.yaml` for the previous V1 deployment.
+NPU, XPU, ROCm and MUSA overrides retain V1. These profiles do not enable
 MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
 set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
 `VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
@@ -640,5 +645,5 @@ set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
 Event-driven orchestration remains independently selectable with
 `VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
 when comparing these modes. Model-runner selection does not change the
-orchestration default or enable experimental reference encoding, chunk ramps,
-generation-output draining or MPS.
+orchestration default. MPS, codec dispatch and graph settings belong to the
+deployment profile; the automatic CUDA default selects MPS when available.

docs/design/feature/sequence_parallel.md

@@ -281,6 +281,10 @@ NOTE: be careful to test adequately when refactoring classes that take this styl
 | `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
 | `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors |
 | `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) |
+| `clone_shard` | bool | Give multi-rank shards their own contiguous storage (default: `False`); single-rank inputs are unchanged |
+
+With `clone_shard=True`, memory savings require callers to drop their references to the full tensor.
+Other references to the same storage also keep the full allocation alive.
 
 **SequenceParallelOutput Parameters:**
 

docs/design/minicpm_o45_mrv2_performance.md

@@ -12,11 +12,25 @@ require new end-to-end performance measurements after the mainline integration.
 Previously reported numbers do not describe this revised codec backend.
 
 The generic MRv2 profile and the opt-in native V1 duplex H200 profile enable
-TF32 matmuls within Code2Wav forward/capture only; the previous process policy
-is restored afterwards. cuDNN's TF32 policy is independent. TF32 changes rounding.
+ordinary TF32 for CFM DiT GEMMs within Code2Wav forward/capture only
+(`torch.backends.cuda.matmul.allow_tf32` on dense QKV/MLP, Triton
+`input_precision="tf32"` on tiled attention). This is not compensated TF32x3.
+The previous process matmul policy is restored afterwards. cuDNN's TF32 policy
+is independent. HiFT stays IEEE FP32. TF32 changes rounding.
 Duplex keeps the mainline V1 session path. Turn results do not establish duplex
 performance or interruption correctness.
 
 The shared asynchronous output snapshot and batched Talker preprocessing also
 apply to V1 when async chunking/scheduling is enabled. Default V1 behavior must
 therefore be included in end-to-end regression validation.
+
+CFM graph execution defaults to the fused DiT body on CUDA when TF32 is
+allowed. NVIDIA SM80+ uses tiled TF32 attention; older CUDA devices use the
+fused body's SDPA fallback. Unsupported DiT layouts or non-FP32 attention
+caches retain the original body. Other platforms retain their existing default.
+Set `cfm_fused_body: false` in connector `extra` to restore the original body.
+Disabling TF32 with `token2wav_allow_tf32: false` or
+`MINICPMO_CODE2WAV_TF32=off` also disables the default fusion choice; an explicit
+`cfm_fused_body: true` overrides that choice and may use TF32 attention.
+Slot pooling and row-offset merging remain opt-in. Fusion changes floating-point
+rounding and does not promise identical waveforms.

docs/design/minicpmo_encoder_cuda_graph.md

@@ -0,0 +1,61 @@
+# MiniCPM-o 4.5 input encoder CUDA graphs
+
+The vision transformer and stateless audio encoder reuse vLLM's encoder graph
+manager for exact input shapes. Each encoder retains at most four graphs by
+default, and captures on the second call with the same shape and CUDA stream.
+Startup profiling can consume slots. Graphs are never evicted: once the cache
+is full, other shapes run eagerly for the lifetime of the model.
+
+Use `--hf-overrides` to configure both encoders:
+
+```json
+{
+  "encoder_cuda_graph": true,
+  "encoder_cuda_graph_max_graphs": 4,
+  "encoder_cuda_graph_min_capture_calls": 2,
+  "encoder_cuda_graph_min_free_bytes": 1073741824,
+  "encoder_cuda_graph_share_pools": true
+}
+```
+
+`max_graphs` is a nonnegative integer per encoder; zero disables capture.
+`min_capture_calls` is an integer of at least two. Increasing it reduces capture
+of short-lived shapes. Admission history is bounded to four times `max_graphs`,
+so an evicted history entry must accumulate its calls again. Increasing the
+graph cap can improve coverage but retains more GPU memory. Existing graphs
+are not replaced, avoiding repeated capture costs when traffic changes.
+
+By default, graphs in the same encoder on the same device and replay stream
+share a memory pool and capture stream. Different encoders and replay streams
+remain isolated. Graph inputs and output buffers are allocated outside capture;
+only transient capture allocations share the pool. Returned outputs remain
+cloned. A late capture waits for earlier replay work and output clones before
+reusing the pool. Set `encoder_cuda_graph_share_pools` to `false` for an A/B
+comparison with a separate pool and capture stream per graph.
+
+The graph count is **not a byte limit**. Sharing pools reuses intermediate
+storage but does not eliminate per-shape static input and output buffers.
+Before a new capture, the adapter checks device-free memory against
+`min_free_bytes`, a nonnegative integer (default: 1 GiB; zero disables the check).
+Below this floor the call runs eagerly, and a later call may try admission
+again. Existing graphs can still replay. This check is conservative about
+reusable allocator memory and does not predict capture allocations, reserve
+memory, or guarantee that capture fits. Size the floor for the deployment;
+disable encoder graphs when their additional memory cannot be accommodated.
+
+Each adapter exposes `get_cumulative_stats()`: graph hits, total misses, hit
+rate, graph count, and separate ineligible, capacity, warmup and memory misses.
+These include calls rejected before the upstream manager, so a full cache's
+lost coverage is visible. Vision uses `vpm._encoder_graph`; audio uses
+`_audio_encoder_graph` on the model. The counters are per adapter, not global.
+
+Capture failures propagate the original exception and make that adapter
+unusable. Subsequent calls raise without retrying capture or running CUDA eager
+work. Restart the worker, optionally with encoder graphs disabled. This is a
+fatal contract, not automatic recovery: a failed CUDA capture can leave the
+device context unusable. `capture_failures` records this terminal state.
+
+`--enforce-eager` always disables these graphs. CPU, grad/autocast, stateful
+streaming audio and the other model-specific eager fallbacks remain unchanged.
+The adapter clones graph outputs through upstream postprocessing so retained
+embeddings survive later replays.

docs/getting_started/installation/gpu/xpu.inc.md

@@ -1,8 +1,11 @@
+<!-- markdownlint-disable MD001 MD025 -->
+
 # --8<-- [start:requirements]
 
 - GPU: Validated on Intel® Arc™ B-Series.
 
 # --8<-- [end:requirements]
+
 # --8<-- [start:set-up-using-python]
 
 vLLM-Omni currently recommends using the Docker image setup steps below.
@@ -26,9 +29,8 @@ DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-s
 ```
 
 This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
-base image, which Docker pulls automatically. The Dockerfile still defaults to
-v0.30.0; the command above explicitly selects v0.31.0 to match this source
-checkout. To target a different vLLM release, pass
+base image, which Docker pulls automatically. The Dockerfile defaults to
+v0.31.0. To target a different vLLM release, pass
 `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
 the base from upstream's own Dockerfile first and point the build at it:
 
@@ -44,7 +46,7 @@ DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-s
 
 ##### Launch with OpenAI API Server
 
-```
+```bash
 docker run -it -d --shm-size 10g \
   --name {container_name} \
   --net=host \

docs/getting_started/installation/npu/npu.inc.md

@@ -10,25 +10,25 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 === "A2"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.31.0
     ```
 
 === "A3"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a3
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.31.0-a3
     ```
 
 === "A5"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a5
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.31.0-a5
     ```
 
 === "310P (Experimental)"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-310p
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.31.0-310p
     ```
 
 ```bash

docs/models/supported_models.md

@@ -29,9 +29,11 @@ th {
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `JoyImageEditPipeline` | JoyAI-Image-Edit | `jdopensource/JoyAI-Image-Edit-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-Image-Edit.md) |
 | `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tongyi-MAI/Z-Image-Turbo.md) |
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) |
@@ -54,12 +56,12 @@ th {
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
 | `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
-| `SeedVR2Pipeline` | [SeedVR2 3B](seedvr2.md) video restoration | `ByteDance-Seed/SeedVR2-3B` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/ByteDance/SeedVR2-RTX-5090.md) |
+| `SeedVR2Pipeline` | [SeedVR2 3B](seedvr2.md) video restoration | `ByteDance-Seed/SeedVR2-3B` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/ByteDance/SeedVR2-RTX-5090.md) |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
-| `Pi05Pipeline` | π0.5 VLA | `lerobot/pi05_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi05.md) |
+| `Pi05Pipeline` | π0.5 VLA | `lerobot/pi05_base` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi05.md) |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -70,7 +72,7 @@ th {
 | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/StabilityAI/Stable-Diffusion-3.5.md) |
 | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) |
+| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) |
 | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
 | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
 | `BreezeForConditionalGeneration` | Breeze-TTS-2 (voice design, voice clone, voice direction, CFG) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
@@ -82,8 +84,8 @@ th {
 | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
-| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | | [Recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/YuE2-3B.md) / [Audio-to-score](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/SheetSage2-H200.md) |
+| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | ✅︎ | — |
+| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | ✅︎ | [Recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/YuE2-3B.md) / [Audio-to-score](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/SheetSage2-H200.md) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -93,7 +95,7 @@ th {
 | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
 | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
 | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
-| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
+| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
 | `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
 | `IndexTTS25TalkerForConditionalGeneration` | IndexTTS-2.5 | native `checkpoints/` bundle | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/IndexTeam/IndexTTS-2_5.md) |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -106,6 +108,7 @@ th {
 | `SenseNovaU1Pipeline` | SenseNova-U1-A3B (MoE) | `sensenova/SenseNova-U1-A3B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.md#a3b-moe-variant) |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) |
+| `Kandinsky6TI2VAPipeline` | Kandinsky 6 TI2VA | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Kandinsky/Kandinsky6-TI2VA.md) |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
@@ -115,7 +118,7 @@ th {
 | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
 | `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
 | `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
-| `AnimaPipeline` | Anima | `circlestone-labs/Anima` | ✅︎ | ✅︎ | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/circlestone-labs/Anima.md) |
+| `AnimaPipeline` | Anima | `circlestone-labs/Anima` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/circlestone-labs/Anima.md) |
 
 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the

docs/serving/README.md

@@ -154,18 +154,16 @@ All six routes are model- or configuration-dependent. In particular,
 pipeline declares a `duplex_plugin` and its deploy configuration sets
 `session_mode: duplex` (with `session_mode: turn` the same model boots the
 ordinary turn-based serving stack instead). On a duplex server a stock
-Realtime client needs no vendor query parameter; `duplex=0` selects the
-turn-based Realtime handler, which a duplex server does not mount, so that
-connection is refused with `Realtime API is not available`. Such a server
-serves the websocket route plus `POST /v1/chat/completions`, `/v1/models`
-and `/health`, and no other turn-based HTTP route. The chat route is the
-ordinary chat service running on the duplex engine: a request is a
-turn-based generation on the same stages, served alongside the live
-websocket sessions; it opens no duplex session and holds no
-`duplex_session.max_sessions` slot -- see [Full Duplex](full_duplex_api.md).
-On a server that is not duplex, an explicit `?duplex=1` is refused rather
-than answered by the turn-based handler, so a client that requested duplex
-never silently gets the other protocol.
+Realtime client needs no vendor query parameter; `duplex=0` opts out of native
+duplex and selects the turn-based Realtime handler. It runs generation on the
+same stages, alongside live websocket sessions, and opens no duplex session
+or consumes a `duplex_session.max_sessions` slot -- see
+[Full Duplex](full_duplex_api.md).
+Such a server serves the websocket route plus `POST /v1/chat/completions`,
+`/v1/models` and `/health`, and no other turn-based HTTP route. On a server
+that is not duplex, an explicit `?duplex=1` is refused rather than answered by
+the turn-based handler, so a client that requested duplex never silently gets
+the other protocol.
 
 ## Related Endpoints
 

docs/serving/speech_api.md

@@ -183,6 +183,16 @@ emit these headers.
 **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
 WAV) as they are decoded.
 
+`pcm` is 16-bit signed little-endian, interleaved, at the model's native sample
+rate and channel count (24 kHz mono for most models; MOSS-TTS-Local v1.5 emits
+48 kHz stereo). Raw streams and non-streaming responses state the format in
+response headers:
+
+| Header | Description |
+| --- | --- |
+| `x-audio-sample-rate` | Sample rate of the returned audio in Hz (after any `sample_rate` resampling). |
+| `x-audio-channels` | Channel count of the returned audio. |
+
 Both streaming modes (`stream_format="audio"` and `"sse"`) require
 `response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted).
 
@@ -193,9 +203,13 @@ Each event has an `event:` line and a JSON `data:` line:
 - `speech.audio.delta` — a base64 audio chunk:
 
     ```json
-    { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm" }
+    { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm",
+      "sample_rate": 24000, "channels": 1 }
     ```
 
+    `sample_rate` and `channels` extend the OpenAI event and describe the
+    decoded chunk.
+
 - `speech.audio.done` — terminal event, carrying token `usage`:
 
     ```json
@@ -340,7 +354,7 @@ Server -> Client:
 
 | Message | Description |
 | --------- | ------------- |
-| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
+| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000, "channels": 1}` | Audio generation starting for the buffered input. With `stream_audio=true` it is sent with the first audio chunk, and for `pcm` it carries that audio's `sample_rate` and `channels` (omitted when no audio was produced). |
 | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
 | `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
 | `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete |

docs/serving/videos_api.md

@@ -237,13 +237,22 @@ At least one mask is required. A nontrivial `video_noise_mask` requires
 `source_audio` or a `source_video` with an audio stream. Mask values are in
 `[0, 1]`: `0` preserves the source, `1` regenerates it, and fractional values
 blend the two behaviors. Exact all-one masks are no-ops and do not require a
-source. Source uploads without a mask are rejected. Masks may be a JSON scalar
-or arrays matching the H3 latent/token grid; pixel-resolution masks must be
-resized or pooled by the client before upload.
+source. Source uploads without a mask are rejected. Masks may be a JSON scalar,
+an array matching the H3 latent/token grid, or (video only) a frame-space
+array that the server resizes to the latent grid.
 
 For an aligned output of `F` frames at `W x H`, the video latent grid is
 `[Tv, H/16, W/16]`, where `Tv = 2 + 5 * ((F - 5) / 17)`. The video mask may be
-a scalar, a flat token vector, `[Tv, H/32, W/32]`, or the full latent grid. For
+a scalar, a flat token vector, `[Tv, H/32, W/32]`, the full latent grid, a
+spatial `[h, w]` mask applied to every frame, or a frame-space `[T, h, w]` mask
+with one slice per 24 fps output frame. A shape that matches the token or full
+latent grid is always read as that grid. Frame-space masks are area-resized to
+`[H/16, W/16]`; along time, a mask shorter than `F` is padded with its last
+slice (as the source video is) and a longer one is truncated, then each latent
+takes the maximum over the frames the causal video VAE folds into it, so any
+regenerated frame regenerates its latent. Because the 8 MiB limit below applies
+to the JSON file, clients should area-downsample frame-space masks by 16
+spatially before upload. For
 the model input, timestep, and velocity, a full-grid mask is max-pooled over
 each 2x2 spatial token and fractional values are rounded upward to 1/256
 levels. The final x0 restore uses the original, unquantized mask, so full-grid

docs/user_guide/diffusion/cpu_offload.md

@@ -42,6 +42,7 @@ higher-priority strategy is selected.
 | --- | --- | --- | :---: | :---: | :---: | --- |
 | Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
 | LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
+| JoyImageEditPipeline | `jdopensource/JoyAI-Image-Edit-Diffusers` | `JoyImageEditTransformer3DModel` | yes | yes | — | `double_blocks` |
 | NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
 | OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
 | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -1,12 +1,12 @@
 # VAE Parallelism Guide
 
-
 ## Table of Content
 
 - [Overview](#overview)
 - [Quick Start](#quick-start)
 - [Example Script](#example-script)
 - [Configuration Parameters](#configuration-parameters)
+- [Batch Parallel Decode](#batch-parallel-decode)
 - [Best Practices](#best-practices)
 - [Troubleshooting](#troubleshooting)
 - [Summary](#summary)
@@ -15,23 +15,22 @@
 
 ## Overview
 
-VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers patch/tile parallelism, batch parallel decode for image VAEs, and Wan spatial-shard decode. Patch/tile parallelism divides images into spatial regions; batch parallel decode assigns complete images to ranks.
 
 This is particularly useful for:
+
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
 - **Memory-constrained environments** where the VAE decode activation peak exceeds available VRAM
 - **Multi-GPU setups** where you want to leverage distributed resources for the VAE stage
 
 See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
 
-
 VAE patch parallelism uses two strategies based on image size:
 
-| Strategy | Use Case | How It Works | Overlap Handling | Output Quality |
-|----------|----------|--------------|------------------|----------------|
-| **Tiled Decode** | Large images (triggers VAE tiling) | Distributes existing VAE tiling computation across ranks. Each rank decodes a subset of overlapping tiles. | Uses VAE's native `blend_v` and `blend_h` functions to seamlessly merge overlapping regions | Bit-identical (same logic as single-GPU tiling) |
-| **Patch Decode** | Small images (no VAE tiling) | Splits latent into spatial patches with halos. Each rank decodes one patch with boundary context. | Halo regions provide edge context; core regions are directly stitched without blending | Near-identical (diff < 0.5%, visually imperceptible) |
-
+| Strategy         | Use Case                           | How It Works                                                                                               | Overlap Handling                                                                            | Output Quality                                       |
+| ---------------- | ---------------------------------- | ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- | ---------------------------------------------------- |
+| **Tiled Decode** | Large images (triggers VAE tiling) | Distributes existing VAE tiling computation across ranks. Each rank decodes a subset of overlapping tiles. | Uses VAE's native `blend_v` and `blend_h` functions to seamlessly merge overlapping regions | Bit-identical (same logic as single-GPU tiling)      |
+| **Patch Decode** | Small images (no VAE tiling)       | Splits latent into spatial patches with halos. Each rank decodes one patch with boundary context.          | Halo regions provide edge context; core regions are directly stitched without blending      | Near-identical (diff < 0.5%, visually imperceptible) |
 
 VAE Patch Parallelism **reuses the DiT process group** (`dit_group`) and does not initialize a separate ProcessGroup. This means:
 
@@ -110,19 +109,80 @@ vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
 
 In `DiffusionParallelConfig`:
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
+| Parameter                 | Type   | Default   | Description                                                                                                                                                       |
+| ------------------------- | ------ | --------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `vae_patch_parallel_size` | int    | 1         | Number of ranks requested for VAE parallelism, including batch mode. Set to 2 or higher to enable. Batch mode caps this degree at the existing worker group size. |
+| `vae_parallel_mode`       | str    | `"tile"`  | `"tile"` for tile/patch decode, `"batch"` for complete-image batch decode, or `"spatial_shard_height"` / `"spatial_shard_width"` for Wan spatial-shard decode.    |
 
 Additional requirements:
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. |
+| Parameter        | Type   | Default   | Description                                             |
+| ---------------- | ------ | --------- | ------------------------------------------------------- |
+| `vae_use_tiling` | bool   | False     | Must be set to `True` when using VAE patch parallelism. |
 
 !!! note "Automatic VAE Tiling"
-    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled.
+    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), tile and spatial-shard modes automatically enable `vae_use_tiling`. Batch mode preserves the requested tiling setting.
+
+---
+
+## Batch Parallel Decode
+
+Set `vae_parallel_mode="batch"` to distribute complete images across ranks during
+VAE decode. This mode supports pipelines that load `DistributedAutoencoderKL`
+or `DistributedAutoencoderKLFlux2`, including FLUX.2-dev and FLUX.2 Klein. Klein
+selects the distributed VAE only for batch mode; its default path is unchanged.
+Selecting this mode for an
+unsupported VAE, including Wan, raises an error at model initialization.
+
+For a pipeline already using two tensor-parallel workers, select batch decode
+with these Python settings:
+
+```python
+parallel_config = DiffusionParallelConfig(
+    tensor_parallel_size=2,
+    vae_patch_parallel_size=2,
+    vae_parallel_mode="batch",
+)
+sampling_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=2)
+```
+
+The corresponding serving flags are `--vae-patch-parallel-size 2
+--vae-parallel-mode batch`. Keep the existing DiT parallelism configuration;
+the VAE degree does not allocate extra workers. Send a request for multiple
+images to use the batch split.
+
+Each rank decodes a contiguous part of the image batch through its native VAE
+decoder. Results are reassembled in the original image order and made available
+on every worker. With two images and four workers, two workers decode one image
+each. With five images and four workers, the assigned batch sizes are 2, 1, 1,
+and 1. Workers without images still participate in the required communication.
+
+**Constraints and behavior:**
+
+- Batch mode and distributed tile/patch mode are mutually exclusive. Encoding
+  uses the native local path, including native tiling when enabled, rather than
+  distributing spatial tiles across workers.
+- Batch size one, VAE degree one, or a single worker uses native decode on each
+  worker. Batch mode does not automatically switch to distributed tile/patch
+  decode for a single image, even with `vae_use_tiling=True`. To distribute a
+  single image across workers, select `vae_parallel_mode="tile"` on a pipeline
+  that supports distributed tile/patch decoding.
+- Data parallelism, pipeline parallelism, and CFG parallelism must each have
+  degree one because this implementation communicates over the worker WORLD
+  group. Tensor and sequence parallelism may supply the shared workers.
+- Batch mode does not force tiling. Explicit native VAE tiling and slicing
+  settings remain in effect within each rank's assigned batch and on the
+  single-image fallback. Native tiling splits an image locally on one worker;
+  it does not add cross-worker tile parallelism to batch mode.
+- It does not merge independent serving requests. The VAE must receive a batch
+  containing multiple images, for example through `num_outputs_per_prompt`.
+
+Splitting the batch can reduce per-rank decoder work and activation memory, but
+collecting decoded images adds communication and buffers. Measure complete
+request latency and memory for the intended model, image count, and resolution.
+This feature has no universal speedup or memory-saving guarantee. Changing the
+decoder batch shape can also change floating-point rounding; validate generated
+outputs when selecting a deployment configuration.
 
 ---
 
@@ -194,7 +254,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 ### Common Issue 1: Model Not Support VAE Patch Parallel
 
 **Symptoms**:
-```
+
+```text
 WARNING: vae_patch_parallel_size=2 is set but VAE patch parallelism is NOT enabled for xxxPipeline; ignoring.
 ```
 
@@ -217,7 +278,6 @@ if vae_pp_size > 1 and not is_distributed_vae:
 
 2. To add support for a new model, implement `DistributedVaeMixin` on its VAE class (contributions are welcome).
 
-
 ### Common Issue 2: `vae_patch_parallel_size` Exceeds DiT Process Group Size
 
 **Symptoms**: Shows warning message, and vae patch parallel size is resized to DiT process group size
@@ -227,6 +287,7 @@ if vae_pp_size > 1 and not is_distributed_vae:
 **Recommendation**: Always set `vae_patch_parallel_size` to be no greater than your DiT process group size.
 
 Note that the size of DiT process group size equals to:
+
 ```text
 dit_parallel_size = data_parallel_size
                   × cfg_parallel_size
@@ -235,7 +296,8 @@ dit_parallel_size = data_parallel_size
                   × tensor_parallel_size
 
 ```
-_sequence_parallel_size = ulysses_degree × ring_degree_
+
+`sequence_parallel_size = ulysses_degree × ring_degree`
 
 ---
 

docs/user_guide/diffusion/vae_fast_path.md

@@ -7,7 +7,8 @@ convolutional network whose reference implementation in diffusers moves every
 activation through memory several times per layer (normalization, activation,
 causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
 installs a fast path on every loaded Wan VAE that fuses this data movement into
-a handful of Triton kernels while leaving the convolutions themselves untouched.
+a handful of Triton kernels. The opt-in `channels_last` level also specializes
+first-frame convolutions that have no history.
 
 ## Levels
 
@@ -15,9 +16,9 @@ The fast path is controlled by `--vae-fast-path` (engine argument
 `vae_fast_path`, deploy-config key `vae_fast_path`):
 
 | Level | Default | Output vs. diffusers | What it does |
-|---|---|---|---|
+| --- | --- | --- | --- |
 | `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
-| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus decoder convolution weights converted to channels-last memory format so cuDNN picks its channels-last kernels, and a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the bias of the preceding `conv1`. |
+| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus channels-last convolution weights, a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the preceding `conv1` bias, and Conv2d for supported single-frame causal convolutions without history. |
 | `off` | no | bit-identical | Reference diffusers decoder. |
 
 ```bash
@@ -61,6 +62,12 @@ stages:
 - The `channels_last` level changes the order in which cuDNN accumulates
   convolutions, so outputs differ from the reference in the last bits. Use
   `lossless` when bitwise reproducibility against diffusers matters.
+- For a supported causal convolution with one input frame and no history,
+  the temporal input is `[0, 0, x]`. The `channels_last` level uses Conv2d with
+  the last temporal weight slice, preserving the convolution formula for
+  finite inputs and the existing cache update. Calls with history retain
+  Conv3d. The weight slice is taken on each call, so weight updates and device
+  moves do not leave a stale packed copy.
 - For VAEs kept in fp32, cuDNN's channels-last convolution algorithms run in
   TF32 under PyTorch's default `torch.backends.cudnn.allow_tf32 = True`, which
   dominates the difference to the reference (about 1e-3). Set
@@ -73,10 +80,10 @@ Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_dec
 (best of 2 runs after warmup):
 
 | `--vae-fast-path` | Decode time | Speedup | Output vs. `off` |
-|---|-------------|---------|---|
-| `off` | 6.00 s      | 1.00x   | reference |
-| `lossless` | 3.26 s      | 1.84x   | bit-identical |
-| `channels_last` | 2.51 s      | 2.39x   | PSNR 62.6 dB, max abs diff 4.6e-2 |
+| --- | ------------- | --------- | --- |
+| `off` | 6.00 s | 1.00x | reference |
+| `lossless` | 3.26 s | 1.84x | bit-identical |
+| `channels_last` | 2.51 s | 2.39x | PSNR 62.6 dB, max abs diff 4.6e-2 |
 
 Peak decode memory was unchanged (about 13 GiB): it is set by the largest
 activations and the cuDNN workspace, not by the output assembly.
@@ -93,6 +100,22 @@ python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano
     --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile
 ```
 
+For a single-GPU comparison that isolates first-frame Conv2d within
+`channels_last`, use `--first-frame-ablation`. It alternates the two paths,
+disables cuDNN TF32 and autotuning, and prints JSON with every timing sample,
+summary statistics, peak allocated memory and output comparisons. Per-call
+weight packing is included in the measured decode time.
+
+```bash
+python benchmarks/diffusion/bench_wan_vae_decode.py \
+    --model Wan-AI/Wan2.1-T2V-1.3B-Diffusers --size 832x480 --frames 1 \
+    --dtype bf16 --first-frame-ablation --warmup 60 --iters 50 --seed 0
+```
+
+The default input is a seeded latent tensor. `--latents /path/to/latents.pt`
+can instead load a saved VAE input after latent mean/std scaling; the benchmark
+selects the temporal prefix requested by `--frames` and checks its shape.
+
 The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
 the decode is timed across all ranks and rank 0 reports:
 

docs/user_guide/diffusion_features.md

@@ -152,6 +152,7 @@ The following tables show which models support each feature:
 | **HiDream-O1-Image**     |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ✅*        |
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **JoyAI-Image-Edit**    |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ✅     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
@@ -163,6 +164,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**       |     ❌     |     ❌      |           ❓           |       ✅        |         ✅         |          ❌          |   ❓    |             ❓             |      ✅ (decode)      |       ✅        |        ✅         |
 | **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **SenseNova-U1.5**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ✅        |        ❌         |
 | **SenseNova-U1-A3B**     |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
@@ -181,7 +183,8 @@ 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. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident.
+> 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`. Ulysses sequence parallelism is wired and ring attention is not supported, but multi-GPU padding e2e and a parity number are not in this PR's H200 results, so SP stays ❓. 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. Online FP8 is documented in the quantization guide with H200 numbers.
+> 9. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident.
 
 ### VideoGen
 
@@ -196,6 +199,7 @@ The following tables show which models support each feature:
 |  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |   ✅ (`--usp`, frame-sharded)   |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
+| **Kandinsky 6 TI2VA**        | ✅* (uncalibrated) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ (FP8) | ✅ (1 request) |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -3,7 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_image>.
 
 
-This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
+This example edits an input image with `Qwen/Qwen-Image-Edit` and `JoyAI-Image-Edit` using the `image_edit.py` CLI.
 
 ## Local CLI Usage
 
@@ -26,6 +26,27 @@ python image_edit.py \
   --cfg-scale 4.0
 ```
 
+### JoyAI-Image-Edit
+
+JoyAI-Image-Edit uses the same `image_edit.py` CLI and supports one input
+image per request. Requested dimensions are snapped to the nearest supported
+Joy bucket; for square outputs, use `1024x1024`.
+
+```bash
+python image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image input.png \
+  --prompt "Change the background to a clean studio while preserving the subject." \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --guidance-scale 4.0 \
+  --output output_joyai_edit.png
+```
+
+JoyAI defaults to guidance scale `4.0`; pass `--guidance-scale` explicitly when
+you need a different value.
+
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
 
 For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or  `Qwen/Qwen-Image-Edit-2511`:
@@ -43,7 +64,7 @@ python image_edit.py \
 
 Key arguments:
 
-- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
+- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. `jdopensource/JoyAI-Image-Edit-Diffusers` supports one input image.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative-prompt`: text description (string).
 - `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.

docs/user_guide/examples/online_serving/image_to_image.md

@@ -3,10 +3,14 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/image_to_image>.
 
 
-This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni.
+This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni.
+
+Supported models include Qwen-Image-Edit, BAGEL, JoyAI-Image-Edit, and other image-to-image pipelines.
 
 For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content.
 
+For **JoyAI-Image-Edit**, use `jdopensource/JoyAI-Image-Edit-Diffusers` and send exactly one input image.
+
 ## Start Server
 
 ### Basic Start
@@ -24,6 +28,12 @@ vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
 vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092
 ```
 
+### JoyAI-Image-Edit
+
+```bash
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers --omni --port 8092
+```
+
 ### Start with Parameters
 
 
@@ -76,6 +86,11 @@ curl -s http://localhost:8092/v1/chat/completions \
   | cut -d',' -f2 | base64 -d > output.png
 ```
 
+For JoyAI-Image-Edit, set classifier-free guidance with `guidance_scale`
+in `extra_body` (default: `4.0`). `true_cfg_scale` remains supported as an
+alias; if both are provided, their values must match. Values at or below `1.0`
+disable CFG.
+
 ### Method 2: Using OpenAI Python SDK
 
 ```python

docs/user_guide/quantization/fp8.md

@@ -1,4 +1,4 @@
-# FP8 Quantization
+# FP8 Quantization
 
 ## Overview
 
@@ -17,7 +17,7 @@ in deep DiT blocks.
 ## Hardware Support
 
 | Device | Support |
-|--------|---------|
+| -------- | --------- |
 | NVIDIA Blackwell GPU (SM 100+) | ✅ |
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ |
@@ -93,8 +93,9 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 ### Diffusion Models
 
 | 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. On H200 (1024×1024, teapot prompt) DiT FP8-all is ~30.4 GB / 6.06 s vs BF16 ~36.9 GB / 7.34 s (~33 dB vs BF16). 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 | |
@@ -110,20 +111,20 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
-| Model | Scope | Format | Status |
-|-------|-------|--------|--------|
+| Model      | Scope                        | Format                                   | Status                              |
+| ---------- | ---------------------------- | ---------------------------------------- | ----------------------------------- |
 | Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction |
-| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
+| Qwen3-TTS  | TTS language-model stage     | Checkpoint config                        | Not validated                       |
 
 Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
 model-specific guide says otherwise.
 
 ### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
 
-| Model | Scope | Status | Notes |
-|-------|-------|--------|-------|
-| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
-| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
+| Model     | Scope                                    | Status        | Notes                                           |
+| --------- | ---------------------------------------- | ------------- | ----------------------------------------------- |
+| BAGEL     | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline          |
 
 ## Configuration
 
@@ -162,7 +163,7 @@ For a pipeline that exposes both a transformer and a quantization-aware text
 encoder, the scope is:
 
 | Configuration | Transformer | Text encoder | Components without supported quantizable layers |
-|---------------|-------------|--------------|-------------------------------------------------|
+| --------------- | ------------- | -------------- | ------------------------------------------------- |
 | `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
 | `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
 | `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
@@ -211,7 +212,7 @@ does not fit the workload.
 ## Parameters
 
 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
 | `method` | str | - | Quantization method (`"fp8"`) |
 | `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
 | `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available |

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.
@@ -39,6 +39,20 @@ Usage (multiple images):
         --cfg-scale 4.0 \
         --guidance-scale 1.0
 
+Usage (JoyAI-Image-Edit, single image):
+    python image_edit.py \
+        --model jdopensource/JoyAI-Image-Edit-Diffusers \
+        --image input.png \
+        --prompt "Change the background to a clean studio while preserving the subject." \
+        --height 1024 \
+        --width 1024 \
+        --num-inference-steps 50 \
+        --guidance-scale 4.0 \
+        --output output_joyai_edit.png
+
+    Note: JoyAI-Image-Edit snaps requested dimensions to the nearest supported
+    Joy bucket; for square outputs, use 1024x1024.
+
 Usage (with cache-dit acceleration):
     python image_edit.py \
         --image input.png \
@@ -209,14 +223,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(
@@ -270,9 +284,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--cfg-scale",
         type=float,
-        default=4.0,
+        default=None,
         help=(
-            "True classifier-free guidance scale (default: 4.0). Guidance scale as defined in Classifier-Free "
+            "True classifier-free guidance scale. When omitted, the pipeline uses its model-specific default. "
+            "Guidance scale as defined in Classifier-Free "
             "Diffusion Guidance. Classifier-free guidance is enabled by setting cfg_scale > 1 and providing "
             "a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, "
             "usually at the expense of lower image quality."
@@ -281,9 +296,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--guidance-scale",
         type=float,
-        default=1.0,
+        default=None,
         help=(
-            "Guidance scale for guidance-distilled models (default: 1.0, disabled). "
+            "Guidance scale for guidance-distilled models. When omitted, the pipeline uses its model-specific default. "
             "Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale "
             "directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
         ),
@@ -368,7 +383,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
@@ -523,6 +538,12 @@ def parse_args() -> argparse.Namespace:
             "silently masking with a possibly-wrong prompt."
         ),
     )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=600,
+        help="Overall pipeline initialization timeout in seconds.",
+    )
     return parser.parse_args()
 
 
@@ -593,6 +614,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
+        init_timeout=args.init_timeout,
     )
     if args.enforce_eager is not None:
         omni_kwargs["enforce_eager"] = args.enforce_eager
@@ -647,6 +669,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/image_to_image.md

@@ -1,6 +1,6 @@
 # Image-To-Image
 
-This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
+This example edits an input image with `Qwen/Qwen-Image-Edit` and `JoyAI-Image-Edit` using the `image_edit.py` CLI.
 
 ## Local CLI Usage
 
@@ -23,6 +23,27 @@ python image_edit.py \
   --cfg-scale 4.0
 ```
 
+### JoyAI-Image-Edit
+
+JoyAI-Image-Edit uses the same `image_edit.py` CLI and supports one input
+image per request. Requested dimensions are snapped to the nearest supported Joy
+bucket; for square outputs, use 1024x1024.
+
+```bash
+python image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image input.png \
+  --prompt "Change the background to a clean studio while preserving the subject." \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --guidance-scale 4.0 \
+  --output output_joyai_edit.png
+```
+
+JoyAI defaults to guidance scale `4.0`; pass `--guidance-scale` explicitly when
+you need a different value.
+
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
 
 For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or  `Qwen/Qwen-Image-Edit-2511`:
@@ -40,7 +61,7 @@ python image_edit.py \
 
 Key arguments:
 
-- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
+- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. `jdopensource/JoyAI-Image-Edit-Diffusers` supports one input image.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative-prompt`: text description (string).
 - `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.

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/image_to_video/README.md

@@ -1,7 +1,7 @@
 # Image-To-Video
 
 This shared example generates videos from images with VACE, Wan2.2, LTX-2,
-HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, and other compatible pipelines.
+HunyuanVideo-1.5, SANA-Video, Cosmos3, MAGI-2, Kandinsky 6, and other compatible pipelines.
 
 - `image_to_video.py`: command-line script for single video generation with advanced options.
 
@@ -30,6 +30,7 @@ This folder provides a unified CLI script for image-to-video generation using vL
 | `sand-ai/MAGI-2-preview` | 512 x 896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
 | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE |
 | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE |
+| `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | 480 x 864 | 125 | 50 | 5.0 | ~75 GiB reserved (H100, CPU offload); gated Hub repo |
 
 MAGI-2 native Preview setup, four-GPU topology and DLO choices, I2VA commands,
 and eight-GPU validation status are documented in the

examples/offline_inference/text_to_image/text_to_image.py

@@ -108,7 +108,8 @@ def build_parser() -> argparse.ArgumentParser:
         "--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, "

examples/offline_inference/text_to_video/text_to_video.md

@@ -14,6 +14,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
 | `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
+| `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | 480x864 | 125 | 50 | 5.0 | ~75 GiB reserved (H100, CPU offload, SDPA) |
 | `sand-ai/MAGI-2-preview` | 512x896 | 125 | 100 | Model-fixed | Native four-GPU TP/SP; resident SP4 default; DLO available |
 | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper |
 | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper |
@@ -249,6 +250,24 @@ python text_to_video.py \
 > conditioning tensors that cannot be passed through the JSON `--extra-body`
 > flag; they are out of scope for this text-to-video example.
 
+### Kandinsky 6 (TI2VA)
+
+The Hub checkpoint is gated; export `HF_TOKEN` if download returns 401.
+`--enable-cpu-offload` is required on an 80 GB H100. Audio is on by default.
+
+```bash
+python text_to_video.py \
+  --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \
+  --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \
+  --enable-cpu-offload \
+  --output kandinsky6_t2va.mp4
+```
+
+Video only: add `--extra-body '{"sample_audio": false}'`. Image-to-video uses
+the shared [`image_to_video.py`](../image_to_video/README.md) example with
+`--image`. Serving numbers and flags are in the
+[`Kandinsky 6 TI2VA recipe`](../../../recipes/Kandinsky/Kandinsky6-TI2VA.md).
+
 ## Key Arguments
 
 ### Common

examples/offline_inference/text_to_video/text_to_video.py

@@ -235,7 +235,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, "
+        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, MAGI-2, Kandinsky 6, "
         "and other text-to-video models."
     )
     parser.add_argument(

examples/online_serving/image_to_image/README.md

@@ -2,10 +2,12 @@
 
 This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni.
 
-Supported models include Qwen-Image-Edit, BAGEL, and other image-to-image pipelines.
+Supported models include Qwen-Image-Edit, BAGEL, JoyAI-Image-Edit, and other image-to-image pipelines.
 
 For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content.
 
+For **JoyAI-Image-Edit**, use `jdopensource/JoyAI-Image-Edit-Diffusers` and send exactly one input image.
+
 ## Start Server
 
 ### Basic Start
@@ -29,6 +31,12 @@ vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
 ```
 
+### JoyAI-Image-Edit
+
+```bash
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers --omni --port 8092
+```
+
 ### Start with Parameters
 
 
@@ -77,6 +85,11 @@ EOF
 curl -s http://localhost:8092/v1/chat/completions   -H "Content-Type: application/json"   -d @request.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
 ```
 
+For JoyAI-Image-Edit, set classifier-free guidance with `guidance_scale`
+in `extra_body` (default: `4.0`). `true_cfg_scale` remains supported as an
+alias; if both are provided, their values must match. Values at or below `1.0`
+disable CFG.
+
 ### Method 2: Using OpenAI Python SDK
 
 ```python

examples/online_serving/qwen3_omni/README.md

@@ -277,61 +277,6 @@ cd examples/online_serving/qwen3_omni
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost"
 ```
 
-#### Realtime WebSocket client (`openai_realtime_client.py`)
-
-[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.output_audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging.
-
-Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation.
-
-**Dependencies:**
-
-```bash
-pip install websockets
-```
-
-**From this directory** (`examples/online_serving/qwen3_omni`):
-
-```bash
-python openai_realtime_client.py \
-  --url ws://localhost:8091/v1/realtime \
-  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --input-wav /path/to/input_16k_mono.wav \
-  --output-wav realtime_output.wav \
-  --server-vad \
-  --delta-dump-dir ./rt_delta_wavs
-```
-
-**Arguments:**
-
-| Flag | Default | Description |
-| ------ | --------- | ------------- |
-| `--url` | `ws://localhost:8091/v1/realtime` | Full WebSocket URL including path |
-| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (sent in `session.update`) |
-| `--input-wav` | *(required)* | Input WAV: mono, 16-bit PCM, **16 kHz** |
-| `--output-wav` | `realtime_output.wav` | Output path for concatenated reply audio |
-| `--output-text` | *(optional)* | If set, write final transcription text to this path |
-| `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) |
-| `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) |
-| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.output_audio.delta` WAV files |
-| `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits |
-| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
-| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
-
-Server VAD uses the pinned Silero v6.2 ONNX artifact from `istupakov/silero-vad-onnx` and never downloads model files while processing a session. Make the artifact available in the Hugging Face cache before startup, or configure a local artifact in the deployment YAML:
-
-```yaml
-base_config: /path/to/vllm-omni/vllm_omni/deploy/qwen3_omni_duplex.yaml
-duplex_session:
-  server_vad_model_path: /models/silero_vad.onnx
-```
-
-For the engine-owned plugin, use an overlay based on
-`vllm_omni/deploy/qwen3_omni_duplex.yaml`, with `duplex_session.server_vad_model_path`
-pointing to your local artifact. See the [complete duplex setup](../realtime_web/README.md#qwen3-with-server-vad-automatic-turns).
-The default `qwen3_omni_moe.yaml` remains a turn deployment. Duplex mode serves
-`/v1/realtime?duplex=1` and `/v1/chat/completions`; the legacy STT handler requires
-a turn deployment.
-
 The Python client supports the following command-line arguments:
 
 - `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
@@ -358,6 +303,13 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
 bash run_curl_multimodal_generation.sh use_image
 ```
 
+### Realtime WebSocket
+
+The [shared Realtime UI](../realtime_web/README.md#qwen3-explicit-turn-realtime)
+supports explicit turns through `/v1/realtime?duplex=0`. The
+[Server VAD setup](../realtime_web/README.md#qwen3-with-server-vad-automatic-turns)
+uses the duplex endpoint.
+
 ### FAQ
 
 ## Modality control
@@ -582,26 +534,6 @@ The gradio script supports the following arguments:
 
 Run these commands from the repository root with the vLLM-Omni environment activated.
 
-### Without VAD: manually submit each turn
-
-Start the backend:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-```
-
-In another terminal, check backend readiness and start the UI:
-
-```bash
-curl --fail http://127.0.0.1:8091/health
-python -m examples.online_serving.qwen3_omni.realtime_web \
-    --backend ws://127.0.0.1:8091 --stt --port 7863
-```
-
-Open `http://localhost:7863`, start a session, speak, and press **Send turn**.
-`--stt` is the default. Each turn uses a separate connection without shared model
-history; the conversation shown in the browser is a local log.
-
 ### With VAD: automatically submit after silence
 
 Start the backend with a Server VAD deployment overlay and a compatible Silero
@@ -610,7 +542,7 @@ ONNX artifact, following the
 After its health check succeeds, start the UI:
 
 ```bash
-python -m examples.online_serving.qwen3_omni.realtime_web \
+python -m examples.online_serving.realtime_web --profile qwen3-turn \
     --backend ws://127.0.0.1:8091 --vad --port 7863
 ```
 
@@ -621,9 +553,7 @@ VAD mode keeps uploading audio during replies and supports speech interruption.
 The engine-owned Qwen duplex plugin maintains conversation history and tracks
 played replies through playback acknowledgements. Each committed utterance starts
 a new generation request; the model does not use native streaming-input KV decoding.
-STT mode pauses microphone upload while the reply is generated and played.
 In VAD mode, click **Camera** to upload sampled frames with your spoken question;
 see [camera setup and limits](../realtime_web/README.md#qwen-vad-camera-input).
-STT remains audio-only. To switch back to manual turns, use the default turn
-deployment, restart the UI with `--stt`, refresh the page, and reconnect. Stop the existing UI before reusing port 7863.
-See the [shared UI guide](../realtime_web/README.md) for HTTPS access and testing.
+For explicit turns, use `--manual` with the shared UI as described in the
+[shared UI guide](../realtime_web/README.md#qwen3-explicit-turn-realtime).

examples/online_serving/realtime_web/README.md

@@ -7,14 +7,14 @@ This implements [RFC #7222](https://github.com/vllm-project/vllm-omni/issues/722
 | Profile | Turn control | Camera | Playback acknowledgement |
 | --- | --- | --- | --- |
 | `minicpm-native` | Model-controlled listen/speak, continuous audio input | Frames accompany audio | Yes |
-| `qwen3-turn --stt` (default) | User presses **Send turn** | No | No |
+| `qwen3-turn --manual` | Press **Send turn**; mic pauses while the reply plays | No | No |
 | `qwen3-turn --vad` | Server detects trailing silence; speech interrupts replies | Sampled frames with each spoken turn | Yes |
 | `aura-ptt` | Hold **Hold to talk**; release commits. Hold also stops local playback | Sticky frames with speech; vision-follow every 2 frames only while unlocked (locked from PTT release until text-final / listen) | Yes |
 
 Qwen3 VAD uses an engine-owned duplex plugin: microphone upload continues
 while replies stream, and speech can interrupt generation and playback. The
 checkpoint still generates committed turns; this is not native streaming-input
-KV decoding. STT remains explicit-turn and audio-only.
+KV decoding.
 
 ## MiniCPM: existing command stays valid
 
@@ -62,45 +62,36 @@ just-committed speech turn on Stage0/1 (`max_num_seqs=1`). Audio still playing
 does not hold the next vision turn once text is done.
 MiniCPM / Qwen profiles do not set `pushToTalk`, so the PTT control stays hidden.
 
-## Qwen3: explicit-turn STT adapter
+## Qwen3: explicit-turn Realtime
 
-Start the backend, then the UI in a second terminal:
+The shared UI sends OpenAI Realtime events to `/v1/realtime?duplex=0` and keeps
+conversation history on the same connection:
+
+Start the turn-based backend:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-
-python -m examples.online_serving.qwen3_omni.realtime_web \
-    --backend ws://127.0.0.1:8091 --stt --port 7863
 ```
 
-Open `http://localhost:7863`, start a session, speak, and press **Send turn**.
-Wait for the answer to finish playing before speaking again. The UI opens a
-fresh STT connection for each turn; the displayed conversation is a local log,
-and **previous turns are not provided as model history**. The STT endpoint
-accepts only model selection, so system-prompt controls are hidden in this mode.
+In another terminal, check readiness and start the UI:
 
-The adapter follows the shipped `qwen3_omni/openai_realtime_client.py` and
-`vllm_omni/entrypoints/openai/realtime_connection.py`:
+```bash
+curl --fail http://127.0.0.1:8091/health
+python -m examples.online_serving.realtime_web --profile qwen3-turn \
+    --backend ws://127.0.0.1:8091 --manual --port 7863
+```
 
-
-1. Explicit `duplex=0` selects the legacy STT handler on a turn deployment.
-   It does not enable STT on a duplex deployment.
-2. Send `{type: "session.update", model: ...}` and `commit(final=false)`.
-3. Stream mono PCM16 at 16 kHz; **Send turn** flushes buffered audio, then sends `commit(final=true)`.
-4. `response.output_audio.delta.audio` contains PCM; `response.output_audio.done`
-   terminates this adapter, including a response with no audio.
-5. The local implementation emits **model output** on `transcription.*`, so the
-   UI displays it as assistant text. This differs from the early RFC's suggested
-   user-transcript mapping. The VAD adapter also supports `response.output_text.*`
-   and audio-transcript events without duplicating the answer.
+Speak, then press **Send turn**. The client commits the audio and requests a
+response; recording pauses while the reply plays.
 
 ## Qwen3 with Server VAD (automatic turns)
 
 Run commands from the repository root with the vLLM-Omni environment activated.
-Prepare the pinned Silero v6.2 ONNX artifact as described in the
-[Qwen3 Server VAD instructions](../qwen3_omni/README.md#realtime-websocket-client-openai_realtime_clientpy).
-The backend requires ONNX Runtime and a compatible local artifact; setting
-`--vad` on the UI alone does not enable VAD on the backend.
+The plugin uses the pinned Silero v6.2 ONNX artifact from
+`istupakov/silero-vad-onnx` and does not download it while serving. Make it
+available in the Hugging Face cache or set `server_vad_model_path` to a local
+copy in the overlay below. The backend requires ONNX Runtime; setting `--vad`
+on the UI alone does not enable VAD.
 
 Create an overlay named `qwen3_vad.yaml`, replacing both paths with absolute paths
 on the backend host:
@@ -124,7 +115,7 @@ In a second terminal, wait until the health check succeeds, then start the UI:
 
 ```bash
 curl --fail http://127.0.0.1:8091/health
-python -m examples.online_serving.qwen3_omni.realtime_web \
+python -m examples.online_serving.realtime_web --profile qwen3-turn \
     --backend ws://127.0.0.1:8091 --vad --port 7863
 ```
 
@@ -148,21 +139,17 @@ The bundled duplex deploy allows four audio inputs and eight images per prompt.
 Use headphones to avoid speaker audio triggering VAD interruption. Reference
 voices and tool calls are not supported by this plugin.
 
-### Switching between VAD and manual turns
+### Using the Qwen3 web UI
 
-| Mode | Backend | UI flag | Submit a turn |
-| --- | --- | --- | --- |
-| Without VAD | Default Qwen turn deployment | `--stt` (default) | Press **Send turn** |
-| With VAD | Server VAD configuration and Silero artifact | `--vad` | Pause after speaking |
+The shared UI supports both explicit turns and Server VAD. Use the manual UI
+command above with the turn-based deployment, or start the VAD deployment above
+and launch the UI with `--vad`.
 
 Stop the existing UI process before starting another on the same port, then
-refresh the browser and reconnect. To use STT, start a backend with the default
-turn deployment and restart the UI with `--stt`. The duplex backend exposes
-Realtime duplex and ChatCompletion; it does not expose the legacy STT handler.
-To enable VAD on a default backend, first restart that backend with the overlay.
-Make `--backend` match the actual backend port; the UI and backend use separate
-ports. `Connection refused` means the target backend is unavailable: check its
-startup log and `/health` before connecting.
+refresh the browser and reconnect. To enable VAD on a default backend, restart
+that backend with the overlay. Make `--backend` match the actual backend port;
+the UI and backend use separate ports. `Connection refused` means the target
+backend is unavailable: check its startup log and `/health` before connecting.
 
 VAD connects with `duplex=1` and sends nested `session.audio.input` configuration
 with `create_response: true` and `interrupt_response: true`. It sends playback
@@ -210,8 +197,7 @@ native camera protocol, where Stage 0 interleaves frames at unit boundaries.
 
 This is a supported subset, not full OpenAI Realtime API compatibility.
 `duplex=1`, `overlap_policy`, `playback.ack`, and `session.close` remain vLLM-Omni
-extensions. The legacy `--stt` mode is unchanged and is not the standard Realtime
-conversation API. This demo does not add WebRTC, semantic VAD, or tool calling.
+extensions. This demo does not add WebRTC, semantic VAD, or tool calling.
 
 ## Shared host options
 
@@ -219,11 +205,14 @@ The shared entry point can select a profile explicitly:
 
 ```bash
 python -m examples.online_serving.realtime_web --profile qwen3-turn \
-    --backend ws://127.0.0.1:8091 --stt --port 7863
+    --backend ws://127.0.0.1:8091 --vad --port 7863
 ```
 
+This Qwen3 UI profile requires the server-VAD deployment above.
+
 - `--backend` / `--ws-backend`: backend WebSocket origin (not the full `/v1/realtime` URL).
 - `--model`: override the wrapper's model name.
+- `--manual` / `--vad`: Qwen turn control; manual turns are the default.
 - `--public-realtime-url`: optional browser-visible WebSocket URL; otherwise the
   static host proxies `/v1/realtime` on the same origin.
 - `--host`, `--port`: UI bind address and port (default port 7862).
@@ -245,7 +234,7 @@ pytest -m cpu tests/model_executor/models/qwen3_omni/test_duplex_plugin.py \
        tests/engine/test_duplex_orchestrator.py
 ```
 
-These CPU tests cover profile messages, two STT turns, VAD terminal/drain ordering,
+These CPU tests cover profile message serialization, VAD terminal/drain ordering,
 backpressure, native capture during playback, server config injection, and shared
 worklets. Backend regressions cover current audio inclusion in history-based
 prompts, delayed playback ACK ordering, and no extra generation during silence

examples/online_serving/realtime_web/app/static/profiles/qwen3_turn.js

@@ -2,24 +2,24 @@
   'use strict';
   const profiles = root.OmniRealtimeProfiles;
   profiles['qwen3-turn'] = (config) => {
-    const vad = config.adapter === 'vad';
+    const vad = config.turnMode === 'vad';
     let interruptOnSpeech = vad;
     let imageSeq = 0;
     let imageItems = [];
     // The engine refuses a ninth image (or 4 MiB of them) for the whole
     // session, so the camera track retires its own oldest frame instead.
     const MAX_IMAGE_ITEMS = 8;
-    const hint = 'This Qwen deploy needs session_mode: duplex and Server VAD support; use --stt or enable Server VAD.';
+    const hint = 'This Qwen deployment needs session_mode: duplex with Server VAD enabled.';
     return {
-      title: 'Qwen3-Omni Voice', eyebrow: vad ? 'Engine-owned duplex call' : 'Turn-based realtime call',
+      title: 'Qwen3-Omni Voice', eyebrow: vad ? 'Engine-owned duplex call' : 'Turn-based Realtime call',
       policy: vad ? 'Server VAD • interruptible responses' : 'Send turn • one response at a time',
       description: vad
         ? 'Pause to send your turn; speak again to interrupt. Microphone upload continues during playback. Requires the Qwen duplex plugin and Silero. Optional camera frames accompany each spoken turn.'
-        : 'Speak, then press Send turn. Microphone upload pauses while the model answers. Each turn uses a fresh connection; conversation history is not carried between turns. No barge-in or camera input.',
+        : 'Speak, then press Send turn. Microphone upload pauses while the model answers. Conversation history stays on the Realtime session. No barge-in or camera input.',
       waiting: 'Waiting for you', camera: vad, cameraMaxDimension: 448, playbackAck: vad, clientCommit: !vad,
-      waitForResponseDone: true, halfDuplex: !vad, deduplicateTranscript: true, closeSession: vad, reconnectEachTurn: !vad,
-      inputSampleRate: vad ? 24000 : 16000,
-      readyEvent: vad ? 'session.updated' : 'session.created', instructions: vad, sendIntervalMs: 200,
+      waitForResponseDone: true, halfDuplex: !vad, deduplicateTranscript: true, closeSession: vad,
+      inputSampleRate: 24000,
+      readyEvent: 'session.updated', instructions: true, sendIntervalMs: 200,
       presets: { assistant: 'You are a helpful assistant. Answer clearly and concisely.' },
       url(config, location) {
         const url = profiles.url(config, location);
@@ -33,22 +33,17 @@
       initialMessages(config, instructions) {
         imageItems = [];
         interruptOnSpeech = vad;
-        if (!vad) return [
-          { type: 'session.update', model: config.model },
-          { type: 'input_audio_buffer.commit', final: false },
-        ];
-        const session = { model: config.model, overlap_policy: 'barge_in_on_speech', audio: { input: {
+        const session = { model: config.model, audio: { input: {
           format: { type: 'audio/pcm', rate: 24000 },
-          turn_detection: { type: 'server_vad', threshold: 0.5, prefix_padding_ms: 300,
-            silence_duration_ms: 500, create_response: true, interrupt_response: true },
+          turn_detection: vad ? { type: 'server_vad', threshold: 0.5, prefix_padding_ms: 300,
+            silence_duration_ms: 500, create_response: true, interrupt_response: true } : null,
         } } };
+        if (vad) session.overlap_policy = 'barge_in_on_speech';
+        else session.type = 'realtime';
         if (instructions) session.instructions = instructions;
         return [{ type: 'session.update', session }];
       },
-      append(audio) {
-        return vad ? { type: 'input_audio_buffer.append', audio }
-          : { type: 'input_audio_buffer.append', audio, format: 'pcm16', sample_rate_hz: 16000 };
-      },
+      append: (audio) => ({ type: 'input_audio_buffer.append', audio }),
       // Qwen is a turn model: a frame is a picture added to the conversation,
       // not a track interleaved into the audio. That is exactly the OpenAI
       // Realtime image interface, so the camera sends conversation items
@@ -68,7 +63,10 @@
         } });
         return messages;
       },
-      commitMessages: () => vad ? [] : [{ type: 'input_audio_buffer.commit', final: true }],
+      commitMessages: () => vad ? [] : [
+        { type: 'input_audio_buffer.commit' },
+        { type: 'response.create' },
+      ],
       ack: (responseId, playedMs) => vad ? { type: 'playback.ack', response_id: responseId,
         item_id: `item_${responseId}`, played_ms: playedMs, committed_ms: playedMs } : null,
       mapEvent(event) {
@@ -87,30 +85,28 @@
           return { kind: 'interrupt' };
         }
         if (event.type === 'playback.acknowledged') return { kind: 'ack', committedMs: (event.event || event).committed_ms || 0 };
-        if (event.type === 'response.output_text.delta' || event.type === 'transcription.delta') {
-          // On the shipped STT route transcription.* is model-generated text,
-          // not a separate ASR transcript of the user. See realtime_connection.py.
+        if (event.type === 'response.output_text.delta') {
           return { kind: 'text', role: 'assistant', channel: 'text', text: event.delta || '' };
         }
-        if (event.type === 'response.output_text.done' || event.type === 'transcription.done') {
+        if (event.type === 'response.output_text.done') {
           return { kind: 'text-final', role: 'assistant', channel: 'text', text: event.text || '' };
         }
         if (vad && (event.type === 'output_audio_buffer.cleared' ||
             (event.type === 'response.done' && event.response?.status === 'cancelled'))) {
           return { kind: 'interrupt', responseId: event.response_id || event.response?.id || null };
         }
-        if (event.type === 'input_audio_buffer.committed') return { kind: 'begin' };
+        if (vad && event.type === 'input_audio_buffer.committed') return { kind: 'begin' };
+        if (!vad && event.type === 'input_audio_buffer.cleared') return { kind: 'ignore' };
         if (event.type === 'input_audio_buffer.cleared') return { kind: 'backpressure', message: 'Input cleared. Model is answering; please wait.' };
-        if (!vad && event.type === 'response.output_audio.done') return { kind: 'done' };
         const action = profiles.event(event);
         if (action.kind === 'error') {
           if (action.code === 'input_backpressure') return { ...action, kind: 'backpressure', message: 'Input buffer is full; please repeat dropped speech after the response.' };
-          if (action.code === 'server_vad_initialization_failed') action.message = 'Server VAD could not load Silero. Configure the server artifact or restart this UI with --stt.';
+          if (action.code === 'server_vad_initialization_failed') action.message = 'Server VAD could not load Silero. Configure the server artifact and restart the UI.';
           else if (vad && action.code === 'unsupported') action.message += ` ${hint}`;
         }
         return action;
       },
-      connectionHint: vad ? hint : 'Check the Qwen /v1/realtime STT endpoint and model name.',
+      connectionHint: vad ? hint : 'Check the Qwen turn-based /v1/realtime deployment and model name.',
     };
   };
 })(globalThis);

examples/online_serving/realtime_web/server.py

@@ -74,12 +74,12 @@ def build_app(
     public_realtime_url: str | None = None,
     ref_audio: str | None = None,
     profile: str = "minicpm-native",
-    adapter: str = "stt",
+    turn_mode: str = "manual",
 ) -> FastAPI:
     if profile not in {"minicpm-native", "qwen3-turn", "aura-ptt"}:
         raise ValueError(f"Unknown profile: {profile}")
-    if adapter not in {"stt", "vad"}:
-        raise ValueError(f"Unknown adapter: {adapter}")
+    if turn_mode not in {"manual", "vad"}:
+        raise ValueError(f"Unknown turn mode: {turn_mode}")
     if profile == "qwen3-turn" and ref_audio:
         raise ValueError("Qwen3 profile does not accept --ref-audio")
     if profile == "aura-ptt" and ref_audio:
@@ -116,7 +116,7 @@ def build_app(
             {
                 "model": model,
                 "profile": profile,
-                "adapter": adapter,
+                "turnMode": turn_mode,
                 "realtimePath": public_realtime_url or "v1/realtime",
                 "refAudio": ref_audio_uri,
                 "appVersion": app_version,
@@ -190,9 +190,9 @@ def main(default_profile: str = "minicpm-native") -> None:
     parser.add_argument("--profile", choices=("minicpm-native", "qwen3-turn", "aura-ptt"), default=default_profile)
     parser.add_argument("--ws-backend", "--backend", dest="ws_backend")
     mode = parser.add_mutually_exclusive_group()
-    mode.add_argument("--stt", dest="adapter", action="store_const", const="stt")
-    mode.add_argument("--vad", dest="adapter", action="store_const", const="vad")
-    parser.set_defaults(adapter="stt")
+    mode.add_argument("--manual", dest="turn_mode", action="store_const", const="manual")
+    mode.add_argument("--vad", dest="turn_mode", action="store_const", const="vad")
+    parser.set_defaults(turn_mode="manual")
     parser.add_argument(
         "--public-realtime-url",
         help="Browser-visible ws:// or wss:// Realtime URL; defaults to the same-origin proxy.",
@@ -225,7 +225,7 @@ def main(default_profile: str = "minicpm-native") -> None:
             public_realtime_url=args.public_realtime_url,
             ref_audio=args.ref_audio,
             profile=args.profile,
-            adapter=args.adapter,
+            turn_mode=args.turn_mode,
         ),
         host=args.host,
         port=args.port,

pyproject.toml

@@ -232,6 +232,11 @@ ignore = [
 # F-correctness rules still apply.
 "vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
 "vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"]  # System/summarizer prompts are long prose lines
+# Kandinsky 6 keeps native/BigVGAN names (T/H/W, `l`) and NVIDIA vocoder
+# docstring length so the port stays diffable against the source modules.
+"vllm_omni/diffusion/models/kandinsky6/kandinsky6_transformer.py" = ["N803", "E741"]
+"vllm_omni/diffusion/models/kandinsky6/pipeline_kandinsky6.py" = ["N803"]
+"vllm_omni/diffusion/models/kandinsky6/modeling_kandinsky6_audio.py" = ["E501", "E741", "N803", "UP008"]
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
@@ -277,6 +282,7 @@ markers = [
     "parallel: Parallelism/distributed tests",
     "sp: Sequence parallelism tests",
     "example: Doc example code tests",
+    "nvidia_only: Tests that require NVIDIA CUDA and must not run on ROCm",
     # platform markers
     # Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms().
     # Add [gpu] on accelerator platforms that should also get pytest.mark.gpu.
@@ -364,6 +370,8 @@ extend-ignore-identifiers-re = [
     # (output feature dim), not "doubt".
     ".*dout.*",
     ".*DOUT.*",
+    # Kandinsky nabla threshold argument (`thr`), not the word "the".
+    "^thr$",
 ]
 extend-ignore-re = [
     "\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -404,3 +412,8 @@ IY = "IY"
 # Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a
 # typo of "convent".
 convnet = "convnet"
+# Upstream Kandinsky system prompt. "promt" and "scren" are part of the
+# conditioning text and must not be corrected.
+thr = "thr"
+promt = "promt"
+scren = "scren"

recipes/JD/JoyAI-Image-Edit.md

@@ -0,0 +1,221 @@
+# JoyAI-Image-Edit
+
+> Text-guided single-image editing with the native JoyAI-Image-Edit diffusion pipeline
+
+## Summary
+
+- Vendor: JD / JoyImage Team
+- Model: `jdopensource/JoyAI-Image-Edit-Diffusers`
+- Task: Text-guided single-image editing
+- Mode: Offline inference and OpenAI-compatible online serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to run JoyAI-Image-Edit through the shared
+image-to-image examples with a known-good baseline configuration. JoyAI-Image-Edit
+is a large single-image editing model: it requires one reference image per
+request, snaps requested sizes to Joy's supported resolution buckets, and has a
+validated 1024x1024 baseline on an 80 GB NVIDIA GPU.
+
+## References
+
+- Upstream model:
+  [`jdopensource/JoyAI-Image-Edit-Diffusers`](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Diffusers)
+- Upstream project:
+  [`jd-opensource/JoyAI-Image`](https://github.com/jd-opensource/JoyAI-Image)
+- Related offline example:
+  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
+- Related online example:
+  [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md)
+- User guide:
+  [`docs/user_guide/examples/offline_inference/image_to_image.md`](../../docs/user_guide/examples/offline_inference/image_to_image.md),
+  [`docs/user_guide/examples/online_serving/image_to_image.md`](../../docs/user_guide/examples/online_serving/image_to_image.md)
+
+## Hardware Support
+
+This recipe documents the CUDA configuration used to validate the PR's JoyAI
+integration. Other GPU types should be treated as unvalidated until you rerun
+the smoke and parity checks below.
+
+## GPU
+
+### 1x NVIDIA H800 PCIe 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: `3.12.3`
+- Driver / runtime: NVIDIA driver `580.82.07`, CUDA runtime `13.0`
+- PyTorch: `2.11.0+cu130`
+- Diffusers: `0.39.0.dev0` with `JoyImageEditPipeline` available
+- vLLM: `0.21.0`
+- vLLM-Omni: use the commit you are deploying from
+- GPU: NVIDIA H800 PCIe, 81559 MiB
+
+#### Online Serving
+
+Start the OpenAI-compatible server from the repository root:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200
+```
+
+The server is ready when the log shows `Application startup complete.`
+
+If you need extra GPU headroom, use one of the memory-saving variants below.
+These variants are useful on tighter 80 GB deployments, but the latency and
+quality numbers in this recipe are from the baseline command above.
+
+```bash
+# Pipeline-level CPU offload for large components.
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200 \
+  --enable-cpu-offload
+```
+
+```bash
+# DiT blockwise offload. This trades latency for lower peak GPU memory.
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200 \
+  --enable-layerwise-offload
+```
+
+#### Online Verification
+
+Download the public JoyAI plate example:
+
+```bash
+wget -O /tmp/joy_plate.jpg \
+  https://raw.githubusercontent.com/jd-opensource/JoyAI-Image/main/test_images/test_1.jpg
+```
+
+Send one image and one prompt. JoyAI-Image-Edit supports exactly one input image
+per request.
+
+```bash
+IMG_B64=$(base64 -w0 /tmp/joy_plate.jpg)
+
+curl -s http://localhost:8092/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d "$(jq -n --arg img "$IMG_B64" '{
+    messages: [{
+      role: "user",
+      content: [
+        {type: "text", text: "Turn the plate blue"},
+        {type: "image_url", image_url: {url: ("data:image/jpeg;base64," + $img)}}
+      ]
+    }],
+    extra_body: {
+      height: 1024,
+      width: 1024,
+      num_inference_steps: 30,
+      guidance_scale: 4.0,
+      seed: 0
+    }
+  }')" \
+  | jq -r '.choices[0].message.content[0].image_url.url' \
+  | cut -d',' -f2- \
+  | base64 -d > /tmp/joyai_image_edit_plate_online.png
+```
+
+Expected result: `/tmp/joyai_image_edit_plate_online.png` is a 1024x1024 PNG
+where the plate color changes to blue while the surrounding scene is mostly
+preserved. Use a client timeout of at least 600 seconds for first-run or cold
+server checks.
+
+#### Offline Command
+
+Run the same smoke case without starting a server:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image /tmp/joy_plate.jpg \
+  --prompt "Turn the plate blue" \
+  --negative-prompt "" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 30 \
+  --guidance-scale 4.0 \
+  --seed 0 \
+  --enforce-eager \
+  --init-timeout 1200 \
+  --output /tmp/joyai_image_edit_plate_offline.png
+```
+
+For a local checkpoint mirror, replace `--model` with the local model directory.
+The PR validation used the same command shape with
+`/autodl-fs/data/joyai_models/JoyAI-Image-Edit-Diffusers`.
+
+#### Validated Results
+
+The table below is from the 1x H800 PCIe environment above. Resolution requests
+were snapped to the Joy 1024x1024 bucket, batch size was 1, and CFG was 4.0.
+`vLLM peak` is the observed peak GPU memory for the vLLM-Omni run.
+
+| Case | Steps | Seed | vLLM latency | Diffusers latency | vLLM peak | Diffusers peak | SSIM | PSNR |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| `plate_seed123` | 30 | 123 | 50.87 s | 52.74 s | 57.52 GiB | 57.28 GiB | 0.9869 | 34.7051 |
+| `plate_seed0` | 30 | 0 | 51.02 s | 52.01 s | 57.52 GiB | 57.28 GiB | 0.9946 | 41.4351 |
+| `astronaut_seed0` | 40 | 0 | 67.90 s | 69.22 s | 57.52 GiB | 57.28 GiB | 0.9948 | 38.2498 |
+| `crane_removal_seed0` | 40 | 0 | 68.00 s | 69.27 s | 57.52 GiB | 57.28 GiB | 0.9975 | 50.0135 |
+| `object_move_seed0` | 40 | 0 | 68.06 s | 69.32 s | 57.52 GiB | 57.28 GiB | 0.9969 | 46.1138 |
+| `object_rotation_seed0` | 40 | 0 | 67.94 s | 69.28 s | 57.52 GiB | 57.28 GiB | 0.9898 | 35.5365 |
+| `camera_control_seed0` | 40 | 0 | 68.15 s | 69.51 s | 57.52 GiB | 57.28 GiB | 0.9952 | 41.7697 |
+
+#### Output Size Behavior
+
+JoyAI-Image-Edit does not use arbitrary image sizes directly. The pipeline maps
+the input or requested `height` / `width` to the nearest supported Joy bucket,
+then resizes and center-crops the reference image to that bucket before
+denoising.
+
+- Provide both `height` and `width`, or omit both.
+- For square outputs, request `height=1024` and `width=1024`.
+- A request such as `512x512` is also square and is snapped to the 1024x1024 Joy
+  bucket, so the final image size will not match the literal request.
+- Non-square requests are snapped to the nearest aspect-ratio bucket supported
+  by the model.
+
+#### Guidance Parameters
+
+Use `guidance_scale` for JoyAI classifier-free guidance, matching Diffusers.
+The pipeline defaults to `4.0` when the parameter is omitted. `true_cfg_scale`
+remains supported as an alias; if both parameters are supplied, their values
+must match. Values greater than `1.0` enable CFG; an omitted negative prompt is
+treated as an empty string. Values at or below `1.0` disable CFG.
+
+The offline example passes `--guidance-scale 4.0` explicitly for clarity; if it
+is omitted, the JoyAI pipeline uses the same `4.0` default.
+
+#### Known Limitations
+
+- Single image only: JoyAI-Image-Edit rejects requests with zero or multiple
+  input images.
+- Single prompt only: batched prompt requests are rejected.
+- HSDP is not supported yet. Starting with HSDP enabled raises a `ValueError`
+  before the generic HSDP setup reaches the transformer.
+- Tensor Parallelism is not supported yet. Keep `tensor_parallel_size=1`;
+  larger values raise a `ValueError` during pipeline initialization.
+- LoRA is not included for this JoyAI pipeline.
+- Cache-DiT and TeaCache are not included for this JoyAI pipeline. Do not treat
+  the generic image-to-image cache flags as validated JoyAI acceleration knobs.
+- Text encoder quantization is not included.
+- The validation above is CUDA-only; ROCm and NPU configurations are not
+  documented as supported by this recipe.

recipes/Kandinsky/Kandinsky6-TI2VA.md

@@ -0,0 +1,247 @@
+# Kandinsky 6 TI2VA — H100 (v1, single GPU, CPU offload)
+
+> Text/image-to-video-and-audio serving with Kandinsky 6 Pro
+
+## Summary
+
+- Vendor: Kandinsky Lab ([`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers))
+- Model: Kandinsky 6 Pro (TI2VA — text/image-to-video-and-audio)
+- Task: Joint text/image-to-video-and-audio generation
+- Mode: Offline (`Omni(...)`) and online serving with the OpenAI-compatible
+  `/v1/videos` API
+- Hardware: 1x NVIDIA H100 80GB HBM3 with component-level CPU offload
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve Kandinsky 6's TI2VA pipeline — a single DiT that
+jointly denoises video and audio latents, conditioned on a text prompt and,
+optionally, a reference image (image-to-video). Tensor parallel, CFG
+parallel, Ulysses, ring, pipeline parallel, HSDP, layerwise offload,
+MagCache, TeaCache, Cache-DiT, FP8, and tiled VAE patch parallel are wired.
+NaviCache and a registered NABLA backend are not. TeaCache coefficients are
+uncalibrated. See "Known limitations".
+
+## Supported model contract
+
+| Task | Entrypoint | Input | Output |
+|---|---|---|---|
+| Text-to-video-and-audio | `vllm serve kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --omni` / `Omni(model="kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers")` | text prompt | synchronized MP4 (H.264 video + AAC 44.1 kHz mono audio) |
+| Image-to-video-and-audio | same | text prompt + reference image | synchronized MP4, first frame conditioned on the reference image |
+
+- Audio generation is on by default (`sample_audio=True`); pass
+  `extra_args.sample_audio=false` to generate video only.
+- Batch size is 1 request per generation call (upstream limitation of the
+  ported denoise loop; concurrent requests are still served by vLLM-Omni's
+  own request scheduler, just not batched together within one DiT forward
+  pass yet).
+- Serving defaults follow the Pro production geometry registered in
+  `vllm_omni/model_extras/kandinsky6.py`: **864x480, 125 frames @ 24 fps
+  (5.2 s), 50 steps, CFG 5.0**. Any `4k+1` frame count and 16-divisible
+  resolution can be requested explicitly via `height`/`width`/`num_frames`/
+  `num_inference_steps` — the model is *not* fixed-duration.
+- Audio is always 44.1 kHz; its length is derived from the video duration.
+- Negative prompt defaults to the pipeline's built-in K6 negative prompt
+  when omitted.
+
+## References
+
+- Hub checkpoint: [`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers)
+- Joint video+audio native-model reference: [`recipes/MiniMaxAI/MiniMax-H3.md`](../MiniMaxAI/MiniMax-H3.md)
+- Offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
+- Image-to-video example: [`examples/offline_inference/image_to_video/image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py)
+
+## Hardware
+
+- Accelerator model and per-device memory: NVIDIA H100 80GB HBM3
+  (driver 570.133.20)
+- Number of devices: 1 (multi-GPU tensor parallelism is implemented in the
+  transformer's linear layers but not yet validated end-to-end)
+- Device interconnect: N/A (single-device path)
+- Host memory: 1.4 TB installed; the bf16 DiT (~56 GB) plus the text
+  encoders are staged on the host under `--enable-cpu-offload`, so budget
+  at least ~90 GB of free host RAM.
+- Qualification scope: single-GPU T2VA at the Pro default geometry and at a
+  reduced smoke geometry, offline and via `vllm serve`. Numbers below were
+  measured once on the shared box described above; treat them as
+  indicative, not as a benchmark.
+
+## Software environment
+
+- OS: Ubuntu 22.04.5 LTS
+- Python: 3.12.13
+- Driver / runtime: CUDA 12.9 (torch 2.13.0+cu129)
+- vLLM version: 0.29.0
+- vLLM-Omni version or commit: this repository, `vllm_omni/diffusion/models/kandinsky6/`
+- diffusers 0.40.0, transformers 5.14.1
+- Attention: FlashAttention-3 on this box (`flash_attn_interface`). The Pro
+  timings below use it. Without FA3, `attention_engine: "auto"` falls back
+  to PyTorch SDPA; an earlier SDPA Pro run on the same GPU was 26.9 s/it
+  and 1377 s.
+
+## Weights
+
+`Kandinsky6TI2VAPipeline._load_components` loads the public Diffusers layout
+from [`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers)
+(`model_index.json`, `transformer/`, `vae/`, `text_encoder/`, `tokenizer/`,
+`text_encoder_2/`, `tokenizer_2/`, `audio_vae/`, `scheduler/`). Pass that
+repo id (or a local snapshot of it) to `vllm serve` / `Omni(model=...)`.
+The Hub repo is gated; export `HF_TOKEN` if download returns 401.
+
+A raw Kandinsky 6 weight tree is not required. If you already have a local
+Diffusers-layout snapshot, point `--model` at that directory instead of the
+Hub id.
+
+Each weight folder is loaded on its own (`transformer/`, `vae/`, `text_encoder/`, `text_encoder_2/`, `audio_vae/`). The DiT is built with initialization skipped, then the framework loader fills it from `transformer/`.
+
+## Command
+
+```bash
+vllm serve kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --omni \
+  --host 127.0.0.1 --port 8091 \
+  --num-gpus 1 --enable-cpu-offload
+```
+
+`--enable-cpu-offload` enables model-level offload with mutual exclusion
+between the transformer and the two text encoders (they never co-reside on
+the GPU). Without it the bf16 DiT (~56 GB) plus Qwen2.5-VL-7B (~16 GB) plus
+the VAEs and activations do not fit an 80 GB device at the Pro geometry.
+
+Offline (writes an MP4 with the audio track muxed in):
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \
+  --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \
+  --seed 42 --enable-cpu-offload --output kandinsky6_t2va.mp4
+# Video only: add --extra-body '{"sample_audio": false}'
+```
+
+Image-to-video-and-audio uses the shared image-to-video example:
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers \
+  --image first_frame.png \
+  --prompt "A golden retriever runs along a sunny beach, waves crashing, cinematic footage" \
+  --seed 42 --enable-cpu-offload --output kandinsky6_i2va.mp4
+```
+
+The generic scripts pick up the Pro defaults from
+`vllm_omni/model_extras/kandinsky6.py`.
+
+## Verification
+
+`/v1/videos` takes multipart form fields (not a JSON body). Submit, poll,
+then download:
+
+```bash
+VID=$(curl -s -X POST http://127.0.0.1:8091/v1/videos \
+  -F prompt="A golden retriever runs along a sunny beach, waves crashing" \
+  -F size=864x480 -F num_frames=125 -F num_inference_steps=50 -F seed=42 \
+  | python3 -c "import sys,json;print(json.load(sys.stdin)['id'])")
+
+# Poll until "status": "completed"
+curl -s http://127.0.0.1:8091/v1/videos/$VID
+
+curl -s -o kandinsky6.mp4 http://127.0.0.1:8091/v1/videos/$VID/content
+# Expect an MP4 with an H.264 video stream (125 frames, 24 fps, 864x480)
+# and an AAC audio stream (44100 Hz, mono) of matching duration.
+```
+
+Smoke-size request for quick validation (~30 s end-to-end on one H100):
+`-F size=512x320 -F num_frames=25 -F num_inference_steps=10`.
+
+## Measurements
+
+Single H100 80GB, seed 42, prompt
+"A golden retriever runs along a sunny beach, waves crashing, cinematic footage",
+negative prompt left at the pipeline default, MagCache off, prompt expansion
+off. Pro geometry: 864x480, 125 frames, 50 steps, CFG 5.0, audio on.
+
+vLLM-Omni: `text_to_video.py --model kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers --enable-cpu-offload --model-class-name Kandinsky6TI2VAPipeline`,
+FlashAttention-3.
+
+| Wall / generation | Step time | Peak GPU | After load |
+|---|---|---|---|
+| 751.7 s generation (12.5 min); engine init 41 s | 14.4 s/it average | 78.39 GiB reserved (78.0 GiB nvidia-smi) | 17.74 GiB |
+
+An earlier SDPA vLLM-Omni Pro run on the same GPU was 26.9 s/it and 1377 s,
+peak reserved 74.97 GiB. A 512x320 / 25-frame / 10-step SDPA smoke was 1.4 s/it
+and 29.5 s at the same reserved peak.
+
+The MP4 is H.264 864x480, 125 frames at 24 fps (5.22 s), plus AAC 44.1 kHz
+mono (230400 samples): a golden retriever on a sunny beach with breaking waves.
+
+vLLM-Omni's VAE decode logged two failed CUDA allocations (7.15 GiB, then
+4.27 GiB) and still wrote a complete file.
+
+Audio is peak-normalized to full scale by default; pass
+`extra_args.audio_normalization="clip"` to keep the raw decoded amplitude
+instead.
+
+## Notes
+
+- Key flags: `extra_args.sample_audio` (bool, default true),
+  `extra_args.audio_normalization` (`"normalize"` default | `"clip"`),
+  `extra_args.visual_cond_scheme` (defaults to `pretrain` for text-only,
+  `tail_cond_first_frame` when an image is supplied).
+- The post-process payload is a flat dict (`video`, `audio` as float32 in
+  `[-1, 1]`, `audio_sample_rate=44100`, `fps`) so the serving muxer picks
+  up the correct sample rate.
+- Acceleration that is wired:
+  - Tensor parallel: the DiT loader narrows full checkpoint tensors onto
+    each rank's `ColumnParallelLinear` / `RowParallelLinear` shard.
+  - CFG parallel (`--cfg-parallel-size 2`): rank 0 runs the conditional
+    DiT, rank 1 the unconditional one, then both all-gather video and
+    audio velocities.
+  - HSDP: `_hsdp_shard_conditions` covers the indexed transformer blocks.
+    Do not combine with tensor parallel.
+  - Sequence parallel: visual tokens (and their RoPE) are sharded;
+    Ulysses all-to-all runs inside visual self-attention, Ring uses the
+    PyTorch ring kernel. Text stays replicated. Audio queries gather the
+    full visual sequence.
+  - Pipeline parallel: `visual_transformer_blocks` are split with
+    `PPMissingLayer`. Embeddings stay on the first stage, output heads on
+    the last. Activations and the final velocity are sent between stages.
+  - VAE patch parallel: Hunyuan decode splits the latent width when
+    `vae_patch_parallel_size > 1` and tiling is on.
+  - MagCache / TeaCache / Cache-DiT are registered. MagCache reuses the
+    Pro `mag_ratios`. TeaCache coefficients are **uncalibrated** (copied
+    from Qwen-Image) so skips are not quality-neutral. Cache-DiT targets
+    `visual_transformer_blocks` only and falls back to the same step cache
+    if the fused `(video, audio)` block return is rejected.
+  - FP8: online quantization runs in the DiT loader via
+    `process_weights_after_loading` after the sharded assign.
+  - `attention_engine: auto` uses FlashAttention-3 when
+    `flash_attn_interface` imports.
+- Known limitations:
+  - NABLA sparse attention is not registered in the framework attention
+    backend. NaviCache is not wired.
+  - Prompt expansion, NF4 Qwen, and the super-resolution cascade are not
+    part of this port.
+  - Distributed layerwise offload still uses one process unless another
+    parallel axis sets `world_size > 1` (for example tensor parallel).
+  - `_load_components` loads the Diffusers checkpoint directly rather than the
+    streamed-prefetch loader.
+
+## Supported features
+
+| Feature | Status | Guide |
+|---|---|---|
+| Text-to-video-and-audio | Supported (validated on H100) | — |
+| Image-to-video-and-audio | Supported (not yet validated on hardware) | — |
+| CPU offload (`--enable-cpu-offload`) | Supported (required on 80 GB) | — |
+| Layerwise offload (`--enable-layerwise-offload`) | Supported (about 36 GiB peak on the smoke geometry) | — |
+| Tensor parallelism | Wired (sharded DiT load) | — |
+| CFG parallelism | Wired (`--cfg-parallel-size 2`) | — |
+| Ulysses / Ring sequence parallel | Wired (visual tokens only) | — |
+| Pipeline parallelism | Wired (visual block split) | — |
+| HSDP | Wired (not with tensor parallel) | — |
+| VAE patch parallel | Wired (width-split decode) | — |
+| MagCache | Wired (Pro ratios, step reuse) | — |
+| TeaCache | Wired, coefficients uncalibrated | — |
+| Cache-DiT | Wired on `visual_transformer_blocks` | — |
+| FP8 | Wired (online quant; smoke peak about 46 GiB vs 75 GiB bf16) | [`FP8 quantization`](../../docs/user_guide/quantization/fp8.md) |
+| NaviCache | Not supported | — |
+| LoRA | Not supported | — |

recipes/NVIDIA/PersonaPlex.md

@@ -84,8 +84,10 @@ duplex (`is_enabled()` is unconditionally true). Turn-based models
 
 ## Hardware Support
 
-Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the
-native modules, so other CUDA GPUs with enough memory are expected to work.
+Verified on GPU (Hopper-class). The shipped deploy config runs the Stage 0
+talker with CUDA graphs and replays the streaming Mimi encode (Stage 0) and
+decode (Stage 1) frame steps from model-local CUDA graphs, bitwise equal to
+eager. Other CUDA GPUs with enough memory are expected to work.
 
 ## GPU
 
@@ -137,9 +139,11 @@ python tests/e2e/online_serving/personaplex_realtime_duplex.py \
 
 #### Notes
 
-- Realtime budget is 80 ms/frame. Verified eager per-tick latency at four
-  concurrent sessions is ~70-74 ms on this hardware class, i.e. all four
-  conversations stay realtime; single-session has comfortable headroom.
+- Realtime budget is 80 ms/frame. The shipped config admits two sessions; on a
+  141 GB card, 64 concurrent 30 s sessions stayed realtime after raising
+  `duplex_session.max_sessions` and both stages' `max_num_seqs` (see the
+  sizing notes in `vllm_omni/deploy/personaplex.yaml`). Longer sessions cost
+  more per step as their context fills.
 - Decoding is greedy only (temperature/top-k knobs are intentionally not
   exposed): greedy is what the parity gates pin against the reference
   implementation, and sampling amplifies sub-bit numeric drift into divergence.

recipes/OpenMOSS/MOSS-TTS.md

@@ -35,9 +35,9 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
 audio-reference generation. Reference transcripts must match the reference
 audio; they are not style instructions.
 
-For the optional CUDA MRV2 runner, see the
+For the CUDA MRV2 runner and platform fallbacks, see the
 [Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
-The default Local deployment continues to use V1.
+Local 1.5 defaults to MRV2 on CUDA and retains V1 on other platforms.
 
 ## References
 
@@ -159,8 +159,52 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 ## Local 1.5 MRV2 and slot attention
 
-`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
-an explicit deploy profile. The default Local profile continues to use V1.
+`MOSS-TTS-Local-Transformer-v1.5` defaults to the native MRV2 pipeline on
+CUDA, with the original Local projection and sampler. When
+`nvidia-cuda-mps-control` is on `PATH`, it starts private full-quota MPS.
+GPUs with at least 140 GiB total memory use the C128 system profile with
+prefix caching, Triton backbone attention, bounded codec first-chunk decode,
+native Torch sampler compilation and dense codec graph buckets through 128. Its Talker KV budget is 32 GiB; smaller GPUs or a failed memory
+query use C64 with utilization-based memory budgets. When the MPS executable
+is unavailable, the automatic default is C64 MRV2 without MPS. NPU, XPU,
+ROCm and MUSA retain V1. Explicit deploy configs override automatic selection.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+Use `--deploy-config vllm_omni/deploy/moss_tts_local_v1.yaml` to select the
+previous V1 profile explicitly. C128's batch-prefill and direct-token switches
+remain confined to the throughput profile; they are not enabled in C64.
+
+The same `moss_tts_local_mrv2_optimized.yaml` system profile can be selected
+explicitly. It retains the original Local depth projection and sampling algorithm.
+The codec admits fast first chunks below 32 active streams; crowded streams use
+the regular decoder with a dispatch target of 16 and a maximum wait of 6 ms.
+This preserves the low-load first-audio path while coalescing high-load work.
+`local_compile_audio_sampler: true` in the Talker HF overrides compiles the native
+Torch sampler; explicit request generators use the original helper. Set the
+override to `false` in a deployment file to disable that compilation. The setting
+is specific to this C128 system profile.
+Both stages share one GPU, with a 32 GiB Talker KV budget; this profile requires H200-class memory and
+`nvidia-cuda-mps-control` on `PATH`.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+Evaluate WER, speaker similarity and streaming latency for your workload;
+throughput gains do not establish quality equivalence. Do not apply partial
+MPS SM quotas: BF16 GEMM outputs were observed incomplete in that configuration
+on the validation environment. MPS uses an owned control socket or an
+explicitly supplied operator socket; only the owned daemon is stopped at shutdown.
+
+The explicit C64 MRV2 and capped-C128 throughput profiles remain available.
 Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
 and the model's sampling defaults.
 

recipes/Qwen/Qwen-Image-2.1.md

@@ -0,0 +1,418 @@
+# 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 H200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`,
+warmup 1 + median of 3): compile+graph **5.93 s**, compile-only **6.88 s**,
+graph-only **9.52 s**, eager **8.28 s** end-to-end per image; combo and
+compile-only vs eager are **38.10 dB** PSNR (compile fusion drift), 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, 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.
+Model-level CPU offload does not disable decode graphs while
+`enable_cuda_graph_decode` stays on (the default). DiT weights stay on fixed
+device storage so captured pointers remain valid, and that storage is not
+released while the text encoder runs. Pass `--no-enable-cuda-graph-decode`
+(or `enable_cuda_graph_decode: false`) for the eager offload path, which
+moves DiT weights off the device and frees that VRAM. Peak memory for the
+default graph-plus-offload combination was not re-measured in this PR.
+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 H200, single GPU, same prompt as Quality Comparison
+(`"A ceramic teapot on a wooden table"`, 1024×1024, 50 steps, seed 42,
+eager):
+
+| Config | Quantized block linears | PSNR vs BF16 | Steady-state time/image | Peak memory |
+| --- | --- | --- | --- | --- |
+| BF16 | 0 / 160 | — | **7.34 s** | **36.9 GB** |
+| FP8 all layers | 160 / 160 | **33.0 dB** | **6.06 s** | **30.4 GB** |
+| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | **34.3 dB** | **6.98 s** | **34.9 GB** |
+
+Keeping `img_mlp` in BF16 is still the higher-fidelity DiT setting vs
+all-layer FP8. On H200, all-layer FP8 is both leaner and faster than BF16
+at this shape.
+
+## 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 H200 at 1024×1024
+(seed 42, 50 steps, same teapot prompt, eager): text-encoder FP8 lowers peak
+GPU memory from **36.9 GB** (BF16) to **30.4 GB** (−6.5 GB); T2I PSNR vs BF16
+**32.1 dB**. DiT FP8-all is **30.4 GB** / **6.06 s**. Combined DiT+text-encoder
+FP8 was not re-measured on H200. 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 one NVIDIA H200, T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0,
+eager, PSNR vs. the bf16 baseline on the same GPU): **`"fp8"` 34.2 dB**,
+**`"fp8_v"` 44.0 dB**. The error is dominated by
+K quantization: post-RoPE keys are the precision-sensitive half of the
+cache, so `"fp8_v"` buys back about 9.8 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. These sizes are calculated at the
+limit (8192 text tokens + 4 condition images, ~24.6k prefix tokens), not
+measured on the 1024 T2I run above: ~12.9 GB per CFG branch in bf16, ~6.6 GB
+in `"fp8"` and ~9.7 GB in `"fp8_v"`. That 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, H200) | generated (eager + default CUDA-graph) | generated | measured |
+| PSNR / LPIPS vs reference | eager vs Diffusers **47.17 dB**; **default compile+graph** vs Diffusers **38.02 dB**; **default compile+graph** vs eager **38.10 dB** | — | measured (H200) |
+| Latency (time/image, same GPU, H200) | **default compile+graph** **6.04 s** (120.8 ms/step); eager **7.31 s** (146.2 ms/step) | **7.74 s** (154.9 ms/step) | measured |
+
+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.
+Eager stays numerically close to Diffusers (~47 dB). The **default** serving
+path stacks regional `torch.compile` with CUDA-graph decode; on H200 that
+combo lands at **~38.1 dB** vs eager (~1.21× on the Quality Comparison timed
+run). Graph-only vs eager is bit-identical; the ~38 dB drift is the compile
+fusion, not CUDA-graph replay. See "CUDA Graph decode" for the four-mode
+H200 breakdown (warmup 1 + median of 3).
+
+Parity / feature measurements (same prompt/seed/steps/CFG, BF16, eager decode
+unless noted; 4× NVIDIA H200):
+
+- [x] TP=2/4 output parity vs single-card — measured (1024×1024): TP1 vs TP2
+  **39.32 dB**, TP1 vs TP4 **38.37 dB**, TP2 vs TP4 **45.77 dB**. Latency:
+  TP1 **7.30 s**, TP2 **4.89 s**, TP4 **5.12 s** per image.
+- [x] VAE tiling vs non-tiled decode — measured: 1024×1024 nontile vs
+  `--vae-use-tiling` **43.55 dB** (peak mem ~37.7 GB → ~32.7 GB); 2048×2048
+  nontile vs tiled **44.14 dB** (peak mem ~57.9 GB → ~33.3 GB; gen ~36.8 s vs
+  ~37.1 s).
+- [x] CUDA-graph four-mode vs eager — measured on H200 (1024×1024, 50 steps,
+  seed 42, BF16, warmup 1 + median of 3): compile+graph **5.93 s**,
+  compile-only **6.88 s**, graph-only **9.52 s**, eager **8.28 s**. Combo /
+  compile-only vs eager **38.10 dB**; graph-only vs eager bit-identical.
+  Quality Comparison table uses a separate single-run A/B vs Diffusers
+  (compile+graph **6.04 s**, eager **7.31 s**).
+
+## Known Limitations
+
+- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;
+  `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`.
+- Ulysses sequence parallelism is wired; ring attention is not a supported path.
+  A short H200 smoke (1024×1024, 8 steps, seed 42, true CFG 1.0, eager) finished:
+  Ulysses SP=2 vs SP=1 is **46.5 dB**. That size's target length is divisible by 4,
+  so SP padding was not exercised, and this is not the 50-step parity next to the
+  TP tables. The features matrix therefore stays ❓. Ring SP=2 also ran, at
+  **29.9 dB** vs SP=1, which is why it stays unsupported.
+- 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/Qwen/Qwen3-Omni.md

@@ -73,8 +73,20 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     "1": {"max_num_seqs": 4},
     "2": {"max_num_seqs": 4}
   }'
+
+# Replay the thinker's image encoder from CUDA graphs (images only)
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{"0": {"compilation_config": {"cudagraph_mm_encoder": true}}}'
 ```
 
+With the encoder graphs on, the thinker captures one graph per token budget
+(64 to 256 image tokens and up to 4 images per replay by default). A step's
+images replay together when one captured graph holds them all; otherwise they
+run through the eager encoder in one call. Set `encoder_cudagraph_token_budgets`
+and `encoder_cudagraph_max_vision_items_per_batch` in the same
+`compilation_config` to change them; each captured budget adds graph memory on
+the thinker's GPU.
+
 ### Stage-based launch (one stage per process)
 
 Use three terminals (one per stage). Start with the default commands below, then

recipes/README.md

@@ -64,8 +64,10 @@ recipes/
 | [`MiniMaxAI/MiniMax-H3-Spark-GB10.md`](./MiniMaxAI/MiniMax-H3-Spark-GB10.md) | MiniMax H3 T2VA and Ref2VA on DGX Spark | 1x or 2x DGX Spark (GB10) |
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB / 1x Intel Arc Pro B70 32GB |
+| [`Kandinsky/Kandinsky6-TI2VA.md`](./Kandinsky/Kandinsky6-TI2VA.md) | Text/image-to-video-and-audio, offline and `/v1/videos` | 1x H100 80GB with CPU offload; 2x H100 for parallel modes |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU / 1x Intel Arc Pro B70 32GB |
 | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified / 1x Intel Arc Pro B70 32GB |
+| [`JD/JoyAI-Image-Edit.md`](./JD/JoyAI-Image-Edit.md) | Text-guided single-image editing | 1x H800 PCIe 80GB |
 | [`m-a-p/SheetSage2-H200.md`](./m-a-p/SheetSage2-H200.md) | Official audio-to-score preprocessing and YuE2 request export | 1x H200 141GB |
 | [`m-a-p/YuE2-3B.md`](./m-a-p/YuE2-3B.md) | Offline and online text-to-music with optional ABC score | 1x RTX 4090 24GB |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
@@ -86,6 +88,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; offline text-to-image on XPU | 1x A100 80GB / 2x B200 / 1x Intel Arc Pro B70 32GB |
+| [`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 H200 |
 | [`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; offline end-to-end text-to-speech on XPU | 1x A100 80GB / 5x Intel Arc Pro B70 32GB |

recipes/Wan-AI/Wan2.2-I2V.md

@@ -80,9 +80,13 @@ parallelism at `--usp 4` so the two CFG branches cover all 8 GPUs.
 - `--use-hsdp` is recommended for memory efficiency on the 14B DiT backbone.
 - `--vae-patch-parallel-size 8` is recommended for the standard 8-card setup;
   disabling VAE patch parallelism can significantly increase VAE latency.
-- The CUDA Triton fused AdaLayerNorm path is not available in the current
-  release. The `VLLM_OMNI_ENABLE_TRITON_ADALN` toggle is reserved for that
-  future implementation and should not be used until the kernel lands.
+- CUDA eager execution uses the fused Triton AdaLayerNorm path for supported
+  contiguous 3D inputs with shared `(C,)`, `(1, C)`, or `(1, 1, C)` modulation,
+  or per-sample `(B, 1, C)` modulation, including supported strided chunk views.
+  Unsupported inputs use the native path; per-token `(B, L, C)` modulation
+  used by TI2V remains native.
+- During `torch.compile`, AdaLayerNorm uses the native operation chain, which
+  Inductor can fuse. The eager Triton path remains available after compilation.
 - In an 8x NVIDIA H20 test with the official model, `num_inference_steps=4`,
   `flow_shift=12.0`, `--use-hsdp`, `--cfg-parallel-size 2`, `--usp 4`, and
   `--vae-patch-parallel-size 8`, an internal Triton AdaLayerNorm prototype
@@ -181,20 +185,20 @@ for complete client examples and request formats.
 #### Notes
 
 - **Key flags:**
-  - `--omni` — enables vLLM-Omni diffusion serving.
-  - `--use-hsdp` — enables Hybrid Sharded Data Parallelism for the DiT model
+    - `--omni` — enables vLLM-Omni diffusion serving.
+    - `--use-hsdp` — enables Hybrid Sharded Data Parallelism for the DiT model
     weights.
-  - `--usp <N>` — Unified Sequence Parallelism degree.
-  - `--cfg <N>` — Classifier-Free Guidance parallelism; set to 2 for models
+    - `--usp <N>` — Unified Sequence Parallelism degree.
+    - `--cfg <N>` — Classifier-Free Guidance parallelism; set to 2 for models
     that require negative-prompt computation, omit for distilled models.
-  - `--vae-patch-parallel-size 8` — parallelizes VAE decoding across all 8
+    - `--vae-patch-parallel-size 8` — parallelizes VAE decoding across all 8
     cards.
-  - `--vae-use-tiling` — enables tiled VAE decoding to reduce peak memory.
+    - `--vae-use-tiling` — enables tiled VAE decoding to reduce peak memory.
 - **Performance tips:**
-  - Installing mindie-sd and enabling Laser Attention
+    - Installing mindie-sd and enabling Laser Attention
     (`MINDIE_SD_FA_TYPE=ascend_laser_attention`) provides up to ~40%
     performance improvement at 720p resolution due to long-sequence attention
     optimization.
 - **Known limitations:**
-  - `MULTI_STREAM_MEMORY_REUSE=2` is required on NPU when using HSDP/FSDP2
+    - `MULTI_STREAM_MEMORY_REUSE=2` is required on NPU when using HSDP/FSDP2
     due to a multi-stream memory reuse bug. This is not needed on CUDA.

tests/benchmarks/patch/test_patch.py

@@ -35,6 +35,7 @@ from vllm_omni.benchmarks.patch.patch import (
     _iter_image_reference_inputs,
     _iter_video_reference_inputs,
     _omni_request_timeout_s,
+    async_request_openai_audio_speech,
     async_request_openai_chat_omni_completions,
     async_request_openai_image_edits_omni,
     async_request_openai_image_generations_omni,
@@ -474,6 +475,48 @@ async def test_streaming_error_chunk_marks_request_failed(mocker: MockerFixture)
     assert output.error == "EngineCore encountered an issue"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("headers", "env", "expected_duration_s"),
+    [
+        # Server-stated format wins over the env fallback.
+        ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {}, 1.0),
+        ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {"rate": "24000", "channels": "1"}, 1.0),
+        # Older servers send no headers: the env (default 24 kHz mono) applies.
+        ({}, {}, 4.0),
+        ({}, {"rate": "48000", "channels": "2"}, 1.0),
+    ],
+)
+async def test_audio_speech_stream_uses_server_pcm_format(
+    mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, headers, env, expected_duration_s
+):
+    from vllm_omni.metrics import definitions as defs
+
+    for key, var in (("rate", defs.AUDIO_SAMPLE_RATE_ENV), ("channels", defs.AUDIO_CHANNELS_ENV)):
+        if key in env:
+            monkeypatch.setenv(var, env[key])
+        else:
+            monkeypatch.delenv(var, raising=False)
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/audio/speech",
+        prompt_len=10,
+        output_len=20,
+    )
+    # 1 s of 48 kHz stereo s16le = 192000 bytes.
+    mock_response = MockResponse(200, [b"\x00" * 96000, b"\x00" * 96000])
+    mock_response.headers = headers
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_audio_speech(request_input, mock_session)
+
+    assert output.success is True
+    assert output.audio_duration == pytest.approx(expected_duration_s)
+
+
 @pytest.mark.asyncio
 async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
     request_input = RequestFuncInput(

tests/buildkite/test_amd_ready_r2_01.py

@@ -6,15 +6,17 @@ from shlex import split
 
 import pytest
 import yaml
+from _pytest.mark.expression import Expression
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
-R2_01_LABEL = "ROCm · Engine & Model Executor GPU Coverage (R2-01)"
+R2_01_LABEL = "ROCm · Model Executor GPU Coverage (R2-01)"
+ENGINE_LABEL = "Engine Test"
 SINGLE_GPU_MARKERS = (
-    "core_model and cuda and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)"
+    "core_model and cuda and not nvidia_only and not "
+    "(cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)"
 )
-EXCLUDED_TEST = "test_code2wav_streaming_batch_matches_ragged_flow_numerics"
 
 
 def _find_step(label: str) -> dict:
@@ -35,14 +37,28 @@ def _find_step(label: str) -> dict:
 
 def _assert_single_gpu_selection(command: str) -> None:
     argv = split(command)
-    assert "tests/engine/" in argv
+    assert "tests/engine/" not in argv
     assert "tests/model_executor/" in argv
     assert "tests/worker/test_batched_omni_output.py" not in argv
     assert argv[argv.index("-m") + 1] == SINGLE_GPU_MARKERS
-    assert argv[argv.index("-k") + 1] == f"not {EXCLUDED_TEST}"
+    assert "-k" not in argv
     assert argv[argv.index("--run-level") + 1] == "core_model"
 
 
+@pytest.mark.parametrize(
+    ("marks", "selected"),
+    [
+        ({"core_model", "cuda", "nvidia_only", "L4", "cards_1"}, False),
+        ({"core_model", "cuda"}, True),
+        ({"core_model", "cuda", "L4", "MI325", "cards_1"}, True),
+        ({"core_model", "cuda", "cards_2"}, False),
+    ],
+)
+def test_r2_01_marker_selection_contract(marks: set[str], selected: bool) -> None:
+    expression = Expression.compile(SINGLE_GPU_MARKERS)
+    assert expression.evaluate(marks.__contains__) is selected
+
+
 def test_r2_01_is_nonblocking_single_gpu_coverage() -> None:
     step = _find_step(R2_01_LABEL)
 
@@ -53,6 +69,11 @@ def test_r2_01_is_nonblocking_single_gpu_coverage() -> None:
     assert step["artifact_paths"] == ["artifacts/rocm-r2-01/**/*"]
 
 
+def test_dedicated_engine_job_preserves_r2_01_engine_coverage() -> None:
+    commands = "\n".join(_find_step(ENGINE_LABEL)["commands"])
+    assert "tests/engine/test_async_omni_engine_abort.py" in commands
+
+
 def test_r2_01_runs_once_fails_closed_and_publishes_results() -> None:
     commands = _find_step(R2_01_LABEL)["commands"]
     pytest_commands = [command for command in commands if command.lstrip().startswith("pytest ")]

tests/buildkite/test_pi05_ci_routing.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+from shlex import split
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
+AMD_NIGHTLY_PIPELINE = Path(".buildkite/amd/test-amd-nightly.yml")
+AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
+CUDA_MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
+CUDA_READY_PIPELINE = Path(".buildkite/cuda/test-ready.yml")
+PI05_LABEL = "Simple · Pi0.5 CPU Test"
+PI05_TEST_PATH = "tests/diffusion/models/pi05/test_pi05_units.py"
+
+
+def _pipeline_steps(pipeline_path: Path) -> list[dict]:
+    pipeline = yaml.safe_load(pipeline_path.read_text(encoding="utf-8"))
+    flattened = []
+
+    def walk(steps: list[dict]) -> None:
+        for step in steps:
+            flattened.append(step)
+            walk(step.get("steps", []))
+
+    walk(pipeline.get("steps", []))
+    return flattened
+
+
+def _find_step(label: str, pipeline_path: Path) -> dict:
+    step = next((step for step in _pipeline_steps(pipeline_path) if step.get("label") == label), None)
+    assert step is not None, f"missing pipeline step: {label}"
+    return step
+
+
+@pytest.mark.parametrize(
+    "pipeline_path",
+    [AMD_READY_PIPELINE, AMD_MERGE_PIPELINE],
+    ids=["ready", "merge"],
+)
+def test_pi05_cpu_suite_is_excluded_from_amd_pr_gates(pipeline_path: Path) -> None:
+    assert all(step.get("label") != PI05_LABEL for step in _pipeline_steps(pipeline_path))
+
+    step = _find_step("Simple · Diffusion Test · Shard %N/%t", pipeline_path)
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    assert f"--ignore={PI05_TEST_PATH}" in split(pytest_command)
+
+
+def test_pi05_cpu_suite_runs_nonblocking_in_amd_nightly() -> None:
+    steps = [step for step in _pipeline_steps(AMD_NIGHTLY_PIPELINE) if step.get("label") == PI05_LABEL]
+    assert len(steps) == 1
+
+    step = steps[0]
+    assert step["agent_pool"] == "mi300_1"
+    assert step["grade"] == "NonBlocking"
+    assert step["timeout_in_minutes"] == 60
+
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    argv = split(pytest_command)
+    assert PI05_TEST_PATH in argv
+    assert argv[argv.index("-m") + 1] == "core_model and cpu"
+
+
+@pytest.mark.parametrize(
+    "pipeline_path",
+    [CUDA_READY_PIPELINE, CUDA_MERGE_PIPELINE],
+    ids=["ready", "merge"],
+)
+def test_cuda_pr_gates_retain_pi05_cpu_coverage(pipeline_path: Path) -> None:
+    step = _find_step("Simple · Diffusion Test", pipeline_path)
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    argv = split(pytest_command)
+
+    assert "source_file_dependencies" not in step
+    assert "tests/diffusion" in argv
+    assert argv[argv.index("-m") + 1] == "core_model and cpu"
+    assert f"--ignore={PI05_TEST_PATH}" not in argv

tests/config/test_config_factory.py

@@ -2985,6 +2985,9 @@ class TestPlatformOverrides:
             ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
             ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
             ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_optimized.yaml", "moss_tts_local"),
+            ("moss_tts_local.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency_mps.yaml", "moss_tts_local"),
         ],
     )
     @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
@@ -3004,9 +3007,142 @@ class TestPlatformOverrides:
             # v2 only engages the native plane on stages declaring support.
             assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
 
+    @pytest.mark.parametrize("memory_gib", [24, 80, 139, 140, 141, None])
+    @pytest.mark.parametrize("mps_available", [False, True])
+    def test_moss_local_cuda_default_selects_mrv2_capacity_and_mps(self, monkeypatch, memory_gib, mps_available):
+        from vllm.platforms import current_platform
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda")
+
+        def memory():
+            if memory_gib is None:
+                raise RuntimeError("NVML unavailable")
+            return memory_gib * 1024**3
+
+        monkeypatch.setattr(current_platform, "get_device_total_memory", memory)
+        monkeypatch.setattr(
+            "shutil.which", lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps_available else None
+        )
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        high_capacity = mps_available and memory_gib is not None and memory_gib >= 140
+        capacity = 128 if high_capacity else 64
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        talker, codec = config.stage_configs
+        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == capacity
+        assert talker.cache_config.kv_cache_memory_bytes == (32 * 1024**3 if high_capacity else None)
+        assert talker.cache_config.gpu_memory_utilization == 0.60
+        assert codec.cache_config.gpu_memory_utilization == 0.30
+        assert talker.cache_config.enable_prefix_caching is high_capacity
+        assert not codec.cache_config.enable_prefix_caching
+        for stage in config.stage_configs:
+            assert stage.model_config.use_v2_model_runner
+            assert stage.runtime_config.cuda_mps is mps_available
+            assert not stage.runtime_config.env
+        overrides = talker.model_config.hf_overrides or {}
+        assert overrides.get("mrv2_batch_prefill", False) is high_capacity
+        assert overrides.get("mrv2_direct_tokens", False) is high_capacity
+        assert overrides.get("local_compile_audio_sampler", False) is high_capacity
+
+        # The serving engine still consumes the legacy representation.
+        stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {})
+        for stage in stages:
+            assert stage.yaml_engine_args["max_num_seqs"] == capacity
+            assert stage.yaml_engine_args["use_v2_model_runner"]
+            assert stage.yaml_runtime.get("cuda_mps", False) is mps_available
+        assert stages[0].yaml_engine_args.get("kv_cache_memory_bytes") == (32 * 1024**3 if high_capacity else None)
+        assert stages[0].yaml_engine_args["enable_prefix_caching"] is high_capacity
+
+    @pytest.mark.parametrize("platform", ["cpu", "npu", "xpu", "rocm", "musa"])
+    def test_moss_local_non_cuda_default_preserves_v1(self, monkeypatch, platform):
+        from vllm.platforms import current_platform
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_omni_platform, "device_name", platform)
+
+        def unexpected_probe(*args):
+            pytest.fail("Non-CUDA selection must not probe CUDA memory or MPS")
+
+        monkeypatch.setattr(current_platform, "get_device_total_memory", unexpected_probe)
+        monkeypatch.setattr("shutil.which", unexpected_probe)
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        for stage in config.stage_configs:
+            assert not stage.model_config.use_v2_model_runner
+            assert not stage.runtime_config.cuda_mps
+            assert stage.scheduler_config.max_num_seqs == 64
+        assert config.stage_configs[0].cache_config.kv_cache_memory_bytes is None
+        stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {})
+        assert all(not stage.yaml_engine_args["use_v2_model_runner"] for stage in stages)
+        assert all(not stage.yaml_runtime.get("cuda_mps", False) for stage in stages)
+
+    @pytest.mark.parametrize(
+        "filename,runner,mps,capacity",
+        [
+            ("moss_tts_local_v1.yaml", False, False, 64),
+            ("moss_tts_local_mrv2.yaml", True, False, 64),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", True, False, 128),
+        ],
+    )
+    def test_moss_local_explicit_profile_overrides_cuda_default(self, monkeypatch, filename, runner, mps, capacity):
+        from vllm.platforms import current_platform
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda")
+        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: 141 * 1024**3)
+        monkeypatch.setattr("shutil.which", lambda command: "/usr/bin/nvidia-cuda-mps-control")
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = get_deploy_config_path(filename)
+        config = VllmOmniConfig.from_pipeline_config(pipeline, deploy_config_path=path)
+        for stage in config.stage_configs:
+            assert stage.model_config.use_v2_model_runner is runner
+            assert stage.runtime_config.cuda_mps is mps
+            assert stage.scheduler_config.max_num_seqs == capacity
+        stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {}, deploy_config_path=path)
+        for stage in stages:
+            assert stage.yaml_engine_args["use_v2_model_runner"] is runner
+            assert stage.yaml_runtime.get("cuda_mps", False) is mps
+            assert stage.yaml_engine_args["max_num_seqs"] == capacity
+
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_moss_local_system_profile_reaches_native_runner_and_cuda_only_mps(self, platform):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_optimized.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        is_cuda = platform == "cuda"
+        assert deploy.cuda_mps is is_cuda
+        assert all(stage.yaml_runtime.get("cuda_mps", False) is is_cuda for stage in stages)
+        assert all(stage.yaml_engine_args["use_v2_model_runner"] is is_cuda for stage in stages)
+        if is_cuda:
+            args = stages[0].yaml_engine_args
+            assert args["enable_prefix_caching"] is True
+            assert args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+            assert args["hf_overrides"]["mrv2_batch_prefill"] is True
+            assert args["hf_overrides"]["mrv2_direct_tokens"] is True
+            assert deploy.connectors["shm"]["extra"]["codec_first_chunk_fast_path"] == 1
+            assert args["hf_overrides"]["local_compile_audio_sampler"] is True
+            extra = deploy.connectors["shm"]["extra"]
+            assert extra["codec_first_chunk_max_active_streams"] == 32
+            assert extra["generation_min_batch_size"] == 16
+            assert extra["generation_max_wait_ms"] == 6
+            assert not stages[0].yaml_runtime.get("env")
+        else:
+            assert all(not stage.yaml_engine_args["enable_prefix_caching"] for stage in stages)
+            assert all(not stage.yaml_runtime.get("env") for stage in stages)
+
+    def test_platform_mps_rejects_non_boolean(self):
+        deploy = load_deploy_config(get_deploy_config_path("moss_tts_local_mrv2_optimized.yaml"))
+        deploy.platforms["cuda"]["cuda_mps"] = "true"
+        with pytest.raises(ValueError, match="platform cuda_mps must be a boolean"):
+            _apply_platform_overrides(deploy, platform="cuda")
+
     def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
         pipeline = resolve_pipeline_config("moss_tts_local")
-        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        base_path = Path(get_deploy_config_path("moss_tts_local_v1.yaml"))
         candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
         base = load_deploy_config(base_path)
         candidate = load_deploy_config(candidate_path)
@@ -3052,7 +3188,7 @@ class TestPlatformOverrides:
         assert compilation["cudagraph_mode"] == "FULL"
         assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
         assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
-        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        base_path = Path(get_deploy_config_path("moss_tts_local_v1.yaml"))
         for platform in ("npu", "xpu", "rocm", "musa"):
             actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
             expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 96,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 7,
     }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 66,
+        ModelEnvironmentVariableDisposition.PROMOTE: 67,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/config/test_omni_config.py

@@ -1887,7 +1887,7 @@ def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
 
 @pytest.mark.parametrize("source", ["default", "topology", "deploy", "stage-cli"])
 @pytest.mark.parametrize("key_container", [list, tuple])
-def test_diffusion_stage_payload_keys_roundtrip(source, key_container):
+def test_diffusion_stage_payload_keys_roundtrip(source, key_container, tmp_path):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
 
@@ -1920,7 +1920,7 @@ def test_diffusion_stage_payload_keys_roundtrip(source, key_container):
     legacy_stage = merge_pipeline_deploy(pipeline, deploy)[0]
     legacy_args = {**legacy_stage.yaml_engine_args, **(override_keys if source == "stage-cli" else {})}
     restored_stage = ForkingPickler.loads(ForkingPickler.dumps(stage))
-    engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model="test-model")
+    engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model=str(tmp_path))
     diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs(
         engine_args, stage_id=restored_stage.stage_id, include_engine_adapter_metadata=True
     )

tests/core/sched/test_generation_scheduler_init.py

@@ -44,10 +44,13 @@ def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128
     )
 
 
-@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize(
+    "profile,min_batch,max_wait_ms",
+    [("high_concurrency", 32, 12), ("low_latency", 32, 12), ("optimized", 16, 6)],
+)
 @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
 def test_moss_profile_generation_constructor_after_platform_resolution(
-    construct_scheduler, monkeypatch, mocker, profile, platform
+    construct_scheduler, monkeypatch, mocker, profile, min_batch, max_wait_ms, platform
 ):
     # merge_pipeline_deploy owns platform resolution. Mock the detected device
     # instead of applying an override first and then resolving the worker's
@@ -66,11 +69,12 @@ def test_moss_profile_generation_constructor_after_platform_resolution(
     warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
     scheduler = construct_scheduler(config)
     # The raw extras persist through fallback; the live constructor must cope.
-    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    assert extras["generation_min_batch_size"] == min_batch
+    assert extras["generation_max_wait_ms"] == max_wait_ms
     if platform == "cuda":
         assert scheduler._native_data_plane and scheduler.input_coordinator is not None
-        assert scheduler._generation_min_batch_size == 32
-        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_min_batch_size == min_batch
+        assert scheduler._generation_max_wait_s == max_wait_ms / 1000
         assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
         warning.assert_not_called()
         if profile == "low_latency":

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -16,6 +16,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm import SamplingParams
 
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
@@ -23,6 +24,9 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
     uses_native_mrv2_data_plane,
 )
 from vllm_omni.core.sched.output import OmniChunkRecvHandle
+from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import _LoadEntry
+from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+from vllm_omni.request import OmniRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -440,3 +444,20 @@ class TestTimeoutDetection(unittest.TestCase):
         self.assertNotIn("r1", coord.requests_with_ready_chunks)
         self.assertEqual([r.request_id for r in coord._waiting_for_chunk_running], ["r2"])
         self.assertEqual([h.request_id for h in coord.pending_chunk_registrations], ["r2"])
+
+
+def test_sender_address_on_the_engine_request_reaches_both_receive_paths():
+    """Ensure the sender address the orchestrator puts on an engine request reaches the receivers."""
+    sender = {"host": "10.0.0.2", "zmq_port": 50071}
+    engine_request = build_engine_core_request_from_tokens(
+        "req-1", {"prompt_token_ids": [0]}, SamplingParams(max_tokens=1)
+    )
+    engine_request.payload_sender_info = sender
+    request = OmniRequest.from_engine_core_request(engine_request, block_hasher=None)
+    # Build the coordinator and process the request with sender info
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.process_pending_full_payload_inputs(MockQueue([request]), stage_recv_req_ids=set())
+
+    # Ensure that the payload send info is accessible on both pending input registrations and a wrapped load entry
+    assert [handle.payload_sender_info for handle in coordinator.pending_input_registrations] == [sender]
+    assert _LoadEntry(request).source_metadata == {"source_host": "10.0.0.2", "source_port": 50071}

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/cache/test_cache_dit.py

@@ -59,6 +59,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
         "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3,
         "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
         "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
+        "Kandinsky6TI2VAPipeline": cd_model_specific.enable_cache_for_kandinsky6,
         "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2,
         "MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2,
     }

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -1,5 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass
-from types import SimpleNamespace
 
 import pytest
 import torch
@@ -25,7 +27,7 @@ class E2EOperator:
         h_size = z.shape[0]
         w_size = z.shape[1]
 
-        tasks = []
+        tasks: list[TileTask] = []
         for i in range(rows_num):
             for j in range(cols_num):
                 tasks.append(
@@ -62,12 +64,14 @@ class E2EOperator:
 @dataclass
 class FakeWorldGroup:
     device_group: object
+    cpu_group: object = None
 
 
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
-        self.distributed_executor = SimpleNamespace(parallel_size=2, group=None)
+        self.distributed_executor = DistributedVaeExecutor()
+        self.distributed_executor.set_parallel_size(2)
 
 
 @pytest.fixture(autouse=True)

tests/diffusion/distributed/test_sp_storage_lifetime.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import weakref
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelPartialInput,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.base import HookRegistry
+from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook, apply_sequence_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
+
+
+@pytest.fixture
+def parallel_state(monkeypatch):
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    def configure(size=4, rank=1):
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    configure()
+    return configure
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@torch.inference_mode()
+def test_cosmos3_real_sp_hook_releases_full_embedding(parallel_state, sequence_length, batch, rank):
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare, Cosmos3VFMTransformer
+
+    parallel_state(rank=rank)
+    model = nn.Module()
+    model.gen_sp_prepare = Cosmos3GenSPPrepare()
+    metadata = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
+    assert metadata[0].clone_shard
+    assert not metadata[1].clone_shard and not metadata[2].clone_shard
+    apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), {"gen_sp_prepare": metadata})
+
+    def prepare():
+        hidden = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_ref = weakref.ref(hidden)
+        full_pointer = hidden.untyped_storage().data_ptr()
+        cos = torch.ones(batch, sequence_length, 1, 2)
+        sin = torch.zeros_like(cos)
+        return model.gen_sp_prepare(hidden, cos, sin), full_ref, full_pointer
+
+    for _ in range(2):
+        with set_forward_context():
+            (hidden, cos, sin), full_ref, full_pointer = prepare()
+            # Check the real backing allocation, not only Tensor.data_ptr():
+            # nonzero-offset and inference-mode views can hide the shared base.
+            assert full_ref() is None
+            assert hidden.untyped_storage().data_ptr() != full_pointer
+            assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+            assert hidden.is_contiguous()
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[rank]
+            torch.testing.assert_close(hidden, expected, rtol=0, atol=0)
+            assert cos.shape[1] == sin.shape[1] == hidden.shape[1]
+
+            # Once the next layer replaces the shard, the hook must not keep
+            # even the smaller initial activation alive through later layers.
+            shard_ref = weakref.ref(hidden)
+            hidden = hidden + 1
+            assert shard_ref() is None
+
+
+@pytest.mark.parametrize("clone_shard", [False, True])
+@pytest.mark.parametrize("world_size", [1, 4])
+def test_shard_storage_policy_is_opt_in_and_skips_single_rank(parallel_state, clone_shard, world_size):
+    parallel_state(size=world_size, rank=0)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.arange(128).reshape(1, 16, 8)
+    spec = SequenceParallelInput(split_dim=1, clone_shard=clone_shard)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, spec)
+    torch.testing.assert_close(shard, tensor.chunk(world_size, dim=1)[0])
+    owns_storage = shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+    assert owns_storage == (clone_shard and world_size > 1)
+    if world_size == 1:
+        assert shard is tensor
+
+
+def test_owned_shard_preserves_gradients(parallel_state):
+    parallel_state(rank=1)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.randn(2, 16, 8, requires_grad=True)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, SequenceParallelInput(split_dim=1, clone_shard=True))
+        assert shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+        assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
+        shard.sum().backward()
+    expected = torch.zeros_like(tensor)
+    expected[:, 4:8] = 1
+    torch.testing.assert_close(tensor.grad, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("shard_group", [None, "video"])
+@torch.inference_mode()
+def test_pre_forward_clone_owns_storage_and_preserves_shard_metadata(
+    parallel_state, sequence_length, batch, keyword, shard_group
+):
+    parallel_state(rank=1)
+    full_refs = []
+    full_pointers = []
+
+    class Module(nn.Module):
+        def forward(self, x):
+            assert x.untyped_storage().data_ptr() != full_pointers[-1]
+            assert x.untyped_storage().nbytes() == x.numel() * x.element_size()
+            assert x.is_contiguous()
+            return x
+
+    module = Module()
+    spec = SequenceParallelInput(
+        split_dim=1,
+        split_output=False,
+        auto_pad=True,
+        clone_shard=True,
+        shard_group=shard_group,
+    )
+    apply_sequence_parallel(module, SequenceParallelConfig(ulysses_degree=4), {"": {"x": spec}})
+
+    def run():
+        tensor = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_refs.append(weakref.ref(tensor))
+        full_pointers.append(tensor.untyped_storage().data_ptr())
+        return module(x=tensor) if keyword else module(tensor)
+
+    for _ in range(2):
+        with set_forward_context():
+            ctx = get_forward_context()
+            if shard_group is not None:
+                ctx.sp_shard_metadata[shard_group] = 999  # Stale metadata must be replaced.
+            shard = run()
+            assert all(ref() is None for ref in full_refs)
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[1]
+            torch.testing.assert_close(shard, expected, rtol=0, atol=0)
+            assert ctx.sp_shard_metadata == ({} if shard_group is None else {shard_group: sequence_length})
+            assert ctx.sp_original_seq_len == (sequence_length if padding and shard_group is None else None)
+
+
+@pytest.mark.parametrize("phase", ["success", "non_tensor", "forward_error", "split_error"])
+@pytest.mark.parametrize("keyword", [False, True])
+@torch.inference_mode()
+def test_split_hook_does_not_retain_arguments_on_any_exit(parallel_state, phase, keyword):
+    class Module(nn.Module):
+        def forward(self, x):
+            if phase == "forward_error":
+                raise RuntimeError("forward failed")
+            return None if phase == "non_tensor" else x
+
+    module = Module()
+    index = 1 if phase == "split_error" else 0
+    hook = SequenceParallelSplitHook(
+        {index: SequenceParallelInput(split_dim=1, split_output=True, clone_shard=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run():
+        tensor = torch.randn(1, 16, 8)
+        refs.append(weakref.ref(tensor))
+        return module(x=tensor) if keyword else module(tensor)
+
+    with set_forward_context():
+        if phase.endswith("error"):
+            with pytest.raises((RuntimeError, ValueError)):
+                run()
+        else:
+            result = run()
+            assert (result is None) == (phase == "non_tensor")
+    assert all(ref() is None for ref in refs)
+
+
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("tensor_source", [False, True])
+@torch.inference_mode()
+def test_partial_output_split_keeps_only_lengths_across_calls(parallel_state, keyword, tensor_source):
+    parallel_state(rank=1)
+
+    class Module(nn.Module):
+        def forward(self, x, text):
+            return x
+
+    module = Module()
+    hook = SequenceParallelSplitHook(
+        {0: SequenceParallelPartialInput(split_dim=0, text_len_source="text", split_output=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run(text_length):
+        tensor = torch.arange(text_length + 8).reshape(-1, 1)
+        text = torch.zeros(text_length, 8) if tensor_source else text_length
+        refs.append(weakref.ref(tensor))
+        if tensor_source:
+            refs.append(weakref.ref(text))
+        return module(x=tensor, text=text) if keyword else module(tensor, text)
+
+    for text_length in (2, 4):
+        with set_forward_context():
+            result = run(text_length)
+        expected = torch.cat([torch.arange(text_length), torch.arange(text_length + 2, text_length + 4)]).reshape(-1, 1)
+        torch.testing.assert_close(result, expected)
+        assert all(ref() is None for ref in refs)
+        assert hook._text_len_cache == {"text": text_length}

tests/diffusion/distributed/test_vae_batch_parallel.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from dataclasses import dataclass
+from datetime import timedelta
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from diffusers import AutoencoderKL, AutoencoderKLFlux2
+from torch import nn
+
+from vllm_omni.config.omni_config import OmniStageDiffusionParallelConfig
+from vllm_omni.diffusion import registry
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl as kl_module
+from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as executor_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_flux2 import DistributedAutoencoderKLFlux2
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeExecutor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
+
+
+class _RecordingDecoder(nn.Module):
+    def __init__(self, dtype):
+        super().__init__()
+        self.conv = nn.Conv2d(1, 3, kernel_size=1, bias=False, dtype=dtype)
+        with torch.no_grad():
+            self.conv.weight.copy_(torch.arange(1, 4, dtype=dtype).reshape(3, 1, 1, 1))
+        self.batch_sizes = []
+
+    def forward(self, z):
+        self.batch_sizes.append(z.shape[0])
+        return self.conv(z)
+
+
+def _make_vae(vae_type=DistributedAutoencoderKLFlux2, dtype=torch.float32):
+    """Keep the native decode methods; replace only the learned decoder."""
+    vae = vae_type.__new__(vae_type)
+    nn.Module.__init__(vae)
+    vae.register_to_config(use_post_quant_conv=False)
+    vae.decoder = _RecordingDecoder(dtype)
+    vae.post_quant_conv = None
+    vae.use_slicing = False
+    vae.use_tiling = False
+    return vae
+
+
+@dataclass
+class _WorldGroup:
+    device_group: dist.ProcessGroup
+    cpu_group: dist.ProcessGroup
+
+
+def _batch_worker(rank, rendezvous):
+    torch.set_num_threads(1)
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=4, timeout=timedelta(seconds=60))
+    try:
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(executor_module, "get_world_group", lambda: _WorldGroup(dist.group.WORLD, dist.group.WORLD))
+            for name in (
+                "get_data_parallel_world_size",
+                "get_pipeline_parallel_world_size",
+                "get_classifier_free_guidance_world_size",
+            ):
+                patch.setattr(kl_module, name, lambda: 1)
+            for vae_type, native_type in (
+                (DistributedAutoencoderKL, AutoencoderKL),
+                (DistributedAutoencoderKLFlux2, AutoencoderKLFlux2),
+            ):
+                for dtype, autocast in ((torch.float32, False), (torch.bfloat16, False), (torch.float32, True)):
+                    vae = _make_vae(vae_type, dtype)
+                    vae.init_distributed()
+                    vae.set_parallel_size(4, mode="batch")
+                    vae.use_tiling = True
+                    assert not vae.is_distributed_enabled()  # Spatial encode/decode stays disabled.
+                    vae.use_tiling = False
+                    # Reuse the same instance with changing batches and contents.
+                    for batch, return_dict, expected_sizes in (
+                        (2, False, (1, 1, 0, 0)),
+                        (5, True, (2, 1, 1, 1)),
+                        (1, True, (1, 1, 1, 1)),
+                    ):
+                        z = torch.arange(batch * 6, dtype=dtype).reshape(batch, 1, 2, 3) + batch
+                        with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast):
+                            expected = native_type.decode(vae, z, return_dict=False)[0]
+                            vae.decoder.batch_sizes.clear()
+                            if vae_type is DistributedAutoencoderKL and not return_dict:
+                                result = vae.decode(z, False, torch.Generator().manual_seed(0))
+                            else:
+                                result = vae.decode(z, return_dict=return_dict)
+                        actual = result.sample if return_dict else result[0]
+                        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                        assert actual.dtype == expected.dtype
+                        assert sum(vae.decoder.batch_sizes) == expected_sizes[rank]
+                        assert not vae.use_tiling
+
+                    # Degree one bypasses collectives and keeps native slicing.
+                    vae.set_parallel_size(1, mode="batch")
+                    vae.use_slicing = True
+                    z = torch.ones(2, 1, 2, 3, dtype=dtype)
+                    vae.decoder.batch_sizes.clear()
+                    if vae_type is DistributedAutoencoderKL:
+                        actual = vae.decode(z, False, torch.Generator().manual_seed(0))[0]
+                    else:
+                        actual = vae.decode(z, return_dict=False)[0]
+                    assert vae.decoder.batch_sizes == [1, 1]
+                    expected = native_type.decode(vae, z, return_dict=False)[0]
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+                # Execute native tiled decode, including the B1 fallback and
+                # uneven batches with idle ranks; preserving the flag alone
+                # does not prove that tiling is still used.
+                vae = _make_vae(vae_type)
+                vae.init_distributed()
+                vae.set_parallel_size(4, mode="batch")
+                vae.use_tiling = True
+                vae.tile_latent_min_size = vae.tile_sample_min_size = 4
+                vae.tile_overlap_factor = 0.5
+                patch.setattr(
+                    vae, "_strategy_select", lambda _: pytest.fail("Batch mode must not distribute spatial tiles")
+                )
+                for batch, expected_sizes in (
+                    (1, (1, 1, 1, 1)),
+                    (2, (1, 1, 0, 0)),
+                    (5, (2, 1, 1, 1)),
+                ):
+                    z = torch.arange(batch * 36, dtype=torch.float32).reshape(batch, 1, 6, 6)
+                    expected = native_type.decode(vae, z, return_dict=False)[0]
+                    vae.decoder.batch_sizes.clear()
+                    actual = vae.decode(z).sample
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                    assigned_batch = expected_sizes[rank]
+                    if assigned_batch:
+                        # More than one decoder call proves native tiling ran.
+                        assert len(vae.decoder.batch_sizes) > 1
+                        assert set(vae.decoder.batch_sizes) == {assigned_batch}
+                    else:
+                        assert vae.decoder.batch_sizes == []
+                    assert vae.use_tiling
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(not dist.is_gloo_available(), reason="Requires the CPU Gloo backend")
+def test_batch_decode_matches_native_on_all_ranks(tmp_path):
+    mp.spawn(_batch_worker, args=(f"file://{tmp_path / 'rendezvous'}",), nprocs=4, join=True)
+
+
+def test_batch_decode_without_distributed_initialization_uses_native(monkeypatch, mocker):
+    vae = _make_vae()
+    executor = mocker.Mock(spec=DistributedVaeExecutor)
+    executor.parallel_mode = "batch"
+    executor.parallel_size = executor.world_size = 4
+    vae.distributed_executor = executor
+    monkeypatch.setattr(dist, "is_initialized", lambda: False)
+    z = torch.ones(2, 1, 2, 3)
+    expected = AutoencoderKLFlux2.decode(vae, z, return_dict=False)[0]
+    torch.testing.assert_close(vae.decode(z).sample, expected, rtol=0, atol=0)
+    executor.execute.assert_not_called()
+
+
+def test_flux2_batch_encode_keeps_native_tiling(monkeypatch, mocker):
+    vae = _make_vae()
+    vae.register_to_config(use_quant_conv=False)
+    vae.encoder = _RecordingDecoder(torch.float32)
+    vae.quant_conv = None
+    vae.use_tiling = True
+    vae.tile_sample_min_size = vae.tile_latent_min_size = 4
+    vae.tile_overlap_factor = 0.5
+    executor = mocker.Mock(spec=DistributedVaeExecutor)
+    executor.parallel_mode = "batch"
+    executor.parallel_size = executor.world_size = 4
+    executor.group = None
+    vae.distributed_executor = executor
+    monkeypatch.setattr(dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(dist, "get_world_size", lambda group: 4)
+    x = torch.arange(36, dtype=torch.float32).reshape(1, 1, 6, 6)
+    expected = AutoencoderKLFlux2._encode(vae, x)
+    native_tiling = mocker.spy(vae, "_tiled_encode")
+    actual = vae._encode(x)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    native_tiling.assert_called_once_with(x)
+    executor.execute.assert_not_called()
+
+
+@pytest.mark.parametrize("config_type", [DiffusionParallelConfig, OmniStageDiffusionParallelConfig])
+@pytest.mark.parametrize("axis", ["data_parallel_size", "pipeline_parallel_size", "cfg_parallel_size"])
+def test_batch_mode_rejects_separate_request_or_pipeline_groups(config_type, axis):
+    with pytest.raises(ValueError, match="batch.*DP, PP, and CFG"):
+        config_type(vae_parallel_mode="batch", **{axis: 2})
+
+
+@pytest.mark.parametrize("config_type", [DiffusionParallelConfig, OmniStageDiffusionParallelConfig])
+def test_batch_mode_accepts_tensor_parallel_group(config_type):
+    config = config_type(tensor_parallel_size=4, vae_patch_parallel_size=4, vae_parallel_mode="batch")
+    assert config.world_size == 4
+
+
+def test_batch_mode_rejects_implicitly_resolved_data_parallelism():
+    config = DiffusionParallelConfig(vae_parallel_mode="batch", data_parallel_size=None)
+    with pytest.raises(ValueError, match="batch.*DP"):
+        config.resolve_data_parallel_size(2)
+
+
+@pytest.mark.parametrize("tiling", [False, True])
+def test_registry_preserves_requested_tiling_in_batch_mode(monkeypatch, mocker, tiling):
+    vae = _make_vae()
+    vae.distributed_executor = mocker.Mock(spec=DistributedVaeExecutor)
+
+    class Pipeline(nn.Module):
+        def __init__(self, od_config):
+            super().__init__()
+            self.vae = vae
+
+    monkeypatch.setattr(registry.DiffusionModelRegistry, "_try_load_model_cls", lambda _: Pipeline)
+    config = OmniDiffusionConfig(
+        model_class_name="test_batch_vae",
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=4, vae_patch_parallel_size=4, vae_parallel_mode="batch"
+        ),
+        vae_use_tiling=tiling,
+        vae_use_slicing=True,
+        vae_fast_path="off",
+    )
+    model = registry.initialize_model(config)
+    assert model.vae.use_tiling is tiling
+    assert model.vae.use_slicing is True
+    assert config.vae_use_tiling is tiling
+    vae.distributed_executor.set_parallel_size.assert_called_once_with(4, mode="batch")
+
+
+@pytest.mark.parametrize("vae_type", [nn.Identity, DistributedAutoencoderKLWan])
+def test_registry_rejects_unsupported_batch_vae(monkeypatch, vae_type):
+    class Pipeline(nn.Module):
+        def __init__(self, od_config):
+            super().__init__()
+            self.vae = vae_type.__new__(vae_type)
+            nn.Module.__init__(self.vae)
+
+    monkeypatch.setattr(registry.DiffusionModelRegistry, "_try_load_model_cls", lambda _: Pipeline)
+    config = OmniDiffusionConfig(
+        model_class_name="unsupported_batch_vae",
+        parallel_config=DiffusionParallelConfig(vae_parallel_mode="batch"),
+    )
+    with pytest.raises(ValueError, match="batch.*AutoencoderKL"):
+        registry.initialize_model(config)
+
+
+@pytest.mark.parametrize("mode", ["tile", "batch"])
+def test_klein_uses_distributed_vae_only_for_batch_mode(monkeypatch, tmp_path, mode):
+    from vllm_omni.diffusion.models.flux2_klein import pipeline_flux2_klein as module
+
+    class StopAtVaeLoadError(Exception):
+        pass
+
+    def load_component(factory, *args, subfolder, **kwargs):
+        if subfolder == "text_encoder":
+            return nn.Identity()
+        expected = DistributedAutoencoderKLFlux2 if mode == "batch" else AutoencoderKLFlux2
+        assert factory.__self__ is expected
+        raise StopAtVaeLoadError
+
+    monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
+    monkeypatch.setattr(module, "from_pretrained_with_prefetch", load_component)
+    monkeypatch.setattr(module.FlowMatchEulerDiscreteScheduler, "from_pretrained", lambda *args, **kwargs: None)
+    monkeypatch.setattr(module.Qwen2TokenizerFast, "from_pretrained", lambda *args, **kwargs: None)
+    config = OmniDiffusionConfig(model=str(tmp_path), parallel_config=DiffusionParallelConfig(vae_parallel_mode=mode))
+    with pytest.raises(StopAtVaeLoadError):
+        module.Flux2KleinPipeline(od_config=config)

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/distributed/test_wan_vae_fastpath_kernels.py

@@ -22,6 +22,7 @@ from tests.diffusion.distributed.wan_vae_fastpath_helpers import (  # noqa: F401
     original_wan_rms_norm,
     unpatched_wan_rms_norm,
 )
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
     install_wan_vae_fastpath,
@@ -788,6 +789,135 @@ def test_cached_conv_chunk_loop_matches_reference_including_mixed_paths(
             assert torch.equal(out, ref), pattern
 
 
+def _first_frame_test_conv(dtype: torch.dtype, **kwargs) -> WanCausalConv3d:
+    options = dict(kernel_size=3, padding=1)
+    options.update(kwargs)
+    conv = WanCausalConv3d(16, 32, **options).eval().to(device="cuda", dtype=dtype)
+    conv.to(memory_format=torch.channels_last_3d)
+    setattr(conv, fp.CFG_ATTR, fp.FastPathConfig(channels_last=True))
+    return conv
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("temporal_only", [False, True])
+@pytest.mark.parametrize("start", [None, "Rep"])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_first_frame_conv2d_preserves_following_cached_frame(dtype, temporal_only, start, monkeypatch) -> None:
+    torch.manual_seed(19)
+    kwargs = dict(kernel_size=(3, 1, 1), padding=(1, 0, 0)) if temporal_only else {}
+    conv = _first_frame_test_conv(dtype, **kwargs)
+    chunks = [_dense_5d((2, 16, 1, 7, 11), dtype, True) for _ in range(2)]
+    calls = []
+    original = fp._first_frame_conv2d
+
+    def spy(*args):
+        pair = original(*args)
+        calls.append(pair is not None)
+        return pair
+
+    monkeypatch.setattr(fp, "_first_frame_conv2d", spy)
+    cache = [start]
+    reference_cache = None
+    # Disable TF32 so the FP32 comparison isolates Conv2d/Conv3d accumulation.
+    with torch.backends.cudnn.flags(enabled=True, allow_tf32=False):
+        for index, chunk in enumerate(chunks):
+            expected = WanCausalConv3d.forward(conv, chunk, reference_cache)
+            actual = fp._run_cached_causal_conv(conv, chunk, cache, 0)
+            assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+            tol = dict(rtol=1e-4, atol=1e-5) if dtype is torch.float32 else dict(rtol=2e-2, atol=2e-2)
+            torch.testing.assert_close(actual, expected, **tol)
+            previous = torch.zeros_like(chunk) if index == 0 else chunks[index - 1]
+            assert _bits_equal(cache[0], torch.cat([previous, chunk], dim=2))
+            assert cache[0].is_contiguous(memory_format=torch.channels_last_3d)
+            reference_cache = chunk.clone() if index == 0 and start is None else torch.cat([previous, chunk], dim=2)
+    assert calls == [True], "A later T=1 chunk with history must retain the temporal convolution"
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("with_bias", [False, True])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_first_frame_conv2d_deferred_bias_matches_autocast(dtype, with_bias) -> None:
+    torch.manual_seed(20)
+    conv = _first_frame_test_conv(torch.float32)
+    if with_bias:
+        conv.bias.copy_(torch.randn_like(conv.bias))
+    else:
+        conv.register_parameter("bias", None)
+    x = _dense_5d((1, 16, 1, 7, 11), torch.float32, True)
+    cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        expected = WanCausalConv3d.forward(conv, x)
+        out, bias = fp._run_cached_causal_conv(conv, x, cache, 0, return_bias=True)
+        assert out.dtype == expected.dtype == dtype
+        if with_bias:
+            assert bias is not None and bias.dtype == dtype
+            assert _bits_equal(bias, conv.bias.to(dtype))
+            out = out + bias.view(1, -1, 1, 1, 1)
+        else:
+            assert bias is None
+    assert isinstance(cache[0], torch.Tensor)
+    assert cache[0].dtype == x.dtype, "Autocast must not round the temporal cache"
+    assert _bits_equal(cache[0], torch.cat([torch.zeros_like(x), x], dim=2))
+    torch.testing.assert_close(out, expected, rtol=2e-2, atol=2e-2)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "case,kwargs",
+    [
+        ("lossless", {}),
+        ("multiframe", {}),
+        ("kernel", dict(kernel_size=(3, 5, 5), padding=(1, 2, 2))),
+        ("stride", dict(stride=(1, 2, 2))),
+        ("dilation", dict(padding=(1, 2, 2))),
+        ("groups", {}),
+        ("temporal_padding", dict(padding=(2, 1, 1))),
+    ],
+)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_first_frame_conv2d_declines_unsupported_calls(case, kwargs) -> None:
+    conv = _first_frame_test_conv(torch.float32, **kwargs)
+    if case == "lossless":
+        setattr(conv, fp.CFG_ATTR, fp.FastPathConfig())
+    elif case == "dilation":
+        conv.dilation = (1, 2, 2)
+    elif case == "groups":
+        conv.groups = 2
+        conv.weight = nn.Parameter(conv.weight[:, :8].contiguous(memory_format=torch.channels_last_3d))
+    x = _dense_5d((1, 16, 2 if case == "multiframe" else 1, 7, 11), torch.float32, True)
+    assert fp._first_frame_conv2d(conv, x, conv.bias) is None
+    expected = WanCausalConv3d.forward(conv, x)
+    actual = fp._run_cached_causal_conv(conv, x, [None], 0)
+    assert _bits_equal(actual, expected), "Declined inputs must keep the existing exact fallback"
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("attribute,value", [("padding_mode", "reflect"), ("padding", (0, 1, 1))])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_first_frame_conv2d_declines_nonzero_native_padding(attribute, value) -> None:
+    conv = _first_frame_test_conv(torch.float32)
+    setattr(conv, attribute, value)
+    x = _dense_5d((1, 16, 1, 7, 11), torch.float32, True)
+    assert fp._first_frame_conv2d(conv, x, conv.bias) is None
+
+
+@torch.no_grad()
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_first_frame_conv2d_uses_reloaded_weights() -> None:
+    torch.manual_seed(21)
+    conv = _first_frame_test_conv(torch.bfloat16)
+    x = _dense_5d((1, 16, 1, 7, 11), torch.bfloat16, True)
+    before = fp._run_cached_causal_conv(conv, x, [None], 0)
+    replacement = {name: -value.clone() for name, value in conv.state_dict().items()}
+    conv.load_state_dict(replacement)
+    after = fp._run_cached_causal_conv(conv, x, [None], 0)
+    expected = WanCausalConv3d.forward(conv, x)
+    assert not torch.equal(before, after)
+    torch.testing.assert_close(after, expected, rtol=2e-2, atol=2e-2)
+
+
 # --------------------------------------------------------------------------- #
 # Whole decoder
 # --------------------------------------------------------------------------- #

tests/diffusion/layers/test_adalayernorm_compile_graph.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compile and CUDA-graph smoke tests for the fused AdaLayerNorm fast path.
+
+The serving stack reaches AdaLayerNorm.forward_cuda through CustomOp dispatch
+from compiled diffusion workloads (lazy regional torch.compile on repeated
+blocks) and can capture CUDA graphs around pipeline stages. These smokes
+verify that torch.compile uses the native chain and leaves the eager fused
+path eligible, and that the fused path can be captured/replayed inside a CUDA
+graph with outputs matching native. They are smoke tests, not performance
+measurements.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.adalayernorm import (
+    _FAILED_ADALN_KEYS,
+    AdaLayerNorm,
+    _adaln_fused_forward,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+def _make(bs=2, seq=512, hidden=3072, dtype=torch.bfloat16, seed=0, framewise=False):
+    if framewise:
+        hidden = 2240
+    m = AdaLayerNorm(hidden).to(device="cuda", dtype=dtype)
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    if framewise:
+        x = torch.randn(bs, 4, seq // 4, hidden, generator=g, device="cuda", dtype=dtype)
+        modulation = torch.randn(bs, 4, 6, hidden, generator=g, device="cuda", dtype=dtype)
+        scale, shift = modulation.chunk(6, dim=2)[:2]
+        return m, x, scale, shift
+    x = torch.randn(bs, seq, hidden, generator=g, device="cuda", dtype=dtype)
+    scale = torch.randn(bs, 1, hidden, generator=g, device="cuda", dtype=dtype)
+    shift = torch.randn(bs, 1, hidden, generator=g, device="cuda", dtype=dtype)
+    return m, x, scale, shift
+
+
+@pytest.mark.parametrize("framewise", [False, True])
+def test_compile_smoke(framewise):
+    # torch.compile over the guarded forward_cuda: with the is_compiling
+    # guard the compiled region takes the NATIVE chain, so this compares
+    # inductor's own LN+modulation fusion against the eager 4-kernel chain.
+    # Cross-implementation BF16 outputs differ at rounding level (up to a few
+    # ulps at O(1-8) magnitudes), hence the loose tolerance. The smoke
+    # verifies the compiled workload stays functional - no crash, finite
+    # outputs, no request failure.
+    m, x, scale, shift = _make(framewise=framewise)
+    compiled = torch.compile(m.forward_cuda, dynamic=False, fullgraph=True)
+    out1 = compiled(x, scale, shift)
+    out2 = compiled(x, scale, shift)
+    native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out1.float(), native.float(), atol=5e-2, rtol=2e-2)
+    torch.testing.assert_close(out2.float(), native.float(), atol=5e-2, rtol=2e-2)
+    runtime_key = (
+        x.device.index,
+        x.shape[-1],
+        x.dtype,
+        m.layernorm.weight is not None,
+        m.layernorm.bias is not None,
+        True,
+        True,
+    )
+    assert runtime_key not in _FAILED_ADALN_KEYS, "compile must not cache an eager launch failure"
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "compile must leave the eager fused path eligible"
+    torch.testing.assert_close(fused.float(), native.float(), atol=2e-2, rtol=2e-2)
+    assert runtime_key not in _FAILED_ADALN_KEYS
+
+
+@pytest.mark.parametrize("framewise", [False, True])
+def test_cuda_graph_capture_replay_smoke(framewise):
+    # CUDA graph capture must prove the FUSED path itself is captured: call
+    # _adaln_fused_forward inside the capture and assert it returned a tensor
+    # (None would mean the guard/fallback silently degraded the graph to the
+    # native chain). Warm the JIT on a side stream so capture contains no
+    # compile-time allocations, then replay and compare against the native
+    # fallback on the same static input buffers.
+    m, x, scale, shift = _make(framewise=framewise)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(3):
+            fused_warm = _adaln_fused_forward(m, x, scale, shift)
+            assert fused_warm is not None
+    torch.cuda.current_stream().wait_stream(s)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out_graph = _adaln_fused_forward(m, x, scale, shift)
+    assert out_graph is not None
+    graph.replay()
+    torch.accelerator.synchronize()
+    native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out_graph.float(), native.float(), atol=2e-2, rtol=2e-2)
+
+    # Replay must consume NEW input values: mutate the captured buffers and
+    # verify the replay output tracks them (a no-op replay would pass the
+    # check above trivially).
+    x2, scale2, shift2 = torch.randn_like(x), torch.randn_like(scale), torch.randn_like(shift)
+    x.copy_(x2)
+    scale.copy_(scale2)
+    shift.copy_(shift2)
+    graph.replay()
+    torch.accelerator.synchronize()
+    native2 = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out_graph.float(), native2.float(), atol=2e-2, rtol=2e-2)
+    assert not torch.equal(native, native2), "replay must consume the new input"

tests/diffusion/layers/test_adalayernorm_cuda_path.py

@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Candidate-added tests for the fused Triton AdaLayerNorm CUDA path.
+
+Complements (never replaces) the frozen suite: exercises the fused fast path
+and the forward_native fallback branches of forward_cuda. Uses the same
+tolerance scheme as the frozen suite; golden = forward_native real behavior.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.adalayernorm import (
+    AdaLayerNorm,
+    _adaln_fused_forward,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+TOL_STRICT = {"bf16": (2e-2, 2e-2), "fp32": (1e-3, 1e-3)}
+# Matches the frozen suite: the double-rounded golden chain deviates from a
+# pure-fp32 reference more than the strict band, so the fp32-reference check
+# uses the same loose band as the frozen tests.
+TOL_LOOSE = {"bf16": (5e-2, 2e-2), "fp32": (5e-3, 5e-3)}
+DTYPES = [torch.bfloat16, torch.float32]
+
+
+def fp32_reference(x, scale, shift, eps):
+    xf = x.float()
+    xn = torch.nn.functional.layer_norm(xf, (x.shape[-1],), None, None, eps)
+    return xn * (1 + scale.float().reshape(1, 1, -1)) + shift.float().reshape(1, 1, -1)
+
+
+def make_module(hidden, affine, eps, device, dtype):
+    m = AdaLayerNorm(hidden, elementwise_affine=affine, eps=eps)
+    if affine:
+        m = m.to(dtype)
+    return m.to(device)
+
+
+def assert_close(a, b, dtype, loose=False):
+    atol, rtol = (TOL_LOOSE if loose else TOL_STRICT)["bf16" if dtype == torch.bfloat16 else "fp32"]
+    torch.testing.assert_close(a.float(), b.float(), atol=atol, rtol=rtol)
+
+
+def make_inputs(bs, seq, hidden, dtype, device, seed=0):
+    g = torch.Generator(device=device).manual_seed(seed)
+    x = torch.randn(bs, seq, hidden, generator=g, device=device, dtype=dtype)
+    scale = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+    return x, scale, shift
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("affine", [False, True])
+@pytest.mark.parametrize("bs,seq,hidden", [(1, 512, 3072), (2, 128, 1536), (1, 3, 1000)])
+@torch.no_grad()
+def test_fused_fast_path_matches_native(dtype, affine, bs, seq, hidden):
+    # hidden=1000 exercises BLOCK_C masking (next_pow2(1000)=1024 > 1000)
+    device = "cuda"
+    m = make_module(hidden, affine, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(bs, seq, hidden, dtype, device)
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "supported inputs must take the fused path so BLOCK_C masking is exercised"
+    out = fused
+    assert out.shape == x.shape and out.dtype == dtype and out.device == x.device
+    assert_close(out, m.forward_native(x, scale, shift), dtype)
+    assert_close(out, fp32_reference(x, scale, shift, 1e-6), dtype, loose=True)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_fused_noncontiguous_x(dtype):
+    # Non-contiguous x routes to the forward_native fallback (the fused kernel
+    # requires contiguous x); result must still match native exactly.
+    device = "cuda"
+    hidden = 1536
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    g = torch.Generator(device=device).manual_seed(3)
+    x_big = torch.randn(1, 1024, hidden, generator=g, device=device, dtype=dtype)
+    x = x_big[:, ::2, :]
+    scale = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+    out = m.forward_cuda(x, scale, shift)
+    assert_close(out, m.forward_native(x, scale, shift), dtype)
+    assert_close(out, fp32_reference(x.contiguous(), scale, shift, 1e-6), dtype, loose=True)
+
+
+def test_normalized_shape_mismatch_preserves_native_error():
+    m = AdaLayerNorm(16, elementwise_affine=False).cuda()
+    x = torch.randn(1, 2, 32, device="cuda")
+    scale = torch.zeros(32, device="cuda")
+    shift = torch.zeros_like(scale)
+    assert _adaln_fused_forward(m, x, scale, shift) is None
+    with pytest.raises(RuntimeError) as native_error:
+        m.forward_native(x, scale, shift)
+    with pytest.raises(RuntimeError) as cuda_error:
+        m.forward_cuda(x, scale, shift)
+    assert str(cuda_error.value) == str(native_error.value)
+
+
+@pytest.mark.parametrize("grad_target", ["x", "scale", "shift", "weight", "bias"])
+def test_grad_fallback_matches_native_backward(grad_target):
+    m = make_module(32, grad_target in ("weight", "bias"), 1e-6, "cuda", torch.float32)
+    m.requires_grad_(False)
+    x, scale, shift = make_inputs(1, 2, 32, torch.float32, "cuda", seed=12)
+    targets = {"x": x, "scale": scale, "shift": shift, "weight": m.layernorm.weight, "bias": m.layernorm.bias}
+    target = targets[grad_target]
+    target.requires_grad_(True)
+    with torch.enable_grad():
+        assert _adaln_fused_forward(m, x, scale, shift) is None
+        native = m.forward_native(x, scale, shift)
+        out = m.forward_cuda(x, scale, shift)
+        assert out.requires_grad
+        torch.testing.assert_close(out, native)
+        upstream = torch.randn_like(out)
+        native_grad = torch.autograd.grad(native, target, upstream)[0]
+        cuda_grad = torch.autograd.grad(out, target, upstream)[0]
+        torch.testing.assert_close(cuda_grad, native_grad)
+
+
+def test_no_grad_keeps_fused_path_for_grad_requiring_tensors():
+    m = make_module(32, True, 1e-6, "cuda", torch.float32)
+    x, scale, shift = make_inputs(1, 2, 32, torch.float32, "cuda", seed=14)
+    for tensor in (x, scale, shift):
+        tensor.requires_grad_(True)
+    with torch.no_grad():
+        fused = _adaln_fused_forward(m, x, scale, shift)
+        assert fused is not None
+        assert not fused.requires_grad
+        torch.testing.assert_close(fused, m.forward_native(x, scale, shift), atol=1e-3, rtol=1e-3)
+
+
+def test_1d_scale_shift_fast_path():
+    # (C,) scale/shift broadcast per-channel exactly like (1, C)
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, torch.bfloat16)
+    g = torch.Generator(device=device).manual_seed(10)
+    x = torch.randn(1, 256, hidden, generator=g, device=device, dtype=torch.bfloat16)
+    scale = torch.randn(hidden, generator=g, device=device, dtype=torch.bfloat16)
+    shift = torch.randn(hidden, generator=g, device=device, dtype=torch.bfloat16)
+    out = _adaln_fused_forward(m, x, scale, shift)
+    assert out is not None, "(C,) scale/shift must take the fused path"
+    assert_close(out, m.forward_native(x, scale, shift), torch.bfloat16)
+    assert_close(out, fp32_reference(x, scale, shift, 1e-6), torch.bfloat16, loose=True)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_per_row_scale_falls_back_to_native(dtype):
+    # (L, C) scale with B=1 is a valid per-row torch broadcast; it is not
+    # per-channel, so forward_cuda must route it to forward_native.
+    device = "cuda"
+    hidden = 512
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    g = torch.Generator(device=device).manual_seed(4)
+    x = torch.randn(1, 64, hidden, generator=g, device=device, dtype=dtype)
+    scale = torch.randn(64, hidden, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(64, hidden, generator=g, device=device, dtype=dtype)
+    out = m.forward_cuda(x, scale, shift)
+    ref = m.forward_native(x, scale, shift)
+    assert out.dtype == ref.dtype
+    assert_close(out, ref, dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_mixed_dtype_falls_back_to_native(dtype):
+    # scale/shift dtype != x.dtype: native type-promotes; fused path requires
+    # matching dtypes, so this must go through forward_native.
+    other = torch.float32 if dtype == torch.bfloat16 else torch.bfloat16
+    device = "cuda"
+    hidden = 256
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    g = torch.Generator(device=device).manual_seed(5)
+    x = torch.randn(1, 32, hidden, generator=g, device=device, dtype=dtype)
+    scale = torch.randn(1, hidden, generator=g, device=device, dtype=other)
+    shift = torch.randn(1, hidden, generator=g, device=device, dtype=other)
+    out = m.forward_cuda(x, scale, shift)
+    ref = m.forward_native(x, scale, shift)
+    assert out.dtype == ref.dtype
+    torch.testing.assert_close(out, ref)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_zero_dim_scale_falls_back_to_native(dtype):
+    # 0-d scalar modulation is a valid torch broadcast; not per-channel.
+    device = "cuda"
+    hidden = 256
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    g = torch.Generator(device=device).manual_seed(6)
+    x = torch.randn(1, 32, hidden, generator=g, device=device, dtype=dtype)
+    scale = torch.tensor(0.5, device=device, dtype=dtype)
+    shift = torch.tensor(-0.25, device=device, dtype=dtype)
+    out = m.forward_cuda(x, scale, shift)
+    assert_close(out, m.forward_native(x, scale, shift), dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_per_sample_scale_raises_like_native(dtype):
+    # Frozen behavior: B>1 with (B, C) scale/shift -> RuntimeError.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    g = torch.Generator(device=device).manual_seed(7)
+    x = torch.randn(4, 64, hidden, generator=g, device=device, dtype=dtype)
+    scale = torch.randn(4, hidden, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(4, hidden, generator=g, device=device, dtype=dtype)
+    with pytest.raises(RuntimeError):
+        m.forward_native(x, scale, shift)
+    with pytest.raises(RuntimeError):
+        m.forward_cuda(x, scale, shift)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_fused_determinism_repeat_calls(dtype):
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(1, 1024, hidden, dtype, device, seed=8)
+    out1 = _adaln_fused_forward(m, x, scale, shift)
+    out2 = _adaln_fused_forward(m, x, scale, shift)
+    assert out1 is not None and out2 is not None, "determinism must exercise the fused path"
+    assert torch.equal(out1, out2)
+
+
+@pytest.mark.parametrize("eps", [1e-5, 1e-3])
+def test_fused_eps_variants(eps):
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, eps, device, torch.bfloat16)
+    x, scale, shift = make_inputs(1, 256, hidden, torch.bfloat16, device, seed=9)
+    out = _adaln_fused_forward(m, x, scale, shift)
+    assert out is not None, "supported eps variants must take the fused path"
+    assert_close(out, m.forward_native(x, scale, shift), torch.bfloat16)
+    assert_close(out, fp32_reference(x, scale, shift, eps), torch.bfloat16, loose=True)

tests/diffusion/layers/test_adalayernorm_launch_fallback.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Direct tests for the synchronous launch-failure fallback (C3a).
+
+Simulates a backend-specific Triton launch failure (the launcher raises
+synchronously) and verifies the two-layer response: the failed config is
+recorded, the caller gets the correct native result, the same config is not
+retried, and a different constexpr variant stays eligible.
+
+The failed-key set is module-global; the autouse fixture snapshots and
+restores it so tests are self-contained and independent of execution order.
+Collection is safe on non-Triton environments: the Triton-only kernel symbol
+is resolved lazily inside each test and skips if unavailable.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers import adalayernorm as adaln_mod
+from vllm_omni.diffusion.layers.adalayernorm import (
+    _FAILED_ADALN_KEYS,
+    AdaLayerNorm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+# Unique channels + eps so the failed keys used here cannot collide with
+# other tests in the same session.
+_HIDDEN = 2048
+_EPS = 1e-5
+
+
+@pytest.fixture(autouse=True)
+def _isolate_failed_keys():
+    saved = set(_FAILED_ADALN_KEYS)
+    _FAILED_ADALN_KEYS.clear()
+    try:
+        yield
+    finally:
+        _FAILED_ADALN_KEYS.clear()
+        _FAILED_ADALN_KEYS.update(saved)
+
+
+def _make(bs=1):
+    m = AdaLayerNorm(_HIDDEN, eps=_EPS).to(device="cuda", dtype=torch.bfloat16)
+    g = torch.Generator(device="cuda").manual_seed(31)
+    x = torch.randn(bs, 512, _HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, _HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn(1, _HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16)
+    return m, x, scale, shift
+
+
+def _get_kernel():
+    kernel = getattr(adaln_mod, "_adaln_scale_shift_layernorm_kernel", None)
+    if kernel is None:
+        pytest.skip("Triton AdaLayerNorm kernel unavailable")
+    return kernel
+
+
+def test_shared_launch_failure_fallback_and_cache(monkeypatch):
+    kernel = _get_kernel()
+    m, x, scale, shift = _make(1)
+    native = m.forward_native(x, scale, shift)
+
+    calls = {"n": 0}
+
+    def raising_run(self, *args, **kwargs):
+        calls["n"] += 1
+        raise RuntimeError("simulated synchronous launch failure")
+
+    monkeypatch.setattr(type(kernel), "run", raising_run)
+
+    # Phase 1: the launch raises synchronously; the caller must receive the
+    # correct native result and the failed key must be recorded.
+    out = m.forward_cuda(x, scale, shift)
+    torch.testing.assert_close(out.float(), native.float())
+    expected_key = (x.device.index, _HIDDEN, x.dtype, False, False, False, False)
+    assert expected_key in _FAILED_ADALN_KEYS
+    assert calls["n"] == 1  # the fused launcher was attempted exactly once
+
+    # Phase 2: the same config must NOT retry the launcher (failed-key
+    # cache) and still return the correct native result.
+    calls["n"] = 0
+    out2 = m.forward_cuda(x, scale, shift)
+    assert calls["n"] == 0
+    torch.testing.assert_close(out2.float(), native.float())
+
+
+def test_shared_failure_then_per_sample_variant_eligible(monkeypatch):
+    # The failed-key cache is per constexpr variant: after the shared
+    # (False, False) variant fails, the per-sample (True, True) variant must
+    # still attempt the launcher and succeed.
+    kernel = _get_kernel()
+    m, x, shared_scale, shared_shift = _make(2)
+    shared_native = m.forward_native(x, shared_scale, shared_shift)
+    shared_key = (x.device.index, _HIDDEN, x.dtype, False, False, False, False)
+    per_sample_key = (x.device.index, _HIDDEN, x.dtype, False, False, True, True)
+    calls = {"n": 0}
+
+    def raising_run(self, *args, **kwargs):
+        calls["n"] += 1
+        raise RuntimeError("simulated synchronous shared launch failure")
+
+    # Establish the shared failure in this test; the autouse fixture clears
+    # failures from other tests, so their cache entries cannot prove isolation.
+    with monkeypatch.context() as shared_patch:
+        shared_patch.setattr(type(kernel), "run", raising_run)
+        shared_out = m.forward_cuda(x, shared_scale, shared_shift)
+        torch.testing.assert_close(shared_out.float(), shared_native.float())
+        assert calls["n"] == 1
+        assert shared_key in _FAILED_ADALN_KEYS
+        assert per_sample_key not in _FAILED_ADALN_KEYS
+
+        m.forward_cuda(x, shared_scale, shared_shift)
+        assert calls["n"] == 1, "the cached shared failure must not retry"
+
+    g = torch.Generator(device="cuda").manual_seed(37)
+    scale = torch.randn(2, 1, _HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn(2, 1, _HIDDEN, generator=g, device="cuda", dtype=torch.bfloat16)
+    calls["n"] = 0
+
+    def counting_run(self, *args, **kwargs):
+        calls["n"] += 1
+        return real_run(self, *args, **kwargs)
+
+    real_run = type(kernel).run
+    monkeypatch.setattr(type(kernel), "run", counting_run)
+
+    out = adaln_mod._adaln_fused_forward(m, x, scale, shift)
+    assert out is not None, "the per-sample variant must remain eligible after a shared failure"
+    native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out.float(), native.float(), atol=2e-2, rtol=2e-2)
+    assert calls["n"] == 1
+    assert shared_key in _FAILED_ADALN_KEYS
+    assert per_sample_key not in _FAILED_ADALN_KEYS

tests/diffusion/layers/test_adalayernorm_matrix.py

@@ -0,0 +1,462 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""P0 correctness matrix for the fused AdaLayerNorm CUDA path.
+
+Covers: fp32/bf16/fp16, B>1 (shared and per-sample modulation), C in
+{1536, 3072, 4096}, varying L, elementwise_affine False/True with
+NON-DEFAULT weight/bias, eps variants, legal non-contiguous inputs,
+broadcast scale/shift shapes ((C,), (1, C), (1, 1, C), (B, 1, C)),
+fallback completeness (fp64, oversized hidden size, cross-device
+parameters), large-offset/small-variance stability, and the frozen
+B > 1 (B, C) native broadcast-error contract.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.adalayernorm import (
+    AdaLayerNorm,
+    _adaln_fused_forward,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+TOL_STRICT = {"bf16": (2e-2, 2e-2), "fp16": (1e-2, 1e-2), "fp32": (1e-3, 1e-3)}
+TOL_LOOSE = {"bf16": (5e-2, 2e-2), "fp16": (2e-2, 1e-2), "fp32": (5e-3, 5e-3)}
+DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+
+
+def tol_key(dtype):
+    if dtype == torch.bfloat16:
+        return "bf16"
+    if dtype == torch.float16:
+        return "fp16"
+    return "fp32"
+
+
+def fp32_reference(x, scale, shift, eps, weight=None, bias=None):
+    xf = x.float()
+    w = weight.float() if weight is not None else None
+    b = bias.float() if bias is not None else None
+    xn = torch.nn.functional.layer_norm(xf, (x.shape[-1],), w, b, eps)
+    return xn * (1 + scale.float()[:, None]) + shift.float()[:, None]
+
+
+def make_module(hidden, elementwise_affine, eps, device, dtype, nondefault_affine=False):
+    m = AdaLayerNorm(hidden, elementwise_affine=elementwise_affine, eps=eps)
+    if elementwise_affine:
+        m = m.to(dtype)
+        if nondefault_affine:
+            with torch.no_grad():
+                g = torch.Generator(device="cpu").manual_seed(42)
+                m.layernorm.weight.copy_(0.5 + torch.rand(hidden, generator=g))
+                m.layernorm.bias.copy_(0.2 * torch.randn(hidden, generator=g))
+    return m.to(device)
+
+
+def make_inputs(bs, seq, hidden, dtype, device, seed=0, mod_shape=None):
+    g = torch.Generator(device=device).manual_seed(seed)
+    x = torch.randn(bs, seq, hidden, generator=g, device=device, dtype=dtype)
+    ms = mod_shape or (1, hidden)
+    scale = torch.randn(ms, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(ms, generator=g, device=device, dtype=dtype)
+    return x, scale, shift
+
+
+def assert_close(a, b, dtype, loose=False):
+    atol, rtol = (TOL_LOOSE if loose else TOL_STRICT)[tol_key(dtype)]
+    torch.testing.assert_close(a.float(), b.float(), atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("chunks", [2, 6])
+@pytest.mark.parametrize("bs,frames,spatial", [(1, 21, 880), (2, 3, 5), (2, 1, 5), (1, 3, 5)])
+def test_matrix_framewise_modulation(dtype, chunks, bs, frames, spatial):
+    hidden = 2240
+    m = make_module(hidden, False, 1e-6, "cuda", dtype)
+    g = torch.Generator(device="cuda").manual_seed(61)
+    x = torch.randn(bs, frames, spatial, hidden, generator=g, device="cuda", dtype=dtype)
+    modulation = torch.randn(bs, frames, chunks, hidden, generator=g, device="cuda", dtype=dtype)
+    scale, shift = modulation.chunk(chunks, dim=2)[:2]
+    # Sana's block/final modulation uses chunk views spanning frames.
+    assert scale.shape == (bs, frames, 1, hidden)
+    assert scale.stride(1) == chunks * hidden
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "framewise chunk views must take the fused path"
+    assert fused.shape == x.shape and fused.dtype == dtype
+    assert_close(fused, m.forward_native(x, scale, shift), dtype)
+    ref = torch.nn.functional.layer_norm(x.float(), (hidden,), eps=m.eps)
+    ref = ref * (1 + scale.float()) + shift.float()
+    assert_close(fused, ref, dtype, loose=True)
+
+
+@pytest.mark.parametrize("shared_scale,shared_shift", [(False, True), (True, False), (True, True)])
+@torch.no_grad()
+def test_matrix_framewise_shared_modulation(shared_scale, shared_shift):
+    dtype = torch.bfloat16
+    hidden = 2240
+    m = make_module(hidden, True, 1e-6, "cuda", dtype, nondefault_affine=True)
+    x = torch.randn(2, 3, 5, hidden, device="cuda", dtype=dtype)
+    scale = torch.randn((hidden,) if shared_scale else (2, 3, 1, hidden), device="cuda", dtype=dtype)
+    shift = torch.randn((1, 1, 1, hidden) if shared_shift else (2, 3, 1, hidden), device="cuda", dtype=dtype)
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None
+    assert_close(fused, m.forward_native(x, scale, shift), dtype)
+
+
+def test_matrix_framewise_irregular_group_stride_falls_back():
+    dtype = torch.bfloat16
+    hidden = 2240
+    m = make_module(hidden, False, 1e-6, "cuda", dtype)
+    x = torch.randn(2, 3, 5, hidden, device="cuda", dtype=dtype)
+    modulation = torch.randn(3, 2, 2, hidden, device="cuda", dtype=dtype).transpose(0, 1)
+    scale, shift = modulation.chunk(2, dim=2)
+    assert scale.stride(0) != x.shape[1] * scale.stride(1)
+    assert _adaln_fused_forward(m, x, scale, shift) is None
+    assert_close(m.forward_cuda(x, scale, shift), m.forward_native(x, scale, shift), dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_matrix_sana_patch_embed_reaches_framewise_fused_path(dtype):
+    from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmPatchEmbedMS3D
+
+    hidden = 2240
+    patch = SanaWmPatchEmbedMS3D((1, 1, 1), 8, hidden).to(device="cuda", dtype=dtype)
+    latents = torch.randn(2, 8, 3, 2, 3, device="cuda", dtype=dtype)
+    with torch.no_grad():
+        projected = patch.proj(latents)
+        tokens, (frames, height, width) = patch.project_with_shape(latents)
+    torch.testing.assert_close(tokens, projected.flatten(2).transpose(1, 2), atol=0, rtol=0)
+    assert tokens.is_contiguous()
+    x = tokens.reshape(2, frames, height * width, hidden)
+    modulation = torch.randn(2, frames, 6, hidden, device="cuda", dtype=dtype)
+    scale, shift = modulation.chunk(6, dim=2)[:2]
+    m = make_module(hidden, False, 1e-6, "cuda", dtype)
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "real Sana patch features must reach the fused path"
+    assert_close(fused, m.forward_native(x, scale, shift), dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("affine", [False, True])
+@pytest.mark.parametrize(
+    "bs,seq,hidden", [(1, 4096, 3072), (2, 1024, 1536), (3, 128, 4096), (1, 1, 3072), (1, 8192, 1536)]
+)
+@torch.no_grad()
+def test_matrix_main(dtype, affine, bs, seq, hidden):
+    device = "cuda"
+    m = make_module(hidden, affine, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(bs, seq, hidden, dtype, device)
+    out_cuda = _adaln_fused_forward(m, x, scale, shift)
+    assert out_cuda is not None, "supported matrix inputs must take the fused path"
+    out_native = m.forward_native(x, scale, shift)
+    assert out_cuda.shape == x.shape and out_cuda.dtype == dtype and out_cuda.device == x.device
+    assert_close(out_cuda, out_native, dtype, loose=False)
+    w = m.layernorm.weight if affine else None
+    b = m.layernorm.bias if affine else None
+    assert_close(out_cuda, fp32_reference(x, scale, shift, 1e-6, w, b), dtype, loose=True)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@torch.no_grad()
+def test_matrix_weight_bias_nondefault(dtype):
+    # weight/bias semantics with non-identity initialization:
+    # out = (w * ln(x) + b) * (1 + scale) + shift
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, True, 1e-6, device, dtype, nondefault_affine=True)
+    x, scale, shift = make_inputs(1, 1024, hidden, dtype, device, seed=13)
+    out_cuda = _adaln_fused_forward(m, x, scale, shift)
+    assert out_cuda is not None, "non-default affine parameters must take the fused path"
+    out_native = m.forward_native(x, scale, shift)
+    assert_close(out_cuda, out_native, dtype)
+    ref = fp32_reference(x, scale, shift, 1e-6, m.layernorm.weight, m.layernorm.bias)
+    assert_close(out_cuda, ref, dtype, loose=True)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("eps", [1e-5, 1e-3, 1e-2])
+def test_matrix_eps_variants(dtype, eps):
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, eps, device, dtype)
+    x, scale, shift = make_inputs(1, 512, hidden, dtype, device, seed=3)
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "supported eps variants must take the fused path"
+    assert_close(fused, m.forward_native(x, scale, shift), dtype)
+    assert_close(
+        fused,
+        fp32_reference(x, scale, shift, eps),
+        dtype,
+        loose=True,
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("mod_shape", ["C", "1x1xC", "1xC"])
+def test_matrix_broadcast_scale_shift(dtype, mod_shape):
+    # (C,), (1, 1, C) and (1, C) are all legal per-channel broadcasts against x (B, L, C)
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, _, _ = make_inputs(2, 512, hidden, dtype, device, seed=5)
+    g = torch.Generator(device=device).manual_seed(6)
+    if mod_shape == "C":
+        scale = torch.randn(hidden, generator=g, device=device, dtype=dtype)
+        shift = torch.randn(hidden, generator=g, device=device, dtype=dtype)
+    elif mod_shape == "1x1xC":
+        scale = torch.randn(1, 1, hidden, generator=g, device=device, dtype=dtype)
+        shift = torch.randn(1, 1, hidden, generator=g, device=device, dtype=dtype)
+    else:
+        scale = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+        shift = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
+    out_cuda = _adaln_fused_forward(m, x, scale, shift)
+    assert out_cuda is not None, f"shared modulation layout {mod_shape} must take the fused path"
+    out_native = m.forward_native(x, scale, shift)
+    assert_close(out_cuda, out_native, dtype)
+    s = scale.float().reshape(1, 1, hidden) if scale.ndim == 1 else scale.float()
+    sh = shift.float().reshape(1, 1, hidden) if shift.ndim == 1 else shift.float()
+    ref = fp32_reference(x, s.reshape(-1, hidden)[0:1], sh.reshape(-1, hidden)[0:1], 1e-6)
+    assert_close(out_cuda, ref, dtype, loose=True)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_per_sample_modulation(dtype):
+    # Qwen-Image's _modulate produces (B, 1, C): one modulation row per sample.
+    # B > 1 must take the fused path (not fall back to native) and match native.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, _, _ = make_inputs(4, 512, hidden, dtype, device, seed=7, mod_shape=(4, hidden))
+    g = torch.Generator(device=device).manual_seed(8)
+    scale = torch.randn(4, 1, hidden, generator=g, device=device, dtype=dtype)
+    shift = torch.randn(4, 1, hidden, generator=g, device=device, dtype=dtype)
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "(B, 1, C) per-sample modulation must take the fused path"
+    out_cuda = fused
+    out_native = m.forward_native(x, scale, shift)
+    assert_close(out_cuda, out_native, dtype)
+    # fp32_reference's [:, None] expects a 2D (B, C) modulation row; the
+    # (B, 1, C) tensor reshapes to it losslessly.
+    ref = fp32_reference(x, scale.reshape(x.shape[0], hidden), shift.reshape(x.shape[0], hidden), 1e-6)
+    assert_close(out_cuda, ref, dtype, loose=True)
+    # Per-sample: different rows must receive different modulation.
+    row_diff = (out_cuda[0] - out_cuda[1]).abs().max().item()
+    assert row_diff > 0
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_chunk_view_modulation(dtype):
+    # Wan2.2-style producer: the modulation projection is (B, 6, C) and the
+    # consumer chunks it along dim=1. The resulting (B, 1, C) views are
+    # NON-CONTIGUOUS (row stride 6*C) - the fused path must consume them
+    # directly via the explicit row stride, without a contiguous copy.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, _, _ = make_inputs(4, 512, hidden, dtype, device, seed=27)
+    g = torch.Generator(device=device).manual_seed(28)
+    source = torch.randn(4, 6, hidden, generator=g, device=device, dtype=dtype) * 0.1
+    chunks = source.chunk(6, dim=1)
+    scale, shift = chunks[1], chunks[2]
+    assert scale.shape == (4, 1, hidden) and shift.shape == (4, 1, hidden)
+    assert not scale.is_contiguous()
+    assert scale.stride(-1) == 1
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is not None, "chunk-view modulation must take the fused path"
+    out_native = m.forward_native(x, scale, shift)
+    assert_close(fused, out_native, dtype)
+    ref = fp32_reference(x, scale.reshape(4, hidden), shift.reshape(4, hidden), 1e-6)
+    assert_close(fused, ref, dtype, loose=True)
+
+
+def test_matrix_3d_modulation_fallback():
+    # Wan2.2 TI2V-style per-token modulation: scale/shift (B, L, C) full 3D.
+    # This is intentionally outside the fused kernel contract - the fast path
+    # must be rejected (None) and forward_cuda must match native.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, torch.bfloat16)
+    x, scale, shift = make_inputs(2, 128, hidden, torch.bfloat16, device, seed=29, mod_shape=(2, 128, hidden))
+    fused = _adaln_fused_forward(m, x, scale, shift)
+    assert fused is None, "(B, L, C) per-token modulation must be rejected"
+    out = m.forward_cuda(x, scale, shift)
+    native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out.float(), native.float(), atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+def test_matrix_production_widths(dtype):
+    # Real consumer widths: Sana-WM 2240 (BLOCK_C 4096, 45.3% mask) and
+    # Wan2.2 A14B 5120 (BLOCK_C 8192, the max supported block, 16-warp
+    # compiled variant). BF16 + elementwise_affine=False.
+    device = "cuda"
+    for hidden in (2240, 5120):
+        m = make_module(hidden, False, 1e-6, device, dtype)
+        x, scale, shift = make_inputs(1, 512, hidden, dtype, device, seed=31)
+        fused = _adaln_fused_forward(m, x, scale, shift)
+        assert fused is not None, f"hidden={hidden} must take the fused path"
+        out_native = m.forward_native(x, scale, shift)
+        assert_close(fused, out_native, dtype)
+        assert_close(fused, fp32_reference(x, scale, shift, 1e-6), dtype, loose=True)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_noncontiguous_fallback(dtype):
+    # Non-contiguous inputs must fall back to native and stay correct
+    # (the fused kernel only serves the contiguous fast path).
+    device = "cuda"
+    hidden = 1536
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x_big, scale, shift = make_inputs(1, 1024, hidden, dtype, device, seed=7)
+    x = x_big[:, ::2, :]
+    assert not x.is_contiguous()
+    out_cuda = m.forward_cuda(x, scale, shift)
+    out_native = m.forward_native(x, scale, shift)
+    assert_close(out_cuda, out_native, dtype)
+    assert_close(out_cuda, fp32_reference(x.contiguous(), scale, shift, 1e-6), dtype, loose=True)
+
+
+def test_matrix_fallback_fp64():
+    # fp64 is not in the kernel's supported dtype list -> must fall back to
+    # native and stay correct (fallback completeness).
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, torch.float64)
+    x, scale, shift = make_inputs(1, 256, hidden, torch.float64, device, seed=9)
+    out_cuda = m.forward_cuda(x, scale, shift)
+    out_native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out_cuda, out_native)
+
+
+def test_matrix_fallback_oversized_hidden():
+    # hidden sizes whose next_power_of_2 exceeds the kernel's supported block
+    # bound must fall back to native and stay correct.
+    device = "cuda"
+    hidden = 10000  # next_power_of_2 = 16384 > MAX_BLOCK_C
+    m = make_module(hidden, False, 1e-6, device, torch.bfloat16)
+    x, scale, shift = make_inputs(1, 128, hidden, torch.bfloat16, device, seed=10)
+    out_cuda = m.forward_cuda(x, scale, shift)
+    out_native = m.forward_native(x, scale, shift)
+    torch.testing.assert_close(out_cuda, out_native)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_zero_scale_shift_identity(dtype):
+    # With scale = shift = 0 the output must equal a pure LayerNorm.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, _, _ = make_inputs(1, 256, hidden, dtype, device, seed=19)
+    scale = torch.zeros(1, hidden, device=device, dtype=dtype)
+    shift = torch.zeros(1, hidden, device=device, dtype=dtype)
+    out_cuda = _adaln_fused_forward(m, x, scale, shift)
+    assert out_cuda is not None, "zero modulation must take the fused path"
+    assert_close(out_cuda, m.forward_native(x, scale, shift), dtype)
+    assert_close(out_cuda, m.layernorm(x), dtype)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_determinism(dtype):
+    # Two calls on the same input must agree bitwise (a hard constraint for
+    # any cached/direct-launch scheme).
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(1, 4096, hidden, dtype, device, seed=11)
+    out1 = _adaln_fused_forward(m, x, scale, shift)
+    out2 = _adaln_fused_forward(m, x, scale, shift)
+    assert out1 is not None and out2 is not None, "determinism must exercise the fused path"
+    assert torch.equal(out1, out2)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_multi_batch_shared_modulation(dtype):
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(4, 512, hidden, dtype, device, seed=15)
+    out = _adaln_fused_forward(m, x, scale, shift)
+    assert out is not None, "multi-batch shared modulation must take the fused path"
+    assert_close(out, m.forward_native(x, scale, shift), dtype)
+    for b in range(4):
+        single = _adaln_fused_forward(m, x[b : b + 1], scale, shift)
+        assert single is not None
+        assert_close(out[b : b + 1], single, dtype)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_bc_modulation_preserves_native_error(dtype):
+    # Frozen native contract: (B, C) modulation with B > 1 raises
+    # RuntimeError (torch broadcasts B against L). The fused path preserves
+    # this behavior - it neither supports nor silently changes it.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(4, 64, hidden, dtype, device, seed=17, mod_shape=(4, hidden))
+    with pytest.raises(RuntimeError):
+        m.forward_native(x, scale, shift)
+    with pytest.raises(RuntimeError):
+        m.forward_cuda(x, scale, shift)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_matrix_cross_device_fallback(dtype):
+    # Raw pointers go to one Triton kernel, which does no cross-device
+    # checking: the fast-path guard must reject cross-device parameters so
+    # they fall back to native, which raises torch's own cross-device error
+    # instead of crashing in the kernel.
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(1, 512, hidden, dtype, device, seed=21)
+    cpu_scale = scale.to("cpu")
+    with pytest.raises(RuntimeError):
+        m.forward_cuda(x, cpu_scale, shift)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_matrix_large_offset_small_variance(dtype):
+    """A large constant offset with a small variance must stay stable.
+
+    Mirrors the fused_adaptive_group_norm_silu case: x ~ 10000 +/- 0.1 has
+    mean ~1e4 and variance ~1e-2, so a plain fp32 sum of large-offset values
+    rounds away the low-order bits that the variance is built from. The fused
+    kernel centers the row on its first element before the tree sum
+    (shift-invariant two-pass), keeping mean/variance accurate.
+
+    For fp32 the meaningful reference is an fp64 computation, NOT the native
+    eager chain: at mean/std = 1e5 the native chain itself carries a
+    ~0.07 absolute deviation from the fp64 truth (its own fp32 dynamic-range
+    floor), while the fused kernel stays within ~2e-6. For fp16/bf16 inputs
+    the +/-0.1 perturbation quantizes away at magnitude 1e4, so the
+    meaningful assertion is that outputs stay finite and match native.
+    """
+    device = "cuda"
+    hidden = 3072
+    m = make_module(hidden, False, 1e-6, device, dtype)
+    x, scale, shift = make_inputs(1, 4096, hidden, dtype, device, seed=23)
+    # Apply the large constant offset AFTER generation: x ~ 10000 +/- 0.1 is
+    # the cancellation case that motivates the shift-invariant two-pass.
+    x = x * 0.1 + 10000
+    out_cuda = _adaln_fused_forward(m, x, scale, shift)
+    assert out_cuda is not None, "large-offset stability must exercise the fused path"
+    out_native = m.forward_native(x, scale, shift)
+    assert torch.isfinite(out_cuda).all()
+
+    # fp64 reference computed from the SAME fp32 input the implementations see
+    # (the kernel cannot recover bits that fp32 quantization never stored).
+    x64 = x.double()
+    mean64 = x64.mean(dim=-1, keepdim=True)
+    var64 = x64.var(dim=-1, unbiased=False, keepdim=True)
+    ref64 = (x64 - mean64) * torch.rsqrt(var64 + 1e-6) * (1 + scale.double()[:, None, :]) + shift.double()[:, None, :]
+
+    if dtype is torch.float32:
+        # fp32 fused output vs the fp64 truth of the same input: within the
+        # fp32 representation floor (~2e-6 measured). The native eager chain
+        # deviates ~0.07 on the same input (its own dynamic-range floor) -
+        # documented, not asserted.
+        torch.testing.assert_close(out_cuda.double(), ref64, atol=1e-4, rtol=1e-4)
+    else:
+        # bf16/fp16 quantize the perturbation away: no NaN, matches native.
+        assert_close(out_cuda, out_native, dtype, loose=True)

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -721,7 +721,7 @@ def test_initialize_model_sets_current_diffusion_config_during_model_constructio
 
     od_config = SimpleNamespace(
         model_class_name="DummyPipeline",
-        parallel_config=SimpleNamespace(vae_patch_parallel_size=1, sequence_parallel_size=1),
+        parallel_config=DiffusionParallelConfig(vae_patch_parallel_size=1, sequence_parallel_size=1),
         vae_use_slicing=False,
         vae_use_tiling=False,
     )

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import weakref
 from types import SimpleNamespace
 from typing import Any
 
@@ -824,6 +825,134 @@ def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.Monke
         assert not hasattr(model, name)
 
 
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@pytest.mark.parametrize("seacache", [False, True])
+@pytest.mark.parametrize("control_count", [0, 1, 2])
+@pytest.mark.parallel
+@torch.inference_mode()
+def test_sp_releases_full_gen_embedding_during_stack(
+    monkeypatch, sequence_length, batch, rank, seacache, control_count
+) -> None:
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.cache.seacache import SeaCacheConfig, apply_sea_cache_hook
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+    from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+    from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+    from vllm_omni.diffusion.hooks import sequence_parallel
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    full_refs = []
+    full_pointers = []
+    layer_calls = 0
+    total_length = sequence_length * (control_count + 1)
+    uses_sp = control_count <= 1
+
+    class CheckingLayer(nn.Module):
+        def forward(self, hidden, **kwargs):
+            nonlocal layer_calls
+            layer_calls += 1
+            if full_refs:
+                # Inspect lifetime inside the production stack while the
+                # caller's prepared state is still live.
+                assert full_refs[-1]() is None
+                if layer_calls % 2 == 1:
+                    assert hidden.untyped_storage().data_ptr() != full_pointers[-1]
+                assert hidden.shape == (batch, (total_length + 3) // 4, 8)
+                assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+                assert kwargs["freqs_cos"].shape[1] == kwargs["freqs_sin"].shape[1] == hidden.shape[1]
+            elif not uses_sp:
+                assert hidden.shape == (batch, total_length, 8)
+            return hidden + 1
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.gen_layers = nn.ModuleList([CheckingLayer(), CheckingLayer()])
+    model.cached_kv = [(torch.empty(0), torch.empty(0)) for _ in model.gen_layers]
+    inputs = {
+        "hidden_states": torch.randn(batch, 2, 1, 1, sequence_length),
+        "timestep": torch.ones(batch),
+        "text_ids": torch.ones(batch, 2, dtype=torch.long),
+        "text_mask": torch.ones(batch, 2, dtype=torch.long),
+        "video_shape": (1, 1, sequence_length),
+        "control_latents": [torch.randn(batch, 2, 1, 1, sequence_length) for _ in range(control_count)],
+    }
+    gathered_reference = []
+    normalized_reference = []
+    handle = model.gen_sp_gather.register_forward_pre_hook(
+        lambda module, args: gathered_reference.append(args[0].clone())
+    )
+    norm_handle = model.norm_moe_gen.register_forward_hook(
+        lambda module, args, output: normalized_reference.append(output.clone())
+    )
+    expected = model(**inputs)
+    handle.remove()
+    norm_handle.remove()
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (4, rank, None))
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    def gather(hidden, dim, validate):
+        shard_length = hidden.shape[1]
+        start = rank * shard_length
+        valid_length = min(shard_length, total_length - start)
+        reference = normalized_reference[0] if seacache else gathered_reference[0]
+        torch.testing.assert_close(hidden[:, :valid_length], reference[:, start : start + valid_length])
+        # Supply the other ranks' outputs without retaining the original input.
+        return torch.nn.functional.pad(reference, (0, 0, 0, -total_length % 4))
+
+    monkeypatch.setattr(sequence_parallel, "sp_gather", gather)
+    sequence_parallel.apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), model._sp_plan)
+
+    def record_full_embedding(module, args):
+        full_refs.append(weakref.ref(args[0]))
+        full_pointers.append(args[0].untyped_storage().data_ptr())
+
+    model.gen_sp_prepare.register_forward_pre_hook(record_full_embedding)
+    metadata = SimpleNamespace(step=0)
+    if seacache:
+        sea_hook = apply_sea_cache_hook(
+            model,
+            SeaCacheConfig(threshold=100.0, residual_order=0),
+            current_step_callback=lambda: metadata.step,
+            current_sigma_callback=lambda: 1.0,
+            num_inference_steps_callback=lambda: 4,
+        )
+        sea_hook.state_manager.set_context("cond")
+    layer_calls = 0
+    for step in range(4 if seacache else 2):
+        metadata.step = step
+        with set_forward_context():
+            output = model(**inputs)
+            torch.testing.assert_close(output, expected)
+            assert get_forward_context()._sp_shard_depth == 0
+            if seacache:
+                if uses_sp:
+                    assert full_refs[-1]() is None
+                history = sea_hook.state_manager.get_state().history
+                assert history
+                residual = history[-1][1]
+                assert residual.shape == (batch, (total_length + 3) // 4 if uses_sp else total_length, 8)
+                assert residual.untyped_storage().nbytes() == residual.numel() * residual.element_size()
+    assert layer_calls == 4
+    assert len(full_refs) == ((4 if seacache else 2) if uses_sp else 0)
+    if seacache:
+        assert sea_hook.full_count == sea_hook.skip_count == 2
+
+
 def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -1394,3 +1523,28 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No
     )
     _, offset_gen_pos = rotary.position_ids
     assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]
+
+
+def test_shard_gen_prep_rejects_inputs_already_in_execution_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    prep = model._gen_preprocess(
+        torch.zeros(1, 2, 1, 2, 2),
+        torch.tensor([1.0]),
+        torch.tensor([[1, 2]], dtype=torch.long),
+        torch.ones(1, 2, dtype=torch.long),
+        (1, 2, 2),
+    )
+
+    sharded = model._shard_gen_prep(prep, defer_gather=True)
+    assert sharded.freqs_gen is not None
+    assert sharded.defer_gen_gather
+    with pytest.raises(RuntimeError, match="already in the execution layout"):
+        model._shard_gen_prep(sharded)
+    # A deferred gather without sharding would gather a full-layout tensor.
+    with pytest.raises(RuntimeError, match="_shard_gen_prep"):
+        model._run_gen_stack(prep._replace(defer_gen_gather=True))

tests/diffusion/models/flux2/test_flux2_klein_local_tp.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU regressions for Klein's packed TP layout and local-head execution."""
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+import torch.nn.functional as F
+from vllm.model_executor import parameter
+from vllm.model_executor.layers import linear
+from vllm.model_executor.layers.utils import default_unquantized_gemm
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as klein
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@contextmanager
+def _default_dtype(dtype):
+    previous = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(previous)
+
+
+class _CPUAttention(torch.nn.Module):
+    """Use SDPA while retaining the model's actual Q/K norm and RoPE."""
+
+    def __init__(self, num_heads, head_size, softmax_scale, causal=False, **kwargs):
+        super().__init__()
+        self.num_heads = num_heads
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value, metadata=None):
+        assert query.shape[2] == self.num_heads
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=None if metadata is None else metadata.attn_mask,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+@pytest.fixture(autouse=True)
+def _cpu_dispatch(monkeypatch):
+    monkeypatch.setattr(linear, "dispatch_unquantized_gemm", lambda *a, **kw: default_unquantized_gemm)
+    monkeypatch.setattr(klein, "Attention", _CPUAttention)
+    monkeypatch.setattr(klein, "get_forward_context", ForwardContext)
+
+
+def _set_tp(monkeypatch, tp, rank):
+    # Keep real vLLM parameters, constructors and checkpoint loaders. Only the
+    # distributed process metadata is substituted for this CPU unit test.
+    for module in (linear, parameter):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: tp)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: rank)
+    monkeypatch.setattr(klein, "get_tensor_model_parallel_world_size", lambda: tp)
+
+
+def _attention(monkeypatch, tp, rank, *, local_tp=True, heads=8, dim_head=8, mlp_ratio=3, sp=1):
+    _set_tp(monkeypatch, tp, rank)
+    with _default_dtype(torch.bfloat16):
+        attention = klein.Flux2ParallelSelfAttention(
+            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp, ulysses_degree=sp),
+            query_dim=heads * dim_head,
+            heads=heads,
+            dim_head=dim_head,
+            mlp_ratio=mlp_ratio,
+            bias=True,
+            local_tp=local_tp,
+        )
+    attention.norm_q.forward = attention.norm_q.forward_native
+    attention.norm_k.forward = attention.norm_k.forward_native
+    attention.rope.forward = attention.rope.forward_native
+    return attention
+
+
+def _rank_rows(tensor, sizes, tp, rank):
+    return torch.cat([part.chunk(tp, dim=0)[rank] for part in tensor.split(sizes, dim=0)])
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+@pytest.mark.parametrize("name", ["weight", "bias"])
+def test_packed_checkpoint_reconstruction_and_reload(monkeypatch, tp, name):
+    sizes = [64, 64, 64, 192, 192]
+    shape = (sum(sizes), 64) if name == "weight" else (sum(sizes),)
+    canonical = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape)
+    shards = []
+    for rank in range(tp):
+        attention = _attention(monkeypatch, tp, rank)
+        assert attention.local_tp == (tp > 1)
+        assert attention.query_num_heads == 8 // tp
+        projection = attention.to_qkv_mlp_proj
+        param = getattr(projection, name)
+        # Exact row markers verify the installed loader, without BF16 rounding
+        # hiding a channel-order error. The constructor still uses BF16.
+        param.data = param.data.float()
+        for weights in (canonical + 1, canonical):
+            param.weight_loader(param, weights)
+            torch.testing.assert_close(param, _rank_rows(weights, sizes, tp, rank), rtol=0, atol=0)
+        shards.append(param.detach().clone())
+    local_sizes = [size // tp for size in sizes]
+    reconstructed = torch.cat(
+        [torch.cat([shard.split(local_sizes, dim=0)[i] for shard in shards]) for i in range(len(sizes))]
+    )
+    torch.testing.assert_close(reconstructed, canonical, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("tp,heads", [(2, 8), (4, 24), (8, 24)])
+@pytest.mark.parametrize("mask_kind", ["none", "3d", "broadcast", "per_head"])
+@torch.inference_mode()
+def test_local_heads_and_gather_order_match_dense_attention(monkeypatch, tp, heads, mask_kind):
+    torch.manual_seed(17)
+    dim = heads * 8
+    baseline = _attention(monkeypatch, 1, 0, local_tp=False, heads=heads)
+    for name, param in baseline.named_parameters():
+        if name.startswith("norm_"):
+            param.fill_(1)
+        else:
+            param.normal_(std=0.03)
+    # A non-contiguous input and batch > 1 exercise the reshape boundaries.
+    hidden = torch.randn(2, dim, 7, dtype=torch.bfloat16).transpose(1, 2)
+    angles = torch.randn(7, 4)
+    rotary = angles.cos(), angles.sin()
+    mask = None
+    if mask_kind != "none":
+        mask_heads = heads if mask_kind == "per_head" else 1
+        mask = torch.rand(2, mask_heads, 7, 7) > 0.25
+        mask[..., 0] = True
+        if mask_kind == "3d":
+            mask = mask[:, 0]
+
+    captured = []
+    hook = baseline.to_out.register_forward_pre_hook(lambda module, args: captured.append(args[0].clone()))
+    expected = baseline(hidden, attention_mask=mask, image_rotary_emb=rotary)
+    hook.remove()
+    full_attention, full_mlp = captured[0].split([dim, 3 * dim], dim=-1)
+    weights = {name: param.detach().clone() for name, param in baseline.named_parameters()}
+
+    for rank in range(tp):
+        local = _attention(monkeypatch, tp, rank, heads=heads)
+        for name, param in local.named_parameters():
+            getattr(param, "weight_loader", default_weight_loader)(param, weights[name])
+        feature_calls: list[int] = []
+        output_calls: list[int] = []
+
+        def gather_features(value, dim=-1):
+            reference = (full_attention, full_mlp)[len(feature_calls)]
+            torch.testing.assert_close(value, reference.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
+            feature_calls.append(value.shape[-1])
+            return reference
+
+        def gather_output(value, dim=-1):
+            torch.testing.assert_close(value, expected.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
+            output_calls.append(value.shape[-1])
+            return expected
+
+        # Each outgoing tensor must match the independently computed dense
+        # result before the collective supplies the other ranks' slices.
+        monkeypatch.setattr(klein, "tensor_model_parallel_all_gather", gather_features)
+        monkeypatch.setattr(linear, "tensor_model_parallel_all_gather", gather_output)
+        actual = local(hidden, attention_mask=mask, image_rotary_emb=rotary)
+        torch.testing.assert_close(actual, expected)
+        assert feature_calls == [dim // tp, 3 * dim // tp]
+        assert output_calls == [dim // tp]
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"tp": 1},
+        {"tp": 4, "local_tp": False},
+        {"tp": 4, "sp": 2},
+        {"tp": 4, "heads": 6},
+        {"tp": 4, "mlp_ratio": 2.1},
+    ],
+)
+def test_ineligible_shapes_keep_original_projection(monkeypatch, kwargs):
+    attention = _attention(monkeypatch, rank=0, **kwargs)
+    assert not attention.local_tp
+    assert type(attention.to_qkv_mlp_proj) is linear.ColumnParallelLinear
+    assert attention.to_qkv_mlp_proj.gather_output
+    assert attention.query_num_heads == attention.heads
+
+
+@pytest.mark.parametrize(
+    "overrides,parallel_overrides,eligible",
+    [
+        ({}, {}, True),
+        ({"dtype": torch.float32}, {}, False),
+        ({"dtype": torch.float16}, {}, False),
+        ({}, {"tensor_parallel_size": 1}, False),
+        ({}, {"ulysses_degree": 2}, False),
+        ({"quantization_config": "fp8"}, {}, False),
+        ({"lora_path": "local-adapter"}, {}, False),
+        ({"enable_cpu_offload": True}, {}, False),
+        ({"enable_layerwise_offload": True}, {}, False),
+        ({"diffusion_offload_config": {"mode": "module", "components": ["dit"]}}, {}, False),
+        ({}, {"tensor_parallel_size": 1, "use_hsdp": True, "hsdp_shard_size": 4}, False),
+    ],
+)
+def test_pipeline_eligibility_uses_resolved_config(tmp_path, overrides, parallel_overrides, eligible):
+    parallel = DiffusionParallelConfig(**({"tensor_parallel_size": 4} | parallel_overrides))
+    config = OmniDiffusionConfig(
+        **({"model": str(tmp_path), "dtype": torch.bfloat16} | overrides),
+        parallel_config=parallel,
+    )
+    assert klein._use_local_single_stream_tp(config) is eligible
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+def test_host_restore_contract_and_model_weight_reload(monkeypatch, tp):
+    _set_tp(monkeypatch, tp, 0)
+    kwargs = dict(
+        in_channels=8,
+        num_layers=1,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=8,
+        joint_attention_dim=8,
+        timestep_guidance_channels=8,
+        axes_dims_rope=(2, 2, 2, 2),
+        guidance_embeds=False,
+    )
+    with _default_dtype(torch.bfloat16):
+        baseline = klein.Flux2Transformer2DModel(**kwargs)
+        model = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
+    assert baseline.host_weight_restore_contract.version == "1"
+    assert model.host_weight_restore_contract.version == ("2-local-single-stream-tp" if tp > 1 else "1")
+    model.validate_restored_host_weights()
+    name = "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight"
+    weights = torch.linspace(-1, 1, 576 * 64).reshape(576, 64).to(torch.bfloat16)
+    for canonical in (weights, weights.neg()):
+        assert name in model.load_weights([(name, canonical)])
+        param = dict(model.named_parameters())[name]
+        torch.testing.assert_close(param, _rank_rows(canonical, [64, 64, 64, 192, 192], tp, 0), rtol=0, atol=0)
+    with _default_dtype(torch.bfloat16):
+        restored = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
+    restored.load_state_dict(model.state_dict(), strict=True)
+    restored.validate_restored_host_weights()
+    torch.testing.assert_close(dict(restored.named_parameters())[name], param, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("sizes", [(8, 24), (8, 8, 8, 24, 24)])
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+@torch.inference_mode()
+def test_fused_lora_global_rows_dynamic_activation(monkeypatch, sizes, tp):
+    from vllm.lora.lora_model import LoRAModel
+    from vllm.lora.lora_weights import LoRALayerWeights
+    from vllm.lora.peft_helper import PEFTHelper
+    from vllm.lora.request import LoRARequest
+
+    from vllm_omni.diffusion.lora.layers.column_parallel_linear import DiffusionMergedColumnParallelLinearWithLoRA
+    from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+
+    torch.manual_seed(29)
+    weights = torch.randn(sum(sizes), 8)
+    lora_a = torch.randn(2, 8)
+    lora_b = torch.randn(sum(sizes), 2)
+    hidden = torch.randn(2, 3, 8)
+    name = "transformer.to_qkv_mlp_proj"
+    peft = PEFTHelper(r=2, lora_alpha=2, target_modules=["to_qkv_mlp_proj"])
+    request = LoRARequest(lora_name="test", lora_int_id=7, lora_path="unused-by-memory-loader")
+    for rank in range(tp):
+        _set_tp(monkeypatch, tp, rank)
+        projection = linear.MergedColumnParallelLinear(8, list(sizes), bias=False, gather_output=False)
+        projection.weight.weight_loader(projection.weight, weights)
+        pipeline = torch.nn.Module()
+        pipeline.transformer = torch.nn.Module()
+        pipeline.transformer.to_qkv_mlp_proj = projection
+        manager = DiffusionLoRAManager(pipeline, device=torch.device("cpu"), dtype=torch.float32)
+        adapter = LoRAModel(7, 2, {name: LoRALayerWeights(name, 2, 2, lora_a, lora_b)})
+        # Adapter disk I/O is outside this regression. Actual manager injection,
+        # five-/two-slice wrapper, rank slicing and activation all run below.
+        monkeypatch.setattr(manager, "_load_adapter", lambda request: (adapter, peft))
+        base_output = F.linear(hidden, _rank_rows(weights, sizes, tp, rank))
+        delta = F.linear(F.linear(hidden, lora_a), _rank_rows(lora_b, sizes, tp, rank))
+        for scale in (0.5, 1.0):
+            manager.set_active_adapter(request, lora_scale=scale)
+            wrapped = pipeline.transformer.to_qkv_mlp_proj
+            assert isinstance(wrapped, DiffusionMergedColumnParallelLinearWithLoRA)
+            assert wrapped.n_slices == len(sizes)
+            actual, _ = wrapped(hidden)
+            torch.testing.assert_close(actual, base_output + scale * delta, rtol=1e-5, atol=1e-5)
+            manager.set_active_adapter(None)
+            actual, _ = wrapped(hidden)
+            torch.testing.assert_close(actual, base_output, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize("image_count", [1, 2])
+def test_native_editing_reaches_normal_input_validation(tmp_path, monkeypatch, image_count):
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import Flux2KleinPipeline
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    class InputsReachedError(Exception):
+        pass
+
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model=str(tmp_path),
+        dtype=torch.bfloat16,
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=4),
+    )
+    assert klein._use_local_single_stream_tp(pipeline.od_config)
+    images = [Image.new("RGB", (16, 16)) for _ in range(image_count)]
+    image_input = images[0] if image_count == 1 else images
+    request = OmniDiffusionRequest(
+        prompt={"prompt": "paint the boat red", "multi_modal_data": {"image": image_input}},
+        request_id="native-editing",
+        sampling_params=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=4),
+    )
+
+    def check_inputs(**kwargs):
+        assert kwargs["prompt"] == "paint the boat red"
+        raise InputsReachedError
+
+    monkeypatch.setattr(pipeline, "check_inputs", check_inputs)
+    with pytest.raises(InputsReachedError):
+        pipeline.forward(DiffusionRequestBatch(requests=[request]))
+    assert request.prompt["multi_modal_data"]["image"] is image_input

tests/diffusion/models/kandinsky6/test_acceleration.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU checks that Kandinsky 6 acceleration hooks are actually wired."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.kandinsky6.cache_accel import K6_PRO_MAG_RATIOS, K6StepCache
+from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Transformer3DModel
+from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import _shard_loaded_weight
+from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_shard_loaded_weight_narrows_column_parallel_output():
+    class _Owner:
+        tp_rank = 1
+
+        def weight_loader(self, param, loaded):
+            del param, loaded
+
+    param = nn.Parameter(torch.empty(2, 4))
+    param.output_dim = 0
+    param.weight_loader = _Owner().weight_loader
+    loaded = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+
+    shard = _shard_loaded_weight(param, loaded)
+
+    assert tuple(shard.shape) == (2, 4)
+    torch.testing.assert_close(shard, loaded[2:4])
+
+
+def test_row_bias_is_not_sharded_when_shapes_match():
+    class _Owner:
+        tp_rank = 1
+
+        def weight_loader(self, param, loaded):
+            del param, loaded
+
+    param = nn.Parameter(torch.empty(4))
+    param.output_dim = 0
+    param.weight_loader = _Owner().weight_loader
+    loaded = torch.arange(4, dtype=torch.float32)
+
+    shard = _shard_loaded_weight(param, loaded)
+
+    torch.testing.assert_close(shard, loaded)
+
+
+def test_hsdp_predicate_matches_visual_blocks_only_when_indexed():
+    assert Kandinsky6Transformer3DModel._is_transformer_block("visual_transformer_blocks.3", nn.Identity())
+    assert Kandinsky6Transformer3DModel._is_transformer_block("video_text_transformer_blocks.0", nn.Identity())
+    assert not Kandinsky6Transformer3DModel._is_transformer_block("visual_embeddings", nn.Identity())
+    assert Kandinsky6Transformer3DModel._hsdp_shard_conditions
+
+
+def test_sp_plan_shards_visual_tokens_and_gathers_before_the_head():
+    plan = Kandinsky6Transformer3DModel._sp_plan
+    assert "_sp_visual_shard" in plan
+    assert "_sp_visual_rope" in plan
+    assert plan["_sp_visual_gather"].gather_dim == 1
+    assert Kandinsky6Transformer3DModel._magcache_block_attrs == ("visual_transformer_blocks",)
+
+
+def test_k6_is_not_on_the_no_cache_list_and_enablers_are_registered():
+    import vllm_omni.diffusion.cache.cachedit  # noqa: F401  registers enablers
+    from vllm_omni.diffusion.cache.cachedit.backend import CUSTOM_DIT_ENABLERS
+    from vllm_omni.diffusion.cache.teacache.backend import CUSTOM_TEACACHE_ENABLERS
+
+    assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION
+    assert "Kandinsky6TI2VAPipeline" in CUSTOM_TEACACHE_ENABLERS
+    assert "Kandinsky6TI2VAPipeline" in CUSTOM_DIT_ENABLERS
+    assert len(K6_PRO_MAG_RATIOS) == 100
+
+
+def test_step_cache_skips_only_after_a_real_step():
+    cache = K6StepCache(
+        kind="mag",
+        threshold=0.24,
+        max_skip_steps=3,
+        retention_ratio=0.0,
+        mag_ratios=(1.0, 1.0),
+    )
+    assert cache.should_skip(0) is False
+    cache.store(torch.zeros(1), torch.zeros(1))
+    assert cache.should_skip(1) is True

tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import numpy as np
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_kandinsky6_pipeline_import_and_registry():
+    from vllm_omni.diffusion.models.kandinsky6 import (
+        AutoencoderKLHunyuanVideo,
+        Kandinsky6AudioVAE,
+        Kandinsky6TI2VAPipeline,
+        Kandinsky6Transformer3DModel,
+        KandinskyFlowMatchScheduler,
+        get_kandinsky6_post_process_func,
+        get_kandinsky6_pre_process_func,
+    )
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+        _DIFFUSION_PRE_PROCESS_FUNCS,
+        _NO_CACHE_ACCELERATION,
+    )
+
+    assert Kandinsky6TI2VAPipeline is not None
+    assert Kandinsky6Transformer3DModel is not None
+    assert AutoencoderKLHunyuanVideo is not None
+    assert Kandinsky6AudioVAE is not None
+    assert KandinskyFlowMatchScheduler is not None
+    assert get_kandinsky6_post_process_func is not None
+    assert get_kandinsky6_pre_process_func is not None
+
+    assert _DIFFUSION_MODELS["Kandinsky6TI2VAPipeline"] == (
+        "kandinsky6",
+        "pipeline_kandinsky6",
+        "Kandinsky6TI2VAPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_post_process_func"
+    assert _DIFFUSION_PRE_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_pre_process_func"
+    assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION
+
+
+def test_kandinsky6_component_discovery_declarations():
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline
+
+    assert Kandinsky6TI2VAPipeline._dit_modules == ["transformer"]
+    assert Kandinsky6TI2VAPipeline._encoder_modules == ["text_encoder", "text_encoder_2"]
+    assert Kandinsky6TI2VAPipeline._vae_modules == ["vae", "audio_vae"]
+    assert Kandinsky6TI2VAPipeline.supports_step_execution is True
+    assert Kandinsky6TI2VAPipeline.support_audio_output is True
+    assert Kandinsky6TI2VAPipeline.support_image_input is True
+
+
+def test_kandinsky6_pipeline_satisfies_capability_protocols():
+    """isinstance checks against the @runtime_checkable protocols the
+    framework actually uses to detect these capabilities (io_support.py's
+    supports_audio_output, etc.) — stronger than just checking the class
+    attribute exists."""
+    from vllm_omni.diffusion.models.interface import (
+        SupportAudioOutput,
+        SupportImageInput,
+        SupportsComponentDiscovery,
+    )
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline
+
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportAudioOutput)
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportImageInput)
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportsComponentDiscovery)
+
+
+class _FakeAudioVAE:
+    downsample_factor = 1024
+
+
+def test_kandinsky6_post_process_func_packages_video_and_audio():
+    """Pure-function test: no model/weights needed. Confirms the post-process
+    payload shape matches what io_support.py -> output_formatter.py ->
+    media_utils.py's PyAV muxing expects (the same flat {"video", "audio",
+    "audio_sample_rate", "fps"} shape MiniMax H3's own post-process function
+    produces), and that int16 PCM from ``postprocess_audio`` is rescaled to
+    the float32 [-1, 1] waveform the muxer consumes."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 4, 8, 8, 3), dtype=np.uint8)
+    audio = np.array([0, 32767, -32767], dtype=np.int16)
+
+    result = post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")
+
+    assert result["video"] is video
+    assert result["audio"].dtype == np.float32
+    np.testing.assert_allclose(result["audio"], [0.0, 1.0, -1.0])
+    assert result["audio_sample_rate"] == 44100
+    assert result["fps"] == 24.0
+
+
+def test_kandinsky6_post_process_sample_rate_survives_output_formatter():
+    """The formatter only lifts ``audio_sample_rate`` into metadata from the
+    flat payload form; an envelope would silently lose it and the MP4 would
+    be muxed at the 24 kHz default."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+    from vllm_omni.diffusion.output_formatter import normalize_diffusion_postprocess_output
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+    audio = np.zeros((10,), dtype=np.int16)
+
+    normalized = normalize_diffusion_postprocess_output(
+        post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")
+    )
+
+    assert normalized.primary_key == "video"
+    assert normalized.metadata["audio"]["sample_rate"] == 44100
+    assert normalized.metadata["video"]["fps"] == 24.0
+    assert "audio_sample_rate" not in normalized.outputs
+
+
+def test_kandinsky6_post_process_func_omits_audio_when_none():
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+
+    result = post_process({"video": video, "audio": None, "audio_sample_rate": None}, output_type="np")
+
+    assert result["video"] is video
+    assert "audio" not in result
+    assert "audio_sample_rate" not in result
+
+
+def test_kandinsky6_post_process_func_unwraps_batched_audio_list():
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+    audio_item = np.full((50,), 16384, dtype=np.int16)
+
+    result = post_process({"video": video, "audio": [audio_item], "audio_sample_rate": 44100}, output_type="np")
+
+    assert result["audio"].shape == (50,)
+    np.testing.assert_allclose(result["audio"], 16384 / 32767, rtol=1e-6)
+
+
+class _StubAudioVAE:
+    """Audio VAE stand-in: identity "decode" that returns the latent as a waveform."""
+
+    scaling_factor = 0.5
+    mean_value = 0.0
+    device = "cpu"
+
+    def wrapped_decode(self, latents):
+        # (1, audio_dim, A) -> (A,) waveform: take the first channel.
+        return latents[0, 0]
+
+
+def test_postprocess_audio_normalizes_by_default_and_clips_on_request():
+    """The default mode must match the k6_video production pipeline, which
+    peak-normalizes each decoded waveform to full scale; ``clip`` keeps the
+    raw amplitude (saturated to [-1, 1]) for callers that want it."""
+    import torch
+
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import LatentBundle, postprocess_audio
+
+    # Latent (A=4, audio_dim=1); the stub scales by 1/scaling_factor -> [0.2, -0.4, 0.6, 3.0].
+    audio = torch.tensor([[0.1], [-0.2], [0.3], [1.5]], dtype=torch.float32)
+    bundle = LatentBundle(
+        video=None,
+        audio=audio,
+        video_cu_seqlens=None,
+        audio_cu_seqlens=torch.tensor([0, 4], dtype=torch.int32),
+    )
+    vae = _StubAudioVAE()
+
+    normalized = postprocess_audio(bundle, vae)
+    assert normalized is not None and len(normalized) == 1
+    assert normalized[0].dtype == np.int16
+    np.testing.assert_allclose(
+        normalized[0].astype(np.float32) / 32767,
+        np.array([0.2, -0.4, 0.6, 3.0], dtype=np.float32) / 3.0,
+        atol=1e-4,
+    )
+
+    clipped = postprocess_audio(bundle, vae, normalization_mode="clip")
+    np.testing.assert_allclose(
+        clipped[0].astype(np.float32) / 32767,
+        np.array([0.2, -0.4, 0.6, 1.0], dtype=np.float32),
+        atol=1e-4,
+    )
+
+    assert postprocess_audio(LatentBundle(None, None, None, None), vae) is None
+    with pytest.raises(ValueError, match="normalization_mode"):
+        postprocess_audio(bundle, vae, normalization_mode="loud")
+
+
+def test_kandinsky6_pre_process_func_is_identity_for_now():
+    """v1 scope: the pre-process hook is registered (matching the framework's
+    convention for image-conditioned models like SanaImageToVideoPipeline)
+    but does not yet validate/transform the request — documented follow-up,
+    not a silent gap."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_pre_process_func
+
+    pre_process = get_kandinsky6_pre_process_func(od_config=None)
+    sentinel = object()
+
+    assert pre_process(sentinel) is sentinel
+
+
+def test_hub_component_keys_map_onto_pipeline_modules():
+    """Hub folder keys from Kandinsky-6.0-Pro-5s-Diffusers land on the modules."""
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import (
+        _WEIGHT_SUBFOLDERS,
+        _adapt_k6_weight_name,
+    )
+
+    assert _WEIGHT_SUBFOLDERS == (
+        ("transformer", "transformer."),
+        ("vae", "vae."),
+        ("text_encoder", "text_encoder."),
+        ("text_encoder_2", "text_encoder_2."),
+        ("audio_vae", "audio_vae."),
+    )
+    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.videoT.self_attention.to_query.weight") == (
+        "transformer.visual_transformer_blocks.0.video_dec_block.self_attention.to_query.weight"
+    )
+    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.audioT.feed_forward.in_layer.weight") == (
+        "transformer.visual_transformer_blocks.0.audio_dec_block.feed_forward.in_layer.weight"
+    )
+    assert _adapt_k6_weight_name("vae.decoder.conv_in.conv.weight") == "vae.decoder.conv_in.conv.weight"
+    assert _adapt_k6_weight_name("text_encoder_2.encoder.layers.0.mlp.fc1.weight") == (
+        "text_encoder_2.encoder.layers.0.mlp.fc1.weight"
+    )
+    assert (
+        _adapt_k6_weight_name("text_encoder.model.layers.0.input_layernorm.weight")
+        == "text_encoder.model.language_model.layers.0.input_layernorm.weight"
+    )
+    assert _adapt_k6_weight_name("text_encoder.visual.blocks.0.norm1.weight") == (
+        "text_encoder.model.visual.blocks.0.norm1.weight"
+    )
+    assert (
+        _adapt_k6_weight_name("audio_vae.vae.decoder.conv_in.weight")
+        == "audio_vae.native.tod.vae.decoder.conv_in.weight"
+    )
+    assert _adapt_k6_weight_name("audio_vae.vocoder.conv_pre.weight") == "audio_vae.native.tod.vocoder.conv_pre.weight"
+    assert _adapt_k6_weight_name("audio_vae.mel_converter.hann_window") == "audio_vae.native.mel_converter.hann_window"

tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed(monkeypatch):
+    """The native transformer uses vLLM parallel linear layers, which require a
+    tensor-parallel group; initialize a single-process group for CPU tests
+    (mirrors tests/diffusion/models/sana_video/test_transformer_sana_video.py)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda *_args, **_kwargs: default_unquantized_gemm,
+    )
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+# model_dim must be an exact multiple of head_dim = sum(axes_dims); this is
+# an architectural constraint of the native model (DiffusionTransformer3D),
+# not specific to this port.
+_TINY_T2V_CONFIG = {
+    "in_visual_dim": 4,
+    "out_visual_dim": 4,
+    "in_text_dim": 8,
+    "in_text_dim2": 6,
+    "time_dim": 16,
+    "patch_size": (1, 2, 2),
+    "model_dim": 24,
+    "ff_dim": 32,
+    "num_text_blocks": 1,
+    "num_visual_blocks": 2,
+    "axes_dims": (4, 4, 4),
+    "visual_cond": False,
+    "is_multimodal": False,
+    "attention_engine": "sdpa",
+}
+
+# time_dim_a is intentionally omitted (defaults to time_dim): the fused
+# block's cross-modal modulation (va_mod/av_mod) consumes the *other*
+# modality's time embedding by default (fix_modulation=False), so time_dim
+# and time_dim_a must match unless fix_modulation=True.
+_TINY_T2VA_CONFIG = {
+    **_TINY_T2V_CONFIG,
+    "is_multimodal": True,
+    "in_audio_dim": 6,
+    "model_dim_a": 12,
+    "ff_dim_a": 16,
+    "axes_dims_a": (2, 2, 2),
+}
+
+
+def _sdpa_config():
+    from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+    return SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
+        parallel_config=SimpleNamespace(ring_degree=1, allgather_degree=1),
+    )
+
+
+def _build_transformer(config):
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel
+
+    with set_current_diffusion_config(_sdpa_config()):
+        model = Kandinsky6Transformer3DModel(**config)
+    model.eval()
+    return model
+
+
+def test_t2v_only_forward_shape():
+    model = _build_transformer(_TINY_T2V_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    x_video = torch.randn(time_frames, height, width, 4)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    text_rope = model.text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        out = model(
+            x_video=x_video,
+            x_audio=None,
+            text_embed=text_embed,
+            pooled_text_embed=pooled,
+            time=time,
+            visual_rope=visual_rope,
+            audio_rope=None,
+            text_rope=text_rope,
+        )
+
+    assert out.shape == (time_frames, height, width, 4)
+
+
+def test_visual_cond_channel_concat_path():
+    """visual_cond=True doubles+1 the visual input channel count (cond +
+    mask channels appended by _build_video_input, stripped again by
+    OutLayer) — the I2V conditioning path."""
+    config = {**_TINY_T2V_CONFIG, "visual_cond": True}
+    model = _build_transformer(config)
+    time_frames, height, width = 2, 4, 4
+    x_video = torch.randn(time_frames, height, width, 2 * 4 + 1)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    text_rope = model.text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        out = model(
+            x_video=x_video,
+            x_audio=None,
+            text_embed=text_embed,
+            pooled_text_embed=pooled,
+            time=time,
+            visual_rope=visual_rope,
+            audio_rope=None,
+            text_rope=text_rope,
+        )
+
+    assert out.shape == (time_frames, height, width, 4)
+
+
+def test_t2va_joint_forward_shapes():
+    """Exercises the fused cross-modal decoder blocks (Kandinsky6FusedTransformerDecoderBlock)."""
+    model = _build_transformer(_TINY_T2VA_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    audio_len = 7
+    x_video = torch.randn(time_frames, height, width, 4)
+    x_audio = torch.randn(audio_len, 6)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    audio_rope = model.audio_rope_embeddings(torch.arange(audio_len))
+    video_text_rope = model.video_text_rope_embeddings(torch.arange(5))
+    audio_text_rope = model.audio_text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        video_out, audio_out = model(
+            x_video=x_video,
+            x_audio=x_audio,
+            text_embed=[text_embed, text_embed],
+            pooled_text_embed=[pooled, pooled],
+            time=[time, time],
+            visual_rope=visual_rope,
+            audio_rope=audio_rope,
+            text_rope=[video_text_rope, audio_text_rope],
+        )
+
+    assert video_out.shape == (time_frames, height, width, 4)
+    assert audio_out.shape == (audio_len, 6)
+
+
+@pytest.mark.parametrize("modality", ["video", "audio"])
+def test_single_modality_through_multimodal_model(modality):
+    """A multimodal (is_multimodal=True) model must still support a pure
+    video-only or audio-only forward call (denoise_loop's partial-sampling
+    path uses this)."""
+    model = _build_transformer(_TINY_T2VA_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    audio_len = 7
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+
+    if modality == "video":
+        x_video = torch.randn(time_frames, height, width, 4)
+        visual_rope = model.visual_rope_embeddings(
+            shape=(time_frames, height // 2, width // 2),
+            pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+            scale_factor=(1.0, 1.0, 1.0),
+        )
+        text_rope = model.video_text_rope_embeddings(torch.arange(5))
+        with torch.no_grad():
+            out = model(
+                x_video=x_video,
+                x_audio=None,
+                text_embed=text_embed,
+                pooled_text_embed=pooled,
+                time=time,
+                visual_rope=visual_rope,
+                audio_rope=None,
+                text_rope=text_rope,
+            )
+        assert out.shape == (time_frames, height, width, 4)
+    else:
+        x_audio = torch.randn(audio_len, 6)
+        audio_rope = model.audio_rope_embeddings(torch.arange(audio_len))
+        text_rope = model.audio_text_rope_embeddings(torch.arange(5))
+        with torch.no_grad():
+            out = model(
+                x_video=None,
+                x_audio=x_audio,
+                text_embed=text_embed,
+                pooled_text_embed=pooled,
+                time=time,
+                visual_rope=None,
+                audio_rope=audio_rope,
+                text_rope=text_rope,
+            )
+        assert out.shape == (audio_len, 6)
+
+
+@pytest.mark.parametrize("batched", [False, True])
+def test_self_attention_attends_over_tokens_for_batched_and_unbatched_inputs(batched):
+    """Self-attention must contract over the *token* axis whether its input is
+    unbatched text ``(S, C)`` or batched=1 visual/audio ``(1, N, C)``.
+
+    Regression: an unconditional ``unsqueeze(0)`` on the already-batched
+    visual path handed SDPA a 5-D ``(1, 1, N, H, D)`` tensor, which it
+    accepts silently but then treats the head axis as the sequence — every
+    generated video decoded to pure noise."""
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Attention, apply_rotary
+
+    torch.manual_seed(0)
+    channels, head_dim, tokens = 24, 12, 7
+    with set_current_diffusion_config(_sdpa_config()):
+        attn = Kandinsky6Attention(channels, head_dim, engine="sdpa").eval()
+    with torch.no_grad():
+        # vLLM parallel linears allocate with torch.empty (no init).
+        for param in attn.parameters():
+            param.normal_(std=0.2)
+    x = torch.randn(tokens, channels)
+    rope = torch.randn(tokens, 1, head_dim // 2, 2, 2)
+    hidden = x.unsqueeze(0) if batched else x
+
+    with torch.no_grad():
+        out = attn(hidden, rotary_emb=rope)
+
+        heads = channels // head_dim
+        q = attn.query_norm(attn.to_query(x).reshape(tokens, heads, head_dim))
+        k = attn.key_norm(attn.to_key(x).reshape(tokens, heads, head_dim))
+        v = attn.to_value(x).reshape(tokens, heads, head_dim)
+        q = apply_rotary(q, rope).type_as(q)
+        k = apply_rotary(k, rope).type_as(k)
+        ref = torch.nn.functional.scaled_dot_product_attention(
+            q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)
+        ).transpose(0, 1)
+        expected = attn.out_layer(ref.reshape(tokens, channels))
+
+    assert out.shape == (hidden.shape[:-1] + (channels,))
+    torch.testing.assert_close(out.reshape(tokens, channels), expected, rtol=1e-4, atol=1e-5)
+
+
+def test_modulation_projections_are_zero_initialized():
+    """Modulation layers are zero-initialized at construction (standard
+    AdaLN-zero init) — confirms the zero-init in Kandinsky6Modulation.__init__
+    reaches the tensor-parallel out_layer correctly (it's a
+    ColumnParallelLinear, not plain nn.Linear)."""
+    model = _build_transformer(_TINY_T2V_CONFIG)
+    for block in model.text_transformer_blocks:
+        assert torch.all(block.text_modulation.out_layer.weight == 0)
+    for block in model.visual_transformer_blocks:
+        assert torch.all(block.visual_modulation.out_layer.weight == 0)
+
+
+def test_dit_constructs_from_a_diffusers_config_dict():
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel
+
+    with set_current_diffusion_config(_sdpa_config()):
+        model = Kandinsky6Transformer3DModel.from_diffusers_config(
+            {**_TINY_T2VA_CONFIG, "patch_size": [1, 2, 2], "axes_dims": [4, 4, 4]}
+        )
+    assert tuple(model.patch_size) == (1, 2, 2)
+    assert model._repeated_blocks == [
+        "Kandinsky6FusedTransformerDecoderBlock",
+        "Kandinsky6TransformerEncoderBlock",
+        "Kandinsky6TransformerDecoderBlock",
+    ]
+    visual = model.visual_transformer_blocks[0].video_dec_block.self_attention.attn
+    audio = model.visual_transformer_blocks[0].audio_dec_block.self_attention.attn
+    text = model.video_text_transformer_blocks[0].self_attention.attn
+    cross = model.visual_transformer_blocks[0].va_cross_attention.attn
+    assert visual.role == "kandinsky6.visual_self"
+    assert visual.skip_sequence_parallel is False
+    assert audio.skip_sequence_parallel is True
+    assert text.role == "kandinsky6.text_self"
+    assert text.skip_sequence_parallel is True
+    assert cross.role == "kandinsky6.video_audio_cross"
+    assert cross.skip_sequence_parallel is True
+
+
+def test_cfg_combine_matches_apply_cfg_without_normalization():
+    from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import apply_cfg
+
+    positive = torch.tensor([1.0, 3.0])
+    negative = torch.tensor([0.25, -1.0])
+    combined = CFGParallelMixin.combine_cfg_noise(object(), positive, negative, 5.0, False)
+    torch.testing.assert_close(combined, apply_cfg(positive, negative, 5.0))

tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Kandinsky 6 VAE patch parallel keeps the single-GPU tile grid."""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from vllm_omni.diffusion.models.kandinsky6.modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo
+
+
+def _stub_vae() -> AutoencoderKLHunyuanVideo:
+    vae = AutoencoderKLHunyuanVideo.__new__(AutoencoderKLHunyuanVideo)
+    vae.spatial_compression_ratio = 8
+    vae.temporal_compression_ratio = 4
+    vae.use_tiling = True
+    vae.use_framewise_decoding = True
+    vae.tile_sample_min_height = 32
+    vae.tile_sample_min_width = 32
+    vae.tile_sample_min_num_frames = 16
+    vae.tile_sample_stride_height = 24
+    vae.tile_sample_stride_width = 24
+    vae.tile_sample_stride_num_frames = 12
+    vae.tile_size = None
+    vae.distributed_executor = None
+    vae.post_quant_conv = lambda tensor: tensor
+    vae._seen_widths = []
+
+    def decoder(tensor: torch.Tensor) -> torch.Tensor:
+        vae._seen_widths.append(int(tensor.shape[-1]))
+        return tensor + 1
+
+    vae.decoder = decoder
+    return vae
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_distributed_decode_tiles_match_local_tiled_decode():
+    vae = _stub_vae()
+    latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
+
+    local = vae.tiled_decode(latent).sample
+    tasks, spec = vae._decode_tile_split(latent)
+    assert len(tasks) > 1
+    assert max(int(task.tensor.shape[-1]) for task in tasks) < latent.shape[-1]
+
+    decoded = {task.grid_coord: vae._decode_tile_exec(task) for task in tasks}
+    merged = vae._decode_tile_merge(decoded, spec)
+
+    torch.testing.assert_close(merged, local)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_decode_still_applies_optimal_tiling_when_patch_parallel_is_on():
+    vae = _stub_vae()
+    seen = {}
+
+    def tiling(shape, device=None):
+        del device
+        seen["shape"] = tuple(shape)
+        return (1, 17, 32, 32), (8, 24, 24)
+
+    from diffusers.models.autoencoders.vae import DecoderOutput
+
+    vae.get_dec_optimal_tiling = tiling
+    vae._decode = lambda z, return_dict=True: DecoderOutput(sample=z)
+    vae.distributed_executor = SimpleNamespace(parallel_size=2)
+    vae.is_distributed_enabled = lambda: True
+
+    out = vae.decode(torch.zeros(1, 16, 8, 8, 12))
+
+    assert seen["shape"] == (1, 16, 8, 8, 12)
+    assert tuple(out.sample.shape) == (1, 16, 8, 8, 12)
+
+
+def _patch_parallel_worker(rank: int, port: str, shapes: dict) -> None:
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = port
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = "2"
+    dist.init_process_group("gloo", rank=rank, world_size=2)
+    try:
+        import vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor as executor_mod
+
+        world = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
+        executor_mod.get_world_group = lambda: world
+
+        vae = _stub_vae()
+        latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
+        reference = vae.tiled_decode(latent).sample.clone()
+        vae._seen_widths.clear()
+        vae.set_parallel_size(2, mode="tile")
+        decoded = vae.tiled_decode(latent).sample
+        torch.testing.assert_close(decoded, reference)
+        if rank == 0:
+            shapes["widths"] = list(vae._seen_widths)
+            shapes["full_width"] = int(latent.shape[-1])
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_patch_parallel_decode_splits_tiles_across_ranks():
+    manager = mp.Manager()
+    shapes = manager.dict()
+    port = str(29500 + (os.getpid() % 1000))
+    mp.spawn(_patch_parallel_worker, args=(port, shapes), nprocs=2, join=True)
+
+    assert shapes["widths"]
+    assert max(shapes["widths"]) < shapes["full_width"]

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -3,6 +3,8 @@
 """Real-shape checks for MammothModa2's fused QK norm/RoPE and shared CUDA
 attention backend against the native shared RMSNorm path."""
 
+import os
+
 import pytest
 import torch
 from diffusers.models.attention_processor import Attention
@@ -10,6 +12,9 @@ from diffusers.models.attention_processor import Attention
 import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
 import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit
 from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
 
@@ -26,6 +31,13 @@ pytestmark = [
 DIM, HEADS, KV_HEADS = 2520, 21, 7
 
 
+@pytest.fixture(autouse=True)
+def _set_diffusion_config():
+    """Construct attention layers with the same config path used in production."""
+    with set_current_diffusion_config(OmniDiffusionConfig()):
+        yield
+
+
 def _norm_attention(head_dim, dtype):
     attn = Attention(query_dim=head_dim * HEADS, heads=HEADS, kv_heads=KV_HEADS, dim_head=head_dim)
     attn.norm_q = RMSNorm(head_dim)
@@ -33,6 +45,24 @@ def _norm_attention(head_dim, dtype):
     return attn.to(device="cuda", dtype=dtype)
 
 
+def test_explicit_attention_backend_environment_reaches_dit_layer():
+    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+    if backend_name is None or backend_name.lower() == "auto":
+        pytest.skip("no explicit diffusion attention backend")
+    assert backend_name is not None
+
+    block = TransformerBlock(
+        48,
+        6,
+        2,
+        multiple_of=8,
+        ffn_dim_multiplier=1.0,
+        norm_eps=1e-5,
+    )
+    expected_backend = DiffusionAttentionBackendEnum[backend_name.upper()].get_class()
+    assert block.attn.omni_attn.attn_backend is expected_backend
+
+
 def test_unpaired_rope_preserves_native_arithmetic(monkeypatch):
     """A generic block accepts independently valued even and odd RoPE lanes."""
     torch.manual_seed(42)

tests/diffusion/models/qwen_image_21/test_batch_compatibility.py

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

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -0,0 +1,318 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import weakref
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
+from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
+from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class DecodeModel(torch.nn.Module):
+    in_channels = 4
+
+    def _decode_graph_forward(self, entry):
+        return (
+            entry.hidden.sin()
+            + entry.k[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.v[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.freqs[None, :, :]
+        )
+
+
+@pytest.fixture
+def graph_case(request, monkeypatch):
+    device = "cuda" if request.node.get_closest_marker("cuda") else "cpu"
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2)
+    monkeypatch.setattr(manager, "eligible", lambda: True)
+    monkeypatch.setattr(manager, "_backend_name", lambda: "test")
+    cleanup = []
+    monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True))
+    return manager, device, cleanup
+
+
+def make_cache(device, value):
+    return [
+        {
+            "cond": {
+                "key": torch.full((1, 2, 1, 4), value, device=device),
+                "value": torch.full((1, 2, 1, 4), value + 1, device=device),
+            }
+        }
+    ]
+
+
+def register(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    manager.register_prefill(
+        kv_cache=cache,
+        cache_branch="cond",
+        prefix_len=2,
+        img_shapes=[layout],
+        target_freqs=torch.full((4, 4), float(layout[1]), device=device),
+        joint_key_valid=None,
+        dtype=torch.float32,
+    )
+
+
+def decode(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    return manager.try_decode(
+        hidden_states=torch.ones(1, 4, 4, device=device),
+        timestep=torch.ones(1, device=device),
+        kv_cache=cache,
+        cache_branch="cond",
+        img_shapes=[layout],
+        img_mask=torch.tensor([[False, False, True]], device=device),
+        encoder_hidden_states_mask=None,
+    )
+
+
+@pytest.mark.cpu
+def test_model_level_offload_without_staging_disables_decode_graphs():
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True)
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_top_level_offload_hook_without_staging_disables_decode_graphs():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model)
+    registry = HookRegistry.get_or_create(model)
+    # A present-but-empty registry (e.g. after hooks were removed) must not disqualify.
+    assert manager._offload_reason() is None
+    registry.register_hook("sequential_offload", ModelHook())
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_persistent_staging_makes_model_level_offload_graph_eligible():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True)
+    hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu"))
+    HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook)
+    # No staging storage yet: a swap would allocate fresh storage per generation.
+    assert manager._offload_reason() is not None
+    # Staging established: fixed device storage keeps captured pointers valid.
+    hook._stager = object()
+    assert manager._offload_reason() is None
+    # A second, non-offload hook on the top-level module still disqualifies.
+    model._hook_registry.register_hook("teacache", ModelHook())
+    assert manager._offload_reason() is not None
+
+
+@pytest.mark.cpu
+def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+    allocated = []
+
+    def fail_allocation(**kwargs):
+        entry = constructor(**kwargs)
+        allocated.append(weakref.ref(entry.k[0]))
+        raise torch.OutOfMemoryError("injected allocation OOM")
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation)
+    register(manager, cache)
+    register(manager, cache)
+    assert len(allocated) == 1
+    assert allocated[0]() is None
+    assert list(manager.entries.values()) == [None]
+    assert len(cleanup) == 1
+    assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4))
+
+
+@pytest.mark.cpu
+def test_evict_before_allocating_replacement(graph_case, monkeypatch):
+    manager, device, _ = graph_case
+    manager.max_entries = 1
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0])
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+
+    def allocate_after_release(**kwargs):
+        assert old_buffer() is None
+        return constructor(**kwargs)
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release)
+    register(manager, cache, (1, 1, 4))
+    assert len(manager.entries) == 1
+
+
+@pytest.mark.cpu
+def test_registration_preserves_request_cache_and_separates_layouts(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    original = first[0]["cond"]["key"]
+    register(manager, first)
+    register(manager, second)
+    assert first[0]["cond"]["key"] is original
+    assert torch.equal(original, torch.ones_like(original))
+    register(manager, second, (1, 1, 4))
+    assert len(manager.entries) == 2
+    entries = list(manager.entries.values())
+    assert not torch.equal(entries[0].freqs, entries[1].freqs)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+@pytest.mark.parametrize("failure_phase", ["warmup", "capture"])
+def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0])
+    del entry
+
+    original = manager.model._decode_graph_forward
+
+    def fail_capture(entry):
+        output = original(entry)
+        if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing():
+            raise torch.OutOfMemoryError("injected capture OOM")
+        return output
+
+    monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture)
+    assert decode(manager, cache) is None
+    assert entry_ref() is None and buffer_ref() is None
+    assert list(manager.entries.values()) == [None]
+    register(manager, cache)
+    assert decode(manager, cache) is None
+    assert len(cleanup) == 1
+    assert bool((cache[0]["cond"]["key"] == 1).all())
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_replay_uses_active_request_and_layout(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    first_key = weakref.ref(first[0]["cond"]["key"])
+    with torch.inference_mode():
+        register(manager, first)
+        # A second live request aliasing the same key is refused: it keeps its
+        # own cache and decodes eagerly while the owner is in flight.
+        register(manager, second)
+        assert decode(manager, second) is None
+        output = decode(manager, first)
+        expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
+        torch.testing.assert_close(output, expected)
+        # Once the owner's cache is released, the key is free to adopt.
+        del first
+        assert first_key() is None
+        register(manager, second)
+        output = decode(manager, second)
+        expected = (
+            torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2
+        )
+        torch.testing.assert_close(output, expected)
+        original = decode(manager, second).clone()
+        register(manager, second, (1, 1, 4))
+        other = decode(manager, second, (1, 1, 4))
+        torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1)
+        torch.testing.assert_close(decode(manager, second), original)
+        assert all(entry.captures == 1 for entry in manager.entries.values())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("inference", [False, True])
+def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference):
+    manager, device, _ = graph_case
+    with torch.inference_mode(inference):
+        cache = make_cache(device, 1.0) + make_cache(device, 2.0)
+        register(manager, cache)
+        entry = next(iter(manager.entries.values()))
+        entry.refresh_prefix(cache)
+        versions = [tensor._version for tensor in entry.k + entry.v]
+        entry.refresh_prefix(cache)
+        assert versions == [tensor._version for tensor in entry.k + entry.v]
+        cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0)
+        entry.refresh_prefix(cache)
+        assert bool((entry.v[1] == 7).all())
+        source = weakref.ref(cache[0]["cond"]["key"])
+        del cache
+        assert source() is None
+
+
+@pytest.mark.cpu
+def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry.refresh_prefix(cache)
+    cache[0]["cond"]["value"].add_(3)
+    entry.refresh_prefix(cache)
+    assert bool((entry.v[0] == 5).all())
+    register(manager, cache)
+    assert not entry.prefix_sources
+
+
+@pytest.mark.cpu
+def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    for valid in (True, False):
+        manager.register_prefill(
+            kv_cache=cache,
+            cache_branch="cond",
+            prefix_len=2,
+            img_shapes=[(1, 2, 2)],
+            target_freqs=torch.ones(4, 4),
+            joint_key_valid=torch.full((1, 6), valid),
+            dtype=torch.float32,
+        )
+        assert list(manager.entries.values()) == [entry]
+        assert entry.attn_metadata is None
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case):
+    manager, device, _ = graph_case
+    first = make_cache(device, 1.0)
+    second = make_cache(device, 4.0)
+    merged = [
+        {"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}}
+    ]
+    with torch.inference_mode():
+        register(manager, first)
+        single = decode(manager, first).clone()
+        register(manager, merged)
+        kwargs = dict(
+            hidden_states=torch.ones(2, 4, 4, device=device),
+            timestep=torch.ones(2, device=device),
+            kv_cache=merged,
+            cache_branch="cond",
+            img_shapes=[(1, 2, 2)],
+            img_mask=torch.tensor([[False, False, True]] * 2, device=device),
+            encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
+        )
+        result = manager.try_decode(**kwargs)
+        assert result is not None and result.shape == (2, 4, 4)
+        expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None]
+        torch.testing.assert_close(result, expected)
+        # A rebuilt cache (e.g. reordered rows) is a different live cache
+        # aliasing the same key: decode refuses the shared entry until the
+        # previous owner's tensors are released.
+        flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}]
+        kwargs["kv_cache"] = flipped
+        assert manager.try_decode(**kwargs) is None
+        del merged
+        register(manager, flipped)
+        torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0))
+        kwargs["encoder_hidden_states_mask"][1, 0] = False
+        assert manager.try_decode(**kwargs) is None
+        torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_pipeline.py

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

tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py

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

tests/diffusion/models/qwen_image_21/test_registry.py

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

tests/diffusion/models/qwen_image_21/test_token_metadata.py

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

tests/diffusion/models/qwen_image_21/test_transformer_numerics.py

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

tests/diffusion/models/qwen_image_21/test_vae.py

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

tests/diffusion/offloader/test_sequential_backend.py

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

tests/diffusion/test_diffusion_config_propagation.py

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