docs: address recipe review feedback

Signed-off-by: hsliu <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

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

recipes/Qwen/Qwen3-Omni.md

@@ -0,0 +1,90 @@
+# Qwen3-Omni for multimodal chat on 1x A100 80GB
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen3-Omni-30B-A3B-Instruct`
+- Task: Multimodal chat with text, image, audio, or video input
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`Qwen/Qwen3-Omni-30B-A3B-Instruct` with vLLM-Omni on a single 80 GB A100 and
+validate the deployment with the existing multimodal client examples in this
+repository.
+
+## References
+
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/qwen3_omni.md`](../../docs/user_guide/examples/online_serving/qwen3_omni.md)
+- Related example under `examples/`:
+  [`examples/online_serving/qwen3_omni/README.md`](../../examples/online_serving/qwen3_omni/README.md)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+
+## Hardware Support
+
+This recipe currently documents one tested-style reference configuration for
+CUDA GPU serving. Add more sections for other hardware as community validation
+lands.
+
+## GPU
+
+### 1x A100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+To enable async chunking, use the bundled stage config:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --omni \
+  --port 8091 \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
+#### Verification
+
+Run one of the existing example clients after the server is ready:
+
+```bash
+python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --query-type use_image \
+  --port 8091 \
+  --host localhost
+```
+
+For a quick API smoke test, request text-only output:
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+    "messages": [{"role": "user", "content": "Describe vLLM in brief."}],
+    "modalities": ["text"]
+  }'
+```
+
+#### Notes
+
+- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads.
+- Key flags: `--omni` is required; `--stage-configs-path` is optional for custom or async-chunk stage configs.
+- Known limitations: This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.

recipes/README.md

@@ -0,0 +1,35 @@
+# Community Recipes
+
+This directory contains community-maintained recipes for answering a
+practical user question:
+
+> How do I run model X on hardware Y for task Z?
+
+Add recipes for this repository under this in-repo `recipes/` directory. To
+keep naming and layout consistent, organize recipes by model vendor in a way
+that is aligned with
+[`vllm-project/recipes`](https://github.com/vllm-project/recipes), but treat
+that external repository as a reference for structure rather than the place to
+add files for this repo. Use one Markdown file per model family by default.
+
+Example layout:
+
+```text
+recipes/
+  Qwen/
+    Qwen3-Omni.md
+    Qwen3-TTS.md
+  Tencent-Hunyuan/
+    HunyuanVideo.md
+```
+
+## Available Recipes
+
+- [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
+  multimodal chat on `1x A100 80GB`
+
+Within a single recipe file, include different hardware support sections such
+as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
+`1x A100 80GB` or `2x L40S` inside those sections when applicable.
+
+See [TEMPLATE.md](./TEMPLATE.md) for the recommended format.

recipes/TEMPLATE.md

@@ -0,0 +1,82 @@
+# Recipe Title
+
+> Example: Qwen3-Omni for speech chat on 1x A100 80GB
+
+## Summary
+
+- Vendor:
+- Model:
+- Task:
+- Mode:
+- Maintainer:
+
+## When to use this recipe
+
+Briefly describe the concrete scenario this recipe covers.
+
+## References
+
+- Upstream or canonical docs:
+- Related example under `examples/`:
+- Related issue or discussion:
+
+## Hardware Support
+
+Add one section per platform, such as `GPU`, `ROCm`, or `NPU`. Under each
+platform section, document one or more tested hardware configurations.
+
+## GPU
+
+### 1x A100 80GB
+
+#### Environment
+
+- OS:
+- Python:
+- Driver / runtime:
+- vLLM version:
+- vLLM-Omni version or commit:
+
+#### Command
+
+```bash
+# Add the exact command(s) here
+```
+
+#### Verification
+
+```bash
+# Add a quick validation command or expected output here
+```
+
+#### Notes
+
+- Memory usage:
+- Key flags:
+- Known limitations:
+
+### 2x L40S
+
+Repeat the same structure for other hardware setups as needed.
+
+## ROCm
+
+### Example hardware configuration
+
+Repeat the same nested structure for ROCm setups as needed:
+
+- `#### Environment`
+- `#### Command`
+- `#### Verification`
+- `#### Notes`
+
+## NPU
+
+### Example hardware configuration
+
+Repeat the same nested structure for NPU setups as needed:
+
+- `#### Environment`
+- `#### Command`
+- `#### Verification`
+- `#### Notes`

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,
                     }
                 )
             )

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -14,8 +14,10 @@ from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Any
 
 import msgspec
+import torch
 import zmq
 import zmq.asyncio
+from PIL import Image
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
@@ -128,6 +130,7 @@ class StageDiffusionProc:
         request_id: str,
         prompt: Any,
         sampling_params_dict: dict,
+        kv_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a diffusion request and run DiffusionEngine.step()."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -136,6 +139,8 @@ class StageDiffusionProc:
             prompts=[prompt],
             sampling_params=sampling_params,
             request_ids=[request_id],
+            request_id=request_id,
+            kv_sender_info=kv_sender_info,
         )
 
         loop = asyncio.get_running_loop()
