delete review-PR

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12)
 if [[ "$has_new_python" -eq 0 ]]; then
     # use new python from docker
     docker pull python:3-slim
-    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
+    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
 fi
 
 echo "Using python interpreter: $PYTHON"
@@ -36,7 +36,7 @@ mkdir -p "$INDICES_OUTPUT_DIR"
 
 # HACK: we do not need regex module here, but it is required by pre-commit hook
 # To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
+sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
 
 # Generate indices -- the version is just the commit hash (not omni/{commit})
 # because relative paths are computed between the index and wheel directories,
@@ -73,15 +73,16 @@ echo "Pure version (without variant): $pure_version"
 
 # re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
 if [[ "$version" != *"dev"* ]]; then
-    echo "Re-generating indices for /omni/$pure_version/"
+    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$s3_version/"
     rm -rf "${INDICES_OUTPUT_DIR:?}"
     mkdir -p "$INDICES_OUTPUT_DIR"
     # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
     $PYTHON .buildkite/scripts/generate-nightly-index.py \
-        --version "$pure_version" \
+        --version "$s3_version" \
         --wheel-dir "$BUILDKITE_COMMIT" \
         --current-objects "$obj_json" \
         --output-dir "$INDICES_OUTPUT_DIR" \
         --comment "version $pure_version"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$s3_version/"
 fi

.buildkite/scripts/generate-nightly-index.py

@@ -4,6 +4,7 @@
 
 import argparse
 import json
+import re
 import sys
 from dataclasses import asdict, dataclass
 from datetime import datetime
@@ -11,8 +12,6 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import quote
 
-import regex as re
-
 
 def normalize_package_name(name: str) -> str:
     """Normalize package name per PEP 503."""

.buildkite/test-nightly-diffusion.yml

@@ -0,0 +1,367 @@
+# Nightly diffusion GPU tests — appended to the main nightly build via
+#   buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
+# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are
+# foldable in the Buildkite UI (Other / Wan / Qwen-Image).
+steps:
+  - group: ":card_index_dividers: Other Model Test"
+    key: nightly-other-model-test-group
+    steps:
+      - label: ":full_moon: Diffusion · Other · Function Test with H100"
+        timeout_in_minutes: 120
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not test_qwen_image" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Other · Function Test with L4"
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Diffusion · Other · Doc Test"
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+          - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  - group: ":card_index_dividers: Wan Series Model Test"
+    key: nightly-wan-model-test-group
+    steps:
+      - label: ":full_moon: Diffusion · Wan · Function Test"
+        timeout_in_minutes: 90
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Wan · Accuracy Test"
+        key: nightly-wan22-i2v-accuracy
+        timeout_in_minutes: 180
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  - group: ":card_index_dividers: Qwen-Image Series Model Test"
+    key: nightly-qwen-image-edit-group
+    steps:
+      - label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100"
+        timeout_in_minutes: 120
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_qwen_image*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · GEBench Accuracy Test"
+        key: nightly-gebench-accuracy
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test"
+        key: nightly-gedit-bench-accuracy
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                        value: "120"
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · Perf Test"
+        key: nightly-qwen-image-performance
+        timeout_in_minutes: 180
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export CACHE_DIT_VERSION=1.3.0
+          - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+          - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
+          - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -1,228 +1,199 @@
 steps:
-  - label: ":full_moon: Omni Model Test with H100"
-    timeout_in_minutes: 90
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - |
-        pytest -s -v \
-          tests/examples/ \
-          tests/e2e/online_serving/test_*_expansion.py \
-          -m "advanced_model and H100 and omni" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
+  # Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
+  - group: ":card_index_dividers: Omni Model Test"
+    key: nightly-omni-test-group
+    steps:
+      - label: ":full_moon: Omni · Function Test with H100"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
                   - name: devshm
-                    mountPath: /dev/shm
+                    emptyDir:
+                      medium: Memory
                   - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
-  - label: ":full_moon: Omni Model Test with L4"
-    timeout_in_minutes: 90
+      - label: ":full_moon: Omni · Function Test with L4"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Omni · Doc Test with L4"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Omni · Doc Test with H100"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Omni · Perf Test"
+        key: nightly-omni-performance
+        timeout_in_minutes: 180
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+          - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+          - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  # Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as
+  # pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the
+  # uploaded YAML (Other / Wan / Qwen-Image).
+  - label: ":card_index_dividers: Diffusion Model Test"
+    key: nightly-diffusion-model-test
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-      - pytest -s -v tests/examples/ tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
+      - buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
     agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      queue: "cpu_queue_premerge"
 
-  - label: ":full_moon: Diffusion Model Test with H100"
+  - label: ":bar_chart: Testcase Statistics"
+    key: nightly-testcase-statistics
     timeout_in_minutes: 120
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Model (Wan) Test with H100"
-    timeout_in_minutes: 90
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Model Test"
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-
-  - label: ":full_moon: Doc Example Code Test with H100"
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
-    key: nightly-omni-performance
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export BENCHMARK_DIR=tests/dfx/perf/results
-      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-      - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-      - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
       - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
       - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
     agents:
@@ -259,189 +230,13 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: GEBench Accuracy Test with H100"
-    key: nightly-gebench-accuracy
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: GEdit-Bench Accuracy Test with H100"
-    key: nightly-gedit-bench-accuracy
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
-                    value: "120"
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Wan22 I2V Accuracy Test with H100"
-    key: nightly-wan22-i2v-accuracy
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Perf Test with H100"
-    key: nightly-qwen-image-performance
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-      - export CACHE_DIT_VERSION=1.3.0
-      - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
-      - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
-      - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 4
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
   # No need to run this step for PRs with label nightly-test
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
       - nightly-omni-performance
       - nightly-qwen-image-performance
+      - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

.buildkite/test-ready.yml

@@ -16,11 +16,10 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Voxtral TTS CUDA Unit Test"
-    timeout_in_minutes: 10
+  - label: "CUDA Unit Test with single card"
     depends_on: upload-ready-pipeline
     commands:
-      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
+      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and not distributed_cuda' --ignore=tests/e2e --ignore=tests/engine/test_async_omni_engine_abort.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -33,6 +32,22 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "CUDA Unit Test with multi cards"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and distributed_cuda' --ignore=tests/e2e --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -152,24 +167,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion GPU Worker Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
 
   - label: "Engine Test"
     depends_on: upload-ready-pipeline

.claude/skills/add-diffusion-model/SKILL.md

@@ -0,0 +1,534 @@
+---
+name: add-diffusion-model
+description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including Cache-DiT acceleration and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating a new diffusion model, porting a diffusers pipeline or a custom model repo to vllm-omni, creating a new DiT transformer adapter, adding diffusion model support, or enabling multi-GPU parallelism and cache acceleration for an existing model.
+---
+
+# Adding a Diffusion Model to vLLM-Omni
+
+## Overview
+
+This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.
+
+## Prerequisites
+
+Before starting, determine:
+
+1. **Model category**: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
+2. **Reference source**: Diffusers pipeline, custom repo, or a combination
+3. **Model HuggingFace ID** or local checkpoint path
+4. **Architecture**: Scheduler, text encoder, VAE, transformer/backbone
+
+## Step 0: Classify the Migration Path
+
+Check the model's HF repo for `model_index.json`. This determines your path:
+
+| Scenario | How to identify | Migration path |
+|----------|----------------|----------------|
+| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
+| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
+| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
+| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
+
+## Path A: Diffusers-Based Model
+
+For models with a standard diffusers layout. See [references/transformer-adaptation.md](references/transformer-adaptation.md) for detailed code patterns.
+
+### A1. Analyze `model_index.json`
+
+Identify components: `transformer`, `scheduler`, `vae`, `text_encoder`, `tokenizer`.
+
+### A2. Create model directory
+
+```
+vllm_omni/diffusion/models/your_model_name/
+├── __init__.py
+├── pipeline_your_model.py
+└── your_model_transformer.py
+```
+
+### A3. Adapt transformer
+
+1. Copy from diffusers source. Remove mixins (`ModelMixin`, `ConfigMixin`, `AttentionModuleMixin`).
+2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
+3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
+4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
+5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
+
+### A4. Adapt pipeline
+
+Inherit from `nn.Module`. The key contract:
+
+```python
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        # Load VAE, text encoder, tokenizer via from_pretrained()
+        # Instantiate transformer (weights loaded later via weights_sources)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model, subfolder="transformer",
+                prefix="transformer.", fall_back_to_pt=True)]
+
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # Encode prompt → prepare latents → denoise loop → VAE decode
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights):
+        return AutoWeightsLoader(self).load_weights(weights)
+```
+
+Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
+
+### A5. Register, test, docs → continue at Step 4 below.
+
+---
+
+## Path B: Custom/Private Repo Model
+
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
+
+### B1. Understand the reference repo
+
+Study the original model's code to identify:
+- **Model architecture files** (transformers, fusion modules, embeddings)
+- **Weight format** (safetensors, `.pth`, custom checkpoint structure)
+- **Weight loading helpers** (custom init functions, checkpoint loaders)
+- **Pre/post-processing** (image/audio transforms, tokenization, VAE encode/decode)
+- **External dependencies** (packages not on PyPI)
+- **Config format** (JSON config files, hardcoded dicts)
+
+### B2. Decide what lives WHERE
+
+This is the key design decision for custom models. Follow these placement rules:
+
+| Code type | Where to place | Example |
+|-----------|---------------|---------|
+| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
+| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
+| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
+| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
+| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
+| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
+| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
+
+### B3. Handle external dependencies
+
+If the model's code lives in a separate git repo:
+
+**Option 1: Import with graceful fallback** (recommended for models with external utils)
+
+```python
+try:
+    from external_model.utils import init_vae, load_checkpoint
+except ImportError:
+    raise ImportError(
+        "Failed to import from dependency 'external_model'. "
+        "Please run the download script first."
+    )
+```
+
+**Option 2: Port the code directly** (preferred when feasible)
+
+Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and adapt them. This avoids external dependencies. BAGEL does this — `autoencoder.py` and `bagel_transformer.py` are ported directly.
+
+**Decision criteria**: Port if the code is self-contained and won't diverge. Use external deps if the model repo is actively maintained and the code is complex.
+
+### B4. Handle custom weight loading
+
+Custom models have two patterns for weight loading:
+
+**Pattern 1: Bypass standard loader** (DreamID-Omni style)
+
+When the original model has complex custom init functions that load weights in `__init__`:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        model = od_config.model
+        # Load everything eagerly in __init__ using custom helpers
+        self.vae = custom_init_vae(model, device=self.device)
+        self.text_encoder = custom_init_text_encoder(model, device=self.device)
+        self.transformer = CustomFusionModel(CONFIG)
+        load_custom_checkpoint(self.transformer,
+            checkpoint_path=os.path.join(model, "model.safetensors"))
+        # NO weights_sources defined — bypasses standard loader
+
+    def load_weights(self, weights):
+        pass  # No-op — all weights loaded in __init__
+```
+
+**Pattern 2: Use standard loader with custom `load_weights`** (BAGEL style)
+
+When weights are in safetensors format but need name remapping:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        # Instantiate model architecture without weights
+        self.bagel = BagelModel(config)
+        self.vae = AutoEncoder(ae_params)
+
+        # Point loader at the safetensors in the model root
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder=None,  # weights at root, not in subfolder
+                prefix="",
+                fall_back_to_pt=False,
+            )
+        ]
+
+    def load_weights(self, weights):
+        # Custom name remapping for non-diffusers weight names
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Remap original weight names to vllm-omni module names
+            name = self._remap_weight_name(name)
+            if name in params:
+                default_weight_loader(params[name], tensor)
+                loaded.add(name)
+        return loaded
+```
+
+### B5. Create the `model_index.json`
+
+Custom models need a `model_index.json` at the model root for vllm-omni to discover them. For custom models, this is minimal:
+
+```json
+{
+    "_class_name": "YourModelPipeline",
+    "custom_key": "path/to/custom_weights.safetensors"
+}
+```
+
+The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
+
+If the model's weights come from multiple HF repos, write a **download script** that:
+1. Downloads from each repo
+2. Assembles into a single directory
+3. Generates `model_index.json`
+4. Installs any external dependencies (git clone + `.pth` file)
+
+Place at: `examples/offline_inference/<name>/download_<name>.py`
+
+### B6. Handle multi-modal inputs
+
+If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from `vllm_omni/diffusion/models/interface.py`:
+
+```python
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
+    # Protocol markers — the engine uses these to enable proper input routing
+    pass
+```
+
+Preprocessing for custom models is typically done **inside `forward()`** rather than via registered pre-process functions, since the logic is often tightly coupled to the model.
+
+### B7. Continue at Step 4 below.
+
+---
+
+## Common Steps (Both Paths)
+
+### Step 4: Register Model in registry.py
+
+Edit `vllm_omni/diffusion/registry.py`:
+
+```python
+_DIFFUSION_MODELS = {
+    "YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
+}
+_DIFFUSION_POST_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_post_process_func",  # if applicable
+}
+_DIFFUSION_PRE_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_pre_process_func",  # if applicable
+}
+```
+
+The registry key is the `_class_name` from `model_index.json`. The tuple is `(folder_name, module_file, class_name)`.
+
+Create `__init__.py` exporting the pipeline class and any factory functions.
+
+### Step 5: Run, Test, Debug
+
+Use the appropriate existing example script:
+
+| Category | Script |
+|----------|--------|
+| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
+| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
+| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
+| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
+| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
+
+For custom/Omni models that don't fit these categories, create a dedicated example script.
+
+**Validation**: No errors, output is meaningful, quality matches reference implementation.
+
+See [references/troubleshooting.md](references/troubleshooting.md) for common errors.
+
+### Step 6: Add Example Scripts
+
+For Omni or custom models, create:
+- `examples/offline_inference/your_model_name/` — offline script + README
+- `examples/online_serving/your_model_name/` — server script + client
+- Download script if weights require assembly from multiple sources
+
+### Step 7: Update Documentation
+
+Required updates:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
+2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
+3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
+4. `examples/offline_inference/xxx/README.md` — offline example docs
+5. `examples/online_serve/xxx/README.md` — online serve docs
+
+### Step 8: Add E2E Tests (Recommended)
+
+Create `tests/e2e/online_serving/test_your_model_expansion.py`.
+
+### Step 9: Add Cache-DiT Acceleration
+
+Cache-DiT accelerates inference by caching intermediate computation results across denoising steps. After your model is working correctly on a single GPU, add cache-dit support.
+
+See [references/cache-dit-patterns.md](references/cache-dit-patterns.md) for detailed code patterns.
+
+#### 9a. Determine your model type
+
+| Model Type | Description | Action |
+|------------|-------------|--------|
+| **Standard single-transformer** | One transformer with one `ModuleList` of blocks | No code needed — `CacheDiTBackend` auto-detects via `enable_cache_for_dit()` |
+| **Multi-block-list** | One transformer with multiple block lists (e.g., `transformer_blocks` + `single_transformer_blocks`) | Write custom enabler with `BlockAdapter` |
+| **Dual-transformer** | Two transformers (e.g., high-noise + low-noise) | Write custom enabler with `BlockAdapter` wrapping both |
+
+#### 9b. Standard models — verify automatic support
+
+For standard single-transformer models, test directly:
+
+```python
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+Check logs for "Cache-dit enabled successfully on xxx". If it works, skip to Step 9e.
+
+#### 9c. Custom architectures — write a custom enabler
+
+For multi-block-list or dual-transformer models, write a custom enabler function:
+
+```python
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+#### 9d. Register the custom enabler
+
+Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourModelPipeline": enable_cache_for_your_model,  # Add here
+}
+```
+
+#### 9e. Test Cache-DiT
+
+```python
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1, "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4, "residual_diff_threshold": 0.24,
+    }
+)
+images = omni.generate("a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50))
+```
+
+**Verify**: 1) logs show cache enabled, 2) 1.5-2x speedup, 3) output quality acceptable vs baseline.
+
+If quality degrades, lower `residual_diff_threshold` (try 0.12-0.18) or increase `max_warmup_steps` (try 6-8).
+
+---
+
+### Step 10: Add Parallelism Support
+
+After the model works on a single GPU, add multi-GPU parallelism. Add each type incrementally, testing after each addition.
+
+See [references/parallelism-patterns.md](references/parallelism-patterns.md) for detailed code patterns and API reference.
+
+**Recommended order**: TP → SP/USP → CFG Parallel → HSDP
+
+#### 10a. Tensor Parallelism (TP)
+
+Shards DiT linear layers across GPUs. Requires code changes in the transformer.
+
+**What to change in the transformer**:
+1. Replace `nn.Linear` with `ColumnParallelLinear` / `RowParallelLinear` / `QKVParallelLinear`
+2. Update `load_weights()` to handle QKV fusion with `stacked_params_mapping`
+3. Use `self.to_qkv.num_heads` (local heads) instead of total heads for split sizes
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear, ColumnParallelLinear,
+)
+
+# Attention: QKV → RowParallel output
+self.to_qkv = QKVParallelLinear(dim, head_dim, num_heads, num_kv_heads)
+self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
+
+# FFN: ColumnParallel → RowParallel
+self.w1 = ColumnParallelLinear(dim, ffn_dim)
+self.w2 = RowParallelLinear(ffn_dim, dim, input_is_parallel=True)
+```
+
+**Constraints**: `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`.
+
+**Test**: `--tensor-parallel-size 2`
+
+#### 10b. Sequence Parallelism (SP / USP)
+
+Splits sequence tokens across GPUs. Non-intrusive via `_sp_plan` on the transformer class — no changes to `forward()`.
+
+**What to change in the transformer**:
+
+Add `_sp_plan` class attribute:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class YourTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+If inline tensor ops (e.g., `torch.cat`) exist between shard/gather points, extract them into `nn.Module` submodules so hooks can intercept them.
+
+For RoPE that needs splitting, add an entry for the RoPE module with `split_output=True`.
+
+**Test**: `--ulysses-degree 2` (offline) or `--usp 2` (online serving)
+
+#### 10c. CFG Parallel
+
+Distributes positive/negative CFG branches across 2 GPUs. Requires the pipeline to inherit `CFGParallelMixin`.
+
+**What to change in the pipeline**:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, ...) -> torch.Tensor:
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {...}
+            negative_kwargs = {...} if do_true_cfg else None
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg, true_cfg_scale=cfg_scale,
+                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+Override `predict_noise()` if your transformer call is non-standard. Override `combine_cfg_noise()` for multi-output models (e.g., video + audio).
+
+**Constraint**: Exactly 2 GPUs. Only for models using classifier-free guidance.
+
+**Test**: `--cfg-parallel-size 2`
+
+#### 10d. HSDP (Hybrid Sharded Data Parallel)
+
+Shards transformer weights via PyTorch FSDP2 to reduce per-GPU VRAM. No code changes to the forward pass — just add a class attribute.
+
+**What to change in the transformer**:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Constraint**: Cannot combine with TP. For standalone HSDP, set `hsdp_shard_size` explicitly.
+
+**Test**: `--use-hsdp` or `DiffusionParallelConfig(use_hsdp=True)`
+
+#### 10e. Update parallelism documentation
+
+After adding parallelism support, update:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
+2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
+
+---
+
+## Iterative Development Tips
+
+1. **Start minimal**: Basic generation first, no parallelism/caching
+2. **Use `--enforce-eager`**: Disable torch.compile during debugging
+3. **Use small models**: Test with smaller variants first
+4. **Check tensor shapes**: Most errors are reshape mismatches in attention
+5. **Add features incrementally**: Single GPU → TP → SP → CFG → HSDP → Cache-DiT
+6. **For custom models**: Get the model running with the original code first, then progressively replace components with vllm-omni equivalents
+7. **Cache-DiT before parallelism tuning**: Cache-DiT is lossy — verify quality at baseline before combining with parallelism
+8. **Combine lossless + lossy**: e.g., TP + SP + Cache-DiT for maximum throughput
+
+## Reference Files
+
+- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
+- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
+- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
+- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
+- [Troubleshooting](references/troubleshooting.md) — common errors and fixes

.claude/skills/add-diffusion-model/references/cache-dit-patterns.md