@@ -150,6 +155,7 @@ class StageDiffusionProc:
         request_id: str,
         prompts: list[Any],
         sampling_params_dict: dict,
+        kv_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a batched diffusion request and run DiffusionEngine.step().
 
@@ -163,7 +169,9 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=prompts,
             sampling_params=sampling_params,
-            request_ids=[request_id] * len(prompts),
+            request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
+            request_id=request_id,
+            kv_sender_info=kv_sender_info,
         )
 
         loop = asyncio.get_running_loop()
@@ -174,8 +182,13 @@ class StageDiffusionProc:
         merged_mm: dict[str, Any] = {}
         merged_metrics: dict[str, Any] = {}
         merged_durations: dict[str, float] = {}
+        merged_custom: dict[str, Any] = {}
         peak_mem = 0.0
         latents = None
+        trajectory_latents: list[torch.Tensor] | None = None
+        trajectory_timesteps: list[torch.Tensor] | None = None
+        trajectory_log_probs: torch.Tensor | None = None
+        trajectory_decoded: list[Image.Image] | None = None
         final_output_type = "image"
 
         for r in results:
@@ -183,9 +196,18 @@ class StageDiffusionProc:
             merged_mm.update(r._multimodal_output)
             merged_metrics.update(r.metrics)
             merged_durations.update(r.stage_durations)
+            merged_custom.update(r._custom_output)
             peak_mem = max(peak_mem, r.peak_memory_mb)
             if latents is None and r.latents is not None:
                 latents = r.latents
+            if trajectory_latents is None:
+                trajectory_latents = r.trajectory_latents
+            if trajectory_timesteps is None:
+                trajectory_timesteps = r.trajectory_timesteps
+            if trajectory_log_probs is None:
+                trajectory_log_probs = r.trajectory_log_probs
+            if trajectory_decoded is None:
+                trajectory_decoded = r.trajectory_decoded
             if r.final_output_type != "image":
                 final_output_type = r.final_output_type
 
@@ -195,6 +217,11 @@ class StageDiffusionProc:
             prompt=prompts[0] if len(prompts) == 1 else None,
             metrics=merged_metrics,
             latents=latents,
+            trajectory_latents=trajectory_latents,
+            trajectory_timesteps=trajectory_timesteps,
+            trajectory_log_probs=trajectory_log_probs,
+            trajectory_decoded=trajectory_decoded,
+            custom_output=merged_custom or None,
             multimodal_output=merged_mm or None,
             final_output_type=final_output_type,
             stage_durations=merged_durations,
@@ -325,10 +352,20 @@ class StageDiffusionProc:
 
         tasks: dict[str, asyncio.Task] = {}
 
-        async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None:
+        async def _dispatch_request(
+            request_id: str,
+            prompt: Any,
+            sampling_params_dict: dict,
+            kv_sender_info: dict[str, Any] | None = None,
+        ) -> None:
             """Process a single diffusion request and send the response."""
             try:
-                result = await self._process_request(request_id, prompt, sampling_params_dict)
+                result = await self._process_request(
+                    request_id,
+                    prompt,
+                    sampling_params_dict,
+                    kv_sender_info=kv_sender_info,
+                )
                 await response_socket.send(encoder.encode({"type": "result", "output": result}))
             except DiffusionRequestAbortedError as e:
                 logger.info(
@@ -363,6 +400,7 @@ class StageDiffusionProc:
                             request_id,
                             msg["prompt"],
                             msg["sampling_params"],
+                            msg.get("kv_sender_info"),
                         )
                     )
                     tasks[request_id] = task
@@ -370,9 +408,19 @@ class StageDiffusionProc:
                 elif msg_type == "add_batch_request":
                     request_id = msg["request_id"]
 
-                    async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None:
+                    async def _dispatch_batch(
+                        rid: str,
+                        prompts: list,
+                        sp_dict: dict,
+                        kv_sender_info: dict[str, Any] | None = None,
+                    ) -> None:
                         try:
-                            result = await self._process_batch_request(rid, prompts, sp_dict)
+                            result = await self._process_batch_request(
+                                rid,
+                                prompts,
+                                sp_dict,
+                                kv_sender_info=kv_sender_info,
+                            )
                             await response_socket.send(encoder.encode({"type": "result", "output": result}))
                         except DiffusionRequestAbortedError as e:
                             logger.info(
@@ -399,6 +447,7 @@ class StageDiffusionProc:
                             request_id,
                             msg["prompts"],
                             msg["sampling_params"],
+                            msg.get("kv_sender_info"),
                         )
                     )
                     tasks[request_id] = task

vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py

@@ -277,13 +277,15 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
         }
 
         self.config = config
-        host_config = config.get("host", "127.0.0.1")
-        # Support "auto" to auto-detect local IP address
-        if host_config.lower() == "auto":
+        host_config = config.get("host")
+        host_value = "auto" if host_config is None else str(host_config)
+        # Default sender/receiver bootstrap to a routable local IP so the
+        # advertised endpoint matches the interface Mooncake binds.
+        if host_value.lower() == "auto" or host_value in {"", "*", "0.0.0.0", "::"}:
             self.host = self._get_local_ip()
             logger.info(f"Auto-detected local IP for RDMA: {self.host}")
         else:
-            self.host = host_config
+            self.host = host_value
         self.zmq_port = config.get("zmq_port", 50051)
         self.protocol = config.get("protocol", "rdma")
 

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unified OmniConnector and KV cache transfer management."""
 
+import json
+import struct
 import time
 from collections.abc import Callable
 from dataclasses import asdict, dataclass
@@ -12,12 +14,36 @@ from vllm.logger import init_logger
 
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec
+from .utils.initialization import KV_TRANSFER_PORT_OFFSET
 from .utils.kv_utils import normalize_layer_kv
 
 logger = init_logger(__name__)
 
 LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
 
+_SAFE_TORCH_DTYPES = {
+    name: dtype
+    for name in (
+        "bool",
+        "uint8",
+        "int8",
+        "int16",
+        "int32",
+        "int64",
+        "float16",
+        "float32",
+        "float64",
+        "bfloat16",
+        "complex64",
+        "complex128",
+        "float8_e4m3fn",
+        "float8_e4m3fnuz",
+        "float8_e5m2",
+        "float8_e5m2fnuz",
+    )
+    if isinstance((dtype := getattr(torch, name, None)), torch.dtype)
+}
+
 
 @dataclass
 class OmniKVCacheConfig:
@@ -46,6 +72,242 @@ class KVCacheTransferData:
         """Convert to dictionary for serialization."""
         return asdict(self)
 
+    def to_bytes(self) -> bytes:
+        """Convert to compact binary format for fast transfer."""
+        tensors_desc: list[dict[str, Any]] = []
+        tensor_bufs: list[bytes] = []
+        data_offset = 0
+
+        for cache_name in ("key_cache", "value_cache"):
+            cache_list = self.layer_blocks.get(cache_name, [])
+            for layer_idx, tensor in enumerate(cache_list):
+                if tensor is None:
+                    tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True})
+                    continue
+
+                t = tensor.detach().cpu().contiguous()
+                dtype_str = str(t.dtype).removeprefix("torch.")
+                raw = t.view(torch.uint8).numpy().tobytes()
+                tensors_desc.append(
+                    {
+                        "n": f"{cache_name}_{layer_idx}",
+                        "i": layer_idx,
+                        "d": dtype_str,
+                        "s": list(t.shape),
+                        "o": data_offset,
+                        "b": len(raw),
+                    }
+                )
+                tensor_bufs.append(raw)
+                data_offset += len(raw)
+
+        header = json.dumps(
+            {
+                "rid": self.request_id,
+                "bids": self.block_ids,
+                "meta": self.metadata,
+                "td": tensors_desc,
+                "nl": len(self.layer_blocks.get("key_cache", [])),
+            },
+            separators=(",", ":"),
+        ).encode("utf-8")
+        return b"".join([struct.pack(">I", len(header)), header] + tensor_bufs)
+
+    def to_gpu_tensor(self) -> torch.Tensor:
+        """Convert to a packed GPU tensor for raw-data connectors."""
+        tensors_desc: list[dict[str, Any]] = []
+        gpu_tensors: list[torch.Tensor] = []
+        data_offset = 0
+        device = None
+
+        for cache_name in ("key_cache", "value_cache"):
+            cache_list = self.layer_blocks.get(cache_name, [])
+            for layer_idx, tensor in enumerate(cache_list):
+                if tensor is None:
+                    tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True})
+                    continue
+
+                t = tensor.detach().contiguous()
+                if device is None and t.is_cuda:
+                    device = t.device
+                dtype_str = str(t.dtype).removeprefix("torch.")
+                nbytes = t.numel() * t.element_size()
+                tensors_desc.append(
+                    {
+                        "n": f"{cache_name}_{layer_idx}",
+                        "i": layer_idx,
+                        "d": dtype_str,
+                        "s": list(t.shape),
+                        "o": data_offset,
+                        "b": nbytes,
+                    }
+                )
+                gpu_tensors.append(t.view(torch.uint8).flatten())
+                data_offset += nbytes
+
+        if device is None:
+            raise RuntimeError("No CUDA tensors found, use to_bytes() instead")
+
+        header = json.dumps(
+            {
+                "rid": self.request_id,
+                "bids": self.block_ids,
+                "meta": self.metadata,
+                "td": tensors_desc,
+                "nl": len(self.layer_blocks.get("key_cache", [])),
+            },
+            separators=(",", ":"),
+        ).encode("utf-8")
+
+        header_prefix = struct.pack(">I", len(header)) + header
+        total_size = len(header_prefix) + data_offset
+        output = torch.empty(total_size, dtype=torch.uint8, device=device)
+        header_tensor = torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)
+        output[: len(header_prefix)].copy_(header_tensor)
+
+        pos = len(header_prefix)
+        for t_flat in gpu_tensors:
+            n = t_flat.numel()
+            output[pos : pos + n].copy_(t_flat)
+            pos += n
+
+        return output
+
+    @staticmethod
+    def _load_header_from_memoryview(raw_mv: memoryview) -> tuple[dict[str, Any], memoryview]:
+        if len(raw_mv) < 4:
+            raise ValueError("Corrupted KV payload: missing 4-byte header length")
+
+        header_len = struct.unpack(">I", raw_mv[:4])[0]
+        if header_len > len(raw_mv) - 4:
+            raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={len(raw_mv)}")
+
+        return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :]
+
+    @staticmethod
+    def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]:
+        if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1:
+            raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor")
+
+        total_bytes = int(gpu_tensor.numel())
+        if total_bytes < 4:
+            raise ValueError("Corrupted KV payload: missing 4-byte header length")
+
+        header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0]
+        if header_len > total_bytes - 4:
+            raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}")
+
+        header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes()
+        return json.loads(header_bytes), 4 + header_len
+
+    @staticmethod
+    def _validate_tensor_span(name: str, info: dict[str, Any], tensor_data_bytes: int) -> tuple[int, int]:
+        offset = info["o"]
+        nbytes = info["b"]
+        if offset < 0 or nbytes < 0 or offset + nbytes > tensor_data_bytes:
+            raise ValueError(
+                f"Corrupted KV payload tensor span for {name}: "
+                f"offset={offset}, bytes={nbytes}, tensor_data_bytes={tensor_data_bytes}"
+            )
+        return offset, nbytes
+
+    @staticmethod
+    def _resolve_torch_dtype(dtype_name: Any) -> torch.dtype:
+        torch_dtype = _SAFE_TORCH_DTYPES.get(str(dtype_name))
+        if torch_dtype is None:
+            raise ValueError(f"Unsupported dtype in KV payload: {dtype_name}")
+        return torch_dtype
+
+    @staticmethod
+    def _resolve_layer_idx(info: dict[str, Any], num_layers: int) -> int:
+        layer_idx = info.get("i")
+        if layer_idx is None:
+            name = info.get("n")
+            if isinstance(name, str) and name.startswith("key_cache_"):
+                layer_idx = int(name.removeprefix("key_cache_"))
+            elif isinstance(name, str) and name.startswith("value_cache_"):
+                layer_idx = int(name.removeprefix("value_cache_"))
+            else:
+                raise ValueError(f"Invalid KV tensor name in payload: {name}")
+
+        if not isinstance(layer_idx, int):
+            raise ValueError(f"Invalid layer index in KV payload: {layer_idx}")
+        if layer_idx < 0 or layer_idx >= num_layers:
+            raise ValueError(f"Invalid layer index in KV payload: {layer_idx} (num_layers={num_layers})")
+        return layer_idx
+
+    @staticmethod
+    def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]:
+        """Reconstruct KV cache data from the packed bytes format."""
+        raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw
+        header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv)
+
+        num_layers = header["nl"]
+        key_cache: list[torch.Tensor | None] = [None] * num_layers
+        value_cache: list[torch.Tensor | None] = [None] * num_layers
+
+        for info in header["td"]:
+            if info.get("x"):
+                continue
+
+            name: str = info["n"]
+            torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"])
+            offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, len(tensor_data_mv))
+            t = (
+                torch.frombuffer(
+                    tensor_data_mv,
+                    dtype=torch.uint8,
+                    offset=offset,
+                    count=nbytes,
+                )
+                .view(torch_dtype)
+                .reshape(info["s"])
+            )
+            layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers)
+            if name.startswith("key_cache_"):
+                key_cache[layer_idx] = t
+            elif name.startswith("value_cache_"):
+                value_cache[layer_idx] = t
+
+        return {
+            "request_id": header["rid"],
+            "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
+            "block_ids": header["bids"],
+            "metadata": header["meta"],
+        }
+
+    @staticmethod
+    def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]:
+        """Reconstruct KV cache data from a packed GPU tensor."""
+        header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor)
+
+        num_layers = header["nl"]
+        key_cache: list[torch.Tensor | None] = [None] * num_layers
+        value_cache: list[torch.Tensor | None] = [None] * num_layers
+        tensor_data_bytes = int(gpu_tensor.numel()) - data_start
+
+        for info in header["td"]:
+            if info.get("x"):
+                continue
+
+            name: str = info["n"]
+            torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"])
+            offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, tensor_data_bytes)
+            t = gpu_tensor[data_start + offset : data_start + offset + nbytes].clone()
+            t = t.view(torch_dtype).reshape(info["s"])
+            layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers)
+            if name.startswith("key_cache_"):
+                key_cache[layer_idx] = t
+            elif name.startswith("value_cache_"):
+                value_cache[layer_idx] = t
+
+        return {
+            "request_id": header["rid"],
+            "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
+            "block_ids": header["bids"],
+            "metadata": header["meta"],
+        }
+
 
 class OmniKVTransferManager:
     """Unified management for OmniConnector and KV cache transfer.