@@ -0,0 +1,254 @@
+# Cache-DiT Patterns Reference
+
+## Overview
+
+Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.
+
+Three caching strategies:
+- **DBCache**: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
+- **TaylorSeer**: Calibration-based prediction using Taylor expansion to estimate block outputs
+- **SCM** (Step Computation Masking): Dynamic step skipping based on configurable policies
+
+**Typical speedup**: 1.5-2.5x depending on model and configuration.
+
+**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
+
+## Architecture
+
+vLLM-Omni integrates cache-dit through `CacheDiTBackend`:
+
+| Component | Purpose |
+|-----------|---------|
+| `CacheDiTBackend` | Unified backend — auto-selects enabler (standard or custom) |
+| `enable_cache_for_dit()` | Default enabler for standard single-transformer models |
+| `CUSTOM_DIT_ENABLERS` dict | Registry of custom enablers keyed by pipeline class name |
+| `BlockAdapter` | Wraps complex architectures (multi-block-list or multi-transformer) |
+| `ForwardPattern` | Specifies block forward signature: `Pattern_0`, `Pattern_1`, `Pattern_2` |
+| `ParamsModifier` | Per-transformer or per-block-list config customization |
+| `DBCacheConfig` | Configuration for DBCache parameters |
+| `cache_dit.refresh_context()` | Updates cache context when `num_inference_steps` changes |
+
+**Source files**:
+- `vllm_omni/diffusion/cache/cache_dit_backend.py` — `CacheDiTBackend`, enablers, `CUSTOM_DIT_ENABLERS`
+- `vllm_omni/diffusion/cache/` — cache backend implementations
+
+## Standard Models: Automatic Support
+
+Most DiT models follow this pattern:
+- Single transformer with one `nn.ModuleList` of blocks
+- Standard forward signature
+- Compatible with cache-dit's automatic detection
+
+**Examples**: Qwen-Image, Z-Image, FLUX
+
+No code changes needed. `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+What happens automatically:
+
+```python
+def enable_cache_for_dit(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+## Custom Architectures: Writing Custom Enablers
+
+### When you need a custom enabler
+
+- Model has multiple block lists in one transformer (e.g., `transformer_blocks` + `single_transformer_blocks`)
+- Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
+- Model uses non-standard block forward signature
+
+### Pattern 1: Multi-Block-List (LongCat-Image style)
+
+Single transformer with two block lists:
+
+```python
+import cache_dit
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+For single transformer with multiple block lists, `refresh_context` works the same as standard models — call it once on the transformer.
+
+### Pattern 2: Dual-Transformer (Wan2.2 style)
+
+Two transformers with separate configs:
+
+```python
+def enable_cache_for_dual_transformer(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(...)
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=[pipeline.transformer, pipeline.transformer_2],
+            blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
+            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+            params_modifiers=[
+                ParamsModifier(...),  # Config for transformer 1
+                ParamsModifier(...),  # Config for transformer 2
+            ],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        high_steps, low_steps = _split_inference_steps(num_inference_steps)
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=high_steps, verbose=verbose
+        )
+        cache_dit.refresh_context(
+            pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+Key difference: `refresh_context` must be called on **each transformer separately** with its own step count.
+
+### Choosing the ForwardPattern
+
+| Pattern | Block forward signature | Example models |
+|---------|------------------------|----------------|
+| `Pattern_0` | `block(hidden_states, **kwargs)` → residual added inside block | Default |
+| `Pattern_1` | `block(hidden_states, **kwargs)` → returns `(hidden_states, ...)` tuple | FLUX-style single blocks |
+| `Pattern_2` | `block(hidden_states, **kwargs)` → `(hidden_states, ...)` with different residual pattern | Wan2.2 blocks |
+
+Inspect your block's `forward()` return type and residual connection pattern to choose the right one. See [Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) for details.
+
+## Registering Custom Enablers
+
+Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourModelPipeline": enable_cache_for_your_model,
+}
+```
+
+The key must match `pipeline.__class__.__name__`.
+
+## Configuration Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `Fn_compute_blocks` | 1 | Number of blocks to always compute at the front |
+| `Bn_compute_blocks` | 0 | Number of blocks to always compute at the back |
+| `max_warmup_steps` | 4 | Steps to run without caching at the beginning |
+| `max_cached_steps` | — | Max total cached steps |
+| `max_continuous_cached_steps` | — | Max consecutive cached steps |
+| `residual_diff_threshold` | 0.24 | Threshold for deciding whether to cache a block |
+
+### Tuning for quality vs speed
+
+| Goal | Adjustments |
+|------|-------------|
+| **More speed, acceptable quality loss** | Higher `residual_diff_threshold` (0.24-0.4), lower `max_warmup_steps` (2-4) |
+| **Better quality, less speed** | Lower `residual_diff_threshold` (0.12-0.18), higher `max_warmup_steps` (6-8), lower `max_continuous_cached_steps` (2) |
+
+## Testing
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.24,
+    }
+)
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI (online serving):
+
+```bash
+vllm serve your-model --omni --port 8098 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'
+```
+
+**Verification checklist**:
+1. Logs show "Cache-dit enabled successfully on xxx"
+2. Performance: 1.5-2x speedup vs no cache
+3. Quality: compare output with `cache_backend=None`
+
+## Excluded Models
+
+Models listed in `_NO_CACHE_ACCELERATION` in `vllm_omni/diffusion/registry.py` do not support cache-dit (e.g., `NextStep11Pipeline`, `StableDiffusionPipeline`). Check this set before attempting to enable cache-dit.
+
+## Reference Implementations
+
+| Model | Path | Notes |
+|-------|------|-------|
+| Standard DiT | `cache_dit_backend.py::enable_cache_for_dit` | Default enabler, automatic |
+| Wan2.2 | `cache_dit_backend.py::enable_cache_for_wan22` | Dual-transformer, auto-detects mode |
+| LongCat | `cache_dit_backend.py::enable_cache_for_longcat_image` | Multi-block-list |
+| BAGEL | `cache_dit_backend.py::enable_cache_for_bagel` | Complex omni model |

.claude/skills/add-diffusion-model/references/custom-model-patterns.md

@@ -0,0 +1,273 @@
+# Custom Model Patterns Reference
+
+Patterns for adding models that don't come from the standard diffusers pipeline format.
+
+## Directory Structure Comparison
+
+### Diffusers-based model (e.g., Wan2.2)
+
+```
+vllm_omni/diffusion/models/wan2_2/
+├── __init__.py                    # Exports pipeline + transformer + helpers
+├── pipeline_wan2_2.py             # Pipeline: loads components via from_pretrained()
+├── pipeline_wan2_2_i2v.py         # Variant pipeline for image-to-video
+└── wan2_2_transformer.py          # Transformer: ported from diffusers, uses Attention layer
+```
+
+The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
+
+### Custom model with external deps (e.g., DreamID-Omni)
+
+```
+vllm_omni/diffusion/models/dreamid_omni/
+├── __init__.py                    # Exports pipeline only
+├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
+├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
+└── wan2_2.py                      # Re-implemented Wan backbone with split API
+
+examples/offline_inference/x_to_video_audio/
+└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
+```
+
+All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
+
+### Custom model with ported code (e.g., BAGEL)
+
+```
+vllm_omni/diffusion/models/bagel/
+├── __init__.py
+├── pipeline_bagel.py              # Pipeline: instantiates models, uses weights_sources
+├── bagel_transformer.py           # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
+└── autoencoder.py                 # Custom VAE ported from original repo
+```
+
+Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
+
+## Weight Loading Patterns
+
+### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
+
+```
+init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
+```
+
+- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
+- `load_weights()` receives `(name, tensor)` pairs from the loader
+- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
+
+### Pattern 2: Custom safetensors at root (BAGEL)
+
+```
+init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
+```
+
+- `weights_sources` points to **root** of model directory, not a subfolder
+- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
+- `load_weights()` does heavy name normalization
+
+```python
+self.weights_sources = [
+    DiffusersPipelineLoader.ComponentSource(
+        model_or_path=od_config.model,
+        subfolder=None,      # root directory
+        prefix="",           # no prefix stripping
+        fall_back_to_pt=False,
+    )
+]
+```
+
+### Pattern 3: Fully custom loading (DreamID-Omni)
+
+```
+init → load ALL weights eagerly via custom helpers → load_weights() = no-op
+```
+
+- No `weights_sources` attribute — standard loader finds nothing to iterate
+- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
+- `load_weights()` is `pass`
+- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
+
+Use this when:
+- The original model has complex, well-tested loading code you don't want to rewrite
+- Weights span multiple HF repos
+- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
+
+## model_index.json for Custom Models
+
+Standard diffusers `model_index.json`:
+```json
+{
+    "_class_name": "WanPipeline",
+    "_diffusers_version": "0.35.0.dev0",
+    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+    "transformer": ["diffusers", "WanTransformer3DModel"],
+    "vae": ["diffusers", "AutoencoderKLWan"]
+}
+```
+
+Custom model `model_index.json` (minimal):
+```json
+{
+    "_class_name": "DreamIDOmniPipeline",
+    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
+}
+```
+
+The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
+
+## External Dependency Management
+
+### Git clone + .pth injection (DreamID-Omni pattern)
+
+```python
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            subprocess.run([
+                "git", "clone", "--depth", "1",
+                REPO_URL, "--branch", BRANCH,
+                str(DEPENDENCY_DIR)
+            ], check=True)
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # Add to Python path via .pth file
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+```
+
+### Direct port (BAGEL pattern)
+
+Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
+
+## Multi-Modal Input/Output Protocols
+
+Custom models that handle images, audio, or video I/O should implement protocol classes:
+
+```python
+from vllm_omni.diffusion.models.interface import (
+    SupportImageInput,    # Model accepts image input
+    SupportAudioInput,    # Model accepts audio input
+    SupportAudioOutput,   # Model produces audio output
+)
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
+    pass  # Protocol markers enable proper engine routing
+```
+
+The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
+
+## Hardcoded Config vs Config Files
+
+Diffusers models use `config.json` in each subfolder. Custom models often use:
+
+**Module-level config dicts** (DreamID-Omni):
+```python
+VIDEO_CONFIG = {
+    "patch_size": [1, 2, 2], "model_type": "ti2v",
+    "dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
+}
+```
+
+**Loaded from custom JSON** (BAGEL):
+```python
+cfg_path = os.path.join(model_path, "config.json")
+with open(cfg_path) as f:
+    bagel_cfg = json.load(f)
+vae_cfg = bagel_cfg.get("vae_config", {})
+```
+
+## Custom Architecture Patterns
+
+### Split forward API (DreamID-Omni)
+
+When a fusion model needs to interleave blocks from two backbones:
+
+```python
+class WanModel(nn.Module):
+    def prepare_transformer_block_kwargs(self, x, t, context, ...):
+        # Patch embed, time embed, text embed, RoPE
+        return x, e, kwargs
+
+    def post_transformer_block_out(self, x, grid_sizes, e):
+        # Output projection, unpatchify
+        return output
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError  # Fusion model handles block iteration
+```
+
+The `FusionModel` then iterates blocks in lock-step:
+```python
+for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
+    video_out = video_block(video_hidden, ...)
+    audio_out = audio_block(audio_hidden, ...)
+    # Cross-attend between modalities
+    video_out = cross_attention(video_out, audio_out)
+    audio_out = cross_attention(audio_out, video_out)
+```
+
+### LLM-as-denoiser (BAGEL)
+
+When the backbone is a language model that also does diffusion:
+
+```python
+class BagelModel(nn.Module):
+    def __init__(self):
+        self.language_model = Qwen2MoTForCausalLM(config)
+        self.vit_model = SiglipVisionModel(vit_config)
+```
+
+The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
+
+## Pre/Post Processing for Custom Models
+
+Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
+
+```python
+def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    # Inline preprocessing
+    image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
+    image_latent = self._vae_encode(image)
+
+    # ... denoising loop ...
+
+    # Inline postprocessing
+    pil_image = self._decode_to_pil(latents)
+    return DiffusionOutput(output=[pil_image])
+```
+
+If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
+
+## Download Script Template
+
+```python
+# examples/offline_inference/<name>/download_<name>.py
+from huggingface_hub import snapshot_download
+import json, os
+
+def main(output_dir):
+    # Download model weights from HF
+    snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
+
+    # Download additional components if from separate repos
+    snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
+        allow_patterns=["*.safetensors"])
+
+    # Generate model_index.json
+    config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
+    with open(os.path.join(output_dir, "model_index.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    # Install external code dependency (if needed)
+    download_dependency()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", default="./your_model")
+    args = parser.parse_args()
+    main(args.output_dir)
+```

.claude/skills/add-diffusion-model/references/parallelism-patterns.md

@@ -0,0 +1,571 @@
+# Parallelism Patterns Reference
+
+## Overview
+
+vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:
+
+| Strategy | Splits | Best For | Constraint |
+|----------|--------|----------|------------|
+| Tensor Parallel (TP) | Model layers across GPUs | Latency reduction, large models | Requires fast GPU interconnect, `num_heads % tp == 0` |
+| Sequence Parallel (SP/USP) | Sequence tokens across GPUs | Long sequences (video, high-res) | Near-linear scaling |
+| CFG Parallel | Positive/negative CFG branches | Models using classifier-free guidance | Exactly 2 GPUs |
+| HSDP | Weight shards via FSDP2 | VRAM reduction | Cannot combine with TP |
+| VAE Patch Parallel | VAE decode spatial tiles | Large VAE outputs | Auto-enables tiling |
+
+**Recommended integration order**: TP → SP → CFG Parallel → HSDP
+
+**Official design docs**:
+- TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
+- SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
+- CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
+- HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp
+
+---
+
+## Tensor Parallelism (TP)
+
+Replace standard `nn.Linear` with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.
+
+### Layer replacement rules
+
+| Pattern | vLLM Layer | When to Use |
+|---------|-----------|-------------|
+| Fan-out (first in FFN) | `ColumnParallelLinear` | Projection that splits output across ranks |
+| Fan-in (second in FFN) | `RowParallelLinear` | Projection that gathers across ranks |
+| QKV projection | `QKVParallelLinear` | Fused Q/K/V for self-attention |
+| Single Q or K or V | `ColumnParallelLinear` | Separate projections (cross-attention) |
+| Attention output | `RowParallelLinear` | Output projection after attention |
+| Must not shard | `ReplicatedLinear` | Layers that must stay replicated |
+
+### MLP Block (Up-Down Pattern)
+
+```python
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear, RowParallelLinear,
+)
+
+class TPFeedForward(nn.Module):
+    def __init__(self, dim, ffn_dim):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
+        self.fc2 = RowParallelLinear(
+            ffn_dim, dim, bias=False,
+            input_is_parallel=True,  # Input already sharded from fc1
+            return_bias=False,
+        )
+
+    def forward(self, x):
+        x, _ = self.fc1(x)
+        x = torch.nn.functional.gelu(x)
+        x, _ = self.fc2(x)
+        return x
+```
+
+### Attention Block (QKV-Out Pattern)
+
+```python
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm_omni.diffusion.attention.layer import Attention
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads, num_kv_heads=None):
+        super().__init__()
+        num_kv_heads = num_kv_heads or num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            return_bias=False,
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads,      # Local heads per GPU
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,  # Local KV heads per GPU
+        )
+
+    def forward(self, x):
+        qkv, _ = self.to_qkv(x)
+        q, k, v = qkv.split(
+            [self.to_qkv.num_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim],
+            dim=-1,
+        )
+        B, S, _ = x.shape
+        q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
+        k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        out = self.attn(q, k, v)
+        out = out.reshape(B, S, -1)
+        out, _ = self.to_out(out)
+        return out
+```
+
+### QKV Fusion in load_weights
+
+When you fuse separate Q/K/V into `QKVParallelLinear`, map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+    return loaded
+```
+
+### RMSNorm with TP
+
+When RMSNorm sits between TP-sharded dimensions, use `DistributedRMSNorm` — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.
+
+### TP Constraints
+
+- `num_heads % tp_size == 0`
+- `num_kv_heads % tp_size == 0`
+- Use `self.to_qkv.num_heads` (local per-GPU count), not total heads, for split sizes
+
+### Testing TP
+
+```bash
+python text_to_image.py --model Your-org/your-model \
+  --tensor-parallel-size 2 --output "tp_test.png"
+```
+
+**Verify**: speedup, memory reduction proportional to TP size, quality matches single-GPU.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+| FLUX | `vllm_omni/diffusion/models/flux/flux_transformer.py` |
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+
+---
+
+## Sequence Parallelism (SP / USP)
+
+SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the `_sp_plan` dict — no changes to `forward()` logic.
+
+### Approach 1: Non-Intrusive `_sp_plan` (Recommended)
+
+The framework automatically registers hooks to shard inputs and gather outputs at `nn.Module` boundaries.
+
+#### Step 1: Identify module boundaries
+
+Find where tensors need sharding/gathering:
+
+```python
+class MyTransformer(nn.Module):
+    def __init__(self):
+        self.patch_embed = PatchEmbed()    # Before blocks
+        self.pos_embed = RoPE()            # RoPE may need splitting
+        self.blocks = nn.ModuleList([...]) # Blocks process sharded x
+        self.norm_out = LayerNorm()
+        self.proj_out = Linear()           # Gather after this
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        pos = self.pos_embed(x)
+        for block in self.blocks:
+            x = block(x, pos)
+        x = self.norm_out(x)
+        return self.proj_out(x)
+```
+
+#### Step 2: Handle inline operations
+
+`_sp_plan` hooks only work at `nn.Module` boundaries. Inline ops like `torch.cat()` must be extracted into submodules:
+
+```python
+# BAD: Inline — hooks can't intercept
+unified = torch.cat([x, cap_feats], dim=1)
+
+# GOOD: Extract into submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)
+
+self.unified_prepare = UnifiedPrepare()
+unified = self.unified_prepare(x, cap_feats)
+```
+
+Common cases: `torch.cat()`, `pad_sequence()`, `tensor.reshape()`, complex preprocessing.
+
+#### Step 3: Write `_sp_plan`
+
+**Pattern 1: Shard at first block, gather at output** (most common)
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 2: Shard RoPE outputs separately**
+
+```python
+class TransformerWithRoPE(nn.Module):
+    _sp_plan = {
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        },
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 3: Dual-stream (shard image, replicate text)**
+
+```python
+class DualStreamTransformer(nn.Module):
+    _sp_plan = {
+        "rope_preparer": {
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+        },
+        "transformer_blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+### API Reference
+
+**SequenceParallelInput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `split_dim` | int | Dimension to split (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+| `split_output` | bool | `False`: shard input params; `True`: shard output tensors |
+| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size |
+
+**SequenceParallelOutput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `gather_dim` | int | Dimension to gather (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+
+**Module naming**:
+
+| Key | Meaning |
+|-----|---------|
+| `"blocks.0"` | First element of ModuleList |
+| `"blocks.*"` | All elements of ModuleList |
+| `"rope"` | Named submodule |
+
+**Dictionary value types**:
+
+| Key type | split_output | Description |
+|----------|-------------|-------------|
+| `"param_name"` (str) | False | Shard input parameter by name |
+| `0, 1, ...` (int) | True | Shard output tuple by index |
+
+### Approach 2: Intrusive Modification (Complex Cases)
+
+For dynamic sharding logic that can't be expressed via `_sp_plan`:
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+    for block in self.blocks:
+        hidden_states = block(hidden_states)
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_gather(hidden_states, dim=1)
+    return hidden_states
+```
+
+Use intrusive modification as a last resort — `_sp_plan` is preferred for maintainability.
+
+### UAA Mode (Experimental)
+
+`ulysses_mode="advanced_uaa"` handles arbitrary sequence lengths and head counts that aren't divisible by `ulysses_degree`. Uses variable all-to-all split sizes and temporary head padding.
+
+### Combining SP methods
+
+Ulysses and Ring can be combined: `ulysses_degree × ring_degree = total SP GPUs`.
+
+```python
+DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
+```
+
+### Testing SP
+
+```bash
+# Offline
+python text_to_image.py --model Your-model --ulysses-degree 2
+
+# Online serving
+vllm serve Your-model --omni --usp 2
+```
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+
+---
+
+## CFG Parallelism
+
+Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.
+
+### Implementation
+
+Inherit `CFGParallelMixin` and implement `diffuse()`:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
+                do_true_cfg, true_cfg_scale, **kwargs):
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": prompt_embeds,
+                "timestep": t,
+            }
+            negative_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": negative_embeds,
+                "timestep": t,
+            } if do_true_cfg else None
+
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=true_cfg_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+### Customization hooks
+
+| Method | Override when |
+|--------|-------------|
+| `predict_noise()` | Non-standard transformer call (e.g., dual-transformer like Wan2.2) |
+| `cfg_normalize_function()` | Custom normalization (e.g., LongCat with clamping) |
+| `combine_cfg_noise()` | Multi-output models (e.g., video + audio: CFG on video, positive-only on audio) |
+
+**Custom predict_noise** (Wan2.2 — selects active transformer):
+
+```python
+def predict_noise(self, current_model=None, **kwargs):
+    if current_model is None:
+        current_model = self.transformer
+    return current_model(**kwargs)[0]
+```
+
+**Custom combine_cfg_noise** (multi-output):
+
+```python
+def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
+    video_pos, audio_pos = positive_pred
+    video_neg, audio_neg = negative_pred
+    video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+    return (video_combined, audio_pos)
+```
+
+### Composite scheduler for multi-output
+
+When each output has its own schedule:
+
+```python
+class VideoAudioScheduler:
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(
+            noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
+        )[0]
+        audio_out = self.audio_scheduler.step(
+            noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
+        )[0]
+        return ((video_out, audio_out),)
+```
+
+### Testing CFG Parallel
+
+```bash
+python text_to_image.py --model Your-model \
+  --cfg-parallel-size 2 --cfg-scale 4.0 \
+  --negative-prompt "ugly, unclear"
+```
+
+**Constraint**: `guidance_scale > 1.0` and negative prompt must be provided.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` |
+| Mixin base | `vllm_omni/diffusion/distributed/cfg_parallel.py` |
+
+---
+
+## HSDP (Hybrid Sharded Data Parallel)
+
+Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.
+
+### Implementation
+
+Add `_hsdp_shard_conditions` to the transformer class:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+For MoE models, add additional conditions:
+
+```python
+class MoETransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name, module):
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name, module):
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+A module is sharded if **any** condition returns `True`.
+
+### Constraints
+
+- Cannot combine with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+- Can combine with SP: HSDP reduces memory while SP distributes sequence
+
+### Testing HSDP
+
+```python
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
+omni = Omni(model="your-model", parallel_config=parallel_config)
+```
+
+Or CLI:
+
+```bash
+vllm serve Your-model --omni --use-hsdp
+```
+
+**Verify**: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| HSDP Core | `vllm_omni/diffusion/distributed/hsdp.py` |
+
+---
+
+## VAE Patch Parallelism
+
+Shards VAE decode spatially across ranks using tiling:
+
+```bash
+python text_to_image.py --model Your-model --vae-patch-parallel-size 4
+```
+
+Auto-enables `--vae-use-tiling`. Uses `DistributedAutoencoderKLWan` or similar distributed VAE. Set `vae_patch_parallel_size` in `DiffusionParallelConfig`.
+
+---
+
+## Combining Parallelism Methods
+
+Common multi-GPU recipes:
+
+```bash
+# 4 GPUs: CFG (2) × Ulysses (2)
+python text_to_image.py --model Qwen/Qwen-Image \
+  --cfg-parallel-size 2 --ulysses-degree 2
+
+# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
+python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8
+
+# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
+vllm serve Your-model --omni --use-hsdp --usp 2
+```
+
+## Discovering Parallelism Support
+
+Check which parallelism methods a model supports:
+
+| Check | How |
+|-------|-----|
+| **Ulysses / Ring SP** | Transformer defines `_sp_plan`. Search: `grep -r '_sp_plan' vllm_omni/diffusion/models/` |
+| **CFG Parallel** | Pipeline inherits `CFGParallelMixin`. Search: `grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/` |
+| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
+| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
+
+The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.

.claude/skills/add-diffusion-model/references/transformer-adaptation.md

@@ -0,0 +1,218 @@
+# Transformer Adaptation Reference
+
+## Adapting a Diffusers Transformer to vLLM-Omni
+
+### Step-by-step Checklist
+
+1. Copy the transformer class from diffusers source
+2. Remove all mixin classes — inherit only from `nn.Module`
+3. Replace attention dispatch with `vllm_omni.diffusion.attention.layer.Attention`
+4. Replace logger with `vllm.logger.init_logger`
+5. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`
+6. Remove training-only code (gradient checkpointing, dropout)
+7. Add `load_weights()` method for weight loading from safetensors
+8. Add class-level attributes for acceleration features
+
+### Mixin Removal
+
+Remove these diffusers mixins (and their imports):
+
+```python
+# Remove all of these:
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import AttentionModuleMixin
+from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
+
+# Replace:
+class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
+# With:
+class MyTransformer(nn.Module):
+```
+
+Also remove `@register_to_config` decorators from `__init__`.
+
+### Attention Replacement
+
+The vLLM-Omni `Attention` layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.
+
+**QKV tensor shape must be `[batch, seq_len, num_heads, head_dim]`.**
+
+#### Self-Attention Pattern
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_q = nn.Linear(dim, dim)
+        self.to_k = nn.Linear(dim, dim)
+        self.to_v = nn.Linear(dim, dim)
+        self.to_out = nn.Linear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=num_heads,
+        )
+
+    def forward(self, x, attn_mask=None):
+        B, S, _ = x.shape
+        q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
+        k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
+        v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)
+
+        attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+        out = self.attn(q, k, v, attn_metadata=attn_metadata)
+        out = out.reshape(B, S, -1)
+        return self.to_out(out)
+```
+
+#### Fused QKV with TP (Advanced)
+
+For tensor parallelism, use vLLM's parallel linear layers:
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear
+)
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_heads,
+        )
+        self.to_out = RowParallelLinear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=num_heads,
+        )
+```
+
+### Logger Replacement
+
+```python
+# Replace:
+from diffusers.utils import logging
+logger = logging.get_logger(__name__)
+
+# With:
+from vllm.logger import init_logger
+logger = init_logger(__name__)
+```
+
+### Custom Layers from vLLM-Omni
+
+Available utility layers:
+
+```python
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+```
+
+### Config Support
+
+```python
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+class MyTransformer(nn.Module):
+    def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config if od_config else None
+        # ... build layers
+```
+
+The transformer config values come from `model_index.json` → `config.json` in the transformer subfolder. The pipeline uses `get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass)` to filter config keys to match the `__init__` signature.
+
+### Weight Loading
+
+The `load_weights` method receives an iterable of `(name, tensor)` from safetensors files, with the prefix (e.g., `"transformer."`) already stripped by the loader.
+
+```python
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+class MyTransformer(nn.Module):
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Optional: remap names from diffusers to vllm-omni naming
+            # e.g., "ff.net.0.proj" -> "ff.net_0.proj"
+
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+        return loaded
+```
+
+#### QKV Fusion in load_weights
+
+If you fused separate Q/K/V into a `QKVParallelLinear`, you need to map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            # Normal loading
+            ...
+    return loaded
+```
+
+### Class-Level Attributes for Features
+
+```python
+class MyTransformer(nn.Module):
+    # torch.compile: list block class names that repeat and can be compiled
+    _repeated_blocks = ["MyTransformerBlock"]
+
+    # CPU offload: attribute name of the nn.ModuleList containing blocks
+    _layerwise_offload_blocks_attr = "blocks"
+
+    # LoRA: mapping of fused param names to original param names
+    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+
+    # Sequence parallelism plan (advanced — add after basic impl works)
+    _sp_plan = {
+        "blocks.0": SequenceParallelInput(split_dim=1),
+        "proj_out": SequenceParallelOutput(gather_dim=1),
+    }
+```

.claude/skills/add-diffusion-model/references/troubleshooting.md

@@ -0,0 +1,178 @@
+# Troubleshooting Reference
+
+## Common Errors When Adding a Diffusion Model
+
+### ImportError / ModuleNotFoundError
+
+**Cause**: Missing or incorrect registration.
+
+**Fix checklist**:
+1. Model registered in `vllm_omni/diffusion/registry.py` `_DIFFUSION_MODELS` dict
+2. `__init__.py` exports the pipeline class
+3. Pipeline file exists at the correct path: `vllm_omni/diffusion/models/{folder}/{file}.py`
+4. Class name in registry matches the actual class name in the file
+
+### Shape Mismatch in Attention
+
+**Symptom**: `RuntimeError: shape mismatch` or `expected 4D tensor`
+
+**Cause**: QKV tensors not reshaped to `[batch, seq_len, num_heads, head_dim]`.
+
+**Fix**: Before calling `self.attn(q, k, v, ...)`, ensure:
+```python
+q = q.view(batch, seq_len, self.num_heads, self.head_dim)
+k = k.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+v = v.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+```
+
+After attention, reshape back:
+```python
+out = out.reshape(batch, seq_len, -1)
+```
+
+### Weight Loading Failures
+
+**Symptom**: `RuntimeError: size mismatch for parameter ...` or missing keys
+
+**Debugging**:
+1. Print diffusers weight names: `safetensors.safe_open(path, "pt").keys()`
+2. Print model parameter names: `dict(model.named_parameters()).keys()`
+3. Compare and add name remappings in `load_weights()`
+
+**Common remappings needed**:
+- `ff.net.0.proj` → `ff.net_0.proj` (PyTorch Sequential indexing)
+- `.to_out.0.` → `.to_out.` (Sequential unwrapping)
+- `scale_shift_table` → moved to a wrapper module
+
+### Black/Blank/Noisy Output
+
+**Possible causes**:
+1. **Wrong latent normalization**: Check VAE expects latents scaled by `vae.config.scaling_factor`
+2. **Wrong scheduler**: Using the wrong scheduler class or wrong `flow_shift`
+3. **Missing CFG**: Some models require `guidance_scale > 1.0` with negative prompt
+4. **Wrong timestep format**: Some schedulers expect float, others expect int/long
+5. **Missing post-processing**: Raw VAE output may need denormalization
+
+**Quick test**: Run with diffusers directly using the same seed and compare latents at each step.
+
+### OOM (Out of Memory)
+
+**Solutions** (in order of preference):
+1. `--enforce-eager` to disable torch.compile (saves compile memory)
+2. `--enable-cpu-offload` for model-level offload
+3. `--enable-layerwise-offload` for block-level offload (better for large models)
+4. `--vae-use-slicing --vae-use-tiling` for VAE memory reduction
+5. Reduce resolution: `--height 480 --width 832`
+6. Use TP: `--tensor-parallel-size 2`
+
+### Different Output vs Diffusers Reference
+
+**Common causes**:
+1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
+2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
+3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
+4. **Scheduler rounding**: Some schedulers have numerical sensitivity
+
+### Tensor Parallel Errors
+
+**Symptom**: `AssertionError: not divisible` or incorrect output with TP>1
+
+**Fix**:
+1. Verify `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`
+2. Ensure `ColumnParallelLinear` / `RowParallelLinear` are used correctly
+3. Check that norms between parallel layers use distributed norm if needed
+4. Verify `load_weights` handles TP sharding for norm weights
+5. Use `self.to_qkv.num_heads` (local heads per GPU) for QKV split sizes, not total heads
+
+**Missing `input_is_parallel=True`**:
+
+`RowParallelLinear` expects sharded input from `ColumnParallelLinear`:
+```python
+self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False)
+self.w2 = RowParallelLinear(hidden_dim, dim, input_is_parallel=True, return_bias=False)
+```
+
+### Sequence Parallel Errors
+
+**Symptom**: Incorrect output or crashes with `--ulysses-degree N` or `--usp N`
+
+**Possible causes**:
+1. **Inline operations between shard/gather points**: `torch.cat()`, `pad_sequence()` etc. not at `nn.Module` boundaries. Fix: extract into submodule.
+2. **Wrong `split_dim`**: Check the tensor shape at the shard point. Sequence dimension is typically `dim=1` for `[B, S, D]` tensors.
+3. **RoPE not sharded**: If RoPE is computed separately, add it to `_sp_plan` with `split_output=True`.
+4. **Sequence not divisible by SP degree**: Use `auto_pad=True` in `SequenceParallelInput` or switch to `ulysses_mode="advanced_uaa"`.
+
+**Debugging**: Add `expected_dims=N` to `SequenceParallelInput`/`Output` for shape validation at runtime.
+
+### CFG Parallel Errors
+
+**Symptom**: CFG parallel not activating, no speedup
+
+**Fix checklist**:
+1. Pipeline inherits `CFGParallelMixin`
+2. `guidance_scale > 1.0`
+3. Negative prompt provided (even if empty string)
+4. `--cfg-parallel-size 2` specified
+5. `diffuse()` method calls `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
+
+**Symptom**: Different output with CFG parallel vs sequential
+
+**Possible cause**: Non-deterministic scheduler. Fix: pass `generator=torch.Generator(device).manual_seed(seed)` to `scheduler_step_maybe_with_cfg()`.
+
+### HSDP Errors
+
+**Symptom**: HSDP not activating or errors during weight loading
+
+**Fix checklist**:
+1. Transformer defines `_hsdp_shard_conditions` class attribute
+2. Shard condition functions return `True` for correct modules (test with `model.named_modules()`)
+3. Not combining with TP (HSDP and TP are incompatible)
+4. For standalone HSDP, `hsdp_shard_size` is specified explicitly
+
+**Verify**: Check logs for "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root".
+
+### Cache-DiT Not Applied
+
+**Symptom**: No speedup, no cache-related log messages
+
+**Fix checklist**:
+1. Model not in `_NO_CACHE_ACCELERATION` in `registry.py`
+2. Pipeline class name matches `CUSTOM_DIT_ENABLERS` key (if using custom enabler)
+3. `cache_backend="cache_dit"` specified
+4. Check logs for "Cache-dit enabled successfully on xxx"
+
+**Verify pipeline name**: `print(pipeline.__class__.__name__)` — must match registry key.
+
+### Cache-DiT Quality Degradation
+
+**Symptom**: Artifacts or lower quality with cache-dit
+
+**Fix**: Reduce aggressiveness:
+```python
+cache_config={
+    "residual_diff_threshold": 0.12,      # Lower from 0.24
+    "max_warmup_steps": 6,                # Increase from 4
+    "max_continuous_cached_steps": 2,      # Reduce if higher
+}
+```
+
+If quality is still poor, the model may need a custom enabler with per-block-list `ParamsModifier` tuning.
+
+### Model Not Detected / Wrong Pipeline Class
+
+**Symptom**: `ValueError: Model class ... not found in diffusion model registry`
+
+**Cause**: The model's `model_index.json` has a `_class_name` for the pipeline that doesn't match registry keys.
+
+**Fix**: The registry key must match the diffusers pipeline class name from `model_index.json`. If using a different name, map it in the registry:
+```python
+"DiffusersPipelineClassName": ("your_folder", "your_file", "YourVllmClassName"),
+```
+
+## Debugging Workflow
+
+1. **Add verbose logging**: Use `logger.info()` to print tensor shapes at each stage
+2. **Compare step-by-step**: Run diffusers and vllm-omni side by side, comparing tensors after each major operation
+3. **Use small configs**: Reduce `num_inference_steps=2`, small resolution for fast iteration
+4. **Test transformer isolation**: Feed the same input to both diffusers and vllm-omni transformers, compare outputs
+5. **Binary search for bugs**: Comment out blocks/layers to isolate where divergence starts

.claude/skills/add-tts-model/SKILL.md