@@ -79,6 +341,13 @@ class OmniKVTransferManager:
             else (None, None)
         )
 
+        if config.need_send_cache and config.connector_config:
+            try:
+                _ = self.connector
+                logger.info("Sender connector eagerly initialized")
+            except Exception as e:
+                logger.warning("Failed to eagerly initialize sender connector: %s", e)
+
     @classmethod
     def _create(cls, cfg: dict | None) -> "OmniKVTransferManager":
         """Create manager from raw config dict."""
@@ -140,8 +409,39 @@ class OmniKVTransferManager:
             cfg = self.config.connector_config
             if cfg and (c_type := cfg.get("type")):
                 try:
-                    logger.info(f"Initializing OmniConnector with config: {cfg}")
                     c_extra = {k: v for k, v in cfg.items() if k != "type"}
+                    if c_type == "MooncakeTransferEngineConnector":
+                        base_port = c_extra.get("zmq_port", 50051)
+                        c_extra["from_stage"] = (
+                            str(self.config.from_stage) if self.config.from_stage is not None else "0"
+                        )
+                        c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1"
+
+                        if self.config.need_send_cache:
+                            c_extra["role"] = "sender"
+                            from_stage = self.config.from_stage
+                            if from_stage is not None:
+                                try:
+                                    c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage)
+                                except (TypeError, ValueError):
+                                    c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET
+                        elif self.config.need_recv_cache:
+                            c_extra["role"] = "receiver"
+                            from_stage = self.config.from_stage
+                            sender_port = base_port + KV_TRANSFER_PORT_OFFSET
+                            if from_stage is not None:
+                                try:
+                                    sender_port = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage)
+                                except (TypeError, ValueError):
+                                    pass
+                            c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1"))
+                            c_extra.setdefault("sender_zmq_port", sender_port)
+
+                    logger.info(
+                        "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s",
+                        cfg,
+                        c_extra.get("role", "N/A"),
+                    )
                     self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
                 except Exception as e:
                     logger.error(f"Failed to initialize OmniConnector: {e}")
@@ -157,6 +457,85 @@ class OmniKVTransferManager:
         """Get connector (compatibility wrapper for existing code)."""
         return self.connector
 
+    def _resolve_sender_info(
+        self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None
+    ) -> dict[str, Any] | None:
+        if not sender_info:
+            return None
+
+        if "host" in sender_info:
+            return sender_info
+
+        if not isinstance(sender_info, dict):
+            return None
+
+        preferred_keys: list[str | int] = []
+        if sender_stage_id is None:
+            recv_from, _ = self.recv_stages
+            sender_stage_id = recv_from
+
+        if sender_stage_id is not None:
+            preferred_keys.append(sender_stage_id)
+            preferred_keys.append(str(sender_stage_id))
+            try:
+                preferred_keys.append(int(sender_stage_id))
+            except (TypeError, ValueError):
+                pass
+
+        for key in dict.fromkeys(preferred_keys):
+            info = sender_info.get(key)
+            if isinstance(info, dict) and "host" in info:
+                return info
+
+        candidates = [info for info in sender_info.values() if isinstance(info, dict) and "host" in info]
+        if len(candidates) == 1:
+            return candidates[0]
+
+        if candidates:
+            logger.warning(
+                "Ambiguous sender_info for sender_stage_id=%s: "
+                "expected caller to resolve a single sender entry, got %s",
+                sender_stage_id,
+                sender_info,
+            )
+        return None
+
+    @staticmethod
+    def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(data, dict) or "layer_blocks" not in data:
+            return data
+
+        layer_blocks = data["layer_blocks"]
+        for cache_name in ("key_cache", "value_cache"):
+            cache_list = layer_blocks.get(cache_name, [])
+            for idx, tensor in enumerate(cache_list):
+                if isinstance(tensor, torch.Tensor):
+                    cache_list[idx] = tensor.clone()
+        return data
+
+    def update_sender_info(self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None) -> None:
+        """Update receiver-side sender info before loading remote KV cache."""
+        if not self.config.need_recv_cache:
+            return
+
+        actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id)
+        if not actual_info or "host" not in actual_info:
+            logger.warning("Invalid sender_info format: %s", sender_info)
+            return
+
+        if self.config.connector_config:
+            self.config.connector_config["sender_host"] = actual_info.get("host")
+            self.config.connector_config["sender_zmq_port"] = actual_info.get("zmq_port")
+
+        if self._connector and hasattr(self._connector, "update_sender_info"):
+            try:
+                self._connector.update_sender_info(actual_info.get("host"), actual_info.get("zmq_port"))
+            except Exception:
+                if hasattr(self._connector, "sender_host"):
+                    self._connector.sender_host = actual_info.get("host")
+                if hasattr(self._connector, "sender_zmq_port"):
+                    self._connector.sender_zmq_port = actual_info.get("zmq_port")
+
     def handle_finished_requests_kv_transfer(
         self,
         finished_reqs: dict[str, dict[str, Any]],
@@ -203,7 +582,8 @@ class OmniKVTransferManager:
 
                 custom_metadata = data.get("custom_metadata")
 
-                # Extract KV cache from GPU blocks -> CPU tensors
+                # Extract KV cache from GPU blocks and keep it on-device when
+                # possible so raw-data connectors can use the fast path.
                 kv_data = self._extract_kv_cache(
                     req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata
                 )
@@ -280,9 +660,8 @@ class OmniKVTransferManager:
                 flat_k = flat_k[:seq_len]
                 flat_v = flat_v[:seq_len]
 
-            # Move to CPU
-            key_cache[layer_idx] = flat_k.detach().cpu().contiguous()
-            value_cache[layer_idx] = flat_v.detach().cpu().contiguous()
+            key_cache[layer_idx] = flat_k.detach().contiguous()
+            value_cache[layer_idx] = flat_v.detach().contiguous()
 
         if not any(k is not None for k in key_cache):
             return None
@@ -311,14 +690,40 @@ class OmniKVTransferManager:
         if not from_stage or not to_stage:
             raise ValueError("Transfer stages (omni_from_stage, omni_to_stage) not configured")
 
-        # Prepare data and transfer with retry
-        data_dict = kv_data.to_dict()
-        data_dict["request_id"] = transfer_req_id
+        kv_data.request_id = transfer_req_id
+        serialization_start = time.perf_counter()
+        transfer_data: torch.Tensor | bytes | dict[str, Any]
+        supports_raw = getattr(self.connector, "supports_raw_data", False)
 
-        success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict)
+        try:
+            if supports_raw:
+                transfer_data = kv_data.to_gpu_tensor()
+            else:
+                raise RuntimeError("Connector does not support raw tensor")
+        except Exception:
+            try:
+                transfer_data = kv_data.to_bytes()
+            except Exception:
+                data_dict = kv_data.to_dict()
+                data_dict["request_id"] = transfer_req_id
+                transfer_data = data_dict
+
+        serialization_ms = (time.perf_counter() - serialization_start) * 1000
+        logger.info("KV cache serialized for %s in %.1f ms", transfer_req_id, serialization_ms)
+
+        transfer_start = time.perf_counter()
+        success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", transfer_data)
+        elapsed = time.perf_counter() - transfer_start
 
         if success:
-            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
+            mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0
+            logger.info(
+                "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s",
+                transfer_req_id,
+                size,
+                elapsed,
+                mbps,
+            )
         else:
             logger.error(f"KV transfer FAILED: {transfer_req_id}")
 
@@ -327,7 +732,7 @@ class OmniKVTransferManager:
         from_stage: str,
         to_stage: str,
         request_id: str,
-        data: dict[str, Any],
+        data: "dict[str, Any] | bytes | torch.Tensor",
         max_retries: int = 3,
     ) -> tuple[bool, int, dict[str, Any] | None]:
         """Transfer data with retry and exponential backoff.
@@ -393,6 +798,8 @@ class OmniKVTransferManager:
 
         timeout = self.config.recv_timeout
         start_time = time.time()
+        poll_interval = 0.01
+        max_poll_interval = 0.5
 
         logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...")
 
@@ -400,33 +807,74 @@ class OmniKVTransferManager:
             while True:
                 # Build the full key for connector
                 full_request_id = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"
+                link_start = time.perf_counter()
                 result = self.connector.get(
                     from_stage=from_stage,
                     to_stage=to_stage,
                     get_key=full_request_id,
                 )
                 if result:
-                    data, size = result
-                    logger.info(f"Successfully received KV cache for {request_id}, {size} bytes")
+                    raw_data, size = result
+                    elapsed = time.time() - start_time
+                    link_ms = (time.perf_counter() - link_start) * 1000
+                    managed_buffer = None
 
-                    # Move tensors to target device if specified
-                    if target_device is not None and isinstance(data, dict) and "layer_blocks" in data:
-                        layer_blocks = data["layer_blocks"]
-                        for cache_list in [
-                            layer_blocks.get("key_cache", []),
-                            layer_blocks.get("value_cache", []),
-                        ]:
-                            for i, tensor in enumerate(cache_list):
-                                if isinstance(tensor, torch.Tensor) and tensor.device != target_device:
-                                    cache_list[i] = tensor.to(target_device).contiguous()
+                    if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"):
+                        managed_buffer = raw_data
+                        try:
+                            buf_tensor = raw_data.tensor
+                            if buf_tensor.is_cuda:
+                                data = KVCacheTransferData.from_bytes_gpu(buf_tensor)
+                                raw_data.release()
+                                managed_buffer = None
+                            else:
+                                data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy()))
+                                data = self._clone_received_payload_tensors(data)
+                                raw_data.release()
+                                managed_buffer = None
+                        except Exception as e:
+                            logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e)
+                            if managed_buffer is not None:
+                                raw_data.release()
+                            return None, 0
+                    elif isinstance(raw_data, (bytes, bytearray)):
+                        data = KVCacheTransferData.from_bytes(raw_data)
+                    elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1:
+                        data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes())
+                    else:
+                        data = raw_data
 
+                    try:
+                        if isinstance(data, dict) and "layer_blocks" in data:
+                            layer_blocks = data["layer_blocks"]
+                            for cache_list in [
+                                layer_blocks.get("key_cache", []),
+                                layer_blocks.get("value_cache", []),
+                            ]:
+                                for i, tensor in enumerate(cache_list):
+                                    if not isinstance(tensor, torch.Tensor):
+                                        continue
+                                    if target_device is not None and tensor.device != target_device:
+                                        cache_list[i] = tensor.to(target_device).contiguous()
+                    finally:
+                        if managed_buffer is not None:
+                            managed_buffer.release()
+
+                    logger.info(
+                        "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms",
+                        request_id,
+                        size,
+                        elapsed,
+                        link_ms,
+                    )
                     return data, size
 
                 if time.time() - start_time > timeout:
                     logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
                     return None, 0
 
-                time.sleep(0.5)
+                time.sleep(poll_interval)
+                poll_interval = min(poll_interval * 2, max_poll_interval)
 
         except Exception as e:
             logger.error(f"Error receiving KV cache for {request_id}: {e}")
@@ -459,6 +907,16 @@ class OmniKVTransferManager:
             if hasattr(req, "sampling_params") and req.sampling_params is not None:
                 req.sampling_params.kv_metadata = data["metadata"]
 
+    @staticmethod
+    def _resolve_request_id(req: Any) -> str | None:
+        """Resolve the logical request ID used for KV transfer lookups."""
+        request_id = getattr(req, "request_id", None)
+        if request_id:
+            return request_id
+        if hasattr(req, "request_ids") and req.request_ids:
+            return req.request_ids[0]
+        return None
+
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
         """Receive KV cache and populate request object (legacy interface).