@@ -0,0 +1,284 @@
+---
+name: add-tts-model
+description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFace reference implementation through production-ready serving with streaming and CUDA graph acceleration. Use when adding a new TTS model, wiring stage separation for speech synthesis, enabling online voice generation serving, debugging TTS integration behavior, or building audio output pipelines."
+---
+
+# TTS Model Integration Workflow
+
+## Overview
+
+```
+HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
+   (Phase 1)      (Phase 2)          (Phase 3)        (Phase 4)     (Phase 5)
+```
+
+## Phase 1: HuggingFace Reference
+
+**Goal**: Understand the reference implementation and verify it produces correct audio.
+
+### Steps
+
+1. **Run the reference model** end-to-end using the official HuggingFace / GitHub code
+2. **Document the architecture**:
+   - What are the sub-models? (AR decoder, codec decoder, vocoder, etc.)
+   - What is the token vocabulary? (semantic codes, RVQ codebooks, special tokens)
+   - What is the output format? (sample rate, channels, codec type)
+3. **Capture reference outputs** for comparison during integration
+4. **Identify the config structure**: `config.json` fields, `model_type`, sub-model configs
+
+### Key Questions
+
+- How many codebooks? What are the codebook sizes?
+- What special tokens exist? (`<|voice|>`, `<|audio_start|>`, `<|im_end|>`, etc.)
+- What is the token-to-ID mapping for codec codes?
+- What is the hop length / frame rate of the codec?
+- Does the model support voice cloning? How? (reference audio encoding, speaker embeddings, etc.)
+
+### Deliverables
+
+- Working reference script that produces audio
+- Architecture diagram / notes
+- Token vocabulary mapping
+- Reference audio samples for regression testing
+
+## Phase 2: Stage Separation (Offline Inference)
+
+**Goal**: Split the model into vLLM-Omni stages and get offline inference working.
+
+### Steps
+
+1. **Register the model** in `vllm_omni/model_executor/models/registry.py`
+2. **Create config classes** (`configuration_<model>.py`) with `model_type` registration
+3. **Implement Stage 0** (AR model):
+   - Subclass appropriate base (e.g., wrap Qwen3 decoder layers)
+   - Implement `forward()` for autoregressive token generation
+   - Handle special token logic (start/stop tokens, codec token mapping)
+   - If dual-AR (like Fish Speech), implement Fast AR as a nested module
+4. **Implement Stage 1** (Decoder):
+   - Load codec weights (may need lazy loading from separate checkpoint)
+   - Implement `forward()`: codec codes -> audio waveform
+   - Return `OmniOutput` with `multimodal_outputs`
+5. **Create stage config YAML** defining both stages, memory allocation, and model paths
+6. **Create stage input processor** for prompt building
+7. **Write end2end.py** test script
+
+### Critical Parameters to Get Right
+
+| Parameter | Impact if Wrong |
+|-----------|----------------|
+| Hop length | Audio duration wrong, streaming noise |
+| Token ID mapping | Garbage codes -> noise output |
+| Codebook count/size | Shape mismatch crashes |
+| Stop token | Generation never stops or stops too early |
+| dtype / autocast | Numerical issues, silent quality degradation |
+| Repetition penalty | Must match reference (often 1.0 for TTS) |
+
+### Debugging Priority (from experience)
+
+When audio output is wrong, check in this order:
+
+1. **RoPE / attention**: Are position encodings correct? Is the attention mask right?
+2. **Normalization**: RMSNorm epsilon, layer norm placement (pre vs post)
+3. **Hop length**: Product of all upsample rates in the codec decoder
+4. **Token mapping**: Are codec IDs correctly offset from the vocabulary base?
+5. **Sampling parameters**: Temperature, top_k, top_p, repetition_penalty
+6. **Tensor layout**: Codebook-major vs frame-major ordering
+7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
+
+### Deliverables
+
+- Model files in `vllm_omni/model_executor/models/<model_name>/`
+- Stage config YAML
+- Working `end2end.py` with correct audio output
+- README.md in the example directory
+
+## Phase 3: Online Serving
+
+**Goal**: Expose the model via `/v1/audio/speech` API endpoint.
+
+### Steps
+
+1. **Register in `serving_speech.py`**:
+   - Add model stage name to `_TTS_MODEL_STAGES` set
+   - Add model detection flag (e.g., `_is_fish_speech`)
+   - Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
+2. **Handle model-specific parameters**:
+   - Voice cloning: `ref_audio` encoding and prompt injection
+   - `max_new_tokens` override in sampling params
+   - Model-specific default values
+3. **Create client scripts**: `speech_client.py`, `run_server.sh`
+4. **Test all response formats**: wav, mp3, flac, pcm
+5. **Add Gradio demo**: Interactive web UI with streaming support
+
+### Voice Cloning Pattern
+
+```python
+import base64
+from pathlib import Path
+
+def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
+    """Build prompt with reference audio for voice cloning in serving_speech.py."""
+    audio_bytes = Path(ref_audio_path).read_bytes()
+    codes = codec.encode(audio_bytes)  # Encode on CPU using model's codec (e.g., DAC)
+    token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
+    return [
+        {"role": "system", "content": f"<|voice|>{''.join(chr(t) for t in token_ids)}"},
+        {"role": "user", "content": text},
+    ]
+```
+
+### Deliverables
+
+- Updated `serving_speech.py` with model-specific prompt builder
+- Client scripts and server launcher
+- Gradio demo with streaming and voice cloning UI
+- Documentation (offline + online serving docs)
+
+## Phase 4: Async Chunk (Streaming)
+
+**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
+
+### Steps
+
+1. **Update stage config YAML**:
+   ```yaml
+   async_chunk: true
+   codec_chunk_frames: 25      # frames per chunk
+   codec_left_context_frames: 25  # overlap for smooth boundaries
+   ```
+2. **Implement chunk handling in Stage 1**:
+   - Accept partial input (chunk of codec codes)
+   - Handle left context for smooth audio boundaries
+   - Return partial audio in `OmniOutput`
+3. **Test streaming**:
+   - Verify audio quality matches non-streaming output
+   - Check for artifacts at chunk boundaries
+   - Measure TTFA (time to first audio)
+4. **Update online serving** to support `stream=true` with PCM output
+
+### Streaming Architecture
+
+```
+Stage 0 (AR)                    Stage 1 (Decoder)
+  |                                |
+  |-- chunk 0 (25 frames) ------> decode -> audio chunk 0 -> client
+  |-- chunk 1 (25 frames) ------> decode -> audio chunk 1 -> client
+  |-- chunk 2 (25 frames) ------> decode -> audio chunk 2 -> client
+  ...
+```
+
+### Key Considerations
+
+- **Left context overlap**: Prevents audible artifacts at chunk boundaries
+- **Hop length matters**: `context_audio_samples = context_frames * hop_length`
+- **First chunk latency**: Can use larger initial chunk for better quality, then smaller chunks
+
+### Deliverables
+
+- Updated stage config with async_chunk enabled
+- Smooth streaming audio without boundary artifacts
+- TTFA metrics
+
+## Phase 5: CUDA Graph Acceleration
+
+**Goal**: Capture the AR loop as a CUDA graph for significant speedup.
+
+### Steps
+
+1. **Identify the hot loop**: The AR decoding loop that runs N steps per token
+2. **Create static buffers**:
+   - KV caches with fixed max sequence length
+   - Pre-built causal masks and position tensors per step
+   - Static input/output tensors
+3. **Implement graph capture**:
+   - Warm up with real data
+   - Capture the forward pass
+   - Replay with updated inputs
+4. **Handle constraints**:
+   - Use `torch.argmax` instead of `torch.multinomial` (graph-safe)
+   - Fixed batch size (fall back to eager for other sizes)
+   - No dynamic control flow inside the graph
+
+### Example: Code Predictor CUDA Graph (Qwen3-TTS)
+
+```python
+import torch
+
+class CodePredictorGraph:
+    """Captures the 16-step code predictor AR loop as a single CUDA graph."""
+
+    def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
+        self.num_steps = 16
+        self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
+        self.positions = torch.arange(self.num_steps, device=device)
+        self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
+        self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
+        self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
+        # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
+
+    def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
+        self.input_buf.copy_(initial_input)
+        self.graph.replay()
+        return self.output_buf.clone()
+```
+
+### Performance Expectations
+
+Based on Qwen3-TTS code predictor experience:
+- **3-5x speedup** for the graphed component
+- Only effective for fixed batch sizes (typically batch_size=1)
+- Falls back to eager mode for unsupported configurations
+
+### Deliverables
+
+- CUDA graph implementation for the AR hot loop
+- Benchmark script comparing eager vs graph performance
+- Documentation of constraints and fallback behavior
+
+## Integration Checklist
+
+Use this checklist when integrating a new TTS model:
+
+### Phase 1: HF Reference
+- [ ] Reference model runs and produces correct audio
+- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
+- [ ] Reference audio samples saved for comparison
+
+### Phase 2: Stage Separation
+- [ ] Model registered in `registry.py`
+- [ ] Config classes created with `model_type` registration
+- [ ] Stage 0 (AR) implemented and generates correct tokens
+- [ ] Stage 1 (Decoder) produces correct audio from tokens
+- [ ] Stage config YAML created
+- [ ] `end2end.py` produces audio matching reference quality
+- [ ] README.md written
+
+### Phase 3: Online Serving
+- [ ] Model added to `serving_speech.py`
+- [ ] Prompt builder handles text input correctly
+- [ ] Voice cloning works (if supported)
+- [ ] All response formats work (wav, mp3, flac, pcm)
+- [ ] Client scripts and server launcher created
+- [ ] Gradio demo working
+- [ ] Documentation added (offline + online docs, nav, supported models)
+
+### Phase 4: Async Chunk
+- [ ] Stage config updated with `async_chunk: true`
+- [ ] Stage 1 handles partial chunks correctly
+- [ ] No audio artifacts at chunk boundaries
+- [ ] Streaming via API (`stream=true`) works
+- [ ] TTFA measured and acceptable
+
+### Phase 5: CUDA Graph
+- [ ] Hot loop identified and profiled
+- [ ] Static buffers allocated
+- [ ] Graph captured and replays correctly
+- [ ] Benchmark shows meaningful speedup
+- [ ] Fallback to eager works for unsupported configs
+
+## References
+
+- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
+- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
+- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
+- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)

.claude/skills/readme.md

@@ -0,0 +1,34 @@
+# Claude Skills for vLLM-Omni
+
+This directory contains Claude Code skills maintained for the `vllm-omni`
+repository. These skills capture repeatable workflows for common contributor
+tasks such as model integration, pull request review, and release note
+generation.
+
+## Directory Structure
+
+Each skill lives in its own directory under `.claude/skills/`. A skill may
+include:
+
+- `SKILL.md`: the main workflow and operating instructions
+- `references/`: focused reference material used by the skill
+- `scripts/`: small helper scripts used by the skill
+
+## Available Skills
+
+- `add-diffusion-model`: guides integration of a new diffusion model into
+  `vllm-omni`
+- `add-omni-model`: covers addition of new omni-modality model support
+- `add-tts-model`: covers integration of new TTS models and related serving
+  workflows
+- `generate-release-note`: helps prepare release notes for repository changes
+- `review-pr`: provides a structured workflow for reviewing pull requests
+
+## Maintenance Guidelines
+
+- Keep skill names short and task-oriented.
+- Prefer repository-local paths, commands, and examples.
+- Avoid hardcoding fast-changing support matrices unless the skill is actively
+  maintained alongside those changes.
+- Treat skills as contributor tooling: optimize for clarity, actionability, and
+  low maintenance overhead.

.gitignore