@@ -470,11 +928,11 @@ class OmniKVTransferManager:
         Returns:
             True if successful, False otherwise
         """
-        request_id = getattr(req, "request_id", None)
-        if not request_id and hasattr(req, "request_ids") and req.request_ids:
-            # Adaptation for new OmniDiffusionRequest which has list of prompts/ids
-            request_id = req.request_ids[0]
+        kv_sender_info = getattr(req, "kv_sender_info", None)
+        if kv_sender_info:
+            self.update_sender_info(kv_sender_info, sender_stage_id=self.recv_stages[0])
 
+        request_id = self._resolve_request_id(req)
         if not request_id:
             logger.warning("Request has no ID, cannot receive KV cache")
             return False
@@ -513,9 +971,7 @@ class OmniKVTransferManager:
 
         cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None)
         if cfg_ids and cfg_kv_collect_func:
-            request_id = getattr(req, "request_id", None) or (
-                req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None
-            )
+            request_id = self._resolve_request_id(req)
             try:
                 cfg_kvs = cfg_kv_collect_func(
                     request_id,

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -19,9 +19,17 @@ else:
 
 logger = get_connector_logger(__name__)
 
+# Reserve a separate port range for KV-transfer sockets so they do not
+# collide with request-forwarding endpoints that share the same base port.
+KV_TRANSFER_PORT_OFFSET = 100
+
 
 def initialize_connectors_from_config(
-    config_path: str | Path | None = None, default_shm_threshold: int = 65536
+    config_path: str | Path | None = None,
+    default_shm_threshold: int = 65536,
+    purpose: str = "request_forwarding",
+    caller_stage_id: int | str | None = None,
+    is_sender: bool | None = None,
 ) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]:
     """
     Initialize connectors from configuration file.
@@ -36,12 +44,20 @@ def initialize_connectors_from_config(
         return None, {}
 
     # create connectors from config
-    connectors = create_connectors_from_config(transfer_config.connectors)
+    connectors = create_connectors_from_config(
+        transfer_config.connectors,
+        purpose=purpose,
+        caller_stage_id=caller_stage_id,
+        is_sender=is_sender,
+    )
     return transfer_config, connectors
 
 
 def create_connectors_from_config(
     connectors_config: dict[tuple[str, str], ConnectorSpec],
+    purpose: str = "request_forwarding",
+    caller_stage_id: int | str | None = None,
+    is_sender: bool | None = None,
 ) -> dict[tuple[str, str], OmniConnectorBase]:
     """
     Create connectors from config.
@@ -52,12 +68,59 @@ def create_connectors_from_config(
     Returns:
         A dictionary of connectors.
     """
+    purpose_port_offsets = {
+        "request_forwarding": 0,
+        "kv_transfer": KV_TRANSFER_PORT_OFFSET,
+    }
+    port_offset = purpose_port_offsets.get(purpose, 0)
+    orchestrator_port_offset = 200
+
     connectors = {}
     for edge_key, connector_spec in connectors_config.items():
+        from_stage, to_stage = edge_key
         try:
-            connector = OmniConnectorFactory.create_connector(connector_spec)
+            if connector_spec.name == "MooncakeTransferEngineConnector":
+                extra = dict(connector_spec.extra) if connector_spec.extra else {}
+                base_port = extra.get("zmq_port", 50051)
+                try:
+                    stage_offset = int(from_stage)
+                except (TypeError, ValueError):
+                    stage_offset = 0
+
+                if str(caller_stage_id) == "orchestrator":
+                    adjusted_port = base_port + orchestrator_port_offset + stage_offset
+                else:
+                    adjusted_port = base_port + port_offset + stage_offset
+                extra["zmq_port"] = adjusted_port
+
+                if is_sender is not None:
+                    extra["role"] = "sender" if is_sender else "receiver"
+                    if not is_sender:
+                        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
+                        extra.setdefault("sender_zmq_port", adjusted_port)
+                elif caller_stage_id is not None:
+                    caller_str = str(caller_stage_id)
+                    if caller_str == from_stage:
+                        extra["role"] = "sender"
+                    elif caller_str == to_stage:
+                        extra["role"] = "receiver"
+                        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
+                        extra.setdefault("sender_zmq_port", adjusted_port)
+                    else:
+                        extra["role"] = "sender"
+                else:
+                    extra["role"] = extra.get("role", "auto")
+
+                connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra))
+            else:
+                connector = OmniConnectorFactory.create_connector(connector_spec)
             connectors[edge_key] = connector
-            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+            logger.info(
+                "Created connector for %s -> %s: %s",
+                from_stage,
+                to_stage,
+                type(connector).__name__,
+            )
         except Exception as e:
             raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e
 
@@ -289,7 +352,11 @@ def initialize_orchestrator_connectors(
     else:
         default_shm_threshold = max(0, shm_threshold_bytes)
     transfer_config, connectors = initialize_connectors_from_config(
-        config_path, default_shm_threshold=default_shm_threshold
+        config_path,
+        default_shm_threshold=default_shm_threshold,
+        purpose="request_forwarding",
+        caller_stage_id="orchestrator",
+        is_sender=True,
     )
     return transfer_config, connectors
 
@@ -316,6 +383,7 @@ def get_stage_connector_config(
 def build_stage_connectors(
     stage_id: int,
     connectors_config: dict[str, Any],
+    purpose: str = "request_forwarding",
 ) -> dict[tuple[str, str], Any] | None:
     """Instantiate OmniConnectors for a stage based on config."""
     if not connectors_config:
@@ -352,7 +420,12 @@ def build_stage_connectors(
 
     try:
         # Use unified connector creation logic
-        connectors = create_connectors_from_config(stage_connector_specs)
+        connectors = create_connectors_from_config(
+            stage_connector_specs,
+            purpose=purpose,
+            caller_stage_id=stage_id,
+            is_sender=False,
+        )
     except Exception as exc:  # pragma: no cover - defensive logging
         # Fail fast so the stage does not start with missing connectors.
         logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)

vllm_omni/engine/async_omni_engine.py

@@ -395,15 +395,17 @@ class AsyncOmniEngine:
                         proc=proc,
                         addresses=addresses,
                     )
+                    logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
+                    # Keep the stage-specific device visibility until vLLM
+                    # finishes starting all child processes.
+                    complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
+                    logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
                 finally:
                     if previous_visible_devices is None:
                         current_omni_platform.unset_device_control_env_var()
                     else:
                         current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
-            logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
-            complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
-            logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
             assert started_stage is not None
             return started_stage
         except Exception:

vllm_omni/engine/orchestrator.py

@@ -79,7 +79,7 @@ def build_engine_core_request_from_tokens(
         sampling_params=sampling_params,
         pooling_params=pooling_params,
         arrival_time=arrival_time,
-        lora_request=None,
+        lora_request=getattr(params, "lora_request", None),
         cache_salt=None,
         data_parallel_rank=None,
         prompt_embeds=prompt_embeds,
@@ -477,6 +477,30 @@ class Orchestrator:
             ),
         )
 
+    def _build_kv_sender_info(self, sender_stage_ids: list[int]) -> dict[int, dict[str, Any]] | None:
+        """Build per-request sender info for diffusion KV-transfer receivers."""
+        sender_infos: dict[int, dict[str, Any]] = {}
+        for sender_stage_id in dict.fromkeys(sender_stage_ids):
+            if sender_stage_id < 0 or sender_stage_id >= self.num_stages:
+                continue
+
+            sender_stage = self.stage_clients[sender_stage_id]
+            get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
+            if not callable(get_sender_info):
+                continue
+
+            sender_info = get_sender_info()
+            if not sender_info:
+                logger.warning(
+                    "[Orchestrator] Stage-%s has no KV sender info available",
+                    sender_stage_id,
+                )
+                continue
+
+            sender_infos[sender_stage_id] = sender_info
+
+        return sender_infos or None
+
     async def _forward_to_next_stage(
         self,
         req_id: str,
@@ -522,14 +546,22 @@ class Orchestrator:
                         req_id,
                     )
 
+            source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [stage_id])
+            kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids)
             if isinstance(diffusion_prompt, list):
                 await next_client.add_batch_request_async(
                     req_id,
                     diffusion_prompt,
                     params,
+                    kv_sender_info=kv_sender_info,
                 )
             else:
-                await next_client.add_request_async(req_id, diffusion_prompt, params)
+                await next_client.add_request_async(
+                    req_id,
+                    diffusion_prompt,
+                    params,
+                    kv_sender_info=kv_sender_info,
+                )
             req_state.stage_submit_ts[next_stage_id] = _time.time()
             return
 
@@ -731,7 +763,14 @@ class Orchestrator:
             params = req_state.sampling_params_list[next_stage_id]
 
             if next_client.stage_type == "diffusion":
-                await next_client.add_request_async(request_id, req_state.prompt, params)
+                source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [next_stage_id - 1])
+                kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids)
+                await next_client.add_request_async(
+                    request_id,
+                    req_state.prompt,
+                    params,
+                    kv_sender_info=kv_sender_info,
+                )
                 req_state.stage_submit_ts[next_stage_id] = _time.time()
                 continue