@@ -158,7 +158,19 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-.claude/
+/.claude/*
+!.claude/skills/
+!.claude/skills/readme.md
+!.claude/skills/add-diffusion-model/
+!.claude/skills/add-diffusion-model/SKILL.md
+!.claude/skills/add-diffusion-model/references/
+!.claude/skills/add-diffusion-model/references/*.md
+!.claude/skills/add-tts-model/
+!.claude/skills/add-tts-model/SKILL.md
+!.claude/skills/review-pr/
+!.claude/skills/review-pr/SKILL.md
+!.claude/skills/review-pr/references/
+!.claude/skills/review-pr/references/*.md
 
 # Codex
 AGENTS.md

docs/.nav.yml

@@ -98,6 +98,7 @@ nav:
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
       - design/feature/cfg_parallel.md
+      - design/feature/expert_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
       - design/feature/vae_parallel.md

docs/design/feature/expert_parallel.md

@@ -0,0 +1,221 @@
+# Expert Parallel
+
+This section describes how to add Expert Parallel (EP) to a diffusion transformer that uses Mixture-of-Experts (MoE) layers.
+We use **HunyuanImage3.0** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Expert Parallel?
+
+**Expert Parallel** is a parallelism strategy in Mixture-of-Experts (MoE) models that distributes different expert networks across distinct computational devices. Each device holds and computes only a subset of experts (local experts), with tokens dispatched to and gathered from remote devices via collective communication operations (e.g., All-to-All, All-Gather).
+
+| Backend | Description |
+|---------|-------------|
+| `allgather_reducescatter` | Default backend based on allgather/reducescatter primitives, suitable for general EP+DP deployments.|
+
+## Configuration
+
+Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as:
+
+```text
+EP_SIZE = TP_SIZE × SP_SIZE × CFG_SIZE × DP_SIZE
+```
+
+
+Where:
+
+- `TP_SIZE`: Tensor parallel size
+- `SP_SIZE`: Sequence parallel size
+- `CFG_SIZE`: Classifier-free guidance parallel size
+- `DP_SIZE`: Data parallel size
+- `EP_SIZE`: Expert parallel size (computed automatically)
+
+Note:
+- Expert parallelism is only applicable to Mixture-of-Experts (MoE) models.
+- The EP group is created **per pipeline stage**, meaning it includes all ranks that participate in model parallelism except pipeline parallelism.
+- The underlying communication pattern for expert parallelism is **All-to-All** among the ranks in the EP group.
+
+For example, consider a configuration with `TP=2`, `SP=1`, `CFG=2`, and `DP=4` (total 2×1×2×4 = 16 GPUs).
+
+- Expert layers are handled by an EP group of size 16.
+
+- Attention layers use tensor parallelism of size 2 within each of the 8 DP groups (because `DP×CFG×SP = 4×2×1 = 8` groups, each containing the 2 TP ranks). Inside each such group, the attention weights are sharded across the 2 GPUs.
+
+
+## Step-by-Step Implementation
+
+### Step 1: Configure Expert Parallelism Settings
+
+Calculate local experts per rank:
+
+```
+ep_size = 8  # Expert Parallel size (typically equals TP size)
+num_experts = 64
+num_local_experts = num_experts // ep_size  # 8 experts per card
+
+# Check divisibility
+assert num_experts % ep_size == 0, "Experts must be divisible by EP size"
+```
+
+### Step 2: Use Sparse MoE Block to enable EP routing.
+
+Example:
+```
+from vllm.model_executor.layers.linear import ReplicatedLinear
+class HunYuanSparseMoeBlock(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_id: int = -1,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.n_routed_experts = config.num_experts  # 64
+
+        # Calculate local experts per rank (key for EP)
+        if self.tp_size > self.n_routed_experts:
+            raise ValueError(f"TP size {self.tp_size} > experts {self.n_routed_experts}")
+
+        # Routing gate (replicated on all ranks, computes scores for all tokens to all experts)
+        self.gate = ReplicatedLinear(
+            config.hidden_size,
+            config.num_experts,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.gate",
+        )
+
+        # EP expert layer (factory loads platform-specific implementation)
+        self.experts = HunyuanFusedMoE(...)
+```
+**Key Points:**
+- gate is **ReplicatedLinear** (replicated on all ranks)
+- experts is created via **HunyuanFusedMoE factory**, which automatically handles EP dispatch
+
+### Step 3: Initialize EP Runtime
+
+Initialize the EP communication context before model loading.
+```
+from vllm.utils.import_utils import resolve_obj_by_qualname
+# Call during __init__ or model loading
+op_name = "hunyuan_fused_moe"
+
+# Prepare EP runtime: establish communication groups, assign local expert indices, init _expert_map
+current_omni_platform.prepare_diffusion_op_runtime(op_name)
+
+# Factory automatically resolves platform implementation (GPU: FusedMoE / NPU: AscendFusedMoE)
+impl = resolve_obj_by_qualname(
+    current_omni_platform.get_diffusion_model_impl_qualname(op_name)
+)
+```
+
+### Step 4: Expert Weight Mapping & Loading
+
+Each rank loads only the expert weights assigned to its local allocation.
+```
+# Get expert parameter mapping (different per rank)
+expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
+    model=self,
+    ckpt_gate_proj_name="gate_proj",
+    ckpt_down_proj_name="down_proj",
+    ckpt_up_proj_name="up_proj",
+    num_experts=64,
+    num_redundant_experts=0,
+)
+# Returns: [(param_name, weight_name, expert_id, shard_id), ...]
+# Note: Each rank only contains mappings for its local expert_ids
+
+# Filter non-local experts during loading
+for name, loaded_weight in weights:
+    if "mlp.experts" in name:
+        # Parse expert_id from weight name (implementation needed)
+        expert_id = parse_expert_id_from_name(name)
+        local_expert_start = (ep_rank) * num_local_experts
+        local_expert_end = (ep_rank + 1) * num_local_experts
+
+        if not (local_expert_start <= expert_id < local_expert_end):
+            continue  # Skip non-local expert weights
+```
+### Step 5: Forward Pass with EP
+
+Example (MoE Forward):
+```
+def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    orig_shape = hidden_states.shape
+    hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+
+    # 1. Global routing computation (all tokens, all expert scores)
+    # hidden_states: [num_tokens, hidden_dim] (full tensor)
+    router_logits, _ = self.gate(hidden_states)  # [num_tokens, num_experts]
+
+    # 2. EP dispatch and compute (HunyuanFusedMoE handles all_to_all internally)
+    # - Dispatch: Send tokens to target ranks based on router_logits
+    # - Local Compute: Each rank processes only its num_local_experts
+    # - Combine: Results returned to original token positions
+    final_hidden_states = self.experts(
+        hidden_states=hidden_states,
+        router_logits=router_logits,
+    )
+
+    # 3. Add shared expert output (not EP, computed on all ranks)
+    if self.shared_mlp is not None:
+        shared_out = self.shared_mlp(hidden_states)
+        final_hidden_states = final_hidden_states + shared_out
+
+    # 4. Tensor Parallel All-Reduce (synchronize across TP group)
+    if self.tp_size > 1:
+        final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
+            final_hidden_states
+        )
+
+    return final_hidden_states.view(orig_shape)
+```
+
+## Testing
+After adding Expert Parallel support, test via command line:
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --output "ep_enabled.png" \
+    --num-inference-steps 50 \
+    --guidance-scale 5.0 \
+    --tensor-parallel-size 8 \
+    --seed 1234 \
+    --enable-expert-parallel
+```
+
+vLLM‑Omni currently focuses on core diffusion model inference acceleration, so the Expert Parallel implementation includes only the basic multi‑GPU expert sharding functionality (enabled via --enable-expert-parallel). Advanced features such as communication backend selection (--all2all-backend), load balancing (--enable-eplb and its configuration), and multi‑node deployment belong to the extended capabilities of the main vLLM project and have not yet been integrated into Omni.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py` | Standard EP | Full implementation with validation |
+| **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance |
+| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
+
+---
+## Summary
+
+Adding Expert Parallel support to diffusion model:
+
+1. **Identify MoE layers** - Locate the router and expert networks in each transformer block.
+2. **Validate EP constraints** – Ensure num_experts is divisible by expert_parallel_size.
+3. **Test** - Run with enable-expert-parallel, check memory reduction, speedup, and output quality against single‑GPU baseline.

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

@@ -33,8 +33,8 @@ runtime:
         zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
         protocol: "rdma"              # "rdma" or "tcp"
         device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
-        memory_pool_size: 2147483648  # 2GB memory pool
-        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
+        memory_pool_size: 4294967296  # 4 GB (CPU); use 2147483648 (2 GB) for GPU
+        memory_pool_device: "cpu"     # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
 ```
 
 Wire stages to the connector:
@@ -64,8 +64,8 @@ stage_args:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
-| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
+| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). |
 
 ### Networking
 
@@ -107,10 +107,10 @@ receiver_connect  = remote_side_channel_port + tp_rank
 
 ## Memory Pool Modes
 
-| Mode | Config | Data Flow | Best For |
-|---|---|---|---|
-| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
-| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+| Mode | Config | Recommended Pool Size | Data Flow | Best For |
+|---|---|---|---|---|
+| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
 
 > **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
 > switch (PIX topology). On systems where they are connected via PXB or NODE,

docs/models/supported_models.md

@@ -60,5 +60,6 @@ th {
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -139,11 +139,15 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 ## Supported Models
 
-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-------------|
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
 | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally

docs/user_guide/diffusion_features.md

@@ -12,7 +12,7 @@
 
 vLLM-Omni supports various advanced features for diffusion models:
 
-- Acceleration: **cache methods**, **parallelism methods**
+- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**
 - Execution modes: **step execution**
@@ -44,6 +44,12 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
 | **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
 
+#### Startup Optimization
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
+
 **Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
 
 ### Memory Optimization
@@ -107,19 +113,19 @@ The following tables show which models support each feature:
 | **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
 | **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
 | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
 | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
 
 > Notes:
@@ -179,6 +185,59 @@ The following tables show which models support each feature:
     6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
 
 
+## Multi-Thread Weight Loading
+
+Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
+
+This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
+
+### Configuration
+
+| Parameter | CLI Flag | Default | Description |
+|-----------|----------|---------|-------------|
+| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
+| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
+
+!!! tip
+    The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
+
+### Online Serving
+
+```bash
+# Default (multi-thread enabled, 4 threads)
+vllm serve Qwen/Qwen-Image --omni --port 8091
+
+# Custom thread count
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
+
+# Disable multi-thread loading
+vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
+```
+
+### Offline Inference
+
+```python
+from vllm_omni import Omni
+
+# Default (multi-thread enabled, 4 threads)
+omni = Omni(model="Qwen/Qwen-Image")
+
+# Custom thread count
+omni = Omni(
+    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+    num_weight_load_threads=8,
+)
+```
+
+### Benchmarks
+
+Measured on NVIDIA H800:
+
+| Model | Before | After | Speedup |
+|-------|--------|-------|---------|
+| **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
+| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
+
 ## Learn More
 
 **Cache Acceleration:**
@@ -204,6 +263,10 @@ The following tables show which models support each feature:
 
 - **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
 
+**Startup Optimization:**
+
+- **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
+
 **Advanced Topics:**
 
 - **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance

examples/offline_inference/dynin_omni/README.md

@@ -0,0 +1,110 @@
+# Dynin-Omni Offline End2End Example
+
+This folder contains a unified offline inference entrypoint:
+
+- `end2end.py`
+
+## 1. Environment Setup
+
+Run from repository root:
+
+```bash
+cd <REPO_ROOT>
+```
+
+If needed, install this repo in editable mode:
+
+```bash
+pip install -e .
+```
+
+## 2. Extra Dependencies (EMOVA)
+
+Install the following packages for EMOVA-related components:
+
+```bash
+pip install \
+  "phonemizer==3.3.0" \
+  "Unidecode==1.4.0" \
+  "hydra-core==1.3.2" \
+  "pytorch-lightning==1.1.0" \
+  "wget==3.2" \
+  "wrapt==2.1.1" \
+  "onnx==1.20.1" \
+  "frozendict==2.4.7" \
+  "inflect==7.5.0" \
+  "braceexpand==0.1.7" \
+  "webdataset==1.0.2" \
+  "torch-stft==0.1.4" \
+  "editdistance==0.8.1"
+```
+
+## 3. Hardware and VRAM Requirements
+
+This example uses a 3-stage pipeline on one GPU by default
+([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
+
+- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
+- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
+- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
+
+### Requested GPU Memory Budget from `gpu_memory_utilization`
+
+| Stage | Utilization | A100 80GB | H200 141GB |
+| :-- | :-- | :-- | :-- |
+| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
+| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
+
+### Observed Runtime Signal (from your log)
+
+- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
+- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
+- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
+
+### GPU Compatibility
+
+- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
+- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
+
+## 4. End2End Run Examples
+
+```bash
+# t2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# i2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
+
+# s2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
+
+# t2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# v2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
+
+# i2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
+
+# t2s
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+```
+
+## 5. Notes
+
+- Outputs are saved under task-specific directories in `/tmp` by default.
+- You can override output path with `--output-dir`.
+- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
+- If you see the warning
+  `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
+  reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).

examples/offline_inference/image_to_image/image_to_image.md

@@ -51,5 +51,6 @@ Key arguments:
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--strength`: **Z-Image only** - controls the denoising start timestep for I2I (default: 0.6). Range: [0.0, 1.0]. Lower values preserve more of the original image; higher values allow more creative changes.
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -146,7 +146,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+        help="Sample rate for audio output when saved (default: 24000).",
     )
     parser.add_argument(
         "--cache-backend",
@@ -471,15 +471,9 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
-    use_ltx2_export = is_ltx2
-    encode_video = None
-    if use_ltx2_export:
-        try:
-            from diffusers.pipelines.ltx2.export_utils import encode_video
-        except ImportError:
-            encode_video = None
+    if audio is not None:
+        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 
-    if use_ltx2_export and encode_video is not None:
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
         elif isinstance(video_array, np.ndarray):
@@ -490,25 +484,24 @@ def main():
         if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
             frames_np = frames_np[..., :3]
 
-        audio_out = None
-        if audio is not None:
-            if isinstance(audio, list):
-                audio = audio[0] if audio else None
-            if isinstance(audio, np.ndarray):
-                audio = torch.from_numpy(audio)
-            if isinstance(audio, torch.Tensor):
-                audio_out = audio
-                if audio_out.dim() > 1:
-                    audio_out = audio_out[0]
-                audio_out = audio_out.float().cpu()
+        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
 
-        encode_video(
-            frames_np,
-            fps=fps,
-            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
-            output_path=str(output_path),
+        audio_np = audio
+        if isinstance(audio_np, list):
+            audio_np = audio_np[0] if audio_np else None
+        if isinstance(audio_np, torch.Tensor):
+            audio_np = audio_np.detach().cpu().float().numpy()
+        if isinstance(audio_np, np.ndarray):
+            audio_np = np.squeeze(audio_np).astype(np.float32)
+
+        video_bytes = mux_video_audio_bytes(
+            frames_u8,
+            audio_np,
+            fps=float(fps),
+            audio_sample_rate=args.audio_sample_rate,
         )
+        with open(str(output_path), "wb") as f:
+            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")

examples/offline_inference/magi_human/end2end.py

@@ -94,16 +94,21 @@ def main():
             print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
 
             audio_waveform = None
-            if hasattr(first, "multimodal_output") and first.multimodal_output:
-                audio_waveform = first.multimodal_output.get("audio")
+            mm = first.multimodal_output or {}
+            if mm:
+                audio_waveform = mm.get("audio")
                 if audio_waveform is not None:
                     print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
 
+            output_fps = float(mm.get("fps", 25))
+            output_sr = int(mm.get("audio_sample_rate", 24000))
+            print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
+
             video_bytes = mux_video_audio_bytes(
                 video_frames,
                 audio_waveform,
-                fps=25.0,
-                audio_sample_rate=44100,
+                fps=output_fps,
+                audio_sample_rate=output_sr,
             )
             with open(args.output, "wb") as f:
                 f.write(video_bytes)

examples/offline_inference/text_to_image/text_to_image.py

@@ -376,7 +376,7 @@ def main():
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
-    print(f"  CPU offload: {args.enable_cpu_offload}")
+    print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
     print(f"  Image size: {args.width}x{args.height}")
     if args.lora_path:
         print(f"  LoRA: scale={args.lora_scale}")

examples/offline_inference/text_to_video/text_to_video.py

@@ -160,7 +160,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+        help="Sample rate for audio output when saved (default: 24000).",
     )
     parser.add_argument(
         "--vae-patch-parallel-size",
@@ -430,17 +430,8 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
-    use_ltx2_export = False
-    if args.model and "ltx" in str(args.model).lower():
-        use_ltx2_export = True
     if audio is not None:
-        use_ltx2_export = True
-
-    if use_ltx2_export:
-        try:
-            from diffusers.pipelines.ltx2.export_utils import encode_video
-        except ImportError:
-            raise ImportError("diffusers is required for LTX2 encode_video.")
+        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
@@ -449,28 +440,24 @@ def main():
         else:
             frames_np = np.asarray(video_array)
 
-        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
-        video_tensor = torch.from_numpy(frames_u8)
+        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
 
-        audio_out = None
-        if audio is not None:
-            if isinstance(audio, list):
-                audio = audio[0] if audio else None
-            if isinstance(audio, np.ndarray):
-                audio = torch.from_numpy(audio)
-            if isinstance(audio, torch.Tensor):
-                audio_out = audio
-                if audio_out.dim() > 1:
-                    audio_out = audio_out[0]
-                audio_out = audio_out.float().cpu()
+        audio_np = audio
+        if isinstance(audio_np, list):
+            audio_np = audio_np[0] if audio_np else None
+        if isinstance(audio_np, torch.Tensor):
+            audio_np = audio_np.detach().cpu().float().numpy()
+        if isinstance(audio_np, np.ndarray):
+            audio_np = np.squeeze(audio_np).astype(np.float32)
 
-        encode_video(
-            video_tensor,
-            fps=args.fps,
-            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
-            output_path=str(output_path),
+        video_bytes = mux_video_audio_bytes(
+            frames_u8,
+            audio_np,
+            fps=float(args.fps),
+            audio_sample_rate=args.audio_sample_rate,
         )
+        with open(str(output_path), "wb") as f:
+            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")

examples/online_serving/dynin_omni/README.md

@@ -0,0 +1,97 @@
+# Dynin-Omni Online Serving Example
+
+## Installation
+
+Please refer to [README.md](../../../README.md).
+
+## Launch the Server
+
+First, find the `transformers_modules` path:
+
+```bash
+python - <<'PY'
+from transformers.utils.hub import HF_MODULES_CACHE
+print(HF_MODULES_CACHE)
+PY
+```
+
+Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
+
+```bash
+export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
+export HF_MODULES_CACHE=<transformers_modules_path>
+```
+
+Run from repository root:
+
+```bash
+vllm-omni serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+If `vllm-omni` is not in PATH, run:
+
+```bash
+PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+Wait until the server logs show both `All stages initialized successfully` and
+`Application startup complete.` before sending requests.
+
+## Send Requests via Python Client
+
+Move to the example directory:
+
+```bash
+cd examples/online_serving/dynin_omni
+```
+
+### Text -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2i \
+  --prompt "A realistic indoor living room with natural daylight."
+```
+
+### Image -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type i2i \
+  --image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
+  --prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
+```
+
+### Text -> Speech
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2s \
+  --prompt "Hello. This is Dynin-omni."
+```
+
+## CLI Arguments
+
+- `--query-type` (`t2i|t2s|i2i`)
+- `--model` (default: `snu-aidas/Dynin-Omni`)
+- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
+- `--prompt` (custom text)
+- `--image-path` (required for `i2i`)
+- `--modalities` (optional output modalities override)
+- `--output-dir` (default: `/tmp/dynin_online_outputs`)
+
+## Notes
+
+- This client currently supports only `t2i`, `t2s`, and `i2i`.
+- `t2t` is intentionally not exposed in this online example.
+- This example intentionally uses the OpenAI-compatible chat completion endpoint.
+- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
+- Outputs are saved under `/tmp/dynin_online_outputs` by default.
+- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
+- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.

examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import mimetypes
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
+DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
+
+QUERY_CHOICES = ("t2i", "t2s", "i2i")
+DEFAULT_PROMPT_BY_QUERY = {
+    "t2i": "A high quality detailed living room interior photo.",
+    "t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
+    "i2i": "Transform this image into a realistic indoor living room while preserving layout.",
+}
+DEFAULT_MODALITIES_BY_QUERY = {
+    "t2i": ["image"],
+    "t2s": ["audio"],
+    "i2i": ["image"],
+}
+OFFLINE_PARITY_STAGE_COUNT = 3
+OFFLINE_PARITY_STAGE_SAMPLING = {
+    "max_tokens": 1,
+    "temperature": 0.0,
+    "top_p": 1.0,
+    "detokenize": False,
+}
+
+
+def _infer_mime_type(path: Path) -> str:
+    mime_type, _ = mimetypes.guess_type(str(path))
+    return mime_type or "application/octet-stream"
+
+
+def _encode_file_as_data_url(path: Path) -> str:
+    mime_type = _infer_mime_type(path)
+    raw = path.read_bytes()
+    encoded = base64.b64encode(raw).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def _to_image_url(path_or_url: str) -> str:
+    value = str(path_or_url)
+    if value.startswith(("http://", "https://", "data:image/")):
+        return value
+    path = Path(value).expanduser().resolve()
+    if not path.exists():
+        raise FileNotFoundError(f"Image file not found: {path}")
+    return _encode_file_as_data_url(path)
+
+
+def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
+    if query_type == "t2i":
+        return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
+
+    if query_type == "t2s":
+        return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
+
+    if query_type == "i2i":
+        if not image_path:
+            raise ValueError("--image-path is required for query type i2i")
+        return [
+            {"type": "text", "text": f"<|i2i|> {prompt}"},
+            {"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
+        ]
+
+    raise ValueError(f"Unsupported query_type: {query_type}")
+
+
+def _collect_text_from_content(content: Any) -> list[str]:
+    texts: list[str] = []
+    if isinstance(content, str):
+        stripped = content.strip()
+        if stripped:
+            texts.append(stripped)
+        return texts
+
+    if isinstance(content, dict):
+        for key in ("text", "content", "value", "output_text"):
+            text_value = content.get(key)
+            if isinstance(text_value, str) and text_value.strip():
+                texts.append(text_value.strip())
+        return texts
+
+    if isinstance(content, list):
+        for item in content:
+            texts.extend(_collect_text_from_content(item))
+        return texts
+
+    content_text = getattr(content, "text", None)
+    if isinstance(content_text, str) and content_text.strip():
+        texts.append(content_text.strip())
+    content_value = getattr(content, "content", None)
+    if isinstance(content_value, str) and content_value.strip():
+        texts.append(content_value.strip())
+    output_text = getattr(content, "output_text", None)
+    if isinstance(output_text, str) and output_text.strip():
+        texts.append(output_text.strip())
+    return texts
+
+
+def _extract_text_outputs(chat_completion: Any) -> list[str]:
+    texts: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        texts.extend(_collect_text_from_content(content))
+        reasoning_content = getattr(message, "reasoning_content", None)
+        if isinstance(reasoning_content, str) and reasoning_content.strip():
+            texts.append(reasoning_content.strip())
+        choice_text = getattr(choice, "text", None)
+        if isinstance(choice_text, str) and choice_text.strip():
+            texts.append(choice_text.strip())
+    top_level_output_text = getattr(chat_completion, "output_text", None)
+    if isinstance(top_level_output_text, str) and top_level_output_text.strip():
+        texts.append(top_level_output_text.strip())
+    return texts
+
+
+def _extract_image_data_urls(chat_completion: Any) -> list[str]:
+    urls: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        if not isinstance(content, list):
+            continue
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            if item.get("type") != "image_url":
+                continue
+            image_url = (item.get("image_url") or {}).get("url")
+            if isinstance(image_url, str) and image_url.startswith("data:image"):
+                urls.append(image_url)
+    return urls
+
+
+def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
+    payloads: list[bytes] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        message_audio = getattr(message, "audio", None)
+        if message_audio is None:
+            continue
+        data_b64 = getattr(message_audio, "data", None)
+        if isinstance(data_b64, str) and data_b64:
+            try:
+                payloads.append(base64.b64decode(data_b64))
+            except Exception:
+                continue
+    return payloads
+
+
+def _decode_data_url(data_url: str) -> tuple[bytes, str]:
+    header, data = data_url.split(",", 1)
+    mime_type = "image/png"
+    if ";" in header and ":" in header:
+        mime_type = header.split(":", 1)[1].split(";", 1)[0]
+    return base64.b64decode(data), mime_type
+
+
+def _image_extension_from_mime(mime_type: str) -> str:
+    if mime_type == "image/jpeg":
+        return ".jpg"
+    if mime_type == "image/webp":
+        return ".webp"
+    if mime_type == "image/gif":
+        return ".gif"
+    return ".png"
+
+
+def _save_outputs(
+    *,
+    query_type: str,
+    chat_completion: Any,
+    output_dir: Path,
+) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+
+    text_outputs = _extract_text_outputs(chat_completion)
+    image_data_urls = _extract_image_data_urls(chat_completion)
+    audio_payloads = _extract_audio_payloads(chat_completion)
+
+    if text_outputs:
+        text_path = output_dir / f"{query_type}_{stamp}.txt"
+        text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
+        print(f"[dynin-online] text saved: {text_path}")
+        print(text_outputs[0])
+
+    for idx, image_url in enumerate(image_data_urls):
+        image_bytes, mime_type = _decode_data_url(image_url)
+        ext = _image_extension_from_mime(mime_type)
+        image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
+        image_path.write_bytes(image_bytes)
+        print(f"[dynin-online] image saved: {image_path}")
+
+    for idx, audio_bytes in enumerate(audio_payloads):
+        audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
+        audio_path.write_bytes(audio_bytes)
+        print(f"[dynin-online] audio saved: {audio_path}")
+
+    if not text_outputs and not image_data_urls and not audio_payloads:
+        print("[dynin-online] no output extracted from response")
+        raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
+        try:
+            if hasattr(chat_completion, "model_dump_json"):
+                serialized = chat_completion.model_dump_json(indent=2)
+            else:
+                if hasattr(chat_completion, "model_dump"):
+                    raw_payload: Any = chat_completion.model_dump(mode="json")
+                else:
+                    raw_payload = chat_completion
+                try:
+                    serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
+                except Exception:
+                    serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
+            raw_path.write_text(serialized + "\n", encoding="utf-8")
+            print(f"[dynin-online] raw response saved: {raw_path}")
+        except Exception:
+            pass
+
+
+def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
+    return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
+
+
+def run_request(args: argparse.Namespace) -> None:
+    from openai import OpenAI
+
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=f"http://{args.host}:{args.port}/v1",
+    )
+    prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
+    user_content = _build_user_content(
+        query_type=args.query_type,
+        prompt=prompt,
+        image_path=args.image_path,
+    )
+    if args.modalities:
+        modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
+    else:
+        modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
+
+    extra_body = {
+        "sampling_params_list": _build_offline_parity_sampling_params_list(),
+    }
+    chat_completion = client.chat.completions.create(
+        model=args.model,
+        messages=[{"role": "user", "content": user_content}],
+        modalities=modalities,
+        extra_body=extra_body,
+    )
+    _save_outputs(
+        query_type=args.query_type,
+        chat_completion=chat_completion,
+        output_dir=Path(args.output_dir).expanduser(),
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="t2i",
+        choices=QUERY_CHOICES,
+        help="Dynin query type",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default=DEFAULT_MODEL,
+        help="Model name/path",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host/IP of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8091,
+        help="Port of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="",
+        help="Custom prompt text",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Image path/URL for i2i",
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default="",
+        help="Comma-separated output modalities override (e.g., text,image,audio)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default=DEFAULT_OUTPUT_DIR,
+        help="Directory to save outputs",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+    run_request(args)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/image_to_image/README.md

@@ -314,6 +314,7 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)            |
 | `negative_prompt`        | str   | None    | Negative prompt                       |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate          |
+| `strength`               | float | 0.6     | **Z-Image only** - Denoising start timestep for I2I. Range: [0.0, 1.0]. Lower preserves more of original image. |
 | `layers`                 | int   | 4       | Number of layers (Qwen-Image-Layered) |
 | `resolution`             | int   | 640     | Resolution, 640 or 1024 (Qwen-Image-Layered) |
 

tests/benchmarks/metrics/test_metrics.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for metrics.py
+"""
+
+import pytest
+from vllm.benchmarks.serve import TaskType
+
+from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
+from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+
+
+def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput:
+    """Build a minimal successful MixRequestFuncOutput for metrics aggregation."""
+    output = MixRequestFuncOutput()
+    output.success = True
+    output.prompt_len = prompt_len
+    output.output_tokens = output_tokens
+    output.generated_text = "x" * output_tokens
+    output.ttft = 0.1
+    output.text_latency = 1.0
+    output.latency = 1.0
+    output.start_time = 0.0
+    output.itl = [0.1] * max(output_tokens - 1, 0)
+    output.audio_ttfp = 0.0
+    output.audio_rtf = 0.0
+    output.audio_duration = 0.0
+    output.audio_frames = 0
+    output.input_audio_duration = 0.0
+    output.error = ""
+    return output
+
+
+# ============================================================================
+# total_input Tests
+# ============================================================================
+
+
+def test_total_input_aggregated_from_output_prompt_len():
+    """Test that total_input sums outputs[i].prompt_len, not input_requests[i].prompt_len."""
+    outputs = [_make_output(4992), _make_output(3000)]
+
+    metrics, _ = calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=10.0,
+        tokenizer=None,
+        selected_percentiles=[99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=[],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    assert metrics.total_input == 7992, (
+        "total_input should aggregate from outputs[i].prompt_len to reflect the true multimodal input token count"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

tests/benchmarks/patch/test_patch.py

@@ -574,5 +574,59 @@ class TestTextLatencyAttribute:
         )
 
 
+# ============================================================================
+# prompt_len Tests
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
+    # Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=100,
+        output_len=20,
+    )
+
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+            }
+        ),
+        # Final usage chunk emitted because stream_options.include_usage=True.
+        create_sse_chunk(
+            {
+                "choices": [],
+                "usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.prompt_len == 4992, (
+        "prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
+    )
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/conftest.py

@@ -16,6 +16,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
 import concurrent.futures
+import contextlib
 import gc
 import multiprocessing
 import socket
@@ -52,6 +53,7 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+
 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
 PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
@@ -337,10 +339,10 @@ def log_test_name_before_test(request):
 
 def _run_pre_test_cleanup(enable_force=False):
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
-        print("GPU cleanup disabled")
+        print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n")
         return
 
-    print("Pre-test GPU status:")
+    print("\nPre-test GPU status:")
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
@@ -1087,6 +1089,22 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
     return merged
 
 
+@contextlib.contextmanager
+def _serialize_whisper_small_model_download():
+    """Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``)."""
+    import fcntl
+
+    lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock"
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    f = open(lock_path, "a+b")
+    try:
+        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+        yield
+    finally:
+        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+        f.close()
+
+
 def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
@@ -1107,7 +1125,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
     else:
         use_accelerator = False
         device = "cpu"
-    model = whisper.load_model("small", device=device)
+    with _serialize_whisper_small_model_download():
+        model = whisper.load_model("small", device=device)
     try:
         text = model.transcribe(
             output_path,
@@ -1126,8 +1145,7 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
 
 
 def convert_audio_file_to_text(output_path: str) -> str:
-    """Convert an audio file to text in an isolated subprocess."""
-    # Import locally to avoid impacting test module import time.
+    """Convert an audio file to text in an isolated subprocess (spawn)."""
     ctx = multiprocessing.get_context("spawn")
     with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
         future = executor.submit(_whisper_transcribe_in_current_process, output_path)

tests/core/sched/test_generation_scheduler_restore.py

@@ -9,6 +9,10 @@ those requests are permanently orphaned.
 import unittest
 from collections import deque
 
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class FakeAdapter:
     """Minimal mock of OmniChunkTransferAdapter tracking restore calls."""

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -0,0 +1,217 @@
+[
+    {
+        "test_name": "test_ltx2_baseline_eager",
+        "description": "Single-device baseline with enforce-eager (no torch.compile)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_torch_compile",
+        "description": "Single-device with torch.compile (default, no enforce-eager)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cfg2_eager",
+        "description": "CFG-parallel=2 with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cfg2_compile",
+        "description": "CFG-parallel=2 with torch.compile",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cache_dit_eager",
+        "description": "CacheDiT with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cache-backend": "cache_dit",
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    }
+]

tests/diffusion/cache/test_teacache_extractors.py

@@ -21,12 +21,13 @@ from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -113,6 +114,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
     def get_sample_inputs(self, sample_inputs):
         return sample_inputs
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_modulated_input_shape(self, flux2_klein_module, sample_inputs):
         """Test that modulated_input has correct shape matching the model's inner_dim.
 
@@ -126,16 +128,19 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         inner_dim = flux2_klein_module.inner_dim
         assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_run_transformer_blocks_callable(self, flux2_klein_module, sample_inputs):
         """Test that run_transformer_blocks is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.run_transformer_blocks)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_postprocess_callable(self, flux2_klein_module, sample_inputs):
         """Test that postprocess is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.postprocess)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_extra_states_contains_full_transformer(self, flux2_klein_module, sample_inputs):
         """Test that extra_states contains run_flux2_full_transformer_with_single."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
@@ -154,6 +159,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         assert context is not None
         assert context.temb is not None
 
+    @pytest.mark.cpu
     def test_invalid_module_raises_error(self):
         """Test that invalid module without transformer_blocks raises ValueError."""
         invalid_module = Mock()

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
+    def __init__(self, *, dtype: torch.dtype):
+        torch.nn.Module.__init__(self)
+        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
+
+
+def test_wan_vae_execution_context_handles_fp32():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.float32
+
+
+def test_wan_vae_execution_context_handles_bf16():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.bfloat16
+
+
+def test_wan_vae_execution_context_uses_platform_autocast(mocker):
+    sentinel = object()
+    platform = mocker.Mock()
+    platform.create_autocast_context.return_value = sentinel
+    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
+
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+
+    assert model._execution_context() is sentinel
+    platform.create_autocast_context.assert_called_once_with(
+        device_type=model.dummy_weight.device.type,
+        dtype=torch.bfloat16,
+        enabled=True,
+    )

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -11,6 +11,8 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     TileTask,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class E2EOperator:
     """tiles with (2, 3) -- (H,W)"""

tests/diffusion/distributed/test_ulysses_uaa_perf.py

@@ -17,6 +17,7 @@ import pytest
 import torch
 import torch.distributed as dist
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.attention.parallel.ulysses import (
     _all_gather_int,
     _ulysses_all_to_all_any_o,
@@ -69,6 +70,8 @@ PERF_CASES: list[_PerfCase] = [
 
 
 @pytest.mark.parametrize("case", PERF_CASES)
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
 def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
     available_gpus = current_omni_platform.get_device_count()
     if available_gpus < case.world_size:

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BAGEL trajectory recording in the denoising loop."""
+
+import types
+from dataclasses import dataclass
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (
+    Bagel,
+    NaiveCache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_TOKENS = 8
+HIDDEN_DIM = 16
+NUM_TIMESTEPS = 5
+# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1
+EXPECTED_STEPS = NUM_TIMESTEPS - 1
+
+
+def _make_mock_bagel():
+    """Create a mock Bagel with forward returning constant velocity."""
+    mock = MagicMock(spec=Bagel)
+    mock._sp_size = 1
+
+    # forward returns a small constant velocity so x_t changes each step
+    def fake_forward(self, x_t, **kwargs):
+        return torch.ones_like(x_t) * 0.1
+
+    mock.forward = types.MethodType(fake_forward, mock)
+    # _merge_naive_caches is called in the batched CFG path
+    mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock)
+
+    # Bind the real generate_image to our mock
+    mock.generate_image = types.MethodType(Bagel.generate_image, mock)
+    return mock
+
+
+def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False):
+    """Tensor arguments for generate_image.
+
+    Args:
+        cfg: If True, enable batched CFG path (cfg_text_scale > 1.0).
+    """
+    seq_len = num_tokens + 2  # packed_seqlens includes 2 extra tokens
+    base = dict(
+        packed_text_ids=torch.zeros(2, dtype=torch.long),
+        packed_text_indexes=torch.tensor([0, 1], dtype=torch.long),
+        packed_init_noises=torch.randn(num_tokens, hidden_dim),
+        packed_vae_position_ids=torch.arange(num_tokens, dtype=torch.long),
+        packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
+        packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
+        packed_position_ids=torch.arange(seq_len, dtype=torch.long),
+        packed_indexes=torch.arange(seq_len, dtype=torch.long),
+        past_key_values=NaiveCache(1),
+        key_values_lens=torch.tensor([0], dtype=torch.int),
+        packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
+        num_timesteps=NUM_TIMESTEPS,
+        timestep_shift=1.0,
+        cfg_text_scale=1.0,
+        cfg_img_scale=1.0,
+    )
+    if cfg:
+        base |= dict(
+            cfg_text_scale=4.0,
+            cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
+            cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
+            cfg_text_past_key_values=NaiveCache(1),
+            cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
+            cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
+        )
+    return base
+
+
+@pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"])
+def bagel_and_args(request):
+    """Mock Bagel instance and generate_image arguments.
+
+    Parametrized over CFG mode so every test runs on both the no-CFG
+    and batched-CFG code paths.
+    """
+    cfg = request.param
+    with patch(
+        "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
+        return_value=1,
+    ):
+        yield _make_mock_bagel(), _make_generate_args(cfg=cfg)
+
+
+class TestTrajectoryRecording:
+    """Tests for trajectory latent/timestep recording in generate_image."""
+
+    def test_trajectory_disabled_returns_none(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        unpacked, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=False
+        )
+
+        assert isinstance(unpacked, (list, tuple))
+        assert len(unpacked) == 1  # one sequence
+        assert trajectory_latents is None
+        assert trajectory_timesteps is None
+        assert trajectory_log_probs is None
+
+    def test_trajectory_enabled_returns_correct_count(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True
+        )
+
+        assert trajectory_latents is not None
+        assert trajectory_timesteps is not None
+        assert len(trajectory_latents) == EXPECTED_STEPS
+        assert len(trajectory_timesteps) == EXPECTED_STEPS
+        # log_probs is None without a scheduler (default ODE path)
+        assert trajectory_log_probs is None
+
+    def test_trajectory_latents_shape_matches_input(self, bagel_and_args):
+        bagel, args = bagel_and_args
+        expected_shape = args["packed_init_noises"].shape
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i, lat in enumerate(trajectory_latents):
+            assert lat.shape == expected_shape, f"Step {i}: expected {expected_shape}, got {lat.shape}"
+
+    def test_trajectory_latents_are_distinct(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i in range(1, len(trajectory_latents)):
+            assert not torch.equal(trajectory_latents[i], trajectory_latents[i - 1]), (
+                f"Steps {i - 1} and {i} should differ"
+            )
+
+    def test_trajectory_timesteps_are_decreasing(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i in range(1, len(trajectory_timesteps)):
+            assert trajectory_timesteps[i] < trajectory_timesteps[i - 1], (
+                f"Timestep {i} ({trajectory_timesteps[i]:.4f}) should be less than "
+                f"timestep {i - 1} ({trajectory_timesteps[i - 1]:.4f})"
+            )
+
+    def test_trajectory_final_latent_matches_output(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        unpacked, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        # Reconstruct the full final latent from unpacked pieces
+        final_latent = torch.cat(unpacked, dim=0)
+        assert torch.allclose(trajectory_latents[-1], final_latent, atol=1e-6), (
+            "Last trajectory latent should match the final output"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Mock scheduler for log-prob tests
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _MockStepOutput:
+    prev_sample: torch.Tensor
+    log_prob: torch.Tensor
+
+
+class _MockScheduler:
+    """Minimal scheduler: Euler step + constant log-prob per step."""
+
+    def step(self, model_output, sigma, sample, dt, **kwargs):
+        prev_sample = sample - model_output * dt
+        log_prob = torch.tensor(-1.0)
+        return _MockStepOutput(prev_sample=prev_sample, log_prob=log_prob)
+
+
+class TestTrajectoryLogProbs:
+    """Tests for log-prob recording when a scheduler is provided."""
+
+    @pytest.fixture()
+    def bagel_scheduler_args(self):
+        with patch(
+            "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
+            return_value=1,
+        ):
+            yield _make_mock_bagel(), _make_generate_args(), _MockScheduler()
+
+    def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args):
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True, scheduler=scheduler
+        )
+
+        assert trajectory_log_probs is not None
+        assert len(trajectory_log_probs) == EXPECTED_STEPS
+
+    def test_log_probs_are_finite(self, bagel_scheduler_args):
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True, scheduler=scheduler
+        )
+
+        for i, lp in enumerate(trajectory_log_probs):
+            assert torch.isfinite(lp).all(), f"Step {i}: log_prob is not finite"
+
+    def test_log_probs_none_without_scheduler(self, bagel_scheduler_args):
+        bagel, args, _ = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
+
+        assert trajectory_log_probs is None
+
+    def test_scheduler_updates_latents(self, bagel_scheduler_args):
+        """Verify the scheduler's prev_sample is used (not the raw Euler step)."""
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, traj_with_sched, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=scheduler)
+        _, traj_without, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
+
+        # Mock scheduler does the same Euler step, so latents should match
+        for i in range(len(traj_with_sched)):
+            assert torch.allclose(traj_with_sched[i], traj_without[i], atol=1e-5), (
+                f"Step {i}: scheduler and ODE paths should produce same latents"
+            )

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.models.flux2.flux2_transformer import (
     Flux2PosEmbed,
     Flux2Transformer2DModel,
@@ -24,6 +25,8 @@ def setup_tp_group():
 class TestFlux2TransformerWeightLoading:
     """Test Flux2Transformer weight loading functionality"""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_weight_loading_tp2(self, setup_tp_group):
         """Verify weights load correctly with TP=2"""
         # Prepare test data
@@ -78,6 +81,8 @@ class TestFlux2TransformerWeightLoading:
 class TestFlux2RopePositionEmbedding:
     """Test Flux2 RoPE position embedding functionality"""
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_rope_position_embedding(self):
         """Verify RoPE produces correct embeddings for 4D coordinates"""
         # Prepare test data - use model default configuration
@@ -132,6 +137,8 @@ class TestFlux2RopePositionEmbedding:
 class TestFlux2PackedModuleMapping:
     """Test Flux2 packed module mapping functionality"""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_packed_module_mapping(self, setup_tp_group):
         """Verify to_qkv packing matches HF checkpoint"""
         model = Flux2Transformer2DModel(
@@ -208,6 +215,8 @@ class TestFlux2PackedModuleMapping:
             f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}"
         )
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_packed_mapping_edge_cases(self, setup_tp_group):
         """Test edge cases for packed mapping"""
         model = Flux2Transformer2DModel(

tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py

@@ -0,0 +1,58 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
+    pipeline = object.__new__(LTX2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae_temporal_compression_ratio = 4
+    pipeline.audio_vae_mel_compression_ratio = 4
+    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
+    # Mock audio_vae with identity normalization (mean=0, std=1) so
+    # _normalize_audio_latents is a no-op and test values are preserved.
+    pipeline.audio_vae = SimpleNamespace(
+        latents_mean=torch.tensor(0.0),
+        latents_std=torch.tensor(1.0),
+    )
+    return pipeline
+
+
+def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
+    pipeline = _make_pipeline(sequence_parallel_size=4)
+    latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+
+    padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        num_mel_bins=8,
+        audio_latent_length=10,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=latents,
+    )
+
+    assert original_num_frames == 10
+    assert padded_num_frames == 12
+    assert padded.shape == (1, 12, 4)
+    torch.testing.assert_close(padded[:, :10], latents)
+    torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
+
+
+def test_unpad_audio_latents_restores_original_frames_before_unpack():
+    pipeline = _make_pipeline()
+    original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+    padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
+
+    unpadded = pipeline._unpad_audio_latents(padded, 10)
+    unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
+    expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
+
+    assert unpacked.shape == (1, 2, 10, 2)
+    assert not (unpacked == 999.0).any()
+    torch.testing.assert_close(unpacked, expected)

tests/diffusion/quantization/test_fp8_config.py

@@ -5,7 +5,7 @@
 import pytest
 from torch import nn
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_build_quant_config_fp8():

tests/diffusion/test_diffusion_model_runner.py

@@ -8,9 +8,10 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from tests.utils import hardware_test
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.diffusion]
 
 
 @contextmanager
@@ -64,6 +65,8 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     return runner
 
 
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
     """Guard cache diagnostics with runtime backend state to avoid stale-config crashes."""
     runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
@@ -84,6 +87,8 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
     assert cache_summary_calls == []
 
 
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
     class _EnabledCacheBackend:
         def is_enabled(self):
@@ -107,6 +112,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
     assert cache_summary_calls == [(runner.pipeline, True)]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):

tests/diffusion/test_stage_diffusion_proc.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
+    async def run_test():
+        captured = {}
+
+        def step(request):
+            captured["request"] = request
+            return [
+                SimpleNamespace(
+                    images=["img-1"],
+                    _multimodal_output={},
+                    _custom_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
+                    final_output_type="image",
+                ),
+                SimpleNamespace(
+                    images=["img-2"],
+                    _multimodal_output={},
+                    _custom_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
+                    final_output_type="image",
+                ),
+            ]
+
+        proc = object.__new__(StageDiffusionProc)
+        proc._engine = SimpleNamespace(step=step)
+        proc._executor = ThreadPoolExecutor(max_workers=1)
+
+        try:
+            result = await proc._process_batch_request(
+                request_id="req-parent",
+                prompts=["hello", "world"],
+                sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
+                kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}},
+            )
+        finally:
+            proc._executor.shutdown(wait=True)
+
+        request = captured["request"]
+        assert request.request_id == "req-parent"
+        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
+        assert result.request_id == "req-parent"
+        assert result.images == ["img-1", "img-2"]
+
+    asyncio.run(run_test())

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -9,7 +9,7 @@ from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
 
-# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_basic_serialization():
@@ -120,3 +120,61 @@ def test_get_invalid_metadata(shm_connector):
 
     result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
     assert result is None
+
+
+def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
+
+    class _FakePool:
+        is_cuda = False
+
+        def pin_memory(self):
+            return self
+
+        def data_ptr(self):
+            return 1234
+
+    class _FakeTransferEngine:
+        def initialize(self, host, mode, protocol, device_name):
+            self.host = host
+            self.mode = mode
+            self.protocol = protocol
+            self.device_name = device_name
+            return 0
+
+        def get_rpc_port(self):
+            return 23456
+
+        def register_memory(self, base_ptr, pool_size):
+            del base_ptr, pool_size
+            return 0
+
+        def unregister_memory(self, base_ptr):
+            del base_ptr
+            return 0
+
+    monkeypatch.setattr(mooncake_module, "TransferEngine", _FakeTransferEngine)
+    monkeypatch.setattr(mooncake_module.torch, "empty", lambda *args, **kwargs: _FakePool())
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_get_local_ip",
+        lambda self: "10.20.30.40",
+    )
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_zmq_listener_loop",
+        lambda self: self._listener_ready.set(),
+    )
+
+    connector = mooncake_module.MooncakeTransferEngineConnector(
+        {
+            "zmq_port": 50051,
+            "memory_pool_size": 4096,
+        }
+    )
+    try:
+        assert connector.host == "10.20.30.40"
+        assert connector.engine.host == "10.20.30.40"
+        assert connector.get_connection_info()["host"] == "10.20.30.40"
+    finally:
+        connector.close()

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,8 +1,14 @@
+import json
+import struct
+
+import numpy as np
 import pytest
 import torch
 
+import vllm_omni.distributed.omni_connectors.kv_transfer_manager as kv_transfer_manager_module
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    KVCacheTransferData,
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
@@ -60,6 +66,35 @@ def common_constants():
     }
 
 
+def _decode_stored_payload(data):
+    if isinstance(data, torch.Tensor) and data.dtype == torch.uint8 and data.dim() == 1:
+        return KVCacheTransferData.from_bytes(data.cpu().numpy().tobytes())
+
+    if isinstance(data, (bytes, bytearray, memoryview)):
+        return KVCacheTransferData.from_bytes(data)
+
+    return data
+
+
+def _make_serialized_payload() -> tuple[bytes, torch.Tensor]:
+    key_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    payload = KVCacheTransferData(
+        request_id="req-payload",
+        layer_blocks={"key_cache": [key_tensor], "value_cache": [None]},
+        block_ids=[1],
+        metadata={"seq_len": 3},
+    ).to_bytes()
+    return payload, key_tensor
+
+
+def _rewrite_serialized_header(payload: bytes, mutate_header) -> bytes:
+    header_len = struct.unpack(">I", payload[:4])[0]
+    header = json.loads(payload[4 : 4 + header_len])
+    mutate_header(header)
+    new_header = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    return struct.pack(">I", len(new_header)) + new_header + payload[4 + header_len :]
+
+
 def test_manager_extraction(kv_config, mock_connector, common_constants):
     """Test extraction and sending logic in OmniKVTransferManager."""
     num_layers = common_constants["num_layers"]
@@ -95,7 +130,7 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     assert data["request_id"] == req_id
     assert "layer_blocks" in data
     assert len(data["layer_blocks"]["key_cache"]) == num_layers
@@ -106,6 +141,116 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
 
 
+def test_from_bytes_rejects_out_of_bounds_header_len():
+    payload, _ = _make_serialized_payload()
+    bad_payload = struct.pack(">I", len(payload)) + payload[4:]
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_out_of_bounds_tensor_span():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"o": 4096}))
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_unsupported_dtype():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"d": "cuda"}))
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_uses_explicit_layer_index_descriptor():
+    payload, key_tensor = _make_serialized_payload()
+    payload_with_explicit_index = _rewrite_serialized_header(
+        payload,
+        lambda header: header["td"][0].update({"n": "key_cache_extra_suffix", "i": 0}),
+    )
+
+    data = KVCacheTransferData.from_bytes(payload_with_explicit_index)
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_update_sender_info_uses_configured_source_stage():
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        stage_id=2,
+        engine_input_source=[1],
+        need_recv_cache=True,
+    )
+    manager = OmniKVTransferManager(config)
+
+    manager.update_sender_info(
+        {
+            0: {"host": "10.0.0.1", "zmq_port": 50151},
+            1: {"host": "10.0.0.2", "zmq_port": 50152},
+        }
+    )
+
+    assert manager.config.connector_config["sender_host"] == "10.0.0.2"
+    assert manager.config.connector_config["sender_zmq_port"] == 50152
+
+
+def test_clone_received_payload_tensors_breaks_buffer_alias():
+    payload, key_tensor = _make_serialized_payload()
+    raw = np.frombuffer(bytearray(payload), dtype=np.uint8)
+    data = KVCacheTransferData.from_bytes(memoryview(raw))
+
+    OmniKVTransferManager._clone_received_payload_tensors(data)
+    raw[:] = 0
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_receive_kv_cache_uses_exponential_backoff(monkeypatch):
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=0.3,
+    )
+    manager = OmniKVTransferManager(config)
+
+    class _NeverReadyConnector:
+        def get(self, **kwargs):
+            del kwargs
+            return None
+
+    manager._connector = _NeverReadyConnector()
+
+    now = {"value": 0.0}
+    sleep_intervals = []
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "time", lambda: now["value"])
+
+    def _fake_sleep(interval: float) -> None:
+        sleep_intervals.append(interval)
+        now["value"] += interval
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "sleep", _fake_sleep)
+
+    data, size = manager.receive_kv_cache_for_request("req-backoff")
+
+    assert (data, size) == (None, 0)
+    assert sleep_intervals == pytest.approx([0.01, 0.02, 0.04, 0.08, 0.16])
+
+
 def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
     """Test extraction with tuple layout."""
     num_layers = common_constants["num_layers"]
@@ -135,7 +280,7 @@ def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_const
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (seq_len, num_heads, head_dim)
     for idx in range(len(kv_caches)):
         assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
@@ -165,7 +310,7 @@ def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (2 * block_size, num_heads, head_dim)
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
     assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
@@ -254,6 +399,82 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     assert req.kv_metadata["seq_len"] == seq_len
 
 
+def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_config, mock_connector, common_constants):
+    """Batched diffusion requests must fetch KV using the parent/global request ID."""
+    num_layers = common_constants["num_layers"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    parent_req_id = common_constants["req_id"]
+
+    expected_shape = (seq_len, num_heads, head_dim)
+    key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+    value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+
+    data_to_receive = {
+        "request_id": parent_req_id,
+        "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
+        "metadata": {"seq_len": seq_len},
+        "block_ids": [],
+    }
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{parent_req_id}"
+    store_key = f"stage1->stage2:{full_request_id}"
+    mock_connector.store[store_key] = data_to_receive
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
+        request_id=parent_req_id,
+    )
+
+    success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
+
+    assert success
+    assert req.kv_metadata["seq_len"] == seq_len
+    assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
+
+
+def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_config):
+    manager = OmniKVTransferManager(kv_config)
+
+    seen = {}
+
+    def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device):
+        seen["request_id"] = request_id
+        seen["cfg_request_ids"] = cfg_request_ids
+        seen["kv_transfer_manager"] = kv_transfer_manager
+        seen["target_device"] = target_device
+        return {"cfg_text_kv_metadata": {"ok": True}}
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=["req-parent-0", "req-parent-1"],
+        request_id="req-parent",
+    )
+    req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
+
+    manager.receive_kv_cache = lambda request, target_device=None: request is req
+
+    success = manager.receive_multi_kv_cache(
+        req,
+        cfg_kv_collect_func=collect_cfg,
+        target_device=torch.device("cpu"),
+    )
+
+    assert success
+    assert seen["request_id"] == "req-parent"
+    assert seen["cfg_request_ids"] == {"cfg_text": "req-parent__cfg_text"}
+    assert seen["kv_transfer_manager"] is manager
+    assert seen["target_device"] == torch.device("cpu")
+    assert req.sampling_params.cfg_text_kv_metadata == {"ok": True}
+
+
 def test_integration_flow(common_constants):
     """Simulate extraction -> connector -> reception."""
     num_layers = common_constants["num_layers"]

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -3,12 +3,16 @@
 
 from time import time
 
+import pytest
+
 from vllm_omni.distributed.omni_coordinator import (
     InstanceInfo,
     RandomBalancer,
     StageStatus,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def test_load_balancer_select_returns_valid_index():
     """Verify RandomBalancer.select() returns a valid index for instances."""

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -12,6 +12,8 @@ from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForHub,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _bind_pub() -> tuple[zmq.Context, zmq.Socket, str]:
     ctx = zmq.Context.instance()

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -17,6 +17,8 @@ from vllm_omni.distributed.omni_coordinator import (
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
     ctx = zmq.Context.instance()

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -4,6 +4,7 @@
 import json
 import time
 
+import pytest
 import zmq
 from vllm.v1.utils import get_engine_client_zmq_addr
 
@@ -13,6 +14,8 @@ from vllm_omni.distributed.omni_coordinator import (
     StageStatus,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
     """Receive InstanceList JSON from SUB socket. Returns None on timeout."""

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -6,7 +6,8 @@
 This pipeline follows the structure of the user's reference implementation:
 - supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
 - uses an SDE scheduler that can return step logprobs
-- returns rich custom_output fields for testing
+- returns structured trajectory_* fields (latents, timesteps, log_probs)
+  consistent with the BAGEL trajectory recording design
 """
 
 from __future__ import annotations
@@ -393,10 +394,10 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
 
         return DiffusionOutput(
             output=_maybe_to_cpu(image),
+            trajectory_latents=_maybe_to_cpu(all_latents),
+            trajectory_log_probs=_maybe_to_cpu(all_log_probs),
+            trajectory_timesteps=_maybe_to_cpu(all_timesteps),
             custom_output={
-                "all_latents": _maybe_to_cpu(all_latents),
-                "all_log_probs": _maybe_to_cpu(all_log_probs),
-                "all_timesteps": _maybe_to_cpu(all_timesteps),
                 "prompt_embeds": _maybe_to_cpu(prompt_embeds),
                 "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
                 "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -1,7 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""E2E tests for AsyncOmni Qwen-Image generation flow (no Ray, no HTTP server)."""
+"""E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields.
+
+Validates that the custom Qwen-Image pipeline returns structured trajectory
+outputs (latents, timesteps, log_probs) via OmniRequestOutput's trajectory_*
+fields instead of the legacy custom_output dict.
+"""
 
 from __future__ import annotations
 
@@ -191,10 +196,17 @@ async def test_async_omni_generate_with_logprobs():
 
         _assert_valid_image_output(output)
 
-        all_log_probs = output.custom_output.get("all_log_probs")
-        assert all_log_probs is not None, "all_log_probs should be present when logprobs=True"
-        assert hasattr(all_log_probs, "shape")
-        assert all_log_probs.numel() > 0
+        assert output.trajectory_latents is not None, "trajectory_latents should be present"
+        assert hasattr(output.trajectory_latents, "shape")
+        assert output.trajectory_latents.numel() > 0
+
+        assert output.trajectory_timesteps is not None, "trajectory_timesteps should be present"
+        assert hasattr(output.trajectory_timesteps, "shape")
+        assert output.trajectory_timesteps.numel() > 0
+
+        assert output.trajectory_log_probs is not None, "trajectory_log_probs should be present when logprobs=True"
+        assert hasattr(output.trajectory_log_probs, "shape")
+        assert output.trajectory_log_probs.numel() > 0
 
 
 @pytest.mark.core_model

tests/e2e/offline_inference/test_bagel_lora.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for BAGEL LoRA support (Stage 1 / DiT).
+
+Validates that LoRA adapters are correctly loaded, applied with controllable
+scale, and cleanly deactivated.  Uses a synthetic rank-1 adapter targeting the
+first decoder layer's QKV projection.
+
+Assertions:
+  (a) LoRA at scale=1.0 visibly changes the output  (diff > 0.5)
+  (b) scale=2.0 produces a larger delta than scale=1.0  (linearity)
+  (c) The delta is bounded  (diff < 80, not corrupted)
+  (d) Deactivating LoRA exactly restores the baseline  (diff == 0)
+"""
+
+import json
+import os
+
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
+
+
+# ---------------------------------------------------------------------------
+# Helpers (reused from test_bagel_text2img.py patterns)
+# ---------------------------------------------------------------------------
+
+
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 10) -> list[OmniSamplingParams]:
+    params_list = omni.default_sampling_params_list
+    if len(params_list) > 1:
+        params_list[1].num_inference_steps = num_inference_steps
+        params_list[1].extra_args = {
+            "cfg_text_scale": 4.0,
+            "cfg_img_scale": 1.5,
+        }
+    return params_list
+
+
+def _extract_generated_image(omni_outputs: list[OmniRequestOutput]) -> Image.Image | None:
+    for req_output in omni_outputs:
+        if req_output.images:
+            return req_output.images[0]
+    return None
+
+
+def _generate_bagel_image(omni: Omni) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = None
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated"
+    return img
+
+
+def _generate_bagel_image_with_lora(
+    omni: Omni,
+    lora_request: LoRARequest,
+    lora_scale: float = 1.0,
+) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = lora_request
+    params_list[1].lora_scale = lora_scale
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated with LoRA"
+    return img
+
+
+# BAGEL uses GQA: hidden_size=3584, 28 Q heads, 4 KV heads, head_dim=128
+# QKV packed dim = 28*128 + 4*128 + 4*128 = 3584 + 512 + 512 = 4608
+_LORA_DIM = 3584
+_LORA_QKV_DIM = 4608
+_LORA_MODULE = "bagel.language_model.model.layers.0.self_attn.qkv_proj"
+_LORA_RANK = 4
+
+
+def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
+    """Write synthetic adapter to disk and return a file-backed LoRARequest."""
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+    gen = torch.Generator().manual_seed(42)
+    lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
+    lora_b = torch.randn((_LORA_QKV_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
+    save_file(
+        {
+            f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
+            f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
+        encoding="utf-8",
+    )
+    lora_dir = str(adapter_dir)
+    return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
+
+
+# ---------------------------------------------------------------------------
+# Test
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_lora_scale_and_deactivation(run_level, tmp_path):
+    """Validate LoRA effect, bounded perturbation, and clean deactivation."""
+    config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level)
+    omni = Omni(model=MODEL, stage_configs_path=config_path, stage_init_timeout=300)
+    try:
+        lora_request = _make_file_lora_request(tmp_path / "bagel_lora")
+
+        # 1) Baseline (no LoRA)
+        baseline = _generate_bagel_image(omni)
+
+        # 2) LoRA with scale=1.0
+        img_1x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=1.0)
+
+        # 3) LoRA with scale=2.0
+        img_2x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=2.0)
+
+        # 4) No LoRA again (deactivation)
+        restored = _generate_bagel_image(omni)
+
+        baseline_arr = np.array(baseline, dtype=np.int16)
+        img_1x_arr = np.array(img_1x, dtype=np.int16)
+        img_2x_arr = np.array(img_2x, dtype=np.int16)
+        restored_arr = np.array(restored, dtype=np.int16)
+
+        diff_1x = np.abs(baseline_arr - img_1x_arr).mean()
+        diff_2x = np.abs(baseline_arr - img_2x_arr).mean()
+        diff_restored = np.abs(baseline_arr - restored_arr).mean()
+
+        # (a) Adapter has visible effect at both scales
+        assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
+        assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
+
+        # (b) Different scales produce different outputs
+        assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
+            f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
+        )
+
+        # (c) Output is not corrupted
+        assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
+        assert diff_2x < 80, f"LoRA output looks corrupted: diff_2x={diff_2x}"
+
+        # (d) Deactivation fully restores base model
+        assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"
+    finally:
+        omni.close()

tests/e2e/offline_inference/test_dynin_omni.py

@@ -0,0 +1,419 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline smoke tests for Dynin-Omni.
+
+- model: "snu-aidas/Dynin-Omni"
+- stage config: tests/e2e/stage_configs/dynin_omni_ci.yaml
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from transformers import AutoTokenizer
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None
+_DEFAULT_STAGE_CONFIG_PATH = _REPO_ROOT / "tests" / "e2e" / "stage_configs" / "dynin_omni_ci.yaml"
+
+models = ["snu-aidas/Dynin-Omni"]
+stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)]
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.omni,
+]
+
+
+# prompting util
+def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]:
+    encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True)
+    token_ids = [int(v) for v in encoded["input_ids"][0].tolist()]
+    attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()]
+    additional_information: dict[str, Any] = {
+        "task": ["mmu"],
+        "detok_id": [0],
+        "prompt_length": [len(token_ids)],
+        "attention_mask": [attention_mask],
+        "max_new_tokens": [64],
+        "steps": [64],
+        "block_length": [16],
+        "temperature": [0.0],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": token_ids,
+        "additional_information": additional_information,
+        "modalities": ["text"],
+    }
+
+
+def _build_mmu_multimodal_prompt(
+    tokenizer: Any,
+    question: str,
+    dynin_config_path: str | None,
+    *,
+    image: Any | None = None,
+    audio: tuple[np.ndarray, int] | None = None,
+) -> dict[str, Any]:
+    if image is None and audio is None:
+        raise ValueError("At least one multimodal input (image or audio) must be provided.")
+
+    prefix_chunks: list[str] = []
+    mm_data: dict[str, Any] = {}
+    if image is not None:
+        prefix_chunks.append("<|soi|><|image|><|eoi|>")
+        mm_data["image"] = image
+    if audio is not None:
+        prefix_chunks.append("<|soa|><|audio|><|eoa|>")
+        mm_data["audio"] = audio
+
+    prefixed_question = " ".join(prefix_chunks + [question]).strip()
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question=prefixed_question,
+        dynin_config_path=dynin_config_path,
+    )
+    prompt["multi_modal_data"] = mm_data
+    prompt["modalities"] = ["text"]
+    return prompt
+
+
+def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray:
+    x = np.linspace(0, 255, width, dtype=np.uint8)
+    y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
+    red = np.tile(x, (height, 1))
+    green = np.tile(y, (1, width))
+    blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8)
+    return np.stack([red, green, blue], axis=-1)
+
+
+def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]:
+    t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32)
+    waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t)
+    return waveform.astype(np.float32), sample_rate
+
+
+# prompting util
+def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->audio decode path.
+    generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2s"],
+        "detok_id": [1],
+        "generated_token_ids": [generated_audio_token_ids],
+        "audio_codebook_size": [4096],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["audio"],
+    }
+
+
+# prompting util
+def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->image decode path.
+    # MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32.
+    generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2i"],
+        "detok_id": [2],
+        "generated_token_ids": [generated_image_token_ids],
+        "codebook_size": [8192],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["image"],
+    }
+
+
+def _configure_dynin_config_env() -> None:
+    if DYNIN_CONFIG_PATH:
+        os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH)
+    else:
+        os.environ.pop("DYNIN_CONFIG_PATH", None)
+
+
+def _is_finished_request_output(request_output: Any) -> bool:
+    if request_output is None:
+        return False
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        if req is not None and bool(getattr(req, "finished", False)):
+            return True
+    return False
+
+
+def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
+    matched = [
+        stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type
+    ]
+    if not matched:
+        return None
+
+    # Prefer the latest finished chunk to avoid picking an intermediate stream output.
+    for stage_output in reversed(matched):
+        if _is_finished_request_output(getattr(stage_output, "request_output", None)):
+            return stage_output
+    return matched[-1]
+
+
+def _to_token_list(value: Any) -> list[int]:
+    if value is None:
+        return []
+    if hasattr(value, "detach"):
+        value = value.detach()
+    if hasattr(value, "cpu"):
+        value = value.cpu()
+    if hasattr(value, "flatten"):
+        value = value.flatten().tolist()
+    if isinstance(value, tuple):
+        value = list(value)
+    if not isinstance(value, list):
+        return []
+    out: list[int] = []
+    for token in value:
+        if isinstance(token, bool):
+            continue
+        try:
+            out.append(int(token))
+        except Exception:
+            continue
+    return out
+
+
+def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return ""
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = (
+            getattr(completion, "multimodal_output", None)
+            or getattr(req, "multimodal_output", None)
+            or getattr(stage_output, "multimodal_output", None)
+            or {}
+        )
+        text = mm_out.get("text")
+        if isinstance(text, list) and text:
+            text = text[-1]
+        if isinstance(text, str) and text.strip():
+            return text.strip()
+        if tokenizer is not None:
+            for key in ("text_tokens", "token_ids"):
+                token_ids = _to_token_list(mm_out.get(key))
+                if not token_ids:
+                    continue
+                decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
+                if isinstance(decoded, str) and decoded.strip():
+                    return decoded.strip()
+        fallback = getattr(completion, "text", None)
+        if isinstance(fallback, str) and fallback.strip():
+            return fallback.strip()
+    return ""
+
+
+def _extract_audio(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "audio" in mm_out:
+            return mm_out["audio"]
+    return None
+
+
+def _extract_image(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "image" in mm_out:
+            return mm_out["image"]
+    return None
+
+
+def _numel(value: Any) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, torch.Tensor):
+        return int(value.numel())
+    shape = getattr(value, "shape", None)
+    if shape is not None:
+        try:
+            total = 1
+            for dim in shape:
+                total *= int(dim)
+            return int(total)
+        except Exception:
+            pass
+    if isinstance(value, (list, tuple)):
+        return len(value)
+    return 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2i_decode_to_image(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    image_output = _find_stage_output(outputs, "image")
+    assert image_output is not None
+    image_value = _extract_image(image_output)
+    assert image_value is not None
+    assert _numel(image_value) > 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_mmu_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question="What is 2 + 2? Answer in one short sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_image_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Describe the image briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        image=_generate_synthetic_image(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_speech_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Transcribe the audio briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        audio=_generate_synthetic_audio(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2s_decode_to_audio(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    audio_output = _find_stage_output(outputs, "audio")
+    assert audio_output is not None
+    audio_value = _extract_audio(audio_output)
+    assert audio_value is not None
+    assert _numel(audio_value) > 0

tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import hashlib
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from tests.utils import hardware_test
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
+I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
+
+T2V_PROMPT = (
+    "At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
+    "soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
+    "the ripples, cinematic, realistic, highly detailed."
+)
+T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
+I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
+I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+def _get_ltx2_model() -> str:
+    return os.environ.get("VLLM_TEST_LTX2_MODEL", "Lightricks/LTX-2")
+
+
+def _md5(path: Path) -> str:
+    digest = hashlib.md5(usedforsecurity=False)
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _make_deterministic_test_image(path: Path) -> None:
+    """Create a deterministic 256x256 test image for I2V tests."""
+    rng = np.random.RandomState(42)
+    img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
+    img.save(path)
+
+
+def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
+    result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
+    assert result.returncode == 0, (
+        f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+    )
+    generated_md5 = _md5(output_path)
+    assert generated_md5 == expected_md5, (
+        f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+    )
+
+
+# ── T2V tests ──
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
+    """T2V with CFG=4.0, cfg-parallel-size=2."""
+    output = tmp_path / "t2v_cfg4.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
+    cmd = [
+        sys.executable,
+        str(T2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--prompt",
+        T2V_PROMPT,
+        "--negative-prompt",
+        T2V_NEGATIVE_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "145",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "4.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--cfg-parallel-size",
+        "2",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx2_t2v_no_cfg(tmp_path: Path):
+    """T2V with CFG=1.0 (no classifier-free guidance)."""
+    output = tmp_path / "t2v_nocfg.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
+    cmd = [
+        sys.executable,
+        str(T2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--prompt",
+        T2V_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "145",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "1.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
+
+
+# ── I2V tests ──
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
+    """I2V with CFG=4.0, cfg-parallel-size=2."""
+    test_image = tmp_path / "test_input.png"
+    _make_deterministic_test_image(test_image)
+    output = tmp_path / "i2v_cfg4.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
+    cmd = [
+        sys.executable,
+        str(I2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--model-class-name",
+        "LTX2ImageToVideoPipeline",
+        "--image",
+        str(test_image),
+        "--prompt",
+        I2V_PROMPT,
+        "--negative-prompt",
+        I2V_NEGATIVE_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "73",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "4.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--cfg-parallel-size",
+        "2",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx2_i2v_no_cfg(tmp_path: Path):
+    """I2V with CFG=1.0 (no classifier-free guidance)."""
+    test_image = tmp_path / "test_input.png"
+    _make_deterministic_test_image(test_image)
+    output = tmp_path / "i2v_nocfg.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
+    cmd = [
+        sys.executable,
+        str(I2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--model-class-name",
+        "LTX2ImageToVideoPipeline",
+        "--image",
+        str(test_image),
+        "--prompt",
+        I2V_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "73",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "1.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")

tests/e2e/offline_inference/test_magi_human.py

@@ -110,16 +110,31 @@ def test_magi_human_e2e(run_level):
         assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
         assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
 
-        audio_waveform = None
-        if hasattr(first, "multimodal_output") and first.multimodal_output:
-            audio_waveform = first.multimodal_output.get("audio")
+        mm = first.multimodal_output
+        assert mm, "multimodal_output is empty or missing"
+
+        audio_waveform = mm.get("audio")
         assert audio_waveform is not None, "No audio waveform in multimodal_output"
 
+        audio_sample_rate = mm.get("audio_sample_rate")
+        assert audio_sample_rate is not None, (
+            "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
+        )
+        assert isinstance(audio_sample_rate, (int, float)), (
+            f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
+        )
+        assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
+
+        fps = mm.get("fps")
+        assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
+        assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
+        assert int(fps) > 0, f"fps must be positive, got {fps}"
+
         video_bytes = mux_video_audio_bytes(
             video_frames,
             audio_waveform,
-            fps=25.0,
-            audio_sample_rate=44100,
+            fps=float(fps),
+            audio_sample_rate=int(audio_sample_rate),
         )
         assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
         assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Example online tests for Dynin-Omni model.
+"""
+
+import base64
+import gc
+import os
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+from vllm.assets.image import ImageAsset
+
+from tests import conftest as tests_conftest
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "snu-aidas/Dynin-Omni"
+STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml")
+_WHISPER_SAMPLE_RATE_HZ = 16_000
+
+T2I_PROMPT = "A high quality detailed living room interior photo."
+T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving."
+I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
+
+TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)]
+_STAGE_COUNT = 3
+_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
+
+
+def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray:
+    """Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper."""
+    if samplerate <= 0:
+        raise ValueError(f"Invalid audio sample rate: {samplerate}")
+
+    waveform = np.asarray(audio_data, dtype=np.float32)
+    if waveform.ndim == 0:
+        raise ValueError("Audio waveform must have at least one dimension")
+    if waveform.ndim > 1:
+        waveform = np.mean(waveform, axis=1)
+    if waveform.size == 0:
+        raise ValueError("Empty audio waveform")
+
+    if samplerate != _WHISPER_SAMPLE_RATE_HZ:
+        target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1)
+        source_positions = np.arange(waveform.shape[0], dtype=np.float64)
+        target_positions = np.linspace(
+            0.0,
+            max(waveform.shape[0] - 1, 0),
+            num=target_num_samples,
+            dtype=np.float64,
+        )
+        waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32)
+
+    return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32)
+
+
+def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str:
+    """Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg."""
+    import whisper
+
+    data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True)
+    audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate)
+
+    model = whisper.load_model("small", device="cpu")
+    try:
+        transcript = model.transcribe(
+            audio_waveform,
+            temperature=0.0,
+            word_timestamps=True,
+            condition_on_previous_text=False,
+        )["text"]
+    finally:
+        del model
+        gc.collect()
+
+    return transcript or ""
+
+
+@pytest.fixture
+def dynin_t2s_openai_client(openai_client, monkeypatch):
+    monkeypatch.setattr(
+        tests_conftest,
+        "convert_audio_bytes_to_text",
+        _convert_audio_bytes_to_text_without_ffmpeg,
+    )
+    return openai_client
+
+
+def _build_t2i_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}]
+
+
+def _build_t2s_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}]
+
+
+def _build_i2i_messages(prompt: str) -> list[dict]:
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"<|i2i|> {prompt}"},
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
+            ],
+        }
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_i2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_i2i_messages(I2I_PROMPT),
+        "modalities": ["image"],
+        "extra_body": {
+            "sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2i_messages(T2I_PROMPT),
+        "modalities": ["image"],
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2s_messages(T2S_PROMPT),
+        "modalities": ["audio"],
+    }
+    dynin_t2s_openai_client.send_omni_request(request_config)

tests/e2e/stage_configs/dynin_omni_ci.yaml

@@ -0,0 +1,84 @@
+# stage config for running dynin_omni with a 3-stage architecture.
+# this config is intended for e2e smoke tests.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2text
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2image
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
+    final_output: true
+    final_output_type: image
+
+  - stage_id: 2
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2audio
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
+    final_output: true
+    final_output_type: audio
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

tests/engine/test_cross_stage_lora.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for cross-stage LoRA routing in the orchestrator."""
+
+from __future__ import annotations
+
+import pytest
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestBuildEngineCoreRequestLoRA:
+    """Verify build_engine_core_request_from_tokens passes LoRA from params."""
+
+    def test_lora_extracted_from_diffusion_params(self):
+        lr = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
+        params = OmniDiffusionSamplingParams(lora_request=lr)
+
+        # OmniDiffusionSamplingParams is not a SamplingParams, so
+        # build_engine_core_request_from_tokens takes the pooling path.
+        # We only care that lora_request is extracted via getattr.
+        request = build_engine_core_request_from_tokens(
+            request_id="req-1",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is lr
+
+    def test_no_lora_on_sampling_params(self):
+        params = SamplingParams(max_tokens=10)
+
+        request = build_engine_core_request_from_tokens(
+            request_id="req-2",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is None

tests/engine/test_orchestrator.py

@@ -0,0 +1,510 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import queue
+import threading
+import time
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any
+
+import janus
+import pytest
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class OrchestratorFixture:
+    orchestrator: Orchestrator
+    request_sync_q: Any
+    output_sync_q: Any
+    queues: tuple[janus.Queue, ...]
+    thread: threading.Thread
+    result_future: concurrent.futures.Future[None]
+
+
+class FakeStageClient:
+    def __init__(
+        self,
+        *,
+        stage_type: str = "llm",
+        final_output: bool = False,
+        final_output_type: str = "text",
+        next_inputs: list[dict] | None = None,
+    ) -> None:
+        self.stage_type = stage_type
+        self.final_output = final_output
+        self.final_output_type = final_output_type
+        self.next_inputs = list(next_inputs or [])
+        self.custom_process_input_func = None
+        self.add_request_calls: list[tuple] = []
+        self.abort_calls: list[list[str]] = []
+        self.shutdown_calls = 0
+        self._engine_core_outputs = queue.Queue()
+        self._diffusion_outputs = queue.Queue()
+
+    # Orchestrator-facing interface.
+    async def add_request_async(self, *args, **_kwargs) -> None:
+        self.add_request_calls.append(args)
+
+    async def get_output_async(self):
+        try:
+            return self._engine_core_outputs.get_nowait()
+        except queue.Empty:
+            return SimpleNamespace(outputs=[])
+
+    def get_diffusion_output_nowait(self):
+        try:
+            return self._diffusion_outputs.get_nowait()
+        except queue.Empty:
+            return None
+
+    def set_engine_outputs(self, outputs) -> None:
+        return None
+
+    def process_engine_inputs(self, stage_list, prompt=None):
+        return list(self.next_inputs)
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        self.abort_calls.append(list(request_ids))
+
+    def shutdown(self) -> None:
+        self.shutdown_calls += 1
+
+    # Test helpers for seeding fake stage outputs.
+    def push_engine_core_outputs(self, outputs) -> None:
+        self._engine_core_outputs.put_nowait(outputs)
+
+    def push_diffusion_output(self, output) -> None:
+        self._diffusion_outputs.put_nowait(output)
+
+
+class FakeOutputProcessor:
+    def __init__(self, *, request_outputs: list[object] | None = None) -> None:
+        self.request_outputs = list(request_outputs or [])
+
+    def add_request(self, *_args, **_kwargs) -> None:
+        return None
+
+    def process_outputs(self, *_args, **_kwargs):
+        return SimpleNamespace(
+            request_outputs=list(self.request_outputs),
+            reqs_to_abort=[],
+        )
+
+    def update_scheduler_stats(self, _scheduler_stats) -> None:
+        return None
+
+
+def _sampling_params(max_tokens: int = 4) -> SamplingParams:
+    return SamplingParams(max_tokens=max_tokens)
+
+
+def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
+    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
+
+
+def _build_request_output(
+    request_id: str,
+    *,
+    token_ids: list[int] | None = None,
+    prompt_token_ids: list[int] | None = None,
+    finished: bool = True,
+    text: str = "test",
+) -> RequestOutput:
+    completion = CompletionOutput(
+        index=0,
+        text=text,
+        token_ids=list(token_ids or [1, 2]),
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop" if finished else None,
+        stop_reason=None,
+    )
+    return RequestOutput(
+        request_id=request_id,
+        prompt="prompt",
+        prompt_token_ids=list(prompt_token_ids or [10, 11]),
+        prompt_logprobs=None,
+        outputs=[completion],
+        finished=finished,
+        metrics=None,
+        lora_request=None,
+    )
+
+
+def _build_harness(
+    stage_clients: list[object],
+    *,
+    output_processors: list[object] | None = None,
+    stage_vllm_configs: list[object] | None = None,
+    async_chunk: bool = False,
+) -> OrchestratorFixture:
+    if output_processors is None:
+        output_processors = [FakeOutputProcessor() for _ in stage_clients]
+    if stage_vllm_configs is None:
+        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+
+    ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = (
+        concurrent.futures.Future()
+    )
+    result_future: concurrent.futures.Future[None] = concurrent.futures.Future()
+
+    def _runner() -> None:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        async def _run() -> None:
+            request_queue = janus.Queue()
+            output_queue = janus.Queue()
+            rpc_queue = janus.Queue()
+            orchestrator = Orchestrator(
+                request_async_queue=request_queue.async_q,
+                output_async_queue=output_queue.async_q,
+                rpc_async_queue=rpc_queue.async_q,
+                stage_clients=stage_clients,
+                output_processors=output_processors,
+                stage_vllm_configs=stage_vllm_configs,
+                async_chunk=async_chunk,
+            )
+            ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
+            await orchestrator.run()
+
+        try:
+            loop.run_until_complete(_run())
+            result_future.set_result(None)
+        except Exception as exc:
+            result_future.set_exception(exc)
+        finally:
+            try:
+                pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            finally:
+                asyncio.set_event_loop(None)
+                loop.close()
+
+    thread = threading.Thread(target=_runner, daemon=True, name="test-orchestrator")
+    thread.start()
+
+    orchestrator, request_queue, output_queue, rpc_queue = ready_future.result(timeout=5)
+    return OrchestratorFixture(
+        orchestrator=orchestrator,
+        request_sync_q=request_queue.sync_q,
+        output_sync_q=output_queue.sync_q,
+        queues=(request_queue, output_queue, rpc_queue),
+        thread=thread,
+        result_future=result_future,
+    )
+
+
+async def _shutdown_orchestrator(orchestrator_fixture: OrchestratorFixture) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+    await asyncio.to_thread(orchestrator_fixture.thread.join, 5)
+    if orchestrator_fixture.thread.is_alive():
+        raise AssertionError("Timed out waiting for orchestrator thread shutdown")
+    orchestrator_fixture.result_future.result(timeout=0)
+
+
+async def _wait_for(predicate, *, timeout: float = 2.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for predicate")
+        await asyncio.sleep(0.01)
+
+
+async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+    deadline = time.monotonic() + timeout
+    while True:
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for orchestrator output")
+        try:
+            msg = orchestrator_fixture.output_sync_q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.01)
+            continue
+        if msg.get("type") == "output":
+            return msg
+
+
+async def _enqueue_add_request(
+    orchestrator_fixture: OrchestratorFixture,
+    *,
+    request_id: str,
+    prompt,
+    original_prompt,
+    sampling_params_list,
+    final_stage_id: int,
+) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait(
+        {
+            "type": "add_request",
+            "request_id": request_id,
+            "prompt": prompt,
+            "original_prompt": original_prompt,
+            "sampling_params_list": sampling_params_list,
+            "final_stage_id": final_stage_id,
+        }
+    )
+
+
+async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, request_ids: list[str]) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait(
+        {
+            "type": "abort",
+            "request_ids": request_ids,
+        }
+    )
+
+
+@pytest.fixture
+def orchestrator_factory():
+    fixtures: list[OrchestratorFixture] = []
+
+    def _factory(*args, **kwargs) -> OrchestratorFixture:
+        fixture = _build_harness(*args, **kwargs)
+        fixtures.append(fixture)
+        return fixture
+
+    yield _factory
+
+    for fixture in fixtures:
+        if fixture.thread.is_alive():
+            fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            fixture.thread.join(timeout=5)
+        for q in fixture.queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+async def test_run_two_stage_llm(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(
+        stage_type="llm",
+        final_output=True,
+        next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
+    )
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[3, 4], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[10, 11], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
+    request = SimpleNamespace(request_id="req-llm", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-llm",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        stage1_request = stage1.add_request_calls[0][0]
+        assert stage1_request.request_id == "req-llm"
+        assert stage1_request.prompt_token_ids == [7, 8, 9]
+
+        stage1.push_engine_core_outputs(_engine_core_outputs("stage1-raw", 2.0))
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-llm"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-llm"
+        assert "req-llm" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-diff",
+            prompt={"prompt": "draw a cat"},
+            original_prompt={"prompt": "draw a cat"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_diffusion(
+                request_id="req-diff",
+                images=[],
+                final_output_type="image",
+            )
+        )
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-diff"
+        assert output_msg["stage_id"] == 0
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-diff"
+        assert "req-diff" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-img", token_ids=[3, 4], finished=True)]),
+        FakeOutputProcessor(),
+    ]
+    orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
+    request = SimpleNamespace(request_id="req-img", prompt_token_ids=[1, 2, 3])
+    params = OmniDiffusionSamplingParams()
+    original_prompt = {"prompt": "draw a fox"}
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-img",
+            prompt=request,
+            original_prompt=original_prompt,
+            sampling_params_list=[_sampling_params(), params],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        assert stage1.add_request_calls[0] == ("req-img", original_prompt, params)
+
+        stage1.push_diffusion_output(
+            OmniRequestOutput.from_diffusion(
+                request_id="req-img",
+                images=[],
+                final_output_type="image",
+            )
+        )
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-img"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-img"
+        assert "req-img" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_async_chunk(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[20, 21], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(
+        [stage0, stage1],
+        output_processors=processors,
+        async_chunk=True,
+    )
+    request = SimpleNamespace(request_id="req-async", prompt_token_ids=[1, 2, 3, 4])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-async",
+            prompt=request,
+            original_prompt={"prompt": "hello async"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        prewarmed_request = stage1.add_request_calls[0][0]
+        assert prewarmed_request.request_id == "req-async"
+        assert prewarmed_request.prompt_token_ids
+        assert all(token_id == 0 for token_id in prewarmed_request.prompt_token_ids)
+
+        stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0))
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-async"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert "req-async" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_shutdown(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image"),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages)
+
+    await _shutdown_orchestrator(orchestrator_fixture)
+
+    assert not orchestrator_fixture.thread.is_alive()
+    for stage in stages:
+        assert stage.shutdown_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_run_abort(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[2], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
+    request = SimpleNamespace(request_id="req-abort", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-abort",
+            prompt=request,
+            original_prompt={"prompt": "cancel me"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+
+        await _enqueue_abort_request(orchestrator_fixture, ["req-abort"])
+        await _wait_for(lambda: all(stage.abort_calls for stage in stages))
+
+        for stage in stages:
+            assert stage.abort_calls == [["req-abort"]]
+        assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -0,0 +1,207 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummySenderStage:
+    def __init__(self, sender_info):
+        self._sender_info = sender_info
+        self.engine_outputs = None
+
+    def set_engine_outputs(self, outputs):
+        self.engine_outputs = outputs
+
+    def get_kv_sender_info(self):
+        return self._sender_info
+
+
+class _DummyDiffusionStage:
+    stage_type = "diffusion"
+    custom_process_input_func = None
+
+    def __init__(self, engine_input_source=None):
+        self.engine_input_source = engine_input_source or [0]
+        self.calls = []
+
+    async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None):
+        self.calls.append(
+            {
+                "request_id": request_id,
+                "prompt": prompt,
+                "sampling_params": sampling_params,
+                "kv_sender_info": kv_sender_info,
+            }
+        )
+
+
+def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50151,
+    }
+
+
+def test_stage_engine_core_client_falls_back_to_detected_ip_for_loopback(monkeypatch):
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.client_addresses = {"input_address": "tcp://127.0.0.1:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    monkeypatch.setattr(client, "_detect_local_ip", lambda: "192.168.0.12")
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "192.168.0.12",
+        "zmq_port": 50152,
+    }
+
+
+def test_stage_engine_core_client_uses_connector_config_for_sender_port():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 3
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "3",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "10.20.30.99",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.99",
+        "zmq_port": 51103,
+    }
+
+
+def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 2
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "2",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "127.0.0.1",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "127.0.0.1",
+        "zmq_port": 51102,
+    }
+
+
+def test_forward_to_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 2
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None]
+    orchestrator.output_processors = [None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-1",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), params],
+        final_stage_id=1,
+    )
+
+    output = SimpleNamespace(request_id="req-1", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
+
+    assert sender_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["request_id"] == "req-1"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0
+
+
+def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 3
+    orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None, None]
+    orchestrator.output_processors = [None, None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-3",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params],
+        final_stage_id=2,
+    )
+
+    output = SimpleNamespace(request_id="req-3", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
+
+    assert previous_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+
+
+def test_prewarm_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.num_stages = 2
+
+    req_state = OrchestratorRequestState(
+        request_id="req-2",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()],
+        final_stage_id=1,
+    )
+
+    stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3])
+    asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state))
+
+    assert diffusion_stage.calls[0]["request_id"] == "req-2"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.3", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0

tests/engine/test_output_modality.py

@@ -12,6 +12,7 @@ import pytest
 import torch
 
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 _ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine"
 

tests/entrypoints/openai_api/test_image_server.py

@@ -178,7 +178,7 @@ def test_client(mock_async_diffusion):
     )
     app.state.args = Namespace(
         default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        max_generated_image_size=1024 * 1792,
     )
 
     return TestClient(app)
@@ -245,7 +245,7 @@ def async_omni_stage_configs_only_client():
     # AsyncOmni exposes stage_configs on the engine instance.
     app.state.args = Namespace(
         default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        max_generated_image_size=1024 * 1792,
     )
     return TestClient(app)
 
@@ -392,6 +392,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def test_generate_images_max_size_rejected(async_omni_test_client):
+    """Test that a size exceeding max_generated_image_size returns 400."""
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a cat",
+            "size": "2048x2048",  # 4,194,304 pixels > max_generated_image_size (1,048,576)
+        },
+    )
+    assert response.status_code == 400
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(
@@ -982,12 +994,13 @@ def test_image_edit_parameter_default_single_stage(test_client):
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
 
+    # Size exceeding max_generated_image_size (1024*1792) returns 400
     response = test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes_1)],
         data={
             "prompt": "hello world.",
-            "size": "96x96",
+            "size": "2048x2048",
         },
     )
     assert response.status_code == 400

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -385,3 +385,18 @@ class TestStreamingSpeechWebSocket:
 
         speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
         assert websocket.send_json.await_count == 2
+
+
+class TestGeneratePcmChunksContract:
+    """Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech.
+
+    The WebSocket handler calls speech_service._generate_pcm_chunks()
+    at runtime. If the method is removed, all WS TTS streaming breaks
+    with an AttributeError. This test catches that at CI time.
+    """
+
+    def test_generate_pcm_chunks_defined(self):
+        assert hasattr(OmniOpenAIServingSpeech, "_generate_pcm_chunks")
+        assert asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._generate_pcm_chunks) or callable(
+            OmniOpenAIServingSpeech._generate_pcm_chunks
+        )

tests/entrypoints/openai_api/test_text_splitter.py

@@ -4,7 +4,7 @@ import pytest
 
 from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
 
-pytestmark = [pytest.mark.openai, pytest.mark.speech]
+pytestmark = [pytest.mark.openai, pytest.mark.speech, pytest.mark.core_model, pytest.mark.cpu]
 
 
 class TestSentenceSplitterEnglish:

tests/entrypoints/test_stage_utils.py

@@ -6,8 +6,6 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.stage_utils import set_stage_devices
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 
 def _make_dummy_torch(call_log):
     class _Props:
@@ -55,6 +53,8 @@ def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_
     return mock_platform
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -75,6 +75,8 @@ def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypat
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -90,9 +92,10 @@ def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, mon
         mock_platform,
     )
 
-    # Raise since we need 4 GPUs, but we only have 2 visible
-    with pytest.raises(ValueError):
-        set_stage_devices(stage_id=0, devices="0,1,2,3")
+    # Keep the logical mapping and resolve to the visible subset.
+    set_stage_devices(stage_id=0, devices="0,1,2,3")
+
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -8,6 +8,8 @@ import pytest
 import torch
 import torch.nn as nn
 
+from tests.utils import hardware_test
+
 
 class TestPreLookaheadLayer:
     """Tests for PreLookaheadLayer."""
@@ -18,6 +20,8 @@ class TestPreLookaheadLayer:
 
         return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3)
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_forward_shape(self, layer):
         """Test that output shape matches input shape."""
         batch, seq_len, channels = 2, 10, 512
@@ -27,6 +31,8 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_forward_with_context(self, layer):
         """Test forward with context for streaming."""
         batch, seq_len, channels = 1, 10, 512
@@ -38,6 +44,8 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_residual_connection(self, layer):
         """Test that residual connection is applied."""
         batch, seq_len, channels = 1, 5, 512
@@ -59,6 +67,8 @@ class TestDiTAttention:
 
         return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, attention):
         """Test attention output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -68,6 +78,8 @@ class TestDiTAttention:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_with_mask(self, attention):
         """Test attention with mask."""
         batch, seq_len, dim = 2, 16, 512
@@ -81,6 +93,8 @@ class TestDiTAttention:
         # Masked positions should be zero
         assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:]))
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_qkv_projections(self, attention):
         """Test that Q/K/V projections exist and have correct dimensions."""
         assert hasattr(attention, "to_q")
@@ -100,6 +114,8 @@ class TestDiTBlock:
 
         return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, block):
         """Test block output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -110,6 +126,8 @@ class TestDiTBlock:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_adalayernorm_modulation(self, block):
         """Test that AdaLayerNorm modulates based on timestep."""
         batch, seq_len, dim = 1, 8, 512
@@ -144,6 +162,8 @@ class TestDiT:
             long_skip_connection=True,
         )
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, dit):
         """Test DiT forward output shape."""
         batch, mel_dim, seq_len = 1, 80, 32
@@ -158,6 +178,8 @@ class TestDiT:
 
         assert out.shape == (batch, mel_dim, seq_len)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_timestep_embedding(self, dit):
         """Test that different timesteps produce different outputs."""
         batch, mel_dim, seq_len = 1, 80, 16
@@ -190,6 +212,8 @@ class TestCFM:
 
         return DummyEstimator()
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_causal_conditional_cfm_forward(self, dummy_estimator):
         """Test CausalConditionalCFM forward pass."""
         from omegaconf import DictConfig
@@ -228,6 +252,8 @@ class TestCFM:
 class TestSDPAFallback:
     """Test SDPA fallback for float32 inputs."""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_float32_uses_sdpa(self):
         """Test that float32 inputs use SDPA fallback."""
         from vllm_omni.diffusion.attention.layer import Attention

tests/model_executor/models/cosyvoice3/test_cosyvoice3_utils.py

@@ -5,6 +5,8 @@
 import pytest
 import torch
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class TestMakePadMask:
     """Tests for make_pad_mask utility."""

tests/test_diffusion_config_propagation.py

@@ -7,6 +7,7 @@ Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
 
 from collections.abc import Mapping
 
+import pytest
 import torch
 
 from vllm_omni.config.stage_config import StageConfigFactory
@@ -15,6 +16,8 @@ from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     """Simulate the real path: create_default_diffusion → OmniDiffusionConfig.

tests/worker/test_cudagraph_wrapper_perf.py

@@ -1,185 +0,0 @@
-"""Tests for CUDAGraphWrapper.__getattr__ performance optimization.
-
-This module tests that the patched CUDAGraphWrapper avoids expensive __repr__
-calls when hasattr() is used for non-existent attributes. The original vLLM
-implementation includes {self.runnable} in the AttributeError message, which
-triggers model tree traversal and can take ~6ms on large models.
-"""
-
-import time
-
-import pytest
-import torch
-import torch.nn as nn
-
-from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class SlowReprModel(nn.Module):
-    """A mock model with artificially slow __repr__ to detect unwanted calls."""
-
-    def __init__(self, repr_delay_ms: float = 10.0):
-        super().__init__()
-        self.linear = nn.Linear(16, 16)
-        self.repr_delay_ms = repr_delay_ms
-        self.repr_call_count = 0
-
-    def forward(self, x):
-        return self.linear(x)
-
-    def __repr__(self):
-        self.repr_call_count += 1
-        # Simulate expensive repr by sleeping
-        time.sleep(self.repr_delay_ms / 1000.0)
-        return f"SlowReprModel(delay={self.repr_delay_ms}ms)"
-
-
-class MockCUDAGraphWrapper:
-    """A minimal mock that mimics CUDAGraphWrapper structure for CPU testing."""
-
-    def __init__(self, runnable):
-        # Store in __dict__ directly to avoid triggering __getattr__
-        object.__setattr__(self, "runnable", runnable)
-
-    def __getattr__(self, key: str):
-        # This is the optimized implementation we're testing
-        runnable = object.__getattribute__(self, "runnable")
-        if hasattr(runnable, key):
-            return getattr(runnable, key)
-        # Key optimization: DO NOT include {self.runnable} in error message
-        # as it triggers expensive __repr__ on large models
-        raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper")
-
-
-def test_hasattr_nonexistent_does_not_trigger_repr():
-    """Verify that hasattr for non-existent attributes doesn't call __repr__."""
-    model = SlowReprModel(repr_delay_ms=100.0)  # Very slow repr
-    wrapper = MockCUDAGraphWrapper(model)
-
-    # Reset counter
-    model.repr_call_count = 0
-
-    # Call hasattr for non-existent attribute multiple times
-    for _ in range(10):
-        result = hasattr(wrapper, "nonexistent_attribute_xyz")
-        assert result is False
-
-    # __repr__ should never have been called
-    assert model.repr_call_count == 0, (
-        f"__repr__ was called {model.repr_call_count} times when checking "
-        "for non-existent attributes. This indicates the AttributeError "
-        "message contains {self.runnable} which triggers expensive repr."
-    )
-
-
-def test_hasattr_nonexistent_is_fast():
-    """Verify that hasattr for non-existent attributes is fast (<1ms per call)."""
-    model = SlowReprModel(repr_delay_ms=100.0)
-    wrapper = MockCUDAGraphWrapper(model)
-
-    num_iterations = 100
-    start = time.perf_counter()
-    for _ in range(num_iterations):
-        hasattr(wrapper, "nonexistent_attribute_xyz")
-    elapsed_ms = (time.perf_counter() - start) * 1000
-
-    avg_ms = elapsed_ms / num_iterations
-    # If __repr__ were being called, each would take ~100ms
-    # We expect <1ms per call with the fix
-    assert avg_ms < 1.0, (
-        f"hasattr for non-existent attribute took {avg_ms:.2f}ms on average. "
-        "Expected <1ms. This suggests __repr__ is being triggered."
-    )
-
-
-def test_hasattr_existing_attribute_works():
-    """Verify that hasattr for existing attributes returns True and works correctly."""
-    model = SlowReprModel()
-    wrapper = MockCUDAGraphWrapper(model)
-
-    # 'forward' exists on nn.Module
-    assert hasattr(wrapper, "forward") is True
-
-    # 'linear' exists on our model
-    assert hasattr(wrapper, "linear") is True
-
-    # Can actually access the attribute
-    linear = wrapper.linear
-    assert isinstance(linear, nn.Linear)
-
-
-def test_getattr_existing_attribute_returns_value():
-    """Verify that getattr for existing attributes returns the correct value."""
-    model = SlowReprModel()
-    wrapper = MockCUDAGraphWrapper(model)
-
-    # Access forward method
-    forward_method = wrapper.forward
-    assert callable(forward_method)
-
-    # Access linear layer
-    linear = wrapper.linear
-    assert isinstance(linear, nn.Linear)
-    assert linear.in_features == 16
-    assert linear.out_features == 16
-
-
-def test_getattr_nonexistent_raises_attribute_error():
-    """Verify that getattr for non-existent attributes raises AttributeError."""
-    model = SlowReprModel()
-    wrapper = MockCUDAGraphWrapper(model)
-
-    with pytest.raises(AttributeError) as exc_info:
-        _ = wrapper.nonexistent_attribute
-
-    # Verify error message format (should NOT contain model repr)
-    error_msg = str(exc_info.value)
-    assert "nonexistent_attribute" in error_msg
-    assert "cudagraph wrapper" in error_msg
-    # Should NOT contain the slow repr output
-    assert "SlowReprModel(delay=" not in error_msg
-
-
-def test_attribute_error_message_does_not_contain_runnable_repr():
-    """Explicitly verify the error message doesn't trigger runnable repr."""
-    model = SlowReprModel(repr_delay_ms=100.0)
-    wrapper = MockCUDAGraphWrapper(model)
-    model.repr_call_count = 0
-
-    try:
-        _ = wrapper.nonexistent_attr
-    except AttributeError:
-        pass
-
-    # __repr__ should not have been called during error construction
-    assert model.repr_call_count == 0, (
-        "AttributeError message construction triggered __repr__. The error message should not include {self.runnable}."
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_real_cudagraph_wrapper_hasattr_performance():
-    """Test the actual CUDAGraphWrapper from vllm_omni (requires CUDA)."""
-    from vllm.config import CUDAGraphMode
-
-    model = SlowReprModel(repr_delay_ms=50.0).cuda()
-    model.repr_call_count = 0
-
-    # Create actual CUDAGraphWrapper
-    try:
-        wrapper = CUDAGraphWrapper(model, runtime_mode=CUDAGraphMode.NONE)
-    except Exception:
-        pytest.skip("Could not create CUDAGraphWrapper")
-
-    # Test hasattr performance
-    num_iterations = 50
-    start = time.perf_counter()
-    for _ in range(num_iterations):
-        hasattr(wrapper, "nonexistent_xyz")
-    elapsed_ms = (time.perf_counter() - start) * 1000
-
-    avg_ms = elapsed_ms / num_iterations
-    assert avg_ms < 1.0, f"Real CUDAGraphWrapper hasattr took {avg_ms:.2f}ms avg. Expected <1ms with the optimization."
-    assert model.repr_call_count == 0, f"__repr__ called {model.repr_call_count} times"

vllm_omni/__init__.py

@@ -24,11 +24,27 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
 
 from .config import OmniModelConfig
-from .entrypoints import AsyncOmni, Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 
 
+def __getattr__(name: str):
+    # Lazy import for AsyncOmni and Omni to avoid pulling in heavy
+    # dependencies (vllm model_loader → fused_moe → pynvml) at package
+    # import time.  This prevents crashes in lightweight subprocesses
+    # (e.g. model-architecture inspection) that lack a CUDA context.
+    # See: https://github.com/vllm-project/vllm-omni/issues/1793
+    if name == "AsyncOmni":
+        from .entrypoints.async_omni import AsyncOmni
+
+        return AsyncOmni
+    if name == "Omni":
+        from .entrypoints.omni import Omni
+
+        return Omni
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "__version__",
     "__version_tuple__",

vllm_omni/benchmarks/metrics/metrics.py

@@ -185,7 +185,7 @@ def calculate_metrics(
                     # Note : this may inflate the output token count slightly
                     output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
             actual_output_lens.append(output_len)
-            total_input += input_requests[i].prompt_len
+            total_input += outputs[i].prompt_len
             tpot = 0
             if output_len > 1:
                 latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft

vllm_omni/benchmarks/patch/patch.py

@@ -190,6 +190,10 @@ async def async_request_openai_chat_omni_completions(
                                 if metrics := data.get("metrics"):
                                     output.output_tokens = metrics.get("num_tokens_out", 0)
 
+                                if usage := data.get("usage"):
+                                    if (pt := usage.get("prompt_tokens")) is not None:
+                                        output.prompt_len = pt
+
                     output.latency = timestamp - st
                     output.generated_text = generated_text
                     if generated_audio is not None:

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -209,3 +209,13 @@ class FlashAttentionImpl(AttentionImpl):
             layout="BNSD",
         )
         return output
+
+    def forward_musa(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # XXX (MUSA): MUSA uses the same implementation as XPU (mate only provides flash_attn_varlen_func)
+        return self.forward_xpu(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
+from functools import lru_cache
+
 import torch
 import torch.nn.functional as F
 
@@ -38,8 +40,10 @@ elif current_omni_platform.is_xpu():
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_musa():
-    # XXX (MUSA): Add MUSA-specific Flash Attention when available
-    pass
+    try:
+        from mate import flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
 else:
     # CUDA: try FA3 -> FA2 fallback chain
     # Try FA3 from fa3-fwd PyPI package
@@ -76,6 +80,12 @@ else:
 HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
+@lru_cache(maxsize=1)
+def is_mate_available() -> bool:
+    """Check if MATE (MUSA Flash Attention) is available."""
+    return current_omni_platform.is_musa() and flash_attn_varlen_func is not None
+
+
 def _index_first_axis(tensor, indices):
     """
     A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis,

vllm_omni/diffusion/cache/teacache/backend.py

@@ -48,16 +48,7 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
         coefficients=config.coefficients,
     )
     transformer = pipeline.bagel
-    original_forward_flow = transformer._forward_flow
-
-    import types
-
-    def forward_alias(self, *args, **kwargs):
-        return original_forward_flow(*args, **kwargs)
-
-    transformer.forward = types.MethodType(forward_alias, transformer)
     apply_teacache_hook(transformer, teacache_config)
-    transformer._forward_flow = transformer.forward
     pipeline.transformer = transformer
 
     logger.info(

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import types
 from typing import Any
 
 import numpy as np
@@ -74,15 +73,8 @@ class BagelAdapter:
 
     @staticmethod
     def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
-        original_forward_flow = transformer._forward_flow
-
-        def forward_alias(self, *args, **kwargs):
-            return original_forward_flow(*args, **kwargs)
-
-        transformer.forward = types.MethodType(forward_alias, transformer)
         registry = HookRegistry.get_or_create(transformer)
         registry.register_hook(hook._HOOK_NAME, hook)
-        transformer._forward_flow = transformer.forward
 
 
 class StableAudioAdapter:

vllm_omni/diffusion/data.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass, field, fields
 from typing import TYPE_CHECKING, Any
 
 import torch
+from PIL import Image
 from pydantic import model_validator
 from typing_extensions import Self
 from vllm.config.utils import config
@@ -701,10 +702,12 @@ class DiffusionOutput:
     Final output (after pipeline completion)
     """
 
-    output: torch.Tensor | None = None
-    trajectory_timesteps: list[torch.Tensor] | None = None
-    trajectory_latents: torch.Tensor | None = None
-    trajectory_decoded: list[torch.Tensor] | None = None
+    # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm
+    output: torch.Tensor | dict | None = None
+    trajectory_timesteps: torch.Tensor | dict | None = None
+    trajectory_latents: torch.Tensor | dict | None = None
+    trajectory_log_probs: torch.Tensor | dict | None = None
+    trajectory_decoded: list[Image.Image] | None = None
     error: str | None = None
     aborted: bool = False
     abort_message: str | None = None

vllm_omni/diffusion/diffusion_engine.py

@@ -145,8 +145,12 @@ class DiffusionEngine:
         postprocess_start_time = time.perf_counter()
         outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
         audio_payload = None
+        model_audio_sample_rate = None
+        model_fps = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")
+            model_audio_sample_rate = outputs.get("audio_sample_rate")
+            model_fps = outputs.get("fps")
             outputs = outputs.get("video", outputs)
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
@@ -192,6 +196,10 @@ class DiffusionEngine:
                         prompt=prompt,
                         metrics=metrics,
                         latents=output.trajectory_latents,
+                        trajectory_latents=output.trajectory_latents,
+                        trajectory_timesteps=output.trajectory_timesteps,
+                        trajectory_log_probs=output.trajectory_log_probs,
+                        trajectory_decoded=output.trajectory_decoded,
                         multimodal_output={"audio": request_audio_payload},
                         final_output_type="audio",
                         stage_durations=output.stage_durations,
@@ -202,6 +210,10 @@ class DiffusionEngine:
                 mm_output = {}
                 if audio_payload is not None:
                     mm_output["audio"] = audio_payload
+                if model_audio_sample_rate is not None:
+                    mm_output["audio_sample_rate"] = model_audio_sample_rate
+                if model_fps is not None:
+                    mm_output["fps"] = model_fps
                 return [
                     OmniRequestOutput.from_diffusion(
                         request_id=request_id,
@@ -209,6 +221,10 @@ class DiffusionEngine:
                         prompt=prompt,
                         metrics=metrics,
                         latents=output.trajectory_latents,
+                        trajectory_latents=output.trajectory_latents,
+                        trajectory_timesteps=output.trajectory_timesteps,
+                        trajectory_log_probs=output.trajectory_log_probs,
+                        trajectory_decoded=output.trajectory_decoded,
                         custom_output=output.custom_output or {},
                         multimodal_output=mm_output,
                         stage_durations=output.stage_durations,
@@ -240,6 +256,10 @@ class DiffusionEngine:
                             prompt=prompt,
                             metrics=metrics,
                             latents=output.trajectory_latents,
+                            trajectory_latents=output.trajectory_latents,
+                            trajectory_timesteps=output.trajectory_timesteps,
+                            trajectory_log_probs=output.trajectory_log_probs,
+                            trajectory_decoded=output.trajectory_decoded,
                             multimodal_output={"audio": request_audio_payload},
                             final_output_type="audio",
                             stage_durations=output.stage_durations,
@@ -260,6 +280,10 @@ class DiffusionEngine:
                                 if num_outputs == 1:
                                     sliced_audio = sliced_audio[0]
                         mm_output["audio"] = sliced_audio
+                    if model_audio_sample_rate is not None:
+                        mm_output["audio_sample_rate"] = model_audio_sample_rate
+                    if model_fps is not None:
+                        mm_output["fps"] = model_fps
                     results.append(
                         OmniRequestOutput.from_diffusion(
                             request_id=request_id,
@@ -267,6 +291,10 @@ class DiffusionEngine:
                             prompt=prompt,
                             metrics=metrics,
                             latents=output.trajectory_latents,
+                            trajectory_latents=output.trajectory_latents,
+                            trajectory_timesteps=output.trajectory_timesteps,
+                            trajectory_log_probs=output.trajectory_log_probs,
+                            trajectory_decoded=output.trajectory_decoded,
                             custom_output=output.custom_output or {},
                             multimodal_output=mm_output,
                             stage_durations=output.stage_durations,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
+class OmniAutoencoderKLWan(AutoencoderKLWan):
+    def _execution_context(self):
+        try:
+            first_param = next(self.parameters())
+        except StopIteration:
+            return nullcontext()
+
+        dtype = first_param.dtype
+        if dtype not in (torch.float16, torch.bfloat16):
+            return nullcontext()
+
+        return current_omni_platform.create_autocast_context(
+            device_type=first_param.device.type,
+            dtype=dtype,
+            enabled=True,
+        )
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().encode(x, return_dict=return_dict)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().decode(z, return_dict=return_dict)
+
+
+class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
     @classmethod
     def from_pretrained(cls, *args: Any, **kwargs: Any):
         model = super().from_pretrained(*args, **kwargs)
@@ -84,11 +112,12 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         """Decode a single latent tile into RGB space."""
         self.clear_cache()
         time = []
-        for k in range(len(task.tensor)):
-            self._conv_idx = [0]
-            tile = self.post_quant_conv(task.tensor[k])
-            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
-            time.append(decoded)
+        with self._execution_context():
+            for k in range(len(task.tensor)):
+                self._conv_idx = [0]
+                tile = self.post_quant_conv(task.tensor[k])
+                decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+                time.append(decoded)
         result = torch.cat(time, dim=2)
         return result
 

vllm_omni/diffusion/envs.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.attention.backends.utils.fa import is_mate_available
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -52,6 +53,10 @@ class PackagesEnvChecker:
         """Check if flash attention is available and compatible."""
         platform = current_omni_platform
 
+        # MUSA uses MATE for flash attention
+        if platform.is_musa():
+            return is_mate_available()
+
         # Flash attention requires CUDA-like platforms (CUDA or ROCm)
         if not platform.is_cuda_alike():
             return False

vllm_omni/diffusion/ipc.py

@@ -78,13 +78,29 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     return tensor
 
 
+def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict:
+    """Replace a tensor with an SHM handle if it exceeds the threshold."""
+    if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD:
+        return _tensor_to_shm(val)
+    return val
+
+
+def _unpack_if_shm_handle(val: object) -> object:
+    """Reconstruct a tensor from an SHM handle dict, or return as-is."""
+    if isinstance(val, dict) and val.get("__tensor_shm__"):
+        return _tensor_from_shm(val)
+    return val
+
+
 def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     if output.output is not None and isinstance(output.output, torch.Tensor):
-        if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
-            output.output = _tensor_to_shm(output.output)
+        output.output = _pack_tensor_if_large(output.output)
     if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
-        if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
-            output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
+        output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents)
+    if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor):
+        output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps)
+    if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
+        output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs)
     return output
 
 
@@ -104,10 +120,10 @@ def pack_diffusion_output_shm(output: object) -> object:
 
 
 def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
-    if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
-        output.output = _tensor_from_shm(output.output)
-    if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
-        output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
+    output.output = _unpack_if_shm_handle(output.output)
+    output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents)
+    output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps)
+    output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs)
     return output
 
 

vllm_omni/diffusion/layers/rope.py

@@ -145,6 +145,14 @@ class RotaryEmbedding(CustomOp):
     ) -> torch.Tensor:
         return self.forward_native(x, cos, sin)
 
+    def forward_musa(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x, cos, sin)
+
     def forward_native(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1655,6 +1655,9 @@ class Bagel(nn.Module):
         cfg_img_past_key_values: NaiveCache | None = None,
         cfg_img_key_values_lens: torch.IntTensor | None = None,
         cfg_img_packed_key_value_indexes: torch.LongTensor | None = None,
+        return_trajectory_latents: bool = False,
+        scheduler: object | None = None,
+        scheduler_kwargs: dict | None = None,
     ):
         x_t = packed_init_noises
 
@@ -1663,6 +1666,14 @@ class Bagel(nn.Module):
         dts = timesteps[:-1] - timesteps[1:]
         timesteps = timesteps[:-1]
 
+        # Optional trajectory recording for RL rollout data collection
+        trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
+        trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
+        trajectory_log_probs: list[torch.Tensor] | None = (
+            [] if (return_trajectory_latents and scheduler is not None) else None
+        )
+        _sched_kw = scheduler_kwargs or {}
+
         use_cfg_text = cfg_text_scale > 1.0
         use_cfg_img = cfg_img_scale > 1.0
 
@@ -1699,6 +1710,9 @@ class Bagel(nn.Module):
                 cfg_img_past_key_values=cfg_img_past_key_values,
                 cfg_img_key_values_lens=cfg_img_key_values_lens,
                 cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes,
+                return_trajectory_latents=return_trajectory_latents,
+                scheduler=scheduler,
+                scheduler_kwargs=scheduler_kwargs,
             )
 
         # ── SP + CFG: sequential single-branch forwards ──
@@ -1720,7 +1734,7 @@ class Bagel(nn.Module):
                     packed_seqlens=packed_seqlens,
                 )
 
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     **common,
                     packed_indexes=packed_indexes,
                     packed_position_ids=packed_position_ids,
@@ -1730,7 +1744,7 @@ class Bagel(nn.Module):
                 )
 
                 if cfg_text_scale_ > 1.0:
-                    cfg_text_v_t = self._forward_flow_single_branch(
+                    cfg_text_v_t = self.forward_single_branch(
                         **common,
                         packed_indexes=cfg_text_packed_query_indexes,
                         packed_position_ids=cfg_text_packed_position_ids,
@@ -1740,7 +1754,7 @@ class Bagel(nn.Module):
                     )
                     cfg_img_v_t = None
                     if cfg_img_scale_ > 1.0:
-                        cfg_img_v_t = self._forward_flow_single_branch(
+                        cfg_img_v_t = self.forward_single_branch(
                             **common,
                             packed_indexes=cfg_img_packed_query_indexes,
                             packed_position_ids=cfg_img_packed_position_ids,
@@ -1758,16 +1772,25 @@ class Bagel(nn.Module):
                         cfg_renorm_min,
                     )
 
-                x_t = x_t - v_t.to(x_t.device) * dts[i]
+                if scheduler is not None:
+                    out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
+                    x_t = out.prev_sample
+                    if trajectory_log_probs is not None and out.log_prob is not None:
+                        trajectory_log_probs.append(out.log_prob)
+                else:
+                    x_t = x_t - v_t.to(x_t.device) * dts[i]
+                if return_trajectory_latents:
+                    trajectory_latents.append(x_t.clone())
+                    trajectory_timesteps.append(timesteps[i] - dts[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
-            return unpacked_latent
+            return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
         # ── SP without CFG: direct single-branch loop ──
         if use_sp:
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -1781,10 +1804,20 @@ class Bagel(nn.Module):
                     past_key_values=past_key_values,
                     packed_key_value_indexes=packed_key_value_indexes,
                 )
-                x_t = x_t - v_t.to(x_t.device) * dts[i]
+                if scheduler is not None:
+                    out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
+                    x_t = out.prev_sample
+                    out_log_prob = getattr(out, "log_prob", None)
+                    if trajectory_log_probs is not None and out_log_prob is not None:
+                        trajectory_log_probs.append(out_log_prob)
+                else:
+                    x_t = x_t - v_t.to(x_t.device) * dts[i]
+                if return_trajectory_latents:
+                    trajectory_latents.append(x_t.clone())
+                    trajectory_timesteps.append(timesteps[i] - dts[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
-            return unpacked_latent
+            return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
         # ── Batched CFG mode (cfg_parallel_size=1, no SP) ──
         cfg_batched = None
@@ -1850,7 +1883,7 @@ class Bagel(nn.Module):
             else:
                 cfg_text_scale_ = 1.0
                 cfg_img_scale_ = 1.0
-            v_t = self._forward_flow(
+            v_t = self.forward(
                 x_t=x_t,
                 timestep=timestep,
                 packed_vae_token_indexes=packed_vae_token_indexes,
@@ -1870,10 +1903,19 @@ class Bagel(nn.Module):
                 cfg_batched=cfg_batched,
             )
 
-            x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
+            if scheduler is not None:
+                out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
+                x_t = out.prev_sample
+                if trajectory_log_probs is not None and out.log_prob is not None:
+                    trajectory_log_probs.append(out.log_prob)
+            else:
+                x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
+            if return_trajectory_latents:
+                trajectory_latents.append(x_t.clone())
+                trajectory_timesteps.append(timesteps[i] - dts[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
-        return unpacked_latent
+        return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
     def _generate_image_parallel(
         self,
@@ -1905,6 +1947,9 @@ class Bagel(nn.Module):
         cfg_img_past_key_values: NaiveCache | None,
         cfg_img_key_values_lens: torch.IntTensor | None,
         cfg_img_packed_key_value_indexes: torch.LongTensor | None,
+        return_trajectory_latents: bool = False,
+        scheduler: object | None = None,
+        scheduler_kwargs: dict | None = None,
     ):
         """CFG parallel denoising loop: each rank computes one CFG branch.
 
@@ -1961,13 +2006,20 @@ class Bagel(nn.Module):
         else:
             raise RuntimeError(f"Unexpected cfg_rank={cfg_rank} for Bagel 3-branch CFG parallel")
 
+        trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
+        trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
+        trajectory_log_probs: list[torch.Tensor] | None = (
+            [] if (return_trajectory_latents and scheduler is not None) else None
+        )
+        _sched_kw = scheduler_kwargs or {}
+
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
             use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
 
             if use_cfg_this_step:
                 # CFG interval: each rank computes its own branch
-                local_v_t = self._forward_flow_single_branch(
+                local_v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -1994,7 +2046,7 @@ class Bagel(nn.Module):
                 )
             else:
                 # Outside CFG interval: all ranks compute with gen inputs, no comm
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -2009,10 +2061,19 @@ class Bagel(nn.Module):
                     packed_key_value_indexes=packed_key_value_indexes,
                 )
 
-            x_t = x_t - v_t.to(x_t.device) * dts[i]
+            if scheduler is not None:
+                out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
+                x_t = out.prev_sample
+                if trajectory_log_probs is not None and out.log_prob is not None:
+                    trajectory_log_probs.append(out.log_prob)
+            else:
+                x_t = x_t - v_t.to(x_t.device) * dts[i]
+            if return_trajectory_latents:
+                trajectory_latents.append(x_t.clone())
+                trajectory_timesteps.append(timesteps[i] - dts[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
-        return unpacked_latent
+        return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
     @staticmethod
     def _combine_cfg(
@@ -2067,7 +2128,7 @@ class Bagel(nn.Module):
 
         return v_t
 
-    def _forward_flow_single_branch(
+    def forward_single_branch(
         self,
         x_t: torch.Tensor,
         timestep: torch.LongTensor,
@@ -2197,7 +2258,7 @@ class Bagel(nn.Module):
         v_t = v_t[packed_vae_token_indexes]
         return v_t
 
-    def _forward_flow(
+    def forward(
         self,
         x_t: torch.Tensor,
         timestep: torch.LongTensor,

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -159,6 +159,9 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self.od_config = od_config
         self.device = get_local_device()
 
+        self.scheduler: object | None = None
+        self.scheduler_kwargs: dict = {}
+
         model = od_config.model
         local_files_only = os.path.exists(model)
         if local_files_only:
@@ -630,7 +633,7 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             enabled=self.device.type != "cpu",
             dtype=self.od_config.dtype,
         ):
-            latents = self.bagel.generate_image(
+            latents, trajectory_latents, trajectory_timesteps, trajectory_log_probs = self.bagel.generate_image(
                 past_key_values=gen_context["past_key_values"],
                 cfg_text_past_key_values=cfg_text_context["past_key_values"],
                 cfg_img_past_key_values=cfg_img_context["past_key_values"],
@@ -650,11 +653,36 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"],
                 cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
                 cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
+                return_trajectory_latents=req.sampling_params.return_trajectory_latents,
+                scheduler=self.scheduler,
+                scheduler_kwargs=self.scheduler_kwargs,
             )
 
         img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
+
+        # Build trajectory output when requested
+        trajectory_latents_stacked: torch.Tensor | None = None
+        trajectory_timesteps_stacked: torch.Tensor | None = None
+        trajectory_decoded: list[Image.Image] | None = None
+        if trajectory_latents:
+            trajectory_latents_stacked = torch.stack(trajectory_latents)
+            trajectory_timesteps_stacked = torch.stack(trajectory_timesteps)
+            if req.sampling_params.return_trajectory_decoded:
+                trajectory_decoded = [
+                    self._decode_image_from_latent(self.bagel, self.vae, lat, image_shape) for lat in trajectory_latents
+                ]
+
+        trajectory_log_probs_stacked: torch.Tensor | None = None
+        if trajectory_log_probs:
+            trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
+
         return DiffusionOutput(
-            output=img, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+            output=img,
+            trajectory_latents=trajectory_latents_stacked,
+            trajectory_timesteps=trajectory_timesteps_stacked,
+            trajectory_log_probs=trajectory_log_probs_stacked,
+            trajectory_decoded=trajectory_decoded,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py

@@ -582,6 +582,7 @@ class LongCatImageTransformer2DModel(nn.Module):
     """
 
     _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     # Sequence Parallelism for LongCat (following diffusers' _cp_plan pattern)
     _sp_plan = {

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -28,8 +28,6 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
-    get_cfg_group,
-    get_classifier_free_guidance_rank,
     get_classifier_free_guidance_world_size,
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -122,6 +120,31 @@ def calculate_shift(
     return mu
 
 
+class _VideoAudioScheduler:
+    """Composite scheduler dispatching to video and audio schedulers."""
+
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(
+            noise_pred[0],
+            t[0],
+            latents[0],
+            return_dict=False,
+            generator=generator,
+        )[0]
+        audio_out = self.audio_scheduler.step(
+            noise_pred[1],
+            t[1],
+            latents[1],
+            return_dict=False,
+            generator=generator,
+        )[0]
+        return ((video_out, audio_out),)
+
+
 class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
     def __init__(
         self,
@@ -542,6 +565,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
         return latents
 
+    @staticmethod
+    def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
+        return latents[:, :num_frames]
+
     def prepare_latents(
         self,
         batch_size: int = 1,
@@ -597,25 +624,49 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         noise_scale: float = 0.0,
         dtype: torch.dtype | None = None,
         device: torch.device | None = None,
-        generator: torch.Generator | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, int]:
+    ) -> tuple[torch.Tensor, int, int]:
+        original_latent_length = audio_latent_length
+        padded_latent_length = original_latent_length
+
+        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
+
+        sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
+        if sp_size > 1:
+            padded_latent_length += (sp_size - (original_latent_length % sp_size)) % sp_size
+
         if latents is not None:
             if latents.ndim == 4:
                 # latents are of shape [B, C, L, M], need to be packed
                 latents = self._pack_audio_latents(latents)
             if latents.ndim != 3:
                 raise ValueError(
-                    f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]."  # noqa
+                    f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is "
+                    "[batch_size, num_seq, num_features] or [batch_size, num_channels, audio_length, mel_bins]."
                 )
             latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std)
             latents = self._create_noised_state(latents, noise_scale, generator)
-            return latents.to(device=device, dtype=dtype)
 
-        # TODO: confirm whether this logic is correct
-        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
+            if latents.shape[1] not in {original_latent_length, padded_latent_length}:
+                raise ValueError(
+                    "Provided `audio_latents` has incompatible audio frame count "
+                    f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
+                )
 
-        shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
+            if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
+                padding = torch.zeros(
+                    latents.shape[0],
+                    padded_latent_length - original_latent_length,
+                    latents.shape[2],
+                    dtype=latents.dtype,
+                    device=latents.device,
+                )
+                latents = torch.cat([latents, padding], dim=1)
+
+            return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
+
+        shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
 
         if isinstance(generator, list) and len(generator) != batch_size:
             raise ValueError(
@@ -625,7 +676,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
 
         latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         latents = self._pack_audio_latents(latents)
-        return latents
+        return latents, original_latent_length, padded_latent_length
 
     @property
     def guidance_scale(self):
@@ -655,147 +706,44 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
     def interrupt(self):
         return self._interrupt
 
-    def _is_cfg_parallel_enabled(self, do_true_cfg: bool) -> bool:
-        return do_true_cfg and get_classifier_free_guidance_world_size() > 1
-
     def _transformer_cache_context(self, context_name: str):
         cache_context = getattr(self.transformer, "cache_context", None)
         if callable(cache_context):
             return cache_context(context_name)
         return nullcontext()
 
-    def _predict_noise_av(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+    def predict_noise(self, **kwargs):
         with self._transformer_cache_context("cond_uncond"):
             noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
-        return noise_pred_video, noise_pred_audio
-
-    def predict_noise_av_maybe_with_cfg(
-        self,
-        do_true_cfg: bool,
-        true_cfg_scale: float,
-        positive_kwargs: dict[str, Any],
-        negative_kwargs: dict[str, Any] | None,
-        guidance_rescale: float = 0.0,
-        cfg_normalize: bool = False,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        if do_true_cfg:
-            cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1
-
-            if cfg_parallel_ready:
-                cfg_group = get_cfg_group()
-                cfg_rank = get_classifier_free_guidance_rank()
-
-                if cfg_rank == 0:
-                    noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
-                else:
-                    noise_pred_video, noise_pred_audio = self._predict_noise_av(**negative_kwargs)
-
-                noise_pred_video = noise_pred_video.float()
-                noise_pred_audio = noise_pred_audio.float()
-
-                gathered_video = cfg_group.all_gather(noise_pred_video, separate_tensors=True)
-                gathered_audio = cfg_group.all_gather(noise_pred_audio, separate_tensors=True)
-
-                if cfg_rank == 0:
-                    noise_pred_video_text = gathered_video[0]
-                    noise_pred_video_uncond = gathered_video[1]
-                    noise_pred_audio_text = gathered_audio[0]
-                    noise_pred_audio_uncond = gathered_audio[1]
-
-                    noise_pred_video = self.combine_cfg_noise(
-                        noise_pred_video_text,
-                        noise_pred_video_uncond,
-                        true_cfg_scale,
-                        cfg_normalize,
-                    )
-                    noise_pred_audio = self.combine_cfg_noise(
-                        noise_pred_audio_text,
-                        noise_pred_audio_uncond,
-                        true_cfg_scale,
-                        cfg_normalize,
-                    )
-
-                    if guidance_rescale > 0:
-                        noise_pred_video = rescale_noise_cfg(
-                            noise_pred_video,
-                            noise_pred_video_text,
-                            guidance_rescale=guidance_rescale,
-                        )
-                        noise_pred_audio = rescale_noise_cfg(
-                            noise_pred_audio,
-                            noise_pred_audio_text,
-                            guidance_rescale=guidance_rescale,
-                        )
-                    return noise_pred_video, noise_pred_audio
-                return None, None
-
-            noise_pred_video_text, noise_pred_audio_text = self._predict_noise_av(**positive_kwargs)
-            noise_pred_video_uncond, noise_pred_audio_uncond = self._predict_noise_av(**negative_kwargs)
-
-            noise_pred_video_text = noise_pred_video_text.float()
-            noise_pred_audio_text = noise_pred_audio_text.float()
-            noise_pred_video_uncond = noise_pred_video_uncond.float()
-            noise_pred_audio_uncond = noise_pred_audio_uncond.float()
-
-            noise_pred_video = self.combine_cfg_noise(
-                noise_pred_video_text,
-                noise_pred_video_uncond,
-                true_cfg_scale,
-                cfg_normalize,
-            )
-            noise_pred_audio = self.combine_cfg_noise(
-                noise_pred_audio_text,
-                noise_pred_audio_uncond,
-                true_cfg_scale,
-                cfg_normalize,
-            )
-
-            if guidance_rescale > 0:
-                noise_pred_video = rescale_noise_cfg(
-                    noise_pred_video,
-                    noise_pred_video_text,
-                    guidance_rescale=guidance_rescale,
-                )
-                noise_pred_audio = rescale_noise_cfg(
-                    noise_pred_audio,
-                    noise_pred_audio_text,
-                    guidance_rescale=guidance_rescale,
-                )
-
-            return noise_pred_video, noise_pred_audio
-
-        noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
         return noise_pred_video.float(), noise_pred_audio.float()
 
-    def _scheduler_step_video_audio_maybe_with_cfg(
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize=False):
+        """Per-element CFG combine with guidance_rescale support."""
+        (video_pos, audio_pos) = positive_noise_pred
+        (video_neg, audio_neg) = negative_noise_pred
+        video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize)
+        audio_combined = super().combine_cfg_noise(audio_pos, audio_neg, true_cfg_scale, cfg_normalize)
+        if self._guidance_rescale and self._guidance_rescale > 0:
+            video_combined = rescale_noise_cfg(video_combined, video_pos, guidance_rescale=self._guidance_rescale)
+            audio_combined = rescale_noise_cfg(audio_combined, audio_pos, guidance_rescale=self._guidance_rescale)
+        return (video_combined, audio_combined)
+
+    def _synchronize_cfg_parallel_step_output(
         self,
-        noise_pred_video: torch.Tensor | None,
-        noise_pred_audio: torch.Tensor | None,
-        t: torch.Tensor,
-        latents: torch.Tensor,
-        audio_latents: torch.Tensor,
-        audio_scheduler: FlowMatchEulerDiscreteScheduler,
+        latents: tuple[torch.Tensor, torch.Tensor],
         do_true_cfg: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        cfg_parallel_ready = self._is_cfg_parallel_enabled(do_true_cfg)
+        if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
+            return latents
 
-        if cfg_parallel_ready:
-            cfg_group = get_cfg_group()
-            cfg_rank = get_classifier_free_guidance_rank()
-
-            if cfg_rank == 0:
-                latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
-                audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
-
-            latents = latents.contiguous()
-            audio_latents = audio_latents.contiguous()
-            cfg_group.broadcast(latents, src=0)
-            cfg_group.broadcast(audio_latents, src=0)
-            return latents, audio_latents
-
-        latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
-        audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
-        return latents, audio_latents
+        # Without this sync, CUDA async execution causes non-deterministic
+        # numerical drift across denoising steps in CFG parallel mode,
+        # producing different video outputs across runs.
+        latents = tuple(tensor.contiguous() for tensor in latents)
+        device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
+        if device is not None:
+            torch.cuda.current_stream(device).synchronize()
+        return latents
 
     @torch.no_grad()
     def forward(
@@ -828,6 +776,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         attention_kwargs: dict[str, Any] | None = None,
         max_sequence_length: int | None = None,
     ) -> DiffusionOutput:
+        # Extract prompt/negative_prompt from request.
+        # Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
@@ -869,6 +819,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             else req.sampling_params.extra_args.get("audio_latents", audio_latents)
         )
 
+        # Override with pre-computed embeddings if provided in request.
         req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
         if any(p is not None for p in req_prompt_embeds):
             prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore[arg-type]
@@ -939,20 +890,17 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             max_sequence_length=max_sequence_length,
             device=device,
         )
-        cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
+        # Compute positive prompt connectors
         additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
             prompt_embeds, additive_attention_mask, additive_mask=True
         )
 
+        # Compute negative prompt connectors when CFG is enabled
         negative_connector_prompt_embeds = None
         negative_connector_audio_prompt_embeds = None
         negative_connector_attention_mask = None
-        if cfg_parallel_ready:
+        if self.do_classifier_free_guidance:
             negative_additive_attention_mask = (
                 1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
             ) * -1000000.0
@@ -1027,20 +975,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         num_channels_latents_audio = (
             self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
         )
-
-        # padding audio_latents if needed
-        sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
-        if sp_size > 1:
-            pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
-            if pad_len > 0:
-                if audio_latents is not None:
-                    pad_shape = list(audio_latents.shape)
-                    pad_shape[2] = pad_len
-                    padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
-                    audio_latents = torch.cat([audio_latents, padding], dim=2)
-                audio_num_frames += pad_len
-
-        audio_latents = self.prepare_audio_latents(
+        audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
             batch_size * num_videos_per_prompt,
             num_channels_latents=num_channels_latents_audio,
             audio_latent_length=audio_num_frames,
@@ -1061,6 +996,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             self.scheduler.config.get("max_shift", 2.05),
         )
         audio_scheduler = copy.deepcopy(self.scheduler)
+        video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler)
         _ = retrieve_timesteps(
             audio_scheduler,
             num_inference_steps,
@@ -1083,12 +1019,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
         )
         audio_coords = self.transformer.audio_rope.prepare_audio_coords(
-            audio_latents.shape[0], audio_num_frames, audio_latents.device
+            audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
         )
-        # Duplicate the positional ids as well if using CFG
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1))  # Repeat twice in batch dim
-            audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
+        # No coord duplication needed: mixin handles CFG via separate forward calls,
+        # not batch=2. Each forward gets batch=1 coords directly.
 
         with self.progress_bar(total=len(timesteps)) as pbar:
             for i, t in enumerate(timesteps):
@@ -1097,119 +1031,60 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
 
                 self._current_timestep = t
 
-                if cfg_parallel_ready:
-                    latent_model_input = latents.to(prompt_embeds.dtype)
-                    audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
-                    timestep = t.expand(latent_model_input.shape[0])
+                latent_model_input = latents.to(prompt_embeds.dtype)
+                audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
+                timestep = t.expand(latent_model_input.shape[0])
+                do_true_cfg = self.do_classifier_free_guidance
 
-                    positive_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "audio_hidden_states": audio_latent_model_input,
-                        "encoder_hidden_states": connector_prompt_embeds,
-                        "audio_encoder_hidden_states": connector_audio_prompt_embeds,
-                        "timestep": timestep,
-                        "encoder_attention_mask": connector_attention_mask,
-                        "audio_encoder_attention_mask": connector_attention_mask,
-                        "num_frames": latent_num_frames,
-                        "height": latent_height,
-                        "width": latent_width,
-                        "fps": frame_rate,
-                        "audio_num_frames": audio_num_frames,
-                        "video_coords": video_coords,
-                        "audio_coords": audio_coords,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
-                    }
-                    negative_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "audio_hidden_states": audio_latent_model_input,
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "audio_hidden_states": audio_latent_model_input,
+                    "encoder_hidden_states": connector_prompt_embeds,
+                    "audio_encoder_hidden_states": connector_audio_prompt_embeds,
+                    "timestep": timestep,
+                    "encoder_attention_mask": connector_attention_mask,
+                    "audio_encoder_attention_mask": connector_attention_mask,
+                    "num_frames": latent_num_frames,
+                    "height": latent_height,
+                    "width": latent_width,
+                    "fps": frame_rate,
+                    "audio_num_frames": padded_audio_num_frames,
+                    "video_coords": video_coords,
+                    "audio_coords": audio_coords,
+                    "attention_kwargs": attention_kwargs,
+                    "return_dict": False,
+                }
+                negative_kwargs = (
+                    {
+                        **positive_kwargs,
                         "encoder_hidden_states": negative_connector_prompt_embeds,
                         "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
-                        "timestep": timestep,
                         "encoder_attention_mask": negative_connector_attention_mask,
                         "audio_encoder_attention_mask": negative_connector_attention_mask,
-                        "num_frames": latent_num_frames,
-                        "height": latent_height,
-                        "width": latent_width,
-                        "fps": frame_rate,
-                        "audio_num_frames": audio_num_frames,
-                        "video_coords": video_coords,
-                        "audio_coords": audio_coords,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
                     }
+                    if do_true_cfg
+                    else None
+                )
 
-                    noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
-                        do_true_cfg=True,
-                        true_cfg_scale=guidance_scale,
-                        positive_kwargs=positive_kwargs,
-                        negative_kwargs=negative_kwargs,
-                        guidance_rescale=guidance_rescale,
-                        cfg_normalize=False,
-                    )
+                noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=False,
+                )
 
-                    latents, audio_latents = self._scheduler_step_video_audio_maybe_with_cfg(
-                        noise_pred_video,
-                        noise_pred_audio,
-                        t,
-                        latents,
-                        audio_latents,
-                        audio_scheduler,
-                        do_true_cfg=True,
-                    )
-                else:
-                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                    latent_model_input = latent_model_input.to(prompt_embeds.dtype)
-                    audio_latent_model_input = (
-                        torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
-                    )
-                    audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
-
-                    timestep = t.expand(latent_model_input.shape[0])
-
-                    with self._transformer_cache_context("cond_uncond"):
-                        noise_pred_video, noise_pred_audio = self.transformer(
-                            hidden_states=latent_model_input,
-                            audio_hidden_states=audio_latent_model_input,
-                            encoder_hidden_states=connector_prompt_embeds,
-                            audio_encoder_hidden_states=connector_audio_prompt_embeds,
-                            timestep=timestep,
-                            encoder_attention_mask=connector_attention_mask,
-                            audio_encoder_attention_mask=connector_attention_mask,
-                            num_frames=latent_num_frames,
-                            height=latent_height,
-                            width=latent_width,
-                            fps=frame_rate,
-                            audio_num_frames=audio_num_frames,
-                            video_coords=video_coords,
-                            audio_coords=audio_coords,
-                            attention_kwargs=attention_kwargs,
-                            return_dict=False,
-                        )
-                    noise_pred_video = noise_pred_video.float()
-                    noise_pred_audio = noise_pred_audio.float()
-
-                    if self.do_classifier_free_guidance:
-                        noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
-                        noise_pred_video = noise_pred_video_uncond + guidance_scale * (
-                            noise_pred_video_text - noise_pred_video_uncond
-                        )
-
-                        noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
-                        noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
-                            noise_pred_audio_text - noise_pred_audio_uncond
-                        )
-
-                        if guidance_rescale > 0:
-                            noise_pred_video = rescale_noise_cfg(
-                                noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
-                            )
-                            noise_pred_audio = rescale_noise_cfg(
-                                noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
-                            )
-
-                    latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
-                    audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
+                latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                    (noise_pred_video, noise_pred_audio),
+                    (t, t),
+                    (latents, audio_latents),
+                    do_true_cfg=do_true_cfg,
+                    per_request_scheduler=video_audio_scheduler,
+                )
+                latents, audio_latents = self._synchronize_cfg_parallel_step_output(
+                    (latents, audio_latents),
+                    do_true_cfg=do_true_cfg,
+                )
 
                 pbar.update()
 
@@ -1225,10 +1100,15 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
         )
 
+        audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
         audio_latents = self._denormalize_audio_latents(
             audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
         )
-        audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
+        audio_latents = self._unpack_audio_latents(
+            audio_latents,
+            original_audio_num_frames,
+            num_mel_bins=latent_mel_bins,
+        )
 
         if output_type == "latent":
             video = latents

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py

@@ -14,7 +14,7 @@ import torch
 import torch.nn as nn
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
@@ -22,7 +22,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.parallel_state import get_cfg_group, get_classifier_free_guidance_rank
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -46,6 +45,32 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return _get_ltx2_post_process_func(od_config)
 
 
+class _I2VVideoAudioScheduler:
+    """Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio."""
+
+    def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
+        self.video_scheduler = pipeline.scheduler
+        self.audio_scheduler = audio_scheduler
+        self._pipeline = pipeline
+        self._latent_num_frames = latent_num_frames
+        self._latent_height = latent_height
+        self._latent_width = latent_width
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self._pipeline._step_video_latents_i2v(
+            noise_pred[0],
+            latents[0],
+            t[0],
+            self._latent_num_frames,
+            self._latent_height,
+            self._latent_width,
+        )
+        audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[
+            0
+        ]
+        return ((video_out, audio_out),)
+
+
 class LTX2ImageToVideoPipeline(LTX2Pipeline):
     support_image_input = True
 
@@ -287,6 +312,8 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
         attention_kwargs: dict[str, Any] | None = None,
         max_sequence_length: int | None = None,
     ) -> DiffusionOutput:
+        # Extract prompt/negative_prompt from request.
+        # Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
@@ -328,6 +355,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
             else req.sampling_params.extra_args.get("audio_latents", audio_latents)
         )
 
+        # Override with pre-computed embeddings if provided in request.
         req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
         if any(p is not None for p in req_prompt_embeds):
             prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore[arg-type]
@@ -429,20 +457,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
             max_sequence_length=max_sequence_length,
             device=device,
         )
-        cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
+        # Compute positive prompt connectors
         additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
             prompt_embeds, additive_attention_mask, additive_mask=True
         )
 
+        # Compute negative prompt connectors when CFG is enabled
         negative_connector_prompt_embeds = None
         negative_connector_audio_prompt_embeds = None
         negative_connector_attention_mask = None
-        if cfg_parallel_ready:
+        if self.do_classifier_free_guidance:
             negative_additive_attention_mask = (
                 1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
             ) * -1000000.0
@@ -500,8 +525,6 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
             generator,
             latents,
         )
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
 
         duration_s = num_frames / frame_rate
         audio_latents_per_second = (
@@ -529,20 +552,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
         num_channels_latents_audio = (
             self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
         )
-
-        # padding audio_latents if needed
-        sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
-        if sp_size > 1:
-            pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
-            if pad_len > 0:
-                if audio_latents is not None:
-                    pad_shape = list(audio_latents.shape)
-                    pad_shape[2] = pad_len
-                    padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
-                    audio_latents = torch.cat([audio_latents, padding], dim=2)
-                audio_num_frames += pad_len
-
-        audio_latents = self.prepare_audio_latents(
+        audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
             batch_size * num_videos_per_prompt,
             num_channels_latents=num_channels_latents_audio,
             audio_latent_length=audio_num_frames,
@@ -585,12 +595,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
             latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
         )
         audio_coords = self.transformer.audio_rope.prepare_audio_coords(
-            audio_latents.shape[0], audio_num_frames, audio_latents.device
+            audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
         )
-        # Duplicate the positional ids as well if using CFG
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1))  # Repeat twice in batch dim
-            audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
+
+        i2v_scheduler = _I2VVideoAudioScheduler(
+            pipeline=self,
+            audio_scheduler=audio_scheduler,
+            latent_num_frames=latent_num_frames,
+            latent_height=latent_height,
+            latent_width=latent_width,
+        )
+        # No coord duplication needed: mixin handles CFG via separate forward calls.
 
         with self.progress_bar(total=len(timesteps)) as pbar:
             for i, t in enumerate(timesteps):
@@ -599,140 +614,62 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
 
                 self._current_timestep = t
 
-                if cfg_parallel_ready:
-                    latent_model_input = latents.to(prompt_embeds.dtype)
-                    audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
+                latent_model_input = latents.to(prompt_embeds.dtype)
+                audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
+                timestep = t.expand(latent_model_input.shape[0])
+                video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
+                do_true_cfg = self.do_classifier_free_guidance
 
-                    timestep = t.expand(latent_model_input.shape[0])
-                    video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
-
-                    positive_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "audio_hidden_states": audio_latent_model_input,
-                        "encoder_hidden_states": connector_prompt_embeds,
-                        "audio_encoder_hidden_states": connector_audio_prompt_embeds,
-                        "timestep": video_timestep,
-                        "audio_timestep": timestep,
-                        "encoder_attention_mask": connector_attention_mask,
-                        "audio_encoder_attention_mask": connector_attention_mask,
-                        "num_frames": latent_num_frames,
-                        "height": latent_height,
-                        "width": latent_width,
-                        "fps": frame_rate,
-                        "audio_num_frames": audio_num_frames,
-                        "video_coords": video_coords,
-                        "audio_coords": audio_coords,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
-                    }
-                    negative_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "audio_hidden_states": audio_latent_model_input,
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "audio_hidden_states": audio_latent_model_input,
+                    "encoder_hidden_states": connector_prompt_embeds,
+                    "audio_encoder_hidden_states": connector_audio_prompt_embeds,
+                    "timestep": video_timestep,
+                    "audio_timestep": timestep,
+                    "encoder_attention_mask": connector_attention_mask,
+                    "audio_encoder_attention_mask": connector_attention_mask,
+                    "num_frames": latent_num_frames,
+                    "height": latent_height,
+                    "width": latent_width,
+                    "fps": frame_rate,
+                    "audio_num_frames": padded_audio_num_frames,
+                    "video_coords": video_coords,
+                    "audio_coords": audio_coords,
+                    "attention_kwargs": attention_kwargs,
+                    "return_dict": False,
+                }
+                negative_kwargs = (
+                    {
+                        **positive_kwargs,
                         "encoder_hidden_states": negative_connector_prompt_embeds,
                         "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
-                        "timestep": video_timestep,
-                        "audio_timestep": timestep,
                         "encoder_attention_mask": negative_connector_attention_mask,
                         "audio_encoder_attention_mask": negative_connector_attention_mask,
-                        "num_frames": latent_num_frames,
-                        "height": latent_height,
-                        "width": latent_width,
-                        "fps": frame_rate,
-                        "audio_num_frames": audio_num_frames,
-                        "video_coords": video_coords,
-                        "audio_coords": audio_coords,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
                     }
+                    if do_true_cfg
+                    else None
+                )
 
-                    noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
-                        do_true_cfg=True,
-                        true_cfg_scale=guidance_scale,
-                        positive_kwargs=positive_kwargs,
-                        negative_kwargs=negative_kwargs,
-                        guidance_rescale=guidance_rescale,
-                        cfg_normalize=False,
-                    )
+                noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=False,
+                )
 
-                    if get_classifier_free_guidance_rank() == 0:
-                        latents = self._step_video_latents_i2v(
-                            noise_pred_video,
-                            latents,
-                            t,
-                            latent_num_frames,
-                            latent_height,
-                            latent_width,
-                        )
-                        audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
-
-                    cfg_group = get_cfg_group()
-                    latents = latents.contiguous()
-                    audio_latents = audio_latents.contiguous()
-                    cfg_group.broadcast(latents, src=0)
-                    cfg_group.broadcast(audio_latents, src=0)
-                else:
-                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                    latent_model_input = latent_model_input.to(prompt_embeds.dtype)
-                    audio_latent_model_input = (
-                        torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
-                    )
-                    audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
-
-                    timestep = t.expand(latent_model_input.shape[0])
-                    video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
-
-                    with self._transformer_cache_context("cond_uncond"):
-                        noise_pred_video, noise_pred_audio = self.transformer(
-                            hidden_states=latent_model_input,
-                            audio_hidden_states=audio_latent_model_input,
-                            encoder_hidden_states=connector_prompt_embeds,
-                            audio_encoder_hidden_states=connector_audio_prompt_embeds,
-                            timestep=video_timestep,
-                            audio_timestep=timestep,
-                            encoder_attention_mask=connector_attention_mask,
-                            audio_encoder_attention_mask=connector_attention_mask,
-                            num_frames=latent_num_frames,
-                            height=latent_height,
-                            width=latent_width,
-                            fps=frame_rate,
-                            audio_num_frames=audio_num_frames,
-                            video_coords=video_coords,
-                            audio_coords=audio_coords,
-                            attention_kwargs=attention_kwargs,
-                            return_dict=False,
-                        )
-                    noise_pred_video = noise_pred_video.float()
-                    noise_pred_audio = noise_pred_audio.float()
-
-                    if self.do_classifier_free_guidance:
-                        noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
-                        noise_pred_video = noise_pred_video_uncond + guidance_scale * (
-                            noise_pred_video_text - noise_pred_video_uncond
-                        )
-
-                        noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
-                        noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
-                            noise_pred_audio_text - noise_pred_audio_uncond
-                        )
-
-                        if guidance_rescale > 0:
-                            noise_pred_video = rescale_noise_cfg(
-                                noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
-                            )
-                            noise_pred_audio = rescale_noise_cfg(
-                                noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
-                            )
-
-                    latents = self._step_video_latents_i2v(
-                        noise_pred_video,
-                        latents,
-                        t,
-                        latent_num_frames,
-                        latent_height,
-                        latent_width,
-                    )
-
-                    audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
+                latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                    (noise_pred_video, noise_pred_audio),
+                    (t, t),
+                    (latents, audio_latents),
+                    do_true_cfg=do_true_cfg,
+                    per_request_scheduler=i2v_scheduler,
+                )
+                latents, audio_latents = self._synchronize_cfg_parallel_step_output(
+                    (latents, audio_latents),
+                    do_true_cfg=do_true_cfg,
+                )
 
                 pbar.update()
 
@@ -748,10 +685,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
             latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
         )
 
+        audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
         audio_latents = self._denormalize_audio_latents(
             audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
         )
-        audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
+        audio_latents = self._unpack_audio_latents(
+            audio_latents,
+            original_audio_num_frames,
+            num_mel_bins=latent_mel_bins,
+        )
 
         if output_type == "latent":
             video = latents

vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

@@ -1624,7 +1624,12 @@ def get_magi_human_post_process_func(*args, **kwargs):
     def post_process(output):
         if isinstance(output, tuple) and len(output) == 2:
             video, audio = output
-            return {"video": video, "audio": audio}
+            return {
+                "video": video,
+                "audio": audio,
+                "audio_sample_rate": 44100,
+                "fps": 25,
+            }
         return output
 
     return post_process

vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep.py

@@ -114,6 +114,8 @@ class NextStepConfig(LlamaConfig):
 
 
 class NextStepModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["layers"]
+
     def __init__(self, config: NextStepConfig):
         super().__init__()
         self.config = config

vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py

@@ -366,6 +366,7 @@ class OvisImageTransformer2DModel(nn.Module):
     """
 
     _repeated_blocks = ["OvisImageTransformerBlock", "OvisImageSingleTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     def __init__(
         self,

vllm_omni/diffusion/models/sd3/sd3_transformer.py

@@ -387,6 +387,7 @@ class SD3Transformer2DModel(nn.Module):
     """
 
     _repeated_blocks = ["SD3TransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
 
     def __init__(
         self,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -272,7 +272,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -217,7 +217,7 @@ class Wan22I2VPipeline(
 
         # VAE
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Transformers (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -24,13 +24,13 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLWan
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -174,8 +174,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
         ).to(self.device)
 
         # VAE
-        self.vae = AutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = OmniAutoencoderKLWan.from_pretrained(
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Single transformer (TI2V uses dense 5B model, not MoE)

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -21,9 +21,10 @@ import os
 from collections.abc import Callable, Iterable
 from typing import Any
 
+import PIL.Image
 import torch
 import torch.nn as nn
-from diffusers.image_processor import VaeImageProcessor
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
@@ -59,7 +60,7 @@ def get_post_process_func(
         vae_config = json.load(f)
         vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
 
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
 
     def post_process_func(
         images: torch.Tensor,
@@ -83,6 +84,20 @@ def calculate_shift(
     return mu
 
 
+# Copied from diffusers
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
 def retrieve_timesteps(
     scheduler,
@@ -187,6 +202,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
+
     def encode_prompt(
         self,
         prompt: str | list[str],
@@ -282,12 +299,45 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         device,
         generator,
         latents=None,
+        image=None,
+        timestep=None,
     ):
         height = 2 * (int(height) // (self.vae_scale_factor * 2))
         width = 2 * (int(width) // (self.vae_scale_factor * 2))
 
         shape = (batch_size, num_channels_latents, height, width)
 
+        if image is not None:
+            if latents is not None:
+                return latents.to(device=device, dtype=dtype)
+
+            image = image.to(device=device, dtype=dtype)
+            if image.shape[1] != num_channels_latents:
+                if isinstance(generator, list):
+                    image_latents = [
+                        retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
+                        for i in range(image.shape[0])
+                    ]
+                    image_latents = torch.cat(image_latents, dim=0)
+                else:
+                    image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+                image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+            else:
+                image_latents = image
+
+            if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+                additional_image_per_prompt = batch_size // image_latents.shape[0]
+                image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+            elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+                raise ValueError(
+                    f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+                )
+
+            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self.scheduler.scale_noise(image_latents, timestep, noise)
+            return latents
+
         if latents is None:
             latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         else:
@@ -296,6 +346,14 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
             latents = latents.to(device)
         return latents
 
+    def get_timesteps(self, num_inference_steps, strength, device):
+        init_timestep = min(num_inference_steps * strength, num_inference_steps)
+        t_start = int(max(num_inference_steps - init_timestep, 0))
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(t_start * self.scheduler.order)
+        return timesteps, num_inference_steps - t_start
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -320,6 +378,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self,
         req: OmniDiffusionRequest,
         prompt: str | list[str] | None = None,
+        image: PipelineImageInput = None,
+        strength: float = 0.6,
         height: int = 1024,
         width: int = 1024,
         num_inference_steps: int = 50,
@@ -347,6 +407,11 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
             prompt (`str` or `list[str]`, *optional*):
                 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                 instead.
+            image (`PipelineImageInput`, *optional*):
+                The image to use for img2img generation. If provided, the pipeline
+                will perform img2img instead of text-to-image.
+            strength (`float`, *optional*, defaults to 0.6):
+                Indicates extent to transform the reference `image`. Must be between 0 and 1.
             height (`int`, *optional*, defaults to 1024):
                 The height in pixels of the generated image.
             width (`int`, *optional*, defaults to 1024):
@@ -425,6 +490,34 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
+        # Handle img2img: extract image from request
+        if image is None and req.prompts:
+            if len(req.prompts) > 1:
+                logger.warning(
+                    "This model only supports a single prompt for img2img, not a batched request. "
+                    "Taking only the first image for now."
+                )
+            first_prompt = req.prompts[0]
+            if not isinstance(first_prompt, str):
+                raw_image = first_prompt.get("multi_modal_data", {}).get("image")
+                if raw_image is not None:
+                    if isinstance(raw_image, list):
+                        image = [PIL.Image.open(im) if isinstance(im, str) else raw_image[0] for im in raw_image[:1]]
+                    else:
+                        image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image
+
+        # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter
+        strength = req.sampling_params.strength if req.sampling_params.strength is not None else strength
+        if strength is not None and image is None:
+            logger.warning(
+                "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. "
+                "It will be ignored for text-to-image (T2I) generation.",
+                strength,
+            )
+            strength = None
+        if image is not None and strength is not None and (strength < 0 or strength > 1):
+            raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}")
+
         height = req.sampling_params.height or height
         width = req.sampling_params.width or width
         num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
@@ -491,16 +584,71 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.in_channels
 
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
+        # img2img mode: prepare latents from input image
+        if image is not None:
+            # Handle image list - take first image
+            if isinstance(image, list):
+                image = image[0]
+
+            # Prepare image for VAE encoding using image_processor
+            if not isinstance(image, torch.Tensor):
+                init_image = self.image_processor.preprocess(image, height, width)
+                image = init_image.to(dtype=torch.float32, device=device)
+
+            # Initialize scheduler kwargs for img2img
+            mu = calculate_shift(
+                (height // self.vae_scale_factor // 2) * (width // self.vae_scale_factor // 2),
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            )
+            self.scheduler.sigma_min = 0.0
+            scheduler_kwargs = {"mu": mu}
+
+            # First initialize timesteps in scheduler
+            timesteps, num_inference_steps = retrieve_timesteps(
+                self.scheduler,
+                num_inference_steps,
+                device,
+                sigmas=sigmas,
+                **scheduler_kwargs,
+            )
+
+            # Then adjust timesteps based on strength
+            timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+
+            if num_inference_steps < 1:
+                raise ValueError(
+                    f"After adjusting the num_inference_steps by strength parameter: "
+                    f"{strength}, the number of pipeline steps is {num_inference_steps} "
+                    f"which is < 1 and not appropriate for this pipeline."
+                )
+            latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+            latents = self.prepare_latents(
+                batch_size * num_images_per_prompt,
+                num_channels_latents,
+                height,
+                width,
+                prompt_embeds[0].dtype,
+                device,
+                generator,
+                latents,
+                image,
+                latent_timestep,
+            )
+        else:
+            latents = self.prepare_latents(
+                batch_size * num_images_per_prompt,
+                num_channels_latents,
+                height,
+                width,
+                torch.float32,
+                device,
+                generator,
+                latents,
+            )
 
         # Repeat prompt_embeds for num_images_per_prompt
         if num_images_per_prompt > 1:
@@ -509,25 +657,28 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
 
         actual_batch_size = batch_size * num_images_per_prompt
-        image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
 
         # 5. Prepare timesteps
-        mu = calculate_shift(
-            image_seq_len,
-            self.scheduler.config.get("base_image_seq_len", 256),
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_shift", 0.5),
-            self.scheduler.config.get("max_shift", 1.15),
-        )
-        self.scheduler.sigma_min = 0.0
-        scheduler_kwargs = {"mu": mu}
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            sigmas=sigmas,
-            **scheduler_kwargs,
-        )
+        if image is None:
+            image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            )
+            self.scheduler.sigma_min = 0.0
+            scheduler_kwargs = {"mu": mu}
+
+            timesteps, num_inference_steps = retrieve_timesteps(
+                self.scheduler,
+                num_inference_steps,
+                device,
+                sigmas=sigmas,
+                **scheduler_kwargs,
+            )
+
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
         self._num_timesteps = len(timesteps)
 

vllm_omni/diffusion/offloader/module_collector.py

@@ -21,7 +21,7 @@ class PipelineModules:
 class ModuleDiscovery:
     """Discovers pipeline components for offloading"""
 
-    DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks"]
+    DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"]
     ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]
     VAE_ATTRS = ["vae", "audio_vae"]
 

vllm_omni/diffusion/request.py

@@ -26,6 +26,8 @@ class OmniDiffusionRequest:
     sampling_params: OmniDiffusionSamplingParams
 
     request_ids: list[str] = field(default_factory=list)
+    request_id: str | None = None
+    kv_sender_info: dict | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/stage_diffusion_client.py

@@ -179,6 +179,7 @@ class StageDiffusionClient:
         request_id: str,
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         self._request_socket.send(
             self._encoder.encode(
@@ -187,6 +188,7 @@ class StageDiffusionClient:
                     "request_id": request_id,
                     "prompt": prompt,
                     "sampling_params": self._sampling_params_to_dict(sampling_params),
+                    "kv_sender_info": kv_sender_info,
                 }
             )
         )
@@ -198,6 +200,7 @@ class StageDiffusionClient:
         request_id: str,
         prompts: list[OmniPromptType],
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         """Submit a list of prompts as a single batched engine call.
 
@@ -206,7 +209,12 @@ class StageDiffusionClient:
         *request_id*.
         """
         task = asyncio.create_task(
-            self._run_batch(request_id, prompts, sampling_params),
+            self._run_batch(
+                request_id,
+                prompts,
+                sampling_params,
+                kv_sender_info,
+            ),
             name=f"diffusion-batch-{request_id}",
         )
         self._tasks[request_id] = task
@@ -216,6 +224,7 @@ class StageDiffusionClient:
         request_id: str,
         prompts: list[OmniPromptType],
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         try:
             self._request_socket.send(
@@ -225,6 +234,7 @@ class StageDiffusionClient:
                         "request_id": request_id,
                         "prompts": prompts,
                         "sampling_params": self._sampling_params_to_dict(sampling_params),
+                        "kv_sender_info": kv_sender_info,
                     }
                 )
             )